diff --git a/backend/tests/test_agent_loop.py b/backend/tests/test_agent_loop.py new file mode 100644 index 00000000..320de4cd --- /dev/null +++ b/backend/tests/test_agent_loop.py @@ -0,0 +1,712 @@ +"""Unit tests for `backend.apps.agents.agent_loop.AgentLoop`. + +The agent loop is the provider-agnostic streaming + tool-use + HITL +core: it drives `BaseProvider.stream_message`, accumulates content +blocks, executes tools (with HITL gating), emits the WebSocket events +the frontend consumes, and persists the final messages. + +The thinking-block path is already covered by `test_phase1_stress.py`. +This file focuses on the rest of the surface — the `run()` control +flow, `_execute_tools` (HITL deny / updated_input / executor errors / +truncation / multi-tool / non-tool skip), `_stream_and_collect` +JSON handling and stop_reason routing, `_emit_collected_messages` +formatting, and token-usage accumulation. + +Pure unit tests: no FastAPI client, no network. We script provider +output via `_StubProvider` and capture WS emissions via `_WSRecorder`. +The conftest bootstrap (run automatically by virtue of living in +`backend/tests/`) redirects `OPENSWARM_DATA_DIR` and mocks PostHog so +nothing here ever touches the real disk or external services. +""" + +from __future__ import annotations + +from typing import Any +from unittest.mock import AsyncMock + +import pytest + +from backend.apps.agents.agent_loop import AgentLoop +from backend.apps.agents.providers.base import ( + ContentBlock, + ModelResponse, + ProviderMessage, + StreamEvent, + ToolCall, +) + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +class _StubProvider: + """Scripted `BaseProvider` stand-in for the agent loop. + + Each call to `stream_message` consumes one entry from `turns` and + yields its `StreamEvent`s in order. The messages list passed in is + snapshotted into `self.calls[i]` so tests can assert on the + conversation-history shape the agent loop sent on each turn. + + Duck-typed (not a `BaseProvider` subclass) — matches the + convention in `test_phase1_stress.py` and avoids having to stub + out `create_message` / `get_model_id` that the loop never calls. + """ + + def __init__(self, turns: list[list[StreamEvent]]): + self._turns = list(turns) + self.calls: list[list[ProviderMessage]] = [] + + async def stream_message( + self, + *, + model: str, + system: str | None, + messages: list[ProviderMessage], + tools: list, + ): + # Snapshot at call time — the loop mutates `self.messages` + # between turns, so a reference would get clobbered. + self.calls.append(list(messages)) + if not self._turns: + return + for ev in self._turns.pop(0): + yield ev + + def format_user_message(self, content: Any) -> ProviderMessage: + return ProviderMessage(role="user", content=content) + + def format_assistant_message(self, response: ModelResponse) -> ProviderMessage: + # Shape doesn't matter for these tests — the loop just appends + # the message; we never feed it back through a real provider. + return ProviderMessage(role="assistant", content=response.content) + + def format_tool_result(self, tool_use_id: str, content: list[dict]) -> dict: + return {"type": "tool_result", "tool_use_id": tool_use_id, "content": content} + + +class _WSRecorder: + """Async-callable that captures every (event, payload) pair the + loop emits. Tests assert on `events` directly or via `of_type`.""" + + def __init__(self): + self.events: list[tuple[str, dict]] = [] + + async def __call__(self, event_type: str, data: dict) -> None: + self.events.append((event_type, data)) + + def of_type(self, event_type: str) -> list[dict]: + return [d for (e, d) in self.events if e == event_type] + + +# Block factories. A "block" is the start/delta/stop triple for one +# content_block in the model's output; a "turn" is one full +# stream_message call (one or more blocks + optional usage + +# message_stop). + + +def _text_block(text: str, *, index: int = 0) -> list[StreamEvent]: + return [ + StreamEvent(type="content_block_start", index=index, block_type="text"), + StreamEvent( + type="content_block_delta", + index=index, + delta_type="text_delta", + text=text, + ), + StreamEvent(type="content_block_stop", index=index), + ] + + +def _tool_block( + name: str, + tool_id: str, + json_input: str, + *, + index: int = 0, +) -> list[StreamEvent]: + return [ + StreamEvent( + type="content_block_start", + index=index, + block_type="tool_use", + tool_name=name, + tool_id=tool_id, + ), + StreamEvent( + type="content_block_delta", + index=index, + delta_type="input_json_delta", + text=json_input, + ), + StreamEvent(type="content_block_stop", index=index), + ] + + +def _turn( + *blocks: list[StreamEvent], + usage: dict[str, int] | None = None, +) -> list[StreamEvent]: + """Compose 1+ blocks into a complete turn (with usage + message_stop).""" + events: list[StreamEvent] = [] + for block in blocks: + events.extend(block) + if usage: + events.append(StreamEvent(type="usage", usage=usage)) + events.append(StreamEvent(type="message_stop")) + return events + + +def _text_turn(text: str, *, usage: dict[str, int] | None = None) -> list[StreamEvent]: + return _turn(_text_block(text), usage=usage) + + +def _tool_turn( + name: str, + tool_id: str, + json_input: str, + *, + usage: dict[str, int] | None = None, +) -> list[StreamEvent]: + return _turn(_tool_block(name, tool_id, json_input), usage=usage) + + +def _make_loop( + *, + provider: _StubProvider | None = None, + executor=None, + hitl=None, + max_turns: int | None = None, + system_prompt: str | None = None, +) -> tuple[AgentLoop, _WSRecorder]: + """Build an `AgentLoop` plus its WS recorder. Sensible defaults + for tests that don't care about a particular callback.""" + ws = _WSRecorder() + loop = AgentLoop( + session_id="test-session", + provider=provider if provider is not None else _StubProvider([]), + model="sonnet", + system_prompt=system_prompt, + tools=[], + tool_executor=executor + if executor is not None + else AsyncMock(return_value=[{"type": "text", "text": "ok"}]), + hitl_handler=hitl + if hitl is not None + else AsyncMock(return_value=(True, None)), + ws_emitter=ws, + max_turns=max_turns, + ) + return loop, ws + + +# --------------------------------------------------------------------------- +# Group 1 — run() control flow +# --------------------------------------------------------------------------- + + +async def test_end_turn_exits_after_one_iteration(): + """No tool_use → loop calls the provider once and stops with + [user, assistant] in history.""" + provider = _StubProvider([_text_turn("hello!")]) + loop, _ = _make_loop(provider=provider) + + await loop.run("hi") + + assert len(provider.calls) == 1 + assert [m.role for m in loop.messages] == ["user", "assistant"] + + +async def test_tool_use_continues_loop_then_terminates_on_end_turn(): + """tool_use → execute → second model call → end_turn. Verifies the + full conversation grows to [user, assistant, tool_result, assistant] + AND that turn 2 sees the prior tool_result.""" + provider = _StubProvider([ + _tool_turn("Read", "t1", '{"path": "/x"}'), + _text_turn("done"), + ]) + executor = AsyncMock(return_value=[{"type": "text", "text": "file contents"}]) + loop, _ = _make_loop(provider=provider, executor=executor) + + await loop.run("hi") + + assert len(provider.calls) == 2 + assert [m.role for m in loop.messages] == [ + "user", "assistant", "tool_result", "assistant", + ] + # The second model call must see the tool_result we appended. + roles_seen_on_turn_2 = [m.role for m in provider.calls[1]] + assert roles_seen_on_turn_2 == ["user", "assistant", "tool_result"] + executor.assert_awaited_once() + + +async def test_max_turns_halts_loop_before_second_model_call(): + """`max_turns=1` allows exactly one model call + tool execution, + then the top-of-loop guard breaks before turn 2. The unused + second turn stays scripted but un-consumed.""" + provider = _StubProvider([ + _tool_turn("Read", "t1", "{}"), + _tool_turn("Read", "t2", "{}"), # would be consumed if guard failed + ]) + executor = AsyncMock(return_value=[{"type": "text", "text": "ok"}]) + loop, _ = _make_loop(provider=provider, executor=executor, max_turns=1) + + await loop.run("hi") + + assert len(provider.calls) == 1, "max_turns=1 must cap provider calls" + assert executor.await_count == 1, "tool exec runs in turn 1, not gated by max_turns" + + +async def test_no_tool_results_breaks_loop(monkeypatch): + """Defensive guard: if `_execute_tools` returns `[]` for any + reason, the loop must exit without appending an empty + `tool_result` and without re-calling the provider.""" + provider = _StubProvider([ + _tool_turn("Read", "t1", "{}"), + _text_turn("never reached"), + ]) + loop, _ = _make_loop(provider=provider) + + async def _empty_results(self, response): + return [] + + monkeypatch.setattr(AgentLoop, "_execute_tools", _empty_results) + + await loop.run("hi") + + assert len(provider.calls) == 1 + assert [m.role for m in loop.messages] == ["user", "assistant"] + + +# --------------------------------------------------------------------------- +# Group 2 — _execute_tools +# --------------------------------------------------------------------------- + + +def _tool_use_response(*calls: tuple[str, str, dict]) -> ModelResponse: + """Build a ModelResponse(stop_reason='tool_use') from (id, name, input) tuples.""" + return ModelResponse( + content=[ + ContentBlock( + type="tool_use", + tool_call=ToolCall(id=tid, name=name, input=inp), + ) + for (tid, name, inp) in calls + ], + stop_reason="tool_use", + ) + + +async def test_hitl_denial_skips_executor_and_returns_denial_text(): + executor = AsyncMock() + hitl = AsyncMock(return_value=(False, None)) + loop, ws = _make_loop(executor=executor, hitl=hitl) + + response = _tool_use_response(("t1", "Read", {"path": "/x"})) + results = await loop._execute_tools(response) + + executor.assert_not_called() + assert len(results) == 1, "denied tools still produce a tool_result for the model" + + tool_result_msgs = [ + d["message"] for d in ws.of_type("agent:message") + if d["message"]["role"] == "tool_result" + ] + assert len(tool_result_msgs) == 1 + assert tool_result_msgs[0]["content"]["text"] == "Tool use was denied by the user." + assert tool_result_msgs[0]["content"]["tool_name"] == "Read" + + +async def test_hitl_updated_input_passed_to_executor(): + """When HITL approves with an `updated_input`, the executor must + see that dict — not the model's original input.""" + executor = AsyncMock(return_value=[{"type": "text", "text": "ok"}]) + hitl = AsyncMock(return_value=(True, {"path": "/y"})) + loop, _ = _make_loop(executor=executor, hitl=hitl) + + response = _tool_use_response(("t1", "Read", {"path": "/x"})) + await loop._execute_tools(response) + + executor.assert_awaited_once_with("Read", {"path": "/y"}) + + +async def test_executor_exception_is_caught_and_surfaced_as_error_text(): + async def boom(name, inp): + raise RuntimeError("disk on fire") + + loop, ws = _make_loop(executor=boom) + response = _tool_use_response(("t1", "Read", {})) + + results = await loop._execute_tools(response) + + # Loop survives, returns a formatted error result for the provider. + assert len(results) == 1 + tool_results = [ + d["message"] for d in ws.of_type("agent:message") + if d["message"]["role"] == "tool_result" + ] + assert tool_results[0]["content"]["text"] == "Error executing Read: disk on fire" + + +async def test_tool_result_text_truncated_to_15000_chars_in_emitted_message(): + """The model gets the full tool output; the WS message bubble + that the UI renders is sliced to 15K to avoid jank.""" + huge = "x" * 20_000 + executor = AsyncMock(return_value=[{"type": "text", "text": huge}]) + loop, ws = _make_loop(executor=executor) + + response = _tool_use_response(("t1", "Read", {})) + results = await loop._execute_tools(response) + + # WS-emitted snippet capped at 15K. + tr_msg = next( + d["message"] for d in ws.of_type("agent:message") + if d["message"]["role"] == "tool_result" + ) + assert len(tr_msg["content"]["text"]) == 15_000 + + # Provider-bound result is the raw, untruncated content. + assert results[0]["content"][0]["text"] == huge + + +async def test_multiple_tool_calls_in_one_response_all_execute(): + executor = AsyncMock(return_value=[{"type": "text", "text": "ok"}]) + loop, ws = _make_loop(executor=executor) + + response = _tool_use_response( + ("t1", "Read", {"path": "/a"}), + ("t2", "Edit", {"path": "/b"}), + ) + results = await loop._execute_tools(response) + + assert executor.await_count == 2 + assert [r["tool_use_id"] for r in results] == ["t1", "t2"] + tool_result_msgs = [ + d["message"] for d in ws.of_type("agent:message") + if d["message"]["role"] == "tool_result" + ] + assert [m["content"]["tool_name"] for m in tool_result_msgs] == ["Read", "Edit"] + + +async def test_non_tool_use_blocks_are_skipped_in_executor(): + """A response with text + tool_use should only execute the tool block.""" + executor = AsyncMock(return_value=[{"type": "text", "text": "ok"}]) + loop, _ = _make_loop(executor=executor) + + response = ModelResponse( + content=[ + ContentBlock(type="text", text="thinking aloud"), + ContentBlock( + type="tool_use", + tool_call=ToolCall(id="t1", name="Read", input={}), + ), + ], + stop_reason="tool_use", + ) + + results = await loop._execute_tools(response) + assert len(results) == 1 + executor.assert_awaited_once() + + +# --------------------------------------------------------------------------- +# Group 3 — _stream_and_collect +# --------------------------------------------------------------------------- + + +async def test_invalid_tool_input_json_falls_back_to_empty_dict(): + """Malformed JSON in input_json_delta must not crash; the + resulting ToolCall.input is `{}`.""" + provider = _StubProvider([_tool_turn("Read", "t1", "{not valid json")]) + loop, _ = _make_loop(provider=provider) + + response = await loop._stream_and_collect() + + tool_blocks = [b for b in response.content if b.type == "tool_use"] + assert len(tool_blocks) == 1 + assert tool_blocks[0].tool_call is not None + assert tool_blocks[0].tool_call.input == {} + + +async def test_tool_use_input_assembled_from_multiple_json_deltas(): + """Real Anthropic streams ship tool input in multiple + input_json_delta chunks — they must concatenate into one JSON + parse.""" + provider = _StubProvider([[ + StreamEvent( + type="content_block_start", index=0, block_type="tool_use", + tool_name="Read", tool_id="t1", + ), + StreamEvent( + type="content_block_delta", index=0, + delta_type="input_json_delta", text='{"pa', + ), + StreamEvent( + type="content_block_delta", index=0, + delta_type="input_json_delta", text='th": "/x", "n": 7}', + ), + StreamEvent(type="content_block_stop", index=0), + StreamEvent(type="message_stop"), + ]]) + loop, _ = _make_loop(provider=provider) + + response = await loop._stream_and_collect() + + assert response.content[0].tool_call.input == {"path": "/x", "n": 7} + + +async def test_stop_reason_routes_on_presence_of_tool_use_block(): + """`response.stop_reason` is `"tool_use"` iff any collected + block is tool_use, else `"end_turn"`.""" + text_provider = _StubProvider([_text_turn("hi")]) + loop1, _ = _make_loop(provider=text_provider) + resp_text = await loop1._stream_and_collect() + assert resp_text.stop_reason == "end_turn" + + tool_provider = _StubProvider([_tool_turn("Read", "t1", "{}")]) + loop2, _ = _make_loop(provider=tool_provider) + resp_tool = await loop2._stream_and_collect() + assert resp_tool.stop_reason == "tool_use" + + +# --------------------------------------------------------------------------- +# Group 3b — _stream_and_collect WS emissions +# --------------------------------------------------------------------------- +# Thinking-block WS emissions are covered in `test_phase1_stress.py`. The +# tests below cover the text and tool_use streaming paths plus the +# routing rule that determines WHERE `agent:stream_end` fires for each +# block type — text waits for `message_stop`, tool_use/thinking close at +# `content_block_stop`. + + +async def test_text_block_streams_deltas_and_ends_at_message_stop(): + """Text block emits one stream_start (role=assistant), one + stream_delta per text_delta event (carrying the same message_id), + and one stream_end deferred to `message_stop` — NOT to + `content_block_stop`.""" + provider = _StubProvider([[ + StreamEvent(type="content_block_start", index=0, block_type="text"), + StreamEvent( + type="content_block_delta", index=0, + delta_type="text_delta", text="hel", + ), + StreamEvent( + type="content_block_delta", index=0, + delta_type="text_delta", text="lo", + ), + StreamEvent(type="content_block_stop", index=0), + StreamEvent(type="message_stop"), + ]]) + loop, ws = _make_loop(provider=provider) + + await loop._stream_and_collect() + + starts = ws.of_type("agent:stream_start") + assert len(starts) == 1 + assert starts[0]["role"] == "assistant" + text_msg_id = starts[0]["message_id"] + + deltas = ws.of_type("agent:stream_delta") + assert [d["delta"] for d in deltas] == ["hel", "lo"] + assert all(d["message_id"] == text_msg_id for d in deltas) + + ends = ws.of_type("agent:stream_end") + assert len(ends) == 1 + assert ends[0] == {"message_id": text_msg_id} + + # Stream-end placement: text's stream_end must come AFTER the + # content_block_stop has already been processed — i.e. it's tied + # to message_stop. Concretely, no agent:stream_delta or new + # agent:stream_start can follow it for this same message_id. + types_in_order = [e for (e, _) in ws.events] + end_idx = types_in_order.index("agent:stream_end") + assert "agent:stream_delta" not in types_in_order[end_idx + 1:] + + +async def test_tool_use_block_streams_deltas_and_ends_at_block_stop(): + """Tool_use block emits stream_start (role=tool_call, tool_name), + one stream_delta per input_json_delta, and stream_end at + `content_block_stop` so the UI can finalize the tool-call bubble + before any subsequent text streams in.""" + provider = _StubProvider([[ + StreamEvent( + type="content_block_start", index=0, block_type="tool_use", + tool_name="Read", tool_id="t1", + ), + StreamEvent( + type="content_block_delta", index=0, + delta_type="input_json_delta", text='{"path"', + ), + StreamEvent( + type="content_block_delta", index=0, + delta_type="input_json_delta", text=': "/x"}', + ), + StreamEvent(type="content_block_stop", index=0), + StreamEvent(type="message_stop"), + ]]) + loop, ws = _make_loop(provider=provider) + + await loop._stream_and_collect() + + starts = ws.of_type("agent:stream_start") + assert len(starts) == 1 + assert starts[0]["role"] == "tool_call" + assert starts[0]["tool_name"] == "Read" + tool_msg_id = starts[0]["message_id"] + + deltas = ws.of_type("agent:stream_delta") + assert [d["delta"] for d in deltas] == ['{"path"', ': "/x"}'] + assert all(d["message_id"] == tool_msg_id for d in deltas) + + ends = ws.of_type("agent:stream_end") + assert len(ends) == 1 + assert ends[0]["message_id"] == tool_msg_id + + # tool_use stream_end fires at content_block_stop — i.e. BEFORE + # the message_stop housekeeping. Verify there's no subsequent + # delta/end for this id and that an agent:message (the persisted + # tool_call bubble) follows in `_emit_collected_messages`. + types_in_order = [e for (e, _) in ws.events] + stop_idx = types_in_order.index("agent:stream_end") + assert "agent:stream_delta" not in types_in_order[stop_idx + 1:] + assert "agent:message" in types_in_order[stop_idx + 1:], ( + "the persisted tool_call message must be emitted after stream_end" + ) + + +# --------------------------------------------------------------------------- +# Group 4 — _emit_collected_messages +# --------------------------------------------------------------------------- + + +async def test_text_blocks_joined_with_single_newline_and_msg_id_preserved(): + """Multiple text blocks → one assistant Message whose `id` is the + streamed msg_id (so the client can dedupe its optimistic bubble).""" + loop, ws = _make_loop() + content = [ + ContentBlock(type="text", text="line A"), + ContentBlock(type="text", text="line B"), + ContentBlock(type="text", text=""), # empty blocks dropped + ] + + await loop._emit_collected_messages( + content, text_msg_id="text-id-123", tool_msg_ids={}, + ) + + assistants = [ + d["message"] for d in ws.of_type("agent:message") + if d["message"]["role"] == "assistant" + ] + assert len(assistants) == 1 + assert assistants[0]["content"] == "line A\nline B" + assert assistants[0]["id"] == "text-id-123" + + +async def test_thinking_blocks_joined_with_double_newline_and_metadata(): + """Multiple thinking blocks → one persisted thinking Message + joined by `\\n\\n`. `elapsed_ms` and `tokens` are derived from + the server-stamped accumulators.""" + loop, ws = _make_loop() + content = [ + ContentBlock(type="thinking", text="step one"), + ContentBlock(type="thinking", text="step two"), + ContentBlock(type="text", text="answer"), + ] + + await loop._emit_collected_messages( + content, + text_msg_id="t1", + tool_msg_ids={}, + thinking_elapsed_ms=1234, + thinking_total_chars=20, + ) + + msgs = [d["message"] for d in ws.of_type("agent:message")] + thinking = [m for m in msgs if m["role"] == "thinking"] + assert len(thinking) == 1 + assert thinking[0]["content"] == "step one\n\nstep two" + assert thinking[0]["elapsed_ms"] == 1234 + assert thinking[0]["tokens"] == max(1, round(20 / 3.6)) + + +async def test_thinking_message_omits_tokens_and_elapsed_when_zero(): + """`thinking_elapsed_ms=0` → `elapsed_ms=None`, and + `thinking_total_chars=0` → `tokens=None`. Matches the + `... or None` / `if thinking_total_chars` guards.""" + loop, ws = _make_loop() + content = [ContentBlock(type="thinking", text="thoughts")] + + await loop._emit_collected_messages( + content, + text_msg_id=None, + tool_msg_ids={}, + thinking_elapsed_ms=0, + thinking_total_chars=0, + ) + + thinking = next( + d["message"] for d in ws.of_type("agent:message") + if d["message"]["role"] == "thinking" + ) + assert thinking["elapsed_ms"] is None + assert thinking["tokens"] is None + + +async def test_tool_call_messages_use_stream_msg_ids_in_index_order(): + """`tool_msg_ids` is keyed by stream block index. Emission must + sort by index so the first emitted tool_call gets the id from + index 0, even if the dict was inserted out of order.""" + loop, ws = _make_loop() + content = [ + ContentBlock( + type="tool_use", + tool_call=ToolCall(id="t1", name="Read", input={"x": 1}), + ), + ContentBlock( + type="tool_use", + tool_call=ToolCall(id="t2", name="Edit", input={"y": 2}), + ), + ] + + # Insert higher index first to guard against accidental + # insertion-order semantics in the future. + await loop._emit_collected_messages( + content, + text_msg_id=None, + tool_msg_ids={1: "ID-B", 0: "ID-A"}, + ) + + tool_calls = [ + d["message"] for d in ws.of_type("agent:message") + if d["message"]["role"] == "tool_call" + ] + assert [m["id"] for m in tool_calls] == ["ID-A", "ID-B"] + assert tool_calls[0]["content"] == {"id": "t1", "tool": "Read", "input": {"x": 1}} + assert tool_calls[1]["content"] == {"id": "t2", "tool": "Edit", "input": {"y": 2}} + + +# --------------------------------------------------------------------------- +# Group 5 — usage tracking +# --------------------------------------------------------------------------- + + +async def test_token_usage_accumulates_across_turns(): + """Per-turn `usage` events must sum into `total_input_tokens` / + `total_output_tokens` over the whole `run()`.""" + provider = _StubProvider([ + _tool_turn( + "Read", "t1", "{}", + usage={"input_tokens": 100, "output_tokens": 50}, + ), + _text_turn( + "done", + usage={"input_tokens": 30, "output_tokens": 20}, + ), + ]) + loop, _ = _make_loop(provider=provider) + + await loop.run("hi") + + assert loop.total_input_tokens == 130 + assert loop.total_output_tokens == 70