parallelize graceful agent drains, serialize container stops on fast lane; unify shutdown+checkpoint+compact prompt
This commit is contained in:
parent
f6d9ca7f99
commit
fc42f97691
5 changed files with 189 additions and 60 deletions
|
|
@ -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<S: Surface>(
|
|||
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<S: Surface>(
|
|||
}
|
||||
},
|
||||
};
|
||||
let ctrl = handle_turn::<S>(socket, &bus, stats.as_ref(), files, &turn_lock, next).await;
|
||||
let ctrl = handle_turn::<S>(
|
||||
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<S: Surface>(
|
|||
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<S: Surface>(
|
|||
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);
|
||||
|
|
|
|||
|
|
@ -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<u64> {
|
||||
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
|
||||
|
|
|
|||
Loading…
Reference in a new issue