From d47b1a6d2b21e108d0b13e00c51a6fb7b02d7752 Mon Sep 17 00:00:00 2001 From: wanghao <1562495626@qq.com> Date: Wed, 8 Jul 2026 15:35:11 +0800 Subject: [PATCH 1/2] fix(sight): persist idle streams and tool results Drain idle in-flight streamed calls so interrupted conversations remain visible as pending evidence. Preserve Anthropic tool_result blocks as structured input messages and add coverage for oversized request-body eviction. Assisted-by: Codex Signed-off-by: wanghao <1562495626@qq.com> --- .../src/aggregator/http/aggregator.rs | 148 ++++++++++++++++-- src/agentsight/src/aggregator/unified.rs | 8 + src/agentsight/src/genai/call_builder.rs | 138 +++++++++++++++- src/agentsight/src/probes/probes.rs | 2 +- src/agentsight/src/probes/sslsniff.rs | 2 +- src/agentsight/src/storage/sqlite/token.rs | 2 +- src/agentsight/src/storage/unified.rs | 2 +- src/agentsight/src/unified.rs | 52 ++++++ 8 files changed, 334 insertions(+), 20 deletions(-) diff --git a/src/agentsight/src/aggregator/http/aggregator.rs b/src/agentsight/src/aggregator/http/aggregator.rs index 25a0a943f..39baa9fbf 100644 --- a/src/agentsight/src/aggregator/http/aggregator.rs +++ b/src/agentsight/src/aggregator/http/aggregator.rs @@ -173,11 +173,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 +196,20 @@ 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); + 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 +221,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 } @@ -241,6 +242,50 @@ impl HttpConnectionAggregator { } } + /// Drain non-idle connections that exceeded `idle_timeout`. + /// + /// These states can represent manually interrupted or abandoned LLM + /// streams. Returning them lets the caller persist a pending GenAI row + /// instead of silently dropping the session evidence. + pub fn drain_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 { + self.sse_continuation_buffers.pop(&key); + self.last_appended_src_ptr.pop(&key); + self.last_activity.pop(&key); + if let Some(state) = self.connections.pop(&key) { + match state { + ConnectionState::Idle => {} + _ => result.push((key, state)), + } + } + } + + if !result.is_empty() { + log::info!( + "[HttpAggregator] drained {} idle in-flight connection(s) (timeout={}s)", + result.len(), + timeout.as_secs() + ); + } + + result + } + /// Record activity timestamp for a connection. fn touch(&mut self, key: &ConnectionId) { self.last_activity.push(*key, Instant::now()); @@ -2374,4 +2419,81 @@ mod tests { agg.evict_idle_and_oversized(); assert!(agg.connections.peek(&conn_id).is_none()); } + + #[test] + fn test_idle_eviction_preserves_sse_active_for_drain() { + let mut agg = HttpConnectionAggregator::with_limits(10, 8192, Duration::from_millis(50)); + let event = create_mock_ssl_event(1234, 0x7000); + 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)); + agg.evict_idle_and_oversized(); + let drained = agg.drain_connections_for_pid(1234); + + assert_eq!(drained.len(), 1); + assert!(matches!(drained[0].1, ConnectionState::SseActive { .. })); + } + + #[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..e4c28e6f7 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() } + + /// Drain 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 drain_idle_connections(&mut self) -> Vec<(ConnectionId, ConnectionState)> { + self.http.drain_idle_connections() + } } diff --git a/src/agentsight/src/genai/call_builder.rs b/src/agentsight/src/genai/call_builder.rs index 1dc45e08d..0b0cf54a0 100644 --- a/src/agentsight/src/genai/call_builder.rs +++ b/src/agentsight/src/genai/call_builder.rs @@ -291,9 +291,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 +401,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 +1005,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/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..fc4538c3e 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(); + // Drain 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,56 @@ impl AgentSight { } } + /// Drain 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 drained = self.aggregator.drain_idle_connections(); + if drained.is_empty() { + return; + } + + use crate::aggregator::ConnectionState; + + for (conn_id, state) in drained { + let (_state_name, request) = match state { + ConnectionState::RequestPending { request } => ("RequestPending", request), + ConnectionState::SseActive { + request: Some(req), .. + } => ("SseActive", req), + _ => continue, + }; + + if let Some(pending) = self.genai_builder.build_pending_from_request( + &request, + &conn_id, + &self.pid_agent_name_cache, + ) { + 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. From 183a6979ee7536859e73c5e2def465d5401dccaa Mon Sep 17 00:00:00 2001 From: wanghao <1562495626@qq.com> Date: Thu, 9 Jul 2026 15:43:03 +0800 Subject: [PATCH 2/2] fix(sight): preserve idle stream snapshots Idle SSE persistence now writes source metadata and a deterministic pending match key so timeout snapshots stay distinguishable from crash drains. The aggregator snapshots idle connections without removing in-memory stream state. Resumed streams can promote the stored pending row instead of losing the active SSE flow. The match key uses request metadata and body bytes. If the parser cannot reconstruct the same request shape, completion still falls back to the existing insert path. Assisted-by: Codex:GPT-5 Signed-off-by: wanghao <1562495626@qq.com> --- .../src/aggregator/http/aggregator.rs | 78 +++++-- src/agentsight/src/aggregator/unified.rs | 6 +- src/agentsight/src/background.rs | 4 +- src/agentsight/src/genai/builder.rs | 53 ++++- src/agentsight/src/genai/call_builder.rs | 8 + .../src/storage/sqlite/genai/mod.rs | 2 +- .../src/storage/sqlite/genai/pending.rs | 156 +++++++++++-- .../src/storage/sqlite/genai/schema.rs | 13 ++ .../src/storage/sqlite/genai/tests.rs | 215 +++++++++++++++++- src/agentsight/src/storage/sqlite/mod.rs | 2 +- src/agentsight/src/unified.rs | 16 +- 11 files changed, 506 insertions(+), 47 deletions(-) diff --git a/src/agentsight/src/aggregator/http/aggregator.rs b/src/agentsight/src/aggregator/http/aggregator.rs index 39baa9fbf..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, } @@ -203,6 +206,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); evicted_idle += 1; } } @@ -232,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!( @@ -242,12 +247,12 @@ impl HttpConnectionAggregator { } } - /// Drain non-idle connections that exceeded `idle_timeout`. + /// Snapshot non-idle connections that exceeded `idle_timeout`. /// /// These states can represent manually interrupted or abandoned LLM - /// streams. Returning them lets the caller persist a pending GenAI row - /// instead of silently dropping the session evidence. - pub fn drain_idle_connections(&mut self) -> Vec<(ConnectionId, ConnectionState)> { + /// 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 @@ -264,20 +269,23 @@ impl HttpConnectionAggregator { let mut result = Vec::new(); for key in keys { - self.sse_continuation_buffers.pop(&key); - self.last_appended_src_ptr.pop(&key); - self.last_activity.pop(&key); - if let Some(state) = self.connections.pop(&key) { - match state { + if self.idle_snapshotted.peek(&key).is_some() { + continue; + } + if let Some(state) = self.connections.peek(&key) { + match state.clone() { ConnectionState::Idle => {} - _ => result.push((key, state)), + state => { + self.idle_snapshotted.push(key, ()); + result.push((key, state)); + } } } } if !result.is_empty() { log::info!( - "[HttpAggregator] drained {} idle in-flight connection(s) (timeout={}s)", + "[HttpAggregator] snapshotted {} idle in-flight connection(s) (timeout={}s)", result.len(), timeout.as_secs() ); @@ -286,6 +294,11 @@ impl HttpConnectionAggregator { 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()); @@ -435,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 @@ -951,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) @@ -960,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() } @@ -984,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 => {} _ => { @@ -1044,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 @@ -2421,9 +2447,13 @@ mod tests { } #[test] - fn test_idle_eviction_preserves_sse_active_for_drain() { + fn test_idle_snapshot_keeps_sse_active_for_resumed_stream() { let mut agg = HttpConnectionAggregator::with_limits(10, 8192, Duration::from_millis(50)); - let event = create_mock_ssl_event(1234, 0x7000); + 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(), @@ -2450,11 +2480,23 @@ mod tests { assert!(agg.process_response(response).is_none()); std::thread::sleep(Duration::from_millis(60)); - agg.evict_idle_and_oversized(); - let drained = agg.drain_connections_for_pid(1234); + let snapshots = agg.snapshot_idle_connections(); + + assert_eq!(snapshots.len(), 1); + assert!(matches!(snapshots[0].1, ConnectionState::SseActive { .. })); - assert_eq!(drained.len(), 1); - assert!(matches!(drained[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] diff --git a/src/agentsight/src/aggregator/unified.rs b/src/agentsight/src/aggregator/unified.rs index e4c28e6f7..159e8a624 100644 --- a/src/agentsight/src/aggregator/unified.rs +++ b/src/agentsight/src/aggregator/unified.rs @@ -178,11 +178,11 @@ impl Aggregator { self.http.drain_dead_pid_connections() } - /// Drain in-flight HTTP connections that exceeded the idle timeout. + /// 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 drain_idle_connections(&mut self) -> Vec<(ConnectionId, ConnectionState)> { - self.http.drain_idle_connections() + 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 0b0cf54a0..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 }, }) 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/unified.rs b/src/agentsight/src/unified.rs index fc4538c3e..9e9159b3d 100644 --- a/src/agentsight/src/unified.rs +++ b/src/agentsight/src/unified.rs @@ -909,7 +909,7 @@ impl AgentSight { self.flush_expired_pending_genai(); // Drain orphaned connections from dead PIDs and persist as pending self.drain_and_persist_dead_connections(); - // Drain idle in-flight streams whose owning process is still alive. + // 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(); @@ -1280,19 +1280,20 @@ impl AgentSight { } } - /// Drain idle in-flight streams and persist them as `pending` records. + /// 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 drained = self.aggregator.drain_idle_connections(); - if drained.is_empty() { + 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 drained { + for (conn_id, state) in snapshots { let (_state_name, request) = match state { ConnectionState::RequestPending { request } => ("RequestPending", request), ConnectionState::SseActive { @@ -1301,11 +1302,12 @@ impl AgentSight { _ => continue, }; - if let Some(pending) = self.genai_builder.build_pending_from_request( + 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) { @@ -2043,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, } }