diff --git a/python/packages/foundry/agent_framework_foundry/_agent.py b/python/packages/foundry/agent_framework_foundry/_agent.py index b473c787e5..2453025c26 100644 --- a/python/packages/foundry/agent_framework_foundry/_agent.py +++ b/python/packages/foundry/agent_framework_foundry/_agent.py @@ -198,7 +198,9 @@ def __init__( self._should_close_client = True # Get OpenAI client from project - async_client = self.project_client.get_openai_client() + async_client = self.project_client.get_openai_client( + default_headers=dict(default_headers) if default_headers else None, + ) super().__init__( async_client=async_client, diff --git a/python/packages/foundry/agent_framework_foundry/_chat_client.py b/python/packages/foundry/agent_framework_foundry/_chat_client.py index 4428d69dc6..1049679c25 100644 --- a/python/packages/foundry/agent_framework_foundry/_chat_client.py +++ b/python/packages/foundry/agent_framework_foundry/_chat_client.py @@ -206,7 +206,9 @@ def __init__( super().__init__( model=resolved_model, - async_client=project_client.get_openai_client(), + async_client=project_client.get_openai_client( + default_headers=dict(default_headers) if default_headers else None, + ), default_headers=default_headers, instruction_role=instruction_role, compaction_strategy=compaction_strategy, diff --git a/python/packages/foundry/tests/foundry/test_foundry_agent.py b/python/packages/foundry/tests/foundry/test_foundry_agent.py index 829af6ab87..6bec41561e 100644 --- a/python/packages/foundry/tests/foundry/test_foundry_agent.py +++ b/python/packages/foundry/tests/foundry/test_foundry_agent.py @@ -80,6 +80,35 @@ def test_raw_foundry_agent_chat_client_init_uses_explicit_parameters() -> None: assert all(parameter.kind != inspect.Parameter.VAR_KEYWORD for parameter in signature.parameters.values()) +def test_raw_foundry_agent_chat_client_init_forwards_default_headers_to_openai_client() -> None: + """default_headers must be forwarded into get_openai_client() so the underlying + AsyncOpenAI client sends them on every outbound request.""" + custom_headers = {"x-custom-header": "test-value"} + mock_project = MagicMock() + mock_project.get_openai_client.return_value = MagicMock() + + RawFoundryAgentChatClient( + project_client=mock_project, + agent_name="test-agent", + default_headers=custom_headers, + ) + + mock_project.get_openai_client.assert_called_once_with(default_headers=custom_headers) + + +def test_raw_foundry_agent_chat_client_init_without_default_headers_passes_none_to_openai_client() -> None: + """When no default_headers are provided, get_openai_client() receives None.""" + mock_project = MagicMock() + mock_project.get_openai_client.return_value = MagicMock() + + RawFoundryAgentChatClient( + project_client=mock_project, + agent_name="test-agent", + ) + + mock_project.get_openai_client.assert_called_once_with(default_headers=None) + + def test_raw_foundry_agent_chat_client_get_agent_reference_with_version() -> None: """Test agent reference includes version when provided.""" diff --git a/python/packages/foundry/tests/foundry/test_foundry_chat_client.py b/python/packages/foundry/tests/foundry/test_foundry_chat_client.py index a7b57029ba..bd04b58529 100644 --- a/python/packages/foundry/tests/foundry/test_foundry_chat_client.py +++ b/python/packages/foundry/tests/foundry/test_foundry_chat_client.py @@ -180,6 +180,37 @@ def test_init_with_default_header() -> None: assert client.default_headers[key] == value +def test_init_forwards_default_headers_to_openai_client() -> None: + """default_headers must be forwarded into get_openai_client() so the underlying + AsyncOpenAI client sends them on every outbound request.""" + custom_headers = {"x-custom-header": "test-value"} + mock_openai_client = _make_mock_openai_client() + project_client = MagicMock() + project_client.get_openai_client.return_value = mock_openai_client + + FoundryChatClient( + project_client=project_client, + model=_TEST_FOUNDRY_MODEL, + default_headers=custom_headers, + ) + + project_client.get_openai_client.assert_called_once_with(default_headers=custom_headers) + + +def test_init_without_default_headers_passes_none_to_openai_client() -> None: + """When no default_headers are provided, get_openai_client() receives None.""" + mock_openai_client = _make_mock_openai_client() + project_client = MagicMock() + project_client.get_openai_client.return_value = mock_openai_client + + FoundryChatClient( + project_client=project_client, + model=_TEST_FOUNDRY_MODEL, + ) + + project_client.get_openai_client.assert_called_once_with(default_headers=None) + + def test_init_with_project_endpoint_creates_project_client() -> None: credential = MagicMock() mock_openai_client = _make_mock_openai_client() diff --git a/python/packages/openai/agent_framework_openai/_shared.py b/python/packages/openai/agent_framework_openai/_shared.py index 7fb12ad14e..0d5dac4297 100644 --- a/python/packages/openai/agent_framework_openai/_shared.py +++ b/python/packages/openai/agent_framework_openai/_shared.py @@ -199,6 +199,8 @@ def load_openai_service_settings( if resolved_model := _resolve_named_setting(openai_settings, openai_model_fields): openai_settings["model"] = resolved_model if client: + if merged_headers: + client = client.with_options(default_headers=merged_headers) return openai_settings, client, False # type: ignore[return-value] if openai_settings.get("api_key") is not None or api_key_callable is not None: resolved_model = _resolve_named_setting(openai_settings, openai_model_fields) @@ -260,6 +262,8 @@ def load_openai_service_settings( f"or the {deployment_env_guidance} environment variable." ) if client: + if merged_headers: + client = client.with_options(default_headers=merged_headers) return azure_settings, client, True # type: ignore[return-value] client_args["default_headers"] = merged_headers if endpoint := azure_settings.get("endpoint"): diff --git a/python/packages/openai/tests/openai/test_openai_shared.py b/python/packages/openai/tests/openai/test_openai_shared.py index 86d43bc43b..3a148626a5 100644 --- a/python/packages/openai/tests/openai/test_openai_shared.py +++ b/python/packages/openai/tests/openai/test_openai_shared.py @@ -12,6 +12,7 @@ AZURE_OPENAI_TOKEN_SCOPE, _ensure_async_token_provider, _resolve_azure_credential_to_token_provider, + load_openai_service_settings, ) @@ -76,3 +77,66 @@ async def async_provider() -> str: result = await wrapper() assert result == "async-token" + + +def test_load_openai_service_settings_applies_default_headers_to_prebuilt_client() -> None: + """When a pre-built client is provided, default_headers must be applied to it. + + load_openai_service_settings used to early-return the pre-built client + without applying merged_headers, silently dropping any custom headers the + caller passed. + """ + pre_built = MagicMock() + new_client = MagicMock() + pre_built.with_options.return_value = new_client + + _, client, _ = load_openai_service_settings( + model="gpt-4o", + api_key=None, + credential=None, + org_id=None, + base_url=None, + endpoint=None, + api_version=None, + default_azure_api_version="2024-05-01-preview", + default_headers={"x-custom-header": "test-value"}, + client=pre_built, + env_file_path=None, + env_file_encoding=None, + ) + + pre_built.with_options.assert_called_once() + call_kwargs = pre_built.with_options.call_args.kwargs + assert call_kwargs.get("default_headers", {}).get("x-custom-header") == "test-value" + assert client is new_client + + +def test_load_openai_service_settings_no_headers_still_applies_app_info() -> None: + """Even with no default_headers, APP_INFO telemetry headers are applied via with_options.""" + pre_built = MagicMock() + new_client = MagicMock() + pre_built.with_options.return_value = new_client + + with ( + patch("agent_framework_openai._shared.APP_INFO", {"agent-framework-version": "python/test-version"}), + patch("agent_framework._telemetry.AGENT_FRAMEWORK_USER_AGENT", "agent-framework-python/test-version"), + ): + _, client, _ = load_openai_service_settings( + model="gpt-4o", + api_key=None, + credential=None, + org_id=None, + base_url=None, + endpoint=None, + api_version=None, + default_azure_api_version="2024-05-01-preview", + default_headers=None, + client=pre_built, + env_file_path=None, + env_file_encoding=None, + ) + + pre_built.with_options.assert_called_once() + call_kwargs = pre_built.with_options.call_args.kwargs + assert call_kwargs.get("default_headers", {}).get("User-Agent") == "agent-framework-python/test-version" + assert client is new_client