refactor(proxy): replace panic-prone unwrap/expect with safe patterns

Harden the proxy module against panics by removing optimistic
unwrap()/expect() calls in favor of pattern matching and graceful
fallbacks:

- copilot_optimizer/cache_injector: bind Value::Array/String directly
  instead of is_array()+as_array().unwrap(); use is_none_or and in-place
  string mutation
- hyper_client: gate the raw-write path with if-let + filter instead of
  has_cases + unwrap()
- gemini_shadow: recover poisoned RwLock via into_inner() rather than
  panicking, with warn logging
- streaming_codex_chat: replace expect("tool state exists") with
  let-else (return/continue)
- merge_tool_results: early-return when messages is absent
- sse: fall back to from_utf8_lossy on the UTF-8 boundary slice

No behavior change on the happy path; all proxy tests pass.
This commit is contained in:
Jason
2026-05-22 00:08:37 +08:00
parent 279b9eabde
commit c12d20efd0
7 changed files with 76 additions and 67 deletions
+5 -2
View File
@@ -52,8 +52,11 @@ pub fn inject(body: &mut Value, config: &OptimizerConfig) {
// (b) system 末尾 // (b) system 末尾
if budget > 0 { if budget > 0 {
// 字符串 system → 转为数组 // 字符串 system → 转为数组
if body.get("system").and_then(|s| s.as_str()).is_some() { if let Some(text) = body
let text = body["system"].as_str().unwrap().to_string(); .get("system")
.and_then(|s| s.as_str())
.map(str::to_string)
{
body["system"] = json!([{"type": "text", "text": text}]); body["system"] = json!([{"type": "text", "text": text}]);
} }
+30 -34
View File
@@ -87,14 +87,13 @@ pub fn classify_request(
// copilot-api 通过先 mergetext 吸收进 tool_result)再 classify 实现同等效果; // copilot-api 通过先 mergetext 吸收进 tool_result)再 classify 实现同等效果;
// 直接在分类层处理更稳健,不依赖 merge 启用状态和执行顺序。 // 直接在分类层处理更稳健,不依赖 merge 启用状态和执行顺序。
let is_user_initiated = match last_msg.get("content") { let is_user_initiated = match last_msg.get("content") {
Some(content) if content.is_array() => { Some(Value::Array(blocks)) => {
let blocks = content.as_array().unwrap();
// 含有 tool_result → 工具续写(agent),否则 → 用户发起(user) // 含有 tool_result → 工具续写(agent),否则 → 用户发起(user)
!blocks !blocks
.iter() .iter()
.any(|block| block.get("type").and_then(|t| t.as_str()) == Some("tool_result")) .any(|block| block.get("type").and_then(|t| t.as_str()) == Some("tool_result"))
} }
Some(content) if content.is_string() => true, Some(Value::String(_)) => true,
_ => false, _ => false,
}; };
@@ -124,7 +123,9 @@ fn is_warmup_request(body: &Value, has_anthropic_beta: bool, is_compact: bool) -
return false; return false;
} }
// 无工具定义 // 无工具定义
!matches!(body.get("tools"), Some(tools) if tools.is_array() && !tools.as_array().unwrap().is_empty()) body.get("tools")
.and_then(|tools| tools.as_array())
.is_none_or(|tools| tools.is_empty())
} }
/// 检测是否为 Claude Code 上下文压缩/compact 请求。 /// 检测是否为 Claude Code 上下文压缩/compact 请求。
@@ -228,7 +229,10 @@ pub fn merge_tool_results(mut body: Value) -> Value {
} }
// Phase 2: 跨消息合并 — 连续的 tool_result-only 用户消息合并 // Phase 2: 跨消息合并 — 连续的 tool_result-only 用户消息合并
let messages = body["messages"].as_array().unwrap().clone(); let messages = match body.get("messages").and_then(|m| m.as_array()) {
Some(messages) => messages.clone(),
None => return body,
};
if messages.len() <= 1 { if messages.len() <= 1 {
return body; return body;
} }
@@ -422,10 +426,8 @@ pub fn sanitize_orphan_tool_results(mut body: Value) -> Value {
// 空 tool_use_id 或不在紧邻 assistant 的 tool_use 中 → orphan // 空 tool_use_id 或不在紧邻 assistant 的 tool_use 中 → orphan
if tool_use_id.is_empty() || !prev_tool_use_ids.contains(tool_use_id) { if tool_use_id.is_empty() || !prev_tool_use_ids.contains(tool_use_id) {
let content_text = match block.get("content") { let content_text = match block.get("content") {
Some(c) if c.is_string() => c.as_str().unwrap_or("").to_string(), Some(Value::String(text)) => text.clone(),
Some(c) if c.is_array() => c Some(Value::Array(blocks)) => blocks
.as_array()
.unwrap()
.iter() .iter()
.filter_map(|b| b.get("text").and_then(|t| t.as_str())) .filter_map(|b| b.get("text").and_then(|t| t.as_str()))
.collect::<Vec<_>>() .collect::<Vec<_>>()
@@ -483,10 +485,8 @@ pub fn strip_thinking_blocks(mut body: Value) -> Value {
/// 从请求体的 `system` 字段提取文本(处理 string/array 两种格式)。 /// 从请求体的 `system` 字段提取文本(处理 string/array 两种格式)。
fn extract_system_text(body: &Value) -> String { fn extract_system_text(body: &Value) -> String {
match body.get("system") { match body.get("system") {
Some(s) if s.is_string() => s.as_str().unwrap_or("").to_string(), Some(Value::String(text)) => text.clone(),
Some(arr) if arr.is_array() => arr Some(Value::Array(blocks)) => blocks
.as_array()
.unwrap()
.iter() .iter()
.filter_map(|b| b.get("text").and_then(|t| t.as_str())) .filter_map(|b| b.get("text").and_then(|t| t.as_str()))
.collect::<Vec<_>>() .collect::<Vec<_>>()
@@ -573,13 +573,12 @@ fn append_text_to_tool_result(tool_result: &mut Value, text_block: &Value) {
} }
// tool_result 的 content 可以是字符串或数组 // tool_result 的 content 可以是字符串或数组
match tool_result.get("content") { match tool_result.get_mut("content") {
Some(c) if c.is_string() => { Some(Value::String(existing)) => {
let existing = c.as_str().unwrap_or(""); existing.push('\n');
tool_result["content"] = Value::String(format!("{existing}\n{text}")); existing.push_str(text);
} }
Some(c) if c.is_array() => { Some(Value::Array(arr)) => {
let arr = tool_result["content"].as_array_mut().unwrap();
arr.push(serde_json::json!({"type": "text", "text": text})); arr.push(serde_json::json!({"type": "text", "text": text}));
} }
_ => { _ => {
@@ -592,21 +591,18 @@ fn append_text_to_tool_result(tool_result: &mut Value, text_block: &Value) {
/// 从消息中提取文本内容 /// 从消息中提取文本内容
fn extract_text_from_message(msg: &Value) -> String { fn extract_text_from_message(msg: &Value) -> String {
match msg.get("content") { match msg.get("content") {
Some(content) if content.is_string() => content.as_str().unwrap_or("").to_string(), Some(Value::String(text)) => text.clone(),
Some(content) if content.is_array() => { Some(Value::Array(blocks)) => blocks
let blocks = content.as_array().unwrap(); .iter()
blocks .filter_map(|block| {
.iter() if block.get("type").and_then(|t| t.as_str()) == Some("text") {
.filter_map(|block| { block.get("text").and_then(|t| t.as_str())
if block.get("type").and_then(|t| t.as_str()) == Some("text") { } else {
block.get("text").and_then(|t| t.as_str()) None
} else { }
None })
} .collect::<Vec<_>>()
}) .join(" "),
.collect::<Vec<_>>()
.join(" ")
}
_ => String::new(), _ => String::new(),
} }
} }
+5 -9
View File
@@ -250,18 +250,14 @@ pub async fn send_request(
proxy_url, proxy_url,
); );
if has_cases { if let Some(original_cases) = original_cases
.as_ref()
.filter(|cases| !cases.cases.is_empty())
{
// Primary path: use raw write + hyper handshake for exact header casing // Primary path: use raw write + hyper handshake for exact header casing
let result = tokio::time::timeout( let result = tokio::time::timeout(
timeout, timeout,
send_raw_request( send_raw_request(&uri, &method, &headers, original_cases, &body, proxy_url),
&uri,
&method,
&headers,
original_cases.as_ref().unwrap(),
&body,
proxy_url,
),
) )
.await .await
.map_err(|_| ProxyError::Timeout(format!("请求超时: {}s", timeout.as_secs())))?; .map_err(|_| ProxyError::Timeout(format!("请求超时: {}s", timeout.as_secs())))?;
+20 -10
View File
@@ -6,7 +6,7 @@
use serde_json::Value; use serde_json::Value;
use std::collections::{HashMap, VecDeque}; use std::collections::{HashMap, VecDeque};
use std::sync::RwLock; use std::sync::{RwLock, RwLockReadGuard, RwLockWriteGuard};
/// Composite key for a Gemini shadow session. /// Composite key for a Gemini shadow session.
#[derive(Debug, Clone, PartialEq, Eq, Hash)] #[derive(Debug, Clone, PartialEq, Eq, Hash)]
@@ -145,7 +145,7 @@ impl GeminiShadowStore {
let key = GeminiShadowKey::new(provider_id, session_id); let key = GeminiShadowKey::new(provider_id, session_id);
let turn = GeminiAssistantTurn::new(assistant_content, tool_calls); let turn = GeminiAssistantTurn::new(assistant_content, tool_calls);
let mut inner = self.inner.write().expect("gemini shadow lock poisoned"); let mut inner = self.write_inner();
Self::touch_session_order(&mut inner.session_order, &key); Self::touch_session_order(&mut inner.session_order, &key);
let snapshot = { let snapshot = {
@@ -193,7 +193,7 @@ impl GeminiShadowStore {
session_id: &str, session_id: &str,
) -> Option<GeminiShadowSessionSnapshot> { ) -> Option<GeminiShadowSessionSnapshot> {
let key = GeminiShadowKey::new(provider_id, session_id); let key = GeminiShadowKey::new(provider_id, session_id);
let mut inner = self.inner.write().expect("gemini shadow lock poisoned"); let mut inner = self.write_inner();
let snapshot = inner let snapshot = inner
.sessions .sessions
.get(&key) .get(&key)
@@ -208,7 +208,7 @@ impl GeminiShadowStore {
#[allow(dead_code)] #[allow(dead_code)]
pub fn clear_session(&self, provider_id: &str, session_id: &str) -> bool { pub fn clear_session(&self, provider_id: &str, session_id: &str) -> bool {
let key = GeminiShadowKey::new(provider_id, session_id); let key = GeminiShadowKey::new(provider_id, session_id);
let mut inner = self.inner.write().expect("gemini shadow lock poisoned"); let mut inner = self.write_inner();
let removed = inner.sessions.remove(&key).is_some(); let removed = inner.sessions.remove(&key).is_some();
if removed { if removed {
Self::remove_key_from_order(&mut inner.session_order, &key); Self::remove_key_from_order(&mut inner.session_order, &key);
@@ -219,7 +219,7 @@ impl GeminiShadowStore {
/// Remove all sessions for a provider. /// Remove all sessions for a provider.
#[allow(dead_code)] #[allow(dead_code)]
pub fn clear_provider(&self, provider_id: &str) -> usize { pub fn clear_provider(&self, provider_id: &str) -> usize {
let mut inner = self.inner.write().expect("gemini shadow lock poisoned"); let mut inner = self.write_inner();
let keys: Vec<_> = inner let keys: Vec<_> = inner
.sessions .sessions
.keys() .keys()
@@ -236,11 +236,21 @@ impl GeminiShadowStore {
/// Number of tracked sessions. /// Number of tracked sessions.
#[allow(dead_code)] #[allow(dead_code)]
pub fn session_count(&self) -> usize { pub fn session_count(&self) -> usize {
self.inner self.read_inner().sessions.len()
.read() }
.expect("gemini shadow lock poisoned")
.sessions fn read_inner(&self) -> RwLockReadGuard<'_, GeminiShadowInner> {
.len() self.inner.read().unwrap_or_else(|poisoned| {
log::warn!("[GeminiShadow] recovering poisoned read lock");
poisoned.into_inner()
})
}
fn write_inner(&self) -> RwLockWriteGuard<'_, GeminiShadowInner> {
self.inner.write().unwrap_or_else(|poisoned| {
log::warn!("[GeminiShadow] recovering poisoned write lock");
poisoned.into_inner()
})
} }
fn snapshot_session( fn snapshot_session(
@@ -442,7 +442,9 @@ impl ChatToResponsesState {
if should_add { if should_add {
let assigned = self.next_output_index(); let assigned = self.next_output_index();
let state = self.tools.get_mut(&chat_index).expect("tool state exists"); let Some(state) = self.tools.get_mut(&chat_index) else {
return events;
};
state.added = true; state.added = true;
if state.call_id.is_empty() { if state.call_id.is_empty() {
state.call_id = format!("call_{chat_index}"); state.call_id = format!("call_{chat_index}");
@@ -655,7 +657,9 @@ impl ChatToResponsesState {
.unwrap_or(false) .unwrap_or(false)
{ {
let assigned = self.next_output_index(); let assigned = self.next_output_index();
let state = self.tools.get_mut(&key).expect("tool state exists"); let Some(state) = self.tools.get_mut(&key) else {
continue;
};
state.added = true; state.added = true;
if state.call_id.is_empty() { if state.call_id.is_empty() {
state.call_id = format!("call_{key}"); state.call_id = format!("call_{key}");
@@ -687,7 +691,9 @@ impl ChatToResponsesState {
events.push(event); events.push(event);
} }
let state = self.tools.get_mut(&key).expect("tool state exists"); let Some(state) = self.tools.get_mut(&key) else {
continue;
};
let output_index = state.output_index.unwrap_or(0); let output_index = state.output_index.unwrap_or(0);
let arguments = canonicalize_tool_arguments_str(&state.arguments); let arguments = canonicalize_tool_arguments_str(&state.arguments);
let item = response_function_call_item( let item = response_function_call_item(
+5 -4
View File
@@ -66,10 +66,11 @@ pub(crate) fn append_utf8_safe(buffer: &mut String, remainder: &mut Vec<u8>, new
} }
Err(e) => { Err(e) => {
let valid_up_to = pos + e.valid_up_to(); let valid_up_to = pos + e.valid_up_to();
buffer.push_str( let valid_slice = &input[pos..valid_up_to];
// Safety: from_utf8 guarantees [pos..valid_up_to] is valid UTF-8. match std::str::from_utf8(valid_slice) {
std::str::from_utf8(&input[pos..valid_up_to]).unwrap(), Ok(valid) => buffer.push_str(valid),
); Err(_) => buffer.push_str(&String::from_utf8_lossy(valid_slice)),
}
if let Some(invalid_len) = e.error_len() { if let Some(invalid_len) = e.error_len() {
// Genuinely invalid byte(s) emit U+FFFD and continue. // Genuinely invalid byte(s) emit U+FFFD and continue.
buffer.push('\u{FFFD}'); buffer.push('\u{FFFD}');
+2 -5
View File
@@ -75,15 +75,12 @@ pub fn optimize(body: &mut Value, config: &OptimizerConfig) {
/// 追加 beta 标识到 anthropic_beta 数组(去重) /// 追加 beta 标识到 anthropic_beta 数组(去重)
fn append_beta(body: &mut Value, beta: &str) { fn append_beta(body: &mut Value, beta: &str) {
match body.get("anthropic_beta") { match body.get_mut("anthropic_beta") {
Some(Value::Array(arr)) => { Some(Value::Array(arr)) => {
if arr.iter().any(|v| v.as_str() == Some(beta)) { if arr.iter().any(|v| v.as_str() == Some(beta)) {
return; return;
} }
body["anthropic_beta"] arr.push(json!(beta));
.as_array_mut()
.unwrap()
.push(json!(beta));
} }
Some(Value::Null) | None => { Some(Value::Null) | None => {
body["anthropic_beta"] = json!([beta]); body["anthropic_beta"] = json!([beta]);