From 88ee2b25be039a401574781a58d60ec7ea792310 Mon Sep 17 00:00:00 2001 From: dawn <90008@klbr.net> Date: Wed, 24 Jun 2026 20:34:09 +0300 Subject: [PATCH] fix cache hits add more metrics etc --- .gitignore | 2 + Cargo.toml | 3 + antigravity-bridge/src/auth.rs | 7 ++ antigravity-bridge/src/main.rs | 1 + antigravity-bridge/src/mapping.rs | 63 +++++++++++---- antigravity-bridge/src/routes.rs | 130 ++++++++++++++++++++++++++---- klbr-core/src/agent.rs | 8 +- klbr-core/src/context.rs | 33 +++++++- klbr-core/src/lib.rs | 6 ++ klbr-core/src/models.rs | 67 ++++++++++++++- klbr-daemon/src/daemon.rs | 12 +++ klbr-discord/src/lib.rs | 7 +- klbr-ipc/src/lib.rs | 12 +++ klbr-web/src/App.svelte | 33 ++++++++ klbr-web/src/lib/protocol.ts | 2 +- scripts/capture_agy.sh | 31 +++++++ scripts/capture_tool_call.sh | 29 +++++++ scripts/dump_sse.py | 35 ++++++++ scripts/inspect_role.py | 38 +++++++++ scripts/read_flow.py | 32 ++++++++ scripts/send_msg.py | 39 +++++++++ scripts/show_stream_generate.py | 36 +++++++++ 22 files changed, 582 insertions(+), 44 deletions(-) create mode 100755 scripts/capture_agy.sh create mode 100755 scripts/capture_tool_call.sh create mode 100644 scripts/dump_sse.py create mode 100644 scripts/inspect_role.py create mode 100644 scripts/read_flow.py create mode 100644 scripts/send_msg.py create mode 100644 scripts/show_stream_generate.py diff --git a/.gitignore b/.gitignore index 05e0df7..6b339bc 100644 --- a/.gitignore +++ b/.gitignore @@ -7,3 +7,5 @@ !/benchmarks/runs/README.md klbr.json klbr.kdl +test_agent.db +test_klbr.kdl diff --git a/Cargo.toml b/Cargo.toml index c4e98fa..0430663 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -8,3 +8,6 @@ members = [ "antigravity-bridge", ] resolver = "2" + +[profile.dev] +opt-level = 2 diff --git a/antigravity-bridge/src/auth.rs b/antigravity-bridge/src/auth.rs index 6b9ad3d..679d964 100644 --- a/antigravity-bridge/src/auth.rs +++ b/antigravity-bridge/src/auth.rs @@ -4,6 +4,12 @@ use std::time::{SystemTime, UNIX_EPOCH}; pub const BINARY_PATH: &str = "/home/mayer/.local/bin/agy"; +#[derive(Clone)] +pub struct CachedSessionHashes { + pub trajectory_id: String, + pub hashes: Vec<(serde_json::Value, String)>, +} + pub struct BridgeState { pub client_id: String, pub client_secret: String, @@ -12,6 +18,7 @@ pub struct BridgeState { pub project_id: Option, pub thought_signature_cache: HashMap, pub sessions_cache: HashMap, + pub hash_cache: Vec, } pub type SharedState = std::sync::Arc>; diff --git a/antigravity-bridge/src/main.rs b/antigravity-bridge/src/main.rs index cd9e9ed..4ed9cc4 100644 --- a/antigravity-bridge/src/main.rs +++ b/antigravity-bridge/src/main.rs @@ -54,6 +54,7 @@ async fn main() -> anyhow::Result<()> { project_id: None, thought_signature_cache: auth::load_signatures(), sessions_cache: auth::load_sessions(), + hash_cache: Vec::new(), })); use axum::extract::DefaultBodyLimit; diff --git a/antigravity-bridge/src/mapping.rs b/antigravity-bridge/src/mapping.rs index 0cce1cb..fc91b5f 100644 --- a/antigravity-bridge/src/mapping.rs +++ b/antigravity-bridge/src/mapping.rs @@ -193,7 +193,8 @@ pub fn map_openai_messages_to_gemini( let mut parts = vec![serde_json::json!({ "functionResponse": { "name": fn_name, - "response": response_obj + "response": response_obj, + "id": tc_id } })]; @@ -222,7 +223,7 @@ pub fn map_openai_messages_to_gemini( } contents.push(serde_json::json!({ - "role": "user", + "role": "model", "parts": parts })); } @@ -253,7 +254,8 @@ pub fn map_openai_messages_to_gemini( parts.push(serde_json::json!({ "functionCall": { "name": func_name, - "args": args_val + "args": args_val, + "id": tc_id }, "thoughtSignature": signature })); @@ -287,10 +289,19 @@ pub fn map_openai_messages_to_gemini( "image_url" => { if let Some(url) = item["image_url"]["url"].as_str() { if let Some((mime, data)) = parse_data_url(url) { - let cleaned_data: String = data.chars().filter(|c| !c.is_whitespace()).collect(); - use base64::{engine::general_purpose, Engine as _}; - let final_mime = if let Ok(decoded) = general_purpose::STANDARD.decode(cleaned_data.as_bytes()) { - sniff_or_validate_mime(&mime, &decoded).unwrap_or(mime) + let cleaned_data: String = data + .chars() + .filter(|c| !c.is_whitespace()) + .collect(); + use base64::{ + engine::general_purpose, Engine as _, + }; + let final_mime = if let Ok(decoded) = + general_purpose::STANDARD + .decode(cleaned_data.as_bytes()) + { + sniff_or_validate_mime(&mime, &decoded) + .unwrap_or(mime) } else { mime }; @@ -315,10 +326,21 @@ pub fn map_openai_messages_to_gemini( .and_then(|t| t.as_str()), source.get("data").and_then(|d| d.as_str()), ) { - let cleaned_data: String = data.chars().filter(|c| !c.is_whitespace()).collect(); - use base64::{engine::general_purpose, Engine as _}; - let final_mime = if let Ok(decoded) = general_purpose::STANDARD.decode(cleaned_data.as_bytes()) { - sniff_or_validate_mime(media_type, &decoded).unwrap_or_else(|| media_type.to_string()) + let cleaned_data: String = data + .chars() + .filter(|c| !c.is_whitespace()) + .collect(); + use base64::{ + engine::general_purpose, Engine as _, + }; + let final_mime = if let Ok(decoded) = + general_purpose::STANDARD + .decode(cleaned_data.as_bytes()) + { + sniff_or_validate_mime(media_type, &decoded) + .unwrap_or_else(|| { + media_type.to_string() + }) } else { media_type.to_string() }; @@ -335,10 +357,19 @@ pub fn map_openai_messages_to_gemini( let url_key = format!("{}_url", item_type); if let Some(url) = item[&url_key]["url"].as_str() { if let Some((mime, data)) = parse_data_url(url) { - let cleaned_data: String = data.chars().filter(|c| !c.is_whitespace()).collect(); - use base64::{engine::general_purpose, Engine as _}; - let final_mime = if let Ok(decoded) = general_purpose::STANDARD.decode(cleaned_data.as_bytes()) { - sniff_or_validate_mime(&mime, &decoded).unwrap_or(mime) + let cleaned_data: String = data + .chars() + .filter(|c| !c.is_whitespace()) + .collect(); + use base64::{ + engine::general_purpose, Engine as _, + }; + let final_mime = if let Ok(decoded) = + general_purpose::STANDARD + .decode(cleaned_data.as_bytes()) + { + sniff_or_validate_mime(&mime, &decoded) + .unwrap_or(mime) } else { mime }; @@ -611,7 +642,7 @@ mod tests { let (contents, _) = map_openai_messages_to_gemini(&messages, &HashMap::new()); assert_eq!(contents.len(), 2); let tool_response_msg = &contents[1]; - assert_eq!(tool_response_msg["role"], "user"); + assert_eq!(tool_response_msg["role"], "model"); let parts = tool_response_msg["parts"].as_array().unwrap(); assert_eq!(parts.len(), 2); assert_eq!(parts[0]["functionResponse"]["name"], "read_media"); diff --git a/antigravity-bridge/src/routes.rs b/antigravity-bridge/src/routes.rs index f56fce7..0c5ebb2 100644 --- a/antigravity-bridge/src/routes.rs +++ b/antigravity-bridge/src/routes.rs @@ -51,32 +51,81 @@ impl LineDecoder { } } -fn hash_messages_incremental(messages: &[Value]) -> Vec { +fn hash_messages_incremental_cached( + messages: &[Value], + hash_cache: &mut Vec, +) -> (Vec, Vec<(Value, String)>) { use sha2::{Digest, Sha256}; + + let mut best_match_index = None; + let mut best_match_len = 0; + + for (cache_idx, cached) in hash_cache.iter().enumerate() { + let mut match_len = 0; + for (i, msg) in messages.iter().enumerate() { + if i < cached.hashes.len() && cached.hashes[i].0 == *msg { + match_len += 1; + } else { + break; + } + } + if match_len > best_match_len { + best_match_len = match_len; + best_match_index = Some(cache_idx); + } + } + let mut hashes = Vec::with_capacity(messages.len()); let mut current_hash = [0u8; 32]; - - for msg in messages { + let mut new_cache_hashes = Vec::with_capacity(messages.len()); + + if let Some(cache_idx) = best_match_index { + let cached = &hash_cache[cache_idx]; + for i in 0..best_match_len { + let hash_str = &cached.hashes[i].1; + hashes.push(hash_str.clone()); + new_cache_hashes.push((messages[i].clone(), hash_str.clone())); + } + if best_match_len > 0 { + let last_hex = &cached.hashes[best_match_len - 1].1; + if let Ok(bytes) = hex::decode(last_hex) { + if bytes.len() == 32 { + current_hash.copy_from_slice(&bytes); + } + } + } + } + + for i in best_match_len..messages.len() { + let msg = &messages[i]; let msg_json = serde_json::to_string(msg).unwrap_or_default(); let mut msg_hasher = Sha256::new(); msg_hasher.update(msg_json.as_bytes()); let msg_hash = msg_hasher.finalize(); - + let mut chain_hasher = Sha256::new(); chain_hasher.update(¤t_hash); chain_hasher.update(&msg_hash); let final_hash = chain_hasher.finalize(); current_hash.copy_from_slice(&final_hash); - - hashes.push(hex::encode(final_hash)); + + let hash_str = hex::encode(final_hash); + hashes.push(hash_str.clone()); + new_cache_hashes.push((msg.clone(), hash_str)); + } + + if let Some(cache_idx) = best_match_index { + hash_cache.remove(cache_idx); } - hashes + + (hashes, new_cache_hashes) } pub async fn chat_completions( State(state_arc): State, Json(body): Json, ) -> impl IntoResponse { + let start_bridge = std::time::Instant::now(); let (token, project) = match get_valid_token_and_project(&state_arc).await { Ok(vals) => vals, Err(e) => { @@ -87,10 +136,12 @@ pub async fn chat_completions( .into_response(); } }; + let t_auth = start_bridge.elapsed(); let model_name = body.model.as_deref().unwrap_or("gemini-3.5-flash"); let model_id = map_model_name(model_name, body.thinking_level.as_deref()); + let t_start_mapping = std::time::Instant::now(); let signature_cache = { let state = state_arc.read().await; state.thought_signature_cache.clone() @@ -98,6 +149,7 @@ pub async fn chat_completions( let messages_val = serde_json::json!(body.messages); let (contents, system_instruction) = map_openai_messages_to_gemini(&messages_val, &signature_cache); + let t_mapping = t_start_mapping.elapsed(); let mut budget = 0; if let Some(ref thinking_level) = body.thinking_level { @@ -116,8 +168,12 @@ pub async fn chat_completions( } } + let t_start_hashing = std::time::Instant::now(); let messages_arr = body.messages.as_slice(); - let prefix_hashes = hash_messages_incremental(messages_arr); + let (prefix_hashes, new_cache_hashes) = { + let mut state = state_arc.write().await; + hash_messages_incremental_cached(messages_arr, &mut state.hash_cache) + }; let mut traj_id = None; let mut sess_id = None; @@ -139,16 +195,44 @@ pub async fn chat_completions( let trajectory_id = traj_id.unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); let session_id = sess_id.unwrap_or_else(|| rand::Rng::gen::(&mut rand::thread_rng()).to_string()); + let t_hashing = t_start_hashing.elapsed(); + let t_start_disk = std::time::Instant::now(); { let mut state = state_arc.write().await; + + state.hash_cache.insert( + 0, + crate::auth::CachedSessionHashes { + trajectory_id: trajectory_id.clone(), + hashes: new_cache_hashes, + }, + ); + state.hash_cache.truncate(10); + if let Some(current_hash) = prefix_hashes.last() { - state - .sessions_cache - .insert(current_hash.clone(), (trajectory_id.clone(), session_id.clone())); - crate::auth::save_sessions(&state.sessions_cache); + state.sessions_cache.insert( + current_hash.clone(), + (trajectory_id.clone(), session_id.clone()), + ); + let sessions_clone = state.sessions_cache.clone(); + tokio::spawn(async move { + crate::auth::save_sessions(&sessions_clone); + }); } } + let t_disk = t_start_disk.elapsed(); + + let total_setup = start_bridge.elapsed(); + let setup_ms = total_setup.as_millis() as u64; + tracing::info!( + "bridge setup finished in {}ms (auth: {}ms, mapping: {}ms, hashing: {}ms, disk_io: {}ms)", + setup_ms, + t_auth.as_millis(), + t_mapping.as_millis(), + t_hashing.as_millis(), + t_disk.as_millis() + ); let mut payload = serde_json::json!({ "project": project, @@ -282,6 +366,7 @@ pub async fn chat_completions( // Usage metadata if let Some(usage_meta) = response_part.get("usageMetadata") { + tracing::info!("USAGE METADATA: {}", usage_meta); let out = serde_json::json!({ "id": request_id, "object": "chat.completion.chunk", @@ -291,7 +376,8 @@ pub async fn chat_completions( "usage": { "prompt_tokens": usage_meta.get("promptTokenCount").unwrap_or(&serde_json::json!(0)), "completion_tokens": usage_meta.get("candidatesTokenCount").unwrap_or(&serde_json::json!(0)), - "total_tokens": usage_meta.get("totalTokenCount").unwrap_or(&serde_json::json!(0)) + "total_tokens": usage_meta.get("totalTokenCount").unwrap_or(&serde_json::json!(0)), + "cached_prompt_tokens": usage_meta.get("cachedContentTokenCount").unwrap_or(&serde_json::json!(0)) } }); yield Ok(Event::default().data(out.to_string())); @@ -333,14 +419,21 @@ pub async fn chat_completions( args_val.as_str().unwrap_or("{}").to_string() }; - let tc_id = format!("call_{}", &uuid::Uuid::new_v4().simple().to_string()[..12]); + let tc_id = f["id"].as_str() + .map(|s| s.to_string()) + .unwrap_or_else(|| { + format!("call_{}", &uuid::Uuid::new_v4().simple().to_string()[..12]) + }); let thought_sig = part.get("thoughtSignature") .or_else(|| part.get("thought_signature")); if let Some(ts_val) = thought_sig { if let Some(ts_str) = ts_val.as_str() { let mut state = state_arc.write().await; state.thought_signature_cache.insert(tc_id.clone(), ts_str.to_string()); - crate::auth::save_signatures(&state.thought_signature_cache); + let sigs_clone = state.thought_signature_cache.clone(); + tokio::spawn(async move { + crate::auth::save_signatures(&sigs_clone); + }); } } @@ -416,7 +509,12 @@ pub async fn chat_completions( yield Ok(Event::default().data("[DONE]")); }; - Sse::new(stream).into_response() + let mut response = Sse::new(stream).into_response(); + response.headers_mut().insert( + "x-bridge-setup-ms", + axum::http::HeaderValue::from_str(&setup_ms.to_string()).unwrap(), + ); + response } pub async fn list_models() -> impl IntoResponse { diff --git a/klbr-core/src/agent.rs b/klbr-core/src/agent.rs index c8776f1..aacaf4d 100644 --- a/klbr-core/src/agent.rs +++ b/klbr-core/src/agent.rs @@ -745,7 +745,7 @@ impl Agent { let _ = self.output.send(AgentEvent::ScratchToken(tok)); } Some(LlmEvent::Usage(usage)) => { - ctx.update_tokens(usage.total_tokens); + ctx.update_tokens(&usage); } Some(LlmEvent::ToolCalls(calls)) => { tool_calls = calls; @@ -1195,6 +1195,12 @@ impl Agent { turn_count: *turn_count, context_tokens: ctx.total_tokens, watermark: self.config.watermark_tokens, + prompt_tokens: ctx.prompt_tokens, + completion_tokens: ctx.completion_tokens, + cached_prompt_tokens: ctx.cached_prompt_tokens, + prompt_processing_ms: ctx.prompt_processing_ms, + generation_ms: ctx.generation_ms, + bridge_processing_ms: ctx.bridge_processing_ms, }; *self.snapshot.write().await = Some(metrics.clone()); let _ = self.output.send(AgentEvent::Metrics(metrics)); diff --git a/klbr-core/src/context.rs b/klbr-core/src/context.rs index 546db07..07a3795 100644 --- a/klbr-core/src/context.rs +++ b/klbr-core/src/context.rs @@ -23,6 +23,12 @@ pub struct Context { /// rolling conversation turns pub turns: Vec, pub total_tokens: usize, + pub prompt_tokens: usize, + pub completion_tokens: usize, + pub cached_prompt_tokens: usize, + pub prompt_processing_ms: u64, + pub generation_ms: u64, + pub bridge_processing_ms: u64, passive_recall_log: Vec<(usize, i64)>, } @@ -40,6 +46,12 @@ impl Context { system: vec![Message::system(system_content)], turns: vec![], total_tokens: 0, + prompt_tokens: 0, + completion_tokens: 0, + cached_prompt_tokens: 0, + prompt_processing_ms: 0, + generation_ms: 0, + bridge_processing_ms: 0, passive_recall_log: Vec::new(), } } @@ -147,6 +159,11 @@ impl Context { pub fn clear(&mut self) { self.turns.clear(); self.total_tokens = 0; + self.prompt_tokens = 0; + self.completion_tokens = 0; + self.cached_prompt_tokens = 0; + self.prompt_processing_ms = 0; + self.generation_ms = 0; self.passive_recall_log.clear(); } @@ -255,12 +272,24 @@ impl Context { return vec![]; } self.total_tokens = 0; + self.prompt_tokens = 0; + self.completion_tokens = 0; + self.cached_prompt_tokens = 0; + self.prompt_processing_ms = 0; + self.generation_ms = 0; + self.bridge_processing_ms = 0; self.drain_passive_recall_log(safe_cut); self.turns.drain(..safe_cut).collect() } - pub fn update_tokens(&mut self, tokens: usize) { - self.total_tokens = tokens; + pub fn update_tokens(&mut self, usage: &crate::models::Usage) { + self.total_tokens = usage.total_tokens; + self.prompt_tokens = usage.prompt_tokens; + self.completion_tokens = usage.completion_tokens; + self.cached_prompt_tokens = usage.cached_prompt_tokens; + self.prompt_processing_ms = usage.prompt_processing_ms; + self.generation_ms = usage.generation_ms; + self.bridge_processing_ms = usage.bridge_processing_ms; } pub fn passively_recalled_ids(&self) -> HashSet { diff --git a/klbr-core/src/lib.rs b/klbr-core/src/lib.rs index 0ebaec2..903380c 100644 --- a/klbr-core/src/lib.rs +++ b/klbr-core/src/lib.rs @@ -22,6 +22,12 @@ pub struct AgentMetrics { pub turn_count: usize, pub context_tokens: usize, pub watermark: usize, + pub prompt_tokens: usize, + pub completion_tokens: usize, + pub cached_prompt_tokens: usize, + pub prompt_processing_ms: u64, + pub generation_ms: u64, + pub bridge_processing_ms: u64, } #[derive(Debug, Clone, Serialize, Deserialize)] diff --git a/klbr-core/src/models.rs b/klbr-core/src/models.rs index 77b63e1..7e13363 100644 --- a/klbr-core/src/models.rs +++ b/klbr-core/src/models.rs @@ -157,6 +157,14 @@ pub struct Usage { pub prompt_tokens: usize, pub completion_tokens: usize, pub total_tokens: usize, + #[serde(default)] + pub cached_prompt_tokens: usize, + #[serde(default)] + pub prompt_processing_ms: u64, + #[serde(default)] + pub generation_ms: u64, + #[serde(default)] + pub bridge_processing_ms: u64, } #[derive(Debug, Clone, Serialize, Deserialize)] @@ -482,6 +490,10 @@ impl LlmClient { tool_names ); + let start_time = std::time::Instant::now(); + let mut first_response_time: Option = None; + let mut final_usage: Option = None; + let res = self .with_auth(self.client.post(&endpoint), &self.config.llm) .json(&body) @@ -496,6 +508,12 @@ impl LlmClient { } }; tracing::info!("stream response status: {}", res.status()); + let bridge_setup_ms = res + .headers() + .get("x-bridge-setup-ms") + .and_then(|h| h.to_str().ok()) + .and_then(|s| s.parse::().ok()) + .unwrap_or(0); let res = Self::error_for_status_with_body(res).await?; let mut res = res.bytes_stream(); @@ -550,6 +568,20 @@ impl LlmClient { let calls: Vec = calls.into_iter().map(|(_, c)| c).collect(); let _ = tok_tx.send(LlmEvent::ToolCalls(calls)).await; } + + // Send final usage with latency calculations + let elapsed = start_time.elapsed(); + let mut u = final_usage.clone().unwrap_or_default(); + if let Some(ttfc) = first_response_time { + u.prompt_processing_ms = ttfc.as_millis() as u64; + u.generation_ms = elapsed.saturating_sub(ttfc).as_millis() as u64; + } else { + u.prompt_processing_ms = elapsed.as_millis() as u64; + u.generation_ms = 0; + } + u.bridge_processing_ms = bridge_setup_ms; + let _ = tok_tx.send(LlmEvent::Usage(u)).await; + return Ok(()); } @@ -568,7 +600,16 @@ impl LlmClient { // tracing::info!("got stream chunk: {v:#?}"); if v["usage"].is_object() { - let u: Usage = serde_json::from_value(v["usage"].clone())?; + let mut u: Usage = serde_json::from_value(v["usage"].clone())?; + let elapsed = start_time.elapsed(); + if let Some(ttfc) = first_response_time { + u.prompt_processing_ms = ttfc.as_millis() as u64; + u.generation_ms = elapsed.saturating_sub(ttfc).as_millis() as u64; + } else { + u.prompt_processing_ms = elapsed.as_millis() as u64; + } + u.bridge_processing_ms = bridge_setup_ms; + final_usage = Some(u.clone()); let _ = tok_tx.send(LlmEvent::Usage(u)).await; continue; } @@ -583,6 +624,9 @@ impl LlmClient { // accumulate tool call deltas if let Some(tc_arr) = delta["tool_calls"].as_array() { + if first_response_time.is_none() { + first_response_time = Some(start_time.elapsed()); + } if !tool_call_started_sent { let _ = tok_tx.send(LlmEvent::ToolCallStarted).await; tool_call_started_sent = true; @@ -603,7 +647,6 @@ impl LlmClient { continue; } - // reasoning_content comes before content during thinking if let Some(t) = delta .get("reasoning_content") // this is for openrouter @@ -611,6 +654,9 @@ impl LlmClient { .and_then(|v| v.as_str()) { if !t.is_empty() { + if first_response_time.is_none() { + first_response_time = Some(start_time.elapsed()); + } tracing::debug!("thinking token: {}", t); if tok_tx .send(LlmEvent::ThinkToken(t.to_string())) @@ -623,6 +669,9 @@ impl LlmClient { } if let Some(t) = delta["content"].as_str() { if !t.is_empty() { + if first_response_time.is_none() { + first_response_time = Some(start_time.elapsed()); + } tracing::debug!("content token: {}", t); if tok_tx.send(LlmEvent::Token(t.to_string())).await.is_err() { return Ok(()); @@ -632,6 +681,18 @@ impl LlmClient { } } + let elapsed = start_time.elapsed(); + let mut u = final_usage.unwrap_or_default(); + if let Some(ttfc) = first_response_time { + u.prompt_processing_ms = ttfc.as_millis() as u64; + u.generation_ms = elapsed.saturating_sub(ttfc).as_millis() as u64; + } else { + u.prompt_processing_ms = elapsed.as_millis() as u64; + u.generation_ms = 0; + } + u.bridge_processing_ms = bridge_setup_ms; + let _ = tok_tx.send(LlmEvent::Usage(u)).await; + Ok(()) } @@ -906,8 +967,6 @@ impl LlmClient { } } - - fn replace_media_urls_with_placeholders(text: &str, limit_bytes: usize) -> String { let mut result = String::new(); let mut current_idx = 0; diff --git a/klbr-daemon/src/daemon.rs b/klbr-daemon/src/daemon.rs index 3479b5c..dfebffe 100644 --- a/klbr-daemon/src/daemon.rs +++ b/klbr-daemon/src/daemon.rs @@ -102,6 +102,12 @@ async fn handle( turn_count: metrics.turn_count, context_tokens: metrics.context_tokens, watermark: metrics.watermark, + prompt_tokens: metrics.prompt_tokens, + completion_tokens: metrics.completion_tokens, + cached_prompt_tokens: metrics.cached_prompt_tokens, + prompt_processing_ms: metrics.prompt_processing_ms, + generation_ms: metrics.generation_ms, + bridge_processing_ms: metrics.bridge_processing_ms, }, ) .await?; @@ -242,6 +248,12 @@ async fn handle( turn_count: m.turn_count, context_tokens: m.context_tokens, watermark: m.watermark, + prompt_tokens: m.prompt_tokens, + completion_tokens: m.completion_tokens, + cached_prompt_tokens: m.cached_prompt_tokens, + prompt_processing_ms: m.prompt_processing_ms, + generation_ms: m.generation_ms, + bridge_processing_ms: m.bridge_processing_ms, }, AgentEvent::UserTurn(entry) => ServerMsg::Turn { entry: map_history(entry), diff --git a/klbr-discord/src/lib.rs b/klbr-discord/src/lib.rs index 08b0168..a3e6032 100644 --- a/klbr-discord/src/lib.rs +++ b/klbr-discord/src/lib.rs @@ -719,10 +719,9 @@ impl DiscordRuntime { Some("gif") => "image/gif".to_string(), Some("webp") => "image/webp".to_string(), Some("pdf") => "application/pdf".to_string(), - _ => attachment - .content_type - .clone() - .unwrap_or_else(|| "application/octet-stream".to_string()), + _ => attachment.content_type.clone().unwrap_or_else(|| { + "application/octet-stream".to_string() + }), }; if !final_content.is_empty() { diff --git a/klbr-ipc/src/lib.rs b/klbr-ipc/src/lib.rs index a101aef..4f7816a 100644 --- a/klbr-ipc/src/lib.rs +++ b/klbr-ipc/src/lib.rs @@ -119,6 +119,18 @@ pub enum ServerMsg { turn_count: usize, context_tokens: usize, watermark: usize, + #[serde(default)] + prompt_tokens: usize, + #[serde(default)] + completion_tokens: usize, + #[serde(default)] + cached_prompt_tokens: usize, + #[serde(default)] + prompt_processing_ms: u64, + #[serde(default)] + generation_ms: u64, + #[serde(default)] + bridge_processing_ms: u64, }, History { turns: Vec, diff --git a/klbr-web/src/App.svelte b/klbr-web/src/App.svelte index d11863b..0b89a98 100644 --- a/klbr-web/src/App.svelte +++ b/klbr-web/src/App.svelte @@ -74,6 +74,12 @@ turnCount: number; contextTokens: number; watermark: number; + promptTokens?: number; + completionTokens?: number; + cachedPromptTokens?: number; + promptProcessingMs?: number; + generationMs?: number; + bridgeProcessingMs?: number; } interface ToolCallItem { @@ -821,6 +827,12 @@ turnCount: msg.turn_count, contextTokens: msg.context_tokens, watermark: msg.watermark, + promptTokens: msg.prompt_tokens, + completionTokens: msg.completion_tokens, + cachedPromptTokens: msg.cached_prompt_tokens, + promptProcessingMs: msg.prompt_processing_ms, + generationMs: msg.generation_ms, + bridgeProcessingMs: msg.bridge_processing_ms, }; break; case "history": @@ -1661,9 +1673,30 @@ `${contextRemaining(session).toLocaleString()} tok left`, ); } + if (session.metrics.promptTokens) { + let promptStr = `prompt ${session.metrics.promptTokens.toLocaleString()}`; + if (session.metrics.cachedPromptTokens) { + promptStr += ` (cached ${session.metrics.cachedPromptTokens.toLocaleString()})`; + } + parts.push(promptStr); + } + if (session.metrics.completionTokens) { + parts.push(`completion ${session.metrics.completionTokens.toLocaleString()}`); + } if (session.lastTps) { parts.push(`${session.lastTps.toFixed(1)} tok/s`); } + if (session.metrics.promptProcessingMs) { + parts.push(`ttfc ${(session.metrics.promptProcessingMs / 1000).toFixed(2)}s`); + } + if (session.metrics.generationMs) { + parts.push(`gen ${(session.metrics.generationMs / 1000).toFixed(2)}s`); + } + if (session.metrics.bridgeProcessingMs) { + const ms = session.metrics.bridgeProcessingMs; + const str = ms < 1000 ? `${ms}ms` : `${(ms / 1000).toFixed(2)}s`; + parts.push(`bridge ${str}`); + } if (!parts.length) { parts.push("no metrics yet"); } diff --git a/klbr-web/src/lib/protocol.ts b/klbr-web/src/lib/protocol.ts index 51329cd..98c4d6a 100644 --- a/klbr-web/src/lib/protocol.ts +++ b/klbr-web/src/lib/protocol.ts @@ -67,7 +67,7 @@ export type ServerMsg = | { type: "error"; content: string } | { type: "status"; content: string } | { type: "system_nudge"; content: string } - | { type: "metrics"; turn_count: number; context_tokens: number; watermark: number } + | { type: "metrics"; turn_count: number; context_tokens: number; watermark: number; prompt_tokens?: number; completion_tokens?: number; cached_prompt_tokens?: number; prompt_processing_ms?: number; generation_ms?: number; bridge_processing_ms?: number } | { type: "history"; turns: HistoryEntry[] } | { type: "turn"; entry: HistoryEntry } | { type: "external_event"; source: string; conversation_id: string; content: string } diff --git a/scripts/capture_agy.sh b/scripts/capture_agy.sh new file mode 100755 index 0000000..9491c2e --- /dev/null +++ b/scripts/capture_agy.sh @@ -0,0 +1,31 @@ +#!/usr/bin/env bash +set -e + +# Start mitmdump in the background on port 8080, writing flows to traffic.flow +mitmdump -p 8080 -w traffic.flow > mitmdump.log 2>&1 & +MITM_PID=$! + +# Cleanup function to ensure mitmdump is killed even if script fails +cleanup() { + echo "Killing mitmdump (PID: $MITM_PID)..." + kill $MITM_PID || true +} +trap cleanup EXIT + +# Wait for mitmdump to start and generate certificates +echo "Waiting for mitmdump to start..." +sleep 3 + +# Export proxy env vars +export HTTPS_PROXY=http://127.0.0.1:8080 +export HTTP_PROXY=http://127.0.0.1:8080 +export SSL_CERT_FILE=$HOME/.mitmproxy/mitmproxy-ca-cert.pem +export NIX_SSL_CERT_FILE=$HOME/.mitmproxy/mitmproxy-ca-cert.pem + +echo "Running first prompt..." +~/.local/bin/agy -p "im puppy woof" + +echo "Running second prompt (continue)..." +~/.local/bin/agy -c -p "idk just chatting ig :3" + +echo "Done! Captured flows are in traffic.flow." diff --git a/scripts/capture_tool_call.sh b/scripts/capture_tool_call.sh new file mode 100755 index 0000000..e2eca77 --- /dev/null +++ b/scripts/capture_tool_call.sh @@ -0,0 +1,29 @@ +#!/usr/bin/env bash +set -e + +# Start mitmdump in the background on port 8080, writing flows to traffic.flow +nix-shell -p mitmproxy --run "mitmdump -p 8080 -w traffic.flow" > mitmdump.log 2>&1 & +MITM_PID=$! + +# Cleanup function to ensure mitmdump is killed even if script fails +cleanup() { + echo "Killing mitmdump (PID: $MITM_PID)..." + # Find child process since nix-shell spawns a child + pkill -P $MITM_PID || true + kill $MITM_PID || true +} +trap cleanup EXIT + +echo "Waiting for mitmdump to start..." +sleep 5 + +# Export proxy env vars +export HTTPS_PROXY=http://127.0.0.1:8080 +export HTTP_PROXY=http://127.0.0.1:8080 +export SSL_CERT_FILE=$HOME/.mitmproxy/mitmproxy-ca-cert.pem +export NIX_SSL_CERT_FILE=$HOME/.mitmproxy/mitmproxy-ca-cert.pem + +echo "Running agy to read a file and trigger a tool call..." +~/.local/bin/agy --dangerously-skip-permissions -p "please read docs/README.md and summarize it in one sentence." + +echo "Done! Captured flows are in traffic.flow." diff --git a/scripts/dump_sse.py b/scripts/dump_sse.py new file mode 100644 index 0000000..76ffa15 --- /dev/null +++ b/scripts/dump_sse.py @@ -0,0 +1,35 @@ +from mitmproxy import io +import json +import sys + +with open("traffic.flow", "rb") as logfile: + freader = io.FlowReader(logfile) + for flow in freader.stream(): + if "daily-cloudcode-pa.googleapis.com" in flow.request.host: + if "streamGenerateContent" in flow.request.path: + print(f"\n==========================================") + print(f"URL: {flow.request.url}") + print(f"Request Body: {flow.request.content.decode('utf-8', errors='ignore')[:300]}") + if flow.response: + content = flow.response.content.decode("utf-8", errors="ignore") + print(f"Total Response Length: {len(content)}") + print("First 5 lines:") + for line in content.splitlines()[:5]: + print(line[:300]) + print("Last 5 lines:") + for line in content.splitlines()[-5:]: + print(line[:300]) + + # Find all lines starting with "data: " and print them + print("\nAll data lines containing functionCall or thoughtSignature:") + for line in content.splitlines(): + if line.startswith("data: "): + data_str = line[len("data: "):] + try: + val = json.loads(data_str) + parts = val.get("response", {}).get("candidates", [{}])[0].get("content", {}).get("parts", []) + for p in parts: + if "functionCall" in p or "thoughtSignature" in p: + print(f"PART: {json.dumps(p, indent=2)}") + except Exception as e: + pass diff --git a/scripts/inspect_role.py b/scripts/inspect_role.py new file mode 100644 index 0000000..6320afa --- /dev/null +++ b/scripts/inspect_role.py @@ -0,0 +1,38 @@ +import json +from mitmproxy import io +from mitmproxy.exceptions import FlowReadException + +def main(): + flow_file = "traffic.flow" + with open(flow_file, "rb") as logfile: + freader = io.FlowReader(logfile) + try: + for flow in freader.stream(): + if "streamGenerateContent" in flow.request.url: + try: + req_data = json.loads(flow.request.content) + contents = req_data.get("request", {}).get("contents", []) + print("="*80) + print(f"URL: {flow.request.url}") + print(f"Total messages in contents: {len(contents)}") + for idx, msg in enumerate(contents): + role = msg.get("role") + parts = msg.get("parts", []) + part_types = [] + for p in parts: + if "text" in p: + part_types.append("text") + elif "functionCall" in p: + part_types.append(f"functionCall({p['functionCall']['name']})") + elif "functionResponse" in p: + part_types.append(f"functionResponse({p['functionResponse']['name']})") + else: + part_types.append(str(list(p.keys()))) + print(f" Message {idx}: role={role}, parts={part_types}") + except Exception as e: + print("Failed to parse flow:", e) + except FlowReadException as e: + print(f"Flow file reading failed: {e}") + +if __name__ == "__main__": + main() diff --git a/scripts/read_flow.py b/scripts/read_flow.py new file mode 100644 index 0000000..d51e968 --- /dev/null +++ b/scripts/read_flow.py @@ -0,0 +1,32 @@ +from mitmproxy import io +from mitmproxy.exceptions import FlowReadException +import json +import sys + +with open("traffic.flow", "rb") as logfile: + freader = io.FlowReader(logfile) + try: + for flow in freader.stream(): + if "daily-cloudcode-pa.googleapis.com" in flow.request.host: + path = flow.request.path + if "streamGenerateContent" in path or "lo" in path or "se" in path or "fe" in path: + print(f"\n==========================================") + print(f"REQUEST: {flow.request.method} {flow.request.url}") + headers = dict(flow.request.headers) + print(f"Headers: {json.dumps(headers, indent=2)}") + try: + req_json = json.loads(flow.request.content.decode("utf-8", errors="ignore")) + print(f"Body: {json.dumps(req_json, indent=2)}") + except Exception: + print(f"Body (raw): {flow.request.content.decode('utf-8', errors='ignore')}") + + print(f"\nRESPONSE STATUS: {flow.response.status_code if flow.response else 'N/A'}") + if flow.response: + try: + resp_json = json.loads(flow.response.content.decode("utf-8", errors="ignore")) + print(f"Response Body: {json.dumps(resp_json, indent=2)}") + except Exception: + # Might be SSE stream, print raw or truncated + print(f"Response Body (raw):\n{flow.response.content.decode('utf-8', errors='ignore')[:2000]}") + except FlowReadException as e: + print(f"Flow file reading error: {e}") diff --git a/scripts/send_msg.py b/scripts/send_msg.py new file mode 100644 index 0000000..7390ad7 --- /dev/null +++ b/scripts/send_msg.py @@ -0,0 +1,39 @@ +import asyncio +import websockets +import json +import sys + +async def send(): + uri = "ws://127.0.0.1:8765" + # Set max message size to 50MB + async with websockets.connect(uri, max_size=50 * 1024 * 1024) as websocket: + print("Connected. Waiting for initial server messages...") + while True: + msg = await websocket.recv() + print(f"Received: {msg[:100]}...") + try: + msg_data = json.loads(msg) + if msg_data.get("type") == "history": + break + except Exception: + pass + + print("\nSending message to daemon...") + payload = { + "type": "message", + "source": "tui", + "content": "hey there, just a test message" + } + await websocket.send(json.dumps(payload)) + + print("Waiting for response events...") + try: + while True: + msg = await websocket.recv() + print(f"Agent Event: {msg[:200]}") + if "done" in msg: + break + except Exception as e: + print(f"Error: {e}") + +asyncio.run(send()) diff --git a/scripts/show_stream_generate.py b/scripts/show_stream_generate.py new file mode 100644 index 0000000..c9fc03a --- /dev/null +++ b/scripts/show_stream_generate.py @@ -0,0 +1,36 @@ +import json +import sys +from mitmproxy import io +from mitmproxy.exceptions import FlowReadException + +def main(): + flow_file = "traffic.flow" + print(f"Reading {flow_file}...") + with open(flow_file, "rb") as logfile: + freader = io.FlowReader(logfile) + try: + for flow in freader.stream(): + if "streamGenerateContent" in flow.request.url: + print("="*80) + print(f"URL: {flow.request.url}") + print("REQUEST JSON:") + try: + req_data = json.loads(flow.request.content) + print(json.dumps(req_data, indent=2)) + except Exception as e: + print("Failed to parse request content:", e) + print(flow.request.content) + print("-"*80) + print("RESPONSE SNIPPET:") + try: + resp_text = flow.response.content.decode("utf-8", errors="ignore") + print(resp_text[:3000]) + if len(resp_text) > 3000: + print("... (truncated)") + except Exception as e: + print("Failed to decode response content:", e) + except FlowReadException as e: + print(f"Flow file reading failed: {e}") + +if __name__ == "__main__": + main() -- 2.51.2