heartbit_core/agent/flow/
ctx.rs

1//! [`WorkflowCtx`] — the shared run context threaded into every combinator.
2//!
3//! Holds the run-wide concurrency cap (a single [`Semaphore`] shared across the
4//! whole run), the runaway agent backstop, a cancellation token, an optional
5//! event sink, and the "default phase" applied to subsequently-issued agents.
6//! It is cheap to [`Clone`] (an `Arc` inside) so combinator thunks can capture
7//! their own handle.
8
9use std::sync::Arc;
10use std::sync::atomic::{AtomicU64, Ordering};
11use std::sync::{Mutex, RwLock};
12
13use tokio::sync::Semaphore;
14use tokio_util::sync::CancellationToken;
15
16use crate::error::Error;
17use crate::llm::BoxedProvider;
18
19use super::budget::Budget;
20use super::event::{OnWorkflowEvent, WorkflowEvent};
21use super::journal::{ResumeMode, RunJournal};
22
23/// Host-supplied provider construction for per-call model selection: maps a
24/// model ROLE ("fast", "frontier") or a raw model id to a ready provider.
25/// String-keyed on purpose — the host decides what names mean.
26pub type ProviderFactory = dyn Fn(&str) -> Result<Arc<BoxedProvider>, Error> + Send + Sync;
27
28/// Hard upper bound on concurrently-running agents, matching Claude Code's
29/// `min(16, cores - 2)`.
30const MAX_CONCURRENCY_CAP: usize = 16;
31
32/// Default total-agent runaway backstop (Claude Code uses 1000 per run).
33const DEFAULT_MAX_AGENTS: u64 = 1000;
34
35/// Default concurrency cap: `min(16, available_parallelism - 2)`, at least 1.
36///
37/// Mirrors the idiom in [`crate::agent::batch`] but clamps to the workflow cap.
38pub(crate) fn default_concurrency() -> usize {
39    std::thread::available_parallelism()
40        .map(|n| n.get())
41        .unwrap_or(2)
42        .saturating_sub(2)
43        .clamp(1, MAX_CONCURRENCY_CAP)
44}
45
46/// A run-level limit breach that must halt the whole run, distinct from an
47/// agent-domain failure. When one occurs inside a combinator the breach is
48/// recorded (set-once) and the run-wide cancellation token is fired, so the
49/// breach survives the combinator's `Err -> None` collapse and the caller can
50/// detect it via [`WorkflowCtx::control_breach`].
51#[derive(Debug, Clone, Copy, PartialEq, Eq)]
52#[non_exhaustive]
53pub enum ControlBreach {
54    /// The shared token budget was exhausted.
55    Budget {
56        /// Weighted cost recorded when admission was refused.
57        used: u64,
58        /// The configured ceiling.
59        limit: u64,
60    },
61    /// The runaway agent backstop was hit.
62    AgentBackstop {
63        /// The configured maximum agents per run.
64        limit: u64,
65    },
66}
67
68/// Shared, cheaply-cloneable workflow run context.
69#[derive(Clone)]
70pub struct WorkflowCtx {
71    inner: Arc<CtxInner>,
72}
73
74struct CtxInner {
75    /// Type-erased provider shared by every agent leaf.
76    provider: Arc<BoxedProvider>,
77    /// Optional default tool set wired into every agent leaf that does not
78    /// supply its own per-call tools. Shared with nested workflows.
79    base_tools: Option<Arc<Vec<Arc<dyn crate::tool::Tool>>>>,
80    /// Global concurrency limiter; permits acquired only at the agent() leaf.
81    sem: Arc<Semaphore>,
82    /// Monotonic count of agents ever issued (runaway backstop; never
83    /// decremented). `Arc` so a nested workflow shares the same counter.
84    spawned: Arc<AtomicU64>,
85    /// Maximum total agents per run.
86    max_agents: u64,
87    /// Shared hard-ceiling token-equivalent spend pool.
88    budget: Budget,
89    /// First run-level limit breach, if any (set-once). `Arc` so a nested
90    /// workflow's breach is visible on the parent ctx.
91    control: Arc<Mutex<Option<ControlBreach>>>,
92    /// Optional resume journal: memoizes agent outputs for deterministic replay.
93    journal: Option<Arc<RunJournal>>,
94    /// Optional workflow event sink.
95    events: Option<Arc<OnWorkflowEvent>>,
96    /// Optional [`AgentEvent`](crate::agent::events::AgentEvent) sink wired
97    /// into every inner [`AgentRunner`](crate::agent::AgentRunner) the leaves
98    /// build — surfaces recipe-internal tool calls / LLM responses to a host
99    /// (e.g. the TUI trace). Distinct from `events`, which carries the
100    /// workflow-level lifecycle.
101    agent_events: Option<Arc<crate::agent::events::OnEvent>>,
102    /// Optional per-call model resolution (see [`ProviderFactory`]). Absent →
103    /// `AgentCall::model()` degrades to the shared provider with a log line.
104    provider_factory: Option<Arc<ProviderFactory>>,
105    /// Optional workspace (git repo root) — required by leaves that ask for
106    /// [`Isolation::Worktree`](super::agent::Isolation). Shared with nested
107    /// workflows.
108    workspace: Option<std::path::PathBuf>,
109    /// Cooperative cancellation (pause/stop); P1 only races it to `Ok(None)`.
110    cancel: CancellationToken,
111    /// Default phase for subsequently-issued agents. `std` lock: never held
112    /// across `.await` (snapshotted at `AgentCall` construction). NOT shared
113    /// with a nested workflow — each level groups its own agents.
114    default_phase: RwLock<Option<Arc<str>>>,
115    /// Nesting depth: 0 for a top-level run, 1 for a `workflow()` child. A
116    /// second level is rejected by [`WorkflowCtx::nested`].
117    depth: u8,
118}
119
120impl std::fmt::Debug for WorkflowCtx {
121    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
122        f.debug_struct("WorkflowCtx")
123            .field("max_agents", &self.inner.max_agents)
124            .field("permits", &self.inner.sem.available_permits())
125            .field("spawned", &self.inner.spawned.load(Ordering::Relaxed))
126            .field("budget_spent", &self.inner.budget.spent())
127            .field("budget_total", &self.inner.budget.total())
128            .finish()
129    }
130}
131
132impl WorkflowCtx {
133    /// Start building a context from a type-erased provider.
134    pub fn builder(provider: Arc<BoxedProvider>) -> WorkflowCtxBuilder {
135        WorkflowCtxBuilder {
136            provider,
137            base_tools: None,
138            max_concurrency: None,
139            max_agents: None,
140            budget: None,
141            journal: None,
142            events: None,
143            agent_events: None,
144            provider_factory: None,
145            workspace: None,
146            cancel: None,
147        }
148    }
149
150    /// Whether the run has been cancelled.
151    pub fn is_cancelled(&self) -> bool {
152        self.inner.cancel.is_cancelled()
153    }
154
155    /// Stop the run: fires the shared cancellation token so in-flight agents
156    /// wind down at their next leaf cancel-check (and combinator tasks abort).
157    /// Idempotent. Restarting a stopped run is done by re-running with a resume
158    /// [`journal`](WorkflowCtxBuilder::journal): completed agents replay, the
159    /// stopped ones run live.
160    pub fn stop(&self) {
161        self.inner.cancel.cancel();
162    }
163
164    /// A clone of the run's cancellation token (for external pause/stop wiring).
165    pub fn cancellation_token(&self) -> CancellationToken {
166        self.inner.cancel.clone()
167    }
168
169    /// Total agents issued so far (monotonic; includes rejected admissions).
170    pub fn spawned(&self) -> u64 {
171        self.inner.spawned.load(Ordering::Relaxed)
172    }
173
174    /// The configured runaway agent backstop.
175    pub fn max_agents(&self) -> u64 {
176        self.inner.max_agents
177    }
178
179    /// The shared token-equivalent budget pool.
180    pub fn budget(&self) -> &Budget {
181        &self.inner.budget
182    }
183
184    /// Remaining budget headroom, or `u64::MAX` if unbounded. Drives
185    /// `loop-until-budget`: `while ctx.remaining() > threshold { … }`.
186    pub fn remaining(&self) -> u64 {
187        self.inner.budget.remaining()
188    }
189
190    /// The first run-level limit breach, if one has occurred.
191    pub fn control_breach(&self) -> Option<ControlBreach> {
192        self.inner.control.lock().ok().and_then(|g| *g)
193    }
194
195    // ----- pub(crate) seams consumed by the agent() leaf (slice 5) -----
196
197    /// A clone of the shared, type-erased provider for building an agent leaf.
198    pub(crate) fn provider(&self) -> Arc<BoxedProvider> {
199        Arc::clone(&self.inner.provider)
200    }
201
202    /// The ctx's default tool set (cloned `Vec`), if any. Used by an agent leaf
203    /// that supplies no per-call tools.
204    pub(crate) fn agent_events(&self) -> Option<Arc<crate::agent::events::OnEvent>> {
205        self.inner.agent_events.clone()
206    }
207
208    pub(crate) fn provider_factory(&self) -> Option<Arc<ProviderFactory>> {
209        self.inner.provider_factory.clone()
210    }
211
212    pub(crate) fn workspace(&self) -> Option<std::path::PathBuf> {
213        self.inner.workspace.clone()
214    }
215
216    pub(crate) fn base_tools(&self) -> Option<Vec<Arc<dyn crate::tool::Tool>>> {
217        self.inner.base_tools.as_ref().map(|t| t.as_ref().clone())
218    }
219
220    /// A clone of the global concurrency limiter; acquired only at the leaf.
221    pub(crate) fn semaphore(&self) -> Arc<Semaphore> {
222        Arc::clone(&self.inner.sem)
223    }
224
225    /// Borrow the run's cancellation token (for the leaf's `select!` race).
226    pub(crate) fn cancel_token(&self) -> &CancellationToken {
227        &self.inner.cancel
228    }
229
230    /// Reserve one slot against the runaway backstop. Monotonic: a rejected
231    /// admission still counts (it is a backstop, not a live gauge). On breach,
232    /// records the control breach and fires run-wide cancellation.
233    pub(crate) fn register_agent(&self) -> Result<(), Error> {
234        let prior = self.inner.spawned.fetch_add(1, Ordering::Relaxed);
235        if prior >= self.inner.max_agents {
236            let limit = self.inner.max_agents;
237            self.record_breach(ControlBreach::AgentBackstop { limit });
238            return Err(Error::AgentBudgetExceeded { limit });
239        }
240        Ok(())
241    }
242
243    /// Admit one agent against the shared budget. On exhaustion, records the
244    /// control breach, fires run-wide cancellation, and returns the error.
245    pub(crate) fn admit_budget(&self) -> Result<(), Error> {
246        match self.inner.budget.check_admit() {
247            Ok(()) => Ok(()),
248            Err(err) => {
249                if let Error::BudgetExceeded { used, limit } = err {
250                    self.record_breach(ControlBreach::Budget { used, limit });
251                }
252                Err(err)
253            }
254        }
255    }
256
257    /// Record one agent's completed cost against the shared budget.
258    pub(crate) fn record_spend(&self, usage: &crate::llm::types::TokenUsage) {
259        self.inner.budget.record(usage);
260    }
261
262    /// A clone of the resume journal handle, if journaling is enabled.
263    pub(crate) fn journal_arc(&self) -> Option<Arc<RunJournal>> {
264        self.inner.journal.clone()
265    }
266
267    /// Create a one-level-deeper child context that **shares** this run's
268    /// provider, concurrency limiter, runaway backstop counter, budget pool,
269    /// breach state, journal, event sink, and cancellation token — but gets a
270    /// fresh default-phase. Rejects a second level with [`Error::Config`].
271    pub(crate) fn nested(&self) -> Result<WorkflowCtx, Error> {
272        if self.inner.depth >= 1 {
273            return Err(Error::Config(
274                "workflow() may nest only one level deep".into(),
275            ));
276        }
277        Ok(WorkflowCtx {
278            inner: Arc::new(CtxInner {
279                provider: Arc::clone(&self.inner.provider),
280                base_tools: self.inner.base_tools.clone(),
281                sem: Arc::clone(&self.inner.sem),
282                spawned: Arc::clone(&self.inner.spawned),
283                max_agents: self.inner.max_agents,
284                budget: self.inner.budget.clone(),
285                control: Arc::clone(&self.inner.control),
286                journal: self.inner.journal.clone(),
287                events: self.inner.events.clone(),
288                agent_events: self.inner.agent_events.clone(),
289                provider_factory: self.inner.provider_factory.clone(),
290                workspace: self.inner.workspace.clone(),
291                cancel: self.inner.cancel.clone(),
292                default_phase: RwLock::new(None),
293                depth: self.inner.depth + 1,
294            }),
295        })
296    }
297
298    /// Record a run-level breach (set-once) and fire run-wide cancellation so
299    /// in-flight combinator tasks wind down. Idempotent on the stored breach.
300    pub(crate) fn record_breach(&self, breach: ControlBreach) {
301        if let Ok(mut guard) = self.inner.control.lock()
302            && guard.is_none()
303        {
304            *guard = Some(breach);
305        }
306        self.inner.cancel.cancel();
307    }
308
309    /// Emit a workflow event if a sink is installed.
310    pub(crate) fn emit(&self, event: WorkflowEvent) {
311        if let Some(cb) = &self.inner.events {
312            cb(event);
313        }
314    }
315
316    /// Snapshot the current default phase (the phase newly-issued agents adopt).
317    pub(crate) fn current_phase(&self) -> Option<Arc<str>> {
318        self.inner.default_phase.read().ok().and_then(|g| g.clone())
319    }
320
321    /// Replace the default phase, returning the prior value (for RAII restore).
322    pub(crate) fn swap_default_phase(&self, next: Option<Arc<str>>) -> Option<Arc<str>> {
323        match self.inner.default_phase.write() {
324            Ok(mut g) => std::mem::replace(&mut *g, next),
325            Err(_) => None,
326        }
327    }
328}
329
330/// Builder for [`WorkflowCtx`].
331pub struct WorkflowCtxBuilder {
332    provider: Arc<BoxedProvider>,
333    base_tools: Option<Arc<Vec<Arc<dyn crate::tool::Tool>>>>,
334    max_concurrency: Option<usize>,
335    max_agents: Option<u64>,
336    budget: Option<Budget>,
337    journal: Option<Arc<RunJournal>>,
338    events: Option<Arc<OnWorkflowEvent>>,
339    agent_events: Option<Arc<crate::agent::events::OnEvent>>,
340    provider_factory: Option<Arc<ProviderFactory>>,
341    workspace: Option<std::path::PathBuf>,
342    cancel: Option<CancellationToken>,
343}
344
345impl WorkflowCtxBuilder {
346    /// Set the default tool set wired into every agent leaf that does not supply
347    /// its own per-call tools (see [`AgentCall::tools`](super::agent::AgentCall::tools)).
348    pub fn base_tools(mut self, tools: Vec<Arc<dyn crate::tool::Tool>>) -> Self {
349        self.base_tools = Some(Arc::new(tools));
350        self
351    }
352
353    /// Override the concurrency cap (must be >= 1).
354    pub fn max_concurrency(mut self, n: usize) -> Self {
355        self.max_concurrency = Some(n);
356        self
357    }
358
359    /// Override the runaway agent backstop (must be >= 1).
360    pub fn max_agents(mut self, n: u64) -> Self {
361        self.max_agents = Some(n);
362        self
363    }
364
365    /// Set a hard token-equivalent budget ceiling for the run (`0` ⇒ unbounded).
366    /// Default is unbounded.
367    pub fn budget(mut self, total: u64) -> Self {
368        self.budget = Some(Budget::with_total(total));
369        self
370    }
371
372    /// Share an existing [`Budget`] pool (e.g. a parent run's, for nested
373    /// workflows). This and [`budget`](Self::budget) write the same slot, so
374    /// whichever is called *last* wins (last-write-wins, not unconditional
375    /// precedence).
376    pub fn budget_pool(mut self, budget: Budget) -> Self {
377        self.budget = Some(budget);
378        self
379    }
380
381    /// Enable deterministic resume by journaling agent outputs to `path` (a
382    /// `.jsonl` file). [`ResumeMode::Fresh`] starts a new journal;
383    /// [`ResumeMode::Resume`] replays matching calls from an existing one.
384    /// Returns an error if the journal file cannot be opened.
385    ///
386    /// See [`RunJournal`] for the soundness scope (return values, not side
387    /// effects; single process; fixed provider).
388    pub fn journal(
389        mut self,
390        path: impl AsRef<std::path::Path>,
391        mode: ResumeMode,
392    ) -> Result<Self, Error> {
393        self.journal = Some(Arc::new(RunJournal::open(path.as_ref(), mode)?));
394        Ok(self)
395    }
396
397    /// Share an existing [`RunJournal`] handle (e.g. a parent run's).
398    pub fn journal_handle(mut self, journal: Arc<RunJournal>) -> Self {
399        self.journal = Some(journal);
400        self
401    }
402
403    /// Install a workflow event sink.
404    pub fn on_event(mut self, callback: Arc<OnWorkflowEvent>) -> Self {
405        self.events = Some(callback);
406        self
407    }
408
409    /// Install an [`AgentEvent`](crate::agent::events::AgentEvent) sink: every
410    /// inner `AgentRunner` built by the agent leaves forwards its full event
411    /// stream (tool calls, LLM responses, run lifecycle) here. Lets a host
412    /// (e.g. the TUI trace) observe recipe-internal activity that would
413    /// otherwise stay invisible behind the tool boundary.
414    pub fn on_agent_event(mut self, callback: Arc<crate::agent::events::OnEvent>) -> Self {
415        self.agent_events = Some(callback);
416        self
417    }
418
419    /// Install the per-call model resolver (see [`ProviderFactory`]): lets
420    /// `agent(&ctx, …).model("fast")` run a leaf on a different provider.
421    pub fn provider_factory(mut self, factory: Arc<ProviderFactory>) -> Self {
422        self.provider_factory = Some(factory);
423        self
424    }
425
426    /// Set the workspace (git repo root) — required for leaves that ask for
427    /// [`Isolation::Worktree`](super::agent::Isolation).
428    pub fn workspace(mut self, root: impl Into<std::path::PathBuf>) -> Self {
429        self.workspace = Some(root.into());
430        self
431    }
432
433    /// Provide an external cancellation token (defaults to a fresh one).
434    pub fn cancellation_token(mut self, token: CancellationToken) -> Self {
435        self.cancel = Some(token);
436        self
437    }
438
439    /// Build the context.
440    ///
441    /// Rejects a zero concurrency cap or zero agent backstop (consistent with
442    /// the zero-rejection in [`crate::agent::batch`] and the workflow agents).
443    pub fn build(self) -> Result<WorkflowCtx, Error> {
444        let max_concurrency = self.max_concurrency.unwrap_or_else(default_concurrency);
445        if max_concurrency == 0 {
446            return Err(Error::Config(
447                "WorkflowCtx max_concurrency must be at least 1".into(),
448            ));
449        }
450        let max_agents = self.max_agents.unwrap_or(DEFAULT_MAX_AGENTS);
451        if max_agents == 0 {
452            return Err(Error::Config(
453                "WorkflowCtx max_agents must be at least 1".into(),
454            ));
455        }
456        Ok(WorkflowCtx {
457            inner: Arc::new(CtxInner {
458                provider: self.provider,
459                base_tools: self.base_tools,
460                sem: Arc::new(Semaphore::new(max_concurrency)),
461                spawned: Arc::new(AtomicU64::new(0)),
462                max_agents,
463                budget: self.budget.unwrap_or_default(),
464                control: Arc::new(Mutex::new(None)),
465                journal: self.journal,
466                events: self.events,
467                agent_events: self.agent_events,
468                provider_factory: self.provider_factory,
469                workspace: self.workspace,
470                cancel: self.cancel.unwrap_or_default(),
471                default_phase: RwLock::new(None),
472                depth: 0,
473            }),
474        })
475    }
476}
477
478#[cfg(test)]
479mod tests {
480    use super::*;
481    use crate::agent::test_helpers::MockProvider;
482
483    fn provider() -> Arc<BoxedProvider> {
484        Arc::new(BoxedProvider::new(MockProvider::new(vec![])))
485    }
486
487    #[test]
488    fn default_concurrency_in_range() {
489        let c = default_concurrency();
490        assert!((1..=MAX_CONCURRENCY_CAP).contains(&c), "got {c}");
491    }
492
493    #[test]
494    fn builder_rejects_zero_concurrency() {
495        let result = WorkflowCtx::builder(provider()).max_concurrency(0).build();
496        assert!(result.is_err());
497    }
498
499    #[test]
500    fn builder_rejects_zero_max_agents() {
501        let result = WorkflowCtx::builder(provider()).max_agents(0).build();
502        assert!(result.is_err());
503    }
504
505    #[test]
506    fn builder_accepts_defaults() {
507        let ctx = WorkflowCtx::builder(provider()).build().expect("build");
508        assert_eq!(ctx.max_agents(), DEFAULT_MAX_AGENTS);
509        assert!(!ctx.is_cancelled());
510    }
511
512    // ----- P6: nesting + stop() -----
513
514    fn budgeted(total: u64) -> WorkflowCtx {
515        WorkflowCtx::builder(provider())
516            .budget(total)
517            .build()
518            .expect("build")
519    }
520
521    #[test]
522    fn nested_shares_budget_pool() {
523        let parent = budgeted(1000);
524        let child = parent.nested().expect("nest");
525        // A spend recorded through the child is visible on the parent's budget.
526        child.record_spend(&crate::llm::types::TokenUsage {
527            input_tokens: 30,
528            ..Default::default()
529        });
530        assert_eq!(parent.budget().spent(), 30);
531        assert_eq!(child.budget().spent(), 30);
532    }
533
534    #[test]
535    fn nested_shares_runaway_backstop() {
536        let parent = WorkflowCtx::builder(provider())
537            .max_agents(2)
538            .build()
539            .expect("build");
540        let child = parent.nested().expect("nest");
541        // Two registrations across parent+child exhaust the shared backstop.
542        assert!(parent.register_agent().is_ok()); // prior 0
543        assert!(child.register_agent().is_ok()); // prior 1
544        assert!(child.register_agent().is_err()); // prior 2 >= 2 -> breach
545        // The monotonic counter is shared, not per-ctx.
546        assert_eq!(parent.spawned(), 3);
547        assert_eq!(child.spawned(), 3);
548    }
549
550    #[test]
551    fn child_breach_cancels_parent() {
552        let parent = budgeted(1000);
553        let child = parent.nested().expect("nest");
554        child.record_breach(ControlBreach::AgentBackstop { limit: 5 });
555        // Breach + cancellation propagate to the shared parent ctx.
556        assert!(parent.is_cancelled());
557        assert!(matches!(
558            parent.control_breach(),
559            Some(ControlBreach::AgentBackstop { limit: 5 })
560        ));
561    }
562
563    #[test]
564    fn nesting_is_one_level_only() {
565        let parent = budgeted(1000);
566        let child = parent.nested().expect("first level ok");
567        let grandchild = child.nested();
568        assert!(
569            grandchild.is_err(),
570            "a second nesting level must be rejected"
571        );
572    }
573
574    #[test]
575    fn child_has_independent_default_phase() {
576        let parent = budgeted(1000);
577        parent.swap_default_phase(Some(std::sync::Arc::from("parent-phase")));
578        let child = parent.nested().expect("nest");
579        // The child starts with no default phase of its own.
580        assert!(child.current_phase().is_none());
581        // And setting the child's phase does not leak back to the parent.
582        child.swap_default_phase(Some(std::sync::Arc::from("child-phase")));
583        assert_eq!(parent.current_phase().as_deref(), Some("parent-phase"));
584        assert_eq!(child.current_phase().as_deref(), Some("child-phase"));
585    }
586
587    #[test]
588    fn stop_sets_cancelled() {
589        let ctx = budgeted(1000);
590        assert!(!ctx.is_cancelled());
591        ctx.stop();
592        assert!(ctx.is_cancelled());
593    }
594}
heartbit_core/agent/flow/ctx.rs

heartbit_core/agent/flow/
ctx.rs
