fix(stella-core,stella-pipeline): close the sub-agent bracket on cancellation (#1954)

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<CommittedTally>,
     ) -> 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<String> {
     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<f64>,
+}
+
+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<AtomicUsize>) -> 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<CommittedTally>) -> 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<AtomicUsize>) -> 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<CommittedTally>,
+    /// 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> {

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<Vec<Result<CompletionResult, ProviderError>>>,
+    hang_reached: std::sync::Arc<tokio::sync::Notify>,
+}
+
+#[async_trait]
+impl Provider for HangAfterScript {
+    fn id(&self) -> &str {
+        "hanging"
+    }
+
+    async fn complete_ref(
+        &self,
+        _request: CompletionRequestRef<'_>,
+    ) -> Result<CompletionResult, ProviderError> {
+        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()
+    );
+}

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())