"""Tests for the single-capture content layer (src/llm/capture.py).""" from __future__ import annotations from collections.abc import AsyncIterator import pytest from src.llm import capture from src.llm.backend import CompletionResult, ToolCallResult from src.llm.capture import ( CapturedLLMCall, build_captured_call, build_captured_messages, canonical_json, clip_for_trace, compute_content_hash, ) from src.llm.types import ( HonchoLLMCallStreamChunk, LLMTelemetryContext, StreamingResponseWithMetadata, ) async def _chunks( texts: list[str], *, raise_after: BaseException | None = None ) -> AsyncIterator[HonchoLLMCallStreamChunk]: for text in texts: yield HonchoLLMCallStreamChunk(content=text) if raise_after is not None: raise raise_after def _wrapper( stream: AsyncIterator[HonchoLLMCallStreamChunk], recorder: list[tuple[str, str]], ): return StreamingResponseWithMetadata( stream=stream, tool_calls_made=[], input_tokens=0, output_tokens=0, cache_creation_input_tokens=0, cache_read_input_tokens=0, capture_finalizer=lambda text, reason: recorder.append((text, reason)), ) class TestStreamingCaptureFinalizer: async def test_clean_drain_captures_stop(self): recorded: list[tuple[str, str]] = [] wrapper = _wrapper(_chunks(["hel", "lo"]), recorded) async for _ in wrapper: pass assert recorded == [("hello", "stop")] async def test_error_drain_captures_error_and_partial_text(self): recorded: list[tuple[str, str]] = [] wrapper = _wrapper(_chunks(["par"], raise_after=RuntimeError("boom")), recorded) with pytest.raises(RuntimeError): async for _ in wrapper: pass # Partial text still captured, tagged error. assert recorded == [("par", "error")] async def test_cancelled_drain_captures_cancelled(self): import asyncio recorded: list[tuple[str, str]] = [] wrapper = _wrapper( _chunks(["x"], raise_after=asyncio.CancelledError()), recorded ) with pytest.raises(asyncio.CancelledError): async for _ in wrapper: pass assert recorded == [("x", "cancelled")] class TestContentHash: def test_is_deterministic_and_prefixed(self): h1 = compute_content_hash("user", "hello", None) h2 = compute_content_hash("user", "hello", None) assert h1 == h2 assert h1.startswith("sha256:") def test_role_is_inside_the_hash(self): # Identical text under different roles must never collide — role lives # inside the hash, closing the role-in-hash collision bug. assert compute_content_hash("user", "hi", None) != compute_content_hash( "assistant", "hi", None ) def test_tool_call_id_is_inside_the_hash(self): assert compute_content_hash("tool", "ok", "call_1") != compute_content_hash( "tool", "ok", "call_2" ) def test_canonical_json_is_order_independent(self): assert canonical_json({"a": 1, "b": 2}) == canonical_json({"b": 2, "a": 1}) class TestClipForTrace: def test_leaves_small_content_untouched(self): content, truncated = clip_for_trace("short") assert content == "short" assert truncated is False def test_clips_oversized_string(self, monkeypatch: pytest.MonkeyPatch): from src.config import settings monkeypatch.setattr(settings.TELEMETRY, "TRACE_MAX_BYTES", 32) content, truncated = clip_for_trace("x" * 1000) assert truncated is True assert content.endswith("…[truncated]") assert len(content.encode("utf-8")) <= settings.TELEMETRY.TRACE_MAX_BYTES def test_leaves_structured_content_intact(self, monkeypatch: pytest.MonkeyPatch): from src.config import settings monkeypatch.setattr(settings.TELEMETRY, "TRACE_MAX_BYTES", 4) blocks = [{"type": "text", "text": "a long block of structured content"}] content, truncated = clip_for_trace(blocks) assert content == blocks assert truncated is False class TestBuildCapturedMessages: def test_hashes_each_message(self): messages = [ {"role": "user", "content": "q"}, {"role": "assistant", "content": "a"}, ] captured, truncated = build_captured_messages(messages, memo=None) assert [m.role for m in captured] == ["user", "assistant"] assert all(m.content_hash.startswith("sha256:") for m in captured) assert truncated is False def test_memo_makes_hashing_on(self, monkeypatch: pytest.MonkeyPatch): # The conversation is append-only and message dicts are reused, so with # a shared memo each message is hashed exactly once across iterations. calls = {"n": 0} real = compute_content_hash def counting( role: str, content: object, tool_call_id: str | None, tool_calls: list[dict[str, object]] | None = None, ) -> str: calls["n"] += 1 return real(role, content, tool_call_id, tool_calls) monkeypatch.setattr(capture, "compute_content_hash", counting) m1 = {"role": "user", "content": "q1"} m2 = {"role": "assistant", "content": "a1"} m3 = {"role": "user", "content": "q2"} memo: dict[int, capture.CapturedMessage] = {} build_captured_messages([m1, m2], memo) assert calls["n"] == 2 # both hashed build_captured_messages([m1, m2, m3], memo) assert calls["n"] == 3 # only the newly-appended m3 hashed (not re-hashed) class TestNormalizeToolCalls: """Tool calls live outside `content` for openai/gemini — capture must lift them into the unified `tool_calls` shape (the PR concern).""" def test_openai_assistant_tool_calls_captured(self): msg = { "role": "assistant", "content": None, "tool_calls": [ { "id": "call_1", "type": "function", "function": { "name": "search_memory", "arguments": '{"query": "coffee"}', }, } ], } captured, _ = build_captured_messages([msg], memo=None, transport="openai") assert captured[0].tool_calls == [ {"id": "call_1", "name": "search_memory", "input": {"query": "coffee"}} ] def test_gemini_model_parts_captured(self): msg = { "role": "model", "parts": [ {"text": "let me look"}, {"function_call": {"name": "grep_messages", "args": {"text": "x"}}}, ], } captured, _ = build_captured_messages([msg], memo=None, transport="gemini") assert captured[0].content == "let me look" assert captured[0].tool_calls == [ {"id": None, "name": "grep_messages", "input": {"text": "x"}} ] def test_gemini_tool_result_recovered(self): # Gemini tool results live in `parts` (no `content` key) and were dropped. msg = { "role": "user", "parts": [ { "function_response": { "name": "grep_messages", "response": {"result": "3 hits"}, } } ], } captured, _ = build_captured_messages([msg], memo=None, transport="gemini") assert captured[0].content == "3 hits" assert captured[0].tool_call_id == "grep_messages" def test_anthropic_tool_use_blocks_normalized(self): msg = { "role": "assistant", "content": [ {"type": "text", "text": "searching"}, { "type": "tool_use", "id": "tu_1", "name": "search_memory", "input": {"q": "x"}, }, ], } captured, _ = build_captured_messages([msg], memo=None, transport="anthropic") assert captured[0].content == "searching" assert captured[0].tool_calls == [ {"id": "tu_1", "name": "search_memory", "input": {"q": "x"}} ] def test_hash_distinguishes_tool_calls(self): # Two empty-content assistant turns with different tool calls must not # collide in the dedup store (they did before tool_calls entered the hash). base = {"role": "assistant", "content": None} a = { **base, "tool_calls": [ { "id": "c1", "type": "function", "function": {"name": "search_memory", "arguments": "{}"}, } ], } b = { **base, "tool_calls": [ { "id": "c2", "type": "function", "function": {"name": "search_messages", "arguments": "{}"}, } ], } (ca,), _ = build_captured_messages([a], memo=None, transport="openai") (cb,), _ = build_captured_messages([b], memo=None, transport="openai") assert ca.content_hash != cb.content_hash class TestBuildCapturedCall: def test_maps_telemetry_and_result(self): telemetry = LLMTelemetryContext( workspace_name="ws", call_purpose="dialectic.answer", parent_category="dialectic", run_id="r1", trace_id="r1", span_id="r1", session_id="sess_abc", iteration=2, step_seq=2, ) result = CompletionResult( content="answer", input_tokens=10, output_tokens=5, finish_reason="stop", tool_calls=[ToolCallResult(id="t1", name="search", input={"q": "x"})], ) call = build_captured_call( telemetry=telemetry, transport="anthropic", provider_label=None, model="claude-x", messages=[{"role": "user", "content": "q"}], tools=None, tool_choice=None, result=result, attempt=1, was_fallback=False, was_stream=False, finish_reason="stop", ) assert isinstance(call, CapturedLLMCall) assert call.trace_id == "r1" and call.span_id == "r1" assert call.iteration == 2 and call.step_seq == 2 assert call.output_content == "answer" assert call.output_tool_calls == [ {"id": "t1", "name": "search", "input": {"q": "x"}} ] assert call.input_tokens == 10 and call.output_tokens == 5 assert call.session_id == "sess_abc" assert len(call.input_messages) == 1 assert call.input_messages[0].content_hash.startswith("sha256:") def test_session_id_defaults_none_without_telemetry(self): # Sessionless calls (and the no-telemetry path) carry session_id=None so # the Langfuse projection emits no session grouping for them. telemetry = LLMTelemetryContext(run_id="r1", trace_id="r1", span_id="r1") call = build_captured_call( telemetry=telemetry, transport="anthropic", provider_label=None, model="claude-x", messages=[{"role": "user", "content": "q"}], tools=None, tool_choice=None, result=CompletionResult(content="a", finish_reason="stop"), attempt=1, was_fallback=False, was_stream=False, finish_reason="stop", ) assert call.session_id is None def test_self_parent_is_normalized_to_none(self): # The tool loop sets parent_span_id == span_id on the run span (it # doubles as the Langfuse "inside a run" signal). A span that is its own # parent is a root, so the EXPORTED parent_span_id must be None — else # span-tree consumers file the root as a child of itself. telemetry = LLMTelemetryContext( run_id="r1", trace_id="r1", span_id="r1", parent_span_id="r1" ) call = build_captured_call( telemetry=telemetry, transport="anthropic", provider_label=None, model="claude-x", messages=[{"role": "user", "content": "q"}], tools=None, tool_choice=None, result=CompletionResult(content="a", finish_reason="stop"), attempt=1, was_fallback=False, was_stream=False, finish_reason="stop", ) assert call.span_id == "r1" assert call.parent_span_id is None # A genuine distinct parent is preserved. telemetry.parent_span_id = "parent-span" assert telemetry.exported_parent_span_id() == "parent-span" def test_error_path_collapses_output(self): call = build_captured_call( telemetry=None, transport="anthropic", provider_label=None, model="claude-x", messages=[{"role": "user", "content": "q"}], tools=None, tool_choice=None, result=None, attempt=2, was_fallback=True, was_stream=False, finish_reason="error", ) assert call.output_content is None assert call.output_tool_calls == [] assert call.finish_reason == "error" assert call.attempt == 2 and call.was_fallback is True class TestExporterRegistry: def test_register_dispatch_and_clear(self): capture.clear_exporters() assert capture.has_exporters() is False seen: list[CapturedLLMCall] = [] class _Spy: def export(self, call: CapturedLLMCall) -> None: seen.append(call) capture.register_exporter(_Spy()) assert capture.has_exporters() is True call = build_captured_call( telemetry=None, transport="anthropic", provider_label=None, model="m", messages=[], tools=None, tool_choice=None, result=None, attempt=1, was_fallback=False, was_stream=False, finish_reason="stop", ) capture.dispatch_captured_call(call) assert seen == [call] capture.clear_exporters() assert capture.has_exporters() is False def test_dispatch_swallows_exporter_errors(self): capture.clear_exporters() class _Boom: def export(self, call: CapturedLLMCall) -> None: raise RuntimeError(f"nope: {call.model}") capture.register_exporter(_Boom()) call = build_captured_call( telemetry=None, transport="anthropic", provider_label=None, model="m", messages=[], tools=None, tool_choice=None, result=None, attempt=1, was_fallback=False, was_stream=False, finish_reason="stop", ) # Must not raise — telemetry never breaks the LLM path. capture.dispatch_captured_call(call) capture.clear_exporters() class TestThoughtSignatureSerialization: """Gemini `thought_signature` is bytes; it must not break trace serialization.""" def test_bytes_signature_base64_encoded_and_serializes(self): import json from src.telemetry.events.trace import LLMCallTracedEvent result = CompletionResult( content=None, finish_reason="STOP", tool_calls=[ ToolCallResult( id="call_1", name="grep_messages", input={"text": "coffee"}, thought_signature=b"\x0a\x1f\x88\xff\x00sig", ) ], ) call = build_captured_call( telemetry=LLMTelemetryContext(trace_id="t1", span_id="s1"), transport="gemini", provider_label=None, model="gemini-2.5-flash", messages=[{"role": "user", "content": "q"}], tools=None, tool_choice=None, result=result, attempt=1, was_fallback=False, was_stream=False, finish_reason="STOP", ) sig = call.output_tool_calls[0]["thought_signature"] assert isinstance(sig, str) # base64, not raw bytes # The traced event must serialize to JSON without raising (the emit path # calls model_dump(mode="json"), which threw UnicodeDecodeError on bytes). event = LLMCallTracedEvent( model="gemini-2.5-flash", transport="gemini", output_tool_calls=call.output_tool_calls, ) json.dumps(event.model_dump(mode="json"))