From e582eb73d2d94aebf8a664e362eaa334de437da1 Mon Sep 17 00:00:00 2001 From: zyycn Date: Mon, 27 Jul 2026 02:26:02 +0800 Subject: [PATCH] fix(xai): preserve unknown fields, scope history sanitize, guard cache-route tools, batch selector reads Rebuild history items with in-place Map surgery so unknown upstream fields survive instead of evaporating through a field whitelist. Narrow history sanitization to the known Grok-internal injection sites rather than stripping every null and any key named phase at any depth. Ensure the prompt-cache route's injected web_search/x_search tools can never be actually invoked when the client sent its own tools. Batch the selector's per-account observed-model-support and cooldown reads instead of a serial Redis round trip each. One plan scope failing no longer aborts the whole provider catalog refresh. Co-authored-by: Claude --- .../providers/xai/src/credential/catalog.rs | 30 +++- .../providers/xai/src/credential/selector.rs | 44 +++-- .../providers/xai/src/transport/request.rs | 134 +++++++++------- .../providers/xai/tests/credential/catalog.rs | 59 ++++++- .../providers/xai/tests/transport/request.rs | 151 ++++++++++++++++-- 5 files changed, 327 insertions(+), 91 deletions(-) diff --git a/backend/crates/providers/xai/src/credential/catalog.rs b/backend/crates/providers/xai/src/credential/catalog.rs index e0aff4323..d24df203d 100644 --- a/backend/crates/providers/xai/src/credential/catalog.rs +++ b/backend/crates/providers/xai/src/credential/catalog.rs @@ -1008,22 +1008,46 @@ impl GrokCredentialCatalogService { .filter(eligible_catalog_candidate) .collect::>(); let groups = catalog_candidates_by_scope(candidates)?; + // 单个套餐失败只跳过该套餐，不阻断其他套餐的目录同步。 let mut fetched = Vec::with_capacity(groups.len()); + let mut last_error = None; for (scope, candidates) in groups { - fetched.push(self.fetch_scope_catalog(scope, candidates).await?); + match self.fetch_scope_catalog(scope.clone(), candidates).await { + Ok(catalog) => fetched.push(catalog), + Err(error) => { + tracing::warn!( + scope = scope.as_str(), + error = %error, + "skipping failed xAI plan catalog scope" + ); + last_error = Some(error); + } + } + } + if fetched.is_empty() { + return Err(last_error.unwrap_or(GrokCredentialCatalogError::NoEligibleCredential)); } fetched.sort_by(|left, right| left.scope.cmp(&right.scope)); let models = strict_model_union(&fetched)?; let observed_at = Utc::now(); for fetched in fetched { - self.cache + let scope = fetched.scope.clone(); + if self + .cache .replace(GrokPlanCatalog { scope: fetched.scope, observed_at, seed: fetched.seed, }) .await - .map_err(|_| GrokCredentialCatalogError::Cache)?; + .is_err() + { + // cache 是可重建 TTL 数据，单套餐写入失败不阻断目录发布。 + tracing::warn!( + scope = scope.as_str(), + "failed to cache xAI plan catalog scope" + ); + } } self.publish_models(&models); diff --git a/backend/crates/providers/xai/src/credential/selector.rs b/backend/crates/providers/xai/src/credential/selector.rs index c9fac09ca..8ad962aa1 100644 --- a/backend/crates/providers/xai/src/credential/selector.rs +++ b/backend/crates/providers/xai/src/credential/selector.rs @@ -72,17 +72,30 @@ impl GrokAccountSessionSelector { .await .map_err(|_| GrokSessionSelectorError::Unavailable)?; self.quota.prepare_scheduling(&accounts).await; + // 目录支持度按套餐 scope 去重后并发批量读取，避免逐账号串行往返。 + let account_scopes = accounts + .iter() + .map(|account| GrokCatalogScope::for_account(account).ok()) + .collect::>(); + let unique_scopes = account_scopes + .iter() + .flatten() + .cloned() + .collect::>(); + let support_reads = futures::future::join_all(unique_scopes.iter().map(|scope| { + self.catalog_cache + .observed_model_support(scope, request.upstream_model().as_str()) + })) + .await; + let support_by_scope = unique_scopes + .into_iter() + .zip(support_reads) + .collect::>(); let mut catalog_eligible = Vec::new(); - for account in accounts { - let support = match GrokCatalogScope::for_account(&account) { - Ok(scope) => { - self.catalog_cache - .observed_model_support(&scope, request.upstream_model().as_str()) - .await - } - Err(error) => Err(error), - }; - if !matches!(support, Ok(Some(false))) { + for (account, scope) in accounts.into_iter().zip(account_scopes) { + let unsupported = scope + .is_some_and(|scope| matches!(support_by_scope.get(&scope), Some(Ok(Some(false))))); + if !unsupported { catalog_eligible.push(account); } } @@ -220,9 +233,16 @@ impl GrokAccountSessionSelector { accounts: &[ProviderAccount], ) -> std::collections::BTreeMap { let now = SystemTime::now(); + // 冷却状态并发批量读取，避免逐账号串行往返。 + let reads = futures::future::join_all( + accounts + .iter() + .map(|account| self.cooldowns.read(account.id())), + ) + .await; let mut cooled = std::collections::BTreeMap::new(); - for account in accounts { - let Ok(Some(cooldown)) = self.cooldowns.read(account.id()).await else { + for (account, read) in accounts.iter().zip(reads) { + let Ok(Some(cooldown)) = read else { continue; }; if cooldown.credential_revision() != account.revision() { diff --git a/backend/crates/providers/xai/src/transport/request.rs b/backend/crates/providers/xai/src/transport/request.rs index d09960d9d..afe782398 100644 --- a/backend/crates/providers/xai/src/transport/request.rs +++ b/backend/crates/providers/xai/src/transport/request.rs @@ -59,6 +59,10 @@ const SESSION_FIELDS: &[&str] = &[ const FOREIGN_CLIENT_METADATA_FIELDS: &[&str] = &["x-openai-subagent"]; const MAX_SESSION_SEED_BYTES: usize = 1_024; const GROK_CACHE_ROUTE_TOOLS: &[&str] = &["web_search", "x_search"]; +/// Grok CLI 在历史条目上注入、Grok Build 无法反序列化的内部键；只在已知 +/// 注入位置剥离，避免误伤工具 schema 或输出中恰好同名的键。 +const GROK_INTERNAL_HISTORY_KEYS: &[&str] = + &["phase", "internal_chat_message_metadata_passthrough"]; /// 保留客户端 OpenAI Responses object 的 xAI 上游请求。 pub struct GrokResponsesRequest { @@ -324,17 +328,16 @@ fn enable_grok_prompt_cache_route(body: &mut Map) { let Value::Array(tools) = tools else { return; }; - let had_client_tools = !tools.is_empty(); + // 缓存路由不得改变模型可调用的工具集：客户端带工具时无法既保留其 + // tool_choice 语义又阻止注入的原生搜索被调用，因此直接不注入。 + if !tools.is_empty() { + return; + } for tool in GROK_CACHE_ROUTE_TOOLS { - if !has_tool_type(tools, tool) { - tools.push(json_object([("type", Value::String((*tool).to_owned()))])); - } - } - // 无客户端工具时禁用工具调用，使原生搜索声明只承担缓存路由作用； - // 混合工具请求保留客户端原有的 tool_choice 语义。 - if !had_client_tools { - body.insert("tool_choice".to_owned(), Value::String("none".to_owned())); + tools.push(json_object([("type", Value::String((*tool).to_owned()))])); } + // 禁用工具调用，使原生搜索声明只承担缓存路由作用。 + body.insert("tool_choice".to_owned(), Value::String("none".to_owned())); } fn explicit_session_seed(request: &GenerateRequest, body: &Map) -> Option { @@ -2212,11 +2215,14 @@ impl ToolNormalizer { role => role, }; let content = self.normalize_message_content(item.get("content"), role)?; - Ok(Map::from_iter([ - ("type".to_owned(), Value::String("message".to_owned())), - ("role".to_owned(), Value::String(role.to_owned())), - ("content".to_owned(), content), - ])) + // 透明代理合法形态：只覆盖归一化字段并剥离 Grok 注入键，未知官方 + // 字段原样保留。 + let mut converted = item.clone(); + strip_grok_internal_keys(&mut converted); + converted.insert("type".to_owned(), Value::String("message".to_owned())); + converted.insert("role".to_owned(), Value::String(role.to_owned())); + converted.insert("content".to_owned(), content); + Ok(converted) } fn normalize_message_content( @@ -2322,12 +2328,14 @@ impl ToolNormalizer { if !namespace.is_empty() { name = self.alias(ToolIdentity::new(ToolKind::Function, namespace, &name)); } - Ok(Map::from_iter([ - ("type".to_owned(), Value::String("function_call".to_owned())), - ("call_id".to_owned(), Value::String(call_id.to_owned())), - ("name".to_owned(), Value::String(name)), - ("arguments".to_owned(), Value::String(arguments)), - ])) + let mut converted = item.clone(); + strip_grok_internal_keys(&mut converted); + converted.remove("namespace"); + converted.insert("type".to_owned(), Value::String("function_call".to_owned())); + converted.insert("call_id".to_owned(), Value::String(call_id.to_owned())); + converted.insert("name".to_owned(), Value::String(name)); + converted.insert("arguments".to_owned(), Value::String(arguments)); + Ok(converted) } fn normalize_custom_tool_call_input( @@ -2348,12 +2356,15 @@ impl ToolNormalizer { )])) .map_err(|_| GrokRequestEncodeError::InvalidRequestNormalization)?; let alias = self.alias(identity); - Ok(Map::from_iter([ - ("type".to_owned(), Value::String("function_call".to_owned())), - ("call_id".to_owned(), Value::String(call_id.to_owned())), - ("name".to_owned(), Value::String(alias)), - ("arguments".to_owned(), Value::String(arguments)), - ])) + let mut converted = item.clone(); + strip_grok_internal_keys(&mut converted); + converted.remove("namespace"); + converted.remove("input"); + converted.insert("type".to_owned(), Value::String("function_call".to_owned())); + converted.insert("call_id".to_owned(), Value::String(call_id.to_owned())); + converted.insert("name".to_owned(), Value::String(alias)); + converted.insert("arguments".to_owned(), Value::String(arguments)); + Ok(converted) } fn normalize_function_call_output_input( @@ -2370,14 +2381,15 @@ impl ToolNormalizer { } output => Value::String(encode_tool_output(output)?), }; - Ok(Map::from_iter([ - ( - "type".to_owned(), - Value::String("function_call_output".to_owned()), - ), - ("call_id".to_owned(), Value::String(call_id.to_owned())), - ("output".to_owned(), output), - ])) + let mut converted = item.clone(); + strip_grok_internal_keys(&mut converted); + converted.insert( + "type".to_owned(), + Value::String("function_call_output".to_owned()), + ); + converted.insert("call_id".to_owned(), Value::String(call_id.to_owned())); + converted.insert("output".to_owned(), output); + Ok(converted) } fn normalize_function_output_blocks( @@ -2606,6 +2618,8 @@ fn require_content_source( fn sanitize_reasoning_input(item: &Map) -> Map { let mut converted = copy_non_null_history_fields(item, &["id", "summary", "content", "encrypted_content"]); + // Grok CLI 会在 summary/content 条目上注入 phase 等内部键。 + strip_grok_internal_entry_keys(&mut converted, &["summary", "content"]); converted.insert("type".to_owned(), Value::String("reasoning".to_owned())); if converted .get("encrypted_content") @@ -2679,6 +2693,14 @@ fn sanitize_native_history_input(item: &Map, item_type: &str) -> _ => &[], }; let mut converted = copy_non_null_history_fields(item, fields); + // Grok CLI 会在 shell_call 的 action object 里注入内部键与 null 占位 + // 字段（如 `timeout_ms: null`）；只在这一层剥离，深层内容原样保留。 + if item_type == "shell_call" + && let Some(Value::Object(action)) = converted.get_mut("action") + { + strip_grok_internal_keys(action); + action.retain(|_, value| !value.is_null()); + } converted.insert("type".to_owned(), Value::String(item_type.to_owned())); converted } @@ -2688,36 +2710,28 @@ fn copy_non_null_history_fields(item: &Map, fields: &[&str]) -> M .iter() .filter_map(|field| { item.get(*field) - .and_then(sanitize_history_json_value) - .map(|value| ((*field).to_owned(), value)) + .filter(|value| !value.is_null()) + .map(|value| ((*field).to_owned(), value.clone())) }) .collect() } -fn sanitize_history_json_value(value: &Value) -> Option { - match value { - Value::Null => None, - Value::Object(object) => Some(Value::Object( - object - .iter() - .filter(|(key, _)| { - !matches!( - key.as_str(), - "phase" | "internal_chat_message_metadata_passthrough" - ) - }) - .filter_map(|(key, value)| { - sanitize_history_json_value(value).map(|value| (key.clone(), value)) - }) - .collect(), - )), - Value::Array(values) => Some(Value::Array( - values - .iter() - .filter_map(sanitize_history_json_value) - .collect(), - )), - value => Some(value.clone()), +fn strip_grok_internal_keys(object: &mut Map) { + for key in GROK_INTERNAL_HISTORY_KEYS { + object.remove(*key); + } +} + +fn strip_grok_internal_entry_keys(converted: &mut Map, fields: &[&str]) { + for field in fields { + let Some(Value::Array(entries)) = converted.get_mut(*field) else { + continue; + }; + for entry in entries { + if let Some(object) = entry.as_object_mut() { + strip_grok_internal_keys(object); + } + } } } diff --git a/backend/crates/providers/xai/tests/credential/catalog.rs b/backend/crates/providers/xai/tests/credential/catalog.rs index 250e25086..9a76cc147 100644 --- a/backend/crates/providers/xai/tests/credential/catalog.rs +++ b/backend/crates/providers/xai/tests/credential/catalog.rs @@ -512,7 +512,7 @@ async fn quota_exhausted_account_remains_eligible_for_catalog_discovery() { } #[tokio::test] -async fn one_upstream_failure_rejects_the_whole_catalog_cycle() { +async fn sole_plan_failure_rejects_the_catalog_cycle() { let (_, repository) = repository_with_accounts(&[("failed", "subject-failed")]).await; let cache_port: Arc = MemoryGrokCatalogCache::shared(); let service = crate::support::grok_catalog_service( @@ -526,6 +526,63 @@ async fn one_upstream_failure_rejects_the_whole_catalog_cycle() { )); } +#[tokio::test] +async fn failed_plan_scope_is_skipped_and_surviving_plans_still_cache() { + let store = MemoryProviderAccountStore::shared(); + let account_store: Arc = store.clone(); + let repository = GrokCredentialRepository::new(account_store); + // "plan:pro" 按 scope 排序在 "plan:standard" 之前，先消费失败响应。 + let mut pro = create_input("plan-pro", "subject-pro"); + pro.account.plan_type = Some("pro".to_owned()); + seed_input(&store, &pro).await.expect("create pro account"); + seed_input(&store, &create_input("plan-standard", "subject-standard")) + .await + .expect("create standard account"); + let transport = QueueCatalogTransport::from_results([ + Err(GrokModelCatalogTransportError::new( + GrokModelCatalogTransportErrorKind::Unavailable, + )), + Ok(GrokModelCatalogTransportResponse::new( + OFFICIAL_FIXTURE.to_vec(), + None, + )), + ]); + let cache = MemoryGrokCatalogCache::shared(); + let cache_port: Arc = cache.clone(); + let service = crate::support::grok_catalog_service(repository, transport.clone(), cache_port); + + let models = service + .query_models() + .await + .expect("surviving plan still refreshes the catalog"); + + assert_eq!(models.len(), 1); + assert_eq!(models[0].request_model().as_str(), "grok-4.5"); + assert_eq!(transport.calls(), 2); + let standard = store + .account(&account_id("plan-standard")) + .expect("created account"); + let standard_scope = GrokCatalogScope::for_account(&standard).expect("catalog scope"); + assert_eq!( + cache + .observed_model_support(&standard_scope, "grok-4.5") + .await + .expect("cache lookup"), + Some(true) + ); + let pro = store + .account(&account_id("plan-pro")) + .expect("created account"); + let pro_scope = GrokCatalogScope::for_account(&pro).expect("catalog scope"); + assert_eq!( + cache + .observed_model_support(&pro_scope, "grok-4.5") + .await + .expect("cache lookup"), + None + ); +} + #[tokio::test] async fn conflicting_facts_for_same_slug_fail_closed() { let (store, repository) = diff --git a/backend/crates/providers/xai/tests/transport/request.rs b/backend/crates/providers/xai/tests/transport/request.rs index 6d5e46930..868079a35 100644 --- a/backend/crates/providers/xai/tests/transport/request.rs +++ b/backend/crates/providers/xai/tests/transport/request.rs @@ -241,7 +241,7 @@ fn explicit_session_should_enable_the_noop_native_cache_route() { } #[test] -fn explicit_session_should_merge_client_tools_into_the_native_cache_route() { +fn explicit_session_should_not_add_invokable_cache_tools_to_client_tool_requests() { let request = raw_request(json!({ "model": "client", "prompt_cache_key": "conversation-42", @@ -258,17 +258,14 @@ fn explicit_session_should_merge_client_tools_into_the_native_cache_route() { GrokResponsesRequest::encode(&request, "grok-4.5", &client_key()).expect("request"); let body = Value::Object(encoded.body().clone()); + // 缓存路由不得扩大模型可调用的工具集：带客户端工具时不注入原生搜索。 assert_eq!( body.pointer("/tools"), - Some(&json!([ - { - "type": "function", - "name": "read_file", - "parameters": {"type": "object"} - }, - {"type": "web_search"}, - {"type": "x_search"} - ])) + Some(&json!([{ + "type": "function", + "name": "read_file", + "parameters": {"type": "object"} + }])) ); assert_eq!(body.pointer("/tool_choice"), Some(&json!("auto"))); } @@ -372,7 +369,7 @@ fn function_parameters_should_report_the_invalid_nullable_root_field() { } #[test] -fn explicit_session_should_enable_cache_after_codex_additional_tools_normalization() { +fn explicit_session_should_not_inject_cache_tools_after_codex_additional_tools_normalization() { let request = raw_request(json!({ "model": "client", "prompt_cache_key": "conversation-42", @@ -405,8 +402,8 @@ fn explicit_session_should_enable_cache_after_codex_additional_tools_normalizati Some(&json!(["patch"])) ); assert_eq!(body.pointer("/tools/1/name"), Some(&json!("read_file"))); - assert_eq!(body.pointer("/tools/2/type"), Some(&json!("web_search"))); - assert_eq!(body.pointer("/tools/3/type"), Some(&json!("x_search"))); + // 归一化装载的会话工具视同客户端工具：不再追加可被调用的原生搜索。 + assert_eq!(body.pointer("/tools/2"), None); assert_eq!(body.pointer("/tool_choice"), Some(&json!("auto"))); } @@ -660,7 +657,7 @@ fn history_should_rebuild_codex_calls_outputs_shell_and_private_fields() { {"type": "local_shell"} ], "input": [ - {"type": "message", "role": "assistant", "id": "drop", "content": [ + {"type": "message", "role": "assistant", "id": "msg_1", "content": [ {"type": "output_text", "text": "done"}, {"type": "refusal", "refusal": "no"} ]}, @@ -680,7 +677,7 @@ fn history_should_rebuild_codex_calls_outputs_shell_and_private_fields() { let body = Value::Object(encoded.body().clone()); assert_eq!(body.pointer("/input/0/content"), Some(&json!("done\nno"))); - assert_eq!(body.pointer("/input/0/id"), None); + assert_eq!(body.pointer("/input/0/id"), Some(&json!("msg_1"))); assert_eq!(body.pointer("/input/1/type"), Some(&json!("function_call"))); assert_eq!( body.pointer("/input/1/arguments"), @@ -712,6 +709,130 @@ fn history_should_rebuild_codex_calls_outputs_shell_and_private_fields() { assert_eq!(body.pointer("/input/8/role"), Some(&json!("developer"))); } +#[test] +fn history_rebuild_should_preserve_unknown_fields_and_strip_grok_internal_keys() { + let request = raw_request(json!({ + "model": "client", + "input": [ + { + "type": "message", "role": "assistant", "status": "completed", + "phase": "final_answer", + "internal_chat_message_metadata_passthrough": {"turn_id": "t1"}, + "future_message_field": {"keep": true}, + "content": [{"type": "output_text", "text": "hi"}] + }, + { + "type": "function_call", "id": "fc_1", "call_id": "call_1", + "name": "read_file", "arguments": "{}", + "future_call_field": 7, + "internal_chat_message_metadata_passthrough": {"turn_id": "t1"} + }, + { + "type": "function_call_output", "id": "fco_1", "call_id": "call_1", + "output": "ok", "future_output_field": true, + "internal_chat_message_metadata_passthrough": {"turn_id": "t1"} + }, + { + "type": "custom_tool_call", "id": "ctc_1", "call_id": "call_2", + "name": "render", "input": "raw", "status": "completed", + "future_custom_field": "keep", + "internal_chat_message_metadata_passthrough": {"turn_id": "t1"} + } + ] + })); + + let encoded = GrokResponsesRequest::encode(&request, "grok-4.5", &client_key()) + .expect("history normalization"); + let body = Value::Object(encoded.body().clone()); + + for pointer in [ + "/input/0/phase", + "/input/0/internal_chat_message_metadata_passthrough", + "/input/1/internal_chat_message_metadata_passthrough", + "/input/2/internal_chat_message_metadata_passthrough", + "/input/3/internal_chat_message_metadata_passthrough", + "/input/3/input", + ] { + assert_eq!(body.pointer(pointer), None, "field survived at {pointer}"); + } + assert_eq!( + body.pointer("/input/0/future_message_field"), + Some(&json!({"keep": true})) + ); + assert_eq!(body.pointer("/input/0/status"), Some(&json!("completed"))); + assert_eq!(body.pointer("/input/0/content"), Some(&json!("hi"))); + assert_eq!(body.pointer("/input/1/id"), Some(&json!("fc_1"))); + assert_eq!(body.pointer("/input/1/future_call_field"), Some(&json!(7))); + assert_eq!( + body.pointer("/input/2/future_output_field"), + Some(&json!(true)) + ); + assert_eq!(body.pointer("/input/2/output"), Some(&json!("ok"))); + assert_eq!(body.pointer("/input/3/type"), Some(&json!("function_call"))); + assert_eq!( + body.pointer("/input/3/future_custom_field"), + Some(&json!("keep")) + ); + assert_eq!( + body.pointer("/input/3/arguments"), + Some(&json!("{\"input\":\"raw\"}")) + ); +} + +#[test] +fn history_sanitizer_should_only_strip_known_grok_injection_sites() { + let request = raw_request(json!({ + "model": "client", + "input": [ + {"type": "mcp_list_tools", "id": "ml_1", "server_label": "srv", "tools": [ + { + "name": "deploy", + "input_schema": { + "properties": {"phase": {"type": "string"}}, + "required": ["phase"] + } + } + ]}, + { + "type": "mcp_call", "id": "mc_1", "name": "deploy", + "server_label": "srv", "arguments": "{}", + "output": {"result": null, "phase": "keep"} + }, + { + "type": "shell_call", "id": "sh_1", "call_id": "call_9", + "status": "completed", + "action": { + "commands": ["pwd"], + "timeout_ms": null, + "internal_chat_message_metadata_passthrough": {"turn_id": "t1"} + } + } + ] + })); + + let encoded = GrokResponsesRequest::encode(&request, "grok-4.5", &client_key()) + .expect("history normalization"); + let body = Value::Object(encoded.body().clone()); + + // 工具 schema 里恰好叫 phase 的属性与语义 null 不再被误删。 + assert_eq!( + body.pointer("/input/0/tools/0/input_schema/properties/phase"), + Some(&json!({"type": "string"})) + ); + assert_eq!(body.pointer("/input/1/output/result"), Some(&Value::Null)); + assert_eq!(body.pointer("/input/1/output/phase"), Some(&json!("keep"))); + // shell_call action 是已知注入点：内部键与 null 占位字段仍被剥离。 + assert_eq!( + body.pointer("/input/2/action/commands/0"), + Some(&json!("pwd")) + ); + assert_eq!(body.pointer("/input/2/action/timeout_ms"), None); + assert_eq!( + body.pointer("/input/2/action/internal_chat_message_metadata_passthrough"), + None + ); +} + #[test] fn custom_apply_patch_history_should_wrap_raw_input_in_patch_field() { let patch = concat!(