fabro(01KQT1V2W1R6ZH72CFT2QDJ39Q): implement (succeeded)

Fabro-Run: 01KQT1V2W1R6ZH72CFT2QDJ39Q
Fabro-Completed: 5
Fabro-Checkpoint: 89ea5f732a

⚒️ Generated with [Fabro](https://fabro.sh)
This commit is contained in:
Fabro 2026-05-04 19:13:15 +00:00
parent 5edca0b698
commit 1c3eea1053
55 changed files with 1965 additions and 717 deletions

View file

@ -42,7 +42,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
let config = SessionOptions::default();
let mut session = Session::new(client, profile, sandbox, config);
session.initialize().await;
session.initialize().await?;
// Subscribe to events before sending input
let mut events = session.subscribe();
@ -936,4 +936,4 @@ Register it on the client:
```rust
client.register_provider(Arc::new(MyProvider)).await?;
```
```

View file

@ -140,7 +140,7 @@ let config = SessionConfig {
// 5. Create and initialize the session
let mut session = Session::new(client, profile, env, config, None);
session.initialize().await;
session.initialize().await?;
// 6. Subscribe to events (for UI rendering)
let mut rx = session.subscribe();
@ -232,4 +232,4 @@ profile.register_subagent_tools(manager, factory, 0);
- **Tool output truncation** -- Per-tool character and line limits with head/tail or tail-only truncation modes
- **Environment variable filtering** -- `LocalSandbox` strips secrets (`*_API_KEY`, `*_SECRET`, `*_TOKEN`, `*_PASSWORD`, `*_CREDENTIAL`) from subprocess environments
- **Command timeouts** -- Configurable per-command with process group cleanup (SIGTERM then SIGKILL)
- **Project doc discovery** -- Automatically discovers `AGENTS.md`, `CLAUDE.md`, `GEMINI.md`, or `.codex/instructions.md` based on provider, with a 32KB budget
- **Project doc discovery** -- Automatically discovers `AGENTS.md`, `CLAUDE.md`, `GEMINI.md`, or `.codex/instructions.md` based on provider, with a 32KB budget

View file

@ -724,7 +724,7 @@ pub async fn run_with_args_and_client(
});
// Initialize and run
session.initialize().await;
session.initialize().await?;
let result = session.process_input(&args.prompt).await;
if matches!(output_format, OutputFormat::Text) {

View file

@ -1,8 +1,10 @@
use std::collections::HashSet;
use fabro_model::Provider;
use tokio_util::sync::CancellationToken;
use tracing::{debug, info, warn};
use crate::error::{Error, InterruptReason};
use crate::sandbox::Sandbox;
const BUDGET_BYTES: usize = 32768;
@ -12,7 +14,8 @@ pub async fn discover_memory(
git_root: &str,
working_dir: &str,
provider: Provider,
) -> Vec<String> {
cancel_token: &CancellationToken,
) -> Result<Vec<String>, Error> {
let directories = build_directory_walk(git_root, working_dir);
let candidate_filenames: Vec<&str> = match provider {
@ -34,8 +37,15 @@ pub async fn discover_memory(
for dir in &directories {
for filename in &candidate_filenames {
if cancel_token.is_cancelled() {
return Err(Error::Interrupted(InterruptReason::Cancelled));
}
let path = format!("{dir}/{filename}");
if let Ok(content) = env.read_file(&path, None, None).await {
let read_result = env.read_file(&path, None, None).await;
if cancel_token.is_cancelled() {
return Err(Error::Interrupted(InterruptReason::Cancelled));
}
if let Ok(content) = read_result {
if content.is_empty() {
warn!(path = %path, "Project doc file empty, skipping");
continue;
@ -68,7 +78,7 @@ pub async fn discover_memory(
let total_bytes: usize = results.iter().map(std::string::String::len).sum();
info!(files = results.len(), total_bytes, "Project docs loaded");
results
Ok(results)
}
fn build_directory_walk(git_root: &str, working_dir: &str) -> Vec<String> {
@ -117,6 +127,8 @@ mod tests {
use std::collections::HashMap;
use std::sync::Arc;
use tokio_util::sync::CancellationToken;
use super::*;
use crate::sandbox::Sandbox;
use crate::test_support::MockSandbox;
@ -129,7 +141,15 @@ mod tests {
files,
..Default::default()
});
let docs = discover_memory(env.as_ref(), "/repo", "/repo", Provider::Anthropic).await;
let docs = discover_memory(
env.as_ref(),
"/repo",
"/repo",
Provider::Anthropic,
&CancellationToken::new(),
)
.await
.unwrap();
assert_eq!(docs.len(), 1);
assert_eq!(docs[0], "Agent instructions");
}
@ -146,8 +166,15 @@ mod tests {
files: files.clone(),
..Default::default()
});
let anthropic_docs =
discover_memory(env.as_ref(), "/repo", "/repo", Provider::Anthropic).await;
let anthropic_docs = discover_memory(
env.as_ref(),
"/repo",
"/repo",
Provider::Anthropic,
&CancellationToken::new(),
)
.await
.unwrap();
assert_eq!(anthropic_docs.len(), 2);
assert_eq!(anthropic_docs[0], "agents");
assert_eq!(anthropic_docs[1], "claude");
@ -156,7 +183,15 @@ mod tests {
files: files.clone(),
..Default::default()
});
let openai_docs = discover_memory(env.as_ref(), "/repo", "/repo", Provider::OpenAi).await;
let openai_docs = discover_memory(
env.as_ref(),
"/repo",
"/repo",
Provider::OpenAi,
&CancellationToken::new(),
)
.await
.unwrap();
assert_eq!(openai_docs.len(), 2);
assert_eq!(openai_docs[0], "agents");
assert_eq!(openai_docs[1], "copilot");
@ -165,7 +200,15 @@ mod tests {
files,
..Default::default()
});
let gemini_docs = discover_memory(env.as_ref(), "/repo", "/repo", Provider::Gemini).await;
let gemini_docs = discover_memory(
env.as_ref(),
"/repo",
"/repo",
Provider::Gemini,
&CancellationToken::new(),
)
.await
.unwrap();
assert_eq!(gemini_docs.len(), 2);
assert_eq!(gemini_docs[0], "agents");
assert_eq!(gemini_docs[1], "gemini");
@ -184,7 +227,15 @@ mod tests {
files,
..Default::default()
});
let docs = discover_memory(env.as_ref(), "/repo", "/repo", Provider::Anthropic).await;
let docs = discover_memory(
env.as_ref(),
"/repo",
"/repo",
Provider::Anthropic,
&CancellationToken::new(),
)
.await
.unwrap();
assert_eq!(docs.len(), 2);
assert_eq!(docs[0], large_content);
// Second doc should be truncated to fit remaining budget
@ -201,7 +252,15 @@ mod tests {
files,
..Default::default()
});
let docs = discover_memory(env.as_ref(), "/repo", "/repo", Provider::Anthropic).await;
let docs = discover_memory(
env.as_ref(),
"/repo",
"/repo",
Provider::Anthropic,
&CancellationToken::new(),
)
.await
.unwrap();
assert_eq!(docs.len(), 1);
assert_eq!(docs[0], "shared instructions");
}
@ -215,7 +274,15 @@ mod tests {
files,
..Default::default()
});
let docs = discover_memory(env.as_ref(), "/repo", "/repo/src", Provider::Anthropic).await;
let docs = discover_memory(
env.as_ref(),
"/repo",
"/repo/src",
Provider::Anthropic,
&CancellationToken::new(),
)
.await
.unwrap();
assert_eq!(docs.len(), 1);
assert_eq!(docs[0], "shared instructions");
}
@ -231,8 +298,15 @@ mod tests {
files,
..Default::default()
});
let docs =
discover_memory(env.as_ref(), "/repo", "/repo/src/app", Provider::Anthropic).await;
let docs = discover_memory(
env.as_ref(),
"/repo",
"/repo/src/app",
Provider::Anthropic,
&CancellationToken::new(),
)
.await
.unwrap();
assert_eq!(docs.len(), 3);
assert_eq!(docs[0], "root agents");
assert_eq!(docs[1], "src agents");

View file

@ -130,13 +130,24 @@ impl Session {
/// Initialize session by discovering project docs and capturing environment
/// context. Call before `process_input`.
pub async fn initialize(&mut self) {
///
/// # Errors
///
/// Returns `Error::Interrupted(InterruptReason::Cancelled)` if the
/// session's cancel token fires during initialization.
pub async fn initialize(&mut self) -> Result<(), Error> {
let cancel_token = self.cancel_token.clone();
self.event_emitter
.emit(self.id.clone(), AgentEvent::SessionStarted {
provider: Some(self.provider_profile.provider().to_string()),
model: Some(self.provider_profile.model().to_string()),
});
if cancel_token.is_cancelled() {
return Err(Error::Interrupted(InterruptReason::Cancelled));
}
let doc_root = self
.config
.git_root
@ -147,8 +158,9 @@ impl Session {
&doc_root,
self.sandbox.working_directory(),
self.provider_profile.provider(),
&cancel_token,
)
.await;
.await?;
// Discover skills
let skill_dirs = if let Some(dirs) = &self.config.skill_dirs {
@ -158,7 +170,7 @@ impl Session {
let skills_str = skills_dir.to_string_lossy().to_string();
default_skill_dirs(Some(&skills_str), self.config.git_root.as_deref())
};
self.skills = discover_skills(self.sandbox.as_ref(), &skill_dirs).await;
self.skills = discover_skills(self.sandbox.as_ref(), &skill_dirs, &cancel_token).await?;
debug!(skill_count = self.skills.len(), "Skills discovered");
// Register use_skill tool when skills are available
@ -175,7 +187,7 @@ impl Session {
if !self.config.mcp_servers.is_empty() {
// Resolve Sandbox transports: start the server inside the sandbox,
// then rewrite the config to Http using the sandbox's preview URL.
let mcp_servers = self.resolve_sandbox_mcp_servers().await;
let mcp_servers = self.resolve_sandbox_mcp_servers(&cancel_token).await?;
let mut manager = McpConnectionManager::new();
let results = manager.start_servers(&mcp_servers).await;
@ -209,7 +221,7 @@ impl Session {
}
// Populate environment context
self.env_context = self.build_env_context().await;
self.env_context = self.build_env_context(&cancel_token).await?;
debug!(
is_git_repo = self.env_context.is_git_repo,
model = %self.env_context.model,
@ -224,19 +236,30 @@ impl Session {
self.config.user_instructions.as_deref(),
&self.skills,
);
Ok(())
}
/// Resolve `McpTransport::Sandbox` configs by starting the MCP server
/// inside the sandbox and rewriting the transport to `Http` with the
/// sandbox's preview URL.
async fn resolve_sandbox_mcp_servers(&self) -> Vec<McpServerSettings> {
async fn resolve_sandbox_mcp_servers(
&self,
cancel_token: &CancellationToken,
) -> Result<Vec<McpServerSettings>, Error> {
let mut resolved = Vec::with_capacity(self.config.mcp_servers.len());
for config in &self.config.mcp_servers {
if cancel_token.is_cancelled() {
return Err(Error::Interrupted(InterruptReason::Cancelled));
}
match &config.transport {
McpTransport::Sandbox { command, port, env } => {
let port = *port;
match self.start_sandbox_mcp_server(command, port, env).await {
match self
.start_sandbox_mcp_server(command, port, env, cancel_token)
.await?
{
Ok((url, headers)) => {
info!(
server = %config.name,
@ -268,17 +291,24 @@ impl Session {
}
}
resolved
Ok(resolved)
}
/// Start an MCP server inside the sandbox and return (url, headers) for
/// HTTP connection.
///
/// The outer `Result` surfaces fatal cancellation as
/// `Error::Interrupted(InterruptReason::Cancelled)` (the running MCP
/// process group is terminated before returning). The inner `Result`
/// captures non-fatal startup failures that the caller logs and turns
/// into an `McpServerFailed` event.
async fn start_sandbox_mcp_server(
&self,
command: &[String],
port: u16,
env: &std::collections::HashMap<String, String>,
) -> Result<(String, std::collections::HashMap<String, String>), String> {
cancel_token: &CancellationToken,
) -> Result<Result<(String, std::collections::HashMap<String, String>), String>, Error> {
let sandbox = self.sandbox.as_ref();
let cmd_str = command
@ -296,27 +326,63 @@ impl Session {
quoted = fabro_sandbox::shell_quote(&inner)
);
let env_ref = if env.is_empty() { None } else { Some(env) };
let launch_result = sandbox
.exec_command(&launch_script, 30_000, None, env_ref, None)
.await
.map_err(|e| format!("Failed to launch MCP server: {}", e.display_with_causes()))?;
let pid = launch_result.stdout.trim();
info!(pid, port, "MCP server process launched in sandbox");
if cancel_token.is_cancelled() {
return Err(Error::Interrupted(InterruptReason::Cancelled));
}
let launch_result = match sandbox
.exec_command(
&launch_script,
30_000,
None,
env_ref,
Some(cancel_token.child_token()),
)
.await
{
Ok(result) => result,
Err(e) => {
if cancel_token.is_cancelled() {
return Err(Error::Interrupted(InterruptReason::Cancelled));
}
return Ok(Err(format!(
"Failed to launch MCP server: {}",
e.display_with_causes()
)));
}
};
let pid = launch_result.stdout.trim().to_string();
info!(pid = %pid, port, "MCP server process launched in sandbox");
// Wait for the server to start listening on the port
let poll_cmd = format!(
"for i in $(seq 1 30); do ss -tln | grep -q ':{port} ' && echo ready && exit 0; sleep 1; done; echo timeout"
);
let poll_result = sandbox
.exec_command(&poll_cmd, 60_000, None, None, None)
.await
.map_err(|e| {
format!(
.exec_command(
&poll_cmd,
60_000,
None,
None,
Some(cancel_token.child_token()),
)
.await;
if cancel_token.is_cancelled() {
kill_mcp_pid(sandbox, &pid).await;
return Err(Error::Interrupted(InterruptReason::Cancelled));
}
let poll_result = match poll_result {
Ok(result) => result,
Err(e) => {
return Ok(Err(format!(
"Failed to poll MCP server readiness: {}",
e.display_with_causes()
)
})?;
)));
}
};
if poll_result.stdout.trim() != "ready" {
// Grab stderr for debugging
@ -326,51 +392,80 @@ impl Session {
10_000,
None,
None,
None,
Some(cancel_token.child_token()),
)
.await
.map(|r| r.stdout)
.unwrap_or_default();
return Err(format!(
return Ok(Err(format!(
"MCP server did not start listening on port {port} within 30s. stderr:\n{stderr}"
));
)));
}
// Get the preview URL for the port, or fall back to localhost for local
// sandboxes
if let Some(url_and_headers) = sandbox
.get_preview_url(port)
.await
.map_err(|e| e.display_with_causes())?
{
Ok(url_and_headers)
let preview = match sandbox.get_preview_url(port).await {
Ok(p) => p,
Err(e) => return Ok(Err(e.display_with_causes())),
};
if cancel_token.is_cancelled() {
kill_mcp_pid(sandbox, &pid).await;
return Err(Error::Interrupted(InterruptReason::Cancelled));
}
if let Some(url_and_headers) = preview {
Ok(Ok(url_and_headers))
} else {
info!(port, "No preview URL available, using localhost");
Ok((
Ok(Ok((
format!("http://localhost:{port}"),
std::collections::HashMap::new(),
))
)))
}
}
async fn build_env_context(&self) -> EnvContext {
async fn build_env_context(
&self,
cancel_token: &CancellationToken,
) -> Result<EnvContext, Error> {
let today = chrono::Local::now().format("%Y-%m-%d").to_string();
let model_name = self.provider_profile.model().to_string();
if cancel_token.is_cancelled() {
return Err(Error::Interrupted(InterruptReason::Cancelled));
}
// Detect git info via sandbox
let git_branch = self
.sandbox
.exec_command("git rev-parse --abbrev-ref HEAD", 5000, None, None, None)
.exec_command(
"git rev-parse --abbrev-ref HEAD",
5000,
None,
None,
Some(cancel_token.child_token()),
)
.await
.ok()
.filter(fabro_sandbox::ExecResult::is_success)
.map(|r| r.stdout.trim().to_string());
if cancel_token.is_cancelled() {
return Err(Error::Interrupted(InterruptReason::Cancelled));
}
let is_git_repo = git_branch.is_some();
let git_status_short = if is_git_repo {
self.sandbox
.exec_command("git status --short", 5000, None, None, None)
.exec_command(
"git status --short",
5000,
None,
None,
Some(cancel_token.child_token()),
)
.await
.ok()
.filter(fabro_sandbox::ExecResult::is_success)
@ -380,9 +475,19 @@ impl Session {
None
};
if cancel_token.is_cancelled() {
return Err(Error::Interrupted(InterruptReason::Cancelled));
}
let git_recent_commits = if is_git_repo {
self.sandbox
.exec_command("git log --oneline -10", 5000, None, None, None)
.exec_command(
"git log --oneline -10",
5000,
None,
None,
Some(cancel_token.child_token()),
)
.await
.ok()
.filter(fabro_sandbox::ExecResult::is_success)
@ -392,7 +497,11 @@ impl Session {
None
};
EnvContext {
if cancel_token.is_cancelled() {
return Err(Error::Interrupted(InterruptReason::Cancelled));
}
Ok(EnvContext {
git_branch,
is_git_repo,
current_date: today,
@ -400,7 +509,7 @@ impl Session {
knowledge_cutoff: self.provider_profile.knowledge_cutoff().unwrap_or_default(),
git_status_short,
git_recent_commits,
}
})
}
#[must_use]
@ -1031,6 +1140,23 @@ const fn is_auth_error(err: &LlmError) -> bool {
)
}
/// Best-effort kill of a sandbox MCP server process group. Used when
/// `start_sandbox_mcp_server` is cancelled after spawning a detached
/// `setsid` child but before reporting readiness. Errors from the sandbox
/// are logged and swallowed; the caller is already returning a Cancelled
/// error.
async fn kill_mcp_pid(sandbox: &dyn Sandbox, pid: &str) {
let pid = pid.trim();
if pid.is_empty() {
return;
}
let script =
format!("kill -TERM -{pid} 2>/dev/null; sleep 1; kill -KILL -{pid} 2>/dev/null; true");
if let Err(err) = sandbox.exec_command(&script, 5_000, None, None, None).await {
warn!(pid, error = %err.display_with_causes(), "Failed to kill MCP server process group during cancellation");
}
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
@ -1297,7 +1423,7 @@ mod tests {
let mut session = make_session(vec![text_response("Hello")]).await;
let mut rx = session.subscribe();
session.initialize().await;
session.initialize().await.unwrap();
session.process_input("Hi").await.unwrap();
session.close();
@ -1825,7 +1951,7 @@ mod tests {
let mut session = make_session(responses).await;
let mut rx = session.subscribe();
session.initialize().await;
session.initialize().await.unwrap();
session.process_input("one").await.unwrap();
session.process_input("two").await.unwrap();
session.close();
@ -1858,7 +1984,7 @@ mod tests {
..Default::default()
};
let mut session = Session::new(client, profile, env, config, None);
session.initialize().await;
session.initialize().await.unwrap();
session.process_input("test").await.unwrap();
// Verify user instructions are included in the system prompt
@ -2647,7 +2773,7 @@ mod tests {
let mut rx = session.subscribe();
// Initialize starts the MCP server and registers tools
session.initialize().await;
session.initialize().await.unwrap();
// Verify McpServerReady event was emitted
let mut mcp_ready = false;
@ -2842,7 +2968,7 @@ mod tests {
#[tokio::test]
async fn process_input_emits_processing_end_on_idle_transition() {
let mut session = make_session(vec![text_response("Hello")]).await;
session.initialize().await;
session.initialize().await.unwrap();
let mut rx = session.subscribe();
session.process_input("Hi").await.unwrap();

View file

@ -1,7 +1,9 @@
use std::sync::Arc;
use fabro_llm::types::ToolDefinition;
use tokio_util::sync::CancellationToken;
use crate::error::{Error, InterruptReason};
use crate::sandbox::Sandbox;
use crate::tool_registry::RegisteredTool;
use crate::tools::required_str;
@ -224,17 +226,35 @@ pub fn default_skill_dirs(fabro_skills_dir: Option<&str>, git_root: Option<&str>
dirs
}
pub async fn discover_skills(env: &dyn Sandbox, dirs: &[String]) -> Vec<Skill> {
pub async fn discover_skills(
env: &dyn Sandbox,
dirs: &[String],
cancel_token: &CancellationToken,
) -> Result<Vec<Skill>, Error> {
let mut skills_by_name: std::collections::HashMap<String, Skill> =
std::collections::HashMap::new();
for dir in dirs {
let Ok(paths) = env.glob("*/SKILL.md", Some(dir)).await else {
if cancel_token.is_cancelled() {
return Err(Error::Interrupted(InterruptReason::Cancelled));
}
let glob_result = env.glob("*/SKILL.md", Some(dir)).await;
if cancel_token.is_cancelled() {
return Err(Error::Interrupted(InterruptReason::Cancelled));
}
let Ok(paths) = glob_result else {
continue;
};
for path in paths {
let Ok(content) = env.read_file(&path, None, None).await else {
if cancel_token.is_cancelled() {
return Err(Error::Interrupted(InterruptReason::Cancelled));
}
let read_result = env.read_file(&path, None, None).await;
if cancel_token.is_cancelled() {
return Err(Error::Interrupted(InterruptReason::Cancelled));
}
let Ok(content) = read_result else {
continue;
};
@ -246,7 +266,7 @@ pub async fn discover_skills(env: &dyn Sandbox, dirs: &[String]) -> Vec<Skill> {
let mut skills: Vec<Skill> = skills_by_name.into_values().collect();
skills.sort_by(|a, b| a.name.cmp(&b.name));
skills
Ok(skills)
}
#[cfg(test)]
@ -454,7 +474,9 @@ name: trimmed
..Default::default()
};
let skills = discover_skills(&env, &["/skills".into()]).await;
let skills = discover_skills(&env, &["/skills".into()], &CancellationToken::new())
.await
.unwrap();
assert_eq!(skills.len(), 1);
assert_eq!(skills[0].name, "commit");
assert_eq!(skills[0].description, "Make a commit");
@ -477,7 +499,9 @@ name: trimmed
..Default::default()
};
let skills = discover_skills(&env, &["/skills".into()]).await;
let skills = discover_skills(&env, &["/skills".into()], &CancellationToken::new())
.await
.unwrap();
assert_eq!(skills.len(), 1);
assert_eq!(skills[0].name, "good");
}
@ -485,7 +509,9 @@ name: trimmed
#[tokio::test]
async fn discover_empty_dirs() {
let env = MockSandbox::default();
let skills = discover_skills(&env, &[]).await;
let skills = discover_skills(&env, &[], &CancellationToken::new())
.await
.unwrap();
assert!(skills.is_empty());
}
@ -514,7 +540,13 @@ name: trimmed
};
// discover_skills iterates dirs in order; later dirs override earlier names
let skills = discover_skills(&env, &["/global".into(), "/project".into()]).await;
let skills = discover_skills(
&env,
&["/global".into(), "/project".into()],
&CancellationToken::new(),
)
.await
.unwrap();
assert_eq!(skills.len(), 1);
assert_eq!(skills[0].description, "Project commit");
}

View file

@ -111,7 +111,7 @@ impl SubAgentManager {
let task_prompt_for_spawn = task_prompt.clone();
let task = tokio::spawn(async move {
session.initialize().await;
session.initialize().await?;
session.process_input(&task_prompt_for_spawn).await?;
let turns = session.history().turns();
let last_text = turns.iter().rev().find_map(|t| match t {

View file

@ -1466,7 +1466,7 @@ def farewell(name):
SessionOptions::default(),
None,
);
session.initialize().await;
session.initialize().await.unwrap();
session
.process_input("Update the greeting functions")
.await

View file

@ -170,7 +170,7 @@ macro_rules! provider_test {
async fn [<$prefix _ $scenario>]() {
let tmp = tempfile::tempdir().expect("failed to create tempdir");
let mut session = make_session($provider, $model, tmp.path(), None).await;
session.initialize().await;
session.initialize().await.unwrap();
[<scenario_ $scenario>](&mut session, tmp.path()).await;
}
}
@ -195,7 +195,7 @@ macro_rules! openai_twin_provider_test {
tmp.path(),
Some(twin),
).await;
session.initialize().await;
session.initialize().await.unwrap();
[<scenario_ $scenario>](&mut session, tmp.path()).await;
}
}
@ -670,7 +670,7 @@ macro_rules! reasoning_effort_tests {
};
let mut session =
make_session_with_config($provider, $model, tmp.path(), config, None).await;
session.initialize().await;
session.initialize().await.unwrap();
session
.process_input("Say hello")
.await
@ -749,7 +749,7 @@ macro_rules! loop_detection_tests {
};
let mut session =
make_session_with_config($provider, $model, tmp.path(), config, None).await;
session.initialize().await;
session.initialize().await.unwrap();
session
.process_input("Repeatedly read the file /dev/null")
.await

View file

@ -8,7 +8,6 @@ use std::collections::HashMap;
use std::io::{BufRead as StdBufRead, BufReader as StdBufReader};
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Duration;
use anyhow::{Context, Result, anyhow};
@ -34,6 +33,7 @@ use fabro_workflow::runtime_store::{RunStoreBackend, RunStoreHandle};
use tokio::signal::unix::{SignalKind, signal};
use tokio::sync::{Mutex, RwLock as AsyncRwLock, mpsc};
use tokio::time::sleep;
use tokio_util::sync::CancellationToken;
use crate::args::RunWorkerMode;
use crate::server_client;
@ -86,10 +86,10 @@ pub(crate) async fn execute(
worker_token.to_owned(),
)));
let interviewer = Arc::new(ControlInterviewer::new());
let cancel_token = Arc::new(AtomicBool::new(false));
spawn_worker_control_stream(Arc::clone(&interviewer), Arc::clone(&cancel_token))?;
let cancel_token = CancellationToken::new();
spawn_worker_control_stream(Arc::clone(&interviewer), cancel_token.clone())?;
let run_control = RunControlState::new();
install_signal_handlers(Arc::clone(&run_control), Arc::clone(&cancel_token))?;
install_signal_handlers(Arc::clone(&run_control), cancel_token.clone())?;
let vault = load_worker_vault(storage_dir.as_deref())?;
let github_app = {
let vault_guard = match &vault {
@ -100,7 +100,7 @@ pub(crate) async fn execute(
};
let services = StartServices {
run_id,
cancel_token: Some(Arc::clone(&cancel_token)),
cancel_token: cancel_token.clone(),
emitter: Arc::new(Emitter::new(run_id)),
interviewer,
run_store: run_store.clone(),
@ -162,7 +162,7 @@ enum WorkerControlStreamEvent {
)]
fn spawn_worker_control_stream(
interviewer: Arc<ControlInterviewer>,
cancel_token: Arc<AtomicBool>,
cancel_token: CancellationToken,
) -> Result<()> {
let (event_tx, event_rx) = mpsc::unbounded_channel();
tokio::spawn(handle_worker_control_stream_events(
@ -205,7 +205,7 @@ fn read_worker_control_stream_blocking<R>(
async fn handle_worker_control_stream_events(
interviewer: Arc<ControlInterviewer>,
cancel_token: Arc<AtomicBool>,
cancel_token: CancellationToken,
mut event_rx: mpsc::UnboundedReceiver<WorkerControlStreamEvent>,
) {
while let Some(event) = event_rx.recv().await {
@ -225,7 +225,7 @@ async fn handle_worker_control_stream_events(
async fn apply_worker_control_line(
interviewer: &ControlInterviewer,
cancel_token: &AtomicBool,
cancel_token: &CancellationToken,
line: &str,
) {
if line.trim().is_empty() {
@ -243,7 +243,7 @@ async fn apply_worker_control_line(
.await;
}
WorkerControlMessage::RunCancel => {
cancel_token.store(true, Ordering::SeqCst);
cancel_token.cancel();
interviewer.interrupt_all().await;
}
}
@ -561,7 +561,7 @@ fn clone_sandbox_requires_github_credentials(provider: &str) -> bool {
fn install_signal_handlers(
run_control: Arc<RunControlState>,
cancel_token: Arc<AtomicBool>,
cancel_token: CancellationToken,
) -> Result<()> {
#[cfg(unix)]
{
@ -581,17 +581,17 @@ fn install_signal_handlers(
});
let mut terminate = signal(SignalKind::terminate())?;
let terminate_cancel = Arc::clone(&cancel_token);
let terminate_cancel = cancel_token.clone();
tokio::spawn(async move {
while terminate.recv().await.is_some() {
terminate_cancel.store(true, Ordering::SeqCst);
terminate_cancel.cancel();
}
});
let mut interrupt = signal(SignalKind::interrupt())?;
tokio::spawn(async move {
while interrupt.recv().await.is_some() {
cancel_token.store(true, Ordering::SeqCst);
cancel_token.cancel();
}
});
}
@ -606,7 +606,6 @@ fn install_signal_handlers(
)]
mod tests {
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use chrono::Utc;
use fabro_auth::{AuthCredential, AuthDetails};
@ -623,6 +622,7 @@ mod tests {
};
use fabro_vault::{SecretType, Vault};
use fabro_workflow::event::RunEventSink;
use tokio_util::sync::CancellationToken;
use super::{
WorkerControlStreamEvent, WorkerTitlePhase, apply_worker_control_line,
@ -823,7 +823,7 @@ mod tests {
#[tokio::test]
async fn worker_control_line_routes_answer_by_question_id() {
let interviewer = Arc::new(ControlInterviewer::new());
let cancel_token = Arc::new(AtomicBool::new(false));
let cancel_token = CancellationToken::new();
let mut question = Question::new("Approve?", QuestionType::YesNo);
question.id = "q-1".to_string();
let ask_interviewer = Arc::clone(&interviewer);
@ -838,13 +838,13 @@ mod tests {
let answer = answer_task.await.unwrap().answer;
assert_eq!(answer.value, AnswerValue::Yes);
assert!(!cancel_token.load(Ordering::SeqCst));
assert!(!cancel_token.is_cancelled());
}
#[tokio::test]
async fn worker_control_line_cancel_sets_cancel_token_and_interrupts_pending_interviews() {
let interviewer = Arc::new(ControlInterviewer::new());
let cancel_token = Arc::new(AtomicBool::new(false));
let cancel_token = CancellationToken::new();
let mut question = Question::new("Approve?", QuestionType::YesNo);
question.id = "q-1".to_string();
let ask_interviewer = Arc::clone(&interviewer);
@ -860,7 +860,7 @@ mod tests {
let answer = answer_task.await.unwrap().answer;
assert_eq!(answer.value, AnswerValue::Interrupted);
assert!(cancel_token.load(Ordering::SeqCst));
assert!(cancel_token.is_cancelled());
}
#[tokio::test]
@ -893,7 +893,7 @@ mod tests {
#[tokio::test]
async fn worker_control_event_loop_eof_interrupts_pending_interviews() {
let interviewer = Arc::new(ControlInterviewer::new());
let cancel_token = Arc::new(AtomicBool::new(false));
let cancel_token = CancellationToken::new();
let mut question = Question::new("Approve?", QuestionType::YesNo);
question.id = "q-1".to_string();
let ask_interviewer = Arc::clone(&interviewer);
@ -905,14 +905,14 @@ mod tests {
handle_worker_control_stream_events(
Arc::clone(&interviewer),
Arc::clone(&cancel_token),
cancel_token.clone(),
event_rx,
)
.await;
let answer = answer_task.await.unwrap().answer;
assert_eq!(answer.value, AnswerValue::Interrupted);
assert!(!cancel_token.load(Ordering::SeqCst));
assert!(!cancel_token.is_cancelled());
}
#[tokio::test]

View file

@ -34,6 +34,7 @@ async fn run_real_cli_test(provider: Provider, model: &str) {
&emitter,
&env,
None,
tokio_util::sync::CancellationToken::new(),
)
.await
.unwrap_or_else(|_| panic!("CLI backend ({provider}/{model}) should succeed"));

View file

@ -1,5 +1,6 @@
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
#[cfg(test)]
use std::sync::atomic::Ordering;
use std::time::Instant;
use tokio::time::sleep;
@ -18,7 +19,7 @@ use crate::state::ExecutionState;
#[derive(Default)]
pub struct ExecutorOptions {
pub cancel_token: Option<Arc<AtomicBool>>,
pub cancel_token: Option<CancellationToken>,
pub stall_token: Option<CancellationToken>,
pub max_node_visits: Option<usize>,
}
@ -58,7 +59,7 @@ impl<G: Graph + 'static> ExecutorBuilder<G> {
}
#[must_use]
pub fn cancel_token(mut self, token: Arc<AtomicBool>) -> Self {
pub fn cancel_token(mut self, token: CancellationToken) -> Self {
self.options.cancel_token = Some(token);
self
}
@ -95,7 +96,7 @@ impl<G: Graph + 'static> Executor<G> {
loop {
// Check cancellation
if let Some(ref token) = self.options.cancel_token {
if token.load(Ordering::Relaxed) {
if token.is_cancelled() {
state.cancelled = true;
let outcome = Outcome::fail("run cancelled");
self.lifecycle.on_run_end(&outcome, &state).await;
@ -500,7 +501,8 @@ mod tests {
#[tokio::test]
async fn executor_builder_sets_cancel_token() {
let token = Arc::new(AtomicBool::new(true)); // already cancelled
let token = CancellationToken::new();
token.cancel(); // already cancelled
let g = linear_graph(&["start", "end"]);
let state = ExecutionState::new(&g).unwrap();
let executor =
@ -511,6 +513,38 @@ mod tests {
assert!(matches!(result, Err(Error::Cancelled)));
}
#[tokio::test]
async fn executor_cancel_token_fired_during_run_returns_cancelled() {
// Cancel token fired by a handler during the first node; the executor
// checks cancellation at the next node boundary and returns Cancelled.
let token = CancellationToken::new();
let token_clone = token.clone();
struct CancellingHandler(CancellationToken);
#[async_trait]
impl NodeHandler<TestGraph> for CancellingHandler {
async fn execute(
&self,
_node: &TestNode,
_context: &Context,
_g: &TestGraph,
) -> Result<Outcome> {
self.0.cancel();
Ok(Outcome::success())
}
}
let g = linear_graph(&["start", "work", "end"]);
let state = ExecutionState::new(&g).unwrap();
let executor = ExecutorBuilder::new(
Arc::new(CancellingHandler(token_clone)) as Arc<dyn NodeHandler<TestGraph>>
)
.cancel_token(token)
.build();
let result = executor.run(&g, state).await;
assert!(matches!(result, Err(Error::Cancelled)));
}
// ---- Step 9: Terminal nodes, goal gates, visit limits ----
#[tokio::test]
@ -908,10 +942,10 @@ mod tests {
#[tokio::test]
async fn executor_cancellation_stops_run() {
let token = Arc::new(AtomicBool::new(false));
let token = CancellationToken::new();
let token_clone = token.clone();
struct CancellingHandler(Arc<AtomicBool>);
struct CancellingHandler(CancellationToken);
#[async_trait]
impl NodeHandler<TestGraph> for CancellingHandler {
async fn execute(
@ -921,7 +955,7 @@ mod tests {
_g: &TestGraph,
) -> Result<Outcome> {
// Cancel after first node
self.0.store(true, Ordering::Relaxed);
self.0.cancel();
Ok(Outcome::success())
}
}

View file

@ -5,6 +5,7 @@ use std::time::Duration;
use tokio::sync::Notify;
use tokio::task::JoinHandle;
use tokio::time::sleep;
use tokio_util::sync::CancellationToken;
/// Trait for receiving stall timeout notifications.
pub trait ActivityMonitor: Send + Sync {
@ -16,11 +17,11 @@ pub trait ActivityMonitor: Send + Sync {
/// Watches for inactivity and fires a stall timeout if no activity is
/// reported within the configured duration.
pub struct StallWatchdog {
timeout: Duration,
cancel_token: Arc<AtomicBool>,
activity: Arc<Notify>,
shutdown: Arc<AtomicBool>,
monitor: Arc<dyn ActivityMonitor>,
timeout: Duration,
stall_token: CancellationToken,
activity: Arc<Notify>,
shutdown: Arc<AtomicBool>,
monitor: Arc<dyn ActivityMonitor>,
}
/// Guard that resets the stall timer on activity. Drop to stop watching.
@ -33,12 +34,12 @@ pub struct StallGuard {
impl StallWatchdog {
pub fn new(
timeout: Duration,
cancel_token: Arc<AtomicBool>,
stall_token: CancellationToken,
monitor: Arc<dyn ActivityMonitor>,
) -> Self {
Self {
timeout,
cancel_token,
stall_token,
activity: Arc::new(Notify::new()),
shutdown: Arc::new(AtomicBool::new(false)),
monitor,
@ -51,7 +52,7 @@ impl StallWatchdog {
let activity = self.activity.clone();
let shutdown = self.shutdown.clone();
let timeout = self.timeout;
let cancel_token = self.cancel_token;
let stall_token = self.stall_token;
let monitor = self.monitor;
let handle = tokio::spawn(async move {
@ -66,7 +67,7 @@ impl StallWatchdog {
"Stall timeout: no activity detected"
);
monitor.on_stall_timeout(timeout);
cancel_token.store(true, Ordering::Relaxed);
stall_token.cancel();
return;
}
() = activity.notified() => {
@ -136,7 +137,7 @@ mod tests {
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn stall_watchdog_cancels_on_inactivity() {
let cancel = Arc::new(AtomicBool::new(false));
let cancel = CancellationToken::new();
let monitor = TestMonitor::new();
let watchdog =
StallWatchdog::new(Duration::from_millis(50), cancel.clone(), monitor.clone());
@ -145,13 +146,13 @@ mod tests {
// Wait for timeout to fire
sleep(Duration::from_millis(100)).await;
assert!(cancel.load(Ordering::Relaxed));
assert!(cancel.is_cancelled());
assert_eq!(monitor.stalls(), 1);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn stall_watchdog_resets_on_activity() {
let cancel = Arc::new(AtomicBool::new(false));
let cancel = CancellationToken::new();
let monitor = TestMonitor::new();
let watchdog =
StallWatchdog::new(Duration::from_millis(80), cancel.clone(), monitor.clone());
@ -164,17 +165,17 @@ mod tests {
// After another 50ms (100ms total, but only 50ms since activity), should not
// have timed out
sleep(Duration::from_millis(50)).await;
assert!(!cancel.load(Ordering::Relaxed));
assert!(!cancel.is_cancelled());
// Wait long enough for timeout after last activity (80ms + margin)
sleep(Duration::from_millis(60)).await;
assert!(cancel.load(Ordering::Relaxed));
assert!(cancel.is_cancelled());
assert_eq!(monitor.stalls(), 1);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn stall_watchdog_clean_shutdown_on_success() {
let cancel = Arc::new(AtomicBool::new(false));
let cancel = CancellationToken::new();
let monitor = TestMonitor::new();
let watchdog =
StallWatchdog::new(Duration::from_millis(50), cancel.clone(), monitor.clone());
@ -187,13 +188,13 @@ mod tests {
sleep(Duration::from_millis(100)).await;
// Should NOT have triggered
assert!(!cancel.load(Ordering::Relaxed));
assert!(!cancel.is_cancelled());
assert_eq!(monitor.stalls(), 0);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn stall_guard_cleanup_on_drop() {
let cancel = Arc::new(AtomicBool::new(false));
let cancel = CancellationToken::new();
let monitor = TestMonitor::new();
let watchdog =
StallWatchdog::new(Duration::from_millis(50), cancel.clone(), monitor.clone());
@ -206,6 +207,6 @@ mod tests {
sleep(Duration::from_millis(150)).await;
// Cancel should not be set
assert!(!cancel.load(Ordering::Relaxed));
assert!(!cancel.is_cancelled());
}
}

View file

@ -204,7 +204,10 @@ pub async fn run_retro_agent(
// Optionally forward agent events via the callback
let event_forwarder_handle = event_callback.map(|cb| spawn_retro_event_forwarder(&session, cb));
session.initialize().await;
session
.initialize()
.await
.context("Retro agent session initialization failed")?;
let prompt = build_retro_prompt(RETRO_DATA_DIR);

View file

@ -1307,7 +1307,7 @@ impl Sandbox for DaytonaSandbox {
async fn exec_command_streaming(
&self,
command: &str,
timeout_ms: u64,
timeout_ms: Option<u64>,
working_dir: Option<&str>,
env_vars: Option<&HashMap<String, String>>,
cancel_token: Option<CancellationToken>,
@ -1397,7 +1397,7 @@ impl Sandbox for DaytonaSandbox {
&session,
&command_id,
session_exec.exit_code,
Duration::from_millis(timeout_ms),
timeout_ms,
cancel_token.unwrap_or_default(),
&mut stream_task,
)
@ -1730,7 +1730,7 @@ async fn wait_for_completion(
session: &DaytonaSession,
command_id: &str,
initial_exit_code: Option<i32>,
timeout: Duration,
timeout_ms: Option<u64>,
cancel_token: CancellationToken,
stream_task: &mut JoinHandle<Result<(), DaytonaError>>,
) -> crate::Result<WaitOutcome> {
@ -1742,8 +1742,13 @@ async fn wait_for_completion(
});
}
let timeout_sleep = time::sleep(timeout);
tokio::pin!(timeout_sleep);
let timeout_future = async {
match timeout_ms {
Some(ms) => time::sleep(Duration::from_millis(ms)).await,
None => std::future::pending::<()>().await,
}
};
tokio::pin!(timeout_future);
loop {
tokio::select! {
() = time::sleep(Duration::from_millis(250)) => {
@ -1765,7 +1770,7 @@ async fn wait_for_completion(
});
}
}
() = &mut timeout_sleep => {
() = &mut timeout_future => {
return Ok(WaitOutcome {
exit_code: None,
termination: CommandTermination::TimedOut,

View file

@ -362,7 +362,7 @@ impl DockerSandbox {
async fn docker_exec_shell_streaming(
&self,
command: &str,
timeout_ms: u64,
timeout_ms: Option<u64>,
working_dir: Option<&str>,
env_vars: Option<&HashMap<String, String>>,
cancel_token: Option<CancellationToken>,
@ -380,7 +380,13 @@ impl DockerSandbox {
controlled_command,
];
let timeout_duration = Duration::from_millis(timeout_ms);
let timeout_future = async {
match timeout_ms {
Some(ms) => time::sleep(Duration::from_millis(ms)).await,
None => std::future::pending::<()>().await,
}
};
tokio::pin!(timeout_future);
let token = cancel_token.unwrap_or_default();
let container_id = self.container_id()?.to_string();
@ -399,7 +405,7 @@ impl DockerSandbox {
joined
.map_err(|e| crate::Error::context("Docker exec stream task failed", e))??
}
() = time::sleep(timeout_duration) => {
() = &mut timeout_future => {
termination = CommandTermination::TimedOut;
self.request_docker_exec_stop(&stop_file).await?;
output_task
@ -1192,7 +1198,7 @@ impl Sandbox for DockerSandbox {
async fn exec_command_streaming(
&self,
command: &str,
timeout_ms: u64,
timeout_ms: Option<u64>,
working_dir: Option<&str>,
env_vars: Option<&HashMap<String, String>>,
cancel_token: Option<CancellationToken>,

View file

@ -1,5 +1,5 @@
use std::path::{Path, PathBuf};
use std::time::Instant;
use std::time::{Duration, Instant};
use async_trait::async_trait;
use fabro_static::EnvVars;
@ -327,7 +327,7 @@ impl Sandbox for LocalSandbox {
async fn exec_command_streaming(
&self,
command: &str,
timeout_ms: u64,
timeout_ms: Option<u64>,
working_dir: Option<&str>,
env_vars: Option<&std::collections::HashMap<String, String>>,
cancel_token: Option<CancellationToken>,
@ -367,7 +367,13 @@ impl Sandbox for LocalSandbox {
.spawn()
.map_err(|e| crate::Error::context("Failed to spawn command", e))?;
let timeout_duration = std::time::Duration::from_millis(timeout_ms);
let timeout_future = async {
match timeout_ms {
Some(ms) => time::sleep(Duration::from_millis(ms)).await,
None => std::future::pending::<()>().await,
}
};
tokio::pin!(timeout_future);
let token = cancel_token.unwrap_or_default();
let stdout_pipe = child.stdout.take();
@ -387,7 +393,7 @@ impl Sandbox for LocalSandbox {
.map_err(|e| crate::Error::context("Failed to wait for process", e))?;
(CommandTermination::Exited, status.code())
}
() = time::sleep(timeout_duration) => {
() = &mut timeout_future => {
sigterm_then_kill(&mut child).await;
(CommandTermination::TimedOut, None)
}

View file

@ -93,7 +93,7 @@ macro_rules! delegate_sandbox {
async fn exec_command_streaming(
&self,
command: &str,
timeout_ms: u64,
timeout_ms: Option<u64>,
working_dir: Option<&str>,
env_vars: Option<&std::collections::HashMap<String, String>>,
cancel_token: Option<tokio_util::sync::CancellationToken>,
@ -607,14 +607,21 @@ pub trait Sandbox: Send + Sync {
async fn exec_command_streaming(
&self,
command: &str,
timeout_ms: u64,
timeout_ms: Option<u64>,
working_dir: Option<&str>,
env_vars: Option<&std::collections::HashMap<String, String>>,
cancel_token: Option<CancellationToken>,
output_callback: CommandOutputCallback,
) -> crate::Result<ExecStreamingResult> {
let fallback_timeout_ms = timeout_ms.unwrap_or(u64::MAX);
let result = self
.exec_command(command, timeout_ms, working_dir, env_vars, cancel_token)
.exec_command(
command,
fallback_timeout_ms,
working_dir,
env_vars,
cancel_token,
)
.await?;
if !result.stdout.is_empty() {
output_callback(

View file

@ -236,7 +236,7 @@ impl Sandbox for WorktreeSandbox {
async fn exec_command_streaming(
&self,
command: &str,
timeout_ms: u64,
timeout_ms: Option<u64>,
working_dir: Option<&str>,
env_vars: Option<&HashMap<String, String>>,
cancel_token: Option<CancellationToken>,

View file

@ -63,7 +63,7 @@ mod daytona_streaming_live {
sandbox_for_exec
.exec_command_streaming(
"printf 'live-out\\n'; printf 'live-err\\n' >&2; sleep 30",
60_000,
Some(60_000),
None,
None,
Some(cancel_for_exec),
@ -186,7 +186,14 @@ mod daytona_streaming_live {
let chunks = Arc::new(Mutex::new(Vec::new()));
let callback = capture_callback(Arc::clone(&chunks));
let result = sandbox
.exec_command_streaming(command, timeout_ms, None, None, cancel_token, callback)
.exec_command_streaming(
command,
Some(timeout_ms),
None,
None,
cancel_token,
callback,
)
.await?;
let chunks = chunks.lock().await.clone();

View file

@ -49,7 +49,7 @@ async fn streaming_timeout_terminates_docker_exec_before_returning() {
let result = sandbox
.exec_command_streaming(
&format!("trap '' HUP TERM; echo start; sleep 5 # {marker}"),
200,
Some(200),
None,
None,
None,

View file

@ -56,6 +56,7 @@ globset.workspace = true
tower = "0.5"
tower-http = { version = "0.6", features = ["trace"] }
tokio-stream = { workspace = true, features = ["sync"] }
tokio-util.workspace = true
base64.workspace = true
jsonwebtoken.workspace = true
hkdf.workspace = true
@ -108,4 +109,4 @@ tokio-util.workspace = true
fabro-macros = { path = "../fabro-macros" }
fabro-sandbox = { path = "../fabro-sandbox", features = ["test-support"] }
fabro-test = { workspace = true }
fabro-types = { path = "../fabro-types", features = ["test-support"] }
fabro-types = { path = "../fabro-types", features = ["test-support"] }

View file

@ -111,6 +111,7 @@ use tokio::task::spawn_blocking;
use tokio::time::{sleep, timeout};
use tokio_stream::StreamExt;
use tokio_stream::wrappers::{BroadcastStream, UnboundedReceiverStream};
use tokio_util::sync::CancellationToken;
use tower::{ServiceExt, service_fn};
use tracing::{Instrument, debug, error, info, warn};
use ulid::Ulid;
@ -197,7 +198,7 @@ struct ManagedRun {
event_tx: Option<broadcast::Sender<RunEvent>>,
checkpoint: Option<Checkpoint>,
cancel_tx: Option<oneshot::Sender<()>>,
cancel_token: Option<Arc<AtomicBool>>,
cancel_token: Option<CancellationToken>,
worker_pid: Option<u32>,
worker_pgid: Option<u32>,
run_dir: Option<std::path::PathBuf>,
@ -1488,7 +1489,7 @@ async fn delete_run_internal(
if let Some(mut managed_run) = managed_run {
if let Some(token) = &managed_run.cancel_token {
token.store(true, Ordering::SeqCst);
token.cancel();
}
if let Some(answer_transport) = managed_run.answer_transport.clone() {
let _ = answer_transport.cancel_run().await;
@ -2581,12 +2582,12 @@ async fn execute_run_in_process(state: Arc<AppState>, run_id: RunId) {
};
let (cancel_tx, cancel_rx) = oneshot::channel::<()>();
let cancel_token = Arc::new(AtomicBool::new(false));
let cancel_token = CancellationToken::new();
let (event_tx, _) = broadcast::channel(256);
managed_run.status = RunStatus::Starting;
managed_run.cancel_tx = Some(cancel_tx);
managed_run.cancel_token = Some(Arc::clone(&cancel_token));
managed_run.cancel_token = Some(cancel_token.clone());
managed_run.event_tx = Some(event_tx);
(
@ -2679,7 +2680,7 @@ async fn execute_run_in_process(state: Arc<AppState>, run_id: RunId) {
};
let server_settings = state.server_settings();
let github_settings = &server_settings.server.integrations.github;
if cancel_token.load(Ordering::SeqCst) {
if cancel_token.is_cancelled() {
finish_cancelled_run_before_execution(&state, run_id).await;
return;
}
@ -2714,7 +2715,7 @@ async fn execute_run_in_process(state: Arc<AppState>, run_id: RunId) {
let github_app = match github_app_result {
Ok(github_app) => github_app,
Err(e) => {
if cancel_token.load(Ordering::SeqCst) {
if cancel_token.is_cancelled() {
finish_cancelled_run_before_execution(&state, run_id).await;
return;
}
@ -2741,7 +2742,7 @@ async fn execute_run_in_process(state: Arc<AppState>, run_id: RunId) {
.collect();
let services = operations::StartServices {
run_id,
cancel_token: Some(Arc::clone(&cancel_token)),
cancel_token: cancel_token.clone(),
emitter: Arc::clone(&emitter),
interviewer: Arc::clone(&interview_runtime),
run_store: run_store.clone().into(),
@ -2765,7 +2766,7 @@ async fn execute_run_in_process(state: Arc<AppState>, run_id: RunId) {
let result = tokio::select! {
result = execution => ExecutionResult::Completed(Box::new(result)),
_ = cancel_rx => {
cancel_token.store(true, Ordering::SeqCst);
cancel_token.cancel();
ExecutionResult::CancelledBySignal
}
};

View file

@ -1,13 +1,13 @@
use std::sync::Arc;
use super::super::{
ApiError, AppState, FailureReason, ForkRequest, ForkResponse, IntoResponse, Json, Ordering,
Path, Principal, RequiredUser, Response, RewindRequest, RewindResponse, Router,
RunAnswerTransport, RunControlAction, RunExecutionMode, RunId, RunStatus, RunStatusResponse,
StartRunRequest, State, StatusCode, Storage, TimelineEntryResponse, WORKER_CANCEL_GRACE,
WorkflowError, append_control_request, get, load_pending_control, managed_run, operations,
parse_run_id_path, persist_cancelled_run_status, post, reject_if_archived, sleep,
update_live_run_from_event, workflow_event,
ApiError, AppState, FailureReason, ForkRequest, ForkResponse, IntoResponse, Json, Path,
Principal, RequiredUser, Response, RewindRequest, RewindResponse, Router, RunAnswerTransport,
RunControlAction, RunExecutionMode, RunId, RunStatus, RunStatusResponse, StartRunRequest,
State, StatusCode, Storage, TimelineEntryResponse, WORKER_CANCEL_GRACE, WorkflowError,
append_control_request, get, load_pending_control, managed_run, operations, parse_run_id_path,
persist_cancelled_run_status, post, reject_if_archived, sleep, update_live_run_from_event,
workflow_event,
};
pub(super) fn routes() -> Router<Arc<AppState>> {
@ -249,7 +249,7 @@ async fn cancel_run(
}
if let Some(token) = &cancel_token {
token.store(true, Ordering::SeqCst);
token.cancel();
}
let sent_cancel_signal = if let Some(cancel_tx) = cancel_tx {
let _ = cancel_tx.send(());

View file

@ -239,6 +239,20 @@ pub struct AgentCliCompletedProps {
pub duration_ms: u64,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct AgentCliCancelledProps {
pub stdout: String,
pub stderr: String,
pub duration_ms: u64,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct AgentCliTimedOutProps {
pub stdout: String,
pub stderr: String,
pub duration_ms: u64,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct PullRequestCreatedProps {
pub pr_url: String,

View file

@ -252,6 +252,10 @@ pub enum EventBody {
AgentCliStarted(AgentCliStartedProps),
#[serde(rename = "agent.cli.completed")]
AgentCliCompleted(AgentCliCompletedProps),
#[serde(rename = "agent.cli.cancelled")]
AgentCliCancelled(AgentCliCancelledProps),
#[serde(rename = "agent.cli.timed_out")]
AgentCliTimedOut(AgentCliTimedOutProps),
#[serde(rename = "pull_request.created")]
PullRequestCreated(PullRequestCreatedProps),
#[serde(rename = "pull_request.failed")]
@ -433,6 +437,8 @@ impl EventBody {
Self::CommandCompleted(_) => "command.completed",
Self::AgentCliStarted(_) => "agent.cli.started",
Self::AgentCliCompleted(_) => "agent.cli.completed",
Self::AgentCliCancelled(_) => "agent.cli.cancelled",
Self::AgentCliTimedOut(_) => "agent.cli.timed_out",
Self::PullRequestCreated(_) => "pull_request.created",
Self::PullRequestFailed(_) => "pull_request.failed",
Self::DevcontainerResolved(_) => "devcontainer.resolved",

View file

@ -1,5 +1,3 @@
use std::sync::Arc;
use std::sync::atomic::AtomicBool;
use std::time::Instant;
use fabro_agent::sandbox::Sandbox;
@ -7,10 +5,10 @@ use fabro_devcontainer::DevcontainerSpec;
use fabro_sandbox::daytona::{DaytonaSnapshotConfig, DockerfileSource};
use futures::future::try_join_all;
use sha2::{Digest, Sha256};
use tokio_util::sync::CancellationToken;
use crate::error::Error;
use crate::event::{Emitter, Event};
use crate::handler::sandbox_cancel_token;
/// Compute a deterministic snapshot name from Dockerfile content.
pub fn snapshot_name_for_dockerfile(dockerfile: &str) -> String {
@ -39,7 +37,7 @@ pub async fn run_devcontainer_lifecycle(
phase: &str,
commands: &[fabro_devcontainer::Command],
timeout_ms: u64,
cancel_requested: Option<Arc<AtomicBool>>,
cancel_token: CancellationToken,
) -> Result<(), Error> {
if commands.is_empty() {
return Ok(());
@ -61,7 +59,7 @@ pub async fn run_devcontainer_lifecycle(
&format!("sh -c {}", shlex::try_quote(s).unwrap_or_else(|_| s.into())),
index,
timeout_ms,
cancel_requested.clone(),
cancel_token.clone(),
)
.await?;
}
@ -78,7 +76,7 @@ pub async fn run_devcontainer_lifecycle(
&joined,
index,
timeout_ms,
cancel_requested.clone(),
cancel_token.clone(),
)
.await?;
}
@ -92,7 +90,7 @@ pub async fn run_devcontainer_lifecycle(
);
let phase = phase.to_string();
let name = name.clone();
let cancel_requested = cancel_requested.clone();
let cancel_token = cancel_token.clone();
async move {
let cmd_start = Instant::now();
emitter.emit(&Event::DevcontainerLifecycleCommandStarted {
@ -100,14 +98,14 @@ pub async fn run_devcontainer_lifecycle(
command: name.clone(),
index,
});
let cancel_token = sandbox_cancel_token(cancel_requested);
let child_token = cancel_token.child_token();
let result = sandbox
.exec_command(
&command,
timeout_ms,
None,
None,
cancel_token.clone(),
Some(child_token.clone()),
)
.await
.map_err(|e| {
@ -115,12 +113,10 @@ pub async fn run_devcontainer_lifecycle(
"Devcontainer {phase} parallel command '{name}' failed: {e}"
))
})?;
if let Some(token) = &cancel_token {
if token.is_cancelled() {
return Err(Error::Cancelled);
}
token.cancel();
if cancel_token.is_cancelled() {
return Err(Error::Cancelled);
}
child_token.cancel();
let cmd_duration = crate::millis_u64(cmd_start.elapsed());
if !result.is_success() {
let exit_code = result.display_exit_code();
@ -175,7 +171,7 @@ async fn run_single_lifecycle_command(
command: &str,
index: usize,
timeout_ms: u64,
cancel_requested: Option<Arc<AtomicBool>>,
cancel_token: CancellationToken,
) -> Result<(), Error> {
emitter.emit(&Event::DevcontainerLifecycleCommandStarted {
phase: phase.to_string(),
@ -183,19 +179,17 @@ async fn run_single_lifecycle_command(
index,
});
let cmd_start = Instant::now();
let cancel_token = sandbox_cancel_token(cancel_requested);
let child_token = cancel_token.child_token();
let result = sandbox
.exec_command(command, timeout_ms, None, None, cancel_token.clone())
.exec_command(command, timeout_ms, None, None, Some(child_token.clone()))
.await
.map_err(|e| {
Error::engine_with_source(format!("Devcontainer {phase} command failed"), &e)
})?;
if let Some(token) = &cancel_token {
if token.is_cancelled() {
return Err(Error::Cancelled);
}
token.cancel();
if cancel_token.is_cancelled() {
return Err(Error::Cancelled);
}
child_token.cancel();
let cmd_duration = crate::millis_u64(cmd_start.elapsed());
if !result.is_success() {
let exit_code = result.display_exit_code();
@ -227,7 +221,6 @@ async fn run_single_lifecycle_command(
#[cfg(test)]
mod tests {
use std::collections::HashMap;
use std::sync::atomic::AtomicBool;
use std::sync::{Arc, Mutex};
use async_trait::async_trait;
@ -437,9 +430,16 @@ mod tests {
let sandbox = TestSandbox::new();
let emitter = Emitter::default();
let commands = vec![fabro_devcontainer::Command::Shell("echo hi".to_string())];
run_devcontainer_lifecycle(&sandbox, &emitter, "on_create", &commands, 300_000, None)
.await
.unwrap();
run_devcontainer_lifecycle(
&sandbox,
&emitter,
"on_create",
&commands,
300_000,
CancellationToken::new(),
)
.await
.unwrap();
let captured = sandbox.captured_commands();
assert_eq!(captured.len(), 1);
assert!(captured[0].contains("echo hi"), "command: {}", captured[0]);
@ -453,9 +453,16 @@ mod tests {
"echo".to_string(),
"hi".to_string(),
])];
run_devcontainer_lifecycle(&sandbox, &emitter, "on_create", &commands, 300_000, None)
.await
.unwrap();
run_devcontainer_lifecycle(
&sandbox,
&emitter,
"on_create",
&commands,
300_000,
CancellationToken::new(),
)
.await
.unwrap();
let captured = sandbox.captured_commands();
assert_eq!(captured.len(), 1);
assert!(
@ -475,9 +482,16 @@ mod tests {
});
let sandbox = TestSandbox::new();
let commands = vec![fabro_devcontainer::Command::Shell("echo hi".to_string())];
run_devcontainer_lifecycle(&sandbox, &emitter, "on_create", &commands, 300_000, None)
.await
.unwrap();
run_devcontainer_lifecycle(
&sandbox,
&emitter,
"on_create",
&commands,
300_000,
CancellationToken::new(),
)
.await
.unwrap();
let events = events.lock().unwrap();
let started = events[0].properties().unwrap();
assert_eq!(events[0].event_name(), "devcontainer.lifecycle.started");
@ -515,9 +529,15 @@ mod tests {
});
let sandbox = TestSandbox::with_exit_code(1);
let commands = vec![fabro_devcontainer::Command::Shell("false".to_string())];
let result =
run_devcontainer_lifecycle(&sandbox, &emitter, "on_create", &commands, 300_000, None)
.await;
let result = run_devcontainer_lifecycle(
&sandbox,
&emitter,
"on_create",
&commands,
300_000,
CancellationToken::new(),
)
.await;
assert!(result.is_err());
let events = events.lock().unwrap();
let failed = events
@ -550,9 +570,16 @@ mod tests {
events_clone.lock().unwrap().push(event.clone());
});
let sandbox = TestSandbox::new();
run_devcontainer_lifecycle(&sandbox, &emitter, "on_create", &[], 300_000, None)
.await
.unwrap();
run_devcontainer_lifecycle(
&sandbox,
&emitter,
"on_create",
&[],
300_000,
CancellationToken::new(),
)
.await
.unwrap();
assert!(events.lock().unwrap().is_empty());
}
@ -564,9 +591,16 @@ mod tests {
map.insert("install".to_string(), "npm install".to_string());
map.insert("build".to_string(), "npm run build".to_string());
let commands = vec![fabro_devcontainer::Command::Parallel(map)];
run_devcontainer_lifecycle(&sandbox, &emitter, "post_create", &commands, 300_000, None)
.await
.unwrap();
run_devcontainer_lifecycle(
&sandbox,
&emitter,
"post_create",
&commands,
300_000,
CancellationToken::new(),
)
.await
.unwrap();
let captured = sandbox.captured_commands();
assert_eq!(captured.len(), 2);
}
@ -576,7 +610,8 @@ mod tests {
let sandbox = TestSandbox::waiting_for_cancel();
let emitter = Emitter::default();
let commands = vec![fabro_devcontainer::Command::Shell("sleep 5".to_string())];
let cancel_requested = Arc::new(AtomicBool::new(true));
let cancel_token = CancellationToken::new();
cancel_token.cancel();
let result = run_devcontainer_lifecycle(
&sandbox,
@ -584,7 +619,7 @@ mod tests {
"on_create",
&commands,
300_000,
Some(cancel_requested),
cancel_token,
)
.await;
@ -600,7 +635,8 @@ mod tests {
map.insert("install".to_string(), "sleep 5".to_string());
map.insert("build".to_string(), "sleep 5".to_string());
let commands = vec![fabro_devcontainer::Command::Parallel(map)];
let cancel_requested = Arc::new(AtomicBool::new(true));
let cancel_token = CancellationToken::new();
cancel_token.cancel();
let result = run_devcontainer_lifecycle(
&sandbox,
@ -608,7 +644,7 @@ mod tests {
"post_create",
&commands,
300_000,
Some(cancel_requested),
cancel_token,
)
.await;

View file

@ -1006,6 +1006,26 @@ fn event_body_from_event(event: &Event) -> EventBody {
exit_code: *exit_code,
duration_ms: *duration_ms,
}),
Event::AgentCliCancelled {
stdout,
stderr,
duration_ms,
..
} => EventBody::AgentCliCancelled(fabro_types::AgentCliCancelledProps {
stdout: stdout.clone(),
stderr: stderr.clone(),
duration_ms: *duration_ms,
}),
Event::AgentCliTimedOut {
stdout,
stderr,
duration_ms,
..
} => EventBody::AgentCliTimedOut(fabro_types::AgentCliTimedOutProps {
stdout: stdout.clone(),
stderr: stderr.clone(),
duration_ms: *duration_ms,
}),
Event::PullRequestCreated {
pr_url,
pr_number,
@ -1809,6 +1829,48 @@ mod tests {
});
}
#[test]
fn agent_cli_cancelled_maps_to_event_body_with_node_id() {
let stored = to_run_event(&fixtures::RUN_1, &Event::AgentCliCancelled {
node_id: "code".to_string(),
stdout: "out".to_string(),
stderr: "err".to_string(),
duration_ms: 42,
});
assert_eq!(stored.event_name(), "agent.cli.cancelled");
assert_eq!(stored.node_id.as_deref(), Some("code"));
match &stored.body {
EventBody::AgentCliCancelled(props) => {
assert_eq!(props.stdout, "out");
assert_eq!(props.stderr, "err");
assert_eq!(props.duration_ms, 42);
}
other => panic!("expected AgentCliCancelled, got {other:?}"),
}
}
#[test]
fn agent_cli_timed_out_maps_to_event_body_with_node_id() {
let stored = to_run_event(&fixtures::RUN_1, &Event::AgentCliTimedOut {
node_id: "code".to_string(),
stdout: "out".to_string(),
stderr: "err".to_string(),
duration_ms: 99,
});
assert_eq!(stored.event_name(), "agent.cli.timed_out");
assert_eq!(stored.node_id.as_deref(), Some("code"));
match &stored.body {
EventBody::AgentCliTimedOut(props) => {
assert_eq!(props.stdout, "out");
assert_eq!(props.stderr, "err");
assert_eq!(props.duration_ms, 99);
}
other => panic!("expected AgentCliTimedOut, got {other:?}"),
}
}
#[test]
fn stall_watchdog_timeout_populates_watchdog_actor() {
let stored = to_run_event(&fixtures::RUN_1, &Event::StallWatchdogTimeout {

View file

@ -530,6 +530,18 @@ pub enum Event {
exit_code: i32,
duration_ms: u64,
},
AgentCliCancelled {
node_id: String,
stdout: String,
stderr: String,
duration_ms: u64,
},
AgentCliTimedOut {
node_id: String,
stdout: String,
stderr: String,
duration_ms: u64,
},
PullRequestCreated {
pr_url: String,
pr_number: u64,
@ -1247,6 +1259,20 @@ impl Event {
} => {
debug!(node_id, exit_code, duration_ms, "Agent CLI completed");
}
Self::AgentCliCancelled {
node_id,
duration_ms,
..
} => {
debug!(node_id, duration_ms, "Agent CLI cancelled");
}
Self::AgentCliTimedOut {
node_id,
duration_ms,
..
} => {
debug!(node_id, duration_ms, "Agent CLI timed out");
}
Self::PullRequestCreated {
pr_url,
pr_number,

View file

@ -116,6 +116,8 @@ pub fn event_name(event: &Event) -> &'static str {
Event::CommandCompleted { .. } => "command.completed",
Event::AgentCliStarted { .. } => "agent.cli.started",
Event::AgentCliCompleted { .. } => "agent.cli.completed",
Event::AgentCliCancelled { .. } => "agent.cli.cancelled",
Event::AgentCliTimedOut { .. } => "agent.cli.timed_out",
Event::PullRequestCreated { .. } => "pull_request.created",
Event::PullRequestFailed { .. } => "pull_request.failed",
Event::DevcontainerResolved { .. } => "devcontainer.resolved",

View file

@ -116,7 +116,9 @@ fn stored_event_fields_for_variant(event: &Event) -> StoredEventFields {
| Event::CommandStarted { node_id, .. }
| Event::CommandCompleted { node_id, .. }
| Event::AgentCliStarted { node_id, .. }
| Event::AgentCliCompleted { node_id, .. } => node_stored_fields(Some(node_id.clone())),
| Event::AgentCliCompleted { node_id, .. }
| Event::AgentCliCancelled { node_id, .. }
| Event::AgentCliTimedOut { node_id, .. } => node_stored_fields(Some(node_id.clone())),
Event::Agent {
stage,
visit,

View file

@ -7,6 +7,7 @@ use fabro_agent::Sandbox;
use fabro_graphviz::graph::{Graph, Node};
use fabro_template::{TemplateContext, render as render_template};
use fabro_types::RunId;
use tokio_util::sync::CancellationToken;
use super::{EngineServices, Handler};
use crate::context::{Context, WorkflowContext, keys};
@ -44,6 +45,7 @@ pub trait CodergenBackend: Send + Sync {
emitter: &Arc<Emitter>,
sandbox: &Arc<dyn Sandbox>,
tool_hooks: Option<Arc<dyn fabro_agent::ToolHookCallback>>,
cancel_token: CancellationToken,
) -> Result<CodergenResult, Error>;
/// Run a single LLM call with no tools (one_shot mode).
@ -297,6 +299,7 @@ impl Handler for AgentHandler {
&services.run.emitter,
&services.run.sandbox,
tool_hooks,
services.run.cancel_token(),
)
.await;
match result {
@ -307,6 +310,7 @@ impl Handler for AgentHandler {
files_touched,
last_file_touched,
}) => (text, usage, files_touched, last_file_touched),
Err(Error::Cancelled) => return Err(Error::Cancelled),
Err(e) if e.is_retryable() => {
return Err(e);
}
@ -615,6 +619,7 @@ mod tests {
_emitter: &Arc<Emitter>,
_sandbox: &Arc<dyn fabro_agent::Sandbox>,
_tool_hooks: Option<Arc<dyn fabro_agent::ToolHookCallback>>,
_cancel_token: CancellationToken,
) -> Result<CodergenResult, Error> {
Ok(CodergenResult::Text {
text:
@ -675,6 +680,7 @@ mod tests {
_emitter: &Arc<Emitter>,
_sandbox: &Arc<dyn fabro_agent::Sandbox>,
_tool_hooks: Option<Arc<dyn fabro_agent::ToolHookCallback>>,
_cancel_token: CancellationToken,
) -> Result<CodergenResult, Error> {
Ok(CodergenResult::Text {
text: "Done writing results.".to_string(),
@ -736,6 +742,7 @@ mod tests {
emitter: &Arc<Emitter>,
_sandbox: &Arc<dyn fabro_agent::Sandbox>,
_tool_hooks: Option<Arc<dyn fabro_agent::ToolHookCallback>>,
_cancel_token: CancellationToken,
) -> Result<CodergenResult, Error> {
let scope = StageScope::for_handler(context, &node.id);
emitter.emit_scoped(
@ -847,6 +854,7 @@ mod tests {
_emitter: &Arc<Emitter>,
_sandbox: &Arc<dyn Sandbox>,
_tool_hooks: Option<Arc<dyn fabro_agent::ToolHookCallback>>,
_cancel_token: CancellationToken,
) -> Result<CodergenResult, Error> {
*self.captured_thread_id.lock().unwrap() = Some(thread_id.map(String::from));
Ok(CodergenResult::Text {
@ -899,6 +907,7 @@ mod tests {
_emitter: &Arc<Emitter>,
_sandbox: &Arc<dyn Sandbox>,
_tool_hooks: Option<Arc<dyn fabro_agent::ToolHookCallback>>,
_cancel_token: CancellationToken,
) -> Result<CodergenResult, Error> {
*self.captured_thread_id.lock().unwrap() = Some(thread_id.map(String::from));
Ok(CodergenResult::Text {
@ -946,6 +955,7 @@ mod tests {
_emitter: &Arc<Emitter>,
_sandbox: &Arc<dyn Sandbox>,
_tool_hooks: Option<Arc<dyn fabro_agent::ToolHookCallback>>,
_cancel_token: CancellationToken,
) -> Result<CodergenResult, Error> {
Err(Error::handler("Request timed out".to_string()))
}
@ -1093,6 +1103,7 @@ Some text in between.
_emitter: &Arc<Emitter>,
_sandbox: &Arc<dyn Sandbox>,
_tool_hooks: Option<Arc<dyn fabro_agent::ToolHookCallback>>,
_cancel_token: CancellationToken,
) -> Result<CodergenResult, Error> {
Err(Error::Validation("bad config".to_string()))
}
@ -1133,6 +1144,7 @@ Some text in between.
_emitter: &Arc<Emitter>,
_sandbox: &Arc<dyn Sandbox>,
_tool_hooks: Option<Arc<dyn fabro_agent::ToolHookCallback>>,
_cancel_token: CancellationToken,
) -> Result<CodergenResult, Error> {
*self.captured_prompt.lock().unwrap() = Some(prompt.to_string());
Ok(CodergenResult::Text {
@ -1202,6 +1214,7 @@ Some text in between.
_emitter: &Arc<Emitter>,
_sandbox: &Arc<dyn Sandbox>,
_tool_hooks: Option<Arc<dyn fabro_agent::ToolHookCallback>>,
_cancel_token: CancellationToken,
) -> Result<CodergenResult, Error> {
*self.captured_prompt.lock().unwrap() = Some(prompt.to_string());
Ok(CodergenResult::Text {

View file

@ -109,7 +109,7 @@ impl Handler for CommandHandler {
} else {
Some(&services.env)
};
let cancel_token = services.run.sandbox_cancel_token();
let cancel_token = services.run.cancel_token().child_token();
let stage_id = stage_scope.stage_id();
let recorder = CommandLogRecorder::create(run_dir, &stage_id).await?;
let output_callback: CommandOutputCallback = {
@ -130,16 +130,14 @@ impl Handler for CommandHandler {
.sandbox
.exec_command_streaming(
&command,
timeout_ms,
Some(timeout_ms),
None,
env_vars,
cancel_token.clone(),
Some(cancel_token.clone()),
output_callback,
)
.await;
if let Some(token) = cancel_token {
token.cancel();
}
cancel_token.cancel();
let streaming = match result {
Ok(streaming) => streaming,
Err(err) => {
@ -237,7 +235,6 @@ fn tail_bytes(text: &str, max_bytes: usize) -> String {
#[cfg(test)]
mod tests {
use std::sync::Arc;
use std::sync::atomic::AtomicBool;
use std::time::Duration;
use bytes::Bytes;
@ -1126,7 +1123,7 @@ mod tests {
let mut services = make_spy_services(spy.clone());
services.run = services
.run
.with_cancel_requested(Some(Arc::new(AtomicBool::new(false))));
.with_cancel_token(tokio_util::sync::CancellationToken::new());
handler
.execute(&node, &context, &graph, run_dir.path(), &services)

View file

@ -4,6 +4,7 @@ use std::sync::Arc;
use async_trait::async_trait;
use fabro_agent::Sandbox;
use fabro_graphviz::graph::{Graph, Node};
use tokio_util::sync::CancellationToken;
use super::agent::{CodergenBackend, CodergenResult};
use super::{EngineServices, Handler};
@ -86,6 +87,7 @@ impl Handler for FanInHandler {
&node.id,
&services.run.emitter,
&services.run.sandbox,
services.run.cancel_token(),
)
.await?
} else {
@ -223,6 +225,7 @@ async fn llm_evaluate(
node_id: &str,
emitter: &Arc<Emitter>,
sandbox: &Arc<dyn Sandbox>,
cancel_token: CancellationToken,
) -> Result<Candidate, Error> {
let results_text =
serde_json::to_string_pretty(results).unwrap_or_else(|_| results.to_string());
@ -259,6 +262,7 @@ async fn llm_evaluate(
emitter,
sandbox,
None,
cancel_token,
)
.await
{
@ -474,6 +478,7 @@ mod tests {
_emitter: &Arc<Emitter>,
_sandbox: &Arc<dyn Sandbox>,
_tool_hooks: Option<Arc<dyn fabro_agent::ToolHookCallback>>,
_cancel_token: CancellationToken,
) -> Result<CodergenResult, Error> {
// Return text that contains the ID "branch_b"
Ok(CodergenResult::Text {

View file

@ -325,12 +325,7 @@ impl Handler for HumanHandler {
// 5. Handle unanswered / interrupted interview sessions.
if answer.value == AnswerValue::Interrupted {
if services
.run
.cancel_requested
.as_ref()
.is_some_and(|flag| flag.load(Ordering::SeqCst))
{
if services.run.cancel_token().is_cancelled() {
return Err(Error::Cancelled);
}
self.emit(

View file

@ -14,6 +14,8 @@ use fabro_llm::types::{Message, Request, TokenCounts};
use fabro_mcp::config::McpServerSettings;
use fabro_model::{FallbackTarget, Provider};
use tokio::sync::Mutex as TokioMutex;
use tokio::task::JoinHandle;
use tokio_util::sync::CancellationToken;
use super::super::agent::{CodergenBackend, CodergenResult};
use crate::context::keys::Fidelity;
@ -22,6 +24,96 @@ use crate::error::Error;
use crate::event::{Emitter, Event, StageScope};
use crate::outcome::billed_model_usage_from_llm;
/// Spawn a task that, when the run-level token cancels, sets the agent
/// `Session`'s interrupt reason to `Cancelled` and cancels the session token.
///
/// Factored out of `SessionCancelBridgeGuard::replace` so it can be unit-tested
/// without constructing a real `Session`.
fn spawn_bridge_task(
run_token: CancellationToken,
interrupt_reason: Arc<Mutex<Option<fabro_agent::InterruptReason>>>,
session_token: CancellationToken,
) -> JoinHandle<()> {
tokio::spawn(async move {
run_token.cancelled().await;
{
let mut guard = interrupt_reason
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if guard.is_none() {
*guard = Some(fabro_agent::InterruptReason::Cancelled);
}
}
session_token.cancel();
})
}
/// Per-invocation guard that maps a run-level `CancellationToken` to an agent
/// `Session`'s interrupt reason and cancel token.
///
/// Dropping the guard aborts the spawned bridge task so a still-cached session
/// (after success) is not left wired to a stale run token.
struct SessionCancelBridgeGuard {
handle: Option<JoinHandle<()>>,
}
impl SessionCancelBridgeGuard {
fn new() -> Self {
Self { handle: None }
}
fn replace(&mut self, run_token: CancellationToken, session: &Session) {
self.abort();
self.handle = Some(spawn_bridge_task(
run_token,
session.interrupt_reason_handle(),
session.cancel_token(),
));
}
fn abort(&mut self) {
if let Some(handle) = self.handle.take() {
handle.abort();
}
}
}
impl Drop for SessionCancelBridgeGuard {
fn drop(&mut self) {
self.abort();
}
}
/// Classification of an `fabro_agent::Error` for the API backend's `run` path.
enum AgentApiErrorDisposition {
/// Session was interrupted via cancellation; surface as `Error::Cancelled`.
Cancelled,
/// Underlying LLM error eligible for provider failover.
FailoverEligible(fabro_llm::Error),
/// Terminal error; abort the invocation with this workflow `Error`.
Terminal(Error),
}
fn classify_agent_error(err: fabro_agent::Error, allow_failover: bool) -> AgentApiErrorDisposition {
match err {
fabro_agent::Error::Interrupted(fabro_agent::InterruptReason::Cancelled) => {
AgentApiErrorDisposition::Cancelled
}
fabro_agent::Error::Interrupted(fabro_agent::InterruptReason::WallClockTimeout) => {
AgentApiErrorDisposition::Terminal(Error::Precondition(
"Agent session hit its wall-clock timeout".to_string(),
))
}
fabro_agent::Error::Llm(err) if allow_failover && err.failover_eligible() => {
AgentApiErrorDisposition::FailoverEligible(err)
}
fabro_agent::Error::Llm(err) => AgentApiErrorDisposition::Terminal(Error::Llm(err)),
other => AgentApiErrorDisposition::Terminal(Error::handler(format!(
"Agent session failed: {other}"
))),
}
}
fn build_profile(model: &str, provider: Provider) -> Box<dyn AgentProfile> {
match provider {
Provider::OpenAi => Box::new(OpenAiProfile::new(model)),
@ -426,6 +518,7 @@ impl CodergenBackend for AgentApiBackend {
emitter: &Arc<Emitter>,
sandbox: &Arc<dyn Sandbox>,
tool_hooks: Option<Arc<dyn fabro_agent::ToolHookCallback>>,
cancel_token: CancellationToken,
) -> Result<CodergenResult, Error> {
let actual_model = node.model().unwrap_or(&self.model).to_string();
let _actual_provider = node
@ -440,7 +533,14 @@ impl CodergenBackend for AgentApiBackend {
None
};
// Take a cached session if reusing, otherwise create a new one.
let mut bridge = SessionCancelBridgeGuard::new();
// Take a cached session if reusing, otherwise create a new one. Cancel
// checks bracket `Client::from_source(...)` so cancellation arriving
// during credential refresh is not lost.
if cancel_token.is_cancelled() {
return Err(Error::Cancelled);
}
let (mut session, is_reused) = if let Some(ref key) = reuse_key {
let existing = self.sessions.lock().unwrap().remove(key);
if let Some(s) = existing {
@ -459,6 +559,10 @@ impl CodergenBackend for AgentApiBackend {
false,
)
};
if cancel_token.is_cancelled() {
return Err(Error::Cancelled);
}
bridge.replace(cancel_token.clone(), &session);
tracing::info!(
node = %node.id,
@ -487,99 +591,157 @@ impl CodergenBackend for AgentApiBackend {
// Record turn count before processing so we only aggregate new usage.
let turns_before = session.history().turns().len();
if !is_reused {
session.initialize().await;
}
let result = session.process_input(prompt).await;
// On failover-eligible error, try fallback providers.
let result = match result {
Ok(()) => Ok(()),
Err(fabro_agent::Error::Llm(ref sdk_err))
if sdk_err.failover_eligible() && !self.fallback_chain.is_empty() =>
{
let error_msg = sdk_err.to_string();
let from_provider = self.provider.to_string();
let from_model = self.model.clone();
let mut last_err = Error::Llm(sdk_err.clone());
let mut succeeded = false;
for target in &self.fallback_chain {
emitter.emit_scoped(
&Event::Failover {
stage: node.id.clone(),
from_provider: from_provider.clone(),
from_model: from_model.clone(),
to_provider: target.provider.clone(),
to_model: target.model.clone(),
error: error_msg.clone(),
},
&stage_scope,
);
let target_provider: Provider = match target.provider.parse() {
Ok(p) => p,
Err(_) => continue,
};
let new_session = match Self::create_session_for(
&target.model,
target_provider,
node,
sandbox,
self.source.as_ref(),
&self.env,
tool_hooks.clone(),
self.mcp_servers.clone(),
)
.await
{
Ok(s) => s,
Err(e) => {
last_err = e;
continue;
}
};
session = new_session;
// Re-subscribe to forward events + track files from the new session
spawn_event_forwarder(
&session,
node.id.clone(),
stage_scope.clone(),
Arc::clone(emitter),
Arc::clone(&file_tracking),
);
session.initialize().await;
match session.process_input(prompt).await {
Ok(()) => {
succeeded = true;
break;
}
Err(fabro_agent::Error::Llm(err)) if err.failover_eligible() => {
last_err = Error::Llm(err);
}
Err(fabro_agent::Error::Llm(err)) => return Err(Error::Llm(err)),
Err(fabro_agent::Error::Interrupted(_)) => {
return Err(Error::Cancelled);
}
Err(other) => {
return Err(Error::handler(format!("Agent session failed: {other}")));
}
let allow_failover_primary = !self.fallback_chain.is_empty();
let init_result = if is_reused {
Ok(())
} else {
match session.initialize().await {
Ok(()) => Ok(()),
Err(err) => match classify_agent_error(err, allow_failover_primary) {
AgentApiErrorDisposition::Cancelled => {
bridge.abort();
return Err(Error::Cancelled);
}
}
if succeeded { Ok(()) } else { Err(last_err) }
AgentApiErrorDisposition::Terminal(err) => {
bridge.abort();
return Err(err);
}
AgentApiErrorDisposition::FailoverEligible(sdk_err) => {
Err(fabro_agent::Error::Llm(sdk_err))
}
},
}
Err(fabro_agent::Error::Llm(sdk_err)) => Err(Error::Llm(sdk_err)),
Err(fabro_agent::Error::Interrupted(_)) => Err(Error::Cancelled),
Err(other) => Err(Error::handler(format!("Agent session failed: {other}"))),
};
// On error, drop the session (don't cache failed state).
// If initialize failed with a failover-eligible error, treat as a
// process_input failover trigger; otherwise run process_input.
let result = match init_result {
Ok(()) => session.process_input(prompt).await,
Err(err) => Err(err),
};
// On failover-eligible error, try fallback providers.
let result: Result<(), Error> = match result {
Ok(()) => Ok(()),
Err(err) => match classify_agent_error(err, allow_failover_primary) {
AgentApiErrorDisposition::Cancelled => {
bridge.abort();
return Err(Error::Cancelled);
}
AgentApiErrorDisposition::Terminal(err) => {
bridge.abort();
return Err(err);
}
AgentApiErrorDisposition::FailoverEligible(sdk_err) => {
let error_msg = sdk_err.to_string();
let from_provider = self.provider.to_string();
let from_model = self.model.clone();
let mut last_err = Error::Llm(sdk_err);
let mut succeeded = false;
for (index, target) in self.fallback_chain.iter().enumerate() {
emitter.emit_scoped(
&Event::Failover {
stage: node.id.clone(),
from_provider: from_provider.clone(),
from_model: from_model.clone(),
to_provider: target.provider.clone(),
to_model: target.model.clone(),
error: error_msg.clone(),
},
&stage_scope,
);
let target_provider: Provider = match target.provider.parse() {
Ok(p) => p,
Err(_) => continue,
};
// Detach the bridge from the failing session before
// refreshing credentials and building a new one.
bridge.abort();
if cancel_token.is_cancelled() {
return Err(Error::Cancelled);
}
let new_session = match Self::create_session_for(
&target.model,
target_provider,
node,
sandbox,
self.source.as_ref(),
&self.env,
tool_hooks.clone(),
self.mcp_servers.clone(),
)
.await
{
Ok(s) => s,
Err(e) => {
last_err = e;
continue;
}
};
if cancel_token.is_cancelled() {
return Err(Error::Cancelled);
}
session = new_session;
bridge.replace(cancel_token.clone(), &session);
// Re-subscribe to forward events + track files from the new session
spawn_event_forwarder(
&session,
node.id.clone(),
stage_scope.clone(),
Arc::clone(emitter),
Arc::clone(&file_tracking),
);
let allow_failover_next = index + 1 < self.fallback_chain.len();
if let Err(err) = session.initialize().await {
match classify_agent_error(err, allow_failover_next) {
AgentApiErrorDisposition::Cancelled => {
bridge.abort();
return Err(Error::Cancelled);
}
AgentApiErrorDisposition::Terminal(err) => {
bridge.abort();
return Err(err);
}
AgentApiErrorDisposition::FailoverEligible(sdk_err) => {
last_err = Error::Llm(sdk_err);
continue;
}
}
}
match session.process_input(prompt).await {
Ok(()) => {
succeeded = true;
break;
}
Err(err) => match classify_agent_error(err, allow_failover_next) {
AgentApiErrorDisposition::Cancelled => {
bridge.abort();
return Err(Error::Cancelled);
}
AgentApiErrorDisposition::Terminal(err) => {
bridge.abort();
return Err(err);
}
AgentApiErrorDisposition::FailoverEligible(sdk_err) => {
last_err = Error::Llm(sdk_err);
}
},
}
}
if succeeded { Ok(()) } else { Err(last_err) }
}
},
};
// On error, drop the session (don't cache failed state). The bridge's
// `Drop` will abort the spawned task on early return.
result?;
// Aggregate token usage only from new turns (prevents double-counting on
@ -622,8 +784,10 @@ impl CodergenBackend for AgentApiBackend {
(v, s.last.clone())
};
// Cache session back for reuse on success.
// Cache session back for reuse on success. Detach the bridge first so
// the cached session is not left wired to this run's cancel token.
if let Some(key) = reuse_key {
bridge.abort();
self.sessions.lock().unwrap().insert(key, session);
}
@ -640,6 +804,7 @@ impl CodergenBackend for AgentApiBackend {
mod tests {
use fabro_agent::subagent::SessionFactory;
use fabro_auth::{AuthCredential, AuthDetails, VaultCredentialSource};
use fabro_llm::{Error as LlmError, ProviderErrorDetail, ProviderErrorKind};
use fabro_vault::{SecretType, Vault};
use tokio::sync::RwLock as AsyncRwLock;
@ -820,4 +985,238 @@ mod tests {
assert_eq!(client.provider_names(), vec!["anthropic"]);
}
// --- Bridge guard tests ---
fn failover_eligible_llm_error() -> LlmError {
LlmError::Network {
message: "boom".into(),
source: None,
}
}
fn non_failover_llm_error() -> LlmError {
LlmError::Provider {
kind: ProviderErrorKind::Authentication,
detail: Box::new(ProviderErrorDetail {
message: "bad key".into(),
provider: "openai".into(),
status_code: Some(401),
error_code: None,
retry_after: None,
raw: None,
}),
}
}
#[tokio::test]
async fn spawn_bridge_task_sets_cancelled_and_cancels_session_token() {
let run_token = CancellationToken::new();
let interrupt_reason = Arc::new(Mutex::new(None));
let session_token = CancellationToken::new();
let handle = spawn_bridge_task(
run_token.clone(),
Arc::clone(&interrupt_reason),
session_token.clone(),
);
assert!(!session_token.is_cancelled());
assert!(interrupt_reason.lock().unwrap().is_none());
run_token.cancel();
handle.await.unwrap();
assert!(session_token.is_cancelled());
assert_eq!(
*interrupt_reason.lock().unwrap(),
Some(fabro_agent::InterruptReason::Cancelled)
);
}
#[tokio::test]
async fn spawn_bridge_task_preserves_existing_interrupt_reason() {
let run_token = CancellationToken::new();
let interrupt_reason = Arc::new(Mutex::new(Some(
fabro_agent::InterruptReason::WallClockTimeout,
)));
let session_token = CancellationToken::new();
let handle = spawn_bridge_task(
run_token.clone(),
Arc::clone(&interrupt_reason),
session_token.clone(),
);
run_token.cancel();
handle.await.unwrap();
// Existing reason wins; the bridge does not overwrite a wall-clock
// timeout already recorded by the session.
assert_eq!(
*interrupt_reason.lock().unwrap(),
Some(fabro_agent::InterruptReason::WallClockTimeout)
);
assert!(session_token.is_cancelled());
}
#[tokio::test]
async fn bridge_guard_drop_aborts_pending_task() {
let run_token = CancellationToken::new();
let interrupt_reason = Arc::new(Mutex::new(None));
let session_token = CancellationToken::new();
{
let mut guard = SessionCancelBridgeGuard::new();
guard.handle = Some(spawn_bridge_task(
run_token.clone(),
Arc::clone(&interrupt_reason),
session_token.clone(),
));
// guard dropped here
}
// Trigger the run token after the guard has been dropped. The aborted
// task must not write to interrupt_reason or cancel session_token.
run_token.cancel();
// Yield enough times for any errant task to run.
for _ in 0..10 {
tokio::task::yield_now().await;
}
assert!(interrupt_reason.lock().unwrap().is_none());
assert!(!session_token.is_cancelled());
}
#[tokio::test]
async fn bridge_guard_replace_aborts_prior_task() {
// First (prior) bridge wiring.
let prior_run_token = CancellationToken::new();
let prior_interrupt_reason = Arc::new(Mutex::new(None));
let prior_session_token = CancellationToken::new();
// Second (replacement) bridge wiring.
let new_run_token = CancellationToken::new();
let new_interrupt_reason = Arc::new(Mutex::new(None));
let new_session_token = CancellationToken::new();
let mut guard = SessionCancelBridgeGuard::new();
guard.handle = Some(spawn_bridge_task(
prior_run_token.clone(),
Arc::clone(&prior_interrupt_reason),
prior_session_token.clone(),
));
// Replace with a new task pointing at different handles.
guard.handle = {
// Manually mirror `replace` semantics: abort then install.
if let Some(h) = guard.handle.take() {
h.abort();
}
Some(spawn_bridge_task(
new_run_token.clone(),
Arc::clone(&new_interrupt_reason),
new_session_token.clone(),
))
};
// Cancelling the prior run token must not affect anything because the
// prior task was aborted by `replace`.
prior_run_token.cancel();
for _ in 0..10 {
tokio::task::yield_now().await;
}
assert!(prior_interrupt_reason.lock().unwrap().is_none());
assert!(!prior_session_token.is_cancelled());
// The replacement task must still be alive and react to its own token.
new_run_token.cancel();
guard.handle.take().unwrap().await.unwrap();
assert_eq!(
*new_interrupt_reason.lock().unwrap(),
Some(fabro_agent::InterruptReason::Cancelled)
);
assert!(new_session_token.is_cancelled());
}
// --- classify_agent_error tests ---
#[test]
fn classify_interrupted_cancelled_is_cancelled() {
let err = fabro_agent::Error::Interrupted(fabro_agent::InterruptReason::Cancelled);
assert!(matches!(
classify_agent_error(err, true),
AgentApiErrorDisposition::Cancelled
));
}
#[test]
fn classify_interrupted_wall_clock_is_terminal_precondition() {
let err = fabro_agent::Error::Interrupted(fabro_agent::InterruptReason::WallClockTimeout);
match classify_agent_error(err, true) {
AgentApiErrorDisposition::Terminal(Error::Precondition(msg)) => {
assert!(msg.contains("wall-clock"));
}
_ => panic!("expected Terminal(Error::Precondition) for WallClockTimeout"),
}
}
#[test]
fn classify_failover_eligible_llm_returns_failover_when_allowed() {
let err = fabro_agent::Error::Llm(failover_eligible_llm_error());
assert!(matches!(
classify_agent_error(err, true),
AgentApiErrorDisposition::FailoverEligible(_)
));
}
#[test]
fn classify_failover_eligible_llm_returns_terminal_when_not_allowed() {
let err = fabro_agent::Error::Llm(failover_eligible_llm_error());
match classify_agent_error(err, false) {
AgentApiErrorDisposition::Terminal(Error::Llm(_)) => {}
_ => panic!("expected Terminal(Error::Llm) when failover disallowed"),
}
}
#[test]
fn classify_non_failover_eligible_llm_is_terminal_llm() {
let err = fabro_agent::Error::Llm(non_failover_llm_error());
match classify_agent_error(err, true) {
AgentApiErrorDisposition::Terminal(Error::Llm(_)) => {}
_ => panic!("expected Terminal(Error::Llm) for non-failover-eligible LLM error"),
}
}
#[test]
fn classify_session_closed_is_terminal_handler() {
let err = fabro_agent::Error::SessionClosed;
match classify_agent_error(err, true) {
AgentApiErrorDisposition::Terminal(Error::Handler { message, .. }) => {
assert!(message.contains("Agent session failed"));
}
_ => panic!("expected Terminal(Error::Handler) for SessionClosed"),
}
}
#[test]
fn classify_invalid_state_is_terminal_handler() {
let err = fabro_agent::Error::InvalidState("oops".into());
match classify_agent_error(err, true) {
AgentApiErrorDisposition::Terminal(Error::Handler { message, .. }) => {
assert!(message.contains("Agent session failed"));
}
_ => panic!("expected Terminal(Error::Handler) for InvalidState"),
}
}
#[test]
fn classify_tool_execution_is_terminal_handler() {
let err = fabro_agent::Error::ToolExecution("tool blew up".into());
match classify_agent_error(err, true) {
AgentApiErrorDisposition::Terminal(Error::Handler { message, .. }) => {
assert!(message.contains("Agent session failed"));
}
_ => panic!("expected Terminal(Error::Handler) for ToolExecution"),
}
}
}

View file

@ -3,14 +3,14 @@ use std::sync::Arc;
use async_trait::async_trait;
use fabro_agent::Sandbox;
use fabro_agent::sandbox::ExecResult;
use fabro_auth::{CliAgentKind, CredentialResolver, CredentialUsage, ResolvedCredential};
use fabro_graphviz::graph::Node;
use fabro_llm::types::TokenCounts;
use fabro_model::Provider;
use fabro_types::CommandTermination;
use fabro_types::{CommandOutputStream, CommandTermination};
use fabro_util::time::elapsed_ms;
use tokio::time::sleep;
use tokio::sync::Mutex as TokioMutex;
use tokio_util::sync::CancellationToken;
use super::super::agent::{CodergenBackend, CodergenResult};
use crate::context::Context;
@ -66,6 +66,7 @@ async fn ensure_cli(
provider: Provider,
sandbox: &Arc<dyn Sandbox>,
emitter: &Arc<Emitter>,
cancel_token: &CancellationToken,
) -> Result<(), Error> {
let start = std::time::Instant::now();
let cli_name = cli.name();
@ -84,7 +85,7 @@ async fn ensure_cli(
30_000,
None,
None,
None,
Some(cancel_token.child_token()),
)
.await
.map_err(|e| {
@ -112,7 +113,13 @@ async fn ensure_cli(
cli.npm_package()
);
let install_result = sandbox
.exec_command(&install_cmd, 180_000, None, None, None)
.exec_command(
&install_cmd,
180_000,
None,
None,
Some(cancel_token.child_token()),
)
.await
.map_err(|e| Error::handler_with_source(format!("Failed to install {cli_name}"), &e))?;
@ -477,6 +484,7 @@ impl CodergenBackend for AgentCliBackend {
emitter: &Arc<Emitter>,
sandbox: &Arc<dyn Sandbox>,
_tool_hooks: Option<Arc<dyn fabro_agent::ToolHookCallback>>,
cancel_token: CancellationToken,
) -> Result<CodergenResult, Error> {
// 1. Snapshot git state before the CLI run
let files_before = self.detect_changed_files(sandbox).await;
@ -485,9 +493,6 @@ impl CodergenBackend for AgentCliBackend {
let run_id = uuid::Uuid::new_v4().to_string();
let tmp_prefix = format!("/tmp/fabro_cli_{run_id}");
let prompt_path = format!("{tmp_prefix}_prompt.txt");
let stdout_path = format!("{tmp_prefix}_stdout.log");
let stderr_path = format!("{tmp_prefix}_stderr.log");
let exit_code_path = format!("{tmp_prefix}_exit_code");
let env_path = format!("{tmp_prefix}_env.sh");
sandbox
@ -504,7 +509,7 @@ impl CodergenBackend for AgentCliBackend {
// Ensure the CLI tool is installed in the sandbox
let cli = AgentCli::for_provider(provider);
ensure_cli(cli, provider, sandbox, emitter).await?;
ensure_cli(cli, provider, sandbox, emitter, &cancel_token).await?;
let command = cli_command_for_provider(provider, model, &prompt_path);
let stage_scope = StageScope::for_handler(context, &node.id);
@ -521,11 +526,8 @@ impl CodergenBackend for AgentCliBackend {
);
// Forward provider API key and custom env vars so the CLI tool can
// authenticate. Build a HashMap to pass via exec_command's env_vars
// parameter — this prepends `export` statements directly into the
// base64-encoded command, avoiding filesystem-to-process race
// conditions that can occur when writing an env file via the fs API and
// sourcing it via the process API.
// authenticate. Resolve credentials and run any pre-login command
// before the main CLI invocation.
let cli_agent = match cli {
AgentCli::Claude => CliAgentKind::Claude,
AgentCli::Codex => CliAgentKind::Codex,
@ -541,7 +543,13 @@ impl CodergenBackend for AgentCliBackend {
};
if let Some(login_cmd) = &cli_credential.login_command {
let login_result = sandbox
.exec_command(login_cmd, 30_000, None, None, None)
.exec_command(
login_cmd,
30_000,
None,
None,
Some(cancel_token.child_token()),
)
.await
.map_err(|e| Error::handler_with_source("codex login failed", &e))?;
if !login_result.is_success() {
@ -566,102 +574,176 @@ impl CodergenBackend for AgentCliBackend {
launch_env.insert(name.clone(), val.clone());
}
// Also write env file as fallback for commands that source it (e.g. ensure_cli
// PATH)
// Write env file so the inner shell that runs the CLI command picks up
// PATH and provider env vars; we still pass `launch_env` to
// `exec_command_streaming` for parity.
let mut env_lines: Vec<String> = vec!["export PATH=\"$HOME/.local/bin:$PATH\"".to_string()];
env_lines.extend(
launch_env
.iter()
.map(|(k, v)| format!("export {k}={}", shell_quote(v))),
);
{
sandbox
.write_file(&env_path, &env_lines.join("\n"))
.await
.map_err(|e| Error::handler_with_source("Failed to write env file", &e))?;
}
sandbox
.write_file(&env_path, &env_lines.join("\n"))
.await
.map_err(|e| Error::handler_with_source("Failed to write env file", &e))?;
// 3a. Disable auto-stop so the sandbox stays alive during long CLI runs
// Disable auto-stop so the sandbox stays alive during long CLI runs.
if let Err(e) = sandbox.set_autostop_interval(0).await {
tracing::warn!("Failed to disable sandbox auto-stop: {e}");
}
// 3b. Launch CLI command in background (env file is always written)
let inner_command = format!(". {env_path} && {command}");
// Use setsid (if available) to create a new session so the child process is
// fully detached from the shell. Without this, Daytona's POST /process/execute
// blocks until ALL descendant processes exit, causing a 60s HTTP timeout.
// $SID is empty on macOS (where setsid doesn't exist but isn't needed since
// the local exec implementation doesn't wait for grandchildren).
let bg_command = format!(
"SID=$(command -v setsid || true)\n$SID sh -c '{inner_command} > {stdout_path} 2>{stderr_path}; echo $? > {exit_code_path}' </dev/null >/dev/null 2>&1 &\necho $!"
);
let launch_start = std::time::Instant::now();
// Stream the CLI command directly: the previous detached `setsid &`
// launcher could not be cancelled mid-flight. By running through
// `exec_command_streaming` the run-level cancel token (and node
// timeout, when set) terminate the CLI and its descendants.
let outer_command = format!(". {env_path} && {command}");
let stdout_buffer: Arc<TokioMutex<Vec<u8>>> = Arc::new(TokioMutex::new(Vec::new()));
let stderr_buffer: Arc<TokioMutex<Vec<u8>>> = Arc::new(TokioMutex::new(Vec::new()));
let stdout_buf_cb = Arc::clone(&stdout_buffer);
let stderr_buf_cb = Arc::clone(&stderr_buffer);
let emitter_for_callback = Arc::clone(emitter);
let output_callback: fabro_agent::CommandOutputCallback = Arc::new(move |stream, bytes| {
let stdout_buf = Arc::clone(&stdout_buf_cb);
let stderr_buf = Arc::clone(&stderr_buf_cb);
let emitter = Arc::clone(&emitter_for_callback);
Box::pin(async move {
// Touch the stall watchdog whenever the CLI emits output
// so long-running invocations don't trip stall timeout.
emitter.touch();
match stream {
CommandOutputStream::Stdout => {
stdout_buf.lock().await.extend_from_slice(&bytes);
}
CommandOutputStream::Stderr => {
stderr_buf.lock().await.extend_from_slice(&bytes);
}
}
Ok(())
})
});
let launch_env_ref = if launch_env.is_empty() {
None
} else {
Some(&launch_env)
};
let launch_result = sandbox
.exec_command(&bg_command, 30_000, None, launch_env_ref, None)
.await
.map_err(|e| Error::handler_with_source("Failed to launch CLI command", &e))?;
let pid = launch_result.stdout.trim();
tracing::info!(pid, "CLI process launched in background");
let timeout_ms = node.timeout().map(crate::millis_u64);
let invocation_token = cancel_token.child_token();
let launch_start = std::time::Instant::now();
let streaming_result = sandbox
.exec_command_streaming(
&outer_command,
timeout_ms,
None,
launch_env_ref,
Some(invocation_token.clone()),
output_callback,
)
.await;
// 3c. Poll for completion
let poll_command =
format!("[ -f {exit_code_path} ] && cat {exit_code_path} || echo running");
let poll_interval = self.poll_interval;
let exit_code: i32 = loop {
sleep(poll_interval).await;
emitter.touch(); // keep the stall watchdog alive while polling
let poll_result = sandbox
.exec_command(&poll_command, 30_000, None, None, None)
.await
.map_err(|e| Error::handler_with_source("Failed to poll CLI command", &e))?;
let status = poll_result.stdout.trim();
if status != "running" {
break status.parse::<i32>().unwrap_or(-1);
let cleanup_temp_files = || {
let sandbox = Arc::clone(sandbox);
let cleanup_cmd = format!("rm -f {tmp_prefix}_*");
async move {
let _ = sandbox
.exec_command(&cleanup_cmd, 30_000, None, None, None)
.await;
}
};
// 3d. Read results
let duration_ms = u64::try_from(launch_start.elapsed().as_millis()).unwrap_or(u64::MAX);
let stdout_result = sandbox
.exec_command(&format!("cat {stdout_path}"), 60_000, None, None, None)
.await
.map_err(|e| Error::handler_with_source("Failed to read stdout", &e))?;
let stderr_result = sandbox
.exec_command(&format!("cat {stderr_path}"), 60_000, None, None, None)
.await
.map_err(|e| Error::handler_with_source("Failed to read stderr", &e))?;
let result = ExecResult {
stdout: stdout_result.stdout,
stderr: stderr_result.stdout,
exit_code: Some(exit_code),
termination: CommandTermination::Exited,
duration_ms,
let streaming = match streaming_result {
Ok(streaming) => streaming,
Err(err) => {
cleanup_temp_files().await;
return Err(Error::handler_with_source(
"Failed to run CLI command",
&err,
));
}
};
emitter.emit_scoped(
&Event::AgentCliCompleted {
node_id: node.id.clone(),
stdout: result.stdout.clone(),
stderr: result.stderr.clone(),
exit_code: result.exit_code.unwrap_or(-1),
duration_ms: result.duration_ms,
},
&stage_scope,
);
let result = streaming.result;
// Prefer the buffered streaming output (live chunks); fall back to the
// result struct for sandboxes that bundle output at the end.
let buffered_stdout = String::from_utf8_lossy(&stdout_buffer.lock().await).into_owned();
let buffered_stderr = String::from_utf8_lossy(&stderr_buffer.lock().await).into_owned();
let stdout = if buffered_stdout.is_empty() {
result.stdout.clone()
} else {
buffered_stdout
};
let stderr = if buffered_stderr.is_empty() {
result.stderr.clone()
} else {
buffered_stderr
};
let duration_ms = u64::try_from(launch_start.elapsed().as_millis()).unwrap_or(u64::MAX);
// 3e. Cleanup temp files
let _ = sandbox
.exec_command(&format!("rm -f {tmp_prefix}_*"), 30_000, None, None, None)
.await;
match result.termination {
CommandTermination::Cancelled => {
emitter.emit_scoped(
&Event::AgentCliCancelled {
node_id: node.id.clone(),
stdout: stdout.clone(),
stderr: stderr.clone(),
duration_ms,
},
&stage_scope,
);
cleanup_temp_files().await;
return Err(Error::Cancelled);
}
CommandTermination::TimedOut => {
emitter.emit_scoped(
&Event::AgentCliTimedOut {
node_id: node.id.clone(),
stdout: stdout.clone(),
stderr: stderr.clone(),
duration_ms,
},
&stage_scope,
);
cleanup_temp_files().await;
let tail = |s: &str, n: usize| -> String {
s.chars()
.rev()
.take(n)
.collect::<Vec<_>>()
.into_iter()
.rev()
.collect()
};
let stderr_tail = tail(&stderr, 500);
let stdout_tail = tail(&stdout, 500);
let detail = match (stderr_tail.is_empty(), stdout_tail.is_empty()) {
(false, false) => format!("{stderr_tail}\nstdout: {stdout_tail}"),
(false, true) => stderr_tail,
(true, false) => format!("stdout: {stdout_tail}"),
(true, true) => format!("command: {command}"),
};
return Err(Error::handler(format!(
"CLI command timed out after {duration_ms} ms: {detail}"
)));
}
CommandTermination::Exited => {
emitter.emit_scoped(
&Event::AgentCliCompleted {
node_id: node.id.clone(),
stdout: stdout.clone(),
stderr: stderr.clone(),
exit_code: result.exit_code.unwrap_or(-1),
duration_ms,
},
&stage_scope,
);
}
}
if !result.is_success() {
// Cleanup temp files (Exited path).
cleanup_temp_files().await;
let exited_success =
result.termination == CommandTermination::Exited && result.exit_code == Some(0);
if !exited_success {
let tail = |s: &str, n: usize| -> String {
s.chars()
.rev()
@ -671,8 +753,8 @@ impl CodergenBackend for AgentCliBackend {
.rev()
.collect()
};
let stderr_tail = tail(&result.stderr, 500);
let stdout_tail = tail(&result.stdout, 500);
let stderr_tail = tail(&stderr, 500);
let stdout_tail = tail(&stdout, 500);
let detail = match (stderr_tail.is_empty(), stdout_tail.is_empty()) {
(false, false) => format!("{stderr_tail}\nstdout: {stdout_tail}"),
(false, true) => stderr_tail,
@ -681,12 +763,14 @@ impl CodergenBackend for AgentCliBackend {
};
return Err(Error::handler(format!(
"CLI command exited with code {}: {detail}",
result.display_exit_code(),
result
.exit_code
.map_or_else(|| "<unknown>".to_string(), |c| c.to_string()),
)));
}
// 4. Parse the CLI output
let parsed = parse_cli_response(provider, &result.stdout)
let parsed = parse_cli_response(provider, &stdout)
.ok_or_else(|| Error::handler("Failed to parse CLI output".to_string()))?;
// 5. Detect changed files
@ -789,17 +873,32 @@ impl CodergenBackend for BackendRouter {
emitter: &Arc<Emitter>,
sandbox: &Arc<dyn Sandbox>,
tool_hooks: Option<Arc<dyn fabro_agent::ToolHookCallback>>,
cancel_token: CancellationToken,
) -> Result<CodergenResult, Error> {
if self.should_use_cli(node) {
self.cli_backend
.run(
node, prompt, context, thread_id, emitter, sandbox, tool_hooks,
node,
prompt,
context,
thread_id,
emitter,
sandbox,
tool_hooks,
cancel_token,
)
.await
} else {
self.api_backend
.run(
node, prompt, context, thread_id, emitter, sandbox, tool_hooks,
node,
prompt,
context,
thread_id,
emitter,
sandbox,
tool_hooks,
cancel_token,
)
.await
}
@ -820,6 +919,7 @@ impl CodergenBackend for BackendRouter {
mod tests {
use std::path::Path;
use fabro_agent::sandbox::ExecResult;
use fabro_graphviz::graph::AttrValue;
use super::*;
@ -999,7 +1099,14 @@ mod tests {
));
let emitter = Arc::new(Emitter::default());
let result = ensure_cli(AgentCli::Claude, Provider::Anthropic, &sandbox, &emitter).await;
let result = ensure_cli(
AgentCli::Claude,
Provider::Anthropic,
&sandbox,
&emitter,
&CancellationToken::new(),
)
.await;
assert!(result.is_ok());
let commands = commands.lock().unwrap();
@ -1020,7 +1127,14 @@ mod tests {
));
let emitter = Arc::new(Emitter::default());
let result = ensure_cli(AgentCli::Claude, Provider::Anthropic, &sandbox, &emitter).await;
let result = ensure_cli(
AgentCli::Claude,
Provider::Anthropic,
&sandbox,
&emitter,
&CancellationToken::new(),
)
.await;
assert!(result.is_ok());
let commands = commands.lock().unwrap();
@ -1045,7 +1159,14 @@ mod tests {
move |event| events.lock().unwrap().push(event.clone())
});
let result = ensure_cli(AgentCli::Claude, Provider::Anthropic, &sandbox, &emitter).await;
let result = ensure_cli(
AgentCli::Claude,
Provider::Anthropic,
&sandbox,
&emitter,
&CancellationToken::new(),
)
.await;
assert!(result.is_err());
let error = result.unwrap_err().to_string();
assert!(error.contains("install exited with code 1"));
@ -1276,6 +1397,7 @@ mod tests {
_emitter: &Arc<Emitter>,
_sandbox: &Arc<dyn Sandbox>,
_tool_hooks: Option<Arc<dyn fabro_agent::ToolHookCallback>>,
_cancel_token: CancellationToken,
) -> Result<CodergenResult, Error> {
Ok(CodergenResult::Text {
text: "stub".to_string(),
@ -1285,4 +1407,233 @@ mod tests {
})
}
}
/// Sandbox stub whose `exec_command_streaming` returns a configurable
/// `CommandTermination` so we can exercise the cancel/timeout paths in
/// `AgentCliBackend::run` without spawning real processes.
struct StreamingCliMock {
commands: Arc<Mutex<Vec<String>>>,
termination: CommandTermination,
exit_code: Option<i32>,
}
#[async_trait]
impl Sandbox for StreamingCliMock {
async fn read_file(
&self,
_path: &str,
_offset: Option<usize>,
_limit: Option<usize>,
) -> fabro_sandbox::Result<String> {
Ok(String::new())
}
async fn write_file(&self, _path: &str, _content: &str) -> fabro_sandbox::Result<()> {
Ok(())
}
async fn delete_file(&self, _path: &str) -> fabro_sandbox::Result<()> {
Ok(())
}
async fn file_exists(&self, _path: &str) -> fabro_sandbox::Result<bool> {
Ok(false)
}
async fn list_directory(
&self,
_path: &str,
_depth: Option<usize>,
) -> fabro_sandbox::Result<Vec<fabro_agent::sandbox::DirEntry>> {
Ok(vec![])
}
async fn exec_command(
&self,
command: &str,
_timeout_ms: u64,
_working_dir: Option<&str>,
_env_vars: Option<&std::collections::HashMap<String, String>>,
_cancel_token: Option<CancellationToken>,
) -> fabro_sandbox::Result<ExecResult> {
self.commands.lock().unwrap().push(command.to_string());
// Default: success for git/version/cat/rm/ls.
if command.contains("--version") {
return Ok(ok_result());
}
Ok(ExecResult {
stdout: String::new(),
stderr: String::new(),
exit_code: Some(0),
termination: CommandTermination::Exited,
duration_ms: 1,
})
}
async fn exec_command_streaming(
&self,
command: &str,
_timeout_ms: Option<u64>,
_working_dir: Option<&str>,
_env_vars: Option<&std::collections::HashMap<String, String>>,
_cancel_token: Option<CancellationToken>,
_output_callback: fabro_agent::CommandOutputCallback,
) -> fabro_sandbox::Result<fabro_sandbox::ExecStreamingResult> {
self.commands.lock().unwrap().push(command.to_string());
Ok(fabro_sandbox::ExecStreamingResult {
result: ExecResult {
stdout: String::new(),
stderr: String::new(),
exit_code: self.exit_code,
termination: self.termination,
duration_ms: 5,
},
streams_separated: true,
live_streaming: true,
})
}
async fn grep(
&self,
_pattern: &str,
_path: &str,
_options: &fabro_agent::sandbox::GrepOptions,
) -> fabro_sandbox::Result<Vec<String>> {
Ok(vec![])
}
async fn glob(
&self,
_pattern: &str,
_path: Option<&str>,
) -> fabro_sandbox::Result<Vec<String>> {
Ok(vec![])
}
async fn download_file_to_local(&self, _: &str, _: &Path) -> fabro_sandbox::Result<()> {
Ok(())
}
async fn upload_file_from_local(&self, _: &Path, _: &str) -> fabro_sandbox::Result<()> {
Ok(())
}
async fn initialize(&self) -> fabro_sandbox::Result<()> {
Ok(())
}
async fn cleanup(&self) -> fabro_sandbox::Result<()> {
Ok(())
}
fn working_directory(&self) -> &str {
"/workspace"
}
fn platform(&self) -> &str {
"linux"
}
fn os_version(&self) -> String {
"Ubuntu 22.04".into()
}
async fn set_autostop_interval(&self, _minutes: i32) -> fabro_sandbox::Result<()> {
Ok(())
}
}
fn collect_events(emitter: &Arc<Emitter>) -> Arc<Mutex<Vec<fabro_types::RunEvent>>> {
let events = Arc::new(Mutex::new(Vec::new()));
let events_clone = Arc::clone(&events);
emitter.on_event(move |event| events_clone.lock().unwrap().push(event.clone()));
events
}
#[tokio::test]
async fn agent_cli_backend_run_emits_cancelled_event_and_returns_cancelled() {
let commands = Arc::new(Mutex::new(Vec::new()));
let sandbox: Arc<dyn Sandbox> = Arc::new(StreamingCliMock {
commands: Arc::clone(&commands),
termination: CommandTermination::Cancelled,
exit_code: None,
});
let backend = AgentCliBackend::new_from_env("claude-opus-4-6".into(), Provider::Anthropic);
let node = Node::new("step");
let context = Context::new();
let emitter = Arc::new(Emitter::default());
let events = collect_events(&emitter);
let result = backend
.run(
&node,
"Do something",
&context,
None,
&emitter,
&sandbox,
None,
CancellationToken::new(),
)
.await;
let Err(err) = result else {
panic!("cancelled streaming should bubble Error::Cancelled");
};
assert!(matches!(err, Error::Cancelled));
let events = events.lock().unwrap();
let names: Vec<String> = events
.iter()
.map(|e| e.body.event_name().to_string())
.collect();
assert!(
names.iter().any(|n| n == "agent.cli.cancelled"),
"expected agent.cli.cancelled, got events: {names:?}"
);
assert!(
!names.iter().any(|n| n == "agent.cli.completed"),
"should not emit agent.cli.completed on cancellation"
);
// Cleanup `rm -f` ran.
let cmds = commands.lock().unwrap();
assert!(
cmds.iter().any(|c| c.starts_with("rm -f /tmp/fabro_cli_")),
"expected temp cleanup, got commands: {cmds:?}"
);
}
#[tokio::test]
async fn agent_cli_backend_run_emits_timed_out_event_and_returns_handler_error() {
let commands = Arc::new(Mutex::new(Vec::new()));
let sandbox: Arc<dyn Sandbox> = Arc::new(StreamingCliMock {
commands: Arc::clone(&commands),
termination: CommandTermination::TimedOut,
exit_code: None,
});
let backend = AgentCliBackend::new_from_env("claude-opus-4-6".into(), Provider::Anthropic);
let node = Node::new("step");
let context = Context::new();
let emitter = Arc::new(Emitter::default());
let events = collect_events(&emitter);
let result = backend
.run(
&node,
"Do something slow",
&context,
None,
&emitter,
&sandbox,
None,
CancellationToken::new(),
)
.await;
let Err(err) = result else {
panic!("timeout streaming should produce a handler error");
};
assert!(
matches!(err, Error::Handler { .. }),
"expected handler error on timeout, got {err:?}"
);
let events = events.lock().unwrap();
let names: Vec<String> = events
.iter()
.map(|e| e.body.event_name().to_string())
.collect();
assert!(
names.iter().any(|n| n == "agent.cli.timed_out"),
"expected agent.cli.timed_out, got events: {names:?}"
);
assert!(
!names.iter().any(|n| n == "agent.cli.completed"),
"should not emit agent.cli.completed on timeout"
);
}
}

View file

@ -1,7 +1,6 @@
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Duration;
use async_trait::async_trait;
@ -197,13 +196,12 @@ impl Handler for SubWorkflowHandler {
let child_logs = run_dir.join(format!("stages/{}@{visit}/child", node.id));
let _ = fs::create_dir_all(&child_logs).await;
let cancel_token = Arc::new(AtomicBool::new(false));
let child_cancel = Arc::clone(&cancel_token);
let child_run_token = services.run.cancel_token().child_token();
let child_run_options = RunOptions {
settings: WorkflowSettings::default(),
run_dir: child_logs,
cancel_token: Some(cancel_token),
cancel_token: child_run_token.clone(),
// Child workflows are part of the parent run's event stream.
run_id: services.run.emitter.run_id(),
labels: HashMap::new(),
@ -246,11 +244,14 @@ impl Handler for SubWorkflowHandler {
.map_err(|err| Error::engine(err.to_string()))?;
let artifact_store = ArtifactStore::new(object_store, "artifacts");
// Spawn child engine
// Spawn child engine. Child runs receive a derived cancel token from
// the parent run; parent cancellation propagates parent-to-child via
// `child_token()`, but child cancellation does not cancel the parent.
let child_run_token_for_services = child_run_token.clone();
let mut child_handle = tokio::spawn(async move {
let child_run = parent_run
.with_run_store(run_store.into())
.with_cancel_requested(None);
.with_cancel_token(child_run_token_for_services);
let initialized = Initialized {
graph: child_graph,
source: String::new(),
@ -319,7 +320,7 @@ impl Handler for SubWorkflowHandler {
if !stop_condition.is_empty() {
let dummy_outcome = Outcome::success();
if evaluate_condition(stop_condition, &dummy_outcome, context) {
child_cancel.store(true, Ordering::Relaxed);
child_run_token.cancel();
// Give child a moment to wind down
let _ = timeout(
Duration::from_millis(100),
@ -337,7 +338,7 @@ impl Handler for SubWorkflowHandler {
}
// Max cycles exceeded — cancel child
child_cancel.store(true, Ordering::Relaxed);
child_run_token.cancel();
let _ = timeout(Duration::from_millis(100), &mut child_handle).await;
Ok(Outcome::fail_classify(format!(

View file

@ -23,7 +23,6 @@ use fabro_interview::Interviewer;
use crate::context::Context;
use crate::error::Error;
use crate::outcome::{Outcome, OutcomeExt};
pub(crate) use crate::services::sandbox_cancel_token;
pub use crate::services::{EngineServices, RunServices};
/// The handler interface for node execution.

View file

@ -463,6 +463,9 @@ impl Handler for ParallelHandler {
Ok(Ok(result)) => {
results.push(result);
}
Ok(Err(Error::Cancelled)) => {
return Err(Error::Cancelled);
}
Ok(Err(e)) => {
results.push(BranchResult {
id: String::new(),

View file

@ -71,8 +71,10 @@ impl Handler for PromptHandler {
working_dir,
working_dir,
provider,
&services.run.cancel_token(),
)
.await;
.await
.unwrap_or_default();
if docs.is_empty() {
None
@ -115,6 +117,7 @@ impl Handler for PromptHandler {
files_touched,
..
}) => (text, usage, files_touched),
Err(Error::Cancelled) => return Err(Error::Cancelled),
Err(e) if e.is_retryable() => {
return Err(e);
}
@ -185,6 +188,7 @@ mod tests {
use fabro_types::fixtures;
use object_store::memory::InMemory;
use tempfile::TempDir;
use tokio_util::sync::CancellationToken;
use super::*;
@ -270,6 +274,7 @@ mod tests {
_emitter: &Arc<crate::event::Emitter>,
_sandbox: &Arc<dyn Sandbox>,
_tool_hooks: Option<Arc<dyn fabro_agent::ToolHookCallback>>,
_cancel_token: CancellationToken,
) -> Result<CodergenResult, Error> {
panic!("run() should not be called for prompt handler");
}
@ -330,6 +335,7 @@ mod tests {
_emitter: &Arc<crate::event::Emitter>,
_sandbox: &Arc<dyn Sandbox>,
_tool_hooks: Option<Arc<dyn fabro_agent::ToolHookCallback>>,
_cancel_token: CancellationToken,
) -> Result<CodergenResult, Error> {
panic!("run() should not be called for prompt handler");
}
@ -387,6 +393,7 @@ mod tests {
_emitter: &Arc<crate::event::Emitter>,
_sandbox: &Arc<dyn fabro_agent::Sandbox>,
_tool_hooks: Option<Arc<dyn fabro_agent::ToolHookCallback>>,
_cancel_token: CancellationToken,
) -> Result<CodergenResult, Error> {
panic!("run() should not be called for prompt handler");
}

View file

@ -601,7 +601,7 @@ mod tests {
Arc::new(RunOptions {
settings: WorkflowSettings::default(),
run_dir: run_dir.to_path_buf(),
cancel_token: None,
cancel_token: tokio_util::sync::CancellationToken::new(),
run_id: fixtures::RUN_1,
labels: HashMap::new(),
workflow_slug: Some("metadata".to_string()),
@ -998,7 +998,7 @@ mod tests {
repo_dir.path().to_path_buf(),
)),
None,
None,
tokio_util::sync::CancellationToken::new(),
fabro_model::Provider::Anthropic,
Arc::new(fabro_auth::EnvCredentialSource::new()),
Arc::new(SandboxGitRuntime::new()),

View file

@ -250,6 +250,8 @@ fn replay_event_for_fork_projection(body: &EventBody) -> bool {
| EventBody::InterviewInterrupted(_)
| EventBody::AgentSessionStarted(_)
| EventBody::AgentCliStarted(_)
| EventBody::AgentCliCancelled(_)
| EventBody::AgentCliTimedOut(_)
| EventBody::CommandStarted(_)
| EventBody::CommandCompleted(_)
| EventBody::ParallelCompleted(_)

View file

@ -1,6 +1,5 @@
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
@ -30,6 +29,7 @@ use fabro_types::settings::run::{
use fabro_vault::Vault;
use tokio::runtime::Handle;
use tokio::sync::RwLock as AsyncRwLock;
use tokio_util::sync::CancellationToken;
use crate::ManifestPath;
use crate::artifact_upload::ArtifactSink;
@ -54,7 +54,7 @@ use crate::runtime_store::RunStoreHandle;
use crate::workflow_bundle::{RunDefinition, WorkflowBundle};
struct RunSession {
cancel_token: Option<Arc<AtomicBool>>,
cancel_token: CancellationToken,
emitter: Arc<Emitter>,
sandbox: SandboxSpec,
llm: LlmSpec,
@ -86,7 +86,7 @@ struct RunSession {
pub struct StartServices {
pub run_id: RunId,
pub cancel_token: Option<Arc<AtomicBool>>,
pub cancel_token: CancellationToken,
pub emitter: Arc<Emitter>,
pub interviewer: Arc<dyn Interviewer>,
pub run_store: RunStoreHandle,
@ -831,7 +831,7 @@ impl RunSession {
struct DetachedRunBootstrapGuard {
run_id: RunId,
event_sink: RunEventSink,
cancel_token: Option<Arc<AtomicBool>>,
cancel_token: CancellationToken,
active: bool,
}
@ -840,7 +840,7 @@ impl DetachedRunBootstrapGuard {
run_id: RunId,
_run_dir: &Path,
event_sink: RunEventSink,
cancel_token: Option<Arc<AtomicBool>>,
cancel_token: CancellationToken,
) -> Self {
Self {
run_id,
@ -858,10 +858,7 @@ impl DetachedRunBootstrapGuard {
impl Drop for DetachedRunBootstrapGuard {
fn drop(&mut self) {
if self.active {
let cancelled = self
.cancel_token
.as_ref()
.is_some_and(|token| token.load(Ordering::SeqCst));
let cancelled = self.cancel_token.is_cancelled();
let reason = if cancelled {
FailureReason::Cancelled
} else {
@ -891,12 +888,12 @@ const POSTRUN_CANCELLED_MESSAGE: &str = "Run cancelled before post-run finalizat
struct DetachedRunCompletionGuard {
event_sink: RunEventSink,
run_id: RunId,
cancel_token: Option<Arc<AtomicBool>>,
cancel_token: CancellationToken,
active: bool,
}
impl DetachedRunCompletionGuard {
fn arm(run_id: RunId, event_sink: RunEventSink, cancel_token: Option<Arc<AtomicBool>>) -> Self {
fn arm(run_id: RunId, event_sink: RunEventSink, cancel_token: CancellationToken) -> Self {
Self {
event_sink,
run_id,
@ -916,10 +913,7 @@ impl Drop for DetachedRunCompletionGuard {
return;
}
let cancelled = self
.cancel_token
.as_ref()
.is_some_and(|token| token.load(Ordering::SeqCst));
let cancelled = self.cancel_token.is_cancelled();
let reason = if cancelled {
FailureReason::Cancelled
} else {
@ -1108,7 +1102,7 @@ mod tests {
) -> StartServices {
StartServices {
run_id: fixtures::RUN_1,
cancel_token: None,
cancel_token: CancellationToken::new(),
emitter,
interviewer: Arc::new(fabro_interview::AutoApproveInterviewer::engine()),
run_store: store.open_run(&fixtures::RUN_1).await.unwrap().into(),

View file

@ -243,9 +243,7 @@ pub async fn execute(init: Initialized) -> Executed {
let mut builder = ExecutorBuilder::new(handler as Arc<dyn NodeHandler<WorkflowGraph>>)
.lifecycle(Box::new(lifecycle));
if let Some(ref cancel) = run_options.cancel_token {
builder = builder.cancel_token(cancel.clone());
}
builder = builder.cancel_token(run_options.cancel_token.clone());
if let Some(token) = stall_token.clone() {
builder = builder.stall_token(token);
}

View file

@ -7,7 +7,7 @@
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
use std::sync::atomic::{AtomicU32, Ordering};
use std::time::Duration;
use async_trait::async_trait;
@ -91,7 +91,7 @@ fn test_emitter_arc(label: &str) -> Arc<Emitter> {
fn test_run_options(run_dir: &Path, run_id: &str) -> RunOptions {
RunOptions {
run_dir: run_dir.to_path_buf(),
cancel_token: None,
cancel_token: tokio_util::sync::CancellationToken::new(),
run_id: test_run_id(run_id),
settings: WorkflowSettings::default(),
git: None,
@ -864,16 +864,16 @@ async fn execute_cancelled_mid_run() {
g.edges.push(Edge::new("start", "work"));
g.edges.push(Edge::new("work", "exit"));
let cancel_token = Arc::new(AtomicBool::new(false));
let cancel_token_clone = Arc::clone(&cancel_token);
let cancel_token = tokio_util::sync::CancellationToken::new();
let cancel_token_clone = cancel_token.clone();
let mut registry = make_registry();
registry.register("slow", Box::new(SlowHandler { sleep_ms: 200 }));
let mut run_options = test_run_options(dir.path(), "test-run");
run_options.cancel_token = Some(cancel_token);
run_options.cancel_token = cancel_token;
tokio::spawn(async move {
tokio::time::sleep(Duration::from_millis(50)).await;
cancel_token_clone.store(true, Ordering::Relaxed);
cancel_token_clone.cancel();
});
let result = run_graph(
@ -901,16 +901,16 @@ async fn execute_cancelled_mid_run_persists_cancelled_status() {
g.edges.push(Edge::new("start", "work"));
g.edges.push(Edge::new("work", "exit"));
let cancel_token = Arc::new(AtomicBool::new(false));
let cancel_token_clone = Arc::clone(&cancel_token);
let cancel_token = tokio_util::sync::CancellationToken::new();
let cancel_token_clone = cancel_token.clone();
let mut registry = make_registry();
registry.register("slow", Box::new(SlowHandler { sleep_ms: 200 }));
let mut run_options = test_run_options(dir.path(), "test-run");
run_options.cancel_token = Some(cancel_token);
run_options.cancel_token = cancel_token;
tokio::spawn(async move {
tokio::time::sleep(Duration::from_millis(50)).await;
cancel_token_clone.store(true, Ordering::Relaxed);
cancel_token_clone.cancel();
});
let executed = execute_test_run_with_options(run_options, g, Some(Arc::new(registry))).await;

View file

@ -621,7 +621,7 @@ mod tests {
RunOptions {
settings: WorkflowSettings::default(),
run_dir: run_dir.to_path_buf(),
cancel_token: None,
cancel_token: tokio_util::sync::CancellationToken::new(),
run_id: test_run_id(),
labels: HashMap::new(),
workflow_slug: None,
@ -816,7 +816,7 @@ mod tests {
emitter,
sandbox,
None,
None,
tokio_util::sync::CancellationToken::new(),
fabro_model::Provider::Anthropic,
Arc::new(fabro_auth::EnvCredentialSource::new()),
Arc::new(SandboxGitRuntime::new()),
@ -842,7 +842,7 @@ mod tests {
std::env::current_dir().unwrap(),
)),
None,
None,
tokio_util::sync::CancellationToken::new(),
fabro_model::Provider::Anthropic,
Arc::new(fabro_auth::EnvCredentialSource::new()),
Arc::new(SandboxGitRuntime::new()),

View file

@ -30,7 +30,7 @@ use crate::error::Error;
use crate::event::{Emitter, Event, RunNoticeLevel};
use crate::git::RUN_BRANCH_PREFIX;
use crate::handler::llm::{AgentApiBackend, AgentCliBackend, BackendRouter};
use crate::handler::{HandlerRegistry, default_registry, sandbox_cancel_token};
use crate::handler::{HandlerRegistry, default_registry};
use crate::run_metadata::{
RunMetadataRuntime, build_metadata_writer, metadata_branch_name, mint_token,
};
@ -637,23 +637,21 @@ pub async fn initialize(
index,
});
let cmd_start = Instant::now();
let cancel_token = sandbox_cancel_token(options.run_options.cancel_token.clone());
let cancel_token = options.run_options.cancel_token.child_token();
let result = sandbox
.exec_command(
command,
options.lifecycle.setup_command_timeout_ms,
None,
None,
cancel_token.clone(),
Some(cancel_token.clone()),
)
.await
.map_err(|e| Error::engine_with_source("Setup command failed", &e))?;
if let Some(token) = &cancel_token {
if token.is_cancelled() {
return Err(Error::Cancelled);
}
token.cancel();
if options.run_options.cancel_token.is_cancelled() {
return Err(Error::Cancelled);
}
cancel_token.cancel();
let duration_ms = crate::millis_u64(cmd_start.elapsed());
if !result.is_success() {
let exit_code = result.display_exit_code();
@ -753,7 +751,6 @@ pub async fn initialize(
mod tests {
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::AtomicBool;
use std::time::Duration;
use fabro_auth::{AuthCredential, AuthDetails};
@ -846,7 +843,7 @@ mod tests {
RunOptions {
settings: WorkflowSettings::default(),
run_dir: run_dir.to_path_buf(),
cancel_token: None,
cancel_token: tokio_util::sync::CancellationToken::new(),
run_id: test_run_id(),
labels: HashMap::new(),
workflow_slug: None,
@ -1257,9 +1254,10 @@ mod tests {
std::fs::create_dir_all(&run_dir).unwrap();
let (graph, source) = simple_graph();
let persisted = test_persisted(graph, source, &run_dir);
let cancel_token = Arc::new(AtomicBool::new(true));
let cancel_token = tokio_util::sync::CancellationToken::new();
cancel_token.cancel();
let mut run_options = test_settings(&run_dir);
run_options.cancel_token = Some(cancel_token);
run_options.cancel_token = cancel_token;
let result = initialize(persisted, InitOptions {
run_id: test_run_id(),
@ -1318,9 +1316,10 @@ mod tests {
std::fs::create_dir_all(&run_dir).unwrap();
let (graph, source) = simple_graph();
let persisted = test_persisted(graph, source, &run_dir);
let cancel_token = Arc::new(AtomicBool::new(true));
let cancel_token = tokio_util::sync::CancellationToken::new();
cancel_token.cancel();
let mut run_options = test_settings(&run_dir);
run_options.cancel_token = Some(cancel_token);
run_options.cancel_token = cancel_token;
let result = initialize(persisted, InitOptions {
run_id: test_run_id(),

View file

@ -304,7 +304,7 @@ mod tests {
RunOptions {
settings: WorkflowSettings::default(),
run_dir: run_dir.to_path_buf(),
cancel_token: None,
cancel_token: tokio_util::sync::CancellationToken::new(),
run_id: test_run_id(),
labels: HashMap::new(),
workflow_slug: None,
@ -340,7 +340,7 @@ mod tests {
Arc::clone(&emitter),
Arc::clone(&sandbox),
None,
None,
tokio_util::sync::CancellationToken::new(),
fabro_llm::Provider::Anthropic,
test_llm_source(),
Arc::new(crate::sandbox_git_runtime::SandboxGitRuntime::new()),
@ -395,7 +395,7 @@ mod tests {
std::env::current_dir().unwrap(),
)),
None,
None,
tokio_util::sync::CancellationToken::new(),
fabro_llm::Provider::Anthropic,
test_llm_source(),
Arc::new(crate::sandbox_git_runtime::SandboxGitRuntime::new()),

View file

@ -642,7 +642,7 @@ mod tests {
RunOptions {
settings: WorkflowSettings::default(),
run_dir: tempfile::tempdir().unwrap().path().to_path_buf(),
cancel_token: None,
cancel_token: tokio_util::sync::CancellationToken::new(),
run_id: fabro_types::fixtures::RUN_1,
labels: HashMap::new(),
workflow_slug: Some("metadata".to_string()),

View file

@ -1,10 +1,9 @@
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::Arc;
use std::sync::atomic::AtomicBool;
use fabro_types::settings::run::RunMode;
use fabro_types::{ForkSourceRef, GitContext, RunId, WorkflowSettings};
use tokio_util::sync::CancellationToken;
use crate::git::{GitAuthor, git_author_from_settings};
@ -21,7 +20,10 @@ pub struct GitCheckpointOptions {
pub struct RunOptions {
pub settings: WorkflowSettings,
pub run_dir: PathBuf,
pub cancel_token: Option<Arc<AtomicBool>>,
/// Cancellation token for this run. Cancelling this token cancels the
/// run and propagates to handlers, sandbox commands, and child runs.
/// Default constructors should use `CancellationToken::new()`.
pub cancel_token: CancellationToken,
/// Unique identifier for this workflow run.
pub run_id: RunId,
/// User-defined key-value labels for this run.

View file

@ -2,7 +2,7 @@ use std::collections::HashMap;
#[cfg(test)]
use std::path::PathBuf;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
#[cfg(test)]
use std::time::Duration;
use fabro_agent::Sandbox;
@ -11,7 +11,6 @@ use fabro_auth::CredentialSource;
use fabro_auth::ResolvedCredentials;
use fabro_hooks::{HookContext, HookDecision, HookRunner};
use fabro_model::Provider;
use tokio::time;
use tokio_util::sync::CancellationToken;
use crate::ManifestPath;
@ -24,13 +23,20 @@ use crate::sandbox_git_runtime::SandboxGitRuntime;
use crate::workflow_bundle::WorkflowBundle;
/// Services shared across workflow phases.
///
/// Production construction is expected to happen from pipeline initialization
/// with the run's root cancellation token. Use
/// [`RunServices::with_cancel_token`] only with the same root token or a
/// `child_token()` derived from it. The token semantically means "cancel this
/// run or child run," not a generic shutdown signal — dropping a `RunServices`
/// does NOT count as cancellation.
#[derive(Clone)]
pub struct RunServices {
pub run_store: RunStoreHandle,
pub emitter: Arc<Emitter>,
pub sandbox: Arc<dyn Sandbox>,
pub hook_runner: Option<Arc<HookRunner>>,
pub cancel_requested: Option<Arc<AtomicBool>>,
pub(crate) cancel_token: CancellationToken,
pub provider: Provider,
pub llm_source: Arc<dyn CredentialSource>,
pub(crate) sandbox_git: Arc<SandboxGitRuntime>,
@ -45,7 +51,7 @@ impl RunServices {
emitter: Arc<Emitter>,
sandbox: Arc<dyn Sandbox>,
hook_runner: Option<Arc<HookRunner>>,
cancel_requested: Option<Arc<AtomicBool>>,
cancel_token: CancellationToken,
provider: Provider,
llm_source: Arc<dyn CredentialSource>,
sandbox_git: Arc<SandboxGitRuntime>,
@ -57,7 +63,7 @@ impl RunServices {
emitter,
sandbox,
hook_runner,
cancel_requested,
cancel_token,
provider,
llm_source,
sandbox_git,
@ -66,10 +72,11 @@ impl RunServices {
})
}
/// Bridge the core executor's atomic cancel flag to sandbox command
/// cancellation.
pub fn sandbox_cancel_token(&self) -> Option<CancellationToken> {
sandbox_cancel_token(self.cancel_requested.clone())
/// The run-level cancellation token. Cancel this to terminate the run.
/// Derive child tokens via `cancel_token().child_token()` for sandbox
/// command invocations.
pub fn cancel_token(&self) -> CancellationToken {
self.cancel_token.clone()
}
/// Run lifecycle hooks and return the merged decision.
@ -107,13 +114,15 @@ impl RunServices {
})
}
/// Replace the cancellation token. Use only with the same root token or
/// a child derived from it via `child_token()`.
#[must_use]
pub fn with_cancel_requested(
pub(crate) fn with_cancel_token(
self: &Arc<Self>,
cancel_requested: Option<Arc<AtomicBool>>,
cancel_token: CancellationToken,
) -> Arc<Self> {
Arc::new(Self {
cancel_requested,
cancel_token,
..self.as_ref().clone()
})
}
@ -209,7 +218,7 @@ impl EngineServices {
std::env::current_dir().unwrap_or_else(|_| PathBuf::from(".")),
)),
None,
None,
CancellationToken::new(),
Provider::Anthropic,
Arc::new(StubCredentialSource),
Arc::new(SandboxGitRuntime::new()),
@ -227,34 +236,6 @@ impl EngineServices {
}
}
pub(crate) fn sandbox_cancel_token(
cancel_requested: Option<Arc<AtomicBool>>,
) -> Option<CancellationToken> {
let cancel_requested = cancel_requested?;
let token = CancellationToken::new();
if cancel_requested.load(Ordering::Relaxed) {
token.cancel();
return Some(token);
}
let token_clone = token.clone();
tokio::spawn(async move {
loop {
if token_clone.is_cancelled() {
return;
}
if cancel_requested.load(Ordering::Relaxed) {
token_clone.cancel();
return;
}
time::sleep(Duration::from_millis(10)).await;
}
});
Some(token)
}
#[cfg(test)]
mod tests {
use super::EngineServices;

View file

@ -41,6 +41,7 @@ use fabro_workflow::records::Checkpoint;
use fabro_workflow::run_options::{GitCheckpointOptions, RunOptions};
use fabro_workflow::test_support::{WorkflowRunner, test_store_dir};
use object_store::local::LocalFileSystem;
use tokio_util::sync::CancellationToken;
use ulid::Ulid;
fn test_run_id(label: &str) -> RunId {
@ -249,7 +250,7 @@ async fn daytona_exec_command_cancelled() {
let env = create_env_with_github_app(Some(creds)).await;
env.initialize().await.unwrap();
let token = tokio_util::sync::CancellationToken::new();
let token = CancellationToken::new();
let token_clone = token.clone();
// Cancel the token shortly after starting
@ -513,7 +514,7 @@ async fn daytona_pipeline_artifact_offload_and_sync() {
let run_options = RunOptions {
settings: WorkflowSettings::default(),
run_dir: dir.path().to_path_buf(),
cancel_token: None,
cancel_token: CancellationToken::new(),
run_id: test_run_id("test-run"),
labels: std::collections::HashMap::new(),
workflow_slug: None,
@ -698,7 +699,7 @@ async fn daytona_git_checkpoint_remote_emits_events() {
let run_options = RunOptions {
settings: WorkflowSettings::default(),
run_dir: dir.path().to_path_buf(),
cancel_token: None,
cancel_token: CancellationToken::new(),
run_id: test_run_id("git-cp-test"),
labels: std::collections::HashMap::new(),
workflow_slug: None,
@ -871,7 +872,7 @@ async fn daytona_parallel_git_branching_e2e() {
let run_options = RunOptions {
settings: WorkflowSettings::default(),
run_dir: run_tmp.path().to_path_buf(),
cancel_token: None,
cancel_token: CancellationToken::new(),
run_id,
labels: std::collections::HashMap::new(),
workflow_slug: None,
@ -1078,6 +1079,7 @@ async fn run_daytona_cli_test(provider: Provider, model: &str, install_command:
&emitter,
&env,
None,
CancellationToken::new(),
)
.await;
@ -1209,7 +1211,7 @@ async fn daytona_git_checkpoint_with_shadow_branch() {
let run_options = RunOptions {
settings: WorkflowSettings::default(),
run_dir: dir.path().to_path_buf(),
cancel_token: None,
cancel_token: CancellationToken::new(),
run_id,
labels: std::collections::HashMap::new(),
workflow_slug: None,
@ -1366,7 +1368,7 @@ async fn daytona_asset_collection() {
..WorkflowSettings::default()
},
run_dir: dir.path().to_path_buf(),
cancel_token: None,
cancel_token: CancellationToken::new(),
run_id: test_run_id("artifact-test-daytona"),
labels: std::collections::HashMap::new(),
workflow_slug: None,
@ -1634,7 +1636,7 @@ async fn daytona_git_push_run_branch_to_origin() {
let run_options = RunOptions {
settings: WorkflowSettings::default(),
run_dir: dir.path().to_path_buf(),
cancel_token: None,
cancel_token: CancellationToken::new(),
run_id,
labels: std::collections::HashMap::new(),
workflow_slug: None,

View file

@ -21,6 +21,7 @@ use fabro_workflow::handler::exit::ExitHandler;
use fabro_workflow::handler::start::StartHandler;
use fabro_workflow::run_options::{GitCheckpointOptions, RunOptions};
use fabro_workflow::test_support::run_graph;
use tokio_util::sync::CancellationToken;
fn assert_success(output: &Output, context: &str) {
assert!(
@ -154,7 +155,7 @@ fn make_registry() -> HandlerRegistry {
fn test_run_options(run_dir: &Path) -> RunOptions {
RunOptions {
run_dir: run_dir.to_path_buf(),
cancel_token: None,
cancel_token: CancellationToken::new(),
run_id: fixtures::RUN_2,
settings: WorkflowSettings::default(),
git: None,

File diff suppressed because it is too large Load diff