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

feat: initial commit — antigravity proxy with MITM, standalone LS, and snapshot tooling

Browse Source
This commit is contained in:
Nikketryhard
2026-02-14 02:24:35 -06:00
commit d5e7f09225
30 changed files with 9980 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

271
src/mitm/intercept.rs Normal file
View File

@@ -0,0 +1,271 @@
//! API response interceptor: parses Anthropic/Google API responses to extract usage data.
//!
//! Handles both streaming (SSE) and non-streaming (JSON) responses.
use super::store::ApiUsage;
use serde_json::Value;
use tracing::{debug, trace};
/// Parse a complete (non-streaming) Anthropic Messages API response body.
///
/// Response format:
/// ```json
/// {
/// "id": "msg_...",
/// "type": "message",
/// "model": "claude-sonnet-4-20250514",
/// "usage": {
/// "input_tokens": 1234,
/// "output_tokens": 567,
/// "cache_creation_input_tokens": 0,
/// "cache_read_input_tokens": 890
/// },
/// "stop_reason": "end_turn"
/// }
/// ```
pub fn parse_non_streaming_response(body: &[u8]) -> Option<ApiUsage> {
let json: Value = serde_json::from_slice(body).ok()?;
extract_usage_from_message(&json)
}
/// Parse SSE events from a streaming Anthropic response body chunk.
///
/// Events of interest:
/// - `message_start` — contains `message.usage.input_tokens` + cache tokens
/// - `message_delta` — contains `usage.output_tokens`
/// - `message_stop` — marks end (no usage data)
///
/// Returns accumulated usage across all events in this chunk.
pub fn parse_streaming_chunk(chunk: &str, accumulator: &mut StreamingAccumulator) {
for line in chunk.lines() {
if let Some(data) = line.strip_prefix("data: ") {
if data.trim() == "[DONE]" {
continue;
}
if let Ok(event) = serde_json::from_str::<Value>(data) {
accumulator.process_event(&event);
}
}
}
}
/// Accumulates usage data across streaming SSE events.
#[derive(Debug, Default)]
pub struct StreamingAccumulator {
pub input_tokens: u64,
pub output_tokens: u64,
pub cache_creation_input_tokens: u64,
pub cache_read_input_tokens: u64,
pub model: Option<String>,
pub stop_reason: Option<String>,
pub is_complete: bool,
}
impl StreamingAccumulator {
pub fn new() -> Self {
Self::default()
}
/// Process a single SSE event.
pub fn process_event(&mut self, event: &Value) {
let event_type = event["type"].as_str().unwrap_or("");
match event_type {
"message_start" => {
// message_start contains the initial usage (input tokens + cache)
if let Some(usage) = event.get("message").and_then(|m| m.get("usage")) {
self.input_tokens = usage["input_tokens"].as_u64().unwrap_or(0);
self.cache_creation_input_tokens = usage["cache_creation_input_tokens"].as_u64().unwrap_or(0);
self.cache_read_input_tokens = usage["cache_read_input_tokens"].as_u64().unwrap_or(0);
}
if let Some(model) = event.get("message").and_then(|m| m["model"].as_str()) {
self.model = Some(model.to_string());
}
trace!(
input = self.input_tokens,
cache_read = self.cache_read_input_tokens,
cache_create = self.cache_creation_input_tokens,
"SSE message_start: captured input usage"
);
}
"message_delta" => {
// message_delta contains the output usage
if let Some(usage) = event.get("usage") {
self.output_tokens = usage["output_tokens"].as_u64().unwrap_or(self.output_tokens);
}
if let Some(reason) = event["delta"]["stop_reason"].as_str() {
self.stop_reason = Some(reason.to_string());
}
trace!(output = self.output_tokens, "SSE message_delta: updated output tokens");
}
"message_stop" => {
self.is_complete = true;
debug!(
input = self.input_tokens,
output = self.output_tokens,
cache_read = self.cache_read_input_tokens,
model = ?self.model,
"SSE message_stop: stream complete"
);
}
"content_block_start" | "content_block_delta" | "content_block_stop" | "ping" => {
// Content events — no usage data, just pass through
}
_ => {
trace!(event_type, "SSE: unknown event type");
}
}
}
/// Convert accumulated data to an ApiUsage.
pub fn into_usage(self) -> ApiUsage {
ApiUsage {
input_tokens: self.input_tokens,
output_tokens: self.output_tokens,
cache_creation_input_tokens: self.cache_creation_input_tokens,
cache_read_input_tokens: self.cache_read_input_tokens,
thinking_output_tokens: 0,
response_output_tokens: 0,
total_cost_usd: None,
model: self.model,
stop_reason: self.stop_reason,
api_provider: Some("anthropic".to_string()),
grpc_method: None,
captured_at: std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs(),
}
}
}
/// Extract usage from a complete Message JSON object.
fn extract_usage_from_message(msg: &Value) -> Option<ApiUsage> {
let usage = msg.get("usage")?;
Some(ApiUsage {
input_tokens: usage["input_tokens"].as_u64().unwrap_or(0),
output_tokens: usage["output_tokens"].as_u64().unwrap_or(0),
cache_creation_input_tokens: usage["cache_creation_input_tokens"].as_u64().unwrap_or(0),
cache_read_input_tokens: usage["cache_read_input_tokens"].as_u64().unwrap_or(0),
thinking_output_tokens: 0,
response_output_tokens: 0,
total_cost_usd: None,
model: msg["model"].as_str().map(|s| s.to_string()),
stop_reason: msg["stop_reason"].as_str().map(|s| s.to_string()),
api_provider: Some("anthropic".to_string()),
grpc_method: None,
captured_at: std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs(),
})
}
/// Try to identify a cascade ID from the request body.
///
/// The LS includes cascade-related metadata in its API requests (as part of
/// the system prompt or metadata field). We try to find it.
pub fn extract_cascade_hint(request_body: &[u8]) -> Option<String> {
let json: Value = serde_json::from_slice(request_body).ok()?;
// Check for metadata field (some API configurations include it)
if let Some(metadata) = json.get("metadata") {
if let Some(user_id) = metadata["user_id"].as_str() {
// The LS often sets user_id to the cascadeId
return Some(user_id.to_string());
}
}
// Check system prompt for cascade/workspace markers
if let Some(system) = json.get("system") {
let system_str = match system {
Value::String(s) => s.clone(),
Value::Array(arr) => {
// Array of content blocks
arr.iter()
.filter_map(|b| b["text"].as_str())
.collect::<Vec<_>>()
.join(" ")
}
_ => return None,
};
// Look for workspace_id or cascade_id patterns
if let Some(pos) = system_str.find("workspace_id") {
let rest = &system_str[pos..];
// Extract the value after workspace_id
if let Some(val) = rest.split_whitespace().nth(1) {
return Some(val.to_string());
}
}
}
None
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_parse_non_streaming() {
let body = r#"{
"id": "msg_123",
"type": "message",
"model": "claude-sonnet-4-20250514",
"usage": {
"input_tokens": 100,
"output_tokens": 50,
"cache_creation_input_tokens": 10,
"cache_read_input_tokens": 30
},
"stop_reason": "end_turn"
}"#;
let usage = parse_non_streaming_response(body.as_bytes()).unwrap();
assert_eq!(usage.input_tokens, 100);
assert_eq!(usage.output_tokens, 50);
assert_eq!(usage.cache_creation_input_tokens, 10);
assert_eq!(usage.cache_read_input_tokens, 30);
assert_eq!(usage.model.as_deref(), Some("claude-sonnet-4-20250514"));
}
#[test]
fn test_streaming_accumulator() {
let mut acc = StreamingAccumulator::new();
// message_start
let start = serde_json::json!({
"type": "message_start",
"message": {
"model": "claude-sonnet-4-20250514",
"usage": {
"input_tokens": 200,
"cache_creation_input_tokens": 5,
"cache_read_input_tokens": 50
}
}
});
acc.process_event(&start);
assert_eq!(acc.input_tokens, 200);
assert_eq!(acc.cache_read_input_tokens, 50);
// message_delta
let delta = serde_json::json!({
"type": "message_delta",
"delta": { "stop_reason": "end_turn" },
"usage": { "output_tokens": 75 }
});
acc.process_event(&delta);
assert_eq!(acc.output_tokens, 75);
// message_stop
let stop = serde_json::json!({ "type": "message_stop" });
acc.process_event(&stop);
assert!(acc.is_complete);
let usage = acc.into_usage();
assert_eq!(usage.input_tokens, 200);
assert_eq!(usage.output_tokens, 75);
}
}