diff --git a/crates/worker/src/bootstrap.rs b/crates/worker/src/bootstrap.rs index cb7771a86..d7871660b 100644 --- a/crates/worker/src/bootstrap.rs +++ b/crates/worker/src/bootstrap.rs @@ -185,13 +185,6 @@ where .await .map_err(WorkerBootstrapError::Worker)?; - crate::feature::authorize_requested_manifest( - worker.workspace_client_handle().as_ref(), - worker.manifest(), - ) - .map_err(|error| { - WorkerBootstrapError::Worker(WorkerError::FeatureLifecycle(error.to_string())) - })?; if let Some(workdir_session) = self.workdir_session { worker.bind_single_workdir_session(Some(workdir_session)); } diff --git a/crates/worker/src/controller.rs b/crates/worker/src/controller.rs index 22e528770..df6a0475d 100644 --- a/crates/worker/src/controller.rs +++ b/crates/worker/src/controller.rs @@ -1403,14 +1403,11 @@ where // Worker-immutable snapshots taken before the mutable worker borrow // below so the worker borrow doesn't conflict with reads on `worker`. let feature_config = worker.manifest().feature.clone(); - // First boundary: before any Feature initialization, Tool/WIP/service - // registration, prompt contribution, or discovery side effect. Host derives - // the live caller grant; the Profile only requests these switches. - crate::feature::authorize_requested_manifest( - worker.workspace_client_handle().as_ref(), - worker.manifest(), - ) - .map_err(|error| std::io::Error::new(std::io::ErrorKind::PermissionDenied, error))?; + // New activation requires admission before any initialization side effect. + // Restore reconstructs saved composition; operation authority stays live. + worker + .authorize_feature_configuration() + .map_err(|error| std::io::Error::new(std::io::ErrorKind::PermissionDenied, error))?; let workdir_sessions = worker.workdir_sessions(); let restrict_workdir_tools_to_current_capabilities = inherited_workdir_tool_broker.is_some(); let workdir_tool_broker = Some( diff --git a/crates/worker/src/feature.rs b/crates/worker/src/feature.rs index 6c5bf2fdd..096cbe998 100644 --- a/crates/worker/src/feature.rs +++ b/crates/worker/src/feature.rs @@ -1534,8 +1534,18 @@ impl ProtocolProviderRegistrar<'_> { } } +/// Restore preserves the saved Feature composition; only new activation asks +/// for admission. Neither mode grants authority for protected operations. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub(crate) enum FeatureInstallation { + #[default] + Activate, + Restore, +} + /// Install-time context provided to a feature module. pub struct FeatureInstallContext<'a> { + installation: FeatureInstallation, feature_id: &'a FeatureId, declarations: &'a FeatureContributionDeclarations, pending_tools: &'a mut Vec, @@ -1548,6 +1558,10 @@ pub struct FeatureInstallContext<'a> { } impl FeatureInstallContext<'_> { + pub(crate) fn is_restoring(&self) -> bool { + self.installation == FeatureInstallation::Restore + } + pub fn feature_id(&self) -> &FeatureId { self.feature_id } @@ -1751,6 +1765,7 @@ pub enum FeaturePlanError { /// Builder/installer for enabled feature modules. pub struct FeatureRegistryBuilder { modules: Vec>, + installation: FeatureInstallation, activation_guard: Option Result<(), String> + Send + Sync>>, } @@ -1998,6 +2013,7 @@ impl Default for FeatureRegistryBuilder { fn default() -> Self { Self { modules: Vec::new(), + installation: FeatureInstallation::Activate, activation_guard: None, } } @@ -2008,7 +2024,13 @@ impl FeatureRegistryBuilder { Self::default() } - /// Add a trusted, live preflight that runs before *any* module install. Worker + pub(crate) fn with_installation(mut self, installation: FeatureInstallation) -> Self { + self.installation = installation; + self + } + + /// Add a trusted, live preflight that runs before *any* new module activation. + /// Restoring saved composition does not repeat admission checks. Worker /// hosts use this in addition to the earlier pre-construction preflight. pub fn with_activation_guard( mut self, @@ -2095,10 +2117,11 @@ impl FeatureRegistryBuilder { hook_builder: &mut HookRegistryBuilder, mut installed_tool_names: HashMap, ) -> FeatureRegistryInstallReport { - if let Some(error) = self - .activation_guard - .as_ref() - .and_then(|guard| guard().err()) + if self.installation == FeatureInstallation::Activate + && let Some(error) = self + .activation_guard + .as_ref() + .and_then(|guard| guard().err()) { return FeatureRegistryInstallReport { reports: { @@ -2231,6 +2254,7 @@ impl FeatureRegistryBuilder { let install_result = { let mut context = FeatureInstallContext { + installation: self.installation, feature_id: &descriptor.id, declarations: &declarations, pending_tools, diff --git a/crates/worker/src/feature/builtin/ticket.rs b/crates/worker/src/feature/builtin/ticket.rs index 704c5744d..b8adff036 100644 --- a/crates/worker/src/feature/builtin/ticket.rs +++ b/crates/worker/src/feature/builtin/ticket.rs @@ -822,14 +822,15 @@ impl FeatureModule for TicketFeature { fn install(&self, context: &mut FeatureInstallContext<'_>) -> Result<(), FeatureInstallError> { let workspace_client = self.workspace_client(); - // Direct registry installs do not necessarily have a matching manifest - // switch. The public module identity still requires live activation - // authority before any Tool, instruction or service contribution. - workspace_client - .authorize_feature_activation(FEATURE_ID) - .map_err(|error| { - FeatureInstallError::Install(format!("Ticket activation rejected: {error}")) - })?; + // Direct new registry installs may not have a matching manifest switch. + // Restore retains saved composition; operations still require live authority. + if !context.is_restoring() { + workspace_client + .authorize_feature_activation(FEATURE_ID) + .map_err(|error| { + FeatureInstallError::Install(format!("Ticket activation rejected: {error}")) + })?; + } let ticket_service: Arc = Arc::new(WorkspaceTicketService { backend: WorkspaceHttpTicketBackend::new(workspace_client.clone()), access: self.access, @@ -4370,6 +4371,42 @@ mod tests { } } + #[test] + fn restored_ticket_keeps_tools_and_services_but_rejects_revoked_operation() { + let client = Arc::new(CreateAuthorizationClient { + workspace: Some("workspace"), + activation_allowed: false, + allowed: false, + backend_status: 200, + requests: Default::default(), + }); + let mut tools = Vec::new(); + let report = FeatureRegistryBuilder::new() + .with_module(ticket_tools_feature( + client.clone(), + TicketFeatureAccess::workspace_authoring(), + )) + .with_installation(crate::feature::FeatureInstallation::Restore) + .install_into_pending(&mut tools, &mut HookRegistryBuilder::default()); + assert!(!report.has_errors(), "{}", report.error_message()); + assert!(tools.iter().any(|tool| tool().0.name == "TicketCreate")); + assert!(!report.installed_instruction_contributions().is_empty()); + assert!(client.requests.lock().unwrap().is_empty()); + let service = report + .services + .service::(&ServiceId::builtin(TICKET_SERVICE_ID)) + .unwrap(); + let error = service + .create(NewTicket::new("Revoked operation")) + .unwrap_err(); + assert!(error.to_string().contains("HTTP status 403"), "{error}"); + let requests = client.requests.lock().unwrap(); + assert_eq!(requests.len(), 1, "denial must precede the backend effect"); + assert!(requests[0].path.ends_with("/authorization/check")); + let body: Value = serde_json::from_str(requests[0].body.as_deref().unwrap()).unwrap(); + assert_eq!(body["operation"], "ticket.create"); + } + #[test] fn installed_ticket_service_readonly_profile_cannot_create_with_host_permission() { let client = Arc::new(CreateAuthorizationClient { diff --git a/crates/worker/src/worker.rs b/crates/worker/src/worker.rs index 9db91a7fa..ac9d44d6e 100644 --- a/crates/worker/src/worker.rs +++ b/crates/worker/src/worker.rs @@ -3488,6 +3488,7 @@ pub struct Worker { prepared_initial_seed: Option>, prepared_initial_prompt_completion: Option<(String, PromptContributionCompletion)>, subjektiv_host: Option>, + feature_installation: crate::feature::FeatureInstallation, /// One-shot generic Feature lifecycle notification required after restore. restore_feature_lifecycle_pending: bool, /// Complete system prompt replacement installed by an enabled Feature. @@ -3759,6 +3760,7 @@ impl Worker { prepared_initial_seed: None, prepared_initial_prompt_completion: None, subjektiv_host: None, + feature_installation: crate::feature::FeatureInstallation::Activate, restore_feature_lifecycle_pending: false, feature_system_prompt_override: None, user_segments: Vec::new(), @@ -4202,6 +4204,18 @@ impl Worker { self.feature_invocations.clone() } + pub(crate) fn authorize_feature_configuration(&self) -> Result<(), WorkspaceClientError> { + match self.feature_installation { + crate::feature::FeatureInstallation::Activate => { + crate::feature::authorize_requested_manifest( + self.workspace_client_handle().as_ref(), + self.manifest(), + ) + } + crate::feature::FeatureInstallation::Restore => Ok(()), + } + } + /// Install enabled feature modules into the Worker host surfaces. pub fn install_features( &mut self, @@ -4210,10 +4224,12 @@ impl Worker { self.backend_job.get_or_init(|| None); let client = self.workspace_client_handle(); let manifest = self.manifest().clone(); - let registry = registry.with_activation_guard(move || { - crate::feature::authorize_requested_manifest(client.as_ref(), &manifest) - .map_err(|error| error.to_string()) - }); + let registry = registry + .with_installation(self.feature_installation) + .with_activation_guard(move || { + crate::feature::authorize_requested_manifest(client.as_ref(), &manifest) + .map_err(|error| error.to_string()) + }); let worker = self.engine.as_mut().expect("worker taken during run"); let report = registry.install_into_engine(worker, &mut self.hook_builder); if report.has_errors() { @@ -8277,6 +8293,11 @@ where // session_id + segment_id are allocated here so the worker-allocation // registration can record them from the start. Pending lifecycle // recovery reuses its already persisted Session identity. + let feature_installation = if existing_session_id.is_some() { + crate::feature::FeatureInstallation::Restore + } else { + crate::feature::FeatureInstallation::Activate + }; let session_id = existing_session_id.unwrap_or_else(session_store::new_session_id); let segment_id = session_store::new_segment_id(); @@ -8360,6 +8381,7 @@ where prepared_initial_seed: None, prepared_initial_prompt_completion: None, subjektiv_host: None, + feature_installation, restore_feature_lifecycle_pending: false, feature_system_prompt_override: None, user_segments: Vec::new(), @@ -8760,6 +8782,7 @@ where prepared_initial_seed: None, prepared_initial_prompt_completion: None, subjektiv_host: None, + feature_installation: crate::feature::FeatureInstallation::Activate, restore_feature_lifecycle_pending: false, feature_system_prompt_override: None, user_segments: Vec::new(), @@ -8893,6 +8916,7 @@ where prepared_initial_seed: None, prepared_initial_prompt_completion: None, subjektiv_host: None, + feature_installation: crate::feature::FeatureInstallation::Activate, restore_feature_lifecycle_pending: false, feature_system_prompt_override: None, user_segments: Vec::new(), @@ -9305,6 +9329,7 @@ where prepared_initial_seed: None, prepared_initial_prompt_completion: None, subjektiv_host: None, + feature_installation: crate::feature::FeatureInstallation::Restore, restore_feature_lifecycle_pending: true, feature_system_prompt_override: None, user_segments: state.user_segments, @@ -18195,6 +18220,234 @@ permission = "write" assert!(writer.state.append_lock.try_lock().is_ok()); } + // Production bootstrap, Controller registration, and durable stores; only the + // Workspace transport and model are replaced. No live Server is needed. + #[derive(Debug, Default)] + struct FeatureAdmissionClient { + requests: Mutex>, + response_status: Mutex>, + } + + impl WorkspaceClient for FeatureAdmissionClient { + fn workspace_id(&self) -> Option<&str> { + Some("feature-restore-workspace") + } + + fn kind(&self) -> &str { + "feature-admission-test" + } + + fn is_available(&self) -> bool { + true + } + + fn execute( + &self, + request: WorkspaceRequest, + ) -> Result { + assert!(request.path.ends_with("/authorization/check")); + self.requests.lock().unwrap().push(request); + match *self.response_status.lock().unwrap() { + Some(status) => Ok(WorkspaceResponse { + status, + body: if status == 200 { + r#"{"scope":null}"#.into() + } else { + r#"{"message":"current grant denied"}"#.into() + }, + }), + None => Err(WorkspaceClientError::Unavailable("Host not serving".into())), + } + } + } + + fn feature_restore_bootstrap( + temp: &tempfile::TempDir, + manifest: WorkerManifest, + store: BodyPublicationStore, + client: Arc, + ) -> crate::WorkerBootstrap { + crate::WorkerBootstrap::new( + manifest, + store, + PromptCatalogSource::builtins_only(), + WorkerWorkspaceContext::with_client( + Some(WorkspaceId::new("feature-restore-workspace").unwrap()), + client, + ), + WorkerFilesystemAuthority::None, + crate::WorkerBootstrapLayout::RuntimeManagedRun { + run_dir: temp + .path() + .join("runtime/runs") + .join(uuid::Uuid::now_v7().to_string()), + }, + crate::WorkerControllerTransport::InProcess, + ) + .with_model_client(NoopClient) + } + + fn feature_restore_fixture(temp: &tempfile::TempDir) -> (WorkerManifest, BodyPublicationStore) { + let mut manifest = minimal_manifest(); + manifest.worker.name = format!("feature-restore-{}", uuid::Uuid::now_v7()); + manifest.scope.allow.clear(); + manifest.feature.task.enabled = true; + manifest.feature.ticket.enabled = true; + manifest.profile = Some(manifest::ProfileManifestSnapshot { + source: manifest::ProfileSource::Registry { + source: manifest::ProfileRegistrySource::Builtin, + name: "default".into(), + path: None, + provenance: None, + }, + profile: None, + }); + let store = session_store::CombinedStore::new( + session_store::FsStore::new(temp.path().join("sessions")).unwrap(), + session_store::FsWorkerStore::new(temp.path().join("workers")).unwrap(), + ); + (manifest, store) + } + + #[tokio::test] + async fn new_worker_rejects_denied_feature_before_persisting_metadata() { + let temp = tempfile::tempdir().unwrap(); + let _sandbox = worker_allocation::test_util::RuntimeDirSandbox::new(temp.path()); + let (manifest, store) = feature_restore_fixture(&temp); + let client = Arc::new(FeatureAdmissionClient::default()); + *client.response_status.lock().unwrap() = Some(403); + let result = + feature_restore_bootstrap(&temp, manifest.clone(), store.clone(), client.clone()) + .prepare() + .await; + assert!(matches!( + result, + Err(crate::WorkerBootstrapError::Worker( + WorkerError::FeatureLifecycle(_) + )) + )); + assert!(store.read_by_name(&manifest.worker.name).unwrap().is_none()); + assert!(!client.requests.lock().unwrap().is_empty()); + assert!(!temp.path().join("runtime").exists()); + } + + #[tokio::test] + async fn new_worker_rechecks_feature_admission_before_controller_publication() { + let temp = tempfile::tempdir().unwrap(); + let _sandbox = worker_allocation::test_util::RuntimeDirSandbox::new(temp.path()); + let (manifest, store) = feature_restore_fixture(&temp); + let client = Arc::new(FeatureAdmissionClient::default()); + *client.response_status.lock().unwrap() = Some(200); + let prepared = feature_restore_bootstrap(&temp, manifest, store, client.clone()) + .prepare() + .await + .unwrap(); + client.requests.lock().unwrap().clear(); + *client.response_status.lock().unwrap() = Some(403); + match prepared.start().await { + Err(crate::WorkerBootstrapError::Controller { source, .. }) => { + assert_eq!(source.kind(), std::io::ErrorKind::PermissionDenied); + assert!(source.to_string().contains("HTTP status 403")); + } + Err(other) => panic!("unexpected rejection: {other:?}"), + Ok(controller) => { + stop_body_publication_controller(controller).await; + panic!("new activation must not bypass admission"); + } + } + assert!(!client.requests.lock().unwrap().is_empty()); + } + + #[tokio::test] + async fn restore_preserves_saved_features_and_prompt_without_feature_readmission() { + for materialized in [false, true] { + for response_status in [Some(403), None] { + let temp = tempfile::tempdir().unwrap(); + let _sandbox = worker_allocation::test_util::RuntimeDirSandbox::new(temp.path()); + let (manifest, store) = feature_restore_fixture(&temp); + let client = Arc::new(FeatureAdmissionClient::default()); + *client.response_status.lock().unwrap() = Some(200); + let prepared = feature_restore_bootstrap( + &temp, + manifest.clone(), + store.clone(), + client.clone(), + ) + .prepare() + .await + .unwrap(); + let session = prepared.worker().session_id(); + let mut saved = Vec::new(); + let mut saved_prompt = None; + let mut saved_segment = None; + if materialized { + let segment = prepared.worker().segment_id(); + let controller = prepared.start().await.unwrap(); + saved = store.read_all(session, segment).unwrap(); + saved_prompt = session_store::collect_state(&saved).system_prompt; + assert!(saved_prompt.is_some()); + saved_segment = Some(segment); + stop_body_publication_controller(controller).await; + } else { + drop(prepared); + assert!(store.list_segments(session).unwrap().is_empty()); + } + client.requests.lock().unwrap().clear(); + *client.response_status.lock().unwrap() = response_status; + let mut fallback = manifest.clone(); + fallback.feature.task.enabled = false; + fallback.feature.ticket.enabled = false; + fallback.feature.workdir_catalog.enabled = false; + fallback.engine.instruction = "must not replace saved instructions".into(); + let restored = + feature_restore_bootstrap(&temp, fallback, store.clone(), client.clone()) + .prepare_restored(&manifest.worker.name) + .await + .unwrap(); + assert_eq!(restored.worker().session_id(), session); + assert_eq!(restored.worker().manifest().feature, manifest.feature); + assert_eq!( + restored.worker().manifest().engine.instruction, + manifest.engine.instruction + ); + if let Some(prompt) = &saved_prompt { + assert_eq!( + restored.worker().engine().get_system_prompt(), + Some(prompt.as_str()) + ); + assert_eq!(Some(restored.worker().segment_id()), saved_segment); + } + assert!(client.requests.lock().unwrap().is_empty()); + let controller = restored.start().await.unwrap(); + let installed: WorkerManifest = + toml::from_str(&controller.handle.shared_state.manifest_toml).unwrap(); + assert_eq!(installed.feature, manifest.feature); + let committed = controller.handle.committed_entries(); + if materialized { + assert_eq!( + serde_json::to_value(&committed[..saved.len()]).unwrap(), + serde_json::to_value(&saved).unwrap(), + ); + assert_eq!( + session_store::collect_state(&committed).system_prompt, + saved_prompt + ); + } else { + assert!( + session_store::collect_state(&committed) + .system_prompt + .is_some() + ); + } + assert!( + client.requests.lock().unwrap().is_empty(), + "restore must not ask for Feature admission" + ); + stop_body_publication_controller(controller).await; + } + } + } + fn minimal_manifest() -> WorkerManifest { let toml_str = r#" [worker]