From 5ec55cd161ebb2e29b50e52ccf18cd965e189dd7 Mon Sep 17 00:00:00 2001 From: Mac Anderson Date: Thu, 6 Aug 2026 18:43:09 -0700 Subject: [PATCH] fix(stella-core,stella-pipeline): close the sub-agent bracket on cancellation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A caller that drops a sub-agent future mid-flight — a latency ceiling, a hard cancel — left `Started` open forever, so every ceiling-bearing caller had to forge its own `Finished` and could only guess `steps: 0`: the committed-call count lived inside the dropped turn. `CancelBracket` moves that obligation into the primitive. It is armed between `Started` and the normal `Finished`, and drop order does the sequencing — the turn future drops first, so the engine's cancel guard has already emitted the abandoned call's `UsageIncomplete { Cancelled }` envelope and `SettleChildOnDrop` has folded the money back before the bracket closes. It therefore reports only committed steps and cost, tallied by `child_sender` as each `StepUsage` crosses the boundary. The research stage drops its forged bracket accordingly; the next caller with a ceiling inherits the fix instead of repeating the bug. Closes #1954 --- crates/stella-core/src/subagent.rs | 134 +++++++++++-- crates/stella-core/src/subagent/tests.rs | 129 +++++++++++++ .../src/pipeline/research_stage.rs | 36 +--- .../src/pipeline/tests/research.rs | 177 ++++++++++++++++++ 4 files changed, 439 insertions(+), 37 deletions(-) diff --git a/crates/stella-core/src/subagent.rs b/crates/stella-core/src/subagent.rs index 2b060a15..710754ad 100644 --- a/crates/stella-core/src/subagent.rs +++ b/crates/stella-core/src/subagent.rs @@ -84,6 +84,12 @@ //! it is the metering record, and dropping it is exactly how child cost //! would vanish from `stella stats` and quietly falsify `$/resolved task`. //! +//! The bracket survives cancellation (#1954): a caller that drops the +//! future mid-flight — a latency ceiling, a hard cancel — still gets a +//! `Finished` carrying the committed step count and cost (`CancelBracket`), +//! after the engine's own drop guards have emitted the abandoned call's +//! `UsageIncomplete { Cancelled }` envelope and settled the money. +//! //! # Nesting //! //! [`SubAgentSpec::depth`] is checked against [`MAX_SUB_AGENT_DEPTH`] before @@ -564,11 +570,30 @@ impl Engine<'_> { depth: spec.depth, }, }); + // The committed tally, owned HERE rather than inside the child turn + // so the cancel bracket below can report it after the turn future is + // gone (#1954). Written by `child_sender` as each `StepUsage` passes + // the boundary. + let tally = Arc::new(CommittedTally::default()); + // Armed between `Started` and the normal `Finished`: a caller that + // drops this future mid-flight (a latency ceiling, a hard cancel) + // still owes the stream a balanced bracket, and only this frame can + // pay it — the dropped turn future cannot. See `CancelBracket`. + let mut bracket = CancelBracket { + events: events.clone(), + agent_id: spec.agent_id.clone(), + tally: tally.clone(), + armed: true, + }; let outcome = match refusal(spec, &carve) { Some(reason) => SubAgentOutcome::Refused { reason }, - None => self.run_child_turn(host, spec, carve, budget, events).await, + None => { + self.run_child_turn(host, spec, carve, budget, events, &tally) + .await + } }; + bracket.armed = false; let _ = events.send(AgentEvent::SubAgent { phase: SubAgentPhase::Finished { @@ -597,6 +622,7 @@ impl Engine<'_> { mut carve: BudgetGuard, budget: &mut BudgetGuard, events: &EventSender, + tally: &Arc, ) -> SubAgentOutcome { // Attribution is entered before anything the child could emit and // released by drop, so an unwind cannot leave the parent's later @@ -680,8 +706,7 @@ impl Engine<'_> { messages.push(CompletionMessage::user(spec.instruction.clone())); let seeded = messages.len(); - let steps = Arc::new(AtomicUsize::new(0)); - let child_events = child_sender(events.clone(), steps.clone()); + let child_events = child_sender(events.clone(), tally.clone()); // The carve is handed to the turn through a guard that settles it on // DROP, not on return (#1850). `settle_child` used to be a statement // after the await, so any exit that was not a return skipped it: a @@ -719,7 +744,7 @@ impl Engine<'_> { }); let absorbed_messages = messages.len().saturating_sub(seeded); - let steps = steps.load(Ordering::Relaxed); + let steps = tally.steps(); let build = |text: &str| { let (summary, truncated) = truncate_marked(text.trim(), spec.max_report_chars); SubAgentReport { @@ -786,16 +811,48 @@ fn refusal(spec: &SubAgentSpec, carve: &BudgetGuard) -> Option { None } +/// What a child has actually *committed*: one model call counted, and its +/// cost added, as each `StepUsage` crosses the boundary. +/// +/// It lives outside the child turn because that is the only place it survives +/// a cancel (#1954). When a caller drops the turn future the outcome never +/// exists, and this is the sole committed record [`CancelBracket`] can close +/// the bracket with — which is also why the two numbers travel as one type: +/// a bracket that reported a step count without the cost that produced it +/// would be half an answer. +#[derive(Default)] +struct CommittedTally { + steps: AtomicUsize, + cost_usd: Mutex, +} + +impl CommittedTally { + /// Record one committed model call. + fn observe(&self, cost_usd: f64) { + self.steps.fetch_add(1, Ordering::Relaxed); + *self.cost_usd.lock().unwrap_or_else(|p| p.into_inner()) += cost_usd; + } + + fn steps(&self) -> usize { + self.steps.load(Ordering::Relaxed) + } + + fn cost_usd(&self) -> f64 { + *self.cost_usd.lock().unwrap_or_else(|p| p.into_inner()) + } +} + /// The child's event sender: drops what must not cross ([`forwards_to_parent`]) -/// and counts committed model calls on the way past. +/// and tallies committed model calls — count and cost — on the way past. /// -/// Counting here rather than from the turn outcome is what makes `steps` -/// truthful on an abort too — `StepUsage` is emitted per committed call, so -/// a child that died on step 5 of 16 reports 5. -fn child_sender(parent: EventSender, steps: Arc) -> EventSender { +/// Tallying here rather than from the turn outcome is what makes the numbers +/// truthful on an abort too — `StepUsage` is emitted per committed call, so a +/// child that died on step 5 of 16 reports 5. See [`CommittedTally`] for why +/// it is also the only record that survives a cancel. +fn child_sender(parent: EventSender, tally: Arc) -> EventSender { EventSender::from_fn(move |event| { - if matches!(event, AgentEvent::StepUsage { .. }) { - steps.fetch_add(1, Ordering::Relaxed); + if let AgentEvent::StepUsage { cost_usd, .. } = &event { + tally.observe(*cost_usd); } if forwards_to_parent(&event) { parent.send(event) @@ -805,6 +862,61 @@ fn child_sender(parent: EventSender, steps: Arc) -> EventSender { }) } +/// Balances the `Started`/`Finished` bracket when a caller drops the +/// sub-agent future mid-flight — a latency ceiling, a hard cancel (#1954). +/// +/// The bracket contract ("delivered exactly once, on `Finished`") used to +/// hold only on paths that returned: a dropped future left `Started` open +/// forever, so every ceiling-bearing caller had to forge its own `Finished` — +/// and could only guess `steps: 0`, because the committed-call count lived +/// inside the dropped turn. Owning the close here, in the primitive, is the +/// same argument that moved the goal verifier onto [`Engine::run_sub_agent`]: +/// the next caller with a ceiling inherits the fix instead of repeating the +/// bug. +/// +/// Drop order does the sequencing: the in-flight turn future — declared +/// after this guard — drops first, so the engine's own `CancelUsageGuard` +/// has already emitted the `UsageIncomplete { Cancelled }` envelope for the +/// abandoned call and `SettleChildOnDrop` has already folded the money back +/// by the time this closes the bracket. This guard therefore reports only +/// what was **committed** ([`CommittedTally`], as `child_sender` recorded it); +/// the in-flight call's usage rides its own envelope, never a guess here. +struct CancelBracket { + events: EventSender, + agent_id: String, + tally: Arc, + /// True between `Started` and the normal `Finished`; the completion path + /// disarms before emitting its own bracket, so this never double-closes. + armed: bool, +} + +impl Drop for CancelBracket { + fn drop(&mut self) { + if !self.armed { + return; + } + let _ = self.events.send(AgentEvent::SubAgent { + phase: SubAgentPhase::Finished { + agent_id: self.agent_id.clone(), + status: SubAgentStatus::Incomplete, + summary: String::new(), + truncated: false, + cost_usd: self.tally.cost_usd(), + steps: self.tally.steps(), + // The transcript died with the future; 0 is the honest floor, + // not a claim that the child absorbed nothing. + absorbed_messages: 0, + reason: Some( + "cancelled: the caller dropped this sub-agent mid-flight; \ + committed steps and cost only — the abandoned call's usage \ + rides its own incomplete-usage envelope" + .to_string(), + ), + }, + }); + } +} + /// The last assistant text in a transcript, for salvaging an aborted child's /// work. Skips empty assistant turns (a step that only called tools). fn last_assistant_text(messages: &[CompletionMessage]) -> Option<&str> { diff --git a/crates/stella-core/src/subagent/tests.rs b/crates/stella-core/src/subagent/tests.rs index 14afe88d..cc2b3d90 100644 --- a/crates/stella-core/src/subagent/tests.rs +++ b/crates/stella-core/src/subagent/tests.rs @@ -1290,3 +1290,132 @@ fn the_ledger_accumulates_and_drains_to_zero() { assert!((drain_sub_agent_spend(&ledger) - 0.03).abs() < 1e-9); assert_eq!(drain_sub_agent_spend(&ledger), 0.0); } + +// ---- cancellation (#1954) -------------------------------------------- + +/// Serves its script, then hangs forever — and says so on `hang_reached`, +/// which is what lets a test cancel the child at a *deterministic* point +/// instead of racing a wall-clock timeout. +struct HangAfterScript { + script: Mutex>>, + hang_reached: std::sync::Arc, +} + +#[async_trait] +impl Provider for HangAfterScript { + fn id(&self) -> &str { + "hanging" + } + + async fn complete_ref( + &self, + _request: CompletionRequestRef<'_>, + ) -> Result { + let next = self.script.lock().unwrap().pop(); + match next { + Some(result) => result, + None => { + self.hang_reached.notify_one(); + std::future::pending().await + } + } + } +} + +/// #1954 witness: a caller that drops the sub-agent future mid-flight still +/// gets a **balanced** bracket whose `Finished` carries the committed step +/// count and cost, after the abandoned call's `UsageIncomplete { Cancelled }` +/// envelope — and the money still settles into the parent's guard. Before +/// `CancelBracket`, the `Started` bracket stayed open forever and every +/// ceiling-bearing caller had to forge a `Finished` it could only fill with +/// `steps: 0`. +#[tokio::test] +async fn a_cancelled_child_closes_its_bracket_with_committed_steps_and_cost() { + let parent_provider = ScriptedProvider::new(vec![]); + let hang_reached = std::sync::Arc::new(tokio::sync::Notify::new()); + // One committed step (a tool call), then the second model call hangs. + let child_provider = HangAfterScript { + script: Mutex::new(vec![Ok(tool_call_result("read_file", "c1", 0.002))]), + hang_reached: hang_reached.clone(), + }; + let tools = MixedTools::default(); + let parent = Engine::with_sleeper(&parent_provider, &tools, EngineConfig::default(), &NoSleep); + let mut budget = BudgetGuard::new(BudgetMode::Observed, None, None); + let (tx, mut rx) = mpsc::unbounded_channel(); + let spec = SubAgentSpec::read_only("search-1", "find it"); + + { + let fut = parent.run_sub_agent(SubAgentHost::new(&child_provider), &spec, &mut budget, &tx); + let mut fut = std::pin::pin!(fut); + tokio::select! { + _ = &mut fut => unreachable!("a hanging child cannot complete"), + _ = hang_reached.notified() => {} + } + // `fut` drops here: the cancel every latency-ceiling caller performs. + } + + let events = drain(&mut rx); + let started = events + .iter() + .filter(|e| { + matches!( + e, + AgentEvent::SubAgent { + phase: SubAgentPhase::Started { .. } + } + ) + }) + .count(); + let finished: Vec<_> = events + .iter() + .filter_map(|e| match e { + AgentEvent::SubAgent { + phase: phase @ SubAgentPhase::Finished { .. }, + } => Some(phase.clone()), + _ => None, + }) + .collect(); + assert_eq!(started, 1); + assert_eq!( + finished.len(), + 1, + "the bracket must close exactly once on a cancel: {events:?}" + ); + match &finished[0] { + SubAgentPhase::Finished { + status, + steps, + cost_usd, + reason, + .. + } => { + assert_eq!(*status, SubAgentStatus::Incomplete); + assert_eq!( + *steps, 1, + "the committed step count, not a forged zero (#1954)" + ); + assert!( + (*cost_usd - 0.002).abs() < 1e-9, + "the committed cost: {cost_usd}" + ); + let reason = reason.as_deref().unwrap_or_default(); + assert!(reason.contains("cancelled"), "the close says why: {reason}"); + } + SubAgentPhase::Started { .. } => unreachable!(), + } + assert!( + events.iter().any(|e| matches!( + e, + AgentEvent::UsageIncomplete { + reason: stella_protocol::UsageIncompleteReason::Cancelled, + .. + } + )), + "the abandoned in-flight call owes its envelope: {events:?}" + ); + assert!( + (budget.session_spent_usd() - 0.002).abs() < 1e-9, + "the committed spend still settles on the drop path: {}", + budget.session_spent_usd() + ); +} diff --git a/crates/stella-pipeline/src/pipeline/research_stage.rs b/crates/stella-pipeline/src/pipeline/research_stage.rs index 003384b8..e4c0e95a 100644 --- a/crates/stella-pipeline/src/pipeline/research_stage.rs +++ b/crates/stella-pipeline/src/pipeline/research_stage.rs @@ -27,7 +27,6 @@ use super::*; use stella_core::subagent::{SubAgentHost, SubAgentOutcome, SubAgentSpec}; -use stella_protocol::{SubAgentPhase, SubAgentStatus}; use crate::candidate_fanout::FanOutBudget; use crate::research::{ @@ -119,10 +118,13 @@ impl Pipeline<'_> { }; // The ceiling is per child, INSIDE the future, so a timed-out // child settles its spend through the sub-agent primitive's - // drop guard and the stage still returns — research degrading + // drop guards and the stage still returns — research degrading // to fewer findings must never wedge the turn. The dropped - // child's `Finished` bracket is emitted here, because the - // cancelled future can no longer balance its own `Started`. + // child's stream stays whole without help here (#1954): the + // primitive's `CancelBracket` closes the `Started`/`Finished` + // bracket with the committed step count and cost, and the + // engine's own cancel guard emits the abandoned call's + // `UsageIncomplete { Cancelled }` envelope. let outcome = tokio::time::timeout( ceiling, engine.run_sub_agent_with_sender( @@ -140,28 +142,10 @@ impl Pipeline<'_> { answer: report.summary, }) } - // Refusals, aborts, and empty answers are missing - // findings, not errors — partial work is not evidence - // worth planning on. - Ok(_) => None, - Err(_elapsed) => { - self.emit(AgentEvent::SubAgent { - phase: SubAgentPhase::Finished { - agent_id: spec.agent_id.clone(), - status: SubAgentStatus::Incomplete, - summary: String::new(), - truncated: false, - cost_usd: child_budget.session_spent_usd(), - steps: 0, - absorbed_messages: 0, - reason: Some(format!( - "research latency ceiling ({}s) elapsed", - ceiling.as_secs() - )), - }, - }); - None - } + // Refusals, aborts, empty answers, and a child past the + // ceiling are missing findings, not errors — partial work + // is not evidence worth planning on. + Ok(_) | Err(_) => None, }; fan.settle(&child_budget); (finding, child_budget.session_spent_usd()) diff --git a/crates/stella-pipeline/src/pipeline/tests/research.rs b/crates/stella-pipeline/src/pipeline/tests/research.rs index 002f3e41..b1de3860 100644 --- a/crates/stella-pipeline/src/pipeline/tests/research.rs +++ b/crates/stella-pipeline/src/pipeline/tests/research.rs @@ -227,6 +227,183 @@ async fn no_questions_means_no_stage_no_sub_agents_no_section() { assert!(!planner_prompt(&shapes).contains("## Research findings")); } +/// One scripted entry of [`HangTailProvider`]: serve a result, or hang. +enum HangScript { + Serve(CompletionResult), + Hang, +} + +/// Serves its script in order; a [`HangScript::Hang`] entry parks that call +/// forever (the research latency ceiling is what ends it), and later calls +/// keep serving the rest of the script — so the run can continue past the +/// cancelled child. +struct HangTailProvider { + script: TokioMutex>, +} + +#[async_trait] +impl Provider for HangTailProvider { + fn id(&self) -> &str { + "hang-tail" + } + + async fn complete_ref( + &self, + _req: CompletionRequestRef<'_>, + ) -> Result { + let next = self.script.lock().await.pop_front(); + match next { + Some(HangScript::Serve(result)) => Ok(result), + Some(HangScript::Hang) => std::future::pending().await, + None => Err(ProviderError::Terminal("script exhausted".into())), + } + } +} + +/// #1954's witness, verbatim: a research child that never answers within the +/// ceiling produces a **balanced** Started/Finished bracket, a +/// `UsageIncomplete` with reason `cancelled`, and `Finished.steps` equal to +/// the child's committed `StepUsage` count — while the turn itself degrades +/// to a missing finding and completes. Before the primitive owned the +/// cancel bracket, the stage forged `Finished { steps: 0 }` and this failed. +#[tokio::test] +async fn a_child_past_the_ceiling_closes_its_bracket_with_committed_steps() { + // Sequence: triage (one question) → child call 1 commits a tool step → + // child call 2 hangs until the ceiling cancels it → plan → close-out. + let mut committed_step = text_result(""); + committed_step.tool_calls = vec![ToolCall { + call_id: "r1".into(), + name: "read_file".into(), + input: serde_json::json!({ "path": "src/lib.rs" }), + }]; + committed_step.cost_usd = 0.002; + let provider = HangTailProvider { + script: TokioMutex::new(VecDeque::from([ + HangScript::Serve(text_result( + "CLASS: multi\nWITNESS: yes\nVERIFIER: yes\nRESEARCH: Which module owns retries?", + )), + HangScript::Serve(committed_step), + HangScript::Hang, + HangScript::Serve(text_result(r#"["update retry.rs"]"#)), + HangScript::Serve(text_result("PLAN COMPLETE: done.")), + ])), + }; + let resolver = OneHangProvider(&provider); + let runner = ScriptedRunner::new(vec![false, true, true], "@@ -1 +1 @@\n-old\n+new"); + let tools = EmptyTools; + let recall = NoContextRecall; + let repo = NoRepoStructure; + let repo_status = NoRepoStatus; + let approvals = AutoApproveGate; + let sleeper = NoopSleeper; + let router = router(); + let (tx, mut rx) = mpsc::unbounded_channel(); + let pipeline = Pipeline::new( + PipelinePorts { + router: &router, + providers: &resolver, + tools: &tools, + recall: &recall, + repo: &repo, + repo_status: &repo_status, + touches: &NoFileTouches, + diagnostics: &runner, + tests: &runner, + lint: None, + mutation: None, + coverage: None, + approvals: &approvals, + sleeper: &sleeper, + hooks: None, + candidate_workspaces: None, + mcp_prefetch: None, + steering: None, + }, + tx, + PipelineConfig { + test_command: Some("cargo test -p x".into()), + diff_diagnostic: Some(DiagnosticInvocation::GitDiff), + research_latency_ceiling: std::time::Duration::from_millis(200), + ..PipelineConfig::default() + }, + ); + let mut messages = vec![CompletionMessage::system("sys")]; + let mut budget = BudgetGuard::new(BudgetMode::Off, None, None); + let outcome = pipeline + .run( + "Refactor the retry layer end to end", + &mut messages, + &mut budget, + ) + .await + .expect("run succeeds"); + let events = drain(&mut rx); + + assert_eq!( + outcome.status, + PipelineStatus::Completed, + "a cancelled child degrades to a missing finding, never a wedged turn" + ); + let started = events + .iter() + .filter(|e| { + matches!( + e, + AgentEvent::SubAgent { + phase: stella_protocol::SubAgentPhase::Started { .. } + } + ) + }) + .count(); + let finished: Vec<_> = events + .iter() + .filter_map(|e| match e { + AgentEvent::SubAgent { + phase: + stella_protocol::SubAgentPhase::Finished { + steps, + status, + reason, + .. + }, + } => Some((*steps, *status, reason.clone())), + _ => None, + }) + .collect(); + assert_eq!(started, 1); + assert_eq!(finished.len(), 1, "a balanced bracket: {events:?}"); + let (steps, status, reason) = &finished[0]; + assert_eq!( + *steps, 1, + "Finished.steps is the committed StepUsage count, not a forged zero" + ); + assert_eq!(*status, stella_protocol::SubAgentStatus::Incomplete); + assert!( + reason.as_deref().unwrap_or_default().contains("cancelled"), + "the close says why: {reason:?}" + ); + assert!( + events.iter().any(|e| matches!( + e, + AgentEvent::UsageIncomplete { + reason: stella_protocol::UsageIncompleteReason::Cancelled, + .. + } + )), + "the abandoned in-flight call owes its envelope: {events:?}" + ); +} + +/// A resolver over the hanging double — [`OneProvider`] is typed to the +/// scripted one, and the harness needs exactly the same everything-is-this +/// behavior here. +struct OneHangProvider<'p>(&'p HangTailProvider); +impl ProviderResolver for OneHangProvider<'_> { + fn provider_for(&self, _model: &ModelRef) -> Option<&dyn Provider> { + Some(self.0) + } +} + /// A research round that produces nothing usable — children answering empty — /// degrades to exactly the no-research planner prompt: the stage may not /// leave a half-empty section behind, and the turn still completes.