Python: preserve Responses replay metadata across AG-UI continuations

python/packages/ag-ui/agent_framework_ag_ui/_a2ui/_agent.py

@@ -351,6 +351,15 @@ def default_options(self) -> Any:
     def context_providers(self) -> Any:
         return getattr(self.inner_agent, "context_providers", [])
 
+    @property
+    def service_session_state_keys(self) -> Any:
+        return getattr(self.inner_agent, "service_session_state_keys", ())
+
+    def __getattr__(self, name: str) -> Any:
+        if name == "create_conversation":
+            return getattr(self.inner_agent, name)
+        raise AttributeError(name)
+
     # -- public run -------------------------------------------------------
 
     def run(self, messages: Any = None, *, stream: bool = False, **kwargs: Any) -> Any:

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

@@ -2109,12 +2109,20 @@ def _text_events_to_snapshot_messages(events: list[BaseEvent]) -> list[dict[str,
     return [message for message in messages if message.get("content")]
 
 
-def _restore_session_continuation_state(session: AgentSession, snapshot: AGUIThreadSnapshot | None) -> None:
+def _restore_session_continuation_state(
+    session: AgentSession,
+    snapshot: AGUIThreadSnapshot | None,
+    *,
+    restore_service_session_id: bool,
+    excluded_state_keys: set[str],
+) -> None:
     """Restore typed private state from trusted snapshot storage."""
     if snapshot is None or snapshot.session_state is None:
         return
     serialized_state = copy.deepcopy(snapshot.session_state)
     service_session_id = serialized_state.pop(_PROVIDER_SERVICE_SESSION_ID_STATE_KEY, None)
+    for key in excluded_state_keys:
+        serialized_state.pop(key, None)
     try:
         restored = AgentSession.from_dict(
             {
@@ -2130,7 +2138,7 @@ def _restore_session_continuation_state(session: AgentSession, snapshot: AGUIThr
             session.session_id,
         )
         return
-    if service_session_id is not None:
+    if restore_service_session_id and service_session_id is not None:
         session.service_session_id = restored.service_session_id
     session.state.update(restored.state)
 
@@ -2176,14 +2184,24 @@ def _request_state_protected_keys(agent: SupportsAgentRun) -> set[str]:
         InMemoryHistoryProvider.DEFAULT_SOURCE_ID,
         MESSAGE_INJECTION_PENDING_MESSAGES_STATE_KEY,
         *(provider.source_id for provider in context_providers),
+        *_provider_service_session_state_keys(agent),
     }
 
 
+def _provider_service_session_state_keys(agent: SupportsAgentRun) -> set[str]:
+    """Return provider-owned session-state keys that must not cross stateless runs."""
+    keys = getattr(agent, "service_session_state_keys", ())
+    if not isinstance(keys, (list, tuple, set, frozenset)):
+        return set()
+    return {key for key in keys if isinstance(key, str)}
+
+
 def _serialize_session_continuation_state(
     session: AgentSession,
     agent: SupportsAgentRun,
     *,
     shared_state_keys: set[str],
+    include_service_session_id: bool,
 ) -> dict[str, Any] | None:
     """Serialize server-owned state while preserving each AG-UI State Authority."""
     context_providers = cast(list[Any], getattr(agent, "context_providers", []))
@@ -2193,13 +2211,16 @@ def _serialize_session_continuation_state(
         _PROVIDER_SERVICE_SESSION_ID_STATE_KEY,
         *(provider.source_id for provider in context_providers if isinstance(provider, HistoryProvider)),
     }
+    if not include_service_session_id:
+        excluded_keys.update(_provider_service_session_state_keys(agent))
     continuation_state = {key: value for key, value in session.state.items() if key not in excluded_keys}
-    if not continuation_state and session.service_session_id is None:
+    service_session_id = session.service_session_id if include_service_session_id else None
+    if not continuation_state and service_session_id is None:
         return None
 
     serialized_session = AgentSession(
         session_id=session.session_id,
-        service_session_id=session.service_session_id,
+        service_session_id=service_session_id,
     )
     serialized_session.state.update(continuation_state)
     serialized_payload = serialized_session.to_dict()
@@ -2214,13 +2235,15 @@ def _safe_serialize_session_continuation_state(
     agent: SupportsAgentRun,
     *,
     shared_state_keys: set[str],
+    include_service_session_id: bool,
 ) -> dict[str, Any] | None:
     """Return JSON-safe continuation state without failing a completed run."""
     try:
         serialized_state = _serialize_session_continuation_state(
             session,
             agent,
             shared_state_keys=shared_state_keys,
+            include_service_session_id=include_service_session_id,
         )
         if serialized_state is None:
             return None
@@ -2532,6 +2555,11 @@ async def run_agent_stream(
 
     # Create session (with service session support)
     if config.use_service_session:
+        if isinstance(default_options, dict) and default_options.get("store") is False:
+            raise ValueError(
+                "use_service_session=True requires provider storage. Set agent default_options['store']=True "
+                "or disable use_service_session."
+            )
         if not config.service_session_id_from_thread_id and not snapshot_session.enabled:
             raise ValueError(
                 "use_service_session=True requires snapshot persistence unless service_session_id_from_thread_id=True."
@@ -2557,7 +2585,12 @@ async def run_agent_stream(
             session = created_session
     else:
         session = AgentSession(session_id=thread_id)
-    _restore_session_continuation_state(session, stored_snapshot)
+    _restore_session_continuation_state(
+        session,
+        stored_snapshot,
+        restore_service_session_id=config.use_service_session,
+        excluded_state_keys=set() if config.use_service_session else _provider_service_session_state_keys(agent),
+    )
     protected_session_state_keys = _request_state_protected_keys(agent)
     session.state.update(
         {
@@ -2693,6 +2726,7 @@ async def run_agent_stream(
                 session,
                 agent,
                 shared_state_keys=set(flow.current_state).difference(protected_session_state_keys),
+                include_service_session_id=config.use_service_session,
             ),
         )
         _save_tool_approval_state(session, approval_state_store, approval_thread_id)
@@ -3051,6 +3085,7 @@ async def run_agent_stream(
             session,
             agent,
             shared_state_keys=set(flow.current_state).difference(protected_session_state_keys),
+            include_service_session_id=config.use_service_session,
         ),
     )
     _save_tool_approval_state(session, approval_state_store, approval_thread_id)

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

@@ -12,7 +12,7 @@
 
 import asyncio
 import json
-from typing import Any
+from typing import Any, cast
 
 import pytest
 
@@ -1189,12 +1189,84 @@ class _Inner:
         client = object()
         default_options: dict[str, Any] = {"tools": []}
         context_providers = ["cp"]
+        service_session_state_keys = frozenset({"provider_session"})
 
     inner = _Inner()
     runner = A2UIAgent(inner, _RenderSub())
     assert runner.client is inner.client
     assert runner.default_options is inner.default_options
     assert runner.context_providers is inner.context_providers
+    assert runner.service_session_state_keys is inner.service_session_state_keys
+
+
+def test_a2ui_agent_preserves_provider_state_authority_boundary():
+    from agent_framework import AgentSession
+
+    from agent_framework_ag_ui import AGUIThreadSnapshot
+    from agent_framework_ag_ui._agent_run import (
+        _provider_service_session_state_keys,
+        _request_state_protected_keys,
+        _restore_session_continuation_state,
+        _serialize_session_continuation_state,
+    )
+
+    inner = type(
+        "_Inner",
+        (),
+        {
+            "id": "i",
+            "name": "n",
+            "description": "d",
+            "context_providers": [],
+            "service_session_state_keys": frozenset({"provider_session"}),
+        },
+    )()
+    runner = A2UIAgent(inner, _RenderSub())
+
+    assert "provider_session" in _request_state_protected_keys(cast(Any, runner))
+    source = AgentSession()
+    source.state.update({"provider_session": "provider-id", "private": "preserved"})
+    serialized = _serialize_session_continuation_state(
+        source,
+        cast(Any, runner),
+        shared_state_keys=set(),
+        include_service_session_id=False,
+    )
+    assert serialized == {"private": "preserved"}
+
+    restored = AgentSession()
+    _restore_session_continuation_state(
+        restored,
+        AGUIThreadSnapshot(
+            session_state={
+                "provider_session": "client-or-stale-id",
+                "private": "preserved",
+            }
+        ),
+        restore_service_session_id=False,
+        excluded_state_keys=_provider_service_session_state_keys(cast(Any, runner)),
+    )
+    assert restored.state == {"private": "preserved"}
+
+
+def test_a2ui_agent_conditionally_delegates_conversation_creation():
+    async def create_conversation(*, session_id: str) -> str:
+        return session_id
+
+    inner = type(
+        "_Inner",
+        (),
+        {
+            "id": "i",
+            "name": "n",
+            "description": "d",
+            "create_conversation": staticmethod(create_conversation),
+        },
+    )()
+    runner = A2UIAgent(inner, _RenderSub())
+
+    assert runner.create_conversation is getattr(inner, "create_conversation")
+    assert not hasattr(A2UIAgent(type("_NoConversation", (), {})(), _RenderSub()), "create_conversation")
 
 
 def test_a2ui_agent_uses_per_request_context_over_constructor():

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

@@ -2,7 +2,7 @@
 
 """Tests for _agent_run.py helper functions and FlowState."""
 
-from typing import cast
+from typing import Any, cast
 
 import pytest
 from ag_ui.core import (
@@ -2658,3 +2658,72 @@ async def test_provider_owned_service_session_requires_snapshot_persistence():
                 }
             )
         ]
+
+
+async def test_service_session_rejects_disabled_provider_storage():
+    """Service-session continuation cannot work when provider storage is disabled."""
+    from conftest import StubAgent  # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports]
+
+    from agent_framework_ag_ui import AgentFrameworkAgent, InMemoryAGUIThreadSnapshotStore
+
+    agent = AgentFrameworkAgent(
+        agent=StubAgent(default_options={"store": False}),
+        use_service_session=True,
+        snapshot_store=InMemoryAGUIThreadSnapshotStore(),
+    )
+
+    with pytest.raises(ValueError, match="requires provider storage"):
+        _ = [
+            event
+            async for event in agent.run(
+                {
+                    "thread_id": "frontend-thread",
+                    "run_id": "run-store-false",
+                    "__ag_ui_snapshot_scope": "test",
+                    "messages": [{"role": "user", "content": "Hello"}],
+                }
+            )
+        ]
+
+
+async def test_stateless_snapshot_excludes_only_provider_service_session_state():
+    """Stateless runs restore unrelated private state but not provider-owned continuation."""
+    from conftest import StubAgent  # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports]
+
+    from agent_framework_ag_ui import AgentFrameworkAgent, InMemoryAGUIThreadSnapshotStore
+
+    stub = StubAgent()
+    setattr(stub, "service_session_state_keys", frozenset({"provider_continuation"}))
+    observed_state: list[dict[str, Any]] = []
+    original_run = stub.run
+
+    def capture_state(*args: Any, **kwargs: Any) -> Any:
+        session = kwargs["session"]
+        observed_state.append(dict(session.state))
+        session.state["provider_continuation"] = "provider-session"
+        session.state["private"] = "preserved"
+        return original_run(*args, **kwargs)
+
+    stub.run = capture_state  # type: ignore[assignment, method-assign]  # ty: ignore[invalid-assignment]
+    store = InMemoryAGUIThreadSnapshotStore()
+    agent = AgentFrameworkAgent(agent=stub, snapshot_store=store)
+    payload = {
+        "thread_id": "frontend-thread",
+        "__ag_ui_snapshot_scope": "test",
+        "messages": [{"role": "user", "content": "Hello"}],
+        "state": {
+            "provider_continuation": "client-injected",
+            "client_value": "available",
+        },
+    }
+
+    _ = [event async for event in agent.run(payload)]
+    first_snapshot = await store.get(scope="test", thread_id="frontend-thread")
+    assert first_snapshot is not None
+    _ = [event async for event in agent.run(payload)]
+
+    assert first_snapshot.session_state == {"private": "preserved"}
+    assert observed_state == [
+        {"client_value": "available"},
+        {"private": "preserved", "client_value": "available"},
+    ]

python/packages/foundry/agent_framework_foundry/_agent.py

@@ -659,6 +659,8 @@ class RawFoundryAgent(
             result = await agent.run("Hello!")
     """
 
+    service_session_state_keys: ClassVar[frozenset[str]] = frozenset({FOUNDRY_HOSTED_AGENT_SESSION_ID_KEY})
+
     def __init__(
         self,
         *,