mirror of
https://github.com/block/buzz.git
synced 2026-08-18 06:50:31 +02:00
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:
co-authored by
Tyler Longwell
parent
6d4700b673
commit
4719362107
+472
-30
@@ -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
@@ -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());
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user