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

feat: Add LICENSE file and refactor MITM response handling and tracing.

Browse Source
This commit is contained in:
Nikketryhard
2026-02-18 02:43:05 -06:00
parent c0c12de83c
commit ad0aa1556c
26 changed files with 1132 additions and 569 deletions
Download Patch File Download Diff File Expand all files Collapse all files

280
src/api/responses.rs
View File

@@ -142,10 +142,6 @@ fn extract_responses_input(
(final_text, tool_results, image)
}
/// Response-specific data for building a Response object.
struct ResponseData {
id: String,
@@ -270,7 +266,7 @@ pub(crate) async fn handle_responses(
// ── 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| {
let has_web_search = body.tools.as_ref().is_some_and(|tools| {
tools.iter().any(|t| {
let t_type = t["type"].as_str().unwrap_or("");
t_type == "web_search_preview" || t_type == "web_search"
@@ -280,14 +276,14 @@ pub(crate) async fn handle_responses(
// 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 {
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)
});
let tool_config = body.tool_choice.as_ref().map(openai_tool_choice_to_gemini);
// Build generation params locally
let (response_mime_type, response_schema, text_format) = if let Some(ref text_val) = body.text {
@@ -372,7 +368,10 @@ pub(crate) async fn handle_responses(
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
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,
@@ -381,7 +380,9 @@ pub(crate) async fn handle_responses(
}
// Start debug trace
let trace = state.trace.start(&cascade_id, "POST /v1/responses", &model.name, body.stream);
let trace = state
.trace
.start(&cascade_id, "POST /v1/responses", model.name, body.stream);
if let Some(ref t) = trace {
t.set_client_request(crate::trace::ClientRequestSummary {
message_count: if is_tool_result_turn { 0 } else { 1 },
@@ -391,34 +392,43 @@ pub(crate) async fn handle_responses(
user_text_preview: user_text.chars().take(200).collect(),
system_prompt: body.instructions.is_some(),
has_image: image.is_some(),
}).await;
})
.await;
t.start_turn().await;
}
let mitm_gate = std::sync::Arc::new(tokio::sync::Notify::new());
let mitm_gate_clone = mitm_gate.clone();
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(),
gate: mitm_gate_clone,
trace_handle: trace.clone(),
trace_turn: 0,
}).await;
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(),
gate: mitm_gate_clone,
trace_handle: trace.clone(),
trace_turn: 0,
})
.await;
// Send REAL user text to LS
match state
.backend
.send_message_with_image(&cascade_id, &format!(".<cid:{}>", 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, _)) => {
@@ -448,15 +458,16 @@ pub(crate) async fn handle_responses(
// Wait for MITM gate: 5s → 502 if MITM enabled
let gate_start = std::time::Instant::now();
let gate_matched = tokio::time::timeout(
std::time::Duration::from_secs(5),
mitm_gate.notified(),
).await;
let gate_matched =
tokio::time::timeout(std::time::Duration::from_secs(5), mitm_gate.notified()).await;
let gate_wait_ms = gate_start.elapsed().as_millis() as u64;
if gate_matched.is_err() {
if state.mitm_enabled {
state.mitm_store.remove_request(&cascade_id).await;
if let Some(ref t) = trace { t.record_error("MITM gate timeout (5s)".to_string()).await; t.finish("mitm_timeout").await; }
if let Some(ref t) = trace {
t.record_error("MITM gate timeout (5s)".to_string()).await;
t.finish("mitm_timeout").await;
}
return err_response(
StatusCode::BAD_GATEWAY,
"MITM proxy did not match request within 5s".to_string(),
@@ -466,7 +477,9 @@ pub(crate) async fn handle_responses(
warn!(cascade = %cascade_id, "MITM gate timeout (--no-mitm mode)");
} else {
debug!(cascade = %cascade_id, gate_wait_ms, "MITM gate signaled — request matched");
if let Some(ref t) = trace { t.record_mitm_match(0, gate_wait_ms).await; }
if let Some(ref t) = trace {
t.record_mitm_match(0, gate_wait_ms).await;
}
}
// Capture request params for response building
@@ -655,12 +668,22 @@ async fn handle_responses_sync(
while let Some(event) = tokio::time::timeout(
std::time::Duration::from_secs(timeout.saturating_sub(start.elapsed().as_secs())),
rx.recv(),
).await.ok().flatten() {
)
.await
.ok()
.flatten()
{
use crate::mitm::store::MitmEvent;
match event {
MitmEvent::ThinkingDelta(t) => { acc_thinking = Some(t); }
MitmEvent::TextDelta(t) => { acc_text = t; }
MitmEvent::Usage(u) => { _last_usage = Some(u); }
MitmEvent::ThinkingDelta(t) => {
acc_thinking = Some(t);
}
MitmEvent::TextDelta(t) => {
acc_text = t;
}
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 {
@@ -672,38 +695,57 @@ async fn handle_responses_sync(
for fc in &calls {
let call_id = format!(
"call_{}",
uuid::Uuid::new_v4().to_string().replace('-', "")[..24].to_string()
&uuid::Uuid::new_v4().to_string().replace('-', "")[..24]
);
state.mitm_store.register_call_id(&cascade_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));
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,
&cascade_id,
&None,
&params.user_text,
"",
)
.await;
state.mitm_store.remove_request(&cascade_id).await;
// Record trace before usage is moved
if let Some(ref t) = trace {
let fc_summaries: Vec<crate::trace::FunctionCallSummary> = calls.iter().map(|fc| {
crate::trace::FunctionCallSummary {
let fc_summaries: Vec<crate::trace::FunctionCallSummary> = calls
.iter()
.map(|fc| crate::trace::FunctionCallSummary {
name: fc.name.clone(),
args_preview: serde_json::to_string(&fc.args).unwrap_or_default().chars().take(200).collect(),
}
}).collect();
t.record_response(0, crate::trace::ResponseSummary {
text_len: 0,
thinking_len: 0,
text_preview: String::new(),
finish_reason: Some("tool_calls".to_string()),
function_calls: fc_summaries,
grounding: false,
}).await;
args_preview: serde_json::to_string(&fc.args)
.unwrap_or_default()
.chars()
.take(200)
.collect(),
})
.collect();
t.record_response(
0,
crate::trace::ResponseSummary {
text_len: 0,
thinking_len: 0,
text_preview: String::new(),
finish_reason: Some("tool_calls".to_string()),
function_calls: fc_summaries,
grounding: false,
},
)
.await;
t.set_usage(crate::trace::TrackedUsage {
input_tokens: usage.input_tokens,
output_tokens: usage.output_tokens,
thinking_tokens: usage.output_tokens_details.reasoning_tokens,
cache_read: usage.input_tokens_details.cached_tokens,
}).await;
})
.await;
t.finish("tool_call").await;
}
let resp = build_response_object(
@@ -731,7 +773,7 @@ async fn handle_responses_sync(
// Reinstall channel and unblock gate.
let (new_tx, new_rx) = tokio::sync::mpsc::channel(64);
state.mitm_store.set_channel(&cascade_id, new_tx).await;
let _ = state.mitm_store.take_any_function_calls().await;
rx = new_rx;
debug!(
@@ -741,33 +783,44 @@ async fn handle_responses_sync(
continue;
}
let (usage, _) = usage_from_poll(
&state.mitm_store, &cascade_id, &None, &params.user_text, &acc_text,
).await;
&state.mitm_store,
&cascade_id,
&None,
&params.user_text,
&acc_text,
)
.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 {
output_items.push(build_reasoning_output(t));
}
let msg_id = format!("msg_{}", uuid::Uuid::new_v4().to_string().replace('-', ""));
let msg_id =
format!("msg_{}", uuid::Uuid::new_v4().to_string().replace('-', ""));
output_items.push(build_message_output(&msg_id, &acc_text));
// Record trace before usage is moved
if let Some(ref t) = trace {
t.record_response(0, crate::trace::ResponseSummary {
text_len: acc_text.len(),
thinking_len: acc_thinking.as_ref().map_or(0, |s| s.len()),
text_preview: acc_text.chars().take(200).collect(),
finish_reason: Some("stop".to_string()),
function_calls: Vec::new(),
grounding: false,
}).await;
t.record_response(
0,
crate::trace::ResponseSummary {
text_len: acc_text.len(),
thinking_len: acc_thinking.as_ref().map_or(0, |s| s.len()),
text_preview: acc_text.chars().take(200).collect(),
finish_reason: Some("stop".to_string()),
function_calls: Vec::new(),
grounding: false,
},
)
.await;
t.set_usage(crate::trace::TrackedUsage {
input_tokens: usage.input_tokens,
output_tokens: usage.output_tokens,
thinking_tokens: usage.output_tokens_details.reasoning_tokens,
cache_read: usage.input_tokens_details.cached_tokens,
}).await;
})
.await;
t.finish("completed").await;
}
let resp = build_response_object(
@@ -787,7 +840,14 @@ async fn handle_responses_sync(
}
MitmEvent::UpstreamError(err) => {
state.mitm_store.remove_request(&cascade_id).await;
if let Some(ref t) = trace { t.record_error(format!("Upstream: {}", err.message.as_deref().unwrap_or("unknown"))).await; t.finish("upstream_error").await; }
if let Some(ref t) = trace {
t.record_error(format!(
"Upstream: {}",
err.message.as_deref().unwrap_or("unknown")
))
.await;
t.finish("upstream_error").await;
}
return upstream_err_response(&err);
}
}
@@ -795,7 +855,10 @@ async fn handle_responses_sync(
// Timeout
state.mitm_store.remove_request(&cascade_id).await;
if let Some(ref t) = trace { t.record_error(format!("Timeout: {}s", timeout)).await; t.finish("timeout").await; }
if let Some(ref t) = trace {
t.record_error(format!("Timeout: {}s", timeout)).await;
t.finish("timeout").await;
}
return err_response(
StatusCode::GATEWAY_TIMEOUT,
format!("Timeout: no response from Google API after {timeout}s"),
@@ -834,7 +897,7 @@ async fn handle_responses_sync(
for fc in calls {
let call_id = format!(
"call_{}",
uuid::Uuid::new_v4().to_string().replace('-', "")[..24].to_string()
&uuid::Uuid::new_v4().to_string().replace('-', "")[..24]
);
// Register call_id → name mapping for tool result routing
state
@@ -858,26 +921,36 @@ async fn handle_responses_sync(
// Record trace before usage is moved
if let Some(ref t) = trace {
let fc_summaries: Vec<crate::trace::FunctionCallSummary> = calls.iter().map(|fc| {
crate::trace::FunctionCallSummary {
let fc_summaries: Vec<crate::trace::FunctionCallSummary> = calls
.iter()
.map(|fc| crate::trace::FunctionCallSummary {
name: fc.name.clone(),
args_preview: serde_json::to_string(&fc.args).unwrap_or_default().chars().take(200).collect(),
}
}).collect();
t.record_response(0, crate::trace::ResponseSummary {
text_len: poll_result.text.len(),
thinking_len: poll_result.thinking.as_ref().map_or(0, |s| s.len()),
text_preview: String::new(),
finish_reason: Some("tool_calls".to_string()),
function_calls: fc_summaries,
grounding: false,
}).await;
args_preview: serde_json::to_string(&fc.args)
.unwrap_or_default()
.chars()
.take(200)
.collect(),
})
.collect();
t.record_response(
0,
crate::trace::ResponseSummary {
text_len: poll_result.text.len(),
thinking_len: poll_result.thinking.as_ref().map_or(0, |s| s.len()),
text_preview: String::new(),
finish_reason: Some("tool_calls".to_string()),
function_calls: fc_summaries,
grounding: false,
},
)
.await;
t.set_usage(crate::trace::TrackedUsage {
input_tokens: usage.input_tokens,
output_tokens: usage.output_tokens,
thinking_tokens: usage.output_tokens_details.reasoning_tokens,
cache_read: usage.input_tokens_details.cached_tokens,
}).await;
})
.await;
t.finish("tool_call").await;
}
@@ -920,20 +993,25 @@ async fn handle_responses_sync(
// Record trace before usage is moved
if let Some(ref t) = trace {
t.record_response(0, crate::trace::ResponseSummary {
text_len: poll_result.text.len(),
thinking_len: thinking_text.as_ref().map_or(0, |s| s.len()),
text_preview: poll_result.text.chars().take(200).collect(),
finish_reason: Some("stop".to_string()),
function_calls: Vec::new(),
grounding: false,
}).await;
t.record_response(
0,
crate::trace::ResponseSummary {
text_len: poll_result.text.len(),
thinking_len: thinking_text.as_ref().map_or(0, |s| s.len()),
text_preview: poll_result.text.chars().take(200).collect(),
finish_reason: Some("stop".to_string()),
function_calls: Vec::new(),
grounding: false,
},
)
.await;
t.set_usage(crate::trace::TrackedUsage {
input_tokens: usage.input_tokens,
output_tokens: usage.output_tokens,
thinking_tokens: usage.output_tokens_details.reasoning_tokens,
cache_read: usage.input_tokens_details.cached_tokens,
}).await;
})
.await;
t.finish("completed").await;
}
@@ -1184,7 +1262,7 @@ async fn handle_responses_stream(
for (i, fc) in calls.iter().enumerate() {
let call_id = format!(
"call_{}",
uuid::Uuid::new_v4().to_string().replace('-', "")[..24].to_string()
&uuid::Uuid::new_v4().to_string().replace('-', "")[..24]
);
let arguments = serde_json::to_string(&fc.args).unwrap_or_default();
state.mitm_store.register_call_id(&cascade_id, call_id.clone(), fc.name.clone()).await;
@@ -1229,7 +1307,7 @@ async fn handle_responses_stream(
for fc in &calls {
let call_id = format!(
"call_{}",
uuid::Uuid::new_v4().to_string().replace('-', "")[..24].to_string()
&uuid::Uuid::new_v4().to_string().replace('-', "")[..24]
);
let arguments = serde_json::to_string(&fc.args).unwrap_or_default();
output_items.push(build_function_call_output(&call_id, &fc.name, &arguments));
@@ -1317,7 +1395,7 @@ async fn handle_responses_stream(
// Create a new channel and unblock the gate.
let (new_tx, new_rx) = tokio::sync::mpsc::channel(64);
state.mitm_store.set_channel(&cascade_id, new_tx).await;
let _ = state.mitm_store.take_any_function_calls().await;
rx = new_rx;
debug!(