From fb516505c432d2b7f6ec00403c64a2a0d29ed868 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 13 Oct 2026 14:18:55 +0000 Subject: [PATCH 1/8] feat(agui): add a shared reader for handler streams Handlers can yield the AG-UI text and tool subset. config().stream() still yields chunk then done. stream() is the opt-in entry. Co-authored-by: alexis --- .../client/src/launchdarkly_ai_server/agui.py | 584 ++++++++++++++++++ .../src/launchdarkly_ai_server/client.py | 58 +- .../src/launchdarkly_ai_server/graph.py | 31 +- .../src/launchdarkly_ai_server/tracking.py | 37 +- .../src/launchdarkly_ai_server/types.py | 7 +- packages/client/tests/test_agui.py | 469 ++++++++++++++ 6 files changed, 1150 insertions(+), 36 deletions(-) create mode 100644 packages/client/src/launchdarkly_ai_server/agui.py create mode 100644 packages/client/tests/test_agui.py diff --git a/packages/client/src/launchdarkly_ai_server/agui.py b/packages/client/src/launchdarkly_ai_server/agui.py new file mode 100644 index 00000000..bdf7ec17 --- /dev/null +++ b/packages/client/src/launchdarkly_ai_server/agui.py @@ -0,0 +1,584 @@ +"""Opt-in AG-UI reader for one AI Config turn. + +Handlers may yield ``chunk`` or the text and tool subset below, then one +``done``. :func:`stream` is the opt-in entry: it brackets that subset with +``RUN_STARTED`` and ``RUN_FINISHED``. ``config().stream()`` still yields +``chunk`` then ``done``. The SDK runs the tool loop. This stream only reports it. + +Argument deltas and tool results are content. :class:`TurnEmitter` omits them +when the handler was built with ``capture_content`` off, and still emits the +tool lifecycle. +""" + +from __future__ import annotations + +import json +import math +import uuid +from collections.abc import AsyncGenerator +from typing import Any + +from .conversation import bind_conversation_id +from .sdk_usage import report_usage +from .types import JudgeResult, LDContext + +PROTOCOL_VERSION = "1.0" + +TEXT_MESSAGE_START = "TEXT_MESSAGE_START" +TEXT_MESSAGE_CONTENT = "TEXT_MESSAGE_CONTENT" +TEXT_MESSAGE_END = "TEXT_MESSAGE_END" +TOOL_CALL_START = "TOOL_CALL_START" +TOOL_CALL_ARGS = "TOOL_CALL_ARGS" +TOOL_CALL_END = "TOOL_CALL_END" +TOOL_CALL_RESULT = "TOOL_CALL_RESULT" +RUN_STARTED = "RUN_STARTED" +RUN_FINISHED = "RUN_FINISHED" + +TEXT_TYPES = frozenset({TEXT_MESSAGE_START, TEXT_MESSAGE_CONTENT, TEXT_MESSAGE_END}) +TOOL_TYPES = frozenset( + {TOOL_CALL_START, TOOL_CALL_ARGS, TOOL_CALL_END, TOOL_CALL_RESULT} +) +HANDLER_RICH_TYPES = TEXT_TYPES | TOOL_TYPES + +_FAILURE_CONTENT = "An error occurred." + + +def argument_text(value: Any) -> str | None: + """The argument text for one ``TOOL_CALL_ARGS`` delta. + + A string is that string. Any other JSON value is compact JSON. ``None`` and + a value that is not JSON under :func:`json.dumps` with ``allow_nan=False`` + produce no delta. + """ + if value is None: + return None + if isinstance(value, str): + return value or None + if isinstance(value, (bytes, bytearray, memoryview)): + return None + if isinstance(value, float) and (math.isnan(value) or math.isinf(value)): + return None + try: + return json.dumps( + value, ensure_ascii=False, separators=(",", ":"), allow_nan=False + ) + except (TypeError, ValueError): + return None + + +def safe_tool_content(value: Any) -> tuple[str, bool]: + """``(content, failed)`` for a tool result. + + A string is not encoded again. Any other JSON value is compact JSON. + A value that is not a string and not serializable under that rule is a + failed result: the caller masks it. The model-facing value is unchanged. + """ + if isinstance(value, str): + return value, False + if isinstance(value, (bytes, bytearray, memoryview)): + return _FAILURE_CONTENT, True + if isinstance(value, float) and (math.isnan(value) or math.isinf(value)): + return _FAILURE_CONTENT, True + try: + return ( + json.dumps( + value, ensure_ascii=False, separators=(",", ":"), allow_nan=False + ), + False, + ) + except (TypeError, ValueError): + return _FAILURE_CONTENT, True + + +class TurnEmitter: + """Yields the text and tool subset for one handler run. + + ``capture_content`` gates argument deltas and result content. The tool name + and call id are always present on ``TOOL_CALL_START``. A closed call is not + reopened. + """ + + def __init__(self, *, capture_content: bool) -> None: + self.capture_content = capture_content + self._open: str | None = None + self._seq = 0 + self._started: set[str] = set() + self._args: set[str] = set() + self._ended: set[str] = set() + self._resulted: set[str] = set() + + def _mint(self, prefix: str) -> str: + self._seq += 1 + return f"{prefix}-{self._seq}" + + def text(self, delta: str) -> list[dict[str, Any]]: + """Append assistant text. The first slice opens the message.""" + if not delta: + return [] + events: list[dict[str, Any]] = [] + if self._open is None: + self._open = self._mint("msg") + events.append({"type": TEXT_MESSAGE_START, "messageId": self._open}) + events.append( + { + "type": TEXT_MESSAGE_CONTENT, + "messageId": self._open, + "delta": delta, + } + ) + return events + + def end_text(self) -> list[dict[str, Any]]: + if self._open is None: + return [] + event = {"type": TEXT_MESSAGE_END, "messageId": self._open} + self._open = None + return [event] + + def tool_start( + self, tool_call_id: str, tool_call_name: str + ) -> list[dict[str, Any]]: + if not tool_call_id or not tool_call_name or tool_call_id in self._started: + return [] + self._started.add(tool_call_id) + event: dict[str, Any] = { + "type": TOOL_CALL_START, + "toolCallId": tool_call_id, + "toolCallName": tool_call_name, + } + if self._open is not None: + event["parentMessageId"] = self._open + return [event] + + def tool_args(self, tool_call_id: str, delta: str) -> list[dict[str, Any]]: + if ( + not self.capture_content + or not delta + or tool_call_id not in self._started + or tool_call_id in self._ended + ): + return [] + self._args.add(tool_call_id) + return [{"type": TOOL_CALL_ARGS, "toolCallId": tool_call_id, "delta": delta}] + + def tool_end( + self, tool_call_id: str, *, arguments: Any = None + ) -> list[dict[str, Any]]: + """Close the argument text. One delta is emitted when none arrived.""" + if tool_call_id not in self._started or tool_call_id in self._ended: + return [] + events: list[dict[str, Any]] = [] + if ( + self.capture_content + and tool_call_id not in self._args + and arguments is not None + ): + text = argument_text(arguments) + if text: + self._args.add(tool_call_id) + events.append( + {"type": TOOL_CALL_ARGS, "toolCallId": tool_call_id, "delta": text} + ) + self._ended.add(tool_call_id) + events.append({"type": TOOL_CALL_END, "toolCallId": tool_call_id}) + return events + + def tool_result( + self, tool_call_id: str, value: Any = None, *, failed: bool = False + ) -> list[dict[str, Any]]: + """Report a finished call, including failure. Does not run the tool.""" + if tool_call_id not in self._started or tool_call_id in self._resulted: + return [] + events = self.tool_end(tool_call_id) if tool_call_id not in self._ended else [] + self._resulted.add(tool_call_id) + event: dict[str, Any] = { + "type": TOOL_CALL_RESULT, + "messageId": self._mint("toolmsg"), + "toolCallId": tool_call_id, + } + if failed: + event["metadata"] = {"launchdarkly": {"toolFailed": True}} + if self.capture_content: + event["content"] = _FAILURE_CONTENT + elif self.capture_content: + content, bad = safe_tool_content(value) + if bad: + event["content"] = _FAILURE_CONTENT + event["metadata"] = {"launchdarkly": {"toolFailed": True}} + else: + event["content"] = content + events.append(event) + return events + + +class StableProjector: + """Projects a handler stream onto ``chunk`` then one internal ``done``. + + ``TEXT_MESSAGE_CONTENT`` becomes a chunk. Tool events and any other type are + dropped and do not close the turn. Only the handler's ``done`` does. + """ + + def __init__(self) -> None: + self.saw_text = False + self.done: dict[str, Any] | None = None + + def feed(self, event: dict[str, Any]) -> list[dict[str, Any]]: + typ = event.get("type") + if typ in TEXT_TYPES: + self.saw_text = True + if typ == "done": + self.done = event + return [] + if typ == "chunk": + if self.saw_text: + return [] + text = event.get("text", "") + if not isinstance(text, str): + return [] + return [{"type": "chunk", "text": text}] + if typ == TEXT_MESSAGE_CONTENT: + delta = event.get("delta") or "" + if not delta: + return [] + return [{"type": "chunk", "text": delta}] + return [] + + +def _token_usage( + parsed: dict[str, Any], *, provider: str, model: str +) -> dict[str, Any]: + usage: dict[str, Any] = { + "inputTokens": int(parsed.get("input") or 0), + "outputTokens": int(parsed.get("output") or 0), + "totalTokens": int(parsed.get("total") or 0), + } + details = parsed.get("input_details") + if isinstance(details, dict): + usage["cachedInputTokens"] = int(details.get("cache_read") or 0) + usage["cacheWriteInputTokens"] = int(details.get("cache_creation") or 0) + if provider: + usage["provider"] = provider + if model: + usage["model"] = model + return usage + + +def _judge_payload(results: dict[str, JudgeResult]) -> dict[str, Any] | None: + if not results: + return None + payload: dict[str, Any] = {} + for key, result in results.items(): + usage: dict[str, Any] = { + "input": result.usage.input, + "output": result.usage.output, + "total": result.usage.total, + } + details = result.usage.input_details + if details is not None: + usage["input_details"] = { + "uncached": details.uncached, + "cache_read": details.cache_read, + "cache_creation": details.cache_creation, + } + payload[key] = { + "usage": usage, + "response": result.response, + "score": result.score, + } + return payload + + +class RunReader: + """Maps one handler stream onto one AG-UI run. + + Legacy ``chunk`` events become one text message when the handler yielded no + ``TEXT_MESSAGE_*`` event. ``RUN_STARTED`` and ``RUN_FINISHED`` are minted + here. The handler's ``done`` is consumed, not forwarded. + """ + + def __init__(self) -> None: + self.thread_id = str(uuid.uuid4()) + self.run_id = str(uuid.uuid4()) + self._mode: str | None = None + self._open: str | None = None + self._deltas: list[str] = [] + self._started: set[str] = set() + self._ended: set[str] = set() + self._resulted: set[str] = set() + self._seq = 0 + + def start(self) -> dict[str, Any]: + return { + "type": RUN_STARTED, + "threadId": self.thread_id, + "runId": self.run_id, + "protocolVersion": PROTOCOL_VERSION, + } + + def _mint(self, prefix: str) -> str: + self._seq += 1 + return f"{prefix}-{self._seq}" + + def feed(self, event: dict[str, Any]) -> list[dict[str, Any]]: + typ = event.get("type") + if typ == "chunk": + return self._chunk(event.get("text") or "") + if typ == TEXT_MESSAGE_START: + return self._text_start(event.get("messageId")) + if typ == TEXT_MESSAGE_CONTENT: + return self._text_content(event.get("messageId"), event.get("delta")) + if typ == TEXT_MESSAGE_END: + return self._text_end(event.get("messageId")) + if typ == TOOL_CALL_START: + return self._tool_start(event) + if typ == TOOL_CALL_ARGS: + return self._tool_args(event) + if typ == TOOL_CALL_END: + return self._tool_end(event) + if typ == TOOL_CALL_RESULT: + return self._tool_result(event) + return [] + + def finish( + self, + done: dict[str, Any], + judge_results: dict[str, JudgeResult] | None, + *, + provider: str, + model: str, + ) -> list[dict[str, Any]]: + events: list[dict[str, Any]] = [] + if self._open is not None: + events.append({"type": TEXT_MESSAGE_END, "messageId": self._open}) + self._open = None + response = done.get("response") + if not isinstance(response, str): + response = "".join(self._deltas) + finished: dict[str, Any] = { + "type": RUN_FINISHED, + "threadId": self.thread_id, + "runId": self.run_id, + } + if response: + finished["result"] = response + parsed = done.get("usage") + if done.get("_had_usage") and isinstance(parsed, dict): + finished["usage"] = [_token_usage(parsed, provider=provider, model=model)] + judges = _judge_payload(judge_results or {}) + if judges is not None: + finished["metadata"] = {"launchdarkly": {"judgeResults": judges}} + events.append(finished) + return events + + def _chunk(self, text: str) -> list[dict[str, Any]]: + if self._mode == "message" or not text: + return [] + self._mode = "chunk" + events: list[dict[str, Any]] = [] + if self._open is None: + self._open = self._mint("msg") + events.append({"type": TEXT_MESSAGE_START, "messageId": self._open}) + events.append( + {"type": TEXT_MESSAGE_CONTENT, "messageId": self._open, "delta": text} + ) + self._deltas.append(text) + return events + + def _close_open(self, message_id: str) -> list[dict[str, Any]]: + if self._open is None or self._open == message_id: + return [] + event = {"type": TEXT_MESSAGE_END, "messageId": self._open} + self._open = None + return [event] + + def _text_start(self, message_id: Any) -> list[dict[str, Any]]: + if not isinstance(message_id, str) or not message_id: + return [] + self._mode = "message" + if self._open == message_id: + return [] + events = self._close_open(message_id) + self._open = message_id + events.append({"type": TEXT_MESSAGE_START, "messageId": message_id}) + return events + + def _text_content(self, message_id: Any, delta: Any) -> list[dict[str, Any]]: + if not isinstance(delta, str): + return [] + self._mode = "message" + mid = message_id if isinstance(message_id, str) and message_id else None + events: list[dict[str, Any]] = [] + if self._open is None: + if mid is None: + mid = self._mint("msg") + self._open = mid + events.append({"type": TEXT_MESSAGE_START, "messageId": mid}) + elif mid is not None and mid != self._open: + events.extend(self._close_open(mid)) + self._open = mid + events.append({"type": TEXT_MESSAGE_START, "messageId": mid}) + events.append( + { + "type": TEXT_MESSAGE_CONTENT, + "messageId": self._open, + "delta": delta, + } + ) + self._deltas.append(delta) + return events + + def _text_end(self, message_id: Any) -> list[dict[str, Any]]: + if self._open is None: + return [] + if isinstance(message_id, str) and message_id and message_id != self._open: + return [] + event = {"type": TEXT_MESSAGE_END, "messageId": self._open} + self._open = None + return [event] + + def _tool_start(self, event: dict[str, Any]) -> list[dict[str, Any]]: + call_id = event.get("toolCallId") + name = event.get("toolCallName") + if ( + not isinstance(call_id, str) + or not call_id + or not isinstance(name, str) + or not name + or call_id in self._started + ): + return [] + self._started.add(call_id) + started: dict[str, Any] = { + "type": TOOL_CALL_START, + "toolCallId": call_id, + "toolCallName": name, + } + parent = event.get("parentMessageId") + if isinstance(parent, str) and parent: + started["parentMessageId"] = parent + return [started] + + def _tool_args(self, event: dict[str, Any]) -> list[dict[str, Any]]: + call_id = event.get("toolCallId") + delta = event.get("delta") + if ( + not isinstance(call_id, str) + or call_id not in self._started + or call_id in self._ended + or not isinstance(delta, str) + ): + return [] + return [{"type": TOOL_CALL_ARGS, "toolCallId": call_id, "delta": delta}] + + def _tool_end(self, event: dict[str, Any]) -> list[dict[str, Any]]: + call_id = event.get("toolCallId") + if ( + not isinstance(call_id, str) + or call_id not in self._started + or call_id in self._ended + ): + return [] + self._ended.add(call_id) + return [{"type": TOOL_CALL_END, "toolCallId": call_id}] + + def _tool_result(self, event: dict[str, Any]) -> list[dict[str, Any]]: + call_id = event.get("toolCallId") + if ( + not isinstance(call_id, str) + or call_id not in self._started + or call_id in self._resulted + ): + return [] + self._resulted.add(call_id) + message_id = event.get("messageId") + if not isinstance(message_id, str) or not message_id: + message_id = self._mint("toolmsg") + result: dict[str, Any] = { + "type": TOOL_CALL_RESULT, + "messageId": message_id, + "toolCallId": call_id, + } + content = event.get("content") + if isinstance(content, str): + result["content"] = content + metadata = event.get("metadata") + if isinstance(metadata, dict): + launchdarkly = metadata.get("launchdarkly") + if ( + isinstance(launchdarkly, dict) + and launchdarkly.get("toolFailed") is True + ): + result["metadata"] = {"launchdarkly": {"toolFailed": True}} + return [result] + + +def stream( + instance: Any, + user_input: str | None, + context: LDContext, + *, + variables: dict[str, Any] | None = None, + history: list[dict[str, Any]] | None = None, +) -> AsyncGenerator[dict[str, Any], None]: + """Stream one config turn as AG-UI events. + + Same arguments as ``ConfigInstance.stream``. ``RUN_STARTED`` is first and + ``RUN_FINISHED`` is last. A failure, including cancellation, raises and does + not yield ``RUN_FINISHED``. + """ + report_usage("client.agui.stream") + return bind_conversation_id( + _stream_events(instance, user_input, context, variables, history) + ) + + +async def _stream_events( + instance: Any, + user_input: str | None, + context: LDContext, + variables: dict[str, Any] | None, + history: list[dict[str, Any]] | None, +) -> AsyncGenerator[dict[str, Any], None]: + from .judges import run_judges + + sink: dict[str, Any] = {} + reader = RunReader() + started = False + done: dict[str, Any] | None = None + async for event in instance._execute_stream( + user_input, context, variables, history, sink + ): + if not started: + yield reader.start() + started = True + if event.get("type") == "done": + done = event + continue + for item in reader.feed(event): + yield item + if not started: + yield reader.start() + if done is None: + done = {"response": "", "usage": {}, "_had_usage": False} + judge_results: dict[str, JudgeResult] = {} + if not instance._skip_judges and sink: + judge_results = await run_judges( + config=sink["config"], + user_context=context, + handler=sink["handler"], + handlers=sink["handlers"], + user_input=user_input, + trajectory=done.get("trajectory", ""), + llm_response=done.get("response", "") + if isinstance(done.get("response"), str) + else "", + base_track_data=done.get("track_data", {}), + tool_handlers=sink["tools"], + ) + for item in reader.finish( + done, + judge_results, + provider=str(sink.get("provider") or ""), + model=str(sink.get("model") or ""), + ): + yield item diff --git a/packages/client/src/launchdarkly_ai_server/client.py b/packages/client/src/launchdarkly_ai_server/client.py index 0f77383f..9ea95962 100644 --- a/packages/client/src/launchdarkly_ai_server/client.py +++ b/packages/client/src/launchdarkly_ai_server/client.py @@ -4,6 +4,7 @@ from collections.abc import AsyncGenerator, Callable from typing import Any +from .agui import StableProjector from .conversation import bind_conversation_id from .judges import _build_judge_tasks, run_judges from .lifecycle import extract_variation @@ -192,13 +193,20 @@ def _stream( self._stream_events(user_input, context, variables, history) ) - async def _stream_events( + async def _execute_stream( self, user_input: str | None, context: LDContext, - variables: dict[str, Any] | None = None, - history: list[dict[str, Any]] | None = None, - ) -> AsyncGenerator[StreamEvent, None]: + variables: dict[str, Any] | None, + history: list[dict[str, Any]] | None, + sink: dict[str, Any], + ) -> AsyncGenerator[dict[str, Any], None]: + """Yield handler events and the internal ``done``. + + Fills *sink* with the resolved config and handler before the first event, + so both ``config().stream()`` and the AG-UI reader run judges from one + pass. + """ resolved_handler_list = resolve_handlers( self._registry, self._normalize_handlers() ) @@ -211,8 +219,17 @@ async def _stream_events( if not resolved_handler_list: raise ValueError("No handlers provided to config()") handler = select_handler(config, meta, resolved_handler_list) - - done_event: dict[str, Any] | None = None + provider = "" + model = "" + if isinstance(config, dict): + provider = str((config.get("provider") or {}).get("name") or "") + model = str((config.get("model") or {}).get("name") or "") + sink["config"] = config + sink["handler"] = handler + sink["handlers"] = resolved_handler_list + sink["tools"] = resolved_tools + sink["provider"] = provider + sink["model"] = model async for event in execute_and_stream( config_key=self._key, @@ -225,10 +242,23 @@ async def _stream_events( variables=variables, history=history, ): - if event.get("type") == "chunk": - yield event - else: - done_event = event + yield event + + async def _stream_events( + self, + user_input: str | None, + context: LDContext, + variables: dict[str, Any] | None = None, + history: list[dict[str, Any]] | None = None, + ) -> AsyncGenerator[StreamEvent, None]: + sink: dict[str, Any] = {} + projector = StableProjector() + async for event in self._execute_stream( + user_input, context, variables, history, sink + ): + for stable in projector.feed(event): + yield stable + done_event = projector.done if done_event: track_data = done_event.get("track_data", {}) @@ -236,15 +266,15 @@ async def _stream_events( {} if self._skip_judges else await run_judges( - config=config, + config=sink["config"], user_context=context, - handler=handler, - handlers=resolved_handler_list, + handler=sink["handler"], + handlers=sink["handlers"], user_input=user_input, trajectory=done_event.get("trajectory", ""), llm_response=done_event.get("response", ""), base_track_data=track_data, - tool_handlers=resolved_tools, + tool_handlers=sink["tools"], ) ) yield { diff --git a/packages/client/src/launchdarkly_ai_server/graph.py b/packages/client/src/launchdarkly_ai_server/graph.py index beb39d97..63e68736 100644 --- a/packages/client/src/launchdarkly_ai_server/graph.py +++ b/packages/client/src/launchdarkly_ai_server/graph.py @@ -10,6 +10,7 @@ from collections.abc import AsyncGenerator, Callable from typing import Any +from .agui import StableProjector from .conversation import bind_conversation_id, bind_span_context from .registry import resolve_handlers, resolve_tools from .types import ( @@ -509,6 +510,7 @@ async def stream_node( "graphKey": key, } + projector = StableProjector() async for event in execute_and_stream( config_key=node.key, config=node.config, @@ -521,17 +523,17 @@ async def stream_node( graph_key=key, history=opts.get("history"), ): - if event.get("type") == "chunk": + for stable in projector.feed(event): yield { "type": "chunk", - "text": event["text"], + "text": stable["text"], "nodeKey": node.key, } - else: - response = event.get("response", "") - usage = event.get("usage") or usage - track_data = event.get("track_data") or track_data - trajectory = event.get("trajectory") or trajectory + if projector.done: + response = projector.done.get("response", "") + usage = projector.done.get("usage") or usage + track_data = projector.done.get("track_data") or track_data + trajectory = projector.done.get("trajectory") or trajectory judge_results = await run_judges( config=node.config, @@ -648,6 +650,7 @@ async def stream_route( merged_tool_handlers = {**(tool_handlers or {}), **handoff_handlers} + projector = StableProjector() async for event in execute_and_stream( config_key=node.key, config=routed_config, @@ -660,17 +663,17 @@ async def stream_route( graph_key=key, history=opts.get("history"), ): - if event.get("type") == "chunk": + for stable in projector.feed(event): yield { "type": "chunk", - "text": event["text"], + "text": stable["text"], "nodeKey": node.key, } - else: - response = event.get("response", "") - usage = event.get("usage") or usage - track_data = event.get("track_data") or track_data - trajectory = event.get("trajectory") or trajectory + if projector.done: + response = projector.done.get("response", "") + usage = projector.done.get("usage") or usage + track_data = projector.done.get("track_data") or track_data + trajectory = projector.done.get("trajectory") or trajectory # Judge against the node's original config, not the routing-augmented one. judge_results = await run_judges( diff --git a/packages/client/src/launchdarkly_ai_server/tracking.py b/packages/client/src/launchdarkly_ai_server/tracking.py index 6e7a5638..a2d3454d 100644 --- a/packages/client/src/launchdarkly_ai_server/tracking.py +++ b/packages/client/src/launchdarkly_ai_server/tracking.py @@ -7,6 +7,7 @@ from collections.abc import AsyncGenerator, Callable from typing import Any +from .agui import HANDLER_RICH_TYPES, TEXT_MESSAGE_CONTENT, TEXT_TYPES from .trajectory import HANDOFF_TOOL_PREFIX, TrajectoryRecorder, render_trajectory from .types import ( NATIVE_TOOL_KEY, @@ -250,9 +251,10 @@ async def execute_and_stream( history: list[dict[str, Any]] | None = None, ) -> AsyncGenerator[ExecuteStreamEvent, None]: """ - Streaming counterpart to ``execute_and_track``. Yields ``chunk`` events - from the handler, then a single ``done`` event once the stream completes - and LD telemetry has been flushed. + Streaming counterpart to ``execute_and_track``. Yields ``chunk`` events, + the AG-UI text and tool subset when the handler emits it, then a single + ``done`` event once the stream completes and LD telemetry has been flushed. + ``config().stream()`` projects that subset back to ``chunk`` then ``done``. Falls back to calling the blocking handler when the handler has no ``stream`` implementation. @@ -297,6 +299,8 @@ async def execute_and_stream( start_time = time.monotonic() full_text = "" raw_usage: dict[str, Any] = {} + had_usage = False + saw_text = False try: if handler.has_stream: @@ -304,20 +308,38 @@ async def execute_and_stream( config, user_input, tracked_tool_handlers, merged_variables, history ) async for event in stream_gen: - if event.get("type") == "chunk": + typ = event.get("type") + if typ in TEXT_TYPES: + saw_text = True + if typ == "chunk": + if saw_text: + continue full_text += event.get("text", "") yield {"type": "chunk", "text": event["text"]} - elif event.get("type") == "done": + elif typ == TEXT_MESSAGE_CONTENT: + delta = event.get("delta") or "" + if isinstance(delta, str): + full_text += delta + yield event + elif typ in HANDLER_RICH_TYPES: + yield event + elif typ == "done": if event.get("output") is not None: full_text = event["output"] - raw_usage = event.get("usage", {}) + usage_value = event.get("usage") + if isinstance(usage_value, dict) and usage_value: + raw_usage = usage_value + had_usage = True else: result = await handler( config, user_input, tracked_tool_handlers, merged_variables, history ) out = result.get("output") full_text = str(out) if out is not None else "" - raw_usage = result.get("usage") or {} + usage_value = result.get("usage") + if isinstance(usage_value, dict) and usage_value: + raw_usage = usage_value + had_usage = True if full_text: yield {"type": "chunk", "text": full_text} except Exception: @@ -348,4 +370,5 @@ async def execute_and_stream( observable_tools=exposed_tool_keys, omitted=recorder.omitted, ), + "_had_usage": had_usage, } diff --git a/packages/client/src/launchdarkly_ai_server/types.py b/packages/client/src/launchdarkly_ai_server/types.py index 44c876fc..09455788 100644 --- a/packages/client/src/launchdarkly_ai_server/types.py +++ b/packages/client/src/launchdarkly_ai_server/types.py @@ -114,7 +114,12 @@ class Message: class HandlerStreamEvent: - """Union of chunk / done events yielded by a handler's stream method.""" + """Events a handler stream yields. + + ``chunk`` and ``done`` remain valid. A handler may instead yield the AG-UI + text and tool subset (``TEXT_MESSAGE_*``, ``TOOL_CALL_*``) and one ``done``. + ``config().stream()`` still yields ``chunk`` then ``done``. + """ pass diff --git a/packages/client/tests/test_agui.py b/packages/client/tests/test_agui.py new file mode 100644 index 00000000..28ed548b --- /dev/null +++ b/packages/client/tests/test_agui.py @@ -0,0 +1,469 @@ +"""The opt-in AG-UI reader, and the stable chunk/done projection.""" + +from __future__ import annotations + +from collections.abc import AsyncGenerator +from typing import Any +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +import launchdarkly_ai_server.lifecycle as lifecycle_module +from launchdarkly_ai_server import ProviderHandler, config +from launchdarkly_ai_server.agui import ( + RUN_FINISHED, + RUN_STARTED, + TEXT_MESSAGE_CONTENT, + TEXT_MESSAGE_END, + TEXT_MESSAGE_START, + TOOL_CALL_ARGS, + TOOL_CALL_END, + TOOL_CALL_RESULT, + TOOL_CALL_START, + RunReader, + StableProjector, + TurnEmitter, + stream, +) +from launchdarkly_ai_server.types import InputTokenDetails, JudgeResult, UsageDict + +CONTEXT = {"kind": "user", "key": "u1"} + + +def _client() -> MagicMock: + client = MagicMock() + client.track = MagicMock() + client.flush = AsyncMock() + client.close = AsyncMock() + client.variation = AsyncMock( + return_value={ + "model": {"name": "gpt-4"}, + "provider": {"name": "TestProvider"}, + "instructions": "Be helpful.", + "_ldMeta": { + "enabled": True, + "variationKey": "v1", + "version": 1, + "mode": "messages", + }, + } + ) + return client + + +@pytest.fixture +def mock_ld_client() -> MagicMock: + client = _client() + lifecycle_module._set_client_for_testing(client) + yield client + lifecycle_module._reset_for_testing() + + +def _handler( + events: list[dict[str, Any]], + *, + capture_content: bool = False, +) -> ProviderHandler: + async def fn(cfg, user_input, tool_handlers, variables, history=None) -> dict: # type: ignore[override] + return {"output": "", "usage": {}} + + async def stream_fn( + cfg, user_input, tool_handlers, variables, history=None + ) -> AsyncGenerator[dict[str, Any], None]: + for event in events: + yield event + + return ProviderHandler( + fn=fn, + provides_for=("TestProvider", "messages"), + stream_fn=stream_fn, + capture_content=capture_content, + ) + + +def _types(events: list[dict[str, Any]]) -> list[str]: + return [event["type"] for event in events] + + +class TestTurnEmitter: + def test_capture_on_includes_args_and_result(self) -> None: + emitter = TurnEmitter(capture_content=True) + events: list[dict[str, Any]] = [] + events.extend(emitter.text("Looking up")) + events.extend(emitter.tool_start("a", "search")) + events.extend(emitter.tool_args("a", '{"q":')) + events.extend(emitter.tool_end("a")) + events.extend(emitter.end_text()) + events.extend(emitter.tool_result("a", {"ok": True})) + assert _types(events) == [ + TEXT_MESSAGE_START, + TEXT_MESSAGE_CONTENT, + TOOL_CALL_START, + TOOL_CALL_ARGS, + TOOL_CALL_END, + TEXT_MESSAGE_END, + TOOL_CALL_RESULT, + ] + assert events[3]["delta"] == '{"q":' + assert events[2]["parentMessageId"] == events[0]["messageId"] + assert events[6]["content"] == '{"ok":true}' + assert events[6]["messageId"] != "a" + + def test_capture_off_keeps_the_lifecycle_without_payloads(self) -> None: + emitter = TurnEmitter(capture_content=False) + events: list[dict[str, Any]] = [] + events.extend(emitter.tool_start("a", "search")) + events.extend(emitter.tool_args("a", '{"q":')) + events.extend(emitter.tool_end("a", arguments={"q": "x"})) + events.extend(emitter.tool_result("a", "secret")) + assert _types(events) == [TOOL_CALL_START, TOOL_CALL_END, TOOL_CALL_RESULT] + assert events[0]["toolCallName"] == "search" + assert "parentMessageId" not in events[0] + assert "content" not in events[2] + assert "delta" not in str(events) + + def test_failure_is_masked_and_a_string_is_not_reencoded(self) -> None: + emitter = TurnEmitter(capture_content=True) + events: list[dict[str, Any]] = [] + events.extend(emitter.tool_start("a", "search")) + events.extend(emitter.tool_result("a", failed=True)) + events.extend(emitter.tool_start("b", "search")) + events.extend(emitter.tool_result("b", "plain")) + failed = events[2] + assert failed["content"] == "An error occurred." + assert failed["metadata"] == {"launchdarkly": {"toolFailed": True}} + assert events[5]["content"] == "plain" + + def test_unserializable_result_is_a_failed_result(self) -> None: + emitter = TurnEmitter(capture_content=True) + events = emitter.tool_start("a", "search") + events.extend(emitter.tool_result("a", float("nan"))) + assert events[-1]["content"] == "An error occurred." + assert events[-1]["metadata"]["launchdarkly"]["toolFailed"] is True + + def test_a_second_start_for_the_same_id_is_dropped(self) -> None: + emitter = TurnEmitter(capture_content=True) + assert emitter.tool_start("a", "search") + assert emitter.tool_start("a", "search") == [] + + +class TestRunReader: + def test_legacy_chunks_become_one_text_message(self) -> None: + reader = RunReader() + events = [reader.start()] + events.extend(reader.feed({"type": "chunk", "text": "Hi"})) + events.extend(reader.feed({"type": "chunk", "text": "!"})) + events.extend( + reader.finish( + { + "response": "Hi!", + "usage": {"input": 2, "output": 1, "total": 3}, + "_had_usage": True, + }, + None, + provider="OpenAI", + model="gpt-4", + ) + ) + assert _types(events) == [ + RUN_STARTED, + TEXT_MESSAGE_START, + TEXT_MESSAGE_CONTENT, + TEXT_MESSAGE_CONTENT, + TEXT_MESSAGE_END, + RUN_FINISHED, + ] + assert events[0]["protocolVersion"] == "1.0" + assert events[0]["threadId"] == events[-1]["threadId"] + assert events[-1]["result"] == "Hi!" + assert events[-1]["usage"] == [ + { + "inputTokens": 2, + "outputTokens": 1, + "totalTokens": 3, + "provider": "OpenAI", + "model": "gpt-4", + } + ] + assert "metadata" not in events[-1] + + def test_forwards_the_subset_and_drops_a_third_type(self) -> None: + reader = RunReader() + events = [reader.start()] + events.extend(reader.feed({"type": TEXT_MESSAGE_START, "messageId": "m1"})) + events.extend( + reader.feed( + {"type": TEXT_MESSAGE_CONTENT, "messageId": "m1", "delta": "Hi"} + ) + ) + events.extend(reader.feed({"type": "widget", "text": "nope"})) + events.extend(reader.feed({"type": "chunk", "text": "dropped"})) + events.extend( + reader.feed( + { + "type": TOOL_CALL_START, + "toolCallId": "a", + "toolCallName": "search", + "parentMessageId": "m1", + } + ) + ) + events.extend( + reader.feed({"type": TOOL_CALL_ARGS, "toolCallId": "a", "delta": "{}"}) + ) + events.extend(reader.feed({"type": TOOL_CALL_END, "toolCallId": "a"})) + events.extend(reader.feed({"type": TEXT_MESSAGE_END, "messageId": "m1"})) + events.extend( + reader.feed( + { + "type": TOOL_CALL_RESULT, + "messageId": "tm1", + "toolCallId": "a", + "content": "ok", + } + ) + ) + events.extend( + reader.feed( + { + "type": TOOL_CALL_START, + "toolCallId": "a", + "toolCallName": "search", + } + ) + ) + events.extend( + reader.finish( + {"response": "Hi", "_had_usage": False}, + None, + provider="", + model="", + ) + ) + assert "widget" not in _types(events) + assert "chunk" not in _types(events) + assert "dropped" not in str(events) + assert _types(events).count(TOOL_CALL_START) == 1 + assert events[-1]["type"] == RUN_FINISHED + assert "usage" not in events[-1] + assert events[-1]["result"] == "Hi" + + def test_cache_split_and_judge_metadata(self) -> None: + reader = RunReader() + finished = reader.finish( + { + "response": "", + "_had_usage": True, + "usage": { + "input": 5, + "output": 1, + "total": 6, + "input_details": { + "uncached": 3, + "cache_read": 1, + "cache_creation": 1, + }, + }, + }, + { + "quality": JudgeResult( + usage=UsageDict( + input=1, + output=1, + total=2, + input_details=InputTokenDetails( + uncached=1, cache_read=0, cache_creation=0 + ), + ), + response="fine", + score=1.0, + ) + }, + provider="", + model="", + ) + usage = finished[-1]["usage"][0] + assert usage["cachedInputTokens"] == 1 + assert usage["cacheWriteInputTokens"] == 1 + assert "provider" not in usage + assert "result" not in finished[-1] + assert ( + finished[-1]["metadata"]["launchdarkly"]["judgeResults"]["quality"]["score"] + == 1.0 + ) + + +class TestStableProjector: + def test_text_content_is_a_chunk_and_a_third_type_is_not_done(self) -> None: + projector = StableProjector() + assert projector.feed({"type": "widget"}) == [] + assert projector.done is None + assert projector.feed( + {"type": TEXT_MESSAGE_CONTENT, "messageId": "m", "delta": "Hi"} + ) == [{"type": "chunk", "text": "Hi"}] + assert projector.feed({"type": "chunk", "text": "nope"}) == [] + assert ( + projector.feed( + {"type": TOOL_CALL_START, "toolCallId": "a", "toolCallName": "search"} + ) + == [] + ) + done = {"type": "done", "response": "Hi", "usage": {"input": 1}} + assert projector.feed(done) == [] + assert projector.done is done + + +class TestConfigAndReader: + async def test_config_stream_stays_chunk_then_done( + self, mock_ld_client: MagicMock + ) -> None: + handler = _handler( + [ + {"type": TEXT_MESSAGE_START, "messageId": "m1"}, + {"type": TEXT_MESSAGE_CONTENT, "messageId": "m1", "delta": "Hi"}, + { + "type": TOOL_CALL_START, + "toolCallId": "a", + "toolCallName": "search", + "parentMessageId": "m1", + }, + {"type": TOOL_CALL_ARGS, "toolCallId": "a", "delta": '{"q":"x"}'}, + {"type": TOOL_CALL_END, "toolCallId": "a"}, + {"type": TEXT_MESSAGE_END, "messageId": "m1"}, + { + "type": TOOL_CALL_RESULT, + "messageId": "tm", + "toolCallId": "a", + "content": "secret", + }, + {"type": "widget", "name": "nope"}, + { + "type": "done", + "output": "Hi", + "usage": {"input_tokens": 3, "output_tokens": 1}, + }, + ] + ) + events = [ + event + async for event in config(key="flag", handler=handler).stream("q", CONTEXT) + ] + assert events == [ + {"type": "chunk", "text": "Hi"}, + { + "type": "done", + "response": "Hi", + "usage": {"input": 3, "output": 1, "total": 4}, + "judge_results": None, + }, + ] + + async def test_reader_brackets_the_subset(self, mock_ld_client: MagicMock) -> None: + handler = _handler( + [ + {"type": "chunk", "text": "Hi"}, + { + "type": "done", + "output": "Hi", + "usage": {"input_tokens": 3, "output_tokens": 1}, + }, + ] + ) + events = [ + event + async for event in stream(config(key="flag", handler=handler), "q", CONTEXT) + ] + assert _types(events) == [ + RUN_STARTED, + TEXT_MESSAGE_START, + TEXT_MESSAGE_CONTENT, + TEXT_MESSAGE_END, + RUN_FINISHED, + ] + assert events[-1]["result"] == "Hi" + assert events[-1]["usage"][0]["provider"] == "TestProvider" + assert events[-1]["usage"][0]["model"] == "gpt-4" + assert events[-1]["usage"][0]["inputTokens"] == 3 + + async def test_a_raised_turn_has_no_run_finished( + self, mock_ld_client: MagicMock + ) -> None: + async def fn(cfg, user_input, tool_handlers, variables, history=None) -> dict: # type: ignore[override] + return {"output": "", "usage": {}} + + async def stream_fn( + cfg, user_input, tool_handlers, variables, history=None + ) -> AsyncGenerator[dict[str, Any], None]: + yield {"type": "chunk", "text": "Hi"} + raise RuntimeError("stream-fail") + yield # pragma: no cover + + handler = ProviderHandler( + fn=fn, + provides_for=("TestProvider", "messages"), + stream_fn=stream_fn, + ) + seen: list[str] = [] + with pytest.raises(RuntimeError, match="stream-fail"): + async for event in stream( + config(key="flag", handler=handler), "q", CONTEXT + ): + seen.append(event["type"]) + assert seen == [RUN_STARTED, TEXT_MESSAGE_START, TEXT_MESSAGE_CONTENT] + assert RUN_FINISHED not in seen + + async def test_judge_results_land_on_run_finished( + self, mock_ld_client: MagicMock + ) -> None: + handler = _handler( + [ + {"type": "chunk", "text": "Hi"}, + { + "type": "done", + "output": "Hi", + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + ] + ) + judged = { + "quality": JudgeResult( + usage=UsageDict(input=1, output=1, total=2), + response="fine", + score=0.5, + ) + } + with patch( + "launchdarkly_ai_server.judges.run_judges", + new=AsyncMock(return_value=judged), + ): + events = [ + event + async for event in stream( + config(key="flag", handler=handler), "q", CONTEXT + ) + ] + assert ( + events[-1]["metadata"]["launchdarkly"]["judgeResults"]["quality"]["score"] + == 0.5 + ) + + async def test_closing_early_does_not_finish_the_run( + self, mock_ld_client: MagicMock + ) -> None: + handler = _handler( + [ + {"type": "chunk", "text": "Hi"}, + {"type": "chunk", "text": " there"}, + { + "type": "done", + "output": "Hi there", + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + ] + ) + gen = stream(config(key="flag", handler=handler), "q", CONTEXT) + first = await gen.__anext__() + second = await gen.__anext__() + await gen.aclose() + assert first["type"] == RUN_STARTED + assert second["type"] == TEXT_MESSAGE_START From b7a9f3acde08367167f7f839500c799f3e06dac5 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 13 Oct 2026 14:21:05 +0000 Subject: [PATCH 2/8] feat(agui): close an unfinished tool call as a failure A handler that raises mid-call can still report the lifecycle without the tool result. Co-authored-by: alexis --- packages/client/src/launchdarkly_ai_server/agui.py | 8 ++++++++ packages/client/tests/test_agui.py | 8 ++++++++ 2 files changed, 16 insertions(+) diff --git a/packages/client/src/launchdarkly_ai_server/agui.py b/packages/client/src/launchdarkly_ai_server/agui.py index bdf7ec17..cb750553 100644 --- a/packages/client/src/launchdarkly_ai_server/agui.py +++ b/packages/client/src/launchdarkly_ai_server/agui.py @@ -210,6 +210,14 @@ def tool_result( events.append(event) return events + def fail_open(self) -> list[dict[str, Any]]: + """Close every call that started and did not finish, as a masked failure.""" + events: list[dict[str, Any]] = [] + for call_id in list(self._started): + if call_id not in self._resulted: + events.extend(self.tool_result(call_id, failed=True)) + return events + class StableProjector: """Projects a handler stream onto ``chunk`` then one internal ``done``. diff --git a/packages/client/tests/test_agui.py b/packages/client/tests/test_agui.py index 28ed548b..b104f266 100644 --- a/packages/client/tests/test_agui.py +++ b/packages/client/tests/test_agui.py @@ -141,6 +141,14 @@ def test_unserializable_result_is_a_failed_result(self) -> None: assert events[-1]["content"] == "An error occurred." assert events[-1]["metadata"]["launchdarkly"]["toolFailed"] is True + def test_fail_open_masks_a_call_that_never_finished(self) -> None: + emitter = TurnEmitter(capture_content=False) + emitter.tool_start("a", "search") + events = emitter.fail_open() + assert [event["type"] for event in events] == [TOOL_CALL_END, TOOL_CALL_RESULT] + assert "content" not in events[-1] + assert events[-1]["metadata"]["launchdarkly"]["toolFailed"] is True + def test_a_second_start_for_the_same_id_is_dropped(self) -> None: emitter = TurnEmitter(capture_content=True) assert emitter.tool_start("a", "search") From b2f253f27b922a3ec929ca391baa5b85896d81cc Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 13 Oct 2026 14:30:37 +0000 Subject: [PATCH 3/8] feat(claude): emit AG-UI text and tool events Messages and agents map Claude stream events and the tool loop onto the shared subset. Argument text and results follow capture_content. Co-authored-by: alexis --- .../launchdarkly_ai_claude_agents/handler.py | 100 ++++++++++++++++-- packages/claude-agents/tests/test_handler.py | 33 ++++-- .../handler.py | 64 +++++++++-- .../claude-messages/tests/test_handler.py | 57 +++++++++- 4 files changed, 223 insertions(+), 31 deletions(-) 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 eb7665e1..7399f8bd 100644 --- a/packages/claude-agents/src/launchdarkly_ai_claude_agents/handler.py +++ b/packages/claude-agents/src/launchdarkly_ai_claude_agents/handler.py @@ -48,6 +48,7 @@ set_output_content_attributes, set_tool_call_content_attributes, ) +from launchdarkly_ai_server.agui import TurnEmitter from ._version import PACKAGE_NAME, __version__ from .spans import ( @@ -197,6 +198,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 ]: @@ -216,6 +220,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 @@ -257,6 +266,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") ) @@ -710,6 +724,49 @@ 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, 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, @@ -751,6 +808,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 @@ -771,7 +832,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( @@ -787,20 +852,23 @@ 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, ResultMessage): # Before the root, so the children it parents are already closed. inference.finish() @@ -829,6 +897,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 @@ -852,12 +926,22 @@ 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: cancelled = True raise except Exception as exc: + while pending: + yield pending.pop(0) + for rich in emitter.fail_open(): + yield rich if tool_telemetry is not None: tool_telemetry(exc) if not root_usage_written and inference.run_usage["reported"]: diff --git a/packages/claude-agents/tests/test_handler.py b/packages/claude-agents/tests/test_handler.py index 83aef4de..c498e350 100644 --- a/packages/claude-agents/tests/test_handler.py +++ b/packages/claude-agents/tests/test_handler.py @@ -999,8 +999,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( @@ -1050,7 +1055,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 @@ -1095,9 +1100,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 @@ -1111,8 +1115,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 @@ -1204,7 +1213,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 @@ -1928,7 +1937,11 @@ 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"] in { + "TEXT_MESSAGE_START", + "TEXT_MESSAGE_CONTENT", + "TOOL_CALL_START", + } await gen.aclose() tools = [s for s in spans() if s.name.startswith("execute_tool ")] 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 9d01c14f..e9a92fd6 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 ._version import PACKAGE_NAME, __version__ from .spans import ( @@ -477,6 +478,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, @@ -512,6 +544,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 @@ -567,15 +601,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() @@ -614,6 +643,8 @@ async def _stream_gen( raise if final_msg.stop_reason != "tool_use": + for rich in emitter.end_text(): + yield rich break if steps >= _MAX_STEPS: @@ -623,6 +654,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": @@ -652,10 +694,14 @@ async def _stream_gen( except Exception as exc: fail_span(tool_span, exc, ended) open_tool_span = None + for rich in emitter.tool_result(str(block.id), failed=True): + yield rich raise # 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", diff --git a/packages/claude-messages/tests/test_handler.py b/packages/claude-messages/tests/test_handler.py index 0a0353d9..c90b77b7 100644 --- a/packages/claude-messages/tests/test_handler.py +++ b/packages/claude-messages/tests/test_handler.py @@ -1303,10 +1303,59 @@ 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 "content" not in result + + 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_all_chunks_before_done(self, mock_anthropic: MagicMock) -> None: import launchdarkly_ai_claude_messages.spans as spans_mod From cae50a1e49e5b52c8fcb2bad4eeec7a589b954e2 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 13 Oct 2026 15:59:43 +0000 Subject: [PATCH 4/8] fix(claude): close tool calls the messages stream does not finish A streamed tool_use that never becomes a tool turn, and any sibling still open when a call raises, now get a failed TOOL_CALL_RESULT. Co-authored-by: alexis --- .../handler.py | 4 ++ .../claude-messages/tests/test_handler.py | 37 +++++++++++++++++++ 2 files changed, 41 insertions(+) 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 e9a92fd6..5a2c11cf 100644 --- a/packages/claude-messages/src/launchdarkly_ai_claude_messages/handler.py +++ b/packages/claude-messages/src/launchdarkly_ai_claude_messages/handler.py @@ -645,6 +645,8 @@ async def _stream_gen( 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: @@ -738,6 +740,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 c90b77b7..d3909f49 100644 --- a/packages/claude-messages/tests/test_handler.py +++ b/packages/claude-messages/tests/test_handler.py @@ -1357,6 +1357,43 @@ async def _events(capture: bool) -> list[dict[str, Any]]: ) 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 + started.content_block = MagicMock(type="tool_use", id="c1", name="search") + _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")] + assert any( + event["type"] == "TOOL_CALL_RESULT" and event["toolCallId"] == "c1" + for event in events + ) + async def test_all_chunks_before_done(self, mock_anthropic: MagicMock) -> None: import launchdarkly_ai_claude_messages.spans as spans_mod from launchdarkly_ai_claude_messages import create_claude_messages_handler From d8f1166c886aba21e7e6b5bff95e056660beecd5 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 13 Oct 2026 16:37:30 +0000 Subject: [PATCH 5/8] fix(claude): close every messages tool call that the turn abandons A stream error, a handler that raises, and a streamed call the final message never runs now emit TOOL_CALL_RESULT before the turn ends. Co-authored-by: alexis --- .../handler.py | 6 ++ .../claude-messages/tests/test_handler.py | 72 +++++++++++++++++++ 2 files changed, 78 insertions(+) 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 5a2c11cf..50a5f0c7 100644 --- a/packages/claude-messages/src/launchdarkly_ai_claude_messages/handler.py +++ b/packages/claude-messages/src/launchdarkly_ai_claude_messages/handler.py @@ -640,6 +640,8 @@ async def _stream_gen( except Exception as exc: fail_span(model_span, exc, ended) open_model_span = None + for rich in emitter.fail_open(): + yield rich raise if final_msg.stop_reason != "tool_use": @@ -698,6 +700,8 @@ async def _stream_gen( open_tool_span = None for rich in emitter.tool_result(str(block.id), failed=True): yield rich + for rich in emitter.fail_open(): + yield rich raise # 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 @@ -711,6 +715,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( diff --git a/packages/claude-messages/tests/test_handler.py b/packages/claude-messages/tests/test_handler.py index d3909f49..42a4c789 100644 --- a/packages/claude-messages/tests/test_handler.py +++ b/packages/claude-messages/tests/test_handler.py @@ -1394,6 +1394,78 @@ async def _mgr() -> AsyncGenerator[Any, None]: 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 + started.content_block = MagicMock(type="tool_use", id="c1", name="search") + + 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 from launchdarkly_ai_claude_messages import create_claude_messages_handler From 580a8229c9b0971358b3be9bc37aa840fadd9069 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 13 Oct 2026 18:35:19 +0000 Subject: [PATCH 6/8] fix(claude): fail the run before closing tool calls Closing tool events inside an error handler let a consumer stop the stream before the root span was marked failed, so the run looked abandoned. Streamed MCP tools now use the config name. Co-authored-by: alexis --- .../launchdarkly_ai_claude_agents/handler.py | 12 ++++++---- packages/claude-agents/tests/test_handler.py | 23 +++++++++++++++++++ .../handler.py | 6 ----- .../claude-messages/tests/test_handler.py | 15 ++++++++++-- 4 files changed, 43 insertions(+), 13 deletions(-) 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 7399f8bd..fb730e81 100644 --- a/packages/claude-agents/src/launchdarkly_ai_claude_agents/handler.py +++ b/packages/claude-agents/src/launchdarkly_ai_claude_agents/handler.py @@ -740,7 +740,9 @@ def _map_claude_stream_event( call_id = str(block.get("id") or "") if isinstance(index, int): block_ids[index] = call_id - return emitter.tool_start(call_id, str(block.get("name") or "")) + 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): @@ -938,16 +940,16 @@ async def _stream_gen( cancelled = True raise except Exception as exc: - while pending: - yield pending.pop(0) - for rich in emitter.fail_open(): - yield rich if tool_telemetry is not None: tool_telemetry(exc) if not root_usage_written and inference.run_usage["reported"]: 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 c498e350..e6b229b9 100644 --- a/packages/claude-agents/tests/test_handler.py +++ b/packages/claude-agents/tests/test_handler.py @@ -2242,3 +2242,26 @@ async def _query(**kwargs: Any) -> AsyncIterator[Any]: tool = named("execute_tool ")[0] assert tool.end_time is not None + + +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 50a5f0c7..6dfc77aa 100644 --- a/packages/claude-messages/src/launchdarkly_ai_claude_messages/handler.py +++ b/packages/claude-messages/src/launchdarkly_ai_claude_messages/handler.py @@ -640,8 +640,6 @@ async def _stream_gen( except Exception as exc: fail_span(model_span, exc, ended) open_model_span = None - for rich in emitter.fail_open(): - yield rich raise if final_msg.stop_reason != "tool_use": @@ -698,10 +696,6 @@ async def _stream_gen( except Exception as exc: fail_span(tool_span, exc, ended) open_tool_span = None - for rich in emitter.tool_result(str(block.id), failed=True): - yield rich - for rich in emitter.fail_open(): - yield rich raise # 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 diff --git a/packages/claude-messages/tests/test_handler.py b/packages/claude-messages/tests/test_handler.py index 42a4c789..95d12cea 100644 --- a/packages/claude-messages/tests/test_handler.py +++ b/packages/claude-messages/tests/test_handler.py @@ -1366,7 +1366,11 @@ async def test_a_streamed_tool_without_tool_use_is_closed( started = MagicMock() started.type = "content_block_start" started.index = 0 - started.content_block = MagicMock(type="tool_use", id="c1", name="search") + 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" @@ -1389,6 +1393,9 @@ async def _mgr() -> AsyncGenerator[Any, None]: 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 @@ -1435,7 +1442,11 @@ async def test_a_stream_error_closes_an_opened_tool( started = MagicMock() started.type = "content_block_start" started.index = 0 - started.content_block = MagicMock(type="tool_use", id="c1", name="search") + block = MagicMock() + block.type = "tool_use" + block.id = "c1" + block.name = "search" + started.content_block = block class _Stream: def __aiter__(self) -> AsyncIterator[Any]: From c8d198c3f6253a8c4eab45d32d459fb1614c502b Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 13 Oct 2026 18:51:47 +0000 Subject: [PATCH 7/8] fix(agui): keep each run valid for an AG-UI client Message ids are unique across runs, tool results always carry content, and a run cannot finish while a tool call is open. Argument text is completed only when it extends what was streamed. Disconnects close the handler, and callers can echo threadId and runId. Co-authored-by: alexis --- .../client/src/launchdarkly_ai_server/agui.py | 186 ++++++++++++------ .../src/launchdarkly_ai_server/client.py | 58 +++--- .../src/launchdarkly_ai_server/tracking.py | 51 ++--- packages/client/tests/test_agui.py | 99 ++++++++-- 4 files changed, 265 insertions(+), 129 deletions(-) diff --git a/packages/client/src/launchdarkly_ai_server/agui.py b/packages/client/src/launchdarkly_ai_server/agui.py index cb750553..6b302e65 100644 --- a/packages/client/src/launchdarkly_ai_server/agui.py +++ b/packages/client/src/launchdarkly_ai_server/agui.py @@ -5,9 +5,20 @@ ``RUN_STARTED`` and ``RUN_FINISHED``. ``config().stream()`` still yields ``chunk`` then ``done``. The SDK runs the tool loop. This stream only reports it. -Argument deltas and tool results are content. :class:`TurnEmitter` omits them -when the handler was built with ``capture_content`` off, and still emits the -tool lifecycle. +Contract for one run: + +* ``RUN_STARTED`` is first and ``RUN_FINISHED`` is last. A failure or + cancellation raises and does not yield ``RUN_FINISHED`` or ``RUN_ERROR``. +* Message and tool-result ids are unique across runs. A chat UI that keeps + messages by id will not append a later answer onto an earlier one. +* ``TOOL_CALL_RESULT.content`` is always a string. Argument deltas and result + text are omitted (sent as ``""``) when ``capture_content`` is off. A failed + call still sets ``metadata.launchdarkly.toolFailed``. +* ``RUN_FINISHED`` is not sent while a tool call is open. An unfinished call + is closed first. +* Tool-result text is the string the model was given, not a second encoding. +* Optional ``thread_id`` and ``run_id`` are echoed on ``RUN_STARTED`` and + ``RUN_FINISHED``. When omitted, this reader mints them. """ from __future__ import annotations @@ -16,6 +27,7 @@ import math import uuid from collections.abc import AsyncGenerator +from contextlib import aclosing from typing import Any from .conversation import bind_conversation_id @@ -101,19 +113,22 @@ class TurnEmitter: def __init__(self, *, capture_content: bool) -> None: self.capture_content = capture_content self._open: str | None = None - self._seq = 0 self._started: set[str] = set() self._args: set[str] = set() + self._arg_text: dict[str, str] = {} self._ended: set[str] = set() self._resulted: set[str] = set() def _mint(self, prefix: str) -> str: - self._seq += 1 - return f"{prefix}-{self._seq}" + return f"{prefix}-{uuid.uuid4()}" + + def has_argument_deltas(self, tool_call_id: str) -> bool: + """True once a streamed argument slice was recorded for this call.""" + return tool_call_id in self._args def text(self, delta: str) -> list[dict[str, Any]]: """Append assistant text. The first slice opens the message.""" - if not delta: + if not isinstance(delta, str) or not delta: return [] events: list[dict[str, Any]] = [] if self._open is None: @@ -159,25 +174,31 @@ def tool_args(self, tool_call_id: str, delta: str) -> list[dict[str, Any]]: ): return [] self._args.add(tool_call_id) + self._arg_text[tool_call_id] = self._arg_text.get(tool_call_id, "") + delta return [{"type": TOOL_CALL_ARGS, "toolCallId": tool_call_id, "delta": delta}] def tool_end( self, tool_call_id: str, *, arguments: Any = None ) -> list[dict[str, Any]]: - """Close the argument text. One delta is emitted when none arrived.""" + """Close the argument text. + + When nothing was streamed, one delta carries the complete arguments. + When the complete text starts with what was already streamed, the + missing suffix is appended. Otherwise the streamed text wins: a + differently encoded complete value is not concatenated onto it. + """ if tool_call_id not in self._started or tool_call_id in self._ended: return [] events: list[dict[str, Any]] = [] - if ( - self.capture_content - and tool_call_id not in self._args - and arguments is not None - ): + if self.capture_content and arguments is not None: text = argument_text(arguments) - if text: + have = self._arg_text.get(tool_call_id, "") + extra = text[len(have) :] if text and text.startswith(have) else "" + if extra: self._args.add(tool_call_id) + self._arg_text[tool_call_id] = have + extra events.append( - {"type": TOOL_CALL_ARGS, "toolCallId": tool_call_id, "delta": text} + {"type": TOOL_CALL_ARGS, "toolCallId": tool_call_id, "delta": extra} ) self._ended.add(tool_call_id) events.append({"type": TOOL_CALL_END, "toolCallId": tool_call_id}) @@ -198,8 +219,7 @@ def tool_result( } if failed: event["metadata"] = {"launchdarkly": {"toolFailed": True}} - if self.capture_content: - event["content"] = _FAILURE_CONTENT + event["content"] = _FAILURE_CONTENT if self.capture_content else "" elif self.capture_content: content, bad = safe_tool_content(value) if bad: @@ -207,6 +227,8 @@ def tool_result( event["metadata"] = {"launchdarkly": {"toolFailed": True}} else: event["content"] = content + else: + event["content"] = "" events.append(event) return events @@ -246,7 +268,7 @@ def feed(self, event: dict[str, Any]) -> list[dict[str, Any]]: return [{"type": "chunk", "text": text}] if typ == TEXT_MESSAGE_CONTENT: delta = event.get("delta") or "" - if not delta: + if not isinstance(delta, str) or not delta: return [] return [{"type": "chunk", "text": delta}] return [] @@ -304,16 +326,17 @@ class RunReader: here. The handler's ``done`` is consumed, not forwarded. """ - def __init__(self) -> None: - self.thread_id = str(uuid.uuid4()) - self.run_id = str(uuid.uuid4()) + def __init__( + self, *, thread_id: str | None = None, run_id: str | None = None + ) -> None: + self.thread_id = thread_id or str(uuid.uuid4()) + self.run_id = run_id or str(uuid.uuid4()) self._mode: str | None = None self._open: str | None = None self._deltas: list[str] = [] self._started: set[str] = set() self._ended: set[str] = set() self._resulted: set[str] = set() - self._seq = 0 def start(self) -> dict[str, Any]: return { @@ -324,8 +347,7 @@ def start(self) -> dict[str, Any]: } def _mint(self, prefix: str) -> str: - self._seq += 1 - return f"{prefix}-{self._seq}" + return f"{prefix}-{uuid.uuid4()}" def feed(self, event: dict[str, Any]) -> list[dict[str, Any]]: typ = event.get("type") @@ -359,6 +381,10 @@ def finish( if self._open is not None: events.append({"type": TEXT_MESSAGE_END, "messageId": self._open}) self._open = None + for call_id in list(self._started): + if call_id not in self._ended: + self._ended.add(call_id) + events.append({"type": TOOL_CALL_END, "toolCallId": call_id}) response = done.get("response") if not isinstance(response, str): response = "".join(self._deltas) @@ -497,6 +523,10 @@ def _tool_result(self, event: dict[str, Any]) -> list[dict[str, Any]]: or call_id in self._resulted ): return [] + events: list[dict[str, Any]] = [] + if call_id not in self._ended: + self._ended.add(call_id) + events.append({"type": TOOL_CALL_END, "toolCallId": call_id}) self._resulted.add(call_id) message_id = event.get("messageId") if not isinstance(message_id, str) or not message_id: @@ -505,10 +535,10 @@ def _tool_result(self, event: dict[str, Any]) -> list[dict[str, Any]]: "type": TOOL_CALL_RESULT, "messageId": message_id, "toolCallId": call_id, + "content": event.get("content") + if isinstance(event.get("content"), str) + else "", } - content = event.get("content") - if isinstance(content, str): - result["content"] = content metadata = event.get("metadata") if isinstance(metadata, dict): launchdarkly = metadata.get("launchdarkly") @@ -517,7 +547,8 @@ def _tool_result(self, event: dict[str, Any]) -> list[dict[str, Any]]: and launchdarkly.get("toolFailed") is True ): result["metadata"] = {"launchdarkly": {"toolFailed": True}} - return [result] + events.append(result) + return events def stream( @@ -527,16 +558,59 @@ def stream( *, variables: dict[str, Any] | None = None, history: list[dict[str, Any]] | None = None, + thread_id: str | None = None, + run_id: str | None = None, ) -> AsyncGenerator[dict[str, Any], None]: """Stream one config turn as AG-UI events. - Same arguments as ``ConfigInstance.stream``. ``RUN_STARTED`` is first and - ``RUN_FINISHED`` is last. A failure, including cancellation, raises and does - not yield ``RUN_FINISHED``. + Same arguments as ``ConfigInstance.stream``, plus optional ``thread_id`` + and ``run_id`` echoed from the caller's AG-UI run. ``RUN_STARTED`` is first + and ``RUN_FINISHED`` is last. A failure, including cancellation, raises and + does not yield ``RUN_FINISHED``. """ report_usage("client.agui.stream") return bind_conversation_id( - _stream_events(instance, user_input, context, variables, history) + _stream_events( + instance, + user_input, + context, + variables, + history, + thread_id=thread_id, + run_id=run_id, + ) + ) + + +async def judge_completed_stream( + instance: Any, + sink: dict[str, Any], + done: dict[str, Any], + user_input: str | None, + context: LDContext, +) -> dict[str, JudgeResult]: + """Run the turn's judges once the handler has finished. + + A non-string response is not sent to the judge. Both ``config().stream()`` + and :func:`stream` use this, so the two copies cannot drift. + """ + from .judges import run_judges + + if getattr(instance, "_skip_judges", False) or not sink: + return {} + response = done.get("response", "") + if not isinstance(response, str): + response = "" + return await run_judges( + config=sink["config"], + user_context=context, + handler=sink["handler"], + handlers=sink["handlers"], + user_input=user_input, + trajectory=done.get("trajectory", ""), + llm_response=response, + base_track_data=done.get("track_data", {}), + tool_handlers=sink["tools"], ) @@ -546,43 +620,33 @@ async def _stream_events( context: LDContext, variables: dict[str, Any] | None, history: list[dict[str, Any]] | None, + *, + thread_id: str | None = None, + run_id: str | None = None, ) -> AsyncGenerator[dict[str, Any], None]: - from .judges import run_judges - sink: dict[str, Any] = {} - reader = RunReader() + reader = RunReader(thread_id=thread_id, run_id=run_id) started = False done: dict[str, Any] | None = None - async for event in instance._execute_stream( - user_input, context, variables, history, sink - ): - if not started: - yield reader.start() - started = True - if event.get("type") == "done": - done = event - continue - for item in reader.feed(event): - yield item + async with aclosing( + instance._execute_stream(user_input, context, variables, history, sink) + ) as source: + async for event in source: + if not started: + yield reader.start() + started = True + if event.get("type") == "done": + done = event + continue + for item in reader.feed(event): + yield item if not started: yield reader.start() if done is None: done = {"response": "", "usage": {}, "_had_usage": False} - judge_results: dict[str, JudgeResult] = {} - if not instance._skip_judges and sink: - judge_results = await run_judges( - config=sink["config"], - user_context=context, - handler=sink["handler"], - handlers=sink["handlers"], - user_input=user_input, - trajectory=done.get("trajectory", ""), - llm_response=done.get("response", "") - if isinstance(done.get("response"), str) - else "", - base_track_data=done.get("track_data", {}), - tool_handlers=sink["tools"], - ) + judge_results = await judge_completed_stream( + instance, sink, done, user_input, context + ) for item in reader.finish( done, judge_results, diff --git a/packages/client/src/launchdarkly_ai_server/client.py b/packages/client/src/launchdarkly_ai_server/client.py index 9ea95962..da73be45 100644 --- a/packages/client/src/launchdarkly_ai_server/client.py +++ b/packages/client/src/launchdarkly_ai_server/client.py @@ -2,9 +2,10 @@ import json from collections.abc import AsyncGenerator, Callable +from contextlib import aclosing from typing import Any -from .agui import StableProjector +from .agui import StableProjector, judge_completed_stream from .conversation import bind_conversation_id from .judges import _build_judge_tasks, run_judges from .lifecycle import extract_variation @@ -231,18 +232,21 @@ async def _execute_stream( sink["provider"] = provider sink["model"] = model - async for event in execute_and_stream( - config_key=self._key, - config=config, - meta=meta, - user_context=context, - handler=handler, - user_input=user_input, - tool_handlers=resolved_tools, - variables=variables, - history=history, - ): - yield event + async with aclosing( + execute_and_stream( + config_key=self._key, + config=config, + meta=meta, + user_context=context, + handler=handler, + user_input=user_input, + tool_handlers=resolved_tools, + variables=variables, + history=history, + ) + ) as source: + async for event in source: + yield event async def _stream_events( self, @@ -253,29 +257,17 @@ async def _stream_events( ) -> AsyncGenerator[StreamEvent, None]: sink: dict[str, Any] = {} projector = StableProjector() - async for event in self._execute_stream( - user_input, context, variables, history, sink - ): - for stable in projector.feed(event): - yield stable + async with aclosing( + self._execute_stream(user_input, context, variables, history, sink) + ) as source: + async for event in source: + for stable in projector.feed(event): + yield stable done_event = projector.done if done_event: - track_data = done_event.get("track_data", {}) - judge_results = ( - {} - if self._skip_judges - else await run_judges( - config=sink["config"], - user_context=context, - handler=sink["handler"], - handlers=sink["handlers"], - user_input=user_input, - trajectory=done_event.get("trajectory", ""), - llm_response=done_event.get("response", ""), - base_track_data=track_data, - tool_handlers=sink["tools"], - ) + judge_results = await judge_completed_stream( + self, sink, done_event, user_input, context ) yield { "type": "done", diff --git a/packages/client/src/launchdarkly_ai_server/tracking.py b/packages/client/src/launchdarkly_ai_server/tracking.py index a2d3454d..7dca2d57 100644 --- a/packages/client/src/launchdarkly_ai_server/tracking.py +++ b/packages/client/src/launchdarkly_ai_server/tracking.py @@ -5,6 +5,7 @@ import time import uuid from collections.abc import AsyncGenerator, Callable +from contextlib import aclosing from typing import Any from .agui import HANDLER_RICH_TYPES, TEXT_MESSAGE_CONTENT, TEXT_TYPES @@ -307,29 +308,33 @@ async def execute_and_stream( stream_gen = await handler.stream( config, user_input, tracked_tool_handlers, merged_variables, history ) - async for event in stream_gen: - typ = event.get("type") - if typ in TEXT_TYPES: - saw_text = True - if typ == "chunk": - if saw_text: - continue - full_text += event.get("text", "") - yield {"type": "chunk", "text": event["text"]} - elif typ == TEXT_MESSAGE_CONTENT: - delta = event.get("delta") or "" - if isinstance(delta, str): - full_text += delta - yield event - elif typ in HANDLER_RICH_TYPES: - yield event - elif typ == "done": - if event.get("output") is not None: - full_text = event["output"] - usage_value = event.get("usage") - if isinstance(usage_value, dict) and usage_value: - raw_usage = usage_value - had_usage = True + async with aclosing(stream_gen) as agen: + async for event in agen: + typ = event.get("type") + if typ in TEXT_TYPES: + saw_text = True + if typ == "chunk": + if saw_text: + continue + text = event.get("text", "") + if not isinstance(text, str): + continue + full_text += text + yield {"type": "chunk", "text": text} + elif typ == TEXT_MESSAGE_CONTENT: + delta = event.get("delta") or "" + if isinstance(delta, str): + full_text += delta + yield event + elif typ in HANDLER_RICH_TYPES: + yield event + elif typ == "done": + if event.get("output") is not None: + full_text = event["output"] + usage_value = event.get("usage") + if isinstance(usage_value, dict) and usage_value: + raw_usage = usage_value + had_usage = True else: result = await handler( config, user_input, tracked_tool_handlers, merged_variables, history diff --git a/packages/client/tests/test_agui.py b/packages/client/tests/test_agui.py index b104f266..915e6944 100644 --- a/packages/client/tests/test_agui.py +++ b/packages/client/tests/test_agui.py @@ -119,7 +119,7 @@ def test_capture_off_keeps_the_lifecycle_without_payloads(self) -> None: assert _types(events) == [TOOL_CALL_START, TOOL_CALL_END, TOOL_CALL_RESULT] assert events[0]["toolCallName"] == "search" assert "parentMessageId" not in events[0] - assert "content" not in events[2] + assert events[2]["content"] == "" assert "delta" not in str(events) def test_failure_is_masked_and_a_string_is_not_reencoded(self) -> None: @@ -146,7 +146,7 @@ def test_fail_open_masks_a_call_that_never_finished(self) -> None: emitter.tool_start("a", "search") events = emitter.fail_open() assert [event["type"] for event in events] == [TOOL_CALL_END, TOOL_CALL_RESULT] - assert "content" not in events[-1] + assert events[-1]["content"] == "" assert events[-1]["metadata"]["launchdarkly"]["toolFailed"] is True def test_a_second_start_for_the_same_id_is_dropped(self) -> None: @@ -154,6 +154,25 @@ def test_a_second_start_for_the_same_id_is_dropped(self) -> None: assert emitter.tool_start("a", "search") assert emitter.tool_start("a", "search") == [] + def test_message_ids_differ_across_runs(self) -> None: + first = TurnEmitter(capture_content=False).text("a")[0]["messageId"] + second = TurnEmitter(capture_content=False).text("b")[0]["messageId"] + assert first != second + + def test_tool_end_extends_a_matching_prefix_only(self) -> None: + emitter = TurnEmitter(capture_content=True) + events = emitter.tool_start("a", "search") + events.extend(emitter.tool_args("a", '{"q":')) + events.extend(emitter.tool_end("a", arguments={"q": "x"})) + deltas = [event["delta"] for event in events if event["type"] == TOOL_CALL_ARGS] + assert "".join(deltas) == '{"q":"x"}' + + spaced = TurnEmitter(capture_content=True) + spaced.tool_start("a", "search") + spaced.tool_args("a", '{"q": "') + closed = spaced.tool_end("a", arguments={"q": "xy"}) + assert not any(event["type"] == TOOL_CALL_ARGS for event in closed) + class TestRunReader: def test_legacy_chunks_become_one_text_message(self) -> None: @@ -300,6 +319,40 @@ def test_cache_split_and_judge_metadata(self) -> None: == 1.0 ) + def test_finish_closes_an_open_tool_before_the_run(self) -> None: + reader = RunReader(thread_id="thread-1", run_id="run-1") + events = [reader.start()] + events.extend( + reader.feed( + { + "type": TOOL_CALL_START, + "toolCallId": "a", + "toolCallName": "search", + } + ) + ) + events.extend(reader.finish({"response": ""}, None, provider="", model="")) + assert [event["type"] for event in events] == [ + RUN_STARTED, + TOOL_CALL_START, + TOOL_CALL_END, + RUN_FINISHED, + ] + assert events[0]["threadId"] == "thread-1" + assert events[-1]["runId"] == "run-1" + + def test_a_result_without_an_end_closes_the_call_first(self) -> None: + reader = RunReader() + reader.feed( + {"type": TOOL_CALL_START, "toolCallId": "a", "toolCallName": "search"} + ) + events = reader.feed({"type": TOOL_CALL_RESULT, "toolCallId": "a"}) + assert [event["type"] for event in events] == [ + TOOL_CALL_END, + TOOL_CALL_RESULT, + ] + assert events[-1]["content"] == "" + class TestStableProjector: def test_text_content_is_a_chunk_and_a_third_type_is_not_done(self) -> None: @@ -458,20 +511,42 @@ async def test_judge_results_land_on_run_finished( async def test_closing_early_does_not_finish_the_run( self, mock_ld_client: MagicMock ) -> None: - handler = _handler( - [ - {"type": "chunk", "text": "Hi"}, - {"type": "chunk", "text": " there"}, - { - "type": "done", - "output": "Hi there", - "usage": {"input_tokens": 1, "output_tokens": 1}, - }, - ] + closed = False + + async def fn_unused( + cfg, user_input, tool_handlers, variables, history=None + ) -> dict: + return {"output": "", "usage": {}} + + async def stream_fn( + cfg, user_input, tool_handlers, variables, history=None + ) -> AsyncGenerator[dict[str, Any], None]: + nonlocal closed + try: + for event in [ + {"type": "chunk", "text": "Hi"}, + {"type": "chunk", "text": " there"}, + { + "type": "done", + "output": "Hi there", + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + ]: + yield event + finally: + closed = True + + handler = ProviderHandler( + fn=fn_unused, + provides_for=("TestProvider", "messages"), + stream_fn=stream_fn, + capture_content=False, ) + gen = stream(config(key="flag", handler=handler), "q", CONTEXT) first = await gen.__anext__() second = await gen.__anext__() await gen.aclose() assert first["type"] == RUN_STARTED assert second["type"] == TEXT_MESSAGE_START + assert closed is True From bbde74e38f40d2f6fb75ace2d7997f250f9ce7d6 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 13 Oct 2026 18:57:35 +0000 Subject: [PATCH 8/8] fix(claude): close built-in tool calls from their results A Read or other built-in with no configured hook was marked failed at the end of the run. The user message's tool result now closes it, and the stream test covers the full tool lifecycle. Co-authored-by: alexis --- .../launchdarkly_ai_claude_agents/handler.py | 16 ++++ packages/claude-agents/tests/test_handler.py | 73 +++++++++++++++++-- .../claude-messages/tests/test_handler.py | 2 +- 3 files changed, 85 insertions(+), 6 deletions(-) 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 fb730e81..e56c6851 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, ) @@ -871,6 +873,20 @@ async def _stream_gen( 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() diff --git a/packages/claude-agents/tests/test_handler.py b/packages/claude-agents/tests/test_handler.py index e6b229b9..9ff0ba31 100644 --- a/packages/claude-agents/tests/test_handler.py +++ b/packages/claude-agents/tests/test_handler.py @@ -1937,11 +1937,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"] in { - "TEXT_MESSAGE_START", - "TEXT_MESSAGE_CONTENT", - "TOOL_CALL_START", - } + assert (await gen.__anext__())["type"] == "TOOL_CALL_START" await gen.aclose() tools = [s for s in spans() if s.name.startswith("execute_tool ")] @@ -2244,6 +2240,73 @@ async def _query(**kwargs: Any) -> AsyncIterator[Any]: 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 diff --git a/packages/claude-messages/tests/test_handler.py b/packages/claude-messages/tests/test_handler.py index 95d12cea..96bf4b20 100644 --- a/packages/claude-messages/tests/test_handler.py +++ b/packages/claude-messages/tests/test_handler.py @@ -1347,7 +1347,7 @@ async def _events(capture: bool) -> list[dict[str, Any]]: 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 "content" not in result + assert result["content"] == "" shown = await _events(True) args = next(event for event in shown if event["type"] == "TOOL_CALL_ARGS")