feat(claude): emit AG-UI text and tool events

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.

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"