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

fix: preserve saved feature configuration during worker restore

Browse Source
This commit is contained in:
Hare
2026-10-20 22:00:39 +09:00
parent 70d153b365
commit 3cfe04be3a
5 changed files with 336 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/worker/src/bootstrap.rs
-7
View File
@@ -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));
}
crates/worker/src/controller.rs
+4 -7
View File
@@ -1403,13 +1403,10 @@ 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(),
)
// 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();
crates/worker/src/feature.rs
+26 -2
View File
@@ -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<ToolDefinition>,
@@ -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<Arc<dyn FeatureModule>>,
installation: FeatureInstallation,
activation_guard: Option<Arc<dyn Fn() -> 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,7 +2117,8 @@ impl FeatureRegistryBuilder {
hook_builder: &mut HookRegistryBuilder,
mut installed_tool_names: HashMap<String, FeatureId>,
) -> FeatureRegistryInstallReport {
if let Some(error) = self
if self.installation == FeatureInstallation::Activate
&& let Some(error) = self
.activation_guard
.as_ref()
.and_then(|guard| guard().err())
@@ -2231,6 +2254,7 @@ impl FeatureRegistryBuilder {
let install_result = {
let mut context = FeatureInstallContext {
installation: self.installation,
feature_id: &descriptor.id,
declarations: &declarations,
pending_tools,
crates/worker/src/feature/builtin/ticket.rs
+40 -3
View File
@@ -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.
// 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<dyn TicketService> = 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::<dyn TicketService>(&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 {
crates/worker/src/worker.rs
+254 -1
View File
@@ -3488,6 +3488,7 @@ pub struct Worker<C: LlmClient, St: Store> {
prepared_initial_seed: Option<Vec<LogEntry>>,
prepared_initial_prompt_completion: Option<(String, PromptContributionCompletion)>,
subjektiv_host: Option<Arc<dyn crate::subjektiv::SubjektivHost>>,
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<C: LlmClient + 'static, St: Store> Worker<C, St> {
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<C: LlmClient + 'static, St: Store> Worker<C, St> {
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,7 +4224,9 @@ impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
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 || {
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())
});
@@ -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<Vec<WorkspaceRequest>>,
response_status: Mutex<Option<u16>>,
}
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<WorkspaceResponse, WorkspaceClientError> {
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<FeatureAdmissionClient>,
) -> crate::WorkerBootstrap<BodyPublicationStore> {
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]