From fa05f747d57da3698a2a6f8b1d4093856ad17b2b Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 13 Oct 2026 14:18:55 +0000 Subject: [PATCH 1/4] 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 2/4] 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 3/4] 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 4/4] 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: