From de4f56ff299501a34beff10afe34520632554f5b Mon Sep 17 00:00:00 2001 From: "Slai.life" Date: Fri, 30 Oct 2026 06:51:24 +0800 Subject: [PATCH 1/3] Fix checkpoint response identity persistence --- .../responses/hosting/_orchestrator.py | 51 ++++++++++++++++++- .../tests/unit/test_checkpoint.py | 38 ++++++++++++++ 2 files changed, 87 insertions(+), 2 deletions(-) diff --git a/sdk/agentserver/azure-ai-agentserver-responses/azure/ai/agentserver/responses/hosting/_orchestrator.py b/sdk/agentserver/azure-ai-agentserver-responses/azure/ai/agentserver/responses/hosting/_orchestrator.py index e3e4468ffb18..338ed722354b 100644 --- a/sdk/agentserver/azure-ai-agentserver-responses/azure/ai/agentserver/responses/hosting/_orchestrator.py +++ b/sdk/agentserver/azure-ai-agentserver-responses/azure/ai/agentserver/responses/hosting/_orchestrator.py @@ -330,6 +330,10 @@ async def _do_checkpoint_persist( response_id: str, last_snapshot: "bytes | None", terminal_seen: bool, + agent_reference: "generated_models.AgentReference | dict[str, Any] | None" = None, + model: str | None = None, + agent_session_id: str | None = None, + conversation_id: str | None = None, ) -> "bytes | None": """Persist a developer checkpoint snapshot (spec 025 §A.3). @@ -355,6 +359,14 @@ async def _do_checkpoint_persist( :paramtype last_snapshot: bytes | None :keyword terminal_seen: Whether a terminal event has already been processed. :paramtype terminal_seen: bool + :keyword agent_reference: Canonical agent identity to stamp on the checkpoint. + :paramtype agent_reference: AgentReference | dict[str, Any] | None + :keyword model: Model name to stamp when the checkpoint has none. + :paramtype model: str | None + :keyword agent_session_id: Resolved session ID to stamp on the checkpoint. + :paramtype agent_session_id: str | None + :keyword conversation_id: Resolved conversation ID to stamp on the checkpoint. + :paramtype conversation_id: str | None :returns: The new ``last_snapshot`` bytes (unchanged when nothing persisted). :rtype: bytes | None """ @@ -367,8 +379,35 @@ async def _do_checkpoint_persist( response = event.response if response is None or provider is None: return last_snapshot + + # Checkpoints may be emitted by a stream created with only ``response_id``. + # Normalize the payload with the same identity defaults used by terminal + # persistence before handing it to the provider, so a checkpoint is valid + # on its own rather than relying on a later terminal write to fill fields. + checkpoint_response = deepcopy(dict(response)) + normalized_events: list[generated_models.ResponseStreamEvent] = [ + cast( + "generated_models.ResponseStreamEvent", + { + "type": "response.in_progress", + "response": checkpoint_response, + }, + ) + ] + normalized_snapshot = _extract_response_snapshot_from_events( + normalized_events, + response_id=response_id, + agent_reference=agent_reference or {}, + model=model, + agent_session_id=agent_session_id, + conversation_id=conversation_id, + ) + checkpoint_response = cast("dict[str, Any]", normalized_snapshot) + # Keep checkpoint snapshots independently readable by providers that use + # the persisted mode flags for background-response validation. + checkpoint_response["background"] = background try: - snapshot_bytes = json.dumps(dict(response), sort_keys=True, default=str).encode("utf-8") + snapshot_bytes = json.dumps(checkpoint_response, sort_keys=True, default=str).encode("utf-8") except Exception: # pylint: disable=broad-exception-caught logger.debug("checkpoint() snapshot serialisation failed for %s", response_id, exc_info=True) return last_snapshot @@ -376,7 +415,7 @@ async def _do_checkpoint_persist( return last_snapshot # idempotent — nothing changed since the last checkpoint result = last_snapshot try: - await provider.update_response(response, context=context) + await provider.update_response(checkpoint_response, context=context) result = snapshot_bytes except Exception as exc: # pylint: disable=broad-exception-caught setattr(exc, PLATFORM_ERROR_TAG, True) @@ -1039,6 +1078,10 @@ async def _bg_drain_handler_events( response_id=response_id, last_snapshot=st.checkpoint_snapshot, terminal_seen=st.terminal_seen, + agent_reference=agent_reference, + model=model, + agent_session_id=agent_session_id, + conversation_id=conversation_id, ) continue # Client-initiated cancel → discard and force cancelled. @@ -2342,6 +2385,10 @@ async def _persist_checkpoint( response_id=ctx.response_id, last_snapshot=state.last_persisted_snapshot, terminal_seen=state.pending_terminal is not None, + agent_reference=ctx.agent_reference, + model=ctx.model, + agent_session_id=ctx.agent_session_id, + conversation_id=ctx.conversation_id, ) async def _emit_standalone_error( diff --git a/sdk/agentserver/azure-ai-agentserver-responses/tests/unit/test_checkpoint.py b/sdk/agentserver/azure-ai-agentserver-responses/tests/unit/test_checkpoint.py index bf4087ad24ee..5da0e7150d3e 100644 --- a/sdk/agentserver/azure-ai-agentserver-responses/tests/unit/test_checkpoint.py +++ b/sdk/agentserver/azure-ai-agentserver-responses/tests/unit/test_checkpoint.py @@ -169,6 +169,44 @@ async def test_t21_status_as_is_in_snapshot(): assert p.updates[0]["metadata"]["_internal_metadata"] == '{"cp":1}' +@pytest.mark.asyncio +async def test_checkpoint_stamps_identity_for_an_unseeded_stream(): + p = _RecordingProvider() + ev = _event(cp=1) + + await _do_checkpoint_persist( + ev, + provider=p, + runtime_options=_opts(True), + store=True, + background=True, + context=None, + response_id="r1", + last_snapshot=None, + terminal_seen=False, + agent_reference={"type": "agent_reference", "name": "payments", "version": "22"}, + model="gpt-4.1", + agent_session_id="session_456", + conversation_id="conversation_789", + ) + + assert p.updates == [ + { + "id": "r1", + "response_id": "r1", + "object": "response", + "status": "in_progress", + "output": [], + "model": "m", + "agent_reference": {"type": "agent_reference", "name": "payments", "version": "22"}, + "agent_session_id": "session_456", + "conversation": {"id": "conversation_789"}, + "background": True, + "metadata": {"_internal_metadata": '{"cp":1}'}, + } + ] + + @pytest.mark.asyncio async def test_t22_failure_swallowed_and_tagged(): from azure.ai.agentserver.core._platform_headers import PLATFORM_ERROR_TAG # noqa: E501 From e7809a09745ea81fa752eb4cbcf273e7ccdd9c43 Mon Sep 17 00:00:00 2001 From: "Slai.life" Date: Fri, 30 Oct 2026 07:09:52 +0800 Subject: [PATCH 2/3] Handle checkpoint snapshot preparation failures --- .../responses/hosting/_orchestrator.py | 47 ++++++++++--------- .../tests/unit/test_checkpoint.py | 26 ++++++++++ 2 files changed, 51 insertions(+), 22 deletions(-) diff --git a/sdk/agentserver/azure-ai-agentserver-responses/azure/ai/agentserver/responses/hosting/_orchestrator.py b/sdk/agentserver/azure-ai-agentserver-responses/azure/ai/agentserver/responses/hosting/_orchestrator.py index 338ed722354b..4d47054a042b 100644 --- a/sdk/agentserver/azure-ai-agentserver-responses/azure/ai/agentserver/responses/hosting/_orchestrator.py +++ b/sdk/agentserver/azure-ai-agentserver-responses/azure/ai/agentserver/responses/hosting/_orchestrator.py @@ -384,29 +384,32 @@ async def _do_checkpoint_persist( # Normalize the payload with the same identity defaults used by terminal # persistence before handing it to the provider, so a checkpoint is valid # on its own rather than relying on a later terminal write to fill fields. - checkpoint_response = deepcopy(dict(response)) - normalized_events: list[generated_models.ResponseStreamEvent] = [ - cast( - "generated_models.ResponseStreamEvent", - { - "type": "response.in_progress", - "response": checkpoint_response, - }, - ) - ] - normalized_snapshot = _extract_response_snapshot_from_events( - normalized_events, - response_id=response_id, - agent_reference=agent_reference or {}, - model=model, - agent_session_id=agent_session_id, - conversation_id=conversation_id, - ) - checkpoint_response = cast("dict[str, Any]", normalized_snapshot) - # Keep checkpoint snapshots independently readable by providers that use - # the persisted mode flags for background-response validation. - checkpoint_response["background"] = background + # Preparation is part of snapshot serialisation: extensions can contain + # values such as generators that cannot be deep-copied, and checkpoint + # failures must never fail the live response. try: + checkpoint_response = deepcopy(dict(response)) + normalized_events: list[generated_models.ResponseStreamEvent] = [ + cast( + "generated_models.ResponseStreamEvent", + { + "type": "response.in_progress", + "response": checkpoint_response, + }, + ) + ] + normalized_snapshot = _extract_response_snapshot_from_events( + normalized_events, + response_id=response_id, + agent_reference=agent_reference or {}, + model=model, + agent_session_id=agent_session_id, + conversation_id=conversation_id, + ) + checkpoint_response = cast("dict[str, Any]", normalized_snapshot) + # Keep checkpoint snapshots independently readable by providers that use + # the persisted mode flags for background-response validation. + checkpoint_response["background"] = background snapshot_bytes = json.dumps(checkpoint_response, sort_keys=True, default=str).encode("utf-8") except Exception: # pylint: disable=broad-exception-caught logger.debug("checkpoint() snapshot serialisation failed for %s", response_id, exc_info=True) diff --git a/sdk/agentserver/azure-ai-agentserver-responses/tests/unit/test_checkpoint.py b/sdk/agentserver/azure-ai-agentserver-responses/tests/unit/test_checkpoint.py index 5da0e7150d3e..7f319da00bd8 100644 --- a/sdk/agentserver/azure-ai-agentserver-responses/tests/unit/test_checkpoint.py +++ b/sdk/agentserver/azure-ai-agentserver-responses/tests/unit/test_checkpoint.py @@ -207,6 +207,32 @@ async def test_checkpoint_stamps_identity_for_an_unseeded_stream(): ] +@pytest.mark.asyncio +async def test_checkpoint_preparation_failure_is_swallowed(): + p = _RecordingProvider() + ev = _event() + generator = (value for value in [1]) + ev.response["extension"] = {"generator": generator} + + try: + snap = await _do_checkpoint_persist( + ev, + provider=p, + runtime_options=_opts(True), + store=True, + background=True, + context=None, + response_id="r1", + last_snapshot=b"previous", + terminal_seen=False, + ) + finally: + generator.close() + + assert snap == b"previous" + assert p.updates == [] + + @pytest.mark.asyncio async def test_t22_failure_swallowed_and_tagged(): from azure.ai.agentserver.core._platform_headers import PLATFORM_ERROR_TAG # noqa: E501 From 660a5759d1014ef22ed650000a355604ce696380 Mon Sep 17 00:00:00 2001 From: "Slai.life" Date: Fri, 30 Oct 2026 08:43:00 +0800 Subject: [PATCH 3/3] fix: type checkpoint response persistence --- .../azure/ai/agentserver/responses/hosting/_orchestrator.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/sdk/agentserver/azure-ai-agentserver-responses/azure/ai/agentserver/responses/hosting/_orchestrator.py b/sdk/agentserver/azure-ai-agentserver-responses/azure/ai/agentserver/responses/hosting/_orchestrator.py index 4d47054a042b..2c6e52fd630f 100644 --- a/sdk/agentserver/azure-ai-agentserver-responses/azure/ai/agentserver/responses/hosting/_orchestrator.py +++ b/sdk/agentserver/azure-ai-agentserver-responses/azure/ai/agentserver/responses/hosting/_orchestrator.py @@ -418,7 +418,9 @@ async def _do_checkpoint_persist( return last_snapshot # idempotent — nothing changed since the last checkpoint result = last_snapshot try: - await provider.update_response(checkpoint_response, context=context) + await provider.update_response( + cast("generated_models.ResponseObject", checkpoint_response), context=context + ) result = snapshot_bytes except Exception as exc: # pylint: disable=broad-exception-caught setattr(exc, PLATFORM_ERROR_TAG, True)