ai-agent: inherit child proc lifecycle

Author: kylewanginchina

agent/src/common/proc_event/linux.rs

@@ -238,6 +238,9 @@ const PROC_LC_UID_OFF: usize = 9;
 const PROC_LC_GID_OFF: usize = 13;
 const PROC_LC_TS_OFF: usize = 17;
 const PROC_LC_COMM_OFF: usize = 25;
+pub const PROC_LIFECYCLE_FORK: u8 = 1;
+pub const PROC_LIFECYCLE_EXEC: u8 = 2;
+pub const PROC_LIFECYCLE_EXIT: u8 = 3;
 
 struct ProcLifecycleEventData {
     lifecycle_type: u8,
@@ -249,6 +252,13 @@ struct ProcLifecycleEventData {
     comm: Vec<u8>,
 }
 
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub struct ProcLifecycleInfo {
+    pub lifecycle_type: u8,
+    pub pid: u32,
+    pub parent_pid: u32,
+}
+
 impl TryFrom<&[u8]> for ProcLifecycleEventData {
     type Error = Error;
 
@@ -440,6 +450,17 @@ impl ProcEvent {
 
         Ok(BoxedProcEvents(Box::new(proc_event)))
     }
+
+    pub fn proc_lifecycle_info(&self) -> Option<ProcLifecycleInfo> {
+        match &self.event_data {
+            EventData::ProcLifecycleEvent(data) => Some(ProcLifecycleInfo {
+                lifecycle_type: data.lifecycle_type,
+                pid: data.pid,
+                parent_pid: data.parent_pid,
+            }),
+            _ => None,
+        }
+    }
 }
 
 #[derive(Debug)]
@@ -491,3 +512,37 @@ impl Sendable for BoxedProcEvents {
         SendMessageType::ProcEvents
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_proc_lifecycle_info_extracts_fields() {
+        let event_data = ProcLifecycleEventData {
+            lifecycle_type: 1,
+            pid: 4321,
+            parent_pid: 1234,
+            uid: 0,
+            gid: 0,
+            timestamp: 42,
+            comm: b"sleep".to_vec(),
+        };
+        let proc_event = ProcEvent {
+            pid: 1234,
+            pod_id: 0,
+            thread_id: 0,
+            coroutine_id: 0,
+            process_kname: b"python3".to_vec(),
+            start_time: 42,
+            end_time: 43,
+            event_type: EventType::ProcLifecycleEvent,
+            event_data: EventData::ProcLifecycleEvent(event_data),
+        };
+
+        let info = proc_event.proc_lifecycle_info().expect("missing info");
+        assert_eq!(info.lifecycle_type, 1);
+        assert_eq!(info.pid, 4321);
+        assert_eq!(info.parent_pid, 1234);
+    }
+}

agent/src/ebpf_dispatcher.rs

@@ -61,7 +61,7 @@ use crate::common::l7_protocol_log::{
     get_all_protocol, L7ProtocolBitmap, L7ProtocolParserInterface,
 };
 use crate::common::meta_packet::{MetaPacket, SegmentFlags};
-use crate::common::proc_event::{BoxedProcEvents, EventType, ProcEvent};
+use crate::common::proc_event::{BoxedProcEvents, EventType, ProcEvent, PROC_LIFECYCLE_FORK};
 use crate::common::{FlowAclListener, FlowAclListenerId};
 use crate::config::handler::{CollectorAccess, EbpfAccess, EbpfConfig, LogParserAccess};
 use crate::config::FlowAccess;
@@ -118,6 +118,19 @@ pub struct SyncEbpfCounter {
     counter: Arc<EbpfCounter>,
 }
 
+#[cfg(feature = "enterprise")]
+fn register_ai_agent_child(event: &BoxedProcEvents) {
+    if let Some(info) = event.0.proc_lifecycle_info() {
+        if info.lifecycle_type != PROC_LIFECYCLE_FORK {
+            return;
+        }
+        if let Some(registry) = enterprise_utils::ai_agent::global_registry() {
+            let now = Duration::from_nanos(event.0.start_time);
+            registry.register_child(info.parent_pid, info.pid, now);
+        }
+    }
+}
+
 impl OwnedCountable for SyncEbpfCounter {
     fn get_counters(&self) -> Vec<Counter> {
         let rx = self.counter.rx.swap(0, Ordering::Relaxed);
@@ -644,6 +657,8 @@ impl EbpfCollector {
                 if let Some(policy) = POLICY_GETTER.as_ref() {
                     event.0.pod_id = policy.lookup_pod_id(&container_id);
                 }
+                #[cfg(feature = "enterprise")]
+                register_ai_agent_child(&event);
                 if let Err(e) = PROC_EVENT_SENDER.as_mut().unwrap().send(event) {
                     warn!("event send ebpf error: {:?}", e);
                 }