feat: persist standalone sessions for restore
This commit is contained in:
@@ -10,7 +10,10 @@ use agen::llm_client::types::Request;
|
||||
use async_trait::async_trait;
|
||||
use futures::{Stream, stream};
|
||||
use protocol::{Event, Method};
|
||||
use standalone::{StandaloneHost, StandaloneLaunchConfig};
|
||||
use standalone::{
|
||||
StaleLeasePolicy, StandaloneHost, StandaloneLaunchConfig, StandaloneListScope,
|
||||
StandaloneSessionStatus, StandaloneSessionStore, StandaloneStartupError, StandaloneStoreError,
|
||||
};
|
||||
use uuid::Uuid;
|
||||
|
||||
#[derive(Clone)]
|
||||
@@ -159,7 +162,7 @@ async fn state_store_failure_is_redacted_and_starts_no_controller() {
|
||||
assert_eq!(error, standalone::StandaloneStartupError::StateStore);
|
||||
assert_eq!(
|
||||
error.to_string(),
|
||||
"the standalone state store could not be opened"
|
||||
"the standalone state store could not be opened or validated"
|
||||
);
|
||||
assert!(!error.to_string().contains("secret-name"));
|
||||
assert!(
|
||||
@@ -210,3 +213,261 @@ fn launch_rejects_path_profile_before_worker_startup() {
|
||||
standalone::StandaloneLaunchError::PathProfileUnsupported
|
||||
);
|
||||
}
|
||||
|
||||
type TestResult = Result<(), Box<dyn std::error::Error>>;
|
||||
|
||||
#[tokio::test]
|
||||
async fn standalone_restore_preserves_history_tasks_notifications_and_cwd_scope() -> TestResult {
|
||||
let temp = tempfile::tempdir()?;
|
||||
let cwd = temp.path().join("project");
|
||||
let state_dir = temp.path().join("client").join("standalone-sessions");
|
||||
std::fs::create_dir_all(&cwd)?;
|
||||
let launch = StandaloneLaunchConfig::new(
|
||||
&cwd,
|
||||
&state_dir,
|
||||
manifest::ProfileSelector::Default,
|
||||
"display-name-is-not-session-identity",
|
||||
)
|
||||
.resolve()?;
|
||||
let first_client = ScriptedClient::new(vec![
|
||||
vec![
|
||||
LlmEvent::tool_use_start(0, "task-1", "TaskCreate"),
|
||||
LlmEvent::tool_input_delta(
|
||||
0,
|
||||
r#"{"subject":"persisted task","description":"survives restore"}"#,
|
||||
),
|
||||
LlmEvent::tool_use_stop(0),
|
||||
],
|
||||
vec![
|
||||
LlmEvent::text_block_start(0),
|
||||
LlmEvent::text_delta(0, "first answer"),
|
||||
LlmEvent::text_block_stop(0, Some(StopReason::EndTurn)),
|
||||
],
|
||||
vec![
|
||||
LlmEvent::text_block_start(0),
|
||||
LlmEvent::text_delta(0, "notification acknowledged"),
|
||||
LlmEvent::text_block_stop(0, Some(StopReason::EndTurn)),
|
||||
],
|
||||
]);
|
||||
let host = StandaloneHost::start_with_model_client(launch, first_client).await?;
|
||||
let session_id = host.session_id();
|
||||
let mut events = host.subscribe();
|
||||
host.send(Method::run_text("first request")).await?;
|
||||
wait_for_run_end(&mut events).await?;
|
||||
host.send(Method::Notify {
|
||||
message: "persisted notification".to_string(),
|
||||
auto_run: true,
|
||||
})
|
||||
.await?;
|
||||
wait_for_run_end(&mut events).await?;
|
||||
host.shutdown().await?;
|
||||
|
||||
let store = StandaloneSessionStore::open(&state_dir)?;
|
||||
let current = store.list(&cwd, StandaloneListScope::CurrentCwd, 100)?;
|
||||
assert_eq!(current.len(), 1);
|
||||
assert_eq!(current[0].session_id, session_id);
|
||||
assert_eq!(current[0].status, StandaloneSessionStatus::Stopped);
|
||||
let other_cwd = temp.path().join("other");
|
||||
std::fs::create_dir(&other_cwd)?;
|
||||
assert!(
|
||||
store
|
||||
.list(&other_cwd, StandaloneListScope::CurrentCwd, 100)?
|
||||
.is_empty()
|
||||
);
|
||||
assert_eq!(
|
||||
store.list(&other_cwd, StandaloneListScope::All, 100)?.len(),
|
||||
1
|
||||
);
|
||||
|
||||
let second_client = ScriptedClient::new(vec![vec![
|
||||
LlmEvent::text_block_start(0),
|
||||
LlmEvent::text_delta(0, "second answer"),
|
||||
LlmEvent::text_block_stop(0, Some(StopReason::EndTurn)),
|
||||
]]);
|
||||
let second_inspection = second_client.clone();
|
||||
let host =
|
||||
StandaloneHost::restore_with_model_client(state_dir.clone(), session_id, second_client)
|
||||
.await?;
|
||||
let snapshot = format!("{:?}", host.snapshot());
|
||||
assert!(snapshot.contains("first request"), "{snapshot}");
|
||||
assert!(snapshot.contains("first answer"), "{snapshot}");
|
||||
assert!(snapshot.contains("persisted task"), "{snapshot}");
|
||||
assert!(snapshot.contains("persisted notification"), "{snapshot}");
|
||||
|
||||
let mut events = host.subscribe();
|
||||
host.send(Method::run_text("continue after restore"))
|
||||
.await?;
|
||||
wait_for_run_end(&mut events).await?;
|
||||
let request = second_inspection
|
||||
.requests()
|
||||
.into_iter()
|
||||
.next()
|
||||
.expect("restored run request");
|
||||
let projected = format!("{:?}", request.items);
|
||||
assert!(projected.contains("first answer"), "{projected}");
|
||||
assert!(projected.contains("persisted notification"), "{projected}");
|
||||
assert!(projected.contains("persisted task"), "{projected}");
|
||||
host.shutdown().await?;
|
||||
|
||||
store.delete(session_id)?;
|
||||
assert!(cwd.exists(), "deleting session state must not mutate cwd");
|
||||
assert!(matches!(
|
||||
store.load(session_id),
|
||||
Err(StandaloneStoreError::SessionNotFound(_))
|
||||
));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn standalone_restore_rejects_concurrent_lease_and_missing_cwd() -> TestResult {
|
||||
let temp = tempfile::tempdir()?;
|
||||
let cwd = temp.path().join("project");
|
||||
let moved = temp.path().join("moved-project");
|
||||
let state_dir = temp.path().join("state");
|
||||
std::fs::create_dir(&cwd)?;
|
||||
let launch = StandaloneLaunchConfig::new(
|
||||
&cwd,
|
||||
&state_dir,
|
||||
manifest::ProfileSelector::Default,
|
||||
"standalone-lease-test",
|
||||
)
|
||||
.resolve()?;
|
||||
let host =
|
||||
StandaloneHost::start_with_model_client(launch, ScriptedClient::new(Vec::new())).await?;
|
||||
let session_id = host.session_id();
|
||||
let store = StandaloneSessionStore::open(&state_dir)?;
|
||||
assert!(matches!(
|
||||
store.acquire_lease(session_id, StaleLeasePolicy::Recover),
|
||||
Err(StandaloneStoreError::SessionLeased(id)) if id == session_id
|
||||
));
|
||||
let restore = StandaloneHost::restore_with_model_client(
|
||||
state_dir.clone(),
|
||||
session_id,
|
||||
ScriptedClient::new(Vec::new()),
|
||||
)
|
||||
.await;
|
||||
assert!(matches!(
|
||||
restore,
|
||||
Err(StandaloneStartupError::SessionActive)
|
||||
));
|
||||
host.shutdown().await?;
|
||||
|
||||
std::fs::rename(&cwd, &moved)?;
|
||||
let restore = StandaloneHost::restore_with_model_client(
|
||||
state_dir,
|
||||
session_id,
|
||||
ScriptedClient::new(Vec::new()),
|
||||
)
|
||||
.await;
|
||||
assert!(matches!(
|
||||
restore,
|
||||
Err(StandaloneStartupError::WorkingDirectoryUnavailable)
|
||||
));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn standalone_restore_recovers_only_a_proven_stale_lease() -> TestResult {
|
||||
let temp = tempfile::tempdir()?;
|
||||
let state_dir = temp.path().join("state");
|
||||
let mut launch = StandaloneLaunchConfig::new(
|
||||
temp.path(),
|
||||
&state_dir,
|
||||
manifest::ProfileSelector::Default,
|
||||
"standalone-stale-lease-test",
|
||||
)
|
||||
.resolve()?;
|
||||
launch.profile.manifest.profile = Some(manifest::ProfileManifestSnapshot {
|
||||
source: manifest::ProfileSource::Registry {
|
||||
source: manifest::ProfileRegistrySource::User,
|
||||
name: "user-standalone".to_string(),
|
||||
path: None,
|
||||
provenance: Some("user-config-revision-7".to_string()),
|
||||
},
|
||||
profile: Some(manifest::ProfileMetadata {
|
||||
name: Some("User standalone".to_string()),
|
||||
description: None,
|
||||
format: None,
|
||||
}),
|
||||
});
|
||||
let host =
|
||||
StandaloneHost::start_with_model_client(launch, ScriptedClient::new(Vec::new())).await?;
|
||||
let session_id = host.session_id();
|
||||
host.shutdown().await?;
|
||||
let store = StandaloneSessionStore::open(&state_dir)?;
|
||||
assert!(matches!(
|
||||
store.load(session_id)?.manifest.profile,
|
||||
Some(manifest::ProfileManifestSnapshot {
|
||||
source: manifest::ProfileSource::Registry {
|
||||
source: manifest::ProfileRegistrySource::User,
|
||||
..
|
||||
},
|
||||
..
|
||||
})
|
||||
));
|
||||
let session_dir = state_dir.join(session_id.to_string());
|
||||
std::fs::write(
|
||||
session_dir.join("lease.json"),
|
||||
serde_json::to_vec(&serde_json::json!({
|
||||
"lease_id": uuid::Uuid::now_v7(),
|
||||
"pid": u32::MAX,
|
||||
"process_start_marker": 1,
|
||||
"acquired_at_unix_ms": 1
|
||||
}))?,
|
||||
)?;
|
||||
|
||||
let host = StandaloneHost::restore_with_model_client(
|
||||
state_dir,
|
||||
session_id,
|
||||
ScriptedClient::new(Vec::new()),
|
||||
)
|
||||
.await?;
|
||||
host.shutdown().await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn standalone_metadata_fails_closed_on_incomplete_or_newer_records() -> TestResult {
|
||||
let temp = tempfile::tempdir()?;
|
||||
let state_dir = temp.path().join("state");
|
||||
let launch = StandaloneLaunchConfig::new(
|
||||
temp.path(),
|
||||
&state_dir,
|
||||
manifest::ProfileSelector::Default,
|
||||
"standalone-schema-test",
|
||||
)
|
||||
.resolve()?;
|
||||
let host =
|
||||
StandaloneHost::start_with_model_client(launch, ScriptedClient::new(Vec::new())).await?;
|
||||
let session_id = host.session_id();
|
||||
host.shutdown().await?;
|
||||
let store = StandaloneSessionStore::open(&state_dir)?;
|
||||
let session_dir = state_dir.join(session_id.to_string());
|
||||
std::fs::write(session_dir.join("commit.pending"), b"interrupted\n")?;
|
||||
assert!(matches!(
|
||||
store.load(session_id),
|
||||
Err(StandaloneStoreError::IncompleteCommit(id)) if id == session_id
|
||||
));
|
||||
std::fs::remove_file(session_dir.join("commit.pending"))?;
|
||||
let record_path = session_dir.join("record.json");
|
||||
let mut record: serde_json::Value = serde_json::from_slice(&std::fs::read(&record_path)?)?;
|
||||
record["schema_version"] = serde_json::json!(u32::MAX);
|
||||
std::fs::write(&record_path, serde_json::to_vec_pretty(&record)?)?;
|
||||
assert!(matches!(
|
||||
store.load(session_id),
|
||||
Err(StandaloneStoreError::NewerSchema { id, .. }) if id == session_id
|
||||
));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn wait_for_run_end(events: &mut tokio::sync::broadcast::Receiver<Event>) -> TestResult {
|
||||
tokio::time::timeout(Duration::from_secs(10), async {
|
||||
loop {
|
||||
if matches!(events.recv().await, Ok(Event::RunEnd { .. })) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
})
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user