From b324ad632ab2e2a5f9e393909108d6a87dc9bed3 Mon Sep 17 00:00:00 2001 From: Yufeng He <40085740+he-yufeng@users.noreply.github.com> Date: Fri, 22 May 2026 23:51:44 +0800 Subject: [PATCH 1/2] fix: pass Foundry agent default headers --- .../foundry/agent_framework_foundry/_agent.py | 6 +++ .../tests/foundry/test_foundry_agent.py | 42 +++++++++++++++++++ 2 files changed, 48 insertions(+) diff --git a/python/packages/foundry/agent_framework_foundry/_agent.py b/python/packages/foundry/agent_framework_foundry/_agent.py index 056d3977af..25b3b388d9 100644 --- a/python/packages/foundry/agent_framework_foundry/_agent.py +++ b/python/packages/foundry/agent_framework_foundry/_agent.py @@ -610,6 +610,7 @@ def __init__( credential: AzureCredentialTypes | None = None, project_client: AIProjectClient | None = None, allow_preview: bool | None = None, + default_headers: Mapping[str, str] | None = None, tools: FunctionTool | Callable[..., Any] | Sequence[FunctionTool | Callable[..., Any]] | None = None, context_providers: Sequence[ContextProvider] | None = None, middleware: Sequence[MiddlewareTypes] | None = None, @@ -639,6 +640,7 @@ def __init__( credential: Azure credential for authentication. project_client: An existing AIProjectClient to use. allow_preview: Enables preview opt-in on internally-created AIProjectClient. + default_headers: Additional HTTP headers for requests made through the OpenAI client. tools: Function tools to provide to the agent. Only ``FunctionTool`` objects are accepted. context_providers: Optional context providers for injecting dynamic context. middleware: Optional agent-level middleware. @@ -672,6 +674,7 @@ def __init__( "credential": credential, "project_client": project_client, "allow_preview": allow_preview, + "default_headers": default_headers, "env_file_path": env_file_path, "env_file_encoding": env_file_encoding, } @@ -894,6 +897,7 @@ def __init__( credential: AzureCredentialTypes | None = None, project_client: AIProjectClient | None = None, allow_preview: bool | None = None, + default_headers: Mapping[str, str] | None = None, tools: FunctionTool | Callable[..., Any] | Sequence[FunctionTool | Callable[..., Any]] | None = None, context_providers: Sequence[ContextProvider] | None = None, middleware: Sequence[MiddlewareTypes] | None = None, @@ -936,6 +940,7 @@ def __init__( Set this to ``True`` for HostedAgents that need preview-only session APIs, including lazy service session creation from ``isolation_key``. + default_headers: Additional HTTP headers for requests made through the OpenAI client. tools: Function tools to provide to the agent. Only ``FunctionTool`` objects are accepted. context_providers: Optional context providers. middleware: Optional agent-level middleware. @@ -963,6 +968,7 @@ def __init__( credential=credential, project_client=project_client, allow_preview=allow_preview, + default_headers=default_headers, tools=tools, context_providers=context_providers, middleware=middleware, diff --git a/python/packages/foundry/tests/foundry/test_foundry_agent.py b/python/packages/foundry/tests/foundry/test_foundry_agent.py index 6ae8a433ca..3ca7016566 100644 --- a/python/packages/foundry/tests/foundry/test_foundry_agent.py +++ b/python/packages/foundry/tests/foundry/test_foundry_agent.py @@ -505,6 +505,46 @@ def test_raw_foundry_agent_init_creates_client() -> None: assert agent.client.agent_name == "test-agent" +def test_raw_foundry_agent_init_passes_default_headers_to_client() -> None: + """Test that RawFoundryAgent passes default_headers to the underlying client.""" + + mock_project = MagicMock() + mock_project.get_openai_client.return_value = MagicMock() + default_headers = {"x-ms-user-isolation-key": "user-1"} + + RawFoundryAgent( + project_client=mock_project, + agent_name="hosted-agent", + allow_preview=True, + default_headers=default_headers, + ) + + mock_project.get_openai_client.assert_called_once_with( + default_headers=default_headers, + agent_name="hosted-agent", + ) + + +def test_foundry_agent_init_passes_default_headers_to_client() -> None: + """Test that FoundryAgent passes default_headers to the underlying client.""" + + mock_project = MagicMock() + mock_project.get_openai_client.return_value = MagicMock() + default_headers = {"x-ms-user-isolation-key": "user-1"} + + FoundryAgent( + project_client=mock_project, + agent_name="hosted-agent", + allow_preview=True, + default_headers=default_headers, + ) + + mock_project.get_openai_client.assert_called_once_with( + default_headers=default_headers, + agent_name="hosted-agent", + ) + + def test_raw_foundry_agent_init_with_custom_client_type() -> None: """Test that client_type parameter is respected.""" @@ -523,6 +563,7 @@ def test_raw_foundry_agent_init_with_custom_client_type() -> None: def test_raw_foundry_agent_init_uses_explicit_parameters() -> None: signature = inspect.signature(RawFoundryAgent.__init__) + assert "default_headers" in signature.parameters assert "instructions" in signature.parameters assert "default_options" in signature.parameters assert "compaction_strategy" in signature.parameters @@ -534,6 +575,7 @@ def test_raw_foundry_agent_init_uses_explicit_parameters() -> None: def test_foundry_agent_init_uses_explicit_parameters() -> None: signature = inspect.signature(FoundryAgent.__init__) + assert "default_headers" in signature.parameters assert "instructions" in signature.parameters assert "default_options" in signature.parameters assert "compaction_strategy" in signature.parameters From 91770df522071d0ab603ea602f7d03ca7d0a43f0 Mon Sep 17 00:00:00 2001 From: Yufeng He <40085740+he-yufeng@users.noreply.github.com> Date: Thu, 28 May 2026 01:59:31 +0800 Subject: [PATCH 2/2] test: loosen Foundry default header assertions --- .../foundry/tests/foundry/test_foundry_agent.py | 14 ++++---------- 1 file changed, 4 insertions(+), 10 deletions(-) diff --git a/python/packages/foundry/tests/foundry/test_foundry_agent.py b/python/packages/foundry/tests/foundry/test_foundry_agent.py index 3ca7016566..44bc744f64 100644 --- a/python/packages/foundry/tests/foundry/test_foundry_agent.py +++ b/python/packages/foundry/tests/foundry/test_foundry_agent.py @@ -515,14 +515,11 @@ def test_raw_foundry_agent_init_passes_default_headers_to_client() -> None: RawFoundryAgent( project_client=mock_project, agent_name="hosted-agent", - allow_preview=True, default_headers=default_headers, ) - mock_project.get_openai_client.assert_called_once_with( - default_headers=default_headers, - agent_name="hosted-agent", - ) + mock_project.get_openai_client.assert_called_once() + assert mock_project.get_openai_client.call_args.kwargs["default_headers"] == default_headers def test_foundry_agent_init_passes_default_headers_to_client() -> None: @@ -535,14 +532,11 @@ def test_foundry_agent_init_passes_default_headers_to_client() -> None: FoundryAgent( project_client=mock_project, agent_name="hosted-agent", - allow_preview=True, default_headers=default_headers, ) - mock_project.get_openai_client.assert_called_once_with( - default_headers=default_headers, - agent_name="hosted-agent", - ) + mock_project.get_openai_client.assert_called_once() + assert mock_project.get_openai_client.call_args.kwargs["default_headers"] == default_headers def test_raw_foundry_agent_init_with_custom_client_type() -> None: