Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 12 additions & 0 deletions agent/auxiliary_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -6077,6 +6077,18 @@ def _wrap_if_needed(client_obj, final_model_str: str, base_url_str: str = "",
raw_base_for_wrap = custom_base
_clean_base2, _dq2 = _extract_url_query_params(openai_base)
_extra2 = {"default_query": _dq2} if _dq2 else {}
# Named custom providers may carry credentials in
# ``extra_headers`` (for example Cloudflare Access service
# headers). Preserve them on auxiliary clients just as the
# main client path does. Values are intentionally never
# logged.
_provider_headers2 = custom_entry.get("extra_headers")
if isinstance(_provider_headers2, dict) and _provider_headers2:
_extra2["default_headers"] = {
str(key): str(value)
for key, value in _provider_headers2.items()
if value is not None
}
_headers2 = _apply_user_default_headers(_extra2.get("default_headers"))
if _headers2:
_extra2["default_headers"] = _headers2
Expand Down
44 changes: 44 additions & 0 deletions tests/agent/test_auxiliary_named_custom_providers.py
Original file line number Diff line number Diff line change
Expand Up @@ -109,6 +109,50 @@ def test_named_custom_provider(self, tmp_path):
assert model == "my-model"
assert "beans.local" in str(client.base_url)

def test_named_custom_provider_extra_headers_reach_auxiliary_client(self, tmp_path):
"""Named-provider headers must be installed on auxiliary OpenAI clients."""
_write_config(tmp_path, {
"providers": {
"llama": {
"name": "llama",
"base_url": "https://llama.example/v1",
"api_key": "no-key-required",
"extra_headers": {
"CF-Access-Client-Id": "client.access",
"CF-Access-Client-Secret": "secret-value",
},
},
},
})
from agent import auxiliary_client

with patch.object(auxiliary_client, "OpenAI") as mock_openai:
mock_openai.return_value = MagicMock()
client, model = auxiliary_client.resolve_provider_client(
"llama", "gemma-test"
)

assert client is not None
assert model == "gemma-test"
kwargs = mock_openai.call_args.kwargs
assert kwargs["default_headers"] == {
"CF-Access-Client-Id": "client.access",
"CF-Access-Client-Secret": "secret-value",
}

def test_named_custom_provider_default_model(self, tmp_path):
_write_config(tmp_path, {
"model": {"default": "main-model"},
"custom_providers": [
{"name": "beans", "base_url": "http://beans.local/v1", "api_key": "k"},
],
})
from agent.auxiliary_client import resolve_provider_client
client, model = resolve_provider_client("beans")
assert client is not None
# Should use _read_main_model() fallback
assert model == "main-model"


def test_named_custom_no_api_key_uses_fallback(self, tmp_path):
_write_config(tmp_path, {
Expand Down
Loading