From 00c8564055a5e419ea09846ae536ae0760106702 Mon Sep 17 00:00:00 2001 From: Bryan Helmkamp Date: Wed, 1 Apr 2026 22:23:28 -0400 Subject: [PATCH] Persist provider metadata from LLM handlers --- .../fabro-workflow/src/handler/agent.rs | 97 +++++++++++++++++++ .../fabro-workflow/src/handler/prompt.rs | 89 ++++++++++++++++- 2 files changed, 185 insertions(+), 1 deletion(-) diff --git a/lib/crates/fabro-workflow/src/handler/agent.rs b/lib/crates/fabro-workflow/src/handler/agent.rs index 9bc6f9c94..e4921f181 100644 --- a/lib/crates/fabro-workflow/src/handler/agent.rs +++ b/lib/crates/fabro-workflow/src/handler/agent.rs @@ -200,6 +200,33 @@ pub(crate) fn truncate(s: &str, max_chars: usize) -> &str { } } +pub(crate) async fn sync_provider_used_to_store( + stage_dir: &Path, + node_ref: &NodeVisitRef<'_>, + services: &EngineServices, +) -> Result<(), FabroError> { + let Some(ref store) = services.run_store else { + return Ok(()); + }; + + let path = stage_dir.join("provider_used.json"); + let json = match fs::read_to_string(&path).await { + Ok(json) => json, + Err(err) if err.kind() == std::io::ErrorKind::NotFound => return Ok(()), + Err(err) => { + return Err(FabroError::handler(format!( + "Failed to read provider_used.json: {err}" + ))); + } + }; + let value: serde_json::Value = serde_json::from_str(&json) + .map_err(|err| FabroError::handler(format!("Failed to parse provider_used.json: {err}")))?; + store + .put_node_provider_used(node_ref, &value) + .await + .map_err(|err| FabroError::handler(err.to_string())) +} + /// Shared simulate implementation for LLM-backed handlers (agent & prompt). /// Produces a simulated outcome with standard context updates. pub(crate) fn simulate_llm_handler(node: &Node) -> Outcome { @@ -303,6 +330,7 @@ impl Handler for AgentHandler { .await; match result { Ok(CodergenResult::Full(outcome)) => { + sync_provider_used_to_store(&stage_dir, &node_ref, services).await?; let status_json = serde_json::to_string_pretty(&outcome) .unwrap_or_else(|_| "{}".to_string()); fs::write(stage_dir.join("status.json"), &status_json).await?; @@ -339,6 +367,7 @@ impl Handler for AgentHandler { } else { fs::write(stage_dir.join("response.md"), &response_text).await?; } + sync_provider_used_to_store(&stage_dir, &node_ref, services).await?; // 7. Build and write status let mut outcome = Outcome::success(); @@ -402,13 +431,28 @@ mod tests { use super::*; use crate::event::EventEmitter; use fabro_graphviz::graph::AttrValue; + use fabro_store::{InMemoryStore, RunStore, Store}; use fabro_types::fixtures; + use std::sync::Arc; use tempfile::TempDir; fn make_services() -> EngineServices { EngineServices::test_default() } + async fn make_services_with_run_store() -> (EngineServices, Arc) { + let store = InMemoryStore::default(); + let run_store = store + .create_run(&fixtures::RUN_1, chrono::Utc::now(), None) + .await + .unwrap(); + let services = EngineServices { + run_store: Some(Arc::clone(&run_store)), + ..EngineServices::test_default() + }; + (services, run_store) + } + fn test_context() -> Context { let context = Context::new(); context.set( @@ -668,6 +712,59 @@ mod tests { ); } + #[tokio::test] + async fn codergen_handler_persists_provider_used_in_run_store() { + struct ProviderUsedBackend; + + #[async_trait] + impl CodergenBackend for ProviderUsedBackend { + async fn run( + &self, + _node: &Node, + _prompt: &str, + _context: &Context, + _thread_id: Option<&str>, + _emitter: &Arc, + stage_dir: &Path, + _sandbox: &Arc, + _tool_hooks: Option>, + ) -> Result { + std::fs::write( + stage_dir.join("provider_used.json"), + r#"{"mode":"agent","provider":"openai","model":"gpt-5.4"}"#, + ) + .unwrap(); + Ok(CodergenResult::Text { + text: "done".to_string(), + usage: None, + files_touched: Vec::new(), + last_file_touched: None, + }) + } + } + + let handler = AgentHandler::new(Some(Box::new(ProviderUsedBackend))); + let node = Node::new("step"); + let context = test_context(); + let graph = Graph::new("test"); + let tmp = TempDir::new().unwrap(); + let (services, run_store) = make_services_with_run_store().await; + + handler + .execute(&node, &context, &graph, tmp.path(), &services) + .await + .unwrap(); + + let snapshot = run_store + .get_node(&NodeVisitRef { + node_id: "step", + visit: 1, + }) + .await + .unwrap(); + assert_eq!(snapshot.provider_used.unwrap()["provider"], "openai"); + } + #[test] fn expand_variables_replaces_goal() { let mut graph = Graph::new("test"); diff --git a/lib/crates/fabro-workflow/src/handler/prompt.rs b/lib/crates/fabro-workflow/src/handler/prompt.rs index ea64b0270..7bd136870 100644 --- a/lib/crates/fabro-workflow/src/handler/prompt.rs +++ b/lib/crates/fabro-workflow/src/handler/prompt.rs @@ -15,7 +15,8 @@ use fabro_graphviz::graph::{Graph, Node}; use tokio::fs; use super::agent::{ - CodergenBackend, CodergenResult, expand_variables, extract_status_fields, truncate, + CodergenBackend, CodergenResult, expand_variables, extract_status_fields, + sync_provider_used_to_store, truncate, }; use super::{EngineServices, Handler}; @@ -114,6 +115,7 @@ impl Handler for PromptHandler { .await; match result { Ok(CodergenResult::Full(outcome)) => { + sync_provider_used_to_store(&stage_dir, &node_ref, services).await?; let status_json = serde_json::to_string_pretty(&outcome) .unwrap_or_else(|_| "{}".to_string()); fs::write(stage_dir.join("status.json"), &status_json).await?; @@ -168,6 +170,7 @@ impl Handler for PromptHandler { } else { fs::write(stage_dir.join("response.md"), &response_text).await?; } + sync_provider_used_to_store(&stage_dir, &node_ref, services).await?; // 5. Build and write status let mut outcome = Outcome::success(); @@ -200,6 +203,8 @@ impl Handler for PromptHandler { mod tests { use super::*; use fabro_graphviz::graph::AttrValue; + use fabro_store::{InMemoryStore, NodeVisitRef, RunStore, Store}; + use fabro_types::fixtures; use std::sync::Arc; use tempfile::TempDir; @@ -207,6 +212,19 @@ mod tests { EngineServices::test_default() } + async fn make_services_with_run_store() -> (EngineServices, Arc) { + let store = InMemoryStore::default(); + let run_store = store + .create_run(&fixtures::RUN_1, chrono::Utc::now(), None) + .await + .unwrap(); + let services = EngineServices { + run_store: Some(Arc::clone(&run_store)), + ..EngineServices::test_default() + }; + (services, run_store) + } + #[tokio::test] async fn prompt_handler_simulate() { let handler = PromptHandler::new(None); @@ -306,6 +324,75 @@ mod tests { assert_eq!(response_content, "one-shot response"); } + #[tokio::test] + async fn prompt_handler_persists_provider_used_in_run_store() { + use fabro_agent::Sandbox; + + struct ProviderOneShotBackend; + + #[async_trait] + impl CodergenBackend for ProviderOneShotBackend { + async fn run( + &self, + _node: &Node, + _prompt: &str, + _context: &Context, + _thread_id: Option<&str>, + _emitter: &Arc, + _stage_dir: &Path, + _sandbox: &Arc, + _tool_hooks: Option>, + ) -> Result { + panic!("run() should not be called for prompt handler"); + } + + async fn one_shot( + &self, + _node: &Node, + _prompt: &str, + _system_prompt: Option<&str>, + stage_dir: &Path, + ) -> Result { + std::fs::write( + stage_dir.join("provider_used.json"), + r#"{"mode":"prompt","provider":"openai","model":"gpt-5.4"}"#, + ) + .unwrap(); + Ok(CodergenResult::Text { + text: "one-shot response".to_string(), + usage: None, + files_touched: Vec::new(), + last_file_touched: None, + }) + } + } + + let handler = PromptHandler::new(Some(Box::new(ProviderOneShotBackend))); + let mut node = Node::new("classify"); + node.attrs.insert( + "prompt".to_string(), + AttrValue::String("Classify this".to_string()), + ); + let context = Context::new(); + let graph = Graph::new("test"); + let tmp = TempDir::new().unwrap(); + let (services, run_store) = make_services_with_run_store().await; + + handler + .execute(&node, &context, &graph, tmp.path(), &services) + .await + .unwrap(); + + let snapshot = run_store + .get_node(&NodeVisitRef { + node_id: "classify", + visit: 1, + }) + .await + .unwrap(); + assert_eq!(snapshot.provider_used.unwrap()["mode"], "prompt"); + } + struct OneShotCapturingBackend { captured_prompt: Arc>>, captured_system_prompt: Arc>>>,