Hare/yoi
1
1
Fork 0
Code Actions Packages Releases Activity

runtime: canonicalize Worker aggregates

Browse Source
This commit is contained in:
Hare
2026-08-11 22:38:29 +09:00
parent 33d3d1a43f
commit 34f8949e85
16 changed files with 2013 additions and 201 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/session-store/src/lib.rs
+5 -3
View File
@@ -38,6 +38,7 @@ pub mod segment_log;
pub mod store;
pub mod system_item;
pub mod worker_metadata;
pub mod worker_session_store;
pub use event_trace::{TraceEntry, TracePayload};
pub use fs_store::FsStore;
@@ -54,10 +55,11 @@ pub use segment_log::{LogEntry, RestoredState, SegmentOrigin, SessionExtension,
pub use store::{Store, StoreError};
pub use system_item::{SystemItem, SystemReminder, SystemReminderSource, render_worker_event};
pub use worker_metadata::{
CombinedStore, FsWorkerStore, WorkerActiveSegmentRef, WorkerMetadata, WorkerMetadataStore,
WorkerPeer, WorkerReclaimedChild, WorkerSpawnedChild, WorkerSpawnedScopeRule, WorkerStoreError,
validate_worker_name,
CombinedStore, FsWorkerStore, WorkerActiveSegmentRef, WorkerAggregateStore, WorkerMetadata,
WorkerMetadataStore, WorkerPeer, WorkerReclaimedChild, WorkerSpawnedChild,
WorkerSpawnedScopeRule, WorkerStoreError, validate_worker_name,
};
pub use worker_session_store::WorkerSessionStore;
/// Session identifier — the fork-tree root. UUID v7 (time-ordered).
///
crates/session-store/src/worker_metadata.rs
+130
View File
@@ -319,6 +319,115 @@ pub trait WorkerMetadataStore: Send + Sync {
}
}
/// Metadata store for one canonical Worker aggregate.
///
/// The Worker key is fixed by Runtime authority and metadata lives directly at
/// `<worker-aggregate>/metadata.json`. This store cannot address another
/// Worker's global metadata root.
#[derive(Clone)]
pub struct WorkerAggregateStore {
root: PathBuf,
worker_name: String,
}
impl WorkerAggregateStore {
pub fn new(
root: impl Into<PathBuf>,
worker_name: impl Into<String>,
) -> Result<Self, WorkerStoreError> {
let root = root.into();
let worker_name = worker_name.into();
validate_worker_name(&worker_name)?;
fs::create_dir_all(&root)?;
Ok(Self { root, worker_name })
}
fn validate_name(&self, worker_name: &str) -> Result<(), WorkerStoreError> {
validate_worker_name(worker_name)?;
if worker_name == self.worker_name {
Ok(())
} else {
Err(WorkerStoreError::InvalidWorkerName(format!(
"aggregate owns `{}`; requested `{worker_name}`",
self.worker_name
)))
}
}
fn metadata_path(&self) -> PathBuf {
self.root.join("metadata.json")
}
}
impl WorkerMetadataStore for WorkerAggregateStore {
fn write(&self, metadata: &WorkerMetadata) -> Result<(), WorkerStoreError> {
self.validate_name(&metadata.worker_name)?;
let mut content = serde_json::to_vec_pretty(metadata)?;
content.push(b'\n');
let path = self.metadata_path();
let temp = self.root.join(format!(
".metadata.json.tmp-{}-{}",
std::process::id(),
uuid::Uuid::now_v7()
));
let result = (|| -> Result<(), WorkerStoreError> {
use std::io::Write;
let mut file = std::fs::OpenOptions::new()
.write(true)
.create_new(true)
.open(&temp)?;
file.write_all(&content)?;
file.sync_all()?;
drop(file);
fs::rename(&temp, &path)?;
std::fs::File::open(&self.root)?.sync_all()?;
Ok(())
})();
if result.is_err() {
let _ = fs::remove_file(temp);
}
result
}
fn read_by_name(&self, worker_name: &str) -> Result<Option<WorkerMetadata>, WorkerStoreError> {
self.validate_name(worker_name)?;
let content = match fs::read(self.metadata_path()) {
Ok(content) => content,
Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(None),
Err(error) => return Err(error.into()),
};
let metadata: WorkerMetadata = serde_json::from_slice(&content)?;
if metadata.worker_name != self.worker_name {
return Err(WorkerStoreError::InvalidWorkerName(format!(
"aggregate identity mismatch: expected `{}`, found `{}`",
self.worker_name, metadata.worker_name
)));
}
Ok(Some(metadata))
}
fn list_names(&self) -> Result<Vec<String>, WorkerStoreError> {
Ok(if self.metadata_path().is_file() {
vec![self.worker_name.clone()]
} else {
Vec::new()
})
}
fn root_dir(&self) -> Option<PathBuf> {
Some(self.root.clone())
}
fn delete_by_name(&self, worker_name: &str) -> Result<(), WorkerStoreError> {
self.validate_name(worker_name)?;
match fs::remove_file(self.metadata_path()) {
Ok(()) => Ok(()),
Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
Err(error) => Err(error.into()),
}
}
}
/// Filesystem-backed Worker metadata store.
#[derive(Clone)]
pub struct FsWorkerStore {
@@ -555,6 +664,27 @@ mod tests {
assert_eq!(restored, metadata);
}
#[test]
fn worker_aggregate_store_writes_one_fixed_metadata_identity() {
let tmp = tempfile::tempdir().unwrap();
let store = WorkerAggregateStore::new(tmp.path(), "worker-runtime-7").unwrap();
let metadata = WorkerMetadata::new("worker-runtime-7", None);
store.write(&metadata).unwrap();
assert!(tmp.path().join("metadata.json").is_file());
assert_eq!(store.list_names().unwrap(), vec!["worker-runtime-7"]);
assert_eq!(
store
.read_by_name("worker-runtime-7")
.unwrap()
.unwrap()
.worker_name,
"worker-runtime-7"
);
let error = store.read_by_name("worker-runtime-8").unwrap_err();
assert!(error.to_string().contains("aggregate owns"));
}
#[test]
fn fs_store_writes_under_worker_state_root_only() {
let tmp = tempfile::TempDir::new().unwrap();
crates/session-store/src/worker_session_store.rs
+422
View File
@@ -0,0 +1,422 @@
//! Filesystem store for the canonical `1 Worker = 1 Session` aggregate.
//!
//! Layout under one Worker aggregate:
//! - `session/session.json` — immutable Session identity
//! - `session/segments/<segment_id>.jsonl`
//! - `session/segments/<segment_id>.trace.jsonl`
//!
//! Unlike [`crate::FsStore`], this store cannot enumerate or switch between
//! arbitrary Sessions. The first segment materializes the sole Session identity;
//! every later operation must use that same ID.
use crate::event_trace::TraceEntry;
use crate::segment_log::LogEntry;
use crate::store::{Store, StoreError};
use crate::{SegmentId, SessionId};
use serde::{Deserialize, Serialize};
use std::fs::{self, File, OpenOptions};
use std::io::{Read, Seek, SeekFrom, Write};
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};
use std::time::SystemTime;
const SESSION_SCHEMA_VERSION: u32 = 1;
const SESSION_FILE: &str = "session.json";
const SEGMENTS_DIR: &str = "segments";
#[derive(Clone)]
pub struct WorkerSessionStore {
root: PathBuf,
session_id: Arc<Mutex<Option<SessionId>>>,
append_lock: Arc<Mutex<()>>,
}
#[derive(Debug, Serialize, Deserialize)]
struct SessionManifest {
schema_version: u32,
session_id: SessionId,
}
impl WorkerSessionStore {
/// Open the Session store rooted at `<worker-aggregate>/session`.
pub fn new(root: impl Into<PathBuf>) -> Result<Self, StoreError> {
let root = root.into();
fs::create_dir_all(root.join(SEGMENTS_DIR))?;
let session_id = match fs::read(root.join(SESSION_FILE)) {
Ok(bytes) => {
let manifest: SessionManifest = serde_json::from_slice(&bytes)?;
if manifest.schema_version != SESSION_SCHEMA_VERSION {
return Err(StoreError::Corrupt {
line: 0,
message: format!(
"unsupported Worker Session schema version {}, expected {}",
manifest.schema_version, SESSION_SCHEMA_VERSION
),
});
}
Some(manifest.session_id)
}
Err(error) if error.kind() == std::io::ErrorKind::NotFound => None,
Err(error) => return Err(error.into()),
};
Ok(Self {
root,
session_id: Arc::new(Mutex::new(session_id)),
append_lock: Arc::new(Mutex::new(())),
})
}
pub fn root_dir(&self) -> &Path {
&self.root
}
pub fn session_id(&self) -> Result<Option<SessionId>, StoreError> {
self.session_id
.lock()
.map(|session_id| *session_id)
.map_err(|_| std::io::Error::other("Worker Session identity lock was poisoned").into())
}
pub fn session_modified_at(&self) -> Result<Option<SystemTime>, StoreError> {
let metadata = match fs::metadata(&self.root) {
Ok(metadata) => metadata,
Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(None),
Err(error) => return Err(error.into()),
};
let mut latest = Some(metadata.modified()?);
for entry in fs::read_dir(self.root.join(SEGMENTS_DIR))? {
let modified = entry?.metadata()?.modified()?;
if latest.map(|current| modified > current).unwrap_or(true) {
latest = Some(modified);
}
}
Ok(latest)
}
fn ensure_session(&self, requested: SessionId, materialize: bool) -> Result<(), StoreError> {
let mut session_id = self
.session_id
.lock()
.map_err(|_| std::io::Error::other("Worker Session identity lock was poisoned"))?;
match *session_id {
Some(existing) if existing == requested => Ok(()),
Some(existing) => Err(StoreError::Corrupt {
line: 0,
message: format!(
"Worker aggregate owns Session {existing}; cannot attach or switch to Session {requested}"
),
}),
None if !materialize => Err(StoreError::Corrupt {
line: 0,
message: format!(
"Worker aggregate has no materialized Session; requested Session {requested}"
),
}),
None => {
let manifest = SessionManifest {
schema_version: SESSION_SCHEMA_VERSION,
session_id: requested,
};
atomic_write_json(&self.root.join(SESSION_FILE), &manifest)?;
*session_id = Some(requested);
Ok(())
}
}
}
fn log_path(&self, segment_id: SegmentId) -> PathBuf {
self.root
.join(SEGMENTS_DIR)
.join(format!("{segment_id}.jsonl"))
}
fn trace_path(&self, segment_id: SegmentId) -> PathBuf {
self.root
.join(SEGMENTS_DIR)
.join(format!("{segment_id}.trace.jsonl"))
}
fn append_line(&self, path: &Path, line: &str) -> Result<(), StoreError> {
let _guard = self
.append_lock
.lock()
.map_err(|_| std::io::Error::other("Worker Session append lock was poisoned"))?;
let mut file = OpenOptions::new()
.create(true)
.read(true)
.write(true)
.append(true)
.open(path)?;
let committed_len = truncate_uncommitted_tail(&mut file)?;
let mut record = Vec::with_capacity(line.len() + 1);
record.extend_from_slice(line.as_bytes());
record.push(b'\n');
if let Err(write_error) = file.write_all(&record) {
return match file.set_len(committed_len) {
Ok(()) => Err(write_error.into()),
Err(rollback_error) => Err(std::io::Error::new(
rollback_error.kind(),
format!(
"session append failed ({write_error}) and rollback failed: {rollback_error}"
),
)
.into()),
};
}
Ok(())
}
}
impl Store for WorkerSessionStore {
fn append(
&self,
session_id: SessionId,
segment_id: SegmentId,
entry: &LogEntry,
) -> Result<(), StoreError> {
self.ensure_session(session_id, true)?;
self.append_line(&self.log_path(segment_id), &serde_json::to_string(entry)?)
}
fn read_all(
&self,
session_id: SessionId,
segment_id: SegmentId,
) -> Result<Vec<LogEntry>, StoreError> {
self.ensure_session(session_id, false)?;
let path = self.log_path(segment_id);
if !path.exists() {
return Err(StoreError::NotFound(segment_id));
}
parse_jsonl(&fs::read(path)?)
}
fn list_sessions(&self) -> Result<Vec<SessionId>, StoreError> {
Ok(self.session_id()?.into_iter().collect())
}
fn list_segments(&self, session_id: SessionId) -> Result<Vec<SegmentId>, StoreError> {
self.ensure_session(session_id, false)?;
let mut segments: Vec<SegmentId> = Vec::new();
for entry in fs::read_dir(self.root.join(SEGMENTS_DIR))? {
let path = entry?.path();
let name = path
.file_name()
.and_then(|name| name.to_str())
.unwrap_or("");
if name.ends_with(".jsonl")
&& !name.ends_with(".trace.jsonl")
&& let Ok(segment_id) = name.trim_end_matches(".jsonl").parse()
{
segments.push(segment_id);
}
}
segments.sort_by(|left, right| right.cmp(left));
Ok(segments)
}
fn lookup_session_of(&self, segment_id: SegmentId) -> Result<Option<SessionId>, StoreError> {
let session_id = self.session_id()?;
Ok(session_id.filter(|_| self.log_path(segment_id).exists()))
}
fn create_segment(
&self,
session_id: SessionId,
segment_id: SegmentId,
entries: &[LogEntry],
) -> Result<(), StoreError> {
self.ensure_session(session_id, true)?;
let mut content = Vec::new();
for entry in entries {
serde_json::to_writer(&mut content, entry)?;
content.push(b'\n');
}
atomic_write_bytes(&self.log_path(segment_id), &content)?;
Ok(())
}
fn exists(&self, session_id: SessionId, segment_id: SegmentId) -> Result<bool, StoreError> {
self.ensure_session(session_id, false)?;
Ok(self.log_path(segment_id).exists())
}
fn read_entry_count(
&self,
session_id: SessionId,
segment_id: SegmentId,
) -> Result<usize, StoreError> {
self.ensure_session(session_id, false)?;
let path = self.log_path(segment_id);
if !path.exists() {
return Err(StoreError::NotFound(segment_id));
}
let content = fs::read(path)?;
let complete = complete_jsonl_prefix(&content);
let complete = std::str::from_utf8(complete).map_err(|error| StoreError::Corrupt {
line: complete[..error.valid_up_to()]
.iter()
.filter(|byte| **byte == b'\n')
.count()
+ 1,
message: error.to_string(),
})?;
Ok(complete
.lines()
.filter(|line| !line.trim().is_empty())
.count())
}
fn append_trace(
&self,
session_id: SessionId,
segment_id: SegmentId,
entry: &TraceEntry,
) -> Result<(), StoreError> {
self.ensure_session(session_id, true)?;
self.append_line(&self.trace_path(segment_id), &serde_json::to_string(entry)?)
}
}
fn atomic_write_json<T: Serialize>(path: &Path, value: &T) -> Result<(), StoreError> {
let mut bytes = serde_json::to_vec_pretty(value)?;
bytes.push(b'\n');
atomic_write_bytes(path, &bytes)
}
fn atomic_write_bytes(path: &Path, bytes: &[u8]) -> Result<(), StoreError> {
let parent = path
.parent()
.ok_or_else(|| std::io::Error::other("Worker Session path has no parent"))?;
fs::create_dir_all(parent)?;
let tmp = path.with_file_name(format!(
".{}.tmp-{}-{}",
path.file_name()
.and_then(|name| name.to_str())
.unwrap_or("session"),
std::process::id(),
uuid::Uuid::now_v7()
));
let result = (|| -> Result<(), StoreError> {
let mut file = OpenOptions::new().create_new(true).write(true).open(&tmp)?;
file.write_all(bytes)?;
file.sync_all()?;
drop(file);
fs::rename(&tmp, path)?;
File::open(parent)?.sync_all()?;
Ok(())
})();
if result.is_err() {
let _ = fs::remove_file(&tmp);
}
result
}
fn complete_jsonl_prefix(content: &[u8]) -> &[u8] {
if content.last() == Some(&b'\n') {
return content;
}
content
.iter()
.rposition(|byte| *byte == b'\n')
.map(|index| &content[..=index])
.unwrap_or(&[])
}
fn parse_jsonl<T: serde::de::DeserializeOwned>(content: &[u8]) -> Result<Vec<T>, StoreError> {
let complete = complete_jsonl_prefix(content);
let content = std::str::from_utf8(complete).map_err(|error| StoreError::Corrupt {
line: complete[..error.valid_up_to()]
.iter()
.filter(|byte| **byte == b'\n')
.count()
+ 1,
message: error.to_string(),
})?;
content
.lines()
.enumerate()
.filter(|(_, line)| !line.trim().is_empty())
.map(|(index, line)| {
serde_json::from_str(line).map_err(|error| StoreError::Corrupt {
line: index + 1,
message: error.to_string(),
})
})
.collect()
}
fn truncate_uncommitted_tail(file: &mut File) -> std::io::Result<u64> {
const SCAN_BYTES: usize = 8 * 1024;
let len = file.metadata()?.len();
if len == 0 {
return Ok(0);
}
file.seek(SeekFrom::End(-1))?;
let mut last = [0_u8; 1];
file.read_exact(&mut last)?;
if last[0] == b'\n' {
return Ok(len);
}
let mut end = len;
let mut buffer = [0_u8; SCAN_BYTES];
while end > 0 {
let start = end.saturating_sub(SCAN_BYTES as u64);
let chunk_len = (end - start) as usize;
file.seek(SeekFrom::Start(start))?;
file.read_exact(&mut buffer[..chunk_len])?;
if let Some(index) = buffer[..chunk_len].iter().rposition(|byte| *byte == b'\n') {
let committed_len = start + index as u64 + 1;
file.set_len(committed_len)?;
return Ok(committed_len);
}
end = start;
}
file.set_len(0)?;
Ok(0)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{Store, new_segment_id, new_session_id};
#[test]
fn canonical_layout_and_single_session_invariant() {
let root = tempfile::tempdir().unwrap();
let store = WorkerSessionStore::new(root.path().join("session")).unwrap();
let session_id = new_session_id();
let segment_id = new_segment_id();
store.create_segment(session_id, segment_id, &[]).unwrap();
assert!(root.path().join("session/session.json").is_file());
assert!(
root.path()
.join(format!("session/segments/{segment_id}.jsonl"))
.is_file()
);
assert_eq!(store.list_sessions().unwrap(), vec![session_id]);
let other = new_session_id();
let error = store
.create_segment(other, new_segment_id(), &[])
.unwrap_err();
assert!(error.to_string().contains("cannot attach or switch"));
assert_eq!(store.list_sessions().unwrap(), vec![session_id]);
}
#[test]
fn reopen_preserves_session_and_segment_ids() {
let root = tempfile::tempdir().unwrap();
let session_id = new_session_id();
let segment_id = new_segment_id();
WorkerSessionStore::new(root.path())
.unwrap()
.create_segment(session_id, segment_id, &[])
.unwrap();
let reopened = WorkerSessionStore::new(root.path()).unwrap();
assert_eq!(reopened.session_id().unwrap(), Some(session_id));
assert!(reopened.exists(session_id, segment_id).unwrap());
}
}
crates/tui/src/ui.rs
+1 -1
View File
@@ -1524,7 +1524,7 @@ fn render_compact(lines: &mut Vec<Line<'static>>, evt: &CompactEvent, width: u16
.collect::<String>();
let elapsed = elapsed_suffix(*elapsed_secs);
(
format!("[compact] done (new session {short}){elapsed}"),
format!("[compact] done (new segment {short}){elapsed}"),
MessageKind::NoticeWarn,
)
}
crates/worker-runtime/src/execution.rs
+4
View File
@@ -245,6 +245,8 @@ impl fmt::Debug for WorkerExecutionContext {
#[derive(Clone, Debug)]
pub struct WorkerExecutionSpawnRequest {
pub worker_ref: WorkerRef,
/// Monotonic execution generation reserved durably before launch.
pub run_generation: u64,
pub request: crate::catalog::CreateWorkerRequest,
pub context: WorkerExecutionContext,
pub working_directory: Option<WorkingDirectoryBinding>,
@@ -255,6 +257,8 @@ pub struct WorkerExecutionSpawnRequest {
#[derive(Clone, Debug)]
pub struct WorkerExecutionRestoreRequest {
pub worker_ref: WorkerRef,
/// Monotonic execution generation reserved durably before restore.
pub run_generation: u64,
pub request: crate::catalog::CreateWorkerRequest,
pub context: WorkerExecutionContext,
pub previous_working_directory: Option<WorkingDirectoryStatus>,
crates/worker-runtime/src/fs_store.rs
+1062 -1
View File
File diff suppressed because it is too large Load Diff
crates/worker-runtime/src/main.rs
+13 -3
View File
@@ -18,7 +18,7 @@ use worker_runtime::auth::{
RuntimeHttpAuthConfig, RuntimeIdentityMaterial, TrustedServerKey, decode_public_key,
};
use worker_runtime::error::RuntimeError;
use worker_runtime::fs_store::FsRuntimeStoreOptions;
use worker_runtime::fs_store::{FsRuntimeStore, FsRuntimeStoreOptions};
use worker_runtime::http_server::{
RuntimeHttpServerConfig, RuntimeHttpServerError, RuntimeHttpStoreSelection,
};
@@ -80,9 +80,13 @@ fn run() -> Result<(), ProcessError> {
fn build_runtime(config: &ProcessConfig) -> Result<Runtime, ProcessError> {
let fs_paths = config.resolved_fs_paths();
let runtime_store_dir = match &config.http.store {
RuntimeHttpStoreSelection::Memory => fs_paths.runtime_dir.clone(),
RuntimeHttpStoreSelection::Fs { root } => root.clone(),
_ => fs_paths.runtime_dir.clone(),
};
let mut factory = ProfileRuntimeWorkerFactory::new(fs_paths.worker_dir.join("worker-root"))
.with_store_dir(fs_paths.worker_dir.join("sessions"))
.with_worker_metadata_dir(fs_paths.worker_dir.join("metadata"));
.with_runtime_store_dir(runtime_store_dir);
if let Some(endpoint) = config.backend_resource_endpoint.clone() {
factory = factory.with_resource_client(Arc::new(
worker_runtime::resource::HttpBackendResourceClient::new(
@@ -105,6 +109,12 @@ fn build_runtime(config: &ProcessConfig) -> Result<Runtime, ProcessError> {
.map_err(ProcessError::Runtime)
}
RuntimeHttpStoreSelection::Fs { root } => {
FsRuntimeStore::migrate_legacy_worker_aggregates(
root,
fs_paths.worker_dir.join("sessions"),
fs_paths.worker_dir.join("metadata"),
)
.map_err(ProcessError::Runtime)?;
let mut options = FsRuntimeStoreOptions::new(root.clone());
options.display_name = config.http.display_name.clone();
Runtime::with_fs_store_and_execution_backend(options, backend)
crates/worker-runtime/src/runtime.rs
+88 -29
View File
@@ -532,12 +532,16 @@ impl Runtime {
status: WorkerStatus::Stopped,
workspace_id: scope.map(|scope| scope.workspace_id.clone()),
request: request.clone(),
run_generation: 1,
working_directory: None,
execution_handle: None,
};
state.workers.insert(worker_id, record);
state.persist_runtime_snapshot()?;
state.persist_worker(&worker_ref.worker_id)?;
let spawn_request = WorkerExecutionSpawnRequest {
worker_ref: worker_ref.clone(),
run_generation: 1,
request,
context: self.execution_context(worker_ref.clone()),
working_directory: None,
@@ -859,35 +863,46 @@ impl Runtime {
/// present this is idempotent; otherwise the configured backend is tried.
pub fn restore_worker(&self, worker_ref: &WorkerRef) -> Result<WorkerDetail, RuntimeError> {
let (backend, request) = {
let state = self.lock()?;
let mut state = self.lock()?;
state.ensure_running()?;
let worker = state.worker(worker_ref)?;
if worker.execution_handle.is_some() {
return Ok(worker.detail());
}
if worker.status == WorkerStatus::Cancelled {
return Err(RuntimeError::InvalidRequest(format!(
"worker {} is cancelled",
worker_ref.worker_id
)));
}
let (worker_request, previous_working_directory, config_bundle, run_generation) = {
let worker = state.worker(worker_ref)?;
if worker.execution_handle.is_some() {
return Ok(worker.detail());
}
if worker.status == WorkerStatus::Cancelled {
return Err(RuntimeError::InvalidRequest(format!(
"worker {} is cancelled",
worker_ref.worker_id
)));
}
let config_bundle = worker
.request
.config_bundle
.as_ref()
.and_then(|bundle_ref| state.config_bundles.get(&bundle_ref.id))
.cloned();
(
worker.request.clone(),
worker.working_directory.clone(),
config_bundle,
worker.run_generation.saturating_add(1).max(1),
)
};
let backend = state.execution_backend.clone().ok_or_else(|| {
RuntimeError::WorkerExecutionUnavailable {
worker_id: worker_ref.worker_id.clone(),
message: "runtime has no execution backend".to_string(),
}
})?;
let config_bundle = worker
.request
.config_bundle
.as_ref()
.and_then(|bundle_ref| state.config_bundles.get(&bundle_ref.id))
.cloned();
state.worker_mut(worker_ref)?.run_generation = run_generation;
state.persist_worker(&worker_ref.worker_id)?;
let request = WorkerExecutionRestoreRequest {
worker_ref: worker_ref.clone(),
request: worker.request.clone(),
run_generation,
request: worker_request,
context: self.execution_context(worker_ref.clone()),
previous_working_directory: worker.working_directory.clone(),
previous_working_directory,
working_directory: None,
config_bundle,
};
@@ -1514,34 +1529,64 @@ impl Runtime {
struct RestoreCandidate {
worker_ref: WorkerRef,
request: CreateWorkerRequest,
run_generation: u64,
previous_working_directory: Option<CatalogWorkingDirectoryStatus>,
config_bundle: Option<ConfigBundle>,
}
let candidates = {
let state = self.lock()?;
let mut state = self.lock()?;
if state.execution_backend.is_none() {
return Ok(());
}
state
let worker_ids = state
.workers
.values()
.filter(|worker| worker.execution_handle.is_none())
.map(|worker| {
.map(|worker| worker.worker_id)
.collect::<Vec<_>>();
let mut candidates = Vec::with_capacity(worker_ids.len());
for worker_id in worker_ids {
let (
worker_ref,
request,
previous_working_directory,
config_bundle,
run_generation,
) = {
let worker = state
.workers
.get(&worker_id)
.expect("collected Worker exists");
let config_bundle = worker
.request
.config_bundle
.as_ref()
.and_then(|bundle_ref| state.config_bundles.get(&bundle_ref.id))
.cloned();
RestoreCandidate {
worker_ref: worker.worker_ref.clone(),
request: worker.request.clone(),
previous_working_directory: worker.working_directory.clone(),
(
worker.worker_ref.clone(),
worker.request.clone(),
worker.working_directory.clone(),
config_bundle,
}
})
.collect::<Vec<_>>()
worker.run_generation.saturating_add(1).max(1),
)
};
state
.workers
.get_mut(&worker_id)
.expect("collected Worker exists")
.run_generation = run_generation;
state.persist_worker(&worker_id)?;
candidates.push(RestoreCandidate {
worker_ref,
request,
run_generation,
previous_working_directory,
config_bundle,
});
}
candidates
};
for candidate in candidates {
@@ -1554,6 +1599,7 @@ impl Runtime {
};
let request = WorkerExecutionRestoreRequest {
worker_ref: candidate.worker_ref.clone(),
run_generation: candidate.run_generation,
request: candidate.request,
context: self.execution_context(candidate.worker_ref.clone()),
previous_working_directory: candidate.previous_working_directory,
@@ -1723,6 +1769,7 @@ impl RuntimeState {
status: WorkerStatus::Stopped,
workspace_id: worker.workspace_id,
request: worker.request,
run_generation: worker.run_generation,
working_directory: worker.working_directory,
execution_handle: None,
},
@@ -2281,6 +2328,7 @@ struct WorkerRecord {
status: WorkerStatus,
workspace_id: Option<String>,
request: CreateWorkerRequest,
run_generation: u64,
working_directory: Option<CatalogWorkingDirectoryStatus>,
execution_handle: Option<WorkerExecutionHandle>,
}
@@ -2324,6 +2372,7 @@ impl WorkerRecord {
worker_ref: self.worker_ref.clone(),
worker_id: self.worker_id.clone(),
request: self.request.clone(),
run_generation: self.run_generation,
workspace_id: self.workspace_id.clone(),
working_directory: self.working_directory.clone(),
}
@@ -2627,6 +2676,7 @@ mod tests {
dispatch_result: Mutex<Option<WorkerExecutionResult>>,
restore_result: Mutex<Option<WorkerExecutionSpawnResult>>,
restore_count: Mutex<u64>,
run_generations: Mutex<Vec<u64>>,
contexts: Mutex<BTreeMap<WorkerId, WorkerExecutionContext>>,
dispatched_inputs: Mutex<Vec<WorkerInput>>,
preserve_commit_ack_submission_id: AtomicBool,
@@ -2670,6 +2720,10 @@ mod tests {
}
fn spawn_worker(&self, request: WorkerExecutionSpawnRequest) -> WorkerExecutionSpawnResult {
self.run_generations
.lock()
.unwrap()
.push(request.run_generation);
self.contexts
.lock()
.unwrap()
@@ -2689,6 +2743,10 @@ mod tests {
request: WorkerExecutionRestoreRequest,
) -> WorkerExecutionSpawnResult {
*self.restore_count.lock().unwrap() += 1;
self.run_generations
.lock()
.unwrap()
.push(request.run_generation);
if let Some(result) = self.restore_result.lock().unwrap().clone() {
return result;
}
@@ -3551,6 +3609,7 @@ mod tests {
.unwrap();
assert_eq!(*backend.restore_count.lock().unwrap(), 1);
assert_eq!(*backend.run_generations.lock().unwrap(), vec![1, 2]);
assert_eq!(
runtime.worker_detail(&detail.worker_ref).unwrap().status,
WorkerStatus::Idle
crates/worker-runtime/src/worker_backend.rs
+118 -126
View File
@@ -10,7 +10,7 @@
use std::collections::HashMap;
use std::future::Future;
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex, mpsc};
use std::time::Duration;
@@ -30,9 +30,12 @@ use crate::working_directory::{
WorkingDirectoryBinding, WorkingDirectoryDiagnostic, WorkingDirectoryMaterializer,
};
use async_trait::async_trait;
use manifest::paths;
use protocol::{Event, Method, Segment, WorkerStatus};
use session_store::{CombinedStore, FsStore, FsWorkerStore, LogEntry, collect_state};
use session_store::{
CombinedStore, LogEntry, WorkerAggregateStore, WorkerSessionStore, collect_state,
};
#[cfg(test)]
use session_store::{FsStore, FsWorkerStore};
use tokio::runtime::Runtime;
#[cfg(feature = "ws-server")]
use tokio::sync::broadcast;
@@ -57,7 +60,6 @@ const RUNTIME_TASK_TIMEOUT: Duration = Duration::from_secs(10);
// Keep this below the adapter task timeout so a failed acknowledgement task
// returns a typed execution error instead of leaving the outer waiter to time out.
const USER_INPUT_COMMIT_TIMEOUT: Duration = Duration::from_secs(9);
static NEXT_RUNTIME_ARTIFACT_ROOT: AtomicU64 = AtomicU64::new(1);
fn user_input_has_submission(entry: &LogEntry, submission_id: &str) -> bool {
let LogEntry::UserInput { extensions, .. } = entry else {
@@ -69,43 +71,9 @@ fn user_input_has_submission(entry: &LogEntry, submission_id: &str) -> bool {
})
}
#[derive(Clone)]
enum RuntimeArtifactRoot {
Owned(Arc<OwnedRuntimeArtifactRoot>),
External(PathBuf),
}
impl RuntimeArtifactRoot {
fn owned() -> Self {
let sequence = NEXT_RUNTIME_ARTIFACT_ROOT.fetch_add(1, Ordering::Relaxed);
Self::Owned(Arc::new(OwnedRuntimeArtifactRoot {
path: std::env::temp_dir().join(format!(
"yoi-worker-runtime-artifacts-{}-{sequence}",
std::process::id()
)),
}))
}
fn path(&self) -> &std::path::Path {
match self {
Self::Owned(root) => &root.path,
Self::External(path) => path,
}
}
}
struct OwnedRuntimeArtifactRoot {
path: PathBuf,
}
impl Drop for OwnedRuntimeArtifactRoot {
fn drop(&mut self) {
let _ = std::fs::remove_dir_all(&self.path);
}
}
pub struct RuntimeWorkerController {
pub handle: WorkerHandle,
pub shutdown: Arc<tokio::sync::Mutex<Option<worker::ShutdownReceiver>>>,
pub workspace_client: Arc<dyn WorkspaceClient>,
}
@@ -244,9 +212,7 @@ impl WorkerObservationProvider for RuntimeGrantedWorkerObservationProvider {
pub struct ProfileRuntimeWorkerFactory {
observation_hub: Arc<RuntimeWorkerObservationHub>,
profile_base_dir: PathBuf,
store_dir: Option<PathBuf>,
worker_metadata_dir: Option<PathBuf>,
runtime_base_dir: RuntimeArtifactRoot,
worker_aggregate_root: Option<PathBuf>,
resource_client: Option<Arc<dyn BackendResourceClient>>,
profile_archive_cache: Arc<ProfileSourceArchiveCache>,
}
@@ -257,26 +223,14 @@ impl ProfileRuntimeWorkerFactory {
Self {
observation_hub: Arc::new(RuntimeWorkerObservationHub::default()),
profile_base_dir,
store_dir: None,
worker_metadata_dir: None,
runtime_base_dir: RuntimeArtifactRoot::owned(),
worker_aggregate_root: None,
resource_client: None,
profile_archive_cache: Arc::new(ProfileSourceArchiveCache::default()),
}
}
pub fn with_store_dir(mut self, store_dir: impl Into<PathBuf>) -> Self {
self.store_dir = Some(store_dir.into());
self
}
pub fn with_worker_metadata_dir(mut self, worker_metadata_dir: impl Into<PathBuf>) -> Self {
self.worker_metadata_dir = Some(worker_metadata_dir.into());
self
}
pub fn with_runtime_base_dir(mut self, runtime_base_dir: impl Into<PathBuf>) -> Self {
self.runtime_base_dir = RuntimeArtifactRoot::External(runtime_base_dir.into());
pub fn with_runtime_store_dir(mut self, runtime_store_dir: impl Into<PathBuf>) -> Self {
self.worker_aggregate_root = Some(runtime_store_dir.into().join("workers"));
self
}
@@ -285,28 +239,16 @@ impl ProfileRuntimeWorkerFactory {
self
}
fn store_dir(&self) -> Result<PathBuf, String> {
self.store_dir
.clone()
.or_else(paths::sessions_dir)
fn worker_aggregate_dir(&self, worker_ref: &WorkerRef) -> Result<PathBuf, String> {
self.worker_aggregate_root
.as_ref()
.map(|root| root.join(worker_ref.worker_id.to_string()))
.ok_or_else(|| {
"could not resolve sessions directory (set YOI_DATA_DIR, YOI_HOME, XDG_DATA_HOME, or HOME)"
"Runtime Worker aggregate root is not configured; global Session/metadata roots are migration-only"
.to_string()
})
}
fn worker_metadata_dir(&self, store_dir: &std::path::Path) -> PathBuf {
self.worker_metadata_dir
.clone()
.or_else(|| paths::data_dir().map(|data_dir| data_dir.join("workers")))
.or_else(|| store_dir.parent().map(|parent| parent.join("workers")))
.unwrap_or_else(|| PathBuf::from("workers"))
}
fn runtime_base_dir(&self) -> Result<PathBuf, String> {
Ok(self.runtime_base_dir.path().to_path_buf())
}
fn runtime_worker_name_for_ref(worker_ref: &crate::identity::WorkerRef) -> String {
format!("worker-runtime-{}", worker_ref.worker_id)
}
@@ -575,20 +517,23 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory {
};
let flow_transition_enabled = manifest.feature.flow.enabled;
let store_dir = self.store_dir()?;
let session_store = FsStore::new(&store_dir).map_err(|err| {
let worker_aggregate_dir = self.worker_aggregate_dir(&request.worker_ref)?;
let session_dir = worker_aggregate_dir.join("session");
let session_store = WorkerSessionStore::new(&session_dir).map_err(|err| {
format!(
"failed to initialize session store at {}: {err}",
store_dir.display()
)
})?;
let worker_metadata_dir = self.worker_metadata_dir(&store_dir);
let worker_metadata_store = FsWorkerStore::new(&worker_metadata_dir).map_err(|err| {
format!(
"failed to initialize worker metadata store at {}: {err}",
worker_metadata_dir.display()
"failed to initialize canonical Worker Session store at {}: {err}",
session_dir.display()
)
})?;
let worker_metadata_store =
WorkerAggregateStore::new(&worker_aggregate_dir, worker_name.clone()).map_err(
|err| {
format!(
"failed to initialize canonical Worker metadata store at {}: {err}",
worker_aggregate_dir.display()
)
},
)?;
let store = CombinedStore::new(session_store, worker_metadata_store);
let mut worker = Worker::from_manifest_with_context(
@@ -642,10 +587,17 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory {
}
let workspace_client = worker.workspace_client_handle();
let runtime_base = self.runtime_base_dir()?;
let (handle, _shutdown_rx) = WorkerController::spawn_runtime_managed(worker, &runtime_base)
let run_dir = worker_aggregate_dir
.join("runs")
.join(request.run_generation.to_string());
let (handle, shutdown_rx) = WorkerController::spawn_runtime_managed_run(worker, &run_dir)
.await
.map_err(|err| format!("failed to spawn Worker controller: {err}"))?;
.map_err(|err| {
format!(
"failed to spawn Worker controller in {}: {err}",
run_dir.display()
)
})?;
if flow_transition_enabled {
handle.shared_state.enable_flow_transition();
}
@@ -656,6 +608,7 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory {
);
Ok(RuntimeWorkerController {
handle,
shutdown: Arc::new(tokio::sync::Mutex::new(Some(shutdown_rx))),
workspace_client,
})
}
@@ -692,26 +645,29 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory {
let workspace_context = workspace_backend_ref.worker_context(&request.worker_ref);
let (manifest, loader) = Self::restore_fallback_manifest(&worker_name)?;
let store_dir = self.store_dir()?;
let session_store = FsStore::new(&store_dir).map_err(|err| {
let worker_aggregate_dir = self.worker_aggregate_dir(&request.worker_ref)?;
let session_dir = worker_aggregate_dir.join("session");
let session_store = WorkerSessionStore::new(&session_dir).map_err(|err| {
format!(
"failed to initialize session store at {}: {err}",
store_dir.display()
)
})?;
let worker_metadata_dir = self.worker_metadata_dir(&store_dir);
let worker_metadata_store = FsWorkerStore::new(&worker_metadata_dir).map_err(|err| {
format!(
"failed to initialize worker metadata store at {}: {err}",
worker_metadata_dir.display()
"failed to initialize canonical Worker Session store at {}: {err}",
session_dir.display()
)
})?;
let worker_metadata_store =
WorkerAggregateStore::new(&worker_aggregate_dir, worker_name.clone()).map_err(
|err| {
format!(
"failed to initialize canonical Worker metadata store at {}: {err}",
worker_aggregate_dir.display()
)
},
)?;
let store = CombinedStore::new(session_store, worker_metadata_store);
let mut worker = match Worker::restore_from_worker_metadata_with_context(
&worker_name,
manifest.clone(),
store,
store.clone(),
loader.clone(),
workspace_context.clone(),
filesystem_authority.clone(),
@@ -722,20 +678,6 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory {
Err(WorkerError::WorkerMetadataPending { .. })
if request.request.initial_input.is_none() =>
{
let session_store = FsStore::new(&store_dir).map_err(|err| {
format!(
"failed to initialize session store at {}: {err}",
store_dir.display()
)
})?;
let worker_metadata_store =
FsWorkerStore::new(&worker_metadata_dir).map_err(|err| {
format!(
"failed to initialize worker metadata store at {}: {err}",
worker_metadata_dir.display()
)
})?;
let store = CombinedStore::new(session_store, worker_metadata_store);
Worker::restore_pending_from_worker_metadata_with_context(
&worker_name,
manifest.clone(),
@@ -792,10 +734,17 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory {
}
let workspace_client = worker.workspace_client_handle();
let runtime_base = self.runtime_base_dir()?;
let (handle, _shutdown_rx) = WorkerController::spawn_runtime_managed(worker, &runtime_base)
let run_dir = worker_aggregate_dir
.join("runs")
.join(request.run_generation.to_string());
let (handle, shutdown_rx) = WorkerController::spawn_runtime_managed_run(worker, &run_dir)
.await
.map_err(|err| format!("failed to spawn restored Worker controller: {err}"))?;
.map_err(|err| {
format!(
"failed to spawn restored Worker controller in {}: {err}",
run_dir.display()
)
})?;
if flow_transition_enabled {
handle.shared_state.enable_flow_transition();
}
@@ -806,6 +755,7 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory {
);
Ok(RuntimeWorkerController {
handle,
shutdown: Arc::new(tokio::sync::Mutex::new(Some(shutdown_rx))),
workspace_client,
})
}
@@ -813,6 +763,7 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory {
struct RuntimeWorkerExecution {
handle: WorkerHandle,
shutdown: Arc<tokio::sync::Mutex<Option<worker::ShutdownReceiver>>>,
busy: Arc<AtomicBool>,
workspace_client: Option<Arc<dyn WorkspaceClient>>,
}
@@ -828,7 +779,10 @@ pub struct WorkerRuntimeExecutionBackend<F = ProfileRuntimeWorkerFactory> {
impl WorkerRuntimeExecutionBackend<ProfileRuntimeWorkerFactory> {
pub fn from_workspace(workspace_root: impl Into<PathBuf>) -> Result<Self, String> {
Self::new(ProfileRuntimeWorkerFactory::new(workspace_root))
let workspace_root = workspace_root.into();
let factory = ProfileRuntimeWorkerFactory::new(&workspace_root)
.with_runtime_store_dir(workspace_root.join(".yoi/runtime-store"));
Self::new(factory)
}
}
@@ -1094,6 +1048,7 @@ where
worker_ref: crate::identity::WorkerRef,
bridge_context: crate::execution::WorkerExecutionContext,
handle: WorkerHandle,
shutdown: Arc<tokio::sync::Mutex<Option<worker::ShutdownReceiver>>>,
working_directory: Option<WorkingDirectoryBinding>,
workspace_client: Option<Arc<dyn WorkspaceClient>>,
) -> WorkerExecutionSpawnResult {
@@ -1157,6 +1112,7 @@ where
worker_ref.clone(),
RuntimeWorkerExecution {
handle,
shutdown,
busy,
workspace_client,
},
@@ -1393,6 +1349,7 @@ where
worker_ref,
bridge_context,
controller.handle,
controller.shutdown,
working_directory,
Some(controller.workspace_client),
)
@@ -1489,6 +1446,7 @@ where
worker_ref,
bridge_context,
controller.handle,
controller.shutdown,
working_directory,
Some(controller.workspace_client),
)
@@ -1698,12 +1656,28 @@ where
"execution handle does not reference a live Worker",
);
};
self.send_method(
let shutdown = execution.shutdown.clone();
let result = self.send_method(
WorkerExecutionOperation::Stop,
execution.handle,
Method::Shutdown,
WorkerExecutionRunState::Stopped,
)
);
if result.outcome != crate::execution::WorkerExecutionOutcome::Accepted {
return result;
}
match self.run_on_adapter_runtime(async move {
let receiver = shutdown.lock().await.take();
if let Some(receiver) = receiver {
receiver
.await
.map_err(|_| "Worker shutdown completion channel closed".to_string())?;
}
Ok(())
}) {
Ok(()) => result,
Err(message) => WorkerExecutionResult::errored(WorkerExecutionOperation::Stop, message),
}
}
fn cancel_worker(&self, handle: &WorkerExecutionHandle) -> WorkerExecutionResult {
@@ -1924,12 +1898,13 @@ mod tests {
)
.await
.map_err(|err| err.to_string())?;
let (handle, _shutdown_rx) =
let (handle, shutdown_rx) =
WorkerController::spawn_runtime_managed(worker, &self.runtime_base)
.await
.map_err(|err| err.to_string())?;
Ok(RuntimeWorkerController {
handle,
shutdown: Arc::new(tokio::sync::Mutex::new(Some(shutdown_rx))),
workspace_client,
})
}
@@ -1939,6 +1914,7 @@ mod tests {
) -> Result<RuntimeWorkerController, String> {
let request = WorkerExecutionSpawnRequest {
worker_ref: request.worker_ref,
run_generation: request.run_generation,
request: request.request,
context: request.context,
working_directory: request.working_directory,
@@ -2210,6 +2186,7 @@ mod tests {
let worker_ref = crate::identity::WorkerRef::new(crate::identity::WorkerId::new(1));
let request = WorkerExecutionSpawnRequest {
worker_ref: worker_ref.clone(),
run_generation: 1,
request: create_request("1"),
context: test_execution_context(worker_ref),
working_directory: None,
@@ -2263,9 +2240,9 @@ mod tests {
#[tokio::test]
async fn restore_pending_worker_uses_saved_manifest_snapshot() {
let root = tempfile::tempdir().unwrap();
let store_dir = root.path().join("sessions");
let worker_metadata_dir = root.path().join("workers");
let runtime_store_dir = root.path().join("runtime");
let worker_ref = WorkerRef::new(crate::identity::WorkerId::new(1));
let worker_aggregate_dir = runtime_store_dir.join("workers/1");
let worker_name = ProfileRuntimeWorkerFactory::runtime_worker_name_for_ref(&worker_ref);
let session_id = session_store::new_session_id();
let manifest = manifest::WorkerManifest::from_toml(&format!(
@@ -2294,7 +2271,7 @@ mod tests {
root.path().display(),
))
.unwrap();
FsWorkerStore::new(&worker_metadata_dir)
WorkerAggregateStore::new(&worker_aggregate_dir, &worker_name)
.unwrap()
.set_active(
&worker_name,
@@ -2312,10 +2289,10 @@ mod tests {
runtime_id: Some("runtime-restore".to_string()),
});
let controller = ProfileRuntimeWorkerFactory::new(root.path())
.with_store_dir(&store_dir)
.with_worker_metadata_dir(&worker_metadata_dir)
.with_runtime_store_dir(&runtime_store_dir)
.restore_controller(WorkerExecutionRestoreRequest {
worker_ref: worker_ref.clone(),
run_generation: 1,
request,
context: test_execution_context(worker_ref),
previous_working_directory: None,
@@ -2325,8 +2302,23 @@ mod tests {
.await
.expect("pending restore should use the saved manifest snapshot");
assert!(controller.handle.shared_state.flow_transition_enabled());
let run_dir = runtime_store_dir.join("workers/1/runs/1");
assert!(run_dir.join("worker.sock").exists());
assert!(run_dir.join("worker.out.log").is_file());
assert!(run_dir.join("worker.err.log").is_file());
assert!(run_dir.join("artifacts").is_dir());
assert!(run_dir.join("spawned").is_dir());
let shutdown = controller.shutdown.clone();
controller.handle.send(Method::Shutdown).await.unwrap();
if let Some(receiver) = shutdown.lock().await.take() {
receiver.await.unwrap();
}
assert!(
run_dir.is_dir(),
"run evidence remains until a separate retention policy disposes it"
);
assert!(!run_dir.join("worker.sock").exists());
}
#[test]
crates/worker/src/controller.rs
+29 -4
View File
@@ -211,7 +211,7 @@ impl WorkerController {
C: LlmClient + Clone + 'static,
St: Store + WorkerMetadataStore + Clone + Send + Sync + 'static,
{
Self::spawn_inner(worker, runtime_base, false).await
Self::spawn_inner(worker, runtime_base, false, None).await
}
/// Spawn a Worker owned by `worker-runtime`.
@@ -227,20 +227,37 @@ impl WorkerController {
C: LlmClient + Clone + 'static,
St: Store + WorkerMetadataStore + Clone + Send + Sync + 'static,
{
Self::spawn_inner(worker, runtime_base, true).await
Self::spawn_inner(worker, runtime_base, true, None).await
}
/// Spawn into an exact persistent `runs/<generation>` directory.
pub async fn spawn_runtime_managed_run<C, St>(
worker: Worker<C, St>,
run_dir: &Path,
) -> Result<(WorkerHandle, ShutdownReceiver), std::io::Error>
where
C: LlmClient + Clone + 'static,
St: Store + WorkerMetadataStore + Clone + Send + Sync + 'static,
{
let parent = run_dir
.parent()
.ok_or_else(|| std::io::Error::other("run path has no parent"))?;
Self::spawn_inner(worker, parent, true, Some(run_dir)).await
}
async fn spawn_inner<C, St>(
worker: Worker<C, St>,
runtime_base: &Path,
runtime_managed: bool,
runtime_run: Option<&Path>,
) -> Result<(WorkerHandle, ShutdownReceiver), std::io::Error>
where
C: LlmClient + Clone + 'static,
St: Store + WorkerMetadataStore + Clone + Send + Sync + 'static,
{
let session = worker.workdir_session().cloned();
let result = Self::spawn_initialized(worker, runtime_base, runtime_managed).await;
let result =
Self::spawn_initialized(worker, runtime_base, runtime_managed, runtime_run).await;
if result.is_err()
&& let Some(session) = session
&& let Err(error) = session.close().await
@@ -254,6 +271,7 @@ impl WorkerController {
mut worker: Worker<C, St>,
runtime_base: &Path,
runtime_managed: bool,
runtime_run: Option<&Path>,
) -> Result<(WorkerHandle, ShutdownReceiver), std::io::Error>
where
C: LlmClient + Clone + 'static,
@@ -273,7 +291,9 @@ impl WorkerController {
// the spawn-tool factories need its socket path, and before the
// initial status/history writes consume the greeting we build
// after registration is complete.
let runtime_dir = Arc::new(if runtime_managed {
let runtime_dir = Arc::new(if let Some(run_dir) = runtime_run {
RuntimeDir::create_worker_run(run_dir).await?
} else if runtime_managed {
RuntimeDir::create_transient(runtime_base, &worker.manifest().worker.name).await?
} else {
RuntimeDir::create(runtime_base, &worker.manifest().worker.name).await?
@@ -1298,6 +1318,11 @@ async fn controller_loop<C, St>(
}
}
drop(_socket_server);
if let Err(error) = runtime_dir.close_socket().await {
tracing::warn!(%error, "Worker runtime socket cleanup failed");
}
// Background memory jobs own extract/consolidate workers after a
// turn completes. Join them before closing the Workdir session so no
// Worker-owned task can outlive its operation attachment.
crates/worker/src/runtime/dir.rs
+47 -2
View File
@@ -42,6 +42,8 @@ pub struct SpawnedWorkerRecord {
pub struct RuntimeDir {
path: PathBuf,
write_legacy_snapshots: bool,
preserve_on_drop: bool,
socket_file_name: &'static str,
}
impl RuntimeDir {
@@ -56,6 +58,8 @@ impl RuntimeDir {
Ok(Self {
path,
write_legacy_snapshots: true,
preserve_on_drop: false,
socket_file_name: "sock",
})
}
@@ -69,6 +73,36 @@ impl RuntimeDir {
Ok(Self {
path,
write_legacy_snapshots: false,
preserve_on_drop: false,
socket_file_name: "sock",
})
}
/// Create an exact, persistent generation-scoped Worker run directory.
/// Existing directories are rejected so stale artifacts cannot be reused.
pub async fn create_worker_run(path: &Path) -> Result<Self, io::Error> {
let parent = path
.parent()
.ok_or_else(|| io::Error::other("run path has no parent"))?;
fs::create_dir_all(parent).await?;
fs::create_dir(path).await?;
fs::create_dir(path.join("artifacts")).await?;
fs::create_dir(path.join("spawned")).await?;
for log in ["worker.out.log", "worker.err.log"] {
let file = fs::OpenOptions::new()
.write(true)
.create_new(true)
.open(path.join(log))
.await?;
file.sync_all().await?;
}
std::fs::File::open(path)?.sync_all()?;
std::fs::File::open(parent)?.sync_all()?;
Ok(Self {
path: path.to_path_buf(),
write_legacy_snapshots: false,
preserve_on_drop: true,
socket_file_name: "worker.sock",
})
}
@@ -116,13 +150,24 @@ impl RuntimeDir {
/// that only know the worker name (e.g. the TUI's attach flow)
/// predict the same path via [`manifest::paths::worker_socket_path`].
pub fn socket_path(&self) -> PathBuf {
self.path.join("sock")
self.path.join(self.socket_file_name)
}
pub async fn close_socket(&self) -> Result<(), io::Error> {
match fs::remove_file(self.socket_path()).await {
Ok(()) => Ok(()),
Err(error) if error.kind() == io::ErrorKind::NotFound => Ok(()),
Err(error) => Err(error),
}
}
}
impl Drop for RuntimeDir {
fn drop(&mut self) {
let _ = std::fs::remove_dir_all(&self.path);
let _ = std::fs::remove_file(self.socket_path());
if !self.preserve_on_drop {
let _ = std::fs::remove_dir_all(&self.path);
}
}
}
crates/worker/src/runtime/worker_allocation/lifecycle.rs
+4 -4
View File
@@ -120,15 +120,15 @@ pub fn adopt_allocation(
/// The Worker's in-memory `segment_id` can change underneath the
/// allocation in two normal places:
///
/// - `Worker::compact` mints a fresh session and swaps it in.
/// - `session_store::ensure_head_or_fork` auto-forks when another
/// - `Worker::compact` mints a fresh Segment in the same Session.
/// - `session_store::ensure_head_or_fork` auto-forks within that Session when another
/// writer has advanced the store head behind our back.
///
/// Both paths must call this so subsequent [`lookup_segment`] queries
/// find the live session id, not the old one. Without this update a
/// find the live Segment id, not the old one. Without this update a
/// concurrent `restore_from_manifest(new_id)` would see "no live
/// writer" and proceed to register a competing allocation on the
/// session this Worker just moved into.
/// Segment lineage this Worker just moved into.
///
/// The lock is opened once and the allocation is rewritten inside the
/// guard, so the segment_id collision check is atomic with the
crates/worker/src/worker.rs
+2 -2
View File
@@ -3067,13 +3067,13 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
/// Compact the current session by summarising history via a
/// disposable Engine, then replacing history with
/// `[summary, ...recent_turns]` and creating a new session.
/// `[summary, ...recent_turns]` in a new Segment of the same Session.
///
/// The summary Engine uses:
/// - `compaction.model` from the manifest if configured, or
/// - a clone of the main LlmClient via `clone_boxed()`.
///
/// Returns the new session ID.
/// Returns the new Segment ID. The Worker keeps its Session ID.
pub async fn compact(&mut self, retained_tokens: u64) -> Result<SegmentId, WorkerError> {
use crate::compact::worker::{
CompactWorkerContext, CompactWorkerInterceptor, add_reference_tool,
crates/worker/tests/compact_events_test.rs
+7 -3
View File
@@ -421,7 +421,11 @@ async fn pre_run_compact_success_broadcasts_start_and_done() {
// Drain run events so only compact events remain in `rx`.
let _ = drain(&mut rx);
let session_before = worker.session_id();
let segment_before = worker.segment_id();
worker.try_pre_run_compact().await;
assert_eq!(worker.session_id(), session_before);
assert_ne!(worker.segment_id(), segment_before);
let events = drain(&mut rx);
let kinds: Vec<&str> = events
@@ -442,7 +446,7 @@ async fn pre_run_compact_success_broadcasts_start_and_done() {
"unexpected CompactFailed in {kinds:?}"
);
// CompactDone carries the new session id.
// CompactDone carries the new Segment ID; the Session ID is unchanged.
let new_id_in_event = events.iter().find_map(|e| match e {
Event::CompactDone { new_segment_id } => Some(*new_segment_id),
_ => None,
@@ -583,11 +587,11 @@ async fn compact_resets_extract_pointer_so_extract_can_fire_again() {
);
// Compact runs. Without the fix the in-memory pointer would still
// reference the old session's history_len.
// reference the old Segment's history_len.
worker.try_pre_run_compact().await;
assert!(
worker.extract_pointer().is_none(),
"extract_pointer must be reset to None after compact (matches cold-restore on the new session)"
"extract_pointer must be reset to None after compact (matches cold-restore on the new Segment)"
);
}
crates/workspace-server/src/server.rs
+18
View File
@@ -269,6 +269,7 @@ impl WorkspaceApi {
let resource_broker = BackendResourceBroker::default();
let execution_backend = WorkerRuntimeExecutionBackend::new(
ProfileRuntimeWorkerFactory::new(config.workspace_root.clone())
.with_runtime_store_dir(config.embedded_runtime_store_root.clone())
.with_resource_client(Arc::new(resource_broker.clone())),
)
.map_err(|err| {
@@ -318,6 +319,23 @@ impl WorkspaceApi {
.await?;
import_configured_repositories(store.as_ref(), &config)?;
config.repositories = load_configured_repositories_from_store(store.as_ref(), &config)?;
let default_embedded_root =
ServerConfig::default_embedded_runtime_store_root(&config.workspace_id);
if config.embedded_runtime_store_root == default_embedded_root
&& let (Some(legacy_sessions), Some(data_dir)) =
(manifest::paths::sessions_dir(), manifest::paths::data_dir())
{
worker_runtime::fs_store::FsRuntimeStore::migrate_legacy_worker_aggregates(
&config.embedded_runtime_store_root,
legacy_sessions,
data_dir.join("workers"),
)
.map_err(|error| {
crate::Error::Store(format!(
"failed to migrate embedded Runtime Worker aggregates: {error}"
))
})?;
}
let embedded_runtime = EmbeddedWorkerRuntime::new_fs_store_with_execution_backend(
config.workspace_id.clone(),
config.embedded_runtime_store_root.clone(),