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