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

fix: prevent subject worker launch stalls

Browse Source
This commit is contained in:
Hare
2026-10-05 01:22:58 +09:00
parent bd79b93c1d
commit 685b8aa65f
4 changed files with 170 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/worker-runtime/src/worker_source.rs
+32 -5
View File
@@ -381,10 +381,10 @@ impl RuntimeOwnedWorkspaceClient {
self
}
fn execute_with_permission(
fn execute_with_optional_timeout(
&self,
request: WorkspaceRequest,
permission: &'static str,
timeout: Option<Duration>,
) -> Result<WorkspaceResponse, WorkspaceClientError> {
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<WorkspaceResponse, WorkspaceClientError> {
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<WorkspaceResponse, WorkspaceClientError> {
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));
crates/worker/src/feature/builtin/memory.rs
+42 -6
View File
@@ -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<server_api::SubjektivMemoryBackendResponse, WorkspaceMemoryBackendError> {
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<crate::worker::WorkspaceResponse, WorkspaceClientError>,
) -> Result<server_api::SubjektivMemoryBackendResponse, WorkspaceMemoryBackendError> {
let workspace_id =
client
@@ -163,11 +179,14 @@ async fn execute_subjektiv_memory_operation(
client.kind()
),
})?;
let response = client.execute(WorkspaceRequest::json(
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<String>,
scope_allowed: bool,
paths: Mutex<Vec<String>>,
timeouts: Mutex<Vec<Duration>>,
}
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<WorkspaceResponse, WorkspaceClientError> {
self.timeouts.lock().unwrap().push(timeout);
self.execute(request)
}
}
fn resident_client(content: &str) -> Arc<dyn WorkspaceClient> {
@@ -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,
crates/worker/src/worker.rs
+11
View File
@@ -738,6 +738,17 @@ pub trait WorkspaceClient: std::fmt::Debug + Send + Sync {
fn execute(&self, request: WorkspaceRequest)
-> Result<WorkspaceResponse, WorkspaceClientError>;
/// 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<WorkspaceResponse, WorkspaceClientError> {
self.execute(request)
}
fn execute_server_operation(
&self,
operation: WorkspaceServerOperation,
crates/workspace-server/src/server.rs
+82 -6
View File
@@ -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<WorkerSpawnResult> {
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]