Skip to content

Register agent tasks behind use_agent_identity #17387

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
adrian-openai merged 18 commits into from
Apr 16, 2026
Merged

Register agent tasks behind use_agent_identity #17387

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
18 commits
Select commit Hold shift + click to select a range
a338dfb
Register agent tasks behind use_agent_identity
adrian-openai Apr 9, 2026
ceb0593
Use canonical agent task registration URL
adrian-openai Mar 24, 2026
a597680
Register agent tasks behind use_agent_identity
adrian-openai Mar 24, 2026
7b8fcfb
Align agent task registration auth
adrian-openai Apr 10, 2026
663ccd2
Fix agent identity task rebase
adrian-openai Apr 11, 2026
ef856d0
Update Bazel lockfile for agent identity task deps
adrian-openai Apr 11, 2026
e9ec3e1
Revalidate cached agent identity tasks
adrian-openai Apr 11, 2026
f702951
Use ChatGPT base URL for agent task registration
adrian-openai Apr 13, 2026
1bbaf34
Fix agent task registration clippy
adrian-openai Apr 13, 2026
14730d1
Fix agent task biscuit target method
adrian-openai Apr 15, 2026
1fff58d
Annotate agent task test binding argument
adrian-openai Apr 15, 2026
43214ec
Make agent task registration failure retryable
adrian-openai Apr 15, 2026
dfc5e05
codex: fix CI timeout flakes on PR3
adrian-openai Apr 16, 2026
ac8ef69
codex: fix task registration auth test rebase
adrian-openai Apr 16, 2026
df29a13
codex: make realtime startup context test deterministic
adrian-openai Apr 16, 2026
f1ad347
codex: relax realtime shell tool test timeout
adrian-openai Apr 16, 2026
024dfb2
codex: seed realtime recent-thread fixture with user event
adrian-openai Apr 16, 2026
cb42e36
codex: stabilize auth refresh status tests
adrian-openai Apr 16, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions MODULE.bazel.lock
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

43 changes: 43 additions & 0 deletions codex-rs/Cargo.lock
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions codex-rs/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -224,6 +224,7 @@ color-eyre = "0.6.3"
constant_time_eq = "0.3.1"
crossbeam-channel = "0.5.15"
crossterm = "0.28.1"
crypto_box = { version = "0.9.1", features = ["seal"] }
csv = "1.3.1"
ctor = "0.6.3"
deno_core_icudata = "0.77.0"
Expand Down
8 changes: 6 additions & 2 deletions codex-rs/app-server/tests/suite/auth.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -351,14 +351,16 @@ async fn get_auth_status_omits_token_after_proactive_refresh_failure() -> Result
)?;

let server = MockServer::start().await;
// App-server startup may proactively read stale auth before this test sends
// getAuthStatus; require the refresh path without depending on that race.
Mock::given(method("POST"))
.and(path("/oauth/token"))
.respond_with(ResponseTemplate::new(401).set_body_json(serde_json::json!({
"error": {
"code": "refresh_token_reused"
}
})))
.expect(2)
.expect(1..=2)
.mount(&server)
.await;

Expand Down Expand Up @@ -418,14 +420,16 @@ async fn get_auth_status_returns_token_after_proactive_refresh_recovery() -> Res
)?;

let server = MockServer::start().await;
// App-server startup may proactively read stale auth before this test sends
// getAuthStatus; require the refresh path without depending on that race.
Mock::given(method("POST"))
.and(path("/oauth/token"))
.respond_with(ResponseTemplate::new(401).set_body_json(serde_json::json!({
"error": {
"code": "refresh_token_reused"
}
})))
.expect(2)
.expect(1..=2)
.mount(&server)
.await;

Expand Down
3 changes: 2 additions & 1 deletion codex-rs/app-server/tests/suite/v2/plugin_list.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,8 @@ use wiremock::matchers::method;
use wiremock::matchers::path;
use wiremock::matchers::query_param;

const DEFAULT_TIMEOUT: Duration = Duration::from_secs(10);
// These tests start full app-server processes and can also run plugin startup warmers.
const DEFAULT_TIMEOUT: Duration = Duration::from_secs(60);
const TEST_CURATED_PLUGIN_SHA: &str = "0123456789abcdef0123456789abcdef01234567";
const STARTUP_REMOTE_PLUGIN_SYNC_MARKER_FILE: &str = ".tmp/app-server-remote-plugin-sync-v1";

Expand Down
5 changes: 4 additions & 1 deletion codex-rs/app-server/tests/suite/v2/realtime_conversation.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -72,6 +72,7 @@ use wiremock::matchers::path;
use wiremock::matchers::path_regex;

