fix: use InternalContext for debug breakpoints in llamaindex runtime

the breakpoint wrapper was calling wait_for_event on the workflow-level
ExternalContext instead of a per-step InternalContext. this caused
ContextStateError on every uipath debug run (uipath run was unaffected
because it skips breakpoint injection).

fix: create an InternalContext via Context._create_internal(workflow)
inside the wrapper, which runs within the step worker where the
framework has already set up the required context variables.

also adds a pexpect-based integration test (testcases/debug-breakpoints)
that exercises single/multiple breakpoints, step mode, and quit against
the debug-agent sample.

Author: radu-mocanu

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: