diff --git a/crates/buzz-agent/tests/fake_llm.rs b/crates/buzz-agent/tests/fake_llm.rs index cd242aed451..32b9f167de5 100644 --- a/crates/buzz-agent/tests/fake_llm.rs +++ b/crates/buzz-agent/tests/fake_llm.rs @@ -985,7 +985,17 @@ async fn steer_rejected_on_run_id_mismatch() { json!({"sessionId": sid, "prompt": [{"type":"text","text":"go"}]}), ) .await; - let _live_run = recv_active_run_id(&mut h).await; + + // Single ordered log from the prompt onward — nothing is discarded between + // the prompt send and the assertions (the activeRunId advert is read + // through the same log, so even pre-advert frames would be kept). + let mut log = FrameLog::new(); + let _live_run = log + .wait_until(&mut h, |v| { + v.get("method") == Some(&json!("session/update")) + && v["params"]["update"]["_meta"]["goose"]["activeRunId"].is_string() + }) + .await; let s_id = h .send( @@ -998,21 +1008,66 @@ async fn steer_rejected_on_run_id_mismatch() { ) .await; - let mut saw_reject = false; - for _ in 0..40 { - let v = h.recv().await; - if v["id"] == json!(s_id) { - assert_eq!( - v["error"]["code"], -32602, - "mismatched runId must be rejected" - ); - saw_reject = true; - } else if v["id"] == json!(p_id) { - // Turn finishes normally regardless of the rejected steer. - break; - } - } - assert!(saw_reject, "run-id mismatch was not rejected"); + // Wait for the steer response itself. The turn may finish before the + // rejection — its prompt response lands in the log and is matched there + // first, so the follow-up wait below can never hang on a second copy. + let v = log.wait_until(&mut h, |v| v["id"] == json!(s_id)).await; + assert_eq!( + v["error"]["code"], -32602, + "mismatched runId must be rejected" + ); + // Turn finishes normally regardless of the rejected steer. Its response + // may already be in the log — found, not re-read. + let response = log.wait_until(&mut h, |v| v["id"] == json!(p_id)).await; + assert_eq!( + response["result"]["stopReason"], "end_turn", + "turn must finish normally despite the rejected steer" + ); + h.shutdown().await; +} + +/// Deterministic regression (prompt-before-rejection): the turn's prompt +/// response is fully collected BEFORE the steer is even sent, so the steer +/// rejection necessarily arrives after it. The rejection must still be +/// invalid_params (no active run), and the prompt-response wait must be +/// satisfied from the existing log rather than blocking on another read. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn steer_rejection_after_prompt_response_uses_existing_frames() { + let url = spawn_fake_llm(vec![openai_text("all done")]).await; + let mut h = Harness::spawn(&url).await; + let sid = init_session(&mut h).await; + + let p_id = h + .send( + "session/prompt", + json!({"sessionId": sid, "prompt": [{"type":"text","text":"go"}]}), + ) + .await; + let mut log = FrameLog::new(); + let response = log.wait_until(&mut h, |v| v["id"] == json!(p_id)).await; + assert_eq!(response["result"]["stopReason"], "end_turn"); + + // Steer now targets a turn that already ended. + let s_id = h + .send( + "_goose/unstable/session/steer", + json!({ + "sessionId": sid, + "expectedRunId": "run_already_finished", + "prompt": [{"type":"text","text":"too late"}], + }), + ) + .await; + let v = log.wait_until(&mut h, |v| v["id"] == json!(s_id)).await; + assert_eq!( + v["error"]["code"], -32602, + "steer at an ended run must be rejected invalid_params" + ); + + // The prompt response is already in the log; this must return it without + // reading further (a second wait would hang to the recv timeout). + let again = log.wait_until(&mut h, |v| v["id"] == json!(p_id)).await; + assert_eq!(again, response, "prompt response must be found in the log"); h.shutdown().await; } @@ -1064,6 +1119,80 @@ fn is_usage_update(v: &Value) -> bool { && v["params"]["update"]["sessionUpdate"] == "usage_update" } +/// Ordered, accumulating collector for the agent's stdout stream. +/// +/// `recv_until`/`recv_until_with_drain` answer one question at a time and +/// either discard what they see or window-split it. That is exactly how the +/// steer/cancel races slipped through: a prompt response drained during the +/// cancel-ack wait was forgotten, so the follow-up wait for it hung to the +/// recv timeout; and frames read after an early prompt response were folded +/// into the same "before" bucket, letting a late usage_update falsely satisfy +/// ordering. +/// +/// `FrameLog` keeps every frame in arrival order and answers queries against +/// the whole history: `wait_until` matches already-collected frames BEFORE +/// reading more, so a response seen in an earlier window terminates the wait +/// immediately instead of blocking on a second copy that will never come. +struct FrameLog { + frames: Vec, +} + +impl FrameLog { + fn new() -> Self { + Self { frames: Vec::new() } + } + + /// Wait until a frame matching `pred` exists in the log, appending every + /// frame read along the way (bounded by `Harness::recv`'s timeout). + /// Already-collected frames are matched first. + async fn wait_until bool>( + &mut self, + h: &mut Harness, + mut pred: F, + ) -> Value { + if let Some(v) = self.frames.iter().find(|f| pred(f)) { + return v.clone(); + } + loop { + let v = h.recv().await; + let matched = pred(&v); + self.frames.push(v.clone()); + if matched { + return v; + } + } + } + + /// Index of the first frame matching `pred`, in arrival order. + fn first_index bool>(&self, pred: F) -> Option { + self.frames.iter().position(pred) + } +} + +/// Invariant: buzz-agent emits the turn's usage_update notification strictly +/// BEFORE the `session/prompt` response for that turn (buzz-acp's UsageTracker +/// consumes the notification while the turn is still in flight, before the +/// response triggers take_turn_usage()). +/// +/// Asserted over the ordered frame log: a usage_update must exist at an index +/// strictly less than the index of the FIRST frame whose id matches `p_id`. +/// Usage updates at or after the first prompt response do NOT satisfy the +/// invariant — the response has already gone out — and anchoring on the FIRST +/// response matters when the log spans multiple turns: usage arriving before a +/// later turn's response cannot retroactively fix ordering for the first. +fn assert_usage_strictly_before_prompt_response(log: &FrameLog, p_id: i64) { + let first_response = log + .first_index(|v| v["id"] == json!(p_id)) + .unwrap_or_else(|| panic!("no session/prompt response for id {p_id} in frame log")); + let usage_at = log.first_index(is_usage_update); + assert!( + matches!(usage_at, Some(u) if u < first_response), + "usage_update must arrive strictly before the FIRST session/prompt response \ + (usage_update index {usage_at:?}, response index {first_response}); frames: {:#?}", + log.frames, + ); +} + /// Collect every frame that arrives BEFORE the message matching `until_pred`, /// then return (frames_before, matching_frame). async fn recv_until_with_drain(h: &mut Harness, mut until_pred: F) -> (Vec, Value) @@ -1420,59 +1549,149 @@ async fn cancelled_turn_with_usage_emits_notification_before_response() { ) .await; - // Wait for the activeRunId advert (agent is live). - let _run_id = recv_active_run_id(&mut h).await; + // Wait for the activeRunId advert (agent is live) — collected, not + // discarded, so the usage_update racing the advert is never dropped. + let mut log = FrameLog::new(); + log.wait_until(&mut h, |v| { + v.get("method") == Some(&json!("session/update")) + && v["params"]["update"]["_meta"]["goose"]["activeRunId"].is_string() + }) + .await; // Wait for tool_call_update — proves round 1 LLM response is fully processed - // and tokens are captured before we send cancel. - h.recv_until(|v| { + // and tokens are captured before we send cancel. Every frame seen along the + // way stays in the ordered log and feeds the assertions below. + log.wait_until(&mut h, |v| { v.get("method") == Some(&json!("session/update")) && v["params"]["update"]["sessionUpdate"] == "tool_call_update" }) .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. + // Now send cancel and release the round-2 gate. Waiting for the cancel ack + // first is what makes this deterministic: the handler awaits + // `cancel_session` before replying, so once the ack lands the cancellation + // is registered and round 2 cannot beat it to the turn's end. let c_id = h.send("session/cancel", json!({"sessionId": sid})).await; + let cancel_ack = log.wait_until(&mut h, |v| v["id"] == json!(c_id)).await; + assert!( + cancel_ack.get("result").is_some() && cancel_ack.get("error").is_none(), + "session/cancel must be acknowledged with a success response" + ); 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) { - saw_usage = true; - if !saw_prompt_response { - saw_usage_before_prompt_response = true; - } - } else if v["id"] == json!(p_id) { - saw_prompt_response = true; - // The gate guarantees stopReason: cancelled — not a race-driven error. - assert_eq!( - v["result"]["stopReason"], "cancelled", - "turn must end with stopReason: cancelled" - ); - } - if saw_usage && saw_prompt_response && saw_cancel_ok { - break; - } - } - assert!(saw_cancel_ok, "session/cancel was not acknowledged"); - assert!( - saw_usage, - "expected usage_update notification for cancelled turn with observed tokens" + // Drain everything up to and including the prompt response. If the prompt + // response already landed in an earlier window it is FOUND here, not + // re-read — the previous unconditional second wait could hang to the recv + // timeout when the response raced ahead of the cancel ack. + let prompt_response = log.wait_until(&mut h, |v| v["id"] == json!(p_id)).await; + // The gate guarantees stopReason: cancelled — not a race-driven error. + assert_eq!( + prompt_response["result"]["stopReason"], "cancelled", + "turn must end with stopReason: cancelled" ); + // Ordering is asserted over the ordered log: a usage_update must sit + // strictly BEFORE the FIRST prompt response. A usage frame collected after + // an early prompt response cannot satisfy this. + assert_usage_strictly_before_prompt_response(&log, p_id); + h.shutdown().await; +} + +/// Deterministic regression (prompt-before-ack): cancel is sent AFTER the +/// turn has already completed, so the prompt response is necessarily already +/// in the frame log when the cancel ack arrives. The ack wait must not +/// discard it, and the follow-up prompt-response wait must be satisfied from +/// the log without a further read. Cancel on an idle session still acks +/// success (send on a dead cancel_tx is ignored). +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn cancel_after_prompt_response_uses_existing_frames() { + let url = spawn_fake_llm(vec![openai_text_with_usage("quick turn", 7, 3)]).await; + let mut h = Harness::spawn(&url).await; + let sid = init_session(&mut h).await; + + let p_id = h + .send( + "session/prompt", + json!({"sessionId": sid, "prompt": [{"type":"text","text":"go"}]}), + ) + .await; + let mut log = FrameLog::new(); + let response = log.wait_until(&mut h, |v| v["id"] == json!(p_id)).await; + assert_eq!(response["result"]["stopReason"], "end_turn"); + + // Late cancel: the turn is over, the response is collected, but the cancel + // must still be acknowledged with success. + let c_id = h.send("session/cancel", json!({"sessionId": sid})).await; + let ack = log.wait_until(&mut h, |v| v["id"] == json!(c_id)).await; assert!( - saw_usage_before_prompt_response, - "usage_update must arrive before the session/prompt response" + ack.get("result").is_some() && ack.get("error").is_none(), + "session/cancel on an idle session must still be acked with success" ); + // Found in the log — no further read (a second wait would hang). + let again = log.wait_until(&mut h, |v| v["id"] == json!(p_id)).await; + assert_eq!(again, response, "prompt response must be found in the log"); + assert_usage_strictly_before_prompt_response(&log, p_id); h.shutdown().await; } +/// Unit check of the ordering assertion itself: a usage_update that arrives +/// AFTER the first session/prompt response must be rejected, and one that +/// arrives before a LATER turn's response cannot retroactively fix ordering +/// for the first turn. +#[test] +fn usage_ordering_assertion_rejects_usage_after_first_prompt_response() { + let usage = json!({ + "method": "_goose/unstable/session/update", + "params": {"update": {"sessionUpdate": "usage_update"}} + }); + let resp = json!({"id": 42, "result": {"stopReason": "cancelled"}}); + + // usage after the first response → must be rejected. + let mut log = FrameLog::new(); + log.frames = vec![resp.clone(), usage.clone()]; + let r = std::panic::catch_unwind(|| { + let log = FrameLog { + frames: log.frames.clone(), + }; + assert_usage_strictly_before_prompt_response(&log, 42); + }); + assert!(r.is_err(), "usage after the first response must fail"); + + // usage before the response → must pass. + let mut log = FrameLog::new(); + log.frames = vec![usage.clone(), resp.clone()]; + assert_usage_strictly_before_prompt_response(&log, 42); + + // usage before a SECOND turn's response cannot fix the first turn's + // ordering: anchored on the FIRST response, this must be rejected. + let mut log = FrameLog::new(); + log.frames = vec![ + resp.clone(), + usage.clone(), + json!({"id": 43, "result": {"stopReason": "end_turn"}}), + ]; + let r = std::panic::catch_unwind(|| { + let log = FrameLog { + frames: log.frames.clone(), + }; + assert_usage_strictly_before_prompt_response(&log, 42); + }); + assert!( + r.is_err(), + "usage before a later response must not satisfy the first turn's ordering" + ); + + // no response for the id at all → must be rejected. + let mut log = FrameLog::new(); + log.frames = vec![usage.clone()]; + let r = std::panic::catch_unwind(|| { + let log = FrameLog { + frames: log.frames.clone(), + }; + assert_usage_strictly_before_prompt_response(&log, 42); + }); + assert!(r.is_err(), "missing prompt response must fail"); +} + /// A tool-call OpenAI response with a `usage` block. Used to capture tokens in /// round 1 before a cancel fires at the round boundary. fn openai_tool_call_with_usage(