Skip to content

Commit e96336f

Browse files
committed
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.
1 parent feccaf7 commit e96336f

8 files changed

Lines changed: 258 additions & 28 deletions

File tree

packages/uipath-llamaindex/pyproject.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
[project]
22
name = "uipath-llamaindex"
3-
version = "0.5.8"
3+
version = "0.5.9"
44
description = "Python SDK that enables developers to build and deploy LlamaIndex agents to the UiPath Cloud Platform"
55
readme = { file = "README.md", content-type = "text/markdown" }
66
requires-python = ">=3.11"

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

Lines changed: 19 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
from __future__ import annotations
66

77
import functools
8-
from typing import Any, Protocol, cast
8+
from typing import Any, cast
99

1010
from workflows import Context, Workflow
1111
from workflows.decorators import StepFunction
@@ -17,10 +17,6 @@
1717
from uipath_llamaindex.runtime.schema import get_step_config
1818

1919

20-
class DebuggableWorkflow(Protocol):
21-
context: Context | None = None
22-
23-
2420
class BreakpointEvent(InputRequiredEvent):
2521
"""Event emitted when a breakpoint is hit (before step execution)."""
2622

@@ -74,25 +70,28 @@ def make_wrapper(
7470
) -> StepFunction[..., Any]:
7571
"""
7672
Return a wrapped step function that pauses on breakpoints.
73+
74+
The wrapper creates an InternalContext via ``Context._create_internal``
75+
to call ``wait_for_event``. This works because the wrapper executes
76+
inside the step worker where the framework has already set the
77+
``StepWorkerStateContextVar``.
7778
"""
7879

7980
@functools.wraps(original)
8081
async def wrapper(self, *args: Any, **kwargs: Any) -> Any:
81-
# Grab ctx from the workflow, as wired by UiPathLlamaIndexRuntime
82-
ctx: Context | None = getattr(self, "context", None)
83-
84-
if isinstance(ctx, Context):
85-
bp_event = BreakpointEvent(
86-
breakpoint_node=step_name,
87-
prefix=f"Breakpoint at {step_name}",
88-
)
89-
# Suspend until debugger resumes
90-
await ctx.wait_for_event(
91-
BreakpointResumeEvent,
92-
waiter_event=bp_event,
93-
waiter_id=f"bp_{step_name}",
94-
timeout=None,
95-
)
82+
ctx = Context._create_internal(workflow=self)
83+
84+
bp_event = BreakpointEvent(
85+
breakpoint_node=step_name,
86+
prefix=f"Breakpoint at {step_name}",
87+
)
88+
# Suspend until debugger resumes
89+
await ctx.wait_for_event(
90+
BreakpointResumeEvent,
91+
waiter_event=bp_event,
92+
waiter_id=f"bp_{step_name}",
93+
timeout=None,
94+
)
9695

9796
# Continue original step logic
9897
return await original(self, *args, **kwargs)

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

Lines changed: 1 addition & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
import asyncio
44
import json
55
import logging
6-
from typing import Any, AsyncGenerator, cast
6+
from typing import Any, AsyncGenerator
77
from uuid import uuid4
88

