diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_toolbox.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_toolbox.py index 0696e70d8c1..46ae6d89ea6 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_toolbox.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_toolbox.py @@ -22,6 +22,7 @@ ) from agent_framework._telemetry import mark_feature_used from azure.ai.agentserver.core import get_request_context +from azure.ai.agentserver.core.platform_headers import FOUNDRY_CALL_ID from typing_extensions import override from ._feature_usage import FeatureIndex @@ -103,6 +104,8 @@ def __init__(self, credential: AzureCredentialTypes, scope: str) -> None: def _apply_headers(self, request: httpx.Request, token: AccessToken) -> None: request.headers["Authorization"] = f"Bearer {token.token}" + # A Request may be retried or reused after an earlier auth flow, so absence must clear prior caller context. + request.headers.pop(FOUNDRY_CALL_ID, None) for key, value in get_request_context().platform_headers().items(): request.headers[key] = value diff --git a/python/packages/foundry_hosting/tests/test_toolbox.py b/python/packages/foundry_hosting/tests/test_toolbox.py index 444b3007260..9ff23376c0e 100644 --- a/python/packages/foundry_hosting/tests/test_toolbox.py +++ b/python/packages/foundry_hosting/tests/test_toolbox.py @@ -217,6 +217,38 @@ async def test_auth_flow_omits_call_id_when_absent() -> None: assert "x-agent-foundry-call-id" not in prepared.headers +async def test_async_auth_flow_removes_call_id_from_reused_request_when_context_absent() -> None: + auth = _ToolboxAuth(_FakeCredential(), "scope") # type: ignore + request = httpx.Request("POST", "https://h/toolboxes/tb/mcp") + + token = set_request_context(FoundryAgentRequestContext(call_id="canary-call-id")) + try: + prepared = await anext(auth.async_auth_flow(request)) + assert prepared.headers["x-agent-foundry-call-id"] == "canary-call-id" + finally: + reset_request_context(token) + + prepared = await anext(auth.async_auth_flow(request)) + + assert "x-agent-foundry-call-id" not in prepared.headers + + +def test_sync_auth_flow_removes_call_id_from_reused_request_when_context_absent() -> None: + auth = _ToolboxAuth(_FakeCredential(), "scope") # type: ignore + request = httpx.Request("POST", "https://h/toolboxes/tb/mcp") + + token = set_request_context(FoundryAgentRequestContext(call_id="canary-call-id")) + try: + prepared = next(auth.sync_auth_flow(request)) + assert prepared.headers["x-agent-foundry-call-id"] == "canary-call-id" + finally: + reset_request_context(token) + + prepared = next(auth.sync_auth_flow(request)) + + assert "x-agent-foundry-call-id" not in prepared.headers + + async def test_close_closes_owned_http_client() -> None: toolbox = FoundryToolbox( _FakeCredential(), # type: ignore