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

310
src/api/completions.rs
View File

@@ -18,15 +18,9 @@ use super::util::{err_response, now_unix, upstream_err_response};
use super::AppState;
use crate::mitm::store::{CapturedFunctionCall, PendingToolResult, ToolRound};
/// Extract a conversation/session ID from a flexible JSON value.
/// Accepts a plain string or an object with an "id" 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,
}
}
/// System fingerprint for completions responses (derived from crate version at compile time).
fn system_fingerprint() -> String {
@@ -187,10 +181,7 @@ pub(crate) async fn handle_completions(
model_name, body.stream
);
// Diagnostic: dump OpenCode's raw request
if let Ok(pretty) = serde_json::to_string_pretty(&body) {
let _ = std::fs::write("/tmp/opencode-request.json", &pretty);
}
let model = match lookup_model(model_name) {
Some(m) => m,
@@ -204,35 +195,28 @@ pub(crate) async fn handle_completions(
}
};
// Store client tools from this request (or clear stale ones from other endpoints)
if let Some(ref tools) = body.tools {
let gemini_tools = crate::mitm::modify::openai_tools_to_gemini(tools);
if !gemini_tools.is_empty() {
state.mitm_store.set_tools(gemini_tools).await;
if let Some(ref choice) = body.tool_choice {
let gemini_config = crate::mitm::modify::openai_tool_choice_to_gemini(choice);
state.mitm_store.set_tool_config(gemini_config).await;
}
info!(
count = tools.len(),
"Completions: stored client tools for MITM injection"
);
} else {
state.mitm_store.clear_tools().await;
// ── Build per-request state locally ──────────────────────────────────
// Convert OpenAI tools to Gemini format
let tools = body.tools.as_ref().and_then(|t| {
let gemini_tools = crate::mitm::modify::openai_tools_to_gemini(t);
if gemini_tools.is_empty() { None } else {
info!(count = t.len(), "Completions: client tools for MITM injection");
Some(gemini_tools)
}
} else {
state.mitm_store.clear_tools().await;
}
});
let tool_config = body.tools.as_ref().and_then(|_| {
body.tool_choice.as_ref().map(|choice| {
crate::mitm::modify::openai_tool_choice_to_gemini(choice)
})
});
// ── Extract tool results from messages for MITM injection ──────────
// When OpenCode sends back tool results, the messages array contains:
// 1. assistant message with tool_calls (the model's previous function calls)
// 2. tool messages with results (the executed tool outputs)
// We build ToolRounds: each round pairs one assistant's tool_calls with
// the subsequent tool result messages. This enables correct per-turn
// history rewriting for multi-step tool use.
// Build ToolRounds from message history: each round pairs assistant tool_calls
// with subsequent tool result messages. Local call_id_to_name mapping.
let mut tool_rounds: Vec<ToolRound> = Vec::new();
let mut call_id_to_name: std::collections::HashMap<String, String> = std::collections::HashMap::new();
{
let mut rounds: Vec<ToolRound> = Vec::new();
let mut current_round: Option<ToolRound> = None;
for msg in &body.messages {
@@ -241,7 +225,7 @@ pub(crate) async fn handle_completions(
// Finalize any open round
if let Some(round) = current_round.take() {
if !round.calls.is_empty() {
rounds.push(round);
tool_rounds.push(round);
}
}
// Start new round if this assistant has tool_calls
@@ -255,14 +239,15 @@ pub(crate) async fn handle_completions(
.unwrap_or(serde_json::json!({}));
let call_id = tc["id"].as_str().unwrap_or("").to_string();
// Register call_id → name for lookup
// Register call_id → name locally
if !call_id.is_empty() {
state.mitm_store.register_call_id(call_id, name.clone()).await;
call_id_to_name.insert(call_id, name.clone());
}
calls.push(CapturedFunctionCall {
name,
args,
thought_signature: None,
captured_at: std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
@@ -281,16 +266,13 @@ pub(crate) async fn handle_completions(
"tool" => {
let text = extract_message_text(&msg.content);
if let Some(ref call_id) = msg.tool_call_id {
// Look up function name from call_id, fall back to
// positional index within the current round's calls
let result_index = current_round
.as_ref()
.map(|r| r.results.len())
.unwrap_or(0);
let name = state
.mitm_store
.lookup_call_id(call_id)
.await
let name = call_id_to_name
.get(call_id.as_str())
.cloned()
.unwrap_or_else(|| {
current_round
.as_ref()
@@ -314,7 +296,7 @@ pub(crate) async fn handle_completions(
// Any other role (user, system) finalizes the current round
if let Some(round) = current_round.take() {
if !round.calls.is_empty() {
rounds.push(round);
tool_rounds.push(round);
}
}
}
@@ -323,69 +305,86 @@ pub(crate) async fn handle_completions(
// Finalize last round
if let Some(round) = current_round.take() {
if !round.calls.is_empty() {
rounds.push(round);
tool_rounds.push(round);
}
}
if !rounds.is_empty() {
if !tool_rounds.is_empty() {
info!(
round_count = rounds.len(),
calls = ?rounds.iter().map(|r| r.calls.iter().map(|c| &c.name).collect::<Vec<_>>()).collect::<Vec<_>>(),
"Completions: stored {} tool round(s) for MITM history rewrite",
rounds.len(),
round_count = tool_rounds.len(),
calls = ?tool_rounds.iter().map(|r| r.calls.iter().map(|c| &c.name).collect::<Vec<_>>()).collect::<Vec<_>>(),
"Completions: {} tool round(s) for MITM history rewrite",
tool_rounds.len(),
);
// Also set last_function_calls from the latest round for proxy.rs recording compat
if let Some(last_round) = rounds.last() {
state.mitm_store.set_last_function_calls(last_round.calls.clone()).await;
// Merge thought_signatures from MITM-captured function calls.
// OpenAI format doesn't carry thought signatures, but Google requires
// them when injecting functionCall parts back into history.
let sigs = state.mitm_store.peek_thought_signatures().await;
if !sigs.is_empty() {
let mut merged = 0usize;
for round in &mut tool_rounds {
for fc in &mut round.calls {
if fc.thought_signature.is_none() {
if let Some(sig) = sigs.get(&fc.name) {
fc.thought_signature = Some(sig.clone());
merged += 1;
}
}
}
}
if merged > 0 {
info!(
merged_count = merged,
"Completions: merged {} thought_signature(s) from MITM capture",
merged,
);
}
}
state.mitm_store.set_tool_rounds(rounds).await;
}
}
// Store generation parameters for MITM injection
// Build generation parameters locally
use crate::mitm::store::GenerationParams;
let (response_mime_type, response_schema) = match body.response_format.as_ref() {
Some(rf) => match rf.format_type.as_str() {
"json_object" | "json" => (Some("application/json".to_string()), None),
"json_schema" => {
let schema = rf.json_schema.as_ref().and_then(|js| js.schema.clone());
(Some("application/json".to_string()), schema)
}
_ => (None, None),
},
None => (None, None),
};
let gp = GenerationParams {
temperature: body.temperature,
top_p: body.top_p,
top_k: None,
max_output_tokens: body.max_tokens.or(body.max_completion_tokens),
stop_sequences: body.stop.clone().map(|s| s.into_vec()),
frequency_penalty: body.frequency_penalty,
presence_penalty: body.presence_penalty,
reasoning_effort: body.reasoning_effort.clone(),
response_mime_type,
response_schema,
google_search: body.web_search,
};
let generation_params = if gp.temperature.is_some()
|| gp.top_p.is_some()
|| gp.max_output_tokens.is_some()
|| gp.frequency_penalty.is_some()
|| gp.presence_penalty.is_some()
|| gp.reasoning_effort.is_some()
|| gp.stop_sequences.is_some()
|| gp.response_mime_type.is_some()
|| gp.response_schema.is_some()
|| gp.google_search
{
use crate::mitm::store::GenerationParams;
let (response_mime_type, response_schema) = match body.response_format.as_ref() {
Some(rf) => match rf.format_type.as_str() {
"json_object" | "json" => (Some("application/json".to_string()), None),
"json_schema" => {
let schema = rf.json_schema.as_ref().and_then(|js| js.schema.clone());
(Some("application/json".to_string()), schema)
}
_ => (None, None),
},
None => (None, None),
};
let gp = GenerationParams {
temperature: body.temperature,
top_p: body.top_p,
top_k: None, // OpenAI doesn't have top_k
max_output_tokens: body.max_tokens.or(body.max_completion_tokens),
stop_sequences: body.stop.clone().map(|s| s.into_vec()),
frequency_penalty: body.frequency_penalty,
presence_penalty: body.presence_penalty,
reasoning_effort: body.reasoning_effort.clone(),
response_mime_type,
response_schema,
google_search: body.web_search,
};
// Only store if at least one param is set
if gp.temperature.is_some()
|| gp.top_p.is_some()
|| gp.max_output_tokens.is_some()
|| gp.frequency_penalty.is_some()
|| gp.presence_penalty.is_some()
|| gp.reasoning_effort.is_some()
|| gp.stop_sequences.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;
}
}
Some(gp)
} else {
None
};
let token = state.backend.oauth_token().await;
if token.is_empty() {
@@ -410,23 +409,8 @@ pub(crate) async fn handle_completions(
warn!("n={n} requested with streaming — streaming only supports n=1, ignoring n");
}
// Session/conversation: reuse cascade if conversation ID provided
let session_id_str = extract_conversation_id(&body.conversation);
// Helper to create a cascade (reuses session or creates fresh)
let create_cascade = |state: Arc<AppState>, session_id: Option<String>| async move {
if let Some(ref sid) = session_id {
state
.sessions
.get_or_create(Some(sid), || state.backend.create_cascade())
.await
.map(|sr| sr.cascade_id)
} else {
state.backend.create_cascade().await
}
};
let cascade_id = match create_cascade(Arc::clone(&state), session_id_str.clone()).await {
// 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(
@@ -437,40 +421,54 @@ pub(crate) async fn handle_completions(
}
};
// Send message on primary cascade
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 {
// Image for MITM injection
let pending_image = image.as_ref().map(|img| {
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;
}
crate::mitm::store::PendingImage {
base64_data: base64::engine::general_purpose::STANDARD.encode(&img.data),
mime_type: img.mime_type.clone(),
}
});
// Pre-flight: install channel BEFORE send_message so the MITM proxy
// can grab it when the LS fires its API call.
// Only for streaming — sync paths use poll_for_response (legacy store).
let has_custom_tools = state.mitm_store.get_tools().await.is_some();
let mitm_rx = if has_custom_tools && body.stream {
state.mitm_store.clear_response_async().await;
state.mitm_store.clear_upstream_error().await;
let _ = state.mitm_store.take_any_function_calls().await;
// Get last calls from the latest tool round (if any) for proxy recording compat
let last_function_calls = tool_rounds.last()
.map(|r| r.calls.clone())
.unwrap_or_default();
// Build event channel for streaming
let has_custom_tools = tools.is_some();
let (mitm_rx, event_tx) = if has_custom_tools && body.stream {
let (tx, rx) = tokio::sync::mpsc::channel(64);
state.mitm_store.set_channel(tx).await;
Some(rx)
(Some(rx), Some(tx))
} else {
None
(None, None)
};
// Build pending tool results from latest round
let pending_tool_results = tool_rounds.last()
.map(|r| r.results.clone())
.unwrap_or_default();
// 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,
call_id_to_name,
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, _)) => {
@@ -481,7 +479,7 @@ pub(crate) async fn handle_completions(
});
}
Ok((status, _)) => {
state.mitm_store.drop_channel().await;
state.mitm_store.remove_request(&cascade_id).await;
return err_response(
StatusCode::BAD_GATEWAY,
format!("Backend returned {status}"),
@@ -489,7 +487,7 @@ pub(crate) async fn handle_completions(
);
}
Err(e) => {
state.mitm_store.drop_channel().await;
state.mitm_store.remove_request(&cascade_id).await;
return err_response(
StatusCode::BAD_GATEWAY,
format!("Send failed: {e}"),
@@ -537,7 +535,7 @@ pub(crate) async fn handle_completions(
// Send the same message on each extra cascade
match state
.backend
.send_message_with_image(&cid, ".", model.model_enum, image.as_ref())
.send_message_with_image(&cid, &format!(".<cid:{}>", cid), model.model_enum, image.as_ref())
.await
{
Ok((200, _)) => {
@@ -775,7 +773,7 @@ async fn chat_completions_stream(
)));
}
yield Ok(Event::default().data("[DONE]"));
state.mitm_store.drop_channel().await;
state.mitm_store.remove_request(&cascade_id).await;
return;
}
MitmEvent::ResponseComplete => {
@@ -803,15 +801,15 @@ async fn chat_completions_stream(
)));
}
yield Ok(Event::default().data("[DONE]"));
state.mitm_store.drop_channel().await;
state.mitm_store.remove_request(&cascade_id).await;
return;
} else if !acc_thinking.is_empty() && !did_unblock_ls {
// Thinking-only response — LS needs follow-up API calls.
// Create a new channel and unblock the gate.
did_unblock_ls = true;
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!(
@@ -845,7 +843,7 @@ async fn chat_completions_stream(
)));
}
yield Ok(Event::default().data("[DONE]"));
state.mitm_store.drop_channel().await;
state.mitm_store.remove_request(&cascade_id).await;
return;
}
// Don't break — wait for more channel events
@@ -861,7 +859,7 @@ async fn chat_completions_stream(
None,
)));
yield Ok(Event::default().data("[DONE]"));
state.mitm_store.drop_channel().await;
state.mitm_store.remove_request(&cascade_id).await;
return;
}
continue 'channel_loop;
@@ -878,7 +876,7 @@ async fn chat_completions_stream(
}
})).unwrap()));
yield Ok(Event::default().data("[DONE]"));
state.mitm_store.drop_channel().await;
state.mitm_store.remove_request(&cascade_id).await;
return;
}
MitmEvent::Usage(u) => {
@@ -891,7 +889,7 @@ async fn chat_completions_stream(
}
// Channel closed or timeout — clean up
state.mitm_store.drop_channel().await;
state.mitm_store.remove_request(&cascade_id).await;
// If we got here from timeout with content, emit what we have
if !last_text.is_empty() || last_thinking_len > 0 {
@@ -1026,7 +1024,7 @@ async fn chat_completions_stream(
}
})).unwrap()));
// Always clear in-flight flag when stream ends
state.mitm_store.drop_channel().await;
state.mitm_store.remove_request(&cascade_id).await;
yield Ok(Event::default().data("[DONE]"));
};