feat: persist standalone sessions for restore
This commit is contained in:
Generated
+4
@@ -4622,9 +4622,12 @@ version = "0.1.0"
|
|||||||
dependencies = [
|
dependencies = [
|
||||||
"agen",
|
"agen",
|
||||||
"async-trait",
|
"async-trait",
|
||||||
|
"fs4",
|
||||||
"futures",
|
"futures",
|
||||||
"manifest",
|
"manifest",
|
||||||
"protocol",
|
"protocol",
|
||||||
|
"serde",
|
||||||
|
"serde_json",
|
||||||
"session-store",
|
"session-store",
|
||||||
"tempfile",
|
"tempfile",
|
||||||
"thiserror 2.0.18",
|
"thiserror 2.0.18",
|
||||||
@@ -5333,6 +5336,7 @@ dependencies = [
|
|||||||
"session-store",
|
"session-store",
|
||||||
"standalone",
|
"standalone",
|
||||||
"tempfile",
|
"tempfile",
|
||||||
|
"thiserror 2.0.18",
|
||||||
"ticket",
|
"ticket",
|
||||||
"tokio",
|
"tokio",
|
||||||
"toml",
|
"toml",
|
||||||
|
|||||||
@@ -7,11 +7,15 @@ license.workspace = true
|
|||||||
|
|
||||||
[dependencies]
|
[dependencies]
|
||||||
agen.workspace = true
|
agen.workspace = true
|
||||||
|
fs4.workspace = true
|
||||||
manifest.workspace = true
|
manifest.workspace = true
|
||||||
protocol.workspace = true
|
protocol.workspace = true
|
||||||
|
serde.workspace = true
|
||||||
|
serde_json.workspace = true
|
||||||
session-store.workspace = true
|
session-store.workspace = true
|
||||||
thiserror.workspace = true
|
thiserror.workspace = true
|
||||||
tokio = { workspace = true, features = ["rt", "sync", "time"] }
|
tokio = { workspace = true, features = ["rt", "sync", "time"] }
|
||||||
|
uuid = { workspace = true, features = ["v7"] }
|
||||||
worker.workspace = true
|
worker.workspace = true
|
||||||
|
|
||||||
[dev-dependencies]
|
[dev-dependencies]
|
||||||
@@ -19,4 +23,3 @@ async-trait.workspace = true
|
|||||||
futures.workspace = true
|
futures.workspace = true
|
||||||
tempfile.workspace = true
|
tempfile.workspace = true
|
||||||
tokio = { workspace = true, features = ["macros", "rt-multi-thread", "time"] }
|
tokio = { workspace = true, features = ["macros", "rt-multi-thread", "time"] }
|
||||||
uuid = { workspace = true, features = ["v7"] }
|
|
||||||
|
|||||||
+285
-34
@@ -1,33 +1,49 @@
|
|||||||
|
use std::path::PathBuf;
|
||||||
use std::time::Duration;
|
use std::time::Duration;
|
||||||
|
|
||||||
use agen::llm_client::client::LlmClient;
|
use agen::llm_client::client::LlmClient;
|
||||||
use protocol::{Event, Method};
|
use protocol::{Event, Method};
|
||||||
use session_store::{CombinedStore, FsStore, FsWorkerStore};
|
use session_store::{
|
||||||
|
CombinedStore, FsStore, FsWorkerStore, WorkerActiveSegmentRef, WorkerMetadataStore,
|
||||||
|
};
|
||||||
use thiserror::Error;
|
use thiserror::Error;
|
||||||
use tokio::sync::broadcast;
|
use tokio::sync::broadcast;
|
||||||
use worker::bootstrap::{WorkerBootstrap, WorkerBootstrapError, WorkerBootstrapLayout};
|
use worker::bootstrap::{WorkerBootstrap, WorkerBootstrapError, WorkerBootstrapLayout};
|
||||||
use worker::controller::WorkerControllerTransport;
|
use worker::controller::WorkerControllerTransport;
|
||||||
use worker::{WorkerError, WorkerFilesystemAuthority, WorkerWorkspaceContext};
|
use worker::{BootstrappedWorker, WorkerError, WorkerFilesystemAuthority, WorkerWorkspaceContext};
|
||||||
|
|
||||||
use crate::launch::ResolvedStandaloneLaunch;
|
use crate::launch::ResolvedStandaloneLaunch;
|
||||||
|
use crate::store::{
|
||||||
|
StaleLeasePolicy, StandaloneSessionId, StandaloneSessionLease, StandaloneSessionRecord,
|
||||||
|
StandaloneSessionStore, StandaloneShutdownReason, StandaloneStoreError,
|
||||||
|
};
|
||||||
|
|
||||||
const DEFAULT_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(10);
|
const DEFAULT_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(10);
|
||||||
|
type StandaloneBackingStore = CombinedStore<FsStore, FsWorkerStore>;
|
||||||
|
|
||||||
/// 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
|
/// The host deliberately exposes the existing typed Worker protocol rather than owning an
|
||||||
/// than owning an HTTP/WebSocket server or a second execution model.
|
/// HTTP/WebSocket server or creating Runtime/Workspace/Ticket/Workdir domain records.
|
||||||
pub struct StandaloneHost {
|
pub struct StandaloneHost {
|
||||||
handle: worker::WorkerHandle,
|
handle: worker::WorkerHandle,
|
||||||
shutdown: worker::controller::ShutdownReceiver,
|
shutdown: Option<worker::controller::ShutdownReceiver>,
|
||||||
shutdown_timeout: Duration,
|
shutdown_timeout: Duration,
|
||||||
|
store: StandaloneSessionStore,
|
||||||
|
worker_store: FsWorkerStore,
|
||||||
|
record: StandaloneSessionRecord,
|
||||||
|
lease: Option<StandaloneSessionLease>,
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Error)]
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Error)]
|
||||||
pub enum StandaloneStartupError {
|
pub enum StandaloneStartupError {
|
||||||
#[error("the standalone state store could not be opened")]
|
#[error("the standalone state store could not be opened or validated")]
|
||||||
StateStore,
|
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,
|
WorkerConfiguration,
|
||||||
#[error("the configured model provider is unavailable")]
|
#[error("the configured model provider is unavailable")]
|
||||||
ModelProvider,
|
ModelProvider,
|
||||||
@@ -49,6 +65,8 @@ pub enum StandaloneShutdownError {
|
|||||||
DeadlineExceeded,
|
DeadlineExceeded,
|
||||||
#[error("the standalone Worker shutdown confirmation was lost")]
|
#[error("the standalone Worker shutdown confirmation was lost")]
|
||||||
ConfirmationLost,
|
ConfirmationLost,
|
||||||
|
#[error("the standalone session final state could not be committed")]
|
||||||
|
StateStore,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl StandaloneHost {
|
impl StandaloneHost {
|
||||||
@@ -67,24 +85,36 @@ impl StandaloneHost {
|
|||||||
}
|
}
|
||||||
|
|
||||||
async fn start_with_optional_model_client(
|
async fn start_with_optional_model_client(
|
||||||
launch: ResolvedStandaloneLaunch,
|
mut launch: ResolvedStandaloneLaunch,
|
||||||
model_client: Option<Box<dyn LlmClient>>,
|
model_client: Option<Box<dyn LlmClient>>,
|
||||||
) -> Result<Self, StandaloneStartupError> {
|
) -> Result<Self, StandaloneStartupError> {
|
||||||
std::fs::create_dir_all(&launch.state_dir)
|
let store = StandaloneSessionStore::open(&launch.state_dir)
|
||||||
.map_err(|_| StandaloneStartupError::StateStore)?;
|
.map_err(classify_store_startup_error)?;
|
||||||
let session_store = FsStore::new(launch.state_dir.join("sessions"))
|
let allocation = store
|
||||||
.map_err(|_| StandaloneStartupError::StateStore)?;
|
.allocate(&launch.cwd, StaleLeasePolicy::Reject)
|
||||||
let worker_store = FsWorkerStore::new(launch.state_dir.join("workers"))
|
.map_err(classify_store_startup_error)?;
|
||||||
.map_err(|_| StandaloneStartupError::StateStore)?;
|
let id = allocation.id();
|
||||||
let store = CombinedStore::new(session_store, worker_store);
|
|
||||||
|
// 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 =
|
let filesystem_authority =
|
||||||
WorkerFilesystemAuthority::local(launch.cwd.clone(), launch.cwd.clone());
|
WorkerFilesystemAuthority::local(launch.cwd.clone(), launch.cwd.clone());
|
||||||
let workspace_context = WorkerWorkspaceContext::local_filesystem(None);
|
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(
|
let mut bootstrap = WorkerBootstrap::new(
|
||||||
launch.profile.manifest,
|
manifest.clone(),
|
||||||
store,
|
backing_store,
|
||||||
launch.prompt_catalog,
|
launch.prompt_catalog,
|
||||||
workspace_context,
|
workspace_context,
|
||||||
filesystem_authority,
|
filesystem_authority,
|
||||||
@@ -94,13 +124,155 @@ impl StandaloneHost {
|
|||||||
if let Some(model_client) = model_client {
|
if let Some(model_client) = model_client {
|
||||||
bootstrap = bootstrap.with_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, StandaloneStartupError> {
|
||||||
|
Self::restore_with_optional_model_client(state_dir, session_id, None).await
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn restore_with_model_client<C>(
|
||||||
|
state_dir: PathBuf,
|
||||||
|
session_id: StandaloneSessionId,
|
||||||
|
model_client: C,
|
||||||
|
) -> Result<Self, StandaloneStartupError>
|
||||||
|
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<Box<dyn LlmClient>>,
|
||||||
|
) -> Result<Self, StandaloneStartupError> {
|
||||||
|
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,
|
handle: started.handle,
|
||||||
shutdown: started.shutdown,
|
shutdown: Some(started.shutdown),
|
||||||
shutdown_timeout: DEFAULT_SHUTDOWN_TIMEOUT,
|
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> {
|
pub async fn send(&self, method: Method) -> Result<(), StandaloneRequestError> {
|
||||||
@@ -123,18 +295,97 @@ impl StandaloneHost {
|
|||||||
self
|
self
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn shutdown(self) -> Result<(), StandaloneShutdownError> {
|
pub async fn shutdown(mut self) -> Result<(), StandaloneShutdownError> {
|
||||||
let StandaloneHost {
|
let _ = self.handle.send(Method::Shutdown).await;
|
||||||
handle,
|
let Some(shutdown) = self.shutdown.take() else {
|
||||||
shutdown,
|
self.retain_lease();
|
||||||
shutdown_timeout,
|
return Err(StandaloneShutdownError::ConfirmationLost);
|
||||||
} = self;
|
};
|
||||||
let _ = handle.send(Method::Shutdown).await;
|
match tokio::time::timeout(self.shutdown_timeout, shutdown).await {
|
||||||
match tokio::time::timeout(shutdown_timeout, shutdown).await {
|
Ok(Ok(())) => {}
|
||||||
Ok(Ok(())) => Ok(()),
|
Ok(Err(_)) => {
|
||||||
Ok(Err(_)) => Err(StandaloneShutdownError::ConfirmationLost),
|
self.retain_lease();
|
||||||
Err(_) => Err(StandaloneShutdownError::DeadlineExceeded),
|
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<WorkerActiveSegmentRef, StandaloneStartupError> {
|
||||||
|
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,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -6,8 +6,14 @@
|
|||||||
|
|
||||||
pub mod host;
|
pub mod host;
|
||||||
pub mod launch;
|
pub mod launch;
|
||||||
|
pub mod store;
|
||||||
|
|
||||||
pub use host::{
|
pub use host::{
|
||||||
StandaloneHost, StandaloneRequestError, StandaloneShutdownError, StandaloneStartupError,
|
StandaloneHost, StandaloneRequestError, StandaloneShutdownError, StandaloneStartupError,
|
||||||
};
|
};
|
||||||
pub use launch::{ResolvedStandaloneLaunch, StandaloneLaunchConfig, StandaloneLaunchError};
|
pub use launch::{ResolvedStandaloneLaunch, StandaloneLaunchConfig, StandaloneLaunchError};
|
||||||
|
pub use store::{
|
||||||
|
StaleLeasePolicy, StandaloneCwdIdentity, StandaloneListScope, StandaloneSessionId,
|
||||||
|
StandaloneSessionRecord, StandaloneSessionStatus, StandaloneSessionStore,
|
||||||
|
StandaloneShutdownReason, StandaloneStoreError,
|
||||||
|
};
|
||||||
|
|||||||
@@ -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<Self, Self::Err> {
|
||||||
|
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<u64>,
|
||||||
|
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||||
|
pub inode: Option<u64>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl StandaloneCwdIdentity {
|
||||||
|
pub fn capture(path: impl AsRef<Path>) -> Result<Self, StandaloneStoreError> {
|
||||||
|
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<PathBuf, StandaloneStoreError> {
|
||||||
|
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<SegmentId>,
|
||||||
|
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<StandaloneShutdownReason>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[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<PathBuf>) -> Result<Self, StandaloneStoreError> {
|
||||||
|
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<Path>,
|
||||||
|
policy: StaleLeasePolicy,
|
||||||
|
) -> Result<StandaloneSessionAllocation, StandaloneStoreError> {
|
||||||
|
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<SegmentId>,
|
||||||
|
) -> Result<StandaloneSessionRecord, StandaloneStoreError> {
|
||||||
|
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<StandaloneSessionRecord, StandaloneStoreError> {
|
||||||
|
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<Path>,
|
||||||
|
scope: StandaloneListScope,
|
||||||
|
limit: usize,
|
||||||
|
) -> Result<Vec<StandaloneSessionRecord>, 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<StandaloneSessionLease, StandaloneStoreError> {
|
||||||
|
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<SegmentId>,
|
||||||
|
) -> Result<StandaloneSessionRecord, StandaloneStoreError> {
|
||||||
|
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<SegmentId>,
|
||||||
|
reason: StandaloneShutdownReason,
|
||||||
|
) -> Result<StandaloneSessionRecord, StandaloneStoreError> {
|
||||||
|
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<u64>,
|
||||||
|
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<StandaloneSessionRecord, StandaloneStoreError> {
|
||||||
|
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<Self, StandaloneStoreError> {
|
||||||
|
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<u64>,
|
||||||
|
acquired_at_unix_ms: u64,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl LeaseRecord {
|
||||||
|
fn current() -> Result<Self, StandaloneStoreError> {
|
||||||
|
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<LeaseRecord, StandaloneStoreError> {
|
||||||
|
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<u64> {
|
||||||
|
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<u64> {
|
||||||
|
(pid == std::process::id()).then_some(0)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn now_unix_ms() -> Result<u64, StandaloneStoreError> {
|
||||||
|
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),
|
||||||
|
}
|
||||||
@@ -10,7 +10,10 @@ use agen::llm_client::types::Request;
|
|||||||
use async_trait::async_trait;
|
use async_trait::async_trait;
|
||||||
use futures::{Stream, stream};
|
use futures::{Stream, stream};
|
||||||
use protocol::{Event, Method};
|
use protocol::{Event, Method};
|
||||||
use standalone::{StandaloneHost, StandaloneLaunchConfig};
|
use standalone::{
|
||||||
|
StaleLeasePolicy, StandaloneHost, StandaloneLaunchConfig, StandaloneListScope,
|
||||||
|
StandaloneSessionStatus, StandaloneSessionStore, StandaloneStartupError, StandaloneStoreError,
|
||||||
|
};
|
||||||
use uuid::Uuid;
|
use uuid::Uuid;
|
||||||
|
|
||||||
#[derive(Clone)]
|
#[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, standalone::StandaloneStartupError::StateStore);
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
error.to_string(),
|
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!(!error.to_string().contains("secret-name"));
|
||||||
assert!(
|
assert!(
|
||||||
@@ -210,3 +213,261 @@ fn launch_rejects_path_profile_before_worker_startup() {
|
|||||||
standalone::StandaloneLaunchError::PathProfileUnsupported
|
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(())
|
||||||
|
}
|
||||||
|
|||||||
@@ -130,6 +130,29 @@ where
|
|||||||
Ok(PreparedWorker::new(worker, self.layout, self.transport))
|
Ok(PreparedWorker::new(worker, self.layout, self.transport))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub async fn prepare_restored(
|
||||||
|
self,
|
||||||
|
worker_name: &str,
|
||||||
|
) -> Result<PreparedWorker<Box<dyn LlmClient>, 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<BootstrappedWorker, WorkerBootstrapError> {
|
pub async fn start(self) -> Result<BootstrappedWorker, WorkerBootstrapError> {
|
||||||
self.prepare().await?.start().await
|
self.prepare().await?.start().await
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -5515,6 +5515,27 @@ where
|
|||||||
loader: PromptCatalogSource,
|
loader: PromptCatalogSource,
|
||||||
workspace_context: WorkerWorkspaceContext,
|
workspace_context: WorkerWorkspaceContext,
|
||||||
filesystem_authority: WorkerFilesystemAuthority,
|
filesystem_authority: WorkerFilesystemAuthority,
|
||||||
|
) -> Result<Self, WorkerError> {
|
||||||
|
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<Box<dyn LlmClient>>,
|
||||||
) -> Result<Self, WorkerError> {
|
) -> Result<Self, WorkerError> {
|
||||||
let metadata =
|
let metadata =
|
||||||
store
|
store
|
||||||
@@ -5535,7 +5556,7 @@ where
|
|||||||
worker_name: worker_name.to_string(),
|
worker_name: worker_name.to_string(),
|
||||||
})?;
|
})?;
|
||||||
if let Some(segment_id) = active.segment_id {
|
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,
|
active.session_id,
|
||||||
segment_id,
|
segment_id,
|
||||||
restore_manifest_from_worker_metadata_snapshot(
|
restore_manifest_from_worker_metadata_snapshot(
|
||||||
@@ -5547,6 +5568,7 @@ where
|
|||||||
loader,
|
loader,
|
||||||
workspace_context,
|
workspace_context,
|
||||||
filesystem_authority,
|
filesystem_authority,
|
||||||
|
model_client,
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
}
|
}
|
||||||
@@ -5557,12 +5579,13 @@ where
|
|||||||
})?;
|
})?;
|
||||||
let manifest =
|
let manifest =
|
||||||
restore_manifest_from_worker_metadata_snapshot(worker_name, Some(snapshot), fallback)?;
|
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,
|
manifest,
|
||||||
store,
|
store,
|
||||||
loader,
|
loader,
|
||||||
workspace_context,
|
workspace_context,
|
||||||
filesystem_authority,
|
filesystem_authority,
|
||||||
|
model_client,
|
||||||
)
|
)
|
||||||
.await
|
.await
|
||||||
}
|
}
|
||||||
@@ -5614,6 +5637,29 @@ where
|
|||||||
loader: PromptCatalogSource,
|
loader: PromptCatalogSource,
|
||||||
workspace_context: WorkerWorkspaceContext,
|
workspace_context: WorkerWorkspaceContext,
|
||||||
filesystem_authority: WorkerFilesystemAuthority,
|
filesystem_authority: WorkerFilesystemAuthority,
|
||||||
|
) -> Result<Self, WorkerError> {
|
||||||
|
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<Box<dyn LlmClient>>,
|
||||||
) -> Result<Self, WorkerError> {
|
) -> Result<Self, WorkerError> {
|
||||||
// Read raw entries once so we can both reconstruct state and
|
// Read raw entries once so we can both reconstruct state and
|
||||||
// seed the broadcast sink's mirror with the same prefix that
|
// seed the broadcast sink's mirror with the same prefix that
|
||||||
@@ -5629,13 +5675,14 @@ where
|
|||||||
let mirror_entries: Vec<LogEntry> = raw_entries.clone();
|
let mirror_entries: Vec<LogEntry> = raw_entries.clone();
|
||||||
let scope_config = effective_restore_scope_config(&store, &manifest)?;
|
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,
|
&manifest,
|
||||||
&loader,
|
&loader,
|
||||||
/* parse_template */ false,
|
/* parse_template */ false,
|
||||||
workspace_context,
|
workspace_context,
|
||||||
filesystem_authority,
|
filesystem_authority,
|
||||||
scope_config,
|
scope_config,
|
||||||
|
model_client,
|
||||||
)?;
|
)?;
|
||||||
|
|
||||||
// Atomic: register_worker inside install_top_level rejects when
|
// Atomic: register_worker inside install_top_level rejects when
|
||||||
|
|||||||
Reference in New Issue
Block a user