From a3fc6f8dd6417327b578b387febe25a9175cd24a Mon Sep 17 00:00:00 2001 From: Jiangtian Feng Date: Wed, 10 Jun 2026 16:04:01 +0800 Subject: [PATCH] feat(sight): add agent activity monitor via schedmon BPF MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace the cgroup-based idle-burst scheduler (which caused CFS weight starvation on wakeup) with a pure observation-only activity monitor. The schedmon BPF probe attaches to tp_btf/sched_switch and tp_btf/sched_wakeup to track per-thread sleep/wakeup events for traced agent processes. The userspace ActivityMonitor aggregates these into per-family idle/active state with configurable debounce threshold. No cgroup operations, no cpu.idle toggle, no weight manipulation — CPU scheduling policy belongs in the container spec, not in the observability layer. New files: - schedmon.bpf.c / schedmon.h: BPF tracepoint programs - probes/schedmon.rs: Rust wrapper with map reuse - scheduler/mod.rs: ActivityMonitor state machine (9 tests) Config: activity_monitor.enabled + idle_threshold_ms in agentsight.json Co-Authored-By: Claude Opus 4.6 (1M context) --- src/agentsight/build.rs | 4 + src/agentsight/src/bpf/common.h | 1 + src/agentsight/src/bpf/schedmon.bpf.c | 88 +++++++ src/agentsight/src/bpf/schedmon.h | 31 +++ src/agentsight/src/config.rs | 109 +++++++-- src/agentsight/src/event.rs | 9 +- src/agentsight/src/lib.rs | 75 +++--- src/agentsight/src/parser/unified.rs | 27 ++- src/agentsight/src/probes/mod.rs | 24 +- src/agentsight/src/probes/probes.rs | 108 ++++++--- src/agentsight/src/probes/schedmon.rs | 158 ++++++++++++ src/agentsight/src/scheduler/mod.rs | 331 ++++++++++++++++++++++++++ src/agentsight/src/unified.rs | 163 +++++++++---- 13 files changed, 957 insertions(+), 171 deletions(-) create mode 100644 src/agentsight/src/bpf/schedmon.bpf.c create mode 100644 src/agentsight/src/bpf/schedmon.h create mode 100644 src/agentsight/src/probes/schedmon.rs create mode 100644 src/agentsight/src/scheduler/mod.rs diff --git a/src/agentsight/build.rs b/src/agentsight/build.rs index 7d28903e0..5a70728f1 100644 --- a/src/agentsight/build.rs +++ b/src/agentsight/build.rs @@ -61,6 +61,10 @@ fn main() { // Generate tcpsniff skeleton (no header — reuses sslsniff.h event format) generate_skeleton(&mut out, "tcpsniff"); + + // Generate schedmon skeleton and bindings + generate_skeleton(&mut out, "schedmon"); + generate_header(&mut out, "schedmon"); // generate_header(&mut out, "frametypes"); // generate_header(&mut out, "errors"); diff --git a/src/agentsight/src/bpf/common.h b/src/agentsight/src/bpf/common.h index 77620464a..114ddeb1c 100644 --- a/src/agentsight/src/bpf/common.h +++ b/src/agentsight/src/bpf/common.h @@ -23,6 +23,7 @@ typedef enum { EVENT_SOURCE_FILEWATCH = 4, // File watch events (filewatch) EVENT_SOURCE_FILEWRITE = 5, // File write events (filewrite) EVENT_SOURCE_UDPDNS = 6, // UDP DNS query events (udpdns) + EVENT_SOURCE_SCHED = 7, // Scheduler monitor events (schedmon) } event_source_t; // Common event header - every ringbuffer event MUST start with this diff --git a/src/agentsight/src/bpf/schedmon.bpf.c b/src/agentsight/src/bpf/schedmon.bpf.c new file mode 100644 index 000000000..3ade30a01 --- /dev/null +++ b/src/agentsight/src/bpf/schedmon.bpf.c @@ -0,0 +1,88 @@ +// SPDX-License-Identifier: (LGPL-2.1 OR BSD-2-Clause) +// Copyright (c) 2025 AgentSight Project +// +// Scheduler monitor BPF program — detects idle/active state transitions for +// traced Agent processes via BTF-typed sched_switch / sched_wakeup tracepoints. +#include "vmlinux.h" +#include +#include +#include +#include "schedmon.h" +#include "common.h" + +struct pid_sched_state { + u8 last_state; +}; + +struct { + __uint(type, BPF_MAP_TYPE_LRU_HASH); + __uint(max_entries, MAX_TRACED_PROCESSES); + __type(key, u32); + __type(value, struct pid_sched_state); +} pid_sched_state SEC(".maps"); + +static __always_inline int +emit_sched_event(u32 tgid, u32 tid, u64 now, u8 event_type) +{ + struct pid_sched_state *st = bpf_map_lookup_elem(&pid_sched_state, &tid); + if (st) { + if (st->last_state == event_type) + return 0; + st->last_state = event_type; + } else { + struct pid_sched_state new_st = { + .last_state = event_type, + }; + bpf_map_update_elem(&pid_sched_state, &tid, &new_st, BPF_ANY); + } + + struct sched_event *ev = bpf_ringbuf_reserve(&rb, sizeof(*ev), 0); + if (!ev) + return 0; + + ev->source = EVENT_SOURCE_SCHED; + ev->tgid = tgid; + ev->tid = tid; + ev->timestamp_ns = now; + ev->event_type = event_type; + ev->pad[0] = 0; + ev->pad[1] = 0; + ev->pad[2] = 0; + + bpf_ringbuf_submit(ev, 0); + return 0; +} + +SEC("tp_btf/sched_switch") +int BPF_PROG(handle_sched_switch, bool preempt, struct task_struct *prev, + struct task_struct *next) +{ + if (preempt) + return 0; + + unsigned int state = BPF_CORE_READ(prev, __state); + if (state == 0) + return 0; + + u32 tgid = BPF_CORE_READ(prev, tgid); + u32 *traced = bpf_map_lookup_elem(&traced_processes, &tgid); + if (!traced) + return 0; + + u32 tid = BPF_CORE_READ(prev, pid); + return emit_sched_event(tgid, tid, bpf_ktime_get_ns(), SCHED_EVENT_SLEEP); +} + +SEC("tp_btf/sched_wakeup") +int BPF_PROG(handle_sched_wakeup, struct task_struct *p) +{ + u32 tgid = BPF_CORE_READ(p, tgid); + u32 *traced = bpf_map_lookup_elem(&traced_processes, &tgid); + if (!traced) + return 0; + + u32 tid = BPF_CORE_READ(p, pid); + return emit_sched_event(tgid, tid, bpf_ktime_get_ns(), SCHED_EVENT_WAKEUP); +} + +char LICENSE[] SEC("license") = "GPL"; diff --git a/src/agentsight/src/bpf/schedmon.h b/src/agentsight/src/bpf/schedmon.h new file mode 100644 index 000000000..6c80895ca --- /dev/null +++ b/src/agentsight/src/bpf/schedmon.h @@ -0,0 +1,31 @@ +// SPDX-License-Identifier: (LGPL-2.1 OR BSD-2-Clause) +// Copyright (c) 2025 AgentSight Project +#ifndef __SCHEDMON_H +#define __SCHEDMON_H + +#define TASK_COMM_LEN 16 + +typedef signed char s8; +typedef unsigned char u8; +typedef signed short s16; +typedef unsigned short u16; +typedef signed int s32; +typedef unsigned int u32; +typedef signed long long s64; +typedef unsigned long long u64; + +enum sched_event_type { + SCHED_EVENT_SLEEP = 1, + SCHED_EVENT_WAKEUP = 2, +}; + +struct sched_event { + u32 source; // EVENT_SOURCE_SCHED + u32 tgid; + u32 tid; + u64 timestamp_ns; + u8 event_type; // enum sched_event_type + u8 pad[3]; +}; + +#endif /* __SCHEDMON_H */ diff --git a/src/agentsight/src/config.rs b/src/agentsight/src/config.rs index 041da7ded..d52ab2444 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. @@ -215,7 +218,6 @@ impl FromStr for TcpTarget { } } - /// Internal JSON structures for parsing the config file (same format as FFI). #[derive(serde::Deserialize)] struct JsonFullConfig { @@ -241,6 +243,16 @@ struct JsonFullConfig { cgroup_filter_enabled: Option, #[serde(default)] cgroup_ids: Option>, + #[serde(default)] + activity_monitor: Option, +} + +#[derive(serde::Deserialize, Clone, Debug)] +struct JsonActivityMonitor { + #[serde(default)] + enabled: Option, + #[serde(default)] + idle_threshold_ms: Option, } /// DeadLoop 检测配置区段 @@ -331,7 +343,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(), + }); } } } @@ -357,13 +371,14 @@ 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))?; +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. @@ -384,8 +399,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 } @@ -502,6 +517,12 @@ pub struct AgentsightConfig { pub deadloop_kill_enabled: bool, /// 触发 kill 的循环次数阈值（检测到 N 次后 kill） pub deadloop_kill_after_count: usize, + + // --- Activity Monitor Configuration --- + /// Enable scheduler activity monitor (BPF schedmon probe) + pub enable_activity_monitor: bool, + /// Idle threshold in ms before a family transitions to Idle state + pub activity_idle_threshold_ms: u64, } impl Default for AgentsightConfig { @@ -542,8 +563,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(), @@ -562,6 +588,8 @@ impl Default for AgentsightConfig { // DeadLoop auto-kill defaults (disabled by default) deadloop_kill_enabled: false, deadloop_kill_after_count: 3, + enable_activity_monitor: false, + activity_idle_threshold_ms: 50, } } } @@ -645,8 +673,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; @@ -711,6 +739,15 @@ impl AgentsightConfig { self.cgroup_ids = ids; } + if let Some(ref am) = parsed.activity_monitor { + if let Some(enabled) = am.enabled { + self.enable_activity_monitor = enabled; + } + if let Some(ms) = am.idle_threshold_ms { + self.activity_idle_threshold_ms = ms; + } + } + let (cmdline_rules, https_rules, http_targets) = extract_rules(&parsed); self.cmdline_rules.extend(cmdline_rules); self.https_rules.extend(https_rules); @@ -770,7 +807,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() } } @@ -1015,14 +1055,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] @@ -1042,7 +1088,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); } @@ -1061,7 +1110,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"); @@ -1080,8 +1131,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); } @@ -1093,7 +1148,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")); @@ -1103,7 +1159,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()); // https rules: dashscope.aliyuncs.com configured by default assert_eq!(https_rules.len(), 1); diff --git a/src/agentsight/src/event.rs b/src/agentsight/src/event.rs index 57fe03566..81100fb2a 100644 --- a/src/agentsight/src/event.rs +++ b/src/agentsight/src/event.rs @@ -1,8 +1,9 @@ -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::schedmon::SchedEvent; +use crate::probes::sslsniff::SslEvent; use crate::probes::udpdns::UdpDnsEvent; /// Unified event type that can represent any probe event @@ -16,6 +17,7 @@ pub enum Event { FileWatch(FileWatchEvent), FileWrite(FileWriteEvent), UdpDns(UdpDnsEvent), + Sched(SchedEvent), } impl Event { @@ -28,6 +30,7 @@ impl Event { Event::FileWatch(_) => "FileWatch", Event::FileWrite(_) => "FileWrite", Event::UdpDns(_) => "UdpDns", + Event::Sched(_) => "Sched", } } } diff --git a/src/agentsight/src/lib.rs b/src/agentsight/src/lib.rs index c730a3d66..c0de2a694 100644 --- a/src/agentsight/src/lib.rs +++ b/src/agentsight/src/lib.rs @@ -24,63 +24,54 @@ //! 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 use config::{default_base_path, AgentsightConfig}; 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 utils; +pub mod parser; +pub mod response_map; +pub mod scheduler; #[cfg(feature = "server")] pub mod server; +pub mod skill_metrics; +pub mod storage; +pub mod tokenizer; mod unified; -pub mod ffi; +pub mod utils; // 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::{next_flow_id, ns_to_us, ChromeTraceEvent, ToChromeTraceEvent, TraceArgs}; +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, - format_tokens, format_tokens_with_commas, + format_tokens, format_tokens_with_commas, AuditStore, HttpStore, SqliteConfig, SqliteStore, + Storage, StorageBackend, TimePeriod, TokenBreakdown, TokenComparison, TokenQuery, + TokenQueryResult, TokenStore, Trend, }; // Re-export unified entry point @@ -93,12 +84,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/unified.rs b/src/agentsight/src/parser/unified.rs index 70e62ff8b..bd73d5c5b 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,21 @@ 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(), + }, + Event::Sched(_) => ParseResult { + messages: Vec::new(), + }, } } diff --git a/src/agentsight/src/probes/mod.rs b/src/agentsight/src/probes/mod.rs index ebd022ea1..c1e5d1c87 100644 --- a/src/agentsight/src/probes/mod.rs +++ b/src/agentsight/src/probes/mod.rs @@ -1,20 +1,20 @@ - - -pub mod sslsniff; -pub mod proctrace; -pub mod procmon; pub mod filewatch; pub mod filewrite; -pub mod udpdns; -pub mod tcpsniff; pub mod probes; +pub mod procmon; +pub mod proctrace; +pub mod schedmon; +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 schedmon::{SchedEvent, SchedMon}; +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 a47552f8a..36f39696f 100644 --- a/src/agentsight/src/probes/probes.rs +++ b/src/agentsight/src/probes/probes.rs @@ -8,22 +8,26 @@ use anyhow::{Context, Result}; use libbpf_rs::{MapHandle, RingBufferBuilder}; use std::{ mem, - sync::{Arc, atomic::{AtomicBool, Ordering}}, + sync::{ + atomic::{AtomicBool, Ordering}, + Arc, + }, 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::schedmon::{RawSchedEvent, SchedMon}; +use super::sslsniff::bpf::probe_SSL_data_t as RawSslEvent; +use super::sslsniff::SslSniff; use super::tcpsniff::TcpSniff; +use super::udpdns::{RawUdpDnsEvent, UdpDns}; +use crate::config::TcpTarget; const POLL_TIMEOUT_MS: u64 = 100; @@ -34,9 +38,10 @@ const EVENT_SOURCE_PROCMON: u32 = 3; const EVENT_SOURCE_FILEWATCH: u32 = 4; const EVENT_SOURCE_FILEWRITE: u32 = 5; const EVENT_SOURCE_UDPDNS: u32 = 6; +const EVENT_SOURCE_SCHED: u32 = 7; /// 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 @@ -57,6 +62,8 @@ pub struct Probes { udpdns: Option, /// TCP sniff probe (captures plain HTTP traffic on configured ports, optional) tcpsniff: Option, + /// Scheduler monitor probe (detects idle/active transitions, optional) + schedmon: Option, /// Shared ring buffer handle (cloned from proctrace) for polling rb_handle: MapHandle, /// Unified event channel - events are converted to Event type inside the poller @@ -87,6 +94,7 @@ impl Probes { enable_udpdns, tcp_targets, false, + false, ) } @@ -104,6 +112,7 @@ impl Probes { enable_udpdns: bool, tcp_targets: &[TcpTarget], cgroup_filter_enabled: bool, + enable_schedmon: bool, ) -> Result { // Create proctrace first - it owns the traced_processes map, the ring // buffer, and (when enabled) the cgroup_filter map. @@ -117,10 +126,10 @@ impl Probes { .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")?; // Only fetch a cgroup_filter handle when the feature is on; when off, // we let each probe load its own private (unused) cgroup_filter map @@ -141,8 +150,7 @@ impl Probes { .context("failed to create sslsniff")?; // Create procmon - it reuses the ring buffer (no cgroup filter: full audit) - 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 { @@ -181,8 +189,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) @@ -191,8 +199,21 @@ impl Probes { None }; + let schedmon = if enable_schedmon { + match SchedMon::new_with_maps(&map_handle, &rb_handle) { + Ok(sm) => Some(sm), + Err(e) => { + log::warn!("SchedMon probe disabled (requires BTF tp_btf): {e}"); + None + } + } + } else { + log::info!("SchedMon probe disabled"); + None + }; + let (event_tx, event_rx) = crossbeam_channel::unbounded(); - + Ok(Self { proctrace, sslsniff, @@ -201,6 +222,7 @@ impl Probes { filewrite, udpdns, tcpsniff, + schedmon, rb_handle, event_tx, event_rx, @@ -210,26 +232,29 @@ 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")?; + } + // Attach schedmon for scheduler activity monitoring (if enabled) + if let Some(ref mut sm) = self.schedmon { + sm.attach().context("failed to attach schedmon")?; } // sslsniff uses uprobes attached per-process via attach_process() Ok(()) @@ -242,7 +267,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(()) } @@ -258,6 +284,7 @@ impl Probes { let filewatch_event_size = mem::size_of::(); let filewrite_event_size = mem::size_of::(); let udpdns_event_size = mem::size_of::(); + let sched_event_size = mem::size_of::(); let event_tx = self.event_tx.clone(); let stop_flag = Arc::new(AtomicBool::new(false)); @@ -324,13 +351,19 @@ impl Probes { None } } + EVENT_SOURCE_SCHED => { + if data.len() >= sched_event_size { + super::schedmon::SchedEvent::from_bytes(data).map(Event::Sched) + } else { + None + } + } _ => { - // Unknown source - ignore log::warn!("probes: unknown event source {source}"); None } }; - + if let Some(e) = event { let _ = event_tx.send(e); } @@ -377,13 +410,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") } @@ -396,7 +431,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(()) } } @@ -414,13 +452,15 @@ impl Probes { /// proctrace / filewatch / filewrite. sslsniff, udpdns, and procmon are /// unaffected. pub fn add_traced_cgroup(&mut self, cgroup_id: u64) -> Result<()> { - self.proctrace.add_traced_cgroup(cgroup_id) + self.proctrace + .add_traced_cgroup(cgroup_id) .context("failed to add traced cgroup") } /// Remove a cgroup inode id from the shared cgroup_filter map at runtime. pub fn remove_traced_cgroup(&mut self, cgroup_id: u64) -> Result<()> { - self.proctrace.remove_traced_cgroup(cgroup_id) + self.proctrace + .remove_traced_cgroup(cgroup_id) .context("failed to remove traced cgroup") } } diff --git a/src/agentsight/src/probes/schedmon.rs b/src/agentsight/src/probes/schedmon.rs new file mode 100644 index 000000000..310314660 --- /dev/null +++ b/src/agentsight/src/probes/schedmon.rs @@ -0,0 +1,158 @@ +// SPDX-License-Identifier: (LGPL-2.1 OR BSD-2-Clause) +// Copyright (c) 2025 AgentSight Project + +use crate::config; +use anyhow::{Context, Result}; +use libbpf_rs::{ + skel::{OpenSkel, SkelBuilder}, + Link, MapHandle, +}; +use std::{mem::MaybeUninit, os::fd::AsFd}; + +#[allow( + non_camel_case_types, + non_upper_case_globals, + dead_code, + non_snake_case +)] +mod bpf { + include!(concat!(env!("OUT_DIR"), "/schedmon.skel.rs")); + include!(concat!(env!("OUT_DIR"), "/schedmon.rs")); +} +use bpf::*; + +pub type RawSchedEvent = bpf::sched_event; + +pub const SCHED_EVENT_SLEEP: u8 = 1; +pub const SCHED_EVENT_WAKEUP: u8 = 2; + +#[derive(Debug, Clone)] +pub struct SchedEvent { + pub tgid: u32, + pub tid: u32, + pub timestamp_ns: u64, + pub event_type: u8, +} + +impl SchedEvent { + pub fn from_bytes(data: &[u8]) -> Option { + if data.len() < std::mem::size_of::() { + return None; + } + + let raw = unsafe { &*(data.as_ptr() as *const RawSchedEvent) }; + + Some(SchedEvent { + tgid: raw.tgid, + tid: raw.tid, + timestamp_ns: config::ktime_to_unix_ns(raw.timestamp_ns), + event_type: raw.event_type, + }) + } + + pub fn is_sleep(&self) -> bool { + self.event_type == SCHED_EVENT_SLEEP + } + + pub fn is_wakeup(&self) -> bool { + self.event_type == SCHED_EVENT_WAKEUP + } +} + +pub struct SchedMon { + _open_object: Box>, + skel: Box>, + _links: Vec, +} + +impl SchedMon { + pub fn new_with_maps(traced_processes: &MapHandle, rb: &MapHandle) -> Result { + let mut builder = SchedmonSkelBuilder::default(); + builder.obj_builder.debug(config::verbose()); + + let open_object = Box::new(MaybeUninit::::uninit()); + let mut open_skel = builder + .open() + .context("failed to open schedmon BPF object")?; + + open_skel + .maps_mut() + .traced_processes() + .reuse_fd(traced_processes.as_fd()) + .context("failed to reuse traced_processes map")?; + + open_skel + .maps_mut() + .rb() + .reuse_fd(rb.as_fd()) + .context("failed to reuse rb map")?; + + let skel = open_skel + .load() + .context("failed to load schedmon BPF object")?; + + let skel = + unsafe { Box::from_raw(Box::into_raw(Box::new(skel)) as *mut SchedmonSkel<'static>) }; + + Ok(Self { + _open_object: open_object, + skel, + _links: Vec::new(), + }) + } + + pub fn attach(&mut self) -> Result<()> { + let mut links = Vec::new(); + + let link = self + .skel + .progs_mut() + .handle_sched_switch() + .attach() + .context("failed to attach sched_switch tracepoint")?; + links.push(link); + + let link = self + .skel + .progs_mut() + .handle_sched_wakeup() + .attach() + .context("failed to attach sched_wakeup tracepoint")?; + links.push(link); + + self._links = links; + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_sched_event_from_bytes_too_short() { + let data = [0u8; 4]; + assert!(SchedEvent::from_bytes(&data).is_none()); + } + + #[test] + fn test_sched_event_is_sleep_wakeup() { + let ev = SchedEvent { + tgid: 100, + tid: 100, + timestamp_ns: 12345, + event_type: SCHED_EVENT_SLEEP, + }; + assert!(ev.is_sleep()); + assert!(!ev.is_wakeup()); + + let ev2 = SchedEvent { + tgid: 100, + tid: 101, + timestamp_ns: 12345, + event_type: SCHED_EVENT_WAKEUP, + }; + assert!(ev2.is_wakeup()); + assert!(!ev2.is_sleep()); + } +} diff --git a/src/agentsight/src/scheduler/mod.rs b/src/agentsight/src/scheduler/mod.rs new file mode 100644 index 000000000..b2dbf4b44 --- /dev/null +++ b/src/agentsight/src/scheduler/mod.rs @@ -0,0 +1,331 @@ +use std::collections::{HashMap, HashSet}; +use std::time::Instant; + +use crate::probes::schedmon::{SCHED_EVENT_SLEEP, SCHED_EVENT_WAKEUP}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ActivityState { + Active, + Idle, +} + +#[derive(Debug, Clone)] +pub struct ActivityConfig { + pub enabled: bool, + pub idle_threshold_ms: u64, +} + +impl Default for ActivityConfig { + fn default() -> Self { + Self { + enabled: false, + idle_threshold_ms: 50, + } + } +} + +#[derive(Debug, Clone, Default)] +pub struct FamilyMetrics { + pub idle_to_active_count: u64, + pub active_to_idle_count: u64, + pub last_active_duration_ns: u64, + pub last_idle_duration_ns: u64, +} + +struct FamilyState { + member_pids: HashSet, + active_tids: HashSet, + state: ActivityState, + idle_since: Option, + last_transition: Instant, + metrics: FamilyMetrics, +} + +pub struct ActivityMonitor { + config: ActivityConfig, + families: HashMap, + pid_to_root: HashMap, +} + +impl ActivityMonitor { + pub fn new(config: ActivityConfig) -> Self { + Self { + config, + families: HashMap::new(), + pid_to_root: HashMap::new(), + } + } + + pub fn add_process(&mut self, pid: u32, root_pid: u32) { + self.pid_to_root.insert(pid, root_pid); + + let family = self + .families + .entry(root_pid) + .or_insert_with(|| FamilyState { + member_pids: HashSet::new(), + active_tids: HashSet::new(), + state: ActivityState::Active, + idle_since: None, + last_transition: Instant::now(), + metrics: FamilyMetrics::default(), + }); + + family.member_pids.insert(pid); + family.active_tids.insert(pid); + family.idle_since = None; + } + + pub fn remove_process(&mut self, pid: u32) { + let root_pid = match self.pid_to_root.remove(&pid) { + Some(r) => r, + None => return, + }; + + let should_remove = if let Some(family) = self.families.get_mut(&root_pid) { + family.member_pids.remove(&pid); + family.active_tids.remove(&pid); + family.member_pids.is_empty() + } else { + false + }; + + if should_remove { + if let Some(family) = self.families.remove(&root_pid) { + log::debug!( + "activity: family {root_pid} removed (idle_to_active={}, active_to_idle={})", + family.metrics.idle_to_active_count, + family.metrics.active_to_idle_count, + ); + } + } + } + + pub fn on_sched_event(&mut self, tgid: u32, tid: u32, event_type: u8) { + let root_pid = match self.pid_to_root.get(&tgid) { + Some(&r) => r, + None => return, + }; + + let family = match self.families.get_mut(&root_pid) { + Some(f) => f, + None => return, + }; + + match event_type { + SCHED_EVENT_WAKEUP => { + family.active_tids.insert(tid); + } + SCHED_EVENT_SLEEP => { + family.active_tids.remove(&tid); + } + _ => return, + } + + if family.active_tids.is_empty() { + if family.state == ActivityState::Active && family.idle_since.is_none() { + family.idle_since = Some(Instant::now()); + } + return; + } + + family.idle_since = None; + if family.state == ActivityState::Idle { + let idle_duration = family.last_transition.elapsed(); + family.state = ActivityState::Active; + family.last_transition = Instant::now(); + family.metrics.idle_to_active_count += 1; + family.metrics.last_idle_duration_ns = idle_duration.as_nanos() as u64; + log::debug!( + "activity: family {root_pid} IDLE -> ACTIVE (idle_ms={})", + idle_duration.as_millis(), + ); + } + } + + pub fn tick(&mut self) { + let threshold = std::time::Duration::from_millis(self.config.idle_threshold_ms); + + for (&root_pid, family) in &mut self.families { + if family.state == ActivityState::Idle || !family.active_tids.is_empty() { + continue; + } + match family.idle_since { + Some(t) if t.elapsed() >= threshold => {} + _ => continue, + } + + let active_duration = family.last_transition.elapsed(); + family.state = ActivityState::Idle; + family.last_transition = Instant::now(); + family.metrics.active_to_idle_count += 1; + family.metrics.last_active_duration_ns = active_duration.as_nanos() as u64; + log::debug!( + "activity: family {root_pid} ACTIVE -> IDLE (active_ms={})", + active_duration.as_millis(), + ); + } + } + + pub fn family_count(&self) -> usize { + self.families.len() + } + + pub fn family_state(&self, root_pid: u32) -> Option { + self.families.get(&root_pid).map(|f| f.state) + } + + pub fn family_metrics(&self, root_pid: u32) -> Option<&FamilyMetrics> { + self.families.get(&root_pid).map(|f| &f.metrics) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::time::Duration; + + fn test_config() -> ActivityConfig { + ActivityConfig { + enabled: true, + idle_threshold_ms: 10, + } + } + + fn make_monitor() -> ActivityMonitor { + ActivityMonitor::new(test_config()) + } + + fn force_idle_window(monitor: &mut ActivityMonitor, root_pid: u32) { + if let Some(family) = monitor.families.get_mut(&root_pid) { + family.idle_since = Some(Instant::now() - Duration::from_millis(20)); + } + } + + #[test] + fn test_add_remove_process() { + let mut mon = make_monitor(); + + mon.add_process(100, 100); + assert_eq!(mon.family_count(), 1); + assert_eq!(mon.pid_to_root.get(&100), Some(&100)); + + mon.add_process(200, 100); + assert_eq!(mon.family_count(), 1); + + mon.remove_process(200); + assert_eq!(mon.family_count(), 1); + + mon.remove_process(100); + assert_eq!(mon.family_count(), 0); + } + + #[test] + fn test_sched_event_active() { + let mut mon = make_monitor(); + mon.add_process(100, 100); + + assert_eq!(mon.family_state(100), Some(ActivityState::Active)); + + mon.on_sched_event(100, 100, SCHED_EVENT_SLEEP); + assert_eq!(mon.family_state(100), Some(ActivityState::Active)); + + mon.on_sched_event(100, 100, SCHED_EVENT_WAKEUP); + assert_eq!(mon.family_state(100), Some(ActivityState::Active)); + } + + #[test] + fn test_sched_event_idle_after_threshold() { + let mut mon = make_monitor(); + mon.add_process(100, 100); + + mon.on_sched_event(100, 100, SCHED_EVENT_SLEEP); + force_idle_window(&mut mon, 100); + + mon.tick(); + assert_eq!(mon.family_state(100), Some(ActivityState::Idle)); + assert_eq!(mon.family_metrics(100).unwrap().active_to_idle_count, 1); + } + + #[test] + fn test_wakeup_cancels_pending_idle() { + let mut mon = make_monitor(); + mon.add_process(100, 100); + + mon.on_sched_event(100, 100, SCHED_EVENT_SLEEP); + mon.on_sched_event(100, 100, SCHED_EVENT_WAKEUP); + force_idle_window(&mut mon, 100); + mon.tick(); + assert_eq!(mon.family_state(100), Some(ActivityState::Active)); + } + + #[test] + fn test_multithreaded_active_while_any_thread_runs() { + let mut mon = make_monitor(); + mon.add_process(100, 100); + + mon.on_sched_event(100, 101, SCHED_EVENT_WAKEUP); + mon.on_sched_event(100, 100, SCHED_EVENT_SLEEP); + assert_eq!(mon.family_state(100), Some(ActivityState::Active)); + + mon.on_sched_event(100, 101, SCHED_EVENT_SLEEP); + force_idle_window(&mut mon, 100); + mon.tick(); + assert_eq!(mon.family_state(100), Some(ActivityState::Idle)); + + mon.on_sched_event(100, 101, SCHED_EVENT_WAKEUP); + assert_eq!(mon.family_state(100), Some(ActivityState::Active)); + assert_eq!(mon.family_metrics(100).unwrap().idle_to_active_count, 1); + } + + #[test] + fn test_family_active_if_any_active() { + let mut mon = make_monitor(); + mon.add_process(100, 100); + mon.add_process(200, 100); + + assert_eq!(mon.family_state(100), Some(ActivityState::Active)); + + mon.on_sched_event(100, 100, SCHED_EVENT_SLEEP); + assert_eq!(mon.family_state(100), Some(ActivityState::Active)); + + mon.on_sched_event(200, 200, SCHED_EVENT_SLEEP); + assert_eq!(mon.family_state(100), Some(ActivityState::Active)); + + force_idle_window(&mut mon, 100); + mon.tick(); + assert_eq!(mon.family_state(100), Some(ActivityState::Idle)); + + mon.on_sched_event(200, 200, SCHED_EVENT_WAKEUP); + assert_eq!(mon.family_state(100), Some(ActivityState::Active)); + } + + #[test] + fn test_remove_nonexistent_process() { + let mut mon = make_monitor(); + mon.remove_process(999); + assert_eq!(mon.family_count(), 0); + } + + #[test] + fn test_sched_event_unknown_pid() { + let mut mon = make_monitor(); + mon.on_sched_event(999, 999, SCHED_EVENT_WAKEUP); + assert_eq!(mon.family_count(), 0); + } + + #[test] + fn test_multiple_families() { + let mut mon = make_monitor(); + mon.add_process(100, 100); + mon.add_process(200, 200); + assert_eq!(mon.family_count(), 2); + + mon.on_sched_event(100, 100, SCHED_EVENT_SLEEP); + force_idle_window(&mut mon, 100); + mon.tick(); + + assert_eq!(mon.family_state(100), Some(ActivityState::Idle)); + assert_eq!(mon.family_state(200), Some(ActivityState::Active)); + } +} diff --git a/src/agentsight/src/unified.rs b/src/agentsight/src/unified.rs index 1b57c76a6..6b3298dd2 100644 --- a/src/agentsight/src/unified.rs +++ b/src/agentsight/src/unified.rs @@ -21,8 +21,8 @@ use anyhow::{Context, Result}; use std::collections::{HashMap, HashSet}; use std::path::{Path, PathBuf}; -use std::sync::Arc; use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::Arc; use std::sync::Mutex; use crate::aggregator::Aggregator; @@ -33,10 +33,11 @@ use crate::event::Event; use crate::ffi::{FfiEvent, FfiEventSender}; use crate::genai::semantic::GenAISemanticEvent; use crate::genai::{GenAIBuilder, GenAIExporter, GenAIStore, LogtailExporter}; -use crate::interruption::{DetectorConfig, InterruptionDetector, recover_oom_events}; +use crate::interruption::{recover_oom_events, DetectorConfig, InterruptionDetector}; use crate::parser::Parser; use crate::probes::{FileWatchEvent, FileWriteEvent, Probes, ProbesPoller}; use crate::response_map::ResponseSessionMapper; +use crate::scheduler::{ActivityConfig, ActivityMonitor}; use crate::storage::sqlite::{GenAISqliteStore, InterruptionStore}; use crate::storage::{SqliteConfig, Storage, TimePeriod, TokenQuery, TokenQueryResult}; use crate::tokenizer::LlmTokenizer; @@ -103,6 +104,8 @@ pub struct AgentSight { deadloop_kill_enabled: bool, /// DeadLoop auto-kill: trigger threshold (kill after N detections) deadloop_kill_after_count: usize, + /// Activity monitor for idle/active state tracking via schedmon BPF + activity_monitor: Option, } /// GenAI events waiting for session_id resolution via ResponseSessionMapper. @@ -213,16 +216,16 @@ 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_with_cgroup_filter( - &[], - config.target_uid, - config.enable_filewatch, - enable_udpdns, - &tcp_targets, - config.cgroup_filter_enabled, - ) - .context("Failed to create probes")?; + let mut probes = Probes::new_with_cgroup_filter( + &[], + config.target_uid, + config.enable_filewatch, + enable_udpdns, + &tcp_targets, + config.cgroup_filter_enabled, + config.enable_activity_monitor, + ) + .context("Failed to create probes")?; // Attach procmon for process monitoring probes.attach().context("Failed to attach probes")?; @@ -230,7 +233,8 @@ impl AgentSight { // Seed cgroup_filter map with pre-configured cgroup inode IDs if config.cgroup_filter_enabled && !config.cgroup_ids.is_empty() { for &cg_id in &config.cgroup_ids { - probes.add_traced_cgroup(cg_id) + probes + .add_traced_cgroup(cg_id) .context("Failed to register cgroup_id")?; log::info!("Registered cgroup_id {}", cg_id); } @@ -504,6 +508,14 @@ impl AgentSight { pending_logtail, deadloop_kill_enabled: config.deadloop_kill_enabled, deadloop_kill_after_count: config.deadloop_kill_after_count, + activity_monitor: if config.enable_activity_monitor { + Some(ActivityMonitor::new(ActivityConfig { + enabled: true, + idle_threshold_ms: config.activity_idle_threshold_ms, + })) + } else { + None + }, }) } @@ -540,6 +552,9 @@ impl AgentSight { /// Attach SSL probes to a specific agent process pub fn attach_process(&mut self, pid: u32, agent_name: &str) { Self::attach_process_internal(&mut self.probes, pid, agent_name); + if let Some(ref mut monitor) = self.activity_monitor { + monitor.add_process(pid, pid); + } } /// Internal helper to attach SSL probes to a process @@ -562,6 +577,9 @@ impl AgentSight { log::error!("failed to delete {pid} from traced pid map: {e}"); }); self.probes.detach_ssl_probes(pid); + if let Some(ref mut monitor) = self.activity_monitor { + monitor.remove_process(pid); + } } /// Try to receive and process the next event (non-blocking) @@ -626,14 +644,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 + ); } } } @@ -641,7 +664,8 @@ impl AgentSight { Err(e) => { log::warn!( "[UDP-DNS] DNS resolve failed for http domain {}: {}", - dns_event.domain, e + dns_event.domain, + e ); } } @@ -650,6 +674,14 @@ impl AgentSight { return None; } + // Handle scheduler activity events + if let Event::Sched(ref sched_event) = event { + if let Some(ref mut monitor) = self.activity_monitor { + monitor.on_sched_event(sched_event.tgid, sched_event.tid, sched_event.event_type); + } + return None; + } + // Parse the event let result = self.parser.parse_event(event); @@ -671,7 +703,10 @@ impl AgentSight { 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() + t.agent = self + .pid_agent_name_cache + .get(&t.pid) + .cloned() .or_else(|| Some(t.comm.clone())); } } @@ -831,6 +866,10 @@ impl AgentSight { self.drain_and_persist_dead_connections(); // Check if config watcher deposited a new LogtailExporter self.check_pending_logtail(); + // Tick activity monitor to check idle thresholds + if let Some(ref mut monitor) = self.activity_monitor { + monitor.tick(); + } std::thread::sleep(std::time::Duration::from_millis(10)); } } @@ -860,7 +899,10 @@ impl AgentSight { fn check_pending_logtail(&mut self) { if let Ok(mut guard) = self.pending_logtail.try_lock() { if let Some(exporter) = guard.take() { - log::info!("Registering dynamically-activated LogtailExporter: '{}'", exporter.name()); + log::info!( + "Registering dynamically-activated LogtailExporter: '{}'", + exporter.name() + ); self.genai_exporters.push(exporter); } } @@ -879,7 +921,7 @@ impl AgentSight { encryption_pem: Option, trace_enabled: bool, ) { - use notify::{RecommendedWatcher, RecursiveMode, Watcher, Event as NotifyEvent, EventKind}; + use notify::{Event as NotifyEvent, EventKind, RecommendedWatcher, RecursiveMode, Watcher}; let watch_path = config_path.clone(); std::thread::Builder::new() @@ -924,9 +966,10 @@ impl AgentSight { } // Filter by filename - let is_target = event.paths.iter().any(|p| { - p.file_name().map(|f| f.to_os_string()) == target_filename - }); + let is_target = event + .paths + .iter() + .any(|p| p.file_name().map(|f| f.to_os_string()) == target_filename); if !is_target { continue; } @@ -1061,26 +1104,36 @@ impl AgentSight { cid, ie.interruption_type.as_str(), ); - if count >= 5 && ie.interruption_type != crate::interruption::InterruptionType::RetryStorm { - let storm_event = crate::interruption::InterruptionEvent::new( - crate::interruption::InterruptionType::RetryStorm, - ie.session_id.clone(), - ie.trace_id.clone(), - ie.conversation_id.clone(), - ie.call_id.clone(), - ie.pid, - ie.agent_name.clone(), - llm_call.end_timestamp_ns as i64, - Some(serde_json::json!({ - "repeated_type": ie.interruption_type.as_str(), - "count": count, - })), - ); - if !istore.exists_for_conversation(cid, &crate::interruption::InterruptionType::RetryStorm, None) { + if count >= 5 + && ie.interruption_type + != crate::interruption::InterruptionType::RetryStorm + { + let storm_event = + crate::interruption::InterruptionEvent::new( + crate::interruption::InterruptionType::RetryStorm, + ie.session_id.clone(), + ie.trace_id.clone(), + ie.conversation_id.clone(), + ie.call_id.clone(), + ie.pid, + ie.agent_name.clone(), + llm_call.end_timestamp_ns as i64, + Some(serde_json::json!({ + "repeated_type": ie.interruption_type.as_str(), + "count": count, + })), + ); + if !istore.exists_for_conversation( + cid, + &crate::interruption::InterruptionType::RetryStorm, + None, + ) { let _ = istore.insert(&storm_event); log::warn!( "RetryStorm detected: {} × {:?} in conversation {}", - count, ie.interruption_type, cid + count, + ie.interruption_type, + cid ); } } @@ -1136,10 +1189,13 @@ impl AgentSight { &recent, ) { let _ = istore.insert(&loop_event); - crate::genai::logtail::export_interruption_events(std::slice::from_ref(&loop_event)); + crate::genai::logtail::export_interruption_events( + std::slice::from_ref(&loop_event), + ); log::warn!( "DeadLoop detected in conversation {}: {:?}", - cid, loop_event.detail + cid, + loop_event.detail ); // ── Auto-kill 止血 ── @@ -1203,7 +1259,13 @@ impl AgentSight { use crate::genai::GenAIBuilder; // Track persisted pending calls: (pid, call_id, session_id, agent_name, conversation_id) - let mut persisted_pending: Vec<(u32, String, Option, Option, Option)> = Vec::new(); + let mut persisted_pending: Vec<( + u32, + String, + Option, + Option, + Option, + )> = Vec::new(); for (conn_id, state) in drained { // Destructure to capture both request AND sse_events @@ -1414,7 +1476,9 @@ impl AgentSight { // than the HealthChecker (30s cycle in serve process). if !persisted_pending.is_empty() { if let Some(ref istore) = self.interruption_store { - use crate::interruption::{InterruptionEvent, InterruptionType, was_pid_oom_killed}; + use crate::interruption::{ + was_pid_oom_killed, InterruptionEvent, InterruptionType, + }; let now_ns = std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) @@ -1427,7 +1491,8 @@ impl AgentSight { continue; // already checked this PID } if was_pid_oom_killed(*pid as i32) { - let call_ids: Vec<&str> = persisted_pending.iter() + let call_ids: Vec<&str> = persisted_pending + .iter() .filter(|(p, _, _, _, _)| *p == *pid) .map(|(_, c, _, _, _)| c.as_str()) .collect(); @@ -1454,13 +1519,19 @@ impl AgentSight { Some(detail), ); if let Err(e) = istore.insert(&event) { - log::warn!("[DrainCheck] Failed to record OOM agent_crash for pid={}: {}", pid, e); + log::warn!( + "[DrainCheck] Failed to record OOM agent_crash for pid={}: {}", + pid, + e + ); } else { log::info!("[DrainCheck] Recorded OOM agent_crash for pid={}", pid); } // Mark all pending calls for this PID as interrupted if let Some(ref store) = self.genai_sqlite_store { - if let Err(e) = store.mark_pending_interrupted_for_pid(*pid as i32, "oom_crash") { + if let Err(e) = + store.mark_pending_interrupted_for_pid(*pid as i32, "oom_crash") + { log::warn!("[DrainCheck] Failed to mark pending interrupted for pid={}: {}", pid, e); } }