diff --git a/joi/src/core/Badge/index.tsx b/joi/src/core/Badge/index.tsx
index d9b04fd2b..5aeb19631 100644
--- a/joi/src/core/Badge/index.tsx
+++ b/joi/src/core/Badge/index.tsx
@@ -49,4 +49,4 @@ const Badge = ({ className, theme, size, variant, ...props }: BadgeProps) => {
   )
 }
 
-export { Badge }
+export { Badge, badgeVariants }
diff --git a/web/containers/ModelDropdown/index.tsx b/web/containers/ModelDropdown/index.tsx
index 46154ce34..daef58e16 100644
--- a/web/containers/ModelDropdown/index.tsx
+++ b/web/containers/ModelDropdown/index.tsx
@@ -11,7 +11,7 @@ import Image from 'next/image'
 
 import { EngineConfig, InferenceEngine } from '@janhq/core'
 import {
-  Badge,
+  badgeVariants,
   Button,
   Input,
   ScrollArea,
@@ -318,11 +318,13 @@ const ModelDropdown = ({
     >
       <div className="flex [&>div]:w-full" ref={setToggle}>
         {chatInputMode ? (
-          <Badge
+          <button
             data-testid="model-selector-badge"
-            theme="secondary"
-            variant={open ? 'solid' : 'outline'}
             className={twMerge(
+              badgeVariants({
+                theme: 'secondary',
+                variant: open ? 'solid' : 'outline',
+              }),
               'inline-block max-w-[200px] cursor-pointer overflow-hidden text-ellipsis whitespace-nowrap',
               open && 'border border-transparent'
             )}
@@ -335,7 +337,7 @@ const ModelDropdown = ({
             >
               {selectedModel?.name || 'Select a model'}
             </span>
-          </Badge>
+          </button>
         ) : (
           <Input
             value={selectedModel?.name || ''}
diff --git a/web/screens/Thread/ThreadCenterPanel/ChatInput/index.tsx b/web/screens/Thread/ThreadCenterPanel/ChatInput/index.tsx
index 23b137be8..44145330f 100644
--- a/web/screens/Thread/ThreadCenterPanel/ChatInput/index.tsx
+++ b/web/screens/Thread/ThreadCenterPanel/ChatInput/index.tsx
@@ -3,7 +3,13 @@ import { useEffect, useRef, useState } from 'react'
 
 import { InferenceEngine } from '@janhq/core'
 
-import { TextArea, Button, Tooltip, useClickOutside, Badge } from '@janhq/joi'
+import {
+  TextArea,
+  Button,
+  Tooltip,
+  useClickOutside,
+  badgeVariants,
+} from '@janhq/joi'
 import { useAtom, useAtomValue } from 'jotai'
 import {
   FileTextIcon,
@@ -350,16 +356,17 @@ const ChatInput = () => {
         >
           <div className="flex items-center gap-x-2">
             <ModelDropdown chatInputMode />
-            <Badge
-              theme="secondary"
+            <button
               className={twMerge(
+                badgeVariants({
+                  theme: 'secondary',
+                  variant:
+                    activeTabThreadRightPanel === 'model' ? 'solid' : 'outline',
+                }),
                 'flex cursor-pointer items-center gap-x-1',
                 activeTabThreadRightPanel === 'model' &&
                   'border border-transparent'
               )}
-              variant={
-                activeTabThreadRightPanel === 'model' ? 'solid' : 'outline'
-              }
               onClick={() => {
                 // TODO @faisal: should be refactor later and better experience beetwen tab and toggle button
                 if (showRightPanel && activeTabThreadRightPanel !== 'model') {
@@ -384,7 +391,7 @@ const ChatInput = () => {
                 size={16}
                 className="flex-shrink-0 cursor-pointer text-[hsla(var(--text-secondary))]"
               />
-            </Badge>
+            </button>
           </div>
           {selectedModel && (
             <Button