atagen

roundabout

public
Star0
Watch0
9117440444eb609ecea8476b5a62b03917077591
1
ssh://brewery@pls.lobotomise.me:33/atagen/roundabout.git
src/hook.rsRust425 lines13.7 KB
1
use std::{
2
    io::{self, Read},
3
    path::{Path, PathBuf},
4
};
5
6
use serde::Deserialize;
7
use serde_json::{Value, json};
8
9
use crate::{
10
    artifacts, client,
11
    config::{Config, StopPolicy, Unwrapped},
12
    history, paths,
13
    proto::{self, Event, LeaseInfo, Reader, Request},
14
    units::Dur,
15
};
16
17
#[derive(Deserialize)]
18
struct Input {
19
    #[serde(default)]
20
    session_id: String,
21
    #[serde(default)]
22
    cwd: Option<PathBuf>,
23
    #[serde(default)]
24
    tool_name: Option<String>,
25
    #[serde(default, rename = "tool_input")]
26
    tool_args: Option<Value>,
27
    #[serde(default)]
28
    stop_hook_active: bool,
29
}
30
31
fn read_input() -> Option<Input> {
32
    let mut buf = String::new();
33
    io::stdin().read_to_string(&mut buf).ok()?;
34
    serde_json::from_str(&buf).ok()
35
}
36
37
fn emit(v: &Value) {
38
    println!("{v}");
39
}
40
41
/// quote-aware split on `&& || ; | & \n`. subshells, heredocs and friends get
42
/// treated as words, which errs toward seeing a build that isn't there
43
fn segments(cmd: &str) -> Vec<Vec<String>> {
44
    let mut out = Vec::new();
45
    let mut words = Vec::new();
46
    let mut word = String::new();
47
    let mut quote: Option<char> = None;
48
    let mut chars = cmd.chars().peekable();
49
    let flush_word = |word: &mut String, words: &mut Vec<String>| {
50
        if !word.is_empty() {
51
            words.push(std::mem::take(word));
52
        }
53
    };
54
    while let Some(c) = chars.next() {
55
        match (quote, c) {
56
            (Some(q), c) if c == q => quote = None,
57
            (Some('"') | None, '\\') => {
58
                if let Some(n) = chars.next() {
59
                    word.push(n);
60
                }
61
            }
62
            (None, '\'' | '"') => quote = Some(c),
63
            (None, '&') if word.ends_with('>') || chars.peek() == Some(&'>') => word.push(c),
64
            (None, '&' | '|' | ';' | '\n' | '(' | ')') => {
65
                flush_word(&mut word, &mut words);
66
                if !words.is_empty() {
67
                    out.push(std::mem::take(&mut words));
68
                }
69
            }
70
            (None, c) if c.is_whitespace() => flush_word(&mut word, &mut words),
71
            (_, c) => word.push(c),
72
        }
73
    }
74
    flush_word(&mut word, &mut words);
75
    if !words.is_empty() {
76
        out.push(words);
77
    }
78
    out
79
}
80
81
/// peels env assignments and wrappers that don't change what runs
82
fn head(seg: &[String]) -> &[String] {
83
    let mut s = seg;
84
    loop {
85
        match s.first().map(String::as_str) {
86
            Some(w) if w.contains('=') && !w.starts_with('-') && !w.starts_with('=') => s = &s[1..],
87
            Some("time" | "nice" | "nohup" | "command" | "exec" | "env" | "sudo" | "doas") => {
88
                s = &s[1..];
89
            }
90
            Some("timeout") => {
91
                s = &s[1..];
92
                while s.first().is_some_and(|w| w.starts_with('-')) {
93
                    s = &s[1..];
94
                }
95
                s = s.get(1..).unwrap_or_default();
96
            }
97
            _ => return s,
98
        }
99
    }
100
}
101
102
fn is_roundabout(word: &str) -> bool {
103
    word == "roundabout" || word.ends_with("/roundabout")
104
}
105
106
fn matches_pattern(seg: &[String], pattern: &str) -> bool {
107
    let want: Vec<&str> = pattern.split_whitespace().collect();
108
    let first_ok = |w: &str, p: &str| w == p || w.rsplit('/').next() == Some(p);
109
    !want.is_empty()
110
        && seg.len() >= want.len()
111
        && first_ok(&seg[0], want[0])
112
        && seg[1..want.len()]
113
            .iter()
114
            .zip(&want[1..])
115
            .all(|(w, p)| w == p)
116
}
117
118
enum Scan {
119
    Clean,
120
    Unwrapped(Vec<String>),
121
}
122
123
fn scan(cfg: &Config, cmd: &str) -> Scan {
124
    let mut found = Vec::new();
125
    for seg in segments(cmd) {
126
        let seg = head(&seg);
127
        if seg.first().is_some_and(|w| is_roundabout(w)) {
128
            continue;
129
        }
130
        if cfg.detect_patterns().any(|p| matches_pattern(seg, p)) {
131
            let words: Vec<&str> = seg
132
                .iter()
133
                .map(String::as_str)
134
                .take_while(|w| !w.contains('>') && !w.contains('<'))
135
                .collect();
136
            found.push(words.join(" "));
137
        }
138
    }
139
    if found.is_empty() {
140
        Scan::Clean
141
    } else {
142
        Scan::Unwrapped(found)
143
    }
144
}
145
146
fn guess(cfg: &Config, cwd: &Path, cmd: &[String]) -> String {
147
    let key = history::key(&paths::project_root(cwd), cmd);
148
    history::lookup(&key).map_or_else(
149
        || format!("--size {}", cfg.enforce.default_size),
150
        |rec| {
151
            crate::run::base(cfg).map_or_else(
152
                |_| format!("--size {}", cfg.enforce.default_size),
153
                |b| {
154
                    let n = history::suggest(&rec, b, cmd);
155
                    format!("--ram {} --cores {} --time {}", n.ram, n.cores, Dur(n.time))
156
                },
157
            )
158
        },
159
    )
160
}
161
162
fn shell_quote(s: &str) -> String {
163
    format!("'{}'", s.replace('\'', r"'\''"))
164
}
165
166
fn pre_tool_use(cfg: &Config, input: Input) {
167
    if input.tool_name.as_deref() != Some("Bash") {
168
        return;
169
    }
170
    let Some(mut tool_input) = input.tool_args else {
171
        return;
172
    };
173
    let Some(cmd) = tool_input
174
        .get("command")
175
        .and_then(Value::as_str)
176
        .map(str::to_owned)
177
    else {
178
        return;
179
    };
180
    let cwd = input.cwd.unwrap_or_else(|| "/".into());
181
182
    match (scan(cfg, &cmd), cfg.enforce.unwrapped) {
183
        (Scan::Clean, _) | (_, Unwrapped::Allow) => {}
184
        (Scan::Unwrapped(found), Unwrapped::Deny) => {
185
            let first: Vec<String> = found[0].split_whitespace().map(Into::into).collect();
186
            let hint = guess(cfg, &cwd, &first);
187
            let reason = format!(
188
                "builds go through roundabout here so parallel agents don't oom the box. wrap each build \
189
                 in `roundabout run <size> -- <cmd>`, e.g. `roundabout run {hint} -- {}`. sizes are {}, or \
190
                 give --ram/--disk/--cores/--time yourself. for anything long use \
191
                 run_in_background and start `roundabout watch <lease>` in the background too so you \
192
                 hear if it drags. see the roundabout skill",
193
                found[0],
194
                cfg.sizes.keys().cloned().collect::<Vec<_>>().join("/"),
195
            );
196
            emit(&json!({
197
                "hookSpecificOutput": {
198
                    "hookEventName": "PreToolUse",
199
                    "permissionDecision": "deny",
200
                    "permissionDecisionReason": reason,
201
                }
202
            }));
203
        }
204
        (Scan::Unwrapped(_), Unwrapped::Wrap) => {
205
            let hint = guess(cfg, &cwd, std::slice::from_ref(&cmd));
206
            tool_input["command"] =
207
                Value::String(format!("roundabout run {hint} -- bash -c {}", shell_quote(&cmd)));
208
            // no permissionDecision on purpose. allow would skip the user's
209
            // prompt for whatever the agent wrote, rewriting shouldn't
210
            emit(&json!({
211
                "hookSpecificOutput": {
212
                    "hookEventName": "PreToolUse",
213
                    "updatedInput": tool_input,
214
                }
215
            }));
216
        }
217
    }
218
}
219
220
fn codex_pre_tool_use(cfg: &Config, input: Input) -> Option<Value> {
221
    if input.tool_name.as_deref() != Some("Bash") {
222
        return None;
223
    }
224
    let command = input.tool_args?.get("command")?.as_str()?.to_owned();
225
    let cwd = input.cwd.unwrap_or_else(|| "/".into());
226
    let Scan::Unwrapped(found) = scan(cfg, &command) else {
227
        return None;
228
    };
229
    if cfg.enforce.unwrapped == Unwrapped::Allow {
230
        return None;
231
    }
232
    let first: Vec<String> = found[0].split_whitespace().map(Into::into).collect();
233
    let hint = guess(cfg, &cwd, &first);
234
    let reason = format!(
235
        "builds go through roundabout here so parallel agents don't oom the box. run `roundabout run {hint} -- {}` instead. see the roundabout skill",
236
        found[0]
237
    );
238
    Some(json!({
239
        "hookSpecificOutput": {
240
            "hookEventName": "PreToolUse",
241
            "permissionDecision": "deny",
242
            "permissionDecisionReason": reason,
243
        }
244
    }))
245
}
246
247
fn overdue(session: &str) -> Vec<LeaseInfo> {
248
    if session.is_empty() {
249
        return Vec::new();
250
    }
251
    let Some(mut s) = client::connect() else {
252
        return Vec::new();
253
    };
254
    if proto::send(
255
        &mut s,
256
        &Request::Overdue {
257
            session: session.to_owned(),
258
        },
259
    )
260
    .is_err()
261
    {
262
        return Vec::new();
263
    }
264
    match Reader::new(&s).and_then(|mut r| r.recv::<Event>()) {
265
        Ok(Some(Event::Overdue { leases })) => leases,
266
        _ => Vec::new(),
267
    }
268
}
269
270
/// same source `roundabout run` reads. after `--continue` the env var can keep the
271
/// startup id while the json moves on, and we need both ends to agree
272
fn session(input: &Input, codex: bool) -> String {
273
    if codex {
274
        return input.session_id.clone();
275
    }
276
    std::env::var("CLAUDE_CODE_SESSION_ID")
277
        .ok()
278
        .filter(|s| !s.is_empty())
279
        .unwrap_or_else(|| input.session_id.clone())
280
}
281
282
/// both ids, since around /clear the env var and the json can disagree about
283
/// which session just ended. never spawns, a missing daemon holds no intents
284
fn session_end(input: &Input, codex: bool) {
285
    let mut ids = vec![session(input, codex), input.session_id.clone()];
286
    ids.dedup();
287
    for session in ids.into_iter().filter(|s| !s.is_empty()) {
288
        if let Some(mut s) = client::connect() {
289
            let _ = proto::send(
290
                &mut s,
291
                &Request::Expect {
292
                    session,
293
                    need: None,
294
                },
295
            );
296
            let _ = Reader::new(&s).and_then(|mut r| r.recv::<Event>());
297
        }
298
    }
299
}
300
301
fn remind(event: &str, input: &Input, codex: bool) {
302
    let late = overdue(&session(input, codex));
303
    if late.is_empty() {
304
        return;
305
    }
306
    let lines: Vec<String> = late
307
        .iter()
308
        .map(|l| {
309
            format!(
310
                "roundabout lease {} (`{}`) has run {} against a {} estimate. `roundabout kill {}` to stop it, \
311
                 `roundabout extend {} <dur>` if it's worth the wait",
312
                l.id,
313
                l.cmd.join(" "),
314
                Dur::secs(l.elapsed),
315
                Dur(l.need.time),
316
                l.id,
317
                l.id
318
            )
319
        })
320
        .collect();
321
    emit(&json!({
322
        "hookSpecificOutput": {
323
            "hookEventName": event,
324
            "additionalContext": lines.join("\n"),
325
        }
326
    }));
327
}
328
329
fn stop(cfg: &Config, event: &str, input: &Input) {
330
    if input.stop_hook_active || cfg.artifacts.stop == StopPolicy::Off {
331
        return;
332
    }
333
    let root = input.cwd.as_deref().map(paths::project_root);
334
    let report = artifacts::check(cfg, root.as_deref(), false);
335
    if !report.over() {
336
        return;
337
    }
338
    let msg = format!(
339
        "{}. run `{}`, or `roundabout artifacts list` and clear out what you don't need",
340
        report.describe(cfg),
341
        report.remedy()
342
    );
343
    let out = match cfg.artifacts.stop {
344
        StopPolicy::Block => json!({ "decision": "block", "reason": msg }),
345
        StopPolicy::Nudge => json!({
346
            "hookSpecificOutput": { "hookEventName": event, "additionalContext": msg }
347
        }),
348
        StopPolicy::Warn => json!({ "systemMessage": msg }),
349
        StopPolicy::Off => return,
350
    };
351
    emit(&out);
352
}
353
354
/// hooks never fail loudly. a broken roundabout shouldn't wedge the agent
355
pub fn main(cfg: &Config, event: &str) -> i32 {
356
    let Some(input) = read_input() else {
357
        return 0;
358
    };
359
    match event {
360
        "codex-pre-tool-use" => {
361
            if let Some(output) = codex_pre_tool_use(cfg, input) {
362
                emit(&output);
363
            }
364
        }
365
        "codex-post-tool-use" | "omp-post-tool-use" => remind("PostToolUse", &input, true),
366
        "codex-prompt-submit" | "omp-prompt-submit" => remind("UserPromptSubmit", &input, true),
367
        "codex-stop" | "omp-stop" | "stop" => stop(cfg, "Stop", &input),
368
        "codex-subagent-stop" | "omp-subagent-stop" | "subagent-stop" => {
369
            stop(cfg, "SubagentStop", &input);
370
        }
371
        "codex-session-end" | "omp-session-end" => session_end(&input, true),
372
        "pre-tool-use" | "omp-pre-tool-use" => pre_tool_use(cfg, input),
373
        "post-tool-use" => remind("PostToolUse", &input, false),
374
        "prompt-submit" => remind("UserPromptSubmit", &input, false),
375
        "session-end" => session_end(&input, false),
376
        other => eprintln!("roundabout: unknown hook {other}"),
377
    }
378
    0
379
}
380
381
#[cfg(test)]
382
mod tests {
383
    use super::*;
384
385
    fn bash(command: &str) -> Input {
386
        Input {
387
            session_id: String::new(),
388
            cwd: Some("/".into()),
389
            tool_name: Some("Bash".into()),
390
            tool_args: Some(json!({ "command": command })),
391
            stop_hook_active: false,
392
        }
393
    }
394
395
    #[test]
396
    fn codex_denies_unwrapped_builds() {
397
        let output = codex_pre_tool_use(&Config::default(), bash("cargo build")).unwrap();
398
        assert_eq!(output["hookSpecificOutput"]["permissionDecision"], "deny");
399
        assert!(
400
            output["hookSpecificOutput"]["permissionDecisionReason"]
401
                .as_str()
402
                .unwrap()
403
                .contains("roundabout run")
404
        );
405
    }
406
407
    #[test]
408
    fn codex_does_not_skip_approval_to_wrap_builds() {
409
        let mut cfg = Config::default();
410
        cfg.enforce.unwrapped = Unwrapped::Wrap;
411
        let output = codex_pre_tool_use(&cfg, bash("nix build .#example")).unwrap();
412
        assert_eq!(output["hookSpecificOutput"]["permissionDecision"], "deny");
413
        assert!(output["hookSpecificOutput"].get("updatedInput").is_none());
414
    }
415
416
    #[test]
417
    fn codex_ignores_other_tools_and_wrapped_builds() {
418
        assert!(
419
            codex_pre_tool_use(&Config::default(), bash("roundabout run -s m -- cargo build")).is_none()
420
        );
421
        let mut input = bash("cargo build");
422
        input.tool_name = Some("apply_patch".into());
423
        assert!(codex_pre_tool_use(&Config::default(), input).is_none());
424
    }
425
}