mirror of
https://github.com/fabro-sh/fabro.git
synced 2026-10-09 03:20:56 +00:00
Sync artifact files to remote execution environments
After offloading large context values to local disk, check whether each artifact file is accessible in the execution environment. When it isn't (Docker/Daytona), read the local file and upload it via write_file, rewriting the pointer to the remote path. Local envs skip the sync since file_exists returns true for local paths. Also switches the pointer prefix from artifact:// to file://. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
parent
8e6acde1b3
commit
648418a1b0
3 changed files with 369 additions and 8 deletions
|
|
@ -6,6 +6,8 @@ use chrono::{DateTime, Utc};
|
|||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
|
||||
use agent::ExecutionEnvironment;
|
||||
|
||||
use crate::error::{AttractorError, Result};
|
||||
|
||||
/// Threshold above which artifacts are stored on disk instead of in memory (100KB).
|
||||
|
|
@ -194,12 +196,12 @@ impl ArtifactStore {
|
|||
}
|
||||
|
||||
/// Prefix used to identify artifact pointer strings in context values.
|
||||
const ARTIFACT_POINTER_PREFIX: &str = "artifact://";
|
||||
const ARTIFACT_POINTER_PREFIX: &str = "file://";
|
||||
|
||||
/// Offload context values exceeding the file-backing threshold into the artifact store.
|
||||
///
|
||||
/// For each entry in `updates` whose serialized JSON exceeds `FILE_BACKING_THRESHOLD`,
|
||||
/// the value is stored in `store` and replaced with an `"artifact://{path}"` pointer.
|
||||
/// the value is stored in `store` and replaced with a `"file://{path}"` pointer.
|
||||
/// Small values are left untouched.
|
||||
///
|
||||
/// # Errors
|
||||
|
|
@ -225,7 +227,7 @@ pub fn offload_large_values(
|
|||
|
||||
/// Extract the file path from an artifact pointer value.
|
||||
///
|
||||
/// Returns `Some(path)` if the value is a string starting with `"artifact://"`,
|
||||
/// Returns `Some(path)` if the value is a string starting with `"file://"`,
|
||||
/// `None` otherwise.
|
||||
#[must_use]
|
||||
pub fn artifact_path(value: &Value) -> Option<&str> {
|
||||
|
|
@ -234,7 +236,7 @@ pub fn artifact_path(value: &Value) -> Option<&str> {
|
|||
.and_then(|s| s.strip_prefix(ARTIFACT_POINTER_PREFIX))
|
||||
}
|
||||
|
||||
/// Returns `true` if `path` looks like an artifact pointer path (starts with `"artifact://"`).
|
||||
/// Returns `true` if `path` looks like an artifact pointer path (starts with `"file://"`).
|
||||
#[must_use]
|
||||
pub fn is_artifact_pointer(value: &Value) -> bool {
|
||||
artifact_path(value).is_some()
|
||||
|
|
@ -242,12 +244,65 @@ pub fn is_artifact_pointer(value: &Value) -> bool {
|
|||
|
||||
/// Resolve an artifact pointer to the base name displayed in preamble rendering.
|
||||
///
|
||||
/// Given `"artifact:///tmp/logs/artifacts/response.plan.json"`, returns `"See: /tmp/logs/artifacts/response.plan.json"`.
|
||||
/// Given `"file:///tmp/logs/artifacts/response.plan.json"`, returns `"See: /tmp/logs/artifacts/response.plan.json"`.
|
||||
#[must_use]
|
||||
pub fn format_artifact_reference(path: &str) -> String {
|
||||
format!("See: {path}")
|
||||
}
|
||||
|
||||
/// Sync artifact files to a remote execution environment.
|
||||
///
|
||||
/// For each `file://` pointer in `updates`, checks whether the file is accessible
|
||||
/// in `env`. If not, reads the local file and uploads it via `env.write_file`,
|
||||
/// placing it at `{working_directory}/.attractor/artifacts/{filename}`. The pointer
|
||||
/// is rewritten to reference the remote path.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns an error if reading a local artifact or writing to the remote env fails.
|
||||
pub async fn sync_artifacts_to_env(
|
||||
updates: &mut HashMap<String, Value>,
|
||||
env: &dyn ExecutionEnvironment,
|
||||
) -> Result<()> {
|
||||
for value in updates.values_mut() {
|
||||
let local_path = match artifact_path(value) {
|
||||
Some(p) => p.to_string(),
|
||||
None => continue,
|
||||
};
|
||||
|
||||
match env.file_exists(&local_path).await {
|
||||
Ok(true) => continue,
|
||||
Ok(false) => {}
|
||||
Err(e) => {
|
||||
return Err(AttractorError::Engine(format!(
|
||||
"failed to check artifact existence: {e}"
|
||||
)));
|
||||
}
|
||||
}
|
||||
|
||||
let content = std::fs::read_to_string(&local_path).map_err(|e| {
|
||||
AttractorError::Engine(format!("failed to read local artifact {local_path}: {e}"))
|
||||
})?;
|
||||
|
||||
let filename = std::path::Path::new(&local_path)
|
||||
.file_name()
|
||||
.and_then(|f| f.to_str())
|
||||
.unwrap_or("artifact.json");
|
||||
|
||||
let remote_path = format!(
|
||||
"{}/.attractor/artifacts/{filename}",
|
||||
env.working_directory()
|
||||
);
|
||||
|
||||
env.write_file(&remote_path, &content).await.map_err(|e| {
|
||||
AttractorError::Engine(format!("failed to write artifact to remote env: {e}"))
|
||||
})?;
|
||||
|
||||
*value = Value::String(format!("{ARTIFACT_POINTER_PREFIX}{remote_path}"));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
|
@ -415,7 +470,7 @@ mod tests {
|
|||
|
||||
#[test]
|
||||
fn artifact_path_extracts_path_from_pointer() {
|
||||
let value = serde_json::json!("artifact:///tmp/logs/artifacts/response.plan.json");
|
||||
let value = serde_json::json!("file:///tmp/logs/artifacts/response.plan.json");
|
||||
assert_eq!(
|
||||
artifact_path(&value),
|
||||
Some("/tmp/logs/artifacts/response.plan.json")
|
||||
|
|
@ -433,4 +488,153 @@ mod tests {
|
|||
let value = serde_json::json!(42);
|
||||
assert_eq!(artifact_path(&value), None);
|
||||
}
|
||||
|
||||
// --- sync_artifacts_to_env tests ---
|
||||
|
||||
use std::sync::Mutex;
|
||||
|
||||
struct TestSyncEnv {
|
||||
accessible: bool,
|
||||
written: Mutex<Vec<(String, String)>>,
|
||||
working_dir: String,
|
||||
}
|
||||
|
||||
impl TestSyncEnv {
|
||||
fn new(accessible: bool, working_dir: &str) -> Self {
|
||||
Self {
|
||||
accessible,
|
||||
written: Mutex::new(Vec::new()),
|
||||
working_dir: working_dir.to_string(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl ExecutionEnvironment for TestSyncEnv {
|
||||
async fn read_file(&self, _path: &str, _offset: Option<usize>, _limit: Option<usize>) -> std::result::Result<String, String> {
|
||||
Err("not implemented".to_string())
|
||||
}
|
||||
|
||||
async fn write_file(&self, path: &str, content: &str) -> std::result::Result<(), String> {
|
||||
self.written
|
||||
.lock()
|
||||
.unwrap()
|
||||
.push((path.to_string(), content.to_string()));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn delete_file(&self, _path: &str) -> std::result::Result<(), String> {
|
||||
Err("not implemented".to_string())
|
||||
}
|
||||
|
||||
async fn file_exists(&self, _path: &str) -> std::result::Result<bool, String> {
|
||||
Ok(self.accessible)
|
||||
}
|
||||
|
||||
async fn list_directory(&self, _path: &str, _depth: Option<usize>) -> std::result::Result<Vec<agent::DirEntry>, String> {
|
||||
Err("not implemented".to_string())
|
||||
}
|
||||
|
||||
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<tokio_util::sync::CancellationToken>,
|
||||
) -> std::result::Result<agent::ExecResult, String> {
|
||||
Err("not implemented".to_string())
|
||||
}
|
||||
|
||||
async fn grep(&self, _pattern: &str, _path: &str, _options: &agent::GrepOptions) -> std::result::Result<Vec<String>, String> {
|
||||
Err("not implemented".to_string())
|
||||
}
|
||||
|
||||
async fn glob(&self, _pattern: &str, _path: Option<&str>) -> std::result::Result<Vec<String>, String> {
|
||||
Err("not implemented".to_string())
|
||||
}
|
||||
|
||||
async fn initialize(&self) -> std::result::Result<(), String> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn cleanup(&self) -> std::result::Result<(), String> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn working_directory(&self) -> &str {
|
||||
&self.working_dir
|
||||
}
|
||||
|
||||
fn platform(&self) -> &str {
|
||||
"linux"
|
||||
}
|
||||
|
||||
fn os_version(&self) -> String {
|
||||
"Linux 5.15".to_string()
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn sync_uploads_artifact_when_not_accessible() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let artifact_file = dir.path().join("response.plan.json");
|
||||
std::fs::write(&artifact_file, r#""hello from artifact""#).unwrap();
|
||||
|
||||
let pointer = format!("file://{}", artifact_file.display());
|
||||
let mut updates = HashMap::new();
|
||||
updates.insert("response.plan".to_string(), Value::String(pointer));
|
||||
|
||||
let env = TestSyncEnv::new(false, "/workspace");
|
||||
sync_artifacts_to_env(&mut updates, &env).await.unwrap();
|
||||
|
||||
let written = env.written.lock().unwrap();
|
||||
assert_eq!(written.len(), 1);
|
||||
assert_eq!(
|
||||
written[0].0,
|
||||
"/workspace/.attractor/artifacts/response.plan.json"
|
||||
);
|
||||
assert_eq!(written[0].1, r#""hello from artifact""#);
|
||||
|
||||
let new_pointer = updates["response.plan"].as_str().unwrap();
|
||||
assert_eq!(
|
||||
new_pointer,
|
||||
"file:///workspace/.attractor/artifacts/response.plan.json"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn sync_skips_when_artifact_already_accessible() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let artifact_file = dir.path().join("data.json");
|
||||
std::fs::write(&artifact_file, "{}").unwrap();
|
||||
|
||||
let pointer = format!("file://{}", artifact_file.display());
|
||||
let mut updates = HashMap::new();
|
||||
updates.insert("key".to_string(), Value::String(pointer.clone()));
|
||||
|
||||
let env = TestSyncEnv::new(true, "/workspace");
|
||||
sync_artifacts_to_env(&mut updates, &env).await.unwrap();
|
||||
|
||||
let written = env.written.lock().unwrap();
|
||||
assert!(written.is_empty());
|
||||
assert_eq!(updates["key"].as_str().unwrap(), &pointer);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn sync_ignores_non_artifact_values() {
|
||||
let mut updates = HashMap::new();
|
||||
updates.insert("name".to_string(), serde_json::json!("Alice"));
|
||||
updates.insert("count".to_string(), serde_json::json!(42));
|
||||
updates.insert("nested".to_string(), serde_json::json!({"a": 1}));
|
||||
|
||||
let env = TestSyncEnv::new(false, "/workspace");
|
||||
sync_artifacts_to_env(&mut updates, &env).await.unwrap();
|
||||
|
||||
let written = env.written.lock().unwrap();
|
||||
assert!(written.is_empty());
|
||||
assert_eq!(updates["name"], serde_json::json!("Alice"));
|
||||
assert_eq!(updates["count"], serde_json::json!(42));
|
||||
assert_eq!(updates["nested"], serde_json::json!({"a": 1}));
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -10,7 +10,7 @@ use chrono::Utc;
|
|||
use futures::FutureExt;
|
||||
use rand::Rng;
|
||||
|
||||
use crate::artifact::{offload_large_values, ArtifactStore};
|
||||
use crate::artifact::{offload_large_values, sync_artifacts_to_env, ArtifactStore};
|
||||
use crate::checkpoint::Checkpoint;
|
||||
use crate::condition::evaluate_condition;
|
||||
use crate::context::Context;
|
||||
|
|
@ -935,6 +935,14 @@ impl PipelineEngine {
|
|||
context.append_log(format!("artifact offload failed: {e}"));
|
||||
}
|
||||
|
||||
// Sync artifact files to the execution environment (no-op for local envs)
|
||||
if let Err(e) = sync_artifacts_to_env(
|
||||
&mut outcome.context_updates,
|
||||
&*self.services.execution_env,
|
||||
).await {
|
||||
context.append_log(format!("artifact sync failed: {e}"));
|
||||
}
|
||||
|
||||
// Step 3: Record completion
|
||||
completed_nodes.push(node.id.clone());
|
||||
node_outcomes.insert(node.id.clone(), outcome.clone());
|
||||
|
|
|
|||
|
|
@ -6899,7 +6899,7 @@ async fn large_context_values_are_offloaded_to_artifact_store() {
|
|||
.expect("context should have response.big_output");
|
||||
let pointer_str = pointer_value.as_str().expect("pointer should be a string");
|
||||
assert!(
|
||||
pointer_str.starts_with("artifact://"),
|
||||
pointer_str.starts_with("file://"),
|
||||
"value should be an artifact pointer, got: {pointer_str}"
|
||||
);
|
||||
|
||||
|
|
@ -6933,4 +6933,153 @@ async fn large_context_values_are_offloaded_to_artifact_store() {
|
|||
"artifact_count should be > 0, got {artifact_count}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Artifact sync to remote execution environments
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// A mock execution environment where `file_exists` always returns false,
|
||||
/// simulating a remote container that doesn't have local artifact files.
|
||||
struct RemoteMockEnv {
|
||||
working_dir: String,
|
||||
written: std::sync::Mutex<Vec<(String, String)>>,
|
||||
}
|
||||
|
||||
impl RemoteMockEnv {
|
||||
fn new(working_dir: &str) -> Self {
|
||||
Self {
|
||||
working_dir: working_dir.to_string(),
|
||||
written: std::sync::Mutex::new(Vec::new()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl agent::ExecutionEnvironment for RemoteMockEnv {
|
||||
async fn read_file(&self, _path: &str, _offset: Option<usize>, _limit: Option<usize>) -> std::result::Result<String, String> {
|
||||
Err("not implemented".to_string())
|
||||
}
|
||||
|
||||
async fn write_file(&self, path: &str, content: &str) -> std::result::Result<(), String> {
|
||||
self.written
|
||||
.lock()
|
||||
.unwrap()
|
||||
.push((path.to_string(), content.to_string()));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn delete_file(&self, _path: &str) -> std::result::Result<(), String> {
|
||||
Err("not implemented".to_string())
|
||||
}
|
||||
|
||||
async fn file_exists(&self, _path: &str) -> std::result::Result<bool, String> {
|
||||
Ok(false)
|
||||
}
|
||||
|
||||
async fn list_directory(&self, _path: &str, _depth: Option<usize>) -> std::result::Result<Vec<agent::DirEntry>, String> {
|
||||
Err("not implemented".to_string())
|
||||
}
|
||||
|
||||
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<tokio_util::sync::CancellationToken>,
|
||||
) -> std::result::Result<agent::ExecResult, String> {
|
||||
Err("not implemented".to_string())
|
||||
}
|
||||
|
||||
async fn grep(&self, _pattern: &str, _path: &str, _options: &agent::GrepOptions) -> std::result::Result<Vec<String>, String> {
|
||||
Err("not implemented".to_string())
|
||||
}
|
||||
|
||||
async fn glob(&self, _pattern: &str, _path: Option<&str>) -> std::result::Result<Vec<String>, String> {
|
||||
Err("not implemented".to_string())
|
||||
}
|
||||
|
||||
async fn initialize(&self) -> std::result::Result<(), String> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn cleanup(&self) -> std::result::Result<(), String> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn working_directory(&self) -> &str {
|
||||
&self.working_dir
|
||||
}
|
||||
|
||||
fn platform(&self) -> &str {
|
||||
"linux"
|
||||
}
|
||||
|
||||
fn os_version(&self) -> String {
|
||||
"Linux 5.15".to_string()
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn artifact_pointers_rewritten_for_remote_execution_env() {
|
||||
// Pipeline: start -> big_output -> exit
|
||||
// big_output uses LargeOutputHandler which returns a >100KB context_update.
|
||||
// RemoteMockEnv simulates a container where local files don't exist.
|
||||
let mut graph = make_graph_with_start_exit("ArtifactSync");
|
||||
graph.attrs.insert(
|
||||
"goal".to_string(),
|
||||
AttrValue::String("Test artifact sync to remote env".to_string()),
|
||||
);
|
||||
|
||||
let mut big_output = Node::new("big_output");
|
||||
big_output.attrs.insert(
|
||||
"label".to_string(),
|
||||
AttrValue::String("Big Output".to_string()),
|
||||
);
|
||||
graph.nodes.insert("big_output".to_string(), big_output);
|
||||
|
||||
graph.edges.push(Edge::new("start", "big_output"));
|
||||
graph.edges.push(Edge::new("big_output", "exit"));
|
||||
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let mut registry = HandlerRegistry::new(Box::new(LargeOutputHandler));
|
||||
registry.register("start", Box::new(StartHandler));
|
||||
registry.register("exit", Box::new(ExitHandler));
|
||||
|
||||
let remote_env = Arc::new(RemoteMockEnv::new("/sandbox"));
|
||||
let engine = PipelineEngine::new(registry, Arc::new(EventEmitter::new()), remote_env.clone());
|
||||
let config = RunConfig {
|
||||
logs_root: dir.path().to_path_buf(),
|
||||
cancel_token: None,
|
||||
dry_run: false,
|
||||
};
|
||||
|
||||
let outcome = engine
|
||||
.run(&graph, &config)
|
||||
.await
|
||||
.expect("pipeline should succeed");
|
||||
assert_eq!(outcome.status, StageStatus::Success);
|
||||
|
||||
// The checkpoint context should contain a pointer rewritten for the remote env
|
||||
let checkpoint = Checkpoint::load(&dir.path().join("checkpoint.json"))
|
||||
.expect("checkpoint should load");
|
||||
let pointer_value = checkpoint
|
||||
.context_values
|
||||
.get("response.big_output")
|
||||
.expect("context should have response.big_output");
|
||||
let pointer_str = pointer_value.as_str().expect("pointer should be a string");
|
||||
assert!(
|
||||
pointer_str.starts_with("file:///sandbox/.attractor/artifacts/"),
|
||||
"pointer should reference remote path, got: {pointer_str}"
|
||||
);
|
||||
|
||||
// The RemoteMockEnv should have received exactly one write with >100KB content
|
||||
let written = remote_env.written.lock().unwrap();
|
||||
assert_eq!(written.len(), 1, "should have written 1 artifact");
|
||||
assert!(
|
||||
written[0].1.len() > 100 * 1024,
|
||||
"written content should be >100KB, got {} bytes",
|
||||
written[0].1.len()
|
||||
);
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue