bamboo_engine/runtime/execution/
agent_spawn.rs

1//! Core agent execution spawning logic.
2//!
3//! Provides [`spawn_session_execution`] which handles the full lifecycle of a
4//! background agent run: spawn task → execute → finalize runner → persist session.
5
6use std::collections::{BTreeSet, HashMap};
7use std::future::Future;
8use std::pin::Pin;
9use std::sync::Arc;
10
11use tokio::sync::{mpsc, RwLock};
12use tokio_util::sync::CancellationToken;
13use tracing::Instrument;
14
15use bamboo_agent_core::tools::ToolExecutor;
16use bamboo_agent_core::{AgentError, AgentEvent, Session};
17use bamboo_domain::ReasoningEffort;
18use bamboo_llm::LLMProvider;
19
20use crate::runtime::config::{
21    AuxiliaryModelConfig, BashCompletionSink, BashResumeHook, GoldConfig, GuardianConfig,
22    GuardianSpawner, ImageFallbackConfig,
23};
24use crate::runtime::execution::runner_lifecycle::finalize_runner;
25use crate::runtime::execution::runner_state::AgentRunner;
26use crate::runtime::model_roster::ModelRoster;
27use crate::runtime::Agent;
28use crate::runtime::{ExecuteRequest, ExecuteRequestBuilder};
29
30/// Shared, per-session-locked session cache.
31///
32/// A `DashMap` gives each session id its own shard-level lock (so unrelated
33/// sessions never contend), and the inner `parking_lot::RwLock` is a *sync*
34/// lock held only to briefly clone-out or mutate-and-write-back a single
35/// `Session` — never across an `.await`. Using a sync lock makes "no guard
36/// across await" a compile-time guarantee in Send futures.
37pub type SessionCache = std::sync::Arc<
38    dashmap::DashMap<String, std::sync::Arc<parking_lot::RwLock<bamboo_agent_core::Session>>>,
39>;
40
41/// Read a session out of the in-memory cache, cloning it out from under the
42/// brief sync read-lock. Returns `None` on a cache miss.
43///
44/// This is the single canonical cache-read used everywhere a caller holds a
45/// `SessionCache` (HTTP handlers, server tools, the app-state loader). It
46/// replaced ~13 verbatim copies of the
47/// `cache.get(id).map(|e| e.value().clone()).map(|a| a.read().clone())` idiom.
48pub fn read_cached_session(cache: &SessionCache, id: &str) -> Option<bamboo_agent_core::Session> {
49    cache
50        .get(id)
51        .map(|e| e.value().clone())
52        .map(|a| a.read().clone())
53}
54
55const SKILL_CONTEXT_START_MARKER: &str = "<!-- BAMBOO_SKILL_CONTEXT_START -->";
56const TOOL_GUIDE_START_MARKER: &str = "<!-- BAMBOO_TOOL_GUIDE_START -->";
57const EXTERNAL_MEMORY_START_MARKER: &str = "<!-- BAMBOO_EXTERNAL_MEMORY_START -->";
58const TASK_LIST_START_MARKER: &str = "<!-- BAMBOO_TASK_LIST_START -->";
59
60/// Outcome of an agent execution, handed to an optional
61/// [`SessionCompletionHook`].
62///
63/// Deliberately decoupled from the runtime's internal error type so the hook
64/// API stays stable across crates and callers don't need to match on engine
65/// error variants.
66pub struct SessionExecutionOutcome {
67    /// The run finished without error.
68    pub success: bool,
69    /// The run ended because it was cancelled (a non-success subset).
70    pub cancelled: bool,
71    /// Stringified error, present when `!success`.
72    pub error: Option<String>,
73}
74
75impl SessionExecutionOutcome {
76    fn from_result(result: &Result<(), AgentError>) -> Self {
77        match result {
78            Ok(()) => Self {
79                success: true,
80                cancelled: false,
81                error: None,
82            },
83            Err(error) => Self {
84                success: false,
85                cancelled: error.is_cancelled(),
86                error: Some(error.to_string()),
87            },
88        }
89    }
90}
91
92/// Optional post-execution hook for [`spawn_session_execution`].
93///
94/// Invoked after the runner is finalized but **before** the session is
95/// persisted, so a caller can record bespoke terminal bookkeeping (e.g. a
96/// scheduled-run status) and/or append a closing message that is then saved
97/// with the session. Receives the execution outcome plus a mutable handle to
98/// the session. This is how callers with extra finalization (the schedule
99/// manager) route through the single canonical execution path instead of
100/// forking their own spawn + `execute` + persist sequence.
101pub type SessionCompletionHook = Box<
102    dyn for<'a> FnOnce(
103            SessionExecutionOutcome,
104            &'a mut Session,
105        ) -> Pin<Box<dyn Future<Output = ()> + Send + 'a>>
106        + Send,
107>;
108
109/// Arguments for spawning a background agent execution.
110///
111/// This is the crate-agnostic equivalent of the server's `SpawnAgentExecution`.
112/// It holds everything needed to run the agent loop and persist the result,
113/// without depending on HTTP types or `AppState`.
114pub struct SessionExecutionArgs {
115    // Core execution.
116    pub agent: Arc<Agent>,
117    pub session_id: String,
118    pub session: Session,
119
120    // Execution parameters.
121    pub tools_override: Option<Arc<dyn ToolExecutor>>,
122    pub provider_override: Option<Arc<dyn LLMProvider>>,
123    /// Cohesive primary + auxiliary model/provider selection. The primary
124    /// `model` is required for a spawn (see [`ModelRoster::model`]); the three
125    /// auxiliary roles default to their `Config::get_*` fallbacks when `None`.
126    pub model_roster: ModelRoster,
127    pub reasoning_effort: Option<ReasoningEffort>,
128    pub reasoning_effort_source: String,
129    pub auxiliary_model_resolver:
130        Option<Arc<dyn Fn() -> crate::runtime::config::AuxiliaryModelConfig + Send + Sync>>,
131    /// Optional per-round live resolver for the disabled tool/skill sets (#136).
132    /// When `None` the per-run snapshot is used (sub-agent spawns pass `None`, so
133    /// short-lived children keep the spawn-time snapshot — by design).
134    pub disabled_filter_resolver:
135        Option<Arc<dyn Fn() -> (BTreeSet<String>, BTreeSet<String>) + Send + Sync>>,
136    pub disabled_tools: Option<BTreeSet<String>>,
137    pub disabled_skill_ids: Option<BTreeSet<String>>,
138    pub selected_skill_ids: Option<Vec<String>>,
139    pub selected_skill_mode: Option<String>,
140    pub cancel_token: CancellationToken,
141    pub mpsc_tx: mpsc::Sender<AgentEvent>,
142    pub image_fallback: Option<ImageFallbackConfig>,
143    pub gold_config: Option<GoldConfig>,
144    /// Optional guardian adversarial-review gate configuration.
145    pub guardian_config: Option<GuardianConfig>,
146    /// Late-bound guardian reviewer spawner (server-provided; the runner cannot
147    /// construct a child directly).
148    pub guardian_spawner: Option<Arc<dyn GuardianSpawner>>,
149    /// Late-bound bash self-resume hook (issue #84 Phase 2b).
150    pub bash_resume_hook: Option<Arc<dyn BashResumeHook>>,
151    /// Late-bound bash completion sink (issue #84 Phase 2b follow-up).
152    pub bash_completion_sink: Option<Arc<dyn BashCompletionSink>>,
153    pub app_data_dir: Option<std::path::PathBuf>,
154
155    // Post-execution resources.
156    pub runners: Arc<RwLock<HashMap<String, AgentRunner>>>,
157    pub sessions_cache: SessionCache,
158
159    /// Optional bespoke finalization, run after the runner is finalized and
160    /// before the session is persisted. See [`SessionCompletionHook`].
161    pub on_complete: Option<SessionCompletionHook>,
162}
163
164/// The per-request parameter subset of [`SessionExecutionArgs`] — everything
165/// that maps onto an [`ExecuteRequest`], minus the three required positional
166/// fields (`initial_message`, `event_tx`, `cancel_token`) and the post-execution
167/// resources (runners, sessions cache, completion hook).
168///
169/// Grouping these here lets [`build_execute_request`] perform the
170/// `SessionExecutionArgs` → [`ExecuteRequest`] mapping through the canonical
171/// [`ExecuteRequestBuilder`] in one place, instead of a hand-written struct
172/// literal that must be kept field-aligned with [`ExecuteRequest`] by hand.
173struct ExecuteRequestParams {
174    tools: Option<Arc<dyn ToolExecutor>>,
175    provider_override: Option<Arc<dyn LLMProvider>>,
176    model_roster: ModelRoster,
177    reasoning_effort: Option<ReasoningEffort>,
178    auxiliary_model_resolver: Option<Arc<dyn Fn() -> AuxiliaryModelConfig + Send + Sync>>,
179    disabled_filter_resolver:
180        Option<Arc<dyn Fn() -> (BTreeSet<String>, BTreeSet<String>) + Send + Sync>>,
181    disabled_tools: Option<BTreeSet<String>>,
182    disabled_skill_ids: Option<BTreeSet<String>>,
183    selected_skill_ids: Option<Vec<String>>,
184    selected_skill_mode: Option<String>,
185    image_fallback: Option<ImageFallbackConfig>,
186    gold_config: Option<GoldConfig>,
187    guardian_config: Option<GuardianConfig>,
188    guardian_spawner: Option<Arc<dyn GuardianSpawner>>,
189    bash_resume_hook: Option<Arc<dyn BashResumeHook>>,
190    bash_completion_sink: Option<Arc<dyn BashCompletionSink>>,
191    app_data_dir: Option<std::path::PathBuf>,
192}
193
194/// Assemble an [`ExecuteRequest`] from the resolved spawn parameters via the
195/// canonical [`ExecuteRequestBuilder`].
196///
197/// Centralizing this mapping keeps every optional field threaded with exactly
198/// the same value the old struct literal carried (the builder defaults each
199/// unset field to `None`), while removing the field-by-field duplication.
200fn build_execute_request(
201    initial_message: String,
202    event_tx: mpsc::Sender<AgentEvent>,
203    cancel_token: CancellationToken,
204    params: ExecuteRequestParams,
205) -> ExecuteRequest {
206    let ExecuteRequestParams {
207        tools,
208        provider_override,
209        model_roster,
210        reasoning_effort,
211        auxiliary_model_resolver,
212        disabled_filter_resolver,
213        disabled_tools,
214        disabled_skill_ids,
215        selected_skill_ids,
216        selected_skill_mode,
217        image_fallback,
218        gold_config,
219        guardian_config,
220        guardian_spawner,
221        bash_resume_hook,
222        bash_completion_sink,
223        app_data_dir,
224    } = params;
225
226    let mut builder = ExecuteRequestBuilder::new(initial_message, event_tx, cancel_token)
227        .model_roster(model_roster)
228        .gold_config(gold_config)
229        .guardian_config(guardian_config)
230        .guardian_spawner(guardian_spawner)
231        .bash_resume_hook(bash_resume_hook)
232        .bash_completion_sink(bash_completion_sink);
233
234    if let Some(tools) = tools {
235        builder = builder.tools(tools);
236    }
237    if let Some(provider_override) = provider_override {
238        builder = builder.provider_override(provider_override);
239    }
240    if let Some(reasoning_effort) = reasoning_effort {
241        builder = builder.reasoning_effort(reasoning_effort);
242    }
243    if let Some(disabled_filter_resolver) = disabled_filter_resolver {
244        builder = builder.disabled_filter_resolver(disabled_filter_resolver);
245    }
246    if let Some(auxiliary_model_resolver) = auxiliary_model_resolver {
247        builder = builder.auxiliary_model_resolver(auxiliary_model_resolver);
248    }
249    if let Some(disabled_tools) = disabled_tools {
250        builder = builder.disabled_tools(disabled_tools);
251    }
252    if let Some(disabled_skill_ids) = disabled_skill_ids {
253        builder = builder.disabled_skill_ids(disabled_skill_ids);
254    }
255    if let Some(selected_skill_ids) = selected_skill_ids {
256        builder = builder.selected_skill_ids(selected_skill_ids);
257    }
258    if let Some(selected_skill_mode) = selected_skill_mode {
259        builder = builder.selected_skill_mode(selected_skill_mode);
260    }
261    if let Some(image_fallback) = image_fallback {
262        builder = builder.image_fallback(image_fallback);
263    }
264    if let Some(app_data_dir) = app_data_dir {
265        builder = builder.app_data_dir(app_data_dir);
266    }
267
268    builder.build()
269}
270
271/// Spawn a background agent execution task.
272///
273/// This function spawns a tokio task that:
274/// 1. Executes the agent loop via `agent.execute()`
275/// 2. Sends a terminal error event if the execution fails
276/// 3. Finalizes the runner status
277/// 4. Persists the session via merge-save (preserves concurrent UI title/pin edits)
278/// 5. Updates the in-memory session cache
279pub fn spawn_session_execution(args: SessionExecutionArgs) {
280    let span_session_id = args.session_id.clone();
281    let session_span = tracing::info_span!("agent_execution", session_id = %span_session_id);
282
283    tokio::spawn(
284        async move {
285            let SessionExecutionArgs {
286                agent,
287                session_id,
288                mut session,
289                tools_override,
290                provider_override,
291                model_roster,
292                reasoning_effort,
293                reasoning_effort_source,
294                auxiliary_model_resolver,
295                disabled_filter_resolver,
296                disabled_tools,
297                disabled_skill_ids,
298                selected_skill_ids,
299                selected_skill_mode,
300                cancel_token,
301                mpsc_tx,
302                image_fallback,
303                gold_config,
304                guardian_config,
305                guardian_spawner,
306                bash_resume_hook,
307                bash_completion_sink,
308                app_data_dir,
309                runners,
310                sessions_cache,
311                on_complete,
312            } = args;
313
314            // The primary model is required for a spawn; the roster stores it as
315            // `Option<String>` for uniformity, so recover the owned String here
316            // for session attribution / logging (same value the caller set).
317            let model = model_roster.model.clone().unwrap_or_default();
318
319            let initial_message = initial_user_message_for_session(&session);
320            let selected_skill_ids =
321                selected_skill_ids.or_else(|| selected_skill_ids_for_session(&session));
322            let selected_skill_mode =
323                selected_skill_mode.or_else(|| selected_skill_mode_for_session(&session));
324
325            tracing::info!(
326                "[{}] Using resolved session model: {}, reasoning_effort={}, reasoning_source={}",
327                session_id,
328                model,
329                reasoning_effort
330                    .map(ReasoningEffort::as_str)
331                    .unwrap_or("none"),
332                reasoning_effort_source
333            );
334
335            // Set the resolved model via the single authoritative pre-execution
336            // mutation point. The caller already placed the system prompt on the
337            // session, so pass `None` for `system_prompt` (the subsequent
338            // `system_prompt_for_session` read below sees the caller's message).
339            // This must run before that read / logging so the observable
340            // sequence (model set, then prompt snapshot) is identical.
341            crate::session_app::execution_prep::prepare_session_for_execution(
342                &mut session,
343                None,
344                Some(&model),
345            );
346
347            let system_prompt = system_prompt_for_session(&session);
348            if let Some(prompt) = system_prompt.as_ref() {
349                log_base_system_prompt_snapshot(&session_id, prompt);
350            }
351
352            let execute_request = build_execute_request(
353                initial_message,
354                mpsc_tx.clone(),
355                cancel_token,
356                ExecuteRequestParams {
357                    tools: tools_override,
358                    provider_override,
359                    model_roster,
360                    reasoning_effort,
361                    auxiliary_model_resolver,
362                    disabled_filter_resolver,
363                    disabled_tools,
364                    disabled_skill_ids,
365                    selected_skill_ids,
366                    selected_skill_mode,
367                    image_fallback,
368                    gold_config,
369                    guardian_config,
370                    guardian_spawner,
371                    bash_resume_hook,
372                    bash_completion_sink,
373                    app_data_dir,
374                },
375            );
376
377            let result = agent.execute(&mut session, execute_request).await;
378
379            // Send terminal event for all error cases (including cancellation).
380            if let Some(error_event) = terminal_error_event_for_result(&result) {
381                let _ = mpsc_tx.send(error_event).await;
382            }
383
384            // Record the terminal run status on the session BEFORE persisting so
385            // the session summary reports a real `last_run_status`. Top-level
386            // sessions otherwise never set it, so summaries show
387            // `last_run_status: null`; the frontend then cannot confirm the run
388            // finished and falls back on a ~5s optimistic-settle window, leaving
389            // a phantom "thinking" indicator after the reply is already done
390            // (notably on a session's first turn). Same Ok/cancelled/error
391            // mapping `status_from_execution_result` applies to the runner.
392            //
393            // A suspended run also returns `Ok(())` but is NOT terminal: it
394            // stamped `runtime.suspend_reason` (awaiting_clarification /
395            // waiting_for_children / waiting_for_bash / awaiting_parent_approval)
396            // and will resume later (which removes the reason — see respond.rs /
397            // child_completion_coordinator). Mark it "suspended" rather than
398            // "completed" so a session waiting on the user or on children isn't
399            // reported as finished (mirrors the child path in `sdk::spawn`).
400            let suspended_non_terminal = result.is_ok()
401                && session
402                    .metadata
403                    .get("runtime.suspend_reason")
404                    .is_some_and(|reason| !reason.trim().is_empty());
405            match &result {
406                Ok(()) if suspended_non_terminal => {
407                    session.set_last_run_status("suspended");
408                    session.clear_last_run_error();
409                }
410                Ok(()) => {
411                    session.set_last_run_status("completed");
412                    session.clear_last_run_error();
413                }
414                Err(error) if error.is_cancelled() => {
415                    session.set_last_run_status("cancelled");
416                    session.set_last_run_error(error.to_string());
417                }
418                Err(error) => {
419                    session.set_last_run_status("error");
420                    session.set_last_run_error(error.to_string());
421                }
422            }
423
424            // Bespoke terminal bookkeeping (e.g. a scheduled-run status) runs
425            // here — before persistence — so any closing message the hook
426            // appends is saved with the session below.
427            if let Some(on_complete) = on_complete {
428                on_complete(SessionExecutionOutcome::from_result(&result), &mut session).await;
429            }
430
431            // Save session via merge-save so any concurrent UI edits to
432            // title / pinned / title_version are preserved (the runtime is not
433            // an authoritative title writer).
434            if let Err(error) = agent.persistence().save_runtime_session(&mut session).await {
435                tracing::warn!("[{}] Failed to save session: {}", session_id, error);
436            }
437
438            // Flip the runner registry to a terminal status (which makes session
439            // summaries report `is_running: false`) ONLY AFTER the run status is
440            // persisted above, so `is_running` and `last_run_status` become
441            // visible together and the frontend settles immediately instead of
442            // lingering in its optimistic-settle window.
443            finalize_runner(&runners, &session_id, &result).await;
444
445            // Update memory cache.
446            sessions_cache.insert(
447                session_id.clone(),
448                Arc::new(parking_lot::RwLock::new(session)),
449            );
450
451            tracing::info!("[{}] Agent execution completed", session_id);
452        }
453        .instrument(session_span),
454    );
455}
456
457/// Log a snapshot of the base system prompt for debugging.
458pub fn log_base_system_prompt_snapshot(session_id: &str, prompt: &str) {
459    tracing::info!(
460        "[{}] Base system prompt snapshot: len={} chars, has_skill={}, has_tool_guide={}, has_external_memory={}, has_task_list={}",
461        session_id,
462        prompt.len(),
463        prompt.contains(SKILL_CONTEXT_START_MARKER),
464        prompt.contains(TOOL_GUIDE_START_MARKER),
465        prompt.contains(EXTERNAL_MEMORY_START_MARKER),
466        prompt.contains(TASK_LIST_START_MARKER),
467    );
468
469    tracing::debug!(
470        "[{}] ========== BASE SYSTEM PROMPT SNAPSHOT ==========",
471        session_id
472    );
473    tracing::debug!("[{}] Snapshot length: {} chars", session_id, prompt.len());
474    tracing::debug!("[{}] -----------------------------------", session_id);
475    tracing::debug!("[{}] {}", session_id, prompt);
476    tracing::debug!(
477        "[{}] ========== END BASE SYSTEM PROMPT SNAPSHOT ==========",
478        session_id
479    );
480}
481
482/// Map an execution result to a terminal error event.
483pub fn terminal_error_event_for_result(result: &Result<(), AgentError>) -> Option<AgentEvent> {
484    match result {
485        Ok(_) => None,
486        Err(error) if error.is_cancelled() => Some(AgentEvent::Error {
487            message: "Agent execution cancelled by user".to_string(),
488        }),
489        Err(error) => Some(AgentEvent::Error {
490            message: error.to_string(),
491        }),
492    }
493}
494
495// Session metadata helpers (pure functions, no server dependency).
496
497fn system_prompt_for_session(session: &Session) -> Option<String> {
498    session
499        .messages
500        .iter()
501        .find(|message| matches!(message.role, bamboo_agent_core::Role::System))
502        .map(|message| message.content.clone())
503}
504
505fn initial_user_message_for_session(session: &Session) -> String {
506    session
507        .messages
508        .last()
509        .filter(|message| matches!(message.role, bamboo_agent_core::Role::User))
510        .map(|message| message.content.clone())
511        .unwrap_or_default()
512}
513
514fn selected_skill_ids_for_session(session: &Session) -> Option<Vec<String>> {
515    session
516        .metadata
517        .get("selected_skill_ids")
518        .and_then(|raw| bamboo_skills::selection::parse_selected_skill_ids_metadata(raw))
519}
520
521fn selected_skill_mode_for_session(session: &Session) -> Option<String> {
522    let value = session
523        .metadata
524        .get("skill_mode")
525        .or_else(|| session.metadata.get("mode"))?;
526    let trimmed = value.trim();
527    if trimmed.is_empty() {
528        None
529    } else {
530        Some(trimmed.to_string())
531    }
532}
bamboo_engine/runtime/execution/agent_spawn.rs

bamboo_engine/runtime/execution/
agent_spawn.rs
