Skip to content

feat: implement user agent handling across providers and normalize headers #8445

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Soulter wants to merge 1 commit into
base: master
Choose a base branch
from
Open

feat: implement user agent handling across providers and normalize headers #8445

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
7 changes: 4 additions & 3 deletions astrbot/core/config/default.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
from astrbot.core.utils.astrbot_path import get_astrbot_data_path

VERSION = "4.25.2"
ASTRBOT_USER_AGENT = f"astrbot/{VERSION.removeprefix('v')}"
DB_PATH = os.path.join(get_astrbot_data_path(), "data_v4.db")
PERSONAL_WECHAT_CONFIG_METADATA = {
"weixin_oc_base_url": {
Expand Down Expand Up @@ -1199,7 +1200,7 @@
"api_base": "https://api.kimi.com/coding",
"timeout": 120,
"proxy": "",
"custom_headers": {"User-Agent": "claude-code/0.1.0"},
"custom_headers": {"User-Agent": ASTRBOT_USER_AGENT},
"anth_thinking_config": {"type": "", "budget": 0, "effort": ""},
},
"Moonshot": {
Expand Down Expand Up @@ -1236,7 +1237,7 @@
"api_base": "https://api.minimaxi.com/anthropic",
"timeout": 120,
"proxy": "",
"custom_headers": {"User-Agent": "claude-code/0.1.0"},
"custom_headers": {"User-Agent": ASTRBOT_USER_AGENT},
"anth_thinking_config": {"type": "", "budget": 0, "effort": ""},
},
"Xiaomi": {
Expand All @@ -1261,7 +1262,7 @@
"api_base": "https://token-plan-cn.xiaomimimo.com/anthropic",
"timeout": 120,
"proxy": "",
"custom_headers": {"User-Agent": "claude-code/0.1.0"},
"custom_headers": {"User-Agent": ASTRBOT_USER_AGENT},
"anth_thinking_config": {"type": "", "budget": 0, "effort": ""},
},
"xAI": {
Expand Down
22 changes: 12 additions & 10 deletions astrbot/core/provider/sources/anthropic_source.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,9 +13,11 @@
from astrbot import logger
from astrbot.api.provider import Provider
from astrbot.core.agent.message import AudioURLPart, ContentPart, ImageURLPart, TextPart
from astrbot.core.config.default import ASTRBOT_USER_AGENT
from astrbot.core.exceptions import EmptyModelOutputError
from astrbot.core.provider.entities import LLMResponse, TokenUsage
from astrbot.core.provider.func_tool_manager import ToolSet
from astrbot.core.utils.http_headers import apply_default_headers, normalize_headers
from astrbot.core.utils.io import download_image_by_url
from astrbot.core.utils.network_utils import (
create_proxy_client,
Expand Down Expand Up @@ -50,13 +52,12 @@ def _ensure_usable_response(

@staticmethod
def _normalize_custom_headers(provider_config: dict) -> dict[str, str] | None:
custom_headers = provider_config.get("custom_headers", {})
if not isinstance(custom_headers, dict) or not custom_headers:
normalized_headers = normalize_headers(
provider_config.get("custom_headers", {})
)
if not normalized_headers:
return None
normalized_headers: dict[str, str] = {}
for key, value in custom_headers.items():
normalized_headers[str(key)] = str(value)
return normalized_headers or None
return normalized_headers

@classmethod
def _resolve_custom_headers(
Expand All @@ -67,9 +68,7 @@ def _resolve_custom_headers(
) -> dict[str, str] | None:
merged_headers = cls._normalize_custom_headers(provider_config) or {}
if required_headers:
for header_name, header_value in required_headers.items():
if not merged_headers.get(header_name, "").strip():
merged_headers[header_name] = header_value
merged_headers = apply_default_headers(merged_headers, required_headers)
return merged_headers or None

def __init__(
Expand All @@ -89,7 +88,10 @@ def __init__(
if isinstance(self.timeout, str):
self.timeout = int(self.timeout)
self.thinking_config = provider_config.get("anth_thinking_config", {})
self.custom_headers = self._resolve_custom_headers(provider_config)
self.custom_headers = self._resolve_custom_headers(
provider_config,
required_headers={"User-Agent": ASTRBOT_USER_AGENT},
)

if use_api_key:
self._init_api_key(provider_config)
Expand Down
36 changes: 34 additions & 2 deletions astrbot/core/provider/sources/gemini_source.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,11 +18,13 @@
from astrbot import logger
from astrbot.api.provider import Provider
from astrbot.core.agent.message import AudioURLPart, ContentPart, ImageURLPart, TextPart
from astrbot.core.config.default import ASTRBOT_USER_AGENT
from astrbot.core.exceptions import EmptyModelOutputError
from astrbot.core.message.message_event_result import MessageChain
from astrbot.core.provider.entities import LLMResponse, TokenUsage
from astrbot.core.provider.func_tool_manager import ToolSet
from astrbot.core.utils.astrbot_path import get_astrbot_temp_path
from astrbot.core.utils.http_headers import apply_default_headers, normalize_headers
from astrbot.core.utils.io import download_file, download_image_by_url
from astrbot.core.utils.media_utils import ensure_wav
from astrbot.core.utils.network_utils import is_connection_error, log_connection_failure
Expand Down Expand Up @@ -76,24 +78,49 @@ def __init__(
if self.api_base and self.api_base.endswith("/"):
self.api_base = self.api_base[:-1]

self.custom_headers = self._resolve_custom_headers(provider_config)
Copy link
Copy Markdown
Contributor

@sourcery-ai sourcery-ai Bot May 30, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

issue (complexity): Consider simplifying header handling by computing normalized headers once and reusing them everywhere instead of mutating Gemini client internals after creation.

You can simplify this without losing any functionality by:

  1. Building the final, normalized headers once.
  2. Passing them into both HttpOptions and httpx.AsyncClient.
  3. Dropping the post-hoc mutation of Gemini client internals.

That removes indirection and reflection and keeps the header behavior centralized.

1. Simplify _resolve_custom_headers

Normalize keys (including user-agent) and apply defaults in one pass:

@staticmethod
def _resolve_custom_headers(provider_config: dict) -> dict[str, str]:
    # Normalize keys/values
    raw = normalize_headers(provider_config.get("custom_headers", {}))

    # Normalize the user-agent key while building the dict
    normalized = {
        ("user-agent" if key.lower() == "user-agent" else key): value
        for key, value in raw.items()
    }

    # Apply defaults with the normalized UA key
    return apply_default_headers(
        normalized,
        {"user-agent": ASTRBOT_USER_AGENT},
    )

This removes the extra comprehension after apply_default_headers.

2. Remove _set_gemini_user_agent and avoid private internals

You already pass headers into HttpOptions, which are used by the Gemini client. If you ensure self.custom_headers["user-agent"] is always the final value (as above), you don’t need to mutate http_options.headers via private attributes:

# Delete this helper entirely
# @staticmethod
# def _set_gemini_user_agent(...):
#     ...

Then simplify _init_client to rely only on the precomputed, final headers:

def _init_client(self) -> None:
    """初始化Gemini客户端"""
    proxy = self.provider_config.get("proxy", "")

    http_options = types.HttpOptions(
        base_url=self.api_base,
        headers=self.custom_headers,
        timeout=self.timeout * 1000,  # 毫秒
    )

    async_client_kwargs: dict = {
        "base_url": self.api_base,
        "headers": self.custom_headers,
        "timeout": self.timeout,
    }
    if proxy:
        async_client_kwargs["proxy"] = proxy
        async_client_kwargs["trust_env"] = False
        logger.info("[Gemini] 使用代理")
    else:
        async_client_kwargs["trust_env"] = True

    if self._http_client is not None:
        self._stale_http_clients = [self._http_client]

    self._http_client = httpx.AsyncClient(**async_client_kwargs)
    http_options.httpx_async_client = self._http_client

    self.client = genai.Client(
        api_key=self.chosen_api_key,
        http_options=http_options,
    ).aio

This keeps:

  • Single source of truth for headers (self.custom_headers).
  • Consistent headers between httpx.AsyncClient and HttpOptions.
  • No reliance on private _api_client / _http_options.
  • No triple handling or post-hoc header mutation.

self._http_client: httpx.AsyncClient | None = None
self._stale_http_clients: list[httpx.AsyncClient] = []
self._init_client()
self.set_model(provider_config.get("model", "unknown"))
self._init_safety_settings()

@staticmethod
def _resolve_custom_headers(provider_config: dict) -> dict[str, str]:
headers = apply_default_headers(
normalize_headers(provider_config.get("custom_headers", {})),
{"user-agent": ASTRBOT_USER_AGENT},
)
return {
"user-agent" if key.lower() == "user-agent" else key: value
for key, value in headers.items()
}

@staticmethod
def _set_gemini_user_agent(client: object, user_agent: str) -> None:
api_client = getattr(client, "_api_client", None)
http_options = getattr(api_client, "_http_options", None)
if http_options is None or http_options.headers is None:
return
Comment on lines +100 to +104
Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist Bot May 30, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

To prevent potential AttributeError or type mismatch issues if the internal structure of the Gemini SDK client changes in future versions, we should add a defensive type check to ensure http_options.headers is indeed a mutable dictionary before attempting to modify it.

Suggested change
def _set_gemini_user_agent(client: object, user_agent: str) -> None:
api_client = getattr(client, "_api_client", None)
http_options = getattr(api_client, "_http_options", None)
if http_options is None or http_options.headers is None:
return
@staticmethod
def _set_gemini_user_agent(client: object, user_agent: str) -> None:
api_client = getattr(client, "_api_client", None)
http_options = getattr(api_client, "_http_options", None)
if http_options is None or not isinstance(http_options.headers, dict):
return

for key in list(http_options.headers):
if key.lower() == "user-agent":
http_options.headers.pop(key)
http_options.headers["user-agent"] = user_agent

def _init_client(self) -> None:
"""初始化Gemini客户端"""
proxy = self.provider_config.get("proxy", "")
http_options = types.HttpOptions(
base_url=self.api_base,
headers=dict(self.custom_headers),
timeout=self.timeout * 1000, # 毫秒
)

# 强制使用 httpx 作为异步 HTTP 后端，避免 aiohttp 响应类型兼容问题 (#7564)
# httpx.AsyncClient 的 timeout 单位为秒（与 HttpOptions 的毫秒不同）
async_client_kwargs: dict = {
"base_url": self.api_base,
"headers": dict(self.custom_headers),
"timeout": self.timeout,
}
if proxy:
Expand All @@ -112,10 +139,15 @@ def _init_client(self) -> None:
self._http_client = httpx.AsyncClient(**async_client_kwargs)
http_options.httpx_async_client = self._http_client

self.client = genai.Client(
genai_client = genai.Client(
api_key=self.chosen_api_key,
http_options=http_options,
).aio
)
self._set_gemini_user_agent(
genai_client,
self.custom_headers["user-agent"],
)
self.client = genai_client.aio

def _init_safety_settings(self) -> None:
"""初始化安全设置"""
Expand Down
4 changes: 3 additions & 1 deletion astrbot/core/provider/sources/kimi_code_source.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,9 +1,11 @@
from astrbot.core.config.default import ASTRBOT_USER_AGENT

from ..register import register_provider_adapter
from .anthropic_source import ProviderAnthropic

KIMI_CODE_API_BASE = "https://api.kimi.com/coding"
KIMI_CODE_DEFAULT_MODEL = "kimi-for-coding"
KIMI_CODE_USER_AGENT = "claude-code/0.1.0"
KIMI_CODE_USER_AGENT = ASTRBOT_USER_AGENT


@register_provider_adapter(
Expand Down
17 changes: 10 additions & 7 deletions astrbot/core/provider/sources/openai_source.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,10 +34,12 @@
TextPart,
)
from astrbot.core.agent.tool import ToolSet
from astrbot.core.config.default import ASTRBOT_USER_AGENT
from astrbot.core.exceptions import EmptyModelOutputError
from astrbot.core.message.message_event_result import MessageChain
from astrbot.core.provider.entities import LLMResponse, TokenUsage, ToolCallsResult
from astrbot.core.utils.astrbot_path import get_astrbot_temp_path
from astrbot.core.utils.http_headers import apply_default_headers, normalize_headers
from astrbot.core.utils.io import download_file, download_image_by_url
from astrbot.core.utils.media_utils import ensure_wav
from astrbot.core.utils.network_utils import (
Expand Down Expand Up @@ -68,6 +70,13 @@ class ProviderOpenAIOfficial(Provider):
"AVIF": "image/avif",
}

@staticmethod
def _resolve_custom_headers(provider_config: dict) -> dict[str, str]:
return apply_default_headers(
normalize_headers(provider_config.get("custom_headers", {})),
{"User-Agent": ASTRBOT_USER_AGENT},
)

@classmethod
def _truncate_error_text_candidate(cls, text: str) -> str:
if len(text) <= cls._ERROR_TEXT_CANDIDATE_MAX_CHARS:
Expand Down Expand Up @@ -498,16 +507,10 @@ def __init__(self, provider_config, provider_settings) -> None:
self.api_keys: list = super().get_keys()
self.chosen_api_key = self.api_keys[0] if len(self.api_keys) > 0 else None
self.timeout = provider_config.get("timeout", 120)
self.custom_headers = provider_config.get("custom_headers", {})
self.custom_headers = self._resolve_custom_headers(provider_config)
if isinstance(self.timeout, str):
self.timeout = int(self.timeout)

if not isinstance(self.custom_headers, dict) or not self.custom_headers:
self.custom_headers = None
else:
for key in self.custom_headers:
self.custom_headers[key] = str(self.custom_headers[key])

if "api_version" in provider_config:
# Using Azure OpenAI API
self.client = AsyncAzureOpenAI(
Expand Down
31 changes: 31 additions & 0 deletions astrbot/core/utils/http_headers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,31 @@
from collections.abc import Mapping


def normalize_headers(headers: object) -> dict[str, str]:
if not isinstance(headers, dict):
return {}
return {str(key): str(value) for key, value in headers.items()}
Comment on lines +4 to +7
Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist Bot May 30, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

In normalize_headers, if a header value is None (for example, if custom_headers contains {"User-Agent": null} in JSON), str(value) will convert it to the literal string "None". This can lead to invalid headers like User-Agent: None being sent to LLM providers, which might cause request failures or unexpected behavior.

We should defensively convert None values to an empty string "" so that apply_default_headers can correctly recognize them as empty and fall back to the default header values.

def normalize_headers(headers: object) -> dict[str, str]:
    if not isinstance(headers, dict):
        return {}
    return {
        str(key): str(value) if value is not None else ""
        for key, value in headers.items()
    }



def apply_default_headers(
headers: dict[str, str],
default_headers: Mapping[str, str],
) -> dict[str, str]:
merged_headers = dict(headers)
for default_name, default_value in default_headers.items():
existing_name = next(
(
header_name
for header_name in merged_headers
if header_name.lower() == default_name.lower()
),
None,
)
if existing_name is None:
merged_headers[default_name] = default_value
continue
if merged_headers[existing_name].strip():
continue
merged_headers.pop(existing_name)
merged_headers[default_name] = default_value
return merged_headers
20 changes: 20 additions & 0 deletions tests/test_anthropic_kimi_code_provider.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@

import astrbot.core.provider.sources.anthropic_source as anthropic_source
import astrbot.core.provider.sources.kimi_code_source as kimi_code_source
from astrbot.core.config.default import ASTRBOT_USER_AGENT
from astrbot.core.exceptions import EmptyModelOutputError
from astrbot.core.provider.entities import LLMResponse

Expand All @@ -16,6 +17,25 @@ async def close(self):
return None


def test_anthropic_provider_uses_astrbot_default_user_agent(monkeypatch):
monkeypatch.setattr(anthropic_source, "AsyncAnthropic", _FakeAsyncAnthropic)

provider = anthropic_source.ProviderAnthropic(
provider_config={
"id": "anthropic-test",
"type": "anthropic_chat_completion",
"model": "claude-test",
"key": ["test-key"],
},
Comment on lines +20 to +29
Copy link
Copy Markdown
Contributor

@sourcery-ai sourcery-ai Bot May 30, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

suggestion (testing): Extend Anthropics tests to cover custom headers with/without user agent and type normalization.

Consider adding a couple of nearby tests to fully cover the new header behavior:

  • One where provider_config.custom_headers includes its own User-Agent plus extra headers, to verify a non-empty custom UA is preserved (and the default UA is not also applied) and that additional headers remain.
  • Optionally, one where custom_headers includes non-string values (e.g., integers) to confirm normalization still coerces values to strings as before.

This will better demonstrate that existing Anthropics header behavior is preserved with the new shared helpers.

Suggested implementation:

    assert provider.client.kwargs["default_headers"] == {
        "User-Agent": ASTRBOT_USER_AGENT,
    }


def test_anthropic_provider_preserves_custom_user_agent_and_headers(monkeypatch):
    monkeypatch.setattr(anthropic_source, "AsyncAnthropic", _FakeAsyncAnthropic)

    custom_headers = {
        "User-Agent": "MyApp/1.0 (test)",
        "X-Extra-Header": "extra-value",
    }

    provider = anthropic_source.ProviderAnthropic(
        provider_config={
            "id": "anthropic-test",
            "type": "anthropic_chat_completion",
            "model": "claude-test",
            "key": ["test-key"],
            "custom_headers": custom_headers,
        },
        provider_settings={},
    )

    # Custom non-empty User-Agent should be preserved, and the default ASTRBOT_USER_AGENT
    # should not be added on top.
    assert provider.custom_headers == custom_headers
    assert provider.client.kwargs["default_headers"] == custom_headers


def test_anthropic_provider_normalizes_custom_header_values_to_strings(monkeypatch):
    monkeypatch.setattr(anthropic_source, "AsyncAnthropic", _FakeAsyncAnthropic)

    custom_headers = {
        "User-Agent": "MyApp/2.0 (test)",
        "X-Int-Header": 123,
        "X-Bool-Header": True,
    }

    provider = anthropic_source.ProviderAnthropic(
        provider_config={
            "id": "anthropic-test",
            "type": "anthropic_chat_completion",
            "model": "claude-test",
            "key": ["test-key"],
            "custom_headers": custom_headers,
        },
        provider_settings={},
    )

    expected_headers = {
        "User-Agent": "MyApp/2.0 (test)",
        "X-Int-Header": "123",
        "X-Bool-Header": "True",
    }

    assert provider.custom_headers == expected_headers
    assert provider.client.kwargs["default_headers"] == expected_headers


def test_anthropic_provider_passes_custom_headers_via_default_headers(monkeypatch):
    monkeypatch.setattr(anthropic_source, "AsyncAnthropic", _FakeAsyncAnthropic)

If ProviderAnthropic currently expects custom_headers to be passed via provider_settings instead of provider_config, adjust the tests to move the custom_headers dict from provider_config={...} into provider_settings={"custom_headers": custom_headers} to match the existing convention in the rest of the file.

provider_settings={},
)

assert provider.custom_headers == {"User-Agent": ASTRBOT_USER_AGENT}
assert provider.client.kwargs["default_headers"] == {
"User-Agent": ASTRBOT_USER_AGENT,
}


def test_anthropic_provider_passes_custom_headers_via_default_headers(monkeypatch):
monkeypatch.setattr(anthropic_source, "AsyncAnthropic", _FakeAsyncAnthropic)

Expand Down
38 changes: 38 additions & 0 deletions tests/test_gemini_source.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,10 +1,48 @@
import pytest

from astrbot.core.config.default import ASTRBOT_USER_AGENT
from astrbot.core.exceptions import EmptyModelOutputError
import astrbot.core.provider.sources.gemini_source as gemini_source
from astrbot.core.provider.entities import LLMResponse
from astrbot.core.provider.sources.gemini_source import ProviderGoogleGenAI


class _FakeGenAIClient:
def __init__(self, **kwargs):
self.kwargs = kwargs
self._api_client = type(
"FakeAPIClient",
(),
{"_http_options": kwargs["http_options"]},
)()
self.aio = type("FakeAioClient", (), {"_api_client": self._api_client})()


@pytest.mark.asyncio
async def test_gemini_provider_uses_astrbot_default_user_agent(monkeypatch):
monkeypatch.setattr(gemini_source.genai, "Client", _FakeGenAIClient)

provider = ProviderGoogleGenAI(
provider_config={
"id": "gemini-test",
"type": "googlegenai_chat_completion",
"model": "gemini-test",
"key": ["test-key"],
"api_base": "https://generativelanguage.googleapis.com/",
},
provider_settings={},
)

try:
assert provider.custom_headers["user-agent"] == ASTRBOT_USER_AGENT
assert provider.client._api_client._http_options.headers["user-agent"] == (
ASTRBOT_USER_AGENT
)
assert provider._http_client.headers["user-agent"] == ASTRBOT_USER_AGENT
finally:
await provider.terminate()


def test_gemini_empty_output_raises_empty_model_output_error():
llm_response = LLMResponse(role="assistant")

Expand Down
Loading
Loading