[AgentServer] Platform headers, persistence resilience, x-request-id, logging fix (#46429)

* [AgentServer] Platform headers, x-request-id middleware, logging fixes

- Added _platform_headers module centralizing all platform HTTP header
  name constants (x-request-id, x-platform-server, x-agent-session-id,
  isolation keys, traceparent, x-ms-client-request-id).

- Added RequestIdMiddleware (core) that sets x-request-id response header
  on every HTTP response. Value resolved: OTEL trace ID > incoming header > UUID.

- Error responses (4xx/5xx) with JSON error body are automatically enriched
  with error.additionalInfo.request_id matching the x-request-id header.

- Foundry storage logging now includes traceparent in all log messages.

- Fixed FoundryStorageLoggingPolicy crash on transport-level failures
  (DNS, connection refused) when no response is available. Transport
  failures logged at ERROR level; original exception propagates cleanly.

- Removed x-ms-request-id from Foundry storage response logging.

- Migrated all header string literals to _platform_headers constants.

* feat(responses): persistence failure resilience — buffer-then-persist-then-yield

Implement persistence failure resilience per spec. Terminal SSE events
are now buffered, persistence attempted, and on failure the terminal
is replaced with response.failed carrying error_code='storage_error'.

Key changes:
- _process_handler_events: buffer terminal events to state.pending_terminal
- _persist_and_resolve_terminal: new method — attempts persistence, replaces
  terminal on failure with storage_error response.failed
- _apply_storage_error_replacement: helper for storage error substitution
- _finalize_stream: stripped of persistence (moved to resolve_terminal),
  eviction guarded by not persistence_failed
- run_sync: persistence failure → _HandlerError → HTTP 500
- _run_background_non_stream: Phase 2 failure sets persistence_failed
- _register_bg_execution: Phase 1 failure → standalone error event (§3.3)
- ResponseExecution: added persistence_failed/persistence_exception fields
- build_failed_response: added error_code parameter

Tests: 8 new contract tests in test_persistence_failure.py covering all
modes (streaming, sync, bg+stream, bg+non-stream) and failure scenarios.
All 942 tests pass.

* fix: address PR review feedback — hardening and test improvements

- _request_id.py: guard scope['state'] with isinstance(MutableMapping),
  filter existing x-request-id header to prevent duplicates
- _validation.py: fix docstrings to use camelCase 'additionalInfo',
  guard _enrich_error_payload against non-dict additionalInfo values
- _endpoint_handler.py: wire _get_scope_request_id into handle_create
  error responses for consistent error body enrichment
- _orchestrator.py: replace terminal event in-place (preserve seq num),
  evict runtime record on sync persistence failure to avoid memory leak,
  evict on bg+stream Phase 1 failure for §3.3 compliance,
  defer subject.publish for terminal events until persistence resolves
- test_request_id_middleware.py: remove unused PlainTextResponse import,
  fix _error_plain docstring, add id1 != id2 assertion,
  monkeypatch _get_trace_id in fallback test for deterministic assertion
- test_foundry_logging_policy.py: rename test to match actual behavior
- test_persistence_failure.py: add error payload assertions (code check)

* fix: resolve pylint errors — add missing docstring types and disable too-many-branches

Author: RaviPidaparthi

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

@@ -4,6 +4,8 @@
 
 ### Features Added
 
+- `RequestIdMiddleware` — pure-ASGI middleware that sets an `x-request-id` response header on every response. The request ID is resolved from the OpenTelemetry trace ID, an incoming `x-request-id` header, or a generated UUID (in that priority). The resolved value is stored in ASGI scope state under the well-known key `agentserver.request_id` for use by sibling protocol packages. Automatically wired into `AgentServerHost`.
+
 ### Breaking Changes
 
 ### Bugs Fixed

sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/__init__.py

@@ -27,6 +27,7 @@
 from ._config import AgentConfig
 from ._errors import create_error_response
 from ._middleware import InboundRequestLoggingMiddleware
+from ._request_id import RequestIdMiddleware
 from ._server_version import build_server_version
 from ._tracing import (
     configure_observability,
@@ -43,6 +44,7 @@
     "AgentConfig",
     "AgentServerHost",
     "InboundRequestLoggingMiddleware",
+    "RequestIdMiddleware",
     "build_server_version",
     "configure_observability",
     "create_error_response",

sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_base.py

@@ -25,6 +25,7 @@
 
 from . import _config, _tracing
 from ._middleware import InboundRequestLoggingMiddleware
+from ._request_id import RequestIdMiddleware as _RequestIdMiddleware
 from ._server_version import build_server_version
 from ._version import VERSION as _CORE_VERSION
 
@@ -278,6 +279,7 @@ async def _lifespan(_app: Starlette) -> AsyncGenerator[None, None]:  # noqa: RUF
                     _PlatformHeaderMiddleware,
                     get_server_version=self._build_server_version,
                 ),
+                Middleware(_RequestIdMiddleware),  # type: ignore[arg-type]
             ],
             **kwargs,
         )

