diff --git a/client/src/pages/platform/workflow-editor/components/datapills/DataPillPanelBodyPropertiesItem.tsx b/client/src/pages/platform/workflow-editor/components/datapills/DataPillPanelBodyPropertiesItem.tsx
index 3bd7df49731..0872661fa7e 100644
--- a/client/src/pages/platform/workflow-editor/components/datapills/DataPillPanelBodyPropertiesItem.tsx
+++ b/client/src/pages/platform/workflow-editor/components/datapills/DataPillPanelBodyPropertiesItem.tsx
@@ -1,4 +1,4 @@
-import {Button} from '@/components/ui/button';
+import Button from '@/components/Button/Button';
 import DataPill from '@/pages/platform/workflow-editor/components/datapills/DataPill';
 import getNestedObject from '@/pages/platform/workflow-editor/utils/getNestedObject';
 import {NodeDataType, PropertyAllType} from '@/shared/types';
diff --git a/client/src/pages/platform/workflow-editor/components/properties/ArrayProperty.tsx b/client/src/pages/platform/workflow-editor/components/properties/ArrayProperty.tsx
index 31a7b9ecb6e..b5501da5666 100644
--- a/client/src/pages/platform/workflow-editor/components/properties/ArrayProperty.tsx
+++ b/client/src/pages/platform/workflow-editor/components/properties/ArrayProperty.tsx
@@ -1,4 +1,4 @@
-import {Button} from '@/components/ui/button';
+import Button from '@/components/Button/Button';
 import ArrayPropertyItem from '@/pages/platform/workflow-editor/components/properties/components/ArrayPropertyItem';
 import SubPropertyPopover from '@/pages/platform/workflow-editor/components/properties/components/SubPropertyPopover';
 import {useWorkflowEditor} from '@/pages/platform/workflow-editor/providers/workflowEditorProvider';
@@ -391,16 +391,14 @@ const ArrayProperty = ({onDeleteClick, parentArrayItems, path, property}: ArrayP
                 />
             ) : (
                 <Button
-                    className="mt-3 rounded-sm bg-gray-100 text-xs font-medium hover:bg-gray-200"
+                    className="mt-3 rounded-sm"
+                    icon={<PlusIcon />}
                     key={`${path}_${name}_addPropertyPopoverButton`}
+                    label={property.placeholder || 'Add array item'}
                     onClick={handleAddItemClick}
                     size="sm"
-                    variant="ghost"
-                >
-                    <PlusIcon className="size-4" />
-
-                    {property.placeholder || 'Add array item'}
-                </Button>
+                    variant="secondary"
+                />
             )}
         </Fragment>
     );
diff --git a/client/src/pages/platform/workflow-editor/components/properties/components/DeletePropertyButton.tsx b/client/src/pages/platform/workflow-editor/components/properties/components/DeletePropertyButton.tsx
index 07314fc2a05..9273c33d0f2 100644
--- a/client/src/pages/platform/workflow-editor/components/properties/components/DeletePropertyButton.tsx
+++ b/client/src/pages/platform/workflow-editor/components/properties/components/DeletePropertyButton.tsx
@@ -1,4 +1,4 @@
-import {Button} from '@/components/ui/button';
+import Button from '@/components/Button/Button';
 import {Tooltip, TooltipContent, TooltipTrigger} from '@/components/ui/tooltip';
 import {TooltipPortal} from '@radix-ui/react-tooltip';
 import {XIcon} from 'lucide-react';
