use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicU64, Ordering};
use sha2::Digest;
use shepherd_cli::shepherd::{
RunState,
registry::{
DispatchSingletonInput, DispatchSingletonPublicationInput, Error, Registry,
SingletonPublicationState,
},
};
use shepherd_cli::{DispatchStore, PublicationFault};
static NEXT: AtomicU64 = AtomicU64::new(0);
fn fixture(label: &str) -> PathBuf {
let ordinal = NEXT.fetch_add(1, Ordering::Relaxed);
let root = std::env::temp_dir().join(format!(
"shepherd-publication-{label}-{}-{ordinal}",
std::process::id()
));
std::fs::create_dir_all(&root).expect("create fixture");
std::fs::canonicalize(root).expect("canonical fixture")
}
fn claim(agent_id: &str) -> DispatchSingletonInput {
DispatchSingletonInput {
project_id: "0192f6e8-7b2c-7abc-8def-0123456789ab".into(),
run_id: "v645".into(),
role: "engineer".into(),
lane_id: None,
agent_id: agent_id.into(),
harness: "claude".into(),
agent_type: "shepherd:engineer".into(),
parent_agent_id: None,
session_id: format!("session-{agent_id}"),
write_scope: vec![".shepherd/runs/v645/plan.md".into()],
claimed_at: 1,
resumes_agent_id: None,
}
}
fn input(agent_id: &str, nonce: &str) -> DispatchSingletonPublicationInput {
let record_json =
format!("{{\"schema\":\"shepherd.dispatch/3\",\"agent_id\":\"{agent_id}\"}}\n");
let record_sha256 = sha256(record_json.as_bytes());
DispatchSingletonPublicationInput {
nonce: nonce.into(),
claim: claim(agent_id),
record_path: format!("v645/dispatch/{agent_id}.json"),
record_sha256,
record_json,
prepared_at: 1,
}
}
fn sha256(bytes: &[u8]) -> String {
let digest = sha2::Sha256::digest(bytes);
digest.iter().map(|byte| format!("{byte:02x}")).collect()
}
fn setup(root: &Path) -> (PathBuf, PathBuf) {
let runs = root.join(".shepherd/runs");
let state: RunState = serde_json::from_value(serde_json::json!({
"run": "v645",
"status": "executing"
}))
.expect("run state");
state.store(&runs.join("v645/run.json")).expect("store run");
let db = root.join(".shepherd/shepherd.db");
let registry = Registry::open_migrated(&db).expect("migrate registry");
registry
.execute(
"INSERT INTO projects (id, name, created_at, updated_at) VALUES (?1, ?2, ?3, ?3)",
("0192f6e8-7b2c-7abc-8def-0123456789ab", "fixture", 1_i64),
)
.expect("seed fixture project");
(runs, db)
}
#[test]
fn preparing_filesystem_publication_replays_after_injected_failure() {
let root = fixture("replay");
let (runs, db) = setup(&root);
let mut registry = Registry::open_migrated(&db).expect("registry");
let publication = registry
.transaction_immediate::<_, Error, _>(|tx| {
tx.prepare_dispatch_singleton(&input("engineer-a", "nonce-replay"))
})
.expect("prepare intent");
let store = DispatchStore::new(&runs).with_publication_fault(PublicationFault::AfterPreparing);
assert!(store.publish_singleton_prepared(&publication).is_err());
assert!(
runs.join("v645/dispatch/.singleton.nonce-replay.preparing")
.is_file()
);
assert!(!runs.join("v645/dispatch/engineer-a.json").exists());
let replay = DispatchStore::new(&runs)
.reconcile_singletons(&mut registry, 0)
.expect("reconcile prepared publication despite a clock rollback");
assert_eq!(replay.published, 1);
assert_eq!(
std::fs::read(runs.join("v645/dispatch/engineer-a.json")).expect("read final"),
publication.record_json.as_bytes()
);
let loaded = registry
.load_dispatch_publication("nonce-replay")
.expect("load intent")
.expect("intent exists");
assert_eq!(loaded.state, SingletonPublicationState::Published);
assert!(
!runs
.join("v645/dispatch/.singleton.nonce-replay.preparing")
.exists()
);
drop(registry);
std::fs::remove_dir_all(root).expect("cleanup");
}
#[test]
fn filesystem_publish_failure_after_final_name_is_replayable() {
let root = fixture("after-filesystem");
let (runs, db) = setup(&root);
let mut registry = Registry::open_migrated(&db).expect("registry");
let publication = registry
.transaction_immediate::<_, Error, _>(|tx| {
tx.prepare_dispatch_singleton(&input("engineer-a", "nonce-after-fs"))
})
.expect("prepare intent");
let store =
DispatchStore::new(&runs).with_publication_fault(PublicationFault::AfterFilesystemPublish);
assert!(store.publish_singleton_prepared(&publication).is_err());
assert!(runs.join("v645/dispatch/engineer-a.json").is_file());
let replay = DispatchStore::new(&runs)
.reconcile_singletons(&mut registry, 2)
.expect("reconcile final filesystem publication");
assert_eq!(replay.published, 1);
assert_eq!(
registry
.load_dispatch_publication("nonce-after-fs")
.expect("load publication")
.expect("publication exists")
.state,
SingletonPublicationState::Published
);
drop(registry);
std::fs::remove_dir_all(root).expect("cleanup");
}
#[test]
fn corrupt_published_bytes_are_quarantined_and_release_the_claim() {
let root = fixture("corruption");
let (runs, db) = setup(&root);
let mut registry = Registry::open_migrated(&db).expect("registry");
let publication = registry
.transaction_immediate::<_, Error, _>(|tx| {
tx.prepare_dispatch_singleton(&input("engineer-a", "nonce-corrupt"))
})
.expect("prepare intent");
DispatchStore::new(&runs)
.publish_singleton_prepared(&publication)
.expect("publish fs");
registry
.transaction_immediate::<_, Error, _>(|tx| {
tx.mark_dispatch_singleton_published("nonce-corrupt", 2)
})
.expect("mark published");
std::fs::write(runs.join("v645/dispatch/engineer-a.json"), b"forged\n").expect("corrupt final");
let replay = DispatchStore::new(&runs)
.reconcile_singletons(&mut registry, 3)
.expect("reconcile corruption");
assert_eq!(replay.quarantined, 1);
let loaded = registry
.load_dispatch_publication("nonce-corrupt")
.expect("load intent")
.expect("intent exists");
assert_eq!(loaded.state, SingletonPublicationState::Quarantined);
assert!(
registry
.load_dispatch_singleton("project-1", "v645", "engineer", "__run__")
.expect("load claim")
.is_none()
);
assert!(
runs.join("v645/dispatch/quarantine/nonce-corrupt.preparing")
.exists()
);
assert!(
runs.join("v645/dispatch/quarantine/nonce-corrupt.published")
.exists()
);
assert!(!runs.join("v645/dispatch/engineer-a.json").exists());
drop(registry);
std::fs::remove_dir_all(root).expect("cleanup");
}