diff --git a/packages/claude-agents/src/launchdarkly_ai_claude_agents/handler.py b/packages/claude-agents/src/launchdarkly_ai_claude_agents/handler.py index 5256e5ea..3fe55f3d 100644 --- a/packages/claude-agents/src/launchdarkly_ai_claude_agents/handler.py +++ b/packages/claude-agents/src/launchdarkly_ai_claude_agents/handler.py @@ -24,6 +24,8 @@ HookMatcher, ResultMessage, StreamEvent, + ToolResultBlock, + UserMessage, query, ) @@ -48,6 +50,7 @@ set_output_content_attributes, set_tool_call_content_attributes, ) +from launchdarkly_ai_server.agui import TurnEmitter from launchdarkly_ai_server.lifecycle import effective_capture_content from ._version import PACKAGE_NAME, __version__ @@ -198,6 +201,9 @@ def build_tool_hooks( native_tool_map: dict[str, Any], parent_context: Any, capture_content: bool, + *, + emitter: TurnEmitter | None = None, + pending: list[dict[str, Any]] | None = None, ) -> tuple[ dict[str, list[HookMatcher]], ToolTelemetry, AbandonTelemetry, CancelTelemetry ]: @@ -217,6 +223,11 @@ def _finish( ) -> None: if tool_use_id is None: return + if emitter is not None and pending is not None: + if error is None: + pending.extend(emitter.tool_result(tool_use_id, result)) + else: + pending.extend(emitter.tool_result(tool_use_id, failed=True)) span = tool_spans.pop(tool_use_id, None) if span is None: return @@ -258,6 +269,11 @@ async def _pre_tool_use( # out of an unfiled span leaves it open with nothing tracking it: close_open_spans and the # teardown both walk this dict, so a span missing from it is a span that never exports. tool_spans[use_id] = span + if emitter is not None and pending is not None and use_id: + pending.extend(emitter.tool_start(use_id, display_name)) + pending.extend( + emitter.tool_end(use_id, arguments=input_data.get("tool_input")) + ) set_tool_call_content_attributes( span, capture_content, arguments=input_data.get("tool_input") ) @@ -720,6 +736,51 @@ def _stream_impl( ) +def _map_claude_stream_event( + event: dict[str, Any], + emitter: TurnEmitter, + block_ids: dict[int, str], + saw_args: set[str], +) -> list[dict[str, Any]]: + """Map one Claude Agent SDK stream event onto the AG-UI subset.""" + kind = event.get("type") + if kind == "content_block_start": + block = event.get("content_block") or {} + if not isinstance(block, dict) or block.get("type") != "tool_use": + return [] + index = event.get("index") + call_id = str(block.get("id") or "") + if isinstance(index, int): + block_ids[index] = call_id + return emitter.tool_start( + call_id, tool_display_name(str(block.get("name") or "")) + ) + if kind == "content_block_delta": + delta = event.get("delta") or {} + if not isinstance(delta, dict): + return [] + if delta.get("type") == "text_delta": + text = delta.get("text") or "" + if not isinstance(text, str) or not text: + return [] + return emitter.text(text) + if delta.get("type") == "input_json_delta": + index = event.get("index") + call_id = block_ids.get(index, "") if isinstance(index, int) else "" + partial = delta.get("partial_json") or "" + if isinstance(partial, str) and partial and call_id: + saw_args.add(call_id) + return emitter.tool_args( + call_id, partial if isinstance(partial, str) else "" + ) + if kind == "content_block_stop": + index = event.get("index") + call_id = block_ids.get(index, "") if isinstance(index, int) else "" + if call_id and call_id in saw_args: + return emitter.tool_end(call_id) + return [] + + async def _stream_gen( config: AiConfigRep, user_input: str, @@ -761,6 +822,10 @@ async def _stream_gen( # stream; a CancelledError means something cancelled the run, usually a timeout, and the # consumer chose nothing. The blocking path already tells these apart. cancelled = False + emitter = TurnEmitter(capture_content=capture_content) + pending: list[dict[str, Any]] = [] + block_ids: dict[int, str] = {} + saw_args: set[str] = set() try: # Inside the guard, because serialising the prompt raises on anything that is not # JSON-serialisable. A raise out here would leave the root open with the `finally` never @@ -781,7 +846,11 @@ async def _stream_gen( hooks = None if native_tool_names or mcp_allowed_tools: hooks, tool_telemetry, abandon_tool_spans, _ = build_tool_hooks( - native_tool_map, parent, capture_content + native_tool_map, + parent, + capture_content, + emitter=emitter, + pending=pending, ) options = _build_query_options( @@ -797,20 +866,37 @@ async def _stream_gen( full_output = "" gen = query(prompt=query_prompt, options=options) async for message in gen: + while pending: + yield pending.pop(0) record_conversation_id(span, message) record_native_tools(span, message, capture_content, catalog) inference.record(message) if isinstance(message, StreamEvent): event = message.event - if ( - event.get("type") == "content_block_delta" - and event.get("delta", {}).get("type") == "text_delta" + if event.get("type") == "message_start": + for rich in emitter.end_text(): + yield rich + for rich in _map_claude_stream_event( + event, emitter, block_ids, saw_args ): - text = event["delta"].get("text", "") - if text: - yield {"type": "chunk", "text": text} - full_output += text + if rich.get("type") == "TEXT_MESSAGE_CONTENT": + full_output += rich.get("delta") or "" + yield rich + elif isinstance(message, UserMessage): + blocks = message.content if isinstance(message.content, list) else [] + for block in blocks: + if not isinstance(block, ToolResultBlock): + continue + # Built-in tools have no hook when the config lists none. + # The SDK still yields the result, and that is what closes + # the call instead of fail_open marking it failed. + for rich in emitter.tool_result( + str(block.tool_use_id), + block.content, + failed=bool(getattr(block, "is_error", False)), + ): + yield rich elif isinstance(message, ResultMessage): # Before the root, so the children it parents are already closed. inference.finish() @@ -839,6 +925,12 @@ async def _stream_gen( ) mark_ok(span) end_span_once(span, ended) + while pending: + yield pending.pop(0) + for rich in emitter.end_text(): + yield rich + for rich in emitter.fail_open(): + yield rich yield {"type": "done", "output": final_output, "usage": raw_usage} return @@ -862,6 +954,12 @@ async def _stream_gen( ) mark_ok(span) end_span_once(span, ended) + while pending: + yield pending.pop(0) + for rich in emitter.end_text(): + yield rich + for rich in emitter.fail_open(): + yield rich yield {"type": "done", "output": full_output, "usage": streamed_usage["total"]} except asyncio.CancelledError: @@ -874,6 +972,10 @@ async def _stream_gen( finish_root_span(span, config, inference.run_usage["total"]) root_usage_written = True fail_span(span, exc, ended) + while pending: + yield pending.pop(0) + for rich in emitter.fail_open(): + yield rich raise finally: # A no-op on the success path, where the result message already closed the last response. diff --git a/packages/claude-agents/tests/test_handler.py b/packages/claude-agents/tests/test_handler.py index 765367d9..74f64e24 100644 --- a/packages/claude-agents/tests/test_handler.py +++ b/packages/claude-agents/tests/test_handler.py @@ -1024,8 +1024,13 @@ async def test_yields_chunks_then_one_done_event( ) h = create_claude_agents_handler() events = await _collect(await h.stream(BASE_CONFIG, "q", {}, {})) - assert [e["type"] for e in events] == ["chunk", "done"] - assert events[0]["text"] == "Hi" + assert [e["type"] for e in events] == [ + "TEXT_MESSAGE_START", + "TEXT_MESSAGE_CONTENT", + "TEXT_MESSAGE_END", + "done", + ] + assert events[1]["delta"] == "Hi" assert events[-1]["output"] == "Hi" async def test_emits_same_span_tree_as_blocking_path( @@ -1075,7 +1080,7 @@ async def _slow_query(**_kwargs: Any) -> AsyncIterator[Any]: gen = await h.stream(BASE_CONFIG, "q", {}, {}) # Consume only the first chunk, then abandon the generator without exhausting it. first = await gen.__anext__() - assert first["type"] == "chunk" + assert first["type"] == "TEXT_MESSAGE_START" await gen.aclose() assert root().status.status_code == StatusCode.UNSET @@ -1120,9 +1125,8 @@ async def test_yields_chunk_events_for_text_deltas( h = create_claude_agents_handler() events = [e async for e in await h.stream(BASE_CONFIG, "hi")] - chunks = [e for e in events if e.get("type") == "chunk"] - assert len(chunks) == 2 - assert chunks[0]["text"] == "hello " + deltas = [e["delta"] for e in events if e.get("type") == "TEXT_MESSAGE_CONTENT"] + assert deltas == ["hello ", "world"] async def test_all_chunks_before_done( self, monkeypatch: pytest.MonkeyPatch @@ -1136,8 +1140,13 @@ async def test_all_chunks_before_done( events = [e async for e in await h.stream(BASE_CONFIG, "hi")] done_idx = next(i for i, e in enumerate(events) if e.get("type") == "done") - chunk_indices = [i for i, e in enumerate(events) if e.get("type") == "chunk"] - assert all(ci < done_idx for ci in chunk_indices) + text_indices = [ + i + for i, e in enumerate(events) + if e.get("type") + in {"TEXT_MESSAGE_START", "TEXT_MESSAGE_CONTENT", "TEXT_MESSAGE_END"} + ] + assert all(ci < done_idx for ci in text_indices) async def test_yields_exactly_one_done_event( self, monkeypatch: pytest.MonkeyPatch @@ -1229,7 +1238,7 @@ async def _query(**_kwargs: Any) -> AsyncIterator[Any]: h = create_claude_agents_handler() gen = await h.stream(BASE_CONFIG, "q", {}, {}) first = await gen.__anext__() - assert first["type"] == "chunk" + assert first["type"] == "TEXT_MESSAGE_START" await gen.aclose() assert closed["value"] is True @@ -1953,7 +1962,7 @@ async def _query_that_opens_a_tool(**kwargs: Any) -> AsyncIterator[Any]: monkeypatch.setattr(handler_mod, "query", _query_that_opens_a_tool) h = create_claude_agents_handler() gen = await h.stream(config, "q", {"webSearch": stub}, {}) - assert (await gen.__anext__())["type"] == "chunk" + assert (await gen.__anext__())["type"] == "TOOL_CALL_START" await gen.aclose() tools = [s for s in spans() if s.name.startswith("execute_tool ")] @@ -2254,3 +2263,93 @@ async def _query(**kwargs: Any) -> AsyncIterator[Any]: tool = named("execute_tool ")[0] assert tool.end_time is not None + + +async def test_a_builtin_tool_result_is_not_marked_failed( + monkeypatch: pytest.MonkeyPatch, +) -> None: + started = StreamEvent( + uuid="u1", + session_id="sess-1", + event={ + "type": "content_block_start", + "index": 0, + "content_block": {"type": "tool_use", "id": "tu-1", "name": "Read"}, + }, + ) + args = StreamEvent( + uuid="u2", + session_id="sess-1", + event={ + "type": "content_block_delta", + "index": 0, + "delta": {"type": "input_json_delta", "partial_json": '{"file":"a"}'}, + }, + ) + stopped = StreamEvent( + uuid="u3", + session_id="sess-1", + event={"type": "content_block_stop", "index": 0}, + ) + monkeypatch.setattr( + handler_mod, + "query", + _fake_query( + [ + started, + args, + stopped, + tool_result_user_message("tu-1", "file body"), + result_message("done"), + ] + ), + ) + + async def _events(capture: bool) -> list[dict[str, Any]]: + handler = create_claude_agents_handler(capture_content=capture) + return [event async for event in await handler.stream(BASE_CONFIG, "q", {}, {})] + + shown = await _events(True) + types = [event["type"] for event in shown if str(event["type"]).startswith("TOOL_")] + assert types == [ + "TOOL_CALL_START", + "TOOL_CALL_ARGS", + "TOOL_CALL_END", + "TOOL_CALL_RESULT", + ] + assert shown[0]["toolCallName"] == "Read" or any( + event.get("toolCallName") == "Read" for event in shown + ) + result = next(event for event in shown if event["type"] == "TOOL_CALL_RESULT") + assert result["content"] == "file body" + assert "metadata" not in result + + hidden = await _events(False) + assert not any(event["type"] == "TOOL_CALL_ARGS" for event in hidden) + hidden_result = next( + event for event in hidden if event["type"] == "TOOL_CALL_RESULT" + ) + assert hidden_result["content"] == "" + + +def test_a_streamed_mcp_tool_uses_the_config_name() -> None: + from launchdarkly_ai_claude_agents.handler import _map_claude_stream_event + from launchdarkly_ai_server.agui import TurnEmitter + + emitter = TurnEmitter(capture_content=True) + events = _map_claude_stream_event( + { + "type": "content_block_start", + "index": 0, + "content_block": { + "type": "tool_use", + "id": "tu-1", + "name": "mcp__tool-mcp__search", + }, + }, + emitter, + {}, + set(), + ) + assert events[0]["toolCallName"] == "search" + assert events[0]["toolCallId"] == "tu-1" diff --git a/packages/claude-messages/src/launchdarkly_ai_claude_messages/handler.py b/packages/claude-messages/src/launchdarkly_ai_claude_messages/handler.py index 391ca5df..69e82fdd 100644 --- a/packages/claude-messages/src/launchdarkly_ai_claude_messages/handler.py +++ b/packages/claude-messages/src/launchdarkly_ai_claude_messages/handler.py @@ -25,6 +25,7 @@ set_tool_call_content_attributes, to_semconv_finish_reason, ) +from launchdarkly_ai_server.agui import TurnEmitter from launchdarkly_ai_server.lifecycle import effective_capture_content from ._version import PACKAGE_NAME, __version__ @@ -485,6 +486,37 @@ def _stream_impl( ) +def _map_claude_event( + event: Any, emitter: TurnEmitter, block_ids: dict[int, str] +) -> list[dict[str, Any]]: + """Map one Messages API stream event onto the AG-UI subset.""" + kind = getattr(event, "type", None) + if kind == "content_block_start": + block = getattr(event, "content_block", None) + index = getattr(event, "index", None) + if getattr(block, "type", None) != "tool_use" or not isinstance(index, int): + return [] + call_id = str(getattr(block, "id", "") or "") + block_ids[index] = call_id + return emitter.tool_start(call_id, str(getattr(block, "name", "") or "")) + if kind == "content_block_delta": + delta = getattr(event, "delta", None) + delta_type = getattr(delta, "type", None) + if delta_type == "text_delta": + text = getattr(delta, "text", "") + if not isinstance(text, str) or not text: + return [] + return emitter.text(text) + if delta_type == "input_json_delta": + index = getattr(event, "index", None) + call_id = block_ids.get(index, "") if isinstance(index, int) else "" + partial = getattr(delta, "partial_json", "") + if not isinstance(partial, str): + return [] + return emitter.tool_args(call_id, partial) + return [] + + async def _stream_gen( client: Any, config: AiConfigRep, @@ -520,6 +552,8 @@ async def _stream_gen( conversation = list(messages) full_output = "" steps = 0 + emitter = TurnEmitter(capture_content=capture_content) + block_ids: dict[int, str] = {} ended: set[int] = set() open_model_span: Any = None @@ -575,15 +609,10 @@ async def _stream_gen( stream = client.messages.stream(**kwargs) async with stream as s: async for event in s: - if ( - hasattr(event, "type") - and event.type == "content_block_delta" - and hasattr(event, "delta") - and getattr(event.delta, "type", None) == "text_delta" - ): - text = event.delta.text - full_output += text - yield {"type": "chunk", "text": text} + for rich in _map_claude_event(event, emitter, block_ids): + if rich.get("type") == "TEXT_MESSAGE_CONTENT": + full_output += rich.get("delta") or "" + yield rich final_msg = await s.get_final_message() @@ -622,6 +651,10 @@ async def _stream_gen( raise if final_msg.stop_reason != "tool_use": + for rich in emitter.end_text(): + yield rich + for rich in emitter.fail_open(): + yield rich break if steps >= _MAX_STEPS: @@ -631,6 +664,17 @@ async def _stream_gen( steps += 1 conversation.append({"role": "assistant", "content": final_msg.content}) + for block in final_msg.content: + if getattr(block, "type", None) != "tool_use": + continue + for rich in emitter.tool_start(str(block.id), str(block.name)): + yield rich + for rich in emitter.tool_end( + str(block.id), arguments=getattr(block, "input", None) + ): + yield rich + for rich in emitter.end_text(): + yield rich tool_results = [] for block in final_msg.content: if block.type != "tool_use": @@ -664,6 +708,8 @@ async def _stream_gen( # Cleared on both paths that end the span, and deliberately not in a `finally`: # a `finally` would also clear it for a BaseException, which is the one case where # the span is still open and the outer `finally` is the only thing left to close it. + for rich in emitter.tool_result(str(block.id), result): + yield rich tool_results.append( { "type": "tool_result", @@ -671,6 +717,8 @@ async def _stream_gen( "content": str(result), } ) + for rich in emitter.fail_open(): + yield rich conversation.append({"role": "user", "content": tool_results}) set_output_content_attributes( @@ -700,6 +748,8 @@ async def _stream_gen( if run_usage.reported: finish_root_span(span, config, run_usage.total) fail_span(span, exc, ended) + for rich in emitter.fail_open(): + yield rich raise finally: # A no-op on the success and failure paths, because both already ended their spans through diff --git a/packages/claude-messages/tests/test_handler.py b/packages/claude-messages/tests/test_handler.py index e107a7a5..bd2e3a90 100644 --- a/packages/claude-messages/tests/test_handler.py +++ b/packages/claude-messages/tests/test_handler.py @@ -1303,10 +1303,179 @@ async def test_yields_chunk_events_for_text_deltas( with patch.object(spans_mod, "_HAS_OTEL", False): h = create_claude_messages_handler() events = [e async for e in await h.stream(CONFIG, "q")] - chunks = [e for e in events if e.get("type") == "chunk"] - assert len(chunks) == 2 - assert chunks[0]["text"] == "hello " - assert chunks[1]["text"] == "world" + deltas = [e["delta"] for e in events if e.get("type") == "TEXT_MESSAGE_CONTENT"] + assert deltas == ["hello ", "world"] + assert events[0]["type"] == "TEXT_MESSAGE_START" + assert events[-2]["type"] == "TEXT_MESSAGE_END" + + async def test_tool_lifecycle_hides_payloads_until_capture_is_on( + self, mock_anthropic: MagicMock + ) -> None: + import launchdarkly_ai_claude_messages.spans as spans_mod + from launchdarkly_ai_claude_messages import create_claude_messages_handler + + block = MagicMock() + block.type = "tool_use" + block.id = "c1" + block.name = "search" + block.input = {"q": "x"} + + async def search(_args: dict[str, Any]) -> str: + return "secret" + + async def _events(capture: bool) -> list[dict[str, Any]]: + first_ctx, first_final = _make_stream_context(["Looking"]) + first_final.stop_reason = "tool_use" + first_final.content = [block] + second_ctx, _second_final = _make_stream_context(["Done"]) + mock_anthropic.messages.stream = MagicMock( + side_effect=[first_ctx, second_ctx] + ) + with patch.object(spans_mod, "_HAS_OTEL", False): + handler = create_claude_messages_handler(capture_content=capture) + return [ + event + async for event in await handler.stream( + CONFIG, "q", {"search": search} + ) + ] + + hidden = await _events(False) + start = next(event for event in hidden if event["type"] == "TOOL_CALL_START") + assert start["toolCallId"] == "c1" + assert start["toolCallName"] == "search" + assert start["parentMessageId"] + assert not any(event["type"] == "TOOL_CALL_ARGS" for event in hidden) + result = next(event for event in hidden if event["type"] == "TOOL_CALL_RESULT") + assert result["content"] == "" + + shown = await _events(True) + args = next(event for event in shown if event["type"] == "TOOL_CALL_ARGS") + assert args["delta"] == '{"q":"x"}' + shown_result = next( + event for event in shown if event["type"] == "TOOL_CALL_RESULT" + ) + assert shown_result["content"] == "secret" + + async def test_a_streamed_tool_without_tool_use_is_closed( + self, mock_anthropic: MagicMock + ) -> None: + import launchdarkly_ai_claude_messages.spans as spans_mod + from launchdarkly_ai_claude_messages import create_claude_messages_handler + + started = MagicMock() + started.type = "content_block_start" + started.index = 0 + block = MagicMock() + block.type = "tool_use" + block.id = "c1" + block.name = "search" + started.content_block = block + _ctx, final = _make_stream_context([]) + final.stop_reason = "end_turn" + + class _Stream: + def __aiter__(self) -> AsyncIterator[Any]: + return self._iter() + + async def _iter(self) -> AsyncIterator[Any]: + yield started + yield _make_stream_event("Hi") + + async def get_final_message(self) -> Any: + return final + + @asynccontextmanager + async def _mgr() -> AsyncGenerator[Any, None]: + yield _Stream() + + mock_anthropic.messages.stream = MagicMock(return_value=_mgr()) + with patch.object(spans_mod, "_HAS_OTEL", False): + handler = create_claude_messages_handler() + events = [event async for event in await handler.stream(CONFIG, "q")] + start = next(event for event in events if event["type"] == "TOOL_CALL_START") + assert start["toolCallId"] == "c1" + assert start["toolCallName"] == "search" + assert any( + event["type"] == "TOOL_CALL_RESULT" and event["toolCallId"] == "c1" + for event in events + ) + + async def test_a_failed_tool_closes_its_siblings( + self, mock_anthropic: MagicMock + ) -> None: + import launchdarkly_ai_claude_messages.spans as spans_mod + from launchdarkly_ai_claude_messages import create_claude_messages_handler + + ctx, final = _make_stream_context([]) + final.stop_reason = "tool_use" + final.content = [ + _tool_use_block("boom", "c1"), + _tool_use_block("later", "c2"), + ] + mock_anthropic.messages.stream = MagicMock(return_value=ctx) + + def boom(_args: dict[str, Any]) -> str: + raise RuntimeError("nope") + + seen: list[dict[str, Any]] = [] + with ( + patch.object(spans_mod, "_HAS_OTEL", False), + pytest.raises(RuntimeError, match="nope"), + ): + handler = create_claude_messages_handler() + async for event in await handler.stream( + CONFIG, "q", {"boom": boom, "later": lambda _args: "ok"} + ): + seen.append(event) + results = { + event["toolCallId"] for event in seen if event["type"] == "TOOL_CALL_RESULT" + } + assert results == {"c1", "c2"} + + async def test_a_stream_error_closes_an_opened_tool( + self, mock_anthropic: MagicMock + ) -> None: + import launchdarkly_ai_claude_messages.spans as spans_mod + from launchdarkly_ai_claude_messages import create_claude_messages_handler + + started = MagicMock() + started.type = "content_block_start" + started.index = 0 + block = MagicMock() + block.type = "tool_use" + block.id = "c1" + block.name = "search" + started.content_block = block + + class _Stream: + def __aiter__(self) -> AsyncIterator[Any]: + return self._iter() + + async def _iter(self) -> AsyncIterator[Any]: + yield started + raise RuntimeError("stream broke") + + async def get_final_message(self) -> Any: + raise AssertionError("final message is not read") + + @asynccontextmanager + async def _mgr() -> AsyncGenerator[Any, None]: + yield _Stream() + + mock_anthropic.messages.stream = MagicMock(return_value=_mgr()) + seen: list[dict[str, Any]] = [] + with ( + patch.object(spans_mod, "_HAS_OTEL", False), + pytest.raises(RuntimeError, match="stream broke"), + ): + handler = create_claude_messages_handler() + async for event in await handler.stream(CONFIG, "q"): + seen.append(event) + assert any( + event["type"] == "TOOL_CALL_RESULT" and event["toolCallId"] == "c1" + for event in seen + ) async def test_all_chunks_before_done(self, mock_anthropic: MagicMock) -> None: import launchdarkly_ai_claude_messages.spans as spans_mod