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

the breakpoint wrapper was reading the workflow-level ExternalContext via
self.context instead of the per-step InternalContext passed by the
workflows framework as a kwarg. wait_for_event requires InternalContext,
so debug mode always raised ContextStateError.

Author: radu-mocanu

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)."""
 
@@ -75,11 +71,13 @@ def make_wrapper(
     """
     Return a wrapped step function that pauses on breakpoints.
     """
+    step_config = original._step_config
+    ctx_param: str | None = step_config.context_parameter if step_config else None
 
     @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)
+        # Grab the internal context passed by the workflows framework to this step.
+        ctx: Context | None = kwargs.get(ctx_param) if ctx_param else None
 
         if isinstance(ctx, Context):
             bp_event = BreakpointEvent(

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: