Python: fix(python): separate provider input from UI snapshot in AG-UI service-session mode

python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py

@@ -2195,6 +2195,60 @@ def _safe_serialize_session_continuation_state(
     return None
 
 
+def _split_service_session_input(
+    stored_snapshot_messages: list[dict[str, Any]],
+    current_turn_messages: list[dict[str, Any]],
+) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
+    """Splits current turn messages into provider suffix and full snapshot messages.
+
+    Uses identity (message ID) matching with a role-based fallback for assistant
+    messages (whose agent-generated IDs may differ from what the UI relays) to
+    safely handle stale snapshots, incremental requests, and truncated histories
+    without relying on brittle length cursors.
+
+    Args:
+        stored_snapshot_messages: The messages persisted from the previous turn.
+        current_turn_messages: The incoming messages from the UI for the current turn.
+
+    Returns:
+        A tuple containing the provider suffix and the reconstructed snapshot messages.
+    """
+
+    def _get_msg_id(msg: Mapping[str, Any]) -> str | None:
+        return msg.get("id") or msg.get("message_id")
+
+    def _get_msg_role(msg: Mapping[str, Any]) -> str | None:
+        return msg.get("role")
+
+    i = 0
+    j = 0
+    while i < len(stored_snapshot_messages) and j < len(current_turn_messages):
+        stored_msg = stored_snapshot_messages[i]
+        current_msg = current_turn_messages[j]
+        stored_id = _get_msg_id(stored_msg)
+        current_id = _get_msg_id(current_msg)
+        stored_role = _get_msg_role(stored_msg)
+        current_role = _get_msg_role(current_msg)
+
+        if stored_id and current_id and stored_id == current_id:
+            i += 1
+            j += 1
+        elif stored_role == "assistant" and current_role == "assistant":
+            i += 1
+            j += 1
+        else:
+            break
+
+    provider_suffix = current_turn_messages[j:]
+
+    if j > 0:
+        snapshot_messages = list(current_turn_messages)
+    else:
+        snapshot_messages = list(stored_snapshot_messages) + list(current_turn_messages)
+
+    return provider_suffix, snapshot_messages
+
+
 async def run_agent_stream(
     input_data: dict[str, Any],
     agent: SupportsAgentRun,
@@ -2266,16 +2320,32 @@ async def run_agent_stream(
             yield event
         return
 
+    snapshot_seed_messages: list[dict[str, Any]] | None = None
+
     if stored_snapshot is not None:
         if resume_payload is not None and stored_pending_approval_interrupt_ids:
-            raw_messages = snapshot_session.resume_seeded_messages(raw_messages)
             seeded_resume_from_snapshot = True
-        else:
+
+            if not config.use_service_session:
+                raw_messages = snapshot_session.resume_seeded_messages(raw_messages)
+            else:
+                provider_suffix, snapshot_seed_messages = _split_service_session_input(
+                    stored_snapshot_messages=stored_snapshot.messages,
+                    current_turn_messages=raw_messages,
+                )
+                raw_messages = provider_suffix
+        elif not config.use_service_session:
             raw_messages = _reconstruct_messages_from_thread_snapshot(
                 stored_messages=stored_snapshot.messages,
                 incoming_messages=raw_messages,
                 stored_interrupt=stored_snapshot.interrupt,
             )
+        else:
+            provider_suffix, snapshot_seed_messages = _split_service_session_input(
+                stored_snapshot_messages=stored_snapshot.messages,
+                current_turn_messages=raw_messages,
+            )
+            raw_messages = provider_suffix
 
     # Initialize flow state with stored state plus request-provided overrides;
     # endpoint-deferred defaults apply only to keys missing from both.
@@ -2378,6 +2448,20 @@ async def run_agent_stream(
         protected_tool_call_ids=protected_tool_call_ids,
     )
 
+    if config.use_service_session and snapshot_seed_messages is not None:
+        _, snapshot_messages = normalize_agui_input_messages(
+            snapshot_seed_messages,
+            protected_tool_call_ids=protected_tool_call_ids,
+        )
+    elif config.use_service_session and stored_snapshot is not None:
+        if seeded_resume_from_snapshot:
+            snapshot_messages = snapshot_session.resume_seeded_messages(snapshot_messages)
+        else:
+            snapshot_messages = _reconstruct_messages_from_thread_snapshot(
+                stored_messages=stored_snapshot.messages,
+                incoming_messages=snapshot_messages,
+                stored_interrupt=stored_snapshot.interrupt,
+            )
     # Check for structured output mode (skip text content)
     skip_text = False
     response_format: type[Any] | None = None
@@ -2840,6 +2924,7 @@ async def run_agent_stream(
         flow.pending_tool_calls or flow.tool_results or flow.accumulated_text or flow.reasoning_messages
     )
     latest_messages_snapshot = snapshot_messages
+
     if should_emit_snapshot:
         # Always fold this turn's output into the persisted snapshot, even when the
         # outbound MESSAGES_SNAPSHOT event is suppressed for predictive tools.

python/packages/ag-ui/tests/ag_ui/test_snapshots.py

@@ -2,9 +2,19 @@
 
 """Tests for AG-UI thread snapshot storage primitives."""
 
+from collections.abc import AsyncGenerator
 from dataclasses import fields
+from typing import Any, cast
 
-from agent_framework_ag_ui import AGUIThreadSnapshot, AGUIThreadSnapshotStore, InMemoryAGUIThreadSnapshotStore
+import pytest
+from agent_framework import AgentResponseUpdate, Content, SupportsAgentRun
+
+from agent_framework_ag_ui import (
+    AgentFrameworkAgent,
+    AGUIThreadSnapshot,
+    AGUIThreadSnapshotStore,
+    InMemoryAGUIThreadSnapshotStore,
+)
 
 
 def test_thread_snapshot_model_contains_replayable_and_private_snapshot_fields() -> None:
@@ -204,3 +214,166 @@ async def test_in_memory_snapshot_store_rejects_invalid_keys() -> None:
         await store.delete(scope=None, thread_id="thread-1")  # type: ignore[arg-type]  # ty: ignore[invalid-argument-type]
     with pytest.raises(ValueError):
         await store.clear(scope="")
+
+
+class _InputSpyAgent:
+    """Minimal agent implementation for spying on provider input."""
+
+    name = "spy"
+    description = ""
+    default_options: dict = {}
+    context_providers: list = []
+
+    def __init__(self) -> None:
+        self.calls: list[dict] = []
+
+    def run(self, messages: Any, *, session: Any, stream: bool = False, **kwargs: Any) -> Any:
+        async def updates() -> AsyncGenerator[Any, None]:
+            self.calls.append(
+                {
+                    "roles": [m.role for m in messages],
+                    "service_session_id": session.service_session_id,
+                }
+            )
+            yield AgentResponseUpdate(
+                contents=[Content.from_text("ACK")],
+                role="assistant",
+                response_id=f"resp-{len(self.calls)}",
+            )
+
+        return updates()
+
+
+async def _drain(runner, body):
+    return [event async for event in runner.run(body)]
+
+
+@pytest.mark.asyncio
+async def test_service_session_snapshot_split_authority() -> None:
+    """Verify use_service_session + snapshot_store separates provider input from UI hydration.
+
+    - Provider must receive ONLY incremental input
+    - Snapshot store must retain FULL history for UI hydration
+    """
+    agent = _InputSpyAgent()
+    store = InMemoryAGUIThreadSnapshotStore()
+    runner = AgentFrameworkAgent(
+        agent=cast(SupportsAgentRun, agent),
+        use_service_session=True,
+        snapshot_store=store,
+    )
+
+    first_turn = {
+        "threadId": "conv_FHA_SESSION",
+        "__ag_ui_snapshot_scope": "split-auth-test",
+        "messages": [{"id": "u1", "role": "user", "content": "first"}],
+    }
+    second_turn = {
+        "threadId": "conv_FHA_SESSION",
+        "__ag_ui_snapshot_scope": "split-auth-test",
+        "messages": [
+            {"id": "u1", "role": "user", "content": "first"},
+            {"id": "a1", "role": "assistant", "content": "ACK"},
+            {"id": "u2", "role": "user", "content": "second"},
+        ],
+    }
+
+    await _drain(runner, first_turn)
+    await _drain(runner, second_turn)
+
+    assert agent.calls[1]["roles"] == ["user"], (
+        f"Expected incremental-only input for service-session mode, got: {agent.calls[1]['roles']}"
+    )
+    assert agent.calls[1]["service_session_id"] == "conv_FHA_SESSION"
+
+    snapshot = await store.get(scope="split-auth-test", thread_id="conv_FHA_SESSION")
+    assert snapshot is not None, "Snapshot should exist after two turns"
+    roles = [m.get("role") for m in snapshot.messages]
+    assert roles == ["user", "assistant", "user", "assistant"], (
+        f"Snapshot must contain full transcript for UI hydration, got: {roles}"
+    )
+
+
+@pytest.mark.asyncio
+async def test_service_session_snapshot_incremental_request() -> None:
+    """Verify identity matching handles incremental requests where UI sends only new messages.
+
+    - Stored snapshot has [u1, a1]
+    - UI sends ONLY [u2] (incremental, not full history)
+    - Provider should receive [u2]
+    - Snapshot store should retain [u1, a1, u2, a2]
+    """
+    agent = _InputSpyAgent()
+    store = InMemoryAGUIThreadSnapshotStore()
+    runner = AgentFrameworkAgent(
+        agent=cast(SupportsAgentRun, agent),
+        use_service_session=True,
+        snapshot_store=store,
+    )
+
+    first_turn = {
+        "threadId": "conv_INCREMENTAL",
+        "__ag_ui_snapshot_scope": "incremental-test",
+        "messages": [{"id": "u1", "role": "user", "content": "first"}],
+    }
+    await _drain(runner, first_turn)
+
+    incremental_turn = {
+        "threadId": "conv_INCREMENTAL",
+        "__ag_ui_snapshot_scope": "incremental-test",
+        "messages": [{"id": "u2", "role": "user", "content": "second"}],
+    }
+    await _drain(runner, incremental_turn)
+
+    assert agent.calls[1]["roles"] == ["user"], (
+        f"Provider should receive incremental input, got: {agent.calls[1]['roles']}"
+    )
+
+    snapshot = await store.get(scope="incremental-test", thread_id="conv_INCREMENTAL")
+    assert snapshot is not None
+    roles = [m.get("role") for m in snapshot.messages]
+    assert roles == ["user", "assistant", "user", "assistant"], f"Snapshot must contain full transcript, got: {roles}"
+
+
+@pytest.mark.asyncio
+async def test_service_session_snapshot_stale_snapshot() -> None:
+    """Verify identity matching handles stale snapshots where UI and store diverge.
+
+    - Stored snapshot has [u1, a1]
+    - UI sends [u1, u2] (UI is ahead, missing a1)
+    - Provider should receive [u2] (identity match skips u1)
+    - Snapshot store should retain UI's view [u1, u2, a2]
+    """
+    agent = _InputSpyAgent()
+    store = InMemoryAGUIThreadSnapshotStore()
+    runner = AgentFrameworkAgent(
+        agent=cast(SupportsAgentRun, agent),
+        use_service_session=True,
+        snapshot_store=store,
+    )
+
+    first_turn = {
+        "threadId": "conv_STALE",
+        "__ag_ui_snapshot_scope": "stale-test",
+        "messages": [{"id": "u1", "role": "user", "content": "first"}],
+    }
+    await _drain(runner, first_turn)
+
+    stale_turn = {
+        "threadId": "conv_STALE",
+        "__ag_ui_snapshot_scope": "stale-test",
+        "messages": [
+            {"id": "u1", "role": "user", "content": "first"},
+            {"id": "u2", "role": "user", "content": "second"},
+        ],
+    }
+    await _drain(runner, stale_turn)
+
+    assert agent.calls[1]["roles"] == ["user"], (
+        f"Provider should receive only new messages, got: {agent.calls[1]['roles']}"
+    )
+
+    snapshot = await store.get(scope="stale-test", thread_id="conv_STALE")
+    assert snapshot is not None
+    msg_ids = [m.get("id") for m in snapshot.messages]
+    assert "u1" in msg_ids and "u2" in msg_ids, f"Snapshot should contain UI's messages, got: {msg_ids}"