From 038a35645475de87324c7b2cf4ad4ae2e569a280 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Thu, 27 Aug 2026 09:19:16 +0200 Subject: [PATCH 1/8] fix(python): isolate stateless AG-UI sessions Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../ag-ui/agent_framework_ag_ui/_agent_run.py | 25 ++++++-- .../foundry_hosting/tests/test_responses.py | 60 +++++++++++++++++++ 2 files changed, 80 insertions(+), 5 deletions(-) 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 24ba9a2641..d180231240 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 @@ -2109,7 +2109,12 @@ 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, +) -> None: """Restore typed private state from trusted snapshot storage.""" if snapshot is None or snapshot.session_state is None: return @@ -2130,7 +2135,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) @@ -2184,6 +2189,7 @@ def _serialize_session_continuation_state( 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", [])) @@ -2194,12 +2200,13 @@ def _serialize_session_continuation_state( *(provider.source_id for provider in context_providers if isinstance(provider, HistoryProvider)), } 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,6 +2221,7 @@ 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: @@ -2221,6 +2229,7 @@ def _safe_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 @@ -2557,7 +2566,11 @@ 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, + ) protected_session_state_keys = _request_state_protected_keys(agent) session.state.update( { @@ -2693,6 +2706,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 +3065,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) diff --git a/python/packages/foundry_hosting/tests/test_responses.py b/python/packages/foundry_hosting/tests/test_responses.py index cd7cdef316..21d99b5d76 100644 --- a/python/packages/foundry_hosting/tests/test_responses.py +++ b/python/packages/foundry_hosting/tests/test_responses.py @@ -48,6 +48,7 @@ tool, ) from agent_framework.ag_ui import AgentFrameworkAgent, InMemoryAGUIThreadSnapshotStore +from agent_framework.exceptions import ChatClientException from agent_framework.openai import OpenAIChatClient from azure.ai.agentserver.core import get_request_context from azure.ai.agentserver.responses import ( @@ -551,6 +552,65 @@ async def test_agui_service_storage_response_mode_persists_provider_continuation assert [message["role"] for message in second_snapshot] == ["user", "assistant", "user", "assistant"] +async def test_agui_stateless_store_true_does_not_restore_provider_continuation() -> None: + """Stateless snapshot replay must not combine full history with a stored response id.""" + hosted_agent_backend = _make_agent( + response=AgentResponse(messages=[Message(role="assistant", contents=[Content.from_text("ACK")])]) + ) + transport = _CapturingASGITransport(_make_server(hosted_agent_backend)) + responses_client = AsyncOpenAI( + api_key="test-key", + base_url="http://test", + http_client=httpx.AsyncClient(transport=transport), + max_retries=0, + ) + store = InMemoryAGUIThreadSnapshotStore() + runner = AgentFrameworkAgent( + agent=Agent( + client=OpenAIChatClient( # ty: ignore[invalid-argument-type] + model="test-model", + async_client=responses_client, + ), + default_options={"store": True}, + ), + snapshot_store=store, + ) + thread_id = "stateless-thread" + + try: + first_events = [ + event + async for event in runner.run({ + "thread_id": thread_id, + "__ag_ui_snapshot_scope": "test", + "messages": [{"role": "user", "content": "first"}], + }) + ] + first_snapshot = next( + event.model_dump(by_alias=True)["messages"] + for event in reversed(first_events) + if getattr(event, "type", None) == "MESSAGES_SNAPSHOT" + ) + with pytest.raises(ChatClientException, match="schema validation"): + _ = [ + event + async for event in runner.run({ + "thread_id": thread_id, + "__ag_ui_snapshot_scope": "test", + "messages": [*first_snapshot, {"role": "user", "content": "second"}], + }) + ] + finally: + await responses_client.close() + + assert all("conversation" not in payload for payload in transport.payloads) + assert all("previous_response_id" not in payload for payload in transport.payloads) + assert [item["role"] for item in transport.payloads[1]["input"]] == ["user", "assistant", "user"] + stored = await store.get(scope="test", thread_id=thread_id) + assert stored is not None + assert stored.session_state is None + + def _sse_event_types(events: list[dict[str, Any]]) -> list[str]: """Extract event type strings from parsed SSE events.""" return [e["event"] for e in events] From 390f70c9068ae2f19b8e143761236be312ce9175 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Thu, 27 Aug 2026 09:55:59 +0200 Subject: [PATCH 2/8] fix(python): reject disabled service storage Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../ag-ui/agent_framework_ag_ui/_agent_run.py | 5 ++++ python/packages/ag-ui/tests/ag_ui/test_run.py | 26 +++++++++++++++++++ 2 files changed, 31 insertions(+) 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 d180231240..ea5a00211f 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 @@ -2541,6 +2541,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." 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 3408dc4257..9d01839d91 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_run.py +++ b/python/packages/ag-ui/tests/ag_ui/test_run.py @@ -2658,3 +2658,29 @@ 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"}], + } + ) + ] From 0b6c49af707fc1443b7ede819c75f4ac41056855 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Thu, 27 Aug 2026 11:43:51 +0200 Subject: [PATCH 3/8] fix(python): preserve Responses replay metadata Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../foundry/test_agui_provider_matrix.py | 244 ++++++++++++++++++ .../agent_framework_openai/_chat_client.py | 13 +- .../tests/openai/test_agui_provider_matrix.py | 215 +++++++++++++++ .../tests/openai/test_openai_chat_client.py | 50 ++++ 4 files changed, 521 insertions(+), 1 deletion(-) create mode 100644 python/packages/foundry/tests/foundry/test_agui_provider_matrix.py create mode 100644 python/packages/openai/tests/openai/test_agui_provider_matrix.py diff --git a/python/packages/foundry/tests/foundry/test_agui_provider_matrix.py b/python/packages/foundry/tests/foundry/test_agui_provider_matrix.py new file mode 100644 index 0000000000..2a23d51236 --- /dev/null +++ b/python/packages/foundry/tests/foundry/test_agui_provider_matrix.py @@ -0,0 +1,244 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Live AG-UI multi-turn coverage for Foundry provider continuation modes.""" + +from __future__ import annotations + +import json +import os +from typing import Any, Literal, cast +from uuid import uuid4 + +import pytest +from agent_framework import Agent, AgentSession, Message +from agent_framework_ag_ui import AgentFrameworkAgent, InMemoryAGUIThreadSnapshotStore +from azure.ai.projects import models as projects_models +from azure.ai.projects.aio import AIProjectClient +from azure.identity import AzureCliCredential + +from agent_framework_foundry import FoundryAgent, FoundryChatClient + +skip_if_foundry_integration_tests_disabled = pytest.mark.skipif( + os.getenv("FOUNDRY_PROJECT_ENDPOINT", "") in ("", "https://test-project.services.ai.azure.com/") + or os.getenv("FOUNDRY_MODEL", "") == "", + reason="No real FOUNDRY_PROJECT_ENDPOINT or FOUNDRY_MODEL provided; skipping integration tests.", +) +skip_if_foundry_hosted_agent_integration_tests_disabled = pytest.mark.skipif( + os.getenv("FOUNDRY_PROJECT_ENDPOINT", "") in ("", "https://test-project.services.ai.azure.com/") + or os.getenv("FOUNDRY_AGENT_NAME", "") == "", + reason="No real FOUNDRY_PROJECT_ENDPOINT or FOUNDRY_AGENT_NAME provided; skipping integration tests.", +) + +_Mode = Literal["stateless", "conversation", "previous_response"] +_PROVIDER_SESSION_KEY = "__ag_ui_provider_service_session_id" +_MODES = [ + pytest.param("stateless", False, id="stateless-snapshot-replay"), + pytest.param("conversation", True, id="conversation"), + pytest.param("previous_response", True, id="previous-response-id"), +] +_HOSTED_AGENT_MODES = [ + # TODO: Remove this strict xfail when Foundry Hosted Agents accept standard + # output_text replay without requiring fabricated logprobs. + pytest.param( + "stateless", + False, + id="stateless-snapshot-replay", + marks=pytest.mark.xfail( + strict=True, + reason="Foundry Hosted Agent rejects output_text replay when logprobs are absent", + ), + ), + pytest.param("conversation", True, id="conversation"), + pytest.param("previous_response", True, id="previous-response-id"), +] + + +class _CapturingAgent: + """Delegate to an Agent without exposing its conversation factory.""" + + def __init__(self, agent: Any) -> None: + self._agent = agent + self.inputs: list[list[Message]] = [] + self.service_session_ids: list[Any] = [] + self.created_conversation_ids: list[str] = [] + + def run(self, messages: Any = None, **kwargs: Any) -> Any: + self.inputs.append(list(messages) if isinstance(messages, list) else [messages]) + session = kwargs.get("session") + self.service_session_ids.append(session.service_session_id if session is not None else None) + return self._agent.run(messages, **kwargs) + + def __getattr__(self, name: str) -> Any: + if name == "create_conversation": + raise AttributeError(name) + return getattr(self._agent, name) + + +class _ConversationCapturingAgent(_CapturingAgent): + """Expose backend conversation creation to AgentFrameworkAgent.""" + + async def create_conversation(self, *, session_id: str | None = None) -> AgentSession: + conversation = await self._agent.client.client.conversations.create() + self.created_conversation_ids.append(conversation.id) + return self._agent.get_session(conversation.id, session_id=session_id) + + +async def _exercise_agui_case(agent: Any, mode: _Mode) -> None: + marker = f"AF-AGUI-{uuid4().hex}" + follow_up = "Return only the exact marker from my previous message." + thread_id = f"agui-thread-{uuid4().hex}" + scope = f"agui-scope-{uuid4().hex}" + snapshot_store = InMemoryAGUIThreadSnapshotStore() + capturing = _ConversationCapturingAgent(agent) if mode == "conversation" else _CapturingAgent(agent) + + runner = AgentFrameworkAgent( + agent=cast(Any, capturing), + use_service_session=mode != "stateless", + service_session_id_from_thread_id=False, + snapshot_store=snapshot_store, + ) + + try: + first_events = [ + event + async for event in runner.run({ + "threadId": thread_id, + "runId": f"run-1-{uuid4().hex}", + "__ag_ui_snapshot_scope": scope, + "messages": [ + { + "role": "user", + "content": f"Remember this exact marker for my next message: {marker}", + } + ], + }) + ] + first_snapshot = next( + event for event in reversed(first_events) if getattr(event, "type", None) == "MESSAGES_SNAPSHOT" + ) + replay_messages = cast(list[dict[str, Any]], first_snapshot.model_dump(by_alias=True)["messages"]) + first_stored = await snapshot_store.get(scope=scope, thread_id=thread_id) + assert first_stored is not None + + second_events = [ + event + async for event in runner.run({ + "threadId": thread_id, + "runId": f"run-2-{uuid4().hex}", + "__ag_ui_snapshot_scope": scope, + "messages": [*replay_messages, {"role": "user", "content": follow_up}], + }) + ] + + assert not [event for event in second_events if getattr(event, "type", None) == "RUN_ERROR"] + assert [[message.role for message in messages] for messages in capturing.inputs[:1]] == [["user"]] + if mode == "stateless": + assert [message.role for message in capturing.inputs[1]] == ["user", "assistant", "user"] + assert capturing.service_session_ids == [None, None] + assert first_stored.session_state is None + else: + assert [(message.role, message.text) for message in capturing.inputs[1]] == [("user", follow_up)] + assert first_stored.session_state is not None + provider_id = first_stored.session_state[_PROVIDER_SESSION_KEY] + assert isinstance(provider_id, str) + assert provider_id != thread_id + assert capturing.service_session_ids[1] == provider_id + assert provider_id not in json.dumps(first_stored.messages) + if mode == "conversation": + assert provider_id.startswith("conv_") + assert capturing.service_session_ids == [provider_id, provider_id] + assert capturing.created_conversation_ids == [provider_id] + else: + assert provider_id.startswith(("resp_", "caresp_")) + assert capturing.service_session_ids[0] is None + assert not capturing.created_conversation_ids + + assert marker in capturing.inputs[0][0].text + assert capturing.inputs[1][-1].text == follow_up + response_text = "".join( + str(getattr(event, "delta", "")) + for event in second_events + if getattr(event, "type", None) == "TEXT_MESSAGE_CONTENT" + ) + assert marker in response_text + + final_stored = await snapshot_store.get(scope=scope, thread_id=thread_id) + assert final_stored is not None + assert [message.get("role") for message in final_stored.messages].count("user") == 2 + assert [message.get("role") for message in final_stored.messages].count("assistant") >= 2 + finally: + for conversation_id in capturing.created_conversation_ids: + await agent.client.client.conversations.delete(conversation_id) + + +@pytest.mark.flaky +@pytest.mark.integration +@skip_if_foundry_integration_tests_disabled +@pytest.mark.parametrize(("mode", "store"), _MODES) +async def test_foundry_chat_client_agui_provider_matrix(mode: _Mode, store: bool) -> None: + credential = AzureCliCredential() + client = FoundryChatClient(credential=cast(Any, credential)) + agent = Agent(client=cast(Any, client), default_options=cast(Any, {"store": store})) + try: + await _exercise_agui_case(agent, mode) + finally: + await client.client.close() + await client.project_client.close() + credential.close() + + +@pytest.mark.flaky +@pytest.mark.integration +@skip_if_foundry_integration_tests_disabled +@pytest.mark.parametrize(("mode", "store"), _MODES) +async def test_foundry_prompt_agent_agui_provider_matrix(mode: _Mode, store: bool) -> None: + credential = AzureCliCredential() + project_client = AIProjectClient( + endpoint=os.environ["FOUNDRY_PROJECT_ENDPOINT"], + credential=cast(Any, credential), + allow_preview=True, + ) + created_agent: Any | None = None + prompt_agent: FoundryAgent | None = None + try: + created_agent = await project_client.agents.create_version( + agent_name=f"af-agui-{uuid4().hex[:12]}", + definition=projects_models.PromptAgentDefinition( + model=os.environ["FOUNDRY_MODEL"], + instructions="Follow the user instructions exactly and answer concisely.", + ), + ) + prompt_agent = FoundryAgent( + project_client=project_client, + agent_name=created_agent.name, + agent_version=created_agent.version, + allow_preview=False, + default_options=cast(Any, {"store": store}), + ) + await _exercise_agui_case(prompt_agent, mode) + finally: + if created_agent is not None: + await project_client.agents.delete(agent_name=created_agent.name, force=True) + if prompt_agent is not None: + await cast(Any, prompt_agent.client).client.close() + await project_client.close() + credential.close() + + +@pytest.mark.flaky +@pytest.mark.integration +@skip_if_foundry_hosted_agent_integration_tests_disabled +@pytest.mark.parametrize(("mode", "store"), _HOSTED_AGENT_MODES) +async def test_foundry_hosted_agent_agui_provider_matrix(mode: _Mode, store: bool) -> None: + credential = AzureCliCredential() + hosted_agent = FoundryAgent( + credential=cast(Any, credential), + allow_preview=True, + default_options=cast(Any, {"store": store}), + ) + try: + await _exercise_agui_case(hosted_agent, mode) + finally: + await cast(Any, hosted_agent.client).client.close() + await cast(Any, hosted_agent.client).close() + credential.close() diff --git a/python/packages/openai/agent_framework_openai/_chat_client.py b/python/packages/openai/agent_framework_openai/_chat_client.py index 55ff20950a..9607d4de7e 100644 --- a/python/packages/openai/agent_framework_openai/_chat_client.py +++ b/python/packages/openai/agent_framework_openai/_chat_client.py @@ -1866,11 +1866,16 @@ def _prepare_content_for_openai( if role == "assistant": # Assistant history is represented as output text items; Azure validation # requires `annotations` to be present for this type. - return { + output_text = { "type": "output_text", "text": content.text, "annotations": _annotations_to_output_text(getattr(content, "annotations", None)), } + if "logprobs" in content.additional_properties: + output_text["logprobs"] = self._serialize_provider_payload( + content.additional_properties["logprobs"] + ) + return output_text return _attach_prompt_cache_breakpoint( { "type": "input_text", @@ -2614,8 +2619,14 @@ def _parse_response_from_openai( for message_content in item.content: # type: ignore[reportMissingTypeArgument] match message_content.type: case "output_text": + logprobs = getattr(message_content, "logprobs", None) text_content = Content.from_text( text=message_content.text, + additional_properties=( + {"logprobs": self._serialize_provider_payload(logprobs)} + if logprobs is not None + else None + ), raw_representation=message_content, ) metadata.update(self._get_metadata_from_response(message_content)) diff --git a/python/packages/openai/tests/openai/test_agui_provider_matrix.py b/python/packages/openai/tests/openai/test_agui_provider_matrix.py new file mode 100644 index 0000000000..f2df7c8f14 --- /dev/null +++ b/python/packages/openai/tests/openai/test_agui_provider_matrix.py @@ -0,0 +1,215 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Live AG-UI multi-turn coverage for OpenAI provider continuation modes.""" + +from __future__ import annotations + +import json +import os +from typing import Any, Literal, cast +from uuid import uuid4 + +import pytest +from agent_framework import Agent, AgentSession, Message +from agent_framework_ag_ui import AgentFrameworkAgent, InMemoryAGUIThreadSnapshotStore + +from agent_framework_openai import OpenAIChatClient, OpenAIChatCompletionClient + +skip_if_openai_integration_tests_disabled = pytest.mark.skipif( + os.getenv("OPENAI_API_KEY", "") in ("", "test-dummy-key"), + reason="No real OPENAI_API_KEY provided; skipping integration tests.", +) + +_Mode = Literal["stateless", "conversation", "previous_response"] +_PROVIDER_SESSION_KEY = "__ag_ui_provider_service_session_id" + + +class _CapturingAgent: + """Delegate to an Agent without exposing its conversation factory.""" + + def __init__(self, agent: Any) -> None: + self._agent = agent + self.inputs: list[list[Message]] = [] + self.service_session_ids: list[Any] = [] + self.created_conversation_ids: list[str] = [] + + def run(self, messages: Any = None, **kwargs: Any) -> Any: + self.inputs.append(list(messages) if isinstance(messages, list) else [messages]) + session = kwargs.get("session") + self.service_session_ids.append(session.service_session_id if session is not None else None) + return self._agent.run(messages, **kwargs) + + def __getattr__(self, name: str) -> Any: + if name == "create_conversation": + raise AttributeError(name) + return getattr(self._agent, name) + + +class _ConversationCapturingAgent(_CapturingAgent): + """Expose backend conversation creation to AgentFrameworkAgent.""" + + async def create_conversation(self, *, session_id: str | None = None) -> AgentSession: + conversation = await self._agent.client.client.conversations.create() + self.created_conversation_ids.append(conversation.id) + return self._agent.get_session(conversation.id, session_id=session_id) + + +async def _exercise_agui_case(agent: Any, mode: _Mode) -> None: + marker = f"AF-AGUI-{uuid4().hex}" + follow_up = "Return only the exact marker from my previous message." + thread_id = f"agui-thread-{uuid4().hex}" + scope = f"agui-scope-{uuid4().hex}" + snapshot_store = InMemoryAGUIThreadSnapshotStore() + capturing = _ConversationCapturingAgent(agent) if mode == "conversation" else _CapturingAgent(agent) + + runner = AgentFrameworkAgent( + agent=cast(Any, capturing), + use_service_session=mode != "stateless", + service_session_id_from_thread_id=False, + snapshot_store=snapshot_store, + ) + + try: + first_events = [ + event + async for event in runner.run({ + "threadId": thread_id, + "runId": f"run-1-{uuid4().hex}", + "__ag_ui_snapshot_scope": scope, + "messages": [ + { + "role": "user", + "content": f"Remember this exact marker for my next message: {marker}", + } + ], + }) + ] + first_snapshot = next( + event for event in reversed(first_events) if getattr(event, "type", None) == "MESSAGES_SNAPSHOT" + ) + replay_messages = cast(list[dict[str, Any]], first_snapshot.model_dump(by_alias=True)["messages"]) + first_stored = await snapshot_store.get(scope=scope, thread_id=thread_id) + assert first_stored is not None + + second_events = [ + event + async for event in runner.run({ + "threadId": thread_id, + "runId": f"run-2-{uuid4().hex}", + "__ag_ui_snapshot_scope": scope, + "messages": [*replay_messages, {"role": "user", "content": follow_up}], + }) + ] + + assert not [event for event in second_events if getattr(event, "type", None) == "RUN_ERROR"] + assert [[message.role for message in messages] for messages in capturing.inputs[:1]] == [["user"]] + if mode == "stateless": + assert [message.role for message in capturing.inputs[1]] == ["user", "assistant", "user"] + assert capturing.service_session_ids == [None, None] + assert first_stored.session_state is None + else: + assert [(message.role, message.text) for message in capturing.inputs[1]] == [("user", follow_up)] + assert first_stored.session_state is not None + provider_id = first_stored.session_state[_PROVIDER_SESSION_KEY] + assert isinstance(provider_id, str) + assert provider_id != thread_id + assert capturing.service_session_ids[1] == provider_id + assert provider_id not in json.dumps(first_stored.messages) + if mode == "conversation": + assert provider_id.startswith("conv_") + assert capturing.service_session_ids == [provider_id, provider_id] + assert capturing.created_conversation_ids == [provider_id] + else: + assert provider_id.startswith("resp_") + assert capturing.service_session_ids[0] is None + assert not capturing.created_conversation_ids + + assert marker in capturing.inputs[0][0].text + assert capturing.inputs[1][-1].text == follow_up + response_text = "".join( + str(getattr(event, "delta", "")) + for event in second_events + if getattr(event, "type", None) == "TEXT_MESSAGE_CONTENT" + ) + assert marker in response_text + + final_stored = await snapshot_store.get(scope=scope, thread_id=thread_id) + assert final_stored is not None + assert [message.get("role") for message in final_stored.messages].count("user") == 2 + assert [message.get("role") for message in final_stored.messages].count("assistant") >= 2 + finally: + for conversation_id in capturing.created_conversation_ids: + await agent.client.client.conversations.delete(conversation_id) + + +@pytest.mark.flaky +@pytest.mark.integration +@skip_if_openai_integration_tests_disabled +@pytest.mark.parametrize( + ("mode", "store"), + [ + pytest.param("stateless", False, id="stateless-snapshot-replay"), + pytest.param("conversation", True, id="conversation"), + pytest.param("previous_response", True, id="previous-response-id"), + ], +) +async def test_openai_responses_agui_provider_matrix(mode: _Mode, store: bool) -> None: + client = OpenAIChatClient() + agent = Agent(client=cast(Any, client), default_options=cast(Any, {"store": store})) + try: + await _exercise_agui_case(agent, mode) + finally: + await client.client.close() + + +@pytest.mark.flaky +@pytest.mark.integration +@skip_if_openai_integration_tests_disabled +@pytest.mark.parametrize("store", [pytest.param(False, id="store-false"), pytest.param(True, id="store-true")]) +async def test_openai_chat_completions_agui_provider_matrix(store: bool) -> None: + client = OpenAIChatCompletionClient() + agent = Agent(client=cast(Any, client), default_options=cast(Any, {"store": store})) + try: + await _exercise_agui_case(agent, "stateless") + finally: + await client.client.close() + + +@pytest.mark.flaky +@pytest.mark.integration +@skip_if_openai_integration_tests_disabled +async def test_openai_responses_replays_real_assistant_logprobs() -> None: + """Real provider logprobs survive direct assistant-message replay without fabrication.""" + client = OpenAIChatClient() + follow_up = Message(role="user", contents=["Reply with exactly: done"]) + try: + first = await client.get_response( + [Message(role="user", contents=["Reply with exactly: hello"])], + options={ + "store": False, + "include": ["message.output_text.logprobs"], + "top_logprobs": 2, + }, + ) + assistant_content = next( + content + for message in first.messages + for content in message.contents + if message.role == "assistant" and content.type == "text" + ) + real_logprobs = assistant_content.additional_properties.get("logprobs") + assert isinstance(real_logprobs, list) + assert real_logprobs + + _, run_options, _ = await client._prepare_request( + [*first.messages, follow_up], + {"store": False}, + ) + assistant_input = next(item for item in run_options["input"] if item.get("role") == "assistant") + replayed_logprobs = assistant_input["content"][0]["logprobs"] + + assert replayed_logprobs == real_logprobs + second = await client.get_response([*first.messages, follow_up], options={"store": False}) + assert "done" in second.text.lower() + finally: + await client.client.close() diff --git a/python/packages/openai/tests/openai/test_openai_chat_client.py b/python/packages/openai/tests/openai/test_openai_chat_client.py index aab0fc1fe2..da52a9b6db 100644 --- a/python/packages/openai/tests/openai/test_openai_chat_client.py +++ b/python/packages/openai/tests/openai/test_openai_chat_client.py @@ -2881,10 +2881,60 @@ def test_prepare_content_for_openai_text_uses_role_specific_type() -> None: assert user_result["type"] == "input_text" assert assistant_result["type"] == "output_text" assert assistant_result["annotations"] == [] + assert "logprobs" not in assistant_result assert user_result["text"] == "hello" assert assistant_result["text"] == "hello" +def test_prepare_content_for_openai_replays_real_assistant_logprobs() -> None: + """Assistant history replays provider logprobs only when the response supplied them.""" + client = OpenAIChatClient(model="test-model", api_key="test-key") + logprobs = [ + { + "token": "hello", + "bytes": [104, 101, 108, 108, 111], + "logprob": -0.1, + "top_logprobs": [], + } + ] + text_content = Content.from_text(text="hello", additional_properties={"logprobs": logprobs}) + + result = client._prepare_content_for_openai("assistant", text_content) + + assert result["logprobs"] == logprobs + + +def test_parse_and_replay_preserves_real_assistant_logprobs() -> None: + """Parsed Responses logprobs remain attached to assistant content for direct replay.""" + client = OpenAIChatClient(model="test-model", api_key="test-key") + logprobs = [ + { + "token": "hello", + "bytes": [104, 101, 108, 108, 111], + "logprob": -0.1, + "top_logprobs": [], + } + ] + output_text = MagicMock(type="output_text", text="hello", annotations=[], logprobs=logprobs) + output_message = MagicMock(type="message", content=[output_text]) + response = MagicMock( + output_parsed=None, + output=[output_message], + metadata={}, + usage=None, + id="resp-test", + created_at=1_000_000_000, + model="test-model", + ) + + parsed = client._parse_response_from_openai(response, options={"store": False}) + text_content = parsed.messages[0].contents[0] + replayed = client._prepare_content_for_openai("assistant", text_content) + + assert text_content.additional_properties["logprobs"] == logprobs + assert replayed["logprobs"] == logprobs + + def test_prepare_messages_for_openai_assistant_history_uses_output_text_with_annotations() -> None: """Assistant history should be output_text and include required annotations.""" client = OpenAIChatClient(model="test-model", api_key="test-key") From b7ffc316ad641ab04e1b8cea13a1f75c081f7834 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Thu, 27 Aug 2026 14:24:38 +0200 Subject: [PATCH 4/8] fix(python): tighten AG-UI provider state replay Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../ag-ui/agent_framework_ag_ui/_agent_run.py | 15 +++ python/packages/ag-ui/tests/ag_ui/test_run.py | 45 ++++++++- .../foundry/agent_framework_foundry/_agent.py | 24 +++-- .../agent_framework_foundry/_chat_client.py | 11 ++- .../foundry/test_agui_provider_matrix.py | 92 +++++++++++-------- .../agent_framework_openai/_chat_client.py | 43 ++++++++- .../tests/openai/test_openai_chat_client.py | 55 +++++++++++ 7 files changed, 237 insertions(+), 48 deletions(-) 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 ea5a00211f..91641d497f 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 @@ -2114,12 +2114,15 @@ def _restore_session_continuation_state( 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( { @@ -2181,9 +2184,18 @@ 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, @@ -2199,6 +2211,8 @@ 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} service_session_id = session.service_session_id if include_service_session_id else None if not continuation_state and service_session_id is None: @@ -2575,6 +2589,7 @@ async def run_agent_stream( 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( 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 9d01839d91..b593f0ef69 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_run.py +++ b/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 ( @@ -2684,3 +2684,46 @@ async def test_service_session_rejects_disabled_provider_storage(): } ) ] + + +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"}, + ] diff --git a/python/packages/foundry/agent_framework_foundry/_agent.py b/python/packages/foundry/agent_framework_foundry/_agent.py index 0407f9c48f..fc218e9721 100644 --- a/python/packages/foundry/agent_framework_foundry/_agent.py +++ b/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, *, diff --git a/python/packages/foundry/agent_framework_foundry/_chat_client.py b/python/packages/foundry/agent_framework_foundry/_chat_client.py index cbc76852d1..69294e64e4 100644 --- a/python/packages/foundry/agent_framework_foundry/_chat_client.py +++ b/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, diff --git a/python/packages/foundry/tests/foundry/test_agui_provider_matrix.py b/python/packages/foundry/tests/foundry/test_agui_provider_matrix.py index 2a23d51236..8a13ea79c3 100644 --- a/python/packages/foundry/tests/foundry/test_agui_provider_matrix.py +++ b/python/packages/foundry/tests/foundry/test_agui_provider_matrix.py @@ -11,6 +11,7 @@ import pytest from agent_framework import Agent, AgentSession, Message +from agent_framework.exceptions import ChatClientException from agent_framework_ag_ui import AgentFrameworkAgent, InMemoryAGUIThreadSnapshotStore from azure.ai.projects import models as projects_models from azure.ai.projects.aio import AIProjectClient @@ -36,21 +37,6 @@ pytest.param("conversation", True, id="conversation"), pytest.param("previous_response", True, id="previous-response-id"), ] -_HOSTED_AGENT_MODES = [ - # TODO: Remove this strict xfail when Foundry Hosted Agents accept standard - # output_text replay without requiring fabricated logprobs. - pytest.param( - "stateless", - False, - id="stateless-snapshot-replay", - marks=pytest.mark.xfail( - strict=True, - reason="Foundry Hosted Agent rejects output_text replay when logprobs are absent", - ), - ), - pytest.param("conversation", True, id="conversation"), - pytest.param("previous_response", True, id="previous-response-id"), -] class _CapturingAgent: @@ -83,13 +69,19 @@ async def create_conversation(self, *, session_id: str | None = None) -> AgentSe return self._agent.get_session(conversation.id, session_id=session_id) -async def _exercise_agui_case(agent: Any, mode: _Mode) -> None: +async def _exercise_agui_case( + agent: Any, + mode: _Mode, + *, + expect_hosted_stateless_failure: bool = False, +) -> None: marker = f"AF-AGUI-{uuid4().hex}" follow_up = "Return only the exact marker from my previous message." thread_id = f"agui-thread-{uuid4().hex}" scope = f"agui-scope-{uuid4().hex}" snapshot_store = InMemoryAGUIThreadSnapshotStore() capturing = _ConversationCapturingAgent(agent) if mode == "conversation" else _CapturingAgent(agent) + provider_id: str | None = None runner = AgentFrameworkAgent( agent=cast(Any, capturing), @@ -119,40 +111,62 @@ async def _exercise_agui_case(agent: Any, mode: _Mode) -> None: replay_messages = cast(list[dict[str, Any]], first_snapshot.model_dump(by_alias=True)["messages"]) first_stored = await snapshot_store.get(scope=scope, thread_id=thread_id) assert first_stored is not None - - second_events = [ - event - async for event in runner.run({ - "threadId": thread_id, - "runId": f"run-2-{uuid4().hex}", - "__ag_ui_snapshot_scope": scope, - "messages": [*replay_messages, {"role": "user", "content": follow_up}], - }) - ] - - assert not [event for event in second_events if getattr(event, "type", None) == "RUN_ERROR"] - assert [[message.role for message in messages] for messages in capturing.inputs[:1]] == [["user"]] + assert [[message.role for message in messages] for messages in capturing.inputs] == [["user"]] if mode == "stateless": - assert [message.role for message in capturing.inputs[1]] == ["user", "assistant", "user"] - assert capturing.service_session_ids == [None, None] + assert capturing.service_session_ids == [None] assert first_stored.session_state is None else: - assert [(message.role, message.text) for message in capturing.inputs[1]] == [("user", follow_up)] assert first_stored.session_state is not None provider_id = first_stored.session_state[_PROVIDER_SESSION_KEY] assert isinstance(provider_id, str) assert provider_id != thread_id - assert capturing.service_session_ids[1] == provider_id assert provider_id not in json.dumps(first_stored.messages) if mode == "conversation": assert provider_id.startswith("conv_") - assert capturing.service_session_ids == [provider_id, provider_id] + assert capturing.service_session_ids == [provider_id] assert capturing.created_conversation_ids == [provider_id] else: assert provider_id.startswith(("resp_", "caresp_")) - assert capturing.service_session_ids[0] is None + assert capturing.service_session_ids == [None] assert not capturing.created_conversation_ids + try: + second_events = [ + event + async for event in runner.run({ + "threadId": thread_id, + "runId": f"run-2-{uuid4().hex}", + "__ag_ui_snapshot_scope": scope, + "messages": [*replay_messages, {"role": "user", "content": follow_up}], + }) + ] + except ChatClientException as exc: + if not expect_hosted_stateless_failure: + raise + error = str(exc) + assert "request body failed schema validation" in error + assert "Expected one of: string, array; got array" in error + assert "'param': '$.input'" in error + # TODO: Remove this expected-failure branch when Foundry Hosted Agents + # accept standard output_text replay without fabricated logprobs. + pytest.xfail("Foundry Hosted Agent rejects output_text replay when logprobs are absent") + + if expect_hosted_stateless_failure: + pytest.fail("Foundry Hosted Agent replay now works; remove the expected-failure branch") + + assert not [event for event in second_events if getattr(event, "type", None) == "RUN_ERROR"] + if mode == "stateless": + assert [message.role for message in capturing.inputs[1]] == ["user", "assistant", "user"] + assert capturing.service_session_ids == [None, None] + else: + assert [(message.role, message.text) for message in capturing.inputs[1]] == [("user", follow_up)] + assert provider_id is not None + assert capturing.service_session_ids[1] == provider_id + if mode == "conversation": + assert capturing.service_session_ids == [provider_id, provider_id] + else: + assert capturing.service_session_ids[0] is None + assert marker in capturing.inputs[0][0].text assert capturing.inputs[1][-1].text == follow_up response_text = "".join( @@ -228,7 +242,7 @@ async def test_foundry_prompt_agent_agui_provider_matrix(mode: _Mode, store: boo @pytest.mark.flaky @pytest.mark.integration @skip_if_foundry_hosted_agent_integration_tests_disabled -@pytest.mark.parametrize(("mode", "store"), _HOSTED_AGENT_MODES) +@pytest.mark.parametrize(("mode", "store"), _MODES) async def test_foundry_hosted_agent_agui_provider_matrix(mode: _Mode, store: bool) -> None: credential = AzureCliCredential() hosted_agent = FoundryAgent( @@ -237,7 +251,11 @@ async def test_foundry_hosted_agent_agui_provider_matrix(mode: _Mode, store: boo default_options=cast(Any, {"store": store}), ) try: - await _exercise_agui_case(hosted_agent, mode) + await _exercise_agui_case( + hosted_agent, + mode, + expect_hosted_stateless_failure=mode == "stateless", + ) finally: await cast(Any, hosted_agent.client).client.close() await cast(Any, hosted_agent.client).close() diff --git a/python/packages/openai/agent_framework_openai/_chat_client.py b/python/packages/openai/agent_framework_openai/_chat_client.py index 9607d4de7e..51edc5904d 100644 --- a/python/packages/openai/agent_framework_openai/_chat_client.py +++ b/python/packages/openai/agent_framework_openai/_chat_client.py @@ -710,6 +710,7 @@ def _inner_get_response( if stream: function_call_ids: dict[int, tuple[str, str]] = {} seen_reasoning_delta_item_ids: set[str] = set() + output_text_logprobs: dict[str, list[Any]] = {} validated_options: dict[str, Any] | None = None # Captured once request options are validated/prepared so the streaming finalizer can # still parse the aggregated response into structured output after the stream completes. @@ -748,6 +749,7 @@ async def _stream() -> AsyncIterable[ChatResponseUpdate]: options=validated_options, function_call_ids=function_call_ids, seen_reasoning_delta_item_ids=seen_reasoning_delta_item_ids, + output_text_logprobs=output_text_logprobs, ) if served_model is not None: update.model = served_model @@ -778,6 +780,7 @@ async def _stream() -> AsyncIterable[ChatResponseUpdate]: options=validated_options, function_call_ids=function_call_ids, seen_reasoning_delta_item_ids=seen_reasoning_delta_item_ids, + output_text_logprobs=output_text_logprobs, ) else: raw_create_response = await client.responses.with_raw_response.create( @@ -792,6 +795,7 @@ async def _stream() -> AsyncIterable[ChatResponseUpdate]: options=validated_options, function_call_ids=function_call_ids, seen_reasoning_delta_item_ids=seen_reasoning_delta_item_ids, + output_text_logprobs=output_text_logprobs, ) if served_model is not None: update.model = served_model @@ -2619,7 +2623,7 @@ def _parse_response_from_openai( for message_content in item.content: # type: ignore[reportMissingTypeArgument] match message_content.type: case "output_text": - logprobs = getattr(message_content, "logprobs", None) + logprobs = getattr(cast(Any, message_content), "logprobs", None) text_content = Content.from_text( text=message_content.text, additional_properties=( @@ -2903,6 +2907,7 @@ 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 an OpenAI Responses API streaming event into a ChatResponseUpdate.""" metadata: dict[str, Any] = {} @@ -2914,6 +2919,20 @@ def _parse_chunk_from_openai( continuation_token: OpenAIContinuationToken | None = None finish_reason: FinishReason | None = None model = self.model + + def output_text_properties(output: Any, item_id: str) -> dict[str, Any] | None: + logprobs = getattr(output, "logprobs", None) + if logprobs is None: + return None + serialized = self._serialize_provider_payload(logprobs) + if not isinstance(serialized, list): + return None + if output_text_logprobs is None: + return {"logprobs": serialized} + accumulated = output_text_logprobs.setdefault(item_id, []) + accumulated.extend(cast(list[Any], serialized)) + return {"logprobs": accumulated} + match event.type: # types: # ResponseAudioDeltaEvent, @@ -2973,14 +2992,32 @@ def _parse_chunk_from_openai( event_part = event.part match event_part.type: case "output_text": - contents.append(Content.from_text(text=event_part.text, raw_representation=event)) + contents.append( + Content.from_text( + text=event_part.text, + additional_properties=output_text_properties( + cast(Any, event_part), + cast(Any, event).item_id, + ), + raw_representation=event, + ) + ) metadata.update(self._get_metadata_from_response(event_part)) case "refusal": contents.append(Content.from_text(text=event_part.refusal, raw_representation=event)) case _: pass case "response.output_text.delta": - contents.append(Content.from_text(text=event.delta, raw_representation=event)) + contents.append( + Content.from_text( + text=event.delta, + additional_properties=output_text_properties( + cast(Any, event), + cast(Any, event).item_id, + ), + raw_representation=event, + ) + ) metadata.update(self._get_metadata_from_response(event)) case "response.reasoning_text.delta": if seen_reasoning_delta_item_ids is not None: diff --git a/python/packages/openai/tests/openai/test_openai_chat_client.py b/python/packages/openai/tests/openai/test_openai_chat_client.py index da52a9b6db..6431cea8da 100644 --- a/python/packages/openai/tests/openai/test_openai_chat_client.py +++ b/python/packages/openai/tests/openai/test_openai_chat_client.py @@ -2935,6 +2935,61 @@ def test_parse_and_replay_preserves_real_assistant_logprobs() -> None: assert replayed["logprobs"] == logprobs +def test_streaming_parse_and_replay_preserves_all_real_assistant_logprobs() -> None: + """Streamed token logprobs accumulate on assistant content for direct replay.""" + client = OpenAIChatClient(model="test-model", api_key="test-key") + first_logprob = { + "token": "hel", + "bytes": [104, 101, 108], + "logprob": -0.1, + "top_logprobs": [], + } + second_logprob = { + "token": "lo", + "bytes": [108, 111], + "logprob": -0.2, + "top_logprobs": [], + } + events = [ + ResponseTextDeltaEvent( + type="response.output_text.delta", + content_index=0, + item_id="msg-1", + output_index=0, + sequence_number=1, + logprobs=[first_logprob], # type: ignore[list-item] + delta="hel", + ), + ResponseTextDeltaEvent( + type="response.output_text.delta", + content_index=0, + item_id="msg-1", + output_index=0, + sequence_number=2, + logprobs=[second_logprob], # type: ignore[list-item] + delta="lo", + ), + ] + accumulated_logprobs: dict[str, list[Any]] = {} + updates = [ + client._parse_chunk_from_openai( + event, + options={}, + function_call_ids={}, + output_text_logprobs=accumulated_logprobs, + ) + for event in events + ] + + response = ChatResponse.from_updates(updates) + text_content = response.messages[0].contents[0] + replayed = client._prepare_content_for_openai("assistant", text_content) + + assert text_content.text == "hello" + assert text_content.additional_properties["logprobs"] == [first_logprob, second_logprob] + assert replayed["logprobs"] == [first_logprob, second_logprob] + + def test_prepare_messages_for_openai_assistant_history_uses_output_text_with_annotations() -> None: """Assistant history should be output_text and include required annotations.""" client = OpenAIChatClient(model="test-model", api_key="test-key") From 60fdf056ac9903f2472d6c365e39a58e4326705d Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Mon, 31 Aug 2026 10:16:16 +0200 Subject: [PATCH 5/8] fix(python): preserve wrapped provider replay state Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../agent_framework_ag_ui/_a2ui/_agent.py | 4 ++ .../packages/ag-ui/tests/ag_ui/test_a2ui.py | 54 ++++++++++++++++++- .../agent_framework_openai/_chat_client.py | 46 +++++++++------- .../tests/openai/test_openai_chat_client.py | 6 +-- 4 files changed, 86 insertions(+), 24 deletions(-) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_a2ui/_agent.py b/python/packages/ag-ui/agent_framework_ag_ui/_a2ui/_agent.py index 285b2ac0f9..d4a68213ec 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_a2ui/_agent.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_a2ui/_agent.py @@ -351,6 +351,10 @@ 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", ()) + # -- public run ------------------------------------------------------- def run(self, messages: Any = None, *, stream: bool = False, **kwargs: Any) -> Any: diff --git a/python/packages/ag-ui/tests/ag_ui/test_a2ui.py b/python/packages/ag-ui/tests/ag_ui/test_a2ui.py index 1a28334ddd..2919758f05 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_a2ui.py +++ b/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,64 @@ 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_uses_per_request_context_over_constructor(): diff --git a/python/packages/openai/agent_framework_openai/_chat_client.py b/python/packages/openai/agent_framework_openai/_chat_client.py index 51edc5904d..76aee49dad 100644 --- a/python/packages/openai/agent_framework_openai/_chat_client.py +++ b/python/packages/openai/agent_framework_openai/_chat_client.py @@ -710,7 +710,6 @@ def _inner_get_response( if stream: function_call_ids: dict[int, tuple[str, str]] = {} seen_reasoning_delta_item_ids: set[str] = set() - output_text_logprobs: dict[str, list[Any]] = {} validated_options: dict[str, Any] | None = None # Captured once request options are validated/prepared so the streaming finalizer can # still parse the aggregated response into structured output after the stream completes. @@ -749,7 +748,6 @@ async def _stream() -> AsyncIterable[ChatResponseUpdate]: options=validated_options, function_call_ids=function_call_ids, seen_reasoning_delta_item_ids=seen_reasoning_delta_item_ids, - output_text_logprobs=output_text_logprobs, ) if served_model is not None: update.model = served_model @@ -780,7 +778,6 @@ async def _stream() -> AsyncIterable[ChatResponseUpdate]: options=validated_options, function_call_ids=function_call_ids, seen_reasoning_delta_item_ids=seen_reasoning_delta_item_ids, - output_text_logprobs=output_text_logprobs, ) else: raw_create_response = await client.responses.with_raw_response.create( @@ -795,7 +792,6 @@ async def _stream() -> AsyncIterable[ChatResponseUpdate]: options=validated_options, function_call_ids=function_call_ids, seen_reasoning_delta_item_ids=seen_reasoning_delta_item_ids, - output_text_logprobs=output_text_logprobs, ) if served_model is not None: update.model = served_model @@ -865,7 +861,28 @@ def _finalize_response_updates( """Finalize streamed updates and add post-stream Azure AI Search citation metadata.""" self._enrich_streamed_azure_ai_search_citations(updates) self._enrich_mcp_search_citations([content for update in updates for content in update.contents]) - return super()._finalize_response_updates(updates, response_format=response_format) + response = super()._finalize_response_updates(updates, response_format=response_format) + logprobs = [ + logprob + for update in updates + for content in update.contents + if content.type == "text" + for logprob in content.additional_properties.get("logprobs", []) + ] + if logprobs: + assistant_text = next( + ( + content + for message in response.messages + if message.role == "assistant" + for content in message.contents + if content.type == "text" + ), + None, + ) + if assistant_text is not None: + assistant_text.additional_properties["logprobs"] = logprobs + return response @classmethod def _extract_served_model(cls, headers: Any) -> str | None: @@ -2907,7 +2924,6 @@ 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 an OpenAI Responses API streaming event into a ChatResponseUpdate.""" metadata: dict[str, Any] = {} @@ -2920,18 +2936,14 @@ def _parse_chunk_from_openai( finish_reason: FinishReason | None = None model = self.model - def output_text_properties(output: Any, item_id: str) -> dict[str, Any] | None: + def output_text_properties(output: Any) -> dict[str, Any] | None: logprobs = getattr(output, "logprobs", None) if logprobs is None: return None serialized = self._serialize_provider_payload(logprobs) if not isinstance(serialized, list): return None - if output_text_logprobs is None: - return {"logprobs": serialized} - accumulated = output_text_logprobs.setdefault(item_id, []) - accumulated.extend(cast(list[Any], serialized)) - return {"logprobs": accumulated} + return {"logprobs": serialized} match event.type: # types: @@ -2995,10 +3007,7 @@ def output_text_properties(output: Any, item_id: str) -> dict[str, Any] | None: contents.append( Content.from_text( text=event_part.text, - additional_properties=output_text_properties( - cast(Any, event_part), - cast(Any, event).item_id, - ), + additional_properties=output_text_properties(cast(Any, event_part)), raw_representation=event, ) ) @@ -3011,10 +3020,7 @@ def output_text_properties(output: Any, item_id: str) -> dict[str, Any] | None: contents.append( Content.from_text( text=event.delta, - additional_properties=output_text_properties( - cast(Any, event), - cast(Any, event).item_id, - ), + additional_properties=output_text_properties(cast(Any, event)), raw_representation=event, ) ) diff --git a/python/packages/openai/tests/openai/test_openai_chat_client.py b/python/packages/openai/tests/openai/test_openai_chat_client.py index 6431cea8da..8e98ca451f 100644 --- a/python/packages/openai/tests/openai/test_openai_chat_client.py +++ b/python/packages/openai/tests/openai/test_openai_chat_client.py @@ -2970,18 +2970,18 @@ def test_streaming_parse_and_replay_preserves_all_real_assistant_logprobs() -> N delta="lo", ), ] - accumulated_logprobs: dict[str, list[Any]] = {} updates = [ client._parse_chunk_from_openai( event, options={}, function_call_ids={}, - output_text_logprobs=accumulated_logprobs, ) for event in events ] - response = ChatResponse.from_updates(updates) + assert updates[0].contents[0].additional_properties["logprobs"] == [first_logprob] + assert updates[1].contents[0].additional_properties["logprobs"] == [second_logprob] + response = client._finalize_response_updates(updates) text_content = response.messages[0].contents[0] replayed = client._prepare_content_for_openai("assistant", text_content) From ca70704de9337f730b19bb8c9b069e2c5254b1d8 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Mon, 31 Aug 2026 10:16:39 +0200 Subject: [PATCH 6/8] refactor(python): keep stream parsing state internal Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../foundry/agent_framework_foundry/_agent.py | 22 +++++-------------- .../agent_framework_foundry/_chat_client.py | 11 +--------- 2 files changed, 7 insertions(+), 26 deletions(-) diff --git a/python/packages/foundry/agent_framework_foundry/_agent.py b/python/packages/foundry/agent_framework_foundry/_agent.py index fc218e9721..cd44b46088 100644 --- a/python/packages/foundry/agent_framework_foundry/_agent.py +++ b/python/packages/foundry/agent_framework_foundry/_agent.py @@ -439,26 +439,16 @@ 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: - 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, - ) + update = super()._parse_chunk_from_openai( + event, + options, + function_call_ids, + seen_reasoning_delta_item_ids, + ) if agent_session_id := _extract_foundry_hosted_agent_session_id(getattr(event, "response", None)): if update.additional_properties is None: update.additional_properties = {} diff --git a/python/packages/foundry/agent_framework_foundry/_chat_client.py b/python/packages/foundry/agent_framework_foundry/_chat_client.py index 69294e64e4..cbc76852d1 100644 --- a/python/packages/foundry/agent_framework_foundry/_chat_client.py +++ b/python/packages/foundry/agent_framework_foundry/_chat_client.py @@ -292,21 +292,12 @@ 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 - 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, - ) + return super()._parse_chunk_from_openai(event, options, function_call_ids, seen_reasoning_delta_item_ids) async def configure_azure_monitor( self, From 84774126411b80b12840f12bdd72504d1c21416b Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Mon, 31 Aug 2026 10:31:58 +0200 Subject: [PATCH 7/8] fix(python): preserve A2UI provider session behavior Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../agent_framework_ag_ui/_a2ui/_agent.py | 5 +++++ .../packages/ag-ui/tests/ag_ui/test_a2ui.py | 20 +++++++++++++++++++ .../foundry_hosting/tests/test_responses.py | 19 +++++++++--------- 3 files changed, 34 insertions(+), 10 deletions(-) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_a2ui/_agent.py b/python/packages/ag-ui/agent_framework_ag_ui/_a2ui/_agent.py index d4a68213ec..6e444b1c88 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_a2ui/_agent.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_a2ui/_agent.py @@ -355,6 +355,11 @@ def context_providers(self) -> Any: 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: diff --git a/python/packages/ag-ui/tests/ag_ui/test_a2ui.py b/python/packages/ag-ui/tests/ag_ui/test_a2ui.py index 2919758f05..d89b222f00 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_a2ui.py +++ b/python/packages/ag-ui/tests/ag_ui/test_a2ui.py @@ -1249,6 +1249,26 @@ def test_a2ui_agent_preserves_provider_state_authority_boundary(): 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(): # A reused runner must serve the CURRENT request's catalog, not a stale constructor one. old = build_ag_ui_context_slice( diff --git a/python/packages/foundry_hosting/tests/test_responses.py b/python/packages/foundry_hosting/tests/test_responses.py index 21d99b5d76..258ffce8a8 100644 --- a/python/packages/foundry_hosting/tests/test_responses.py +++ b/python/packages/foundry_hosting/tests/test_responses.py @@ -48,7 +48,6 @@ tool, ) from agent_framework.ag_ui import AgentFrameworkAgent, InMemoryAGUIThreadSnapshotStore -from agent_framework.exceptions import ChatClientException from agent_framework.openai import OpenAIChatClient from azure.ai.agentserver.core import get_request_context from azure.ai.agentserver.responses import ( @@ -591,21 +590,21 @@ async def test_agui_stateless_store_true_does_not_restore_provider_continuation( for event in reversed(first_events) if getattr(event, "type", None) == "MESSAGES_SNAPSHOT" ) - with pytest.raises(ChatClientException, match="schema validation"): - _ = [ - event - async for event in runner.run({ - "thread_id": thread_id, - "__ag_ui_snapshot_scope": "test", - "messages": [*first_snapshot, {"role": "user", "content": "second"}], - }) - ] + second_events = [ + event + async for event in runner.run({ + "thread_id": thread_id, + "__ag_ui_snapshot_scope": "test", + "messages": [*first_snapshot, {"role": "user", "content": "second"}], + }) + ] finally: await responses_client.close() assert all("conversation" not in payload for payload in transport.payloads) assert all("previous_response_id" not in payload for payload in transport.payloads) assert [item["role"] for item in transport.payloads[1]["input"]] == ["user", "assistant", "user"] + assert not [event for event in second_events if getattr(event, "type", None) == "RUN_ERROR"] stored = await store.get(scope="test", thread_id=thread_id) assert stored is not None assert stored.session_state is None From c7ef49d0787f932ff3d2831b54364706366fa078 Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Mon, 31 Aug 2026 12:22:13 +0200 Subject: [PATCH 8/8] test(python): update AG-UI merge coverage Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../foundry/test_agui_provider_matrix.py | 41 +++++-------------- .../tests/openai/test_agui_provider_matrix.py | 2 +- 2 files changed, 11 insertions(+), 32 deletions(-) diff --git a/python/packages/foundry/tests/foundry/test_agui_provider_matrix.py b/python/packages/foundry/tests/foundry/test_agui_provider_matrix.py index 8a13ea79c3..681a0af4a1 100644 --- a/python/packages/foundry/tests/foundry/test_agui_provider_matrix.py +++ b/python/packages/foundry/tests/foundry/test_agui_provider_matrix.py @@ -11,7 +11,6 @@ import pytest from agent_framework import Agent, AgentSession, Message -from agent_framework.exceptions import ChatClientException from agent_framework_ag_ui import AgentFrameworkAgent, InMemoryAGUIThreadSnapshotStore from azure.ai.projects import models as projects_models from azure.ai.projects.aio import AIProjectClient @@ -72,8 +71,6 @@ async def create_conversation(self, *, session_id: str | None = None) -> AgentSe async def _exercise_agui_case( agent: Any, mode: _Mode, - *, - expect_hosted_stateless_failure: bool = False, ) -> None: marker = f"AF-AGUI-{uuid4().hex}" follow_up = "Return only the exact marker from my previous message." @@ -130,29 +127,15 @@ async def _exercise_agui_case( assert capturing.service_session_ids == [None] assert not capturing.created_conversation_ids - try: - second_events = [ - event - async for event in runner.run({ - "threadId": thread_id, - "runId": f"run-2-{uuid4().hex}", - "__ag_ui_snapshot_scope": scope, - "messages": [*replay_messages, {"role": "user", "content": follow_up}], - }) - ] - except ChatClientException as exc: - if not expect_hosted_stateless_failure: - raise - error = str(exc) - assert "request body failed schema validation" in error - assert "Expected one of: string, array; got array" in error - assert "'param': '$.input'" in error - # TODO: Remove this expected-failure branch when Foundry Hosted Agents - # accept standard output_text replay without fabricated logprobs. - pytest.xfail("Foundry Hosted Agent rejects output_text replay when logprobs are absent") - - if expect_hosted_stateless_failure: - pytest.fail("Foundry Hosted Agent replay now works; remove the expected-failure branch") + second_events = [ + event + async for event in runner.run({ + "threadId": thread_id, + "runId": f"run-2-{uuid4().hex}", + "__ag_ui_snapshot_scope": scope, + "messages": [*replay_messages, {"role": "user", "content": follow_up}], + }) + ] assert not [event for event in second_events if getattr(event, "type", None) == "RUN_ERROR"] if mode == "stateless": @@ -251,11 +234,7 @@ async def test_foundry_hosted_agent_agui_provider_matrix(mode: _Mode, store: boo default_options=cast(Any, {"store": store}), ) try: - await _exercise_agui_case( - hosted_agent, - mode, - expect_hosted_stateless_failure=mode == "stateless", - ) + await _exercise_agui_case(hosted_agent, mode) finally: await cast(Any, hosted_agent.client).client.close() await cast(Any, hosted_agent.client).close() diff --git a/python/packages/openai/tests/openai/test_agui_provider_matrix.py b/python/packages/openai/tests/openai/test_agui_provider_matrix.py index f2df7c8f14..b3cb99617b 100644 --- a/python/packages/openai/tests/openai/test_agui_provider_matrix.py +++ b/python/packages/openai/tests/openai/test_agui_provider_matrix.py @@ -180,7 +180,7 @@ async def test_openai_chat_completions_agui_provider_matrix(store: bool) -> None @skip_if_openai_integration_tests_disabled async def test_openai_responses_replays_real_assistant_logprobs() -> None: """Real provider logprobs survive direct assistant-message replay without fabrication.""" - client = OpenAIChatClient() + client = OpenAIChatClient(model=os.environ["OPENAI_CHAT_COMPLETION_MODEL"]) follow_up = Message(role="user", contents=["Reply with exactly: done"]) try: first = await client.get_response(