sdk/agentserver/azure-ai-agentserver-core/azure/ai/agentserver/core/_request_id.py

@@ -0,0 +1,103 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+"""Request ID middleware for Azure AI Agent Server hosts.
+
+A pure-ASGI middleware that sets the ``x-request-id`` response header on
+every HTTP response.  The value is resolved in priority order:
+
+1. Current OTEL trace ID.
+2. Incoming ``x-request-id`` request header (client-provided correlation ID).
+3. A new UUID as fallback.
+
+The resolved value is also stored in ``scope["state"]["agentserver.request_id"]``
+for use by downstream handlers (e.g., error body enrichment in protocol
+packages).
+"""
+
+from __future__ import annotations
+
+import uuid
+from typing import Any, MutableMapping
+
+from starlette.types import ASGIApp, Receive, Scope, Send
+
+from ._middleware import _extract_header, _get_trace_id
+
+# Key used to store the resolved request ID in ASGI scope state.
+REQUEST_ID_STATE_KEY = "agentserver.request_id"
+
+
+class RequestIdMiddleware:
+    """Pure-ASGI middleware that sets ``x-request-id`` on every HTTP response.
+
+    The resolved request ID is stored in ``scope["state"]`` under
+    :data:`REQUEST_ID_STATE_KEY` so that protocol-specific code (e.g. the
+    Responses package) can read it for error body enrichment.
+
+    Unlike ``BaseHTTPMiddleware``, this passes the ``receive`` callable
+    through to the inner application untouched, preserving
+    ``request.is_disconnected()`` behaviour.
+
+    :param app: The inner ASGI application.
+    :type app: ASGIApp
+    """
+
+    def __init__(self, app: ASGIApp) -> None:
+        self.app = app
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        if scope["type"] != "http":
+            await self.app(scope, receive, send)
+            return
+
+        raw_headers: list[tuple[bytes, bytes]] = scope.get("headers", [])
+        request_id = _resolve_request_id(raw_headers)
+
+        # Store in scope state for downstream access.
+        state = scope.get("state")
+        if not isinstance(state, MutableMapping):
+            state = {}
+            scope["state"] = state
+        state[REQUEST_ID_STATE_KEY] = request_id
+
+        async def _send_with_request_id(message: MutableMapping[str, Any]) -> None:
+            if message["type"] == "http.response.start":
+                # Filter any existing x-request-id to avoid duplicates, then add ours.
+                headers = [
+                    (name, value)
+                    for name, value in message.get("headers", [])
+                    if name.lower() != b"x-request-id"
+                ]
+                headers.append((b"x-request-id", request_id.encode()))
+                message = {**message, "headers": headers}
+            await send(message)
+
+        await self.app(scope, receive, _send_with_request_id)
+
+
+def _resolve_request_id(raw_headers: list[tuple[bytes, bytes]]) -> str:
+    """Resolve the request ID from available sources.
+
+    Priority:
+    1. OTEL trace ID from current activity.
+    2. Incoming ``x-request-id`` request header.
+    3. New UUID fallback.
+
+    :param raw_headers: Raw ASGI header tuples.
+    :type raw_headers: list[tuple[bytes, bytes]]
+    :return: The resolved request ID string.
+    :rtype: str
+    """
+    # Priority 1: OTEL trace ID
+    trace_id = _get_trace_id(raw_headers)
+    if trace_id and trace_id != "0" * 32:
+        return trace_id
+
+    # Priority 2: Incoming x-request-id header
+    incoming = _extract_header(raw_headers, b"x-request-id")
+    if incoming:
+        return incoming
+
+    # Priority 3: New UUID
+    return uuid.uuid4().hex

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

@@ -4,12 +4,16 @@
 
 ### Features Added
 
+- All HTTP responses now include an `x-request-id` header for request correlation, inherited from `RequestIdMiddleware` in `azure-ai-agentserver-core>=2.0.0b3`. The value is resolved from the OpenTelemetry trace ID, an incoming `x-request-id` header, or a generated UUID.
+
 ### Breaking Changes
 
 ### Bugs Fixed
 
 ### Other Changes
 
