From fd258335657491fde02108b070f6424432512be0 Mon Sep 17 00:00:00 2001 From: ciregenz Date: Tue, 23 Jun 2026 00:46:35 -0700 Subject: [PATCH] [eric] agents: streaming harness driving real _run_agent_loop against mocked SDK (Phase 0 safety net) --- backend/tests/test_streaming_harness.py | 143 ++++++++++++++++++++++++ 1 file changed, 143 insertions(+) create mode 100644 backend/tests/test_streaming_harness.py diff --git a/backend/tests/test_streaming_harness.py b/backend/tests/test_streaming_harness.py new file mode 100644 index 00000000..0b8522f3 --- /dev/null +++ b/backend/tests/test_streaming_harness.py @@ -0,0 +1,143 @@ +"""Streaming harness: drive the real _run_agent_loop with a MOCKED claude_agent_sdk.query +that yields a controlled SDK message sequence, and assert the session state + emitted WS +events. This is the safety net for restructuring the streaming loop (it had no isolated +coverage), so it pins the observable contract: streamed text lands as an assistant message, +tool calls are recorded, and the turn completes.""" + +import asyncio + +import claude_agent_sdk +from claude_agent_sdk import AssistantMessage, ResultMessage +from claude_agent_sdk.types import TextBlock, ToolUseBlock, ThinkingBlock + +from backend.apps.agents.agent_manager import AgentManager +import backend.apps.agents.core.ws_manager as ws_mod + + +def _mock_query_yielding(*messages): + async def _q(*args, **kwargs): + for m in messages: + yield m + return _q + + +def _drive(monkeypatch, messages, prompt="hi"): + """Run one _run_agent_loop turn against a mocked SDK message stream; return (session, ws_events).""" + events = [] + + async def fake_send(session_id, event, data): + events.append((event, data)) + + monkeypatch.setattr(ws_mod.ws_manager, "send_to_session", fake_send, raising=True) + monkeypatch.setattr(claude_agent_sdk, "query", _mock_query_yielding(*messages), raising=True) + + mgr = AgentManager() + from backend.apps.agents.core.models import AgentSession + session = AgentSession(name="t", model="sonnet", dashboard_id="d") + mgr.sessions[session.id] = session + asyncio.run(mgr._run_agent_loop(session.id, prompt)) + return session, events + + +def _result(**kw): + base = dict(subtype="success", duration_ms=100, duration_api_ms=80, is_error=False, + num_turns=1, session_id="sdk-1", usage={"input_tokens": 10, "output_tokens": 5}) + base.update(kw) + return ResultMessage(**base) + + +def _assistant(blocks, **kw): + base = dict(content=blocks, model="sonnet", message_id="m1", stop_reason="end_turn", + session_id="sdk-1", usage={"input_tokens": 10, "output_tokens": 5}) + base.update(kw) + return AssistantMessage(**base) + + +def test_streamed_text_lands_as_assistant_message(monkeypatch): + session, events = _drive(monkeypatch, [ + _assistant([TextBlock(text="Hello there")]), + _result(), + ]) + assert any(m.role == "assistant" and "Hello there" in str(m.content) for m in session.messages) + assert session.status == "completed" + # the assistant reply + the token meter are broadcast to the UI + assert any(e == "agent:message" for e, _ in events) + assert any(e == "agent:context_update" for e, _ in events) + + +def test_tool_use_is_recorded(monkeypatch): + session, events = _drive(monkeypatch, [ + _assistant([ToolUseBlock(id="tu1", name="Read", input={"file_path": "/x.py"})]), + _result(), + ]) + assert any(m.role == "tool_call" for m in session.messages) + # the tool name survives onto the recorded call + assert any("Read" in str(m.content) for m in session.messages if m.role == "tool_call") + assert session.status == "completed" + + +def test_text_then_tool_in_one_turn(monkeypatch): + session, events = _drive(monkeypatch, [ + _assistant([TextBlock(text="Let me read it."), ToolUseBlock(id="tu1", name="Read", input={"file_path": "/x.py"})]), + _result(), + ]) + roles = [m.role for m in session.messages] + assert "assistant" in roles and "tool_call" in roles + assert session.status == "completed" + + +def test_completes_even_with_no_content(monkeypatch): + # an empty assistant turn (e.g. a pure stop) must still finish cleanly, not hang + session, events = _drive(monkeypatch, [_assistant([]), _result()]) + assert session.status == "completed" + + +def test_thinking_block_before_text_is_handled(monkeypatch): + # a ThinkingBlock mutates the separate thinking-state cluster; the turn must still + # surface the final answer and complete (pins the thinking path for the restructuring) + session, events = _drive(monkeypatch, [ + _assistant([ThinkingBlock(thinking="let me reason about this", signature="sig-1"), + TextBlock(text="the answer is 42")]), + _result(), + ]) + assert any(m.role == "assistant" and "the answer is 42" in str(m.content) for m in session.messages) + assert session.status == "completed" + + +def test_transient_capacity_error_is_retried_then_succeeds(monkeypatch): + # the capacity-retry while-loop: first query() raises a transient error, the loop + # backs off (sleep mocked to no-op) and re-queries, which succeeds. This is the exact + # behavior the streaming restructuring must preserve. + real_sleep = asyncio.sleep # capture before patching to avoid self-recursion + + async def _fast_sleep(*a, **k): + await real_sleep(0) # still yields to the loop, but no real backoff delay + + monkeypatch.setattr(asyncio, "sleep", _fast_sleep) + + state = {"n": 0} + + async def flaky_query(*args, **kwargs): + state["n"] += 1 + if state["n"] == 1: + raise Exception("No pool capacity available. Try again shortly.") + yield _assistant([TextBlock(text="Recovered after backoff")]) + yield _result() + + events = [] + + async def fake_send(session_id, event, data): + events.append((event, data)) + + monkeypatch.setattr(ws_mod.ws_manager, "send_to_session", fake_send, raising=True) + monkeypatch.setattr(claude_agent_sdk, "query", flaky_query, raising=True) + + mgr = AgentManager() + from backend.apps.agents.core.models import AgentSession + session = AgentSession(name="t", model="sonnet", dashboard_id="d") + mgr.sessions[session.id] = session + asyncio.run(mgr._run_agent_loop(session.id, "hi")) + + assert state["n"] == 2 # retried exactly once + assert any(m.role == "assistant" and "Recovered" in str(m.content) for m in session.messages) + assert session.status == "completed"