diff --git a/crates/worker-runtime/src/worker_source.rs b/crates/worker-runtime/src/worker_source.rs index 7454d53e..15760540 100644 --- a/crates/worker-runtime/src/worker_source.rs +++ b/crates/worker-runtime/src/worker_source.rs @@ -399,6 +399,7 @@ impl WorkspaceClient for RuntimeOwnedWorkspaceClient { fn current_prompt_projection( &self, + minimum_revision: Option, ) -> Result, WorkspaceClientError> { let Some(cache) = self.prompt_projection_cache.as_ref() else { return Ok(None); @@ -406,6 +407,11 @@ impl WorkspaceClient for RuntimeOwnedWorkspaceClient { if let Some(resolution) = cache .active(&self.workspace_id) .map_err(WorkspaceClientError::Request)? + .filter(|resolution| { + minimum_revision + .map(|minimum| resolution.projection.config_revision >= minimum) + .unwrap_or(true) + }) { return Ok(Some((*resolution).clone())); } @@ -420,6 +426,11 @@ impl WorkspaceClient for RuntimeOwnedWorkspaceClient { if let Some(resolution) = cache .active(&self.workspace_id) .map_err(WorkspaceClientError::Request)? + .filter(|resolution| { + minimum_revision + .map(|minimum| resolution.projection.config_revision >= minimum) + .unwrap_or(true) + }) { return Ok(Some((*resolution).clone())); } @@ -445,10 +456,18 @@ impl WorkspaceClient for RuntimeOwnedWorkspaceClient { self.workspace_id, projection.workspace_id ))); } - let projection = cache + let resolution = cache .observe(projection) .map_err(WorkspaceClientError::Request)?; - Ok(Some((*projection).clone())) + if let Some(minimum_revision) = minimum_revision + && resolution.projection.config_revision < minimum_revision + { + return Err(WorkspaceClientError::Request(format!( + "active Workspace Prompt projection is stale: required revision {minimum_revision}, got {}", + resolution.projection.config_revision + ))); + } + Ok(Some((*resolution).clone())) } fn execute_worker_remove( @@ -616,14 +635,60 @@ mod tests { ) .with_prompt_projection_cache(cache); - let projection = client.current_prompt_projection().unwrap().unwrap(); - let second = client.current_prompt_projection().unwrap().unwrap(); + let projection = client.current_prompt_projection(None).unwrap().unwrap(); + let second = client.current_prompt_projection(None).unwrap().unwrap(); assert_eq!(projection.projection.config_revision, 3); assert_eq!(projection.projection.source_digest, "source-3"); assert!(Arc::ptr_eq(&projection.catalog, &second.catalog)); } + #[test] + fn prompt_projection_minimum_revision_rejects_stale_server_response() { + use std::io::{Read, Write}; + use std::net::TcpListener; + + let catalog = worker::EffectivePromptCatalog::new( + std::collections::BTreeMap::from([("default".to_string(), "stale prompt".to_string())]), + 3, + "schema", + "toolchain", + ) + .unwrap(); + let projection = WorkspacePromptProjection::new( + "workspace-a", + "source-3", + catalog.catalog_digest.clone(), + catalog, + ) + .unwrap(); + let body = serde_json::to_string(&projection).unwrap(); + let cache = Arc::new(WorkspacePromptProjectionCache::default()); + cache.observe(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 client = + RuntimeOwnedWorkspaceClient::new("workspace-a", base_url, "runtime-a", "worker-a") + .with_prompt_projection_cache(cache); + + let error = client.current_prompt_projection(Some(4)).unwrap_err(); + server.join().unwrap(); + + assert!(error.to_string().contains("required revision 4, got 3")); + } + #[test] fn concurrent_prompt_projection_miss_fetches_once_and_shares_catalog() { use std::io::{Read, Write}; @@ -673,7 +738,7 @@ mod tests { let barrier = barrier.clone(); std::thread::spawn(move || { barrier.wait(); - client.current_prompt_projection().unwrap().unwrap() + client.current_prompt_projection(None).unwrap().unwrap() }) }) .collect::>(); @@ -732,7 +797,7 @@ mod tests { RuntimeOwnedWorkspaceClient::new("workspace-a", base_url, "runtime-a", "worker-a") .with_prompt_projection_cache(Arc::new(WorkspacePromptProjectionCache::default())); - let error = client.current_prompt_projection().unwrap_err(); + let error = client.current_prompt_projection(None).unwrap_err(); server.join().unwrap(); assert!(error.to_string().contains("scope mismatch")); diff --git a/crates/worker/src/worker.rs b/crates/worker/src/worker.rs index f618980a..374a496f 100644 --- a/crates/worker/src/worker.rs +++ b/crates/worker/src/worker.rs @@ -261,6 +261,7 @@ pub trait WorkspaceClient: std::fmt::Debug + Send + Sync { /// launch/session state; this hook never reconstructs historical prompts. fn current_prompt_projection( &self, + _minimum_revision: Option, ) -> Result, WorkspaceClientError> { Ok(None) } @@ -328,8 +329,9 @@ impl WorkspaceClient for ReviewerChildWorkspaceClient { fn current_prompt_projection( &self, + minimum_revision: Option, ) -> Result, WorkspaceClientError> { - self.inner.current_prompt_projection() + self.inner.current_prompt_projection(minimum_revision) } fn execute( @@ -1293,7 +1295,7 @@ impl Worker { let Some(resolution) = self .workspace_context .client() - .current_prompt_projection() + .current_prompt_projection(None) .map_err(|source| WorkerError::WorkspacePromptProjection { message: source.to_string(), })?