From c5e52ab7647e71bbd7b41cafce05e8d2c1079129 Mon Sep 17 00:00:00 2001 From: Lennart Kotzur Date: Sun, 19 Jul 2026 11:33:16 +0200 Subject: [PATCH 1/2] fix(omp): soft-stop nested task teardown without killing parent identity OMP task subagents share the parent process and were binding/stopping the parent on session_shutdown. Soft-finalize on owner shutdown, skip nested bind via a process-wide identity latch, recover bindings from HCOM_INSTANCE_NAME, and split the omp hook module under 1k lines. --- src/hooks/omp.rs | 1024 ------------------------------------- src/hooks/omp/handlers.rs | 365 +++++++++++++ src/hooks/omp/mod.rs | 19 + src/hooks/omp/plugin.rs | 148 ++++++ src/hooks/omp/tests.rs | 670 ++++++++++++++++++++++++ src/instance_binding.rs | 38 ++ src/omp_plugin/hcom.ts | 62 ++- 7 files changed, 1295 insertions(+), 1031 deletions(-) delete mode 100644 src/hooks/omp.rs create mode 100644 src/hooks/omp/handlers.rs create mode 100644 src/hooks/omp/mod.rs create mode 100644 src/hooks/omp/plugin.rs create mode 100644 src/hooks/omp/tests.rs diff --git a/src/hooks/omp.rs b/src/hooks/omp.rs deleted file mode 100644 index fd1fb311..00000000 --- a/src/hooks/omp.rs +++ /dev/null @@ -1,1024 +0,0 @@ -//! Omp Coding Agent hook handlers — argv-based lifecycle plus TypeScript plugin. - -use std::time::Instant; - -use serde_json::Value; - -use crate::bootstrap; -use crate::db::HcomDb; -use crate::instance_binding; -use crate::instance_lifecycle as lifecycle; -use crate::instances; -use crate::log::{log_error, log_info}; -use crate::shared::ST_LISTENING; -use crate::shared::context::HcomContext; - -use super::common; -use super::common::finalize_session; - -fn parse_flag(argv: &[String], flag: &str) -> Option { - argv.iter() - .position(|a| a == flag) - .and_then(|i| argv.get(i + 1)) - .cloned() -} - -fn has_flag(argv: &[String], flag: &str) -> bool { - argv.iter().any(|a| a == flag) -} - -fn upsert_plugin_notify_endpoint(db: &HcomDb, instance_name: &str, port: u16) { - if let Err(e) = db.upsert_notify_endpoint(instance_name, "plugin", port) { - log_error( - "native", - "omp.register_notify_fail", - &format!( - "Failed to register plugin notify port for {}: {}", - instance_name, e - ), - ); - return; - } - - crate::notify::wake(db, instance_name, crate::notify::WakeKind::DELIVERY_LOOPS); -} - -fn initialize_last_event_id(db: &HcomDb, instance_name: &str) { - if let Ok(Some(existing)) = db.get_instance_full(instance_name) - && existing.last_event_id == 0 - { - let launch_event_id: Option = std::env::var("HCOM_LAUNCH_EVENT_ID") - .ok() - .and_then(|s| s.parse().ok()); - let current_max = db.get_last_event_id(); - let new_id = match launch_event_id { - Some(lei) if lei <= current_max => lei, - _ => current_max, - }; - let mut updates = serde_json::Map::new(); - updates.insert("last_event_id".into(), serde_json::json!(new_id)); - instances::update_instance_position(db, instance_name, &updates); - } -} - -fn bootstrap_for(ctx: &HcomContext, db: &HcomDb, instance_name: &str) -> String { - let tag = db - .get_instance_full(instance_name) - .ok() - .flatten() - .and_then(|d| d.tag.clone()) - .unwrap_or_default(); - let hcom_config = crate::config::HcomConfig::load(None).unwrap_or_default(); - let relay_enabled = crate::relay::is_relay_enabled(&hcom_config); - let effective_tag = if tag.is_empty() { - &hcom_config.tag - } else { - &tag - }; - bootstrap::get_bootstrap( - db, - &ctx.hcom_dir, - instance_name, - "omp", - ctx.is_background, - ctx.is_launched, - &ctx.notes, - effective_tag, - relay_enabled, - ctx.background_name.as_deref(), - ) -} - -fn handle_start(ctx: &HcomContext, db: &HcomDb, argv: &[String]) -> (i32, String) { - // Plugin RPC returns JSON errors on exit 0 so the extension can handle - // setup failures without Pi treating the hook itself as failed. - let session_id = match parse_flag(argv, "--session-id") { - Some(sid) => sid, - None => return (0, r#"{"error":"Missing --session-id"}"#.to_string()), - }; - let transcript_path = parse_flag(argv, "--transcript-path"); - let cwd = parse_flag(argv, "--cwd"); - let notify_port: Option = parse_flag(argv, "--notify-port").and_then(|s| s.parse().ok()); - - let process_id = match &ctx.process_id { - Some(pid) => pid.clone(), - None => return (0, r#"{"error":"HCOM_PROCESS_ID not set"}"#.to_string()), - }; - - let instance_name = - match instance_binding::bind_session_to_process(db, &session_id, Some(&process_id)) { - Some(name) => name, - None => { - return ( - 0, - r#"{"error":"No instance bound to this process"}"#.to_string(), - ); - } - }; - - initialize_last_event_id(db, &instance_name); - lifecycle::set_status( - db, - &instance_name, - ST_LISTENING, - "start", - Default::default(), - ); - instance_binding::capture_and_store_launch_context(db, &instance_name); - - let mut updates = serde_json::Map::new(); - updates.insert("tool".into(), serde_json::json!("omp")); - updates.insert("session_id".into(), serde_json::json!(&session_id)); - if let Some(path) = transcript_path.as_ref().filter(|p| !p.is_empty()) { - updates.insert("transcript_path".into(), serde_json::json!(path)); - } - let cwd_value = cwd - .as_deref() - .filter(|p| !p.is_empty()) - .or_else(|| ctx.cwd.to_str()); - if let Some(cwd) = cwd_value { - updates.insert("directory".into(), serde_json::json!(cwd)); - } - instances::update_instance_position(db, &instance_name, &updates); - if let Some(port) = notify_port { - upsert_plugin_notify_endpoint(db, &instance_name, port); - } - log_info( - "hooks", - "omp-start.bind", - &format!("instance={} session_id={}", instance_name, session_id), - ); - crate::relay::worker::ensure_worker(true); - - let response = serde_json::json!({ - "name": instance_name, - "session_id": session_id, - "bootstrap": bootstrap_for(ctx, db, &instance_name), - }); - (0, response.to_string()) -} - -fn handle_status(db: &HcomDb, argv: &[String]) -> (i32, String) { - let name = match parse_flag(argv, "--name") { - Some(n) => n, - None => return (0, r#"{"error":"Missing --name or --status"}"#.to_string()), - }; - let status = match parse_flag(argv, "--status") { - Some(s) => s, - None => return (0, r#"{"error":"Missing --name or --status"}"#.to_string()), - }; - let context = parse_flag(argv, "--context").unwrap_or_default(); - let detail = parse_flag(argv, "--detail").unwrap_or_default(); - let was_listening = db - .get_instance_full(&name) - .ok() - .flatten() - .is_some_and(|inst| inst.status == ST_LISTENING); - - lifecycle::set_status( - db, - &name, - &status, - &context, - lifecycle::StatusUpdate { - detail: &detail, - ..Default::default() - }, - ); - if status == ST_LISTENING && !was_listening { - crate::notify::wake(db, &name, &[]); - } - (0, r#"{"ok":true}"#.to_string()) -} - -fn handle_read(db: &HcomDb, argv: &[String]) -> (i32, String) { - let name = match parse_flag(argv, "--name") { - Some(n) => n, - None => return (0, r#"{"error":"Missing --name"}"#.to_string()), - }; - let format_mode = has_flag(argv, "--format"); - let check_mode = has_flag(argv, "--check"); - let ack_mode = has_flag(argv, "--ack"); - - let raw_messages = db.get_unread_messages(&name); - let messages: Vec = raw_messages.iter().map(common::message_to_value).collect(); - - if format_mode { - if messages.is_empty() { - return (0, String::new()); - } - let deliver = common::limit_delivery_messages(&messages); - return ( - 0, - common::format_messages_json_for_instance(db, &deliver, &name), - ); - } - if ack_mode { - if let Some(up_to) = parse_flag(argv, "--up-to") { - let Ok(ack_id) = up_to.parse::() else { - return ( - 0, - serde_json::json!({"error": format!("Invalid --up-to: {}", up_to)}).to_string(), - ); - }; - let mut updates = serde_json::Map::new(); - updates.insert("last_event_id".into(), serde_json::json!(ack_id)); - instances::update_instance_position(db, &name, &updates); - return (0, serde_json::json!({"acked_to": ack_id}).to_string()); - } - if messages.is_empty() { - return (0, r#"{"acked":0}"#.to_string()); - } - let ack_id = messages - .iter() - .filter_map(|m| m.get("event_id").and_then(|v| v.as_i64())) - .max() - .filter(|id| *id > 0) - .unwrap_or_else(|| db.get_last_event_id()); - if ack_id > 0 { - let mut updates = serde_json::Map::new(); - updates.insert("last_event_id".into(), serde_json::json!(ack_id)); - instances::update_instance_position(db, &name, &updates); - } - return (0, serde_json::json!({"acked": messages.len()}).to_string()); - } - if check_mode { - return ( - 0, - if messages.is_empty() { "false" } else { "true" }.to_string(), - ); - } - ( - 0, - serde_json::to_string(&messages).unwrap_or_else(|_| "[]".to_string()), - ) -} - -fn handle_beforetool(db: &HcomDb, argv: &[String]) -> (i32, String) { - let name = match parse_flag(argv, "--name") { - Some(n) => n, - None => return (0, r#"{"decision":"allow"}"#.to_string()), - }; - let tool_name = parse_flag(argv, "--tool").unwrap_or_default(); - let input = parse_flag(argv, "--input-json") - .and_then(|raw| serde_json::from_str::(&raw).ok()) - .unwrap_or_else(|| serde_json::json!({})); - if !tool_name.is_empty() { - common::update_tool_status(db, &name, "omp", &tool_name, &input); - } - (0, r#"{"decision":"allow"}"#.to_string()) -} - -fn handle_stop(db: &HcomDb, argv: &[String]) -> (i32, String) { - let name = match parse_flag(argv, "--name") { - Some(n) => n, - None => return (0, r#"{"error":"Missing --name"}"#.to_string()), - }; - let reason = parse_flag(argv, "--reason").unwrap_or_else(|| "unknown".to_string()); - finalize_session(db, &name, &reason, None); - (0, r#"{"ok":true}"#.to_string()) -} - -pub fn dispatch_omp_hook(hook_name: &str, argv: &[String]) -> (i32, String) { - let start = Instant::now(); - let ctx = HcomContext::from_os(); - crate::paths::ensure_hcom_directories_at(&ctx.hcom_dir); - let db = match HcomDb::open() { - Ok(db) => db, - Err(e) => { - log_error( - "hooks", - "hook.error", - &format!("hook={} op=db_open err={}", hook_name, e), - ); - return ( - 0, - serde_json::json!({"error": format!("DB open failed: {}", e)}).to_string(), - ); - } - }; - if !common::hook_gate_check(&ctx, &db) { - return (0, String::new()); - } - let handler_argv: Vec = if !argv.is_empty() && argv[0] == hook_name { - argv[1..].to_vec() - } else { - argv.to_vec() - }; - let hook_name_owned = hook_name.to_string(); - let handler_start = Instant::now(); - let (exit_code, output) = common::dispatch_with_panic_guard( - "omp", - &hook_name_owned, - ( - 0, - serde_json::json!({"error": "internal panic"}).to_string(), - ), - || match hook_name_owned.as_str() { - "omp-start" => handle_start(&ctx, &db, &handler_argv), - "omp-status" => handle_status(&db, &handler_argv), - "omp-read" => handle_read(&db, &handler_argv), - "omp-beforetool" => handle_beforetool(&db, &handler_argv), - "omp-stop" => handle_stop(&db, &handler_argv), - _ => ( - 0, - serde_json::json!({"error": format!("Unknown Omp hook: {}", hook_name_owned)}) - .to_string(), - ), - }, - ); - log_info( - "hooks", - "omp.dispatch.timing", - &format!( - "hook={} handler_ms={:.2} total_ms={:.2} exit_code={}", - hook_name, - handler_start.elapsed().as_secs_f64() * 1000.0, - start.elapsed().as_secs_f64() * 1000.0, - exit_code - ), - ); - (exit_code, output) -} - -pub const PLUGIN_SOURCE: &str = include_str!("../omp_plugin/hcom.ts"); -const PLUGIN_FILENAME: &str = "hcom.ts"; - -fn current_home_dir() -> std::path::PathBuf { - std::env::var("HOME") - .map(std::path::PathBuf::from) - .unwrap_or_else(|_| dirs::home_dir().unwrap_or_default()) -} - -fn omp_plugin_dir() -> std::path::PathBuf { - let tool_root = crate::runtime_env::tool_config_root(); - let home = current_home_dir(); - if tool_root == home { - if let Ok(dir) = std::env::var("PI_CODING_AGENT_DIR") - && !dir.is_empty() - { - return std::path::PathBuf::from(dir).join("extensions"); - } - home.join(".omp").join("agent").join("extensions") - } else { - tool_root.join(".omp").join("extensions") - } -} - -pub fn get_omp_plugin_path() -> std::path::PathBuf { - omp_plugin_dir().join(PLUGIN_FILENAME) -} - -pub fn extension_inject_args() -> Vec { - vec![ - "-e".to_string(), - get_omp_plugin_path().to_string_lossy().to_string(), - ] -} - -/// Remove hcom's managed OMP extension injection (`-e ` / -/// `--extension …`, incl. the `=` forms) from a stored or replayed launch-arg -/// vector, preserving every user-supplied extension and its ordering. An entry -/// is treated as managed when its path is the current plugin path, an existing -/// hcom-owned file, or — for a moved/missing managed file — a narrow lexical -/// match (basename `hcom.ts` directly under an `extensions` directory). -/// -/// Idempotent. Callers strip stored args before snapshotting and reinjecting so -/// a stale plugin path from an older hcom/config layout is not replayed -/// alongside the freshly injected current path (which could fail startup or load -/// hcom twice). A genuine `-e other.ts` user extension always survives. -pub fn strip_managed_extension_args(args: &mut Vec) { - let current = get_omp_plugin_path(); - let is_managed = |value: &str| -> bool { - let path = std::path::Path::new(value); - if path == current.as_path() { - return true; - } - if is_hcom_owned(path) { - return true; - } - // Moved/missing managed file only: basename hcom.ts under an `extensions` - // dir. Gated on !exists so an EXISTING user `-e …/extensions/hcom.ts` - // with unrelated contents (which is_hcom_owned already rejected) is kept - // — only exact-current and hcom-owned files are removed when present. - !path.exists() - && path.file_name().and_then(|n| n.to_str()) == Some(PLUGIN_FILENAME) - && path - .parent() - .and_then(|p| p.file_name()) - .and_then(|n| n.to_str()) - == Some("extensions") - }; - let mut out: Vec = Vec::with_capacity(args.len()); - let mut i = 0; - while i < args.len() { - let tok = args[i].as_str(); - // Two-token forms: `-e PATH` / `--extension PATH`. - if (tok == "-e" || tok == "--extension") && i + 1 < args.len() { - if is_managed(&args[i + 1]) { - i += 2; - continue; - } - out.push(args[i].clone()); - out.push(args[i + 1].clone()); - i += 2; - continue; - } - // Equals forms: `--extension=PATH` / `-e=PATH`. - if let Some(value) = tok - .strip_prefix("--extension=") - .or_else(|| tok.strip_prefix("-e=")) - && is_managed(value) - { - i += 1; - continue; - } - out.push(args[i].clone()); - i += 1; - } - *args = out; -} - -fn plugin_matches_source(path: &std::path::Path) -> bool { - match std::fs::read_to_string(path) { - Ok(content) => content == PLUGIN_SOURCE, - Err(_) => false, - } -} - -pub fn verify_omp_plugin_installed() -> bool { - plugin_matches_source(&get_omp_plugin_path()) -} - -fn is_hcom_owned(path: &std::path::Path) -> bool { - std::fs::read_to_string(path) - .map(|content| content.contains("customType: \"hcom-bootstrap\"")) - .unwrap_or(false) -} - -pub fn install_omp_plugin() -> std::io::Result { - let target_dir = omp_plugin_dir(); - let target = target_dir.join(PLUGIN_FILENAME); - std::fs::create_dir_all(&target_dir)?; - if target.is_symlink() || target.exists() { - if !plugin_matches_source(&target) && !is_hcom_owned(&target) { - return Err(std::io::Error::new( - std::io::ErrorKind::AlreadyExists, - "A non-hcom hcom.ts file already exists and will not be overwritten", - )); - } - std::fs::remove_file(&target)?; - } - std::fs::write(&target, PLUGIN_SOURCE)?; - Ok(true) -} - -pub fn ensure_omp_plugin_installed() -> bool { - if verify_omp_plugin_installed() { - return true; - } - install_omp_plugin().unwrap_or(false) -} - -pub fn remove_omp_plugin() -> std::io::Result<()> { - let path = get_omp_plugin_path(); - if (path.exists() || path.is_symlink()) - && (plugin_matches_source(&path) || is_hcom_owned(&path)) - { - std::fs::remove_file(&path)?; - } - Ok(()) -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::shared::{ST_ACTIVE, ST_LISTENING}; - use std::io::ErrorKind; - use std::net::TcpListener; - use std::path::PathBuf; - use std::time::{Duration, Instant}; - - fn setup_test_db() -> (HcomDb, PathBuf) { - use std::sync::atomic::{AtomicU64, Ordering}; - static COUNTER: AtomicU64 = AtomicU64::new(0); - - let temp_dir = std::env::temp_dir(); - let test_id = COUNTER.fetch_add(1, Ordering::Relaxed); - let db_path = temp_dir.join(format!( - "test_omp_hooks_{}_{}.db", - std::process::id(), - test_id - )); - - let db = HcomDb::open_at(&db_path).unwrap(); - (db, db_path) - } - - fn cleanup(path: PathBuf) { - let _ = std::fs::remove_file(&path); - let _ = std::fs::remove_file(path.with_extension("db-wal")); - let _ = std::fs::remove_file(path.with_extension("db-shm")); - } - - fn save_test_instance(db: &HcomDb, name: &str, status: &str) { - let mut row = serde_json::Map::new(); - row.insert("name".into(), serde_json::json!(name)); - row.insert("tool".into(), serde_json::json!("omp")); - row.insert("status".into(), serde_json::json!(status)); - row.insert("status_context".into(), serde_json::json!("")); - row.insert("status_detail".into(), serde_json::json!("")); - row.insert("created_at".into(), serde_json::json!(1.0)); - db.save_instance_named(name, &row).unwrap(); - } - - fn bind_probe() -> TcpListener { - let listener = TcpListener::bind("127.0.0.1:0").unwrap(); - listener.set_nonblocking(true).unwrap(); - listener - } - - fn await_connect(listener: &TcpListener, timeout: Duration) -> bool { - let deadline = Instant::now() + timeout; - loop { - match listener.accept() { - Ok(_) => return true, - Err(e) if e.kind() == ErrorKind::WouldBlock => { - if Instant::now() >= deadline { - return false; - } - std::thread::sleep(Duration::from_millis(5)); - } - Err(_) => return false, - } - } - } - - #[test] - fn strip_managed_extension_removes_only_hcom_entry() { - let current = get_omp_plugin_path().to_string_lossy().to_string(); - - // Current managed path (two-token) is removed; user extension survives. - let mut args = vec![ - "--model".into(), - "opus".into(), - "-e".into(), - current.clone(), - "-e".into(), - "/home/u/mine.ts".into(), - ]; - strip_managed_extension_args(&mut args); - assert_eq!( - args, - vec!["--model", "opus", "-e", "/home/u/mine.ts"] - .into_iter() - .map(String::from) - .collect::>() - ); - - // Idempotent. - let once = args.clone(); - strip_managed_extension_args(&mut args); - assert_eq!(args, once); - - // Legacy/moved managed path (missing file) via lexical fallback: - // basename hcom.ts under an `extensions` dir, incl. the `=` form. - let mut legacy = vec![ - "--extension=/old/place/.omp/agent/extensions/hcom.ts".into(), - "--extension".into(), - "/old/place/extensions/hcom.ts".into(), - "--extension=/home/u/other.ts".into(), - ]; - strip_managed_extension_args(&mut legacy); - assert_eq!(legacy, vec!["--extension=/home/u/other.ts".to_string()]); - - // A user extension merely named hcom.ts but NOT under `extensions/` is - // preserved (narrow matcher). - let mut keep = vec!["-e".into(), "/home/u/project/hcom.ts".into()]; - strip_managed_extension_args(&mut keep); - assert_eq!( - keep, - vec!["-e".to_string(), "/home/u/project/hcom.ts".to_string()] - ); - - // An EXISTING non-hcom file at extensions/hcom.ts must survive: the - // lexical fallback applies only to moved/missing files. - let dir = tempfile::tempdir().unwrap(); - let ext_dir = dir.path().join("extensions"); - std::fs::create_dir_all(&ext_dir).unwrap(); - let user_file = ext_dir.join("hcom.ts"); - std::fs::write(&user_file, "export default () => {}; // not hcom").unwrap(); - let user_arg = user_file.to_string_lossy().to_string(); - let mut existing = vec!["-e".into(), user_arg.clone()]; - strip_managed_extension_args(&mut existing); - assert_eq!(existing, vec!["-e".to_string(), user_arg]); - } - - #[test] - fn plugin_bootstraps_via_hidden_message() { - assert!(PLUGIN_SOURCE.contains("before_agent_start")); - assert!(PLUGIN_SOURCE.contains("customType: \"hcom-bootstrap\"")); - assert!(PLUGIN_SOURCE.contains("display: false")); - assert!(!PLUGIN_SOURCE.contains("text: `${bootstrapText}\\n\\n${event.text}`")); - } - - #[test] - fn plugin_reconcile_does_not_report_active_polling_status() { - assert!(!PLUGIN_SOURCE.contains( - "reportStatus(currentCtx, currentCtx.isIdle() ? \"listening\" : \"active\")" - )); - assert!(PLUGIN_SOURCE.contains("pi.on(\"agent_end\"")); - assert!(PLUGIN_SOURCE.contains("IDLE_DEBOUNCE_MS")); - assert!(PLUGIN_SOURCE.contains("currentCtx?.isIdle()")); - assert!(!PLUGIN_SOURCE.contains("pi.on(\"turn_end\", async (_event, ctx) => {\n\t\tcurrentCtx = ctx;\n\t\tawait reportStatus(ctx, \"listening\");")); - } - - #[test] - fn plugin_delivery_reports_active_edge() { - assert!(PLUGIN_SOURCE.contains("reportStatus(ctx, \"active\"")); - assert!(PLUGIN_SOURCE.contains("`deliver:${sender}`")); - } - - // The embedded plugin is include_str!'d and never tsc'd, so these guard the - // delivery-correctness invariants that upstream API/lifecycle drift silently - // broke before (see PR review). They pin behavior, not just strings. - - #[test] - fn plugin_acks_transform_submission_in_before_agent_start() { - // omp applies the bodyless-wake transform inline (no source:"extension" - // re-emit), so the transform-path ack must happen in before_agent_start. - // Without it pendingAckId stays set and deliverPending jams forever. - let idx = PLUGIN_SOURCE - .find("pi.on(\"before_agent_start\"") - .expect("before_agent_start handler present"); - assert!( - PLUGIN_SOURCE[idx..].contains("ackPending(\"before_agent_start\")"), - "before_agent_start must ack the inline transform submission" - ); - assert!(PLUGIN_SOURCE.contains("if (pendingAckId !== null) await ackPending")); - } - - #[test] - fn plugin_replays_wakes_dropped_during_in_flight_window() { - // In-flight/pending-ack wakes must be queued and replayed, not dropped. - assert!(PLUGIN_SOURCE.contains("deliveryPending")); - assert!(PLUGIN_SOURCE.contains("schedulePendingDelivery")); - assert!(PLUGIN_SOURCE.contains("drainPendingDelivery")); - // ackPending drains so the transform-path ack replays queued wakes. - let idx = PLUGIN_SOURCE - .find("async function ackPending") - .expect("ackPending present"); - assert!(PLUGIN_SOURCE[idx..].contains("drainPendingDelivery(\"post_ack_wake\")")); - } - - #[test] - fn plugin_keeps_ack_gate_until_command_succeeds() { - let idx = PLUGIN_SOURCE - .find("async function ackPending") - .expect("ackPending present"); - let ack = &PLUGIN_SOURCE[idx..]; - let command = ack.find("await hcom([\"omp-read\"").expect("ack command"); - let clear = ack.find("pendingAckId = null").expect("pending ack clear"); - assert!( - command < clear, - "pendingAckId must remain set while the ack command is in flight" - ); - assert!(ack.contains("if (result.code !== 0)")); - assert!(ack.contains("plugin.delivery_ack_failed")); - assert!(PLUGIN_SOURCE.contains("ackInFlight")); - assert!(PLUGIN_SOURCE.contains("await ackPending(\"reconcile\")")); - } - - #[test] - fn plugin_rebinds_identity_on_session_branch() { - // /branch mints a new session id and emits only session_branch (not - // session_switch); the plugin must reset+rebind or delivery dies. - let idx = PLUGIN_SOURCE - .find("pi.on(\"session_branch\"") - .expect("session_branch handler present"); - let handler = &PLUGIN_SOURCE[idx..]; - assert!(handler.contains("resetBinding()")); - assert!(handler.contains("bindIdentity(ctx)")); - } - - #[test] - fn start_handler_registering_plugin_notify_wakes_pty_delivery_loop() { - let (db, path) = setup_test_db(); - let temp = tempfile::TempDir::new().unwrap(); - save_test_instance(&db, "luna", ST_ACTIVE); - db.set_process_binding("pid-omp", "", "luna").unwrap(); - - let pty_listener = bind_probe(); - let pty_port = pty_listener.local_addr().unwrap().port(); - db.upsert_notify_endpoint("luna", "pty", pty_port).unwrap(); - - let plugin_listener = bind_probe(); - let plugin_port = plugin_listener.local_addr().unwrap().port(); - - let env = std::collections::HashMap::from([ - ("HCOM_PROCESS_ID".to_string(), "pid-omp".to_string()), - ("HCOM_LAUNCHED".to_string(), "1".to_string()), - ("HCOM_TOOL".to_string(), "omp".to_string()), - ]); - let ctx = HcomContext::from_env(&env, temp.path().to_path_buf()); - - let (code, output) = handle_start( - &ctx, - &db, - &[ - "--session-id".to_string(), - "sid-omp".to_string(), - "--notify-port".to_string(), - plugin_port.to_string(), - "--cwd".to_string(), - temp.path().to_string_lossy().to_string(), - ], - ); - - assert_eq!(code, 0); - let response: serde_json::Value = serde_json::from_str(&output).unwrap(); - assert_eq!(response.get("name").and_then(|v| v.as_str()), Some("luna")); - assert!(db.has_notify_endpoint_kind("luna", "plugin")); - - let stored_plugin_port: i64 = db - .conn() - .query_row( - "SELECT port FROM notify_endpoints WHERE instance = 'luna' AND kind = 'plugin'", - [], - |row| row.get(0), - ) - .unwrap(); - assert_eq!(stored_plugin_port, i64::from(plugin_port)); - assert!( - await_connect(&pty_listener, Duration::from_millis(500)), - "successful plugin bind must wake the PTY delivery loop so launch readiness is observed promptly" - ); - - drop(plugin_listener); - cleanup(path); - } - - #[test] - fn plugin_notify_registration_failure_does_not_wake_delivery_loop() { - let (db, path) = setup_test_db(); - save_test_instance(&db, "luna", ST_ACTIVE); - - let pty_listener = bind_probe(); - let pty_port = pty_listener.local_addr().unwrap().port(); - db.upsert_notify_endpoint("luna", "pty", pty_port).unwrap(); - - db.conn() - .execute_batch( - "CREATE TRIGGER fail_plugin_notify_insert - BEFORE INSERT ON notify_endpoints - WHEN NEW.kind = 'plugin' - BEGIN - SELECT RAISE(ABORT, 'plugin registration blocked'); - END;", - ) - .unwrap(); - - let plugin_listener = bind_probe(); - let plugin_port = plugin_listener.local_addr().unwrap().port(); - upsert_plugin_notify_endpoint(&db, "luna", plugin_port); - - assert!(!db.has_notify_endpoint_kind("luna", "plugin")); - assert!( - !await_connect(&pty_listener, Duration::from_millis(100)), - "failed plugin bind must not wake delivery loops" - ); - - drop(plugin_listener); - cleanup(path); - } - - #[test] - fn status_handler_wakes_plugin_only_when_entering_listening() { - let (db, path) = setup_test_db(); - save_test_instance(&db, "luna", ST_LISTENING); - - let listener = TcpListener::bind("127.0.0.1:0").unwrap(); - listener.set_nonblocking(true).unwrap(); - let port = listener.local_addr().unwrap().port(); - db.upsert_notify_endpoint("luna", "plugin", port).unwrap(); - - let argv = vec![ - "--name".to_string(), - "luna".to_string(), - "--status".to_string(), - ST_LISTENING.to_string(), - ]; - let (code, _) = handle_status(&db, &argv); - assert_eq!(code, 0); - std::thread::sleep(Duration::from_millis(20)); - assert!(listener.accept().is_err()); - - let mut updates = serde_json::Map::new(); - updates.insert("status".into(), serde_json::json!(ST_ACTIVE)); - instances::update_instance_position(&db, "luna", &updates); - - let (code, _) = handle_status(&db, &argv); - assert_eq!(code, 0); - let mut accepted = false; - for _ in 0..10 { - if listener.accept().is_ok() { - accepted = true; - break; - } - std::thread::sleep(Duration::from_millis(10)); - } - assert!(accepted); - - cleanup(path); - } - - #[test] - fn start_handler_uses_central_binding_for_existing_session() { - let (db, path) = setup_test_db(); - let temp = tempfile::TempDir::new().unwrap(); - - let mut canonical = serde_json::Map::new(); - canonical.insert("name".into(), serde_json::json!("miso")); - canonical.insert("tool".into(), serde_json::json!("omp")); - canonical.insert("session_id".into(), serde_json::json!("sid-123")); - canonical.insert("status".into(), serde_json::json!(ST_LISTENING)); - canonical.insert("status_context".into(), serde_json::json!("")); - canonical.insert("status_detail".into(), serde_json::json!("")); - canonical.insert("last_event_id".into(), serde_json::json!(42)); - canonical.insert("created_at".into(), serde_json::json!(1.0)); - db.save_instance_named("miso", &canonical).unwrap(); - db.rebind_session("sid-123", "miso").unwrap(); - - let mut placeholder = serde_json::Map::new(); - placeholder.insert("name".into(), serde_json::json!("temp")); - placeholder.insert("tool".into(), serde_json::json!("omp")); - placeholder.insert("status".into(), serde_json::json!("pending")); - placeholder.insert("status_context".into(), serde_json::json!("new")); - placeholder.insert("status_detail".into(), serde_json::json!("")); - placeholder.insert("created_at".into(), serde_json::json!(1.0)); - db.save_instance_named("temp", &placeholder).unwrap(); - db.set_process_binding("pid-123", "", "temp").unwrap(); - - let env = std::collections::HashMap::from([ - ("HCOM_PROCESS_ID".to_string(), "pid-123".to_string()), - ("HCOM_LAUNCHED".to_string(), "1".to_string()), - ("HCOM_TOOL".to_string(), "omp".to_string()), - ]); - let ctx = HcomContext::from_env(&env, temp.path().to_path_buf()); - - let (code, output) = handle_start( - &ctx, - &db, - &[ - "--session-id".to_string(), - "sid-123".to_string(), - "--cwd".to_string(), - temp.path().to_string_lossy().to_string(), - ], - ); - assert_eq!(code, 0); - let response: serde_json::Value = serde_json::from_str(&output).unwrap(); - assert_eq!(response.get("name").and_then(|v| v.as_str()), Some("miso")); - assert!(db.get_instance_full("temp").unwrap().is_none()); - assert_eq!( - db.get_process_binding("pid-123").unwrap(), - Some("miso".to_string()) - ); - - let rebound = db.get_instance_full("miso").unwrap().unwrap(); - assert_eq!(rebound.last_event_id, 42); - assert_eq!(rebound.directory, temp.path().to_string_lossy()); - - cleanup(path); - } - - // ── Plugin install/remove safety ────────────────────────────────── - - /// Helper: run a closure with a temp HOME + HCOM_DIR (via isolated_test_env), - /// Runs a test with isolated HCOM_DIR and HOME, Config reset, - /// and PI_CODING_AGENT_DIR explicitly unset so the default ~/.omp path is used. - fn with_isolated_omp_env(f: impl FnOnce(&std::path::Path)) { - let (_dir, _hcom, home, _guard) = crate::hooks::test_helpers::isolated_test_env(); - unsafe { - std::env::remove_var("PI_CODING_AGENT_DIR"); - } - f(&home); - } - #[test] - #[serial_test::serial] - fn plugin_dir_respects_pi_coding_agent_dir() { - let (_dir, _hcom, home, _guard) = crate::hooks::test_helpers::isolated_test_env(); - let custom = home.join("custom-omp"); - unsafe { - std::env::set_var("PI_CODING_AGENT_DIR", &custom); - } - - let path = get_omp_plugin_path(); - assert_eq!(path, custom.join("extensions").join("hcom.ts")); - } - - #[test] - fn extension_inject_args_contains_absolute_plugin_path() { - with_isolated_omp_env(|_| { - let args = extension_inject_args(); - assert_eq!(args.len(), 2); - assert_eq!(args[0], "-e"); - let path = std::path::Path::new(&args[1]); - assert!(path.is_absolute()); - assert_eq!(path.file_name().and_then(|n| n.to_str()), Some("hcom.ts")); - }); - } - - #[test] - fn plugin_source_uses_omp_cli_commands_only() { - assert!(PLUGIN_SOURCE.contains("[\"omp-read\"")); - assert!(PLUGIN_SOURCE.contains("[\"omp-status\"")); - assert!(PLUGIN_SOURCE.contains("[\"omp-stop\"")); - assert!(!PLUGIN_SOURCE.contains("[\"pi-read\"")); - assert!(!PLUGIN_SOURCE.contains("[\"pi-status\"")); - assert!(!PLUGIN_SOURCE.contains("[\"pi-stop\"")); - } - - #[test] - fn plugin_source_matches_omp_input_result_shape() { - assert!(PLUGIN_SOURCE.contains("return {}")); - assert!(PLUGIN_SOURCE.contains("return { text:")); - assert!(PLUGIN_SOURCE.contains("return { handled: true }")); - assert!(!PLUGIN_SOURCE.contains("action: \"continue\"")); - assert!(!PLUGIN_SOURCE.contains("action: \"transform\"")); - assert!(!PLUGIN_SOURCE.contains("action: \"handled\"")); - assert!(!PLUGIN_SOURCE.contains("streamingBehavior")); - } - - #[test] - fn plugin_source_handles_omp_session_switch_and_shutdown_shape() { - assert!(PLUGIN_SOURCE.contains("pi.on(\"session_switch\"")); - assert!(PLUGIN_SOURCE.contains("\"--reason\", \"shutdown\"")); - assert!(!PLUGIN_SOURCE.contains("event.reason")); - } - - #[test] - fn install_writes_plugin_source() { - with_isolated_omp_env(|_| { - assert!(install_omp_plugin().unwrap()); - let content = std::fs::read_to_string(get_omp_plugin_path()).unwrap(); - assert_eq!(content, PLUGIN_SOURCE); - assert!(verify_omp_plugin_installed()); - }); - } - - #[test] - fn install_refuses_to_overwrite_non_hcom_file() { - with_isolated_omp_env(|_| { - let path = get_omp_plugin_path(); - std::fs::create_dir_all(path.parent().unwrap()).unwrap(); - std::fs::write(&path, "// user's custom plugin").unwrap(); - - let result = install_omp_plugin(); - assert!(result.is_err()); - assert_eq!( - std::fs::read_to_string(&path).unwrap(), - "// user's custom plugin", - ); - assert!(!verify_omp_plugin_installed()); - }); - } - - #[test] - fn install_upgrades_stale_hcom_owned_plugin() { - with_isolated_omp_env(|_| { - let path = get_omp_plugin_path(); - std::fs::create_dir_all(path.parent().unwrap()).unwrap(); - // Old hcom plugin: has the ownership marker but doesn't match current source. - std::fs::write(&path, r#"const x = customType: "hcom-bootstrap";"#).unwrap(); - - assert!(install_omp_plugin().unwrap()); - assert_eq!(std::fs::read_to_string(&path).unwrap(), PLUGIN_SOURCE,); - assert!(verify_omp_plugin_installed()); - }); - } - - #[test] - fn remove_deletes_hcom_plugin() { - with_isolated_omp_env(|_| { - install_omp_plugin().unwrap(); - let path = get_omp_plugin_path(); - assert!(path.exists()); - - remove_omp_plugin().unwrap(); - assert!(!path.exists()); - }); - } - - #[test] - fn remove_preserves_non_hcom_file() { - with_isolated_omp_env(|_| { - let path = get_omp_plugin_path(); - std::fs::create_dir_all(path.parent().unwrap()).unwrap(); - std::fs::write(&path, "// user's custom plugin").unwrap(); - - remove_omp_plugin().unwrap(); - assert!(path.exists(), "non-hcom file must not be removed"); - }); - } -} diff --git a/src/hooks/omp/handlers.rs b/src/hooks/omp/handlers.rs new file mode 100644 index 00000000..8a3e40d6 --- /dev/null +++ b/src/hooks/omp/handlers.rs @@ -0,0 +1,365 @@ +//! Omp Coding Agent hook handlers — argv-based lifecycle plus TypeScript plugin. + +use std::time::Instant; + +use serde_json::Value; + +use crate::bootstrap; +use crate::db::HcomDb; +use crate::instance_binding; +use crate::instance_lifecycle as lifecycle; +use crate::instances; +use crate::log::{log_error, log_info}; +use crate::shared::ST_LISTENING; +use crate::shared::context::HcomContext; + +use crate::hooks::common; +use crate::hooks::common::finalize_session; + +fn parse_flag(argv: &[String], flag: &str) -> Option { + argv.iter() + .position(|a| a == flag) + .and_then(|i| argv.get(i + 1)) + .cloned() +} + +fn has_flag(argv: &[String], flag: &str) -> bool { + argv.iter().any(|a| a == flag) +} + +pub(crate) fn upsert_plugin_notify_endpoint(db: &HcomDb, instance_name: &str, port: u16) { + if let Err(e) = db.upsert_notify_endpoint(instance_name, "plugin", port) { + log_error( + "native", + "omp.register_notify_fail", + &format!( + "Failed to register plugin notify port for {}: {}", + instance_name, e + ), + ); + return; + } + + crate::notify::wake(db, instance_name, crate::notify::WakeKind::DELIVERY_LOOPS); +} + +fn initialize_last_event_id(db: &HcomDb, instance_name: &str) { + if let Ok(Some(existing)) = db.get_instance_full(instance_name) + && existing.last_event_id == 0 + { + let launch_event_id: Option = std::env::var("HCOM_LAUNCH_EVENT_ID") + .ok() + .and_then(|s| s.parse().ok()); + let current_max = db.get_last_event_id(); + let new_id = match launch_event_id { + Some(lei) if lei <= current_max => lei, + _ => current_max, + }; + let mut updates = serde_json::Map::new(); + updates.insert("last_event_id".into(), serde_json::json!(new_id)); + instances::update_instance_position(db, instance_name, &updates); + } +} + +fn instance_name_from_env(ctx: &HcomContext) -> Option { + ctx.raw_env + .get("HCOM_INSTANCE_NAME") + .or_else(|| ctx.raw_env.get("HCOM_NAME")) + .filter(|s| !s.is_empty()) + .cloned() +} + +fn bootstrap_for(ctx: &HcomContext, db: &HcomDb, instance_name: &str) -> String { + let tag = db + .get_instance_full(instance_name) + .ok() + .flatten() + .and_then(|d| d.tag.clone()) + .unwrap_or_default(); + let hcom_config = crate::config::HcomConfig::load(None).unwrap_or_default(); + let relay_enabled = crate::relay::is_relay_enabled(&hcom_config); + let effective_tag = if tag.is_empty() { + &hcom_config.tag + } else { + &tag + }; + bootstrap::get_bootstrap( + db, + &ctx.hcom_dir, + instance_name, + "omp", + ctx.is_background, + ctx.is_launched, + &ctx.notes, + effective_tag, + relay_enabled, + ctx.background_name.as_deref(), + ) +} + +pub(crate) fn handle_start(ctx: &HcomContext, db: &HcomDb, argv: &[String]) -> (i32, String) { + // Plugin RPC returns JSON errors on exit 0 so the extension can handle + // setup failures without Pi treating the hook itself as failed. + let session_id = match parse_flag(argv, "--session-id") { + Some(sid) => sid, + None => return (0, r#"{"error":"Missing --session-id"}"#.to_string()), + }; + let transcript_path = parse_flag(argv, "--transcript-path"); + let cwd = parse_flag(argv, "--cwd"); + let notify_port: Option = parse_flag(argv, "--notify-port").and_then(|s| s.parse().ok()); + + let process_id = match &ctx.process_id { + Some(pid) => pid.clone(), + None => return (0, r#"{"error":"HCOM_PROCESS_ID not set"}"#.to_string()), + }; + + let instance_name = + match instance_binding::bind_session_to_process(db, &session_id, Some(&process_id)) { + Some(name) => name, + None => match instance_name_from_env(ctx).and_then(|name| { + instance_binding::recover_process_binding_for_instance( + db, + &name, + &session_id, + &process_id, + ) + }) { + Some(name) => name, + None => { + return ( + 0, + r#"{"error":"No instance bound to this process"}"#.to_string(), + ); + } + }, + }; + + initialize_last_event_id(db, &instance_name); + lifecycle::set_status( + db, + &instance_name, + ST_LISTENING, + "start", + Default::default(), + ); + instance_binding::capture_and_store_launch_context(db, &instance_name); + + let mut updates = serde_json::Map::new(); + updates.insert("tool".into(), serde_json::json!("omp")); + updates.insert("session_id".into(), serde_json::json!(&session_id)); + if let Some(path) = transcript_path.as_ref().filter(|p| !p.is_empty()) { + updates.insert("transcript_path".into(), serde_json::json!(path)); + } + let cwd_value = cwd + .as_deref() + .filter(|p| !p.is_empty()) + .or_else(|| ctx.cwd.to_str()); + if let Some(cwd) = cwd_value { + updates.insert("directory".into(), serde_json::json!(cwd)); + } + instances::update_instance_position(db, &instance_name, &updates); + if let Some(port) = notify_port { + upsert_plugin_notify_endpoint(db, &instance_name, port); + } + log_info( + "hooks", + "omp-start.bind", + &format!("instance={} session_id={}", instance_name, session_id), + ); + crate::relay::worker::ensure_worker(true); + + let response = serde_json::json!({ + "name": instance_name, + "session_id": session_id, + "bootstrap": bootstrap_for(ctx, db, &instance_name), + }); + (0, response.to_string()) +} + +pub(crate) fn handle_status(db: &HcomDb, argv: &[String]) -> (i32, String) { + let name = match parse_flag(argv, "--name") { + Some(n) => n, + None => return (0, r#"{"error":"Missing --name or --status"}"#.to_string()), + }; + let status = match parse_flag(argv, "--status") { + Some(s) => s, + None => return (0, r#"{"error":"Missing --name or --status"}"#.to_string()), + }; + let context = parse_flag(argv, "--context").unwrap_or_default(); + let detail = parse_flag(argv, "--detail").unwrap_or_default(); + let was_listening = db + .get_instance_full(&name) + .ok() + .flatten() + .is_some_and(|inst| inst.status == ST_LISTENING); + + lifecycle::set_status( + db, + &name, + &status, + &context, + lifecycle::StatusUpdate { + detail: &detail, + ..Default::default() + }, + ); + if status == ST_LISTENING && !was_listening { + crate::notify::wake(db, &name, &[]); + } + (0, r#"{"ok":true}"#.to_string()) +} + +fn handle_read(db: &HcomDb, argv: &[String]) -> (i32, String) { + let name = match parse_flag(argv, "--name") { + Some(n) => n, + None => return (0, r#"{"error":"Missing --name"}"#.to_string()), + }; + let format_mode = has_flag(argv, "--format"); + let check_mode = has_flag(argv, "--check"); + let ack_mode = has_flag(argv, "--ack"); + + let raw_messages = db.get_unread_messages(&name); + let messages: Vec = raw_messages.iter().map(common::message_to_value).collect(); + + if format_mode { + if messages.is_empty() { + return (0, String::new()); + } + let deliver = common::limit_delivery_messages(&messages); + return ( + 0, + common::format_messages_json_for_instance(db, &deliver, &name), + ); + } + if ack_mode { + if let Some(up_to) = parse_flag(argv, "--up-to") { + let Ok(ack_id) = up_to.parse::() else { + return ( + 0, + serde_json::json!({"error": format!("Invalid --up-to: {}", up_to)}).to_string(), + ); + }; + let mut updates = serde_json::Map::new(); + updates.insert("last_event_id".into(), serde_json::json!(ack_id)); + instances::update_instance_position(db, &name, &updates); + return (0, serde_json::json!({"acked_to": ack_id}).to_string()); + } + if messages.is_empty() { + return (0, r#"{"acked":0}"#.to_string()); + } + let ack_id = messages + .iter() + .filter_map(|m| m.get("event_id").and_then(|v| v.as_i64())) + .max() + .filter(|id| *id > 0) + .unwrap_or_else(|| db.get_last_event_id()); + if ack_id > 0 { + let mut updates = serde_json::Map::new(); + updates.insert("last_event_id".into(), serde_json::json!(ack_id)); + instances::update_instance_position(db, &name, &updates); + } + return (0, serde_json::json!({"acked": messages.len()}).to_string()); + } + if check_mode { + return ( + 0, + if messages.is_empty() { "false" } else { "true" }.to_string(), + ); + } + ( + 0, + serde_json::to_string(&messages).unwrap_or_else(|_| "[]".to_string()), + ) +} + +fn handle_beforetool(db: &HcomDb, argv: &[String]) -> (i32, String) { + let name = match parse_flag(argv, "--name") { + Some(n) => n, + None => return (0, r#"{"decision":"allow"}"#.to_string()), + }; + let tool_name = parse_flag(argv, "--tool").unwrap_or_default(); + let input = parse_flag(argv, "--input-json") + .and_then(|raw| serde_json::from_str::(&raw).ok()) + .unwrap_or_else(|| serde_json::json!({})); + if !tool_name.is_empty() { + common::update_tool_status(db, &name, "omp", &tool_name, &input); + } + (0, r#"{"decision":"allow"}"#.to_string()) +} + +pub(crate) fn handle_stop(db: &HcomDb, argv: &[String]) -> (i32, String) { + let name = match parse_flag(argv, "--name") { + Some(n) => n, + None => return (0, r#"{"error":"Missing --name"}"#.to_string()), + }; + let reason = parse_flag(argv, "--reason").unwrap_or_else(|| "unknown".to_string()); + if has_flag(argv, "--soft") { + common::soft_finalize_session(db, &name, &reason, None); + (0, r#"{"ok":true,"soft":true}"#.to_string()) + } else { + finalize_session(db, &name, &reason, None); + (0, r#"{"ok":true}"#.to_string()) + } +} + +pub fn dispatch_omp_hook(hook_name: &str, argv: &[String]) -> (i32, String) { + let start = Instant::now(); + let ctx = HcomContext::from_os(); + crate::paths::ensure_hcom_directories_at(&ctx.hcom_dir); + let db = match HcomDb::open() { + Ok(db) => db, + Err(e) => { + log_error( + "hooks", + "hook.error", + &format!("hook={} op=db_open err={}", hook_name, e), + ); + return ( + 0, + serde_json::json!({"error": format!("DB open failed: {}", e)}).to_string(), + ); + } + }; + if !common::hook_gate_check(&ctx, &db) { + return (0, String::new()); + } + let handler_argv: Vec = if !argv.is_empty() && argv[0] == hook_name { + argv[1..].to_vec() + } else { + argv.to_vec() + }; + let hook_name_owned = hook_name.to_string(); + let handler_start = Instant::now(); + let (exit_code, output) = common::dispatch_with_panic_guard( + "omp", + &hook_name_owned, + ( + 0, + serde_json::json!({"error": "internal panic"}).to_string(), + ), + || match hook_name_owned.as_str() { + "omp-start" => handle_start(&ctx, &db, &handler_argv), + "omp-status" => handle_status(&db, &handler_argv), + "omp-read" => handle_read(&db, &handler_argv), + "omp-beforetool" => handle_beforetool(&db, &handler_argv), + "omp-stop" => handle_stop(&db, &handler_argv), + _ => ( + 0, + serde_json::json!({"error": format!("Unknown Omp hook: {}", hook_name_owned)}) + .to_string(), + ), + }, + ); + log_info( + "hooks", + "omp.dispatch.timing", + &format!( + "hook={} handler_ms={:.2} total_ms={:.2} exit_code={}", + hook_name, + handler_start.elapsed().as_secs_f64() * 1000.0, + start.elapsed().as_secs_f64() * 1000.0, + exit_code + ), + ); + (exit_code, output) +} diff --git a/src/hooks/omp/mod.rs b/src/hooks/omp/mod.rs new file mode 100644 index 00000000..5084af05 --- /dev/null +++ b/src/hooks/omp/mod.rs @@ -0,0 +1,19 @@ +//! Omp Coding Agent hook handlers — argv-based lifecycle plus TypeScript plugin. + +mod handlers; +mod plugin; + +#[cfg(test)] +mod tests; + +pub use handlers::dispatch_omp_hook; +pub use plugin::{ + PLUGIN_SOURCE, ensure_omp_plugin_installed, extension_inject_args, get_omp_plugin_path, + install_omp_plugin, remove_omp_plugin, strip_managed_extension_args, + verify_omp_plugin_installed, +}; + +#[cfg(test)] +pub(crate) use handlers::{ + handle_start, handle_status, handle_stop, upsert_plugin_notify_endpoint, +}; diff --git a/src/hooks/omp/plugin.rs b/src/hooks/omp/plugin.rs new file mode 100644 index 00000000..75886630 --- /dev/null +++ b/src/hooks/omp/plugin.rs @@ -0,0 +1,148 @@ +pub const PLUGIN_SOURCE: &str = include_str!("../../omp_plugin/hcom.ts"); +const PLUGIN_FILENAME: &str = "hcom.ts"; + +fn current_home_dir() -> std::path::PathBuf { + std::env::var("HOME") + .map(std::path::PathBuf::from) + .unwrap_or_else(|_| dirs::home_dir().unwrap_or_default()) +} + +fn omp_plugin_dir() -> std::path::PathBuf { + let tool_root = crate::runtime_env::tool_config_root(); + let home = current_home_dir(); + if tool_root == home { + if let Ok(dir) = std::env::var("PI_CODING_AGENT_DIR") + && !dir.is_empty() + { + return std::path::PathBuf::from(dir).join("extensions"); + } + home.join(".omp").join("agent").join("extensions") + } else { + tool_root.join(".omp").join("extensions") + } +} + +pub fn get_omp_plugin_path() -> std::path::PathBuf { + omp_plugin_dir().join(PLUGIN_FILENAME) +} + +pub fn extension_inject_args() -> Vec { + vec![ + "-e".to_string(), + get_omp_plugin_path().to_string_lossy().to_string(), + ] +} + +/// Remove hcom's managed OMP extension injection (`-e ` / +/// `--extension …`, incl. the `=` forms) from a stored or replayed launch-arg +/// vector, preserving every user-supplied extension and its ordering. An entry +/// is treated as managed when its path is the current plugin path, an existing +/// hcom-owned file, or — for a moved/missing managed file — a narrow lexical +/// match (basename `hcom.ts` directly under an `extensions` directory). +/// +/// Idempotent. Callers strip stored args before snapshotting and reinjecting so +/// a stale plugin path from an older hcom/config layout is not replayed +/// alongside the freshly injected current path (which could fail startup or load +/// hcom twice). A genuine `-e other.ts` user extension always survives. +pub fn strip_managed_extension_args(args: &mut Vec) { + let current = get_omp_plugin_path(); + let is_managed = |value: &str| -> bool { + let path = std::path::Path::new(value); + if path == current.as_path() { + return true; + } + if is_hcom_owned(path) { + return true; + } + // Moved/missing managed file only: basename hcom.ts under an `extensions` + // dir. Gated on !exists so an EXISTING user `-e …/extensions/hcom.ts` + // with unrelated contents (which is_hcom_owned already rejected) is kept + // — only exact-current and hcom-owned files are removed when present. + !path.exists() + && path.file_name().and_then(|n| n.to_str()) == Some(PLUGIN_FILENAME) + && path + .parent() + .and_then(|p| p.file_name()) + .and_then(|n| n.to_str()) + == Some("extensions") + }; + let mut out: Vec = Vec::with_capacity(args.len()); + let mut i = 0; + while i < args.len() { + let tok = args[i].as_str(); + // Two-token forms: `-e PATH` / `--extension PATH`. + if (tok == "-e" || tok == "--extension") && i + 1 < args.len() { + if is_managed(&args[i + 1]) { + i += 2; + continue; + } + out.push(args[i].clone()); + out.push(args[i + 1].clone()); + i += 2; + continue; + } + // Equals forms: `--extension=PATH` / `-e=PATH`. + if let Some(value) = tok + .strip_prefix("--extension=") + .or_else(|| tok.strip_prefix("-e=")) + && is_managed(value) + { + i += 1; + continue; + } + out.push(args[i].clone()); + i += 1; + } + *args = out; +} + +fn plugin_matches_source(path: &std::path::Path) -> bool { + match std::fs::read_to_string(path) { + Ok(content) => content == PLUGIN_SOURCE, + Err(_) => false, + } +} + +pub fn verify_omp_plugin_installed() -> bool { + plugin_matches_source(&get_omp_plugin_path()) +} + +fn is_hcom_owned(path: &std::path::Path) -> bool { + std::fs::read_to_string(path) + .map(|content| content.contains("customType: \"hcom-bootstrap\"")) + .unwrap_or(false) +} + +pub fn install_omp_plugin() -> std::io::Result { + let target_dir = omp_plugin_dir(); + let target = target_dir.join(PLUGIN_FILENAME); + std::fs::create_dir_all(&target_dir)?; + if target.is_symlink() || target.exists() { + if !plugin_matches_source(&target) && !is_hcom_owned(&target) { + return Err(std::io::Error::new( + std::io::ErrorKind::AlreadyExists, + "A non-hcom hcom.ts file already exists and will not be overwritten", + )); + } + std::fs::remove_file(&target)?; + } + std::fs::write(&target, PLUGIN_SOURCE)?; + Ok(true) +} + +pub fn ensure_omp_plugin_installed() -> bool { + if verify_omp_plugin_installed() { + return true; + } + install_omp_plugin().unwrap_or(false) +} + +pub fn remove_omp_plugin() -> std::io::Result<()> { + let path = get_omp_plugin_path(); + if (path.exists() || path.is_symlink()) + && (plugin_matches_source(&path) || is_hcom_owned(&path)) + { + std::fs::remove_file(&path)?; + } + Ok(()) +} diff --git a/src/hooks/omp/tests.rs b/src/hooks/omp/tests.rs new file mode 100644 index 00000000..ebb8ea5d --- /dev/null +++ b/src/hooks/omp/tests.rs @@ -0,0 +1,670 @@ +use super::*; +use crate::db::HcomDb; +use crate::instances; +use crate::shared::context::HcomContext; +use crate::shared::{ST_ACTIVE, ST_LISTENING}; +use std::io::ErrorKind; +use std::net::TcpListener; +use std::path::PathBuf; +use std::time::{Duration, Instant}; + +fn setup_test_db() -> (HcomDb, PathBuf) { + use std::sync::atomic::{AtomicU64, Ordering}; + static COUNTER: AtomicU64 = AtomicU64::new(0); + + let temp_dir = std::env::temp_dir(); + let test_id = COUNTER.fetch_add(1, Ordering::Relaxed); + let db_path = temp_dir.join(format!( + "test_omp_hooks_{}_{}.db", + std::process::id(), + test_id + )); + + let db = HcomDb::open_at(&db_path).unwrap(); + (db, db_path) +} + +fn cleanup(path: PathBuf) { + let _ = std::fs::remove_file(&path); + let _ = std::fs::remove_file(path.with_extension("db-wal")); + let _ = std::fs::remove_file(path.with_extension("db-shm")); +} + +fn save_test_instance(db: &HcomDb, name: &str, status: &str) { + let mut row = serde_json::Map::new(); + row.insert("name".into(), serde_json::json!(name)); + row.insert("tool".into(), serde_json::json!("omp")); + row.insert("status".into(), serde_json::json!(status)); + row.insert("status_context".into(), serde_json::json!("")); + row.insert("status_detail".into(), serde_json::json!("")); + row.insert("created_at".into(), serde_json::json!(1.0)); + db.save_instance_named(name, &row).unwrap(); +} + +fn bind_probe() -> TcpListener { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + listener.set_nonblocking(true).unwrap(); + listener +} + +fn await_connect(listener: &TcpListener, timeout: Duration) -> bool { + let deadline = Instant::now() + timeout; + loop { + match listener.accept() { + Ok(_) => return true, + Err(e) if e.kind() == ErrorKind::WouldBlock => { + if Instant::now() >= deadline { + return false; + } + std::thread::sleep(Duration::from_millis(5)); + } + Err(_) => return false, + } + } +} + +#[test] +#[serial_test::serial] +fn strip_managed_extension_removes_only_hcom_entry() { + with_isolated_omp_env(|_| { + let current = get_omp_plugin_path().to_string_lossy().to_string(); + + // Current managed path (two-token) is removed; user extension survives. + let mut args = vec![ + "--model".into(), + "opus".into(), + "-e".into(), + current.clone(), + "-e".into(), + "/home/u/mine.ts".into(), + ]; + strip_managed_extension_args(&mut args); + assert_eq!( + args, + vec!["--model", "opus", "-e", "/home/u/mine.ts"] + .into_iter() + .map(String::from) + .collect::>() + ); + + // Idempotent. + let once = args.clone(); + strip_managed_extension_args(&mut args); + assert_eq!(args, once); + + // Legacy/moved managed path (missing file) via lexical fallback: + // basename hcom.ts under an `extensions` dir, incl. the `=` form. + let mut legacy = vec![ + "--extension=/old/place/.omp/agent/extensions/hcom.ts".into(), + "--extension".into(), + "/old/place/extensions/hcom.ts".into(), + "--extension=/home/u/other.ts".into(), + ]; + strip_managed_extension_args(&mut legacy); + assert_eq!(legacy, vec!["--extension=/home/u/other.ts".to_string()]); + + // A user extension merely named hcom.ts but NOT under `extensions/` is + // preserved (narrow matcher). + let mut keep = vec!["-e".into(), "/home/u/project/hcom.ts".into()]; + strip_managed_extension_args(&mut keep); + assert_eq!( + keep, + vec!["-e".to_string(), "/home/u/project/hcom.ts".to_string()] + ); + + // An EXISTING non-hcom file at extensions/hcom.ts must survive: the + // lexical fallback applies only to moved/missing files. + let dir = tempfile::tempdir().unwrap(); + let ext_dir = dir.path().join("extensions"); + std::fs::create_dir_all(&ext_dir).unwrap(); + let user_file = ext_dir.join("hcom.ts"); + std::fs::write(&user_file, "export default () => {}; // not hcom").unwrap(); + let user_arg = user_file.to_string_lossy().to_string(); + let mut existing = vec!["-e".into(), user_arg.clone()]; + strip_managed_extension_args(&mut existing); + assert_eq!(existing, vec!["-e".to_string(), user_arg]); + }); +} + +#[test] +fn plugin_bootstraps_via_hidden_message() { + assert!(PLUGIN_SOURCE.contains("before_agent_start")); + assert!(PLUGIN_SOURCE.contains("customType: \"hcom-bootstrap\"")); + assert!(PLUGIN_SOURCE.contains("display: false")); + assert!(!PLUGIN_SOURCE.contains("text: `${bootstrapText}\\n\\n${event.text}`")); +} + +#[test] +fn plugin_reconcile_does_not_report_active_polling_status() { + assert!( + !PLUGIN_SOURCE + .contains("reportStatus(currentCtx, currentCtx.isIdle() ? \"listening\" : \"active\")") + ); + assert!(PLUGIN_SOURCE.contains("pi.on(\"agent_end\"")); + assert!(PLUGIN_SOURCE.contains("IDLE_DEBOUNCE_MS")); + assert!(PLUGIN_SOURCE.contains("currentCtx?.isIdle()")); + assert!(!PLUGIN_SOURCE.contains("pi.on(\"turn_end\", async (_event, ctx) => {\n\t\tcurrentCtx = ctx;\n\t\tawait reportStatus(ctx, \"listening\");")); +} + +#[test] +fn plugin_delivery_reports_active_edge() { + assert!(PLUGIN_SOURCE.contains("reportStatus(ctx, \"active\"")); + assert!(PLUGIN_SOURCE.contains("`deliver:${sender}`")); +} + +// The embedded plugin is include_str!'d and never tsc'd, so these guard the +// delivery-correctness invariants that upstream API/lifecycle drift silently +// broke before (see PR review). They pin behavior, not just strings. + +#[test] +fn plugin_acks_transform_submission_in_before_agent_start() { + // omp applies the bodyless-wake transform inline (no source:"extension" + // re-emit), so the transform-path ack must happen in before_agent_start. + // Without it pendingAckId stays set and deliverPending jams forever. + let idx = PLUGIN_SOURCE + .find("pi.on(\"before_agent_start\"") + .expect("before_agent_start handler present"); + assert!( + PLUGIN_SOURCE[idx..].contains("ackPending(\"before_agent_start\")"), + "before_agent_start must ack the inline transform submission" + ); + assert!(PLUGIN_SOURCE.contains("if (pendingAckId !== null) await ackPending")); +} + +#[test] +fn plugin_replays_wakes_dropped_during_in_flight_window() { + // In-flight/pending-ack wakes must be queued and replayed, not dropped. + assert!(PLUGIN_SOURCE.contains("deliveryPending")); + assert!(PLUGIN_SOURCE.contains("schedulePendingDelivery")); + assert!(PLUGIN_SOURCE.contains("drainPendingDelivery")); + // ackPending drains so the transform-path ack replays queued wakes. + let idx = PLUGIN_SOURCE + .find("async function ackPending") + .expect("ackPending present"); + assert!(PLUGIN_SOURCE[idx..].contains("drainPendingDelivery(\"post_ack_wake\")")); +} + +#[test] +fn plugin_keeps_ack_gate_until_command_succeeds() { + let idx = PLUGIN_SOURCE + .find("async function ackPending") + .expect("ackPending present"); + let ack = &PLUGIN_SOURCE[idx..]; + let command = ack.find("await hcom([\"omp-read\"").expect("ack command"); + let clear = ack.find("pendingAckId = null").expect("pending ack clear"); + assert!( + command < clear, + "pendingAckId must remain set while the ack command is in flight" + ); + assert!(ack.contains("if (result.code !== 0)")); + assert!(ack.contains("plugin.delivery_ack_failed")); + assert!(PLUGIN_SOURCE.contains("ackInFlight")); + assert!(PLUGIN_SOURCE.contains("await ackPending(\"reconcile\")")); +} + +#[test] +fn plugin_rebinds_identity_on_session_branch() { + // /branch mints a new session id and emits only session_branch (not + // session_switch); the plugin must reset+rebind or delivery dies. + // keepOwner retains process-env identity ownership so nested task + // extensions still skip bind while this instance rebinds. + let idx = PLUGIN_SOURCE + .find("pi.on(\"session_branch\"") + .expect("session_branch handler present"); + let handler = &PLUGIN_SOURCE[idx..]; + assert!(handler.contains("resetBinding({ keepOwner: true })")); + assert!(handler.contains("bindIdentity(ctx)")); +} + +#[test] +fn start_handler_registering_plugin_notify_wakes_pty_delivery_loop() { + let (db, path) = setup_test_db(); + let temp = tempfile::TempDir::new().unwrap(); + save_test_instance(&db, "luna", ST_ACTIVE); + db.set_process_binding("pid-omp", "", "luna").unwrap(); + + let pty_listener = bind_probe(); + let pty_port = pty_listener.local_addr().unwrap().port(); + db.upsert_notify_endpoint("luna", "pty", pty_port).unwrap(); + + let plugin_listener = bind_probe(); + let plugin_port = plugin_listener.local_addr().unwrap().port(); + + let env = std::collections::HashMap::from([ + ("HCOM_PROCESS_ID".to_string(), "pid-omp".to_string()), + ("HCOM_LAUNCHED".to_string(), "1".to_string()), + ("HCOM_TOOL".to_string(), "omp".to_string()), + ]); + let ctx = HcomContext::from_env(&env, temp.path().to_path_buf()); + + let (code, output) = handle_start( + &ctx, + &db, + &[ + "--session-id".to_string(), + "sid-omp".to_string(), + "--notify-port".to_string(), + plugin_port.to_string(), + "--cwd".to_string(), + temp.path().to_string_lossy().to_string(), + ], + ); + + assert_eq!(code, 0); + let response: serde_json::Value = serde_json::from_str(&output).unwrap(); + assert_eq!(response.get("name").and_then(|v| v.as_str()), Some("luna")); + assert!(db.has_notify_endpoint_kind("luna", "plugin")); + + let stored_plugin_port: i64 = db + .conn() + .query_row( + "SELECT port FROM notify_endpoints WHERE instance = 'luna' AND kind = 'plugin'", + [], + |row| row.get(0), + ) + .unwrap(); + assert_eq!(stored_plugin_port, i64::from(plugin_port)); + assert!( + await_connect(&pty_listener, Duration::from_millis(500)), + "successful plugin bind must wake the PTY delivery loop so launch readiness is observed promptly" + ); + + drop(plugin_listener); + cleanup(path); +} + +#[test] +fn plugin_notify_registration_failure_does_not_wake_delivery_loop() { + let (db, path) = setup_test_db(); + save_test_instance(&db, "luna", ST_ACTIVE); + + let pty_listener = bind_probe(); + let pty_port = pty_listener.local_addr().unwrap().port(); + db.upsert_notify_endpoint("luna", "pty", pty_port).unwrap(); + + db.conn() + .execute_batch( + "CREATE TRIGGER fail_plugin_notify_insert + BEFORE INSERT ON notify_endpoints + WHEN NEW.kind = 'plugin' + BEGIN + SELECT RAISE(ABORT, 'plugin registration blocked'); + END;", + ) + .unwrap(); + + let plugin_listener = bind_probe(); + let plugin_port = plugin_listener.local_addr().unwrap().port(); + upsert_plugin_notify_endpoint(&db, "luna", plugin_port); + + assert!(!db.has_notify_endpoint_kind("luna", "plugin")); + assert!( + !await_connect(&pty_listener, Duration::from_millis(100)), + "failed plugin bind must not wake delivery loops" + ); + + drop(plugin_listener); + cleanup(path); +} + +#[test] +fn status_handler_wakes_plugin_only_when_entering_listening() { + let (db, path) = setup_test_db(); + save_test_instance(&db, "luna", ST_LISTENING); + + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + listener.set_nonblocking(true).unwrap(); + let port = listener.local_addr().unwrap().port(); + db.upsert_notify_endpoint("luna", "plugin", port).unwrap(); + + let argv = vec![ + "--name".to_string(), + "luna".to_string(), + "--status".to_string(), + ST_LISTENING.to_string(), + ]; + let (code, _) = handle_status(&db, &argv); + assert_eq!(code, 0); + std::thread::sleep(Duration::from_millis(20)); + assert!(listener.accept().is_err()); + + let mut updates = serde_json::Map::new(); + updates.insert("status".into(), serde_json::json!(ST_ACTIVE)); + instances::update_instance_position(&db, "luna", &updates); + + let (code, _) = handle_status(&db, &argv); + assert_eq!(code, 0); + let mut accepted = false; + for _ in 0..10 { + if listener.accept().is_ok() { + accepted = true; + break; + } + std::thread::sleep(Duration::from_millis(10)); + } + assert!(accepted); + + cleanup(path); +} + +#[test] +fn start_handler_uses_central_binding_for_existing_session() { + let (db, path) = setup_test_db(); + let temp = tempfile::TempDir::new().unwrap(); + + let mut canonical = serde_json::Map::new(); + canonical.insert("name".into(), serde_json::json!("miso")); + canonical.insert("tool".into(), serde_json::json!("omp")); + canonical.insert("session_id".into(), serde_json::json!("sid-123")); + canonical.insert("status".into(), serde_json::json!(ST_LISTENING)); + canonical.insert("status_context".into(), serde_json::json!("")); + canonical.insert("status_detail".into(), serde_json::json!("")); + canonical.insert("last_event_id".into(), serde_json::json!(42)); + canonical.insert("created_at".into(), serde_json::json!(1.0)); + db.save_instance_named("miso", &canonical).unwrap(); + db.rebind_session("sid-123", "miso").unwrap(); + + let mut placeholder = serde_json::Map::new(); + placeholder.insert("name".into(), serde_json::json!("temp")); + placeholder.insert("tool".into(), serde_json::json!("omp")); + placeholder.insert("status".into(), serde_json::json!("pending")); + placeholder.insert("status_context".into(), serde_json::json!("new")); + placeholder.insert("status_detail".into(), serde_json::json!("")); + placeholder.insert("created_at".into(), serde_json::json!(1.0)); + db.save_instance_named("temp", &placeholder).unwrap(); + db.set_process_binding("pid-123", "", "temp").unwrap(); + + let env = std::collections::HashMap::from([ + ("HCOM_PROCESS_ID".to_string(), "pid-123".to_string()), + ("HCOM_LAUNCHED".to_string(), "1".to_string()), + ("HCOM_TOOL".to_string(), "omp".to_string()), + ]); + let ctx = HcomContext::from_env(&env, temp.path().to_path_buf()); + + let (code, output) = handle_start( + &ctx, + &db, + &[ + "--session-id".to_string(), + "sid-123".to_string(), + "--cwd".to_string(), + temp.path().to_string_lossy().to_string(), + ], + ); + assert_eq!(code, 0); + let response: serde_json::Value = serde_json::from_str(&output).unwrap(); + assert_eq!(response.get("name").and_then(|v| v.as_str()), Some("miso")); + assert!(db.get_instance_full("temp").unwrap().is_none()); + assert_eq!( + db.get_process_binding("pid-123").unwrap(), + Some("miso".to_string()) + ); + + let rebound = db.get_instance_full("miso").unwrap().unwrap(); + assert_eq!(rebound.last_event_id, 42); + assert_eq!(rebound.directory, temp.path().to_string_lossy()); + + cleanup(path); +} + +#[test] +fn soft_stop_keeps_instance_row_clears_process_binding() { + let (db, path) = setup_test_db(); + let now = chrono::Utc::now().timestamp() as f64; + db.conn() + .execute( + "INSERT INTO instances (name, status, created_at, tool, session_id) + VALUES ('luna', 'listening', ?1, 'omp', 'sid-soft')", + rusqlite::params![now], + ) + .unwrap(); + db.set_process_binding("pid-soft", "sid-soft", "luna") + .unwrap(); + db.rebind_session("sid-soft", "luna").unwrap(); + + let (code, output) = handle_stop( + &db, + &[ + "--name".to_string(), + "luna".to_string(), + "--soft".to_string(), + "--reason".to_string(), + "turn_end".to_string(), + ], + ); + + assert_eq!(code, 0); + let response: serde_json::Value = serde_json::from_str(&output).unwrap(); + assert_eq!(response.get("ok").and_then(|v| v.as_bool()), Some(true)); + assert_eq!(response.get("soft").and_then(|v| v.as_bool()), Some(true)); + assert!(db.get_instance_full("luna").unwrap().is_some()); + assert_eq!( + db.get_status("luna").unwrap().map(|(s, _)| s), + Some(crate::shared::ST_INACTIVE.to_string()) + ); + assert_eq!(db.get_process_binding("pid-soft").unwrap(), None); + assert_eq!(db.get_session_binding("sid-soft").unwrap(), None); + + cleanup(path); +} + +#[test] +fn start_recovers_binding_after_soft_stop_via_instance_name_env() { + let (db, path) = setup_test_db(); + let temp = tempfile::TempDir::new().unwrap(); + save_test_instance(&db, "miso", ST_LISTENING); + db.conn() + .execute( + "UPDATE instances SET session_id = 'sid-old' WHERE name = 'miso'", + [], + ) + .unwrap(); + db.set_process_binding("pid-recover", "sid-old", "miso") + .unwrap(); + db.rebind_session("sid-old", "miso").unwrap(); + + let (stop_code, _) = handle_stop( + &db, + &[ + "--name".to_string(), + "miso".to_string(), + "--soft".to_string(), + ], + ); + assert_eq!(stop_code, 0); + assert_eq!(db.get_process_binding("pid-recover").unwrap(), None); + + let env = std::collections::HashMap::from([ + ("HCOM_PROCESS_ID".to_string(), "pid-recover".to_string()), + ("HCOM_INSTANCE_NAME".to_string(), "miso".to_string()), + ("HCOM_LAUNCHED".to_string(), "1".to_string()), + ("HCOM_TOOL".to_string(), "omp".to_string()), + ]); + let ctx = HcomContext::from_env(&env, temp.path().to_path_buf()); + + let (code, output) = handle_start( + &ctx, + &db, + &[ + "--session-id".to_string(), + "sid-new".to_string(), + "--cwd".to_string(), + temp.path().to_string_lossy().to_string(), + ], + ); + + assert_eq!(code, 0); + let response: serde_json::Value = serde_json::from_str(&output).unwrap(); + assert_eq!(response.get("name").and_then(|v| v.as_str()), Some("miso")); + assert_eq!( + db.get_process_binding("pid-recover").unwrap(), + Some("miso".to_string()) + ); + assert_eq!( + db.get_session_binding("sid-new").unwrap(), + Some("miso".to_string()) + ); + let rebound = db.get_instance_full("miso").unwrap().unwrap(); + assert_eq!(rebound.session_id.as_deref(), Some("sid-new")); + assert_eq!(rebound.status, ST_LISTENING); + + cleanup(path); +} + +// ── Plugin install/remove safety ────────────────────────────────── + +/// Helper: run a closure with a temp HOME + HCOM_DIR (via isolated_test_env), +/// Runs a test with isolated HCOM_DIR and HOME, Config reset, +/// and PI_CODING_AGENT_DIR explicitly unset so the default ~/.omp path is used. +fn with_isolated_omp_env(f: impl FnOnce(&std::path::Path)) { + let (_dir, _hcom, home, _guard) = crate::hooks::test_helpers::isolated_test_env(); + unsafe { + std::env::remove_var("PI_CODING_AGENT_DIR"); + } + f(&home); +} +#[test] +#[serial_test::serial] +fn plugin_dir_respects_pi_coding_agent_dir() { + let (_dir, _hcom, home, _guard) = crate::hooks::test_helpers::isolated_test_env(); + let custom = home.join("custom-omp"); + unsafe { + std::env::set_var("PI_CODING_AGENT_DIR", &custom); + } + + let path = get_omp_plugin_path(); + assert_eq!(path, custom.join("extensions").join("hcom.ts")); +} + +#[test] +fn extension_inject_args_contains_absolute_plugin_path() { + with_isolated_omp_env(|_| { + let args = extension_inject_args(); + assert_eq!(args.len(), 2); + assert_eq!(args[0], "-e"); + let path = std::path::Path::new(&args[1]); + assert!(path.is_absolute()); + assert_eq!(path.file_name().and_then(|n| n.to_str()), Some("hcom.ts")); + }); +} + +#[test] +fn plugin_source_uses_omp_cli_commands_only() { + assert!(PLUGIN_SOURCE.contains("[\"omp-read\"")); + assert!(PLUGIN_SOURCE.contains("[\"omp-status\"")); + assert!(PLUGIN_SOURCE.contains("[\"omp-stop\"")); + assert!(!PLUGIN_SOURCE.contains("[\"pi-read\"")); + assert!(!PLUGIN_SOURCE.contains("[\"pi-status\"")); + assert!(!PLUGIN_SOURCE.contains("[\"pi-stop\"")); +} + +#[test] +fn plugin_source_matches_omp_input_result_shape() { + assert!(PLUGIN_SOURCE.contains("return {}")); + assert!(PLUGIN_SOURCE.contains("return { text:")); + assert!(PLUGIN_SOURCE.contains("return { handled: true }")); + assert!(!PLUGIN_SOURCE.contains("action: \"continue\"")); + assert!(!PLUGIN_SOURCE.contains("action: \"transform\"")); + assert!(!PLUGIN_SOURCE.contains("action: \"handled\"")); + assert!(!PLUGIN_SOURCE.contains("streamingBehavior")); +} + +#[test] +fn plugin_source_handles_omp_session_switch_and_shutdown_shape() { + assert!(PLUGIN_SOURCE.contains("pi.on(\"session_switch\"")); + assert!(PLUGIN_SOURCE.contains("pi.on(\"session_shutdown\"")); + // Soft-finalize on bound shutdown; never hard-delete from this path. + assert!(PLUGIN_SOURCE.contains("\"--soft\"")); + assert!(PLUGIN_SOURCE.contains("plugin.session_shutdown_skipped")); + assert!(PLUGIN_SOURCE.contains("nested_session")); + assert!(PLUGIN_SOURCE.contains("HCOM_OMP_IDENTITY_OWNER")); + assert!(PLUGIN_SOURCE.contains("plugin.bind_skipped_nested")); + assert!(PLUGIN_SOURCE.contains("ownsIdentity")); + assert!(PLUGIN_SOURCE.contains("keepOwner: true")); + // session_switch must also keepOwner (same race as session_branch). + let switch_idx = PLUGIN_SOURCE + .find("pi.on(\"session_switch\"") + .expect("session_switch handler present"); + assert!( + PLUGIN_SOURCE[switch_idx..].contains("resetBinding({ keepOwner: true })"), + "session_switch must keepOwner across rebind" + ); + // Dead session-id gates removed (SessionShutdownEvent has no sessionId). + assert!(!PLUGIN_SOURCE.contains("foreign_session")); + assert!(!PLUGIN_SOURCE.contains("rootSessionId")); + assert!(!PLUGIN_SOURCE.contains("shutdownSessionIdFromEvent")); + assert!(!PLUGIN_SOURCE.contains("shutdownReasonFromEvent")); + // Hard stop without --soft must not be the session_shutdown path. + let idx = PLUGIN_SOURCE + .find("pi.on(\"session_shutdown\"") + .expect("session_shutdown handler present"); + let handler = &PLUGIN_SOURCE[idx..]; + let soft = handler.find("\"--soft\"").expect("soft stop in shutdown"); + let stop = handler.find("[\"omp-stop\"").expect("omp-stop in shutdown"); + assert!(stop < soft, "omp-stop invocation must include --soft"); +} + +#[test] +fn install_writes_plugin_source() { + with_isolated_omp_env(|_| { + assert!(install_omp_plugin().unwrap()); + let content = std::fs::read_to_string(get_omp_plugin_path()).unwrap(); + assert_eq!(content, PLUGIN_SOURCE); + assert!(verify_omp_plugin_installed()); + }); +} + +#[test] +fn install_refuses_to_overwrite_non_hcom_file() { + with_isolated_omp_env(|_| { + let path = get_omp_plugin_path(); + std::fs::create_dir_all(path.parent().unwrap()).unwrap(); + std::fs::write(&path, "// user's custom plugin").unwrap(); + + let result = install_omp_plugin(); + assert!(result.is_err()); + assert_eq!( + std::fs::read_to_string(&path).unwrap(), + "// user's custom plugin", + ); + assert!(!verify_omp_plugin_installed()); + }); +} + +#[test] +fn install_upgrades_stale_hcom_owned_plugin() { + with_isolated_omp_env(|_| { + let path = get_omp_plugin_path(); + std::fs::create_dir_all(path.parent().unwrap()).unwrap(); + // Old hcom plugin: has the ownership marker but doesn't match current source. + std::fs::write(&path, r#"const x = customType: "hcom-bootstrap";"#).unwrap(); + + assert!(install_omp_plugin().unwrap()); + assert_eq!(std::fs::read_to_string(&path).unwrap(), PLUGIN_SOURCE,); + assert!(verify_omp_plugin_installed()); + }); +} + +#[test] +fn remove_deletes_hcom_plugin() { + with_isolated_omp_env(|_| { + install_omp_plugin().unwrap(); + let path = get_omp_plugin_path(); + assert!(path.exists()); + + remove_omp_plugin().unwrap(); + assert!(!path.exists()); + }); +} + +#[test] +fn remove_preserves_non_hcom_file() { + with_isolated_omp_env(|_| { + let path = get_omp_plugin_path(); + std::fs::create_dir_all(path.parent().unwrap()).unwrap(); + std::fs::write(&path, "// user's custom plugin").unwrap(); + + remove_omp_plugin().unwrap(); + assert!(path.exists(), "non-hcom file must not be removed"); + }); +} diff --git a/src/instance_binding.rs b/src/instance_binding.rs index 77f442e8..33884a50 100644 --- a/src/instance_binding.rs +++ b/src/instance_binding.rs @@ -632,6 +632,44 @@ pub fn bind_session_to_process( None } +/// Rebind process/session after soft-finalize cleared bindings but left the +/// instance row (typically inactive). Used when `bind_session_to_process` finds +/// no process binding, but the caller still knows the instance name (e.g. via +/// `HCOM_INSTANCE_NAME` / `HCOM_NAME` in a live OMP process). +pub fn recover_process_binding_for_instance( + db: &HcomDb, + instance_name: &str, + session_id: &str, + process_id: &str, +) -> Option { + if instance_name.is_empty() || session_id.is_empty() || process_id.is_empty() { + return None; + } + db.get_instance_full(instance_name).ok().flatten()?; + + if let Err(e) = db.clear_session_id_from_other_instances(session_id, instance_name) { + crate::log::log_error("binding", "recover.clear_session", &format!("{e}")); + } + let mut updates = serde_json::Map::new(); + updates.insert("session_id".into(), serde_json::json!(session_id)); + update_instance_position(db, instance_name, &updates); + if let Err(e) = db.rebind_session(session_id, instance_name) { + crate::log::log_error("binding", "recover.rebind_session", &format!("{e}")); + } + if let Err(e) = db.set_process_binding(process_id, session_id, instance_name) { + crate::log::log_error("binding", "recover.set_process_binding", &format!("{e}")); + } + crate::log::log_info( + "binding", + "recover_process_binding_for_instance", + &format!( + "instance={} session_id={} process_id={}", + instance_name, session_id, process_id + ), + ); + Some(instance_name.to_string()) +} + /// Initialize the DB row and default bindings for an instance identity. /// /// This is the shared setup path used by launch, resume, and orphan recovery. diff --git a/src/omp_plugin/hcom.ts b/src/omp_plugin/hcom.ts index f8a19192..5ff121c3 100644 --- a/src/omp_plugin/hcom.ts +++ b/src/omp_plugin/hcom.ts @@ -72,9 +72,16 @@ function isBodylessWake(text: string): boolean { return trimmed === "" || trimmed === ""; } +// Same-process latch: OMP task subagents load a fresh extension instance in the +// parent Node process (SessionShutdownEvent has no sessionId). The first binder +// owns hcom identity; nested instances skip bind/stop so dispose cannot soft-stop +// the parent (lefo / task repro). ExtensionContext does not expose taskDepth. +const IDENTITY_OWNER_ENV = "HCOM_OMP_IDENTITY_OWNER"; + export default function hcomExtension(pi: ExtensionAPI) { let instanceName: string | null = null; let sessionId: string | null = null; + let ownsIdentity = false; let bootstrapText: string | null = null; let bindingPromise: Promise | null = null; let notifyServer: Server | null = null; @@ -142,6 +149,13 @@ export default function hcomExtension(pi: ExtensionAPI) { currentCtx = ctx; if (instanceName || bindingPromise) return bindingPromise ?? Promise.resolve(); if (process.env.HCOM_LAUNCHED !== "1") return; + // Nested task extension instance: parent already claimed identity in this process. + if (process.env[IDENTITY_OWNER_ENV] && !ownsIdentity) { + log("INFO", "plugin.bind_skipped_nested", null, { + owner: process.env[IDENTITY_OWNER_ENV], + }); + return; + } bindingPromise = (async () => { try { const sid = ctx.sessionManager.getSessionId(); @@ -164,6 +178,8 @@ export default function hcomExtension(pi: ExtensionAPI) { } instanceName = json.name; sessionId = json.session_id || sid; + ownsIdentity = true; + process.env[IDENTITY_OWNER_ENV] = instanceName ?? "1"; bootstrapText = typeof json.bootstrap === "string" ? json.bootstrap : null; log("INFO", "plugin.bound", instanceName, { session_id: sessionId, @@ -355,7 +371,7 @@ export default function hcomExtension(pi: ExtensionAPI) { if (!reconcileTimer) reconcileTimer = setInterval(() => void reconcile(), 5_000); } - function resetBinding(): void { + function resetBinding(opts?: { keepOwner?: boolean }): void { stopNotifyServer(); bindingGeneration++; instanceName = null; @@ -372,6 +388,12 @@ export default function hcomExtension(pi: ExtensionAPI) { lastPendingPollAt = 0; agentActive = false; clearIdleTimer(); + // keepOwner: session_branch rebinds in the same extension instance — retain + // process-env ownership so nested task extensions still skip bind. + if (!opts?.keepOwner && ownsIdentity) { + delete process.env[IDENTITY_OWNER_ENV]; + ownsIdentity = false; + } } pi.on("session_start", async (_event, ctx) => { @@ -381,16 +403,41 @@ export default function hcomExtension(pi: ExtensionAPI) { startReconcileTimer(); }); - pi.on("session_shutdown", async (_event) => { - if (instanceName) { - await hcom(["omp-stop", "--name", instanceName, "--reason", "shutdown"]); + // SessionShutdownEvent is only `{ type: "session_shutdown" }` — no sessionId. + // Soft-stop only when THIS extension instance owns the identity (nested task + // instances never bind, so they never stop the parent). + pi.on("session_shutdown", async () => { + if (instanceName && ownsIdentity) { + const reason = "shutdown"; + try { + const result = await hcom(["omp-stop", "--name", instanceName, "--reason", reason, "--soft"]); + if (result.code !== 0) { + log("WARN", "plugin.session_shutdown_soft_stop_failed", instanceName, { + exit_code: result.code, + reason, + stderr: result.stderr.slice(0, 300), + }); + } + } catch (error) { + log("ERROR", "plugin.session_shutdown_soft_stop_error", instanceName, { + error: String(error), + reason, + }); + } + } else if (process.env[IDENTITY_OWNER_ENV] && !ownsIdentity) { + log("INFO", "plugin.session_shutdown_skipped", null, { + reason: "nested_session", + owner: process.env[IDENTITY_OWNER_ENV], + }); } resetBinding(); }); pi.on("session_switch", async (_event, ctx) => { currentCtx = ctx; - resetBinding(); + // Keep process-env ownership across switch rebind so a live nested task + // extension cannot claim identity in the window between clear and omp-start. + resetBinding({ keepOwner: true }); await bindIdentity(ctx); }); @@ -398,11 +445,12 @@ export default function hcomExtension(pi: ExtensionAPI) { // which mints a NEW session id + file and emits only session_branch — not // session_switch. Without rebinding here the cached sessionId stays stale, // isBoundSession() fails against the new id, and every later deliverPending - // silently returns false (delivery dead after branch). Rebind like a switch. + // silently returns false (delivery dead after branch). Rebind like a switch, + // but keep process-env ownership so nested task extensions still skip bind. // session_tree does NOT mint a new session id/file, so it needs no rebind. pi.on("session_branch", async (_event, ctx) => { currentCtx = ctx; - resetBinding(); + resetBinding({ keepOwner: true }); await bindIdentity(ctx); }); From a80463a2198c7fec4f1e13f7417d77745426c8ad Mon Sep 17 00:00:00 2001 From: Lennart Kotzur Date: Wed, 22 Jul 2026 21:43:48 +0200 Subject: [PATCH 2/2] fix(omp): harden nested identity latch and soft-stop recover --- src/delivery.rs | 58 ++++++++++++--- src/delivery/antigravity.rs | 9 +-- src/hooks/common.rs | 81 +++++++++++++++++--- src/hooks/gemini.rs | 2 +- src/hooks/omp/handlers.rs | 3 +- src/hooks/omp/tests.rs | 133 +++++++++++++++++++++++++++++++-- src/instance_binding.rs | 25 ++++++- src/omp_plugin/hcom.ts | 142 ++++++++++++++++++++++++++++++------ 8 files changed, 391 insertions(+), 62 deletions(-) diff --git a/src/delivery.rs b/src/delivery.rs index 7c8406cf..ed4c032a 100644 --- a/src/delivery.rs +++ b/src/delivery.rs @@ -2005,7 +2005,10 @@ pub fn run_delivery_loop( let owns_instance = instance_owns_process_binding(db, &process_id, ¤t_name); - if matches!(Tool::from_str(&config.tool), Ok(Tool::Antigravity)) { + if matches!( + Tool::from_str(&config.tool), + Ok(Tool::Antigravity | Tool::Omp) + ) { antigravity::cleanup_antigravity_pty_exit(db, ¤t_name, &process_id, owns_instance); } else { cleanup_pty_exit_default(db, ¤t_name, &process_id, owns_instance); @@ -2085,6 +2088,25 @@ pub(crate) fn cleanup_deleted_instance(db: &mut HcomDb, current_name: &str) { } } +/// Log why PTY exit cleanup was skipped when this thread no longer owns the instance. +pub(crate) fn log_pty_cleanup_skipped(db: &HcomDb, current_name: &str) { + let reason = if db + .get_status(current_name) + .ok() + .flatten() + .is_some_and(|(status, _)| status == ST_INACTIVE) + { + "instance inactive (soft-finalize); process binding cleared or reassigned" + } else { + "name reassigned to new process" + }; + log_info( + "native", + "delivery.cleanup_skipped", + &format!("Skipping instance cleanup for {current_name} — {reason}"), + ); +} + fn cleanup_pty_exit_default( db: &mut HcomDb, current_name: &str, @@ -2094,14 +2116,7 @@ fn cleanup_pty_exit_default( if owns_instance { cleanup_deleted_instance(db, current_name); } else { - log_info( - "native", - "delivery.cleanup_skipped", - &format!( - "Skipping instance cleanup for {} — name reassigned to new process", - current_name - ), - ); + log_pty_cleanup_skipped(db, current_name); } if !process_id.is_empty() @@ -2231,6 +2246,31 @@ mod tests { assert_eq!(events, vec![("samu".to_string(), "killed".to_string())]); } + #[test] + fn soft_stopped_instance_survives_pty_exit_cleanup() { + let dir = tempfile::tempdir().unwrap(); + let db_path = dir.path().join("test.db"); + let mut db = HcomDb::open_raw(&db_path).unwrap(); + db.init_db().unwrap(); + db.conn() + .execute( + "INSERT INTO instances (name, tool, status, status_context, status_time, created_at, session_id) + VALUES ('luna', 'omp', 'inactive', 'exit:turn_end', 0, 0, 'sid-soft')", + [], + ) + .unwrap(); + db.set_process_binding("pid-soft", "sid-soft", "luna") + .unwrap(); + + antigravity::cleanup_antigravity_pty_exit(&mut db, "luna", "pid-soft", true); + + assert!(db.get_instance_full("luna").unwrap().is_some()); + assert_eq!( + db.get_status("luna").unwrap().map(|(s, _)| s), + Some(ST_INACTIVE.to_string()) + ); + } + // ---- phase-1 ownership tests ---- #[test] diff --git a/src/delivery/antigravity.rs b/src/delivery/antigravity.rs index 28414688..14e815ee 100644 --- a/src/delivery/antigravity.rs +++ b/src/delivery/antigravity.rs @@ -18,14 +18,7 @@ pub(crate) fn cleanup_antigravity_pty_exit( owns_instance: bool, ) { if !owns_instance { - log_info( - "native", - "delivery.cleanup_skipped", - &format!( - "Skipping instance cleanup for {} — name reassigned to new process", - current_name - ), - ); + super::log_pty_cleanup_skipped(db, current_name); } else { let already_inactive = db .get_status(current_name) diff --git a/src/hooks/common.rs b/src/hooks/common.rs index bf3c4bd0..b54ea2d2 100644 --- a/src/hooks/common.rs +++ b/src/hooks/common.rs @@ -1199,13 +1199,17 @@ fn stop_instance_inner( /// agy's real teardown is the PTY exit (`cleanup_antigravity_pty_exit`), which sees /// the inactive status and preserves the row for `hcom r`. /// -/// Clears session/process bindings and logs a stopped life event with snapshot, but -/// does not delete the instance row. +/// Clears session bindings (and process bindings unless `keep_process_binding`), +/// and logs a stopped life event with snapshot, but does not delete the instance row. +/// +/// OMP soft-stop passes `keep_process_binding: true` so the live process can rebind +/// via `bind_session_to_process` on the next turn. Antigravity passes `false`. pub fn soft_finalize_session( db: &HcomDb, instance_name: &str, reason: &str, updates: Option<&serde_json::Map>, + keep_process_binding: bool, ) { log::log_info( "hooks", @@ -1256,17 +1260,21 @@ pub fn soft_finalize_session( "DELETE FROM session_bindings WHERE session_id = ?", params![session_id], ); - let _ = db.conn().execute( - "DELETE FROM process_bindings WHERE session_id = ?", - params![session_id], - ); + if !keep_process_binding { + let _ = db.conn().execute( + "DELETE FROM process_bindings WHERE session_id = ?", + params![session_id], + ); + } } let _ = db.delete_notify_endpoints(instance_name); - let _ = db.conn().execute( - "DELETE FROM process_bindings WHERE instance_name = ?", - params![instance_name], - ); + if !keep_process_binding { + let _ = db.conn().execute( + "DELETE FROM process_bindings WHERE instance_name = ?", + params![instance_name], + ); + } let _ = db.cleanup_subscriptions(instance_name); if let Err(e) = db.log_life_event( @@ -1945,7 +1953,15 @@ mod tests { ) .unwrap(); - soft_finalize_session(&db, "vine", "unknown", None); + db.conn() + .execute( + "INSERT INTO process_bindings (process_id, session_id, instance_name, updated_at) + VALUES ('pid-soft', 'sess-soft-1', 'vine', ?1)", + rusqlite::params![now], + ) + .unwrap(); + + soft_finalize_session(&db, "vine", "unknown", None, false); assert!(db.get_instance_full("vine").unwrap().is_some()); let status = db.get_status("vine").unwrap().map(|(s, _)| s); @@ -1957,5 +1973,48 @@ mod tests { .as_deref(), Some("vine") ); + assert_eq!(db.get_process_binding("pid-soft").unwrap(), None); + } + + #[test] + fn soft_finalize_session_can_keep_process_binding() { + let dir = tempfile::tempdir().unwrap(); + let db_path = dir.path().join("test.db"); + let db = crate::db::HcomDb::open_raw(&db_path).unwrap(); + db.init_db().unwrap(); + let now = chrono::Utc::now().timestamp() as f64; + db.conn() + .execute( + "INSERT INTO instances (name, status, created_at, tool, session_id) + VALUES ('luna', 'listening', ?1, 'omp', 'sess-keep')", + rusqlite::params![now], + ) + .unwrap(); + db.conn() + .execute( + "INSERT INTO session_bindings (session_id, instance_name, created_at) + VALUES ('sess-keep', 'luna', ?1)", + rusqlite::params![now], + ) + .unwrap(); + db.conn() + .execute( + "INSERT INTO process_bindings (process_id, session_id, instance_name, updated_at) + VALUES ('pid-keep', 'sess-keep', 'luna', ?1)", + rusqlite::params![now], + ) + .unwrap(); + + soft_finalize_session(&db, "luna", "turn_end", None, true); + + assert_eq!( + db.get_process_binding("pid-keep").unwrap(), + Some("luna".to_string()) + ); + assert_eq!(db.get_session_binding("sess-keep").unwrap(), None); + assert_eq!( + db.get_status("luna").unwrap().map(|(s, _)| s), + Some(ST_INACTIVE.to_string()) + ); } } diff --git a/src/hooks/gemini.rs b/src/hooks/gemini.rs index 84ffc10f..5b625e6c 100644 --- a/src/hooks/gemini.rs +++ b/src/hooks/gemini.rs @@ -637,7 +637,7 @@ fn handle_sessionend(db: &HcomDb, _ctx: &HcomContext, payload: &HookPayload) -> // that's about to take another turn. Real teardown is the PTY exit. Other gemini // tools have a genuine SessionEnd == process death, so they hard-finalize. if is_agy { - common::soft_finalize_session(db, &instance.name, &reason, None); + common::soft_finalize_session(db, &instance.name, &reason, None, false); } else { common::finalize_session(db, &instance.name, &reason, None); } diff --git a/src/hooks/omp/handlers.rs b/src/hooks/omp/handlers.rs index 8a3e40d6..fbe5a35d 100644 --- a/src/hooks/omp/handlers.rs +++ b/src/hooks/omp/handlers.rs @@ -64,7 +64,6 @@ fn initialize_last_event_id(db: &HcomDb, instance_name: &str) { fn instance_name_from_env(ctx: &HcomContext) -> Option { ctx.raw_env .get("HCOM_INSTANCE_NAME") - .or_else(|| ctx.raw_env.get("HCOM_NAME")) .filter(|s| !s.is_empty()) .cloned() } @@ -294,7 +293,7 @@ pub(crate) fn handle_stop(db: &HcomDb, argv: &[String]) -> (i32, String) { }; let reason = parse_flag(argv, "--reason").unwrap_or_else(|| "unknown".to_string()); if has_flag(argv, "--soft") { - common::soft_finalize_session(db, &name, &reason, None); + common::soft_finalize_session(db, &name, &reason, None, true); (0, r#"{"ok":true,"soft":true}"#.to_string()) } else { finalize_session(db, &name, &reason, None); diff --git a/src/hooks/omp/tests.rs b/src/hooks/omp/tests.rs index ebb8ea5d..d32cbf3f 100644 --- a/src/hooks/omp/tests.rs +++ b/src/hooks/omp/tests.rs @@ -408,7 +408,7 @@ fn start_handler_uses_central_binding_for_existing_session() { } #[test] -fn soft_stop_keeps_instance_row_clears_process_binding() { +fn soft_stop_keeps_instance_row_and_process_binding() { let (db, path) = setup_test_db(); let now = chrono::Utc::now().timestamp() as f64; db.conn() @@ -442,14 +442,17 @@ fn soft_stop_keeps_instance_row_clears_process_binding() { db.get_status("luna").unwrap().map(|(s, _)| s), Some(crate::shared::ST_INACTIVE.to_string()) ); - assert_eq!(db.get_process_binding("pid-soft").unwrap(), None); + assert_eq!( + db.get_process_binding("pid-soft").unwrap(), + Some("luna".to_string()) + ); assert_eq!(db.get_session_binding("sid-soft").unwrap(), None); cleanup(path); } #[test] -fn start_recovers_binding_after_soft_stop_via_instance_name_env() { +fn start_rebinds_via_process_binding_after_soft_stop() { let (db, path) = setup_test_db(); let temp = tempfile::TempDir::new().unwrap(); save_test_instance(&db, "miso", ST_LISTENING); @@ -472,11 +475,13 @@ fn start_recovers_binding_after_soft_stop_via_instance_name_env() { ], ); assert_eq!(stop_code, 0); - assert_eq!(db.get_process_binding("pid-recover").unwrap(), None); + assert_eq!( + db.get_process_binding("pid-recover").unwrap(), + Some("miso".to_string()) + ); let env = std::collections::HashMap::from([ ("HCOM_PROCESS_ID".to_string(), "pid-recover".to_string()), - ("HCOM_INSTANCE_NAME".to_string(), "miso".to_string()), ("HCOM_LAUNCHED".to_string(), "1".to_string()), ("HCOM_TOOL".to_string(), "omp".to_string()), ]); @@ -511,6 +516,124 @@ fn start_recovers_binding_after_soft_stop_via_instance_name_env() { cleanup(path); } +#[test] +fn start_recovers_binding_via_instance_name_when_process_binding_cleared() { + let (db, path) = setup_test_db(); + let temp = tempfile::TempDir::new().unwrap(); + save_test_instance(&db, "miso", ST_LISTENING); + db.conn() + .execute( + "UPDATE instances SET session_id = 'sid-old', status = 'inactive' WHERE name = 'miso'", + [], + ) + .unwrap(); + db.rebind_session("sid-old", "miso").unwrap(); + + let env = std::collections::HashMap::from([ + ("HCOM_PROCESS_ID".to_string(), "pid-recover".to_string()), + ("HCOM_INSTANCE_NAME".to_string(), "miso".to_string()), + ("HCOM_LAUNCHED".to_string(), "1".to_string()), + ("HCOM_TOOL".to_string(), "omp".to_string()), + ]); + let ctx = HcomContext::from_env(&env, temp.path().to_path_buf()); + + let (code, output) = handle_start( + &ctx, + &db, + &[ + "--session-id".to_string(), + "sid-new".to_string(), + "--cwd".to_string(), + temp.path().to_string_lossy().to_string(), + ], + ); + + assert_eq!(code, 0); + let response: serde_json::Value = serde_json::from_str(&output).unwrap(); + assert_eq!(response.get("name").and_then(|v| v.as_str()), Some("miso")); + assert_eq!( + db.get_process_binding("pid-recover").unwrap(), + Some("miso".to_string()) + ); + assert_eq!( + db.get_session_binding("sid-new").unwrap(), + Some("miso".to_string()) + ); + + cleanup(path); +} + +#[test] +fn recover_rejects_non_omp_tool() { + let (db, path) = setup_test_db(); + save_test_instance(&db, "luna", crate::shared::ST_INACTIVE); + db.conn() + .execute( + "UPDATE instances SET tool = 'claude' WHERE name = 'luna'", + [], + ) + .unwrap(); + + assert_eq!( + crate::instance_binding::recover_process_binding_for_instance( + &db, "luna", "sid-new", "pid-1" + ), + None + ); + + cleanup(path); +} + +#[test] +fn recover_rejects_active_instance() { + let (db, path) = setup_test_db(); + save_test_instance(&db, "luna", ST_LISTENING); + + assert_eq!( + crate::instance_binding::recover_process_binding_for_instance( + &db, "luna", "sid-new", "pid-1" + ), + None + ); + + cleanup(path); +} + +#[test] +fn plugin_source_pins_soft_stop_polish_markers() { + assert!(PLUGIN_SOURCE.contains("Symbol.for(\"hcom.omp.identity\")")); + assert!(PLUGIN_SOURCE.contains("stopReconcileTimer")); + assert!(PLUGIN_SOURCE.contains("tearingDown")); + assert!(PLUGIN_SOURCE.contains("HCOM_TIMEOUT_MS")); + assert!(PLUGIN_SOURCE.contains("1800")); + let shutdown_idx = PLUGIN_SOURCE + .find("pi.on(\"session_shutdown\"") + .expect("session_shutdown handler present"); + let handler = &PLUGIN_SOURCE[shutdown_idx..]; + assert!( + handler.contains("keepOwner = !softStopOk"), + "failed soft-stop path must keepOwner" + ); + assert!( + handler.contains("reg.tearingDown = false"), + "failed soft-stop keepOwner must clear tearingDown so owner can rebind" + ); + assert!( + PLUGIN_SOURCE.contains("reg.tearingDown && !ownsIdentity"), + "tearingDown must not block the owning extension from rebinding" + ); +} + +/// Manual/CI scenario: nested OMP task extension must not steal parent identity +/// after soft-stop + session_branch rebind. Run with: +/// `cargo test omp_nested_task_identity_survives_soft_stop -- --ignored` +#[test] +#[ignore = "requires live OMP nested task; run manually"] +fn omp_nested_task_identity_survives_soft_stop() { + // Launch OMP via hcom, soft-stop parent turn, spawn nested task extension, + // verify child does not bind and parent recovers on next turn. +} + // ── Plugin install/remove safety ────────────────────────────────── /// Helper: run a closure with a temp HOME + HCOM_DIR (via isolated_test_env), diff --git a/src/instance_binding.rs b/src/instance_binding.rs index 33884a50..8a9dd823 100644 --- a/src/instance_binding.rs +++ b/src/instance_binding.rs @@ -634,8 +634,8 @@ pub fn bind_session_to_process( /// Rebind process/session after soft-finalize cleared bindings but left the /// instance row (typically inactive). Used when `bind_session_to_process` finds -/// no process binding, but the caller still knows the instance name (e.g. via -/// `HCOM_INSTANCE_NAME` / `HCOM_NAME` in a live OMP process). +/// no process binding, but the caller still knows the instance name via +/// `HCOM_INSTANCE_NAME` in a live OMP process. pub fn recover_process_binding_for_instance( db: &HcomDb, instance_name: &str, @@ -645,20 +645,37 @@ pub fn recover_process_binding_for_instance( if instance_name.is_empty() || session_id.is_empty() || process_id.is_empty() { return None; } - db.get_instance_full(instance_name).ok().flatten()?; + + let instance = db.get_instance_full(instance_name).ok().flatten()?; + + if instance.tool != "omp" { + return None; + } + if instance.status != ST_INACTIVE { + return None; + } if let Err(e) = db.clear_session_id_from_other_instances(session_id, instance_name) { crate::log::log_error("binding", "recover.clear_session", &format!("{e}")); + return None; } + let mut updates = serde_json::Map::new(); updates.insert("session_id".into(), serde_json::json!(session_id)); - update_instance_position(db, instance_name, &updates); + if let Err(e) = db.update_instance_fields(instance_name, &updates) { + crate::log::log_error("binding", "recover.update_session_id", &format!("{e}")); + return None; + } + if let Err(e) = db.rebind_session(session_id, instance_name) { crate::log::log_error("binding", "recover.rebind_session", &format!("{e}")); + return None; } if let Err(e) = db.set_process_binding(process_id, session_id, instance_name) { crate::log::log_error("binding", "recover.set_process_binding", &format!("{e}")); + return None; } + crate::log::log_info( "binding", "recover_process_binding_for_instance", diff --git a/src/omp_plugin/hcom.ts b/src/omp_plugin/hcom.ts index 5ff121c3..0df28ac8 100644 --- a/src/omp_plugin/hcom.ts +++ b/src/omp_plugin/hcom.ts @@ -34,11 +34,27 @@ function log( } catch {} } +const HCOM_TIMEOUT_MS = 1800; + function hcom(args: string[]): Promise { return new Promise((resolve) => { const child = spawn("hcom", args, { stdio: ["ignore", "pipe", "pipe"] }); let stdout = ""; let stderr = ""; + let settled = false; + const finish = (code: number) => { + if (settled) return; + settled = true; + clearTimeout(timer); + resolve({ code, stdout, stderr }); + }; + const timer = setTimeout(() => { + try { + child.kill("SIGTERM"); + } catch {} + finish(124); + }, HCOM_TIMEOUT_MS); + timer.unref?.(); child.stdout.setEncoding("utf8"); child.stderr.setEncoding("utf8"); child.stdout.on("data", (chunk) => { @@ -47,8 +63,11 @@ function hcom(args: string[]): Promise { child.stderr.on("data", (chunk) => { stderr += chunk; }); - child.on("error", (error) => resolve({ code: 127, stdout, stderr: String(error) })); - child.on("close", (code) => resolve({ code: code ?? 0, stdout, stderr })); + child.on("error", (error) => { + stderr = stderr || String(error); + finish(127); + }); + child.on("close", (code) => finish(code === null ? 1 : code)); }); } @@ -76,12 +95,39 @@ function isBodylessWake(text: string): boolean { // parent Node process (SessionShutdownEvent has no sessionId). The first binder // owns hcom identity; nested instances skip bind/stop so dispose cannot soft-stop // the parent (lefo / task repro). ExtensionContext does not expose taskDepth. +const IDENTITY_REGISTRY_KEY = Symbol.for("hcom.omp.identity"); const IDENTITY_OWNER_ENV = "HCOM_OMP_IDENTITY_OWNER"; +type OmpIdentityRegistry = { + owner: string | null; + tearingDown: boolean; +}; + +function getIdentityRegistry(): OmpIdentityRegistry { + const g = globalThis as Record; + if (!g[IDENTITY_REGISTRY_KEY]) { + g[IDENTITY_REGISTRY_KEY] = { owner: null, tearingDown: false }; + } + return g[IDENTITY_REGISTRY_KEY]!; +} + +function syncIdentityOwnerEnv(owner: string | null): void { + if (owner) process.env[IDENTITY_OWNER_ENV] = owner; + else delete process.env[IDENTITY_OWNER_ENV]; +} + +function clearIdentityOwnership(): void { + const reg = getIdentityRegistry(); + reg.owner = null; + reg.tearingDown = false; + syncIdentityOwnerEnv(null); +} + export default function hcomExtension(pi: ExtensionAPI) { let instanceName: string | null = null; let sessionId: string | null = null; let ownsIdentity = false; + let nestedOptOut = false; let bootstrapText: string | null = null; let bindingPromise: Promise | null = null; let notifyServer: Server | null = null; @@ -145,19 +191,42 @@ export default function hcomExtension(pi: ExtensionAPI) { notifyPort = null; } + function nestedSkipReason(): string | null { + const reg = getIdentityRegistry(); + if (nestedOptOut) return "sticky_nested_opt_out"; + // Owner may rebind after a failed soft-stop while keepOwner retained the latch; + // tearingDown only blocks nested/non-owner extensions. + if (reg.tearingDown && !ownsIdentity) return "tearing_down"; + if (reg.owner && !ownsIdentity) return "nested_registry"; + if (!reg.owner && process.env[IDENTITY_OWNER_ENV] && !ownsIdentity) return "nested_env"; + return null; + } + async function bindIdentity(ctx: ExtensionContext): Promise { currentCtx = ctx; if (instanceName || bindingPromise) return bindingPromise ?? Promise.resolve(); if (process.env.HCOM_LAUNCHED !== "1") return; - // Nested task extension instance: parent already claimed identity in this process. - if (process.env[IDENTITY_OWNER_ENV] && !ownsIdentity) { + const skipReason = nestedSkipReason(); + if (skipReason) { + nestedOptOut = true; + const reg = getIdentityRegistry(); log("INFO", "plugin.bind_skipped_nested", null, { - owner: process.env[IDENTITY_OWNER_ENV], + reason: skipReason, + owner: reg.owner ?? process.env[IDENTITY_OWNER_ENV] ?? null, }); return; } bindingPromise = (async () => { try { + const reg = getIdentityRegistry(); + if ((reg.tearingDown && !ownsIdentity) || (reg.owner && !ownsIdentity)) { + nestedOptOut = true; + log("INFO", "plugin.bind_skipped_nested", null, { + reason: reg.tearingDown && !ownsIdentity ? "tearing_down" : "nested_registry", + owner: reg.owner, + }); + return; + } const sid = ctx.sessionManager.getSessionId(); const transcriptPath = ctx.sessionManager.getSessionFile(); const port = await startNotifyServer(); @@ -179,8 +248,10 @@ export default function hcomExtension(pi: ExtensionAPI) { instanceName = json.name; sessionId = json.session_id || sid; ownsIdentity = true; - process.env[IDENTITY_OWNER_ENV] = instanceName ?? "1"; + reg.owner = instanceName; + syncIdentityOwnerEnv(instanceName ?? "1"); bootstrapText = typeof json.bootstrap === "string" ? json.bootstrap : null; + startReconcileTimer(); log("INFO", "plugin.bound", instanceName, { session_id: sessionId, notify_port: port, @@ -368,10 +439,19 @@ export default function hcomExtension(pi: ExtensionAPI) { } function startReconcileTimer(): void { - if (!reconcileTimer) reconcileTimer = setInterval(() => void reconcile(), 5_000); + stopReconcileTimer(); + reconcileTimer = setInterval(() => void reconcile(), 5_000); + } + + function stopReconcileTimer(): void { + if (reconcileTimer) { + clearInterval(reconcileTimer); + reconcileTimer = null; + } } function resetBinding(opts?: { keepOwner?: boolean }): void { + stopReconcileTimer(); stopNotifyServer(); bindingGeneration++; instanceName = null; @@ -389,9 +469,9 @@ export default function hcomExtension(pi: ExtensionAPI) { agentActive = false; clearIdleTimer(); // keepOwner: session_branch rebinds in the same extension instance — retain - // process-env ownership so nested task extensions still skip bind. + // process-local ownership so nested task extensions still skip bind. if (!opts?.keepOwner && ownsIdentity) { - delete process.env[IDENTITY_OWNER_ENV]; + clearIdentityOwnership(); ownsIdentity = false; } } @@ -400,42 +480,60 @@ export default function hcomExtension(pi: ExtensionAPI) { currentCtx = ctx; resetBinding(); await bindIdentity(ctx); - startReconcileTimer(); }); // SessionShutdownEvent is only `{ type: "session_shutdown" }` — no sessionId. // Soft-stop only when THIS extension instance owns the identity (nested task // instances never bind, so they never stop the parent). pi.on("session_shutdown", async () => { + let keepOwner = false; if (instanceName && ownsIdentity) { + const reg = getIdentityRegistry(); + reg.tearingDown = true; const reason = "shutdown"; + const stopName = instanceName; + let softStopOk = false; try { - const result = await hcom(["omp-stop", "--name", instanceName, "--reason", reason, "--soft"]); - if (result.code !== 0) { - log("WARN", "plugin.session_shutdown_soft_stop_failed", instanceName, { + const result = await hcom(["omp-stop", "--name", stopName, "--reason", reason, "--soft"]); + if (result.code === 0) { + softStopOk = true; + } else { + log("WARN", "plugin.session_shutdown_soft_stop_failed", stopName, { exit_code: result.code, reason, stderr: result.stderr.slice(0, 300), }); } } catch (error) { - log("ERROR", "plugin.session_shutdown_soft_stop_error", instanceName, { + log("ERROR", "plugin.session_shutdown_soft_stop_error", stopName, { error: String(error), reason, }); } - } else if (process.env[IDENTITY_OWNER_ENV] && !ownsIdentity) { - log("INFO", "plugin.session_shutdown_skipped", null, { - reason: "nested_session", - owner: process.env[IDENTITY_OWNER_ENV], - }); + keepOwner = !softStopOk; + // Shutdown attempt finished. If we retain ownership after a failed stop, + // drop tearingDown so the owner can re-omp-start; nested skip still uses reg.owner. + if (keepOwner) { + reg.tearingDown = false; + } + } else { + const skipReason = nestedSkipReason(); + if (skipReason) { + nestedOptOut = true; + const reg = getIdentityRegistry(); + log("INFO", "plugin.session_shutdown_skipped", null, { + reason: "nested_session", + skip_reason: skipReason, + owner: reg.owner ?? process.env[IDENTITY_OWNER_ENV] ?? null, + }); + } } - resetBinding(); + resetBinding({ keepOwner }); }); pi.on("session_switch", async (_event, ctx) => { currentCtx = ctx; - // Keep process-env ownership across switch rebind so a live nested task + // Keep process-local ownership across switch rebind so a live nested task // extension cannot claim identity in the window between clear and omp-start. resetBinding({ keepOwner: true }); await bindIdentity(ctx); @@ -446,7 +544,7 @@ export default function hcomExtension(pi: ExtensionAPI) { // session_switch. Without rebinding here the cached sessionId stays stale, // isBoundSession() fails against the new id, and every later deliverPending // silently returns false (delivery dead after branch). Rebind like a switch, - // but keep process-env ownership so nested task extensions still skip bind. + // but keep process-local ownership so nested task extensions still skip bind. // session_tree does NOT mint a new session id/file, so it needs no rebind. pi.on("session_branch", async (_event, ctx) => { currentCtx = ctx;