External/buzz
1
0
Fork 0
mirror of https://github.com/block/buzz.git synced 2026-08-18 06:50:31 +02:00
Code Issues Packages Projects Releases Wiki Activity

fix(buzz-acp): clear steer_rx on all run_prompt_task exit paths (#1391)

Browse Source 
Signed-off-by: Will Pfleger <pfleger.will@gmail.com>
Co-authored-by: npub1mn7jgtj4w2pd0g0zeuhxsa6jy6p0rewxz4kujt98my82ahfmp72sxjexk7 <dcfd242e557282d7a1e2cf2e6877522682f1e5c6156dc92ca7d90eaedd3b0f95@sprout-oss.stage.blox.sqprod.co>
This commit is contained in:
Will Pfleger
2026-06-30 11:41:17 -04:00
committed by GitHub
co-authored by npub1mn7jgtj4w2pd0g0zeuhxsa6jy6p0rewxz4kujt98my82ahfmp72sxjexk7
parent 0827171d5e
commit 10aaa72f01
2 changed files with 273 additions and 88 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/buzz-acp/src/acp.rs
+19
View File
@@ -522,6 +522,25 @@ impl AcpClient {
self.steer_rx = Some(rx);
}
/// Clear any installed steer receiver without consuming it.
///
/// Called by `send_prompt_result` on every exit path of `run_prompt_task`
/// so that `install_steer_rx`'s `is_none()` invariant holds for the next
/// dispatch even when the turn ended before the read loop ran `take()`.
/// Idempotent — safe to call when `steer_rx` is already `None`.
pub fn clear_steer_rx(&mut self) {
self.steer_rx = None;
}
/// Returns `true` if no steer receiver is currently installed.
///
/// Test-only: used by `pool` tests to assert the post-return invariant
/// without exposing the private field directly.
#[cfg(test)]
pub fn steer_rx_is_none(&self) -> bool {
self.steer_rx.is_none()
}
/// Cancel a turn cleanly, handling any pending permission request first.
///
/// Steps:
crates/buzz-acp/src/pool.rs
+254 -88
View File
@@ -996,6 +996,35 @@ fn with_core(framed: Option<String>, core: Option<&str>) -> Option<String> {
}
}
/// Return `agent` to the pool via `result_tx`, clearing any steer receiver first.
///
/// Every path that returns an `OwnedAgent` to the pool via `PromptResult` goes
/// through this function. Panic/abort paths do not — and don't need to, since a
/// panicked task's agent is never sent back via `PromptResult`.
///
/// Clearing `steer_rx` here — rather than per-arm — makes the `install_steer_rx`
/// invariant (`steer_rx.is_none()` at dispatch) structurally unviolatable: a receiver
/// installed for a turn that ends before the read loop's `take()` (e.g. session-create
/// error) is always dropped before the agent re-enters the pool, so the next dispatch
/// can never trigger the assert.
///
/// On the happy path the read loop has already called `take()`, so this is a no-op.
fn send_prompt_result(
result_tx: &mpsc::UnboundedSender<PromptResult>,
mut agent: OwnedAgent,
source: PromptSource,
outcome: PromptOutcome,
batch: Option<FlushBatch>,
) {
agent.acp.clear_steer_rx();
let _ = result_tx.send(PromptResult {
agent,
source,
outcome,
batch,
});
}
/// Core async function spawned for each prompt.
///
/// Lifecycle:
@@ -1153,21 +1182,23 @@ pub async fn run_prompt_task(
}
Err(AcpError::AgentExited) => {
agent.state.invalidate_all();
let _ = result_tx.send(PromptResult {
send_prompt_result(
&result_tx,
agent,
source,
outcome: PromptOutcome::AgentExited,
batch: requeue_batch_if_queue(&ctx, batch),
});
PromptOutcome::AgentExited,
requeue_batch_if_queue(&ctx, batch),
);
return;
}
Err(e) => {
let _ = result_tx.send(PromptResult {
send_prompt_result(
&result_tx,
agent,
source,
outcome: PromptOutcome::Error(e),
batch: requeue_batch_if_queue(&ctx, batch),
});
PromptOutcome::Error(e),
requeue_batch_if_queue(&ctx, batch),
);
return;
}
}
@@ -1189,21 +1220,23 @@ pub async fn run_prompt_task(
}
Err(AcpError::AgentExited) => {
agent.state.invalidate_all();
let _ = result_tx.send(PromptResult {
send_prompt_result(
&result_tx,
agent,
source,
outcome: PromptOutcome::AgentExited,
batch: None,
});
PromptOutcome::AgentExited,
None,
);
return;
}
Err(e) => {
let _ = result_tx.send(PromptResult {
send_prompt_result(
&result_tx,
agent,
source,
outcome: PromptOutcome::Error(e),
batch: None,
});
PromptOutcome::Error(e),
None,
);
return;
}
}
@@ -1254,12 +1287,13 @@ pub async fn run_prompt_task(
}
Err(AcpError::AgentExited) => {
agent.state.invalidate_all();
let _ = result_tx.send(PromptResult {
send_prompt_result(
&result_tx,
agent,
source,
outcome: PromptOutcome::AgentExited,
batch: requeue_batch_if_queue(&ctx, batch),
});
PromptOutcome::AgentExited,
requeue_batch_if_queue(&ctx, batch),
);
return;
}
Err(AcpError::IdleTimeout(_)) => {
@@ -1278,12 +1312,13 @@ pub async fn run_prompt_task(
}
Err(AcpError::AgentExited) => {
agent.state.invalidate_all();
let _ = result_tx.send(PromptResult {
send_prompt_result(
&result_tx,
agent,
source,
outcome: PromptOutcome::AgentExited,
batch: requeue_batch_if_queue(&ctx, batch),
});
PromptOutcome::AgentExited,
requeue_batch_if_queue(&ctx, batch),
);
return;
}
Err(e) => {
@@ -1294,12 +1329,13 @@ pub async fn run_prompt_task(
agent.state.invalidate(&source);
}
}
let _ = result_tx.send(PromptResult {
send_prompt_result(
&result_tx,
agent,
source,
outcome: PromptOutcome::Timeout,
batch: requeue_batch_if_queue(&ctx, batch),
});
PromptOutcome::Timeout,
requeue_batch_if_queue(&ctx, batch),
);
return;
}
Err(AcpError::HardTimeout) => {
@@ -1309,12 +1345,13 @@ pub async fn run_prompt_task(
ctx.max_turn_duration.as_secs()
);
agent.state.invalidate_all();
let _ = result_tx.send(PromptResult {
send_prompt_result(
&result_tx,
agent,
source,
outcome: PromptOutcome::Timeout,
batch: requeue_batch_if_queue(&ctx, batch),
});
PromptOutcome::Timeout,
requeue_batch_if_queue(&ctx, batch),
);
return;
}
Err(e) => {
@@ -1323,12 +1360,13 @@ pub async fn run_prompt_task(
"initial_message failed for channel {cid}: {e} — invalidating session"
);
agent.state.invalidate(&source);
let _ = result_tx.send(PromptResult {
send_prompt_result(
&result_tx,
agent,
source,
outcome: PromptOutcome::Error(e),
batch: requeue_batch_if_queue(&ctx, batch),
});
PromptOutcome::Error(e),
requeue_batch_if_queue(&ctx, batch),
);
return;
}
}
@@ -1396,12 +1434,13 @@ pub async fn run_prompt_task(
// Should not happen — batch is None only for heartbeats which have prompt_text.
// Return the agent to the pool to prevent a permanent slot leak.
tracing::error!("run_prompt_task: no batch and no prompt_text — returning agent");
let _ = result_tx.send(PromptResult {
send_prompt_result(
&result_tx,
agent,
source,
outcome: PromptOutcome::Error(AcpError::Protocol("no batch and no prompt_text".into())),
batch: None,
});
PromptOutcome::Error(AcpError::Protocol("no batch and no prompt_text".into())),
None,
);
return;
};
@@ -1500,12 +1539,13 @@ pub async fn run_prompt_task(
let retry_batch =
requeue_cancelled_batch(&ctx, control_signal, batch);
let _ = result_tx.send(PromptResult {
send_prompt_result(
&result_tx,
agent,
source,
outcome: PromptOutcome::Cancelled,
batch: retry_batch,
});
PromptOutcome::Cancelled,
retry_batch,
);
return;
}
Err(AcpError::AgentExited) => {
@@ -1513,12 +1553,13 @@ pub async fn run_prompt_task(
let retry_batch =
requeue_cancelled_batch(&ctx, control_signal, batch);
let _ = result_tx.send(PromptResult {
send_prompt_result(
&result_tx,
agent,
source,
outcome: PromptOutcome::AgentExited,
batch: retry_batch,
});
PromptOutcome::AgentExited,
retry_batch,
);
return;
}
Err(AcpError::IdleTimeout(_) | AcpError::HardTimeout) => {
@@ -1527,12 +1568,13 @@ pub async fn run_prompt_task(
let retry_batch =
requeue_cancelled_batch(&ctx, control_signal, batch);
let _ = result_tx.send(PromptResult {
send_prompt_result(
&result_tx,
agent,
source,
outcome: PromptOutcome::Timeout,
batch: retry_batch,
});
PromptOutcome::Timeout,
retry_batch,
);
return;
}
Err(e) => {
@@ -1540,12 +1582,13 @@ pub async fn run_prompt_task(
let retry_batch =
requeue_cancelled_batch(&ctx, control_signal, batch);
let _ = result_tx.send(PromptResult {
send_prompt_result(
&result_tx,
agent,
source,
outcome: PromptOutcome::Error(e),
batch: retry_batch,
});
PromptOutcome::Error(e),
retry_batch,
);
return;
}
}
@@ -1583,12 +1626,13 @@ pub async fn run_prompt_task(
&source,
&control_signal,
);
let _ = result_tx.send(PromptResult {
send_prompt_result(
&result_tx,
agent,
source,
outcome: PromptOutcome::Ok(StopReason::EndTurn),
batch: None, // turn succeeded — batch was processed, no requeue
});
PromptOutcome::Ok(StopReason::EndTurn),
None, // turn succeeded — batch was processed, no requeue
);
return;
}
}
@@ -1632,22 +1676,24 @@ pub async fn run_prompt_task(
agent.state.invalidate(&source);
}
let _ = result_tx.send(PromptResult {
send_prompt_result(
&result_tx,
agent,
source,
outcome: PromptOutcome::Ok(stop_reason),
batch: None,
});
PromptOutcome::Ok(stop_reason),
None,
);
}
Err(AcpError::AgentExited) => {
tracing::error!(target: "pool::prompt", "agent {} exited during prompt", agent.index);
agent.state.invalidate_all();
let _ = result_tx.send(PromptResult {
send_prompt_result(
&result_tx,
agent,
source,
outcome: PromptOutcome::AgentExited,
batch: requeue_batch_if_queue(&ctx, batch),
});
PromptOutcome::AgentExited,
requeue_batch_if_queue(&ctx, batch),
);
}
Err(AcpError::IdleTimeout(_)) => {
tracing::warn!(
@@ -1664,12 +1710,13 @@ pub async fn run_prompt_task(
log_stop_reason(&source, &stop_reason);
// Timeout triggers respawn in handle_prompt_result —
// session state will be discarded with the old agent.
let _ = result_tx.send(PromptResult {
send_prompt_result(
&result_tx,
agent,
source,
outcome: PromptOutcome::Timeout,
batch: requeue_batch_if_queue(&ctx, batch),
});
PromptOutcome::Timeout,
requeue_batch_if_queue(&ctx, batch),
);
}
Err(AcpError::AgentExited) => {
tracing::error!(
@@ -1678,12 +1725,13 @@ pub async fn run_prompt_task(
agent.index
);
agent.state.invalidate_all();
let _ = result_tx.send(PromptResult {
send_prompt_result(
&result_tx,
agent,
source,
outcome: PromptOutcome::AgentExited,
batch: requeue_batch_if_queue(&ctx, batch),
});
PromptOutcome::AgentExited,
requeue_batch_if_queue(&ctx, batch),
);
}
Err(e) => {
tracing::error!(
@@ -1691,12 +1739,13 @@ pub async fn run_prompt_task(
"cancel_with_cleanup error: {e} — invalidating session"
);
agent.state.invalidate(&source);
let _ = result_tx.send(PromptResult {
send_prompt_result(
&result_tx,
agent,
source,
outcome: PromptOutcome::Timeout,
batch: requeue_batch_if_queue(&ctx, batch),
});
PromptOutcome::Timeout,
requeue_batch_if_queue(&ctx, batch),
);
}
}
}
@@ -1707,12 +1756,13 @@ pub async fn run_prompt_task(
ctx.max_turn_duration.as_secs()
);
agent.state.invalidate_all();
let _ = result_tx.send(PromptResult {
send_prompt_result(
&result_tx,
agent,
source,
outcome: PromptOutcome::Timeout,
batch: requeue_batch_if_queue(&ctx, batch),
});
PromptOutcome::Timeout,
requeue_batch_if_queue(&ctx, batch),
);
}
Err(e) => {
tracing::error!(target: "pool::prompt", "session_prompt error: {e}");
@@ -1722,12 +1772,13 @@ pub async fn run_prompt_task(
if !matches!(e, AcpError::AgentError(_)) {
agent.state.invalidate(&source);
}
let _ = result_tx.send(PromptResult {
send_prompt_result(
&result_tx,
agent,
source,
outcome: PromptOutcome::Error(e),
batch: requeue_batch_if_queue(&ctx, batch),
});
PromptOutcome::Error(e),
requeue_batch_if_queue(&ctx, batch),
);
}
}
// _reaction_guard drops here → spawns clear_reactions for all exit paths.
@@ -3447,4 +3498,119 @@ mod tests {
}
// No observer to assert against — reaching here without panic is the test.
}
// ── steer_rx invariant tests ──────────────────────────────────────────
//
// These pin the `send_prompt_result` invariant: `steer_rx` is always
// `None` on any agent returned to the pool, regardless of which exit
// path fired.
//
// Test 1 (session-create-error path): installs a receiver, then calls
// `send_prompt_result` without the read loop running `take()` — simulating
// any early-return arm (e.g. session-create failure). The receiver must be
// cleared and the next `install_steer_rx` must not panic.
//
// Test 2 (post-read-loop path): receiver is already `None` (the read loop
// already consumed it via `take()`). `send_prompt_result` is idempotent —
// `steer_rx` stays `None` and the next `install_steer_rx` still does not
// panic.
/// After an early-return path (receiver installed but read loop never ran),
/// the returned agent's `steer_rx` is `None` and a subsequent
/// `install_steer_rx` does not panic.
#[tokio::test]
async fn test_send_prompt_result_clears_steer_rx_on_early_return() {
let acp = AcpClient::spawn("bash", &["-c".to_string(), "sleep 10".to_string()], &[])
.await
.expect("failed to spawn test agent");
let mut agent = OwnedAgent {
index: 0,
acp,
state: SessionState::default(),
model_capabilities: None,
desired_model: None,
model_overridden: false,
protocol_version: 2,
};
// Simulate dispatch: install a steer receiver (normally done by
// `dispatch_pending` before `run_prompt_task` is spawned).
let (_steer_tx, steer_rx) = tokio::sync::mpsc::channel::<SteerRequest>(1);
agent.acp.install_steer_rx(steer_rx);
// Simulate session-create error: early-return path calls
// `send_prompt_result` without the read loop ever running `take()`.
let (result_tx, mut result_rx) = tokio::sync::mpsc::unbounded_channel::<PromptResult>();
let source = PromptSource::Heartbeat;
send_prompt_result(
&result_tx,
agent,
source,
PromptOutcome::Error(AcpError::Protocol("simulated session-create error".into())),
None,
);
// Receive the PromptResult back from the channel.
let mut result = result_rx.recv().await.expect("PromptResult must be sent");
// steer_rx must be cleared even though the read loop never ran take().
assert!(
result.agent.acp.steer_rx_is_none(),
"steer_rx must be None after send_prompt_result on error path"
);
// The next dispatch can now install a fresh receiver without panicking.
let (_steer_tx2, steer_rx2) = tokio::sync::mpsc::channel::<SteerRequest>(1);
result.agent.acp.install_steer_rx(steer_rx2);
// Reaching here without a panic is the test.
}
/// After a successful prompt (read loop already consumed `steer_rx` via
/// `take()`), `send_prompt_result` is a no-op — `steer_rx` stays `None`
/// and the next `install_steer_rx` does not panic.
#[tokio::test]
async fn test_send_prompt_result_is_noop_when_steer_rx_already_consumed() {
let acp = AcpClient::spawn("bash", &["-c".to_string(), "sleep 10".to_string()], &[])
.await
.expect("failed to spawn test agent");
let agent = OwnedAgent {
index: 0,
acp,
state: SessionState::default(),
model_capabilities: None,
desired_model: None,
model_overridden: false,
protocol_version: 2,
};
// Simulate a completed turn: `steer_rx` was consumed by the read loop
// (`take()` was called), so it is already `None` when the turn ends.
assert!(
agent.acp.steer_rx_is_none(),
"precondition: steer_rx starts as None"
);
let (result_tx, mut result_rx) = tokio::sync::mpsc::unbounded_channel::<PromptResult>();
let source = PromptSource::Heartbeat;
send_prompt_result(
&result_tx,
agent,
source,
PromptOutcome::Ok(StopReason::EndTurn),
None,
);
let mut result = result_rx.recv().await.expect("PromptResult must be sent");
// Still None — clear_steer_rx on an already-None field is idempotent.
assert!(
result.agent.acp.steer_rx_is_none(),
"steer_rx must remain None after send_prompt_result on happy path"
);
// The next dispatch can install a fresh receiver without panicking.
let (_steer_tx, steer_rx) = tokio::sync::mpsc::channel::<SteerRequest>(1);
result.agent.acp.install_steer_rx(steer_rx);
// Reaching here without a panic is the test.
}
}