mirror of
https://github.com/openai/codex.git
synced 2026-10-19 16:17:49 +00:00
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:
Generated
+2
@@ -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"
|
||||
},
|
||||
|
||||
+15
@@ -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": [
|
||||
|
||||
+15
@@ -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
@@ -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": [
|
||||
|
||||
BIN
Binary file not shown.
BIN
Binary file not shown.
Generated
+5
-1
@@ -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, };
|
||||
|
||||
+6
-1
@@ -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, };
|
||||
|
||||
@@ -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!({
|
||||
|
||||
@@ -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 }
|
||||
|
||||
@@ -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`.
|
||||
|
||||
@@ -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(¤t.config_layer_stack, ¤t.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,
|
||||
},
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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."),
|
||||
|
||||
@@ -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()?;
|
||||
|
||||
@@ -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()?)
|
||||
}
|
||||
@@ -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;
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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 [
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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 }
|
||||
|
||||
@@ -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):
|
||||
|
||||
Reference in New Issue
Block a user