fix: share compiled prompt projection cache

crates/worker-runtime/src/worker_backend.rs

@@ -219,14 +219,26 @@ impl WorkerObservationProvider for RuntimeGrantedWorkerObservationProvider {
 
 #[derive(Debug, Default)]
 pub(crate) struct WorkspacePromptProjectionCache {
-    active: Mutex<HashMap<String, Arc<worker::WorkspacePromptProjection>>>,
+    active: Mutex<HashMap<String, Arc<worker::WorkspacePromptCatalogResolution>>>,
+    fetch_gates: Mutex<HashMap<String, Arc<Mutex<()>>>>,
 }
 
 impl WorkspacePromptProjectionCache {
+    pub(crate) fn fetch_gate(&self, workspace_id: &str) -> Result<Arc<Mutex<()>>, String> {
+        let mut gates = self
+            .fetch_gates
+            .lock()
+            .map_err(|_| "Workspace Prompt projection fetch gates lock was poisoned".to_string())?;
+        Ok(gates
+            .entry(workspace_id.to_string())
+            .or_insert_with(|| Arc::new(Mutex::new(())))
+            .clone())
+    }
+
     pub(crate) fn active(
         &self,
         workspace_id: &str,
-    ) -> Result<Option<Arc<worker::WorkspacePromptProjection>>, String> {
+    ) -> Result<Option<Arc<worker::WorkspacePromptCatalogResolution>>, String> {
         self.active
             .lock()
             .map(|active| active.get(workspace_id).cloned())
@@ -236,31 +248,40 @@ impl WorkspacePromptProjectionCache {
     pub(crate) fn observe(
         &self,
         projection: worker::WorkspacePromptProjection,
-    ) -> Result<Arc<worker::WorkspacePromptProjection>, String> {
+    ) -> Result<Arc<worker::WorkspacePromptCatalogResolution>, String> {
         projection.validate().map_err(|error| error.to_string())?;
         let workspace_id = projection.workspace_id.clone();
-        let projection = Arc::new(projection);
+        let resolution = Arc::new(
+            worker::WorkspacePromptCatalogResolution::new(projection)
+                .map_err(|error| error.to_string())?,
+        );
         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 {
+            if current.projection.config_revision > resolution.projection.config_revision {
                 return Ok(current.clone());
             }
-            if current.config_revision == projection.config_revision
+            if current.projection.config_revision == resolution.projection.config_revision
-                && (current.source_digest != projection.source_digest
+                && (current.projection.source_digest != resolution.projection.source_digest
-                    || current.projection_digest != projection.projection_digest
+                    || current.projection.projection_digest
-                    || current.catalog.catalog_digest != projection.catalog.catalog_digest)
+                        != resolution.projection.projection_digest
+                    || current.projection.catalog.catalog_digest
+                        != resolution.projection.catalog.catalog_digest
+                    || current.projection.catalog.schema_fingerprint
+                        != resolution.projection.catalog.schema_fingerprint
+                    || current.projection.catalog.toolchain_fingerprint
+                        != resolution.projection.catalog.toolchain_fingerprint)
             {
                 return Err(format!(
                     "Workspace Prompt projection identity changed without a config revision transition: workspace={workspace_id} revision={}",
-                    projection.config_revision
+                    resolution.projection.config_revision
                 ));
             }
         }
-        active.insert(workspace_id, projection.clone());
+        active.insert(workspace_id, resolution.clone());
-        Ok(projection)
+        Ok(resolution)
     }
 }
 
@@ -677,7 +698,7 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory {
                 )
                 .map_err(|error| error.to_string())?;
                 let projection = self.prompt_projection_cache.observe(projection)?;
-                loader = loader.with_effective_catalog(projection.catalog.clone());
+                loader = loader.with_effective_catalog(projection.projection.catalog.clone());
             }
         }
         let flow_transition_enabled = manifest.feature.flow.enabled;
@@ -1995,8 +2016,11 @@ mod tests {
         cache.observe(projection_v2).unwrap();
 
         let active = cache.active("workspace-a").unwrap().unwrap();
-        assert_eq!(active.config_revision, 9);
+        assert_eq!(active.projection.config_revision, 9);
-        assert_eq!(active.catalog.catalog_digest, catalog_v2.catalog_digest);
+        assert_eq!(
+            active.projection.catalog.catalog_digest,
+            catalog_v2.catalog_digest
+        );
     }
 
     #[test]
@@ -2033,6 +2057,35 @@ mod tests {
         );
     }
 
