diff --git a/src/agentsight/src/aggregator/http/aggregator.rs b/src/agentsight/src/aggregator/http/aggregator.rs index 25a0a943f..be7817d12 100644 --- a/src/agentsight/src/aggregator/http/aggregator.rs +++ b/src/agentsight/src/aggregator/http/aggregator.rs @@ -85,6 +85,8 @@ pub struct HttpConnectionAggregator { last_appended_src_ptr: LruCache, /// Last activity timestamp per connection, used for idle timeout eviction. last_activity: LruCache, + /// Connections already snapshotted after idle timeout. + idle_snapshotted: LruCache, /// Maximum bytes buffered per connection (request body or compressed SSE). max_body_bytes: usize, /// Idle timeout before a connection state is forcibly dropped. @@ -142,6 +144,7 @@ impl HttpConnectionAggregator { sse_continuation_buffers: LruCache::new(cap), last_appended_src_ptr: LruCache::new(cap), last_activity: LruCache::new(cap), + idle_snapshotted: LruCache::new(cap), max_body_bytes: max_body_bytes.max(1024), idle_timeout, } @@ -173,11 +176,11 @@ impl HttpConnectionAggregator { } } - /// Evict connections that have been idle for longer than `self.idle_timeout` - /// or whose buffered body exceeds `self.max_body_bytes`. + /// Evict discardable connections that have been idle for longer than + /// `self.idle_timeout` or whose buffered body exceeds `self.max_body_bytes`. /// - /// Called periodically from the main event loop (via `UnifiedAggregator`) - /// to prevent stale or oversized connection states from accumulating. + /// In-flight request/response states are preserved here so callers can + /// drain and persist them instead of losing a manually interrupted stream. pub fn evict_idle_and_oversized(&mut self) { let now = Instant::now(); let timeout = self.idle_timeout; @@ -196,16 +199,21 @@ impl HttpConnectionAggregator { }) .collect(); + let mut evicted_idle = 0usize; for key in &to_evict { - self.connections.pop(key); - self.sse_continuation_buffers.pop(key); - self.last_appended_src_ptr.pop(key); - self.last_activity.pop(key); + if matches!(self.connections.peek(key), Some(ConnectionState::Idle)) { + self.connections.pop(key); + self.sse_continuation_buffers.pop(key); + self.last_appended_src_ptr.pop(key); + self.last_activity.pop(key); + self.idle_snapshotted.pop(key); + evicted_idle += 1; + } } - if !to_evict.is_empty() { + if evicted_idle > 0 { log::info!( "[HttpAggregator] evicted {} idle connections (timeout={}s)", - to_evict.len(), + evicted_idle, timeout.as_secs() ); } @@ -217,9 +225,6 @@ impl HttpConnectionAggregator { .filter_map(|(k, state)| { let body_len = match state { ConnectionState::RequestBodyPending { body_buffer, .. } => body_buffer.len(), - ConnectionState::SseActive { - compressed_buffer, .. - } => compressed_buffer.as_ref().map_or(0, |b| b.len()), _ => 0, }; if body_len > max_bytes { Some(*k) } else { None } @@ -231,6 +236,7 @@ impl HttpConnectionAggregator { self.sse_continuation_buffers.pop(key); self.last_appended_src_ptr.pop(key); self.last_activity.pop(key); + self.idle_snapshotted.pop(key); } if !oversized.is_empty() { log::warn!( @@ -241,6 +247,58 @@ impl HttpConnectionAggregator { } } + /// Snapshot non-idle connections that exceeded `idle_timeout`. + /// + /// These states can represent manually interrupted or abandoned LLM + /// streams. Returning cloned states lets the caller persist a pending GenAI + /// row while keeping the live connection available if the stream resumes. + pub fn snapshot_idle_connections(&mut self) -> Vec<(ConnectionId, ConnectionState)> { + let now = Instant::now(); + let timeout = self.idle_timeout; + let keys: Vec = self + .last_activity + .iter() + .filter_map(|(k, ts)| { + if now.duration_since(*ts) > timeout { + Some(*k) + } else { + None + } + }) + .collect(); + + let mut result = Vec::new(); + for key in keys { + if self.idle_snapshotted.peek(&key).is_some() { + continue; + } + if let Some(state) = self.connections.peek(&key) { + match state.clone() { + ConnectionState::Idle => {} + state => { + self.idle_snapshotted.push(key, ()); + result.push((key, state)); + } + } + } + } + + if !result.is_empty() { + log::info!( + "[HttpAggregator] snapshotted {} idle in-flight connection(s) (timeout={}s)", + result.len(), + timeout.as_secs() + ); + } + + result + } + + /// Backward-compatible alias for idle snapshots. + pub fn drain_idle_connections(&mut self) -> Vec<(ConnectionId, ConnectionState)> { + self.snapshot_idle_connections() + } + /// Record activity timestamp for a connection. fn touch(&mut self, key: &ConnectionId) { self.last_activity.push(*key, Instant::now()); @@ -390,6 +448,7 @@ impl HttpConnectionAggregator { /// Process HTTP Request (from HTTP Parser) pub fn process_request(&mut self, request: ParsedRequest) { let connection_id = ConnectionId::from_ssl_event(&request.source_event); + self.idle_snapshotted.pop(&connection_id); // Check if body is complete by comparing with Content-Length let content_length: Option = request @@ -906,6 +965,10 @@ impl HttpConnectionAggregator { /// Clear all connections pub fn clear(&mut self) { self.connections.clear(); + self.sse_continuation_buffers.clear(); + self.last_appended_src_ptr.clear(); + self.last_activity.clear(); + self.idle_snapshotted.clear(); } /// Drain all connections (for force complete) @@ -915,7 +978,11 @@ impl HttpConnectionAggregator { .map(|(k, v)| (*k, v.clone())) .collect::>() .into_iter() - .map(|(k, _)| (k, self.connections.pop(&k).unwrap())) + .map(|(k, _)| { + self.last_activity.pop(&k); + self.idle_snapshotted.pop(&k); + (k, self.connections.pop(&k).unwrap()) + }) .collect() } @@ -939,6 +1006,8 @@ impl HttpConnectionAggregator { let mut result = Vec::new(); for key in keys { if let Some(state) = self.connections.pop(&key) { + self.last_activity.pop(&key); + self.idle_snapshotted.pop(&key); match state { ConnectionState::Idle => {} _ => { @@ -999,6 +1068,8 @@ impl HttpConnectionAggregator { let mut result = Vec::new(); for key in dead_keys { if let Some(state) = self.connections.pop(&key) { + self.last_activity.pop(&key); + self.idle_snapshotted.pop(&key); match state { ConnectionState::Idle => { // Silently discard idle entries @@ -2374,4 +2445,97 @@ mod tests { agg.evict_idle_and_oversized(); assert!(agg.connections.peek(&conn_id).is_none()); } + + #[test] + fn test_idle_snapshot_keeps_sse_active_for_resumed_stream() { + let mut agg = HttpConnectionAggregator::with_limits(10, 8192, Duration::from_millis(50)); + let conn_id = ConnectionId { + pid: 1234, + ssl_ptr: 0x7000, + }; + let event = create_mock_ssl_event(conn_id.pid, conn_id.ssl_ptr); + let request = ParsedRequest { + method: "POST".to_string(), + path: "/v1/messages".to_string(), + version: 11, + headers: HashMap::new(), + body_offset: 0, + body_len: 0, + source_event: event.clone(), + reassembled_body: None, + }; + agg.process_request(request); + + let mut headers = HashMap::new(); + headers.insert("content-type".to_string(), "text/event-stream".to_string()); + let response = ParsedResponse { + version: 11, + status_code: 200, + reason: "OK".to_string(), + headers, + body_offset: 0, + body_len: 0, + source_event: event, + }; + assert!(agg.process_response(response).is_none()); + + std::thread::sleep(Duration::from_millis(60)); + let snapshots = agg.snapshot_idle_connections(); + + assert_eq!(snapshots.len(), 1); + assert!(matches!(snapshots[0].1, ConnectionState::SseActive { .. })); + + let done_event = create_mock_ssl_event_with_buf( + conn_id.pid, + conn_id.ssl_ptr, + b"data: [DONE]\n\n".to_vec(), + 0, + ); + let done = ParsedSseEvent::new(None, None, None, 6, 6, done_event); + let result = agg.process_sse_event(&conn_id, done); + assert!( + matches!(result, Some(AggregatedResult::SseComplete(_))), + "idle snapshot must not remove the in-flight stream state" + ); + } + + #[test] + fn test_oversized_request_body_pending_is_evicted() { + let mut agg = HttpConnectionAggregator::with_limits(10, 1024, Duration::from_secs(60)); + let conn_id = ConnectionId { + pid: 1234, + ssl_ptr: 0x7100, + }; + let event = create_mock_ssl_event(conn_id.pid, conn_id.ssl_ptr); + let request = ParsedRequest { + method: "POST".to_string(), + path: "/v1/messages".to_string(), + version: 11, + headers: HashMap::new(), + body_offset: 0, + body_len: 0, + source_event: event, + reassembled_body: None, + }; + + agg.connections.push( + conn_id, + ConnectionState::RequestBodyPending { + request, + expected_body_len: Some(4096), + body_buffer: vec![b'x'; 2048], + }, + ); + agg.last_activity.push(conn_id, Instant::now()); + agg.sse_continuation_buffers + .push(conn_id, b"stale".to_vec()); + agg.last_appended_src_ptr.push(conn_id, 42); + + agg.evict_idle_and_oversized(); + + assert!(agg.connections.peek(&conn_id).is_none()); + assert!(agg.last_activity.peek(&conn_id).is_none()); + assert!(agg.sse_continuation_buffers.peek(&conn_id).is_none()); + assert!(agg.last_appended_src_ptr.peek(&conn_id).is_none()); + } } diff --git a/src/agentsight/src/aggregator/unified.rs b/src/agentsight/src/aggregator/unified.rs index d27c46b87..159e8a624 100644 --- a/src/agentsight/src/aggregator/unified.rs +++ b/src/agentsight/src/aggregator/unified.rs @@ -177,4 +177,12 @@ impl Aggregator { pub fn drain_dead_pid_connections(&mut self) -> Vec<(ConnectionId, ConnectionState)> { self.http.drain_dead_pid_connections() } + + /// Snapshot in-flight HTTP connections that exceeded the idle timeout. + /// + /// Used to persist evidence for manually interrupted streams where the + /// agent process remains alive, so dead-PID draining would never run. + pub fn snapshot_idle_connections(&mut self) -> Vec<(ConnectionId, ConnectionState)> { + self.http.snapshot_idle_connections() + } } diff --git a/src/agentsight/src/background.rs b/src/agentsight/src/background.rs index 7c674e62d..aeded5f5a 100644 --- a/src/agentsight/src/background.rs +++ b/src/agentsight/src/background.rs @@ -271,7 +271,7 @@ mod tests { #[test] fn test_stale_scanner_loop_runs_body_then_stops() { - use crate::storage::sqlite::PendingCallInfo; + use crate::storage::sqlite::{PendingCallInfo, PendingOrigin}; let dir = tmp_dir("stale2"); let store = Arc::new(GenAISqliteStore::new_with_path(&dir.join("test.db")).unwrap()); @@ -297,6 +297,8 @@ mod tests { model: None, provider: None, call_kind: "main".to_string(), + pending_origin: PendingOrigin::RequestCapture, + pending_match_key: None, }) .unwrap(); diff --git a/src/agentsight/src/genai/builder.rs b/src/agentsight/src/genai/builder.rs index 30be4c0fa..44ca92c1b 100644 --- a/src/agentsight/src/genai/builder.rs +++ b/src/agentsight/src/genai/builder.rs @@ -11,7 +11,8 @@ use crate::analyzer::AnalysisResult; use crate::analyzer::token::TokenParser; use crate::parser::sse::ParsedSseEvent; use crate::response_map::ResponseSessionMapper; -use crate::storage::sqlite::{PendingCallInfo, SseEnrichment}; +use crate::storage::sqlite::{PendingCallInfo, PendingOrigin, SseEnrichment}; +use sha2::{Digest, Sha256}; use std::sync::atomic::{AtomicU64, Ordering}; use std::time::{SystemTime, UNIX_EPOCH}; @@ -57,6 +58,44 @@ impl GenAIBuilder { } } + pub(super) fn pending_match_key( + pid: u32, + start_timestamp_ns: u64, + method: &str, + path: &str, + request_body: Option<&str>, + ) -> String { + let mut hasher = Sha256::new(); + hasher.update(b"agentsight-pending-v1\0"); + hasher.update(pid.to_string().as_bytes()); + hasher.update(b"\0"); + hasher.update(start_timestamp_ns.to_string().as_bytes()); + hasher.update(b"\0"); + hasher.update(method.as_bytes()); + hasher.update(b"\0"); + hasher.update(path.as_bytes()); + hasher.update(b"\0"); + if let Some(body) = request_body { + hasher.update(body.as_bytes()); + } + format!("pmk-{:x}", hasher.finalize()) + } + + fn parsed_request_match_body( + request: &ParsedRequest, + body: Option<&serde_json::Value>, + ) -> Option { + body.map(|v| serde_json::to_string(v).unwrap_or_default()) + .or_else(|| { + let raw = request.body(); + if raw.is_empty() { + None + } else { + Some(String::from_utf8_lossy(raw).to_string()) + } + }) + } + /// Build GenAI semantic events AND a `PendingCallInfo` to be written to DB /// before the response arrives. /// @@ -156,6 +195,8 @@ impl GenAIBuilder { .get("call_kind") .cloned() .unwrap_or_else(|| "main".to_string()), + pending_origin: PendingOrigin::RequestCapture, + pending_match_key: llm_call.metadata.get("pending_match_key").cloned(), }); events.push(GenAISemanticEvent::LLMCall(llm_call)); @@ -201,6 +242,14 @@ impl GenAIBuilder { let call_id = self.generate_id(); let body = request.json_body(); + let match_body = Self::parsed_request_match_body(request, body.as_ref()); + let pending_match_key = Self::pending_match_key( + conn_id.pid, + request.source_event.timestamp_ns, + &request.method, + &request.path, + match_body.as_deref(), + ); // Determine if streaming let is_sse = body @@ -364,6 +413,8 @@ impl GenAIBuilder { model, provider, call_kind: call_kind.to_string(), + pending_origin: PendingOrigin::DeadPidDrain, + pending_match_key: Some(pending_match_key), }) } diff --git a/src/agentsight/src/genai/call_builder.rs b/src/agentsight/src/genai/call_builder.rs index 1dc45e08d..120afbae7 100644 --- a/src/agentsight/src/genai/call_builder.rs +++ b/src/agentsight/src/genai/call_builder.rs @@ -42,6 +42,13 @@ impl GenAIBuilder { // Need at least HttpRecord to build LLMCall let http = http_record?; + let pending_match_key = Self::pending_match_key( + http.pid, + http.timestamp_ns, + &http.method, + &http.path, + http.request_body.as_deref(), + ); // 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); @@ -255,6 +262,7 @@ impl GenAIBuilder { if let Some(ref sid) = session_id { meta.insert("session_id".to_string(), sid.clone()); } + meta.insert("pending_match_key".to_string(), pending_match_key); meta }, }) @@ -291,9 +299,7 @@ impl GenAIBuilder { let role = format!("{:?}", m.role).to_lowercase(); InputMessage { role, - parts: vec![MessagePart::Text { - content: m.content.as_text(), - }], + parts: Self::anthropic_message_content_to_parts(&m.content), name: None, } }) @@ -403,6 +409,77 @@ impl GenAIBuilder { // openai_msg_to_output / parse_openai_tool_call_value / parse_sse_response_body / // extract_parts_from_sse_body live in `openai_parse.rs` (same impl block). + fn anthropic_message_content_to_parts( + content: &crate::analyzer::message::AnthropicMessageContent, + ) -> Vec { + match content { + crate::analyzer::message::AnthropicMessageContent::Text(text) => { + if text.is_empty() { + Vec::new() + } else { + vec![MessagePart::Text { + content: text.clone(), + }] + } + } + crate::analyzer::message::AnthropicMessageContent::Blocks(blocks) => blocks + .iter() + .filter_map(Self::anthropic_content_block_to_part) + .collect(), + } + } + + fn anthropic_content_block_to_part( + block: &crate::analyzer::message::AnthropicContentBlock, + ) -> Option { + match block { + crate::analyzer::message::AnthropicContentBlock::Text { text, .. } => { + if text.is_empty() { + None + } else { + Some(MessagePart::Text { + content: text.clone(), + }) + } + } + crate::analyzer::message::AnthropicContentBlock::ToolUse { id, name, input } => { + Some(MessagePart::ToolCall { + id: Some(id.clone()), + name: name.clone(), + arguments: Some(input.clone()), + }) + } + crate::analyzer::message::AnthropicContentBlock::ToolResult { + tool_use_id, + content, + is_error, + } => { + let response = match (content.clone(), *is_error) { + (Some(value), Some(is_error)) => { + serde_json::json!({"content": value, "is_error": is_error}) + } + (Some(value), None) => value, + (None, Some(is_error)) => serde_json::json!({"is_error": is_error}), + (None, None) => serde_json::Value::Null, + }; + Some(MessagePart::ToolCallResponse { + id: Some(tool_use_id.clone()), + response, + }) + } + crate::analyzer::message::AnthropicContentBlock::Thinking { thinking, .. } => { + if thinking.is_empty() { + None + } else { + Some(MessagePart::Reasoning { + content: thinking.clone(), + }) + } + } + _ => None, + } + } + /// Build LLMResponse from parsed message or HTTP record fn build_response( &self, @@ -936,6 +1013,69 @@ mod tests { assert!(call.request.tools.is_some()); } + #[test] + fn test_build_request_anthropic_preserves_tool_result_blocks() { + let builder = GenAIBuilder::new(); + let anth_req = AnthropicRequest { + model: "claude-3".to_string(), + messages: vec![ + AnthMsg { + role: MessageRole::Assistant, + content: AnthropicMessageContent::Blocks(vec![ + AnthropicContentBlock::ToolUse { + id: "toolu_1".to_string(), + name: "Bash".to_string(), + input: serde_json::json!({"command": "cat /tmp/missing.txt"}), + }, + ]), + }, + AnthMsg { + role: MessageRole::User, + content: AnthropicMessageContent::Blocks(vec![ + AnthropicContentBlock::ToolResult { + tool_use_id: "toolu_1".to_string(), + content: Some(serde_json::json!( + "Exit code 1\ncat: /tmp/missing.txt: No such file or directory" + )), + is_error: Some(true), + }, + ]), + }, + ], + max_tokens: 200, + system: None, + stream: Some(false), + temperature: None, + top_p: None, + top_k: None, + stop_sequences: None, + metadata: None, + tools: None, + tool_choice: None, + }; + let parsed = ParsedApiMessage::AnthropicMessage { + request: Some(anth_req), + response: None, + }; + let http = make_http("/v1/messages", None, None); + let call = build_call( + &builder, + &[AnalysisResult::Http(http), AnalysisResult::Message(parsed)], + ) + .unwrap(); + + assert!(call.request.messages.iter().any(|message| { + message.parts.iter().any(|part| { + matches!( + part, + MessagePart::ToolCallResponse { id, response } + if id.as_deref() == Some("toolu_1") + && response.get("is_error").and_then(|v| v.as_bool()) == Some(true) + ) + }) + })); + } + #[test] fn test_build_request_sysom_full() { let builder = GenAIBuilder::new(); diff --git a/src/agentsight/src/probes/probes.rs b/src/agentsight/src/probes/probes.rs index 89da014af..a863f0dfe 100644 --- a/src/agentsight/src/probes/probes.rs +++ b/src/agentsight/src/probes/probes.rs @@ -345,7 +345,7 @@ impl Probes { } ChannelPolicy::Sample(n) => { let idx = drop_counter.fetch_add(1, Ordering::Relaxed); - if idx % n == 0 { + if idx.is_multiple_of(n) { if event_tx.try_send(e).is_err() { log::warn!( "Probes event channel full (capacity={}); dropping sampled event", diff --git a/src/agentsight/src/probes/sslsniff.rs b/src/agentsight/src/probes/sslsniff.rs index 6d00ded54..83304ca9d 100644 --- a/src/agentsight/src/probes/sslsniff.rs +++ b/src/agentsight/src/probes/sslsniff.rs @@ -622,7 +622,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).next_back() { + match hs_matches.iter().rfind(|&&o| o < read_off) { Some(&o) => o, None => { if verbose { diff --git a/src/agentsight/src/storage/sqlite/genai/mod.rs b/src/agentsight/src/storage/sqlite/genai/mod.rs index 7ac5af58f..61d4db0ed 100644 --- a/src/agentsight/src/storage/sqlite/genai/mod.rs +++ b/src/agentsight/src/storage/sqlite/genai/mod.rs @@ -27,7 +27,7 @@ use crate::config::BatchConfig; // Re-export public types from sub-modules pub use events::TraceEventDetail; -pub use pending::{PendingCallInfo, SseEnrichment}; +pub use pending::{PendingCallInfo, PendingOrigin, SseEnrichment}; pub use session::{SavingsSessionSummary, SessionSummary, ToolCallTurnInfo, TraceSummary}; pub use stats::{AgentTokenSummary, ModelTimeseriesBucket, TimeseriesBucket}; diff --git a/src/agentsight/src/storage/sqlite/genai/pending.rs b/src/agentsight/src/storage/sqlite/genai/pending.rs index 7a1f85cdb..c38ea05c4 100644 --- a/src/agentsight/src/storage/sqlite/genai/pending.rs +++ b/src/agentsight/src/storage/sqlite/genai/pending.rs @@ -7,6 +7,29 @@ use crate::genai::semantic::GenAISemanticEvent; // ─── Query result types ──────────────────────────────────────────────────────── +/// Source that created a pending GenAI row. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub enum PendingOrigin { + /// Request was captured before the response arrived. + #[default] + RequestCapture, + /// In-flight request was drained after its owner PID disappeared. + DeadPidDrain, + /// In-flight stream was snapshotted after an idle timeout while the PID lived. + IdleDrain, +} + +impl PendingOrigin { + /// Stable string stored in SQLite. + pub fn as_str(self) -> &'static str { + match self { + Self::RequestCapture => "request_capture", + Self::DeadPidDrain => "dead_pid_drain", + Self::IdleDrain => "idle_drain", + } + } +} + /// Lightweight info needed to write a PENDING record when a request is first seen pub struct PendingCallInfo { /// Unique call ID (same one that will be used in the complete record) @@ -43,6 +66,10 @@ pub struct PendingCallInfo { pub provider: Option, /// Call kind classification: "main" | "recap" | "web_search" pub call_kind: String, + /// Source that created this pending record. + pub pending_origin: PendingOrigin, + /// Deterministic key used to reconcile idle snapshots with resumed streams. + pub pending_match_key: Option, } /// Data extracted from captured SSE events for enriching a pending record. @@ -73,14 +100,14 @@ impl GenAISqliteStore { start_timestamp_ns, pid, process_name, agent_name, http_method, http_path, is_sse, input_messages, system_instructions, user_query, - model, provider, call_kind, + model, provider, call_kind, pending_origin, pending_match_key, event_json ) VALUES ( 'llm_call', 'pending', ?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, - ?16, ?17, ?18, + ?16, ?17, ?18, ?19, ?20, '{}' )", params![ @@ -102,6 +129,8 @@ impl GenAISqliteStore { info.model, info.provider, info.call_kind, + info.pending_origin.as_str(), + info.pending_match_key, ], )?; Ok(()) @@ -245,38 +274,38 @@ impl GenAISqliteStore { call.metadata.get("session_id"), call.end_timestamp_ns as i64, call.duration_ns as i64, - call.provider, - call.model, - call.model, - call.model, + call.provider.as_str(), + call.model.as_str(), + call.model.as_str(), + call.model.as_str(), call.request.temperature, call.request.max_tokens.map(|v| v as i64), call.request.top_p, call.request.frequency_penalty, call.request.presence_penalty, - finish_reasons, + finish_reasons.as_deref(), call.metadata.get("server.address"), input_tokens, output_tokens, total_tokens, cache_creation, cache_read, - system_instructions, - input_messages, - output_messages, + system_instructions.as_deref(), + input_messages.as_deref(), + output_messages.as_deref(), call.metadata .get("status_code") .and_then(|s| s.parse::().ok()), call.metadata .get("sse_event_count") .and_then(|s| s.parse::().ok()), - event_json, - tool_call_ids, + event_json.as_str(), + tool_call_ids.as_deref(), call.metadata .get("call_kind") .map(|s| s.as_str()) .unwrap_or("main"), - call.call_id, + call.call_id.as_str(), ], )?; @@ -304,6 +333,100 @@ impl GenAISqliteStore { "[GenAI] No row for call_id={}, inserting directly", call.call_id ); + + if let Some(match_key) = call.metadata.get("pending_match_key") { + let updated = conn.execute( + "UPDATE genai_events SET + status = 'complete', + trace_id = ?1, + conversation_id = ?2, + session_id = ?3, + end_timestamp_ns = ?4, + duration_ns = ?5, + provider = ?6, + model = ?7, + request_model = ?8, + response_model = ?9, + temperature = ?10, + max_tokens = ?11, + top_p = ?12, + frequency_penalty = ?13, + presence_penalty = ?14, + finish_reasons = ?15, + server_address = ?16, + input_tokens = ?17, + output_tokens = ?18, + total_tokens = ?19, + cache_creation_tokens = ?20, + cache_read_tokens = ?21, + system_instructions = ?22, + input_messages = ?23, + output_messages = ?24, + status_code = ?25, + sse_event_count = ?26, + event_json = ?27, + tool_call_ids = ?28, + call_kind = ?29, + call_id = ?30 + WHERE id = ( + SELECT id FROM genai_events + WHERE event_type = 'llm_call' + AND status IN ('pending', 'interrupted') + AND pending_origin = 'idle_drain' + AND pending_match_key = ?31 + ORDER BY start_timestamp_ns DESC + LIMIT 1 + )", + params![ + call.metadata.get("response_id"), + call.metadata.get("conversation_id"), + call.metadata.get("session_id"), + call.end_timestamp_ns as i64, + call.duration_ns as i64, + call.provider.as_str(), + call.model.as_str(), + call.model.as_str(), + call.model.as_str(), + call.request.temperature, + call.request.max_tokens.map(|v| v as i64), + call.request.top_p, + call.request.frequency_penalty, + call.request.presence_penalty, + finish_reasons.as_deref(), + call.metadata.get("server.address"), + input_tokens, + output_tokens, + total_tokens, + cache_creation, + cache_read, + system_instructions.as_deref(), + input_messages.as_deref(), + output_messages.as_deref(), + call.metadata + .get("status_code") + .and_then(|s| s.parse::().ok()), + call.metadata + .get("sse_event_count") + .and_then(|s| s.parse::().ok()), + event_json.as_str(), + tool_call_ids.as_deref(), + call.metadata + .get("call_kind") + .map(|s| s.as_str()) + .unwrap_or("main"), + call.call_id.as_str(), + match_key.as_str(), + ], + )?; + + if updated > 0 { + log::debug!( + "[GenAI] Promoted idle pending snapshot→complete for call_id={}", + call.call_id + ); + return Ok(()); + } + } } // Fallback: store_event handles the full INSERT path self.store_event(event) @@ -444,7 +567,8 @@ impl GenAISqliteStore { FROM genai_events WHERE event_type = 'llm_call' AND status = 'pending' - AND pid = ?1", + AND pid = ?1 + AND pending_origin != 'idle_drain'", )?; let rows = stmt .query_map(params![pid], |row| { @@ -475,7 +599,8 @@ impl GenAISqliteStore { SET status = 'interrupted', interruption_type = ?1 WHERE event_type = 'llm_call' AND status = 'pending' - AND pid = ?2", + AND pid = ?2 + AND pending_origin != 'idle_drain'", params![itype, pid], )?; if updated > 0 { @@ -507,6 +632,7 @@ impl GenAISqliteStore { FROM genai_events WHERE event_type = 'llm_call' AND status = 'pending' + AND pending_origin != 'idle_drain' AND pid IN ({placeholders})" ); let params_vec: Vec> = pids diff --git a/src/agentsight/src/storage/sqlite/genai/schema.rs b/src/agentsight/src/storage/sqlite/genai/schema.rs index 8c55b5665..8c6e515ab 100644 --- a/src/agentsight/src/storage/sqlite/genai/schema.rs +++ b/src/agentsight/src/storage/sqlite/genai/schema.rs @@ -88,6 +88,9 @@ impl GenAISqliteStore { interruption_type TEXT, -- Call kind classification (main/recap/web_search) call_kind TEXT NOT NULL DEFAULT 'main', + -- Pending row provenance and reconciliation key + pending_origin TEXT NOT NULL DEFAULT 'request_capture', + pending_match_key TEXT, -- Full event as JSON (fallback) event_json TEXT NOT NULL, created_at DATETIME DEFAULT CURRENT_TIMESTAMP @@ -185,6 +188,16 @@ impl GenAISqliteStore { "idx_genai_call_kind" ); + // v7: pending provenance for idle/drain lifecycle handling + ensure_col!( + "pending_origin", + "TEXT NOT NULL DEFAULT 'request_capture'", + "idx_genai_pending_origin" + ); + + // v8: stable key used to reconcile idle stream snapshots on completion + ensure_col!("pending_match_key", "TEXT", "idx_genai_pending_match_key"); + Ok(()) } diff --git a/src/agentsight/src/storage/sqlite/genai/tests.rs b/src/agentsight/src/storage/sqlite/genai/tests.rs index fd977e47a..3baac2b30 100644 --- a/src/agentsight/src/storage/sqlite/genai/tests.rs +++ b/src/agentsight/src/storage/sqlite/genai/tests.rs @@ -1,5 +1,7 @@ use super::*; -use crate::genai::semantic::{GenAISemanticEvent, LLMCall, LLMRequest}; +use crate::genai::semantic::{ + GenAISemanticEvent, LLMCall, LLMRequest, LLMResponse, MessagePart, OutputMessage, +}; /// Integration test: store_event (post-fix, no per-insert VACUUM) still /// persists data correctly and the row is immediately readable. @@ -659,6 +661,8 @@ fn test_insert_pending() { model: Some("gpt-4".to_string()), provider: Some("openai".to_string()), call_kind: "main".to_string(), + pending_origin: PendingOrigin::RequestCapture, + pending_match_key: None, }; store.insert_pending(&info).unwrap(); let conn = store.conn.lock().unwrap(); @@ -674,6 +678,156 @@ fn test_insert_pending() { cleanup_db(&path); } +#[test] +fn test_insert_pending_records_idle_origin_and_match_key() { + let path = + std::env::temp_dir().join(format!("test_genai_idle_origin_{}.db", std::process::id())); + cleanup_db(&path); + let store = GenAISqliteStore::new_with_path(&path).unwrap(); + let info = PendingCallInfo { + call_id: "idle-temp".to_string(), + trace_id: None, + conversation_id: Some("c-idle".to_string()), + session_id: Some("s-idle".to_string()), + start_timestamp_ns: BASE_NS as u64, + pid: 42, + process_name: "claude".to_string(), + agent_name: Some("claude".to_string()), + http_method: Some("POST".to_string()), + http_path: Some("/v1/messages".to_string()), + input_messages: None, + system_instructions: None, + user_query: Some("hello".to_string()), + is_sse: true, + model: Some("claude-sonnet".to_string()), + provider: Some("anthropic".to_string()), + call_kind: "main".to_string(), + pending_origin: PendingOrigin::IdleDrain, + pending_match_key: Some("match-idle-1".to_string()), + }; + store.insert_pending(&info).unwrap(); + + let conn = store.conn.lock().unwrap(); + let (origin, match_key): (String, String) = conn + .query_row( + "SELECT pending_origin, pending_match_key FROM genai_events WHERE call_id = 'idle-temp'", + [], + |r| Ok((r.get(0)?, r.get(1)?)), + ) + .unwrap(); + assert_eq!(origin, "idle_drain"); + assert_eq!(match_key, "match-idle-1"); + drop(conn); + cleanup_db(&path); +} + +#[test] +fn test_complete_pending_promotes_idle_snapshot_by_match_key() { + let path = + std::env::temp_dir().join(format!("test_genai_idle_promote_{}.db", std::process::id())); + cleanup_db(&path); + let store = GenAISqliteStore::new_with_path(&path).unwrap(); + let info = PendingCallInfo { + call_id: "idle-temp".to_string(), + trace_id: None, + conversation_id: Some("c-idle".to_string()), + session_id: Some("s-idle".to_string()), + start_timestamp_ns: BASE_NS as u64, + pid: 42, + process_name: "claude".to_string(), + agent_name: Some("claude".to_string()), + http_method: Some("POST".to_string()), + http_path: Some("/v1/messages".to_string()), + input_messages: None, + system_instructions: None, + user_query: Some("hello".to_string()), + is_sse: true, + model: Some("claude-sonnet".to_string()), + provider: Some("anthropic".to_string()), + call_kind: "main".to_string(), + pending_origin: PendingOrigin::IdleDrain, + pending_match_key: Some("match-idle-2".to_string()), + }; + store.insert_pending(&info).unwrap(); + + let request = LLMRequest { + messages: vec![], + temperature: None, + max_tokens: None, + frequency_penalty: None, + presence_penalty: None, + top_p: None, + top_k: None, + seed: None, + stop_sequences: None, + stream: true, + tools: None, + raw_body: None, + }; + let mut call = LLMCall::new( + "real-response-id".to_string(), + BASE_NS as u64, + "anthropic".to_string(), + "claude-sonnet".to_string(), + request, + 42, + "claude".to_string(), + ); + call.set_response( + LLMResponse { + messages: vec![OutputMessage { + role: "assistant".to_string(), + parts: vec![MessagePart::Text { + content: "done".to_string(), + }], + name: None, + finish_reason: Some("stop".to_string()), + }], + streamed: true, + raw_body: None, + }, + (BASE_NS + STEP_NS) as u64, + ); + call.metadata + .insert("response_id".to_string(), "real-response-id".to_string()); + call.metadata + .insert("pending_match_key".to_string(), "match-idle-2".to_string()); + call.metadata + .insert("conversation_id".to_string(), "c-real".to_string()); + call.metadata + .insert("session_id".to_string(), "s-real".to_string()); + call.metadata + .insert("status_code".to_string(), "200".to_string()); + call.metadata + .insert("sse_event_count".to_string(), "2".to_string()); + call.metadata + .insert("call_kind".to_string(), "main".to_string()); + + store + .complete_pending(&GenAISemanticEvent::LLMCall(call)) + .unwrap(); + + let conn = store.conn.lock().unwrap(); + let total: i64 = conn + .query_row("SELECT COUNT(*) FROM genai_events", [], |r| r.get(0)) + .unwrap(); + assert_eq!(total, 1, "complete must update the idle snapshot row"); + + let (status, call_id, trace_id, origin): (String, String, String, String) = conn + .query_row( + "SELECT status, call_id, trace_id, pending_origin FROM genai_events", + [], + |r| Ok((r.get(0)?, r.get(1)?, r.get(2)?, r.get(3)?)), + ) + .unwrap(); + assert_eq!(status, "complete"); + assert_eq!(call_id, "real-response-id"); + assert_eq!(trace_id, "real-response-id"); + assert_eq!(origin, "idle_drain"); + drop(conn); + cleanup_db(&path); +} + #[test] fn test_mark_interrupted_stale() { let (store, path) = create_populated_store("mark_stale"); @@ -733,6 +887,63 @@ fn test_mark_pending_interrupted_for_pid() { cleanup_db(&path); } +#[test] +fn test_crash_sweep_ignores_idle_drain_pending() { + let path = + std::env::temp_dir().join(format!("test_genai_idle_sweep_{}.db", std::process::id())); + cleanup_db(&path); + let store = GenAISqliteStore::new_with_path(&path).unwrap(); + + for (call_id, origin) in [ + ("dead-drain", PendingOrigin::DeadPidDrain), + ("idle-drain", PendingOrigin::IdleDrain), + ] { + let info = PendingCallInfo { + call_id: call_id.to_string(), + trace_id: None, + conversation_id: Some("c-sweep".to_string()), + session_id: Some("s-sweep".to_string()), + start_timestamp_ns: BASE_NS as u64, + pid: 42, + process_name: "claude".to_string(), + agent_name: Some("claude".to_string()), + http_method: Some("POST".to_string()), + http_path: Some("/v1/messages".to_string()), + input_messages: None, + system_instructions: None, + user_query: Some("hello".to_string()), + is_sse: true, + model: Some("claude-sonnet".to_string()), + provider: Some("anthropic".to_string()), + call_kind: "main".to_string(), + pending_origin: origin, + pending_match_key: Some(format!("match-{call_id}")), + }; + store.insert_pending(&info).unwrap(); + } + + let listed = store.list_pending_for_pid(42).unwrap(); + assert_eq!(listed.len(), 1); + assert_eq!(listed[0].0, "dead-drain"); + + let updated = store + .mark_pending_interrupted_for_pid(42, "agent_crash") + .unwrap(); + assert_eq!(updated, 1); + + let conn = store.conn.lock().unwrap(); + let idle_status: String = conn + .query_row( + "SELECT status FROM genai_events WHERE call_id = 'idle-drain'", + [], + |r| r.get(0), + ) + .unwrap(); + assert_eq!(idle_status, "pending"); + drop(conn); + cleanup_db(&path); +} + #[test] fn test_enrich_pending_from_sse() { let (store, path) = create_populated_store("enrich_sse"); @@ -830,6 +1041,8 @@ fn make_store_with_pending(records: &[(&str, &str, &str, &str, i64)]) -> GenAISq model: Some("gpt-4".to_string()), provider: Some("openai".to_string()), call_kind: kind.to_string(), + pending_origin: PendingOrigin::RequestCapture, + pending_match_key: None, }; store.insert_pending(&info).unwrap(); } diff --git a/src/agentsight/src/storage/sqlite/mod.rs b/src/agentsight/src/storage/sqlite/mod.rs index 36ad0a2cd..6c968fdb6 100644 --- a/src/agentsight/src/storage/sqlite/mod.rs +++ b/src/agentsight/src/storage/sqlite/mod.rs @@ -35,7 +35,7 @@ pub use token_consumption::{ pub use http::HttpStore; // Re-export GenAI SQLite storage -pub use genai::{GenAISqliteStore, PendingCallInfo, SseEnrichment}; +pub use genai::{GenAISqliteStore, PendingCallInfo, PendingOrigin, SseEnrichment}; // Re-export Interruption SQLite storage pub use interruption::{InterruptionRecord, InterruptionStore, InterruptionTypeStat}; diff --git a/src/agentsight/src/storage/sqlite/token.rs b/src/agentsight/src/storage/sqlite/token.rs index 4b79e6df6..6c82c9492 100644 --- a/src/agentsight/src/storage/sqlite/token.rs +++ b/src/agentsight/src/storage/sqlite/token.rs @@ -208,7 +208,7 @@ pub fn format_tokens_with_commas(count: u64) -> String { let s = count.to_string(); let mut result = String::new(); for (i, c) in s.chars().enumerate() { - if i > 0 && (s.len() - i) % 3 == 0 { + if i > 0 && (s.len() - i).is_multiple_of(3) { result.push(','); } result.push(c); diff --git a/src/agentsight/src/storage/unified.rs b/src/agentsight/src/storage/unified.rs index d275091ab..a6e029b52 100644 --- a/src/agentsight/src/storage/unified.rs +++ b/src/agentsight/src/storage/unified.rs @@ -272,7 +272,7 @@ impl Storage { // 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 count.is_multiple_of(self.purge_interval) { if let Err(e) = self.purge_expired() { log::warn!("Auto-purge failed: {e}"); } diff --git a/src/agentsight/src/unified.rs b/src/agentsight/src/unified.rs index e7a5798a6..9e9159b3d 100644 --- a/src/agentsight/src/unified.rs +++ b/src/agentsight/src/unified.rs @@ -909,6 +909,8 @@ impl AgentSight { self.flush_expired_pending_genai(); // Drain orphaned connections from dead PIDs and persist as pending self.drain_and_persist_dead_connections(); + // Snapshot idle in-flight streams whose owning process is still alive. + self.drain_and_persist_idle_connections(); // Check if config watcher deposited a new LogtailExporter self.check_pending_logtail(); // Periodically purge old/oversized interruption DB entries @@ -1278,6 +1280,58 @@ impl AgentSight { } } + /// Snapshot idle in-flight streams and persist them as `pending` records. + /// + /// This covers manual output interruption where the agent process stays + /// alive, so dead-PID draining cannot discover the abandoned stream. + fn drain_and_persist_idle_connections(&mut self) { + let snapshots = self.aggregator.snapshot_idle_connections(); + if snapshots.is_empty() { + return; + } + + use crate::aggregator::ConnectionState; + use crate::storage::sqlite::PendingOrigin; + + for (conn_id, state) in snapshots { + let (_state_name, request) = match state { + ConnectionState::RequestPending { request } => ("RequestPending", request), + ConnectionState::SseActive { + request: Some(req), .. + } => ("SseActive", req), + _ => continue, + }; + + if let Some(mut pending) = self.genai_builder.build_pending_from_request( + &request, + &conn_id, + &self.pid_agent_name_cache, + ) { + pending.pending_origin = PendingOrigin::IdleDrain; + if let Some(ref store) = self.genai_sqlite_store { + let call_id = pending.call_id.clone(); + if let Err(e) = store.insert_pending(&pending) { + log::warn!("[IdleDrain] Failed to persist pending call: {e}"); + continue; + } + log::info!( + "[IdleDrain] persisted idle in-flight call as pending: call_id={} pid={} path={}", + call_id, + conn_id.pid, + request.path, + ); + } + } else { + log::debug!( + "[IdleDrain] build_pending returned None: pid={} path={} body_len={}", + conn_id.pid, + request.path, + request.body_len + ); + } + } + } + /// Drain aggregator connections whose PID is no longer alive and persist /// them as `pending` records in `genai_events`. Rate-limited to once per /// second to avoid excessive `/proc` scanning. @@ -1991,6 +2045,8 @@ mod tests { model: Some("gpt-4".to_string()), provider: Some("openai".to_string()), call_kind: "main".to_string(), + pending_origin: crate::storage::sqlite::genai::PendingOrigin::RequestCapture, + pending_match_key: None, } }