Skip to main content

caboodle/agent_memory/
hooks.rs

1use super::{
2    lifecycle::{load, private_write},
3    BUNDLE_NAME,
4};
5use anyhow::{bail, Context, Result};
6use serde_json::{json, Value};
7use std::{
8    fs,
9    io::{Read, Write},
10    path::Path,
11    process::{Command, Stdio},
12    time::{Duration, Instant, SystemTime, UNIX_EPOCH},
13};
14
15pub fn quote(value: &str) -> String {
16    format!("'{}'", value.replace('\'', "'\"'\"'"))
17}
18pub fn bundle(directory: &Path, mode: super::RedirectMode) -> Value {
19    let mut hooks = Vec::new();
20    for harness in ["claude", "codex"] {
21        for event in [
22            "SessionStart",
23            "UserPromptSubmit",
24            "PreToolUse",
25            "Stop",
26            "SessionEnd",
27        ] {
28            if (harness == "codex" && event == "SessionEnd")
29                || (mode == super::RedirectMode::Off && event == "PreToolUse")
30            {
31                continue;
32            }
33            let base = format!(
34                "tambo memory hook --plan {} --harness {harness} --event {event}",
35                quote(&directory.display().to_string())
36            );
37            let evidence = format!(
38                "tambo memory evidence --plan {} --harness {harness} --event {event}",
39                quote(&directory.display().to_string())
40            );
41            let mut hook = json!({"event":event,"command":base,"timeout":10,"harnesses":[harness],
42                "evidence":{"command":evidence,"max_age_seconds":3600}});
43            if event == "PreToolUse" {
44                hook["matcher"] = json!(if harness == "claude" {
45                    "Write|Edit|MultiEdit"
46                } else {
47                    "apply_patch"
48                });
49            }
50            hooks.push(hook);
51        }
52    }
53    json!({"schema":"st.hook-bundle/1","name":BUNDLE_NAME,"version":"1.0.0","owner":"tambo","roles":["*"],"hooks":hooks})
54}
55fn valid(harness: &str, event: &str) -> bool {
56    ["claude", "codex"].contains(&harness)
57        && [
58            "SessionStart",
59            "UserPromptSubmit",
60            "PreToolUse",
61            "Stop",
62            "SessionEnd",
63        ]
64        .contains(&event)
65        && !(harness == "codex" && event == "SessionEnd")
66}
67fn receipt_path(directory: &Path, harness: &str, event: &str) -> Result<std::path::PathBuf> {
68    if !valid(harness, event) {
69        bail!("unsupported memory hook event/harness");
70    }
71    Ok(directory
72        .join("receipts")
73        .join(format!("{harness}-{event}.json")))
74}
75fn applied_at(directory: &Path) -> Result<f64> {
76    let state: Value = serde_json::from_slice(&fs::read(directory.join("state.json"))?)?;
77    if state["schema"] != 1 || state["phase"] != "applied" {
78        bail!("memory plan is not applied");
79    }
80    state["at"]
81        .as_f64()
82        .filter(|n| n.is_finite() && *n > 0.0)
83        .context("no application time")
84}
85pub fn record(
86    directory: &Path,
87    harness: &str,
88    event: &str,
89    session: &str,
90    exit: i32,
91) -> Result<()> {
92    let plan = load(directory)?;
93    applied_at(directory)?;
94    if session.is_empty() || session.len() > 128 || session.chars().any(char::is_control) {
95        bail!("hook receipt has no usable session identity");
96    }
97    if exit != 0 && !(event == "PreToolUse" && exit == 2) {
98        bail!("unsuccessful hook cannot produce firing evidence");
99    }
100    let path = receipt_path(directory, harness, event)?;
101    let now = SystemTime::now().duration_since(UNIX_EPOCH)?.as_secs_f64();
102    private_write(
103        &path,
104        &serde_json::to_vec(
105            &json!({"schema":1,"plan":plan.id,"harness":harness,"event":event,
106        "session":session,"completed_at":now,"exit_code":exit}),
107        )?,
108    )
109}
110pub fn evidence(directory: &Path, harness: &str, event: &str) -> Result<f64> {
111    let plan = load(directory)?;
112    let path = receipt_path(directory, harness, event)?;
113    let meta = fs::symlink_metadata(&path)?;
114    if !meta.is_file() || meta.file_type().is_symlink() || meta.len() > 4096 {
115        bail!("unreadable firing receipt");
116    }
117    let value: Value = serde_json::from_slice(&fs::read(path)?)?;
118    let now = SystemTime::now().duration_since(UNIX_EPOCH)?.as_secs_f64();
119    let stamp = value["completed_at"]
120        .as_f64()
121        .context("receipt has no timestamp")?;
122    let code = value["exit_code"]
123        .as_i64()
124        .context("receipt has no exit code")?;
125    if value["schema"] != 1
126        || value["plan"] != plan.id
127        || value["harness"] != harness
128        || value["event"] != event
129        || !value["session"]
130            .as_str()
131            .is_some_and(|s| !s.is_empty() && s.len() <= 128)
132        || !stamp.is_finite()
133        || stamp < applied_at(directory)?
134        || stamp > now + 5.0
135        || !(code == 0 || event == "PreToolUse" && code == 2)
136    {
137        bail!("firing receipt does not match the applied memory plan");
138    }
139    Ok(stamp)
140}
141pub(crate) fn producer_calls(
142    harness: &str,
143    event: &str,
144    mode: super::RedirectMode,
145) -> Result<Vec<Vec<&'static str>>> {
146    if !valid(harness, event) {
147        bail!("unsupported memory hook");
148    }
149    if event == "PreToolUse" && mode == super::RedirectMode::Off {
150        return Ok(vec![]);
151    }
152    let mut calls = match event {
153        "SessionStart" => vec![vec!["flush", "--hook"]],
154        "UserPromptSubmit" => vec![vec!["hook", "user-prompt-submit"]],
155        "PreToolUse" => vec![vec!["hook", "memory-redirect", "--mode", mode.as_str()]],
156        "Stop" => vec![vec!["hook", "stop"]],
157        "SessionEnd" => vec![vec!["hook", "session-end"]],
158        _ => bail!("unsupported event"),
159    };
160    // Codex has no observed SessionEnd hook. Consolidate on its supported
161    // Stop event, using the producer's existing bounded detached path.
162    if harness == "codex" && event == "Stop" {
163        calls.push(vec!["hook", "session-end"]);
164    }
165    Ok(calls)
166}
167
168/// Run only supported producer hooks. The receipt is metadata, never payload
169/// text. A child failure/timeout is fail-open and cannot create success evidence.
170pub fn run(directory: &Path, harness: &str, event: &str) -> i32 {
171    fn inner(directory: &Path, harness: &str, event: &str) -> Result<i32> {
172        let plan = load(directory)?;
173        applied_at(directory)?;
174        if event == "PreToolUse" && plan.redirect_mode == super::RedirectMode::Off {
175            return Ok(0);
176        }
177        receipt_path(directory, harness, event)?;
178        let mut bytes = Vec::new();
179        std::io::stdin()
180            .take(4 * 1024 * 1024 + 1)
181            .read_to_end(&mut bytes)?;
182        if bytes.len() > 4 * 1024 * 1024 {
183            bail!("hook input exceeds bound");
184        }
185        let payload: Value = serde_json::from_slice(&bytes)?;
186        let session = payload["session_id"]
187            .as_str()
188            .or_else(|| payload["sessionId"].as_str())
189            .context("no session identity")?;
190        let calls = producer_calls(harness, event, plan.redirect_mode)?;
191        let deadline = Instant::now() + Duration::from_secs(8);
192        let mut code = 0;
193        for args in calls {
194            let mut child = Command::new("cy")
195                .args(args)
196                .stdin(Stdio::piped())
197                .stdout(Stdio::inherit())
198                .stderr(Stdio::inherit())
199                .spawn()?;
200            let mut input = child.stdin.take().context("child stdin unavailable")?;
201            let body = bytes.clone();
202            let writer = std::thread::spawn(move || input.write_all(&body));
203            loop {
204                if let Some(status) = child.try_wait()? {
205                    code = status.code().unwrap_or(1);
206                    break;
207                }
208                if Instant::now() >= deadline {
209                    let _ = child.kill();
210                    let _ = child.wait();
211                    let _ = writer.join();
212                    bail!("producer hook deadline");
213                }
214                std::thread::sleep(Duration::from_millis(20));
215            }
216            let _ = writer.join();
217            if code != 0 {
218                break;
219            }
220        }
221        record(directory, harness, event, session, code)?;
222        Ok(if event == "PreToolUse" && code == 2 {
223            2
224        } else {
225            0
226        })
227    }
228    inner(directory, harness, event).unwrap_or(0)
229}