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_tools: None,
425 disabled_skill_ids: None,
426 selected_skill_ids: None,
427 selected_skill_mode: None,
428 cancel_token,
429 mpsc_tx,
430 image_fallback: None,
431 gold_config: resolved.gold_config.clone(),
432 guardian_config: None,
434 guardian_spawner: None,
435 bash_resume_hook: None,
447 app_data_dir: ctx.app_data_dir.clone(),
448 runners: ctx.agent_runners.clone(),
449 sessions_cache: ctx.sessions_cache.clone(),
450 on_complete: Some(on_complete),
451 });
452
453 Ok(ScheduleRunLifecycleResult::BackgroundExecutionInProgress)
454}
455
456pub fn build_schedule_context(
462 base: ScheduleContext,
463 config: std::sync::Arc<tokio::sync::RwLock<bamboo_llm::Config>>,
464 provider_registry: Arc<bamboo_llm::ProviderRegistry>,
465) -> ScheduleContext {
466 ScheduleContext {
467 schedule_store: base.schedule_store,
468 agent: base.agent,
469 tools: base.tools,
470 sessions_cache: base.sessions_cache,
471 agent_runners: base.agent_runners,
472 session_event_senders: base.session_event_senders,
473 account_feed_inbox: base.account_feed_inbox,
474 app_data_dir: base.app_data_dir,
475 trigger_engine: base.trigger_engine,
476 persistence: base.persistence,
477 resolve_run_config: std::sync::Arc::new(move |job: &ScheduleRunJob| {
478 resolve_run_config_from_config(job, &config, &provider_registry)
479 }),
480 }
481}
482
483fn resolve_run_config_from_config(
484 job: &ScheduleRunJob,
485 config: &std::sync::Arc<tokio::sync::RwLock<bamboo_llm::Config>>,
486 provider_registry: &Arc<bamboo_llm::ProviderRegistry>,
487) -> ResolvedRunConfig {
488 let config_snapshot = config.try_read().map(|g| g.clone()).unwrap_or_default();
489
490 let requested_model = job
491 .run_config
492 .model
493 .as_deref()
494 .map(str::trim)
495 .filter(|v| !v.is_empty())
496 .map(|v| v.to_string());
497
498 let model = if let Some(m) = requested_model {
499 m
500 } else {
501 bamboo_engine::model_config_helper::get_schedule_model_from_config(&config_snapshot)
502 .unwrap_or_default()
503 };
504
505 let provider_name = Some(config_snapshot.effective_default_provider().to_string());
506 let provider_type = provider_name.as_deref().and_then(|name| {
507 bamboo_engine::model_config_helper::resolve_provider_type(
508 &config_snapshot,
509 name,
510 provider_registry,
511 )
512 });
513
514 let capability_provider_name = provider_name
515 .as_deref()
516 .unwrap_or(config_snapshot.effective_default_provider());
517 let areas = bamboo_engine::model_areas::resolve_global_area_models(
519 &config_snapshot,
520 capability_provider_name,
521 provider_registry,
522 );
523
524 let requested_reasoning_effort = job.run_config.reasoning_effort;
525 let reasoning_effort = requested_reasoning_effort.or(config_snapshot.get_reasoning_effort());
526
527 let global_default_prompt =
528 bamboo_engine::prompt_defaults::read_global_default_system_prompt_template();
529 let base_system_prompt = job
530 .run_config
531 .system_prompt
532 .as_deref()
533 .map(str::trim)
534 .filter(|v| !v.is_empty())
535 .unwrap_or(global_default_prompt.as_str());
536
537 let workspace_path = job
538 .run_config
539 .workspace_path
540 .as_deref()
541 .map(str::trim)
542 .filter(|v| !v.is_empty())
543 .map(ToString::to_string)
544 .or_else(|| {
545 config_snapshot
546 .get_default_work_area_path()
547 .map(|path| bamboo_config::paths::path_to_display_string(&path))
548 });
549
550 let enhance_prompt = job
551 .run_config
552 .enhance_prompt
553 .as_deref()
554 .map(str::trim)
555 .filter(|v| !v.is_empty());
556
557 let system_prompt = bamboo_engine::context::assemble_system_prompt(
558 base_system_prompt,
559 enhance_prompt,
560 workspace_path.as_deref(),
561 );
562
563 let model_roster =
564 bamboo_engine::ModelRoster::from_areas(Some(model), provider_name, provider_type, areas);
565
566 ResolvedRunConfig {
567 model_roster,
568 reasoning_effort,
569 gold_config: bamboo_engine::model_config_helper::resolve_gold_config(
570 &config_snapshot,
571 None,
572 ),
573 system_prompt,
574 base_system_prompt: base_system_prompt.to_string(),
575 workspace_path,
576 }
577}
578
579#[cfg(test)]
580mod build_context_tests {
581 use super::resolve_run_config_from_config;
582 use super::ScheduleRunJob;
583 use bamboo_config::DefaultsConfig;
584 use bamboo_config::{OpenAIConfig, ProviderConfigs};
585 use bamboo_domain::{ProviderModelRef, ScheduleRunConfig};
586 use bamboo_llm::{Config, ProviderRegistry};
587 use std::collections::HashMap;
588 use std::sync::Arc;
589 use tokio::sync::RwLock;
590
591 fn test_job() -> ScheduleRunJob {
592 ScheduleRunJob {
593 run_id: "run-1".to_string(),
594 schedule_id: "schedule-1".to_string(),
595 schedule_name: "nightly".to_string(),
596 run_config: ScheduleRunConfig::default(),
597 scheduled_for: chrono::Utc::now(),
598 claimed_at: chrono::Utc::now(),
599 was_catch_up: false,
600 }
601 }
602
603 #[test]
604 fn resolve_run_config_from_config_prefers_fast_model() {
605 let config = Config {
606 provider: "openai".to_string(),
607 defaults: None,
608 features: bamboo_config::FeatureFlags {
609 provider_model_ref: false,
610 ..Default::default()
611 },
612 providers: ProviderConfigs {
613 openai: Some(OpenAIConfig {
614 api_key: "test".to_string(),
615 api_key_encrypted: None,
616 base_url: None,
617 model: Some("gpt-4o".to_string()),
618 fast_model: Some("gpt-4o-mini".to_string()),
619 vision_model: None,
620 reasoning_effort: None,
621 responses_only_models: vec![],
622 request_overrides: None,
623 extra: Default::default(),
624 }),
625 ..ProviderConfigs::default()
626 },
627 ..Config::default()
628 };
629
630 let registry = Arc::new(ProviderRegistry::new(
631 Default::default(),
632 "openai".to_string(),
633 ));
634 let resolved =
635 resolve_run_config_from_config(&test_job(), &Arc::new(RwLock::new(config)), ®istry);
636 assert_eq!(resolved.model_roster.model.as_deref(), Some("gpt-4o-mini"));
637 }
638
639 #[test]
640 fn resolve_run_config_from_config_falls_back_to_default_model_when_fast_missing() {
641 let config = Config {
642 provider: "openai".to_string(),
643 defaults: Some(DefaultsConfig {
644 chat: ProviderModelRef::new("openai", "gpt-chat"),
645 fast: None,
646 task_summary: None,
647 vision: None,
648 memory_background: None,
649 planning: None,
650 search: None,
651 code_review: None,
652 sub_agent: None,
653 subagent_models: HashMap::new(),
654 }),
655 features: bamboo_config::FeatureFlags {
656 provider_model_ref: true,
657 ..Default::default()
658 },
659 providers: ProviderConfigs::default(),
660 ..Config::default()
661 };
662
663 let registry = Arc::new(ProviderRegistry::new(
664 Default::default(),
665 "openai".to_string(),
666 ));
667 let resolved =
668 resolve_run_config_from_config(&test_job(), &Arc::new(RwLock::new(config)), ®istry);
669 assert_eq!(resolved.model_roster.model.as_deref(), Some("gpt-chat"));
670 }
671}