malvin 0.2.4

Non-interactive research and coding agent
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149

use std::collections::HashSet;
use std::path::Path;
use std::sync::Arc;
use std::time::Duration;

use tracing::warn;

const POLL_INTERVAL: Duration = if cfg!(test) {
    Duration::from_millis(10)
} else {
    Duration::from_millis(500)
};
const MAX_CONSECUTIVE_RSS_SAMPLE_FAILURES: u32 = 3;

pub struct MemWatchHandles {
    pub reader_dead: Arc<std::sync::atomic::AtomicBool>,
    pub pgid: Option<u32>,
    pub limit_bytes: u64,
    pub spawn_pid_baseline: HashSet<u32>,
    pub run_dir: Option<std::path::PathBuf>,
}

#[cfg(unix)]
pub async fn watch_process_group_memory(handles: MemWatchHandles) {
    watch_process_group_memory_with_rss_sampler(handles, |pgid, baseline| {
        crate::malvin_sandbox::malvin_session_rss_bytes(pgid, baseline)
    })
    .await;
}

#[cfg(unix)]
pub async fn watch_process_group_memory_with_rss_sampler(
    handles: MemWatchHandles,
    sample_rss: fn(Option<u32>, &HashSet<u32>) -> Option<u64>,
) {
    watch_session_memory_with_rss_sampler(handles, sample_rss).await;
}

#[cfg(unix)]
async fn watch_session_memory_with_rss_sampler(
    handles: MemWatchHandles,
    sample_rss: fn(Option<u32>, &HashSet<u32>) -> Option<u64>,
) {
    let MemWatchHandles {
        reader_dead,
        pgid: watch_pgid,
        limit_bytes,
        spawn_pid_baseline,
        run_dir,
    } = handles;
    let mut consecutive_rss_failures = 0u32;
    loop {
        if !crate::malvin_sandbox::sandbox_still_alive(watch_pgid, &spawn_pid_baseline) {
            return;
        }
        let rss = sample_rss(watch_pgid, &spawn_pid_baseline);
        let allow_fail_closed = !reader_dead.load(std::sync::atomic::Ordering::SeqCst);
        if memory_watch_should_terminate(
            rss,
            limit_bytes,
            &mut consecutive_rss_failures,
            allow_fail_closed,
        ) {
            let (reason, rss_bytes) = rss.map_or_else(
                || {
                    warn!(
                        limit_bytes,
                        pgid = watch_pgid,
                        consecutive_failures = consecutive_rss_failures,
                        "malvin sandbox cannot measure memory; terminating (fail-closed)"
                    );
                    (crate::sandbox_oom::OOM_REASON_MEASUREMENT_FAIL_CLOSED, None)
                },
                |rss_bytes| {
                    warn!(
                        rss_bytes,
                        limit_bytes,
                        pgid = watch_pgid,
                        "malvin sandbox exceeded memory limit; terminating"
                    );
                    (crate::sandbox_oom::OOM_REASON_MEMORY_LIMIT, Some(rss_bytes))
                },
            );
            record_sandbox_oom_marker(
                run_dir.as_deref(),
                crate::sandbox_oom::SandboxOomKillFacts {
                    reason,
                    rss_bytes,
                    limit_bytes,
                    pgid: watch_pgid.unwrap_or(0),
                },
            );
            crate::acp::unix_process_group_teardown::terminate_agent_process_group(
                watch_pgid,
                &spawn_pid_baseline,
            )
            .await;
            return;
        }
        tokio::time::sleep(POLL_INTERVAL).await;
    }
}

#[cfg(unix)]
fn record_sandbox_oom_marker(
    run_dir: Option<&Path>,
    facts: crate::sandbox_oom::SandboxOomKillFacts<'_>,
) {
    let Some(run_dir) = run_dir else {
        return;
    };
    let gate_iteration = crate::gate_loop_session::active_gate_iteration().unwrap_or(0);
    let record = crate::sandbox_oom::SandboxOomKillRecord::from_facts(gate_iteration, facts);
    if let Err(e) = crate::sandbox_oom::record_sandbox_oom_kill(run_dir, record) {
        warn!(error = %e, "failed to write sandbox OOM marker");
    }
}

#[cfg(unix)]
#[allow(clippy::missing_const_for_fn)]
fn memory_watch_should_terminate(
    rss: Option<u64>,
    limit_bytes: u64,
    consecutive_failures: &mut u32,
    allow_fail_closed: bool,
) -> bool {
    if let Some(bytes) = rss {
        *consecutive_failures = 0;
        bytes > limit_bytes
    } else if allow_fail_closed {
        *consecutive_failures = consecutive_failures.saturating_add(1);
        *consecutive_failures >= MAX_CONSECUTIVE_RSS_SAMPLE_FAILURES
    } else {
        *consecutive_failures = 0;
        false
    }
}

#[cfg(test)]
#[path = "process_group_mem_watch_tests.rs"]
mod process_group_mem_watch_tests;

#[cfg(all(test, unix))]
#[path = "process_group_mem_watch_policy_tests.rs"]
mod process_group_mem_watch_policy_tests;

#[cfg(test)]
#[path = "process_group_mem_watch_oom_marker_tests.rs"]
mod process_group_mem_watch_oom_marker_tests;