diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py index 436607114a9..40f1d02a7e2 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py @@ -3710,6 +3710,22 @@ async def save_safe_point_snapshot() -> None: interrupt=snapshot_interrupts(), ) + # Received input is durable independently of the provider's response. Save + # before invoking it so an exception cannot erase this turn's user messages + # or attachments. Later safe points still exclude unfinalized model output. + # Service-session history must remain atomic with terminal provider continuation. + if snapshot_session.enabled and not config.use_service_session: + persisted_input = snapshot_messages + if resume_payload is not None and not seeded_resume_from_snapshot and snapshot_seed_messages is None: + persisted_input = snapshot_session.resume_seeded_messages(persisted_input) + await save_thread_snapshot( + persisted_messages=persisted_input, + state=latest_state_snapshot, + interrupt=merge_snapshot_interrupts( + remaining_stored_interrupts(retired_approval_interrupt_ids), flow.interrupts + ), + ) + initial_state_snapshot = flow.current_state if state_schema and flow.current_state else None # With both IDs supplied there is nothing to wait for, so start the run before @@ -3763,6 +3779,8 @@ async def save_safe_point_snapshot() -> None: # IDs remain authoritative for lifecycle correlation and thread-scoped persistence. if not run_started_emitted: if supplied_thread_id is None and conv_id: + if not config.use_service_session: + await snapshot_session.migrate_thread_id(conv_id) thread_id = conv_id snapshot_session.rebind_thread_id(thread_id) if supplied_run_id is None and update.response_id: diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py b/python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py index 41adb4a7d2d..300d71273f6 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py @@ -593,6 +593,12 @@ def _legacy_binary_part(part: dict[str, Any]) -> dict[str, Any]: """Convert draft/legacy multimodal parts to AG-UI snapshot binary shape.""" normalized: dict[str, Any] = {"type": "binary"} url, data, binary_id, mime_type = _extract_multimodal_source_fields(part) + metadata = part.get("metadata") + filename = part.get("filename") + if not filename and isinstance(metadata, dict): + filename = metadata.get("filename") + if isinstance(filename, str) and filename: + normalized["filename"] = filename if isinstance(mime_type, str) and mime_type: normalized["mimeType"] = mime_type diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py b/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py index a946b1461da..9d080877d1b 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py @@ -1472,8 +1472,12 @@ def _reconstruct_messages_from_thread_snapshot( latest_user_index = _latest_user_message_index(incoming_messages) incoming_suffix = incoming_messages if latest_user_index is None else incoming_messages[latest_user_index:] + # Client projections can omit or regroup assistant output. A contiguous + # overlap then ends before later messages that are already stored. Their + # IDs still identify backend-owned records, even if the client edited them. + stored_ids = {str(message["id"]) for message in stored_messages if message.get("id")} incoming_suffix = _filter_untrusted_suffix( - incoming_suffix, + [message for message in incoming_suffix if not message.get("id") or str(message["id"]) not in stored_ids], stored_messages=stored_messages, stored_interrupt=stored_interrupt, ) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_snapshot_session.py b/python/packages/ag-ui/agent_framework_ag_ui/_snapshot_session.py index 85546830245..37b4ba51881 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_snapshot_session.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_snapshot_session.py @@ -107,6 +107,21 @@ def rebind_thread_id(self, thread_id: str) -> None: """ self._thread_id = thread_id + async def migrate_thread_id(self, thread_id: str) -> None: + """Move a provisional checkpoint before publishing a provider's thread ID. + + Save the destination before deleting the source. Storage failures propagate + so callers cannot announce an ID whose checkpoint was not migrated. + """ + if thread_id == self._thread_id: + return + if self._store is not None and self._scope is not None: + checkpoint = await self._store.get(scope=self._scope, thread_id=self._thread_id) + if checkpoint is not None: + await self._store.save(scope=self._scope, thread_id=thread_id, snapshot=checkpoint) + await self._store.delete(scope=self._scope, thread_id=self._thread_id) + self.rebind_thread_id(thread_id) + async def hydrate_events(self, *, run_id: str) -> AsyncGenerator[BaseEvent]: """Replay the stored snapshot as a complete run without invoking the agent.""" yield RunStartedEvent(run_id=run_id, thread_id=self._thread_id) diff --git a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py index d6cd10765d9..83d61fd3404 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_endpoint.py +++ b/python/packages/ag-ui/tests/ag_ui/test_endpoint.py @@ -1819,7 +1819,9 @@ async def stream_fn( await asyncio.wait_for(source_closed.wait(), timeout=5) assert not source_completed - assert await store.get(scope="tenant-a", thread_id="thread-cancelled") is None + snapshot = await store.get(scope="tenant-a", thread_id="thread-cancelled") + assert snapshot is not None + assert [(message["role"], message["content"]) for message in snapshot.messages] == [("user", "Start")] async def test_endpoint_detached_workflow_runner_completes_after_disconnect() -> None: @@ -11871,8 +11873,8 @@ async def stream_fn(messages: Any, options: Any, **kwargs: Any): assert _interrupt_metadata_value(hydrated_interrupts[0])["function_call"]["name"] == "draft_steps" -async def test_agent_endpoint_run_error_does_not_overwrite_previous_snapshot(streaming_chat_client_stub): - """A failing agent turn leaves the last good AG-UI Thread Snapshot available for hydration.""" +async def test_agent_endpoint_run_error_preserves_previous_history_and_received_input(streaming_chat_client_stub): + """A failing agent turn retains prior output and the newly received user turn.""" app = FastAPI() call_count = 0 @@ -11917,7 +11919,7 @@ async def stream_fn(messages: Any, options: Any, **kwargs: Any): assert call_count == 2 messages = _latest_messages_snapshot(hydrate_response) assert any(message.get("role") == "assistant" and message.get("content") == "Stable reply" for message in messages) - assert not any(message.get("content") == "Break the run" for message in messages) + assert sum(message.get("content") == "Break the run" for message in messages) == 1 async def test_workflow_endpoint_hydrates_emitted_snapshots_without_invoking_workflow(): @@ -13329,6 +13331,8 @@ async def stream_fn( ) -> AsyncIterator[ChatResponseUpdate]: del options, kwargs provider_result = next(message.text for message in messages if message.role == "system") + if any(message.text == "Second" for message in messages): + store.fail_next_save = True yield ChatResponseUpdate(contents=[Content.from_text(text=provider_result)]) store = _FailNextSaveStore() @@ -13353,7 +13357,6 @@ async def stream_fn( "/snapshot-failure", json={"thread_id": "thread-1", "messages": [{"role": "user", "content": "First"}]}, ) - store.fail_next_save = True failed_save_response = client.post( "/snapshot-failure", json={"thread_id": "thread-1", "messages": [{"role": "user", "content": "Second"}]}, @@ -14050,3 +14053,200 @@ async def stream_fn( if content.type == "function_result" and content.call_id == call_id ] assert provider_results == ["Action executed by server"] + + +@pytest.mark.parametrize("existing_thread", [False, True]) +async def test_agent_endpoint_retains_received_media_when_provider_fails(existing_thread: bool): + """A downstream failure keeps accepted input, without committing partial model output.""" + from agent_framework import AgentResponse, ResponseInvalidatedException, ResponseStream + + from agent_framework_ag_ui._snapshot_session import ThreadSnapshotSession + + store = InMemoryAGUIThreadSnapshotStore() + prior = [{"id": "prior", "role": "user", "content": "Earlier turn"}] if existing_thread else [] + if existing_thread: + await store.save(scope="tenant-a", thread_id="media-thread", snapshot=AGUIThreadSnapshot(messages=prior)) + agent = StubAgent() + + def failing_run(*args: Any, **kwargs: Any) -> Any: + async def updates() -> AsyncIterator[AgentResponseUpdate]: + yield AgentResponseUpdate(contents=[Content.from_text(text="Unfinalized output")], role="assistant") + raise ResponseInvalidatedException("provider rejected media") + + return ResponseStream(updates(), finalizer=AgentResponse.from_updates) + + agent.run = failing_run # type: ignore[assignment, method-assign] # ty: ignore[invalid-assignment] + app = FastAPI() + add_agent_framework_fastapi_endpoint( + app, agent, path="/media-failure", snapshot_store=store, snapshot_scope_resolver=lambda request: "tenant-a" + ) + media = [ + { + "type": kind, + "source": {"type": "data", "data": "aGVsbG8=", "mimeType": mime}, + "metadata": {"filename": filename}, + } + for kind, mime, filename in [ + ("image", "image/png", "original.png"), + ("document", "text/plain", "original.txt"), + ("audio", "audio/wav", "original.wav"), + ("video", "video/mp4", "original.mp4"), + ] + ] + received = {"id": "media-user", "role": "user", "content": media} + response = TestClient(app).post( + "/media-failure", + json={"threadId": "media-thread", "runId": "failed-run", "messages": [*prior, received], "state": {"count": 7}}, + ) + events = _decode_sse_events(response) + assert [event["code"] for event in events if event["type"] == "RUN_ERROR"] == ["ResponseInvalidatedException"] + assert not any(event["type"] == "RUN_FINISHED" for event in events) + cold = await ThreadSnapshotSession.open(store=store, scope="tenant-a", thread_id="media-thread") + assert cold.stored is not None + expected = { + **received, + "content": [ + { + "type": "binary", + "data": part["source"]["data"], + "mimeType": part["source"]["mimeType"], + "filename": part["metadata"]["filename"], + } + for part in media + ], + } + assert cold.stored.messages == [*prior, expected] + assert cold.stored.state == {"count": 7} + hydrated = [event async for event in cold.hydrate_events(run_id="cold-run")] + replay = next(event for event in hydrated if isinstance(event, MessagesSnapshotEvent)) + assert replay.messages[-1].model_dump(by_alias=True, exclude_none=True) == expected + + +@pytest.mark.parametrize("approved", [True, False]) +async def test_native_approval_replay_gap_preserves_canonical_history_once(approved: bool): + """A later client projection cannot duplicate an earlier native approval outcome.""" + import copy + + store = InMemoryAGUIThreadSnapshotStore() + client, _, executed_cities = _build_weather_approval_endpoint(snapshot_store=store) + resumed = client.post( + "/approval", + json={ + "threadId": "thread-weather", + "messages": [], + "resume": [{"interruptId": "call_get_weather", "status": "resolved", "payload": {"approved": approved}}], + }, + ) + assert not any(event["type"] == "RUN_ERROR" for event in _decode_sse_events(resumed)) + first = await store.get(scope="tenant-a", thread_id="thread-weather") + assert first is not None + result = [ + message + for message in first.messages + if message.get("toolCallId", message.get("tool_call_id")) == "call_get_weather" + ] + assert len(result) == 1 + assert result[0]["content"] == ( + "Sunny in Seattle" if approved else "Error: Tool call invocation was rejected by user." + ) + for index in range(2): + before = await store.get(scope="tenant-a", thread_id="thread-weather") + assert before is not None + replay = copy.deepcopy(before.messages) + assistant = next(message for message in replay if message["role"] == "assistant") + assistant["id"] = f"client-projection-{index}" + assistant["content"] = "Client regrouped assistant output" + new_user = {"id": f"followup-{index}", "role": "user", "content": "Continue"} + response = client.post("/approval", json={"threadId": "thread-weather", "messages": [*replay, new_user]}) + assert not any(event["type"] == "RUN_ERROR" for event in _decode_sse_events(response)) + after = await store.get(scope="tenant-a", thread_id="thread-weather") + assert after is not None + assert after.messages[: len(before.messages)] == before.messages + assert after.messages[len(before.messages)] == new_user + assert len({message["id"] for message in after.messages}) == len(after.messages) + assert [ + message + for message in after.messages + if message.get("toolCallId", message.get("tool_call_id")) == "call_get_weather" + ] == result + assert executed_cities == (["Seattle"] if approved else []) + + +async def test_input_checkpoint_preserves_unconsumed_interrupt_on_provider_failure(): + """A failed ordinary turn cannot consume a stored confirmation prompt.""" + from agent_framework import AgentResponse, ResponseInvalidatedException, ResponseStream + + store = InMemoryAGUIThreadSnapshotStore() + interrupt = {"id": "confirm-one", "reason": "confirm_changes", "payload": {"steps": ["Review"]}} + prior = {"id": "prior", "role": "user", "content": "Draft a plan"} + await store.save( + scope="tenant-a", + thread_id="pending-thread", + snapshot=AGUIThreadSnapshot(messages=[prior], interrupt=[interrupt]), + ) + agent = StubAgent() + + def failing_run(*args: Any, **kwargs: Any) -> Any: + async def updates() -> AsyncIterator[AgentResponseUpdate]: + yield AgentResponseUpdate(contents=[Content.from_text(text="Partial")], role="assistant") + raise ResponseInvalidatedException("provider failed") + + return ResponseStream(updates(), finalizer=AgentResponse.from_updates) + + agent.run = failing_run # type: ignore[assignment, method-assign] # ty: ignore[invalid-assignment] + app = FastAPI() + add_agent_framework_fastapi_endpoint( + app, agent, path="/pending", snapshot_store=store, snapshot_scope_resolver=lambda request: "tenant-a" + ) + received = {"id": "next", "role": "user", "content": "Explain the plan"} + response = TestClient(app).post( + "/pending", json={"threadId": "pending-thread", "runId": "failed", "messages": [prior, received]} + ) + assert any(event["type"] == "RUN_ERROR" for event in _decode_sse_events(response)) + saved = await store.get(scope="tenant-a", thread_id="pending-thread") + assert saved is not None + assert saved.messages == [prior, received] + assert saved.interrupt == [interrupt] + + +@pytest.mark.parametrize("provider_fails", [False, True]) +async def test_input_checkpoint_follows_provider_thread_id(provider_fails: bool): + """The public fallback thread ID owns input, without a provisional orphan.""" + from unittest.mock import AsyncMock + + from agent_framework import AgentResponse, ResponseInvalidatedException, ResponseStream + + store = InMemoryAGUIThreadSnapshotStore() + original_save = store.save + store.save = AsyncMock(wraps=original_save) # type: ignore[method-assign] # ty: ignore[invalid-assignment] + agent = StubAgent() + + def provider_run(*args: Any, **kwargs: Any) -> Any: + async def updates() -> AsyncIterator[AgentResponseUpdate]: + yield AgentResponseUpdate( + contents=[Content.from_text(text="Partial")], + role="assistant", + raw_representation=ChatResponseUpdate(conversation_id="provider-thread"), + ) + if provider_fails: + raise ResponseInvalidatedException("provider failed") + + return ResponseStream(updates(), finalizer=AgentResponse.from_updates) + + agent.run = provider_run # type: ignore[assignment, method-assign] # ty: ignore[invalid-assignment] + app = FastAPI() + add_agent_framework_fastapi_endpoint( + app, agent, path="/fallback", snapshot_store=store, snapshot_scope_resolver=lambda request: "tenant-a" + ) + received = {"id": "input", "role": "user", "content": "Remember this"} + response = TestClient(app).post("/fallback", json={"messages": [received]}) + events = _decode_sse_events(response) + assert next(event for event in events if event["type"] == "RUN_STARTED")["threadId"] == "provider-thread" + assert any(event["type"] == "RUN_ERROR" for event in events) == provider_fails + saved = await store.get(scope="tenant-a", thread_id="provider-thread") + assert saved is not None + assert saved.messages[0] == received + for call in store.save.call_args_list: + saved_id = call.kwargs["thread_id"] + if saved_id != "provider-thread": + assert await store.get(scope="tenant-a", thread_id=saved_id) is None diff --git a/python/packages/ag-ui/tests/ag_ui/test_message_adapters.py b/python/packages/ag-ui/tests/ag_ui/test_message_adapters.py index 22bf0d38be8..d4341f9341f 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_message_adapters.py +++ b/python/packages/ag-ui/tests/ag_ui/test_message_adapters.py @@ -3470,3 +3470,28 @@ def test_parse_multimodal_media_part_url_field_backward_compat(): ) assert from_uri is not None assert from_uri.uri == "https://example.com/b.png" + + +@pytest.mark.parametrize( + "kind,mime", [("image", "image/png"), ("document", "text/plain"), ("audio", "audio/wav"), ("video", "video/mp4")] +) +def test_snapshot_media_retains_filename_across_normalization(kind: str, mime: str) -> None: + """Draft metadata filenames survive the legacy snapshot shape and subsequent turns.""" + incoming = [ + { + "id": "upload", + "role": "user", + "content": [ + { + "type": kind, + "source": {"type": "data", "data": "aGVsbG8=", "mimeType": mime}, + "metadata": {"filename": "original-name"}, + } + ], + } + ] + normalized = agui_messages_to_snapshot_format(incoming) + assert normalized[0]["content"] == [ + {"type": "binary", "data": "aGVsbG8=", "mimeType": mime, "filename": "original-name"} + ] + assert agui_messages_to_snapshot_format(normalized) == normalized diff --git a/python/packages/ag-ui/tests/ag_ui/test_run.py b/python/packages/ag-ui/tests/ag_ui/test_run.py index 6aa7aab5afc..54228d0ba0e 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_run.py +++ b/python/packages/ag-ui/tests/ag_ui/test_run.py @@ -3241,7 +3241,10 @@ def invalidate_result(response: AgentResponse[Any]) -> None: snapshot = await store.get(scope="tenant-a", thread_id="invalidated-thread") assert snapshot is not None - assert snapshot.messages == [{"id": "baseline", "role": "user", "content": "Persisted baseline"}] + assert snapshot.messages[0] == {"id": "baseline", "role": "user", "content": "Persisted baseline"} + assert len(snapshot.messages) == 2 + assert snapshot.messages[1]["role"] == "user" + assert snapshot.messages[1]["content"] == "Start" assert "c1" not in str(snapshot.messages) diff --git a/python/packages/ag-ui/tests/ag_ui/test_run_common.py b/python/packages/ag-ui/tests/ag_ui/test_run_common.py index 69822014420..a8f9c456d64 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_run_common.py +++ b/python/packages/ag-ui/tests/ag_ui/test_run_common.py @@ -1108,3 +1108,37 @@ def test_new_turn_gets_new_reasoning_id(self): assert id1 is not None assert id2 is not None assert id1 != id2, "New turn should get a new reasoning message_id" + + +@pytest.mark.parametrize("outcome", ["approved", "declined", None]) +def test_snapshot_replay_after_projection_gap_keeps_original_occurrences(outcome: str | None) -> None: + """A changed assistant projection must not append later stored user/results again.""" + stored: list[dict[str, Any]] = [ + {"id": "u1", "role": "user", "content": "Start"}, + {"id": "a1", "role": "assistant", "content": "Canonical answer"}, + {"id": "u2", "role": "user", "content": "Approve the meeting"}, + { + "id": "call-message", + "role": "assistant", + "toolCalls": [{"id": "meeting", "type": "function", "function": {"name": "schedule", "arguments": "{}"}}], + }, + ] + if outcome is not None: + stored.append({"id": "result", "role": "tool", "toolCallId": "meeting", "content": outcome}) + # Mix a completed native approval with the parallel frontend completed/pending call. + stored[3]["toolCalls"].append( + {"id": "native-call", "type": "function", "function": {"name": "publish", "arguments": "{}"}} + ) + stored.append( + {"id": "native-result", "role": "tool", "toolCallId": "native-call", "content": "Published after approval"} + ) + incoming = [dict(message) for message in stored] + incoming[1] = {"id": "client-a1", "role": "assistant", "content": "Different live projection"} + incoming[2] = {**incoming[2], "content": "Client edit must not replace canonical history"} + new_messages = [{"id": "u3", "role": "user", "content": "Continue"}] + incoming.extend(new_messages) + + actual = _reconstruct_messages_from_thread_snapshot(stored_messages=stored, incoming_messages=incoming) + + assert actual == stored + new_messages + assert len({message["id"] for message in actual}) == len(actual) diff --git a/python/packages/ag-ui/tests/ag_ui/test_snapshot_session.py b/python/packages/ag-ui/tests/ag_ui/test_snapshot_session.py index ef72e585812..64e5cb7ec82 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_snapshot_session.py +++ b/python/packages/ag-ui/tests/ag_ui/test_snapshot_session.py @@ -367,3 +367,39 @@ async def test_clears_all_interrupts_when_ids_omitted(self) -> None: async def test_disabled_session_clear_is_a_noop(self) -> None: session = await ThreadSnapshotSession.open(store=None, scope=None, thread_id="t1") await session.clear_interrupts(interrupt_ids={"int-1"}) + + +@pytest.mark.parametrize("failed_operation", ["save", "delete"]) +async def test_migrate_thread_id_keeps_source_when_storage_fails(failed_operation: str) -> None: + """A migration failure keeps the recoverable source and does not publish the new ID.""" + from unittest.mock import AsyncMock, patch + + snapshot = AGUIThreadSnapshot(messages=[{"id": "input", "role": "user", "content": "Keep me"}]) + store = await make_store_with("tenant", "provisional", snapshot) + session = await ThreadSnapshotSession.open(store=store, scope="tenant", thread_id="provisional") + with patch.object(store, failed_operation, AsyncMock(side_effect=OSError("storage unavailable"))): + with pytest.raises(OSError, match="storage unavailable"): + await session.migrate_thread_id("provider") + assert await store.get(scope="tenant", thread_id="provisional") == snapshot + events = [event async for event in session.hydrate_events(run_id="replay")] + assert events[0].thread_id == "provisional" + if failed_operation == "delete": + assert await store.get(scope="tenant", thread_id="provider") == snapshot + + +async def test_migrate_thread_id_preserves_snapshot_and_scope() -> None: + """Migration moves the saved checkpoint without changing any of its payloads.""" + snapshot = AGUIThreadSnapshot( + messages=[{"id": "input", "role": "user", "content": "Keep me"}], + state={"counter": 2}, + interrupt=[{"id": "confirm", "reason": "confirm_changes"}], + session_state={"private": "continuation"}, + ) + store = await make_store_with("tenant", "provisional", snapshot) + await store.save(scope="other", thread_id="provisional", snapshot=snapshot) + session = await ThreadSnapshotSession.open(store=store, scope="tenant", thread_id="provisional") + await session.migrate_thread_id("provider") + await session.migrate_thread_id("provider") + assert await store.get(scope="tenant", thread_id="provider") == snapshot + assert await store.get(scope="tenant", thread_id="provisional") is None + assert await store.get(scope="other", thread_id="provisional") == snapshot