+- Bumped minimum `azure-ai-agentserver-core` dependency to `>=2.0.0b3`.
+
 ## 1.0.0b2 (2026-04-17)
 
 ### Features Added

sdk/agentserver/azure-ai-agentserver-invocations/pyproject.toml

@@ -21,7 +21,7 @@ classifiers = [
 keywords = ["azure", "azure sdk", "agent", "agentserver", "invocations"]
 
 dependencies = [
-    "azure-ai-agentserver-core>=2.0.0b2",
+    "azure-ai-agentserver-core>=2.0.0b3",
 ]
 
 [dependency-groups]

sdk/agentserver/azure-ai-agentserver-invocations/tests/test_request_id.py

@@ -0,0 +1,82 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+"""Tests for x-request-id header on invocations responses.
+
+The ``RequestIdMiddleware`` is wired in ``AgentServerHost`` (core) and
+inherited by ``InvocationAgentServerHost``.  These tests verify the
+header appears on invocations endpoints.
+
+Note: Error body enrichment (``additionalInfo.request_id``) is a
+Responses-only feature and is NOT applied to invocations error responses.
+"""
+import uuid
+
+import pytest
+
+
+# ---------------------------------------------------------------------------
+# Header presence — success responses
+# ---------------------------------------------------------------------------
+
+@pytest.mark.asyncio
+async def test_invoke_returns_request_id_header(echo_client):
+    """POST /invocations success response includes x-request-id."""
+    resp = await echo_client.post("/invocations", content=b"hello")
+    assert "x-request-id" in resp.headers
+    assert resp.headers["x-request-id"]  # non-empty
+
+
+@pytest.mark.asyncio
+async def test_request_id_is_stable_per_request(echo_client):
+    """Each request gets a unique x-request-id."""
+    ids = set()
+    for _ in range(5):
+        resp = await echo_client.post("/invocations", content=b"x")
+        ids.add(resp.headers["x-request-id"])
+    assert len(ids) == 5
+
+
+@pytest.mark.asyncio
+async def test_request_id_echoes_incoming_header(echo_client):
+    """When the client sends x-request-id, the same value is returned."""
+    custom_id = uuid.uuid4().hex
+    resp = await echo_client.post(
+        "/invocations",
+        content=b"test",
+        headers={"x-request-id": custom_id},
+    )
+    assert resp.headers["x-request-id"] == custom_id
+
+
+@pytest.mark.asyncio
+async def test_readiness_returns_request_id(echo_client):
+    """GET /readiness also gets x-request-id (middleware applies to all routes)."""
+    resp = await echo_client.get("/readiness")
+    assert resp.status_code == 200
+    assert "x-request-id" in resp.headers
+
+
+# ---------------------------------------------------------------------------
+# Error responses — header present, but NO body enrichment
+# ---------------------------------------------------------------------------
+
+@pytest.mark.asyncio
+async def test_error_response_has_request_id_header(failing_client):
+    """500 error response includes x-request-id header."""
+    resp = await failing_client.post("/invocations", content=b"boom")
+    assert resp.status_code == 500
+    assert "x-request-id" in resp.headers
+    assert resp.headers["x-request-id"]
+
+
+@pytest.mark.asyncio
+async def test_error_response_no_additional_info_enrichment(failing_client):
+    """Invocations error bodies do NOT get additionalInfo.request_id (Responses-only)."""
+    resp = await failing_client.post("/invocations", content=b"boom")
+    assert resp.status_code == 500
+
+    body = resp.json()
+    error = body.get("error", {})
+    # core's create_error_response doesn't include additionalInfo
+    assert "additionalInfo" not in error

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

@@ -4,12 +4,22 @@
 
 ### Features Added
 
+- Added `_platform_headers` module centralizing all platform HTTP header name constants (`x-request-id`, `x-platform-server`, `x-agent-session-id`, isolation keys, `traceparent`, `x-ms-client-request-id`). All header references now use shared constants instead of scattered string literals.
+- Added `RequestIdMiddleware` (in `azure-ai-agentserver-core`) that sets the `x-request-id` response header on every HTTP response. Value is resolved in priority order: OTEL trace ID → incoming `x-request-id` header → new UUID.
+- Error responses (4xx/5xx) with a JSON `error` body are automatically enriched with `error.additionalInfo.request_id` matching the `x-request-id` response header, enabling client-side error correlation.
+- Foundry storage logging now includes the `traceparent` header (W3C distributed trace ID) in all log messages, enabling correlation between SDK log entries and backend distributed traces.
+
 ### Breaking Changes
 
 ### Bugs Fixed
 
+- Fixed crash in `FoundryStorageLoggingPolicy` when a transport-level failure (DNS resolution, connection refused, timeout) occurs before any HTTP response is received. The policy previously attempted to access `response.headers` unconditionally, raising an unrelated exception that masked the real transport error. Transport failures are now logged at ERROR level and the original exception propagates cleanly.
+
 ### Other Changes
 
+- Removed `x-ms-request-id` from Foundry storage response logging (unused service header).
+- Migrated all header string literals to use `_platform_headers` constants.
+
 ## 1.0.0b4 (2026-04-19)
 
 ### Bugs Fixed

sdk/agentserver/azure-ai-agentserver-responses/azure/ai/agentserver/responses/_platform_headers.py

@@ -0,0 +1,95 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+"""Platform HTTP header name constants used across the AgentServer packages.
+
+These headers form the wire contract between the Foundry platform, agent
+containers, and downstream storage services.  All header name constants
+are defined here to eliminate scattered string literals and ensure
+consistency across logging policies, middleware, and endpoint handlers.
+
+**Response headers** (set by the server on every response):
+
+- :data:`REQUEST_ID` — request correlation ID.
+- :data:`SERVER_VERSION` — server SDK identity.
+- :data:`SESSION_ID` — resolved session ID (when applicable).
+
+**Request headers** (set by the platform or client):
+
+- :data:`REQUEST_ID` — client-provided correlation ID (echoed back on the response).
+- :data:`USER_ISOLATION_KEY` / :data:`CHAT_ISOLATION_KEY` — platform isolation keys.
+- :data:`CLIENT_HEADER_PREFIX` — prefix for pass-through client headers.
+- :data:`TRACEPARENT` — W3C Trace Context propagation header.
+- :data:`CLIENT_REQUEST_ID` — Azure SDK client correlation header.
+"""
+
+from __future__ import annotations
+
+# -- Response/request correlation -------------------------------------------
+
+REQUEST_ID: str = "x-request-id"
+"""The ``x-request-id`` header — carries the request correlation ID.
+
+On responses, the server always sets this header
+(OTEL trace ID → incoming header → UUID).
+On requests, clients may set it to provide their own correlation ID.
+"""
+
+SERVER_VERSION: str = "x-platform-server"
+"""The ``x-platform-server`` header — identifies the server SDK stack
+(hosting version, protocol versions, language, and runtime).
+Set on every response by ``_PlatformHeaderMiddleware``.
+"""
+
+SESSION_ID: str = "x-agent-session-id"
+"""The ``x-agent-session-id`` header — the resolved session ID for the request.
+Set on responses by protocol-specific session resolution logic.
+"""
+
+# -- Platform isolation -----------------------------------------------------
+
+USER_ISOLATION_KEY: str = "x-agent-user-isolation-key"
+"""The ``x-agent-user-isolation-key`` header — the platform-injected
+partition key for user-private state.
+"""
+
+CHAT_ISOLATION_KEY: str = "x-agent-chat-isolation-key"
+"""The ``x-agent-chat-isolation-key`` header — the platform-injected
+partition key for conversation-scoped state.
+"""
+
+# -- Client pass-through ---------------------------------------------------
+
+CLIENT_HEADER_PREFIX: str = "x-client-"
+"""The prefix ``x-client-`` for pass-through client headers.
+
+All request headers starting with this prefix are extracted and forwarded
+to the handler via the invocation context.
+"""
+
+# -- Tracing & diagnostics -------------------------------------------------
+
+TRACEPARENT: str = "traceparent"
+"""The ``traceparent`` header — W3C Trace Context propagation header.
+Used for distributed tracing correlation on outbound storage requests.
+"""
+
+CLIENT_REQUEST_ID: str = "x-ms-client-request-id"
+"""The ``x-ms-client-request-id`` header — Azure SDK client correlation header.
+Logged for diagnostic correlation with upstream Azure SDK callers.
+"""
+
+# -- Storage diagnostics (response headers from Foundry) --------------------
+
+APIM_REQUEST_ID: str = "apim-request-id"
+"""The ``apim-request-id`` header — APIM gateway correlation header.
+Extracted from Foundry storage responses for diagnostic logging.
+"""
+
+# -- HttpContext item key ---------------------------------------------------
+
+REQUEST_ID_ITEM_KEY: str = "agentserver.request_id"
+"""Key used to store the resolved request ID in ASGI scope state.
+
+Downstream handlers and middleware can read this value to correlate the
+request ID without re-resolving it.
+"""