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"