diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py index 422665cbe1..b09471d541 100644 --- a/python/packages/core/agent_framework/_types.py +++ b/python/packages/core/agent_framework/_types.py @@ -1511,6 +1511,12 @@ def _add_text_reasoning_content(self, other: Content) -> Content: # Concatenate text, handling None values self_text = self.text or "" other_text = other.text or "" + if ( + self_text + and other_text + and ("reasoning_text" in self.additional_properties) != ("reasoning_text" in other.additional_properties) + ): + raise AdditionItemMismatch("Cannot merge reasoning text with a reasoning summary") combined_text = self_text + other_text if (self_text or other_text) else None # Handle protected_data replacement diff --git a/python/packages/core/tests/core/test_function_invocation_logic.py b/python/packages/core/tests/core/test_function_invocation_logic.py index 6426fade63..852c0ee50a 100644 --- a/python/packages/core/tests/core/test_function_invocation_logic.py +++ b/python/packages/core/tests/core/test_function_invocation_logic.py @@ -3823,6 +3823,7 @@ def ai_func(arg1: str) -> str: assert response.messages[1].role == "tool" assert response.messages[1].contents[0].type == "function_result" assert response.messages[1].contents[0].result == "terminated by middleware" + assert response.messages[1].contents[0].additional_properties == {} # Verify the second response is still in the queue (wasn't consumed) assert len(chat_client_base.run_responses) == 1 # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] @@ -3891,6 +3892,11 @@ def terminating_func(arg1: str) -> str: assert response.messages[1].role == "tool" # Both function results should be present assert len(response.messages[1].contents) == 2 + assert [result.result for result in response.messages[1].contents] == [ + "Normal value1", + "terminated by middleware", + ] + assert all(result.additional_properties == {} for result in response.messages[1].contents) # Verify the second response is still in the queue (wasn't consumed) assert len(chat_client_base.run_responses) == 1 # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] @@ -3939,6 +3945,12 @@ def ai_func(arg1: str) -> str: # Should have function call update and function result update # The loop should NOT have continued to call the LLM again assert len(updates) == 2 + function_results = [ + content for update in updates for content in update.contents if content.type == "function_result" + ] + assert len(function_results) == 1 + assert function_results[0].result == "terminated by middleware" + assert function_results[0].additional_properties == {} # Verify the second streaming response is still in the queue (wasn't consumed) assert len(chat_client_base.streaming_responses) == 1 # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] diff --git a/python/packages/foundry/agent_framework_foundry/_agent.py b/python/packages/foundry/agent_framework_foundry/_agent.py index 7f080df9f9..8e26bd47aa 100644 --- a/python/packages/foundry/agent_framework_foundry/_agent.py +++ b/python/packages/foundry/agent_framework_foundry/_agent.py @@ -330,6 +330,8 @@ async def _prepare_options( **kwargs: Any, ) -> dict[str, Any]: """Prepare options for the Responses API and validate client-side tools.""" + caller_requested_encrypted_reasoning = "reasoning.encrypted_content" in (options.get("include") or []) + # Validate tools — only FunctionTool allowed tools = options.get("tools", []) if tools: @@ -347,6 +349,16 @@ async def _prepare_options( # Call parent prepare_options (OpenAI Responses API format) run_options = await super()._prepare_options(prepared_messages, options, **kwargs) + # Foundry Agent deployments can reject the OpenAI client's automatic encrypted-reasoning + # opt-in even when the configured model otherwise supports reasoning. Preserve an explicit + # caller request, but do not add this provider capability implicitly. + if not caller_requested_encrypted_reasoning and isinstance(run_options.get("include"), list): + include = [item for item in run_options["include"] if item != "reasoning.encrypted_content"] + if include: + run_options["include"] = include + else: + run_options.pop("include") + # Apply Azure AI schema transforms if "input" in run_options and isinstance(run_options["input"], list): run_options["input"] = self._transform_input_for_azure_ai(cast(list[dict[str, Any]], run_options["input"])) diff --git a/python/packages/foundry/tests/foundry/test_foundry_agent.py b/python/packages/foundry/tests/foundry/test_foundry_agent.py index 914866d4e1..6d335a17e4 100644 --- a/python/packages/foundry/tests/foundry/test_foundry_agent.py +++ b/python/packages/foundry/tests/foundry/test_foundry_agent.py @@ -3,23 +3,31 @@ from __future__ import annotations import inspect +import json import os import sys +from collections.abc import Awaitable, Callable from types import SimpleNamespace from typing import Any, cast from unittest.mock import AsyncMock, MagicMock, patch from uuid import uuid4 +import httpx import pytest from agent_framework import ( Agent, + AgentExecutor, AgentResponse, AgentSession, ChatContext, ChatMiddleware, ChatResponse, ChatResponseUpdate, + FunctionInvocationContext, + FunctionMiddleware, Message, + MiddlewareTermination, + WorkflowBuilder, tool, ) from agent_framework_openai._chat_client import RawOpenAIChatClient @@ -27,6 +35,7 @@ from azure.core.exceptions import ResourceNotFoundError from azure.identity import AzureCliCredential from azure.identity.aio import AzureCliCredential as AsyncAzureCliCredential +from openai import AsyncOpenAI from agent_framework_foundry._agent import ( FoundryAgent, @@ -110,6 +119,70 @@ def test_raw_foundry_agent_chat_client_init_with_agent_name() -> None: mock_project.get_openai_client.assert_called_once_with() +async def test_foundry_agent_basic_call_does_not_request_unsupported_encrypted_reasoning() -> None: + """A Foundry agent call must not opt into encrypted reasoning unless the caller requests it.""" + mock_response = MagicMock() + mock_response.id = "response_123" + mock_response.model = "test-model" + mock_response.created_at = 1000000000 + mock_response.output_parsed = None + mock_response.metadata = {} + mock_response.output = [] + mock_response.usage = None + mock_response.finish_reason = None + mock_response.conversation = None + mock_response.status = "completed" + mock_response.parse.return_value = mock_response + mock_response.headers = {} + + async def create_response(**kwargs: Any) -> Any: + if "reasoning.encrypted_content" in kwargs.get("include", []): + raise ValueError("Encrypted content is not supported with this model.") + return mock_response + + mock_openai = MagicMock() + mock_openai.responses.with_raw_response.create = AsyncMock(side_effect=create_response) + mock_project = MagicMock() + mock_project.get_openai_client.return_value = mock_openai + client = RawFoundryAgentChatClient(project_client=mock_project, agent_name="test-agent") + + response = await client.get_response([Message(role="user", contents=["Hello"])]) + + assert response.response_id == "response_123" + + +async def test_foundry_agent_preserves_caller_requested_encrypted_reasoning() -> None: + """A caller can explicitly opt into encrypted reasoning on a capable Foundry deployment.""" + mock_response = MagicMock() + mock_response.id = "response_123" + mock_response.model = "test-model" + mock_response.created_at = 1000000000 + mock_response.output_parsed = None + mock_response.metadata = {} + mock_response.output = [] + mock_response.usage = None + mock_response.finish_reason = None + mock_response.conversation = None + mock_response.status = "completed" + mock_response.parse.return_value = mock_response + mock_response.headers = {} + + mock_openai = MagicMock() + mock_openai.responses.with_raw_response.create = AsyncMock(return_value=mock_response) + mock_project = MagicMock() + mock_project.get_openai_client.return_value = mock_openai + client = RawFoundryAgentChatClient(project_client=mock_project, agent_name="test-agent") + + await client.get_response( + [Message(role="user", contents=["Hello"])], + options={"include": ["reasoning.encrypted_content"]}, + ) + + await_args = mock_openai.responses.with_raw_response.create.await_args + assert await_args is not None + assert await_args.kwargs["include"] == ["reasoning.encrypted_content"] + + def test_agent_accepts_raw_foundry_agent_chat_client() -> None: mock_project = MagicMock() mock_project.get_openai_client.return_value = MagicMock() @@ -1194,6 +1267,176 @@ def _import_with_missing_azure_monitor( await agent.configure_azure_monitor() +@pytest.mark.parametrize("terminate_tool_loop", [False, True]) +async def test_foundry_agent_workflow_replays_parallel_reasoning_function_group( + terminate_tool_loop: bool, +) -> None: + """Stateless cross-agent replay keeps a parallel reasoning/function batch atomic.""" + summary_inputs: list[list[dict[str, Any]]] = [] + + def _message(message_id: str, text: str) -> dict[str, Any]: + return { + "id": message_id, + "content": [{"annotations": [], "text": text, "type": "output_text"}], + "role": "assistant", + "status": "completed", + "type": "message", + } + + def _response(response_id: str, output: list[dict[str, Any]]) -> dict[str, Any]: + return { + "id": response_id, + "created_at": 0, + "model": "gpt-5.4", + "object": "response", + "output": output, + "parallel_tool_calls": True, + "status": "completed", + "tool_choice": "auto", + "tools": [], + } + + async def foundry_responses_boundary(request: httpx.Request) -> httpx.Response: + payload = json.loads(request.content) + agent_name = payload["agent_reference"]["name"] + + if agent_name == "research-agent" and "previous_response_id" not in payload: + return httpx.Response( + 200, + json=_response( + "resp_research_tool", + [ + { + "encrypted_content": "encrypted-reasoning", + "id": "rs_required", + "summary": [{"text": "I need both local tools.", "type": "summary_text"}], + "type": "reasoning", + }, + { + "arguments": '{"query":"Agent Framework"}', + "call_id": "call_paired", + "id": "fc_paired", + "name": "lookup_docs", + "status": "completed", + "type": "function_call", + }, + { + "arguments": '{"query":"stateless replay"}', + "call_id": "call_guarded", + "id": "fc_guarded", + "name": "guarded_lookup", + "status": "completed", + "type": "function_call", + }, + ], + ), + ) + + if agent_name == "research-agent": + return httpx.Response( + 200, + json=_response( + "resp_research_final", + [_message("msg_research", "Microsoft Agent Framework")], + ), + ) + + input_items = payload["input"] + summary_inputs.append(input_items) + input_types = {item.get("type") for item in input_items if isinstance(item, dict)} + if "function_call" in input_types and "reasoning" not in input_types: + return httpx.Response( + 400, + json={ + "error": { + "message": ( + "Item 'fc_paired' of type 'function_call' was provided without its " + "required 'reasoning' item: 'rs_required'." + ), + "type": "invalid_request_error", + "code": "invalid_request_error", + } + }, + ) + + return httpx.Response( + 200, + json=_response( + "resp_summary", + [_message("msg_summary", "The research result was Microsoft Agent Framework.")], + ), + ) + + @tool(name="lookup_docs", approval_mode="never_require") + def lookup_docs(query: str) -> str: + return f"Found documentation for {query}" + + @tool(name="guarded_lookup", approval_mode="never_require") + def guarded_lookup(query: str) -> str: + return f"Found guarded documentation for {query}" + + class TerminateToolLoopMiddleware(FunctionMiddleware): + async def process( + self, + context: FunctionInvocationContext, + call_next: Callable[[], Awaitable[None]], + ) -> None: + if context.function.name == "guarded_lookup": + context.result = "Blocked by policy" + raise MiddlewareTermination("Policy blocked tool execution") + await call_next() + + transport = httpx.MockTransport(foundry_responses_boundary) + responses_client = AsyncOpenAI( + api_key="test-key", + http_client=httpx.AsyncClient(transport=transport), + max_retries=0, + ) + project_client = MagicMock() + project_client.get_openai_client.return_value = responses_client + + research_agent = FoundryAgent( + project_client=project_client, + agent_name="research-agent", + tools=[lookup_docs, guarded_lookup], + middleware=[TerminateToolLoopMiddleware()] if terminate_tool_loop else None, + ) + summary_agent = FoundryAgent( + project_client=project_client, + agent_name="summary-agent", + ) + research_executor = AgentExecutor(research_agent, id="research") + summary_executor = AgentExecutor(summary_agent, id="summary") + workflow = ( + WorkflowBuilder(start_executor=research_executor, output_from=[summary_executor]) + .add_edge(research_executor, summary_executor) + .build() + ) + + result = await workflow.run("Research Agent Framework and summarize the result") + + outputs = result.get_outputs() + assert outputs + assert outputs[-1].text == "The research result was Microsoft Agent Framework." + assert len(summary_inputs) == 1 + assert [item["type"] for item in summary_inputs[0]] == [ + "message", + "reasoning", + "function_call", + "function_call", + "function_call_output", + "function_call_output", + *(["message"] if not terminate_tool_loop else []), + ] + assert summary_inputs[0][1]["encrypted_content"] == "encrypted-reasoning" + assert [item["call_id"] for item in summary_inputs[0][2:4]] == ["call_paired", "call_guarded"] + assert [item["call_id"] for item in summary_inputs[0][4:6]] == ["call_paired", "call_guarded"] + assert summary_inputs[0][4]["output"] == "Found documentation for Agent Framework" + assert summary_inputs[0][5]["output"] == ( + "Blocked by policy" if terminate_tool_loop else "Found guarded documentation for stateless replay" + ) + + @pytest.mark.flaky @pytest.mark.integration @skip_if_foundry_agent_integration_tests_disabled diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py index 5466e7adb5..114ff4f005 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py @@ -105,10 +105,10 @@ TextContent, ) from azure.ai.agentserver.responses.streaming._builders import ( + OutputItemBuilder, OutputItemFunctionCallBuilder, OutputItemMcpCallBuilder, OutputItemMessageBuilder, - OutputItemReasoningItemBuilder, ReasoningSummaryPartBuilder, TextContentBuilder, ) @@ -860,8 +860,9 @@ def __init__(self, stream: ResponseEventStream) -> None: # Builder state — only one is active at a time self._message_item: OutputItemMessageBuilder | None = None self._text_content: TextContentBuilder | None = None - self._reasoning_item: OutputItemReasoningItemBuilder | None = None + self._reasoning_item: OutputItemBuilder | None = None self._summary_part: ReasoningSummaryPartBuilder | None = None + self._reasoning_encrypted_content: str | None = None self._fc_builder: OutputItemFunctionCallBuilder | None = None self._mcp_builder: OutputItemMcpCallBuilder | None = None self.needs_async = False @@ -881,12 +882,15 @@ def handle(self, content: Content) -> Generator[ResponseStreamEvent]: yield self._text_content.emit_delta(content.text) elif content.type == "text_reasoning": - if self._active_type != "text_reasoning": + if self._active_type != "text_reasoning" or (content.id is not None and content.id != self._active_id): yield from self._close() - yield from self._open_reasoning() - self._accumulated.append(content.text or "") - if self._summary_part is not None: - yield self._summary_part.emit_text_delta(content.text or "") + yield from self._open_reasoning(content) + if encrypted_content := _reasoning_encrypted_content(content): + self._reasoning_encrypted_content = encrypted_content + if content.text: + self._accumulated.append(content.text) + if self._summary_part is not None: + yield self._summary_part.emit_text_delta(content.text) elif content.type == "function_call" and content.call_id is not None: if self._active_type != "function_call" or self._active_id != content.call_id: @@ -943,12 +947,28 @@ def _open_message(self) -> Generator[ResponseStreamEvent]: yield self._message_item.emit_added() yield self._text_content.emit_added() - def _open_reasoning(self) -> Generator[ResponseStreamEvent]: - self._reasoning_item = self._stream.add_output_item_reasoning_item() - self._summary_part = self._reasoning_item.add_summary_part() + def _open_reasoning(self, content: Content) -> Generator[ResponseStreamEvent]: + item_id = content.id + if not item_id or not IdGenerator.is_valid(item_id)[0]: + item_id = IdGenerator.new_id("rs") + self._reasoning_item = self._stream.add_output_item(item_id) + self._summary_part = ReasoningSummaryPartBuilder( + self._stream, + self._reasoning_item.output_index, + 0, + item_id, + ) + self._reasoning_encrypted_content = _reasoning_encrypted_content(content) self._active_type = "text_reasoning" - self._active_id = None - yield self._reasoning_item.emit_added() + self._active_id = item_id + yield self._reasoning_item.emit_added( + _reasoning_output_item( + item_id=item_id, + summary_texts=[], + encrypted_content=None, + status="in_progress", + ) + ) yield self._summary_part.emit_added() def _open_function_call(self, content: Content) -> Generator[ResponseStreamEvent]: @@ -983,9 +1003,17 @@ def _close(self) -> Generator[ResponseStreamEvent]: elif self._active_type == "text_reasoning" and self._summary_part and self._reasoning_item: yield self._summary_part.emit_text_done(accumulated) yield self._summary_part.emit_done() - yield self._reasoning_item.emit_done() + yield self._reasoning_item.emit_done( + _reasoning_output_item( + item_id=self._reasoning_item.item_id, + summary_texts=[accumulated], + encrypted_content=self._reasoning_encrypted_content, + status="completed", + ) + ) self._summary_part = None self._reasoning_item = None + self._reasoning_encrypted_content = None elif self._active_type == "function_call" and self._fc_builder: yield self._fc_builder.emit_arguments_done(accumulated) @@ -1064,6 +1092,21 @@ async def _items_to_messages( return messages +def _reasoning_item_to_contents(reasoning: ItemReasoningItem | OutputItemReasoningItem) -> list[Content]: + """Convert a hosted reasoning item without losing its stateless replay metadata.""" + encrypted_content = getattr(reasoning, "encrypted_content", None) + if reasoning.summary: + return [ + Content.from_text_reasoning( + id=reasoning.id, + text=summary.text, + protected_data=encrypted_content if index == 0 else None, + ) + for index, summary in enumerate(reasoning.summary) + ] + return [Content.from_text_reasoning(id=reasoning.id, protected_data=encrypted_content)] + + async def _item_to_message(item: Item, *, approval_storage: ApprovalStorage | None = None) -> Message: """Converts an Item to a Message. @@ -1113,13 +1156,7 @@ async def _item_to_message(item: Item, *, approval_storage: ApprovalStorage | No if item.type == "reasoning": reasoning = cast(ItemReasoningItem, item) - reason_contents: list[Content] = [] - if reasoning.summary: - for summary in reasoning.summary: - reason_contents.append(Content.from_text_reasoning(id=reasoning.id, text=summary.text)) - else: - reason_contents.append(Content.from_text_reasoning(id=reasoning.id)) - return Message(role="assistant", contents=reason_contents) + return Message(role="assistant", contents=_reasoning_item_to_contents(reasoning)) if item.type == "mcp_call": mcp = cast(ItemMcpToolCall, item) @@ -1404,13 +1441,7 @@ async def _output_item_to_message(item: OutputItem, *, approval_storage: Approva if item.type == "reasoning": reasoning = cast(OutputItemReasoningItem, item) - contents: list[Content] = [] - if reasoning.summary: - for summary in reasoning.summary: - contents.append(Content.from_text_reasoning(id=reasoning.id, text=summary.text)) - else: - contents.append(Content.from_text_reasoning(id=reasoning.id)) - return Message(role="assistant", contents=contents) + return Message(role="assistant", contents=_reasoning_item_to_contents(reasoning)) if item.type == "mcp_call": mcp = cast(OutputItemMcpToolCall, item) @@ -1768,6 +1799,68 @@ def _arguments_to_str(arguments: Any | None) -> str: return json.dumps(arguments, default=_argument_json_default) +def _reasoning_encrypted_content(content: Content) -> str | None: + """Return the opaque reasoning payload used for stateless replay.""" + encrypted_content = content.protected_data or content.additional_properties.get("encrypted_content") + return encrypted_content if isinstance(encrypted_content, str) else None + + +def _reasoning_output_item( + *, + item_id: str, + summary_texts: Sequence[str], + encrypted_content: str | None, + status: Literal["in_progress", "completed"], +) -> OutputItemReasoningItem: + """Build a hosted reasoning item while retaining provider replay metadata.""" + return OutputItemReasoningItem({ + "type": "reasoning", + "id": item_id, + "summary": [{"type": "summary_text", "text": text} for text in summary_texts], + "encrypted_content": encrypted_content, + "status": status, + }) + + +def _emit_reasoning_output( + stream: ResponseEventStream, + contents: Sequence[Content], +) -> Generator[ResponseStreamEvent]: + """Emit one reasoning output item for contents sharing a provider reasoning ID.""" + first = contents[0] + item_id = first.id + if not item_id or not IdGenerator.is_valid(item_id)[0]: + item_id = IdGenerator.new_id("rs") + summary_texts = [content.text or "" for content in contents] + encrypted_content = next( + (value for content in contents if (value := _reasoning_encrypted_content(content))), + None, + ) + builder = stream.add_output_item(item_id) + yield builder.emit_added( + _reasoning_output_item( + item_id=item_id, + summary_texts=[], + encrypted_content=None, + status="in_progress", + ) + ) + for summary_index, summary_text in enumerate(summary_texts): + summary_part = ReasoningSummaryPartBuilder(stream, builder.output_index, summary_index, item_id) + yield summary_part.emit_added() + yield summary_part.emit_text_delta(summary_text) + yield summary_part.emit_text_done(summary_text) + yield summary_part.emit_done() + yield builder.emit_done( + _reasoning_output_item( + item_id=item_id, + summary_texts=summary_texts, + encrypted_content=encrypted_content, + status="completed", + ) + ) + + async def _to_outputs( stream: ResponseEventStream, content: Content, @@ -1791,7 +1884,7 @@ async def _to_outputs( async for event in stream.aoutput_item_message(content.text): yield event elif content.type == "text_reasoning": - async for event in stream.aoutput_item_reasoning_item(content.text or ""): + for event in _emit_reasoning_output(stream, [content]): yield event elif content.type == "function_call": async for event in stream.aoutput_item_function_call( @@ -1943,10 +2036,20 @@ async def _to_outputs_for_messages( call/result content are encountered, or - standard output items for all other content types. """ + pending_reasoning: list[Content] = [] pending_mcp_call: Content | None = None for message in messages: for content in message.contents: + if pending_reasoning: + reasoning_id = pending_reasoning[0].id + if content.type == "text_reasoning" and reasoning_id is not None and content.id == reasoning_id: + pending_reasoning.append(content) + continue + for event in _emit_reasoning_output(stream, pending_reasoning): + yield event + pending_reasoning.clear() + if pending_mcp_call is not None: if content.type == "mcp_server_tool_result" and content.call_id == pending_mcp_call.call_id: for event in _emit_completed_mcp_call( @@ -1963,6 +2066,10 @@ async def _to_outputs_for_messages( yield event pending_mcp_call = None + if content.type == "text_reasoning": + pending_reasoning.append(content) + continue + if content.type == "mcp_server_tool_call" and content.call_id: pending_mcp_call = content continue @@ -1970,6 +2077,10 @@ async def _to_outputs_for_messages( async for event in _to_outputs(stream, content, approval_storage=approval_storage): yield event + if pending_reasoning: + for event in _emit_reasoning_output(stream, pending_reasoning): + yield event + if pending_mcp_call is not None: async for event in _to_outputs(stream, pending_mcp_call, approval_storage=approval_storage): yield event diff --git a/python/packages/foundry_hosting/tests/test_responses.py b/python/packages/foundry_hosting/tests/test_responses.py index 498caa65c5..321ac37308 100644 --- a/python/packages/foundry_hosting/tests/test_responses.py +++ b/python/packages/foundry_hosting/tests/test_responses.py @@ -333,13 +333,19 @@ async def test_hosted_mcp_call_and_result_persist_as_single_mcp_call(self) -> No assert mcp_items[0]["output"] == "found 10 cats" async def test_reasoning_content(self) -> None: + reasoning_id = "rs_576d207b35d96b3200pkcXkMwXAij920Wcv7WhRXiMPiLdOA63" agent = _make_agent( response=AgentResponse( messages=[ Message( role="assistant", contents=[ - Content.from_text_reasoning(text="Let me think..."), + Content.from_text_reasoning( + id=reasoning_id, + text="Let me ", + protected_data="encrypted-reasoning", + ), + Content.from_text_reasoning(id=reasoning_id, text="think..."), Content.from_text("The answer is 42"), ], ), @@ -356,6 +362,11 @@ async def test_reasoning_content(self) -> None: types = [item["type"] for item in body["output"]] assert "reasoning" in types assert "message" in types + reasoning_items = [item for item in body["output"] if item["type"] == "reasoning"] + assert len(reasoning_items) == 1 + assert reasoning_items[0]["id"] == reasoning_id + assert reasoning_items[0]["encrypted_content"] == "encrypted-reasoning" + assert [part["text"] for part in reasoning_items[0]["summary"]] == ["Let me ", "think..."] async def test_empty_response(self) -> None: agent = _make_agent(response=AgentResponse(messages=[])) @@ -561,11 +572,26 @@ async def test_alternating_text_and_function_call(self) -> None: assert args_done[0]["data"]["arguments"] == '{"q": "x"}' async def test_reasoning_then_text_streaming(self) -> None: + reasoning_id = "rs_576d207b35d96b3200pkcXkMwXAij920Wcv7WhRXiMPiLdOA63" agent = _make_agent( stream_updates=[ # Reasoning deltas - AgentResponseUpdate(contents=[Content.from_text_reasoning(text="Let me ")], role="assistant"), - AgentResponseUpdate(contents=[Content.from_text_reasoning(text="think...")], role="assistant"), + AgentResponseUpdate( + contents=[Content.from_text_reasoning(id=reasoning_id, text="Let me ")], role="assistant" + ), + AgentResponseUpdate( + contents=[Content.from_text_reasoning(id=reasoning_id, text="think...")], role="assistant" + ), + AgentResponseUpdate( + contents=[ + Content.from_text_reasoning( + id=reasoning_id, + text="", + protected_data="encrypted-reasoning", + ) + ], + role="assistant", + ), # Text deltas AgentResponseUpdate(contents=[Content.from_text("The answer ")], role="assistant"), AgentResponseUpdate(contents=[Content.from_text("is 42")], role="assistant"), @@ -589,6 +615,14 @@ async def test_reasoning_then_text_streaming(self) -> None: text_done = [e for e in events if e["event"] == "response.output_text.done"] assert len(text_done) == 1 assert text_done[0]["data"]["text"] == "The answer is 42" + reasoning_done = [ + event + for event in events + if event["event"] == "response.output_item.done" and event["data"]["item"]["type"] == "reasoning" + ] + assert len(reasoning_done) == 1 + assert reasoning_done[0]["data"]["item"]["id"] == reasoning_id + assert reasoning_done[0]["data"]["item"]["encrypted_content"] == "encrypted-reasoning" async def test_empty_streaming(self) -> None: agent = _make_agent(stream_updates=[]) @@ -809,6 +843,7 @@ async def test_reasoning(self) -> None: item = OutputItemReasoningItem({ "type": "reasoning", "id": "r-1", + "encrypted_content": "encrypted-reasoning", "summary": [SummaryTextContent({"type": "summary_text", "text": "thinking hard"})], }) msg = await _output_item_to_message(item) @@ -817,6 +852,7 @@ async def test_reasoning(self) -> None: assert msg.contents[0].type == "text_reasoning" assert msg.contents[0].id == "r-1" assert msg.contents[0].text == "thinking hard" + assert msg.contents[0].protected_data == "encrypted-reasoning" async def test_reasoning_no_summary(self) -> None: from azure.ai.agentserver.responses.models import OutputItemReasoningItem @@ -1304,6 +1340,7 @@ async def test_reasoning_with_summary(self) -> None: item = ItemReasoningItem({ "type": "reasoning", "id": "r-1", + "encrypted_content": "encrypted-reasoning", "summary": [SummaryTextContent({"type": "summary_text", "text": "thinking hard"})], }) msg = await _item_to_message(item) @@ -1313,6 +1350,7 @@ async def test_reasoning_with_summary(self) -> None: assert msg.contents[0].type == "text_reasoning" assert msg.contents[0].id == "r-1" assert msg.contents[0].text == "thinking hard" + assert msg.contents[0].protected_data == "encrypted-reasoning" async def test_reasoning_no_summary(self) -> None: from azure.ai.agentserver.responses.models import ItemReasoningItem diff --git a/python/packages/foundry_hosting/tests/test_responses_int.py b/python/packages/foundry_hosting/tests/test_responses_int.py index 4e599ea7d6..953510ce91 100644 --- a/python/packages/foundry_hosting/tests/test_responses_int.py +++ b/python/packages/foundry_hosting/tests/test_responses_int.py @@ -3,8 +3,9 @@ """Integration tests for ResponsesHostServer with a real Foundry endpoint. These tests exercise the full HTTP pipeline using httpx.AsyncClient with -ASGITransport — no real server process is started. The agent talks to a real -Foundry project endpoint so every test requires valid credentials. +ASGITransport — no real server process is started. Most tests talk to a real +Foundry project endpoint. Deterministic cross-package regressions replace only +the external Responses HTTP boundary. Required environment variables: FOUNDRY_PROJECT_ENDPOINT - The Microsoft Foundry project endpoint URL. @@ -18,13 +19,15 @@ import os from pathlib import Path from typing import Annotated, Any +from unittest.mock import MagicMock import httpx import pytest -from agent_framework import Agent, tool +from agent_framework import Agent, SlidingWindowStrategy, tool from agent_framework.foundry import FoundryChatClient from azure.ai.agentserver.responses import InMemoryResponseProvider from azure.identity import AzureCliCredential +from openai import AsyncOpenAI from agent_framework_foundry_hosting import ResponsesHostServer @@ -38,7 +41,6 @@ reason="No real FOUNDRY_PROJECT_ENDPOINT or FOUNDRY_MODEL provided; skipping integration tests.", ) - # --------------------------------------------------------------------------- # Fixtures # --------------------------------------------------------------------------- @@ -491,6 +493,175 @@ async def test_multi_turn_streaming(self, server: ResponsesHostServer) -> None: assert "42" in done_events[0]["data"]["text"] +class TestReasoningHostedMcpReplay: + """Regression coverage for stateless reasoning + hosted MCP replay.""" + + async def test_second_turn_replays_mcp_call_with_encrypted_reasoning(self) -> None: + """A hosted agent replays an encrypted reasoning and MCP pair when store is disabled.""" + call_count = 0 + reasoning_id = "rs_576d207b35d96b3200pkcXkMwXAij920Wcv7WhRXiMPiLdOA63" + provider_payloads: list[dict[str, Any]] = [] + + def _message(message_id: str) -> dict[str, Any]: + return { + "id": message_id, + "content": [{"annotations": [], "text": "Microsoft Agent Framework", "type": "output_text"}], + "role": "assistant", + "status": "completed", + "type": "message", + } + + def _response(response_id: str, output: list[dict[str, Any]]) -> dict[str, Any]: + return { + "id": response_id, + "created_at": 0, + "model": "gpt-5.4", + "object": "response", + "output": output, + "parallel_tool_calls": True, + "tool_choice": "auto", + "tools": [], + "status": "completed", + } + + async def foundry_responses_boundary(request: httpx.Request) -> httpx.Response: + nonlocal call_count + call_count += 1 + payload = json.loads(request.content) + provider_payloads.append(payload) + if call_count == 1: + return httpx.Response( + 200, + json=_response( + "resp_first", + [ + { + "encrypted_content": "encrypted-reasoning", + "id": reasoning_id, + "summary": [{"text": "The MCP server has the answer.", "type": "summary_text"}], + "type": "reasoning", + }, + { + "id": "mcp_paired", + "arguments": '{"query":"Agent Framework overview"}', + "name": "microsoft_docs_search", + "server_label": "Microsoft_Learn", + "type": "mcp_call", + "output": "Microsoft Agent Framework", + "status": "completed", + }, + _message("msg_first"), + ], + ), + ) + + input_items = payload["input"] + reasoning_items = [item for item in input_items if item.get("type") == "reasoning"] + mcp_calls = [item for item in input_items if item.get("type") == "mcp_call"] + if ( + len(reasoning_items) != 1 + or reasoning_items[0].get("encrypted_content") != "encrypted-reasoning" + or len(mcp_calls) != 1 + or mcp_calls[0].get("output") != "Microsoft Agent Framework" + ): + return httpx.Response( + 400, + json={ + "error": { + "message": ( + "The stateless request did not replay the complete encrypted " + "reasoning and hosted MCP call/result group." + ), + "type": "invalid_request_error", + "code": "invalid_request_error", + } + }, + ) + + return httpx.Response( + 200, + json=_response( + "resp_second", + [_message("msg_second")], + ), + ) + + transport = httpx.MockTransport(foundry_responses_boundary) + responses_client = AsyncOpenAI( + api_key="test-key", + http_client=httpx.AsyncClient(transport=transport), + max_retries=0, + ) + project_client = MagicMock() + project_client.get_openai_client.return_value = responses_client + client = FoundryChatClient( + project_client=project_client, + model="gpt-5.4", + compaction_strategy=SlidingWindowStrategy(keep_last_groups=4), + ) + learn_mcp = client.get_mcp_tool( + name="Microsoft Learn", + url="https://learn.microsoft.com/api/mcp", + allowed_tools=["microsoft_docs_search"], + approval_mode="never_require", + ) + agent = Agent( + client=client, # ty: ignore[invalid-argument-type] + instructions=( + "Always use the Microsoft Learn MCP tool to answer documentation questions. " + "Keep the final answer to one short sentence." + ), + tools=[learn_mcp], + default_options={ # pyrefly: ignore[bad-argument-type] + "store": False, + "reasoning": {"effort": "low", "summary": "auto"}, + }, + ) + server = ResponsesHostServer(agent, store=InMemoryResponseProvider()) + + first = await _post_json( + server, + { + "input": "Use Microsoft Learn MCP to find the official Agent Framework overview and state its title.", + "stream": False, + }, + ) + + assert first.status_code == 200 + first_body = first.json() + assert first_body["status"] == "completed", first_body.get("error") + first_output_types = {item["type"] for item in first_body["output"]} + assert {"reasoning", "mcp_call"} <= first_output_types + first_reasoning = next(item for item in first_body["output"] if item["type"] == "reasoning") + assert first_reasoning["id"] == reasoning_id + assert first_reasoning["encrypted_content"] == "encrypted-reasoning" + assert "reasoning.encrypted_content" in provider_payloads[0]["include"] + + second = await _post_json( + server, + { + "input": "Which Microsoft framework did you just look up? Reply with only its name.", + "stream": False, + "previous_response_id": first_body["id"], + }, + ) + + assert second.status_code == 200 + second_body = second.json() + assert second_body["status"] == "completed", second_body.get("error") + assert call_count == 2 + + second_input = provider_payloads[1]["input"] + reasoning_items = [item for item in second_input if item.get("type") == "reasoning"] + mcp_calls = [item for item in second_input if item.get("type") == "mcp_call"] + assert len(reasoning_items) == 1 + assert reasoning_items[0]["id"] == reasoning_id + assert reasoning_items[0]["encrypted_content"] == "encrypted-reasoning" + assert len(mcp_calls) == 1 + assert mcp_calls[0]["id"] == "mcp_paired" + assert mcp_calls[0]["output"] == "Microsoft Agent Framework" + + # --------------------------------------------------------------------------- # Tests — tool calling # --------------------------------------------------------------------------- diff --git a/python/packages/openai/agent_framework_openai/_chat_client.py b/python/packages/openai/agent_framework_openai/_chat_client.py index de703509da..fbefa379d2 100644 --- a/python/packages/openai/agent_framework_openai/_chat_client.py +++ b/python/packages/openai/agent_framework_openai/_chat_client.py @@ -29,7 +29,13 @@ ) from agent_framework._clients import BaseChatClient -from agent_framework._compaction import CompactionStrategy, TokenizerProtocol +from agent_framework._compaction import ( + GROUP_ANNOTATION_KEY, + GROUP_HAS_REASONING_KEY, + GROUP_ID_KEY, + CompactionStrategy, + TokenizerProtocol, +) from agent_framework._middleware import ChatAndFunctionMiddlewareTypes, ChatMiddlewareLayer from agent_framework._settings import SecretString from agent_framework._telemetry import USER_AGENT_KEY @@ -134,34 +140,6 @@ _AF_MCP_PENDING_OUTPUT_KEY = "__af_pending_mcp_result__" -def _mcp_call_ids_paired_with_reasoning(messages: Sequence[Message]) -> set[str]: - paired_call_ids: set[str] = set() - pending_reasoning_prefix = False - - for message in messages: - has_reasoning = any(content.type == "text_reasoning" for content in message.contents) - has_mcp_call = False - for content in message.contents: - if content.type != "mcp_server_tool_call": - continue - has_mcp_call = True - if (has_reasoning or pending_reasoning_prefix) and content.call_id: - paired_call_ids.add(content.call_id) - - if has_mcp_call: - pending_reasoning_prefix = False - elif ( - message.role == "assistant" - and message.contents - and all(content.type == "text_reasoning" for content in message.contents) - ): - pending_reasoning_prefix = True - else: - pending_reasoning_prefix = False - - return paired_call_ids - - class OpenAIContinuationToken(ContinuationToken): """Continuation token for OpenAI Responses API background operations.""" @@ -1394,6 +1372,12 @@ async def _prepare_options( if isinstance(value, str) and value: request_uses_service_side_storage = True break + if not request_uses_service_side_storage: + include = list(run_options.get("include", [])) + if "reasoning.encrypted_content" not in include: + include.append("reasoning.encrypted_content") + run_options["include"] = include + request_input = self._prepare_messages_for_openai( messages, request_uses_service_side_storage=request_uses_service_side_storage, @@ -1506,18 +1490,18 @@ def _prepare_messages_for_openai( Returns: The prepared chat messages for a request. """ - drops_reasoning_without_storage = not request_uses_service_side_storage and any( - content.type == "text_reasoning" for message in chat_messages for content in message.contents - ) - drop_mcp_call_ids: set[str] = set() - if drops_reasoning_without_storage: - drop_mcp_call_ids = _mcp_call_ids_paired_with_reasoning(chat_messages) + 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) + serialized_reasoning_ids: set[str] = set() list_of_list = [ self._prepare_message_for_openai( message, request_uses_service_side_storage=request_uses_service_side_storage, - drop_mcp_call_ids=drop_mcp_call_ids, + reasoning_items=reasoning_items, + serialized_reasoning_ids=serialized_reasoning_ids, ) for message in chat_messages ] @@ -1527,12 +1511,96 @@ 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: + """Reject reasoning-bound tool groups that cannot be reconstructed.""" + group_reasoning_contents: dict[str, list[Content]] = {} + group_call_ids: dict[str, list[str]] = {} + groups_with_reasoning: list[str] = [] + seen_reasoning_groups: set[str] = set() + pending_reasoning_group: str | None = None + active_tool_group: str | None = None + + for index, message in enumerate(chat_messages): + raw_annotation = message.additional_properties.get(GROUP_ANNOTATION_KEY) + annotation = cast("Mapping[str, Any]", raw_annotation) if isinstance(raw_annotation, Mapping) else None + annotated_group_id = annotation.get(GROUP_ID_KEY) if annotation is not None else None + reasoning_contents = [content for content in message.contents if content.type == "text_reasoning"] + call_ids = [ + content.call_id + for content in message.contents + if content.type + in { + "function_call", + "function_result", + "mcp_server_tool_call", + "mcp_server_tool_result", + } + and content.call_id + ] + + if isinstance(annotated_group_id, str): + group_id = annotated_group_id + elif message.role == "assistant" and reasoning_contents and not call_ids: + group_id = pending_reasoning_group or f"message_{index}" + pending_reasoning_group = group_id + elif message.role == "assistant" and call_ids: + group_id = pending_reasoning_group or f"message_{index}" + pending_reasoning_group = None + active_tool_group = group_id + elif message.role == "tool" and active_tool_group: + group_id = active_tool_group + else: + group_id = f"message_{index}" + pending_reasoning_group = None + active_tool_group = None + + if ( + reasoning_contents or (annotation is not None and annotation.get(GROUP_HAS_REASONING_KEY) is True) + ) and group_id not in seen_reasoning_groups: + groups_with_reasoning.append(group_id) + seen_reasoning_groups.add(group_id) + group_reasoning_contents.setdefault(group_id, []).extend(reasoning_contents) + group_call_ids.setdefault(group_id, []).extend(call_ids) + + invalid_groups: list[str] = [] + for group_id in groups_with_reasoning: + call_ids = list(dict.fromkeys(group_call_ids.get(group_id, []))) + 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 "" + for content in reasoning_contents + if not content.id or content.id not in replayable_reasoning_ids + ) + ) + if not reasoning_contents: + missing_reasoning_ids.append(f"") + if missing_reasoning_ids: + invalid_groups.append( + f"reasoning item(s) {', '.join(missing_reasoning_ids)} for call(s) {', '.join(call_ids)}" + ) + + 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." + ) + def _prepare_message_for_openai( self, message: Message, *, request_uses_service_side_storage: bool = True, - drop_mcp_call_ids: set[str] | None = None, + reasoning_items: Mapping[str, dict[str, Any]] | None = None, + serialized_reasoning_ids: set[str] | None = None, ) -> list[dict[str, Any]]: """Prepare a chat message for the OpenAI Responses API format.""" all_messages: list[dict[str, Any]] = [] @@ -1551,14 +1619,18 @@ def _prepare_message_for_openai( # flag, not a message-level one: HistoryProvider-attributed messages # (replays_local_storage) still need stripping when the request also carries a continuation # marker, since the server-stored items would otherwise duplicate the inline ones. Without - # storage, standalone reasoning items are invalid per the API ("reasoning was provided - # without its required following item"), so the reasoning branch always drops. When that - # happens, `_prepare_messages_for_openai` also drops the paired hosted-MCP IDs across - # message boundaries rather than replaying bare MCP items. - drop_mcp_call_ids = drop_mcp_call_ids or set() + # storage, replayable reasoning items are reconstructed from their encrypted payload and + # emitted once per provider reasoning id by `_prepare_messages_for_openai`. + reasoning_items = reasoning_items or {} + serialized_reasoning_ids = serialized_reasoning_ids if serialized_reasoning_ids is not None else set() for content in message.contents: match content.type: case "text_reasoning": + if request_uses_service_side_storage or not content.id or content.id in serialized_reasoning_ids: + continue + if reasoning_item := reasoning_items.get(content.id): + all_messages.append(reasoning_item) + serialized_reasoning_ids.add(content.id) continue case "function_result": if request_uses_service_side_storage: @@ -1611,9 +1683,7 @@ def _prepare_message_for_openai( # the prior response's items (#3295). Drop the call here; the # orphan result is dropped by the coalesce step that follows. # - # Without storage, a reasoning + hosted-MCP pair cannot be replayed - # partially: reasoning is stripped above, and a bare mcp_call is rejected. - if request_uses_service_side_storage or content.call_id in drop_mcp_call_ids: + if request_uses_service_side_storage: continue prepared_mcp = self._prepare_content_for_openai( message.role, @@ -1636,6 +1706,52 @@ def _prepare_message_for_openai( all_messages.append(args) return all_messages + def _prepare_reasoning_items_for_openai( + self, + chat_messages: Sequence[Message], + ) -> dict[str, dict[str, Any]]: + """Reconstruct one replayable Responses reasoning item per provider id.""" + grouped_contents: dict[str, list[Content]] = {} + for message in chat_messages: + for content in message.contents: + if content.type == "text_reasoning" and content.id: + grouped_contents.setdefault(content.id, []).append(content) + + reasoning_items: dict[str, dict[str, Any]] = {} + for reasoning_id, contents in grouped_contents.items(): + encrypted_content = next( + ( + content.protected_data or content.additional_properties.get("encrypted_content") + for content in contents + if content.protected_data or content.additional_properties.get("encrypted_content") + ), + None, + ) + if not encrypted_content: + continue + + item: dict[str, Any] = { + "type": "reasoning", + "id": reasoning_id, + "summary": [], + "encrypted_content": encrypted_content, + } + reasoning_texts: list[dict[str, str]] = [] + for content in contents: + props = content.additional_properties + if status := props.get("status"): + item["status"] = status + if reasoning_text_marker := props.get("reasoning_text"): + reasoning_text = content.text if reasoning_text_marker is True else reasoning_text_marker + if isinstance(reasoning_text, str) and reasoning_text: + reasoning_texts.append({"type": "reasoning_text", "text": reasoning_text}) + elif content.text: + item["summary"].append({"type": "summary_text", "text": content.text}) + if reasoning_texts: + item["content"] = reasoning_texts + reasoning_items[reasoning_id] = item + return reasoning_items + def _prepare_content_for_openai( self, role: Role | str, @@ -1664,14 +1780,17 @@ def _prepare_content_for_openai( if content.id: ret["id"] = content.id props: dict[str, Any] | None = getattr(content, "additional_properties", None) + reasoning_text_marker: Any = None if props: if status := props.get("status"): ret["status"] = status - if reasoning_text := props.get("reasoning_text"): - ret["content"] = [{"type": "reasoning_text", "text": reasoning_text}] + if reasoning_text_marker := props.get("reasoning_text"): + reasoning_text = content.text if reasoning_text_marker is True else reasoning_text_marker + if isinstance(reasoning_text, str) and reasoning_text: + ret["content"] = [{"type": "reasoning_text", "text": reasoning_text}] if encrypted_content := props.get("encrypted_content"): ret["encrypted_content"] = encrypted_content - if content.text: + if content.text and reasoning_text_marker is not True: ret["summary"].append({"type": "summary_text", "text": content.text}) return ret case "data" | "uri": @@ -2474,15 +2593,17 @@ def _parse_response_from_openai( ) case "reasoning": # ResponseOutputReasoning added_reasoning = False + encrypted_content = getattr(item, "encrypted_content", None) if item_content := getattr(item, "content", None): for index, reasoning_content in enumerate(item_content): - additional_properties: dict[str, Any] = {} + additional_properties: dict[str, Any] = {"reasoning_text": True} if hasattr(item, "summary") and item.summary and index < len(item.summary): additional_properties["summary"] = item.summary[index] contents.append( Content.from_text_reasoning( id=item.id, text=reasoning_content.text, + protected_data=encrypted_content if not added_reasoning else None, raw_representation=reasoning_content, additional_properties=additional_properties or None, ) @@ -2494,6 +2615,7 @@ def _parse_response_from_openai( Content.from_text_reasoning( id=item.id, text=summary.text, + protected_data=encrypted_content if not added_reasoning else None, raw_representation=summary, ) ) @@ -2502,12 +2624,11 @@ def _parse_response_from_openai( # Reasoning item with no visible text (e.g. encrypted reasoning). # Always emit an empty marker so co-occurrence detection can be done additional_properties_empty: dict[str, Any] = {} - if encrypted := getattr(item, "encrypted_content", None): - additional_properties_empty["encrypted_content"] = encrypted contents.append( Content.from_text_reasoning( id=item.id, text="", + protected_data=encrypted_content, raw_representation=item, additional_properties=additional_properties_empty or None, ) @@ -2755,6 +2876,7 @@ def _parse_chunk_from_openai( id=event.item_id, text=event.delta, raw_representation=event, + additional_properties={"reasoning_text": True}, ) ) metadata.update(self._get_metadata_from_response(event)) @@ -2768,6 +2890,7 @@ def _parse_chunk_from_openai( id=event.item_id, text=event.text, raw_representation=event, + additional_properties={"reasoning_text": True}, ) ) metadata.update(self._get_metadata_from_response(event)) @@ -2964,7 +3087,7 @@ def _parse_chunk_from_openai( added_reasoning = False if hasattr(event_item, "content") and event_item.content: for index, reasoning_content in enumerate(event_item.content): - additional_properties: dict[str, Any] = {} + additional_properties: dict[str, Any] = {"reasoning_text": True} if ( hasattr(event_item, "summary") and event_item.summary @@ -2975,6 +3098,7 @@ def _parse_chunk_from_openai( Content.from_text_reasoning( id=reasoning_id or None, text=reasoning_content.text, + protected_data=getattr(event_item, "encrypted_content", None), raw_representation=reasoning_content, additional_properties=additional_properties or None, ) @@ -2983,15 +3107,12 @@ def _parse_chunk_from_openai( if not added_reasoning: # Reasoning item with no visible text (e.g. encrypted reasoning). # Always emit an empty marker so co-occurrence detection can occur. - additional_properties_empty: dict[str, Any] = {} - if encrypted := getattr(event_item, "encrypted_content", None): - additional_properties_empty["encrypted_content"] = encrypted contents.append( Content.from_text_reasoning( id=reasoning_id or None, text="", + protected_data=getattr(event_item, "encrypted_content", None), raw_representation=event_item, - additional_properties=additional_properties_empty or None, ) ) case "web_search_call" | "file_search_call": @@ -3152,7 +3273,18 @@ def _get_ann_value(key: str) -> Any: logger.debug("Unparsed annotation type in streaming: %s", ann_type) case "response.output_item.done": done_item = event.item - if getattr(done_item, "type", None) == "mcp_call": + if getattr(done_item, "type", None) == "reasoning": + encrypted_content = getattr(done_item, "encrypted_content", None) + if encrypted_content: + contents.append( + Content.from_text_reasoning( + id=getattr(done_item, "id", None), + text="", + protected_data=encrypted_content, + raw_representation=done_item, + ) + ) + elif getattr(done_item, "type", None) == "mcp_call": call_id = getattr(done_item, "id", None) or getattr(done_item, "call_id", None) or "" output_text = getattr(done_item, "output", None) parsed_output: list[Content] | 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 0497d3ddf0..10e23bab87 100644 --- a/python/packages/openai/tests/openai/test_openai_chat_client.py +++ b/python/packages/openai/tests/openai/test_openai_chat_client.py @@ -34,6 +34,7 @@ InMemoryHistoryProvider, SessionContext, ) +from agent_framework._workflows._checkpoint_encoding import decode_checkpoint_value, encode_checkpoint_value from agent_framework.exceptions import ( ChatClientException, ChatClientInvalidRequestException, @@ -57,8 +58,8 @@ from pydantic import BaseModel from pytest import param -from agent_framework_openai import OpenAIChatClient -from agent_framework_openai._chat_client import OPENAI_LOCAL_SHELL_CALL_ITEM_ID_KEY, RawOpenAIChatClient +from agent_framework_openai import OpenAIChatClient, OpenAIChatOptions, RawOpenAIChatClient +from agent_framework_openai._chat_client import OPENAI_LOCAL_SHELL_CALL_ITEM_ID_KEY from agent_framework_openai._exceptions import OpenAIContentFilterException skip_if_openai_integration_tests_disabled = pytest.mark.skipif( @@ -976,6 +977,332 @@ async def __aexit__( assert update.model == "test-model" +async def test_streamed_reasoning_function_group_survives_serialization_for_stateless_replay() -> None: + client = OpenAIChatClient(model="test-model", api_key="test-key") + + reasoning_done = MagicMock() + reasoning_done.type = "response.output_item.done" + reasoning_done.item = MagicMock() + reasoning_done.item.type = "reasoning" + reasoning_done.item.id = "rs_streamed" + reasoning_done.item.encrypted_content = "encrypted-streamed-reasoning" + + function_added = MagicMock() + function_added.type = "response.output_item.added" + function_added.output_index = 1 + function_added.item = MagicMock() + function_added.item.type = "function_call" + function_added.item.call_id = "call_streamed" + function_added.item.name = "get_weather" + + function_arguments = MagicMock() + function_arguments.type = "response.function_call_arguments.delta" + function_arguments.output_index = 1 + function_arguments.item_id = "fc_streamed" + function_arguments.delta = '{"location":"Seattle"}' + + events = [ + ResponseReasoningSummaryTextDeltaEvent( + type="response.reasoning_summary_text.delta", + item_id="rs_streamed", + output_index=0, + sequence_number=1, + summary_index=0, + delta="I should check the weather.", + ), + reasoning_done, + function_added, + function_arguments, + ] + fake_stream = _FakeAsyncEventStream(events) + + with ( + patch.object(client, "_prepare_request", new=AsyncMock(return_value=(client.client, {}, {}))), + patch.object(client.client.responses, "create", new=AsyncMock(return_value=fake_stream)), + patch.object(client, "_get_metadata_from_response", return_value={}), + ): + stream = client.get_response( + messages=[Message(role="user", contents=["What's the weather?"])], + options={}, + stream=True, + ) + response = await stream.get_final_response() + + session_restored_message = Message.from_json(response.messages[0].to_json()) + workflow_payload = json.loads(json.dumps(encode_checkpoint_value(session_restored_message))) + restored_message = decode_checkpoint_value(workflow_payload) + assert isinstance(restored_message, Message) + reasoning_contents = [content for content in restored_message.contents if content.type == "text_reasoning"] + assert [(content.id, content.text) for content in reasoning_contents] == [ + ("rs_streamed", "I should check the weather.") + ] + assert reasoning_contents[0].protected_data == "encrypted-streamed-reasoning" + + messages = [ + Message(role="user", contents=["What's the weather?"]), + restored_message, + Message( + role="tool", + contents=[Content.from_function_result(call_id="call_streamed", result="Sunny")], + ), + ] + _, run_options, _ = await client._prepare_request(messages, {"store": False}) + + assert run_options["input"] == [ + { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "What's the weather?"}], + }, + { + "type": "reasoning", + "id": "rs_streamed", + "summary": [{"type": "summary_text", "text": "I should check the weather."}], + "encrypted_content": "encrypted-streamed-reasoning", + }, + { + "call_id": "call_streamed", + "id": "fc_streamed", + "type": "function_call", + "name": "get_weather", + "arguments": '{"location":"Seattle"}', + }, + { + "call_id": "call_streamed", + "type": "function_call_output", + "output": "Sunny", + }, + ] + + +def test_streamed_reasoning_text_replays_as_reasoning_content() -> None: + client = OpenAIChatClient(model="test-model", api_key="test-key") + function_call_ids: dict[int, tuple[str, str]] = {} + seen_reasoning_delta_item_ids: set[str] = set() + reasoning_delta = ResponseReasoningTextDeltaEvent( + type="response.reasoning_text.delta", + content_index=0, + item_id="rs_reasoning_text", + output_index=0, + sequence_number=1, + delta="Private reasoning text", + ) + reasoning_done = MagicMock() + reasoning_done.type = "response.output_item.done" + reasoning_done.item = MagicMock() + reasoning_done.item.type = "reasoning" + reasoning_done.item.id = "rs_reasoning_text" + reasoning_done.item.encrypted_content = "encrypted-reasoning-text" + + response = ChatResponse.from_updates([ + client._parse_chunk_from_openai( + event, + {}, + function_call_ids, + seen_reasoning_delta_item_ids, + ) + for event in (reasoning_delta, reasoning_done) + ]) + prepared = client._prepare_messages_for_openai( + response.messages, + request_uses_service_side_storage=False, + ) + + assert prepared == [ + { + "type": "reasoning", + "id": "rs_reasoning_text", + "summary": [], + "content": [{"type": "reasoning_text", "text": "Private reasoning text"}], + "encrypted_content": "encrypted-reasoning-text", + } + ] + + +def _streamed_reasoning_text_events(event_kind: str) -> tuple[list[object], str, str, str]: + cases = { + "delta": ("rs_streamed_reasoning", "Streamed private reasoning", "encrypted-streamed-reasoning"), + "fallback": ("rs_fallback_reasoning", "Fallback private reasoning", "encrypted-fallback-reasoning"), + "snapshot": ("rs_snapshot_reasoning", "Snapshot private reasoning", "encrypted-snapshot-reasoning"), + } + item_id, text, encrypted_content = cases[event_kind] + + if event_kind == "snapshot": + reasoning_content = MagicMock() + reasoning_content.text = text + reasoning_item = MagicMock() + reasoning_item.type = "reasoning" + reasoning_item.id = item_id + reasoning_item.content = [reasoning_content] + reasoning_item.summary = [] + reasoning_item.encrypted_content = encrypted_content + reasoning_added = MagicMock() + reasoning_added.type = "response.output_item.added" + reasoning_added.output_index = 0 + reasoning_added.item = reasoning_item + return [reasoning_added], item_id, text, encrypted_content + + if event_kind == "delta": + text_event: object = ResponseReasoningTextDeltaEvent( + type="response.reasoning_text.delta", + content_index=0, + item_id=item_id, + output_index=0, + sequence_number=1, + delta=text, + ) + else: + text_event = ResponseReasoningTextDoneEvent( + type="response.reasoning_text.done", + content_index=0, + item_id=item_id, + output_index=0, + sequence_number=1, + text=text, + ) + reasoning_done = MagicMock() + reasoning_done.type = "response.output_item.done" + reasoning_done.item = MagicMock() + reasoning_done.item.type = "reasoning" + reasoning_done.item.id = item_id + reasoning_done.item.encrypted_content = encrypted_content + return [text_event, reasoning_done], item_id, text, encrypted_content + + +@pytest.mark.parametrize("event_kind", ["delta", "fallback", "snapshot"]) +async def test_streamed_reasoning_text_is_stored_once_and_replayed(event_kind: str) -> None: + """The public streaming client stores reasoning text once and replays it as reasoning content.""" + client = OpenAIChatClient(model="test-model", api_key="test-key") + events, item_id, text, encrypted_content = _streamed_reasoning_text_events(event_kind) + streamed_response = _FakeAsyncEventStream(events) + + follow_up_response = MagicMock() + follow_up_response.output_parsed = None + follow_up_response.metadata = {} + follow_up_response.usage = None + follow_up_response.id = "resp-follow-up" + follow_up_response.model = "test-model" + follow_up_response.created_at = 1000000001 + follow_up_response.status = "completed" + follow_up_response.finish_reason = "stop" + follow_up_response.incomplete = None + follow_up_response.conversation = None + follow_up_response.output = [] + + with patch.object( + client.client.responses, + "create", + new=AsyncMock(side_effect=[streamed_response, _as_raw(follow_up_response)]), + ) as mock_create: + stream = client.get_response( + [Message(role="user", contents=["Think about this"])], + options={"store": False}, + stream=True, + ) + response = await stream.get_final_response() + + reasoning = response.messages[0].contents[0] + assert reasoning.text == text + assert reasoning.additional_properties == {"reasoning_text": True} + + await client.get_response( + [ + Message(role="user", contents=["Think about this"]), + *response.messages, + Message(role="user", contents=["Continue"]), + ], + options={"store": False}, + ) + + replayed_reasoning = [ + item for item in mock_create.call_args_list[1].kwargs["input"] if item.get("type") == "reasoning" + ] + assert replayed_reasoning == [ + { + "type": "reasoning", + "id": item_id, + "summary": [], + "encrypted_content": encrypted_content, + "content": [{"type": "reasoning_text", "text": text}], + } + ] + + +def test_streamed_encrypted_reasoning_without_visible_text_is_replayable() -> None: + client = OpenAIChatClient(model="test-model", api_key="test-key") + reasoning_done = MagicMock() + reasoning_done.type = "response.output_item.done" + reasoning_done.item = MagicMock() + reasoning_done.item.type = "reasoning" + reasoning_done.item.id = "rs_encrypted_only" + reasoning_done.item.encrypted_content = "encrypted-only" + + update = client._parse_chunk_from_openai(reasoning_done, {}, {}) + response = ChatResponse.from_updates([update]) + + assert [(content.id, content.text, content.protected_data) for content in response.messages[0].contents] == [ + ("rs_encrypted_only", "", "encrypted-only") + ] + + +def test_streamed_summary_and_reasoning_text_replay_as_one_provider_item() -> None: + client = OpenAIChatClient(model="test-model", api_key="test-key") + function_call_ids: dict[int, tuple[str, str]] = {} + seen_reasoning_delta_item_ids: set[str] = set() + summary_delta = ResponseReasoningSummaryTextDeltaEvent( + type="response.reasoning_summary_text.delta", + item_id="rs_mixed", + output_index=0, + sequence_number=1, + summary_index=0, + delta="Visible summary", + ) + reasoning_delta = ResponseReasoningTextDeltaEvent( + type="response.reasoning_text.delta", + content_index=0, + item_id="rs_mixed", + output_index=0, + sequence_number=2, + delta="Private reasoning", + ) + reasoning_done = MagicMock() + reasoning_done.type = "response.output_item.done" + reasoning_done.item = MagicMock() + reasoning_done.item.type = "reasoning" + reasoning_done.item.id = "rs_mixed" + reasoning_done.item.encrypted_content = "encrypted-mixed" + + response = ChatResponse.from_updates([ + client._parse_chunk_from_openai( + event, + {}, + function_call_ids, + seen_reasoning_delta_item_ids, + ) + for event in (summary_delta, reasoning_delta, reasoning_done) + ]) + reasoning_contents = response.messages[0].contents + assert [(content.text, content.additional_properties) for content in reasoning_contents] == [ + ("Visible summary", {}), + ("Private reasoning", {"reasoning_text": True}), + ] + assert reasoning_contents[1].protected_data == "encrypted-mixed" + + prepared = client._prepare_messages_for_openai( + response.messages, + request_uses_service_side_storage=False, + ) + assert prepared == [ + { + "type": "reasoning", + "id": "rs_mixed", + "summary": [{"type": "summary_text", "text": "Visible summary"}], + "content": [{"type": "reasoning_text", "text": "Private reasoning"}], + "encrypted_content": "encrypted-mixed", + } + ] + + async def test_streaming_text_format_preserves_final_structured_output() -> None: """Streaming structured output should still parse into the final ChatResponse value.""" client = OpenAIChatClient(model="test-model", api_key="test-key") @@ -1161,6 +1488,344 @@ def test_response_content_creation_with_reasoning() -> None: assert response.messages[0].contents[0].text == "Reasoning step" +async def test_non_streaming_reasoning_function_group_round_trips_for_stateless_replay() -> None: + client = OpenAIChatClient(model="test-model", api_key="test-key") + + mock_response = MagicMock() + mock_response.output_parsed = None + mock_response.metadata = {} + mock_response.usage = None + mock_response.id = "resp-1" + mock_response.model = "test-model" + mock_response.created_at = 1000000000 + mock_response.status = "completed" + mock_response.finish_reason = "tool_calls" + mock_response.incomplete = None + mock_response.conversation = None + + mock_reasoning_content = MagicMock() + mock_reasoning_content.text = "Reasoning step" + mock_reasoning_item = MagicMock() + mock_reasoning_item.type = "reasoning" + mock_reasoning_item.id = "rs_123" + mock_reasoning_item.content = [mock_reasoning_content] + mock_reasoning_item.summary = [Summary(text="Visible summary", type="summary_text")] + mock_reasoning_item.encrypted_content = "encrypted-reasoning" + + mock_function_call_item = MagicMock() + mock_function_call_item.type = "function_call" + mock_function_call_item.id = "fc_123" + mock_function_call_item.call_id = "call_123" + mock_function_call_item.name = "get_weather" + mock_function_call_item.arguments = '{"location":"Amsterdam"}' + mock_function_call_item.status = "completed" + mock_response.output = [mock_reasoning_item, mock_function_call_item] + + response = client._parse_response_from_openai(mock_response, options={}) # type: ignore[arg-type] + + assert [(content.type, content.text) for content in response.messages[0].contents[:2]] == [ + ("text_reasoning", "Reasoning step"), + ("text_reasoning", "Visible summary"), + ] + assert response.messages[0].contents[0].protected_data == "encrypted-reasoning" + + messages = [ + Message(role="user", contents=["What's the weather?"]), + *response.messages, + Message( + role="tool", + contents=[Content.from_function_result(call_id="call_123", result="Sunny")], + ), + ] + _, run_options, _ = await client._prepare_request( + messages, + {"store": False, "include": ["message.output_text.logprobs"]}, + ) + + assert run_options["include"] == ["message.output_text.logprobs", "reasoning.encrypted_content"] + assert run_options["input"] == [ + { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "What's the weather?"}], + }, + { + "type": "reasoning", + "id": "rs_123", + "summary": [{"type": "summary_text", "text": "Visible summary"}], + "content": [{"type": "reasoning_text", "text": "Reasoning step"}], + "encrypted_content": "encrypted-reasoning", + }, + { + "call_id": "call_123", + "id": "fc_123", + "type": "function_call", + "name": "get_weather", + "arguments": '{"location":"Amsterdam"}', + "status": "completed", + }, + { + "call_id": "call_123", + "type": "function_call_output", + "output": "Sunny", + }, + ] + + +async def test_non_streaming_reasoning_text_is_stored_once_and_replayed() -> None: + """The public client stores reasoning text in Content.text and replays it as reasoning content.""" + client = OpenAIChatClient(model="test-model", api_key="test-key") + + first_response = MagicMock() + first_response.output_parsed = None + first_response.metadata = {} + first_response.usage = None + first_response.id = "resp-reasoning" + first_response.model = "test-model" + first_response.created_at = 1000000000 + first_response.status = "completed" + first_response.finish_reason = "stop" + first_response.incomplete = None + first_response.conversation = None + + reasoning_content = MagicMock() + reasoning_content.text = "Private reasoning text" + reasoning_item = MagicMock() + reasoning_item.type = "reasoning" + reasoning_item.id = "rs_reasoning" + reasoning_item.content = [reasoning_content] + reasoning_item.summary = [] + reasoning_item.encrypted_content = "encrypted-reasoning" + first_response.output = [reasoning_item] + + second_response = MagicMock() + second_response.output_parsed = None + second_response.metadata = {} + second_response.usage = None + second_response.id = "resp-follow-up" + second_response.model = "test-model" + second_response.created_at = 1000000001 + second_response.status = "completed" + second_response.finish_reason = "stop" + second_response.incomplete = None + second_response.conversation = None + second_response.output = [] + + with patch.object( + client.client.responses, + "create", + side_effect=[_as_raw(first_response), _as_raw(second_response)], + ) as mock_create: + response = await client.get_response( + [Message(role="user", contents=["Think about this"])], + options={"store": False}, + ) + + reasoning = response.messages[0].contents[0] + assert reasoning.text == "Private reasoning text" + assert reasoning.additional_properties == {"reasoning_text": True} + + await client.get_response( + [ + Message(role="user", contents=["Think about this"]), + *response.messages, + Message(role="user", contents=["Continue"]), + ], + options={"store": False}, + ) + + replayed_reasoning = [ + item for item in mock_create.call_args_list[1].kwargs["input"] if item.get("type") == "reasoning" + ] + assert replayed_reasoning == [ + { + "type": "reasoning", + "id": "rs_reasoning", + "summary": [], + "encrypted_content": "encrypted-reasoning", + "content": [{"type": "reasoning_text", "text": "Private reasoning text"}], + } + ] + + +async def test_prepare_request_does_not_duplicate_encrypted_reasoning_include() -> None: + client = OpenAIChatClient(model="test-model", api_key="test-key") + + _, run_options, _ = await client._prepare_request( + [Message(role="user", contents=["Hello"])], + {"include": ["reasoning.encrypted_content", "message.output_text.logprobs"]}, + ) + + assert run_options["include"] == ["reasoning.encrypted_content", "message.output_text.logprobs"] + + +async def test_stateless_reasoning_group_without_encrypted_content_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_missing", text="I need both tools."), + Content.from_function_call(call_id="call_one", name="first_tool", arguments="{}"), + Content.from_function_call(call_id="call_two", name="second_tool", arguments="{}"), + ], + ) + ] + + 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_missing" in message + assert "call_one" in message + assert "call_two" in message + assert "service-side continuation" in message + assert "atomic compaction" 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 + return True + + client = OpenAIChatClient( + model="test-model", + api_key="test-key", + compaction_strategy=exclude_reasoning_message, + ) + create = AsyncMock() + messages = [ + Message( + role="assistant", + contents=[ + Content.from_text_reasoning( + id="rs_compacted", + text="I need a tool.", + protected_data="encrypted-reasoning", + ) + ], + ), + Message( + role="assistant", + contents=[Content.from_function_call(call_id="call_compacted", name="tool", arguments="{}")], + ), + ] + + 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 "group_msg_0" in message + assert "call_compacted" in message + assert "atomic compaction" in message + create.assert_not_awaited() + + +async def test_split_reasoning_group_without_encrypted_content_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_split", text="I need a tool.")], + ), + Message( + role="assistant", + contents=[Content.from_function_call(call_id="call_split", name="tool", arguments="{}")], + ), + ] + + with ( + patch.object(client.client.responses.with_raw_response, "create", new=create), + pytest.raises(ChatClientInvalidRequestException, match="rs_split.*call_split"), + ): + await client.get_response(messages, options={"store": False}) + + create.assert_not_awaited() + + +async def test_fully_compacted_reasoning_group_continues_with_remaining_messages() -> None: + async def exclude_reasoning_group(messages: list[Message]) -> bool: + for message in messages: + if message.role in {"assistant", "tool"}: + message.additional_properties["_excluded"] = True + return True + + client = OpenAIChatClient( + model="test-model", + api_key="test-key", + compaction_strategy=exclude_reasoning_group, + ) + mock_response = MagicMock( + id="response_123", + model="test-model", + created_at=1000000000, + metadata={}, + output_parsed=None, + output=[], + usage=None, + finish_reason=None, + conversation=None, + status="completed", + incomplete_details=None, + ) + messages = [ + Message(role="user", contents=["Keep this request."]), + Message( + role="assistant", + contents=[ + Content.from_text_reasoning(id="rs_excluded", text="Legacy reasoning."), + Content.from_function_call(call_id="call_excluded", name="tool", arguments="{}"), + ], + ), + Message( + role="tool", + contents=[Content.from_function_result(call_id="call_excluded", result="Excluded result.")], + ), + ] + + with patch.object(client.client.responses, "create", return_value=_as_raw(mock_response)) as create: + await client.get_response(messages, options={"store": False}) + + create.assert_awaited_once() + assert create.await_args is not None + assert create.await_args.kwargs["input"] == [ + { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "Keep this request."}], + } + ] + + +async def test_encrypted_reasoning_capability_rejection_is_not_retried_lossily() -> None: + client = OpenAIChatClient(model="test-model", api_key="test-key") + service_error = BadRequestError( + message="Encrypted reasoning is not supported", + response=MagicMock(), + body={"error": {"code": "invalid_request", "message": "Encrypted reasoning is not supported"}}, + ) + service_error.code = "invalid_request" + + with ( + patch.object(client.client.responses, "create", side_effect=service_error) as create, + pytest.raises(ChatClientException, match="Encrypted reasoning is not supported"), + ): + await client.get_response([Message(role="user", contents=["Hello"])], options={"store": False}) + + create.assert_awaited_once() + assert create.await_args is not None + assert create.await_args.kwargs["include"] == ["reasoning.encrypted_content"] + + def test_response_content_keeps_reasoning_and_function_calls_in_one_message() -> None: """Reasoning + function calls should parse into one assistant message.""" client = OpenAIChatClient(model="test-model", api_key="test-key") @@ -1540,8 +2205,8 @@ def local_exec(command: str) -> str: assert len(local_shell_outputs) == 0 -async def test_tool_loop_store_false_omits_reasoning_items_from_second_request() -> None: - """Stateless tool-loop replay must omit response-scoped reasoning items.""" +async def test_tool_loop_store_false_replays_encrypted_reasoning_group() -> None: + """The public client replays an encrypted reasoning/call/result group.""" client = OpenAIChatClient(model="test-model", api_key="test-key") mock_response1 = MagicMock() @@ -1561,7 +2226,7 @@ async def test_tool_loop_store_false_omits_reasoning_items_from_second_request() mock_reasoning_item.id = "rs_local_only" mock_reasoning_item.content = [] mock_reasoning_item.summary = [] - mock_reasoning_item.encrypted_content = None + mock_reasoning_item.encrypted_content = "encrypted-reasoning" mock_function_call_item = MagicMock() mock_function_call_item.type = "function_call" @@ -1607,9 +2272,19 @@ async def test_tool_loop_store_false_omits_reasoning_items_from_second_request() assert response.text == "The weather in Amsterdam is sunny." assert mock_create.call_count == 2 + assert mock_create.call_args_list[0].kwargs["include"] == ["reasoning.encrypted_content"] + assert mock_create.call_args_list[1].kwargs["include"] == ["reasoning.encrypted_content"] second_call_input = mock_create.call_args_list[1].kwargs["input"] - assert not any(item.get("type") == "reasoning" for item in second_call_input) + reasoning_items = [item for item in second_call_input if item.get("type") == "reasoning"] + assert reasoning_items == [ + { + "type": "reasoning", + "id": "rs_local_only", + "summary": [], + "encrypted_content": "encrypted-reasoning", + } + ] function_calls = [item for item in second_call_input if item.get("type") == "function_call"] assert len(function_calls) == 1 @@ -1620,6 +2295,43 @@ async def test_tool_loop_store_false_omits_reasoning_items_from_second_request() assert function_outputs[0]["call_id"] == "call_123" +async def test_stateless_request_rejects_non_replayable_reasoning_bound_mcp_output() -> None: + client = OpenAIChatClient(model="test-model", api_key="test-key") + messages = [ + Message(role="user", contents=["Search the API specifications."]), + Message( + role="assistant", + contents=[ + Content.from_text_reasoning(id="rs_mcp", text="Use the hosted MCP server."), + Content.from_mcp_server_tool_call( + call_id="mcp_search", + tool_name="search", + server_name="api_specs", + arguments='{"q":"cats"}', + ), + ], + ), + Message( + role="tool", + contents=[ + Content.from_mcp_server_tool_result( + call_id="mcp_search", + output=[Content.from_text(text="found 10 cats")], + ) + ], + ), + ] + + create = AsyncMock() + with ( + patch.object(client.client.responses.with_raw_response, "create", new=create), + pytest.raises(ChatClientInvalidRequestException, match="rs_mcp.*mcp_search"), + ): + await client.get_response(messages=messages, options={"store": False}) + + create.assert_not_awaited() + + def test_response_content_creation_with_shell_call() -> None: """Test _parse_response_from_openai with shell_call output.""" client = OpenAIChatClient(model="test-model", api_key="test-key") @@ -2575,18 +3287,14 @@ def test_prepare_message_for_openai_with_function_approval_response() -> None: assert prepared_message["approve"] is True -def test_prepare_message_for_openai_includes_reasoning_with_function_call() -> None: - """Test _prepare_message_for_openai includes reasoning items alongside function_calls. - - Reasoning models require reasoning items to be present in the input when - function_call items are included. Stripping reasoning causes a 400 error: - "function_call was provided without its required reasoning item". - """ +def test_prepare_messages_for_openai_keeps_active_function_call_for_tool_loop() -> None: + """An active tool loop retains its current function call until the model produces a follow-up.""" client = OpenAIChatClient(model="test-model", api_key="test-key") reasoning = Content.from_text_reasoning( id="rs_abc123", text="Let me analyze the request", + protected_data="encrypted-reasoning", additional_properties={"status": "completed"}, ) function_call = Content.from_function_call( @@ -2597,20 +3305,94 @@ def test_prepare_message_for_openai_includes_reasoning_with_function_call() -> N message = Message(role="assistant", contents=[reasoning, function_call]) - # Storage-on path strips both server-issued reasoning (rs_*) and function_call items - # because the server already has them via previous_response_id (#3295). - storage_on_result = client._prepare_message_for_openai(message, request_uses_service_side_storage=True) + storage_on_result = client._prepare_messages_for_openai([message], request_uses_service_side_storage=True) storage_on_types = [item["type"] for item in storage_on_result] assert "reasoning" not in storage_on_types assert "function_call" not in storage_on_types - # Storage-off path keeps function_call inline so the server sees the call. Reasoning items - # cannot be replayed inline against a server that has no record of the prior response, so - # they remain dropped on this path as well. - storage_off_result = client._prepare_message_for_openai(message, request_uses_service_side_storage=False) + storage_off_result = client._prepare_messages_for_openai([message], request_uses_service_side_storage=False) storage_off_types = [item["type"] for item in storage_off_result] + assert "reasoning" in storage_off_types assert "function_call" in storage_off_types - assert "reasoning" not in storage_off_types + + +def test_prepare_messages_for_openai_replays_middleware_terminated_function_group() -> None: + """A terminated function loop replays through its ordinary result contents.""" + client = OpenAIChatClient(model="test-model", api_key="test-key") + + messages = [ + Message( + role="assistant", + contents=[ + Content.from_text_reasoning( + id="rs_terminated", + text="I need to call the guarded tool", + protected_data="encrypted-reasoning", + additional_properties={"status": "completed"}, + ), + Content.from_function_call( + call_id="call_terminated", + name="guarded_tool", + arguments="{}", + ), + ], + ), + Message( + role="tool", + contents=[ + Content.from_function_result( + call_id="call_terminated", + result="Blocked by policy", + ) + ], + ), + ] + + result = client._prepare_messages_for_openai(messages, request_uses_service_side_storage=False) + + types = [item.get("type") for item in result] + assert types == ["reasoning", "function_call", "function_call_output"] + assert result[0]["encrypted_content"] == "encrypted-reasoning" + assert result[2]["output"] == "Blocked by policy" + + +def test_prepare_messages_for_openai_replays_active_parallel_function_group() -> None: + """An active parallel batch retains pending calls and completed siblings.""" + client = OpenAIChatClient(model="test-model", api_key="test-key") + + messages = [ + Message( + role="assistant", + contents=[ + Content.from_text_reasoning( + id="rs_parallel", + text="I need both tools.", + protected_data="encrypted-reasoning", + ), + Content.from_function_call(call_id="call_done", name="first_tool", arguments="{}"), + Content.from_function_call(call_id="call_pending", name="second_tool", arguments="{}"), + ], + ), + Message( + role="tool", + contents=[Content.from_function_result(call_id="call_done", result="first result")], + ), + ] + + result = client._prepare_messages_for_openai(messages, request_uses_service_side_storage=False) + + assert [item["type"] for item in result] == [ + "reasoning", + "function_call", + "function_call", + "function_call_output", + ] + assert [item["call_id"] for item in result[1:3]] == ["call_done", "call_pending"] + assert result[3] == { + "type": "function_call_output", + "call_id": "call_done", + "output": "first result", + } def test_prepare_messages_for_openai_full_conversation_with_reasoning() -> None: @@ -2630,6 +3412,7 @@ def test_prepare_messages_for_openai_full_conversation_with_reasoning() -> None: Content.from_text_reasoning( id="rs_test123", text="I need to search for hotels", + protected_data="encrypted-reasoning", additional_properties={"status": "completed"}, ), Content.from_function_call( @@ -2655,19 +3438,18 @@ def test_prepare_messages_for_openai_full_conversation_with_reasoning() -> None: ), ] - # Storage-off path: function_call kept inline (server has no record of it), - # function_call_output kept. Reasoning is still dropped because rs_* response-scoped IDs - # cannot be replayed against a server that has no record of the originating response. result = client._prepare_messages_for_openai(messages, request_uses_service_side_storage=False) types = [item.get("type") for item in result] assert "message" in types, "User/assistant messages should be present" - assert "function_call" in types, "Function call items must be present without storage" - assert "function_call_output" in types, "Function call output must be present" - - # Verify function_call has id - fc_items = [item for item in result if item.get("type") == "function_call"] - assert fc_items[0]["id"] == "fc_test456" + assert types == ["message", "reasoning", "function_call", "function_call_output", "message"] + assert result[1] == { + "type": "reasoning", + "id": "rs_test123", + "summary": [{"type": "summary_text", "text": "I need to search for hotels"}], + "encrypted_content": "encrypted-reasoning", + "status": "completed", + } def test_prepare_message_for_openai_filters_error_content() -> None: @@ -5230,7 +6012,7 @@ async def test_prepare_options_store_parameter_handling() -> None: assert "previous_response_id" not in options -async def test_prepare_options_store_false_omits_reasoning_items_for_stateless_replay() -> None: +async def test_prepare_options_store_false_rejects_non_replayable_reasoning_items() -> None: client = OpenAIChatClient(model="test-model", api_key="test-key") messages = [ Message(role="user", contents=[Content.from_text(text="search for hotels")]), @@ -5261,11 +6043,8 @@ async def test_prepare_options_store_false_omits_reasoning_items_for_stateless_r ), ] - options = await client._prepare_options(messages, ChatOptions(store=False)) # type: ignore[arg-type] - - assert not any(item.get("type") == "reasoning" for item in options["input"]) - assert any(item.get("type") == "function_call" for item in options["input"]) - assert any(item.get("type") == "function_call_output" for item in options["input"]) + with pytest.raises(ChatClientInvalidRequestException, match="rs_test123.*call_1"): + await client._prepare_options(messages, ChatOptions(store=False)) # type: ignore[arg-type] async def test_prepare_options_with_conversation_id_strips_server_issued_items() -> None: @@ -5825,6 +6604,68 @@ def get_test_image() -> Content: assert "house" in response.text.lower(), f"Model did not describe the house image. Response: {response.text}" +@pytest.mark.flaky +@pytest.mark.integration +@skip_if_openai_integration_tests_disabled +async def test_integration_stateless_reasoning_survives_json_and_checkpoint_round_trip() -> None: + """Encrypted reasoning can be restored from durable storage and replayed on a later request.""" + marker = "STATELESS-REASONING-ROUND-TRIP-7233" + + @tool(name="get_round_trip_marker", approval_mode="never_require") + def get_round_trip_marker() -> str: + """Return the marker that must be repeated in the final answer.""" + return marker + + client = RawOpenAIChatClient(model="gpt-5-mini") + initial_message = Message( + role="user", + contents=["Call get_round_trip_marker, then answer with exactly the value returned by the tool."], + ) + first_options: OpenAIChatOptions[None] = { + "store": False, + "reasoning": {"effort": "low", "summary": "auto"}, + "tools": [get_round_trip_marker], + "tool_choice": {"mode": "required", "required_function_name": "get_round_trip_marker"}, + } + first_response: ChatResponse[Any] = await client.get_response( + [initial_message], + options=first_options, + ) + + first_message = first_response.messages[0] + reasoning_contents = [content for content in first_message.contents if content.type == "text_reasoning"] + assert reasoning_contents + assert any(content.protected_data for content in reasoning_contents) + function_call = next(content for content in first_message.contents if content.type == "function_call") + call_id = function_call.call_id + assert call_id is not None + + message_restored_from_json = Message.from_json(first_message.to_json()) + checkpoint_payload = json.loads(json.dumps(encode_checkpoint_value(message_restored_from_json))) + restored_message = decode_checkpoint_value(checkpoint_payload) + assert isinstance(restored_message, Message) + + final_options: OpenAIChatOptions[None] = { + "store": False, + "reasoning": {"effort": "low", "summary": "auto"}, + "tools": [get_round_trip_marker], + "tool_choice": "none", + } + final_response: ChatResponse[Any] = await client.get_response( + [ + initial_message, + restored_message, + Message( + role="tool", + contents=[Content.from_function_result(call_id=call_id, result=marker)], + ), + ], + options=final_options, + ) + + assert marker in final_response.text + + @pytest.mark.flaky @pytest.mark.integration @skip_if_openai_integration_tests_disabled @@ -6718,14 +7559,18 @@ def test_prepare_messages_for_openai_coalesces_mcp_call_and_result_into_single_i assert fco_items == [], f"unexpected orphan function_call_output items: {fco_items}" -def test_prepare_messages_for_openai_drops_mcp_call_when_paired_reasoning_is_stripped() -> None: +def test_prepare_messages_for_openai_replays_completed_reasoning_bound_mcp_call() -> None: client = OpenAIChatClient(model="test-model", api_key="test-key") messages = [ Message( role="assistant", contents=[ - Content.from_text_reasoning(id="rs_abc123", text="Need the MCP server."), + Content.from_text_reasoning( + id="rs_abc123", + text="Need the MCP server.", + protected_data="encrypted-reasoning", + ), Content.from_mcp_server_tool_call( call_id="mcp_abc123", tool_name="search", @@ -6747,19 +7592,37 @@ def test_prepare_messages_for_openai_drops_mcp_call_when_paired_reasoning_is_str result = client._prepare_messages_for_openai(messages, request_uses_service_side_storage=False) - types = [item.get("type") for item in result if isinstance(item, dict)] - assert "reasoning" not in types - assert "mcp_call" not in types - assert "function_call_output" not in types + assert result == [ + { + "type": "reasoning", + "id": "rs_abc123", + "summary": [{"type": "summary_text", "text": "Need the MCP server."}], + "encrypted_content": "encrypted-reasoning", + }, + { + "type": "mcp_call", + "id": "mcp_abc123", + "server_label": "api_specs", + "name": "search", + "arguments": '{"q": "cats"}', + "output": "found 10 cats", + }, + ] -def test_prepare_messages_for_openai_drops_mcp_call_across_reasoning_messages() -> None: +def test_prepare_messages_for_openai_replays_active_mcp_call_across_reasoning_messages() -> None: client = OpenAIChatClient(model="test-model", api_key="test-key") messages = [ Message( role="assistant", - contents=[Content.from_text_reasoning(id="rs_abc123", text="Need a tool call.")], + contents=[ + Content.from_text_reasoning( + id="rs_abc123", + text="Need a tool call.", + protected_data="encrypted-reasoning", + ) + ], ), Message( role="assistant", @@ -6772,70 +7635,47 @@ def test_prepare_messages_for_openai_drops_mcp_call_across_reasoning_messages() ) ], ), - Message( - role="tool", - contents=[ - Content.from_mcp_server_tool_result( - call_id="mcp_abc123", - output=[Content.from_text(text="found 10 cats")], - ) - ], - ), ] result = client._prepare_messages_for_openai(messages, request_uses_service_side_storage=False) - types = [item.get("type") for item in result if isinstance(item, dict)] - assert "reasoning" not in types - assert "mcp_call" not in types - assert "function_call_output" not in types + assert [item["type"] for item in result] == ["reasoning", "mcp_call"] + assert result[1]["id"] == "mcp_abc123" + assert "output" not in result[1] -def test_prepare_messages_for_openai_keeps_unpaired_mcp_when_reasoning_is_stripped() -> None: +def test_prepare_messages_for_openai_replays_all_mcp_calls_for_one_reasoning_item() -> None: client = OpenAIChatClient(model="test-model", api_key="test-key") messages = [ Message( role="assistant", contents=[ + Content.from_text_reasoning( + id="rs_abc123", + text="Search both indexes.", + protected_data="encrypted-reasoning", + ), Content.from_mcp_server_tool_call( - call_id="mcp_keep", + call_id="mcp_dogs", tool_name="search", server_name="api_specs", arguments='{"q": "dogs"}', - ) - ], - ), - Message( - role="tool", - contents=[ - Content.from_mcp_server_tool_result( - call_id="mcp_keep", - output=[Content.from_text(text="found 5 dogs")], - ) - ], - ), - Message( - role="assistant", - contents=[Content.from_text_reasoning(id="rs_abc123", text="Need a tool call.")], - ), - Message( - role="assistant", - contents=[ + ), Content.from_mcp_server_tool_call( - call_id="mcp_drop", + call_id="mcp_cats", tool_name="search", server_name="api_specs", arguments='{"q": "cats"}', - ) + ), ], ), Message( role="tool", contents=[ Content.from_mcp_server_tool_result( - call_id="mcp_drop", - output=[Content.from_text(text="found 10 cats")], + call_id="mcp_dogs", + output=[Content.from_text(text="found 5 dogs")], ) ], ), @@ -6844,8 +7684,9 @@ def test_prepare_messages_for_openai_keeps_unpaired_mcp_when_reasoning_is_stripp result = client._prepare_messages_for_openai(messages, request_uses_service_side_storage=False) mcp_items = [item for item in result if isinstance(item, dict) and item.get("type") == "mcp_call"] - assert [item["id"] for item in mcp_items] == ["mcp_keep"] + assert [item["id"] for item in mcp_items] == ["mcp_dogs", "mcp_cats"] assert mcp_items[0]["output"] == "found 5 dogs" + assert "output" not in mcp_items[1] def test_prepare_messages_for_openai_drops_orphan_mcp_server_tool_result() -> None: