b3nw/zerogravity
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor: decompose large functions and remove dead code

Browse Source 
- Decompose modify_request() into 7 single-responsibility helpers
- Decompose handle_http_over_tls(): extract read_full_request, dispatch_stream_events
- Promote connect_upstream/resolve_upstream to module-level functions
- Split standalone.rs (1238 lines) into 4 submodules:
  standalone/mod.rs, spawn.rs, discovery.rs, stub.rs
- Extract proto wire primitives into proto/wire.rs
- Remove 6 dead MitmStore methods
- Remove dead SessionResult, DEFAULT_SESSION, get_or_create
- Remove dead decode_varint_at, extract_conversation_id
- Clean all unused imports across 10 files
- Suppress structural dead_code warnings on deserialization fields

Warnings: 20 -> 0. All 43 tests pass.
This commit is contained in:
Nikketryhard
2026-02-17 22:27:26 -06:00
parent 637fbc0e54
commit 48674f65da
21 changed files with 3099 additions and 3346 deletions
Download Patch File Download Diff File Expand all files Collapse all files

296
src/api/responses.rs
View File

@@ -142,14 +142,9 @@ fn extract_responses_input(
(final_text, tool_results, image)
}
/// Extract conversation/session ID from Responses API `conversation` field.
fn extract_conversation_id(conv: &Option<serde_json::Value>) -> Option<String> {
match conv {
Some(serde_json::Value::String(s)) => Some(s.clone()),
Some(obj) => obj["id"].as_str().map(|s| s.to_string()),
None => None,
}
}
/// Response-specific data for building a Response object.
struct ResponseData {
@@ -241,47 +236,26 @@ pub(crate) async fn handle_responses(
// Handle tool result submission (function_call_output in input)
let is_tool_result_turn = !tool_results.is_empty();
if is_tool_result_turn {
let mut pending: Vec<PendingToolResult> = Vec::new();
for tr in &tool_results {
// Look up function name from call_id
let name = state
.mitm_store
.lookup_call_id(&tr.call_id)
.await
.unwrap_or_else(|| "unknown_function".to_string());
let mut pending_tool_results: Vec<PendingToolResult> = Vec::new();
if is_tool_result_turn {
for tr in &tool_results {
// For tool result turns, we use the call_id as the name directly.
// The proxy captured function calls (with real names) are paired in
// the ToolRound when we know the cascade_id later.
let name = tr.call_id.clone();
// Parse the output as JSON, fall back to string wrapper
let result_value = serde_json::from_str::<serde_json::Value>(&tr.output)
.unwrap_or_else(|_| serde_json::json!({"result": tr.output}));
// Also store as pending (legacy compat)
state
.mitm_store
.add_tool_result(PendingToolResult {
name: name.clone(),
result: result_value.clone(),
})
.await;
pending.push(PendingToolResult {
pending_tool_results.push(PendingToolResult {
name,
result: result_value,
});
}
// Build a ToolRound from the MITM-captured function calls + client results.
// get_last_function_calls() has the calls from Google's previous response.
// We take existing accumulated rounds and append this new round.
let last_calls = state.mitm_store.get_last_function_calls().await;
let mut rounds = state.mitm_store.take_tool_rounds().await;
rounds.push(crate::mitm::store::ToolRound {
calls: last_calls,
results: pending,
});
state.mitm_store.set_tool_rounds(rounds).await;
info!(
count = tool_results.len(),
"Stored tool results for MITM injection (built tool round)"
"Tool results for MITM injection (will build tool round after cascade_id)"
);
}
@@ -293,7 +267,8 @@ pub(crate) async fn handle_responses(
);
}
// Store client tools in MitmStore for MITM injection
// ── Build per-request state locally ──────────────────────────────────
// Detect web_search_preview tool (OpenAI spec) → enable Google Search grounding
let has_web_search = body.tools.as_ref().map_or(false, |tools| {
tools.iter().any(|t| {
@@ -301,27 +276,20 @@ pub(crate) async fn handle_responses(
t_type == "web_search_preview" || t_type == "web_search"
})
});
if let Some(ref tools) = body.tools {
let gemini_tools = openai_tools_to_gemini(tools);
if !gemini_tools.is_empty() {
state.mitm_store.set_tools(gemini_tools).await;
info!(
count = tools.len(),
"Stored client tools for MITM injection"
);
} else {
state.mitm_store.clear_tools().await;
}
} else {
state.mitm_store.clear_tools().await;
}
if let Some(ref choice) = body.tool_choice {
let gemini_config = openai_tool_choice_to_gemini(choice);
state.mitm_store.set_tool_config(gemini_config).await;
}
// Store generation parameters for MITM injection
// Extract text.format for structured output (json_schema)
// Convert OpenAI tools to Gemini format
let tools = body.tools.as_ref().and_then(|t| {
let gemini_tools = openai_tools_to_gemini(t);
if gemini_tools.is_empty() { None } else {
info!(count = t.len(), "Client tools for MITM injection");
Some(gemini_tools)
}
});
let tool_config = body.tool_choice.as_ref().map(|choice| {
openai_tool_choice_to_gemini(choice)
});
// Build generation params locally
let (response_mime_type, response_schema, text_format) = if let Some(ref text_val) = body.text {
let fmt_type = text_val["format"]["type"].as_str().unwrap_or("text");
if fmt_type == "json_schema" {
@@ -345,100 +313,98 @@ pub(crate) async fn handle_responses(
} else {
(None, None, TextFormat::default())
};
{
use crate::mitm::store::GenerationParams;
let gp = GenerationParams {
temperature: body.temperature,
top_p: body.top_p,
top_k: None,
max_output_tokens: body.max_output_tokens,
stop_sequences: None,
frequency_penalty: None,
presence_penalty: None,
reasoning_effort: body.reasoning_effort.clone(),
response_mime_type,
response_schema,
google_search: has_web_search,
};
if gp.temperature.is_some()
|| gp.top_p.is_some()
|| gp.max_output_tokens.is_some()
|| gp.reasoning_effort.is_some()
|| gp.response_mime_type.is_some()
|| gp.response_schema.is_some()
|| gp.google_search
{
state.mitm_store.set_generation_params(gp).await;
} else {
state.mitm_store.clear_generation_params().await;
}
}
let response_id = format!("resp_{}", uuid::Uuid::new_v4().to_string().replace('-', ""));
// Session/conversation management
let session_id_str = extract_conversation_id(&body.conversation);
let cascade_id = if let Some(ref sid) = session_id_str {
match state
.sessions
.get_or_create(Some(sid), || state.backend.create_cascade())
.await
{
Ok(sr) => sr.cascade_id,
Err(e) => {
return err_response(
StatusCode::BAD_GATEWAY,
format!("StartCascade failed: {e}"),
"server_error",
);
}
}
} else {
match state.backend.create_cascade().await {
Ok(cid) => cid,
Err(e) => {
return err_response(
StatusCode::BAD_GATEWAY,
format!("StartCascade failed: {e}"),
"server_error",
);
}
}
use crate::mitm::store::GenerationParams;
let gp = GenerationParams {
temperature: body.temperature,
top_p: body.top_p,
top_k: None,
max_output_tokens: body.max_output_tokens,
stop_sequences: None,
frequency_penalty: None,
presence_penalty: None,
reasoning_effort: body.reasoning_effort.clone(),
response_mime_type,
response_schema,
google_search: has_web_search,
};
// Send message
state.mitm_store.set_active_cascade(&cascade_id).await;
// Store real user text for MITM injection — LS gets a dummy prompt
state.mitm_store.set_pending_user_text(user_text.clone()).await;
// Store image for MITM injection (LS doesn't forward images to Google API)
if let Some(ref img) = image {
use base64::Engine;
state
.mitm_store
.set_pending_image(crate::mitm::store::PendingImage {
base64_data: base64::engine::general_purpose::STANDARD.encode(&img.data),
mime_type: img.mime_type.clone(),
})
.await;
}
// Pre-flight: install channel BEFORE send_message so the MITM proxy
// can grab it when the LS fires its API call.
let has_custom_tools = state.mitm_store.get_tools().await.is_some();
let mitm_rx = if has_custom_tools {
state.mitm_store.clear_response_async().await;
state.mitm_store.clear_upstream_error().await;
let _ = state.mitm_store.take_any_function_calls().await;
let (tx, rx) = tokio::sync::mpsc::channel(64);
state.mitm_store.set_channel(tx).await;
Some(rx)
let generation_params = if gp.temperature.is_some()
|| gp.top_p.is_some()
|| gp.max_output_tokens.is_some()
|| gp.reasoning_effort.is_some()
|| gp.response_mime_type.is_some()
|| gp.response_schema.is_some()
|| gp.google_search
{
Some(gp)
} else {
None
};
let response_id = format!("resp_{}", uuid::Uuid::new_v4().to_string().replace('-', ""));
// Always create a new cascade for every request
let cascade_id = match state.backend.create_cascade().await {
Ok(cid) => cid,
Err(e) => {
return err_response(
StatusCode::BAD_GATEWAY,
format!("StartCascade failed: {e}"),
"server_error",
);
}
};
// Image for MITM injection
let pending_image = image.as_ref().map(|img| {
use base64::Engine;
crate::mitm::store::PendingImage {
base64_data: base64::engine::general_purpose::STANDARD.encode(&img.data),
mime_type: img.mime_type.clone(),
}
});
// Build event channel
let has_custom_tools = tools.is_some();
let (mitm_rx, event_tx) = if has_custom_tools {
let (tx, rx) = tokio::sync::mpsc::channel(64);
(Some(rx), Some(tx))
} else {
(None, None)
};
// Build tool rounds now that cascade_id is known
let mut tool_rounds: Vec<crate::mitm::store::ToolRound> = Vec::new();
if is_tool_result_turn && !pending_tool_results.is_empty() {
// Get last captured function calls from the previous request context
let last_calls = state.mitm_store.take_function_calls(&cascade_id).await
.unwrap_or_default();
tool_rounds.push(crate::mitm::store::ToolRound {
calls: last_calls,
results: pending_tool_results.clone(),
});
}
// Register all per-request state atomically
state.mitm_store.register_request(crate::mitm::store::RequestContext {
cascade_id: cascade_id.clone(),
pending_user_text: user_text.clone(),
event_channel: event_tx,
generation_params,
pending_image,
tools,
tool_config,
pending_tool_results,
tool_rounds,
last_function_calls: Vec::new(),
call_id_to_name: std::collections::HashMap::new(),
created_at: std::time::Instant::now(),
}).await;
// Send REAL user text to LS
match state
.backend
.send_message_with_image(&cascade_id, ".", model.model_enum, image.as_ref())
.send_message_with_image(&cascade_id, &format!(".<cid:{}>", cascade_id), model.model_enum, image.as_ref())
.await
{
Ok((200, _)) => {
@@ -449,7 +415,7 @@ pub(crate) async fn handle_responses(
});
}
Ok((status, _)) => {
state.mitm_store.drop_channel().await;
state.mitm_store.remove_request(&cascade_id).await;
return err_response(
StatusCode::BAD_GATEWAY,
format!("Antigravity returned {status}"),
@@ -457,7 +423,7 @@ pub(crate) async fn handle_responses(
);
}
Err(e) => {
state.mitm_store.drop_channel().await;
state.mitm_store.remove_request(&cascade_id).await;
return err_response(
StatusCode::BAD_GATEWAY,
format!("Send message failed: {e}"),
@@ -644,7 +610,7 @@ async fn handle_responses_sync(
let mut acc_text = String::new();
let mut acc_thinking: Option<String> = None;
let mut last_usage: Option<crate::mitm::store::ApiUsage> = None;
let mut _last_usage: Option<crate::mitm::store::ApiUsage> = None;
while let Some(event) = tokio::time::timeout(
std::time::Duration::from_secs(timeout.saturating_sub(start.elapsed().as_secs())),
@@ -654,7 +620,7 @@ async fn handle_responses_sync(
match event {
MitmEvent::ThinkingDelta(t) => { acc_thinking = Some(t); }
MitmEvent::TextDelta(t) => { acc_text = t; }
MitmEvent::Usage(u) => { last_usage = Some(u); }
MitmEvent::Usage(u) => { _last_usage = Some(u); }
MitmEvent::Grounding(_) => {} // stored by proxy directly
MitmEvent::FunctionCall(raw_calls) => {
let calls: Vec<_> = if let Some(max) = params.max_tool_calls {
@@ -668,14 +634,14 @@ async fn handle_responses_sync(
"call_{}",
uuid::Uuid::new_v4().to_string().replace('-', "")[..24].to_string()
);
state.mitm_store.register_call_id(call_id.clone(), fc.name.clone()).await;
state.mitm_store.register_call_id(&cascade_id, call_id.clone(), fc.name.clone()).await;
let arguments = serde_json::to_string(&fc.args).unwrap_or_default();
output_items.push(build_function_call_output(&call_id, &fc.name, &arguments));
}
let (usage, _) = usage_from_poll(
&state.mitm_store, &cascade_id, &None, &params.user_text, "",
).await;
state.mitm_store.drop_channel().await;
state.mitm_store.remove_request(&cascade_id).await;
let resp = build_response_object(
ResponseData {
id: response_id,
@@ -700,8 +666,8 @@ async fn handle_responses_sync(
// Thinking-only — LS needs to make a follow-up request.
// Reinstall channel and unblock gate.
let (new_tx, new_rx) = tokio::sync::mpsc::channel(64);
state.mitm_store.set_channel(new_tx).await;
state.mitm_store.clear_request_in_flight();
state.mitm_store.set_channel(&cascade_id, new_tx).await;
let _ = state.mitm_store.take_any_function_calls().await;
rx = new_rx;
debug!(
@@ -713,7 +679,7 @@ async fn handle_responses_sync(
let (usage, _) = usage_from_poll(
&state.mitm_store, &cascade_id, &None, &params.user_text, &acc_text,
).await;
state.mitm_store.drop_channel().await;
state.mitm_store.remove_request(&cascade_id).await;
let mut output_items: Vec<serde_json::Value> = Vec::new();
if let Some(ref t) = acc_thinking {
@@ -738,14 +704,14 @@ async fn handle_responses_sync(
return Json(resp).into_response();
}
MitmEvent::UpstreamError(err) => {
state.mitm_store.drop_channel().await;
state.mitm_store.remove_request(&cascade_id).await;
return upstream_err_response(&err);
}
}
}
// Timeout
state.mitm_store.drop_channel().await;
state.mitm_store.remove_request(&cascade_id).await;
return err_response(
StatusCode::GATEWAY_TIMEOUT,
format!("Timeout: no response from Google API after {timeout}s"),
@@ -789,7 +755,7 @@ async fn handle_responses_sync(
// Register call_id → name mapping for tool result routing
state
.mitm_store
.register_call_id(call_id.clone(), fc.name.clone())
.register_call_id(&cascade_id, call_id.clone(), fc.name.clone())
.await;
// Stringify args (OpenAI sends arguments as JSON string)
@@ -1092,7 +1058,7 @@ async fn handle_responses_stream(
uuid::Uuid::new_v4().to_string().replace('-', "")[..24].to_string()
);
let arguments = serde_json::to_string(&fc.args).unwrap_or_default();
state.mitm_store.register_call_id(call_id.clone(), fc.name.clone()).await;
state.mitm_store.register_call_id(&cascade_id, call_id.clone(), fc.name.clone()).await;
let fc_item_id = format!("fc_{}", uuid::Uuid::new_v4().to_string().replace('-', ""));
yield Ok(responses_sse_event(
@@ -1166,7 +1132,7 @@ async fn handle_responses_stream(
"response": response_to_json(&final_resp),
}),
));
state.mitm_store.drop_channel().await;
state.mitm_store.remove_request(&cascade_id).await;
return;
}
MitmEvent::ResponseComplete => {
@@ -1184,14 +1150,14 @@ async fn handle_responses_stream(
) {
yield Ok(evt);
}
state.mitm_store.drop_channel().await;
state.mitm_store.remove_request(&cascade_id).await;
return;
} else if !last_thinking.is_empty() {
// Thinking-only response — LS needs follow-up API calls.
// Create a new channel and unblock the gate.
let (new_tx, new_rx) = tokio::sync::mpsc::channel(64);
state.mitm_store.set_channel(new_tx).await;
state.mitm_store.clear_request_in_flight();
state.mitm_store.set_channel(&cascade_id, new_tx).await;
let _ = state.mitm_store.take_any_function_calls().await;
rx = new_rx;
debug!(
@@ -1220,7 +1186,7 @@ async fn handle_responses_stream(
},
}),
));
state.mitm_store.drop_channel().await;
state.mitm_store.remove_request(&cascade_id).await;
return;
}
MitmEvent::Usage(_) | MitmEvent::Grounding(_) => {
@@ -1230,7 +1196,7 @@ async fn handle_responses_stream(
}
// Timeout in channel mode
state.mitm_store.drop_channel().await;
state.mitm_store.remove_request(&cascade_id).await;
yield Ok(responses_sse_event(
"response.failed",
serde_json::json!({