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

149
src/api/completions.rs
View File

@@ -18,10 +18,6 @@ use super::util::{err_response, now_unix, upstream_err_response};
use super::AppState;
use crate::mitm::store::{CapturedFunctionCall, PendingToolResult, ToolRound};
/// System fingerprint for completions responses (derived from crate version at compile time).
fn system_fingerprint() -> String {
format!("fp_{}", env!("CARGO_PKG_VERSION").replace('.', ""))
@@ -181,8 +177,6 @@ pub(crate) async fn handle_completions(
model_name, body.stream
);
let model = match lookup_model(model_name) {
Some(m) => m,
None => {
@@ -200,22 +194,28 @@ pub(crate) async fn handle_completions(
// 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");
if gemini_tools.is_empty() {
None
} else {
info!(
count = t.len(),
"Completions: client tools for MITM injection"
);
Some(gemini_tools)
}
});
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)
})
body.tool_choice
.as_ref()
.map(crate::mitm::modify::openai_tool_choice_to_gemini)
});
// ── Extract tool results from messages for MITM injection ──────────
// 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 call_id_to_name: std::collections::HashMap<String, String> =
std::collections::HashMap::new();
{
let mut current_round: Option<ToolRound> = None;
@@ -266,10 +266,8 @@ pub(crate) async fn handle_completions(
"tool" => {
let text = extract_message_text(&msg.content);
if let Some(ref call_id) = msg.tool_call_id {
let result_index = current_round
.as_ref()
.map(|r| r.results.len())
.unwrap_or(0);
let result_index =
current_round.as_ref().map(|r| r.results.len()).unwrap_or(0);
let name = call_id_to_name
.get(call_id.as_str())
.cloned()
@@ -336,8 +334,7 @@ pub(crate) async fn handle_completions(
if merged > 0 {
info!(
merged_count = merged,
"Completions: merged {} thought_signature(s) from MITM capture",
merged,
"Completions: merged {} thought_signature(s) from MITM capture", merged,
);
}
}
@@ -431,7 +428,8 @@ pub(crate) async fn handle_completions(
});
// Get last calls from the latest tool round (if any) for proxy recording compat
let last_function_calls = tool_rounds.last()
let last_function_calls = tool_rounds
.last()
.map(|r| r.calls.clone())
.unwrap_or_default();
@@ -440,12 +438,18 @@ pub(crate) async fn handle_completions(
let (mitm_rx, event_tx) = (Some(rx), tx);
// Build pending tool results from latest round
let pending_tool_results = tool_rounds.last()
let pending_tool_results = tool_rounds
.last()
.map(|r| r.results.clone())
.unwrap_or_default();
// Start debug trace
let trace = state.trace.start(&cascade_id, "POST /v1/chat/completions", model_name, body.stream);
let trace = state.trace.start(
&cascade_id,
"POST /v1/chat/completions",
model_name,
body.stream,
);
if let Some(ref t) = trace {
t.set_client_request(crate::trace::ClientRequestSummary {
message_count: body.messages.len(),
@@ -455,35 +459,44 @@ pub(crate) async fn handle_completions(
user_text_preview: user_text.chars().take(200).collect(),
system_prompt: body.messages.iter().any(|m| m.role == "system"),
has_image: image.is_some(),
}).await;
})
.await;
// Start turn 0
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,
call_id_to_name,
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,
call_id_to_name,
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, _)) => {
@@ -495,7 +508,10 @@ pub(crate) async fn handle_completions(
}
Ok((status, _)) => {
state.mitm_store.remove_request(&cascade_id).await;
if let Some(ref t) = trace { t.record_error(format!("Backend returned {status}")).await; t.finish("backend_error").await; }
if let Some(ref t) = trace {
t.record_error(format!("Backend returned {status}")).await;
t.finish("backend_error").await;
}
return err_response(
StatusCode::BAD_GATEWAY,
format!("Backend returned {status}"),
@@ -504,7 +520,10 @@ pub(crate) async fn handle_completions(
}
Err(e) => {
state.mitm_store.remove_request(&cascade_id).await;
if let Some(ref t) = trace { t.record_error(format!("Send failed: {e}")).await; t.finish("send_error").await; }
if let Some(ref t) = trace {
t.record_error(format!("Send failed: {e}")).await;
t.finish("send_error").await;
}
return err_response(
StatusCode::BAD_GATEWAY,
format!("Send failed: {e}"),
@@ -515,10 +534,8 @@ pub(crate) async fn handle_completions(
// 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 {
@@ -549,7 +566,7 @@ pub(crate) async fn handle_completions(
let include_usage = body
.stream_options
.as_ref()
.map_or(false, |o| o.include_usage);
.is_some_and(|o| o.include_usage);
if body.stream {
chat_completions_stream(
@@ -582,7 +599,12 @@ pub(crate) async fn handle_completions(
// Send the same message on each extra cascade
match state
.backend
.send_message_with_image(&cid, &format!(".<cid:{}>", cid), model.model_enum, image.as_ref())
.send_message_with_image(
&cid,
&format!(".<cid:{}>", cid),
model.model_enum,
image.as_ref(),
)
.await
{
Ok((200, _)) => {
@@ -783,7 +805,7 @@ async fn chat_completions_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();
tool_calls.push(serde_json::json!({
@@ -885,7 +907,7 @@ async fn chat_completions_stream(
did_unblock_ls = true;
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!(
@@ -1111,7 +1133,7 @@ async fn chat_completions_stream(
// Keep-alive comment every ~5 iterations
keepalive_counter += 1;
if keepalive_counter % 5 == 0 {
if keepalive_counter.is_multiple_of(5) {
yield Ok(Event::default().comment("keepalive"));
}
@@ -1193,21 +1215,26 @@ async fn chat_completions_sync(
// Record trace data
if let Some(ref t) = trace {
t.record_response(0, crate::trace::ResponseSummary {
text_len: result.text.len(),
thinking_len: result.thinking.as_ref().map_or(0, |s| s.len()),
text_preview: result.text.chars().take(200).collect(),
finish_reason: Some(finish_reason.to_string()),
function_calls: Vec::new(),
grounding: false,
}).await;
t.record_response(
0,
crate::trace::ResponseSummary {
text_len: result.text.len(),
thinking_len: result.thinking.as_ref().map_or(0, |s| s.len()),
text_preview: result.text.chars().take(200).collect(),
finish_reason: Some(finish_reason.to_string()),
function_calls: Vec::new(),
grounding: false,
},
)
.await;
if prompt_tokens > 0 || completion_tokens > 0 {
t.set_usage(crate::trace::TrackedUsage {
input_tokens: prompt_tokens,
output_tokens: completion_tokens,
thinking_tokens: thinking_tokens,
thinking_tokens,
cache_read: cached_tokens,
}).await;
})
.await;
}
t.finish("completed").await;
}