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 d0344705374..99a9bef5171 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_toolbox.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_toolbox.py @@ -42,6 +42,22 @@ DEFAULT_TOOLBOX_SCOPE = "https://ai.azure.com/.default" # Default timeout (seconds) for toolbox MCP requests. _DEFAULT_TIMEOUT = 120.0 +# Environment variable used to inject platform-provided toolbox feature flags. +_TOOLSET_FEATURES_ENV_VAR = "FOUNDRY_AGENT_TOOLSET_FEATURES" +# Mandatory preview feature flag for Foundry toolbox requests. +_MANDATORY_TOOLBOX_FEATURE = "Toolboxes=V1Preview" + + +def _build_toolbox_features_header(additional_features: str | None) -> str: + """Merge platform-provided features with the mandatory toolbox feature.""" + if additional_features is None or not additional_features.strip(): + return _MANDATORY_TOOLBOX_FEATURE + if any( + feature.strip().casefold() == _MANDATORY_TOOLBOX_FEATURE.casefold() + for feature in additional_features.split(",") + ): + return additional_features + return f"{_MANDATORY_TOOLBOX_FEATURE},{additional_features}" def _resolve_toolbox_endpoint() -> str: @@ -81,7 +97,7 @@ def _toolbox_name_from_endpoint(endpoint: str) -> str: class _ToolboxAuth(httpx.Auth): - """Injects a fresh bearer token and the platform call-id on every request. + """Injects a fresh bearer token, feature flags, and the platform call-id on every request. Both the synchronous (``sync_auth_flow``) and asynchronous (``async_auth_flow``) httpx auth hooks are implemented, so the same auth works regardless of which @@ -99,11 +115,14 @@ class _ToolboxAuth(httpx.Auth): def __init__(self, credential: AzureCredentialTypes, scope: str) -> None: self._credential = credential self._scope = scope + # Feature flags are startup configuration, matching the .NET toolbox service. + self._features_header = _build_toolbox_features_header(os.environ.get(_TOOLSET_FEATURES_ENV_VAR)) def _apply_headers(self, request: httpx.Request, token: AccessToken) -> None: request.headers["Authorization"] = f"Bearer {token.token}" for key, value in get_request_context().platform_headers().items(): request.headers[key] = value + request.headers["Foundry-Features"] = self._features_header def sync_auth_flow(self, request: httpx.Request) -> Generator[httpx.Request, httpx.Response, None]: # azure-core credentials cache the token internally and only refresh near @@ -138,7 +157,8 @@ class FoundryToolbox(MCPStreamableHTTPTool): ``MCPStreamableHTTPTool`` by hand it: - resolves the toolbox endpoint and tool name from the environment when not given, - - authenticates every request with a bearer token from ``credential``, and + - authenticates every request with a bearer token from ``credential``, + - sends the mandatory toolbox preview feature plus platform-provided feature flags, and - forwards the platform per-request call-id (``x-agent-foundry-call-id``) so the Foundry MCP proxy can resolve the caller context server-side. diff --git a/python/packages/foundry_hosting/tests/test_toolbox.py b/python/packages/foundry_hosting/tests/test_toolbox.py index 93442d48cb5..22dca726275 100644 --- a/python/packages/foundry_hosting/tests/test_toolbox.py +++ b/python/packages/foundry_hosting/tests/test_toolbox.py @@ -163,6 +163,32 @@ async def test_auth_flow_injects_bearer_token_async_credential() -> None: assert cred.scopes == ["https://ai.azure.com/.default"] +@pytest.mark.parametrize( + ("additional_features", "expected"), + [ + (None, "Toolboxes=V1Preview"), + (" ", "Toolboxes=V1Preview"), + ("FeatureOne=Enabled,FeatureTwo=Enabled", "Toolboxes=V1Preview,FeatureOne=Enabled,FeatureTwo=Enabled"), + ("FeatureOne=Enabled, toolboxes=v1preview ", "FeatureOne=Enabled, toolboxes=v1preview "), + ], +) +async def test_auth_flow_injects_foundry_features_header( + monkeypatch: pytest.MonkeyPatch, + additional_features: str | None, + expected: str, +) -> None: + if additional_features is None: + monkeypatch.delenv("FOUNDRY_AGENT_TOOLSET_FEATURES", raising=False) + else: + monkeypatch.setenv("FOUNDRY_AGENT_TOOLSET_FEATURES", additional_features) + auth = _ToolboxAuth(_FakeCredential(), "scope") # type: ignore + request = httpx.Request("POST", "https://h/toolboxes/tb/mcp") + + prepared = await anext(auth.async_auth_flow(request)) + + assert prepared.headers["Foundry-Features"] == expected + + def test_sync_auth_flow_injects_bearer_token() -> None: cred = _FakeCredential("sync123") auth = _ToolboxAuth(cred, "https://ai.azure.com/.default") # type: ignore @@ -174,6 +200,16 @@ def test_sync_auth_flow_injects_bearer_token() -> None: assert cred.scopes == ["https://ai.azure.com/.default"] +def test_sync_auth_flow_injects_foundry_features_header(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("FOUNDRY_AGENT_TOOLSET_FEATURES", "FeatureOne=Enabled") + auth = _ToolboxAuth(_FakeCredential(), "scope") # type: ignore + request = httpx.Request("POST", "https://h/toolboxes/tb/mcp") + + prepared = next(auth.sync_auth_flow(request)) + + assert prepared.headers["Foundry-Features"] == "Toolboxes=V1Preview,FeatureOne=Enabled" + + def test_sync_auth_flow_rejects_async_credential() -> None: auth = _ToolboxAuth(_FakeAsyncCredential(), "scope") # type: ignore request = httpx.Request("POST", "https://h/toolboxes/tb/mcp")