Enable enterprise MCP sign-in and account-scoped grant cleanup (#49276)

## Why

The app server previously rejected OAuth login for enterprise-managed MCP servers. These servers need an explicit sign-in flow that ties saved authorization to the active Codex user and workspace and prevents canceled or superseded attempts from saving credentials.

## What changed

- Enable enterprise sign-in through `mcpServer/oauth/login` for configured servers with a connected thread, trusted IdP registration, and `features.use_xaa` enabled. Scope shared enterprise credentials by user, workspace, issuer, and client ID.
- Add `loginId` to MCP OAuth login responses and completion notifications, preserving compatibility with older servers that omit it.
- Support enterprise login replacement and `account/login/cancel`, wait for callback listener shutdown before replacement, and recheck account and configuration before saving grants. Fence stale attempts across processes and account changes.
- Remove shared enterprise authorization during account logout and `codex mcp logout`, including when the feature is disabled. Continue account logout if enterprise cleanup fails, and invalidate MCP runtimes on account changes.

Start a fresh session after successful enterprise sign-in to use the saved grant.

## Testing

Add regression coverage for login replacement and cancellation, callback port reuse, account changes during setup, cross-process logout invalidation, persisted workspace policy selection, cleanup failures, and responses from older servers without `loginId`.

GitOrigin-RevId: 7ed4caaab4659ed1e67064262cf5cee48e3a26ac
This commit is contained in:
Nick Steele
2026-09-29 14:46:33 +00:00
committed by copyberry
parent 677b07eca0
commit 4994306e9f
43 changed files with 1994 additions and 147 deletions
+2
View File
@@ -2144,6 +2144,7 @@ dependencies = [
"futures",
"hmac 0.12.1",
"http 1.4.0",
"keyring",
"opentelemetry",
"opentelemetry_sdk",
"pretty_assertions",
@@ -4395,6 +4396,7 @@ dependencies = [
"thiserror 2.0.18",
"tiny_http",
"tokio",
"tokio-util",
"tracing",
"tracing-subscriber",
"tracing-test",
@@ -2791,6 +2791,13 @@
"null"
]
},
"loginId": {
"description": "Identifies the explicit login attempt. Older servers omit this field.",
"type": [
"string",
"null"
]
},
"name": {
"type": "string"
},
@@ -14838,6 +14838,13 @@
"null"
]
},
"loginId": {
"description": "Identifies the explicit login attempt. Older servers omit this field.",
"type": [
"string",
"null"
]
},
"name": {
"type": "string"
},
@@ -14909,6 +14916,14 @@
"properties": {
"authorizationUrl": {
"type": "string"
},
"loginId": {
"default": null,
"description": "Identifies this login attempt across the response and completion notification. Older servers omit this field; current servers always return it.",
"type": [
"string",
"null"
]
}
},
"required": [
@@ -10870,6 +10870,13 @@
"null"
]
},
"loginId": {
"description": "Identifies the explicit login attempt. Older servers omit this field.",
"type": [
"string",
"null"
]
},
"name": {
"type": "string"
},
@@ -10941,6 +10948,14 @@
"properties": {
"authorizationUrl": {
"type": "string"
},
"loginId": {
"default": null,
"description": "Identifies this login attempt across the response and completion notification. Older servers omit this field; current servers always return it.",
"type": [
"string",
"null"
]
}
},
"required": [
@@ -7,6 +7,13 @@
"null"
]
},
"loginId": {
"description": "Identifies the explicit login attempt. Older servers omit this field.",
"type": [
"string",
"null"
]
},
"name": {
"type": "string"
},
@@ -3,6 +3,14 @@
"properties": {
"authorizationUrl": {
"type": "string"
},
"loginId": {
"default": null,
"description": "Identifies this login attempt across the response and completion notification. Older servers omit this field; current servers always return it.",
"type": [
"string",
"null"
]
}
},
"required": [
@@ -2,4 +2,8 @@
// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually.
export type McpServerOauthLoginCompletedNotification = { name: string, threadId: string | null, success: boolean, error?: string, };
export type McpServerOauthLoginCompletedNotification = { name: string, threadId: string | null,
/**
* Identifies the explicit login attempt. Older servers omit this field.
*/
loginId?: string | null, success: boolean, error?: string, };
@@ -2,4 +2,9 @@
// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually.
export type McpServerOauthLoginResponse = { authorizationUrl: string, };
export type McpServerOauthLoginResponse = { authorizationUrl: string,
/**
* Identifies this login attempt across the response and completion notification.
* Older servers omit this field; current servers always return it.
*/
loginId?: string, };
+13 -5
View File
@@ -2193,6 +2193,12 @@ mod tests {
.ok_or_else(|| anyhow::anyhow!("missing account usage response fixture"))?,
)?;
assert!(account_usage_response_ts.contains("threadUsage?: ThreadUsage | null"));
let mcp_login_completion_ts = std::str::from_utf8(
fixture_tree
.get(Path::new("v2/McpServerOauthLoginCompletedNotification.ts"))
.ok_or_else(|| anyhow::anyhow!("missing MCP login completion fixture"))?,
)?;
assert!(mcp_login_completion_ts.contains("loginId?: string | null"));
let server_request_ts = std::str::from_utf8(
fixture_tree
.get(Path::new("ServerRequest.ts"))
@@ -2386,14 +2392,16 @@ mod tests {
// If the last non-whitespace before ':' is '?', then this is an
// optional field with a nullable type (i.e., "?: T | null").
// These are only allowed in *Params types, except the additive stable usage
// response field, which older servers omit and newer servers return as null.
let legacy_account_usage_response = path
// These are only allowed in *Params types, except additive stable fields
// that older servers omit and newer servers may return as null.
let legacy_optional_nullable_field = (path
== Path::new("v2/GetAccountTokenUsageResponse.ts")
&& field_prefix.trim() == "threadUsage?";
&& field_prefix.trim() == "threadUsage?")
|| (path == Path::new("v2/McpServerOauthLoginCompletedNotification.ts")
&& field_prefix.trim() == "loginId?");
if field_prefix.chars().rev().find(|c| !c.is_whitespace()) == Some('?')
&& !allow_optional_nullable
&& !legacy_account_usage_response
&& !legacy_optional_nullable_field
{
let line_number =
contents[..abs_idx].chars().filter(|c| *c == '\n').count() + 1;
@@ -321,6 +321,11 @@ pub enum McpServerOauthClientRegistration {
#[ts(export_to = "v2/")]
pub struct McpServerOauthLoginResponse {
pub authorization_url: String,
/// Identifies this login attempt across the response and completion notification.
/// Older servers omit this field; current servers always return it.
#[serde(default)]
#[ts(optional)]
pub login_id: Option<String>,
}
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)]
#[serde(rename_all = "camelCase")]
@@ -338,6 +343,9 @@ pub struct McpToolCallProgressNotification {
pub struct McpServerOauthLoginCompletedNotification {
pub name: String,
pub thread_id: Option<String>,
/// Identifies the explicit login attempt. Older servers omit this field.
#[ts(optional, as = "Option<Option<String>>")]
pub login_id: Option<String>,
pub success: bool,
#[serde(default, skip_serializing_if = "Option::is_none")]
#[ts(optional)]
@@ -99,6 +99,22 @@ fn managed_hooks_requirements_default_interrupt_to_empty() {
assert_eq!(parsed.interrupt, Vec::new());
}
#[test]
fn mcp_oauth_login_response_accepts_older_servers_without_login_id() {
let response = serde_json::from_value::<McpServerOauthLoginResponse>(json!({
"authorizationUrl": "https://example.com/authorize",
}))
.expect("older login response should deserialize");
assert_eq!(
response,
McpServerOauthLoginResponse {
authorization_url: "https://example.com/authorize".to_string(),
login_id: None,
}
);
}
#[test]
fn external_agent_config_detect_response_defaults_connectors_for_older_servers() {
let response = serde_json::from_value::<ExternalAgentConfigDetectResponse>(json!({
+1
View File
@@ -140,6 +140,7 @@ core_test_support = { workspace = true }
flate2 = { workspace = true }
hmac = { workspace = true }
http = { workspace = true }
keyring = { workspace = true }
opentelemetry = { workspace = true }
opentelemetry_sdk = { workspace = true }
pretty_assertions = { workspace = true }
+9
View File
@@ -352,6 +352,15 @@ layers with enterprise MCP disabled. Other planned thread refreshes are processe
before the rejection is reported, so an error does not imply that no changes were
applied. Correct the policy before retrying the reload.
# Enterprise sign-in
Call `mcpServer/oauth/login` with a directly configured server's `name` and its
connected `threadId`. Open the returned `authorizationUrl` and match
`mcpServer/oauthLogin/completed` by `loginId`. Starting another enterprise sign-in
cancels the previous attempt and waits for its callback listener to close. Use
`account/login/cancel` to cancel explicitly. Start a fresh session after success to
use the saved grant.
# Selected workspace routing
The experimental `account/read.workspaceRouting` response field returns the selected ChatGPT workspace's `chatgptAccountId`, resolved HTTPS `backendOrigin`, and backend-provided `accountRoutingOverride`. The routing value is `us`, `us_cr`, or the explicit `NO_CONSTRAINT` value. API-only and signed-out accounts return `null` and do not need `accounts/check`.
+3 -1
View File
@@ -473,6 +473,7 @@ impl MessageProcessor {
thread_state_manager.clone(),
outgoing.clone(),
config_manager.clone(),
account_processor.clone(),
);
let plugin_processor = PluginRequestProcessor::new(
auth_manager.clone(),
@@ -1680,7 +1681,8 @@ impl MessageProcessor {
self.turn_processor.review_start(&request_id, params).await
}
ClientRequest::McpServerOauthLogin { params, .. } => {
self.mcp_processor.mcp_server_oauth_login(params).await
// Keep MCP discovery's large future from inflating every request.
Box::pin(self.mcp_processor.mcp_server_oauth_login(params)).await
}
ClientRequest::McpServerRefresh { params, .. } => {
self.mcp_processor.mcp_server_refresh(params).await
@@ -11,13 +11,19 @@ use codex_app_server_protocol::DesktopOnboardingEntrypoint;
use codex_app_server_protocol::GetAccountRateLimitsParams;
use codex_login::LoginOnboardingEntrypoint;
use codex_login::login_with_bedrock_access_keys;
use codex_mcp::ema_auth_scope;
use codex_model_provider::is_supported_amazon_bedrock_region;
use codex_rmcp_client::EnterpriseOAuthCredentialGuard;
mod bedrock_setup;
mod enterprise_login;
mod gateway_oauth;
mod rate_limit_resets;
mod workspace_routing;
pub(super) use enterprise_login::EnterpriseLoginCompletion;
pub(super) use enterprise_login::EnterpriseLoginTarget;
// Duration before a browser ChatGPT login attempt is abandoned.
const LOGIN_CHATGPT_TIMEOUT: Duration = Duration::from_secs(10 * 60);
const ACCOUNT_TOKEN_USAGE_FETCH_TIMEOUT: Duration = Duration::from_secs(/*secs*/ 10);
@@ -100,6 +106,7 @@ pub(crate) struct AccountRequestProcessor {
gateway_login: Arc<std::sync::Mutex<Option<gateway_oauth::ActiveGatewayLogin>>>,
gateway_client: Arc<std::sync::Mutex<Option<Arc<codex_login::GatewayAuthManager>>>>,
_gateway_notifications: Arc<tokio_util::task::AbortOnDropHandle<()>>,
pub(super) enterprise_login: Arc<enterprise_login::EnterpriseLoginState>,
}
impl AccountRequestProcessor {
@@ -115,6 +122,11 @@ impl AccountRequestProcessor {
config_manager.clone(),
Arc::clone(&outgoing),
);
let enterprise_login = Arc::new(enterprise_login::EnterpriseLoginState::new(
Arc::clone(&auth_manager),
Arc::clone(&thread_manager),
config_manager.clone(),
));
let processor = Arc::new(Self {
_gateway_notifications: Arc::new(gateway_notifications),
auth_manager,
@@ -128,6 +140,7 @@ impl AccountRequestProcessor {
workspace_routing: Arc::new(Mutex::new(None)),
workspace_routing_fetches: Arc::new(Mutex::new(HashMap::new())),
workspace_routing_shutdown: CancellationToken::new(),
enterprise_login,
});
let resolver: Arc<dyn codex_login::WorkspaceRoutingResolver> = processor.clone();
processor
@@ -210,6 +223,7 @@ impl AccountRequestProcessor {
pub(crate) async fn cancel_active_login(&self) {
self.cancel_gateway_login();
self.enterprise_login.cancel(/*login_id*/ None).await;
let mut guard = self.active_login.lock().await;
if let Some(active_login) = guard.take() {
drop(active_login);
@@ -318,6 +332,7 @@ impl AccountRequestProcessor {
if self.auth_manager.is_workload_identity_selected() {
return Err(self.configured_auth_owned_by_host_error());
}
self.enterprise_login.cancel(/*login_id*/ None).await;
match params {
LoginAccountParams::ApiKey { api_key } => {
self.login_api_key_v2(request_id, LoginApiKeyParams { api_key })
@@ -801,6 +816,11 @@ impl AccountRequestProcessor {
let login_id = params.login_id;
let uuid = Uuid::parse_str(&login_id)
.map_err(|_| invalid_request(format!("invalid login id: {login_id}")))?;
if self.enterprise_login.cancel(Some(&login_id)).await {
return Ok(CancelLoginAccountResponse {
status: CancelLoginAccountStatus::Canceled,
});
}
let status = match self.cancel_login_chatgpt_common(uuid).await {
Ok(()) => CancelLoginAccountStatus::Canceled,
Err(CancelLoginError::NotFound) => CancelLoginAccountStatus::NotFound,
@@ -884,6 +904,7 @@ impl AccountRequestProcessor {
}
async fn send_login_success_notifications(&self, login_id: Option<Uuid>) {
self.thread_manager.invalidate_mcp_runtimes().await;
self.send_account_login_notifications(AccountLoginCompletedNotification {
login_id: login_id.map(|id| id.to_string()),
success: true,
@@ -945,6 +966,7 @@ impl AccountRequestProcessor {
self.auth_manager.reload().await;
let auth_changes = self.auth_manager.auth_change_state_receiver();
let owner_generation = auth_changes.borrow().owner_generation;
self.thread_manager.invalidate_mcp_runtimes().await;
self.config_manager.replace_cloud_config_bundle_loader(
self.auth_manager.clone(),
self.config.chatgpt_base_url.clone(),
@@ -965,16 +987,65 @@ impl AccountRequestProcessor {
if self.auth_manager.is_workload_identity_selected() {
return Err(self.configured_auth_owned_by_host_error());
}
let config = self.load_latest_config().await;
// Another process may have changed the persisted workspace. Reload both
// account authority and its policy before selecting a grant to remove.
self.auth_manager.reload().await;
self.config_manager.replace_cloud_config_bundle_loader(
Arc::clone(&self.auth_manager),
self.config.chatgpt_base_url.clone(),
self.config.http_client_factory(),
);
let config = self
.config_manager
.load_latest_config(/*fallback_cwd*/ None)
.await;
let scope = ema_auth_scope(self.auth_manager.auth_cached().as_ref());
let enterprise_policy_failed = scope.is_some() && config.is_err();
// Startup policy may belong to another workspace. Never use that fallback
// to select a credential; a policy failure must not block primary logout.
let enterprise_grant = config.as_ref().ok().and_then(|config| {
scope
.as_ref()
.zip(config.mcp_enterprise_managed_auth.as_ref())
.map(|(scope, profile)| {
(
profile.idp.credential_name(scope),
profile.idp.issuer.clone(),
config.auth_keyring_backend_kind(),
)
})
});
let config = config.unwrap_or_else(|_| self.config.as_ref().clone());
self.cancel_active_login().await;
// Retain the credential lock through primary logout. Otherwise a second
// process can commit after deletion but before the account is removed.
let enterprise_guard =
if let Some((credential_name, issuer, keyring_backend)) = enterprise_grant {
EnterpriseOAuthCredentialGuard::acquire(&credential_name, &issuer, keyring_backend)
.await
.map(Some)
} else {
Ok(None)
};
let cleanup_failed = match &enterprise_guard {
Ok(Some(guard)) => guard.delete_tokens().is_err(),
Ok(None) => false,
Err(_) => true,
};
if enterprise_policy_failed || cleanup_failed {
tracing::warn!("Failed to remove enterprise authorization; continuing account logout");
}
match self.auth_manager.logout_with_revoke().await {
Ok(_) => {}
Err(err) => {
return Err(internal_error(format!("logout failed: {err}")));
}
}
drop(enterprise_guard);
self.thread_manager.invalidate_mcp_runtimes().await;
self.config_manager.clear_cloud_config_bundle_loader();
@@ -0,0 +1,274 @@
//! Owns one explicit enterprise sign-in attempt for a configured MCP server.
//! Start, replacement, cancellation, and local commit admission all pass through this state.
//! The credential store separately fences logins invalidated by another process.
use super::*;
use codex_config::McpEmaRegistration;
use codex_config::types::AuthKeyringBackendKind;
use codex_mcp::ema_auth_scope;
use codex_rmcp_client::EnterpriseOAuthLoginRequest;
use codex_rmcp_client::perform_enterprise_oauth_login_return_url;
use futures::FutureExt;
use tokio::sync::watch;
pub(crate) struct EnterpriseLoginTarget {
pub(crate) thread_id: String,
pub(crate) server_name: String,
pub(crate) registration: McpEmaRegistration,
}
#[derive(Clone, Copy, Eq, PartialEq)]
pub(crate) enum EnterpriseLoginCompletion {
Pending,
Succeeded,
Failed,
}
impl EnterpriseLoginCompletion {
pub(crate) fn is_complete(self) -> bool {
matches!(self, Self::Succeeded | Self::Failed)
}
}
pub(crate) struct EnterpriseLogin {
pub(crate) login_id: String,
pub(crate) authorization_url: String,
pub(crate) completion: watch::Receiver<EnterpriseLoginCompletion>,
}
pub(crate) struct EnterpriseLoginState {
auth_manager: Arc<AuthManager>,
thread_manager: Arc<ThreadManager>,
config_manager: ConfigManager,
// Cancellation and commit admission use this same lock. A cancelled attempt cannot commit.
active: Mutex<Option<EnterpriseLoginAttempt>>,
// Start and cancel hold this through callback shutdown; the worker never acquires it.
start_gate: Mutex<()>,
// Account changes cancel queued and running setup before waiting for start_gate.
setup_cancel: Mutex<CancellationToken>,
}
struct EnterpriseLoginAttempt {
login_id: String,
cancel: CancellationToken,
completion: watch::Receiver<EnterpriseLoginCompletion>,
}
impl Drop for EnterpriseLoginAttempt {
fn drop(&mut self) {
self.cancel.cancel();
}
}
impl EnterpriseLoginState {
pub(super) fn new(
auth_manager: Arc<AuthManager>,
thread_manager: Arc<ThreadManager>,
config_manager: ConfigManager,
) -> Self {
Self {
auth_manager,
thread_manager,
config_manager,
active: Mutex::new(None),
start_gate: Mutex::new(()),
setup_cancel: Mutex::new(CancellationToken::new()),
}
}
#[expect(
clippy::await_holding_invalid_type,
reason = "serialize start/cancel through callback shutdown; the worker uses a separate state lock"
)]
pub(super) async fn cancel(&self, login_id: Option<&str>) -> bool {
if login_id.is_none() {
let mut setup_cancel = self.setup_cancel.lock().await;
setup_cancel.cancel();
*setup_cancel = CancellationToken::new();
}
let (login_id, mut completion) = {
let active = self.active.lock().await;
if login_id.is_some_and(|id| active.as_ref().is_none_or(|a| a.login_id != id)) {
return false;
}
let Some(attempt) = active.as_ref() else {
return false;
};
attempt.cancel.cancel();
(attempt.login_id.clone(), attempt.completion.clone())
};
// A partial callback request can stall the blocking HTTP worker. Cancellation
// already fences commit; primary logout must not wait indefinitely for I/O.
if tokio::time::timeout(std::time::Duration::from_secs(10), async {
let _start = self.start_gate.lock().await;
completion.wait_for(|state| state.is_complete()).await
})
.await
.is_ok()
{
self.clear(&login_id).await;
}
// Retain a timed-out attempt so a replacement still waits for port cleanup.
true
}
async fn clear(&self, login_id: &str) {
let mut active = self.active.lock().await;
if active
.as_ref()
.is_some_and(|attempt| attempt.login_id == login_id)
{
*active = None;
}
}
async fn configuration_is_current(
&self,
target: &EnterpriseLoginTarget,
keyring_backend: AuthKeyringBackendKind,
) -> bool {
let Ok(thread_id) = ThreadId::from_string(&target.thread_id) else {
return false;
};
let Ok(thread) = self.thread_manager.get_thread(thread_id).await else {
return false;
};
let current = thread.config().await;
let Ok(config) = self
.config_manager
.load_latest_config_with_session_layers(&current.config_layer_stack, &current.cwd)
.await
else {
return false;
};
let config = thread.runtime_mcp_config_and_context(&config).await.0;
let auth = self.auth_manager.auth_cached();
let servers = codex_mcp::effective_mcp_servers(&config, auth.as_ref());
config.xaa_enabled
&& config.auth_keyring_backend_kind == keyring_backend
&& servers.get(&target.server_name).is_some_and(|server| {
server.config().enabled
&& server.config().ema_registration() == Some(&target.registration)
})
}
#[expect(
clippy::await_holding_invalid_type,
reason = "serialize start/cancel through callback shutdown; the worker uses a separate state lock"
)]
pub(crate) async fn start(
self: &Arc<Self>,
request: EnterpriseOAuthLoginRequest<'_>,
scope: codex_config::McpEmaAuthScope,
target: EnterpriseLoginTarget,
) -> Result<EnterpriseLogin, JSONRPCErrorError> {
// MCP requests are not in the account request queue. Serialize start and
// cancellation here; the separate commit gate stays responsive to the worker.
let setup_cancel = self.setup_cancel.lock().await.child_token();
let _start = tokio::select! {
biased;
_ = setup_cancel.cancelled() => return Err(internal_error("enterprise sign-in was cancelled")),
guard = self.start_gate.lock() => guard,
};
let previous = {
let active = self.active.lock().await;
// Keep the cancelled attempt installed until its callback worker closes.
// A cancelled replacement request leaves the next start responsible for waiting.
active.as_ref().map(|attempt| {
attempt.cancel.cancel();
(attempt.login_id.clone(), attempt.completion.clone())
})
};
if let Some((login_id, mut completion)) = previous {
// An abandoned setup request drops its flow and closes this channel.
let _ = tokio::select! {
biased;
_ = setup_cancel.cancelled() => return Err(internal_error("enterprise sign-in was cancelled")),
result = tokio::time::timeout(
std::time::Duration::from_secs(10),
completion.wait_for(|state| state.is_complete()),
) => result.map_err(|_| internal_error("previous enterprise sign-in did not close"))?,
};
self.clear(&login_id).await;
}
let keyring_backend = request.keyring_backend_kind;
// Keep commit cancellation under `active`; cancelling the setup epoch must
// not race an already-admitted synchronous credential write.
let cancel = CancellationToken::new();
let login_id = Uuid::new_v4().to_string();
let (completed, completion) = watch::channel(EnterpriseLoginCompletion::Pending);
*self.active.lock().await = Some(EnterpriseLoginAttempt {
login_id: login_id.clone(),
cancel: cancel.clone(),
completion: completion.clone(),
});
let handle = tokio::select! {
biased;
_ = setup_cancel.cancelled() => Err(internal_error("enterprise sign-in was cancelled")),
_ = cancel.cancelled() => Err(internal_error("enterprise sign-in was cancelled")),
result = perform_enterprise_oauth_login_return_url(request) =>
result.map_err(|_| internal_error("failed to start enterprise sign-in")),
};
let handle = match handle {
Ok(handle) => handle,
Err(error) => {
self.clear(&login_id).await;
completed.send_replace(EnterpriseLoginCompletion::Failed);
return Err(error);
}
};
let callback_closed = handle.callback_closed().shared();
let authorization_url = handle.authorization_url();
let lifecycle = Arc::clone(self);
let completion_id = login_id.clone();
tokio::spawn(async move {
let result = tokio::select! {
biased;
_ = cancel.cancelled() => Err(anyhow::anyhow!("Enterprise sign-in was cancelled.")),
result = async {
let credentials = handle.wait().await?;
callback_closed.clone().await;
let mut active = credentials.commit_if(|| async {
lifecycle.auth_manager.reload().await;
let eligible = lifecycle.configuration_is_current(&target, keyring_backend).await;
// Reread persisted account authority after acquiring the credential
// lock, including changes made by another app-server process.
lifecycle.auth_manager.reload().await;
if !eligible {
return None;
}
// Keep cancellation responsive during the async checks above.
// Once admitted, retain this guard through the synchronous write.
let active = lifecycle.active.lock().await;
(active.as_ref().is_some_and(|attempt| {
attempt.login_id == completion_id && !attempt.cancel.is_cancelled()
})
&& ema_auth_scope(lifecycle.auth_manager.auth_cached().as_ref()).as_ref() == Some(&scope))
.then_some(active)
}).await?;
// A successful commit is no longer cancelable. Clear it while
// still holding the commit gate.
*active = None;
Ok(())
} => result,
};
// The canceled wait future has dropped its listener guard. Also wait
// for the blocking callback worker before allowing a fixed-port retry.
callback_closed.await;
if result.is_err() {
lifecycle.clear(&completion_id).await;
}
let success = result.is_ok();
completed.send_replace(if success {
EnterpriseLoginCompletion::Succeeded
} else {
EnterpriseLoginCompletion::Failed
});
});
Ok(EnterpriseLogin {
login_id,
authorization_url,
completion,
})
}
}
@@ -1,8 +1,13 @@
use super::account_processor::EnterpriseLoginCompletion;
use super::account_processor::EnterpriseLoginTarget;
use super::thread_input::ensure_direct_input_allowed;
use super::*;
use codex_config::types::McpServerAuth;
use codex_mcp::McpServerSource;
use codex_mcp::ReadResourceRequestParams;
use codex_mcp::ema_auth_scope;
use codex_mcp::resolve_oauth_callback;
use codex_rmcp_client::EnterpriseOAuthLoginRequest;
use std::future::Future;
use crate::thread_state::ThreadStateManager;
@@ -15,6 +20,7 @@ pub(crate) struct McpRequestProcessor {
thread_manager: Arc<ThreadManager>,
pub(super) outgoing: Arc<OutgoingMessageSender>,
config_manager: ConfigManager,
account_processor: Arc<AccountRequestProcessor>,
pub(super) thread_state_manager: ThreadStateManager,
}
@@ -25,12 +31,14 @@ impl McpRequestProcessor {
thread_state_manager: ThreadStateManager,
outgoing: Arc<OutgoingMessageSender>,
config_manager: ConfigManager,
account_processor: Arc<AccountRequestProcessor>,
) -> Self {
Self {
auth_manager,
thread_manager,
outgoing,
config_manager,
account_processor,
thread_state_manager,
}
}
@@ -134,12 +142,6 @@ impl McpRequestProcessor {
scopes,
timeout_secs,
} = params;
let client_registration = match client_registration.unwrap_or_default() {
McpServerOauthClientRegistration::Auto => McpOAuthClientRegistration::Auto,
McpServerOauthClientRegistration::Cimd => McpOAuthClientRegistration::Cimd,
McpServerOauthClientRegistration::Dcr => McpOAuthClientRegistration::Dcr,
};
let auth = self.auth_manager.auth().await;
let (mcp_config, runtime_context) = match thread_id.as_deref() {
Some(thread_id) => {
@@ -174,84 +176,175 @@ impl McpRequestProcessor {
StreamableHttpRedirectMode::Legacy
};
let server = server.config();
if matches!(server.auth, codex_config::McpServerAuth::EmaAuth) {
return Err(invalid_request(
"EMA MCP connections are not enabled in this version",
));
}
let (url, http_headers, env_http_headers) = match &server.transport {
McpServerTransportConfig::StreamableHttp {
url,
http_headers,
env_http_headers,
..
} => (url.clone(), http_headers.clone(), env_http_headers.clone()),
_ => {
let enterprise = matches!(server.auth, McpServerAuth::EmaAuth);
let (authorization_url, login_id, completion): (
String,
String,
futures::future::BoxFuture<'static, anyhow::Result<()>>,
) = if enterprise {
let thread_id = thread_id.clone().ok_or_else(|| {
invalid_request("enterprise MCP login requires a connected thread")
})?;
if !mcp_config.xaa_enabled {
return Err(invalid_request(
"OAuth login is only supported for streamable HTTP servers.",
"enterprise MCP authentication requires `[features].use_xaa = true`",
));
}
};
let http_client = runtime_context
.resolve_http_client(&name, server)
.map_err(|err| {
internal_error(format!("failed to resolve MCP server runtime: {err}"))
if !server.enabled {
return Err(invalid_request(format!("MCP server '{name}' is disabled.")));
}
server
.validate_ema_auth_transport()
.map_err(|reason| invalid_request(format!("MCP server '{name}' {reason}")))?;
if client_registration.is_some() {
return Err(invalid_request(
"enterprise MCP login uses the configured enterprise client and does not support clientRegistration overrides",
));
}
// Resource scopes and per-server callbacks do not alter the shared IdP registration.
let scope = ema_auth_scope(auth.as_ref()).ok_or_else(|| {
invalid_request(
"enterprise MCP authentication requires a signed-in Codex user and workspace",
)
})?;
let idp = server.oauth_idp().ok_or_else(|| {
invalid_request(format!(
"MCP server '{name}' has no configured enterprise IdP"
))
})?;
let registration = server.ema_registration().cloned().ok_or_else(|| {
invalid_request(format!(
"MCP server '{name}' has no validated enterprise registration"
))
})?;
let credential_name = idp.credential_name(&scope);
let http_client =
runtime_context
.resolve_http_client(&name, server)
.map_err(|err| {
internal_error(format!("failed to resolve MCP server runtime: {err}"))
})?;
let login = self
.account_processor
.enterprise_login
.start(
EnterpriseOAuthLoginRequest {
credential_name: &credential_name,
issuer: &idp.issuer,
client_id: &idp.client_id,
keyring_backend_kind: mcp_config.auth_keyring_backend_kind,
callback_port: mcp_config.mcp_oauth_callback_port,
callback_url: mcp_config.mcp_oauth_callback_url.as_deref(),
timeout_secs,
http_client,
redirect_mode,
},
scope,
EnterpriseLoginTarget {
thread_id,
server_name: name.clone(),
registration,
},
)
.await?;
let mut completion = login.completion;
(
login.authorization_url,
login.login_id,
Box::pin(async move {
match completion.wait_for(|state| state.is_complete()).await {
Ok(result) if *result == EnterpriseLoginCompletion::Succeeded => Ok(()),
_ => Err(anyhow::anyhow!("Enterprise sign-in failed.")),
}
}),
)
} else {
let (url, http_headers, env_http_headers) = match &server.transport {
McpServerTransportConfig::StreamableHttp {
url,
http_headers,
env_http_headers,
..
} => (url.clone(), http_headers.clone(), env_http_headers.clone()),
_ => {
return Err(invalid_request(
"OAuth login is only supported for streamable HTTP servers.",
));
}
};
let discovered_scopes = if scopes.is_none() && server.scopes.is_none() {
discover_supported_scopes(
&server.transport,
Arc::clone(&http_client),
codex_rmcp_client::OAuthDiscoveryTimeout::Requested,
let http_client =
runtime_context
.resolve_http_client(&name, server)
.map_err(|err| {
internal_error(format!("failed to resolve MCP server runtime: {err}"))
})?;
let discovered_scopes = if scopes.is_none() && server.scopes.is_none() {
discover_supported_scopes(
&server.transport,
Arc::clone(&http_client),
codex_rmcp_client::OAuthDiscoveryTimeout::Requested,
redirect_mode,
)
.await
} else {
None
};
let resolved_scopes =
resolve_oauth_scopes(scopes, server.scopes.clone(), discovered_scopes);
let oauth_credential_name = server.oauth_credential_name(&name);
let client_registration = match client_registration.unwrap_or_default() {
McpServerOauthClientRegistration::Auto => McpOAuthClientRegistration::Auto,
McpServerOauthClientRegistration::Cimd => McpOAuthClientRegistration::Cimd,
McpServerOauthClientRegistration::Dcr => McpOAuthClientRegistration::Dcr,
};
let callback_url =
resolve_oauth_callback(server, &url, mcp_config.mcp_oauth_callback_url.as_deref())
.map_err(|err| {
internal_error(format!("failed to resolve MCP OAuth callback: {err}"))
})?;
let handle = perform_oauth_login_return_url(
oauth_credential_name.as_ref(),
&url,
mcp_config.mcp_oauth_credentials_store_mode,
mcp_config.auth_keyring_backend_kind,
http_headers,
env_http_headers,
&resolved_scopes.scopes,
server.oauth.as_ref(),
client_registration,
server.oauth_resource.as_deref(),
timeout_secs,
server.oauth_callback_port(mcp_config.mcp_oauth_callback_port),
callback_url.as_deref(),
mcp_config.mcp_oauth_callback_url.as_deref(),
http_client,
redirect_mode,
)
.await
} else {
None
.map_err(|err| {
internal_error(format!("failed to login to MCP server '{name}': {err}"))
})?;
(
handle.authorization_url().to_string(),
Uuid::now_v7().to_string(),
Box::pin(handle.wait()),
)
};
let resolved_scopes =
resolve_oauth_scopes(scopes, server.scopes.clone(), discovered_scopes);
let oauth_credential_name = server.oauth_credential_name(&name);
let callback_url =
resolve_oauth_callback(server, &url, mcp_config.mcp_oauth_callback_url.as_deref())
.map_err(|err| {
internal_error(format!("failed to resolve MCP OAuth callback: {err}"))
})?;
let handle = perform_oauth_login_return_url(
oauth_credential_name.as_ref(),
&url,
mcp_config.mcp_oauth_credentials_store_mode,
mcp_config.auth_keyring_backend_kind,
http_headers,
env_http_headers,
&resolved_scopes.scopes,
server.oauth.as_ref(),
client_registration,
server.oauth_resource.as_deref(),
timeout_secs,
server.oauth_callback_port(mcp_config.mcp_oauth_callback_port),
callback_url.as_deref(),
mcp_config.mcp_oauth_callback_url.as_deref(),
http_client,
redirect_mode,
)
.await
.map_err(|err| internal_error(format!("failed to login to MCP server '{name}': {err}")))?;
let authorization_url = handle.authorization_url().to_string();
let notification_name = name.clone();
let notification_thread_id = thread_id;
let notification_login_id = login_id.clone();
let outgoing = Arc::clone(&self.outgoing);
let thread_manager = Arc::clone(&self.thread_manager);
tokio::spawn(async move {
let (success, error) = match handle.wait().await {
let (success, error) = match completion.await {
Ok(()) => (true, None),
Err(err) => (false, Some(err.to_string())),
};
if success {
if success && !enterprise {
thread_manager.invalidate_mcp_runtimes().await;
}
@@ -259,6 +352,7 @@ impl McpRequestProcessor {
McpServerOauthLoginCompletedNotification {
name: notification_name,
thread_id: notification_thread_id,
login_id: Some(notification_login_id),
success,
error,
},
@@ -266,7 +360,10 @@ impl McpRequestProcessor {
outgoing.send_server_notification(notification).await;
});
Ok(McpServerOauthLoginResponse { authorization_url })
Ok(McpServerOauthLoginResponse {
authorization_url,
login_id: Some(login_id),
})
}
async fn list_mcp_server_status(
@@ -1954,6 +1954,7 @@ impl PluginRequestProcessor {
McpServerOauthLoginCompletedNotification {
name: notification_name,
thread_id: None,
login_id: None,
success,
error,
},
+50 -8
View File
@@ -92,6 +92,9 @@ fn expected_workspace_routing(
})
}
#[path = "account_enterprise_tests.rs"]
mod enterprise_tests;
// Helper to create a minimal config.toml for the app server
#[derive(Default)]
struct CreateConfigTomlParams {
@@ -281,23 +284,62 @@ async fn mock_oauth_token(server: &MockServer, id_token: &str) {
.await;
}
#[test_case(false; "api_key")]
#[test_case(true; "enterprise_cleanup_failure")]
#[tokio::test]
async fn logout_account_removes_auth_and_notifies() -> Result<()> {
async fn logout_account_removes_auth_and_notifies(enterprise_cleanup_failure: bool) -> Result<()> {
let oauth_server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/oauth/revoke"))
.respond_with(ResponseTemplate::new(200))
.expect(u64::from(enterprise_cleanup_failure))
.mount(&oauth_server)
.await;
let codex_home = TempDir::new()?;
create_config_toml(codex_home.path(), CreateConfigTomlParams::default())?;
login_with_api_key(
codex_home.path(),
"sk-test-key",
AuthCredentialsStoreMode::File,
AuthKeyringBackendKind::default(),
)?;
if enterprise_cleanup_failure {
write_chatgpt_auth(
codex_home.path(),
ChatGptAuthFixture::new("access-token")
.account_id(WORKSPACE_ID_INITIAL)
.chatgpt_user_id("enterprise-user"),
AuthCredentialsStoreMode::File,
)?;
let mut config = read_config_toml(codex_home.path())?;
config.as_table_mut().expect("config table").insert(
"mcp_enterprise_managed_auth".into(),
toml::toml! {
[idp]
issuer = "https://idp.example"
client_id = "enterprise-client"
}
.into(),
);
std::fs::write(
codex_home.path().join("config.toml"),
toml::to_string(&config)?,
)?;
// Fail credential cleanup before touching the platform keyring.
std::fs::write(codex_home.path().join("mcp-oauth-locks"), "not a directory")?;
} else {
login_with_api_key(
codex_home.path(),
"sk-test-key",
AuthCredentialsStoreMode::File,
AuthKeyringBackendKind::default(),
)?;
}
assert!(codex_home.path().join("auth.json").exists());
let refresh_url = format!("{}/oauth/token", oauth_server.uri());
let mut mcp = TestAppServer::builder()
.with_codex_home(codex_home.path())
.without_auto_env()
.with_env_overrides(&[("OPENAI_API_KEY", None)])
.with_env_overrides(&[
("OPENAI_API_KEY", None),
(REFRESH_TOKEN_URL_OVERRIDE_ENV_VAR, Some(&refresh_url)),
])
.build_initialized_with_timeout(DEFAULT_READ_TIMEOUT)
.await?;
@@ -0,0 +1,814 @@
//! Account API regressions for enterprise policy, persisted authority, and login attempts.
use super::*;
use std::any::Any;
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::Mutex;
use anyhow::Context;
use app_test_support::MockResponsesConfig;
use base64::Engine;
use base64::engine::general_purpose::URL_SAFE_NO_PAD;
use codex_app_server::in_process;
use codex_app_server::in_process::InProcessClientHandle;
use codex_app_server::in_process::InProcessServerEvent;
use codex_app_server::in_process::InProcessStartArgs;
use codex_app_server_protocol::InitializeParams;
use codex_arg0::Arg0DispatchPaths;
use codex_config::CloudConfigBundleLoader;
use codex_config::LoaderOverrides;
use codex_config::McpEmaAuthScope;
use codex_config::McpServerIdpOAuthConfig;
use codex_config::NoopThreadConfigLoader;
use codex_config::test_support::CloudConfigBundleFixture;
use codex_config::types::OAuthCredentialsStoreMode;
use codex_core::config::ConfigBuilder;
use codex_exec_server::EnvironmentManager;
use codex_features::Feature;
use codex_feedback::CodexFeedback;
use codex_login::CODEX_ACCESS_TOKEN_ENV_VAR;
use codex_protocol::protocol::SessionSource;
use codex_rmcp_client::stored_oauth_credentials;
use core_test_support::skip_if_no_network;
use keyring::credential::Credential;
use keyring::credential::CredentialApi;
use keyring::credential::CredentialBuilderApi;
use keyring::credential::CredentialPersistence;
use keyring::mock::MockCredential;
use pretty_assertions::assert_eq;
use serde_json::Value;
use test_case::test_case;
use tokio::io::AsyncReadExt;
use tokio::io::AsyncWriteExt;
use tokio::process::Command;
use wiremock::matchers::body_partial_json;
use wiremock::matchers::header;
const WORKSPACE_ID_REFRESHED: &str = "123e4567-e89b-42d3-a456-426614174012";
#[test_case(false; "current_policy")]
#[test_case(true; "policy_load_failure")]
#[tokio::test]
async fn logout_reloads_persisted_workspace_and_its_enterprise_policy(
policy_load_failure: bool,
) -> Result<()> {
let server = MockServer::start().await;
for (workspace, issuer) in [
(WORKSPACE_ID_INITIAL, "https://initial-idp.example"),
(WORKSPACE_ID_REFRESHED, "https://next-idp.example"),
] {
let contents = if policy_load_failure && workspace == WORKSPACE_ID_REFRESHED {
"invalid = [".to_string()
} else {
format!(
r#"[mcp_enterprise_managed_auth.idp]
issuer = "{issuer}"
client_id = "enterprise-client""#
)
};
Mock::given(method("GET"))
.and(path("/backend-api/wham/config/bundle"))
.and(header("chatgpt-account-id", workspace))
.respond_with(
ResponseTemplate::new(200).set_body_json(
CloudConfigBundleFixture::enterprise_config(contents).into_bundle(),
),
)
.expect(1)
.mount(&server)
.await;
}
Mock::given(method("POST"))
.and(path("/oauth/revoke"))
.and(body_partial_json(json!({
"token": "next-refresh-token",
"token_type_hint": "refresh_token",
})))
.respond_with(ResponseTemplate::new(200))
.expect(1)
.mount(&server)
.await;
let codex_home = TempDir::new()?;
create_config_toml(
codex_home.path(),
CreateConfigTomlParams {
chatgpt_base_url: Some(format!("{}/backend-api", server.uri())),
..Default::default()
},
)?;
write_chatgpt_auth(
codex_home.path(),
ChatGptAuthFixture::new("initial-token")
.refresh_token("initial-refresh-token")
.account_id(WORKSPACE_ID_INITIAL)
.chatgpt_user_id("enterprise-user")
.plan_type("business"),
AuthCredentialsStoreMode::File,
)?;
// Keep the real platform keyring out of this API test. Enterprise cleanup
// failure must still permit revocation and removal of the selected account.
if !policy_load_failure {
std::fs::write(codex_home.path().join("mcp-oauth-locks"), "not a directory")?;
}
let refresh_url = format!("{}/oauth/token", server.uri());
let mut app = TestAppServer::builder()
.with_codex_home(codex_home.path())
.without_auto_env()
.with_env_overrides(&[
("OPENAI_API_KEY", None),
(REFRESH_TOKEN_URL_OVERRIDE_ENV_VAR, Some(&refresh_url)),
])
.build_initialized_with_timeout(DEFAULT_READ_TIMEOUT)
.await?;
// Simulate a second process switching workspaces while this server caches A.
write_chatgpt_auth(
codex_home.path(),
ChatGptAuthFixture::new("next-token")
.refresh_token("next-refresh-token")
.account_id(WORKSPACE_ID_REFRESHED)
.chatgpt_user_id("enterprise-user")
.plan_type("business"),
AuthCredentialsStoreMode::File,
)?;
let id = app.send_logout_account_request().await?;
let response: LogoutAccountResponse =
timeout(DEFAULT_READ_TIMEOUT, app.read_response(id)).await??;
assert_eq!(response, LogoutAccountResponse {});
let notification = timeout(
DEFAULT_READ_TIMEOUT,
app.read_stream_until_notification_message("account/updated"),
)
.await??;
let notification: ServerNotification = notification.try_into()?;
let ServerNotification::AccountUpdated(account) = notification else {
bail!("unexpected notification: {notification:?}");
};
assert_eq!(
account,
AccountUpdatedNotification {
auth_mode: None,
plan_type: None,
}
);
assert!(!codex_home.path().join("auth.json").exists());
if policy_load_failure {
assert!(
!codex_home.path().join("mcp-oauth-locks").exists(),
"invalid current policy must not select any enterprise credential for deletion"
);
}
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn enterprise_login_attempt_lifecycle() -> Result<()> {
skip_if_no_network!(Ok(()));
const CHILD: &str = "CODEX_ENTERPRISE_ATTEMPT_TEST_CHILD";
if std::env::var_os(CHILD).is_none() {
// The app server and callback worker share only this child's mock keyring.
let home = TempDir::new()?;
let oauth_server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/oauth/revoke"))
.respond_with(ResponseTemplate::new(200))
.expect(1)
.mount(&oauth_server)
.await;
let output = Command::new(std::env::current_exe()?)
.args([
"--exact",
"suite::v2::account::enterprise_tests::enterprise_login_attempt_lifecycle",
"--nocapture",
])
.env(CHILD, "1")
.env("CODEX_HOME", home.path())
.env(
REFRESH_TOKEN_URL_OVERRIDE_ENV_VAR,
format!("{}/oauth/token", oauth_server.uri()),
)
.env_remove("OPENAI_API_KEY")
.env_remove("CODEX_API_KEY")
.env_remove(CODEX_ACCESS_TOKEN_ENV_VAR)
.current_dir(home.path())
.output()
.await?;
let stdout = String::from_utf8_lossy(&output.stdout);
assert!(
output.status.success() && stdout.contains("1 passed; 0 failed"),
"{stdout}\n{}",
String::from_utf8_lossy(&output.stderr)
);
return Ok(());
}
keyring::set_default_credential_builder(Box::<TestKeyring>::default());
let server = MockServer::builder().start().await;
let origin = server.uri();
let issuer = format!("{origin}/idp");
Mock::given(method("GET"))
.and(path("/backend-api/wham/config/bundle"))
.respond_with(
ResponseTemplate::new(200)
.set_body_json(CloudConfigBundleFixture::default().into_bundle()),
)
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/.well-known/oauth-authorization-server/idp"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"issuer": issuer,
"authorization_endpoint": format!("{issuer}/authorize"),
"token_endpoint": format!("{issuer}/token"),
"token_endpoint_auth_methods_supported": ["none"],
})))
.mount(&server)
.await;
let assertion = format!(
"{}.{}.signature",
URL_SAFE_NO_PAD.encode(r#"{"alg":"ES256"}"#),
URL_SAFE_NO_PAD.encode(
json!({"iss":issuer,"sub":"enterprise-user","aud":"idp-client","exp":4102444800_u64})
.to_string()
)
);
Mock::given(method("POST"))
.and(path("/idp/token"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"access_token":"enterprise-access", "refresh_token":"enterprise-refresh",
"id_token":assertion, "token_type":"Bearer",
})))
.mount(&server)
.await;
let home = TempDir::new()?;
let listener = std::net::TcpListener::bind("127.0.0.1:0")?;
let callback_port = listener.local_addr()?.port();
MockResponsesConfig::new(&origin)
.enable_feature(Feature::UseXaa)
.disable_feature(Feature::SecretAuthStorage)
.disable_feature(Feature::Apps)
.with_root_config(&format!(
r#"cli_auth_credentials_store = "file"
chatgpt_base_url = "{origin}/backend-api"
mcp_oauth_callback_port = {callback_port}
analytics = {{ enabled = false }}"#
))
.with_provider_config("requires_openai_auth = true")
.with_extra_config(&format!(
r#"[mcp_enterprise_managed_auth.idp]
issuer = "{issuer}"
client_id = "idp-client"
[mcp_servers.enterprise]
url = "{origin}/mcp"
auth = "ema_auth"
scopes = ["files.read"]
oauth = {{ client_id = "mcp-client", authorization_server_issuer = "{origin}/as" }}
[mcp_servers.disabled-enterprise]
url = "{origin}/mcp"
enabled = false
auth = "ema_auth"
oauth = {{ client_id = "mcp-client", authorization_server_issuer = "{origin}/as" }}"#
))
.write(home.path())?;
write_chatgpt_auth(
home.path(),
ChatGptAuthFixture::new("account-access")
.account_id(WORKSPACE_ID_INITIAL)
.chatgpt_user_id("enterprise-user")
.plan_type("business"),
AuthCredentialsStoreMode::File,
)?;
let loader_overrides = LoaderOverrides::without_managed_config_for_tests();
let config = Arc::new(
ConfigBuilder::default()
.codex_home(home.path().to_path_buf())
.fallback_cwd(Some(home.path().to_path_buf()))
.loader_overrides(loader_overrides.clone())
.build()
.await?,
);
let mut client = in_process::start(InProcessStartArgs {
arg0_paths: Arg0DispatchPaths::default(),
config,
cli_overrides: Vec::new(),
loader_overrides,
strict_config: false,
cloud_config_bundle: CloudConfigBundleLoader::default(),
thread_config_loader: Arc::new(NoopThreadConfigLoader),
feedback: CodexFeedback::new(),
log_db: None,
state_db: None,
environment_manager: Arc::new(EnvironmentManager::default_for_tests()),
config_warnings: Vec::new(),
embedded_network_policy: Default::default(),
session_source: SessionSource::Cli,
enable_codex_api_key_env: false,
initialize: InitializeParams {
client_info: ClientInfo {
name: "codex-app-server-tests".into(),
title: None,
version: "0.1.0".into(),
},
capabilities: Some(InitializeCapabilities {
experimental_api: true,
..Default::default()
}),
},
channel_capacity: in_process::DEFAULT_IN_PROCESS_CHANNEL_CAPACITY,
})
.await?;
drop(listener);
let thread = enterprise_rpc(&client, "thread/start", json!({})).await?;
let thread_id = thread["thread"]["id"].as_str().context("thread ID")?;
assert_eq!(
enterprise_rpc(
&client,
"mcpServer/oauth/login",
json!({"name":"disabled-enterprise","threadId":thread_id}),
)
.await
.expect_err("disabled enterprise MCP login must reject immediately")
.to_string(),
"mcpServer/oauth/login: MCP server 'disabled-enterprise' is disabled."
);
let first = enterprise_rpc(
&client,
"mcpServer/oauth/login",
json!({"name":"enterprise","threadId":thread_id}),
)
.await?;
enterprise_callback(&first, &issuer, callback_port)?;
let replacement = enterprise_rpc(
&client,
"mcpServer/oauth/login",
json!({"name":"enterprise","threadId":thread_id}),
)
.await?;
assert_ne!(replacement["loginId"], first["loginId"]);
enterprise_callback(&replacement, &issuer, callback_port)?;
assert_eq!(
enterprise_completion(&mut client).await?,
json!({"name":"enterprise","threadId":thread_id,"loginId":first["loginId"],"success":false,"error":"Enterprise sign-in failed."})
);
assert_eq!(
enterprise_rpc(
&client,
"account/login/cancel",
json!({"loginId":replacement["loginId"]}),
)
.await?,
json!({"status":"canceled"})
);
assert_eq!(
enterprise_completion(&mut client).await?,
json!({"name":"enterprise","threadId":thread_id,"loginId":replacement["loginId"],"success":false,"error":"Enterprise sign-in failed."})
);
let idp = McpServerIdpOAuthConfig {
issuer: issuer.clone(),
client_id: "idp-client".into(),
};
let initial_scope = McpEmaAuthScope::new("enterprise-user".into(), WORKSPACE_ID_INITIAL.into())
.expect("initial scope");
assert_eq!(
stored_oauth_credentials(
&idp.credential_name(&initial_scope),
&issuer,
OAuthCredentialsStoreMode::Keyring,
AuthKeyringBackendKind::Direct,
)?,
None
);
let previous = enterprise_rpc(
&client,
"mcpServer/oauth/login",
json!({"name":"enterprise","threadId":thread_id}),
)
.await?;
let previous_callback = enterprise_callback(&previous, &issuer, callback_port)?;
write_chatgpt_auth(
home.path(),
ChatGptAuthFixture::new("account-access")
.account_id(WORKSPACE_ID_REFRESHED)
.chatgpt_user_id("enterprise-user")
.plan_type("business"),
AuthCredentialsStoreMode::File,
)?;
let browser = HttpClientBuilder::new()
.without_request_logging()
.without_redirects()
.build_direct()?;
// Each attempt owns a one-shot listener on this port. Do not reuse an HTTP
// connection from the previous callback after that listener has retired.
let response = timeout(
DEFAULT_READ_TIMEOUT,
browser
.get(previous_callback)
.header(http::header::CONNECTION, "close")
.send(),
)
.await??;
assert_eq!(response.status(), http::StatusCode::OK);
assert_eq!(
enterprise_completion(&mut client).await?,
json!({"name":"enterprise","threadId":thread_id,"loginId":previous["loginId"],"success":false,"error":"Enterprise sign-in failed."})
);
let retry = enterprise_rpc(
&client,
"mcpServer/oauth/login",
json!({"name":"enterprise","threadId":thread_id}),
)
.await?;
let retry_callback = enterprise_callback(&retry, &issuer, callback_port)?;
let response = timeout(
DEFAULT_READ_TIMEOUT,
browser
.get(retry_callback)
.header(http::header::CONNECTION, "close")
.send(),
)
.await??;
assert_eq!(response.status(), http::StatusCode::OK);
assert_eq!(
enterprise_completion(&mut client).await?,
json!({"name":"enterprise","threadId":thread_id,"loginId":retry["loginId"],"success":true})
);
let replacement_scope =
McpEmaAuthScope::new("enterprise-user".into(), WORKSPACE_ID_REFRESHED.into())
.expect("replacement scope");
for (scope, present) in [(initial_scope, false), (replacement_scope, true)] {
assert_eq!(
stored_oauth_credentials(
&idp.credential_name(&scope),
&issuer,
OAuthCredentialsStoreMode::Keyring,
AuthKeyringBackendKind::Direct,
)?
.is_some(),
present
);
}
let stalled = enterprise_rpc(
&client,
"mcpServer/oauth/login",
json!({"name":"enterprise","threadId":thread_id}),
)
.await?;
let mut unfinished = tokio::net::TcpStream::connect(("127.0.0.1", callback_port)).await?;
unfinished
.write_all(b"POST /unfinished HTTP/1.1\r\nHost: 127.0.0.1\r\nContent-Length: 2048\r\n\r\n")
.await?;
// Reading the response proves the callback worker has accepted this request.
// It cannot finish draining the declared body while this socket stays open.
let mut response = [0; 256];
assert!(timeout(DEFAULT_READ_TIMEOUT, unfinished.read(&mut response)).await?? > 0);
let logout = timeout(
std::time::Duration::from_secs(15),
enterprise_rpc(&client, "account/logout", json!({})),
)
.await;
// Release the worker even when the regression fails, so child shutdown cannot hang.
drop(unfinished);
logout.context("an unfinished callback must not prevent primary logout")??;
assert!(!home.path().join("auth.json").exists());
assert!(
stored_oauth_credentials(
&idp.credential_name(
&McpEmaAuthScope::new("enterprise-user".into(), WORKSPACE_ID_REFRESHED.into(),)
.expect("replacement scope")
),
&issuer,
OAuthCredentialsStoreMode::Keyring,
AuthKeyringBackendKind::Direct,
)?
.is_none()
);
assert_eq!(
enterprise_completion(&mut client).await?,
json!({"name":"enterprise","threadId":thread_id,"loginId":stalled["loginId"],"success":false,"error":"Enterprise sign-in failed."})
);
client.shutdown().await?;
Ok(())
}
fn enterprise_callback(start: &Value, issuer: &str, callback_port: u16) -> Result<Url> {
let query = Url::parse(
start["authorizationUrl"]
.as_str()
.context("authorization URL")?,
)?
.query_pairs()
.into_owned()
.collect::<HashMap<_, _>>();
let mut callback = Url::parse(query.get("redirect_uri").context("redirect_uri")?)?;
assert_eq!(
(callback.scheme(), callback.host_str(), callback.port()),
("http", Some("127.0.0.1"), Some(callback_port))
);
callback
.query_pairs_mut()
.append_pair("code", "enterprise-code")
.append_pair("state", query.get("state").context("state")?)
.append_pair("iss", issuer);
Ok(callback)
}
async fn enterprise_rpc(
client: &InProcessClientHandle,
method: &str,
params: Value,
) -> Result<Value> {
let request = serde_json::from_value(
json!({"id":uuid::Uuid::new_v4().to_string(),"method":method,"params":params}),
)?;
let response = timeout(DEFAULT_READ_TIMEOUT, client.request(request))
.await
.with_context(|| format!("timed out waiting for {method}"))??;
response.map_err(|error| anyhow::anyhow!("{method}: {}", error.message))
}
async fn enterprise_completion(client: &mut InProcessClientHandle) -> Result<Value> {
timeout(DEFAULT_READ_TIMEOUT, async {
loop {
match client
.next_event()
.await
.context("app-server notification")?
{
InProcessServerEvent::ServerNotification(notification) => {
if let ServerNotification::McpServerOauthLoginCompleted(notification) =
*notification
{
return Ok(serde_json::to_value(notification)?);
}
}
event => bail!("unexpected app-server event: {event:?}"),
}
}
})
.await
.context("timed out waiting for enterprise login completion")?
}
#[derive(Default)]
struct TestKeyring(Mutex<HashMap<(String, String), Arc<MockCredential>>>);
struct TestCredential(Arc<MockCredential>);
impl CredentialApi for TestCredential {
fn set_secret(&self, secret: &[u8]) -> keyring::Result<()> {
self.0.set_secret(secret)
}
fn get_secret(&self) -> keyring::Result<Vec<u8>> {
self.0.get_secret()
}
fn delete_credential(&self) -> keyring::Result<()> {
self.0.delete_credential()
}
fn as_any(&self) -> &dyn Any {
self
}
}
impl CredentialBuilderApi for TestKeyring {
fn build(
&self,
_target: Option<&str>,
service: &str,
user: &str,
) -> keyring::Result<Box<Credential>> {
let credential = self
.0
.lock()
.expect("mock keyring")
.entry((service.into(), user.into()))
.or_default()
.clone();
Ok(Box::new(TestCredential(credential)))
}
fn as_any(&self) -> &dyn Any {
self
}
fn persistence(&self) -> CredentialPersistence {
CredentialPersistence::ProcessOnly
}
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn enterprise_login_setup_is_cancelled_by_account_changes() -> Result<()> {
skip_if_no_network!(Ok(()));
const CHILD: &str = "CODEX_ENTERPRISE_SETUP_TEST_CHILD";
if std::env::var_os(CHILD).is_none() {
// The app server and callback worker share only this child's mock keyring.
let home = TempDir::new()?;
let oauth_server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/oauth/revoke"))
.respond_with(ResponseTemplate::new(200))
.expect(0..=1)
.mount(&oauth_server)
.await;
let output = Command::new(std::env::current_exe()?)
.args([
"--exact",
"suite::v2::account::enterprise_tests::enterprise_login_setup_is_cancelled_by_account_changes",
"--nocapture",
])
.env(CHILD, "1")
.env("CODEX_HOME", home.path())
.env(
REFRESH_TOKEN_URL_OVERRIDE_ENV_VAR,
format!("{}/oauth/token", oauth_server.uri()),
)
.env_remove("OPENAI_API_KEY")
.env_remove("CODEX_API_KEY")
.env_remove(CODEX_ACCESS_TOKEN_ENV_VAR)
.current_dir(home.path())
.output()
.await?;
let stdout = String::from_utf8_lossy(&output.stdout);
assert!(
output.status.success() && stdout.contains("1 passed; 0 failed"),
"{stdout}\n{}",
String::from_utf8_lossy(&output.stderr)
);
return Ok(());
}
keyring::set_default_credential_builder(Box::<TestKeyring>::default());
for (account_method, account_params) in [
("account/logout", json!({})),
(
"account/login/start",
json!({"type":"apiKey","apiKey":"replacement-api-key"}),
),
] {
let server = MockServer::builder().start().await;
let origin = server.uri();
let issuer = format!("{origin}/idp");
Mock::given(method("GET"))
.and(path("/backend-api/wham/config/bundle"))
.respond_with(
ResponseTemplate::new(200)
.set_body_json(CloudConfigBundleFixture::default().into_bundle()),
)
.mount(&server)
.await;
let discovery_started = Arc::new(tokio::sync::Notify::new());
let entered = Arc::clone(&discovery_started);
Mock::given(method("GET"))
.and(path("/.well-known/oauth-authorization-server/idp"))
.respond_with(move |_: &wiremock::Request| {
entered.notify_one();
ResponseTemplate::new(200).set_delay(std::time::Duration::from_secs(60))
})
.mount(&server)
.await;
let assertion = format!(
"{}.{}.signature",
URL_SAFE_NO_PAD.encode(r#"{"alg":"ES256"}"#),
URL_SAFE_NO_PAD.encode(
json!({"iss":issuer,"sub":"enterprise-user","aud":"idp-client","exp":4102444800_u64})
.to_string()
)
);
Mock::given(method("POST"))
.and(path("/idp/token"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"access_token":"enterprise-access", "refresh_token":"enterprise-refresh",
"id_token":assertion, "token_type":"Bearer",
})))
.mount(&server)
.await;
let home = TempDir::new()?;
let listener = std::net::TcpListener::bind("127.0.0.1:0")?;
let callback_port = listener.local_addr()?.port();
MockResponsesConfig::new(&origin)
.enable_feature(Feature::UseXaa)
.disable_feature(Feature::SecretAuthStorage)
.disable_feature(Feature::Apps)
.with_root_config(&format!(
r#"cli_auth_credentials_store = "file"
chatgpt_base_url = "{origin}/backend-api"
mcp_oauth_callback_port = {callback_port}
analytics = {{ enabled = false }}"#
))
.with_provider_config("requires_openai_auth = true")
.with_extra_config(&format!(
r#"[mcp_enterprise_managed_auth.idp]
issuer = "{issuer}"
client_id = "idp-client"
[mcp_servers.enterprise]
url = "{origin}/mcp"
auth = "ema_auth"
scopes = ["files.read"]
oauth = {{ client_id = "mcp-client", authorization_server_issuer = "{origin}/as" }}
[mcp_servers.disabled-enterprise]
url = "{origin}/mcp"
enabled = false
auth = "ema_auth"
oauth = {{ client_id = "mcp-client", authorization_server_issuer = "{origin}/as" }}"#
))
.write(home.path())?;
write_chatgpt_auth(
home.path(),
ChatGptAuthFixture::new("account-access")
.account_id(WORKSPACE_ID_INITIAL)
.chatgpt_user_id("enterprise-user")
.plan_type("business"),
AuthCredentialsStoreMode::File,
)?;
let loader_overrides = LoaderOverrides::without_managed_config_for_tests();
let config = Arc::new(
ConfigBuilder::default()
.codex_home(home.path().to_path_buf())
.fallback_cwd(Some(home.path().to_path_buf()))
.loader_overrides(loader_overrides.clone())
.build()
.await?,
);
let client = in_process::start(InProcessStartArgs {
arg0_paths: Arg0DispatchPaths::default(),
config,
cli_overrides: Vec::new(),
loader_overrides,
strict_config: false,
cloud_config_bundle: CloudConfigBundleLoader::default(),
thread_config_loader: Arc::new(NoopThreadConfigLoader),
feedback: CodexFeedback::new(),
log_db: None,
state_db: None,
environment_manager: Arc::new(EnvironmentManager::default_for_tests()),
config_warnings: Vec::new(),
embedded_network_policy: Default::default(),
session_source: SessionSource::Cli,
enable_codex_api_key_env: false,
initialize: InitializeParams {
client_info: ClientInfo {
name: "codex-app-server-tests".into(),
title: None,
version: "0.1.0".into(),
},
capabilities: Some(InitializeCapabilities {
experimental_api: true,
..Default::default()
}),
},
channel_capacity: in_process::DEFAULT_IN_PROCESS_CHANNEL_CAPACITY,
})
.await?;
drop(listener);
let thread = enterprise_rpc(&client, "thread/start", json!({})).await?;
let thread_id = thread["thread"]["id"].as_str().context("thread ID")?;
let (setup_result, account_result) = tokio::join!(
enterprise_rpc(
&client,
"mcpServer/oauth/login",
json!({"name":"enterprise","threadId":thread_id}),
),
async {
timeout(DEFAULT_READ_TIMEOUT, discovery_started.notified()).await?;
timeout(
std::time::Duration::from_secs(15),
enterprise_rpc(&client, account_method, account_params),
)
.await
.context("account changes must not wait for IdP discovery")?
},
);
account_result.context("stalled enterprise setup must not block account changes")?;
assert_eq!(
setup_result
.expect_err("account changes cancel pending enterprise setup")
.to_string(),
"mcpServer/oauth/login: enterprise sign-in was cancelled",
);
let idp = McpServerIdpOAuthConfig {
issuer: issuer.clone(),
client_id: "idp-client".into(),
};
let scope =
McpEmaAuthScope::new("enterprise-user".into(), WORKSPACE_ID_INITIAL.into()).unwrap();
assert_eq!(
stored_oauth_credentials(
&idp.credential_name(&scope),
&issuer,
OAuthCredentialsStoreMode::Keyring,
AuthKeyringBackendKind::Direct,
)?,
None
);
client.shutdown().await?;
}
Ok(())
}
@@ -1310,6 +1310,11 @@ startup_timeout_sec = 10
McpServerOauthLoginCompletedNotification {
name: OAUTH_MCP_SERVER_NAME.to_string(),
thread_id: Some(selected_thread.clone()),
login_id: Some(
response
.login_id
.expect("login response should contain an ID")
),
success: true,
error: None,
}
@@ -372,6 +372,15 @@ async fn oauth_login_validates_callback_issuer_and_uses_http_headers_helper(
mcp.read_notification("mcpServer/oauthLogin/completed"),
)
.await??;
assert_eq!(
completed.login_id.as_deref(),
Some(
response
.login_id
.as_deref()
.expect("login response should contain an ID")
)
);
assert_eq!(
(
completed.name.as_str(),
@@ -629,11 +638,21 @@ async fn oauth_login_automatically_selects_callback_specific_cimd_without_metada
app_server.read_notification("mcpServer/oauthLogin/completed"),
)
.await??;
assert_eq!(
completed.login_id.as_deref(),
Some(
response
.login_id
.as_deref()
.expect("login response should contain an ID")
)
);
assert_eq!(
completed,
McpServerOauthLoginCompletedNotification {
name: "cimd".to_string(),
thread_id: None,
login_id: response.login_id,
success: true,
error: None,
}
@@ -2037,11 +2037,21 @@ async fn plugin_oauth_login_preserves_registered_callbacks_or_uses_legacy_fallba
mcp.read_notification("mcpServer/oauthLogin/completed"),
)
.await??;
assert_eq!(
completed.login_id.as_deref(),
Some(
response
.login_id
.as_deref()
.expect("login response should contain an ID")
)
);
assert_eq!(
completed,
McpServerOauthLoginCompletedNotification {
name: "sample-mcp".to_string(),
thread_id: None,
login_id: response.login_id,
success: true,
error: None,
}
+34
View File
@@ -22,8 +22,10 @@ use codex_login::login_with_api_key;
use codex_login::logout_with_revoke;
use codex_login::run_device_code_login;
use codex_login::run_login_server;
use codex_mcp::ema_auth_scope;
use codex_protocol::auth::AuthMode;
use codex_protocol::config_types::ForcedLoginMethod;
use codex_rmcp_client::EnterpriseOAuthCredentialGuard;
use codex_utils_cli::CliConfigOverrides;
use std::fs::OpenOptions;
use std::io::IsTerminal;
@@ -511,6 +513,37 @@ pub async fn run_logout(cli_config_overrides: CliConfigOverrides) -> ! {
let config = load_config_or_exit(cli_config_overrides).await;
let auth_route_config = config.auth_route_config();
// Cleanup is independent of the current feature gate. Retain the grant lock
// through account removal so another process cannot commit a late login.
let enterprise_guard = async {
let Some(profile) = &config.mcp_enterprise_managed_auth else {
return Ok::<_, anyhow::Error>(None);
};
let auth = config
.auth_config()
.load_auth(/*enable_codex_api_key_env*/ false)
.await?;
let Some(scope) = ema_auth_scope(auth.as_ref()) else {
return Ok(None);
};
EnterpriseOAuthCredentialGuard::acquire(
&profile.idp.credential_name(&scope),
&profile.idp.issuer,
config.auth_keyring_backend_kind(),
)
.await
.map(Some)
}
.await;
let cleanup_failed = match &enterprise_guard {
Ok(Some(guard)) => guard.delete_tokens().is_err(),
Ok(None) => false,
Err(_) => true,
};
if cleanup_failed {
eprintln!("Warning: failed to remove enterprise authorization; continuing account logout");
}
let logged_out = match logout_with_revoke(
&config.codex_home,
config.cli_auth_credentials_store_mode,
@@ -525,6 +558,7 @@ pub async fn run_logout(cli_config_overrides: CliConfigOverrides) -> ! {
std::process::exit(1);
}
};
drop(enterprise_guard);
let cleared_bedrock_config =
if let Some(paths) = ConfigEditsBuilder::bedrock_provider_config_paths_to_clear(&config) {
+33
View File
@@ -8,6 +8,7 @@ use anyhow::bail;
use clap::ArgGroup;
use clap::builder::TypedValueParser;
use codex_config::types::AppToolApproval;
use codex_config::types::McpServerAuth;
use codex_config::types::McpServerConfig;
use codex_config::types::McpServerOAuthConfig;
use codex_config::types::McpServerTransportConfig;
@@ -28,6 +29,7 @@ use codex_mcp::McpRuntimeContext;
use codex_mcp::apply_http_headers_helper;
use codex_mcp::compute_auth_statuses;
use codex_mcp::discover_supported_scopes;
use codex_mcp::ema_auth_scope;
use codex_mcp::oauth_login_support;
use codex_mcp::resolve_oauth_callback;
use codex_mcp::resolve_oauth_scopes;
@@ -36,6 +38,7 @@ use codex_rmcp_client::McpOAuthCallbackMode;
use codex_rmcp_client::McpOAuthClientRegistration;
use codex_rmcp_client::OAuthDiscoveryTimeout;
use codex_rmcp_client::StreamableHttpRedirectMode;
use codex_rmcp_client::delete_enterprise_oauth_tokens;
use codex_rmcp_client::delete_oauth_tokens;
use codex_rmcp_client::resolve_mcp_oauth_callback_url;
use codex_utils_cli::CliConfigOverrides;
@@ -639,6 +642,36 @@ async fn run_logout(config: &Config, logout_args: LogoutArgs) -> Result<()> {
.get(&name)
.ok_or_else(|| anyhow!("No MCP server named '{name}' found in configuration."))?;
if matches!(server.auth, McpServerAuth::EmaAuth) {
let auth_manager =
AuthManager::shared_from_config(config, /*enable_codex_api_key_env*/ true).await?;
let auth = auth_manager.auth().await;
let scope = ema_auth_scope(auth.as_ref())
.context("Sign in to your Codex account before removing enterprise authorization")?;
// Cleanup must remain available even when this server is no longer eligible
// for EMA (for example, after disabling the feature or changing its transport).
let profile = config
.mcp_enterprise_managed_auth
.as_ref()
.context("EMA logout requires a trusted enterprise IdP profile")?;
let idp = &profile.idp;
let credential_name = idp.credential_name(&scope);
match delete_enterprise_oauth_tokens(
&credential_name,
&idp.issuer,
config.auth_keyring_backend_kind(),
)
.await
{
Ok(true) => println!(
"Removed the shared enterprise authorization used by EMA MCP servers (selected via '{name}')."
),
Ok(false) => println!("No shared enterprise authorization is stored."),
Err(_) => return Err(anyhow!("failed to delete enterprise authorization")),
}
return Ok(());
}
let url = match &server.transport {
McpServerTransportConfig::StreamableHttp { url, .. } => url.clone(),
_ => bail!("OAuth logout is only supported for streamable_http transports."),
+73
View File
@@ -183,6 +183,79 @@ fn logout_clears_only_the_selected_bedrock_provider() -> Result<()> {
Ok(())
}
#[tokio::test]
async fn logout_survives_enterprise_cleanup_failure_with_xaa_disabled() -> Result<()> {
let server = MockServer::start().await;
Mock::given(path("/backend-api/wham/config/bundle"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"config_toml": {"enterprise_managed": []},
"requirements_toml": {"enterprise_managed": []},
})))
.mount(&server)
.await;
Mock::given(method("POST"))
.and(path("/oauth/revoke"))
.respond_with(ResponseTemplate::new(200))
.expect(1)
.mount(&server)
.await;
let home = TempDir::new()?;
std::fs::write(
home.path().join("config.toml"),
format!(
"cli_auth_credentials_store = \"file\"\nchatgpt_base_url = \"{}/backend-api\"\n[features]\nuse_xaa = false\n[mcp_enterprise_managed_auth.idp]\nissuer = \"https://idp.example\"\nclient_id = \"enterprise-client\"\n[mcp_servers.enterprise]\nurl = \"https://resource.example/mcp\"\nauth = \"ema_auth\"\nbearer_token_env_var = \"UNUSED_TOKEN\"\n",
server.uri()
),
)?;
write_chatgpt_auth(
home.path(),
ChatGptAuthFixture::new("account-access")
.account_id("workspace")
.chatgpt_user_id("user"),
AuthCredentialsStoreMode::File,
)?;
// Fail before touching the real keyring; the subprocess owns this isolated home.
std::fs::write(home.path().join("mcp-oauth-locks"), "not a directory")?;
for enabled in [false, true] {
codex_command(home.path())?
.current_dir(home.path())
.env("NO_PROXY", "127.0.0.1,localhost")
.env("no_proxy", "127.0.0.1,localhost")
.env_remove("CODEX_API_KEY")
.env_remove("OPENAI_API_KEY")
.env_remove(CODEX_ACCESS_TOKEN_ENV_VAR)
.args([
"-c",
&format!("features.use_xaa={enabled}"),
"mcp",
"logout",
"enterprise",
])
.assert()
.failure()
.stderr(contains("failed to delete enterprise authorization"));
assert!(home.path().join("auth.json").exists());
}
codex_command(home.path())?
.current_dir(home.path())
.env("NO_PROXY", "127.0.0.1,localhost")
.env("no_proxy", "127.0.0.1,localhost")
.env(
REVOKE_TOKEN_URL_OVERRIDE_ENV_VAR,
format!("{}/oauth/revoke", server.uri()),
)
.env_remove("CODEX_API_KEY")
.env_remove("OPENAI_API_KEY")
.env_remove(CODEX_ACCESS_TOKEN_ENV_VAR)
.args(["logout"])
.assert()
.success()
.stderr(contains("continuing account logout"))
.stderr(contains("Successfully logged out"));
assert!(!home.path().join("auth.json").exists());
Ok(())
}
#[test]
fn login_with_access_token_rejects_invalid_jwt() -> Result<()> {
let codex_home = TempDir::new()?;
+9
View File
@@ -0,0 +1,9 @@
//! Derives the enterprise credential owner from the active Codex account.
use codex_config::McpEmaAuthScope;
use codex_login::CodexAuth;
pub fn ema_auth_scope(auth: Option<&CodexAuth>) -> Option<McpEmaAuthScope> {
let auth = auth?;
McpEmaAuthScope::new(auth.get_chatgpt_user_id()?, auth.get_account_id()?)
}
+2
View File
@@ -8,6 +8,7 @@ pub use elicitation::ElicitationLifecycle;
pub use elicitation::ElicitationReviewRequest;
pub use elicitation::ElicitationReviewer;
pub use elicitation::ElicitationReviewerHandle;
pub use ema_auth::ema_auth_scope;
pub use event_stream::McpEventStreamOpener;
pub use resource_client::CodexAppsResourceListParams;
pub use resource_client::McpEventCatalogSnapshot;
@@ -110,6 +111,7 @@ mod client_tool_catalog;
pub(crate) mod codex_apps;
pub(crate) mod connection_manager;
pub(crate) mod elicitation;
mod ema_auth;
mod event_stream;
mod executor_environment_http_client;
pub(crate) mod mcp;
+1
View File
@@ -171,6 +171,7 @@ pub use marketplace_edit::record_user_marketplace;
pub use marketplace_edit::remove_user_marketplace;
pub use marketplace_edit::remove_user_marketplace_config;
pub use mcp_edit::load_global_mcp_servers;
pub use mcp_ema::McpEmaAuthScope;
pub use mcp_ema::McpEmaRegistration;
pub use mcp_ema::McpEnterpriseManagedAuthConfig;
pub use mcp_ema::McpServerIdpOAuthConfig;
+39 -1
View File
@@ -1,13 +1,17 @@
//! Enterprise MCP configuration and registration provenance.
//! Enterprise MCP configuration and account-scoped credential names.
//!
//! Only host, user, or managed configuration selects the IdP. Plugin declarations
//! and project settings may not redirect the enterprise credential source.
use std::io;
use base64::Engine;
use base64::engine::general_purpose::URL_SAFE_NO_PAD;
use schemars::JsonSchema;
use serde::Deserialize;
use serde::Serialize;
use sha2::Digest;
use sha2::Sha256;
use codex_features::Feature;
@@ -70,6 +74,40 @@ impl McpEmaRegistration {
}
}
/// Account identity that owns an enterprise OAuth credential.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct McpEmaAuthScope {
user_id: String,
workspace_id: String,
}
impl McpEmaAuthScope {
pub fn new(user_id: String, workspace_id: String) -> Option<Self> {
(!user_id.trim().is_empty() && !workspace_id.trim().is_empty()).then_some(Self {
user_id,
workspace_id,
})
}
}
impl McpServerIdpOAuthConfig {
/// Shared by allowed MCP resources, but never by another account or IdP client.
pub fn credential_name(&self, scope: &McpEmaAuthScope) -> String {
let mut digest = Sha256::new();
digest.update(b"codex-ema-idp-credential-v1\0");
for part in [
&scope.user_id,
&scope.workspace_id,
&self.issuer,
&self.client_id,
] {
digest.update((part.len() as u64).to_be_bytes());
digest.update(part.as_bytes());
}
format!("ema-idp:{}", URL_SAFE_NO_PAD.encode(digest.finalize()))
}
}
impl McpEnterpriseManagedAuthConfig {
/// Resolves the trusted profile and checks provenance at a config load boundary.
pub fn resolve(
+39 -1
View File
@@ -1,4 +1,4 @@
//! Enterprise registration provenance checks.
//! Enterprise registration provenance and credential-boundary checks.
use super::*;
use crate::AbsolutePathBuf;
@@ -215,6 +215,44 @@ fn xaa_opt_in_requires_a_non_project_source() {
);
}
#[test]
fn ema_credential_names_bind_user_workspace_issuer_and_client() {
let mut names = std::collections::HashSet::new();
for user in ["a", "b"] {
for workspace in ["a", "b"] {
for issuer in ["https://one.example", "https://two.example"] {
for client in ["a", "b"] {
let scope = McpEmaAuthScope::new(user.into(), workspace.into()).unwrap();
let idp = McpServerIdpOAuthConfig {
issuer: issuer.into(),
client_id: client.into(),
};
assert!(names.insert(idp.credential_name(&scope)));
}
}
}
}
assert_eq!(McpEmaAuthScope::new(" ".into(), "workspace".into()), None);
let identifiers = [
"sensitive-user@example.com",
"sensitive-workspace-id",
"https://sensitive-idp.example",
"sensitive-oauth-client",
];
let scope = McpEmaAuthScope::new(identifiers[0].into(), identifiers[1].into()).unwrap();
let name = McpServerIdpOAuthConfig {
issuer: identifiers[2].into(),
client_id: identifiers[3].into(),
}
.credential_name(&scope);
assert!(name.starts_with("ema-idp:"));
for identifier in identifiers {
assert!(!name.contains(identifier));
assert!(!name.contains(&URL_SAFE_NO_PAD.encode(identifier)));
}
}
#[test]
fn ema_rejects_alternate_credentials_and_executor_custody() {
for extra in [
+1
View File
@@ -3,6 +3,7 @@
// Note this file should generally be restricted to simple struct/enum
// definitions that do not contain business logic.
pub use crate::mcp_ema::McpEmaAuthScope;
pub use crate::mcp_ema::McpEnterpriseManagedAuthConfig;
pub use crate::mcp_ema::McpServerIdpOAuthConfig;
pub use crate::mcp_types::AppToolApproval;
+1
View File
@@ -61,6 +61,7 @@ tokio = { workspace = true, features = [
"io-std",
"time",
] }
tokio-util = { workspace = true, features = ["rt"] }
tracing = { workspace = true, features = ["log"] }
url = { workspace = true }
urlencoding = { workspace = true }
+19 -14
View File
@@ -336,10 +336,14 @@ mod tests {
handle: JoinHandle<()>,
}
fn test_http_client() -> Arc<dyn HttpClient> {
Arc::new(RouteAwareHttpClient::new(HttpClientFactory::new(
OutboundProxyPolicy::ReqwestDefault,
)))
async fn test_http_client() -> Arc<dyn HttpClient> {
let client: Arc<dyn HttpClient> = Arc::new(RouteAwareHttpClient::new(
HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
));
crate::oauth::test_support::warm_http_client(client.as_ref())
.await
.expect("warm production HTTP client before discovery deadlines");
client
}
impl Drop for TestServer {
@@ -487,7 +491,7 @@ mod tests {
/*env_http_headers*/ None,
OAuthCredentialsStoreMode::Keyring,
AuthKeyringBackendKind::default(),
test_http_client(),
Arc::new(RecordingHttpClient::default()),
OAuthDiscoveryTimeout::Requested,
StreamableHttpRedirectMode::Legacy,
)
@@ -512,7 +516,7 @@ mod tests {
)])),
OAuthCredentialsStoreMode::Keyring,
AuthKeyringBackendKind::default(),
test_http_client(),
Arc::new(RecordingHttpClient::default()),
OAuthDiscoveryTimeout::Requested,
StreamableHttpRedirectMode::Legacy,
)
@@ -567,7 +571,7 @@ mod tests {
&url,
/*http_headers*/ None,
/*env_http_headers*/ None,
test_http_client(),
test_http_client().await,
OAuthDiscoveryTimeout::LOCAL,
StreamableHttpRedirectMode::Legacy,
)
@@ -610,7 +614,7 @@ mod tests {
"sensitive-key".to_string(),
)])),
/*env_http_headers*/ None,
test_http_client(),
test_http_client().await,
OAuthDiscoveryTimeout::LOCAL,
StreamableHttpRedirectMode::Legacy,
)
@@ -632,6 +636,7 @@ mod tests {
#[tokio::test]
async fn determine_auth_status_preserves_transient_http_errors() {
let _home = crate::oauth::test_support::TempCodexHome::new();
for status in [
StatusCode::REQUEST_TIMEOUT,
StatusCode::TOO_EARLY,
@@ -653,7 +658,7 @@ mod tests {
/*env_http_headers*/ None,
OAuthCredentialsStoreMode::File,
AuthKeyringBackendKind::default(),
test_http_client(),
test_http_client().await,
OAuthDiscoveryTimeout::LOCAL,
StreamableHttpRedirectMode::Legacy,
)
@@ -685,7 +690,7 @@ mod tests {
&server.url,
/*http_headers*/ None,
/*env_http_headers*/ None,
test_http_client(),
test_http_client().await,
OAuthDiscoveryTimeout::LOCAL,
StreamableHttpRedirectMode::Legacy,
)
@@ -716,7 +721,7 @@ mod tests {
&server.url,
/*http_headers*/ None,
/*env_http_headers*/ None,
test_http_client(),
test_http_client().await,
OAuthDiscoveryTimeout::LOCAL,
StreamableHttpRedirectMode::Legacy,
)
@@ -872,7 +877,7 @@ mod tests {
&resource_server.url,
/*http_headers*/ None,
/*env_http_headers*/ None,
test_http_client(),
test_http_client().await,
OAuthDiscoveryTimeout::LOCAL,
StreamableHttpRedirectMode::Legacy,
)
@@ -899,7 +904,7 @@ mod tests {
&server.url,
/*http_headers*/ None,
/*env_http_headers*/ None,
test_http_client(),
test_http_client().await,
OAuthDiscoveryTimeout::LOCAL,
StreamableHttpRedirectMode::Legacy,
)
@@ -922,7 +927,7 @@ mod tests {
&server.url,
/*http_headers*/ None,
/*env_http_headers*/ None,
test_http_client(),
test_http_client().await,
OAuthDiscoveryTimeout::LOCAL,
StreamableHttpRedirectMode::Legacy,
)
@@ -119,6 +119,12 @@ impl EnterpriseOAuthLoginHandle {
self.flow.authorization_url()
}
/// Capture before waiting; after the handle or wait future is dropped, await
/// this to ensure the callback worker released its listener before rebinding.
pub fn callback_closed(&self) -> impl Future<Output = ()> + Send + 'static {
self.flow.callback_closed()
}
pub async fn wait(self) -> Result<EnterpriseOAuthCredentials> {
let stored = self
.flow
@@ -187,22 +193,17 @@ impl EnterpriseOAuthCredentials {
pub async fn perform_enterprise_oauth_login_return_url(
request: EnterpriseOAuthLoginRequest<'_>,
) -> Result<EnterpriseOAuthLoginHandle> {
// Capture before discovery or browser setup, then release the lock while the user signs in.
let generation = {
let observed_generation = {
let guard = EnterpriseOAuthCredentialGuard::acquire(
request.credential_name,
request.issuer,
request.keyring_backend_kind,
)
.await?;
match guard.generation_file.current() {
Ok(Some(generation)) => generation,
Ok(None) => guard
.generation_file
.replace()
.map_err(|_| anyhow!("failed to initialize enterprise login generation"))?,
Err(_) => bail!("failed to read enterprise login generation"),
}
guard
.generation_file
.current()
.map_err(|_| anyhow!("failed to read enterprise login generation"))?
};
let flow = OauthLoginFlow::new(
request.credential_name,
@@ -232,6 +233,37 @@ pub async fn perform_enterprise_oauth_login_return_url(
.with_subscriber(tracing::subscriber::NoSubscriber::default())
.await
.map_err(|_| anyhow!("failed to start enterprise IdP authorization"))?;
// Keep the current attempt valid until replacement setup succeeds. A logout
// or another completed setup during this work must still invalidate this start.
let generation = async {
let guard = EnterpriseOAuthCredentialGuard::acquire(
request.credential_name,
request.issuer,
request.keyring_backend_kind,
)
.await?;
let current_generation = guard
.generation_file
.current()
.map_err(|_| anyhow!("failed to read enterprise login generation"))?;
if current_generation != observed_generation {
bail!("enterprise login changed during replacement setup");
}
guard
.generation_file
.replace()
.map_err(|_| anyhow!("failed to start enterprise login generation"))
}
.await;
let generation = match generation {
Ok(generation) => generation,
Err(error) => {
let callback_closed = flow.callback_closed();
drop(flow);
callback_closed.await;
return Err(error);
}
};
Ok(EnterpriseOAuthLoginHandle {
flow,
keyring_backend: request.keyring_backend_kind,
@@ -19,7 +19,6 @@ use keyring::mock::MockCredential;
use oauth2::TokenResponse;
use pretty_assertions::assert_eq;
use serde_json::json;
use sha2::Digest;
use tokio::io::AsyncReadExt;
use tokio::io::AsyncWriteExt;
use tracing_test::traced_test;
@@ -239,12 +238,6 @@ async fn enterprise_public_api_storage_and_privacy() -> Result<()> {
let server = MockServer::start().await;
let issuer = format!("{}/idp", server.uri());
metadata(&server, &issuer).await;
assert!(
login(&issuer, Some(&format!("{}/callback", server.uri())))
.await
.is_err(),
"an occupied registered callback port must not silently move"
);
let assertion = format!(
"{}.{}.signature",
URL_SAFE_NO_PAD.encode(r#"{"alg":"ES256"}"#),
@@ -284,31 +277,31 @@ async fn enterprise_public_api_storage_and_privacy() -> Result<()> {
})
.collect::<Vec<_>>();
assert_eq!(token_requests.len(), login_index + 1);
let token_form = url::form_urlencoded::parse(&token_requests[login_index].body)
let token_request = &token_requests[login_index];
assert!(!token_request.headers.contains_key("authorization"));
let mut token_form = url::form_urlencoded::parse(&token_request.body)
.into_owned()
.collect::<HashMap<_, _>>();
assert_eq!(
token_form.get("grant_type").map(String::as_str),
Some("authorization_code")
);
assert!(!token_form.contains_key("resource"));
assert_eq!(
token_form.get("redirect_uri"),
Some(
authorization_query
.get("redirect_uri")
.expect("authorization redirect URI")
)
let verifier = oauth2::PkceCodeVerifier::new(
token_form.remove("code_verifier").expect("PKCE verifier"),
);
assert_eq!(
authorization_query
.get("code_challenge_method")
.map(String::as_str),
Some("S256")
oauth2::PkceCodeChallenge::from_code_verifier_sha256(&verifier).as_str(),
authorization_query["code_challenge"],
);
assert_eq!(authorization_query["code_challenge_method"], "S256");
assert_eq!(
token_form,
HashMap::from([
("grant_type".to_string(), "authorization_code".to_string()),
("code".to_string(), SECRET.to_string()),
("client_id".to_string(), "enterprise-client".to_string()),
(
"redirect_uri".to_string(),
authorization_query["redirect_uri"].clone()
),
])
);
let verifier = token_form.get("code_verifier").expect("PKCE verifier");
let challenge = URL_SAFE_NO_PAD.encode(sha2::Sha256::digest(verifier.as_bytes()));
assert_eq!(authorization_query.get("code_challenge"), Some(&challenge));
assert!(
keyring
.values
@@ -381,13 +374,23 @@ async fn enterprise_public_api_storage_and_privacy() -> Result<()> {
.all(|value| value.get_secret().is_err())
);
// Rejected old attempts neither write nor delete a newer grant.
// Failed replacement setup leaves the existing attempt valid.
let retained = complete_login(&issuer).await?;
assert!(
login(&issuer, Some(&format!("{}/callback", server.uri())))
.await
.is_err(),
"an occupied registered callback port must not silently move"
);
retained.commit_if(|| async { Some(()) }).await?;
// A newer start fences a delayed callback even across Codex processes.
let old = complete_login(&issuer).await?;
complete_login(&issuer)
.await?
.commit_if(|| async { Some(()) })
.await?;
assert!(old.commit_if(|| async { None::<()> }).await.is_err());
assert!(old.commit_if(|| async { Some(()) }).await.is_err());
assert!(
keyring
.values
@@ -3,9 +3,16 @@
use std::fs;
use std::path::PathBuf;
use codex_exec_server::ExecServerError;
use codex_exec_server::HttpRequestParams;
use codex_exec_server::HttpRequestResponse;
use codex_exec_server::HttpResponseBodyStream;
use futures::FutureExt;
use futures::future::BoxFuture;
use pretty_assertions::assert_eq;
use sha2::Digest;
use sha2::Sha256;
use tokio::sync::Notify;
use super::*;
@@ -46,6 +53,48 @@ async fn logout_invalidates_pending_login_across_processes() -> Result<()> {
.mount(&server)
.await;
// A logout during provider setup must invalidate the in-flight start.
let paused_client = Arc::new(PausedSetupHttpClient {
inner: Arc::new(RouteAwareHttpClient::new(HttpClientFactory::new(
OutboundProxyPolicy::ReqwestDefault,
))),
should_pause: AtomicBool::new(true),
entered: Notify::new(),
release: Notify::new(),
});
let starting_client = Arc::clone(&paused_client);
let starting_issuer = issuer.clone();
let starting = tokio::spawn(async move {
perform_enterprise_oauth_login_return_url(EnterpriseOAuthLoginRequest {
credential_name: CREDENTIAL_NAME,
issuer: &starting_issuer,
client_id: "enterprise-client",
keyring_backend_kind: AuthKeyringBackendKind::Direct,
callback_port: None,
callback_url: None,
timeout_secs: Some(5),
http_client: starting_client,
redirect_mode: StreamableHttpRedirectMode::Legacy,
})
.await
});
tokio::time::timeout(
std::time::Duration::from_secs(5),
paused_client.entered.notified(),
)
.await?;
delete_enterprise_oauth_tokens(CREDENTIAL_NAME, &issuer, AuthKeyringBackendKind::Direct)
.await?;
paused_client.release.notify_one();
let error = tokio::time::timeout(std::time::Duration::from_secs(5), starting)
.await??
.err()
.expect("logout must invalidate a start during provider setup");
assert_eq!(
error.to_string(),
"enterprise login changed during replacement setup"
);
// The first logout has no grant to delete. The second deletes the fresh grant
// from the previous iteration. Both must invalidate pending callbacks and staged grants.
for _ in 0..2 {
@@ -127,6 +176,46 @@ async fn logout_invalidates_pending_login_across_processes() -> Result<()> {
Ok(())
}
struct PausedSetupHttpClient {
inner: Arc<dyn HttpClient>,
should_pause: AtomicBool,
entered: Notify,
release: Notify,
}
impl PausedSetupHttpClient {
async fn pause_once(&self) {
if self.should_pause.swap(false, Ordering::SeqCst) {
self.entered.notify_one();
self.release.notified().await;
}
}
}
impl HttpClient for PausedSetupHttpClient {
fn http_request(
&self,
params: HttpRequestParams,
) -> BoxFuture<'_, Result<HttpRequestResponse, ExecServerError>> {
async move {
self.pause_once().await;
self.inner.http_request(params).await
}
.boxed()
}
fn http_request_stream(
&self,
params: HttpRequestParams,
) -> BoxFuture<'_, Result<(HttpRequestResponse, HttpResponseBodyStream), ExecServerError>> {
async move {
self.pause_once().await;
self.inner.http_request_stream(params).await
}
.boxed()
}
}
fn stored(issuer: &str) -> Result<Option<StoredOAuthTokens>> {
crate::stored_oauth_credentials(
CREDENTIAL_NAME,
@@ -1,4 +1,4 @@
//! Persistent, non-secret logout generations for staged enterprise logins.
//! Persistent, non-secret generations for staged enterprise logins and logout.
//! All reads and writes require the credential lock; missing state never admits an old attempt.
use std::fs::File;
@@ -386,6 +386,10 @@ async fn oauth_registration_redirects_never_forward_resource_only_headers() -> R
#[tokio::test]
async fn same_origin_redirects_preserve_timeout_and_response_body_limits() -> Result<()> {
let http_client: Arc<dyn codex_exec_server::HttpClient> = Arc::new(RouteAwareHttpClient::new(
HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
));
crate::oauth::test_support::warm_http_client(http_client.as_ref()).await?;
for oversized_redirect_body in [false, true] {
let server = MockServer::start().await;
Mock::given(method("GET"))
@@ -417,9 +421,7 @@ async fn same_origin_redirects_preserve_timeout_and_response_body_limits() -> Re
.await;
let adapter = OAuthHttpClientAdapter::new(
Arc::new(RouteAwareHttpClient::new(HttpClientFactory::new(
OutboundProxyPolicy::ReqwestDefault,
))),
Arc::clone(&http_client),
build_default_headers(
Some(HashMap::from([(
"X-Api-Key".to_string(),
@@ -1,4 +1,5 @@
use std::collections::HashMap;
use std::future::Future;
use std::net::SocketAddr;
use std::string::String;
use std::sync::Arc;
@@ -282,9 +283,13 @@ fn spawn_callback_server(
server: Arc<Server>,
tx: oneshot::Sender<CallbackResult>,
expected_callback_path: String,
closed: tokio_util::sync::CancellationToken,
) {
tokio::task::spawn_blocking(move || {
while let Ok(request) = server.recv() {
// Drop the listener before publishing closure, including if this task unwinds.
let _closed = closed.drop_guard();
let listener = server;
while let Ok(request) = listener.recv() {
let path = request.url().to_string();
match parse_oauth_callback(&path, &expected_callback_path) {
CallbackOutcome::Success(OauthCallbackResult {
@@ -442,6 +447,7 @@ pub(crate) struct OauthLoginFlow {
authorization_server_issuer: Option<String>,
rx: oneshot::Receiver<CallbackResult>,
guard: CallbackServerGuard,
callback_closed: tokio_util::sync::CancellationToken,
server_name: String,
server_url: String,
store_mode: OAuthCredentialsStoreMode,
@@ -639,7 +645,25 @@ impl OauthLoginFlow {
// Port zero asks the OS for a free ephemeral port; the resolved
// redirect receives that port after the listener has been bound.
let bind_addr = SocketAddr::new(bind_ip, callback_port.unwrap_or(0));
let server = Arc::new(Server::http(bind_addr).map_err(|err| anyhow!(err))?);
let bind_deadline = tokio::time::Instant::now() + Duration::from_secs(/*secs*/ 1);
let server = loop {
match Server::http(bind_addr) {
Ok(server) => break Arc::new(server),
Err(err)
if is_enterprise_idp
&& callback_port.is_some()
&& err
.downcast_ref::<std::io::Error>()
.is_some_and(|err| err.kind() == std::io::ErrorKind::AddrInUse)
&& tokio::time::Instant::now() < bind_deadline =>
{
// tiny_http drops its accept loop asynchronously after Server::drop.
// Only retry the fixed listener bind, never OAuth requests or callbacks.
tokio::time::sleep(Duration::from_millis(/*millis*/ 10)).await;
}
Err(err) => return Err(anyhow!(err)),
}
};
let guard = CallbackServerGuard {
server: Arc::clone(&server),
};
@@ -701,8 +725,6 @@ impl OauthLoginFlow {
.await?
};
let callback_path = callback_path_from_redirect_uri(&redirect_uri)?;
let (tx, rx) = oneshot::channel();
spawn_callback_server(server, tx, callback_path);
let auth_url = append_query_param(
&oauth_state.get_authorization_url().await?,
"resource",
@@ -710,6 +732,11 @@ impl OauthLoginFlow {
);
let timeout_secs = timeout_secs.unwrap_or(DEFAULT_OAUTH_TIMEOUT_SECS).max(1);
let timeout = Duration::from_secs(timeout_secs as u64);
// No fallible work or suspension after spawning: every callback worker
// must be owned by a returned flow, including during constructor cancellation.
let (tx, rx) = oneshot::channel();
let callback_closed = tokio_util::sync::CancellationToken::new();
spawn_callback_server(server, tx, callback_path, callback_closed.clone());
Ok(Self {
auth_url,
@@ -718,6 +745,7 @@ impl OauthLoginFlow {
authorization_server_issuer,
rx,
guard,
callback_closed,
server_name: server_name.to_string(),
server_url: server_url.to_string(),
store_mode,
@@ -731,6 +759,10 @@ impl OauthLoginFlow {
self.auth_url.clone()
}
pub(crate) fn callback_closed(&self) -> impl Future<Output = ()> + Send + 'static {
self.callback_closed.clone().cancelled_owned()
}
async fn finish(self, emit_browser_url: bool) -> Result<()> {
let store_mode = self.store_mode;
let keyring_backend_kind = self.keyring_backend_kind;
@@ -2584,6 +2584,13 @@ class McpServerOauthLoginCompletedNotification(BaseModel):
populate_by_name=True,
)
error: str | None = None
login_id: Annotated[
str | None,
Field(
alias="loginId",
description="Identifies the explicit login attempt. Older servers omit this field.",
),
] = None
name: str
success: bool
thread_id: Annotated[str | None, Field(alias="threadId")] = None
@@ -2611,6 +2618,13 @@ class McpServerOauthLoginResponse(BaseModel):
populate_by_name=True,
)
authorization_url: Annotated[str, Field(alias="authorizationUrl")]
login_id: Annotated[
str | None,
Field(
alias="loginId",
description="Identifies this login attempt across the response and completion notification. Older servers omit this field; current servers always return it.",
),
] = None
class McpServerRefreshResponse(BaseModel):