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

fix: block ALL LS follow-up requests across connections

Browse Source 
Move the in-flight blocking check to the top of the LLM request flow,
BEFORE request modification. This catches follow-ups on ALL connections
(the LS opens multiple parallel TLS connections). Only the very first
modified request reaches Google — all others get fake STOP responses.

Previously, each new connection independently allowed one request
through before blocking, letting 4-5 requests leak per turn.
This commit is contained in:
Nikketryhard
2026-02-16 00:57:33 -06:00
parent a8f3c8915f
commit 3fdd0368a0
23 changed files with 992 additions and 568 deletions
Download Patch File Download Diff File Expand all files Collapse all files

105
src/api/completions.rs
View File

@@ -10,9 +10,11 @@ use std::sync::Arc;
use tracing::{debug, info, warn};
use super::models::{lookup_model, DEFAULT_MODEL, MODELS};
use super::polling::{extract_response_text, extract_thinking_content, is_response_done, poll_for_response};
use super::polling::{
extract_response_text, extract_thinking_content, is_response_done, poll_for_response,
};
use super::types::*;
use super::util::{err_response, upstream_err_response, now_unix};
use super::util::{err_response, now_unix, upstream_err_response};
use super::AppState;
/// Extract a conversation/session ID from a flexible JSON value.
@@ -33,7 +35,8 @@ fn system_fingerprint() -> String {
/// Build a streaming chunk JSON with all required OpenAI fields.
/// Includes system_fingerprint, service_tier, and logprobs:null in choices.
fn chunk_json(
id: &str, model: &str,
id: &str,
model: &str,
choices: serde_json::Value,
usage: Option<serde_json::Value>,
) -> String {
@@ -53,7 +56,11 @@ fn chunk_json(
}
/// Build a single choice for a streaming chunk (delta + finish_reason + logprobs).
fn chunk_choice(index: u32, delta: serde_json::Value, finish_reason: Option<&str>) -> serde_json::Value {
fn chunk_choice(
index: u32,
delta: serde_json::Value,
finish_reason: Option<&str>,
) -> serde_json::Value {
serde_json::json!({
"index": index,
"delta": delta,
@@ -70,7 +77,9 @@ fn chunk_choice(index: u32, delta: serde_json::Value, finish_reason: Option<&str
fn google_to_openai_finish_reason(stop_reason: Option<&str>) -> &'static str {
match stop_reason {
Some("MAX_TOKENS") => "length",
Some("SAFETY") | Some("RECITATION") | Some("BLOCKLIST") | Some("PROHIBITED_CONTENT") => "content_filter",
Some("SAFETY") | Some("RECITATION") | Some("BLOCKLIST") | Some("PROHIBITED_CONTENT") => {
"content_filter"
}
_ => "stop",
}
}
@@ -84,7 +93,9 @@ fn google_to_openai_finish_reason(stop_reason: Option<&str>) -> &'static str {
/// sends the entire messages array to the model.
fn extract_chat_input(messages: &[CompletionMessage]) -> (String, Option<crate::proto::ImageData>) {
// Extract image from last user message content array
let image = messages.iter().rev()
let image = messages
.iter()
.rev()
.find(|m| m.role == "user")
.and_then(|m| super::util::extract_first_image(&m.content));
// Always build the full conversation
@@ -141,10 +152,7 @@ fn build_conversation_with_tools(messages: &[CompletionMessage]) -> String {
if let Some(func) = tc.get("function") {
let name = func["name"].as_str().unwrap_or("unknown");
let args = func["arguments"].as_str().unwrap_or("{}");
parts.push(format!(
"[Tool call: {}({})]",
name, args
));
parts.push(format!("[Tool call: {}({})]", name, args));
}
}
}
@@ -153,10 +161,7 @@ fn build_conversation_with_tools(messages: &[CompletionMessage]) -> String {
let text = extract_message_text(&msg.content);
let tool_id = msg.tool_call_id.as_deref().unwrap_or("unknown");
if !text.is_empty() {
parts.push(format!(
"[Tool result ({})]:\n{}",
tool_id, text
));
parts.push(format!("[Tool result ({})]:\n{}", tool_id, text));
}
}
_ => {}
@@ -202,7 +207,10 @@ pub(crate) async fn handle_completions(
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");
info!(
count = tools.len(),
"Completions: stored client tools for MITM injection"
);
} else {
state.mitm_store.clear_tools().await;
}
@@ -239,10 +247,15 @@ pub(crate) async fn handle_completions(
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()
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;
@@ -306,12 +319,13 @@ pub(crate) async fn handle_completions(
// 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 {
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;
})
.await;
}
match state
.backend
@@ -346,7 +360,10 @@ pub(crate) async fn handle_completions(
uuid::Uuid::new_v4().to_string().replace('-', "")
);
let include_usage = body.stream_options.as_ref().map_or(false, |o| o.include_usage);
let include_usage = body
.stream_options
.as_ref()
.map_or(false, |o| o.include_usage);
if body.stream {
chat_completions_stream(
@@ -374,11 +391,17 @@ pub(crate) async fn handle_completions(
match state.backend.create_cascade().await {
Ok(cid) => {
// Send the same message on each extra cascade
match state.backend.send_message_with_image(&cid, &user_text, model.model_enum, image.as_ref()).await {
match state
.backend
.send_message_with_image(&cid, &user_text, model.model_enum, image.as_ref())
.await
{
Ok((200, _)) => {
let bg = Arc::clone(&state.backend);
let cid2 = cid.clone();
tokio::spawn(async move { let _ = bg.update_annotations(&cid2).await; });
tokio::spawn(async move {
let _ = bg.update_annotations(&cid2).await;
});
extra_cascade_ids.push(cid);
}
_ => {} // Skip failed cascades
@@ -420,7 +443,12 @@ pub(crate) async fn handle_completions(
mitm.as_ref().and_then(|u| u.stop_reason.as_deref()),
);
let (pt, ct, cached, thinking) = if let Some(ref mu) = mitm {
(mu.input_tokens, mu.output_tokens, mu.cache_read_input_tokens, mu.thinking_output_tokens)
(
mu.input_tokens,
mu.output_tokens,
mu.cache_read_input_tokens,
mu.thinking_output_tokens,
)
} else if let Some(u) = &result.usage {
(u.input_tokens, u.output_tokens, 0, 0)
} else {
@@ -874,15 +902,22 @@ async fn chat_completions_sync(
None => state.mitm_store.take_usage("_latest").await,
};
let finish_reason = google_to_openai_finish_reason(mitm.as_ref().and_then(|u| u.stop_reason.as_deref()));
let finish_reason =
google_to_openai_finish_reason(mitm.as_ref().and_then(|u| u.stop_reason.as_deref()));
let (prompt_tokens, completion_tokens, cached_tokens, thinking_tokens) = if let Some(ref mitm_usage) = mitm {
(mitm_usage.input_tokens, mitm_usage.output_tokens, mitm_usage.cache_read_input_tokens, mitm_usage.thinking_output_tokens)
} else if let Some(u) = &result.usage {
(u.input_tokens, u.output_tokens, 0, 0)
} else {
(0, 0, 0, 0)
};
let (prompt_tokens, completion_tokens, cached_tokens, thinking_tokens) =
if let Some(ref mitm_usage) = mitm {
(
mitm_usage.input_tokens,
mitm_usage.output_tokens,
mitm_usage.cache_read_input_tokens,
mitm_usage.thinking_output_tokens,
)
} else if let Some(u) = &result.usage {
(u.input_tokens, u.output_tokens, 0, 0)
} else {
(0, 0, 0, 0)
};
// Build message object, including reasoning_content if thinking is present
let mut message = serde_json::json!({