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

feat: unify worker lifecycle restore semantics

Browse Source
This commit is contained in:
Hare
2026-09-03 13:23:31 +09:00
parent 6b1b8a8846
commit 85e1ea320a
21 changed files with 681 additions and 393 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/agen/examples/engine_cancel_demo.rs
+2 -2
View File
@@ -4,7 +4,7 @@
use agen::llm_client::scheme::{Scheme, anthropic::AnthropicScheme};
use agen::llm_client::transport::{HttpTransport, ResolvedAuth};
use agen::{Engine, EngineRunExit, StopReason};
use agen::{Engine, EngineRunExit, RunInterruptionReason};
use std::time::Duration;
#[tokio::main]
@@ -51,7 +51,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
EngineRunExit::Finished => println!("✅ Task completed normally"),
EngineRunExit::Paused => println!("⏸️ Task paused"),
EngineRunExit::Yielded => println!("↩️ Task yielded"),
EngineRunExit::Interrupted(StopReason::LimitReached) => {
EngineRunExit::Interrupted(RunInterruptionReason::LimitReached) => {
println!("🔒 Turn limit reached")
}
EngineRunExit::Interrupted(reason) => println!("❌ Task interrupted: {reason:?}"),
crates/agen/examples/engine_cli.rs
+4 -3
View File
@@ -39,7 +39,7 @@ use tracing::info;
use tracing_subscriber::EnvFilter;
use agen::{
Engine, EngineRunExit, StopReason,
Engine, EngineRunExit, RunInterruptionReason,
interceptor::{Interceptor, PostToolAction, ToolResultInfo},
llm_client::{
LlmClient,
@@ -478,7 +478,8 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
// One-shot mode
if let Some(prompt) = args.prompt {
let output = engine.run(&mut history, &prompt).await;
if let EngineRunExit::Interrupted(StopReason::Unexpected(error)) = output.result {
if let EngineRunExit::Interrupted(RunInterruptionReason::Unexpected(error)) = output.result
{
eprintln!("\n❌ Error: {error}");
}
@@ -518,7 +519,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
break;
}
if let EngineRunExit::Interrupted(StopReason::Unexpected(error)) =
if let EngineRunExit::Interrupted(RunInterruptionReason::Unexpected(error)) =
locked.run(&mut history, input).await
{
eprintln!("\n❌ Error: {error}");
crates/agen/src/engine.rs
+9 -7
View File
@@ -147,12 +147,12 @@ pub enum EngineRunExit {
Finished,
Paused,
Yielded,
Interrupted(StopReason),
Interrupted(RunInterruptionReason),
}
/// A typed reason why an engine run could not finish normally.
#[derive(Debug)]
pub enum StopReason {
pub enum RunInterruptionReason {
LimitReached,
ContextWindowExceeded,
Cancelled,
@@ -165,13 +165,15 @@ impl From<Result<EngineResult, EngineError>> for EngineRunExit {
Ok(EngineResult::Finished) => Self::Finished,
Ok(EngineResult::Paused) => Self::Paused,
Ok(EngineResult::Yielded) => Self::Yielded,
Ok(EngineResult::LimitReached) => Self::Interrupted(StopReason::LimitReached),
Err(EngineError::Client(ClientError::ContextWindowExceeded)) => {
Self::Interrupted(StopReason::ContextWindowExceeded)
Ok(EngineResult::LimitReached) => {
Self::Interrupted(RunInterruptionReason::LimitReached)
}
Err(EngineError::Cancelled) => Self::Interrupted(StopReason::Cancelled),
Err(EngineError::Client(ClientError::ContextWindowExceeded)) => {
Self::Interrupted(RunInterruptionReason::ContextWindowExceeded)
}
Err(EngineError::Cancelled) => Self::Interrupted(RunInterruptionReason::Cancelled),
Err(EngineError::PauseRequested) => Self::Paused,
Err(error) => Self::Interrupted(StopReason::Unexpected(error)),
Err(error) => Self::Interrupted(RunInterruptionReason::Unexpected(error)),
}
}
}
crates/agen/src/lib.rs
+1 -1
View File
@@ -22,7 +22,7 @@ pub use agen_macros::{description, tool, tool_registry};
pub use callback::{TextBlockScope, ThinkingBlockScope, ToolUseBlockScope};
pub use engine::{
Engine, EngineConfig, EngineError, EngineResult, EngineRunExit, EngineRunOutput,
LlmRetryNotice, StopReason, ToolRegistryError,
LlmRetryNotice, RunInterruptionReason, ToolRegistryError,
};
pub use handler::ToolUseBlockStart;
pub use history::{History, HistoryEntry};
crates/agen/tests/engine_state_test.rs
+5 -5
View File
@@ -14,7 +14,7 @@ use agen::interceptor::{
};
use agen::llm_client::event::{Event, ResponseStatus, StatusEvent};
use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
use agen::{Engine, EngineError, EngineRunExit, History, StopReason};
use agen::{Engine, EngineError, EngineRunExit, History, RunInterruptionReason};
use async_trait::async_trait;
use common::MockLlmClient;
@@ -205,7 +205,7 @@ async fn history_append_failure_stops_before_tool_execution() {
let exit = engine.run(&mut history, "use the tool").await;
assert!(
matches!(exit, EngineRunExit::Interrupted(StopReason::Unexpected(EngineError::HistoryAppend(ref message))) if message == "simulated ENOSPC")
matches!(exit, EngineRunExit::Interrupted(RunInterruptionReason::Unexpected(EngineError::HistoryAppend(ref message))) if message == "simulated ENOSPC")
);
assert_eq!(tool.call_count(), 0);
assert_eq!(history.len(), 1);
@@ -730,7 +730,7 @@ async fn paused_tool_resume_does_not_reset_the_consumed_turn_budget() {
assert!(matches!(
engine.resume(&mut history).await,
EngineRunExit::Interrupted(StopReason::LimitReached)
EngineRunExit::Interrupted(RunInterruptionReason::LimitReached)
));
assert_eq!(engine.turn_count(), 1);
assert_eq!(engine.active_run_turn_count(), None);
@@ -785,7 +785,7 @@ async fn interceptor_continuation_consumes_the_logical_run_budget() {
assert!(matches!(
engine.run(&mut history, "start").await,
EngineRunExit::Interrupted(StopReason::LimitReached)
EngineRunExit::Interrupted(RunInterruptionReason::LimitReached)
));
assert_eq!(engine.turn_count(), 1);
assert_eq!(engine.llm_call_count(), 1);
@@ -803,7 +803,7 @@ async fn restored_active_run_budget_is_enforced_before_another_llm_call() {
assert!(matches!(
engine.resume(&mut history).await,
EngineRunExit::Interrupted(StopReason::LimitReached)
EngineRunExit::Interrupted(RunInterruptionReason::LimitReached)
));
assert_eq!(engine.turn_count(), 7);
assert_eq!(engine.llm_call_count(), 0);
crates/agen/tests/parallel_execution_test.rs
+2 -2
View File
@@ -580,7 +580,7 @@ async fn cooperative_cancellation_commits_bounded_terminal_output() {
);
assert!(matches!(
output.result,
agen::EngineRunExit::Interrupted(agen::StopReason::Cancelled)
agen::EngineRunExit::Interrupted(agen::RunInterruptionReason::Cancelled)
));
}
@@ -1214,7 +1214,7 @@ async fn post_tool_abort_commits_confirmed_result_before_stopping_run() {
);
assert!(matches!(
output.result,
agen::EngineRunExit::Interrupted(agen::StopReason::Unexpected(
agen::EngineRunExit::Interrupted(agen::RunInterruptionReason::Unexpected(
agen::EngineError::Aborted(ref reason)
)) if reason == "policy stopped the run"
));
crates/protocol/src/subscription.rs
+19 -1
View File
@@ -557,7 +557,6 @@ pub enum SubscriptionWorkerState {
Running,
Paused,
Stopped,
Cancelled,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
@@ -1110,6 +1109,25 @@ mod tests {
);
}
#[test]
fn worker_subscription_state_has_exactly_four_lifecycle_values() {
for (state, wire) in [
(SubscriptionWorkerState::Idle, "idle"),
(SubscriptionWorkerState::Running, "running"),
(SubscriptionWorkerState::Paused, "paused"),
(SubscriptionWorkerState::Stopped, "stopped"),
] {
assert_eq!(
serde_json::to_value(state).unwrap(),
serde_json::json!(wire)
);
}
assert!(
serde_json::from_value::<SubscriptionWorkerState>(serde_json::json!("cancelled"))
.is_err()
);
}
#[test]
fn client_selector_has_no_workspace_scope_field() {
let json = serde_json::to_value(EventSubscriptionSelector::WorkspaceWorkers).unwrap();
crates/session-store/tests/session_test.rs
+1 -1
View File
@@ -195,7 +195,7 @@ async fn run_and_persist(
)
.unwrap();
}
agen::EngineRunExit::Interrupted(agen::StopReason::LimitReached) => {
agen::EngineRunExit::Interrupted(agen::RunInterruptionReason::LimitReached) => {
session_store::save_run_completed(
store,
session_id,
crates/worker-runtime/src/catalog.rs
+11 -1
View File
@@ -250,6 +250,10 @@ pub struct CreateWorkerRequest {
}
/// Worker lifecycle status for the in-memory embedded runtime.
///
/// Run termination details are carried separately by the Worker protocol. In
/// particular, cancellation returns a Worker to `Idle`; it is not a lifecycle
/// state of its own.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum WorkerStatus {
@@ -257,7 +261,6 @@ pub enum WorkerStatus {
Running,
Paused,
Stopped,
Cancelled,
}
impl WorkerStatus {
@@ -266,6 +269,13 @@ impl WorkerStatus {
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub(crate) enum WorkerRestoreIntent {
Automatic,
Explicit,
}
/// Lightweight catalog row.
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct WorkerSummary {
crates/worker-runtime/src/fs_store.rs
+114 -32
View File
@@ -1,4 +1,6 @@
use crate::catalog::{CreateWorkerRequest, WorkingDirectoryStatus};
use crate::catalog::{
CreateWorkerRequest, WorkerRestoreIntent, WorkerStatus, WorkingDirectoryStatus,
};
use crate::config_bundle::ConfigBundle;
use crate::diagnostics::{DiagnosticSeverity, RuntimeDiagnostic};
use crate::error::RuntimeError;
@@ -13,7 +15,7 @@ use std::io::{BufReader, Write};
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicU64, Ordering};
const SCHEMA_VERSION: u32 = 3;
const SCHEMA_VERSION: u32 = 4;
const RUNTIME_FILE: &str = "runtime.json";
const WORKERS_DIR: &str = "workers";
const WORKER_FILE: &str = "worker.json";
@@ -274,13 +276,24 @@ pub(crate) struct PersistedRuntimeState {
pub(crate) diagnostics: Vec<RuntimeDiagnostic>,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub(crate) struct PersistedWorkerExecutionBinding {
pub(crate) run_generation: u64,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub(crate) struct PersistedWorkerExecution {
pub(crate) binding: Option<PersistedWorkerExecutionBinding>,
pub(crate) restore_intent: WorkerRestoreIntent,
}
#[derive(Clone, Debug)]
pub(crate) struct PersistedWorkerRecord {
pub(crate) worker_ref: WorkerRef,
pub(crate) worker_id: WorkerId,
pub(crate) request: CreateWorkerRequest,
/// Last generation durably reserved for this Worker's execution.
pub(crate) run_generation: u64,
pub(crate) status: WorkerStatus,
pub(crate) execution: PersistedWorkerExecution,
pub(crate) workspace_id: Option<String>,
pub(crate) working_directory: Option<WorkingDirectoryStatus>,
}
@@ -357,8 +370,8 @@ fn plan_runtime_store_migration(
format!("Runtime store schema version {schema_version} is out of range"),
)
})?;
let staging = migration_sibling(root, "schema-v3-staging")?;
let backup = migration_sibling(root, "pre-schema-v3-backup")?;
let staging = migration_sibling(root, "schema-v4-staging")?;
let backup = migration_sibling(root, "pre-schema-v4-backup")?;
if staging.exists() || backup.exists() {
return Err(runtime_store_corrupt(
root,
@@ -384,11 +397,11 @@ fn plan_runtime_store_migration(
};
return Ok((plan, Vec::new()));
}
if !matches!(current_schema_version, 1 | 2) {
if current_schema_version != 3 {
return Err(runtime_store_corrupt(
&runtime_path,
format!(
"unsupported Runtime store schema version {schema_version}; expected 1, 2, or {SCHEMA_VERSION}"
"unsupported Runtime store schema version {schema_version}; expected 3 or {SCHEMA_VERSION}"
),
));
}
@@ -448,7 +461,7 @@ fn plan_runtime_store_migration(
let worker_id = name.parse::<WorkerId>().map_err(|_| {
runtime_store_corrupt(
&source_dir,
format!("schema-v2 Worker directory name must be a UUIDv7, found {name}"),
format!("pre-v4 Worker directory name must be a UUIDv7, found {name}"),
)
})?;
(worker_id, None, None)
@@ -610,7 +623,7 @@ fn migrate_worker_document(
snapshot_path: &Path,
) -> Result<serde_json::Value, RuntimeError> {
if source_schema_version == 1 {
return migrate_v1_worker_document(
document = migrate_v1_worker_document(
document,
mapping.ok_or_else(|| {
runtime_store_corrupt(
@@ -619,7 +632,7 @@ fn migrate_worker_document(
)
})?,
snapshot_path,
);
)?;
}
let object = document.as_object_mut().ok_or_else(|| {
runtime_store_corrupt(
@@ -627,10 +640,46 @@ fn migrate_worker_document(
"Worker snapshot must be an object".to_string(),
)
})?;
let run_generation = object
.remove("run_generation")
.map(|value| {
value.as_u64().ok_or_else(|| {
runtime_store_corrupt(
snapshot_path,
"Worker snapshot run_generation must be an unsigned integer".to_string(),
)
})
})
.transpose()?
.filter(|generation| *generation > 0);
let legacy_execution = object.remove("execution");
if !object.contains_key("working_directory") {
if let Some(working_directory) = legacy_execution
.as_ref()
.and_then(serde_json::Value::as_object)
.and_then(|execution| execution.get("working_directory"))
.cloned()
{
object.insert("working_directory".to_string(), working_directory);
}
}
object.insert(
"schema_version".to_string(),
serde_json::Value::from(SCHEMA_VERSION),
);
object.insert(
"status".to_string(),
serde_json::Value::String("stopped".to_string()),
);
object.insert(
"execution".to_string(),
serde_json::json!({
"binding": run_generation.map(|run_generation| {
serde_json::json!({ "run_generation": run_generation })
}),
"restore_intent": "explicit",
}),
);
Ok(document)
}
@@ -1005,8 +1054,8 @@ fn migrate_runtime_store(
if !plan.migration_required {
return Ok(plan);
}
let staging = migration_sibling(root, "schema-v3-staging")?;
let backup = migration_sibling(root, "pre-schema-v3-backup")?;
let staging = migration_sibling(root, "schema-v4-staging")?;
let backup = migration_sibling(root, "pre-schema-v4-backup")?;
if staging.exists() || backup.exists() {
return Err(runtime_store_corrupt(
root,
@@ -1236,22 +1285,12 @@ struct WorkerSnapshot {
worker_ref: WorkerRef,
worker_id: WorkerId,
request: CreateWorkerRequest,
#[serde(default)]
run_generation: u64,
status: WorkerStatus,
execution: PersistedWorkerExecution,
#[serde(default, skip_serializing_if = "Option::is_none")]
workspace_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
working_directory: Option<WorkingDirectoryStatus>,
/// One-way migration input for schema-v1 snapshots. New snapshots never
/// write the removed execution projection.
#[serde(default, rename = "execution", skip_serializing)]
legacy_execution: Option<LegacyWorkerExecutionProjection>,
}
#[derive(Clone, Debug, Deserialize)]
struct LegacyWorkerExecutionProjection {
#[serde(default)]
working_directory: Option<WorkingDirectoryStatus>,
}
impl WorkerSnapshot {
@@ -1261,10 +1300,10 @@ impl WorkerSnapshot {
worker_ref: worker.worker_ref.clone(),
worker_id: worker.worker_id.clone(),
request: worker.request.clone(),
run_generation: worker.run_generation,
status: worker.status,
execution: worker.execution.clone(),
workspace_id: worker.workspace_id.clone(),
working_directory: worker.working_directory.clone(),
legacy_execution: None,
}
}
@@ -1289,6 +1328,51 @@ impl WorkerSnapshot {
),
});
}
match (self.status, self.execution.restore_intent) {
(status, WorkerRestoreIntent::Automatic) if status.is_active() => {
let Some(binding) = self.execution.binding.as_ref() else {
return Err(RuntimeError::StoreCorrupt {
operation: "read worker snapshot",
path: path.to_path_buf(),
message: "automatic restore intent requires an execution binding"
.to_string(),
});
};
if binding.run_generation == 0 {
return Err(RuntimeError::StoreCorrupt {
operation: "read worker snapshot",
path: path.to_path_buf(),
message: "execution binding run_generation must be greater than zero"
.to_string(),
});
}
}
(WorkerStatus::Stopped, WorkerRestoreIntent::Explicit) => {
if self
.execution
.binding
.as_ref()
.is_some_and(|binding| binding.run_generation == 0)
{
return Err(RuntimeError::StoreCorrupt {
operation: "read worker snapshot",
path: path.to_path_buf(),
message: "execution binding run_generation must be greater than zero"
.to_string(),
});
}
}
_ => {
return Err(RuntimeError::StoreCorrupt {
operation: "read worker snapshot",
path: path.to_path_buf(),
message: format!(
"worker status {:?} conflicts with restore intent {:?}",
self.status, self.execution.restore_intent
),
});
}
}
Ok(())
}
@@ -1303,12 +1387,10 @@ impl WorkerSnapshot {
worker_ref: self.worker_ref,
worker_id: self.worker_id,
request: self.request,
run_generation: self.run_generation,
status: self.status,
execution: self.execution,
workspace_id,
working_directory: self.working_directory.or_else(|| {
self.legacy_execution
.and_then(|execution| execution.working_directory)
}),
working_directory: self.working_directory,
}
}
}
crates/worker-runtime/src/main.rs
+4 -4
View File
@@ -1028,14 +1028,14 @@ mod tests {
}
#[test]
fn migration_dry_run_accepts_real_v1_document_without_workers_field() {
fn migration_dry_run_accepts_previous_schema_document_without_workers_field() {
let temp = tempfile::tempdir().unwrap();
let root = temp.path().join("runtime");
std::fs::create_dir_all(root.join("workers")).unwrap();
std::fs::write(
root.join("runtime.json"),
serde_json::to_vec_pretty(&serde_json::json!({
"schema_version": 1,
"schema_version": 3,
"display_name": "local",
"backend": "fs_store",
"status": "running",
@@ -1067,14 +1067,14 @@ mod tests {
}
#[test]
fn migration_dry_run_rejects_v1_document_that_cannot_decode_as_v3() {
fn migration_dry_run_rejects_previous_schema_document_that_cannot_decode_as_v4() {
let temp = tempfile::tempdir().unwrap();
let root = temp.path().join("runtime");
std::fs::create_dir_all(root.join("workers")).unwrap();
std::fs::write(
root.join("runtime.json"),
serde_json::to_vec_pretty(&serde_json::json!({
"schema_version": 1,
"schema_version": 3,
"display_name": "local",
"backend": "fs_store",
"status": 3,
crates/worker-runtime/src/management.rs
-1
View File
@@ -43,7 +43,6 @@ pub struct RuntimeSummary {
pub worker_count: usize,
pub active_worker_count: usize,
pub stopped_worker_count: usize,
pub cancelled_worker_count: usize,
pub diagnostic_count: usize,
#[serde(default = "unknown_platform_component")]
pub os: String,
crates/worker-runtime/src/retention.rs
+40 -8
View File
@@ -303,7 +303,12 @@ impl FsWorkerRetentionProvider {
));
continue;
}
match self.inventory(workspace_id, runtime_id, worker_id, snapshot.run_generation) {
match self.inventory(
workspace_id,
runtime_id,
worker_id,
snapshot.run_generation(),
) {
Ok(item) => workers.push(item),
Err(_) => diagnostics.push(runtime_aggregate_diagnostic(
&bounded_id,
@@ -388,10 +393,11 @@ impl WorkerRetentionProvider for FsWorkerRetentionProvider {
if worker.workspace_id.as_deref() != Some(workspace_id) {
return Err(RuntimeError::WorkerNotFound { worker_id });
}
if worker.run_generation != run_generation {
let current_run_generation = worker.run_generation();
if current_run_generation != run_generation {
return Err(RuntimeError::InvalidRequest(format!(
"Worker retention inventory expected generation {run_generation}, current generation is {}",
worker.run_generation
current_run_generation
)));
}
let session_dir = worker_dir.join("session");
@@ -498,10 +504,11 @@ impl WorkerRetentionProvider for FsWorkerRetentionProvider {
worker_id: request.worker_id,
});
}
if snapshot.run_generation != request.expected_run_generation {
let run_generation = snapshot.run_generation();
if run_generation != request.expected_run_generation {
return Err(RuntimeError::InvalidRequest(format!(
"Worker retention plan expected generation {}, current generation is {}",
request.expected_run_generation, snapshot.run_generation
request.expected_run_generation, run_generation
)));
}
@@ -568,10 +575,29 @@ impl WorkerRetentionProvider for FsWorkerRetentionProvider {
struct WorkerGenerationSnapshot {
#[serde(default)]
workspace_id: Option<String>,
#[serde(default)]
execution: WorkerGenerationExecution,
}
#[derive(Deserialize)]
struct WorkerGenerationExecution {
binding: Option<WorkerGenerationBinding>,
}
#[derive(Deserialize)]
struct WorkerGenerationBinding {
run_generation: u64,
}
impl WorkerGenerationSnapshot {
fn run_generation(&self) -> u64 {
self.execution
.binding
.as_ref()
.map(|binding| binding.run_generation)
.unwrap_or(0)
}
}
#[derive(Deserialize)]
struct CanonicalSessionManifest {
session_id: String,
@@ -1264,7 +1290,10 @@ mod tests {
let worker = root.join("workers").join(worker_id.to_string());
write_json(
&worker.join("worker.json"),
&serde_json::json!({"workspace_id": "workspace-a", "run_generation": generation}),
&serde_json::json!({
"workspace_id": "workspace-a",
"execution": {"binding": {"run_generation": generation}}
}),
);
write_json(
&worker.join("session/session.json"),
@@ -1462,7 +1491,10 @@ mod tests {
.join("workers")
.join(other_worker.to_string())
.join("worker.json"),
&serde_json::json!({"workspace_id": "other-workspace", "run_generation": 1}),
&serde_json::json!({
"workspace_id": "other-workspace",
"execution": {"binding": {"run_generation": 1}}
}),
);
fs::create_dir_all(temp.path().join("workers/not-a-worker")).unwrap();
fs::write(
crates/worker-runtime/src/runtime.rs
+433 -277
View File
@@ -1,7 +1,8 @@
use crate::catalog::{
ConfigBundleRef, CreateWorkerRequest, ProfileSelector, WorkerDetail, WorkerLifecycleAck,
WorkerStatus, WorkerSummary, WorkingDirectoryRepositoryAccessRequest, WorkingDirectoryRequest,
WorkingDirectoryStatus as CatalogWorkingDirectoryStatus, WorkspaceApiRef,
WorkerRestoreIntent, WorkerStatus, WorkerSummary, WorkingDirectoryRepositoryAccessRequest,
WorkingDirectoryRequest, WorkingDirectoryStatus as CatalogWorkingDirectoryStatus,
WorkspaceApiRef,
};
use crate::config_bundle::{
ConfigBundle, ConfigBundleAvailability, ConfigBundleSummary, validate_config_bundle,
@@ -17,7 +18,8 @@ use crate::execution::{
};
#[cfg(feature = "fs-store")]
use crate::fs_store::{
FsRuntimeStore, FsRuntimeStoreOptions, PersistedRuntimeState, PersistedWorkerRecord,
FsRuntimeStore, FsRuntimeStoreOptions, PersistedRuntimeState, PersistedWorkerExecution,
PersistedWorkerExecutionBinding, PersistedWorkerRecord,
};
use crate::identity::{WorkerId, WorkerRef};
use crate::interaction::{WorkerInput, WorkerInputKind, WorkerInteractionAck};
@@ -229,14 +231,12 @@ impl Runtime {
let state = self.lock()?;
let mut active_worker_count = 0;
let mut stopped_worker_count = 0;
let mut cancelled_worker_count = 0;
for worker in state.workers.values() {
match worker.status {
WorkerStatus::Idle | WorkerStatus::Running | WorkerStatus::Paused => {
active_worker_count += 1;
}
WorkerStatus::Stopped => stopped_worker_count += 1,
WorkerStatus::Cancelled => cancelled_worker_count += 1,
}
}
@@ -247,7 +247,6 @@ impl Runtime {
worker_count: state.workers.len(),
active_worker_count,
stopped_worker_count,
cancelled_worker_count,
diagnostic_count: state.diagnostics.len(),
os: std::env::consts::OS.to_string(),
arch: std::env::consts::ARCH.to_string(),
@@ -343,17 +342,6 @@ impl Runtime {
return Ok(());
}
state.status = RuntimeStatus::Stopped;
let mut stopped = Vec::new();
for (worker_id, worker) in &mut state.workers {
if worker.status.is_active() {
worker.status = WorkerStatus::Stopped;
worker.internal_workers.clear();
stopped.push(*worker_id);
}
}
for worker_id in stopped {
state.publish_worker_upsert(worker_id)?;
}
state.persist_runtime_snapshot()?;
state.persist_workers()?;
Ok(())
@@ -683,6 +671,8 @@ impl Runtime {
workspace_id: scope.map(|scope| scope.workspace_id.clone()),
request: durable_request,
run_generation: 1,
execution_bound: true,
restore_intent: WorkerRestoreIntent::Explicit,
working_directory: None,
execution_handle: None,
internal_workers: BTreeMap::new(),
@@ -1019,12 +1009,6 @@ impl Runtime {
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
)));
}
(
worker.request.clone(),
worker.working_directory.clone(),
@@ -1037,7 +1021,12 @@ impl Runtime {
message: "runtime has no execution backend".to_string(),
}
})?;
state.worker_mut(worker_ref)?.run_generation = run_generation;
{
let worker = state.worker_mut(worker_ref)?;
worker.run_generation = run_generation;
worker.execution_bound = true;
worker.restore_intent = WorkerRestoreIntent::Automatic;
}
state.persist_worker(&worker_ref.worker_id)?;
let workspace_scope = worker_request.workspace_api.as_ref().and_then(|api| {
state
@@ -1091,18 +1080,20 @@ impl Runtime {
}
fn ensure_worker_execution(&self, worker_ref: &WorkerRef) -> Result<(), RuntimeError> {
let has_handle = {
let state = self.lock()?;
state
.worker(worker_ref)?
.execution_handle
.as_ref()
.is_some()
};
if has_handle {
let state = self.lock()?;
let worker = state.worker(worker_ref)?;
if worker.execution_handle.is_some() {
return Ok(());
}
self.restore_worker(worker_ref).map(|_| ())
let message = if worker.status == WorkerStatus::Stopped {
"stopped worker requires an explicit restore"
} else {
"worker has no live execution handle"
};
Err(RuntimeError::WorkerExecutionUnavailable {
worker_id: worker_ref.worker_id,
message: message.to_string(),
})
}
/// Accept input into a Worker through a workspace-scoped Runtime authorization context.
@@ -1455,7 +1446,9 @@ impl Runtime {
let detail = {
let worker = state.worker_mut(worker_ref)?;
worker.execution_handle = Some(handle);
worker.execution_bound = true;
worker.status = worker_status_from_run_state(run_state);
worker.restore_intent = restore_intent_for_status(worker.status);
worker.working_directory = working_directory;
worker.detail()
};
@@ -1484,8 +1477,13 @@ impl Runtime {
) -> Result<(), RuntimeError> {
let mut state = self.lock()?;
if result.is_accepted() {
state.worker_mut(worker_ref)?.status = worker_status_from_run_state(result.run_state);
let status = worker_status_from_run_state(result.run_state);
let worker = state.worker_mut(worker_ref)?;
worker.status = status;
worker.restore_intent = restore_intent_for_status(status);
state.publish_worker_upsert(worker_ref.worker_id)?;
state.persist_runtime_snapshot()?;
state.persist_worker(&worker_ref.worker_id)?;
}
Ok(())
}
@@ -1568,15 +1566,26 @@ impl Runtime {
self.cancel_worker(worker_ref, reason)
}
/// Cancel a Worker. Repeated cancels are idempotent.
/// Cancel the current run while keeping the Worker session available.
pub fn cancel_worker(
&self,
worker_ref: &WorkerRef,
reason: Option<String>,
) -> Result<WorkerLifecycleAck, RuntimeError> {
let current = {
let state = self.lock()?;
state.ensure_running()?;
state.worker(worker_ref)?.status
};
if matches!(current, WorkerStatus::Idle | WorkerStatus::Stopped) {
return Ok(WorkerLifecycleAck {
worker_ref: worker_ref.clone(),
status: current,
});
}
self.dispatch_lifecycle_to_backend(worker_ref, WorkerExecutionOperation::Cancel)?;
let _ = reason;
self.transition_worker(worker_ref, WorkerStatus::Cancelled)
self.transition_worker_preserving_execution(worker_ref, WorkerStatus::Idle)
}
/// Delete a non-running Worker through a workspace-scoped Runtime authorization context.
@@ -1726,6 +1735,10 @@ impl Runtime {
if status_changed || activity_changed {
state.publish_worker_upsert(worker_ref.worker_id)?;
}
if status_changed {
state.persist_runtime_snapshot()?;
state.persist_worker(&worker_ref.worker_id)?;
}
let event = state.push_worker_observation_event(worker_ref.clone(), payload);
Ok(event)
}
@@ -1758,6 +1771,26 @@ impl Runtime {
Ok(())
}
fn transition_worker_preserving_execution(
&self,
worker_ref: &WorkerRef,
status: WorkerStatus,
) -> Result<WorkerLifecycleAck, RuntimeError> {
let mut state = self.lock()?;
state.ensure_running()?;
let worker = state.worker_mut(worker_ref)?;
worker.status = status;
worker.restore_intent = restore_intent_for_status(status);
let status = worker.status;
state.publish_worker_upsert(worker_ref.worker_id)?;
state.persist_runtime_snapshot()?;
state.persist_worker(&worker_ref.worker_id)?;
Ok(WorkerLifecycleAck {
worker_ref: worker_ref.clone(),
status,
})
}
fn transition_worker(
&self,
worker_ref: &WorkerRef,
@@ -1779,6 +1812,7 @@ impl Runtime {
let worker = state.worker_mut(worker_ref)?;
worker.status = status;
worker.restore_intent = restore_intent_for_status(status);
worker.execution_handle = None;
worker.internal_workers.clear();
let status = worker.status;
@@ -1834,7 +1868,12 @@ impl Runtime {
let worker_ids = state
.workers
.values()
.filter(|worker| worker.execution_handle.is_none())
.filter(|worker| {
worker.execution_handle.is_none()
&& worker.execution_bound
&& worker.status.is_active()
&& worker.restore_intent == WorkerRestoreIntent::Automatic
})
.map(|worker| worker.worker_id)
.collect::<Vec<_>>();
let mut candidates = Vec::with_capacity(worker_ids.len());
@@ -1925,7 +1964,9 @@ impl Runtime {
{
let worker = state.worker_mut(worker_ref)?;
worker.execution_handle = Some(handle);
worker.execution_bound = true;
worker.status = worker_status_from_run_state(run_state);
worker.restore_intent = restore_intent_for_status(worker.status);
worker.working_directory = working_directory;
}
state.publish_worker_upsert(worker_ref.worker_id)?;
@@ -2194,15 +2235,23 @@ impl RuntimeState {
let diagnostics = persisted.diagnostics;
let next_diagnostic_id = persisted.next_diagnostic_id;
for (worker_id, worker) in persisted.workers {
let run_generation = worker
.execution
.binding
.as_ref()
.map(|binding| binding.run_generation)
.unwrap_or(0);
workers.insert(
worker_id,
WorkerRecord {
worker_ref: worker.worker_ref,
worker_id: worker.worker_id,
status: WorkerStatus::Stopped,
status: worker.status,
workspace_id: worker.workspace_id,
request: worker.request,
run_generation: worker.run_generation,
run_generation,
execution_bound: worker.execution.binding.is_some(),
restore_intent: worker.execution.restore_intent,
working_directory: worker.working_directory,
execution_handle: None,
internal_workers: BTreeMap::new(),
@@ -2675,9 +2724,11 @@ impl RuntimeState {
let worker = self.worker_mut(worker_ref)?;
worker.execution_handle = None;
worker.status = WorkerStatus::Stopped;
worker.restore_intent = WorkerRestoreIntent::Explicit;
worker.internal_workers.clear();
self.publish_worker_upsert(worker_ref.worker_id)?;
self.persist_runtime_snapshot()?;
self.persist_worker(&worker_ref.worker_id)?;
Ok(())
}
@@ -2893,6 +2944,7 @@ impl RuntimeState {
if let Some(next_status) = next_status {
let changed = worker.status != next_status;
worker.status = next_status;
worker.restore_intent = restore_intent_for_status(next_status);
changed
} else {
false
@@ -2914,6 +2966,8 @@ struct WorkerRecord {
workspace_id: Option<String>,
request: CreateWorkerRequest,
run_generation: u64,
execution_bound: bool,
restore_intent: WorkerRestoreIntent,
working_directory: Option<CatalogWorkingDirectoryStatus>,
execution_handle: Option<WorkerExecutionHandle>,
internal_workers: BTreeMap<String, InternalWorkerActivity>,
@@ -2958,13 +3012,29 @@ impl WorkerRecord {
worker_ref: self.worker_ref.clone(),
worker_id: self.worker_id.clone(),
request: self.request.clone(),
run_generation: self.run_generation,
status: self.status,
execution: PersistedWorkerExecution {
binding: self
.execution_bound
.then_some(PersistedWorkerExecutionBinding {
run_generation: self.run_generation,
}),
restore_intent: self.restore_intent,
},
workspace_id: self.workspace_id.clone(),
working_directory: self.working_directory.clone(),
}
}
}
fn restore_intent_for_status(status: WorkerStatus) -> WorkerRestoreIntent {
if status.is_active() {
WorkerRestoreIntent::Automatic
} else {
WorkerRestoreIntent::Explicit
}
}
fn worker_status_from_run_state(run_state: WorkerExecutionRunState) -> WorkerStatus {
match run_state {
WorkerExecutionRunState::Idle => WorkerStatus::Idle,
@@ -3164,7 +3234,6 @@ fn subscription_worker_state(status: WorkerStatus) -> SubscriptionWorkerState {
WorkerStatus::Running => SubscriptionWorkerState::Running,
WorkerStatus::Paused => SubscriptionWorkerState::Paused,
WorkerStatus::Stopped => SubscriptionWorkerState::Stopped,
WorkerStatus::Cancelled => SubscriptionWorkerState::Cancelled,
}
}
@@ -4842,15 +4911,22 @@ mod tests {
}
#[test]
fn input_restores_stopped_worker_without_persisted_connection_state() {
fn stopped_worker_rejects_input_until_explicitly_restored() {
let (runtime, backend) = runtime_and_backend();
let detail = runtime
.create_worker(task_request("restore on input"))
.create_worker(task_request("restore explicitly"))
.unwrap();
runtime
.send_protocol_method(&detail.worker_ref, Method::Shutdown)
.unwrap();
assert!(matches!(
runtime.send_input(&detail.worker_ref, WorkerInput::user("do not wake")),
Err(RuntimeError::WorkerExecutionUnavailable { .. })
));
assert_eq!(*backend.restore_count.lock().unwrap(), 0);
runtime.restore_worker(&detail.worker_ref).unwrap();
runtime
.send_input(&detail.worker_ref, WorkerInput::user("wake up"))
.unwrap();
@@ -4998,10 +5074,13 @@ mod tests {
}
#[test]
fn stop_and_cancel_workers_update_summary() {
fn stop_and_cancel_workers_keep_four_state_summary() {
let runtime = runtime_with_backend();
let stopped = runtime.create_worker(task_request("stop me")).unwrap();
let cancelled = runtime.create_worker(task_request("cancel me")).unwrap();
runtime
.send_input(&cancelled.worker_ref, WorkerInput::user("start"))
.unwrap();
let stop_ack = runtime
.stop_worker(&stopped.worker_ref, Some("done".to_string()))
@@ -5011,13 +5090,13 @@ mod tests {
let cancel_ack = runtime
.cancel_worker(&cancelled.worker_ref, Some("abort".to_string()))
.unwrap();
assert_eq!(cancel_ack.status, WorkerStatus::Cancelled);
assert_eq!(cancel_ack.status, WorkerStatus::Idle);
let summary = runtime.summary().unwrap();
assert_eq!(summary.worker_count, 2);
assert_eq!(summary.active_worker_count, 0);
assert_eq!(summary.active_worker_count, 1);
assert_eq!(summary.stopped_worker_count, 1);
assert_eq!(summary.cancelled_worker_count, 1);
assert!(serde_json::from_value::<WorkerStatus>(serde_json::json!("cancelled")).is_err());
}
#[test]
@@ -5064,14 +5143,16 @@ mod tests {
let summary = runtime.summary().unwrap();
assert_eq!(summary.active_worker_count, 0);
assert_eq!(summary.stopped_worker_count, 1);
assert_eq!(summary.cancelled_worker_count, 0);
}
#[test]
fn cancel_then_stop_preserves_cancelled_terminal_state() {
fn cancel_then_stop_transitions_idle_session_to_stopped() {
let runtime = runtime_with_backend();
let worker = runtime
.create_worker(task_request("stable cancelled"))
.create_worker(task_request("cancel then stop"))
.unwrap();
runtime
.send_input(&worker.worker_ref, WorkerInput::user("start"))
.unwrap();
let cancel_ack = runtime
@@ -5081,17 +5162,16 @@ mod tests {
.stop_worker(&worker.worker_ref, Some("late stop".to_string()))
.unwrap();
assert_eq!(cancel_ack.status, WorkerStatus::Cancelled);
assert_eq!(stop_ack.status, WorkerStatus::Cancelled);
assert_eq!(cancel_ack.status, WorkerStatus::Idle);
assert_eq!(stop_ack.status, WorkerStatus::Stopped);
assert_eq!(
runtime.worker_detail(&worker.worker_ref).unwrap().status,
WorkerStatus::Cancelled
WorkerStatus::Stopped
);
let summary = runtime.summary().unwrap();
assert_eq!(summary.active_worker_count, 0);
assert_eq!(summary.stopped_worker_count, 0);
assert_eq!(summary.cancelled_worker_count, 1);
assert_eq!(summary.stopped_worker_count, 1);
}
#[cfg(feature = "fs-store")]
@@ -5119,239 +5199,38 @@ mod tests {
#[cfg(feature = "fs-store")]
#[test]
fn fs_store_migrates_legacy_numeric_worker_identity_to_workspace_uuid() {
let root = fs_store_root("worker-id-v1");
let runtime_id = "arcadia";
let runtime = Runtime::with_fs_store_and_execution_backend(
crate::fs_store::FsRuntimeStoreOptions {
root: root.clone(),
runtime_id: runtime_id.to_string(),
display_name: None,
},
Arc::new(TestExecutionBackend::default()),
)
fn fs_store_rejects_schema_older_than_previous_release() {
let root = fs_store_root("unsupported-old-schema");
let runtime = Runtime::with_fs_store(crate::fs_store::FsRuntimeStoreOptions {
root: root.clone(),
runtime_id: "test-runtime".to_string(),
display_name: None,
})
.unwrap();
runtime.store_config_bundle(test_bundle()).unwrap();
let worker = runtime
.create_worker_scoped(
&RuntimeWorkspaceScope::new("workspace-a", "server"),
scoped_task_request("legacy", "workspace-a"),
)
.unwrap();
drop(runtime);
let current_dir = root.join("workers").join(worker.worker_id.to_string());
let legacy_dir = root.join("workers").join("7");
std::fs::rename(&current_dir, &legacy_dir).unwrap();
let worker_path = legacy_dir.join("worker.json");
let mut worker_json: serde_json::Value =
serde_json::from_slice(&std::fs::read(&worker_path).unwrap()).unwrap();
worker_json["schema_version"] = serde_json::json!(1);
worker_json["worker_id"] = serde_json::json!(7);
worker_json["worker_ref"]["worker_id"] = serde_json::json!(7);
let request = worker_json["request"].as_object_mut().unwrap();
request.remove("worker_id");
request.remove("create_fingerprint");
request.insert("idempotency_key".to_string(), serde_json::Value::Null);
request.insert(
"idempotency_fingerprint".to_string(),
serde_json::Value::Null,
);
std::fs::write(
&worker_path,
serde_json::to_vec_pretty(&worker_json).unwrap(),
)
.unwrap();
let legacy_worker_name = "worker-runtime-7";
let legacy_manifest = manifest::WorkerManifest::from_toml(&format!(
r#"
[worker]
name = "{legacy_worker_name}"
[model]
scheme = "anthropic"
model_id = "test-model"
[engine]
[[scope.allow]]
target = "/tmp"
permission = "write"
"#,
))
.unwrap();
std::fs::write(
legacy_dir.join("metadata.json"),
serde_json::to_vec_pretty(&serde_json::json!({
"worker_name": legacy_worker_name,
"workspace_id": "workspace-a",
"resolved_manifest_snapshot": legacy_manifest
}))
.unwrap(),
)
.unwrap();
let runtime_path = root.join("runtime.json");
let mut runtime_json: serde_json::Value =
serde_json::from_slice(&std::fs::read(&runtime_path).unwrap()).unwrap();
runtime_json["schema_version"] = serde_json::json!(1);
runtime_json["workers"] = serde_json::json!({"legacy": "ignored"});
runtime_json["next_worker_sequence"] = serde_json::json!(8);
runtime_json["next_diagnostic_id"] = serde_json::json!(3);
runtime_json["diagnostics"] = serde_json::json!([
{
"id": 1,
"severity": "warning",
"code": "mapped_legacy_worker",
"message": "mapped diagnostic",
"worker_ref": {"worker_id": 7}
},
{
"id": 2,
"severity": "warning",
"code": "deleted_legacy_worker",
"message": "unmapped diagnostic",
"worker_ref": {"worker_id": 6}
}
]);
runtime_json["schema_version"] = serde_json::json!(2);
std::fs::write(
&runtime_path,
serde_json::to_vec_pretty(&runtime_json).unwrap(),
)
.unwrap();
#[cfg(unix)]
{
let run_dir = legacy_dir.join("runs").join("6");
std::fs::create_dir_all(&run_dir).unwrap();
let socket =
std::os::unix::net::UnixListener::bind(run_dir.join("worker.sock")).unwrap();
drop(socket);
}
let runtime_options = crate::fs_store::FsRuntimeStoreOptions {
let error = Runtime::with_fs_store(crate::fs_store::FsRuntimeStoreOptions {
root: root.clone(),
runtime_id: runtime_id.to_string(),
runtime_id: "test-runtime".to_string(),
display_name: None,
};
let runtime_before_dry_run = std::fs::read(&runtime_path).unwrap();
let plan = crate::fs_store::FsRuntimeStore::migration_plan(&runtime_options).unwrap();
assert!(plan.migration_required);
assert_eq!(plan.worker_count, 1);
assert_eq!(plan.migrated_worker_aggregate_count, 1);
assert_eq!(plan.migrated_diagnostic_worker_ref_count, 1);
assert_eq!(plan.cleared_diagnostic_worker_ref_count, 1);
assert_eq!(plan.mappings[0].legacy_worker_id, 7);
#[cfg(unix)]
assert_eq!(
plan.excluded_ephemeral_paths,
vec!["workers/7/runs/6/worker.sock"]
);
assert_eq!(
std::fs::read(&runtime_path).unwrap(),
runtime_before_dry_run
);
assert!(legacy_dir.exists());
let restored = Runtime::with_fs_store(runtime_options.clone()).unwrap();
let expected = WorkerId::from_legacy_binding("workspace-a", runtime_id, 7);
let detail = restored.worker_detail(&WorkerRef::new(expected)).unwrap();
assert_eq!(detail.worker_id, expected);
assert_eq!(detail.worker_ref.worker_id, expected);
let expected_worker_dir = root.join("workers").join(expected.to_string());
assert!(expected_worker_dir.exists());
#[cfg(unix)]
assert!(!expected_worker_dir.join("runs/6/worker.sock").exists());
assert!(!legacy_dir.exists());
let migrated_runtime: serde_json::Value =
serde_json::from_slice(&std::fs::read(&runtime_path).unwrap()).unwrap();
assert_eq!(migrated_runtime["schema_version"], serde_json::json!(3));
assert!(migrated_runtime.get("workers").is_none());
assert!(migrated_runtime.get("next_worker_sequence").is_none());
assert_eq!(
migrated_runtime["diagnostics"][0]["worker_ref"]["worker_id"],
serde_json::json!(expected.to_string())
);
})
.unwrap_err();
assert!(
migrated_runtime["diagnostics"][1]
.get("worker_ref")
.is_none()
);
let diagnostics = restored.diagnostics().unwrap();
assert_eq!(
diagnostics
.iter()
.find(|diagnostic| diagnostic.code == "mapped_legacy_worker")
.and_then(|diagnostic| diagnostic.worker_ref.as_ref()),
Some(&WorkerRef::new(expected))
);
assert!(
diagnostics
.iter()
.find(|diagnostic| diagnostic.code == "deleted_legacy_worker")
.is_some_and(|diagnostic| diagnostic.worker_ref.is_none())
);
let metadata_path = expected_worker_dir.join("metadata.json");
let mut migrated_metadata: serde_json::Value =
serde_json::from_slice(&std::fs::read(&metadata_path).unwrap()).unwrap();
let expected_worker_name = format!("worker-runtime-{expected}");
assert_eq!(
migrated_metadata["worker_name"],
serde_json::json!(expected_worker_name)
);
assert_eq!(
migrated_metadata["resolved_manifest_snapshot"]["worker"]["name"],
serde_json::json!(expected_worker_name)
error
.to_string()
.contains("unsupported Runtime store schema version 2; expected 3 or 4")
);
drop(restored);
let mut schema_v2_runtime: serde_json::Value =
serde_json::from_slice(&std::fs::read(&runtime_path).unwrap()).unwrap();
schema_v2_runtime["schema_version"] = serde_json::json!(2);
std::fs::write(
&runtime_path,
serde_json::to_vec_pretty(&schema_v2_runtime).unwrap(),
)
.unwrap();
let migrated_worker_path = expected_worker_dir.join("worker.json");
let mut schema_v2_worker: serde_json::Value =
serde_json::from_slice(&std::fs::read(&migrated_worker_path).unwrap()).unwrap();
schema_v2_worker["schema_version"] = serde_json::json!(2);
std::fs::write(
&migrated_worker_path,
serde_json::to_vec_pretty(&schema_v2_worker).unwrap(),
)
.unwrap();
migrated_metadata["worker_name"] = serde_json::json!(legacy_worker_name);
migrated_metadata["resolved_manifest_snapshot"]["worker"]["name"] =
serde_json::json!(legacy_worker_name);
std::fs::write(
&metadata_path,
serde_json::to_vec_pretty(&migrated_metadata).unwrap(),
)
.unwrap();
let recovery_plan =
crate::fs_store::FsRuntimeStore::migration_plan(&runtime_options).unwrap();
assert_eq!(recovery_plan.current_schema_version, 2);
assert_eq!(recovery_plan.target_schema_version, 3);
assert!(recovery_plan.migration_required);
assert_eq!(recovery_plan.worker_count, 1);
assert_eq!(recovery_plan.migrated_worker_aggregate_count, 1);
assert!(recovery_plan.mappings.is_empty());
let recovered = Runtime::with_fs_store(runtime_options).unwrap();
let recovered_metadata: serde_json::Value =
serde_json::from_slice(&std::fs::read(metadata_path).unwrap()).unwrap();
assert_eq!(
recovered_metadata["worker_name"],
serde_json::json!(expected_worker_name)
);
assert_eq!(
recovered_metadata["resolved_manifest_snapshot"]["worker"]["name"],
serde_json::json!(expected_worker_name)
);
drop(recovered);
let _ = std::fs::remove_dir_all(root);
}
@@ -5391,8 +5270,16 @@ mod tests {
let worker_snapshot: serde_json::Value =
serde_json::from_slice(&std::fs::read(worker_store_dir.join("worker.json")).unwrap())
.unwrap();
assert!(worker_snapshot.get("status").is_none());
assert!(worker_snapshot.get("execution").is_none());
assert_eq!(worker_snapshot["schema_version"], serde_json::json!(4));
assert_eq!(worker_snapshot["status"], serde_json::json!("stopped"));
assert_eq!(
worker_snapshot["execution"]["binding"]["run_generation"],
serde_json::json!(1)
);
assert_eq!(
worker_snapshot["execution"]["restore_intent"],
serde_json::json!("explicit")
);
assert!(!root.join("events.jsonl").exists());
std::fs::write(
worker_store_dir.join("observations.jsonl"),
@@ -5549,8 +5436,8 @@ mod tests {
display_name: None,
})
.unwrap();
let stopped_worker = backendless.worker_detail(&worker.worker_ref).unwrap();
assert_eq!(stopped_worker.status, WorkerStatus::Stopped);
let persisted_worker = backendless.worker_detail(&worker.worker_ref).unwrap();
assert_eq!(persisted_worker.status, WorkerStatus::Idle);
drop(backendless);
let restoring_backend = Arc::new(TestExecutionBackend::default());
@@ -5574,6 +5461,272 @@ mod tests {
let _ = std::fs::remove_dir_all(root);
}
#[cfg(feature = "fs-store")]
#[test]
fn fs_store_automatically_restores_every_active_lifecycle_state() {
for status in ["idle", "running", "paused"] {
let root = fs_store_root(&format!("automatic-{status}"));
let options = crate::fs_store::FsRuntimeStoreOptions {
root: root.clone(),
runtime_id: "test-runtime".to_string(),
display_name: None,
};
let runtime = Runtime::with_fs_store_and_execution_backend(
options.clone(),
Arc::new(TestExecutionBackend::default()),
)
.unwrap();
runtime.store_config_bundle(test_bundle()).unwrap();
let worker = runtime
.create_worker(task_request(&format!("restore {status}")))
.unwrap();
drop(runtime);
let worker_path = root
.join("workers")
.join(worker.worker_id.to_string())
.join("worker.json");
let mut worker_json: serde_json::Value =
serde_json::from_slice(&std::fs::read(&worker_path).unwrap()).unwrap();
worker_json["status"] = serde_json::json!(status);
std::fs::write(
&worker_path,
serde_json::to_vec_pretty(&worker_json).unwrap(),
)
.unwrap();
let backend = Arc::new(TestExecutionBackend::default());
let restored =
Runtime::with_fs_store_and_execution_backend(options, backend.clone()).unwrap();
assert_eq!(*backend.restore_count.lock().unwrap(), 1, "status={status}");
assert_eq!(
restored.worker_detail(&worker.worker_ref).unwrap().status,
WorkerStatus::Idle,
"status={status}"
);
let _ = std::fs::remove_dir_all(root);
}
}
#[cfg(feature = "fs-store")]
#[test]
fn fs_store_current_schema_requires_lifecycle_authority() {
let root = fs_store_root("current-schema-requires-lifecycle");
let options = crate::fs_store::FsRuntimeStoreOptions {
root: root.clone(),
runtime_id: "test-runtime".to_string(),
display_name: None,
};
let runtime = Runtime::with_fs_store_and_execution_backend(
options.clone(),
Arc::new(TestExecutionBackend::default()),
)
.unwrap();
runtime.store_config_bundle(test_bundle()).unwrap();
let worker = runtime
.create_worker(task_request("missing lifecycle authority"))
.unwrap();
drop(runtime);
let worker_path = root
.join("workers")
.join(worker.worker_id.to_string())
.join("worker.json");
let mut worker_json: serde_json::Value =
serde_json::from_slice(&std::fs::read(&worker_path).unwrap()).unwrap();
worker_json.as_object_mut().unwrap().remove("status");
std::fs::write(
&worker_path,
serde_json::to_vec_pretty(&worker_json).unwrap(),
)
.unwrap();
let restored = Runtime::with_fs_store_and_execution_backend(
options,
Arc::new(TestExecutionBackend::default()),
)
.unwrap();
assert!(restored.list_workers().unwrap().is_empty());
assert!(
restored
.diagnostics()
.unwrap()
.iter()
.any(|diagnostic| diagnostic.code == "worker_snapshot_ignored")
);
let _ = std::fs::remove_dir_all(root);
}
#[cfg(feature = "fs-store")]
#[test]
fn fs_store_shutdown_preserves_active_worker_restore_intent() {
let root = fs_store_root("shutdown-preserves-worker");
let runtime = Runtime::with_fs_store_and_execution_backend(
crate::fs_store::FsRuntimeStoreOptions {
root: root.clone(),
runtime_id: "test-runtime".to_string(),
display_name: None,
},
Arc::new(TestExecutionBackend::default()),
)
.unwrap();
runtime.store_config_bundle(test_bundle()).unwrap();
let worker = runtime
.create_worker(task_request("preserve active worker"))
.unwrap();
runtime.stop_runtime().unwrap();
let snapshot: serde_json::Value = serde_json::from_slice(
&std::fs::read(
root.join("workers")
.join(worker.worker_id.to_string())
.join("worker.json"),
)
.unwrap(),
)
.unwrap();
assert_eq!(snapshot["status"], serde_json::json!("idle"));
assert_eq!(
snapshot["execution"]["restore_intent"],
serde_json::json!("automatic")
);
let _ = std::fs::remove_dir_all(root);
}
#[cfg(feature = "fs-store")]
#[test]
fn fs_store_stopped_worker_requires_explicit_restore() {
let root = fs_store_root("stopped-explicit-restore");
let runtime = Runtime::with_fs_store_and_execution_backend(
crate::fs_store::FsRuntimeStoreOptions {
root: root.clone(),
runtime_id: "test-runtime".to_string(),
display_name: None,
},
Arc::new(TestExecutionBackend::default()),
)
.unwrap();
runtime.store_config_bundle(test_bundle()).unwrap();
let worker = runtime
.create_worker(task_request("explicit restore only"))
.unwrap();
runtime
.stop_worker(&worker.worker_ref, Some("operator stop".to_string()))
.unwrap();
drop(runtime);
let backend = Arc::new(TestExecutionBackend::default());
let restored = Runtime::with_fs_store_and_execution_backend(
crate::fs_store::FsRuntimeStoreOptions {
root: root.clone(),
runtime_id: "test-runtime".to_string(),
display_name: None,
},
backend.clone(),
)
.unwrap();
assert_eq!(*backend.restore_count.lock().unwrap(), 0);
assert_eq!(
restored.worker_detail(&worker.worker_ref).unwrap().status,
WorkerStatus::Stopped
);
assert!(matches!(
restored.send_input(&worker.worker_ref, WorkerInput::user("implicit restore")),
Err(RuntimeError::WorkerExecutionUnavailable { .. })
));
assert_eq!(*backend.restore_count.lock().unwrap(), 0);
restored.restore_worker(&worker.worker_ref).unwrap();
assert_eq!(*backend.restore_count.lock().unwrap(), 1);
assert_eq!(
restored.worker_detail(&worker.worker_ref).unwrap().status,
WorkerStatus::Idle
);
let _ = std::fs::remove_dir_all(root);
}
#[cfg(feature = "fs-store")]
#[test]
fn fs_store_migrates_schema_v3_workers_to_stopped_explicit_restore() {
let root = fs_store_root("schema-v3-restore-intent");
let options = crate::fs_store::FsRuntimeStoreOptions {
root: root.clone(),
runtime_id: "test-runtime".to_string(),
display_name: None,
};
let runtime = Runtime::with_fs_store_and_execution_backend(
options.clone(),
Arc::new(TestExecutionBackend::default()),
)
.unwrap();
runtime.store_config_bundle(test_bundle()).unwrap();
let worker = runtime
.create_worker(task_request("schema v3 worker"))
.unwrap();
drop(runtime);
let runtime_path = root.join("runtime.json");
let worker_path = root
.join("workers")
.join(worker.worker_id.to_string())
.join("worker.json");
let mut runtime_json: serde_json::Value =
serde_json::from_slice(&std::fs::read(&runtime_path).unwrap()).unwrap();
runtime_json["schema_version"] = serde_json::json!(3);
std::fs::write(
&runtime_path,
serde_json::to_vec_pretty(&runtime_json).unwrap(),
)
.unwrap();
let mut worker_json: serde_json::Value =
serde_json::from_slice(&std::fs::read(&worker_path).unwrap()).unwrap();
worker_json["schema_version"] = serde_json::json!(3);
worker_json.as_object_mut().unwrap().remove("status");
worker_json.as_object_mut().unwrap().remove("execution");
worker_json["run_generation"] = serde_json::json!(7);
std::fs::write(
&worker_path,
serde_json::to_vec_pretty(&worker_json).unwrap(),
)
.unwrap();
let backend = Arc::new(TestExecutionBackend::default());
let migrated =
Runtime::with_fs_store_and_execution_backend(options, backend.clone()).unwrap();
assert_eq!(*backend.restore_count.lock().unwrap(), 0);
assert_eq!(
migrated.worker_detail(&worker.worker_ref).unwrap().status,
WorkerStatus::Stopped
);
let migrated_json: serde_json::Value =
serde_json::from_slice(&std::fs::read(&worker_path).unwrap()).unwrap();
assert_eq!(migrated_json["schema_version"], serde_json::json!(4));
assert_eq!(migrated_json["status"], serde_json::json!("stopped"));
assert_eq!(
migrated_json["execution"]["binding"]["run_generation"],
serde_json::json!(7)
);
assert_eq!(
migrated_json["execution"]["restore_intent"],
serde_json::json!("explicit")
);
assert!(matches!(
migrated.send_input(&worker.worker_ref, WorkerInput::notify("do not restore")),
Err(RuntimeError::WorkerExecutionUnavailable { .. })
));
migrated.restore_worker(&worker.worker_ref).unwrap();
assert_eq!(*backend.restore_count.lock().unwrap(), 1);
assert_eq!(backend.run_generations.lock().unwrap().as_slice(), &[8]);
let _ = std::fs::remove_dir_all(root);
}
#[cfg(feature = "fs-store")]
#[test]
fn fs_store_stops_worker_and_reports_when_execution_restore_fails() {
@@ -5627,7 +5780,10 @@ mod tests {
WorkerInput::user("after failed restore"),
)
.unwrap_err();
assert!(matches!(err, RuntimeError::WorkerExecutionRejected { .. }));
assert!(matches!(
err,
RuntimeError::WorkerExecutionUnavailable { .. }
));
let _ = std::fs::remove_dir_all(root);
}
crates/worker/src/worker.rs
+24 -25
View File
@@ -10,8 +10,8 @@ use agen::llm_client::client::LlmClient;
use agen::llm_client::types::Role;
use agen::state::Mutable;
use agen::{
Engine, EngineError, EngineResult, EngineRunExit, History, HistoryEntry, Item, StopReason,
ToolExecutionPolicy, ToolOutputLimits, UsageRecord,
Engine, EngineError, EngineResult, EngineRunExit, History, HistoryEntry, Item,
RunInterruptionReason, ToolExecutionPolicy, ToolOutputLimits, UsageRecord,
};
use arc_swap::ArcSwap;
use session_store::{
@@ -2605,7 +2605,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
) -> bool {
if !matches!(
result,
EngineRunExit::Paused | EngineRunExit::Interrupted(StopReason::Cancelled)
EngineRunExit::Paused | EngineRunExit::Interrupted(RunInterruptionReason::Cancelled)
) {
return false;
}
@@ -3403,15 +3403,15 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
self.last_run_interrupted = true;
Ok(WorkerRunResult::Paused)
}
EngineRunExit::Interrupted(StopReason::LimitReached) => {
EngineRunExit::Interrupted(RunInterruptionReason::LimitReached) => {
self.last_run_interrupted = false;
Ok(WorkerRunResult::LimitReached)
}
EngineRunExit::Interrupted(reason) => {
self.last_run_interrupted = true;
Ok(WorkerRunResult::Interrupted {
code: stop_reason_error_code(&reason),
message: stop_reason_message(&reason),
code: run_interruption_reason_error_code(&reason),
message: run_interruption_reason_message(&reason),
})
}
EngineRunExit::Yielded => unreachable!("yielded handled above"),
@@ -3731,9 +3731,9 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
result,
EngineRunExit::Paused
| EngineRunExit::Yielded
| EngineRunExit::Interrupted(StopReason::Cancelled)
| EngineRunExit::Interrupted(StopReason::ContextWindowExceeded)
| EngineRunExit::Interrupted(StopReason::Unexpected(_))
| EngineRunExit::Interrupted(RunInterruptionReason::Cancelled)
| EngineRunExit::Interrupted(RunInterruptionReason::ContextWindowExceeded)
| EngineRunExit::Interrupted(RunInterruptionReason::Unexpected(_))
);
let active_run_turn_count = self.engine.as_ref().unwrap().active_run_turn_count();
match result {
@@ -3751,7 +3751,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
active_run_turn_count,
})?;
}
EngineRunExit::Interrupted(StopReason::LimitReached) => {
EngineRunExit::Interrupted(RunInterruptionReason::LimitReached) => {
self.commit_entry(LogEntry::RunCompleted {
ts: segment_log::now_millis(),
interrupted: false,
@@ -3763,7 +3763,7 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
self.commit_entry(LogEntry::RunErrored {
ts: segment_log::now_millis(),
interrupted,
message: stop_reason_message(reason),
message: run_interruption_reason_message(reason),
})?;
}
}
@@ -6122,15 +6122,14 @@ fn restore_manifest_from_worker_metadata_snapshot(
}
}
fn stop_reason_error_code(reason: &StopReason) -> ErrorCode {
fn run_interruption_reason_error_code(reason: &RunInterruptionReason) -> ErrorCode {
match reason {
StopReason::ContextWindowExceeded | StopReason::Unexpected(EngineError::Client(_)) => {
ErrorCode::ProviderError
}
StopReason::Unexpected(EngineError::Tool(_)) => ErrorCode::ToolError,
StopReason::LimitReached
| StopReason::Cancelled
| StopReason::Unexpected(
RunInterruptionReason::ContextWindowExceeded
| RunInterruptionReason::Unexpected(EngineError::Client(_)) => ErrorCode::ProviderError,
RunInterruptionReason::Unexpected(EngineError::Tool(_)) => ErrorCode::ToolError,
RunInterruptionReason::LimitReached
| RunInterruptionReason::Cancelled
| RunInterruptionReason::Unexpected(
EngineError::Aborted(_)
| EngineError::Cancelled
| EngineError::PauseRequested
@@ -6141,12 +6140,12 @@ fn stop_reason_error_code(reason: &StopReason) -> ErrorCode {
}
}
fn stop_reason_message(reason: &StopReason) -> String {
fn run_interruption_reason_message(reason: &RunInterruptionReason) -> String {
match reason {
StopReason::LimitReached => "engine turn limit reached".to_string(),
StopReason::ContextWindowExceeded => "model context window reached".to_string(),
StopReason::Cancelled => "engine run cancelled".to_string(),
StopReason::Unexpected(error) => format!("unexpected engine failure: {error}"),
RunInterruptionReason::LimitReached => "engine turn limit reached".to_string(),
RunInterruptionReason::ContextWindowExceeded => "model context window reached".to_string(),
RunInterruptionReason::Cancelled => "engine run cancelled".to_string(),
RunInterruptionReason::Unexpected(error) => format!("unexpected engine failure: {error}"),
}
}
@@ -8695,7 +8694,7 @@ mod build_summary_prompt_tests {
]);
let _ = worker
.handle_worker_result(
EngineRunExit::Interrupted(StopReason::Cancelled),
EngineRunExit::Interrupted(RunInterruptionReason::Cancelled),
worker.history().len(),
)
.await
crates/workspace-server/src/hosts.rs
+3 -4
View File
@@ -3772,7 +3772,6 @@ fn embedded_worker_status_label(status: EmbeddedWorkerStatus) -> &'static str {
EmbeddedWorkerStatus::Running => "running",
EmbeddedWorkerStatus::Paused => "paused",
EmbeddedWorkerStatus::Stopped => "stopped",
EmbeddedWorkerStatus::Cancelled => "cancelled",
}
}
@@ -5678,7 +5677,7 @@ mod tests {
json!({
"workers": [
worker_json_with_status("remote:primary", &worker_ids[0], "stopped"),
worker_json_with_status("remote:primary", &worker_ids[1], "cancelled"),
worker_json_with_status("remote:primary", &worker_ids[1], "running"),
worker_json_with_status("remote:primary", &worker_ids[2], "paused"),
worker_json_with_status("remote:primary", &worker_ids[3], "idle")
]
@@ -5717,11 +5716,11 @@ mod tests {
let workers = registry.list_workers(10);
assert_eq!(workers.items.len(), 4);
assert!(!workers.items[0].capabilities.can_stop);
assert!(!workers.items[1].capabilities.can_stop);
assert!(workers.items[1].capabilities.can_stop);
assert!(workers.items[2].capabilities.can_stop);
assert!(workers.items[3].capabilities.can_stop);
assert_eq!(workers.items[0].state, "stopped");
assert_eq!(workers.items[1].state, "cancelled");
assert_eq!(workers.items[1].state, "running");
assert_eq!(workers.items[2].state, "paused");
assert_eq!(workers.items[3].state, "idle");
crates/workspace-server/src/server.rs
+5 -14
View File
@@ -13199,22 +13199,13 @@ fn compensate_failed_worker_spawn(
let cancellation = api
.runtime
.cancel_worker(&worker.worker, lifecycle_request.clone());
let cancellation_accepted = cancellation
let stop = api.runtime.stop_worker(&worker.worker, lifecycle_request);
let stop_accepted = stop
.as_ref()
.is_ok_and(|result| result.state == WorkerOperationState::Accepted);
let stop = (!cancellation_accepted)
.then(|| api.runtime.stop_worker(&worker.worker, lifecycle_request));
let stop_accepted = stop.as_ref().is_some_and(|result| {
result
.as_ref()
.is_ok_and(|result| result.state == WorkerOperationState::Accepted)
});
let termination_detail = (!cancellation_accepted && !stop_accepted).then(|| {
let termination_detail = (!stop_accepted).then(|| {
let cancellation = lifecycle_failure_detail("cancel", &cancellation);
let stop = stop
.as_ref()
.map(|result| lifecycle_failure_detail("stop", result))
.unwrap_or_else(|| "stop was not attempted".to_string());
let stop = lifecycle_failure_detail("stop", &stop);
format!("{cancellation}; {stop}")
});
@@ -20287,7 +20278,7 @@ mod tests {
.iter()
.any(|assignment| assignment.role == "coder")
);
assert_eq!(api.runtime.worker(&worker).unwrap().state, "cancelled");
assert_eq!(api.runtime.worker(&worker).unwrap().state, "idle");
let Json(replayed) = scoped_cancel_ticket_implementation(State(api), path(), request())
.await
web/workspace/src/lib/generated/protocol.ts
+1 -1
View File
@@ -178,7 +178,7 @@ export type SubscriptionWorkdirId = string;
export type SubscriptionWorkerIds = Array<SubscriptionWorkerId>;
export type SubscriptionWorkerState = "idle" | "running" | "paused" | "stopped" | "cancelled";
export type SubscriptionWorkerState = "idle" | "running" | "paused" | "stopped";
export type EventSubscriptionSelector = { "topic": "runtime_workers" } | { "topic": "worker_lifecycle", worker_ids: SubscriptionWorkerIds, } | { "topic": "worker_protocol", worker_id: SubscriptionWorkerId, runtime_id?: string | null, } | { "topic": "workspace_workers" } | { "topic": "workspace_workdirs" };
web/workspace/src/lib/workspace/sidebar/worker-actions.ts
+1 -1
View File
@@ -46,7 +46,7 @@ function diagnosticMessage(
}
export function canDeleteSidebarWorker(worker: Worker): boolean {
return worker.state === "stopped" || worker.state === "cancelled";
return worker.state === "stopped";
}
export async function stopSidebarWorker(
web/workspace/src/lib/workspace/sidebar/worker-subscription.ts
+1 -1
View File
@@ -94,7 +94,7 @@ function projectWorker(worker: SubscriptionWorker): SidebarWorker {
display_hint: 'Workspace-authorized Runtime Worker',
},
capabilities: {
can_stop: worker.state !== 'stopped' && worker.state !== 'cancelled',
can_stop: worker.state !== 'stopped',
can_spawn_followup: false,
},
repository_key: worker.repository_key ?? null,
web/workspace/test/sidebar/worker-actions.test.ts
+1 -2
View File
@@ -155,10 +155,9 @@ Deno.test("sidebar Delete reports cleanup-plan blocking reasons", async () => {
);
});
Deno.test("sidebar Delete is enabled only for terminal Worker states", () => {
Deno.test("sidebar Delete is enabled only for stopped Workers", () => {
assert(!canDeleteSidebarWorker(worker));
assert(canDeleteSidebarWorker({ ...worker, state: "stopped" }));
assert(canDeleteSidebarWorker({ ...worker, state: "cancelled" }));
});
Deno.test("Worker navigation exposes an accessible hover action menu", async () => {