mirror of
https://github.com/openai/codex.git
synced 2026-09-20 12:47:38 +00:00
0df39752cb
## What changed Advertise the experimental `codex/auth-change` capability for stdio MCP connections with an auth manager. When the server opts in, send `notifications/codex/authChanged` after initialization and on subsequent auth changes, with credential and owner generation counters and no credentials. Track owner changes separately from credential refreshes so consumers can detect login, logout, or user, workspace, and auth-mode changes even when notifications coalesce. Treat credential changes with incomplete owner identity as owner changes as well. Tie the notification watcher to the managed client's lifetime, limit each send to five seconds, and close the connection if a subsequent notification fails. ## Testing Add tests for credential refreshes versus owner changes, coalesced logout and account switches, capability opt-in, notification payloads, and watcher cleanup. GitOrigin-RevId: 35a99f5252bd48ca60a71fbdd1f708a8158e812c
63 lines
1.9 KiB
Rust
63 lines
1.9 KiB
Rust
//! Forwards auth invalidations without credentials. The managed client owns the watcher.
|
|
|
|
use std::sync::Arc;
|
|
use std::time::Duration;
|
|
|
|
use anyhow::Result;
|
|
use codex_login::AuthChangeState;
|
|
use codex_rmcp_client::RmcpClient;
|
|
use rmcp::model::ServerCapabilities;
|
|
use serde_json::json;
|
|
use tokio::sync::watch;
|
|
use tokio_util::task::AbortOnDropHandle;
|
|
|
|
pub(crate) const CAPABILITY: &str = "codex/auth-change";
|
|
const NOTIFICATION: &str = "notifications/codex/authChanged";
|
|
const SEND_TIMEOUT: Duration = Duration::from_secs(5);
|
|
|
|
pub(crate) async fn start(
|
|
client: Arc<RmcpClient>,
|
|
capabilities: &ServerCapabilities,
|
|
changes: Option<watch::Receiver<AuthChangeState>>,
|
|
) -> Result<Option<Arc<AbortOnDropHandle<()>>>> {
|
|
let Some(mut changes) = changes.filter(|_| {
|
|
capabilities
|
|
.experimental
|
|
.as_ref()
|
|
.is_some_and(|capabilities| capabilities.contains_key(CAPABILITY))
|
|
}) else {
|
|
return Ok(None);
|
|
};
|
|
|
|
notify(&client, &mut changes).await?;
|
|
let task = tokio::spawn(async move {
|
|
while changes.changed().await.is_ok() {
|
|
if notify(&client, &mut changes).await.is_err() {
|
|
tracing::warn!("MCP auth invalidation delivery failed; closing connection");
|
|
client.shutdown().await;
|
|
break;
|
|
}
|
|
}
|
|
});
|
|
Ok(Some(Arc::new(AbortOnDropHandle::new(task))))
|
|
}
|
|
|
|
async fn notify(client: &RmcpClient, changes: &mut watch::Receiver<AuthChangeState>) -> Result<()> {
|
|
let state = *changes.borrow_and_update();
|
|
tokio::time::timeout(
|
|
SEND_TIMEOUT,
|
|
client.send_custom_notification(
|
|
NOTIFICATION,
|
|
Some(json!({
|
|
"generation": state.generation,
|
|
"ownerGeneration": state.owner_generation,
|
|
})),
|
|
),
|
|
)
|
|
.await?
|
|
}
|
|
|
|
#[cfg(test)]
|
|
#[path = "auth_changes_tests.rs"]
|
|
mod tests;
|