diff --git a/src/frontend/src/components/flowToolbarComponent/components/playground-button.tsx b/src/frontend/src/components/flowToolbarComponent/components/playground-button.tsx
new file mode 100644
index 000000000..261dfe56f
--- /dev/null
+++ b/src/frontend/src/components/flowToolbarComponent/components/playground-button.tsx
@@ -0,0 +1,50 @@
+import ForwardedIconComponent from "@/components/genericIconComponent";
+import ShadTooltip from "@/components/shadTooltipComponent";
+import IOModal from "@/modals/IOModal/newModal";
+
+const PlaygroundButton = ({ hasIO, open, setOpen, canvasOpen }) => {
+  const PlayIcon = () => (
+    <ForwardedIconComponent name="Play" className="h-4 w-4 transition-all" />
+  );
+
+  const ButtonLabel = () => <span className="hidden md:block">Playground</span>;
+
+  const ActiveButton = () => (
+    <div
+      data-testid="playground-btn-flow-io"
+      className="playground-btn-flow-toolbar hover:bg-accent"
+    >
+      <PlayIcon />
+      <ButtonLabel />
+    </div>
+  );
+
+  const DisabledButton = () => (
+    <div
+      className="playground-btn-flow-toolbar cursor-not-allowed text-muted-foreground duration-150"
+      data-testid="playground-btn-flow"
+    >
+      <PlayIcon />
+      <ButtonLabel />
+    </div>
+  );
+
+  return hasIO ? (
+    <IOModal
+      open={open}
+      setOpen={setOpen}
+      disable={!hasIO}
+      canvasOpen={canvasOpen}
+    >
+      <ActiveButton />
+    </IOModal>
+  ) : (
+    <ShadTooltip content="Add a Chat Input or Chat Output to use the playground">
+      <div>
+        <DisabledButton />
+      </div>
+    </ShadTooltip>
+  );
+};
+
+export default PlaygroundButton;
diff --git a/src/frontend/src/components/flowToolbarComponent/index.tsx b/src/frontend/src/components/flowToolbarComponent/index.tsx
index 0784f5e98..fedc7f259 100644
--- a/src/frontend/src/components/flowToolbarComponent/index.tsx
+++ b/src/frontend/src/components/flowToolbarComponent/index.tsx
@@ -15,6 +15,7 @@ import { useShortcutsStore } from "../../stores/shortcuts";
 import { useStoreStore } from "../../stores/storeStore";
 import { classNames, isThereModal } from "../../utils/utils";
 import ForwardedIconComponent from "../genericIconComponent";
+import PlaygroundButton from "./components/playground-button";
 
 export default function FlowToolbar(): JSX.Element {
   const preventDefault = true;
@@ -125,38 +126,12 @@ export default function FlowToolbar(): JSX.Element {
         >
           <div className="flex gap-1.5">
             <div className="flex h-full w-full gap-1.5 rounded-sm transition-all">
-              {hasIO ? (
-                <IOModal
-                  open={open}
-                  setOpen={setOpen}
-                  disable={!hasIO}
-                  canvasOpen
-                >
-                  <div
-                    data-testid="playground-btn-flow-io"
-                    className="relative inline-flex h-8 w-full items-center justify-center gap-1.5 rounded px-3 py-1.5 text-sm font-semibold transition-all duration-500 ease-in-out hover:bg-accent"
-                  >
-                    <ForwardedIconComponent
-                      name="Play"
-                      className={"h-4 w-4 transition-all"}
-                    />
-                    <span className="hidden md:block">Playground</span>
-                  </div>
-                </IOModal>
-              ) : (
-                <ShadTooltip content="Add a Chat Input or Chat Output to use the playground">
-                  <div
-                    className={`relative inline-flex h-8 w-full cursor-not-allowed items-center justify-center gap-1 px-5 py-3 text-sm font-semibold text-muted-foreground transition-all duration-150 ease-in-out`}
-                    data-testid="playground-btn-flow"
-                  >
-                    <ForwardedIconComponent
-                      name="BotMessageSquareIcon"
-                      className={"h-5 w-5 transition-all"}
-                    />
-                    <span className="hidden md:block">Playground</span>
-                  </div>
-                </ShadTooltip>
-              )}
+              <PlaygroundButton
+                hasIO={hasIO}
+                open={open}
+                setOpen={setOpen}
+                canvasOpen
+              />
             </div>
             {ENABLE_API && (
               <>
diff --git a/src/frontend/src/style/applies.css b/src/frontend/src/style/applies.css
index e4089c842..42cadc8f6 100644
--- a/src/frontend/src/style/applies.css
+++ b/src/frontend/src/style/applies.css
@@ -1268,6 +1268,10 @@
   .toolbar-wrapper {
     @apply flex h-10 items-center gap-1 rounded-xl border border-border bg-background p-1 shadow-sm;
   }
+
+  .playground-btn-flow-toolbar {
+    @apply relative inline-flex h-8 w-full items-center justify-center gap-1.5 rounded px-3 py-1.5 text-sm font-semibold transition-all duration-500 ease-in-out;
+  }
 }
 
 /* Gradient background */