From fb516505c432d2b7f6ec00403c64a2a0d29ed868 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 13 Oct 2026 14:18:55 +0000 Subject: [PATCH 01/14] 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/14] 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 2312bf2117eab8c7297d9ef4f050731eb2dd3a8d Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 13 Oct 2026 14:34:08 +0000 Subject: [PATCH 03/14] feat(langchain): emit AG-UI text and tool events Messages and agents map model chunks and the tool loop onto the shared subset. Argument text and results follow capture_content. Co-authored-by: alexis --- .../handler.py | 30 +++++++++++- .../launchdarkly_ai_langchain_agents/spans.py | 25 ++++++++++ .../langchain-agents/tests/test_handler.py | 9 +++- .../handler.py | 47 ++++++++++++++++++- .../langchain-messages/tests/test_handler.py | 17 ++++--- 5 files changed, 118 insertions(+), 10 deletions(-) diff --git a/packages/langchain-agents/src/launchdarkly_ai_langchain_agents/handler.py b/packages/langchain-agents/src/launchdarkly_ai_langchain_agents/handler.py index fde27d23..3004d260 100644 --- a/packages/langchain-agents/src/launchdarkly_ai_langchain_agents/handler.py +++ b/packages/langchain-agents/src/launchdarkly_ai_langchain_agents/handler.py @@ -31,10 +31,12 @@ set_input_content_attributes, set_output_content_attributes, ) +from launchdarkly_ai_server.agui import TurnEmitter from ._version import PACKAGE_NAME, __version__ from .messages import to_lang_chain_messages from .spans import ( + _tool_result_text, build_span_callbacks, fail_span, finish_root_span, @@ -492,6 +494,7 @@ async def _stream_gen( run_usage = create_run_usage() full_output = "" + emitter = TurnEmitter(capture_content=capture_content) # agent.astream() yields state updates per graph step: { [node_name]: { messages: [...] } } async for step_state in agent.astream( @@ -513,8 +516,33 @@ async def _stream_gen( if getattr(msg, "type", None) == "ai": text = lang_chain_content_text(msg.content) if text: - yield {"type": "chunk", "text": text} + for rich in emitter.text(text): + yield rich full_output = text + for tc in getattr(msg, "tool_calls", None) or []: + if not isinstance(tc, dict): + continue + call_id = str(tc.get("id") or tc.get("name") or "") + name = str(tc.get("name") or "") + for rich in emitter.tool_start(call_id, name): + yield rich + for rich in emitter.tool_end( + call_id, arguments=tc.get("args") + ): + yield rich + for call_id, name, arguments, output, failed in list( + span_callbacks.finished_calls + ): + for rich in emitter.tool_start(call_id, name): + yield rich + for rich in emitter.tool_end(call_id, arguments=arguments): + yield rich + value = _tool_result_text(output) if not failed else None + for rich in emitter.tool_result(call_id, value, failed=failed): + yield rich + span_callbacks.finished_calls.clear() + for rich in emitter.end_text(): + yield rich # The same reconciliation the blocking path does, for the same reason. This walk reads # usage_metadata off the astream payloads, and sees nothing at all unless the graph is diff --git a/packages/langchain-agents/src/launchdarkly_ai_langchain_agents/spans.py b/packages/langchain-agents/src/launchdarkly_ai_langchain_agents/spans.py index 3621de06..a7b43181 100644 --- a/packages/langchain-agents/src/launchdarkly_ai_langchain_agents/spans.py +++ b/packages/langchain-agents/src/launchdarkly_ai_langchain_agents/spans.py @@ -304,6 +304,8 @@ def __init__( self.run_usage = run_usage self.model_spans: dict[str, Any] = {} self.tool_spans: dict[str, Any] = {} + self._open_calls: dict[str, tuple[str, str, Any]] = {} + self.finished_calls: list[tuple[str, str, Any, Any, bool]] = [] def _start_model(self, run_id: Any, messages: Any = None) -> None: key = str(run_id) @@ -404,12 +406,23 @@ async def on_tool_start( # Tracked before the content write, not after. Serialising the arguments can raise, and a # span created but never inserted is unreachable by every cleanup path there is. self.tool_spans[str(run_id)] = span + call_id = str(tool_call_id or run_id) + self._open_calls[str(run_id)] = ( + call_id, + str(tool_name), + kwargs.get("inputs") or input_str, + ) set_tool_call_content_attributes( span, self._capture_content, arguments=kwargs.get("inputs") or input_str ) async def on_tool_end(self, output: Any, *, run_id: Any, **kwargs: Any) -> None: key = str(run_id) + recorded = self._open_calls.pop(key, None) + if recorded is not None: + self.finished_calls.append( + (recorded[0], recorded[1], recorded[2], output, False) + ) span = self.tool_spans.pop(key, None) if span is None: return @@ -427,6 +440,11 @@ async def on_tool_error( self, error: BaseException, *, run_id: Any, **kwargs: Any ) -> None: key = str(run_id) + recorded = self._open_calls.pop(key, None) + if recorded is not None: + self.finished_calls.append( + (recorded[0], recorded[1], recorded[2], None, True) + ) span = self.tool_spans.pop(key, None) if span is None: return @@ -492,6 +510,13 @@ def __init__( self.run_usage = run_usage self._handler = handler + @property + def finished_calls(self) -> list[tuple[str, str, Any, Any, bool]]: + if self._handler is None: + return [] + calls: list[tuple[str, str, Any, Any, bool]] = self._handler.finished_calls + return calls + def close_open_spans(self, error: BaseException) -> None: if self._handler is not None: self._handler.close_open_spans(error) diff --git a/packages/langchain-agents/tests/test_handler.py b/packages/langchain-agents/tests/test_handler.py index 4571a79b..238c125a 100644 --- a/packages/langchain-agents/tests/test_handler.py +++ b/packages/langchain-agents/tests/test_handler.py @@ -1427,7 +1427,11 @@ async def _mock_astream(*a: Any, **kw: Any) -> AsyncIterator[Any]: llm=MagicMock() ).stream(_make_config(), "q") ] - assert {"type": "chunk", "text": "visible answer"} in events + assert any( + event.get("type") == "TEXT_MESSAGE_CONTENT" + and event.get("delta") == "visible answer" + for event in events + ) assert events[-1]["output"] == "visible answer" @pytest.mark.asyncio @@ -2564,5 +2568,6 @@ def factory(config: Any) -> Any: ] assert seen[0]["model"]["parameters"] == {"temperature": 0.2} assert any( - e.get("type") == "chunk" and e.get("text") == "streamed" for e in events + e.get("type") == "TEXT_MESSAGE_CONTENT" and e.get("delta") == "streamed" + for e in events ) diff --git a/packages/langchain-messages/src/launchdarkly_ai_langchain_messages/handler.py b/packages/langchain-messages/src/launchdarkly_ai_langchain_messages/handler.py index d41876ba..d982004f 100644 --- a/packages/langchain-messages/src/launchdarkly_ai_langchain_messages/handler.py +++ b/packages/langchain-messages/src/launchdarkly_ai_langchain_messages/handler.py @@ -32,6 +32,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 ( @@ -755,6 +756,7 @@ async def _stream_gen( conversation_messages = list(initial_messages) full_output: Any = "" steps = 0 + emitter = TurnEmitter(capture_content=capture_content) ended: set[int] = set() open_model_span: Any = None @@ -824,8 +826,33 @@ async def _stream_gen( async for chunk in chunk_stream: text = lang_chain_content_text(chunk.content) if text: - yield {"type": "chunk", "text": text} + for rich in emitter.text(text): + yield rich accumulated_content += text + for piece in getattr(chunk, "tool_call_chunks", None) or []: + piece_name = ( + piece.get("name") + if isinstance(piece, dict) + else getattr(piece, "name", None) + ) + piece_id = ( + piece.get("id") + if isinstance(piece, dict) + else getattr(piece, "id", None) + ) + piece_args = ( + piece.get("args") + if isinstance(piece, dict) + else getattr(piece, "args", None) + ) + if piece_name and piece_id: + for rich in emitter.tool_start( + str(piece_id), str(piece_name) + ): + yield rich + if piece_id and isinstance(piece_args, str) and piece_args: + for rich in emitter.tool_args(str(piece_id), piece_args): + yield rich usage = getattr(chunk, "usage_metadata", None) if usage: usage_reported = True @@ -888,6 +915,8 @@ async def _stream_gen( raise if not accumulated_tool_calls: + for rich in emitter.end_text(): + yield rich full_output = (full_output or "") + accumulated_content break @@ -898,6 +927,14 @@ async def _stream_gen( steps += 1 full_output = (full_output or "") + accumulated_content + for tc in accumulated_tool_calls: + call_id = str(tc.get("id") or tc["name"]) + for rich in emitter.tool_start(call_id, str(tc["name"])): + yield rich + for rich in emitter.tool_end(call_id, arguments=tc.get("args")): + yield rich + for rich in emitter.end_text(): + yield rich assistant_msg = AIMessage( content=accumulated_content, tool_calls=accumulated_tool_calls ) @@ -934,10 +971,18 @@ 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.get("id") or tc["name"]), 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.get("id") or tc["name"]), result_val + ): + yield rich tool_results.append( ToolMessage( tool_call_id=tc.get("id") or tc["name"], content=str(result_val) diff --git a/packages/langchain-messages/tests/test_handler.py b/packages/langchain-messages/tests/test_handler.py index 0b337c71..c4824a6d 100644 --- a/packages/langchain-messages/tests/test_handler.py +++ b/packages/langchain-messages/tests/test_handler.py @@ -1178,10 +1178,8 @@ async def test_yields_chunk_events(self) -> None: llm = self._make_streaming_llm(["hello ", "world"]) h = create_langchain_messages_handler(llm=llm) 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"] async def test_streams_text_while_ignoring_thinking_blocks(self) -> None: from launchdarkly_ai_langchain_messages import create_langchain_messages_handler @@ -1200,7 +1198,13 @@ async def test_streams_text_while_ignoring_thinking_blocks(self) -> None: CONFIG, "q" ) ] - assert {"type": "chunk", "text": "visible answer"} in events + assert { + "type": "TEXT_MESSAGE_CONTENT", + "messageId": events[1]["messageId"] + if events[0]["type"] == "TEXT_MESSAGE_START" + else events[0]["messageId"], + "delta": "visible answer", + } in events assert events[-1]["output"] == "visible answer" async def test_yields_exactly_one_done_event(self) -> None: @@ -2583,5 +2587,6 @@ def factory(config: Any) -> Any: ] assert seen[0]["model"]["parameters"] == {"temperature": 0.2} assert any( - e.get("type") == "chunk" and e.get("text") == "streamed" for e in events + e.get("type") == "TEXT_MESSAGE_CONTENT" and e.get("delta") == "streamed" + for e in events ) From e7c3ecf2cbc48d2d239b472d5215e8d4af182c3a Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 13 Oct 2026 15:58:20 +0000 Subject: [PATCH 04/14] fix(langchain): keep argument deltas that omit the call id Later LangChain tool_call_chunks carry only an index. Remember the id from the first piece so those slices are not dropped. Co-authored-by: alexis --- .../handler.py | 37 ++++++++++++++++--- .../langchain-messages/tests/test_handler.py | 25 +++++++++++++ 2 files changed, 56 insertions(+), 6 deletions(-) diff --git a/packages/langchain-messages/src/launchdarkly_ai_langchain_messages/handler.py b/packages/langchain-messages/src/launchdarkly_ai_langchain_messages/handler.py index d982004f..b3b13c47 100644 --- a/packages/langchain-messages/src/launchdarkly_ai_langchain_messages/handler.py +++ b/packages/langchain-messages/src/launchdarkly_ai_langchain_messages/handler.py @@ -794,6 +794,7 @@ async def _stream_gen( accumulated_content = "" accumulated_tool_calls: list[Any] = [] + tool_chunk_ids: dict[int, str] = {} # The cache breakdown has to be accumulated alongside the scalars. LangChain reports it # per chunk in `usage_metadata.input_token_details`, and synthesizing a usage bag # without it would make the streaming span emit `cache_read = 0` where the blocking path @@ -840,18 +841,42 @@ async def _stream_gen( if isinstance(piece, dict) else getattr(piece, "id", None) ) + piece_index = ( + piece.get("index") + if isinstance(piece, dict) + else getattr(piece, "index", None) + ) piece_args = ( piece.get("args") if isinstance(piece, dict) else getattr(piece, "args", None) ) - if piece_name and piece_id: - for rich in emitter.tool_start( - str(piece_id), str(piece_name) - ): + # Later chunks often omit id and carry only index plus an args slice. + if ( + isinstance(piece_id, str) + and piece_id + and isinstance(piece_index, int) + ): + tool_chunk_ids[piece_index] = piece_id + elif ( + not isinstance(piece_id, str) or not piece_id + ) and isinstance(piece_index, int): + piece_id = tool_chunk_ids.get(piece_index) + if ( + isinstance(piece_name, str) + and piece_name + and isinstance(piece_id, str) + and piece_id + ): + for rich in emitter.tool_start(piece_id, piece_name): yield rich - if piece_id and isinstance(piece_args, str) and piece_args: - for rich in emitter.tool_args(str(piece_id), piece_args): + if ( + isinstance(piece_id, str) + and piece_id + and isinstance(piece_args, str) + and piece_args + ): + for rich in emitter.tool_args(piece_id, piece_args): yield rich usage = getattr(chunk, "usage_metadata", None) if usage: diff --git a/packages/langchain-messages/tests/test_handler.py b/packages/langchain-messages/tests/test_handler.py index c4824a6d..24a5711c 100644 --- a/packages/langchain-messages/tests/test_handler.py +++ b/packages/langchain-messages/tests/test_handler.py @@ -1181,6 +1181,31 @@ async def test_yields_chunk_events(self) -> None: deltas = [e["delta"] for e in events if e.get("type") == "TEXT_MESSAGE_CONTENT"] assert deltas == ["hello ", "world"] + async def test_later_tool_arg_chunks_keep_the_first_id(self) -> None: + from launchdarkly_ai_langchain_messages import create_langchain_messages_handler + + first = FakeAIMessage("") + first.tool_call_chunks = [ + {"name": "search", "id": "c1", "index": 0, "args": '{"q":'} + ] + second = FakeAIMessage("") + second.tool_call_chunks = [{"index": 0, "args": '"x"}'}] + llm = MagicMock() + llm.bind_tools = MagicMock(return_value=llm) + + async def _astream(_msgs: Any) -> AsyncGenerator[Any, None]: + yield first + yield second + + llm.astream = _astream + llm.ainvoke = AsyncMock(return_value=FakeAIMessage("")) + handler = create_langchain_messages_handler(llm=llm, capture_content=True) + events = [event async for event in await handler.stream(CONFIG, "q")] + args = "".join( + event["delta"] for event in events if event["type"] == "TOOL_CALL_ARGS" + ) + assert args == '{"q":"x"}' + async def test_streams_text_while_ignoring_thinking_blocks(self) -> None: from launchdarkly_ai_langchain_messages import create_langchain_messages_handler From 462209b1fb15b23432c91c925b131b0c77316323 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 13 Oct 2026 16:39:31 +0000 Subject: [PATCH 05/14] fix(langchain): keep every tool-argument slice and finish open calls Later chunks that omit the call id, and a tool_end that used to skip the complete arguments after a prefix, no longer publish truncated JSON. A failed turn also closes the calls it already opened. Co-authored-by: alexis --- .../client/src/launchdarkly_ai_server/agui.py | 22 ++++--- packages/client/tests/test_agui.py | 10 ++++ .../handler.py | 30 ++++++---- .../langchain-messages/tests/test_handler.py | 59 +++++++++++++++++++ 4 files changed, 103 insertions(+), 18 deletions(-) diff --git a/packages/client/src/launchdarkly_ai_server/agui.py b/packages/client/src/launchdarkly_ai_server/agui.py index cb750553..f717e4ae 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() @@ -159,25 +160,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/langchain-messages/src/launchdarkly_ai_langchain_messages/handler.py b/packages/langchain-messages/src/launchdarkly_ai_langchain_messages/handler.py index b3b13c47..f926fbf6 100644 --- a/packages/langchain-messages/src/launchdarkly_ai_langchain_messages/handler.py +++ b/packages/langchain-messages/src/launchdarkly_ai_langchain_messages/handler.py @@ -795,6 +795,7 @@ async def _stream_gen( accumulated_content = "" accumulated_tool_calls: list[Any] = [] tool_chunk_ids: dict[int, str] = {} + seen_call_ids: list[str] = [] # The cache breakdown has to be accumulated alongside the scalars. LangChain reports it # per chunk in `usage_metadata.input_token_details`, and synthesizing a usage bag # without it would make the streaming span emit `cache_read = 0` where the blocking path @@ -851,17 +852,20 @@ async def _stream_gen( if isinstance(piece, dict) else getattr(piece, "args", None) ) - # Later chunks often omit id and carry only index plus an args slice. - if ( - isinstance(piece_id, str) - and piece_id - and isinstance(piece_index, int) - ): - tool_chunk_ids[piece_index] = piece_id + # Later chunks often omit id. They carry an index, or, + # when only one call is open, just the args slice. + if isinstance(piece_id, str) and piece_id: + if piece_id not in seen_call_ids: + seen_call_ids.append(piece_id) + if isinstance(piece_index, int): + tool_chunk_ids[piece_index] = piece_id elif ( - not isinstance(piece_id, str) or not piece_id - ) and isinstance(piece_index, int): - piece_id = tool_chunk_ids.get(piece_index) + isinstance(piece_index, int) + and piece_index in tool_chunk_ids + ): + piece_id = tool_chunk_ids[piece_index] + elif len(seen_call_ids) == 1: + piece_id = seen_call_ids[0] if ( isinstance(piece_name, str) and piece_name @@ -901,6 +905,8 @@ async def _stream_gen( # path. fail_span(model_span, exc, ended) open_model_span = None + for rich in emitter.fail_open(): + yield rich raise # The already-mapped figures, not a bag rebuilt from them: this path summed the chunks @@ -1000,6 +1006,8 @@ async def _stream_gen( str(tc.get("id") or tc["name"]), 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 @@ -1053,6 +1061,8 @@ async def _stream_gen( if run_usage.reported: finish_root_span(span, config, run_usage.total) fail_span(span, exc, ended) + for rich in emitter.fail_open(): + yield rich raise finally: # A no-op on the success and failure paths, because both already ended their spans through diff --git a/packages/langchain-messages/tests/test_handler.py b/packages/langchain-messages/tests/test_handler.py index 24a5711c..8bb382e9 100644 --- a/packages/langchain-messages/tests/test_handler.py +++ b/packages/langchain-messages/tests/test_handler.py @@ -1206,6 +1206,65 @@ async def _astream(_msgs: Any) -> AsyncGenerator[Any, None]: ) assert args == '{"q":"x"}' + async def test_a_later_arg_chunk_without_an_index_keeps_the_only_id(self) -> None: + from launchdarkly_ai_langchain_messages import create_langchain_messages_handler + + first = FakeAIMessage("") + first.tool_call_chunks = [{"name": "search", "id": "c1", "args": '{"q":'}] + second = FakeAIMessage("") + second.tool_call_chunks = [{"args": '"x"}'}] + llm = MagicMock() + llm.bind_tools = MagicMock(return_value=llm) + + async def _astream(_msgs: Any) -> AsyncGenerator[Any, None]: + yield first + yield second + + llm.astream = _astream + llm.ainvoke = AsyncMock(return_value=FakeAIMessage("")) + handler = create_langchain_messages_handler(llm=llm, capture_content=True) + events = [event async for event in await handler.stream(CONFIG, "q")] + args = "".join( + event["delta"] for event in events if event["type"] == "TOOL_CALL_ARGS" + ) + assert args == '{"q":"x"}' + + async def test_tool_end_completes_a_truncated_argument_prefix(self) -> None: + from launchdarkly_ai_langchain_messages import create_langchain_messages_handler + + first = FakeAIMessage("") + first.tool_call_chunks = [{"name": "search", "id": "c1", "args": '{"q":'}] + second = FakeAIMessage( + "", + tool_calls=[{"name": "search", "id": "c1", "args": {"q": "x"}}], + ) + done = FakeAIMessage("ok") + llm = MagicMock() + llm.bind_tools = MagicMock(return_value=llm) + turns = [first, second] + + async def _astream(_msgs: Any) -> AsyncGenerator[Any, None]: + if turns: + yield turns.pop(0) + if turns: + yield turns.pop(0) + return + yield done + + llm.astream = _astream + llm.ainvoke = AsyncMock(return_value=FakeAIMessage("")) + handler = create_langchain_messages_handler(llm=llm, capture_content=True) + events = [ + event + async for event in await handler.stream( + CONFIG, "q", {"search": lambda _args: "found"} + ) + ] + args = "".join( + event["delta"] for event in events if event["type"] == "TOOL_CALL_ARGS" + ) + assert args == '{"q":"x"}' + async def test_streams_text_while_ignoring_thinking_blocks(self) -> None: from launchdarkly_ai_langchain_messages import create_langchain_messages_handler From 23a0548cece1391361ac1dd91c14ad23792d1d7f Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 13 Oct 2026 18:30:01 +0000 Subject: [PATCH 06/14] fix(langchain): run gathered tool calls and close a failed agent turn The messages loop used the last chunk's tool_calls, which drops earlier parallel calls. Chunks are merged the way LangChain adds them, and a failed agent stream closes the tool calls it already opened. Co-authored-by: alexis --- .../handler.py | 4 +- .../langchain-agents/tests/test_handler.py | 28 ++++++++ .../handler.py | 34 ++++++---- .../langchain-messages/tests/test_handler.py | 68 +++++++++++++++++++ 4 files changed, 121 insertions(+), 13 deletions(-) diff --git a/packages/langchain-agents/src/launchdarkly_ai_langchain_agents/handler.py b/packages/langchain-agents/src/launchdarkly_ai_langchain_agents/handler.py index 3004d260..606a5a7c 100644 --- a/packages/langchain-agents/src/launchdarkly_ai_langchain_agents/handler.py +++ b/packages/langchain-agents/src/launchdarkly_ai_langchain_agents/handler.py @@ -470,6 +470,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: # Inside the guard, because serialising the prompt raises on anything that is not # JSON-serialisable. A raise out here would leave the root open with the `finally` never @@ -494,7 +495,6 @@ async def _stream_gen( run_usage = create_run_usage() full_output = "" - emitter = TurnEmitter(capture_content=capture_content) # agent.astream() yields state updates per graph step: { [node_name]: { messages: [...] } } async for step_state in agent.astream( @@ -584,6 +584,8 @@ async def _stream_gen( if span_callbacks.run_usage.reported: finish_root_span(span, config, span_callbacks.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/langchain-agents/tests/test_handler.py b/packages/langchain-agents/tests/test_handler.py index 238c125a..476efa19 100644 --- a/packages/langchain-agents/tests/test_handler.py +++ b/packages/langchain-agents/tests/test_handler.py @@ -1451,6 +1451,34 @@ async def _bad_astream(*a: Any, **kw: Any) -> AsyncIterator[Any]: async for _ in await h.stream(_make_config(), "hi"): pass + @pytest.mark.asyncio + async def test_a_failed_run_closes_started_tool_calls(self) -> None: + mocks = _make_langchain_mock() + msg = _make_ai_msg("looking") + msg.tool_calls = [ + {"name": "search", "id": "c1", "args": {"q": "x"}}, + {"name": "other", "id": "d1", "args": {"n": 1}}, + ] + + async def _astream(*_a: Any, **_kw: Any) -> AsyncIterator[Any]: + yield {"agent": {"messages": [msg]}} + raise RuntimeError("tool failed") + + mocks["_agent"].astream = _astream + seen: list[dict[str, Any]] = [] + with ( + _patch_lc(mocks), + patch.object(spans_mod, "_HAS_OTEL", False), + pytest.raises(RuntimeError, match="tool failed"), + ): + handler = create_langchain_agents_handler(llm=MagicMock()) + async for event in await handler.stream(_make_config(), "q"): + seen.append(event) + results = { + event["toolCallId"] for event in seen if event["type"] == "TOOL_CALL_RESULT" + } + assert results == {"c1", "d1"} + # --------------------------------------------------------------------------- # §1.5 Streaming telemetry (Appendix A.5 — do not patch _HAS_OTEL=False) diff --git a/packages/langchain-messages/src/launchdarkly_ai_langchain_messages/handler.py b/packages/langchain-messages/src/launchdarkly_ai_langchain_messages/handler.py index f926fbf6..bdb29d28 100644 --- a/packages/langchain-messages/src/launchdarkly_ai_langchain_messages/handler.py +++ b/packages/langchain-messages/src/launchdarkly_ai_langchain_messages/handler.py @@ -151,6 +151,21 @@ def append_turn(turn: dict[str, Any]) -> None: return messages +def _merge_stream_chunk(gathered: Any, chunk: Any) -> Any: + """Fold one streamed chunk into the turn the way LangChain does. + + ``AIMessageChunk.tool_calls`` is parsed from that chunk alone. Keeping only + the last chunk drops every earlier parallel call and every argument slice + that chunk did not repeat. Addition merges chunks by index. + """ + if gathered is None: + return chunk + try: + return gathered + chunk + except TypeError: + return chunk + + def _assistant_output_messages( content: Any, tool_calls: list[Any] | None ) -> list[SpanMessage]: @@ -793,7 +808,7 @@ async def _stream_gen( open_model_span = model_span accumulated_content = "" - accumulated_tool_calls: list[Any] = [] + gathered: Any = None tool_chunk_ids: dict[int, str] = {} seen_call_ids: list[str] = [] # The cache breakdown has to be accumulated alongside the scalars. LangChain reports it @@ -894,9 +909,7 @@ async def _stream_gen( turn_usage.cache_creation += number_or_zero( details.get("cache_creation") ) - chunk_tools = getattr(chunk, "tool_calls", None) or [] - if chunk_tools: - accumulated_tool_calls = chunk_tools + gathered = _merge_stream_chunk(gathered, chunk) finish_reasons = lang_chain_finish_reasons(chunk) or finish_reasons open_chunk_stream = None except Exception as exc: @@ -905,8 +918,6 @@ async def _stream_gen( # path. fail_span(model_span, exc, ended) open_model_span = None - for rich in emitter.fail_open(): - yield rich raise # The already-mapped figures, not a bag rebuilt from them: this path summed the chunks @@ -923,6 +934,11 @@ async def _stream_gen( # billed these tokens: a content failure is our problem and must not report the run as # having spent less than it did. run_usage.add(turn_usage if usage_reported else None) + accumulated_tool_calls = [ + tc + for tc in (getattr(gathered, "tool_calls", None) or []) + if isinstance(tc, dict) and tc.get("name") + ] # Inside a guard for the same reason as the blocking path: a raise while serialising # completion content would otherwise leave this span for `finally` to end as abandoned, @@ -1002,12 +1018,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.get("id") or tc["name"]), 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/langchain-messages/tests/test_handler.py b/packages/langchain-messages/tests/test_handler.py index 8bb382e9..5a7b411e 100644 --- a/packages/langchain-messages/tests/test_handler.py +++ b/packages/langchain-messages/tests/test_handler.py @@ -1265,6 +1265,74 @@ async def _astream(_msgs: Any) -> AsyncGenerator[Any, None]: ) assert args == '{"q":"x"}' + async def test_parallel_chunks_use_the_gathered_tool_calls(self) -> None: + from langchain_core.messages import AIMessageChunk + + from launchdarkly_ai_langchain_messages import create_langchain_messages_handler + + first = AIMessageChunk( + content="", + tool_call_chunks=[ + {"name": "search", "args": "", "id": "call_1", "index": 0} + ], + ) + second = AIMessageChunk( + content="", + tool_call_chunks=[ + {"name": None, "args": '{"q":', "id": None, "index": 0}, + {"name": "other", "args": "", "id": "call_2", "index": 1}, + ], + ) + third = AIMessageChunk( + content="", + tool_call_chunks=[ + {"name": None, "args": '"x"}', "id": None, "index": 0}, + {"name": None, "args": "{}", "id": None, "index": 1}, + ], + ) + seen: list[dict[str, Any]] = [] + llm = MagicMock() + llm.bind_tools = MagicMock(return_value=llm) + first_turn = True + + async def _astream(_msgs: Any) -> AsyncGenerator[Any, None]: + nonlocal first_turn + if first_turn: + first_turn = False + yield first + yield second + yield third + return + yield FakeAIMessage("ok") + + llm.astream = _astream + llm.ainvoke = AsyncMock(return_value=FakeAIMessage("")) + + def search(args: dict[str, Any]) -> str: + seen.append(args) + return "found" + + def other(args: dict[str, Any]) -> str: + seen.append(args) + return "also" + + handler = create_langchain_messages_handler(llm=llm, capture_content=True) + events = [ + event + async for event in await handler.stream( + CONFIG, "q", {"search": search, "other": other} + ) + ] + starts = [ + event["toolCallId"] + for event in events + if event["type"] == "TOOL_CALL_START" + ] + assert starts == ["call_1", "call_2"] + assert seen == [{"q": "x"}, {}] + assert events[-1]["type"] == "done" + assert "ok" in events[-1]["output"] + async def test_streams_text_while_ignoring_thinking_blocks(self) -> None: from launchdarkly_ai_langchain_messages import create_langchain_messages_handler From c8d198c3f6253a8c4eab45d32d459fb1614c502b Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 13 Oct 2026 18:51:47 +0000 Subject: [PATCH 07/14] 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 3555490153f86c94a827054a5708f6df3a5d13af Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 13 Oct 2026 18:59:48 +0000 Subject: [PATCH 08/14] fix(langchain): close tool calls from the graph's tool messages The agent stream reports the tool result the graph already returned, instead of a second copy recorded by the span callback. Co-authored-by: alexis --- .../handler.py | 21 ++++---- .../launchdarkly_ai_langchain_agents/spans.py | 25 --------- .../langchain-agents/tests/test_handler.py | 54 +++++++++++++++++++ 3 files changed, 63 insertions(+), 37 deletions(-) diff --git a/packages/langchain-agents/src/launchdarkly_ai_langchain_agents/handler.py b/packages/langchain-agents/src/launchdarkly_ai_langchain_agents/handler.py index 606a5a7c..70911eb7 100644 --- a/packages/langchain-agents/src/launchdarkly_ai_langchain_agents/handler.py +++ b/packages/langchain-agents/src/launchdarkly_ai_langchain_agents/handler.py @@ -36,7 +36,6 @@ from ._version import PACKAGE_NAME, __version__ from .messages import to_lang_chain_messages from .spans import ( - _tool_result_text, build_span_callbacks, fail_span, finish_root_span, @@ -530,17 +529,15 @@ async def _stream_gen( call_id, arguments=tc.get("args") ): yield rich - for call_id, name, arguments, output, failed in list( - span_callbacks.finished_calls - ): - for rich in emitter.tool_start(call_id, name): - yield rich - for rich in emitter.tool_end(call_id, arguments=arguments): - yield rich - value = _tool_result_text(output) if not failed else None - for rich in emitter.tool_result(call_id, value, failed=failed): - yield rich - span_callbacks.finished_calls.clear() + elif getattr(msg, "type", None) == "tool": + call_id = str(getattr(msg, "tool_call_id", "") or "") + if call_id: + for rich in emitter.tool_result( + call_id, + getattr(msg, "content", None), + failed=getattr(msg, "status", None) == "error", + ): + yield rich for rich in emitter.end_text(): yield rich diff --git a/packages/langchain-agents/src/launchdarkly_ai_langchain_agents/spans.py b/packages/langchain-agents/src/launchdarkly_ai_langchain_agents/spans.py index a7b43181..3621de06 100644 --- a/packages/langchain-agents/src/launchdarkly_ai_langchain_agents/spans.py +++ b/packages/langchain-agents/src/launchdarkly_ai_langchain_agents/spans.py @@ -304,8 +304,6 @@ def __init__( self.run_usage = run_usage self.model_spans: dict[str, Any] = {} self.tool_spans: dict[str, Any] = {} - self._open_calls: dict[str, tuple[str, str, Any]] = {} - self.finished_calls: list[tuple[str, str, Any, Any, bool]] = [] def _start_model(self, run_id: Any, messages: Any = None) -> None: key = str(run_id) @@ -406,23 +404,12 @@ async def on_tool_start( # Tracked before the content write, not after. Serialising the arguments can raise, and a # span created but never inserted is unreachable by every cleanup path there is. self.tool_spans[str(run_id)] = span - call_id = str(tool_call_id or run_id) - self._open_calls[str(run_id)] = ( - call_id, - str(tool_name), - kwargs.get("inputs") or input_str, - ) set_tool_call_content_attributes( span, self._capture_content, arguments=kwargs.get("inputs") or input_str ) async def on_tool_end(self, output: Any, *, run_id: Any, **kwargs: Any) -> None: key = str(run_id) - recorded = self._open_calls.pop(key, None) - if recorded is not None: - self.finished_calls.append( - (recorded[0], recorded[1], recorded[2], output, False) - ) span = self.tool_spans.pop(key, None) if span is None: return @@ -440,11 +427,6 @@ async def on_tool_error( self, error: BaseException, *, run_id: Any, **kwargs: Any ) -> None: key = str(run_id) - recorded = self._open_calls.pop(key, None) - if recorded is not None: - self.finished_calls.append( - (recorded[0], recorded[1], recorded[2], None, True) - ) span = self.tool_spans.pop(key, None) if span is None: return @@ -510,13 +492,6 @@ def __init__( self.run_usage = run_usage self._handler = handler - @property - def finished_calls(self) -> list[tuple[str, str, Any, Any, bool]]: - if self._handler is None: - return [] - calls: list[tuple[str, str, Any, Any, bool]] = self._handler.finished_calls - return calls - def close_open_spans(self, error: BaseException) -> None: if self._handler is not None: self._handler.close_open_spans(error) diff --git a/packages/langchain-agents/tests/test_handler.py b/packages/langchain-agents/tests/test_handler.py index 476efa19..85ce6e00 100644 --- a/packages/langchain-agents/tests/test_handler.py +++ b/packages/langchain-agents/tests/test_handler.py @@ -1479,6 +1479,60 @@ async def _astream(*_a: Any, **_kw: Any) -> AsyncIterator[Any]: } assert results == {"c1", "d1"} + @pytest.mark.asyncio + async def test_tool_messages_close_parallel_calls(self) -> None: + mocks = _make_langchain_mock() + ai = _make_ai_msg("looking") + ai.tool_calls = [ + {"name": "search", "id": "c1", "args": {"q": "x"}}, + {"name": "other", "id": "d1", "args": {"n": 1}}, + ] + + def tool(call_id: str, content: str) -> Any: + message = MagicMock() + message.type = "tool" + message.tool_call_id = call_id + message.content = content + message.status = "success" + message.usage_metadata = None + return message + + async def _astream(*_a: Any, **_kw: Any) -> AsyncIterator[Any]: + yield {"agent": {"messages": [ai]}} + yield {"tools": {"messages": [tool("c1", "found"), tool("d1", "also")]}} + + mocks["_agent"].astream = _astream + with _patch_lc(mocks), patch.object(spans_mod, "_HAS_OTEL", False): + events = [ + event + async for event in await create_langchain_agents_handler( + llm=MagicMock(), capture_content=True + ).stream(_make_config(), "q") + ] + by_id: dict[str, list[str]] = {} + for event in events: + if not str(event.get("type", "")).startswith("TOOL_CALL"): + continue + by_id.setdefault(event["toolCallId"], []).append(event["type"]) + assert by_id["c1"] == [ + "TOOL_CALL_START", + "TOOL_CALL_ARGS", + "TOOL_CALL_END", + "TOOL_CALL_RESULT", + ] + assert by_id["d1"] == [ + "TOOL_CALL_START", + "TOOL_CALL_ARGS", + "TOOL_CALL_END", + "TOOL_CALL_RESULT", + ] + contents = { + event["toolCallId"]: event["content"] + for event in events + if event["type"] == "TOOL_CALL_RESULT" + } + assert contents == {"c1": "found", "d1": "also"} + # --------------------------------------------------------------------------- # §1.5 Streaming telemetry (Appendix A.5 — do not patch _HAS_OTEL=False) From 6b1393c41701973f34a0f4fcce71346113bffdf5 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 13 Oct 2026 19:16:11 +0000 Subject: [PATCH 09/14] fix(langchain): close tool calls the gathered message omits A streamed chunk can open a tool call that never appears in gathered.tool_calls. The success path then ended the turn without TOOL_CALL_END or TOOL_CALL_RESULT. Those calls are now closed as masked failures before the next turn, and a provider error closes the open text message after the span is failed. Co-authored-by: alexis --- .../handler.py | 11 +++ .../langchain-messages/tests/test_handler.py | 71 +++++++++++++++++++ 2 files changed, 82 insertions(+) diff --git a/packages/langchain-messages/src/launchdarkly_ai_langchain_messages/handler.py b/packages/langchain-messages/src/launchdarkly_ai_langchain_messages/handler.py index bdb29d28..a86b5c23 100644 --- a/packages/langchain-messages/src/launchdarkly_ai_langchain_messages/handler.py +++ b/packages/langchain-messages/src/launchdarkly_ai_langchain_messages/handler.py @@ -964,6 +964,11 @@ async def _stream_gen( if not accumulated_tool_calls: for rich in emitter.end_text(): yield rich + # A chunk can open a call that never becomes a complete tool + # call. Close it before the turn ends, or the client keeps it + # open across RUN_FINISHED. + for rich in emitter.fail_open(): + yield rich full_output = (full_output or "") + accumulated_content break @@ -1031,6 +1036,10 @@ async def _stream_gen( tool_call_id=tc.get("id") or tc["name"], content=str(result_val) ) ) + # Chunks can start a call the gathered message then omits. The + # calls we ran are already closed; this finishes the rest. + for rich in emitter.fail_open(): + yield rich conversation_messages.extend(tool_results) full_output_str = ( @@ -1071,6 +1080,8 @@ async def _stream_gen( if run_usage.reported: finish_root_span(span, config, run_usage.total) fail_span(span, exc, ended) + for rich in emitter.end_text(): + yield rich for rich in emitter.fail_open(): yield rich raise diff --git a/packages/langchain-messages/tests/test_handler.py b/packages/langchain-messages/tests/test_handler.py index 5a7b411e..6d71647d 100644 --- a/packages/langchain-messages/tests/test_handler.py +++ b/packages/langchain-messages/tests/test_handler.py @@ -1228,6 +1228,77 @@ async def _astream(_msgs: Any) -> AsyncGenerator[Any, None]: event["delta"] for event in events if event["type"] == "TOOL_CALL_ARGS" ) assert args == '{"q":"x"}' + result = next(event for event in events if event["type"] == "TOOL_CALL_RESULT") + assert result["toolCallId"] == "c1" + assert result["metadata"]["launchdarkly"]["toolFailed"] is True + assert events.index(result) < len(events) - 1 + assert events[-1]["type"] == "done" + + async def test_a_gathered_call_list_that_omits_a_started_call_closes_it( + self, + ) -> None: + from launchdarkly_ai_langchain_messages import create_langchain_messages_handler + + first = FakeAIMessage("") + first.tool_call_chunks = [ + {"name": "search", "id": "c1", "args": "{}"}, + {"name": "other", "id": "c2", "args": "{}"}, + ] + second = FakeAIMessage( + "", + tool_calls=[{"name": "search", "id": "c1", "args": {}}], + ) + done = FakeAIMessage("ok") + llm = MagicMock() + llm.bind_tools = MagicMock(return_value=llm) + turns = [[first, second], [done]] + + async def _astream(_msgs: Any) -> AsyncGenerator[Any, None]: + for chunk in turns.pop(0): + yield chunk + + llm.astream = _astream + handler = create_langchain_messages_handler(llm=llm, capture_content=True) + events = [ + event + async for event in await handler.stream( + CONFIG, "q", {"search": lambda _args: "found"} + ) + ] + results = { + event["toolCallId"]: event + for event in events + if event["type"] == "TOOL_CALL_RESULT" + } + assert results["c1"]["content"] == "found" + assert "metadata" not in results["c1"] + assert results["c2"]["metadata"]["launchdarkly"]["toolFailed"] is True + ok = next( + event + for event in events + if event["type"] == "TEXT_MESSAGE_CONTENT" and event["delta"] == "ok" + ) + assert events.index(results["c2"]) < events.index(ok) + + async def test_a_provider_error_closes_the_open_text(self) -> None: + from launchdarkly_ai_langchain_messages import create_langchain_messages_handler + + llm = MagicMock() + llm.bind_tools = MagicMock(return_value=llm) + + async def _astream(_msgs: Any) -> AsyncGenerator[Any, None]: + yield FakeAIMessage("Hello") + raise RuntimeError("stream broke") + + llm.astream = _astream + seen: list[dict[str, Any]] = [] + handler = create_langchain_messages_handler(llm=llm) + with pytest.raises(RuntimeError, match="stream broke"): + 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_tool_end_completes_a_truncated_argument_prefix(self) -> None: from launchdarkly_ai_langchain_messages import create_langchain_messages_handler From fa05f747d57da3698a2a6f8b1d4093856ad17b2b Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 13 Oct 2026 14:18:55 +0000 Subject: [PATCH 10/14] 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 11/14] 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 12/14] 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 13/14] 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: From b76235cba3c9a7370f1e09038acd31d34f832851 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 14 Oct 2026 21:52:39 +0000 Subject: [PATCH 14/14] fix(langchain): close tool calls a successful agent run leaves open A graph step can start a tool call and never send a tool message. The success path ended the text and yielded done with that call still open. It now closes those calls as masked failures before the run finishes. Co-authored-by: alexis --- .../handler.py | 5 ++++ .../langchain-agents/tests/test_handler.py | 23 +++++++++++++++++++ 2 files changed, 28 insertions(+) diff --git a/packages/langchain-agents/src/launchdarkly_ai_langchain_agents/handler.py b/packages/langchain-agents/src/launchdarkly_ai_langchain_agents/handler.py index bb2c25c3..4f14b1ae 100644 --- a/packages/langchain-agents/src/launchdarkly_ai_langchain_agents/handler.py +++ b/packages/langchain-agents/src/launchdarkly_ai_langchain_agents/handler.py @@ -547,6 +547,11 @@ async def _stream_gen( for rich in emitter.end_text(): yield rich + # A step can start a call the graph never answers with a tool message. + # Close it before the run finishes, or the client keeps it open. + for rich in emitter.fail_open(): + yield rich + # The same reconciliation the blocking path does, for the same reason. This walk reads # usage_metadata off the astream payloads, and sees nothing at all unless the graph is # streaming updates rather than value snapshots. The callbacks read the LLMResult and also diff --git a/packages/langchain-agents/tests/test_handler.py b/packages/langchain-agents/tests/test_handler.py index 5f65d6e1..07bf45a2 100644 --- a/packages/langchain-agents/tests/test_handler.py +++ b/packages/langchain-agents/tests/test_handler.py @@ -1533,6 +1533,29 @@ async def _astream(*_a: Any, **_kw: Any) -> AsyncIterator[Any]: } assert contents == {"c1": "found", "d1": "also"} + @pytest.mark.asyncio + async def test_a_successful_run_closes_a_call_with_no_tool_message(self) -> None: + mocks = _make_langchain_mock() + ai = _make_ai_msg("looking") + ai.tool_calls = [{"name": "search", "id": "c1", "args": {"q": "x"}}] + + async def _astream(*_a: Any, **_kw: Any) -> AsyncIterator[Any]: + yield {"agent": {"messages": [ai]}} + + mocks["_agent"].astream = _astream + with _patch_lc(mocks), patch.object(spans_mod, "_HAS_OTEL", False): + events = [ + event + async for event in await create_langchain_agents_handler( + llm=MagicMock() + ).stream(_make_config(), "q") + ] + result = next(event for event in events if event["type"] == "TOOL_CALL_RESULT") + assert result["toolCallId"] == "c1" + assert result["metadata"]["launchdarkly"]["toolFailed"] is True + assert events[-1]["type"] == "done" + assert events.index(result) < len(events) - 1 + # --------------------------------------------------------------------------- # §1.5 Streaming telemetry (Appendix A.5 — do not patch _HAS_OTEL=False)