From 685b8aa65f0e37f757034719e3fb7b9742029f62 Mon Sep 17 00:00:00 2001 From: Hare Date: Mon, 5 Oct 2026 01:22:58 +0900 Subject: [PATCH] fix: prevent subject worker launch stalls --- crates/worker-runtime/src/worker_source.rs | 37 +++++++-- crates/worker/src/feature/builtin/memory.rs | 54 ++++++++++--- crates/worker/src/worker.rs | 11 +++ crates/workspace-server/src/server.rs | 88 +++++++++++++++++++-- 4 files changed, 170 insertions(+), 20 deletions(-) diff --git a/crates/worker-runtime/src/worker_source.rs b/crates/worker-runtime/src/worker_source.rs index bb700ae41..089ed9f8f 100644 --- a/crates/worker-runtime/src/worker_source.rs +++ b/crates/worker-runtime/src/worker_source.rs @@ -381,10 +381,10 @@ impl RuntimeOwnedWorkspaceClient { self } - fn execute_with_permission( + fn execute_with_optional_timeout( &self, request: WorkspaceRequest, - permission: &'static str, + timeout: Option, ) -> Result { let method = match request.method { WorkspaceRequestMethod::Get => reqwest::Method::GET, @@ -409,9 +409,9 @@ impl RuntimeOwnedWorkspaceClient { path_and_query: request.path, body, headers, - permission: permission.to_string(), + permission: WORKSPACE_REQUEST_PERMISSION.to_string(), worker_id: Some(self.worker_id.clone()), - timeout: self.request_timeout, + timeout, max_response_bytes: 8 * 1024 * 1024, }) .map_err(|error| { @@ -478,7 +478,15 @@ impl WorkspaceClient for RuntimeOwnedWorkspaceClient { &self, request: WorkspaceRequest, ) -> Result { - self.execute_with_permission(request, WORKSPACE_REQUEST_PERMISSION) + self.execute_with_optional_timeout(request, self.request_timeout) + } + + fn execute_with_timeout( + &self, + request: WorkspaceRequest, + timeout: Duration, + ) -> Result { + self.execute_with_optional_timeout(request, Some(timeout)) } fn execute_server_operation( @@ -1341,6 +1349,25 @@ mod tests { server.join().unwrap(); } + #[test] + fn runtime_owned_workspace_client_honors_per_request_timeouts() { + let (base_url, server) = delayed_workspace_response(Duration::from_millis(75)); + let client = + RuntimeOwnedWorkspaceClient::new("workspace-a", base_url, "runtime-a", "worker-a"); + + let error = client + .execute_with_timeout( + WorkspaceRequest::json(WorkspaceRequestMethod::Post, "/api/test", "{}"), + Duration::from_millis(20), + ) + .unwrap_err(); + assert_eq!(client.request_timeout, None); + let message = error.to_string(); + assert!(message.contains("POST /api/test"), "{message}"); + assert!(message.contains("timed out"), "{message}"); + server.join().unwrap(); + } + #[test] fn runtime_owned_workspace_client_reports_request_timeouts() { let (base_url, server) = delayed_workspace_response(Duration::from_millis(75)); diff --git a/crates/worker/src/feature/builtin/memory.rs b/crates/worker/src/feature/builtin/memory.rs index 62411f540..d1ed3689d 100644 --- a/crates/worker/src/feature/builtin/memory.rs +++ b/crates/worker/src/feature/builtin/memory.rs @@ -5,6 +5,7 @@ //! workspace backend instead of resolving `.yoi/memory` from a Worker workdir. use std::sync::Arc; +use std::time::Duration; use agen::tool::{Tool, ToolDefinition, ToolError, ToolExecutionContext, ToolMeta, ToolOutput}; use arc_swap::ArcSwap; @@ -150,9 +151,24 @@ impl dyn WorkspaceClient + '_ { } } +const RESIDENT_SUMMARY_REQUEST_TIMEOUT: Duration = Duration::from_secs(5); + async fn execute_subjektiv_memory_operation( client: &dyn WorkspaceClient, operation: server_api::SubjektivMemoryBackendOperation, +) -> Result { + execute_subjektiv_memory_operation_with(client, operation, |client, request| { + client.execute(request) + }) +} + +fn execute_subjektiv_memory_operation_with( + client: &dyn WorkspaceClient, + operation: server_api::SubjektivMemoryBackendOperation, + execute: impl FnOnce( + &dyn WorkspaceClient, + WorkspaceRequest, + ) -> Result, ) -> Result { let workspace_id = client @@ -163,11 +179,14 @@ async fn execute_subjektiv_memory_operation( client.kind() ), })?; - let response = client.execute(WorkspaceRequest::json( - WorkspaceRequestMethod::Post, - format!("/api/w/{workspace_id}/subjektiv/memory"), - serde_json::to_string(&server_api::SubjektivMemoryBackendRequest { operation })?, - ))?; + let response = execute( + client, + WorkspaceRequest::json( + WorkspaceRequestMethod::Post, + format!("/api/w/{workspace_id}/subjektiv/memory"), + serde_json::to_string(&server_api::SubjektivMemoryBackendRequest { operation })?, + ), + )?; let status = reqwest::StatusCode::from_u16(response.status) .unwrap_or(reqwest::StatusCode::INTERNAL_SERVER_ERROR); if !response.is_success() { @@ -699,12 +718,13 @@ impl SystemPromptContributionSource for WorkspaceResidentSummarySource { #[async_trait] impl SystemPromptContributionSource for WorkspaceSubjektivResidentSummarySource { async fn load(&self) -> SystemPromptContribution { - match execute_subjektiv_memory_operation( + match execute_subjektiv_memory_operation_with( self.client.as_ref(), server_api::SubjektivMemoryBackendOperation::ResidentSummary(Default::default()), - ) - .await - { + |client, request| { + client.execute_with_timeout(request, RESIDENT_SUMMARY_REQUEST_TIMEOUT) + }, + ) { Ok(server_api::SubjektivMemoryBackendResponse::ResidentSummary(output)) => { match resident_summary_contribution(output) { SystemPromptContribution::Ready(content) if content.trim().is_empty() => { @@ -1113,6 +1133,7 @@ mod tests { content: Option, scope_allowed: bool, paths: Mutex>, + timeouts: Mutex>, } impl SubjectResidentClient { @@ -1125,6 +1146,7 @@ mod tests { content: content.map(str::to_string), scope_allowed: true, paths: Mutex::new(Vec::new()), + timeouts: Mutex::new(Vec::new()), } } @@ -1134,6 +1156,7 @@ mod tests { content: Some("must not be injected".into()), scope_allowed: false, paths: Mutex::new(Vec::new()), + timeouts: Mutex::new(Vec::new()), } } } @@ -1182,6 +1205,15 @@ mod tests { .unwrap(), }) } + + fn execute_with_timeout( + &self, + request: WorkspaceRequest, + timeout: Duration, + ) -> Result { + self.timeouts.lock().unwrap().push(timeout); + self.execute(request) + } } fn resident_client(content: &str) -> Arc { @@ -1490,6 +1522,10 @@ permission = "write" ready_client.paths.lock().unwrap().clone(), vec!["/api/w/workspace/subjektiv/memory".to_string()] ); + assert_eq!( + ready_client.timeouts.lock().unwrap().as_slice(), + &[RESIDENT_SUMMARY_REQUEST_TIMEOUT] + ); let ready_empty_client = Arc::new(SubjectResidentClient::new( memory::backend::MemoryResidentSummaryAvailability::Ready, diff --git a/crates/worker/src/worker.rs b/crates/worker/src/worker.rs index 348fd6705..c3e7257f7 100644 --- a/crates/worker/src/worker.rs +++ b/crates/worker/src/worker.rs @@ -738,6 +738,17 @@ pub trait WorkspaceClient: std::fmt::Debug + Send + Sync { fn execute(&self, request: WorkspaceRequest) -> Result; + /// Executes one bounded Workspace request when the caller is on a + /// latency-sensitive lifecycle boundary. Implementations without transport + /// timeout support retain their existing execution semantics. + fn execute_with_timeout( + &self, + request: WorkspaceRequest, + _timeout: Duration, + ) -> Result { + self.execute(request) + } + fn execute_server_operation( &self, operation: WorkspaceServerOperation, diff --git a/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index e0bd7bc16..c0124335e 100644 --- a/crates/workspace-server/src/server.rs +++ b/crates/workspace-server/src/server.rs @@ -162,11 +162,11 @@ use crate::hosts::{ RuntimeRegistryError, RuntimeRegistryUnregisterResult, TicketWorkerRole, WorkerCompletionsRequest, WorkerCompletionsResult, WorkerControlOperation, WorkerCreateBinding, WorkerInputDisposition, WorkerInputKind, WorkerInputRequest, WorkerInputResult, - WorkerLifecycleRequest, WorkerLifecycleResult, WorkerSpawnAcceptanceRequirement, - WorkerSpawnIntent, WorkerSpawnRequest, WorkerSpawnResult, WorkerSpawnWorkingDirectoryRequest, - WorkerTicketAssignmentRequest, WorkspaceRuntimeAuthorization, - is_disallowed_remote_runtime_address, is_loopback_runtime_origin, - worker_spawn_create_fingerprint, workspace_worker_summary, + WorkerLifecycleRequest, WorkerLifecycleResult, WorkerSpawnAcceptanceEvidence, + WorkerSpawnAcceptanceRequirement, WorkerSpawnIntent, WorkerSpawnRequest, WorkerSpawnResult, + WorkerSpawnWorkingDirectoryRequest, WorkerTicketAssignmentRequest, + WorkspaceRuntimeAuthorization, is_disallowed_remote_runtime_address, + is_loopback_runtime_origin, worker_spawn_create_fingerprint, workspace_worker_summary, }; use crate::memory_backend::execute_memory_backend_operation_with_authority; use crate::memory_staging::{ @@ -27610,6 +27610,11 @@ async fn create_workspace_worker_inner( Ok(result) => result, Err(error) => return Err(error.into()), }; + let result = if resolved_subjektiv_attached { + restore_reused_subject_worker(&api, result)? + } else { + result + }; Ok(Json(record_browser_worker_spawn( &api, runtime_id, @@ -27619,6 +27624,52 @@ async fn create_workspace_worker_inner( )?)) } +fn restore_reused_subject_worker( + api: &WorkspaceApi, + mut result: WorkerSpawnResult, +) -> ApiResult { + if result.state != InternalWorkerOperationState::Accepted { + return Ok(result); + } + let Some(worker) = result.worker.as_ref() else { + return Ok(result); + }; + if worker.state != "stopped" { + return Ok(result); + } + + let worker_ref = worker.worker.clone(); + let restored = api.restore_workspace_worker(&worker_ref)?; + if restored.state != server_api::WorkerRestoreState::Accepted { + let message = runtime_diagnostics_message(&restored.diagnostics); + return Err(ApiError::with_diagnostics( + Error::RuntimeOperationFailed { + runtime_id: worker_ref.runtime_id, + code: "subject_worker_restore_rejected".to_string(), + message, + }, + restored.diagnostics, + )); + } + let worker = restored + .worker + .ok_or_else(|| Error::RuntimeOperationFailed { + runtime_id: worker_ref.runtime_id, + code: "subject_worker_restore_missing_summary".to_string(), + message: "Runtime restored the existing subject Worker without returning its summary" + .to_string(), + })?; + result.worker = Some(worker); + result.diagnostics.extend(restored.diagnostics); + result + .acceptance_evidence + .push(WorkerSpawnAcceptanceEvidence { + kind: "subject_worker_restored".to_string(), + detail: "the existing stopped subject Worker was restored".to_string(), + }); + Ok(result) +} + fn record_browser_worker_spawn( api: &WorkspaceApi, requested_runtime_id: String, @@ -39361,7 +39412,7 @@ mod tests { async fn ordinary_worker_launch_validates_optional_subjektiv_connection() { let workspace = tempfile::tempdir().unwrap(); init_clean_git_workspace(workspace.path()); - let api = test_api(workspace.path()).await; + let (api, execution) = test_api_with_recording_backend(workspace.path()).await; let store = open_subjektiv_store(&api).unwrap(); let active = store .create_subject(crate::subjektiv::SubjectRole::new("companion").unwrap()) @@ -39456,6 +39507,31 @@ mod tests { first.worker.singleton_key.as_deref(), Some(format!("subjektiv:{}", active.id).as_str()) ); + + let worker = RuntimeWorkerRef::new(&first.runtime_id, &first.worker_id); + let stopped = api + .runtime + .stop_worker( + &worker, + WorkerLifecycleRequest { + reason: Some("test subject Worker restart".to_string()), + ticket_assignment: None, + }, + ) + .unwrap(); + assert_eq!(stopped.state, InternalWorkerOperationState::Accepted); + assert_eq!(api.runtime.worker(&worker).unwrap().state, "stopped"); + execution.accept_restores(); + + let Json(restored) = create_workspace_worker( + State(api.clone()), + HeaderMap::new(), + Json(request(Some(&active.id), "builtin:companion", None)), + ) + .await + .unwrap(); + assert_eq!(restored.worker_id, first.worker_id); + assert_eq!(restored.worker.state, "idle"); } #[tokio::test]