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:
@@ -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");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user