mirror of
https://github.com/openai/codex.git
synced 2026-11-14 06:48:26 +00:00
feat: use run agent task auth for inference (#19051)
## Stack This is PR 3 of the simplified HAI single-run-task stack: - [#19047](https://github.com/openai/codex/pull/19047) Agent Identity assertion and task-registration primitives, including the shared run-task helper used by existing Agent Identity JWT auth. - [#19049](https://github.com/openai/codex/pull/19049) Disabled-by-default ChatGPT auth opt-in that provisions/reuses persisted Agent Identity runtime auth and its single run task. - [#19051](https://github.com/openai/codex/pull/19051) Run-scoped provider auth that uses one backend-owned task id for first-party inference and compaction requests. [#19054](https://github.com/openai/codex/pull/19054) collapsed out of the active stack because the simplified design no longer needs a separate background/control-plane task helper. ## Summary This PR moves Agent Identity usage into provider auth resolution. That keeps `AgentAssertion` auth tied to first-party OpenAI provider requests instead of applying a late session-wide override that could affect local, custom, Bedrock, API-key, or external-bearer providers. What changed: - adds a small `ProviderAuthScope` struct carrying the run auth policy and session source needed by provider-scoped auth resolution - lets `Session` opt the existing `ModelClient` into `ChatGptAuth` policy when `use_agent_identity` is enabled, without adding a second model-client constructor - resolves Agent Identity only for first-party OpenAI provider auth paths - uses the persisted run task id from the `AgentIdentityAuth` record to build `AgentAssertion` auth for Responses requests - routes shared request setup through scoped provider auth so unary compact requests use the same run-task assertion path as inference turns - keeps local/custom/Bedrock/env-key/external-bearer provider auth unchanged - lets missing run-task state surface through the existing model-request error path instead of silently falling back to bearer auth This PR intentionally does not create thread-scoped, target-scoped, or background-scoped task identities. The run task is the only task Codex registers in this POC shape. ## Testing - `just test -p codex-model-provider` - `just test -p codex-core client::tests::provider_auth_scope_uses` - `just test -p codex-core remote_compact_uses_agent_identity_assertion`
This commit is contained in:
Adrian authored and GitHub committed
1 parent
f66d793a2d
commit
51864b0b4b
22 files changed
+1020
-20
No files matched your search
@@ -1,6 +1,8 @@
|
||||
use super::AuthRequestTelemetryContext;
|
||||
use super::CompactConversationRequestSettings;
|
||||
use super::ModelClient;
|
||||
use super::PendingUnauthorizedRetry;
|
||||
use super::Prompt;
|
||||
use super::UnauthorizedRecoveryExecution;
|
||||
use super::X_CODEX_INSTALLATION_ID_HEADER;
|
||||
use super::X_CODEX_PARENT_THREAD_ID_HEADER;
|
||||
@@ -13,11 +15,15 @@ use crate::GenerateAttestationFuture;
|
||||
use crate::responses_metadata::CodexResponsesMetadata;
|
||||
use crate::test_support::TestCodexResponsesRequestKind;
|
||||
use crate::test_support::responses_metadata as test_responses_metadata;
|
||||
use codex_api::AgentIdentityTelemetry;
|
||||
use codex_api::ApiError;
|
||||
use codex_api::ResponseEvent;
|
||||
use codex_api::TransportError;
|
||||
use codex_login::AuthCredentialsStoreMode;
|
||||
use codex_login::AuthKeyringBackendKind;
|
||||
use codex_login::AuthManager;
|
||||
use codex_login::CodexAuth;
|
||||
use codex_login::auth::AgentIdentityAuthPolicy;
|
||||
use codex_model_provider::BearerAuthProvider;
|
||||
use codex_model_provider::SharedModelProvider;
|
||||
use codex_model_provider::create_model_provider;
|
||||
@@ -28,6 +34,7 @@ use codex_model_provider_info::create_oss_provider_with_base_url;
|
||||
use codex_otel::SessionTelemetry;
|
||||
use codex_protocol::ThreadId;
|
||||
use codex_protocol::auth::AuthMode;
|
||||
use codex_protocol::models::BaseInstructions;
|
||||
use codex_protocol::models::ContentItem;
|
||||
use codex_protocol::models::ResponseItem;
|
||||
use codex_protocol::openai_models::ModelInfo;
|
||||
@@ -35,6 +42,7 @@ use codex_protocol::openai_models::ReasoningEffort;
|
||||
use codex_protocol::protocol::InternalSessionSource;
|
||||
use codex_protocol::protocol::SessionSource;
|
||||
use codex_protocol::protocol::SubAgentSource;
|
||||
use codex_rollout_trace::CompactionTraceContext;
|
||||
use codex_rollout_trace::ExecutionStatus;
|
||||
use codex_rollout_trace::InferenceTraceAttempt;
|
||||
use codex_rollout_trace::InferenceTraceContext;
|
||||
@@ -65,14 +73,27 @@ use tracing_subscriber::layer::Context as LayerContext;
|
||||
use tracing_subscriber::layer::SubscriberExt;
|
||||
use tracing_subscriber::registry::LookupSpan;
|
||||
use tracing_subscriber::util::SubscriberInitExt;
|
||||
use wiremock::Mock;
|
||||
use wiremock::MockServer;
|
||||
use wiremock::ResponseTemplate;
|
||||
use wiremock::matchers::method;
|
||||
use wiremock::matchers::path;
|
||||
|
||||
const TEST_CHATGPT_ID_TOKEN: &str = "eyJhbGciOiJub25lIiwidHlwIjoiSldUIn0.eyJlbWFpbCI6InVzZXJAZXhhbXBsZS5jb20iLCJlbWFpbF92ZXJpZmllZCI6dHJ1ZSwiaHR0cHM6Ly9hcGkub3BlbmFpLmNvbS9hdXRoIjp7ImNoYXRncHRfdXNlcl9pZCI6InVzZXItMTIzNDUiLCJ1c2VyX2lkIjoidXNlci0xMjM0NSIsImNoYXRncHRfcGxhbl90eXBlIjoicHJvIiwiY2hhdGdwdF9hY2NvdW50X2lkIjoiYWNjb3VudC0xMjMifX0.c2ln";
|
||||
const TEST_INSTALLATION_ID: &str = "11111111-1111-4111-8111-111111111111";
|
||||
|
||||
fn test_model_client(session_source: SessionSource) -> ModelClient {
|
||||
test_model_client_with_thread_id(ThreadId::new(), session_source)
|
||||
}
|
||||
|
||||
fn test_model_client_with_thread_id(
|
||||
thread_id: ThreadId,
|
||||
session_source: SessionSource,
|
||||
) -> ModelClient {
|
||||
let provider = create_oss_provider_with_base_url("https://example.com/v1", WireApi::Responses);
|
||||
let thread_id = ThreadId::new();
|
||||
ModelClient::new(
|
||||
/*auth_manager*/ None,
|
||||
AgentIdentityAuthPolicy::JwtOnly,
|
||||
thread_id,
|
||||
provider,
|
||||
session_source,
|
||||
@@ -86,6 +107,115 @@ fn test_model_client(session_source: SessionSource) -> ModelClient {
|
||||
)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn compact_uses_bearer_after_agent_identity_session_fallback() -> anyhow::Result<()> {
|
||||
let server = MockServer::start().await;
|
||||
let registration_count = Arc::new(AtomicUsize::new(0));
|
||||
let response_count = Arc::clone(®istration_count);
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/v1/agent/register"))
|
||||
.respond_with(move |_request: &wiremock::Request| {
|
||||
response_count.fetch_add(1, Ordering::SeqCst);
|
||||
ResponseTemplate::new(/*status*/ 503)
|
||||
})
|
||||
.mount(&server)
|
||||
.await;
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/v1/responses/compact"))
|
||||
.respond_with(ResponseTemplate::new(/*status*/ 200).set_body_json(json!({
|
||||
"output": []
|
||||
})))
|
||||
.expect(/*requests*/ 1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let codex_home = TempDir::new()?;
|
||||
let auth_manager = chatgpt_auth_manager(&codex_home, server.uri()).await;
|
||||
let mut provider = ModelProviderInfo::create_openai_provider(/*base_url*/ None);
|
||||
provider.base_url = Some(format!("{}/v1", server.uri()));
|
||||
provider.supports_websockets = false;
|
||||
let thread_id = ThreadId::new();
|
||||
let client = ModelClient::new(
|
||||
Some(auth_manager),
|
||||
AgentIdentityAuthPolicy::ChatGptAuth,
|
||||
thread_id,
|
||||
provider,
|
||||
SessionSource::Cli,
|
||||
"test_originator".to_string(),
|
||||
/*model_verbosity*/ None,
|
||||
/*enable_request_compression*/ false,
|
||||
/*include_timing_metrics*/ false,
|
||||
/*beta_features_header*/ None,
|
||||
/*item_ids_enabled*/ false,
|
||||
/*attestation_provider*/ None,
|
||||
);
|
||||
let prompt = Prompt {
|
||||
input: vec![ResponseItem::Message {
|
||||
id: None,
|
||||
role: "user".to_string(),
|
||||
content: vec![ContentItem::InputText {
|
||||
text: "please compact".to_string(),
|
||||
}],
|
||||
phase: None,
|
||||
internal_chat_message_metadata_passthrough: None,
|
||||
}],
|
||||
base_instructions: BaseInstructions {
|
||||
text: "base instructions".to_string(),
|
||||
},
|
||||
..Default::default()
|
||||
};
|
||||
let responses_metadata = test_responses_metadata_for_client(
|
||||
&client,
|
||||
/*turn_id*/ None,
|
||||
format!("{}:0", client.state.thread_id),
|
||||
/*parent_thread_id*/ None,
|
||||
TestCodexResponsesRequestKind::Turn,
|
||||
);
|
||||
|
||||
let output = client
|
||||
.compact_conversation_history(
|
||||
&prompt,
|
||||
&test_model_info(),
|
||||
/*turn_state*/ None,
|
||||
CompactConversationRequestSettings {
|
||||
effort: None,
|
||||
summary: codex_protocol::config_types::ReasoningSummary::None,
|
||||
service_tier: None,
|
||||
},
|
||||
&test_session_telemetry(),
|
||||
&CompactionTraceContext::disabled(),
|
||||
&responses_metadata,
|
||||
)
|
||||
.await?;
|
||||
|
||||
assert!(output.is_empty());
|
||||
assert_eq!(registration_count.load(Ordering::SeqCst), 3);
|
||||
let requests = server
|
||||
.received_requests()
|
||||
.await
|
||||
.expect("server should record requests");
|
||||
let compact_request = requests
|
||||
.iter()
|
||||
.find(|request| request.url.path() == "/v1/responses/compact")
|
||||
.expect("compact request should be captured");
|
||||
assert_eq!(
|
||||
compact_request
|
||||
.headers
|
||||
.get(http::header::AUTHORIZATION)
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("Bearer test-access-token")
|
||||
);
|
||||
assert_eq!(
|
||||
compact_request
|
||||
.headers
|
||||
.get("ChatGPT-Account-ID")
|
||||
.and_then(|value| value.to_str().ok()),
|
||||
Some("account-123")
|
||||
);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn test_model_provider() -> SharedModelProvider {
|
||||
test_model_client(SessionSource::Cli).state.provider.clone()
|
||||
}
|
||||
@@ -169,6 +299,45 @@ fn ultra_reasoning_uses_max_for_requests() {
|
||||
);
|
||||
}
|
||||
|
||||
fn write_chatgpt_auth_json(codex_home: &std::path::Path) {
|
||||
let auth_json = json!({
|
||||
"tokens": {
|
||||
"id_token": TEST_CHATGPT_ID_TOKEN,
|
||||
"access_token": "test-access-token",
|
||||
"refresh_token": "test-refresh-token",
|
||||
"account_id": "account-123"
|
||||
},
|
||||
"last_refresh": "2099-01-01T00:00:00Z"
|
||||
});
|
||||
std::fs::write(
|
||||
codex_home.join("auth.json"),
|
||||
serde_json::to_string_pretty(&auth_json).expect("serialize auth.json"),
|
||||
)
|
||||
.expect("write auth.json");
|
||||
}
|
||||
|
||||
async fn chatgpt_auth_manager(
|
||||
codex_home: &TempDir,
|
||||
agent_identity_authapi_base_url: String,
|
||||
) -> Arc<AuthManager> {
|
||||
write_chatgpt_auth_json(codex_home.path());
|
||||
let auth_manager = AuthManager::shared(
|
||||
codex_home.path().to_path_buf(),
|
||||
/*enable_codex_api_key_env*/ false,
|
||||
AuthCredentialsStoreMode::File,
|
||||
/*forced_chatgpt_workspace_id*/ None,
|
||||
/*chatgpt_base_url*/ None,
|
||||
AuthKeyringBackendKind::default(),
|
||||
/*auth_route_config*/ None,
|
||||
)
|
||||
.await;
|
||||
let auth = auth_manager.auth().await.expect("auth should load");
|
||||
AuthManager::from_auth_for_testing_with_agent_identity_authapi_base_url(
|
||||
auth,
|
||||
agent_identity_authapi_base_url,
|
||||
)
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct TagCollectorVisitor {
|
||||
tags: BTreeMap<String, String>,
|
||||
@@ -572,6 +741,7 @@ fn auth_request_telemetry_context_tracks_attached_auth_and_retry_phase() {
|
||||
let auth_context = AuthRequestTelemetryContext::new(
|
||||
Some(AuthMode::Chatgpt),
|
||||
&BearerAuthProvider::for_test(Some("access-token"), Some("workspace-123")),
|
||||
/*agent_identity_telemetry*/ None,
|
||||
PendingUnauthorizedRetry::from_recovery(UnauthorizedRecoveryExecution {
|
||||
mode: "managed",
|
||||
phase: "refresh_token",
|
||||
@@ -586,6 +756,27 @@ fn auth_request_telemetry_context_tracks_attached_auth_and_retry_phase() {
|
||||
assert_eq!(auth_context.recovery_phase, Some("refresh_token"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn auth_request_telemetry_context_tracks_agent_identity_ids() {
|
||||
let auth_context = AuthRequestTelemetryContext::new(
|
||||
Some(AuthMode::Chatgpt),
|
||||
&BearerAuthProvider::for_test(/*token*/ None, /*account_id*/ None),
|
||||
Some(AgentIdentityTelemetry {
|
||||
agent_id: "agent-runtime-context".to_string(),
|
||||
task_id: "task-run-context".to_string(),
|
||||
}),
|
||||
PendingUnauthorizedRetry::default(),
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
auth_context.agent_identity_telemetry(),
|
||||
Some(&AgentIdentityTelemetry {
|
||||
agent_id: "agent-runtime-context".to_string(),
|
||||
task_id: "task-run-context".to_string(),
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
fn model_client_with_counting_attestation(
|
||||
include_attestation: bool,
|
||||
) -> (ModelClient, Arc<AtomicUsize>) {
|
||||
@@ -623,6 +814,7 @@ fn model_client_with_counting_attestation(
|
||||
};
|
||||
let model_client = ModelClient::new(
|
||||
auth_manager,
|
||||
AgentIdentityAuthPolicy::JwtOnly,
|
||||
ThreadId::new(),
|
||||
provider,
|
||||
SessionSource::Exec,
|
||||
|
||||
Reference in new issue
Block a user