feat: capture function calls from Google + block follow-up quota waste

When MITM strips LS tools and injects custom tools:
- Google returns functionCall → captured in MitmStore
- Follow-up LS requests are blocked with fake SSE response
- Proxy consumes captured calls and clears the flag
- Result: 1 real Google API call instead of 5+ per tool call

Flow: Client → Proxy → LS → MITM(inject tool) → Google
      Google returns functionCall → MITM captures it
      LS tries follow-up → MITM blocks (fake response)
      Proxy reads captured functionCall → returns to client
This commit is contained in:
Nikketryhard
2026-02-14 22:37:28 -06:00
parent 146be139a2
commit 8455aa674f
5 changed files with 161 additions and 5 deletions

View File

@@ -2,9 +2,9 @@
//!
//! Handles both streaming (SSE) and non-streaming (JSON) responses.
use super::store::ApiUsage;
use super::store::{ApiUsage, CapturedFunctionCall};
use serde_json::Value;
use tracing::{debug, trace};
use tracing::{debug, info, trace};
/// Parse a complete (non-streaming) Anthropic Messages API response body.
///
@@ -66,6 +66,8 @@ pub struct StreamingAccumulator {
pub stop_reason: Option<String>,
pub is_complete: bool,
pub api_provider: Option<String>,
/// Captured function calls from Google's response.
pub function_calls: Vec<CapturedFunctionCall>,
}
impl StreamingAccumulator {
@@ -96,6 +98,24 @@ impl StreamingAccumulator {
self.thinking_text.push_str(text);
}
}
// Detect functionCall from Google (tool call response)
else if let Some(fc) = part.get("functionCall") {
let name = fc["name"].as_str().unwrap_or("unknown").to_string();
let args = fc["args"].clone();
info!(
tool_name = %name,
tool_args = %args,
"MITM: Google returned functionCall!"
);
self.function_calls.push(CapturedFunctionCall {
name,
args,
captured_at: std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs(),
});
}
// Capture non-thinking response text (skip thoughtSignature parts)
else if part.get("thoughtSignature").is_none() {
if let Some(text) = part["text"].as_str() {
@@ -112,6 +132,10 @@ impl StreamingAccumulator {
if reason == "STOP" {
self.is_complete = true;
}
// Log non-STOP finish reasons
if reason != "STOP" {
info!(finish_reason = reason, "MITM: non-STOP finish reason");
}
}
}
}