Improve shared harness logic

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<String> {
-        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<String> {
@@ -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<String> {
         }
         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,

src/harness/codex.rs

@@ -4,17 +4,12 @@
 //! correlates rollout files under `<codex-home>/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<String> {
@@ -108,10 +102,7 @@ impl Harness for Codex {
 
     fn correlate_fs(&self, cwd: &Path, spawned: SystemTime, home: Option<&Path>) -> Option<String> {
         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)]

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<String> {
-        None
-    }
-
     fn scrape_exit(&self, text: &str) -> Option<String> {
         // 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<PathBuf>, p: PathBuf) {
-    if !found.contains(&p) {
-        found.push(p);
-    }
-}
-
 fn encoded_dir(sessions: &Path, cwd: &Path) -> Option<PathBuf> {
     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<String> {
             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

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<String>;
+    /// 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<String> {
+        None
+    }
 
     /// Extract a session ID from final terminal text, including scrollback.
     fn scrape_exit(&self, text: &str) -> Option<String>;
@@ -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<String> {
+    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> {
     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<u128> {
+    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<String>) -> Option<String> {
+    match candidates.as_slice() {
+        [only] if is_uuid(only) => Some(only.clone()),
+        _ => None,
+    }
+}
+
+/// Append `x` unless an equal entry is present.
+fn push_unique<T: PartialEq>(v: &mut Vec<T>, 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<Vec<Option<jzon::JsonValue>>> {
+    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('\'', "'\\''"))

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<String> {
-        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<String> {
         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))
         });
     }