99
from llama_index.core.agent.workflow.workflow_events import (
@@ -40,7 +40,6 @@
4040
from uipath_llamaindex.runtime.breakpoints import (
4141
BreakpointEvent,
4242
BreakpointResumeEvent,
43-
DebuggableWorkflow,
4443
inject_breakpoints,
4544
)
4645
from uipath_llamaindex.runtime.chat import UiPathChatMessagesMapper
@@ -144,11 +143,6 @@ async def _run_workflow(
144143

145144
self._context = await self._load_context()
146145

147-
# Make the Context discoverable from inside steps
148-
if self.debug_mode and self._context is not None:
149-
debug_workflow = cast(DebuggableWorkflow, self.workflow)
150-
debug_workflow.context = self._context
151-
152146
if is_resuming:
153147
handler: WorkflowHandler = self.workflow.run(ctx=self._context)
154148
if workflow_input:
Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
{"ticket_id": "T-12345", "customer_message": "The payment system is broken!", "customer_tier": "premium"}
Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
[project]
2+
name = "debug-breakpoints-test"
3+
version = "0.0.1"
4+
description = "Test case for debug breakpoint functionality with LlamaIndex workflows"
5+
authors = [{ name = "UiPath", email = "test@uipath.com" }]
6+
requires-python = ">=3.11"
7+
dependencies = [
8+
"uipath-llamaindex",
9+
"pexpect>=4.9.0",
10+
"pyte>=0.8.0",
11+
"pytest>=8.0.0",
12+
"pytest-asyncio>=0.24.0",
13+
]
14+
15+
[tool.uv.sources]
16+
uipath-llamaindex = { path = "../../", editable = true }
17+
18+
[tool.pytest.ini_options]
19+
asyncio_mode = "auto"
20+
asyncio_default_fixture_loop_scope = "function"
Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
#!/bin/bash
2+
set -e
3+
4+
SAMPLE_DIR="../../samples/debug-agent"
5+
6+
echo "Copying agent files from debug-agent sample..."
7+
cp "$SAMPLE_DIR/main.py" main.py
8+
cp "$SAMPLE_DIR/llama_index.json" llama_index.json
9+
cp "$SAMPLE_DIR/uipath.json" uipath.json
10+
11+
echo "Syncing dependencies..."
12+
uv sync
13+
14+
echo "Initializing the project..."
15+
uv run uipath init
16+
17+
echo "=== Running debug breakpoint tests with pexpect ==="
18+
uv run pytest src/test_debug.py -v -s
Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
# Debug breakpoints tests are run via pytest in test_debug.py
2+
# This file is a placeholder for the testcase runner convention.
3+
4+
print("Debug breakpoints tests completed via pytest.")
Lines changed: 194 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,194 @@
1+
"""
2+
Pexpect-based tests for uipath debug command with LlamaIndex workflows.
3+
4+
Tests the interactive debugger functionality including:
5+
- Single breakpoint
6+
- Multiple breakpoints
7+
- Step mode (s command)
8+
- Quit debugger (q command)
9+
10+
Regression test for: ContextStateError when using wait_for_event in
11+
breakpoint wrapper (the wrapper must use an InternalContext, not the
12+
workflow-level ExternalContext).
13+
14+
See: https://github.com/UiPath/uipath-integrations-python/pull/274
15+
"""
16+
17+
import re
18+
import sys
19+
from pathlib import Path
20+
from typing import Optional
21+
22+
import pexpect
23+
import pytest
24+
25+
26+
# ---------------------------------------------------------------------------
27+
# Minimal PromptTest helper (mirrors uipath-langchain-python/testcases/common)
28+
# ---------------------------------------------------------------------------
29+
30+
def strip_ansi(text: str) -> str:
31+
return re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])').sub('', text)
32+
33+
34+
class PromptTest:
35+
def __init__(self, command: str, test_name: str, prompt: str = "> ", timeout: int = 60):
36+
self.command = command
37+
self.test_name = test_name
38+
self.prompt = prompt
39+
self.timeout = timeout
40+
self.child: Optional[pexpect.spawn] = None
41+
self._log_handle = None
42+
self._log_path = Path(f"{test_name}.log")
43+
44+
def start(self):
45+
self.child = pexpect.spawn(self.command, encoding="utf-8", timeout=self.timeout)
46+
self._log_handle = open(self._log_path, "w")
47+
self.child.logfile_read = self._log_handle
48+
49+
def send_command(self, command: str, expect: Optional[str] = None):
50+
self.child.expect(self.prompt)
51+
self.child.sendline(command)
52+
if expect:
53+
self.child.expect(expect)
54+
55+
def expect_eof(self):
56+
self.child.expect(pexpect.EOF, timeout=self.timeout)
57+
58+
def get_output(self) -> str:
59+
if self._log_path.exists():
60+
if self._log_handle:
61+
self._log_handle.flush()
62+
with open(self._log_path, "r", encoding="utf-8") as f:
63+
return strip_ansi(f.read())
64+
return ""
65+
66+
@property
67+
def before(self) -> str:
68+
return self.child.before if self.child else ""
69+
70+
def close(self):
71+
if self._log_handle:
72+
self._log_handle.close()
73+
self._log_handle = None
74+
if self.child:
75+
self.child.close()
76+
self.child = None
77+
78+
79+
# ---------------------------------------------------------------------------
80+
# Test configuration
81+
# ---------------------------------------------------------------------------
82+
83+
COMMAND = "uv run uipath debug agent --file input.json"
84+
PROMPT = r"> "
85+
TIMEOUT = 60
86+
87+
88+
# ---------------------------------------------------------------------------
89+
# Tests
90+
# ---------------------------------------------------------------------------
91+
92+
def test_single_breakpoint():
93+
"""Test setting and hitting a single breakpoint."""
94+
test = PromptTest(command=COMMAND, test_name="debug_single_breakpoint", prompt=PROMPT, timeout=TIMEOUT)
95+
try:
96+
test.start()
97+
98+
test.send_command("b classify_category", expect=r"Breakpoint set at: classify_category")
99+
test.send_command("c", expect=r"BREAKPOINT.*classify_category.*before")
100+
test.send_command("c", expect=r"Debug session completed")
101+
102+
test.expect_eof()
103+
104+
output = test.get_output()
105+
assert "ticket_id" in output, "Expected ticket_id in output"
106+
107+
except (pexpect.exceptions.TIMEOUT, pexpect.exceptions.EOF) as e:
108+
print(f"\nERROR: {type(e).__name__}", file=sys.stderr)
109+
print(f"\n--- Output before failure ---\n{test.before}", file=sys.stderr)
110+
pytest.fail(f"Test failed: {e}")
111+
finally:
112+
test.close()
113+
114+
115+
def test_multiple_breakpoints():
116+
"""Test setting and hitting multiple breakpoints."""
117+
test = PromptTest(command=COMMAND, test_name="debug_multiple_breakpoints", prompt=PROMPT, timeout=TIMEOUT)
118+
try:
119+
test.start()
120+
121+
test.send_command("b analyze_sentiment", expect=r"Breakpoint set at: analyze_sentiment")
122+
test.send_command("b determine_priority", expect=r"Breakpoint set at: determine_priority")
123+
test.send_command("c", expect=r"BREAKPOINT.*analyze_sentiment.*before")
124+
test.send_command("c", expect=r"BREAKPOINT.*determine_priority.*before")
125+
test.send_command("c", expect=r"Debug session completed")
126+
127+
test.expect_eof()
128+
129+
output = test.get_output()
130+
breakpoint_count = output.count("BREAKPOINT")
131+
assert breakpoint_count >= 2, f"Expected at least 2 breakpoints hit, got {breakpoint_count}"
132+
133+
except (pexpect.exceptions.TIMEOUT, pexpect.exceptions.EOF) as e:
134+
print(f"\nERROR: {type(e).__name__}", file=sys.stderr)
135+
print(f"\n--- Output before failure ---\n{test.before}", file=sys.stderr)
136+
pytest.fail(f"Test failed: {e}")
137+
finally:
138+
test.close()
139+
140+
141+
def test_step_mode():
142+
"""Test step mode - breaks on every node."""
143+
test = PromptTest(command=COMMAND, test_name="debug_step_mode", prompt=PROMPT, timeout=TIMEOUT)
144+
try:
145+
test.start()
146+
147+
# Step through all 9 workflow steps
148+
test.send_command("s", expect=r"BREAKPOINT.*analyze_sentiment.*before")
149+
test.send_command("s", expect=r"BREAKPOINT.*classify_category.*before")
150+
test.send_command("s", expect=r"BREAKPOINT.*check_urgency.*before")
151+
test.send_command("s", expect=r"BREAKPOINT.*determine_priority.*before")
152+
test.send_command("s", expect=r"BREAKPOINT.*check_escalation.*before")
153+
test.send_command("s", expect=r"BREAKPOINT.*route_to_department.*before")
154+
test.send_command("s", expect=r"BREAKPOINT.*assign_queue.*before")
155+
test.send_command("s", expect=r"BREAKPOINT.*generate_response.*before")
156+
test.send_command("s", expect=r"BREAKPOINT.*finalize_ticket.*before")
157+
test.send_command("s", expect=r"Debug session completed")
158+
159+
test.expect_eof()
160+
161+
output = test.get_output()
162+
breakpoint_count = output.count("BREAKPOINT")
163+
assert breakpoint_count >= 9, f"Expected at least 9 breakpoints in step mode, got {breakpoint_count}"
164+
165+
except (pexpect.exceptions.TIMEOUT, pexpect.exceptions.EOF) as e:
166+
print(f"\nERROR: {type(e).__name__}", file=sys.stderr)
167+
print(f"\n--- Output before failure ---\n{test.before}", file=sys.stderr)
168+
pytest.fail(f"Test failed: {e}")
169+
finally:
170+
test.close()
171+
172+
173+
def test_quit_debugger():
174+
"""Test quitting the debugger early with 'q' command."""
175+
test = PromptTest(command=COMMAND, test_name="debug_quit", prompt=PROMPT, timeout=TIMEOUT)
176+
try:
177+
test.start()
178+
179+
test.send_command("b check_urgency", expect=r"Breakpoint set at: check_urgency")
180+
test.send_command("c", expect=r"BREAKPOINT.*check_urgency.*before")
181+
test.send_command("q")
182+
183+
test.expect_eof()
184+
185+
except (pexpect.exceptions.TIMEOUT, pexpect.exceptions.EOF) as e:
186+
print(f"\nERROR: {type(e).__name__}", file=sys.stderr)
187+
print(f"\n--- Output before failure ---\n{test.before}", file=sys.stderr)
188+
pytest.fail(f"Test failed: {e}")
189+
finally:
190+
test.close()
191+
192+
193+
if __name__ == "__main__":
194+
pytest.main([__file__, "-v"])

0 commit comments

Comments
 (0)