use std::collections::{BTreeMap, HashMap}; use std::hash::{Hash, Hasher}; use std::path::PathBuf; use std::sync::Arc; use std::time::Duration; use fabro_agent::Sandbox; use fabro_auth::{CredentialSource, test_support as auth_test_support}; use fabro_graphviz::graph::Graph as GvGraph; use fabro_interview::AutoApproveInterviewer; use fabro_model::Catalog; #[cfg(feature = "test-support")] use fabro_model::ProviderId; use fabro_store::{ArtifactStore, RunProjection, test_support as store_test_support}; use object_store::local::LocalFileSystem; use crate::artifact_upload::ArtifactSink; use crate::error::{Error, Result}; use crate::event::{Emitter, Event, StoreProgressLogger, append_event}; use crate::handler::HandlerRegistry; use crate::outcome::Outcome; use crate::pipeline; use crate::pipeline::types::{Executed, Initialized}; use crate::pipeline::{billing_from_projection, build_terminal_event}; use crate::records::Checkpoint; use crate::run_metadata::RunMetadataRuntime; use crate::run_options::RunOptions; use crate::sandbox_git_runtime::SandboxGitRuntime; use crate::services::{EngineServices, RunLocations, RunServices}; use crate::stage_execution::StageExecutionTracker; #[cfg(feature = "test-support")] pub(crate) fn test_configured_provider_ids( catalog: &Catalog, configured_provider_ids: Vec, assume_ready: bool, ) -> Vec { if assume_ready { catalog.all_provider_ids().into_iter().collect() } else { configured_provider_ids } } /// These helpers stop at EXECUTE, so they emit the terminal event here to /// keep test consumers seeing the same end-of-run signal as production /// (FINALIZE). /// /// The first flush is needed because `StoreProgressLogger` forwards events /// through an mpsc channel — without it, billing would read from a stale /// checkpoint. The second flush ensures the just-emitted terminal event is /// persisted before tests reopen the run store. async fn execute_and_emit_terminal(initialized: InitializedState) -> Executed { let executed = Box::pin(pipeline::execute(initialized.initialized)).await; initialized .store_logger .flush() .await .expect("test run events should persist"); let state = executed.engine.run.run_store.state().await.ok(); let billing = state.as_ref().and_then(billing_from_projection); let event = build_terminal_event( &executed.outcome, fabro_types::RunTiming::wall_only(executed.wall_time_ms), 0, None, None, None, billing, ); executed.engine.run.emitter.emit(&event); initialized .store_logger .flush() .await .expect("test run events should persist"); executed } /// Construct a fully-populated `BilledModelUsage` for tests. Centralised so /// callers don't keep rebuilding the same JSON skeleton. #[must_use] pub fn test_usage( model_id: &str, input_tokens: i64, output_tokens: i64, ) -> fabro_types::BilledModelUsage { serde_json::from_value(serde_json::json!({ "input": { "usage": { "model": { "provider": "openai", "model_id": model_id }, "tokens": { "input_tokens": input_tokens, "output_tokens": output_tokens } }, "facts": { "algorithm": "openai" } }, "total_usd_micros": input_tokens + output_tokens })) .expect("test_usage JSON must deserialise") } /// Append the `RunStartRequested → RunRunnable → RunStarting → RunRunning` /// sequence so subsequent calls observe the run as live. pub async fn mark_run_running(run_store: &fabro_store::RunDatabase, run_id: &fabro_types::RunId) { append_event(run_store, run_id, &Event::RunStartRequested { resume: false, actor: None, }) .await .expect("seed run.start_requested"); append_event(run_store, run_id, &Event::RunRunnable { source: fabro_types::RunRunnableSource::StartRequested, actor: None, }) .await .expect("seed run.runnable"); append_event(run_store, run_id, &Event::RunStarting) .await .expect("seed run.starting"); append_event(run_store, run_id, &Event::RunRunning) .await .expect("seed run.running"); } /// Record every event the emitter publishes, for assertions after a run. pub fn collect_events(emitter: &Emitter) -> Arc>> { let events = Arc::new(std::sync::Mutex::new(Vec::new())); let captured = Arc::clone(&events); emitter.on_event(move |event| captured.lock().unwrap().push(event.clone())); events } pub fn test_store_dir(run_dir: &std::path::Path) -> PathBuf { let mut hasher = std::collections::hash_map::DefaultHasher::new(); std::process::id().hash(&mut hasher); run_dir.hash(&mut hasher); std::env::temp_dir() .join("fabro-test-run-stores") .join(format!("{:016x}", hasher.finish())) } struct InitializedOptions { hook_runner: Option>, env: HashMap, checkpoint: Option, llm_source: Option>, } struct InitializedState { initialized: Initialized, store_logger: StoreProgressLogger, } fn bound_emitter(run_id: fabro_types::RunId, observer: &Arc) -> Arc { let emitter = Arc::new(Emitter::new(run_id)); let observer_clone = Arc::clone(observer); emitter.on_event(move |event| observer_clone.dispatch_run_event(event)); emitter } async fn initialized( registry: HandlerRegistry, emitter: Arc, sandbox: Arc, graph: &GvGraph, run_options: &RunOptions, options: InitializedOptions, ) -> InitializedState { std::fs::create_dir_all(&run_options.run_dir).expect("failed to create run dir"); let store_dir = test_store_dir(&run_options.run_dir); let _ = std::fs::remove_dir_all(&store_dir); for database_path in [ store_test_support::test_blob_store_path(&store_dir), store_test_support::test_run_summary_store_path(&store_dir), ] { for suffix in ["", "-wal", "-shm"] { let mut sibling = database_path.clone().into_os_string(); sibling.push(suffix); let _ = std::fs::remove_file(sibling); } } std::fs::create_dir_all(&store_dir).expect("failed to create local test run store dir"); let store = Arc::new(store_test_support::test_database_at( Arc::new( LocalFileSystem::new_with_prefix(&store_dir) .expect("failed to create local test run store"), ), "", Duration::from_millis(1), None, &store_dir, )); let inner_store = store .create_run(&run_options.run_id) .await .expect("failed to create slate-backed test run store"); let run_store = inner_store; append_event(&run_store, &run_options.run_id, &Event::RunCreated { run_id: run_options.run_id, title: None, settings: serde_json::to_value(&run_options.settings) .expect("failed to serialize settings"), graph: serde_json::to_value(graph).expect("failed to serialize graph"), workflow_source: None, labels: run_options .labels .clone() .into_iter() .collect::>(), source_directory: Some(sandbox.working_directory().to_string()), workflow_slug: run_options.workflow_slug.clone(), workflow_version_id: None, target: None, automation: None, provenance: fabro_types::RunProvenance { server: None, client: None, subject: fabro_types::Principal::System { system_kind: fabro_types::SystemActorKind::Engine, }, }, manifest_blob: None, spec_blob: None, git: run_options.pre_run_git.clone(), fork_source_ref: run_options.fork_source_ref.clone(), retried_from: None, parent_id: None, web_url: None, }) .await .expect("failed to seed run.created event in run store"); append_event(&run_store, &run_options.run_id, &Event::RunRunnable { source: fabro_types::RunRunnableSource::StartRequested, actor: None, }) .await .expect("failed to seed run.runnable event in run store"); append_event(&run_store, &run_options.run_id, &Event::RunStarting) .await .expect("failed to seed run.starting event in run store"); let emitter = bound_emitter(run_options.run_id, &emitter); let store_logger = StoreProgressLogger::new(run_store.clone()); store_logger.register(emitter.as_ref()); let artifact_store = ArtifactStore::new( Arc::new( LocalFileSystem::new_with_prefix(&store_dir) .expect("failed to create local test artifact store"), ), "artifacts", ); let locations = RunLocations::for_sandbox(None, sandbox.as_ref(), run_options.run_dir.clone()); InitializedState { initialized: Initialized { graph: graph.clone(), source: String::new(), run_options: run_options.clone(), checkpoint: options.checkpoint, seed_context: None, on_node: None, artifact_sink: Some(ArtifactSink::Store(artifact_store)), run_control: None, engine: Arc::new(EngineServices { run: RunServices::new( run_store.into(), emitter, sandbox, options.hook_runner, locations, run_options.cancel_token.clone(), fabro_model::ProviderId::anthropic(), "claude-sonnet-4-6".to_string(), options .llm_source .unwrap_or_else(auth_test_support::vault_only_credential_source), Arc::new(Catalog::from_builtin().expect("default catalog should build")), Arc::new(SandboxGitRuntime::new()), Arc::new(RunMetadataRuntime::new()), None, StageExecutionTracker::default(), ), registry: Arc::new(registry), interviewer: Arc::new(AutoApproveInterviewer::engine()), base_env: options.env, github_token: None, inputs: run_options.settings.run.inputs.clone(), dry_run: run_options.dry_run_enabled(), workflow_path: None, workflow_bundle: None, }), model: String::new(), }, store_logger, } } pub async fn run_graph( registry: HandlerRegistry, emitter: Arc, sandbox: Arc, graph: &GvGraph, run_options: &RunOptions, ) -> Result { let initialized = initialized( registry, emitter, sandbox, graph, run_options, InitializedOptions { hook_runner: None, env: HashMap::new(), checkpoint: None, llm_source: None, }, ) .await; let executed = execute_and_emit_terminal(initialized).await; executed.outcome } pub async fn run_graph_with_state( registry: HandlerRegistry, emitter: Arc, sandbox: Arc, graph: &GvGraph, run_options: &RunOptions, ) -> Result<(Outcome, RunProjection)> { let initialized = initialized( registry, emitter, sandbox, graph, run_options, InitializedOptions { hook_runner: None, env: HashMap::new(), checkpoint: None, llm_source: None, }, ) .await; let executed = execute_and_emit_terminal(initialized).await; let outcome = executed.outcome?; let state = executed .engine .run .run_store .state() .await .map_err(|err| Error::engine(err.to_string()))?; Ok((outcome, state)) } pub async fn run_graph_with_hooks( registry: HandlerRegistry, emitter: Arc, sandbox: Arc, graph: &GvGraph, run_options: &RunOptions, hook_runner: Arc, env: Option>, ) -> Result { let initialized = initialized( registry, emitter, sandbox, graph, run_options, InitializedOptions { hook_runner: Some(hook_runner), env: env.unwrap_or_default(), checkpoint: None, llm_source: None, }, ) .await; let executed = execute_and_emit_terminal(initialized).await; executed.outcome } pub async fn run_graph_with_hooks_and_state( registry: HandlerRegistry, emitter: Arc, sandbox: Arc, graph: &GvGraph, run_options: &RunOptions, hook_runner: Arc, env: Option>, ) -> Result<(Outcome, RunProjection)> { let initialized = initialized( registry, emitter, sandbox, graph, run_options, InitializedOptions { hook_runner: Some(hook_runner), env: env.unwrap_or_default(), checkpoint: None, llm_source: None, }, ) .await; let executed = execute_and_emit_terminal(initialized).await; let outcome = executed.outcome?; let state = executed .engine .run .run_store .state() .await .map_err(|err| Error::engine(err.to_string()))?; Ok((outcome, state)) } pub async fn run_graph_from_checkpoint( registry: HandlerRegistry, emitter: Arc, sandbox: Arc, graph: &GvGraph, run_options: &RunOptions, checkpoint: &Checkpoint, ) -> Result { let initialized = initialized( registry, emitter, sandbox, graph, run_options, InitializedOptions { hook_runner: None, env: HashMap::new(), checkpoint: Some(checkpoint.clone()), llm_source: None, }, ) .await; let executed = execute_and_emit_terminal(initialized).await; executed.outcome } pub async fn run_graph_from_checkpoint_with_state( registry: HandlerRegistry, emitter: Arc, sandbox: Arc, graph: &GvGraph, run_options: &RunOptions, checkpoint: &Checkpoint, ) -> Result<(Outcome, RunProjection)> { let initialized = initialized( registry, emitter, sandbox, graph, run_options, InitializedOptions { hook_runner: None, env: HashMap::new(), checkpoint: Some(checkpoint.clone()), llm_source: None, }, ) .await; let executed = execute_and_emit_terminal(initialized).await; let outcome = executed.outcome?; let state = executed .engine .run .run_store .state() .await .map_err(|err| Error::engine(err.to_string()))?; Ok((outcome, state)) } pub async fn run_graph_with_state_and_llm_source( registry: HandlerRegistry, emitter: Arc, sandbox: Arc, graph: &GvGraph, run_options: &RunOptions, llm_source: Arc, ) -> Result<(Outcome, RunProjection)> { let initialized = initialized( registry, emitter, sandbox, graph, run_options, InitializedOptions { hook_runner: None, env: HashMap::new(), checkpoint: None, llm_source: Some(llm_source), }, ) .await; let executed = pipeline::execute(initialized.initialized).await; initialized .store_logger .flush() .await .expect("test run events should persist"); let outcome = executed.outcome?; let state = executed .engine .run .run_store .state() .await .map_err(|err| Error::engine(err.to_string()))?; Ok((outcome, state)) } pub struct WorkflowRunner { registry: std::sync::Mutex>, emitter: Arc, sandbox: Arc, } impl WorkflowRunner { #[must_use] pub fn new( registry: HandlerRegistry, emitter: Arc, sandbox: Arc, ) -> Self { Self { registry: std::sync::Mutex::new(Some(registry)), emitter, sandbox, } } pub async fn run(&self, graph: &GvGraph, run_options: &RunOptions) -> Result { let registry = self .registry .lock() .unwrap() .take() .expect("WorkflowRunner may only be used once"); Box::pin(run_graph( registry, Arc::clone(&self.emitter), Arc::clone(&self.sandbox), graph, run_options, )) .await } pub async fn run_with_state( &self, graph: &GvGraph, run_options: &RunOptions, ) -> Result<(Outcome, RunProjection)> { let registry = self .registry .lock() .unwrap() .take() .expect("WorkflowRunner may only be used once"); Box::pin(run_graph_with_state( registry, Arc::clone(&self.emitter), Arc::clone(&self.sandbox), graph, run_options, )) .await } pub async fn run_with_state_and_llm_source( &self, graph: &GvGraph, run_options: &RunOptions, llm_source: Arc, ) -> Result<(Outcome, RunProjection)> { let registry = self .registry .lock() .unwrap() .take() .expect("WorkflowRunner may only be used once"); Box::pin(run_graph_with_state_and_llm_source( registry, Arc::clone(&self.emitter), Arc::clone(&self.sandbox), graph, run_options, llm_source, )) .await } pub async fn run_from_checkpoint( &self, graph: &GvGraph, run_options: &RunOptions, checkpoint: &Checkpoint, ) -> Result { let registry = self .registry .lock() .unwrap() .take() .expect("WorkflowRunner may only be used once"); Box::pin(run_graph_from_checkpoint( registry, Arc::clone(&self.emitter), Arc::clone(&self.sandbox), graph, run_options, checkpoint, )) .await } pub async fn run_from_checkpoint_with_state( &self, graph: &GvGraph, run_options: &RunOptions, checkpoint: &Checkpoint, ) -> Result<(Outcome, RunProjection)> { let registry = self .registry .lock() .unwrap() .take() .expect("WorkflowRunner may only be used once"); Box::pin(run_graph_from_checkpoint_with_state( registry, Arc::clone(&self.emitter), Arc::clone(&self.sandbox), graph, run_options, checkpoint, )) .await } }