@@ -11,9 +11,7 @@ interface DeletePropertyButtonProps {
 const DeletePropertyButton = ({onClick, propertyName}: DeletePropertyButtonProps) => (
     <Tooltip>
         <TooltipTrigger asChild>
-            <Button className="size-auto p-0.5" onClick={() => onClick()} size="icon" variant="ghost">
-                <XIcon className="size-4 cursor-pointer text-destructive" />
-            </Button>
+            <Button icon={<XIcon />} onClick={() => onClick()} size="iconXs" variant="destructiveGhost" />
         </TooltipTrigger>
 
         <TooltipPortal>
diff --git a/client/src/pages/platform/workflow-editor/components/properties/components/InputTypeSwitchButton.tsx b/client/src/pages/platform/workflow-editor/components/properties/components/InputTypeSwitchButton.tsx
index 76c3df22411..59c9b9cbb26 100644
--- a/client/src/pages/platform/workflow-editor/components/properties/components/InputTypeSwitchButton.tsx
+++ b/client/src/pages/platform/workflow-editor/components/properties/components/InputTypeSwitchButton.tsx
@@ -1,4 +1,4 @@
-import {Button} from '@/components/ui/button';
+import Button from '@/components/Button/Button';
 import {Tooltip, TooltipContent, TooltipTrigger} from '@/components/ui/tooltip';
 import {TooltipPortal} from '@radix-ui/react-tooltip';
 import {FormInputIcon, FunctionSquareIcon} from 'lucide-react';
@@ -6,13 +6,12 @@ import {FormInputIcon, FunctionSquareIcon} from 'lucide-react';
 const InputTypeSwitchButton = ({handleClick, mentionInput}: {handleClick: () => void; mentionInput: boolean}) => (
     <Tooltip>
         <TooltipTrigger asChild>
-            <Button className="size-auto p-0.5" onClick={handleClick} size="icon" variant="ghost">
-                {mentionInput ? (
-                    <FormInputIcon className="size-5 text-gray-600 hover:text-gray-800" />
-                ) : (
-                    <FunctionSquareIcon className="size-5 text-gray-600 hover:text-gray-800" />
-                )}
-            </Button>
+            <Button
+                icon={mentionInput ? <FormInputIcon /> : <FunctionSquareIcon />}
+                onClick={handleClick}
+                size="iconXs"
+                variant="secondary"
+            />
         </TooltipTrigger>
 
         <TooltipPortal>
diff --git a/client/src/pages/platform/workflow-editor/components/properties/components/PropertyComboBox.tsx b/client/src/pages/platform/workflow-editor/components/properties/components/PropertyComboBox.tsx
index 1bd97f74c3b..5f4316f67b7 100644
--- a/client/src/pages/platform/workflow-editor/components/properties/components/PropertyComboBox.tsx
+++ b/client/src/pages/platform/workflow-editor/components/properties/components/PropertyComboBox.tsx
@@ -1,6 +1,6 @@
+import Button from '@/components/Button/Button';
 import LoadingIcon from '@/components/LoadingIcon';
 import RequiredMark from '@/components/RequiredMark';
-import {Button} from '@/components/ui/button';
 import {Command, CommandEmpty, CommandGroup, CommandInput, CommandItem, CommandList} from '@/components/ui/command';
 import {Label} from '@/components/ui/label';
 import {Popover, PopoverContent, PopoverTrigger} from '@/components/ui/popover';
diff --git a/client/src/pages/platform/workflow-editor/components/properties/components/SubPropertyPopover.tsx b/client/src/pages/platform/workflow-editor/components/properties/components/SubPropertyPopover.tsx
index 2b52706565c..1f7b60a7383 100644
--- a/client/src/pages/platform/workflow-editor/components/properties/components/SubPropertyPopover.tsx
+++ b/client/src/pages/platform/workflow-editor/components/properties/components/SubPropertyPopover.tsx
@@ -1,4 +1,4 @@
-import {Button} from '@/components/ui/button';
+import Button from '@/components/Button/Button';
 import {Popover, PopoverContent, PopoverTrigger} from '@/components/ui/popover';
 import PropertySelect from '@/pages/platform/workflow-editor/components/properties/components/PropertySelect';
 import PropertyInput from '@/pages/platform/workflow-editor/components/properties/components/property-input/PropertyInput';
@@ -46,14 +46,12 @@ const SubPropertyPopover = ({
         <Popover>
             <PopoverTrigger asChild>
                 <Button
-                    className="mt-3 rounded-sm bg-gray-100 text-xs font-medium hover:bg-gray-200"
+                    className="mt-3 rounded-sm"
+                    icon={<PlusIcon />}
+                    label={buttonLabel || `Add ${array ? 'array item' : 'object property'}`}
                     size="sm"
-                    variant="ghost"
-                >
-                    <PlusIcon className="size-4" />
-
-                    {buttonLabel || `Add ${array ? 'array item' : 'object property'}`}
-                </Button>
+                    variant="secondary"
+                />
             </PopoverTrigger>
 
             <PopoverContent className="min-w-sub-property-popover-width space-y-4 p-4">
@@ -124,14 +122,7 @@ const SubPropertyPopover = ({
 
                 <footer className="flex items-center justify-end space-x-2">
                     <PopoverClose asChild>
-                        <Button
-                            className="cursor-pointer"
-                            disabled={!array && !newPropertyName}
-                            onClick={handleClick}
-                            size="sm"
-                        >
-                            Add
-                        </Button>
+                        <Button disabled={!array && !newPropertyName} label="Add" onClick={handleClick} size="sm" />
                     </PopoverClose>
                 </footer>
             </PopoverContent>
diff --git a/client/src/pages/platform/workflow-editor/components/properties/components/property-code-editor/PropertyCodeEditor.tsx b/client/src/pages/platform/workflow-editor/components/properties/components/property-code-editor/PropertyCodeEditor.tsx
index 9bee8b9ad82..6d32d026c33 100644
--- a/client/src/pages/platform/workflow-editor/components/properties/components/property-code-editor/PropertyCodeEditor.tsx
+++ b/client/src/pages/platform/workflow-editor/components/properties/components/property-code-editor/PropertyCodeEditor.tsx
@@ -1,5 +1,5 @@
+import Button from '@/components/Button/Button';
 import RequiredMark from '@/components/RequiredMark';
-import {Button} from '@/components/ui/button';
 import {Label} from '@/components/ui/label';
 import {Tooltip, TooltipContent, TooltipTrigger} from '@/components/ui/tooltip';
 import PropertyCodeEditorSheet from '@/pages/platform/workflow-editor/components/properties/components/property-code-editor/PropertyCodeEditorSheet';
@@ -85,12 +85,11 @@ const PropertyCodeEditor = forwardRef<HTMLButtonElement, PropertyCodeEditorProps
 
                             <Button
                                 className="ml-10 flex-1 rounded-l-none"
+                                label="Open Code Editor"
                                 onClick={() => setShowPropertyCodeEditorSheet(true)}
                                 ref={ref}
                                 variant="outline"
-                            >
-                                Open Code Editor
-                            </Button>
+                            />
 
                             {error && (
                                 <div className="pointer-events-none absolute inset-y-0 right-0 flex items-center pr-3">
diff --git a/client/src/pages/platform/workflow-editor/components/properties/components/property-code-editor/PropertyCodeEditorSheet.tsx b/client/src/pages/platform/workflow-editor/components/properties/components/property-code-editor/PropertyCodeEditorSheet.tsx
index a3fbd483668..a04870cc3f6 100644
--- a/client/src/pages/platform/workflow-editor/components/properties/components/property-code-editor/PropertyCodeEditorSheet.tsx
+++ b/client/src/pages/platform/workflow-editor/components/properties/components/property-code-editor/PropertyCodeEditorSheet.tsx
@@ -1,3 +1,4 @@
+import Button from '@/components/Button/Button';
 import LoadingDots from '@/components/LoadingDots';
 import {
     AlertDialog,
@@ -9,7 +10,6 @@ import {
     AlertDialogHeader,
     AlertDialogTitle,
 } from '@/components/ui/alert-dialog';
-import {Button} from '@/components/ui/button';
 import {ResizableHandle, ResizablePanel, ResizablePanelGroup} from '@/components/ui/resizable';
 import {Sheet, SheetCloseButton, SheetContent, SheetHeader, SheetTitle} from '@/components/ui/sheet';
 import {Tooltip, TooltipContent, TooltipTrigger} from '@/components/ui/tooltip';
@@ -107,15 +107,13 @@ const PropertyCodeEditorSheet = ({
                             <Tooltip>
                                 <TooltipTrigger asChild>
                                     <Button
-                                        className="[&_svg]:size-5"
                                         disabled={!dirty}
+                                        icon={<SaveIcon />}
                                         onClick={() => onChange(newValue)}
                                         size="icon"
                                         type="submit"
                                         variant="ghost"
-                                    >
-                                        <SaveIcon />
-                                    </Button>
+                                    />
                                 </TooltipTrigger>
 
                                 <TooltipContent>Save current workflow</TooltipContent>
@@ -126,14 +124,12 @@ const PropertyCodeEditorSheet = ({
                                     <TooltipTrigger asChild>
                                         <span tabIndex={0}>
                                             <Button
-                                                className="[&_svg]:size-5"
                                                 disabled={dirty}
+                                                icon={<PlayIcon className="text-success" />}
                                                 onClick={handleRunClick}
                                                 size="icon"
                                                 variant="ghost"
-                                            >
-                                                <PlayIcon className="text-success" />
-                                            </Button>
+                                            />
                                         </span>
                                     </TooltipTrigger>
 
@@ -143,15 +139,13 @@ const PropertyCodeEditorSheet = ({
 
                             {scriptIsRunning && (
                                 <Button
-                                    className="[&_svg]:size-5"
+                                    icon={<SquareIcon />}
                                     onClick={() => {
                                         // TODO
                                     }}
                                     size="icon"
                                     variant="destructive"
-                                >
-                                    <SquareIcon />
-                                </Button>
+                                />
                             )}
 
                             <CopilotButton parameters={{language}} source={Source.CODE_EDITOR} />
diff --git a/client/src/pages/platform/workflow-editor/components/properties/components/property-code-editor/PropertyCodeEditorSheetRightPanelConnections.tsx b/client/src/pages/platform/workflow-editor/components/properties/components/property-code-editor/PropertyCodeEditorSheetRightPanelConnections.tsx
index f7523c75363..e0de91e943e 100644
--- a/client/src/pages/platform/workflow-editor/components/properties/components/property-code-editor/PropertyCodeEditorSheetRightPanelConnections.tsx
+++ b/client/src/pages/platform/workflow-editor/components/properties/components/property-code-editor/PropertyCodeEditorSheetRightPanelConnections.tsx
@@ -1,4 +1,4 @@
-import {Button} from '@/components/ui/button';
+import Button from '@/components/Button/Button';
 import {Card, CardContent, CardHeader, CardTitle} from '@/components/ui/card';
 import PropertyCodeEditorSheetRightPanelConnectionsLabel from '@/pages/platform/workflow-editor/components/properties/components/property-code-editor/PropertyCodeEditorSheetRightPanelConnectionsLabel';
 import PropertyCodeEditorSheetRightPanelConnectionsPopover, {
@@ -199,7 +199,7 @@ const PropertyCodeEditorSheetRightPanelConnections = ({
                                 <div className="mt-6">
                                     <PropertyCodeEditorSheetRightPanelConnectionsPopover
                                         onSubmit={handleOnSubmit}
-                                        triggerNode={<Button>Add Component</Button>}
+                                        triggerNode={<Button label="Add Component" />}
                                     />
                                 </div>
                             </div>
@@ -210,13 +210,14 @@ const PropertyCodeEditorSheetRightPanelConnections = ({
                                 <div className="flex items-center pb-2">
                                     <span className="font-medium">Note</span>
 
-                                    <button
-                                        className="ml-auto p-0"
+                                    <Button
+                                        className="ml-auto hover:bg-transparent active:bg-transparent active:text-content-neutral-primary"
+                                        icon={<XIcon aria-hidden="true" />}
                                         onClick={() => setShowConnectionNote(false)}
+                                        size="iconXs"
                                         title="Close the note"
-                                    >
-                                        <XIcon aria-hidden="true" className="size-4 cursor-pointer" />
-                                    </button>
+                                        variant="ghost"
+                                    />
                                 </div>
 
                                 <p className="text-sm text-gray-800">
diff --git a/client/src/pages/platform/workflow-editor/components/properties/components/property-code-editor/PropertyCodeEditorSheetRightPanelConnectionsLabel.tsx b/client/src/pages/platform/workflow-editor/components/properties/components/property-code-editor/PropertyCodeEditorSheetRightPanelConnectionsLabel.tsx
index a9bf1d5673a..51d29f83d2e 100644
--- a/client/src/pages/platform/workflow-editor/components/properties/components/property-code-editor/PropertyCodeEditorSheetRightPanelConnectionsLabel.tsx
+++ b/client/src/pages/platform/workflow-editor/components/properties/components/property-code-editor/PropertyCodeEditorSheetRightPanelConnectionsLabel.tsx
@@ -1,5 +1,5 @@
+import Button from '@/components/Button/Button';
 import RequiredMark from '@/components/RequiredMark';
-import {Button} from '@/components/ui/button';
 import {Label} from '@/components/ui/label';
 import {Tooltip, TooltipContent, TooltipTrigger} from '@/components/ui/tooltip';
 import {ComponentConnection} from '@/shared/middleware/platform/configuration';
@@ -38,9 +38,13 @@ const PropertyCodeEditorSheetRightPanelConnectionsLabel = ({
                 </Tooltip>
             </div>
 
-            <Button className="px-0 text-destructive" onClick={onRemoveClick} size="sm" variant="link">
-                Remove
-            </Button>
+            <Button
+                className="px-0 text-content-destructive-primary hover:text-content-destructive-primary active:text-content-destructive-primary"
+                label="Remove"
+                onClick={onRemoveClick}
+                size="sm"
+                variant="link"
+            />
         </div>
     );
 };
diff --git a/client/src/pages/platform/workflow-editor/components/properties/components/property-code-editor/PropertyCodeEditorSheetRightPanelConnectionsPopover.tsx b/client/src/pages/platform/workflow-editor/components/properties/components/property-code-editor/PropertyCodeEditorSheetRightPanelConnectionsPopover.tsx
index f852dccfccd..75775ca2286 100644
--- a/client/src/pages/platform/workflow-editor/components/properties/components/property-code-editor/PropertyCodeEditorSheetRightPanelConnectionsPopover.tsx
+++ b/client/src/pages/platform/workflow-editor/components/properties/components/property-code-editor/PropertyCodeEditorSheetRightPanelConnectionsPopover.tsx
@@ -1,5 +1,5 @@
+import Button from '@/components/Button/Button';
 import ComboBox from '@/components/ComboBox';
-import {Button} from '@/components/ui/button';
 import {Form, FormControl, FormField, FormItem, FormLabel, FormMessage} from '@/components/ui/form';
 import {Input} from '@/components/ui/input';
 import {Popover, PopoverContent, PopoverTrigger} from '@/components/ui/popover';
@@ -46,13 +46,7 @@ const PropertyCodeEditorSheetRightPanelConnectionsPopover = ({
     return (
         <Popover onOpenChange={setOpen} open={open}>
             <PopoverTrigger asChild>
-                {triggerNode ? (
-                    triggerNode
-                ) : (
-                    <Button size="sm" variant="secondary">
-                        Add Component
-                    </Button>
-                )}
+                {triggerNode ? triggerNode : <Button label="Add Component" size="sm" variant="secondary" />}
             </PopoverTrigger>
 
             <PopoverContent align="end" className="min-w-property-code-editor-sheet-connections-sheet-width">
@@ -133,7 +127,7 @@ const PropertyCodeEditorSheetRightPanelConnectionsPopover = ({
                         </main>
 
                         <footer className="flex items-center justify-end space-x-2">
-                            <Button type="submit">Add</Button>
+                            <Button label="Add" type="submit" />
                         </footer>
                     </form>
                 </Form>
diff --git a/client/src/pages/platform/workflow-editor/components/properties/components/property-code-editor/PropertyCodeEditorSheetRightPanelConnectionsSelect.tsx b/client/src/pages/platform/workflow-editor/components/properties/components/property-code-editor/PropertyCodeEditorSheetRightPanelConnectionsSelect.tsx
index 52fe0d74603..7dae8f2d439 100644
--- a/client/src/pages/platform/workflow-editor/components/properties/components/property-code-editor/PropertyCodeEditorSheetRightPanelConnectionsSelect.tsx
+++ b/client/src/pages/platform/workflow-editor/components/properties/components/property-code-editor/PropertyCodeEditorSheetRightPanelConnectionsSelect.tsx
@@ -1,4 +1,4 @@
-import {Button} from '@/components/ui/button';
+import Button from '@/components/Button/Button';
 import {Select, SelectContent, SelectItem, SelectTrigger, SelectValue} from '@/components/ui/select';
 import {useWorkflowEditor} from '@/pages/platform/workflow-editor/providers/workflowEditorProvider';
 import EnvironmentBadge from '@/shared/components/EnvironmentBadge';
@@ -95,13 +95,13 @@ const PropertyCodeEditorSheetRightPanelConnectionsSelect = ({
                     </SelectTrigger>
 
                     <Button
-                        className="mt-auto p-2"
+                        className="mt-auto"
+                        icon={<PlusIcon />}
                         onClick={() => setShowNewConnectionDialog(true)}
+                        size="icon"
                         title="Create a new connection"
                         variant="outline"
-                    >
-                        <PlusIcon className="size-5" />
-                    </Button>
+                    />
                 </div>
 
                 <SelectContent>
diff --git a/client/src/pages/platform/workflow-editor/components/properties/components/property-json-schema-builder/PropertyJsonSchemaBuilder.tsx b/client/src/pages/platform/workflow-editor/components/properties/components/property-json-schema-builder/PropertyJsonSchemaBuilder.tsx
index 2a39192e8c8..5223fb67b51 100755
--- a/client/src/pages/platform/workflow-editor/components/properties/components/property-json-schema-builder/PropertyJsonSchemaBuilder.tsx
+++ b/client/src/pages/platform/workflow-editor/components/properties/components/property-json-schema-builder/PropertyJsonSchemaBuilder.tsx
@@ -1,6 +1,6 @@
+import Button from '@/components/Button/Button';
 import {SchemaRecordType} from '@/components/JsonSchemaBuilder/utils/types';
 import RequiredMark from '@/components/RequiredMark';
-import {Button} from '@/components/ui/button';
 import {Label} from '@/components/ui/label';
 import {Tooltip, TooltipContent, TooltipTrigger} from '@/components/ui/tooltip';
 import InputTypeSwitchButton from '@/pages/platform/workflow-editor/components/properties/components/InputTypeSwitchButton';
@@ -86,12 +86,11 @@ const PropertyJsonSchemaBuilder = forwardRef<HTMLButtonElement, PropertyJsonSche
 
                             <Button
                                 className="ml-10 flex-1 rounded-l-none"
+                                label="Open JSON Schema Builder"
                                 onClick={() => setShowPropertyJsonSchemaBuilder(true)}
                                 ref={ref}
                                 variant="outline"
-                            >
-                                Open JSON Schema Builder
-                            </Button>
+                            />
 
                             {error && (
                                 <div className="pointer-events-none absolute inset-y-0 right-0 flex items-center pr-3">
diff --git a/client/src/pages/platform/workflow-editor/components/properties/components/property-json-schema-builder/PropertyJsonSchemaBuilderSampleDataDialog.tsx b/client/src/pages/platform/workflow-editor/components/properties/components/property-json-schema-builder/PropertyJsonSchemaBuilderSampleDataDialog.tsx
index 962f31ac501..750094d7ad2 100644
--- a/client/src/pages/platform/workflow-editor/components/properties/components/property-json-schema-builder/PropertyJsonSchemaBuilderSampleDataDialog.tsx
+++ b/client/src/pages/platform/workflow-editor/components/properties/components/property-json-schema-builder/PropertyJsonSchemaBuilderSampleDataDialog.tsx
@@ -1,5 +1,5 @@
+import Button from '@/components/Button/Button';
 import {SchemaRecordType} from '@/components/JsonSchemaBuilder/utils/types';
-import {Button} from '@/components/ui/button';
 import {
     Dialog,
     DialogCloseButton,
@@ -94,7 +94,7 @@ const PropertyJsonSchemaBuilderSampleDataDialog = ({onChange}: {onChange?: (newS
     return (
         <Dialog onOpenChange={handleOpenChange} open={open}>
             <DialogTrigger asChild>
-                <Button variant="outline">Generate</Button>
+                <Button label="Generate" variant="outline" />
             </DialogTrigger>
 
             <DialogContent className="max-w-output-tab-sample-data-dialog-width">
@@ -130,9 +130,7 @@ const PropertyJsonSchemaBuilderSampleDataDialog = ({onChange}: {onChange?: (newS
                 </div>
 
                 <DialogFooter>
-                    <Button disabled={!curSchema} onClick={handleOnSubmit} type="submit">
-                        Generate
-                    </Button>
+                    <Button disabled={!curSchema} label="Generate" onClick={handleOnSubmit} type="submit" />
                 </DialogFooter>
             </DialogContent>
         </Dialog>
diff --git a/client/src/pages/platform/workflow-editor/components/properties/components/property-mentions-input/PropertyMentionsInputEditorSheet.tsx b/client/src/pages/platform/workflow-editor/components/properties/components/property-mentions-input/PropertyMentionsInputEditorSheet.tsx
index 734bd0a040a..a08f70b7d00 100644
--- a/client/src/pages/platform/workflow-editor/components/properties/components/property-mentions-input/PropertyMentionsInputEditorSheet.tsx
+++ b/client/src/pages/platform/workflow-editor/components/properties/components/property-mentions-input/PropertyMentionsInputEditorSheet.tsx
@@ -1,4 +1,4 @@
-import {Button} from '@/components/ui/button';
+import Button from '@/components/Button/Button';
 import {Sheet, SheetCloseButton, SheetContent, SheetHeader, SheetTitle, SheetTrigger} from '@/components/ui/sheet';
 import PropertyMentionsInputEditor from '@/pages/platform/workflow-editor/components/properties/components/property-mentions-input/PropertyMentionsInputEditor';
 import {
@@ -46,9 +46,7 @@ const PropertyMentionsInputEditorSheet = ({
         }}
     >
         <SheetTrigger asChild>
-            <Button className="size-auto p-0.5" size="icon" variant="ghost">
-                <MaximizeIcon className="size-4" />
-            </Button>
+            <Button icon={<MaximizeIcon />} size="iconXs" variant="secondary" />
         </SheetTrigger>
 
         <SheetContent className="flex w-11/12 flex-col gap-0 p-4 sm:max-w-screen-md">
diff --git a/client/src/pages/platform/workflow-editor/edges/BranchCaseLabel.tsx b/client/src/pages/platform/workflow-editor/edges/BranchCaseLabel.tsx
index 0b0d2e3cd88..2099308c5d9 100644
--- a/client/src/pages/platform/workflow-editor/edges/BranchCaseLabel.tsx
+++ b/client/src/pages/platform/workflow-editor/edges/BranchCaseLabel.tsx
@@ -1,4 +1,4 @@
-import {Button} from '@/components/ui/button';
+import Button from '@/components/Button/Button';
 import {Input} from '@/components/ui/input';
 import {BranchCaseType, NodeDataType} from '@/shared/types';
 import {useQueryClient} from '@tanstack/react-query';
@@ -216,18 +216,20 @@ export default function BranchCaseLabel({caseKey, edgeId, sourceY, targetX}: Bra
 
                         <Button
                             className="absolute right-2 top-1/2 size-4 -translate-y-1/2 cursor-pointer text-content-neutral-primary/50 hover:bg-transparent hover:text-content-neutral-primary [&_svg]:size-3"
+                            icon={
+                                isCaseKeyEditable ? <CheckIcon className="text-content-brand-primary" /> : <PenIcon />
+                            }
                             onClick={isCaseKeyEditable ? handleSaveCaseClick : handleEditCaseClick}
                             size="icon"
                             variant="ghost"
-                        >
-                            {isCaseKeyEditable ? <CheckIcon className="text-content-brand-primary" /> : <PenIcon />}
-                        </Button>
+                        />
                     </div>
                 )}
 
                 {!isDefaultCase && branchCases?.length > 1 && (
                     <Button
                         className="ml-1 size-auto cursor-pointer p-1 text-content-destructive/50 hover:bg-surface-destructive-secondary hover:text-content-destructive [&_svg]:size-4"
+                        icon={isDeleteConfirmationVisible ? <CheckIcon /> : <TrashIcon />}
                         onClick={
                             isDeleteConfirmationVisible
                                 ? () => handleDeleteCaseClick(caseKey)
@@ -235,20 +237,17 @@ export default function BranchCaseLabel({caseKey, edgeId, sourceY, targetX}: Bra
                         }
                         size="icon"
                         variant="ghost"
-                    >
-                        {isDeleteConfirmationVisible ? <CheckIcon /> : <TrashIcon />}
-                    </Button>
+                    />
                 )}
 
                 {isLastCase && (
                     <Button
                         className="ml-1 size-auto cursor-pointer p-1 text-content-neutral-primary/50 hover:bg-surface-neutral-primary-hover hover:text-content-neutral-primary [&_svg]:size-4"
+                        icon={<PlusIcon />}
                         onClick={handleCreateCaseClick}
                         size="icon"
                         variant="ghost"
-                    >
-                        <PlusIcon />
-                    </Button>
+                    />
                 )}
             </div>
         </EdgeLabelRenderer>