From fc42f9769171e5a89749a00b44010aaedb4bf758 Mon Sep 17 00:00:00 2001 From: damocles Date: Mon, 29 Jun 2026 19:19:20 +0200 Subject: [PATCH] parallelize graceful agent drains, serialize container stops on fast lane; unify shutdown+checkpoint+compact prompt --- docs/coordinator.md | 2 +- docs/turn-loop.md | 6 ++ hive-ag3nt/src/bin/hive.rs | 26 ++++++-- hive-ag3nt/src/turn.rs | 114 ++++++++++++++++++++++++++------- hive-c0re/src/rebuild_queue.rs | 101 +++++++++++++++++++---------- 5 files changed, 189 insertions(+), 60 deletions(-) diff --git a/docs/coordinator.md b/docs/coordinator.md index cc0f6374..5ae49cc7 100644 --- a/docs/coordinator.md +++ b/docs/coordinator.md @@ -40,7 +40,7 @@ somewhere." | `Destroy` | For future use (`destroy --purge` does real I/O). Variant exists so the wire shape doesn't change later; not currently routed through the queue. | | `Restart` | Stop + start a container without touching config (~5-10s). Routed through the queue so it serialises against in-flight rebuilds for the same agent — prevents a restart racing a rebuild mid-flight. Sources: dashboard ↺ button, the `restart` MCP tool. | | `PermChange` | Write a tool-group or capability change to the shared JSON file (`tool-groups.json` / `capabilities.json`), then rebuild the agent so the updated `HIVE_TOOL_GROUPS` / `HIVE_CAPABILITIES` env var takes effect. Serialising the file write through the queue prevents concurrent dashboard batch-apply actions from racing on the shared file. After a successful file write, emits `CapabilitiesChanged` or `ToolGroupsChanged` SSE snapshot so the P3RM1SS10NS tab updates live. | -| `GracefulStop` | Quiesce then stop a container (the `?graceful=true` path on `/api/kill/`). Signals the harness (its next `Recv` returns `GracefulStop` — the inbound fence — so it runs one stop-checkpoint turn to flush durable `/state`, then exits), waits for it to drain (bounded by a 3-min timeout → hard-stop fallback), then runs the normal container-stop teardown. Queued so it can't race an in-flight rebuild for the same agent. | +| `GracefulStop` | Quiesce then stop a container (the `?graceful=true` path on `/api/kill/`). Signals the harness (its next `Recv` returns `GracefulStop` — the inbound fence — so it runs **one** stop-checkpoint turn that flushes durable `/state` and compacts the session if it crossed the watermark, then exits) and **immediately releases the build lane**, spawning a detached watcher that holds the `Stopping` transient across the drain (bounded by a 3-min timeout → hard-stop fallback) and then enqueues a fast-lane `Stop` (`parent_id` = this entry) for the actual `nixos-container stop`. Net: a whole-hive graceful stop signals every agent up front, drains overlap, and only the container teardowns serialise (on the fast lane). Queued so the signal can't race an in-flight rebuild for the same agent. | **Intentionally not queued** (sub-second ops): the *hard* `start`, `stop`, `kill`. (A *graceful* stop is the `GracefulStop` kind above — it takes a checkpoint turn, so it rides the queue.) diff --git a/docs/turn-loop.md b/docs/turn-loop.md index 942762b1..c30fd865 100644 --- a/docs/turn-loop.md +++ b/docs/turn-loop.md @@ -517,6 +517,12 @@ ttl_seconds?, to?)`, `answer(id, answer)`. explicit `from: "graceful-stop"` message instead of an empty inbox. This unmissably directs the agent to flush durable state (`/state` files) and end the turn — the container exits when the turn completes. + This is a **single** wake: the graceful-stop turn *is* the checkpoint + (its prompt carries the full notes/CLAUDE.md/TODO.md flush guidance), + so the harness compacts directly afterwards if the context crossed the + watermark (`CompactionMode::CompactOnly`) rather than waking a second + dedicated checkpoint turn. Compacting before shutdown keeps a later + cold start cheap instead of re-uploading a huge transcript. - `ask` — surface a structured question to the operator (default) or a peer agent (`to: ""`). Non-blocking — returns a question id; the answer arrives as a `question_answered` system event in the diff --git a/hive-ag3nt/src/bin/hive.rs b/hive-ag3nt/src/bin/hive.rs index 2a67da85..61012d2e 100644 --- a/hive-ag3nt/src/bin/hive.rs +++ b/hive-ag3nt/src/bin/hive.rs @@ -191,8 +191,13 @@ fn graceful_stop_message() -> hive_sh4re::DeliveredMessage { hive_sh4re::DeliveredMessage { from: "graceful-stop".into(), body: "You are being gracefully stopped — the container will shut down after this turn, \ - and new inbound messages are already fenced. Flush anything worth keeping to your \ - durable /state files now, then end your turn. Do not start new long-running work." + and new inbound messages are already fenced. This is your one checkpoint turn: \ + flush anything worth keeping into your durable /state files now — update your \ + notes / CLAUDE.md / TODO.md with in-flight task state, decisions made, important \ + file paths, and whatever you'd need to resume cleanly later with only a summary \ + of this conversation to go on. Do not start new work or reply to anyone; just \ + write your notes and end your turn. The session may be compacted after this turn \ + so a later cold start resumes cheaply." .into(), id: 0, redelivered: false, @@ -541,6 +546,9 @@ async fn serve_loop( stats.as_ref(), files, &turn_lock, + // The graceful prompt IS the checkpoint — compact + // directly afterwards if needed, no second wake. + turn::CompactionMode::CompactOnly, graceful_stop_message(), ) .await; @@ -549,7 +557,16 @@ async fn serve_loop( } }, }; - let ctrl = handle_turn::(socket, &bus, stats.as_ref(), files, &turn_lock, next).await; + let ctrl = handle_turn::( + socket, + &bus, + stats.as_ref(), + files, + &turn_lock, + turn::CompactionMode::CheckpointThenCompact, + next, + ) + .await; if ctrl.auth_failed { *login_state.lock().unwrap() = LoginState::NeedsLogin; turn::wait_for_login( @@ -578,6 +595,7 @@ async fn handle_turn( stats: Option<&TurnStats>, files: &turn::TurnFiles, turn_lock: &TurnLock, + compaction: turn::CompactionMode, first: hive_sh4re::DeliveredMessage, ) -> TurnControl { let from = first.from; @@ -598,7 +616,7 @@ async fn handle_turn( let prompt = serve_common::format_wake_prompt(&from, &body, unread, redelivered); let outcome = { let _guard = turn_lock.lock().await; - turn::drive_turn(&prompt, files, bus).await + turn::drive_turn_with(&prompt, files, bus, compaction).await }; turn::emit_turn_end(bus, &outcome); bus.set_state(TurnState::Idle); diff --git a/hive-ag3nt/src/turn.rs b/hive-ag3nt/src/turn.rs index e44c732b..3503dcc0 100644 --- a/hive-ag3nt/src/turn.rs +++ b/hive-ag3nt/src/turn.rs @@ -263,6 +263,20 @@ fn compact_watermark_tokens(bus: &Bus) -> u64 { effective_context_window(bus) * 3 / 4 } +/// Post-turn compaction behaviour for [`drive_turn`]. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum CompactionMode { + /// Normal turns: if the context crossed the watermark, run a dedicated + /// notes-checkpoint turn so the agent can flush durable state, THEN + /// `/compact`. The checkpoint turn is a separate wake. + CheckpointThenCompact, + /// Graceful-stop turn: the turn that just ran WAS the checkpoint (its + /// prompt already told the agent to flush state), so compact directly + /// if the watermark was crossed — no second wake. Keeps a cold-start + /// resume cheap without waking the agent twice. + CompactOnly, +} + /// Drive one turn end-to-end. Three paths layer on top of the raw `run_turn`: /// /// - **Auto-reset (pre-turn)** — context is large AND the prompt cache has @@ -278,10 +292,24 @@ fn compact_watermark_tokens(bus: &Bus) -> u64 { /// size has crept past the watermark: while the session is still healthy we /// give the agent one dedicated turn to checkpoint its `/state` notes, then /// compact. This keeps a later turn from hitting the reactive path (where -/// there is no chance to save anything first). +/// there is no chance to save anything first). The graceful-stop path uses +/// [`CompactionMode::CompactOnly`] to skip the dedicated checkpoint turn. /// /// Called once per turn by the `hive` serve loop (every agent role). pub async fn drive_turn(prompt: &str, files: &TurnFiles, bus: &Bus) -> TurnOutcome { + drive_turn_with(prompt, files, bus, CompactionMode::CheckpointThenCompact).await +} + +/// As [`drive_turn`] but with explicit control over the post-turn +/// compaction behaviour — see [`CompactionMode`]. The graceful-stop path +/// passes [`CompactionMode::CompactOnly`] so the agent isn't woken a +/// second time for a checkpoint it already did. +pub async fn drive_turn_with( + prompt: &str, + files: &TurnFiles, + bus: &Bus, + compaction: CompactionMode, +) -> TurnOutcome { maybe_auto_reset(bus); let outcome = match run_turn(prompt, files, bus).await { TurnOutcome::PromptTooLong => { @@ -312,12 +340,22 @@ pub async fn drive_turn(prompt: &str, files: &TurnFiles, bus: &Bus) -> TurnOutco other => other, }; // Proactive: a turn just completed on a still-healthy session. If its - // context crossed the watermark, checkpoint + compact before a later - // turn overflows into the reactive path. Best-effort — never changes - // the outcome of the turn that already succeeded, but records it as - // `Compacted` so turn stats can distinguish it from a plain `Ok`. - if matches!(outcome, TurnOutcome::Ok) && maybe_checkpoint_and_compact(files, bus).await { - return TurnOutcome::Compacted; + // context crossed the watermark, compact before a later turn overflows + // into the reactive path. In the normal mode this first runs a separate + // notes-checkpoint turn so the agent can flush durable state; in + // `CompactOnly` mode (graceful stop) the turn that just ran was itself + // the checkpoint, so we compact directly without a second wake. + // Best-effort — never changes the outcome of the turn that already + // succeeded, but records it as `Compacted` so turn stats can distinguish + // it from a plain `Ok`. + if matches!(outcome, TurnOutcome::Ok) { + let compacted = match compaction { + CompactionMode::CheckpointThenCompact => maybe_checkpoint_and_compact(files, bus).await, + CompactionMode::CompactOnly => maybe_compact(files, bus).await, + }; + if compacted { + return TurnOutcome::Compacted; + } } outcome } @@ -329,16 +367,10 @@ pub async fn drive_turn(prompt: &str, files: &TurnFiles, bus: &Bus) -> TurnOutco /// fails the turn that already succeeded. Returns `true` if compaction /// was attempted (watermark crossed), `false` if skipped. async fn maybe_checkpoint_and_compact(files: &TurnFiles, bus: &Bus) -> bool { - let watermark = compact_watermark_tokens(bus); - if watermark == 0 { - return false; // proactive compaction disabled - } - let Some(used) = bus.last_ctx_usage().map(|u| u.context_tokens()) else { - return false; // no usage reading yet — nothing to compare against - }; - if used < watermark { + let Some(used) = watermark_crossed(bus) else { return false; - } + }; + let watermark = compact_watermark_tokens(bus); bus.emit(LiveEvent::Note { text: format!( "context at {used} tokens (watermark {watermark}) — running a \ @@ -363,11 +395,50 @@ async fn maybe_checkpoint_and_compact(files: &TurnFiles, bus: &Bus) -> bool { text: format!("checkpoint turn failed ({e:#}) — compacting anyway"), }), } - // Best-effort: never changes the outcome of the turn that already - // succeeded. Mirror the checkpoint-turn handling above — emit a Note - // for each failure mode and move on; the next real turn will surface - // the underlying issue (rate-limit / 401 / etc.) through the normal - // path anyway. + do_compact(files, bus).await; + true +} + +/// Compact-only proactive path: if the context crossed the watermark, +/// `/compact` directly — NO preceding checkpoint turn. Used by the +/// graceful-stop path, where the turn that just ran already flushed +/// durable state (its prompt said so), so a second checkpoint wake would +/// be redundant. Compacting before the container stops keeps a later +/// cold-start resume cheap rather than re-uploading a huge transcript. +/// Returns `true` if compaction was attempted. +async fn maybe_compact(files: &TurnFiles, bus: &Bus) -> bool { + let Some(used) = watermark_crossed(bus) else { + return false; + }; + let watermark = compact_watermark_tokens(bus); + bus.emit(LiveEvent::Note { + text: format!( + "context at {used} tokens (watermark {watermark}) — compacting before \ + graceful stop so a later cold start resumes cheap" + ), + }); + do_compact(files, bus).await; + true +} + +/// Returns `Some(used_tokens)` when proactive compaction is enabled AND the +/// last inference's context size has reached the watermark; `None` when +/// compaction is disabled (watermark 0), there's no usage reading yet, or +/// the context is still below the watermark. +fn watermark_crossed(bus: &Bus) -> Option { + let watermark = compact_watermark_tokens(bus); + if watermark == 0 { + return None; // proactive compaction disabled + } + let used = bus.last_ctx_usage().map(|u| u.context_tokens())?; + (used >= watermark).then_some(used) +} + +/// Run `/compact`, surfacing each failure mode as a best-effort Note. +/// Never changes the outcome of the turn that already succeeded — the +/// next real turn surfaces any underlying issue (rate-limit / 401 / etc.) +/// through the normal path anyway. +async fn do_compact(files: &TurnFiles, bus: &Bus) { match compact_session(files, bus).await { TurnOutcome::Ok | TurnOutcome::Compacted => {} TurnOutcome::PromptTooLong => bus.emit(LiveEvent::Note { @@ -386,7 +457,6 @@ async fn maybe_checkpoint_and_compact(files: &TurnFiles, bus: &Bus) -> bool { }); } } - true } /// Pre-turn auto-reset check. If context is large AND the prompt cache has diff --git a/hive-c0re/src/rebuild_queue.rs b/hive-c0re/src/rebuild_queue.rs index e40cad20..deddb90e 100644 --- a/hive-c0re/src/rebuild_queue.rs +++ b/hive-c0re/src/rebuild_queue.rs @@ -43,8 +43,11 @@ pub enum QueueKind { /// actions for different agents never race on the shared JSON file. PermChange, /// Gracefully stop a container: signal the harness to run one - /// stop-checkpoint turn (flush durable `/state`), wait for it to drain, - /// then `nixos-container stop`. Falls back to a hard stop on timeout. + /// stop-checkpoint turn (flush durable `/state`), then hand the drain-wait + /// to a detached watcher (freeing the build lane) which, once the agent + /// drains or `GRACEFUL_STOP_TIMEOUT` elapses, enqueues a fast-lane `Stop` + /// for the actual `nixos-container stop`. The build worker only does the + /// cheap signal, so whole-hive graceful stops overlap every agent's drain. GracefulStop, /// Start a stopped container (`lifecycle::start`). Routed through the /// queue so the dashboard shows a visible queued→running transient — a @@ -77,7 +80,8 @@ impl QueueKind { /// serial fast worker concurrently with the build lane (so a stop/start /// never waits behind another container's slow build). `GracefulStop` /// and `Restart` are deliberately NOT fast — they go through the build - /// lane (`GracefulStop` holds the worker while the harness drains; + /// lane (`GracefulStop` does the cheap harness signal then detaches the + /// drain-wait, enqueueing a fast-lane `Stop` for the real container stop; /// `Restart` is a stop+start). pub fn is_fast(self) -> bool { matches!(self, QueueKind::Start | QueueKind::Stop) @@ -752,10 +756,12 @@ fn lane_clear(entries: &VecDeque, e: &QueueEntry) -> bool { /// signal the worker exits after its current entry finishes; pending /// `Queued` entries are dropped (they'll either be replayed by the /// startup sweep on next boot or left for an operator to re-queue). -/// Max time the `GracefulStop` worker waits for the harness to run its +/// Max time the `GracefulStop` drain watcher waits for the harness to run its /// stop-checkpoint turn + drain before falling back to a hard container stop. /// Generous — a checkpoint turn can take a while — but bounded so a wedged -/// agent never blocks the stop indefinitely. +/// agent never blocks the stop indefinitely. The wait runs in a detached +/// watcher task (not the build worker), so a whole-hive graceful stop overlaps +/// every agent's drain instead of serialising N × this timeout. const GRACEFUL_STOP_TIMEOUT: std::time::Duration = std::time::Duration::from_mins(3); /// Run one claimed queue entry to completion: snapshot, dispatch, mark @@ -1000,7 +1006,7 @@ async fn dispatch( crate::auto_update::current_flake_rev(&coord.hyperhive_flake).unwrap_or_default(); crate::auto_update::rebuild_agent(coord, name, ¤t_rev, Some(entry.id), true).await } - (QueueKind::GracefulStop, _) => run_graceful_stop(coord, entry).await, + (QueueKind::GracefulStop, _) => run_graceful_stop(coord, entry), (QueueKind::Start, _) => run_start(coord, entry).await, (QueueKind::Stop, _) => run_stop(coord, entry).await, } @@ -1044,40 +1050,69 @@ async fn run_stop( /// Run one `GracefulStop` entry: signal the harness to quiesce (it returns /// `GracefulStop` on its next `Recv`, runs one stop-checkpoint turn to flush -/// durable `/state`, then exits), wait for it to drain — bounded by -/// `GRACEFUL_STOP_TIMEOUT` so a wedged agent can't block forever — then stop -/// the container with the same teardown as a plain kill. -async fn run_graceful_stop( +/// durable `/state`, then exits), then hand the drain-wait + container stop to +/// a detached watcher and return — freeing the build lane immediately. +/// +/// This is the concurrency split: the build worker only does the cheap signal, +/// so a whole-hive graceful stop signals every agent up front and their +/// checkpoint drains overlap. The watcher waits for this agent's drain +/// (bounded by `GRACEFUL_STOP_TIMEOUT` so a wedged agent can't block forever), +/// then enqueues a fast-lane `Stop` for the actual `nixos-container stop`. +/// Routing the real stop through the fast lane means the container stops +/// serialise there (one stop at a time) while the drains ran in parallel. +// Returns `Result` for symmetry with the other `run_*` dispatch arms even +// though the fallible drain now lives in the detached watcher and this body +// is infallible. +#[allow(clippy::unnecessary_wraps)] +fn run_graceful_stop( coord: &std::sync::Arc, entry: &QueueEntry, ) -> anyhow::Result<()> { - let name = &entry.agent; - let _guard = coord.transient_guard(name, crate::coordinator::TransientKind::Stopping); + let name = entry.agent.clone(); + let parent_id = entry.id; + let source = entry.source; // Signal the harness; the kick breaks an idle long-poll so it's seen promptly. coord.set_queue_step(Some(entry.id), "graceful stop: signalling agent"); - coord.mark_graceful_stop(name); - coord.kick_agent(name, "graceful stop requested"); - // Wait for the harness to drain (it clears the flag via `GracefulStopComplete`) - // or fall back to a hard stop after the timeout. The single queue worker is - // intentionally held for the duration — graceful stops are infrequent. - coord.set_queue_step(Some(entry.id), "graceful stop: waiting for agent to drain"); - let deadline = std::time::Instant::now() + GRACEFUL_STOP_TIMEOUT; - while coord.is_graceful_stop_pending(name) { - if std::time::Instant::now() >= deadline { - tracing::warn!(agent = %name, "graceful stop: drain timed out — hard-stopping"); - break; + coord.mark_graceful_stop(&name); + coord.kick_agent(&name, "graceful stop requested"); + // Detached watcher: wait for the drain (or timeout), then enqueue the + // container stop on the fast lane. The build entry itself is now Done — + // the dashboard groups the follow-up `Stop` under it via `parent_id`. + let coord = std::sync::Arc::clone(coord); + tokio::spawn(async move { + // Hold the `Stopping` transient across the drain so the dashboard keeps + // showing the agent quiescing; dropped before the fast `Stop` is + // enqueued (its `run_stop` re-establishes the transient) so the two + // never clobber each other's clear-on-drop. + let guard = coord.transient_guard(&name, crate::coordinator::TransientKind::Stopping); + // Wait for the harness to drain (it clears the flag via + // `GracefulStopComplete`) or fall back to a hard stop after the timeout. + let deadline = std::time::Instant::now() + GRACEFUL_STOP_TIMEOUT; + while coord.is_graceful_stop_pending(&name) { + if std::time::Instant::now() >= deadline { + tracing::warn!(agent = %name, "graceful stop: drain timed out — hard-stopping"); + break; + } + tokio::time::sleep(std::time::Duration::from_millis(500)).await; } - tokio::time::sleep(std::time::Duration::from_millis(500)).await; - } - coord.clear_graceful_stop(name); - // Stop the container — same teardown as a plain kill. - coord.set_queue_step(Some(entry.id), "nixos-container stop"); - crate::lifecycle::kill(name).await?; - coord.unregister_agent(name); - coord.notify_manager(&hive_sh4re::HelperEvent::Killed { - agent: name.clone(), + coord.clear_graceful_stop(&name); + drop(guard); + // Enqueue the actual container stop on the fast lane (same teardown as a + // plain kill — `run_stop`). `parent_id` links it to the graceful entry + // for dashboard grouping. `enqueue_full` nudges the fast worker itself. + coord.rebuild_queue.enqueue_full(FullEnqueue { + kind: QueueKind::Stop, + agent: name.clone(), + source, + reason: format!("container stop after graceful drain of {name}"), + parent_id: Some(parent_id), + inputs: Vec::new(), + approval_id: None, + perm_payload: None, + depends_on: Vec::new(), + }); + coord.emit_rebuild_queue_snapshot(); }); - coord.rescan_containers_and_emit().await; Ok(()) }