diff --git a/docs/CONSUMING.md b/docs/CONSUMING.md index 87de154..1fbbe93 100644 --- a/docs/CONSUMING.md +++ b/docs/CONSUMING.md @@ -822,7 +822,14 @@ learns the one it was given from `RUN_STARTED`, which is always the first event. read; the rest of `messages` is ignored and the checkpointer's transcript is the truth. That diverges from AG-UI's client-is-authoritative convention on purpose — the values session state exists to keep out of the model's context would -otherwise have to live in the browser and be posted back every turn. +otherwise have to live in the browser and be posted back every turn. The stream +says so rather than leaving a client to discover it: a `MESSAGES_SNAPSHOT` +closes every run with the thread as the server holds it. It closes rather than +opens the run because a snapshot drops every local message it does not name, and +one sent before the turn was checkpointed would take the user's own question off +their screen. Ids line up — the question keeps the client's `id`, the answer +carries the id the thread will store — so a client reconciles in place rather +than rebuilding its list. **The read routes are what the stream deliberately leaves out.** A `STATE_SNAPSHOT` carries `{kind, tool, bytes}` per key and never the payload, so @@ -902,7 +909,9 @@ nothing else. `STATE_SNAPSHOT` carries metadata only — `kind`, `tool`, `bytes`, and `seq` once known — never the stored value, which a frontend fetches when it actually wants -to draw it. `seq` is assigned when a write is merged, so mid-turn snapshots omit +to draw it. It sits under `toolState` inside AG-UI's state object, so read +`snapshot.toolState`; the rest of that object is the client's, to hold whatever +state of its own it wants to keep there. `seq` is assigned when a write is merged, so mid-turn snapshots omit it and the snapshot closing the turn carries it. Snapshots are also **partial mid-turn**: one names what its node wrote, so a client merges them into what it holds rather than replacing. diff --git a/examples/agui-events/README.md b/examples/agui-events/README.md index 939d756..7890eff 100644 --- a/examples/agui-events/README.md +++ b/examples/agui-events/README.md @@ -78,14 +78,15 @@ reader can only demonstrate that a format is self-consistent. | `RUN_STARTED` / `RUN_FINISHED` / `RUN_ERROR` | run lifecycle | | `TEXT_MESSAGE_START` / `CONTENT` / `END` | the answer, streamed | | `TOOL_CALL_START` / `ARGS` / `END` / `RESULT` | the tool-call lifecycle | +| `MESSAGES_SNAPSHOT` | the thread as the server holds it, closing each run | | `STATE_SNAPSHOT` | the state channel | | `ACTIVITY_SNAPSHOT` | a first-class message role in AG-UI, which is what receipts ride | Those are the only event types this server emits, out of the 33 AG-UI defines. -No `STATE_DELTA`, no `MESSAGES_SNAPSHOT`, no `STEP_*`, `REASONING_*`, `CUSTOM` -or `RAW`. Snapshots only, each complete in itself, because a delta for an -unknown `messageId` is dropped silently by the client and would make the wire -depend on a patch having applied. +No `STATE_DELTA`, no `STEP_*`, `REASONING_*`, `CUSTOM` or `RAW`. Snapshots only, +each complete in itself, because a delta for an unknown `messageId` is dropped +silently by the client and would make the wire depend on a patch having +applied. ### What a consumer has to know that the protocol does not tell it @@ -107,6 +108,10 @@ what session state exists to keep out of the conversation. Fetching one is `GET /threads/{id}/state/{key}`, which is outside the protocol entirely. A stock client showing "state" will show sizes and kinds and think it has everything. +The metadata sits under a **`toolState`** key rather than at the root of the +state object, so the rest of that object stays the client's own — read +`snapshot.toolState`, not `snapshot`. + **History is the server's.** AG-UI's convention is client-authoritative: `RunAgentInput.messages` is the conversation, and the client owns it. `HttpAgent` duly posts its whole array every turn — and this server reads the @@ -114,7 +119,39 @@ trailing user message and discards the rest, because the values kept out of the model's context would otherwise have to live in the browser and be posted back. The consequence for a consumer is concrete: **mutating `agent.messages` does not edit the thread.** Editing a message, branching, or dropping a turn are -client-side illusions here. `GET /threads/{id}` is the truth. +client-side illusions here. + +The server now says so on the wire: a **`MESSAGES_SNAPSHOT`** closes every +run, carrying the thread as the server holds it, so a client that had diverged +is corrected rather than drifting. + +It closes the run rather than opening it, and that matters. A snapshot is +applied by dropping every local message it does not name — sent up front, before +the turn is checkpointed, it would take the question the user just typed off the +screen and leave the answer under nothing. At the end everything the turn +produced is in the thread. + +**Ids line up on purpose.** The question keeps the `id` the client gave it, and +the answer is labelled with the id the thread will store, taken off the +provider's own stream. So the snapshot reconciles a client's list in place +instead of dropping every message and re-appending the server's — which would +leave activities stranded at the top. One caveat the server enforces for you: an +id the thread already holds is not reused, because the message reducer matches on +id and would replace that message rather than add one. + +`GET /threads/{id}` remains the way to read a thread without running one — and +this client now uses it. The thread id is in the URL as `?thread=`, so +**reloading the page brings the conversation back** rather than starting a fresh +one. Two routes rebuild it: `/threads/{id}` for the transcript and +`/threads/{id}/turns` for what state held at the end of each turn, joined on the +question, since the stream carries no turn boundary a reloaded client could have +seen. + +What does *not* come back is the annotation. Receipts, views and citations are +activity messages, and the server does not rebuild past turns' activities — so a +restored thread shows what was said and what is in session state, but not where +a tool's arguments came from, and the cross-highlighting is empty until the next +turn. **Views are a second protocol.** `mcp.view` names a `ui://` URI and carries the tool's structured content; the HTML comes from `GET /views/{toolset}/{view}`, diff --git a/examples/agui-events/web/src/agui.ts b/examples/agui-events/web/src/agui.ts index 77f6caa..79c3482 100644 --- a/examples/agui-events/web/src/agui.ts +++ b/examples/agui-events/web/src/agui.ts @@ -48,11 +48,31 @@ export async function readState( return (await response.json()) as StateValue; } +/** A thread's messages, for a client that reloaded. + * + * The conversation, and nothing around it: receipts, views and the rest are + * activities, and the server does not rebuild past turns' activities. So a + * restored thread shows what was said but not where each tool's arguments came + * from — see the README. + */ +export async function readThread(threadId: string) { + const response = await fetch(`/api/threads/${threadId}`); + // 404 is the ordinary answer for a thread id that has never run, which is + // what a hand-edited URL produces. The caller starts fresh instead. + if (response.status === 404) return null; + if (!response.ok) throw new Error(`${response.status}`); + return (await response.json()) as { + threadId: string; + messages: { id: string; role: string; content?: string | null }[]; + state: Record; + }; +} + /** A thread's turns, and what session state held at the end of each. * - * Not used by the live client — it builds turns from the events as they - * arrive — but this is how a client that reloaded would get them back, and it - * is the route the panel's per-turn view is really made of. + * The live client builds turns from the events as they arrive; this is how one + * that reloaded gets them back, and it is the route the panel's per-turn view + * is really made of. */ export async function readTurns(threadId: string) { const response = await fetch(`/api/threads/${threadId}/turns`); diff --git a/examples/agui-events/web/src/chat.tsx b/examples/agui-events/web/src/chat.tsx index 4ec7f5a..6fea797 100644 --- a/examples/agui-events/web/src/chat.tsx +++ b/examples/agui-events/web/src/chat.tsx @@ -2,7 +2,7 @@ import { HttpAgent, type Message } from "@ag-ui/client"; import { useEffect, useMemo, useRef, useState } from "react"; import Markdown from "react-markdown"; -import { readState } from "./agui"; +import { readState, readThread, readTurns } from "./agui"; /** Session state as the stream describes it: no payloads, one line per key. */ type StateEntry = { @@ -15,6 +15,13 @@ type StateEntry = { /** Every key the thread holds, which is what `STATE_SNAPSHOT` carries. */ type Snapshot = Record; +/** + * Key inside AG-UI's `state` object holding session-state metadata. The rest + * of that object belongs to the client, so this reads the one key rather than + * treating the whole snapshot as the agent's. + */ +const TOOL_STATE = "toolState"; + /** Where a key came from, read off the `state.published` that announced it. */ type Origin = { toolCallId: string; tool: string; activityId: string }; @@ -171,9 +178,16 @@ export function Chat() { // `RunAgentInput`, runs the SSE through `verifyEvents`, and applies each // event to `messages` and `state`. If this server emitted anything the // protocol disallows, the run would fail here rather than render wrongly. + // `?thread=` if the URL names one, so a reload comes back to the same + // conversation rather than a fresh one — the thread lives in the + // checkpointer, and the id is the only thing a client needs to keep. + const [threadId] = useState( + () => + new URLSearchParams(location.search).get("thread") || crypto.randomUUID(), + ); const agent = useMemo( - () => new HttpAgent({ url: "/api/runs", threadId: crypto.randomUUID() }), - [], + () => new HttpAgent({ url: "/api/runs", threadId }), + [threadId], ); const log = useRef(null); const [messages, setMessages] = useState([]); @@ -209,6 +223,59 @@ export function Chat() { log.current?.scrollTo({ top: log.current.scrollHeight }); }, [messages]); + // Put the thread in the URL, so reloading the page restores it. Replace + // rather than push: this is not a navigation, and a back button that stepped + // through thread ids would be nonsense. + useEffect(() => { + const url = new URL(location.href); + if (url.searchParams.get("thread") === threadId) return; + url.searchParams.set("thread", threadId); + history.replaceState(null, "", url); + }, [threadId]); + + /** Rebuild the conversation from the thread id alone. + * + * Two routes, because the stream has no turn boundary a reloaded client + * could have seen: `/threads/{id}` is the transcript, `/threads/{id}/turns` + * is what state held at the end of each turn. They are joined on the + * question — turn *n* starts at the *n*th user message. + * + * **Activities do not come back.** Receipts, views and citations are + * activity messages, and the server does not rebuild past turns' — so a + * restored thread shows what was said and what is in state, but not where a + * tool's arguments came from. `published` is empty for the same reason: it + * is read off `state.published`, which is an activity. + */ + useEffect(() => { + let cancelled = false; + (async () => { + const thread = await readThread(threadId).catch(() => null); + if (cancelled || !thread || thread.messages.length === 0) return; + const past = await readTurns(threadId).catch(() => null); + if (cancelled) return; + + const questions = thread.messages + .map((message, index) => ({ message, index })) + .filter(({ message }) => message.role === "user"); + const restored: Turn[] = questions.map(({ message, index }, n) => ({ + n: n + 1, + question: message.content || "", + questionId: message.id, + from: index, + state: (past?.history[n]?.state ?? {}) as Snapshot, + published: {}, + })); + + agent.setMessages(thread.messages as unknown as Message[]); + setMessages([...agent.messages]); + setTurns(restored); + setShowing(Math.max(restored.length - 1, 0)); + })(); + return () => { + cancelled = true; + }; + }, [agent, threadId]); + /** Put the question that started a turn at the top of the log. * * `scrollTo` on the log rather than `scrollIntoView` on the message, which @@ -293,8 +360,10 @@ export function Chat() { setMessages([...messages]); patch((turn) => ({ ...turn, published: origins(messages, turn.from) })); }, - // Session state is ours, not AG-UI's `state` — see the README. It - // arrives on the standard channel carrying `{kind, tool, bytes, seq}`. + // Session state arrives on AG-UI's standard `state` channel, under + // `toolState` — the rest of that object is the client's, so read the + // one key rather than the whole snapshot. Each entry carries + // `{kind, tool, bytes, seq}`; see the README. // // Merged rather than assigned: a mid-turn snapshot is built from what // that node wrote, so it names only those keys, and assigning it would @@ -306,7 +375,12 @@ export function Chat() { onStateSnapshotEvent: ({ event }) => { patch((turn) => ({ ...turn, - state: { ...turn.state, ...((event.snapshot ?? {}) as Snapshot) }, + state: { + ...turn.state, + ...(((event.snapshot as Record | undefined)?.[ + TOOL_STATE + ] ?? {}) as Snapshot), + }, })); }, }); diff --git a/src/mcp_agent/streaming.py b/src/mcp_agent/streaming.py index c3ce971..546ddb2 100644 --- a/src/mcp_agent/streaming.py +++ b/src/mcp_agent/streaming.py @@ -43,9 +43,18 @@ @dataclass class AnswerChunk: - """A piece of the model's answer, as it was generated.""" + """A piece of the model's answer, as it was generated. + + ``id`` is the id the completed message will carry in the thread — the + provider's own, taken off the chunk and unchanged by the checkpointer. A + host labelling the streamed message with it names the same message the + thread does, so a later readback lines up with what was rendered instead of + looking like a different message with the same words. ``None`` where the + chunk carried none, and a host then has to mint one. + """ text: str + id: str | None = None @dataclass @@ -57,11 +66,16 @@ class ToolStarted: first point the arguments are known to be whole. A parameter filled from state is absent here — that is the point of it — and shows up in the matching :class:`ToolFinished`'s ``received``. + + ``id`` is the tool call's; ``message_id`` is the assistant message that + asked for it, as the thread stores it. See :class:`AnswerChunk` for why a + host wants the thread's own ids. """ id: str name: str arguments: dict[str, Any] + message_id: str | None = None @dataclass @@ -74,11 +88,15 @@ class ToolFinished: ``published`` maps a field of the tool's own return to the key it was stored under. ``artifact`` is the whole tool artifact, which a host needs to build view props. + + ``id`` is the tool call this answers; ``message_id`` is the tool message + itself, as the thread stores it. """ id: str name: str content: str + message_id: str | None = None artifact: Any = None received: dict[str, Receipt] = field(default_factory=dict) published: dict[str, str] = field(default_factory=dict) @@ -123,11 +141,13 @@ def _published(artifact: Any) -> dict[str, str]: def _tool_calls(message: BaseMessage) -> list[ToolStarted]: + parent = getattr(message, "id", None) return [ ToolStarted( id=str(call.get("id") or ""), name=str(call.get("name") or ""), arguments=dict(call.get("args") or {}), + message_id=str(parent) if parent else None, ) for call in getattr(message, "tool_calls", None) or [] ] @@ -135,8 +155,10 @@ def _tool_calls(message: BaseMessage) -> list[ToolStarted]: def _tool_result(message: BaseMessage) -> ToolFinished: artifact = getattr(message, "artifact", None) + identifier = getattr(message, "id", None) return ToolFinished( id=str(getattr(message, "tool_call_id", "") or ""), + message_id=str(identifier) if identifier else None, name=str(getattr(message, "name", "") or ""), # ``.text`` for the same reason the answer uses it below: an MCP server # returning content blocks makes ``.content`` a list, and ``str()`` on @@ -149,8 +171,8 @@ def _tool_result(message: BaseMessage) -> ToolFinished: ) -def _is_token(mode: str, payload: Any) -> str | None: - """The answer text in a message-channel payload, or ``None``. +def _is_token(mode: str, payload: Any) -> AnswerChunk | None: + """The answer chunk in a message-channel payload, or ``None``. A tool result reaches this channel as well as the update channel, and it carries content — so the node and the chunk type both have to agree before @@ -164,7 +186,10 @@ def _is_token(mode: str, payload: Any) -> str | None: if (metadata or {}).get("langgraph_node") != MODEL_NODE: return None text = str(getattr(chunk, "text", "") or "") - return text or None + if not text: + return None + identifier = getattr(chunk, "id", None) + return AnswerChunk(text, str(identifier) if identifier else None) async def stream_turn( @@ -172,6 +197,7 @@ async def stream_turn( text: str, thread_id: str, config: dict[str, Any] | None = None, + message_id: str | None = None, ) -> AsyncIterator[TurnEvent]: """Run one chat turn on ``thread_id``, yielding each part as it arrives. @@ -183,6 +209,15 @@ async def stream_turn( ``config`` is the runnable config, for a host attaching per-turn callbacks or metadata; ``thread_id`` is merged into its ``configurable`` and wins over any set there. + + ``message_id`` labels the message this turn adds. A host whose client + already has an id for the question should pass it, so the thread and the + client agree on what to call it; left unset, the checkpointer assigns one + and the client has no way to match its own copy to what a readback returns. + An id the thread already holds is **ignored**: LangGraph's message reducer + matches on id, so reusing one replaces that message instead of adding this + one — a client that numbers its messages per session, or retries with the + same id, would silently rewrite its own history. """ merged = dict(config or {}) merged["configurable"] = { @@ -190,20 +225,27 @@ async def stream_turn( "thread_id": thread_id, } - # The thread is read before the turn only to know where this turn's messages + # The thread is read before the turn to know where this turn's messages # begin, exactly as run_turn does; it is cheap next to the model call. before = await agent.aget_state(cast(Any, merged)) - seen = len((getattr(before, "values", None) or {}).get("messages") or []) + existing: list[BaseMessage] = (getattr(before, "values", None) or {}).get( + "messages" + ) or [] + seen = len(existing) + if message_id is not None and any( + getattr(message, "id", None) == message_id for message in existing + ): + message_id = None last_ai: BaseMessage | None = None running: dict[str, StateEntry] = {} async for mode, payload in agent.astream( - cast(Any, {"messages": [HumanMessage(text)]}), + cast(Any, {"messages": [HumanMessage(text, id=message_id)]}), cast(Any, merged), stream_mode=["updates", "messages"], ): if (token := _is_token(mode, payload)) is not None: - yield AnswerChunk(token) + yield token continue if mode != "updates" or not isinstance(payload, dict): continue diff --git a/src/mcp_agent_api/__init__.py b/src/mcp_agent_api/__init__.py index 0032fa6..6715bb9 100644 --- a/src/mcp_agent_api/__init__.py +++ b/src/mcp_agent_api/__init__.py @@ -28,6 +28,7 @@ ANSWER_CITATIONS, MCP_VIEW, STATE_CONSUMED, + STATE_NAMESPACE, STATE_PUBLISHED, TOOLS_WITHHELD, agui_events, @@ -49,6 +50,7 @@ "ANSWER_CITATIONS", "MCP_VIEW", "STATE_CONSUMED", + "STATE_NAMESPACE", "STATE_PUBLISHED", "TOOLS_WITHHELD", "Built", diff --git a/src/mcp_agent_api/events.py b/src/mcp_agent_api/events.py index 016348a..5ae33d5 100644 --- a/src/mcp_agent_api/events.py +++ b/src/mcp_agent_api/events.py @@ -37,13 +37,15 @@ """ import json -from collections.abc import AsyncIterator, Mapping, Sequence +from collections.abc import AsyncIterator, Awaitable, Callable, Mapping, Sequence from dataclasses import asdict from typing import Any from ag_ui.core import ( ActivitySnapshotEvent, BaseEvent, + Message, + MessagesSnapshotEvent, RunErrorEvent, RunFinishedEvent, RunStartedEvent, @@ -79,6 +81,25 @@ MCP_VIEW = "mcp.view" ANSWER_CITATIONS = "answer.citations" +#: Key inside AG-UI's ``state`` object under which every ``STATE_SNAPSHOT`` +#: here carries session-state metadata. +#: +#: The object is the *client's* as much as ours — the protocol has it hold +#: whatever a client and agent share, and ``STATE_DELTA`` patches it by JSON +#: Pointer. Writing our map at the root leaves a client nowhere to keep its own +#: keys, and overwrites any it sent. Under a namespace both fit, and a patch +#: has a stable path to address. +STATE_NAMESPACE = "toolState" + + +def _state_snapshot(state: Mapping[str, StateEntry] | None) -> StateSnapshotEvent: + """One ``STATE_SNAPSHOT``, with the metadata under :data:`STATE_NAMESPACE`. + + The read routes serve :func:`state_metadata` unwrapped: there the body is + ours alone, and there is nothing to share it with. + """ + return StateSnapshotEvent(snapshot={STATE_NAMESPACE: state_metadata(state)}) + def state_metadata(state: Mapping[str, StateEntry] | None) -> dict[str, Any]: """``tool_state`` with the values replaced by what describes them. @@ -202,6 +223,7 @@ async def agui_events( run_id: str, tools: Mapping[str, BaseTool] | None = None, withheld: Sequence[Unsatisfiable] = (), + history: Callable[[], Awaitable[Sequence[Message]]] | None = None, ) -> AsyncIterator[BaseEvent]: """Map one turn onto AG-UI, in the order a client can render. @@ -210,6 +232,23 @@ async def agui_events( once at the top of the run so a client can explain a capability it does not have rather than appearing to ignore the request. + ``history`` is awaited once the turn has finished, and what it returns goes + out as a ``MESSAGES_SNAPSHOT`` before ``RUN_FINISHED``. History here is the + server's — a client posts its whole array and only the trailing user + message is read — so one that edited its own copy is otherwise wrong with + nothing on the wire to say so. + + It is deliberately the *end* of the run and not the start. A client renders + its question the moment it is typed, and a snapshot is applied by dropping + every local message the snapshot does not name: sent up front, before this + turn is checkpointed, it would take the question off the screen and leave + the answer under nothing. At the end everything this turn produced is in + the thread, and — because the ids match, see ``AnswerChunk.id`` — the + client reconciles in place rather than rebuilding its list. + + Omit it and no snapshot is sent, which is what a consumer driving this + without a checkpointer wants. + Exceptions from the turn become ``RUN_ERROR`` and end the stream: a client that opened an SSE connection gets told, rather than watching it close. """ @@ -265,7 +304,9 @@ def ready() -> list[BaseEvent]: open_message = None calls[event.id] = event yield ToolCallStartEvent( - tool_call_id=event.id, tool_call_name=event.name + tool_call_id=event.id, + tool_call_name=event.name, + parent_message_id=event.message_id, ) yield ToolCallArgsEvent( tool_call_id=event.id, @@ -275,7 +316,7 @@ def ready() -> list[BaseEvent]: case ToolFinished(): yield ToolCallResultEvent( - message_id=next_id("tool"), + message_id=event.message_id or next_id("tool"), tool_call_id=event.id, content=event.content, ) @@ -308,7 +349,7 @@ def ready() -> list[BaseEvent]: case StateChanged(): state = dict(event.state) - yield StateSnapshotEvent(snapshot=state_metadata(state)) + yield _state_snapshot(state) # The write that was missing is in now, so a view held # back at its tool can be filled and sent. for view in ready(): @@ -322,7 +363,12 @@ def ready() -> list[BaseEvent]: # the answer instead of beside its call. for view in ready(): yield view - open_message = next_id("msg") + # The provider's own id where there is one, so the + # message this stream created and the message a + # readback returns are the same message. A snapshot + # then reconciles a client's copy in place instead of + # dropping it and re-appending the server's. + open_message = event.id or next_id("msg") yield TextMessageStartEvent(message_id=open_message) yield TextMessageContentEvent( message_id=open_message, delta=event.text @@ -340,9 +386,7 @@ def ready() -> list[BaseEvent]: # they are assembled from what each node wrote, before the # reducer has assigned write order. if event.result.sidecar: - yield StateSnapshotEvent( - snapshot=state_metadata(event.result.sidecar) - ) + yield _state_snapshot(event.result.sidecar) if event.result.citations: yield _activity( next_id("act"), @@ -361,4 +405,9 @@ def ready() -> list[BaseEvent]: yield RunErrorEvent(message=str(error) or type(error).__name__) return + if history is not None: + # Not inside the `try`: a failed turn leaves the thread mid-write, and + # a snapshot of that would tell a client to adopt a transcript the + # server may itself discard. A run that errored says so and stops. + yield MessagesSnapshotEvent(messages=list(await history())) yield RunFinishedEvent(thread_id=thread_id, run_id=run_id) diff --git a/src/mcp_agent_api/routes.py b/src/mcp_agent_api/routes.py index eea327c..cb67776 100644 --- a/src/mcp_agent_api/routes.py +++ b/src/mcp_agent_api/routes.py @@ -256,6 +256,16 @@ class StateValueResponse(BaseModel): value: Any = Field(description="The payload the event stream left out.") +def latest_user_message( + messages: Iterable[Mapping[str, Any]], +) -> Mapping[str, Any] | None: + """The last user message, or ``None`` where the request carried none.""" + for message in reversed(list(messages)): + if message.get("role") == "user": + return message + return None + + def latest_user_text(messages: Iterable[Mapping[str, Any]]) -> str: """The text of the last user message, flattened out of its content parts. @@ -263,18 +273,16 @@ def latest_user_text(messages: Iterable[Mapping[str, Any]]) -> str: image alongside its question still has a question in there, so the text parts are joined rather than the whole thing rejected. """ - for message in reversed(list(messages)): - if message.get("role") != "user": - continue - content = message.get("content") - if isinstance(content, str): - return content - if isinstance(content, list): - return "".join( - str(part.get("text", "")) - for part in content - if isinstance(part, dict) and part.get("type") == "text" - ) + message = latest_user_message(messages) + content = message.get("content") if message else None + if isinstance(content, str): + return content + if isinstance(content, list): + return "".join( + str(part.get("text", "")) + for part in content + if isinstance(part, dict) and part.get("type") == "text" + ) return "" @@ -424,6 +432,19 @@ async def thread_values(thread_id: str) -> dict[str, Any]: raise HTTPException(404, f"no thread {thread_id!r}") return values + async def thread_snapshot(thread_id: str) -> list[Message]: + """The thread as it stands, for the run's closing ``MESSAGES_SNAPSHOT``. + + Deliberately not :func:`thread_values`: an unknown thread is a 404 to + someone asking to read one, and an empty list to a run that has just + opened one — a thread whose turn failed before writing anything is a + thread with no messages, not a client mistake. + """ + config = {"configurable": {"thread_id": thread_id}} + snapshot = await built().agent.aget_state(cast(Any, config)) + values = dict(getattr(snapshot, "values", None) or {}) + return thread_messages(values.get("messages") or []) + @router.post( "/runs", response_class=EventStreamResponse, @@ -440,7 +461,12 @@ async def thread_values(thread_id: str) -> dict[str, Any]: "`TOOL_CALL_*` per tool, `STATE_SNAPSHOT`, and " "`ACTIVITY_SNAPSHOT` for what AG-UI has no vocabulary for " "— where a tool's arguments came from and which `ui://` " - "view renders its result. Ends with `RUN_FINISHED`, or " + "view renders its result. Every `STATE_SNAPSHOT` carries " + "its metadata under `toolState`, leaving the rest of the " + "state object to the client. A `MESSAGES_SNAPSHOT` closes " + "the run with the thread as the server holds it — history " + "is the server's, so reconcile against this rather than " + "trusting a local copy. Ends with `RUN_FINISHED`, or " "`RUN_ERROR` if the turn failed after the stream opened. " "The schema below is every event the protocol defines; this " "server emits the subset named above." @@ -460,11 +486,15 @@ async def create_run(body: RunRequest, request: Request) -> StreamingResponse: already does for anything the turn raises. """ agent = built() - # Dumped back to mappings for the helper, which reads AG-UI content - # parts and is shared with callers that never had models. - question = latest_user_text( - message.model_dump(by_alias=True) for message in body.messages - ) + # Dumped back to mappings for the helpers, which read AG-UI content + # parts and are shared with callers that never had models. + posted = [message.model_dump(by_alias=True) for message in body.messages] + question = latest_user_text(posted) + asked = latest_user_message(posted) + # The client's own id for the question becomes the thread's, so the + # message it is already showing and the message a readback returns are + # the same message rather than two with the same words. + question_id = str((asked or {}).get("id") or "") or None if not question: # A turn on an empty string would run the model, cost a call and # answer nothing. The client dropped its own question. @@ -485,13 +515,16 @@ async def frames() -> AsyncIterator[str]: else nullcontext(None) ) with user_credentials(credentials or None), around as config: - turn = stream_turn(agent.agent, question, thread_id, config) + turn = stream_turn( + agent.agent, question, thread_id, config, message_id=question_id + ) async for event in agui_events( turn, thread_id=thread_id, run_id=run_id, tools={tool.name: tool for tool in agent.tools}, withheld=agent.withheld, + history=lambda: thread_snapshot(thread_id), ): yield encoder.encode(event) diff --git a/tests/mcp_agent/test_streaming.py b/tests/mcp_agent/test_streaming.py index 78c7cbb..491cc97 100644 --- a/tests/mcp_agent/test_streaming.py +++ b/tests/mcp_agent/test_streaming.py @@ -209,8 +209,15 @@ def _agent(script: list[BaseMessage] | None = None) -> Any: return agent -async def _collect(agent: Any, text: str = "clip chirps", thread: str = "t1") -> list: - return [event async for event in stream_turn(agent, text, thread)] +async def _collect( + agent: Any, + text: str = "clip chirps", + thread: str = "t1", + message_id: str | None = None, +) -> list: + return [ + event async for event in stream_turn(agent, text, thread, message_id=message_id) + ] async def test_a_tool_actually_ran(): @@ -437,6 +444,36 @@ async def test_a_second_turn_continues_the_thread(): assert [m.text for m in second[-1].result.new_messages] == ["Second answer."] +async def test_the_question_keeps_the_id_its_client_gave_it(): + """So a client's own copy and a readback are the same message, which is what + lets a MESSAGES_SNAPSHOT reconcile rather than rebuild.""" + agent = _agent([AIMessage(content="Answer.")]) + + events = [ + event async for event in stream_turn(agent, "one", "t1", message_id="client-1") + ] + + human = [m for m in events[-1].result.history if m.type == "human"] + assert [m.id for m in human] == ["client-1"] + + +async def test_an_id_the_thread_already_holds_is_not_reused(): + """LangGraph's reducer matches on id, so reusing one *replaces* that message + rather than adding this one. A client numbering messages per session, or + retrying with the same id, would silently rewrite its own history.""" + agent = _agent([AIMessage(content="First."), AIMessage(content="Second.")]) + + await _collect(agent, "one", "t1", message_id="same") + events = [ + event async for event in stream_turn(agent, "two", "t1", message_id="same") + ] + + human = [m for m in events[-1].result.history if m.type == "human"] + assert [m.text for m in human] == ["one", "two"], "the first must survive" + assert human[0].id == "same" + assert human[1].id != "same" + + async def test_a_turn_with_no_tools_yields_only_text(): events = await _collect(_agent([AIMessage(content="No tools needed.")])) diff --git a/tests/mcp_agent_api/test_events.py b/tests/mcp_agent_api/test_events.py index 9dcc5d6..e02df69 100644 --- a/tests/mcp_agent_api/test_events.py +++ b/tests/mcp_agent_api/test_events.py @@ -8,6 +8,7 @@ import json from typing import Any +from ag_ui.core import UserMessage from ag_ui.encoder import EventEncoder from langchain_core.messages import AIMessage from langchain_core.tools import BaseTool, StructuredTool @@ -21,6 +22,7 @@ ANSWER_CITATIONS, MCP_VIEW, STATE_CONSUMED, + STATE_NAMESPACE, STATE_PUBLISHED, TOOLS_WITHHELD, agui_events, @@ -180,12 +182,89 @@ async def test_what_a_tool_published_gets_its_own_activity(): assert published["tool"] == "search" +async def test_history_closes_the_run_rather_than_opening_it(): + """A client renders its question the moment it is typed, and a snapshot is + applied by dropping every local message it does not name. Sent up front, + before this turn is checkpointed, it would take the question off the screen + and leave the answer under nothing.""" + + async def history() -> list[UserMessage]: + return [UserMessage(id="u0", content="the question")] + + types = _types(await _events(history=history)) + + assert types.count("MESSAGES_SNAPSHOT") == 1 + assert types[-2:] == ["MESSAGES_SNAPSHOT", "RUN_FINISHED"] + + +async def test_no_history_callable_sends_no_snapshot(): + """The pure layer stays usable by a consumer with no checkpointer to read.""" + assert "MESSAGES_SNAPSHOT" not in _types(await _events()) + + +async def test_a_failed_turn_snapshots_nothing(): + """A turn that raised left the thread mid-write. Telling a client to adopt + that as the transcript would hand it a state the server may itself drop.""" + + async def boom() -> Any: + raise RuntimeError("nope") + yield # pragma: no cover - never reached + + async def history() -> list[UserMessage]: # pragma: no cover - not called + raise AssertionError("history must not be read after a failure") + + types = _types( + [ + event + async for event in agui_events( + boom(), thread_id="t1", run_id="r1", history=history + ) + ] + ) + + assert types == ["RUN_STARTED", "RUN_ERROR"] + + +async def test_the_answer_is_labelled_with_the_id_the_thread_will_keep(): + """The whole reason a closing snapshot reconciles instead of rebuilding: the + message this stream created and the message a readback returns are the same + message, so a client keeps its copy in place.""" + agent = _agent() + events = [ + event + async for event in agui_events( + stream_turn(agent, "clip chirps", "t1"), thread_id="t1", run_id="r1" + ) + ] + opened = [e for e in events if e.type.value == "TEXT_MESSAGE_START"] + assert opened + + state = await agent.aget_state({"configurable": {"thread_id": "t1"}}) + stored = { + message.id + for message in state.values["messages"] + if type(message).__name__ == "AIMessage" + } + assert {event.message_id for event in opened} <= stored + + +async def test_the_state_object_leaves_room_for_the_client(): + """AG-UI's `state` is shared with the client, not ours to own. Our metadata + sits under one key so a client's own state can sit beside it, and so a + STATE_DELTA has a stable path to patch.""" + events = await _events() + + snapshot = [e for e in events if e.type.value == "STATE_SNAPSHOT"][-1].snapshot + assert list(snapshot) == [STATE_NAMESPACE] + assert STATE_KEY in snapshot[STATE_NAMESPACE] + + async def test_state_snapshots_carry_metadata_and_never_the_value(): events = await _events() snapshots = [e for e in events if e.type.value == "STATE_SNAPSHOT"] assert snapshots, "a tool published, so state changed" - entry = snapshots[-1].snapshot[STATE_KEY] + entry = snapshots[-1].snapshot[STATE_NAMESPACE][STATE_KEY] assert entry["kind"] == "geojson.AreaOfInterest" assert entry["tool"] == "search" assert entry["bytes"] > 0 @@ -200,8 +279,8 @@ async def test_the_last_state_snapshot_is_the_merged_one(): events = await _events() snapshots = [e for e in events if e.type.value == "STATE_SNAPSHOT"] - assert "seq" not in snapshots[0].snapshot[STATE_KEY] - assert snapshots[-1].snapshot[STATE_KEY]["seq"] == 1 + assert "seq" not in snapshots[0].snapshot[STATE_NAMESPACE][STATE_KEY] + assert snapshots[-1].snapshot[STATE_NAMESPACE][STATE_KEY]["seq"] == 1 assert _types(events).index("STATE_SNAPSHOT") < _types(events).index( "TEXT_MESSAGE_START" )