Improve AgentServer warm response latency and reduce repeated work

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

@@ -10,6 +10,8 @@
 
 ### Other Changes
 
+- Generate fallback invocation and session IDs only when no valid supplied ID is available.
+
 ## 1.2.0b1 (2026-09-03)
 
 ### Features Added

sdk/agentserver/azure-ai-agentserver-invocations/azure/ai/agentserver/invocations/_invocation.py

@@ -128,22 +128,23 @@ def _ensure_log_filter() -> None:
         _log_filter_installed = True
 
 
-def _sanitize_id(value: str, fallback: str) -> str:
+def _sanitize_id(value: str, fallback: str | None = None) -> str:
     """Validate a user-provided ID string.
 
-    Returns *value* unchanged when it passes validation, otherwise returns
-    *fallback*.  This prevents excessively long or malformed IDs from
+    Returns *value* unchanged when it passes validation, otherwise uses
+    *fallback* or generates a UUID when it is omitted. This prevents malformed IDs from
     propagating into headers, span attributes, and log messages.
 
     :param value: The raw ID from a header or query parameter.
     :type value: str
-    :param fallback: A safe fallback value (typically a generated UUID).
-    :type fallback: str
+    :param fallback: A safe fallback value, or ``None`` to generate a UUID only
+        when *value* fails validation.
+    :type fallback: str | None
     :return: The validated ID or the fallback.
     :rtype: str
     """
     if not value or len(value) > _MAX_ID_LENGTH or not _VALID_ID_RE.match(value):
-        return fallback
+        return fallback if fallback is not None else str(uuid.uuid4())
     return value
 
 
@@ -517,14 +518,12 @@ async def _wrapped_body() -> AsyncIterator[Any]:
     async def _create_invocation_endpoint(self, request: Request) -> Response:
         invocation_id = _sanitize_id(
             request.headers.get(InvocationConstants.INVOCATION_ID_HEADER) or "",
-            str(uuid.uuid4()),
         )
         request.state.invocation_id = invocation_id
 
         # Session ID: query param overrides env var / generated UUID
         session_id = _sanitize_id(
             request.query_params.get("agent_session_id") or self.config.session_id or "",
-            str(uuid.uuid4()),
         )
         request.state.session_id = session_id
 

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

@@ -0,0 +1,104 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+"""UUID generation occurs only when an invocation/session ID needs a fallback."""
+
+import asyncio
+import uuid
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, patch
+
+import pytest
+from starlette.requests import Request
+from starlette.responses import JSONResponse
+
+from azure.ai.agentserver.invocations import _invocation
+
+
+@pytest.mark.parametrize("value", ["valid-id", "a" * _invocation._MAX_ID_LENGTH])
+def test_valid_identifier_never_generates_uuid(value):
+    with patch.object(_invocation.uuid, "uuid4") as generate:
+        assert _invocation._sanitize_id(value) == value
+        assert _invocation._sanitize_id(value, "fallback") == value
+    generate.assert_not_called()
+
+
+@pytest.mark.parametrize("value", ["", "bad id!", "a" * (_invocation._MAX_ID_LENGTH + 1)])
+def test_invalid_identifier_generates_one_uuid_when_no_fallback_supplied(value):
+    generated = uuid.UUID("11111111-1111-4111-8111-111111111111")
+    with patch.object(_invocation.uuid, "uuid4", return_value=generated) as generate:
+        assert _invocation._sanitize_id(value) == str(generated)
+    generate.assert_called_once_with()
+
+
+@pytest.mark.parametrize("fallback", ["", "literal-invalid fallback"])
+def test_explicit_fallback_retains_existing_get_cancel_behavior(fallback):
+    with patch.object(_invocation.uuid, "uuid4") as generate:
+        assert _invocation._sanitize_id("bad id!", fallback) == fallback
+    generate.assert_not_called()
+
+
+def _request(invocation_id, query):
+    return Request(
+        {
+            "type": "http",
+            "method": "POST",
+            "path": "/invocations",
+            "headers": [(_invocation.InvocationConstants.INVOCATION_ID_HEADER.encode(), invocation_id.encode())],
+            "query_string": query.encode(),
+        }
+    )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "invocation_id,query,configured,expected_calls,expected_session",
+    [
+        ("valid-id", "agent_session_id=query-session", "config-session", 0, "query-session"),
+        ("", "agent_session_id=query-session", "", 1, "query-session"),
+        ("valid-id", "", "config-session", 0, "config-session"),
+        ("valid-id", "agent_session_id=", "config-session", 0, "config-session"),
+        ("valid-id", "agent_session_id=bad%20id!", "config-session", 1, None),
+        ("bad id!", "", "", 2, None),
+    ],
+)
+async def test_post_endpoint_generates_only_needed_fallbacks(
+    invocation_id, query, configured, expected_calls, expected_session
+):
+    request = _request(invocation_id, query)
+    host = SimpleNamespace(
+        config=SimpleNamespace(session_id=configured),
+        _dispatch_invoke=AsyncMock(return_value=JSONResponse({"ok": True})),
+    )
+    real_uuid = uuid.uuid4
+    with patch.object(_invocation.uuid, "uuid4", wraps=real_uuid) as generate:
+        result = await _invocation.InvocationAgentServerHost._create_invocation_endpoint(host, request)
+    assert result.status_code == 200
+    assert generate.call_count == expected_calls
+    assert result.headers[_invocation.InvocationConstants.SESSION_ID_HEADER] == request.state.session_id
+    if expected_session is not None:
+        assert request.state.session_id == expected_session
+    else:
+        assert str(uuid.UUID(request.state.session_id)) == request.state.session_id
+    if invocation_id == "valid-id":
+        assert request.state.invocation_id == invocation_id
+    else:
+        assert str(uuid.UUID(request.state.invocation_id)) == request.state.invocation_id
+
+
+@pytest.mark.asyncio
+async def test_concurrent_requests_receive_distinct_fallback_ids():
+    host = SimpleNamespace(
+        config=SimpleNamespace(session_id=""),
+        _dispatch_invoke=AsyncMock(side_effect=lambda _: JSONResponse({"ok": True})),
+    )
+    first, second = _request("", ""), _request("", "")
+
+    await asyncio.gather(
+        _invocation.InvocationAgentServerHost._create_invocation_endpoint(host, first),
+        _invocation.InvocationAgentServerHost._create_invocation_endpoint(host, second),
+    )
+
+    assert (
+        len({first.state.invocation_id, first.state.session_id, second.state.invocation_id, second.state.session_id})
+        == 4
+    )

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

