jerry/buzz
1
0
Fork 0
mirror of https://github.com/block/buzz.git synced 2026-11-04 00:12:02 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(acp): fence steer completion and planned replacement

Browse Source 
Signed-off-by: Alex Rosenzweig <arosenzweig@squareup.com>
This commit is contained in:
Alex Rosenzweig
2026-10-13 18:49:14 -05:00
parent 26fc4480f3
commit 79536da07c
7 changed files with 713 additions and 121 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
crates/buzz-acp/src/acp.rs
View File

@@ -150,6 +150,7 @@ pub struct AcpClient {
child: Child,
/// Retirement is independent of the terminal result of its last prompt.
shutdown: bool,
retired_for_untracked_steer: bool,
/// Sole stdin writer; None after an interrupted/failed frame closes it.
stdin: Option<ChildStdin>,
/// Framed reader over the agent's stdout pipe (line-oriented, bounded).
@@ -432,6 +433,11 @@ impl AcpClient {
self.shutdown
}
/// Whether a steering race intentionally retired this worker for replacement.
pub fn retired_for_untracked_steer(&self) -> bool {
self.retired_for_untracked_steer
}
/// Kill the agent subprocess and wait for it to exit (no zombies).
///
/// `Drop` only calls `start_kill()` (sends SIGKILL but doesn't reap).
@@ -627,6 +633,7 @@ impl AcpClient {
Ok(Self {
child,
shutdown: false,
retired_for_untracked_steer: false,
stdin: Some(stdin),
reader: FramedRead::new(stdout, LinesCodec::new_with_max_length(MAX_LINE_SIZE)),
next_id: 0,
@@ -1708,6 +1715,7 @@ impl AcpClient {
// prompt request whose completion we can await.
// Stop it before releasing the payload for retry.
// Recovery is at-least-once: work may have begun.
self.retired_for_untracked_steer = true;
self.shutdown().await;
let _ = ack_tx.send(crate::pool::SteerAck::Err(
crate::pool::SteerError::Transport(
@@ -1766,6 +1774,7 @@ impl AcpClient {
) {
// The old prompt can finish before the extension's response.
// Until its outcome is known a detached turn may be running.
self.retired_for_untracked_steer = true;
self.shutdown().await;
if let Some((_, _, ack_tx)) = pending_steer.take() {
let _ = ack_tx

220
crates/buzz-acp/src/lib.rs
View File

@@ -2374,6 +2374,7 @@ struct SteerAckEvent {
/// and deadline extension target this, not the whole channel.
scope: scope::SessionScope,
event_id: String,
turn_id: String,
/// `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
@@ -4007,7 +4008,9 @@ async fn run_harness(
// not the whole channel — a sibling thread still running in the
// same channel must keep its indicator.
if let Some(scope) = result.source.scope() {
typing_channels.remove(scope);
if !pool.has_newer_scope_owner(scope, result.agent.index) {
typing_channels.remove(scope);
}
}
if handle_prompt_result(
&mut pool,
@@ -4080,78 +4083,8 @@ async fn run_harness(
typing_channels.insert(scope, thread_tags);
}
}
Some(PoolEvent::SteerAck(SteerAckEvent {
channel_id,
scope,
event_id,
ack,
})) => {
// Acceptance and prompt completion arrive independently. Keep
// the payload until both succeed; release it on either failure.
// Unsupported steering also requests the cancel-and-merge path.
let (release_withheld, accepted_native_steer, signal_fallback) = match &ack {
Ok(pool::SteerAck::Success { .. }) => (false, true, false),
// -32601 = method_not_found: agent does not implement the
// steer extension. Fire cancel+merge so the message still
// reaches the agent.
Ok(pool::SteerAck::Err(pool::SteerError::AgentError { code, .. }))
if *code == -32601 =>
{
(true, false, true)
}
// AgentError: write landed, agent rejected it at the
// application level (e.g. wrong run id). Release for
// normal dispatch; no fallback signal (the turn is still
// running or just ended — either way there is nothing to
// cancel).
Ok(pool::SteerAck::Err(pool::SteerError::AgentError { .. })) => {
(true, false, false)
}
// Transport / ExpectedRunIdMissing / OutcomeRejected: the
// steer did not land. Release and fire the cancel+merge
// fallback so the message still reaches the agent.
Ok(pool::SteerAck::Err(_)) => (true, false, true),
Ok(pool::SteerAck::PromptCompletedNeutral) => (true, false, false),
Err(_recv_err) => (true, false, false),
};
tracing::info!(
channel = %channel_id,
event_id = %event_id,
?ack,
release_withheld,
accepted_native_steer,
signal_fallback,
"non-cancelling steer ack received"
);
if let Ok(pool::SteerAck::Success { session_id }) = &ack {
queue.extend_in_flight_deadline(&scope, config.max_turn_duration_secs);
if queue.acknowledge_native_steer(&scope, &event_id)
&& !pool.record_successful_steer(
&scope,
event_id.clone(),
session_id.clone(),
)
{
tracing::warn!(
channel = %channel_id,
event_id = %event_id,
"successful steer lost its in-flight delivery ledger"
);
}
}
if release_withheld {
queue.release_native_steer(&scope, &event_id);
}
if signal_fallback {
// Universal cancel+merge fallback. Note: the
// queued event has already been released to the
// front of `queues[scope]`, so the cancel will pick
// it up as part of the merged batch and re-prompt the
// agent. Scope-exact so the fallback cancels the
// steered event's OWN thread, not a sibling thread
// in the same channel.
signal_in_flight_task_for_scope(&mut pool, &scope, ControlSignal::Steer);
}
Some(PoolEvent::SteerAck(event)) => {
handle_steer_ack(&mut pool, &mut queue, event, config.max_turn_duration_secs);
// After releasing a withheld event, give dispatch a chance
// to re-flush. If the prompt is still in flight, the
// channel stays `in_flight_channels` and `flush_next`
@@ -4458,10 +4391,11 @@ fn signal_in_flight_task_for_scope(
scope: &scope::SessionScope,
mode: ControlSignal,
) -> bool {
let turn_id = pool.current_scope_turn(scope).map(str::to_owned);
let entry = pool
.task_map_mut()
.values_mut()
.find(|m| m.scope.as_ref() == Some(scope));
.find(|m| m.scope.as_ref() == Some(scope) && Some(&m.turn_id) == turn_id.as_ref());
if let Some(meta) = entry {
if let Some(tx) = meta.control_tx.take() {
@@ -4478,6 +4412,84 @@ fn signal_in_flight_task_for_scope(
false
}
fn handle_steer_ack(
pool: &mut AgentPool,
queue: &mut EventQueue,
event: SteerAckEvent,
max_turn_duration_secs: u64,
) {
let SteerAckEvent {
channel_id,
scope,
event_id,
turn_id,
ack,
} = event;
if !queue.owns_native_steer(&scope, &event_id, &turn_id) {
return;
}
// Acceptance and prompt completion arrive independently. Keep
// the payload until both succeed; release it on either failure.
// Unsupported steering also requests the cancel-and-merge path.
let (release_withheld, accepted_native_steer, signal_fallback) = match &ack {
Ok(pool::SteerAck::Success { .. }) => (false, true, false),
// -32601 = method_not_found: agent does not implement the
// steer extension. Fire cancel+merge so the message still
// reaches the agent.
Ok(pool::SteerAck::Err(pool::SteerError::AgentError { code, .. })) if *code == -32601 => {
(true, false, true)
}
// AgentError: write landed, agent rejected it at the
// application level (e.g. wrong run id). Release for
// normal dispatch; no fallback signal (the turn is still
// running or just ended — either way there is nothing to
// cancel).
Ok(pool::SteerAck::Err(pool::SteerError::AgentError { .. })) => (true, false, false),
// Transport / ExpectedRunIdMissing / OutcomeRejected: the
// steer did not land. Release and fire the cancel+merge
// fallback so the message still reaches the agent.
Ok(pool::SteerAck::Err(_)) => (true, false, true),
Ok(pool::SteerAck::PromptCompletedNeutral) => (true, false, false),
Err(_recv_err) => (true, false, false),
};
tracing::info!(
channel = %channel_id,
event_id = %event_id,
?ack,
release_withheld,
accepted_native_steer,
signal_fallback,
"non-cancelling steer ack received"
);
if let Ok(pool::SteerAck::Success { session_id }) = &ack {
if pool.is_current_scope_turn(&scope, &turn_id) {
queue.extend_in_flight_deadline(&scope, max_turn_duration_secs);
}
if queue.acknowledge_native_steer(&scope, &event_id)
&& !pool.record_successful_steer(&scope, &turn_id, event_id.clone(), session_id.clone())
{
tracing::warn!(
channel = %channel_id,
event_id = %event_id,
"successful steer lost its in-flight delivery ledger"
);
}
}
if release_withheld {
queue.release_native_steer(&scope, &event_id);
}
if signal_fallback && pool.is_current_scope_turn(&scope, &turn_id) {
// Universal cancel+merge fallback. Note: the
// queued event has already been released to the
// front of `queues[scope]`, so the cancel will pick
// it up as part of the merged batch and re-prompt the
// agent. Scope-exact so the fallback cancels the
// steered event's OWN thread, not a sibling thread
// in the same channel.
signal_in_flight_task_for_scope(pool, &scope, ControlSignal::Steer);
}
}
/// Attempt the non-cancelling (ACP) steer for a freshly-queued event.
///
/// Caller invariants:
@@ -4540,13 +4552,13 @@ fn try_native_steer(
};
match pool.send_steer(&scope, request) {
Ok(()) => {
Ok(turn_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
// re-deliver the event via normal dispatch. See
// `EventQueue::mark_native_steer_pending` docs at queue.rs:606.
let withheld = queue.mark_native_steer_pending(&scope, &event_id_hex);
let withheld = queue.mark_native_steer_pending(&scope, &event_id_hex, &turn_id);
if !withheld {
// Race: the event was already drained out of the queue
// before we got here (e.g. a concurrent flush picked it
@@ -4571,6 +4583,7 @@ fn try_native_steer(
channel_id,
scope: scope_for_watcher,
event_id: event_id_for_watcher,
turn_id,
ack,
});
});
@@ -5223,7 +5236,7 @@ fn handle_prompt_result(
result.outcome,
PromptOutcome::Cancelled | PromptOutcome::CancelDrainTimeout(_)
) && result.batch.is_none());
queue.finish_native_steers(scope, completed);
queue.finish_native_steers(scope, &result.turn_id, completed);
// The task may have invalidated this session before returning. Never
// resurrect delivery state for a dead session; its replacement must
// receive fresh standing context and history.
@@ -5383,7 +5396,11 @@ fn handle_prompt_result(
}
match &result.source {
PromptSource::Channel(scope) => queue.mark_complete(scope.clone()),
PromptSource::Channel(scope) => {
if !pool.has_newer_scope_owner(scope, agent_index) {
queue.mark_complete(scope.clone());
}
}
PromptSource::Heartbeat => *heartbeat_in_flight = false,
}
@@ -5443,7 +5460,11 @@ fn handle_prompt_result(
// A late steer may force worker retirement after its original prompt
// has already completed. Replace the worker without retrying that work
// or reporting the successful turn as an error.
PromptOutcome::Ok(_) if result.agent.acp.is_shutdown() => {
PromptOutcome::Ok(_) | PromptOutcome::AgentExited
if result.agent.acp.retired_for_untracked_steer()
|| (matches!(result.outcome, PromptOutcome::Ok(_))
&& result.agent.acp.is_shutdown()) =>
{
let slot_history = &mut crash_history[agent_index];
if !spawn_respawn_task(
result.agent,
@@ -5661,7 +5682,7 @@ fn recover_panicked_agent(
};
let i = meta.agent_index;
if let Some(scope) = &meta.scope {
queue.finish_native_steers(scope, false);
queue.finish_native_steers(scope, &meta.turn_id, false);
}
// Requeue BEFORE mark_complete (same rationale as handle_prompt_result).
@@ -5691,8 +5712,10 @@ fn recover_panicked_agent(
Some(scope) => {
// Clear the panicked turn's exact scope so a sibling thread in
// the same channel keeps its typing indicator.
typing_channels.remove(scope);
queue.mark_complete(scope.clone());
if !pool.has_newer_scope_owner(scope, i) {
typing_channels.remove(scope);
queue.mark_complete(scope.clone());
}
}
None => {
typing_channels.retain(|scope, _| scope.channel_id() != ch);
@@ -5975,19 +5998,24 @@ fn spawn_respawn_task(
) -> bool {
let index = old_agent.index;
// Circuit breaker: record crash, decide whether to respawn.
let delay = match slot.record_crash() {
CrashVerdict::CircuitOpen => {
tracing::error!(agent = index, "circuit open — not respawning");
return false;
}
CrashVerdict::HalfOpenProbe => {
tracing::info!(agent = index, "circuit half-open — probe respawn");
Duration::ZERO
}
CrashVerdict::Respawn(d) => {
tracing::info!(agent = index, delay_ms = d.as_millis(), "respawn backoff");
d
// Planned steering retirement replaces a healthy worker without charging
// its crash budget. Generic shutdown after failure still counts as a crash.
let delay = if old_agent.acp.retired_for_untracked_steer() {
Duration::ZERO
} else {
match slot.record_crash() {
CrashVerdict::CircuitOpen => {
tracing::error!(agent = index, "circuit open — not respawning");
return false;
}
CrashVerdict::HalfOpenProbe => {
tracing::info!(agent = index, "circuit half-open — probe respawn");
Duration::ZERO
}
CrashVerdict::Respawn(d) => {
tracing::info!(agent = index, delay_ms = d.as_millis(), "respawn backoff");
d
}
}
};
@@ -11593,6 +11621,7 @@ mod error_outcome_emission_tests {
assert!(pool.record_successful_steer(
&scope::SessionScope::Conversation { channel_id },
"t",
steer_event_id.into(),
"live-session".into(),
));
@@ -11620,6 +11649,7 @@ mod error_outcome_emission_tests {
assert!(!pool.record_successful_steer(
&scope::SessionScope::Conversation { channel_id },
"t",
"stale-event".into(),
"old-session".into(),
));

51
crates/buzz-acp/src/pool.rs
View File

@@ -1112,6 +1112,43 @@ impl AgentPool {
self.agents[idx] = Some(agent);
}
/// Whether another worker can still own this scope's in-flight guard.
/// Invalidation may remove the session owner during cancellation; until
/// overlapping tasks settle, absence of ownership is not permission to clear.
pub(crate) fn has_newer_scope_owner(&self, scope: &SessionScope, agent_index: usize) -> bool {
match self.session_owners.get(scope) {
Some(owner) => owner.agent_index != agent_index,
None => self
.task_map
.values()
.any(|meta| meta.scope.as_ref() == Some(scope) && meta.agent_index != agent_index),
}
}
fn current_scope_task(&self, scope: &SessionScope) -> Option<&TaskMeta> {
let mut tasks = self
.task_map
.values()
.filter(|meta| meta.scope.as_ref() == Some(scope));
if let Some(owner) = self.session_owners.get(scope) {
tasks.find(|meta| meta.agent_index == owner.agent_index)
} else {
let only = tasks.next()?;
tasks.next().is_none().then_some(only)
}
}
/// The exact current task, or none when invalidation left ownership ambiguous.
pub(crate) fn current_scope_turn(&self, scope: &SessionScope) -> Option<&str> {
self.current_scope_task(scope)
.map(|meta| meta.turn_id.as_str())
}
/// Whether a steering acknowledgement still refers to the current task.
pub(crate) fn is_current_scope_turn(&self, scope: &SessionScope, turn_id: &str) -> bool {
self.current_scope_turn(scope) == Some(turn_id)
}
/// Whether any agent is currently idle (sitting in its slot).
pub fn any_idle(&self) -> bool {
self.agents.iter().any(|slot| slot.is_some())
@@ -1202,7 +1239,7 @@ impl AgentPool {
/// Try to send a goose-native steer request to the in-flight task for
/// `channel_id`.
///
/// Returns `Ok(())` if the request was accepted by the read loop's
/// Returns `Ok(turn_id)` if the request was accepted by the read loop's
/// receiver (capacity-1 mpsc; one slot is the single in-flight steer
/// write). Returns `Err(SteerError::Transport(_))` on `Full`/`Closed`
/// (already-in-flight write, or read loop torn down); the `Transport`
@@ -1227,18 +1264,16 @@ impl AgentPool {
&mut self,
scope: &SessionScope,
request: SteerRequest,
) -> Result<(), SteerError> {
) -> Result<String, SteerError> {
let meta = self
.task_map
.values_mut()
.find(|m| m.scope.as_ref() == Some(scope))
.current_scope_task(scope)
.ok_or(SteerError::PromptCompleted)?;
let tx = meta
.steer_tx
.as_ref()
.ok_or_else(|| SteerAdmissionReason::SenderAbsent.transport_refusal())?;
match tx.try_send(request) {
Ok(()) => Ok(()),
Ok(()) => Ok(meta.turn_id.clone()),
// Capacity-1 mailbox already holds one in-flight steer write.
// The refused request (including its oneshot ack) comes back
// inside the error and is dropped here, exactly as before.
@@ -1259,13 +1294,14 @@ impl AgentPool {
pub fn record_successful_steer(
&mut self,
scope: &SessionScope,
turn_id: &str,
event_id: String,
session_id: String,
) -> bool {
if let Some(meta) = self
.task_map
.values_mut()
.find(|meta| meta.scope.as_ref() == Some(scope))
.find(|meta| meta.scope.as_ref() == Some(scope) && meta.turn_id == turn_id)
{
meta.successful_steer_deliveries
.insert(SuccessfulSteerDelivery {
@@ -7908,6 +7944,7 @@ printf '%s\n' '{{"jsonrpc":"2.0","id":0,"result":{{"stopReason":"end_turn"}}}}'"
let mut pool = AgentPool::from_slots(vec![Some(agent)]);
assert!(pool.record_successful_steer(
&conv(channel_id),
"t",
steered_event_id.clone(),
"live-session".into(),
));

21
crates/buzz-acp/src/queue.rs
View File

@@ -6116,8 +6116,8 @@ mod tests {
//
// Side-table semantics: `mark_native_steer_pending` moves an event out of
// `queues` into `withheld_native_steer`, making it invisible to
// `flush_next` / `has_flushable_work` / contiguous drain. `Success` ack
// drops it via `remove_event`; `Err` / `PromptCompletedNeutral` ack
// `flush_next` / `has_flushable_work` / contiguous drain. Successful ack
// and matching turn completion drop it; `Err` / `PromptCompletedNeutral` ack
// restores it to the queue front via `release_native_steer`. The
// `in_flight_deadline` expiry bulk-recovers withheld events so they
// are never permanently orphaned.
@@ -6135,7 +6135,7 @@ mod tests {
let event_id = qe.event.id.to_hex();
q.push(qe);
assert!(q.mark_native_steer_pending(ch, &event_id));
assert!(q.mark_native_steer_pending(ch, &event_id, "turn"));
assert!(
q.flush_next().is_none(),
@@ -6175,7 +6175,7 @@ mod tests {
q.push(e3);
// Steer in flight for e3 — withhold it from normal dispatch.
assert!(q.mark_native_steer_pending(ch, &e3_id));
assert!(q.mark_native_steer_pending(ch, &e3_id, "turn"));
// Earlier events flush as a normal batch; e3 is invisible.
let batch = q
@@ -6220,7 +6220,7 @@ mod tests {
q.in_flight_scopes.insert(conv(ch));
q.in_flight_deadlines.insert(conv(ch), Instant::now());
q.in_flight_batch_sizes.insert(conv(ch), 1);
assert!(q.mark_native_steer_pending(ch, &event_id));
assert!(q.mark_native_steer_pending(ch, &event_id, "turn"));
// Force the in-flight deadline to be in the past, simulating the
// steer ack never arriving and the read loop hanging long enough
@@ -6275,9 +6275,9 @@ mod tests {
// tests. What matters here is that the bulk-recovery path
// (reverse iter + push_front) composes to original FIFO at the
// queue front.
assert!(q.mark_native_steer_pending(ch, &e1_id));
assert!(q.mark_native_steer_pending(ch, &e2_id));
assert!(q.mark_native_steer_pending(ch, &e3_id));
assert!(q.mark_native_steer_pending(ch, &e1_id, "turn"));
assert!(q.mark_native_steer_pending(ch, &e2_id, "turn"));
assert!(q.mark_native_steer_pending(ch, &e3_id, "turn"));
assert_eq!(pending_count(&q), 0);
assert_eq!(
q.withheld_native_steer.get(&conv(ch)).map(|v| v.len()),
@@ -6782,7 +6782,7 @@ mod tests {
let withheld = queued_edit(ch, &withheld_target, None);
let event_id = withheld.event.id.to_hex();
q.push(withheld);
assert!(q.mark_native_steer_pending(ch, &event_id));
assert!(q.mark_native_steer_pending(ch, &event_id, "turn"));
let drained: HashSet<String> = q.drain_channel(ch).into_iter().collect();
assert_eq!(drained, HashSet::from([queued_target, withheld_target]));
@@ -7268,3 +7268,6 @@ mod tests {
);
}
}
#[cfg(test)]
mod steer_ownership_tests;

26
crates/buzz-acp/src/queue/native_steer.rs
View File

@@ -3,6 +3,7 @@ use super::*;
pub(super) struct WithheldSteer {
pub(super) queued: QueuedEvent,
turn_id: String,
acknowledged: bool,
completed: Option<bool>,
}
@@ -18,6 +19,17 @@ impl EventQueue {
< MAX_PENDING_PER_CHANNEL
}
/// Whether this acknowledgement still owns the retained delivery attempt.
pub fn owns_native_steer(&self, scope: &SessionScope, event_id: &str, turn_id: &str) -> bool {
self.withheld_native_steer
.get(scope)
.is_some_and(|entries| {
entries.iter().any(|entry| {
entry.turn_id == turn_id && entry.queued.event.id.to_hex() == event_id
})
})
}
/// Record acceptance without discarding the payload until its turn settles.
/// Returns false when a failed turn requires ordinary redelivery instead.
pub fn acknowledge_native_steer(&mut self, scope: &SessionScope, event_id: &str) -> bool {
@@ -46,14 +58,14 @@ impl EventQueue {
/// Settle only this turn's follow-ups. The ack and result may arrive in
/// either order; entries already settled by an earlier turn stay unchanged.
pub fn finish_native_steers(&mut self, scope: &SessionScope, completed: bool) {
pub fn finish_native_steers(&mut self, scope: &SessionScope, turn_id: &str, completed: bool) {
let Some(entries) = self.withheld_native_steer.get_mut(scope) else {
return;
};
let ready: Vec<_> = entries
.iter_mut()
.filter_map(|entry| {
if entry.completed.is_some() {
if entry.turn_id != turn_id || entry.completed.is_some() {
return None;
}
entry.completed = Some(completed);
@@ -79,10 +91,15 @@ impl EventQueue {
/// the event may have already been drained, removed, or never queued).
///
/// Must be called synchronously from the mode-gate fork immediately
/// after `pool.send_steer` returns `Ok(())` and before any watcher task
/// after `pool.send_steer` returns `Ok(turn_id)` and before any watcher task
/// is spawned, so the withhold is established before `mark_complete` /
/// any subsequent `flush_next` tick can run.
pub fn mark_native_steer_pending<K: IntoScope>(&mut self, scope: K, event_id: &str) -> bool {
pub fn mark_native_steer_pending<K: IntoScope>(
&mut self,
scope: K,
event_id: &str,
turn_id: &str,
) -> bool {
let scope = scope.into_scope();
let Some(q) = self.queues.get_mut(&scope) else {
return false;
@@ -101,6 +118,7 @@ impl EventQueue {
.or_default()
.push(WithheldSteer {
queued: qe,
turn_id: turn_id.to_owned(),
acknowledged: false,
completed: None,
});

316
crates/buzz-acp/src/queue/steer_ownership_tests.rs Normal file
View File

@@ -0,0 +1,316 @@
//! Real queue expiry and production admission/ACK/result seams, without a journal.
use super::*;
use crate::error_outcome_emission_tests::{bind_agent_scope_owner, dummy_agent, test_config};
use crate::*;
fn event(scope: &SessionScope) -> QueuedEvent {
QueuedEvent {
channel_id: scope.channel_id(),
scope: scope.clone(),
event: nostr::EventBuilder::new(nostr::Kind::Custom(9), "follow-up")
.sign_with_keys(&nostr::Keys::generate())
.unwrap(),
received_at: Instant::now(),
prompt_tag: "mention".into(),
edit: None,
}
}
fn task(
pool: &mut AgentPool,
agent: &mut OwnedAgent,
scope: &SessionScope,
turn: &str,
) -> (
tokio::task::Id,
mpsc::Receiver<pool::SteerRequest>,
tokio::sync::oneshot::Receiver<ControlSignal>,
) {
bind_agent_scope_owner(pool, agent, scope.clone());
let id = pool.join_set.spawn(std::future::pending()).id();
let (steer_tx, rx) = mpsc::channel(1);
let (control_tx, control_rx) = tokio::sync::oneshot::channel();
pool.task_map_mut().insert(
id,
pool::TaskMeta {
agent_index: agent.index,
channel_id: Some(scope.channel_id()),
scope: Some(scope.clone()),
turn_id: turn.into(),
recoverable_batch: None,
control_tx: Some(control_tx),
steer_tx: Some(steer_tx),
successful_steer_deliveries: HashSet::new(),
},
);
(id, rx, control_rx)
}
async fn admit(
pool: &mut AgentPool,
queue: &mut EventQueue,
event: QueuedEvent,
rx: &mut mpsc::Receiver<pool::SteerRequest>,
) -> (pool::SteerRequest, mpsc::UnboundedReceiver<SteerAckEvent>) {
let scope = event.scope.clone();
let batch_event = BatchEvent {
event: event.event.clone(),
received_at: event.received_at,
prompt_tag: event.prompt_tag.clone(),
edit: None,
};
queue.push(event);
let (tx, ack_rx) = mpsc::unbounded_channel();
assert!(try_native_steer(pool, queue, scope, batch_event, &tx));
(
rx.try_recv().expect("current task receives the steer"),
ack_rx,
)
}
fn circuit() -> SlotCircuit {
SlotCircuit {
crash_times: vec![],
open_until: None,
respawn_in_flight: false,
}
}
#[tokio::test]
async fn expired_turn_result_or_task_failure_cannot_settle_replacement_follow_up() {
for (panic, invalidate_owner) in [(false, false), (true, false), (false, true), (true, true)] {
let scope = SessionScope::Conversation {
channel_id: Uuid::new_v4(),
};
let mut queue = EventQueue::new(DedupMode::Queue);
queue.push(event(&scope));
let _original = queue.flush_next().unwrap();
let mut old = dummy_agent(0).await;
let mut next = dummy_agent(1).await;
let mut pool = AgentPool::from_slots(vec![None, None]);
let (old_id, mut old_rx, _old_control) = task(&mut pool, &mut old, &scope, "old");
let (pending, _old_ack) = admit(&mut pool, &mut queue, event(&scope), &mut old_rx).await;
queue
.in_flight_deadlines
.insert(scope.clone(), Instant::now());
let recovered = queue
.flush_next()
.expect("real expiry recovers original steer");
assert_eq!(recovered.events.len(), 1);
let (_, mut next_rx, _next_control) = task(&mut pool, &mut next, &scope, "next");
let follow_up = event(&scope);
let follow_up_id = follow_up.event.id.to_hex();
let (accepted, mut ack_rx) = admit(&mut pool, &mut queue, follow_up, &mut next_rx).await;
accepted
.ack_tx
.send(pool::SteerAck::Success {
session_id: "next-session".into(),
})
.unwrap();
let ack = tokio::time::timeout(Duration::from_secs(2), ack_rx.recv())
.await
.unwrap()
.unwrap();
assert_eq!(ack.turn_id, "next");
handle_steer_ack(&mut pool, &mut queue, ack, 60);
if invalidate_owner {
pool.invalidate_scope_session(&scope);
let (ack_tx, _ack_rx) = tokio::sync::oneshot::channel();
assert!(matches!(
pool.send_steer(
&scope,
pool::SteerRequest {
prompt_blocks: vec!["ambiguous".into()],
ack_tx,
}
),
Err(pool::SteerError::PromptCompleted)
));
assert!(!signal_in_flight_task_for_scope(
&mut pool,
&scope,
ControlSignal::Steer
));
assert!(old_rx.try_recv().is_err());
assert!(next_rx.try_recv().is_err());
assert!(pool
.task_map()
.values()
.all(|meta| meta.control_tx.is_some()));
}
let mut history = [circuit(), circuit()];
let (tx, _rx) = mpsc::channel(4);
let mut respawns = tokio::task::JoinSet::new();
if panic {
// Preserve the real task identity so production panic cleanup runs.
pool.join_set.abort_all();
let mut old_error = None;
while let Some(joined) = pool.join_set.join_next().await {
let error = joined.unwrap_err();
if error.id() == old_id {
old_error = Some(error);
}
}
let mut typing = HashMap::from([(scope.clone(), ThreadTags::default())]);
recover_panicked_agent(
&mut pool,
&mut queue,
&test_config(),
old_error.unwrap(),
&mut false,
&HashSet::new(),
&mut typing,
&mut history,
&tx,
&mut respawns,
None,
);
assert!(
typing.contains_key(&scope),
"old panic must not clear replacement typing"
);
old.acp.shutdown().await;
} else {
handle_prompt_result(
&mut pool,
&mut queue,
&test_config(),
PromptResult {
agent: old,
source: PromptSource::Channel(scope.clone()),
turn_id: "old".into(),
outcome: PromptOutcome::Ok(acp::StopReason::EndTurn),
batch: None,
},
&mut false,
&HashSet::new(),
&mut history,
&tx,
&mut respawns,
None,
None,
&mut tokio::time::Instant::now(),
);
}
assert!(
queue.is_scope_in_flight(&scope),
"old completion cannot clear replacement guard"
);
assert!(
queue.owns_native_steer(&scope, &follow_up_id, "next"),
"old completion cannot discard new follow-up"
);
assert!(queue.flush_next().is_none());
next.acp.shutdown().await;
handle_prompt_result(
&mut pool,
&mut queue,
&test_config(),
PromptResult {
agent: next,
source: PromptSource::Channel(scope.clone()),
turn_id: "next".into(),
outcome: PromptOutcome::AgentExited,
batch: None,
},
&mut false,
&HashSet::new(),
&mut history,
&tx,
&mut respawns,
None,
None,
&mut tokio::time::Instant::now(),
);
let retry = queue
.flush_next()
.expect("replacement failure retains accepted follow-up");
assert_eq!(retry.events.len(), 1);
assert_eq!(retry.events[0].event.id.to_hex(), follow_up_id);
drop(pending);
respawns.shutdown().await;
shutdown_agent_pool(&mut pool).await;
}
}
#[tokio::test]
async fn expired_turn_ack_cannot_settle_or_cancel_the_same_event_in_a_new_turn() {
for late_ack in [
pool::SteerAck::Success {
session_id: "old-session".into(),
},
pool::SteerAck::Err(pool::SteerError::Transport("fixture failure".into())),
pool::SteerAck::PromptCompletedNeutral,
] {
let scope = SessionScope::Conversation {
channel_id: Uuid::new_v4(),
};
let mut queue = EventQueue::new(DedupMode::Queue);
queue.push(event(&scope));
queue.flush_next().unwrap();
let mut old = dummy_agent(0).await;
let mut next = dummy_agent(1).await;
let mut pool = AgentPool::from_slots(vec![None, None]);
let (_, mut old_rx, mut old_control) = task(&mut pool, &mut old, &scope, "old");
let follow_up = event(&scope);
let follow_up_id = follow_up.event.id.to_hex();
let (pending, mut old_ack_rx) =
admit(&mut pool, &mut queue, follow_up.clone(), &mut old_rx).await;
queue
.in_flight_deadlines
.insert(scope.clone(), Instant::now());
queue.flush_next().unwrap();
let (_, mut next_rx, mut next_control) = task(&mut pool, &mut next, &scope, "next");
let (accepted, mut next_ack_rx) =
admit(&mut pool, &mut queue, follow_up, &mut next_rx).await;
let deadline = queue.in_flight_deadlines[&scope];
pending.ack_tx.send(late_ack).unwrap();
let ack = tokio::time::timeout(Duration::from_secs(2), old_ack_rx.recv())
.await
.unwrap()
.unwrap();
handle_steer_ack(&mut pool, &mut queue, ack, 7200);
assert_eq!(queue.in_flight_deadlines[&scope], deadline);
assert!(queue.owns_native_steer(&scope, &follow_up_id, "next"));
assert!(matches!(
next_control.try_recv(),
Err(tokio::sync::oneshot::error::TryRecvError::Empty)
));
// Completing this turn before its own ack cannot consume the stale ack.
queue.finish_native_steers(&scope, "next", true);
assert!(queue.owns_native_steer(&scope, &follow_up_id, "next"));
accepted
.ack_tx
.send(pool::SteerAck::Success {
session_id: "next-session".into(),
})
.unwrap();
let ack = tokio::time::timeout(Duration::from_secs(2), next_ack_rx.recv())
.await
.unwrap()
.unwrap();
handle_steer_ack(&mut pool, &mut queue, ack, 60);
assert!(!queue.owns_native_steer(&scope, &follow_up_id, "next"));
let (rejected, mut fallback_rx) =
admit(&mut pool, &mut queue, event(&scope), &mut next_rx).await;
rejected
.ack_tx
.send(pool::SteerAck::Err(pool::SteerError::Transport(
"fallback".into(),
)))
.unwrap();
let ack = tokio::time::timeout(Duration::from_secs(2), fallback_rx.recv())
.await
.unwrap()
.unwrap();
handle_steer_ack(&mut pool, &mut queue, ack, 60);
assert!(matches!(next_control.try_recv(), Ok(ControlSignal::Steer)));
assert!(matches!(
old_control.try_recv(),
Err(tokio::sync::oneshot::error::TryRecvError::Empty)
));
old.acp.shutdown().await;
next.acp.shutdown().await;
shutdown_agent_pool(&mut pool).await;
}
}

191
crates/buzz-acp/src/steer_recovery_tests.rs
View File

@@ -40,7 +40,7 @@ async fn settle(ack_first: bool, success: bool) {
let event = follow_up(&scope);
let event_id = event.event.id.to_hex();
queue.push(event);
assert!(queue.mark_native_steer_pending(&scope, &event_id));
assert!(queue.mark_native_steer_pending(&scope, &event_id, "turn"));
if ack_first {
assert!(queue.acknowledge_native_steer(&scope, &event_id));
}
@@ -132,7 +132,7 @@ async fn acknowledged_native_follow_up_survives_prompt_task_panic() {
let event = follow_up(&scope);
let id = event.event.id.to_hex();
queue.push(event);
assert!(queue.mark_native_steer_pending(&scope, &id));
assert!(queue.mark_native_steer_pending(&scope, &id, "turn"));
assert!(queue.acknowledge_native_steer(&scope, &id));
let failure = pool.join_set.join_next().await.unwrap().unwrap_err();
let (tx, _rx) = mpsc::channel(1);
@@ -167,9 +167,9 @@ fn a_later_turn_does_not_overwrite_a_failed_turn_waiting_for_its_ack() {
let event = follow_up(&scope);
let id = event.event.id.to_hex();
queue.push(event);
assert!(queue.mark_native_steer_pending(&scope, &id));
queue.finish_native_steers(&scope, false);
queue.finish_native_steers(&scope, true);
assert!(queue.mark_native_steer_pending(&scope, &id, "turn"));
queue.finish_native_steers(&scope, "turn", false);
queue.finish_native_steers(&scope, "later-turn", true);
assert!(!queue.acknowledge_native_steer(&scope, &id));
assert_eq!(queue.flush_next().unwrap().events[0].event.id.to_hex(), id);
}
@@ -187,7 +187,7 @@ async fn original_completion_before_steer_ack_retries_only_the_follow_up() {
let follow_up = follow_up(&scope);
let follow_up_id = follow_up.event.id.to_hex();
queue.push(follow_up);
assert!(queue.mark_native_steer_pending(&scope, &follow_up_id));
assert!(queue.mark_native_steer_pending(&scope, &follow_up_id, "original"));
// Complete the original request only after the steer is on the wire,
// then report a detached follow-up. The original work is definitively done.
@@ -322,3 +322,182 @@ cat >/dev/null"#;
shutdown_agent_pool(&mut pool).await;
server.abort();
}
/// Drive both wire orderings through real subprocess replacement three times,
/// then prove actual failures still open the same slot's circuit.
#[tokio::test]
async fn planned_steer_retirements_preserve_availability_and_genuine_crashes_count() {
for original_first in [true, false] {
let original = r#"{"jsonrpc":"2.0","id":1,"result":{"stopReason":"end_turn"}}"#;
let steer = r#"{"jsonrpc":"2.0","id":2,"result":{"outcome":"startedNewTurn"}}"#;
let (first, second) = if original_first {
(original, steer)
} else {
(steer, original)
};
let script = format!("read initialize\necho '{{\"jsonrpc\":\"2.0\",\"id\":0,\"result\":{{\"_meta\":{{\"steering\":{{\"supported\":true}}}}}}}}'\nread prompt\nread steer\necho '{first}'\necho '{second}'\ncat >/dev/null");
let mut config = test_config();
config.agent_command = "/bin/sh".into();
config.agent_args = vec!["-c".into(), script];
let mut agent = dummy_agent(0).await;
agent.acp.shutdown().await;
agent.acp = spawn_and_init(
&config.agent_command,
&config.agent_args,
&[],
false,
0,
None,
)
.await
.unwrap()
.0;
let mut pool = AgentPool::from_slots(vec![Some(agent)]);
let mut history = [SlotCircuit {
crash_times: vec![],
open_until: None,
respawn_in_flight: false,
}];
let (tx, mut rx) = mpsc::channel(1);
let mut tasks = tokio::task::JoinSet::new();
for round in 0..6 {
let planned = round < 3;
let scope = scope::SessionScope::Conversation {
channel_id: Uuid::new_v4(),
};
let turn_id = format!("turn-{round}");
let mut agent = pool
.try_claim(None)
.expect("replacement must remain available");
bind_agent_scope_owner(&mut pool, &mut agent, scope.clone());
let task_id = pool.join_set.spawn(async {}).id();
pool.task_map_mut().insert(
task_id,
pool::TaskMeta {
agent_index: 0,
channel_id: Some(scope.channel_id()),
scope: Some(scope.clone()),
turn_id: turn_id.clone(),
recoverable_batch: None,
control_tx: None,
steer_tx: None,
successful_steer_deliveries: HashSet::new(),
},
);
let mut queue = EventQueue::new(DedupMode::Queue);
let event = follow_up(&scope);
let event_id = event.event.id.to_hex();
queue.push(event);
assert!(queue.mark_native_steer_pending(&scope, &event_id, &turn_id));
let outcome = if planned {
let (steer_tx, steer_rx) = mpsc::channel(1);
agent.acp.install_steer_rx(steer_rx);
let (ack_tx, ack_rx) = tokio::sync::oneshot::channel();
steer_tx
.send(pool::SteerRequest {
prompt_blocks: vec!["follow-up".into()],
ack_tx,
})
.await
.unwrap();
let result = tokio::time::timeout(
Duration::from_secs(5),
agent.acp.session_prompt_with_idle_timeout(
"session",
"original",
Duration::from_secs(3),
Duration::from_secs(3),
),
)
.await
.unwrap();
assert!(agent.acp.retired_for_untracked_steer());
assert!(agent.acp.is_shutdown());
assert!(matches!(
ack_rx.await.unwrap(),
pool::SteerAck::PromptCompletedNeutral | pool::SteerAck::Err(_)
));
queue.release_native_steer(&scope, &event_id);
if original_first {
PromptOutcome::Ok(result.expect("known original success preserved"))
} else {
assert!(matches!(result, Err(acp::AcpError::AgentExited)));
PromptOutcome::AgentExited
}
} else {
// A failure path can also have already shut down the client.
// Only the explicitly planned retirement flag exempts a crash.
agent.acp.shutdown().await;
assert!(!agent.acp.retired_for_untracked_steer());
PromptOutcome::AgentExited
};
let action = handle_prompt_result(
&mut pool,
&mut queue,
&config,
PromptResult {
agent,
source: PromptSource::Channel(scope.clone()),
turn_id,
outcome,
batch: None,
},
&mut false,
&HashSet::new(),
&mut history,
&tx,
&mut tasks,
None,
None,
&mut tokio::time::Instant::now(),
);
if planned {
assert!(
history[0].crash_times.is_empty(),
"planned retirement must not charge crash history"
);
assert!(history[0].open_until.is_none());
assert_eq!(
queue.flush_next().unwrap().events[0].event.id.to_hex(),
event_id
);
} else {
assert_eq!(history[0].crash_times.len(), round - 2);
}
if round == 5 {
assert!(matches!(action, LoopAction::Exit));
assert!(
history[0].open_until.is_some(),
"three real failures must still open the circuit"
);
} else {
assert!(matches!(action, LoopAction::Continue));
assert!(history[0].respawn_in_flight);
let replacement = tokio::time::timeout(Duration::from_secs(5), rx.recv())
.await
.unwrap()
.unwrap();
let (acp, protocol_version, agent_name) =
replacement.result.expect("replacement handshake succeeds");
history[0].respawn_in_flight = false;
pool.return_agent(OwnedAgent {
index: 0,
acp,
protocol_version,
agent_name,
state: Default::default(),
model_capabilities: None,
desired_model: None,
model_overridden: false,
desired_model_request_id: None,
desired_model_pending_ack: false,
startup_effort: None,
goose_system_prompt_supported: None,
});
assert!(pool.any_idle());
}
}
tasks.shutdown().await;
shutdown_agent_pool(&mut pool).await;
}
}