diff --git a/crates/buzz-agent/tests/fake_llm.rs b/crates/buzz-agent/tests/fake_llm.rs index 4253ef329c..88050b5b1e 100644 --- a/crates/buzz-agent/tests/fake_llm.rs +++ b/crates/buzz-agent/tests/fake_llm.rs @@ -1266,7 +1266,7 @@ async fn cancelled_turn_with_usage_emits_notification_before_response() { // in-flight TCP request can resolve. The queue is empty for round 2, so the // agent receives the fallback "no canned response" body which it treats as // an LLM error; the cancel check at the round boundary fires first because - // the gate is only released after cancel is enqueued. + // the gate is only released after cancel is acknowledged. let responses = vec![openai_tool_call_with_usage( "call_cancel_test", "fake__noop", @@ -1346,20 +1346,25 @@ async fn cancelled_turn_with_usage_emits_notification_before_response() { }) .await; - // Now send cancel and release the round-2 gate. Cancel is enqueued before - // round 2 can respond, so the turn exits with stopReason: cancelled. + // Wait for the cancel acknowledgement before releasing the round-2 gate. + // Sending only proves the request was written to stdin; the acknowledgement + // proves the agent processed it before round 2 can respond. let c_id = h.send("session/cancel", json!({"sessionId": sid})).await; + let (frames_before_cancel_ack, _) = + recv_until_with_drain(&mut h, |v| v["id"] == json!(c_id)).await; let _ = gate_tx.send(()); // unblock round 2 let mut saw_usage_before_prompt_response = false; let mut saw_usage = false; - let mut saw_cancel_ok = false; let mut saw_prompt_response = false; - for _ in 0..40 { - let v = h.recv().await; - if v["id"] == json!(c_id) { - saw_cancel_ok = true; - } else if is_usage_update(&v) { + let mut pending_frames = VecDeque::from(frames_before_cancel_ack); + let frame_budget = 40 + pending_frames.len(); + for _ in 0..frame_budget { + let v = match pending_frames.pop_front() { + Some(v) => v, + None => h.recv().await, + }; + if is_usage_update(&v) { saw_usage = true; if !saw_prompt_response { saw_usage_before_prompt_response = true; @@ -1372,11 +1377,10 @@ async fn cancelled_turn_with_usage_emits_notification_before_response() { "turn must end with stopReason: cancelled" ); } - if saw_usage && saw_prompt_response && saw_cancel_ok { + if saw_usage && saw_prompt_response { break; } } - assert!(saw_cancel_ok, "session/cancel was not acknowledged"); assert!( saw_usage, "expected usage_update notification for cancelled turn with observed tokens"