From 6037edd0fa65f7548a5041d8623641d91c16b961 Mon Sep 17 00:00:00 2001 From: its-amann Date: Tue, 27 Oct 2026 12:46:01 +0530 Subject: [PATCH 1/4] fix(foundry-hosting): preserve response request metadata --- .../_responses.py | 8 +-- .../foundry_hosting/tests/test_responses.py | 57 +++++++++++++++++++ 2 files changed, 61 insertions(+), 4 deletions(-) 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..284e463643 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,13 @@ 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 _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, response_id=context.response_id) def _agent_response_updates(response: AgentResponse[Any], response_id: str) -> list[AgentResponseUpdate]: @@ -1128,7 +1128,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() diff --git a/python/packages/foundry_hosting/tests/test_responses.py b/python/packages/foundry_hosting/tests/test_responses.py index aa9a5f0ce9..3b17ce47ea 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,7 @@ _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] _is_allowed_oauth_consent_link, # pyright: ignore[reportPrivateUsage] _item_to_message, # pyright: ignore[reportPrivateUsage] _json_safe_to_str, # pyright: ignore[reportPrivateUsage] @@ -1104,6 +1106,61 @@ 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 + + class TestSerializationHelpers: def test_json_safe_to_str_preserves_structured_conversion_and_falls_back_to_string(self) -> None: @dataclass From bbdb8c2ed0f9525a7e9ad1a81a5276c875bb46ad Mon Sep 17 00:00:00 2001 From: Aman Agnihotri Date: Tue, 27 Oct 2026 09:32:09 +0000 Subject: [PATCH 2/4] fix(foundry-hosting): keep recovery metadata server-owned Drop a client-supplied _internal_metadata bag when a fresh response is seeded, and keep that bag only for a trusted persisted snapshot. When public metadata already holds 16 entries, still record a truncating finish reason so the response ends incomplete instead of failed. --- .../_responses.py | 109 +++++++++-- .../_workflow_responses.py | 9 +- .../foundry_hosting/tests/test_responses.py | 181 ++++++++++++++++++ 3 files changed, 278 insertions(+), 21 deletions(-) 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 284e463643..eecf460f0b 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,75 @@ def _reject_busy_conversation(session: AgentSession) -> None: ) +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(request=request, response_id=context.response_id) + return ResponseEventStream( + request=_request_without_reserved_metadata(request), + response_id=context.response_id, + ) + + +def _write_hosted_internal(stream: ResponseEventStream, key: str, value: Any) -> None: + """Persist one host-owned recovery value without displacing public metadata. + + The SDK stores this bag in the public metadata map and rejects a 17th entry. A full map of + caller metadata must still keep every public entry. The incomplete reason then lives on + ``incomplete_details``, which is outside that quota, so a truncated turn can still finish + ``incomplete``. Other host cursors are skipped rather than failing the response; the public + map is left as the caller sent it. + """ + try: + stream.internal_metadata[key] = value + except ValueError: + if key == _INCOMPLETE_REASON_KEY and isinstance(value, str): + stream.response["incomplete_details"] = {"reason": value} + return + logger.debug( + "Skipping host metadata %s because the response metadata map cannot store another entry.", + key, + ) + + +def _read_persisted_incomplete_reason(stream: ResponseEventStream) -> str | None: + """Return a truncating finish reason restored from a trusted response snapshot.""" + persisted = stream.internal_metadata.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 +612,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 @@ -1770,8 +1836,8 @@ 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 + _write_hosted_internal(response_event_stream, _HOSTED_PROVIDER_OUTPUT_COUNT_KEY, output_count) + _write_hosted_internal(response_event_stream, _HOSTED_PROVIDER_USAGE_KEY, tracker.usage_details) if self._resilient_background: yield response_event_stream.checkpoint() else: @@ -2268,7 +2334,7 @@ def snapshot_response( ): return yield from tracker.close() - response_event_stream.internal_metadata[_LATEST_CHECKPOINT_ID_KEY] = checkpoint_id + _write_hosted_internal(response_event_stream, _LATEST_CHECKPOINT_ID_KEY, checkpoint_id) yield response_event_stream.checkpoint() main_iter = _SignalledIterator( @@ -2428,12 +2494,12 @@ 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) @@ -2524,7 +2590,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 +2599,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 + _write_hosted_internal(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.""" @@ -3120,14 +3187,18 @@ 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) return - self._stream.internal_metadata[_PENDING_CONTAINER_FILE_CITATIONS_KEY] = [ - { - "container_id": citation.container_id, - "file_id": citation.file_id, - "filename": citation.filename, - } - for citation in self._pending_container_file_citations.values() - ] + _write_hosted_internal( + self._stream, + _PENDING_CONTAINER_FILE_CITATIONS_KEY, + [ + { + "container_id": citation.container_id, + "file_id": citation.file_id, + "filename": citation.filename, + } + for citation in self._pending_container_file_citations.values() + ], + ) # endregion 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..cfca38eddf 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 @@ -55,7 +55,9 @@ _agent_response_updates, # pyright: ignore[reportPrivateUsage] _item_to_message, # pyright: ignore[reportPrivateUsage] _OutputItemTracker, # pyright: ignore[reportPrivateUsage] + _request_without_reserved_metadata, # pyright: ignore[reportPrivateUsage] _SignalledIterator, # pyright: ignore[reportPrivateUsage] + _write_hosted_internal, # pyright: ignore[reportPrivateUsage] ) from ._scope import FoundryRequestScope from ._state_store import CheckpointStoreProvider, ContextScopedStoreProvider @@ -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 + _write_hosted_internal(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 3b17ce47ea..a614bac019 100644 --- a/python/packages/foundry_hosting/tests/test_responses.py +++ b/python/packages/foundry_hosting/tests/test_responses.py @@ -107,6 +107,7 @@ _OutputItemTracker, # pyright: ignore[reportPrivateUsage] _SignalledIterator, # pyright: ignore[reportPrivateUsage] _stringify_mcp_output, # pyright: ignore[reportPrivateUsage] + _write_hosted_internal, # pyright: ignore[reportPrivateUsage] consent_url_from_error, ) from agent_framework_foundry_hosting._state_store import ( @@ -790,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} @@ -805,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]] = [] @@ -1161,6 +1175,117 @@ async def test_recovery_preserves_persisted_metadata_or_seeds_request(has_persis 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 = CreateResponse(input="Hi", metadata=cast(Metadata, request_metadata)) + 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 + + +def test_recovery_keeps_trusted_internal_metadata_when_the_request_forges_it() -> 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" + 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.CONTENT_FILTER + + +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 = json.loads(stream.response["metadata"]["_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") + _write_hosted_internal(stream, _LATEST_CHECKPOINT_ID_KEY, "ckpt_skipped") + + assert tracker.incomplete_reason == ResponseIncompleteReason.CONTENT_FILTER + assert stream.response["metadata"] == metadata + assert "_internal_metadata" not in stream.response["metadata"] + assert stream.internal_metadata.get(_LATEST_CHECKPOINT_ID_KEY) is None + assert stream.response["incomplete_details"] == {"reason": "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 @@ -8403,6 +8528,62 @@ 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 + async def test_full_public_metadata_content_filter_stays_incomplete(self) -> None: + """Sixteen caller entries are valid and must not turn a truncated turn into a failure.""" + metadata = _public_metadata(16) + server = _make_server(self._filtered_agent(finish_reason="content_filter")) + + resp = await _post(server, input_text="hello", stream=False, metadata=metadata, store=True) + assert resp.status_code == 200 + body = resp.json() + assert body["status"] == "incomplete" + assert body["incomplete_details"] == {"reason": "content_filter"} + assert body.get("error") is None + assert body.get("metadata") == metadata + + 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": "content_filter"} + assert stored_body.get("metadata") == metadata + + async def test_full_public_metadata_streaming_length_stays_incomplete(self) -> None: + metadata = _public_metadata(16) + server = _make_server(self._filtered_agent(finish_reason="length")) + + resp = await _post(server, input_text="hello", stream=True, metadata=metadata) + assert resp.status_code == 200 + events = _parse_sse_events(resp.text) + assert _sse_event_types(events)[-1] == "response.incomplete" + assert "response.failed" not in _sse_event_types(events) + incomplete = events[-1]["data"]["response"] + assert incomplete["status"] == "incomplete" + assert incomplete["incomplete_details"] == {"reason": "max_output_tokens"} + assert incomplete.get("metadata") == metadata + + 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) From 7007e73ee76987c59d5462eba66691498c230565 Mon Sep 17 00:00:00 2001 From: Aman Agnihotri Date: Tue, 27 Oct 2026 10:24:30 +0000 Subject: [PATCH 3/4] test(foundry-hosting): cover 16-key truncation in both response modes A full public metadata map still ends length and content_filter turns incomplete, with the assistant text and all 16 public keys kept, for both streaming and non-streaming responses. --- .../foundry_hosting/tests/test_responses.py | 62 ++++++++++++------- 1 file changed, 41 insertions(+), 21 deletions(-) diff --git a/python/packages/foundry_hosting/tests/test_responses.py b/python/packages/foundry_hosting/tests/test_responses.py index a614bac019..13e3a91bb5 100644 --- a/python/packages/foundry_hosting/tests/test_responses.py +++ b/python/packages/foundry_hosting/tests/test_responses.py @@ -8528,39 +8528,59 @@ 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 - async def test_full_public_metadata_content_filter_stays_incomplete(self) -> None: - """Sixteen caller entries are valid and must not turn a truncated turn into a failure.""" + @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) - server = _make_server(self._filtered_agent(finish_reason="content_filter")) + 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=False, metadata=metadata, store=True) + resp = await _post(server, input_text="hello", stream=stream, metadata=metadata, store=True) assert resp.status_code == 200 - body = resp.json() + 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": "content_filter"} + 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"] + 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": "content_filter"} + assert stored_body["incomplete_details"] == {"reason": expected_reason} + assert stored_body.get("error") is None assert stored_body.get("metadata") == metadata - - async def test_full_public_metadata_streaming_length_stays_incomplete(self) -> None: - metadata = _public_metadata(16) - server = _make_server(self._filtered_agent(finish_reason="length")) - - resp = await _post(server, input_text="hello", stream=True, metadata=metadata) - assert resp.status_code == 200 - events = _parse_sse_events(resp.text) - assert _sse_event_types(events)[-1] == "response.incomplete" - assert "response.failed" not in _sse_event_types(events) - incomplete = events[-1]["data"]["response"] - assert incomplete["status"] == "incomplete" - assert incomplete["incomplete_details"] == {"reason": "max_output_tokens"} - assert incomplete.get("metadata") == metadata + 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( From 2e81d2523bbec374132a03e3a9c14d797b79aa92 Mon Sep 17 00:00:00 2001 From: its-amann Date: Thu, 29 Oct 2026 20:59:14 +0530 Subject: [PATCH 4/4] fix(foundry-hosting): persist recovery state outside public metadata --- .../_responses.py | 76 ++++---- .../_workflow_responses.py | 4 +- .../foundry_hosting/tests/test_responses.py | 169 ++++++++++++++++-- 3 files changed, 195 insertions(+), 54 deletions(-) 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 eecf460f0b..3524206225 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py @@ -377,30 +377,25 @@ def _create_response_event_stream(context: ResponseContext, request: CreateRespo ) -def _write_hosted_internal(stream: ResponseEventStream, key: str, value: Any) -> None: - """Persist one host-owned recovery value without displacing public metadata. - - The SDK stores this bag in the public metadata map and rejects a 17th entry. A full map of - caller metadata must still keep every public entry. The incomplete reason then lives on - ``incomplete_details``, which is outside that quota, so a truncated turn can still finish - ``incomplete``. Other host cursors are skipped rather than failing the response; the public - map is left as the caller sent it. +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. """ - try: - stream.internal_metadata[key] = value - except ValueError: - if key == _INCOMPLETE_REASON_KEY and isinstance(value, str): - stream.response["incomplete_details"] = {"reason": value} - return - logger.debug( - "Skipping host metadata %s because the response metadata map cannot store another entry.", - key, - ) + 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 = stream.internal_metadata.get(_INCOMPLETE_REASON_KEY) + persisted = _hosted_internal_metadata(stream).get(_INCOMPLETE_REASON_KEY) if isinstance(persisted, str): return persisted details = stream.response.get("incomplete_details") @@ -1812,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( @@ -1836,8 +1833,9 @@ async def _handle_inner_agent( continue for event in tracker.close(): yield event - _write_hosted_internal(response_event_stream, _HOSTED_PROVIDER_OUTPUT_COUNT_KEY, output_count) - _write_hosted_internal(response_event_stream, _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: @@ -2213,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( @@ -2329,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() - _write_hosted_internal(response_event_stream, _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( @@ -2462,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) @@ -2502,7 +2500,7 @@ def __init__( 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) @@ -2599,7 +2597,7 @@ def record_finish_reason(self, finish_reason: str | None) -> None: self._incomplete_reason = ResponseIncompleteReason.MAX_OUTPUT_TOKENS else: return - _write_hosted_internal(self._stream, _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.""" @@ -3185,20 +3183,16 @@ 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 - _write_hosted_internal( - self._stream, - _PENDING_CONTAINER_FILE_CITATIONS_KEY, - [ - { - "container_id": citation.container_id, - "file_id": citation.file_id, - "filename": citation.filename, - } - for citation in self._pending_container_file_citations.values() - ], - ) + _hosted_internal_metadata(self._stream)[_PENDING_CONTAINER_FILE_CITATIONS_KEY] = [ + { + "container_id": citation.container_id, + "file_id": citation.file_id, + "filename": citation.filename, + } + for citation in self._pending_container_file_citations.values() + ] # endregion 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 cfca38eddf..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,11 +53,11 @@ 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] - _write_hosted_internal, # pyright: ignore[reportPrivateUsage] ) from ._scope import FoundryRequestScope from ._state_store import CheckpointStoreProvider, ContextScopedStoreProvider @@ -509,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()) - _write_hosted_internal(draft, _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 13e3a91bb5..2ba80fd457 100644 --- a/python/packages/foundry_hosting/tests/test_responses.py +++ b/python/packages/foundry_hosting/tests/test_responses.py @@ -98,6 +98,7 @@ _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] @@ -107,7 +108,6 @@ _OutputItemTracker, # pyright: ignore[reportPrivateUsage] _SignalledIterator, # pyright: ignore[reportPrivateUsage] _stringify_mcp_output, # pyright: ignore[reportPrivateUsage] - _write_hosted_internal, # pyright: ignore[reportPrivateUsage] consent_url_from_error, ) from agent_framework_foundry_hosting._state_store import ( @@ -1187,7 +1187,10 @@ def test_fresh_request_strips_forged_internal_metadata() -> None: "_foundry_provider_output_count": 4, }) request_metadata = {"source": "client", "_internal_metadata": forged} - request = CreateResponse(input="Hi", metadata=cast(Metadata, request_metadata)) + 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) @@ -1209,7 +1212,8 @@ def test_fresh_request_with_only_reserved_metadata_keeps_public_metadata_empty() assert _OutputItemTracker(stream).incomplete_reason is None -def test_recovery_keeps_trusted_internal_metadata_when_the_request_forges_it() -> 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", @@ -1218,6 +1222,11 @@ def test_recovery_keeps_trusted_internal_metadata_when_the_request_forges_it() - 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) @@ -1226,7 +1235,21 @@ def test_recovery_keeps_trusted_internal_metadata_when_the_request_forges_it() - 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.CONTENT_FILTER + 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: @@ -1256,7 +1279,7 @@ def test_reserved_metadata_key_does_not_consume_the_host_incomplete_reason_slot( tracker.record_finish_reason("content_filter") assert tracker.incomplete_reason == ResponseIncompleteReason.CONTENT_FILTER - stored = json.loads(stream.response["metadata"]["_internal_metadata"]) + 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()) @@ -1270,13 +1293,13 @@ def test_full_public_metadata_records_incomplete_reason_outside_the_quota() -> N tracker.record_finish_reason("length") tracker.record_finish_reason("content_filter") - _write_hosted_internal(stream, _LATEST_CHECKPOINT_ID_KEY, "ckpt_skipped") + _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.internal_metadata.get(_LATEST_CHECKPOINT_ID_KEY) is None - assert stream.response["incomplete_details"] == {"reason": "content_filter"} + 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") @@ -2955,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: @@ -8517,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) @@ -8567,6 +8712,7 @@ async def test_sixteen_public_metadata_keys_keep_truncated_output( 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 @@ -8578,6 +8724,7 @@ async def test_sixteen_public_metadata_keys_keep_truncated_output( 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 @@ -9805,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"] @@ -9820,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"