diff --git a/crates/worker-runtime/src/runtime.rs b/crates/worker-runtime/src/runtime.rs index 239f1c63..ca54eab1 100644 --- a/crates/worker-runtime/src/runtime.rs +++ b/crates/worker-runtime/src/runtime.rs @@ -900,8 +900,6 @@ impl Runtime { worker.run_generation.saturating_add(1).max(1), ) }; - let config_bundle = - state.resolve_config_bundle_ref(worker_request.config_bundle.as_ref())?; let backend = state.execution_backend.clone().ok_or_else(|| { RuntimeError::WorkerExecutionUnavailable { worker_id: worker_ref.worker_id.clone(), @@ -924,7 +922,7 @@ impl Runtime { context: self.execution_context(worker_ref.clone()), previous_working_directory, working_directory: None, - config_bundle, + config_bundle: None, }; (backend, request) }; @@ -1580,8 +1578,6 @@ impl Runtime { worker.run_generation.saturating_add(1).max(1), ) }; - let config_bundle = - state.resolve_config_bundle_ref(request.config_bundle.as_ref())?; state .workers .get_mut(&worker_id) @@ -1593,7 +1589,7 @@ impl Runtime { request, run_generation, previous_working_directory, - config_bundle, + config_bundle: None, }); } candidates @@ -3476,12 +3472,12 @@ mod tests { runtime.restore_worker(&detail.worker_ref).unwrap(); assert_eq!( backend.config_bundles.lock().unwrap().as_slice(), - &[Some(bundle.clone()), Some(bundle)] + &[Some(bundle), None] ); } #[test] - fn restore_fails_closed_when_recorded_config_bundle_is_missing_or_mismatched() { + fn restore_does_not_require_recorded_config_bundle() { let (runtime, backend) = runtime_and_backend(); let bundle = test_bundle(); let detail = runtime @@ -3489,11 +3485,11 @@ mod tests { .unwrap(); runtime.stop_worker(&detail.worker_ref, None).unwrap(); runtime.lock().unwrap().config_bundles.clear(); - assert!(matches!( - runtime.restore_worker(&detail.worker_ref), - Err(RuntimeError::ConfigBundleMissing { .. }) - )); - assert_eq!(backend.config_bundles.lock().unwrap().len(), 1); + runtime.restore_worker(&detail.worker_ref).unwrap(); + assert_eq!( + backend.config_bundles.lock().unwrap().as_slice(), + &[Some(bundle), None] + ); let (runtime, backend) = runtime_and_backend(); let bundle = test_bundle(); @@ -3509,11 +3505,11 @@ mod tests { .unwrap() .config_bundles .insert(replacement.metadata.id.clone(), replacement); - assert!(matches!( - runtime.restore_worker(&detail.worker_ref), - Err(RuntimeError::ConfigBundleDigestMismatch { .. }) - )); - assert_eq!(backend.config_bundles.lock().unwrap().len(), 1); + runtime.restore_worker(&detail.worker_ref).unwrap(); + assert_eq!( + backend.config_bundles.lock().unwrap().as_slice(), + &[Some(bundle), None] + ); } #[test] diff --git a/crates/worker-runtime/src/worker_backend.rs b/crates/worker-runtime/src/worker_backend.rs index bf1f3076..33016bb9 100644 --- a/crates/worker-runtime/src/worker_backend.rs +++ b/crates/worker-runtime/src/worker_backend.rs @@ -210,6 +210,71 @@ impl WorkerObservationProvider for RuntimeGrantedWorkerObservationProvider { } } +#[derive(Debug, Default)] +struct WorkspacePromptProjectionCache { + active: Mutex>>, +} + +impl WorkspacePromptProjectionCache { + fn observe( + &self, + projection: worker::WorkspacePromptProjection, + ) -> Result, String> { + projection.validate().map_err(|error| error.to_string())?; + let workspace_id = projection.workspace_id.clone(); + let projection = Arc::new(projection); + let mut active = self + .active + .lock() + .map_err(|_| "Workspace Prompt projection cache lock was poisoned".to_string())?; + if let Some(current) = active.get(&workspace_id) { + if current.config_revision > projection.config_revision { + return Ok(current.clone()); + } + if current.config_revision == projection.config_revision + && (current.projection_digest != projection.projection_digest + || current.catalog.catalog_digest != projection.catalog.catalog_digest) + { + return Err(format!( + "Workspace Prompt projection identity changed without a config revision transition: workspace={workspace_id} revision={}", + projection.config_revision + )); + } + } + active.insert(workspace_id, projection.clone()); + Ok(projection) + } + + fn fetch_current( + &self, + workspace_client: &Arc, + workspace_id: &worker::WorkspaceId, + ) -> Result, String> { + let response = workspace_client + .execute(worker::WorkspaceRequest { + method: worker::WorkspaceRequestMethod::Get, + path: format!( + "/api/w/{}/config/projections/prompts", + workspace_id.as_str() + ), + body: None, + }) + .map_err(|error| error.to_string())?; + if !(200..300).contains(&response.status) { + return Err(format!( + "Workspace Prompt projection request failed with status {}", + response.status + )); + } + let projection: worker::WorkspacePromptProjection = + serde_json::from_str(&response.body).map_err(|error| error.to_string())?; + if projection.workspace_id != workspace_id.as_str() { + return Err("Workspace Prompt projection returned a mismatched workspace".to_string()); + } + self.observe(projection) + } +} + #[derive(Clone)] pub struct ProfileRuntimeWorkerFactory { observation_hub: Arc, @@ -217,6 +282,7 @@ pub struct ProfileRuntimeWorkerFactory { worker_aggregate_root: Option, resource_client: Option>, profile_archive_cache: Arc, + prompt_projection_cache: Arc, runtime_id: Option, worker_mutation_identity: Option, embedded_worker_mutation_dispatcher: Option>, @@ -232,6 +298,7 @@ impl ProfileRuntimeWorkerFactory { worker_aggregate_root: None, resource_client: None, profile_archive_cache: Arc::new(ProfileSourceArchiveCache::default()), + prompt_projection_cache: Arc::new(WorkspacePromptProjectionCache::default()), runtime_id: None, worker_mutation_identity: None, embedded_worker_mutation_dispatcher: None, @@ -583,12 +650,30 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory { )? } }; - if let Some(prompt_catalog) = request - .config_bundle - .as_ref() - .and_then(|bundle| bundle.prompt_catalog.clone()) - { - loader = loader.with_effective_catalog(prompt_catalog); + if let Some(bundle) = request.config_bundle.as_ref() { + if let Some(prompt_catalog) = bundle.prompt_catalog.clone() { + let source_digest = bundle + .metadata + .provenance + .detail + .as_deref() + .and_then(|detail| { + detail + .split(';') + .find_map(|part| part.strip_prefix("source_tree_digest=")) + }) + .unwrap_or(&prompt_catalog.catalog_digest) + .to_string(); + let projection = worker::WorkspacePromptProjection::new( + bundle.metadata.workspace_id.clone(), + source_digest, + prompt_catalog.catalog_digest.clone(), + prompt_catalog, + ) + .map_err(|error| error.to_string())?; + let projection = self.prompt_projection_cache.observe(projection)?; + loader = loader.with_effective_catalog(projection.catalog.clone()); + } } let flow_transition_enabled = manifest.feature.flow.enabled; @@ -726,12 +811,11 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory { self.embedded_worker_mutation_dispatcher.as_ref(), ); let (manifest, mut loader) = Self::restore_fallback_manifest(&worker_name)?; - if let Some(prompt_catalog) = request - .config_bundle - .as_ref() - .and_then(|bundle| bundle.prompt_catalog.clone()) - { - loader = loader.with_effective_catalog(prompt_catalog); + if let Some(workspace_id) = workspace_context.workspace_id() { + let projection = self + .prompt_projection_cache + .fetch_current(&workspace_context.client_handle(), workspace_id)?; + loader = loader.with_effective_catalog(projection.catalog.clone()); } let worker_aggregate_dir = self.worker_aggregate_dir(&request.worker_ref)?; @@ -1859,6 +1943,75 @@ mod tests { use manifest::{Scope, WorkerManifest}; use session_store::{LogEntry, WorkerMetadataStore}; + #[derive(Debug)] + struct PromptProjectionWorkspaceClient { + response: String, + calls: AtomicUsize, + } + + impl worker::WorkspaceClient for PromptProjectionWorkspaceClient { + fn kind(&self) -> &'static str { + "prompt-projection-test" + } + + fn workspace_id(&self) -> Option<&str> { + Some("workspace-a") + } + + fn is_available(&self) -> bool { + true + } + + fn execute( + &self, + request: worker::WorkspaceRequest, + ) -> Result { + assert_eq!( + request.path, + "/api/w/workspace-a/config/projections/prompts" + ); + self.calls.fetch_add(1, Ordering::SeqCst); + Ok(worker::WorkspaceResponse { + status: 200, + body: self.response.clone(), + }) + } + } + + #[test] + fn workspace_prompt_projection_cache_fetches_and_validates_current_projection() { + let catalog = worker::EffectivePromptCatalog::new( + BTreeMap::from([("default".to_string(), "prompt".to_string())]), + 8, + "schema", + "toolchain", + ) + .unwrap(); + let projection = worker::WorkspacePromptProjection::new( + "workspace-a", + "source-digest", + catalog.catalog_digest.clone(), + catalog, + ) + .unwrap(); + let client = Arc::new(PromptProjectionWorkspaceClient { + response: serde_json::to_string(&projection).unwrap(), + calls: AtomicUsize::new(0), + }); + let cache = WorkspacePromptProjectionCache::default(); + let workspace_id = worker::WorkspaceId::new("workspace-a".to_string()).unwrap(); + + let resolved = cache + .fetch_current( + &(client.clone() as Arc), + &workspace_id, + ) + .unwrap(); + + assert_eq!(resolved.as_ref(), &projection); + assert_eq!(client.calls.load(Ordering::SeqCst), 1); + } + #[test] fn restart_restore_reconstructs_runtime_owned_worker_mutation_client() { let identity = RuntimeIdentityMaterial::generate("runtime-source").unwrap(); @@ -2512,23 +2665,14 @@ mod tests { ) .unwrap(); - let mut request = create_request("restore"); - request.workspace_api = Some(crate::catalog::WorkspaceApiRef { - workspace_id: "workspace-restore".to_string(), - base_url: "http://workspace.invalid".to_string(), - }); - let identity = RuntimeIdentityMaterial::generate("runtime-restore").unwrap(); + let request = create_request("restore"); let controller = ProfileRuntimeWorkerFactory::new(root.path()) - .with_remote_worker_mutation_identity(identity) .with_runtime_store_dir(&runtime_store_dir) .restore_controller(WorkerExecutionRestoreRequest { worker_ref: worker_ref.clone(), run_generation: 1, request, - workspace_scope: Some(crate::runtime::RuntimeWorkspaceScope::new( - "workspace-restore", - "server-main", - )), + workspace_scope: None, context: test_execution_context(worker_ref), previous_working_directory: None, working_directory: None, diff --git a/crates/worker/src/lib.rs b/crates/worker/src/lib.rs index 675efb78..ad4fd207 100644 --- a/crates/worker/src/lib.rs +++ b/crates/worker/src/lib.rs @@ -33,7 +33,8 @@ pub use manifest::{ }; pub use model_client::{ProviderError, build_client}; pub use prompt::catalog::{ - CatalogError, EffectivePromptCatalog, PromptCatalog, WorkerPrompt, prompt_schema_source, + CatalogError, EffectivePromptCatalog, PromptCatalog, WorkerPrompt, WorkspacePromptProjection, + prompt_schema_source, }; pub use prompt::source::PromptCatalogSource; pub use prompt::system::{SystemPromptContext, SystemPromptError, SystemPromptTemplate}; diff --git a/crates/worker/src/prompt/catalog.rs b/crates/worker/src/prompt/catalog.rs index 42a7ee95..28e748e9 100644 --- a/crates/worker/src/prompt/catalog.rs +++ b/crates/worker/src/prompt/catalog.rs @@ -171,6 +171,77 @@ pub enum CatalogError { DigestMismatch { expected: String, actual: String }, } +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct WorkspacePromptProjection { + pub workspace_id: String, + pub config_revision: u64, + pub source_digest: String, + pub projection_digest: String, + pub schema_fingerprint: String, + pub toolchain_fingerprint: String, + pub catalog: EffectivePromptCatalog, +} + +impl WorkspacePromptProjection { + pub fn new( + workspace_id: impl Into, + source_digest: impl Into, + projection_digest: impl Into, + catalog: EffectivePromptCatalog, + ) -> Result { + let workspace_id = workspace_id.into(); + let source_digest = source_digest.into(); + let projection_digest = projection_digest.into(); + catalog.verify_digest()?; + if workspace_id.trim().is_empty() { + return Err(CatalogError::InvalidTemplateCatalog( + "Workspace Prompt projection workspace_id must not be empty".to_string(), + )); + } + if source_digest.trim().is_empty() { + return Err(CatalogError::InvalidTemplateCatalog( + "Workspace Prompt projection source digest must not be empty".to_string(), + )); + } + if projection_digest.trim().is_empty() { + return Err(CatalogError::InvalidTemplateCatalog( + "Workspace Prompt projection digest must not be empty".to_string(), + )); + } + Ok(Self { + workspace_id, + config_revision: catalog.config_revision, + source_digest, + projection_digest, + schema_fingerprint: catalog.schema_fingerprint.clone(), + toolchain_fingerprint: catalog.toolchain_fingerprint.clone(), + catalog, + }) + } + + pub fn catalog(&self) -> &EffectivePromptCatalog { + &self.catalog + } + + pub fn validate(&self) -> Result<(), CatalogError> { + let rebuilt = Self::new( + self.workspace_id.clone(), + self.source_digest.clone(), + self.projection_digest.clone(), + self.catalog.clone(), + )?; + if rebuilt.config_revision != self.config_revision + || rebuilt.schema_fingerprint != self.schema_fingerprint + || rebuilt.toolchain_fingerprint != self.toolchain_fingerprint + { + return Err(CatalogError::InvalidTemplateCatalog( + "Workspace Prompt projection metadata does not match its catalog".to_string(), + )); + } + Ok(()) + } +} + pub struct PromptCatalog { env: Environment<'static>, projection: EffectivePromptCatalog, @@ -550,6 +621,32 @@ mod tests { )); } + #[test] + fn workspace_prompt_projection_round_trips_and_rejects_tampered_metadata() { + let catalog = EffectivePromptCatalog::new( + BTreeMap::from([("default".to_string(), "PROMPT".to_string())]), + 8, + "schema", + "toolchain", + ) + .unwrap(); + let projection = WorkspacePromptProjection::new( + "workspace-a", + "source-digest", + catalog.catalog_digest.clone(), + catalog, + ) + .unwrap(); + let serialized = serde_json::to_string(&projection).unwrap(); + let restored: WorkspacePromptProjection = serde_json::from_str(&serialized).unwrap(); + assert_eq!(restored, projection); + restored.validate().unwrap(); + + let mut tampered = restored; + tampered.config_revision += 1; + assert!(tampered.validate().is_err()); + } + #[test] fn catalog_source_preserves_workspace_projection_for_subworkers() { let templates = BTreeMap::from([("template".to_string(), "OVERRIDE".to_string())]); diff --git a/crates/workspace-server/src/profile_settings.rs b/crates/workspace-server/src/profile_settings.rs index c0ff7708..f76aff77 100644 --- a/crates/workspace-server/src/profile_settings.rs +++ b/crates/workspace-server/src/profile_settings.rs @@ -252,15 +252,25 @@ pub fn selector_for_workspace_candidate( fn validate_prompt_projection_matches_state( workspace_id: &str, state: &WorkspaceConfigState, - projection: &crate::prompt_settings::WorkspacePromptProjection, + projection: &worker::WorkspacePromptProjection, ) -> Result<()> { - if !projection.matches(workspace_id, state) { - return Err(Error::Config( - "Prompt projection identity does not match Workspace config state".to_string(), - )); - } + projection + .validate() + .map_err(|error| Error::Config(error.to_string()))?; let prompt_catalog = projection.catalog(); let mismatches = [ + ( + projection.workspace_id != workspace_id, + "workspace identity", + ), + ( + projection.source_digest != state.snapshot.digest, + "source digest", + ), + ( + projection.projection_digest != prompt_catalog.catalog_digest, + "projection digest", + ), ( prompt_catalog.config_revision != state.snapshot.revision, "config revision", @@ -353,7 +363,7 @@ pub fn build_virtual_profile_config_bundle_with_prompt_projection( workspace_id: &str, workspace_created_at: &str, selector: &str, - prompt_projection: &crate::prompt_settings::WorkspacePromptProjection, + prompt_projection: &worker::WorkspacePromptProjection, ) -> Result> { validate_prompt_projection_matches_state(workspace_id, state, prompt_projection)?; let prompt_catalog = prompt_projection.catalog().clone(); diff --git a/crates/workspace-server/src/prompt_settings.rs b/crates/workspace-server/src/prompt_settings.rs index 8f920e61..376e53e1 100644 --- a/crates/workspace-server/src/prompt_settings.rs +++ b/crates/workspace-server/src/prompt_settings.rs @@ -2,7 +2,7 @@ use std::collections::BTreeMap; use std::sync::{Arc, Mutex, OnceLock}; use config_source::{ConfigProjectionValidator, ConfigSchemaContribution}; -use worker::{EffectivePromptCatalog, prompt_schema_source}; +use worker::{EffectivePromptCatalog, WorkspacePromptProjection, prompt_schema_source}; use crate::config_source::{ WorkspaceConfigSchemaProvider, WorkspaceConfigState, evaluate_workspace_config_state, @@ -13,6 +13,7 @@ use crate::{Error, Result}; struct PromptProjectionCacheKey { workspace_id: String, config_revision: u64, + source_digest: String, projection_digest: String, schema_fingerprint: String, toolchain_fingerprint: String, @@ -23,6 +24,7 @@ impl PromptProjectionCacheKey { Self { workspace_id: workspace_id.to_string(), config_revision: state.snapshot.revision, + source_digest: state.snapshot.digest.clone(), projection_digest: state.projection_digest.clone(), schema_fingerprint: state.contract.schema_bundle.fingerprint.clone(), toolchain_fingerprint: state.contract.fingerprint.clone(), @@ -30,22 +32,6 @@ impl PromptProjectionCacheKey { } } -#[derive(Debug, Clone)] -pub struct WorkspacePromptProjection { - identity: PromptProjectionCacheKey, - catalog: EffectivePromptCatalog, -} - -impl WorkspacePromptProjection { - pub fn catalog(&self) -> &EffectivePromptCatalog { - &self.catalog - } - - pub fn matches(&self, workspace_id: &str, state: &WorkspaceConfigState) -> bool { - self.identity == PromptProjectionCacheKey::new(workspace_id, state) - } -} - type PromptProjectionCell = OnceLock, String>>; #[derive(Debug, Default)] @@ -242,10 +228,14 @@ pub fn project_workspace_prompt_projection( workspace_id: &str, state: &WorkspaceConfigState, ) -> Result { - Ok(WorkspacePromptProjection { - identity: PromptProjectionCacheKey::new(workspace_id, state), - catalog: project_prompts_from_workspace_config(state)?, - }) + let catalog = project_prompts_from_workspace_config(state)?; + WorkspacePromptProjection::new( + workspace_id, + state.snapshot.digest.clone(), + catalog.catalog_digest.clone(), + catalog, + ) + .map_err(|error| Error::RegistryInconsistency(error.to_string())) } #[cfg(test)] diff --git a/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index 7f44904d..fc50dae1 100644 --- a/crates/workspace-server/src/server.rs +++ b/crates/workspace-server/src/server.rs @@ -1243,6 +1243,10 @@ pub fn build_router(api: WorkspaceApi) -> Router { "/api/w/{workspace_id}/config/source-tree", get(scoped_get_workspace_config_tree), ) + .route( + "/api/w/{workspace_id}/config/projections/prompts", + get(scoped_get_prompt_projection), + ) .route( "/api/w/{workspace_id}/config/source-tree/commit", post(scoped_commit_workspace_config_tree), @@ -2581,6 +2585,25 @@ struct WorkspaceConfigEntryPath { path: String, } +async fn scoped_get_prompt_projection( + State(api): State, + AxumPath(path): AxumPath, +) -> ApiResult> { + validate_workspace_scope(&api, &path.workspace_id)?; + let state = api + .config_store + .load_workspace_config(&path.workspace_id)? + .ok_or_else(|| { + ApiError::from(Error::InvalidInput( + "Workspace config is not initialized".to_string(), + )) + })?; + let projection = api + .prompt_projection_cache + .resolve(&path.workspace_id, &state)?; + Ok(Json(projection.as_ref().clone())) +} + #[derive(Debug, Serialize)] struct WorkspaceConfigTreeResponse { snapshot: ConfigTreeSnapshot,