From 3633cd3c98e6be12af30e1a9507d98f67ade624c Mon Sep 17 00:00:00 2001 From: ran Date: Thu, 30 Jul 2026 13:18:58 +0300 Subject: [PATCH] feat(crewai): START/CONTENT/END triples as the default wire shape (PNI-136, part 2) Text and tool-call output now streams as canonical START/CONTENT/END triples; emission_shape="chunks" (or AGUI_CREWAI_EMISSION_SHAPE=chunks) keeps the previous CHUNK form for consumers pinned to it. Chunks were never universally applicable: apply/default.ts throws on an untransformed chunk, so any raw-SSE reader could not consume CrewAI's stream, and CrewAI's own progressive-state demos break the client-side chunk transform when a state update interleaves a tool call. Triples are the form the server can produce authoritatively because it owns the message-boundary state. Both transports route through one EmissionShaper, so shape and payload never depend on the installed crewai version. The shaper owns the message and tool-call lifecycle; the attribution BoundaryTracker owns steps. A shared state machine tracks the open message, the parallel tool calls (each routed by id), and closes every open sequence before a step, snapshot, or terminal boundary. A late delta for a closed tool call is dropped rather than reopened (a second START would duplicate the call client-side); an interleaved progressive STATE_SNAPSHOT/CUSTOM closes an open text message but leaves tool calls open, because litellm stamps the tool id on the first delta only. sdk streams tool calls keyed by delta index with accumulated identity and a parent message id, so parallel calls and split-identity providers survive and the live TOOL_CALL_ARGS match the final ModelResponse. parent_message_id is preserved on both transports and both shapes. get_capabilities reports the resolved wire shape. Built fresh on current main and integrated with the merged RAW passthrough, get_capabilities, MCP, checkpointing, and multi-agent attribution. --- integrations/crew-ai/python/README.md | 29 ++ .../python/ag_ui_crewai/_capabilities.py | 38 +- .../crew-ai/python/ag_ui_crewai/_config.py | 35 ++ .../crew-ai/python/ag_ui_crewai/_frames.py | 333 ++++++++++++++---- .../crew-ai/python/ag_ui_crewai/endpoint.py | 86 ++++- .../crew-ai/python/ag_ui_crewai/sdk.py | 111 ++++-- .../tests/test_capability_declaration.py | 63 ++++ .../crew-ai/python/tests/test_streaming.py | 317 ++++++++++++++--- 8 files changed, 848 insertions(+), 164 deletions(-) diff --git a/integrations/crew-ai/python/README.md b/integrations/crew-ai/python/README.md index ffefd4fbf..0b7d41db4 100644 --- a/integrations/crew-ai/python/README.md +++ b/integrations/crew-ai/python/README.md @@ -52,6 +52,35 @@ add_crewai_flow_fastapi_endpoint(app, MyFlow(), "/flow") ## Protocol surface +### Wire shape: START / CONTENT / END triples (default) + +Text and tool-call output is emitted as `TEXT_MESSAGE_START` / `TEXT_MESSAGE_CONTENT` +/ `TEXT_MESSAGE_END` and `TOOL_CALL_START` / `TOOL_CALL_ARGS` / `TOOL_CALL_END` — the +protocol's canonical discrete form. `emission_shape="chunks"` (or +`AGUI_CREWAI_EMISSION_SHAPE=chunks`) opts back into the previous +`TEXT_MESSAGE_CHUNK` / `TOOL_CALL_CHUNK` form. + +```python +add_crewai_flow_fastapi_endpoint(app, MyFlow(), "/flow") # triples +add_crewai_flow_fastapi_endpoint(app, MyFlow(), "/flow", emission_shape="chunks") +``` + +The two shapes are **not** equivalent. In `chunks` mode a `copilotkit_emit_state` / +`copilotkit_predict_state` call that lands between two tool-call argument deltas makes +`@ag-ui/client`'s chunk transform close the call and then throw on the next delta; +triples keep the call open server-side, because the server knows more deltas are +coming and the client does not. Triples are also the form any consumer can apply +directly — `apply/default.ts` throws if a chunk reaches it untransformed — so a raw +SSE reader (the Python SDK, conformance tooling, custom clients) needs no +chunk-transform stage. + +Both transports (the crewai >= 1.6 `StreamFrame` path and the legacy +event-bus-listener fallback) route through one `EmissionShaper`, so the event shape +and payload never depend on the installed crewai version. A run never ends with an +open sequence: any open message, tool call, or step is closed before `RUN_FINISHED`. +MCP tool executions always use triples regardless of this setting — their name, args +and result arrive together rather than streamed. + ### RAW passthrough (opt-in, default OFF) `emit_raw_events=True` mirrors the crewai events this bridge does **not** map onto diff --git a/integrations/crew-ai/python/ag_ui_crewai/_capabilities.py b/integrations/crew-ai/python/ag_ui_crewai/_capabilities.py index 431ab9493..a7cac504a 100644 --- a/integrations/crew-ai/python/ag_ui_crewai/_capabilities.py +++ b/integrations/crew-ai/python/ag_ui_crewai/_capabilities.py @@ -696,6 +696,7 @@ def _reasoning_capability(llm: Any, *, raw_enabled: bool = False) -> dict: def get_capabilities( *, llm: Any = None, + emission_shape: str | None = None, emit_raw_events: bool | None = None, ) -> dict: """Return the CrewAI bridge's capability declaration. @@ -739,9 +740,32 @@ def get_capabilities( # ``_config`` is a leaf (``_env`` + stdlib only), imported locally purely to # keep this module's "crewai / litellm / stdlib only" property for every path # that never calls ``get_capabilities``. - from ._config import DEFAULT_EMIT_RAW_EVENTS, resolve_emit_raw_events + from ._config import ( + DEFAULT_EMIT_RAW_EVENTS, + resolve_emission_shape, + resolve_emit_raw_events, + ) resolved_raw = resolve_emit_raw_events(emit_raw_events) + resolved_shape = resolve_emission_shape(emission_shape) + text_events = ( + [EventType.TEXT_MESSAGE_CHUNK.value] + if resolved_shape == "chunks" + else [ + EventType.TEXT_MESSAGE_START.value, + EventType.TEXT_MESSAGE_CONTENT.value, + EventType.TEXT_MESSAGE_END.value, + ] + ) + tool_events = ( + [EventType.TOOL_CALL_CHUNK.value] + if resolved_shape == "chunks" + else [ + EventType.TOOL_CALL_START.value, + EventType.TOOL_CALL_ARGS.value, + EventType.TOOL_CALL_END.value, + ] + ) return { "identity": {"type": "crewai", "crewaiVersion": CAPABILITIES.crewai_version}, "humanInTheLoop": { @@ -768,12 +792,12 @@ def get_capabilities( "streamFrames": CAPABILITIES.stream_frame_available, }, "wireShape": { - # This build streams LLM text / tool calls as CHUNK events. MCP tool - # executions are the exception: their name, args and result arrive - # together, so they already emit canonical TOOL_CALL_* triples. - "emissionShape": "chunks", - "textMessages": [EventType.TEXT_MESSAGE_CHUNK.value], - "toolCalls": [EventType.TOOL_CALL_CHUNK.value], + # START/CONTENT/END triples by default; "chunks" is a compatibility + # opt-out. MCP tool executions always use triples (name, args and result + # arrive together, not streamed), independent of this setting. + "emissionShape": resolved_shape, + "textMessages": text_events, + "toolCalls": tool_events, "mcpToolCalls": [ EventType.TOOL_CALL_START.value, EventType.TOOL_CALL_ARGS.value, diff --git a/integrations/crew-ai/python/ag_ui_crewai/_config.py b/integrations/crew-ai/python/ag_ui_crewai/_config.py index 664b51698..cbfe6d9ff 100644 --- a/integrations/crew-ai/python/ag_ui_crewai/_config.py +++ b/integrations/crew-ai/python/ag_ui_crewai/_config.py @@ -26,6 +26,15 @@ EMIT_RAW_EVENTS_ENV_VAR = "AGUI_CREWAI_EMIT_RAW_EVENTS" +# Wire shape for streamed text / tool-call output. Triples (START/CONTENT/END) is +# the canonical discrete form and the shipped default; "chunks" is a compatibility +# opt-out. Kept here (a leaf module) so the capability declaration can report it +# without importing the streaming stack. +SUPPORTED_EMISSION_SHAPES = frozenset({"triples", "chunks"}) +DEFAULT_EMISSION_SHAPE = "triples" + +EMISSION_SHAPE_ENV_VAR = "AGUI_CREWAI_EMISSION_SHAPE" + # Vocabulary ``_parse_env_bool`` accepts, so the "was this value used?" check stays # in step with the parser instead of duplicating its token list. _BOOL_TOKENS = _TRUE_VALUES | _FALSE_VALUES @@ -76,3 +85,29 @@ def resolve_emit_raw_events(emit_raw_events: bool | None) -> bool: used = raw is not None and raw.strip().casefold() in _BOOL_TOKENS _warn_if_env_value_ignored(EMIT_RAW_EVENTS_ENV_VAR, raw, used) return resolved + + +def resolve_emission_shape(emission_shape: str | None) -> str: + """Resolve the wire shape: explicit argument > env var > shipped default.""" + if emission_shape is not None: + if not isinstance(emission_shape, str): + raise ValueError( + f"emission_shape must be a string, got " + f"{type(emission_shape).__name__} ({emission_shape!r})" + ) + normalized = emission_shape.strip().casefold() + if normalized not in SUPPORTED_EMISSION_SHAPES: + raise ValueError( + f"Unknown emission_shape {emission_shape!r}; " + f"expected one of {sorted(SUPPORTED_EMISSION_SHAPES)}" + ) + return normalized + raw = os.environ.get(EMISSION_SHAPE_ENV_VAR) + resolved = DEFAULT_EMISSION_SHAPE + used = False + if raw is not None: + token = raw.strip().casefold() + if token in SUPPORTED_EMISSION_SHAPES: + resolved, used = token, True + _warn_if_env_value_ignored(EMISSION_SHAPE_ENV_VAR, raw, used) + return resolved diff --git a/integrations/crew-ai/python/ag_ui_crewai/_frames.py b/integrations/crew-ai/python/ag_ui_crewai/_frames.py index c50ebacd7..9412119bc 100644 --- a/integrations/crew-ai/python/ag_ui_crewai/_frames.py +++ b/integrations/crew-ai/python/ag_ui_crewai/_frames.py @@ -48,15 +48,13 @@ proportional to in-flight (not total) frames. If a future crewai release makes frame retention opt-out, revisit the session consumption in ``endpoint.py``. -SWAPPABLE EMISSION SHAPE: CrewAI is currently -the only integration emitting TEXT_MESSAGE_CHUNK / TOOL_CALL_CHUNK (chunks) -rather than the START/CONTENT/END triples the six other integrations emit. That -final choice belongs to a Parity-lane ticket, not this migration. The translator -therefore routes LLM text / LLM-tool-call emission through a single -``emission_shape`` strategy that DEFAULTS to ``"chunks"`` so this migration is -byte-for-byte behavior-preserving on that channel. The ``"triples"`` strategy is -a deliberate NotImplementedError placeholder — the Parity ticket owns wiring it -up and flipping the default. +EMISSION SHAPE: LLM text / tool-call output ships as START/CONTENT/END triples by +default, the canonical discrete form any AG-UI consumer can apply without a +chunk-transform stage. ``emission_shape="chunks"`` opts back into the previous +CHUNK form. The message / tool-call open-close lifecycle lives in a shared +``EmissionShaper`` so both transports emit the same shape; STEP lifecycle is owned +separately by the attribution ``BoundaryTracker``. Under ``"chunks"`` the shaper is +an identity passthrough, so that mode is byte-for-byte the previous behaviour. MCP EVENTS are the ONE exception to "chunks-only": crewai's discrete MCP tool executions (name + full args + result arrive together, not streamed) @@ -70,6 +68,7 @@ from __future__ import annotations import logging +import uuid from collections.abc import Callable from typing import Any @@ -81,7 +80,13 @@ from ag_ui.core.events import ( RawEvent, TextMessageChunkEvent, + TextMessageStartEvent, + TextMessageContentEvent, + TextMessageEndEvent, ToolCallChunkEvent, + ToolCallStartEvent, + ToolCallArgsEvent, + ToolCallEndEvent, StepFinishedEvent, MessagesSnapshotEvent, StateSnapshotEvent, @@ -165,6 +170,206 @@ def _agent_role(event: Any) -> str: return _coerce_name(role, "agent") +class EmissionShaper: + """Message- and tool-call-level open/close state for START/CONTENT/END triples. + + Shared by both transports so the wire shape never depends on which runs: the + StreamFrame translator delegates its text / tool handling here, and the legacy + driver runs its already-wire events through :meth:`reshape`. STEP lifecycle is + NOT handled here — the attribution ``BoundaryTracker`` owns it. Under + ``"chunks"`` every method is an identity passthrough. + """ + + def __init__( + self, + shape: str = "triples", + *, + thread_id: str | None = None, + run_id: str | None = None, + ) -> None: + if shape not in _SUPPORTED_EMISSION_SHAPES: + raise ValueError( + f"Unknown emission_shape {shape!r}; " + f"expected one of {sorted(_SUPPORTED_EMISSION_SHAPES)}" + ) + self.shape = shape + self._thread_id = thread_id + self._run_id = run_id + self._text_open = False + self._open_message_id: str | None = None + self._open_tool_calls: list[str] = [] + self._closed_tool_calls: set[str] = set() + + @property + def open_tool_calls(self) -> tuple[str, ...]: + return tuple(self._open_tool_calls) + + def text(self, event: Any) -> list[Any]: + message_id = getattr(event, "message_id", None) + role = getattr(event, "role", None) + delta = getattr(event, "delta", None) + if self.shape == "chunks": + return [ + TextMessageChunkEvent( + type=EventType.TEXT_MESSAGE_CHUNK, + message_id=message_id, + role=role, + delta=delta, + ) + ] + out: list[Any] = [] + # A different message id, or an open tool call, is a sequence boundary. + if ( + self._text_open and message_id is not None + and message_id != self._open_message_id + ) or self._open_tool_calls: + out.extend(self.flush()) + if not self._text_open: + self._text_open = True + # message_id is a required str on the triple events; a producer that + # omits it (litellm always sets it, but the shaper is a public seam) + # gets a generated id so START/CONTENT/END stay valid and paired. + self._open_message_id = message_id or uuid.uuid4().hex + out.append( + TextMessageStartEvent( + type=EventType.TEXT_MESSAGE_START, + message_id=self._open_message_id, + role=role or "assistant", + ) + ) + if delta is not None: + out.append( + TextMessageContentEvent( + type=EventType.TEXT_MESSAGE_CONTENT, + message_id=self._open_message_id, + delta=delta, + ) + ) + return out + + def tool(self, event: Any) -> list[Any]: + tool_call_id = getattr(event, "tool_call_id", None) + tool_call_name = getattr(event, "tool_call_name", None) + delta = getattr(event, "delta", None) + if self.shape == "chunks": + return [ + ToolCallChunkEvent( + type=EventType.TOOL_CALL_CHUNK, + tool_call_id=tool_call_id, + tool_call_name=tool_call_name, + delta=delta, + parent_message_id=getattr(event, "parent_message_id", None), + ) + ] + out: list[Any] = [] + if self._text_open: + out.extend(self.flush(tools=False)) + if tool_call_id is not None and tool_call_id in self._open_tool_calls: + if delta is not None: + out.append( + ToolCallArgsEvent( + type=EventType.TOOL_CALL_ARGS, + tool_call_id=tool_call_id, + delta=delta, + ) + ) + return out + if tool_call_id is not None and tool_call_id in self._closed_tool_calls: + _LOGGER.error( + "ag-ui-crewai dropped a TOOL_CALL_CHUNK for the already-closed call " + "%r: reopening it would duplicate the tool call client-side " + "(thread=%s run=%s)", + tool_call_id, + self._thread_id, + self._run_id, + ) + return out + if tool_call_id is None or tool_call_name is None: + _LOGGER.error( + "ag-ui-crewai dropped a TOOL_CALL_CHUNK with no open call to attach " + "to: the first chunk must carry tool_call_id and tool_call_name " + "(got id=%r name=%r, thread=%s run=%s)", + tool_call_id, + tool_call_name, + self._thread_id, + self._run_id, + ) + return out + self._open_tool_calls.append(tool_call_id) + out.append( + ToolCallStartEvent( + type=EventType.TOOL_CALL_START, + tool_call_id=tool_call_id, + tool_call_name=tool_call_name, + parent_message_id=getattr(event, "parent_message_id", None), + ) + ) + if delta is not None: + out.append( + ToolCallArgsEvent( + type=EventType.TOOL_CALL_ARGS, + tool_call_id=tool_call_id, + delta=delta, + ) + ) + return out + + def flush(self, *, tools: bool = True) -> list[Any]: + """Close the open text message, and (when ``tools``) the open tool calls.""" + if self.shape == "chunks": + return [] + out: list[Any] = [] + if self._text_open: + out.append( + TextMessageEndEvent( + type=EventType.TEXT_MESSAGE_END, + message_id=self._open_message_id, + ) + ) + self._text_open = False + self._open_message_id = None + if tools and self._open_tool_calls: + for tool_call_id in reversed(self._open_tool_calls): + out.append( + ToolCallEndEvent( + type=EventType.TOOL_CALL_END, tool_call_id=tool_call_id + ) + ) + self._closed_tool_calls.add(tool_call_id) + self._open_tool_calls = [] + return out + + def reshape(self, event: Any) -> list[Any]: + """Reshape one already-wire event (the legacy driver's queue output). + + Chunk events become triples; a boundary event flushes the open message / + tool sequence first, so the legacy transport emits the same shape as the + StreamFrame path. STEP events pass through (the listener already emits the + attributed steps) after a flush. Passthrough under ``"chunks"``. + """ + if self.shape == "chunks": + return [event] + event_type = getattr(event, "type", None) + if event_type == EventType.TEXT_MESSAGE_CHUNK: + return self.text(event) + if event_type == EventType.TOOL_CALL_CHUNK: + return self.tool(event) + if event_type == EventType.MESSAGES_SNAPSHOT: + # Only the method-finish path emits this; its message + tool calls are + # complete, so a full flush closes tools BEFORE the snapshot. + return [*self.flush(), event] + if event_type in (EventType.STATE_SNAPSHOT, EventType.CUSTOM): + # Progressive side-channel: close an open TEXT message, keep tools open. + return [*self.flush(tools=False), event] + if event_type in ( + EventType.STEP_STARTED, + EventType.STEP_FINISHED, + EventType.RUN_FINISHED, + ): + return [*self.flush(), event] + return [event] + + class StreamFrameTranslator: """Stateless-ish mapper from a RAW crewai/bridge event to AG-UI events. @@ -186,7 +391,7 @@ def __init__( thread_id: str, run_id: str, state_provider: Callable[[], Any], - emission_shape: str = "chunks", + emission_shape: str = "triples", ) -> None: if emission_shape not in _SUPPORTED_EMISSION_SHAPES: raise ValueError( @@ -197,6 +402,10 @@ def __init__( self._run_id = run_id self._state_provider = state_provider self.emission_shape = emission_shape + # Message / tool-call triple lifecycle (STEP lifecycle is the tracker's). + self._shaper = EmissionShaper( + emission_shape, thread_id=thread_id, run_id=run_id + ) # One ordered boundary stack per run, driven in emit order by # ``translate``: ``enter`` on a start frame, ``exit`` on the matching # finish, ``drain_all`` at run end so every STEP_STARTED is closed. @@ -244,10 +453,10 @@ def translate(self, event: Any) -> list[Any]: # a run that never opened. return [] self._run_finished_emitted = True - # Close every STEP_STARTED still open before RUN_FINISHED so a - # boundary whose finish frame never arrived does not dangle. - # ``drain_all`` returns them deepest-first for balanced closes. - events: list[Any] = [ + # Close open message / tool sequences, then every STEP still open, + # before RUN_FINISHED (``drain_all`` returns steps deepest-first). + events: list[Any] = list(self._shaper.flush()) + events += [ step_finished_event(b) for b in self._tracker.drain_all() ] events.append( @@ -259,24 +468,27 @@ def translate(self, event: Any) -> list[Any]: ) return events if event_type == _METHOD_STARTED: - return self._method_started_events(event) + return [*self._shaper.flush(), *self._method_started_events(event)] if event_type == _METHOD_FINISHED: - return self._method_finished_events(event) + return [*self._shaper.flush(), *self._method_finished_events(event)] if event_type == _METHOD_FAILED: # Close the method boundary so a failed flow method (the flow may # continue) does not leave a dangling STEP_STARTED. No snapshots. method_name = _coerce_name(getattr(event, "method_name", None), "method") - return self._close_boundaries( - self._tracker.exit(FLOW_METHOD, method_name), _METHOD_FAILED - ) + return [ + *self._shaper.flush(), + *self._close_boundaries( + self._tracker.exit(FLOW_METHOD, method_name), _METHOD_FAILED + ), + ] if event_type == _CREW_STARTED: - return self._crew_started_events(event) + return [*self._shaper.flush(), *self._crew_started_events(event)] if event_type in (_CREW_COMPLETED, _CREW_FAILED): - return self._crew_finished_events(event) + return [*self._shaper.flush(), *self._crew_finished_events(event)] if event_type == _AGENT_STARTED: - return self._agent_started_events(event) + return [*self._shaper.flush(), *self._agent_started_events(event)] if event_type in (_AGENT_COMPLETED, _AGENT_ERROR): - return self._agent_finished_events(event) + return [*self._shaper.flush(), *self._agent_finished_events(event)] # crewai's first-class MCP events (crewai >= 1.4). Emitted with # the agent/crew as source (not the flow), so the driver's sink parks @@ -288,23 +500,25 @@ def translate(self, event: Any) -> list[Any]: # Bridge-emitted events carry the AG-UI EventType string as ``.type`` # and the verbatim payload on typed attributes (no ``to_serializable``). if event_type == EventType.TEXT_MESSAGE_CHUNK: - return self._text_events(event) + return self._shaper.text(event) if event_type == EventType.TOOL_CALL_CHUNK: - return self._tool_events(event) + return self._shaper.tool(event) if event_type == EventType.CUSTOM: return [ + *self._shaper.flush(tools=False), CustomEvent( type=EventType.CUSTOM, name=getattr(event, "name", None), value=getattr(event, "value", None), - ) + ), ] if event_type == EventType.STATE_SNAPSHOT: return [ + *self._shaper.flush(tools=False), StateSnapshotEvent( type=EventType.STATE_SNAPSHOT, snapshot=getattr(event, "snapshot", None), - ) + ), ] # crewai native llm / tools / messages / lifecycle events and internal @@ -324,10 +538,10 @@ def finalize(self) -> list[Any]: """ if self._run_started_emitted and not self._run_finished_emitted: self._run_finished_emitted = True - # Same drain-before-terminal discipline as the ``flow_finished`` - # branch: close every open boundary (deepest-first) so the client - # never sees a dangling STEP_STARTED. - events: list[Any] = [ + # Same drain-before-terminal discipline as ``flow_finished``: close + # message / tool sequences, then every open boundary (deepest-first). + events: list[Any] = list(self._shaper.flush()) + events += [ step_finished_event(b) for b in self._tracker.drain_all() ] events.append( @@ -477,45 +691,13 @@ def _method_finished_events(self, event: Any) -> list[Any]: ) return events - # -- emission-shape strategy (default "chunks") ----------------------- + def close_pending(self) -> list[Any]: + """Close open message / tool sequences before a terminal RUN_ERROR. - def _text_events(self, event: Any) -> list[Any]: - if self.emission_shape == "chunks": - return [ - TextMessageChunkEvent( - type=EventType.TEXT_MESSAGE_CHUNK, - message_id=getattr(event, "message_id", None), - role=getattr(event, "role", None), - delta=getattr(event, "delta", None), - ) - ] - # TODO(Parity lane): emit TEXT_MESSAGE_START / _CONTENT / _END - # triples here to match the other six integrations. Do NOT flip the - # default in this migration (it must stay behavior-preserving on the - # wire). - raise NotImplementedError( - "emission_shape='triples' is a Parity-lane placeholder; " - "the StreamFrame migration ships the behavior-preserving " - "'chunks' shape only." - ) - - def _tool_events(self, event: Any) -> list[Any]: - if self.emission_shape == "chunks": - return [ - ToolCallChunkEvent( - type=EventType.TOOL_CALL_CHUNK, - tool_call_id=getattr(event, "tool_call_id", None), - tool_call_name=getattr(event, "tool_call_name", None), - delta=getattr(event, "delta", None), - ) - ] - # TODO(Parity lane): TOOL_CALL_START / _ARGS / _END triples — see - # the note in ``_text_events``. - raise NotImplementedError( - "emission_shape='triples' is a Parity-lane placeholder; " - "the StreamFrame migration ships the behavior-preserving " - "'chunks' shape only." - ) + Open STEPS are deliberately left open: the error is the explanation, and + claiming a step finished when it did not would be a lie on the wire. + """ + return self._shaper.flush() # Event types the translator recognises. Used to decide whether an event is @@ -527,11 +709,13 @@ def _tool_events(self, event: Any) -> list[Any]: _FLOW_FINISHED, _METHOD_STARTED, _METHOD_FINISHED, + _METHOD_FAILED, EventType.TEXT_MESSAGE_CHUNK, EventType.TOOL_CALL_CHUNK, EventType.CUSTOM, EventType.STATE_SNAPSHOT, } + | CREW_AGENT_LIFECYCLE_TYPES ) # Source tag on emitted RAW events, so a consumer can tell CrewAI passthrough apart @@ -556,8 +740,15 @@ def log_raw_loss(message: str, *args: Any) -> None: def is_recognized_event(event: Any) -> bool: - """True when the translator maps this event, so RAW must not duplicate it.""" - return getattr(event, "type", None) in _RECOGNIZED_EVENT_TYPES + """True when the translator maps this event, so RAW must not duplicate it. + + Includes the crew/agent lifecycle and MCP events that ``translate`` maps: with + ``emit_raw_events`` on, a recognized event must NOT also be mirrored as RAW. + """ + return ( + getattr(event, "type", None) in _RECOGNIZED_EVENT_TYPES + or is_mcp_event(event) + ) def raw_event_for(event: Any) -> RawEvent | None: diff --git a/integrations/crew-ai/python/ag_ui_crewai/endpoint.py b/integrations/crew-ai/python/ag_ui_crewai/endpoint.py index 3b6042038..2068b1078 100644 --- a/integrations/crew-ai/python/ag_ui_crewai/endpoint.py +++ b/integrations/crew-ai/python/ag_ui_crewai/endpoint.py @@ -32,8 +32,16 @@ reset_stream_sinks, ) from ._checkpoint import build_checkpoint_kwargs -from ._frames import StreamFrameTranslator, CREW_AGENT_LIFECYCLE_TYPES -from ._config import resolve_emit_raw_events as _resolve_emit_raw_events +from ._frames import ( + CREW_AGENT_LIFECYCLE_TYPES, + EmissionShaper, + StreamFrameTranslator, +) +from ._config import ( + DEFAULT_EMISSION_SHAPE as _DEFAULT_EMISSION_SHAPE, + resolve_emission_shape as _resolve_emission_shape, + resolve_emit_raw_events as _resolve_emit_raw_events, +) from ._frames import is_recognized_event, log_raw_loss, raw_event_for from .mcp import is_mcp_event, register_mcp_listeners, translate_mcp_event from .attribution import flat_method_attribution @@ -1032,6 +1040,10 @@ def _(source, event): tool_call_id=event.tool_call_id, tool_call_name=event.tool_call_name, delta=event.delta, + # Forwarded so the shaper can stamp TOOL_CALL_START.parentMessageId + # on the legacy transport too; without it the client renders each + # tool call as a separate assistant message. + parent_message_id=getattr(event, "parent_message_id", None), ) ) @crewai_event_bus.on(BridgedCustomEvent) @@ -1209,6 +1221,7 @@ async def _run_flow_event_stream( timeout: float | None, checkpoint_kwargs: dict | None = None, emit_raw_events: bool = False, + emission_shape: str = _DEFAULT_EMISSION_SHAPE, ): """Drive a single flow kickoff and yield encoded AG-UI events. @@ -1242,6 +1255,13 @@ async def _run_flow_event_stream( # main ``try:`` block, the registered queue is orphaned — nothing deletes # it. Wrap both in a narrow ``try/except`` that ``delete_queue``'s on # failure so the registration is symmetric. + # One shaper for the whole run: the legacy listener enqueues final wire events + # (chunks + lifecycle), and ``shaper.reshape`` turns them into the same shape + # the StreamFrame path emits, so the wire shape never depends on the transport. + shaper = EmissionShaper( + emission_shape, thread_id=input_data.thread_id, run_id=input_data.run_id + ) + queue = await create_queue(flow_copy) try: token = flow_context.set(flow_copy) @@ -1364,12 +1384,13 @@ async def _drain_queue_until_sentinel_or_empty(): # for today's extras events (StepStarted, # MessagesSnapshot, ...) since they don't declare the # fields. - _stamp_correlation_ids( - item_local, - thread_id=input_data.thread_id, - run_id=input_data.run_id, - ) - yield encoder.encode(item_local) + for shaped in shaper.reshape(item_local): + _stamp_correlation_ids( + shaped, + thread_id=input_data.thread_id, + run_id=input_data.run_id, + ) + yield encoder.encode(shaped) # Budget exhausted: exit regardless of what the # current pass produced. Log only when we cut a @@ -1550,13 +1571,13 @@ async def _drain_queue_until_sentinel_or_empty(): # fields (see _stamp_correlation_ids). RUN_STARTED / # RUN_FINISHED always do; future correlated events are covered # automatically. - _stamp_correlation_ids( - item, - thread_id=input_data.thread_id, - run_id=input_data.run_id, - ) - - yield encoder.encode(item) + for shaped in shaper.reshape(item): + _stamp_correlation_ids( + shaped, + thread_id=input_data.thread_id, + run_id=input_data.run_id, + ) + yield encoder.encode(shaped) except _KickoffCancelled: # Kickoff task was cancelled externally (not by our teardown path, @@ -1575,6 +1596,8 @@ async def _drain_queue_until_sentinel_or_empty(): # the client-facing message all agree on "kickoff". f"CrewAI kickoff was cancelled" ) + for _pending in shaper.flush(): + yield encoder.encode(_pending) yield encoder.encode( RunErrorEvent( message=message, @@ -1603,6 +1626,8 @@ async def _drain_queue_until_sentinel_or_empty(): f"thread={input_data.thread_id} run={input_data.run_id}: " f"CrewAI flow exceeded ceiling={ceiling_display}" ) + for _pending in shaper.flush(): + yield encoder.encode(_pending) yield encoder.encode( RunErrorEvent( message=message, @@ -1637,6 +1662,8 @@ async def _drain_queue_until_sentinel_or_empty(): f"CrewAI upstream timeout during kickoff " f"(ceiling={ceiling_display} did not fire)" ) + for _pending in shaper.flush(): + yield encoder.encode(_pending) yield encoder.encode( RunErrorEvent( message=message, @@ -1666,6 +1693,8 @@ async def _drain_queue_until_sentinel_or_empty(): # ``^[A-Z][A-Z0-9_]+$`` convention peer events follow and break # downstream regex-matchers. sanitized_name = _sanitize_exception_code(type(e).__name__) + for _pending in shaper.flush(): + yield encoder.encode(_pending) yield encoder.encode( RunErrorEvent( message=message, @@ -1807,6 +1836,7 @@ async def _run_flow_frame_stream( timeout: float | None, checkpoint_kwargs: dict | None = None, emit_raw_events: bool = False, + emission_shape: str = _DEFAULT_EMISSION_SHAPE, ): """StreamFrame-path driver: drive ``flow.astream`` and yield encoded AG-UI events. @@ -1854,6 +1884,7 @@ async def _run_flow_frame_stream( thread_id=input_data.thread_id, run_id=input_data.run_id, state_provider=lambda: getattr(flow_copy, "state", {}), + emission_shape=emission_shape, ) # Raw-event lookup buffer, populated by our scoped sink below. Keyed by # ``event.event_id`` (== ``StreamFrame.id``). Only OUTER-flow events land @@ -2105,6 +2136,8 @@ def _sink(source: Any, event: Any) -> None: f"thread={input_data.thread_id} run={input_data.run_id}: " f"CrewAI flow exceeded ceiling={ceiling_display}" ) + for _pending in translator.close_pending(): + yield encoder.encode(_pending) yield encoder.encode( RunErrorEvent( message=message, @@ -2132,6 +2165,8 @@ def _sink(source: Any, event: Any) -> None: f"CrewAI upstream timeout during kickoff " f"(ceiling={ceiling_display} did not fire)" ) + for _pending in translator.close_pending(): + yield encoder.encode(_pending) yield encoder.encode( RunErrorEvent( message=message, @@ -2151,6 +2186,8 @@ def _sink(source: Any, event: Any) -> None: f"CrewAI flow failed; see server logs" ) sanitized_name = _sanitize_exception_code(type(e).__name__) + for _pending in translator.close_pending(): + yield encoder.encode(_pending) yield encoder.encode( RunErrorEvent( message=message, @@ -2185,6 +2222,7 @@ def _run_flow_stream( timeout: float | None, checkpoint_kwargs: dict | None = None, emit_raw_events: bool = False, + emission_shape: str = _DEFAULT_EMISSION_SHAPE, ): """Select the StreamFrame path (crewai >= 1.6 + a real ``astream`` flow) or the legacy bus-listener path, returning the chosen async generator. @@ -2206,6 +2244,7 @@ def _run_flow_stream( timeout=timeout, checkpoint_kwargs=checkpoint_kwargs, emit_raw_events=emit_raw_events, + emission_shape=emission_shape, ) return _run_flow_event_stream( flow_copy=flow_copy, @@ -2215,6 +2254,7 @@ def _run_flow_stream( timeout=timeout, checkpoint_kwargs=checkpoint_kwargs, emit_raw_events=emit_raw_events, + emission_shape=emission_shape, ) @@ -2224,14 +2264,21 @@ def add_crewai_flow_fastapi_endpoint( path: str = "/", *, emit_raw_events: bool | None = None, + emission_shape: str | None = None, ): """Adds a CrewAI endpoint to the FastAPI app. + ``emission_shape`` selects the wire shape for text / tool-call output: + ``"triples"`` (START/CONTENT/END, the default) or ``"chunks"`` (the previous + CHUNK form) — ``None`` reads ``AGUI_CREWAI_EMISSION_SHAPE``. + ``emit_raw_events`` opts into RAW passthrough: the crewai events this bridge does not map (its llm / agent / task / tool channels, nested-flow lifecycle, internals) - are mirrored onto AG-UI ``RAW`` events. Off by default. Resolved at registration, - so a non-bool raises here rather than per request; ``None`` reads + are mirrored onto AG-UI ``RAW`` events. Off by default; ``None`` reads ``AGUI_CREWAI_EMIT_RAW_EVENTS``. + + Both resolve at registration, so a bad value fails once at startup rather than + per request. """ global GLOBAL_EVENT_LISTENER # pylint: disable=global-statement @@ -2256,6 +2303,7 @@ def add_crewai_flow_fastapi_endpoint( # Resolved HERE, not per request: a wrong-typed argument should fail once at # startup rather than on every call. resolved_emit_raw_events = _resolve_emit_raw_events(emit_raw_events) + resolved_emission_shape = _resolve_emission_shape(emission_shape) @app.post(path) async def agentic_chat_endpoint(input_data: RunAgentInput, request: Request): @@ -2293,6 +2341,7 @@ async def agentic_chat_endpoint(input_data: RunAgentInput, request: Request): timeout=timeout, checkpoint_kwargs=checkpoint_kwargs, emit_raw_events=resolved_emit_raw_events, + emission_shape=resolved_emission_shape, ), media_type=encoder.get_content_type(), ) @@ -2304,6 +2353,7 @@ def add_crewai_crew_fastapi_endpoint( path: str = "/", *, emit_raw_events: bool | None = None, + emission_shape: str | None = None, ): """Adds a CrewAI crew endpoint to the FastAPI app. @@ -2325,6 +2375,7 @@ def add_crewai_crew_fastapi_endpoint( GLOBAL_EVENT_LISTENER = FastAPICrewFlowEventListener() resolved_emit_raw_events = _resolve_emit_raw_events(emit_raw_events) + resolved_emission_shape = _resolve_emission_shape(emission_shape) _cached_flow = None # Dedicated per-endpoint lock so two concurrent first-requests cannot @@ -2375,6 +2426,7 @@ async def crew_endpoint(input_data: RunAgentInput, request: Request): timeout=timeout, checkpoint_kwargs=checkpoint_kwargs, emit_raw_events=resolved_emit_raw_events, + emission_shape=resolved_emission_shape, ), media_type=encoder.get_content_type(), ) diff --git a/integrations/crew-ai/python/ag_ui_crewai/sdk.py b/integrations/crew-ai/python/ag_ui_crewai/sdk.py index a69abf836..81ace7f5b 100644 --- a/integrations/crew-ai/python/ag_ui_crewai/sdk.py +++ b/integrations/crew-ai/python/ag_ui_crewai/sdk.py @@ -3,6 +3,7 @@ """ import copy +import logging import uuid from dataclasses import dataclass from typing import ( @@ -39,6 +40,8 @@ ) from .utils import yield_control +_LOGGER = logging.getLogger(__name__) + class CopilotKitProperties(BaseModel): """CopilotKit properties""" actions: List[Any] = Field(default_factory=list) @@ -308,13 +311,15 @@ async def _copilotkit_stream_custom_stream_wrapper(response: CustomStreamWrapper flow = flow_context.get(None) message_id: Optional[str] = None - tool_call_id: str = "" content = "" created = 0 model = "" system_fingerprint = "" finish_reason=None - all_tool_calls = [] + # Tool calls keyed by the delta's ``index`` (litellm/OpenAI stream each parallel + # call under its own stable index). Reading only ``tool_calls[0]`` lost parallel + # calls; keying by index keeps them apart. + tool_calls_by_index: dict[Any, dict] = {} async for chunk in response: if message_id is None: @@ -338,36 +343,75 @@ async def _copilotkit_stream_custom_stream_wrapper(response: CustomStreamWrapper # yield control to the event loop await yield_control() - # Stream tool calls - tool_calls = chunk["choices"][0]["delta"]["tool_calls"] or None - tool_call_id = tool_calls[0].id if tool_calls is not None else None - tool_call_arguments = tool_calls[0].function["arguments"] if tool_calls is not None else None - tool_call_name = tool_calls[0].function["name"] if tool_calls is not None else None - - if tool_call_id is not None: - all_tool_calls.append( - { - "id": tool_call_id, - "name": tool_call_name, - "arguments": "", - } + # Stream tool calls. Every entry in the delta is processed, not just the + # first, so parallel calls in one delta are not lost. + for position, tool_call_delta in enumerate( + chunk["choices"][0]["delta"]["tool_calls"] or [] + ): + index = getattr(tool_call_delta, "index", None) + try: + index = position if index is None else int(index) + except (TypeError, ValueError): + index = position + function = getattr(tool_call_delta, "function", None) or {} + delta_id = getattr(tool_call_delta, "id", None) + delta_name = ( + function.get("name") if hasattr(function, "get") + else getattr(function, "name", None) + ) + tool_call_arguments = ( + function.get("arguments") if hasattr(function, "get") + else getattr(function, "arguments", None) ) - # Checked on whichever chunk carries the name (some providers stream the - # tool id and name in separate deltas), not only the id-bearing chunk. - if tool_call_name is not None: - _mark_predicted_tool_streamed(flow, tool_call_name) - - if tool_call_arguments is not None: - # add to the current tool call - all_tool_calls[-1]["arguments"] += tool_call_arguments + # Checked on whichever delta carries the name (some providers stream the + # tool id and name separately), not only the id-bearing one. + if delta_name is not None: + _mark_predicted_tool_streamed(flow, delta_name) + + entry = tool_calls_by_index.get(index) + if entry is None: + if delta_id is None and delta_name is None and tool_call_arguments is None: + # A truly empty delta (a bare ``{"index": n}``) carries nothing: + # creating an entry would manufacture a phantom call. A delta + # with ARGUMENTS but no identity yet DOES open an entry so the + # buffered prefix is not lost when the id/name arrive later. + continue + entry = {"id": delta_id, "name": delta_name, "arguments": ""} + tool_calls_by_index[index] = entry + else: + if delta_id is not None: + entry["id"] = delta_id + if delta_name is not None: + entry["name"] = delta_name + if tool_call_arguments is not None: + entry["arguments"] += tool_call_arguments + + # Emit only once both id and name are known: the triples shaper needs + # both to open a TOOL_CALL_START, and it stamps the ACCUMULATED id on + # every chunk so a provider that splits identity across deltas still + # streams (rather than shipping unattachable id-less chunks). + if entry["id"] is None or entry["name"] is None: + continue + if not entry.get("streamed"): + # First emit for this call. If argument fragments arrived on the + # earlier (identity-less) deltas, stream the ACCUMULATED prefix now, + # not just this delta, so the live TOOL_CALL_ARGS match the final + # ModelResponse. + entry["streamed"] = True + delta_out = entry["arguments"] or None + else: + delta_out = tool_call_arguments crewai_event_bus.emit( flow, BridgedToolCallChunkEvent( type=EventType.TOOL_CALL_CHUNK, - tool_call_id=tool_call_id, - tool_call_name=tool_call_name, - delta=tool_call_arguments, + tool_call_id=entry["id"], + tool_call_name=entry["name"], + delta=delta_out, + # Set, not just forwarded: without a parent, apply/default.ts + # creates a separate assistant message per tool call. + parent_message_id=message_id, ) ) # yield control to the event loop @@ -382,6 +426,16 @@ async def _copilotkit_stream_custom_stream_wrapper(response: CustomStreamWrapper if finish_reason is not None: break + incomplete = [ + e for e in tool_calls_by_index.values() + if e["id"] is None or e["name"] is None + ] + if incomplete: + _LOGGER.error( + "ag-ui-crewai dropped %d incomplete tool call(s) that never received " + "both an id and a name", + len(incomplete), + ) tool_calls = [ ChatCompletionMessageToolCall( function=LiteLLMFunction( @@ -391,7 +445,10 @@ async def _copilotkit_stream_custom_stream_wrapper(response: CustomStreamWrapper id=tool_call["id"], type="function" ) - for tool_call in all_tool_calls + # Insertion order preserves the provider's ordering; keys are + # heterogeneous so are not sortable. + for tool_call in tool_calls_by_index.values() + if tool_call["id"] is not None and tool_call["name"] is not None ] return ModelResponse( id=message_id, diff --git a/integrations/crew-ai/python/tests/test_capability_declaration.py b/integrations/crew-ai/python/tests/test_capability_declaration.py index a42601e66..472a117c7 100644 --- a/integrations/crew-ai/python/tests/test_capability_declaration.py +++ b/integrations/crew-ai/python/tests/test_capability_declaration.py @@ -19,6 +19,7 @@ def _clean_protocol_env(monkeypatch): """Clear the RAW env var: otherwise an exported AGUI_CREWAI_EMIT_RAW_EVENTS makes these tests assert the ambient environment rather than the shipped defaults.""" + monkeypatch.delenv(config_mod.EMISSION_SHAPE_ENV_VAR, raising=False) monkeypatch.delenv(config_mod.EMIT_RAW_EVENTS_ENV_VAR, raising=False) @@ -362,3 +363,65 @@ def kickoff_async(self, inputs=None): # pragma: no cover - never called ep.add_crewai_flow_fastapi_endpoint( FastAPI(), _Flow(), "/flow", emit_raw_events="false" ) + + +def test_emission_shape_resolution_precedence_and_validation(monkeypatch): + """Explicit argument > env var > shipped default (triples); a wrong-typed or + unknown value raises rather than silently mis-shaping the wire.""" + monkeypatch.delenv(config_mod.EMISSION_SHAPE_ENV_VAR, raising=False) + assert config_mod.resolve_emission_shape(None) == "triples" + assert config_mod.resolve_emission_shape("Chunks") == "chunks" + + monkeypatch.setenv(config_mod.EMISSION_SHAPE_ENV_VAR, "chunks") + assert config_mod.resolve_emission_shape(None) == "chunks" + # Explicit argument wins over the env var. + assert config_mod.resolve_emission_shape("triples") == "triples" + + for bad in ("bogus", 123): + with pytest.raises(ValueError): + config_mod.resolve_emission_shape(bad) + + +def test_unrecognised_emission_shape_env_is_warned_not_silently_ignored( + monkeypatch, caplog +): + import logging + + monkeypatch.setenv(config_mod.EMISSION_SHAPE_ENV_VAR, "tripples") + monkeypatch.setattr(config_mod, "_ENV_WARN_SEEN", set()) + with caplog.at_level(logging.WARNING, logger="ag_ui_crewai._config"): + assert config_mod.resolve_emission_shape(None) == "triples" + assert any("tripples" in r.getMessage() for r in caplog.records), caplog.text + + +def test_get_capabilities_reports_the_resolved_wire_shape(): + """The declaration reflects the shape the endpoint will actually emit.""" + triples = get_capabilities()["wireShape"] + assert triples["emissionShape"] == "triples" + assert triples["textMessages"] == [ + "TEXT_MESSAGE_START", "TEXT_MESSAGE_CONTENT", "TEXT_MESSAGE_END" + ] + assert triples["toolCalls"] == [ + "TOOL_CALL_START", "TOOL_CALL_ARGS", "TOOL_CALL_END" + ] + # MCP tool executions are triples regardless of the streaming shape. + assert triples["mcpToolCalls"][0] == "TOOL_CALL_START" + + chunks = get_capabilities(emission_shape="chunks")["wireShape"] + assert chunks["emissionShape"] == "chunks" + assert chunks["textMessages"] == ["TEXT_MESSAGE_CHUNK"] + assert chunks["toolCalls"] == ["TOOL_CALL_CHUNK"] + + with pytest.raises(ValueError): + get_capabilities(emission_shape="bogus") + + +def test_endpoint_factory_rejects_a_bad_emission_shape_at_registration(): + class _Flow: + def kickoff_async(self, inputs=None): # pragma: no cover - never called + raise AssertionError + + with pytest.raises(ValueError): + ep.add_crewai_flow_fastapi_endpoint( + FastAPI(), _Flow(), "/flow", emission_shape="bogus" + ) diff --git a/integrations/crew-ai/python/tests/test_streaming.py b/integrations/crew-ai/python/tests/test_streaming.py index 7a8ef41f9..bf5c3b128 100644 --- a/integrations/crew-ai/python/tests/test_streaming.py +++ b/integrations/crew-ai/python/tests/test_streaming.py @@ -447,9 +447,9 @@ async def kickoff_async(self, inputs=None): # -- translator wire shape (default = chunks) ------------------------------- -def test_translator_produces_current_chunk_wire_shape(): - """The default translator maps bridge/lifecycle events onto exactly the - events the legacy listener produced — chunks, not triples.""" +def test_translator_produces_triples_wire_shape(): + """The default translator maps bridge/lifecycle events onto START/CONTENT/END + triples, closing each open sequence before the next boundary.""" state = {"messages": [{"role": "assistant", "content": "hi", "id": "m1"}]} tr = frames_mod.StreamFrameTranslator( thread_id="t-1", run_id="r-1", state_provider=lambda: state, @@ -460,69 +460,80 @@ def test_translator_produces_current_chunk_wire_shape(): ] assert tr.run_started is True start_ev = tr.translate(_ev("method_execution_started", method_name="chat")) - assert start_ev[0].type == EventType.STEP_STARTED + assert [e.type for e in start_ev] == [EventType.STEP_STARTED] assert start_ev[0].step_name == "chat" text = tr.translate(_ev( "TEXT_MESSAGE_CHUNK", message_id="m1", role="assistant", delta="hi", )) - assert len(text) == 1 - assert text[0].type == EventType.TEXT_MESSAGE_CHUNK - assert (text[0].message_id, text[0].role, text[0].delta) == ("m1", "assistant", "hi") + assert [e.type for e in text] == [ + EventType.TEXT_MESSAGE_START, EventType.TEXT_MESSAGE_CONTENT, + ] + assert text[0].message_id == "m1" and text[0].role == "assistant" + assert text[1].delta == "hi" + # Opening a tool call closes the open text message first. tool = tr.translate(_ev( "TOOL_CALL_CHUNK", tool_call_id="tc1", tool_call_name="searchTool", delta='{"q":1}', )) - assert tool[0].type == EventType.TOOL_CALL_CHUNK - assert (tool[0].tool_call_id, tool[0].tool_call_name, tool[0].delta) == ( - "tc1", "searchTool", '{"q":1}' - ) + assert [e.type for e in tool] == [ + EventType.TEXT_MESSAGE_END, + EventType.TOOL_CALL_START, + EventType.TOOL_CALL_ARGS, + ] + assert (tool[1].tool_call_id, tool[1].tool_call_name) == ("tc1", "searchTool") + assert tool[2].delta == '{"q":1}' + # A side-channel CUSTOM / STATE_SNAPSHOT does NOT close the open tool call. custom = tr.translate(_ev("CUSTOM", name="PredictState", value=[1])) - assert custom[0].type == EventType.CUSTOM - assert (custom[0].name, custom[0].value) == ("PredictState", [1]) - + assert [e.type for e in custom] == [EventType.CUSTOM] snap = tr.translate(_ev("STATE_SNAPSHOT", snapshot={"p": 5})) - assert snap[0].type == EventType.STATE_SNAPSHOT - assert snap[0].snapshot == {"p": 5} + assert [e.type for e in snap] == [EventType.STATE_SNAPSHOT] + assert tr._shaper.open_tool_calls == ("tc1",) + # method_finished closes the open tool call, then the snapshots + STEP_FINISHED. finished = tr.translate(_ev("method_execution_finished", method_name="chat")) assert [e.type for e in finished] == [ - EventType.MESSAGES_SNAPSHOT, EventType.STATE_SNAPSHOT, EventType.STEP_FINISHED, + EventType.TOOL_CALL_END, + EventType.MESSAGES_SNAPSHOT, + EventType.STATE_SNAPSHOT, + EventType.STEP_FINISHED, ] - assert finished[0].messages[0].id == "m1" - assert finished[2].step_name == "chat" + assert finished[1].messages[0].id == "m1" + assert finished[3].step_name == "chat" fin = tr.translate(_ev("flow_finished")) - assert fin[0].type == EventType.RUN_FINISHED + assert [e.type for e in fin] == [EventType.RUN_FINISHED] assert tr.run_finished is True - # Idempotent: a second flow_finished never re-emits RUN_FINISHED. assert tr.translate(_ev("flow_finished")) == [] - # Native crewai channels / unknown events are dropped (behavior-preserving). assert tr.translate(_ev("llm_stream_chunk", chunk="x")) == [] assert tr.translate(_ev("cc_env")) == [] -def test_translator_emission_shape_is_swappable_and_defaults_to_chunks(): - """The emission shape is a single seam defaulting to chunks; the parity - 'triples' shape is a documented NotImplementedError placeholder.""" +def test_translator_emission_shape_defaults_to_triples_with_chunks_opt_out(): + """The wire shape is a single seam: triples by default, chunks on opt-out, + unknown values rejected at construction.""" tr = frames_mod.StreamFrameTranslator( thread_id="t", run_id="r", state_provider=dict, ) - assert tr.emission_shape == "chunks" + assert tr.emission_shape == "triples" with pytest.raises(ValueError): frames_mod.StreamFrameTranslator( thread_id="t", run_id="r", state_provider=dict, emission_shape="bogus", ) - triples = frames_mod.StreamFrameTranslator( - thread_id="t", run_id="r", state_provider=dict, emission_shape="triples", + chunks = frames_mod.StreamFrameTranslator( + thread_id="t", run_id="r", state_provider=dict, emission_shape="chunks", ) - with pytest.raises(NotImplementedError): - triples.translate(_ev("TEXT_MESSAGE_CHUNK", message_id="m", delta="x")) + out = chunks.translate(_ev("TEXT_MESSAGE_CHUNK", message_id="m", delta="x")) + assert [e.type for e in out] == [EventType.TEXT_MESSAGE_CHUNK] + tool = chunks.translate(_ev( + "TOOL_CALL_CHUNK", tool_call_id="c", tool_call_name="fn", delta="{}", + )) + assert [e.type for e in tool] == [EventType.TOOL_CALL_CHUNK] # -- end-to-end through a REAL crewai Flow via astream ---------------------- @@ -548,11 +559,10 @@ async def chat(self): @requires_stream_frames -async def test_frame_path_end_to_end_matches_legacy_wire_shape(): +async def test_frame_path_end_to_end_emits_triples(): """Driving a real Flow through the StreamFrame path yields RUN_STARTED, - STEP_STARTED, per-delta TEXT_MESSAGE_CHUNK, a CUSTOM, then - MESSAGES/STATE snapshot + STEP_FINISHED + RUN_FINISHED — the same wire - shape the legacy listener produced.""" + STEP_STARTED, a TEXT_MESSAGE_START / _CONTENT+ / _END triple, a CUSTOM, then + MESSAGES/STATE snapshot + STEP_FINISHED + RUN_FINISHED.""" from ag_ui.encoder import EventEncoder flow = _BridgeEmittingFlow() @@ -569,20 +579,25 @@ async def test_frame_path_end_to_end_matches_legacy_wire_shape(): assert types[0] == "RUN_STARTED" assert types[-1] == "RUN_FINISHED" - assert types.count("TEXT_MESSAGE_CHUNK") == 2 + # One text message: a single START/END pair around two CONTENT deltas. + assert types.count("TEXT_MESSAGE_START") == 1 + assert types.count("TEXT_MESSAGE_END") == 1 + assert types.count("TEXT_MESSAGE_CONTENT") == 2 + assert "TEXT_MESSAGE_CHUNK" not in types assert "STEP_STARTED" in types assert "STEP_FINISHED" in types assert "CUSTOM" in types assert "MESSAGES_SNAPSHOT" in types - # No START/CONTENT/END triples — behavior-preserving chunk shape. - assert not any(t.endswith("_START") or t.endswith("_END") for t in types - if t not in ("RUN_STARTED",)) # Correlation ids are stamped, not the listener's "?" placeholders. run_started = next(p for p in payloads if p["type"] == "RUN_STARTED") assert run_started["threadId"] == "t-1" assert run_started["runId"] == "r-1" - text_deltas = [p["delta"] for p in payloads if p["type"] == "TEXT_MESSAGE_CHUNK"] + text_deltas = [ + p["delta"] for p in payloads if p["type"] == "TEXT_MESSAGE_CONTENT" + ] assert text_deltas == ["Hello ", "world"] + # The text message closes before the run ends. + assert types.index("TEXT_MESSAGE_END") < types.index("RUN_FINISHED") def _make_run_input(thread_id="t-1", run_id="r-1"): @@ -662,10 +677,10 @@ async def test_frame_path_two_completions_emit_single_run_lifecycle(): assert types[-1] == "RUN_FINISHED" # The follow-up text reaches the client, inside the run. - assert "TEXT_MESSAGE_CHUNK" in types, types - follow = next(p for p in payloads if p["type"] == "TEXT_MESSAGE_CHUNK") + assert "TEXT_MESSAGE_CONTENT" in types, types + follow = next(p for p in payloads if p["type"] == "TEXT_MESSAGE_CONTENT") assert follow["delta"] == "Crew is done." - assert types.index("TEXT_MESSAGE_CHUNK") < types.index("RUN_FINISHED") + assert types.index("TEXT_MESSAGE_CONTENT") < types.index("RUN_FINISHED") # -- Review invariants: raw-payload fidelity, nested non-leak, terminal @@ -1431,3 +1446,221 @@ def test_mapped_events_are_never_duplicated_as_raw(): assert frames_mod.is_recognized_event(_ev("flow_started")) is True assert frames_mod.is_recognized_event(_ev("TEXT_MESSAGE_CHUNK")) is True assert frames_mod.is_recognized_event(_ev("llm_thinking_chunk")) is False + + +# -- Wire-shape (triples) state-machine invariants -------------------------- + +def _shaper(shape="triples"): + return frames_mod.EmissionShaper(shape, thread_id="t", run_id="r") + + +def test_parallel_tool_calls_stay_separate_and_close_in_order(): + """crewai streams parallel calls; each id gets its own START/ARGS, and a flush + closes them innermost-first. A single-slot model mis-attributed the arguments.""" + sh = _shaper() + a = sh.tool(_ev("TOOL_CALL_CHUNK", tool_call_id="a", tool_call_name="fa", delta='{"x":')) + b = sh.tool(_ev("TOOL_CALL_CHUNK", tool_call_id="b", tool_call_name="fb", delta='{"y":')) + assert [e.type for e in a] == [EventType.TOOL_CALL_START, EventType.TOOL_CALL_ARGS] + assert [e.type for e in b] == [EventType.TOOL_CALL_START, EventType.TOOL_CALL_ARGS] + a2 = sh.tool(_ev("TOOL_CALL_CHUNK", tool_call_id="a", delta='1}')) + assert [e.type for e in a2] == [EventType.TOOL_CALL_ARGS] and a2[0].tool_call_id == "a" + ends = sh.flush() + assert [(e.type, e.tool_call_id) for e in ends] == [ + (EventType.TOOL_CALL_END, "b"), (EventType.TOOL_CALL_END, "a"), + ] + + +def test_side_channel_events_do_not_close_open_tool_calls(): + """A STATE_SNAPSHOT / CUSTOM between two argument deltas must NOT close the + call: litellm stamps the id on the first delta only, so a reopened call could + not carry its identity and its arguments would truncate.""" + sh = _shaper() + sh.tool(_ev("TOOL_CALL_CHUNK", tool_call_id="a", tool_call_name="fa", delta='{"x":')) + # A side-channel event reshaped through the shaper leaves the tool call open. + assert [e.type for e in sh.reshape(_ev("STATE_SNAPSHOT", snapshot={}))] == [ + EventType.STATE_SNAPSHOT + ] + assert sh.open_tool_calls == ("a",) + cont = sh.tool(_ev("TOOL_CALL_CHUNK", tool_call_id="a", delta='1}')) + assert [e.type for e in cont] == [EventType.TOOL_CALL_ARGS] + + +def test_late_delta_for_a_closed_tool_call_is_dropped_not_reopened(): + """Reopening a closed id would emit a second TOOL_CALL_START for it, which the + client turns into a duplicate tool call.""" + sh = _shaper() + sh.tool(_ev("TOOL_CALL_CHUNK", tool_call_id="a", tool_call_name="fa", delta="{}")) + sh.flush() # closes "a" + assert sh.tool(_ev("TOOL_CALL_CHUNK", tool_call_id="a", delta="!")) == [] + + +def test_switching_from_a_tool_call_to_text_closes_the_call(): + sh = _shaper() + sh.tool(_ev("TOOL_CALL_CHUNK", tool_call_id="a", tool_call_name="fa", delta="{}")) + out = sh.text(_ev("TEXT_MESSAGE_CHUNK", message_id="m", delta="hi")) + assert [e.type for e in out] == [ + EventType.TOOL_CALL_END, + EventType.TEXT_MESSAGE_START, + EventType.TEXT_MESSAGE_CONTENT, + ] + + +def test_run_finished_reshape_closes_an_open_message_first(): + """The shaper owns the message / tool lifecycle: a RUN_FINISHED reshaped on the + legacy path closes an open text message before the terminal. (Closing open + STEPS before RUN_FINISHED is the attribution tracker's job, exercised via the + translator in the e2e tests.)""" + sh = _shaper() + sh.text(_ev("TEXT_MESSAGE_CHUNK", message_id="m", delta="partial")) + from ag_ui.core import RunFinishedEvent + out = sh.reshape(RunFinishedEvent( + type=EventType.RUN_FINISHED, thread_id="t", run_id="r")) + assert [e.type for e in out] == [ + EventType.TEXT_MESSAGE_END, EventType.RUN_FINISHED, + ] + + +def test_chunks_opt_out_is_pure_passthrough(): + sh = _shaper("chunks") + t = sh.text(_ev("TEXT_MESSAGE_CHUNK", message_id="m", delta="x")) + assert [e.type for e in t] == [EventType.TEXT_MESSAGE_CHUNK] + assert sh.flush() == [] + from ag_ui.core import RunFinishedEvent + rf = sh.reshape(RunFinishedEvent( + type=EventType.RUN_FINISHED, thread_id="t", run_id="r")) + assert [e.type for e in rf] == [EventType.RUN_FINISHED] + + +def test_both_transports_emit_identical_triples_for_one_stream(): + """The whole point of the flip: shape does not depend on the transport. Drive + the same logical stream through the frame translator and the legacy reshaper + and assert identical output.""" + logical = [ + ("flow_started", {}), + ("method_execution_started", {"method_name": "chat"}), + ("TEXT_MESSAGE_CHUNK", {"message_id": "m1", "role": "assistant", "delta": "Hi"}), + ("TOOL_CALL_CHUNK", {"tool_call_id": "c1", "tool_call_name": "fn", "delta": "{}"}), + ("method_execution_finished", {"method_name": "chat"}), + ("flow_finished", {}), + ] + # frame path: raw events through the translator + tr = frames_mod.StreamFrameTranslator( + thread_id="t", run_id="r", state_provider=lambda: {"messages": []}, + ) + frame_types = [] + for t, kw in logical: + frame_types += [e.type for e in tr.translate(_ev(t, **kw))] + + # legacy path: the listener's WIRE events through the shaper's reshape + from ag_ui.core import ( + RunStartedEvent, RunFinishedEvent, StepStartedEvent, StepFinishedEvent, + MessagesSnapshotEvent, StateSnapshotEvent, + ) + from ag_ui.core.events import TextMessageChunkEvent, ToolCallChunkEvent + wire = [ + RunStartedEvent(type=EventType.RUN_STARTED, thread_id="t", run_id="r"), + StepStartedEvent(type=EventType.STEP_STARTED, step_name="chat"), + TextMessageChunkEvent(type=EventType.TEXT_MESSAGE_CHUNK, message_id="m1", role="assistant", delta="Hi"), + ToolCallChunkEvent(type=EventType.TOOL_CALL_CHUNK, tool_call_id="c1", tool_call_name="fn", delta="{}"), + MessagesSnapshotEvent(type=EventType.MESSAGES_SNAPSHOT, messages=[]), + StateSnapshotEvent(type=EventType.STATE_SNAPSHOT, snapshot={}), + StepFinishedEvent(type=EventType.STEP_FINISHED, step_name="chat"), + RunFinishedEvent(type=EventType.RUN_FINISHED, thread_id="t", run_id="r"), + ] + sh = _shaper() + legacy_types = [] + for e in wire: + legacy_types += [x.type for x in sh.reshape(e)] + + assert frame_types == legacy_types, (frame_types, legacy_types) + assert EventType.TEXT_MESSAGE_START in frame_types + assert EventType.TOOL_CALL_END in frame_types + + +# -- R1 review fixes: double-emit, buffered args, None-id message ------------ + +def test_mapped_events_are_not_double_emitted_as_raw(): + """is_recognized_event must cover every type translate() maps — including the + crew/agent lifecycle, method_failed, and MCP — or emit_raw_events mirrors them + as RAW alongside the translated STEP/TOOL/CUSTOM events.""" + for t in ( + "flow_started", "flow_finished", "method_execution_started", + "method_execution_finished", "method_execution_failed", + "crew_kickoff_started", "crew_kickoff_completed", "agent_execution_started", + "agent_execution_completed", + ): + assert frames_mod.is_recognized_event(_ev(t)) is True, t + # MCP events are recognized dynamically. + from ag_ui_crewai import mcp as mcp_mod + mcp_ev = _ev("mcp_tool_execution_started") + assert frames_mod.is_recognized_event(mcp_ev) is mcp_mod.is_mcp_event(mcp_ev) + # An unmapped native event is still eligible for RAW. + assert frames_mod.is_recognized_event(_ev("llm_stream_chunk")) is False + + +async def test_tool_args_streamed_before_identity_are_flushed_when_identity_arrives(): + """A provider that streams argument fragments before the tool id/name must not + lose the prefix on the wire: the first emit sends the accumulated arguments, + and the final ModelResponse agrees.""" + flow_context.set(None) + emitted = [] + from ag_ui_crewai._capabilities import crewai_event_bus + with crewai_event_bus.scoped_handlers(): + @crewai_event_bus.on(BridgedToolCallChunkEvent) + def _on(source, event): # pylint: disable=unused-argument + emitted.append(event) + + def d(index, call_id, name, args): + item = _tool_call_delta(call_id=call_id, name=name, arguments=args) + item.index = index + return item + + async def _gen(): + # args first, identity later — all under index 0 + yield _stream_chunk("msg-1", tool_calls=[d(0, None, None, '{"q":')]) + yield _stream_chunk("msg-1", tool_calls=[d(0, "c1", "searchTool", '1}')]) + yield _stream_chunk("msg-1", finish_reason="tool_calls") + + resp = await copilotkit_stream(_FakeStreamWrapper(_gen())) + await _settle_bus() + + streamed = "".join(e.delta for e in emitted if e.delta) + assert streamed == '{"q":1}', [e.delta for e in emitted] + assert all(e.tool_call_id == "c1" for e in emitted) + assert resp.choices[0].message.tool_calls[0].function.arguments == '{"q":1}' + + +def test_shaper_text_with_none_message_id_opens_once_and_closes(): + """A chunk stream with no message_id must still produce exactly one + START/CONTENT/END, not a fresh START per delta with no END.""" + sh = _shaper() + a = sh.text(_ev("TEXT_MESSAGE_CHUNK", message_id=None, delta="one")) + b = sh.text(_ev("TEXT_MESSAGE_CHUNK", message_id=None, delta="two")) + assert [e.type for e in a] == [EventType.TEXT_MESSAGE_START, EventType.TEXT_MESSAGE_CONTENT] + assert [e.type for e in b] == [EventType.TEXT_MESSAGE_CONTENT] + assert [e.type for e in sh.flush()] == [EventType.TEXT_MESSAGE_END] + + +def test_parent_message_id_is_preserved_on_every_path_and_shape(): + """A tool call must attach to the streamed assistant message via + parent_message_id on both transports and in both shapes; dropping it makes the + client render the call as a separate message.""" + chunk_ev = _ev( + "TOOL_CALL_CHUNK", tool_call_id="c1", tool_call_name="fn", delta="{}", + parent_message_id="m1", + ) + # triples: TOOL_CALL_START carries it + start = _shaper("triples").tool(chunk_ev)[0] + assert start.type == EventType.TOOL_CALL_START and start.parent_message_id == "m1" + # chunks: the passthrough chunk carries it + chunk = _shaper("chunks").tool(chunk_ev)[0] + assert chunk.type == EventType.TOOL_CALL_CHUNK and chunk.parent_message_id == "m1" + # legacy listener rebuild forwards it (so the shaper can stamp it downstream) + from ag_ui.core.events import ToolCallChunkEvent + rebuilt = ToolCallChunkEvent( + type=EventType.TOOL_CALL_CHUNK, tool_call_id="c1", tool_call_name="fn", + delta="{}", parent_message_id="m1", + ) + reshaped_start = _shaper("triples").reshape(rebuilt)[0] + assert reshaped_start.type == EventType.TOOL_CALL_START + assert reshaped_start.parent_message_id == "m1"