const DEFAULT_TIMEOUT: Duration = Duration::from_secs(10);
const DELEGATED_SHELL_TOOL_TIMEOUT_MS: u64 = 30_000;
const STARTUP_CONTEXT_HEADER: &str = "Startup context from Codex.";
const V2_STEERING_ACKNOWLEDGEMENT: &str =
"This was sent to steer the previous background agent task.";
Expand Down Expand Up @@ -1717,7 +1718,9 @@ async fn webrtc_v2_tool_call_delegated_turn_can_execute_shell_tool() -> Result<(
create_shell_command_sse_response(
realtime_tool_ok_command(),
/*workdir*/ None,
Some(5000),
// Windows CI can spend several seconds starting the nested PowerShell command. This
// test verifies delegated shell-tool plumbing, not timeout enforcement.
Some(DELEGATED_SHELL_TOOL_TIMEOUT_MS),
"shell_call",
)?,
create_final_assistant_message_sse_response("shell tool finished")?,
Expand Down
2 changes: 2 additions & 0 deletions codex-rs/core/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@ codex-connectors = { workspace = true }
codex-config = { workspace = true }
codex-core-plugins = { workspace = true }
codex-core-skills = { workspace = true }
crypto_box = { workspace = true }
codex-exec-server = { workspace = true }
codex-features = { workspace = true }
codex-feedback = { workspace = true }
Expand Down Expand Up @@ -99,6 +100,7 @@ rmcp = { workspace = true, default-features = false, features = [
serde = { workspace = true, features = ["derive"] }
serde_json = { workspace = true }
sha1 = { workspace = true }
sha2 = { workspace = true }
shlex = { workspace = true }
similar = { workspace = true }
tempfile = { workspace = true }
Expand Down
81 changes: 62 additions & 19 deletions codex-rs/core/src/agent_identity.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,10 @@ use tracing::debug;
use tracing::info;
use tracing::warn;

mod task_registration;

pub(crate) use task_registration::RegisteredAgentTask;

use crate::config::Config;

const AGENT_REGISTRATION_TIMEOUT: Duration = Duration::from_secs(15);
Expand Down Expand Up @@ -119,32 +123,58 @@ impl AgentIdentityManager {
return Ok(None);
}

let Some(auth) = self.auth_manager.auth().await else {
debug!("skipping agent identity registration because no auth is available");
let Some((auth, binding)) = self.current_auth_binding().await else {
return Ok(None);
};

let Some(binding) =
AgentIdentityBinding::from_auth(&auth, self.auth_manager.forced_chatgpt_workspace_id())
else {
debug!("skipping agent identity registration because ChatGPT auth is unavailable");
return Ok(None);
};
self.ensure_registered_identity_for_binding(&auth, &binding)
.await
.map(Some)
}

async fn ensure_registered_identity_for_binding(
&self,
auth: &CodexAuth,
binding: &AgentIdentityBinding,
) -> Result<StoredAgentIdentity> {
let _guard = self.ensure_lock.lock().await;

if let Some(stored_identity) = self.load_stored_identity(&auth, &binding)? {
if let Some(stored_identity) = self.load_stored_identity(auth, binding)? {
info!(
agent_runtime_id = %stored_identity.agent_runtime_id,
binding_id = %binding.binding_id,
"reusing stored agent identity"
);
return Ok(Some(stored_identity));
return Ok(stored_identity);
}

let stored_identity = self.register_agent_identity(&binding).await?;
self.store_identity(&auth, &stored_identity)?;
Ok(Some(stored_identity))
let stored_identity = self.register_agent_identity(binding).await?;
self.store_identity(auth, &stored_identity)?;
Ok(stored_identity)
}

pub(crate) async fn task_matches_current_binding(&self, task: &RegisteredAgentTask) -> bool {
if !self.feature_enabled {
return false;
}

self.current_auth_binding()
.await
.is_some_and(|(_, binding)| task.matches_binding(&binding))
}

async fn current_auth_binding(&self) -> Option<(CodexAuth, AgentIdentityBinding)> {
let Some(auth) = self.auth_manager.auth().await else {
debug!("skipping agent identity flow because no auth is available");
return None;
};

let binding =
AgentIdentityBinding::from_auth(&auth, self.auth_manager.forced_chatgpt_workspace_id());
if binding.is_none() {
debug!("skipping agent identity flow because ChatGPT auth is unavailable");
}
binding.map(|binding| (auth, binding))
}

async fn register_agent_identity(
Expand Down Expand Up @@ -351,12 +381,11 @@ impl StoredAgentIdentity {
}

fn matches_binding(&self, binding: &AgentIdentityBinding) -> bool {
self.binding_id == binding.binding_id
&& self.chatgpt_account_id == binding.chatgpt_account_id
&& match binding.chatgpt_user_id.as_deref() {
Some(chatgpt_user_id) => self.chatgpt_user_id.as_deref() == Some(chatgpt_user_id),
None => true,
}
binding.matches_parts(
&self.binding_id,
&self.chatgpt_account_id,
self.chatgpt_user_id.as_deref(),
)
}

fn validate_key_material(&self) -> Result<()> {
Expand All @@ -375,6 +404,20 @@ impl StoredAgentIdentity {
}

impl AgentIdentityBinding {
fn matches_parts(
&self,
binding_id: &str,
chatgpt_account_id: &str,
chatgpt_user_id: Option<&str>,
) -> bool {
binding_id == self.binding_id
&& chatgpt_account_id == self.chatgpt_account_id
&& match self.chatgpt_user_id.as_deref() {
Some(expected_user_id) => chatgpt_user_id == Some(expected_user_id),
None => true,
}
}

fn from_auth(auth: &CodexAuth, forced_workspace_id: Option<String>) -> Option<Self> {
if !auth.is_chatgpt_auth() {
return None;
Expand Down
Loading
Loading