From fb516505c432d2b7f6ec00403c64a2a0d29ed868 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 13 Oct 2026 14:18:55 +0000 Subject: [PATCH 01/15] 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 02/15] 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 6e013a97e65b9dee1a05b274632a304d04d662dc Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 13 Oct 2026 14:25:09 +0000 Subject: [PATCH 03/15] feat(openai): emit AG-UI text and tool events Messages and agents map Responses stream events and the tool loop onto the shared subset. Argument text and results follow capture_content. Co-authored-by: alexis --- .../launchdarkly_ai_openai_agents/handler.py | 128 ++++++++++++++++-- packages/openai-agents/tests/test_handler.py | 50 ++++++- .../handler.py | 113 +++++++++++++++- .../openai-messages/tests/test_handler.py | 53 +++++++- 4 files changed, 325 insertions(+), 19 deletions(-) diff --git a/packages/openai-agents/src/launchdarkly_ai_openai_agents/handler.py b/packages/openai-agents/src/launchdarkly_ai_openai_agents/handler.py index 355aaab9..f6824159 100644 --- a/packages/openai-agents/src/launchdarkly_ai_openai_agents/handler.py +++ b/packages/openai-agents/src/launchdarkly_ai_openai_agents/handler.py @@ -45,6 +45,7 @@ set_tool_call_content_attributes, text_message, ) +from launchdarkly_ai_server.agui import TurnEmitter from ._version import PACKAGE_NAME, __version__ from .spans import ( @@ -249,10 +250,14 @@ def __init__( self.run_usage = run_usage self.open_model_span: Any = None self.open_tool_spans: dict[str, Any] = {} + self.emitter = TurnEmitter(capture_content=capture_content) + self.pending: list[dict[str, Any]] = [] + self.item_to_call: dict[str, str] = {} async def on_llm_start( self, context: Any, agent: Any, system_prompt: str | None, input_items: Any ) -> None: + self.pending.extend(self.emitter.end_text()) span = start_model_span(self.config, self.parent) self.open_model_span = span if self.capture_content: @@ -319,6 +324,9 @@ async def on_tool_start(self, context: Any, agent: Any, tool: Any) -> None: # abandon_open_spans both walk this dict, so a span missing from it never ends and never # exports. self.open_tool_spans[str(call_id)] = span + self.pending.extend(self.emitter.tool_start(str(call_id), str(name))) + tool_args = getattr(context, "tool_arguments", None) + self.pending.extend(self.emitter.tool_end(str(call_id), arguments=tool_args)) if self.capture_content: # `tool_arguments` is the raw JSON args string the model produced. Parsed here rather # than passed through, per TELEMETRY-CONTRACT.md section 12: arguments hold the object @@ -335,7 +343,9 @@ async def on_tool_start(self, context: Any, agent: Any, tool: Any) -> None: async def on_tool_end( self, context: Any, agent: Any, tool: Any, result: Any ) -> None: - span = self.open_tool_spans.pop(self._tool_span_key(context, tool), None) + call_id = self._tool_span_key(context, tool) + self.pending.extend(self.emitter.tool_result(str(call_id), result)) + span = self.open_tool_spans.pop(call_id, None) if span is None: return # Once popped, ending it is this method's job alone: nothing else knows the span exists. A @@ -576,6 +586,95 @@ def _stream_impl( ) +_RESPONSES_ARG_DELTA = { + "response.function_call_arguments.delta": "delta", + "response.mcp_call_arguments.delta": "delta", + "response.custom_tool_call_input.delta": "delta", +} +_RESPONSES_ARG_DONE = { + "response.function_call_arguments.done": "arguments", + "response.mcp_call_arguments.done": "arguments", + "response.custom_tool_call_input.done": "input", +} +_RESPONSES_OPEN_TYPES = frozenset({"function_call", "mcp_call", "custom_tool_call"}) +_RESPONSES_COMPLETED = { + "response.file_search_call.completed": "file_search", + "response.code_interpreter_call.completed": "code_interpreter", +} + + +def _responses_call_id(event: Any, item_to_call: dict[str, str]) -> str: + direct = getattr(event, "call_id", None) + if isinstance(direct, str) and direct: + return direct + item_id = getattr(event, "item_id", None) + if isinstance(item_id, str) and item_id in item_to_call: + return item_to_call[item_id] + return item_id if isinstance(item_id, str) else "" + + +def _map_responses_event( + event: Any, emitter: TurnEmitter, item_to_call: dict[str, str] +) -> list[dict[str, Any]]: + """Map one Responses stream event onto the AG-UI subset. Unknown events are ignored.""" + kind = getattr(event, "type", None) + if kind == "response.output_text.delta": + delta = getattr(event, "delta", "") + if not isinstance(delta, str) or not delta: + return [] + return emitter.text(delta) + if kind == "response.output_item.added": + item = getattr(event, "item", None) + if getattr(item, "type", None) not in _RESPONSES_OPEN_TYPES: + return [] + call_id = getattr(item, "call_id", None) or getattr(item, "id", None) or "" + name = getattr(item, "name", None) or "" + item_id = getattr(item, "id", None) + if isinstance(item_id, str) and isinstance(call_id, str) and call_id: + item_to_call[item_id] = call_id + return emitter.tool_start(str(call_id), str(name)) + if kind in _RESPONSES_ARG_DELTA: + delta = getattr(event, "delta", "") + if not isinstance(delta, str): + return [] + return emitter.tool_args(_responses_call_id(event, item_to_call), delta) + if kind in _RESPONSES_ARG_DONE: + arguments = getattr(event, _RESPONSES_ARG_DONE[kind], None) + return emitter.tool_end( + _responses_call_id(event, item_to_call), arguments=arguments + ) + if kind in ("response.mcp_call.completed", "response.mcp_call.failed"): + call_id = _responses_call_id(event, item_to_call) + item = getattr(event, "item", None) + if not call_id and item is not None: + raw_id = getattr(item, "call_id", None) or getattr(item, "id", None) + call_id = raw_id if isinstance(raw_id, str) else "" + if kind == "response.mcp_call.failed": + return emitter.tool_result(call_id, failed=True) + output = getattr(event, "output", None) + if output is None and item is not None: + output = getattr(item, "output", None) + return emitter.tool_result(call_id, output) + if kind in _RESPONSES_COMPLETED: + item = getattr(event, "item", None) + raw_id = "" + if item is not None: + candidate = getattr(item, "call_id", None) or getattr(item, "id", None) + raw_id = candidate if isinstance(candidate, str) else "" + if not raw_id: + candidate = getattr(event, "item_id", None) + raw_id = candidate if isinstance(candidate, str) else "" + name = _RESPONSES_COMPLETED[kind] + events = emitter.tool_start(raw_id, name) + events.extend(emitter.tool_end(raw_id)) + output = getattr(item, "results", None) if item is not None else None + if output is None and item is not None: + output = getattr(item, "outputs", None) + events.extend(emitter.tool_result(raw_id, output)) + return events + return [] + + def _stringify_output(value: Any) -> str: if isinstance(value, str): return value @@ -644,16 +743,23 @@ async def _stream_gen( full_output = "" async for event in streamed.stream_events(): + while hooks.pending: + yield hooks.pending.pop(0) if event.type == "raw_response_event": raw = getattr(event, "data", None) - if ( - raw is not None - and getattr(raw, "type", None) == "response.output_text.delta" - ): - delta = getattr(raw, "delta", "") - if isinstance(delta, str) and delta: - yield {"type": "chunk", "text": delta} - full_output += delta + if raw is not None: + for rich in _map_responses_event( + raw, hooks.emitter, hooks.item_to_call + ): + if rich.get("type") == "TEXT_MESSAGE_CONTENT": + full_output += rich.get("delta") or "" + yield rich + while hooks.pending: + yield hooks.pending.pop(0) + for rich in hooks.emitter.end_text(): + yield rich + for rich in hooks.emitter.fail_open(): + yield rich stream_completed = True final_output = streamed.final_output @@ -685,6 +791,10 @@ async def _stream_gen( raise except Exception as exc: hooks.close_open_spans(exc) + while hooks.pending: + yield hooks.pending.pop(0) + for rich in hooks.emitter.fail_open(): + yield rich _write_failed_run_usage(span, config, exc, run_usage) fail_span(span, exc, ended) raise diff --git a/packages/openai-agents/tests/test_handler.py b/packages/openai-agents/tests/test_handler.py index fe0d4af2..38226526 100644 --- a/packages/openai-agents/tests/test_handler.py +++ b/packages/openai-agents/tests/test_handler.py @@ -1325,12 +1325,58 @@ async def test_yields_chunks_then_exactly_one_done_event(self) -> None: CONFIG, "q", {}, {} ) ] - chunks = [e["text"] for e in events if e["type"] == "chunk"] - assert chunks == ["Hello", " world"] + deltas = [e["delta"] for e in events if e["type"] == "TEXT_MESSAGE_CONTENT"] + assert deltas == ["Hello", " world"] + assert events[0]["type"] == "TEXT_MESSAGE_START" + assert any(e["type"] == "TEXT_MESSAGE_END" for e in events) done = [e for e in events if e["type"] == "done"] assert len(done) == 1 assert done[0]["output"] == "Hello world" + async def test_tool_lifecycle_hides_payloads_until_capture_is_on(self) -> None: + turns = [ + { + "deltas": ["Looking"], + "tool_calls": [ + { + "name": "search", + "id": "c1", + "args": {"q": "x"}, + "result": "secret", + } + ], + "output": [], + } + ] + + async def _events(capture: bool) -> list[dict[str, Any]]: + agents_mod = _fake_agents_module( + run_streamed=_make_run_streamed(turns, "done") + ) + with _patched_agents(agents_mod): + return [ + event + async for event in await create_openai_agent_handler( + capture_content=capture + ).stream(CONFIG, "q", {}, {}) + ] + + 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 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"] == json.dumps({"q": "x"}) + shown_result = next( + event for event in shown if event["type"] == "TOOL_CALL_RESULT" + ) + assert shown_result["content"] == "secret" + async def test_opens_the_same_root_span_name_as_the_blocking_path(self) -> None: ctx, rec = _recording() turns = [{"deltas": ["hi"], "output": _text_output("hi")}] diff --git a/packages/openai-messages/src/launchdarkly_ai_openai_messages/handler.py b/packages/openai-messages/src/launchdarkly_ai_openai_messages/handler.py index 6afc2b49..366b4de4 100644 --- a/packages/openai-messages/src/launchdarkly_ai_openai_messages/handler.py +++ b/packages/openai-messages/src/launchdarkly_ai_openai_messages/handler.py @@ -27,6 +27,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 ( @@ -433,6 +434,95 @@ def _final_output_messages(output: str) -> list[SpanMessage]: ] +_RESPONSES_ARG_DELTA = { + "response.function_call_arguments.delta": "delta", + "response.mcp_call_arguments.delta": "delta", + "response.custom_tool_call_input.delta": "delta", +} +_RESPONSES_ARG_DONE = { + "response.function_call_arguments.done": "arguments", + "response.mcp_call_arguments.done": "arguments", + "response.custom_tool_call_input.done": "input", +} +_RESPONSES_OPEN_TYPES = frozenset({"function_call", "mcp_call", "custom_tool_call"}) +_RESPONSES_COMPLETED = { + "response.file_search_call.completed": "file_search", + "response.code_interpreter_call.completed": "code_interpreter", +} + + +def _responses_call_id(event: Any, item_to_call: dict[str, str]) -> str: + direct = getattr(event, "call_id", None) + if isinstance(direct, str) and direct: + return direct + item_id = getattr(event, "item_id", None) + if isinstance(item_id, str) and item_id in item_to_call: + return item_to_call[item_id] + return item_id if isinstance(item_id, str) else "" + + +def _map_responses_event( + event: Any, emitter: TurnEmitter, item_to_call: dict[str, str] +) -> list[dict[str, Any]]: + """Map one Responses stream event onto the AG-UI subset. Unknown events are ignored.""" + kind = getattr(event, "type", None) + if kind == "response.output_text.delta": + delta = getattr(event, "delta", "") + if not isinstance(delta, str) or not delta: + return [] + return emitter.text(delta) + if kind == "response.output_item.added": + item = getattr(event, "item", None) + if getattr(item, "type", None) not in _RESPONSES_OPEN_TYPES: + return [] + call_id = getattr(item, "call_id", None) or getattr(item, "id", None) or "" + name = getattr(item, "name", None) or "" + item_id = getattr(item, "id", None) + if isinstance(item_id, str) and isinstance(call_id, str) and call_id: + item_to_call[item_id] = call_id + return emitter.tool_start(str(call_id), str(name)) + if kind in _RESPONSES_ARG_DELTA: + delta = getattr(event, "delta", "") + if not isinstance(delta, str): + return [] + return emitter.tool_args(_responses_call_id(event, item_to_call), delta) + if kind in _RESPONSES_ARG_DONE: + arguments = getattr(event, _RESPONSES_ARG_DONE[kind], None) + return emitter.tool_end( + _responses_call_id(event, item_to_call), arguments=arguments + ) + if kind in ("response.mcp_call.completed", "response.mcp_call.failed"): + call_id = _responses_call_id(event, item_to_call) + item = getattr(event, "item", None) + if not call_id and item is not None: + raw_id = getattr(item, "call_id", None) or getattr(item, "id", None) + call_id = raw_id if isinstance(raw_id, str) else "" + if kind == "response.mcp_call.failed": + return emitter.tool_result(call_id, failed=True) + output = getattr(event, "output", None) + if output is None and item is not None: + output = getattr(item, "output", None) + return emitter.tool_result(call_id, output) + if kind in _RESPONSES_COMPLETED: + item = getattr(event, "item", None) + raw_id = "" + if item is not None: + candidate = getattr(item, "call_id", None) or getattr(item, "id", None) + raw_id = candidate if isinstance(candidate, str) else "" + if not raw_id: + candidate = getattr(event, "item_id", None) + raw_id = candidate if isinstance(candidate, str) else "" + name = _RESPONSES_COMPLETED[kind] + events = emitter.tool_start(raw_id, name) + events.extend(emitter.tool_end(raw_id)) + output = getattr(item, "results", None) if item is not None else None + if output is None and item is not None: + output = getattr(item, "outputs", None) + events.extend(emitter.tool_result(raw_id, output)) + return events + return [] + + async def _stream_gen( client: Any, config: AiConfigRep, @@ -485,6 +575,8 @@ async def _stream_gen( previous_response_id: str | None = None current_input: Any = input_messages steps = 0 + emitter = TurnEmitter(capture_content=capture_content) + item_to_call: dict[str, str] = {} while True: model_span = start_model_span(config, parent) @@ -515,10 +607,10 @@ async def _stream_gen( stream = client.responses.stream(**stream_params) async with stream as s: async for event in s: - if getattr(event, "type", None) == "response.output_text.delta": - text = getattr(event, "delta", "") - full_output += text - yield {"type": "chunk", "text": text} + for rich in _map_responses_event(event, emitter, item_to_call): + if rich.get("type") == "TEXT_MESSAGE_CONTENT": + full_output += rich.get("delta") or "" + yield rich final_resp = await s.get_final_response() @@ -548,6 +640,8 @@ async def _stream_gen( if getattr(item, "type", None) == "function_call" ] if not tool_calls: + for rich in emitter.end_text(): + yield rich break if steps >= _MAX_STEPS: @@ -557,6 +651,13 @@ async def _stream_gen( steps += 1 previous_response_id = getattr(final_resp, "id", None) + for tc in tool_calls: + for rich in emitter.tool_start(str(tc.call_id), str(tc.name)): + yield rich + for rich in emitter.tool_end(str(tc.call_id), arguments=tc.arguments): + yield rich + for rich in emitter.end_text(): + yield rich tool_outputs = [] for tc in tool_calls: tool_span = start_tool_span(tc.name, tc.call_id, parent) @@ -585,10 +686,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(tc.call_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(tc.call_id), result): + yield rich tool_outputs.append( { "type": "function_call_output", diff --git a/packages/openai-messages/tests/test_handler.py b/packages/openai-messages/tests/test_handler.py index 35da2f9d..80e7718d 100644 --- a/packages/openai-messages/tests/test_handler.py +++ b/packages/openai-messages/tests/test_handler.py @@ -1281,10 +1281,55 @@ async def test_yields_chunk_events(self, mock_openai: MagicMock) -> None: with patch.object(spans_mod, "_HAS_OTEL", False): h = create_openai_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" + assert events[-1]["type"] == "done" + + async def test_tool_lifecycle_hides_payloads_until_capture_is_on( + self, mock_openai: MagicMock + ) -> None: + import launchdarkly_ai_openai_messages.spans as spans_mod + from launchdarkly_ai_openai_messages import create_openai_messages_handler + + tool_call = _function_call_item("search", "c1", {"q": "x"}) + + async def search(args: dict[str, Any]) -> str: + return "secret" + + async def _events(capture: bool) -> list[dict[str, Any]]: + mock_openai.responses.stream = MagicMock( + side_effect=[ + _make_openai_stream_context(["Looking"], output=[tool_call]), + _make_openai_stream_context(["Done"]), + ] + ) + with patch.object(spans_mod, "_HAS_OTEL", False): + handler = create_openai_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"] == json.dumps({"q": "x"}) + shown_result = next( + event for event in shown if event["type"] == "TOOL_CALL_RESULT" + ) + assert shown_result["content"] == "secret" async def test_yields_exactly_one_done_event(self, mock_openai: MagicMock) -> None: import launchdarkly_ai_openai_messages.spans as spans_mod From c35077dc7a3f815a5566bb37463f7250a23d939f Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 13 Oct 2026 15:56:05 +0000 Subject: [PATCH 04/15] fix(openai): keep streamed tool arguments and hosted results Do not close a call from the tool hook once argument deltas have started. Read hosted tool output from response.output_item.done, and close every open call when the messages stream fails. Co-authored-by: alexis --- .../client/src/launchdarkly_ai_server/agui.py | 4 + .../launchdarkly_ai_openai_agents/handler.py | 87 +++++++++++-------- packages/openai-agents/tests/test_handler.py | 22 +++++ .../handler.py | 80 ++++++++++------- .../openai-messages/tests/test_handler.py | 67 ++++++++++++++ 5 files changed, 193 insertions(+), 67 deletions(-) diff --git a/packages/client/src/launchdarkly_ai_server/agui.py b/packages/client/src/launchdarkly_ai_server/agui.py index cb750553..cb5bb37f 100644 --- a/packages/client/src/launchdarkly_ai_server/agui.py +++ b/packages/client/src/launchdarkly_ai_server/agui.py @@ -150,6 +150,10 @@ def tool_start( event["parentMessageId"] = self._open return [event] + 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 tool_args(self, tool_call_id: str, delta: str) -> list[dict[str, Any]]: if ( not self.capture_content diff --git a/packages/openai-agents/src/launchdarkly_ai_openai_agents/handler.py b/packages/openai-agents/src/launchdarkly_ai_openai_agents/handler.py index f6824159..d0f8ba5e 100644 --- a/packages/openai-agents/src/launchdarkly_ai_openai_agents/handler.py +++ b/packages/openai-agents/src/launchdarkly_ai_openai_agents/handler.py @@ -325,8 +325,15 @@ async def on_tool_start(self, context: Any, agent: Any, tool: Any) -> None: # exports. self.open_tool_spans[str(call_id)] = span self.pending.extend(self.emitter.tool_start(str(call_id), str(name))) - tool_args = getattr(context, "tool_arguments", None) - self.pending.extend(self.emitter.tool_end(str(call_id), arguments=tool_args)) + # The Runner can invoke this hook from its background task while + # argument deltas are still unread. Closing now would drop the rest + # and skip the full argument string, because a prefix is already + # recorded. arguments.done closes a call that is still streaming. + if not self.emitter.has_argument_deltas(str(call_id)): + tool_args = getattr(context, "tool_arguments", None) + self.pending.extend( + self.emitter.tool_end(str(call_id), arguments=tool_args) + ) if self.capture_content: # `tool_arguments` is the raw JSON args string the model produced. Parsed here rather # than passed through, per TELEMETRY-CONTRACT.md section 12: arguments hold the object @@ -597,9 +604,14 @@ def _stream_impl( "response.custom_tool_call_input.done": "input", } _RESPONSES_OPEN_TYPES = frozenset({"function_call", "mcp_call", "custom_tool_call"}) -_RESPONSES_COMPLETED = { - "response.file_search_call.completed": "file_search", - "response.code_interpreter_call.completed": "code_interpreter", +_HOSTED_OUTPUT_FIELD = { + "mcp_call": "output", + "file_search_call": "results", + "code_interpreter_call": "outputs", +} +_HOSTED_TOOL_NAME = { + "file_search_call": "file_search", + "code_interpreter_call": "code_interpreter", } @@ -613,6 +625,38 @@ def _responses_call_id(event: Any, item_to_call: dict[str, str]) -> str: return item_id if isinstance(item_id, str) else "" +def _hosted_item_result( + item: Any, emitter: TurnEmitter, item_to_call: dict[str, str] +) -> list[dict[str, Any]]: + """Result for a hosted tool from ``response.output_item.done``. + + The ``*.completed`` events only carry ``item_id``. Reading ``output`` off + them yields ``None``, and with capture on that becomes the string ``null`` + and seals the call before the real item arrives. + """ + item_type = getattr(item, "type", None) + if item_type not in _HOSTED_OUTPUT_FIELD: + return [] + call_id = getattr(item, "call_id", None) or getattr(item, "id", None) or "" + item_id = getattr(item, "id", None) + if isinstance(item_id, str) and isinstance(call_id, str) and call_id: + item_to_call[item_id] = call_id + call_id = str(call_id) + name = getattr(item, "name", None) or _HOSTED_TOOL_NAME.get(item_type, "") + events = emitter.tool_start(call_id, str(name)) + events.extend(emitter.tool_end(call_id)) + status = getattr(item, "status", None) + if status in ("failed", "incomplete") or getattr(item, "error", None): + events.extend(emitter.tool_result(call_id, failed=True)) + else: + events.extend( + emitter.tool_result( + call_id, getattr(item, _HOSTED_OUTPUT_FIELD[item_type], None) + ) + ) + return events + + def _map_responses_event( event: Any, emitter: TurnEmitter, item_to_call: dict[str, str] ) -> list[dict[str, Any]]: @@ -643,35 +687,10 @@ def _map_responses_event( return emitter.tool_end( _responses_call_id(event, item_to_call), arguments=arguments ) - if kind in ("response.mcp_call.completed", "response.mcp_call.failed"): - call_id = _responses_call_id(event, item_to_call) - item = getattr(event, "item", None) - if not call_id and item is not None: - raw_id = getattr(item, "call_id", None) or getattr(item, "id", None) - call_id = raw_id if isinstance(raw_id, str) else "" - if kind == "response.mcp_call.failed": - return emitter.tool_result(call_id, failed=True) - output = getattr(event, "output", None) - if output is None and item is not None: - output = getattr(item, "output", None) - return emitter.tool_result(call_id, output) - if kind in _RESPONSES_COMPLETED: - item = getattr(event, "item", None) - raw_id = "" - if item is not None: - candidate = getattr(item, "call_id", None) or getattr(item, "id", None) - raw_id = candidate if isinstance(candidate, str) else "" - if not raw_id: - candidate = getattr(event, "item_id", None) - raw_id = candidate if isinstance(candidate, str) else "" - name = _RESPONSES_COMPLETED[kind] - events = emitter.tool_start(raw_id, name) - events.extend(emitter.tool_end(raw_id)) - output = getattr(item, "results", None) if item is not None else None - if output is None and item is not None: - output = getattr(item, "outputs", None) - events.extend(emitter.tool_result(raw_id, output)) - return events + if kind == "response.mcp_call.failed": + return emitter.tool_result(_responses_call_id(event, item_to_call), failed=True) + if kind == "response.output_item.done": + return _hosted_item_result(getattr(event, "item", None), emitter, item_to_call) return [] diff --git a/packages/openai-agents/tests/test_handler.py b/packages/openai-agents/tests/test_handler.py index 38226526..6334276e 100644 --- a/packages/openai-agents/tests/test_handler.py +++ b/packages/openai-agents/tests/test_handler.py @@ -1377,6 +1377,28 @@ async def _events(capture: bool) -> list[dict[str, Any]]: ) assert shown_result["content"] == "secret" + async def test_a_streaming_prefix_is_not_closed_by_the_tool_hook(self) -> None: + from launchdarkly_ai_openai_agents.handler import _SpanningHooks + from launchdarkly_ai_server import create_run_usage + + hooks = _SpanningHooks(CONFIG, None, True, create_run_usage()) + hooks.emitter.tool_start("c1", "search") + hooks.emitter.tool_args("c1", '{"q":') + ctx = SimpleNamespace( + tool_name="search", + tool_call_id="c1", + tool_arguments='{"q":"x"}', + ) + await hooks.on_tool_start( + ctx, SimpleNamespace(), SimpleNamespace(name="search") + ) + assert not any(event["type"] == "TOOL_CALL_END" for event in hooks.pending) + rest = hooks.emitter.tool_args("c1", '"}') + assert rest[0]["delta"] == '"}' + assert any( + event["type"] == "TOOL_CALL_END" for event in hooks.emitter.tool_end("c1") + ) + async def test_opens_the_same_root_span_name_as_the_blocking_path(self) -> None: ctx, rec = _recording() turns = [{"deltas": ["hi"], "output": _text_output("hi")}] diff --git a/packages/openai-messages/src/launchdarkly_ai_openai_messages/handler.py b/packages/openai-messages/src/launchdarkly_ai_openai_messages/handler.py index 366b4de4..0f08ed83 100644 --- a/packages/openai-messages/src/launchdarkly_ai_openai_messages/handler.py +++ b/packages/openai-messages/src/launchdarkly_ai_openai_messages/handler.py @@ -445,9 +445,14 @@ def _final_output_messages(output: str) -> list[SpanMessage]: "response.custom_tool_call_input.done": "input", } _RESPONSES_OPEN_TYPES = frozenset({"function_call", "mcp_call", "custom_tool_call"}) -_RESPONSES_COMPLETED = { - "response.file_search_call.completed": "file_search", - "response.code_interpreter_call.completed": "code_interpreter", +_HOSTED_OUTPUT_FIELD = { + "mcp_call": "output", + "file_search_call": "results", + "code_interpreter_call": "outputs", +} +_HOSTED_TOOL_NAME = { + "file_search_call": "file_search", + "code_interpreter_call": "code_interpreter", } @@ -461,6 +466,38 @@ def _responses_call_id(event: Any, item_to_call: dict[str, str]) -> str: return item_id if isinstance(item_id, str) else "" +def _hosted_item_result( + item: Any, emitter: TurnEmitter, item_to_call: dict[str, str] +) -> list[dict[str, Any]]: + """Result for a hosted tool from ``response.output_item.done``. + + The ``*.completed`` events only carry ``item_id``. Reading ``output`` off + them yields ``None``, and with capture on that becomes the string ``null`` + and seals the call before the real item arrives. + """ + item_type = getattr(item, "type", None) + if item_type not in _HOSTED_OUTPUT_FIELD: + return [] + call_id = getattr(item, "call_id", None) or getattr(item, "id", None) or "" + item_id = getattr(item, "id", None) + if isinstance(item_id, str) and isinstance(call_id, str) and call_id: + item_to_call[item_id] = call_id + call_id = str(call_id) + name = getattr(item, "name", None) or _HOSTED_TOOL_NAME.get(item_type, "") + events = emitter.tool_start(call_id, str(name)) + events.extend(emitter.tool_end(call_id)) + status = getattr(item, "status", None) + if status in ("failed", "incomplete") or getattr(item, "error", None): + events.extend(emitter.tool_result(call_id, failed=True)) + else: + events.extend( + emitter.tool_result( + call_id, getattr(item, _HOSTED_OUTPUT_FIELD[item_type], None) + ) + ) + return events + + def _map_responses_event( event: Any, emitter: TurnEmitter, item_to_call: dict[str, str] ) -> list[dict[str, Any]]: @@ -491,35 +528,10 @@ def _map_responses_event( return emitter.tool_end( _responses_call_id(event, item_to_call), arguments=arguments ) - if kind in ("response.mcp_call.completed", "response.mcp_call.failed"): - call_id = _responses_call_id(event, item_to_call) - item = getattr(event, "item", None) - if not call_id and item is not None: - raw_id = getattr(item, "call_id", None) or getattr(item, "id", None) - call_id = raw_id if isinstance(raw_id, str) else "" - if kind == "response.mcp_call.failed": - return emitter.tool_result(call_id, failed=True) - output = getattr(event, "output", None) - if output is None and item is not None: - output = getattr(item, "output", None) - return emitter.tool_result(call_id, output) - if kind in _RESPONSES_COMPLETED: - item = getattr(event, "item", None) - raw_id = "" - if item is not None: - candidate = getattr(item, "call_id", None) or getattr(item, "id", None) - raw_id = candidate if isinstance(candidate, str) else "" - if not raw_id: - candidate = getattr(event, "item_id", None) - raw_id = candidate if isinstance(candidate, str) else "" - name = _RESPONSES_COMPLETED[kind] - events = emitter.tool_start(raw_id, name) - events.extend(emitter.tool_end(raw_id)) - output = getattr(item, "results", None) if item is not None else None - if output is None and item is not None: - output = getattr(item, "outputs", None) - events.extend(emitter.tool_result(raw_id, output)) - return events + if kind == "response.mcp_call.failed": + return emitter.tool_result(_responses_call_id(event, item_to_call), failed=True) + if kind == "response.output_item.done": + return _hosted_item_result(getattr(event, "item", None), emitter, item_to_call) return [] @@ -558,6 +570,7 @@ 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) try: tools = _build_tools(config.get("tools") or {}) input_messages = _build_input_messages(config, user_input, variables, history) @@ -575,7 +588,6 @@ async def _stream_gen( previous_response_id: str | None = None current_input: Any = input_messages steps = 0 - emitter = TurnEmitter(capture_content=capture_content) item_to_call: dict[str, str] = {} while True: @@ -728,6 +740,8 @@ async def _stream_gen( if run_usage.reported: finish_root_span(span, last_response_model, 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/openai-messages/tests/test_handler.py b/packages/openai-messages/tests/test_handler.py index 80e7718d..60b3bc6c 100644 --- a/packages/openai-messages/tests/test_handler.py +++ b/packages/openai-messages/tests/test_handler.py @@ -1331,6 +1331,73 @@ async def _events(capture: bool) -> list[dict[str, Any]]: ) assert shown_result["content"] == "secret" + async def test_a_failed_call_closes_its_siblings( + self, mock_openai: MagicMock + ) -> None: + import launchdarkly_ai_openai_messages.spans as spans_mod + from launchdarkly_ai_openai_messages import create_openai_messages_handler + + calls = [ + _function_call_item("boom", "c1", {}), + _function_call_item("later", "c2", {}), + ] + mock_openai.responses.stream = MagicMock( + return_value=_make_openai_stream_context(["x"], output=calls) + ) + + def boom(_args: dict[str, Any]) -> str: + raise RuntimeError("nope") + + with patch.object(spans_mod, "_HAS_OTEL", False): + handler = create_openai_messages_handler() + seen: list[dict[str, Any]] = [] + with pytest.raises(RuntimeError, match="nope"): + 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"} + + def test_hosted_completion_waits_for_the_output_item(self) -> None: + from types import SimpleNamespace + + from launchdarkly_ai_openai_messages.handler import _map_responses_event + from launchdarkly_ai_server.agui import TurnEmitter + + emitter = TurnEmitter(capture_content=True) + ids: dict[str, str] = {} + added = SimpleNamespace( + type="response.output_item.added", + item=SimpleNamespace( + type="mcp_call", call_id="m1", id="item-1", name="search" + ), + ) + completed = SimpleNamespace( + type="response.mcp_call.completed", item_id="item-1", output=None + ) + done = SimpleNamespace( + type="response.output_item.done", + item=SimpleNamespace( + type="mcp_call", + call_id="m1", + id="item-1", + name="search", + status="completed", + output="found", + ), + ) + events: list[dict[str, Any]] = [] + events.extend(_map_responses_event(added, emitter, ids)) + events.extend(_map_responses_event(completed, emitter, ids)) + events.extend(_map_responses_event(done, emitter, ids)) + assert "null" not in str(events) + result = next(event for event in events if event["type"] == "TOOL_CALL_RESULT") + assert result["content"] == "found" + assert [event["type"] for event in events].count("TOOL_CALL_RESULT") == 1 + async def test_yields_exactly_one_done_event(self, mock_openai: MagicMock) -> None: import launchdarkly_ai_openai_messages.spans as spans_mod from launchdarkly_ai_openai_messages import create_openai_messages_handler From 4d7504c3bedc013f7bf9be2650be30a1d34a2da5 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 13 Oct 2026 16:36:11 +0000 Subject: [PATCH 05/15] fix(openai): serialize hosted results and close every open call File-search results and code-interpreter outputs are SDK models, so json.dumps marked a successful call failed. A streamed prefix no longer blocks the rest of the arguments, and a messages failure or a call that never becomes a function result still gets TOOL_CALL_RESULT. Co-authored-by: alexis --- .../client/src/launchdarkly_ai_server/agui.py | 22 ++-- packages/client/tests/test_agui.py | 10 ++ .../launchdarkly_ai_openai_agents/handler.py | 27 +++- packages/openai-agents/tests/test_handler.py | 33 +++++ .../handler.py | 33 ++++- .../openai-messages/tests/test_handler.py | 124 ++++++++++++++++++ 6 files changed, 239 insertions(+), 10 deletions(-) diff --git a/packages/client/src/launchdarkly_ai_server/agui.py b/packages/client/src/launchdarkly_ai_server/agui.py index cb5bb37f..dfd250e3 100644 --- a/packages/client/src/launchdarkly_ai_server/agui.py +++ b/packages/client/src/launchdarkly_ai_server/agui.py @@ -104,6 +104,7 @@ def __init__(self, *, capture_content: bool) -> 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() @@ -163,25 +164,30 @@ 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. + A streamed prefix is extended with the missing suffix of those + arguments, so a partial delta cannot publish truncated JSON. + """ 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}) diff --git a/packages/client/tests/test_agui.py b/packages/client/tests/test_agui.py index b104f266..e2297666 100644 --- a/packages/client/tests/test_agui.py +++ b/packages/client/tests/test_agui.py @@ -149,6 +149,16 @@ def test_fail_open_masks_a_call_that_never_finished(self) -> None: assert "content" not in events[-1] assert events[-1]["metadata"]["launchdarkly"]["toolFailed"] is True + def test_tool_end_extends_a_streamed_argument_prefix(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"}' + def test_a_second_start_for_the_same_id_is_dropped(self) -> None: emitter = TurnEmitter(capture_content=True) assert emitter.tool_start("a", "search") diff --git a/packages/openai-agents/src/launchdarkly_ai_openai_agents/handler.py b/packages/openai-agents/src/launchdarkly_ai_openai_agents/handler.py index d0f8ba5e..e652fbcc 100644 --- a/packages/openai-agents/src/launchdarkly_ai_openai_agents/handler.py +++ b/packages/openai-agents/src/launchdarkly_ai_openai_agents/handler.py @@ -625,6 +625,30 @@ def _responses_call_id(event: Any, item_to_call: dict[str, str]) -> str: return item_id if isinstance(item_id, str) else "" +def _json_safe(value: Any) -> Any: + """Plain JSON for a hosted tool payload. + + File-search results and code-interpreter outputs are SDK model objects. + ``json.dumps`` rejects those, and a successful call would otherwise be + reported as a failed tool result. + """ + if value is None or isinstance(value, (str, int, float, bool)): + return value + if isinstance(value, dict): + return {str(key): _json_safe(item) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return [_json_safe(item) for item in value] + dump = getattr(value, "model_dump", None) + if callable(dump): + try: + dumped = dump(mode="json") + except TypeError: + dumped = dump() + if dumped is not value: + return _json_safe(dumped) + return value + + def _hosted_item_result( item: Any, emitter: TurnEmitter, item_to_call: dict[str, str] ) -> list[dict[str, Any]]: @@ -651,7 +675,8 @@ def _hosted_item_result( else: events.extend( emitter.tool_result( - call_id, getattr(item, _HOSTED_OUTPUT_FIELD[item_type], None) + call_id, + _json_safe(getattr(item, _HOSTED_OUTPUT_FIELD[item_type], None)), ) ) return events diff --git a/packages/openai-agents/tests/test_handler.py b/packages/openai-agents/tests/test_handler.py index 6334276e..276f6857 100644 --- a/packages/openai-agents/tests/test_handler.py +++ b/packages/openai-agents/tests/test_handler.py @@ -2174,3 +2174,36 @@ async def _drain() -> None: assert rec.root.attributes.get("launchdarkly.run.cancelled") is True assert "launchdarkly.stream.abandoned" not in rec.root.attributes + + +class TestHostedToolResults: + def test_model_objects_are_json_not_failures(self) -> None: + from types import SimpleNamespace + + from launchdarkly_ai_openai_agents.handler import _map_responses_event + from launchdarkly_ai_server.agui import TurnEmitter + + class _Row: + def model_dump(self, **_kwargs: object) -> dict[str, str]: + return {"file_id": "f1", "text": "hit"} + + for item_type, call_id, field in ( + ("file_search_call", "fs1", "results"), + ("code_interpreter_call", "ci1", "outputs"), + ): + emitter = TurnEmitter(capture_content=True) + done = SimpleNamespace( + type="response.output_item.done", + item=SimpleNamespace( + type=item_type, + id=call_id, + status="completed", + **{field: [_Row()]}, + ), + ) + events = _map_responses_event(done, emitter, {}) + result = next( + event for event in events if event["type"] == "TOOL_CALL_RESULT" + ) + assert result["content"] == '[{"file_id":"f1","text":"hit"}]' + assert "metadata" not in result diff --git a/packages/openai-messages/src/launchdarkly_ai_openai_messages/handler.py b/packages/openai-messages/src/launchdarkly_ai_openai_messages/handler.py index 0f08ed83..d56b9239 100644 --- a/packages/openai-messages/src/launchdarkly_ai_openai_messages/handler.py +++ b/packages/openai-messages/src/launchdarkly_ai_openai_messages/handler.py @@ -466,6 +466,30 @@ def _responses_call_id(event: Any, item_to_call: dict[str, str]) -> str: return item_id if isinstance(item_id, str) else "" +def _json_safe(value: Any) -> Any: + """Plain JSON for a hosted tool payload. + + File-search results and code-interpreter outputs are SDK model objects. + ``json.dumps`` rejects those, and a successful call would otherwise be + reported as a failed tool result. + """ + if value is None or isinstance(value, (str, int, float, bool)): + return value + if isinstance(value, dict): + return {str(key): _json_safe(item) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return [_json_safe(item) for item in value] + dump = getattr(value, "model_dump", None) + if callable(dump): + try: + dumped = dump(mode="json") + except TypeError: + dumped = dump() + if dumped is not value: + return _json_safe(dumped) + return value + + def _hosted_item_result( item: Any, emitter: TurnEmitter, item_to_call: dict[str, str] ) -> list[dict[str, Any]]: @@ -492,7 +516,8 @@ def _hosted_item_result( else: events.extend( emitter.tool_result( - call_id, getattr(item, _HOSTED_OUTPUT_FIELD[item_type], None) + call_id, + _json_safe(getattr(item, _HOSTED_OUTPUT_FIELD[item_type], None)), ) ) return events @@ -644,6 +669,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 tool_calls = [ @@ -654,6 +681,8 @@ async def _stream_gen( if not tool_calls: for rich in emitter.end_text(): yield rich + for rich in emitter.fail_open(): + yield rich break if steps >= _MAX_STEPS: @@ -700,6 +729,8 @@ async def _stream_gen( open_tool_span = None for rich in emitter.tool_result(str(tc.call_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/openai-messages/tests/test_handler.py b/packages/openai-messages/tests/test_handler.py index 60b3bc6c..c22e36a5 100644 --- a/packages/openai-messages/tests/test_handler.py +++ b/packages/openai-messages/tests/test_handler.py @@ -1398,6 +1398,130 @@ def test_hosted_completion_waits_for_the_output_item(self) -> None: assert result["content"] == "found" assert [event["type"] for event in events].count("TOOL_CALL_RESULT") == 1 + def test_hosted_model_results_are_json_not_failures(self) -> None: + from types import SimpleNamespace + + from launchdarkly_ai_openai_messages.handler import _map_responses_event + from launchdarkly_ai_server.agui import TurnEmitter + + class _Row: + def model_dump(self, **_kwargs: Any) -> dict[str, str]: + return {"file_id": "f1", "text": "hit"} + + def _done(item_type: str, call_id: str, field: str) -> Any: + return SimpleNamespace( + type="response.output_item.done", + item=SimpleNamespace( + type=item_type, + id=call_id, + status="completed", + **{field: [_Row()]}, + ), + ) + + for item_type, call_id, field in ( + ("file_search_call", "fs1", "results"), + ("code_interpreter_call", "ci1", "outputs"), + ): + emitter = TurnEmitter(capture_content=True) + events = _map_responses_event(_done(item_type, call_id, field), emitter, {}) + result = next( + event for event in events if event["type"] == "TOOL_CALL_RESULT" + ) + assert result["toolCallId"] == call_id + assert result["content"] == '[{"file_id":"f1","text":"hit"}]' + assert "metadata" not in result + + async def test_an_opened_call_without_a_function_result_is_closed( + self, mock_openai: MagicMock + ) -> None: + from types import SimpleNamespace + + import launchdarkly_ai_openai_messages.spans as spans_mod + from launchdarkly_ai_openai_messages import create_openai_messages_handler + + added = SimpleNamespace( + type="response.output_item.added", + item=SimpleNamespace( + type="function_call", call_id="c1", id="item-1", name="search" + ), + ) + + class _Stream: + def __aiter__(self) -> AsyncIterator[Any]: + return self._iter() + + async def _iter(self) -> AsyncIterator[Any]: + yield added + + async def get_final_response(self) -> Any: + final = MagicMock() + final.output = [] + final.status = "completed" + final.incomplete_details = None + final.usage = MagicMock(input_tokens=1, output_tokens=1) + final.usage.input_tokens_details = MagicMock(cached_tokens=0) + final.id = "resp-stream" + final.model = "gpt-4o" + return final + + @asynccontextmanager + async def _mgr() -> AsyncGenerator[Any, None]: + yield _Stream() + + mock_openai.responses.stream = MagicMock(return_value=_mgr()) + with patch.object(spans_mod, "_HAS_OTEL", False): + handler = create_openai_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_a_provider_error_closes_an_opened_call( + self, mock_openai: MagicMock + ) -> None: + from types import SimpleNamespace + + import launchdarkly_ai_openai_messages.spans as spans_mod + from launchdarkly_ai_openai_messages import create_openai_messages_handler + + added = SimpleNamespace( + type="response.output_item.added", + item=SimpleNamespace( + type="function_call", call_id="c1", id="item-1", name="search" + ), + ) + + class _Stream: + def __aiter__(self) -> AsyncIterator[Any]: + return self._iter() + + async def _iter(self) -> AsyncIterator[Any]: + yield added + raise RuntimeError("stream broke") + + async def get_final_response(self) -> Any: + raise AssertionError("final response is not read") + + @asynccontextmanager + async def _mgr() -> AsyncGenerator[Any, None]: + yield _Stream() + + mock_openai.responses.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_openai_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_yields_exactly_one_done_event(self, mock_openai: MagicMock) -> None: import launchdarkly_ai_openai_messages.spans as spans_mod from launchdarkly_ai_openai_messages import create_openai_messages_handler From e8c324ed7a6b1e9552d219b4e9903f61516a0506 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 13 Oct 2026 18:31:41 +0000 Subject: [PATCH 06/15] fix(openai): close tool arguments and text in read order The agents hook can finish a tool before the consumer has read the rest of its argument deltas. The result now closes with the full argument string. A new response closes the open text message from the stream itself, not from the background hook. Co-authored-by: alexis --- .../launchdarkly_ai_openai_agents/handler.py | 14 ++++++- packages/openai-agents/tests/test_handler.py | 39 +++++++++++++++++++ .../handler.py | 8 +--- 3 files changed, 54 insertions(+), 7 deletions(-) diff --git a/packages/openai-agents/src/launchdarkly_ai_openai_agents/handler.py b/packages/openai-agents/src/launchdarkly_ai_openai_agents/handler.py index e652fbcc..a6d0261f 100644 --- a/packages/openai-agents/src/launchdarkly_ai_openai_agents/handler.py +++ b/packages/openai-agents/src/launchdarkly_ai_openai_agents/handler.py @@ -257,7 +257,9 @@ def __init__( async def on_llm_start( self, context: Any, agent: Any, system_prompt: str | None, input_items: Any ) -> None: - self.pending.extend(self.emitter.end_text()) + # Text is closed from response.created, in the order the consumer + # reads raw events. Closing here runs ahead of unread deltas and + # splits one turn across two messages. span = start_model_span(self.config, self.parent) self.open_model_span = span if self.capture_content: @@ -351,6 +353,14 @@ async def on_tool_end( self, context: Any, agent: Any, tool: Any, result: Any ) -> None: call_id = self._tool_span_key(context, tool) + # The hook can run before unread argument deltas. Close with the full + # argument string first, so tool_result does not seal a prefix. + self.pending.extend( + self.emitter.tool_end( + str(call_id), + arguments=getattr(context, "tool_arguments", None), + ) + ) self.pending.extend(self.emitter.tool_result(str(call_id), result)) span = self.open_tool_spans.pop(call_id, None) if span is None: @@ -687,6 +697,8 @@ def _map_responses_event( ) -> list[dict[str, Any]]: """Map one Responses stream event onto the AG-UI subset. Unknown events are ignored.""" kind = getattr(event, "type", None) + if kind == "response.created": + return emitter.end_text() if kind == "response.output_text.delta": delta = getattr(event, "delta", "") if not isinstance(delta, str) or not delta: diff --git a/packages/openai-agents/tests/test_handler.py b/packages/openai-agents/tests/test_handler.py index 276f6857..f2e7be88 100644 --- a/packages/openai-agents/tests/test_handler.py +++ b/packages/openai-agents/tests/test_handler.py @@ -1399,6 +1399,45 @@ async def test_a_streaming_prefix_is_not_closed_by_the_tool_hook(self) -> None: event["type"] == "TOOL_CALL_END" for event in hooks.emitter.tool_end("c1") ) + async def test_tool_end_publishes_arguments_the_consumer_has_not_read(self) -> None: + from launchdarkly_ai_openai_agents.handler import _SpanningHooks + from launchdarkly_ai_server import create_run_usage + + hooks = _SpanningHooks(CONFIG, None, True, create_run_usage()) + hooks.emitter.tool_start("c1", "search") + hooks.emitter.tool_args("c1", '{"q":') + ctx = SimpleNamespace( + tool_name="search", + tool_call_id="c1", + tool_arguments='{"q":"x"}', + ) + await hooks.on_tool_end( + ctx, SimpleNamespace(), SimpleNamespace(name="search"), "found" + ) + deltas = [ + event["delta"] + for event in hooks.pending + if event["type"] == "TOOL_CALL_ARGS" + ] + assert "".join(['{"q":', *deltas]) == '{"q":"x"}' + assert any(event["type"] == "TOOL_CALL_RESULT" for event in hooks.pending) + assert hooks.emitter.tool_args("c1", "more") == [] + + def test_a_new_response_closes_text_in_read_order(self) -> None: + from launchdarkly_ai_openai_agents.handler import _map_responses_event + from launchdarkly_ai_server.agui import TurnEmitter + + emitter = TurnEmitter(capture_content=True) + opened = emitter.text("Let me ") + closed = _map_responses_event( + SimpleNamespace(type="response.created"), emitter, {} + ) + assert closed[0]["type"] == "TEXT_MESSAGE_END" + assert closed[0]["messageId"] == opened[0]["messageId"] + nxt = emitter.text("Answer") + assert nxt[0]["type"] == "TEXT_MESSAGE_START" + assert nxt[0]["messageId"] != opened[0]["messageId"] + async def test_opens_the_same_root_span_name_as_the_blocking_path(self) -> None: ctx, rec = _recording() turns = [{"deltas": ["hi"], "output": _text_output("hi")}] diff --git a/packages/openai-messages/src/launchdarkly_ai_openai_messages/handler.py b/packages/openai-messages/src/launchdarkly_ai_openai_messages/handler.py index d56b9239..4e83a506 100644 --- a/packages/openai-messages/src/launchdarkly_ai_openai_messages/handler.py +++ b/packages/openai-messages/src/launchdarkly_ai_openai_messages/handler.py @@ -528,6 +528,8 @@ def _map_responses_event( ) -> list[dict[str, Any]]: """Map one Responses stream event onto the AG-UI subset. Unknown events are ignored.""" kind = getattr(event, "type", None) + if kind == "response.created": + return emitter.end_text() if kind == "response.output_text.delta": delta = getattr(event, "delta", "") if not isinstance(delta, str) or not delta: @@ -669,8 +671,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 tool_calls = [ @@ -727,10 +727,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(tc.call_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 From 57abadc0b6c800a3879237c15441bafeeb47d150 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 13 Oct 2026 18:45:28 +0000 Subject: [PATCH 07/15] fix(openai): close the open text message when a turn fails Co-authored-by: alexis --- .../launchdarkly_ai_openai_agents/handler.py | 6 ++++-- packages/openai-agents/tests/test_handler.py | 19 +++++++++++++++++++ 2 files changed, 23 insertions(+), 2 deletions(-) diff --git a/packages/openai-agents/src/launchdarkly_ai_openai_agents/handler.py b/packages/openai-agents/src/launchdarkly_ai_openai_agents/handler.py index a6d0261f..7df7f3db 100644 --- a/packages/openai-agents/src/launchdarkly_ai_openai_agents/handler.py +++ b/packages/openai-agents/src/launchdarkly_ai_openai_agents/handler.py @@ -847,12 +847,14 @@ async def _stream_gen( raise except Exception as exc: hooks.close_open_spans(exc) + _write_failed_run_usage(span, config, exc, run_usage) + fail_span(span, exc, ended) while hooks.pending: yield hooks.pending.pop(0) + for rich in hooks.emitter.end_text(): + yield rich for rich in hooks.emitter.fail_open(): yield rich - _write_failed_run_usage(span, config, exc, run_usage) - fail_span(span, exc, ended) raise finally: # A no-op on the success and failure paths: both already ended their spans through diff --git a/packages/openai-agents/tests/test_handler.py b/packages/openai-agents/tests/test_handler.py index f2e7be88..22380ee9 100644 --- a/packages/openai-agents/tests/test_handler.py +++ b/packages/openai-agents/tests/test_handler.py @@ -1510,6 +1510,25 @@ async def test_fails_the_spans_when_the_stream_raises(self) -> None: assert StatusCode.ERROR in rec.root.statuses assert rec.root.ended == 1 + async def test_a_failed_turn_closes_the_open_text_message(self) -> None: + err = RuntimeError("next turn failed") + turns = [{"deltas": ["Let me "], "error": err}] + agents_mod = _fake_agents_module(run_streamed=_make_run_streamed(turns, "")) + seen: list[dict[str, Any]] = [] + with ( + _patched_agents(agents_mod), + pytest.raises(RuntimeError, match="next turn failed"), + ): + async for event in await create_openai_agent_handler().stream( + CONFIG, "q", {}, {} + ): + seen.append(event) + starts = [event for event in seen if event["type"] == "TEXT_MESSAGE_START"] + ends = [event for event in seen if event["type"] == "TEXT_MESSAGE_END"] + assert len(starts) == 1 + assert len(ends) == 1 + assert ends[0]["messageId"] == starts[0]["messageId"] + async def test_emits_no_content_by_default_on_the_streaming_path(self) -> None: turns = [{"deltas": ["hi"], "output": _text_output("hi")}] ctx, rec = _recording() From c8d198c3f6253a8c4eab45d32d459fb1614c502b Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 13 Oct 2026 18:51:47 +0000 Subject: [PATCH 08/15] 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 f43b88889372d97aca735294f4f6cc97e63f0220 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 13 Oct 2026 18:56:08 +0000 Subject: [PATCH 09/15] fix(openai): share the function-tool stream mapper Hosted tool events cannot arrive while only function tools are offered, so that mapping is gone. Both handlers report the string the model was sent, and argument deltas stay on the call they belong to. Co-authored-by: alexis --- .../src/launchdarkly_ai_server/responses.py | 57 ++++++++ .../launchdarkly_ai_openai_agents/handler.py | 133 +---------------- packages/openai-agents/tests/test_handler.py | 39 +---- .../handler.py | 134 +----------------- .../openai-messages/tests/test_handler.py | 111 +++++++-------- 5 files changed, 120 insertions(+), 354 deletions(-) create mode 100644 packages/client/src/launchdarkly_ai_server/responses.py diff --git a/packages/client/src/launchdarkly_ai_server/responses.py b/packages/client/src/launchdarkly_ai_server/responses.py new file mode 100644 index 00000000..ef2665e1 --- /dev/null +++ b/packages/client/src/launchdarkly_ai_server/responses.py @@ -0,0 +1,57 @@ +"""Map OpenAI Responses stream events onto the AG-UI text and tool subset. + +Both OpenAI handler packages call :func:`map_responses_event`. They only +offer function tools, so hosted and MCP item types are ignored until those +tools are wired. +""" + +from __future__ import annotations + +from typing import Any + +from .agui import TurnEmitter + + +def _call_id(event: Any, item_to_call: dict[str, str]) -> str: + direct = getattr(event, "call_id", None) + if isinstance(direct, str) and direct: + return direct + item_id = getattr(event, "item_id", None) + if isinstance(item_id, str) and item_id in item_to_call: + return item_to_call[item_id] + return item_id if isinstance(item_id, str) else "" + + +def map_responses_event( + event: Any, emitter: TurnEmitter, item_to_call: dict[str, str] +) -> list[dict[str, Any]]: + """Map one Responses event. Unknown events, including hosted tools, are ignored.""" + kind = getattr(event, "type", None) + if kind == "response.created": + return emitter.end_text() + if kind == "response.output_text.delta": + delta = getattr(event, "delta", "") + if not isinstance(delta, str) or not delta: + return [] + return emitter.text(delta) + if kind == "response.output_item.added": + item = getattr(event, "item", None) + if getattr(item, "type", None) != "function_call": + return [] + call_id = getattr(item, "call_id", None) or getattr(item, "id", None) or "" + name = getattr(item, "name", None) or "" + item_id = getattr(item, "id", None) + if isinstance(item_id, str) and isinstance(call_id, str) and call_id: + item_to_call[item_id] = call_id + return emitter.tool_start(str(call_id), str(name)) + if kind == "response.function_call_arguments.delta": + delta = getattr(event, "delta", "") + if not isinstance(delta, str): + return [] + return emitter.tool_args(_call_id(event, item_to_call), delta) + if kind == "response.function_call_arguments.done": + return emitter.tool_end( + _call_id(event, item_to_call), + arguments=getattr(event, "arguments", None), + ) + return [] diff --git a/packages/openai-agents/src/launchdarkly_ai_openai_agents/handler.py b/packages/openai-agents/src/launchdarkly_ai_openai_agents/handler.py index 7df7f3db..197d4609 100644 --- a/packages/openai-agents/src/launchdarkly_ai_openai_agents/handler.py +++ b/packages/openai-agents/src/launchdarkly_ai_openai_agents/handler.py @@ -46,6 +46,7 @@ text_message, ) from launchdarkly_ai_server.agui import TurnEmitter +from launchdarkly_ai_server.responses import map_responses_event from ._version import PACKAGE_NAME, __version__ from .spans import ( @@ -361,7 +362,7 @@ async def on_tool_end( arguments=getattr(context, "tool_arguments", None), ) ) - self.pending.extend(self.emitter.tool_result(str(call_id), result)) + self.pending.extend(self.emitter.tool_result(str(call_id), str(result))) span = self.open_tool_spans.pop(call_id, None) if span is None: return @@ -603,134 +604,6 @@ def _stream_impl( ) -_RESPONSES_ARG_DELTA = { - "response.function_call_arguments.delta": "delta", - "response.mcp_call_arguments.delta": "delta", - "response.custom_tool_call_input.delta": "delta", -} -_RESPONSES_ARG_DONE = { - "response.function_call_arguments.done": "arguments", - "response.mcp_call_arguments.done": "arguments", - "response.custom_tool_call_input.done": "input", -} -_RESPONSES_OPEN_TYPES = frozenset({"function_call", "mcp_call", "custom_tool_call"}) -_HOSTED_OUTPUT_FIELD = { - "mcp_call": "output", - "file_search_call": "results", - "code_interpreter_call": "outputs", -} -_HOSTED_TOOL_NAME = { - "file_search_call": "file_search", - "code_interpreter_call": "code_interpreter", -} - - -def _responses_call_id(event: Any, item_to_call: dict[str, str]) -> str: - direct = getattr(event, "call_id", None) - if isinstance(direct, str) and direct: - return direct - item_id = getattr(event, "item_id", None) - if isinstance(item_id, str) and item_id in item_to_call: - return item_to_call[item_id] - return item_id if isinstance(item_id, str) else "" - - -def _json_safe(value: Any) -> Any: - """Plain JSON for a hosted tool payload. - - File-search results and code-interpreter outputs are SDK model objects. - ``json.dumps`` rejects those, and a successful call would otherwise be - reported as a failed tool result. - """ - if value is None or isinstance(value, (str, int, float, bool)): - return value - if isinstance(value, dict): - return {str(key): _json_safe(item) for key, item in value.items()} - if isinstance(value, (list, tuple)): - return [_json_safe(item) for item in value] - dump = getattr(value, "model_dump", None) - if callable(dump): - try: - dumped = dump(mode="json") - except TypeError: - dumped = dump() - if dumped is not value: - return _json_safe(dumped) - return value - - -def _hosted_item_result( - item: Any, emitter: TurnEmitter, item_to_call: dict[str, str] -) -> list[dict[str, Any]]: - """Result for a hosted tool from ``response.output_item.done``. - - The ``*.completed`` events only carry ``item_id``. Reading ``output`` off - them yields ``None``, and with capture on that becomes the string ``null`` - and seals the call before the real item arrives. - """ - item_type = getattr(item, "type", None) - if item_type not in _HOSTED_OUTPUT_FIELD: - return [] - call_id = getattr(item, "call_id", None) or getattr(item, "id", None) or "" - item_id = getattr(item, "id", None) - if isinstance(item_id, str) and isinstance(call_id, str) and call_id: - item_to_call[item_id] = call_id - call_id = str(call_id) - name = getattr(item, "name", None) or _HOSTED_TOOL_NAME.get(item_type, "") - events = emitter.tool_start(call_id, str(name)) - events.extend(emitter.tool_end(call_id)) - status = getattr(item, "status", None) - if status in ("failed", "incomplete") or getattr(item, "error", None): - events.extend(emitter.tool_result(call_id, failed=True)) - else: - events.extend( - emitter.tool_result( - call_id, - _json_safe(getattr(item, _HOSTED_OUTPUT_FIELD[item_type], None)), - ) - ) - return events - - -def _map_responses_event( - event: Any, emitter: TurnEmitter, item_to_call: dict[str, str] -) -> list[dict[str, Any]]: - """Map one Responses stream event onto the AG-UI subset. Unknown events are ignored.""" - kind = getattr(event, "type", None) - if kind == "response.created": - return emitter.end_text() - if kind == "response.output_text.delta": - delta = getattr(event, "delta", "") - if not isinstance(delta, str) or not delta: - return [] - return emitter.text(delta) - if kind == "response.output_item.added": - item = getattr(event, "item", None) - if getattr(item, "type", None) not in _RESPONSES_OPEN_TYPES: - return [] - call_id = getattr(item, "call_id", None) or getattr(item, "id", None) or "" - name = getattr(item, "name", None) or "" - item_id = getattr(item, "id", None) - if isinstance(item_id, str) and isinstance(call_id, str) and call_id: - item_to_call[item_id] = call_id - return emitter.tool_start(str(call_id), str(name)) - if kind in _RESPONSES_ARG_DELTA: - delta = getattr(event, "delta", "") - if not isinstance(delta, str): - return [] - return emitter.tool_args(_responses_call_id(event, item_to_call), delta) - if kind in _RESPONSES_ARG_DONE: - arguments = getattr(event, _RESPONSES_ARG_DONE[kind], None) - return emitter.tool_end( - _responses_call_id(event, item_to_call), arguments=arguments - ) - if kind == "response.mcp_call.failed": - return emitter.tool_result(_responses_call_id(event, item_to_call), failed=True) - if kind == "response.output_item.done": - return _hosted_item_result(getattr(event, "item", None), emitter, item_to_call) - return [] - - def _stringify_output(value: Any) -> str: if isinstance(value, str): return value @@ -804,7 +677,7 @@ async def _stream_gen( if event.type == "raw_response_event": raw = getattr(event, "data", None) if raw is not None: - for rich in _map_responses_event( + for rich in map_responses_event( raw, hooks.emitter, hooks.item_to_call ): if rich.get("type") == "TEXT_MESSAGE_CONTENT": diff --git a/packages/openai-agents/tests/test_handler.py b/packages/openai-agents/tests/test_handler.py index 22380ee9..891196ea 100644 --- a/packages/openai-agents/tests/test_handler.py +++ b/packages/openai-agents/tests/test_handler.py @@ -1367,7 +1367,7 @@ async def _events(capture: bool) -> list[dict[str, Any]]: assert start["toolCallName"] == "search" 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") @@ -1424,12 +1424,12 @@ async def test_tool_end_publishes_arguments_the_consumer_has_not_read(self) -> N assert hooks.emitter.tool_args("c1", "more") == [] def test_a_new_response_closes_text_in_read_order(self) -> None: - from launchdarkly_ai_openai_agents.handler import _map_responses_event from launchdarkly_ai_server.agui import TurnEmitter + from launchdarkly_ai_server.responses import map_responses_event emitter = TurnEmitter(capture_content=True) opened = emitter.text("Let me ") - closed = _map_responses_event( + closed = map_responses_event( SimpleNamespace(type="response.created"), emitter, {} ) assert closed[0]["type"] == "TEXT_MESSAGE_END" @@ -2232,36 +2232,3 @@ async def _drain() -> None: assert rec.root.attributes.get("launchdarkly.run.cancelled") is True assert "launchdarkly.stream.abandoned" not in rec.root.attributes - - -class TestHostedToolResults: - def test_model_objects_are_json_not_failures(self) -> None: - from types import SimpleNamespace - - from launchdarkly_ai_openai_agents.handler import _map_responses_event - from launchdarkly_ai_server.agui import TurnEmitter - - class _Row: - def model_dump(self, **_kwargs: object) -> dict[str, str]: - return {"file_id": "f1", "text": "hit"} - - for item_type, call_id, field in ( - ("file_search_call", "fs1", "results"), - ("code_interpreter_call", "ci1", "outputs"), - ): - emitter = TurnEmitter(capture_content=True) - done = SimpleNamespace( - type="response.output_item.done", - item=SimpleNamespace( - type=item_type, - id=call_id, - status="completed", - **{field: [_Row()]}, - ), - ) - events = _map_responses_event(done, emitter, {}) - result = next( - event for event in events if event["type"] == "TOOL_CALL_RESULT" - ) - assert result["content"] == '[{"file_id":"f1","text":"hit"}]' - assert "metadata" not in result diff --git a/packages/openai-messages/src/launchdarkly_ai_openai_messages/handler.py b/packages/openai-messages/src/launchdarkly_ai_openai_messages/handler.py index 4e83a506..86c7b1fa 100644 --- a/packages/openai-messages/src/launchdarkly_ai_openai_messages/handler.py +++ b/packages/openai-messages/src/launchdarkly_ai_openai_messages/handler.py @@ -28,6 +28,7 @@ set_tool_call_content_attributes, ) from launchdarkly_ai_server.agui import TurnEmitter +from launchdarkly_ai_server.responses import map_responses_event from ._version import PACKAGE_NAME, __version__ from .spans import ( @@ -434,134 +435,6 @@ def _final_output_messages(output: str) -> list[SpanMessage]: ] -_RESPONSES_ARG_DELTA = { - "response.function_call_arguments.delta": "delta", - "response.mcp_call_arguments.delta": "delta", - "response.custom_tool_call_input.delta": "delta", -} -_RESPONSES_ARG_DONE = { - "response.function_call_arguments.done": "arguments", - "response.mcp_call_arguments.done": "arguments", - "response.custom_tool_call_input.done": "input", -} -_RESPONSES_OPEN_TYPES = frozenset({"function_call", "mcp_call", "custom_tool_call"}) -_HOSTED_OUTPUT_FIELD = { - "mcp_call": "output", - "file_search_call": "results", - "code_interpreter_call": "outputs", -} -_HOSTED_TOOL_NAME = { - "file_search_call": "file_search", - "code_interpreter_call": "code_interpreter", -} - - -def _responses_call_id(event: Any, item_to_call: dict[str, str]) -> str: - direct = getattr(event, "call_id", None) - if isinstance(direct, str) and direct: - return direct - item_id = getattr(event, "item_id", None) - if isinstance(item_id, str) and item_id in item_to_call: - return item_to_call[item_id] - return item_id if isinstance(item_id, str) else "" - - -def _json_safe(value: Any) -> Any: - """Plain JSON for a hosted tool payload. - - File-search results and code-interpreter outputs are SDK model objects. - ``json.dumps`` rejects those, and a successful call would otherwise be - reported as a failed tool result. - """ - if value is None or isinstance(value, (str, int, float, bool)): - return value - if isinstance(value, dict): - return {str(key): _json_safe(item) for key, item in value.items()} - if isinstance(value, (list, tuple)): - return [_json_safe(item) for item in value] - dump = getattr(value, "model_dump", None) - if callable(dump): - try: - dumped = dump(mode="json") - except TypeError: - dumped = dump() - if dumped is not value: - return _json_safe(dumped) - return value - - -def _hosted_item_result( - item: Any, emitter: TurnEmitter, item_to_call: dict[str, str] -) -> list[dict[str, Any]]: - """Result for a hosted tool from ``response.output_item.done``. - - The ``*.completed`` events only carry ``item_id``. Reading ``output`` off - them yields ``None``, and with capture on that becomes the string ``null`` - and seals the call before the real item arrives. - """ - item_type = getattr(item, "type", None) - if item_type not in _HOSTED_OUTPUT_FIELD: - return [] - call_id = getattr(item, "call_id", None) or getattr(item, "id", None) or "" - item_id = getattr(item, "id", None) - if isinstance(item_id, str) and isinstance(call_id, str) and call_id: - item_to_call[item_id] = call_id - call_id = str(call_id) - name = getattr(item, "name", None) or _HOSTED_TOOL_NAME.get(item_type, "") - events = emitter.tool_start(call_id, str(name)) - events.extend(emitter.tool_end(call_id)) - status = getattr(item, "status", None) - if status in ("failed", "incomplete") or getattr(item, "error", None): - events.extend(emitter.tool_result(call_id, failed=True)) - else: - events.extend( - emitter.tool_result( - call_id, - _json_safe(getattr(item, _HOSTED_OUTPUT_FIELD[item_type], None)), - ) - ) - return events - - -def _map_responses_event( - event: Any, emitter: TurnEmitter, item_to_call: dict[str, str] -) -> list[dict[str, Any]]: - """Map one Responses stream event onto the AG-UI subset. Unknown events are ignored.""" - kind = getattr(event, "type", None) - if kind == "response.created": - return emitter.end_text() - if kind == "response.output_text.delta": - delta = getattr(event, "delta", "") - if not isinstance(delta, str) or not delta: - return [] - return emitter.text(delta) - if kind == "response.output_item.added": - item = getattr(event, "item", None) - if getattr(item, "type", None) not in _RESPONSES_OPEN_TYPES: - return [] - call_id = getattr(item, "call_id", None) or getattr(item, "id", None) or "" - name = getattr(item, "name", None) or "" - item_id = getattr(item, "id", None) - if isinstance(item_id, str) and isinstance(call_id, str) and call_id: - item_to_call[item_id] = call_id - return emitter.tool_start(str(call_id), str(name)) - if kind in _RESPONSES_ARG_DELTA: - delta = getattr(event, "delta", "") - if not isinstance(delta, str): - return [] - return emitter.tool_args(_responses_call_id(event, item_to_call), delta) - if kind in _RESPONSES_ARG_DONE: - arguments = getattr(event, _RESPONSES_ARG_DONE[kind], None) - return emitter.tool_end( - _responses_call_id(event, item_to_call), arguments=arguments - ) - if kind == "response.mcp_call.failed": - return emitter.tool_result(_responses_call_id(event, item_to_call), failed=True) - if kind == "response.output_item.done": - return _hosted_item_result(getattr(event, "item", None), emitter, item_to_call) - return [] - - async def _stream_gen( client: Any, config: AiConfigRep, @@ -646,7 +519,7 @@ async def _stream_gen( stream = client.responses.stream(**stream_params) async with stream as s: async for event in s: - for rich in _map_responses_event(event, emitter, item_to_call): + for rich in map_responses_event(event, emitter, item_to_call): if rich.get("type") == "TEXT_MESSAGE_CONTENT": full_output += rich.get("delta") or "" yield rich @@ -731,7 +604,8 @@ async def _stream_gen( # Cleared on both paths that end the span, and deliberately not in a `finally`: # a `finally` would also clear it for a BaseException, which is the one case where # the span is still open and the outer `finally` is the only thing left to close it. - for rich in emitter.tool_result(str(tc.call_id), result): + # The model is sent str(result). The stream reports that same text. + for rich in emitter.tool_result(str(tc.call_id), str(result)): yield rich tool_outputs.append( { diff --git a/packages/openai-messages/tests/test_handler.py b/packages/openai-messages/tests/test_handler.py index c22e36a5..69887795 100644 --- a/packages/openai-messages/tests/test_handler.py +++ b/packages/openai-messages/tests/test_handler.py @@ -1321,7 +1321,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") @@ -1361,76 +1361,71 @@ def boom(_args: dict[str, Any]) -> str: } assert results == {"c1", "c2"} - def test_hosted_completion_waits_for_the_output_item(self) -> None: + def test_parallel_argument_deltas_stay_on_their_call(self) -> None: from types import SimpleNamespace - from launchdarkly_ai_openai_messages.handler import _map_responses_event from launchdarkly_ai_server.agui import TurnEmitter + from launchdarkly_ai_server.responses import map_responses_event emitter = TurnEmitter(capture_content=True) ids: dict[str, str] = {} - added = SimpleNamespace( - type="response.output_item.added", - item=SimpleNamespace( - type="mcp_call", call_id="m1", id="item-1", name="search" - ), - ) - completed = SimpleNamespace( - type="response.mcp_call.completed", item_id="item-1", output=None - ) - done = SimpleNamespace( - type="response.output_item.done", - item=SimpleNamespace( - type="mcp_call", - call_id="m1", - id="item-1", - name="search", - status="completed", - output="found", - ), - ) events: list[dict[str, Any]] = [] - events.extend(_map_responses_event(added, emitter, ids)) - events.extend(_map_responses_event(completed, emitter, ids)) - events.extend(_map_responses_event(done, emitter, ids)) - assert "null" not in str(events) - result = next(event for event in events if event["type"] == "TOOL_CALL_RESULT") - assert result["content"] == "found" - assert [event["type"] for event in events].count("TOOL_CALL_RESULT") == 1 - - def test_hosted_model_results_are_json_not_failures(self) -> None: - from types import SimpleNamespace - from launchdarkly_ai_openai_messages.handler import _map_responses_event - from launchdarkly_ai_server.agui import TurnEmitter + def feed(event: Any) -> None: + events.extend(map_responses_event(event, emitter, ids)) - class _Row: - def model_dump(self, **_kwargs: Any) -> dict[str, str]: - return {"file_id": "f1", "text": "hit"} - - def _done(item_type: str, call_id: str, field: str) -> Any: - return SimpleNamespace( - type="response.output_item.done", + feed( + SimpleNamespace( + type="response.output_item.added", item=SimpleNamespace( - type=item_type, - id=call_id, - status="completed", - **{field: [_Row()]}, + type="function_call", id="fc_1", call_id="call_1", name="search" ), ) - - for item_type, call_id, field in ( - ("file_search_call", "fs1", "results"), - ("code_interpreter_call", "ci1", "outputs"), - ): - emitter = TurnEmitter(capture_content=True) - events = _map_responses_event(_done(item_type, call_id, field), emitter, {}) - result = next( - event for event in events if event["type"] == "TOOL_CALL_RESULT" + ) + feed( + SimpleNamespace( + type="response.output_item.added", + item=SimpleNamespace( + type="function_call", id="fc_2", call_id="call_2", name="other" + ), + ) + ) + feed( + SimpleNamespace( + type="response.function_call_arguments.delta", + item_id="fc_1", + delta='{"q":', + ) + ) + feed( + SimpleNamespace( + type="response.function_call_arguments.delta", + item_id="fc_2", + delta='{"n":', + ) + ) + feed( + SimpleNamespace( + type="response.function_call_arguments.delta", + item_id="fc_1", + delta='"x"}', + ) + ) + feed( + SimpleNamespace( + type="response.function_call_arguments.delta", + item_id="fc_2", + delta="1}", + ) + ) + by_id: dict[str, str] = {} + for event in events: + if event["type"] != "TOOL_CALL_ARGS": + continue + by_id[event["toolCallId"]] = ( + by_id.get(event["toolCallId"], "") + event["delta"] ) - assert result["toolCallId"] == call_id - assert result["content"] == '[{"file_id":"f1","text":"hit"}]' - assert "metadata" not in result + assert by_id == {"call_1": '{"q":"x"}', "call_2": '{"n":1}'} async def test_an_opened_call_without_a_function_result_is_closed( self, mock_openai: MagicMock From 3248d3cd64341e5eb281837a3dfc74aa182850d7 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 13 Oct 2026 19:12:23 +0000 Subject: [PATCH 10/15] fix(openai): close text before a tool result and on failure A tool result was emitted while the assistant message was still open, so an AG-UI client could attach the result to the wrong turn. The agents hook now ends that message before the tool result. A provider error in the messages handler ended the span and then raised without TEXT_MESSAGE_END; the failure path now closes the message after the span is failed. Co-authored-by: alexis --- .../launchdarkly_ai_openai_agents/handler.py | 1 + packages/openai-agents/tests/test_handler.py | 20 ++++++++++ .../handler.py | 2 + .../openai-messages/tests/test_handler.py | 38 +++++++++++++++++++ 4 files changed, 61 insertions(+) diff --git a/packages/openai-agents/src/launchdarkly_ai_openai_agents/handler.py b/packages/openai-agents/src/launchdarkly_ai_openai_agents/handler.py index 197d4609..fc6bcd25 100644 --- a/packages/openai-agents/src/launchdarkly_ai_openai_agents/handler.py +++ b/packages/openai-agents/src/launchdarkly_ai_openai_agents/handler.py @@ -356,6 +356,7 @@ async def on_tool_end( call_id = self._tool_span_key(context, tool) # The hook can run before unread argument deltas. Close with the full # argument string first, so tool_result does not seal a prefix. + self.pending.extend(self.emitter.end_text()) self.pending.extend( self.emitter.tool_end( str(call_id), diff --git a/packages/openai-agents/tests/test_handler.py b/packages/openai-agents/tests/test_handler.py index 891196ea..3b88252a 100644 --- a/packages/openai-agents/tests/test_handler.py +++ b/packages/openai-agents/tests/test_handler.py @@ -1423,6 +1423,26 @@ async def test_tool_end_publishes_arguments_the_consumer_has_not_read(self) -> N assert any(event["type"] == "TOOL_CALL_RESULT" for event in hooks.pending) assert hooks.emitter.tool_args("c1", "more") == [] + async def test_a_tool_result_closes_the_open_text_first(self) -> None: + from launchdarkly_ai_openai_agents.handler import _SpanningHooks + from launchdarkly_ai_server import create_run_usage + + hooks = _SpanningHooks(CONFIG, None, True, create_run_usage()) + hooks.emitter.text("Looking") + ctx = SimpleNamespace( + tool_name="search", + tool_call_id="c1", + tool_arguments="{}", + ) + await hooks.on_tool_start( + ctx, SimpleNamespace(), SimpleNamespace(name="search") + ) + await hooks.on_tool_end( + ctx, SimpleNamespace(), SimpleNamespace(name="search"), "found" + ) + types = [event["type"] for event in hooks.pending] + assert types.index("TEXT_MESSAGE_END") < types.index("TOOL_CALL_RESULT") + def test_a_new_response_closes_text_in_read_order(self) -> None: from launchdarkly_ai_server.agui import TurnEmitter from launchdarkly_ai_server.responses import map_responses_event diff --git a/packages/openai-messages/src/launchdarkly_ai_openai_messages/handler.py b/packages/openai-messages/src/launchdarkly_ai_openai_messages/handler.py index 86c7b1fa..11237644 100644 --- a/packages/openai-messages/src/launchdarkly_ai_openai_messages/handler.py +++ b/packages/openai-messages/src/launchdarkly_ai_openai_messages/handler.py @@ -641,6 +641,8 @@ async def _stream_gen( if run_usage.reported: finish_root_span(span, last_response_model, run_usage.total) fail_span(span, exc, ended) + for rich in emitter.end_text(): + yield rich for rich in emitter.fail_open(): yield rich raise diff --git a/packages/openai-messages/tests/test_handler.py b/packages/openai-messages/tests/test_handler.py index 69887795..6409975f 100644 --- a/packages/openai-messages/tests/test_handler.py +++ b/packages/openai-messages/tests/test_handler.py @@ -1517,6 +1517,44 @@ async def _mgr() -> AsyncGenerator[Any, None]: for event in seen ) + async def test_a_provider_error_closes_the_open_text( + self, mock_openai: MagicMock + ) -> None: + from types import SimpleNamespace + + import launchdarkly_ai_openai_messages.spans as spans_mod + from launchdarkly_ai_openai_messages import create_openai_messages_handler + + delta = SimpleNamespace(type="response.output_text.delta", delta="Hello") + + class _Stream: + def __aiter__(self) -> AsyncIterator[Any]: + return self._iter() + + async def _iter(self) -> AsyncIterator[Any]: + yield delta + raise RuntimeError("stream broke") + + async def get_final_response(self) -> Any: + raise AssertionError("final response is not read") + + @asynccontextmanager + async def _mgr() -> AsyncGenerator[Any, None]: + yield _Stream() + + mock_openai.responses.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_openai_messages_handler() + async for event in await handler.stream(CONFIG, "q"): + seen.append(event) + start = next(event for event in seen if event["type"] == "TEXT_MESSAGE_START") + end = next(event for event in seen if event["type"] == "TEXT_MESSAGE_END") + assert end["messageId"] == start["messageId"] + async def test_yields_exactly_one_done_event(self, mock_openai: MagicMock) -> None: import launchdarkly_ai_openai_messages.spans as spans_mod from launchdarkly_ai_openai_messages import create_openai_messages_handler From 3807bbed2d7d5d868c5b9ec75c8e44f141250601 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 13 Oct 2026 19:24:41 +0000 Subject: [PATCH 11/15] fix(openai): apply unread text before a tool result The agents tool hook runs before the consumer has read every text delta from that model call. Closing the message in the hook split the rest of the turn into a second message and put the tool result ahead of it. The stream now keeps that text on the open message, then closes it before the tool result. Co-authored-by: alexis --- .../launchdarkly_ai_openai_agents/handler.py | 58 +++++++++++++----- packages/openai-agents/tests/test_handler.py | 61 +++++++++++++++++-- 2 files changed, 99 insertions(+), 20 deletions(-) diff --git a/packages/openai-agents/src/launchdarkly_ai_openai_agents/handler.py b/packages/openai-agents/src/launchdarkly_ai_openai_agents/handler.py index fc6bcd25..ca0d1ffa 100644 --- a/packages/openai-agents/src/launchdarkly_ai_openai_agents/handler.py +++ b/packages/openai-agents/src/launchdarkly_ai_openai_agents/handler.py @@ -356,7 +356,8 @@ async def on_tool_end( call_id = self._tool_span_key(context, tool) # The hook can run before unread argument deltas. Close with the full # argument string first, so tool_result does not seal a prefix. - self.pending.extend(self.emitter.end_text()) + # Text stays open: this runs before unread text deltas, and closing + # here would start a second message for the rest of the turn. self.pending.extend( self.emitter.tool_end( str(call_id), @@ -411,6 +412,26 @@ def abandon_open_spans(self, ended: set[int], cancelled: bool = False) -> None: self.open_model_span = None +def _take_finished_tools(hooks: _SpanningHooks) -> list[dict[str, Any]]: + """Return tool events recorded by the hook, text closed before each result. + + The hook runs before unread deltas from the same model call. The caller + applies those deltas first. Tool results then follow ``TEXT_MESSAGE_END`` + so the client does not attach them to the next message. + """ + if not hooks.pending: + return [] + earlier: list[dict[str, Any]] = [] + results: list[dict[str, Any]] = [] + for event in hooks.pending: + if event.get("type") == "TOOL_CALL_RESULT": + results.append(event) + else: + earlier.append(event) + hooks.pending.clear() + return [*earlier, *hooks.emitter.end_text(), *results] + + def _write_failed_run_usage( span: Any, config: AiConfigRep, @@ -673,19 +694,24 @@ async def _stream_gen( full_output = "" async for event in streamed.stream_events(): - while hooks.pending: - yield hooks.pending.pop(0) - if event.type == "raw_response_event": - raw = getattr(event, "data", None) - if raw is not None: - for rich in map_responses_event( - raw, hooks.emitter, hooks.item_to_call - ): - if rich.get("type") == "TEXT_MESSAGE_CONTENT": - full_output += rich.get("delta") or "" - yield rich - while hooks.pending: - yield hooks.pending.pop(0) + raw = ( + getattr(event, "data", None) + if event.type == "raw_response_event" + else None + ) + # The next model call is queued only after this turn's tools have + # run. Unread deltas from this turn are still ahead of that event, + # so the tool result waits until they have been applied. + if getattr(raw, "type", None) == "response.created": + for rich in _take_finished_tools(hooks): + yield rich + if raw is not None: + for rich in map_responses_event(raw, hooks.emitter, hooks.item_to_call): + if rich.get("type") == "TEXT_MESSAGE_CONTENT": + full_output += rich.get("delta") or "" + yield rich + for rich in _take_finished_tools(hooks): + yield rich for rich in hooks.emitter.end_text(): yield rich for rich in hooks.emitter.fail_open(): @@ -723,8 +749,8 @@ async def _stream_gen( hooks.close_open_spans(exc) _write_failed_run_usage(span, config, exc, run_usage) fail_span(span, exc, ended) - while hooks.pending: - yield hooks.pending.pop(0) + for rich in _take_finished_tools(hooks): + yield rich for rich in hooks.emitter.end_text(): yield rich for rich in hooks.emitter.fail_open(): diff --git a/packages/openai-agents/tests/test_handler.py b/packages/openai-agents/tests/test_handler.py index 3b88252a..b7a0c62a 100644 --- a/packages/openai-agents/tests/test_handler.py +++ b/packages/openai-agents/tests/test_handler.py @@ -117,6 +117,10 @@ async def stream_events(self) -> AsyncIterator[Any]: await self.hooks.on_llm_start( MagicMock(), self.agent, self.agent.instructions, self.prompt ) + yield SimpleNamespace( + type="raw_response_event", + data=SimpleNamespace(type="response.created"), + ) for delta in turn.get("deltas", []): yield SimpleNamespace( type="raw_response_event", @@ -135,6 +139,13 @@ async def stream_events(self) -> AsyncIterator[Any]: await self.hooks.on_tool_end( ctx, self.agent, tool_obj, call.get("result", "ok") ) + for delta in turn.get("late_deltas", []): + yield SimpleNamespace( + type="raw_response_event", + data=SimpleNamespace( + type="response.output_text.delta", delta=delta + ), + ) if turn.get("error") is not None: raise turn["error"] response = SimpleNamespace( @@ -1376,6 +1387,39 @@ async def _events(capture: bool) -> list[dict[str, Any]]: event for event in shown if event["type"] == "TOOL_CALL_RESULT" ) assert shown_result["content"] == "secret" + types = [event["type"] for event in shown] + assert types.index("TEXT_MESSAGE_CONTENT") < types.index("TEXT_MESSAGE_END") + assert types.index("TEXT_MESSAGE_END") < types.index("TOOL_CALL_RESULT") + + async def test_unread_text_stays_on_the_message_the_tool_result_closes( + self, + ) -> None: + turns = [ + { + "deltas": ["Hel"], + "tool_calls": [ + {"name": "search", "id": "c1", "args": {}, "result": "found"} + ], + "late_deltas": ["lo"], + "output": [], + } + ] + agents_mod = _fake_agents_module(run_streamed=_make_run_streamed(turns, "done")) + with _patched_agents(agents_mod): + events = [ + event + async for event in await create_openai_agent_handler( + capture_content=True + ).stream(CONFIG, "q", {}, {}) + ] + contents = [ + event for event in events if event["type"] == "TEXT_MESSAGE_CONTENT" + ] + assert [event["delta"] for event in contents] == ["Hel", "lo"] + assert contents[0]["messageId"] == contents[1]["messageId"] + types = [event["type"] for event in events] + assert types.index("TEXT_MESSAGE_END") < types.index("TOOL_CALL_RESULT") + assert types.index("TOOL_CALL_START") < types.index("TEXT_MESSAGE_END") async def test_a_streaming_prefix_is_not_closed_by_the_tool_hook(self) -> None: from launchdarkly_ai_openai_agents.handler import _SpanningHooks @@ -1423,12 +1467,15 @@ async def test_tool_end_publishes_arguments_the_consumer_has_not_read(self) -> N assert any(event["type"] == "TOOL_CALL_RESULT" for event in hooks.pending) assert hooks.emitter.tool_args("c1", "more") == [] - async def test_a_tool_result_closes_the_open_text_first(self) -> None: - from launchdarkly_ai_openai_agents.handler import _SpanningHooks + async def test_a_tool_hook_leaves_text_open_until_the_deltas_are_read(self) -> None: + from launchdarkly_ai_openai_agents.handler import ( + _SpanningHooks, + _take_finished_tools, + ) from launchdarkly_ai_server import create_run_usage hooks = _SpanningHooks(CONFIG, None, True, create_run_usage()) - hooks.emitter.text("Looking") + opened = hooks.emitter.text("Hel") ctx = SimpleNamespace( tool_name="search", tool_call_id="c1", @@ -1440,7 +1487,13 @@ async def test_a_tool_result_closes_the_open_text_first(self) -> None: await hooks.on_tool_end( ctx, SimpleNamespace(), SimpleNamespace(name="search"), "found" ) - types = [event["type"] for event in hooks.pending] + assert not any(event["type"] == "TEXT_MESSAGE_END" for event in hooks.pending) + more = hooks.emitter.text("lo") + assert more[0]["type"] == "TEXT_MESSAGE_CONTENT" + assert more[0]["messageId"] == opened[0]["messageId"] + released = _take_finished_tools(hooks) + types = [event["type"] for event in released] + assert types.index("TOOL_CALL_START") < types.index("TEXT_MESSAGE_END") assert types.index("TEXT_MESSAGE_END") < types.index("TOOL_CALL_RESULT") def test_a_new_response_closes_text_in_read_order(self) -> None: From fa05f747d57da3698a2a6f8b1d4093856ad17b2b Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 13 Oct 2026 14:18:55 +0000 Subject: [PATCH 12/15] 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 1a7f6819..75c8008a 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 b104481b6436016124c894898750ac1ef36cc3ee Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 13 Oct 2026 14:21:05 +0000 Subject: [PATCH 13/15] 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 8294f5938972294e2c7ac79a6c60e2f1153893fe Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 13 Oct 2026 18:51:47 +0000 Subject: [PATCH 14/15] 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 37c738785139281c9b5523c9cb47442760043573 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 14 Oct 2026 18:14:20 +0000 Subject: [PATCH 15/15] fix(client): close only streams that implement aclose contextlib.aclosing always awaits aclose. A handler may return a plain async iterator, and that call raised AttributeError after the loop had already succeeded. The tracker now closes the stream only when aclose is present. Co-authored-by: alexis --- .../src/launchdarkly_ai_server/tracking.py | 22 ++++++- packages/client/tests/test_client.py | 59 +++++++++++++++++++ 2 files changed, 78 insertions(+), 3 deletions(-) diff --git a/packages/client/src/launchdarkly_ai_server/tracking.py b/packages/client/src/launchdarkly_ai_server/tracking.py index 7dca2d57..acc4f8b5 100644 --- a/packages/client/src/launchdarkly_ai_server/tracking.py +++ b/packages/client/src/launchdarkly_ai_server/tracking.py @@ -4,8 +4,8 @@ import os import time import uuid -from collections.abc import AsyncGenerator, Callable -from contextlib import aclosing +from collections.abc import AsyncGenerator, AsyncIterator, Callable +from contextlib import asynccontextmanager from typing import Any from .agui import HANDLER_RICH_TYPES, TEXT_MESSAGE_CONTENT, TEXT_TYPES @@ -23,6 +23,22 @@ from .utils import model_stamps_from_meta, parse_usage, to_ld_context +@asynccontextmanager +async def _closing_stream(stream: Any) -> AsyncIterator[Any]: + """Close a handler stream when it implements ``aclose``. + + ``contextlib.aclosing`` always awaits ``aclose``. A handler may return a + plain async iterator, which has no such method, and that call then raises + ``AttributeError`` after the loop has already succeeded. + """ + try: + yield stream + finally: + close = getattr(stream, "aclose", None) + if close is not None: + await close() + + def _try_get_environment_id() -> str | None: """ Returns the LaunchDarkly environment MongoDB ObjectId used to set @@ -308,7 +324,7 @@ async def execute_and_stream( stream_gen = await handler.stream( config, user_input, tracked_tool_handlers, merged_variables, history ) - async with aclosing(stream_gen) as agen: + async with _closing_stream(stream_gen) as agen: async for event in agen: typ = event.get("type") if typ in TEXT_TYPES: diff --git a/packages/client/tests/test_client.py b/packages/client/tests/test_client.py index 1ea9f66a..07ccb3a4 100644 --- a/packages/client/tests/test_client.py +++ b/packages/client/tests/test_client.py @@ -786,6 +786,65 @@ async def _noop_fn( events_tracked = [c[0][0] for c in mock_ld_client.track.call_args_list] assert "$ld:ai:generation:error" in events_tracked + async def test_a_plain_async_iterator_exits_cleanly( + self, mock_ld_client: MagicMock + ) -> None: + class _Iterator: + def __init__(self) -> None: + self._events = [ + {"type": "chunk", "text": "hi"}, + { + "type": "done", + "output": "hi", + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + ] + + def __aiter__(self) -> Any: + return self + + async def __anext__(self) -> dict[str, Any]: + if not self._events: + raise StopAsyncIteration + event = self._events.pop(0) + return event + + def _stream_fn( + cfg: Any, + user_input: Any, + tool_handlers: Any, + variables: Any, + history: Any = None, + ) -> _Iterator: + iterator = _Iterator() + assert not hasattr(iterator, "aclose") + return iterator + + async def _noop_fn( + cfg: Any, + user_input: Any, + tool_handlers: Any, + variables: Any, + history: Any = None, + ) -> dict[str, Any]: + return {"output": "", "usage": {}} + + handler = ProviderHandler( + fn=_noop_fn, + stream_fn=_stream_fn, + provides_for=("TestProvider", "messages"), + ) + events = [ + event + async for event in config(key="flag", handler=handler).stream("q", CONTEXT) + ] + chunks = [event for event in events if event["type"] == "chunk"] + assert chunks == [{"type": "chunk", "text": "hi"}] + assert events[-1]["type"] == "done" + assert events[-1]["response"] == "hi" + tracked = [call[0][0] for call in mock_ld_client.track.call_args_list] + assert "$ld:ai:generation:success" in tracked + async def test_ld_context_injected_into_stream_variables( self, mock_ld_client: MagicMock ) -> None: