fix(buzz-acp): bind steered-👀 cleanup to exact turn; cover shutdown

Review (Wren) found two correctness gaps in the steered-reaction fix:

1. Generation-reuse race: SteerAckEvent carried only (channel_id,
   event_id), and record_steered_event matched TaskMeta by channel. If
   turn A completed and turn B started on the same channel before A's
   delayed SteerAck::Success arrived, A's steered 👀 attached to B and
   survived until the wrong turn ended. send_steer now returns the
   turn's tokio::task::Id; the ack watcher carries it, and
   record_steered_event matches on the exact task id — a stale id
   refuses to bind (returns false) and the ack arm removes the 👀
   immediately.

2. Graceful shutdown bypassed cleanup: the post-loop drain consumes
   PromptResults directly (never via handle_prompt_result) and aborts
   stragglers, so steered ids held in TaskMeta died with the pool and
   their 👀 went stale. tokio_main now drains all steered ids
   (pool.drain_all_steered_event_ids) as the main loop exits and
   bound-awaits a best-effort removal task before process exit.

The three duplicated cleanup blocks are unified in
spawn_steered_eyes_cleanup. New lifecycle tests observe actual wire
behavior against a relay stub — the kind:5 (NIP-09) deletion e-tagging
the 👀 reaction — for normal completion, panic recovery, the shutdown
drain composition, and late-ack immediate cleanup; the generation race
is pinned at the binding level in pool::tests.

