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 {
|
struct SteerAckEvent {
|
||||||
channel_id: Uuid,
|
channel_id: Uuid,
|
||||||
event_id: String,
|
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.
|
/// `Ok` if the read loop sent any of the locked `SteerAck` variants.
|
||||||
/// `Err` if the oneshot was dropped without a send — should not happen
|
/// `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
|
/// 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 {
|
Some(PoolEvent::SteerAck(SteerAckEvent {
|
||||||
channel_id,
|
channel_id,
|
||||||
event_id,
|
event_id,
|
||||||
|
task_id,
|
||||||
ack,
|
ack,
|
||||||
})) => {
|
})) => {
|
||||||
// Goose-native steer attempt resolved. Locked semantics
|
// Goose-native steer attempt resolved. Locked semantics
|
||||||
@@ -2256,16 +2264,14 @@ async fn tokio_main() -> Result<()> {
|
|||||||
queue.remove_event(channel_id, &event_id);
|
queue.remove_event(channel_id, &event_id);
|
||||||
// The event was absorbed into the in-flight turn without
|
// The event was absorbed into the in-flight turn without
|
||||||
// ever entering a FlushBatch, so no ReactionGuard owns
|
// ever entering a FlushBatch, so no ReactionGuard owns
|
||||||
// its 👀 (added at queue-push time). Attach it to the
|
// its 👀 (added at queue-push time). Attach it to that
|
||||||
// turn's TaskMeta so the cleanup fires when that turn
|
// 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
|
// ends. If the turn already ended (result raced ahead
|
||||||
// of the ack), clean up immediately.
|
// of the ack), clean up immediately.
|
||||||
if !pool.record_steered_event(channel_id, &event_id) {
|
if !pool.record_steered_event(task_id, &event_id) {
|
||||||
let rc = ctx.rest_client.clone();
|
spawn_steered_eyes_cleanup(Some(&ctx.rest_client), vec![event_id.clone()]);
|
||||||
let eid = event_id.clone();
|
|
||||||
tokio::spawn(async move {
|
|
||||||
pool::reaction_remove(&rc, &eid, "👀").await;
|
|
||||||
});
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
if release_withheld {
|
if release_withheld {
|
||||||
@@ -2295,6 +2301,17 @@ async fn tokio_main() -> Result<()> {
|
|||||||
}
|
}
|
||||||
|
|
||||||
tracing::info!("shutdown: waiting for in-flight prompts");
|
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
|
// 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.
|
// max_turn_duration here would cause Ctrl+C to hang for up to an hour.
|
||||||
let grace = Duration::from_secs(30);
|
let grace = Duration::from_secs(30);
|
||||||
@@ -2348,6 +2365,19 @@ async fn tokio_main() -> Result<()> {
|
|||||||
}
|
}
|
||||||
drop(pool);
|
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
|
// 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
|
// running spawn_and_init — either way, we don't want them spawning new
|
||||||
// children after the main loop has exited. RespawnGuard::Drop sends a
|
// 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) {
|
match pool.send_steer(channel_id, request) {
|
||||||
Ok(()) => {
|
Ok(task_id) => {
|
||||||
// Withhold the queued event synchronously BEFORE spawning
|
// Withhold the queued event synchronously BEFORE spawning
|
||||||
// the watcher: this closes the race where `mark_complete`
|
// the watcher: this closes the race where `mark_complete`
|
||||||
// clears `in_flight_channels` and a stray `flush_next` could
|
// clears `in_flight_channels` and a stray `flush_next` could
|
||||||
@@ -2559,6 +2589,7 @@ fn try_native_steer(
|
|||||||
let _ = ack_tx_clone.send(SteerAckEvent {
|
let _ = ack_tx_clone.send(SteerAckEvent {
|
||||||
channel_id,
|
channel_id,
|
||||||
event_id: event_id_for_watcher,
|
event_id: event_id_for_watcher,
|
||||||
|
task_id,
|
||||||
ack,
|
ack,
|
||||||
});
|
});
|
||||||
});
|
});
|
||||||
@@ -2667,6 +2698,30 @@ fn dispatch_pending(
|
|||||||
dispatched_channels
|
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)]
|
#[allow(clippy::too_many_arguments)]
|
||||||
fn handle_prompt_result(
|
fn handle_prompt_result(
|
||||||
pool: &mut AgentPool,
|
pool: &mut AgentPool,
|
||||||
@@ -2698,16 +2753,7 @@ fn handle_prompt_result(
|
|||||||
}
|
}
|
||||||
});
|
});
|
||||||
debug_assert_eq!(before, pool.task_map().len() + 1);
|
debug_assert_eq!(before, pool.task_map().len() + 1);
|
||||||
if !steered_ids.is_empty() {
|
spawn_steered_eyes_cleanup(rest_client, steered_ids);
|
||||||
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;
|
|
||||||
}
|
|
||||||
});
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// Requeue BEFORE mark_complete: requeue() sets retry_after with a future
|
// Requeue BEFORE mark_complete: requeue() sets retry_after with a future
|
||||||
// deadline, and mark_complete() checks for it to decide whether to preserve
|
// 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
|
// Steered events never entered a FlushBatch, so the panicked task's
|
||||||
// ReactionGuard never owned their 👀 — clean them up here (same
|
// ReactionGuard never owned their 👀 — clean them up here (same
|
||||||
// contract as handle_prompt_result: eyes clear when the turn ends).
|
// contract as handle_prompt_result: eyes clear when the turn ends).
|
||||||
if !meta.steered_event_ids.is_empty() {
|
spawn_steered_eyes_cleanup(rest_client, meta.steered_event_ids);
|
||||||
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;
|
|
||||||
}
|
|
||||||
});
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// Requeue BEFORE mark_complete (same rationale as handle_prompt_result).
|
// Requeue BEFORE mark_complete (same rationale as handle_prompt_result).
|
||||||
if let Some(batch) = meta.recoverable_batch {
|
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)]
|
#[cfg(test)]
|
||||||
mod observer_payload_trim_tests {
|
mod observer_payload_trim_tests {
|
||||||
use super::*;
|
use super::*;
|
||||||
|
|||||||
+112
-51
@@ -513,36 +513,43 @@ impl AgentPool {
|
|||||||
/// and this call, or the channel was never in flight). This is
|
/// and this call, or the channel was never in flight). This is
|
||||||
/// semantically a soft no-op — the caller should release any withheld
|
/// semantically a soft no-op — the caller should release any withheld
|
||||||
/// event and let normal dispatch handle delivery.
|
/// 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(
|
pub fn send_steer(
|
||||||
&mut self,
|
&mut self,
|
||||||
channel_id: Uuid,
|
channel_id: Uuid,
|
||||||
request: SteerRequest,
|
request: SteerRequest,
|
||||||
) -> Result<(), SteerError> {
|
) -> Result<tokio::task::Id, SteerError> {
|
||||||
let meta = self
|
let (task_id, meta) = self
|
||||||
.task_map
|
.task_map
|
||||||
.values_mut()
|
.iter_mut()
|
||||||
.find(|m| m.channel_id == Some(channel_id))
|
.find(|(_, m)| m.channel_id == Some(channel_id))
|
||||||
.ok_or(SteerError::PromptCompleted)?;
|
.ok_or(SteerError::PromptCompleted)?;
|
||||||
let tx = meta
|
let tx = meta
|
||||||
.steer_tx
|
.steer_tx
|
||||||
.as_ref()
|
.as_ref()
|
||||||
.ok_or_else(|| SteerError::Transport("steer_tx not installed".into()))?;
|
.ok_or_else(|| SteerError::Transport("steer_tx not installed".into()))?;
|
||||||
tx.try_send(request)
|
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
|
/// Record an event id as steered into the turn identified by `task_id`,
|
||||||
/// `channel_id`, so its 👀 is cleared when that turn ends (the event
|
/// so its 👀 is cleared when that turn ends (the event never enters a
|
||||||
/// never enters a `FlushBatch`, so `run_prompt_task`'s `ReactionGuard`
|
/// `FlushBatch`, so `run_prompt_task`'s `ReactionGuard` cannot own it).
|
||||||
/// 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
|
/// Matches on the exact task id — never on channel — so a delayed
|
||||||
/// reaction immediately instead.
|
/// `SteerAck::Success` for a turn that already ended cannot attach its
|
||||||
pub fn record_steered_event(&mut self, channel_id: Uuid, event_id: &str) -> bool {
|
/// event to a successor turn on the same channel (which would keep the
|
||||||
match self
|
/// 👀 alive until the *wrong* turn stopped). Returns `false` if that
|
||||||
.task_map
|
/// task is no longer in flight — the turn ended and the caller must
|
||||||
.values_mut()
|
/// clean up the reaction immediately instead.
|
||||||
.find(|m| m.channel_id == Some(channel_id))
|
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) => {
|
Some(meta) => {
|
||||||
meta.steered_event_ids.push(event_id.to_string());
|
meta.steered_event_ids.push(event_id.to_string());
|
||||||
true
|
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> {
|
pub fn result_tx(&self) -> mpsc::UnboundedSender<PromptResult> {
|
||||||
self.result_tx.clone()
|
self.result_tx.clone()
|
||||||
}
|
}
|
||||||
@@ -3686,23 +3708,20 @@ mod tests {
|
|||||||
// Reaching here without a panic is the test.
|
// 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
|
// Pins the stale-👀 fix: events delivered via SteerAck::Success never
|
||||||
// enter a FlushBatch, so their 👀 must be attached to the in-flight
|
// 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,
|
// turn's TaskMeta (cleared on turn end) — matched by exact task id, so
|
||||||
// the caller must clean up immediately (record returns false).
|
// 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
|
||||||
/// Recording against an in-flight channel stores the id in that turn's
|
// up immediately.
|
||||||
/// 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();
|
|
||||||
|
|
||||||
|
fn insert_meta(pool: &mut AgentPool, channel_id: Uuid) -> tokio::task::Id {
|
||||||
let abort_handle = pool.join_set.spawn(async {});
|
let abort_handle = pool.join_set.spawn(async {});
|
||||||
|
let task_id = abort_handle.id();
|
||||||
pool.task_map_mut().insert(
|
pool.task_map_mut().insert(
|
||||||
abort_handle.id(),
|
task_id,
|
||||||
TaskMeta {
|
TaskMeta {
|
||||||
agent_index: 0,
|
agent_index: 0,
|
||||||
channel_id: Some(channel_id),
|
channel_id: Some(channel_id),
|
||||||
@@ -3712,44 +3731,86 @@ mod tests {
|
|||||||
steered_event_ids: Vec::new(),
|
steered_event_ids: Vec::new(),
|
||||||
},
|
},
|
||||||
);
|
);
|
||||||
|
task_id
|
||||||
|
}
|
||||||
|
|
||||||
assert!(pool.record_steered_event(channel_id, "aaa"));
|
/// Recording against an in-flight task stores the id in that turn's
|
||||||
assert!(pool.record_steered_event(channel_id, "bbb"));
|
/// 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
|
let ids: Vec<String> = pool
|
||||||
.task_map()
|
.task_map()
|
||||||
.values()
|
.get(&task_id)
|
||||||
.find(|m| m.channel_id == Some(channel_id))
|
|
||||||
.expect("task meta must exist")
|
.expect("task meta must exist")
|
||||||
.steered_event_ids
|
.steered_event_ids
|
||||||
.clone();
|
.clone();
|
||||||
assert_eq!(ids, vec!["aaa".to_string(), "bbb".to_string()]);
|
assert_eq!(ids, vec!["aaa".to_string(), "bbb".to_string()]);
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Recording against a channel with no in-flight turn returns false so
|
/// Recording against a task that is no longer in flight returns false
|
||||||
/// the caller cleans up the 👀 immediately.
|
/// so the caller cleans up the 👀 immediately.
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_record_steered_event_returns_false_when_no_turn_in_flight() {
|
async fn test_record_steered_event_returns_false_when_no_turn_in_flight() {
|
||||||
let mut pool = AgentPool::from_slots(vec![]);
|
let mut pool = AgentPool::from_slots(vec![]);
|
||||||
let in_flight = Uuid::new_v4();
|
let channel_id = Uuid::new_v4();
|
||||||
let other = 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 {});
|
assert!(!pool.record_steered_event(task_id, "ccc"));
|
||||||
pool.task_map_mut().insert(
|
}
|
||||||
abort_handle.id(),
|
|
||||||
TaskMeta {
|
/// Generation-reuse race: turn A ends, turn B starts on the SAME
|
||||||
agent_index: 0,
|
/// channel, then A's delayed SteerAck::Success arrives. The stale
|
||||||
channel_id: Some(in_flight),
|
/// task id must NOT attach to B — record returns false (caller
|
||||||
recoverable_batch: None,
|
/// removes the 👀 immediately) and B's meta stays untouched.
|
||||||
control_tx: None,
|
#[tokio::test]
|
||||||
steer_tx: None,
|
async fn test_late_ack_for_ended_turn_does_not_bind_to_successor_turn() {
|
||||||
steered_event_ids: Vec::new(),
|
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"));
|
/// Shutdown drain: every steered id is taken from every in-flight
|
||||||
// Empty pool: no task meta at all.
|
/// TaskMeta exactly once, leaving the metas empty.
|
||||||
let mut empty = AgentPool::from_slots(vec![]);
|
#[tokio::test]
|
||||||
assert!(!empty.record_steered_event(in_flight, "ddd"));
|
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