fix: use step-internal context for debug breakpoints in llamaindex runtime

packages/uipath-llamaindex/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "uipath-llamaindex"
-version = "0.5.8"
+version = "0.5.9"
 description = "Python SDK that enables developers to build and deploy LlamaIndex agents to the UiPath Cloud Platform"
 readme = { file = "README.md", content-type = "text/markdown" }
 requires-python = ">=3.11"

packages/uipath-llamaindex/src/uipath_llamaindex/runtime/breakpoints.py

@@ -5,7 +5,7 @@
 from __future__ import annotations
 
 import functools
-from typing import Any, Protocol, cast
+from typing import Any, cast
 
 from workflows import Context, Workflow
 from workflows.decorators import StepFunction
@@ -17,10 +17,6 @@
 from uipath_llamaindex.runtime.schema import get_step_config
 
 
-class DebuggableWorkflow(Protocol):
-    context: Context | None = None
-
-
 class BreakpointEvent(InputRequiredEvent):
     """Event emitted when a breakpoint is hit (before step execution)."""
 
@@ -74,25 +70,28 @@ def make_wrapper(
 ) -> StepFunction[..., Any]:
     """
     Return a wrapped step function that pauses on breakpoints.
+
+    The wrapper creates an InternalContext via ``Context._create_internal``
+    to call ``wait_for_event``. This works because the wrapper executes
+    inside the step worker where the framework has already set the
+    ``StepWorkerStateContextVar``.
     """
 
     @functools.wraps(original)
     async def wrapper(self, *args: Any, **kwargs: Any) -> Any:
-        # Grab ctx from the workflow, as wired by UiPathLlamaIndexRuntime
-        ctx: Context | None = getattr(self, "context", None)
-
-        if isinstance(ctx, Context):
-            bp_event = BreakpointEvent(
-                breakpoint_node=step_name,
-                prefix=f"Breakpoint at {step_name}",
-            )
-            # Suspend until debugger resumes
-            await ctx.wait_for_event(
-                BreakpointResumeEvent,
-                waiter_event=bp_event,
-                waiter_id=f"bp_{step_name}",
-                timeout=None,
-            )
+        ctx = Context._create_internal(workflow=self)
+
+        bp_event = BreakpointEvent(
+            breakpoint_node=step_name,
+            prefix=f"Breakpoint at {step_name}",
+        )
+        # Suspend until debugger resumes
+        await ctx.wait_for_event(
+            BreakpointResumeEvent,
+            waiter_event=bp_event,
+            waiter_id=f"bp_{step_name}",
+            timeout=None,
+        )
 
         # Continue original step logic
         return await original(self, *args, **kwargs)

packages/uipath-llamaindex/src/uipath_llamaindex/runtime/runtime.py

@@ -3,7 +3,7 @@
 import asyncio
 import json
 import logging
