diff --git a/src/agentsight/build.rs b/src/agentsight/build.rs index f6b420a64..a8f580326 100644 --- a/src/agentsight/build.rs +++ b/src/agentsight/build.rs @@ -38,11 +38,11 @@ fn main() { generate_skeleton(&mut out, "sslsniff"); generate_header(&mut out, "sslsniff"); - + // Generate proctrace skeleton and bindings generate_skeleton(&mut out, "proctrace"); generate_header(&mut out, "proctrace"); - + // Generate procmon skeleton and bindings generate_skeleton(&mut out, "procmon"); generate_header(&mut out, "procmon"); @@ -61,7 +61,7 @@ fn main() { // Generate tcpsniff skeleton (no header — reuses sslsniff.h event format) generate_skeleton(&mut out, "tcpsniff"); - + // generate_header(&mut out, "frametypes"); // generate_header(&mut out, "errors"); // generate_header(&mut out, "stackdeltatypes"); @@ -71,8 +71,7 @@ fn main() { let manifest_dir = env::var("CARGO_MANIFEST_DIR").expect("CARGO_MANIFEST_DIR must be set"); let frontend_dist = PathBuf::from(&manifest_dir).join("frontend-dist"); if !frontend_dist.exists() { - std::fs::create_dir_all(&frontend_dist) - .expect("Failed to create frontend-dist directory"); + std::fs::create_dir_all(&frontend_dist).expect("Failed to create frontend-dist directory"); } // Watch the directory AND each file inside it so cargo detects content changes println!("cargo:rerun-if-changed=frontend-dist"); @@ -84,7 +83,9 @@ fn main() { // Generate C header from src/ffi.rs via cbindgen let crate_dir = env::var("CARGO_MANIFEST_DIR").unwrap(); - let header_path = PathBuf::from(&crate_dir).join("include").join("agentsight.h"); + let header_path = PathBuf::from(&crate_dir) + .join("include") + .join("agentsight.h"); std::fs::create_dir_all(header_path.parent().unwrap()) .expect("Failed to create include/ directory"); cbindgen::Builder::new() diff --git a/src/agentsight/src/aggregator/http/aggregator.rs b/src/agentsight/src/aggregator/http/aggregator.rs index 8ad660819..54c077cdd 100644 --- a/src/agentsight/src/aggregator/http/aggregator.rs +++ b/src/agentsight/src/aggregator/http/aggregator.rs @@ -3,15 +3,15 @@ //! This module implements the HTTP Aggregator specification for correlating //! parsed HTTP requests and responses into complete request/response pairs. -use std::num::NonZeroUsize; -use lru::LruCache; +use super::super::result::AggregatedResult; +use super::pair::HttpPair; +use super::response::AggregatedResponse; use crate::config::DEFAULT_CONNECTION_CAPACITY; -use crate::probes::sslsniff::SslEvent; use crate::parser::http::{ParsedRequest, ParsedResponse}; use crate::parser::sse::{ParsedSseEvent, SseParser}; -use super::response::AggregatedResponse; -use super::pair::HttpPair; -use super::super::result::AggregatedResult; +use crate::probes::sslsniff::SslEvent; +use lru::LruCache; +use std::num::NonZeroUsize; /// Connection identifier - uniquely identifies an SSL connection #[derive(Debug, Clone, Copy, Hash, Eq, PartialEq)] @@ -36,9 +36,7 @@ pub enum ConnectionState { /// Idle - waiting for request Idle, /// Request pending - waiting for response - RequestPending { - request: ParsedRequest, - }, + RequestPending { request: ParsedRequest }, /// Request body pending - body not yet complete, waiting for more data or response RequestBodyPending { request: ParsedRequest, @@ -82,20 +80,20 @@ impl HttpConnectionAggregator { /// Insert connection state, logging if an unrelated entry is evicted by LRU fn insert(&mut self, key: ConnectionId, state: ConnectionState) { - if let Some((evicted_key, evicted_state)) = self.connections.push(key, state) { - if evicted_key != key { - log::warn!( - "[HttpAggregator] LRU evicted conn={:?} state={} | capacity={}", - evicted_key, - match evicted_state { - ConnectionState::Idle => "Idle", - ConnectionState::RequestPending { .. } => "RequestPending", - ConnectionState::RequestBodyPending { .. } => "RequestBodyPending", - ConnectionState::SseActive { .. } => "SseActive", - }, - self.connections.cap(), - ); - } + if let Some((evicted_key, evicted_state)) = self.connections.push(key, state) + && evicted_key != key + { + log::warn!( + "[HttpAggregator] LRU evicted conn={:?} state={} | capacity={}", + evicted_key, + match evicted_state { + ConnectionState::Idle => "Idle", + ConnectionState::RequestPending { .. } => "RequestPending", + ConnectionState::RequestBodyPending { .. } => "RequestBodyPending", + ConnectionState::SseActive { .. } => "SseActive", + }, + self.connections.cap(), + ); } } @@ -165,10 +163,7 @@ impl HttpConnectionAggregator { request.method, request.path, ); - self.insert( - connection_id, - ConnectionState::RequestPending { request }, - ); + self.insert(connection_id, ConnectionState::RequestPending { request }); } else { log::debug!( "[HttpAggregator] State transition: -> RequestBodyPending | conn={:?} | method={} | path={} | body_len={} | content_length={:?}", @@ -192,14 +187,11 @@ impl HttpConnectionAggregator { /// Process HTTP Response (from HTTP Parser) /// Returns completed HttpPair or SSE started signal - pub fn process_response( - &mut self, - response: ParsedResponse, - ) -> Option { + pub fn process_response(&mut self, response: ParsedResponse) -> Option { let connection_id = ConnectionId::from_ssl_event(&response.source_event); - + let state = self.connections.pop(&connection_id)?; - + match state { ConnectionState::RequestBodyPending { request, @@ -232,11 +224,7 @@ impl HttpConnectionAggregator { ); None } else { - let pair = HttpPair::from_parsed( - connection_id, - completed_request, - response, - ); + let pair = HttpPair::from_parsed(connection_id, completed_request, response); Some(AggregatedResult::HttpComplete(pair)) } } @@ -259,7 +247,7 @@ impl HttpConnectionAggregator { sse_events, }, ); - + // Don't return HttpPair yet, wait for SSE events to complete None } else { @@ -268,11 +256,7 @@ impl HttpConnectionAggregator { connection_id, response.status_code, ); - let pair = HttpPair::from_parsed( - connection_id, - request, - response, - ); + let pair = HttpPair::from_parsed(connection_id, request, response); Some(AggregatedResult::HttpComplete(pair)) } } @@ -396,7 +380,7 @@ impl HttpConnectionAggregator { sse_event: ParsedSseEvent, ) -> Option { let state = self.connections.pop(connection_id)?; - + match state { ConnectionState::SseActive { request, @@ -420,11 +404,11 @@ impl HttpConnectionAggregator { "[HttpAggregator] State transition: SseActive -> Complete | conn={:?}", connection_id, ); - + // Build aggregated response with SSE events let mut response = AggregatedResponse::from_parsed(response_headers); response.set_sse_events(sse_events); - + // Return appropriate result based on whether request exists if let Some(req) = request { let parsed = response.parsed.clone(); @@ -447,7 +431,7 @@ impl HttpConnectionAggregator { sse_events, }, ); - + None } } @@ -496,7 +480,8 @@ impl HttpConnectionAggregator { /// Drain all connections (for force complete) pub fn drain_connections(&mut self) -> Vec<(ConnectionId, ConnectionState)> { - self.connections.iter_mut() + self.connections + .iter_mut() .map(|(k, v)| (*k, v.clone())) .collect::>() .into_iter() @@ -515,12 +500,11 @@ impl HttpConnectionAggregator { use std::collections::HashSet; // 1. Collect unique PIDs - let pids: HashSet = self.connections.iter() - .map(|(k, _)| k.pid) - .collect(); + let pids: HashSet = self.connections.iter().map(|(k, _)| k.pid).collect(); // 2. Determine which PIDs are dead - let dead_pids: HashSet = pids.into_iter() + let dead_pids: HashSet = pids + .into_iter() .filter(|pid| !std::path::Path::new(&format!("/proc/{}", pid)).exists()) .collect(); @@ -529,7 +513,9 @@ impl HttpConnectionAggregator { } // 3. Collect keys for dead PIDs (can't mutate while iterating) - let dead_keys: Vec = self.connections.iter() + let dead_keys: Vec = self + .connections + .iter() .filter(|(k, _)| dead_pids.contains(&k.pid)) .map(|(k, _)| *k) .collect(); @@ -545,7 +531,8 @@ impl HttpConnectionAggregator { _ => { log::debug!( "[HttpAggregator] Draining dead-PID connection: pid={} ssl_ptr={:#x}", - key.pid, key.ssl_ptr, + key.pid, + key.ssl_ptr, ); result.push((key, state)); } @@ -568,8 +555,8 @@ impl HttpConnectionAggregator { #[cfg(test)] mod tests { use super::*; - use std::rc::Rc; use std::collections::HashMap; + use std::rc::Rc; fn create_mock_ssl_event(pid: u32, ssl_ptr: u64) -> Rc { Rc::new(SslEvent { @@ -590,7 +577,10 @@ mod tests { #[test] fn test_connection_id() { - let id = ConnectionId { pid: 1234, ssl_ptr: 0x1000 }; + let id = ConnectionId { + pid: 1234, + ssl_ptr: 0x1000, + }; assert_eq!(id.pid, 1234); assert_eq!(id.ssl_ptr, 0x1000); } @@ -599,7 +589,7 @@ mod tests { fn test_process_request_response_pair() { let mut aggregator = HttpConnectionAggregator::new(); let event = create_mock_ssl_event(1234, 0x1000); - + // Process request let request = ParsedRequest { method: "GET".to_string(), @@ -612,9 +602,12 @@ mod tests { reassembled_body: None, }; aggregator.process_request(request); - - assert!(aggregator.has_pending_request(&ConnectionId { pid: 1234, ssl_ptr: 0x1000 })); - + + assert!(aggregator.has_pending_request(&ConnectionId { + pid: 1234, + ssl_ptr: 0x1000 + })); + // Process response let response = ParsedResponse { version: 11, @@ -625,10 +618,10 @@ mod tests { body_len: 0, source_event: event, }; - + let result = aggregator.process_response(response); assert!(result.is_some()); - + if let Some(AggregatedResult::HttpComplete(pair)) = result { assert_eq!(pair.request.method, "GET"); assert_eq!(pair.response.status_code(), 200); @@ -642,7 +635,7 @@ mod tests { fn test_sse_detection() { let mut aggregator = HttpConnectionAggregator::new(); let event = create_mock_ssl_event(1234, 0x1000); - + // Process request let request = ParsedRequest { method: "GET".to_string(), @@ -655,11 +648,11 @@ mod tests { reassembled_body: None, }; aggregator.process_request(request); - + // Process SSE response let mut headers = HashMap::new(); headers.insert("content-type".to_string(), "text/event-stream".to_string()); - + let response = ParsedResponse { version: 11, status_code: 200, @@ -669,15 +662,23 @@ mod tests { body_len: 0, source_event: event, }; - + let result = aggregator.process_response(response); - + // SSE response should not return result immediately, but should activate SSE state assert!(result.is_none()); - assert!(aggregator.is_sse_active(&ConnectionId { pid: 1234, ssl_ptr: 0x1000 })); + assert!(aggregator.is_sse_active(&ConnectionId { + pid: 1234, + ssl_ptr: 0x1000 + })); } - fn create_mock_ssl_event_with_buf(pid: u32, ssl_ptr: u64, buf: Vec, rw: i32) -> Rc { + fn create_mock_ssl_event_with_buf( + pid: u32, + ssl_ptr: u64, + buf: Vec, + rw: i32, + ) -> Rc { Rc::new(SslEvent { source: 0, timestamp_ns: 1000, @@ -700,12 +701,15 @@ mod tests { // Simulate a request with Content-Length: 20 but only 5 bytes in first event let headers_and_partial_body = b"POST /api HTTP/1.1\r\nContent-Length: 20\r\n\r\nhello"; - let event1 = create_mock_ssl_event_with_buf(1234, 0x2000, headers_and_partial_body.to_vec(), 1); + let event1 = + create_mock_ssl_event_with_buf(1234, 0x2000, headers_and_partial_body.to_vec(), 1); // Parse as request (simulating what HttpParser would produce) - let header_end = headers_and_partial_body.windows(4) + let header_end = headers_and_partial_body + .windows(4) .position(|w| w == b"\r\n\r\n") - .unwrap() + 4; + .unwrap() + + 4; let body_len = headers_and_partial_body.len() - header_end; let mut headers = HashMap::new(); @@ -724,27 +728,44 @@ mod tests { // Process request - should enter RequestBodyPending since body_len(5) < content_length(20) aggregator.process_request(request); - let conn_id = ConnectionId { pid: 1234, ssl_ptr: 0x2000 }; + let conn_id = ConnectionId { + pid: 1234, + ssl_ptr: 0x2000, + }; assert!(!aggregator.has_pending_request(&conn_id)); // Send continuation data (10 bytes) let continuation1 = SslEvent { - source: 0, timestamp_ns: 2000, delta_ns: 0, - pid: 1234, tid: 1, uid: 0, len: 10, rw: 1, + source: 0, + timestamp_ns: 2000, + delta_ns: 0, + pid: 1234, + tid: 1, + uid: 0, + len: 10, + rw: 1, comm: String::new(), buf: b" world fir".to_vec(), - is_handshake: false, ssl_ptr: 0x2000, + is_handshake: false, + ssl_ptr: 0x2000, }; let result = aggregator.process_raw_body_data(&continuation1); assert!(result.is_none()); // Still incomplete (15 < 20) // Send final continuation (5 bytes, total = 5 + 10 + 5 = 20) let continuation2 = SslEvent { - source: 0, timestamp_ns: 3000, delta_ns: 0, - pid: 1234, tid: 1, uid: 0, len: 5, rw: 1, + source: 0, + timestamp_ns: 3000, + delta_ns: 0, + pid: 1234, + tid: 1, + uid: 0, + len: 5, + rw: 1, comm: String::new(), buf: b"st!!!".to_vec(), - is_handshake: false, ssl_ptr: 0x2000, + is_handshake: false, + ssl_ptr: 0x2000, }; let result = aggregator.process_raw_body_data(&continuation2); assert!(result.is_none()); // Transitioned to RequestPending @@ -753,8 +774,12 @@ mod tests { assert!(aggregator.has_pending_request(&conn_id)); // Sending a response should complete the pair - let resp_event = create_mock_ssl_event_with_buf(1234, 0x2000, - b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\n\r\nOK".to_vec(), 0); + let resp_event = create_mock_ssl_event_with_buf( + 1234, + 0x2000, + b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\n\r\nOK".to_vec(), + 0, + ); let response = ParsedResponse { version: 1, status_code: 200, @@ -786,9 +811,11 @@ mod tests { let headers_and_partial = b"POST /chat HTTP/1.1\r\nContent-Length: 100\r\n\r\npartial"; let event = create_mock_ssl_event_with_buf(5678, 0x3000, headers_and_partial.to_vec(), 1); - let header_end = headers_and_partial.windows(4) + let header_end = headers_and_partial + .windows(4) .position(|w| w == b"\r\n\r\n") - .unwrap() + 4; + .unwrap() + + 4; let body_len = headers_and_partial.len() - header_end; let mut headers = HashMap::new(); @@ -809,17 +836,24 @@ mod tests { // Send some continuation let cont = SslEvent { - source: 0, timestamp_ns: 2000, delta_ns: 0, - pid: 5678, tid: 1, uid: 0, len: 10, rw: 1, + source: 0, + timestamp_ns: 2000, + delta_ns: 0, + pid: 5678, + tid: 1, + uid: 0, + len: 10, + rw: 1, comm: String::new(), buf: b"_more_data".to_vec(), - is_handshake: false, ssl_ptr: 0x3000, + is_handshake: false, + ssl_ptr: 0x3000, }; aggregator.process_raw_body_data(&cont); // Response arrives before Content-Length is satisfied → force-complete - let resp_event = create_mock_ssl_event_with_buf(5678, 0x3000, - b"HTTP/1.1 200 OK\r\n\r\n{}".to_vec(), 0); + let resp_event = + create_mock_ssl_event_with_buf(5678, 0x3000, b"HTTP/1.1 200 OK\r\n\r\n{}".to_vec(), 0); let response = ParsedResponse { version: 1, status_code: 200, @@ -850,9 +884,11 @@ mod tests { let full_request = b"POST /api HTTP/1.1\r\nContent-Length: 5\r\n\r\nhello"; let event = create_mock_ssl_event_with_buf(1234, 0x4000, full_request.to_vec(), 1); - let header_end = full_request.windows(4) + let header_end = full_request + .windows(4) .position(|w| w == b"\r\n\r\n") - .unwrap() + 4; + .unwrap() + + 4; let body_len = full_request.len() - header_end; let mut headers = HashMap::new(); @@ -871,7 +907,10 @@ mod tests { // Should go directly to RequestPending (no aggregation needed) aggregator.process_request(request); - let conn_id = ConnectionId { pid: 1234, ssl_ptr: 0x4000 }; + let conn_id = ConnectionId { + pid: 1234, + ssl_ptr: 0x4000, + }; assert!(aggregator.has_pending_request(&conn_id)); } @@ -881,11 +920,18 @@ mod tests { // Send raw data for a connection that has no pending body let raw = SslEvent { - source: 0, timestamp_ns: 1000, delta_ns: 0, - pid: 9999, tid: 1, uid: 0, len: 5, rw: 1, + source: 0, + timestamp_ns: 1000, + delta_ns: 0, + pid: 9999, + tid: 1, + uid: 0, + len: 5, + rw: 1, comm: String::new(), buf: b"hello".to_vec(), - is_handshake: false, ssl_ptr: 0x5000, + is_handshake: false, + ssl_ptr: 0x5000, }; let result = aggregator.process_raw_body_data(&raw); assert!(result.is_none()); @@ -900,9 +946,7 @@ mod tests { let partial = b"POST /stream HTTP/1.1\r\nContent-Length: 50\r\n\r\ndata"; let event = create_mock_ssl_event_with_buf(1234, 0x6000, partial.to_vec(), 1); - let header_end = partial.windows(4) - .position(|w| w == b"\r\n\r\n") - .unwrap() + 4; + let header_end = partial.windows(4).position(|w| w == b"\r\n\r\n").unwrap() + 4; let body_len = partial.len() - header_end; let mut headers = HashMap::new(); @@ -922,8 +966,12 @@ mod tests { aggregator.process_request(request); // SSE response arrives → should force-complete body and enter SseActive - let resp_event = create_mock_ssl_event_with_buf(1234, 0x6000, - b"HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\n\r\n".to_vec(), 0); + let resp_event = create_mock_ssl_event_with_buf( + 1234, + 0x6000, + b"HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\n\r\n".to_vec(), + 0, + ); let mut resp_headers = HashMap::new(); resp_headers.insert("content-type".to_string(), "text/event-stream".to_string()); @@ -940,7 +988,10 @@ mod tests { let result = aggregator.process_response(response); // SSE response should not return immediately, should enter SseActive assert!(result.is_none()); - let conn_id = ConnectionId { pid: 1234, ssl_ptr: 0x6000 }; + let conn_id = ConnectionId { + pid: 1234, + ssl_ptr: 0x6000, + }; assert!(aggregator.is_sse_active(&conn_id)); } @@ -980,21 +1031,21 @@ mod tests { h }, body_offset: resp_buf.windows(4).position(|w| w == b"\r\n\r\n").unwrap() + 4, - body_len: resp_buf.len() - (resp_buf.windows(4).position(|w| w == b"\r\n\r\n").unwrap() + 4), + body_len: resp_buf.len() + - (resp_buf.windows(4).position(|w| w == b"\r\n\r\n").unwrap() + 4), source_event: resp_event, }; let result = aggregator.process_response(response); assert!(result.is_none()); - let done_event = create_mock_ssl_event_with_buf( - 4321, - 0x7000, - b"data: [DONE]\n\n".to_vec(), - 0, - ); + let done_event = + create_mock_ssl_event_with_buf(4321, 0x7000, b"data: [DONE]\n\n".to_vec(), 0); let done = ParsedSseEvent::new(None, None, None, 6, 6, done_event); - let conn_id = ConnectionId { pid: 4321, ssl_ptr: 0x7000 }; + let conn_id = ConnectionId { + pid: 4321, + ssl_ptr: 0x7000, + }; let result = aggregator.process_sse_event(&conn_id, done); let pair = match result { Some(AggregatedResult::SseComplete(pair)) => pair, diff --git a/src/agentsight/src/aggregator/http/mod.rs b/src/agentsight/src/aggregator/http/mod.rs index b94ddc70e..84bdcb9fc 100644 --- a/src/agentsight/src/aggregator/http/mod.rs +++ b/src/agentsight/src/aggregator/http/mod.rs @@ -7,7 +7,7 @@ mod pair; mod response; // Re-export main types -pub use aggregator::{HttpConnectionAggregator, ConnectionId, ConnectionState}; +pub use aggregator::{ConnectionId, ConnectionState, HttpConnectionAggregator}; pub use pair::HttpPair; pub use response::AggregatedResponse; diff --git a/src/agentsight/src/aggregator/http/pair.rs b/src/agentsight/src/aggregator/http/pair.rs index 1260b3ca5..1876d57b3 100644 --- a/src/agentsight/src/aggregator/http/pair.rs +++ b/src/agentsight/src/aggregator/http/pair.rs @@ -29,14 +29,14 @@ impl HttpPair { response: ParsedResponse, ) -> Self { let aggregated_response = AggregatedResponse::from_parsed(response); - + HttpPair { connection_id, request, response: aggregated_response, } } - + /// Add SSE event to the response pub fn add_sse_event(&mut self, event: ParsedSseEvent) { self.response.add_sse_event(event); @@ -45,32 +45,29 @@ impl HttpPair { impl ToChromeTraceEvent for HttpPair { /// Convert HttpPair to Chrome Trace Events - /// + /// /// Returns request, response events and flow events linking them fn to_chrome_trace_events(&self) -> Vec { let mut events = Vec::new(); - + // Add request events let req_events = self.request.to_chrome_trace_events(); events.extend(req_events); - + // Add response events let resp_events = self.response.to_chrome_trace_events(); events.extend(resp_events); - + // Create flow events linking request to response // Flow ID is generated on-demand for trace generation if let (Some(req_event), Some(resp_event)) = (events.first(), events.last()) { let flow_id = next_flow_id(); - let (flow_start, flow_end) = ChromeTraceEvent::flow_from_events_with_id( - req_event, - resp_event, - flow_id - ); + let (flow_start, flow_end) = + ChromeTraceEvent::flow_from_events_with_id(req_event, resp_event, flow_id); events.push(flow_start); events.push(flow_end); } - + events } } diff --git a/src/agentsight/src/aggregator/http/response.rs b/src/agentsight/src/aggregator/http/response.rs index 33496cc6c..f039d01a6 100644 --- a/src/agentsight/src/aggregator/http/response.rs +++ b/src/agentsight/src/aggregator/http/response.rs @@ -27,12 +27,13 @@ impl AggregatedResponse { } pub fn body(&self) -> &[u8] { - &self.parsed.body() + self.parsed.body() } pub fn body_string(&self) -> String { let first = std::str::from_utf8(self.body()).unwrap_or(""); - let sse_body: String = self.sse_events + let sse_body: String = self + .sse_events .iter() .map(|event| event.body_str()) .collect::>() diff --git a/src/agentsight/src/aggregator/http2.rs b/src/agentsight/src/aggregator/http2.rs index f65bb3f6d..9b6c65e73 100644 --- a/src/agentsight/src/aggregator/http2.rs +++ b/src/agentsight/src/aggregator/http2.rs @@ -4,14 +4,14 @@ //! by their stream_id and correlating request (client->server) with response (server->client) //! to form complete HTTP/2 request/response pairs. -use std::num::NonZeroUsize; -use lru::LruCache; -use crate::config::DEFAULT_CONNECTION_CAPACITY; -use crate::parser::http2::ParsedHttp2Frame; -use crate::parser::sse::SSEParser; use crate::aggregator::http::ConnectionId; use crate::aggregator::result::AggregatedResult; use crate::chrome_trace::{ChromeTraceEvent, ToChromeTraceEvent, ns_to_us}; +use crate::config::DEFAULT_CONNECTION_CAPACITY; +use crate::parser::http2::ParsedHttp2Frame; +use crate::parser::sse::SSEParser; +use lru::LruCache; +use std::num::NonZeroUsize; /// Stream identifier within an HTTP/2 connection #[derive(Debug, Clone, Copy, Hash, Eq, PartialEq)] @@ -230,22 +230,23 @@ impl Http2Stream { } /// Parse response body as SSE events and return JSON array of event data - /// + /// /// This method parses the response body as SSE (Server-Sent Events) stream /// and returns a JSON array containing each event's data field. /// If the body is not valid SSE format, returns None. pub fn response_sse_json_array(&self) -> Option { let body_str = self.response_body_str()?; - + // Use legacy SSEParser to parse the stream (returns owned data) let sse_events = SSEParser::parse_stream(&body_str); - + if sse_events.events.is_empty() { return None; } - + // Extract JSON data from each event - let json_array: Vec = sse_events.events + let json_array: Vec = sse_events + .events .iter() .filter_map(|event| { // Skip [DONE] marker @@ -256,7 +257,7 @@ impl Http2Stream { serde_json::from_str::(&event.data).ok() }) .collect(); - + if json_array.is_empty() { None } else { @@ -266,10 +267,12 @@ impl Http2Stream { /// Check if response content-type indicates SSE stream pub fn is_response_sse(&self) -> bool { - self.response_headers.as_ref() + self.response_headers + .as_ref() .map(|h| { let headers = h.decode_headers_stateless(); - headers.iter() + headers + .iter() .find(|(name, _)| name.eq_ignore_ascii_case("content-type")) .and_then(|(_, value)| value.clone()) .map(|ct| ct.contains("text/event-stream")) @@ -280,10 +283,12 @@ impl Http2Stream { /// Extract HTTP method from request headers (e.g., "GET", "POST") pub fn method(&self) -> String { - self.request_headers.as_ref() + self.request_headers + .as_ref() .map(|h| { let headers = h.decode_headers_stateless(); - headers.iter() + headers + .iter() .find(|(name, _)| name == ":method") .and_then(|(_, value)| value.clone()) .unwrap_or_else(|| "POST".to_string()) @@ -293,10 +298,12 @@ impl Http2Stream { /// Extract path from request headers (e.g., "/v1/chat/completions") pub fn path(&self) -> String { - self.request_headers.as_ref() + self.request_headers + .as_ref() .map(|h| { let headers = h.decode_headers_stateless(); - headers.iter() + headers + .iter() .find(|(name, _)| name == ":path") .and_then(|(_, value)| value.clone()) .unwrap_or_default() @@ -306,10 +313,12 @@ impl Http2Stream { /// Extract status code from response headers pub fn status_code(&self) -> u16 { - self.response_headers.as_ref() + self.response_headers + .as_ref() .map(|h| { let headers = h.decode_headers_stateless(); - headers.iter() + headers + .iter() .find(|(name, _)| name == ":status") .and_then(|(_, value)| value.clone()) .and_then(|v| v.parse().ok()) @@ -321,7 +330,8 @@ impl Http2Stream { /// Get request headers as JSON string pub fn request_headers_json(&self) -> String { if let Some(ref headers) = self.request_headers { - let decoded = headers.decode_headers_stateless() + let decoded = headers + .decode_headers_stateless() .into_iter() .filter_map(|(name, value)| value.map(|v| (name, v))) .collect::>(); @@ -334,7 +344,8 @@ impl Http2Stream { /// Get response headers as JSON string pub fn response_headers_json(&self) -> String { if let Some(ref headers) = self.response_headers { - let decoded = headers.decode_headers_stateless() + let decoded = headers + .decode_headers_stateless() .into_iter() .filter_map(|(name, value)| value.map(|v| (name, v))) .collect::>(); @@ -346,21 +357,39 @@ impl Http2Stream { /// Get process command name from source event pub fn comm(&self) -> String { - self.request_headers.as_ref() + self.request_headers + .as_ref() .map(|h| h.source_event.comm_str()) - .or_else(|| self.request_data_frames.first().map(|f| f.source_event.comm_str())) - .or_else(|| self.response_headers.as_ref().map(|h| h.source_event.comm_str())) - .or_else(|| self.response_data_frames.first().map(|f| f.source_event.comm_str())) + .or_else(|| { + self.request_data_frames + .first() + .map(|f| f.source_event.comm_str()) + }) + .or_else(|| { + self.response_headers + .as_ref() + .map(|h| h.source_event.comm_str()) + }) + .or_else(|| { + self.response_data_frames + .first() + .map(|f| f.source_event.comm_str()) + }) .unwrap_or_default() } /// Get process ID from source event pub fn pid(&self) -> u32 { - self.request_headers.as_ref() + self.request_headers + .as_ref() .map(|h| h.source_event.pid) .or_else(|| self.request_data_frames.first().map(|f| f.source_event.pid)) .or_else(|| self.response_headers.as_ref().map(|h| h.source_event.pid)) - .or_else(|| self.response_data_frames.first().map(|f| f.source_event.pid)) + .or_else(|| { + self.response_data_frames + .first() + .map(|f| f.source_event.pid) + }) .unwrap_or(0) } } @@ -418,11 +447,9 @@ impl Http2StreamAggregator { // Get or create stream state let state = self.streams.pop(&stream_id); - let mut state = state.unwrap_or_else(|| { - Http2StreamState::WaitingRequestData { - request_headers: None, - request_data_frames: Vec::new(), - } + let mut state = state.unwrap_or_else(|| Http2StreamState::WaitingRequestData { + request_headers: None, + request_data_frames: Vec::new(), }); // Process frame based on current state and direction @@ -450,9 +477,16 @@ impl Http2StreamAggregator { direction: StreamDirection, stream_id: &StreamId, ) -> Http2StreamState { - log::debug!("Processing http/2 frame in state: {}, stream_id: {:?}", state.state_name(), stream_id); + log::debug!( + "Processing http/2 frame in state: {}, stream_id: {:?}", + state.state_name(), + stream_id + ); match state { - Http2StreamState::WaitingRequestData { mut request_headers, mut request_data_frames } => { + Http2StreamState::WaitingRequestData { + mut request_headers, + mut request_data_frames, + } => { if direction == StreamDirection::Request { if frame.is_headers() { request_headers = Some(frame.clone()); @@ -487,18 +521,25 @@ impl Http2StreamAggregator { } } - Http2StreamState::RequestComplete { request_headers, request_data_frames } => { + Http2StreamState::RequestComplete { + request_headers, + request_data_frames, + } => { if direction == StreamDirection::Response { let mut response_headers = None; let mut response_data_frames = Vec::new(); - + if frame.is_headers() { response_headers = Some(frame.clone()); if frame.has_end_stream() { // Response is complete (no body) - let mut stream = Http2Stream::new(*stream_id, - request_headers.as_ref().map(|h| h.source_event.timestamp_ns) - .unwrap_or(frame.source_event.timestamp_ns)); + let mut stream = Http2Stream::new( + *stream_id, + request_headers + .as_ref() + .map(|h| h.source_event.timestamp_ns) + .unwrap_or(frame.source_event.timestamp_ns), + ); stream.request_headers = request_headers; stream.request_data_frames = request_data_frames; stream.request_complete = true; @@ -511,9 +552,13 @@ impl Http2StreamAggregator { response_data_frames.push(frame.clone()); if frame.has_end_stream() { // Response is complete - let mut stream = Http2Stream::new(*stream_id, - request_headers.as_ref().map(|h| h.source_event.timestamp_ns) - .unwrap_or(frame.source_event.timestamp_ns)); + let mut stream = Http2Stream::new( + *stream_id, + request_headers + .as_ref() + .map(|h| h.source_event.timestamp_ns) + .unwrap_or(frame.source_event.timestamp_ns), + ); stream.request_headers = request_headers; stream.request_data_frames = request_data_frames; stream.request_complete = true; @@ -524,7 +569,7 @@ impl Http2StreamAggregator { return Http2StreamState::Complete(stream); } } - + // Continue receiving response data Http2StreamState::ReceivingResponse { request_headers, @@ -552,9 +597,13 @@ impl Http2StreamAggregator { response_headers = Some(frame.clone()); if frame.has_end_stream() { // Response is complete - let mut stream = Http2Stream::new(*stream_id, - request_headers.as_ref().map(|h| h.source_event.timestamp_ns) - .unwrap_or(frame.source_event.timestamp_ns)); + let mut stream = Http2Stream::new( + *stream_id, + request_headers + .as_ref() + .map(|h| h.source_event.timestamp_ns) + .unwrap_or(frame.source_event.timestamp_ns), + ); stream.request_headers = request_headers; stream.request_data_frames = request_data_frames; stream.request_complete = true; @@ -568,9 +617,13 @@ impl Http2StreamAggregator { response_data_frames.push(frame.clone()); if frame.has_end_stream() { // Response is complete - let mut stream = Http2Stream::new(*stream_id, - request_headers.as_ref().map(|h| h.source_event.timestamp_ns) - .unwrap_or(frame.source_event.timestamp_ns)); + let mut stream = Http2Stream::new( + *stream_id, + request_headers + .as_ref() + .map(|h| h.source_event.timestamp_ns) + .unwrap_or(frame.source_event.timestamp_ns), + ); stream.request_headers = request_headers; stream.request_data_frames = request_data_frames; stream.request_complete = true; @@ -618,44 +671,59 @@ impl Http2StreamAggregator { /// Useful for shutdown or forced completion pub fn drain_pending(&mut self) -> Vec { let mut result = Vec::new(); - + // Move all streams from LRU cache while let Some((stream_id, state)) = self.streams.pop_lru() { if let Some(stream) = self.stream_from_state(state, stream_id) { result.push(stream); } } - + result } /// Convert a stream state to a Http2Stream if possible - fn stream_from_state(&self, state: Http2StreamState, stream_id: StreamId) -> Option { + fn stream_from_state( + &self, + state: Http2StreamState, + stream_id: StreamId, + ) -> Option { match state { Http2StreamState::Complete(stream) => Some(stream), - Http2StreamState::RequestComplete { request_headers, request_data_frames } => { - let timestamp_ns = request_headers.as_ref() + Http2StreamState::RequestComplete { + request_headers, + request_data_frames, + } => { + let timestamp_ns = request_headers + .as_ref() .map(|h| h.source_event.timestamp_ns) - .unwrap_or_else(|| request_data_frames.first() - .map(|f| f.source_event.timestamp_ns) - .unwrap_or(0)); + .unwrap_or_else(|| { + request_data_frames + .first() + .map(|f| f.source_event.timestamp_ns) + .unwrap_or(0) + }); let mut stream = Http2Stream::new(stream_id, timestamp_ns); stream.request_headers = request_headers; stream.request_data_frames = request_data_frames; stream.request_complete = true; Some(stream) } - Http2StreamState::ReceivingResponse { - request_headers, - request_data_frames, - response_headers, - response_data_frames + Http2StreamState::ReceivingResponse { + request_headers, + request_data_frames, + response_headers, + response_data_frames, } => { - let timestamp_ns = request_headers.as_ref() + let timestamp_ns = request_headers + .as_ref() .map(|h| h.source_event.timestamp_ns) - .unwrap_or_else(|| request_data_frames.first() - .map(|f| f.source_event.timestamp_ns) - .unwrap_or(0)); + .unwrap_or_else(|| { + request_data_frames + .first() + .map(|f| f.source_event.timestamp_ns) + .unwrap_or(0) + }); let mut stream = Http2Stream::new(stream_id, timestamp_ns); stream.request_headers = request_headers; stream.request_data_frames = request_data_frames; @@ -680,7 +748,10 @@ impl ToChromeTraceEvent for Http2Stream { fn to_chrome_trace_events(&self) -> Vec { let mut events = Vec::new(); let ts_us = ns_to_us(self.start_timestamp_ns); - let dur_us = ns_to_us(self.end_timestamp_ns.saturating_sub(self.start_timestamp_ns)); + let dur_us = ns_to_us( + self.end_timestamp_ns + .saturating_sub(self.start_timestamp_ns), + ); const MIN_DUR_US: u64 = 1_000; let actual_dur = dur_us.max(MIN_DUR_US); @@ -717,8 +788,8 @@ impl ToChromeTraceEvent for Http2Stream { #[cfg(test)] mod tests { use super::*; - use std::rc::Rc; use crate::probes::sslsniff::SslEvent; + use std::rc::Rc; fn create_test_event(pid: u32, ssl_ptr: u64, rw: i32, timestamp_ns: u64) -> Rc { Rc::new(SslEvent { @@ -746,7 +817,7 @@ mod tests { ) -> ParsedHttp2Frame { let payload_offset = 9; // Skip frame header let payload_len = payload.len(); - + // Create a new event with the payload in buf let mut buf = Vec::with_capacity(9 + payload_len); // Frame header @@ -761,7 +832,7 @@ mod tests { buf.push((stream_id & 0xFF) as u8); // Payload buf.extend_from_slice(&payload); - + let event_with_buf = Rc::new(SslEvent { source: event.source, timestamp_ns: event.timestamp_ns, @@ -800,13 +871,16 @@ mod tests { #[test] fn test_aggregator_process_request_response() { let mut aggregator = Http2StreamAggregator::new(); - let _conn_id = ConnectionId { pid: 1234, ssl_ptr: 0x1000 }; + let _conn_id = ConnectionId { + pid: 1234, + ssl_ptr: 0x1000, + }; // Create request HEADERS frame (rw=1, write) with END_STREAM (no body) let req_event = create_test_event(1234, 0x1000, 1, 1000); let req_headers = create_test_frame( - 1, // stream_id - 1, // HEADERS + 1, // stream_id + 1, // HEADERS 0x05, // END_HEADERS | END_STREAM - request has no body b":method: POST\n:path: /api/test".to_vec(), req_event, @@ -820,8 +894,8 @@ mod tests { // Create response HEADERS frame (rw=0, read) let resp_event = create_test_event(1234, 0x1000, 0, 2000); let resp_headers = create_test_frame( - 1, // stream_id - 1, // HEADERS + 1, // stream_id + 1, // HEADERS 0x05, // END_HEADERS | END_STREAM b":status: 200".to_vec(), resp_event, @@ -830,7 +904,7 @@ mod tests { // Process response let completed = aggregator.process_frames(vec![resp_headers]); assert_eq!(completed.len(), 1); - + let stream = &completed[0]; assert_eq!(stream.stream_id.stream_id, 1); assert!(stream.request_complete); @@ -859,7 +933,7 @@ mod tests { let completed = aggregator.process_frames(vec![resp_headers]); assert_eq!(completed.len(), 1); - + let stream = &completed[0]; assert_eq!(stream.request_data_frames.len(), 1); assert_eq!(stream.response_data_frames.len(), 0); diff --git a/src/agentsight/src/aggregator/mod.rs b/src/agentsight/src/aggregator/mod.rs index 23a4baa7a..de2ebca04 100644 --- a/src/agentsight/src/aggregator/mod.rs +++ b/src/agentsight/src/aggregator/mod.rs @@ -32,14 +32,12 @@ pub use result::AggregatedResult; // Re-export HTTP types pub use http::{ - HttpConnectionAggregator, ConnectionId, ConnectionState, - HttpPair, ParsedRequest, AggregatedResponse, + AggregatedResponse, ConnectionId, ConnectionState, HttpConnectionAggregator, HttpPair, + ParsedRequest, }; // Re-export HTTP/2 types -pub use http2::{ - Http2StreamAggregator, Http2Stream, Http2StreamState, StreamId, StreamDirection, -}; +pub use http2::{Http2Stream, Http2StreamAggregator, Http2StreamState, StreamDirection, StreamId}; // Re-export proctrace types -pub use proctrace::{ProcessEventAggregator, AggregatedProcess}; +pub use proctrace::{AggregatedProcess, ProcessEventAggregator}; diff --git a/src/agentsight/src/aggregator/proctrace/aggregator.rs b/src/agentsight/src/aggregator/proctrace/aggregator.rs index 4d275d430..c14ae1b87 100644 --- a/src/agentsight/src/aggregator/proctrace/aggregator.rs +++ b/src/agentsight/src/aggregator/proctrace/aggregator.rs @@ -2,10 +2,10 @@ //! //! Aggregates process events (exec, stdout, exit) by PID into complete process lifecycles. -use std::collections::HashMap; -use crate::probes::proctrace::VariableEvent; -use crate::parser::proctrace::{ParsedProcEvent, ProcEventType}; use super::process::AggregatedProcess; +use crate::parser::proctrace::{ParsedProcEvent, ProcEventType}; +use crate::probes::proctrace::VariableEvent; +use std::collections::HashMap; /// Process Event Aggregator - aggregates process events by PID #[derive(Debug, Clone)] @@ -28,24 +28,30 @@ impl ProcessEventAggregator { /// None otherwise. pub fn process_event(&mut self, event: &VariableEvent) -> Option { match event { - VariableEvent::Exec { header, filename, args } => { + VariableEvent::Exec { + header, + filename, + args, + } => { let pid = header.pid; - let aggregated = self.aggregates - .entry(pid) - .or_insert_with(|| { - AggregatedProcess::new( - pid, - header.tid, - header.ppid, - header.ptid, - event.comm_str(), - header.timestamp_ns, - ) - }); + let aggregated = self.aggregates.entry(pid).or_insert_with(|| { + AggregatedProcess::new( + pid, + header.tid, + header.ppid, + header.ptid, + event.comm_str(), + header.timestamp_ns, + ) + }); aggregated.add_exec(filename.clone(), args.clone(), header.timestamp_ns); None } - VariableEvent::Stdout { header, fd, payload } => { + VariableEvent::Stdout { + header, + fd, + payload, + } => { let pid = header.pid; if let Some(aggregated) = self.aggregates.get_mut(&pid) { if *fd == 2 { @@ -71,7 +77,10 @@ impl ProcessEventAggregator { /// Process multiple events pub fn process_events(&mut self, events: &[VariableEvent]) -> Vec { - events.iter().filter_map(|e| self.process_event(e)).collect() + events + .iter() + .filter_map(|e| self.process_event(e)) + .collect() } /// Process a parsed process event @@ -81,18 +90,16 @@ impl ProcessEventAggregator { pub fn process_parsed_event(&mut self, event: &ParsedProcEvent) -> Option { match event.event_type { ProcEventType::Exec => { - let aggregated = self.aggregates - .entry(event.pid) - .or_insert_with(|| { - AggregatedProcess::new( - event.pid, - event.tid, - event.ppid, - event.ptid, - event.comm.clone(), - event.timestamp_ns, - ) - }); + let aggregated = self.aggregates.entry(event.pid).or_insert_with(|| { + AggregatedProcess::new( + event.pid, + event.tid, + event.ppid, + event.ptid, + event.comm.clone(), + event.timestamp_ns, + ) + }); if let Some(ref args) = event.args { let filename = event.comm.clone(); aggregated.add_exec(filename, args.clone(), event.timestamp_ns); @@ -100,10 +107,10 @@ impl ProcessEventAggregator { None } ProcEventType::Stdout => { - if let Some(aggregated) = self.aggregates.get_mut(&event.pid) { - if let Some(ref data) = event.stdout_data { - aggregated.add_stdout(data.as_bytes(), event.timestamp_ns); - } + if let Some(aggregated) = self.aggregates.get_mut(&event.pid) + && let Some(ref data) = event.stdout_data + { + aggregated.add_stdout(data.as_bytes(), event.timestamp_ns); } None } @@ -120,7 +127,10 @@ impl ProcessEventAggregator { /// Get all incomplete aggregations (running processes) pub fn get_incomplete(&self) -> Vec<&AggregatedProcess> { - self.aggregates.values().filter(|agg| !agg.is_complete).collect() + self.aggregates + .values() + .filter(|agg| !agg.is_complete) + .collect() } /// Clear all aggregations diff --git a/src/agentsight/src/aggregator/proctrace/process.rs b/src/agentsight/src/aggregator/proctrace/process.rs index a4ec91977..d66646686 100644 --- a/src/agentsight/src/aggregator/proctrace/process.rs +++ b/src/agentsight/src/aggregator/proctrace/process.rs @@ -3,7 +3,7 @@ //! Defines the `AggregatedProcess` structure for representing a complete //! process lifecycle with exec, stdout, stderr, and exit events. -use crate::chrome_trace::{ChromeTraceEvent, ToChromeTraceEvent, ns_to_us, next_flow_id}; +use crate::chrome_trace::{ChromeTraceEvent, ToChromeTraceEvent, next_flow_id, ns_to_us}; /// Aggregated process data for a specific PID #[derive(Debug, Clone)] @@ -90,7 +90,8 @@ impl AggregatedProcess { /// Get the duration in nanoseconds pub fn duration_ns(&self) -> u64 { - self.end_timestamp_ns.saturating_sub(self.start_timestamp_ns) + self.end_timestamp_ns + .saturating_sub(self.start_timestamp_ns) } /// Get total stdout data size @@ -115,7 +116,12 @@ impl AggregatedProcess { }; match (stdout_preview.is_empty(), stderr_str.is_empty()) { - (false, false) => format!("process: {} | stdout: {} | stderr: {}", self.comm, stdout_preview, self.stderr_size()), + (false, false) => format!( + "process: {} | stdout: {} | stderr: {}", + self.comm, + stdout_preview, + self.stderr_size() + ), (false, true) => format!("process: {} | stdout: {}", self.comm, stdout_preview), (true, false) => format!("process: {} | stderr: {}", self.comm, self.stderr_size()), (true, true) => format!("process: {}", self.comm), @@ -140,7 +146,7 @@ impl ToChromeTraceEvent for AggregatedProcess { "child_filename": self.filename, "child_args": self.args, }); - + let fork_event = ChromeTraceEvent::complete( format!("fork: {}", self.comm), "process.fork", @@ -156,10 +162,10 @@ impl ToChromeTraceEvent for AggregatedProcess { // 2. Child process lifecycle event let stdout_str = self.stdout_string(); let stderr_str = self.stderr_string(); - + // Build event name with stdout preview (max 50 chars) let name = self.build_event_name(&stdout_str, &stderr_str); - + let args = serde_json::json!({ "pid": self.pid, "ppid": self.ppid, diff --git a/src/agentsight/src/aggregator/result.rs b/src/agentsight/src/aggregator/result.rs index cd0bd7d56..56d8e0c2a 100644 --- a/src/agentsight/src/aggregator/result.rs +++ b/src/agentsight/src/aggregator/result.rs @@ -3,11 +3,11 @@ //! This module defines the `AggregatedResult` enum which represents //! the output of aggregating parsed messages from various sources. -use crate::chrome_trace::{ChromeTraceEvent, ToChromeTraceEvent}; -use crate::parser::http2::ParsedHttp2Frame; -use super::http::{ConnectionId, HttpPair, ParsedRequest, AggregatedResponse}; +use super::http::{AggregatedResponse, ConnectionId, HttpPair, ParsedRequest}; use super::http2::Http2Stream; use super::proctrace::AggregatedProcess; +use crate::chrome_trace::{ChromeTraceEvent, ToChromeTraceEvent}; +use crate::parser::http2::ParsedHttp2Frame; /// Aggregated result from any aggregator #[derive(Debug, Clone)] @@ -52,7 +52,6 @@ impl AggregatedResult { } } - impl ToChromeTraceEvent for AggregatedResult { fn to_chrome_trace_events(&self) -> Vec { match self { @@ -62,17 +61,16 @@ impl ToChromeTraceEvent for AggregatedResult { AggregatedResult::RequestOnly { .. } => { log::warn!("RequestOnly: {:?}", self); vec![] - }, + } AggregatedResult::ResponseOnly { .. } => { log::warn!("ResponseOnly: {:?}", self); vec![] - }, - AggregatedResult::Http2Frames { frames, .. } => { - frames.iter().flat_map(|f| f.to_chrome_trace_events()).collect() - }, - AggregatedResult::Http2StreamComplete(stream) => { - stream.to_chrome_trace_events() - }, + } + AggregatedResult::Http2Frames { frames, .. } => frames + .iter() + .flat_map(|f| f.to_chrome_trace_events()) + .collect(), + AggregatedResult::Http2StreamComplete(stream) => stream.to_chrome_trace_events(), } } } diff --git a/src/agentsight/src/aggregator/unified.rs b/src/agentsight/src/aggregator/unified.rs index a2b628c1c..dd5a5bb02 100644 --- a/src/agentsight/src/aggregator/unified.rs +++ b/src/agentsight/src/aggregator/unified.rs @@ -46,20 +46,20 @@ impl Aggregator { self.http.process_request(req); vec![] } - ParsedMessage::Response(resp) => { - self.http.process_response(resp).into_iter().collect() - } + ParsedMessage::Response(resp) => self.http.process_response(resp).into_iter().collect(), ParsedMessage::SseEvent(sse_event) => { let conn_id = ConnectionId::from_ssl_event(sse_event.source_event()); - self.http.process_sse_event(&conn_id, sse_event).into_iter().collect() - } - ParsedMessage::ProcEvent(proc_event) => { - self.process - .process_parsed_event(&proc_event) - .map(AggregatedResult::ProcessComplete) + self.http + .process_sse_event(&conn_id, sse_event) .into_iter() .collect() } + ParsedMessage::ProcEvent(proc_event) => self + .process + .process_parsed_event(&proc_event) + .map(AggregatedResult::ProcessComplete) + .into_iter() + .collect(), ParsedMessage::Http2Frames(frames) => { // Use HTTP/2 stream aggregator to correlate frames by stream_id let completed_streams = self.http2.process_frames(frames); @@ -68,9 +68,11 @@ impl Aggregator { .map(AggregatedResult::Http2StreamComplete) .collect() } - ParsedMessage::RawData(ssl_event) => { - self.http.process_raw_body_data(&ssl_event).into_iter().collect() - } + ParsedMessage::RawData(ssl_event) => self + .http + .process_raw_body_data(&ssl_event) + .into_iter() + .collect(), } } @@ -91,12 +93,12 @@ impl Aggregator { .into_iter() .flat_map(|msg| self.process_message(msg)) .collect(); - + // Export chrome trace if enabled for r in &results { export_trace_events(r); } - + results } diff --git a/src/agentsight/src/analyzer/audit/analyzer.rs b/src/agentsight/src/analyzer/audit/analyzer.rs index 16f6939d1..671e7a966 100644 --- a/src/agentsight/src/analyzer/audit/analyzer.rs +++ b/src/agentsight/src/analyzer/audit/analyzer.rs @@ -2,12 +2,12 @@ //! //! Pure logic layer, no IO. Converts `AggregatedResult` into `AuditRecord`. -use crate::aggregator::HttpPair; +use super::record::{AuditEventType, AuditExtra, AuditRecord}; use crate::aggregator::AggregatedProcess; use crate::aggregator::AggregatedResult; +use crate::aggregator::HttpPair; use crate::analyzer::HttpRecord; use crate::analyzer::token::TokenRecord; -use super::record::{AuditEventType, AuditExtra, AuditRecord}; /// Analyzes aggregated results and extracts audit records pub struct AuditAnalyzer; @@ -42,14 +42,19 @@ impl AuditAnalyzer { /// Only creates llm_call for SSE responses, which are LLM streaming API calls. /// Non-SSE requests (like npm package queries) are filtered out. /// This method works for both HTTP/1.1 and HTTP/2 uniformly. - pub fn analyze_http(&self, http_record: &HttpRecord, token_record: Option<&TokenRecord>) -> Option { + pub fn analyze_http( + &self, + http_record: &HttpRecord, + token_record: Option<&TokenRecord>, + ) -> Option { // Only create llm_call for SSE responses if !http_record.is_sse { return None; } // Extract model from request body if available - let model = http_record.request_body + let model = http_record + .request_body .as_ref() .and_then(|body| serde_json::from_str::(body).ok()) .and_then(|json| json.get("model")?.as_str().map(|s| s.to_string())); diff --git a/src/agentsight/src/analyzer/message/anthropic.rs b/src/agentsight/src/analyzer/message/anthropic.rs index ad9d4d76d..c358cda49 100644 --- a/src/agentsight/src/analyzer/message/anthropic.rs +++ b/src/agentsight/src/analyzer/message/anthropic.rs @@ -25,7 +25,10 @@ //! } //! ``` -use super::types::{AnthropicRequest, AnthropicResponse, AnthropicContentBlock, AnthropicUsage, MessageRole, AnthropicSseEvent}; +use super::types::{ + AnthropicContentBlock, AnthropicRequest, AnthropicResponse, AnthropicSseEvent, AnthropicUsage, + MessageRole, +}; /// Parser for Anthropic Messages API /// @@ -54,11 +57,13 @@ impl AnthropicParser { /// ``` pub fn parse_request(body: &serde_json::Value) -> Option { // Quick validation - must have model, messages, and max_tokens fields - if !body.get("model").is_some() - || !body.get("messages").is_some() - || !body.get("max_tokens").is_some() + if body.get("model").is_none() + || body.get("messages").is_none() + || body.get("max_tokens").is_none() { - log::trace!("Anthropic request missing required fields: model, messages, or max_tokens"); + log::trace!( + "Anthropic request missing required fields: model, messages, or max_tokens" + ); return None; } @@ -102,9 +107,9 @@ impl AnthropicParser { /// ``` pub fn parse_response(body: &serde_json::Value) -> Option { // Try standard response format first (has id, type="message", content) - if body.get("id").is_some() + if body.get("id").is_some() && body.get("type").and_then(|v| v.as_str()) == Some("message") - && body.get("content").is_some() + && body.get("content").is_some() { match serde_json::from_value::(body.clone()) { Ok(response) => { @@ -146,15 +151,25 @@ impl AnthropicParser { // State for the current content block being streamed enum CurrentBlock { - Text { text: String }, - Thinking { thinking: String, signature: String }, - ToolUse { id: String, name: String, input_json: String }, + Text { + text: String, + }, + Thinking { + thinking: String, + signature: String, + }, + ToolUse { + id: String, + name: String, + input_json: String, + }, } let mut current_block: Option = None; for event_value in events { // Try to parse as AnthropicSseEvent - if let Ok(sse_event) = serde_json::from_value::(event_value.clone()) { + if let Ok(sse_event) = serde_json::from_value::(event_value.clone()) + { match &sse_event { AnthropicSseEvent::MessageStart { message } => { message_start = Some(sse_event.clone()); @@ -178,7 +193,9 @@ impl AnthropicParser { } _ => { // Text or any other block type - Some(CurrentBlock::Text { text: String::new() }) + Some(CurrentBlock::Text { + text: String::new(), + }) } }; } @@ -186,7 +203,9 @@ impl AnthropicParser { use super::types::AnthropicSseDelta; match delta { AnthropicSseDelta::TextDelta { text } => { - if let Some(CurrentBlock::Text { text: ref mut buf }) = current_block { + if let Some(CurrentBlock::Text { text: ref mut buf }) = + current_block + { buf.push_str(text); } else if current_block.is_none() { // Fallback: no ContentBlockStart seen, create text block @@ -194,7 +213,11 @@ impl AnthropicParser { } } AnthropicSseDelta::ThinkingDelta { thinking } => { - if let Some(CurrentBlock::Thinking { thinking: ref mut buf, .. }) = current_block { + if let Some(CurrentBlock::Thinking { + thinking: ref mut buf, + .. + }) = current_block + { buf.push_str(thinking); } else if current_block.is_none() { current_block = Some(CurrentBlock::Thinking { @@ -204,12 +227,19 @@ impl AnthropicParser { } } AnthropicSseDelta::SignatureDelta { signature } => { - if let Some(CurrentBlock::Thinking { signature: ref mut sig, .. }) = current_block { + if let Some(CurrentBlock::Thinking { + signature: ref mut sig, + .. + }) = current_block + { sig.push_str(signature); } } AnthropicSseDelta::InputJsonDelta { partial_json } => { - if let Some(CurrentBlock::ToolUse { ref mut input_json, .. }) = current_block { + if let Some(CurrentBlock::ToolUse { + ref mut input_json, .. + }) = current_block + { input_json.push_str(partial_json); } } @@ -227,17 +257,31 @@ impl AnthropicParser { }); } } - CurrentBlock::Thinking { thinking, signature } => { + CurrentBlock::Thinking { + thinking, + signature, + } => { if !thinking.is_empty() { content_blocks.push(AnthropicContentBlock::Thinking { thinking, - signature: if signature.is_empty() { None } else { Some(signature) }, + signature: if signature.is_empty() { + None + } else { + Some(signature) + }, }); } } - CurrentBlock::ToolUse { id, name, input_json } => { - let input = serde_json::from_str::(&input_json) - .unwrap_or(serde_json::Value::Object(serde_json::Map::new())); + CurrentBlock::ToolUse { + id, + name, + input_json, + } => { + let input = + serde_json::from_str::(&input_json) + .unwrap_or(serde_json::Value::Object( + serde_json::Map::new(), + )); content_blocks.push(AnthropicContentBlock::ToolUse { id, name, @@ -247,14 +291,21 @@ impl AnthropicParser { } } } - AnthropicSseEvent::MessageDelta { delta, usage: delta_usage } => { + AnthropicSseEvent::MessageDelta { + delta, + usage: delta_usage, + } => { stop_reason = delta.stop_reason.clone(); if let Some(du) = delta_usage { usage = Some(AnthropicUsage { input_tokens: usage.as_ref().map(|u| u.input_tokens).unwrap_or(0), output_tokens: du.output_tokens, - cache_creation_input_tokens: usage.as_ref().and_then(|u| u.cache_creation_input_tokens), - cache_read_input_tokens: usage.as_ref().and_then(|u| u.cache_read_input_tokens), + cache_creation_input_tokens: usage + .as_ref() + .and_then(|u| u.cache_creation_input_tokens), + cache_read_input_tokens: usage + .as_ref() + .and_then(|u| u.cache_read_input_tokens), }); } } @@ -274,15 +325,26 @@ impl AnthropicParser { }); } } - CurrentBlock::Thinking { thinking, signature } => { + CurrentBlock::Thinking { + thinking, + signature, + } => { if !thinking.is_empty() { content_blocks.push(AnthropicContentBlock::Thinking { thinking, - signature: if signature.is_empty() { None } else { Some(signature) }, + signature: if signature.is_empty() { + None + } else { + Some(signature) + }, }); } } - CurrentBlock::ToolUse { id, name, input_json } => { + CurrentBlock::ToolUse { + id, + name, + input_json, + } => { let input = serde_json::from_str::(&input_json) .unwrap_or(serde_json::Value::Object(serde_json::Map::new())); content_blocks.push(AnthropicContentBlock::ToolUse { id, name, input }); @@ -312,7 +374,10 @@ impl AnthropicParser { }) } else if !content_blocks.is_empty() { // No message_start but we still parsed content blocks — return with defaults - log::debug!("aggregate_sse_events: no message_start found, returning {} content blocks with defaults", content_blocks.len()); + log::debug!( + "aggregate_sse_events: no message_start found, returning {} content blocks with defaults", + content_blocks.len() + ); Some(AnthropicResponse { id: String::new(), type_: "message".to_string(), @@ -534,7 +599,9 @@ mod tests { #[test] fn test_matches_path() { assert!(AnthropicParser::matches_path("/v1/messages")); - assert!(AnthropicParser::matches_path("https://api.anthropic.com/v1/messages")); + assert!(AnthropicParser::matches_path( + "https://api.anthropic.com/v1/messages" + )); assert!(!AnthropicParser::matches_path("/v1/chat/completions")); assert!(!AnthropicParser::matches_path("/v1/completions")); } diff --git a/src/agentsight/src/analyzer/message/mod.rs b/src/agentsight/src/analyzer/message/mod.rs index 5b5548153..2ebf8f6ea 100644 --- a/src/agentsight/src/analyzer/message/mod.rs +++ b/src/agentsight/src/analyzer/message/mod.rs @@ -294,7 +294,10 @@ mod tests { assert!(result.is_some()); match result.unwrap() { - ParsedApiMessage::OpenAICompletion { request: req, response: resp } => { + ParsedApiMessage::OpenAICompletion { + request: req, + response: resp, + } => { assert!(req.is_some()); assert!(resp.is_none()); } @@ -319,7 +322,10 @@ mod tests { assert!(result.is_some()); match result.unwrap() { - ParsedApiMessage::AnthropicMessage { request: req, response: resp } => { + ParsedApiMessage::AnthropicMessage { + request: req, + response: resp, + } => { assert!(req.is_none()); assert!(resp.is_some()); } @@ -339,12 +345,18 @@ mod tests { #[test] fn test_detect_provider() { - assert_eq!(MessageParser::detect_provider("/v1/messages"), Some("anthropic")); + assert_eq!( + MessageParser::detect_provider("/v1/messages"), + Some("anthropic") + ); assert_eq!( MessageParser::detect_provider("/v1/chat/completions"), Some("openai") ); - assert_eq!(MessageParser::detect_provider("/v1/completions"), Some("openai")); + assert_eq!( + MessageParser::detect_provider("/v1/completions"), + Some("openai") + ); assert_eq!(MessageParser::detect_provider("/v1/embeddings"), None); } diff --git a/src/agentsight/src/analyzer/message/openai.rs b/src/agentsight/src/analyzer/message/openai.rs index 28c84e799..4683aceb6 100644 --- a/src/agentsight/src/analyzer/message/openai.rs +++ b/src/agentsight/src/analyzer/message/openai.rs @@ -26,7 +26,10 @@ //! } //! ``` -use super::types::{OpenAIRequest, OpenAIResponse, OpenAIChoice, OpenAIChatMessage, MessageRole, OpenAIContent, OpenAiSseChunk}; +use super::types::{ + MessageRole, OpenAIChatMessage, OpenAIChoice, OpenAIContent, OpenAIRequest, OpenAIResponse, + OpenAiSseChunk, +}; /// Parser for OpenAI Chat Completions API /// @@ -54,7 +57,7 @@ impl OpenAIParser { /// ``` pub fn parse_request(body: &serde_json::Value) -> Option { // Quick validation - must have model and messages fields - if !body.get("model").is_some() || !body.get("messages").is_some() { + if body.get("model").is_none() || body.get("messages").is_none() { log::trace!("OpenAI request missing required fields: model or messages"); return None; } @@ -141,29 +144,30 @@ impl OpenAIParser { } // Extract content delta for aggregation for choice in &sse_chunk.choices { - if let Some(content) = &choice.delta.content { - if !content.is_empty() { - content_parts.push(content.clone()); - } + if let Some(content) = &choice.delta.content + && !content.is_empty() + { + content_parts.push(content.clone()); } // Extract reasoning_content delta - if let Some(reasoning) = &choice.delta.reasoning_content { - if !reasoning.is_empty() { - reasoning_parts.push(reasoning.clone()); - } + if let Some(reasoning) = &choice.delta.reasoning_content + && !reasoning.is_empty() + { + reasoning_parts.push(reasoning.clone()); } // Extract and merge tool_call deltas by index if let Some(calls) = &choice.delta.tool_calls { for tc in calls { let idx = tc.get("index").and_then(|v| v.as_u64()).unwrap_or(0) as u32; - let entry = tool_call_map.entry(idx) + let entry = tool_call_map + .entry(idx) .or_insert_with(|| (String::new(), String::new(), String::new())); - if let Some(id) = tc.get("id").and_then(|v| v.as_str()) { - if !id.is_empty() { - entry.0 = id.to_string(); - } - // 空字符串不覆盖已有的 id + if let Some(id) = tc.get("id").and_then(|v| v.as_str()) + && !id.is_empty() + { + entry.0 = id.to_string(); } + // 空字符串不覆盖已有的 id if let Some(func) = tc.get("function") { if let Some(name) = func.get("name").and_then(|v| v.as_str()) { entry.1 = name.to_string(); @@ -187,56 +191,65 @@ impl OpenAIParser { } else { let mut sorted_indices: Vec = tool_call_map.keys().cloned().collect(); sorted_indices.sort(); - let merged: Vec = sorted_indices.into_iter().filter_map(|idx| { - tool_call_map.remove(&idx).map(|(id, name, arguments)| { - serde_json::json!({ - "id": id, - "type": "function", - "function": { - "name": name, - "arguments": arguments - } + let merged: Vec = sorted_indices + .into_iter() + .filter_map(|idx| { + tool_call_map.remove(&idx).map(|(id, name, arguments)| { + serde_json::json!({ + "id": id, + "type": "function", + "function": { + "name": name, + "arguments": arguments + } + }) }) }) - }).collect(); - if merged.is_empty() { None } else { Some(merged) } + .collect(); + if merged.is_empty() { + None + } else { + Some(merged) + } }; // Build aggregated response from chunks first_chunk.and_then(|first| { - serde_json::from_value::(first.clone()).ok().map(|chunk| { - let combined_content = content_parts.join(""); - let combined_reasoning = if reasoning_parts.is_empty() { - None - } else { - Some(reasoning_parts.join("")) - }; - OpenAIResponse { - id: chunk.id, - object: "chat.completion".to_string(), - created: chunk.created, - model: chunk.model, - choices: vec![OpenAIChoice { - index: 0, - message: OpenAIChatMessage { - role: MessageRole::Assistant, - content: Some(OpenAIContent::Text(combined_content)), - reasoning_content: combined_reasoning, - refusal: None, - function_call: None, - tool_calls, - tool_call_id: None, - name: None, - annotations: None, - audio: None, - }, - finish_reason, - logprobs: None, - }], - usage: None, - system_fingerprint: chunk.system_fingerprint, - } - }) + serde_json::from_value::(first.clone()) + .ok() + .map(|chunk| { + let combined_content = content_parts.join(""); + let combined_reasoning = if reasoning_parts.is_empty() { + None + } else { + Some(reasoning_parts.join("")) + }; + OpenAIResponse { + id: chunk.id, + object: "chat.completion".to_string(), + created: chunk.created, + model: chunk.model, + choices: vec![OpenAIChoice { + index: 0, + message: OpenAIChatMessage { + role: MessageRole::Assistant, + content: Some(OpenAIContent::Text(combined_content)), + reasoning_content: combined_reasoning, + refusal: None, + function_call: None, + tool_calls, + tool_call_id: None, + name: None, + annotations: None, + audio: None, + }, + finish_reason, + logprobs: None, + }], + usage: None, + system_fingerprint: chunk.system_fingerprint, + } + }) }) } @@ -391,7 +404,9 @@ mod tests { fn test_matches_path() { assert!(OpenAIParser::matches_path("/v1/chat/completions")); assert!(OpenAIParser::matches_path("/v1/completions")); - assert!(OpenAIParser::matches_path("https://api.openai.com/v1/chat/completions")); + assert!(OpenAIParser::matches_path( + "https://api.openai.com/v1/chat/completions" + )); assert!(!OpenAIParser::matches_path("/v1/messages")); assert!(!OpenAIParser::matches_path("/v1/embeddings")); } @@ -429,6 +444,9 @@ mod tests { assert!(response.is_some()); let response = response.unwrap(); - assert_eq!(response.choices[0].finish_reason, Some("tool_calls".to_string())); + assert_eq!( + response.choices[0].finish_reason, + Some("tool_calls".to_string()) + ); } } diff --git a/src/agentsight/src/analyzer/message/sysom.rs b/src/agentsight/src/analyzer/message/sysom.rs index a8dba3dca..aa7bcbafe 100644 --- a/src/agentsight/src/analyzer/message/sysom.rs +++ b/src/agentsight/src/analyzer/message/sysom.rs @@ -174,9 +174,7 @@ impl SysomParser { /// /// Returns `None` if `llmParamString` is absent or cannot be decoded. pub fn parse_request(body: &serde_json::Value) -> Option { - let llm_param_string = body - .get("llmParamString") - .and_then(|v| v.as_str())?; + let llm_param_string = body.get("llmParamString").and_then(|v| v.as_str())?; let params: SysomLlmParams = serde_json::from_str(llm_param_string) .map_err(|e| { diff --git a/src/agentsight/src/analyzer/message/types.rs b/src/agentsight/src/analyzer/message/types.rs index 2534a62f4..93841047b 100644 --- a/src/agentsight/src/analyzer/message/types.rs +++ b/src/agentsight/src/analyzer/message/types.rs @@ -16,6 +16,7 @@ use serde::{Deserialize, Serialize}; /// Unified message role across different LLM providers #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "lowercase")] +#[derive(Default)] pub enum MessageRole { /// System message (instructions/context) /// Note: Some newer OpenAI models use "developer" instead of "system" @@ -24,6 +25,7 @@ pub enum MessageRole { /// Used for developer instructions that should take precedence over user messages Developer, /// User message (human input) + #[default] User, /// Assistant message (LLM response) Assistant, @@ -31,12 +33,6 @@ pub enum MessageRole { Tool, } -impl Default for MessageRole { - fn default() -> Self { - MessageRole::User - } -} - // ============================================================================ // OpenAI Types // ============================================================================ @@ -620,9 +616,7 @@ pub struct OpenAiSseDelta { #[serde(rename_all = "snake_case")] pub enum AnthropicSseEvent { /// Message start event (contains initial message metadata) - MessageStart { - message: AnthropicSseMessageStart, - }, + MessageStart { message: AnthropicSseMessageStart }, /// Content block start event ContentBlockStart { index: u32, @@ -634,9 +628,7 @@ pub enum AnthropicSseEvent { delta: AnthropicSseDelta, }, /// Content block stop event - ContentBlockStop { - index: u32, - }, + ContentBlockStop { index: u32 }, /// Message delta event (stop_reason, usage update) MessageDelta { delta: AnthropicSseMessageDelta, @@ -647,9 +639,7 @@ pub enum AnthropicSseEvent { /// Ping event Ping, /// Error event - Error { - error: serde_json::Value, - }, + Error { error: serde_json::Value }, } /// Anthropic SSE message start data @@ -683,22 +673,16 @@ pub struct AnthropicSseMessageStart { #[serde(rename_all = "snake_case")] pub enum AnthropicSseDelta { /// Text delta - TextDelta { - text: String, - }, + TextDelta { text: String }, /// Thinking delta (extended thinking / chain-of-thought) - ThinkingDelta { - thinking: String, - }, + ThinkingDelta { thinking: String }, /// Signature delta (thinking block signature) SignatureDelta { #[serde(default)] signature: String, }, /// Input JSON delta (for tool use) - InputJsonDelta { - partial_json: String, - }, + InputJsonDelta { partial_json: String }, } /// Anthropic SSE message delta @@ -915,7 +899,9 @@ mod tests { #[test] fn test_openai_content_parts_with_image() { let parts = OpenAIContent::Parts(vec![ - OpenAIContentPart::Text { text: "Look at this:".to_string() }, + OpenAIContentPart::Text { + text: "Look at this:".to_string(), + }, OpenAIContentPart::ImageUrl { image_url: OpenAIImageUrl { url: "https://example.com/img.png".to_string(), @@ -948,7 +934,8 @@ mod tests { #[test] fn test_anthropic_system_prompt_blocks() { - let json_str = r#"[{"type": "text", "text": "Part 1"}, {"type": "text", "text": "Part 2"}]"#; + let json_str = + r#"[{"type": "text", "text": "Part 1"}, {"type": "text", "text": "Part 2"}]"#; let blocks: Vec = serde_json::from_str(json_str).unwrap(); let prompt = AnthropicSystemPrompt::Blocks(blocks); assert_eq!(prompt.as_text(), "Part 1\nPart 2"); @@ -988,8 +975,14 @@ mod tests { assert_eq!(text.as_text(), "simple"); let blocks = AnthropicMessageContent::Blocks(vec![ - AnthropicContentBlock::Text { text: "part1".to_string(), cache_control: None }, - AnthropicContentBlock::Text { text: "part2".to_string(), cache_control: None }, + AnthropicContentBlock::Text { + text: "part1".to_string(), + cache_control: None, + }, + AnthropicContentBlock::Text { + text: "part2".to_string(), + cache_control: None, + }, ]); assert_eq!(blocks.as_text(), "part1part2"); } @@ -1052,12 +1045,22 @@ mod tests { request: Some(OpenAIRequest { model: "gpt-4".to_string(), messages: vec![], - temperature: None, max_tokens: None, stream: None, - top_p: None, n: None, stop: None, - presence_penalty: None, frequency_penalty: None, - user: None, tools: None, tool_choice: None, - response_format: None, seed: None, logprobs: None, - top_logprobs: None, parallel_tool_calls: None, + temperature: None, + max_tokens: None, + stream: None, + top_p: None, + n: None, + stop: None, + presence_penalty: None, + frequency_penalty: None, + user: None, + tools: None, + tool_choice: None, + response_format: None, + seed: None, + logprobs: None, + top_logprobs: None, + parallel_tool_calls: None, }), response: None, }; @@ -1072,8 +1075,11 @@ mod tests { response: Some(OpenAIResponse { id: "chatcmpl-xyz".to_string(), object: "chat.completion".to_string(), - created: 0, model: "gpt-4".to_string(), - choices: vec![], usage: None, system_fingerprint: None, + created: 0, + model: "gpt-4".to_string(), + choices: vec![], + usage: None, + system_fingerprint: None, }), }; assert_eq!(msg.response_id(), Some("chatcmpl-xyz")); diff --git a/src/agentsight/src/analyzer/mod.rs b/src/agentsight/src/analyzer/mod.rs index 04822229e..293b32d8b 100644 --- a/src/agentsight/src/analyzer/mod.rs +++ b/src/agentsight/src/analyzer/mod.rs @@ -8,26 +8,30 @@ pub mod audit; pub mod message; -pub mod token; mod result; +pub mod token; mod unified; // Re-export audit types pub use audit::{AuditAnalyzer, AuditEventType, AuditExtra, AuditRecord, AuditSummary}; // Re-export token types from the token module -pub use token::{TokenParser, TokenUsage, TokenRecord, LLMProvider}; +pub use token::{LLMProvider, TokenParser, TokenRecord, TokenUsage}; // Re-export message types from the message module pub use message::{ - MessageParser, ParsedApiMessage, - OpenAIRequest, OpenAIResponse, OpenAIChatMessage, OpenAIContent, OpenAIUsage, OpenAIChoice, - AnthropicRequest, AnthropicResponse, AnthropicMessage, AnthropicUsage, - MessageRole, + AnthropicMessage, AnthropicRequest, AnthropicResponse, AnthropicUsage, MessageParser, + MessageRole, OpenAIChatMessage, OpenAIChoice, OpenAIContent, OpenAIRequest, OpenAIResponse, + OpenAIUsage, ParsedApiMessage, }; // Re-export analysis result -pub use result::{AnalysisResult, PromptTokenCount, HttpRecord, TokenConsumptionBreakdown, MessageTokenCount, OutputTokenCount}; +pub use result::{ + AnalysisResult, HttpRecord, MessageTokenCount, OutputTokenCount, PromptTokenCount, + TokenConsumptionBreakdown, +}; // Re-export unified analyzer -pub use unified::{Analyzer, count_request_tokens, count_response_tokens, RequestTokenCount, ResponseTokenCount}; +pub use unified::{ + Analyzer, RequestTokenCount, ResponseTokenCount, count_request_tokens, count_response_tokens, +}; diff --git a/src/agentsight/src/analyzer/result.rs b/src/agentsight/src/analyzer/result.rs index fa5fc8440..de9787f32 100644 --- a/src/agentsight/src/analyzer/result.rs +++ b/src/agentsight/src/analyzer/result.rs @@ -6,7 +6,7 @@ use serde::{Deserialize, Serialize}; use std::collections::HashMap; -use super::{AuditRecord, TokenRecord, ParsedApiMessage}; +use super::{AuditRecord, ParsedApiMessage, TokenRecord}; /// Computed prompt token count for a request #[derive(Debug, Clone)] @@ -113,6 +113,7 @@ pub struct MessageTokenCount { /// Unified analysis result from different analyzers #[derive(Debug, Clone)] +#[allow(clippy::large_enum_variant)] pub enum AnalysisResult { /// Audit record from AuditAnalyzer Audit(AuditRecord), diff --git a/src/agentsight/src/analyzer/token/data.rs b/src/agentsight/src/analyzer/token/data.rs index b1c4a836a..5a1d8637b 100644 --- a/src/agentsight/src/analyzer/token/data.rs +++ b/src/agentsight/src/analyzer/token/data.rs @@ -45,7 +45,11 @@ impl TokenData { } /// Add a request message - pub fn add_request_message(mut self, role: impl Into, content: impl Into) -> Self { + pub fn add_request_message( + mut self, + role: impl Into, + content: impl Into, + ) -> Self { self.request_messages.push(MessageTokenData { role: role.into(), content: content.into(), @@ -88,38 +92,38 @@ impl TokenData { /// Get all request text content combined pub fn request_text(&self) -> String { let mut parts = Vec::new(); - + if let Some(ref system) = self.system_prompt { parts.push(format!("system: {}", system)); } - + for msg in &self.request_messages { parts.push(format!("{}: {}", msg.role, msg.content)); } - + for tool in &self.tools { parts.push(format!("tool: {}", tool)); } - + parts.join("\n") } /// Get all response text content combined pub fn response_text(&self) -> String { let mut parts = Vec::new(); - + if let Some(ref reasoning) = self.reasoning_content { parts.push(format!("reasoning: {}", reasoning)); } - + for content in &self.response_content { parts.push(content.content.clone()); } - + for tool_call in &self.tool_calls { parts.push(format!("tool_call: {}", tool_call)); } - + parts.join("\n") } @@ -130,27 +134,24 @@ impl TokenData { /// Get messages grouped by role pub fn messages_by_role(&self) -> std::collections::HashMap> { - let mut map: std::collections::HashMap> = + let mut map: std::collections::HashMap> = std::collections::HashMap::new(); - + for msg in &self.request_messages { - map.entry(msg.role.clone()) - .or_insert_with(Vec::new) - .push(msg); + map.entry(msg.role.clone()).or_default().push(msg); } - + map } /// Count messages by role pub fn count_by_role(&self) -> std::collections::HashMap { - let mut counts: std::collections::HashMap = - std::collections::HashMap::new(); - + let mut counts: std::collections::HashMap = std::collections::HashMap::new(); + for msg in &self.request_messages { *counts.entry(msg.role.clone()).or_insert(0) += 1; } - + counts } @@ -162,31 +163,31 @@ impl TokenData { /// Get total character count (rough estimate for token calculation) pub fn total_chars(&self) -> usize { let mut total = 0; - + if let Some(ref system) = self.system_prompt { total += system.len(); } - + for msg in &self.request_messages { total += msg.content.len(); } - + for tool in &self.tools { total += tool.len(); } - + for content in &self.response_content { total += content.content.len(); } - + if let Some(ref reasoning) = self.reasoning_content { total += reasoning.len(); } - + for tool_call in &self.tool_calls { total += tool_call.len(); } - + total } } @@ -220,7 +221,10 @@ mod tests { assert_eq!(data.provider, "openai"); assert_eq!(data.model, "gpt-4"); - assert_eq!(data.system_prompt, Some("You are a helpful assistant".to_string())); + assert_eq!( + data.system_prompt, + Some("You are a helpful assistant".to_string()) + ); assert_eq!(data.request_messages.len(), 1); assert_eq!(data.response_content.len(), 1); } @@ -242,24 +246,22 @@ mod tests { #[test] fn test_add_tool() { - let data = TokenData::new("openai", "gpt-4") - .add_tool(r#"{"name":"search"}"#); + let data = TokenData::new("openai", "gpt-4").add_tool(r#"{"name":"search"}"#); assert_eq!(data.tools.len(), 1); assert!(data.request_text().contains("tool: {\"name\":\"search\"}")); } #[test] fn test_with_reasoning_content() { - let data = TokenData::new("openai", "qwen") - .with_reasoning_content("Let me think..."); + let data = TokenData::new("openai", "qwen").with_reasoning_content("Let me think..."); assert_eq!(data.reasoning_content, Some("Let me think...".to_string())); assert!(data.response_text().contains("reasoning: Let me think...")); } #[test] fn test_add_tool_call() { - let data = TokenData::new("openai", "gpt-4") - .add_tool_call(r#"get_weather({"city":"Beijing"})"#); + let data = + TokenData::new("openai", "gpt-4").add_tool_call(r#"get_weather({"city":"Beijing"})"#); assert_eq!(data.tool_calls.len(), 1); assert!(data.response_text().contains("tool_call: get_weather")); } @@ -309,12 +311,12 @@ mod tests { #[test] fn test_total_chars() { let data = TokenData::new("openai", "gpt-4") - .with_system_prompt("abc") // 3 + .with_system_prompt("abc") // 3 .add_request_message("user", "de") // 2 - .add_tool("fg") // 2 - .add_response_content("hij") // 3 - .with_reasoning_content("kl") // 2 - .add_tool_call("mno"); // 3 + .add_tool("fg") // 2 + .add_response_content("hij") // 3 + .with_reasoning_content("kl") // 2 + .add_tool_call("mno"); // 3 assert_eq!(data.total_chars(), 15); } @@ -339,8 +341,7 @@ mod tests { #[test] fn test_request_text_no_system() { - let data = TokenData::new("openai", "gpt-4") - .add_request_message("user", "Hello"); + let data = TokenData::new("openai", "gpt-4").add_request_message("user", "Hello"); let text = data.request_text(); assert!(!text.contains("system:")); assert!(text.contains("user: Hello")); diff --git a/src/agentsight/src/analyzer/token/extractor/anthropic.rs b/src/agentsight/src/analyzer/token/extractor/anthropic.rs index 68830b77f..959707ecf 100644 --- a/src/agentsight/src/analyzer/token/extractor/anthropic.rs +++ b/src/agentsight/src/analyzer/token/extractor/anthropic.rs @@ -1,8 +1,8 @@ //! Anthropic token data extraction -use serde_json::Value; -use super::super::data::{TokenData, MessageTokenData, ResponseTokenData}; +use super::super::data::{MessageTokenData, ResponseTokenData, TokenData}; use super::utils::extract_model_from_json; +use serde_json::Value; /// Extract token data from Anthropic format JSON pub fn extract_token_data( @@ -41,7 +41,9 @@ pub fn extract_token_data( if let Some(messages) = req.get("messages").and_then(|m| m.as_array()) { for msg in messages { if let Some((role, content)) = extract_message(msg) { - token_data.request_messages.push(MessageTokenData { role, content }); + token_data + .request_messages + .push(MessageTokenData { role, content }); has_content = true; } } @@ -64,25 +66,30 @@ pub fn extract_token_data( if let Some(content) = resp.get("content").and_then(|c| c.as_array()) { for block in content { let block_type = block.get("type").and_then(|t| t.as_str()); - + match block_type { Some("text") => { - if let Some(text) = block.get("text").and_then(|t| t.as_str()) { - if !text.is_empty() { - token_data.response_content.push(ResponseTokenData { - content: text.to_string(), - }); - has_content = true; - } + if let Some(text) = block.get("text").and_then(|t| t.as_str()) + && !text.is_empty() + { + token_data.response_content.push(ResponseTokenData { + content: text.to_string(), + }); + has_content = true; } } Some("tool_use") => { - let name = block.get("name").and_then(|n| n.as_str()).unwrap_or("unknown"); - if let Some(input) = block.get("input") { - if let Ok(input_str) = serde_json::to_string(input) { - token_data.tool_calls.push(format!("{}: {}", name, input_str)); - has_content = true; - } + let name = block + .get("name") + .and_then(|n| n.as_str()) + .unwrap_or("unknown"); + if let Some(input) = block.get("input") + && let Ok(input_str) = serde_json::to_string(input) + { + token_data + .tool_calls + .push(format!("{}: {}", name, input_str)); + has_content = true; } } _ => {} @@ -91,18 +98,14 @@ pub fn extract_token_data( } } - if has_content { - Some(token_data) - } else { - None - } + if has_content { Some(token_data) } else { None } } /// Extract role and content from Anthropic message JSON fn extract_message(msg: &Value) -> Option<(String, String)> { let role = msg.get("role").and_then(|r| r.as_str())?; let content = extract_content(msg.get("content"))?; - + if content.is_empty() { None } else { @@ -132,12 +135,8 @@ fn extract_content(content: Option<&Value>) -> Option { }) .collect::>() .join(""); - - if text.is_empty() { - None - } else { - Some(text) - } + + if text.is_empty() { None } else { Some(text) } } _ => None, } @@ -199,7 +198,10 @@ mod tests { assert!(token_data.is_some()); let data = token_data.unwrap(); - assert_eq!(data.system_prompt, Some("You are Claude\nBe helpful".to_string())); + assert_eq!( + data.system_prompt, + Some("You are Claude\nBe helpful".to_string()) + ); } #[test] diff --git a/src/agentsight/src/analyzer/token/extractor/mod.rs b/src/agentsight/src/analyzer/token/extractor/mod.rs index 11d63af73..e3477bc7d 100644 --- a/src/agentsight/src/analyzer/token/extractor/mod.rs +++ b/src/agentsight/src/analyzer/token/extractor/mod.rs @@ -13,12 +13,12 @@ //! //! Note: This module is for internal use only and not exposed in the public API. -pub mod openai; mod anthropic; +pub mod openai; mod utils; -use serde_json::Value; use super::data::TokenData; +use serde_json::Value; /// Extract token data from JSON request/response bodies /// diff --git a/src/agentsight/src/analyzer/token/extractor/openai.rs b/src/agentsight/src/analyzer/token/extractor/openai.rs index 8d96a131a..973c85406 100644 --- a/src/agentsight/src/analyzer/token/extractor/openai.rs +++ b/src/agentsight/src/analyzer/token/extractor/openai.rs @@ -1,8 +1,8 @@ //! OpenAI token data extraction -use serde_json::Value; -use super::super::data::{TokenData, MessageTokenData, ResponseTokenData}; +use super::super::data::{MessageTokenData, ResponseTokenData, TokenData}; use super::utils::extract_model_from_json; +use serde_json::Value; /// Extract token data from OpenAI format JSON pub fn extract_token_data( @@ -21,7 +21,9 @@ pub fn extract_token_data( if let Some(messages) = req.get("messages").and_then(|m| m.as_array()) { for msg in messages { if let Some((role, content)) = extract_message(msg) { - token_data.request_messages.push(MessageTokenData { role, content }); + token_data + .request_messages + .push(MessageTokenData { role, content }); has_content = true; } } @@ -41,7 +43,9 @@ pub fn extract_token_data( // Extract from response using shared logic if let Some((content, reasoning, tool_calls)) = extract_response_content(response_json) { if !content.is_empty() { - token_data.response_content.push(ResponseTokenData { content }); + token_data + .response_content + .push(ResponseTokenData { content }); has_content = true; } if let Some(r) = reasoning { @@ -54,15 +58,11 @@ pub fn extract_token_data( } } - if has_content { - Some(token_data) - } else { - None - } + if has_content { Some(token_data) } else { None } } /// Extract response content from OpenAI format response JSON -/// +/// /// Returns a tuple of (content, reasoning_content, tool_calls) /// - content: The main response text /// - reasoning_content: Optional reasoning/thinking content @@ -72,7 +72,7 @@ pub fn extract_response_content( ) -> Option<(String, Option, Vec)> { let resp = response_json?; let choices = resp.get("choices").and_then(|c| c.as_array())?; - + let mut content = String::new(); let mut reasoning = None; let mut tool_calls = Vec::new(); @@ -81,26 +81,26 @@ pub fn extract_response_content( for choice in choices { // Support both "message" (standard response) and "delta" (SSE streaming) formats let msg_or_delta = choice.get("message").or_else(|| choice.get("delta")); - + if let Some(msg) = msg_or_delta { // Extract content - if let Some(c) = msg.get("content").and_then(|c| c.as_str()) { - if !c.is_empty() { - content.push_str(c); - has_data = true; - } + if let Some(c) = msg.get("content").and_then(|c| c.as_str()) + && !c.is_empty() + { + content.push_str(c); + has_data = true; } // Extract reasoning_content - if let Some(r) = msg.get("reasoning_content").and_then(|r| r.as_str()) { - if !r.is_empty() { - // For SSE chunks, accumulate reasoning content - reasoning = match reasoning { - Some(existing) => Some(existing + r), - None => Some(r.to_string()), - }; - has_data = true; - } + if let Some(r) = msg.get("reasoning_content").and_then(|r| r.as_str()) + && !r.is_empty() + { + // For SSE chunks, accumulate reasoning content + reasoning = match reasoning { + Some(existing) => Some(existing + r), + None => Some(r.to_string()), + }; + has_data = true; } // Extract tool_calls - only extract function name and arguments @@ -108,7 +108,8 @@ pub fn extract_response_content( for tool_call in calls { if let Some(func) = tool_call.get("function") { let name = func.get("name").and_then(|n| n.as_str()).unwrap_or(""); - let arguments = func.get("arguments").and_then(|a| a.as_str()).unwrap_or(""); + let arguments = + func.get("arguments").and_then(|a| a.as_str()).unwrap_or(""); let tool_content = format!("{}: {}", name, arguments); if !tool_content.is_empty() { tool_calls.push(tool_content); @@ -131,7 +132,7 @@ pub fn extract_response_content( fn extract_message(msg: &Value) -> Option<(String, String)> { let role = msg.get("role").and_then(|r| r.as_str())?; let content = extract_content(msg.get("content"))?; - + if content.is_empty() { None } else { @@ -161,12 +162,8 @@ fn extract_content(content: Option<&Value>) -> Option { }) .collect::>() .join(""); - - if text.is_empty() { - None - } else { - Some(text) - } + + if text.is_empty() { None } else { Some(text) } } _ => None, } @@ -255,7 +252,10 @@ mod tests { assert!(token_data.is_some()); let data = token_data.unwrap(); - assert_eq!(data.reasoning_content, Some("Let me think about this...".to_string())); + assert_eq!( + data.reasoning_content, + Some("Let me think about this...".to_string()) + ); } #[test] diff --git a/src/agentsight/src/analyzer/token/extractor/utils.rs b/src/agentsight/src/analyzer/token/extractor/utils.rs index a2d02362d..4d0ac23de 100644 --- a/src/agentsight/src/analyzer/token/extractor/utils.rs +++ b/src/agentsight/src/analyzer/token/extractor/utils.rs @@ -1,7 +1,7 @@ //! Utility functions for token data extraction -use serde_json::Value; use super::Provider; +use serde_json::Value; /// Detect provider from API endpoint path and/or JSON content /// @@ -163,7 +163,10 @@ mod tests { "model": "gpt-4", "messages": [{"role": "user", "content": "Hello"}] }); - let provider = detect_provider("/compatible-mode/v1/chat/completions", Some(&openai_request)); + let provider = detect_provider( + "/compatible-mode/v1/chat/completions", + Some(&openai_request), + ); assert_eq!(provider, Some(Provider::OpenAI)); // Compatible mode path with Anthropic format JSON @@ -173,7 +176,10 @@ mod tests { "system": "You are Claude", "messages": [{"role": "user", "content": "Hello"}] }); - let provider = detect_provider("/compatible-mode/v1/chat/completions", Some(&anthropic_request)); + let provider = detect_provider( + "/compatible-mode/v1/chat/completions", + Some(&anthropic_request), + ); assert_eq!(provider, Some(Provider::Anthropic)); } diff --git a/src/agentsight/src/analyzer/token/mod.rs b/src/agentsight/src/analyzer/token/mod.rs index 9033051ea..e2a055221 100644 --- a/src/agentsight/src/analyzer/token/mod.rs +++ b/src/agentsight/src/analyzer/token/mod.rs @@ -23,9 +23,9 @@ //! } //! ``` -mod record; mod data; mod parser; +mod record; // Extractor submodule for JSON token data extraction mod extractor; @@ -125,12 +125,18 @@ pub fn extract_usage_object( let (input_tokens, output_tokens) = match provider { LLMProvider::OpenAI => { let input = usage.get("prompt_tokens").and_then(|v| v.as_u64())?; - let output = usage.get("completion_tokens").and_then(|v| v.as_u64()).unwrap_or(0); + let output = usage + .get("completion_tokens") + .and_then(|v| v.as_u64()) + .unwrap_or(0); (input, output) } LLMProvider::Anthropic => { let input = usage.get("input_tokens").and_then(|v| v.as_u64())?; - let output = usage.get("output_tokens").and_then(|v| v.as_u64()).unwrap_or(0); + let output = usage + .get("output_tokens") + .and_then(|v| v.as_u64()) + .unwrap_or(0); (input, output) } LLMProvider::Gemini => { @@ -212,9 +218,7 @@ pub fn detect_provider_from_endpoint(endpoint: Option<&str>) -> LLMProvider { Some(ep) if ep.contains("anthropic.com") || ep.contains("api.anthropic.com") => { LLMProvider::Anthropic } - Some(ep) if ep.contains("generativelanguage.googleapis.com") - || ep.contains("gemini") => - { + Some(ep) if ep.contains("generativelanguage.googleapis.com") || ep.contains("gemini") => { LLMProvider::Gemini } _ => LLMProvider::Unknown, diff --git a/src/agentsight/src/analyzer/token/parser.rs b/src/agentsight/src/analyzer/token/parser.rs index abc5b0c3b..760a5f978 100644 --- a/src/agentsight/src/analyzer/token/parser.rs +++ b/src/agentsight/src/analyzer/token/parser.rs @@ -21,7 +21,7 @@ //! } //! ``` -use super::{detect_provider_from_usage, extract_usage_object, LLMProvider, TokenUsage}; +use super::{LLMProvider, TokenUsage, detect_provider_from_usage, extract_usage_object}; use crate::parser::sse::ParsedSseEvent; /// Token parser for extracting usage from SSE events @@ -70,19 +70,18 @@ impl TokenParser { /// Internal method to parse JSON and extract token usage pub fn parse_json(&self, json: &serde_json::Value) -> Option { // 1. Check for message_start event (Anthropic streaming) - if json.get("type").and_then(|v| v.as_str()) == Some("message_start") { - if let Some(message) = json.get("message") { - if let Some(usage) = message.get("usage") { - return extract_usage_object(usage, LLMProvider::Anthropic, json); - } - } + if json.get("type").and_then(|v| v.as_str()) == Some("message_start") + && let Some(message) = json.get("message") + && let Some(usage) = message.get("usage") + { + return extract_usage_object(usage, LLMProvider::Anthropic, json); } // 2. Check for message_delta event (Anthropic streaming final) - if json.get("type").and_then(|v| v.as_str()) == Some("message_delta") { - if let Some(usage) = json.get("usage") { - return extract_usage_object(usage, LLMProvider::Anthropic, json); - } + if json.get("type").and_then(|v| v.as_str()) == Some("message_delta") + && let Some(usage) = json.get("usage") + { + return extract_usage_object(usage, LLMProvider::Anthropic, json); } // 3. Check for usage object directly (OpenAI and compatible APIs) @@ -214,7 +213,10 @@ mod tests { let event = create_test_event(data); let usage = parser.parse_event(&event); - assert!(usage.is_some(), "Should extract usage from SSE streaming data"); + assert!( + usage.is_some(), + "Should extract usage from SSE streaming data" + ); let usage = usage.unwrap(); assert_eq!(usage.input_tokens, 61744); diff --git a/src/agentsight/src/analyzer/token/record.rs b/src/agentsight/src/analyzer/token/record.rs index ab437c535..cf36e916b 100644 --- a/src/agentsight/src/analyzer/token/record.rs +++ b/src/agentsight/src/analyzer/token/record.rs @@ -156,8 +156,8 @@ mod tests { #[test] fn test_with_request_id() { - let record = TokenRecord::new(1, "p".to_string(), "o".to_string(), 0, 0) - .with_request_id("req-123"); + let record = + TokenRecord::new(1, "p".to_string(), "o".to_string(), 0, 0).with_request_id("req-123"); assert_eq!(record.request_id, Some("req-123".to_string())); } diff --git a/src/agentsight/src/analyzer/unified.rs b/src/agentsight/src/analyzer/unified.rs index 2fec8d5a3..09625d8b1 100644 --- a/src/agentsight/src/analyzer/unified.rs +++ b/src/agentsight/src/analyzer/unified.rs @@ -9,7 +9,7 @@ //! use agentsight::aggregator::AggregatedResult; //! //! let analyzer = Analyzer::new(); -//! +//! //! // Analyze aggregated result //! for result in analyzer.analyze_aggregated(&aggregated_result) { //! match result { @@ -21,13 +21,16 @@ //! ``` use crate::aggregator::AggregatedResult; +use crate::analyzer::token::extract_response_content; use crate::parser::sse::ParsedSseEvent; use crate::tokenizer::LlmTokenizer; use crate::tokenizer::get_global_tokenizer; -use crate::analyzer::token::extract_response_content; -use super::{AuditAnalyzer, TokenParser, MessageParser, TokenRecord, TokenUsage, ParsedApiMessage, AnalysisResult, HttpRecord}; -use super::result::{TokenConsumptionBreakdown, MessageTokenCount, OutputTokenCount}; +use super::result::{MessageTokenCount, OutputTokenCount, TokenConsumptionBreakdown}; +use super::{ + AnalysisResult, AuditAnalyzer, HttpRecord, MessageParser, ParsedApiMessage, TokenParser, + TokenRecord, TokenUsage, +}; /// Token count result for request messages #[derive(Debug, Clone)] @@ -83,9 +86,8 @@ pub fn count_request_tokens( chat_template: &LlmTokenizer, ) -> Option { // Extract messages - let messages = request_json.get("messages") - .and_then(|m| m.as_array())?; - + let messages = request_json.get("messages").and_then(|m| m.as_array())?; + if messages.is_empty() { return None; } @@ -97,29 +99,34 @@ pub fn count_request_tokens( let template_messages: Vec = messages.to_vec(); // Extract tools JSON array for passing to template - let tools_json: Option> = request_json.get("tools") + let tools_json: Option> = request_json + .get("tools") .and_then(|t| t.as_array()) .map(|arr| arr.to_vec()); // Count tools tokens separately (for informational breakdown) - let mut tools_tokens: usize = tools_json.as_ref() - .map(|arr| arr.iter() - .filter_map(|t| serde_json::to_string(t).ok()) - .filter_map(|s| tokenizer.count(&s).ok()) - .sum()) + let mut tools_tokens: usize = tools_json + .as_ref() + .map(|arr| { + arr.iter() + .filter_map(|t| serde_json::to_string(t).ok()) + .filter_map(|s| tokenizer.count(&s).ok()) + .sum() + }) .unwrap_or(0); // Use apply_chat_template_with_tools to format all messages WITH tools // This ensures the tools instruction text is included in the total count let tools_slice = tools_json.as_deref(); - let total_tokens = match chat_template.apply_chat_template_with_tools(&template_messages, tools_slice, true) { - Ok(formatted) => tokenizer.count(&formatted).unwrap_or(0), - Err(e) => { - log::warn!("Failed to apply chat template with tools: {}", e); - // Fallback: count raw content + tools separately - 0 - } - }; + let total_tokens = + match chat_template.apply_chat_template_with_tools(&template_messages, tools_slice, true) { + Ok(formatted) => tokenizer.count(&formatted).unwrap_or(0), + Err(e) => { + log::warn!("Failed to apply chat template with tools: {}", e); + // Fallback: count raw content + tools separately + 0 + } + }; // Count per-message tokens: first calculate raw token counts, then distribute total_tokens by percentage // Step 1: Calculate raw token count for each message @@ -130,21 +137,26 @@ pub fn count_request_tokens( .ok() .and_then(|s| tokenizer.count(&s).ok()) .unwrap_or(0); - let role = msg.get("role").and_then(|r| r.as_str()).unwrap_or("unknown").to_string(); + let role = msg + .get("role") + .and_then(|r| r.as_str()) + .unwrap_or("unknown") + .to_string(); if role == "tool" { tokens += tools_tokens; tools_tokens = 0; } raw_per_message.push((role, tokens)); } - + // Step 2: Calculate total raw tokens and distribute total_tokens by percentage let raw_total: usize = raw_per_message.iter().map(|(_, t)| *t).sum(); - + if raw_total > 0 { // Distribute total_tokens proportionally based on raw token percentages for (role, raw_tokens) in raw_per_message.iter() { - let actual_tokens = ((*raw_tokens as f64 / raw_total as f64) * total_tokens as f64).round() as usize; + let actual_tokens = + ((*raw_tokens as f64 / raw_total as f64) * total_tokens as f64).round() as usize; *by_role.entry(role.clone()).or_insert(0) += actual_tokens; per_message.push(MessageTokenCount { role: role.clone(), @@ -213,10 +225,10 @@ pub fn count_response_tokens( if !content.is_empty() { all_content.push_str(&content); } - if let Some(r) = reasoning { - if !r.is_empty() { - all_reasoning.push_str(&r); - } + if let Some(r) = reasoning + && !r.is_empty() + { + all_reasoning.push_str(&r); } for tc in tool_calls { if !tc.is_empty() { @@ -227,18 +239,20 @@ pub fn count_response_tokens( } let mut has_content = false; - + // NOTE: API output token count includes: // - Model-generated text markers like ... (these ARE counted) // - NOT special control tokens like <|im_start|>, <|im_end|> (these are NOT counted) - + // Add reasoning content with wrapper (model generates these markers) if !all_reasoning.is_empty() { has_content = true; - + // Format: \n{reasoning}\n\n\n let reasoning_with_tags = format!("\n{}\n\n\n", all_reasoning); - let tokens = tokenizer.count(&reasoning_with_tags).unwrap_or(all_reasoning.len() / 4); + let tokens = tokenizer + .count(&reasoning_with_tags) + .unwrap_or(all_reasoning.len() / 4); *by_type.entry("reasoning".to_string()).or_insert(0) += tokens; per_block.push(OutputTokenCount { content_type: "reasoning".to_string(), @@ -249,8 +263,10 @@ pub fn count_response_tokens( // Add text content if !all_content.is_empty() { has_content = true; - - let tokens = tokenizer.count(&all_content).unwrap_or(all_content.len() / 4); + + let tokens = tokenizer + .count(&all_content) + .unwrap_or(all_content.len() / 4); *by_type.entry("text".to_string()).or_insert(0) += tokens; per_block.push(OutputTokenCount { content_type: "text".to_string(), @@ -267,7 +283,6 @@ pub fn count_response_tokens( // - Following chunks: ": {...}" (colon + arguments fragments) // We need to aggregate all chunks first to get the complete tool call if !all_tool_calls.is_empty() { - // Aggregate all chunks: first chunk has "name: ", rest have ": fragment" let mut aggregated = String::new(); for tc in &all_tool_calls { @@ -280,7 +295,7 @@ pub fn count_response_tokens( aggregated.push_str(fragment); } } - + // Now parse the aggregated "name: arguments" string let (name, arguments) = if let Some(pos) = aggregated.find(": ") { (&aggregated[..pos], &aggregated[pos + 2..]) @@ -290,7 +305,7 @@ pub fn count_response_tokens( } else { ("", aggregated.as_str()) }; - + // Build Qwen tool_call template format: // // @@ -303,7 +318,7 @@ pub fn count_response_tokens( tool_call_str.push_str("\n\n"); - + // Parse arguments JSON and format each parameter if let Ok(args_json) = serde_json::from_str::(arguments) { if let Some(obj) = args_json.as_object() { @@ -324,14 +339,16 @@ pub fn count_response_tokens( } else { // Fallback: use raw arguments string tool_call_str.push_str(arguments); - tool_call_str.push_str("\n"); + tool_call_str.push('\n'); } - + tool_call_str.push_str("\n"); - + has_content = true; - - let tokens = tokenizer.count(&tool_call_str).unwrap_or(tool_call_str.len() / 4); + + let tokens = tokenizer + .count(&tool_call_str) + .unwrap_or(tool_call_str.len() / 4); *by_type.entry("tool_calls".to_string()).or_insert(0) += tokens; per_block.push(OutputTokenCount { content_type: "tool_calls".to_string(), @@ -401,10 +418,7 @@ impl Analyzer { /// let chat_template = ChatTemplateType::Qwen.create_template(); /// let analyzer = Analyzer::with_tokenizer(Box::new(tokenizer), chat_template); /// ``` - pub fn with_tokenizer( - tokenizer: LlmTokenizer, - chat_template: LlmTokenizer, - ) -> Self { + pub fn with_tokenizer(tokenizer: LlmTokenizer, chat_template: LlmTokenizer) -> Self { Analyzer { audit: AuditAnalyzer::new(), token: TokenParser::new(), @@ -459,42 +473,41 @@ impl Analyzer { // 4. HTTP data export - extract raw HTTP request/response data if let Some(http_record) = self.extract_http_record(result) { - if token_result.is_none() && http_record.is_sse { - if let Some(body) = &http_record.response_body { - if let Ok(x) = serde_json::from_str::>(body) { - if let Some(last) = x.last() { - let parser = TokenParser::new(); - if let Some(usage) = parser.parse_json(last) { - let record = TokenRecord::new( - http_record.pid, - http_record.comm.clone(), - usage.provider.to_string(), - usage.input_tokens, - usage.output_tokens, - ) - .with_model(usage.model.clone().unwrap_or_default()) - .with_cache_tokens( - usage.cache_creation_input_tokens.unwrap_or(0), - usage.cache_read_input_tokens.unwrap_or(0), - ); - - token_result = Some(record); - } - } - } + if token_result.is_none() + && http_record.is_sse + && let Some(body) = &http_record.response_body + && let Ok(x) = serde_json::from_str::>(body) + && let Some(last) = x.last() + { + let parser = TokenParser::new(); + if let Some(usage) = parser.parse_json(last) { + let record = TokenRecord::new( + http_record.pid, + http_record.comm.clone(), + usage.provider.to_string(), + usage.input_tokens, + usage.output_tokens, + ) + .with_model(usage.model.clone().unwrap_or_default()) + .with_cache_tokens( + usage.cache_creation_input_tokens.unwrap_or(0), + usage.cache_read_input_tokens.unwrap_or(0), + ); + + token_result = Some(record); } } // Extract audit from HttpRecord (only for SSE responses / LLM calls) // Pass token_result so audit record gets populated token counts - if let Some(audit_record) = self.audit.analyze_http(&http_record, token_result.as_ref()) { + if let Some(audit_record) = self.audit.analyze_http(&http_record, token_result.as_ref()) + { results.push(AnalysisResult::Audit(audit_record)); } results.push(AnalysisResult::Http(http_record)); } - if let Some(record) = token_result { results.push(AnalysisResult::Token(record)); } else { @@ -545,7 +558,8 @@ impl Analyzer { request_body: Option<&serde_json::Value>, sse_events: &[ParsedSseEvent], ) -> Option { - self.message.parse_by_path_with_sse(path, request_body, sse_events) + self.message + .parse_by_path_with_sse(path, request_body, sse_events) .map(AnalysisResult::Message) } @@ -556,7 +570,9 @@ impl Analyzer { pid: u32, comm: &str, ) -> Option { - let usage = sse_events.iter().rev() + let usage = sse_events + .iter() + .rev() .find_map(|e| self.token.parse_event(e))?; let record = TokenRecord::new( @@ -575,7 +591,7 @@ impl Analyzer { // NOTE: tool_calls and reasoning_content extraction from SSE events // is handled in genai::builder via direct SSE response body parsing. - if record.total_tokens() ==0 { + if record.total_tokens() == 0 { return None; } @@ -605,7 +621,8 @@ impl Analyzer { let request_json_ref = request_json.as_ref()?; // Extract model name - let model = request_json_ref.get("model") + let model = request_json_ref + .get("model") .and_then(|m| m.as_str()) .unwrap_or("unknown"); @@ -626,56 +643,64 @@ impl Analyzer { }; // Count input tokens from request messages using chat template - let input_tokens = if let Some(messages) = request_json_ref.get("messages").and_then(|m| m.as_array()) { - if messages.is_empty() { - 0 - } else { - // Clone messages for in-place modification of tool_calls.arguments - let mut msgs = messages.clone(); - - // Process tool_calls arguments: parse JSON string to object in place - for msg in msgs.iter_mut() { - if let Some(tool_calls) = msg.get_mut("tool_calls").and_then(|tc| tc.as_array_mut()) { - for tool_call in tool_calls.iter_mut() { - if let Some(func) = tool_call.get_mut("function") { - if let Some(args) = func.get("arguments") { - if let Some(args_str) = args.as_str() { - // Try to parse arguments string as JSON object - if let Ok(parsed) = serde_json::from_str::(args_str) { - func["arguments"] = parsed; - } + let input_tokens = + if let Some(messages) = request_json_ref.get("messages").and_then(|m| m.as_array()) { + if messages.is_empty() { + 0 + } else { + // Clone messages for in-place modification of tool_calls.arguments + let mut msgs = messages.clone(); + + // Process tool_calls arguments: parse JSON string to object in place + for msg in msgs.iter_mut() { + if let Some(tool_calls) = + msg.get_mut("tool_calls").and_then(|tc| tc.as_array_mut()) + { + for tool_call in tool_calls.iter_mut() { + if let Some(func) = tool_call.get_mut("function") + && let Some(args) = func.get("arguments") + && let Some(args_str) = args.as_str() + { + // Try to parse arguments string as JSON object + if let Ok(parsed) = + serde_json::from_str::(args_str) + { + func["arguments"] = parsed; } } } } } - } - - // Extract tools JSON array for passing to template - let tools_json: Option> = request_json_ref.get("tools") - .and_then(|t| t.as_array()) - .map(|arr| arr.to_vec()); - let tools_slice = tools_json.as_deref(); - - // Apply chat template with tools to get the actual prompt sent to LLM - match tokenizer.apply_chat_template_with_tools(&msgs, tools_slice, true) { - Ok(formatted) => tokenizer.count(&formatted).unwrap_or(0) as u64, - Err(e) => { - log::warn!("Failed to apply chat template: {}, falling back to raw count", e); - // Fallback: count raw message content - let mut total = 0u64; - for msg in &msgs { - if let Ok(msg_str) = serde_json::to_string(msg) { - total += tokenizer.count(&msg_str).unwrap_or(0) as u64; + + // Extract tools JSON array for passing to template + let tools_json: Option> = request_json_ref + .get("tools") + .and_then(|t| t.as_array()) + .map(|arr| arr.to_vec()); + let tools_slice = tools_json.as_deref(); + + // Apply chat template with tools to get the actual prompt sent to LLM + match tokenizer.apply_chat_template_with_tools(&msgs, tools_slice, true) { + Ok(formatted) => tokenizer.count(&formatted).unwrap_or(0) as u64, + Err(e) => { + log::warn!( + "Failed to apply chat template: {}, falling back to raw count", + e + ); + // Fallback: count raw message content + let mut total = 0u64; + for msg in &msgs { + if let Ok(msg_str) = serde_json::to_string(msg) { + total += tokenizer.count(&msg_str).unwrap_or(0) as u64; + } } + total } - total } } - } - } else { - 0 - }; + } else { + 0 + }; // Count output tokens from SSE events content let output_tokens = { @@ -684,20 +709,21 @@ impl Analyzer { let mut all_tool_calls = Vec::new(); for event in sse_events { - if let Some(chunk) = event.json_body() { - if let Some((content, reasoning, tool_calls)) = extract_response_content(Some(&chunk)) { - if !content.is_empty() { - all_content.push_str(&content); - } - if let Some(r) = reasoning { - if !r.is_empty() { - all_reasoning.push_str(&r); - } - } - for tc in tool_calls { - if !tc.is_empty() { - all_tool_calls.push(tc); - } + if let Some(chunk) = event.json_body() + && let Some((content, reasoning, tool_calls)) = + extract_response_content(Some(&chunk)) + { + if !content.is_empty() { + all_content.push_str(&content); + } + if let Some(r) = reasoning + && !r.is_empty() + { + all_reasoning.push_str(&r); + } + for tc in tool_calls { + if !tc.is_empty() { + all_tool_calls.push(tc); } } } @@ -748,20 +774,29 @@ impl Analyzer { let req = &pair.request; let resp = &pair.response; - let request_body = req.json_body() + let request_body = req + .json_body() .map(|v| serde_json::to_string(&v).unwrap_or_default()) .or_else(|| { let body = req.body(); - if body.is_empty() { None } - else { Some(String::from_utf8_lossy(body).to_string()) } + if body.is_empty() { + None + } else { + Some(String::from_utf8_lossy(body).to_string()) + } }); - let response_body = resp.parsed.json_body() + let response_body = resp + .parsed + .json_body() .map(|v| serde_json::to_string(&v).unwrap_or_default()) .or_else(|| { let body = resp.parsed.body(); - if body.is_empty() { None } - else { Some(String::from_utf8_lossy(body).to_string()) } + if body.is_empty() { + None + } else { + Some(String::from_utf8_lossy(body).to_string()) + } }); Some(HttpRecord { @@ -773,9 +808,12 @@ impl Analyzer { status_code: resp.status_code(), request_headers: serde_json::to_string(&req.headers).unwrap_or_default(), request_body, - response_headers: serde_json::to_string(&resp.parsed.headers).unwrap_or_default(), + response_headers: serde_json::to_string(&resp.parsed.headers) + .unwrap_or_default(), response_body, - duration_ns: resp.end_timestamp_ns().saturating_sub(req.source_event.timestamp_ns), + duration_ns: resp + .end_timestamp_ns() + .saturating_sub(req.source_event.timestamp_ns), is_sse: false, sse_event_count: 0, }) @@ -784,12 +822,16 @@ impl Analyzer { let req = &pair.request; let resp = &pair.response; - let request_body = req.json_body() + let request_body = req + .json_body() .map(|v| serde_json::to_string(&v).unwrap_or_default()) .or_else(|| { let body = req.body(); - if body.is_empty() { None } - else { Some(String::from_utf8_lossy(body).to_string()) } + if body.is_empty() { + None + } else { + Some(String::from_utf8_lossy(body).to_string()) + } }); // For SSE responses, aggregate all SSE event JSON payloads @@ -809,20 +851,27 @@ impl Analyzer { status_code: resp.status_code(), request_headers: serde_json::to_string(&req.headers).unwrap_or_default(), request_body, - response_headers: serde_json::to_string(&resp.parsed.headers).unwrap_or_default(), + response_headers: serde_json::to_string(&resp.parsed.headers) + .unwrap_or_default(), response_body, - duration_ns: resp.end_timestamp_ns().saturating_sub(req.source_event.timestamp_ns), + duration_ns: resp + .end_timestamp_ns() + .saturating_sub(req.source_event.timestamp_ns), is_sse: true, sse_event_count: resp.sse_event_count(), }) } AggregatedResult::RequestOnly { request, .. } => { - let request_body = request.json_body() + let request_body = request + .json_body() .map(|v| serde_json::to_string(&v).unwrap_or_default()) .or_else(|| { let body = request.body(); - if body.is_empty() { None } - else { Some(String::from_utf8_lossy(body).to_string()) } + if body.is_empty() { + None + } else { + Some(String::from_utf8_lossy(body).to_string()) + } }); Some(HttpRecord { @@ -846,17 +895,22 @@ impl Analyzer { // Try SSE parsing first, fallback to regular text if it fails // This is more robust than checking content-type header (which may fail due to HPACK) - let (response_body, sse_event_count) = if let Some(sse_json) = stream.response_sse_json_array() { - // Successfully parsed as SSE - let event_count = sse_json.as_array().map(|a| a.len()).unwrap_or(0); - (Some(serde_json::to_string(&sse_json).unwrap_or_default()), event_count) - } else { - // Not SSE, try regular JSON or raw text - let body = stream.response_json_body() - .map(|v| serde_json::to_string(&v).unwrap_or_default()) - .or_else(|| stream.response_body_str()); - (body, 0) - }; + let (response_body, sse_event_count) = + if let Some(sse_json) = stream.response_sse_json_array() { + // Successfully parsed as SSE + let event_count = sse_json.as_array().map(|a| a.len()).unwrap_or(0); + ( + Some(serde_json::to_string(&sse_json).unwrap_or_default()), + event_count, + ) + } else { + // Not SSE, try regular JSON or raw text + let body = stream + .response_json_body() + .map(|v| serde_json::to_string(&v).unwrap_or_default()) + .or_else(|| stream.response_body_str()); + (body, 0) + }; Some(HttpRecord { timestamp_ns: stream.start_timestamp_ns, @@ -869,7 +923,9 @@ impl Analyzer { request_body, response_headers: stream.response_headers_json(), response_body, - duration_ns: stream.end_timestamp_ns.saturating_sub(stream.start_timestamp_ns), + duration_ns: stream + .end_timestamp_ns + .saturating_sub(stream.start_timestamp_ns), is_sse: sse_event_count > 0, sse_event_count, }) @@ -898,19 +954,21 @@ impl Analyzer { comm: &str, ) -> Option { let usage = self.token.parse_event(event)?; - - Some(TokenRecord::new( - pid, - comm.to_string(), - usage.provider.to_string(), - usage.input_tokens, - usage.output_tokens, + + Some( + TokenRecord::new( + pid, + comm.to_string(), + usage.provider.to_string(), + usage.input_tokens, + usage.output_tokens, + ) + .with_model(usage.model.clone().unwrap_or_default()) + .with_cache_tokens( + usage.cache_creation_input_tokens.unwrap_or(0), + usage.cache_read_input_tokens.unwrap_or(0), + ), ) - .with_model(usage.model.clone().unwrap_or_default()) - .with_cache_tokens( - usage.cache_creation_input_tokens.unwrap_or(0), - usage.cache_read_input_tokens.unwrap_or(0), - )) } /// Analyze an SSE event and return AnalysisResult @@ -960,7 +1018,8 @@ impl Analyzer { request_body: Option<&serde_json::Value>, response_body: Option<&serde_json::Value>, ) -> Option { - self.message.parse_by_path(path, request_body, response_body) + self.message + .parse_by_path(path, request_body, response_body) .map(AnalysisResult::Message) } @@ -974,14 +1033,18 @@ impl Analyzer { request_body: Option<&serde_json::Value>, response_body: Option<&serde_json::Value>, ) -> Option { - self.message.parse_by_path(path, request_body, response_body) + self.message + .parse_by_path(path, request_body, response_body) } /// Analyze AggregatedResult and extract token consumption breakdown /// /// This is a convenience method that combines extract_token_data and /// compute_token_consumption into a single call. - pub fn analyze_token_consumption(&self, result: &AggregatedResult) -> Option { + pub fn analyze_token_consumption( + &self, + result: &AggregatedResult, + ) -> Option { // Extract context (pid, comm) from the aggregated result (SseComplete only) let (pid, comm, request_json, response_jsons, path) = match result { AggregatedResult::SseComplete(pair) => ( @@ -1003,7 +1066,9 @@ impl Analyzer { let (tokenizer, chat_template) = match (&self.tokenizer, &self.chat_template) { (Some(t), Some(ct)) => (t, ct), _ => { - log::warn!("Tokenizer or chat template not available, cannot compute accurate token consumption"); + log::warn!( + "Tokenizer or chat template not available, cannot compute accurate token consumption" + ); return None; } }; @@ -1012,7 +1077,8 @@ impl Analyzer { let request_json_ref = request_json.as_ref()?; // Extract model - let model = request_json_ref.get("model") + let model = request_json_ref + .get("model") .and_then(|m| m.as_str()) .unwrap_or("unknown") .to_string(); @@ -1022,7 +1088,8 @@ impl Analyzer { "anthropic" } else { "openai" - }.to_string(); + } + .to_string(); // Count request tokens let request_count = count_request_tokens(request_json_ref, tokenizer, chat_template)?; diff --git a/src/agentsight/src/atif/converter.rs b/src/agentsight/src/atif/converter.rs index 34bd57548..22fa9fddc 100644 --- a/src/agentsight/src/atif/converter.rs +++ b/src/agentsight/src/atif/converter.rs @@ -8,7 +8,7 @@ use std::collections::HashMap; use crate::genai::semantic::{ - GenAISemanticEvent, LLMCall, MessagePart, InputMessage, OutputMessage, + GenAISemanticEvent, InputMessage, LLMCall, MessagePart, OutputMessage, }; use crate::storage::sqlite::genai::TraceEventDetail; @@ -30,41 +30,41 @@ pub fn convert_trace_to_atif( let mut steps = Vec::new(); // 1. System prompt step - if let Some(system_text) = extract_system_prompt(&events[0], parsed[0].as_ref()) { - if !system_text.is_empty() { - step_counter += 1; - steps.push(AtifStep { - step_id: step_counter, - timestamp: Some(ns_to_iso8601(events[0].start_timestamp_ns as u64)), - source: "system".to_string(), - message: Some(system_text), - model_name: None, - reasoning_content: None, - tool_calls: None, - observation: None, - metrics: None, - extra: None, - }); - } + if let Some(system_text) = extract_system_prompt(&events[0], parsed[0].as_ref()) + && !system_text.is_empty() + { + step_counter += 1; + steps.push(AtifStep { + step_id: step_counter, + timestamp: Some(ns_to_iso8601(events[0].start_timestamp_ns as u64)), + source: "system".to_string(), + message: Some(system_text), + model_name: None, + reasoning_content: None, + tool_calls: None, + observation: None, + metrics: None, + extra: None, + }); } // 2. User query step - if let Some(user_text) = extract_user_query(&events[0], parsed[0].as_ref()) { - if !user_text.is_empty() { - step_counter += 1; - steps.push(AtifStep { - step_id: step_counter, - timestamp: Some(ns_to_iso8601(events[0].start_timestamp_ns as u64)), - source: "user".to_string(), - message: Some(user_text), - model_name: None, - reasoning_content: None, - tool_calls: None, - observation: None, - metrics: None, - extra: None, - }); - } + if let Some(user_text) = extract_user_query(&events[0], parsed[0].as_ref()) + && !user_text.is_empty() + { + step_counter += 1; + steps.push(AtifStep { + step_id: step_counter, + timestamp: Some(ns_to_iso8601(events[0].start_timestamp_ns as u64)), + source: "user".to_string(), + message: Some(user_text), + model_name: None, + reasoning_content: None, + tool_calls: None, + observation: None, + metrics: None, + extra: None, + }); } // 3. Agent steps @@ -122,50 +122,47 @@ pub fn convert_session_to_atif( } // System prompt: emit if changed or first time - if let Some(system_text) = - extract_system_prompt(trace_events[0], trace_parsed.first().and_then(|p| p.as_ref())) + if let Some(system_text) = extract_system_prompt( + trace_events[0], + trace_parsed.first().and_then(|p| p.as_ref()), + ) && !system_text.is_empty() + && last_system_text.as_deref() != Some(&system_text) { - if !system_text.is_empty() && last_system_text.as_deref() != Some(&system_text) { - step_counter += 1; - steps.push(AtifStep { - step_id: step_counter, - timestamp: Some(ns_to_iso8601( - trace_events[0].start_timestamp_ns as u64, - )), - source: "system".to_string(), - message: Some(system_text.clone()), - model_name: None, - reasoning_content: None, - tool_calls: None, - observation: None, - metrics: None, - extra: None, - }); - last_system_text = Some(system_text); - } + step_counter += 1; + steps.push(AtifStep { + step_id: step_counter, + timestamp: Some(ns_to_iso8601(trace_events[0].start_timestamp_ns as u64)), + source: "system".to_string(), + message: Some(system_text.clone()), + model_name: None, + reasoning_content: None, + tool_calls: None, + observation: None, + metrics: None, + extra: None, + }); + last_system_text = Some(system_text); } // User query step - if let Some(user_text) = - extract_user_query(trace_events[0], trace_parsed.first().and_then(|p| p.as_ref())) + if let Some(user_text) = extract_user_query( + trace_events[0], + trace_parsed.first().and_then(|p| p.as_ref()), + ) && !user_text.is_empty() { - if !user_text.is_empty() { - step_counter += 1; - steps.push(AtifStep { - step_id: step_counter, - timestamp: Some(ns_to_iso8601( - trace_events[0].start_timestamp_ns as u64, - )), - source: "user".to_string(), - message: Some(user_text), - model_name: None, - reasoning_content: None, - tool_calls: None, - observation: None, - metrics: None, - extra: None, - }); - } + step_counter += 1; + steps.push(AtifStep { + step_id: step_counter, + timestamp: Some(ns_to_iso8601(trace_events[0].start_timestamp_ns as u64)), + source: "user".to_string(), + message: Some(user_text), + model_name: None, + reasoning_content: None, + tool_calls: None, + observation: None, + metrics: None, + extra: None, + }); } // Agent steps @@ -205,10 +202,7 @@ pub fn convert_session_to_atif( /// Parse event_json for all events upfront. Returns a Vec of Option. fn parse_all_events(events: &[TraceEventDetail]) -> Vec> { - events - .iter() - .map(|e| parse_event_json(e)) - .collect() + events.iter().map(parse_event_json).collect() } /// Try to deserialize event_json into an LLMCall. @@ -252,10 +246,7 @@ fn group_by_trace<'a>( } /// Build agent metadata from events. -fn build_agent_metadata( - events: &[TraceEventDetail], - parsed: &[Option], -) -> AtifAgent { +fn build_agent_metadata(events: &[TraceEventDetail], parsed: &[Option]) -> AtifAgent { // Agent name: first non-None agent_name, fallback to process_name let name = events .iter() @@ -301,7 +292,8 @@ fn collect_tool_definitions(parsed: &[Option]) -> Option]) -> Option, -) -> Option { +fn extract_system_prompt(event: &TraceEventDetail, parsed: Option<&LLMCall>) -> Option { // Strategy 1: from parsed LLMCall request messages if let Some(call) = parsed { let text = extract_text_by_role(&call.request.messages, "system"); @@ -339,10 +328,10 @@ fn extract_system_prompt( } } // Might be a plain string - if let Ok(s) = serde_json::from_str::(json) { - if !s.is_empty() { - return Some(s); - } + if let Ok(s) = serde_json::from_str::(json) + && !s.is_empty() + { + return Some(s); } } @@ -350,15 +339,12 @@ fn extract_system_prompt( } /// Extract user query text from event data. -fn extract_user_query( - event: &TraceEventDetail, - parsed: Option<&LLMCall>, -) -> Option { +fn extract_user_query(event: &TraceEventDetail, parsed: Option<&LLMCall>) -> Option { // Strategy 1: user_query column (already cleaned by builder) - if let Some(ref q) = event.user_query { - if !q.is_empty() { - return Some(q.clone()); - } + if let Some(ref q) = event.user_query + && !q.is_empty() + { + return Some(q.clone()); } // Strategy 2: from parsed LLMCall — last user message text @@ -370,12 +356,12 @@ fn extract_user_query( } // Strategy 3: from input_messages column - if let Some(ref json) = event.input_messages { - if let Ok(msgs) = serde_json::from_str::>(json) { - let text = extract_last_user_text_from_input(&msgs); - if let Some(t) = text { - return Some(t); - } + if let Some(ref json) = event.input_messages + && let Ok(msgs) = serde_json::from_str::>(json) + { + let text = extract_last_user_text_from_input(&msgs); + if let Some(t) = text { + return Some(t); } } @@ -414,7 +400,11 @@ fn build_agent_step( reasoning_text.push_str(content); } } - MessagePart::ToolCall { id, name, arguments } => { + MessagePart::ToolCall { + id, + name, + arguments, + } => { let tc_id = id .clone() .unwrap_or_else(|| format!("auto_{}", tool_calls.len())); @@ -432,41 +422,45 @@ fn build_agent_step( } } else { // Fallback: parse output_messages column directly - if let Some(ref json) = event.output_messages { - if let Ok(msgs) = serde_json::from_str::>(json) { - for msg in &msgs { - for part in &msg.parts { - match part { - MessagePart::Text { content } => { - if !content.is_empty() { - if !message_text.is_empty() { - message_text.push('\n'); - } - message_text.push_str(content); + if let Some(ref json) = event.output_messages + && let Ok(msgs) = serde_json::from_str::>(json) + { + for msg in &msgs { + for part in &msg.parts { + match part { + MessagePart::Text { content } => { + if !content.is_empty() { + if !message_text.is_empty() { + message_text.push('\n'); } + message_text.push_str(content); } - MessagePart::Reasoning { content } => { - if !content.is_empty() { - if !reasoning_text.is_empty() { - reasoning_text.push('\n'); - } - reasoning_text.push_str(content); + } + MessagePart::Reasoning { content } => { + if !content.is_empty() { + if !reasoning_text.is_empty() { + reasoning_text.push('\n'); } + reasoning_text.push_str(content); } - MessagePart::ToolCall { id, name, arguments } => { - let tc_id = id + } + MessagePart::ToolCall { + id, + name, + arguments, + } => { + let tc_id = id + .clone() + .unwrap_or_else(|| format!("auto_{}", tool_calls.len())); + tool_calls.push(AtifToolCall { + tool_call_id: tc_id, + function_name: name.clone(), + arguments: arguments .clone() - .unwrap_or_else(|| format!("auto_{}", tool_calls.len())); - tool_calls.push(AtifToolCall { - tool_call_id: tc_id, - function_name: name.clone(), - arguments: arguments - .clone() - .unwrap_or(serde_json::Value::Object(Default::default())), - }); - } - _ => {} + .unwrap_or(serde_json::Value::Object(Default::default())), + }); } + _ => {} } } } @@ -492,16 +486,14 @@ fn build_agent_step( } else { None }, - cached_tokens: event.cache_read_tokens.and_then(|v| { - if v > 0 { Some(v as u32) } else { None } - }), + cached_tokens: event + .cache_read_tokens + .and_then(|v| if v > 0 { Some(v as u32) } else { None }), extra: None, }); // Timestamp: prefer end_timestamp (when response arrived) - let timestamp_ns = event - .end_timestamp_ns - .unwrap_or(event.start_timestamp_ns); + let timestamp_ns = event.end_timestamp_ns.unwrap_or(event.start_timestamp_ns); AtifStep { step_id, @@ -538,7 +530,9 @@ fn latest_round_messages(messages: &[InputMessage]) -> &[InputMessage] { // Find the last assistant message that contains a ToolCall part. let last_assistant_with_tc = messages.iter().rposition(|m| { m.role == "assistant" - && m.parts.iter().any(|p| matches!(p, MessagePart::ToolCall { .. })) + && m.parts + .iter() + .any(|p| matches!(p, MessagePart::ToolCall { .. })) }); if let Some(idx) = last_assistant_with_tc { @@ -578,22 +572,16 @@ fn build_observation( // Strategy 2: from event_json parsed as LLMCall (latest round only) if let Some(call) = parse_event_json(next_event) { let latest = latest_round_messages(&call.request.messages); - collect_tool_responses( - latest, - &tc_ids, - &mut results, - &mut matched_by_id, - ); + collect_tool_responses(latest, &tc_ids, &mut results, &mut matched_by_id); } } // Strategy 3: from input_messages column (already incremental — latest round) - if results.is_empty() { - if let Some(ref json) = next_event.input_messages { - if let Ok(msgs) = serde_json::from_str::>(json) { - collect_tool_responses(&msgs, &tc_ids, &mut results, &mut matched_by_id); - } - } + if results.is_empty() + && let Some(ref json) = next_event.input_messages + && let Ok(msgs) = serde_json::from_str::>(json) + { + collect_tool_responses(&msgs, &tc_ids, &mut results, &mut matched_by_id); } if results.is_empty() { @@ -623,17 +611,16 @@ fn collect_tool_responses( }; // Try to match by ID first - if let Some(tc_id) = id { - if let Some(&idx) = tc_ids.get(tc_id.as_str()) { - if !matched[idx] { - matched[idx] = true; - results.push(AtifObservationResult { - source_call_id: Some(tc_id.clone()), - content: Some(content_str), - }); - continue; - } - } + if let Some(tc_id) = id + && let Some(&idx) = tc_ids.get(tc_id.as_str()) + && !matched[idx] + { + matched[idx] = true; + results.push(AtifObservationResult { + source_call_id: Some(tc_id.clone()), + content: Some(content_str), + }); + continue; } // Fallback: positional matching @@ -691,10 +678,10 @@ fn extract_text_by_role(messages: &[InputMessage], role: &str) -> String { for msg in messages { if msg.role == role { for part in &msg.parts { - if let MessagePart::Text { content } = part { - if !content.is_empty() { - parts.push(content.as_str()); - } + if let MessagePart::Text { content } = part + && !content.is_empty() + { + parts.push(content.as_str()); } } } @@ -763,21 +750,30 @@ mod tests { let messages = vec![ InputMessage { role: "system".to_string(), - parts: vec![MessagePart::Text { content: "You are helpful.".to_string() }], + parts: vec![MessagePart::Text { + content: "You are helpful.".to_string(), + }], name: None, }, InputMessage { role: "user".to_string(), - parts: vec![MessagePart::Text { content: "Hello".to_string() }], + parts: vec![MessagePart::Text { + content: "Hello".to_string(), + }], name: None, }, InputMessage { role: "user".to_string(), - parts: vec![MessagePart::Text { content: "World".to_string() }], + parts: vec![MessagePart::Text { + content: "World".to_string(), + }], name: None, }, ]; - assert_eq!(extract_text_by_role(&messages, "system"), "You are helpful."); + assert_eq!( + extract_text_by_role(&messages, "system"), + "You are helpful." + ); assert_eq!(extract_text_by_role(&messages, "user"), "Hello\nWorld"); assert_eq!(extract_text_by_role(&messages, "assistant"), ""); } @@ -787,44 +783,53 @@ mod tests { let messages = vec![ InputMessage { role: "user".to_string(), - parts: vec![MessagePart::Text { content: "First".to_string() }], + parts: vec![MessagePart::Text { + content: "First".to_string(), + }], name: None, }, InputMessage { role: "assistant".to_string(), - parts: vec![MessagePart::Text { content: "Response".to_string() }], + parts: vec![MessagePart::Text { + content: "Response".to_string(), + }], name: None, }, InputMessage { role: "user".to_string(), - parts: vec![MessagePart::Text { content: "Second".to_string() }], + parts: vec![MessagePart::Text { + content: "Second".to_string(), + }], name: None, }, ]; - assert_eq!(extract_last_user_text(&messages), Some("Second".to_string())); + assert_eq!( + extract_last_user_text(&messages), + Some("Second".to_string()) + ); } #[test] fn test_extract_last_user_text_empty() { - let messages = vec![ - InputMessage { - role: "system".to_string(), - parts: vec![MessagePart::Text { content: "sys".to_string() }], - name: None, - }, - ]; + let messages = vec![InputMessage { + role: "system".to_string(), + parts: vec![MessagePart::Text { + content: "sys".to_string(), + }], + name: None, + }]; assert_eq!(extract_last_user_text(&messages), None); } #[test] fn test_extract_text_from_input_messages() { - let messages = vec![ - InputMessage { - role: "user".to_string(), - parts: vec![MessagePart::Text { content: "Query".to_string() }], - name: None, - }, - ]; + let messages = vec![InputMessage { + role: "user".to_string(), + parts: vec![MessagePart::Text { + content: "Query".to_string(), + }], + name: None, + }]; assert_eq!(extract_text_from_input_messages(&messages, "user"), "Query"); } @@ -833,15 +838,22 @@ mod tests { let messages = vec![ InputMessage { role: "user".to_string(), - parts: vec![MessagePart::Text { content: "Q1".to_string() }], + parts: vec![MessagePart::Text { + content: "Q1".to_string(), + }], name: None, }, InputMessage { role: "user".to_string(), - parts: vec![MessagePart::Text { content: "Q2".to_string() }], + parts: vec![MessagePart::Text { + content: "Q2".to_string(), + }], name: None, }, ]; - assert_eq!(extract_last_user_text_from_input(&messages), Some("Q2".to_string())); + assert_eq!( + extract_last_user_text_from_input(&messages), + Some("Q2".to_string()) + ); } } diff --git a/src/agentsight/src/atif/mod.rs b/src/agentsight/src/atif/mod.rs index 3a36cb3f6..ca47e1184 100644 --- a/src/agentsight/src/atif/mod.rs +++ b/src/agentsight/src/atif/mod.rs @@ -6,13 +6,11 @@ //! This module is independent from the `genai` module — it only depends on //! storage query result types and `genai::semantic` types for deserialization. -pub mod schema; pub mod converter; +pub mod schema; +pub use converter::{convert_session_to_atif, convert_trace_to_atif}; pub use schema::{ - AtifDocument, AtifAgent, AtifStep, AtifToolCall, - AtifObservation, AtifObservationResult, - AtifStepMetrics, AtifFinalMetrics, - SCHEMA_VERSION, + AtifAgent, AtifDocument, AtifFinalMetrics, AtifObservation, AtifObservationResult, AtifStep, + AtifStepMetrics, AtifToolCall, SCHEMA_VERSION, }; -pub use converter::{convert_trace_to_atif, convert_session_to_atif}; diff --git a/src/agentsight/src/atif/schema.rs b/src/agentsight/src/atif/schema.rs index 1b4935dd3..acca93550 100644 --- a/src/agentsight/src/atif/schema.rs +++ b/src/agentsight/src/atif/schema.rs @@ -3,7 +3,7 @@ //! Defines Rust types that serialize to/from the ATIF v1.6 JSON schema. //! See: -use serde::{Serialize, Deserialize}; +use serde::{Deserialize, Serialize}; /// Current ATIF schema version pub const SCHEMA_VERSION: &str = "ATIF-v1.6"; @@ -265,13 +265,11 @@ mod tests { message: None, model_name: Some("claude-3".to_string()), reasoning_content: None, - tool_calls: Some(vec![ - AtifToolCall { - tool_call_id: "tc1".to_string(), - function_name: "search".to_string(), - arguments: serde_json::json!({"query": "rust"}), - }, - ]), + tool_calls: Some(vec![AtifToolCall { + tool_call_id: "tc1".to_string(), + function_name: "search".to_string(), + arguments: serde_json::json!({"query": "rust"}), + }]), observation: Some(AtifObservation { results: vec![AtifObservationResult { source_call_id: Some("tc1".to_string()), diff --git a/src/agentsight/src/bin/agentsight.rs b/src/agentsight/src/bin/agentsight.rs index 91ca0aec3..5aa4f8a16 100644 --- a/src/agentsight/src/bin/agentsight.rs +++ b/src/agentsight/src/bin/agentsight.rs @@ -10,12 +10,19 @@ use structopt::StructOpt; mod cli; -use cli::{token::TokenCommand, trace::TraceCommand, audit::AuditCommand, discover::DiscoverCommand, metrics::MetricsCommand, interruption::InterruptionCommand, skill_metrics::SkillMetricsCommand}; #[cfg(feature = "server")] use cli::serve::ServeCommand; +use cli::{ + audit::AuditCommand, discover::DiscoverCommand, interruption::InterruptionCommand, + metrics::MetricsCommand, skill_metrics::SkillMetricsCommand, token::TokenCommand, + trace::TraceCommand, +}; #[derive(Debug, StructOpt)] -#[structopt(name = "agentsight", about = "AI Agent observability tool - trace processes, SSL traffic, and LLM API calls via eBPF")] +#[structopt( + name = "agentsight", + about = "AI Agent observability tool - trace processes, SSL traffic, and LLM API calls via eBPF" +)] pub enum Command { /// Query token consumption data Token(TokenCommand), diff --git a/src/agentsight/src/bin/cli/audit.rs b/src/agentsight/src/bin/cli/audit.rs index 37b096861..406411dcf 100644 --- a/src/agentsight/src/bin/cli/audit.rs +++ b/src/agentsight/src/bin/cli/audit.rs @@ -44,7 +44,10 @@ impl AuditCommand { } }; - let event_type = self.event_type.as_ref().and_then(|t| t.parse::().ok()); + let event_type = self + .event_type + .as_ref() + .and_then(|t| t.parse::().ok()); if self.summary { self.print_summary(&store); @@ -77,18 +80,21 @@ impl AuditCommand { fn output_records(&self, records: &[agentsight::AuditRecord], scope: &str) { if self.json { - let json_records: Vec = records.iter().map(|r| { - serde_json::json!({ - "id": r.id, - "event_type": r.event_type.to_string(), - "timestamp_ns": r.timestamp_ns, - "pid": r.pid, - "ppid": r.ppid, - "comm": r.comm, - "duration_ns": r.duration_ns, - "extra": r.extra, + let json_records: Vec = records + .iter() + .map(|r| { + serde_json::json!({ + "id": r.id, + "event_type": r.event_type.to_string(), + "timestamp_ns": r.timestamp_ns, + "pid": r.pid, + "ppid": r.ppid, + "comm": r.comm, + "duration_ns": r.duration_ns, + "extra": r.extra, + }) }) - }).collect(); + .collect(); println!("{}", serde_json::to_string_pretty(&json_records).unwrap()); } else { println!("{}: {} audit events", scope, records.len()); diff --git a/src/agentsight/src/bin/cli/discover.rs b/src/agentsight/src/bin/cli/discover.rs index 845432d38..a88df2a0d 100644 --- a/src/agentsight/src/bin/cli/discover.rs +++ b/src/agentsight/src/bin/cli/discover.rs @@ -39,7 +39,10 @@ impl DiscoverCommand { println!(); // Use CmdlineGlobMatcher to list agent info - for matcher in agentsight::default_cmdline_rules().iter().filter_map(|rule| CmdlineGlobMatcher::from_config(rule)) { + for matcher in agentsight::default_cmdline_rules() + .iter() + .filter_map(CmdlineGlobMatcher::from_config) + { let agent = matcher.info(); println!(" {} ({})", agent.name, agent.category); println!(" Process names: {}", agent.process_names.join(", ")); diff --git a/src/agentsight/src/bin/cli/interruption.rs b/src/agentsight/src/bin/cli/interruption.rs index 8e80605d2..6a7b013c1 100644 --- a/src/agentsight/src/bin/cli/interruption.rs +++ b/src/agentsight/src/bin/cli/interruption.rs @@ -55,7 +55,7 @@ //! agentsight interruption list --last 24 --json //! ``` -use agentsight::storage::sqlite::{GenAISqliteStore, InterruptionStore, InterruptionRecord}; +use agentsight::storage::sqlite::{GenAISqliteStore, InterruptionRecord, InterruptionStore}; use structopt::StructOpt; /// Query and manage AI agent session interruption events. @@ -230,7 +230,14 @@ impl InterruptionCommand { match &self.action { InterruptionAction::List { - last, itype, severity, agent, unresolved, resolved, limit, json, + last, + itype, + severity, + agent, + unresolved, + resolved, + limit, + json, } => { let (start_ns, end_ns) = time_range_ns(*last); let resolved_filter = if *unresolved { @@ -242,7 +249,8 @@ impl InterruptionCommand { }; match store.list( - start_ns, end_ns, + start_ns, + end_ns, agent.as_deref(), itype.as_deref(), severity.as_deref(), @@ -263,25 +271,26 @@ impl InterruptionCommand { } } - InterruptionAction::Get { interruption_id, json } => { - match store.get_by_id(interruption_id) { - Ok(Some(record)) => { - if *json { - print_json(&record); - } else { - print_record_detail(&record); - } - } - Ok(None) => { - eprintln!("No interruption found with id: {}", interruption_id); - std::process::exit(1); - } - Err(e) => { - eprintln!("Query error: {}", e); - std::process::exit(1); + InterruptionAction::Get { + interruption_id, + json, + } => match store.get_by_id(interruption_id) { + Ok(Some(record)) => { + if *json { + print_json(&record); + } else { + print_record_detail(&record); } } - } + Ok(None) => { + eprintln!("No interruption found with id: {}", interruption_id); + std::process::exit(1); + } + Err(e) => { + eprintln!("Query error: {}", e); + std::process::exit(1); + } + }, InterruptionAction::Stats { last, json } => { let (start_ns, end_ns) = time_range_ns(*last); @@ -297,7 +306,10 @@ impl InterruptionCommand { println!("{:<20} {:<10} {:>6}", "TYPE", "SEVERITY", "COUNT"); println!("{}", "-".repeat(40)); for s in &stats { - println!("{:<20} {:<10} {:>6}", s.interruption_type, s.severity, s.count); + println!( + "{:<20} {:<10} {:>6}", + s.interruption_type, s.severity, s.count + ); } } } @@ -378,27 +390,28 @@ impl InterruptionCommand { } } - InterruptionAction::Conversation { conversation_id, json } => { - match store.list_by_conversation(conversation_id) { - Ok(rows) => { - if *json { - print_json(&rows); - } else { - if rows.is_empty() { - println!("No interruptions for conversation: {}", conversation_id); - return; - } - println!("Interruptions for conversation {}:", conversation_id); - println!(); - print_records_table(&rows); + InterruptionAction::Conversation { + conversation_id, + json, + } => match store.list_by_conversation(conversation_id) { + Ok(rows) => { + if *json { + print_json(&rows); + } else { + if rows.is_empty() { + println!("No interruptions for conversation: {}", conversation_id); + return; } - } - Err(e) => { - eprintln!("Query error: {}", e); - std::process::exit(1); + println!("Interruptions for conversation {}:", conversation_id); + println!(); + print_records_table(&rows); } } - } + Err(e) => { + eprintln!("Query error: {}", e); + std::process::exit(1); + } + }, InterruptionAction::Resolve { interruption_id } => { match store.resolve(interruption_id) { @@ -453,10 +466,15 @@ fn format_ns(ns: i64) -> String { let hour = total_hours % 24; // Days since epoch to Y-M-D (simplified) - let (year, month, day) = days_to_ymd(total_days as i64); + let (year, month, day) = days_to_ymd(total_days); format!( "{:04}-{:02}-{:02} {:02}:{:02}:{:02}.{:03}", - year, month, day, hour, min, sec, + year, + month, + day, + hour, + min, + sec, nanos_rem / 1_000_000 ) } @@ -494,9 +512,8 @@ fn print_records_table(records: &[InterruptionRecord]) { } println!( - "{:<34} {:<18} {:<10} {:<22} {:<10} {:<14} {:<16} {}", - "INTERRUPTION_ID", "TYPE", "SEVERITY", "OCCURRED_AT", "RESOLVED", - "AGENT", "SESSION_ID", "CONVERSATION_ID" + "{:<34} {:<18} {:<10} {:<22} {:<10} {:<14} {:<16} CONVERSATION_ID", + "INTERRUPTION_ID", "TYPE", "SEVERITY", "OCCURRED_AT", "RESOLVED", "AGENT", "SESSION_ID" ); println!("{}", "-".repeat(140)); @@ -530,19 +547,34 @@ fn print_record_detail(r: &InterruptionRecord) { println!(" ID: {}", r.interruption_id); println!(" Type: {}", r.interruption_type); println!(" Severity: {}", r.severity); - println!(" Occurred At: {} ({}ns)", format_ns(r.occurred_at_ns), r.occurred_at_ns); + println!( + " Occurred At: {} ({}ns)", + format_ns(r.occurred_at_ns), + r.occurred_at_ns + ); println!(" Resolved: {}", if r.resolved { "yes" } else { "no" }); println!(" Session ID: {}", r.session_id.as_deref().unwrap_or("-")); - println!(" Conversation: {}", r.conversation_id.as_deref().unwrap_or("-")); + println!( + " Conversation: {}", + r.conversation_id.as_deref().unwrap_or("-") + ); println!(" Trace ID: {}", r.trace_id.as_deref().unwrap_or("-")); println!(" Call ID: {}", r.call_id.as_deref().unwrap_or("-")); - println!(" PID: {}", r.pid.map(|p| p.to_string()).unwrap_or_else(|| "-".to_string())); + println!( + " PID: {}", + r.pid + .map(|p| p.to_string()) + .unwrap_or_else(|| "-".to_string()) + ); println!(" Agent: {}", r.agent_name.as_deref().unwrap_or("-")); if let Some(ref detail) = r.detail { // Pretty-print JSON detail if let Ok(v) = serde_json::from_str::(detail) { println!(" Detail:"); - println!("{}", serde_json::to_string_pretty(&v).unwrap_or_else(|_| detail.clone())); + println!( + "{}", + serde_json::to_string_pretty(&v).unwrap_or_else(|_| detail.clone()) + ); } else { println!(" Detail: {}", detail); } diff --git a/src/agentsight/src/bin/cli/metrics.rs b/src/agentsight/src/bin/cli/metrics.rs index d31d949d2..ac4b9fcd9 100644 --- a/src/agentsight/src/bin/cli/metrics.rs +++ b/src/agentsight/src/bin/cli/metrics.rs @@ -34,7 +34,9 @@ impl MetricsCommand { // ── Prometheus text format output ────────────────────────────────── - println!("# HELP agentsight_token_input_total Total input tokens consumed by agent (all-time)"); + println!( + "# HELP agentsight_token_input_total Total input tokens consumed by agent (all-time)" + ); println!("# TYPE agentsight_token_input_total counter"); for s in &summaries { println!( @@ -45,7 +47,9 @@ impl MetricsCommand { } println!(); - println!("# HELP agentsight_token_output_total Total output tokens consumed by agent (all-time)"); + println!( + "# HELP agentsight_token_output_total Total output tokens consumed by agent (all-time)" + ); println!("# TYPE agentsight_token_output_total counter"); for s in &summaries { println!( @@ -56,7 +60,9 @@ impl MetricsCommand { } println!(); - println!("# HELP agentsight_token_total_total Total tokens (input+output) consumed by agent (all-time)"); + println!( + "# HELP agentsight_token_total_total Total tokens (input+output) consumed by agent (all-time)" + ); println!("# TYPE agentsight_token_total_total counter"); for s in &summaries { println!( @@ -67,7 +73,9 @@ impl MetricsCommand { } println!(); - println!("# HELP agentsight_llm_requests_total Total LLM requests made by agent (all-time)"); + println!( + "# HELP agentsight_llm_requests_total Total LLM requests made by agent (all-time)" + ); println!("# TYPE agentsight_llm_requests_total counter"); for s in &summaries { println!( @@ -82,6 +90,6 @@ impl MetricsCommand { /// Escape Prometheus label value: backslash → \\, double-quote → \", newline → \n fn escape_label(s: &str) -> String { s.replace('\\', "\\\\") - .replace('"', "\\\"") - .replace('\n', "\\n") + .replace('"', "\\\"") + .replace('\n', "\\n") } diff --git a/src/agentsight/src/bin/cli/mod.rs b/src/agentsight/src/bin/cli/mod.rs index bda716c79..ed4bbec81 100644 --- a/src/agentsight/src/bin/cli/mod.rs +++ b/src/agentsight/src/bin/cli/mod.rs @@ -7,15 +7,15 @@ //! - `discover`: Discover running AI agents //! - `interruption`: Query and manage session interruption events -pub mod token; -pub mod trace; pub mod audit; pub mod discover; -pub mod metrics; pub mod interruption; +pub mod metrics; #[cfg(feature = "server")] pub mod serve; pub mod skill_metrics; +pub mod token; +pub mod trace; /// Parse period string into TimePeriod pub fn parse_period(s: &str) -> agentsight::TimePeriod { diff --git a/src/agentsight/src/bin/cli/serve.rs b/src/agentsight/src/bin/cli/serve.rs index 011014703..4460d736a 100644 --- a/src/agentsight/src/bin/cli/serve.rs +++ b/src/agentsight/src/bin/cli/serve.rs @@ -22,11 +22,12 @@ pub struct ServeCommand { impl ServeCommand { pub fn execute(&self) { - let db_path = self.db + let db_path = self + .db .as_ref() - .map(|p| std::path::PathBuf::from(p)) + .map(std::path::PathBuf::from) // Default to genai_events.db — the same file the tracer writes to - .unwrap_or_else(|| GenAISqliteStore::default_path()); + .unwrap_or_else(GenAISqliteStore::default_path); let host = self.host.clone(); let port = self.port; diff --git a/src/agentsight/src/bin/cli/token.rs b/src/agentsight/src/bin/cli/token.rs index 634f3c1bb..45ca82650 100644 --- a/src/agentsight/src/bin/cli/token.rs +++ b/src/agentsight/src/bin/cli/token.rs @@ -1,8 +1,7 @@ //! Token query subcommand use agentsight::{ - TimePeriod, TokenQueryResult, format_tokens_with_commas, Trend, TokenStore, - SqliteConfig, + SqliteConfig, TimePeriod, TokenQueryResult, TokenStore, Trend, format_tokens_with_commas, }; use structopt::StructOpt; @@ -35,9 +34,10 @@ impl TokenCommand { // Determine data file path // Use the unified database path (agentsight.db) as default, // which is where Storage writes all tables. - let data_path = self.data_file + let data_path = self + .data_file .as_ref() - .map(|p| std::path::PathBuf::from(p)) + .map(std::path::PathBuf::from) .unwrap_or_else(|| SqliteConfig::default().db_path()); self.execute_summary(&data_path); @@ -62,12 +62,10 @@ impl TokenCommand { } else { query.by_period(period) } + } else if self.compare { + query.by_period_with_compare(TimePeriod::Today) } else { - if self.compare { - query.by_period_with_compare(TimePeriod::Today) - } else { - query.by_period(TimePeriod::Today) - } + query.by_period(TimePeriod::Today) }; // Output result @@ -80,10 +78,7 @@ impl TokenCommand { } /// Print human-readable summary output -fn print_human_readable( - result: &TokenQueryResult, - show_compare: bool, -) { +fn print_human_readable(result: &TokenQueryResult, show_compare: bool) { // Main result println!( "{}共消耗 {} tokens。", @@ -92,21 +87,19 @@ fn print_human_readable( ); // Comparison - if show_compare { - if let Some(ref comp) = result.comparison { - let trend = match comp.trend { - Trend::Up => "增长", - Trend::Down => "下降", - Trend::Flat => "持平", - }; + if show_compare && let Some(ref comp) = result.comparison { + let trend = match comp.trend { + Trend::Up => "增长", + Trend::Down => "下降", + Trend::Flat => "持平", + }; - println!( - "比上一时段（{}）{}了 {}。", - format_tokens_with_commas(comp.previous_total), - trend, - comp.formatted_change() - ); - } + println!( + "比上一时段（{}）{}了 {}。", + format_tokens_with_commas(comp.previous_total), + trend, + comp.formatted_change() + ); } // Additional details @@ -119,6 +112,4 @@ fn print_human_readable( format_tokens_with_commas(result.output_tokens) ); } - } - diff --git a/src/agentsight/src/bin/cli/trace.rs b/src/agentsight/src/bin/cli/trace.rs index 8a5dec188..c6ac4ea76 100644 --- a/src/agentsight/src/bin/cli/trace.rs +++ b/src/agentsight/src/bin/cli/trace.rs @@ -1,8 +1,8 @@ //! Trace subcommand - eBPF-based agent activity tracing use agentsight::{AgentSight, AgentsightConfig}; -use structopt::StructOpt; use daemonize::Daemonize; +use structopt::StructOpt; /// Trace subcommand #[derive(Debug, StructOpt, Clone)] @@ -34,20 +34,20 @@ impl TraceCommand { self.run_as_daemon(); return; } - + self.run_tracing(); } - + /// Run as daemon process fn run_as_daemon(&self) { println!("Starting agentsight in daemon mode..."); println!("PID file: {}", self.pid_file); - + let daemonize = Daemonize::new() .pid_file(&self.pid_file) .chown_pid_file(true) .working_directory("/tmp"); - + match daemonize.start() { Ok(_) => { // We're now in the daemon process @@ -59,7 +59,7 @@ impl TraceCommand { } } } - + /// Run the actual tracing logic using AgentSight fn run_tracing(&self) { // Build AgentSight config (empty target_pids means trace all processes). diff --git a/src/agentsight/src/bin/proctrace.rs b/src/agentsight/src/bin/proctrace.rs index d57179876..1ede5f5e9 100644 --- a/src/agentsight/src/bin/proctrace.rs +++ b/src/agentsight/src/bin/proctrace.rs @@ -7,8 +7,8 @@ use agentsight::config; use agentsight::parser::ProcTraceParser; use agentsight::probes::proctrace::ProcTrace; -use structopt::StructOpt; use std::time::Duration; +use structopt::StructOpt; #[derive(Debug, StructOpt)] #[structopt(name = "proctrace", about = "Trace process execution and output")] @@ -16,11 +16,11 @@ pub struct Command { /// Enable verbose/debug output #[structopt(short, long)] verbose: bool, - + /// Target PID to trace (optional, trace all if not specified) #[structopt(short, long)] pid: Option, - + /// Filter by UID #[structopt(long)] uid: Option, diff --git a/src/agentsight/src/bin/sslsniff.rs b/src/agentsight/src/bin/sslsniff.rs index 1ea4b2f13..84852411a 100644 --- a/src/agentsight/src/bin/sslsniff.rs +++ b/src/agentsight/src/bin/sslsniff.rs @@ -8,17 +8,20 @@ use agentsight::config; use agentsight::parser::Parser; use agentsight::probes::sslsniff::SslSniff; -use structopt::StructOpt; use std::rc::Rc; use std::time::Duration; +use structopt::StructOpt; #[derive(Debug, StructOpt)] -#[structopt(name = "sslsniff", about = "Parse and print HTTP/SSE traffic from SSL connections")] +#[structopt( + name = "sslsniff", + about = "Parse and print HTTP/SSE traffic from SSL connections" +)] pub struct Command { /// Enable verbose/debug output #[structopt(short, long)] verbose: bool, - + /// Target PID #[structopt(short, long)] pid: i32, @@ -32,9 +35,11 @@ fn main() { // Create SSL sniffer let mut sniffer = SslSniff::new().expect("Failed to create SSL sniffer"); - + // Attach to target process - sniffer.attach_process(opts.pid).expect("Failed to attach SSL probe"); + sniffer + .attach_process(opts.pid) + .expect("Failed to attach SSL probe"); // Start polling let _poller = sniffer.run().expect("Failed to start SSL poller"); diff --git a/src/agentsight/src/chrome_trace.rs b/src/agentsight/src/chrome_trace.rs index df69d32e9..c63dd84a5 100644 --- a/src/agentsight/src/chrome_trace.rs +++ b/src/agentsight/src/chrome_trace.rs @@ -210,22 +210,22 @@ impl ChromeTraceEvent { } /// Create flow events connecting two Chrome Trace Events - /// + /// /// This generates a pair of flow events (s and f phases) that create /// an arrow in Perfetto from the start event to the end event. - /// + /// /// Flow events only carry the arrow connection, no args. /// The semantic information should be in the accompanying complete/instant events. - /// + /// /// Flow ID is automatically generated using a global atomic counter. - /// + /// /// # Arguments /// * `start` - The source event (arrow starts here) /// * `end` - The target event (arrow ends here) - /// + /// /// # Returns /// A tuple of (flow_start, flow_end, flow_id) events - /// + /// /// # Example /// ```rust,ignore /// let request_event = ChromeTraceEvent::complete("GET /api", "http", pid1, tid1, ts1, dur1); @@ -239,21 +239,25 @@ impl ChromeTraceEvent { } /// Create flow events connecting two Chrome Trace Events with a given flow_id - /// + /// /// This generates a pair of flow events (s and f phases) that create /// an arrow in Perfetto from the start event to the end event. - /// + /// /// Flow events only carry the arrow connection, no args. /// The semantic information should be in the accompanying complete/instant events. - /// + /// /// # Arguments /// * `start` - The source event (arrow starts here) /// * `end` - The target event (arrow ends here) /// * `flow_id` - Unique identifier to link the flow events - /// + /// /// # Returns /// A tuple of (flow_start, flow_end) events - pub fn flow_from_events_with_id(start: &ChromeTraceEvent, end: &ChromeTraceEvent, flow_id: u64) -> (Self, Self) { + pub fn flow_from_events_with_id( + start: &ChromeTraceEvent, + end: &ChromeTraceEvent, + flow_id: u64, + ) -> (Self, Self) { let flow_start = ChromeTraceEvent { name: "flow".to_string(), cat: "flow".to_string(), @@ -266,7 +270,7 @@ impl ChromeTraceEvent { id: Some(flow_id), bp: None, }; - + let flow_end = ChromeTraceEvent { name: "flow".to_string(), cat: "flow".to_string(), @@ -279,7 +283,7 @@ impl ChromeTraceEvent { id: Some(flow_id), bp: Some("e".to_string()), }; - + (flow_start, flow_end) } @@ -356,7 +360,9 @@ fn trace_file_path() -> &'static std::path::PathBuf { static PATH: OnceLock = OnceLock::new(); PATH.get_or_init(|| { let datetime = chrono::Local::now().format("%Y-%m-%d_%H-%M"); - std::env::current_dir().unwrap_or_default().join(format!("trace-{}.json", datetime)) + std::env::current_dir() + .unwrap_or_default() + .join(format!("trace-{}.json", datetime)) }) } @@ -396,7 +402,7 @@ pub fn append_trace_event(event: &ChromeTraceEvent) { } /// Export trace events to file if enabled -/// +/// /// This function checks if chrome trace export is enabled via environment variable /// AGENTSIGHT_CHROME_TRACE, and if so, writes the events to trace.json. pub fn export_trace_events(result: &T) { @@ -559,8 +565,7 @@ mod tests { #[test] fn test_with_trace_args_trait() { - let e = ChromeTraceEvent::instant("t", "c", 1, 1, 0) - .with_trace_args(&MockTraceArgs); + let e = ChromeTraceEvent::instant("t", "c", 1, 1, 0).with_trace_args(&MockTraceArgs); assert_eq!(e.args.unwrap()["custom"], true); } @@ -573,8 +578,7 @@ mod tests { #[test] fn test_with_trace_args_empty_object() { - let e = ChromeTraceEvent::instant("t", "c", 1, 1, 0) - .with_trace_args(&EmptyTraceArgs); + let e = ChromeTraceEvent::instant("t", "c", 1, 1, 0).with_trace_args(&EmptyTraceArgs); // Empty object should not be set assert!(e.args.is_none()); } diff --git a/src/agentsight/src/config.rs b/src/agentsight/src/config.rs index 56029872e..a4d135262 100644 --- a/src/agentsight/src/config.rs +++ b/src/agentsight/src/config.rs @@ -1,8 +1,8 @@ +use anyhow::Context; use std::net::Ipv4Addr; use std::path::{Path, PathBuf}; use std::str::FromStr; use std::sync::atomic::{AtomicBool, Ordering}; -use anyhow::Context; // ==================== Default Constants ==================== @@ -42,7 +42,7 @@ pub const DEFAULT_PURGE_INTERVAL: u64 = 1000; pub const HF_ENDPOINT: &str = "https://hf-mirror.com"; /// Get the HF_HOME path, expanding `~` to the user's home directory. -/// +/// /// Uses `$HOME` on Unix and `$USERPROFILE` on Windows as fallback. /// Returns `./.agentsight/tokenizers` if home directory cannot be determined. pub fn hf_home() -> PathBuf { @@ -171,7 +171,10 @@ impl FromStr for TcpTarget { // Full wildcard shortcuts: "*", "*:*", ":*" if s == "*" || s == "*:*" || s == ":*" { - return Ok(TcpTarget { ip: None, port: None }); + return Ok(TcpTarget { + ip: None, + port: None, + }); } // Helper: parse `"*"` as wildcard, otherwise as IPv4. @@ -195,15 +198,14 @@ impl FromStr for TcpTarget { } }; - if s.starts_with(':') { + if let Some(rest) = s.strip_prefix(':') { // ":port" — port-only - let port = parse_port(&s[1..])?; + let port = parse_port(rest)?; Ok(TcpTarget { ip: None, port }) } else if s.contains(':') { // "ip:port" (either side may be `*`) - let mut parts = s.rsplitn(2, ':'); - let port_str = parts.next().unwrap(); - let ip_str = parts.next().unwrap(); + let (ip_str, port_str) = s.rsplit_once(':').unwrap(); + let ip = parse_ip(ip_str)?; let port = parse_port(port_str)?; Ok(TcpTarget { ip, port }) @@ -215,7 +217,6 @@ impl FromStr for TcpTarget { } } - /// Internal JSON structures for parsing the config file (same format as FFI). #[derive(serde::Deserialize)] struct JsonFullConfig { @@ -304,7 +305,9 @@ fn extract_rules(parsed: &JsonFullConfig) -> (Vec, Vec, for group in https_groups { for pat in &group.rule { if !pat.is_empty() { - https_rules.push(HttpsRule { pattern: pat.clone() }); + https_rules.push(HttpsRule { + pattern: pat.clone(), + }); } } } @@ -330,13 +333,15 @@ fn extract_rules(parsed: &JsonFullConfig) -> (Vec, Vec, /// Parse a JSON config string into cmdline rules, https rules, and http targets. /// /// This is the shared parser for both the config file and FFI's `load_config()`. -pub fn parse_json_rules(json: &str) -> Result<(Vec, Vec, Vec), String> { - let parsed: JsonFullConfig = serde_json::from_str(json) - .map_err(|e| format!("JSON parse error: {}", e))?; +#[allow(clippy::type_complexity)] +pub fn parse_json_rules( + json: &str, +) -> Result<(Vec, Vec, Vec), String> { + let parsed: JsonFullConfig = + serde_json::from_str(json).map_err(|e| format!("JSON parse error: {}", e))?; Ok(extract_rules(&parsed)) } - /// Ensure the agents configuration file exists at the given path. /// /// If the file does not exist, creates it with the embedded default configuration. @@ -357,8 +362,8 @@ pub fn ensure_default_agents_config(path: &Path) -> anyhow::Result<()> { /// Load default cmdline rules (embedded), without touching the filesystem. pub fn default_cmdline_rules() -> Vec { - let (rules, _, _) = parse_json_rules(DEFAULT_AGENTS_JSON) - .expect("embedded DEFAULT_AGENTS_JSON is valid"); + let (rules, _, _) = + parse_json_rules(DEFAULT_AGENTS_JSON).expect("embedded DEFAULT_AGENTS_JSON is valid"); rules } @@ -492,8 +497,13 @@ impl Default for AgentsightConfig { log_path: None, // Tokenizer defaults (read from env vars) - tokenizer_path: std::env::var("AGENTSIGHT_TOKENIZER_PATH").ok().map(PathBuf::from), - tokenizer_url: Some("https://www.modelscope.cn/models/Qwen/Qwen3.5-27B/resolve/master/tokenizer.json".to_owned()), + tokenizer_path: std::env::var("AGENTSIGHT_TOKENIZER_PATH") + .ok() + .map(PathBuf::from), + tokenizer_url: Some( + "https://www.modelscope.cn/models/Qwen/Qwen3.5-27B/resolve/master/tokenizer.json" + .to_owned(), + ), // FFI Rule defaults cmdline_rules: Vec::new(), @@ -577,8 +587,8 @@ impl AgentsightConfig { /// /// Parses `verbose`, `log_path`, `cmdline`, `https` and `http` fields. pub fn load_from_json(&mut self, json: &str) -> Result<(), String> { - let mut parsed: JsonFullConfig = serde_json::from_str(json) - .map_err(|e| format!("JSON parse error: {}", e))?; + let mut parsed: JsonFullConfig = + serde_json::from_str(json).map_err(|e| format!("JSON parse error: {}", e))?; if let Some(t) = parsed.trace_enabled { self.trace_enabled = t; @@ -607,7 +617,8 @@ impl AgentsightConfig { Err(e) => { log::warn!( "Failed to read encryption public_key_path {:?}: {}, encryption disabled", - trimmed, e + trimmed, + e ); } } @@ -674,7 +685,10 @@ impl AgentsightConfig { /// # Panics /// Panics if `config_path` was not set via `set_config_path` (CLI `--config`). pub fn resolve_config_path(&self) -> PathBuf { - assert!(self.config_path.is_some(), "config_path must be set via --config"); + assert!( + self.config_path.is_some(), + "config_path must be set via --config" + ); self.config_path.clone().unwrap() } } @@ -896,14 +910,20 @@ mod tests { fn test_set_tokenizer_path() { let config = AgentsightConfig::new() .set_tokenizer_path(Some(PathBuf::from("/path/to/tokenizer.json"))); - assert_eq!(config.tokenizer_path, Some(PathBuf::from("/path/to/tokenizer.json"))); + assert_eq!( + config.tokenizer_path, + Some(PathBuf::from("/path/to/tokenizer.json")) + ); } #[test] fn test_set_tokenizer_url() { - let config = AgentsightConfig::new() - .set_tokenizer_url(Some("https://example.com/tok.json".into())); - assert_eq!(config.tokenizer_url, Some("https://example.com/tok.json".to_string())); + let config = + AgentsightConfig::new().set_tokenizer_url(Some("https://example.com/tok.json".into())); + assert_eq!( + config.tokenizer_url, + Some("https://example.com/tok.json".to_string()) + ); } #[test] @@ -923,7 +943,10 @@ mod tests { let config = AgentsightConfig::new().add_cmdline_rule(rule); assert_eq!(config.cmdline_rules.len(), 1); assert_eq!(config.cmdline_rules[0].patterns, vec!["node", "*claude*"]); - assert_eq!(config.cmdline_rules[0].agent_name, Some("Claude Code".to_string())); + assert_eq!( + config.cmdline_rules[0].agent_name, + Some("Claude Code".to_string()) + ); assert!(config.cmdline_rules[0].allow); } @@ -942,7 +965,9 @@ mod tests { #[test] fn test_add_https_rule() { - let rule = HttpsRule { pattern: "*.openai.com".to_string() }; + let rule = HttpsRule { + pattern: "*.openai.com".to_string(), + }; let config = AgentsightConfig::new().add_https_rule(rule); assert_eq!(config.https_rules.len(), 1); assert_eq!(config.https_rules[0].pattern, "*.openai.com"); @@ -961,8 +986,12 @@ mod tests { agent_name: Some("Agent2".to_string()), allow: true, }) - .add_https_rule(HttpsRule { pattern: "*.openai.com".to_string() }) - .add_https_rule(HttpsRule { pattern: "*.anthropic.com".to_string() }); + .add_https_rule(HttpsRule { + pattern: "*.openai.com".to_string(), + }) + .add_https_rule(HttpsRule { + pattern: "*.anthropic.com".to_string(), + }); assert_eq!(config.cmdline_rules.len(), 2); assert_eq!(config.https_rules.len(), 2); } @@ -974,7 +1003,8 @@ mod tests { // All should be allow rules assert!(rules.iter().all(|r| r.allow)); // Should contain Hermes, Cosh, OpenClaw agent names - let names: Vec<&str> = rules.iter() + let names: Vec<&str> = rules + .iter() .filter_map(|r| r.agent_name.as_deref()) .collect(); assert!(names.contains(&"Hermes")); @@ -984,7 +1014,8 @@ mod tests { #[test] fn test_default_agents_json_valid() { - let (cmdline_rules, https_rules, http_targets) = parse_json_rules(DEFAULT_AGENTS_JSON).unwrap(); + let (cmdline_rules, https_rules, http_targets) = + parse_json_rules(DEFAULT_AGENTS_JSON).unwrap(); assert!(!cmdline_rules.is_empty()); assert!(https_rules.is_empty()); assert!(http_targets.is_empty()); diff --git a/src/agentsight/src/discovery/connection_scanner.rs b/src/agentsight/src/discovery/connection_scanner.rs index 48aabc3a9..11506f4c1 100644 --- a/src/agentsight/src/discovery/connection_scanner.rs +++ b/src/agentsight/src/discovery/connection_scanner.rs @@ -192,10 +192,10 @@ fn resolve_inodes_to_pids(target_inodes: &HashSet) -> HashMap { for proc in procs.flatten() { if let Ok(fds) = proc.fd() { for fd in fds.flatten() { - if let procfs::process::FDTarget::Socket(inode) = fd.target { - if target_inodes.contains(&inode) { - map.insert(inode, proc.pid() as u32); - } + if let procfs::process::FDTarget::Socket(inode) = fd.target + && target_inodes.contains(&inode) + { + map.insert(inode, proc.pid() as u32); } } } diff --git a/src/agentsight/src/discovery/matcher.rs b/src/agentsight/src/discovery/matcher.rs index 024588e35..a77c96d80 100644 --- a/src/agentsight/src/discovery/matcher.rs +++ b/src/agentsight/src/discovery/matcher.rs @@ -160,7 +160,10 @@ mod tests { #[test] fn test_cmdline_glob_matcher() { - let matcher = CmdlineGlobMatcher::new("Claude Code", vec!["node".to_string(), "*claude*".to_string()]); + let matcher = CmdlineGlobMatcher::new( + "Claude Code", + vec!["node".to_string(), "*claude*".to_string()], + ); let ctx = ProcessContext { comm: "node".to_string(), cmdline_args: vec!["node".to_string(), "/path/claude-code".to_string()], diff --git a/src/agentsight/src/discovery/scanner.rs b/src/agentsight/src/discovery/scanner.rs index f1873db90..80db36205 100644 --- a/src/agentsight/src/discovery/scanner.rs +++ b/src/agentsight/src/discovery/scanner.rs @@ -36,11 +36,11 @@ impl AgentScanner { pub fn from_rules(cmdline_rules: &[CmdlineRule], https_rules: &[HttpsRule]) -> Self { let matchers: Vec = cmdline_rules .iter() - .filter_map(|rule| CmdlineGlobMatcher::from_config(rule)) + .filter_map(CmdlineGlobMatcher::from_config) .collect(); let deny_matchers: Vec = cmdline_rules .iter() - .filter_map(|r| CmdlineGlobMatcher::from_deny_rule(r)) + .filter_map(CmdlineGlobMatcher::from_deny_rule) .collect(); let domain_patterns: Vec = https_rules.iter().map(|r| r.pattern.clone()).collect(); Self { diff --git a/src/agentsight/src/event.rs b/src/agentsight/src/event.rs index fde41b4c1..2d7e5ff85 100644 --- a/src/agentsight/src/event.rs +++ b/src/agentsight/src/event.rs @@ -1,8 +1,8 @@ -use crate::probes::proctrace::VariableEvent as ProcEvent; -use crate::probes::sslsniff::SslEvent; -use crate::probes::procmon::Event as ProcMonEvent; use crate::probes::filewatch::FileWatchEvent; use crate::probes::filewrite::FileWriteEvent; +use crate::probes::procmon::Event as ProcMonEvent; +use crate::probes::proctrace::VariableEvent as ProcEvent; +use crate::probes::sslsniff::SslEvent; use crate::probes::udpdns::UdpDnsEvent; /// Unified event type that can represent any probe event diff --git a/src/agentsight/src/ffi.rs b/src/agentsight/src/ffi.rs index ee0d9433b..797953097 100644 --- a/src/agentsight/src/ffi.rs +++ b/src/agentsight/src/ffi.rs @@ -3,13 +3,14 @@ //! Implements the **eventfd + read** model described in `docs/design-docs/c-ffi-api.md`: //! AgentSight runs a background pipeline thread; completed events are pushed //! into an `mpsc` channel and the caller is notified via `eventfd`. +#![allow(clippy::missing_safety_doc, clippy::large_enum_variant)] //! The caller consumes events by calling `agentsight_read()` with callbacks. use std::ffi::{CStr, CString, c_char, c_int, c_void}; use std::ptr; +use std::sync::Arc; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::mpsc; -use std::sync::Arc; use crate::analyzer::HttpRecord; use crate::config::AgentsightConfig; @@ -256,13 +257,24 @@ fn build_https_data(record: &HttpRecord) -> HttpsDataHolder { fn build_llm_data(call: &LLMCall) -> LlmDataHolder { let response_id = call.metadata.get("response_id").map(|s| safe_cstring(s)); - let conversation_id = call.metadata.get("conversation_id").map(|s| safe_cstring(s)); + let conversation_id = call + .metadata + .get("conversation_id") + .map(|s| safe_cstring(s)); let session_id = call.metadata.get("session_id").map(|s| safe_cstring(s)); let agent_name = call.agent_name.as_ref().map(|s| safe_cstring(s)); // Construct request_url from metadata - let server_addr = call.metadata.get("server.address").cloned().unwrap_or_default(); - let server_port = call.metadata.get("server.port").cloned().unwrap_or_default(); + let server_addr = call + .metadata + .get("server.address") + .cloned() + .unwrap_or_default(); + let server_port = call + .metadata + .get("server.port") + .cloned() + .unwrap_or_default(); let path = call.metadata.get("path").cloned().unwrap_or_default(); let url = if server_port.is_empty() { format!("https://{}{}", server_addr, path) @@ -279,10 +291,7 @@ fn build_llm_data(call: &LLMCall) -> LlmDataHolder { .get("status_code") .and_then(|s| s.parse().ok()) .unwrap_or(0); - let is_sse: bool = call - .metadata - .get("is_sse") - .map_or(false, |s| s == "true"); + let is_sse: bool = call.metadata.get("is_sse").is_some_and(|s| s == "true"); let finish_reason = call .response @@ -304,20 +313,16 @@ fn build_llm_data(call: &LLMCall) -> LlmDataHolder { None => (false, 0, 0, 0, 0, 0), }; - let req_messages_json = - serde_json::to_string(&call.request.messages).unwrap_or_default(); - let resp_messages_json = - serde_json::to_string(&call.response.messages).unwrap_or_default(); + let req_messages_json = serde_json::to_string(&call.request.messages).unwrap_or_default(); + let resp_messages_json = serde_json::to_string(&call.response.messages).unwrap_or_default(); let req_messages = safe_cstring(&req_messages_json); let resp_messages = safe_cstring(&resp_messages_json); // Incremental (latest-round) input messages: the same per-round increment // stored in SQLite (`genai_events.input_messages`). Drops system messages // and keeps everything from the last `user` message onward. - let input_delta = - crate::genai::semantic::latest_round_input_messages(&call.request.messages); - let input_message_delta_json = - serde_json::to_string(&input_delta).unwrap_or_default(); + let input_delta = crate::genai::semantic::latest_round_input_messages(&call.request.messages); + let input_message_delta_json = serde_json::to_string(&input_delta).unwrap_or_default(); let input_message_delta = safe_cstring(&input_message_delta_json); let tools_json = call @@ -414,9 +419,7 @@ unsafe fn dispatch_event( /// The pointer is valid until the next API call on the same thread. #[unsafe(no_mangle)] pub extern "C" fn agentsight_last_error() -> *const c_char { - LAST_ERROR.with(|e| { - e.borrow().as_ref().map_or(ptr::null(), |s| s.as_ptr()) - }) + LAST_ERROR.with(|e| e.borrow().as_ref().map_or(ptr::null(), |s| s.as_ptr())) } // ---- Configuration ---- @@ -587,9 +590,7 @@ pub unsafe extern "C" fn agentsight_config_free(cfg: *mut AgentsightConfigHandle /// Create a new AgentSight handle. Does NOT start the pipeline yet. /// Returns NULL on failure (call `agentsight_last_error()` for details). #[unsafe(no_mangle)] -pub unsafe extern "C" fn agentsight_new( - cfg: *mut AgentsightConfigHandle, -) -> *mut AgentsightHandle { +pub unsafe extern "C" fn agentsight_new(cfg: *mut AgentsightConfigHandle) -> *mut AgentsightHandle { // Create eventfd let efd = unsafe { libc::eventfd(0, libc::EFD_NONBLOCK | libc::EFD_CLOEXEC) }; if efd < 0 { @@ -667,7 +668,10 @@ fn ffi_background_thread( let mut sight = match AgentSight::new(config) { Ok(s) => s, Err(e) => { - log::error!("agentsight background thread: AgentSight::new failed: {}", e); + log::error!( + "agentsight background thread: AgentSight::new failed: {}", + e + ); return; } }; @@ -833,7 +837,10 @@ mod tests { assert!(cfg.load_from_json(json).is_ok()); assert_eq!(cfg.cmdline_rules.len(), 2); assert!(cfg.cmdline_rules[0].allow); - assert_eq!(cfg.cmdline_rules[0].agent_name, Some("Claude Code".to_string())); + assert_eq!( + cfg.cmdline_rules[0].agent_name, + Some("Claude Code".to_string()) + ); assert!(!cfg.cmdline_rules[1].allow); assert!(cfg.cmdline_rules[1].agent_name.is_none()); } diff --git a/src/agentsight/src/genai/builder.rs b/src/agentsight/src/genai/builder.rs index 4a4698bcc..6f4efd895 100644 --- a/src/agentsight/src/genai/builder.rs +++ b/src/agentsight/src/genai/builder.rs @@ -3,25 +3,23 @@ //! This module builds GenAI semantic events from AnalysisResult. //! It reuses already-extracted data to avoid redundant parsing. -use crate::analyzer::{ - AnalysisResult, TokenRecord, ParsedApiMessage, HttpRecord, +use super::semantic::{ + GenAISemanticEvent, InputMessage, LLMCall, LLMRequest, LLMResponse, MessagePart, OutputMessage, + TokenUsage, }; -use crate::analyzer::message::types::OpenAIChatMessage; use crate::aggregator::{ConnectionId, ParsedRequest}; +use crate::analyzer::message::types::OpenAIChatMessage; use crate::analyzer::token::TokenParser; -use crate::discovery::matcher::{ProcessContext, CmdlineGlobMatcher}; +use crate::analyzer::{AnalysisResult, HttpRecord, ParsedApiMessage, TokenRecord}; use crate::config::default_cmdline_rules; +use crate::discovery::matcher::{CmdlineGlobMatcher, ProcessContext}; use crate::parser::sse::ParsedSseEvent; use crate::response_map::ResponseSessionMapper; use crate::storage::sqlite::{PendingCallInfo, SseEnrichment}; -use super::semantic::{ - GenAISemanticEvent, LLMCall, LLMRequest, LLMResponse, - InputMessage, OutputMessage, MessagePart, TokenUsage, -}; +use sha2::{Digest, Sha256}; use std::collections::HashMap; use std::sync::atomic::{AtomicU64, Ordering}; use std::time::{SystemTime, UNIX_EPOCH}; -use sha2::{Sha256, Digest}; /// Output from `GenAIBuilder::build()`, containing built events and deferred resolution info. pub struct BuildOutput { @@ -83,7 +81,8 @@ impl GenAIBuilder { let mut pending: Option = None; let mut pending_response_id = None; - if let Some(llm_call) = self.build_llm_call(results, response_mapper, pid_agent_name_cache) { + if let Some(llm_call) = self.build_llm_call(results, response_mapper, pid_agent_name_cache) + { // Build PendingCallInfo from the same LLMCall before moving it let http_record = results.iter().find_map(|r| match r { AnalysisResult::Http(h) => Some(h.clone()), @@ -92,21 +91,33 @@ impl GenAIBuilder { // Extract input messages for the pending record let (input_messages_json, system_instructions_json) = { - let sys: Vec<_> = llm_call.request.messages.iter() - .filter(|m| m.role == "system").collect(); - let latest = crate::genai::semantic::latest_round_input_messages( - &llm_call.request.messages, - ); + let sys: Vec<_> = llm_call + .request + .messages + .iter() + .filter(|m| m.role == "system") + .collect(); + let latest = + crate::genai::semantic::latest_round_input_messages(&llm_call.request.messages); ( - if latest.is_empty() { None } else { serde_json::to_string(&latest).ok() }, - if sys.is_empty() { None } else { serde_json::to_string(&sys).ok() }, + if latest.is_empty() { + None + } else { + serde_json::to_string(&latest).ok() + }, + if sys.is_empty() { + None + } else { + serde_json::to_string(&sys).ok() + }, ) }; // Determine response_id from call metadata (may come from parsed_message // or SSE body fallback), and check if mapper resolved it. let response_id = llm_call.metadata.get("response_id").cloned(); - let mapper_hit = response_id.as_deref() + let mapper_hit = response_id + .as_deref() .and_then(|rid| response_mapper.get_session_by_response_id(rid)) .is_some(); @@ -142,7 +153,13 @@ impl GenAIBuilder { events.push(GenAISemanticEvent::LLMCall(llm_call)); } - (BuildOutput { events, pending_response_id }, pending) + ( + BuildOutput { + events, + pending_response_id, + }, + pending, + ) } /// Build a `PendingCallInfo` directly from a raw `ParsedRequest` and @@ -164,7 +181,11 @@ impl GenAIBuilder { ) -> Option { // Only process known LLM API paths let path_match = self.is_llm_api_path(&request.path); - let body_str = if request.body_len > 0 { Some(request.body_str().to_string()) } else { None }; + let body_str = if request.body_len > 0 { + Some(request.body_str().to_string()) + } else { + None + }; let body_match = !path_match && Self::is_sysom_pop_request(&body_str); if !path_match && !body_match { return None; @@ -174,7 +195,8 @@ impl GenAIBuilder { let body = request.json_body(); // Determine if streaming - let is_sse = body.as_ref() + let is_sse = body + .as_ref() .and_then(|v| v.get("stream")) .and_then(|v| v.as_bool()) .unwrap_or(false); @@ -190,11 +212,14 @@ impl GenAIBuilder { // "content": [{"type":"text","text":"..."},...] let extract_text = |m: &serde_json::Value| -> Option { let c = m.get("content")?; - if let Some(s) = c.as_str() { - if !s.is_empty() { return Some(s.to_string()); } + if let Some(s) = c.as_str() + && !s.is_empty() + { + return Some(s.to_string()); } if let Some(arr) = c.as_array() { - let text: String = arr.iter() + let text: String = arr + .iter() .filter_map(|item| { // [{"type":"text","text":"..."}] if item.get("type").and_then(|t| t.as_str()) == Some("text") { @@ -205,16 +230,19 @@ impl GenAIBuilder { }) .collect::>() .join("\n"); - if !text.is_empty() { return Some(text); } + if !text.is_empty() { + return Some(text); + } } None }; // session_id: SHA256 of first user message text (same logic // as compute_session_id but operating on raw JSON values) - let first_user_text = messages.iter() + let first_user_text = messages + .iter() .filter(|m| m.get("role").and_then(|r| r.as_str()) == Some("user")) - .find_map(|m| extract_text(m)) + .find_map(&extract_text) .unwrap_or_default(); let session_id = if !first_user_text.is_empty() { @@ -226,10 +254,11 @@ impl GenAIBuilder { // Last user message raw text — used for both conversation_id // (fingerprint hash) and user_query (display text) - let last_user_raw = messages.iter() + let last_user_raw = messages + .iter() .rev() .filter(|m| m.get("role").and_then(|r| r.as_str()) == Some("user")) - .find_map(|m| extract_text(m)); + .find_map(extract_text); // conversation_id: SHA256 of last user message raw text // (same logic as compute_user_query_fingerprint) @@ -239,14 +268,15 @@ impl GenAIBuilder { }); // user_query: last user message text, stripped of metadata prefix - let user_query = last_user_raw.as_deref() - .map(|s| Self::strip_user_query_prefix(s)); + let user_query = last_user_raw.as_deref().map(Self::strip_user_query_prefix); // Serialise message subsets for the pending record - let sys: Vec<_> = messages.iter() + let sys: Vec<_> = messages + .iter() .filter(|m| m.get("role").and_then(|r| r.as_str()) == Some("system")) .collect(); - let non_sys: Vec<_> = messages.iter() + let non_sys: Vec<_> = messages + .iter() .filter(|m| m.get("role").and_then(|r| r.as_str()) != Some("system")) .collect(); @@ -261,7 +291,13 @@ impl GenAIBuilder { serde_json::to_string(&sys).ok() }; - (session_id, conversation_id, user_query, input_messages, system_instructions) + ( + session_id, + conversation_id, + user_query, + input_messages, + system_instructions, + ) } else { // messages key missing or not an array (None, None, None, None, None) @@ -271,7 +307,8 @@ impl GenAIBuilder { }; // Extract model from request body JSON "model" field - let model = body.as_ref() + let model = body + .as_ref() .and_then(|v| v.get("model")) .and_then(|m| m.as_str()) .filter(|s| !s.is_empty()) @@ -281,13 +318,17 @@ impl GenAIBuilder { let provider = self.extract_provider_from_path(&request.path); // Resolve agent_name: check pid→name cache first, then comm-based matching, then comm as fallback - let agent_name = Self::resolve_agent_name_from_comm(&request.source_event.comm, conn_id.pid as u32, pid_agent_name_cache) - .or_else(|| Some(request.source_event.comm_str())); + let agent_name = Self::resolve_agent_name_from_comm( + &request.source_event.comm, + conn_id.pid, + pid_agent_name_cache, + ) + .or_else(|| Some(request.source_event.comm_str())); Some(PendingCallInfo { call_id, - trace_id: None, // LLM API response_id, not available until response - conversation_id, // User query fingerprint hash (from request body) + trace_id: None, // LLM API response_id, not available until response + conversation_id, // User query fingerprint hash (from request body) session_id, start_timestamp_ns: request.source_event.timestamp_ns, pid: conn_id.pid as i32, @@ -330,28 +371,26 @@ impl GenAIBuilder { } if let Some(json) = event.json_body() { // Extract model from first chunk that has it - if model.is_none() { - if let Some(m) = json.get("model").and_then(|v| v.as_str()) { - if !m.is_empty() { - model = Some(m.to_string()); - } - } + if model.is_none() + && let Some(m) = json.get("model").and_then(|v| v.as_str()) + && !m.is_empty() + { + model = Some(m.to_string()); } // Extract response id (trace_id) from first chunk that has it - if trace_id.is_none() { - if let Some(id) = json.get("id").and_then(|v| v.as_str()) { - if !id.is_empty() { - trace_id = Some(id.to_string()); - } - } + if trace_id.is_none() + && let Some(id) = json.get("id").and_then(|v| v.as_str()) + && !id.is_empty() + { + trace_id = Some(id.to_string()); } // Accumulate content deltas if let Some(choices) = json.get("choices").and_then(|v| v.as_array()) { for choice in choices { - if let Some(delta) = choice.get("delta") { - if let Some(c) = delta.get("content").and_then(|v| v.as_str()) { - content_buf.push_str(c); - } + if let Some(delta) = choice.get("delta") + && let Some(c) = delta.get("content").and_then(|v| v.as_str()) + { + content_buf.push_str(c); } } } @@ -359,7 +398,9 @@ impl GenAIBuilder { } // Reverse scan for token usage (usage chunk is near the end) - let usage = sse_events.iter().rev() + let usage = sse_events + .iter() + .rev() .find_map(|e| token_parser.parse_event(e)); let (input_tokens, output_tokens) = match &usage { @@ -368,10 +409,10 @@ impl GenAIBuilder { }; // Use model from usage if not found in content chunks - if model.is_none() { - if let Some(ref u) = usage { - model = u.model.clone(); - } + if model.is_none() + && let Some(ref u) = usage + { + model = u.model.clone(); } // Build output_messages JSON from accumulated content @@ -380,7 +421,8 @@ impl GenAIBuilder { serde_json::to_string(&serde_json::json!([{ "role": "assistant", "parts": [{"Text": {"content": content_buf}}] - }])).ok() + }])) + .ok() } else { None }; @@ -401,7 +443,12 @@ impl GenAIBuilder { /// Build LLMCall from analysis results /// /// Combines data from TokenRecord, HttpRecord, and ParsedApiMessage - fn build_llm_call(&self, results: &[AnalysisResult], response_mapper: &ResponseSessionMapper, pid_agent_name_cache: &std::collections::HashMap) -> Option { + fn build_llm_call( + &self, + results: &[AnalysisResult], + response_mapper: &ResponseSessionMapper, + pid_agent_name_cache: &std::collections::HashMap, + ) -> Option { // Extract components from analysis results let token_record = results.iter().find_map(|r| match r { AnalysisResult::Token(t) => Some(t.clone()), @@ -420,7 +467,7 @@ impl GenAIBuilder { // Need at least HttpRecord to build LLMCall let http = http_record?; - + // Check if this is an LLM API call (path-based or body-based for SysOM POP API) let path_match = self.is_llm_api_path(&http.path); let body_match = !path_match && Self::is_sysom_pop_request(&http.request_body); @@ -447,16 +494,21 @@ impl GenAIBuilder { // Determine provider and model // Priority: path-based (most specific) > body-based > parsed_message > token_record - let provider = self.extract_provider_from_path(&http.path) + let provider = self + .extract_provider_from_path(&http.path) .or_else(|| Self::extract_provider_from_body(&http.request_body)) .or_else(|| parsed_message.as_ref().map(|m| m.provider().to_string())) .or_else(|| token_record.as_ref().map(|t| t.provider.clone())) .unwrap_or_else(|| "unknown".to_string()); // Model priority: parsed_message (most accurate) > token_record > body extraction - let model = self.extract_model_from_message(&parsed_message) - .or_else(|| token_record.as_ref() - .and_then(|t| t.model.as_ref().filter(|m| !m.is_empty()).cloned())) + let model = self + .extract_model_from_message(&parsed_message) + .or_else(|| { + token_record + .as_ref() + .and_then(|t| t.model.as_ref().filter(|m| !m.is_empty()).cloned()) + }) .or_else(|| Self::extract_model_from_body(&http.request_body, &http.response_body)) .unwrap_or_else(|| "unknown".to_string()); @@ -465,11 +517,16 @@ impl GenAIBuilder { let user_query = Self::extract_last_user_query(&request); // session_id: 优先从 agent 自身的 session 获取（通过 response ID → .jsonl UUID 映射）， // fallback 到基于首条 user message 的 hash 计算 - let response_id_val = parsed_message.as_ref().and_then(|m| m.response_id()).map(|s| s.to_string()); - let mapper_session = response_id_val.as_deref() + let response_id_val = parsed_message + .as_ref() + .and_then(|m| m.response_id()) + .map(|s| s.to_string()); + let mapper_session = response_id_val + .as_deref() .and_then(|rid| response_mapper.get_session_by_response_id(rid)) .map(|s| s.to_string()); - let session_id = mapper_session.clone() + let session_id = mapper_session + .clone() .unwrap_or_else(|| Self::compute_session_id(&request)); // 提取 LLM API 的 response_id（如 chatcmpl-xxx），用作 trace_id @@ -512,21 +569,25 @@ impl GenAIBuilder { if let Some(msg) = inner.get("message").and_then(|m| m.as_str()) { return Some(msg.to_string()); } - if let Some(inner_e) = inner.get("error") { - if let Some(msg) = inner_e.get("message").and_then(|m| m.as_str()) { - return Some(msg.to_string()); - } + if let Some(inner_e) = inner.get("error") + && let Some(msg) = + inner_e.get("message").and_then(|m| m.as_str()) + { + return Some(msg.to_string()); } } return Some(s.to_string()); } } // Top-level {"message":"..."} - v.get("message").and_then(|m| m.as_str()).map(|s| s.to_string()) + v.get("message") + .and_then(|m| m.as_str()) + .map(|s| s.to_string()) } let json_str = strip_chunked(body); - serde_json::from_str::(json_str).ok() + serde_json::from_str::(json_str) + .ok() .and_then(|v| extract_message(&v)) .or_else(|| Some(body.clone())) }) @@ -555,16 +616,20 @@ impl GenAIBuilder { meta.insert("path".to_string(), http.path.clone()); meta.insert("status_code".to_string(), http.status_code.to_string()); meta.insert("is_sse".to_string(), http.is_sse.to_string()); - meta.insert("sse_event_count".to_string(), http.sse_event_count.to_string()); + meta.insert( + "sse_event_count".to_string(), + http.sse_event_count.to_string(), + ); // Extract server.address and server.port from Host header - if let Ok(headers) = serde_json::from_str::>(&http.request_headers) { - if let Some(host) = headers.get("host").or_else(|| headers.get("Host")) { - if let Some((addr, port)) = host.rsplit_once(':') { - meta.insert("server.address".to_string(), addr.to_string()); - meta.insert("server.port".to_string(), port.to_string()); - } else { - meta.insert("server.address".to_string(), host.clone()); - } + if let Ok(headers) = + serde_json::from_str::>(&http.request_headers) + && let Some(host) = headers.get("host").or_else(|| headers.get("Host")) + { + if let Some((addr, port)) = host.rsplit_once(':') { + meta.insert("server.address".to_string(), addr.to_string()); + meta.insert("server.port".to_string(), port.to_string()); + } else { + meta.insert("server.address".to_string(), host.clone()); } } // Derive gen_ai.operation.name from path @@ -595,9 +660,7 @@ impl GenAIBuilder { match message { Some(ParsedApiMessage::OpenAICompletion { request, .. }) => { if let Some(req) = request.as_ref() { - let msgs = req.messages.iter().map(|m| { - Self::openai_msg_to_input(m) - }).collect(); + let msgs = req.messages.iter().map(Self::openai_msg_to_input).collect(); return LLMRequest { messages: msgs, temperature: req.temperature, @@ -616,14 +679,20 @@ impl GenAIBuilder { } Some(ParsedApiMessage::AnthropicMessage { request, .. }) => { if let Some(req) = request.as_ref() { - let msgs = req.messages.iter().map(|m| { - let role = format!("{:?}", m.role).to_lowercase(); - InputMessage { - role, - parts: vec![MessagePart::Text { content: m.content.as_text() }], - name: None, - } - }).collect(); + let msgs = req + .messages + .iter() + .map(|m| { + let role = format!("{:?}", m.role).to_lowercase(); + InputMessage { + role, + parts: vec![MessagePart::Text { + content: m.content.as_text(), + }], + name: None, + } + }) + .collect(); return LLMRequest { messages: msgs, temperature: req.temperature, @@ -642,33 +711,48 @@ impl GenAIBuilder { } Some(ParsedApiMessage::SysomMessage { request, .. }) => { if let Some(req) = request.as_ref() { - let msgs = req.params.messages.iter().map(|m| { - let role = m.role.clone(); - let mut parts = Vec::new(); - if role == "tool" { - let response_val = serde_json::from_str::(&m.content) - .unwrap_or_else(|_| serde_json::Value::String(m.content.clone())); - parts.push(MessagePart::ToolCallResponse { - id: m.tool_call_id.clone(), - response: response_val, - }); - } else { - if !m.content.is_empty() { - parts.push(MessagePart::Text { content: m.content.clone() }); - } - } - if let Some(ref tool_calls) = m.tool_calls { - for tc in tool_calls { - let arguments = serde_json::from_str::(&tc.function.arguments).ok(); - parts.push(MessagePart::ToolCall { - id: Some(tc.id.clone()), - name: tc.function.name.clone(), - arguments, + let msgs = req + .params + .messages + .iter() + .map(|m| { + let role = m.role.clone(); + let mut parts = Vec::new(); + if role == "tool" { + let response_val = + serde_json::from_str::(&m.content) + .unwrap_or_else(|_| { + serde_json::Value::String(m.content.clone()) + }); + parts.push(MessagePart::ToolCallResponse { + id: m.tool_call_id.clone(), + response: response_val, + }); + } else if !m.content.is_empty() { + parts.push(MessagePart::Text { + content: m.content.clone(), }); } - } - InputMessage { role, parts, name: m.name.clone() } - }).collect(); + if let Some(ref tool_calls) = m.tool_calls { + for tc in tool_calls { + let arguments = serde_json::from_str::( + &tc.function.arguments, + ) + .ok(); + parts.push(MessagePart::ToolCall { + id: Some(tc.id.clone()), + name: tc.function.name.clone(), + arguments, + }); + } + } + InputMessage { + role, + parts, + name: m.name.clone(), + } + }) + .collect(); return LLMRequest { messages: msgs, temperature: req.params.temperature, @@ -689,10 +773,10 @@ impl GenAIBuilder { } // Fallback: no parsed message — parse request_body directly - if let Some(ref body) = http.request_body { - if let Some(req) = Self::parse_request_body(body) { - return req; - } + if let Some(ref body) = http.request_body + && let Some(req) = Self::parse_request_body(body) + { + return req; } LLMRequest { messages: vec![], @@ -716,52 +800,75 @@ impl GenAIBuilder { let obj = v.as_object()?; // 解析 messages 数组 - let messages = obj.get("messages") + let messages = obj + .get("messages") .and_then(|m| m.as_array()) .map(|arr| { - arr.iter().filter_map(|msg| { - let role = msg.get("role")?.as_str()?.to_string(); - let mut parts = Vec::new(); - - // content 可以是字符串或数组 - if let Some(content) = msg.get("content") { - if let Some(s) = content.as_str() { - if !s.is_empty() { - parts.push(MessagePart::Text { content: s.to_string() }); - } - } else if let Some(arr) = content.as_array() { - for item in arr { - if let Some(text) = item.get("text").and_then(|t| t.as_str()) { - parts.push(MessagePart::Text { content: text.to_string() }); + arr.iter() + .filter_map(|msg| { + let role = msg.get("role")?.as_str()?.to_string(); + let mut parts = Vec::new(); + + // content 可以是字符串或数组 + if let Some(content) = msg.get("content") { + if let Some(s) = content.as_str() { + if !s.is_empty() { + parts.push(MessagePart::Text { + content: s.to_string(), + }); + } + } else if let Some(arr) = content.as_array() { + for item in arr { + if let Some(text) = item.get("text").and_then(|t| t.as_str()) { + parts.push(MessagePart::Text { + content: text.to_string(), + }); + } } } } - } - // tool_call 结果 (role=tool) - if role == "tool" { - if let Some(content) = msg.get("content") { - let id = msg.get("tool_call_id").and_then(|v| v.as_str()).map(|s| s.to_string()); + // tool_call 结果 (role=tool) + if role == "tool" + && let Some(content) = msg.get("content") + { + let id = msg + .get("tool_call_id") + .and_then(|v| v.as_str()) + .map(|s| s.to_string()); parts = vec![MessagePart::ToolCallResponse { id, response: content.clone(), }]; } - } - // tool_calls (role=assistant 发起的 tool calls) - if let Some(tool_calls) = msg.get("tool_calls").and_then(|v| v.as_array()) { - for tc in tool_calls { - let id = tc.get("id").and_then(|v| v.as_str()).map(|s| s.to_string()); - let func = tc.get("function").unwrap_or(&serde_json::Value::Null); - let name = func.get("name").and_then(|v| v.as_str()).unwrap_or("").to_string(); - let arguments = func.get("arguments").map(|v| v.clone()); - parts.push(MessagePart::ToolCall { id, name, arguments }); + // tool_calls (role=assistant 发起的 tool calls) + if let Some(tool_calls) = msg.get("tool_calls").and_then(|v| v.as_array()) { + for tc in tool_calls { + let id = + tc.get("id").and_then(|v| v.as_str()).map(|s| s.to_string()); + let func = tc.get("function").unwrap_or(&serde_json::Value::Null); + let name = func + .get("name") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + let arguments = func.get("arguments").cloned(); + parts.push(MessagePart::ToolCall { + id, + name, + arguments, + }); + } } - } - Some(InputMessage { role, parts, name: None }) - }).collect::>() + Some(InputMessage { + role, + parts, + name: None, + }) + }) + .collect::>() }) .unwrap_or_default(); @@ -769,21 +876,29 @@ impl GenAIBuilder { return None; } - let tools = obj.get("tools") + let tools = obj + .get("tools") .and_then(|v| v.as_array()) .map(|arr| arr.to_vec()); Some(LLMRequest { messages, temperature: obj.get("temperature").and_then(|v| v.as_f64()), - max_tokens: obj.get("max_tokens").and_then(|v| v.as_u64()).map(|v| v as u32), + max_tokens: obj + .get("max_tokens") + .and_then(|v| v.as_u64()) + .map(|v| v as u32), frequency_penalty: obj.get("frequency_penalty").and_then(|v| v.as_f64()), presence_penalty: obj.get("presence_penalty").and_then(|v| v.as_f64()), top_p: obj.get("top_p").and_then(|v| v.as_f64()), top_k: obj.get("top_k").and_then(|v| v.as_f64()), seed: obj.get("seed").and_then(|v| v.as_i64()), stop_sequences: obj.get("stop").and_then(|v| { - v.as_array().map(|arr| arr.iter().filter_map(|s| s.as_str().map(String::from)).collect()) + v.as_array().map(|arr| { + arr.iter() + .filter_map(|s| s.as_str().map(String::from)) + .collect() + }) }), stream: obj.get("stream").and_then(|v| v.as_bool()).unwrap_or(false), tools, @@ -792,62 +907,92 @@ impl GenAIBuilder { } /// Build LLMResponse from parsed message or HTTP record - fn build_response(&self, message: &Option, http: &HttpRecord, _token_record: &Option) -> LLMResponse { + fn build_response( + &self, + message: &Option, + http: &HttpRecord, + _token_record: &Option, + ) -> LLMResponse { // Try to extract from parsed message first let (messages, finish_reason): (Vec, Option) = match message { - Some(ParsedApiMessage::OpenAICompletion { response, .. }) => { - response.as_ref().map(|resp| { - let msgs: Vec = resp.choices.iter().map(|c| { - Self::openai_msg_to_output(&c.message, c.finish_reason.as_deref()) - }).collect(); + Some(ParsedApiMessage::OpenAICompletion { response, .. }) => response + .as_ref() + .map(|resp| { + let msgs: Vec = resp + .choices + .iter() + .map(|c| Self::openai_msg_to_output(&c.message, c.finish_reason.as_deref())) + .collect(); let finish = resp.choices.first().and_then(|c| c.finish_reason.clone()); (msgs, finish) - }).unwrap_or_else(|| (vec![], None)) - } + }) + .unwrap_or_else(|| (vec![], None)), Some(ParsedApiMessage::AnthropicMessage { response, .. }) => { - response.as_ref().map(|resp| { - let mut parts = Vec::new(); - for block in &resp.content { - match block { - crate::analyzer::message::AnthropicContentBlock::Text { text, .. } => { - if !text.is_empty() { - parts.push(MessagePart::Text { content: text.clone() }); + response + .as_ref() + .map(|resp| { + let mut parts = Vec::new(); + for block in &resp.content { + match block { + crate::analyzer::message::AnthropicContentBlock::Text { + text, + .. + } => { + if !text.is_empty() { + parts.push(MessagePart::Text { + content: text.clone(), + }); + } } - } - crate::analyzer::message::AnthropicContentBlock::ToolUse { id, name, input } => { - // Anthropic tool_use: convert to MessagePart::ToolCall - parts.push(MessagePart::ToolCall { - id: Some(id.clone()), - name: name.clone(), - arguments: Some(input.clone()), - }); - } - crate::analyzer::message::AnthropicContentBlock::ToolResult { tool_use_id, content, .. } => { - // Anthropic tool_result: convert to MessagePart::ToolCallResponse - let response_val = content.clone().unwrap_or(serde_json::Value::Null); - parts.push(MessagePart::ToolCallResponse { - id: Some(tool_use_id.clone()), - response: response_val, - }); - } - crate::analyzer::message::AnthropicContentBlock::Thinking { thinking, .. } => { - // Anthropic thinking: convert to MessagePart::Reasoning - if !thinking.is_empty() { - parts.push(MessagePart::Reasoning { content: thinking.clone() }); + crate::analyzer::message::AnthropicContentBlock::ToolUse { + id, + name, + input, + } => { + // Anthropic tool_use: convert to MessagePart::ToolCall + parts.push(MessagePart::ToolCall { + id: Some(id.clone()), + name: name.clone(), + arguments: Some(input.clone()), + }); + } + crate::analyzer::message::AnthropicContentBlock::ToolResult { + tool_use_id, + content, + .. + } => { + // Anthropic tool_result: convert to MessagePart::ToolCallResponse + let response_val = + content.clone().unwrap_or(serde_json::Value::Null); + parts.push(MessagePart::ToolCallResponse { + id: Some(tool_use_id.clone()), + response: response_val, + }); + } + crate::analyzer::message::AnthropicContentBlock::Thinking { + thinking, + .. + } => { + // Anthropic thinking: convert to MessagePart::Reasoning + if !thinking.is_empty() { + parts.push(MessagePart::Reasoning { + content: thinking.clone(), + }); + } } + _ => {} } - _ => {} } - } - let msgs = vec![OutputMessage { - role: "assistant".to_string(), - parts, - name: None, - finish_reason: resp.stop_reason.clone(), - }]; - let finish = resp.stop_reason.clone(); - (msgs, finish) - }).unwrap_or_else(|| (vec![], None)) + let msgs = vec![OutputMessage { + role: "assistant".to_string(), + parts, + name: None, + finish_reason: resp.stop_reason.clone(), + }]; + let finish = resp.stop_reason.clone(); + (msgs, finish) + }) + .unwrap_or_else(|| (vec![], None)) } _ => (vec![], None), }; @@ -855,17 +1000,23 @@ impl GenAIBuilder { // SysOM response handling let (messages, finish_reason) = if messages.is_empty() { match message { - Some(ParsedApiMessage::SysomMessage { response, .. }) => { - response.as_ref().map(|resp| { + Some(ParsedApiMessage::SysomMessage { response, .. }) => response + .as_ref() + .map(|resp| { let choice = resp.choices.first(); let mut parts = Vec::new(); if let Some(choice) = choice { if !choice.message.content.is_empty() { - parts.push(MessagePart::Text { content: choice.message.content.clone() }); + parts.push(MessagePart::Text { + content: choice.message.content.clone(), + }); } if let Some(ref tool_use) = choice.message.tool_use { for item in tool_use { - let arguments = serde_json::from_str::(&item.function.arguments).ok(); + let arguments = serde_json::from_str::( + &item.function.arguments, + ) + .ok(); parts.push(MessagePart::ToolCall { id: Some(item.id.clone()), name: item.function.name.clone(), @@ -885,8 +1036,8 @@ impl GenAIBuilder { }] }; (msgs, Some("stop".to_string())) - }).unwrap_or_else(|| (vec![], None)) - } + }) + .unwrap_or_else(|| (vec![], None)), _ => (messages, finish_reason), } } else { @@ -897,45 +1048,57 @@ impl GenAIBuilder { let messages = if messages.is_empty() && http.is_sse { // No parsed response — reconstruct from SSE response body directly if let Some(ref body) = http.response_body { - Self::parse_sse_response_body(body, finish_reason.as_deref()) - .unwrap_or(messages) + Self::parse_sse_response_body(body, finish_reason.as_deref()).unwrap_or(messages) } else { messages } } else if http.is_sse { // Has parsed response but may be missing reasoning/tool_calls — enrich from SSE body let mut msgs = messages; - if let Some(ref body) = http.response_body { - if let Some(msg) = msgs.first_mut() { - if msg.role == "assistant" { - let has_reasoning = msg.parts.iter().any(|p| matches!(p, MessagePart::Reasoning { .. })); - // Check if any tool_call is missing id - let has_tool_calls_without_id = msg.parts.iter().any(|p| { - matches!(p, MessagePart::ToolCall { id, .. } if id.is_none()) - }); - let has_tool_calls = msg.parts.iter().any(|p| matches!(p, MessagePart::ToolCall { .. })); - - if let Some((extra, sse_finish)) = Self::extract_parts_from_sse_body(body) { - if !has_reasoning { - if let Some(r) = extra.iter().find(|p| matches!(p, MessagePart::Reasoning { .. })) { - msg.parts.insert(0, r.clone()); - } - } - // Always try to enrich tool_calls if missing id or no tool_calls - if !has_tool_calls || has_tool_calls_without_id { - // Remove existing tool_calls without id, replace with SSE ones - if has_tool_calls_without_id { - msg.parts.retain(|p| !matches!(p, MessagePart::ToolCall { id, .. } if id.is_none())); - } - for p in extra.into_iter().filter(|p| matches!(p, MessagePart::ToolCall { .. })) { - msg.parts.push(p); - } - } - // Enrich finish_reason if missing - if msg.finish_reason.is_none() { - msg.finish_reason = sse_finish; - } + if let Some(ref body) = http.response_body + && let Some(msg) = msgs.first_mut() + && msg.role == "assistant" + { + let has_reasoning = msg + .parts + .iter() + .any(|p| matches!(p, MessagePart::Reasoning { .. })); + // Check if any tool_call is missing id + let has_tool_calls_without_id = msg + .parts + .iter() + .any(|p| matches!(p, MessagePart::ToolCall { id, .. } if id.is_none())); + let has_tool_calls = msg + .parts + .iter() + .any(|p| matches!(p, MessagePart::ToolCall { .. })); + + if let Some((extra, sse_finish)) = Self::extract_parts_from_sse_body(body) { + if !has_reasoning + && let Some(r) = extra + .iter() + .find(|p| matches!(p, MessagePart::Reasoning { .. })) + { + msg.parts.insert(0, r.clone()); + } + // Always try to enrich tool_calls if missing id or no tool_calls + if !has_tool_calls || has_tool_calls_without_id { + // Remove existing tool_calls without id, replace with SSE ones + if has_tool_calls_without_id { + msg.parts.retain( + |p| !matches!(p, MessagePart::ToolCall { id, .. } if id.is_none()), + ); } + for p in extra + .into_iter() + .filter(|p| matches!(p, MessagePart::ToolCall { .. })) + { + msg.parts.push(p); + } + } + // Enrich finish_reason if missing + if msg.finish_reason.is_none() { + msg.finish_reason = sse_finish; } } } @@ -953,18 +1116,19 @@ impl GenAIBuilder { /// Check if the path indicates an LLM API call fn is_llm_api_path(&self, path: &str) -> bool { - path.contains("/v1/chat/completions") || - path.contains("/v1/completions") || - path.contains("/v1/messages") || - path.contains("/chat/completions") || - path.contains("/completions") || - path.contains("/api/v1/copilot/generate_copilot") + path.contains("/v1/chat/completions") + || path.contains("/v1/completions") + || path.contains("/v1/messages") + || path.contains("/chat/completions") + || path.contains("/completions") + || path.contains("/api/v1/copilot/generate_copilot") } /// Check if request body contains SysOM POP API markers /// SysOM uses path "/" with action in body (llmParamString field) fn is_sysom_pop_request(request_body: &Option) -> bool { - request_body.as_ref() + request_body + .as_ref() .map(|b| b.contains("llmParamString")) .unwrap_or(false) } @@ -1012,47 +1176,48 @@ impl GenAIBuilder { /// 优先从 request body 取（用户请求的 model）， /// 如果没有则从 response body 取（SSE 响应中的 model） /// 对于 SysOM 请求，需要从 llmParamString 内嵌 JSON 中提取 model - fn extract_model_from_body(request_body: &Option, response_body: &Option) -> Option { + fn extract_model_from_body( + request_body: &Option, + response_body: &Option, + ) -> Option { // 尝试从 request body 获取 - if let Some(body) = request_body { - if let Ok(v) = serde_json::from_str::(body) { - // 标准 OpenAI/Anthropic 格式 - if let Some(model) = v.get("model").and_then(|m| m.as_str()) { - if !model.is_empty() { - return Some(model.to_string()); - } - } - // SysOM 格式：model 嵌套在 llmParamString 中 - if let Some(lps) = v.get("llmParamString").and_then(|v| v.as_str()) { - if let Ok(inner) = serde_json::from_str::(lps) { - if let Some(model) = inner.get("model").and_then(|m| m.as_str()) { - if !model.is_empty() { - return Some(model.to_string()); - } - } - } - } + if let Some(body) = request_body + && let Ok(v) = serde_json::from_str::(body) + { + // 标准 OpenAI/Anthropic 格式 + if let Some(model) = v.get("model").and_then(|m| m.as_str()) + && !model.is_empty() + { + return Some(model.to_string()); + } + // SysOM 格式：model 嵌套在 llmParamString 中 + if let Some(lps) = v.get("llmParamString").and_then(|v| v.as_str()) + && let Ok(inner) = serde_json::from_str::(lps) + && let Some(model) = inner.get("model").and_then(|m| m.as_str()) + && !model.is_empty() + { + return Some(model.to_string()); } } // 尝试从 response body 获取（SSE 响应是 JSON 数组，取第一个 chunk） - if let Some(body) = response_body { - if let Ok(v) = serde_json::from_str::(body) { - // 非 SSE: 直接是 JSON 对象 - if let Some(model) = v.get("model").and_then(|m| m.as_str()) { - if !model.is_empty() { + if let Some(body) = response_body + && let Ok(v) = serde_json::from_str::(body) + { + // 非 SSE: 直接是 JSON 对象 + if let Some(model) = v.get("model").and_then(|m| m.as_str()) + && !model.is_empty() + { + return Some(model.to_string()); + } + // SSE: JSON 数组，取第一个 chunk 的 model + if let Some(arr) = v.as_array() { + for chunk in arr { + if let Some(model) = chunk.get("model").and_then(|m| m.as_str()) + && !model.is_empty() + { return Some(model.to_string()); } } - // SSE: JSON 数组，取第一个 chunk 的 model - if let Some(arr) = v.as_array() { - for chunk in arr { - if let Some(model) = chunk.get("model").and_then(|m| m.as_str()) { - if !model.is_empty() { - return Some(model.to_string()); - } - } - } - } } } None @@ -1064,16 +1229,25 @@ impl GenAIBuilder { /// 1. ParsedApiMessage response.id (OpenAI / Anthropic) /// 2. SSE response body first chunk "id" field /// 3. None (caller should fall back to call_id) - fn extract_response_id(parsed_message: &Option, http: &HttpRecord) -> Option { + fn extract_response_id( + parsed_message: &Option, + http: &HttpRecord, + ) -> Option { // 1. Try parsed message response.id if let Some(msg) = parsed_message { match msg { - ParsedApiMessage::OpenAICompletion { response: Some(resp), .. } => { + ParsedApiMessage::OpenAICompletion { + response: Some(resp), + .. + } => { if !resp.id.is_empty() { return Some(resp.id.clone()); } } - ParsedApiMessage::AnthropicMessage { response: Some(resp), .. } => { + ParsedApiMessage::AnthropicMessage { + response: Some(resp), + .. + } => { if !resp.id.is_empty() { return Some(resp.id.clone()); } @@ -1083,33 +1257,32 @@ impl GenAIBuilder { } // 2. SSE fallback: extract "id" field from response body - if http.is_sse { - if let Some(ref body) = http.response_body { - // Try JSON array format first (from HTTP/2 stream aggregation) - if let Ok(v) = serde_json::from_str::(body) { - if let Some(arr) = v.as_array() { - for chunk in arr { - if let Some(id) = chunk.get("id").and_then(|v| v.as_str()) { - if !id.is_empty() { - return Some(id.to_string()); - } - } - } + if http.is_sse + && let Some(ref body) = http.response_body + { + // Try JSON array format first (from HTTP/2 stream aggregation) + if let Ok(v) = serde_json::from_str::(body) + && let Some(arr) = v.as_array() + { + for chunk in arr { + if let Some(id) = chunk.get("id").and_then(|v| v.as_str()) + && !id.is_empty() + { + return Some(id.to_string()); } } - // Try SSE line format (from HTTP/1.1: "data: {...}" per line) - for line in body.lines() { - let json_str = line.strip_prefix("data: ").unwrap_or(line).trim(); - if json_str.is_empty() || json_str == "[DONE]" { - continue; - } - if let Ok(v) = serde_json::from_str::(json_str) { - if let Some(id) = v.get("id").and_then(|v| v.as_str()) { - if !id.is_empty() { - return Some(id.to_string()); - } - } - } + } + // Try SSE line format (from HTTP/1.1: "data: {...}" per line) + for line in body.lines() { + let json_str = line.strip_prefix("data: ").unwrap_or(line).trim(); + if json_str.is_empty() || json_str == "[DONE]" { + continue; + } + if let Ok(v) = serde_json::from_str::(json_str) + && let Some(id) = v.get("id").and_then(|v| v.as_str()) + && !id.is_empty() + { + return Some(id.to_string()); } } } @@ -1131,12 +1304,18 @@ impl GenAIBuilder { /// - 新会话：时间戳不同 → session_id 不同 fn compute_session_id(request: &LLMRequest) -> String { // 找第一条有实际文本的 user message（原始文本，含时间戳） - let first_user_raw: String = request.messages.iter() + let first_user_raw: String = request + .messages + .iter() .filter(|m| m.role == "user") .find_map(|m| { - let text: String = m.parts.iter() + let text: String = m + .parts + .iter() .filter_map(|p| match p { - MessagePart::Text { content } if !content.is_empty() => Some(content.as_str()), + MessagePart::Text { content } if !content.is_empty() => { + Some(content.as_str()) + } _ => None, }) .collect::>() @@ -1153,13 +1332,19 @@ impl GenAIBuilder { /// /// 跳过 Anthropic 格式中只包含 tool_result 的 user message fn extract_last_user_raw(request: &LLMRequest) -> Option { - request.messages.iter() + request + .messages + .iter() .rev() .filter(|m| m.role == "user") .find_map(|m| { - let text: String = m.parts.iter() + let text: String = m + .parts + .iter() .filter_map(|p| match p { - MessagePart::Text { content } if !content.is_empty() => Some(content.as_str()), + MessagePart::Text { content } if !content.is_empty() => { + Some(content.as_str()) + } _ => None, }) .collect::>() @@ -1170,8 +1355,7 @@ impl GenAIBuilder { /// 提取清理后的 user query（去除 metadata 前缀，用于展示） fn extract_last_user_query(request: &LLMRequest) -> Option { - Self::extract_last_user_raw(request) - .map(|raw| Self::strip_user_query_prefix(&raw)) + Self::extract_last_user_raw(request).map(|raw| Self::strip_user_query_prefix(&raw)) } /// 去除 user message 中的 metadata 前缀，只保留用户实际输入的文本 @@ -1193,7 +1377,9 @@ impl GenAIBuilder { if let Some(bracket_start) = text[..pos].rfind('[') { let bracket_content = &text[bracket_start + 1..pos]; // 简单验证：方括号内包含数字（日期）和冒号（时间） - if bracket_content.contains(':') && bracket_content.chars().any(|c| c.is_ascii_digit()) { + if bracket_content.contains(':') + && bracket_content.chars().any(|c| c.is_ascii_digit()) + { let after = text[pos + 1..].trim_start(); if !after.is_empty() { return after.to_string(); @@ -1203,7 +1389,7 @@ impl GenAIBuilder { } text.to_string() } - + /// 计算 user query 的 fingerprint，用于关联同一个请求的调用链 /// /// 使用原始文本（包含时间戳前缀）计算 hash， @@ -1220,7 +1406,11 @@ impl GenAIBuilder { /// Resolve agent name from comm string only (no /proc access). /// Used for dead-PID drain where the process is already gone. - fn resolve_agent_name_from_comm(comm: &str, pid: u32, cache: &std::collections::HashMap) -> Option { + fn resolve_agent_name_from_comm( + comm: &str, + pid: u32, + cache: &std::collections::HashMap, + ) -> Option { // First check the pid→agent_name cache (works even for dead processes) if let Some(name) = cache.get(&pid) { return Some(name.clone()); @@ -1232,13 +1422,17 @@ impl GenAIBuilder { }; default_cmdline_rules() .iter() - .filter_map(|rule| CmdlineGlobMatcher::from_config(rule)) + .filter_map(CmdlineGlobMatcher::from_config) .find(|m| m.matches(&ctx)) .map(|m| m.info().name.clone()) } /// 通过进程名匹配 agent registry，返回已知 agent 名称 - fn resolve_agent_name(comm: &str, pid: u32, cache: &std::collections::HashMap) -> Option { + fn resolve_agent_name( + comm: &str, + pid: u32, + cache: &std::collections::HashMap, + ) -> Option { // First check the pid→agent_name cache (works even for dead processes) if let Some(name) = cache.get(&pid) { return Some(name.clone()); @@ -1266,7 +1460,7 @@ impl GenAIBuilder { }; default_cmdline_rules() .iter() - .filter_map(|rule| CmdlineGlobMatcher::from_config(rule)) + .filter_map(CmdlineGlobMatcher::from_config) .find(|m| m.matches(&ctx)) .map(|m| m.info().name.clone()) } @@ -1277,20 +1471,24 @@ impl GenAIBuilder { let mut parts = Vec::new(); // Reasoning content first - if let Some(ref rc) = m.reasoning_content { - if !rc.is_empty() { - parts.push(MessagePart::Reasoning { content: rc.clone() }); - } + if let Some(ref rc) = m.reasoning_content + && !rc.is_empty() + { + parts.push(MessagePart::Reasoning { + content: rc.clone(), + }); } // For tool role: content is tool_call_response if role == "tool" { - let response_val = m.content.as_ref() + let response_val = m + .content + .as_ref() .map(|c| { let text = c.as_text(); // Try to parse as JSON, fall back to string serde_json::from_str::(&text) - .unwrap_or_else(|_| serde_json::Value::String(text)) + .unwrap_or(serde_json::Value::String(text)) }) .unwrap_or(serde_json::Value::Null); parts.push(MessagePart::ToolCallResponse { @@ -1316,7 +1514,11 @@ impl GenAIBuilder { } } - InputMessage { role, parts, name: m.name.clone() } + InputMessage { + role, + parts, + name: m.name.clone(), + } } /// Convert OpenAI ChatMessage to parts-based OutputMessage @@ -1325,10 +1527,12 @@ impl GenAIBuilder { let mut parts = Vec::new(); // Reasoning content first - if let Some(ref rc) = m.reasoning_content { - if !rc.is_empty() { - parts.push(MessagePart::Reasoning { content: rc.clone() }); - } + if let Some(ref rc) = m.reasoning_content + && !rc.is_empty() + { + parts.push(MessagePart::Reasoning { + content: rc.clone(), + }); } // Text content @@ -1362,13 +1566,15 @@ impl GenAIBuilder { let name = func.get("name")?.as_str()?.to_string(); let id = tc.get("id").and_then(|v| v.as_str()).map(|s| s.to_string()); // Parse arguments as JSON object (not string) - let arguments = func.get("arguments").and_then(|v| { - match v { - serde_json::Value::String(s) => serde_json::from_str(s).ok(), - other => Some(other.clone()), - } + let arguments = func.get("arguments").and_then(|v| match v { + serde_json::Value::String(s) => serde_json::from_str(s).ok(), + other => Some(other.clone()), }); - Some(MessagePart::ToolCall { id, name, arguments }) + Some(MessagePart::ToolCall { + id, + name, + arguments, + }) } // NOTE: token_record_to_parts and parse_tool_call_strings removed. @@ -1379,14 +1585,17 @@ impl GenAIBuilder { /// /// Merges content/reasoning deltas and tool_call argument fragments by index. /// Extracts finish_reason from the last SSE chunk that has one. - fn parse_sse_response_body(body: &str, fallback_finish_reason: Option<&str>) -> Option> { + fn parse_sse_response_body( + body: &str, + fallback_finish_reason: Option<&str>, + ) -> Option> { let (parts, sse_finish_reason) = Self::extract_parts_from_sse_body(body)?; if parts.is_empty() { return None; } // Prefer finish_reason from SSE, fall back to caller-supplied value - let finish_reason = sse_finish_reason - .or_else(|| fallback_finish_reason.map(|s| s.to_string())); + let finish_reason = + sse_finish_reason.or_else(|| fallback_finish_reason.map(|s| s.to_string())); Some(vec![OutputMessage { role: "assistant".to_string(), parts, @@ -1415,7 +1624,7 @@ impl GenAIBuilder { log::debug!("[GenAI] Parsing SSE body with {} chunks", chunks.len()); - for (_chunk_idx, chunk) in chunks.iter().enumerate() { + for chunk in chunks.iter() { let choices = chunk.get("choices").and_then(|c| c.as_array()); let choices = match choices { Some(c) => c, @@ -1438,14 +1647,15 @@ impl GenAIBuilder { if let Some(calls) = delta.get("tool_calls").and_then(|v| v.as_array()) { for tc in calls { let idx = tc.get("index").and_then(|v| v.as_u64()).unwrap_or(0) as u32; - let entry = tc_map.entry(idx) + let entry = tc_map + .entry(idx) .or_insert_with(|| (String::new(), String::new(), String::new())); - if let Some(id) = tc.get("id").and_then(|v| v.as_str()) { - if !id.is_empty() { - entry.0 = id.to_string(); - } - // 空字符串不覆盖已有的 id + if let Some(id) = tc.get("id").and_then(|v| v.as_str()) + && !id.is_empty() + { + entry.0 = id.to_string(); } + // 空字符串不覆盖已有的 id if let Some(func) = tc.get("function") { if let Some(name) = func.get("name").and_then(|v| v.as_str()) { entry.1 = name.to_string(); @@ -1467,11 +1677,15 @@ impl GenAIBuilder { // Reasoning first if !reasoning_buf.is_empty() { - parts.push(MessagePart::Reasoning { content: reasoning_buf }); + parts.push(MessagePart::Reasoning { + content: reasoning_buf, + }); } // Text content if !content_buf.is_empty() { - parts.push(MessagePart::Text { content: content_buf }); + parts.push(MessagePart::Text { + content: content_buf, + }); } // Merged tool calls if !tc_map.is_empty() { @@ -1493,7 +1707,11 @@ impl GenAIBuilder { } } - if parts.is_empty() { None } else { Some((parts, finish_reason)) } + if parts.is_empty() { + None + } else { + Some((parts, finish_reason)) + } } } @@ -1525,7 +1743,9 @@ mod tests { #[test] fn test_is_sysom_pop_request() { - assert!(GenAIBuilder::is_sysom_pop_request(&Some(r#"{"llmParamString":"{}"}"#.to_string()))); + assert!(GenAIBuilder::is_sysom_pop_request(&Some( + r#"{"llmParamString":"{}"}"#.to_string() + ))); assert!(!GenAIBuilder::is_sysom_pop_request(&Some("{}".to_string()))); assert!(!GenAIBuilder::is_sysom_pop_request(&None)); } @@ -1533,43 +1753,69 @@ mod tests { #[test] fn test_extract_provider_from_path() { let builder = GenAIBuilder::new(); - assert_eq!(builder.extract_provider_from_path("/v1/chat/completions"), Some("openai".to_string())); - assert_eq!(builder.extract_provider_from_path("/v1/messages"), Some("anthropic".to_string())); - assert_eq!(builder.extract_provider_from_path("/api/v1/copilot/generate_copilot"), Some("sysom".to_string())); + assert_eq!( + builder.extract_provider_from_path("/v1/chat/completions"), + Some("openai".to_string()) + ); + assert_eq!( + builder.extract_provider_from_path("/v1/messages"), + Some("anthropic".to_string()) + ); + assert_eq!( + builder.extract_provider_from_path("/api/v1/copilot/generate_copilot"), + Some("sysom".to_string()) + ); assert_eq!(builder.extract_provider_from_path("/unknown"), None); } #[test] fn test_extract_provider_from_body() { assert_eq!( - GenAIBuilder::extract_provider_from_body(&Some(r#"{"llmParamString":"{}"} "#.to_string())), + GenAIBuilder::extract_provider_from_body(&Some( + r#"{"llmParamString":"{}"} "#.to_string() + )), Some("sysom".to_string()) ); - assert_eq!(GenAIBuilder::extract_provider_from_body(&Some("{}".to_string())), None); + assert_eq!( + GenAIBuilder::extract_provider_from_body(&Some("{}".to_string())), + None + ); } #[test] fn test_extract_model_from_body_request() { let body = Some(r#"{"model": "gpt-4", "messages": []}"#.to_string()); - assert_eq!(GenAIBuilder::extract_model_from_body(&body, &None), Some("gpt-4".to_string())); + assert_eq!( + GenAIBuilder::extract_model_from_body(&body, &None), + Some("gpt-4".to_string()) + ); } #[test] fn test_extract_model_from_body_sysom() { let body = Some(r#"{"llmParamString": "{\"model\":\"qwen-max\"}"} "#.to_string()); - assert_eq!(GenAIBuilder::extract_model_from_body(&body, &None), Some("qwen-max".to_string())); + assert_eq!( + GenAIBuilder::extract_model_from_body(&body, &None), + Some("qwen-max".to_string()) + ); } #[test] fn test_extract_model_from_body_response() { let resp = Some(r#"{"model": "claude-3"}"#.to_string()); - assert_eq!(GenAIBuilder::extract_model_from_body(&None, &resp), Some("claude-3".to_string())); + assert_eq!( + GenAIBuilder::extract_model_from_body(&None, &resp), + Some("claude-3".to_string()) + ); } #[test] fn test_extract_model_from_body_sse_array() { let resp = Some(r#"[{"model": "gpt-4o"}, {"model": "gpt-4o"}]"#.to_string()); - assert_eq!(GenAIBuilder::extract_model_from_body(&None, &resp), Some("gpt-4o".to_string())); + assert_eq!( + GenAIBuilder::extract_model_from_body(&None, &resp), + Some("gpt-4o".to_string()) + ); } #[test] @@ -1586,14 +1832,20 @@ mod tests { #[test] fn test_strip_user_query_prefix_no_timestamp() { let text = "plain user input"; - assert_eq!(GenAIBuilder::strip_user_query_prefix(text), "plain user input"); + assert_eq!( + GenAIBuilder::strip_user_query_prefix(text), + "plain user input" + ); } #[test] fn test_strip_user_query_prefix_bracket_no_datetime() { let text = "[not a timestamp] content"; // No ':' and digit in bracket content -> returns original - assert_eq!(GenAIBuilder::strip_user_query_prefix(text), "[not a timestamp] content"); + assert_eq!( + GenAIBuilder::strip_user_query_prefix(text), + "[not a timestamp] content" + ); } #[test] @@ -1601,12 +1853,22 @@ mod tests { let req = LLMRequest { messages: vec![InputMessage { role: "user".to_string(), - parts: vec![MessagePart::Text { content: "hello".to_string() }], + parts: vec![MessagePart::Text { + content: "hello".to_string(), + }], name: None, }], - temperature: None, max_tokens: None, frequency_penalty: None, - presence_penalty: None, top_p: None, top_k: None, seed: None, - stop_sequences: None, stream: false, tools: None, raw_body: None, + temperature: None, + max_tokens: None, + frequency_penalty: None, + presence_penalty: None, + top_p: None, + top_k: None, + seed: None, + stop_sequences: None, + stream: false, + tools: None, + raw_body: None, }; let id1 = GenAIBuilder::compute_session_id(&req); let id2 = GenAIBuilder::compute_session_id(&req); @@ -1619,12 +1881,22 @@ mod tests { let req = LLMRequest { messages: vec![InputMessage { role: "system".to_string(), - parts: vec![MessagePart::Text { content: "sys".to_string() }], + parts: vec![MessagePart::Text { + content: "sys".to_string(), + }], name: None, }], - temperature: None, max_tokens: None, frequency_penalty: None, - presence_penalty: None, top_p: None, top_k: None, seed: None, - stop_sequences: None, stream: false, tools: None, raw_body: None, + temperature: None, + max_tokens: None, + frequency_penalty: None, + presence_penalty: None, + top_p: None, + top_k: None, + seed: None, + stop_sequences: None, + stream: false, + tools: None, + raw_body: None, }; let id = GenAIBuilder::compute_session_id(&req); assert_eq!(id.len(), 32); // still produces 32-char hash of empty string @@ -1636,18 +1908,30 @@ mod tests { messages: vec![ InputMessage { role: "user".to_string(), - parts: vec![MessagePart::Text { content: "first".to_string() }], + parts: vec![MessagePart::Text { + content: "first".to_string(), + }], name: None, }, InputMessage { role: "user".to_string(), - parts: vec![MessagePart::Text { content: "second".to_string() }], + parts: vec![MessagePart::Text { + content: "second".to_string(), + }], name: None, }, ], - temperature: None, max_tokens: None, frequency_penalty: None, - presence_penalty: None, top_p: None, top_k: None, seed: None, - stop_sequences: None, stream: false, tools: None, raw_body: None, + temperature: None, + max_tokens: None, + frequency_penalty: None, + presence_penalty: None, + top_p: None, + top_k: None, + seed: None, + stop_sequences: None, + stream: false, + tools: None, + raw_body: None, }; let fp = GenAIBuilder::compute_user_query_fingerprint(&req); assert_eq!(fp.len(), 32); @@ -1655,12 +1939,22 @@ mod tests { let req2 = LLMRequest { messages: vec![InputMessage { role: "user".to_string(), - parts: vec![MessagePart::Text { content: "second".to_string() }], + parts: vec![MessagePart::Text { + content: "second".to_string(), + }], name: None, }], - temperature: None, max_tokens: None, frequency_penalty: None, - presence_penalty: None, top_p: None, top_k: None, seed: None, - stop_sequences: None, stream: false, tools: None, raw_body: None, + temperature: None, + max_tokens: None, + frequency_penalty: None, + presence_penalty: None, + top_p: None, + top_k: None, + seed: None, + stop_sequences: None, + stream: false, + tools: None, + raw_body: None, }; assert_eq!(fp, GenAIBuilder::compute_user_query_fingerprint(&req2)); } @@ -1671,20 +1965,35 @@ mod tests { messages: vec![ InputMessage { role: "system".to_string(), - parts: vec![MessagePart::Text { content: "sys".to_string() }], + parts: vec![MessagePart::Text { + content: "sys".to_string(), + }], name: None, }, InputMessage { role: "user".to_string(), - parts: vec![MessagePart::Text { content: "[Mon 2026-01-01 10:00 GMT+8] hi".to_string() }], + parts: vec![MessagePart::Text { + content: "[Mon 2026-01-01 10:00 GMT+8] hi".to_string(), + }], name: None, }, ], - temperature: None, max_tokens: None, frequency_penalty: None, - presence_penalty: None, top_p: None, top_k: None, seed: None, - stop_sequences: None, stream: false, tools: None, raw_body: None, + temperature: None, + max_tokens: None, + frequency_penalty: None, + presence_penalty: None, + top_p: None, + top_k: None, + seed: None, + stop_sequences: None, + stream: false, + tools: None, + raw_body: None, }; - assert_eq!(GenAIBuilder::extract_last_user_query(&req), Some("hi".to_string())); + assert_eq!( + GenAIBuilder::extract_last_user_query(&req), + Some("hi".to_string()) + ); } #[test] @@ -1721,10 +2030,18 @@ mod tests { assert_eq!(req.messages.len(), 2); // assistant message has ToolCall part let parts = &req.messages[0].parts; - assert!(parts.iter().any(|p| matches!(p, MessagePart::ToolCall { name, .. } if name == "search"))); + assert!( + parts + .iter() + .any(|p| matches!(p, MessagePart::ToolCall { name, .. } if name == "search")) + ); // tool message has ToolCallResponse part let tool_parts = &req.messages[1].parts; - assert!(tool_parts.iter().any(|p| matches!(p, MessagePart::ToolCallResponse { .. }))); + assert!( + tool_parts + .iter() + .any(|p| matches!(p, MessagePart::ToolCallResponse { .. })) + ); } #[test] @@ -1749,7 +2066,11 @@ mod tests { }); let part = GenAIBuilder::parse_openai_tool_call_value(&tc).unwrap(); match part { - MessagePart::ToolCall { id, name, arguments } => { + MessagePart::ToolCall { + id, + name, + arguments, + } => { assert_eq!(id.unwrap(), "call_abc"); assert_eq!(name, "get_weather"); assert_eq!(arguments.unwrap()["city"], "Beijing"); @@ -1781,7 +2102,9 @@ mod tests { let body = r#"[{"choices":[{"delta":{"reasoning_content":"thinking..."}}]},{"choices":[{"delta":{"content":"answer"}}]}]"#; let (parts, _) = GenAIBuilder::extract_parts_from_sse_body(body).unwrap(); assert_eq!(parts.len(), 2); - assert!(matches!(&parts[0], MessagePart::Reasoning { content } if content == "thinking...")); + assert!( + matches!(&parts[0], MessagePart::Reasoning { content } if content == "thinking...") + ); assert!(matches!(&parts[1], MessagePart::Text { content } if content == "answer")); } @@ -1795,7 +2118,11 @@ mod tests { let (parts, finish) = GenAIBuilder::extract_parts_from_sse_body(body).unwrap(); assert_eq!(parts.len(), 1); match &parts[0] { - MessagePart::ToolCall { id, name, arguments } => { + MessagePart::ToolCall { + id, + name, + arguments, + } => { assert_eq!(id.as_deref(), Some("tc_1")); assert_eq!(name, "search"); assert_eq!(arguments.as_ref().unwrap()["q"], "rust"); @@ -1828,7 +2155,9 @@ mod tests { fn test_openai_msg_to_input_basic() { let msg = OpenAIChatMessage { role: crate::analyzer::message::types::MessageRole::User, - content: Some(crate::analyzer::message::types::OpenAIContent::Text("Hello".to_string())), + content: Some(crate::analyzer::message::types::OpenAIContent::Text( + "Hello".to_string(), + )), reasoning_content: None, refusal: None, function_call: None, @@ -1851,7 +2180,9 @@ mod tests { fn test_openai_msg_to_input_tool_role() { let msg = OpenAIChatMessage { role: crate::analyzer::message::types::MessageRole::Tool, - content: Some(crate::analyzer::message::types::OpenAIContent::Text("result data".to_string())), + content: Some(crate::analyzer::message::types::OpenAIContent::Text( + "result data".to_string(), + )), reasoning_content: None, refusal: None, function_call: None, @@ -1863,14 +2194,18 @@ mod tests { }; let input = GenAIBuilder::openai_msg_to_input(&msg); assert_eq!(input.role, "tool"); - assert!(matches!(&input.parts[0], MessagePart::ToolCallResponse { id, .. } if id.as_deref() == Some("tc_1"))); + assert!( + matches!(&input.parts[0], MessagePart::ToolCallResponse { id, .. } if id.as_deref() == Some("tc_1")) + ); } #[test] fn test_openai_msg_to_output() { let msg = OpenAIChatMessage { role: crate::analyzer::message::types::MessageRole::Assistant, - content: Some(crate::analyzer::message::types::OpenAIContent::Text("Response".to_string())), + content: Some(crate::analyzer::message::types::OpenAIContent::Text( + "Response".to_string(), + )), reasoning_content: Some("thinking".to_string()), refusal: None, function_call: None, @@ -1885,7 +2220,9 @@ mod tests { assert_eq!(output.finish_reason.as_deref(), Some("stop")); // Should have reasoning + text = 2 parts assert_eq!(output.parts.len(), 2); - assert!(matches!(&output.parts[0], MessagePart::Reasoning { content } if content == "thinking")); + assert!( + matches!(&output.parts[0], MessagePart::Reasoning { content } if content == "thinking") + ); assert!(matches!(&output.parts[1], MessagePart::Text { content } if content == "Response")); } @@ -1911,16 +2248,29 @@ mod tests { request: Some(crate::analyzer::message::types::OpenAIRequest { model: "gpt-4-turbo".to_string(), messages: vec![], - temperature: None, max_tokens: None, stream: None, - top_p: None, n: None, stop: None, - presence_penalty: None, frequency_penalty: None, - user: None, tools: None, tool_choice: None, - response_format: None, seed: None, logprobs: None, - top_logprobs: None, parallel_tool_calls: None, + temperature: None, + max_tokens: None, + stream: None, + top_p: None, + n: None, + stop: None, + presence_penalty: None, + frequency_penalty: None, + user: None, + tools: None, + tool_choice: None, + response_format: None, + seed: None, + logprobs: None, + top_logprobs: None, + parallel_tool_calls: None, }), response: None, }); - assert_eq!(builder.extract_model_from_message(&msg), Some("gpt-4-turbo".to_string())); + assert_eq!( + builder.extract_model_from_message(&msg), + Some("gpt-4-turbo".to_string()) + ); assert_eq!(builder.extract_model_from_message(&None), None); } @@ -1936,4 +2286,3 @@ mod tests { assert!(id2.contains('_')); } } - diff --git a/src/agentsight/src/genai/encrypt.rs b/src/agentsight/src/genai/encrypt.rs index 1a706083b..cdb55357d 100644 --- a/src/agentsight/src/genai/encrypt.rs +++ b/src/agentsight/src/genai/encrypt.rs @@ -6,12 +6,12 @@ //! //! 公钥来源：由调用方从 agentsight.json 的 `encryption.public_key` 读取后传入。 -use openssl::rsa::{Rsa, Padding}; -use openssl::pkey::Public; -use openssl::symm::{Cipher, encrypt_aead}; -use openssl::rand::rand_bytes; use base64::Engine; use base64::engine::general_purpose::STANDARD as BASE64; +use openssl::pkey::Public; +use openssl::rand::rand_bytes; +use openssl::rsa::{Padding, Rsa}; +use openssl::symm::{Cipher, encrypt_aead}; /// AES-256 密钥长度（32 字节） const AES_KEY_LEN: usize = 32; @@ -35,7 +35,10 @@ impl MessageEncryptor { pub fn from_pem(pem: &str) -> Option { match Rsa::public_key_from_pem(pem.as_bytes()) { Ok(rsa) => { - log::info!("MessageEncryptor initialized (RSA-{} + AES-256-GCM)", rsa.size() * 8); + log::info!( + "MessageEncryptor initialized (RSA-{} + AES-256-GCM)", + rsa.size() * 8 + ); Some(MessageEncryptor { rsa }) } Err(e) => { @@ -64,25 +67,24 @@ impl MessageEncryptor { Cipher::aes_256_gcm(), &aes_key, Some(&nonce), - &[], // AAD (Additional Authenticated Data) - 不使用 + &[], // AAD (Additional Authenticated Data) - 不使用 plaintext.as_bytes(), &mut tag, - ).map_err(|e| format!("AES-256-GCM encryption failed: {}", e))?; + ) + .map_err(|e| format!("AES-256-GCM encryption failed: {}", e))?; // 4. RSA-OAEP(SHA-256) 加密 AES 密钥 let mut encrypted_key = vec![0u8; self.rsa.size() as usize]; - let encrypted_key_len = self.rsa.public_encrypt( - &aes_key, - &mut encrypted_key, - Padding::PKCS1_OAEP, - ).map_err(|e| format!("RSA-OAEP encryption failed: {}", e))?; + let encrypted_key_len = self + .rsa + .public_encrypt(&aes_key, &mut encrypted_key, Padding::PKCS1_OAEP) + .map_err(|e| format!("RSA-OAEP encryption failed: {}", e))?; encrypted_key.truncate(encrypted_key_len); // 5. 组装二进制输出：[2字节长度] [encrypted_key] [nonce] [ciphertext] [tag] let key_len_bytes = (encrypted_key_len as u16).to_be_bytes(); - let mut output = Vec::with_capacity( - 2 + encrypted_key_len + NONCE_LEN + ciphertext.len() + TAG_LEN - ); + let mut output = + Vec::with_capacity(2 + encrypted_key_len + NONCE_LEN + ciphertext.len() + TAG_LEN); output.extend_from_slice(&key_len_bytes); output.extend_from_slice(&encrypted_key); output.extend_from_slice(&nonce); diff --git a/src/agentsight/src/genai/instance_id.rs b/src/agentsight/src/genai/instance_id.rs index 6f87a1837..2de6cab34 100644 --- a/src/agentsight/src/genai/instance_id.rs +++ b/src/agentsight/src/genai/instance_id.rs @@ -28,15 +28,16 @@ fn metadata_agent() -> ureq::Agent { /// 获取 owner account ID（带缓存）：首次调用请求阿里云 ECS metadata（超时1秒）， /// 失败返回空字符串。后续调用直接返回缓存值。 pub fn get_owner_account_id() -> &'static str { - OWNER_ACCOUNT_ID.get_or_init(|| { - fetch_owner_account_id() - }) + OWNER_ACCOUNT_ID.get_or_init(fetch_owner_account_id) } /// 实际请求 owner-account-id fn fetch_owner_account_id() -> String { let agent = metadata_agent(); - match agent.get("http://100.100.100.200/latest/meta-data/owner-account-id").call() { + match agent + .get("http://100.100.100.200/latest/meta-data/owner-account-id") + .call() + { Ok(resp) => { if let Ok(body) = resp.into_string() { let uid = body.trim().to_string(); @@ -56,16 +57,17 @@ fn fetch_owner_account_id() -> String { /// 获取实例ID（带缓存）：首次调用请求阿里云 ECS metadata（超时1秒）， /// 失败则回退到 hostname。后续调用直接返回缓存值。 pub fn get_instance_id() -> &'static str { - INSTANCE_ID.get_or_init(|| { - fetch_instance_id() - }) + INSTANCE_ID.get_or_init(fetch_instance_id) } /// 实际请求 instance-id fn fetch_instance_id() -> String { // 尝试从 ECS metadata service 获取 instance-id let agent = metadata_agent(); - match agent.get("http://100.100.100.200/latest/meta-data/instance-id").call() { + match agent + .get("http://100.100.100.200/latest/meta-data/instance-id") + .call() + { Ok(resp) => { if let Ok(body) = resp.into_string() { let id = body.trim().to_string(); @@ -76,7 +78,10 @@ fn fetch_instance_id() -> String { } } Err(e) => { - log::debug!("ECS metadata not available, falling back to hostname: {}", e); + log::debug!( + "ECS metadata not available, falling back to hostname: {}", + e + ); } } // 回退: /etc/hostname -> $HOSTNAME -> "unknown" diff --git a/src/agentsight/src/genai/logtail.rs b/src/agentsight/src/genai/logtail.rs index 71daf0aa4..6b729ce21 100644 --- a/src/agentsight/src/genai/logtail.rs +++ b/src/agentsight/src/genai/logtail.rs @@ -7,14 +7,14 @@ //! 仅当该环境变量设置时才启用。 use std::collections::BTreeMap; -use std::path::PathBuf; use std::fs::OpenOptions; -use std::io::{Write, BufWriter}; +use std::io::{BufWriter, Write}; +use std::path::PathBuf; -use super::semantic::GenAISemanticEvent; +use super::encrypt::MessageEncryptor; use super::exporter::GenAIExporter; use super::instance_id; -use super::encrypt::MessageEncryptor; +use super::semantic::GenAISemanticEvent; /// 环境变量名称 pub const LOGTAIL_ENV_VAR: &str = "SLS_LOGTAIL_FILE"; @@ -67,9 +67,15 @@ impl LogtailExporter { log::info!("Logtail exporter: encryption disabled (no public key configured)"); } if !trace_enabled { - log::info!("Logtail exporter: traceEnabled=false, conversation content fields (gen_ai.input.messages, gen_ai.output.messages) will NOT be uploaded"); + log::info!( + "Logtail exporter: traceEnabled=false, conversation content fields (gen_ai.input.messages, gen_ai.output.messages) will NOT be uploaded" + ); } - Some(LogtailExporter { path, encryptor, trace_enabled }) + Some(LogtailExporter { + path, + encryptor, + trace_enabled, + }) } /// 返回导出文件路径 @@ -138,7 +144,11 @@ impl GenAIExporter for LogtailExporter { /// /// `trace_enabled=false` 时跳过 LLMCall 中的对话内容字段 /// (`gen_ai.input.messages` 与 `gen_ai.output.messages`)，token 数量等元数据仍上传。 -pub fn events_to_flat_records(events: &[GenAISemanticEvent], encryptor: Option<&MessageEncryptor>, trace_enabled: bool) -> Vec> { +pub fn events_to_flat_records( + events: &[GenAISemanticEvent], + encryptor: Option<&MessageEncryptor>, + trace_enabled: bool, +) -> Vec> { let hostname = instance_id::get_instance_id(); let uid = instance_id::get_owner_account_id(); let mut records = Vec::with_capacity(events.len()); @@ -165,8 +175,13 @@ pub fn events_to_flat_records(events: &[GenAISemanticEvent], encryptor: Option<& // ── OTel GenAI Required ── m.insert("gen_ai.provider.name".to_string(), call.provider.clone()); m.insert("gen_ai.request.model".to_string(), call.model.clone()); - m.insert("gen_ai.operation.name".to_string(), - call.metadata.get("operation_name").cloned().unwrap_or_else(|| "chat".to_string())); + m.insert( + "gen_ai.operation.name".to_string(), + call.metadata + .get("operation_name") + .cloned() + .unwrap_or_else(|| "chat".to_string()), + ); // ── OTel GenAI Conditionally Required ── if let Some(ref error) = call.error { @@ -183,8 +198,16 @@ pub fn events_to_flat_records(events: &[GenAISemanticEvent], encryptor: Option<& m.insert("gen_ai.response.id".to_string(), call.call_id.clone()); } m.insert("gen_ai.response.model".to_string(), call.model.clone()); - if let Some(reason) = call.response.messages.first().and_then(|msg| msg.finish_reason.as_ref()) { - m.insert("gen_ai.response.finish_reasons".to_string(), format!("[\"{}\"]", reason)); + if let Some(reason) = call + .response + .messages + .first() + .and_then(|msg| msg.finish_reason.as_ref()) + { + m.insert( + "gen_ai.response.finish_reasons".to_string(), + format!("[\"{}\"]", reason), + ); } if let Some(temp) = call.request.temperature { m.insert("gen_ai.request.temperature".to_string(), temp.to_string()); @@ -193,10 +216,16 @@ pub fn events_to_flat_records(events: &[GenAISemanticEvent], encryptor: Option<& m.insert("gen_ai.request.max_tokens".to_string(), max.to_string()); } if let Some(fp) = call.request.frequency_penalty { - m.insert("gen_ai.request.frequency_penalty".to_string(), fp.to_string()); + m.insert( + "gen_ai.request.frequency_penalty".to_string(), + fp.to_string(), + ); } if let Some(pp) = call.request.presence_penalty { - m.insert("gen_ai.request.presence_penalty".to_string(), pp.to_string()); + m.insert( + "gen_ai.request.presence_penalty".to_string(), + pp.to_string(), + ); } if let Some(tp) = call.request.top_p { m.insert("gen_ai.request.top_p".to_string(), tp.to_string()); @@ -208,13 +237,25 @@ pub fn events_to_flat_records(events: &[GenAISemanticEvent], encryptor: Option<& m.insert("gen_ai.request.seed".to_string(), seed.to_string()); } if let Some(ref usage) = call.token_usage { - m.insert("gen_ai.usage.input_tokens".to_string(), usage.input_tokens.to_string()); - m.insert("gen_ai.usage.output_tokens".to_string(), usage.output_tokens.to_string()); + m.insert( + "gen_ai.usage.input_tokens".to_string(), + usage.input_tokens.to_string(), + ); + m.insert( + "gen_ai.usage.output_tokens".to_string(), + usage.output_tokens.to_string(), + ); if let Some(cache_create) = usage.cache_creation_input_tokens { - m.insert("gen_ai.usage.cache_creation.input_tokens".to_string(), cache_create.to_string()); + m.insert( + "gen_ai.usage.cache_creation.input_tokens".to_string(), + cache_create.to_string(), + ); } if let Some(cache_read) = usage.cache_read_input_tokens { - m.insert("gen_ai.usage.cache_read.input_tokens".to_string(), cache_read.to_string()); + m.insert( + "gen_ai.usage.cache_read.input_tokens".to_string(), + cache_read.to_string(), + ); } } if let Some(addr) = call.metadata.get("server.address") { @@ -223,14 +264,19 @@ pub fn events_to_flat_records(events: &[GenAISemanticEvent], encryptor: Option<& m.insert("gen_ai.output.type".to_string(), "text".to_string()); // ── gen_ai.system_instructions (system role messages) ── - let system_msgs: Vec<&super::semantic::InputMessage> = call.request.messages.iter() + let system_msgs: Vec<&super::semantic::InputMessage> = call + .request + .messages + .iter() .filter(|msg| msg.role == "system") .collect(); - if !system_msgs.is_empty() { - if let Ok(json) = serde_json::to_string(&system_msgs) { - m.insert("gen_ai.system_instructions".to_string(), - MessageEncryptor::maybe_encrypt(encryptor, &json)); - } + if !system_msgs.is_empty() + && let Ok(json) = serde_json::to_string(&system_msgs) + { + m.insert( + "gen_ai.system_instructions".to_string(), + MessageEncryptor::maybe_encrypt(encryptor, &json), + ); } // ── gen_ai.input.messages (增量：只取最新一轮) ── @@ -240,21 +286,26 @@ pub fn events_to_flat_records(events: &[GenAISemanticEvent], encryptor: Option<& if trace_enabled { let latest_msgs = super::semantic::latest_round_input_messages(&call.request.messages); - if !latest_msgs.is_empty() { - if let Ok(json) = serde_json::to_string(&latest_msgs) { - m.insert("gen_ai.input.messages".to_string(), - MessageEncryptor::maybe_encrypt(encryptor, &json)); - } + if !latest_msgs.is_empty() + && let Ok(json) = serde_json::to_string(&latest_msgs) + { + m.insert( + "gen_ai.input.messages".to_string(), + MessageEncryptor::maybe_encrypt(encryptor, &json), + ); } } // ── gen_ai.output.messages (parts-based with finish_reason) ── // 同样受 trace_enabled 控制，不上传模型响应内容。 - if trace_enabled && !call.response.messages.is_empty() { - if let Ok(json) = serde_json::to_string(&call.response.messages) { - m.insert("gen_ai.output.messages".to_string(), - MessageEncryptor::maybe_encrypt(encryptor, &json)); - } + if trace_enabled + && !call.response.messages.is_empty() + && let Ok(json) = serde_json::to_string(&call.response.messages) + { + m.insert( + "gen_ai.output.messages".to_string(), + MessageEncryptor::maybe_encrypt(encryptor, &json), + ); } // ── 加密标记字段 ── @@ -264,13 +315,25 @@ pub fn events_to_flat_records(events: &[GenAISemanticEvent], encryptor: Option<& // ── AgentSight extensions ── m.insert("agentsight.pid".to_string(), call.pid.to_string()); - m.insert("agentsight.process_name".to_string(), call.process_name.clone()); + m.insert( + "agentsight.process_name".to_string(), + call.process_name.clone(), + ); if let Some(ref name) = call.agent_name { m.insert("agentsight.agent.name".to_string(), name.clone()); } - m.insert("agentsight.duration_ns".to_string(), call.duration_ns.to_string()); - m.insert("agentsight.start_timestamp_ns".to_string(), call.start_timestamp_ns.to_string()); - m.insert("agentsight.end_timestamp_ns".to_string(), call.end_timestamp_ns.to_string()); + m.insert( + "agentsight.duration_ns".to_string(), + call.duration_ns.to_string(), + ); + m.insert( + "agentsight.start_timestamp_ns".to_string(), + call.start_timestamp_ns.to_string(), + ); + m.insert( + "agentsight.end_timestamp_ns".to_string(), + call.end_timestamp_ns.to_string(), + ); if let Some(method) = call.metadata.get("method") { m.insert("agentsight.http.method".to_string(), method.clone()); } @@ -280,7 +343,13 @@ pub fn events_to_flat_records(events: &[GenAISemanticEvent], encryptor: Option<& if let Some(status) = call.metadata.get("status_code") { m.insert("agentsight.http.status_code".to_string(), status.clone()); } - if call.request.stream || call.metadata.get("is_sse").map(|v| v == "true").unwrap_or(false) { + if call.request.stream + || call + .metadata + .get("is_sse") + .map(|v| v == "true") + .unwrap_or(false) + { m.insert("agentsight.stream".to_string(), "true".to_string()); if let Some(cnt) = call.metadata.get("sse_event_count") { m.insert("agentsight.sse_event_count".to_string(), cnt.clone()); @@ -304,7 +373,10 @@ pub fn events_to_flat_records(events: &[GenAISemanticEvent], encryptor: Option<& if let Some(ref parent_id) = tool.parent_llm_call_id { m.insert("gen_ai.response.id".to_string(), parent_id.clone()); } - m.insert("agentsight.tool.success".to_string(), tool.success.to_string()); + m.insert( + "agentsight.tool.success".to_string(), + tool.success.to_string(), + ); m.insert("agentsight.pid".to_string(), tool.pid.to_string()); if let Some(ref dur) = tool.duration_ns { m.insert("agentsight.duration_ns".to_string(), dur.to_string()); @@ -314,15 +386,30 @@ pub fn events_to_flat_records(events: &[GenAISemanticEvent], encryptor: Option<& } } GenAISemanticEvent::AgentInteraction(interaction) => { - m.insert("gen_ai.operation.name".to_string(), "agent_interaction".to_string()); - m.insert("agentsight.agent.name".to_string(), interaction.agent_name.clone()); - m.insert("agentsight.agent.interaction_type".to_string(), interaction.interaction_type.clone()); + m.insert( + "gen_ai.operation.name".to_string(), + "agent_interaction".to_string(), + ); + m.insert( + "agentsight.agent.name".to_string(), + interaction.agent_name.clone(), + ); + m.insert( + "agentsight.agent.interaction_type".to_string(), + interaction.interaction_type.clone(), + ); m.insert("agentsight.pid".to_string(), interaction.pid.to_string()); } GenAISemanticEvent::StreamChunk(chunk) => { - m.insert("gen_ai.operation.name".to_string(), "stream_chunk".to_string()); + m.insert( + "gen_ai.operation.name".to_string(), + "stream_chunk".to_string(), + ); m.insert("agentsight.stream.id".to_string(), chunk.stream_id.clone()); - m.insert("agentsight.stream.chunk_index".to_string(), chunk.chunk_index.to_string()); + m.insert( + "agentsight.stream.chunk_index".to_string(), + chunk.chunk_index.to_string(), + ); m.insert("agentsight.pid".to_string(), chunk.pid.to_string()); } } @@ -347,12 +434,16 @@ mod tests { messages: vec![ InputMessage { role: "system".to_string(), - parts: vec![MessagePart::Text { content: "you are helpful".to_string() }], + parts: vec![MessagePart::Text { + content: "you are helpful".to_string(), + }], name: None, }, InputMessage { role: "user".to_string(), - parts: vec![MessagePart::Text { content: "hello secret".to_string() }], + parts: vec![MessagePart::Text { + content: "hello secret".to_string(), + }], name: None, }, ], @@ -381,7 +472,9 @@ mod tests { LLMResponse { messages: vec![OutputMessage { role: "assistant".to_string(), - parts: vec![MessagePart::Text { content: "sensitive reply".to_string() }], + parts: vec![MessagePart::Text { + content: "sensitive reply".to_string(), + }], name: None, finish_reason: Some("stop".to_string()), }], @@ -408,11 +501,23 @@ mod tests { let records = events_to_flat_records(&[event], None, true); assert_eq!(records.len(), 1); let r = &records[0]; - assert!(r.contains_key("gen_ai.input.messages"), "input.messages should be uploaded when traceEnabled=true"); - assert!(r.contains_key("gen_ai.output.messages"), "output.messages should be uploaded when traceEnabled=true"); + assert!( + r.contains_key("gen_ai.input.messages"), + "input.messages should be uploaded when traceEnabled=true" + ); + assert!( + r.contains_key("gen_ai.output.messages"), + "output.messages should be uploaded when traceEnabled=true" + ); // token 数量元数据也应存在 - assert_eq!(r.get("gen_ai.usage.input_tokens").map(String::as_str), Some("100")); - assert_eq!(r.get("gen_ai.usage.output_tokens").map(String::as_str), Some("50")); + assert_eq!( + r.get("gen_ai.usage.input_tokens").map(String::as_str), + Some("100") + ); + assert_eq!( + r.get("gen_ai.usage.output_tokens").map(String::as_str), + Some("50") + ); } #[test] @@ -422,16 +527,37 @@ mod tests { let records = events_to_flat_records(&[event], None, false); assert_eq!(records.len(), 1); let r = &records[0]; - assert!(!r.contains_key("gen_ai.input.messages"), "input.messages must NOT be uploaded when traceEnabled=false"); - assert!(!r.contains_key("gen_ai.output.messages"), "output.messages must NOT be uploaded when traceEnabled=false"); + assert!( + !r.contains_key("gen_ai.input.messages"), + "input.messages must NOT be uploaded when traceEnabled=false" + ); + assert!( + !r.contains_key("gen_ai.output.messages"), + "output.messages must NOT be uploaded when traceEnabled=false" + ); // token 消耗与模型元数据仍需上传 - assert_eq!(r.get("gen_ai.usage.input_tokens").map(String::as_str), Some("100")); - assert_eq!(r.get("gen_ai.usage.output_tokens").map(String::as_str), Some("50")); - assert_eq!(r.get("gen_ai.provider.name").map(String::as_str), Some("openai")); - assert_eq!(r.get("gen_ai.request.model").map(String::as_str), Some("gpt-4")); + assert_eq!( + r.get("gen_ai.usage.input_tokens").map(String::as_str), + Some("100") + ); + assert_eq!( + r.get("gen_ai.usage.output_tokens").map(String::as_str), + Some("50") + ); + assert_eq!( + r.get("gen_ai.provider.name").map(String::as_str), + Some("openai") + ); + assert_eq!( + r.get("gen_ai.request.model").map(String::as_str), + Some("gpt-4") + ); assert_eq!(r.get("agentsight.pid").map(String::as_str), Some("42")); - assert_eq!(r.get("agentsight.duration_ns").map(String::as_str), Some("4000")); + assert_eq!( + r.get("agentsight.duration_ns").map(String::as_str), + Some("4000") + ); // 所有名为 gen_ai.*.messages 的字段都应被过滤 for key in r.keys() { assert!( @@ -463,7 +589,10 @@ mod tests { let records = events_to_flat_records(&[event], None, false); assert_eq!(records.len(), 1); let r = &records[0]; - assert_eq!(r.get("gen_ai.operation.name").map(String::as_str), Some("tool_use")); + assert_eq!( + r.get("gen_ai.operation.name").map(String::as_str), + Some("tool_use") + ); assert_eq!(r.get("gen_ai.tool.name").map(String::as_str), Some("shell")); assert_eq!(r.get("agentsight.pid").map(String::as_str), Some("7")); } diff --git a/src/agentsight/src/genai/mod.rs b/src/agentsight/src/genai/mod.rs index 6efbb01aa..9558f7b59 100644 --- a/src/agentsight/src/genai/mod.rs +++ b/src/agentsight/src/genai/mod.rs @@ -3,24 +3,22 @@ //! This module provides GenAI-specific semantic conversion and storage //! for LLM API calls, tool uses, and agent interactions. -pub mod semantic; pub mod builder; +pub mod encrypt; pub mod exporter; -pub mod storage; pub mod instance_id; pub mod logtail; -pub mod encrypt; +pub mod semantic; +pub mod storage; +pub use builder::GenAIBuilder; +pub use exporter::GenAIExporter; +pub use logtail::LogtailExporter; pub use semantic::{ - GenAISemanticEvent, LLMCall, LLMRequest, LLMResponse, - MessagePart, InputMessage, OutputMessage, - TokenUsage, ToolUse, AgentInteraction, StreamChunk, - ToolDefinition, + AgentInteraction, GenAISemanticEvent, InputMessage, LLMCall, LLMRequest, LLMResponse, + MessagePart, OutputMessage, StreamChunk, TokenUsage, ToolDefinition, ToolUse, }; -pub use exporter::GenAIExporter; -pub use builder::GenAIBuilder; pub use storage::{GenAIStore, GenAIStoreStats}; -pub use logtail::LogtailExporter; // Blanket implementation: Arc implements GenAIExporter if T does. // This allows storing an Arc both in genai_exporters and diff --git a/src/agentsight/src/genai/semantic.rs b/src/agentsight/src/genai/semantic.rs index bc97b9c5e..41a8e45f5 100644 --- a/src/agentsight/src/genai/semantic.rs +++ b/src/agentsight/src/genai/semantic.rs @@ -3,11 +3,12 @@ //! This module defines GenAI-specific semantic structures that represent //! LLM interactions at a higher abstraction level than raw HTTP requests/responses. -use serde::{Serialize, Deserialize}; +use serde::{Deserialize, Serialize}; use std::collections::HashMap; /// GenAI semantic event types #[derive(Debug, Clone, Serialize, Deserialize)] +#[allow(clippy::large_enum_variant)] pub enum GenAISemanticEvent { /// LLM API call with request/response LLMCall(LLMCall), @@ -151,13 +152,12 @@ pub fn latest_round_input_messages(messages: &[InputMessage]) -> Vec<&InputMessa // A user message that carries actual text (not just a tool_result). fn is_text_user(m: &InputMessage) -> bool { m.role == "user" - && m.parts.iter().any(|p| { - matches!(p, MessagePart::Text { content } if !content.is_empty()) - }) + && m.parts + .iter() + .any(|p| matches!(p, MessagePart::Text { content } if !content.is_empty())) } - let non_system: Vec<&InputMessage> = - messages.iter().filter(|m| m.role != "system").collect(); + let non_system: Vec<&InputMessage> = messages.iter().filter(|m| m.role != "system").collect(); let last = non_system.iter().rposition(|m| is_text_user(m)); let Some(mut idx) = last else { @@ -336,7 +336,9 @@ mod tests { LLMRequest { messages: vec![InputMessage { role: "user".to_string(), - parts: vec![MessagePart::Text { content: "Hello".to_string() }], + parts: vec![MessagePart::Text { + content: "Hello".to_string(), + }], name: None, }], temperature: Some(0.7), @@ -357,8 +359,13 @@ mod tests { fn test_llm_call_new() { let req = make_request(); let call = LLMCall::new( - "call-1".to_string(), 1000, "openai".to_string(), - "gpt-4".to_string(), req, 100, "test".to_string(), + "call-1".to_string(), + 1000, + "openai".to_string(), + "gpt-4".to_string(), + req, + 100, + "test".to_string(), ); assert_eq!(call.call_id, "call-1"); assert_eq!(call.end_timestamp_ns, 0); @@ -373,13 +380,20 @@ mod tests { fn test_llm_call_set_response() { let req = make_request(); let mut call = LLMCall::new( - "call-2".to_string(), 1000, "anthropic".to_string(), - "claude-3".to_string(), req, 200, "agent".to_string(), + "call-2".to_string(), + 1000, + "anthropic".to_string(), + "claude-3".to_string(), + req, + 200, + "agent".to_string(), ); let resp = LLMResponse { messages: vec![OutputMessage { role: "assistant".to_string(), - parts: vec![MessagePart::Text { content: "Hi".to_string() }], + parts: vec![MessagePart::Text { + content: "Hi".to_string(), + }], name: None, finish_reason: Some("stop".to_string()), }], @@ -396,24 +410,40 @@ mod tests { fn test_llm_call_set_token_usage() { let req = make_request(); let mut call = LLMCall::new( - "call-3".to_string(), 0, "openai".to_string(), - "gpt-4".to_string(), req, 1, "p".to_string(), + "call-3".to_string(), + 0, + "openai".to_string(), + "gpt-4".to_string(), + req, + 1, + "p".to_string(), ); let usage = TokenUsage { - input_tokens: 100, output_tokens: 50, total_tokens: 150, - cache_creation_input_tokens: None, cache_read_input_tokens: Some(10), + input_tokens: 100, + output_tokens: 50, + total_tokens: 150, + cache_creation_input_tokens: None, + cache_read_input_tokens: Some(10), }; call.set_token_usage(usage); assert_eq!(call.token_usage.as_ref().unwrap().total_tokens, 150); - assert_eq!(call.token_usage.as_ref().unwrap().cache_read_input_tokens, Some(10)); + assert_eq!( + call.token_usage.as_ref().unwrap().cache_read_input_tokens, + Some(10) + ); } #[test] fn test_llm_call_set_error() { let req = make_request(); let mut call = LLMCall::new( - "call-4".to_string(), 0, "openai".to_string(), - "gpt-4".to_string(), req, 1, "p".to_string(), + "call-4".to_string(), + 0, + "openai".to_string(), + "gpt-4".to_string(), + req, + 1, + "p".to_string(), ); call.set_error("timeout".to_string()); assert_eq!(call.error.as_ref().unwrap(), "timeout"); @@ -421,7 +451,9 @@ mod tests { #[test] fn test_message_part_serde_text() { - let part = MessagePart::Text { content: "hello world".to_string() }; + let part = MessagePart::Text { + content: "hello world".to_string(), + }; let json = serde_json::to_string(&part).unwrap(); assert!(json.contains("\"type\":\"text\"")); let parsed: MessagePart = serde_json::from_str(&json).unwrap(); @@ -433,7 +465,9 @@ mod tests { #[test] fn test_message_part_serde_reasoning() { - let part = MessagePart::Reasoning { content: "thinking...".to_string() }; + let part = MessagePart::Reasoning { + content: "thinking...".to_string(), + }; let json = serde_json::to_string(&part).unwrap(); assert!(json.contains("\"type\":\"reasoning\"")); let parsed: MessagePart = serde_json::from_str(&json).unwrap(); @@ -453,7 +487,11 @@ mod tests { let json = serde_json::to_string(&part).unwrap(); let parsed: MessagePart = serde_json::from_str(&json).unwrap(); match parsed { - MessagePart::ToolCall { id, name, arguments } => { + MessagePart::ToolCall { + id, + name, + arguments, + } => { assert_eq!(id.unwrap(), "tc-1"); assert_eq!(name, "search"); assert_eq!(arguments.unwrap()["query"], "rust"); @@ -482,8 +520,11 @@ mod tests { #[test] fn test_token_usage_serde_roundtrip() { let usage = TokenUsage { - input_tokens: 500, output_tokens: 200, total_tokens: 700, - cache_creation_input_tokens: Some(100), cache_read_input_tokens: Some(50), + input_tokens: 500, + output_tokens: 200, + total_tokens: 700, + cache_creation_input_tokens: Some(100), + cache_read_input_tokens: Some(50), }; let json = serde_json::to_string(&usage).unwrap(); let parsed: TokenUsage = serde_json::from_str(&json).unwrap(); @@ -591,7 +632,9 @@ mod tests { let input = InputMessage { role: "user".to_string(), parts: vec![ - MessagePart::Text { content: "Hello".to_string() }, + MessagePart::Text { + content: "Hello".to_string(), + }, MessagePart::ToolCallResponse { id: Some("tc".to_string()), response: json!("ok"), @@ -607,7 +650,9 @@ mod tests { let output = OutputMessage { role: "assistant".to_string(), - parts: vec![MessagePart::Text { content: "Hi".to_string() }], + parts: vec![MessagePart::Text { + content: "Hi".to_string(), + }], name: None, finish_reason: Some("stop".to_string()), }; @@ -620,22 +665,35 @@ mod tests { fn test_llm_call_full_serde_roundtrip() { let req = make_request(); let mut call = LLMCall::new( - "call-rt".to_string(), 1000, "openai".to_string(), - "gpt-4o".to_string(), req, 42, "agent".to_string(), + "call-rt".to_string(), + 1000, + "openai".to_string(), + "gpt-4o".to_string(), + req, + 42, + "agent".to_string(), + ); + call.set_response( + LLMResponse { + messages: vec![OutputMessage { + role: "assistant".to_string(), + parts: vec![MessagePart::Text { + content: "world".to_string(), + }], + name: None, + finish_reason: Some("stop".to_string()), + }], + streamed: true, + raw_body: None, + }, + 5000, ); - call.set_response(LLMResponse { - messages: vec![OutputMessage { - role: "assistant".to_string(), - parts: vec![MessagePart::Text { content: "world".to_string() }], - name: None, - finish_reason: Some("stop".to_string()), - }], - streamed: true, - raw_body: None, - }, 5000); call.set_token_usage(TokenUsage { - input_tokens: 10, output_tokens: 5, total_tokens: 15, - cache_creation_input_tokens: None, cache_read_input_tokens: None, + input_tokens: 10, + output_tokens: 5, + total_tokens: 15, + cache_creation_input_tokens: None, + cache_read_input_tokens: None, }); call.metadata.insert("key".to_string(), "value".to_string()); diff --git a/src/agentsight/src/genai/storage.rs b/src/agentsight/src/genai/storage.rs index 7dc6afabe..a10f1dacc 100644 --- a/src/agentsight/src/genai/storage.rs +++ b/src/agentsight/src/genai/storage.rs @@ -3,14 +3,14 @@ //! This module provides storage capabilities for GenAI semantic events, //! including LLM calls, tool uses, and agent interactions. -use std::path::PathBuf; -use std::fs::{File, OpenOptions}; -use std::io::{Write, BufWriter, BufRead, BufReader}; -use serde::{Serialize, Deserialize}; use chrono::{DateTime, Utc}; +use serde::{Deserialize, Serialize}; +use std::fs::{File, OpenOptions}; +use std::io::{BufRead, BufReader, BufWriter, Write}; +use std::path::PathBuf; -use super::semantic::GenAISemanticEvent; use super::exporter::GenAIExporter; +use super::semantic::GenAISemanticEvent; /// Storage for GenAI semantic events pub struct GenAIStore { @@ -20,14 +20,14 @@ pub struct GenAIStore { impl GenAIStore { /// Create a new GenAI store with the given path - pub fn new(path: &PathBuf) -> Self { + pub fn new(path: &std::path::Path) -> Self { // Ensure parent directory exists if let Some(parent) = path.parent() { std::fs::create_dir_all(parent).ok(); } - + GenAIStore { - path: path.clone(), + path: path.to_path_buf(), } } @@ -43,29 +43,32 @@ impl GenAIStore { .create(true) .append(true) .open(&self.path)?; - + let mut writer = BufWriter::new(file); let json_line = serde_json::to_string(event)?; writeln!(writer, "{}", json_line)?; writer.flush()?; - + Ok(()) } /// Add multiple events - pub fn add_batch(&self, events: &[GenAISemanticEvent]) -> Result<(), Box> { + pub fn add_batch( + &self, + events: &[GenAISemanticEvent], + ) -> Result<(), Box> { let file = OpenOptions::new() .create(true) .append(true) .open(&self.path)?; - + let mut writer = BufWriter::new(file); for event in events { let json_line = serde_json::to_string(event)?; writeln!(writer, "{}", json_line)?; } writer.flush()?; - + Ok(()) } @@ -103,19 +106,19 @@ impl GenAIStore { for event in events { if let GenAISemanticEvent::LLMCall(call) = event { let call_time = DateTime::from_timestamp_nanos(call.start_timestamp_ns as i64); - - if let Some(start) = start_time { - if call_time < start { - continue; - } + + if let Some(start) = start_time + && call_time < start + { + continue; } - - if let Some(end) = end_time { - if call_time > end { - continue; - } + + if let Some(end) = end_time + && call_time > end + { + continue; } - + calls.push(call); } } @@ -124,22 +127,26 @@ impl GenAIStore { } /// Query events by process ID - pub fn query_by_pid(&self, pid: i32) -> Result, Box> { + pub fn query_by_pid( + &self, + pid: i32, + ) -> Result, Box> { let events = self.read_all()?; - Ok(events.into_iter().filter(|event| { - match event { + Ok(events + .into_iter() + .filter(|event| match event { GenAISemanticEvent::LLMCall(call) => call.pid == pid, GenAISemanticEvent::ToolUse(tool) => tool.pid == pid, GenAISemanticEvent::AgentInteraction(interaction) => interaction.pid == pid, GenAISemanticEvent::StreamChunk(chunk) => chunk.pid == pid, - } - }).collect()) + }) + .collect()) } /// Get statistics about stored events pub fn get_stats(&self) -> Result> { let events = self.read_all()?; - + let mut stats = GenAIStoreStats { total_events: events.len(), llm_calls: 0, diff --git a/src/agentsight/src/health/checker.rs b/src/agentsight/src/health/checker.rs index c0d84c46b..4c6a53304 100644 --- a/src/agentsight/src/health/checker.rs +++ b/src/agentsight/src/health/checker.rs @@ -100,109 +100,104 @@ impl HealthChecker { // All sessions for that pid are returned and each gets its own event. // - Deduplication key is (agent_name, session_id, conversation_id): at most one // agent_crash event is written per unique (agent, session, conversation) triple. - if !newly_offline.is_empty() { - if let Some(ref istore) = self.interruption_store { - let now_ns = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .map(|d| d.as_nanos() as i64) - .unwrap_or(0); + if !newly_offline.is_empty() + && let Some(ref istore) = self.interruption_store + { + let now_ns = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|d| d.as_nanos() as i64) + .unwrap_or(0); - // Group newly-offline entries by agent_name so multi-process agents - // (e.g. Cosh) are processed together in a single DB query. - let mut by_agent: HashMap> = HashMap::new(); - for offline in &newly_offline { - by_agent - .entry(offline.agent_name.clone()) - .or_default() - .push(offline); - } + // Group newly-offline entries by agent_name so multi-process agents + // (e.g. Cosh) are processed together in a single DB query. + let mut by_agent: HashMap> = HashMap::new(); + for offline in &newly_offline { + by_agent + .entry(offline.agent_name.clone()) + .or_default() + .push(offline); + } - // Dedup key: (agent_name, session_id, conversation_id) - let mut seen_conv: HashSet<(String, Option, Option)> = - HashSet::new(); + // Dedup key: (agent_name, session_id, conversation_id) + let mut seen_conv: HashSet<(String, Option, Option)> = HashSet::new(); - for (agent_name, group) in &by_agent { - let pids: Vec = group.iter().map(|o| o.pid as i32).collect(); - // Use the representative offline entry for metadata (pid shown in detail). - // When multiple pids exist (Cosh), use the largest pid (worker) as it is - // most likely the one that carried LLM traffic. - let rep = group.iter().max_by_key(|o| o.pid).unwrap(); + for (agent_name, group) in &by_agent { + let pids: Vec = group.iter().map(|o| o.pid as i32).collect(); + // Use the representative offline entry for metadata (pid shown in detail). + // When multiple pids exist (Cosh), use the largest pid (worker) as it is + // most likely the one that carried LLM traffic. + let rep = group.iter().max_by_key(|o| o.pid).unwrap(); - // ── Branch A: pending (in-flight) LLM calls ────────────────────────── - let pending_calls = self.get_pending_calls_for_pids(&pids); - if !pending_calls.is_empty() { - let mut by_conv: HashMap< - (Option, Option), - Vec<(String, Option)>, - > = HashMap::new(); - for (call_id, session_id, _trace_id, conversation_id) in &pending_calls { - by_conv - .entry((session_id.clone(), conversation_id.clone())) - .or_default() - .push((call_id.clone(), session_id.clone())); - } - for ((session_id, conversation_id), calls) in &by_conv { - let dedup_key = ( - agent_name.clone(), - session_id.clone(), - conversation_id.clone(), - ); - if !seen_conv.insert(dedup_key) { - log::debug!( - "Skipping duplicate agent_crash for {} session={:?} conversation={:?}", - agent_name, - session_id, - conversation_id - ); - continue; - } - let call_ids: Vec<&str> = - calls.iter().map(|(c, _)| c.as_str()).collect(); - let detail = serde_json::json!({ - "pid": rep.pid, - "agent_name": agent_name, - "exe_path": rep.exe_path.clone(), - "call_ids": call_ids, - }); - let event = InterruptionEvent::new( - InterruptionType::AgentCrash, - session_id.clone(), - None, - conversation_id.clone(), - None, - Some(rep.pid as i32), - Some(agent_name.clone()), - now_ns, - Some(detail), + // ── Branch A: pending (in-flight) LLM calls ────────────────────────── + let pending_calls = self.get_pending_calls_for_pids(&pids); + if !pending_calls.is_empty() { + #[allow(clippy::type_complexity)] + let mut by_conv: HashMap< + (Option, Option), + Vec<(String, Option)>, + > = HashMap::new(); + for (call_id, session_id, _trace_id, conversation_id) in &pending_calls { + by_conv + .entry((session_id.clone(), conversation_id.clone())) + .or_default() + .push((call_id.clone(), session_id.clone())); + } + for ((session_id, conversation_id), calls) in &by_conv { + let dedup_key = ( + agent_name.clone(), + session_id.clone(), + conversation_id.clone(), + ); + if !seen_conv.insert(dedup_key) { + log::debug!( + "Skipping duplicate agent_crash for {} session={:?} conversation={:?}", + agent_name, + session_id, + conversation_id ); - if let Err(e) = istore.insert(&event) { - log::warn!( - "Failed to record agent_crash for pid={}: {}", - rep.pid, - e - ); - } else { - log::info!( - "Recorded agent_crash for {} (pid={}, session={:?}, conversation={:?}, {} call(s))", - agent_name, - rep.pid, - session_id, - conversation_id, - calls.len() - ); - } + continue; } - for o in group { - self.mark_pending_interrupted(o.pid, "agent_crash"); - } - } else { - // No pending calls — treat as normal/graceful shutdown. - log::debug!( - "Agent {} (pids={:?}) exited with no pending calls — treating as normal shutdown", - agent_name, - pids + let call_ids: Vec<&str> = calls.iter().map(|(c, _)| c.as_str()).collect(); + let detail = serde_json::json!({ + "pid": rep.pid, + "agent_name": agent_name, + "exe_path": rep.exe_path.clone(), + "call_ids": call_ids, + }); + let event = InterruptionEvent::new( + InterruptionType::AgentCrash, + session_id.clone(), + None, + conversation_id.clone(), + None, + Some(rep.pid as i32), + Some(agent_name.clone()), + now_ns, + Some(detail), ); + if let Err(e) = istore.insert(&event) { + log::warn!("Failed to record agent_crash for pid={}: {}", rep.pid, e); + } else { + log::info!( + "Recorded agent_crash for {} (pid={}, session={:?}, conversation={:?}, {} call(s))", + agent_name, + rep.pid, + session_id, + conversation_id, + calls.len() + ); + } } + for o in group { + self.mark_pending_interrupted(o.pid, "agent_crash"); + } + } else { + // No pending calls — treat as normal/graceful shutdown. + log::debug!( + "Agent {} (pids={:?}) exited with no pending calls — treating as normal shutdown", + agent_name, + pids + ); } } } @@ -344,6 +339,7 @@ impl HealthChecker { /// Query pending LLM calls for multiple PIDs at once. /// /// Returns a list of (call_id, session_id, trace_id, conversation_id) tuples. + #[allow(clippy::type_complexity)] fn get_pending_calls_for_pids( &self, pids: &[i32], @@ -363,14 +359,14 @@ impl HealthChecker { /// Mark all pending calls for a PID as interrupted in genai_events. fn mark_pending_interrupted(&self, pid: u32, itype: &str) { - if let Some(ref genai_store) = self.genai_store { - if let Err(e) = genai_store.mark_pending_interrupted_for_pid(pid as i32, itype) { - log::warn!( - "Failed to mark pending calls as interrupted for pid={}: {}", - pid, - e - ); - } + if let Some(ref genai_store) = self.genai_store + && let Err(e) = genai_store.mark_pending_interrupted_for_pid(pid as i32, itype) + { + log::warn!( + "Failed to mark pending calls as interrupted for pid={}: {}", + pid, + e + ); } } } diff --git a/src/agentsight/src/health/port_detector.rs b/src/agentsight/src/health/port_detector.rs index a9eef261a..e2f376b0e 100644 --- a/src/agentsight/src/health/port_detector.rs +++ b/src/agentsight/src/health/port_detector.rs @@ -55,10 +55,12 @@ fn collect_socket_inodes(pid: u32) -> std::io::Result> { }; let link_str = link.to_string_lossy(); // Socket symlinks look like "socket:[12345]" - if let Some(inode_str) = link_str.strip_prefix("socket:[").and_then(|s| s.strip_suffix(']')) { - if let Ok(inode) = inode_str.parse::() { - inodes.insert(inode); - } + if let Some(inode_str) = link_str + .strip_prefix("socket:[") + .and_then(|s| s.strip_suffix(']')) + && let Ok(inode) = inode_str.parse::() + { + inodes.insert(inode); } } diff --git a/src/agentsight/src/health/store.rs b/src/agentsight/src/health/store.rs index b7c9e6613..11700b387 100644 --- a/src/agentsight/src/health/store.rs +++ b/src/agentsight/src/health/store.rs @@ -55,6 +55,12 @@ pub struct HealthStore { pub last_scan_time: u64, } +impl Default for HealthStore { + fn default() -> Self { + Self::new() + } +} + impl HealthStore { pub fn new() -> Self { Self { diff --git a/src/agentsight/src/interruption/detector.rs b/src/agentsight/src/interruption/detector.rs index d021c74c1..4a1f063cc 100644 --- a/src/agentsight/src/interruption/detector.rs +++ b/src/agentsight/src/interruption/detector.rs @@ -4,8 +4,8 @@ //! `InterruptionDetector::detect(call)` checks a single call against all //! single-call rules and returns any detected interruption events. -use crate::genai::semantic::LLMCall; use super::types::{InterruptionEvent, InterruptionType}; +use crate::genai::semantic::LLMCall; /// Configuration for the interruption detector pub struct DetectorConfig { @@ -46,23 +46,28 @@ impl InterruptionDetector { let mut events = Vec::new(); let session_id = call.metadata.get("session_id").cloned(); - let trace_id = call.metadata.get("response_id").cloned(); + let trace_id = call.metadata.get("response_id").cloned(); let conversation_id = call.metadata.get("conversation_id").cloned(); - let call_id = Some(call.call_id.clone()); - let pid = Some(call.pid); + let call_id = Some(call.call_id.clone()); + let pid = Some(call.pid); let agent_name = call.agent_name.clone(); - let status_code: u16 = call.metadata.get("status_code") + let status_code: u16 = call + .metadata + .get("status_code") .and_then(|s| s.parse().ok()) .unwrap_or(200); // Helper: scan error message / response body for context-overflow keywords let error_text = call.error.as_deref().unwrap_or(""); - let response_body = call.metadata.get("response_body").map(|s| s.as_str()).unwrap_or(""); + let response_body = call + .metadata + .get("response_body") + .map(|s| s.as_str()) + .unwrap_or(""); let combined_error = format!("{} {}", error_text, response_body).to_ascii_lowercase(); - let is_context_overflow = - combined_error.contains("context_length_exceeded") + let is_context_overflow = combined_error.contains("context_length_exceeded") || combined_error.contains("maximum context length") || combined_error.contains("context window") || combined_error.contains("context_length") @@ -87,8 +92,12 @@ impl InterruptionDetector { }); events.push(InterruptionEvent::new( InterruptionType::ContextOverflow, - session_id.clone(), trace_id.clone(), conversation_id.clone(), call_id.clone(), - pid, agent_name.clone(), + session_id.clone(), + trace_id.clone(), + conversation_id.clone(), + call_id.clone(), + pid, + agent_name.clone(), call.end_timestamp_ns as i64, Some(detail), )); @@ -104,8 +113,12 @@ impl InterruptionDetector { }); events.push(InterruptionEvent::new( InterruptionType::LlmError, - session_id.clone(), trace_id.clone(), conversation_id.clone(), call_id.clone(), - pid, agent_name.clone(), + session_id.clone(), + trace_id.clone(), + conversation_id.clone(), + call_id.clone(), + pid, + agent_name.clone(), call.end_timestamp_ns as i64, Some(detail), )); @@ -113,7 +126,11 @@ impl InterruptionDetector { } // ── 3. SSE truncated ────────────────────────────────────────────────── - let is_sse = call.metadata.get("is_sse").map(|s| s == "true").unwrap_or(false); + let is_sse = call + .metadata + .get("is_sse") + .map(|s| s == "true") + .unwrap_or(false); if is_sse && call.response.messages.is_empty() && call.duration_ns >= self.config.sse_min_duration_ns @@ -125,36 +142,46 @@ impl InterruptionDetector { }); events.push(InterruptionEvent::new( InterruptionType::SseTruncated, - session_id.clone(), trace_id.clone(), conversation_id.clone(), call_id.clone(), - pid, agent_name.clone(), + session_id.clone(), + trace_id.clone(), + conversation_id.clone(), + call_id.clone(), + pid, + agent_name.clone(), call.end_timestamp_ns as i64, Some(detail), )); } // ── 4. Token limit (output capped by max_tokens) ────────────────────── - let finish_reason = call.response.messages.first() + let finish_reason = call + .response + .messages + .first() .and_then(|m| m.finish_reason.as_deref()); - if finish_reason == Some("length") { - if let Some(max_tokens) = call.request.max_tokens { - if let Some(usage) = &call.token_usage { - let ratio = usage.output_tokens as f64 / max_tokens as f64; - if ratio >= self.config.token_limit_ratio { - let detail = serde_json::json!({ - "model": call.model, - "output_tokens": usage.output_tokens, - "max_tokens": max_tokens, - "ratio": ratio, - }); - events.push(InterruptionEvent::new( - InterruptionType::TokenLimit, - session_id.clone(), trace_id.clone(), conversation_id.clone(), call_id.clone(), - pid, agent_name.clone(), - call.end_timestamp_ns as i64, - Some(detail), - )); - } - } + if finish_reason == Some("length") + && let Some(max_tokens) = call.request.max_tokens + && let Some(usage) = &call.token_usage + { + let ratio = usage.output_tokens as f64 / max_tokens as f64; + if ratio >= self.config.token_limit_ratio { + let detail = serde_json::json!({ + "model": call.model, + "output_tokens": usage.output_tokens, + "max_tokens": max_tokens, + "ratio": ratio, + }); + events.push(InterruptionEvent::new( + InterruptionType::TokenLimit, + session_id.clone(), + trace_id.clone(), + conversation_id.clone(), + call_id.clone(), + pid, + agent_name.clone(), + call.end_timestamp_ns as i64, + Some(detail), + )); } } @@ -163,34 +190,36 @@ impl InterruptionDetector { // already exceeds the context window but response still arrives. // Detect via input_tokens >= model context ceiling (heuristic: >90% of // a well-known ceiling, or when input_tokens >> max_tokens). - if finish_reason == Some("length") { - if let Some(usage) = &call.token_usage { - if let Some(max_tokens) = call.request.max_tokens { - // If input tokens are much larger than the output cap, this - // is almost certainly a context-length issue, not output truncation. - if usage.input_tokens > max_tokens * 4 { - let detail = serde_json::json!({ - "model": call.model, - "input_tokens": usage.input_tokens, - "max_tokens": max_tokens, - "finish_reason": "length", - "note": "input_tokens >> max_tokens suggests context overflow", - }); - events.push(InterruptionEvent::new( - InterruptionType::ContextOverflow, - session_id.clone(), trace_id.clone(), conversation_id.clone(), call_id.clone(), - pid, agent_name.clone(), - call.end_timestamp_ns as i64, - Some(detail), - )); - } - } + if finish_reason == Some("length") + && let Some(usage) = &call.token_usage + && let Some(max_tokens) = call.request.max_tokens + { + // If input tokens are much larger than the output cap, this + // is almost certainly a context-length issue, not output truncation. + if usage.input_tokens > max_tokens * 4 { + let detail = serde_json::json!({ + "model": call.model, + "input_tokens": usage.input_tokens, + "max_tokens": max_tokens, + "finish_reason": "length", + "note": "input_tokens >> max_tokens suggests context overflow", + }); + events.push(InterruptionEvent::new( + InterruptionType::ContextOverflow, + session_id.clone(), + trace_id.clone(), + conversation_id.clone(), + call_id.clone(), + pid, + agent_name.clone(), + call.end_timestamp_ns as i64, + Some(detail), + )); } } events } - } #[cfg(test)] @@ -231,9 +260,7 @@ mod tests { pid: 1234, process_name: "agent".to_string(), agent_name: Some("TestAgent".to_string()), - metadata: HashMap::from([ - ("status_code".to_string(), "200".to_string()), - ]), + metadata: HashMap::from([("status_code".to_string(), "200".to_string())]), } } @@ -250,38 +277,54 @@ mod tests { let detector = InterruptionDetector::default(); let mut call = make_base_call(); call.error = Some("context_length_exceeded".to_string()); - call.metadata.insert("status_code".to_string(), "400".to_string()); + call.metadata + .insert("status_code".to_string(), "400".to_string()); let events = detector.detect(&call); assert_eq!(events.len(), 1); - assert_eq!(events[0].interruption_type, InterruptionType::ContextOverflow); + assert_eq!( + events[0].interruption_type, + InterruptionType::ContextOverflow + ); } #[test] fn test_detect_context_overflow_http_413() { let detector = InterruptionDetector::default(); let mut call = make_base_call(); - call.metadata.insert("status_code".to_string(), "413".to_string()); + call.metadata + .insert("status_code".to_string(), "413".to_string()); let events = detector.detect(&call); assert_eq!(events.len(), 1); - assert_eq!(events[0].interruption_type, InterruptionType::ContextOverflow); + assert_eq!( + events[0].interruption_type, + InterruptionType::ContextOverflow + ); } #[test] fn test_detect_context_overflow_response_body() { let detector = InterruptionDetector::default(); let mut call = make_base_call(); - call.metadata.insert("status_code".to_string(), "400".to_string()); - call.metadata.insert("response_body".to_string(), "maximum context length is 128k".to_string()); + call.metadata + .insert("status_code".to_string(), "400".to_string()); + call.metadata.insert( + "response_body".to_string(), + "maximum context length is 128k".to_string(), + ); let events = detector.detect(&call); assert_eq!(events.len(), 1); - assert_eq!(events[0].interruption_type, InterruptionType::ContextOverflow); + assert_eq!( + events[0].interruption_type, + InterruptionType::ContextOverflow + ); } #[test] fn test_detect_llm_error_http_500() { let detector = InterruptionDetector::default(); let mut call = make_base_call(); - call.metadata.insert("status_code".to_string(), "500".to_string()); + call.metadata + .insert("status_code".to_string(), "500".to_string()); let events = detector.detect(&call); assert_eq!(events.len(), 1); assert_eq!(events[0].interruption_type, InterruptionType::LlmError); @@ -301,7 +344,8 @@ mod tests { fn test_detect_sse_truncated() { let detector = InterruptionDetector::default(); let mut call = make_base_call(); - call.metadata.insert("is_sse".to_string(), "true".to_string()); + call.metadata + .insert("is_sse".to_string(), "true".to_string()); call.duration_ns = 2_000_000_000; // > 1 second min // response.messages is empty let events = detector.detect(&call); @@ -313,7 +357,8 @@ mod tests { fn test_no_sse_truncated_short_duration() { let detector = InterruptionDetector::default(); let mut call = make_base_call(); - call.metadata.insert("is_sse".to_string(), "true".to_string()); + call.metadata + .insert("is_sse".to_string(), "true".to_string()); call.duration_ns = 500_000_000; // < 1 second min let events = detector.detect(&call); assert!(events.is_empty()); @@ -384,29 +429,41 @@ mod tests { }]; let events = detector.detect(&call); // Should have context_overflow (from rule 5) - assert!(events.iter().any(|e| e.interruption_type == InterruptionType::ContextOverflow)); + assert!( + events + .iter() + .any(|e| e.interruption_type == InterruptionType::ContextOverflow) + ); } #[test] fn test_context_overflow_supersedes_llm_error() { let detector = InterruptionDetector::default(); let mut call = make_base_call(); - call.metadata.insert("status_code".to_string(), "400".to_string()); + call.metadata + .insert("status_code".to_string(), "400".to_string()); call.error = Some("context_length_exceeded: max 128000 tokens".to_string()); let events = detector.detect(&call); // Should be context_overflow, NOT llm_error assert_eq!(events.len(), 1); - assert_eq!(events[0].interruption_type, InterruptionType::ContextOverflow); + assert_eq!( + events[0].interruption_type, + InterruptionType::ContextOverflow + ); } #[test] fn test_event_metadata_fields() { let detector = InterruptionDetector::default(); let mut call = make_base_call(); - call.metadata.insert("status_code".to_string(), "500".to_string()); - call.metadata.insert("session_id".to_string(), "sess-abc".to_string()); - call.metadata.insert("response_id".to_string(), "trace-xyz".to_string()); - call.metadata.insert("conversation_id".to_string(), "conv-123".to_string()); + call.metadata + .insert("status_code".to_string(), "500".to_string()); + call.metadata + .insert("session_id".to_string(), "sess-abc".to_string()); + call.metadata + .insert("response_id".to_string(), "trace-xyz".to_string()); + call.metadata + .insert("conversation_id".to_string(), "conv-123".to_string()); let events = detector.detect(&call); assert_eq!(events.len(), 1); assert_eq!(events[0].session_id, Some("sess-abc".to_string())); @@ -440,6 +497,10 @@ mod tests { finish_reason: Some("length".to_string()), }]; let events = detector.detect(&call); - assert!(events.iter().any(|e| e.interruption_type == InterruptionType::TokenLimit)); + assert!( + events + .iter() + .any(|e| e.interruption_type == InterruptionType::TokenLimit) + ); } } diff --git a/src/agentsight/src/interruption/mod.rs b/src/agentsight/src/interruption/mod.rs index 39e59ae2d..35761ff83 100644 --- a/src/agentsight/src/interruption/mod.rs +++ b/src/agentsight/src/interruption/mod.rs @@ -1,9 +1,9 @@ //! Interruption module — public API. -pub mod types; pub mod detector; pub mod oom_recovery; +pub mod types; -pub use types::{InterruptionEvent, InterruptionType, Severity}; -pub use detector::{InterruptionDetector, DetectorConfig}; +pub use detector::{DetectorConfig, InterruptionDetector}; pub use oom_recovery::recover_oom_events; +pub use types::{InterruptionEvent, InterruptionType, Severity}; diff --git a/src/agentsight/src/interruption/types.rs b/src/agentsight/src/interruption/types.rs index b22f0f4d1..b299e864a 100644 --- a/src/agentsight/src/interruption/types.rs +++ b/src/agentsight/src/interruption/types.rs @@ -22,20 +22,21 @@ impl InterruptionType { /// String identifier stored in the database pub fn as_str(&self) -> &'static str { match self { - Self::LlmError => "llm_error", - Self::SseTruncated => "sse_truncated", - Self::AgentCrash => "agent_crash", - Self::TokenLimit => "token_limit", + Self::LlmError => "llm_error", + Self::SseTruncated => "sse_truncated", + Self::AgentCrash => "agent_crash", + Self::TokenLimit => "token_limit", Self::ContextOverflow => "context_overflow", } } + #[allow(clippy::should_implement_trait)] pub fn from_str(s: &str) -> Option { match s { - "llm_error" => Some(Self::LlmError), - "sse_truncated" => Some(Self::SseTruncated), - "agent_crash" => Some(Self::AgentCrash), - "token_limit" => Some(Self::TokenLimit), + "llm_error" => Some(Self::LlmError), + "sse_truncated" => Some(Self::SseTruncated), + "agent_crash" => Some(Self::AgentCrash), + "token_limit" => Some(Self::TokenLimit), "context_overflow" => Some(Self::ContextOverflow), _ => None, } @@ -44,11 +45,11 @@ impl InterruptionType { /// Default severity for this interruption type pub fn default_severity(&self) -> Severity { match self { - Self::AgentCrash => Severity::Critical, - Self::LlmError => Severity::High, - Self::SseTruncated => Severity::High, + Self::AgentCrash => Severity::Critical, + Self::LlmError => Severity::High, + Self::SseTruncated => Severity::High, Self::ContextOverflow => Severity::High, - Self::TokenLimit => Severity::Medium, + Self::TokenLimit => Severity::Medium, } } } @@ -67,9 +68,9 @@ impl Severity { pub fn as_str(&self) -> &'static str { match self { Self::Critical => "critical", - Self::High => "high", - Self::Medium => "medium", - Self::Low => "low", + Self::High => "high", + Self::Medium => "medium", + Self::Low => "low", } } @@ -77,9 +78,9 @@ impl Severity { pub fn weight(&self) -> u8 { match self { Self::Critical => 4, - Self::High => 3, - Self::Medium => 2, - Self::Low => 1, + Self::High => 3, + Self::Medium => 2, + Self::Low => 1, } } } @@ -107,6 +108,7 @@ pub struct InterruptionEvent { impl InterruptionEvent { /// Create a new unresolved interruption event with auto-generated ID + #[allow(clippy::too_many_arguments)] pub fn new( itype: InterruptionType, session_id: Option, @@ -159,27 +161,60 @@ mod tests { assert_eq!(InterruptionType::SseTruncated.as_str(), "sse_truncated"); assert_eq!(InterruptionType::AgentCrash.as_str(), "agent_crash"); assert_eq!(InterruptionType::TokenLimit.as_str(), "token_limit"); - assert_eq!(InterruptionType::ContextOverflow.as_str(), "context_overflow"); + assert_eq!( + InterruptionType::ContextOverflow.as_str(), + "context_overflow" + ); } #[test] fn test_interruption_type_from_str() { - assert_eq!(InterruptionType::from_str("llm_error"), Some(InterruptionType::LlmError)); - assert_eq!(InterruptionType::from_str("sse_truncated"), Some(InterruptionType::SseTruncated)); - assert_eq!(InterruptionType::from_str("agent_crash"), Some(InterruptionType::AgentCrash)); - assert_eq!(InterruptionType::from_str("token_limit"), Some(InterruptionType::TokenLimit)); - assert_eq!(InterruptionType::from_str("context_overflow"), Some(InterruptionType::ContextOverflow)); + assert_eq!( + InterruptionType::from_str("llm_error"), + Some(InterruptionType::LlmError) + ); + assert_eq!( + InterruptionType::from_str("sse_truncated"), + Some(InterruptionType::SseTruncated) + ); + assert_eq!( + InterruptionType::from_str("agent_crash"), + Some(InterruptionType::AgentCrash) + ); + assert_eq!( + InterruptionType::from_str("token_limit"), + Some(InterruptionType::TokenLimit) + ); + assert_eq!( + InterruptionType::from_str("context_overflow"), + Some(InterruptionType::ContextOverflow) + ); assert_eq!(InterruptionType::from_str("unknown"), None); assert_eq!(InterruptionType::from_str(""), None); } #[test] fn test_interruption_type_default_severity() { - assert_eq!(InterruptionType::AgentCrash.default_severity(), Severity::Critical); - assert_eq!(InterruptionType::LlmError.default_severity(), Severity::High); - assert_eq!(InterruptionType::SseTruncated.default_severity(), Severity::High); - assert_eq!(InterruptionType::ContextOverflow.default_severity(), Severity::High); - assert_eq!(InterruptionType::TokenLimit.default_severity(), Severity::Medium); + assert_eq!( + InterruptionType::AgentCrash.default_severity(), + Severity::Critical + ); + assert_eq!( + InterruptionType::LlmError.default_severity(), + Severity::High + ); + assert_eq!( + InterruptionType::SseTruncated.default_severity(), + Severity::High + ); + assert_eq!( + InterruptionType::ContextOverflow.default_severity(), + Severity::High + ); + assert_eq!( + InterruptionType::TokenLimit.default_severity(), + Severity::Medium + ); } #[test] @@ -236,7 +271,12 @@ mod tests { fn test_interruption_event_new_no_detail() { let event = InterruptionEvent::new( InterruptionType::AgentCrash, - None, None, None, None, None, None, + None, + None, + None, + None, + None, + None, 500_000, None, ); @@ -276,7 +316,12 @@ mod tests { #[test] fn test_severity_serde_roundtrip() { - let severities = vec![Severity::Critical, Severity::High, Severity::Medium, Severity::Low]; + let severities = vec![ + Severity::Critical, + Severity::High, + Severity::Medium, + Severity::Low, + ]; for s in severities { let json = serde_json::to_string(&s).unwrap(); let back: Severity = serde_json::from_str(&json).unwrap(); diff --git a/src/agentsight/src/lib.rs b/src/agentsight/src/lib.rs index 098e8cd5e..362ca3bb0 100644 --- a/src/agentsight/src/lib.rs +++ b/src/agentsight/src/lib.rs @@ -24,61 +24,51 @@ //! sight.run()?; // blocking event loop //! ``` -pub mod probes; pub mod config; +pub mod probes; // Re-export config types pub use config::{AgentsightConfig, default_base_path}; -pub mod event; -pub mod parser; pub mod aggregator; pub mod analyzer; -pub mod storage; +pub mod atif; pub mod chrome_trace; pub mod discovery; -pub mod health; -pub mod tokenizer; +pub mod event; +pub mod ffi; pub mod genai; -pub mod atif; -pub mod response_map; +pub mod health; pub mod interruption; -pub mod skill_metrics; +pub mod parser; +pub mod response_map; #[cfg(feature = "server")] pub mod server; +pub mod skill_metrics; +pub mod storage; +pub mod tokenizer; mod unified; -pub mod ffi; // Re-export common types for convenience pub use aggregator::{ - Aggregator, AggregatedResult, - HttpConnectionAggregator, ConnectionId, ConnectionState, - HttpPair, - ProcessEventAggregator, AggregatedProcess, - AggregatedResponse, -}; -pub use parser::{ - HttpParser, ParsedHttpMessage, ParsedRequest, ParsedResponse, - SseParser, ParsedSseEvent, - ProcTraceParser, ParsedProcEvent, ProcEventType, - Http2Parser, Http2FrameType, ParsedHttp2Frame, - Parser, ParsedMessage, ParseResult, + AggregatedProcess, AggregatedResponse, AggregatedResult, Aggregator, ConnectionId, + ConnectionState, HttpConnectionAggregator, HttpPair, ProcessEventAggregator, }; pub use analyzer::{ - AuditAnalyzer, AuditEventType, AuditExtra, AuditRecord, AuditSummary, - TokenParser, TokenUsage, TokenRecord, LLMProvider, - MessageParser, ParsedApiMessage, - OpenAIRequest, OpenAIResponse, OpenAIChatMessage, OpenAIContent, OpenAIUsage, - AnthropicRequest, AnthropicResponse, AnthropicMessage, AnthropicUsage, - MessageRole, - AnalysisResult, PromptTokenCount, HttpRecord, Analyzer, + AnalysisResult, Analyzer, AnthropicMessage, AnthropicRequest, AnthropicResponse, + AnthropicUsage, AuditAnalyzer, AuditEventType, AuditExtra, AuditRecord, AuditSummary, + HttpRecord, LLMProvider, MessageParser, MessageRole, OpenAIChatMessage, OpenAIContent, + OpenAIRequest, OpenAIResponse, OpenAIUsage, ParsedApiMessage, PromptTokenCount, TokenParser, + TokenRecord, TokenUsage, +}; +pub use chrome_trace::{ChromeTraceEvent, ToChromeTraceEvent, TraceArgs, next_flow_id, ns_to_us}; +pub use parser::{ + Http2FrameType, Http2Parser, HttpParser, ParseResult, ParsedHttp2Frame, ParsedHttpMessage, + ParsedMessage, ParsedProcEvent, ParsedRequest, ParsedResponse, ParsedSseEvent, Parser, + ProcEventType, ProcTraceParser, SseParser, }; -pub use chrome_trace::{ChromeTraceEvent, TraceArgs, ToChromeTraceEvent, ns_to_us, next_flow_id}; pub use storage::{ - Storage, StorageBackend, SqliteConfig, - SqliteStore, AuditStore, - TokenStore, TokenQuery, - HttpStore, - TimePeriod, TokenQueryResult, TokenBreakdown, TokenComparison, Trend, + AuditStore, HttpStore, SqliteConfig, SqliteStore, Storage, StorageBackend, TimePeriod, + TokenBreakdown, TokenComparison, TokenQuery, TokenQueryResult, TokenStore, Trend, format_tokens, format_tokens_with_commas, }; @@ -92,12 +82,12 @@ pub use probes::FileWatchEvent; pub use response_map::ResponseSessionMapper; // Re-export discovery types -pub use discovery::{AgentInfo, AgentScanner, CmdlineGlobMatcher, DiscoveredAgent, ProcessContext}; pub use config::default_cmdline_rules; +pub use discovery::{AgentInfo, AgentScanner, CmdlineGlobMatcher, DiscoveredAgent, ProcessContext}; // Re-export genai types pub use genai::{ - GenAIBuilder, GenAISemanticEvent, LLMCall, LLMRequest, LLMResponse, - MessagePart, InputMessage, OutputMessage, ToolUse, AgentInteraction, StreamChunk, ToolDefinition, - GenAIStore, GenAIStoreStats, LogtailExporter, GenAIExporter, + AgentInteraction, GenAIBuilder, GenAIExporter, GenAISemanticEvent, GenAIStore, GenAIStoreStats, + InputMessage, LLMCall, LLMRequest, LLMResponse, LogtailExporter, MessagePart, OutputMessage, + StreamChunk, ToolDefinition, ToolUse, }; diff --git a/src/agentsight/src/parser/http/parser.rs b/src/agentsight/src/parser/http/parser.rs index c58e0546a..59a4bfd00 100644 --- a/src/agentsight/src/parser/http/parser.rs +++ b/src/agentsight/src/parser/http/parser.rs @@ -178,7 +178,8 @@ mod tests { #[test] fn test_parse_http_response() { - let data = b"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\n\r\n{\"id\":\"chatcmpl-123\"}"; + let data = + b"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\n\r\n{\"id\":\"chatcmpl-123\"}"; let event = make_ssl_event(data); let parser = HttpParser::new(); let result = parser.parse(event); @@ -187,7 +188,10 @@ mod tests { ParsedHttpMessage::Response(resp) => { assert_eq!(resp.status_code, 200); assert_eq!(resp.reason, "OK"); - assert_eq!(resp.headers.get("content-type").unwrap(), "application/json"); + assert_eq!( + resp.headers.get("content-type").unwrap(), + "application/json" + ); assert!(resp.body_len > 0); } _ => panic!("Expected Response"), @@ -244,7 +248,10 @@ mod tests { assert!(result.is_ok()); match result.unwrap() { ParsedHttpMessage::Response(resp) => { - assert_eq!(resp.headers.get("content-type").unwrap(), "text/event-stream"); + assert_eq!( + resp.headers.get("content-type").unwrap(), + "text/event-stream" + ); assert_eq!(resp.headers.get("transfer-encoding").unwrap(), "chunked"); assert_eq!(resp.headers.get("connection").unwrap(), "keep-alive"); } @@ -255,11 +262,18 @@ mod tests { #[test] fn test_ssl_event_is_http_request() { let event = SslEvent { - source: 0, timestamp_ns: 0, delta_ns: 0, - pid: 1, tid: 1, uid: 0, len: 10, rw: 1, + source: 0, + timestamp_ns: 0, + delta_ns: 0, + pid: 1, + tid: 1, + uid: 0, + len: 10, + rw: 1, comm: String::new(), buf: b"POST /api HTTP/1.1\r\n".to_vec(), - is_handshake: false, ssl_ptr: 0, + is_handshake: false, + ssl_ptr: 0, }; assert!(event.is_http_request()); assert!(event.is_http()); @@ -269,11 +283,18 @@ mod tests { #[test] fn test_ssl_event_is_http_response() { let event = SslEvent { - source: 0, timestamp_ns: 0, delta_ns: 0, - pid: 1, tid: 1, uid: 0, len: 15, rw: 0, + source: 0, + timestamp_ns: 0, + delta_ns: 0, + pid: 1, + tid: 1, + uid: 0, + len: 15, + rw: 0, comm: String::new(), buf: b"HTTP/1.1 200 OK\r\n".to_vec(), - is_handshake: false, ssl_ptr: 0, + is_handshake: false, + ssl_ptr: 0, }; assert!(event.is_http_response()); assert!(event.is_http()); @@ -283,11 +304,18 @@ mod tests { #[test] fn test_ssl_event_is_http2_preface() { let event = SslEvent { - source: 0, timestamp_ns: 0, delta_ns: 0, - pid: 1, tid: 1, uid: 0, len: 24, rw: 1, + source: 0, + timestamp_ns: 0, + delta_ns: 0, + pid: 1, + tid: 1, + uid: 0, + len: 24, + rw: 1, comm: String::new(), buf: b"PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n".to_vec(), - is_handshake: false, ssl_ptr: 0, + is_handshake: false, + ssl_ptr: 0, }; assert!(event.is_http2_preface()); assert!(event.is_http2()); @@ -297,11 +325,18 @@ mod tests { fn test_ssl_event_is_http2_frame() { // Valid HTTP/2 frame: length=0, type=4(SETTINGS), flags=0, stream_id=0 let event = SslEvent { - source: 0, timestamp_ns: 0, delta_ns: 0, - pid: 1, tid: 1, uid: 0, len: 9, rw: 0, + source: 0, + timestamp_ns: 0, + delta_ns: 0, + pid: 1, + tid: 1, + uid: 0, + len: 9, + rw: 0, comm: String::new(), buf: vec![0, 0, 0, 4, 0, 0, 0, 0, 0], - is_handshake: false, ssl_ptr: 0, + is_handshake: false, + ssl_ptr: 0, }; assert!(event.is_http2_frame()); assert!(event.is_http2()); @@ -311,11 +346,18 @@ mod tests { fn test_ssl_event_not_http2_frame_bad_type() { // Frame type > 9 is invalid let event = SslEvent { - source: 0, timestamp_ns: 0, delta_ns: 0, - pid: 1, tid: 1, uid: 0, len: 9, rw: 0, + source: 0, + timestamp_ns: 0, + delta_ns: 0, + pid: 1, + tid: 1, + uid: 0, + len: 9, + rw: 0, comm: String::new(), buf: vec![0, 0, 0, 10, 0, 0, 0, 0, 0], - is_handshake: false, ssl_ptr: 0, + is_handshake: false, + ssl_ptr: 0, }; assert!(!event.is_http2_frame()); } @@ -323,11 +365,18 @@ mod tests { #[test] fn test_ssl_event_payload_and_helpers() { let event = SslEvent { - source: 0, timestamp_ns: 100, delta_ns: 0, - pid: 42, tid: 42, uid: 0, len: 5, rw: 0, + source: 0, + timestamp_ns: 100, + delta_ns: 0, + pid: 42, + tid: 42, + uid: 0, + len: 5, + rw: 0, comm: "curl".to_string(), buf: b"hello".to_vec(), - is_handshake: false, ssl_ptr: 0x2000, + is_handshake: false, + ssl_ptr: 0x2000, }; assert_eq!(event.payload(), Some("hello")); assert_eq!(event.comm_str(), "curl"); diff --git a/src/agentsight/src/parser/http/request.rs b/src/agentsight/src/parser/http/request.rs index e184c4f77..d2a58fe5e 100644 --- a/src/agentsight/src/parser/http/request.rs +++ b/src/agentsight/src/parser/http/request.rs @@ -1,22 +1,22 @@ //! HTTP Request types +use crate::chrome_trace::{ChromeTraceEvent, ToChromeTraceEvent, TraceArgs, ns_to_us}; +use crate::probes::sslsniff::SslEvent; +use serde_json::json; use std::collections::HashMap; use std::fmt; use std::rc::Rc; -use crate::probes::sslsniff::SslEvent; -use crate::chrome_trace::{TraceArgs, ToChromeTraceEvent, ChromeTraceEvent, ns_to_us}; -use serde_json::json; /// 解析后的 HTTP Request #[derive(Clone)] pub struct ParsedRequest { - pub method: String, // GET, POST, etc. - pub path: String, // /api/chat - pub version: u8, // 11 for HTTP/1.1 + pub method: String, // GET, POST, etc. + pub path: String, // /api/chat + pub version: u8, // 11 for HTTP/1.1 pub headers: HashMap, - pub body_offset: usize, // body 在 source_event.buf 中的起始位置 - pub body_len: usize, // body 长度 - pub source_event: Rc, // 原始 SslEvent (Rc 避免拷贝) + pub body_offset: usize, // body 在 source_event.buf 中的起始位置 + pub body_len: usize, // body 长度 + pub source_event: Rc, // 原始 SslEvent (Rc 避免拷贝) /// 重组后的完整 body（跨多事件聚合时使用） pub reassembled_body: Option>, } @@ -34,9 +34,9 @@ impl ParsedRequest { pub fn body_str(&self) -> &str { std::str::from_utf8(self.body()).unwrap_or("") } - + /// 尝试将 body 解析为 JSON - /// + /// /// 如果 body 是有效的 UTF-8 且是有效的 JSON，返回解析后的 Value。 /// 如果直接解析失败，会尝试剥离 HTTP chunked transfer encoding 后再解析。 pub fn json_body(&self) -> Option { @@ -107,7 +107,7 @@ impl ParsedRequest { impl TraceArgs for ParsedRequest { fn to_trace_args(&self) -> serde_json::Value { let mut args = serde_json::Map::new(); - + // Basic request info args.insert("method".to_string(), json!(&self.method)); args.insert("path".to_string(), json!(&self.path)); @@ -120,16 +120,16 @@ impl TraceArgs for ParsedRequest { args.insert("pid".to_string(), json!(self.source_event.pid)); args.insert("tid".to_string(), json!(self.source_event.tid)); args.insert("comm".to_string(), json!(self.source_event.comm_str())); - + // Add headers if present if !self.headers.is_empty() { args.insert("headers".to_string(), json!(&self.headers)); } - + // Add body info if present if self.body_len > 0 { args.insert("body_length".to_string(), json!(self.body_len)); - + // Try to parse as JSON first, fallback to full string if let Some(json_body) = self.json_body() { args.insert("body".to_string(), json_body); @@ -140,7 +140,7 @@ impl TraceArgs for ParsedRequest { } } } - + serde_json::Value::Object(args) } } @@ -148,10 +148,10 @@ impl TraceArgs for ParsedRequest { impl ToChromeTraceEvent for ParsedRequest { fn to_chrome_trace_events(&self) -> Vec { let ts_us = ns_to_us(self.source_event.timestamp_ns); - + // Minimum duration: 10ms = 10,000 microseconds const MIN_DUR_US: u64 = 10_000; - + let event = ChromeTraceEvent::complete( format!("{} {}", self.method, self.path), "http.request", @@ -161,7 +161,7 @@ impl ToChromeTraceEvent for ParsedRequest { MIN_DUR_US, ) .with_trace_args(self); - + vec![event] } } @@ -173,22 +173,22 @@ impl fmt::Debug for ParsedRequest { .field("method", &self.method) .field("path", &self.path) .field("version", &format!("HTTP/1.{}", self.version)); - + // Format headers debug.field("headers", &self.headers); - + // Format body with smart detection let body = self.body(); if !body.is_empty() { debug.field("body", &format_body(body)); } - + // Add metadata from source_event debug .field("pid", &self.source_event.pid) .field("tid", &self.source_event.tid) .field("timestamp_ns", &self.source_event.timestamp_ns); - + debug.finish() } } @@ -205,7 +205,11 @@ fn format_body(data: &[u8]) -> String { format!("(text, {} bytes)\n{}", data.len(), text) } else { // Binary data - show as base64 - format!("(binary, {} bytes)\n{}", data.len(), base64::Engine::encode(&base64::engine::general_purpose::STANDARD, data)) + format!( + "(binary, {} bytes)\n{}", + data.len(), + base64::Engine::encode(&base64::engine::general_purpose::STANDARD, data) + ) } } diff --git a/src/agentsight/src/parser/http/response.rs b/src/agentsight/src/parser/http/response.rs index 7f0fc337b..6ad3dcea6 100644 --- a/src/agentsight/src/parser/http/response.rs +++ b/src/agentsight/src/parser/http/response.rs @@ -1,22 +1,22 @@ //! HTTP Response types +use crate::chrome_trace::{ChromeTraceEvent, ToChromeTraceEvent, TraceArgs, ns_to_us}; +use crate::probes::sslsniff::SslEvent; +use serde_json::json; use std::collections::HashMap; use std::fmt; use std::rc::Rc; -use crate::probes::sslsniff::SslEvent; -use crate::chrome_trace::{TraceArgs, ToChromeTraceEvent, ChromeTraceEvent, ns_to_us}; -use serde_json::json; /// 解析后的 HTTP Response #[derive(Clone)] pub struct ParsedResponse { pub version: u8, - pub status_code: u16, // 200, 404, etc. - pub reason: String, // OK, Not Found, etc. + pub status_code: u16, // 200, 404, etc. + pub reason: String, // OK, Not Found, etc. pub headers: HashMap, - pub body_offset: usize, // body 在 source_event.buf 中的起始位置 - pub body_len: usize, // body 长度 - pub source_event: Rc, // 原始 SslEvent (Rc 避免拷贝) + pub body_offset: usize, // body 在 source_event.buf 中的起始位置 + pub body_len: usize, // body 长度 + pub source_event: Rc, // 原始 SslEvent (Rc 避免拷贝) } impl ParsedResponse { @@ -28,9 +28,9 @@ impl ParsedResponse { pub fn body_str(&self) -> &str { std::str::from_utf8(self.body()).unwrap_or("") } - + /// 尝试将 body 解析为 JSON - /// + /// /// 如果 body 是有效的 UTF-8 且是有效的 JSON，返回解析后的 Value /// 否则返回 null pub fn json_body(&self) -> Option { @@ -54,35 +54,37 @@ impl ParsedResponse { impl TraceArgs for ParsedResponse { fn to_trace_args(&self) -> serde_json::Value { let mut args = serde_json::Map::new(); - + // Basic response info args.insert("status_code".to_string(), json!(self.status_code)); args.insert("reason".to_string(), json!(&self.reason)); - + // SSE indicator if self.is_sse() { args.insert("is_sse".to_string(), json!(true)); } - + // Add body info if present (and not SSE) if !self.is_sse() && self.body_len > 0 { args.insert("body_length".to_string(), json!(self.body_len)); - + // Add body preview (truncated) let body = self.body(); let body_preview = if body.len() > 500 { - format!("{}... ({} bytes total)", - String::from_utf8_lossy(&body[..500]), - body.len()) + format!( + "{}... ({} bytes total)", + String::from_utf8_lossy(&body[..500]), + body.len() + ) } else { String::from_utf8_lossy(body).to_string() }; - + if !body_preview.is_empty() { args.insert("body_preview".to_string(), json!(body_preview)); } } - + serde_json::Value::Object(args) } } @@ -90,10 +92,10 @@ impl TraceArgs for ParsedResponse { impl ToChromeTraceEvent for ParsedResponse { fn to_chrome_trace_events(&self) -> Vec { let ts_us = ns_to_us(self.source_event.timestamp_ns); - + // Minimum duration: 10ms = 10,000 microseconds const MIN_DUR_US: u64 = 10_000; - + let event = ChromeTraceEvent::complete( format!("{} {}", self.status_code, self.reason), "http.response", @@ -103,7 +105,7 @@ impl ToChromeTraceEvent for ParsedResponse { MIN_DUR_US, ) .with_trace_args(self); - + vec![event] } } @@ -114,27 +116,27 @@ impl fmt::Debug for ParsedResponse { debug .field("status", &format!("{} {}", self.status_code, self.reason)) .field("version", &format!("HTTP/1.{}", self.version)); - + // Format headers debug.field("headers", &self.headers); - + // Add SSE indicator if self.is_sse() { debug.field("is_sse", &true); } - + // Format body with smart detection let body = self.body(); if !body.is_empty() { debug.field("body", &format_body(body)); } - + // Add metadata from source_event debug .field("pid", &self.source_event.pid) .field("tid", &self.source_event.tid) .field("timestamp_ns", &self.source_event.timestamp_ns); - + debug.finish() } } @@ -151,6 +153,10 @@ fn format_body(data: &[u8]) -> String { format!("(text, {} bytes)\n{}", data.len(), text) } else { // Binary data - show as base64 - format!("(binary, {} bytes)\n{}", data.len(), base64::Engine::encode(&base64::engine::general_purpose::STANDARD, data)) + format!( + "(binary, {} bytes)\n{}", + data.len(), + base64::Engine::encode(&base64::engine::general_purpose::STANDARD, data) + ) } } diff --git a/src/agentsight/src/parser/http2/frame.rs b/src/agentsight/src/parser/http2/frame.rs index 397e73dfb..ce94c997d 100644 --- a/src/agentsight/src/parser/http2/frame.rs +++ b/src/agentsight/src/parser/http2/frame.rs @@ -3,12 +3,12 @@ //! Defines `Http2FrameType` and `ParsedHttp2Frame` for zero-copy //! HTTP/2 binary frame representation. -use std::fmt; -use std::rc::Rc; +use crate::chrome_trace::{ChromeTraceEvent, ToChromeTraceEvent, TraceArgs, ns_to_us}; use crate::probes::sslsniff::SslEvent; -use crate::chrome_trace::{TraceArgs, ToChromeTraceEvent, ChromeTraceEvent, ns_to_us}; -use serde_json::json; use hpack::Decoder; +use serde_json::json; +use std::fmt; +use std::rc::Rc; /// HTTP/2 frame type (RFC 7540 Section 6) #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -107,14 +107,18 @@ impl ParsedHttp2Frame { /// Check END_STREAM flag (0x01) on DATA/HEADERS frames pub fn has_end_stream(&self) -> bool { - matches!(self.frame_type, Http2FrameType::Data | Http2FrameType::Headers) - && (self.flags & 0x01) != 0 + matches!( + self.frame_type, + Http2FrameType::Data | Http2FrameType::Headers + ) && (self.flags & 0x01) != 0 } /// Check END_HEADERS flag (0x04) on HEADERS/CONTINUATION frames pub fn has_end_headers(&self) -> bool { - matches!(self.frame_type, Http2FrameType::Headers | Http2FrameType::Continuation) - && (self.flags & 0x04) != 0 + matches!( + self.frame_type, + Http2FrameType::Headers | Http2FrameType::Continuation + ) && (self.flags & 0x04) != 0 } /// Human-readable frame type name @@ -127,24 +131,44 @@ impl ParsedHttp2Frame { let mut flags = Vec::new(); match self.frame_type { Http2FrameType::Data => { - if self.flags & 0x01 != 0 { flags.push("END_STREAM"); } - if self.flags & 0x08 != 0 { flags.push("PADDED"); } + if self.flags & 0x01 != 0 { + flags.push("END_STREAM"); + } + if self.flags & 0x08 != 0 { + flags.push("PADDED"); + } } Http2FrameType::Headers => { - if self.flags & 0x01 != 0 { flags.push("END_STREAM"); } - if self.flags & 0x04 != 0 { flags.push("END_HEADERS"); } - if self.flags & 0x08 != 0 { flags.push("PADDED"); } - if self.flags & 0x20 != 0 { flags.push("PRIORITY"); } + if self.flags & 0x01 != 0 { + flags.push("END_STREAM"); + } + if self.flags & 0x04 != 0 { + flags.push("END_HEADERS"); + } + if self.flags & 0x08 != 0 { + flags.push("PADDED"); + } + if self.flags & 0x20 != 0 { + flags.push("PRIORITY"); + } } Http2FrameType::Settings | Http2FrameType::Ping => { - if self.flags & 0x01 != 0 { flags.push("ACK"); } + if self.flags & 0x01 != 0 { + flags.push("ACK"); + } } Http2FrameType::Continuation => { - if self.flags & 0x04 != 0 { flags.push("END_HEADERS"); } + if self.flags & 0x04 != 0 { + flags.push("END_HEADERS"); + } } Http2FrameType::PushPromise => { - if self.flags & 0x04 != 0 { flags.push("END_HEADERS"); } - if self.flags & 0x08 != 0 { flags.push("PADDED"); } + if self.flags & 0x04 != 0 { + flags.push("END_HEADERS"); + } + if self.flags & 0x08 != 0 { + flags.push("PADDED"); + } } _ => {} } @@ -156,13 +180,13 @@ impl ParsedHttp2Frame { } /// Decode HPACK-encoded headers using a provided decoder - /// + /// /// This method requires a stateful HPACK decoder because HTTP/2 header /// compression uses a dynamic table that persists across frames. - /// + /// /// # Arguments /// * `decoder` - A mutable reference to an HPACK decoder (maintains dynamic table state) - /// + /// /// # Returns /// * `Some(Vec<(String, String)>)` - Decoded header name-value pairs on success /// * `None` - If this is not a HEADERS frame or decoding failed @@ -170,12 +194,12 @@ impl ParsedHttp2Frame { if !self.is_headers() && self.frame_type != Http2FrameType::Continuation { return None; } - + let payload = self.payload(); if payload.is_empty() { return Some(Vec::new()); } - + decoder.decode(payload).ok().map(|headers| { headers .into_iter() @@ -189,29 +213,29 @@ impl ParsedHttp2Frame { } /// Decode headers using only the static HPACK table (stateless) - /// + /// /// This method does NOT maintain dynamic table state, so it can only /// decode headers that use static table indices. Useful for quick inspection /// without tracking connection state. - /// + /// /// # Returns /// Decoded header name-value pairs (only static table entries will be resolved) pub fn decode_headers_stateless(&self) -> Vec<(String, Option)> { if !self.is_headers() && self.frame_type != Http2FrameType::Continuation { return Vec::new(); } - + let payload = self.payload(); if payload.is_empty() { return Vec::new(); } - + let mut result = Vec::new(); let mut pos = 0; - + while pos < payload.len() { let first_byte = payload[pos]; - + // Indexed Header Field (1xxxxxxx) - fully indexed in static or dynamic table if first_byte & 0x80 != 0 { let index = (first_byte & 0x7F) as usize; @@ -243,20 +267,25 @@ impl ParsedHttp2Frame { pos += consumed; } } - + result } /// Decode a literal header field - fn decode_literal_header(&self, payload: &[u8], start: usize, _indexed: bool) -> (String, String, usize) { + fn decode_literal_header( + &self, + payload: &[u8], + start: usize, + _indexed: bool, + ) -> (String, String, usize) { let mut pos = start; let first_byte = payload[pos]; pos += 1; - + // Extract name (either from static table or literal) let name: String; let name_index = (first_byte & 0x3F) as usize; - + if name_index > 0 { // Name is in static table if let Some((n, _)) = Self::get_static_table_entry(name_index) { @@ -270,11 +299,11 @@ impl ParsedHttp2Frame { name = lit_name; pos += consumed; } - + // Decode value string let (value, consumed) = Self::decode_literal_string(payload, pos); pos += consumed; - + (name, value, pos - start) } @@ -283,15 +312,15 @@ impl ParsedHttp2Frame { if start >= payload.len() { return (String::new(), 0); } - + let mut pos = start; let first_byte = payload[pos]; let is_huffman = (first_byte & 0x80) != 0; - + // Decode length (variable length integer, 7 bits in first byte) let mut length = (first_byte & 0x7F) as usize; pos += 1; - + // Check if more length bytes follow (this is a simplification) // In full HPACK, length can be multi-byte if length == 0x7F && pos < payload.len() { @@ -306,20 +335,20 @@ impl ParsedHttp2Frame { } } } - + if pos + length > payload.len() { return (String::new(), pos - start); } - + let string_bytes = &payload[pos..pos + length]; - + let result = if is_huffman { // Huffman decode Self::huffman_decode(string_bytes) } else { String::from_utf8_lossy(string_bytes).to_string() }; - + (result, pos + length - start) } @@ -421,7 +450,11 @@ impl TraceArgs for ParsedHttp2Frame { } else { let preview = self.body_str(); if !preview.is_empty() { - let truncated = if preview.len() > 200 { &preview[..200] } else { preview }; + let truncated = if preview.len() > 200 { + &preview[..200] + } else { + preview + }; args.insert("body_preview".to_string(), json!(truncated)); } } @@ -438,7 +471,10 @@ impl TraceArgs for ParsedHttp2Frame { (name, json!(v)) }) .collect(); - args.insert("headers".to_string(), serde_json::Value::Object(headers_json)); + args.insert( + "headers".to_string(), + serde_json::Value::Object(headers_json), + ); } } @@ -478,11 +514,17 @@ impl fmt::Debug for ParsedHttp2Frame { let body = self.payload(); if let Ok(json) = serde_json::from_slice::(body) { let formatted = serde_json::to_string_pretty(&json).unwrap_or_default(); - debug.field("body", &format!("(json, {} bytes)\n{}", body.len(), formatted)); + debug.field( + "body", + &format!("(json, {} bytes)\n{}", body.len(), formatted), + ); } else if let Ok(text) = std::str::from_utf8(body) { let text = text.trim(); if text.len() > 200 { - debug.field("body", &format!("(text, {} bytes)\n{}...", body.len(), &text[..200])); + debug.field( + "body", + &format!("(text, {} bytes)\n{}...", body.len(), &text[..200]), + ); } else { debug.field("body", &format!("(text, {} bytes)\n{}", body.len(), text)); } @@ -497,11 +539,9 @@ impl fmt::Debug for ParsedHttp2Frame { if !headers.is_empty() { let header_strs: Vec = headers .into_iter() - .map(|(name, value)| { - match value { - Some(v) => format!(" {}: {}", name, v), - None => format!(" {}: ", name), - } + .map(|(name, value)| match value { + Some(v) => format!(" {}: {}", name, v), + None => format!(" {}: ", name), }) .collect(); debug.field("headers", &format!("\n{}", header_strs.join("\n"))); @@ -524,7 +564,9 @@ mod tests { #[test] fn test_huffman_decode_rfc7541_vectors() { // RFC 7541 C.4.1: "www.example.com" - let www = [0xf1, 0xe3, 0xc2, 0xe5, 0xf2, 0x3a, 0x6b, 0xa0, 0xab, 0x90, 0xf4, 0xff]; + let www = [ + 0xf1, 0xe3, 0xc2, 0xe5, 0xf2, 0x3a, 0x6b, 0xa0, 0xab, 0x90, 0xf4, 0xff, + ]; assert_eq!(ParsedHttp2Frame::huffman_decode(&www), "www.example.com"); // RFC 7541 C.4.2: "no-cache" diff --git a/src/agentsight/src/parser/http2/parser.rs b/src/agentsight/src/parser/http2/parser.rs index 32328be26..125fee5e4 100644 --- a/src/agentsight/src/parser/http2/parser.rs +++ b/src/agentsight/src/parser/http2/parser.rs @@ -83,7 +83,6 @@ impl Http2Parser { frames.push(frame); } - pos = payload_offset + length; } @@ -198,7 +197,9 @@ mod tests { fn test_parse_incomplete_frame() { // Valid header but truncated payload let mut raw = Vec::new(); - raw.push(0x00); raw.push(0x00); raw.push(0x20); // length = 32 + raw.push(0x00); + raw.push(0x00); + raw.push(0x20); // length = 32 raw.push(0x00); // DATA raw.push(0x00); // no flags raw.extend(&[0x00, 0x00, 0x00, 0x01]); // stream 1 @@ -220,7 +221,8 @@ mod tests { #[test] fn test_data_frame_json_body() { - let json_payload = br#"{"model":"qwen3.5-plus","messages":[{"role":"user","content":"hello"}]}"#; + let json_payload = + br#"{"model":"qwen3.5-plus","messages":[{"role":"user","content":"hello"}]}"#; let raw = build_frame(0, 0x01, 5, json_payload); let event = create_test_event(raw); let parser = Http2Parser::new(); @@ -257,18 +259,15 @@ mod tests { // Contains a HEADERS frame (len=83) and an incomplete DATA frame (len=16384, truncated) // Since we only keep DATA frames, and the DATA frame is truncated, result should be empty let sample_data: Vec = vec![ - 0, 0, 83, 1, 4, 0, 0, 0, 171, 203, 131, 4, 153, 96, 135, 166, - 177, 164, 209, 208, 85, 169, 60, 133, 99, 184, 88, 36, 227, 75, - 4, 61, 53, 208, 84, 152, 245, 35, 135, 202, 201, 200, 199, 198, - 197, 196, 195, 194, 31, 8, 158, 186, 81, 216, 91, 20, 71, 85, - 156, 11, 196, 1, 28, 117, 240, 180, 86, 138, 208, 227, 145, 151, - 218, 142, 87, 136, 65, 133, 185, 25, 143, 193, 192, 191, 190, 15, - 13, 132, 117, 166, 94, 111, 0, 64, 0, 0, 0, 0, 0, 0, 171, 123, - 34, 109, 111, 100, 101, 108, 34, 58, 34, 113, 119, 101, 110, 51, - 46, 53, 45, 112, 108, 117, 115, 34, 44, 34, 109, 101, 115, 115, - 97, 103, 101, 115, 34, 58, 91, 123, 34, 114, 111, 108, 101, 34, - 58, 34, 115, 121, 115, 116, 101, 109, 34, 44, 34, 99, 111, 110, - 116, 101, 110, 116, 34, 58, 34, 89, 111, 117, + 0, 0, 83, 1, 4, 0, 0, 0, 171, 203, 131, 4, 153, 96, 135, 166, 177, 164, 209, 208, 85, + 169, 60, 133, 99, 184, 88, 36, 227, 75, 4, 61, 53, 208, 84, 152, 245, 35, 135, 202, + 201, 200, 199, 198, 197, 196, 195, 194, 31, 8, 158, 186, 81, 216, 91, 20, 71, 85, 156, + 11, 196, 1, 28, 117, 240, 180, 86, 138, 208, 227, 145, 151, 218, 142, 87, 136, 65, 133, + 185, 25, 143, 193, 192, 191, 190, 15, 13, 132, 117, 166, 94, 111, 0, 64, 0, 0, 0, 0, 0, + 0, 171, 123, 34, 109, 111, 100, 101, 108, 34, 58, 34, 113, 119, 101, 110, 51, 46, 53, + 45, 112, 108, 117, 115, 34, 44, 34, 109, 101, 115, 115, 97, 103, 101, 115, 34, 58, 91, + 123, 34, 114, 111, 108, 101, 34, 58, 34, 115, 121, 115, 116, 101, 109, 34, 44, 34, 99, + 111, 110, 116, 101, 110, 116, 34, 58, 34, 89, 111, 117, ]; let event = create_test_event(sample_data); let parser = Http2Parser::new(); diff --git a/src/agentsight/src/parser/mod.rs b/src/agentsight/src/parser/mod.rs index baa31dbf8..694f5edea 100644 --- a/src/agentsight/src/parser/mod.rs +++ b/src/agentsight/src/parser/mod.rs @@ -32,13 +32,13 @@ pub mod http; pub mod http2; -pub mod sse; pub mod proctrace; mod result; +pub mod sse; mod unified; // Re-export result types -pub use result::{ParsedMessage, ParseResult}; +pub use result::{ParseResult, ParsedMessage}; // Re-export unified parser pub use unified::Parser; @@ -47,10 +47,10 @@ pub use unified::Parser; pub use http::{HttpParser, ParsedHttpMessage, ParsedRequest, ParsedResponse}; // Re-export SSE types -pub use sse::{SseParser, ParsedSseEvent}; +pub use sse::{ParsedSseEvent, SseParser}; // Re-export proctrace types -pub use proctrace::{ProcTraceParser, ParsedProcEvent, ProcEventType}; +pub use proctrace::{ParsedProcEvent, ProcEventType, ProcTraceParser}; // Re-export HTTP/2 types -pub use http2::{Http2Parser, Http2FrameType, ParsedHttp2Frame}; +pub use http2::{Http2FrameType, Http2Parser, ParsedHttp2Frame}; diff --git a/src/agentsight/src/parser/proctrace.rs b/src/agentsight/src/parser/proctrace.rs index be9837253..0abc841a1 100644 --- a/src/agentsight/src/parser/proctrace.rs +++ b/src/agentsight/src/parser/proctrace.rs @@ -47,20 +47,24 @@ impl ProcTraceParser { /// Parse a variable-length process event pub fn parse_variable(event: &VariableEvent) -> Option { match event { - VariableEvent::Exec { header, filename: _, args } => { - Some(ParsedProcEvent { - event_type: ProcEventType::Exec, - pid: header.pid, - tid: header.tid, - ppid: header.ppid, - ptid: header.ptid, - comm: event.comm_str(), - timestamp_ns: header.timestamp_ns, - args: Some(args.clone()).filter(|s| !s.is_empty()), - stdout_data: None, - }) - } - VariableEvent::Stdout { header, payload, .. } => { + VariableEvent::Exec { + header, + filename: _, + args, + } => Some(ParsedProcEvent { + event_type: ProcEventType::Exec, + pid: header.pid, + tid: header.tid, + ppid: header.ppid, + ptid: header.ptid, + comm: event.comm_str(), + timestamp_ns: header.timestamp_ns, + args: Some(args.clone()).filter(|s| !s.is_empty()), + stdout_data: None, + }), + VariableEvent::Stdout { + header, payload, .. + } => { let stdout_data = String::from_utf8(payload.clone()).ok(); Some(ParsedProcEvent { event_type: ProcEventType::Stdout, @@ -74,19 +78,17 @@ impl ProcTraceParser { stdout_data, }) } - VariableEvent::Exit { header, .. } => { - Some(ParsedProcEvent { - event_type: ProcEventType::Exit, - pid: header.pid, - tid: header.tid, - ppid: header.ppid, - ptid: header.ptid, - comm: event.comm_str(), - timestamp_ns: header.timestamp_ns, - args: None, - stdout_data: None, - }) - } + VariableEvent::Exit { header, .. } => Some(ParsedProcEvent { + event_type: ProcEventType::Exit, + pid: header.pid, + tid: header.tid, + ppid: header.ppid, + ptid: header.ptid, + comm: event.comm_str(), + timestamp_ns: header.timestamp_ns, + args: None, + stdout_data: None, + }), VariableEvent::Unknown(_) => None, } } @@ -145,23 +147,21 @@ impl ProcTraceParser { bp: None, }) } - ProcEventType::Exit => { - Some(ChromeTraceEvent { - name: format!("exit: {}", parsed.comm), - cat: "process.exit".to_string(), - ph: "i".to_string(), - ts: ts_us, - dur: None, - pid: parsed.pid, - tid: parsed.tid as u64, - args: Some(json!({ - "pid": parsed.pid, - "comm": parsed.comm, - })), - id: None, - bp: None, - }) - } + ProcEventType::Exit => Some(ChromeTraceEvent { + name: format!("exit: {}", parsed.comm), + cat: "process.exit".to_string(), + ph: "i".to_string(), + ts: ts_us, + dur: None, + pid: parsed.pid, + tid: parsed.tid as u64, + args: Some(json!({ + "pid": parsed.pid, + "comm": parsed.comm, + })), + id: None, + bp: None, + }), } } @@ -172,18 +172,21 @@ impl ProcTraceParser { /// Convert multiple variable-length events to Chrome Trace Events pub fn to_chrome_trace_events(events: &[VariableEvent]) -> Vec { - events.iter().filter_map(Self::to_chrome_trace_event).collect() + events + .iter() + .filter_map(Self::to_chrome_trace_event) + .collect() } } impl TraceArgs for ParsedProcEvent { fn to_trace_args(&self) -> serde_json::Value { let mut args = serde_json::Map::new(); - + // Common fields args.insert("pid".to_string(), json!(self.pid)); args.insert("comm".to_string(), json!(&self.comm)); - + // Event type specific fields match self.event_type { ProcEventType::Exec => { @@ -196,7 +199,7 @@ impl TraceArgs for ParsedProcEvent { ProcEventType::Stdout => { if let Some(ref data) = self.stdout_data { args.insert("len".to_string(), json!(data.len())); - + // Add data preview (truncated) let preview = if data.len() > 200 { format!("{}... ({} bytes total)", &data[..200], data.len()) @@ -210,7 +213,7 @@ impl TraceArgs for ParsedProcEvent { // Exit event has minimal args } } - + serde_json::Value::Object(args) } } @@ -232,7 +235,7 @@ impl ParsedProcEvent { } ProcEventType::Exit => format!("exit: {}", self.comm), }; - + let cat = match self.event_type { ProcEventType::Exec => "process.exec", ProcEventType::Stdout => "process.stdout", diff --git a/src/agentsight/src/parser/result.rs b/src/agentsight/src/parser/result.rs index 18b068b8a..01d27462e 100644 --- a/src/agentsight/src/parser/result.rs +++ b/src/agentsight/src/parser/result.rs @@ -3,12 +3,12 @@ //! This module defines the `ParsedMessage` and `ParseResult` types //! representing the output from parsing events. -use std::rc::Rc; use crate::parser::http::{ParsedRequest, ParsedResponse}; -use crate::parser::sse::ParsedSseEvent; -use crate::parser::proctrace::ParsedProcEvent; use crate::parser::http2::ParsedHttp2Frame; +use crate::parser::proctrace::ParsedProcEvent; +use crate::parser::sse::ParsedSseEvent; use crate::probes::sslsniff::SslEvent; +use std::rc::Rc; /// Parsed message from events #[derive(Debug, Clone)] diff --git a/src/agentsight/src/parser/sse/event.rs b/src/agentsight/src/parser/sse/event.rs index dd4e8865a..95cdd6cf8 100644 --- a/src/agentsight/src/parser/sse/event.rs +++ b/src/agentsight/src/parser/sse/event.rs @@ -1,8 +1,8 @@ +use crate::chrome_trace::{ChromeTraceEvent, ns_to_us}; +use crate::probes::sslsniff::SslEvent; use serde::{Deserialize, Serialize}; use std::fmt; use std::rc::Rc; -use crate::chrome_trace::{ChromeTraceEvent, ns_to_us}; -use crate::probes::sslsniff::SslEvent; /// SSE Event - Standard Server-Sent Events message (legacy version with String data) /// Follows the W3C EventSource specification: https://html.spec.whatwg.org/multipage/server-sent-events.html @@ -114,12 +114,11 @@ impl ParsedSseEvent { return true; } // Anthropic style: data contains {"type":"message_stop"} - if trimmed.starts_with('{') { - if let Ok(v) = serde_json::from_str::(trimmed) { - if v.get("type").and_then(|t| t.as_str()) == Some("message_stop") { - return true; - } - } + if trimmed.starts_with('{') + && let Ok(v) = serde_json::from_str::(trimmed) + && v.get("type").and_then(|t| t.as_str()) == Some("message_stop") + { + return true; } false } @@ -138,7 +137,7 @@ impl ParsedSseEvent { impl fmt::Debug for ParsedSseEvent { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { let mut debug = f.debug_struct("ParsedSseEvent"); - + if let Some(ref id) = self.id { debug.field("id", id); } @@ -148,24 +147,24 @@ impl fmt::Debug for ParsedSseEvent { if let Some(retry) = self.retry { debug.field("retry", &retry); } - + // Check if this is a done marker if self.is_done() { debug.field("done", &true); } - + // Format data with smart detection let data = self.data(); if !data.is_empty() { debug.field("data", &format_sse_data(data)); } - + // Add metadata debug .field("data_len", &self.data_len) .field("pid", &self.source_event.pid) .field("timestamp_ns", &self.source_event.timestamp_ns); - + debug.finish() } } @@ -182,7 +181,11 @@ fn format_sse_data(data: &[u8]) -> String { format!("(text, {} bytes)\n{}", data.len(), text) } else { // Binary data - show as base64 - format!("(binary, {} bytes)\n{}", data.len(), base64::Engine::encode(&base64::engine::general_purpose::STANDARD, data)) + format!( + "(binary, {} bytes)\n{}", + data.len(), + base64::Engine::encode(&base64::engine::general_purpose::STANDARD, data) + ) } } @@ -197,6 +200,12 @@ pub struct SSEEvents { pub consumed_bytes: usize, } +impl Default for SSEEvents { + fn default() -> Self { + Self::new() + } +} + impl SSEEvents { /// Create a new empty SSEEvents container pub fn new() -> Self { @@ -256,20 +265,31 @@ impl SSEEvents { // Build args with aggregated information let mut args = serde_json::Map::new(); - args.insert("event_count".to_string(), serde_json::json!(self.events.len())); - args.insert("consumed_bytes".to_string(), serde_json::json!(self.consumed_bytes)); - args.insert("remaining_bytes".to_string(), serde_json::json!(self.remaining.len())); + args.insert( + "event_count".to_string(), + serde_json::json!(self.events.len()), + ); + args.insert( + "consumed_bytes".to_string(), + serde_json::json!(self.consumed_bytes), + ); + args.insert( + "remaining_bytes".to_string(), + serde_json::json!(self.remaining.len()), + ); // Aggregate data from all events let total_data_size: usize = self.events.iter().map(|e| e.data.len()).sum(); - args.insert("total_data_size".to_string(), serde_json::json!(total_data_size)); + args.insert( + "total_data_size".to_string(), + serde_json::json!(total_data_size), + ); // Combine all events' data (no truncation, no limit) - let all_data: Vec = self.events + let all_data: Vec = self + .events .iter() - .map(|e| { - format!("[{}] {}", e.event.as_deref().unwrap_or("message"), e.data) - }) + .map(|e| format!("[{}] {}", e.event.as_deref().unwrap_or("message"), e.data)) .collect(); if !all_data.is_empty() { @@ -277,7 +297,8 @@ impl SSEEvents { } // Collect all event types - let event_types: Vec<&str> = self.events + let event_types: Vec<&str> = self + .events .iter() .filter_map(|e| e.event.as_deref()) .collect(); @@ -316,10 +337,7 @@ impl SSEEvent { /// Check if this is a "ping" or keepalive event (data is empty and no other fields) pub fn is_keepalive(&self) -> bool { - self.data.is_empty() - && self.id.is_none() - && self.event.is_none() - && self.retry.is_none() + self.data.is_empty() && self.id.is_none() && self.event.is_none() && self.retry.is_none() } /// Format as SSE protocol string @@ -354,12 +372,7 @@ impl SSEEvent { /// /// # Returns /// A ChromeTraceEvent suitable for visualization in Perfetto - pub fn to_chrome_trace_event( - &self, - pid: u32, - tid: u64, - timestamp_ns: u64, - ) -> ChromeTraceEvent { + pub fn to_chrome_trace_event(&self, pid: u32, tid: u64, timestamp_ns: u64) -> ChromeTraceEvent { // Build event name based on event type or data preview let name = match &self.event { Some(event_type) => format!("SSE {}", event_type), @@ -376,7 +389,10 @@ impl SSEEvent { self.data.clone() }; args.insert("data".to_string(), serde_json::json!(data_preview)); - args.insert("data_length".to_string(), serde_json::json!(self.data.len())); + args.insert( + "data_length".to_string(), + serde_json::json!(self.data.len()), + ); if let Some(id) = &self.id { args.insert("id".to_string(), serde_json::json!(id)); @@ -412,10 +428,18 @@ mod tests { fn make_event(data: &[u8]) -> Rc { Rc::new(SslEvent { - source: 0, timestamp_ns: 5000, delta_ns: 0, - pid: 1, tid: 1, uid: 0, len: data.len() as u32, - rw: 1, comm: "test".to_string(), - buf: data.to_vec(), is_handshake: false, ssl_ptr: 0x1, + source: 0, + timestamp_ns: 5000, + delta_ns: 0, + pid: 1, + tid: 1, + uid: 0, + len: data.len() as u32, + rw: 1, + comm: "test".to_string(), + buf: data.to_vec(), + is_handshake: false, + ssl_ptr: 0x1, }) } @@ -469,7 +493,7 @@ mod tests { let ev = make_event(data); let parsed = ParsedSseEvent::new( None, - Some("message_stop".to_string()), // event field + Some("message_stop".to_string()), // event field None, 0, data.len(), @@ -482,14 +506,7 @@ mod tests { fn test_is_done_anthropic_event_field_only() { // Even with empty data, event=message_stop should trigger done let ev = make_event(b""); - let parsed = ParsedSseEvent::new( - None, - Some("message_stop".to_string()), - None, - 0, - 0, - ev, - ); + let parsed = ParsedSseEvent::new(None, Some("message_stop".to_string()), None, 0, 0, ev); assert!(parsed.is_done()); } @@ -535,7 +552,12 @@ mod tests { #[test] fn test_sse_event_is_keepalive() { - let e = SSEEvent { id: None, event: None, data: String::new(), retry: None }; + let e = SSEEvent { + id: None, + event: None, + data: String::new(), + retry: None, + }; assert!(e.is_keepalive()); let e2 = SSEEvent::new("data"); @@ -601,8 +623,10 @@ mod tests { container.events.push(SSEEvent::new("data1")); container.events.push(SSEEvent { - id: None, event: Some("delta".to_string()), - data: "data2".to_string(), retry: None, + id: None, + event: Some("delta".to_string()), + data: "data2".to_string(), + retry: None, }); container.consumed_bytes = 100; @@ -641,7 +665,9 @@ mod tests { Some("id1".to_string()), Some("message".to_string()), Some(3000), - 0, data.len(), ev, + 0, + data.len(), + ev, ); let debug = format!("{:?}", parsed); assert!(debug.contains("id1")); diff --git a/src/agentsight/src/parser/sse/parser.rs b/src/agentsight/src/parser/sse/parser.rs index cbb3900bf..e3192bb84 100644 --- a/src/agentsight/src/parser/sse/parser.rs +++ b/src/agentsight/src/parser/sse/parser.rs @@ -1,5 +1,5 @@ -use crate::probes::sslsniff::SslEvent; use super::event::{ParsedSseEvent, SSEEvent, SSEEvents}; +use crate::probes::sslsniff::SslEvent; use std::rc::Rc; /// SSE Parser - parses SSE stream data into events (legacy version) @@ -12,10 +12,10 @@ impl SSEParser { let mut result = SSEEvents::new(); let mut current_event = SSEEvent::new(""); let mut data_lines: Vec = Vec::new(); - let mut lines = buffer.lines().peekable(); + let lines = buffer.lines().peekable(); let mut consumed_len = 0; - while let Some(line) = lines.next() { + for line in lines { consumed_len += line.len() + 1; // +1 for newline if line.is_empty() { @@ -80,7 +80,7 @@ impl SseParser { /// Parse SslEvent and extract SSE events /// Returns Vec of ParsedSseEvent - /// + /// /// Note: For multi-line data fields, data is concatenated with '\n' separators. /// The data_offset points to the first data line, data_len covers all data content /// including internal newlines. @@ -100,7 +100,7 @@ impl SseParser { // Use split_inclusive to properly track byte offsets with \r\n line endings let lines_iter = text.split_inclusive('\n'); - + for line_with_end in lines_iter { // Remove trailing \r\n or \n for parsing, but keep track of original length let line = line_with_end.trim_end_matches('\n').trim_end_matches('\r'); @@ -124,7 +124,11 @@ impl SseParser { let first_line_byte_len = first_line_bytes.len(); // Account for \r if present in original data let has_crlf = data_parts[0].ends_with('\r'); - let adjusted_len = if has_crlf { first_line_byte_len - 1 } else { first_line_byte_len }; + let adjusted_len = if has_crlf { + first_line_byte_len - 1 + } else { + first_line_byte_len + }; (data_start.unwrap_or(0), adjusted_len) } else { (0, 0) @@ -154,17 +158,18 @@ impl SseParser { // "If the line contains a U+003A COLON character (':'), collect the characters // on the line before the first U+003A COLON character (':'), and let field be that string. // Collect the characters on the line after the first U+003A COLON character (':'), - // and let value be that string. If value starts with a U+0020 SPACE character, + // and let value be that string. If value starts with a U+0020 SPACE character, // remove it from value." let has_space_after_colon = value.starts_with(' '); let value_stripped = value.strip_prefix(' ').unwrap_or(value); - + // Calculate value_start in the original buffer // line_start: start of line in buffer // field.len() + 1: skip field and colon // +1 if there was a space after colon - let value_start = line_start + field.len() + 1 + if has_space_after_colon { 1 } else { 0 }; - + let value_start = + line_start + field.len() + 1 + if has_space_after_colon { 1 } else { 0 }; + match field { "id" => current_id = Some(value_stripped.to_string()), "event" => current_event = Some(value_stripped.to_string()), @@ -256,7 +261,7 @@ mod tests { let events = parser.parse(event); assert_eq!(events.len(), 1); - + let evt = &events[0]; assert_eq!(evt.id, None); assert_eq!(evt.event, None); @@ -272,7 +277,7 @@ mod tests { let events = parser.parse(event); assert_eq!(events.len(), 1); - + let evt = &events[0]; assert_eq!(evt.id, Some("123".to_string())); assert_eq!(evt.event, Some("message".to_string())); @@ -287,7 +292,7 @@ mod tests { let events = parser.parse(event); assert_eq!(events.len(), 2); - + assert_eq!(events[0].data(), b"first"); assert_eq!(events[1].data(), b"second"); } @@ -300,7 +305,7 @@ mod tests { let events = parser.parse(event); assert_eq!(events.len(), 1); - + let evt = &events[0]; assert!(evt.is_done()); } @@ -313,7 +318,7 @@ mod tests { let events = parser.parse(event); assert_eq!(events.len(), 1); - + let evt = &events[0]; // Multi-line data: offset points to first line "line1" // data_len is the length of first line only (5) @@ -330,7 +335,7 @@ mod tests { let events = parser.parse(event); assert_eq!(events.len(), 1); - + let evt = &events[0]; assert_eq!(evt.data(), b"hello"); } @@ -428,7 +433,8 @@ mod tests { #[test] fn test_legacy_parse_stream_with_id_event_retry() { - let result = SSEParser::parse_stream("id: 42\nevent: update\nretry: 1000\ndata: payload\n\n"); + let result = + SSEParser::parse_stream("id: 42\nevent: update\nretry: 1000\ndata: payload\n\n"); assert_eq!(result.len(), 1); let evt = &result.events[0]; assert_eq!(evt.id, Some("42".to_string())); diff --git a/src/agentsight/src/parser/unified.rs b/src/agentsight/src/parser/unified.rs index 70e62ff8b..e060818e0 100644 --- a/src/agentsight/src/parser/unified.rs +++ b/src/agentsight/src/parser/unified.rs @@ -12,7 +12,7 @@ use crate::event::Event; use crate::parser::http::{HttpParser, ParsedHttpMessage}; use crate::parser::http2::Http2Parser; use crate::parser::proctrace::ProcTraceParser; -use crate::parser::sse::{SseParser, ParsedSseEvent}; +use crate::parser::sse::{ParsedSseEvent, SseParser}; use crate::probes::proctrace::VariableEvent; use crate::probes::sslsniff::SslEvent; use std::rc::Rc; @@ -88,9 +88,9 @@ impl Parser { let buf = &ssl_event.buf[..buf_size]; if buf == b"0\r\n\r\n" { return ParseResult { - messages: vec![ParsedMessage::SseEvent( - ParsedSseEvent::new_done_marker(Rc::clone(&ssl_event)) - )], + messages: vec![ParsedMessage::SseEvent(ParsedSseEvent::new_done_marker( + Rc::clone(&ssl_event), + ))], }; } } @@ -137,10 +137,18 @@ impl Parser { match event { Event::Ssl(ssl_event) => self.parse_ssl_event(Rc::new(ssl_event)), Event::Proc(proc_event) => self.parse_proc_event(&proc_event), - Event::ProcMon(_) => ParseResult { messages: Vec::new() }, - Event::FileWatch(_) => ParseResult { messages: Vec::new() }, - Event::FileWrite(_) => ParseResult { messages: Vec::new() }, - Event::UdpDns(_) => ParseResult { messages: Vec::new() }, + Event::ProcMon(_) => ParseResult { + messages: Vec::new(), + }, + Event::FileWatch(_) => ParseResult { + messages: Vec::new(), + }, + Event::FileWrite(_) => ParseResult { + messages: Vec::new(), + }, + Event::UdpDns(_) => ParseResult { + messages: Vec::new(), + }, } } diff --git a/src/agentsight/src/probes/filewatch.rs b/src/agentsight/src/probes/filewatch.rs index acedb7971..7da4971ec 100644 --- a/src/agentsight/src/probes/filewatch.rs +++ b/src/agentsight/src/probes/filewatch.rs @@ -9,13 +9,15 @@ use libbpf_rs::{ Link, MapHandle, skel::{OpenSkel, SkelBuilder}, }; -use std::{ - mem::MaybeUninit, - os::fd::AsFd, -}; +use std::{mem::MaybeUninit, os::fd::AsFd}; // ─── Generated skeleton ─────────────────────────────────────────────────────── -#[allow(non_camel_case_types, non_upper_case_globals, dead_code, non_snake_case)] +#[allow( + non_camel_case_types, + non_upper_case_globals, + dead_code, + non_snake_case +)] mod bpf { include!(concat!(env!("OUT_DIR"), "/filewatch.skel.rs")); include!(concat!(env!("OUT_DIR"), "/filewatch.rs")); @@ -49,7 +51,8 @@ impl FileWatchEvent { let raw = unsafe { &*(data.as_ptr() as *const RawFileWatchEvent) }; // Parse comm (null-terminated) - let comm = raw.comm + let comm = raw + .comm .iter() .take_while(|&&c| c != 0) .map(|&c| c as u8) @@ -57,7 +60,8 @@ impl FileWatchEvent { let comm = String::from_utf8_lossy(&comm).into_owned(); // Parse filename (null-terminated) - let filename = raw.filename + let filename = raw + .filename .iter() .take_while(|&&c| c != 0) .map(|&c| c as u8) @@ -94,7 +98,9 @@ impl FileWatch { builder.obj_builder.debug(config::verbose()); let open_object = Box::new(MaybeUninit::::uninit()); - let mut open_skel = builder.open().context("failed to open filewatch BPF object")?; + let mut open_skel = builder + .open() + .context("failed to open filewatch BPF object")?; // Reuse external traced_processes map open_skel @@ -110,9 +116,12 @@ impl FileWatch { .reuse_fd(rb.as_fd()) .context("failed to reuse external rb map for filewatch")?; - let skel = open_skel.load().context("failed to load filewatch BPF object")?; + let skel = open_skel + .load() + .context("failed to load filewatch BPF object")?; // SAFETY: skel borrows open_object which lives in a Box + #[allow(clippy::unnecessary_cast)] let skel = unsafe { Box::from_raw(Box::into_raw(Box::new(skel)) as *mut FilewatchSkel<'static>) }; diff --git a/src/agentsight/src/probes/filewrite.rs b/src/agentsight/src/probes/filewrite.rs index e1dfefed1..56257485d 100644 --- a/src/agentsight/src/probes/filewrite.rs +++ b/src/agentsight/src/probes/filewrite.rs @@ -9,13 +9,15 @@ use libbpf_rs::{ Link, MapHandle, skel::{OpenSkel, SkelBuilder}, }; -use std::{ - mem::MaybeUninit, - os::fd::AsFd, -}; +use std::{mem::MaybeUninit, os::fd::AsFd}; // ─── Generated skeleton ─────────────────────────────────────────────────────── -#[allow(non_camel_case_types, non_upper_case_globals, dead_code, non_snake_case)] +#[allow( + non_camel_case_types, + non_upper_case_globals, + dead_code, + non_snake_case +)] mod bpf { include!(concat!(env!("OUT_DIR"), "/filewrite.skel.rs")); include!(concat!(env!("OUT_DIR"), "/filewrite.rs")); @@ -50,7 +52,8 @@ impl FileWriteEvent { let raw = unsafe { &*(data.as_ptr() as *const RawFileWriteEvent) }; // Parse comm (null-terminated) - let comm = raw.comm + let comm = raw + .comm .iter() .take_while(|&&c| c != 0) .map(|&c| c as u8) @@ -58,7 +61,8 @@ impl FileWriteEvent { let comm = String::from_utf8_lossy(&comm).into_owned(); // Parse filename (null-terminated) - let filename = raw.filename + let filename = raw + .filename .iter() .take_while(|&&c| c != 0) .map(|&c| c as u8) @@ -101,7 +105,9 @@ impl FileWrite { builder.obj_builder.debug(config::verbose()); let open_object = Box::new(MaybeUninit::::uninit()); - let mut open_skel = builder.open().context("failed to open filewrite BPF object")?; + let mut open_skel = builder + .open() + .context("failed to open filewrite BPF object")?; // Reuse external traced_processes map open_skel @@ -117,9 +123,12 @@ impl FileWrite { .reuse_fd(rb.as_fd()) .context("failed to reuse external rb map for filewrite")?; - let skel = open_skel.load().context("failed to load filewrite BPF object")?; + let skel = open_skel + .load() + .context("failed to load filewrite BPF object")?; // SAFETY: skel borrows open_object which lives in a Box + #[allow(clippy::unnecessary_cast)] let skel = unsafe { Box::from_raw(Box::into_raw(Box::new(skel)) as *mut FilewriteSkel<'static>) }; diff --git a/src/agentsight/src/probes/mod.rs b/src/agentsight/src/probes/mod.rs index ebd022ea1..230fccd56 100644 --- a/src/agentsight/src/probes/mod.rs +++ b/src/agentsight/src/probes/mod.rs @@ -1,20 +1,19 @@ - - -pub mod sslsniff; -pub mod proctrace; -pub mod procmon; pub mod filewatch; pub mod filewrite; -pub mod udpdns; -pub mod tcpsniff; +#[allow(clippy::module_inception)] pub mod probes; +pub mod procmon; +pub mod proctrace; +pub mod sslsniff; +pub mod tcpsniff; +pub mod udpdns; // Re-export commonly used types -pub use probes::{Probes, ProbesPoller}; -pub use proctrace::{ProcTrace, ProcPoller, VariableEvent as ProcEvent}; -pub use sslsniff::{SslSniff, SslPoller, SslEvent}; -pub use procmon::{ProcMon, ProcMonEvent, Event as ProcMonEventExt}; pub use filewatch::{FileWatch, FileWatchEvent}; pub use filewrite::{FileWrite as FileWriteProbe, FileWriteEvent}; +pub use probes::{Probes, ProbesPoller}; +pub use procmon::{Event as ProcMonEventExt, ProcMon, ProcMonEvent}; +pub use proctrace::{ProcPoller, ProcTrace, VariableEvent as ProcEvent}; +pub use sslsniff::{SslEvent, SslPoller, SslSniff}; +pub use tcpsniff::TcpSniff; pub use udpdns::{UdpDns, UdpDnsEvent}; -pub use tcpsniff::TcpSniff; \ No newline at end of file diff --git a/src/agentsight/src/probes/probes.rs b/src/agentsight/src/probes/probes.rs index 626e1e2c0..ba2d8d7ca 100644 --- a/src/agentsight/src/probes/probes.rs +++ b/src/agentsight/src/probes/probes.rs @@ -8,22 +8,25 @@ use anyhow::{Context, Result}; use libbpf_rs::{MapHandle, RingBufferBuilder}; use std::{ mem, - sync::{Arc, atomic::{AtomicBool, Ordering}}, + sync::{ + Arc, + atomic::{AtomicBool, Ordering}, + }, thread, time::Duration, }; use crate::event::Event; -use super::proctrace::{ProcTrace, VariableEvent, ProcEventHeader}; -use super::sslsniff::SslSniff; -use super::sslsniff::bpf::probe_SSL_data_t as RawSslEvent; -use super::procmon::{ProcMon, ProcMonEvent}; use super::filewatch::{FileWatch, RawFileWatchEvent}; use super::filewrite::{FileWrite as FileWriteProbe, RawFileWriteEvent}; -use super::udpdns::{UdpDns, RawUdpDnsEvent}; -use crate::config::TcpTarget; +use super::procmon::{ProcMon, ProcMonEvent}; +use super::proctrace::{ProcEventHeader, ProcTrace, VariableEvent}; +use super::sslsniff::SslSniff; +use super::sslsniff::bpf::probe_SSL_data_t as RawSslEvent; use super::tcpsniff::TcpSniff; +use super::udpdns::{RawUdpDnsEvent, UdpDns}; +use crate::config::TcpTarget; const POLL_TIMEOUT_MS: u64 = 100; @@ -36,11 +39,12 @@ const EVENT_SOURCE_FILEWRITE: u32 = 5; const EVENT_SOURCE_UDPDNS: u32 = 6; /// Unified probe manager that coordinates sslsniff and proctrace -/// +/// /// This manager ensures both probes share the same traced_processes map /// and the same ring buffer, allowing coordinated process tracing where: /// - proctrace captures process creation events /// - sslsniff captures SSL traffic from those processes +/// /// Both write to a single shared ring buffer to save memory. pub struct Probes { /// Process trace probe (owns the traced_processes map and ring buffer) @@ -66,28 +70,33 @@ pub struct Probes { impl Probes { /// Create a new unified probe manager - /// + /// /// # Arguments /// * `target_pids` - Initial PIDs to trace (empty means trace all matching UID) /// * `target_uid` - Optional UID filter - pub fn new(target_pids: &[u32], target_uid: Option, enable_filewatch: bool, enable_udpdns: bool, tcp_targets: &[TcpTarget]) -> Result { + pub fn new( + target_pids: &[u32], + target_uid: Option, + enable_filewatch: bool, + enable_udpdns: bool, + tcp_targets: &[TcpTarget], + ) -> Result { // Create proctrace first - it will own the traced_processes map and ring buffer let proctrace = ProcTrace::new_with_target(target_pids, target_uid) .context("failed to create proctrace")?; - + // Get handles to the shared maps for reuse - let map_handle = proctrace.traced_processes_handle() + let map_handle = proctrace + .traced_processes_handle() .context("failed to get traced_processes handle")?; - let rb_handle = proctrace.rb_handle() - .context("failed to get rb handle")?; - + let rb_handle = proctrace.rb_handle().context("failed to get rb handle")?; + // Create sslsniff - it will reuse both the traced_processes map and ring buffer let sslsniff = SslSniff::new_with_traced_processes(Some(&map_handle), Some(&rb_handle)) .context("failed to create sslsniff")?; // Create procmon - it reuses the ring buffer - let procmon = ProcMon::new_with_rb(&rb_handle) - .context("failed to create procmon")?; + let procmon = ProcMon::new_with_rb(&rb_handle).context("failed to create procmon")?; // Optionally create filewatch - it reuses both the traced_processes map and ring buffer let filewatch = if enable_filewatch { @@ -116,8 +125,8 @@ impl Probes { // Optionally create tcpsniff - captures plain HTTP traffic to configured IP/port targets let tcpsniff = if !tcp_targets.is_empty() { - let mut tcp = TcpSniff::new_with_maps(&rb_handle) - .context("failed to create tcpsniff")?; + let mut tcp = + TcpSniff::new_with_maps(&rb_handle).context("failed to create tcpsniff")?; tcp.set_targets(tcp_targets) .context("failed to set tcp targets")?; Some(tcp) @@ -127,7 +136,7 @@ impl Probes { }; let (event_tx, event_rx) = crossbeam_channel::unbounded(); - + Ok(Self { proctrace, sslsniff, @@ -145,26 +154,25 @@ impl Probes { /// Attach all probes pub fn attach(&mut self) -> Result<()> { // Attach procmon for process monitoring - self.procmon.attach() - .context("failed to attach procmon")?; - self.proctrace.attach().context("failed to attach proctrace")?; + self.procmon.attach().context("failed to attach procmon")?; + self.proctrace + .attach() + .context("failed to attach proctrace")?; // Attach filewatch for .jsonl file monitoring (if enabled) if let Some(ref mut fw) = self.filewatch { - fw.attach() - .context("failed to attach filewatch")?; + fw.attach().context("failed to attach filewatch")?; } // Attach filewrite for JSON write monitoring (always enabled) - self.filewrite.attach() + self.filewrite + .attach() .context("failed to attach filewrite")?; // Attach udpdns for DNS query capture (if enabled) if let Some(ref mut dns) = self.udpdns { - dns.attach() - .context("failed to attach udpdns")?; + dns.attach().context("failed to attach udpdns")?; } // Attach tcpsniff for plain HTTP traffic capture (if enabled) if let Some(ref mut tcp) = self.tcpsniff { - tcp.attach() - .context("failed to attach tcpsniff")?; + tcp.attach().context("failed to attach tcpsniff")?; } // sslsniff uses uprobes attached per-process via attach_process() Ok(()) @@ -177,7 +185,8 @@ impl Probes { /// Attach SSL probes to a specific process pub fn attach_ssl_to_process(&mut self, pid: i32) -> Result<()> { - self.sslsniff.attach_process(pid) + self.sslsniff + .attach_process(pid) .context("failed to attach sslsniff to process")?; Ok(()) } @@ -265,7 +274,7 @@ impl Probes { None } }; - + if let Some(e) = event { let _ = event_tx.send(e); } @@ -312,13 +321,15 @@ impl Probes { /// Add a PID to the traced_processes map at runtime pub fn add_traced_pid(&mut self, pid: u32) -> Result<()> { - self.proctrace.add_traced_pid(pid) + self.proctrace + .add_traced_pid(pid) .context("failed to add traced pid") } /// Remove a PID from the traced_processes map at runtime pub fn remove_traced_pid(&mut self, pid: u32) -> Result<()> { - self.proctrace.remove_traced_pid(pid) + self.proctrace + .remove_traced_pid(pid) .context("failed to remove traced pid") } @@ -331,7 +342,10 @@ impl Probes { if let Some(ref mut tcp) = self.tcpsniff { tcp.add_target(target) } else { - log::warn!("TcpSniff not enabled, cannot add runtime target {:?}", target); + log::warn!( + "TcpSniff not enabled, cannot add runtime target {:?}", + target + ); Ok(()) } } diff --git a/src/agentsight/src/probes/procmon.rs b/src/agentsight/src/probes/procmon.rs index c1ea16046..46afbcc17 100644 --- a/src/agentsight/src/probes/procmon.rs +++ b/src/agentsight/src/probes/procmon.rs @@ -9,13 +9,15 @@ use libbpf_rs::{ Link, MapHandle, skel::{OpenSkel, SkelBuilder}, }; -use std::{ - mem::MaybeUninit, - os::fd::AsFd, -}; +use std::{mem::MaybeUninit, os::fd::AsFd}; // ─── Generated skeleton ─────────────────────────────────────────────────────── -#[allow(non_camel_case_types, non_upper_case_globals, dead_code, non_snake_case)] +#[allow( + non_camel_case_types, + non_upper_case_globals, + dead_code, + non_snake_case +)] mod bpf { include!(concat!(env!("OUT_DIR"), "/procmon.skel.rs")); include!(concat!(env!("OUT_DIR"), "/procmon.rs")); @@ -61,7 +63,8 @@ impl Event { let raw = unsafe { &*(data.as_ptr() as *const ProcMonEvent) }; // Parse comm (null-terminated) - let comm = raw.comm + let comm = raw + .comm .iter() .take_while(|&&c| c != 0) .map(|&c| c as u8) @@ -143,6 +146,7 @@ impl ProcMon { let skel = open_skel.load().context("failed to load BPF object")?; // SAFETY: skel borrows open_object which lives in a Box + #[allow(clippy::unnecessary_cast)] let skel = unsafe { Box::from_raw(Box::into_raw(Box::new(skel)) as *mut ProcmonSkel<'static>) }; diff --git a/src/agentsight/src/probes/proctrace.rs b/src/agentsight/src/probes/proctrace.rs index e4aca139f..53066205c 100644 --- a/src/agentsight/src/probes/proctrace.rs +++ b/src/agentsight/src/probes/proctrace.rs @@ -21,7 +21,12 @@ use std::{ }; // ─── Generated skeleton ─────────────────────────────────────────────────────── -#[allow(non_camel_case_types, non_upper_case_globals, dead_code, non_snake_case)] +#[allow( + non_camel_case_types, + non_upper_case_globals, + dead_code, + non_snake_case +)] mod bpf { include!(concat!(env!("OUT_DIR"), "/proctrace.skel.rs")); include!(concat!(env!("OUT_DIR"), "/proctrace.rs")); @@ -72,7 +77,7 @@ impl VariableEvent { // SAFETY: BPF guarantees proper alignment and layout let raw_header = unsafe { &*(data.as_ptr() as *const ProcEventHeader) }; - + // Convert ktime to Unix timestamp let mut header = *raw_header; header.timestamp_ns = config::ktime_to_unix_ns(raw_header.timestamp_ns); @@ -135,10 +140,10 @@ impl VariableEvent { .map(|p| start + p) .unwrap_or(buf.len()); - if let Ok(s) = std::str::from_utf8(&buf[start..end]) { - if !s.is_empty() { - parts.push(s); - } + if let Ok(s) = std::str::from_utf8(&buf[start..end]) + && !s.is_empty() + { + parts.push(s); } start = end + 1; } @@ -286,10 +291,10 @@ impl ProcEvent { .position(|&c| c == 0) .map(|p| start + p) .unwrap_or(buf.len()); - if let Ok(s) = std::str::from_utf8(&buf[start..end]) { - if !s.is_empty() { - parts.push(s); - } + if let Ok(s) = std::str::from_utf8(&buf[start..end]) + && !s.is_empty() + { + parts.push(s); } start = end + 1; } @@ -408,6 +413,7 @@ impl ProcTrace { // SAFETY: skel borrows open_object which lives in a Box // on the heap. We pin both together inside Self and never move either, // so the 'static lifetime cast is sound for the lifetime of Self. + #[allow(clippy::unnecessary_cast)] let skel = unsafe { Box::from_raw(Box::into_raw(Box::new(skel)) as *mut ProctraceSkel<'static>) }; @@ -522,7 +528,7 @@ impl ProcTrace { let mut rb_builder = RingBufferBuilder::new(); let binding = self.skel.maps(); rb_builder - .add(&binding.rb(), move |data: &[u8]| { + .add(binding.rb(), move |data: &[u8]| { if data.len() < min_sz { return 0; } diff --git a/src/agentsight/src/probes/sslsniff.rs b/src/agentsight/src/probes/sslsniff.rs index b8bff9c0a..003357d4a 100644 --- a/src/agentsight/src/probes/sslsniff.rs +++ b/src/agentsight/src/probes/sslsniff.rs @@ -10,8 +10,8 @@ use libbpf_rs::{ Link, MapHandle, RingBufferBuilder, UprobeOpts, skel::{OpenSkel, SkelBuilder}, }; -use std::os::fd::AsFd; use procfs::process::Process; +use std::os::fd::AsFd; use std::{ collections::{HashMap, HashSet}, fs, @@ -27,7 +27,12 @@ use std::{ }; // ─── Generated skeleton ─────────────────────────────────────────────────────── -#[allow(non_camel_case_types, non_upper_case_globals, dead_code, non_snake_case)] +#[allow( + non_camel_case_types, + non_upper_case_globals, + dead_code, + non_snake_case +)] pub mod bpf { include!(concat!(env!("OUT_DIR"), "/sslsniff.skel.rs")); include!(concat!(env!("OUT_DIR"), "/sslsniff.rs")); @@ -39,7 +44,7 @@ const MAX_BUF_SIZE: usize = bpf::MAX_BUF_SIZE as usize; const POLL_TIMEOUT_MS: u64 = 100; /// User-space SslEvent - lightweight version of BPF probe_SSL_data_t -/// +/// /// Unlike the BPF version which has a 512KB fixed-size buffer, this struct /// only stores the actual data received, significantly reducing memory usage. #[derive(Debug, Clone)] @@ -69,22 +74,22 @@ impl SslEvent { let buf = raw.buf[..buf_size.min(MAX_BUF_SIZE)].to_vec(); // Convert ktime (nanoseconds since boot) to Unix timestamp - let ktime_ns = raw.timestamp_ns as u64; + let ktime_ns = raw.timestamp_ns; let unix_ts_ns = config::ktime_to_unix_ns(ktime_ns); Self { - source: raw.source as u32, + source: raw.source, timestamp_ns: unix_ts_ns, - delta_ns: raw.delta_ns as u64, - pid: raw.pid as u32, - tid: raw.tid as u32, - uid: raw.uid as u32, - len: raw.len as u32, + delta_ns: raw.delta_ns, + pid: raw.pid, + tid: raw.tid, + uid: raw.uid, + len: raw.len, rw: raw.rw, comm: Self::parse_comm(&raw.comm), buf, is_handshake: raw.is_handshake != 0, - ssl_ptr: raw.ssl_ptr as u64, + ssl_ptr: raw.ssl_ptr, } } @@ -93,11 +98,7 @@ impl SslEvent { fn parse_comm(comm: &[T; 16]) -> String { debug_assert_eq!(mem::size_of::(), 1); let bytes = unsafe { slice::from_raw_parts(comm.as_ptr() as *const u8, 16) }; - let bytes: Vec = bytes - .iter() - .copied() - .take_while(|&b| b != 0) - .collect(); + let bytes: Vec = bytes.iter().copied().take_while(|&b| b != 0).collect(); String::from_utf8_lossy(&bytes).into_owned() } @@ -113,7 +114,9 @@ impl SslEvent { } pub fn is_http_request(&self) -> bool { - const METHODS: &[&[u8]] = &[b"GET ", b"POST", b"PUT ", b"DELE", b"HEAD", b"OPTI", b"PATC"]; + const METHODS: &[&[u8]] = &[ + b"GET ", b"POST", b"PUT ", b"DELE", b"HEAD", b"OPTI", b"PATC", + ]; METHODS.iter().any(|m| self.buf.starts_with(m)) } @@ -132,9 +135,8 @@ impl SslEvent { return false; } // Parse 3-byte frame length - let length = ((self.buf[0] as usize) << 16) - | ((self.buf[1] as usize) << 8) - | (self.buf[2] as usize); + let length = + ((self.buf[0] as usize) << 16) | ((self.buf[1] as usize) << 8) | (self.buf[2] as usize); // Frame type must be a known type (0..=9) let frame_type = self.buf[3]; if frame_type > 9 { @@ -198,11 +200,14 @@ impl SslSniff { } /// Create a new SslSniff with an optional external traced_processes map and shared ring buffer - /// + /// /// # Arguments /// * `traced_processes` - Optional external MapHandle for traced_processes (for map reuse) /// * `rb` - Optional external MapHandle for shared ring buffer (for map reuse) - pub fn new_with_traced_processes(traced_processes: Option<&MapHandle>, rb: Option<&MapHandle>) -> Result { + pub fn new_with_traced_processes( + traced_processes: Option<&MapHandle>, + rb: Option<&MapHandle>, + ) -> Result { // ── Open + load skeleton ─────────────────────────────────────── let mut builder = SslsniffSkelBuilder::default(); builder.obj_builder.debug(config::verbose()); @@ -233,6 +238,7 @@ impl SslSniff { // SAFETY: skel borrows open_object which lives in a Box // on the heap. We pin both together inside Self and never move either, // so the 'static lifetime cast is sound for the lifetime of Self. + #[allow(clippy::unnecessary_cast)] let skel = unsafe { Box::from_raw(Box::into_raw(Box::new(skel)) as *mut SslsniffSkel<'static>) }; @@ -261,9 +267,12 @@ impl SslSniff { } // Debug: print all libs found - log::debug!("[attach_process] pid={pid}: found {} libs: {:?}", - libs.len(), - libs.iter().map(|(p, i, k)| (p.as_str(), *i, format!("{:?}", k))).collect::>() + log::debug!( + "[attach_process] pid={pid}: found {} libs: {:?}", + libs.len(), + libs.iter() + .map(|(p, i, k)| (p.as_str(), *i, format!("{:?}", k))) + .collect::>() ); let mut attached_inodes = Vec::new(); @@ -282,31 +291,25 @@ impl SslSniff { SslLibKind::OpenSsl => attach_openssl(&mut self.skel, &path, -1), SslLibKind::GnuTls => attach_gnutls(&mut self.skel, &path, -1), SslLibKind::Nss => attach_nss(&mut self.skel, &path, -1), - SslLibKind::Boring => { - match attach_boringssl_by_symbol(&mut self.skel, &path, -1) { - Ok(ls) => Ok(ls), - Err(sym_err) => { - log::debug!( - "[attach_process] pid={pid}: BoringSSL symbol attach failed for {path} ({sym_err:#}), falling back to byte-pattern" - ); - match find_boringssl_offsets(&path) { - Some(off) => attach_boringssl_by_offset( - &mut self.skel, - &path, - &off, - false, - -1, - ), - None => { - log::warn!( - "[attach_process] pid={pid}: BoringSSL detection failed for {path} (no SSL_* in .dynsym and no byte-pattern match), skipping" - ); - continue; - } + SslLibKind::Boring => match attach_boringssl_by_symbol(&mut self.skel, &path, -1) { + Ok(ls) => Ok(ls), + Err(sym_err) => { + log::debug!( + "[attach_process] pid={pid}: BoringSSL symbol attach failed for {path} ({sym_err:#}), falling back to byte-pattern" + ); + match find_boringssl_offsets(&path) { + Some(off) => { + attach_boringssl_by_offset(&mut self.skel, &path, &off, false, -1) + } + None => { + log::warn!( + "[attach_process] pid={pid}: BoringSSL detection failed for {path} (no SSL_* in .dynsym and no byte-pattern match), skipping" + ); + continue; } } } - } + }, }; match result { @@ -339,7 +342,10 @@ impl SslSniff { for inode in &inodes { self.traced_files.remove(inode); } - log::debug!("[detach_process] pid={pid}: removed {} inodes from traced_files", inodes.len()); + log::debug!( + "[detach_process] pid={pid}: removed {} inodes from traced_files", + inodes.len() + ); } } @@ -360,7 +366,7 @@ impl SslSniff { let mut rb_builder = RingBufferBuilder::new(); let binding = self.skel.maps(); rb_builder - .add(&binding.rb(), move |data: &[u8]| { + .add(binding.rb(), move |data: &[u8]| { if data.len() < min_sz { return 0; } @@ -535,7 +541,7 @@ fn find_boringssl_offsets(path: &str) -> Option { hs_matches[0] } else { // Multiple matches: choose the one closest before read_off. - match hs_matches.iter().filter(|&&o| o < read_off).last() { + match hs_matches.iter().filter(|&&o| o < read_off).next_back() { Some(&o) => o, None => { if verbose { @@ -686,6 +692,7 @@ fn ssl_libs_from_maps(pid: i32) -> Result> { // ─── uprobe helpers ─────────────────────────────────────────────────────────── +#[allow(clippy::field_reassign_with_default)] fn make_sym_opts(sym: &str, retprobe: bool) -> UprobeOpts { let mut o = UprobeOpts::default(); o.func_name = sym.to_string(); diff --git a/src/agentsight/src/probes/tcpsniff.rs b/src/agentsight/src/probes/tcpsniff.rs index bffa836ec..2c5f4fc49 100644 --- a/src/agentsight/src/probes/tcpsniff.rs +++ b/src/agentsight/src/probes/tcpsniff.rs @@ -16,17 +16,18 @@ use crate::config::{self, TcpTarget}; use anyhow::{Context, Result}; use libbpf_rs::{ - Link, MapHandle, MapFlags, + Link, MapFlags, MapHandle, skel::{OpenSkel, SkelBuilder}, }; -use std::{ - mem::MaybeUninit, - net::Ipv4Addr, - os::fd::AsFd, -}; +use std::{mem::MaybeUninit, net::Ipv4Addr, os::fd::AsFd}; // --- Generated skeleton --- -#[allow(non_camel_case_types, non_upper_case_globals, dead_code, non_snake_case)] +#[allow( + non_camel_case_types, + non_upper_case_globals, + dead_code, + non_snake_case +)] mod bpf { include!(concat!(env!("OUT_DIR"), "/tcpsniff.skel.rs")); } @@ -56,7 +57,9 @@ impl TcpSniff { builder.obj_builder.debug(config::verbose()); let open_object = Box::new(MaybeUninit::::uninit()); - let mut open_skel = builder.open().context("failed to open tcpsniff BPF object")?; + let mut open_skel = builder + .open() + .context("failed to open tcpsniff BPF object")?; // Reuse the shared ring buffer open_skel @@ -94,9 +97,12 @@ impl TcpSniff { .context("failed to disable old recvmsg fexit")?; } - let skel = open_skel.load().context("failed to load tcpsniff BPF object")?; + let skel = open_skel + .load() + .context("failed to load tcpsniff BPF object")?; // SAFETY: skel borrows open_object which lives in a Box + #[allow(clippy::unnecessary_cast)] let skel = unsafe { Box::from_raw(Box::into_raw(Box::new(skel)) as *mut TcpsniffSkel<'static>) }; diff --git a/src/agentsight/src/probes/udpdns.rs b/src/agentsight/src/probes/udpdns.rs index f0bb9e7cb..137483d08 100644 --- a/src/agentsight/src/probes/udpdns.rs +++ b/src/agentsight/src/probes/udpdns.rs @@ -13,13 +13,15 @@ use libbpf_rs::{ Link, MapHandle, skel::{OpenSkel, SkelBuilder}, }; -use std::{ - mem::MaybeUninit, - os::fd::AsFd, -}; +use std::{mem::MaybeUninit, os::fd::AsFd}; // --- Generated skeleton --- -#[allow(non_camel_case_types, non_upper_case_globals, dead_code, non_snake_case)] +#[allow( + non_camel_case_types, + non_upper_case_globals, + dead_code, + non_snake_case +)] mod bpf { include!(concat!(env!("OUT_DIR"), "/udpdns.skel.rs")); include!(concat!(env!("OUT_DIR"), "/udpdns.rs")); @@ -129,7 +131,8 @@ impl UdpDnsEvent { let raw = unsafe { &*(data.as_ptr() as *const RawUdpDnsEvent) }; // Parse comm (null-terminated) - let comm = raw.comm + let comm = raw + .comm .iter() .take_while(|&&c| c != 0) .map(|&c| c as u8) @@ -186,9 +189,12 @@ impl UdpDns { .reuse_fd(rb.as_fd()) .context("failed to reuse external rb map for udpdns")?; - let skel = open_skel.load().context("failed to load udpdns BPF object")?; + let skel = open_skel + .load() + .context("failed to load udpdns BPF object")?; // SAFETY: skel borrows open_object which lives in a Box + #[allow(clippy::unnecessary_cast)] let skel = unsafe { Box::from_raw(Box::into_raw(Box::new(skel)) as *mut UdpdnsSkel<'static>) }; diff --git a/src/agentsight/src/response_map.rs b/src/agentsight/src/response_map.rs index 2bcf71560..964269999 100644 --- a/src/agentsight/src/response_map.rs +++ b/src/agentsight/src/response_map.rs @@ -30,8 +30,7 @@ static RESPONSE_ID_RE: Lazy = /// Regex to match Anthropic/Claude Code message id format: `"id":"msg_"`. /// Only matches values starting with `msg_` to avoid false positives from other "id" fields. -static ANTHROPIC_MSG_ID_RE: Lazy = - Lazy::new(|| Regex::new(r#""id":"(msg_[^"]+)"#).unwrap()); +static ANTHROPIC_MSG_ID_RE: Lazy = Lazy::new(|| Regex::new(r#""id":"(msg_[^"]+)"#).unwrap()); /// Processes FileWrite events to build an in-memory responseId → sessionId mapping. /// Uses an LRU cache to bound memory usage. @@ -50,9 +49,7 @@ impl ResponseSessionMapper { /// Create a new empty mapper with default capacity. pub fn new() -> Self { ResponseSessionMapper { - map: LruCache::new( - NonZeroUsize::new(MAX_RESPONSE_MAP_ENTRIES).unwrap(), - ), + map: LruCache::new(NonZeroUsize::new(MAX_RESPONSE_MAP_ENTRIES).unwrap()), } } @@ -156,9 +153,8 @@ mod tests { #[test] fn test_extract_session_id_simple() { - let id = ResponseSessionMapper::extract_session_id( - "550e8400-e29b-41d4-a716-446655440000.jsonl", - ); + let id = + ResponseSessionMapper::extract_session_id("550e8400-e29b-41d4-a716-446655440000.jsonl"); assert_eq!(id.as_deref(), Some("550e8400-e29b-41d4-a716-446655440000")); } @@ -182,7 +178,8 @@ mod tests { #[test] fn test_extract_response_id() { - let line = r#"{"responseId":"chatcmpl-03a158a1-8982-90cd-adb1-6c8a1176f1f8","other":"data"}"#; + let line = + r#"{"responseId":"chatcmpl-03a158a1-8982-90cd-adb1-6c8a1176f1f8","other":"data"}"#; let id = ResponseSessionMapper::extract_response_id(line); assert_eq!( id.as_deref(), diff --git a/src/agentsight/src/server/handlers.rs b/src/agentsight/src/server/handlers.rs index 002ab2c60..500eb2363 100644 --- a/src/agentsight/src/server/handlers.rs +++ b/src/agentsight/src/server/handlers.rs @@ -1,12 +1,12 @@ //! API request handlers -use actix_web::{get, post, web, HttpResponse, Responder}; +use actix_web::{HttpResponse, Responder, get, post, web}; use serde::{Deserialize, Serialize}; use super::AppState; use crate::health::AgentHealthStatus; -use crate::storage::sqlite::{GenAISqliteStore}; -use crate::storage::sqlite::genai::{TimeseriesBucket, ModelTimeseriesBucket}; +use crate::storage::sqlite::GenAISqliteStore; +use crate::storage::sqlite::genai::{ModelTimeseriesBucket, TimeseriesBucket}; use crate::storage::sqlite::tokenless::{self, TokenlessStatsStore}; // ─── Prometheus helpers ─────────────────────────────────────────────────────── @@ -15,8 +15,8 @@ use crate::storage::sqlite::tokenless::{self, TokenlessStatsStore}; /// backslash → \\, double-quote → \", newline → \n fn escape_label(s: &str) -> String { s.replace('\\', "\\\\") - .replace('"', "\\\"") - .replace('\n', "\\n") + .replace('"', "\\\"") + .replace('\n', "\\n") } /// GET /health — health check endpoint @@ -51,7 +51,9 @@ pub async fn list_sessions( let db_path = &data.storage_path; let end_ns = query.end_ns.unwrap_or_else(|| now_ns() as i64); - let start_ns = query.start_ns.unwrap_or_else(|| end_ns - 86_400_000_000_000i64); // 24 h + let start_ns = query + .start_ns + .unwrap_or_else(|| end_ns - 86_400_000_000_000i64); // 24 h match GenAISqliteStore::new_with_path(db_path) { Ok(store) => match store.list_sessions(start_ns, end_ns) { @@ -59,8 +61,9 @@ pub async fn list_sessions( Err(e) => HttpResponse::InternalServerError() .json(serde_json::json!({"error": e.to_string()})), }, - Err(e) => HttpResponse::InternalServerError() - .json(serde_json::json!({"error": e.to_string()})), + Err(e) => { + HttpResponse::InternalServerError().json(serde_json::json!({"error": e.to_string()})) + } } } @@ -86,8 +89,9 @@ pub async fn list_traces_by_session( Err(e) => HttpResponse::InternalServerError() .json(serde_json::json!({"error": e.to_string()})), }, - Err(e) => HttpResponse::InternalServerError() - .json(serde_json::json!({"error": e.to_string()})), + Err(e) => { + HttpResponse::InternalServerError().json(serde_json::json!({"error": e.to_string()})) + } } } @@ -108,8 +112,9 @@ pub async fn get_trace_detail( Err(e) => HttpResponse::InternalServerError() .json(serde_json::json!({"error": e.to_string()})), }, - Err(e) => HttpResponse::InternalServerError() - .json(serde_json::json!({"error": e.to_string()})), + Err(e) => { + HttpResponse::InternalServerError().json(serde_json::json!({"error": e.to_string()})) + } } } @@ -130,8 +135,9 @@ pub async fn get_conversation_events( Err(e) => HttpResponse::InternalServerError() .json(serde_json::json!({"error": e.to_string()})), }, - Err(e) => HttpResponse::InternalServerError() - .json(serde_json::json!({"error": e.to_string()})), + Err(e) => { + HttpResponse::InternalServerError().json(serde_json::json!({"error": e.to_string()})) + } } } @@ -165,7 +171,9 @@ pub async fn list_agent_names( ) -> impl Responder { let db_path = &data.storage_path; let end_ns = query.end_ns.unwrap_or_else(|| now_ns() as i64); - let start_ns = query.start_ns.unwrap_or_else(|| end_ns - 86_400_000_000_000i64); + let start_ns = query + .start_ns + .unwrap_or_else(|| end_ns - 86_400_000_000_000i64); match GenAISqliteStore::new_with_path(db_path) { Ok(store) => match store.list_agent_names(start_ns, end_ns) { @@ -173,8 +181,9 @@ pub async fn list_agent_names( Err(e) => HttpResponse::InternalServerError() .json(serde_json::json!({"error": e.to_string()})), }, - Err(e) => HttpResponse::InternalServerError() - .json(serde_json::json!({"error": e.to_string()})), + Err(e) => { + HttpResponse::InternalServerError().json(serde_json::json!({"error": e.to_string()})) + } } } @@ -196,26 +205,38 @@ pub async fn get_timeseries( ) -> impl Responder { let db_path = &data.storage_path; let end_ns = query.end_ns.unwrap_or_else(|| now_ns() as i64); - let start_ns = query.start_ns.unwrap_or_else(|| end_ns - 86_400_000_000_000i64); + let start_ns = query + .start_ns + .unwrap_or_else(|| end_ns - 86_400_000_000_000i64); let buckets = query.buckets.unwrap_or(30); let agent_name = query.agent_name.as_deref(); match GenAISqliteStore::new_with_path(db_path) { Ok(store) => { - let token_series = match store.get_token_timeseries(start_ns, end_ns, agent_name, buckets) { - Ok(v) => v, - Err(e) => return HttpResponse::InternalServerError() - .json(serde_json::json!({"error": e.to_string()})), - }; - let model_series = match store.get_model_timeseries(start_ns, end_ns, agent_name, buckets) { - Ok(v) => v, - Err(e) => return HttpResponse::InternalServerError() - .json(serde_json::json!({"error": e.to_string()})), - }; - HttpResponse::Ok().json(TimeseriesResponse { token_series, model_series }) + let token_series = + match store.get_token_timeseries(start_ns, end_ns, agent_name, buckets) { + Ok(v) => v, + Err(e) => { + return HttpResponse::InternalServerError() + .json(serde_json::json!({"error": e.to_string()})); + } + }; + let model_series = + match store.get_model_timeseries(start_ns, end_ns, agent_name, buckets) { + Ok(v) => v, + Err(e) => { + return HttpResponse::InternalServerError() + .json(serde_json::json!({"error": e.to_string()})); + } + }; + HttpResponse::Ok().json(TimeseriesResponse { + token_series, + model_series, + }) + } + Err(e) => { + HttpResponse::InternalServerError().json(serde_json::json!({"error": e.to_string()})) } - Err(e) => HttpResponse::InternalServerError() - .json(serde_json::json!({"error": e.to_string()})), } } @@ -261,23 +282,29 @@ pub async fn metrics(data: web::Data) -> impl Responder { let mut out = String::with_capacity(512 + summaries.len() * 128); // agentsight_token_input_total - out.push_str("# HELP agentsight_token_input_total Total input tokens consumed by agent (all-time)\n"); + out.push_str( + "# HELP agentsight_token_input_total Total input tokens consumed by agent (all-time)\n", + ); out.push_str("# TYPE agentsight_token_input_total counter\n"); for s in &summaries { out.push_str(&format!( "agentsight_token_input_total{{agent=\"{}\"}} {}\n", - escape_label(&s.agent_name), s.input_tokens + escape_label(&s.agent_name), + s.input_tokens )); } out.push('\n'); // agentsight_token_output_total - out.push_str("# HELP agentsight_token_output_total Total output tokens consumed by agent (all-time)\n"); + out.push_str( + "# HELP agentsight_token_output_total Total output tokens consumed by agent (all-time)\n", + ); out.push_str("# TYPE agentsight_token_output_total counter\n"); for s in &summaries { out.push_str(&format!( "agentsight_token_output_total{{agent=\"{}\"}} {}\n", - escape_label(&s.agent_name), s.output_tokens + escape_label(&s.agent_name), + s.output_tokens )); } out.push('\n'); @@ -288,18 +315,22 @@ pub async fn metrics(data: web::Data) -> impl Responder { for s in &summaries { out.push_str(&format!( "agentsight_token_total_total{{agent=\"{}\"}} {}\n", - escape_label(&s.agent_name), s.total_tokens + escape_label(&s.agent_name), + s.total_tokens )); } out.push('\n'); // agentsight_llm_requests_total - out.push_str("# HELP agentsight_llm_requests_total Total LLM requests made by agent (all-time)\n"); + out.push_str( + "# HELP agentsight_llm_requests_total Total LLM requests made by agent (all-time)\n", + ); out.push_str("# TYPE agentsight_llm_requests_total counter\n"); for s in &summaries { out.push_str(&format!( "agentsight_llm_requests_total{{agent=\"{}\"}} {}\n", - escape_label(&s.agent_name), s.request_count + escape_label(&s.agent_name), + s.request_count )); } out.push('\n'); @@ -368,7 +399,8 @@ pub async fn restart_agent_health( // 从 store 中取出 restart_cmd let restart_cmd = { let store = data.health_store.read().unwrap(); - store.all_agents() + store + .all_agents() .into_iter() .find(|a| a.pid == pid) .and_then(|a| a.restart_cmd) @@ -376,15 +408,15 @@ pub async fn restart_agent_health( let cmd = match restart_cmd { Some(c) if !c.is_empty() => c, - _ => return HttpResponse::BadRequest() - .json(serde_json::json!({"error": "no restart command available for this pid"})), + _ => { + return HttpResponse::BadRequest() + .json(serde_json::json!({"error": "no restart command available for this pid"})); + } }; // Step 1: kill -9 use std::process::Command; - let kill_result = Command::new("kill") - .args(["-9", &pid.to_string()]) - .output(); + let kill_result = Command::new("kill").args(["-9", &pid.to_string()]).output(); if let Err(e) = kill_result { return HttpResponse::InternalServerError() @@ -402,7 +434,9 @@ pub async fn restart_agent_health( let new_pid = child.id(); log::info!( "Restarted agent pid={} -> new pid={}, cmd={:?}", - pid, new_pid, cmd + pid, + new_pid, + cmd ); // 从 store 中删除旧 PID 条目，下次扫描时新 PID 会自动加入 data.health_store.write().unwrap().remove_by_pid(pid); @@ -434,7 +468,7 @@ pub async fn export_atif_trace( Ok(s) => s, Err(e) => { return HttpResponse::InternalServerError() - .json(serde_json::json!({"error": e.to_string()})) + .json(serde_json::json!({"error": e.to_string()})); } }; @@ -442,19 +476,19 @@ pub async fn export_atif_trace( Ok(e) => e, Err(e) => { return HttpResponse::InternalServerError() - .json(serde_json::json!({"error": e.to_string()})) + .json(serde_json::json!({"error": e.to_string()})); } }; if events.is_empty() { - return HttpResponse::NotFound() - .json(serde_json::json!({"error": "trace not found"})); + return HttpResponse::NotFound().json(serde_json::json!({"error": "trace not found"})); } match crate::atif::convert_trace_to_atif(&trace_id, events) { Ok(doc) => HttpResponse::Ok().json(doc), - Err(e) => HttpResponse::InternalServerError() - .json(serde_json::json!({"error": e.to_string()})), + Err(e) => { + HttpResponse::InternalServerError().json(serde_json::json!({"error": e.to_string()})) + } } } @@ -473,7 +507,7 @@ pub async fn export_atif_session( Ok(s) => s, Err(e) => { return HttpResponse::InternalServerError() - .json(serde_json::json!({"error": e.to_string()})) + .json(serde_json::json!({"error": e.to_string()})); } }; @@ -481,19 +515,19 @@ pub async fn export_atif_session( Ok(e) => e, Err(e) => { return HttpResponse::InternalServerError() - .json(serde_json::json!({"error": e.to_string()})) + .json(serde_json::json!({"error": e.to_string()})); } }; if events.is_empty() { - return HttpResponse::NotFound() - .json(serde_json::json!({"error": "session not found"})); + return HttpResponse::NotFound().json(serde_json::json!({"error": "session not found"})); } match crate::atif::convert_session_to_atif(&session_id, events) { Ok(doc) => HttpResponse::Ok().json(doc), - Err(e) => HttpResponse::InternalServerError() - .json(serde_json::json!({"error": e.to_string()})), + Err(e) => { + HttpResponse::InternalServerError().json(serde_json::json!({"error": e.to_string()})) + } } } @@ -512,7 +546,7 @@ pub async fn export_atif_conversation( Ok(s) => s, Err(e) => { return HttpResponse::InternalServerError() - .json(serde_json::json!({"error": e.to_string()})) + .json(serde_json::json!({"error": e.to_string()})); } }; @@ -520,7 +554,7 @@ pub async fn export_atif_conversation( Ok(e) => e, Err(e) => { return HttpResponse::InternalServerError() - .json(serde_json::json!({"error": e.to_string()})) + .json(serde_json::json!({"error": e.to_string()})); } }; @@ -531,8 +565,9 @@ pub async fn export_atif_conversation( match crate::atif::convert_trace_to_atif(&conversation_id, events) { Ok(doc) => HttpResponse::Ok().json(doc), - Err(e) => HttpResponse::InternalServerError() - .json(serde_json::json!({"error": e.to_string()})), + Err(e) => { + HttpResponse::InternalServerError().json(serde_json::json!({"error": e.to_string()})) + } } } @@ -564,12 +599,15 @@ pub async fn list_interruptions( .json(serde_json::json!({"error": "Interruption store not initialized"})); }; - let end_ns = query.end_ns.unwrap_or_else(|| now_ns() as i64); - let start_ns = query.start_ns.unwrap_or_else(|| end_ns - 86_400_000_000_000i64); // 24 h - let limit = query.limit.unwrap_or(200); + let end_ns = query.end_ns.unwrap_or_else(|| now_ns() as i64); + let start_ns = query + .start_ns + .unwrap_or_else(|| end_ns - 86_400_000_000_000i64); // 24 h + let limit = query.limit.unwrap_or(200); match istore.list( - start_ns, end_ns, + start_ns, + end_ns, query.agent_name.as_deref(), query.interruption_type.as_deref(), query.severity.as_deref(), @@ -577,8 +615,9 @@ pub async fn list_interruptions( limit, ) { Ok(rows) => HttpResponse::Ok().json(rows), - Err(e) => HttpResponse::InternalServerError() - .json(serde_json::json!({"error": e.to_string()})), + Err(e) => { + HttpResponse::InternalServerError().json(serde_json::json!({"error": e.to_string()})) + } } } @@ -596,8 +635,10 @@ pub async fn interruption_count( .json(serde_json::json!({"error": "Interruption store not initialized"})); }; - let end_ns = query.end_ns.unwrap_or_else(|| now_ns() as i64); - let start_ns = query.start_ns.unwrap_or_else(|| end_ns - 86_400_000_000_000i64); + let end_ns = query.end_ns.unwrap_or_else(|| now_ns() as i64); + let start_ns = query + .start_ns + .unwrap_or_else(|| end_ns - 86_400_000_000_000i64); match istore.stats(start_ns, end_ns) { Ok(stats) => { @@ -610,9 +651,9 @@ pub async fn interruption_count( total += s.count as u64; match s.severity.as_str() { "critical" => critical += s.count as u64, - "high" => high += s.count as u64, - "medium" => medium += s.count as u64, - _ => low += s.count as u64, + "high" => high += s.count as u64, + "medium" => medium += s.count as u64, + _ => low += s.count as u64, } } HttpResponse::Ok().json(serde_json::json!({ @@ -625,8 +666,9 @@ pub async fn interruption_count( } })) } - Err(e) => HttpResponse::InternalServerError() - .json(serde_json::json!({"error": e.to_string()})), + Err(e) => { + HttpResponse::InternalServerError().json(serde_json::json!({"error": e.to_string()})) + } } } @@ -643,13 +685,16 @@ pub async fn interruption_stats( .json(serde_json::json!({"error": "Interruption store not initialized"})); }; - let end_ns = query.end_ns.unwrap_or_else(|| now_ns() as i64); - let start_ns = query.start_ns.unwrap_or_else(|| end_ns - 86_400_000_000_000i64); + let end_ns = query.end_ns.unwrap_or_else(|| now_ns() as i64); + let start_ns = query + .start_ns + .unwrap_or_else(|| end_ns - 86_400_000_000_000i64); match istore.stats(start_ns, end_ns) { Ok(stats) => HttpResponse::Ok().json(stats), - Err(e) => HttpResponse::InternalServerError() - .json(serde_json::json!({"error": e.to_string()})), + Err(e) => { + HttpResponse::InternalServerError().json(serde_json::json!({"error": e.to_string()})) + } } } @@ -668,15 +713,27 @@ pub async fn interruption_session_counts( .json(serde_json::json!({"error": "Interruption store not initialized"})); }; - let end_ns = query.end_ns.unwrap_or_else(|| now_ns() as i64); - let start_ns = query.start_ns.unwrap_or_else(|| end_ns - 86_400_000_000_000i64); + let end_ns = query.end_ns.unwrap_or_else(|| now_ns() as i64); + let start_ns = query + .start_ns + .unwrap_or_else(|| end_ns - 86_400_000_000_000i64); match istore.count_unresolved_by_session_detailed(start_ns, end_ns) { Ok(rows) => { // Group by session_id - let mut map: std::collections::HashMap, Vec)> = std::collections::HashMap::new(); + #[allow(clippy::type_complexity)] + let mut map: std::collections::HashMap< + String, + ( + i64, + std::collections::HashMap, + Vec, + ), + > = std::collections::HashMap::new(); for (sid, severity, itype, cnt) in rows { - let entry = map.entry(sid).or_insert_with(|| (0, std::collections::HashMap::new(), Vec::new())); + let entry = map + .entry(sid) + .or_insert_with(|| (0, std::collections::HashMap::new(), Vec::new())); entry.0 += cnt; *entry.1.entry(severity.clone()).or_insert(0) += cnt; entry.2.push(serde_json::json!({ @@ -685,23 +742,27 @@ pub async fn interruption_session_counts( "count": cnt, })); } - let json: Vec<_> = map.into_iter().map(|(sid, (total, by_sev, types))| { - serde_json::json!({ - "session_id": sid, - "total": total, - "by_severity": { - "critical": by_sev.get("critical").copied().unwrap_or(0), - "high": by_sev.get("high").copied().unwrap_or(0), - "medium": by_sev.get("medium").copied().unwrap_or(0), - "low": by_sev.get("low").copied().unwrap_or(0), - }, - "types": types, + let json: Vec<_> = map + .into_iter() + .map(|(sid, (total, by_sev, types))| { + serde_json::json!({ + "session_id": sid, + "total": total, + "by_severity": { + "critical": by_sev.get("critical").copied().unwrap_or(0), + "high": by_sev.get("high").copied().unwrap_or(0), + "medium": by_sev.get("medium").copied().unwrap_or(0), + "low": by_sev.get("low").copied().unwrap_or(0), + }, + "types": types, + }) }) - }).collect(); + .collect(); HttpResponse::Ok().json(json) } - Err(e) => HttpResponse::InternalServerError() - .json(serde_json::json!({"error": e.to_string()})), + Err(e) => { + HttpResponse::InternalServerError().json(serde_json::json!({"error": e.to_string()})) + } } } @@ -720,14 +781,26 @@ pub async fn interruption_conversation_counts( .json(serde_json::json!({"error": "Interruption store not initialized"})); }; - let end_ns = query.end_ns.unwrap_or_else(|| now_ns() as i64); - let start_ns = query.start_ns.unwrap_or_else(|| end_ns - 86_400_000_000_000i64); + let end_ns = query.end_ns.unwrap_or_else(|| now_ns() as i64); + let start_ns = query + .start_ns + .unwrap_or_else(|| end_ns - 86_400_000_000_000i64); match istore.count_unresolved_by_conversation_detailed(start_ns, end_ns) { Ok(rows) => { - let mut map: std::collections::HashMap, Vec)> = std::collections::HashMap::new(); + #[allow(clippy::type_complexity)] + let mut map: std::collections::HashMap< + String, + ( + i64, + std::collections::HashMap, + Vec, + ), + > = std::collections::HashMap::new(); for (cid, severity, itype, cnt) in rows { - let entry = map.entry(cid).or_insert_with(|| (0, std::collections::HashMap::new(), Vec::new())); + let entry = map + .entry(cid) + .or_insert_with(|| (0, std::collections::HashMap::new(), Vec::new())); entry.0 += cnt; *entry.1.entry(severity.clone()).or_insert(0) += cnt; entry.2.push(serde_json::json!({ @@ -736,23 +809,27 @@ pub async fn interruption_conversation_counts( "count": cnt, })); } - let json: Vec<_> = map.into_iter().map(|(cid, (total, by_sev, types))| { - serde_json::json!({ - "conversation_id": cid, - "total": total, - "by_severity": { - "critical": by_sev.get("critical").copied().unwrap_or(0), - "high": by_sev.get("high").copied().unwrap_or(0), - "medium": by_sev.get("medium").copied().unwrap_or(0), - "low": by_sev.get("low").copied().unwrap_or(0), - }, - "types": types, + let json: Vec<_> = map + .into_iter() + .map(|(cid, (total, by_sev, types))| { + serde_json::json!({ + "conversation_id": cid, + "total": total, + "by_severity": { + "critical": by_sev.get("critical").copied().unwrap_or(0), + "high": by_sev.get("high").copied().unwrap_or(0), + "medium": by_sev.get("medium").copied().unwrap_or(0), + "low": by_sev.get("low").copied().unwrap_or(0), + }, + "types": types, + }) }) - }).collect(); + .collect(); HttpResponse::Ok().json(json) } - Err(e) => HttpResponse::InternalServerError() - .json(serde_json::json!({"error": e.to_string()})), + Err(e) => { + HttpResponse::InternalServerError().json(serde_json::json!({"error": e.to_string()})) + } } } @@ -772,8 +849,9 @@ pub async fn list_session_interruptions( let session_id = path.into_inner(); match istore.list_by_session(&session_id) { Ok(rows) => HttpResponse::Ok().json(rows), - Err(e) => HttpResponse::InternalServerError() - .json(serde_json::json!({"error": e.to_string()})), + Err(e) => { + HttpResponse::InternalServerError().json(serde_json::json!({"error": e.to_string()})) + } } } @@ -793,8 +871,9 @@ pub async fn list_conversation_interruptions( let conversation_id = path.into_inner(); match istore.list_by_conversation(&conversation_id) { Ok(rows) => HttpResponse::Ok().json(rows), - Err(e) => HttpResponse::InternalServerError() - .json(serde_json::json!({"error": e.to_string()})), + Err(e) => { + HttpResponse::InternalServerError().json(serde_json::json!({"error": e.to_string()})) + } } } @@ -813,11 +892,13 @@ pub async fn resolve_interruption( let interruption_id = path.into_inner(); match istore.resolve(&interruption_id) { - Ok(true) => HttpResponse::Ok().json(serde_json::json!({"status": "resolved"})), - Ok(false) => HttpResponse::NotFound() - .json(serde_json::json!({"error": "Interruption not found"})), - Err(e) => HttpResponse::InternalServerError() - .json(serde_json::json!({"error": e.to_string()})), + Ok(true) => HttpResponse::Ok().json(serde_json::json!({"status": "resolved"})), + Ok(false) => { + HttpResponse::NotFound().json(serde_json::json!({"error": "Interruption not found"})) + } + Err(e) => { + HttpResponse::InternalServerError().json(serde_json::json!({"error": e.to_string()})) + } } } @@ -837,10 +918,12 @@ pub async fn get_interruption( let interruption_id = path.into_inner(); match istore.get_by_id(&interruption_id) { Ok(Some(row)) => HttpResponse::Ok().json(row), - Ok(None) => HttpResponse::NotFound() - .json(serde_json::json!({"error": "Interruption not found"})), - Err(e) => HttpResponse::InternalServerError() - .json(serde_json::json!({"error": e.to_string()})), + Ok(None) => { + HttpResponse::NotFound().json(serde_json::json!({"error": "Interruption not found"})) + } + Err(e) => { + HttpResponse::InternalServerError().json(serde_json::json!({"error": e.to_string()})) + } } } @@ -951,18 +1034,24 @@ pub async fn get_token_savings( ) -> impl Responder { let db_path = &data.storage_path; let end_ns = query.end_ns.unwrap_or_else(|| now_ns() as i64); - let start_ns = query.start_ns.unwrap_or_else(|| end_ns - 86_400_000_000_000i64); + let start_ns = query + .start_ns + .unwrap_or_else(|| end_ns - 86_400_000_000_000i64); let agent_name = query.agent_name.as_deref(); // Step 1: Query sessions from genai_events.db let sessions = match GenAISqliteStore::new_with_path(db_path) { Ok(store) => match store.list_sessions_for_savings(start_ns, end_ns, agent_name) { Ok(s) => s, - Err(e) => return HttpResponse::InternalServerError() - .json(serde_json::json!({"error": e.to_string()})), + Err(e) => { + return HttpResponse::InternalServerError() + .json(serde_json::json!({"error": e.to_string()})); + } }, - Err(e) => return HttpResponse::InternalServerError() - .json(serde_json::json!({"error": e.to_string()})), + Err(e) => { + return HttpResponse::InternalServerError() + .json(serde_json::json!({"error": e.to_string()})); + } }; // Step 2: Open stats.db (read-only, graceful if absent) @@ -974,7 +1063,9 @@ pub async fn get_token_savings( // This gives us all known tool_use_ids and their session membership. let session_ids: Vec<&str> = sessions.iter().map(|s| s.session_id.as_str()).collect(); let turn_indices = match GenAISqliteStore::new_with_path(db_path) { - Ok(store) => store.get_tool_call_turn_indices(&session_ids).unwrap_or_default(), + Ok(store) => store + .get_tool_call_turn_indices(&session_ids) + .unwrap_or_default(), Err(_) => std::collections::HashMap::new(), }; @@ -1056,7 +1147,10 @@ pub async fn get_token_savings( saved_tokens: saved, compounded_saved: compounded, compounding_turns, - before_summary: format!("\u{539f}\u{59cb}\u{5185}\u{5bb9} {} tokens", row.before_tokens), + before_summary: format!( + "\u{539f}\u{59cb}\u{5185}\u{5bb9} {} tokens", + row.before_tokens + ), after_summary: format!("\u{4f18}\u{5316}\u{540e} {} tokens", row.after_tokens), before_text: row.before_text.clone(), after_text: row.after_text.clone(), @@ -1245,10 +1339,10 @@ fn compute_skill_metrics_response( mut options: crate::skill_metrics::MetricOptions, ) -> HttpResponse { // Apply granularity from query params - if let Some(ref g) = query.granularity { - if g == "day" { - options.hotness_granularity = crate::skill_metrics::HotnessGranularity::Day; - } + if let Some(ref g) = query.granularity + && g == "day" + { + options.hotness_granularity = crate::skill_metrics::HotnessGranularity::Day; } let end_ns = query.end_ns.unwrap_or_else(|| now_ns() as i64); diff --git a/src/agentsight/src/server/mod.rs b/src/agentsight/src/server/mod.rs index eb731079c..b764f4498 100644 --- a/src/agentsight/src/server/mod.rs +++ b/src/agentsight/src/server/mod.rs @@ -10,8 +10,8 @@ use std::sync::{Arc, RwLock}; use std::time::{Duration, Instant}; use actix_cors::Cors; -use actix_web::{get, web, App, HttpRequest, HttpResponse, HttpServer, Responder}; -use include_dir::{include_dir, Dir}; +use actix_web::{App, HttpRequest, HttpResponse, HttpServer, Responder, get, web}; +use include_dir::{Dir, include_dir}; use crate::health::{HealthChecker, HealthStore}; use crate::storage::sqlite::InterruptionStore; @@ -56,9 +56,7 @@ async fn serve_frontend(req: HttpRequest) -> impl Responder { } else { mime_for_path(path) }; - HttpResponse::Ok() - .content_type(mime) - .body(f.contents()) + HttpResponse::Ok().content_type(mime).body(f.contents()) } None => { // SPA fallback: return index.html for unmatched paths @@ -66,22 +64,33 @@ async fn serve_frontend(req: HttpRequest) -> impl Responder { Some(index) => HttpResponse::Ok() .content_type("text/html; charset=utf-8") .body(index.contents()), - None => HttpResponse::NotFound().body("Frontend not embedded. Run `npm run build:embed` first."), + None => HttpResponse::NotFound() + .body("Frontend not embedded. Run `npm run build:embed` first."), } } } } fn mime_for_path(path: &str) -> &'static str { - if path.ends_with(".html") { "text/html; charset=utf-8" } - else if path.ends_with(".js") { "application/javascript; charset=utf-8" } - else if path.ends_with(".css") { "text/css; charset=utf-8" } - else if path.ends_with(".json") { "application/json" } - else if path.ends_with(".svg") { "image/svg+xml" } - else if path.ends_with(".png") { "image/png" } - else if path.ends_with(".ico") { "image/x-icon" } - else if path.ends_with(".woff2") { "font/woff2" } - else { "application/octet-stream" } + if path.ends_with(".html") { + "text/html; charset=utf-8" + } else if path.ends_with(".js") { + "application/javascript; charset=utf-8" + } else if path.ends_with(".css") { + "text/css; charset=utf-8" + } else if path.ends_with(".json") { + "application/json" + } else if path.ends_with(".svg") { + "image/svg+xml" + } else if path.ends_with(".png") { + "image/png" + } else if path.ends_with(".ico") { + "image/x-icon" + } else if path.ends_with(".woff2") { + "font/woff2" + } else { + "application/octet-stream" + } } // ─── Server entry point ─────────────────────────────────────────────────────── @@ -142,12 +151,21 @@ pub async fn run_server(host: &str, port: u16, storage_path: PathBuf) -> std::io }); let has_frontend = FRONTEND.get_file("index.html").is_some(); - log::info!("AgentSight API server listening on http://{}:{}", host, port); - eprintln!("AgentSight API server listening on http://{}:{}", host, port); + log::info!( + "AgentSight API server listening on http://{}:{}", + host, + port + ); + eprintln!( + "AgentSight API server listening on http://{}:{}", + host, port + ); if has_frontend { eprintln!("Dashboard UI: http://{}:{}/", host, port); } else { - eprintln!("[WARN] Frontend not embedded. Run `npm run build:embed` in dashboard/ then recompile."); + eprintln!( + "[WARN] Frontend not embedded. Run `npm run build:embed` in dashboard/ then recompile." + ); } HttpServer::new(move || { diff --git a/src/agentsight/src/skill_metrics/extractor.rs b/src/agentsight/src/skill_metrics/extractor.rs index 2972ad49d..b097ccafa 100644 --- a/src/agentsight/src/skill_metrics/extractor.rs +++ b/src/agentsight/src/skill_metrics/extractor.rs @@ -64,9 +64,7 @@ pub fn extract_skill_downloads(event: &TraceEventDetail) -> Vec Result { - let sql = format!( - "DELETE FROM {} WHERE timestamp_ns < ?1", - self.table_name - ); + let sql = format!("DELETE FROM {} WHERE timestamp_ns < ?1", self.table_name); let deleted = self.conn.execute(&sql, params![cutoff_ns as i64])?; Ok(deleted as u64) } @@ -230,12 +217,10 @@ impl AuditStore { std::collections::HashMap::new(); { - let mut stmt = self.conn.prepare( - &format!( - "SELECT extra FROM {} WHERE timestamp_ns >= ?1 AND event_type = 'llm_call'", - self.table_name - ), - )?; + let mut stmt = self.conn.prepare(&format!( + "SELECT extra FROM {} WHERE timestamp_ns >= ?1 AND event_type = 'llm_call'", + self.table_name + ))?; let rows = stmt.query_map(params![since_ns as i64], |row| { let extra_str: String = row.get(0)?; Ok(extra_str) @@ -243,16 +228,16 @@ impl AuditStore { for row in rows.flatten() { if let Ok(extra) = serde_json::from_str::(&row) { - total_input_tokens += - extra.get("input_tokens").and_then(|v| v.as_u64()).unwrap_or(0); + total_input_tokens += extra + .get("input_tokens") + .and_then(|v| v.as_u64()) + .unwrap_or(0); total_output_tokens += extra .get("output_tokens") .and_then(|v| v.as_u64()) .unwrap_or(0); if let Some(provider) = extra.get("provider").and_then(|v| v.as_str()) { - *provider_counts - .entry(provider.to_string()) - .or_insert(0) += 1; + *provider_counts.entry(provider.to_string()).or_insert(0) += 1; } } } @@ -337,16 +322,16 @@ fn extract_cmdline_from_extra(extra_str: &str, comm: &str) -> String { // Try to parse as ProcessAction extra if let Ok(json) = serde_json::from_str::(extra_str) { // Try args first (full command line) - if let Some(args) = json.get("args").and_then(|v| v.as_str()) { - if !args.is_empty() { - return args.to_string(); - } + if let Some(args) = json.get("args").and_then(|v| v.as_str()) + && !args.is_empty() + { + return args.to_string(); } // Try filename as fallback - if let Some(filename) = json.get("filename").and_then(|v| v.as_str()) { - if !filename.is_empty() { - return filename.to_string(); - } + if let Some(filename) = json.get("filename").and_then(|v| v.as_str()) + && !filename.is_empty() + { + return filename.to_string(); } } // Fallback to comm (process name, max 16 chars) diff --git a/src/agentsight/src/storage/sqlite/connection.rs b/src/agentsight/src/storage/sqlite/connection.rs index 5ccc86a42..c62a73402 100644 --- a/src/agentsight/src/storage/sqlite/connection.rs +++ b/src/agentsight/src/storage/sqlite/connection.rs @@ -56,15 +56,15 @@ mod tests { #[test] fn test_create_connection() { let test_path = PathBuf::from("/tmp/test_agentsight_connection.db"); - + // Clean up if exists let _ = fs::remove_file(&test_path); - + let conn = create_connection(&test_path).unwrap(); drop(conn); - + assert!(test_path.exists()); - + // Cleanup fs::remove_file(&test_path).ok(); } diff --git a/src/agentsight/src/storage/sqlite/genai.rs b/src/agentsight/src/storage/sqlite/genai.rs index 40ac0fdaa..2debb25d4 100644 --- a/src/agentsight/src/storage/sqlite/genai.rs +++ b/src/agentsight/src/storage/sqlite/genai.rs @@ -493,9 +493,8 @@ impl GenAISqliteStore { } }; let input_messages: Option = { - let latest = crate::genai::semantic::latest_round_input_messages( - &call.request.messages, - ); + let latest = + crate::genai::semantic::latest_round_input_messages(&call.request.messages); if latest.is_empty() { None } else { @@ -693,6 +692,7 @@ impl GenAISqliteStore { /// Returns (call_id, session_id, trace_id, conversation_id) tuples for all /// PENDING records matching the given PID. Used by HealthChecker to link /// agent_crash events to their associated LLM calls. + #[allow(clippy::type_complexity)] pub fn list_pending_for_pid( &self, pid: i32, @@ -751,6 +751,7 @@ impl GenAISqliteStore { } /// Like `list_pending_for_pid` but accepts multiple PIDs at once. + #[allow(clippy::type_complexity)] pub fn list_pending_for_pids( &self, pids: &[i32], @@ -1015,7 +1016,8 @@ impl GenAISqliteStore { pub fn get_tool_call_turn_indices( &self, session_ids: &[&str], - ) -> Result, Box> { + ) -> Result, Box> + { let conn = self.conn.lock().unwrap(); let mut result = std::collections::HashMap::new(); @@ -1037,20 +1039,26 @@ impl GenAISqliteStore { // Also map the call_id itself (for backward compat with // stats.db that may still store call_id as tool_use_id) - result.insert(call_id.clone(), ToolCallTurnInfo { - turn_index: turn, - session_id: session_id.clone(), - }); + result.insert( + call_id.clone(), + ToolCallTurnInfo { + turn_index: turn, + session_id: session_id.clone(), + }, + ); // Expand each tool_call_id in the JSON array - if let Some(json_str) = tool_call_ids_json { - if let Ok(ids) = serde_json::from_str::>(&json_str) { - for tc_id in ids { - result.insert(tc_id, ToolCallTurnInfo { + if let Some(json_str) = tool_call_ids_json + && let Ok(ids) = serde_json::from_str::>(&json_str) + { + for tc_id in ids { + result.insert( + tc_id, + ToolCallTurnInfo { turn_index: turn, session_id: session_id.clone(), - }); - } + }, + ); } } } @@ -1073,8 +1081,7 @@ impl GenAISqliteStore { // When both start_ns and end_ns are present, rewrite with BETWEEN let sql = if start_ns.is_some() && end_ns.is_some() { - format!( - "SELECT conversation_id, + "SELECT conversation_id, COUNT(*) AS call_count, COALESCE(SUM(input_tokens), 0) AS total_input, COALESCE(SUM(output_tokens), 0) AS total_output, @@ -1089,10 +1096,9 @@ impl GenAISqliteStore { AND start_timestamp_ns BETWEEN ?2 AND ?3 GROUP BY conversation_id ORDER BY start_ns DESC" - ) + .to_string() } else if start_ns.is_some() { - format!( - "SELECT conversation_id, + "SELECT conversation_id, COUNT(*) AS call_count, COALESCE(SUM(input_tokens), 0) AS total_input, COALESCE(SUM(output_tokens), 0) AS total_output, @@ -1107,10 +1113,9 @@ impl GenAISqliteStore { AND start_timestamp_ns >= ?2 GROUP BY conversation_id ORDER BY start_ns DESC" - ) + .to_string() } else if end_ns.is_some() { - format!( - "SELECT conversation_id, + "SELECT conversation_id, COUNT(*) AS call_count, COALESCE(SUM(input_tokens), 0) AS total_input, COALESCE(SUM(output_tokens), 0) AS total_output, @@ -1125,7 +1130,7 @@ impl GenAISqliteStore { AND start_timestamp_ns <= ?2 GROUP BY conversation_id ORDER BY start_ns DESC" - ) + .to_string() } else { String::from( "SELECT conversation_id, @@ -1141,7 +1146,7 @@ impl GenAISqliteStore { AND session_id = ?1 AND conversation_id IS NOT NULL GROUP BY conversation_id - ORDER BY start_ns DESC" + ORDER BY start_ns DESC", ) }; @@ -1628,18 +1633,18 @@ impl GenAISqliteStore { // Check if it's SQLITE_FULL (extended code 13) if let Some(rusqlite::Error::SqliteFailure(err, _)) = e.downcast_ref::() + && err.extended_code == 13 + && retries < MAX_PRUNE_RETRIES { - if err.extended_code == 13 && retries < MAX_PRUNE_RETRIES { - retries += 1; - log::warn!( - "Database full (SQLITE_FULL), pruning old records (attempt {}/{})", - retries, - MAX_PRUNE_RETRIES - ); - self.prune_old_records()?; - self.checkpoint()?; - continue; - } + retries += 1; + log::warn!( + "Database full (SQLITE_FULL), pruning old records (attempt {}/{})", + retries, + MAX_PRUNE_RETRIES + ); + self.prune_old_records()?; + self.checkpoint()?; + continue; } return Err(e); } @@ -1694,9 +1699,8 @@ impl GenAISqliteStore { // Extract input messages (incremental: latest round only) let input_messages: Option = { - let latest = crate::genai::semantic::latest_round_input_messages( - &call.request.messages, - ); + let latest = + crate::genai::semantic::latest_round_input_messages(&call.request.messages); if latest.is_empty() { None } else { diff --git a/src/agentsight/src/storage/sqlite/http.rs b/src/agentsight/src/storage/sqlite/http.rs index da3349513..622264358 100644 --- a/src/agentsight/src/storage/sqlite/http.rs +++ b/src/agentsight/src/storage/sqlite/http.rs @@ -3,11 +3,11 @@ //! Handles table creation, record insertion, and querying for HTTP request/response records. use anyhow::Result; -use rusqlite::{params, Connection}; +use rusqlite::{Connection, params}; use std::path::Path; -use crate::analyzer::HttpRecord; use super::connection::{create_connection, wal_checkpoint}; +use crate::analyzer::HttpRecord; /// SQLite-based HTTP record store pub struct HttpStore { @@ -99,9 +99,7 @@ impl HttpStore { ); let mut stmt = self.conn.prepare(&sql)?; - let rows = stmt.query_map(params![since_ns as i64], |row| { - Ok(row_to_record(row)) - })?; + let rows = stmt.query_map(params![since_ns as i64], |row| Ok(row_to_record(row)))?; let mut records = Vec::new(); for row in rows { @@ -127,9 +125,7 @@ impl HttpStore { ); let mut stmt = self.conn.prepare(&sql)?; - let rows = stmt.query_map(params![pid], |row| { - Ok(row_to_record(row)) - })?; + let rows = stmt.query_map(params![pid], |row| Ok(row_to_record(row)))?; let mut records = Vec::new(); for row in rows { @@ -155,9 +151,7 @@ impl HttpStore { ); let mut stmt = self.conn.prepare(&sql)?; - let rows = stmt.query_map(params![path_pattern], |row| { - Ok(row_to_record(row)) - })?; + let rows = stmt.query_map(params![path_pattern], |row| Ok(row_to_record(row)))?; let mut records = Vec::new(); for row in rows { @@ -182,10 +176,7 @@ impl HttpStore { /// /// Returns the number of deleted rows. pub fn purge_before(&self, cutoff_ns: u64) -> Result { - let sql = format!( - "DELETE FROM {} WHERE timestamp_ns < ?1", - self.table_name - ); + let sql = format!("DELETE FROM {} WHERE timestamp_ns < ?1", self.table_name); let deleted = self.conn.execute(&sql, params![cutoff_ns as i64])?; Ok(deleted as u64) } diff --git a/src/agentsight/src/storage/sqlite/interruption.rs b/src/agentsight/src/storage/sqlite/interruption.rs index 3abf680eb..6b1211bde 100644 --- a/src/agentsight/src/storage/sqlite/interruption.rs +++ b/src/agentsight/src/storage/sqlite/interruption.rs @@ -1,10 +1,10 @@ //! SQLite storage for interruption_events table. -use rusqlite::{params, Connection}; +use rusqlite::{Connection, params}; use std::sync::Mutex; -use crate::interruption::{InterruptionEvent, InterruptionType}; use super::connection::create_connection; +use crate::interruption::{InterruptionEvent, InterruptionType}; // ─── API response types ──────────────────────────────────────────────────────── @@ -43,7 +43,9 @@ pub struct InterruptionStore { impl InterruptionStore { pub fn new_with_path(path: &std::path::Path) -> Result> { let conn = create_connection(path)?; - let store = InterruptionStore { conn: Mutex::new(conn) }; + let store = InterruptionStore { + conn: Mutex::new(conn), + }; store.init_tables()?; Ok(store) } @@ -75,9 +77,8 @@ impl InterruptionStore { CREATE INDEX IF NOT EXISTS idx_interruption_conversation ON interruption_events(conversation_id);", )?; // Migration: add conversation_id column for existing databases - let _ = conn.execute_batch( - "ALTER TABLE interruption_events ADD COLUMN conversation_id TEXT;", - ); + let _ = + conn.execute_batch("ALTER TABLE interruption_events ADD COLUMN conversation_id TEXT;"); Ok(()) } @@ -110,7 +111,10 @@ impl InterruptionStore { } /// Insert multiple events, ignoring duplicates. - pub fn insert_batch(&self, events: &[InterruptionEvent]) -> Result<(), Box> { + pub fn insert_batch( + &self, + events: &[InterruptionEvent], + ) -> Result<(), Box> { for e in events { self.insert(e)?; } @@ -127,7 +131,9 @@ impl InterruptionStore { AND detail LIKE '%\"oom\":true%'", params![pid, occurred_at_ns], |row| row.get::<_, i64>(0), - ).unwrap_or(0) > 0 + ) + .unwrap_or(0) + > 0 } /// Return the maximum occurred_at_ns of OOM-sourced agent_crash events. @@ -139,7 +145,8 @@ impl InterruptionStore { WHERE interruption_type='agent_crash' AND detail LIKE '%\"oom\":true%'", [], |row| row.get::<_, i64>(0), - ).unwrap_or(0) + ) + .unwrap_or(0) } /// Deduplication check: return true if a row with same call_id + type already exists. @@ -149,7 +156,9 @@ impl InterruptionStore { "SELECT COUNT(*) FROM interruption_events WHERE call_id=?1 AND interruption_type=?2", params![call_id, itype.as_str()], |row| row.get::<_, i64>(0), - ).unwrap_or(0) > 0 + ) + .unwrap_or(0) + > 0 } /// Deduplication check: return true if an unresolved row with same @@ -160,7 +169,12 @@ impl InterruptionStore { /// compared via substring containment. This handles cases where the same /// error appears as a clean message in one call and as raw JSON in another. /// When `error_msg` is None, any unresolved row with same (conversation_id, type) matches. - pub fn exists_for_conversation(&self, conversation_id: &str, itype: &InterruptionType, error_msg: Option<&str>) -> bool { + pub fn exists_for_conversation( + &self, + conversation_id: &str, + itype: &InterruptionType, + error_msg: Option<&str>, + ) -> bool { let conn = self.conn.lock().unwrap(); let mut stmt = match conn.prepare( "SELECT detail FROM interruption_events @@ -184,12 +198,12 @@ impl InterruptionStore { } Some(target) => { // Compare normalized error keys (handles nested JSON vs clean message) - if let Some(ref detail_str) = detail_opt { - if let Ok(v) = serde_json::from_str::(detail_str) { - let stored_error = v.get("error").and_then(|e| e.as_str()).unwrap_or(""); - if errors_match(stored_error, target) { - return true; - } + if let Some(ref detail_str) = detail_opt + && let Ok(v) = serde_json::from_str::(detail_str) + { + let stored_error = v.get("error").and_then(|e| e.as_str()).unwrap_or(""); + if errors_match(stored_error, target) { + return true; } } } @@ -211,6 +225,7 @@ impl InterruptionStore { // ─── Query ────────────────────────────────────────────────────────────── /// List interruptions within a time range. + #[allow(clippy::too_many_arguments)] pub fn list( &self, start_ns: i64, @@ -224,13 +239,9 @@ impl InterruptionStore { let conn = self.conn.lock().unwrap(); // Build dynamic WHERE clause - let mut conditions = vec![ - "occurred_at_ns BETWEEN ?1 AND ?2".to_string(), - ]; - let mut args: Vec> = vec![ - Box::new(start_ns), - Box::new(end_ns), - ]; + let mut conditions = vec!["occurred_at_ns BETWEEN ?1 AND ?2".to_string()]; + let mut args: Vec> = + vec![Box::new(start_ns), Box::new(end_ns)]; let mut idx = 3usize; if let Some(a) = agent_name { @@ -267,7 +278,8 @@ impl InterruptionStore { ); args.push(Box::new(limit)); - let params_refs: Vec<&dyn rusqlite::types::ToSql> = args.iter().map(|b| b.as_ref()).collect(); + let params_refs: Vec<&dyn rusqlite::types::ToSql> = + args.iter().map(|b| b.as_ref()).collect(); let mut stmt = conn.prepare(&sql)?; let rows = stmt.query_map(params_refs.as_slice(), |row| { Ok(InterruptionRecord { @@ -287,12 +299,17 @@ impl InterruptionStore { }) })?; let mut result = Vec::new(); - for row in rows { result.push(row?); } + for row in rows { + result.push(row?); + } Ok(result) } /// Get a single interruption by ID. - pub fn get_by_id(&self, interruption_id: &str) -> Result, Box> { + pub fn get_by_id( + &self, + interruption_id: &str, + ) -> Result, Box> { let conn = self.conn.lock().unwrap(); let mut stmt = conn.prepare( "SELECT id, interruption_id, session_id, trace_id, conversation_id, call_id, pid, agent_name, @@ -320,7 +337,10 @@ impl InterruptionStore { } /// Get all interruptions for a session. - pub fn list_by_session(&self, session_id: &str) -> Result, Box> { + pub fn list_by_session( + &self, + session_id: &str, + ) -> Result, Box> { let conn = self.conn.lock().unwrap(); let mut stmt = conn.prepare( "SELECT id, interruption_id, session_id, trace_id, conversation_id, call_id, pid, agent_name, @@ -347,11 +367,16 @@ impl InterruptionStore { }) })?; let mut result = Vec::new(); - for row in rows { result.push(row?); } + for row in rows { + result.push(row?); + } Ok(result) } - pub fn list_by_conversation(&self, conversation_id: &str) -> Result, Box> { + pub fn list_by_conversation( + &self, + conversation_id: &str, + ) -> Result, Box> { let conn = self.conn.lock().unwrap(); let mut stmt = conn.prepare( "SELECT id, interruption_id, session_id, trace_id, conversation_id, call_id, pid, agent_name, @@ -378,7 +403,9 @@ impl InterruptionStore { }) })?; let mut result = Vec::new(); - for row in rows { result.push(row?); } + for row in rows { + result.push(row?); + } Ok(result) } @@ -404,12 +431,15 @@ impl InterruptionStore { }) })?; let mut result = Vec::new(); - for row in rows { result.push(row?); } + for row in rows { + result.push(row?); + } Ok(result) } /// Count unresolved interruptions grouped by (session_id, severity, type). /// Returns detailed rows for building per-severity badges with type tooltips. + #[allow(clippy::type_complexity)] pub fn count_unresolved_by_session_detailed( &self, start_ns: i64, @@ -434,11 +464,14 @@ impl InterruptionStore { )) })?; let mut result = Vec::new(); - for row in rows { result.push(row?); } + for row in rows { + result.push(row?); + } Ok(result) } /// Count unresolved interruptions grouped by (conversation_id, severity, type). + #[allow(clippy::type_complexity)] pub fn count_unresolved_by_conversation_detailed( &self, start_ns: i64, @@ -463,7 +496,9 @@ impl InterruptionStore { )) })?; let mut result = Vec::new(); - for row in rows { result.push(row?); } + for row in rows { + result.push(row?); + } Ok(result) } @@ -489,10 +524,10 @@ fn normalize_error_key(raw: &str) -> String { return msg.to_lowercase(); } // Try to find JSON embedded in the string (e.g. "curl...{\"error\":{...}}") - if let Some(brace_start) = trimmed.find('{') { - if let Some(msg) = extract_message_from_json(&trimmed[brace_start..]) { - return msg.to_lowercase(); - } + if let Some(brace_start) = trimmed.find('{') + && let Some(msg) = extract_message_from_json(&trimmed[brace_start..]) + { + return msg.to_lowercase(); } trimmed.to_lowercase() } @@ -503,7 +538,11 @@ fn normalize_error_key(raw: &str) -> String { fn extract_message_from_json(s: &str) -> Option { let v: serde_json::Value = serde_json::from_str(s).ok()?; // Try "error.message" - if let Some(msg) = v.get("error").and_then(|e| e.get("message")).and_then(|m| m.as_str()) { + if let Some(msg) = v + .get("error") + .and_then(|e| e.get("message")) + .and_then(|m| m.as_str()) + { return Some(msg.to_string()); } // Try top-level "message" diff --git a/src/agentsight/src/storage/sqlite/mod.rs b/src/agentsight/src/storage/sqlite/mod.rs index 9c7ea9fef..4c1b0178a 100644 --- a/src/agentsight/src/storage/sqlite/mod.rs +++ b/src/agentsight/src/storage/sqlite/mod.rs @@ -21,15 +21,14 @@ pub use audit::{AuditStore, SqliteStore}; // Re-export token storage pub use token::{ - TokenStore, TokenQuery, - TimePeriod, TokenQueryResult, TokenBreakdown, TokenComparison, Trend, + TimePeriod, TokenBreakdown, TokenComparison, TokenQuery, TokenQueryResult, TokenStore, Trend, format_tokens, format_tokens_with_commas, }; // Re-export token consumption storage pub use token_consumption::{ - TokenConsumptionStore, TokenConsumptionRecord, - TokenConsumptionFilter, TokenConsumptionQueryResult, + TokenConsumptionFilter, TokenConsumptionQueryResult, TokenConsumptionRecord, + TokenConsumptionStore, }; // Re-export HTTP storage @@ -39,9 +38,7 @@ pub use http::HttpStore; pub use genai::{GenAISqliteStore, PendingCallInfo, SseEnrichment}; // Re-export Interruption SQLite storage -pub use interruption::{ - InterruptionStore, InterruptionRecord, InterruptionTypeStat, -}; +pub use interruption::{InterruptionRecord, InterruptionStore, InterruptionTypeStat}; // Re-export connection utilities pub use connection::{create_connection, default_base_path}; diff --git a/src/agentsight/src/storage/sqlite/token.rs b/src/agentsight/src/storage/sqlite/token.rs index cfc2dc377..5d4ab4d9d 100644 --- a/src/agentsight/src/storage/sqlite/token.rs +++ b/src/agentsight/src/storage/sqlite/token.rs @@ -2,14 +2,14 @@ //! //! Uses SQLite for persistent storage of token usage records. +use chrono::{Datelike, Utc}; +use rusqlite::{Connection, params}; +use serde::{Deserialize, Serialize}; use std::path::PathBuf; use std::time::{SystemTime, UNIX_EPOCH}; -use chrono::{Utc, Datelike}; -use serde::{Deserialize, Serialize}; -use rusqlite::{params, Connection}; -use crate::analyzer::TokenRecord; use super::connection::{create_connection, default_base_path, wal_checkpoint}; +use crate::analyzer::TokenRecord; /// Time period for queries #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] @@ -40,7 +40,7 @@ impl TimePeriod { pub fn time_range(&self) -> (u64, u64) { let now = Utc::now(); let now_naive = now.naive_utc(); - + let (start, end) = match self { TimePeriod::Today => { let start = now_naive.date().and_hms_opt(0, 0, 0).unwrap(); @@ -83,13 +83,13 @@ impl TimePeriod { (start, end) } }; - + let start_ns = start.and_utc().timestamp_nanos_opt().unwrap_or(0) as u64; let end_ns = end.and_utc().timestamp_nanos_opt().unwrap_or(0) as u64; - + (start_ns, end_ns) } - + /// Get previous period for comparison pub fn previous_period(&self) -> TimePeriod { match self { @@ -144,12 +144,12 @@ impl TokenQueryResult { pub fn formatted_total(&self) -> String { format_tokens(self.total_tokens) } - + /// Format input tokens with K/M suffix pub fn formatted_input(&self) -> String { format_tokens(self.input_tokens) } - + /// Format output tokens with K/M suffix pub fn formatted_output(&self) -> String { format_tokens(self.output_tokens) @@ -175,7 +175,7 @@ impl TokenComparison { let sign = if self.change >= 0 { "+" } else { "" }; let change_formatted = format_tokens(self.change.unsigned_abs()); let percent = format!("{:.0}%", self.change_percent.abs()); - + if self.change >= 0 { format!("{}{} (+{}%)", sign, change_formatted, percent) } else { @@ -233,10 +233,10 @@ impl TokenStore { /// Create a new token store with custom table name pub fn with_table(path: impl Into, table_name: &str) -> Self { let path = path.into(); - let conn = create_connection(&path) - .expect("Failed to open SQLite database for token store"); + let conn = + create_connection(&path).expect("Failed to open SQLite database for token store"); let table_name = table_name.to_string(); - + // Create table if not exists let create_table_sql = format!( "CREATE TABLE IF NOT EXISTS {} ( @@ -258,31 +258,39 @@ impl TokenStore { ); conn.execute(&create_table_sql, []) .expect("Failed to create token table"); - + // Create index on timestamp for efficient range queries conn.execute( - &format!("CREATE INDEX IF NOT EXISTS idx_{}_timestamp ON {}(timestamp_ns)", table_name, table_name), + &format!( + "CREATE INDEX IF NOT EXISTS idx_{}_timestamp ON {}(timestamp_ns)", + table_name, table_name + ), [], - ).expect("Failed to create timestamp index"); - + ) + .expect("Failed to create timestamp index"); + // Create index on agent for breakdown queries conn.execute( - &format!("CREATE INDEX IF NOT EXISTS idx_{}_agent ON {}(agent)", table_name, table_name), + &format!( + "CREATE INDEX IF NOT EXISTS idx_{}_agent ON {}(agent)", + table_name, table_name + ), [], - ).expect("Failed to create agent index"); - + ) + .expect("Failed to create agent index"); + TokenStore { conn, table_name } } - + /// Get default storage path pub fn default_path() -> PathBuf { default_base_path().join("tokens.db") } - + /// Insert a token record (unified interface, matches AuditStore) pub fn insert(&self, record: &TokenRecord) -> anyhow::Result { let timestamp_ns = record.timestamp_ns; - + let sql = format!( "INSERT INTO {} ( timestamp_ns, pid, comm, agent, model, provider, @@ -291,34 +299,36 @@ impl TokenStore { ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12)", self.table_name ); - self.conn.execute( - &sql, - params![ - timestamp_ns as i64, - record.pid as i64, - record.comm, - record.agent, - record.model, - record.provider, - record.input_tokens as i64, - record.output_tokens as i64, - record.cache_creation_tokens.map(|v| v as i64), - record.cache_read_tokens.map(|v| v as i64), - record.request_id, - record.endpoint, - ], - ).map_err(|e| anyhow::anyhow!("Failed to insert token record: {}", e))?; - + self.conn + .execute( + &sql, + params![ + timestamp_ns as i64, + record.pid as i64, + record.comm, + record.agent, + record.model, + record.provider, + record.input_tokens as i64, + record.output_tokens as i64, + record.cache_creation_tokens.map(|v| v as i64), + record.cache_read_tokens.map(|v| v as i64), + record.request_id, + record.endpoint, + ], + ) + .map_err(|e| anyhow::anyhow!("Failed to insert token record: {}", e))?; + Ok(self.conn.last_insert_rowid()) } - + /// Add a token record (legacy method, kept for backward compatibility) pub fn add(&mut self, record: TokenRecord) -> Result { let now = SystemTime::now() .duration_since(UNIX_EPOCH) .map(|d| d.as_nanos() as u64) .unwrap_or(0); - + let sql = format!( "INSERT INTO {} ( timestamp_ns, pid, comm, agent, model, provider, @@ -344,10 +354,10 @@ impl TokenStore { record.endpoint, ], )?; - + Ok(self.conn.last_insert_rowid()) } - + /// Get all records (for compatibility, but not recommended for large datasets) pub fn all(&self) -> Vec { let sql = format!( @@ -357,10 +367,11 @@ impl TokenStore { FROM {} ORDER BY timestamp_ns DESC", self.table_name ); - let mut stmt = self.conn + let mut stmt = self + .conn .prepare(&sql) .expect("Failed to prepare statement"); - + stmt.query_map([], |row| { Ok(TokenRecord { id: row.get(0)?, @@ -384,12 +395,12 @@ impl TokenStore { .filter_map(|r| r.ok()) .collect() } - + /// Get records in time range pub fn by_time_range(&self, start_ns: u64, end_ns: u64) -> Vec { self.by_time_range_owned(start_ns, end_ns) } - + /// Get owned records in time range pub fn by_time_range_owned(&self, start_ns: u64, end_ns: u64) -> Vec { let sql = format!( @@ -401,10 +412,11 @@ impl TokenStore { ORDER BY timestamp_ns DESC", self.table_name ); - let mut stmt = self.conn + let mut stmt = self + .conn .prepare(&sql) .expect("Failed to prepare statement"); - + stmt.query_map(params![start_ns as i64, end_ns as i64], |row| { Ok(TokenRecord { id: row.get(0)?, @@ -428,30 +440,35 @@ impl TokenStore { .filter_map(|r| r.ok()) .collect() } - + /// Get records for last N hours pub fn by_last_hours(&self, hours: u64) -> Vec { let now = SystemTime::now() .duration_since(UNIX_EPOCH) .map(|d| d.as_nanos() as u64) .unwrap_or(0); - + let hours_ns = hours * 3600 * 1_000_000_000; let start_ns = now.saturating_sub(hours_ns); - + self.by_time_range_owned(start_ns, now) } - + /// Clear all records pub fn clear(&mut self) -> Result<(), rusqlite::Error> { - self.conn.execute(&format!("DELETE FROM {}", self.table_name), [])?; + self.conn + .execute(&format!("DELETE FROM {}", self.table_name), [])?; Ok(()) } - + /// Get record count pub fn count(&self) -> u64 { self.conn - .query_row(&format!("SELECT COUNT(*) FROM {}", self.table_name), [], |row| row.get::<_, i64>(0)) + .query_row( + &format!("SELECT COUNT(*) FROM {}", self.table_name), + [], + |row| row.get::<_, i64>(0), + ) .unwrap_or(0) as u64 } @@ -459,18 +476,17 @@ impl TokenStore { /// /// Returns the number of deleted rows. pub fn purge_before(&self, cutoff_ns: u64) -> anyhow::Result { - let sql = format!( - "DELETE FROM {} WHERE timestamp_ns < ?1", - self.table_name - ); - let deleted = self.conn.execute(&sql, params![cutoff_ns as i64]) + let sql = format!("DELETE FROM {} WHERE timestamp_ns < ?1", self.table_name); + let deleted = self + .conn + .execute(&sql, params![cutoff_ns as i64]) .map_err(|e| anyhow::anyhow!("Failed to purge token records: {}", e))?; Ok(deleted as u64) } /// Execute WAL checkpoint to flush WAL data back to the main database file pub fn checkpoint(&self) -> anyhow::Result<()> { - wal_checkpoint(&self.conn).map_err(Into::into) + wal_checkpoint(&self.conn) } } @@ -484,20 +500,20 @@ impl<'a> TokenQuery<'a> { pub fn new(store: &'a TokenStore) -> Self { TokenQuery { store } } - + /// Query by time period pub fn by_period(&self, period: TimePeriod) -> TokenQueryResult { let (start_ns, end_ns) = period.time_range(); let records = self.store.by_time_range(start_ns, end_ns); self.build_result(records, period.to_string()) } - + /// Query last N hours pub fn by_hours(&self, hours: u64) -> TokenQueryResult { let records = self.store.by_last_hours(hours); self.build_result(records, format!("最近 {} 小时", hours)) } - + /// Query with comparison pub fn by_period_with_compare(&self, period: TimePeriod) -> TokenQueryResult { let mut result = self.by_period(period); @@ -530,38 +546,41 @@ impl<'a> TokenQuery<'a> { result } - + /// Query with breakdown by agent pub fn by_period_with_breakdown(&self, period: TimePeriod) -> TokenQueryResult { let mut result = self.by_period(period); result.breakdown = self.compute_breakdown(period); result } - + /// Query with comparison and breakdown pub fn full_query(&self, period: TimePeriod) -> TokenQueryResult { let mut result = self.by_period_with_compare(period); result.breakdown = self.compute_breakdown(period); result } - + /// Query hours with comparison pub fn by_hours_with_compare(&self, hours: u64) -> TokenQueryResult { let mut result = self.by_hours(hours); // Get previous period data let prev_records = self.store.by_last_hours(hours * 2); - let prev_records: Vec<_> = prev_records.into_iter().filter(|r| { - // Get records from the earlier half - let now = SystemTime::now() - .duration_since(UNIX_EPOCH) - .map(|d| d.as_nanos() as u64) - .unwrap_or(0); - let hours_ns = hours * 3600 * 1_000_000_000; - let start_ns = now.saturating_sub(hours_ns * 2); - let mid_ns = now.saturating_sub(hours_ns); - r.timestamp_ns >= start_ns && r.timestamp_ns < mid_ns - }).collect(); + let prev_records: Vec<_> = prev_records + .into_iter() + .filter(|r| { + // Get records from the earlier half + let now = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_nanos() as u64) + .unwrap_or(0); + let hours_ns = hours * 3600 * 1_000_000_000; + let start_ns = now.saturating_sub(hours_ns * 2); + let mid_ns = now.saturating_sub(hours_ns); + r.timestamp_ns >= start_ns && r.timestamp_ns < mid_ns + }) + .collect(); let prev_total: u64 = prev_records.iter().map(|r| r.total_tokens()).sum(); @@ -589,14 +608,14 @@ impl<'a> TokenQuery<'a> { result } - + /// Build result from records fn build_result(&self, records: Vec, period: String) -> TokenQueryResult { let input_tokens: u64 = records.iter().map(|r| r.input_tokens).sum(); let output_tokens: u64 = records.iter().map(|r| r.output_tokens).sum(); let total_tokens = input_tokens + output_tokens; let request_count = records.len() as u64; - + TokenQueryResult { period, input_tokens, @@ -607,30 +626,28 @@ impl<'a> TokenQuery<'a> { breakdown: Vec::new(), } } - + /// Compute breakdown by agent fn compute_breakdown(&self, period: TimePeriod) -> Vec { let (start_ns, end_ns) = period.time_range(); let records = self.store.by_time_range(start_ns, end_ns); - + let total_tokens: u64 = records.iter().map(|r| r.total_tokens()).sum(); - + // Group by agent name (or comm if no agent) - let mut agent_totals: std::collections::HashMap = + let mut agent_totals: std::collections::HashMap = std::collections::HashMap::new(); - + for record in records { - let name = record.agent.as_ref() - .unwrap_or(&record.comm) - .clone(); - + let name = record.agent.as_ref().unwrap_or(&record.comm).clone(); + let entry = agent_totals.entry(name).or_insert((0, 0, 0, 0)); entry.0 += record.total_tokens(); entry.1 += record.input_tokens; entry.2 += record.output_tokens; entry.3 += 1; } - + // Convert to breakdown let mut breakdown: Vec = agent_totals .into_iter() @@ -640,7 +657,7 @@ impl<'a> TokenQuery<'a> { } else { 0.0 }; - + TokenBreakdown { name, total_tokens: total, @@ -651,7 +668,7 @@ impl<'a> TokenQuery<'a> { } }) .collect(); - + // Sort by total tokens descending breakdown.sort_by(|a, b| b.total_tokens.cmp(&a.total_tokens)); breakdown @@ -661,55 +678,71 @@ impl<'a> TokenQuery<'a> { #[cfg(test)] mod tests { use super::*; - + #[test] fn test_time_period_range() { let (start, end) = TimePeriod::Today.time_range(); assert!(start < end); assert!(end > 0); } - + #[test] fn test_format_tokens() { assert_eq!(format_tokens(500), "500"); assert_eq!(format_tokens(1500), "1.5K"); assert_eq!(format_tokens(1_500_000), "1.5M"); } - + #[test] fn test_format_tokens_with_commas() { assert_eq!(format_tokens_with_commas(1000), "1,000"); assert_eq!(format_tokens_with_commas(125000), "125,000"); } - + #[test] fn test_token_store() { let mut store = TokenStore::new("/tmp/test_tokens.db"); - + let record = TokenRecord::new(1234, "python".to_string(), "openai".to_string(), 100, 50); let id = store.add(record).unwrap(); assert!(id > 0); - + let records = store.all(); assert!(!records.is_empty()); - + // Cleanup std::fs::remove_file("/tmp/test_tokens.db").ok(); } - + #[test] fn test_token_query() { let mut store = TokenStore::new("/tmp/test_tokens_query.db"); - + // Add some records - store.add(TokenRecord::new(1234, "python".to_string(), "openai".to_string(), 100, 50)).unwrap(); - store.add(TokenRecord::new(1234, "python".to_string(), "anthropic".to_string(), 200, 100)).unwrap(); - + store + .add(TokenRecord::new( + 1234, + "python".to_string(), + "openai".to_string(), + 100, + 50, + )) + .unwrap(); + store + .add(TokenRecord::new( + 1234, + "python".to_string(), + "anthropic".to_string(), + 200, + 100, + )) + .unwrap(); + let query = TokenQuery::new(&store); let result = query.by_period(TimePeriod::Today); - + assert!(result.total_tokens > 0); - + // Cleanup std::fs::remove_file("/tmp/test_tokens_query.db").ok(); } diff --git a/src/agentsight/src/storage/sqlite/token_consumption.rs b/src/agentsight/src/storage/sqlite/token_consumption.rs index a78abffc4..7df3b99ec 100644 --- a/src/agentsight/src/storage/sqlite/token_consumption.rs +++ b/src/agentsight/src/storage/sqlite/token_consumption.rs @@ -6,11 +6,11 @@ use std::collections::HashMap; use std::path::PathBuf; -use rusqlite::{params, Connection}; +use rusqlite::{Connection, params}; use serde::{Deserialize, Serialize}; -use crate::analyzer::TokenConsumptionBreakdown; use super::connection::{create_connection, default_base_path, wal_checkpoint}; +use crate::analyzer::TokenConsumptionBreakdown; /// A row stored in the token_consumption table (excludes per_message and output_per_block) #[derive(Debug, Clone, Serialize, Deserialize)] @@ -156,10 +156,10 @@ impl TokenConsumptionStore { pid: u32, comm: &str, ) -> anyhow::Result { - let by_role_json = serde_json::to_string(&breakdown.by_role) - .unwrap_or_else(|_| "{}".to_string()); - let output_by_type_json = serde_json::to_string(&breakdown.output_by_type) - .unwrap_or_else(|_| "{}".to_string()); + let by_role_json = + serde_json::to_string(&breakdown.by_role).unwrap_or_else(|_| "{}".to_string()); + let output_by_type_json = + serde_json::to_string(&breakdown.output_by_type).unwrap_or_else(|_| "{}".to_string()); let sql = format!( "INSERT INTO {} ( @@ -171,23 +171,24 @@ impl TokenConsumptionStore { self.table_name ); - self.conn.execute( - &sql, - params![ - timestamp_ns as i64, - pid as i64, - comm, - breakdown.provider, - breakdown.model, - breakdown.total_input_tokens as i64, - breakdown.total_output_tokens as i64, - breakdown.tools_tokens as i64, - breakdown.system_prompt_tokens as i64, - by_role_json, - output_by_type_json, - ], - ) - .map_err(|e| anyhow::anyhow!("Failed to insert token_consumption record: {}", e))?; + self.conn + .execute( + &sql, + params![ + timestamp_ns as i64, + pid as i64, + comm, + breakdown.provider, + breakdown.model, + breakdown.total_input_tokens as i64, + breakdown.total_output_tokens as i64, + breakdown.tools_tokens as i64, + breakdown.system_prompt_tokens as i64, + by_role_json, + output_by_type_json, + ], + ) + .map_err(|e| anyhow::anyhow!("Failed to insert token_consumption record: {}", e))?; Ok(self.conn.last_insert_rowid()) } @@ -195,7 +196,10 @@ impl TokenConsumptionStore { /// Query records with the given filter. /// /// Returns individual rows sorted by timestamp descending. - pub fn query(&self, filter: &TokenConsumptionFilter) -> anyhow::Result> { + pub fn query( + &self, + filter: &TokenConsumptionFilter, + ) -> anyhow::Result> { let mut conditions: Vec = Vec::new(); let mut bind_idx = 1usize; @@ -327,7 +331,11 @@ impl TokenConsumptionStore { } /// Query records in a time range - pub fn by_time_range(&self, start_ns: u64, end_ns: u64) -> anyhow::Result> { + pub fn by_time_range( + &self, + start_ns: u64, + end_ns: u64, + ) -> anyhow::Result> { self.query(&TokenConsumptionFilter { start_ns: Some(start_ns), end_ns: Some(end_ns), @@ -346,7 +354,7 @@ impl TokenConsumptionStore { /// Execute WAL checkpoint to flush WAL data back to the main database file pub fn checkpoint(&self) -> anyhow::Result<()> { - wal_checkpoint(&self.conn).map_err(Into::into) + wal_checkpoint(&self.conn) } /// Number of stored records diff --git a/src/agentsight/src/storage/unified.rs b/src/agentsight/src/storage/unified.rs index ada5bfe5d..c5dd958af 100644 --- a/src/agentsight/src/storage/unified.rs +++ b/src/agentsight/src/storage/unified.rs @@ -40,9 +40,9 @@ use std::path::PathBuf; use std::sync::atomic::{AtomicU64, Ordering}; use std::time::{SystemTime, UNIX_EPOCH}; +use super::sqlite::connection::default_base_path; +use super::sqlite::{AuditStore, HttpStore, TokenConsumptionStore, TokenStore}; use crate::analyzer::AnalysisResult; -use super::sqlite::{AuditStore, TokenStore, HttpStore, TokenConsumptionStore}; -use super::sqlite::connection::{default_base_path}; /// Storage backend type #[derive(Debug, Clone, Default)] @@ -221,23 +221,21 @@ impl Storage { Ok(0) } AnalysisResult::Http(record) => self.http_store.insert(record), - AnalysisResult::TokenConsumption(breakdown) => { - self.token_consumption_store.insert( - breakdown, - breakdown.timestamp_ns, - breakdown.pid, - &breakdown.comm, - ) - } + AnalysisResult::TokenConsumption(breakdown) => self.token_consumption_store.insert( + breakdown, + breakdown.timestamp_ns, + breakdown.pid, + &breakdown.comm, + ), }?; // Auto-purge check: trigger every `purge_interval` inserts if self.purge_interval > 0 && self.retention_days > 0 { let count = self.insert_count.fetch_add(1, Ordering::Relaxed) + 1; - if count % self.purge_interval == 0 { - if let Err(e) = self.purge_expired() { - log::warn!("Auto-purge failed: {}", e); - } + if count % self.purge_interval == 0 + && let Err(e) = self.purge_expired() + { + log::warn!("Auto-purge failed: {}", e); } } @@ -274,8 +272,12 @@ impl Storage { if total_deleted > 0 { log::info!( "Purged {} expired records (retention={}d, audit={}, token={}, http={}, consumption={})", - total_deleted, self.retention_days, - audit_deleted, token_deleted, http_deleted, consumption_deleted, + total_deleted, + self.retention_days, + audit_deleted, + token_deleted, + http_deleted, + consumption_deleted, ); } diff --git a/src/agentsight/src/tokenizer/llm_tok.rs b/src/agentsight/src/tokenizer/llm_tok.rs index f20f3fc0e..870d6051e 100644 --- a/src/agentsight/src/tokenizer/llm_tok.rs +++ b/src/agentsight/src/tokenizer/llm_tok.rs @@ -4,12 +4,15 @@ //! [`ChatTemplate`] trait, replacing the previous separate `QwenTokenizer` and //! `QwenChatTemplate` types. -use anyhow::{anyhow, Result}; +use anyhow::{Result, anyhow}; use serde_json::Value; use std::path::Path; use std::sync::Arc; -use llm_tokenizer::{Decoder as _, Encoder as _, HuggingFaceTokenizer, TokenizerTrait, chat_template::ChatTemplateParams}; +use llm_tokenizer::{ + Decoder as _, Encoder as _, HuggingFaceTokenizer, TokenizerTrait, + chat_template::ChatTemplateParams, +}; /// Unified tokenizer + chat template adapter wrapping `llm-tokenizer` crate. /// @@ -38,24 +41,31 @@ impl LlmTokenizer { /// # Arguments /// * `tokenizer_path` - Path to the tokenizer.json file /// * `config_path` - Path to the tokenizer_config.json file (for chat template) - pub fn from_file>( - tokenizer_path: P, - config_path: P, - ) -> Result { + pub fn from_file>(tokenizer_path: P, config_path: P) -> Result { let tokenizer_path = tokenizer_path.as_ref(); let config_path = config_path.as_ref(); - let tokenizer_str = tokenizer_path.to_str() + let tokenizer_str = tokenizer_path + .to_str() .ok_or_else(|| anyhow!("Tokenizer path is not valid UTF-8: {:?}", tokenizer_path))?; - let config_str = config_path.to_str() + let config_str = config_path + .to_str() .ok_or_else(|| anyhow!("Config path is not valid UTF-8: {:?}", config_path))?; // Use HuggingFaceTokenizer with explicit chat template config - let tokenizer = HuggingFaceTokenizer::from_file_with_chat_template(tokenizer_str, Some(config_str)) - .map_err(|e| anyhow!("Failed to load tokenizer from '{}': {}", tokenizer_path.display(), e))?; + let tokenizer = + HuggingFaceTokenizer::from_file_with_chat_template(tokenizer_str, Some(config_str)) + .map_err(|e| { + anyhow!( + "Failed to load tokenizer from '{}': {}", + tokenizer_path.display(), + e + ) + })?; Ok(Self { inner: Arc::new(tokenizer), - model_name: tokenizer_path.file_stem() + model_name: tokenizer_path + .file_stem() .and_then(|s| s.to_str()) .unwrap_or("unknown") .to_string(), @@ -64,14 +74,18 @@ impl LlmTokenizer { /// Encode text with special tokens. pub fn encode_with_special_tokens(&self, text: &str) -> Result> { - let encoding = self.inner.encode(text, true) + let encoding = self + .inner + .encode(text, true) .map_err(|e| anyhow!("Failed to encode text with special tokens: {}", e))?; Ok(encoding.token_ids().to_vec()) } /// Encode text without special tokens. pub fn encode_without_special_tokens(&self, text: &str) -> Result> { - let encoding = self.inner.encode(text, false) + let encoding = self + .inner + .encode(text, false) .map_err(|e| anyhow!("Failed to encode text: {}", e))?; Ok(encoding.token_ids().to_vec()) } @@ -89,20 +103,23 @@ impl LlmTokenizer { add_generation_prompt: bool, ) -> Result { // Use the llm-tokenizer crate's built-in chat template support - self.inner.apply_chat_template( - messages, - ChatTemplateParams { - add_generation_prompt, - tools, - ..Default::default() - }, - ) - .map_err(|e| anyhow!("Failed to apply chat template: {}", e)) + self.inner + .apply_chat_template( + messages, + ChatTemplateParams { + add_generation_prompt, + tools, + ..Default::default() + }, + ) + .map_err(|e| anyhow!("Failed to apply chat template: {}", e)) } /// Count tokens in text pub fn count(&self, text: &str) -> Result { - let encoding = self.inner.encode(text, false) + let encoding = self + .inner + .encode(text, false) .map_err(|e| anyhow!("Failed to encode text: {}", e))?; Ok(encoding.token_ids().len()) } @@ -114,7 +131,8 @@ impl LlmTokenizer { /// Decode token IDs back to text pub fn decode(&self, tokens: &[u32]) -> Result { - self.inner.decode(tokens, false) + self.inner + .decode(tokens, false) .map_err(|e| anyhow!("Failed to decode tokens: {}", e)) } @@ -125,13 +143,19 @@ impl LlmTokenizer { /// Count tokens with special token recognition pub fn count_with_special_tokens(&self, text: &str) -> Result { - let encoding = self.inner.encode(text, true) + let encoding = self + .inner + .encode(text, true) .map_err(|e| anyhow!("Failed to encode text with special tokens: {}", e))?; Ok(encoding.token_ids().len()) } /// Apply chat template to JSON messages - pub fn apply_chat_template(&self, messages: &[Value], add_generation_prompt: bool) -> Result { + pub fn apply_chat_template( + &self, + messages: &[Value], + add_generation_prompt: bool, + ) -> Result { self.do_apply_chat_template(messages, None, add_generation_prompt) } diff --git a/src/agentsight/src/tokenizer/mod.rs b/src/agentsight/src/tokenizer/mod.rs index 4d8955b56..6bce8ffd8 100644 --- a/src/agentsight/src/tokenizer/mod.rs +++ b/src/agentsight/src/tokenizer/mod.rs @@ -10,7 +10,4 @@ pub mod multi_model; // Re-export types pub use llm_tok::LlmTokenizer; pub use model_mapping::map_to_hf_model_id; -pub use multi_model::{ - MultiModelTokenizer, TokenizerEntry, - get_global_tokenizer, -}; +pub use multi_model::{MultiModelTokenizer, TokenizerEntry, get_global_tokenizer}; diff --git a/src/agentsight/src/tokenizer/model_mapping.rs b/src/agentsight/src/tokenizer/model_mapping.rs index 871a198a9..23a7fef16 100644 --- a/src/agentsight/src/tokenizer/model_mapping.rs +++ b/src/agentsight/src/tokenizer/model_mapping.rs @@ -7,13 +7,13 @@ //! (e.g., "gpt-4", "claude-3-opus", "qwen-turbo") that don't match HuggingFace's //! naming convention (e.g., "openai/gpt-4", "anthropic/claude-3-opus"). -use std::collections::HashMap; use once_cell::sync::Lazy; +use std::collections::HashMap; /// Global model name mapping table static MODEL_MAPPING: Lazy> = Lazy::new(|| { let mut m = HashMap::new(); - + // ============== OpenAI Models ============== // GPT-4 series m.insert("gpt-4", "openai/gpt-4"); @@ -23,17 +23,17 @@ static MODEL_MAPPING: Lazy> = Lazy::new(|| { m.insert("gpt-4o-mini", "openai/gpt-4o-mini"); m.insert("gpt-4-32k", "openai/gpt-4-32k"); m.insert("gpt-4-vision-preview", "openai/gpt-4-vision"); - + // GPT-3.5 series m.insert("gpt-3.5-turbo", "openai/gpt-3.5-turbo"); m.insert("gpt-3.5-turbo-16k", "openai/gpt-3.5-turbo-16k"); m.insert("gpt-3.5-turbo-instruct", "openai/gpt-3.5-turbo-instruct"); - + // o1 series m.insert("o1", "openai/o1"); m.insert("o1-mini", "openai/o1-mini"); m.insert("o1-preview", "openai/o1-preview"); - + // ============== Anthropic Models ============== // Claude 3 series m.insert("claude-3-opus", "anthropic/claude-3-opus"); @@ -42,18 +42,18 @@ static MODEL_MAPPING: Lazy> = Lazy::new(|| { m.insert("claude-3-sonnet-20240229", "anthropic/claude-3-sonnet"); m.insert("claude-3-haiku", "anthropic/claude-3-haiku"); m.insert("claude-3-haiku-20240307", "anthropic/claude-3-haiku"); - + // Claude 3.5 series m.insert("claude-3.5-sonnet", "anthropic/claude-3.5-sonnet"); m.insert("claude-3-5-sonnet", "anthropic/claude-3.5-sonnet"); m.insert("claude-3-5-sonnet-20240620", "anthropic/claude-3.5-sonnet"); m.insert("claude-3.5-haiku", "anthropic/claude-3.5-haiku"); m.insert("claude-3-5-haiku", "anthropic/claude-3.5-haiku"); - + // Claude 3.7 series m.insert("claude-3.7-sonnet", "anthropic/claude-3.7-sonnet"); m.insert("claude-3-7-sonnet", "anthropic/claude-3.7-sonnet"); - + // ============== Qwen (Alibaba) Models ============== // Qwen 2.5 series m.insert("qwen2.5-7b-instruct", "Qwen/Qwen2.5-7B-Instruct"); @@ -61,34 +61,37 @@ static MODEL_MAPPING: Lazy> = Lazy::new(|| { m.insert("qwen2.5-32b-instruct", "Qwen/Qwen2.5-32B-Instruct"); m.insert("qwen2.5-72b-instruct", "Qwen/Qwen2.5-72B-Instruct"); m.insert("qwen2.5-math-7b-instruct", "Qwen/Qwen2.5-Math-7B-Instruct"); - m.insert("qwen2.5-coder-7b-instruct", "Qwen/Qwen2.5-Coder-7B-Instruct"); - + m.insert( + "qwen2.5-coder-7b-instruct", + "Qwen/Qwen2.5-Coder-7B-Instruct", + ); + // Qwen 2 series m.insert("qwen2-7b-instruct", "Qwen/Qwen2-7B-Instruct"); m.insert("qwen2-72b-instruct", "Qwen/Qwen2-72B-Instruct"); - + // Qwen 1.5 series m.insert("qwen1.5-7b-instruct", "Qwen/Qwen1.5-7B-Chat"); m.insert("qwen1.5-14b-instruct", "Qwen/Qwen1.5-14B-Chat"); m.insert("qwen1.5-72b-instruct", "Qwen/Qwen1.5-72B-Chat"); - + // Aliyun API model names -> HF mapping m.insert("qwen-turbo", "Qwen/Qwen2.5-7B-Instruct"); m.insert("qwen-plus", "Qwen/Qwen2.5-14B-Instruct"); m.insert("qwen-max", "Qwen/Qwen2.5-72B-Instruct"); m.insert("qwen-long", "Qwen/Qwen2.5-32B-Instruct"); - + // Qwen 3 series (newer) m.insert("qwen3-8b-instruct", "Qwen/Qwen3-8B-Instruct"); m.insert("qwen3-14b-instruct", "Qwen/Qwen3-14B-Instruct"); m.insert("qwen3-32b-instruct", "Qwen/Qwen3-32B-Instruct"); m.insert("qwen3-72b-instruct", "Qwen/Qwen3-72B-Instruct"); - + // Common variations m.insert("qwen3.5-plus", "Qwen/Qwen3.5-397B-A17B"); m.insert("qwen3.5-turbo", "Qwen/Qwen2.5-7B-Instruct"); m.insert("qwen3.5-max", "Qwen/Qwen2.5-72B-Instruct"); - + // ============== DeepSeek Models ============== m.insert("deepseek-chat", "deepseek-ai/DeepSeek-V3"); m.insert("deepseek-coder", "deepseek-ai/DeepSeek-Coder-V2-Instruct"); @@ -96,29 +99,53 @@ static MODEL_MAPPING: Lazy> = Lazy::new(|| { m.insert("deepseek-v2.5", "deepseek-ai/DeepSeek-V2.5"); m.insert("deepseek-v3", "deepseek-ai/DeepSeek-V3"); m.insert("deepseek-r1", "deepseek-ai/DeepSeek-R1"); - m.insert("deepseek-r1-distill-qwen", "deepseek-ai/DeepSeek-R1-Distill-Qwen-32B"); - + m.insert( + "deepseek-r1-distill-qwen", + "deepseek-ai/DeepSeek-R1-Distill-Qwen-32B", + ); + // ============== Llama Models (Meta) ============== m.insert("llama-2-7b-chat", "meta-llama/Llama-2-7b-chat-hf"); m.insert("llama-2-13b-chat", "meta-llama/Llama-2-13b-chat-hf"); m.insert("llama-2-70b-chat", "meta-llama/Llama-2-70b-chat-hf"); m.insert("llama-3-8b-instruct", "meta-llama/Meta-Llama-3-8B-Instruct"); - m.insert("llama-3-70b-instruct", "meta-llama/Meta-Llama-3-70B-Instruct"); - m.insert("llama-3.1-8b-instruct", "meta-llama/Meta-Llama-3.1-8B-Instruct"); - m.insert("llama-3.1-70b-instruct", "meta-llama/Meta-Llama-3.1-70B-Instruct"); - m.insert("llama-3.1-405b-instruct", "meta-llama/Meta-Llama-3.1-405B-Instruct"); + m.insert( + "llama-3-70b-instruct", + "meta-llama/Meta-Llama-3-70B-Instruct", + ); + m.insert( + "llama-3.1-8b-instruct", + "meta-llama/Meta-Llama-3.1-8B-Instruct", + ); + m.insert( + "llama-3.1-70b-instruct", + "meta-llama/Meta-Llama-3.1-70B-Instruct", + ); + m.insert( + "llama-3.1-405b-instruct", + "meta-llama/Meta-Llama-3.1-405B-Instruct", + ); m.insert("llama-3.2-1b-instruct", "meta-llama/Llama-3.2-1B-Instruct"); m.insert("llama-3.2-3b-instruct", "meta-llama/Llama-3.2-3B-Instruct"); - m.insert("llama-3.3-70b-instruct", "meta-llama/Llama-3.3-70B-Instruct"); - + m.insert( + "llama-3.3-70b-instruct", + "meta-llama/Llama-3.3-70B-Instruct", + ); + // ============== Mistral Models ============== m.insert("mistral-7b-instruct", "mistralai/Mistral-7B-Instruct-v0.3"); m.insert("mistral-large", "mistralai/Mistral-Large-Instruct-2407"); - m.insert("mixtral-8x7b-instruct", "mistralai/Mixtral-8x7B-Instruct-v0.1"); - m.insert("mixtral-8x22b-instruct", "mistralai/Mixtral-8x22B-Instruct-v0.1"); + m.insert( + "mixtral-8x7b-instruct", + "mistralai/Mixtral-8x7B-Instruct-v0.1", + ); + m.insert( + "mixtral-8x22b-instruct", + "mistralai/Mixtral-8x22B-Instruct-v0.1", + ); m.insert("mistral-small", "mistralai/Mistral-Small-24B-Instruct-2501"); m.insert("codestral", "mistralai/Codestral-22B-v0.1"); - + // ============== Yi Models (01.AI) ============== m.insert("yi-6b-chat", "01-ai/Yi-6B-Chat"); m.insert("yi-9b-chat", "01-ai/Yi-9B-Chat"); @@ -126,26 +153,26 @@ static MODEL_MAPPING: Lazy> = Lazy::new(|| { m.insert("yi-1.5-9b-chat", "01-ai/Yi-1.5-9B-Chat"); m.insert("yi-1.5-34b-chat", "01-ai/Yi-1.5-34B-Chat"); m.insert("yi-lightning", "01-ai/Yi-1.5-34B-Chat"); - + // ============== GLM Models (Zhipu AI) ============== m.insert("glm-4-9b-chat", "THUDM/glm-4-9b-chat"); m.insert("glm-4", "THUDM/glm-4-9b-chat"); m.insert("chatglm3-6b", "THUDM/chatglm3-6b"); m.insert("chatglm-turbo", "THUDM/glm-4-9b-chat"); m.insert("chatglm_pro", "THUDM/glm-4-9b-chat"); - + // ============== Baichuan Models ============== m.insert("baichuan2-7b-chat", "baichuan-inc/Baichuan2-7B-Chat"); m.insert("baichuan2-13b-chat", "baichuan-inc/Baichuan2-13B-Chat"); m.insert("baichuan-7b", "baichuan-inc/Baichuan-7B"); - + // ============== Moonshot Models ============== m.insert("moonshot-v1-8k", "moonshot-v1-8k"); m.insert("moonshot-v1-32k", "moonshot-v1-32k"); m.insert("moonshot-v1-128k", "moonshot-v1-128k"); // Moonshot uses custom tokenizer, fallback to similar model m.insert("moonshot-v1", "Qwen/Qwen2.5-7B-Instruct"); - + // ============== Kimi Models (Moonshot K2.5 series) ============== // Kimi K2.5 series - uses Qwen2.5 tokenizer as base m.insert("kimi-k2.5", "Qwen/Qwen2.5-7B-Instruct"); @@ -153,19 +180,19 @@ static MODEL_MAPPING: Lazy> = Lazy::new(|| { m.insert("kimi-k2.5-32k", "Qwen/Qwen2.5-7B-Instruct"); m.insert("kimi-k2.5-128k", "Qwen/Qwen2.5-7B-Instruct"); m.insert("kimi-k2.5-longcontext", "Qwen/Qwen2.5-7B-Instruct"); - + // Legacy Kimi models m.insert("kimi-v1", "Qwen/Qwen2.5-7B-Instruct"); m.insert("kimi-v1-8k", "Qwen/Qwen2.5-7B-Instruct"); m.insert("kimi-v1-32k", "Qwen/Qwen2.5-7B-Instruct"); m.insert("kimi-v1-128k", "Qwen/Qwen2.5-7B-Instruct"); - + // ============== Google Gemini Models ============== m.insert("gemini-pro", "google/gemma-2-9b-it"); m.insert("gemini-1.5-pro", "google/gemma-2-9b-it"); m.insert("gemini-1.5-flash", "google/gemma-2-9b-it"); m.insert("gemini-2.0-flash", "google/gemma-2-9b-it"); - + // ============== Gemma Models (Google) ============== m.insert("gemma-2b-it", "google/gemma-2b-it"); m.insert("gemma-7b-it", "google/gemma-7b-it"); @@ -173,35 +200,41 @@ static MODEL_MAPPING: Lazy> = Lazy::new(|| { m.insert("gemma-2-27b-it", "google/gemma-2-27b-it"); m.insert("gemma-3", "google/gemma-3-4b-it"); m.insert("gemma-3-4b-it", "google/gemma-3-4b-it"); - + // ============== Phi Models (Microsoft) ============== m.insert("phi-2", "microsoft/phi-2"); m.insert("phi-3-mini", "microsoft/Phi-3-mini-4k-instruct"); m.insert("phi-3-small", "microsoft/Phi-3-small-8k-instruct"); m.insert("phi-3-medium", "microsoft/Phi-3-medium-4k-instruct"); m.insert("phi-4", "microsoft/phi-4"); - + // ============== Qwen-QwQ Models ============== m.insert("qwq-32b-preview", "Qwen/QwQ-32B-Preview"); m.insert("qwq-32b", "Qwen/QwQ-32B-Preview"); m.insert("qwq", "Qwen/QwQ-32B-Preview"); - + // ============== InternLM Models ============== m.insert("internlm2-7b-chat", "internlm/internlm2-chat-7b"); m.insert("internlm2-20b-chat", "internlm/internlm2-chat-20b"); m.insert("internlm2.5-7b-chat", "internlm/internlm2_5-7b-chat"); m.insert("internlm3-8b-instruct", "internlm/internlm3-8b-instruct"); - + // ============== Command R Models (Cohere) ============== m.insert("command-r", "CohereForAI/c4ai-command-r-v01"); m.insert("command-r-plus", "CohereForAI/c4ai-command-r-plus"); - + // ============== Other Common Models ============== m.insert("starcoder2-7b", "bigcode/starcoder2-7b"); m.insert("starcoder2-15b", "bigcode/starcoder2-15b"); - m.insert("codellama-7b-instruct", "codellama/CodeLlama-7b-Instruct-hf"); - m.insert("codellama-34b-instruct", "codellama/CodeLlama-34b-Instruct-hf"); - + m.insert( + "codellama-7b-instruct", + "codellama/CodeLlama-7b-Instruct-hf", + ); + m.insert( + "codellama-34b-instruct", + "codellama/CodeLlama-34b-Instruct-hf", + ); + m }); @@ -231,7 +264,7 @@ pub fn map_to_hf_model_id(model_name: &str) -> &str { if let Some(&hf_id) = MODEL_MAPPING.get(model_name) { return hf_id; } - + // Try case-insensitive match let lower_name = model_name.to_lowercase(); for (key, value) in MODEL_MAPPING.iter() { @@ -239,12 +272,12 @@ pub fn map_to_hf_model_id(model_name: &str) -> &str { return value; } } - + // If already looks like a HF model ID (contains '/'), return as-is if model_name.contains('/') { return model_name; } - + // Return original name - will fail at download time with clear error "Qwen/Qwen3.5-397B-A17B" } @@ -257,8 +290,10 @@ pub fn map_to_hf_model_id(model_name: &str) -> &str { /// # Returns /// `true` if the model name has a predefined mapping pub fn has_mapping(model_name: &str) -> bool { - MODEL_MAPPING.contains_key(model_name) || - MODEL_MAPPING.keys().any(|k| k.to_lowercase() == model_name.to_lowercase()) + MODEL_MAPPING.contains_key(model_name) + || MODEL_MAPPING + .keys() + .any(|k| k.to_lowercase() == model_name.to_lowercase()) } /// Get all known model name to HF ID mappings. @@ -271,7 +306,7 @@ pub fn all_mappings() -> impl Iterator { #[cfg(test)] mod tests { use super::*; - + #[test] fn test_openai_mapping() { assert_eq!(map_to_hf_model_id("gpt-4"), "openai/gpt-4"); @@ -279,53 +314,83 @@ mod tests { assert_eq!(map_to_hf_model_id("gpt-4o-mini"), "openai/gpt-4o-mini"); assert_eq!(map_to_hf_model_id("gpt-3.5-turbo"), "openai/gpt-3.5-turbo"); } - + #[test] fn test_anthropic_mapping() { - assert_eq!(map_to_hf_model_id("claude-3-opus"), "anthropic/claude-3-opus"); - assert_eq!(map_to_hf_model_id("claude-3.5-sonnet"), "anthropic/claude-3.5-sonnet"); - assert_eq!(map_to_hf_model_id("claude-3-5-sonnet"), "anthropic/claude-3.5-sonnet"); + assert_eq!( + map_to_hf_model_id("claude-3-opus"), + "anthropic/claude-3-opus" + ); + assert_eq!( + map_to_hf_model_id("claude-3.5-sonnet"), + "anthropic/claude-3.5-sonnet" + ); + assert_eq!( + map_to_hf_model_id("claude-3-5-sonnet"), + "anthropic/claude-3.5-sonnet" + ); } - + #[test] fn test_qwen_mapping() { assert_eq!(map_to_hf_model_id("qwen-turbo"), "Qwen/Qwen2.5-7B-Instruct"); assert_eq!(map_to_hf_model_id("qwen-plus"), "Qwen/Qwen2.5-14B-Instruct"); assert_eq!(map_to_hf_model_id("qwen-max"), "Qwen/Qwen2.5-72B-Instruct"); - assert_eq!(map_to_hf_model_id("qwen2.5-7b-instruct"), "Qwen/Qwen2.5-7B-Instruct"); + assert_eq!( + map_to_hf_model_id("qwen2.5-7b-instruct"), + "Qwen/Qwen2.5-7B-Instruct" + ); } - + #[test] fn test_deepseek_mapping() { - assert_eq!(map_to_hf_model_id("deepseek-chat"), "deepseek-ai/DeepSeek-V3"); + assert_eq!( + map_to_hf_model_id("deepseek-chat"), + "deepseek-ai/DeepSeek-V3" + ); assert_eq!(map_to_hf_model_id("deepseek-r1"), "deepseek-ai/DeepSeek-R1"); } - + #[test] fn test_llama_mapping() { - assert_eq!(map_to_hf_model_id("llama-3-8b-instruct"), "meta-llama/Meta-Llama-3-8B-Instruct"); - assert_eq!(map_to_hf_model_id("llama-3.1-70b-instruct"), "meta-llama/Meta-Llama-3.1-70B-Instruct"); + assert_eq!( + map_to_hf_model_id("llama-3-8b-instruct"), + "meta-llama/Meta-Llama-3-8B-Instruct" + ); + assert_eq!( + map_to_hf_model_id("llama-3.1-70b-instruct"), + "meta-llama/Meta-Llama-3.1-70B-Instruct" + ); } - + #[test] fn test_hf_id_passthrough() { // Already a HF ID - should pass through - assert_eq!(map_to_hf_model_id("Qwen/Qwen2.5-7B-Instruct"), "Qwen/Qwen2.5-7B-Instruct"); - assert_eq!(map_to_hf_model_id("meta-llama/Llama-2-7b"), "meta-llama/Llama-2-7b"); + assert_eq!( + map_to_hf_model_id("Qwen/Qwen2.5-7B-Instruct"), + "Qwen/Qwen2.5-7B-Instruct" + ); + assert_eq!( + map_to_hf_model_id("meta-llama/Llama-2-7b"), + "meta-llama/Llama-2-7b" + ); } - + #[test] fn test_unknown_model() { // Unknown model - should return default fallback - assert_eq!(map_to_hf_model_id("some-unknown-model"), "Qwen/Qwen3.5-397B-A17B"); + assert_eq!( + map_to_hf_model_id("some-unknown-model"), + "Qwen/Qwen3.5-397B-A17B" + ); } - + #[test] fn test_case_insensitive() { assert_eq!(map_to_hf_model_id("GPT-4"), "openai/gpt-4"); assert_eq!(map_to_hf_model_id("QWEN-TURBO"), "Qwen/Qwen2.5-7B-Instruct"); } - + #[test] fn test_has_mapping() { assert!(has_mapping("gpt-4")); diff --git a/src/agentsight/src/tokenizer/multi_model.rs b/src/agentsight/src/tokenizer/multi_model.rs index 779da7462..2517b8e76 100644 --- a/src/agentsight/src/tokenizer/multi_model.rs +++ b/src/agentsight/src/tokenizer/multi_model.rs @@ -123,27 +123,27 @@ impl MultiModelTokenizer { // Map model name to HuggingFace model ID let hf_model_id = map_to_hf_model_id(model_name); - + // Try lookup with HF model ID (in case same HF ID was registered under different name) - if hf_model_id != model_name { - if let Some(tokenizer) = self.get(hf_model_id) { - // Cache under original model name too - let entry = self.tokenizers.get(hf_model_id).cloned(); - if let Some(entry) = entry { - self.tokenizers.insert(model_name.to_string(), entry); - } - return Ok(tokenizer); + if hf_model_id != model_name + && let Some(tokenizer) = self.get(hf_model_id) + { + // Cache under original model name too + let entry = self.tokenizers.get(hf_model_id).cloned(); + if let Some(entry) = entry { + self.tokenizers.insert(model_name.to_string(), entry); } + return Ok(tokenizer); } // Register from HuggingFace Hub using mapped ID self.register_from_hf(hf_model_id)?; // If we used a different HF ID, also cache under original model name - if hf_model_id != model_name { - if let Some(entry) = self.tokenizers.get(hf_model_id).cloned() { - self.tokenizers.insert(model_name.to_string(), entry); - } + if hf_model_id != model_name + && let Some(entry) = self.tokenizers.get(hf_model_id).cloned() + { + self.tokenizers.insert(model_name.to_string(), entry); } // Return the tokenizer diff --git a/src/agentsight/src/unified.rs b/src/agentsight/src/unified.rs index 50b3a8ed9..ec4ad0bb6 100644 --- a/src/agentsight/src/unified.rs +++ b/src/agentsight/src/unified.rs @@ -203,8 +203,14 @@ impl AgentSight { // Create probes - agent discovery is handled by AgentScanner via ProcMon events let enable_udpdns = !config.https_rules.is_empty() || !http_domains.is_empty(); - let mut probes = - Probes::new(&[], config.target_uid, config.enable_filewatch, enable_udpdns, &tcp_targets).context("Failed to create probes")?; + let mut probes = Probes::new( + &[], + config.target_uid, + config.enable_filewatch, + enable_udpdns, + &tcp_targets, + ) + .context("Failed to create probes")?; // Attach procmon for process monitoring probes.attach().context("Failed to attach probes")?; @@ -515,14 +521,19 @@ impl AgentSight { if let std::net::IpAddr::V4(ipv4) = addr.ip() { log::info!( "[UDP-DNS] Adding http target {} → {}", - dns_event.domain, ipv4 + dns_event.domain, + ipv4 ); let target = crate::config::TcpTarget { ip: Some(ipv4), port: None, }; if let Err(e) = self.probes.add_tcp_target(&target) { - log::warn!("[UDP-DNS] Failed to add tcp target {}: {}", ipv4, e); + log::warn!( + "[UDP-DNS] Failed to add tcp target {}: {}", + ipv4, + e + ); } } } @@ -530,7 +541,8 @@ impl AgentSight { Err(e) => { log::warn!( "[UDP-DNS] DNS resolve failed for http domain {}: {}", - dns_event.domain, e + dns_event.domain, + e ); } } @@ -558,11 +570,14 @@ impl AgentSight { // Backfill TokenRecord.agent from pid_agent_name_cache, falling back to comm for ar in &mut analysis_results { - if let crate::analyzer::AnalysisResult::Token(t) = ar { - if t.agent.is_none() { - t.agent = self.pid_agent_name_cache.get(&t.pid).cloned() - .or_else(|| Some(t.comm.clone())); - } + if let crate::analyzer::AnalysisResult::Token(t) = ar + && t.agent.is_none() + { + t.agent = self + .pid_agent_name_cache + .get(&t.pid) + .cloned() + .or_else(|| Some(t.comm.clone())); } } @@ -862,10 +877,10 @@ impl AgentSight { // Look up the real session_id from completed records for the same PID. match store.lookup_session_for_pid(pid) { Ok(Some(ref real_session_id)) => { - if pending.session_id.as_deref() != Some(real_session_id.as_str()) { - if let Err(e) = store.update_session_id(&call_id, real_session_id) { - log::warn!("[DrainCheck] FAIL update session_id: {}", e); - } + if pending.session_id.as_deref() != Some(real_session_id.as_str()) + && let Err(e) = store.update_session_id(&call_id, real_session_id) + { + log::warn!("[DrainCheck] FAIL update session_id: {}", e); } } Ok(None) => {} @@ -876,143 +891,120 @@ impl AgentSight { // ── SSE enrichment ──────────────────────────────────── // Parse captured SSE events for model, trace_id, tokens, output content - if !sse_events.is_empty() { - if let Some(mut enrichment) = + if !sse_events.is_empty() + && let Some(mut enrichment) = GenAIBuilder::extract_sse_enrichment(&sse_events) - { - // If SSE didn't carry usage data (stream was interrupted before - // the final chunk), compute tokens via the real tokenizer. - if enrichment.input_tokens.is_none() - || enrichment.output_tokens.is_none() + { + // If SSE didn't carry usage data (stream was interrupted before + // the final chunk), compute tokens via the real tokenizer. + if enrichment.input_tokens.is_none() || enrichment.output_tokens.is_none() { + let model_name = enrichment + .model + .as_deref() + .or(pending.model.as_deref()) + .unwrap_or("unknown"); + if let Ok(tokenizer) = + crate::tokenizer::get_global_tokenizer(model_name) { - let model_name = enrichment - .model - .as_deref() - .or(pending.model.as_deref()) - .unwrap_or("unknown"); - if let Ok(tokenizer) = - crate::tokenizer::get_global_tokenizer(model_name) + // ── input tokens ── + if enrichment.input_tokens.is_none() + && let Some(body) = request.json_body() + && let Some(messages) = + body.get("messages").and_then(|m| m.as_array()) { - // ── input tokens ── - if enrichment.input_tokens.is_none() { - if let Some(body) = request.json_body() { - if let Some(messages) = - body.get("messages").and_then(|m| m.as_array()) - { - let mut msgs = messages.clone(); - // Parse tool_calls.arguments from string to object - for msg in msgs.iter_mut() { - if let Some(tcs) = msg - .get_mut("tool_calls") - .and_then(|tc| tc.as_array_mut()) - { - for tc in tcs.iter_mut() { - if let Some(f) = tc.get_mut("function") - { - if let Some(a) = f - .get("arguments") - .and_then(|a| a.as_str()) - { - if let Ok(p) = - serde_json::from_str::< - serde_json::Value, - >( - a - ) - { - f["arguments"] = p; - } - } - } - } - } - } - let tools_json: Option> = - body.get("tools") - .and_then(|t| t.as_array()) - .map(|a| a.to_vec()); - let count = match tokenizer - .apply_chat_template_with_tools( - &msgs, - tools_json.as_deref(), - true, - ) { - Ok(formatted) => { - tokenizer.count(&formatted).unwrap_or(0) - } - Err(_) => { - // Fallback: raw message count - msgs.iter() - .filter_map(|m| { - serde_json::to_string(m).ok() - }) - .map(|s| { - tokenizer.count(&s).unwrap_or(0) - }) - .sum() - } - }; - if count > 0 { - enrichment.input_tokens = Some(count as i64); + let mut msgs = messages.clone(); + // Parse tool_calls.arguments from string to object + for msg in msgs.iter_mut() { + if let Some(tcs) = msg + .get_mut("tool_calls") + .and_then(|tc| tc.as_array_mut()) + { + for tc in tcs.iter_mut() { + if let Some(f) = tc.get_mut("function") + && let Some(a) = + f.get("arguments").and_then(|a| a.as_str()) + && let Ok(p) = + serde_json::from_str::(a) + { + f["arguments"] = p; } } } } - // ── output tokens ── - if enrichment.output_tokens.is_none() { - use crate::analyzer::token::extract_response_content; - let mut all_content = String::new(); - let mut all_reasoning = String::new(); - let mut all_tool_calls = Vec::new(); - for ev in &sse_events { - if let Some(chunk) = ev.json_body() { - if let Some((content, reasoning, tool_calls)) = - extract_response_content(Some(&chunk)) - { - if !content.is_empty() { - all_content.push_str(&content); - } - if let Some(r) = reasoning { - if !r.is_empty() { - all_reasoning.push_str(&r); - } - } - for tc in tool_calls { - if !tc.is_empty() { - all_tool_calls.push(tc); - } - } + let tools_json: Option> = body + .get("tools") + .and_then(|t| t.as_array()) + .map(|a| a.to_vec()); + let count = match tokenizer.apply_chat_template_with_tools( + &msgs, + tools_json.as_deref(), + true, + ) { + Ok(formatted) => tokenizer.count(&formatted).unwrap_or(0), + Err(_) => { + // Fallback: raw message count + msgs.iter() + .filter_map(|m| serde_json::to_string(m).ok()) + .map(|s| tokenizer.count(&s).unwrap_or(0)) + .sum() + } + }; + if count > 0 { + enrichment.input_tokens = Some(count as i64); + } + } + // ── output tokens ── + if enrichment.output_tokens.is_none() { + use crate::analyzer::token::extract_response_content; + let mut all_content = String::new(); + let mut all_reasoning = String::new(); + let mut all_tool_calls = Vec::new(); + for ev in &sse_events { + if let Some(chunk) = ev.json_body() + && let Some((content, reasoning, tool_calls)) = + extract_response_content(Some(&chunk)) + { + if !content.is_empty() { + all_content.push_str(&content); + } + if let Some(r) = reasoning + && !r.is_empty() + { + all_reasoning.push_str(&r); + } + for tc in tool_calls { + if !tc.is_empty() { + all_tool_calls.push(tc); } } } - let mut total = 0usize; - if !all_reasoning.is_empty() { - let wrapped = - format!("\n{}\n\n\n", all_reasoning); - total += tokenizer.count(&wrapped).unwrap_or(0); - } - if !all_content.is_empty() { - total += tokenizer.count(&all_content).unwrap_or(0); - } - if !all_tool_calls.is_empty() { - total += tokenizer - .count(&all_tool_calls.join("")) - .unwrap_or(0); - } - if total > 0 { - enrichment.output_tokens = Some(total as i64); - } } - } else { - log::warn!( - "[DrainCheck] tokenizer unavailable for model {:?}, skipping token computation", - enrichment.model.as_deref().or(pending.model.as_deref()) - ); + let mut total = 0usize; + if !all_reasoning.is_empty() { + let wrapped = + format!("\n{}\n\n\n", all_reasoning); + total += tokenizer.count(&wrapped).unwrap_or(0); + } + if !all_content.is_empty() { + total += tokenizer.count(&all_content).unwrap_or(0); + } + if !all_tool_calls.is_empty() { + total += + tokenizer.count(&all_tool_calls.join("")).unwrap_or(0); + } + if total > 0 { + enrichment.output_tokens = Some(total as i64); + } } + } else { + log::warn!( + "[DrainCheck] tokenizer unavailable for model {:?}, skipping token computation", + enrichment.model.as_deref().or(pending.model.as_deref()) + ); } - if let Err(e) = store.enrich_pending_from_sse(&call_id, &enrichment) { - log::warn!("[DrainCheck] FAIL enrich SSE: {}", e); - } + } + if let Err(e) = store.enrich_pending_from_sse(&call_id, &enrichment) { + log::warn!("[DrainCheck] FAIL enrich SSE: {}", e); } } }