diff --git a/Cargo.lock b/Cargo.lock index e4ceb03e..45f97992 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4622,9 +4622,12 @@ version = "0.1.0" dependencies = [ "agen", "async-trait", + "fs4", "futures", "manifest", "protocol", + "serde", + "serde_json", "session-store", "tempfile", "thiserror 2.0.18", @@ -5333,6 +5336,7 @@ dependencies = [ "session-store", "standalone", "tempfile", + "thiserror 2.0.18", "ticket", "tokio", "toml", diff --git a/crates/standalone/Cargo.toml b/crates/standalone/Cargo.toml index d23af099..a3bb1813 100644 --- a/crates/standalone/Cargo.toml +++ b/crates/standalone/Cargo.toml @@ -7,11 +7,15 @@ license.workspace = true [dependencies] agen.workspace = true +fs4.workspace = true manifest.workspace = true protocol.workspace = true +serde.workspace = true +serde_json.workspace = true session-store.workspace = true thiserror.workspace = true tokio = { workspace = true, features = ["rt", "sync", "time"] } +uuid = { workspace = true, features = ["v7"] } worker.workspace = true [dev-dependencies] @@ -19,4 +23,3 @@ async-trait.workspace = true futures.workspace = true tempfile.workspace = true tokio = { workspace = true, features = ["macros", "rt-multi-thread", "time"] } -uuid = { workspace = true, features = ["v7"] } diff --git a/crates/standalone/src/host.rs b/crates/standalone/src/host.rs index d0bbdd35..d9839f76 100644 --- a/crates/standalone/src/host.rs +++ b/crates/standalone/src/host.rs @@ -1,33 +1,49 @@ +use std::path::PathBuf; use std::time::Duration; use agen::llm_client::client::LlmClient; use protocol::{Event, Method}; -use session_store::{CombinedStore, FsStore, FsWorkerStore}; +use session_store::{ + CombinedStore, FsStore, FsWorkerStore, WorkerActiveSegmentRef, WorkerMetadataStore, +}; use thiserror::Error; use tokio::sync::broadcast; use worker::bootstrap::{WorkerBootstrap, WorkerBootstrapError, WorkerBootstrapLayout}; use worker::controller::WorkerControllerTransport; -use worker::{WorkerError, WorkerFilesystemAuthority, WorkerWorkspaceContext}; +use worker::{BootstrappedWorker, WorkerError, WorkerFilesystemAuthority, WorkerWorkspaceContext}; use crate::launch::ResolvedStandaloneLaunch; +use crate::store::{ + StaleLeasePolicy, StandaloneSessionId, StandaloneSessionLease, StandaloneSessionRecord, + StandaloneSessionStore, StandaloneShutdownReason, StandaloneStoreError, +}; const DEFAULT_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(10); +type StandaloneBackingStore = CombinedStore; -/// One top-level Worker plus its process-owned local Workdir session. +/// One client-owned top-level Worker and its standalone session authority. /// -/// The host deliberately exposes the existing typed Worker protocol rather -/// than owning an HTTP/WebSocket server or a second execution model. +/// The host deliberately exposes the existing typed Worker protocol rather than owning an +/// HTTP/WebSocket server or creating Runtime/Workspace/Ticket/Workdir domain records. pub struct StandaloneHost { handle: worker::WorkerHandle, - shutdown: worker::controller::ShutdownReceiver, + shutdown: Option, shutdown_timeout: Duration, + store: StandaloneSessionStore, + worker_store: FsWorkerStore, + record: StandaloneSessionRecord, + lease: Option, } #[derive(Debug, Clone, Copy, PartialEq, Eq, Error)] pub enum StandaloneStartupError { - #[error("the standalone state store could not be opened")] + #[error("the standalone state store could not be opened or validated")] StateStore, - #[error("the resolved Worker configuration is invalid")] + #[error("the standalone session is already active")] + SessionActive, + #[error("the standalone session working directory is unavailable or changed")] + WorkingDirectoryUnavailable, + #[error("the resolved Worker configuration or persisted history is invalid")] WorkerConfiguration, #[error("the configured model provider is unavailable")] ModelProvider, @@ -49,6 +65,8 @@ pub enum StandaloneShutdownError { DeadlineExceeded, #[error("the standalone Worker shutdown confirmation was lost")] ConfirmationLost, + #[error("the standalone session final state could not be committed")] + StateStore, } impl StandaloneHost { @@ -67,24 +85,36 @@ impl StandaloneHost { } async fn start_with_optional_model_client( - launch: ResolvedStandaloneLaunch, + mut launch: ResolvedStandaloneLaunch, model_client: Option>, ) -> Result { - std::fs::create_dir_all(&launch.state_dir) - .map_err(|_| StandaloneStartupError::StateStore)?; - let session_store = FsStore::new(launch.state_dir.join("sessions")) - .map_err(|_| StandaloneStartupError::StateStore)?; - let worker_store = FsWorkerStore::new(launch.state_dir.join("workers")) - .map_err(|_| StandaloneStartupError::StateStore)?; - let store = CombinedStore::new(session_store, worker_store); + let store = StandaloneSessionStore::open(&launch.state_dir) + .map_err(classify_store_startup_error)?; + let allocation = store + .allocate(&launch.cwd, StaleLeasePolicy::Reject) + .map_err(classify_store_startup_error)?; + let id = allocation.id(); + + // The standalone session ID is the local identity. A unique internal Worker name avoids + // process-global allocation collisions without creating a Runtime/Workspace Worker ID. + launch.profile.manifest.worker.name = format!("standalone-{id}"); + let manifest = launch.profile.manifest.clone(); + let worker_name = manifest.worker.name.clone(); + let (backing_store, worker_store) = match backing_store(&store, id) { + Ok(stores) => stores, + Err(error) => { + let _ = store.abandon_allocation(allocation); + return Err(error); + } + }; let filesystem_authority = WorkerFilesystemAuthority::local(launch.cwd.clone(), launch.cwd.clone()); let workspace_context = WorkerWorkspaceContext::local_filesystem(None); - let runtime_base = launch.state_dir.join("runtime"); + let runtime_base = store.runtime_dir(id); let mut bootstrap = WorkerBootstrap::new( - launch.profile.manifest, - store, + manifest.clone(), + backing_store, launch.prompt_catalog, workspace_context, filesystem_authority, @@ -94,13 +124,155 @@ impl StandaloneHost { if let Some(model_client) = model_client { bootstrap = bootstrap.with_model_client(model_client); } - let started = bootstrap.start().await.map_err(classify_startup_error)?; + let started = match bootstrap.start().await { + Ok(started) => started, + Err(error) => { + let _ = store.abandon_allocation(allocation); + return Err(classify_startup_error(error)); + } + }; + let active = match active_pointer(&worker_store, &worker_name) { + Ok(active) => active, + Err(error) => { + stop_started_worker(started).await; + let _ = store.abandon_allocation(allocation); + return Err(error); + } + }; + let record = + match store.commit_created(&allocation, manifest, active.session_id, active.segment_id) + { + Ok(record) => record, + Err(_) => { + stop_started_worker(started).await; + let _ = store.abandon_allocation(allocation); + return Err(StandaloneStartupError::StateStore); + } + }; + Ok(Self::from_started( + started, + store, + worker_store, + record, + allocation.into_lease(), + )) + } - Ok(Self { + pub async fn restore( + state_dir: PathBuf, + session_id: StandaloneSessionId, + ) -> Result { + Self::restore_with_optional_model_client(state_dir, session_id, None).await + } + + pub async fn restore_with_model_client( + state_dir: PathBuf, + session_id: StandaloneSessionId, + model_client: C, + ) -> Result + where + C: LlmClient + 'static, + { + Self::restore_with_optional_model_client( + state_dir, + session_id, + Some(Box::new(model_client)), + ) + .await + } + + async fn restore_with_optional_model_client( + state_dir: PathBuf, + session_id: StandaloneSessionId, + model_client: Option>, + ) -> Result { + let store = + StandaloneSessionStore::open(state_dir).map_err(classify_store_startup_error)?; + let record = store + .load(session_id) + .map_err(classify_store_startup_error)?; + record.cwd.verify().map_err(classify_store_startup_error)?; + let lease = store + .acquire_lease(session_id, StaleLeasePolicy::Recover) + .map_err(classify_store_startup_error)?; + let (backing_store, worker_store) = backing_store(&store, session_id)?; + let worker_name = record.worker_name.clone(); + let manifest = record.manifest.clone(); + let filesystem_authority = WorkerFilesystemAuthority::local( + record.cwd.canonical_path.clone(), + record.cwd.canonical_path.clone(), + ); + let workspace_context = WorkerWorkspaceContext::local_filesystem(None); + let runtime_base = store.runtime_dir(session_id); + + let mut bootstrap = WorkerBootstrap::new( + manifest, + backing_store, + worker::PromptCatalogSource::builtins_only(), + workspace_context, + filesystem_authority, + WorkerBootstrapLayout::Direct { runtime_base }, + WorkerControllerTransport::InProcess, + ); + if let Some(model_client) = model_client { + bootstrap = bootstrap.with_model_client(model_client); + } + let prepared = bootstrap + .prepare_restored(&worker_name) + .await + .map_err(classify_startup_error)?; + let started = prepared.start().await.map_err(classify_startup_error)?; + let active = match active_pointer(&worker_store, &worker_name) { + Ok(active) => active, + Err(error) => { + stop_started_worker(started).await; + return Err(error); + } + }; + let record = + match store.update_active_pointer(&record, active.session_id, active.segment_id) { + Ok(record) => record, + Err(_) => { + stop_started_worker(started).await; + lease.retain(); + return Err(StandaloneStartupError::StateStore); + } + }; + Ok(Self::from_started( + started, + store, + worker_store, + record, + lease, + )) + } + + fn from_started( + started: BootstrappedWorker, + store: StandaloneSessionStore, + worker_store: FsWorkerStore, + record: StandaloneSessionRecord, + lease: StandaloneSessionLease, + ) -> Self { + Self { handle: started.handle, - shutdown: started.shutdown, + shutdown: Some(started.shutdown), shutdown_timeout: DEFAULT_SHUTDOWN_TIMEOUT, - }) + store, + worker_store, + record, + lease: Some(lease), + } + } + + #[must_use] + pub fn session_id(&self) -> StandaloneSessionId { + self.record.session_id + } + + #[must_use] + pub fn record(&self) -> &StandaloneSessionRecord { + &self.record } pub async fn send(&self, method: Method) -> Result<(), StandaloneRequestError> { @@ -123,18 +295,97 @@ impl StandaloneHost { self } - pub async fn shutdown(self) -> Result<(), StandaloneShutdownError> { - let StandaloneHost { - handle, - shutdown, - shutdown_timeout, - } = self; - let _ = handle.send(Method::Shutdown).await; - match tokio::time::timeout(shutdown_timeout, shutdown).await { - Ok(Ok(())) => Ok(()), - Ok(Err(_)) => Err(StandaloneShutdownError::ConfirmationLost), - Err(_) => Err(StandaloneShutdownError::DeadlineExceeded), + pub async fn shutdown(mut self) -> Result<(), StandaloneShutdownError> { + let _ = self.handle.send(Method::Shutdown).await; + let Some(shutdown) = self.shutdown.take() else { + self.retain_lease(); + return Err(StandaloneShutdownError::ConfirmationLost); + }; + match tokio::time::timeout(self.shutdown_timeout, shutdown).await { + Ok(Ok(())) => {} + Ok(Err(_)) => { + self.retain_lease(); + return Err(StandaloneShutdownError::ConfirmationLost); + } + Err(_) => { + self.retain_lease(); + return Err(StandaloneShutdownError::DeadlineExceeded); + } } + let active = match active_pointer(&self.worker_store, &self.record.worker_name) { + Ok(active) => active, + Err(_) => { + self.retain_lease(); + return Err(StandaloneShutdownError::StateStore); + } + }; + if self + .store + .mark_stopped( + &self.record, + active.session_id, + active.segment_id, + StandaloneShutdownReason::UserExit, + ) + .is_err() + { + self.retain_lease(); + return Err(StandaloneShutdownError::StateStore); + } + if let Some(lease) = self.lease.take() { + lease + .release() + .map_err(|_| StandaloneShutdownError::StateStore)?; + } + Ok(()) + } + + fn retain_lease(&mut self) { + if let Some(lease) = self.lease.take() { + lease.retain(); + } + } +} + +fn backing_store( + store: &StandaloneSessionStore, + id: StandaloneSessionId, +) -> Result<(StandaloneBackingStore, FsWorkerStore), StandaloneStartupError> { + let session_store = + FsStore::new(store.session_log_dir(id)).map_err(|_| StandaloneStartupError::StateStore)?; + let worker_store = FsWorkerStore::new(store.worker_metadata_dir(id)) + .map_err(|_| StandaloneStartupError::StateStore)?; + Ok(( + CombinedStore::new(session_store, worker_store.clone()), + worker_store, + )) +} + +fn active_pointer( + worker_store: &FsWorkerStore, + worker_name: &str, +) -> Result { + worker_store + .read_by_name(worker_name) + .map_err(|_| StandaloneStartupError::StateStore)? + .and_then(|metadata| metadata.active) + .ok_or(StandaloneStartupError::StateStore) +} + +async fn stop_started_worker(started: BootstrappedWorker) { + let _ = started.handle.send(Method::Shutdown).await; + let _ = tokio::time::timeout(Duration::from_secs(2), started.shutdown).await; +} + +fn classify_store_startup_error(error: StandaloneStoreError) -> StandaloneStartupError { + match error { + StandaloneStoreError::SessionLeased(_) => StandaloneStartupError::SessionActive, + StandaloneStoreError::CwdUnavailable(_) + | StandaloneStoreError::CwdNotDirectory + | StandaloneStoreError::CwdIdentityMismatch => { + StandaloneStartupError::WorkingDirectoryUnavailable + } + _ => StandaloneStartupError::StateStore, } } diff --git a/crates/standalone/src/lib.rs b/crates/standalone/src/lib.rs index b0c65331..19123cb4 100644 --- a/crates/standalone/src/lib.rs +++ b/crates/standalone/src/lib.rs @@ -6,8 +6,14 @@ pub mod host; pub mod launch; +pub mod store; pub use host::{ StandaloneHost, StandaloneRequestError, StandaloneShutdownError, StandaloneStartupError, }; pub use launch::{ResolvedStandaloneLaunch, StandaloneLaunchConfig, StandaloneLaunchError}; +pub use store::{ + StaleLeasePolicy, StandaloneCwdIdentity, StandaloneListScope, StandaloneSessionId, + StandaloneSessionRecord, StandaloneSessionStatus, StandaloneSessionStore, + StandaloneShutdownReason, StandaloneStoreError, +}; diff --git a/crates/standalone/src/store.rs b/crates/standalone/src/store.rs new file mode 100644 index 00000000..6736ce10 --- /dev/null +++ b/crates/standalone/src/store.rs @@ -0,0 +1,674 @@ +use std::fmt; +use std::fs::{self, File, OpenOptions}; +use std::io::{self, Write}; +use std::path::{Path, PathBuf}; +use std::str::FromStr; +use std::time::{SystemTime, UNIX_EPOCH}; + +use fs4::fs_std::FileExt; +use manifest::WorkerManifest; +use serde::{Deserialize, Serialize}; +use session_store::{SegmentId, SessionId}; +use thiserror::Error; +use uuid::Uuid; + +const RECORD_FILE: &str = "record.json"; +const COMMIT_MARKER: &str = "commit.pending"; +const LEASE_FILE: &str = "lease.json"; +const LEASE_LOCK_FILE: &str = "lease.lock"; +const SESSION_DIR: &str = "session"; +const WORKER_DIR: &str = "worker"; +const SCHEMA_VERSION: u32 = 1; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] +#[serde(transparent)] +pub struct StandaloneSessionId(Uuid); + +impl StandaloneSessionId { + #[must_use] + pub fn new() -> Self { + Self(Uuid::now_v7()) + } + + #[must_use] + pub fn short(self) -> String { + let simple = self.0.simple().to_string(); + simple[simple.len() - 12..].to_string() + } +} + +impl Default for StandaloneSessionId { + fn default() -> Self { + Self::new() + } +} + +impl fmt::Display for StandaloneSessionId { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + self.0.fmt(formatter) + } +} + +impl FromStr for StandaloneSessionId { + type Err = uuid::Error; + + fn from_str(value: &str) -> Result { + Uuid::parse_str(value).map(Self) + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct StandaloneCwdIdentity { + pub canonical_path: PathBuf, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub device: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub inode: Option, +} + +impl StandaloneCwdIdentity { + pub fn capture(path: impl AsRef) -> Result { + let canonical_path = + fs::canonicalize(path).map_err(StandaloneStoreError::CwdUnavailable)?; + let metadata = + fs::metadata(&canonical_path).map_err(StandaloneStoreError::CwdUnavailable)?; + if !metadata.is_dir() { + return Err(StandaloneStoreError::CwdNotDirectory); + } + #[cfg(unix)] + let (device, inode) = { + use std::os::unix::fs::MetadataExt; + (Some(metadata.dev()), Some(metadata.ino())) + }; + #[cfg(not(unix))] + let (device, inode) = (None, None); + Ok(Self { + canonical_path, + device, + inode, + }) + } + + pub fn verify(&self) -> Result { + let current = Self::capture(&self.canonical_path)?; + if current != *self { + return Err(StandaloneStoreError::CwdIdentityMismatch); + } + Ok(current.canonical_path) + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum StandaloneSessionStatus { + Active, + Stopped, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum StandaloneShutdownReason { + UserExit, + StartupFailed, + ControllerError, + ProcessInterrupted, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct StandaloneSessionRecord { + pub schema_version: u32, + pub revision: u64, + pub session_id: StandaloneSessionId, + pub worker_name: String, + pub cwd: StandaloneCwdIdentity, + pub manifest: WorkerManifest, + pub active_session_id: SessionId, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub active_segment_id: Option, + pub status: StandaloneSessionStatus, + pub created_at_unix_ms: u64, + pub updated_at_unix_ms: u64, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub shutdown_reason: Option, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum StandaloneListScope { + CurrentCwd, + All, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum StaleLeasePolicy { + Reject, + Recover, +} + +#[derive(Debug, Clone)] +pub struct StandaloneSessionStore { + root: PathBuf, +} + +impl StandaloneSessionStore { + pub fn open(root: impl Into) -> Result { + let root = root.into(); + fs::create_dir_all(&root).map_err(StandaloneStoreError::Io)?; + if !fs::metadata(&root) + .map_err(StandaloneStoreError::Io)? + .is_dir() + { + return Err(StandaloneStoreError::NotDirectory); + } + Ok(Self { root }) + } + + #[must_use] + pub fn root(&self) -> &Path { + &self.root + } + + pub fn allocate( + &self, + cwd: impl AsRef, + policy: StaleLeasePolicy, + ) -> Result { + let id = StandaloneSessionId::new(); + let cwd = StandaloneCwdIdentity::capture(cwd)?; + let dir = self.session_dir(id); + fs::create_dir(&dir).map_err(StandaloneStoreError::Io)?; + fs::create_dir(dir.join(SESSION_DIR)).map_err(StandaloneStoreError::Io)?; + fs::create_dir(dir.join(WORKER_DIR)).map_err(StandaloneStoreError::Io)?; + let lease = self.acquire_lease(id, policy)?; + Ok(StandaloneSessionAllocation { id, cwd, lease }) + } + + pub fn commit_created( + &self, + allocation: &StandaloneSessionAllocation, + manifest: WorkerManifest, + active_session_id: SessionId, + active_segment_id: Option, + ) -> Result { + let now = now_unix_ms()?; + let record = StandaloneSessionRecord { + schema_version: SCHEMA_VERSION, + revision: 1, + session_id: allocation.id, + worker_name: manifest.worker.name.clone(), + cwd: allocation.cwd.clone(), + manifest, + active_session_id, + active_segment_id, + status: StandaloneSessionStatus::Active, + created_at_unix_ms: now, + updated_at_unix_ms: now, + shutdown_reason: None, + }; + self.commit_record(None, &record)?; + Ok(record) + } + + pub fn load( + &self, + id: StandaloneSessionId, + ) -> Result { + let dir = self.session_dir(id); + if dir.join(COMMIT_MARKER).exists() { + return Err(StandaloneStoreError::IncompleteCommit(id)); + } + let bytes = fs::read(dir.join(RECORD_FILE)).map_err(|error| { + if error.kind() == io::ErrorKind::NotFound { + StandaloneStoreError::SessionNotFound(id) + } else { + StandaloneStoreError::Io(error) + } + })?; + let record: StandaloneSessionRecord = serde_json::from_slice(&bytes) + .map_err(|source| StandaloneStoreError::CorruptRecord { id, source })?; + if record.schema_version > SCHEMA_VERSION { + return Err(StandaloneStoreError::NewerSchema { + id, + found: record.schema_version, + supported: SCHEMA_VERSION, + }); + } + if record.schema_version != SCHEMA_VERSION || record.session_id != id { + return Err(StandaloneStoreError::InvalidRecord(id)); + } + Ok(record) + } + + pub fn list( + &self, + cwd: impl AsRef, + scope: StandaloneListScope, + limit: usize, + ) -> Result, StandaloneStoreError> { + let current_cwd = (scope == StandaloneListScope::CurrentCwd) + .then(|| StandaloneCwdIdentity::capture(cwd)) + .transpose()?; + let mut records = Vec::new(); + for entry in fs::read_dir(&self.root).map_err(StandaloneStoreError::Io)? { + let entry = entry.map_err(StandaloneStoreError::Io)?; + if !entry + .file_type() + .map_err(StandaloneStoreError::Io)? + .is_dir() + { + continue; + } + let Ok(id) = entry.file_name().to_string_lossy().parse() else { + continue; + }; + let record = self.load(id)?; + if current_cwd.as_ref().is_none_or(|cwd| &record.cwd == cwd) { + records.push(record); + } + } + records.sort_by(|left, right| { + right + .updated_at_unix_ms + .cmp(&left.updated_at_unix_ms) + .then_with(|| { + right + .session_id + .to_string() + .cmp(&left.session_id.to_string()) + }) + }); + records.truncate(limit); + Ok(records) + } + + pub fn acquire_lease( + &self, + id: StandaloneSessionId, + policy: StaleLeasePolicy, + ) -> Result { + let dir = self.session_dir(id); + let path = dir.join(LEASE_FILE); + let _guard = LeaseMutationGuard::acquire(&dir)?; + let lease = LeaseRecord::current()?; + loop { + match OpenOptions::new().write(true).create_new(true).open(&path) { + Ok(mut file) => { + serde_json::to_writer(&mut file, &lease).map_err(StandaloneStoreError::Json)?; + file.write_all(b"\n").map_err(StandaloneStoreError::Io)?; + file.sync_all().map_err(StandaloneStoreError::Io)?; + sync_directory(&dir)?; + return Ok(StandaloneSessionLease { + path, + lease_id: lease.lease_id, + released: false, + }); + } + Err(error) if error.kind() == io::ErrorKind::AlreadyExists => { + let existing = read_lease(&path, id)?; + if existing.is_live() { + return Err(StandaloneStoreError::SessionLeased(id)); + } + if policy == StaleLeasePolicy::Reject { + return Err(StandaloneStoreError::StaleLease(id)); + } + fs::remove_file(&path).map_err(StandaloneStoreError::Io)?; + sync_directory(&dir)?; + } + Err(error) => return Err(StandaloneStoreError::Io(error)), + } + } + } + + pub fn update_active_pointer( + &self, + record: &StandaloneSessionRecord, + active_session_id: SessionId, + active_segment_id: Option, + ) -> Result { + let mut next = record.clone(); + next.revision = next.revision.saturating_add(1); + next.updated_at_unix_ms = now_unix_ms()?; + next.active_session_id = active_session_id; + next.active_segment_id = active_segment_id; + next.status = StandaloneSessionStatus::Active; + next.shutdown_reason = None; + self.commit_record(Some(record.revision), &next)?; + Ok(next) + } + + pub fn mark_stopped( + &self, + record: &StandaloneSessionRecord, + active_session_id: SessionId, + active_segment_id: Option, + reason: StandaloneShutdownReason, + ) -> Result { + let mut next = record.clone(); + next.revision = next.revision.saturating_add(1); + next.updated_at_unix_ms = now_unix_ms()?; + next.active_session_id = active_session_id; + next.active_segment_id = active_segment_id; + next.status = StandaloneSessionStatus::Stopped; + next.shutdown_reason = Some(reason); + self.commit_record(Some(record.revision), &next)?; + Ok(next) + } + + pub fn delete(&self, id: StandaloneSessionId) -> Result<(), StandaloneStoreError> { + let record = self.load(id)?; + if record.status != StandaloneSessionStatus::Stopped { + return Err(StandaloneStoreError::DeleteActive(id)); + } + let session_dir = self.session_dir(id); + let _guard = LeaseMutationGuard::acquire(&session_dir)?; + let lease_path = session_dir.join(LEASE_FILE); + if lease_path.exists() { + let lease = read_lease(&lease_path, id)?; + return Err(if lease.is_live() { + StandaloneStoreError::SessionLeased(id) + } else { + StandaloneStoreError::StaleLease(id) + }); + } + fs::remove_dir_all(self.session_dir(id)).map_err(StandaloneStoreError::Io)?; + sync_directory(&self.root) + } + + #[must_use] + pub fn session_log_dir(&self, id: StandaloneSessionId) -> PathBuf { + self.session_dir(id).join(SESSION_DIR) + } + + #[must_use] + pub fn worker_metadata_dir(&self, id: StandaloneSessionId) -> PathBuf { + self.session_dir(id).join(WORKER_DIR) + } + + #[must_use] + pub(crate) fn runtime_dir(&self, id: StandaloneSessionId) -> PathBuf { + self.session_dir(id).join("runtime") + } + + pub(crate) fn abandon_allocation( + &self, + allocation: StandaloneSessionAllocation, + ) -> Result<(), StandaloneStoreError> { + let id = allocation.id; + allocation.lease.release()?; + fs::remove_dir_all(self.session_dir(id)).map_err(StandaloneStoreError::Io)?; + sync_directory(&self.root) + } + + fn commit_record( + &self, + expected_revision: Option, + next: &StandaloneSessionRecord, + ) -> Result<(), StandaloneStoreError> { + let dir = self.session_dir(next.session_id); + let marker = dir.join(COMMIT_MARKER); + let mut marker_file = OpenOptions::new() + .write(true) + .create_new(true) + .open(&marker) + .map_err(|error| { + if error.kind() == io::ErrorKind::AlreadyExists { + StandaloneStoreError::IncompleteCommit(next.session_id) + } else { + StandaloneStoreError::Io(error) + } + })?; + writeln!(marker_file, "{}", next.revision).map_err(StandaloneStoreError::Io)?; + marker_file.sync_all().map_err(StandaloneStoreError::Io)?; + sync_directory(&dir)?; + + if let Some(expected) = expected_revision { + let current = self.load_record_while_committing(next.session_id)?; + if current.revision != expected { + let _ = fs::remove_file(&marker); + return Err(StandaloneStoreError::RevisionConflict { + id: next.session_id, + expected, + found: current.revision, + }); + } + } + + let temporary = dir.join(format!("record.{}.tmp", Uuid::now_v7())); + let result = (|| { + let mut file = OpenOptions::new() + .write(true) + .create_new(true) + .open(&temporary) + .map_err(StandaloneStoreError::Io)?; + serde_json::to_writer_pretty(&mut file, next).map_err(StandaloneStoreError::Json)?; + file.write_all(b"\n").map_err(StandaloneStoreError::Io)?; + file.sync_all().map_err(StandaloneStoreError::Io)?; + fs::rename(&temporary, dir.join(RECORD_FILE)).map_err(StandaloneStoreError::Io)?; + sync_directory(&dir)?; + fs::remove_file(&marker).map_err(StandaloneStoreError::Io)?; + sync_directory(&dir) + })(); + if result.is_err() { + let _ = fs::remove_file(&temporary); + } + result + } + + fn load_record_while_committing( + &self, + id: StandaloneSessionId, + ) -> Result { + let bytes = + fs::read(self.session_dir(id).join(RECORD_FILE)).map_err(StandaloneStoreError::Io)?; + serde_json::from_slice(&bytes) + .map_err(|source| StandaloneStoreError::CorruptRecord { id, source }) + } + + fn session_dir(&self, id: StandaloneSessionId) -> PathBuf { + self.root.join(id.to_string()) + } +} + +#[derive(Debug)] +pub struct StandaloneSessionAllocation { + id: StandaloneSessionId, + cwd: StandaloneCwdIdentity, + lease: StandaloneSessionLease, +} + +impl StandaloneSessionAllocation { + #[must_use] + pub fn id(&self) -> StandaloneSessionId { + self.id + } + + #[must_use] + pub fn cwd(&self) -> &StandaloneCwdIdentity { + &self.cwd + } + + pub fn into_lease(self) -> StandaloneSessionLease { + self.lease + } +} + +#[derive(Debug)] +pub struct StandaloneSessionLease { + path: PathBuf, + lease_id: Uuid, + released: bool, +} + +impl StandaloneSessionLease { + pub fn release(mut self) -> Result<(), StandaloneStoreError> { + self.release_inner() + } + + pub(crate) fn retain(mut self) { + self.released = true; + } + + fn release_inner(&mut self) -> Result<(), StandaloneStoreError> { + if self.released { + return Ok(()); + } + if self.path.exists() { + let parent = self.path.parent().expect("lease parent"); + let _guard = LeaseMutationGuard::acquire(parent)?; + let bytes = fs::read(&self.path).map_err(StandaloneStoreError::Io)?; + let current: LeaseRecord = + serde_json::from_slice(&bytes).map_err(StandaloneStoreError::Json)?; + if current.lease_id != self.lease_id { + return Err(StandaloneStoreError::LeaseOwnershipLost); + } + fs::remove_file(&self.path).map_err(StandaloneStoreError::Io)?; + sync_directory(self.path.parent().expect("lease parent"))?; + } + self.released = true; + Ok(()) + } +} + +impl Drop for StandaloneSessionLease { + fn drop(&mut self) { + let _ = self.release_inner(); + } +} + +struct LeaseMutationGuard { + file: File, +} + +impl LeaseMutationGuard { + fn acquire(dir: &Path) -> Result { + let file = OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(false) + .open(dir.join(LEASE_LOCK_FILE)) + .map_err(StandaloneStoreError::Io)?; + file.lock_exclusive().map_err(StandaloneStoreError::Io)?; + Ok(Self { file }) + } +} + +impl Drop for LeaseMutationGuard { + fn drop(&mut self) { + let _ = FileExt::unlock(&self.file); + } +} + +#[derive(Debug, Serialize, Deserialize)] +struct LeaseRecord { + lease_id: Uuid, + pid: u32, + #[serde(default, skip_serializing_if = "Option::is_none")] + process_start_marker: Option, + acquired_at_unix_ms: u64, +} + +impl LeaseRecord { + fn current() -> Result { + Ok(Self { + lease_id: Uuid::now_v7(), + pid: std::process::id(), + process_start_marker: process_start_marker(std::process::id()), + acquired_at_unix_ms: now_unix_ms()?, + }) + } + + fn is_live(&self) -> bool { + process_start_marker(self.pid) + .zip(self.process_start_marker) + .is_some_and(|(current, recorded)| current == recorded) + } +} + +fn read_lease(path: &Path, id: StandaloneSessionId) -> Result { + let bytes = fs::read(path).map_err(StandaloneStoreError::Io)?; + serde_json::from_slice(&bytes) + .map_err(|source| StandaloneStoreError::CorruptLease { id, source }) +} + +#[cfg(target_os = "linux")] +fn process_start_marker(pid: u32) -> Option { + let stat = fs::read_to_string(format!("/proc/{pid}/stat")).ok()?; + let tail = stat.rsplit_once(") ")?.1; + tail.split_whitespace().nth(19)?.parse().ok() +} + +#[cfg(not(target_os = "linux"))] +fn process_start_marker(pid: u32) -> Option { + (pid == std::process::id()).then_some(0) +} + +fn now_unix_ms() -> Result { + let duration = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map_err(|_| StandaloneStoreError::Clock)?; + u64::try_from(duration.as_millis()).map_err(|_| StandaloneStoreError::Clock) +} + +fn sync_directory(path: &Path) -> Result<(), StandaloneStoreError> { + File::open(path) + .and_then(|file| file.sync_all()) + .map_err(StandaloneStoreError::Io) +} + +#[derive(Debug, Error)] +pub enum StandaloneStoreError { + #[error("standalone state path is not a directory")] + NotDirectory, + #[error("standalone cwd is unavailable")] + CwdUnavailable(#[source] io::Error), + #[error("standalone cwd is not a directory")] + CwdNotDirectory, + #[error("standalone cwd identity no longer matches the persisted session")] + CwdIdentityMismatch, + #[error("standalone session {0} was not found")] + SessionNotFound(StandaloneSessionId), + #[error("standalone session {0} has an incomplete metadata commit")] + IncompleteCommit(StandaloneSessionId), + #[error("standalone session {0} has invalid metadata")] + InvalidRecord(StandaloneSessionId), + #[error("standalone session {id} metadata is corrupt")] + CorruptRecord { + id: StandaloneSessionId, + #[source] + source: serde_json::Error, + }, + #[error("standalone session {id} lease is corrupt")] + CorruptLease { + id: StandaloneSessionId, + #[source] + source: serde_json::Error, + }, + #[error("standalone session {id} uses schema {found}, newer than supported schema {supported}")] + NewerSchema { + id: StandaloneSessionId, + found: u32, + supported: u32, + }, + #[error("standalone session {0} is already active")] + SessionLeased(StandaloneSessionId), + #[error("standalone session {0} has a stale lease; explicit recovery is required")] + StaleLease(StandaloneSessionId), + #[error("standalone session lease ownership changed")] + LeaseOwnershipLost, + #[error("standalone session {0} must be stopped before deletion")] + DeleteActive(StandaloneSessionId), + #[error( + "standalone session {id} metadata revision changed (expected {expected}, found {found})" + )] + RevisionConflict { + id: StandaloneSessionId, + expected: u64, + found: u64, + }, + #[error("system clock is before the Unix epoch or out of range")] + Clock, + #[error("standalone metadata serialization failed")] + Json(#[source] serde_json::Error), + #[error("standalone state I/O failed")] + Io(#[source] io::Error), +} diff --git a/crates/standalone/tests/host.rs b/crates/standalone/tests/host.rs index 157f15e6..0e69739c 100644 --- a/crates/standalone/tests/host.rs +++ b/crates/standalone/tests/host.rs @@ -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>; + +#[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) -> TestResult { + tokio::time::timeout(Duration::from_secs(10), async { + loop { + if matches!(events.recv().await, Ok(Event::RunEnd { .. })) { + break; + } + } + }) + .await?; + Ok(()) +} diff --git a/crates/worker/src/bootstrap.rs b/crates/worker/src/bootstrap.rs index 7be14174..c5ec07c5 100644 --- a/crates/worker/src/bootstrap.rs +++ b/crates/worker/src/bootstrap.rs @@ -130,6 +130,29 @@ where Ok(PreparedWorker::new(worker, self.layout, self.transport)) } + pub async fn prepare_restored( + self, + worker_name: &str, + ) -> Result, St>, WorkerBootstrapError> { + let mut worker = + Worker::restore_pending_from_worker_metadata_with_context_and_model_client( + worker_name, + self.manifest, + self.store, + self.prompt_catalog, + self.workspace_context, + self.filesystem_authority, + self.model_client, + ) + .await + .map_err(WorkerBootstrapError::Worker)?; + + if let Some(workdir_session) = self.workdir_session { + worker.bind_workdir_session(Some(workdir_session)); + } + Ok(PreparedWorker::new(worker, self.layout, self.transport)) + } + pub async fn start(self) -> Result { self.prepare().await?.start().await } diff --git a/crates/worker/src/worker.rs b/crates/worker/src/worker.rs index 6aa1d871..c1f15ca0 100644 --- a/crates/worker/src/worker.rs +++ b/crates/worker/src/worker.rs @@ -5515,6 +5515,27 @@ where loader: PromptCatalogSource, workspace_context: WorkerWorkspaceContext, filesystem_authority: WorkerFilesystemAuthority, + ) -> Result { + Self::restore_pending_from_worker_metadata_with_context_and_model_client( + worker_name, + fallback, + store, + loader, + workspace_context, + filesystem_authority, + None, + ) + .await + } + + pub(crate) async fn restore_pending_from_worker_metadata_with_context_and_model_client( + worker_name: &str, + fallback: WorkerManifest, + store: St, + loader: PromptCatalogSource, + workspace_context: WorkerWorkspaceContext, + filesystem_authority: WorkerFilesystemAuthority, + model_client: Option>, ) -> Result { let metadata = store @@ -5535,7 +5556,7 @@ where worker_name: worker_name.to_string(), })?; if let Some(segment_id) = active.segment_id { - return Self::restore_from_manifest_with_context( + return Self::restore_from_manifest_with_context_and_model_client( active.session_id, segment_id, restore_manifest_from_worker_metadata_snapshot( @@ -5547,6 +5568,7 @@ where loader, workspace_context, filesystem_authority, + model_client, ) .await; } @@ -5557,12 +5579,13 @@ where })?; let manifest = restore_manifest_from_worker_metadata_snapshot(worker_name, Some(snapshot), fallback)?; - Self::from_manifest_with_context( + Self::from_manifest_with_context_and_model_client( manifest, store, loader, workspace_context, filesystem_authority, + model_client, ) .await } @@ -5614,6 +5637,29 @@ where loader: PromptCatalogSource, workspace_context: WorkerWorkspaceContext, filesystem_authority: WorkerFilesystemAuthority, + ) -> Result { + Self::restore_from_manifest_with_context_and_model_client( + session_id, + segment_id, + manifest, + store, + loader, + workspace_context, + filesystem_authority, + None, + ) + .await + } + + pub(crate) async fn restore_from_manifest_with_context_and_model_client( + session_id: SessionId, + segment_id: SegmentId, + manifest: WorkerManifest, + store: St, + loader: PromptCatalogSource, + workspace_context: WorkerWorkspaceContext, + filesystem_authority: WorkerFilesystemAuthority, + model_client: Option>, ) -> Result { // Read raw entries once so we can both reconstruct state and // seed the broadcast sink's mirror with the same prefix that @@ -5629,13 +5675,14 @@ where let mirror_entries: Vec = raw_entries.clone(); let scope_config = effective_restore_scope_config(&store, &manifest)?; - let common = prepare_worker_common_with_context( + let common = prepare_worker_common_with_context_and_model_client( &manifest, &loader, /* parse_template */ false, workspace_context, filesystem_authority, scope_config, + model_client, )?; // Atomic: register_worker inside install_top_level rejects when