diff --git a/crates/buzz-acp/src/lib.rs b/crates/buzz-acp/src/lib.rs index 4425257f2..18edd6be4 100644 --- a/crates/buzz-acp/src/lib.rs +++ b/crates/buzz-acp/src/lib.rs @@ -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, +) -> Option> { + 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>>) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base_url = format!("http://{}", listener.local_addr().unwrap()); + let captured: Arc>> = 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::().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::(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>>, + n: usize, + ) -> Vec { + 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::*; diff --git a/crates/buzz-acp/src/pool.rs b/crates/buzz-acp/src/pool.rs index ab8db921d..76e0072df 100644 --- a/crates/buzz-acp/src/pool.rs +++ b/crates/buzz-acp/src/pool.rs @@ -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 { + 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 { + self.task_map + .values_mut() + .flat_map(|m| std::mem::take(&mut m.steered_event_ids)) + .collect() + } + pub fn result_tx(&self) -> mpsc::UnboundedSender { 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 = 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()); } }