github/codex
4
0 Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-13 14:46:14 +00:00
Code Issues Packages Projects Releases Wiki Activity

[codex] Use AgentAssertion downstream behind use_agent_identity (#17980)

Browse Source 
## Summary

This is the AgentAssertion downstream slice for feature-gated agent
identity support, replacing the oversized AgentAssertion slice from PR
#17807.

It isolates task-scoped downstream AgentAssertion wiring on top of the
merged PR3.1 work without re-carrying the earlier agent registration,
task registration, or task-state history.

This PR includes the task-scoped bug-fix call sites from the review:
generic file upload auth, MCP OpenAI file upload auth, and ARC monitor
auth. Broader user/control-plane calls move to PR4.1 and PR4.2.

## Stack

- PR1: https://github.com/openai/codex/pull/17385 - add
`features.use_agent_identity`
- PR2: https://github.com/openai/codex/pull/17386 - register agent
identities when enabled
- PR3: https://github.com/openai/codex/pull/17387 - register agent tasks
when enabled
- PR3.1: https://github.com/openai/codex/pull/17978 - persist and
prewarm registered tasks per thread
- PR4: this PR - use task-scoped `AgentAssertion` downstream when
enabled
- PR4.1: https://github.com/openai/codex/pull/18094 - introduce
AuthManager-owned background/control-plane `AgentAssertion` auth
- PR4.2: https://github.com/openai/codex/pull/18260 - use background
task auth for additional backend/control-plane calls

## What Changed

- add AgentAssertion envelope generation in `codex-core`
- route downstream HTTP and websocket auth through AgentAssertion when
an agent task is present
- extend the model-provider auth provider so non-bearer authorization
schemes can be passed through cleanly
- make generic file uploads attach the full authorization header value
- make MCP OpenAI file uploads use the cached thread agent task
assertion when present
- make ARC monitor calls use the cached thread agent task assertion when
present

## Why

The original PR had drifted ancestry and showed a much larger diff than
the semantic change actually required. Restacking it onto PR3.1 keeps
the reviewable surface down to the downstream assertion slice.

## Validation

- `just fmt`
- `cargo check -p codex-core -p codex-login -p codex-analytics -p
codex-app-server -p codex-cloud-requirements -p codex-cloud-tasks -p
codex-models-manager -p codex-chatgpt -p codex-model-provider -p
codex-mcp -p codex-core-skills`
- `cargo test -p codex-model-provider bearer_auth_provider`
- `cargo test -p codex-core agent_assertion`
- `cargo test -p codex-app-server remote_control`
- `cargo test -p codex-cloud-requirements fetch_cloud_requirements`
- `cargo test -p codex-models-manager manager::tests`
- `cargo test -p codex-chatgpt`
- `cargo test -p codex-cloud-tasks`
- `cargo test -p codex-login agent_identity`
- `just fix -p codex-core -p codex-login -p codex-analytics -p
codex-app-server -p codex-cloud-requirements -p codex-cloud-tasks -p
codex-models-manager -p codex-chatgpt -p codex-model-provider -p
codex-mcp -p codex-core-skills`
- `just fix -p codex-app-server`
- `git diff --check`
This commit is contained in:
Adrian authored and GitHub committed 2026-04-19 23:16:43 -07:00
1 parent 3c75f9b4dd
commit b44d2851cf
19 files changed
+1206 -65
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

codex-rs/core/src/mcp_openai_file.rs
+197 -14
View File
@@ -12,8 +12,10 @@
use crate::session::session::Session;
use crate::session::turn_context::TurnContext;
use codex_api::AuthProvider;
use codex_api::upload_local_file;
use codex_login::CodexAuth;
use codex_model_provider::AuthorizationHeaderAuthProvider;
use codex_model_provider::BearerAuthProvider;
use serde_json::Value as JsonValue;
@@ -40,9 +42,14 @@ pub(crate) async fn rewrite_mcp_tool_arguments_for_openai_files(
let Some(value) = arguments.get(field_name) else {
continue;
};
let Some(uploaded_value) =
rewrite_argument_value_for_openai_files(turn_context, auth.as_ref(), field_name, value)
.await?
let Some(uploaded_value) = rewrite_argument_value_for_openai_files(
sess,
turn_context,
auth.as_ref(),
field_name,
value,
)
.await?
else {
continue;
};
@@ -57,6 +64,7 @@ pub(crate) async fn rewrite_mcp_tool_arguments_for_openai_files(
}
async fn rewrite_argument_value_for_openai_files(
sess: &Session,
turn_context: &TurnContext,
auth: Option<&CodexAuth>,
field_name: &str,
@@ -65,6 +73,7 @@ async fn rewrite_argument_value_for_openai_files(
match value {
JsonValue::String(path_or_file_ref) => {
let rewritten = build_uploaded_local_argument_value(
sess,
turn_context,
auth,
field_name,
@@ -81,6 +90,7 @@ async fn rewrite_argument_value_for_openai_files(
return Ok(None);
};
let rewritten = build_uploaded_local_argument_value(
sess,
turn_context,
auth,
field_name,
@@ -97,6 +107,7 @@ async fn rewrite_argument_value_for_openai_files(
}
async fn build_uploaded_local_argument_value(
sess: &Session,
turn_context: &TurnContext,
auth: Option<&CodexAuth>,
field_name: &str,
@@ -109,17 +120,32 @@ async fn build_uploaded_local_argument_value(
"ChatGPT auth is required to upload local files for Codex Apps tools".to_string(),
);
};
let token_data = auth
.get_token_data()
.map_err(|error| format!("failed to read ChatGPT auth for file upload: {error}"))?;
let upload_auth = BearerAuthProvider {
token: Some(token_data.access_token),
account_id: token_data.account_id,
is_fedramp_account: auth.is_fedramp_account(),
let upload_auth: Box<dyn AuthProvider> = if let Some(authorization_header_value) = sess
.authorization_header_for_current_agent_task()
.await
.map_err(|error| format!("failed to build agent assertion authorization: {error}"))?
{
let mut auth_provider = AuthorizationHeaderAuthProvider::new(
Some(authorization_header_value),
/*account_id*/ None,
);
if auth.is_fedramp_account() {
auth_provider = auth_provider.with_fedramp_routing_header();
}
Box::new(auth_provider)
} else {
let token_data = auth
.get_token_data()
.map_err(|error| format!("failed to read ChatGPT auth for file upload: {error}"))?;
Box::new(BearerAuthProvider {
token: Some(token_data.access_token),
account_id: token_data.account_id,
is_fedramp_account: auth.is_fedramp_account(),
})
};
let uploaded = upload_local_file(
turn_context.config.chatgpt_base_url.trim_end_matches('/'),
&upload_auth,
upload_auth.as_ref(),
&resolved_path,
)
.await
@@ -142,12 +168,82 @@ async fn build_uploaded_local_argument_value(
#[cfg(test)]
mod tests {
use super::*;
use crate::agent_identity::AgentIdentityManager;
use crate::agent_identity::RegisteredAgentTask;
use crate::session::tests::make_session_and_context;
use chrono::Utc;
use codex_login::AuthCredentialsStoreMode;
use codex_login::AuthDotJson;
use codex_login::AuthManager;
use codex_login::save_auth;
use codex_login::token_data::IdTokenInfo;
use codex_login::token_data::TokenData;
use codex_protocol::protocol::SessionSource;
use codex_utils_absolute_path::AbsolutePathBuf;
use pretty_assertions::assert_eq;
use std::sync::Arc;
use tempfile::tempdir;
const TEST_ID_TOKEN: &str = concat!(
"eyJhbGciOiJub25lIiwidHlwIjoiSldUIn0.",
"eyJodHRwczovL2FwaS5vcGVuYWkuY29tL2F1dGgiOnsiY2hhdGdwdF91c2VyX2lk",
"IjpudWxsLCJjaGF0Z3B0X2FjY291bnRfaWQiOiJhY2NvdW50X2lkIn19.",
"c2ln",
);
async fn install_cached_agent_task_auth(
session: &mut Session,
turn_context: &mut TurnContext,
chatgpt_base_url: String,
) {
let auth_dir = tempdir().expect("temp auth dir");
let auth_json = AuthDotJson {
auth_mode: Some(codex_app_server_protocol::AuthMode::Chatgpt),
openai_api_key: None,
tokens: Some(TokenData {
id_token: IdTokenInfo {
email: None,
chatgpt_plan_type: None,
chatgpt_user_id: None,
chatgpt_account_id: Some("account_id".to_string()),
chatgpt_account_is_fedramp: false,
raw_jwt: TEST_ID_TOKEN.to_string(),
},
access_token: "Access Token".to_string(),
refresh_token: "test".to_string(),
account_id: Some("account_id".to_string()),
}),
last_refresh: Some(Utc::now()),
agent_identity: None,
};
save_auth(auth_dir.path(), &auth_json, AuthCredentialsStoreMode::File)
.expect("save test auth");
let auth = CodexAuth::from_auth_storage(auth_dir.path(), AuthCredentialsStoreMode::File)
.expect("load test auth")
.expect("test auth");
let auth_manager = AuthManager::from_auth_for_testing(auth);
let agent_identity_manager = Arc::new(AgentIdentityManager::new_for_tests(
Arc::clone(&auth_manager),
/*feature_enabled*/ true,
chatgpt_base_url,
SessionSource::Exec,
));
let stored_identity = agent_identity_manager
.seed_generated_identity_for_tests("agent-123")
.await
.expect("seed test identity");
session.services.auth_manager = Arc::clone(&auth_manager);
session.services.agent_identity_manager = agent_identity_manager;
turn_context.auth_manager = Some(auth_manager);
session
.cache_agent_task_for_tests(RegisteredAgentTask {
agent_runtime_id: stored_identity.agent_runtime_id,
task_id: "task-123".to_string(),
registered_at: "2026-04-15T00:00:00Z".to_string(),
})
.await;
}
#[tokio::test]
async fn openai_file_argument_rewrite_requires_declared_file_params() {
let (session, turn_context) = make_session_and_context().await;
@@ -212,7 +308,7 @@ mod tests {
.mount(&server)
.await;
let (_, mut turn_context) = make_session_and_context().await;
let (session, mut turn_context) = make_session_and_context().await;
let auth = CodexAuth::create_dummy_chatgpt_auth_for_testing();
let dir = tempdir().expect("temp dir");
let local_path = dir.path().join("file_report.csv");
@@ -226,6 +322,7 @@ mod tests {
turn_context.config = Arc::new(config);
let rewritten = build_uploaded_local_argument_value(
&session,
&turn_context,
Some(&auth),
"file",
@@ -293,7 +390,7 @@ mod tests {
.mount(&server)
.await;
let (_, mut turn_context) = make_session_and_context().await;
let (session, mut turn_context) = make_session_and_context().await;
let auth = CodexAuth::create_dummy_chatgpt_auth_for_testing();
let dir = tempdir().expect("temp dir");
let local_path = dir.path().join("file_report.csv");
@@ -306,6 +403,7 @@ mod tests {
config.chatgpt_base_url = format!("{}/backend-api", server.uri());
turn_context.config = Arc::new(config);
let rewritten = rewrite_argument_value_for_openai_files(
&session,
&turn_context,
Some(&auth),
"file",
@@ -405,7 +503,7 @@ mod tests {
.mount(&server)
.await;
let (_, mut turn_context) = make_session_and_context().await;
let (session, mut turn_context) = make_session_and_context().await;
let auth = CodexAuth::create_dummy_chatgpt_auth_for_testing();
let dir = tempdir().expect("temp dir");
tokio::fs::write(dir.path().join("one.csv"), b"one")
@@ -420,6 +518,7 @@ mod tests {
config.chatgpt_base_url = format!("{}/backend-api", server.uri());
turn_context.config = Arc::new(config);
let rewritten = rewrite_argument_value_for_openai_files(
&session,
&turn_context,
Some(&auth),
"files",
@@ -471,4 +570,88 @@ mod tests {
assert!(error.contains("failed to upload"));
assert!(error.contains("file"));
}
#[tokio::test]
async fn build_uploaded_local_argument_value_uses_agent_assertion_for_cached_task() {
use wiremock::Mock;
use wiremock::MockServer;
use wiremock::ResponseTemplate;
use wiremock::matchers::body_json;
use wiremock::matchers::header_regex;
use wiremock::matchers::method;
use wiremock::matchers::path;
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/backend-api/files"))
.and(header_regex("authorization", r"^AgentAssertion .+"))
.and(body_json(serde_json::json!({
"file_name": "file_report.csv",
"file_size": 5,
"use_case": "codex",
})))
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
"file_id": "file_123",
"upload_url": format!("{}/upload/file_123", server.uri()),
})))
.expect(1)
.mount(&server)
.await;
Mock::given(method("PUT"))
.and(path("/upload/file_123"))
.respond_with(ResponseTemplate::new(200))
.expect(1)
.mount(&server)
.await;
Mock::given(method("POST"))
.and(path("/backend-api/files/file_123/uploaded"))
.and(header_regex("authorization", r"^AgentAssertion .+"))
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
"status": "success",
"download_url": format!("{}/download/file_123", server.uri()),
"file_name": "file_report.csv",
"mime_type": "text/csv",
"file_size_bytes": 5,
})))
.expect(1)
.mount(&server)
.await;
let (mut session, mut turn_context) = make_session_and_context().await;
let auth = CodexAuth::create_dummy_chatgpt_auth_for_testing();
let dir = tempdir().expect("temp dir");
let local_path = dir.path().join("file_report.csv");
tokio::fs::write(&local_path, b"hello")
.await
.expect("write local file");
turn_context.cwd = AbsolutePathBuf::try_from(dir.path()).expect("absolute path");
let mut config = (*turn_context.config).clone();
config.chatgpt_base_url = format!("{}/backend-api", server.uri());
turn_context.config = Arc::new(config);
install_cached_agent_task_auth(&mut session, &mut turn_context, server.uri()).await;
let rewritten = build_uploaded_local_argument_value(
&session,
&turn_context,
Some(&auth),
"file",
/*index*/ None,
"file_report.csv",
)
.await
.expect("rewrite should upload the local file");
assert_eq!(
rewritten,
serde_json::json!({
"download_url": format!("{}/download/file_123", server.uri()),
"file_id": "file_123",
"mime_type": "text/csv",
"file_name": "file_report.csv",
"uri": "sediment://file_123",
"file_size_bytes": 5,
})
);
}
}