diff --git a/src/delivery/grok.rs b/src/delivery/grok.rs index 277be5a9..0441fe3f 100644 --- a/src/delivery/grok.rs +++ b/src/delivery/grok.rs @@ -26,6 +26,11 @@ //! is part of the model's turn, or when the prompt request returns a result. //! A batch that never started (removed from the queue, transport lost) stays //! unread and is queued again later: at-least-once, never silently dropped. +//! +//! When the leader restarts (it does so to apply an auto-update), the bridge +//! reconnects to the new one, replays the session loads and announces it with +//! `x.ai/leader_reconnected`. Requests the old leader had not answered never +//! will be, so hcom handles it as a lost connection. use std::collections::{HashMap, HashSet}; use std::io::{BufRead, BufReader, Write}; @@ -58,7 +63,7 @@ const ROSTER_FAST_WINDOW: Duration = Duration::from_secs(5); const ROSTER_SLOW_POLL: Duration = Duration::from_secs(15); /// Wait before re-queueing a batch whose prompt was dropped before it ran /// (user removed it from the queue, prompt error, transport lost). -const REQUEUE_DELAY: Duration = Duration::from_secs(10); +const REQUEUE_DELAY: Duration = Duration::from_secs(if cfg!(test) { 1 } else { 10 }); /// Wait before loading the bound session again after a failed load. const LOAD_RETRY_DELAY: Duration = Duration::from_secs(2); /// Prompts hcom queues carry this `promptId` prefix. @@ -746,13 +751,20 @@ struct InFlight { /// prompt when its client disconnects, so the batch is followed through /// the reconnected client's queue snapshots, not treated as dropped. detached_at: Option, + /// Whether the latest roster since the drop shows its session idle. + /// Reconnecting brings no queue snapshot, and Grok runs a queued prompt + /// as soon as its session is idle, so only an idle session proves the + /// batch is no longer waiting behind a turn. (A row the user holds for + /// editing at the front can still park it, for up to Grok's edit-hold + /// TTL; that stays a possible duplicate.) + idle_after_detach: bool, } /// How long to wait after `removedFromQueue` for the snapshot showing a merge. const MERGE_EVIDENCE_WAIT: Duration = Duration::from_secs(2); /// How long a batch whose connection dropped waits for a queue snapshot -/// before the transcript decides. -const DETACHED_WAIT: Duration = Duration::from_secs(30); +/// before the transcript decides, once its session is idle. +const DETACHED_WAIT: Duration = Duration::from_secs(if cfg!(test) { 2 } else { 30 }); #[derive(Debug, PartialEq, Eq)] enum Outcome { @@ -840,6 +852,15 @@ fn roster_sessions(result: &Value) -> &[Value] { .unwrap_or_default() } +/// Whether the roster shows `session` in a turn: running, or waiting on an +/// approval or question (`needs_input` wins over `working` mid-turn). +fn roster_busy(sessions: &[Value], session: &str) -> bool { + sessions.iter().any(|entry| { + entry["sessionId"] == session + && matches!(entry["activity"].as_str(), Some("working" | "needs_input")) + }) +} + /// Point the instance at `session`, as a hook-based SessionStart would. /// Returns the instance name, which a switch to a session bound to another /// identity changes. @@ -1070,7 +1091,14 @@ pub(super) fn run( log_warn("native", "grok.acp.roster_failed", &error.to_string()) } None => { - effects.extend(tracker.on_roster(roster_sessions(&value["result"]))) + let sessions = roster_sessions(&value["result"]); + if let Some(flight) = in_flight.as_mut() + && flight.detached_at.is_some() + { + flight.idle_after_detach = + !roster_busy(sessions, &flight.session); + } + effects.extend(tracker.on_roster(sessions)) } }, Some(Pending::Load(session)) => match value.get("error") { @@ -1115,6 +1143,10 @@ pub(super) fn run( } } Event::Notification { method, params } => { + if method == "x.ai/leader_reconnected" { + closed = Some("Grok leader restarted".into()); + break; + } if method == "x.ai/queue/changed" && let Some(flight) = in_flight.as_ref() && params["sessionId"].as_str() == Some(flight.session.as_str()) @@ -1219,6 +1251,7 @@ pub(super) fn run( if let Some(flight) = in_flight.as_mut() { flight.request_id = None; flight.detached_at.get_or_insert_with(Instant::now); + flight.idle_after_detach = false; } log_warn("native", "grok.acp.disconnected", &error); } @@ -1233,6 +1266,7 @@ pub(super) fn run( if outcome.is_none() && in_flight .as_ref() + .filter(|flight| flight.idle_after_detach) .and_then(|flight| flight.detached_at) .is_some_and(|at| at.elapsed() >= DETACHED_WAIT) { @@ -1366,6 +1400,7 @@ pub(super) fn run( ack: prepared.ack, removed_at: None, detached_at: None, + idle_after_detach: false, }); } Err(error) => { @@ -1404,6 +1439,7 @@ mod tests { }, removed_at: removed.then(Instant::now), detached_at: None, + idle_after_detach: false, } } @@ -1464,6 +1500,18 @@ mod tests { )); } + #[test] + fn roster_busy_includes_pending_approvals() { + let roster = |activity: &str| { + vec![json!({"sessionId": "s", "resident": true, "activity": activity})] + }; + assert!(roster_busy(&roster("working"), "s")); + // An approval mid-turn reports needs_input, not working. + assert!(roster_busy(&roster("needs_input"), "s")); + assert!(!roster_busy(&roster("idle"), "s")); + assert!(!roster_busy(&roster("working"), "other")); + } + #[test] fn transcript_proves_a_batch_ran() { let dir = tempfile::tempdir().unwrap(); @@ -1872,4 +1920,178 @@ mod tests { let plain = Launch::new("grok", &[], &["--", "--no-subagents"]).unwrap(); assert!(plain.child_env().is_empty()); } + + /// Stands in for `grok agent --leader --leader-socket stdio`. On + /// the first bridge the leader restarts as hcom's batch goes out: the + /// bridge announces `x.ai/leader_reconnected`, the old leader never + /// answers the prompt, and the new one has an empty queue. With a + /// `queued` flag file beside it, only the bridge drops instead: the + /// batch stays queued behind a turn that ends after `DETACHED_WAIT`. As + /// with Grok, loading a session sends no queue snapshot; only the + /// roster shows the turn. Every prompt the second bridge receives is + /// counted in `requeued`. + #[cfg(unix)] + const FAKE_BRIDGE: &str = r#"#!/bin/sh +dir=$(dirname "$0") +run=$(( $(cat "$dir/runs" 2>/dev/null || echo 0) + 1 )) +echo "$run" > "$dir/runs" +while IFS= read -r line; do + id=$(printf '%s' "$line" | sed -n 's/.*"id":\([0-9][0-9]*\).*/\1/p') + case "$line" in + *'"method":"initialize"'*|*'"method":"session/load"'*) + echo "{\"jsonrpc\":\"2.0\",\"id\":$id,\"result\":{}}" + case "$line" in *session/load*) + if [ -f "$dir/turn" ]; then + # The turn ends and the held batch runs. + held=$(cat "$dir/prompt") + (sleep 4; rm -f "$dir/turn"; echo "{\"jsonrpc\":\"2.0\",\"method\":\"_x.ai/queue/changed\",\"params\":{\"sessionId\":\"s1\",\"runningPromptId\":\"$held\",\"runningKind\":\"prompt\",\"entries\":[]}}") & + fi;; + esac ;; + *'"method":"_x.ai/sessions/list"'*) + activity=idle + [ -f "$dir/turn" ] && activity=working + echo "{\"jsonrpc\":\"2.0\",\"id\":$id,\"result\":{\"sessions\":[{\"sessionId\":\"s1\",\"resident\":true,\"activity\":\"$activity\",\"lastChangeUnixMs\":1}]}}" ;; + *'"method":"session/prompt"'*) + prompt=$(printf '%s' "$line" | sed -n 's/.*"promptId":"\([^"]*\)".*/\1/p') + if [ "$run" = 1 ]; then + echo "$prompt" > "$dir/prompt" + if [ -f "$dir/queued" ]; then + touch "$dir/turn" + exit 0 + fi + echo '{"jsonrpc":"2.0","method":"x.ai/leader_reconnected","params":{}}' + else + echo "$prompt" >> "$dir/requeued" + echo "{\"jsonrpc\":\"2.0\",\"method\":\"_x.ai/queue/changed\",\"params\":{\"sessionId\":\"s1\",\"runningPromptId\":\"$prompt\",\"runningKind\":\"prompt\",\"entries\":[]}}" + fi ;; + esac +done +"#; + + /// Runs one delivery through the fake bridge with the given flag files + /// until the batch is acked or 30 s pass. Returns whether it was acked, + /// how many bridges hcom started, and the prompts the second one got. + #[cfg(unix)] + fn deliver_across_restart(flags: &[&str]) -> (bool, String, Vec) { + use std::os::unix::fs::PermissionsExt; + + let (dir, hcom_dir, _home, _guard) = crate::hooks::test_helpers::isolated_test_env(); + let db_path = hcom_dir.join("test.db"); + let mut db = HcomDb::open_at(&db_path).unwrap(); + let work = dir.path().join("work"); + std::fs::create_dir_all(&work).unwrap(); + db.save_instance_named( + "nova", + &serde_json::Map::from_iter([ + ("tool".into(), json!("grok")), + ("session_id".into(), json!("s1")), + ("directory".into(), json!(work.to_string_lossy())), + ("status".into(), json!(ST_LISTENING)), + ("created_at".into(), json!(1.0)), + ]), + ) + .unwrap(); + db.set_session_binding("s1", "nova").unwrap(); + db.set_process_binding("proc-1", "s1", "nova").unwrap(); + let message = json!({"from": "luna", "text": "are you there?", + "scope": "mentions", "mentions": ["nova"]}); + db.conn() + .execute( + "INSERT INTO events (type, timestamp, instance, data) + VALUES ('message', '2026-01-01T00:00:00Z', 'luna', ?1)", + rusqlite::params![message.to_string()], + ) + .unwrap(); + let message_id = db.conn().last_insert_rowid(); + + let bridge = dir.path().join("grok"); + std::fs::write(&bridge, FAKE_BRIDGE).unwrap(); + std::fs::set_permissions(&bridge, std::fs::Permissions::from_mode(0o755)).unwrap(); + for flag in flags { + std::fs::write(dir.path().join(flag), "").unwrap(); + } + let launch = Launch { + command: bridge.to_string_lossy().into_owned(), + prefix: Vec::new(), + socket: dir + .path() + .join("leader.sock") + .to_string_lossy() + .into_owned(), + no_subagents: false, + }; + + let running = Arc::new(AtomicBool::new(true)); + let watcher = { + let running = running.clone(); + let db_path = db_path.clone(); + std::thread::spawn(move || { + let db = HcomDb::open_at(&db_path).unwrap(); + let deadline = Instant::now() + Duration::from_secs(30); + while db.get_cursor("nova") < message_id && Instant::now() < deadline { + std::thread::sleep(Duration::from_millis(50)); + } + running.store(false, Ordering::Release); + }) + }; + let notify = NotifyServer::new().unwrap(); + let state = DeliveryState { + screen: Arc::new(RwLock::new(Default::default())), + grok_acp: None, + launch_phase_active: Arc::new(AtomicBool::new(false)), + inject_port: 0, + user_activity_cooldown_ms: 0, + }; + let mut name = "nova".to_string(); + run( + &launch, + &running, + &mut db, + ¬ify, + &state, + "proc-1", + &mut name, + &ToolConfig::for_tool(crate::tool::Tool::Grok), + &None, + &None, + &None, + &mut super::super::host_label::HostLabel::resolve(), + &mut LaunchOutcome::Ready, + ); + watcher.join().unwrap(); + + let read = |name: &str| std::fs::read_to_string(dir.path().join(name)).unwrap_or_default(); + ( + db.get_cursor("nova") == message_id, + read("runs").trim().to_string(), + read("requeued").lines().map(str::to_string).collect(), + ) + } + + /// Observed: the leader auto-updated and restarted while idle, and the + /// batch queued across the restart was never answered. It stayed in + /// flight, no later mail was queued, and unread climbed until someone + /// typed into the agent's terminal. + #[cfg(unix)] + #[test] + #[serial_test::serial] + fn batch_lost_in_a_leader_restart_is_queued_again() { + let (acked, runs, requeued) = deliver_across_restart(&[]); + assert!(acked, "the batch was never acked"); + assert_eq!(runs, "2", "hcom should reconnect exactly once"); + assert_eq!(requeued.len(), 1, "the lost batch is queued again once"); + } + + /// Observed: only the bridge dropped while the batch waited behind a long + /// tool call. Reconnecting brought no queue snapshot, hcom queued the + /// batch again, and Grok ran both copies once the tool finished. + #[cfg(unix)] + #[test] + #[serial_test::serial] + fn batch_queued_behind_a_turn_across_a_bridge_drop_is_not_queued_twice() { + let (acked, runs, requeued) = deliver_across_restart(&["queued"]); + assert!(acked, "the batch was never acked"); + assert_eq!(runs, "2", "hcom should reconnect exactly once"); + assert!(requeued.is_empty(), "the queued batch was sent again"); + } }