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 末尾
if budget > 0 {
// 字符串 system → 转为数组
if body.get("system").and_then(|s| s.as_str()).is_some() {
let text = body["system"].as_str().unwrap().to_string();
if let Some(text) = body
.get("system")
.and_then(|s| s.as_str())
.map(str::to_string)
{
body["system"] = json!([{"type": "text", "text": text}]);
}
+21 -25
View File
@@ -87,14 +87,13 @@ pub fn classify_request(
// copilot-api 通过先 mergetext 吸收进 tool_result)再 classify 实现同等效果;
// 直接在分类层处理更稳健,不依赖 merge 启用状态和执行顺序。
let is_user_initiated = match last_msg.get("content") {
Some(content) if content.is_array() => {
let blocks = content.as_array().unwrap();
Some(Value::Array(blocks)) => {
// 含有 tool_result → 工具续写(agent),否则 → 用户发起(user)
!blocks
.iter()
.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,
};
@@ -124,7 +123,9 @@ fn is_warmup_request(body: &Value, has_anthropic_beta: bool, is_compact: bool) -
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 请求。
@@ -228,7 +229,10 @@ pub fn merge_tool_results(mut body: Value) -> Value {
}
// 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 {
return body;
}
@@ -422,10 +426,8 @@ pub fn sanitize_orphan_tool_results(mut body: Value) -> Value {
// 空 tool_use_id 或不在紧邻 assistant 的 tool_use 中 → orphan
if tool_use_id.is_empty() || !prev_tool_use_ids.contains(tool_use_id) {
let content_text = match block.get("content") {
Some(c) if c.is_string() => c.as_str().unwrap_or("").to_string(),
Some(c) if c.is_array() => c
.as_array()
.unwrap()
Some(Value::String(text)) => text.clone(),
Some(Value::Array(blocks)) => blocks
.iter()
.filter_map(|b| b.get("text").and_then(|t| t.as_str()))
.collect::<Vec<_>>()
@@ -483,10 +485,8 @@ pub fn strip_thinking_blocks(mut body: Value) -> Value {
/// 从请求体的 `system` 字段提取文本(处理 string/array 两种格式)。
fn extract_system_text(body: &Value) -> String {
match body.get("system") {
Some(s) if s.is_string() => s.as_str().unwrap_or("").to_string(),
Some(arr) if arr.is_array() => arr
.as_array()
.unwrap()
Some(Value::String(text)) => text.clone(),
Some(Value::Array(blocks)) => blocks
.iter()
.filter_map(|b| b.get("text").and_then(|t| t.as_str()))
.collect::<Vec<_>>()
@@ -573,13 +573,12 @@ fn append_text_to_tool_result(tool_result: &mut Value, text_block: &Value) {
}
// tool_result 的 content 可以是字符串或数组
match tool_result.get("content") {
Some(c) if c.is_string() => {
let existing = c.as_str().unwrap_or("");
tool_result["content"] = Value::String(format!("{existing}\n{text}"));
match tool_result.get_mut("content") {
Some(Value::String(existing)) => {
existing.push('\n');
existing.push_str(text);
}
Some(c) if c.is_array() => {
let arr = tool_result["content"].as_array_mut().unwrap();
Some(Value::Array(arr)) => {
arr.push(serde_json::json!({"type": "text", "text": text}));
}
_ => {
@@ -592,10 +591,8 @@ fn append_text_to_tool_result(tool_result: &mut Value, text_block: &Value) {
/// 从消息中提取文本内容
fn extract_text_from_message(msg: &Value) -> String {
match msg.get("content") {
Some(content) if content.is_string() => content.as_str().unwrap_or("").to_string(),
Some(content) if content.is_array() => {
let blocks = content.as_array().unwrap();
blocks
Some(Value::String(text)) => text.clone(),
Some(Value::Array(blocks)) => blocks
.iter()
.filter_map(|block| {
if block.get("type").and_then(|t| t.as_str()) == Some("text") {
@@ -605,8 +602,7 @@ fn extract_text_from_message(msg: &Value) -> String {
}
})
.collect::<Vec<_>>()
.join(" ")
}
.join(" "),
_ => String::new(),
}
}
+5 -9
View File
@@ -250,18 +250,14 @@ pub async fn send_request(
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
let result = tokio::time::timeout(
timeout,
send_raw_request(
&uri,
&method,
&headers,
original_cases.as_ref().unwrap(),
&body,
proxy_url,
),
send_raw_request(&uri, &method, &headers, original_cases, &body, proxy_url),
)
.await
.map_err(|_| ProxyError::Timeout(format!("请求超时: {}s", timeout.as_secs())))?;
+20 -10
View File
@@ -6,7 +6,7 @@
use serde_json::Value;
use std::collections::{HashMap, VecDeque};
use std::sync::RwLock;
use std::sync::{RwLock, RwLockReadGuard, RwLockWriteGuard};
/// Composite key for a Gemini shadow session.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
@@ -145,7 +145,7 @@ impl GeminiShadowStore {
let key = GeminiShadowKey::new(provider_id, session_id);
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);
let snapshot = {
@@ -193,7 +193,7 @@ impl GeminiShadowStore {
session_id: &str,
) -> Option<GeminiShadowSessionSnapshot> {
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
.sessions
.get(&key)
@@ -208,7 +208,7 @@ impl GeminiShadowStore {
#[allow(dead_code)]
pub fn clear_session(&self, provider_id: &str, session_id: &str) -> bool {
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();
if removed {
Self::remove_key_from_order(&mut inner.session_order, &key);
@@ -219,7 +219,7 @@ impl GeminiShadowStore {
/// Remove all sessions for a provider.
#[allow(dead_code)]
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
.sessions
.keys()
@@ -236,11 +236,21 @@ impl GeminiShadowStore {
/// Number of tracked sessions.
#[allow(dead_code)]
pub fn session_count(&self) -> usize {
self.inner
.read()
.expect("gemini shadow lock poisoned")
.sessions
.len()
self.read_inner().sessions.len()
}
fn read_inner(&self) -> RwLockReadGuard<'_, GeminiShadowInner> {
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(
@@ -442,7 +442,9 @@ impl ChatToResponsesState {
if should_add {
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;
if state.call_id.is_empty() {
state.call_id = format!("call_{chat_index}");
@@ -655,7 +657,9 @@ impl ChatToResponsesState {
.unwrap_or(false)
{
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;
if state.call_id.is_empty() {
state.call_id = format!("call_{key}");
@@ -687,7 +691,9 @@ impl ChatToResponsesState {
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 arguments = canonicalize_tool_arguments_str(&state.arguments);
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) => {
let valid_up_to = pos + e.valid_up_to();
buffer.push_str(
// Safety: from_utf8 guarantees [pos..valid_up_to] is valid UTF-8.
std::str::from_utf8(&input[pos..valid_up_to]).unwrap(),
);
let valid_slice = &input[pos..valid_up_to];
match std::str::from_utf8(valid_slice) {
Ok(valid) => buffer.push_str(valid),
Err(_) => buffer.push_str(&String::from_utf8_lossy(valid_slice)),
}
if let Some(invalid_len) = e.error_len() {
// Genuinely invalid byte(s) emit U+FFFD and continue.
buffer.push('\u{FFFD}');
+2 -5
View File
@@ -75,15 +75,12 @@ pub fn optimize(body: &mut Value, config: &OptimizerConfig) {
/// 追加 beta 标识到 anthropic_beta 数组(去重)
fn append_beta(body: &mut Value, beta: &str) {
match body.get("anthropic_beta") {
match body.get_mut("anthropic_beta") {
Some(Value::Array(arr)) => {
if arr.iter().any(|v| v.as_str() == Some(beta)) {
return;
}
body["anthropic_beta"]
.as_array_mut()
.unwrap()
.push(json!(beta));
arr.push(json!(beta));
}
Some(Value::Null) | None => {
body["anthropic_beta"] = json!([beta]);