[Agent Server] Persist identity fields in checkpoints

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,16 +379,48 @@ 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.
+    # 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:
-        snapshot_bytes = json.dumps(dict(response), sort_keys=True, default=str).encode("utf-8")
+        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)
         return last_snapshot
     if snapshot_bytes == last_snapshot:
         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(
+            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)
@@ -1039,6 +1083,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 +2390,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(

sdk/agentserver/azure-ai-agentserver-responses/tests/unit/test_checkpoint.py

@@ -169,6 +169,70 @@ 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_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