Co-authored-by: Tyler Longwell <tlongwell@block.xyz>
Signed-off-by: Tyler Longwell <tlongwell@block.xyz>
This commit is contained in:
npub1qyvc0c5kl4gqv2fd97fsk46tu378sqgy35vc83rvgfwne90sel7s0ed67d
2026-07-03 11:10:33 -04:00
co-authored by Tyler Longwell
parent 6d4700b673
commit 4719362107
2 changed files with 584 additions and 81 deletions
+472 -30
View File
@@ -988,6 +988,13 @@ struct RespawnResult {
struct SteerAckEvent {
channel_id: Uuid,
event_id: String,
/// The exact turn (`tokio::task::Id`) the steer was sent to, from
/// `pool.send_steer`. Ack-driven bookkeeping must bind to this id, not
/// to "whatever turn currently owns `channel_id`": if the turn ended
/// and a fresh one started on the same channel before a delayed ack
/// arrived, channel-matching would attach the steered event's 👀 to
/// the successor turn — keeping it alive until the wrong turn stopped.
task_id: tokio::task::Id,
/// `Ok` if the read loop sent any of the locked `SteerAck` variants.
/// `Err` if the oneshot was dropped without a send — should not happen
/// under the current read-loop drains, but if it ever does the main
@@ -2163,6 +2170,7 @@ async fn tokio_main() -> Result<()> {
Some(PoolEvent::SteerAck(SteerAckEvent {
channel_id,
event_id,
task_id,
ack,
})) => {
// Goose-native steer attempt resolved. Locked semantics
@@ -2256,16 +2264,14 @@ async fn tokio_main() -> Result<()> {
queue.remove_event(channel_id, &event_id);
// The event was absorbed into the in-flight turn without
// ever entering a FlushBatch, so no ReactionGuard owns
// its 👀 (added at queue-push time). Attach it to the
// turn's TaskMeta so the cleanup fires when that turn
// its 👀 (added at queue-push time). Attach it to that
// exact turn's TaskMeta (matched by task id, not channel
// — a delayed ack must not bind to a successor turn on
// the same channel) so the cleanup fires when the turn
// ends. If the turn already ended (result raced ahead
// of the ack), clean up immediately.
if !pool.record_steered_event(channel_id, &event_id) {
let rc = ctx.rest_client.clone();
let eid = event_id.clone();
tokio::spawn(async move {
pool::reaction_remove(&rc, &eid, "👀").await;
});
if !pool.record_steered_event(task_id, &event_id) {
spawn_steered_eyes_cleanup(Some(&ctx.rest_client), vec![event_id.clone()]);
}
}
if release_withheld {
@@ -2295,6 +2301,17 @@ async fn tokio_main() -> Result<()> {
}
tracing::info!("shutdown: waiting for in-flight prompts");
// Steered-👀 cleanup would otherwise be lost here: the drain below
// consumes PromptResults directly (never via handle_prompt_result), and
// tasks that outlive the grace period are aborted — either way every
// TaskMeta drops with the pool and its steered_event_ids with it. All
// in-flight turns are stopping now, so the user-visible contract (eyes
// clear when the turn stops) says remove them. No new steer acks can be
// recorded after the main loop exits, so draining once here is complete.
// Runs concurrently with the grace-period drain; awaited (bounded, each
// reaction_remove is internally timeout-capped) before exit.
let steered_cleanup =
spawn_steered_eyes_cleanup(Some(&ctx.rest_client), pool.drain_all_steered_event_ids());
// 30 s is generous for in-flight prompts to be cancelled; using
// max_turn_duration here would cause Ctrl+C to hang for up to an hour.
let grace = Duration::from_secs(30);
@@ -2348,6 +2365,19 @@ async fn tokio_main() -> Result<()> {
}
drop(pool);
// Wait (bounded) for the steered-👀 cleanup spawned above — each
// reaction_remove is internally capped at ~2 s of HTTP, so this cannot
// hang shutdown meaningfully. Best-effort: on timeout the reactions
// stay stale, same class of loss as any other kill-during-cleanup.
if let Some(handle) = steered_cleanup {
if tokio::time::timeout(Duration::from_secs(10), handle)
.await
.is_err()
{
tracing::warn!("steered 👀 cleanup did not finish before shutdown deadline");
}
}
// Abort any in-flight respawn tasks. They may be sleeping in backoff or
// running spawn_and_init — either way, we don't want them spawning new
// children after the main loop has exited. RespawnGuard::Drop sends a
@@ -2530,7 +2560,7 @@ fn try_native_steer(
};
match pool.send_steer(channel_id, request) {
Ok(()) => {
Ok(task_id) => {
// Withhold the queued event synchronously BEFORE spawning
// the watcher: this closes the race where `mark_complete`
// clears `in_flight_channels` and a stray `flush_next` could
@@ -2559,6 +2589,7 @@ fn try_native_steer(
let _ = ack_tx_clone.send(SteerAckEvent {
channel_id,
event_id: event_id_for_watcher,
task_id,
ack,
});
});
@@ -2667,6 +2698,30 @@ fn dispatch_pending(
dispatched_channels
}
/// Spawn a best-effort background removal of the 👀 reaction from events
/// that were steered into a turn (`SteerAck::Success`). Steered events never
/// enter a `FlushBatch`, so `run_prompt_task`'s `ReactionGuard` never owns
/// their 👀 — every turn-stopping path (normal completion, panic recovery,
/// graceful shutdown) and the late-ack path clean up through here instead.
///
/// Returns the `JoinHandle` when a task was spawned (non-empty ids and a
/// rest client available) so shutdown can bound-await it; other callers
/// discard the handle.
fn spawn_steered_eyes_cleanup(
rest_client: Option<&relay::RestClient>,
ids: Vec<String>,
) -> Option<tokio::task::JoinHandle<()>> {
if ids.is_empty() {
return None;
}
let rest = rest_client?.clone();
Some(tokio::spawn(async move {
for eid in &ids {
pool::reaction_remove(&rest, eid, "👀").await;
}
}))
}
#[allow(clippy::too_many_arguments)]
fn handle_prompt_result(
pool: &mut AgentPool,
@@ -2698,16 +2753,7 @@ fn handle_prompt_result(
}
});
debug_assert_eq!(before, pool.task_map().len() + 1);
if !steered_ids.is_empty() {
if let Some(rest) = rest_client {
let rest = rest.clone();
tokio::spawn(async move {
for eid in &steered_ids {
pool::reaction_remove(&rest, eid, "👀").await;
}
});
}
}
spawn_steered_eyes_cleanup(rest_client, steered_ids);
// Requeue BEFORE mark_complete: requeue() sets retry_after with a future
// deadline, and mark_complete() checks for it to decide whether to preserve
@@ -2938,17 +2984,7 @@ fn recover_panicked_agent(
// Steered events never entered a FlushBatch, so the panicked task's
// ReactionGuard never owned their 👀 — clean them up here (same
// contract as handle_prompt_result: eyes clear when the turn ends).
if !meta.steered_event_ids.is_empty() {
if let Some(rest) = rest_client {
let rest = rest.clone();
let ids = meta.steered_event_ids.clone();
tokio::spawn(async move {
for eid in &ids {
pool::reaction_remove(&rest, eid, "👀").await;
}
});
}
}
spawn_steered_eyes_cleanup(rest_client, meta.steered_event_ids);
// Requeue BEFORE mark_complete (same rationale as handle_prompt_result).
if let Some(batch) = meta.recoverable_batch {
@@ -4139,6 +4175,412 @@ mod error_outcome_emission_tests {
}
}
#[cfg(test)]
mod steered_eyes_lifecycle_tests {
//! Lifecycle-level regression tests for steered-👀 cleanup.
//!
//! Events delivered mid-turn via the goose-native steer
//! (`SteerAck::Success`) never enter a `FlushBatch`, so
//! `run_prompt_task`'s `ReactionGuard` never owns their 👀. These tests
//! pin that every turn-stopping path actually removes the reaction *on
//! the wire* — a relay stub answers the `POST /query` reaction lookup
//! and captures the signed kind:5 (NIP-09) deletion submitted to
//! `POST /events` — not just that ids are bookkept in `TaskMeta`:
//!
//! - normal completion (`handle_prompt_result`)
//! - panic recovery (`recover_panicked_agent`)
//! - graceful shutdown (the `drain_all_steered_event_ids` +
//! `spawn_steered_eyes_cleanup` composition `tokio_main` runs)
//! - late ack after turn end (immediate cleanup, no turn to attach to)
//!
//! The generation-reuse race (a late ack must not bind to a successor
//! turn on the same channel) is pinned at the binding level in
//! `pool::tests::test_late_ack_for_ended_turn_does_not_bind_to_successor_turn`.
use super::*;
use crate::acp::{AcpClient, StopReason};
use crate::pool::{AgentPool, OwnedAgent, PromptOutcome, PromptResult, PromptSource, TaskMeta};
use std::collections::HashSet;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpListener;
use tokio::sync::Mutex;
/// Minimal relay stub. Answers `POST /query` (the reaction lookup inside
/// `pool::reaction_remove`) with a single 👀 reaction event of id
/// `reaction_id`, and captures every body submitted to `POST /events`
/// (the kind:5 deletions under test). One request per connection.
async fn spawn_mock_relay(
reaction_id: &str,
) -> (relay::RestClient, Arc<Mutex<Vec<serde_json::Value>>>) {
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let base_url = format!("http://{}", listener.local_addr().unwrap());
let captured: Arc<Mutex<Vec<serde_json::Value>>> = Arc::new(Mutex::new(Vec::new()));
let captured_srv = captured.clone();
let reaction_id = reaction_id.to_string();
tokio::spawn(async move {
loop {
let (mut sock, _) = match listener.accept().await {
Ok(p) => p,
Err(_) => return,
};
let captured = captured_srv.clone();
let reaction_id = reaction_id.clone();
tokio::spawn(async move {
let mut buf = Vec::new();
let mut tmp = [0u8; 4096];
while !buf.windows(4).any(|w| w == b"\r\n\r\n") {
match sock.read(&mut tmp).await {
Ok(0) | Err(_) => return,
Ok(n) => buf.extend_from_slice(&tmp[..n]),
}
if buf.len() > 1_000_000 {
return;
}
}
let head_end = buf.windows(4).position(|w| w == b"\r\n\r\n").unwrap() + 4;
let head = String::from_utf8_lossy(&buf[..head_end]).to_string();
let content_length = head
.lines()
.find_map(|l| {
let (k, v) = l.split_once(':')?;
if k.eq_ignore_ascii_case("content-length") {
v.trim().parse::<usize>().ok()
} else {
None
}
})
.unwrap_or(0);
while buf.len() < head_end + content_length {
match sock.read(&mut tmp).await {
Ok(0) | Err(_) => return,
Ok(n) => buf.extend_from_slice(&tmp[..n]),
}
}
let body = &buf[head_end..head_end + content_length];
let path = head
.lines()
.next()
.and_then(|l| l.split_whitespace().nth(1))
.unwrap_or("")
.to_string();
let response_body = if path == "/query" {
serde_json::json!([{ "id": reaction_id, "content": "👀" }]).to_string()
} else {
// `/events` — capture the submitted deletion event.
if let Ok(v) = serde_json::from_slice::<serde_json::Value>(body) {
captured.lock().await.push(v);
}
"{}".to_string()
};
let resp = format!(
"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
response_body.len(),
response_body,
);
let _ = sock.write_all(resp.as_bytes()).await;
let _ = sock.shutdown().await;
});
}
});
let rest = relay::RestClient {
http: reqwest::Client::new(),
base_url,
keys: nostr::Keys::generate(),
auth_tag_json: None,
};
(rest, captured)
}
/// Poll until `captured` holds at least `n` deletion events or ~3 s pass
/// (the cleanup runs on a fire-and-forget spawned task).
async fn wait_for_deletions(
captured: &Arc<Mutex<Vec<serde_json::Value>>>,
n: usize,
) -> Vec<serde_json::Value> {
for _ in 0..300 {
let got = captured.lock().await.clone();
if got.len() >= n {
return got;
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
captured.lock().await.clone()
}
/// The captured event must be a kind:5 (NIP-09) deletion e-tagging the
/// 👀 reaction event the stub advertised.
fn assert_is_deletion_of(event: &serde_json::Value, reaction_id: &str) {
assert_eq!(
event.get("kind").and_then(|k| k.as_u64()),
Some(5),
"must be a kind:5 deletion, got: {event}"
);
let tags = event
.get("tags")
.and_then(|t| t.as_array())
.cloned()
.unwrap_or_default();
assert!(
tags.iter().any(|t| t.as_array().is_some_and(|t| {
t.first().and_then(|v| v.as_str()) == Some("e")
&& t.get(1).and_then(|v| v.as_str()) == Some(reaction_id)
})),
"deletion must e-tag reaction event {reaction_id}, got tags: {tags:?}"
);
}
fn test_config() -> Config {
Config {
keys: nostr::Keys::generate(),
relay_url: "ws://localhost:3000".into(),
// `true` exits cleanly, so the async respawn on the panic path
// fails fast and harmlessly off the JoinSet.
agent_command: "true".into(),
agent_args: vec![],
mcp_command: "test-mcp-server".into(),
idle_timeout_secs: config::DEFAULT_IDLE_TIMEOUT_SECS,
max_turn_duration_secs: 3600,
agents: 1,
heartbeat_interval_secs: 0,
turn_liveness_secs: 10,
heartbeat_prompt: None,
system_prompt: None,
initial_message: None,
subscribe_mode: config::SubscribeMode::All,
dedup_mode: config::DedupMode::Queue,
multiple_event_handling: config::MultipleEventHandling::Queue,
ignore_self: true,
kinds_override: None,
channels_override: None,
no_mention_filter: false,
config_path: std::path::PathBuf::from("./buzz-acp.toml"),
context_message_limit: 12,
max_turns_per_session: 0,
presence_enabled: true,
typing_enabled: true,
memory_enabled: false,
model: None,
permission_mode: config::PermissionMode::BypassPermissions,
respond_to: config::RespondTo::Anyone,
respond_to_allowlist: HashSet::new(),
allowed_respond_to: vec![],
persona_env_vars: vec![],
relay_observer: false,
agent_owner: None,
no_base_prompt: false,
base_prompt_content: None,
}
}
/// Real but inert agent subprocess (`cat`) — the paths under test never
/// talk to it. Same pattern as `error_outcome_emission_tests`.
async fn dummy_agent(index: usize) -> OwnedAgent {
OwnedAgent {
index,
acp: AcpClient::spawn("cat", &[], &[])
.await
.expect("spawn cat as inert agent"),
state: Default::default(),
model_capabilities: None,
desired_model: None,
model_overridden: false,
protocol_version: 1,
}
}
fn fresh_circuit() -> SlotCircuit {
SlotCircuit {
crash_times: Vec::new(),
open_until: None,
respawn_in_flight: false,
}
}
/// Normal completion: a turn that absorbed a steered event ends via
/// `handle_prompt_result` → the steered event's 👀 is removed on the
/// wire (kind:5 deletion submitted to the relay).
#[tokio::test]
async fn completion_removes_steered_eyes_on_the_wire() {
let reaction_id = "ab".repeat(32);
let (rest, captured) = spawn_mock_relay(&reaction_id).await;
let mut pool = AgentPool::from_slots(vec![None]);
let channel_id = Uuid::new_v4();
let task_id = pool.join_set.spawn(async {}).id();
pool.task_map_mut().insert(
task_id,
TaskMeta {
agent_index: 0,
channel_id: Some(channel_id),
recoverable_batch: None,
control_tx: None,
steer_tx: None,
steered_event_ids: Vec::new(),
},
);
assert!(pool.record_steered_event(task_id, &"ef".repeat(32)));
let mut queue = EventQueue::new(config::DedupMode::Queue);
let config = test_config();
let mut heartbeat_in_flight = false;
let removed_channels = HashSet::new();
let mut crash_history = vec![fresh_circuit()];
let (respawn_tx, _respawn_rx) = mpsc::channel(8);
let mut respawn_tasks = tokio::task::JoinSet::new();
handle_prompt_result(
&mut pool,
&mut queue,
&config,
PromptResult {
agent: dummy_agent(0).await,
source: PromptSource::Channel(channel_id),
outcome: PromptOutcome::Ok(StopReason::EndTurn),
batch: None,
},
&mut heartbeat_in_flight,
&removed_channels,
&mut crash_history,
&respawn_tx,
&mut respawn_tasks,
None,
Some(&rest),
);
let deletions = wait_for_deletions(&captured, 1).await;
assert_eq!(
deletions.len(),
1,
"exactly one deletion for one steered event"
);
assert_is_deletion_of(&deletions[0], &reaction_id);
}
/// Panic: a turn that absorbed a steered event panics; recovery removes
/// the 👀 on the wire — same contract as normal completion.
#[tokio::test]
async fn panic_recovery_removes_steered_eyes_on_the_wire() {
let reaction_id = "cd".repeat(32);
let (rest, captured) = spawn_mock_relay(&reaction_id).await;
let mut pool = AgentPool::from_slots(vec![None]);
let channel_id = Uuid::new_v4();
let task_id = pool
.join_set
.spawn(async { panic!("turn panicked mid-steer") })
.id();
pool.task_map_mut().insert(
task_id,
TaskMeta {
agent_index: 0,
channel_id: Some(channel_id),
recoverable_batch: None,
control_tx: None,
steer_tx: None,
steered_event_ids: vec!["12".repeat(32)],
},
);
let join_error = match pool.join_set.join_next().await {
Some(Err(e)) => e,
other => panic!("expected panicked task, got {other:?}"),
};
assert_eq!(join_error.id(), task_id);
let mut queue = EventQueue::new(config::DedupMode::Queue);
let config = test_config();
let mut heartbeat_in_flight = false;
let removed_channels = HashSet::new();
let mut typing_channels = HashMap::new();
let mut crash_history = vec![fresh_circuit()];
let (respawn_tx, _respawn_rx) = mpsc::channel(8);
let mut respawn_tasks = tokio::task::JoinSet::new();
recover_panicked_agent(
&mut pool,
&mut queue,
&config,
join_error,
&mut heartbeat_in_flight,
&removed_channels,
&mut typing_channels,
&mut crash_history,
&respawn_tx,
&mut respawn_tasks,
None,
Some(&rest),
);
let deletions = wait_for_deletions(&captured, 1).await;
assert_eq!(deletions.len(), 1);
assert_is_deletion_of(&deletions[0], &reaction_id);
}
/// Graceful shutdown: the exact composition `tokio_main` runs after the
/// main loop exits — drain every steered id from every in-flight turn,
/// then remove them on the wire — clears all steered 👀 even though no
/// `handle_prompt_result` ever ran for those turns.
#[tokio::test]
async fn shutdown_drain_removes_steered_eyes_on_the_wire() {
let reaction_id = "ee".repeat(32);
let (rest, captured) = spawn_mock_relay(&reaction_id).await;
let mut pool = AgentPool::from_slots(vec![]);
for i in 0..2 {
let task_id = pool.join_set.spawn(std::future::pending()).id();
pool.task_map_mut().insert(
task_id,
TaskMeta {
agent_index: i,
channel_id: Some(Uuid::new_v4()),
recoverable_batch: None,
control_tx: None,
steer_tx: None,
steered_event_ids: Vec::new(),
},
);
assert!(pool.record_steered_event(task_id, &format!("{i}{i}").repeat(32)));
}
let handle = spawn_steered_eyes_cleanup(Some(&rest), pool.drain_all_steered_event_ids())
.expect("two steered ids must spawn a cleanup task");
handle.await.expect("cleanup task must not panic");
let deletions = captured.lock().await.clone();
assert_eq!(deletions.len(), 2, "one deletion per steered event");
for d in &deletions {
assert_is_deletion_of(d, &reaction_id);
}
pool.join_set.shutdown().await;
}
/// Late ack after turn end: when `record_steered_event` refuses the
/// stale task id (see the pool generation-race test), the SteerAck arm
/// falls back to immediate cleanup — which must remove the 👀 on the
/// wire right away rather than waiting on any turn.
#[tokio::test]
async fn late_ack_immediate_cleanup_removes_eyes_on_the_wire() {
let reaction_id = "0f".repeat(32);
let (rest, captured) = spawn_mock_relay(&reaction_id).await;
let handle = spawn_steered_eyes_cleanup(Some(&rest), vec!["34".repeat(32)])
.expect("one id must spawn a cleanup task");
handle.await.expect("cleanup task must not panic");
let deletions = captured.lock().await.clone();
assert_eq!(deletions.len(), 1);
assert_is_deletion_of(&deletions[0], &reaction_id);
}
/// No ids → no task, no HTTP. Pins the no-op path shutdown relies on.
#[tokio::test]
async fn empty_cleanup_spawns_nothing() {
let (rest, captured) = spawn_mock_relay(&"aa".repeat(32)).await;
assert!(spawn_steered_eyes_cleanup(Some(&rest), Vec::new()).is_none());
tokio::time::sleep(Duration::from_millis(50)).await;
assert!(captured.lock().await.is_empty());
}
}
#[cfg(test)]
mod observer_payload_trim_tests {
use super::*;
+112 -51
View File
@@ -513,36 +513,43 @@ impl AgentPool {
/// and this call, or the channel was never in flight). This is
/// semantically a soft no-op — the caller should release any withheld
/// event and let normal dispatch handle delivery.
///
/// On success returns the `tokio::task::Id` of the turn the steer was
/// sent to. The caller must thread this id through the ack watcher so
/// that ack-driven bookkeeping (`record_steered_event`) binds to this
/// exact turn — a late ack matched by channel alone could attach to a
/// *successor* turn on the same channel.
pub fn send_steer(
&mut self,
channel_id: Uuid,
request: SteerRequest,
) -> Result<(), SteerError> {
let meta = self
) -> Result<tokio::task::Id, SteerError> {
let (task_id, meta) = self
.task_map
.values_mut()
.find(|m| m.channel_id == Some(channel_id))
.iter_mut()
.find(|(_, m)| m.channel_id == Some(channel_id))
.ok_or(SteerError::PromptCompleted)?;
let tx = meta
.steer_tx
.as_ref()
.ok_or_else(|| SteerError::Transport("steer_tx not installed".into()))?;
tx.try_send(request)
.map_err(|e| SteerError::Transport(e.to_string()))
.map_err(|e| SteerError::Transport(e.to_string()))?;
Ok(*task_id)
}
/// Record an event id as steered into the in-flight turn for
/// `channel_id`, so its 👀 is cleared when that turn ends (the event
/// never enters a `FlushBatch`, so `run_prompt_task`'s `ReactionGuard`
/// cannot own it). Returns `false` if no task is in flight for the
/// channel — the turn already ended and the caller must clean up the
/// reaction immediately instead.
pub fn record_steered_event(&mut self, channel_id: Uuid, event_id: &str) -> bool {
match self
.task_map
.values_mut()
.find(|m| m.channel_id == Some(channel_id))
{
/// Record an event id as steered into the turn identified by `task_id`,
/// so its 👀 is cleared when that turn ends (the event never enters a
/// `FlushBatch`, so `run_prompt_task`'s `ReactionGuard` cannot own it).
///
/// Matches on the exact task id — never on channel — so a delayed
/// `SteerAck::Success` for a turn that already ended cannot attach its
/// event to a successor turn on the same channel (which would keep the
/// 👀 alive until the *wrong* turn stopped). Returns `false` if that
/// task is no longer in flight — the turn ended and the caller must
/// clean up the reaction immediately instead.
pub fn record_steered_event(&mut self, task_id: tokio::task::Id, event_id: &str) -> bool {
match self.task_map.get_mut(&task_id) {
Some(meta) => {
meta.steered_event_ids.push(event_id.to_string());
true
@@ -551,6 +558,21 @@ impl AgentPool {
}
}
/// Take every steered event id from every in-flight `TaskMeta`.
///
/// Shutdown-only: the graceful-shutdown drain consumes `PromptResult`s
/// directly (never through `handle_prompt_result`) and aborts whatever
/// outlives the grace period, so `TaskMeta`s are dropped with the pool
/// without their steered-👀 cleanup ever firing. Callers drain here and
/// remove the reactions before the process exits — every in-flight turn
/// is, by definition, stopping.
pub fn drain_all_steered_event_ids(&mut self) -> Vec<String> {
self.task_map
.values_mut()
.flat_map(|m| std::mem::take(&mut m.steered_event_ids))
.collect()
}
pub fn result_tx(&self) -> mpsc::UnboundedSender<PromptResult> {
self.result_tx.clone()
}
@@ -3686,23 +3708,20 @@ mod tests {
// Reaching here without a panic is the test.
}
// ── record_steered_event ──────────────────────────────────────────────
// ── record_steered_event / drain_all_steered_event_ids ───────────────
//
// Pins the stale-👀 fix: events delivered via SteerAck::Success never
// enter a FlushBatch, so their 👀 must be attached to the in-flight
// turn's TaskMeta (cleared on turn end) — or, if no turn is in flight,
// the caller must clean up immediately (record returns false).
/// Recording against an in-flight channel stores the id in that turn's
/// TaskMeta; ids accumulate across multiple steers into the same turn.
#[tokio::test]
async fn test_record_steered_event_attaches_to_in_flight_turn() {
let mut pool = AgentPool::from_slots(vec![]);
let channel_id = Uuid::new_v4();
// turn's TaskMeta (cleared on turn end) — matched by exact task id, so
// a delayed ack can never bind to a successor turn on the same channel.
// If the turn is gone, record returns false and the caller must clean
// up immediately.
fn insert_meta(pool: &mut AgentPool, channel_id: Uuid) -> tokio::task::Id {
let abort_handle = pool.join_set.spawn(async {});
let task_id = abort_handle.id();
pool.task_map_mut().insert(
abort_handle.id(),
task_id,
TaskMeta {
agent_index: 0,
channel_id: Some(channel_id),
@@ -3712,44 +3731,86 @@ mod tests {
steered_event_ids: Vec::new(),
},
);
task_id
}
assert!(pool.record_steered_event(channel_id, "aaa"));
assert!(pool.record_steered_event(channel_id, "bbb"));
/// Recording against an in-flight task stores the id in that turn's
/// TaskMeta; ids accumulate across multiple steers into the same turn.
#[tokio::test]
async fn test_record_steered_event_attaches_to_in_flight_turn() {
let mut pool = AgentPool::from_slots(vec![]);
let channel_id = Uuid::new_v4();
let task_id = insert_meta(&mut pool, channel_id);
assert!(pool.record_steered_event(task_id, "aaa"));
assert!(pool.record_steered_event(task_id, "bbb"));
let ids: Vec<String> = pool
.task_map()
.values()
.find(|m| m.channel_id == Some(channel_id))
.get(&task_id)
.expect("task meta must exist")
.steered_event_ids
.clone();
assert_eq!(ids, vec!["aaa".to_string(), "bbb".to_string()]);
}
/// Recording against a channel with no in-flight turn returns false so
/// the caller cleans up the 👀 immediately.
/// Recording against a task that is no longer in flight returns false
/// so the caller cleans up the 👀 immediately.
#[tokio::test]
async fn test_record_steered_event_returns_false_when_no_turn_in_flight() {
let mut pool = AgentPool::from_slots(vec![]);
let in_flight = Uuid::new_v4();
let other = Uuid::new_v4();
let channel_id = Uuid::new_v4();
let task_id = insert_meta(&mut pool, channel_id);
pool.task_map_mut().remove(&task_id);
let abort_handle = pool.join_set.spawn(async {});
pool.task_map_mut().insert(
abort_handle.id(),
TaskMeta {
agent_index: 0,
channel_id: Some(in_flight),
recoverable_batch: None,
control_tx: None,
steer_tx: None,
steered_event_ids: Vec::new(),
},
assert!(!pool.record_steered_event(task_id, "ccc"));
}
/// Generation-reuse race: turn A ends, turn B starts on the SAME
/// channel, then A's delayed SteerAck::Success arrives. The stale
/// task id must NOT attach to B — record returns false (caller
/// removes the 👀 immediately) and B's meta stays untouched.
#[tokio::test]
async fn test_late_ack_for_ended_turn_does_not_bind_to_successor_turn() {
let mut pool = AgentPool::from_slots(vec![]);
let channel_id = Uuid::new_v4();
// Turn A: steer accepted, then the turn completes (meta removed).
let task_a = insert_meta(&mut pool, channel_id);
pool.task_map_mut().remove(&task_a);
// Turn B: a fresh event starts a new turn on the same channel.
let task_b = insert_meta(&mut pool, channel_id);
// A's delayed ack arrives carrying A's task id.
assert!(
!pool.record_steered_event(task_a, "stale"),
"a late ack for an ended turn must demand immediate cleanup"
);
assert!(
pool.task_map()
.get(&task_b)
.expect("turn B in flight")
.steered_event_ids
.is_empty(),
"turn B must not inherit turn A's steered event"
);
}
assert!(!pool.record_steered_event(other, "ccc"));
// Empty pool: no task meta at all.
let mut empty = AgentPool::from_slots(vec![]);
assert!(!empty.record_steered_event(in_flight, "ddd"));
/// Shutdown drain: every steered id is taken from every in-flight
/// TaskMeta exactly once, leaving the metas empty.
#[tokio::test]
async fn test_drain_all_steered_event_ids_takes_everything_once() {
let mut pool = AgentPool::from_slots(vec![]);
let t1 = insert_meta(&mut pool, Uuid::new_v4());
let t2 = insert_meta(&mut pool, Uuid::new_v4());
assert!(pool.record_steered_event(t1, "e1"));
assert!(pool.record_steered_event(t2, "e2"));
assert!(pool.record_steered_event(t2, "e3"));
let mut drained = pool.drain_all_steered_event_ids();
drained.sort();
assert_eq!(drained, vec!["e1", "e2", "e3"]);
assert!(pool.drain_all_steered_event_ids().is_empty());
}
}