Restore durable Responses request context

sdk/agentserver/azure-ai-agentserver-responses/CHANGELOG.md

@@ -22,6 +22,7 @@
 ### Bugs Fixed
 
 - Restored JSON-string encoding for response-level `internal_metadata` so resilient response checkpoints round-trip through Foundry storage.
+- Restored `get_request_context()` identity values while stored Responses handlers run inside durable tasks.
 
 ### Other Changes
 

sdk/agentserver/azure-ai-agentserver-responses/azure/ai/agentserver/responses/hosting/_orchestrator.py

@@ -18,6 +18,11 @@
 
 import anyio
 
+from azure.ai.agentserver.core import (
+    FoundryAgentRequestContext,
+    reset_request_context,
+    set_request_context,
+)
 from azure.ai.agentserver.core.platform_headers import (
     PLATFORM_ERROR_TAG,
 )
@@ -79,6 +84,32 @@
 )
 
 
+async def _iter_handler_with_request_context(
+    create_fn: "Callable[..., AsyncIterator[generated_models.ResponseStreamEvent]]",
+    parsed: "CreateResponse",
+    context: "ResponseContext | None",
+    cancellation_signal: asyncio.Event,
+    agent_session_id: str | None,
+) -> AsyncIterator[generated_models.ResponseStreamEvent]:
+    """Run a developer handler with its reconstructed ambient platform identity."""
+    token = None
+    if context is not None:
+        platform_context = context.platform_context
+        token = set_request_context(
+            FoundryAgentRequestContext(
+                user_id=platform_context.user_id_key or None,
+                call_id=platform_context.call_id or None,
+                session_id=agent_session_id,
+            )
+        )
+    try:
+        async for event in create_fn(parsed, context, cancellation_signal):
+            yield event
+    finally:
+        if token is not None:
+            reset_request_context(token)
+
+
 async def _resolve_input_items_for_persistence(
     context: "ResponseContext | None",
     fallback_items: list[generated_models.OutputItem] | None,
@@ -942,9 +973,14 @@ async def _bg_drain_handler_events(
     :rtype: bool
     """
     try:
-        async for handler_event in _iter_with_winddown(
-            create_fn(parsed, context, cancellation_signal), cancellation_signal
-        ):
+        handler_iterator = _iter_handler_with_request_context(
+            create_fn,
+            parsed,
+            context,
+            cancellation_signal,
+            agent_session_id,
+        )
+        async for handler_event in _iter_with_winddown(handler_iterator, cancellation_signal):
             # Intercept developer ``stream.checkpoint()`` events (spec 025 §A.3):
             # persist (resilient background only) and never forward them.
             if isinstance(handler_event, ResponseCheckpointEvent):
@@ -3831,14 +3867,19 @@ async def _run_resilient_stream_body(
                 exc_info=True,
             )
             state.next_seq = 0
-        handler_iterator = self._create_fn(parsed, context, cancellation_signal)
-
         # Drive the streaming pipeline. Events flow to the per-response
         # stream — the wire iterator on _live_stream's side consumes from
         # the same registry stream independently, and the file-backed
         # backing (when configured) persists every emit to disk for the
         # GET reconnect endpoint.
         try:
+            handler_iterator = _iter_handler_with_request_context(
+                self._create_fn,
+                parsed,
+                context,
+                cancellation_signal,
+                agent_session_id,
+            )
             async for _event in self._process_handler_events(ctx, state, handler_iterator):
                 # Events are emitted to record.subject inside
                 # _process_handler_events; we only need to drain the

sdk/agentserver/azure-ai-agentserver-responses/tests/contract/test_streaming_request_context.py

@@ -16,17 +16,24 @@
 
 from __future__ import annotations
 
+from pathlib import Path
 from typing import Any
 
 from starlette.testclient import TestClient
 
 from azure.ai.agentserver.core import get_request_context
-from azure.ai.agentserver.responses import ResponsesAgentServerHost
+from azure.ai.agentserver.responses import ResponsesAgentServerHost, ResponsesServerOptions
+from azure.ai.agentserver.responses.store._file import FileResponseStore
 from azure.ai.agentserver.responses.store._memory import InMemoryResponseProvider
 from azure.ai.agentserver.responses.streaming import ResponseEventStream
 
 
-def _build_capturing_client(captured: dict[str, Any]) -> TestClient:
+def _build_capturing_client(
+    captured: dict[str, Any],
+    *,
+    resilient: bool = False,
+    storage_dir: Path | None = None,
+) -> TestClient:
     async def _capturing_handler(request: Any, context: Any, cancellation_signal: Any) -> Any:
         # An async-generator handler: its body does not execute until Starlette
         # iterates the streaming response, AFTER handle_create returned and reset
@@ -41,7 +48,10 @@ async def _capturing_handler(request: Any, context: Any, cancellation_signal: An
         yield stream.emit_created()
         yield stream.emit_completed()
 
-    app = ResponsesAgentServerHost(store=InMemoryResponseProvider())
+    app = ResponsesAgentServerHost(
+        options=ResponsesServerOptions(resilient_background=resilient),
+        store=FileResponseStore(storage_dir) if storage_dir is not None else InMemoryResponseProvider(),
+    )
     app.response_handler(_capturing_handler)
     return TestClient(app)
 
@@ -88,3 +98,54 @@ def test_streaming_handler_without_headers_sees_empty_context(self) -> None:
         assert "call_id" in captured
         assert captured.get("call_id") is None
         assert captured.get("user_id") is None
+
+    def test_stored_sync_handler_sees_call_id_and_user_id(self, tmp_path: Path) -> None:
+        captured: dict[str, Any] = {}
+        client = _build_capturing_client(captured, resilient=True, storage_dir=tmp_path / "sync")
+
+        r = client.post(
+            "/responses",
+            json={
+                "model": "m",
+                "input": "hi",
+                "stream": False,
+                "store": True,
+                "agent_session_id": "session_789",
+            },
+            headers={
+                "x-agent-user-id": "user_123",
+                "x-agent-foundry-call-id": "call_456",
+            },
+        )
+
+        assert r.status_code == 200
+        assert captured.get("call_id") == "call_456"
+        assert captured.get("user_id") == "user_123"
+        assert captured.get("session_id") == "session_789"
+
+    def test_stored_streaming_handler_sees_call_id_and_user_id(self, tmp_path: Path) -> None:
+        captured: dict[str, Any] = {}
+        client = _build_capturing_client(captured, resilient=True, storage_dir=tmp_path / "stream")
+
+        with client.stream(
+            "POST",
+            "/responses",
+            json={
+                "model": "m",
+                "input": "hi",
+                "stream": True,
+                "store": True,
+                "agent_session_id": "session_789",
+            },
+            headers={
+                "x-agent-user-id": "user_123",
+                "x-agent-foundry-call-id": "call_456",
+            },
+        ) as r:
+            assert r.status_code == 200
+            for _ in r.iter_lines():
+                pass
+
+        assert captured.get("call_id") == "call_456"
+        assert captured.get("user_id") == "user_123"
+        assert captured.get("session_id") == "session_789"

sdk/agentserver/azure-ai-agentserver-responses/tests/e2e/test_recovery_reconstruction.py

@@ -41,7 +41,8 @@ def _build_params_for_recovery() -> dict:
         disposition="re-invoke",
         agent_reference={"name": "test-agent"},
         agent_session_id="session_xyz",
-        user_id_key=None,
+        user_id_key="user_recovered",
+        call_id="call_recovered",
         client_headers={"client-trace-id": "abc"},
         query_parameters={"q": "1"},
     ).to_task_input()
@@ -75,6 +76,58 @@ def test_reconstruct_from_params_returns_record_and_context() -> None:
     assert context.response_id == "resp_recover_001"
     assert context.conversation_id == "conv_abc"
     assert context.mode_flags.store is True
+    assert context.platform_context.user_id_key == "user_recovered"
+    assert context.platform_context.call_id == "call_recovered"
+
+
+@pytest.mark.asyncio
+async def test_recovered_handler_sees_reconstructed_ambient_request_context() -> None:
+    """Cross-process reconstruction rebinds platform identity while the handler runs."""
+    from azure.ai.agentserver.core import get_request_context
+    from azure.ai.agentserver.responses._options import ResponsesServerOptions
+    from azure.ai.agentserver.responses.hosting._orchestrator import (
+        _iter_handler_with_request_context,
+    )
+    from azure.ai.agentserver.responses.hosting._resilient_orchestrator import (
+        _reconstruct_from_params,
+        _reconstruct_parsed_from_params,
+    )
+
+    params = _build_params_for_recovery()
+    record, context = _reconstruct_from_params(
+        params=params,
+        response_id="resp_recover_001",
+        provider=None,
+        runtime_state=None,
+        runtime_options=ResponsesServerOptions(),
+    )
+    parsed = _reconstruct_parsed_from_params(params)
+    captured: dict[str, str | None] = {}
+
+    async def handler(*_args):
+        request_context = get_request_context()
+        captured["user_id"] = request_context.user_id
+        captured["call_id"] = request_context.call_id
+        captured["session_id"] = request_context.session_id
+        yield {"type": "test"}
+
+    events = [
+        event
+        async for event in _iter_handler_with_request_context(
+            handler,
+            parsed,
+            context,
+            record.cancel_signal,
+            record.agent_session_id,
+        )
+    ]
+
+    assert events == [{"type": "test"}]
+    assert captured == {
+        "user_id": "user_recovered",
+        "call_id": "call_recovered",
+        "session_id": "session_xyz",
+    }
 
 
 def test_reconstruct_preserves_client_headers_and_query() -> None:  # Spec 033 FR-002b