diff --git a/codex-rs/config/src/cloud_config_bundle.rs b/codex-rs/config/src/cloud_config_bundle.rs index 62b65ecb20..0f6daac894 100644 --- a/codex-rs/config/src/cloud_config_bundle.rs +++ b/codex-rs/config/src/cloud_config_bundle.rs @@ -343,6 +343,16 @@ pub enum CloudConfigBundleBindingStatus { } impl CloudConfigBundleBinding { + /// Completes when this captured policy can no longer authorize EMA requests. + pub async fn invalidated(&self) { + let mut changes = self.policy.0.subscribe(); + while self.read().status == CloudConfigBundleBindingStatus::Current { + if changes.changed().await.is_err() { + return; + } + } + } + /// Keeps the policy observation stable until this guard is dropped. pub fn read(&self) -> CloudConfigBundleBindingGuard<'_> { let state = self.policy.0.borrow(); diff --git a/codex-rs/rmcp-client/src/ema_http_client.rs b/codex-rs/rmcp-client/src/ema_http_client.rs index df778a1dce..5ec435e6dc 100644 --- a/codex-rs/rmcp-client/src/ema_http_client.rs +++ b/codex-rs/rmcp-client/src/ema_http_client.rs @@ -37,21 +37,51 @@ pub trait EmaTokenSource: Send + Sync { } /// Revocable authority held while one EMA connection awaits HTTP responses. -#[derive(Clone, Default)] +#[derive(Clone)] pub struct EmaRequestAuthority { revoked: CancellationToken, + is_current: Arc bool + Send + Sync>, responses: Arc>, } +impl Default for EmaRequestAuthority { + fn default() -> Self { + Self::new(CancellationToken::new(), || true) + } +} + impl EmaRequestAuthority { + /// Couples prompt cancellation with a synchronous check at each HTTP boundary. + /// The check must be cheap and must not perform credential or network I/O. + pub fn new( + revoked: CancellationToken, + is_current: impl Fn() -> bool + Send + Sync + 'static, + ) -> Self { + Self { + revoked, + is_current: Arc::new(is_current), + responses: Arc::default(), + } + } + pub fn revoke(&self) { self.revoked.cancel(); } pub fn is_revoked(&self) -> bool { + if !self.revoked.is_cancelled() && !(self.is_current)() { + self.revoke(); + } self.revoked.is_cancelled() } + pub fn check(&self) -> Result<(), ExecServerError> { + if self.is_revoked() { + return Err(revoked_authority_error()); + } + Ok(()) + } + /// Wait until response futures admitted before revocation complete or are cancelled. pub async fn quiesce(&self) { let _responses = self.responses.clone().write_owned().await; @@ -59,9 +89,7 @@ impl EmaRequestAuthority { async fn acquire(&self) -> Result, ExecServerError> { let response = self.responses.clone().read_owned().await; - if self.is_revoked() { - return Err(revoked_authority_error()); - } + self.check()?; Ok(response) } @@ -229,10 +257,12 @@ impl HttpClient for EmaAuthenticatedHttpClient { let request = async { let _authority = self.authority.acquire().await?; let token = self.authenticate(&mut params).await?; + self.authority.check()?; if let Some(deadline) = deadline { params.timeout_ms = Some(remaining_timeout(deadline)?); } let response = self.inner.http_request(params).await?; + self.authority.check()?; self.validate_response(&response, &token)?; Ok(response) }; @@ -256,10 +286,12 @@ impl HttpClient for EmaAuthenticatedHttpClient { let request = async { let _authority = self.authority.acquire().await?; let token = self.authenticate(&mut params).await?; + self.authority.check()?; if let Some(deadline) = deadline { params.timeout_ms = Some(remaining_timeout(deadline)?); } let response = self.inner.http_request_stream(params).await?; + self.authority.check()?; self.validate_response(&response.0, &token)?; Ok(( response.0, diff --git a/codex-rs/rmcp-client/src/ema_http_client_tests.rs b/codex-rs/rmcp-client/src/ema_http_client_tests.rs index 1ff154edf9..82e8668d83 100644 --- a/codex-rs/rmcp-client/src/ema_http_client_tests.rs +++ b/codex-rs/rmcp-client/src/ema_http_client_tests.rs @@ -210,13 +210,19 @@ async fn revocation_during_renewal_prevents_dispatch() -> anyhow::Result<()> { let server = MockServer::start().await; let endpoint = format!("{}/mcp", server.uri()); let (gate, validated_rx, release_tx) = CredentialGate::new(); - let (client, source, params) = fixture(&endpoint); + let (mut client, source, params) = fixture(&endpoint); + let current = Arc::new(AtomicBool::new(true)); + client.authority = EmaRequestAuthority::new(CancellationToken::new(), { + let current = Arc::clone(¤t); + move || current.load(Ordering::SeqCst) + }); *client.cached_token() = None; *source.credential_gate.lock().unwrap() = Some(gate); let authority = client.authority.clone(); let request = tokio::spawn(async move { client.http_request(params).await }); validated_rx.await?; - authority.revoke(); + // A watcher need not have run: dispatch must check the current owner itself. + current.store(false, Ordering::SeqCst); let _ = release_tx.send(()); tokio::time::timeout(Duration::from_secs(1), authority.quiesce()).await?; assert!(matches!( @@ -241,12 +247,13 @@ async fn revocation_ends_an_existing_response_stream() -> anyhow::Result<()> { .respond_with(ResponseTemplate::new(200).set_body_string("account-a-data")) .mount(&server) .await; - let (client, _source, params) = fixture(&format!("{}/mcp", server.uri())); + let (mut client, _source, params) = fixture(&format!("{}/mcp", server.uri())); + let cancellation = CancellationToken::new(); + client.authority = EmaRequestAuthority::new(cancellation.child_token(), || true); crate::oauth::test_support::warm_http_client(client.inner.as_ref()).await?; - let authority = client.authority.clone(); let (_response, mut body) = client.http_request_stream(params).await?; - authority.revoke(); + cancellation.cancel(); assert_eq!(body.recv().await?, None); Ok(())