mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-06 01:17:46 +08:00
Merge origin/main into main
This commit is contained in:
elky
@@ -1930,6 +1930,29 @@ mod tests {
|
||||
assert_eq!(active["reasoning_effort"], "max");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn user_usage_payloads_expose_gemini_thinking_config_reasoning_mapping() {
|
||||
let item = StoredRequestUsageAudit {
|
||||
request_body: Some(json!({
|
||||
"generationConfig": {
|
||||
"thinkingConfig": { "includeThoughts": true, "thinkingLevel": "HIGH" }
|
||||
}
|
||||
})),
|
||||
provider_request_body: Some(json!({
|
||||
"generationConfig": { "thinkingConfig": { "thinkingBudget": 8192 } }
|
||||
})),
|
||||
..sample_usage("completed")
|
||||
};
|
||||
|
||||
let record = build_users_me_usage_record_payload(&item, false, &BTreeMap::new(), false);
|
||||
let active = build_users_me_usage_active_payload(&item);
|
||||
|
||||
assert_eq!(record["requested_reasoning_effort"], "high");
|
||||
assert_eq!(active["requested_reasoning_effort"], "high");
|
||||
assert_eq!(record["reasoning_effort"], "xhigh");
|
||||
assert_eq!(active["reasoning_effort"], "xhigh");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn user_usage_payloads_expose_websocket_transport() {
|
||||
let item = StoredRequestUsageAudit {
|
||||
|
||||
@@ -2681,6 +2681,84 @@ fn windsurf_quota_snapshot_has_stale_cooldown(quota_snapshot: &Map<String, Value
|
||||
has_capacity || !exhausted
|
||||
}
|
||||
|
||||
/// 读取状态快照时，把已越过重置时间点的配额窗口按“已重置”口径归一化。
|
||||
///
|
||||
/// 背景：调度侧早已把到期窗口视为未耗尽（`provider_pool_reset_deadline_elapsed`），
|
||||
/// 账号额度文本也会按到期强制显示 100%，但列表读取此前直接返回存量快照，
|
||||
/// 导致管理端倒计时归零后进度条仍停留在旧的剩余百分比。这里让读取层与
|
||||
/// 调度侧、文本侧使用同一口径，避免三处状态互相矛盾。
|
||||
///
|
||||
/// 覆盖所有带重置时间的提供商窗口（codex/kiro/xai/grok/antigravity/
|
||||
/// gemini_cli/chatgpt_web/windsurf 等）：比例（used_ratio/remaining_ratio）、
|
||||
/// 数值（used_value/remaining_value）与窗口级耗尽标记会一起恢复为“已重置”。
|
||||
fn normalize_expired_quota_windows(snapshot: &mut serde_json::Map<String, Value>) {
|
||||
let Some(quota) = snapshot.get_mut("quota").and_then(Value::as_object_mut) else {
|
||||
return;
|
||||
};
|
||||
let now_unix_secs = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.ok()
|
||||
.map(|duration| duration.as_secs())
|
||||
.unwrap_or(0);
|
||||
let fallback_observed_at = provider_quota_timestamp_unix_secs(quota.get("observed_at"))
|
||||
.or_else(|| provider_quota_timestamp_unix_secs(quota.get("updated_at")));
|
||||
let Some(windows) = quota.get_mut("windows").and_then(Value::as_array_mut) else {
|
||||
return;
|
||||
};
|
||||
|
||||
for window in windows.iter_mut().filter_map(Value::as_object_mut) {
|
||||
// 与同类窗口处理保持一致：window_minutes=0 不是真实配额窗口。
|
||||
if window.get("window_minutes").and_then(Value::as_u64) == Some(0) {
|
||||
continue;
|
||||
}
|
||||
if !aether_provider_pool::provider_pool_reset_deadline_elapsed(
|
||||
window,
|
||||
fallback_observed_at,
|
||||
now_unix_secs,
|
||||
) {
|
||||
continue;
|
||||
}
|
||||
// 只归一化“观测完整”的窗口：要么有比例观测，要么有上限 + 用量数值。
|
||||
// 这样既不会把无数据窗口凭空显示成 100%，也不会出现比例已恢复 100%
|
||||
// 而数值仍停留在旧值的不一致（例如只有 remaining_value 却没有上限的窗口）。
|
||||
let has_ratio_observation = ["used_ratio", "remaining_ratio"]
|
||||
.into_iter()
|
||||
.any(|field| window.get(field).is_some_and(Value::is_number));
|
||||
let limit_value = window
|
||||
.get("limit_value")
|
||||
.and_then(admin_provider_quota_pure::coerce_json_f64)
|
||||
.filter(|value| *value > 0.0);
|
||||
let has_value_observation = limit_value.is_some()
|
||||
&& ["used_value", "remaining_value"]
|
||||
.into_iter()
|
||||
.any(|field| window.get(field).is_some_and(Value::is_number));
|
||||
if !has_ratio_observation && !has_value_observation {
|
||||
continue;
|
||||
}
|
||||
// 比例口径：已用清零、剩余 100%；字段原本为 null 时一并补齐，保证展示口径统一。
|
||||
for (field, value) in [("used_ratio", 0.0), ("remaining_ratio", 1.0)] {
|
||||
if window.contains_key(field) {
|
||||
window.insert(field.to_string(), json!(value));
|
||||
}
|
||||
}
|
||||
// 数值口径：已用清零；有上限时把剩余恢复到上限（用于“剩余 x/y”类文本展示）。
|
||||
if window.contains_key("used_value") {
|
||||
window.insert("used_value".to_string(), json!(0.0));
|
||||
}
|
||||
if window.contains_key("remaining_value") {
|
||||
if let Some(limit_value) = limit_value {
|
||||
window.insert("remaining_value".to_string(), json!(limit_value));
|
||||
}
|
||||
}
|
||||
// 同步清掉窗口级耗尽标记，避免展示与调度口径互相矛盾。
|
||||
for field in ["is_exhausted", "exhausted"] {
|
||||
if let Some(slot) = window.get_mut(field) {
|
||||
*slot = json!(false);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn provider_key_status_snapshot_payload(
|
||||
key: &StoredProviderCatalogKey,
|
||||
provider_type: &str,
|
||||
@@ -2722,9 +2800,8 @@ pub(crate) fn provider_key_status_snapshot_payload(
|
||||
// Legacy snapshots can retain an exhausted summary after a window reset or
|
||||
// newer quota observation. Use the same decision as scheduling so the
|
||||
// account list and its status filter do not keep displaying that stale block.
|
||||
if provider_type.trim().eq_ignore_ascii_case("codex")
|
||||
&& !aether_provider_pool::provider_pool_key_account_quota_exhausted(key, provider_type)
|
||||
{
|
||||
// 对所有提供商生效：适配器判定已是“重置感知”的，与调度口径保持一致。
|
||||
if !aether_provider_pool::provider_pool_key_account_quota_exhausted(key, provider_type) {
|
||||
if let Some(quota) = snapshot.get_mut("quota").and_then(Value::as_object_mut) {
|
||||
quota.insert("exhausted".to_string(), json!(false));
|
||||
if quota.get("code").and_then(Value::as_str) == Some("exhausted") {
|
||||
@@ -2732,6 +2809,9 @@ pub(crate) fn provider_key_status_snapshot_payload(
|
||||
}
|
||||
}
|
||||
}
|
||||
// 读取时归一化已到期的配额窗口，保证列表进度条、额度文字与调度侧、
|
||||
// 账号额度文本使用同一“已重置”口径。
|
||||
normalize_expired_quota_windows(&mut snapshot);
|
||||
snapshot.insert(
|
||||
"oauth".to_string(),
|
||||
build_provider_key_oauth_status_snapshot(key),
|
||||
@@ -3780,6 +3860,421 @@ mod tests {
|
||||
assert_eq!(window.get("reset_seconds"), Some(&json!(3_600u64)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_key_status_snapshot_payload_normalizes_expired_codex_quota_windows() {
|
||||
let mut key = sample_catalog_key();
|
||||
key.status_snapshot = Some(json!({
|
||||
"quota": {
|
||||
"version": 2,
|
||||
"provider_type": "codex",
|
||||
"code": "ok",
|
||||
"exhausted": false,
|
||||
"updated_at": 1_700_000_000u64,
|
||||
"windows": [
|
||||
{
|
||||
"code": "5h",
|
||||
"label": "5H",
|
||||
"scope": "account",
|
||||
"unit": "percent",
|
||||
"used_ratio": 0.88,
|
||||
"remaining_ratio": 0.12,
|
||||
"reset_at": 1_700_003_600u64,
|
||||
"window_minutes": 300
|
||||
},
|
||||
{
|
||||
"code": "weekly",
|
||||
"label": "周",
|
||||
"scope": "account",
|
||||
"unit": "percent",
|
||||
"used_ratio": 0.5,
|
||||
"remaining_ratio": 0.5,
|
||||
"reset_at": 2_000_000_000u64,
|
||||
"window_minutes": 10_080
|
||||
},
|
||||
{
|
||||
"code": "spark_5h",
|
||||
"label": "Spark 5H",
|
||||
"scope": "account",
|
||||
"unit": "percent",
|
||||
"reset_at": 1_700_003_600u64,
|
||||
"window_minutes": 300
|
||||
},
|
||||
{
|
||||
"code": "unlimited",
|
||||
"label": "无限",
|
||||
"scope": "account",
|
||||
"unit": "percent",
|
||||
"used_ratio": 0.3,
|
||||
"remaining_ratio": 0.7,
|
||||
"reset_at": 1_700_003_600u64,
|
||||
"window_minutes": 0
|
||||
}
|
||||
]
|
||||
}
|
||||
}));
|
||||
|
||||
let payload = provider_key_status_snapshot_payload(&key, "codex");
|
||||
|
||||
// 已到期的窗口按“已重置”口径归一化：用量清零、剩余 100%。
|
||||
assert_eq!(payload.pointer("/quota/windows/0/code"), Some(&json!("5h")));
|
||||
assert_eq!(
|
||||
payload.pointer("/quota/windows/0/used_ratio"),
|
||||
Some(&json!(0.0))
|
||||
);
|
||||
assert_eq!(
|
||||
payload.pointer("/quota/windows/0/remaining_ratio"),
|
||||
Some(&json!(1.0))
|
||||
);
|
||||
// 未到期的窗口保留原始观测值。
|
||||
assert_eq!(
|
||||
payload.pointer("/quota/windows/1/code"),
|
||||
Some(&json!("weekly"))
|
||||
);
|
||||
assert_eq!(
|
||||
payload.pointer("/quota/windows/1/used_ratio"),
|
||||
Some(&json!(0.5))
|
||||
);
|
||||
assert_eq!(
|
||||
payload.pointer("/quota/windows/1/remaining_ratio"),
|
||||
Some(&json!(0.5))
|
||||
);
|
||||
// 没有用量观测的窗口不会被凭空补成 100%。
|
||||
assert_eq!(
|
||||
payload.pointer("/quota/windows/2/code"),
|
||||
Some(&json!("spark_5h"))
|
||||
);
|
||||
assert_eq!(payload.pointer("/quota/windows/2/used_ratio"), None);
|
||||
assert_eq!(payload.pointer("/quota/windows/2/remaining_ratio"), None);
|
||||
// window_minutes=0 不是真实配额窗口，与同类 Codex 窗口处理保持一致，不归一化。
|
||||
assert_eq!(
|
||||
payload.pointer("/quota/windows/3/code"),
|
||||
Some(&json!("unlimited"))
|
||||
);
|
||||
assert_eq!(
|
||||
payload.pointer("/quota/windows/3/used_ratio"),
|
||||
Some(&json!(0.3))
|
||||
);
|
||||
assert_eq!(
|
||||
payload.pointer("/quota/windows/3/remaining_ratio"),
|
||||
Some(&json!(0.7))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_key_status_snapshot_payload_clears_expired_codex_window_exhausted_state() {
|
||||
let mut key = sample_catalog_key();
|
||||
key.status_snapshot = Some(json!({
|
||||
"quota": {
|
||||
"version": 2,
|
||||
"provider_type": "codex",
|
||||
"code": "exhausted",
|
||||
"exhausted": true,
|
||||
"usage_ratio": 1.0,
|
||||
"updated_at": 1_700_000_000u64,
|
||||
"windows": [
|
||||
{
|
||||
"code": "weekly",
|
||||
"label": "周",
|
||||
"scope": "account",
|
||||
"unit": "percent",
|
||||
"used_ratio": 1.0,
|
||||
"remaining_ratio": 0.0,
|
||||
"reset_at": 1_700_003_600u64,
|
||||
"window_minutes": 10_080,
|
||||
"is_exhausted": true
|
||||
}
|
||||
]
|
||||
}
|
||||
}));
|
||||
|
||||
let payload = provider_key_status_snapshot_payload(&key, "codex");
|
||||
|
||||
// 汇总标志沿用既有调度口径（窗口已到期 → 不再视为耗尽）。
|
||||
assert_eq!(payload.pointer("/quota/exhausted"), Some(&json!(false)));
|
||||
assert_eq!(payload.pointer("/quota/code"), Some(&json!("ok")));
|
||||
// 窗口级耗尽标记与剩余比例同步归一化。
|
||||
assert_eq!(
|
||||
payload.pointer("/quota/windows/0/is_exhausted"),
|
||||
Some(&json!(false))
|
||||
);
|
||||
assert_eq!(
|
||||
payload.pointer("/quota/windows/0/used_ratio"),
|
||||
Some(&json!(0.0))
|
||||
);
|
||||
assert_eq!(
|
||||
payload.pointer("/quota/windows/0/remaining_ratio"),
|
||||
Some(&json!(1.0))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_key_status_snapshot_payload_normalizes_expired_quota_windows_for_other_providers() {
|
||||
let mut key = sample_catalog_key();
|
||||
key.status_snapshot = Some(json!({
|
||||
"quota": {
|
||||
"version": 2,
|
||||
"provider_type": "kiro",
|
||||
"code": "ok",
|
||||
"exhausted": false,
|
||||
"observed_at": 1_700_000_000u64,
|
||||
"updated_at": 1_700_000_000u64,
|
||||
"windows": [
|
||||
{
|
||||
"code": "usage",
|
||||
"label": "额度",
|
||||
"scope": "account",
|
||||
"unit": "count",
|
||||
"used_ratio": 0.4,
|
||||
"remaining_ratio": 0.6,
|
||||
"used_value": 60.0,
|
||||
"remaining_value": 90.0,
|
||||
"limit_value": 150.0,
|
||||
"reset_at": 1_700_003_600u64,
|
||||
"reset_seconds": 3_600u64
|
||||
},
|
||||
{
|
||||
"code": "usage_null_ratio",
|
||||
"label": "额度",
|
||||
"scope": "account",
|
||||
"unit": "count",
|
||||
"used_ratio": null,
|
||||
"remaining_ratio": null,
|
||||
"used_value": 60.0,
|
||||
"remaining_value": 90.0,
|
||||
"limit_value": 150.0,
|
||||
"reset_at": 1_700_003_600u64,
|
||||
"reset_seconds": 3_600u64
|
||||
},
|
||||
{
|
||||
"code": "usage_active",
|
||||
"label": "额度",
|
||||
"scope": "account",
|
||||
"unit": "count",
|
||||
"used_ratio": 0.4,
|
||||
"remaining_ratio": 0.6,
|
||||
"used_value": 60.0,
|
||||
"remaining_value": 90.0,
|
||||
"limit_value": 150.0,
|
||||
"reset_at": 2_000_000_000u64,
|
||||
"reset_seconds": 300_000_000u64
|
||||
},
|
||||
{
|
||||
"code": "usage_no_observation",
|
||||
"label": "额度",
|
||||
"scope": "account",
|
||||
"unit": "count",
|
||||
"reset_at": 1_700_003_600u64,
|
||||
"reset_seconds": 3_600u64
|
||||
},
|
||||
{
|
||||
"code": "usage_no_deadline",
|
||||
"label": "额度",
|
||||
"scope": "account",
|
||||
"unit": "count",
|
||||
"used_ratio": 0.4,
|
||||
"remaining_ratio": 0.6,
|
||||
"used_value": 60.0,
|
||||
"remaining_value": 90.0,
|
||||
"limit_value": 150.0
|
||||
},
|
||||
{
|
||||
"code": "usage_remaining_without_limit",
|
||||
"label": "额度",
|
||||
"scope": "account",
|
||||
"unit": "count",
|
||||
"used_ratio": null,
|
||||
"remaining_ratio": null,
|
||||
"remaining_value": 0.0,
|
||||
"reset_at": 1_700_003_600u64,
|
||||
"reset_seconds": 3_600u64
|
||||
}
|
||||
]
|
||||
}
|
||||
}));
|
||||
|
||||
let payload = provider_key_status_snapshot_payload(&key, "kiro");
|
||||
|
||||
// 已到期窗口：比例与数值一起按“已重置”归一化（不再只覆盖 Codex）。
|
||||
assert_eq!(
|
||||
payload.pointer("/quota/windows/0/used_ratio"),
|
||||
Some(&json!(0.0))
|
||||
);
|
||||
assert_eq!(
|
||||
payload.pointer("/quota/windows/0/remaining_ratio"),
|
||||
Some(&json!(1.0))
|
||||
);
|
||||
assert_eq!(
|
||||
payload.pointer("/quota/windows/0/used_value"),
|
||||
Some(&json!(0.0))
|
||||
);
|
||||
assert_eq!(
|
||||
payload.pointer("/quota/windows/0/remaining_value"),
|
||||
Some(&json!(150.0))
|
||||
);
|
||||
// 比例字段为 null 时补齐为已重置口径，保证优先读比例的展示实现一致。
|
||||
assert_eq!(
|
||||
payload.pointer("/quota/windows/1/used_ratio"),
|
||||
Some(&json!(0.0))
|
||||
);
|
||||
assert_eq!(
|
||||
payload.pointer("/quota/windows/1/remaining_ratio"),
|
||||
Some(&json!(1.0))
|
||||
);
|
||||
assert_eq!(
|
||||
payload.pointer("/quota/windows/1/used_value"),
|
||||
Some(&json!(0.0))
|
||||
);
|
||||
assert_eq!(
|
||||
payload.pointer("/quota/windows/1/remaining_value"),
|
||||
Some(&json!(150.0))
|
||||
);
|
||||
// 未到期窗口保留原始观测值。
|
||||
assert_eq!(
|
||||
payload.pointer("/quota/windows/2/used_ratio"),
|
||||
Some(&json!(0.4))
|
||||
);
|
||||
assert_eq!(
|
||||
payload.pointer("/quota/windows/2/remaining_value"),
|
||||
Some(&json!(90.0))
|
||||
);
|
||||
// 没有用量观测的窗口不会被凭空补成 100%。
|
||||
assert_eq!(payload.pointer("/quota/windows/3/used_ratio"), None);
|
||||
assert_eq!(payload.pointer("/quota/windows/3/remaining_ratio"), None);
|
||||
// 没有重置时间的窗口无法判定是否已重置，保持原样。
|
||||
assert_eq!(
|
||||
payload.pointer("/quota/windows/4/used_ratio"),
|
||||
Some(&json!(0.4))
|
||||
);
|
||||
assert_eq!(
|
||||
payload.pointer("/quota/windows/4/remaining_ratio"),
|
||||
Some(&json!(0.6))
|
||||
);
|
||||
// 只有 remaining_value 却没有上限的窗口观测不完整，跳过归一化，
|
||||
// 避免比例已恢复 100% 而数值仍停留在旧值。
|
||||
assert_eq!(
|
||||
payload.pointer("/quota/windows/5/remaining_ratio"),
|
||||
Some(&json!(null))
|
||||
);
|
||||
assert_eq!(
|
||||
payload.pointer("/quota/windows/5/remaining_value"),
|
||||
Some(&json!(0.0))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_key_status_snapshot_payload_clears_expired_quota_exhausted_state_for_other_providers(
|
||||
) {
|
||||
let mut key = sample_catalog_key();
|
||||
key.status_snapshot = Some(json!({
|
||||
"quota": {
|
||||
"version": 2,
|
||||
"provider_type": "kiro",
|
||||
"code": "exhausted",
|
||||
"label": "额度耗尽",
|
||||
"reason": "额度已耗尽",
|
||||
"exhausted": true,
|
||||
"usage_ratio": 1.0,
|
||||
"observed_at": 1_700_000_000u64,
|
||||
"updated_at": 1_700_000_000u64,
|
||||
"windows": [
|
||||
{
|
||||
"code": "usage",
|
||||
"label": "额度",
|
||||
"scope": "account",
|
||||
"unit": "count",
|
||||
"used_ratio": 1.0,
|
||||
"remaining_ratio": 0.0,
|
||||
"used_value": 150.0,
|
||||
"remaining_value": 0.0,
|
||||
"limit_value": 150.0,
|
||||
"reset_at": 1_700_003_600u64,
|
||||
"reset_seconds": 3_600u64,
|
||||
"is_exhausted": true
|
||||
}
|
||||
]
|
||||
}
|
||||
}));
|
||||
|
||||
let payload = provider_key_status_snapshot_payload(&key, "kiro");
|
||||
|
||||
// 汇总状态沿用调度口径（窗口已到期 → 不再视为耗尽），对所有提供商生效。
|
||||
assert_eq!(payload.pointer("/quota/exhausted"), Some(&json!(false)));
|
||||
assert_eq!(payload.pointer("/quota/code"), Some(&json!("ok")));
|
||||
// 窗口比例、数值与耗尽标记同步归一化。
|
||||
assert_eq!(
|
||||
payload.pointer("/quota/windows/0/is_exhausted"),
|
||||
Some(&json!(false))
|
||||
);
|
||||
assert_eq!(
|
||||
payload.pointer("/quota/windows/0/used_ratio"),
|
||||
Some(&json!(0.0))
|
||||
);
|
||||
assert_eq!(
|
||||
payload.pointer("/quota/windows/0/remaining_ratio"),
|
||||
Some(&json!(1.0))
|
||||
);
|
||||
assert_eq!(
|
||||
payload.pointer("/quota/windows/0/used_value"),
|
||||
Some(&json!(0.0))
|
||||
);
|
||||
assert_eq!(
|
||||
payload.pointer("/quota/windows/0/remaining_value"),
|
||||
Some(&json!(150.0))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_key_status_snapshot_payload_keeps_active_quota_exhausted_state_for_other_providers()
|
||||
{
|
||||
let mut key = sample_catalog_key();
|
||||
key.status_snapshot = Some(json!({
|
||||
"quota": {
|
||||
"version": 2,
|
||||
"provider_type": "kiro",
|
||||
"code": "exhausted",
|
||||
"label": "额度耗尽",
|
||||
"exhausted": true,
|
||||
"usage_ratio": 1.0,
|
||||
"observed_at": 1_700_000_000u64,
|
||||
"updated_at": 1_700_000_000u64,
|
||||
"windows": [
|
||||
{
|
||||
"code": "usage",
|
||||
"label": "额度",
|
||||
"scope": "account",
|
||||
"unit": "count",
|
||||
"used_ratio": 1.0,
|
||||
"remaining_ratio": 0.0,
|
||||
"used_value": 150.0,
|
||||
"remaining_value": 0.0,
|
||||
"limit_value": 150.0,
|
||||
"reset_at": 2_000_000_000u64,
|
||||
"reset_seconds": 300_000_000u64,
|
||||
"is_exhausted": true
|
||||
}
|
||||
]
|
||||
}
|
||||
}));
|
||||
|
||||
let payload = provider_key_status_snapshot_payload(&key, "kiro");
|
||||
|
||||
// 未到期的耗尽状态保留，不能被读取层提前“恢复”。
|
||||
assert_eq!(payload.pointer("/quota/exhausted"), Some(&json!(true)));
|
||||
assert_eq!(payload.pointer("/quota/code"), Some(&json!("exhausted")));
|
||||
assert_eq!(
|
||||
payload.pointer("/quota/windows/0/is_exhausted"),
|
||||
Some(&json!(true))
|
||||
);
|
||||
assert_eq!(
|
||||
payload.pointer("/quota/windows/0/used_ratio"),
|
||||
Some(&json!(1.0))
|
||||
);
|
||||
assert_eq!(
|
||||
payload.pointer("/quota/windows/0/remaining_ratio"),
|
||||
Some(&json!(0.0))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_key_status_snapshot_payload_backfills_chatgpt_web_image_quota() {
|
||||
let mut key = sample_catalog_key();
|
||||
@@ -3911,14 +4406,14 @@ mod tests {
|
||||
"used_percent": 60.0,
|
||||
"remaining": 60.0,
|
||||
"total": 150.0,
|
||||
"reset_at": 1_778_157_172u64,
|
||||
"reset_at": 2_000_000_000u64,
|
||||
"is_exhausted": false
|
||||
},
|
||||
"quota_heavy": {
|
||||
"display_name": "heavy",
|
||||
"remaining_fraction": 0.0,
|
||||
"used_percent": 100.0,
|
||||
"reset_at": 1_778_157_172u64,
|
||||
"reset_at": 2_000_000_000u64,
|
||||
"is_exhausted": true
|
||||
}
|
||||
}
|
||||
@@ -3941,7 +4436,7 @@ mod tests {
|
||||
assert_eq!(quota.get("pool_tier"), Some(&json!("heavy")));
|
||||
assert_eq!(quota.get("exhausted"), Some(&json!(false)));
|
||||
assert_eq!(quota.get("usage_ratio"), Some(&json!(1.0)));
|
||||
assert_eq!(quota.get("reset_at"), Some(&json!(1_778_157_172u64)));
|
||||
assert_eq!(quota.get("reset_at"), Some(&json!(2_000_000_000u64)));
|
||||
assert_eq!(windows.len(), 2);
|
||||
assert!(windows.iter().any(|window| {
|
||||
window
|
||||
@@ -4054,8 +4549,8 @@ mod tests {
|
||||
"plan_name": "Pro",
|
||||
"daily_remaining_percent": 40.0,
|
||||
"weekly_remaining_percent": 65.0,
|
||||
"daily_reset_at": 1_778_100_000u64,
|
||||
"weekly_reset_at": 1_778_600_000u64,
|
||||
"daily_reset_at": 2_000_000_000u64,
|
||||
"weekly_reset_at": 2_000_600_000u64,
|
||||
"prompt_used": 12.0,
|
||||
"prompt_limit": 100.0,
|
||||
"prompt_remaining": 88.0,
|
||||
@@ -4090,13 +4585,13 @@ mod tests {
|
||||
assert_eq!(quota.get("code"), Some(&json!("ok")));
|
||||
assert_eq!(quota.get("plan_type"), Some(&json!("Pro")));
|
||||
assert_eq!(quota.get("usage_ratio"), Some(&json!(0.6)));
|
||||
assert_eq!(quota.get("reset_at"), Some(&json!(1_778_100_000u64)));
|
||||
assert_eq!(quota.get("reset_at"), Some(&json!(2_000_000_000u64)));
|
||||
assert_eq!(daily.get("remaining_ratio"), Some(&json!(0.4)));
|
||||
assert_eq!(daily.get("used_ratio"), Some(&json!(0.6)));
|
||||
assert_eq!(daily.get("reset_seconds"), Some(&json!(32_754u64)));
|
||||
assert_eq!(daily.get("reset_seconds"), Some(&json!(221_932_754u64)));
|
||||
assert_eq!(weekly.get("remaining_ratio"), Some(&json!(0.65)));
|
||||
assert_eq!(weekly.get("used_ratio"), Some(&json!(0.35)));
|
||||
assert_eq!(weekly.get("reset_seconds"), Some(&json!(532_754u64)));
|
||||
assert_eq!(weekly.get("reset_seconds"), Some(&json!(222_532_754u64)));
|
||||
assert_eq!(quota.get("allowed_models_count"), Some(&json!(82)));
|
||||
}
|
||||
|
||||
|
||||
@@ -92,6 +92,7 @@ mod upstream_admission;
|
||||
mod usage;
|
||||
mod video_tasks;
|
||||
mod wallet_runtime;
|
||||
mod xai_profile;
|
||||
|
||||
pub use self::ai_serving::api::{codex_client_originator, codex_client_user_agent};
|
||||
pub(crate) use self::ai_serving::api::{
|
||||
|
||||
@@ -2527,6 +2527,20 @@ async fn run() -> Result<(), Box<dyn std::error::Error>> {
|
||||
);
|
||||
}
|
||||
}
|
||||
match state.prewarm_xai_client_profile().await {
|
||||
Ok(version) => {
|
||||
info!(
|
||||
xai_client_version = %version,
|
||||
"prewarmed Grok CLI client profile"
|
||||
);
|
||||
}
|
||||
Err(err) => {
|
||||
warn!(
|
||||
error = %err,
|
||||
"failed to refresh Grok CLI client profile; built-in or cached profile remains active"
|
||||
);
|
||||
}
|
||||
}
|
||||
match prewarm_direct_h2c_sender_cache_from_env_for_startup().await {
|
||||
Ok(Some(report)) => {
|
||||
if report.failed_targets > 0 {
|
||||
|
||||
@@ -76,6 +76,7 @@ use crate::maintenance::spawn_stats_hourly_aggregation_worker;
|
||||
use crate::maintenance::spawn_usage_cleanup_worker;
|
||||
use crate::maintenance::spawn_usage_counter_flush_worker;
|
||||
use crate::maintenance::spawn_wallet_daily_usage_aggregation_worker;
|
||||
use crate::xai_profile::spawn_worker as spawn_xai_client_profile_worker;
|
||||
|
||||
const SYSTEM_CONFIG_CACHE_TTL: Duration = Duration::from_secs(30);
|
||||
// Requests may use a stale value after the fresh window until the entry reaches
|
||||
@@ -154,6 +155,10 @@ impl AppState {
|
||||
crate::codex_profile::prewarm(self.runtime_state()).await
|
||||
}
|
||||
|
||||
pub async fn prewarm_xai_client_profile(&self) -> Result<String, String> {
|
||||
crate::xai_profile::prewarm(self.runtime_state()).await
|
||||
}
|
||||
|
||||
pub async fn prewarm_chat_pii_redaction_runtime_config(&self) -> Result<bool, String> {
|
||||
crate::privacy::read_chat_pii_redaction_runtime_config(self)
|
||||
.await
|
||||
@@ -2356,6 +2361,10 @@ impl AppState {
|
||||
crate::task_runtime::TASK_KEY_CODEX_CLIENT_PROFILE,
|
||||
Some(spawn_codex_client_profile_worker(background_state.clone())),
|
||||
);
|
||||
supervise_worker(
|
||||
crate::task_runtime::TASK_KEY_XAI_CLIENT_PROFILE,
|
||||
Some(spawn_xai_client_profile_worker(background_state.clone())),
|
||||
);
|
||||
supervise_worker(
|
||||
crate::task_runtime::TASK_KEY_VIDEO_TASK_POLLER,
|
||||
spawn_video_task_poller(background_state.clone()),
|
||||
|
||||
@@ -25,6 +25,7 @@ pub(crate) const TASK_KEY_USAGE_COUNTER_FLUSH: &str = "usage.counter.flush.worke
|
||||
pub(crate) const TASK_KEY_VIDEO_TASK_POLLER: &str = "video.task.poller";
|
||||
pub(crate) const TASK_KEY_MODEL_FETCH_WORKER: &str = "model.fetch.worker";
|
||||
pub(crate) const TASK_KEY_CODEX_CLIENT_PROFILE: &str = "maintenance.codex.client.profile";
|
||||
pub(crate) const TASK_KEY_XAI_CLIENT_PROFILE: &str = "maintenance.xai.client.profile";
|
||||
pub(crate) const TASK_KEY_PROVIDER_QUOTA_RESET: &str = "provider.quota.reset.worker";
|
||||
pub(crate) const TASK_KEY_ACCOUNT_SELF_CHECK: &str = "account.self_check.worker";
|
||||
pub(crate) const TASK_KEY_POOL_SCORE_REBUILD: &str = "pool.score.rebuild.worker";
|
||||
@@ -211,6 +212,14 @@ const TASK_DEFINITIONS: &[TaskDefinition] = &[
|
||||
true,
|
||||
RETRY_ONCE,
|
||||
),
|
||||
TaskDefinition::new(
|
||||
TASK_KEY_XAI_CLIENT_PROFILE,
|
||||
TaskKind::Scheduled,
|
||||
"interval",
|
||||
true,
|
||||
true,
|
||||
RETRY_ONCE,
|
||||
),
|
||||
TaskDefinition::new(
|
||||
TASK_KEY_PROVIDER_QUOTA_RESET,
|
||||
TaskKind::Scheduled,
|
||||
|
||||
@@ -3044,7 +3044,94 @@ async fn gateway_prefers_status_snapshot_kiro_quota_over_stale_metadata() {
|
||||
assert_eq!(keys[0]["scheduling_status"], json!("available"));
|
||||
assert_eq!(keys[0]["scheduling_reason"], json!("available"));
|
||||
assert_eq!(keys[0]["quota_updated_at"], json!(1_775_553_285u64));
|
||||
assert_eq!(keys[0]["account_quota"], json!("剩余 75.0% (5/20)"));
|
||||
// 窗口重置时间已过：读取层按“已重置”口径归一化（调度侧同样不再视为耗尽），
|
||||
// 同时仍优先使用快照（上限 20）而不是陈旧元数据（上限 100）。
|
||||
assert_eq!(keys[0]["account_quota"], json!("剩余 100.0% (0/20)"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_kiro_pool_filter_clears_stale_exhausted_summary_after_window_reset() {
|
||||
let mut provider = sample_provider("provider-kiro", "kiro", 10);
|
||||
provider.provider_type = "kiro".to_string();
|
||||
provider.config = Some(json!({"pool_advanced": {
|
||||
"reserve_minimum_quota": false, "skip_exhausted_accounts": true
|
||||
}}));
|
||||
let mut key = sample_key(
|
||||
"key-kiro-expired",
|
||||
"provider-kiro",
|
||||
"claude:messages",
|
||||
"oauth-placeholder",
|
||||
);
|
||||
key.name = "kiro expired quota key".to_string();
|
||||
key.auth_type = "oauth".to_string();
|
||||
key.status_snapshot = Some(json!({
|
||||
"quota": {
|
||||
"version": 2,
|
||||
"provider_type": "kiro",
|
||||
"code": "exhausted",
|
||||
"label": "额度耗尽",
|
||||
"reason": "额度已耗尽",
|
||||
"freshness": "fresh",
|
||||
"source": "refresh_api",
|
||||
"observed_at": 1_775_553_285u64,
|
||||
"exhausted": true,
|
||||
"usage_ratio": 1.0,
|
||||
"updated_at": 1_775_553_285u64,
|
||||
"reset_seconds": 0u64,
|
||||
"plan_type": "KIRO PRO+",
|
||||
"windows": [
|
||||
{
|
||||
"code": "usage",
|
||||
"label": "额度",
|
||||
"scope": "account",
|
||||
"unit": "count",
|
||||
"used_ratio": 1.0,
|
||||
"remaining_ratio": 0.0,
|
||||
"used_value": 20.0,
|
||||
"remaining_value": 0.0,
|
||||
"limit_value": 20.0,
|
||||
"reset_at": 1_775_639_685u64,
|
||||
"reset_seconds": 0u64,
|
||||
"is_exhausted": true
|
||||
}
|
||||
]
|
||||
}
|
||||
}));
|
||||
|
||||
let state = AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(GatewayDataState::with_provider_catalog_reader_for_tests(
|
||||
Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![provider],
|
||||
Vec::new(),
|
||||
vec![key],
|
||||
)),
|
||||
));
|
||||
|
||||
for status in ["all", "available", "quota_exhausted"] {
|
||||
let response = local_admin_pool_response(
|
||||
&state,
|
||||
http::Method::GET,
|
||||
&format!("/api/admin/pool/provider-kiro/keys?page=1&page_size=50&status={status}"),
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = serde_json::from_slice(
|
||||
&to_bytes(response.into_body(), usize::MAX)
|
||||
.await
|
||||
.expect("body should read"),
|
||||
)
|
||||
.expect("json body should parse");
|
||||
let keys = payload["keys"].as_array().expect("keys should be array");
|
||||
assert_eq!(keys.len(), usize::from(status != "quota_exhausted"));
|
||||
if let Some(key) = keys.first() {
|
||||
assert_eq!(key["scheduling_status"], json!("available"));
|
||||
assert_eq!(key["status_snapshot"]["quota"]["code"], json!("ok"));
|
||||
assert_eq!(key["status_snapshot"]["quota"]["exhausted"], json!(false));
|
||||
assert_eq!(key["account_quota"], json!("剩余 100.0% (0/20)"));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
|
||||
@@ -0,0 +1,592 @@
|
||||
//! Grok CLI 客户端版本的运行时发布与官方版本刷新。
|
||||
//!
|
||||
//! cli-chat-proxy.grok.com 会对低于最低版本的 `x-grok-client-version` 直接返回 426，
|
||||
//! 因此网关定期读取官方发布渠道并原子替换传输层使用的版本号。
|
||||
|
||||
use std::collections::BTreeMap;
|
||||
use std::future::Future;
|
||||
use std::time::Duration;
|
||||
|
||||
use aether_runtime_state::RuntimeState;
|
||||
use futures_util::StreamExt as _;
|
||||
use reqwest::{redirect::Policy, Client};
|
||||
use semver::Version;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use tracing::{info, warn};
|
||||
|
||||
use crate::provider_transport::{set_xai_client_version, xai_client_version};
|
||||
use crate::AppState;
|
||||
|
||||
/// 官方安装脚本读取的 stable 渠道，响应体是纯文本版本号。
|
||||
const CLI_STABLE_CHANNEL_ENDPOINT: &str = "https://x.ai/cli/stable";
|
||||
/// stable 渠道不可达时（部分部署地区无法直连 x.ai）退回 npm 发布元数据。
|
||||
const CLI_NPM_RELEASE_ENDPOINT: &str = "https://registry.npmjs.org/@xai-official%2Fgrok/latest";
|
||||
const CLI_NPM_PACKAGE: &str = "@xai-official/grok";
|
||||
const PROFILE_CACHE_KEY: &str = "aether:xai:client-profile:v1";
|
||||
const PROFILE_CACHE_TTL: Duration = Duration::from_secs(30 * 24 * 60 * 60);
|
||||
/// xAI 会在发布后很快抬高最低版本，刷新间隔比 Codex 更短。
|
||||
const PROFILE_REFRESH_INTERVAL: Duration = Duration::from_secs(3 * 60 * 60);
|
||||
const RELEASE_CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
|
||||
const RELEASE_REQUEST_TIMEOUT: Duration = Duration::from_secs(30);
|
||||
const MAX_RELEASE_BYTES: usize = 256 * 1024;
|
||||
const CLI_TARGETS: [&str; 6] = [
|
||||
"darwin-arm64",
|
||||
"darwin-x64",
|
||||
"linux-arm64",
|
||||
"linux-x64",
|
||||
"win32-arm64",
|
||||
"win32-x64",
|
||||
];
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct NpmRelease {
|
||||
name: String,
|
||||
version: String,
|
||||
optional_dependencies: BTreeMap<String, String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
struct CachedProfile {
|
||||
version: String,
|
||||
verified_at_unix_secs: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
enum ProfileRefreshError {
|
||||
#[error("Grok CLI release client initialization failed: {0}")]
|
||||
Client(#[from] reqwest::Error),
|
||||
#[error("Grok CLI release request returned HTTP {0}")]
|
||||
HttpStatus(u16),
|
||||
#[error("Grok CLI release response exceeded {MAX_RELEASE_BYTES} bytes")]
|
||||
ResponseTooLarge,
|
||||
#[error("Grok CLI release metadata is invalid")]
|
||||
InvalidMetadata,
|
||||
#[error("Grok CLI release version is older than the active profile")]
|
||||
Rollback,
|
||||
#[error("Grok CLI profile cache operation failed: {0}")]
|
||||
Cache(String),
|
||||
#[error("Grok CLI stable channel failed ({stable}); npm fallback failed ({npm})")]
|
||||
AllSourcesFailed { stable: String, npm: String },
|
||||
}
|
||||
|
||||
fn version_sequence(version: &str) -> Result<u64, ProfileRefreshError> {
|
||||
let parsed = Version::parse(version).map_err(|_| ProfileRefreshError::InvalidMetadata)?;
|
||||
if !parsed.pre.is_empty()
|
||||
|| !parsed.build.is_empty()
|
||||
|| parsed.major > 999
|
||||
|| parsed.minor > 999
|
||||
|| parsed.patch > 999
|
||||
{
|
||||
return Err(ProfileRefreshError::InvalidMetadata);
|
||||
}
|
||||
Ok(1 + parsed.major * 1_000_000 + parsed.minor * 1_000 + parsed.patch)
|
||||
}
|
||||
|
||||
/// stable 渠道只返回一行版本号；任何多余内容都视为异常响应（例如被劫持的 HTML 页面）。
|
||||
fn parse_stable_channel(bytes: &[u8]) -> Result<String, ProfileRefreshError> {
|
||||
if bytes.len() > MAX_RELEASE_BYTES {
|
||||
return Err(ProfileRefreshError::ResponseTooLarge);
|
||||
}
|
||||
let text = std::str::from_utf8(bytes).map_err(|_| ProfileRefreshError::InvalidMetadata)?;
|
||||
let version = text.trim();
|
||||
version_sequence(version)?;
|
||||
Ok(version.to_owned())
|
||||
}
|
||||
|
||||
/// 校验 npm latest 标签及六个平台二进制包来自同一版本发布。
|
||||
fn parse_npm_release(bytes: &[u8]) -> Result<String, ProfileRefreshError> {
|
||||
if bytes.len() > MAX_RELEASE_BYTES {
|
||||
return Err(ProfileRefreshError::ResponseTooLarge);
|
||||
}
|
||||
let release = serde_json::from_slice::<NpmRelease>(bytes)
|
||||
.map_err(|_| ProfileRefreshError::InvalidMetadata)?;
|
||||
version_sequence(&release.version)?;
|
||||
if release.name != CLI_NPM_PACKAGE
|
||||
|| CLI_TARGETS.iter().any(|target| {
|
||||
release
|
||||
.optional_dependencies
|
||||
.get(&format!("{CLI_NPM_PACKAGE}-{target}"))
|
||||
!= Some(&release.version)
|
||||
})
|
||||
{
|
||||
return Err(ProfileRefreshError::InvalidMetadata);
|
||||
}
|
||||
Ok(release.version)
|
||||
}
|
||||
|
||||
fn refresh_enabled_from(value: Option<&str>) -> bool {
|
||||
!value.is_some_and(|value| {
|
||||
matches!(
|
||||
value.trim().to_ascii_lowercase().as_str(),
|
||||
"0" | "false" | "off"
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
fn refresh_enabled() -> bool {
|
||||
refresh_enabled_from(
|
||||
std::env::var("AETHER_XAI_CLIENT_PROFILE_REFRESH")
|
||||
.ok()
|
||||
.as_deref(),
|
||||
)
|
||||
}
|
||||
|
||||
fn fixed_version_from(value: Option<&str>) -> Option<String> {
|
||||
let value = value?.trim();
|
||||
if value.is_empty() || version_sequence(value).is_err() {
|
||||
None
|
||||
} else {
|
||||
Some(value.to_owned())
|
||||
}
|
||||
}
|
||||
|
||||
fn fixed_version_override() -> Option<String> {
|
||||
let value = std::env::var("AETHER_XAI_CLIENT_VERSION").ok()?;
|
||||
let version = fixed_version_from(Some(&value));
|
||||
if version.is_none() {
|
||||
warn!(
|
||||
event_name = "xai_client_profile_fixed_version_invalid",
|
||||
"AETHER_XAI_CLIENT_VERSION is invalid; using cached or built-in profile"
|
||||
);
|
||||
}
|
||||
version
|
||||
}
|
||||
|
||||
fn build_release_client() -> Result<Client, ProfileRefreshError> {
|
||||
Client::builder()
|
||||
.https_only(true)
|
||||
.no_proxy()
|
||||
.redirect(Policy::none())
|
||||
.connect_timeout(RELEASE_CONNECT_TIMEOUT)
|
||||
.timeout(RELEASE_REQUEST_TIMEOUT)
|
||||
.build()
|
||||
.map_err(ProfileRefreshError::Client)
|
||||
}
|
||||
|
||||
async fn fetch_bounded(client: &Client, url: &str) -> Result<Vec<u8>, ProfileRefreshError> {
|
||||
let response = client
|
||||
.get(url)
|
||||
.send()
|
||||
.await
|
||||
.map_err(ProfileRefreshError::Client)?;
|
||||
if !response.status().is_success() {
|
||||
return Err(ProfileRefreshError::HttpStatus(response.status().as_u16()));
|
||||
}
|
||||
if response
|
||||
.content_length()
|
||||
.is_some_and(|length| length > MAX_RELEASE_BYTES as u64)
|
||||
{
|
||||
return Err(ProfileRefreshError::ResponseTooLarge);
|
||||
}
|
||||
|
||||
let mut bytes = Vec::new();
|
||||
let mut stream = response.bytes_stream();
|
||||
while let Some(chunk) = stream.next().await {
|
||||
let chunk = chunk.map_err(ProfileRefreshError::Client)?;
|
||||
if bytes.len().saturating_add(chunk.len()) > MAX_RELEASE_BYTES {
|
||||
return Err(ProfileRefreshError::ResponseTooLarge);
|
||||
}
|
||||
bytes.extend_from_slice(&chunk);
|
||||
}
|
||||
Ok(bytes)
|
||||
}
|
||||
|
||||
async fn fetch_latest_with_fallback<S, SFut, N, NFut>(
|
||||
fetch_stable: S,
|
||||
fetch_npm: N,
|
||||
) -> Result<String, ProfileRefreshError>
|
||||
where
|
||||
S: FnOnce() -> SFut,
|
||||
SFut: Future<Output = Result<String, ProfileRefreshError>>,
|
||||
N: FnOnce() -> NFut,
|
||||
NFut: Future<Output = Result<String, ProfileRefreshError>>,
|
||||
{
|
||||
let stable_error = match fetch_stable().await {
|
||||
Ok(version) => return Ok(version),
|
||||
Err(error) => error,
|
||||
};
|
||||
fetch_npm()
|
||||
.await
|
||||
.map_err(|npm_error| ProfileRefreshError::AllSourcesFailed {
|
||||
stable: stable_error.to_string(),
|
||||
npm: npm_error.to_string(),
|
||||
})
|
||||
}
|
||||
|
||||
async fn fetch_latest_cli_version(client: &Client) -> Result<String, ProfileRefreshError> {
|
||||
fetch_latest_with_fallback(
|
||||
|| async {
|
||||
let bytes = fetch_bounded(client, CLI_STABLE_CHANNEL_ENDPOINT).await?;
|
||||
parse_stable_channel(&bytes)
|
||||
},
|
||||
|| async {
|
||||
let bytes = fetch_bounded(client, CLI_NPM_RELEASE_ENDPOINT).await?;
|
||||
parse_npm_release(&bytes)
|
||||
},
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
fn publish_version(version: &str) -> Result<(), ProfileRefreshError> {
|
||||
set_xai_client_version(version)
|
||||
.map(|_| ())
|
||||
.map_err(|_| ProfileRefreshError::InvalidMetadata)
|
||||
}
|
||||
|
||||
async fn restore_cached_profile(runtime: &RuntimeState) -> Result<(), ProfileRefreshError> {
|
||||
let Some(raw) = runtime
|
||||
.kv_get(PROFILE_CACHE_KEY)
|
||||
.await
|
||||
.map_err(|err| ProfileRefreshError::Cache(err.to_string()))?
|
||||
else {
|
||||
return Ok(());
|
||||
};
|
||||
let cached = serde_json::from_str::<CachedProfile>(&raw)
|
||||
.map_err(|_| ProfileRefreshError::InvalidMetadata)?;
|
||||
if let Some(version) = cached_version_to_restore(&cached, &xai_client_version())? {
|
||||
publish_version(&version)?;
|
||||
info!(
|
||||
event_name = "xai_client_profile_restored",
|
||||
version = %version,
|
||||
verified_at_unix_secs = cached.verified_at_unix_secs,
|
||||
"restored cached Grok CLI profile"
|
||||
);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn cached_version_to_restore(
|
||||
cached: &CachedProfile,
|
||||
active_version: &str,
|
||||
) -> Result<Option<String>, ProfileRefreshError> {
|
||||
let cached_sequence = version_sequence(&cached.version)?;
|
||||
let active_sequence = version_sequence(active_version)?;
|
||||
Ok((cached_sequence > active_sequence).then(|| cached.version.clone()))
|
||||
}
|
||||
|
||||
async fn refresh_once_with_fetch<F, Fut>(
|
||||
runtime: &RuntimeState,
|
||||
fixed_version: Option<&str>,
|
||||
refresh_is_enabled: bool,
|
||||
fetch_latest: F,
|
||||
) -> Result<String, ProfileRefreshError>
|
||||
where
|
||||
F: FnOnce() -> Fut,
|
||||
Fut: Future<Output = Result<String, ProfileRefreshError>>,
|
||||
{
|
||||
if let Some(version) = fixed_version {
|
||||
publish_version(version)?;
|
||||
return Ok(version.to_owned());
|
||||
}
|
||||
|
||||
if let Err(error) = restore_cached_profile(runtime).await {
|
||||
// 缓存损坏或暂时不可用不应阻断官方版本检查；当前进程继续使用旧画像。
|
||||
warn!(
|
||||
event_name = "xai_client_profile_cache_restore_failed",
|
||||
error = %error,
|
||||
"could not restore cached Grok CLI profile"
|
||||
);
|
||||
}
|
||||
if !refresh_is_enabled {
|
||||
return Ok(xai_client_version());
|
||||
}
|
||||
|
||||
let version = fetch_latest().await?;
|
||||
let current = xai_client_version();
|
||||
if version_sequence(&version)? < version_sequence(¤t)? {
|
||||
return Err(ProfileRefreshError::Rollback);
|
||||
}
|
||||
|
||||
let cached = CachedProfile {
|
||||
version: version.clone(),
|
||||
verified_at_unix_secs: chrono::Utc::now().timestamp().max(0) as u64,
|
||||
};
|
||||
let serialized =
|
||||
serde_json::to_string(&cached).map_err(|_| ProfileRefreshError::InvalidMetadata)?;
|
||||
publish_version(&version)?;
|
||||
if let Err(error) = runtime
|
||||
.kv_set(PROFILE_CACHE_KEY, serialized, Some(PROFILE_CACHE_TTL))
|
||||
.await
|
||||
{
|
||||
// 本地版本已经完成原子替换；缓存写失败只影响下次进程启动的恢复。
|
||||
warn!(
|
||||
event_name = "xai_client_profile_cache_write_failed",
|
||||
error = %error,
|
||||
"published Grok CLI profile locally but could not persist the cache"
|
||||
);
|
||||
}
|
||||
Ok(version)
|
||||
}
|
||||
|
||||
async fn refresh_once(runtime: &RuntimeState) -> Result<String, ProfileRefreshError> {
|
||||
let fixed_version = fixed_version_override();
|
||||
refresh_once_with_fetch(
|
||||
runtime,
|
||||
fixed_version.as_deref(),
|
||||
refresh_enabled(),
|
||||
|| async {
|
||||
let client = build_release_client()?;
|
||||
fetch_latest_cli_version(&client).await
|
||||
},
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn prewarm(runtime: &RuntimeState) -> Result<String, String> {
|
||||
refresh_once(runtime).await.map_err(|err| err.to_string())
|
||||
}
|
||||
|
||||
pub(crate) fn spawn_worker(app: AppState) -> tokio::task::JoinHandle<()> {
|
||||
crate::task_runtime::spawn_singleton_worker(
|
||||
app,
|
||||
crate::task_runtime::TASK_KEY_XAI_CLIENT_PROFILE,
|
||||
|app| async move {
|
||||
let mut interval = tokio::time::interval(PROFILE_REFRESH_INTERVAL);
|
||||
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
|
||||
// 启动阶段由 prewarm 完成一次检查；后台任务只负责后续定时刷新，避免重复建连。
|
||||
interval.tick().await;
|
||||
loop {
|
||||
interval.tick().await;
|
||||
match refresh_once(app.runtime_state()).await {
|
||||
Ok(version) => info!(
|
||||
event_name = "xai_client_profile_refreshed",
|
||||
version = %version,
|
||||
"refreshed Grok CLI profile"
|
||||
),
|
||||
Err(error) => warn!(
|
||||
event_name = "xai_client_profile_refresh_failed",
|
||||
error = %error,
|
||||
"keeping the previous Grok CLI profile after refresh failure"
|
||||
),
|
||||
}
|
||||
}
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::{
|
||||
atomic::{AtomicBool, Ordering},
|
||||
Mutex, OnceLock,
|
||||
};
|
||||
use std::time::Duration;
|
||||
|
||||
use aether_runtime_state::{MemoryRuntimeStateConfig, RuntimeState};
|
||||
|
||||
use super::{
|
||||
cached_version_to_restore, fetch_latest_with_fallback, fixed_version_from,
|
||||
parse_npm_release, parse_stable_channel, refresh_enabled_from, refresh_once_with_fetch,
|
||||
CachedProfile, ProfileRefreshError, PROFILE_CACHE_KEY,
|
||||
};
|
||||
use crate::provider_transport::{set_xai_client_version, xai_client_version};
|
||||
|
||||
static PROFILE_TEST_LOCK: OnceLock<Mutex<()>> = OnceLock::new();
|
||||
|
||||
struct VersionRestore(String);
|
||||
|
||||
impl Drop for VersionRestore {
|
||||
fn drop(&mut self) {
|
||||
let _ = set_xai_client_version(&self.0);
|
||||
}
|
||||
}
|
||||
|
||||
fn version_restore_guard() -> (std::sync::MutexGuard<'static, ()>, VersionRestore) {
|
||||
let lock = PROFILE_TEST_LOCK.get_or_init(|| Mutex::new(()));
|
||||
let guard = lock
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||
let restore = VersionRestore(xai_client_version());
|
||||
(guard, restore)
|
||||
}
|
||||
|
||||
fn npm_release(version: &str) -> serde_json::Value {
|
||||
let mut deps = serde_json::Map::new();
|
||||
for target in super::CLI_TARGETS {
|
||||
deps.insert(
|
||||
format!("@xai-official/grok-{target}"),
|
||||
serde_json::Value::String(version.to_string()),
|
||||
);
|
||||
}
|
||||
serde_json::json!({
|
||||
"name": "@xai-official/grok",
|
||||
"version": version,
|
||||
"optionalDependencies": deps,
|
||||
})
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stable_channel_accepts_only_a_bare_release_version() {
|
||||
assert_eq!(parse_stable_channel(b"1.0.46\n").unwrap(), "1.0.46");
|
||||
assert!(parse_stable_channel(b"<html>1.0.46</html>").is_err());
|
||||
assert!(parse_stable_channel(b"1.0.47-alpha.1").is_err());
|
||||
assert!(parse_stable_channel(b"").is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn npm_release_requires_every_platform_binary_at_the_same_version() {
|
||||
let body = npm_release("1.0.46");
|
||||
assert_eq!(
|
||||
parse_npm_release(&serde_json::to_vec(&body).unwrap()).unwrap(),
|
||||
"1.0.46"
|
||||
);
|
||||
|
||||
let mut mismatched = npm_release("1.0.46");
|
||||
mismatched["optionalDependencies"]["@xai-official/grok-linux-x64"] =
|
||||
serde_json::Value::String("1.0.45".to_string());
|
||||
assert!(parse_npm_release(&serde_json::to_vec(&mismatched).unwrap()).is_err());
|
||||
|
||||
let mut wrong_package = npm_release("1.0.46");
|
||||
wrong_package["name"] = serde_json::Value::String("grok".to_string());
|
||||
assert!(parse_npm_release(&serde_json::to_vec(&wrong_package).unwrap()).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn refresh_and_fixed_version_environment_policies_are_strict() {
|
||||
assert!(!refresh_enabled_from(Some("off")));
|
||||
assert!(!refresh_enabled_from(Some(" FALSE ")));
|
||||
assert!(refresh_enabled_from(None));
|
||||
assert_eq!(
|
||||
fixed_version_from(Some(" 1.0.46 ")).as_deref(),
|
||||
Some("1.0.46")
|
||||
);
|
||||
assert!(fixed_version_from(Some("1.0.46-beta.1")).is_none());
|
||||
assert!(fixed_version_from(Some("1.0")).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cached_profile_never_rewinds_active_profile() {
|
||||
let cached = CachedProfile {
|
||||
version: "1.0.50".to_string(),
|
||||
verified_at_unix_secs: 1,
|
||||
};
|
||||
assert_eq!(
|
||||
cached_version_to_restore(&cached, "1.0.46").unwrap(),
|
||||
Some("1.0.50".to_string())
|
||||
);
|
||||
assert_eq!(cached_version_to_restore(&cached, "1.1.0").unwrap(), None);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn npm_is_used_only_when_the_stable_channel_fails() {
|
||||
let npm_called = AtomicBool::new(false);
|
||||
let version = fetch_latest_with_fallback(
|
||||
|| async { Ok("1.0.46".to_string()) },
|
||||
|| async {
|
||||
npm_called.store(true, Ordering::SeqCst);
|
||||
Ok("1.0.45".to_string())
|
||||
},
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(version, "1.0.46");
|
||||
assert!(!npm_called.load(Ordering::SeqCst));
|
||||
|
||||
let version = fetch_latest_with_fallback(
|
||||
|| async { Err(ProfileRefreshError::HttpStatus(503)) },
|
||||
|| async { Ok("1.0.46".to_string()) },
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(version, "1.0.46");
|
||||
|
||||
let result = fetch_latest_with_fallback(
|
||||
|| async { Err(ProfileRefreshError::HttpStatus(503)) },
|
||||
|| async { Err(ProfileRefreshError::InvalidMetadata) },
|
||||
)
|
||||
.await;
|
||||
assert!(matches!(
|
||||
result,
|
||||
Err(ProfileRefreshError::AllSourcesFailed { .. })
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn cache_hit_is_restored_without_network_when_refresh_is_disabled() {
|
||||
let (_lock, _restore) = version_restore_guard();
|
||||
set_xai_client_version("1.0.46").unwrap();
|
||||
let runtime = RuntimeState::memory(MemoryRuntimeStateConfig::default());
|
||||
runtime
|
||||
.kv_set(
|
||||
PROFILE_CACHE_KEY,
|
||||
serde_json::to_string(&CachedProfile {
|
||||
version: "1.0.50".to_string(),
|
||||
verified_at_unix_secs: 1,
|
||||
})
|
||||
.unwrap(),
|
||||
Some(Duration::from_secs(60)),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let result = refresh_once_with_fetch(&runtime, None, false, || async {
|
||||
Err(ProfileRefreshError::HttpStatus(599))
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(result, "1.0.50");
|
||||
assert_eq!(xai_client_version(), "1.0.50");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn refresh_failure_keeps_previous_profile() {
|
||||
let (_lock, _restore) = version_restore_guard();
|
||||
let runtime = RuntimeState::memory(MemoryRuntimeStateConfig::default());
|
||||
let before = xai_client_version();
|
||||
let result = refresh_once_with_fetch(&runtime, None, true, || async {
|
||||
Err(ProfileRefreshError::HttpStatus(503))
|
||||
})
|
||||
.await;
|
||||
|
||||
assert!(matches!(result, Err(ProfileRefreshError::HttpStatus(503))));
|
||||
assert_eq!(xai_client_version(), before);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn successful_refresh_publishes_and_caches_version() {
|
||||
let (_lock, _restore) = version_restore_guard();
|
||||
set_xai_client_version("1.0.46").unwrap();
|
||||
let runtime = RuntimeState::memory(MemoryRuntimeStateConfig::default());
|
||||
let result =
|
||||
refresh_once_with_fetch(&runtime, None, true, || async { Ok("1.0.51".to_string()) })
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(result, "1.0.51");
|
||||
assert_eq!(xai_client_version(), "1.0.51");
|
||||
let cached = runtime.kv_get(PROFILE_CACHE_KEY).await.unwrap().unwrap();
|
||||
assert!(cached.contains("\"1.0.51\""));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn fixed_version_override_skips_network_and_publishes_version() {
|
||||
let (_lock, _restore) = version_restore_guard();
|
||||
let runtime = RuntimeState::memory(MemoryRuntimeStateConfig::default());
|
||||
let fetch_called = AtomicBool::new(false);
|
||||
let result = refresh_once_with_fetch(&runtime, Some("1.0.60"), true, || async {
|
||||
fetch_called.store(true, Ordering::SeqCst);
|
||||
Ok("1.0.61".to_string())
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(result, "1.0.60");
|
||||
assert!(!fetch_called.load(Ordering::SeqCst));
|
||||
assert_eq!(xai_client_version(), "1.0.60");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn rollback_is_rejected_without_replacing_profile() {
|
||||
let (_lock, _restore) = version_restore_guard();
|
||||
set_xai_client_version("1.0.60").unwrap();
|
||||
let runtime = RuntimeState::memory(MemoryRuntimeStateConfig::default());
|
||||
let result =
|
||||
refresh_once_with_fetch(&runtime, None, true, || async { Ok("1.0.59".to_string()) })
|
||||
.await;
|
||||
|
||||
assert!(matches!(result, Err(ProfileRefreshError::Rollback)));
|
||||
assert_eq!(xai_client_version(), "1.0.60");
|
||||
}
|
||||
}
|
||||
@@ -2198,10 +2198,6 @@ struct OpenAIResponsesClientToolResultState {
|
||||
item_started: bool,
|
||||
}
|
||||
|
||||
fn is_responses_web_search_tool(name: &str) -> bool {
|
||||
matches!(name, "web_search" | "web_search_preview")
|
||||
}
|
||||
|
||||
fn web_search_query_from_arguments(arguments: &str) -> String {
|
||||
serde_json::from_str::<Value>(arguments)
|
||||
.ok()
|
||||
@@ -3526,12 +3522,14 @@ impl OpenAIResponsesClientEmitter {
|
||||
.map(|(_, child_name)| child_name.to_string())
|
||||
.unwrap_or_else(|| name.clone());
|
||||
let emitted_namespace = namespaced_tool.map(|(namespace, _)| namespace.to_string());
|
||||
let is_namespaced_tool = namespaced_tool.is_some();
|
||||
let web_search = self
|
||||
.namespace_tool_aliases
|
||||
.emits_hosted_web_search_call(&name);
|
||||
let state = self.tool_calls.entry(index).or_default();
|
||||
state.call_id = call_id.clone();
|
||||
state.name = emitted_name;
|
||||
state.namespace = emitted_namespace;
|
||||
state.web_search = !is_namespaced_tool && is_responses_web_search_tool(&name);
|
||||
state.web_search = web_search;
|
||||
let emitted_call_id = state.call_id.clone();
|
||||
let emitted_name = state.name.clone();
|
||||
let emitted_namespace = state.namespace.clone();
|
||||
@@ -6379,6 +6377,53 @@ mod tests {
|
||||
assert!(!sse.contains("response.function_call_arguments.delta"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openai_responses_client_emitter_keeps_client_declared_web_search_function_as_function_call()
|
||||
{
|
||||
let mut emitter = OpenAIResponsesClientEmitter::with_report_context(&json!({
|
||||
"original_request_body": {
|
||||
"tools": [{
|
||||
"type": "function",
|
||||
"name": "web_search",
|
||||
"parameters": {"type": "object", "properties": {"query": {"type": "string"}}}
|
||||
}]
|
||||
}
|
||||
}));
|
||||
let mut bytes = Vec::new();
|
||||
for event in [
|
||||
CanonicalStreamEvent::ToolCallStart {
|
||||
index: 0,
|
||||
call_id: "call_ws_1".to_string(),
|
||||
name: "web_search".to_string(),
|
||||
},
|
||||
CanonicalStreamEvent::ToolCallArgumentsDelta {
|
||||
index: 0,
|
||||
arguments: r#"{"query":"today tech"}"#.to_string(),
|
||||
},
|
||||
CanonicalStreamEvent::Finish {
|
||||
finish_reason: Some("tool_calls".to_string()),
|
||||
usage: None,
|
||||
},
|
||||
] {
|
||||
bytes.extend(
|
||||
emitter
|
||||
.emit(CanonicalStreamFrame {
|
||||
id: "resp_123".to_string(),
|
||||
model: "gemini-3.8-flash".to_string(),
|
||||
event,
|
||||
})
|
||||
.expect("event should encode"),
|
||||
);
|
||||
}
|
||||
|
||||
let sse = String::from_utf8(bytes).expect("sse should be utf8");
|
||||
assert!(!sse.contains("web_search_call"));
|
||||
assert!(sse.contains(r#""type":"function_call""#));
|
||||
assert!(sse.contains(r#""call_id":"call_ws_1""#));
|
||||
assert!(sse.contains(r#""name":"web_search""#));
|
||||
assert!(sse.contains("response.function_call_arguments.delta"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openai_responses_provider_state_accepts_legacy_outtext_delta_alias() {
|
||||
let mut state = OpenAIResponsesProviderState::default();
|
||||
|
||||
@@ -54,6 +54,7 @@ pub(crate) struct NamespaceToolAliases {
|
||||
by_chat_name: BTreeMap<String, (String, String)>,
|
||||
namespace_tool_indices: BTreeSet<usize>,
|
||||
invalid_namespace_tool_indices: BTreeSet<usize>,
|
||||
client_function_tool_names: BTreeSet<String>,
|
||||
}
|
||||
|
||||
impl NamespaceToolAliases {
|
||||
@@ -171,10 +172,24 @@ impl NamespaceToolAliases {
|
||||
else {
|
||||
return Self::default();
|
||||
};
|
||||
let Some(canonical) = openai_responses_tools_to_canonical(Some(tools)) else {
|
||||
return Self::default();
|
||||
};
|
||||
Self::from_canonical_tools(&canonical)
|
||||
let client_function_tool_names = client_function_tool_names(tools);
|
||||
let mut result = openai_responses_tools_to_canonical(Some(tools))
|
||||
.map(|canonical| Self::from_canonical_tools(&canonical))
|
||||
.unwrap_or_default();
|
||||
result.client_function_tool_names = client_function_tool_names;
|
||||
result
|
||||
}
|
||||
|
||||
/// Whether a tool call named `name` should surface to a Responses client as
|
||||
/// a hosted `web_search_call`. A client that declared its own function or
|
||||
/// custom tool called `web_search` must get a `function_call` back, or it
|
||||
/// cannot answer the call and will echo an unconvertible hosted item.
|
||||
/// When the client declares both a hosted `web_search` tool and a function
|
||||
/// of the same name, the function wins: only the client can answer it.
|
||||
pub(crate) fn emits_hosted_web_search_call(&self, name: &str) -> bool {
|
||||
matches!(name, "web_search" | "web_search_preview")
|
||||
&& self.responses_name(name).is_none()
|
||||
&& !self.client_function_tool_names.contains(name)
|
||||
}
|
||||
|
||||
pub(crate) fn chat_name(&self, namespace: &str, child_name: &str) -> Option<&str> {
|
||||
@@ -227,6 +242,32 @@ impl NamespaceToolAliases {
|
||||
}
|
||||
}
|
||||
|
||||
fn client_function_tool_names(tools: &Value) -> BTreeSet<String> {
|
||||
tools
|
||||
.as_array()
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.filter_map(Value::as_object)
|
||||
.filter(|tool| {
|
||||
tool.get("type")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.is_none_or(|tool_type| {
|
||||
tool_type.eq_ignore_ascii_case("function")
|
||||
|| tool_type.eq_ignore_ascii_case("custom")
|
||||
})
|
||||
})
|
||||
.filter_map(|tool| {
|
||||
non_empty_string(tool.get("name")).or_else(|| {
|
||||
["function", "custom"]
|
||||
.iter()
|
||||
.find_map(|key| non_empty_string(tool.get(*key)?.get("name")))
|
||||
})
|
||||
})
|
||||
.map(ToOwned::to_owned)
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub(crate) fn canonical_tool_is_responses_namespace(tool: &CanonicalToolDefinition) -> bool {
|
||||
raw_responses_tool(tool).is_some_and(|raw| {
|
||||
raw.get("type")
|
||||
@@ -414,6 +455,111 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
fn hosted_web_search(tools: Value, name: &str) -> bool {
|
||||
NamespaceToolAliases::from_report_context(&json!({
|
||||
"original_request_body": {"tools": tools}
|
||||
}))
|
||||
.emits_hosted_web_search_call(name)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn hosted_web_search_call_is_reserved_for_undeclared_search_names() {
|
||||
let schema = json!({"type": "object", "properties": {"query": {"type": "string"}}});
|
||||
let cases = [
|
||||
("no tools", json!([]), "web_search", true),
|
||||
(
|
||||
"hosted tool",
|
||||
json!([{"type": "web_search"}]),
|
||||
"web_search",
|
||||
true,
|
||||
),
|
||||
(
|
||||
"hosted preview tool",
|
||||
json!([{"type": "web_search_preview"}]),
|
||||
"web_search_preview",
|
||||
true,
|
||||
),
|
||||
(
|
||||
"function tool",
|
||||
json!([{"type": "function", "name": "web_search", "parameters": schema}]),
|
||||
"web_search",
|
||||
false,
|
||||
),
|
||||
(
|
||||
"function tool named preview",
|
||||
json!([{"type": "function", "name": "web_search_preview", "parameters": schema}]),
|
||||
"web_search_preview",
|
||||
false,
|
||||
),
|
||||
(
|
||||
"custom tool",
|
||||
json!([{"type": "custom", "name": "web_search"}]),
|
||||
"web_search",
|
||||
false,
|
||||
),
|
||||
(
|
||||
"tool without type",
|
||||
json!([{"name": "web_search", "parameters": schema}]),
|
||||
"web_search",
|
||||
false,
|
||||
),
|
||||
(
|
||||
"chat-shaped function tool",
|
||||
json!([{"type": "function", "function": {"name": "web_search", "parameters": schema}}]),
|
||||
"web_search",
|
||||
false,
|
||||
),
|
||||
(
|
||||
"chat-shaped custom tool",
|
||||
json!([{"type": "custom", "custom": {"name": "web_search"}}]),
|
||||
"web_search",
|
||||
false,
|
||||
),
|
||||
(
|
||||
"hosted and function tool together",
|
||||
json!([
|
||||
{"type": "web_search"},
|
||||
{"type": "function", "name": "web_search", "parameters": schema}
|
||||
]),
|
||||
"web_search",
|
||||
false,
|
||||
),
|
||||
(
|
||||
"unrelated function tool",
|
||||
json!([{"type": "function", "name": "lookup", "parameters": schema}]),
|
||||
"web_search",
|
||||
true,
|
||||
),
|
||||
(
|
||||
"non-search name",
|
||||
json!([{"type": "web_search"}]),
|
||||
"lookup",
|
||||
false,
|
||||
),
|
||||
];
|
||||
|
||||
for (label, tools, name, expected) in cases {
|
||||
assert_eq!(hosted_web_search(tools, name), expected, "{label}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn namespaced_web_search_child_is_not_a_hosted_web_search_call() {
|
||||
assert!(!hosted_web_search(
|
||||
json!([{
|
||||
"type": "namespace",
|
||||
"name": "mcp__search",
|
||||
"description": "Search tools",
|
||||
"tools": [{
|
||||
"type": "function",
|
||||
"name": "web_search",
|
||||
"parameters": {"type": "object", "properties": {}}
|
||||
}]
|
||||
}]),
|
||||
"web_search"
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn namespace_aliases_are_unique_bounded_and_prefix_safe() {
|
||||
let long_namespace = format!("namespace__{}", "n".repeat(120));
|
||||
|
||||
@@ -275,7 +275,7 @@ pub fn to_raw(canonical: &CanonicalResponse, report_context: &Value, compact: bo
|
||||
}));
|
||||
}
|
||||
let namespaced_tool = namespace_tool_aliases.responses_name(name);
|
||||
if namespaced_tool.is_none() && is_responses_web_search_tool(name) {
|
||||
if namespace_tool_aliases.emits_hosted_web_search_call(name) {
|
||||
output.push(json!({
|
||||
"type": "web_search_call",
|
||||
"id": id,
|
||||
@@ -601,10 +601,6 @@ fn openai_responses_output_format_from_mime_type(mime_type: &str) -> String {
|
||||
.to_string()
|
||||
}
|
||||
|
||||
fn is_responses_web_search_tool(name: &str) -> bool {
|
||||
matches!(name, "web_search" | "web_search_preview")
|
||||
}
|
||||
|
||||
fn web_search_query_from_value(input: &Value) -> String {
|
||||
input
|
||||
.get("query")
|
||||
@@ -647,6 +643,70 @@ mod tests {
|
||||
assert!(body["completed_at"].as_i64().is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn responses_response_builder_keeps_client_declared_web_search_function_as_function_call() {
|
||||
let report_context = json!({
|
||||
"original_request_body": {
|
||||
"tools": [{
|
||||
"type": "function",
|
||||
"name": "web_search",
|
||||
"parameters": {"type": "object", "properties": {"query": {"type": "string"}}}
|
||||
}]
|
||||
}
|
||||
});
|
||||
let response = CanonicalResponse {
|
||||
id: "resp_test".to_string(),
|
||||
model: "gemini-3.8-flash".to_string(),
|
||||
content: vec![CanonicalContentBlock::ToolUse {
|
||||
id: "call_ws_1".to_string(),
|
||||
name: "web_search".to_string(),
|
||||
input: json!({"query": "today tech"}),
|
||||
extensions: BTreeMap::new(),
|
||||
}],
|
||||
outputs: Vec::new(),
|
||||
stop_reason: Some(CanonicalStopReason::ToolUse),
|
||||
usage: None,
|
||||
extensions: BTreeMap::new(),
|
||||
};
|
||||
|
||||
let body = to_raw(&response, &report_context, false);
|
||||
|
||||
let item = &body["output"][0];
|
||||
assert_eq!(item["type"], "function_call");
|
||||
assert_eq!(item["name"], "web_search");
|
||||
assert_eq!(item["call_id"], "call_ws_1");
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(item["arguments"].as_str().unwrap()).unwrap(),
|
||||
json!({"query": "today tech"})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn responses_response_builder_emits_web_search_call_for_hosted_web_search_tool() {
|
||||
let report_context = json!({
|
||||
"original_request_body": {"tools": [{"type": "web_search"}]}
|
||||
});
|
||||
let response = CanonicalResponse {
|
||||
id: "resp_test".to_string(),
|
||||
model: "gpt-5-5-low".to_string(),
|
||||
content: vec![CanonicalContentBlock::ToolUse {
|
||||
id: "call_ws_1".to_string(),
|
||||
name: "web_search".to_string(),
|
||||
input: json!({"query": "today tech"}),
|
||||
extensions: BTreeMap::new(),
|
||||
}],
|
||||
outputs: Vec::new(),
|
||||
stop_reason: Some(CanonicalStopReason::ToolUse),
|
||||
usage: None,
|
||||
extensions: BTreeMap::new(),
|
||||
};
|
||||
|
||||
let body = to_raw(&response, &report_context, false);
|
||||
|
||||
assert_eq!(body["output"][0]["type"], "web_search_call");
|
||||
assert_eq!(body["output"][0]["action"]["query"], "today tech");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn responses_response_builder_restores_namespaced_chat_tool_identity() {
|
||||
let report_context = json!({
|
||||
|
||||
@@ -973,6 +973,66 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn streams_gemini_web_search_function_call_to_responses_function_call_for_function_tool() {
|
||||
let mut context = report_context("gemini:generate_content", "openai:responses");
|
||||
context["original_request_body"] = json!({
|
||||
"model": "gemini-3.8-flash",
|
||||
"tools": [{
|
||||
"type": "function",
|
||||
"name": "web_search",
|
||||
"parameters": {"type": "object", "properties": {"query": {"type": "string"}}}
|
||||
}]
|
||||
});
|
||||
let mut matrix = StreamingStandardFormatMatrix::default();
|
||||
let mut output = matrix
|
||||
.transform_line(
|
||||
&context,
|
||||
data_line(json!({
|
||||
"response": {
|
||||
"responseId": "resp_ws_function",
|
||||
"modelVersion": "gemini-3.8-flash",
|
||||
"candidates": [{
|
||||
"index": 0,
|
||||
"content": {
|
||||
"role": "model",
|
||||
"parts": [{
|
||||
"thoughtSignature": "signature",
|
||||
"functionCall": {
|
||||
"name": "web_search",
|
||||
"args": {"query": "conpty newline"},
|
||||
"id": "call_109312"
|
||||
}
|
||||
}]
|
||||
},
|
||||
"finishReason": "STOP"
|
||||
}]
|
||||
}
|
||||
})),
|
||||
)
|
||||
.expect("Gemini function call should transform");
|
||||
output.extend(matrix.finish(&context).expect("stream should finish"));
|
||||
|
||||
let events = json_data_events(&output);
|
||||
let completed = events
|
||||
.iter()
|
||||
.find(|event| event["type"] == "response.completed")
|
||||
.expect("response should complete");
|
||||
let items = completed["response"]["output"]
|
||||
.as_array()
|
||||
.expect("completed response should carry output");
|
||||
assert!(
|
||||
items.iter().all(|item| item["type"] != "web_search_call"),
|
||||
"{items:?}"
|
||||
);
|
||||
let call = items
|
||||
.iter()
|
||||
.find(|item| item["type"] == "function_call")
|
||||
.expect("function tool call should stay a function_call");
|
||||
assert_eq!(call["name"], "web_search");
|
||||
assert_eq!(call["call_id"], "call_109312");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn terminal_observer_marks_malformed_gemini_function_call_as_failure() {
|
||||
let context = report_context("gemini:generate_content", "openai:responses");
|
||||
|
||||
@@ -53,6 +53,75 @@ pub fn extract_provider_reasoning_effort_from_body(value: Option<&Value>) -> Opt
|
||||
.and_then(Value::as_str)
|
||||
})
|
||||
.and_then(normalize_provider_reasoning_effort)
|
||||
.or_else(|| {
|
||||
// Gemini also nests its payload one level down, so both the flat
|
||||
// `generateContent` body and the `v1internal` envelope that carries it are read.
|
||||
extract_gemini_reasoning_effort_from_body(object).or_else(|| {
|
||||
object
|
||||
.get("request")
|
||||
.and_then(Value::as_object)
|
||||
.and_then(extract_gemini_reasoning_effort_from_body)
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
/// Gemini `generateContent` states its reasoning depth inside
|
||||
/// `generationConfig.thinkingConfig`, either as a symbolic `thinkingLevel` or as a token
|
||||
/// `thinkingBudget`. Both camelCase and snake_case spellings are read so that a captured client
|
||||
/// body and a converted provider body resolve to the same label.
|
||||
///
|
||||
/// `includeThoughts` alone is a visibility flag, not a depth, so it never produces a label.
|
||||
fn extract_gemini_reasoning_effort_from_body(
|
||||
object: &serde_json::Map<String, Value>,
|
||||
) -> Option<String> {
|
||||
let generation_config = object
|
||||
.get("generationConfig")
|
||||
.or_else(|| object.get("generation_config"))
|
||||
.and_then(Value::as_object)?;
|
||||
let thinking_config = generation_config
|
||||
.get("thinkingConfig")
|
||||
.or_else(|| generation_config.get("thinking_config"))
|
||||
.and_then(Value::as_object)?;
|
||||
|
||||
if let Some(level) = thinking_config
|
||||
.get("thinkingLevel")
|
||||
.or_else(|| thinking_config.get("thinking_level"))
|
||||
.and_then(Value::as_str)
|
||||
.and_then(normalize_gemini_thinking_level)
|
||||
{
|
||||
return Some(level);
|
||||
}
|
||||
|
||||
thinking_config
|
||||
.get("thinkingBudget")
|
||||
.or_else(|| thinking_config.get("thinking_budget"))
|
||||
.and_then(Value::as_u64)
|
||||
.map(|budget| {
|
||||
// `0` disables reasoning outright. The shared budget ladder collapses it into `low`,
|
||||
// which would report an explicitly disabled request as a shallow one.
|
||||
if budget == 0 {
|
||||
"none".to_string()
|
||||
} else {
|
||||
aether_ai_formats::formats::openai::shared::map_thinking_budget_to_openai_reasoning_effort(budget)
|
||||
.to_string()
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
/// Gemini also emits the protobuf enum spelling (`THINKING_LEVEL_HIGH`); the level itself is what
|
||||
/// the badge vocabulary understands, so the enum prefix is stripped before normalizing.
|
||||
///
|
||||
/// `THINKING_LEVEL_UNSPECIFIED` is the enum's "no explicit level" member, not a depth. It is
|
||||
/// rejected rather than surfaced, otherwise the badge would read `unspecified`.
|
||||
fn normalize_gemini_thinking_level(value: &str) -> Option<String> {
|
||||
let normalized = value.trim().to_ascii_lowercase();
|
||||
let normalized = normalized
|
||||
.strip_prefix("thinking_level_")
|
||||
.unwrap_or(normalized.as_str());
|
||||
if normalized == "unspecified" {
|
||||
return None;
|
||||
}
|
||||
normalize_provider_reasoning_effort(normalized)
|
||||
}
|
||||
|
||||
fn normalize_provider_reasoning_effort(value: &str) -> Option<String> {
|
||||
@@ -3588,6 +3657,157 @@ mod tests {
|
||||
assert_eq!(usage.provider_service_tier(), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gemini_thinking_level_supplies_reasoning_effort_for_both_body_spellings() {
|
||||
let mut usage = sample_usage();
|
||||
usage.provider_request_body = Some(json!({
|
||||
"generationConfig": {
|
||||
"thinkingConfig": { "includeThoughts": true, "thinkingLevel": "HIGH" }
|
||||
}
|
||||
}));
|
||||
|
||||
assert_eq!(usage.provider_reasoning_effort().as_deref(), Some("high"));
|
||||
|
||||
// The converted provider body keeps snake_case keys, and the client body may carry the
|
||||
// protobuf enum spelling. Both must land on the same badge vocabulary.
|
||||
usage.provider_request_body = Some(json!({
|
||||
"generation_config": {
|
||||
"thinking_config": { "thinking_level": "thinking_level_medium" }
|
||||
}
|
||||
}));
|
||||
|
||||
assert_eq!(usage.provider_reasoning_effort().as_deref(), Some("medium"));
|
||||
|
||||
usage.provider_request_body = Some(json!({
|
||||
"generationConfig": {
|
||||
"thinkingConfig": { "thinkingLevel": " low " }
|
||||
}
|
||||
}));
|
||||
|
||||
assert_eq!(usage.provider_reasoning_effort().as_deref(), Some("low"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gemini_thinking_budget_supplies_reasoning_effort_without_collapsing_zero() {
|
||||
let mut usage = sample_usage();
|
||||
usage.provider_request_body = Some(json!({
|
||||
"generationConfig": { "thinkingConfig": { "thinkingBudget": 8192 } }
|
||||
}));
|
||||
|
||||
assert_eq!(usage.provider_reasoning_effort().as_deref(), Some("xhigh"));
|
||||
|
||||
// `0` disables reasoning. The shared budget ladder maps 0..=1664 to `low`, which would
|
||||
// report an explicitly disabled request as shallow, so the Gemini path reports `none`.
|
||||
usage.provider_request_body = Some(json!({
|
||||
"generationConfig": { "thinkingConfig": { "thinkingBudget": 0 } }
|
||||
}));
|
||||
|
||||
assert_eq!(usage.provider_reasoning_effort().as_deref(), Some("none"));
|
||||
|
||||
usage.provider_request_body = Some(json!({
|
||||
"generation_config": { "thinking_config": { "thinking_budget": 1280 } }
|
||||
}));
|
||||
|
||||
assert_eq!(usage.provider_reasoning_effort().as_deref(), Some("low"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gemini_thinking_config_without_level_or_budget_yields_no_reasoning_effort() {
|
||||
let mut usage = sample_usage();
|
||||
usage.provider_request_body = Some(json!({
|
||||
"generationConfig": { "thinkingConfig": { "includeThoughts": true } }
|
||||
}));
|
||||
|
||||
assert_eq!(usage.provider_reasoning_effort(), None);
|
||||
|
||||
// A level-less, budget-less config must not fall back to metadata either: the captured
|
||||
// body is authoritative and it says nothing about depth.
|
||||
usage.request_metadata = Some(json!({ "provider_reasoning_effort": "max" }));
|
||||
assert_eq!(usage.provider_reasoning_effort(), None);
|
||||
|
||||
// `THINKING_LEVEL_UNSPECIFIED` is the enum's "no explicit level" member, not a depth.
|
||||
usage.request_metadata = None;
|
||||
usage.provider_request_body = Some(json!({
|
||||
"generationConfig": {
|
||||
"thinkingConfig": { "thinkingLevel": "THINKING_LEVEL_UNSPECIFIED" }
|
||||
}
|
||||
}));
|
||||
|
||||
assert_eq!(usage.provider_reasoning_effort(), None);
|
||||
}
|
||||
|
||||
/// The v1internal envelope nests the real `generateContent` payload under `request`. This is
|
||||
/// the shape the Antigravity/Gemini CLI transports actually send upstream, so the extraction
|
||||
/// has to descend into it or every converted `openai:chat -> gemini` request loses its badge.
|
||||
#[test]
|
||||
fn gemini_thinking_config_is_read_from_the_v1internal_envelope() {
|
||||
let mut usage = sample_usage();
|
||||
usage.provider_request_body = Some(json!({
|
||||
"model": "gemini-3.8-flash-tiered",
|
||||
"project": "aicode-consumers",
|
||||
"requestId": "req-1",
|
||||
"requestType": "agent",
|
||||
"userAgent": "vscode/1.X.X (Antigravity/4.3.0)",
|
||||
"request": {
|
||||
"contents": [{ "role": "user", "parts": [{ "text": "hi" }] }],
|
||||
"generationConfig": {
|
||||
"maxOutputTokens": 65536,
|
||||
"thinkingConfig": { "includeThoughts": true, "thinkingLevel": "high" }
|
||||
}
|
||||
}
|
||||
}));
|
||||
|
||||
assert_eq!(usage.provider_reasoning_effort().as_deref(), Some("high"));
|
||||
|
||||
usage.provider_request_body = Some(json!({
|
||||
"model": "gemini-3.8-flash-tiered",
|
||||
"request": {
|
||||
"generation_config": {
|
||||
"thinking_config": { "include_thoughts": true, "thinking_budget": 32768 }
|
||||
}
|
||||
}
|
||||
}));
|
||||
|
||||
assert_eq!(usage.provider_reasoning_effort().as_deref(), Some("xhigh"));
|
||||
|
||||
usage.provider_request_body = Some(json!({
|
||||
"model": "gemini-3.8-flash-tiered",
|
||||
"request": {
|
||||
"generationConfig": { "thinkingConfig": { "thinkingBudget": 0 } }
|
||||
}
|
||||
}));
|
||||
|
||||
assert_eq!(usage.provider_reasoning_effort().as_deref(), Some("none"));
|
||||
}
|
||||
|
||||
/// A converted request that carries only `maxOutputTokens` must stay badge-less rather than
|
||||
/// picking up a depth from somewhere else in the envelope.
|
||||
#[test]
|
||||
fn v1internal_envelope_without_thinking_config_yields_no_reasoning_effort() {
|
||||
let mut usage = sample_usage();
|
||||
usage.provider_request_body = Some(json!({
|
||||
"model": "gemini-3.8-flash-tiered",
|
||||
"project": "aicode-consumers",
|
||||
"request": {
|
||||
"contents": [{ "role": "user", "parts": [{ "text": "hi" }] }],
|
||||
"generationConfig": { "maxOutputTokens": 65536 }
|
||||
}
|
||||
}));
|
||||
|
||||
assert_eq!(usage.provider_reasoning_effort(), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gemini_thinking_config_does_not_shadow_explicit_effort_fields() {
|
||||
let mut usage = sample_usage();
|
||||
usage.provider_request_body = Some(json!({
|
||||
"reasoning_effort": "max",
|
||||
"generationConfig": { "thinkingConfig": { "thinkingLevel": "low" } }
|
||||
}));
|
||||
|
||||
assert_eq!(usage.provider_reasoning_effort().as_deref(), Some("max"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn requested_and_provider_reasoning_efforts_remain_independent() {
|
||||
let mut usage = sample_usage();
|
||||
|
||||
@@ -43,7 +43,7 @@ pub use quota::{
|
||||
provider_pool_key_model_quota_hard_blocked, provider_pool_key_quota_hard_blocked,
|
||||
provider_pool_key_scheduling_label, provider_pool_member_quota_snapshot,
|
||||
provider_pool_quota_metadata_provider_type, provider_pool_quota_metadata_updated_at,
|
||||
provider_pool_quota_snapshot_updated_at,
|
||||
provider_pool_quota_snapshot_updated_at, provider_pool_reset_deadline_elapsed,
|
||||
};
|
||||
pub use quota_refresh::ProviderPoolQuotaRequestSpec;
|
||||
pub use service::ProviderPoolService;
|
||||
|
||||
@@ -857,7 +857,14 @@ fn provider_pool_reset_deadline_unix_secs(
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn provider_pool_reset_deadline_elapsed(
|
||||
/// Whether the quota window's reset deadline has already passed.
|
||||
///
|
||||
/// The deadline is taken from `reset_at`/`next_reset_at` when present, otherwise
|
||||
/// derived from `reset_seconds`/`reset_after_seconds` anchored at the window (or
|
||||
/// fallback) observation time. Scheduling ignores exhausted windows once this is
|
||||
/// true; read paths can reuse the same predicate so the displayed quota matches
|
||||
/// the scheduling decision after a reset.
|
||||
pub fn provider_pool_reset_deadline_elapsed(
|
||||
item: &Map<String, Value>,
|
||||
fallback_observed_at: Option<u64>,
|
||||
now_unix_secs: u64,
|
||||
|
||||
@@ -200,7 +200,8 @@ pub use windsurf::{
|
||||
pub use xai::{
|
||||
extract_xai_user_id_from_auth_config, extract_xai_user_id_from_value,
|
||||
insert_cli_identity_headers, insert_cli_identity_headers_if_needed, is_xai_provider_transport,
|
||||
resolved_xai_request_base_url, resolved_xai_upstream_base_url,
|
||||
should_attach_cli_identity_headers, xai_auth_uses_api, xai_uses_official_api, XAI_API_BASE_URL,
|
||||
XAI_CHAT_PROXY_BASE_URL, XAI_PROVIDER_TYPE,
|
||||
resolved_xai_request_base_url, resolved_xai_upstream_base_url, set_xai_client_version,
|
||||
should_attach_cli_identity_headers, xai_auth_uses_api, xai_client_version,
|
||||
xai_uses_official_api, XAI_API_BASE_URL, XAI_CHAT_PROXY_BASE_URL, XAI_DEFAULT_CLIENT_VERSION,
|
||||
XAI_PROVIDER_TYPE,
|
||||
};
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
pub mod video;
|
||||
|
||||
use std::collections::BTreeMap;
|
||||
use std::sync::{OnceLock, RwLock};
|
||||
|
||||
use aether_ai_formats::normalize_api_format_alias;
|
||||
use serde_json::Value;
|
||||
@@ -10,7 +11,10 @@ use crate::snapshot::GatewayProviderTransportSnapshot;
|
||||
pub const XAI_PROVIDER_TYPE: &str = "xai";
|
||||
pub const XAI_CHAT_PROXY_BASE_URL: &str = "https://cli-chat-proxy.grok.com/v1";
|
||||
pub const XAI_API_BASE_URL: &str = "https://api.x.ai/v1";
|
||||
pub const XAI_CLIENT_VERSION: &str = "0.2.120";
|
||||
/// 内置的 Grok CLI 版本；网关后台任务会用官方发布版本覆盖它。
|
||||
///
|
||||
/// cli-chat-proxy 会对过旧的版本直接返回 426，因此这里只作为发布检查不可用时的兜底。
|
||||
pub const XAI_DEFAULT_CLIENT_VERSION: &str = "1.0.46";
|
||||
pub const XAI_TOKEN_AUTH_HEADER: &str = "x-xai-token-auth";
|
||||
pub const XAI_TOKEN_AUTH_VALUE: &str = "xai-grok-cli";
|
||||
pub const XAI_CLIENT_VERSION_HEADER: &str = "x-grok-client-version";
|
||||
@@ -19,8 +23,37 @@ pub const XAI_CLIENT_IDENTIFIER_VALUE: &str = "grok-shell";
|
||||
pub const XAI_AUTHENTICATE_RESPONSE_HEADER: &str = "x-authenticateresponse";
|
||||
pub const XAI_AUTHENTICATE_RESPONSE_VALUE: &str = "authenticate-response";
|
||||
|
||||
static ACTIVE_CLIENT_VERSION: OnceLock<RwLock<String>> = OnceLock::new();
|
||||
|
||||
fn active_client_version() -> &'static RwLock<String> {
|
||||
ACTIVE_CLIENT_VERSION.get_or_init(|| RwLock::new(XAI_DEFAULT_CLIENT_VERSION.to_owned()))
|
||||
}
|
||||
|
||||
/// 返回当前发布的 Grok CLI 版本快照。
|
||||
pub fn xai_client_version() -> String {
|
||||
active_client_version()
|
||||
.read()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
||||
.clone()
|
||||
}
|
||||
|
||||
/// 原子替换当前 Grok CLI 版本，返回替换前的版本；版本校验由发布检查器负责，这里只拒绝明显非法值。
|
||||
pub fn set_xai_client_version(version: &str) -> Result<String, &'static str> {
|
||||
let version = version.trim();
|
||||
if version.is_empty()
|
||||
|| version.len() > 64
|
||||
|| !version.bytes().all(|byte| (33..=126).contains(&byte))
|
||||
{
|
||||
return Err("invalid Grok CLI version");
|
||||
}
|
||||
let mut current = active_client_version()
|
||||
.write()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner);
|
||||
Ok(std::mem::replace(&mut *current, version.to_owned()))
|
||||
}
|
||||
|
||||
pub fn xai_cli_user_agent() -> String {
|
||||
format!("xai-grok-workspace/{XAI_CLIENT_VERSION}")
|
||||
format!("xai-grok-workspace/{}", xai_client_version())
|
||||
}
|
||||
|
||||
pub fn is_xai_provider_transport(transport: &GatewayProviderTransportSnapshot) -> bool {
|
||||
@@ -91,10 +124,11 @@ pub fn should_attach_cli_identity_headers(
|
||||
}
|
||||
|
||||
pub fn insert_cli_identity_headers(headers: &mut BTreeMap<String, String>) {
|
||||
let client_version = xai_client_version();
|
||||
let user_agent = xai_cli_user_agent();
|
||||
for (name, value) in [
|
||||
(XAI_TOKEN_AUTH_HEADER, XAI_TOKEN_AUTH_VALUE),
|
||||
(XAI_CLIENT_VERSION_HEADER, XAI_CLIENT_VERSION),
|
||||
(XAI_CLIENT_VERSION_HEADER, client_version.as_str()),
|
||||
("user-agent", user_agent.as_str()),
|
||||
(XAI_CLIENT_IDENTIFIER_HEADER, XAI_CLIENT_IDENTIFIER_VALUE),
|
||||
(
|
||||
@@ -441,4 +475,31 @@ mod tests {
|
||||
Some(r#"{"api_key":"xai-key","using_api":true}"#)
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cli_identity_headers_follow_published_client_version() {
|
||||
use super::{
|
||||
insert_cli_identity_headers, set_xai_client_version, xai_client_version,
|
||||
XAI_CLIENT_VERSION_HEADER,
|
||||
};
|
||||
|
||||
let previous = xai_client_version();
|
||||
assert!(set_xai_client_version("").is_err());
|
||||
assert!(set_xai_client_version("1.0 .1").is_err());
|
||||
assert_eq!(xai_client_version(), previous);
|
||||
|
||||
set_xai_client_version(" 9.8.7 ").expect("valid version");
|
||||
let mut headers = BTreeMap::new();
|
||||
insert_cli_identity_headers(&mut headers);
|
||||
set_xai_client_version(&previous).expect("restore version");
|
||||
|
||||
assert_eq!(
|
||||
headers.get(XAI_CLIENT_VERSION_HEADER).map(String::as_str),
|
||||
Some("9.8.7")
|
||||
);
|
||||
assert_eq!(
|
||||
headers.get("user-agent").map(String::as_str),
|
||||
Some("xai-grok-workspace/9.8.7")
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -744,6 +744,39 @@ mod tests {
|
||||
assert!(cleared.get("requested_reasoning_effort").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gemini_thinking_config_is_derived_into_client_and_provider_reasoning_metadata() {
|
||||
let client_body = json!({
|
||||
"generationConfig": {
|
||||
"thinkingConfig": { "includeThoughts": true, "thinkingLevel": "HIGH" }
|
||||
}
|
||||
});
|
||||
let provider_body = json!({
|
||||
"generation_config": {
|
||||
"thinking_config": { "thinking_budget": 8192 }
|
||||
}
|
||||
});
|
||||
|
||||
let metadata = attach_client_request_body_metadata(
|
||||
Some(json!({ "trace_id": "trace-1" })),
|
||||
Some(&client_body),
|
||||
)
|
||||
.expect("metadata should remain");
|
||||
assert_eq!(metadata["requested_reasoning_effort"], "high");
|
||||
|
||||
let metadata = attach_provider_request_body_metadata(
|
||||
Some(metadata),
|
||||
Some("gemini:generate_content"),
|
||||
Some("gemini-3.8-flash"),
|
||||
Some("gemini-3.8-flash"),
|
||||
Some(&provider_body),
|
||||
)
|
||||
.expect("metadata should remain");
|
||||
|
||||
assert_eq!(metadata["requested_reasoning_effort"], "high");
|
||||
assert_eq!(metadata["provider_reasoning_effort"], "xhigh");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_request_body_metadata_uses_final_provider_body_as_source_of_truth() {
|
||||
let metadata = Some(json!({
|
||||
|
||||
@@ -635,6 +635,53 @@ describe('RequestDetailDrawer settlement pricing', () => {
|
||||
})
|
||||
})
|
||||
|
||||
it('shows the Gemini thinkingConfig reasoning effort in the model header', async () => {
|
||||
apiMocks.getRequestDetail.mockResolvedValue({
|
||||
...buildEmbeddingDetail(),
|
||||
id: 'usage-gemini-thinking',
|
||||
request_id: 'usage-gemini-thinking',
|
||||
model: 'gemini-3.8-flash',
|
||||
request_type: 'chat',
|
||||
requested_reasoning_effort: 'xhigh',
|
||||
reasoning_effort: 'high',
|
||||
request_body: {
|
||||
generationConfig: {
|
||||
thinkingConfig: { includeThoughts: true, thinkingLevel: 'HIGH' },
|
||||
},
|
||||
},
|
||||
provider_request_body: {
|
||||
generationConfig: { thinkingConfig: { thinkingBudget: 8192 } },
|
||||
},
|
||||
})
|
||||
|
||||
let isOpen!: Ref<boolean>
|
||||
const Host = defineComponent({
|
||||
setup() {
|
||||
isOpen = ref(false)
|
||||
return () => h(RequestDetailDrawer, {
|
||||
isOpen: isOpen.value,
|
||||
requestId: 'usage-gemini-thinking',
|
||||
})
|
||||
},
|
||||
})
|
||||
|
||||
const root = document.createElement('div')
|
||||
document.body.appendChild(root)
|
||||
const app = createApp(Host)
|
||||
app.mount(root)
|
||||
mountedApps.push({ app, root })
|
||||
|
||||
isOpen.value = true
|
||||
await nextTick()
|
||||
|
||||
await vi.waitFor(() => {
|
||||
expect(document.body.querySelector('[data-request-detail-model-display]')?.textContent)
|
||||
.toContain('gemini-3.8-flash')
|
||||
expect(document.body.querySelector('[data-request-detail-model-badge="reasoning"]')?.textContent?.trim())
|
||||
.toBe('xhigh -> high')
|
||||
})
|
||||
})
|
||||
|
||||
it('lets a newer final-provider summary clear facts cached from an earlier candidate', async () => {
|
||||
apiMocks.getRequestDetail.mockResolvedValue({
|
||||
...buildEmbeddingDetail(),
|
||||
|
||||
@@ -440,6 +440,42 @@ describe('UsageRecordsTable', () => {
|
||||
.toBe('Fast')
|
||||
})
|
||||
|
||||
it('shows the Gemini thinkingLevel reasoning effort next to the model name', () => {
|
||||
const root = mountUsageRecordsTable([buildRecord({
|
||||
model: 'gemini-3.8-flash',
|
||||
requested_reasoning_effort: 'high',
|
||||
reasoning_effort: 'high',
|
||||
})])
|
||||
|
||||
expect(root.textContent).toContain('gemini-3.8-flash')
|
||||
const badge = root.querySelector('[data-usage-model-badge="reasoning"]')
|
||||
expect(badge?.textContent?.trim()).toBe('high')
|
||||
expect(badge?.getAttribute('title')).toBe('Reasoning: high')
|
||||
})
|
||||
|
||||
it('shows the Gemini thinkingLevel mapping when request and provider disagree', () => {
|
||||
const root = mountUsageRecordsTable([buildRecord({
|
||||
model: 'gemini-3.8-flash',
|
||||
requested_reasoning_effort: 'xhigh',
|
||||
reasoning_effort: 'high',
|
||||
})])
|
||||
|
||||
expect(root.textContent).toContain('xhigh -> high')
|
||||
expect(root.querySelector('[data-usage-model-badge="reasoning"]')?.textContent?.trim())
|
||||
.toBe('xhigh -> high')
|
||||
})
|
||||
|
||||
it('shows a disabled Gemini thinkingBudget as none', () => {
|
||||
const root = mountUsageRecordsTable([buildRecord({
|
||||
model: 'gemini-3.8-flash',
|
||||
requested_reasoning_effort: null,
|
||||
reasoning_effort: 'none',
|
||||
})])
|
||||
|
||||
expect(root.querySelector('[data-usage-model-badge="reasoning"]')?.textContent?.trim())
|
||||
.toBe('none')
|
||||
})
|
||||
|
||||
it('shows request reasoning effort while the record is pending', () => {
|
||||
const root = mountUsageRecordsTable([buildRecord({
|
||||
status: 'pending',
|
||||
|
||||
@@ -2063,17 +2063,23 @@ const quotaProgressMap = computed<Record<string, QuotaProgressItem[]>>(() => {
|
||||
const quotaProgressDisplayMap = computed<Record<string, QuotaProgressDisplayItem[]>>(() => {
|
||||
const map: Record<string, QuotaProgressDisplayItem[]> = {}
|
||||
for (const key of keyPage.value.keys) {
|
||||
map[key.key_id] = (quotaProgressMap.value[key.key_id] || []).map(item => ({
|
||||
label: getQuotaProgressLabel(item.label),
|
||||
remainingPercent: item.remainingPercent,
|
||||
resetText: getQuotaProgressResetDisplayText(item),
|
||||
meterText: item.numericOnly
|
||||
? item.detail || formatQuotaValue(item.remainingPercent)
|
||||
: getQuotaProgressMeterDisplayText(item),
|
||||
barClass: getQuotaRemainingBarColorByRemaining(item.remainingPercent),
|
||||
meterClass: getQuotaRemainingClassByRemaining(item.remainingPercent),
|
||||
numericOnly: item.numericOnly,
|
||||
}))
|
||||
map[key.key_id] = (quotaProgressMap.value[key.key_id] || []).map(item => {
|
||||
// 倒计时归零表示窗口已越过重置时间点：按“已重置”展示 100%，
|
||||
// 不再显示重置前的旧用量文本，与后端读取口径、调度口径保持一致。
|
||||
const expired = !item.numericOnly && getQuotaProgressCountdown(item)?.isExpired === true
|
||||
const remainingPercent = expired ? 100 : item.remainingPercent
|
||||
return {
|
||||
label: getQuotaProgressLabel(item.label),
|
||||
remainingPercent,
|
||||
resetText: getQuotaProgressResetDisplayText(item),
|
||||
meterText: item.numericOnly
|
||||
? item.detail || formatQuotaValue(remainingPercent)
|
||||
: getQuotaProgressMeterDisplayText(item, remainingPercent, expired),
|
||||
barClass: getQuotaRemainingBarColorByRemaining(remainingPercent),
|
||||
meterClass: getQuotaRemainingClassByRemaining(remainingPercent),
|
||||
numericOnly: item.numericOnly,
|
||||
}
|
||||
})
|
||||
}
|
||||
return map
|
||||
})
|
||||
@@ -3563,10 +3569,15 @@ function getQuotaProgressResetDisplayText(item: QuotaProgressItem): string {
|
||||
return ''
|
||||
}
|
||||
|
||||
function getQuotaProgressMeterDisplayText(item: QuotaProgressItem): string {
|
||||
const detail = item.detail?.trim() || ''
|
||||
function getQuotaProgressMeterDisplayText(
|
||||
item: QuotaProgressItem,
|
||||
remainingPercent = item.remainingPercent,
|
||||
suppressDetail = false,
|
||||
): string {
|
||||
// 窗口已重置时忽略重置前的旧用量文本，直接显示归一化后的剩余百分比。
|
||||
const detail = suppressDetail ? '' : (item.detail?.trim() || '')
|
||||
if (!shouldHideQuotaProgressDetailText(detail) && detail) return detail
|
||||
return `${item.remainingPercent.toFixed(1)}%`
|
||||
return `${remainingPercent.toFixed(1)}%`
|
||||
}
|
||||
|
||||
function getQuotaFallbackText(key: PoolKeyDetail): string | null {
|
||||
|
||||
@@ -35,6 +35,10 @@ const routeMocks = vi.hoisted(() => ({
|
||||
}),
|
||||
}))
|
||||
|
||||
const countdownMocks = vi.hoisted(() => ({
|
||||
isExpired: false,
|
||||
}))
|
||||
|
||||
const proxyStoreMocks = vi.hoisted(() => ({
|
||||
ensureLoaded: vi.fn(),
|
||||
}))
|
||||
@@ -106,7 +110,7 @@ vi.mock('@/composables/useCountdownTimer', async () => {
|
||||
start: vi.fn(),
|
||||
}),
|
||||
getCodexResetCountdown: () => ({
|
||||
isExpired: false,
|
||||
isExpired: countdownMocks.isExpired,
|
||||
text: '1h',
|
||||
}),
|
||||
}
|
||||
@@ -604,6 +608,7 @@ beforeEach(() => {
|
||||
window.sessionStorage.clear()
|
||||
routeMocks.patchQuery.mockClear()
|
||||
proxyStoreMocks.ensureLoaded.mockClear()
|
||||
countdownMocks.isExpired = false
|
||||
|
||||
endpointMocks.getPoolOverview.mockReset()
|
||||
endpointMocks.getPoolSchedulingPresets.mockReset()
|
||||
@@ -736,6 +741,44 @@ describe('PoolManagement Codex cycle stats mode', () => {
|
||||
expect(root.textContent).toContain('生图')
|
||||
})
|
||||
|
||||
it('shows expired Codex quota windows as fully restored', async () => {
|
||||
countdownMocks.isExpired = true
|
||||
const expiredKey = createPoolKey('codex', {
|
||||
status_snapshot: {
|
||||
oauth: { code: 'valid' },
|
||||
account: { code: 'ok', blocked: false },
|
||||
quota: {
|
||||
code: 'ok',
|
||||
exhausted: false,
|
||||
provider_type: 'codex',
|
||||
windows: [
|
||||
{
|
||||
code: '5h',
|
||||
label: '5H',
|
||||
scope: 'account',
|
||||
remaining_ratio: 0.12,
|
||||
reset_at: 1_700_000_000,
|
||||
window_minutes: 300,
|
||||
},
|
||||
],
|
||||
},
|
||||
},
|
||||
})
|
||||
endpointMocks.getPoolOverview.mockResolvedValue({ items: [createOverview('codex')] })
|
||||
endpointMocks.listPoolKeys.mockResolvedValue(createKeyPage(expiredKey))
|
||||
endpointMocks.getProvider.mockResolvedValue(createProvider('codex'))
|
||||
|
||||
const root = mountPoolManagement()
|
||||
await settle()
|
||||
|
||||
// 倒计时归零后进度条与文本按“已重置”显示 100%，不再显示旧百分比。
|
||||
expect(root.querySelector('[data-testid="pool-quota-meter-text"]')?.textContent?.trim()).toBe('100.0%')
|
||||
const track = root.querySelector('[data-testid="pool-quota-progress-track"]')
|
||||
const bar = track?.firstElementChild as HTMLElement | null
|
||||
expect(bar?.style.width).toBe('100%')
|
||||
expect(root.querySelector('[data-testid="pool-quota-reset-text"]')).toBeNull()
|
||||
})
|
||||
|
||||
it('labels Codex quota by the actual refresh window duration', async () => {
|
||||
const monthlyCodexKey = createPoolKey('codex', {
|
||||
status_snapshot: {
|
||||
|
||||
Reference in New Issue
Block a user