diff --git a/astrbot/core/config/default.py b/astrbot/core/config/default.py index 5bf7de7a45..e672680b80 100644 --- a/astrbot/core/config/default.py +++ b/astrbot/core/config/default.py @@ -1360,7 +1360,7 @@ "api_base": "https://api.kimi.com/coding", "timeout": 120, "proxy": "", - "custom_headers": {"User-Agent": "claude-code/0.1.0"}, + "custom_headers": {}, "anth_thinking_config": {"type": "", "budget": 0, "effort": ""}, }, "Moonshot": { @@ -1397,7 +1397,7 @@ "api_base": "https://api.minimaxi.com/anthropic", "timeout": 120, "proxy": "", - "custom_headers": {"User-Agent": "claude-code/0.1.0"}, + "custom_headers": {}, "anth_thinking_config": {"type": "", "budget": 0, "effort": ""}, }, "Xiaomi": { @@ -1422,7 +1422,7 @@ "api_base": "https://token-plan-cn.xiaomimimo.com/anthropic", "timeout": 120, "proxy": "", - "custom_headers": {"User-Agent": "claude-code/0.1.0"}, + "custom_headers": {}, "anth_thinking_config": {"type": "", "budget": 0, "effort": ""}, }, "xAI": { diff --git a/astrbot/core/provider/headers.py b/astrbot/core/provider/headers.py new file mode 100644 index 0000000000..b40d3f94ba --- /dev/null +++ b/astrbot/core/provider/headers.py @@ -0,0 +1,24 @@ +from astrbot import __version__ + +DEFAULT_USER_AGENT = f"astrbot/{__version__}" + + +def build_provider_headers(custom_headers: object = None) -> dict[str, str]: + """Build provider headers with an overridable AstrBot user agent. + + Args: + custom_headers: Optional header mapping from provider configuration. + + Returns: + A new header dictionary with string values and one User-Agent header. + """ + headers = {"User-Agent": DEFAULT_USER_AGENT} + if isinstance(custom_headers, dict): + for name, value in custom_headers.items(): + name, value = str(name), str(value) + if name.lower() == "user-agent": + if value.strip(): + headers["User-Agent"] = value + else: + headers[name] = value + return headers diff --git a/astrbot/core/provider/provider.py b/astrbot/core/provider/provider.py index 891bfdea9e..5e1274c24c 100644 --- a/astrbot/core/provider/provider.py +++ b/astrbot/core/provider/provider.py @@ -12,6 +12,7 @@ RerankResult, ToolCallsResult, ) +from astrbot.core.provider.headers import build_provider_headers from astrbot.core.provider.register import provider_cls_map from astrbot.core.utils.astrbot_path import get_astrbot_path @@ -31,6 +32,9 @@ def __init__(self, provider_config: dict) -> None: super().__init__() self.model_name = "" self.provider_config = provider_config + self.request_headers = build_provider_headers( + provider_config.get("custom_headers") + ) def set_model(self, model_name: str) -> None: """Set the current model name""" diff --git a/astrbot/core/provider/sources/anthropic_source.py b/astrbot/core/provider/sources/anthropic_source.py index c861ded6ba..5a21737b44 100644 --- a/astrbot/core/provider/sources/anthropic_source.py +++ b/astrbot/core/provider/sources/anthropic_source.py @@ -26,6 +26,7 @@ log_connection_failure, ) +from ..headers import build_provider_headers from ..register import register_provider_adapter from .request_retry import retry_provider_request, retry_provider_request_context @@ -71,12 +72,14 @@ def _resolve_custom_headers( *, required_headers: dict[str, str] | None = None, ) -> dict[str, str] | None: - merged_headers = cls._normalize_custom_headers(provider_config) or {} + merged_headers = build_provider_headers( + cls._normalize_custom_headers(provider_config) + ) 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 - return merged_headers or None + return merged_headers def __init__( self, diff --git a/astrbot/core/provider/sources/azure_tts_source.py b/astrbot/core/provider/sources/azure_tts_source.py index 34d8116cf5..6c039122d8 100644 --- a/astrbot/core/provider/sources/azure_tts_source.py +++ b/astrbot/core/provider/sources/azure_tts_source.py @@ -10,7 +10,7 @@ from httpx import AsyncClient, Timeout from astrbot import logger -from astrbot.core.config.default import VERSION +from astrbot.core.provider.headers import build_provider_headers from astrbot.core.utils.astrbot_path import get_astrbot_temp_path from astrbot.core.utils.datetime_utils import generate_timestamp_id @@ -25,6 +25,7 @@ class OTTSProvider: def __init__(self, config: dict) -> None: + self.request_headers = build_provider_headers(config.get("custom_headers")) self.skey = config["OTTS_SKEY"] self.api_url = config["OTTS_URL"] self.auth_time_url = config["OTTS_AUTH_TIME"] @@ -47,7 +48,9 @@ def client(self) -> AsyncClient: async def __aenter__(self): self._client = AsyncClient( - timeout=self.timeout, proxy=self.proxy if self.proxy else None + headers=self.request_headers, + timeout=self.timeout, + proxy=self.proxy if self.proxy else None, ) return self @@ -93,7 +96,7 @@ async def get_audio(self, text: str, voice_params: dict) -> str: "volume": voice_params["volume"], }, headers={ - "User-Agent": f"AstrBot/{VERSION}", + **self.request_headers, "UAK": "AstrBot/AzureTTS", }, ) @@ -148,7 +151,7 @@ def client(self) -> AsyncClient: async def __aenter__(self): self._client = AsyncClient( headers={ - "User-Agent": f"AstrBot/{VERSION}", + **self.request_headers, "Content-Type": "application/ssml+xml", "X-Microsoft-OutputFormat": "riff-48khz-16bit-mono-pcm", }, @@ -194,7 +197,7 @@ async def get_audio(self, text: str) -> str: content=ssml, headers={ "Authorization": f"Bearer {self.token}", - "User-Agent": f"AstrBot/{VERSION}", + **self.request_headers, }, ) response.raise_for_status() @@ -223,6 +226,7 @@ def _parse_provider( raise ValueError("无效的other[...]格式,应形如 other[{...}]") json_str = match.group(1).strip() otts_config = json.loads(json_str) + otts_config.setdefault("custom_headers", config.get("custom_headers")) required = {"OTTS_SKEY", "OTTS_URL", "OTTS_AUTH_TIME"} if missing := required - otts_config.keys(): raise ValueError(f"缺少OTTS参数: {', '.join(missing)}") diff --git a/astrbot/core/provider/sources/bailian_rerank_source.py b/astrbot/core/provider/sources/bailian_rerank_source.py index 030f59dd57..11c8334aac 100644 --- a/astrbot/core/provider/sources/bailian_rerank_source.py +++ b/astrbot/core/provider/sources/bailian_rerank_source.py @@ -70,7 +70,8 @@ def __init__(self, provider_config: dict, provider_settings: dict) -> None: } self.client = aiohttp.ClientSession( - headers=headers, timeout=aiohttp.ClientTimeout(total=self.timeout) + headers={**self.request_headers, **headers}, + timeout=aiohttp.ClientTimeout(total=self.timeout), ) # 设置模型名称 diff --git a/astrbot/core/provider/sources/dashscope_embedding_source.py b/astrbot/core/provider/sources/dashscope_embedding_source.py index 5502016697..56fd14fb38 100644 --- a/astrbot/core/provider/sources/dashscope_embedding_source.py +++ b/astrbot/core/provider/sources/dashscope_embedding_source.py @@ -69,7 +69,10 @@ async def get_embeddings(self, text: list[str]) -> list[list[float]]: or self.model.startswith("tongyi-embedding-vision") ) - kwargs: dict = {"base_address": self.base_url} + kwargs: dict = { + "base_address": self.base_url, + "headers": self.request_headers.copy(), + } if "embedding_dimensions" in self.provider_config: try: dimensions = int(self.provider_config["embedding_dimensions"]) diff --git a/astrbot/core/provider/sources/dashscope_tts.py b/astrbot/core/provider/sources/dashscope_tts.py index 6019ad0932..00c082b6b8 100644 --- a/astrbot/core/provider/sources/dashscope_tts.py +++ b/astrbot/core/provider/sources/dashscope_tts.py @@ -71,6 +71,7 @@ def _call_qwen_tts(self, model: str, text: str): kwargs = { "model": model, + "headers": self.request_headers.copy(), "messages": None, "api_key": self.chosen_api_key, "voice": self.voice or "Cherry", @@ -122,7 +123,9 @@ async def _download_audio_from_url(self, url: str) -> bytes | None: timeout = max(self.timeout_ms / 1000, 1) if self.timeout_ms else 20 try: async with ( - aiohttp.ClientSession() as session, + aiohttp.ClientSession( + headers={"User-Agent": self.request_headers["User-Agent"]} + ) as session, session.get( url, timeout=aiohttp.ClientTimeout(total=timeout), @@ -139,6 +142,9 @@ async def _synthesize_with_cosyvoice( text: str, ) -> tuple[bytes | None, str]: synthesizer = SpeechSynthesizer( + headers={ + name.lower(): value for name, value in self.request_headers.items() + }, model=model, voice=self.voice, format=AudioFormat.WAV_24000HZ_MONO_16BIT, diff --git a/astrbot/core/provider/sources/edge_tts_source.py b/astrbot/core/provider/sources/edge_tts_source.py index f4604be6a3..3cf39883ec 100644 --- a/astrbot/core/provider/sources/edge_tts_source.py +++ b/astrbot/core/provider/sources/edge_tts_source.py @@ -3,8 +3,10 @@ import subprocess import edge_tts +from edge_tts.constants import WSS_HEADERS from astrbot.core import logger +from astrbot.core.provider.headers import DEFAULT_USER_AGENT from astrbot.core.utils.astrbot_path import get_astrbot_temp_path from astrbot.core.utils.datetime_utils import generate_timestamp_id @@ -12,6 +14,9 @@ from ..provider import TTSProvider from ..register import register_provider_adapter +# Edge TTS exposes synthesis headers as a shared SDK default, not a client option. +WSS_HEADERS["User-Agent"] = DEFAULT_USER_AGENT + """ edge_tts 方式,能够免费、快速生成语音,使用需要先安装edge-tts库 ``` diff --git a/astrbot/core/provider/sources/elevenlabs_tts_source.py b/astrbot/core/provider/sources/elevenlabs_tts_source.py index f6977e8b21..3cfa36c313 100644 --- a/astrbot/core/provider/sources/elevenlabs_tts_source.py +++ b/astrbot/core/provider/sources/elevenlabs_tts_source.py @@ -119,6 +119,7 @@ def __init__( if proxy: logger.info(f"[ElevenLabs TTS] 使用代理: {proxy}") self.client = httpx.AsyncClient( + headers=self.request_headers, timeout=timeout, proxy=proxy or None, trust_env=False, diff --git a/astrbot/core/provider/sources/fishaudio_tts_api_source.py b/astrbot/core/provider/sources/fishaudio_tts_api_source.py index 2330f52128..7fe3aa2e9c 100644 --- a/astrbot/core/provider/sources/fishaudio_tts_api_source.py +++ b/astrbot/core/provider/sources/fishaudio_tts_api_source.py @@ -65,6 +65,7 @@ def __init__( if self.proxy: logger.info(f"[FishAudio TTS] 使用代理: {self.proxy}") self.headers = { + **self.request_headers, "Authorization": f"Bearer {self.chosen_api_key}", } # FishAudio API 要求 model 作为 HTTP header 发送,而非请求体字段 diff --git a/astrbot/core/provider/sources/gemini_embedding_source.py b/astrbot/core/provider/sources/gemini_embedding_source.py index 71e9dadc9d..0723d55f4c 100644 --- a/astrbot/core/provider/sources/gemini_embedding_source.py +++ b/astrbot/core/provider/sources/gemini_embedding_source.py @@ -24,7 +24,9 @@ def __init__(self, provider_config: dict, provider_settings: dict) -> None: api_base: str = provider_config["embedding_api_base"] timeout: int = int(provider_config.get("timeout", 20)) - http_options = types.HttpOptions(timeout=timeout * 1000) + http_options = types.HttpOptions( + timeout=timeout * 1000, headers=self.request_headers + ) if api_base: api_base = api_base.removesuffix("/") http_options.base_url = api_base @@ -34,6 +36,8 @@ def __init__(self, provider_config: dict, provider_settings: dict) -> None: logger.info(f"[Gemini Embedding] 使用代理: {proxy}") self.client = genai.Client(api_key=api_key, http_options=http_options).aio + # The SDK adds its own lower-case UA alongside our explicit header. + self.client._api_client._http_options.headers.pop("user-agent", None) self.model = provider_config.get( "embedding_model", diff --git a/astrbot/core/provider/sources/gemini_source.py b/astrbot/core/provider/sources/gemini_source.py index b6b7a97fb2..00268b2c0b 100644 --- a/astrbot/core/provider/sources/gemini_source.py +++ b/astrbot/core/provider/sources/gemini_source.py @@ -85,6 +85,7 @@ def _init_client(self) -> None: """初始化Gemini客户端""" proxy = self.provider_config.get("proxy", "") http_options = types.HttpOptions( + headers=self.request_headers, base_url=self.api_base, timeout=self.timeout * 1000, # 毫秒 ) @@ -114,6 +115,8 @@ def _init_client(self) -> None: api_key=self.chosen_api_key, http_options=http_options, ).aio + # The SDK adds its own lower-case UA alongside our explicit header. + self.client._api_client._http_options.headers.pop("user-agent", None) def _init_safety_settings(self) -> None: """初始化安全设置""" diff --git a/astrbot/core/provider/sources/gemini_tts_source.py b/astrbot/core/provider/sources/gemini_tts_source.py index 03b5405b82..136916aa1d 100644 --- a/astrbot/core/provider/sources/gemini_tts_source.py +++ b/astrbot/core/provider/sources/gemini_tts_source.py @@ -28,7 +28,9 @@ def __init__( api_key: str = provider_config.get("gemini_tts_api_key", "") api_base: str | None = provider_config.get("gemini_tts_api_base") timeout: int = int(provider_config.get("gemini_tts_timeout", 20)) - http_options = types.HttpOptions(timeout=timeout * 1000) + http_options = types.HttpOptions( + timeout=timeout * 1000, headers=self.request_headers + ) if api_base: api_base = api_base.removesuffix("/") @@ -39,6 +41,8 @@ def __init__( logger.info(f"[Gemini TTS] 使用代理: {proxy}") self.client = genai.Client(api_key=api_key, http_options=http_options).aio + # The SDK adds its own lower-case UA alongside our explicit header. + self.client._api_client._http_options.headers.pop("user-agent", None) self.model: str = provider_config.get( "gemini_tts_model", "gemini-2.5-flash-preview-tts", diff --git a/astrbot/core/provider/sources/gsv_selfhosted_source.py b/astrbot/core/provider/sources/gsv_selfhosted_source.py index f9c9076688..f4cac0ce1e 100644 --- a/astrbot/core/provider/sources/gsv_selfhosted_source.py +++ b/astrbot/core/provider/sources/gsv_selfhosted_source.py @@ -42,6 +42,7 @@ def __init__( async def initialize(self) -> None: """异步初始化:在 ProviderManager 中被调用""" self._session = aiohttp.ClientSession( + headers=self.request_headers, timeout=aiohttp.ClientTimeout(total=self.timeout), ) try: diff --git a/astrbot/core/provider/sources/gsvi_tts_source.py b/astrbot/core/provider/sources/gsvi_tts_source.py index f875ecb4b2..d8ae84aa8f 100644 --- a/astrbot/core/provider/sources/gsvi_tts_source.py +++ b/astrbot/core/provider/sources/gsvi_tts_source.py @@ -50,7 +50,7 @@ async def get_audio(self, text: str) -> str: "text_lang": self.text_lang, } - async with aiohttp.ClientSession() as session: + async with aiohttp.ClientSession(headers=self.request_headers) as session: async with session.post(url, json=data, headers=headers) as response: if response.status == 200: resp_json = await response.json() diff --git a/astrbot/core/provider/sources/kimi_code_source.py b/astrbot/core/provider/sources/kimi_code_source.py index 02c200271f..bac18ce190 100644 --- a/astrbot/core/provider/sources/kimi_code_source.py +++ b/astrbot/core/provider/sources/kimi_code_source.py @@ -1,9 +1,10 @@ +from ..headers import DEFAULT_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 = DEFAULT_USER_AGENT @register_provider_adapter( diff --git a/astrbot/core/provider/sources/mimo_stt_api_source.py b/astrbot/core/provider/sources/mimo_stt_api_source.py index 8417e7a988..75117a6be8 100644 --- a/astrbot/core/provider/sources/mimo_stt_api_source.py +++ b/astrbot/core/provider/sources/mimo_stt_api_source.py @@ -84,7 +84,7 @@ async def get_text(self, audio_url: str) -> str: try: response = await self.client.post( build_api_url(self.api_base), - headers=build_headers(self.chosen_api_key), + headers={**build_headers(self.chosen_api_key), **self.request_headers}, json=payload, ) try: diff --git a/astrbot/core/provider/sources/mimo_tts_api_source.py b/astrbot/core/provider/sources/mimo_tts_api_source.py index a9ccaae340..0182615740 100644 --- a/astrbot/core/provider/sources/mimo_tts_api_source.py +++ b/astrbot/core/provider/sources/mimo_tts_api_source.py @@ -103,7 +103,7 @@ def _build_payload(self, text: str) -> dict: async def get_audio(self, text: str) -> str: response = await self.client.post( build_api_url(self.api_base), - headers=build_headers(self.chosen_api_key), + headers={**build_headers(self.chosen_api_key), **self.request_headers}, json=self._build_payload(text), ) diff --git a/astrbot/core/provider/sources/minimax_token_plan_source.py b/astrbot/core/provider/sources/minimax_token_plan_source.py index 8d86c77b73..b16c7e4dc2 100644 --- a/astrbot/core/provider/sources/minimax_token_plan_source.py +++ b/astrbot/core/provider/sources/minimax_token_plan_source.py @@ -47,7 +47,7 @@ async def get_models(self) -> list[str]: logger.warning("No API key configured for MiniMax Token Plan.") return [] try: - async with httpx.AsyncClient() as client: + async with httpx.AsyncClient(headers=self.request_headers) as client: resp = await client.get( "https://api.minimaxi.com/v1/models", headers={"Authorization": f"Bearer {key}"}, diff --git a/astrbot/core/provider/sources/minimax_tts_api_source.py b/astrbot/core/provider/sources/minimax_tts_api_source.py index f86e70b3c6..d815a48f80 100644 --- a/astrbot/core/provider/sources/minimax_tts_api_source.py +++ b/astrbot/core/provider/sources/minimax_tts_api_source.py @@ -103,7 +103,7 @@ async def _call_tts_stream(self, text: str) -> AsyncIterator[str]: """进行流式请求""" try: async with ( - aiohttp.ClientSession() as session, + aiohttp.ClientSession(headers=self.request_headers) as session, session.post( self.concat_base_url, headers=self.headers, diff --git a/astrbot/core/provider/sources/nvidia_embedding_source.py b/astrbot/core/provider/sources/nvidia_embedding_source.py index 475fadeecc..d650c68b4a 100644 --- a/astrbot/core/provider/sources/nvidia_embedding_source.py +++ b/astrbot/core/provider/sources/nvidia_embedding_source.py @@ -49,7 +49,7 @@ async def _get_client(self): } timeout = aiohttp.ClientTimeout(total=self.timeout) self.client = aiohttp.ClientSession( - headers=headers, + headers={**self.request_headers, **headers}, timeout=timeout, ) return self.client diff --git a/astrbot/core/provider/sources/nvidia_rerank_source.py b/astrbot/core/provider/sources/nvidia_rerank_source.py index 6f3f92c1cd..4d241a6545 100644 --- a/astrbot/core/provider/sources/nvidia_rerank_source.py +++ b/astrbot/core/provider/sources/nvidia_rerank_source.py @@ -37,7 +37,8 @@ async def _get_client(self): "Accept": "application/json", } self.client = aiohttp.ClientSession( - headers=headers, timeout=aiohttp.ClientTimeout(total=self.timeout) + headers={**self.request_headers, **headers}, + timeout=aiohttp.ClientTimeout(total=self.timeout), ) return self.client diff --git a/astrbot/core/provider/sources/ollama_embedding_source.py b/astrbot/core/provider/sources/ollama_embedding_source.py index 8982fc51de..411310304f 100644 --- a/astrbot/core/provider/sources/ollama_embedding_source.py +++ b/astrbot/core/provider/sources/ollama_embedding_source.py @@ -42,7 +42,7 @@ async def _get_client(self): } timeout = aiohttp.ClientTimeout(total=self.timeout) self.client = aiohttp.ClientSession( - headers=headers, + headers={**self.request_headers, **headers}, timeout=timeout, ) return self.client diff --git a/astrbot/core/provider/sources/openai_embedding_source.py b/astrbot/core/provider/sources/openai_embedding_source.py index f1e5bc880e..8f29017f4e 100644 --- a/astrbot/core/provider/sources/openai_embedding_source.py +++ b/astrbot/core/provider/sources/openai_embedding_source.py @@ -39,6 +39,7 @@ def __init__(self, provider_config: dict, provider_settings: dict) -> None: ) logger.info(f"[OpenAI Embedding] {provider_id} Using API Base: {api_base}") self.client = AsyncOpenAI( + default_headers=self.request_headers, api_key=provider_config.get("embedding_api_key"), base_url=api_base, timeout=int(provider_config.get("timeout", 20)), diff --git a/astrbot/core/provider/sources/openai_source.py b/astrbot/core/provider/sources/openai_source.py index f7870b7137..0bff46406e 100644 --- a/astrbot/core/provider/sources/openai_source.py +++ b/astrbot/core/provider/sources/openai_source.py @@ -359,16 +359,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.request_headers 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( diff --git a/astrbot/core/provider/sources/openai_tts_api_source.py b/astrbot/core/provider/sources/openai_tts_api_source.py index 27485765fd..6ea2bfba8f 100644 --- a/astrbot/core/provider/sources/openai_tts_api_source.py +++ b/astrbot/core/provider/sources/openai_tts_api_source.py @@ -37,6 +37,7 @@ def __init__( logger.info(f"[OpenAI TTS] 使用代理: {proxy}") http_client = httpx.AsyncClient(proxy=proxy) self.client = AsyncOpenAI( + default_headers=self.request_headers, api_key=self.chosen_api_key, base_url=provider_config.get("api_base"), timeout=timeout, diff --git a/astrbot/core/provider/sources/tei_rerank_source.py b/astrbot/core/provider/sources/tei_rerank_source.py index 9a5d58aaf7..a455c11e34 100644 --- a/astrbot/core/provider/sources/tei_rerank_source.py +++ b/astrbot/core/provider/sources/tei_rerank_source.py @@ -31,7 +31,7 @@ def __init__(self, provider_config: dict, provider_settings: dict) -> None: self.raw_scores = provider_config.get("tei_rerank_raw_scores", False) self.return_text = provider_config.get("tei_rerank_return_text", False) - h = {} + h = self.request_headers.copy() if self.api_key: h["Authorization"] = f"Bearer {self.api_key}" self.client = aiohttp.ClientSession( diff --git a/astrbot/core/provider/sources/vllm_rerank_source.py b/astrbot/core/provider/sources/vllm_rerank_source.py index 6fd44bc589..104fb62858 100644 --- a/astrbot/core/provider/sources/vllm_rerank_source.py +++ b/astrbot/core/provider/sources/vllm_rerank_source.py @@ -26,7 +26,7 @@ def __init__(self, provider_config: dict, provider_settings: dict) -> None: self.timeout = provider_config.get("timeout", 20) self.model = provider_config.get("rerank_model", "BAAI/bge-reranker-base") - h = {} + h = self.request_headers.copy() if self.auth_key: h["Authorization"] = f"Bearer {self.auth_key}" self.client = aiohttp.ClientSession( diff --git a/astrbot/core/provider/sources/volcengine_tts.py b/astrbot/core/provider/sources/volcengine_tts.py index 9e1453df6a..24c4f45be1 100644 --- a/astrbot/core/provider/sources/volcengine_tts.py +++ b/astrbot/core/provider/sources/volcengine_tts.py @@ -75,7 +75,7 @@ async def get_audio(self, text: str) -> str: try: async with ( - aiohttp.ClientSession() as session, + aiohttp.ClientSession(headers=self.request_headers) as session, session.post( self.api_base, data=json.dumps(payload), diff --git a/astrbot/core/provider/sources/whisper_api_source.py b/astrbot/core/provider/sources/whisper_api_source.py index 60aee070b7..cbb06fc9c0 100644 --- a/astrbot/core/provider/sources/whisper_api_source.py +++ b/astrbot/core/provider/sources/whisper_api_source.py @@ -22,6 +22,7 @@ def __init__( self.chosen_api_key = provider_config.get("api_key", "") self.client = AsyncOpenAI( + default_headers=self.request_headers, api_key=self.chosen_api_key, base_url=provider_config.get("api_base"), timeout=provider_config.get("timeout", NOT_GIVEN), diff --git a/astrbot/core/provider/sources/xinference_rerank_source.py b/astrbot/core/provider/sources/xinference_rerank_source.py index 9c3a77c158..88e9fc8522 100644 --- a/astrbot/core/provider/sources/xinference_rerank_source.py +++ b/astrbot/core/provider/sources/xinference_rerank_source.py @@ -44,6 +44,7 @@ async def initialize(self) -> None: else: logger.info("Xinference Rerank: No API key provided.") self.client = Client(self.base_url) + self.client._headers.update(self.request_headers) try: running_models = await self.client.list_models() diff --git a/astrbot/core/provider/sources/xinference_stt_provider.py b/astrbot/core/provider/sources/xinference_stt_provider.py index f0e905e6d8..20729cb4ec 100644 --- a/astrbot/core/provider/sources/xinference_stt_provider.py +++ b/astrbot/core/provider/sources/xinference_stt_provider.py @@ -40,6 +40,7 @@ async def initialize(self) -> None: else: logger.info("Xinference STT: No API key provided.") self.client = Client(self.base_url) + self.client._headers.update(self.request_headers) try: running_models = await self.client.list_models() diff --git a/tests/test_provider_user_agent.py b/tests/test_provider_user_agent.py new file mode 100644 index 0000000000..792bd009a8 --- /dev/null +++ b/tests/test_provider_user_agent.py @@ -0,0 +1,223 @@ +import copy + +import pytest +import pytest_asyncio +from aiohttp import web + +from astrbot import __version__ +from astrbot.core.provider.headers import DEFAULT_USER_AGENT, build_provider_headers +from astrbot.core.provider.sources.anthropic_source import ProviderAnthropic +from astrbot.core.provider.sources.bailian_rerank_source import BailianRerankProvider +from astrbot.core.provider.sources.dashscope_embedding_source import ( + DashScopeEmbeddingProvider, +) +from astrbot.core.provider.sources.gemini_embedding_source import ( + GeminiEmbeddingProvider, +) +from astrbot.core.provider.sources.gemini_source import ProviderGoogleGenAI +from astrbot.core.provider.sources.gemini_tts_source import ProviderGeminiTTSAPI +from astrbot.core.provider.sources.kimi_code_source import ProviderKimiCode +from astrbot.core.provider.sources.nvidia_embedding_source import ( + NvidiaEmbeddingProvider, +) +from astrbot.core.provider.sources.nvidia_rerank_source import NvidiaRerankProvider +from astrbot.core.provider.sources.ollama_embedding_source import ( + OllamaEmbeddingProvider, +) +from astrbot.core.provider.sources.openai_embedding_source import ( + OpenAIEmbeddingProvider, +) +from astrbot.core.provider.sources.openai_responses_source import ( + ProviderOpenAIResponses, +) +from astrbot.core.provider.sources.openai_source import ProviderOpenAIOfficial +from astrbot.core.provider.sources.openai_tts_api_source import ProviderOpenAITTSAPI +from astrbot.core.provider.sources.tei_rerank_source import TEIRerankProvider +from astrbot.core.provider.sources.vllm_rerank_source import VLLMRerankProvider +from astrbot.core.provider.sources.whisper_api_source import ProviderOpenAIWhisperAPI + + +@pytest.mark.parametrize( + "custom_headers", [None, {}, [], "invalid", {"User-Agent": " "}] +) +def test_provider_headers_default_to_current_version(custom_headers): + assert build_provider_headers(custom_headers) == { + "User-Agent": f"astrbot/{__version__}" + } + + +@pytest.mark.parametrize("name", ["User-Agent", "user-agent", "USER-AGENT"]) +def test_provider_headers_preserve_custom_values_without_mutation(name): + custom = {name: "custom/1.0", "X-Trace-Id": 123} + original = copy.deepcopy(custom) + assert build_provider_headers(custom) == { + "User-Agent": "custom/1.0", + "X-Trace-Id": "123", + } + assert custom == original + + +@pytest_asyncio.fixture +async def provider_http_server(unused_tcp_port, monkeypatch): + """Capture actual SDK requests without contacting external providers.""" + requests = [] + monkeypatch.setenv("NO_PROXY", "127.0.0.1") + monkeypatch.setenv("no_proxy", "127.0.0.1") + + async def handle(request): + requests.append(request.headers) + return web.json_response( + { + "object": "list", + "data": [], + "models": [], + "output": {"embeddings": [{"embedding": [0.1], "text_index": 0}]}, + } + ) + + app = web.Application() + app.router.add_route("*", "/{path:.*}", handle) + runner = web.AppRunner(app) + await runner.setup() + try: + await web.TCPSite(runner, "127.0.0.1", unused_tcp_port).start() + yield f"http://127.0.0.1:{unused_tcp_port}", requests + finally: + await runner.cleanup() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "provider_cls", + [ + ProviderOpenAIOfficial, + ProviderOpenAIResponses, + OpenAIEmbeddingProvider, + ProviderOpenAITTSAPI, + ProviderOpenAIWhisperAPI, + ProviderAnthropic, + ProviderKimiCode, + ProviderGoogleGenAI, + GeminiEmbeddingProvider, + ProviderGeminiTTSAPI, + ], +) +@pytest.mark.parametrize("custom_headers", [{}, {"user-agent": "custom/1.0"}]) +async def test_provider_sdk_sends_exactly_one_user_agent( + provider_cls, custom_headers, provider_http_server +): + base_url, requests = provider_http_server + config = { + "id": "test-provider", + "model": "test-model", + "key": ["test-key"], + "api_key": "test-key", + "api_base": base_url, + "embedding_api_key": "test-key", + "embedding_api_base": base_url, + "gemini_tts_api_key": "test-key", + "gemini_tts_api_base": base_url, + "custom_headers": custom_headers, + } + original = copy.deepcopy(config) + provider = provider_cls(config, {}) + try: + await provider.client.models.list() + assert len(requests) == 1 + assert requests[0].getall("User-Agent") == [ + custom_headers.get("user-agent", DEFAULT_USER_AGENT) + ] + assert config == original + finally: + await provider.terminate() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "provider_cls", + [ + BailianRerankProvider, + TEIRerankProvider, + VLLMRerankProvider, + NvidiaEmbeddingProvider, + NvidiaRerankProvider, + OllamaEmbeddingProvider, + ], +) +@pytest.mark.parametrize("custom_headers", [{}, {"USER-AGENT": "custom/1.0"}]) +async def test_aiohttp_provider_sends_user_agent( + provider_cls, custom_headers, provider_http_server +): + base_url, requests = provider_http_server + provider = provider_cls( + { + "embedding_api_key": "test-key", + "embedding_api_base": base_url, + "rerank_api_key": "test-key", + "rerank_api_base": base_url, + "custom_headers": custom_headers, + }, + {}, + ) + try: + client = provider.client or await provider._get_client() + async with client.get(base_url) as response: + assert response.status == 200 + assert requests[0].getall("User-Agent") == [ + custom_headers.get("USER-AGENT", DEFAULT_USER_AGENT) + ] + finally: + await provider.terminate() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model", ["text-embedding-v4", "qwen3-vl-embedding"]) +@pytest.mark.parametrize("custom_headers", [{}, {"User-Agent": "custom/1.0"}]) +async def test_dashscope_sends_user_agent(model, custom_headers, provider_http_server): + base_url, requests = provider_http_server + provider = DashScopeEmbeddingProvider( + { + "embedding_api_key": "test-key", + "embedding_api_base": base_url, + "embedding_model": model, + "custom_headers": custom_headers, + }, + {}, + ) + assert await provider.get_embeddings(["hello"]) == [[0.1]] + assert requests[0].getall("User-Agent") == [ + custom_headers.get("User-Agent", DEFAULT_USER_AGENT) + ] + + +@pytest.mark.asyncio +async def test_edge_tts_synthesis_uses_astrbot_user_agent(monkeypatch, unused_tcp_port): + edge_tts = pytest.importorskip("edge_tts") + from astrbot.core.provider.sources import edge_tts_source # noqa: F401 + + requests = [] + + async def handle(request): + requests.append(request.headers) + websocket = web.WebSocketResponse() + await websocket.prepare(request) + await websocket.receive() + await websocket.receive() + await websocket.close() + return websocket + + app = web.Application() + app.router.add_get("/", handle) + runner = web.AppRunner(app) + await runner.setup() + monkeypatch.setattr( + edge_tts.communicate, "WSS_URL", f"ws://127.0.0.1:{unused_tcp_port}/?test=1" + ) + try: + await web.TCPSite(runner, "127.0.0.1", unused_tcp_port).start() + with pytest.raises(edge_tts.exceptions.NoAudioReceived): + async for _ in edge_tts.Communicate("hello", "en-US-AriaNeural").stream(): + pass + assert requests[0].getall("User-Agent") == [DEFAULT_USER_AGENT] + finally: + await runner.cleanup()