diff --git a/src/components/assistant/DoormileAI/AIMessage.jsx b/src/components/assistant/DoormileAI/AIMessage.jsx
index 5559eb3..c97addc 100644
--- a/src/components/assistant/DoormileAI/AIMessage.jsx
+++ b/src/components/assistant/DoormileAI/AIMessage.jsx
@@ -6,9 +6,9 @@ import { LuChevronDown } from 'react-icons/lu';
 import { HStack } from '@astryxdesign/core/HStack';
 import { VStack } from '@astryxdesign/core/VStack';
 import { Text } from '@astryxdesign/core/Text';
-import { IconButton } from '@astryxdesign/core/IconButton';
-import { Button } from '@/components/ui/button';
 import { ChatToolCalls } from '@astryxdesign/core/Chat';
+import { Button } from '@/components/ui/button';
+import { IconButton } from '@/components/ds';
 
 import { Spark, Metric, StatGrid, StateBlock, AnswerList } from './AIParts';
 import AIFlowStep from './AIFlowStep';
diff --git a/src/components/assistant/DoormileAI/AIPanel.jsx b/src/components/assistant/DoormileAI/AIPanel.jsx
index aed7d1b..93baf57 100644
--- a/src/components/assistant/DoormileAI/AIPanel.jsx
+++ b/src/components/assistant/DoormileAI/AIPanel.jsx
@@ -10,8 +10,8 @@ import { LuChevronDown, LuPanelRightOpen, LuPanelRightClose, LuHistory, LuTrash2
 import { HStack } from '@astryxdesign/core/HStack';
 import { VStack } from '@astryxdesign/core/VStack';
 import { Text } from '@astryxdesign/core/Text';
-import { IconButton } from '@astryxdesign/core/IconButton';
 import { DropdownMenu } from '@astryxdesign/core/DropdownMenu';
+import { IconButton } from '@/components/ds';
 
 import { OpenToast } from 'components/third-party/OpenToast';
 import { STATUS } from 'themes/dt/tokens';
diff --git a/src/components/ds/IconButton.jsx b/src/components/ds/IconButton.jsx
index 8410cdb..1f2553a 100644
--- a/src/components/ds/IconButton.jsx
+++ b/src/components/ds/IconButton.jsx
@@ -19,9 +19,16 @@ export const IconButton = React.forwardRef(
       lg: 'icon-lg',
     };
 
+    const renderIcon = () => {
+      if (!Icon) return null;
+      if (React.isValidElement(Icon)) return Icon;
+      const Comp = Icon;
+      return <Comp aria-hidden="true" />;
+    };
+
     const button = (
       <Button ref={ref} size={sizeMap[size] || 'icon'} aria-label={label} {...props}>
-        <Icon aria-hidden="true" />
+        {renderIcon()}
       </Button>
     );