1use std::collections::HashMap;
2use std::sync::Arc;
3use std::time::Duration;
4
5use chrono::Utc;
6use tokio::sync::{broadcast, mpsc, RwLock};
7
8use bamboo_agent_core::tools::ToolExecutor;
9use bamboo_agent_core::{AgentEvent, Message, Role};
10use bamboo_domain::reasoning::ReasoningEffort;
11use bamboo_engine::config::GoldConfig;
12use bamboo_engine::execution::{
13 create_event_forwarder, get_or_create_event_sender, spawn_session_execution,
14 try_reserve_runner, AgentRunner, RunnerReservation, SessionCompletionHook,
15 SessionExecutionArgs,
16};
17use bamboo_engine::{AuxiliaryModelConfig, ModelRoster};
18use bamboo_storage::LockedSessionStore;
19
20use super::store::{ClaimedScheduleRun, ScheduleStore};
21use super::trigger_engine::DynTriggerEngine;
22use bamboo_domain::{ScheduleRunConfig, ScheduleRunStatus};
23
24#[derive(Debug, Clone)]
25pub struct ScheduleRunJob {
26 pub run_id: String,
27 pub schedule_id: String,
28 pub schedule_name: String,
29 pub run_config: ScheduleRunConfig,
30 pub scheduled_for: chrono::DateTime<chrono::Utc>,
31 pub claimed_at: chrono::DateTime<chrono::Utc>,
32 pub was_catch_up: bool,
33}
34
35#[derive(Clone)]
41pub struct ResolvedRunConfig {
42 pub model_roster: ModelRoster,
45 pub reasoning_effort: Option<ReasoningEffort>,
46 pub gold_config: Option<GoldConfig>,
47 pub system_prompt: String,
48 pub base_system_prompt: String,
49 pub workspace_path: Option<String>,
50}
51
52#[derive(Clone)]
53pub struct ScheduleContext {
54 pub schedule_store: Arc<ScheduleStore>,
55 pub agent: Arc<bamboo_engine::Agent>,
56 pub persistence: Arc<LockedSessionStore>,
57 pub tools: Arc<dyn ToolExecutor>,
58 pub sessions_cache: bamboo_engine::SessionCache,
59 pub agent_runners: Arc<RwLock<HashMap<String, AgentRunner>>>,
60 pub session_event_senders: Arc<RwLock<HashMap<String, broadcast::Sender<AgentEvent>>>>,
61 pub account_feed_inbox: Option<bamboo_engine::execution::AccountFeedInbox>,
63 pub app_data_dir: Option<std::path::PathBuf>,
64 pub trigger_engine: DynTriggerEngine,
65 pub resolve_run_config: Arc<dyn Fn(&ScheduleRunJob) -> ResolvedRunConfig + Send + Sync>,
68}
69
70#[derive(Debug, Clone, Copy, PartialEq, Eq)]
71enum ScheduleRunLifecycleResult {
72 Terminal(ScheduleRunStatus),
73 BackgroundExecutionInProgress,
74}
75
76#[derive(Clone)]
77pub struct ScheduleManager {
78 tx: mpsc::Sender<ScheduleRunJob>,
79}
80
81impl ScheduleManager {
82 pub fn new(ctx: ScheduleContext) -> Self {
83 let (tx, mut rx) = mpsc::channel::<ScheduleRunJob>(128);
84
85 tokio::spawn({
87 let ctx = ctx.clone();
88 async move {
89 while let Some(job) = rx.recv().await {
90 if let Err(error) = ctx
91 .schedule_store
92 .mark_run_started(&job.schedule_id, &job.run_id)
93 .await
94 {
95 tracing::warn!(
96 "failed to mark schedule run started for {} / {}: {}",
97 job.schedule_id,
98 job.run_id,
99 error
100 );
101 }
102 let schedule_id = job.schedule_id.clone();
103 let run_id = job.run_id.clone();
104 match run_schedule_job(ctx.clone(), job).await {
105 Ok(ScheduleRunLifecycleResult::Terminal(status)) => {
106 if let Err(error) = ctx
107 .schedule_store
108 .mark_run_terminal(&schedule_id, &run_id, status, None)
109 .await
110 {
111 tracing::warn!(
112 "failed to mark schedule run terminal state for {} / {}: {}",
113 schedule_id,
114 run_id,
115 error
116 );
117 }
118 }
119 Ok(ScheduleRunLifecycleResult::BackgroundExecutionInProgress) => {}
120 Err(e) => {
121 tracing::warn!("schedule job failed: {e}");
122 if let Err(error) = ctx
123 .schedule_store
124 .mark_run_terminal(
125 &schedule_id,
126 &run_id,
127 ScheduleRunStatus::Failed,
128 Some(e.clone()),
129 )
130 .await
131 {
132 tracing::warn!(
133 "failed to mark schedule run failed state for {} / {}: {}",
134 schedule_id,
135 run_id,
136 error
137 );
138 }
139 }
140 }
141 }
142 }
143 });
144
145 tokio::spawn({
147 let tx = tx.clone();
148 let store = ctx.schedule_store.clone();
149 let trigger_engine = ctx.trigger_engine.clone();
150 async move {
151 let mut ticker = tokio::time::interval(Duration::from_secs(15));
152 loop {
153 ticker.tick().await;
154 let now = Utc::now();
155 let claimed: Vec<ClaimedScheduleRun> = match store
156 .claim_due_runs_with_engine(now, trigger_engine.as_ref())
157 .await
158 {
159 Ok(v) => v,
160 Err(e) => {
161 tracing::warn!("claim_due_runs failed: {e}");
162 continue;
163 }
164 };
165 for c in claimed {
166 let schedule_id = c.schedule_id.clone();
167 let run_id = c.run_id.clone();
168 if tx
169 .send(ScheduleRunJob {
170 run_id: c.run_id,
171 schedule_id: c.schedule_id,
172 schedule_name: c.schedule_name,
173 run_config: c.run_config,
174 scheduled_for: c.scheduled_for,
175 claimed_at: c.claimed_at,
176 was_catch_up: c.was_catch_up,
177 })
178 .await
179 .is_err()
180 {
181 let _ = store
182 .mark_run_dequeued_without_start(
183 &schedule_id,
184 &run_id,
185 Some("schedule manager is not running".to_string()),
186 )
187 .await;
188 }
189 }
190 }
191 }
192 });
193
194 Self { tx }
195 }
196
197 pub async fn enqueue_run_now(&self, job: ScheduleRunJob) -> Result<(), String> {
198 self.tx
199 .send(job)
200 .await
201 .map_err(|_| "schedule manager is not running".to_string())
202 }
203}
204
205async fn run_schedule_job(
206 ctx: ScheduleContext,
207 job: ScheduleRunJob,
208) -> Result<ScheduleRunLifecycleResult, String> {
209 let resolved = (ctx.resolve_run_config)(&job);
210 let resolved_model = resolved.model_roster.model.clone().unwrap_or_default();
214
215 if resolved_model.trim().is_empty() {
217 tracing::warn!(
218 "[schedule:{}] skipping run: resolved model is empty",
219 job.schedule_id
220 );
221 return Ok(ScheduleRunLifecycleResult::Terminal(
222 ScheduleRunStatus::Skipped,
223 ));
224 }
225
226 let requested_model = job
227 .run_config
228 .model
229 .as_deref()
230 .map(str::trim)
231 .filter(|v| !v.is_empty())
232 .map(|v| v.to_string());
233 let requested_reasoning_effort = job.run_config.reasoning_effort;
234
235 let mut session = super::session_factory::create_schedule_session(
236 &job,
237 &resolved_model,
238 &resolved.system_prompt,
239 &resolved.base_system_prompt,
240 resolved.workspace_path.as_deref(),
241 resolved.reasoning_effort,
242 );
243 let session_id = session.id.clone();
244
245 session
250 .agent_runtime_state
251 .get_or_insert_with(bamboo_domain::AgentRuntimeState::default)
252 .no_human_approver = true;
253
254 ctx.persistence
256 .merge_save_runtime(&mut session)
257 .await
258 .map_err(|e| format!("failed to save scheduled session: {e}"))?;
259 if let Err(error) = ctx
260 .schedule_store
261 .bind_run_session(&job.schedule_id, &job.run_id, &session_id)
262 .await
263 {
264 tracing::warn!(
265 "failed to bind session {} to schedule run {} / {}: {}",
266 session_id,
267 job.schedule_id,
268 job.run_id,
269 error
270 );
271 }
272 ctx.sessions_cache.insert(
273 session_id.clone(),
274 Arc::new(parking_lot::RwLock::new(session.clone())),
275 );
276
277 let should_execute = job.run_config.auto_execute
279 && session
280 .messages
281 .last()
282 .map(|m| matches!(m.role, Role::User))
283 .unwrap_or(false);
284
285 tracing::info!(
286 "[schedule:{}] created session {} (auto_execute={}, model={}, model_source={}, reasoning_effort={}, reasoning_source={})",
287 job.schedule_id,
288 session_id,
289 job.run_config.auto_execute,
290 resolved_model,
291 if requested_model.is_some() {
292 "schedule.run_config.model"
293 } else {
294 "resolved"
295 },
296 resolved.reasoning_effort.map(|value| value.as_str()).unwrap_or("none"),
297 if requested_reasoning_effort.is_some() {
298 "schedule.run_config.reasoning_effort"
299 } else {
300 "resolved"
301 }
302 );
303 if !should_execute {
304 return Ok(ScheduleRunLifecycleResult::Terminal(
305 ScheduleRunStatus::Success,
306 ));
307 }
308
309 if resolved_model.trim().is_empty() {
311 let msg = "resolved model is empty".to_string();
312 session.add_message(Message::assistant(format!("❌ {msg}"), None));
313 let _ = ctx.persistence.merge_save_runtime(&mut session).await;
314 return Err(msg);
315 }
316
317 let session_tx = get_or_create_event_sender(&ctx.session_event_senders, &session_id).await;
318
319 let Some(RunnerReservation { cancel_token, .. }) =
321 try_reserve_runner(&ctx.agent_runners, &session_id, &session_tx).await
322 else {
323 return Ok(ScheduleRunLifecycleResult::Terminal(
324 ScheduleRunStatus::Skipped,
325 ));
326 };
327
328 let (mpsc_tx, _forwarder_handle) = create_event_forwarder(
329 session_id.clone(),
330 session_tx.clone(),
331 ctx.agent_runners.clone(),
332 ctx.account_feed_inbox.clone(),
333 );
334
335 let aux_fast_model = resolved.model_roster.fast_model();
342 let aux_fast_provider = resolved.model_roster.fast_model_provider();
343 let aux_background_model = resolved.model_roster.background_model();
344 let aux_background_provider = resolved.model_roster.background_model_provider();
345 let aux_summarization_model = resolved.model_roster.summarization_model();
346 let aux_summarization_provider = resolved.model_roster.summarization_model_provider();
347 let auxiliary_model_resolver = Arc::new(move || AuxiliaryModelConfig {
348 fast_model_name: aux_fast_model.clone(),
349 fast_model_provider: aux_fast_provider.clone(),
350 background_model_name: aux_background_model.clone(),
351 planning_model_name: None,
352 search_model_name: None,
353 summarization_model_name: aux_summarization_model.clone(),
354 background_model_provider: aux_background_provider.clone(),
355 summarization_model_provider: aux_summarization_provider.clone(),
356 });
357
358 let schedule_store = ctx.schedule_store.clone();
359 let schedule_id_for_state = job.schedule_id.clone();
360 let run_id_for_state = job.run_id.clone();
361 let log_session_id = session_id.clone();
362
363 let on_complete: SessionCompletionHook = Box::new(move |outcome, session| {
364 Box::pin(async move {
365 let terminal_status = if outcome.success {
366 tracing::info!(
367 "[schedule:{}][run:{}][session:{}] scheduled run completed",
368 schedule_id_for_state,
369 run_id_for_state,
370 log_session_id
371 );
372 ScheduleRunStatus::Success
373 } else {
374 let detail = outcome.error.as_deref().unwrap_or("unknown error");
375 session.add_message(Message::assistant(
378 format!("❌ Scheduled run failed: {detail}"),
379 None,
380 ));
381 tracing::warn!(
382 "[schedule:{}][run:{}][session:{}] scheduled run failed: {}",
383 schedule_id_for_state,
384 run_id_for_state,
385 log_session_id,
386 detail
387 );
388 if outcome.cancelled {
389 ScheduleRunStatus::Cancelled
390 } else {
391 ScheduleRunStatus::Failed
392 }
393 };
394
395 if let Err(error) = schedule_store
396 .mark_run_terminal(
397 &schedule_id_for_state,
398 &run_id_for_state,
399 terminal_status,
400 None,
401 )
402 .await
403 {
404 tracing::warn!(
405 "failed to mark schedule run terminal state for {} / {}: {}",
406 schedule_id_for_state,
407 run_id_for_state,
408 error
409 );
410 }
411 })
412 });
413
414 spawn_session_execution(SessionExecutionArgs {
415 agent: ctx.agent.clone(),
416 session_id,
417 session,
418 tools_override: Some(ctx.tools.clone()),
419 provider_override: None,
420 model_roster: resolved.model_roster.clone(),
421 reasoning_effort: resolved.reasoning_effort,
422 reasoning_effort_source: "schedule".to_string(),
423 auxiliary_model_resolver: Some(auxiliary_model_resolver),
424 disabled_filter_resolver: None,
427 disabled_tools: None,
428 disabled_skill_ids: None,
429 selected_skill_ids: None,
430 selected_skill_mode: None,
431 cancel_token,
432 mpsc_tx,
433 image_fallback: None,
434 gold_config: resolved.gold_config.clone(),
435 guardian_config: None,
437 guardian_spawner: None,
438 bash_resume_hook: None,
450 bash_completion_sink: None,
453 app_data_dir: ctx.app_data_dir.clone(),
454 runners: ctx.agent_runners.clone(),
455 sessions_cache: ctx.sessions_cache.clone(),
456 on_complete: Some(on_complete),
457 });
458
459 Ok(ScheduleRunLifecycleResult::BackgroundExecutionInProgress)
460}
461
462pub fn build_schedule_context(
468 base: ScheduleContext,
469 config: std::sync::Arc<tokio::sync::RwLock<bamboo_llm::Config>>,
470 provider_registry: Arc<bamboo_llm::ProviderRegistry>,
471) -> ScheduleContext {
472 ScheduleContext {
473 schedule_store: base.schedule_store,
474 agent: base.agent,
475 tools: base.tools,
476 sessions_cache: base.sessions_cache,
477 agent_runners: base.agent_runners,
478 session_event_senders: base.session_event_senders,
479 account_feed_inbox: base.account_feed_inbox,
480 app_data_dir: base.app_data_dir,
481 trigger_engine: base.trigger_engine,
482 persistence: base.persistence,
483 resolve_run_config: std::sync::Arc::new(move |job: &ScheduleRunJob| {
484 resolve_run_config_from_config(job, &config, &provider_registry)
485 }),
486 }
487}
488
489fn resolve_run_config_from_config(
490 job: &ScheduleRunJob,
491 config: &std::sync::Arc<tokio::sync::RwLock<bamboo_llm::Config>>,
492 provider_registry: &Arc<bamboo_llm::ProviderRegistry>,
493) -> ResolvedRunConfig {
494 let config_snapshot = config.try_read().map(|g| g.clone()).unwrap_or_default();
495
496 let requested_model = job
497 .run_config
498 .model
499 .as_deref()
500 .map(str::trim)
501 .filter(|v| !v.is_empty())
502 .map(|v| v.to_string());
503
504 let model = if let Some(m) = requested_model {
505 m
506 } else {
507 bamboo_engine::model_config_helper::get_schedule_model_from_config(&config_snapshot)
508 .unwrap_or_default()
509 };
510
511 let provider_name = Some(config_snapshot.effective_default_provider().to_string());
512 let provider_type = provider_name.as_deref().and_then(|name| {
513 bamboo_engine::model_config_helper::resolve_provider_type(
514 &config_snapshot,
515 name,
516 provider_registry,
517 )
518 });
519
520 let capability_provider_name = provider_name
521 .as_deref()
522 .unwrap_or(config_snapshot.effective_default_provider());
523 let areas = bamboo_engine::model_areas::resolve_global_area_models(
525 &config_snapshot,
526 capability_provider_name,
527 provider_registry,
528 );
529
530 let requested_reasoning_effort = job.run_config.reasoning_effort;
531 let reasoning_effort = requested_reasoning_effort.or(config_snapshot.get_reasoning_effort());
532
533 let global_default_prompt =
534 bamboo_engine::prompt_defaults::read_global_default_system_prompt_template();
535 let base_system_prompt = job
536 .run_config
537 .system_prompt
538 .as_deref()
539 .map(str::trim)
540 .filter(|v| !v.is_empty())
541 .unwrap_or(global_default_prompt.as_str());
542
543 let workspace_path = job
544 .run_config
545 .workspace_path
546 .as_deref()
547 .map(str::trim)
548 .filter(|v| !v.is_empty())
549 .map(ToString::to_string)
550 .or_else(|| {
551 config_snapshot
552 .get_default_work_area_path()
553 .map(|path| bamboo_config::paths::path_to_display_string(&path))
554 });
555
556 let enhance_prompt = job
557 .run_config
558 .enhance_prompt
559 .as_deref()
560 .map(str::trim)
561 .filter(|v| !v.is_empty());
562
563 let system_prompt = bamboo_engine::context::assemble_system_prompt(
564 base_system_prompt,
565 enhance_prompt,
566 workspace_path.as_deref(),
567 );
568
569 let model_roster =
570 bamboo_engine::ModelRoster::from_areas(Some(model), provider_name, provider_type, areas);
571
572 ResolvedRunConfig {
573 model_roster,
574 reasoning_effort,
575 gold_config: bamboo_engine::model_config_helper::resolve_gold_config(
576 &config_snapshot,
577 None,
578 ),
579 system_prompt,
580 base_system_prompt: base_system_prompt.to_string(),
581 workspace_path,
582 }
583}
584
585#[cfg(test)]
586mod build_context_tests {
587 use super::resolve_run_config_from_config;
588 use super::ScheduleRunJob;
589 use bamboo_config::DefaultsConfig;
590 use bamboo_config::{OpenAIConfig, ProviderConfigs};
591 use bamboo_domain::{ProviderModelRef, ScheduleRunConfig};
592 use bamboo_llm::{Config, ProviderRegistry};
593 use std::collections::HashMap;
594 use std::sync::Arc;
595 use tokio::sync::RwLock;
596
597 fn test_job() -> ScheduleRunJob {
598 ScheduleRunJob {
599 run_id: "run-1".to_string(),
600 schedule_id: "schedule-1".to_string(),
601 schedule_name: "nightly".to_string(),
602 run_config: ScheduleRunConfig::default(),
603 scheduled_for: chrono::Utc::now(),
604 claimed_at: chrono::Utc::now(),
605 was_catch_up: false,
606 }
607 }
608
609 #[test]
610 fn resolve_run_config_from_config_prefers_fast_model() {
611 let config = Config {
612 provider: "openai".to_string(),
613 defaults: None,
614 features: bamboo_config::FeatureFlags {
615 provider_model_ref: false,
616 ..Default::default()
617 },
618 providers: ProviderConfigs {
619 openai: Some(OpenAIConfig {
620 api_key: "test".to_string(),
621 api_key_encrypted: None,
622 base_url: None,
623 model: Some("gpt-4o".to_string()),
624 fast_model: Some("gpt-4o-mini".to_string()),
625 vision_model: None,
626 reasoning_effort: None,
627 responses_only_models: vec![],
628 request_overrides: None,
629 extra: Default::default(),
630 }),
631 ..ProviderConfigs::default()
632 },
633 ..Config::default()
634 };
635
636 let registry = Arc::new(ProviderRegistry::new(
637 Default::default(),
638 "openai".to_string(),
639 ));
640 let resolved =
641 resolve_run_config_from_config(&test_job(), &Arc::new(RwLock::new(config)), ®istry);
642 assert_eq!(resolved.model_roster.model.as_deref(), Some("gpt-4o-mini"));
643 }
644
645 #[test]
646 fn resolve_run_config_from_config_falls_back_to_default_model_when_fast_missing() {
647 let config = Config {
648 provider: "openai".to_string(),
649 defaults: Some(DefaultsConfig {
650 chat: ProviderModelRef::new("openai", "gpt-chat"),
651 fast: None,
652 task_summary: None,
653 vision: None,
654 memory_background: None,
655 planning: None,
656 search: None,
657 code_review: None,
658 sub_agent: None,
659 subagent_models: HashMap::new(),
660 }),
661 features: bamboo_config::FeatureFlags {
662 provider_model_ref: true,
663 ..Default::default()
664 },
665 providers: ProviderConfigs::default(),
666 ..Config::default()
667 };
668
669 let registry = Arc::new(ProviderRegistry::new(
670 Default::default(),
671 "openai".to_string(),
672 ));
673 let resolved =
674 resolve_run_config_from_config(&test_job(), &Arc::new(RwLock::new(config)), ®istry);
675 assert_eq!(resolved.model_roster.model.as_deref(), Some("gpt-chat"));
676 }
677}