bamboo_server/app_state/
mod.rs

1//! Unified application state management for the Bamboo server
2//!
3//! This module provides the central AppState struct that consolidates all
4//! server state including sessions, storage, LLM providers, tools, and metrics.
5//!
6//! # Architecture
7//!
8//! The AppState uses a unified design that eliminates the proxy pattern where
9//! web_service created an AgentAppState that called back via HTTP. Instead, it
10//! provides direct access to all components.
11//!
12//! ```text
13//! ┌────────────────────────────────────────────────────┐
14//! │              AppState (Unified)                    │
15//! │                                                    │
16//! │  ┌──────────────┐      ┌──────────────┐          │
17//! │  │   Config     │      │   Provider   │          │
18//! │  │  (Hot-reload)│◄────►│   (LLM)      │          │
19//! │  └──────────────┘      └──────────────┘          │
20//! │                                                    │
21//! │  ┌──────────────┐      ┌──────────────┐          │
22//! │  │   Sessions   │      │   Storage    │          │
23//! │  │  (In-memory) │      │  (Persistent)│          │
24//! │  └──────────────┘      └──────────────┘          │
25//! │                                                    │
26//! │  ┌──────────────┐      ┌──────────────┐          │
27//! │  │    Tools     │      │    Skills    │          │
28//! │  │ (Builtin+MCP)│      │   Manager    │          │
29//! │  └──────────────┘      └──────────────┘          │
30//! │                                                    │
31//! │  ┌──────────────┐      ┌──────────────┐          │
32//! │  │     MCP      │      │   Metrics    │          │
33//! │  │   Manager    │      │   Service    │          │
34//! │  └──────────────┘      └──────────────┘          │
35//! └────────────────────────────────────────────────────┘
36//! ```
37//!
38//! # Key Features
39//!
40//! - **Hot-reloadable configuration**: Config and provider can be reloaded at runtime
41//! - **Direct provider access**: No HTTP proxy overhead
42//! - **Session management**: In-memory session cache with persistent storage
43//! - **Tool composition**: Combines built-in and MCP tools
44//! - **Metrics collection**: Integrated metrics and event tracking
45//!
46//! # Usage Example
47//!
48//! ```rust,no_run
49//! use bamboo_server::app_state::AppState;
50//! use std::path::PathBuf;
51//!
52//! #[tokio::main]
53//! async fn main() {
54//!     // Initialize app state
55//!     let app_data_dir = PathBuf::from("/path/to/.bamboo");
56//!     let state = AppState::new(app_data_dir)
57//!         .await
58//!         .expect("failed to initialize app state");
59//!
60//!     // Access components
61//!     let provider = state.get_provider().await;
62//!     let schemas = state.get_all_tool_schemas();
63//!
64//!     // Hot reload configuration
65//!     state.reload_config().await;
66//!     state.reload_provider().await.ok();
67//! }
68//! ```
69
70use std::collections::HashMap;
71use std::path::PathBuf;
72use std::sync::Arc;
73
74use async_trait::async_trait;
75use tokio::sync::{broadcast, RwLock};
76use tokio_util::sync::CancellationToken;
77
78use crate::error::AppError;
79use crate::metrics_service::MetricsService;
80use crate::schedules::{ScheduleManager, ScheduleStore};
81use crate::spawn_scheduler::SpawnScheduler;
82use bamboo_agent_core::storage::Storage;
83use bamboo_agent_core::AgentEvent;
84use bamboo_agent_core::{tools::ToolSchema, Message};
85use bamboo_engine::McpServerManager;
86use bamboo_engine::SkillManager;
87use bamboo_infrastructure::registry::ProcessRegistry;
88use bamboo_infrastructure::Config;
89use bamboo_infrastructure::SessionStoreV2;
90use bamboo_infrastructure::{LLMError, LLMProvider, LLMStream};
91
92// Context functions moved to bamboo-agent-runtime::context
93pub use bamboo_engine::context::{
94    build_env_prompt_context, build_workspace_prompt_context, workspace_prompt_guidance,
95    DEFAULT_BASE_PROMPT, ENV_CONTEXT_END_MARKER, ENV_CONTEXT_START_MARKER,
96    WORKSPACE_CONTEXT_END_MARKER, WORKSPACE_CONTEXT_PREFIX, WORKSPACE_CONTEXT_START_MARKER,
97};
98
99/// Placeholder provider used when the configured provider cannot be initialized.
100///
101/// This keeps the server usable for configuration/UX flows while ensuring we fail fast
102/// (instead of silently switching to a different provider or model).
103struct UnconfiguredProvider {
104    message: String,
105}
106
107#[async_trait]
108impl LLMProvider for UnconfiguredProvider {
109    async fn chat_stream(
110        &self,
111        _messages: &[Message],
112        _tools: &[ToolSchema],
113        _max_output_tokens: Option<u32>,
114        _model: &str,
115    ) -> bamboo_infrastructure::provider::Result<LLMStream> {
116        Err(LLMError::Auth(format!(
117            "LLM provider is not configured: {}",
118            self.message
119        )))
120    }
121
122    async fn list_models(&self) -> bamboo_infrastructure::provider::Result<Vec<String>> {
123        Err(LLMError::Auth(format!(
124            "LLM provider is not configured: {}",
125            self.message
126        )))
127    }
128}
129
130// Re-export execution types from the runtime crate.
131pub use bamboo_engine::execution::runner_state::{AgentRunner, AgentStatus};
132
133/// Unified application state consolidating web_service and agent/server state
134///
135/// This struct holds all the state needed to run the Bamboo server, including
136/// configuration, LLM providers, sessions, storage, tools, skills, and metrics.
137///
138/// # Design Goals
139///
140/// - **Direct access**: Components are directly accessible without HTTP proxies
141/// - **Hot reload**: Configuration and providers can be reloaded at runtime
142/// - **Thread safety**: Uses Arc<RwLock> for concurrent access
143/// - **Persistence**: Integrates with JsonlStorage for session persistence
144///
145/// # Component Overview
146///
147/// | Component | Purpose | Thread-Safe |
148/// |-----------|---------|--------------|
149/// | `config` | Application configuration | Yes (RwLock) |
150/// | `provider` | Hot-reloadable LLM provider | Yes (RwLock) |
151/// | `sessions` | Active conversation sessions | Yes (RwLock) |
152/// | `storage` | Persistent session storage | Yes (Arc) |
153/// | `tools` | Tool execution (builtin + MCP) | Yes (Arc) |
154/// | `skill_manager` | Skill registry and execution | Yes (Arc) |
155/// | `mcp_manager` | MCP server lifecycle | Yes (Arc) |
156/// | `metrics_service` | Usage metrics collection | Yes (Arc) |
157/// | `agent_runners` | Active agent executions | Yes (RwLock) |
158pub struct AppState {
159    /// Application data directory (configured via `BAMBOO_DATA_DIR`; default `${HOME}/.bamboo`)
160    pub app_data_dir: PathBuf,
161
162    /// Hot-reloadable application configuration
163    ///
164    /// Can be reloaded from disk at runtime using `reload_config()`.
165    pub config: Arc<RwLock<Config>>,
166
167    /// Hot-reloadable LLM provider with direct access
168    ///
169    /// This eliminates the proxy pattern where we created an AgentAppState
170    /// that called back to web_service via HTTP. Now we have direct provider access.
171    pub provider: Arc<RwLock<Arc<dyn LLMProvider>>>,
172
173    /// Stable handle that always delegates to the latest provider in `provider`.
174    ///
175    /// This avoids stale provider snapshots after runtime config updates.
176    provider_handle: Arc<dyn LLMProvider>,
177
178    /// Active conversation sessions (in-memory cache)
179    ///
180    /// Maps session IDs to Session objects. Persisted to storage
181    /// via the `storage` field.
182    pub sessions: Arc<RwLock<HashMap<String, bamboo_agent_core::Session>>>,
183
184    /// Persistent storage backend for sessions (V2).
185    ///
186    /// Implemented as folder-per-session with a global `sessions.json` index.
187    pub storage: Arc<dyn Storage>,
188
189    /// Concrete session store implementation (for index/list/cleanup APIs).
190    pub session_store: Arc<SessionStoreV2>,
191
192    /// Background scheduler for async sub-session spawning.
193    pub spawn_scheduler: Arc<SpawnScheduler>,
194
195    /// Schedule store (timed tasks).
196    pub schedule_store: Arc<ScheduleStore>,
197
198    /// Background schedule manager that triggers scheduled runs.
199    pub schedule_manager: Arc<ScheduleManager>,
200
201    /// Tool surface factory providing pre-built tool executors for each session type.
202    ///
203    /// Use `state.tools_for(ToolSurface::Root)` for root sessions,
204    /// `state.tools_for(ToolSurface::Child)` for child sessions, etc.
205    pub tool_factory: crate::tools::ToolSurfaceFactory,
206
207    /// Cancellation tokens for in-flight requests
208    ///
209    /// Maps request/session IDs to their cancellation tokens,
210    /// allowing graceful shutdown of long-running operations.
211    pub cancel_tokens: Arc<RwLock<HashMap<String, CancellationToken>>>,
212
213    /// Skill manager for prompt-based skill execution
214    ///
215    /// Manages the skill registry and handles skill lookup,
216    /// validation, and execution.
217    pub skill_manager: Arc<SkillManager>,
218
219    /// MCP server manager for external tool servers
220    ///
221    /// Handles lifecycle of Model Context Protocol servers,
222    /// including initialization, tool discovery, and shutdown.
223    pub mcp_manager: Arc<McpServerManager>,
224
225    /// Metrics collection and persistence service
226    ///
227    /// Tracks token usage, costs, and performance metrics
228    /// across all sessions.
229    pub metrics_service: Arc<MetricsService>,
230
231    /// Active agent runners indexed by session ID
232    ///
233    /// Each runner manages event broadcasting and cancellation
234    /// for an active agent execution.
235    pub agent_runners: Arc<RwLock<HashMap<String, AgentRunner>>>,
236
237    /// Session-scoped event streams (long-lived).
238    ///
239    /// Unlike `agent_runners`, these senders exist even when no agent execution is running.
240    /// They are used for:
241    /// - UI subscriptions to `/api/v1/events/{session_id}` (background tasks, etc.)
242    /// - sub-session forwarding (child -> parent)
243    pub session_event_senders: Arc<RwLock<HashMap<String, broadcast::Sender<AgentEvent>>>>,
244
245    /// Registry for tracking external processes.
246    pub process_registry: Arc<ProcessRegistry>,
247
248    /// Optional metrics bus for event streaming
249    ///
250    /// When enabled, allows subscribing to metrics events
251    /// in real-time.
252    pub metrics_bus: Option<bamboo_engine::metrics::bus::MetricsBus>,
253
254    /// Unified agent execution runtime holding shared resources.
255    pub agent: Arc<bamboo_engine::Agent>,
256
257    /// Multi-provider registry (used when features.provider_model_ref is enabled).
258    pub provider_registry: Arc<bamboo_infrastructure::ProviderRegistry>,
259
260    /// Provider/model router (used when features.provider_model_ref is enabled).
261    pub provider_router: Arc<bamboo_infrastructure::ProviderModelRouter>,
262
263    /// Unified model catalog service (used when features.provider_model_ref is enabled).
264    pub model_catalog: Arc<bamboo_infrastructure::ModelCatalogService>,
265}
266
267mod builder;
268mod config_runtime;
269pub mod init;
270mod persistence;
271mod provider_api;
272pub mod resume_adapter;
273pub mod runner_lifecycle;
274pub(crate) mod session_events;
275mod session_loader;
276mod tools;
277
278#[cfg(test)]
279mod tests;
280
281#[derive(Debug, Clone, Copy, Default)]
282pub struct ConfigUpdateEffects {
283    pub reload_provider: bool,
284    pub reconcile_mcp: bool,
285}
bamboo_server/app_state/mod.rs

bamboo_server/app_state/
mod.rs