-from typing import Any, AsyncGenerator, cast
+from typing import Any, AsyncGenerator
 from uuid import uuid4
 
 from llama_index.core.agent.workflow.workflow_events import (
@@ -40,7 +40,6 @@
 from uipath_llamaindex.runtime.breakpoints import (
     BreakpointEvent,
     BreakpointResumeEvent,
-    DebuggableWorkflow,
     inject_breakpoints,
 )
 from uipath_llamaindex.runtime.chat import UiPathChatMessagesMapper
@@ -144,11 +143,6 @@ async def _run_workflow(
 
         self._context = await self._load_context()
 
-        # Make the Context discoverable from inside steps
-        if self.debug_mode and self._context is not None:
-            debug_workflow = cast(DebuggableWorkflow, self.workflow)
-            debug_workflow.context = self._context
-
         if is_resuming:
             handler: WorkflowHandler = self.workflow.run(ctx=self._context)
             if workflow_input:

packages/uipath-llamaindex/testcases/debug-breakpoints/input.json

@@ -0,0 +1 @@
+{"ticket_id": "T-12345", "customer_message": "The payment system is broken!", "customer_tier": "premium"}

packages/uipath-llamaindex/testcases/debug-breakpoints/pyproject.toml

@@ -0,0 +1,20 @@
+[project]
+name = "debug-breakpoints-test"
+version = "0.0.1"
+description = "Test case for debug breakpoint functionality with LlamaIndex workflows"
+authors = [{ name = "UiPath", email = "test@uipath.com" }]
+requires-python = ">=3.11"
+dependencies = [
+    "uipath-llamaindex",
+    "pexpect>=4.9.0",
+    "pyte>=0.8.0",
+    "pytest>=8.0.0",
+    "pytest-asyncio>=0.24.0",
+]
+
+[tool.uv.sources]
+uipath-llamaindex = { path = "../../", editable = true }
+
+[tool.pytest.ini_options]
+asyncio_mode = "auto"
+asyncio_default_fixture_loop_scope = "function"

packages/uipath-llamaindex/testcases/debug-breakpoints/run.sh

@@ -0,0 +1,24 @@
+#!/bin/bash
+set -e
+
+SAMPLE_DIR="../../samples/debug-agent"
+
+echo "Copying agent files from debug-agent sample..."
+cp "$SAMPLE_DIR/main.py" main.py
+cp "$SAMPLE_DIR/llama_index.json" llama_index.json
+cp "$SAMPLE_DIR/uipath.json" uipath.json
+
+echo "Syncing dependencies..."
+uv sync
+
+echo "Authenticating with UiPath..."
+uv run uipath auth --client-id="$CLIENT_ID" --client-secret="$CLIENT_SECRET" --base-url="$BASE_URL"
+
+echo "Initializing the project..."
+uv run uipath init
+
+# Clear job key to force Console mode (not SignalR remote debugging)
+export UIPATH_JOB_KEY=""
+
+echo "=== Running debug breakpoint tests with pexpect ==="
+uv run pytest src/test_debug.py -v -s

packages/uipath-llamaindex/testcases/debug-breakpoints/src/assert.py

@@ -0,0 +1,4 @@
+# Debug breakpoints tests are run via pytest in test_debug.py
+# This file is a placeholder for the testcase runner convention.
+
+print("Debug breakpoints tests completed via pytest.")

packages/uipath-llamaindex/testcases/debug-breakpoints/src/test_debug.py

@@ -0,0 +1,221 @@
+"""
+Pexpect-based tests for uipath debug command with LlamaIndex workflows.
+
+Tests the interactive debugger functionality including:
+- Single breakpoint
+- Multiple breakpoints
+- Step mode (s command)
+- Quit debugger (q command)
+
+Regression test for: ContextStateError when using wait_for_event in
+breakpoint wrapper (the wrapper must use an InternalContext, not the
+workflow-level ExternalContext).
+
+See: https://github.com/UiPath/uipath-integrations-python/pull/274
+"""
+
+import re
+import sys
+from pathlib import Path
+from typing import Optional
+
+import pexpect
+import pytest
+
+# ---------------------------------------------------------------------------
+# Minimal PromptTest helper (mirrors uipath-langchain-python/testcases/common)
+# ---------------------------------------------------------------------------
+
+
+def strip_ansi(text: str) -> str:
+    return re.compile(r"\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])").sub("", text)
+
+
+class PromptTest:
+    def __init__(
+        self, command: str, test_name: str, prompt: str = "> ", timeout: int = 60
+    ):
+        self.command = command
+        self.test_name = test_name
+        self.prompt = prompt
+        self.timeout = timeout
+        self.child: Optional[pexpect.spawn] = None
+        self._log_handle = None
+        self._log_path = Path(f"{test_name}.log")
+
+    def start(self):
+        self.child = pexpect.spawn(self.command, encoding="utf-8", timeout=self.timeout)
+        self._log_handle = open(self._log_path, "w")
+        self.child.logfile_read = self._log_handle
+
+    def send_command(self, command: str, expect: Optional[str] = None):
+        self.child.expect(self.prompt)
+        self.child.sendline(command)
+        if expect:
+            self.child.expect(expect)
+
+    def expect_eof(self):
+        self.child.expect(pexpect.EOF, timeout=self.timeout)
+
+    def get_output(self) -> str:
+        if self._log_path.exists():
+            if self._log_handle:
+                self._log_handle.flush()
+            with open(self._log_path, "r", encoding="utf-8") as f:
+                return strip_ansi(f.read())
+        return ""
+
+    @property
+    def before(self) -> str:
+        return self.child.before if self.child else ""
+
+    def close(self):
+        if self._log_handle:
+            self._log_handle.close()
+            self._log_handle = None
+        if self.child:
+            self.child.close()
+            self.child = None
+
+
+# ---------------------------------------------------------------------------
+# Test configuration
+# ---------------------------------------------------------------------------
+
+COMMAND = "uv run uipath debug agent --file input.json"
+PROMPT = r"> "
+TIMEOUT = 60
+
+
+# ---------------------------------------------------------------------------
+# Tests
+# ---------------------------------------------------------------------------
+
+
+def test_single_breakpoint():
+    """Test setting and hitting a single breakpoint."""
+    test = PromptTest(
+        command=COMMAND,
+        test_name="debug_single_breakpoint",
+        prompt=PROMPT,
+        timeout=TIMEOUT,
+    )
+    try:
+        test.start()
+
+        test.send_command(
+            "b classify_category", expect=r"Breakpoint set at: classify_category"
+        )
+        test.send_command("c", expect=r"BREAKPOINT.*classify_category.*before")
+        test.send_command("c", expect=r"Debug session completed")
+
+        test.expect_eof()
+
+        output = test.get_output()
+        assert "ticket_id" in output, "Expected ticket_id in output"
+
+    except (pexpect.exceptions.TIMEOUT, pexpect.exceptions.EOF) as e:
+        print(f"\nERROR: {type(e).__name__}", file=sys.stderr)
+        print(f"\n--- Output before failure ---\n{test.before}", file=sys.stderr)
+        pytest.fail(f"Test failed: {e}")
+    finally:
+        test.close()
+
+
+def test_multiple_breakpoints():
+    """Test setting and hitting multiple breakpoints."""
+    test = PromptTest(
+        command=COMMAND,
+        test_name="debug_multiple_breakpoints",
+        prompt=PROMPT,
+        timeout=TIMEOUT,
+    )
+    try:
+        test.start()
+
+        test.send_command(
+            "b analyze_sentiment", expect=r"Breakpoint set at: analyze_sentiment"
+        )
+        test.send_command(
+            "b determine_priority", expect=r"Breakpoint set at: determine_priority"
+        )
+        test.send_command("c", expect=r"BREAKPOINT.*analyze_sentiment.*before")
+        test.send_command("c", expect=r"BREAKPOINT.*determine_priority.*before")
+        test.send_command("c", expect=r"Debug session completed")
+
+        test.expect_eof()
+
+        output = test.get_output()
+        breakpoint_count = output.count("BREAKPOINT")
+        assert breakpoint_count >= 2, (
+            f"Expected at least 2 breakpoints hit, got {breakpoint_count}"
+        )
+
+    except (pexpect.exceptions.TIMEOUT, pexpect.exceptions.EOF) as e:
+        print(f"\nERROR: {type(e).__name__}", file=sys.stderr)
+        print(f"\n--- Output before failure ---\n{test.before}", file=sys.stderr)
+        pytest.fail(f"Test failed: {e}")
+    finally:
+        test.close()
+
+
+def test_step_mode():
+    """Test step mode - breaks on every node."""
+    test = PromptTest(
+        command=COMMAND, test_name="debug_step_mode", prompt=PROMPT, timeout=TIMEOUT
+    )
+    try:
+        test.start()
+
+        # Step through all 9 workflow steps
+        test.send_command("s", expect=r"BREAKPOINT.*analyze_sentiment.*before")
+        test.send_command("s", expect=r"BREAKPOINT.*classify_category.*before")
+        test.send_command("s", expect=r"BREAKPOINT.*check_urgency.*before")
+        test.send_command("s", expect=r"BREAKPOINT.*determine_priority.*before")
+        test.send_command("s", expect=r"BREAKPOINT.*check_escalation.*before")
+        test.send_command("s", expect=r"BREAKPOINT.*route_to_department.*before")
+        test.send_command("s", expect=r"BREAKPOINT.*assign_queue.*before")
+        test.send_command("s", expect=r"BREAKPOINT.*generate_response.*before")
+        test.send_command("s", expect=r"BREAKPOINT.*finalize_ticket.*before")
+        test.send_command("s", expect=r"Debug session completed")
+
+        test.expect_eof()
+
+        output = test.get_output()
+        breakpoint_count = output.count("BREAKPOINT")
+        assert breakpoint_count >= 9, (
+            f"Expected at least 9 breakpoints in step mode, got {breakpoint_count}"
+        )
+
+    except (pexpect.exceptions.TIMEOUT, pexpect.exceptions.EOF) as e:
+        print(f"\nERROR: {type(e).__name__}", file=sys.stderr)
+        print(f"\n--- Output before failure ---\n{test.before}", file=sys.stderr)
+        pytest.fail(f"Test failed: {e}")
+    finally:
+        test.close()
+
+
+def test_quit_debugger():
+    """Test quitting the debugger early with 'q' command."""
+    test = PromptTest(
+        command=COMMAND, test_name="debug_quit", prompt=PROMPT, timeout=TIMEOUT
+    )
+    try:
+        test.start()
+
+        test.send_command("b check_urgency", expect=r"Breakpoint set at: check_urgency")
+        test.send_command("c", expect=r"BREAKPOINT.*check_urgency.*before")
+        test.send_command("q")
+
+        test.expect_eof()
+
+    except (pexpect.exceptions.TIMEOUT, pexpect.exceptions.EOF) as e:
+        print(f"\nERROR: {type(e).__name__}", file=sys.stderr)
+        print(f"\n--- Output before failure ---\n{test.before}", file=sys.stderr)
+        pytest.fail(f"Test failed: {e}")
+    finally:
+        test.close()
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])