From 87352e37c193793579448892d5aa0d86c136b9ea Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Wed, 23 Sep 2026 10:29:30 +0200 Subject: [PATCH 1/2] Python: fix Foundry background stateless tool replay --- .../agent_framework_foundry/_chat_client.py | 85 ++++++- .../tests/foundry/test_foundry_chat_client.py | 210 +++++++++++++++++- .../agent_framework_openai/_chat_client.py | 27 ++- 3 files changed, 312 insertions(+), 10 deletions(-) diff --git a/python/packages/foundry/agent_framework_foundry/_chat_client.py b/python/packages/foundry/agent_framework_foundry/_chat_client.py index 3f0d519c025..338a0106f47 100644 --- a/python/packages/foundry/agent_framework_foundry/_chat_client.py +++ b/python/packages/foundry/agent_framework_foundry/_chat_client.py @@ -5,10 +5,11 @@ import logging import sys from collections.abc import Awaitable, Callable, Mapping, Sequence -from typing import TYPE_CHECKING, Any, ClassVar, Generic, Literal +from typing import TYPE_CHECKING, Any, ClassVar, Generic, Literal, cast from agent_framework import ( ChatMiddlewareLayer, + ChatResponse, ChatResponseUpdate, Content, FunctionInvocationConfiguration, @@ -84,6 +85,16 @@ AzureTokenProvider = Callable[[], str | Awaitable[str]] AzureCredentialTypes = TokenCredential | AsyncTokenCredential +_FOUNDRY_REASONING_REPLAY_ITEM_KEY = "__foundry_reasoning_replay_item__" +_FOUNDRY_REASONING_REPLAY_FIELDS = frozenset({ + "type", + "id", + "response_id", + "status", + "summary", + "content", + "encrypted_content", +}) class FoundrySettings(TypedDict, total=False): @@ -270,6 +281,73 @@ async def _prepare_options( run_options.pop("include") return run_options + @override + def _prepare_provider_reasoning_item_for_openai( + self, + reasoning_id: str, + contents: Sequence[Content], + ) -> dict[str, Any] | None: + for content in reversed(contents): + raw_item = content.additional_properties.get(_FOUNDRY_REASONING_REPLAY_ITEM_KEY) + if isinstance(raw_item, Mapping): + typed_raw_item = cast("Mapping[str, Any]", raw_item) + if typed_raw_item.get("id") == reasoning_id: + return dict(typed_raw_item) + return None + + @classmethod + def _prepare_foundry_reasoning_replay_item(cls, item: Any) -> dict[str, Any] | None: + payload = cls._serialize_provider_payload(item) + if not isinstance(payload, Mapping): + return None + typed_payload = cast("Mapping[str, Any]", payload) + if typed_payload.get("type") != "reasoning": + return None + reasoning_id = typed_payload.get("id") + if not isinstance(reasoning_id, str) or not reasoning_id: + return None + response_id = typed_payload.get("response_id") + if not isinstance(response_id, str) or not response_id: + return None + if not isinstance(typed_payload.get("summary"), list) or not isinstance(typed_payload.get("content"), list): + return None + return {key: typed_payload[key] for key in _FOUNDRY_REASONING_REPLAY_FIELDS if key in typed_payload} + + @classmethod + def _attach_foundry_reasoning_replay_item(cls, contents: list[Content], item: Any) -> None: + payload = cls._prepare_foundry_reasoning_replay_item(item) + if payload is None: + return + + reasoning_id = cast(str, payload["id"]) + for content in reversed(contents): + if content.type == "text_reasoning" and content.id == reasoning_id: + content.additional_properties[_FOUNDRY_REASONING_REPLAY_ITEM_KEY] = payload + return + + contents.append( + Content.from_text_reasoning( + id=reasoning_id, + text="", + raw_representation=item, + additional_properties={_FOUNDRY_REASONING_REPLAY_ITEM_KEY: payload}, + ) + ) + + @override + def _parse_response_from_openai( + self, + response: Any, + options: dict[str, Any], + ) -> ChatResponse: + chat_response = super()._parse_response_from_openai(response, options) + if chat_response.messages: + for item in getattr(response, "output", ()): + if getattr(item, "type", None) != "reasoning": + continue + self._attach_foundry_reasoning_replay_item(chat_response.messages[0].contents, item) + return chat_response + @override def _check_model_presence(self, options: dict[str, Any]) -> None: if not options.get("model"): @@ -297,7 +375,10 @@ def _parse_chunk_from_openai( 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) + update = super()._parse_chunk_from_openai(event, options, function_call_ids, seen_reasoning_delta_item_ids) + if event.type == "response.output_item.done" and getattr(event.item, "type", None) == "reasoning": + self._attach_foundry_reasoning_replay_item(update.contents, event.item) + return update async def configure_azure_monitor( self, diff --git a/python/packages/foundry/tests/foundry/test_foundry_chat_client.py b/python/packages/foundry/tests/foundry/test_foundry_chat_client.py index 480242fd017..48fd703fa04 100644 --- a/python/packages/foundry/tests/foundry/test_foundry_chat_client.py +++ b/python/packages/foundry/tests/foundry/test_foundry_chat_client.py @@ -28,7 +28,12 @@ from agent_framework._sessions import AgentSession from agent_framework._telemetry import get_user_agent, mark_feature_used from agent_framework.exceptions import ChatClientException, ChatClientInvalidRequestException -from agent_framework_openai import OpenAIChatClient, OpenAIContentFilterException +from agent_framework_openai import ( + OpenAIChatClient, + OpenAIChatOptions, + OpenAIContentFilterException, + OpenAIContinuationToken, +) from agent_framework_openai._chat_client import RawOpenAIChatClient from azure.ai.projects.models import MCPTool as FoundryMCPTool from azure.core.exceptions import ResourceNotFoundError @@ -53,6 +58,16 @@ class OutputStruct(BaseModel): weather: str | None = None +class FoundryReasoningItem(BaseModel): + """Minimal Foundry reasoning item used to verify provider-native replay.""" + + type: str = "reasoning" + id: str + response_id: str + summary: list[Any] + content: list[Any] + + @pytest.mark.parametrize( "url", [ @@ -526,6 +541,146 @@ async def create_response(**kwargs: Any) -> Any: assert response.response_id == "response_123" +async def test_stateless_replay_uses_foundry_reasoning_item_without_encrypted_content() -> None: + reasoning_item = FoundryReasoningItem( + id="rs_foundry", + response_id="resp_background", + summary=[], + content=[], + ) + function_call = MagicMock( + type="function_call", + id="fc_foundry", + call_id="call_foundry", + arguments='{"step":1}', + status="completed", + ) + function_call.name = "lookup_probe" + mock_response = MagicMock( + id="resp_background", + model="test-model", + created_at=1000000000, + metadata={}, + output_parsed=None, + output=[reasoning_item, function_call], + usage=None, + conversation=None, + status="completed", + incomplete_details=None, + ) + + project_client = MagicMock() + project_client.get_openai_client.return_value = _make_mock_openai_client() + client = FoundryChatClient(project_client=project_client, model="test-model") + response = client._parse_response_from_openai(mock_response, options={"store": False}) + + messages = [ + Message(role="user", contents=["Run the probe."]), + *response.messages, + Message( + role="tool", + contents=[Content.from_function_result(call_id="call_foundry", result="probe-result-step-1")], + ), + ] + _, run_options, _ = await client._prepare_request(messages, {"store": False}) + + assert run_options["input"] == [ + { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "Run the probe."}], + }, + { + "type": "reasoning", + "id": "rs_foundry", + "response_id": "resp_background", + "summary": [], + "content": [], + }, + { + "call_id": "call_foundry", + "id": "fc_foundry", + "type": "function_call", + "name": "lookup_probe", + "arguments": '{"step":1}', + "status": "completed", + }, + { + "call_id": "call_foundry", + "type": "function_call_output", + "output": "probe-result-step-1", + }, + ] + + +def test_streaming_stateless_replay_prefers_completed_foundry_reasoning_item() -> None: + project_client = MagicMock() + project_client.get_openai_client.return_value = _make_mock_openai_client() + client = FoundryChatClient(project_client=project_client, model="test-model") + + added = MagicMock() + added.type = "response.output_item.added" + added.output_index = 0 + added.item = FoundryReasoningItem( + id="rs_foundry", + response_id="resp_incomplete", + summary=[], + content=[], + ) + done = MagicMock() + done.type = "response.output_item.done" + done.output_index = 0 + done.item = FoundryReasoningItem( + id="rs_foundry", + response_id="resp_complete", + summary=[], + content=[], + ) + + response = ChatResponse.from_updates([ + client._parse_chunk_from_openai(added, {}, {}), + client._parse_chunk_from_openai(done, {}, {}), + ]) + + reasoning_items = client._prepare_reasoning_items_for_openai(response.messages) + + assert reasoning_items == { + "rs_foundry": { + "type": "reasoning", + "id": "rs_foundry", + "response_id": "resp_complete", + "summary": [], + "content": [], + } + } + + +async def test_stateless_replay_rejects_incomplete_foundry_reasoning_item() -> None: + project_client = MagicMock() + project_client.get_openai_client.return_value = _make_mock_openai_client() + client = FoundryChatClient(project_client=project_client, model="test-model") + messages = [ + Message( + role="assistant", + contents=[ + Content.from_text_reasoning(id="rs_incomplete", text=""), + Content.from_function_call( + call_id="call_incomplete", + name="lookup_probe", + arguments='{"step":1}', + ), + ], + ), + Message( + role="tool", + contents=[Content.from_function_result(call_id="call_incomplete", result="probe-result-step-1")], + ), + ] + + with pytest.raises(ChatClientInvalidRequestException, match="rs_incomplete"): + await client._prepare_request(messages, {"store": False}) + + async def test_web_search_tool_with_location() -> None: mock_openai_client = _make_mock_openai_client() project_client = MagicMock() @@ -1144,6 +1299,59 @@ async def test_integration_options( assert "location" in response.value +@pytest.mark.flaky +@pytest.mark.integration +@skip_if_foundry_integration_tests_disabled +@pytest.mark.parametrize("store", [True, False]) +@_with_foundry_debug() +async def test_integration_background_local_tool_loop(store: bool) -> None: + executions: list[int] = [] + + @tool(approval_mode="never_require") + def lookup_probe(step: int) -> str: + """Return a deterministic marker for the requested step.""" + executions.append(step) + return f"probe-result-step-{step}" + + client = FoundryChatClient(credential=cast(Any, AzureCliCredential())) + client.function_invocation_configuration["max_iterations"] = 4 + agent = Agent( + client=client, + name=f"background_tool_probe_store_{str(store).lower()}", + instructions=( + "Call lookup_probe with step=1 exactly once. " + "After receiving the tool result, return that exact result and do not call any tool again." + ), + tools=[lookup_probe], + default_options=OpenAIChatOptions( + store=store, + allow_multiple_tool_calls=False, + tool_choice={"mode": "required", "required_function_name": "lookup_probe"}, + ), + ) + session = AgentSession() + + response = await agent.run( + "Run the probe.", + session=session, + options=OpenAIChatOptions(background=True), + ) + deadline = asyncio.get_running_loop().time() + 180 + while response.continuation_token is not None: + if asyncio.get_running_loop().time() >= deadline: + pytest.fail("Background local-tool probe did not finish within 180 seconds.") + await asyncio.sleep(1) + response = await agent.run( + session=session, + options=OpenAIChatOptions( + continuation_token=cast(OpenAIContinuationToken, response.continuation_token), + ), + ) + + assert executions == [1] + assert response.text == "probe-result-step-1" + + @pytest.mark.flaky @pytest.mark.integration @skip_if_foundry_integration_tests_disabled diff --git a/python/packages/openai/agent_framework_openai/_chat_client.py b/python/packages/openai/agent_framework_openai/_chat_client.py index 636fd38893c..71ef265a43a 100644 --- a/python/packages/openai/agent_framework_openai/_chat_client.py +++ b/python/packages/openai/agent_framework_openai/_chat_client.py @@ -1731,8 +1731,11 @@ def _prepare_messages_for_openai( """ reasoning_items: dict[str, dict[str, Any]] = {} if not request_uses_service_side_storage: - self._validate_reasoning_groups_for_stateless_replay(chat_messages) reasoning_items = self._prepare_reasoning_items_for_openai(chat_messages) + self._validate_reasoning_groups_for_stateless_replay( + chat_messages, + replayable_reasoning_ids=set(reasoning_items), + ) serialized_reasoning_ids: set[str] = set() list_of_list = [ @@ -1750,7 +1753,12 @@ def _prepare_messages_for_openai( # items (drop unmatched). See `_AF_MCP_PENDING_OUTPUT_KEY`. return self._coalesce_pending_mcp_results(flat) - def _validate_reasoning_groups_for_stateless_replay(self, chat_messages: Sequence[Message]) -> None: + def _validate_reasoning_groups_for_stateless_replay( + self, + chat_messages: Sequence[Message], + *, + replayable_reasoning_ids: set[str], + ) -> None: """Reject reasoning-bound tool groups that cannot be reconstructed.""" group_reasoning_contents: dict[str, list[Content]] = {} group_call_ids: dict[str, list[str]] = {} @@ -1807,11 +1815,6 @@ def _validate_reasoning_groups_for_stateless_replay(self, chat_messages: Sequenc if not call_ids: continue reasoning_contents = group_reasoning_contents.get(group_id, []) - replayable_reasoning_ids = { - content.id - for content in reasoning_contents - if content.id and (content.protected_data or content.additional_properties.get("encrypted_content")) - } missing_reasoning_ids = list( dict.fromkeys( content.id or "" @@ -1989,6 +1992,8 @@ def _prepare_reasoning_items_for_openai( None, ) if not encrypted_content: + if provider_item := self._prepare_provider_reasoning_item_for_openai(reasoning_id, contents): + reasoning_items[reasoning_id] = provider_item continue item: dict[str, Any] = { @@ -2013,6 +2018,14 @@ def _prepare_reasoning_items_for_openai( reasoning_items[reasoning_id] = item return reasoning_items + def _prepare_provider_reasoning_item_for_openai( + self, + reasoning_id: str, + contents: Sequence[Content], + ) -> dict[str, Any] | None: + """Return a provider-native reasoning item for stateless replay, when supported.""" + return None + def _prepare_content_for_openai( self, role: Role | str, From 422a7dd552f8f2de2be042b88eab2b14eb94199e Mon Sep 17 00:00:00 2001 From: eavanvalkenburg Date: Thu, 24 Sep 2026 09:18:20 +0200 Subject: [PATCH 2/2] Python: harden stateless reasoning replay validation --- .../agent_framework_foundry/_chat_client.py | 6 +-- .../tests/foundry/test_foundry_chat_client.py | 42 +++++++++++++++- .../agent_framework_openai/_chat_client.py | 19 ++++++-- .../tests/openai/test_openai_chat_client.py | 48 +++++++++++++++++++ 4 files changed, 107 insertions(+), 8 deletions(-) diff --git a/python/packages/foundry/agent_framework_foundry/_chat_client.py b/python/packages/foundry/agent_framework_foundry/_chat_client.py index 338a0106f47..a064839cf70 100644 --- a/python/packages/foundry/agent_framework_foundry/_chat_client.py +++ b/python/packages/foundry/agent_framework_foundry/_chat_client.py @@ -290,9 +290,9 @@ def _prepare_provider_reasoning_item_for_openai( for content in reversed(contents): raw_item = content.additional_properties.get(_FOUNDRY_REASONING_REPLAY_ITEM_KEY) if isinstance(raw_item, Mapping): - typed_raw_item = cast("Mapping[str, Any]", raw_item) - if typed_raw_item.get("id") == reasoning_id: - return dict(typed_raw_item) + payload = self._prepare_foundry_reasoning_replay_item(raw_item) + if payload is not None and payload["id"] == reasoning_id: + return payload return None @classmethod diff --git a/python/packages/foundry/tests/foundry/test_foundry_chat_client.py b/python/packages/foundry/tests/foundry/test_foundry_chat_client.py index 48fd703fa04..abfc37c6827 100644 --- a/python/packages/foundry/tests/foundry/test_foundry_chat_client.py +++ b/python/packages/foundry/tests/foundry/test_foundry_chat_client.py @@ -677,7 +677,47 @@ async def test_stateless_replay_rejects_incomplete_foundry_reasoning_item() -> N ), ] - with pytest.raises(ChatClientInvalidRequestException, match="rs_incomplete"): + with pytest.raises( + ChatClientInvalidRequestException, + match="rs_incomplete.*required reasoning replay data is missing or invalid", + ): + await client._prepare_request(messages, {"store": False}) + + +async def test_stateless_replay_rejects_malformed_stored_foundry_reasoning_item() -> None: + project_client = MagicMock() + project_client.get_openai_client.return_value = _make_mock_openai_client() + client = FoundryChatClient(project_client=project_client, model="test-model") + messages = [ + Message( + role="assistant", + contents=[ + Content.from_text_reasoning( + id="rs_malformed", + text="", + additional_properties={ + "__foundry_reasoning_replay_item__": { + "id": "rs_malformed", + } + }, + ), + Content.from_function_call( + call_id="call_malformed", + name="lookup_probe", + arguments='{"step":1}', + ), + ], + ), + Message( + role="tool", + contents=[Content.from_function_result(call_id="call_malformed", result="probe-result-step-1")], + ), + ] + + with pytest.raises( + ChatClientInvalidRequestException, + match="rs_malformed.*required reasoning replay data is missing or invalid", + ): await client._prepare_request(messages, {"store": False}) diff --git a/python/packages/openai/agent_framework_openai/_chat_client.py b/python/packages/openai/agent_framework_openai/_chat_client.py index 71ef265a43a..6ead6a55895 100644 --- a/python/packages/openai/agent_framework_openai/_chat_client.py +++ b/python/packages/openai/agent_framework_openai/_chat_client.py @@ -1810,6 +1810,15 @@ def _validate_reasoning_groups_for_stateless_replay( group_call_ids.setdefault(group_id, []).extend(call_ids) invalid_groups: list[str] = [] + reasoning_id_groups: dict[str, set[str]] = {} + for group_id, reasoning_contents in group_reasoning_contents.items(): + for content in reasoning_contents: + if content.id: + reasoning_id_groups.setdefault(content.id, set()).add(group_id) + reasoning_ids_reused_across_groups = { + reasoning_id for reasoning_id, group_ids in reasoning_id_groups.items() if len(group_ids) > 1 + } + for group_id in groups_with_reasoning: call_ids = list(dict.fromkeys(group_call_ids.get(group_id, []))) if not call_ids: @@ -1819,7 +1828,9 @@ def _validate_reasoning_groups_for_stateless_replay( dict.fromkeys( content.id or "" for content in reasoning_contents - if not content.id or content.id not in replayable_reasoning_ids + if not content.id + or content.id not in replayable_reasoning_ids + or content.id in reasoning_ids_reused_across_groups ) ) if not reasoning_contents: @@ -1831,9 +1842,9 @@ def _validate_reasoning_groups_for_stateless_replay( if invalid_groups: raise ChatClientInvalidRequestException( - f"Stateless replay cannot reconstruct {'; '.join(invalid_groups)} because encrypted reasoning " - "content is missing. Use service-side continuation or explicitly configured atomic compaction to " - "exclude each complete reasoning/tool-call group." + f"Stateless replay cannot reconstruct {'; '.join(invalid_groups)} because required reasoning replay " + "data is missing or invalid. Use service-side continuation or explicitly configured atomic " + "compaction to exclude each complete reasoning/tool-call group." ) def _prepare_message_for_openai( 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 92b68650dac..ea0cdb49113 100644 --- a/python/packages/openai/tests/openai/test_openai_chat_client.py +++ b/python/packages/openai/tests/openai/test_openai_chat_client.py @@ -1899,11 +1899,59 @@ async def test_stateless_reasoning_group_without_encrypted_content_is_rejected_b assert "rs_missing" in message assert "call_one" in message assert "call_two" in message + assert "required reasoning replay data is missing or invalid" in message assert "service-side continuation" in message assert "atomic compaction" in message create.assert_not_awaited() +async def test_stateless_reasoning_id_reused_across_groups_is_rejected_before_transport() -> None: + client = OpenAIChatClient(model="test-model", api_key="test-key") + create = AsyncMock() + messages = [ + Message( + role="assistant", + contents=[ + Content.from_text_reasoning(id="rs_reused", text="First tool group."), + Content.from_function_call(call_id="call_first", name="first_tool", arguments="{}"), + ], + ), + Message( + role="tool", + contents=[Content.from_function_result(call_id="call_first", result="first result")], + ), + Message(role="user", contents=["Start a separate tool group."]), + Message( + role="assistant", + contents=[ + Content.from_text_reasoning( + id="rs_reused", + text="Second tool group.", + protected_data="encrypted-second-group", + ), + Content.from_function_call(call_id="call_second", name="second_tool", arguments="{}"), + ], + ), + Message( + role="tool", + contents=[Content.from_function_result(call_id="call_second", result="second result")], + ), + ] + + with ( + patch.object(client.client.responses.with_raw_response, "create", new=create), + pytest.raises(ChatClientInvalidRequestException) as exc_info, + ): + await client.get_response(messages, options={"store": False}) + + message = str(exc_info.value) + assert "rs_reused" in message + assert "call_first" in message + assert "call_second" in message + assert "required reasoning replay data is missing or invalid" in message + create.assert_not_awaited() + + async def test_partially_compacted_reasoning_group_is_rejected_before_transport() -> None: async def exclude_reasoning_message(messages: list[Message]) -> bool: messages[0].additional_properties["_excluded"] = True