diff --git a/src/harness/claude.rs b/src/harness/claude.rs index 03f5b89..c51a8cb 100644 --- a/src/harness/claude.rs +++ b/src/harness/claude.rs @@ -7,13 +7,13 @@ use std::{ fs, path::{Path, PathBuf}, - time::{SystemTime, UNIX_EPOCH}, + time::SystemTime, }; use super::summary::AWAITING_APPROVAL; use super::{ - CAPTURE_ENV, CapturePaths, Harness, Invocation, SpawnPlan, is_uuid, last_hint, pin_plan, - shell_quote, within_window, within_window_ms, + CAPTURE_ENV, CapturePaths, Harness, Invocation, SpawnPlan, capture_id, is_uuid, last_hint, + pin_plan, same_cwd, shell_quote, sole_id, unix_millis, within_window, within_window_ms, }; pub struct Claude; @@ -50,9 +50,7 @@ impl Harness for Claude { } fn parse_capture(&self, payload: &str) -> Option { - let v = jzon::parse(payload).ok()?; - let id = v["session_id"].as_str()?; - is_uuid(id).then(|| id.to_string()) + capture_id(&jzon::parse(payload).ok()?, "session_id") } fn scrape_exit(&self, text: &str) -> Option { @@ -148,11 +146,10 @@ fn record_for_pid( let dir = Claude.home_root(home)?.join("sessions"); let text = fs::read_to_string(dir.join(format!("{pid}.json"))).ok()?; let rec = parse_record(&text)?; - let spawned_ms = spawned.duration_since(UNIX_EPOCH).ok()?.as_millis(); - // Test literal equality before canonicalizing the task path: identical - // nonexistent paths remain eligible. - let same_cwd = rec.cwd == cwd || cwd.canonicalize().is_ok_and(|c| rec.cwd == c); - (rec.pid == pid && same_cwd && within_window_ms(rec.started_at, spawned_ms)).then_some(rec) + (rec.pid == pid + && same_cwd(&rec.cwd, cwd, cwd.canonicalize().ok().as_deref()) + && within_window_ms(rec.started_at, unix_millis(spawned)?)) + .then_some(rec) } /// Return the UUID stem of the sole `.jsonl` transcript created within @@ -177,10 +174,7 @@ fn unique_in_window(dir: PathBuf, spawned: SystemTime) -> Option { } candidates.push(name.to_string()); } - match candidates.as_slice() { - [only] if is_uuid(only) => Some(only.clone()), - _ => None, - } + sole_id(candidates) } /// Convert an absolute working directory to Claude's project slug by replacing @@ -236,7 +230,7 @@ mod tests { /// Convert epoch milliseconds to [`SystemTime`]. fn at_ms(ms: u64) -> SystemTime { - UNIX_EPOCH + std::time::Duration::from_millis(ms) + std::time::UNIX_EPOCH + std::time::Duration::from_millis(ms) } /// Claude-specific opaque shapes: flags, `--continue`/`-c`, subcommands, diff --git a/src/harness/codex.rs b/src/harness/codex.rs index 0412e3f..385911a 100644 --- a/src/harness/codex.rs +++ b/src/harness/codex.rs @@ -4,17 +4,12 @@ //! correlates rollout files under `/sessions/YYYY/MM/DD/`. //! Missing, empty, and malformed rollouts do not produce a candidate. -use std::{ - fmt::Write as _, - fs, - io::{BufRead, BufReader, Read}, - path::Path, - time::SystemTime, -}; +use std::{fmt::Write as _, fs, path::Path, time::SystemTime}; use super::{ - CAPTURE_ENV, CapturePaths, Harness, Invocation, NOTIFY_CHAIN_ENV, SpawnPlan, is_uuid, - last_hint, leading_uuid, shell_quote, v7_millis, within_window_ms, + CAPTURE_ENV, CapturePaths, Harness, Invocation, NOTIFY_CHAIN_ENV, SpawnPlan, capture_id, + is_uuid, jsonl_head, last_hint, leading_uuid, push_unique, same_cwd, shell_quote, sole_id, + unix_millis, v7_millis, within_window_ms, }; pub struct Codex; @@ -72,8 +67,7 @@ impl Harness for Codex { if v["type"].as_str() != Some("agent-turn-complete") { return None; } - let id = v["thread-id"].as_str()?; - is_uuid(id).then(|| id.to_string()) + capture_id(&v, "thread-id") } fn scrape_exit(&self, text: &str) -> Option { @@ -108,10 +102,7 @@ impl Harness for Codex { fn correlate_fs(&self, cwd: &Path, spawned: SystemTime, home: Option<&Path>) -> Option { let root = self.home_root(home)?; - let spawn_ms = spawned - .duration_since(SystemTime::UNIX_EPOCH) - .ok()? - .as_millis(); + let spawn_ms = unix_millis(spawned)?; // Day directories are named by LOCAL date, which std cannot compute // without a timezone database. The UTC date differs from it by at // most one day, so probing the UTC date ±2 covers local ±1. @@ -157,16 +148,10 @@ impl Harness for Codex { } // Multiple rollouts may name the same thread. Correlation // counts that thread once. - let id = id.to_string(); - if !survivors.contains(&id) { - survivors.push(id); - } + push_unique(&mut survivors, id.to_string()); } } - match survivors.as_slice() { - [only] => Some(only.clone()), - _ => None, - } + sole_id(survivors) } } @@ -298,20 +283,11 @@ fn toml_escape(s: &str) -> String { /// Check the rollout's first record for a matching `cwd` and no explicit /// spawned-thread provenance. Missing and unrecognized `thread_source` values /// remain eligible; `"subagent"` or any `parent_thread_id` rejects the record. -/// Reads stop at 64 KiB because later records do not participate in -/// correlation. fn line1_admits(path: &Path, cwd: &Path) -> bool { - let Ok(file) = fs::File::open(path) else { + let Some(records) = jsonl_head(path, 1) else { return false; }; - let mut line = String::new(); - if BufReader::new(file.take(64 * 1024)) - .read_line(&mut line) - .is_err() - { - return false; - } - let Ok(meta) = jzon::parse(&line) else { + let Some(meta) = records[0].as_ref() else { return false; }; let payload = &meta["payload"]; @@ -322,10 +298,9 @@ fn line1_admits(path: &Path, cwd: &Path) -> bool { } // Rollouts can contain the physical cwd while the task retains a symlinked // path. Canonicalize the task path before rejecting the match. - payload["cwd"].as_str().is_some_and(|c| { - let recorded = Path::new(c); - recorded == cwd || cwd.canonicalize().is_ok_and(|p| p == recorded) - }) + payload["cwd"] + .as_str() + .is_some_and(|c| same_cwd(Path::new(c), cwd, cwd.canonicalize().ok().as_deref())) } #[cfg(test)] diff --git a/src/harness/grok.rs b/src/harness/grok.rs index bfca0e0..9b58a9b 100644 --- a/src/harness/grok.rs +++ b/src/harness/grok.rs @@ -12,7 +12,8 @@ use std::{ }; use super::{ - CapturePaths, Harness, Invocation, SpawnPlan, is_uuid, last_hint, pin_plan, within_window, + CapturePaths, Harness, Invocation, SpawnPlan, last_hint, pin_plan, push_unique, sole_id, + within_window, }; pub struct Grok; @@ -40,11 +41,6 @@ impl Harness for Grok { pin_plan(inv) } - /// Grok has no injected live capture channel. - fn parse_capture(&self, _payload: &str) -> Option { - None - } - fn scrape_exit(&self, text: &str) -> Option { // The last valid short or long resume hint names the conversation. last_hint(text, &["grok -r ", "grok --resume "]) @@ -121,12 +117,6 @@ fn cwd_record(text: &str) -> &str { .unwrap_or(text) } -fn push_unique(found: &mut Vec, p: PathBuf) { - if !found.contains(&p) { - found.push(p); - } -} - fn encoded_dir(sessions: &Path, cwd: &Path) -> Option { let p = sessions.join(encode_cwd(cwd)?); p.is_dir().then_some(p) @@ -171,15 +161,10 @@ fn unique_session(groups: &[PathBuf], spawned: SystemTime) -> Option { let Some(name) = name.to_str() else { continue; }; - if !candidates.iter().any(|c| c == name) { - candidates.push(name.to_string()); - } + push_unique(&mut candidates, name.to_string()); } } - match candidates.as_slice() { - [only] if is_uuid(only) => Some(only.clone()), - _ => None, - } + sole_id(candidates) } /// `YYYY-MM-DDTHH:MM:SS[.frac]Z` as grok writes `created_at`. Any other shape diff --git a/src/harness/mod.rs b/src/harness/mod.rs index 7141aea..9c67167 100644 --- a/src/harness/mod.rs +++ b/src/harness/mod.rs @@ -24,7 +24,7 @@ pub mod summary; use std::{ ffi::OsString, fs::File, - io::Read, + io::{BufRead, BufReader, Read}, path::{Path, PathBuf}, time::{Duration, SystemTime}, }; @@ -94,8 +94,11 @@ pub trait Harness: Sync { home: Option<&Path>, ) -> SpawnPlan; - /// Extract a session ID from hook or notify JSON. - fn parse_capture(&self, payload: &str) -> Option; + /// Extract a session ID from hook or notify JSON. Defaults to `None` for + /// tools without an injected capture channel. + fn parse_capture(&self, _payload: &str) -> Option { + None + } /// Extract a session ID from final terminal text, including scrollback. fn scrape_exit(&self, text: &str) -> Option; @@ -277,6 +280,13 @@ pub fn is_uuid(s: &str) -> bool { }) } +/// Validated session ID at `key` in hook or notify JSON; [`is_uuid`] is the +/// shell-insertion boundary. +fn capture_id(v: &jzon::JsonValue, key: &str) -> Option { + let id = v[key].as_str()?; + is_uuid(id).then(|| id.to_string()) +} + /// Return the strict UUID at the start of `s`. The next byte must end the token; /// an alphanumeric character, `-`, or `_` extends the token and rejects it. fn leading_uuid(s: &str) -> Option<&str> { @@ -362,6 +372,55 @@ fn v7_millis(id: &str) -> Option { u64::from_str_radix(&format!("{}{}", &id[..8], &id[9..13]), 16).ok() } +/// Epoch milliseconds of `t`; `None` before the epoch. +fn unix_millis(t: SystemTime) -> Option { + Some(t.duration_since(SystemTime::UNIX_EPOCH).ok()?.as_millis()) +} + +/// Whether a store-recorded path names the task's working directory: literal +/// equality first, so identical nonexistent paths stay eligible, then the +/// canonical task path `canon` for symlinked invocations. +fn same_cwd(recorded: &Path, cwd: &Path, canon: Option<&Path>) -> bool { + recorded == cwd || canon.is_some_and(|c| recorded == c) +} + +/// The one candidate when exactly one strict UUID survives; any other count +/// or shape yields `None`. +fn sole_id(candidates: Vec) -> Option { + match candidates.as_slice() { + [only] if is_uuid(only) => Some(only.clone()), + _ => None, + } +} + +/// Append `x` unless an equal entry is present. +fn push_unique(v: &mut Vec, x: T) { + if !v.contains(&x) { + v.push(x); + } +} + +/// Parse the first `n` JSONL records of `path`, reading at most 64 KiB so +/// later transcript content cannot affect correlation. Each entry is `None` +/// when its line does not parse, so callers decide whether a malformed record +/// rejects or is skipped. `None` when the file cannot be opened, a read +/// fails, or the file ends before `n` records. +fn jsonl_head(path: &Path, n: usize) -> Option>> { + let file = File::open(path).ok()?; + let mut reader = BufReader::new(file.take(64 * 1024)); + let mut records = Vec::with_capacity(n); + let mut line = String::new(); + for _ in 0..n { + line.clear(); + match reader.read_line(&mut line) { + Ok(len) if len > 0 => {} + _ => return None, + } + records.push(jzon::parse(&line).ok()); + } + Some(records) +} + /// Single-quote `s` for `$SHELL -c`, encoding embedded `'` as `'\''`. fn shell_quote(s: &str) -> String { format!("'{}'", s.replace('\'', "'\\''")) diff --git a/src/harness/omp.rs b/src/harness/omp.rs index 0892f1e..d70f813 100644 --- a/src/harness/omp.rs +++ b/src/harness/omp.rs @@ -25,14 +25,14 @@ use std::{ fs, - io::{BufRead, BufReader, Read}, path::{Path, PathBuf}, time::SystemTime, }; use super::{ - CAPTURE_ENV, CapturePaths, Harness, Invocation, SpawnPlan, is_uuid, leading_uuid, shell_quote, - v7_millis, within_window_ms, + CAPTURE_ENV, CapturePaths, Harness, Invocation, SpawnPlan, capture_id, is_uuid, jsonl_head, + leading_uuid, push_unique, same_cwd, shell_quote, sole_id, unix_millis, v7_millis, + within_window_ms, }; /// Command fragment shared by ordinary exit and recovery hints. @@ -134,9 +134,7 @@ impl Harness for Omp { /// Return `sessionId` from a valid extension payload. fn parse_capture(&self, payload: &str) -> Option { - let v = jzon::parse(payload).ok()?; - let id = v["sessionId"].as_str()?; - is_uuid(id).then(|| id.to_string()) + capture_id(&jzon::parse(payload).ok()?, "sessionId") } /// Return the last trusted exit hint. Unlabelled hints are ordinary exit @@ -167,10 +165,7 @@ impl Harness for Omp { /// stores. The ID follows the last `_` in the filename. fn correlate_fs(&self, cwd: &Path, spawned: SystemTime, home: Option<&Path>) -> Option { let sessions = self.home_root(home)?; - let spawn_ms = spawned - .duration_since(SystemTime::UNIX_EPOCH) - .ok()? - .as_millis(); + let spawn_ms = unix_millis(spawned)?; // Session headers may record either the supplied or canonical path. let canon = cwd.canonicalize().ok(); @@ -210,15 +205,10 @@ impl Harness for Omp { } // The same session may appear in multiple buckets; count its // UUID once. - if !survivors.iter().any(|s| s == id) { - survivors.push(id.to_string()); - } + push_unique(&mut survivors, id.to_string()); } } - match survivors.as_slice() { - [only] => Some(only.clone()), - _ => None, - } + sole_id(survivors) } } @@ -227,26 +217,16 @@ impl Harness for Omp { /// first record is a fixed-width title slot. Nothing later can affect /// correlation, so transcripts are not read beyond the header. fn header_cwd_matches(path: &Path, cwd: &Path, canon: Option<&Path>) -> bool { - let Ok(file) = fs::File::open(path) else { + let Some(records) = jsonl_head(path, 2) else { return false; }; - let mut reader = BufReader::new(file.take(64 * 1024)); - let mut line = String::new(); - for _ in 0..2 { - line.clear(); - if !matches!(reader.read_line(&mut line), Ok(n) if n > 0) { - return false; - } - let Ok(record) = jzon::parse(&line) else { - continue; - }; + for record in records.into_iter().flatten() { if record["type"].as_str() != Some("session") { continue; } return record["cwd"].as_str().is_some_and(|c| { let header = Path::new(c); - header == cwd - || Some(header) == canon + same_cwd(header, cwd, canon) || canon.is_some_and(|canon| header.canonicalize().is_ok_and(|h| h == canon)) }); }