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 54e1323f62..3524206225 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py @@ -340,13 +340,70 @@ def _reject_busy_conversation(session: AgentSession) -> None: ) -def _create_response_event_stream(context: ResponseContext) -> ResponseEventStream: - """Create a response stream seeded from recovery state when available.""" +def _request_without_reserved_metadata(request: CreateResponse) -> CreateResponse: + """Copy a create request with any client-supplied reserved metadata bag removed. + + ``ResponseEventStream`` seeds the response from the request, and host recovery later trusts + that bag for incomplete status, provider cursors, usage, and citations. Only a persisted + response is a trusted source. The caller's request mapping is left unchanged. + """ + raw_metadata = request.get("metadata") + if not isinstance(raw_metadata, Mapping): + return request + # Metadata is a TypedDict of known public keys. Cast so a client-supplied reserved key is visible. + metadata = cast(Mapping[str, Any], raw_metadata) + if _RESERVED_RESPONSE_METADATA_KEY not in metadata: + return request + public_metadata: dict[str, Any] = { + key: value for key, value in metadata.items() if key != _RESERVED_RESPONSE_METADATA_KEY + } + seeded: dict[str, Any] = dict(request) + if public_metadata: + seeded["metadata"] = public_metadata + else: + seeded.pop("metadata", None) + return cast(CreateResponse, seeded) + + +def _create_response_event_stream(context: ResponseContext, request: CreateResponse) -> ResponseEventStream: + """Create a response stream from recovery state or the originating request.""" if context.is_recovery: persisted_response = context.persisted_response if persisted_response is not None: return ResponseEventStream(response=persisted_response, response_id=context.response_id) - return ResponseEventStream(response_id=context.response_id) + return ResponseEventStream( + request=_request_without_reserved_metadata(request), + response_id=context.response_id, + ) + + +def _hosted_internal_metadata(stream: ResponseEventStream) -> dict[str, Any]: + """Keep private state in the durable envelope, outside the public metadata quota. + + The SDK strips its reserved ``internal_metadata`` key at every wire boundary. Older + checkpoints used its quota-limited response metadata bag; copy that trusted state only + when the envelope has no private bag of its own. + """ + metadata = stream.response.get("internal_metadata") + if metadata is None: + metadata = dict(stream.internal_metadata) + stream.response["internal_metadata"] = metadata + if not isinstance(metadata, dict): + raise RuntimeError("The persisted host metadata is invalid.") + return cast(dict[str, Any], metadata) + + +def _read_persisted_incomplete_reason(stream: ResponseEventStream) -> str | None: + """Return a truncating finish reason restored from a trusted response snapshot.""" + persisted = _hosted_internal_metadata(stream).get(_INCOMPLETE_REASON_KEY) + if isinstance(persisted, str): + return persisted + details = stream.response.get("incomplete_details") + if isinstance(details, Mapping): + reason = cast(Mapping[str, Any], details).get("reason") + if isinstance(reason, str): + return reason + return None def _agent_response_updates(response: AgentResponse[Any], response_id: str) -> list[AgentResponseUpdate]: @@ -550,6 +607,10 @@ async def aclose(self) -> None: # ``internal_metadata`` key carrying a truncating finish reason across resilient checkpoints, so a # turn cut short before a crash still ends ``incomplete`` after recovery. _INCOMPLETE_REASON_KEY = "_incomplete_reason" +# SDK-reserved response metadata entry. It round-trips server-owned recovery state (this reason, +# provider cursors, usage, and citations). Clients may send the same key; it is never trusted on a +# fresh request. Foundry counts it toward the 16-entry public metadata limit. +_RESERVED_RESPONSE_METADATA_KEY = "_internal_metadata" # Foundry Toolbox Auth integration @@ -1128,7 +1189,7 @@ async def _handle_response( async for event in events: yield event return - response_event_stream = _create_response_event_stream(context) + response_event_stream = _create_response_event_stream(context, request) if context.is_steered_turn: logger.debug("Serving steered turn (pending_input_count=%d)", context.pending_input_count) yield response_event_stream.emit_created() @@ -1746,7 +1807,9 @@ async def _handle_inner_agent( if provider_background: if session_storage is None or not isinstance(agent, RawAgent): raise RuntimeError("Provider background requires a stored MAF agent session.") - output_count = response_event_stream.internal_metadata.get(_HOSTED_PROVIDER_OUTPUT_COUNT_KEY, 0) + output_count = _hosted_internal_metadata(response_event_stream).get( + _HOSTED_PROVIDER_OUTPUT_COUNT_KEY, 0 + ) if type(output_count) is not int or output_count < 0: raise RuntimeError("The persisted provider output cursor is invalid.") responses = self._provider_background_responses( @@ -1770,8 +1833,9 @@ async def _handle_inner_agent( continue for event in tracker.close(): yield event - response_event_stream.internal_metadata[_HOSTED_PROVIDER_OUTPUT_COUNT_KEY] = output_count - response_event_stream.internal_metadata[_HOSTED_PROVIDER_USAGE_KEY] = tracker.usage_details + private_metadata = _hosted_internal_metadata(response_event_stream) + private_metadata[_HOSTED_PROVIDER_OUTPUT_COUNT_KEY] = output_count + private_metadata[_HOSTED_PROVIDER_USAGE_KEY] = tracker.usage_details if self._resilient_background: yield response_event_stream.checkpoint() else: @@ -2147,7 +2211,7 @@ async def _handle_inner_workflow( # snapshot (recorded in that snapshot's own metadata) -- NOT simply the latest workflow # checkpoint in storage, which may be ahead of what response.output actually reflects if # the crash happened between two response-stream checkpoint() calls. - checkpoint_id = response_event_stream.internal_metadata.get(_LATEST_CHECKPOINT_ID_KEY) + checkpoint_id = _hosted_internal_metadata(response_event_stream).get(_LATEST_CHECKPOINT_ID_KEY) if checkpoint_id is not None: logger.debug("Serving recovery request from workflow checkpoint %s", checkpoint_id) run_stream = self._resume_workflow_from_checkpoint( @@ -2263,12 +2327,12 @@ def snapshot_response( ) -> Generator[ResponseStreamEvent | ResponseCheckpointEvent]: # Pair the response output emitted so far with the workflow checkpoint it corresponds # to, so recovery from that checkpoint replays exactly the updates that came after it. - if checkpoint_id is None or checkpoint_id == response_event_stream.internal_metadata.get( + if checkpoint_id is None or checkpoint_id == _hosted_internal_metadata(response_event_stream).get( _LATEST_CHECKPOINT_ID_KEY ): return yield from tracker.close() - response_event_stream.internal_metadata[_LATEST_CHECKPOINT_ID_KEY] = checkpoint_id + _hosted_internal_metadata(response_event_stream)[_LATEST_CHECKPOINT_ID_KEY] = checkpoint_id yield response_event_stream.checkpoint() main_iter = _SignalledIterator( @@ -2396,7 +2460,7 @@ def __init__( self._stream = stream self._allowed_oauth_consent_origins = allowed_oauth_consent_origins self._usage_details: UsageDetails | None = None - persisted_usage = stream.internal_metadata.get(_HOSTED_PROVIDER_USAGE_KEY) + persisted_usage = _hosted_internal_metadata(stream).get(_HOSTED_PROVIDER_USAGE_KEY) if persisted_usage is not None: if not isinstance(persisted_usage, Mapping) or any( not isinstance(key, str) or (value is not None and type(value) is not int) @@ -2428,15 +2492,15 @@ def __init__( self._message_text_annotations: dict[int, list[ContainerFileCitationBody]] = {} # Set when an agent update reports the model stopped early (content filter, token # limit); the response then ends as ``incomplete`` instead of ``completed`` so callers - # can tell a cut-short turn from a successful one. Mirrored into the stream's - # ``internal_metadata`` so it survives a resilient checkpoint/recovery cycle, which - # rebuilds this tracker from the persisted response. + # can tell a cut-short turn from a successful one. Mirrored onto the response so it + # survives a resilient checkpoint/recovery cycle, which rebuilds this tracker from the + # persisted response. self._incomplete_reason: ResponseIncompleteReason | None = None - persisted_reason = stream.internal_metadata.get(_INCOMPLETE_REASON_KEY) - if isinstance(persisted_reason, str): + persisted_reason = _read_persisted_incomplete_reason(stream) + if persisted_reason is not None: with suppress(ValueError): self._incomplete_reason = ResponseIncompleteReason(persisted_reason) - persisted_citations = stream.internal_metadata.get(_PENDING_CONTAINER_FILE_CITATIONS_KEY) + persisted_citations = _hosted_internal_metadata(stream).get(_PENDING_CONTAINER_FILE_CITATIONS_KEY) if persisted_citations is not None: if not isinstance(persisted_citations, Sequence) or isinstance( persisted_citations, (str, bytes, bytearray) @@ -2524,7 +2588,8 @@ def record_finish_reason(self, finish_reason: str | None) -> None: Only finish reasons that mean the model stopped early are retained, mapped onto the Responses ``incomplete_details.reason`` vocabulary. A content filter is kept in preference to a token limit if both are seen during a multi-step turn, since it is the - more actionable signal for the caller. + more actionable signal for the caller. The reason is recorded even when the caller's + metadata already fills the public map, and that map is left intact. """ if finish_reason == "content_filter": self._incomplete_reason = ResponseIncompleteReason.CONTENT_FILTER @@ -2532,7 +2597,7 @@ def record_finish_reason(self, finish_reason: str | None) -> None: self._incomplete_reason = ResponseIncompleteReason.MAX_OUTPUT_TOKENS else: return - self._stream.internal_metadata[_INCOMPLETE_REASON_KEY] = self._incomplete_reason.value + _hosted_internal_metadata(self._stream)[_INCOMPLETE_REASON_KEY] = self._incomplete_reason.value def _consume_text_update(self, text: str, merge_mode: Literal["append", "replace"] | None) -> str: """Update the authoritative text buffer and return the wire delta still needed.""" @@ -3118,9 +3183,9 @@ def _container_file_annotations_for_text(self, text: str) -> list[ContainerFileC def _persist_pending_container_file_citations(self) -> None: if not self._pending_container_file_citations: - self._stream.internal_metadata.pop(_PENDING_CONTAINER_FILE_CITATIONS_KEY, None) + _hosted_internal_metadata(self._stream).pop(_PENDING_CONTAINER_FILE_CITATIONS_KEY, None) return - self._stream.internal_metadata[_PENDING_CONTAINER_FILE_CITATIONS_KEY] = [ + _hosted_internal_metadata(self._stream)[_PENDING_CONTAINER_FILE_CITATIONS_KEY] = [ { "container_id": citation.container_id, "file_id": citation.file_id, diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_workflow_responses.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_workflow_responses.py index 811ecacbfb..2e3c37ee6e 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_workflow_responses.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_workflow_responses.py @@ -53,8 +53,10 @@ from ._responses import ( _HOSTED_PROVIDER_USAGE_KEY, # pyright: ignore[reportPrivateUsage] _agent_response_updates, # pyright: ignore[reportPrivateUsage] + _hosted_internal_metadata, # pyright: ignore[reportPrivateUsage] _item_to_message, # pyright: ignore[reportPrivateUsage] _OutputItemTracker, # pyright: ignore[reportPrivateUsage] + _request_without_reserved_metadata, # pyright: ignore[reportPrivateUsage] _SignalledIterator, # pyright: ignore[reportPrivateUsage] ) from ._scope import FoundryRequestScope @@ -381,7 +383,10 @@ async def handle_response( cancellation_signal: asyncio.Event, ) -> AsyncGenerator[ResponseStreamEvent | ResponseCheckpointEvent]: """Validate, run, pair, close owned resources, and only then publish terminal success.""" - stream = ResponseEventStream(response_id=context.response_id, request=request) + stream = ResponseEventStream( + response_id=context.response_id, + request=_request_without_reserved_metadata(request), + ) run: HostedWorkflowRun | None = None phase = "preparation" completed = False @@ -504,7 +509,7 @@ async def flush(checkpoint_id: str | None) -> list[ResponseStreamEvent]: for event in buffered: events.extend(await _project(event, draft, next_tracker, approval_projection)) events.extend(next_tracker.close()) - draft.internal_metadata[_HOSTED_PROVIDER_USAGE_KEY] = next_tracker.usage_details + _hosted_internal_metadata(draft)[_HOSTED_PROVIDER_USAGE_KEY] = next_tracker.usage_details await run.stage(_snapshot(draft), approvals=next_aliases, checkpoint_id=checkpoint_id) stream, tracker, aliases = draft, next_tracker, next_aliases buffered.clear() diff --git a/python/packages/foundry_hosting/tests/test_responses.py b/python/packages/foundry_hosting/tests/test_responses.py index aa9a5f0ce9..2ba80fd457 100644 --- a/python/packages/foundry_hosting/tests/test_responses.py +++ b/python/packages/foundry_hosting/tests/test_responses.py @@ -75,6 +75,7 @@ from azure.ai.agentserver.responses.models import ( CreateResponse, Item, + Metadata, OutputItem, ResponseIncompleteReason, ResponseObject, @@ -96,6 +97,8 @@ _agent_response_updates, # pyright: ignore[reportPrivateUsage] _await_before_signal, # pyright: ignore[reportPrivateUsage] _container_file_citations_from_function_result, # pyright: ignore[reportPrivateUsage] + _create_response_event_stream, # pyright: ignore[reportPrivateUsage] + _hosted_internal_metadata, # pyright: ignore[reportPrivateUsage] _is_allowed_oauth_consent_link, # pyright: ignore[reportPrivateUsage] _item_to_message, # pyright: ignore[reportPrivateUsage] _json_safe_to_str, # pyright: ignore[reportPrivateUsage] @@ -788,6 +791,8 @@ async def _post( parallel_tool_calls: bool | None = None, previous_response_id: str | None = None, conversation_id: str | None = None, + metadata: Mapping[str, str] | None = None, + store: bool | None = None, ) -> httpx.Response: """Send a POST /responses request through the ASGI transport.""" payload: dict[str, Any] = {"model": model, "input": input_text, "stream": stream} @@ -803,12 +808,23 @@ async def _post( payload["previous_response_id"] = previous_response_id if conversation_id is not None: payload["conversation"] = conversation_id + if metadata is not None: + payload["metadata"] = dict(metadata) + if store is not None: + payload["store"] = store transport = httpx.ASGITransport(app=server) async with httpx.AsyncClient(transport=transport, base_url="http://test") as client: return await client.post("/responses", json=payload) +async def _get_response(server: ResponsesHostServer, response_id: str) -> httpx.Response: + """GET one stored response through the ASGI transport.""" + transport = httpx.ASGITransport(app=server) + async with httpx.AsyncClient(transport=transport, base_url="http://test") as client: + return await client.get(f"/responses/{response_id}") + + def _parse_sse_events(body: str) -> list[dict[str, Any]]: """Parse SSE text into a list of event dicts with 'event' and 'data' keys.""" events: list[dict[str, Any]] = [] @@ -1104,6 +1120,195 @@ def _sse_event_types(events: list[dict[str, Any]]) -> list[str]: # region Serialization Helpers +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("metadata", [None, {}, {"foo": "bar", "view": "concise"}]) +async def test_request_metadata_is_returned_and_stored(stream: bool, metadata: dict[str, str] | None) -> None: + agent = _make_agent( + response=AgentResponse(messages=[Message(role="assistant", contents=[Content.from_text("Hello!")])]) + ) + server = _make_server(agent) + payload: dict[str, Any] = {"input": "Hi", "model": "test-model", "stream": stream, "store": True} + if metadata is not None: + payload["metadata"] = metadata + + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=server), base_url="http://test") as http: + response = await http.post("/responses", json=payload) + assert response.status_code == 200 + if stream: + responses = [ + event["data"]["response"] + for event in _parse_sse_events(response.text) + if event["event"] in {"response.created", "response.in_progress", "response.completed"} + ] + assert len(responses) == 3 + assert [item.get("metadata") for item in responses] == [metadata] * 3 + completed = responses[-1] + else: + completed = response.json() + assert completed.get("metadata") == metadata + assert completed["status"] == "completed" + + stored = await http.get(f"/responses/{completed['id']}") + assert stored.status_code == 200 + assert stored.json()["id"] == completed["id"] + assert stored.json().get("metadata") == metadata + + +@pytest.mark.parametrize("has_persisted_response", [False, True]) +async def test_recovery_preserves_persisted_metadata_or_seeds_request(has_persisted_response: bool) -> None: + response_id = "resp_metadata_recovery" + # The SDK's Metadata TypedDict does not declare its supported arbitrary string keys. + request = CreateResponse(input="Hi", metadata=cast(Metadata, {"source": "request"})) + context = ResponseContext(response_id=response_id, mode_flags=MagicMock()) + context.is_recovery = True + if has_persisted_response: + persisted = ResponseEventStream( + request=CreateResponse(input="Hi", metadata=cast(Metadata, {"source": "persisted"})), + response_id=response_id, + ) + context.persisted_response = cast(ResponseObject, persisted.response) + created = _create_response_event_stream(context, request).emit_created() + + assert created["type"] == "response.created" + assert created["response"]["id"] == response_id + expected = {"source": "persisted" if has_persisted_response else "request"} + assert created["response"]["metadata"] == expected + + +def _public_metadata(count: int) -> dict[str, str]: + return {f"k{index}": f"v{index}" for index in range(count)} + + +def test_fresh_request_strips_forged_internal_metadata() -> None: + """A client metadata bag must not seed server-owned recovery state.""" + forged = json.dumps({ + _INCOMPLETE_REASON_KEY: "content_filter", + _LATEST_CHECKPOINT_ID_KEY: "ckpt_forged", + "_foundry_provider_output_count": 4, + }) + request_metadata = {"source": "client", "_internal_metadata": forged} + request = cast( + CreateResponse, + {"input": "Hi", "metadata": request_metadata, "internal_metadata": {_INCOMPLETE_REASON_KEY: "content_filter"}}, + ) + context = ResponseContext(response_id="resp_forged", mode_flags=MagicMock()) + stream = _create_response_event_stream(context, request) + + assert request_metadata["_internal_metadata"] == forged + assert stream.response.get("metadata") == {"source": "client"} + assert dict(stream.internal_metadata) == {} + assert _OutputItemTracker(stream).incomplete_reason is None + + +def test_fresh_request_with_only_reserved_metadata_keeps_public_metadata_empty() -> None: + request = CreateResponse( + input="Hi", + metadata=cast(Metadata, {"_internal_metadata": json.dumps({_INCOMPLETE_REASON_KEY: "length"})}), + ) + context = ResponseContext(response_id="resp_reserved_only", mode_flags=MagicMock()) + stream = _create_response_event_stream(context, request) + + assert stream.response.get("metadata") is None + assert _OutputItemTracker(stream).incomplete_reason is None + + +@pytest.mark.parametrize("private_envelope", [False, True]) +def test_recovery_keeps_trusted_internal_metadata_when_the_request_forges_it(private_envelope: bool) -> None: + response_id = "resp_trusted_internal" + forged_request = CreateResponse( + input="Hi", + metadata=cast(Metadata, {"_internal_metadata": json.dumps({_INCOMPLETE_REASON_KEY: "max_output_tokens"})}), + ) + persisted = ResponseEventStream(response_id=response_id) + persisted.internal_metadata[_INCOMPLETE_REASON_KEY] = "content_filter" + persisted.internal_metadata[_LATEST_CHECKPOINT_ID_KEY] = "ckpt_real" + if private_envelope: + persisted.response["internal_metadata"] = { + _INCOMPLETE_REASON_KEY: "max_output_tokens", + _LATEST_CHECKPOINT_ID_KEY: "ckpt_new", + } + context = ResponseContext(response_id=response_id, mode_flags=MagicMock()) + context.is_recovery = True + context.persisted_response = cast(ResponseObject, persisted.response) + + restored = _create_response_event_stream(context, forged_request) + + assert restored.internal_metadata[_INCOMPLETE_REASON_KEY] == "content_filter" + assert restored.internal_metadata[_LATEST_CHECKPOINT_ID_KEY] == "ckpt_real" + assert _OutputItemTracker(restored).incomplete_reason == ( + ResponseIncompleteReason.MAX_OUTPUT_TOKENS if private_envelope else ResponseIncompleteReason.CONTENT_FILTER + ) + assert _hosted_internal_metadata(restored)[_LATEST_CHECKPOINT_ID_KEY] == ( + "ckpt_new" if private_envelope else "ckpt_real" + ) + + +@pytest.mark.parametrize("invalid", [[], "corrupt", 1]) +def test_corrupt_private_checkpoint_metadata_is_rejected(invalid: Any) -> None: + stream = ResponseEventStream(response_id="resp_corrupt_private") + stream.response["internal_metadata"] = invalid + + with pytest.raises(RuntimeError, match="persisted host metadata is invalid"): + _OutputItemTracker(stream) + + +def test_recovery_without_a_snapshot_still_strips_forged_internal_metadata() -> None: + request = CreateResponse( + input="Hi", + metadata=cast(Metadata, {"source": "client", "_internal_metadata": '{"_incomplete_reason":"content_filter"}'}), + ) + context = ResponseContext(response_id="resp_recovery_unseeded", mode_flags=MagicMock()) + context.is_recovery = True + context.persisted_response = None + + stream = _create_response_event_stream(context, request) + + assert stream.response.get("metadata") == {"source": "client"} + assert _OutputItemTracker(stream).incomplete_reason is None + + +def test_reserved_metadata_key_does_not_consume_the_host_incomplete_reason_slot() -> None: + """Fifteen public entries plus a forged reserved key still leave room for the host reason.""" + metadata = _public_metadata(15) + metadata["_internal_metadata"] = json.dumps({_INCOMPLETE_REASON_KEY: "max_output_tokens"}) + request = CreateResponse(input="Hi", metadata=cast(Metadata, metadata)) + context = ResponseContext(response_id="resp_freed_slot", mode_flags=MagicMock()) + stream = _create_response_event_stream(context, request) + tracker = _OutputItemTracker(stream) + + tracker.record_finish_reason("content_filter") + + assert tracker.incomplete_reason == ResponseIncompleteReason.CONTENT_FILTER + stored = stream.response["internal_metadata"] + assert stored[_INCOMPLETE_REASON_KEY] == "content_filter" + assert all(stream.response["metadata"][key] == value for key, value in _public_metadata(15).items()) + + +def test_full_public_metadata_records_incomplete_reason_outside_the_quota() -> None: + metadata = _public_metadata(16) + request = CreateResponse(input="Hi", metadata=cast(Metadata, metadata)) + context = ResponseContext(response_id="resp_full_metadata", mode_flags=MagicMock()) + stream = _create_response_event_stream(context, request) + tracker = _OutputItemTracker(stream) + + tracker.record_finish_reason("length") + tracker.record_finish_reason("content_filter") + _hosted_internal_metadata(stream)[_LATEST_CHECKPOINT_ID_KEY] = "ckpt_saved" + + assert tracker.incomplete_reason == ResponseIncompleteReason.CONTENT_FILTER + assert stream.response["metadata"] == metadata + assert "_internal_metadata" not in stream.response["metadata"] + assert stream.response["internal_metadata"][_LATEST_CHECKPOINT_ID_KEY] == "ckpt_saved" + assert stream.response["internal_metadata"][_INCOMPLETE_REASON_KEY] == "content_filter" + + recovered = _OutputItemTracker( + ResponseEventStream(response=cast(ResponseObject, stream.response), response_id="resp_full_metadata") + ) + assert recovered.incomplete_reason == ResponseIncompleteReason.CONTENT_FILTER + recovered.record_finish_reason("stop") + assert recovered.incomplete_reason == ResponseIncompleteReason.CONTENT_FILTER + + class TestSerializationHelpers: def test_json_safe_to_str_preserves_structured_conversion_and_falls_back_to_string(self) -> None: @dataclass @@ -2773,6 +2978,128 @@ async def complete_with_token() -> str: assert completed is False and result is None operation.assert_not_called() + @pytest.mark.parametrize("metadata_count", [15, 16]) + @pytest.mark.parametrize("crash_stage", ["before-checkpoint", "after-checkpoint"]) + @pytest.mark.parametrize("finish_reason", ["stop", "length", "content_filter"]) + async def test_provider_recovery_keeps_checkpointed_output_once( + self, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + metadata_count: int, + crash_stage: str, + finish_reason: FinishReasonLiteral, + ) -> None: + response_id = IdGenerator.new_response_id() + metadata = _public_metadata(metadata_count) + calls: list[dict[str, Any]] = [] + + async def run( + messages: Any = None, + *, + session: AgentSession, + options: Mapping[str, Any], + **kwargs: Any, + ) -> AgentResponse: + del messages, kwargs + calls.append(dict(options)) + if "continuation_token" not in options: + return AgentResponse( + messages=[], continuation_token=OpenAIContinuationToken(response_id="private-cursor-token") + ) + session.service_session_id = "private-cursor-conversation" + return AgentResponse( + messages=[Message(role="assistant", contents=[Content.from_text("done")])], + usage_details={"input_token_count": 3, "output_token_count": 2, "total_token_count": 5}, + finish_reason=FinishReason(finish_reason), + ) + + def make_host() -> ResponsesHostServer: + agent = Agent(client=_ServiceStorageRecordingClient()) + monkeypatch.setattr(agent, "run", MagicMock(side_effect=run)) + return _make_server( + agent, + history_source="service", + background_source="provider", + options=ResponsesServerOptions(resilient_background=True), + response_store=FileResponseStore(storage_dir=tmp_path), + ) + + request = CreateResponse(input="hello", store=True, background=True, metadata=cast(Metadata, metadata)) + server = make_host() + context = ResponseContext(response_id=response_id, mode_flags=MagicMock()) + response_store = FileResponseStore(storage_dir=tmp_path) + with ( + patch.object(ResponseContext, "get_input_items", new=AsyncMock(return_value=[])), + patch("agent_framework_foundry_hosting._responses.asyncio.sleep", new=AsyncMock()), + ): + handler = cast(AsyncGenerator[Any, None], server._handle_response(request, context, asyncio.Event())) + crashed = False + try: + async for event in handler: + if isinstance(event, ResponseCheckpointEvent): + await response_store.update_response(copy.deepcopy(event.response)) + if crash_stage == "after-checkpoint": + crashed = True + break + elif isinstance(event, Mapping): + if event.get("type") == "response.created": + await response_store.create_response(copy.deepcopy(event["response"]), [], []) + elif crash_stage == "before-checkpoint" and event.get("type") == "response.output_item.done": + crashed = True + break + finally: + await handler.aclose() + assert crashed + checkpoint = await FileResponseStore(storage_dir=tmp_path).get_response(response_id) + assert checkpoint is not None + session_store = AgentSessionStoreProvider().get_store( + config=server.config, platform_context=get_request_context() + ) + saved = await session_store.get(response_id) + assert saved is not None + private_checkpoint = saved.to_dict() + + for _ in range(2): + # Restart independently from the same paired durable checkpoint. + restore_store = AgentSessionStoreProvider().get_store( + config=server.config, platform_context=get_request_context() + ) + await restore_store.get(response_id) + await restore_store.set(response_id, AgentSession.from_dict(copy.deepcopy(private_checkpoint))) + await FileResponseStore(storage_dir=tmp_path).update_response(copy.deepcopy(checkpoint)) + recovered = ResponseContext(response_id=response_id, mode_flags=MagicMock()) + recovered.is_recovery = True + recovered.persisted_response = await FileResponseStore(storage_dir=tmp_path).get_response(response_id) + restarted = make_host() + events = [event async for event in restarted._handle_response(request, recovered, asyncio.Event())] + terminal = events[-1] + assert isinstance(terminal, Mapping) + if finish_reason == "stop": + assert terminal["type"] == "response.completed" + else: + assert terminal["type"] == "response.incomplete" + response = terminal["response"] + assert len(response["output"]) == 1 + message = response["output"][0] + assert message["type"] == "message" + assert len(message["content"]) == 1 + content = message["content"][0] + assert content["type"] == "output_text" + assert content["text"] == "done" + assert response["usage"]["total_tokens"] == 5 + if finish_reason != "stop": + expected_reason = "max_output_tokens" if finish_reason == "length" else "content_filter" + assert response["incomplete_details"] == {"reason": expected_reason} + await response_store.update_response(response) + fetched = await _get_response(restarted, response_id) + assert fetched.status_code == 200, fetched.text + public = fetched.json() + assert public["metadata"] == metadata + assert "internal_metadata" not in public + assert "_internal_metadata" not in public["metadata"] + assert "private-cursor-token" not in str(public) + assert len(calls) == 2 + async def test_final_provider_poll_keeps_token_for_crash_recovery( self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: @@ -8335,7 +8662,7 @@ async def test_incomplete_reason_survives_checkpoint_recovery(self) -> None: stream.emit_in_progress() tracker = _OutputItemTracker(stream) tracker.record_finish_reason("content_filter") - assert stream.internal_metadata[_INCOMPLETE_REASON_KEY] == "content_filter" + assert stream.response["internal_metadata"][_INCOMPLETE_REASON_KEY] == "content_filter" # Simulate recovery: a fresh tracker over the checkpointed response snapshot. recovered = _OutputItemTracker(stream) @@ -8346,6 +8673,84 @@ async def test_incomplete_reason_survives_checkpoint_recovery(self) -> None: # A stream that was never marked restores nothing. assert _OutputItemTracker(ResponseEventStream(response_id="resp_clean")).incomplete_reason is None + @pytest.mark.parametrize( + ("stream", "finish_reason", "expected_reason"), + [ + (False, "content_filter", "content_filter"), + (False, "length", "max_output_tokens"), + (True, "content_filter", "content_filter"), + (True, "length", "max_output_tokens"), + ], + ) + async def test_sixteen_public_metadata_keys_keep_truncated_output( + self, + stream: bool, + finish_reason: FinishReasonLiteral, + expected_reason: str, + ) -> None: + """A full public metadata map must not turn a truncated turn into a failure or drop its text.""" + metadata = _public_metadata(16) + expected_text = "I'm sorry, but I cannot assist with that request." + server = _make_server(self._filtered_agent(finish_reason=finish_reason)) + + resp = await _post(server, input_text="hello", stream=stream, metadata=metadata, store=True) + assert resp.status_code == 200 + if stream: + events = _parse_sse_events(resp.text) + event_types = _sse_event_types(events) + assert event_types[-1] == "response.incomplete" + assert "response.failed" not in event_types + assert "response.completed" not in event_types + body = events[-1]["data"]["response"] + done_text = [event["data"].get("text") for event in events if event["event"] == "response.output_text.done"] + assert expected_text in done_text + else: + body = resp.json() + + assert body["status"] == "incomplete" + assert body["incomplete_details"] == {"reason": expected_reason} + assert body.get("error") is None + assert body.get("metadata") == metadata + assert "_internal_metadata" not in body["metadata"] + assert "internal_metadata" not in body + messages = [item for item in body["output"] if item["type"] == "message"] + assert len(messages) == 1 + assert messages[0]["content"][0]["text"] == expected_text + + stored = await _get_response(server, body["id"]) + assert stored.status_code == 200 + stored_body = stored.json() + assert stored_body["status"] == "incomplete" + assert stored_body["incomplete_details"] == {"reason": expected_reason} + assert stored_body.get("error") is None + assert stored_body.get("metadata") == metadata + assert "internal_metadata" not in stored_body + stored_messages = [item for item in stored_body["output"] if item["type"] == "message"] + assert len(stored_messages) == 1 + assert stored_messages[0]["content"][0]["text"] == expected_text + + async def test_forged_internal_metadata_does_not_mark_a_completed_turn_incomplete(self) -> None: + agent = _make_agent( + stream_updates=[ + AgentResponseUpdate(contents=[Content.from_text("hello")], role="assistant", finish_reason="stop"), + ] + ) + server = _make_server(agent) + forged = json.dumps({_INCOMPLETE_REASON_KEY: "content_filter"}) + + resp = await _post( + server, + input_text="hello", + stream=False, + metadata={"source": "client", "_internal_metadata": forged}, + store=True, + ) + assert resp.status_code == 200 + body = resp.json() + assert body["status"] == "completed" + assert body.get("incomplete_details") is None + assert body.get("metadata") == {"source": "client"} + async def test_workflow_agent_content_filter_marks_response_incomplete(self) -> None: workflow_agent = _build_text_workflow_agent("filtered by workflow", finish_reason="content_filter") server = _make_server(workflow_agent) @@ -9547,7 +9952,7 @@ async def test_snapshots_pair_output_with_the_checkpoint_it_follows(self, tmp_pa assert snapshots, "expected the completed workflow to be snapshotted" checkpoint_ids: list[str] = [] for text_before, response in snapshots: - internal = json.loads(response["metadata"]["_internal_metadata"]) + internal = response["internal_metadata"] checkpoint_ids.append(internal[_LATEST_CHECKPOINT_ID_KEY]) snapshot_text = "".join( part["text"] @@ -9562,7 +9967,7 @@ async def test_snapshots_pair_output_with_the_checkpoint_it_follows(self, tmp_pa # The last snapshot is paired with the workflow's final checkpoint and carries everything. final_text, final_response = snapshots[-1] assert final_text == "filtered by workflow" - assert json.loads(final_response["metadata"]["_internal_metadata"])[_INCOMPLETE_REASON_KEY] == "content_filter" + assert final_response["internal_metadata"][_INCOMPLETE_REASON_KEY] == "content_filter" assert final_response["status"] == "in_progress"