diff --git a/src/cmdforge/gui/dialogs/step_dialog.py b/src/cmdforge/gui/dialogs/step_dialog.py
index 6fcab36..89120e7 100644
--- a/src/cmdforge/gui/dialogs/step_dialog.py
+++ b/src/cmdforge/gui/dialogs/step_dialog.py
@@ -3,12 +3,12 @@
 from PySide6.QtWidgets import (
     QDialog, QVBoxLayout, QFormLayout, QLineEdit,
     QComboBox, QPushButton, QHBoxLayout, QLabel,
-    QPlainTextEdit
+    QPlainTextEdit, QSplitter, QGroupBox, QTextEdit
 )
-from PySide6.QtCore import Qt
+from PySide6.QtCore import Qt, QThread, Signal
 
 from ...tool import PromptStep, CodeStep
-from ...providers import load_providers
+from ...providers import load_providers, call_provider
 from ...profiles import list_profiles
 
 
@@ -136,63 +136,142 @@ class PromptStepDialog(QDialog):
         )
 
 
-class CodeStepDialog(QDialog):
-    """Dialog for editing code steps."""
+class AIGenerateWorker(QThread):
+    """Background worker for AI code generation."""
+    finished = Signal(str)
+    error = Signal(str)
 
-    def __init__(self, parent, step: CodeStep = None):
+    def __init__(self, provider: str, prompt: str):
+        super().__init__()
+        self.provider = provider
+        self.prompt = prompt
+
+    def run(self):
+        try:
+            result = call_provider(self.provider, self.prompt)
+            if result.success:
+                self.finished.emit(result.text)
+            else:
+                self.error.emit(result.error or "Unknown error")
+        except Exception as e:
+            self.error.emit(str(e))
+
+
+class CodeStepDialog(QDialog):
+    """Dialog for editing code steps with AI assist."""
+
+    def __init__(self, parent, step: CodeStep = None, available_vars: list = None):
         super().__init__(parent)
         self.setWindowTitle("Edit Code Step" if step else "Add Code Step")
-        self.setMinimumSize(600, 500)
+        self.setMinimumSize(900, 600)
         self._step = step
+        self._available_vars = available_vars or ["input"]
+        self._worker = None
         self._setup_ui()
 
         if step:
             self._load_step(step)
 
     def _setup_ui(self):
-        """Set up the UI."""
+        """Set up the UI with code editor and AI assist panel."""
         layout = QVBoxLayout(self)
-        layout.setSpacing(16)
+        layout.setSpacing(12)
 
-        # Form
+        # Top: Output variable
         form = QFormLayout()
-        form.setSpacing(12)
-
-        # Output variable
+        form.setSpacing(8)
         self.output_input = QLineEdit()
         self.output_input.setPlaceholderText("result")
         self.output_input.setText("result")
         form.addRow("Output variable:", self.output_input)
-
         layout.addLayout(form)
 
-        # Code editor
-        code_label = QLabel("Code:")
-        layout.addWidget(code_label)
+        # Available variables display
+        vars_text = ", ".join(self._available_vars)
+        vars_label = QLabel(f"Available variables: {vars_text}")
+        vars_label.setStyleSheet("color: #718096; font-size: 11px;")
+        layout.addWidget(vars_label)
+
+        # Main splitter: Code editor | AI Assist
+        splitter = QSplitter(Qt.Horizontal)
+
+        # Left: Code editor
+        code_group = QGroupBox("Code")
+        code_layout = QVBoxLayout(code_group)
 
         self.code_input = QPlainTextEdit()
         self.code_input.setPlaceholderText(
             "# Python code here\n"
-            "# Access input with: input_text\n"
-            "# Access args with their variable names\n"
-            "# Access previous step outputs with their variable names\n\n"
-            "result = input_text.upper()"
+            "# Access input with: input\n"
+            "# Access args with their variable names\n\n"
+            "result = input.upper()"
         )
         font = self.code_input.font()
-        font.setFamily("Courier New, Consolas, monospace")
+        font.setFamily("Consolas, Monaco, monospace")
         self.code_input.setFont(font)
-        layout.addWidget(self.code_input, 1)
+        code_layout.addWidget(self.code_input)
 
-        # Help text
-        help_text = QLabel(
-            "The code will be executed with variables in scope. "
-            "The last expression or the output variable will be captured."
+        splitter.addWidget(code_group)
+
+        # Right: AI Assist panel
+        ai_group = QGroupBox("AI Assisted Code Generation")
+        ai_layout = QVBoxLayout(ai_group)
+        ai_layout.setSpacing(8)
+
+        # Provider selector
+        provider_layout = QHBoxLayout()
+        provider_layout.addWidget(QLabel("Provider:"))
+        self.ai_provider_combo = QComboBox()
+        providers = load_providers()
+        for provider in sorted(providers, key=lambda p: p.name):
+            self.ai_provider_combo.addItem(provider.name)
+        # Add common defaults if not present
+        for default in ["claude", "gpt", "mock"]:
+            if self.ai_provider_combo.findText(default) < 0:
+                self.ai_provider_combo.addItem(default)
+        self.ai_provider_combo.setMinimumWidth(150)
+        provider_layout.addWidget(self.ai_provider_combo)
+        provider_layout.addStretch()
+        ai_layout.addLayout(provider_layout)
+
+        # Prompt editor
+        prompt_label = QLabel("Describe what you want the code to do:")
+        ai_layout.addWidget(prompt_label)
+
+        self.ai_prompt_input = QPlainTextEdit()
+        self.ai_prompt_input.setPlaceholderText(
+            "Example: Parse the input as JSON and extract the 'name' field"
         )
-        help_text.setStyleSheet("color: #718096; font-size: 11px;")
-        help_text.setWordWrap(True)
-        layout.addWidget(help_text)
+        # Set default prompt template
+        self._set_default_prompt()
+        ai_layout.addWidget(self.ai_prompt_input, 2)
 
-        # Buttons
+        # Generate button
+        btn_layout = QHBoxLayout()
+        self.btn_generate = QPushButton("Generate Code")
+        self.btn_generate.clicked.connect(self._generate_code)
+        self.btn_generate.setMinimumHeight(32)
+        btn_layout.addStretch()
+        btn_layout.addWidget(self.btn_generate)
+        btn_layout.addStretch()
+        ai_layout.addLayout(btn_layout)
+
+        # Output/feedback area
+        feedback_label = QLabel("AI Response:")
+        ai_layout.addWidget(feedback_label)
+
+        self.ai_output = QTextEdit()
+        self.ai_output.setReadOnly(True)
+        self.ai_output.setPlaceholderText("AI response will appear here...")
+        self.ai_output.setMaximumHeight(100)
+        ai_layout.addWidget(self.ai_output, 1)
+
+        splitter.addWidget(ai_group)
+        splitter.setSizes([450, 450])
+
+        layout.addWidget(splitter, 1)
+
+        # Bottom: OK/Cancel buttons
         buttons = QHBoxLayout()
         buttons.addStretch()
 
@@ -207,6 +286,81 @@ class CodeStepDialog(QDialog):
 
         layout.addLayout(buttons)
 
+    def _set_default_prompt(self):
+        """Set the default AI prompt template."""
+        vars_formatted = ', '.join(f'"{{{v}}}"' for v in self._available_vars)
+        default_prompt = f"""Write inline Python code (NOT a function definition) according to my instruction.
+
+The code runs directly with variable substitution. Assign any available variables to a local variable first using triple quotes for multi-line content safety.
+
+Example:
+my_var = \"\"\"{{input}}\"\"\"
+result = my_var.upper()
+
+INSTRUCTION: [Describe what you want the code to do]
+
+CURRENT CODE:
+{{code}}
+
+AVAILABLE VARIABLES: {vars_formatted}
+
+IMPORTANT: Return ONLY executable inline Python code. No function definitions, no markdown fencing, no explanations - just the code."""
+        self.ai_prompt_input.setPlainText(default_prompt)
+
+    def _generate_code(self):
+        """Generate code using AI."""
+        prompt_template = self.ai_prompt_input.toPlainText().strip()
+        if not prompt_template:
+            self.ai_output.setHtml("<span style='color: #e53e3e;'>Please enter a prompt</span>")
+            return
+
+        # Replace {code} placeholder with current code
+        current_code = self.code_input.toPlainText().strip() or "# No code yet"
+        prompt = prompt_template.replace("{code}", current_code)
+
+        provider = self.ai_provider_combo.currentText()
+
+        # Disable button and show loading state
+        self.btn_generate.setEnabled(False)
+        self.btn_generate.setText("Generating...")
+        self.ai_output.setHtml(f"<span style='color: #718096;'>Calling {provider}...</span>")
+
+        # Start worker thread
+        self._worker = AIGenerateWorker(provider, prompt)
+        self._worker.finished.connect(self._on_generate_finished)
+        self._worker.error.connect(self._on_generate_error)
+        self._worker.start()
+
+    def _on_generate_finished(self, result: str):
+        """Handle successful AI generation."""
+        self.btn_generate.setEnabled(True)
+        self.btn_generate.setText("Generate Code")
+
+        # Clean up the result - strip markdown code fences if present
+        code = result.strip()
+        if code.startswith("```python"):
+            code = code[9:]
+        elif code.startswith("```"):
+            code = code[3:]
+        if code.endswith("```"):
+            code = code[:-3]
+        code = code.strip()
+
+        # Update code editor
+        self.code_input.setPlainText(code)
+
+        # Show success message
+        self.ai_output.setHtml(
+            f"<span style='color: #38a169;'>Code generated successfully!</span><br>"
+            f"<span style='color: #718096;'>Response length: {len(result)} chars</span>"
+        )
+
+    def _on_generate_error(self, error: str):
+        """Handle AI generation error."""
+        self.btn_generate.setEnabled(True)
+        self.btn_generate.setText("Generate Code")
+        self.ai_output.setHtml(f"<span style='color: #e53e3e;'>Error: {error}</span>")
+
     def _load_step(self, step: CodeStep):
         """Load step data into form."""
         self.output_input.setText(step.output_var)
diff --git a/src/cmdforge/gui/pages/tool_builder_page.py b/src/cmdforge/gui/pages/tool_builder_page.py
index bfed676..d05be40 100644
--- a/src/cmdforge/gui/pages/tool_builder_page.py
+++ b/src/cmdforge/gui/pages/tool_builder_page.py
@@ -261,6 +261,34 @@ class ToolBuilderPage(QWidget):
         del self._tool.arguments[idx]
         self._refresh_arguments()
 
+    def _get_available_vars(self, up_to_step: int = -1) -> list:
+        """Get available variables for a step.
+
+        Args:
+            up_to_step: Include variables from steps before this index.
+                       -1 means include all existing steps (for new steps).
+        """
+        vars_list = ["input"]
+
+        # Add argument variables
+        if self._tool and self._tool.arguments:
+            for arg in self._tool.arguments:
+                if arg.variable and arg.variable not in vars_list:
+                    vars_list.append(arg.variable)
+
+        # Add output variables from previous steps
+        if self._tool and self._tool.steps:
+            steps_to_check = self._tool.steps if up_to_step < 0 else self._tool.steps[:up_to_step]
+            for step in steps_to_check:
+                if hasattr(step, 'output_var') and step.output_var:
+                    # Handle comma-separated output vars
+                    for var in step.output_var.split(','):
+                        var = var.strip()
+                        if var and var not in vars_list:
+                            vars_list.append(var)
+
+        return vars_list
+
     def _add_prompt_step(self):
         """Add a prompt step."""
         from ..dialogs.step_dialog import PromptStepDialog
@@ -275,7 +303,8 @@ class ToolBuilderPage(QWidget):
     def _add_code_step(self):
         """Add a code step."""
         from ..dialogs.step_dialog import CodeStepDialog
-        dialog = CodeStepDialog(self)
+        available_vars = self._get_available_vars()
+        dialog = CodeStepDialog(self, available_vars=available_vars)
         if dialog.exec():
             step = dialog.get_step()
             if not self._tool:
@@ -297,7 +326,8 @@ class ToolBuilderPage(QWidget):
             dialog = PromptStepDialog(self, step)
         elif isinstance(step, CodeStep):
             from ..dialogs.step_dialog import CodeStepDialog
-            dialog = CodeStepDialog(self, step)
+            available_vars = self._get_available_vars(up_to_step=idx)
+            dialog = CodeStepDialog(self, step, available_vars=available_vars)
         else:
             return