diff --git a/crates/workspace-server/src/hosts.rs b/crates/workspace-server/src/hosts.rs index 0b41a9a47..429ece684 100644 --- a/crates/workspace-server/src/hosts.rs +++ b/crates/workspace-server/src/hosts.rs @@ -1297,8 +1297,8 @@ pub trait WorkspaceWorkerRuntime: Send + Sync { fn observation_source( &self, _worker_id: &str, - ) -> Option { - None + ) -> Result, RuntimeDiagnostic> { + Ok(None) } fn send_input(&self, worker_id: &str, _request: WorkerInputRequest) -> WorkerInputResult { @@ -2236,6 +2236,11 @@ impl RuntimeRegistry { let runtime = self.runtime(runtime_id)?; runtime .observation_source(worker_id) + .map_err(|diagnostic| RuntimeRegistryError::RuntimeOperationFailed { + runtime_id: runtime_id.to_owned(), + code: diagnostic.code, + message: diagnostic.message, + })? .ok_or_else(|| RuntimeRegistryError::UnknownWorker { worker: worker.clone(), }) @@ -3203,17 +3208,21 @@ impl WorkspaceWorkerRuntime for EmbeddedWorkerRuntime { fn observation_source( &self, worker_id: &str, - ) -> Option { - let worker_ref = self.worker_ref(worker_id)?; + ) -> Result, RuntimeDiagnostic> { + let Some(worker_ref) = self.worker_ref(worker_id) else { + return Ok(None); + }; if self.runtime.worker_detail(&worker_ref).is_err() { - return None; + return Ok(None); } - Some(crate::observation::RuntimeObservationSource::embedded( - crate::observation::EmbeddedRuntimeObservationSource { - worker: RuntimeWorkerRef::new(&self.runtime_id, worker_id), - runtime: self.runtime.clone(), - worker_ref, - }, + Ok(Some( + crate::observation::RuntimeObservationSource::embedded( + crate::observation::EmbeddedRuntimeObservationSource { + worker: RuntimeWorkerRef::new(&self.runtime_id, worker_id), + runtime: self.runtime.clone(), + worker_ref, + }, + ), )) } @@ -5324,20 +5333,22 @@ impl WorkspaceWorkerRuntime for RemoteWorkerRuntime { fn observation_source( &self, worker_id: &str, - ) -> Option { + ) -> Result, RuntimeDiagnostic> { let path = format!("/v1/workers/{worker_id}/protocol/ws"); let bearer_token = match &self.workspace_authorization { - Some(authorization) => authorization - .issue("GET", &path, "workers:protocol", Some(worker_id), &[]) - .ok(), + Some(authorization) => { + Some(authorization.issue("GET", &path, "workers:protocol", Some(worker_id), &[])?) + } None => self.bearer_token.clone(), }; - Some(crate::observation::RuntimeObservationSource::remote_ws( - crate::observation::RuntimeObservationSourceConfig { - worker: RuntimeWorkerRef::new(&self.runtime_id, worker_id), - endpoint: self.ws_endpoint(worker_id), - bearer_token, - }, + Ok(Some( + crate::observation::RuntimeObservationSource::remote_ws( + crate::observation::RuntimeObservationSourceConfig { + worker: RuntimeWorkerRef::new(&self.runtime_id, worker_id), + endpoint: self.ws_endpoint(worker_id), + bearer_token, + }, + ), )) } diff --git a/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index 46e52a8e6..8376cdbfc 100644 --- a/crates/workspace-server/src/server.rs +++ b/crates/workspace-server/src/server.rs @@ -31468,20 +31468,7 @@ pub(crate) async fn connect_workspace_worker_protocol( worker: &RuntimeWorkerRef, input_source: Option<&protocol::AuthenticatedInputSource>, ) -> Result { - let source = match api.observation_proxy.source(worker) { - Ok(source) => source, - Err(ObservationProxyError::WorkerNotFound(_)) => api - .runtime - .observation_source(worker) - .map_err(|error| error.into_error())?, - Err(error) => { - return Err(Error::RuntimeOperationFailed { - runtime_id: worker.runtime_id.clone(), - code: error.code().to_string(), - message: error.message().to_string(), - }); - } - }; + let source = workspace_worker_protocol_source(api, worker).await?; match source { RuntimeObservationSource::RemoteWs(config) => { connect_remote_worker_protocol(config, input_source).await @@ -31494,6 +31481,37 @@ pub(crate) async fn connect_workspace_worker_protocol( } } +// Cold source discovery includes live authority/signing reads. Keep the whole +// synchronous callback on its admitted DB owner, not the Tokio WS task. +async fn workspace_worker_protocol_source( + api: &WorkspaceApi, + worker: &RuntimeWorkerRef, +) -> Result { + let source = match api.observation_proxy.source(worker) { + Ok(source) => source, + Err(ObservationProxyError::WorkerNotFound(_)) => { + let runtime = api.runtime.clone(); + let worker = worker.clone(); + api.store + .db_executor() + .run(move || { + runtime + .observation_source(&worker) + .map_err(|error| error.into_error()) + }) + .await? + } + Err(error) => { + return Err(Error::RuntimeOperationFailed { + runtime_id: worker.runtime_id.clone(), + code: error.code().to_string(), + message: error.message().to_string(), + }); + } + }; + Ok(source) +} + fn insert_authenticated_input_source_header( headers: &mut HeaderMap, input_source: Option<&protocol::AuthenticatedInputSource>, diff --git a/crates/workspace-server/src/server_db_execution_tests.rs b/crates/workspace-server/src/server_db_execution_tests.rs index 24e613341..53d94da0f 100644 --- a/crates/workspace-server/src/server_db_execution_tests.rs +++ b/crates/workspace-server/src/server_db_execution_tests.rs @@ -848,7 +848,70 @@ fn install_remote_fixture( binding } - +#[tokio::test(flavor = "current_thread")] +async fn cold_protocol_source_signing_waits_off_tokio_and_never_falls_back_on_overload_or_revocation() + { + let fixture = Fixture::new().await; + let binding = install_remote_fixture(&fixture, "protocol-runtime", "http://127.0.0.1:9", true); + let worker = RuntimeWorkerRef::new(&binding.runtime_id, "source-worker"); + let (mut busy, waiting) = BusyObservation::install_store(fixture.store.clone()); + let mut lock = WriteLock::acquire(fixture.api.config.database_path.clone()); + let writer = fixture.store.clone(); + fixture + .executor + .dispatch(move || { + writer + .mark_api_token_used("missing-token", "2026-01-01T00:00:00Z") + .unwrap() + }) + .unwrap(); + observed(waiting).await.unwrap(); // SQL holds the actual shared connection mutex. + let mut source = Box::pin(workspace_worker_protocol_source(&fixture.api, &worker)); + assert!(futures::poll!(&mut source).is_pending()); + runtime_progress_while_locked(&lock).await; + assert!(futures::poll!(&mut source).is_pending()); + assert!( + matches!( + workspace_worker_protocol_source(&fixture.api, &worker).await, + Err(Error::DbCapacityExceeded) + ), + "overload must not return an unsigned or legacy-bearer source" + ); + lock.release(); + busy.resume(); + let RuntimeObservationSource::RemoteWs(config) = observed(source).await.unwrap() else { + panic!("remote source expected") + }; + let first = config + .bearer_token + .expect("fresh Workspace capability is required"); + assert_ne!(first, "must-not-fall-back"); + let RuntimeObservationSource::RemoteWs(second) = + workspace_worker_protocol_source(&fixture.api, &worker) + .await + .unwrap() + else { + panic!("remote source expected") + }; + assert_ne!( + first, + second.bearer_token.unwrap(), + "cold lookup must sign a fresh request, not cache authentication success" + ); + fixture + .store + .revoke_workspace_runtime_binding(WORKSPACE, &binding.runtime_id, "2026-01-02T00:00:00Z") + .unwrap(); + let error = workspace_worker_protocol_source(&fixture.api, &worker) + .await + .err() + .expect("revoked authority cannot produce any source"); + assert!( + matches!(error, Error::RuntimeOperationFailed { ref code, .. } if code == "workspace_runtime_authorization_stale"), + "{error:?}" + ); + fixture.executor.drain().await.unwrap(); +} // Loopback fake at the Session/Worker Runtime HTTP boundary. It owns one reply, // released by an explicit event; Drop unblocks accept/read and joins on failure.