Python: preserve Responses replay metadata across AG-UI continuations

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_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

@@ -439,16 +439,26 @@ def _parse_chunk_from_openai(
         options: dict[str, Any],
         function_call_ids: dict[int, tuple[str, str]],
         seen_reasoning_delta_item_ids: set[str] | None = None,
+        output_text_logprobs: dict[str, list[Any]] | None = None,
     ) -> ChatResponseUpdate:
         """Parse streaming events while preserving hosted-agent session state."""
         update = try_parse_oauth_consent_event(event, self.model)
         if update is None:
-            update = super()._parse_chunk_from_openai(
-                event,
-                options,
-                function_call_ids,
-                seen_reasoning_delta_item_ids,
-            )
+            if output_text_logprobs is None:
+                update = super()._parse_chunk_from_openai(
+                    event,
+                    options,
+                    function_call_ids,
+                    seen_reasoning_delta_item_ids,
+                )
+            else:
+                update = super()._parse_chunk_from_openai(
+                    event,
+                    options,
+                    function_call_ids,
+                    seen_reasoning_delta_item_ids,
+                    output_text_logprobs,
+                )
         if agent_session_id := _extract_foundry_hosted_agent_session_id(getattr(event, "response", None)):
             if update.additional_properties is None:
                 update.additional_properties = {}
@@ -659,6 +669,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,
         *,

python/packages/foundry/agent_framework_foundry/_chat_client.py

@@ -292,12 +292,21 @@ def _parse_chunk_from_openai(
         options: dict[str, Any],
         function_call_ids: dict[int, tuple[str, str]],
         seen_reasoning_delta_item_ids: set[str] | None = None,
+        output_text_logprobs: dict[str, list[Any]] | None = None,
     ) -> ChatResponseUpdate:
         """Parse streaming event, intercepting oauth_consent_request items."""
         update = try_parse_oauth_consent_event(event, self.model)
         if update is not None:
             return update
-        return super()._parse_chunk_from_openai(event, options, function_call_ids, seen_reasoning_delta_item_ids)
+        if output_text_logprobs is None:
+            return super()._parse_chunk_from_openai(event, options, function_call_ids, seen_reasoning_delta_item_ids)
+        return super()._parse_chunk_from_openai(
+            event,
+            options,
+            function_call_ids,
+            seen_reasoning_delta_item_ids,
+            output_text_logprobs,
+        )
 
     async def configure_azure_monitor(
         self,