From 40ed0b5ae4f087f6ce65dc96fa23505c5fe8f863 Mon Sep 17 00:00:00 2001 From: Bryan Helmkamp Date: Mon, 11 May 2026 16:04:44 -0400 Subject: [PATCH] refactor: simplify ACP and one-shot plumbing Collapse CodergenBackend one-shot arguments into a request object, reuse the shared fake ACP fixture in workflow tests, and centralize Docker exec create/start boilerplate. --- lib/crates/fabro-acp/src/test_support.rs | 3 + lib/crates/fabro-sandbox/src/docker.rs | 102 +++++++----- lib/crates/fabro-workflow/Cargo.toml | 1 + .../fabro-workflow/src/handler/agent.rs | 23 +-- .../fabro-workflow/src/handler/llm/acp.rs | 122 +++++--------- .../fabro-workflow/src/handler/llm/api.rs | 19 +-- .../fabro-workflow/src/handler/llm/cli.rs | 152 +++++------------- .../fabro-workflow/src/handler/prompt.rs | 51 ++---- .../fabro-workflow/src/pipeline/initialize.rs | 40 +---- .../fabro-workflow/tests/it/integration.rs | 59 +------ 10 files changed, 193 insertions(+), 379 deletions(-) diff --git a/lib/crates/fabro-acp/src/test_support.rs b/lib/crates/fabro-acp/src/test_support.rs index 167d639d6..854c22a9c 100644 --- a/lib/crates/fabro-acp/src/test_support.rs +++ b/lib/crates/fabro-acp/src/test_support.rs @@ -54,6 +54,9 @@ for line in sys.stdin: record.write(json.dumps(message.get("params", {}), separators=(",", ":"))) respond(message, {"sessionId": session_id}) elif method == "session/prompt": + if os.environ.get("ACP_PROMPT_RECORD"): + with open(os.environ["ACP_PROMPT_RECORD"], "w", encoding="utf-8") as record: + record.write(json.dumps(message.get("params", {}))) mode = os.environ.get("ACP_MODE", "normal") if mode == "timeout": time.sleep(60) diff --git a/lib/crates/fabro-sandbox/src/docker.rs b/lib/crates/fabro-sandbox/src/docker.rs index 6259547d7..d85bc138a 100644 --- a/lib/crates/fabro-sandbox/src/docker.rs +++ b/lib/crates/fabro-sandbox/src/docker.rs @@ -219,17 +219,15 @@ impl DockerSandbox { ..Default::default() }; - let exec_instance = self - .docker - .create_exec(container_id, exec_opts) - .await - .map_err(|e| crate::Error::context("Failed to create exec", e))?; - - let start_result = self - .docker - .start_exec(&exec_instance.id, None) - .await - .map_err(|e| crate::Error::context("Failed to start exec", e))?; + let (exec_id, start_result) = create_and_start_exec( + &self.docker, + container_id, + exec_opts, + None, + "Failed to create exec", + "Failed to start exec", + ) + .await?; let mut stdout = String::new(); let mut stderr = String::new(); @@ -253,7 +251,7 @@ impl DockerSandbox { let inspect = self .docker - .inspect_exec(&exec_instance.id) + .inspect_exec(&exec_id) .await .map_err(|e| crate::Error::context("Failed to inspect exec", e))?; @@ -281,15 +279,15 @@ impl DockerSandbox { ..Default::default() }; - let exec_instance = docker - .create_exec(&container_id, exec_opts) - .await - .map_err(|e| crate::Error::context("Failed to create exec", e))?; - - let start_result = docker - .start_exec(&exec_instance.id, None) - .await - .map_err(|e| crate::Error::context("Failed to start exec", e))?; + let (exec_id, start_result) = create_and_start_exec( + &docker, + &container_id, + exec_opts, + None, + "Failed to create exec", + "Failed to start exec", + ) + .await?; let mut stdout = Vec::new(); let mut stderr = Vec::new(); @@ -314,7 +312,7 @@ impl DockerSandbox { } let inspect = docker - .inspect_exec(&exec_instance.id) + .inspect_exec(&exec_id) .await .map_err(|e| crate::Error::context("Failed to inspect exec", e))?; @@ -820,6 +818,27 @@ fn docker_stdio_exec_options( ) } +async fn create_and_start_exec( + docker: &Docker, + container_id: &str, + exec_options: CreateExecOptions, + start_options: Option, + create_context: &'static str, + start_context: &'static str, +) -> crate::Result<(String, StartExecResults)> { + let exec_instance = docker + .create_exec(container_id, exec_options) + .await + .map_err(|err| crate::Error::context(create_context, err))?; + let exec_id = exec_instance.id; + let start_result = docker + .start_exec(&exec_id, start_options) + .await + .map_err(|err| crate::Error::context(start_context, err))?; + + Ok((exec_id, start_result)) +} + async fn request_docker_exec_stop_with( docker: &Docker, container_id: &str, @@ -833,14 +852,15 @@ async fn request_docker_exec_stop_with( working_dir: Some("/".to_string()), ..Default::default() }; - let exec_instance = docker - .create_exec(container_id, exec_opts) - .await - .map_err(|e| crate::Error::context("Failed to create Docker exec stop request", e))?; - let start_result = docker - .start_exec(&exec_instance.id, None) - .await - .map_err(|e| crate::Error::context("Failed to start Docker exec stop request", e))?; + let (exec_id, start_result) = create_and_start_exec( + docker, + container_id, + exec_opts, + None, + "Failed to create Docker exec stop request", + "Failed to start Docker exec stop request", + ) + .await?; let mut stdout = String::new(); let mut stderr = String::new(); @@ -865,7 +885,7 @@ async fn request_docker_exec_stop_with( } let inspect = docker - .inspect_exec(&exec_instance.id) + .inspect_exec(&exec_id) .await .map_err(|e| crate::Error::context("Failed to inspect Docker exec stop request", e))?; let exit_code = inspect @@ -1529,17 +1549,15 @@ impl Sandbox for DockerSandbox { docker_stdio_exec_options(controlled_command, effective_dir, env); let container_id = self.container_id()?.to_string(); - let exec_instance = self - .docker - .create_exec(&container_id, create_opts) - .await - .map_err(|e| crate::Error::context("Failed to create Docker stdio exec", e))?; - let exec_id = exec_instance.id.clone(); - let start_result = self - .docker - .start_exec(&exec_id, Some(start_opts)) - .await - .map_err(|e| crate::Error::context("Failed to start Docker stdio exec", e))?; + let (exec_id, start_result) = create_and_start_exec( + &self.docker, + &container_id, + create_opts, + Some(start_opts), + "Failed to create Docker stdio exec", + "Failed to start Docker stdio exec", + ) + .await?; let StartExecResults::Attached { mut output, input } = start_result else { return Err(crate::Error::message( diff --git a/lib/crates/fabro-workflow/Cargo.toml b/lib/crates/fabro-workflow/Cargo.toml index a2f106c3c..d535328b0 100644 --- a/lib/crates/fabro-workflow/Cargo.toml +++ b/lib/crates/fabro-workflow/Cargo.toml @@ -71,6 +71,7 @@ toml.workspace = true fabro-vault = { path = "../fabro-vault" } [dev-dependencies] base64.workspace = true +fabro-acp = { path = "../fabro-acp", features = ["test-support"] } fabro-sandbox = { path = "../fabro-sandbox", features = ["daytona", "docker", "test-support"] } fabro-mcp = { path = "../fabro-mcp" } tokio = { workspace = true, features = ["test-util", "macros"] } diff --git a/lib/crates/fabro-workflow/src/handler/agent.rs b/lib/crates/fabro-workflow/src/handler/agent.rs index 936d89d79..c5a2561d2 100644 --- a/lib/crates/fabro-workflow/src/handler/agent.rs +++ b/lib/crates/fabro-workflow/src/handler/agent.rs @@ -28,10 +28,20 @@ pub enum CodergenResult { Full(Outcome), } +pub struct OneShotRequest<'a> { + pub node: &'a Node, + pub prompt: &'a str, + pub system_prompt: Option<&'a str>, + pub emitter: &'a Arc, + pub stage_scope: &'a StageScope, + pub sandbox: &'a Arc, + pub cancel_token: CancellationToken, +} + /// Backend interface for LLM execution in codergen nodes. #[allow( clippy::too_many_arguments, - reason = "Codergen backends need the node, prompt, context, and runtime handles separately." + reason = "Codergen run mode needs the node, prompt, context, and runtime handles separately." )] #[async_trait] pub trait CodergenBackend: Send + Sync { @@ -49,16 +59,7 @@ pub trait CodergenBackend: Send + Sync { ) -> Result; /// Run a single LLM call with no tools (one_shot mode). - async fn one_shot( - &self, - _node: &Node, - _prompt: &str, - _system_prompt: Option<&str>, - _emitter: &Arc, - _stage_scope: &StageScope, - _sandbox: &Arc, - _cancel_token: CancellationToken, - ) -> Result { + async fn one_shot(&self, _request: OneShotRequest<'_>) -> Result { Err(Error::Validation( "one_shot mode not supported by this backend".into(), )) diff --git a/lib/crates/fabro-workflow/src/handler/llm/acp.rs b/lib/crates/fabro-workflow/src/handler/llm/acp.rs index fdae41a76..03197c6e0 100644 --- a/lib/crates/fabro-workflow/src/handler/llm/acp.rs +++ b/lib/crates/fabro-workflow/src/handler/llm/acp.rs @@ -12,7 +12,7 @@ use fabro_model::Provider; use fabro_util::time::elapsed_ms; use tokio_util::sync::CancellationToken; -use super::super::agent::{CodergenBackend, CodergenResult}; +use super::super::agent::{CodergenBackend, CodergenResult, OneShotRequest}; use super::changed_files; use super::cli::AgentCli; use super::launch_env::{AgentLaunchEnvRequest, resolve_agent_launch_env}; @@ -222,22 +222,20 @@ impl CodergenBackend for AgentAcpBackend { .await } - async fn one_shot( - &self, - node: &Node, - prompt: &str, - system_prompt: Option<&str>, - emitter: &Arc, - stage_scope: &StageScope, - sandbox: &Arc, - cancel_token: CancellationToken, - ) -> Result { - let prompt = match system_prompt.filter(|prompt| !prompt.is_empty()) { - Some(system_prompt) => format!("System:\n{system_prompt}\n\nUser:\n{prompt}"), - None => prompt.to_string(), + async fn one_shot(&self, request: OneShotRequest<'_>) -> Result { + let prompt = match request.system_prompt.filter(|prompt| !prompt.is_empty()) { + Some(system_prompt) => format!("System:\n{system_prompt}\n\nUser:\n{}", request.prompt), + None => request.prompt.to_string(), }; - self.run_turn(node, prompt, emitter, stage_scope, sandbox, cancel_token) - .await + self.run_turn( + request.node, + prompt, + request.emitter, + request.stage_scope, + request.sandbox, + request.cancel_token, + ) + .await } } @@ -285,6 +283,7 @@ mod tests { use std::collections::HashMap; use std::sync::{Arc, Mutex}; + use fabro_acp::test_support::fake_acp_agent_script; use fabro_agent::{LocalSandbox, Sandbox, shell_quote}; use fabro_graphviz::graph::{AttrValue, Node}; use fabro_model::Provider; @@ -295,14 +294,14 @@ mod tests { use super::AgentAcpBackend; use crate::context::Context; use crate::event::{Emitter, StageScope}; - use crate::handler::agent::{CodergenBackend, CodergenResult}; + use crate::handler::agent::{CodergenBackend, CodergenResult, OneShotRequest}; #[tokio::test] async fn acp_backend_run_sends_prompt_and_returns_text() { let tempdir = tempfile::tempdir().unwrap(); init_git(tempdir.path()); let script_path = tempdir.path().join("fake_acp_agent.py"); - tokio::fs::write(&script_path, fake_agent_script()) + tokio::fs::write(&script_path, fake_acp_agent_script()) .await .unwrap(); @@ -325,7 +324,10 @@ mod tests { )), ); - let backend = AgentAcpBackend::new_from_env("fake-acp".to_string(), Provider::OpenAi); + let backend = + AgentAcpBackend::new_from_env("fake-acp".to_string(), Provider::OpenAi).with_env( + HashMap::from([("ACP_MODE".to_string(), "write_file".to_string())]), + ); let sandbox: Arc = Arc::new(LocalSandbox::new(tempdir.path().to_path_buf())); let result = backend .run( @@ -358,7 +360,7 @@ mod tests { let tempdir = tempfile::tempdir().unwrap(); let script_path = tempdir.path().join("fake_acp_agent.py"); let prompt_record_path = tempdir.path().join("prompt.json"); - tokio::fs::write(&script_path, fake_agent_script()) + tokio::fs::write(&script_path, fake_acp_agent_script()) .await .unwrap(); @@ -378,21 +380,27 @@ mod tests { ); let backend = AgentAcpBackend::new_from_env("fake-acp".to_string(), Provider::OpenAi) - .with_env(HashMap::from([( - "ACP_PROMPT_RECORD".to_string(), - prompt_record_path.to_string_lossy().into_owned(), - )])); + .with_env(HashMap::from([ + ( + "ACP_PROMPT_RECORD".to_string(), + prompt_record_path.to_string_lossy().into_owned(), + ), + ("ACP_MODE".to_string(), "write_file".to_string()), + ])); let sandbox: Arc = Arc::new(LocalSandbox::new(tempdir.path().to_path_buf())); + let emitter = Arc::new(Emitter::default()); + let context = Context::new(); + let stage_scope = StageScope::for_handler(&context, "prompt"); let result = backend - .one_shot( - &node, - "User prompt", - Some("System prompt"), - &Arc::new(Emitter::default()), - &StageScope::for_handler(&Context::new(), "prompt"), - &sandbox, - CancellationToken::new(), - ) + .one_shot(OneShotRequest { + node: &node, + prompt: "User prompt", + system_prompt: Some("System prompt"), + emitter: &emitter, + stage_scope: &stage_scope, + sandbox: &sandbox, + cancel_token: CancellationToken::new(), + }) .await .unwrap(); @@ -411,7 +419,7 @@ mod tests { async fn acp_backend_cancelled_stop_reason_maps_to_cancelled_error() { let tempdir = tempfile::tempdir().unwrap(); let script_path = tempdir.path().join("fake_acp_agent.py"); - tokio::fs::write(&script_path, fake_agent_script()) + tokio::fs::write(&script_path, fake_acp_agent_script()) .await .unwrap(); @@ -456,7 +464,7 @@ mod tests { async fn acp_started_event_omits_json_command_env_values() { let tempdir = tempfile::tempdir().unwrap(); let script_path = tempdir.path().join("fake_acp_agent.py"); - tokio::fs::write(&script_path, fake_agent_script()) + tokio::fs::write(&script_path, fake_acp_agent_script()) .await .unwrap(); @@ -561,50 +569,6 @@ mod tests { ); } - fn fake_agent_script() -> &'static str { - r#" -import json -import os -import sys - -session_id = "sess-1" - -def send(message): - print(json.dumps(message), flush=True) - -def respond(message, result): - send({"jsonrpc": "2.0", "id": message["id"], "result": result}) - -for line in sys.stdin: - message = json.loads(line) - method = message.get("method") - if method == "initialize": - respond(message, {"protocolVersion": 1, "agentCapabilities": {}}) - elif method == "session/new": - respond(message, {"sessionId": session_id}) - elif method == "session/prompt": - if os.environ.get("ACP_PROMPT_RECORD"): - with open(os.environ["ACP_PROMPT_RECORD"], "w", encoding="utf-8") as record: - record.write(json.dumps(message.get("params", {}))) - with open("hello.txt", "w", encoding="utf-8") as file: - file.write("hello from sandbox\n") - for text in ["hello ", "from acp"]: - send({ - "jsonrpc": "2.0", - "method": "session/update", - "params": { - "sessionId": session_id, - "update": { - "sessionUpdate": "agent_message_chunk", - "content": {"type": "text", "text": text} - } - } - }) - respond(message, {"stopReason": os.environ.get("ACP_STOP_REASON", "end_turn")}) - break -"# - } - #[expect( clippy::disallowed_methods, reason = "unit test initializes an isolated git repository with the system git binary" diff --git a/lib/crates/fabro-workflow/src/handler/llm/api.rs b/lib/crates/fabro-workflow/src/handler/llm/api.rs index 4c65fe462..51aae4d42 100644 --- a/lib/crates/fabro-workflow/src/handler/llm/api.rs +++ b/lib/crates/fabro-workflow/src/handler/llm/api.rs @@ -19,7 +19,7 @@ use tokio::sync::Mutex as TokioMutex; use tokio::task::JoinHandle; use tokio_util::sync::CancellationToken; -use super::super::agent::{CodergenBackend, CodergenResult}; +use super::super::agent::{CodergenBackend, CodergenResult, OneShotRequest}; use super::activation_lease::{ActivationLease, ActivationLeaseOptions}; use crate::context::keys::Fidelity; use crate::context::{Context, WorkflowContext}; @@ -476,16 +476,13 @@ impl CodergenBackend for AgentApiBackend { self.shutdown_cached_sessions(emitter); } - async fn one_shot( - &self, - node: &Node, - prompt: &str, - system_prompt: Option<&str>, - emitter: &Arc, - stage_scope: &StageScope, - _sandbox: &Arc, - _cancel_token: CancellationToken, - ) -> Result { + async fn one_shot(&self, request: OneShotRequest<'_>) -> Result { + let node = request.node; + let prompt = request.prompt; + let system_prompt = request.system_prompt; + let emitter = request.emitter; + let stage_scope = request.stage_scope; + let client = Client::from_source(self.source.as_ref()) .await .map_err(|e| Error::handler_with_source("Failed to create LLM client", &e))?; diff --git a/lib/crates/fabro-workflow/src/handler/llm/cli.rs b/lib/crates/fabro-workflow/src/handler/llm/cli.rs index e5bd49f85..f721f3cac 100644 --- a/lib/crates/fabro-workflow/src/handler/llm/cli.rs +++ b/lib/crates/fabro-workflow/src/handler/llm/cli.rs @@ -38,7 +38,7 @@ fn cli_failure_detail(stdout: &str, stderr: &str, command: &str) -> String { } } -use super::super::agent::{CodergenBackend, CodergenResult}; +use super::super::agent::{CodergenBackend, CodergenResult, OneShotRequest}; use super::acp::AgentAcpBackend; use super::changed_files; use super::launch_env::{AgentLaunchEnvRequest, resolve_agent_launch_env}; @@ -779,43 +779,10 @@ impl CodergenBackend for BackendRouter { } } - async fn one_shot( - &self, - node: &Node, - prompt: &str, - system_prompt: Option<&str>, - emitter: &Arc, - stage_scope: &StageScope, - sandbox: &Arc, - cancel_token: CancellationToken, - ) -> Result { - match Self::select_one_shot_backend(node)? { - LlmBackend::Acp => { - self.acp - .one_shot( - node, - prompt, - system_prompt, - emitter, - stage_scope, - sandbox, - cancel_token, - ) - .await - } - LlmBackend::Api | LlmBackend::Cli => { - self.api - .one_shot( - node, - prompt, - system_prompt, - emitter, - stage_scope, - sandbox, - cancel_token, - ) - .await - } + async fn one_shot(&self, request: OneShotRequest<'_>) -> Result { + match Self::select_one_shot_backend(request.node)? { + LlmBackend::Acp => self.acp.one_shot(request).await, + LlmBackend::Api | LlmBackend::Cli => self.api.one_shot(request).await, } } @@ -862,6 +829,7 @@ mod tests { use std::collections::VecDeque; use std::sync::Mutex; + use fabro_acp::test_support::fake_acp_agent_script; use fabro_agent::sandbox::{DirEntry, GrepOptions}; /// Mock sandbox that returns pre-configured ExecResults in FIFO order. @@ -1286,16 +1254,18 @@ mod tests { let context = Context::new(); let router = test_router(); + let emitter = Arc::new(Emitter::default()); + let stage_scope = StageScope::for_handler(&context, "test"); let result = router - .one_shot( - &node, - "prompt", - None, - &Arc::new(Emitter::default()), - &StageScope::for_handler(&context, "test"), - &sandbox, - CancellationToken::new(), - ) + .one_shot(OneShotRequest { + node: &node, + prompt: "prompt", + system_prompt: None, + emitter: &emitter, + stage_scope: &stage_scope, + sandbox: &sandbox, + cancel_token: CancellationToken::new(), + }) .await .unwrap(); @@ -1313,17 +1283,19 @@ mod tests { )); let context = Context::new(); let router = test_router(); + let emitter = Arc::new(Emitter::default()); + let stage_scope = StageScope::for_handler(&context, "test"); let result = router - .one_shot( - &node, - "prompt", - None, - &Arc::new(Emitter::default()), - &StageScope::for_handler(&context, "test"), - &sandbox, - CancellationToken::new(), - ) + .one_shot(OneShotRequest { + node: &node, + prompt: "prompt", + system_prompt: None, + emitter: &emitter, + stage_scope: &stage_scope, + sandbox: &sandbox, + cancel_token: CancellationToken::new(), + }) .await .unwrap(); @@ -1343,17 +1315,19 @@ mod tests { )); let context = Context::new(); let router = test_router(); + let emitter = Arc::new(Emitter::default()); + let stage_scope = StageScope::for_handler(&context, "test"); let result = router - .one_shot( - &node, - "prompt", - None, - &Arc::new(Emitter::default()), - &StageScope::for_handler(&context, "test"), - &sandbox, - CancellationToken::new(), - ) + .one_shot(OneShotRequest { + node: &node, + prompt: "prompt", + system_prompt: None, + emitter: &emitter, + stage_scope: &stage_scope, + sandbox: &sandbox, + cancel_token: CancellationToken::new(), + }) .await .unwrap(); @@ -1369,43 +1343,6 @@ mod tests { BackendRouter::new(Box::new(StubBackend), cli_backend, acp_backend) } - fn fake_acp_agent_script() -> &'static str { - r#" -import json -import sys - -session_id = "sess-1" - -def send(message): - print(json.dumps(message), flush=True) - -def respond(message, result): - send({"jsonrpc": "2.0", "id": message["id"], "result": result}) - -for line in sys.stdin: - message = json.loads(line) - method = message.get("method") - if method == "initialize": - respond(message, {"protocolVersion": 1, "agentCapabilities": {}}) - elif method == "session/new": - respond(message, {"sessionId": session_id}) - elif method == "session/prompt": - send({ - "jsonrpc": "2.0", - "method": "session/update", - "params": { - "sessionId": session_id, - "update": { - "sessionUpdate": "agent_message_chunk", - "content": {"type": "text", "text": "hello from acp"} - } - } - }) - respond(message, {"stopReason": "end_turn"}) - break -"# - } - /// Minimal stub backend for testing routing logic. struct StubBackend; @@ -1430,16 +1367,7 @@ for line in sys.stdin: }) } - async fn one_shot( - &self, - _node: &Node, - _prompt: &str, - _system_prompt: Option<&str>, - _emitter: &Arc, - _stage_scope: &StageScope, - _sandbox: &Arc, - _cancel_token: CancellationToken, - ) -> Result { + async fn one_shot(&self, _request: OneShotRequest<'_>) -> Result { Ok(CodergenResult::Text { text: "api one-shot".to_string(), usage: None, diff --git a/lib/crates/fabro-workflow/src/handler/prompt.rs b/lib/crates/fabro-workflow/src/handler/prompt.rs index 5c59041fc..7d440087b 100644 --- a/lib/crates/fabro-workflow/src/handler/prompt.rs +++ b/lib/crates/fabro-workflow/src/handler/prompt.rs @@ -6,7 +6,8 @@ use fabro_graphviz::graph::{Graph, Node}; use fabro_model::Provider; use super::agent::{ - CodergenBackend, CodergenResult, expand_variables, extract_status_fields, truncate, + CodergenBackend, CodergenResult, OneShotRequest, expand_variables, extract_status_fields, + truncate, }; use super::{EngineServices, Handler}; use crate::context::{Context, WorkflowContext, keys}; @@ -120,15 +121,15 @@ impl Handler for PromptHandler { let (response_text, stage_usage, backend_files_touched) = if let Some(backend) = &self.backend { let result = backend - .one_shot( + .one_shot(OneShotRequest { node, - &prompt, - system_prompt.as_deref(), - &services.run.emitter, - &stage_scope, - &services.run.sandbox, - services.run.cancel_token(), - ) + prompt: &prompt, + system_prompt: system_prompt.as_deref(), + emitter: &services.run.emitter, + stage_scope: &stage_scope, + sandbox: &services.run.sandbox, + cancel_token: services.run.cancel_token(), + }) .await; match result { Ok(CodergenResult::Full(outcome)) => return Ok(outcome), @@ -332,13 +333,7 @@ mod tests { async fn one_shot( &self, - _node: &Node, - _prompt: &str, - _system_prompt: Option<&str>, - _emitter: &Arc, - _stage_scope: &StageScope, - _sandbox: &Arc, - _cancel_token: CancellationToken, + _request: OneShotRequest<'_>, ) -> Result { Ok(CodergenResult::Text { text: "one-shot response".to_string(), @@ -397,13 +392,7 @@ mod tests { async fn one_shot( &self, - _node: &Node, - _prompt: &str, - _system_prompt: Option<&str>, - _emitter: &Arc, - _stage_scope: &StageScope, - _sandbox: &Arc, - _cancel_token: CancellationToken, + _request: OneShotRequest<'_>, ) -> Result { Ok(CodergenResult::Text { text: "one-shot response".to_string(), @@ -457,18 +446,10 @@ mod tests { panic!("run() should not be called for prompt handler"); } - async fn one_shot( - &self, - _node: &Node, - prompt: &str, - system_prompt: Option<&str>, - _emitter: &Arc, - _stage_scope: &StageScope, - _sandbox: &Arc, - _cancel_token: CancellationToken, - ) -> Result { - *self.captured_prompt.lock().unwrap() = Some(prompt.to_string()); - *self.captured_system_prompt.lock().unwrap() = Some(system_prompt.map(String::from)); + async fn one_shot(&self, request: OneShotRequest<'_>) -> Result { + *self.captured_prompt.lock().unwrap() = Some(request.prompt.to_string()); + *self.captured_system_prompt.lock().unwrap() = + Some(request.system_prompt.map(String::from)); Ok(CodergenResult::Text { text: "classified".to_string(), usage: None, diff --git a/lib/crates/fabro-workflow/src/pipeline/initialize.rs b/lib/crates/fabro-workflow/src/pipeline/initialize.rs index f77245c43..6d472bbdf 100644 --- a/lib/crates/fabro-workflow/src/pipeline/initialize.rs +++ b/lib/crates/fabro-workflow/src/pipeline/initialize.rs @@ -709,6 +709,7 @@ mod tests { use std::sync::Arc; use std::time::Duration; + use fabro_acp::test_support::fake_acp_agent_script; use fabro_auth::{AuthCredential, AuthDetails}; use fabro_graphviz::graph::{AttrValue, Edge, Graph, Node}; use fabro_interview::AutoApproveInterviewer; @@ -1192,7 +1193,7 @@ mod tests { assert_eq!( outcome.context_updates.get(&keys::response_key("writer")), - Some(&serde_json::json!("hello from initialized acp")) + Some(&serde_json::json!("hello from acp")) ); assert!( seen.lock() @@ -1278,43 +1279,6 @@ mod tests { ); } - fn fake_acp_agent_script() -> &'static str { - r#" -import json -import sys - -session_id = "sess-1" - -def send(message): - print(json.dumps(message), flush=True) - -def respond(message, result): - send({"jsonrpc": "2.0", "id": message["id"], "result": result}) - -for line in sys.stdin: - message = json.loads(line) - method = message.get("method") - if method == "initialize": - respond(message, {"protocolVersion": 1, "agentCapabilities": {}}) - elif method == "session/new": - respond(message, {"sessionId": session_id}) - elif method == "session/prompt": - send({ - "jsonrpc": "2.0", - "method": "session/update", - "params": { - "sessionId": session_id, - "update": { - "sessionUpdate": "agent_message_chunk", - "content": {"type": "text", "text": "hello from initialized acp"} - } - } - }) - respond(message, {"stopReason": "end_turn"}) - break -"# - } - #[tokio::test] async fn initialize_setup_failure_preserves_stderr_and_adds_exec_tail() { let (result, events) = diff --git a/lib/crates/fabro-workflow/tests/it/integration.rs b/lib/crates/fabro-workflow/tests/it/integration.rs index 7a451e09a..2d50a638c 100644 --- a/lib/crates/fabro-workflow/tests/it/integration.rs +++ b/lib/crates/fabro-workflow/tests/it/integration.rs @@ -23,6 +23,7 @@ use std::path::{Path, PathBuf}; use std::sync::Arc; use std::time::Duration; +use fabro_acp::test_support::fake_acp_agent_script; use fabro_config::RunScratch; use fabro_graphviz::graph::{AttrValue, Edge, Graph, Node}; use fabro_graphviz::parser::parse; @@ -6284,7 +6285,9 @@ mod real_llm { use fabro_types::WorkflowSettings; use fabro_workflow::context::Context; use fabro_workflow::error::Error; - use fabro_workflow::handler::agent::{AgentHandler, CodergenBackend, CodergenResult}; + use fabro_workflow::handler::agent::{ + AgentHandler, CodergenBackend, CodergenResult, OneShotRequest, + }; use tokio_util::sync::CancellationToken; struct LlmCodergenBackend { @@ -6309,17 +6312,8 @@ mod real_llm { self.complete(prompt).await } - async fn one_shot( - &self, - _node: &Node, - prompt: &str, - _system_prompt: Option<&str>, - _emitter: &Arc, - _stage_scope: &fabro_workflow::event::StageScope, - _sandbox: &Arc, - _cancel_token: tokio_util::sync::CancellationToken, - ) -> Result { - self.complete(prompt).await + async fn one_shot(&self, request: OneShotRequest<'_>) -> Result { + self.complete(request.prompt).await } } @@ -10420,7 +10414,7 @@ async fn backend_router_delegates_to_cli_for_backend_attr() { async fn backend_router_delegates_to_acp_for_acp_node() { let tempdir = tempfile::tempdir().unwrap(); let script_path = tempdir.path().join("fake_acp_agent.py"); - tokio::fs::write(&script_path, fake_router_acp_agent_script()) + tokio::fs::write(&script_path, fake_acp_agent_script()) .await .unwrap(); let env: Arc = @@ -10471,7 +10465,7 @@ async fn backend_router_delegates_to_acp_for_acp_node() { match result { CodergenResult::Text { text, .. } => { assert_eq!( - text, "ACP response", + text, "hello from acp", "should route to ACP backend for backend=acp" ); } @@ -10479,43 +10473,6 @@ async fn backend_router_delegates_to_acp_for_acp_node() { } } -fn fake_router_acp_agent_script() -> &'static str { - r#" -import json -import sys - -session_id = "sess-1" - -def send(message): - print(json.dumps(message), flush=True) - -def respond(message, result): - send({"jsonrpc": "2.0", "id": message["id"], "result": result}) - -for line in sys.stdin: - message = json.loads(line) - method = message.get("method") - if method == "initialize": - respond(message, {"protocolVersion": 1, "agentCapabilities": {}}) - elif method == "session/new": - respond(message, {"sessionId": session_id}) - elif method == "session/prompt": - send({ - "jsonrpc": "2.0", - "method": "session/update", - "params": { - "sessionId": session_id, - "update": { - "sessionUpdate": "agent_message_chunk", - "content": {"type": "text", "text": "ACP response"} - } - } - }) - respond(message, {"stopReason": "end_turn"}) - break -"# -} - // -- Full pipeline e2e with BackendRouter -- #[tokio::test]