From fd25a2c79f15029511274581079a7732cbe7d003 Mon Sep 17 00:00:00 2001
From: alhudz <al.hudz.k@gmail.com>
Date: Sat, 13 Jun 2026 23:37:06 +0530
Subject: [PATCH] sanitise the invocation id path param with a safe fallback

---
 .../CHANGELOG.md                              |  6 ++
 .../ai/agentserver/invocations/_invocation.py |  2 +-
 .../ai/agentserver/invocations/_version.py    |  2 +-
 .../tests/test_invocation_id_sanitization.py  | 60 +++++++++++++++++++
 4 files changed, 68 insertions(+), 2 deletions(-)
 create mode 100644 sdk/agentserver/azure-ai-agentserver-invocations/tests/test_invocation_id_sanitization.py

diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/CHANGELOG.md b/sdk/agentserver/azure-ai-agentserver-invocations/CHANGELOG.md
index 5d0d19f060a7..868499718d0c 100644
--- a/sdk/agentserver/azure-ai-agentserver-invocations/CHANGELOG.md
+++ b/sdk/agentserver/azure-ai-agentserver-invocations/CHANGELOG.md
@@ -1,5 +1,11 @@
 # Release History
 
+## 1.0.0b6 (Unreleased)
+
+### Bugs Fixed
+
+- `GET /invocations/{id}` and `POST /invocations/{id}/cancel` no longer reflect a malformed `invocation_id` path parameter (illegal characters or longer than `_MAX_ID_LENGTH`) into the `x-agent-invocation-id` response header and the request-scoped log and span fields. The id is now validated with a generated fallback, matching the create path.
+
 ## 1.0.0b5 (2026-06-12)
 
 ### Bugs Fixed
diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_invocation.py b/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_invocation.py
index 116ff0d62546..956f599c33e2 100644
--- a/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_invocation.py
+++ b/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_invocation.py
@@ -496,7 +496,7 @@ async def _traced_invocation_endpoint(
         dispatch: Callable[[Request], Awaitable[Response]],
     ) -> Response:
         raw_invocation_id = request.path_params["invocation_id"]
-        invocation_id = _sanitize_id(raw_invocation_id, raw_invocation_id)
+        invocation_id = _sanitize_id(raw_invocation_id, str(uuid.uuid4()))
         request.state.invocation_id = invocation_id
 
         raw_session_id = request.query_params.get("agent_session_id", "")
diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_version.py b/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_version.py
index eecd2a8e450f..ffa055f43119 100644
--- a/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_version.py
+++ b/sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_version.py
@@ -2,4 +2,4 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # ---------------------------------------------------------
 
-VERSION = "1.0.0b5"
+VERSION = "1.0.0b6"
diff --git a/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_invocation_id_sanitization.py b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_invocation_id_sanitization.py
new file mode 100644
index 000000000000..eee78ef06e8a
--- /dev/null
+++ b/sdk/agentserver/azure-ai-agentserver-invocations/tests/test_invocation_id_sanitization.py
@@ -0,0 +1,60 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+"""Tests that the ``{invocation_id}`` path param is sanitised before it is
+reflected into the ``x-agent-invocation-id`` response header."""
+import pytest
+from httpx import ASGITransport, AsyncClient
+from starlette.requests import Request
+from starlette.responses import Response
+
+from azure.ai.agentserver.invocations import InvocationAgentServerHost
+from azure.ai.agentserver.invocations._constants import InvocationConstants
+from azure.ai.agentserver.invocations._invocation import _MAX_ID_LENGTH, _VALID_ID_RE
+
+_HEADER = InvocationConstants.INVOCATION_ID_HEADER
+
+
+def _build_app() -> InvocationAgentServerHost:
+    app = InvocationAgentServerHost()
+
+    @app.invoke_handler
+    async def handle(request: Request) -> Response:
+        return Response(content=b"ok")
+
+    @app.get_invocation_handler
+    async def get_handler(request: Request) -> Response:
+        return Response(content=b"got")
+
+    return app
+
+
+async def _echoed_id(path_id: str) -> str:
+    transport = ASGITransport(app=_build_app())
+    async with AsyncClient(transport=transport, base_url="http://testserver") as client:
+        resp = await client.get(f"/invocations/{path_id}")
+    return resp.headers[_HEADER]
+
+
+@pytest.mark.asyncio
+async def test_invalid_char_id_not_reflected_in_header():
+    """A path id with characters outside the id allow-list is replaced by a
+    safe fallback rather than echoed back verbatim."""
+    echoed = await _echoed_id("bad~id")
+    assert echoed != "bad~id"
+    assert _VALID_ID_RE.match(echoed)
+
+
+@pytest.mark.asyncio
+async def test_overlong_id_not_reflected_in_header():
+    """An over-length path id does not bypass the ``_MAX_ID_LENGTH`` cap."""
+    echoed = await _echoed_id("a" * (_MAX_ID_LENGTH + 50))
+    assert len(echoed) <= _MAX_ID_LENGTH
+    assert _VALID_ID_RE.match(echoed)
+
+
+@pytest.mark.asyncio
+async def test_valid_id_passes_through_unchanged():
+    """A well-formed path id is preserved so lookups keep working."""
+    echoed = await _echoed_id("valid-id_123.abc")
+    assert echoed == "valid-id_123.abc"