+    #[test]
+    fn workspace_prompt_projection_cache_rejects_same_revision_schema_drift() {
+        let templates = BTreeMap::from([("default".to_string(), "prompt".to_string())]);
+        let first_catalog =
+            worker::EffectivePromptCatalog::new(templates.clone(), 8, "schema-a", "toolchain")
+                .unwrap();
+        let drifted_catalog =
+            worker::EffectivePromptCatalog::new(templates, 8, "schema-b", "toolchain").unwrap();
+        let first = worker::WorkspacePromptProjection::new(
+            "workspace-a",
+            "source-a",
+            first_catalog.catalog_digest.clone(),
+            first_catalog,
+        )
+        .unwrap();
+        let drifted = worker::WorkspacePromptProjection::new(
+            "workspace-a",
+            "source-a",
+            drifted_catalog.catalog_digest.clone(),
+            drifted_catalog,
+        )
+        .unwrap();
+        let cache = WorkspacePromptProjectionCache::default();
+
+        cache.observe(first).unwrap();
+        let error = cache.observe(drifted).unwrap_err();
+        assert!(error.contains("without a config revision transition"));
+    }
+
     #[test]
     fn restart_restore_reconstructs_runtime_owned_worker_mutation_client() {
         let identity = RuntimeIdentityMaterial::generate("runtime-source").unwrap();

crates/worker-runtime/src/worker_source.rs

@@ -2,8 +2,8 @@ use std::sync::Arc;
 use std::time::{Duration, SystemTime, UNIX_EPOCH};
 
 use worker::{
-    WorkspaceClient, WorkspaceClientError, WorkspacePromptProjection, WorkspaceRequest,
+    WorkspaceClient, WorkspaceClientError, WorkspacePromptCatalogResolution,
-    WorkspaceRequestMethod, WorkspaceResponse,
+    WorkspacePromptProjection, WorkspaceRequest, WorkspaceRequestMethod, WorkspaceResponse,
 };
 
 use crate::auth::{
@@ -287,6 +287,7 @@ fn execute_remote_worker_remove_http_blocking(
     Ok(WorkspaceResponse { status, body })
 }
 
+#[derive(Clone)]
 pub struct RuntimeOwnedWorkspaceClient {
     workspace_id: String,
     base_url: String,
@@ -398,15 +399,29 @@ impl WorkspaceClient for RuntimeOwnedWorkspaceClient {
 
     fn current_prompt_projection(
         &self,
-    ) -> Result<Option<WorkspacePromptProjection>, WorkspaceClientError> {
+    ) -> Result<Option<WorkspacePromptCatalogResolution>, WorkspaceClientError> {
         let Some(cache) = self.prompt_projection_cache.as_ref() else {
             return Ok(None);
         };
-        if let Some(projection) = cache
+        if let Some(resolution) = cache
             .active(&self.workspace_id)
             .map_err(WorkspaceClientError::Request)?
         {
-            return Ok(Some((*projection).clone()));
+            return Ok(Some((*resolution).clone()));
+        }
+        let fetch_gate = cache
+            .fetch_gate(&self.workspace_id)
+            .map_err(WorkspaceClientError::Request)?;
+        let _fetch_guard = fetch_gate.lock().map_err(|_| {
+            WorkspaceClientError::Request(
+                "Workspace Prompt projection fetch gate was poisoned".to_string(),
+            )
+        })?;
+        if let Some(resolution) = cache
+            .active(&self.workspace_id)
+            .map_err(WorkspaceClientError::Request)?
+        {
+            return Ok(Some((*resolution).clone()));
         }
         let response = self.execute(WorkspaceRequest::get(format!(
             "/api/w/{}/config-sources/active/prompt-projection",
@@ -602,9 +617,78 @@ mod tests {
         .with_prompt_projection_cache(cache);
 
         let projection = client.current_prompt_projection().unwrap().unwrap();
+        let second = client.current_prompt_projection().unwrap().unwrap();
 
-        assert_eq!(projection.config_revision, 3);
+        assert_eq!(projection.projection.config_revision, 3);
-        assert_eq!(projection.source_digest, "source-3");
+        assert_eq!(projection.projection.source_digest, "source-3");
+        assert!(Arc::ptr_eq(&projection.catalog, &second.catalog));
+    }
+
+    #[test]
+    fn concurrent_prompt_projection_miss_fetches_once_and_shares_catalog() {
+        use std::io::{Read, Write};
+        use std::net::TcpListener;
+        use std::sync::Barrier;
+
+        let catalog = worker::EffectivePromptCatalog::new(
+            std::collections::BTreeMap::from([(
+                "default".to_string(),
+                "shared prompt".to_string(),
+            )]),
+            5,
+            "schema",
+            "toolchain",
+        )
+        .unwrap();
+        let projection = WorkspacePromptProjection::new(
+            "workspace-a",
+            "source-5",
+            catalog.catalog_digest.clone(),
+            catalog,
+        )
+        .unwrap();
+        let body = serde_json::to_string(&projection).unwrap();
+        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
+        let base_url = format!("http://{}", listener.local_addr().unwrap());
+        let server = std::thread::spawn(move || {
+            let (mut stream, _) = listener.accept().unwrap();
+            let mut request = [0_u8; 4096];
+            let _ = stream.read(&mut request).unwrap();
+            write!(
+                stream,
+                "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
+                body.len(),
+                body
+            )
+            .unwrap();
+        });
+        let cache = Arc::new(WorkspacePromptProjectionCache::default());
+        let client =
+            RuntimeOwnedWorkspaceClient::new("workspace-a", base_url, "runtime-a", "worker-a")
+                .with_prompt_projection_cache(cache);
+        let barrier = Arc::new(Barrier::new(8));
+        let threads = (0..8)
+            .map(|_| {
+                let client = client.clone();
+                let barrier = barrier.clone();
+                std::thread::spawn(move || {
+                    barrier.wait();
+                    client.current_prompt_projection().unwrap().unwrap()
+                })
+            })
+            .collect::<Vec<_>>();
+        let resolutions = threads
+            .into_iter()
+            .map(|thread| thread.join().unwrap())
+            .collect::<Vec<_>>();
+        server.join().unwrap();
+
+        let first = &resolutions[0].catalog;
+        assert!(
+            resolutions
+                .iter()
+                .all(|resolution| Arc::ptr_eq(first, &resolution.catalog))
+        );
     }
 
     #[test]

crates/worker/src/lib.rs

@@ -45,7 +45,7 @@ pub use shared_state::WorkerSharedState;
 pub use worker::{
     LocalWorkingDirectory, WORKER_INPUT_SUBMISSION_EXTENSION_DOMAIN, Worker, WorkerError,
     WorkerFilesystemAuthority, WorkerRunResult, WorkerWorkspaceContext, WorkspaceClient,
-    WorkspaceClientError, WorkspaceId, WorkspaceIdError, WorkspaceRequest, WorkspaceRequestMethod,
+    WorkspaceClientError, WorkspaceId, WorkspaceIdError, WorkspacePromptCatalogResolution,
-    WorkspaceResponse, apply_worker_manifest, marker_workspace_client,
+    WorkspaceRequest, WorkspaceRequestMethod, WorkspaceResponse, apply_worker_manifest,
-    unavailable_workspace_client,
+    marker_workspace_client, unavailable_workspace_client,
 };

crates/worker/src/worker.rs

@@ -212,6 +212,38 @@ pub enum WorkspaceClientError {
     Request(String),
 }
 
+#[derive(Clone)]
+pub struct WorkspacePromptCatalogResolution {
+    pub projection: Arc<WorkspacePromptProjection>,
+    pub catalog: Arc<PromptCatalog>,
+}
+
+impl std::fmt::Debug for WorkspacePromptCatalogResolution {
+    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        formatter
+            .debug_struct("WorkspacePromptCatalogResolution")
+            .field("workspace_id", &self.projection.workspace_id)
+            .field("config_revision", &self.projection.config_revision)
+            .field("source_digest", &self.projection.source_digest)
+            .field("projection_digest", &self.projection.projection_digest)
+            .finish_non_exhaustive()
+    }
+}
+
+impl WorkspacePromptCatalogResolution {
+    pub fn new(projection: WorkspacePromptProjection) -> Result<Self, CatalogError> {
+        projection.validate()?;
+        let catalog = PromptCatalog::load(
+            &PromptCatalogSource::builtins_only()
+                .with_effective_catalog(projection.catalog.clone()),
+        )?;
+        Ok(Self {
+            projection: Arc::new(projection),
+            catalog,
+        })
+    }
+}
+
 /// Path-free Workspace operation authority injected by Runtime/host code.
 ///
 /// Workers receive this trait object rather than a Backend URL. The concrete
@@ -229,7 +261,7 @@ pub trait WorkspaceClient: std::fmt::Debug + Send + Sync {
     /// launch/session state; this hook never reconstructs historical prompts.
     fn current_prompt_projection(
         &self,
-    ) -> Result<Option<WorkspacePromptProjection>, WorkspaceClientError> {
+    ) -> Result<Option<WorkspacePromptCatalogResolution>, WorkspaceClientError> {
         Ok(None)
     }
 
@@ -296,7 +328,7 @@ impl WorkspaceClient for ReviewerChildWorkspaceClient {
 
     fn current_prompt_projection(
         &self,
-    ) -> Result<Option<WorkspacePromptProjection>, WorkspaceClientError> {
+    ) -> Result<Option<WorkspacePromptCatalogResolution>, WorkspaceClientError> {
         self.inner.current_prompt_projection()
     }
 
@@ -1258,7 +1290,7 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
         if self.system_prompt_template.is_some() {
             return Ok(());
         }
-        let Some(projection) = self
+        let Some(resolution) = self
             .workspace_context
             .client()
             .current_prompt_projection()
@@ -1268,25 +1300,16 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
         else {
             return Ok(());
         };
-        projection
+        let projection = &resolution.projection;
-            .validate()
-            .map_err(|source| WorkerError::WorkspacePromptProjection {
-                message: source.to_string(),
-            })?;
         let current = self.prompts.load();
         if current.projection().config_revision == projection.config_revision
             && current.projection().source_digest == projection.source_digest
             && current.projection().catalog_digest == projection.projection_digest
+            && Arc::ptr_eq(&current, &resolution.catalog)
         {
             return Ok(());
         }
-        let catalog = PromptCatalog::load(
+        self.prompts.store(resolution.catalog);
-            &PromptCatalogSource::builtins_only().with_effective_catalog(projection.catalog),
-        )
-        .map_err(|source| WorkerError::WorkspacePromptProjection {
-            message: source.to_string(),
-        })?;
-        self.prompts.store(catalog);
         Ok(())
     }