From e7f4fd8551a54893bdbecb7b311b515ab2f3f690 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] 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