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

feat: forward generation params via MITM + add usageMetadata to Gemini

Browse Source 
- Add GenerationParams struct to MitmStore for temperature, top_p,
  top_k, max_output_tokens, stop_sequences, frequency/presence_penalty
- MITM modify_request injects params into request.generationConfig
- All 3 endpoints (Completions, Responses, Gemini) store client params
- Add usageMetadata to Gemini sync responses (promptTokenCount,
  candidatesTokenCount, totalTokenCount, thoughtsTokenCount)
- Add generation param fields to GeminiRequest (temperature, topP, etc.)
- Completions stream_options.include_usage emits final usage chunk
- Completions reasoning_tokens in completion_tokens_details
- Update endpoint gap analysis doc (all high-priority gaps resolved)
This commit is contained in:
Nikketryhard
2026-02-15 14:23:05 -06:00
parent 735c3e357d
commit b1bd57ab5e
9 changed files with 1216 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files

242
src/api/completions.rs
View File

@@ -10,7 +10,7 @@ use std::sync::Arc;
use tracing::{debug, info, warn};
use super::models::{lookup_model, DEFAULT_MODEL, MODELS};
use super::polling::{extract_response_text, 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, now_unix};
use super::AppState;
@@ -176,6 +176,28 @@ pub(crate) async fn handle_completions(
}
state.mitm_store.clear_active_function_call();
// Store generation parameters for MITM injection
{
use crate::mitm::store::GenerationParams;
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: None, // TODO: body.stop
frequency_penalty: body.frequency_penalty,
presence_penalty: body.presence_penalty,
};
// 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()
{
state.mitm_store.set_generation_params(gp).await;
} else {
state.mitm_store.clear_generation_params().await;
}
}
let token = state.backend.oauth_token().await;
if token.is_empty() {
return err_response(
@@ -241,6 +263,8 @@ 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);
if body.stream {
chat_completions_stream(
state,
@@ -248,6 +272,7 @@ pub(crate) async fn handle_completions(
model_name.to_string(),
cascade_id,
body.timeout,
include_usage,
)
.await
} else {
@@ -271,6 +296,7 @@ async fn chat_completions_stream(
model_name: String,
cascade_id: String,
timeout: u64,
include_usage: bool,
) -> axum::response::Response {
let stream = async_stream::stream! {
let start = std::time::Instant::now();
@@ -294,6 +320,7 @@ async fn chat_completions_stream(
})).unwrap_or_default()));
let mut keepalive_counter: u64 = 0;
let mut last_thinking_len: usize = 0;
while start.elapsed().as_secs() < timeout {
// ── Check for MITM-captured function calls FIRST ──
@@ -342,6 +369,27 @@ async fn chat_completions_stream(
"finish_reason": "tool_calls",
}],
})).unwrap_or_default()));
if include_usage {
let mitm = state.mitm_store.take_usage(&cascade_id).await
.or(state.mitm_store.take_usage("_latest").await);
let (pt, ct, crt, tt) = if let Some(ref u) = mitm {
(u.input_tokens, u.output_tokens, u.cache_read_input_tokens, u.thinking_output_tokens)
} else { (0, 0, 0, 0) };
yield Ok(Event::default().data(serde_json::to_string(&serde_json::json!({
"id": completion_id,
"object": "chat.completion.chunk",
"created": now_unix(),
"model": model_name,
"choices": [],
"usage": {
"prompt_tokens": pt,
"completion_tokens": ct,
"total_tokens": pt + ct,
"prompt_tokens_details": { "cached_tokens": crt },
"completion_tokens_details": { "reasoning_tokens": tt },
},
})).unwrap_or_default()));
}
yield Ok(Event::default().data("[DONE]"));
return;
}
@@ -349,7 +397,37 @@ async fn chat_completions_stream(
// ── Check for MITM-captured response text (bypass LS) ──
if has_custom_tools {
if let Some(text) = state.mitm_store.peek_response_text().await {
let peek = state.mitm_store.peek_response_text().await;
let complete = state.mitm_store.is_response_complete();
let has_fc = state.mitm_store.has_active_function_call();
if keepalive_counter % 10 == 0 || peek.is_some() || complete || has_fc {
debug!(
"Completions bypass poll: peek={}, complete={}, has_fc={}, last_text_len={}",
peek.as_ref().map(|t| t.len()).unwrap_or(0),
complete, has_fc, last_text.len()
);
}
// Stream thinking text as reasoning_content deltas (MITM bypass)
if let Some(tc) = state.mitm_store.peek_thinking_text().await {
if tc.len() > last_thinking_len {
let delta = &tc[last_thinking_len..];
last_thinking_len = tc.len();
yield Ok(Event::default().data(serde_json::to_string(&serde_json::json!({
"id": completion_id,
"object": "chat.completion.chunk",
"created": now_unix(),
"model": model_name,
"choices": [{
"index": 0,
"delta": {"reasoning_content": delta},
"finish_reason": serde_json::Value::Null,
}],
})).unwrap_or_default()));
}
}
if let Some(text) = peek {
if !text.is_empty() && text != last_text {
let delta = if text.len() > last_text.len() && text.starts_with(&*last_text) {
text[last_text.len()..].to_string()
@@ -387,12 +465,33 @@ async fn chat_completions_stream(
"finish_reason": "stop",
}],
})).unwrap_or_default()));
if include_usage {
let mitm = state.mitm_store.take_usage(&cascade_id).await
.or(state.mitm_store.take_usage("_latest").await);
let (pt, ct, crt, tt) = if let Some(ref u) = mitm {
(u.input_tokens, u.output_tokens, u.cache_read_input_tokens, u.thinking_output_tokens)
} else { (0, 0, 0, 0) };
yield Ok(Event::default().data(serde_json::to_string(&serde_json::json!({
"id": completion_id,
"object": "chat.completion.chunk",
"created": now_unix(),
"model": model_name,
"choices": [],
"usage": {
"prompt_tokens": pt,
"completion_tokens": ct,
"total_tokens": pt + ct,
"prompt_tokens_details": { "cached_tokens": crt },
"completion_tokens_details": { "reasoning_tokens": tt },
},
})).unwrap_or_default()));
}
yield Ok(Event::default().data("[DONE]"));
return;
}
} else if state.mitm_store.is_response_complete() {
// Response complete but no text — might be a tool call we already handled
// or an empty response. Give it a moment then bail.
} else if complete {
// Response complete but no text — might be a tool call arriving shortly,
// stale state from a previous request, or an empty response.
tokio::time::sleep(tokio::time::Duration::from_millis(200)).await;
// Re-check function calls one more time
let final_check = state.mitm_store.take_any_function_calls().await;
@@ -437,10 +536,35 @@ async fn chat_completions_stream(
"finish_reason": "tool_calls",
}],
})).unwrap_or_default()));
if include_usage {
let mitm = state.mitm_store.take_usage(&cascade_id).await
.or(state.mitm_store.take_usage("_latest").await);
let (pt, ct, crt, tt) = if let Some(ref u) = mitm {
(u.input_tokens, u.output_tokens, u.cache_read_input_tokens, u.thinking_output_tokens)
} else { (0, 0, 0, 0) };
yield Ok(Event::default().data(serde_json::to_string(&serde_json::json!({
"id": completion_id,
"object": "chat.completion.chunk",
"created": now_unix(),
"model": model_name,
"choices": [],
"usage": {
"prompt_tokens": pt,
"completion_tokens": ct,
"total_tokens": pt + ct,
"prompt_tokens_details": { "cached_tokens": crt },
"completion_tokens_details": { "reasoning_tokens": tt },
},
})).unwrap_or_default()));
}
yield Ok(Event::default().data("[DONE]"));
return;
}
}
// No text and no function calls but complete=true: stale state.
// Clear the flag so we wait for the real response from this request.
warn!("Completions: stale response_complete detected (no text, no FC) — clearing");
state.mitm_store.clear_response_async().await;
}
// When using bypass mode, skip LS step polling
@@ -457,6 +581,26 @@ async fn chat_completions_stream(
if let Ok((status, data)) = state.backend.get_steps(&cascade_id).await {
if status == 200 {
if let Some(steps) = data["steps"].as_array() {
// Stream thinking deltas (reasoning_content)
if let Some(tc) = extract_thinking_content(steps) {
if tc.len() > last_thinking_len {
let delta = &tc[last_thinking_len..];
last_thinking_len = tc.len();
yield Ok(Event::default().data(serde_json::to_string(&serde_json::json!({
"id": completion_id,
"object": "chat.completion.chunk",
"created": now_unix(),
"model": model_name,
"choices": [{
"index": 0,
"delta": {"reasoning_content": delta},
"finish_reason": serde_json::Value::Null,
}],
})).unwrap_or_default()));
}
}
let text = extract_response_text(steps);
if !text.is_empty() && text != last_text {
@@ -496,6 +640,27 @@ async fn chat_completions_stream(
"finish_reason": "stop",
}],
})).unwrap_or_default()));
if include_usage {
let mitm = state.mitm_store.take_usage(&cascade_id).await
.or(state.mitm_store.take_usage("_latest").await);
let (pt, ct, crt, tt) = if let Some(ref u) = mitm {
(u.input_tokens, u.output_tokens, u.cache_read_input_tokens, u.thinking_output_tokens)
} else { (0, 0, 0, 0) };
yield Ok(Event::default().data(serde_json::to_string(&serde_json::json!({
"id": completion_id,
"object": "chat.completion.chunk",
"created": now_unix(),
"model": model_name,
"choices": [],
"usage": {
"prompt_tokens": pt,
"completion_tokens": ct,
"total_tokens": pt + ct,
"prompt_tokens_details": { "cached_tokens": crt },
"completion_tokens_details": { "reasoning_tokens": tt },
},
})).unwrap_or_default()));
}
yield Ok(Event::default().data("[DONE]"));
return;
}
@@ -519,6 +684,27 @@ async fn chat_completions_stream(
"finish_reason": "stop",
}],
})).unwrap_or_default()));
if include_usage {
let mitm = state.mitm_store.take_usage(&cascade_id).await
.or(state.mitm_store.take_usage("_latest").await);
let (pt, ct, crt, tt) = if let Some(ref u) = mitm {
(u.input_tokens, u.output_tokens, u.cache_read_input_tokens, u.thinking_output_tokens)
} else { (0, 0, 0, 0) };
yield Ok(Event::default().data(serde_json::to_string(&serde_json::json!({
"id": completion_id,
"object": "chat.completion.chunk",
"created": now_unix(),
"model": model_name,
"choices": [],
"usage": {
"prompt_tokens": pt,
"completion_tokens": ct,
"total_tokens": pt + ct,
"prompt_tokens_details": { "cached_tokens": crt },
"completion_tokens_details": { "reasoning_tokens": tt },
},
})).unwrap_or_default()));
}
yield Ok(Event::default().data("[DONE]"));
return;
}
@@ -553,6 +739,27 @@ async fn chat_completions_stream(
"finish_reason": "stop",
}],
})).unwrap_or_default()));
if include_usage {
let mitm = state.mitm_store.take_usage(&cascade_id).await
.or(state.mitm_store.take_usage("_latest").await);
let (pt, ct, crt, tt) = if let Some(ref u) = mitm {
(u.input_tokens, u.output_tokens, u.cache_read_input_tokens, u.thinking_output_tokens)
} else { (0, 0, 0, 0) };
yield Ok(Event::default().data(serde_json::to_string(&serde_json::json!({
"id": completion_id,
"object": "chat.completion.chunk",
"created": now_unix(),
"model": model_name,
"choices": [],
"usage": {
"prompt_tokens": pt,
"completion_tokens": ct,
"total_tokens": pt + ct,
"prompt_tokens_details": { "cached_tokens": crt },
"completion_tokens_details": { "reasoning_tokens": tt },
},
})).unwrap_or_default()));
}
yield Ok(Event::default().data("[DONE]"));
};
@@ -582,14 +789,23 @@ async fn chat_completions_sync(
Some(u) => Some(u),
None => state.mitm_store.take_usage("_latest").await,
};
let (prompt_tokens, completion_tokens, cached_tokens) = if let Some(mitm_usage) = mitm {
(mitm_usage.input_tokens, mitm_usage.output_tokens, mitm_usage.cache_read_input_tokens)
let (prompt_tokens, completion_tokens, cached_tokens, thinking_tokens) = if let Some(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)
(u.input_tokens, u.output_tokens, 0, 0)
} else {
(0, 0, 0)
(0, 0, 0, 0)
};
// Build message object, including reasoning_content if thinking is present
let mut message = serde_json::json!({
"role": "assistant",
"content": result.text,
});
if let Some(ref thinking) = result.thinking {
message["reasoning_content"] = serde_json::json!(thinking);
}
Json(serde_json::json!({
"id": completion_id,
"object": "chat.completion",
@@ -597,10 +813,7 @@ async fn chat_completions_sync(
"model": model_name,
"choices": [{
"index": 0,
"message": {
"role": "assistant",
"content": result.text,
},
"message": message,
"finish_reason": "stop",
}],
"usage": {
@@ -610,6 +823,9 @@ async fn chat_completions_sync(
"prompt_tokens_details": {
"cached_tokens": cached_tokens,
},
"completion_tokens_details": {
"reasoning_tokens": thinking_tokens,
},
},
}))
.into_response()