mirror of
https://github.com/farion1231/cc-switch.git
synced 2026-07-28 08:44:41 +08:00
Jason
8b925c2f2f
Extract a shared `parse_custom_user_agent` helper in provider.rs returning `Result<Option<HeaderValue>>`, and reuse it in the forwarder, stream check, and model fetch paths so detection, forwarding, and model listing all apply the same provider-level User-Agent. Previously only the forwarder honored it, so stream check could fail (or model listing 403) on UA-gated upstreams that the proxy itself handled fine. - stream_check injects the provider's custom UA on the claude/codex paths and still skips the GitHub Copilot fingerprint UA. - model_fetch service + command and the model-fetch.ts wrapper thread an optional UA through to GET /v1/models. - runtime callers silently ignore invalid values via `.ok().flatten()` (no save-time block, so deeplink imports stay lenient).
478 lines
16 KiB
Rust
478 lines
16 KiB
Rust
//! 模型列表获取服务
|
||
//!
|
||
//! 通过 OpenAI 兼容的 GET /v1/models 端点获取供应商可用模型列表。
|
||
//! 主要面向第三方聚合站（硅基流动、OpenRouter 等），以及把 Anthropic
|
||
//! 协议挂在兼容子路径上的官方供应商（DeepSeek、Kimi、智谱 GLM 等）。
|
||
|
||
use reqwest::header::{HeaderValue, USER_AGENT};
|
||
use reqwest::StatusCode;
|
||
use serde::{Deserialize, Serialize};
|
||
use std::time::Duration;
|
||
|
||
/// 获取到的模型信息
|
||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||
#[serde(rename_all = "camelCase")]
|
||
pub struct FetchedModel {
|
||
pub id: String,
|
||
pub owned_by: Option<String>,
|
||
}
|
||
|
||
/// OpenAI 兼容的 /v1/models 响应格式
|
||
#[derive(Debug, Deserialize)]
|
||
struct ModelsResponse {
|
||
data: Option<Vec<ModelEntry>>,
|
||
}
|
||
|
||
#[derive(Debug, Deserialize)]
|
||
struct ModelEntry {
|
||
id: String,
|
||
owned_by: Option<String>,
|
||
}
|
||
|
||
const FETCH_TIMEOUT_SECS: u64 = 15;
|
||
|
||
/// 404/405 响应体截断长度：避免把几十 KB HTML 404 页整页保留到错误串里。
|
||
const ERROR_BODY_MAX_CHARS: usize = 512;
|
||
|
||
/// 已知的「Anthropic 协议兼容子路径」后缀；按长度降序，最长前缀优先匹配。
|
||
/// baseURL 命中这些后缀时，候选列表会追加「剥离后缀再拼 /v1/models / /models」的版本。
|
||
const KNOWN_COMPAT_SUFFIXES: &[&str] = &[
|
||
"/api/claudecode",
|
||
"/api/anthropic",
|
||
"/apps/anthropic",
|
||
"/api/coding",
|
||
"/claudecode",
|
||
"/anthropic",
|
||
"/step_plan",
|
||
"/coding",
|
||
"/claude",
|
||
];
|
||
|
||
/// 获取供应商的可用模型列表
|
||
///
|
||
/// 使用 OpenAI 兼容的 GET /v1/models 端点，按候选列表顺序尝试。
|
||
pub async fn fetch_models(
|
||
base_url: &str,
|
||
api_key: &str,
|
||
is_full_url: bool,
|
||
models_url_override: Option<&str>,
|
||
user_agent: Option<HeaderValue>,
|
||
) -> Result<Vec<FetchedModel>, String> {
|
||
if api_key.is_empty() {
|
||
return Err("API Key is required to fetch models".to_string());
|
||
}
|
||
|
||
let candidates = build_models_url_candidates(base_url, is_full_url, models_url_override)?;
|
||
let client = crate::proxy::http_client::get();
|
||
let mut last_err: Option<String> = None;
|
||
|
||
for url in &candidates {
|
||
log::debug!("[ModelFetch] Trying endpoint: {url}");
|
||
let mut request = client
|
||
.get(url)
|
||
.header("Authorization", format!("Bearer {api_key}"))
|
||
.timeout(Duration::from_secs(FETCH_TIMEOUT_SECS));
|
||
// 自定义 User-Agent：部分 /models 端点同样有 UA 白名单（如 Kimi Coding Plan），
|
||
// 与转发 / 检测路径共用同一 UA，避免"代理可用但取模型失败"。
|
||
if let Some(ua) = &user_agent {
|
||
request = request.header(USER_AGENT, ua.clone());
|
||
}
|
||
let response = match request.send().await {
|
||
Ok(r) => r,
|
||
Err(e) => {
|
||
return Err(format!("Request failed: {e}"));
|
||
}
|
||
};
|
||
|
||
let status = response.status();
|
||
|
||
if status.is_success() {
|
||
let resp: ModelsResponse = response
|
||
.json()
|
||
.await
|
||
.map_err(|e| format!("Failed to parse response: {e}"))?;
|
||
|
||
let mut models: Vec<FetchedModel> = resp
|
||
.data
|
||
.unwrap_or_default()
|
||
.into_iter()
|
||
.map(|m| FetchedModel {
|
||
id: m.id,
|
||
owned_by: m.owned_by,
|
||
})
|
||
.collect();
|
||
|
||
models.sort_by(|a, b| a.id.cmp(&b.id));
|
||
return Ok(models);
|
||
}
|
||
|
||
if status == StatusCode::NOT_FOUND || status == StatusCode::METHOD_NOT_ALLOWED {
|
||
let body = truncate_body(response.text().await.unwrap_or_default());
|
||
last_err = Some(format!("HTTP {status}: {body}"));
|
||
continue;
|
||
}
|
||
|
||
let body = truncate_body(response.text().await.unwrap_or_default());
|
||
return Err(format!("HTTP {status}: {body}"));
|
||
}
|
||
|
||
Err(format!(
|
||
"All candidates failed: {}",
|
||
last_err.unwrap_or_else(|| "no candidates".to_string())
|
||
))
|
||
}
|
||
|
||
/// 构造「模型列表端点」的候选 URL 列表
|
||
///
|
||
/// 候选顺序：
|
||
/// 1. `models_url_override` 非空 → 只返回它
|
||
/// 2. baseURL 拼 `/v1/models`；若已以版本段 `/v{N}` 结尾（`/v1`、智谱
|
||
/// `/api/coding/paas/v4` 等），版本号已在路径里，改拼 `/models`
|
||
/// 3. 版本段非 `/v1`（如 `/v4`）时再追加 `/v1/models` 作为兜底次候选
|
||
/// 4. 若 baseURL 命中 [`KNOWN_COMPAT_SUFFIXES`]，剥离后缀再拼 `/v1/models`、`/models`
|
||
///
|
||
/// 结果已去重且保持首次出现顺序。
|
||
pub fn build_models_url_candidates(
|
||
base_url: &str,
|
||
is_full_url: bool,
|
||
models_url_override: Option<&str>,
|
||
) -> Result<Vec<String>, String> {
|
||
if let Some(raw) = models_url_override {
|
||
let trimmed = raw.trim();
|
||
if !trimmed.is_empty() {
|
||
return Ok(vec![trimmed.to_string()]);
|
||
}
|
||
}
|
||
|
||
let trimmed = base_url.trim().trim_end_matches('/');
|
||
if trimmed.is_empty() {
|
||
return Err("Base URL is empty".to_string());
|
||
}
|
||
|
||
let mut candidates: Vec<String> = Vec::new();
|
||
|
||
if is_full_url {
|
||
if let Some(idx) = trimmed.find("/v1/") {
|
||
candidates.push(format!("{}/v1/models", &trimmed[..idx]));
|
||
} else if let Some(idx) = trimmed.rfind('/') {
|
||
let root = &trimmed[..idx];
|
||
if root.contains("://") && root.len() > root.find("://").unwrap() + 3 {
|
||
candidates.push(format!("{root}/v1/models"));
|
||
}
|
||
}
|
||
if candidates.is_empty() {
|
||
return Err("Cannot derive models endpoint from full URL".to_string());
|
||
}
|
||
return Ok(candidates);
|
||
}
|
||
|
||
// baseURL 已以版本段 /v{N} 结尾时（如 `/v1`、智谱 `/api/coding/paas/v4`），
|
||
// OpenAI 惯例的模型端点是 `{base}/models`，不能再补 `/v1`
|
||
// （否则 .../coding/paas/v4/v1/models → 404）。
|
||
if ends_with_version_segment(trimmed) {
|
||
candidates.push(format!("{trimmed}/models"));
|
||
// 版本段非 /v1 时，保留旧的 /v1/models 作为兜底次候选（正确路径已在前）。
|
||
if !trimmed.ends_with("/v1") {
|
||
candidates.push(format!("{trimmed}/v1/models"));
|
||
}
|
||
} else {
|
||
candidates.push(format!("{trimmed}/v1/models"));
|
||
}
|
||
|
||
if let Some(stripped) = strip_compat_suffix(trimmed) {
|
||
let root = stripped.trim_end_matches('/');
|
||
if !root.is_empty() && root.contains("://") {
|
||
candidates.push(format!("{root}/v1/models"));
|
||
candidates.push(format!("{root}/models"));
|
||
}
|
||
}
|
||
|
||
// 候选最多 3 条，线性去重即可，不值得上 HashSet。
|
||
let mut unique: Vec<String> = Vec::with_capacity(candidates.len());
|
||
for url in candidates {
|
||
if !unique.iter().any(|u| u == &url) {
|
||
unique.push(url);
|
||
}
|
||
}
|
||
|
||
Ok(unique)
|
||
}
|
||
|
||
/// 截断响应体到 [`ERROR_BODY_MAX_CHARS`] 字符，避免 HTML 404 页占用错误串。
|
||
fn truncate_body(body: String) -> String {
|
||
if body.chars().count() <= ERROR_BODY_MAX_CHARS {
|
||
body
|
||
} else {
|
||
let mut s: String = body.chars().take(ERROR_BODY_MAX_CHARS).collect();
|
||
s.push('…');
|
||
s
|
||
}
|
||
}
|
||
|
||
/// 若 baseURL 以任一已知兼容子路径结尾，返回剥离后的剩余部分；否则 `None`。
|
||
///
|
||
/// 依赖 [`KNOWN_COMPAT_SUFFIXES`] 按长度降序排列，确保最长前缀优先命中
|
||
/// （否则 `/anthropic` 会提前匹配掉 `/api/anthropic` 的场景）。
|
||
fn strip_compat_suffix(base_url: &str) -> Option<&str> {
|
||
for suffix in KNOWN_COMPAT_SUFFIXES {
|
||
if base_url.ends_with(*suffix) {
|
||
return Some(&base_url[..base_url.len() - suffix.len()]);
|
||
}
|
||
}
|
||
None
|
||
}
|
||
|
||
/// 判断 baseURL 是否以 OpenAI 风格的版本段 `/v{N}` 结尾（`N` 为一个或多个数字），
|
||
/// 例如 `/v1`、`.../paas/v4`。这类 URL 版本号已在路径中，模型端点应为
|
||
/// `{base}/models`，不能再补 `/v1`（智谱 Coding Plan 即 `.../coding/paas/v4`）。
|
||
fn ends_with_version_segment(url: &str) -> bool {
|
||
let last = url.rsplit('/').next().unwrap_or("");
|
||
last.strip_prefix('v')
|
||
.is_some_and(|digits| !digits.is_empty() && digits.bytes().all(|b| b.is_ascii_digit()))
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
|
||
#[test]
|
||
fn test_candidates_plain_root() {
|
||
let c = build_models_url_candidates("https://api.siliconflow.cn", false, None).unwrap();
|
||
assert_eq!(c, vec!["https://api.siliconflow.cn/v1/models"]);
|
||
}
|
||
|
||
#[test]
|
||
fn test_candidates_trailing_slash() {
|
||
let c = build_models_url_candidates("https://api.example.com/", false, None).unwrap();
|
||
assert_eq!(c, vec!["https://api.example.com/v1/models"]);
|
||
}
|
||
|
||
#[test]
|
||
fn test_candidates_with_v1() {
|
||
let c = build_models_url_candidates("https://api.example.com/v1", false, None).unwrap();
|
||
assert_eq!(c, vec!["https://api.example.com/v1/models"]);
|
||
}
|
||
|
||
#[test]
|
||
fn test_candidates_zhipu_coding_paas_v4() {
|
||
// 智谱 Coding Plan 端点以 /v4 版本段结尾：模型端点是 {base}/models，
|
||
// 正确路径必须排在 .../v4/v1/models（404）之前。
|
||
let c =
|
||
build_models_url_candidates("https://open.bigmodel.cn/api/coding/paas/v4", false, None)
|
||
.unwrap();
|
||
assert_eq!(
|
||
c,
|
||
vec![
|
||
"https://open.bigmodel.cn/api/coding/paas/v4/models",
|
||
"https://open.bigmodel.cn/api/coding/paas/v4/v1/models",
|
||
]
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn test_candidates_zai_coding_paas_v4() {
|
||
let c = build_models_url_candidates("https://api.z.ai/api/coding/paas/v4", false, None)
|
||
.unwrap();
|
||
assert_eq!(
|
||
c,
|
||
vec![
|
||
"https://api.z.ai/api/coding/paas/v4/models",
|
||
"https://api.z.ai/api/coding/paas/v4/v1/models",
|
||
]
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn test_ends_with_version_segment() {
|
||
assert!(ends_with_version_segment("https://x.com/v1"));
|
||
assert!(ends_with_version_segment(
|
||
"https://open.bigmodel.cn/api/coding/paas/v4"
|
||
));
|
||
assert!(ends_with_version_segment("https://x.com/v10"));
|
||
assert!(!ends_with_version_segment("https://x.com/api"));
|
||
assert!(!ends_with_version_segment("https://x.com/vX"));
|
||
assert!(!ends_with_version_segment("https://x.com/models"));
|
||
assert!(!ends_with_version_segment("https://api.siliconflow.cn"));
|
||
}
|
||
|
||
#[test]
|
||
fn test_candidates_full_url() {
|
||
let c = build_models_url_candidates(
|
||
"https://proxy.example.com/v1/chat/completions",
|
||
true,
|
||
None,
|
||
)
|
||
.unwrap();
|
||
assert_eq!(c, vec!["https://proxy.example.com/v1/models"]);
|
||
}
|
||
|
||
#[test]
|
||
fn test_candidates_empty() {
|
||
assert!(build_models_url_candidates("", false, None).is_err());
|
||
}
|
||
|
||
#[test]
|
||
fn test_candidates_override_returns_single() {
|
||
let c = build_models_url_candidates(
|
||
"https://api.deepseek.com/anthropic",
|
||
false,
|
||
Some("https://api.deepseek.com/models"),
|
||
)
|
||
.unwrap();
|
||
assert_eq!(c, vec!["https://api.deepseek.com/models"]);
|
||
}
|
||
|
||
#[test]
|
||
fn test_candidates_override_empty_falls_through() {
|
||
let c =
|
||
build_models_url_candidates("https://api.siliconflow.cn", false, Some(" ")).unwrap();
|
||
assert_eq!(c, vec!["https://api.siliconflow.cn/v1/models"]);
|
||
}
|
||
|
||
#[test]
|
||
fn test_candidates_deepseek_strip_anthropic() {
|
||
let c =
|
||
build_models_url_candidates("https://api.deepseek.com/anthropic", false, None).unwrap();
|
||
assert_eq!(
|
||
c,
|
||
vec![
|
||
"https://api.deepseek.com/anthropic/v1/models",
|
||
"https://api.deepseek.com/v1/models",
|
||
"https://api.deepseek.com/models",
|
||
]
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn test_candidates_zhipu_strip_api_anthropic() {
|
||
let c = build_models_url_candidates("https://open.bigmodel.cn/api/anthropic", false, None)
|
||
.unwrap();
|
||
assert_eq!(
|
||
c,
|
||
vec![
|
||
"https://open.bigmodel.cn/api/anthropic/v1/models",
|
||
"https://open.bigmodel.cn/v1/models",
|
||
"https://open.bigmodel.cn/models",
|
||
]
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn test_candidates_bailian_strip_apps_anthropic() {
|
||
let c = build_models_url_candidates(
|
||
"https://dashscope.aliyuncs.com/apps/anthropic",
|
||
false,
|
||
None,
|
||
)
|
||
.unwrap();
|
||
assert_eq!(
|
||
c,
|
||
vec![
|
||
"https://dashscope.aliyuncs.com/apps/anthropic/v1/models",
|
||
"https://dashscope.aliyuncs.com/v1/models",
|
||
"https://dashscope.aliyuncs.com/models",
|
||
]
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn test_candidates_stepfun_strip_step_plan() {
|
||
let c =
|
||
build_models_url_candidates("https://api.stepfun.com/step_plan", false, None).unwrap();
|
||
assert_eq!(
|
||
c,
|
||
vec![
|
||
"https://api.stepfun.com/step_plan/v1/models",
|
||
"https://api.stepfun.com/v1/models",
|
||
"https://api.stepfun.com/models",
|
||
]
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn test_candidates_doubao_strip_api_coding() {
|
||
let c = build_models_url_candidates(
|
||
"https://ark.cn-beijing.volces.com/api/coding",
|
||
false,
|
||
None,
|
||
)
|
||
.unwrap();
|
||
assert_eq!(
|
||
c,
|
||
vec![
|
||
"https://ark.cn-beijing.volces.com/api/coding/v1/models",
|
||
"https://ark.cn-beijing.volces.com/v1/models",
|
||
"https://ark.cn-beijing.volces.com/models",
|
||
]
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn test_candidates_rightcode_strip_claude() {
|
||
let c = build_models_url_candidates("https://www.right.codes/claude", false, None).unwrap();
|
||
assert_eq!(
|
||
c,
|
||
vec![
|
||
"https://www.right.codes/claude/v1/models",
|
||
"https://www.right.codes/v1/models",
|
||
"https://www.right.codes/models",
|
||
]
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn test_candidates_longer_suffix_wins() {
|
||
// baseURL 以 /api/anthropic 结尾时，应剥离整个 /api/anthropic，
|
||
// 而不是只剥离 /anthropic（那样会得到残缺的 https://.../api 根）。
|
||
let c = build_models_url_candidates("https://api.z.ai/api/anthropic", false, None).unwrap();
|
||
assert_eq!(
|
||
c,
|
||
vec![
|
||
"https://api.z.ai/api/anthropic/v1/models",
|
||
"https://api.z.ai/v1/models",
|
||
"https://api.z.ai/models",
|
||
]
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn test_candidates_no_suffix_no_strip() {
|
||
let c = build_models_url_candidates("https://openrouter.ai/api", false, None).unwrap();
|
||
assert_eq!(c, vec!["https://openrouter.ai/api/v1/models"]);
|
||
}
|
||
|
||
#[test]
|
||
fn test_candidates_deduplicate() {
|
||
// 虚构 case：baseURL 就是 "scheme://host"，剥不出子路径，应只有一个候选。
|
||
let c = build_models_url_candidates("https://host.example.com", false, None).unwrap();
|
||
assert_eq!(c.len(), 1);
|
||
}
|
||
|
||
#[test]
|
||
fn test_parse_response() {
|
||
let json = r#"{"object":"list","data":[{"id":"gpt-4","object":"model","owned_by":"openai"},{"id":"claude-3-sonnet","object":"model","owned_by":"anthropic"}]}"#;
|
||
let resp: ModelsResponse = serde_json::from_str(json).unwrap();
|
||
let data = resp.data.unwrap();
|
||
assert_eq!(data.len(), 2);
|
||
assert_eq!(data[0].id, "gpt-4");
|
||
assert_eq!(data[0].owned_by.as_deref(), Some("openai"));
|
||
assert_eq!(data[1].id, "claude-3-sonnet");
|
||
}
|
||
|
||
#[test]
|
||
fn test_parse_response_no_owned_by() {
|
||
let json = r#"{"object":"list","data":[{"id":"my-model","object":"model"}]}"#;
|
||
let resp: ModelsResponse = serde_json::from_str(json).unwrap();
|
||
let data = resp.data.unwrap();
|
||
assert_eq!(data[0].id, "my-model");
|
||
assert!(data[0].owned_by.is_none());
|
||
}
|
||
|
||
#[test]
|
||
fn test_parse_response_empty_data() {
|
||
let json = r#"{"object":"list","data":[]}"#;
|
||
let resp: ModelsResponse = serde_json::from_str(json).unwrap();
|
||
assert!(resp.data.unwrap().is_empty());
|
||
}
|
||
}
|