From cc66bd5052d1a0b9cd1cfb80ead1219683f7f082 Mon Sep 17 00:00:00 2001 From: Hare Date: Sun, 18 Oct 2026 02:00:06 +0900 Subject: [PATCH] test: connect restored Runtime parent tools to child session --- crates/workspace-server/src/server.rs | 1 + .../tests/restored_parent_control_tests.rs | 512 ++++++++++++++++++ 2 files changed, 513 insertions(+) create mode 100644 crates/workspace-server/src/server/tests/restored_parent_control_tests.rs diff --git a/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index c97fabbd4..e025a40c8 100644 --- a/crates/workspace-server/src/server.rs +++ b/crates/workspace-server/src/server.rs @@ -36969,6 +36969,7 @@ mod tests { mod embedded_owned_execution_tests; mod external_checkout_tests; mod protocol_backpressure_tests; + mod restored_parent_control_tests; mod subject_spawn_tests; mod subjektiv_jobs_tests; mod subjektiv_journal_tests; diff --git a/crates/workspace-server/src/server/tests/restored_parent_control_tests.rs b/crates/workspace-server/src/server/tests/restored_parent_control_tests.rs new file mode 100644 index 000000000..0c5ca25dc --- /dev/null +++ b/crates/workspace-server/src/server/tests/restored_parent_control_tests.rs @@ -0,0 +1,512 @@ +//! Restored Runtime parent client -> installed public Worker tools -> signed loopback +//! HTTP -> real SQLite grants -> live Runtime child -> canonical Session storage. +//! Only the external model is fixed. P is restored without submitting a model turn; +//! C stays live (this does not test child restore). Run with the normal Server lib target. +use super::*; +use agen::tool::{Tool, ToolExecutionContext, ToolOutput}; +use futures::FutureExt; +use std::{panic::AssertUnwindSafe, sync::Mutex, time::Duration}; +use worker_runtime::worker_backend::{ + ProfileRuntimeWorkerFactory, RuntimeWorkerController, RuntimeWorkerFactory, + WorkerRuntimeExecutionBackend, +}; + +const WATCHDOG: Duration = Duration::from_secs(45); + +#[derive(Clone)] +struct FixedModel; +#[async_trait::async_trait] +impl agen::llm_client::LlmClient for FixedModel { + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } + async fn stream( + &self, + _: agen::llm_client::Request, + ) -> std::result::Result + { + Ok(Box::pin(futures::stream::empty())) + } +} + +#[derive(Default)] +struct Captures { + handles: Mutex>, + restored_clients: Mutex< + Vec<( + worker_runtime::identity::WorkerRef, + Arc, + )>, + >, +} + +// Keep the production factory/restore/signers. The archive chooses an unauthenticated +// model so restoring an idle parent never depends on developer credentials. The spawn +// hook fixes only C's external model boundary; it is deliberately not a restore hook. +struct CapturingFactory { + inner: ProfileRuntimeWorkerFactory, + captures: Arc, +} +#[async_trait::async_trait] +impl RuntimeWorkerFactory for CapturingFactory { + async fn spawn_controller( + &self, + mut request: worker_runtime::execution::WorkerExecutionSpawnRequest, + ) -> std::result::Result { + let archive = worker_runtime::profile_archive::ProfileSourceArchive::build( + worker_runtime::profile_archive::ProfileSourceArchiveInput { + id: "profile-source-archive:restored-parent".into(), + entrypoints: BTreeMap::from([("builtin:companion".into(), "profiles/default.dcdl".into())]), + imports: BTreeMap::new(), + sources: BTreeMap::from([("profiles/default.dcdl".into(), r#"{ + slug = "default"; description = "Restored parent connection"; scope = "workspace_read"; + model = { scheme = "anthropic"; model_id = "test-model"; auth = { kind = "none"; }; }; + engine = { max_tokens = 100; }; + }"#.into())]), + }, + ).map_err(|e| e.to_string())?; + request.request.profile_source = + worker_runtime::catalog::ProfileSourceArchiveSource::Embedded { archive }; + let worker_ref = request.worker_ref.clone(); + let controller = self.inner.spawn_controller(request).await?; + self.captures + .handles + .lock() + .unwrap() + .push((worker_ref, controller.handle.clone())); + Ok(controller) + } + async fn preflight_restore( + &self, + request: &worker_runtime::execution::WorkerExecutionRestoreRequest, + ) -> std::result::Result<(), String> { + self.inner.preflight_restore(request).await + } + async fn restore_controller( + &self, + request: worker_runtime::execution::WorkerExecutionRestoreRequest, + ) -> std::result::Result { + let worker_ref = request.worker_ref.clone(); + let controller = self.inner.restore_controller(request).await?; + self.captures + .handles + .lock() + .unwrap() + .push((worker_ref.clone(), controller.handle.clone())); + self.captures + .restored_clients + .lock() + .unwrap() + .push((worker_ref, controller.workspace_client.clone())); + Ok(controller) + } +} + +struct Resources { + execution: Arc>, + captures: Arc, + servers: Vec>, + api: Option, +} +impl Resources { + async fn close(&mut self) { + for server in self.servers.drain(..) { + server.abort(); + assert!(server.await.unwrap_err().is_cancelled()); + } + let execution = self.execution.clone(); + let workers = self + .captures + .handles + .lock() + .unwrap() + .iter() + .map(|(id, _)| id.clone()) + .collect::>(); + // stop_worker waits for shutdown completion and joins all execution-owned tasks. + tokio::task::spawn_blocking(move || { + use worker_runtime::execution::WorkerExecutionBackend; + for worker in workers { + let result = execution.stop_worker(&worker); + assert!(result.is_accepted(), "cleanup {worker:?}: {result:?}"); + } + }) + .await + .unwrap(); + if let Some(api) = self.api.take() { + shutdown_workspace_feature_storage(&api).await.unwrap(); + } + } +} + +async fn spawn_pair_member(api: &WorkspaceApi, name: &str) -> RuntimeWorkerRef { + let Json(created) = create_workspace_worker( + State(api.clone()), + HeaderMap::new(), + Json(CreateWorkspaceWorkerRequest { + runtime_id: EMBEDDED_RUNTIME_ID.into(), + display_name: name.into(), + singleton_key: None, + profile: Some("builtin:companion".into()), + ticket_assignment: None, + initial_submit: vec![], + workdir_attachments: vec![], + feature_connections: Default::default(), + control_operation_id: None, + }), + ) + .await + .unwrap(); + RuntimeWorkerRef::new(created.runtime_id, created.worker_id) +} + +fn handle(captures: &Captures, worker: &RuntimeWorkerRef) -> worker::WorkerHandle { + captures + .handles + .lock() + .unwrap() + .iter() + .rev() + .find(|(id, _)| id.worker_id.to_string() == worker.worker_id) + .unwrap() + .1 + .clone() +} + +async fn finished_submission(events: &mut tokio::sync::broadcast::Receiver) { + let mut accepted_id = None; + let mut finished = false; + loop { + match events + .recv() + .await + .expect("controller event stream closed before Submit completion") + { + protocol::Event::SubmissionAccepted { + submission_request_id, + .. + } => { + assert!(accepted_id.is_none(), "unexpected additional Submit"); + accepted_id = Some(submission_request_id); + } + protocol::Event::SubmissionRejected { message, .. } => { + panic!("Submit rejected: {message}") + } + protocol::Event::RunEnd { result } => { + assert!(accepted_id.is_some()); + assert_eq!(result, protocol::RunResult::Finished); + finished = true; + } + protocol::Event::WorkerState { snapshot } + if finished + && snapshot.state == protocol::WorkerState::Idle + && snapshot.last_finished_submission_request_id == accepted_id => + { + return; + } + _ => {} + } + } +} + +// Public Feature installation needs a Worker tool host. This non-controller host +// uses the *restored controller's* client unchanged, not a fabricated parent signer. +async fn tools_for_restored_client( + root: &Path, + client: Arc, +) -> BTreeMap> { + let manifest = worker::WorkerManifest::from_toml( + "[worker]\nname=\"control-tools\"\n[model]\nscheme=\"anthropic\"\nmodel_id=\"unused\"\n[engine]\n[scope]\nallow=[]\n", + ).unwrap(); + let engine = + agen::Engine::<_, agen::state::Mutable, worker::SessionHistoryMetadata>::new_annotated( + FixedModel, + ); + let mut host = worker::Worker::new( + manifest, + engine, + session_store::FsStore::new(root.join("tool-host")).unwrap(), + worker::WorkerWorkspaceContext::with_client( + Some(worker::WorkspaceId::new(TEST_WORKSPACE_ID).unwrap()), + client.clone(), + ), + worker::WorkerFilesystemAuthority::None, + manifest::Scope::empty(), + ) + .await + .unwrap(); + let report = host.install_features(worker::feature::FeatureRegistryBuilder::new().with_module( + worker::feature::builtin::manage_worker::manage_worker_feature(client, None, false), + )); + assert!(!report.has_errors(), "{}", report.error_message()); + let server = host.engine().tool_server_handle(); + server.flush_pending(); + ["WorkerList", "WorkerSendInput"] + .into_iter() + .map(|name| (name.into(), server.get_tool(name).unwrap().1)) + .collect() +} +async fn call(tools: &BTreeMap>, name: &str, args: Value) -> ToolOutput { + tools[name] + .execute( + &args.to_string(), + ToolExecutionContext::new(uuid::Uuid::now_v7().to_string(), "restored-parent", 0), + ) + .await + .unwrap_or_else(|error| panic!("{name} from restored parent failed: {error}")) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn restored_runtime_parent_lists_and_inputs_existing_child() { + let root = tempfile::tempdir().unwrap(); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base = format!("http://{}", listener.local_addr().unwrap()); + let mut config = test_server_config(root.path()); + config.repositories.clear(); + config.backend_base_url = Some(base.clone()); + let factory = ProfileRuntimeWorkerFactory::new(root.path()) + .with_runtime_store_dir(config.embedded_runtime_store_root.clone()) + .with_remote_worker_mutation_identity(EMBEDDED_RUNTIME_REQUEST_IDENTITY.clone()) + .with_workspace_request_client( + worker_runtime::workspace_request::RuntimeWorkspaceRequestClient::new( + TEST_WORKSPACE_ID, + &base, + EMBEDDED_RUNTIME_ID, + ) + .with_runtime_request_source(&EMBEDDED_RUNTIME_REQUEST_IDENTITY, &base), + ) + .with_spawn_model_client_factory(|| Box::new(FixedModel)); + let captures = Arc::new(Captures::default()); + let execution = Arc::new( + WorkerRuntimeExecutionBackend::new(CapturingFactory { + inner: factory.clone(), + captures: captures.clone(), + }) + .unwrap(), + ); + let mut resources = Resources { + execution, + captures: captures.clone(), + servers: vec![], + api: None, + }; + let exercise = async { + let store = Arc::new(SqliteWorkspaceStore::open_for_test(&config.database_path).unwrap()); + let old_store = Arc::downgrade(&store); + let mut api = WorkspaceApi::new_with_execution_backend( + config.clone(), + store.clone(), + resources.execution.clone(), + ) + .await + .unwrap(); + configure_runtime_request_auth( + &mut api, + &EMBEDDED_RUNTIME_REQUEST_IDENTITY, + EMBEDDED_RUNTIME_ID, + ); + api.config.backend_base_url = Some(base.clone()); + let sources = api.config.remote_runtime_sources.clone(); + resources.api = Some(api.clone()); + let server_api = api.clone(); + resources.servers.push(tokio::spawn(async move { + axum::serve(listener, build_router(server_api)) + .await + .unwrap() + })); + let parent = spawn_pair_member(&api, "Parent").await; + let child = spawn_pair_member(&api, "Child").await; + let child_handle = handle(&captures, &child); + assert_eq!( + child_handle.shared_state.catalog_status(), + protocol::WorkerStatus::Idle + ); + // Materialize P's Session/manifest before stopping, through a real Submit. + let mut parent_events = handle(&captures, &parent).subscribe(); + let accepted = api + .runtime + .send_input( + &parent, + WorkerInputRequest { + kind: WorkerInputKind::UserIfIdle, + content: "parent bootstrap".into(), + submission_request_id: None, + segments: None, + }, + ) + .unwrap(); + assert_eq!(accepted.disposition, WorkerInputDisposition::Accepted); + finished_submission(&mut parent_events).await; + let grant = WorkerControlGrantRecord { + workspace_id: TEST_WORKSPACE_ID.into(), + grant_id: "restored-parent-child".into(), + controller: parent.clone(), + subject: child.clone(), + relation: "spawned".into(), + origin: "test".into(), + permissions: vec!["send_input".into(), "observe".into()], + operation_id: "restored-parent-child".into(), + created_at: now_registry_timestamp(), + revoked_at: None, + }; + api.store.create_worker_control_grant(&grant).unwrap(); + let grants_before = api + .store + .list_active_worker_control_grants(TEST_WORKSPACE_ID, &parent, 10) + .unwrap(); + let stopped = api + .runtime + .stop_worker( + &parent, + WorkerLifecycleRequest { + reason: Some("restart parent".into()), + ticket_assignment: None, + }, + ) + .unwrap(); + assert_eq!(stopped.state, InternalWorkerOperationState::Accepted); + assert!( + handle(&captures, &parent).protocol_is_closed(), + "stop must join the old parent controller" + ); + let server = resources.servers.pop().unwrap(); + server.abort(); + assert!(server.await.unwrap_err().is_cancelled()); + shutdown_workspace_feature_storage(&api).await.unwrap(); + resources.api.take(); + drop(api); + drop(store); + assert!( + old_store.upgrade().is_none(), + "old API/store must be released before reopen" + ); + + // Same DB and Runtime FS. No post-restart registry or grant seed. + let store = Arc::new(SqliteWorkspaceStore::open_for_test(&config.database_path).unwrap()); + let mut api = WorkspaceApi::new_with_execution_backend( + config.clone(), + store, + resources.execution.clone(), + ) + .await + .unwrap(); + api.config.remote_runtime_sources = sources; + resources.api = Some(api.clone()); + let listener = tokio::net::TcpListener::bind(base.strip_prefix("http://").unwrap()) + .await + .unwrap(); + let server_api = api.clone(); + resources.servers.push(tokio::spawn(async move { + axum::serve(listener, build_router(server_api)) + .await + .unwrap() + })); + let observation = api.runtime.worker(&parent).unwrap(); + let restored = api + .runtime + .restore_worker( + &parent, + runtime_api::WorkerRestoreRequest { + expected_observation_token: observation.restore_observation_token.unwrap(), + request_id: "restore-parent".into(), + preparation: Some(runtime_api::WorkerRestorePreparation::default()), + }, + ) + .unwrap(); + assert_eq!( + restored.state, + server_api::WorkerRestoreState::Accepted, + "{restored:?}" + ); + let restored_clients = captures.restored_clients.lock().unwrap().clone(); + assert_eq!( + restored_clients.len(), + 1, + "only P is restored; C remains the original live controller" + ); + assert_eq!( + restored_clients[0].0.worker_id.to_string(), + parent.worker_id + ); + let client = restored_clients[0].1.clone(); + assert_eq!(client.kind(), "runtime-owned-workspace-client"); + let tools = tools_for_restored_client(root.path(), client).await; + let output = call(&tools, "WorkerList", json!({})).await; + let list: Value = serde_json::from_str(output.content.as_deref().unwrap()).unwrap(); + let expected_subject = json!({"kind":"runtime_worker", "runtime_id": child.runtime_id, "worker_id": child.worker_id}); + assert_eq!(list["items"].as_array().unwrap().len(), 1, "{list}"); + assert_eq!( + list["items"][0]["subject"], expected_subject, + "restored parent's list must expose the complete child subject" + ); + assert_eq!(list["items"][0]["permissions"], json!(grant.permissions)); + assert_eq!(list["items"][0]["summary"]["runtime_id"], child.runtime_id); + assert_eq!(list["items"][0]["summary"]["worker_id"], child.worker_id); + assert_eq!( + api.store + .list_active_worker_control_grants(TEST_WORKSPACE_ID, &parent, 10) + .unwrap(), + grants_before + ); + let subject = list["items"][0]["subject"].clone(); + let marker = format!("restored-parent-marker-{}", uuid::Uuid::now_v7()); + let mut child_events = child_handle.subscribe(); + let output = call( + &tools, + "WorkerSendInput", + json!({"subject":subject,"content":marker}), + ) + .await; + let input: Value = serde_json::from_str(output.content.as_deref().unwrap()).unwrap(); + assert_eq!(input["state"], "accepted", "{input}"); + finished_submission(&mut child_events).await; + let child_ref = worker_runtime::identity::WorkerRef::try_from(&child).unwrap(); + let before = factory.retained_session_snapshot(&child_ref).unwrap(); + let bytes = serde_json::to_value(&before.snapshot).unwrap(); + assert_eq!( + before + .snapshot + .entries + .iter() + .filter(|entry| matches!( + &entry.data, + protocol::SessionSnapshotEntryData::UserInput { segments } + if segments == &vec![protocol::Segment::text(marker.clone())] + )) + .count(), + 1, + "canonical Session must contain the marker exactly once" + ); + api.store + .revoke_worker_control_grant( + TEST_WORKSPACE_ID, + &grant.grant_id, + &now_registry_timestamp(), + ) + .unwrap(); + let refused = tools["WorkerSendInput"] + .execute( + &json!({"subject": subject, "content": marker}).to_string(), + ToolExecutionContext::new("revoked-input", "restored-parent", 1), + ) + .await; + let error = refused.expect_err("revoked grant must reject the next public input"); + assert!(error.to_string().contains("unknown worker"), "{error}"); + let after = factory.retained_session_snapshot(&child_ref).unwrap(); + assert_eq!( + serde_json::to_value(&after.snapshot).unwrap(), + bytes, + "rejected input must not change C's canonical Session" + ); + }; + let outcome = AssertUnwindSafe(tokio::time::timeout(WATCHDOG, exercise)) + .catch_unwind() + .await; + tokio::time::timeout(WATCHDOG, resources.close()) + .await + .expect("fixture shutdown/join timed out"); + match outcome { + Ok(result) => result.expect("restored parent connection watchdog expired"), + Err(panic) => std::panic::resume_unwind(panic), + } +}