starred/codex-proxy-rs
1
0
Fork 0
mirror of https://github.com/zyycn/codex-proxy-rs.git synced 2026-11-06 08:06:19 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(xai): preserve unknown fields, scope history sanitize, guard cache-route tools, batch selector reads

Browse Source 
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 <noreply@anthropic.com>
This commit is contained in:
zyycn
2026-07-27 02:26:02 +08:00
co-authored by Claude
parent ec252b3b46
commit e582eb73d2
5 changed files with 327 additions and 91 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/crates/providers/xai/src/credential/catalog.rs
+27 -3
View File
@@ -1008,22 +1008,46 @@ impl GrokCredentialCatalogService {
.filter(eligible_catalog_candidate)
.collect::<Vec<_>>();
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);
backend/crates/providers/xai/src/credential/selector.rs
+32 -12
View File
@@ -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::<Vec<_>>();
let unique_scopes = account_scopes
.iter()
.flatten()
.cloned()
.collect::<std::collections::BTreeSet<_>>();
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::<std::collections::BTreeMap<_, _>>();
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<ProviderAccountId, SystemTime> {
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() {
backend/crates/providers/xai/src/transport/request.rs
+74 -60
View File
@@ -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<String, Value>) {
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<String, Value>) -> Option<String> {
@@ -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<String, Value>) -> Map<String, Value> {
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<String, Value>, 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<String, Value>, 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<Value> {
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<String, Value>) {
for key in GROK_INTERNAL_HISTORY_KEYS {
object.remove(*key);
}
}
fn strip_grok_internal_entry_keys(converted: &mut Map<String, Value>, 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);
}
}
}
}
backend/crates/providers/xai/tests/credential/catalog.rs
+58 -1
View File
@@ -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<dyn GrokCredentialCatalogCache> = 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<dyn ProviderAccountStore> = 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<dyn GrokCredentialCatalogCache> = 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) =
backend/crates/providers/xai/tests/transport/request.rs
+136 -15
View File
@@ -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!(