diff --git a/sim/app/w/[id]/components/workflow-block/components/sub-block/components/code.tsx b/sim/app/w/[id]/components/workflow-block/components/sub-block/components/code.tsx
index be0ff54aed..612301f62f 100644
--- a/sim/app/w/[id]/components/workflow-block/components/sub-block/components/code.tsx
+++ b/sim/app/w/[id]/components/workflow-block/components/sub-block/components/code.tsx
@@ -1,6 +1,6 @@
 import { useEffect, useRef, useState } from 'react'
 import type { ReactElement } from 'react'
-import { SparklesIcon } from 'lucide-react'
+import { Wand2 } from 'lucide-react'
 import { highlight, languages } from 'prismjs'
 import 'prismjs/components/prism-javascript'
 import 'prismjs/themes/prism.css'
@@ -304,9 +304,9 @@ export function Code({
               onClick={isPromptVisible ? hidePromptInline : showPromptInline}
               disabled={isAiLoading || isAiStreaming}
               aria-label="Generate code with AI"
-              className="h-8 w-8 text-muted-foreground hover:text-foreground"
+              className="h-8 w-8 rounded-full bg-muted/80 hover:bg-muted shadow-sm hover:shadow text-muted-foreground hover:text-primary transition-all duration-200 border border-transparent hover:border-primary/20"
             >
-              <SparklesIcon className="h-4 w-4" />
+              <Wand2 className="h-4 w-4" />
             </Button>
           )}
 
diff --git a/sim/app/w/[id]/components/workflow-block/components/sub-block/components/tool-input/components/code-editor/code-editor.tsx b/sim/app/w/[id]/components/workflow-block/components/sub-block/components/tool-input/components/code-editor/code-editor.tsx
index a123dd424f..af99f5b8d3 100644
--- a/sim/app/w/[id]/components/workflow-block/components/sub-block/components/tool-input/components/code-editor/code-editor.tsx
+++ b/sim/app/w/[id]/components/workflow-block/components/sub-block/components/tool-input/components/code-editor/code-editor.tsx
@@ -181,12 +181,12 @@ export function CodeEditor({
         }}
       >
         {code.length === 0 && placeholder && (
-          <div
-            className="absolute top-[12px] text-muted-foreground/50 select-none pointer-events-none"
-            style={{ left: `calc(${gutterWidth} + 12px)` }}
+          <pre
+            className="absolute top-[12px] text-muted-foreground/50 select-none pointer-events-none whitespace-pre-wrap overflow-visible"
+            style={{ left: `calc(${gutterWidth} + 12px)`, fontFamily: 'inherit', margin: 0 }}
           >
             {placeholder}
-          </div>
+          </pre>
         )}
 
         <Editor
diff --git a/sim/app/w/[id]/components/workflow-block/components/sub-block/components/tool-input/components/custom-tool-modal/custom-tool-modal.tsx b/sim/app/w/[id]/components/workflow-block/components/sub-block/components/tool-input/components/custom-tool-modal/custom-tool-modal.tsx
index b05861b6f1..1725225dd4 100644
--- a/sim/app/w/[id]/components/workflow-block/components/sub-block/components/tool-input/components/custom-tool-modal/custom-tool-modal.tsx
+++ b/sim/app/w/[id]/components/workflow-block/components/sub-block/components/tool-input/components/custom-tool-modal/custom-tool-modal.tsx
@@ -1,5 +1,5 @@
 import { useEffect, useMemo, useRef, useState } from 'react'
-import { Code, FileJson, SparklesIcon, X } from 'lucide-react'
+import { Code, FileJson, Wand2, X } from 'lucide-react'
 import { Button } from '@/components/ui/button'
 import {
   Dialog,
@@ -528,7 +528,7 @@ export function CustomToolModal({
                   <Button
                     variant="ghost"
                     size="icon"
-                    className="h-5 w-5 p-0 text-muted-foreground hover:text-foreground"
+                    className="h-5 w-5 p-0 rounded-full bg-muted/80 hover:bg-muted shadow-sm hover:shadow text-muted-foreground hover:text-primary transition-all duration-200 border border-transparent hover:border-primary/20"
                     onClick={() => {
                       logger.debug('Schema AI button clicked')
                       logger.debug(
@@ -542,7 +542,7 @@ export function CustomToolModal({
                     disabled={schemaGeneration.isLoading || schemaGeneration.isStreaming}
                     aria-label="Generate schema with AI"
                   >
-                    <SparklesIcon className="h-3 w-3" />
+                    <Wand2 className="h-3 w-3" />
                   </Button>
                 </div>
                 {schemaError &&
@@ -571,7 +571,7 @@ export function CustomToolModal({
     }
   }
 }`}
-                minHeight="340px"
+                minHeight="360px"
                 className={cn(
                   schemaError && !schemaGeneration.isStreaming ? 'border-red-500' : '',
                   (schemaGeneration.isLoading || schemaGeneration.isStreaming) &&
@@ -598,7 +598,7 @@ export function CustomToolModal({
                   <Button
                     variant="ghost"
                     size="icon"
-                    className="h-5 w-5 p-0 text-muted-foreground hover:text-foreground"
+                    className="h-5 w-5 p-0 rounded-full bg-muted/80 hover:bg-muted shadow-sm hover:shadow text-muted-foreground hover:text-primary transition-all duration-200 border border-transparent hover:border-primary/20"
                     onClick={() => {
                       logger.debug('Code AI button clicked')
                       logger.debug(
@@ -612,7 +612,7 @@ export function CustomToolModal({
                     disabled={codeGeneration.isLoading || codeGeneration.isStreaming}
                     aria-label="Generate code with AI"
                   >
-                    <SparklesIcon className="h-3 w-3" />
+                    <Wand2 className="h-3 w-3" />
                   </Button>
                 </div>
                 {codeError &&
@@ -626,7 +626,7 @@ export function CustomToolModal({
                   onChange={handleFunctionCodeChange}
                   language="javascript"
                   placeholder={`// This code will be executed when the tool is called. You can use environment variables with {{VARIABLE_NAME}}.`}
-                  minHeight="340px"
+                  minHeight="360px"
                   className={cn(
                     codeError && !codeGeneration.isStreaming ? 'border-red-500' : '',
                     (codeGeneration.isLoading || codeGeneration.isStreaming) &&