@@ -11,6 +11,12 @@
 
 ### Other Changes
 
+- Reuse request-scoped history lookups and concurrent input-reference resolution without caching failed or cancelled reads.
+- Flush streaming telemetry after request-owned handler and iterator cleanup,
+  including on disconnects, and before HTTP completion instead of delaying the first stream event.
+- Construct generated model types on demand while preserving real TypedDict contracts and public exports.
+- Avoid redundant event and recovery-seed copies while retaining validation and caller-owned mutation isolation.
+
 - Raised the minimum `azure-ai-agentserver-core` dependency to `>=2.2.0b1`,
   which provides the session GUID configuration and legacy task lookup used by
   resilient Responses.

sdk/agentserver/azure-ai-agentserver-responses/_scripts/extract_model_contracts.py

@@ -9,7 +9,6 @@
 import shutil
 from pathlib import Path
 
-
 ROOT_INIT_PREFIX = (
     "# coding=utf-8\n"
     "# --------------------------------------------------------------------------\n"
@@ -22,6 +21,7 @@
     "from .types import *  # type: ignore # noqa: F401,F403\n"
 )
 
+
 def _remove_pycache(root: Path) -> None:
     for pycache in root.rglob("__pycache__"):
         shutil.rmtree(pycache)
@@ -31,9 +31,7 @@ def _find_emitted_models_root(emitter_output_root: Path) -> Path:
     candidates = sorted(
         path
         for path in emitter_output_root.rglob("types.py")
-        if path.parent.name == "models"
-        and (path.parent / "_unions.py").exists()
-        and (path.parent / "models").is_dir()
+        if path.parent.name == "models" and (path.parent / "_unions.py").exists() and (path.parent / "models").is_dir()
     )
     if not candidates:
         raise FileNotFoundError(f"Could not find emitted TypedDict model package under {emitter_output_root}")
@@ -118,6 +116,11 @@ def finalize(emitter_output_root: Path, generated_root: Path) -> None:
         shutil.copy2(emitted_root / "models" / file_name, models_root / file_name)
 
     (generated_root / "__init__.py").write_text(ROOT_INIT_PREFIX, encoding="utf-8")
+    if __package__:
+        from .lazy_model_emitter import emit
+    else:
+        from lazy_model_emitter import emit
+    emit(generated_root)
     _remove_pycache(generated_root)