diff --git a/MODULE.bazel.lock b/MODULE.bazel.lock index 0e2232fcbf1..0aaa7a6f528 100644 --- a/MODULE.bazel.lock +++ b/MODULE.bazel.lock @@ -663,6 +663,7 @@ "bitflags_1.3.2": "{\"dependencies\":[{\"name\":\"compiler_builtins\",\"optional\":true,\"req\":\"^0.1.2\"},{\"name\":\"core\",\"optional\":true,\"package\":\"rustc-std-workspace-core\",\"req\":\"^1.0.0\"},{\"kind\":\"dev\",\"name\":\"rustversion\",\"req\":\"^1.0\"},{\"kind\":\"dev\",\"name\":\"serde\",\"req\":\"^1.0\"},{\"kind\":\"dev\",\"name\":\"serde_derive\",\"req\":\"^1.0\"},{\"kind\":\"dev\",\"name\":\"serde_json\",\"req\":\"^1.0\"},{\"kind\":\"dev\",\"name\":\"trybuild\",\"req\":\"^1.0\"},{\"kind\":\"dev\",\"name\":\"walkdir\",\"req\":\"^2.3\"}],\"features\":{\"default\":[],\"example_generated\":[],\"rustc-dep-of-std\":[\"core\",\"compiler_builtins\"]}}", "bitflags_2.10.0": "{\"dependencies\":[{\"name\":\"arbitrary\",\"optional\":true,\"req\":\"^1.0\"},{\"features\":[\"derive\"],\"kind\":\"dev\",\"name\":\"arbitrary\",\"req\":\"^1.0\"},{\"name\":\"bytemuck\",\"optional\":true,\"req\":\"^1.12\"},{\"features\":[\"derive\"],\"kind\":\"dev\",\"name\":\"bytemuck\",\"req\":\"^1.12.2\"},{\"kind\":\"dev\",\"name\":\"rustversion\",\"req\":\"^1.0\"},{\"default_features\":false,\"name\":\"serde_core\",\"optional\":true,\"req\":\"^1.0.228\"},{\"kind\":\"dev\",\"name\":\"serde_json\",\"req\":\"^1.0\"},{\"features\":[\"derive\"],\"kind\":\"dev\",\"name\":\"serde_lib\",\"package\":\"serde\",\"req\":\"^1.0.103\"},{\"kind\":\"dev\",\"name\":\"serde_test\",\"req\":\"^1.0.19\"},{\"kind\":\"dev\",\"name\":\"trybuild\",\"req\":\"^1.0.18\"},{\"features\":[\"derive\"],\"kind\":\"dev\",\"name\":\"zerocopy\",\"req\":\"^0.8\"}],\"features\":{\"example_generated\":[],\"serde\":[\"serde_core\"],\"std\":[]}}", "bitflags_2.11.0": "{\"dependencies\":[{\"name\":\"arbitrary\",\"optional\":true,\"req\":\"^1.0\"},{\"features\":[\"derive\"],\"kind\":\"dev\",\"name\":\"arbitrary\",\"req\":\"^1.0\"},{\"name\":\"bytemuck\",\"optional\":true,\"req\":\"^1.12\"},{\"features\":[\"derive\"],\"kind\":\"dev\",\"name\":\"bytemuck\",\"req\":\"^1.12.2\"},{\"kind\":\"dev\",\"name\":\"rustversion\",\"req\":\"^1.0\"},{\"default_features\":false,\"name\":\"serde_core\",\"optional\":true,\"req\":\"^1.0.228\"},{\"kind\":\"dev\",\"name\":\"serde_json\",\"req\":\"^1.0\"},{\"features\":[\"derive\"],\"kind\":\"dev\",\"name\":\"serde_lib\",\"package\":\"serde\",\"req\":\"^1.0.103\"},{\"kind\":\"dev\",\"name\":\"serde_test\",\"req\":\"^1.0.19\"},{\"kind\":\"dev\",\"name\":\"trybuild\",\"req\":\"^1.0.18\"},{\"features\":[\"derive\"],\"kind\":\"dev\",\"name\":\"zerocopy\",\"req\":\"^0.8\"}],\"features\":{\"example_generated\":[],\"serde\":[\"serde_core\"],\"std\":[]}}", + "blake2_0.10.6": "{\"dependencies\":[{\"features\":[\"mac\"],\"name\":\"digest\",\"req\":\"^0.10.3\"},{\"features\":[\"dev\"],\"kind\":\"dev\",\"name\":\"digest\",\"req\":\"^0.10.3\"},{\"kind\":\"dev\",\"name\":\"hex-literal\",\"req\":\"^0.2.2\"}],\"features\":{\"default\":[\"std\"],\"reset\":[],\"simd\":[],\"simd_asm\":[\"simd_opt\"],\"simd_opt\":[\"simd\"],\"size_opt\":[],\"std\":[\"digest/std\"]}}", "block-buffer_0.10.4": "{\"dependencies\":[{\"name\":\"generic-array\",\"req\":\"^0.14\"}],\"features\":{}}", "block-padding_0.3.3": "{\"dependencies\":[{\"name\":\"generic-array\",\"req\":\"^0.14\"}],\"features\":{\"std\":[]}}", "block2_0.6.2": "{\"dependencies\":[{\"default_features\":false,\"features\":[\"std\"],\"name\":\"objc2\",\"req\":\">=0.6.2, <0.8.0\"}],\"features\":{\"alloc\":[],\"compiler-rt\":[\"objc2/unstable-compiler-rt\"],\"default\":[\"std\"],\"gnustep-1-7\":[\"objc2/gnustep-1-7\"],\"gnustep-1-8\":[\"gnustep-1-7\",\"objc2/gnustep-1-8\"],\"gnustep-1-9\":[\"gnustep-1-8\",\"objc2/gnustep-1-9\"],\"gnustep-2-0\":[\"gnustep-1-9\",\"objc2/gnustep-2-0\"],\"gnustep-2-1\":[\"gnustep-2-0\",\"objc2/gnustep-2-1\"],\"std\":[\"alloc\"],\"unstable-coerce-pointee\":[],\"unstable-objfw\":[],\"unstable-private\":[],\"unstable-winobjc\":[\"gnustep-1-8\"]}}", @@ -752,6 +753,8 @@ "crossterm_winapi_0.9.1": "{\"dependencies\":[{\"features\":[\"winbase\",\"consoleapi\",\"processenv\",\"handleapi\",\"synchapi\",\"impl-default\"],\"name\":\"winapi\",\"req\":\"^0.3.8\",\"target\":\"cfg(windows)\"}],\"features\":{}}", "crunchy_0.2.4": "{\"dependencies\":[],\"features\":{\"default\":[\"limit_128\"],\"limit_1024\":[],\"limit_128\":[],\"limit_2048\":[],\"limit_256\":[],\"limit_512\":[],\"limit_64\":[],\"std\":[]}}", "crypto-common_0.1.7": "{\"dependencies\":[{\"features\":[\"more_lengths\"],\"name\":\"generic-array\",\"req\":\"=0.14.7\"},{\"name\":\"rand_core\",\"optional\":true,\"req\":\"^0.6\"},{\"name\":\"typenum\",\"req\":\"^1.14\"}],\"features\":{\"getrandom\":[\"rand_core/getrandom\"],\"std\":[]}}", + "crypto_box_0.9.1": "{\"dependencies\":[{\"default_features\":false,\"name\":\"aead\",\"req\":\"^0.5.2\"},{\"kind\":\"dev\",\"name\":\"bincode\",\"req\":\"^1\"},{\"default_features\":false,\"name\":\"blake2\",\"optional\":true,\"req\":\"^0.10\"},{\"name\":\"chacha20\",\"optional\":true,\"req\":\"^0.9\"},{\"default_features\":false,\"name\":\"crypto_secretbox\",\"req\":\"^0.1.1\"},{\"default_features\":false,\"features\":[\"zeroize\"],\"name\":\"curve25519-dalek\",\"req\":\"^4\"},{\"kind\":\"dev\",\"name\":\"hex-literal\",\"req\":\"^0.4\"},{\"kind\":\"dev\",\"name\":\"rand\",\"req\":\"^0.8\"},{\"kind\":\"dev\",\"name\":\"rmp-serde\",\"req\":\"^1\"},{\"name\":\"salsa20\",\"optional\":true,\"req\":\"^0.10\"},{\"default_features\":false,\"name\":\"serdect\",\"optional\":true,\"req\":\"^0.2\"},{\"default_features\":false,\"name\":\"subtle\",\"req\":\"^2\"},{\"default_features\":false,\"name\":\"zeroize\",\"req\":\"^1\"}],\"features\":{\"alloc\":[\"aead/alloc\"],\"chacha20\":[\"dep:chacha20\",\"crypto_secretbox/chacha20\"],\"default\":[\"alloc\",\"getrandom\",\"salsa20\"],\"getrandom\":[\"aead/getrandom\",\"rand_core\"],\"heapless\":[\"aead/heapless\"],\"rand_core\":[\"aead/rand_core\"],\"salsa20\":[\"dep:salsa20\",\"crypto_secretbox/salsa20\"],\"seal\":[\"dep:blake2\",\"alloc\"],\"serde\":[\"dep:serdect\"],\"std\":[\"aead/std\"]}}", + "crypto_secretbox_0.1.1": "{\"dependencies\":[{\"default_features\":false,\"name\":\"aead\",\"req\":\"^0.5\"},{\"features\":[\"zeroize\"],\"name\":\"chacha20\",\"optional\":true,\"req\":\"^0.9\"},{\"default_features\":false,\"name\":\"cipher\",\"req\":\"^0.4\"},{\"default_features\":false,\"features\":[\"zeroize\"],\"name\":\"generic-array\",\"req\":\"^0.14.7\"},{\"kind\":\"dev\",\"name\":\"hex-literal\",\"req\":\"^0.4\"},{\"name\":\"poly1305\",\"req\":\"^0.8\"},{\"features\":[\"zeroize\"],\"name\":\"salsa20\",\"optional\":true,\"req\":\"^0.10\"},{\"default_features\":false,\"name\":\"subtle\",\"req\":\"^2\"},{\"default_features\":false,\"name\":\"zeroize\",\"req\":\"^1\"}],\"features\":{\"alloc\":[\"aead/alloc\"],\"default\":[\"alloc\",\"getrandom\",\"salsa20\"],\"getrandom\":[\"aead/getrandom\",\"rand_core\"],\"heapless\":[\"aead/heapless\"],\"rand_core\":[\"aead/rand_core\"],\"std\":[\"aead/std\",\"alloc\"],\"stream\":[\"aead/stream\"]}}", "csv-core_0.1.13": "{\"dependencies\":[{\"default_features\":false,\"kind\":\"dev\",\"name\":\"arrayvec\",\"req\":\"^0.5\"},{\"default_features\":false,\"name\":\"memchr\",\"req\":\"^2\"}],\"features\":{\"default\":[],\"libc\":[\"memchr/libc\"]}}", "csv_1.4.0": "{\"dependencies\":[{\"default_features\":false,\"features\":[\"alloc\",\"serde\"],\"kind\":\"dev\",\"name\":\"bstr\",\"req\":\"^1.7.0\"},{\"name\":\"csv-core\",\"req\":\"^0.1.11\"},{\"name\":\"itoa\",\"req\":\"^1\"},{\"name\":\"ryu\",\"req\":\"^1\"},{\"features\":[\"derive\"],\"kind\":\"dev\",\"name\":\"serde\",\"req\":\"^1.0.221\"},{\"name\":\"serde_core\",\"req\":\"^1.0.221\"}],\"features\":{}}", "ctor-proc-macro_0.0.7": "{\"dependencies\":[],\"features\":{\"default\":[]}}", @@ -818,6 +821,8 @@ "dylint_linting_5.0.0": "{\"dependencies\":[{\"kind\":\"dev\",\"name\":\"assert_cmd\",\"req\":\"^2.0\"},{\"name\":\"cargo_metadata\",\"req\":\"^0.23\"},{\"features\":[\"config\"],\"name\":\"dylint_internal\",\"req\":\"=5.0.0\"},{\"name\":\"paste\",\"req\":\"^1.0\"},{\"kind\":\"dev\",\"name\":\"rustc_version\",\"req\":\"^0.4\"},{\"name\":\"rustversion\",\"req\":\"^1.0\"},{\"name\":\"serde\",\"req\":\"^1.0\"},{\"kind\":\"dev\",\"name\":\"tempfile\",\"req\":\"^3.23\"},{\"name\":\"thiserror\",\"req\":\"^2.0\"},{\"name\":\"toml\",\"req\":\"^0.9\"},{\"kind\":\"build\",\"name\":\"toml\",\"req\":\"^0.9\"}],\"features\":{\"constituent\":[]}}", "dylint_testing_5.0.0": "{\"dependencies\":[{\"name\":\"anyhow\",\"req\":\"^1.0\"},{\"name\":\"cargo_metadata\",\"req\":\"^0.23\"},{\"name\":\"compiletest_rs\",\"req\":\"^0.11\"},{\"name\":\"dylint\",\"req\":\"=5.0.0\"},{\"name\":\"dylint_internal\",\"req\":\"=5.0.0\"},{\"name\":\"env_logger\",\"req\":\"^0.11\"},{\"name\":\"once_cell\",\"req\":\"^1.21\"},{\"name\":\"regex\",\"req\":\"^1.11\"},{\"name\":\"serde_json\",\"req\":\"^1.0\"},{\"name\":\"tempfile\",\"req\":\"^3.23\"}],\"features\":{\"default\":[],\"deny_warnings\":[]}}", "dyn-clone_1.0.20": "{\"dependencies\":[{\"kind\":\"dev\",\"name\":\"rustversion\",\"req\":\"^1.0\"},{\"features\":[\"diff\"],\"kind\":\"dev\",\"name\":\"trybuild\",\"req\":\"^1.0.66\"}],\"features\":{}}", + "ed25519-dalek_2.2.0": "{\"dependencies\":[{\"kind\":\"dev\",\"name\":\"bincode\",\"req\":\"^1.0\"},{\"kind\":\"dev\",\"name\":\"blake2\",\"req\":\"^0.10\"},{\"features\":[\"html_reports\"],\"kind\":\"dev\",\"name\":\"criterion\",\"req\":\"^0.5\"},{\"default_features\":false,\"features\":[\"digest\"],\"name\":\"curve25519-dalek\",\"req\":\"^4\"},{\"default_features\":false,\"features\":[\"digest\",\"rand_core\"],\"kind\":\"dev\",\"name\":\"curve25519-dalek\",\"req\":\"^4\"},{\"default_features\":false,\"name\":\"ed25519\",\"req\":\">=2.2, <2.3\"},{\"kind\":\"dev\",\"name\":\"hex\",\"req\":\"^0.4\"},{\"kind\":\"dev\",\"name\":\"hex-literal\",\"req\":\"^0.4\"},{\"default_features\":false,\"name\":\"merlin\",\"optional\":true,\"req\":\"^3\"},{\"kind\":\"dev\",\"name\":\"rand\",\"req\":\"^0.8\"},{\"default_features\":false,\"name\":\"rand_core\",\"optional\":true,\"req\":\"^0.6.4\"},{\"default_features\":false,\"kind\":\"dev\",\"name\":\"rand_core\",\"req\":\"^0.6.4\"},{\"default_features\":false,\"name\":\"serde\",\"optional\":true,\"req\":\"^1.0\"},{\"features\":[\"derive\"],\"kind\":\"dev\",\"name\":\"serde\",\"req\":\"^1.0\"},{\"kind\":\"dev\",\"name\":\"serde_json\",\"req\":\"^1.0\"},{\"default_features\":false,\"name\":\"sha2\",\"req\":\"^0.10\"},{\"kind\":\"dev\",\"name\":\"sha3\",\"req\":\"^0.10\"},{\"default_features\":false,\"name\":\"signature\",\"optional\":true,\"req\":\">=2.0, <2.3\"},{\"default_features\":false,\"name\":\"subtle\",\"req\":\"^2.3.0\"},{\"kind\":\"dev\",\"name\":\"toml\",\"req\":\"^0.7\"},{\"default_features\":false,\"features\":[\"static_secrets\"],\"kind\":\"dev\",\"name\":\"x25519-dalek\",\"req\":\"^2\"},{\"default_features\":false,\"name\":\"zeroize\",\"optional\":true,\"req\":\"^1.5\"}],\"features\":{\"alloc\":[\"curve25519-dalek/alloc\",\"ed25519/alloc\",\"serde?/alloc\",\"zeroize/alloc\"],\"asm\":[\"sha2/asm\"],\"batch\":[\"alloc\",\"merlin\",\"rand_core\"],\"default\":[\"fast\",\"std\",\"zeroize\"],\"digest\":[\"signature/digest\"],\"fast\":[\"curve25519-dalek/precomputed-tables\"],\"hazmat\":[],\"legacy_compatibility\":[\"curve25519-dalek/legacy_compatibility\"],\"pem\":[\"alloc\",\"ed25519/pem\",\"pkcs8\"],\"pkcs8\":[\"ed25519/pkcs8\"],\"rand_core\":[\"dep:rand_core\"],\"serde\":[\"dep:serde\",\"ed25519/serde\"],\"std\":[\"alloc\",\"ed25519/std\",\"serde?/std\",\"sha2/std\"],\"zeroize\":[\"dep:zeroize\",\"curve25519-dalek/zeroize\"]}}", + "ed25519_2.2.3": "{\"dependencies\":[{\"kind\":\"dev\",\"name\":\"bincode\",\"req\":\"^1\"},{\"features\":[\"rand_core\"],\"kind\":\"dev\",\"name\":\"ed25519-dalek\",\"req\":\"^2\"},{\"kind\":\"dev\",\"name\":\"hex-literal\",\"req\":\"^0.4\"},{\"name\":\"pkcs8\",\"optional\":true,\"req\":\"^0.10\"},{\"features\":[\"std\"],\"kind\":\"dev\",\"name\":\"rand_core\",\"req\":\"^0.6\"},{\"default_features\":false,\"features\":[\"signature\"],\"kind\":\"dev\",\"name\":\"ring-compat\",\"req\":\"^0.8\"},{\"default_features\":false,\"name\":\"serde\",\"optional\":true,\"req\":\"^1\"},{\"name\":\"serde_bytes\",\"optional\":true,\"req\":\"^0.11\"},{\"default_features\":false,\"name\":\"signature\",\"req\":\"^2\"},{\"default_features\":false,\"name\":\"zeroize\",\"optional\":true,\"req\":\"^1\"}],\"features\":{\"alloc\":[\"pkcs8?/alloc\"],\"default\":[\"std\"],\"pem\":[\"alloc\",\"pkcs8/pem\"],\"serde_bytes\":[\"serde\",\"dep:serde_bytes\"],\"std\":[\"pkcs8?/std\",\"signature/std\"]}}", "either_1.15.0": "{\"dependencies\":[{\"default_features\":false,\"features\":[\"alloc\",\"derive\"],\"name\":\"serde\",\"optional\":true,\"req\":\"^1.0.95\"},{\"kind\":\"dev\",\"name\":\"serde_json\",\"req\":\"^1.0.0\"}],\"features\":{\"default\":[\"std\"],\"std\":[],\"use_std\":[\"std\"]}}", "ena_0.14.3": "{\"dependencies\":[{\"name\":\"dogged\",\"optional\":true,\"req\":\"^0.2.0\"},{\"name\":\"log\",\"req\":\"^0.4\"}],\"features\":{\"bench\":[],\"persistent\":[\"dogged\"]}}", "encode_unicode_1.0.0": "{\"dependencies\":[{\"default_features\":false,\"name\":\"ascii\",\"optional\":true,\"req\":\"^1.0.0\"},{\"kind\":\"dev\",\"name\":\"lazy_static\",\"req\":\"^1.0\",\"target\":\"cfg(unix)\"},{\"features\":[\"https-native\"],\"kind\":\"dev\",\"name\":\"minreq\",\"req\":\"^2.6\"}],\"features\":{\"default\":[\"std\"],\"std\":[]}}", diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 6c101b940d3..bec80224c0b 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -940,6 +940,15 @@ dependencies = [ "serde_core", ] +[[package]] +name = "blake2" +version = "0.10.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "46502ad458c9a52b69d4d4d32775c788b7a1b85e8bc9d482d92250fc0e3f8efe" +dependencies = [ + "digest", +] + [[package]] name = "block-buffer" version = "0.10.4" @@ -1949,10 +1958,12 @@ dependencies = [ "codex-windows-sandbox", "core-foundation 0.9.4", "core_test_support", + "crypto_box", "csv", "ctor 0.6.3", "dirs", "dunce", + "ed25519-dalek", "env-flags", "eventsource-stream", "futures", @@ -1978,6 +1989,7 @@ dependencies = [ "serde_json", "serial_test", "sha1", + "sha2", "shlex", "similar", "tempfile", @@ -3558,9 +3570,40 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a" dependencies = [ "generic-array", + "rand_core 0.6.4", "typenum", ] +[[package]] +name = "crypto_box" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "16182b4f39a82ec8a6851155cc4c0cda3065bb1db33651726a29e1951de0f009" +dependencies = [ + "aead", + "blake2", + "crypto_secretbox", + "curve25519-dalek", + "salsa20", + "subtle", + "zeroize", +] + +[[package]] +name = "crypto_secretbox" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9d6cf87adf719ddf43a805e92c6870a531aedda35ff640442cbaf8674e141e1" +dependencies = [ + "aead", + "cipher", + "generic-array", + "poly1305", + "salsa20", + "subtle", + "zeroize", +] + [[package]] name = "csv" version = "1.4.0" @@ -3617,6 +3660,7 @@ dependencies = [ "cfg-if", "cpufeatures", "curve25519-dalek-derive", + "digest", "fiat-crypto", "rustc_version", "subtle", @@ -4087,7 +4131,7 @@ dependencies = [ "libc", "option-ext", "redox_users 0.5.2", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -4191,6 +4235,30 @@ version = "1.0.20" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d0881ea181b1df73ff77ffaaf9c7544ecc11e82fba9b5f27b262a3c73a332555" +[[package]] +name = "ed25519" +version = "2.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "115531babc129696a58c64a4fef0a8bf9e9698629fb97e9e40767d235cfbcd53" +dependencies = [ + "pkcs8", + "signature", +] + +[[package]] +name = "ed25519-dalek" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "70e796c081cee67dc755e1a36a0a172b897fab85fc3f6bc48307991f64e4eca9" +dependencies = [ + "curve25519-dalek", + "ed25519", + "serde", + "sha2", + "subtle", + "zeroize", +] + [[package]] name = "either" version = "1.15.0" @@ -4819,6 +4887,7 @@ checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" dependencies = [ "typenum", "version_check", + "zeroize", ] [[package]] @@ -6689,7 +6758,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -11813,7 +11882,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.48.0", + "windows-sys 0.61.2", ] [[package]] diff --git a/codex-rs/Cargo.toml b/codex-rs/Cargo.toml index 32ae50bfb71..c610ee34e76 100644 --- a/codex-rs/Cargo.toml +++ b/codex-rs/Cargo.toml @@ -216,6 +216,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" derive_more = "2" @@ -223,6 +224,7 @@ diffy = "0.4.2" dirs = "6" dotenvy = "0.15.7" dunce = "1.0.4" +ed25519-dalek = { version = "2.2.0", features = ["pkcs8"] } encoding_rs = "0.8.35" env-flags = "0.1.1" env_logger = "0.11.9" diff --git a/codex-rs/app-server-protocol/src/protocol/thread_history.rs b/codex-rs/app-server-protocol/src/protocol/thread_history.rs index d2296c75c5d..6c85612bc99 100644 --- a/codex-rs/app-server-protocol/src/protocol/thread_history.rs +++ b/codex-rs/app-server-protocol/src/protocol/thread_history.rs @@ -218,7 +218,9 @@ impl ThreadHistoryBuilder { RolloutItem::EventMsg(event) => self.handle_event(event), RolloutItem::Compacted(payload) => self.handle_compacted(payload), RolloutItem::ResponseItem(item) => self.handle_response_item(item), - RolloutItem::TurnContext(_) | RolloutItem::SessionMeta(_) => {} + RolloutItem::TurnContext(_) + | RolloutItem::SessionMeta(_) + | RolloutItem::SessionState(_) => {} } } diff --git a/codex-rs/app-server/src/transport/remote_control/tests.rs b/codex-rs/app-server/src/transport/remote_control/tests.rs index 21808a3a18f..6dd6978fa5e 100644 --- a/codex-rs/app-server/src/transport/remote_control/tests.rs +++ b/codex-rs/app-server/src/transport/remote_control/tests.rs @@ -96,6 +96,7 @@ fn remote_control_auth_dot_json(account_id: Option<&str>) -> AuthDotJson { account_id: account_id.map(str::to_string), }), last_refresh: Some(chrono::Utc::now()), + agent_identity: None, } } diff --git a/codex-rs/app-server/src/transport/remote_control/websocket.rs b/codex-rs/app-server/src/transport/remote_control/websocket.rs index 42924c2d34f..d168ec0e180 100644 --- a/codex-rs/app-server/src/transport/remote_control/websocket.rs +++ b/codex-rs/app-server/src/transport/remote_control/websocket.rs @@ -971,6 +971,7 @@ mod tests { account_id: Some("account_id".to_string()), }), last_refresh: Some(Utc::now()), + agent_identity: None, } } diff --git a/codex-rs/app-server/tests/common/auth_fixtures.rs b/codex-rs/app-server/tests/common/auth_fixtures.rs index 99334f07706..86f0fb456dd 100644 --- a/codex-rs/app-server/tests/common/auth_fixtures.rs +++ b/codex-rs/app-server/tests/common/auth_fixtures.rs @@ -163,6 +163,7 @@ pub fn write_chatgpt_auth( openai_api_key: None, tokens: Some(tokens), last_refresh, + agent_identity: None, }; save_auth(codex_home, &auth, cli_auth_credentials_store_mode).context("write auth.json") diff --git a/codex-rs/app-server/tests/suite/v2/app_list.rs b/codex-rs/app-server/tests/suite/v2/app_list.rs index 57a27961af6..dbe61524f5e 100644 --- a/codex-rs/app-server/tests/suite/v2/app_list.rs +++ b/codex-rs/app-server/tests/suite/v2/app_list.rs @@ -117,6 +117,7 @@ async fn list_apps_returns_empty_with_api_key_auth() -> Result<()> { openai_api_key: Some("test-api-key".to_string()), tokens: None, last_refresh: None, + agent_identity: None, }, AuthCredentialsStoreMode::File, )?; diff --git a/codex-rs/app-server/tests/suite/v2/client_metadata.rs b/codex-rs/app-server/tests/suite/v2/client_metadata.rs index c85febd7d46..e343e08470b 100644 --- a/codex-rs/app-server/tests/suite/v2/client_metadata.rs +++ b/codex-rs/app-server/tests/suite/v2/client_metadata.rs @@ -18,7 +18,7 @@ use std::path::Path; use tempfile::TempDir; use tokio::time::timeout; -const DEFAULT_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10); +const DEFAULT_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(20); #[tokio::test] async fn turn_start_forwards_client_metadata_to_responses_request_v2() -> Result<()> { diff --git a/codex-rs/app-server/tests/suite/v2/collaboration_mode_list.rs b/codex-rs/app-server/tests/suite/v2/collaboration_mode_list.rs index 7c36827e6dd..0dd4e8174a8 100644 --- a/codex-rs/app-server/tests/suite/v2/collaboration_mode_list.rs +++ b/codex-rs/app-server/tests/suite/v2/collaboration_mode_list.rs @@ -21,7 +21,7 @@ use pretty_assertions::assert_eq; use tempfile::TempDir; use tokio::time::timeout; -const DEFAULT_TIMEOUT: Duration = Duration::from_secs(10); +const DEFAULT_TIMEOUT: Duration = Duration::from_secs(20); /// Confirms the server returns the default collaboration mode presets in a stable order. #[tokio::test] diff --git a/codex-rs/app-server/tests/suite/v2/connection_handling_websocket.rs b/codex-rs/app-server/tests/suite/v2/connection_handling_websocket.rs index 30caa13761c..e6818eabd7d 100644 --- a/codex-rs/app-server/tests/suite/v2/connection_handling_websocket.rs +++ b/codex-rs/app-server/tests/suite/v2/connection_handling_websocket.rs @@ -47,7 +47,7 @@ use tokio_tungstenite::tungstenite::http::HeaderValue; use tokio_tungstenite::tungstenite::http::header::AUTHORIZATION; use tokio_tungstenite::tungstenite::http::header::ORIGIN; -pub(super) const DEFAULT_READ_TIMEOUT: Duration = Duration::from_secs(10); +pub(super) const DEFAULT_READ_TIMEOUT: Duration = Duration::from_secs(20); pub(super) type WsClient = WebSocketStream>; type HmacSha256 = Hmac; diff --git a/codex-rs/app-server/tests/suite/v2/initialize.rs b/codex-rs/app-server/tests/suite/v2/initialize.rs index 165160468f7..cfb08a17187 100644 --- a/codex-rs/app-server/tests/suite/v2/initialize.rs +++ b/codex-rs/app-server/tests/suite/v2/initialize.rs @@ -24,7 +24,7 @@ use std::time::Duration; use tempfile::TempDir; use tokio::time::timeout; -const DEFAULT_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10); +const DEFAULT_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(20); #[tokio::test] async fn initialize_uses_client_info_name_as_originator() -> Result<()> { diff --git a/codex-rs/app-server/tests/suite/v2/realtime_conversation.rs b/codex-rs/app-server/tests/suite/v2/realtime_conversation.rs index efc167c9d06..4c1b80498c3 100644 --- a/codex-rs/app-server/tests/suite/v2/realtime_conversation.rs +++ b/codex-rs/app-server/tests/suite/v2/realtime_conversation.rs @@ -69,7 +69,7 @@ use wiremock::matchers::method; use wiremock::matchers::path; use wiremock::matchers::path_regex; -const DEFAULT_TIMEOUT: Duration = Duration::from_secs(10); +const DEFAULT_TIMEOUT: Duration = Duration::from_secs(20); const STARTUP_CONTEXT_HEADER: &str = "Startup context from Codex."; const V2_STEERING_ACKNOWLEDGEMENT: &str = "This was sent to steer the previous background agent task."; @@ -1574,7 +1574,7 @@ 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), + Some(10_000), "shell_call", )?, create_final_assistant_message_sse_response("shell tool finished")?, diff --git a/codex-rs/codex-api/src/api_bridge.rs b/codex-rs/codex-api/src/api_bridge.rs index 0ad2b139795..740f7cb783d 100644 --- a/codex-rs/codex-api/src/api_bridge.rs +++ b/codex-rs/codex-api/src/api_bridge.rs @@ -178,13 +178,33 @@ struct UsageErrorBody { pub struct CoreAuthProvider { pub token: Option, pub account_id: Option, + authorization_header_override: Option, } impl CoreAuthProvider { + pub fn from_bearer_token(token: Option, account_id: Option) -> Self { + Self { + token, + account_id, + authorization_header_override: None, + } + } + + pub fn from_authorization_header_value( + authorization_header_value: Option, + account_id: Option, + ) -> Self { + Self { + token: None, + account_id, + authorization_header_override: authorization_header_value, + } + } + pub fn auth_header_attached(&self) -> bool { - self.token + self.authorization_header_value() .as_ref() - .is_some_and(|token| http::HeaderValue::from_str(&format!("Bearer {token}")).is_ok()) + .is_some_and(|value| http::HeaderValue::from_str(value).is_ok()) } pub fn auth_header_name(&self) -> Option<&'static str> { @@ -195,8 +215,20 @@ impl CoreAuthProvider { Self { token: token.map(str::to_string), account_id: account_id.map(str::to_string), + authorization_header_override: None, } } + + #[cfg(test)] + pub fn for_test_authorization_header( + authorization_header_value: Option<&str>, + account_id: Option<&str>, + ) -> Self { + Self::from_authorization_header_value( + authorization_header_value.map(str::to_string), + account_id.map(str::to_string), + ) + } } impl ApiAuthProvider for CoreAuthProvider { @@ -204,6 +236,12 @@ impl ApiAuthProvider for CoreAuthProvider { self.token.clone() } + fn authorization_header_value(&self) -> Option { + self.authorization_header_override + .clone() + .or_else(|| self.bearer_token().map(|token| format!("Bearer {token}"))) + } + fn account_id(&self) -> Option { self.account_id.clone() } diff --git a/codex-rs/codex-api/src/api_bridge_tests.rs b/codex-rs/codex-api/src/api_bridge_tests.rs index 71d3889915c..51c7c8d2fba 100644 --- a/codex-rs/codex-api/src/api_bridge_tests.rs +++ b/codex-rs/codex-api/src/api_bridge_tests.rs @@ -133,11 +133,27 @@ fn map_api_error_extracts_identity_auth_details_from_headers() { #[test] fn core_auth_provider_reports_when_auth_header_will_attach() { - let auth = CoreAuthProvider { - token: Some("access-token".to_string()), - account_id: None, - }; + let auth = CoreAuthProvider::from_bearer_token( + Some("access-token".to_string()), + /*account_id*/ None, + ); assert!(auth.auth_header_attached()); assert_eq!(auth.auth_header_name(), Some("authorization")); } + +#[test] +fn core_auth_provider_supports_non_bearer_authorization_headers() { + let auth = CoreAuthProvider::for_test_authorization_header( + Some("AgentAssertion opaque-token"), + /*account_id*/ None, + ); + + assert!(auth.auth_header_attached()); + assert_eq!(auth.auth_header_name(), Some("authorization")); + assert_eq!(auth.bearer_token(), None); + assert_eq!( + auth.authorization_header_value(), + Some("AgentAssertion opaque-token".to_string()) + ); +} diff --git a/codex-rs/codex-api/src/auth.rs b/codex-rs/codex-api/src/auth.rs index f649062db1f..4f27264a97a 100644 --- a/codex-rs/codex-api/src/auth.rs +++ b/codex-rs/codex-api/src/auth.rs @@ -9,14 +9,17 @@ use http::HeaderValue; /// reach this interface. pub trait AuthProvider: Send + Sync { fn bearer_token(&self) -> Option; + fn authorization_header_value(&self) -> Option { + self.bearer_token().map(|token| format!("Bearer {token}")) + } fn account_id(&self) -> Option { None } } pub(crate) fn add_auth_headers_to_header_map(auth: &A, headers: &mut HeaderMap) { - if let Some(token) = auth.bearer_token() - && let Ok(header) = HeaderValue::from_str(&format!("Bearer {token}")) + if let Some(authorization) = auth.authorization_header_value() + && let Ok(header) = HeaderValue::from_str(&authorization) { let _ = headers.insert(http::header::AUTHORIZATION, header); } diff --git a/codex-rs/core/Cargo.toml b/codex-rs/core/Cargo.toml index 55ce13afdc3..e45bf3d2588 100644 --- a/codex-rs/core/Cargo.toml +++ b/codex-rs/core/Cargo.toml @@ -34,6 +34,7 @@ codex-code-mode = { workspace = true } codex-connectors = { workspace = true } codex-config = { workspace = true } codex-core-skills = { workspace = true } +crypto_box = { workspace = true } codex-exec-server = { workspace = true } codex-features = { workspace = true } codex-feedback = { workspace = true } @@ -41,6 +42,7 @@ codex-login = { workspace = true } codex-mcp = { workspace = true } codex-model-provider-info = { workspace = true } codex-models-manager = { workspace = true } +ed25519-dalek = { workspace = true } codex-shell-command = { workspace = true } codex-execpolicy = { workspace = true } codex-git-utils = { workspace = true } @@ -96,6 +98,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 } diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index 509a578f5c9..d7a6403c937 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -116,6 +116,7 @@ fn keep_forked_rollout_item(item: &RolloutItem) -> bool { | ResponseItem::Compaction { .. } | ResponseItem::Other, ) => false, + RolloutItem::SessionState(_) => false, RolloutItem::Compacted(_) | RolloutItem::EventMsg(_) | RolloutItem::SessionMeta(_) diff --git a/codex-rs/core/src/agent/control_tests.rs b/codex-rs/core/src/agent/control_tests.rs index d332853167f..e3070c64015 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/codex-rs/core/src/agent/control_tests.rs @@ -186,7 +186,7 @@ async fn wait_for_subagent_notification(parent_thread: &Arc) -> boo sleep(Duration::from_millis(25)).await; } }; - timeout(Duration::from_secs(2), wait).await.is_ok() + timeout(Duration::from_secs(5), wait).await.is_ok() } async fn persist_thread_for_tree_resume(thread: &Arc, message: &str) { diff --git a/codex-rs/core/src/agent_identity.rs b/codex-rs/core/src/agent_identity.rs new file mode 100644 index 00000000000..e16ea49ee03 --- /dev/null +++ b/codex-rs/core/src/agent_identity.rs @@ -0,0 +1,795 @@ +use std::sync::Arc; +use std::time::Duration; + +use anyhow::Context; +use anyhow::Result; +use base64::Engine as _; +use base64::engine::general_purpose::STANDARD as BASE64_STANDARD; +use base64::engine::general_purpose::URL_SAFE_NO_PAD; +use chrono::SecondsFormat; +use chrono::Utc; +use codex_features::Feature; +use codex_login::AgentIdentityAuthRecord; +use codex_login::AuthManager; +use codex_login::CodexAuth; +use codex_login::default_client::create_client; +use codex_protocol::protocol::SessionSource; +use ed25519_dalek::SigningKey; +use ed25519_dalek::VerifyingKey; +use ed25519_dalek::pkcs8::DecodePrivateKey; +use ed25519_dalek::pkcs8::EncodePrivateKey; +use rand::TryRngCore; +use rand::rngs::OsRng; +use serde::Deserialize; +use serde::Serialize; +use tokio::sync::Mutex; +use tracing::debug; +use tracing::info; +use tracing::warn; + +use crate::config::Config; + +mod assertion; +mod task_registration; + +#[cfg(test)] +pub(crate) use assertion::AgentAssertionEnvelope; +pub(crate) use assertion::AgentTaskRuntimeMismatch; +pub(crate) use task_registration::RegisteredAgentTask; + +const AGENT_REGISTRATION_TIMEOUT: Duration = Duration::from_secs(15); +const AGENT_IDENTITY_BISCUIT_TIMEOUT: Duration = Duration::from_secs(15); + +#[derive(Clone)] +pub(crate) struct AgentIdentityManager { + auth_manager: Arc, + chatgpt_base_url: String, + feature_enabled: bool, + abom: AgentBillOfMaterials, + ensure_lock: Arc>, +} + +impl std::fmt::Debug for AgentIdentityManager { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("AgentIdentityManager") + .field("chatgpt_base_url", &self.chatgpt_base_url) + .field("feature_enabled", &self.feature_enabled) + .field("abom", &self.abom) + .finish_non_exhaustive() + } +} + +#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)] +pub(crate) struct StoredAgentIdentity { + pub(crate) binding_id: String, + pub(crate) chatgpt_account_id: String, + pub(crate) chatgpt_user_id: Option, + pub(crate) agent_runtime_id: String, + pub(crate) private_key_pkcs8_base64: String, + pub(crate) public_key_ssh: String, + pub(crate) registered_at: String, + pub(crate) abom: AgentBillOfMaterials, +} + +#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)] +pub(crate) struct AgentBillOfMaterials { + pub(crate) agent_version: String, + pub(crate) agent_harness_id: String, + pub(crate) running_location: String, +} + +#[derive(Debug, Serialize)] +struct RegisterAgentRequest { + abom: AgentBillOfMaterials, + agent_public_key: String, + capabilities: Vec, +} + +#[derive(Debug, Deserialize)] +struct RegisterAgentResponse { + agent_runtime_id: String, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +struct AgentIdentityBinding { + binding_id: String, + chatgpt_account_id: String, + chatgpt_user_id: Option, + access_token: String, +} + +struct GeneratedAgentKeyMaterial { + private_key_pkcs8_base64: String, + public_key_ssh: String, +} + +impl AgentIdentityManager { + pub(crate) fn new( + config: &Config, + auth_manager: Arc, + session_source: SessionSource, + ) -> Self { + Self { + auth_manager, + chatgpt_base_url: config.chatgpt_base_url.clone(), + feature_enabled: config.features.enabled(Feature::UseAgentIdentity), + abom: build_abom(session_source), + ensure_lock: Arc::new(Mutex::new(())), + } + } + + pub(crate) fn is_enabled(&self) -> bool { + self.feature_enabled + } + + pub(crate) async fn ensure_registered_identity(&self) -> Result> { + if !self.feature_enabled { + return Ok(None); + } + + let Some((auth, binding)) = self.current_auth_binding().await else { + 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 { + let _guard = self.ensure_lock.lock().await; + + 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(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_identity(&self, task: &RegisteredAgentTask) -> bool { + if !self.feature_enabled { + return false; + } + + self.current_stored_identity() + .await + .is_some_and(|stored_identity| { + stored_identity.agent_runtime_id == task.agent_runtime_id + }) + } + + 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 current_stored_identity(&self) -> Option { + let (auth, binding) = self.current_auth_binding().await?; + self.load_stored_identity(&auth, &binding).ok().flatten() + } + + async fn register_agent_identity( + &self, + binding: &AgentIdentityBinding, + ) -> Result { + let key_material = generate_agent_key_material()?; + let request_body = RegisterAgentRequest { + abom: self.abom.clone(), + agent_public_key: key_material.public_key_ssh.clone(), + capabilities: Vec::new(), + }; + + let url = agent_registration_url(&self.chatgpt_base_url); + let human_biscuit = self.mint_human_biscuit(binding, &url).await?; + let client = create_client(); + let response = client + .post(&url) + .header("X-OpenAI-Authorization", human_biscuit) + .json(&request_body) + .timeout(AGENT_REGISTRATION_TIMEOUT) + .send() + .await + .with_context(|| { + format!("failed to send agent identity registration request to {url}") + })?; + + if response.status().is_success() { + let response_body = response + .json::() + .await + .with_context(|| format!("failed to parse agent identity response from {url}"))?; + let stored_identity = StoredAgentIdentity { + binding_id: binding.binding_id.clone(), + chatgpt_account_id: binding.chatgpt_account_id.clone(), + chatgpt_user_id: binding.chatgpt_user_id.clone(), + agent_runtime_id: response_body.agent_runtime_id, + private_key_pkcs8_base64: key_material.private_key_pkcs8_base64, + public_key_ssh: key_material.public_key_ssh, + registered_at: Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true), + abom: self.abom.clone(), + }; + info!( + agent_runtime_id = %stored_identity.agent_runtime_id, + binding_id = %binding.binding_id, + "registered agent identity" + ); + return Ok(stored_identity); + } + + let status = response.status(); + let body = response.text().await.unwrap_or_default(); + anyhow::bail!("agent identity registration failed with status {status} from {url}: {body}") + } + + async fn mint_human_biscuit( + &self, + binding: &AgentIdentityBinding, + target_url: &str, + ) -> Result { + let url = agent_identity_biscuit_url(&self.chatgpt_base_url); + let request_id = agent_identity_request_id()?; + let client = create_client(); + let response = client + .get(&url) + .bearer_auth(&binding.access_token) + .header("X-Request-Id", request_id.clone()) + .header("X-Original-Method", "GET") + .header("X-Original-Url", target_url) + .timeout(AGENT_IDENTITY_BISCUIT_TIMEOUT) + .send() + .await + .with_context(|| format!("failed to send agent identity biscuit request to {url}"))?; + + if response.status().is_success() { + let human_biscuit = response + .headers() + .get("x-openai-authorization") + .context("agent identity biscuit response did not include x-openai-authorization")? + .to_str() + .context("agent identity biscuit response header was not valid UTF-8")? + .to_string(); + info!( + request_id = %request_id, + "minted human biscuit for agent identity registration" + ); + return Ok(human_biscuit); + } + + let status = response.status(); + let body = response.text().await.unwrap_or_default(); + anyhow::bail!( + "agent identity biscuit minting failed with status {status} from {url}: {body}" + ) + } + + fn load_stored_identity( + &self, + auth: &CodexAuth, + binding: &AgentIdentityBinding, + ) -> Result> { + let Some(record) = auth.get_agent_identity(&binding.chatgpt_account_id) else { + return Ok(None); + }; + + let stored_identity = + match StoredAgentIdentity::from_auth_record(binding, record, self.abom.clone()) { + Ok(stored_identity) => stored_identity, + Err(error) => { + warn!( + binding_id = %binding.binding_id, + error = %error, + "stored agent identity is invalid; deleting cached value" + ); + auth.remove_agent_identity()?; + return Ok(None); + } + }; + + if !stored_identity.matches_binding(binding) { + warn!( + binding_id = %binding.binding_id, + "stored agent identity binding no longer matches current auth; deleting cached value" + ); + auth.remove_agent_identity()?; + return Ok(None); + } + + if let Err(error) = stored_identity.validate_key_material() { + warn!( + agent_runtime_id = %stored_identity.agent_runtime_id, + binding_id = %binding.binding_id, + error = %error, + "stored agent identity key material is invalid; deleting cached value" + ); + auth.remove_agent_identity()?; + return Ok(None); + } + + Ok(Some(stored_identity)) + } + + fn store_identity( + &self, + auth: &CodexAuth, + stored_identity: &StoredAgentIdentity, + ) -> Result<()> { + auth.set_agent_identity(stored_identity.to_auth_record())?; + Ok(()) + } + + #[cfg(test)] + pub(crate) fn new_for_tests( + auth_manager: Arc, + feature_enabled: bool, + chatgpt_base_url: String, + session_source: SessionSource, + ) -> Self { + Self { + auth_manager, + chatgpt_base_url, + feature_enabled, + abom: build_abom(session_source), + ensure_lock: Arc::new(Mutex::new(())), + } + } + + #[cfg(test)] + pub(crate) async fn seed_generated_identity_for_tests( + &self, + agent_runtime_id: &str, + ) -> Result { + let (auth, binding) = self + .current_auth_binding() + .await + .context("test agent identity requires ChatGPT auth")?; + let key_material = generate_agent_key_material()?; + let stored_identity = StoredAgentIdentity { + binding_id: binding.binding_id.clone(), + chatgpt_account_id: binding.chatgpt_account_id.clone(), + chatgpt_user_id: binding.chatgpt_user_id, + agent_runtime_id: agent_runtime_id.to_string(), + private_key_pkcs8_base64: key_material.private_key_pkcs8_base64, + public_key_ssh: key_material.public_key_ssh, + registered_at: Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true), + abom: self.abom.clone(), + }; + self.store_identity(&auth, &stored_identity)?; + Ok(stored_identity) + } +} + +impl StoredAgentIdentity { + fn from_auth_record( + binding: &AgentIdentityBinding, + record: AgentIdentityAuthRecord, + abom: AgentBillOfMaterials, + ) -> Result { + if record.workspace_id != binding.chatgpt_account_id { + anyhow::bail!( + "stored agent identity workspace {:?} does not match current workspace {:?}", + record.workspace_id, + binding.chatgpt_account_id + ); + } + let signing_key = signing_key_from_private_key_pkcs8_base64(&record.agent_private_key)?; + Ok(Self { + binding_id: binding.binding_id.clone(), + chatgpt_account_id: binding.chatgpt_account_id.clone(), + chatgpt_user_id: binding.chatgpt_user_id.clone(), + agent_runtime_id: record.agent_runtime_id, + private_key_pkcs8_base64: record.agent_private_key, + public_key_ssh: encode_ssh_ed25519_public_key(&signing_key.verifying_key()), + registered_at: record.registered_at, + abom, + }) + } + + fn to_auth_record(&self) -> AgentIdentityAuthRecord { + AgentIdentityAuthRecord { + workspace_id: self.chatgpt_account_id.clone(), + agent_runtime_id: self.agent_runtime_id.clone(), + agent_private_key: self.private_key_pkcs8_base64.clone(), + registered_at: self.registered_at.clone(), + } + } + + fn matches_binding(&self, binding: &AgentIdentityBinding) -> bool { + binding.matches_parts( + &self.binding_id, + &self.chatgpt_account_id, + self.chatgpt_user_id.as_deref(), + ) + } + + fn validate_key_material(&self) -> Result<()> { + let signing_key = self.signing_key()?; + let derived_public_key = encode_ssh_ed25519_public_key(&signing_key.verifying_key()); + anyhow::ensure!( + self.public_key_ssh == derived_public_key, + "stored public key does not match the private key" + ); + Ok(()) + } + + pub(crate) fn signing_key(&self) -> Result { + signing_key_from_private_key_pkcs8_base64(&self.private_key_pkcs8_base64) + } +} + +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) -> Option { + // AgentAssertion is currently supported only for ChatGPT-backed Codex sessions. API-key + // sessions keep using their API key until the registration service supports API-key + // identity binding. + if !auth.is_chatgpt_auth() { + return None; + } + + let token_data = auth.get_token_data().ok()?; + let resolved_account_id = + forced_workspace_id + .filter(|value| !value.is_empty()) + .or(token_data + .account_id + .clone() + .filter(|value| !value.is_empty()))?; + + Some(Self { + binding_id: format!("chatgpt-account-{resolved_account_id}"), + chatgpt_account_id: resolved_account_id, + chatgpt_user_id: token_data + .id_token + .chatgpt_user_id + .filter(|value| !value.is_empty()), + access_token: token_data.access_token, + }) + } +} + +fn build_abom(session_source: SessionSource) -> AgentBillOfMaterials { + AgentBillOfMaterials { + agent_version: env!("CARGO_PKG_VERSION").to_string(), + agent_harness_id: match &session_source { + SessionSource::VSCode => "codex-app".to_string(), + SessionSource::Cli + | SessionSource::Exec + | SessionSource::Mcp + | SessionSource::Custom(_) + | SessionSource::SubAgent(_) + | SessionSource::Unknown => "codex-cli".to_string(), + }, + running_location: format!("{}-{}", session_source, std::env::consts::OS), + } +} + +fn generate_agent_key_material() -> Result { + let mut secret_key_bytes = [0u8; 32]; + OsRng + .try_fill_bytes(&mut secret_key_bytes) + .context("failed to generate agent identity private key bytes")?; + let signing_key = SigningKey::from_bytes(&secret_key_bytes); + let private_key_pkcs8 = signing_key + .to_pkcs8_der() + .context("failed to encode agent identity private key as PKCS#8")?; + + Ok(GeneratedAgentKeyMaterial { + private_key_pkcs8_base64: BASE64_STANDARD.encode(private_key_pkcs8.as_bytes()), + public_key_ssh: encode_ssh_ed25519_public_key(&signing_key.verifying_key()), + }) +} + +fn encode_ssh_ed25519_public_key(verifying_key: &VerifyingKey) -> String { + let mut blob = Vec::with_capacity(4 + 11 + 4 + 32); + append_ssh_string(&mut blob, b"ssh-ed25519"); + append_ssh_string(&mut blob, verifying_key.as_bytes()); + format!("ssh-ed25519 {}", BASE64_STANDARD.encode(blob)) +} + +fn append_ssh_string(buf: &mut Vec, value: &[u8]) { + buf.extend_from_slice(&(value.len() as u32).to_be_bytes()); + buf.extend_from_slice(value); +} + +fn agent_registration_url(chatgpt_base_url: &str) -> String { + let trimmed = chatgpt_base_url.trim_end_matches('/'); + format!("{trimmed}/v1/agent/register") +} + +fn signing_key_from_private_key_pkcs8_base64(private_key_pkcs8_base64: &str) -> Result { + let private_key = BASE64_STANDARD + .decode(private_key_pkcs8_base64) + .context("stored agent identity private key is not valid base64")?; + SigningKey::from_pkcs8_der(&private_key) + .context("stored agent identity private key is not valid PKCS#8") +} + +fn agent_identity_biscuit_url(chatgpt_base_url: &str) -> String { + let trimmed = chatgpt_base_url.trim_end_matches('/'); + format!("{trimmed}/authenticate_app_v2") +} + +fn agent_identity_request_id() -> Result { + let mut request_id_bytes = [0u8; 16]; + OsRng + .try_fill_bytes(&mut request_id_bytes) + .context("failed to generate agent identity request id")?; + Ok(format!( + "codex-agent-identity-{}", + URL_SAFE_NO_PAD.encode(request_id_bytes) + )) +} + +#[cfg(test)] +mod tests { + use super::*; + + use base64::engine::general_purpose::URL_SAFE_NO_PAD; + use codex_app_server_protocol::AuthMode as ApiAuthMode; + use codex_login::AuthCredentialsStoreMode; + use codex_login::AuthDotJson; + use codex_login::save_auth; + use codex_login::token_data::IdTokenInfo; + use codex_login::token_data::TokenData; + use pretty_assertions::assert_eq; + use wiremock::Mock; + use wiremock::MockServer; + use wiremock::ResponseTemplate; + use wiremock::matchers::header; + use wiremock::matchers::method; + use wiremock::matchers::path; + + #[tokio::test] + async fn ensure_registered_identity_skips_when_feature_is_disabled() { + let auth_manager = + AuthManager::from_auth_for_testing(make_chatgpt_auth("account-123", Some("user-123"))); + let manager = AgentIdentityManager::new_for_tests( + auth_manager, + /*feature_enabled*/ false, + "https://chatgpt.com/backend-api/".to_string(), + SessionSource::Cli, + ); + + assert_eq!(manager.ensure_registered_identity().await.unwrap(), None); + } + + #[tokio::test] + async fn ensure_registered_identity_skips_for_api_key_auth() { + let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("test-key")); + let manager = AgentIdentityManager::new_for_tests( + auth_manager, + /*feature_enabled*/ true, + "https://chatgpt.com/backend-api/".to_string(), + SessionSource::Cli, + ); + + assert_eq!(manager.ensure_registered_identity().await.unwrap(), None); + } + + #[tokio::test] + async fn ensure_registered_identity_registers_and_reuses_cached_identity() { + let server = MockServer::start().await; + let chatgpt_base_url = server.uri(); + mount_human_biscuit(&server, &chatgpt_base_url).await; + Mock::given(method("POST")) + .and(path("/v1/agent/register")) + .and(header("x-openai-authorization", "human-biscuit")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "agent_runtime_id": "agent-123", + }))) + .expect(1) + .mount(&server) + .await; + + let auth_manager = + AuthManager::from_auth_for_testing(make_chatgpt_auth("account-123", Some("user-123"))); + let manager = AgentIdentityManager::new_for_tests( + auth_manager, + /*feature_enabled*/ true, + chatgpt_base_url, + SessionSource::Cli, + ); + + let first = manager + .ensure_registered_identity() + .await + .unwrap() + .expect("identity should be registered"); + let second = manager + .ensure_registered_identity() + .await + .unwrap() + .expect("identity should be reused"); + + assert_eq!(first.agent_runtime_id, "agent-123"); + assert_eq!(first, second); + assert_eq!(first.abom.agent_harness_id, "codex-cli"); + assert_eq!(first.chatgpt_account_id, "account-123"); + assert_eq!(first.chatgpt_user_id.as_deref(), Some("user-123")); + } + + #[tokio::test] + async fn ensure_registered_identity_deletes_invalid_cached_identity_and_reregisters() { + let server = MockServer::start().await; + let chatgpt_base_url = server.uri(); + mount_human_biscuit(&server, &chatgpt_base_url).await; + Mock::given(method("POST")) + .and(path("/v1/agent/register")) + .and(header("x-openai-authorization", "human-biscuit")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "agent_runtime_id": "agent-456", + }))) + .expect(1) + .mount(&server) + .await; + + let auth = make_chatgpt_auth("account-123", Some("user-123")); + let auth_manager = AuthManager::from_auth_for_testing(auth.clone()); + let manager = AgentIdentityManager::new_for_tests( + auth_manager, + /*feature_enabled*/ true, + chatgpt_base_url, + SessionSource::Cli, + ); + + let binding = + AgentIdentityBinding::from_auth(&auth, /*forced_workspace_id*/ None).expect("binding"); + auth.set_agent_identity(AgentIdentityAuthRecord { + workspace_id: "account-123".to_string(), + agent_runtime_id: "agent_invalid".to_string(), + agent_private_key: "not-valid-base64".to_string(), + registered_at: "2026-01-01T00:00:00Z".to_string(), + }) + .expect("seed invalid identity"); + + let stored = manager + .ensure_registered_identity() + .await + .unwrap() + .expect("identity should be registered"); + + assert_eq!(stored.agent_runtime_id, "agent-456"); + let persisted = auth + .get_agent_identity(&binding.chatgpt_account_id) + .expect("stored identity"); + assert_eq!(persisted.agent_runtime_id, "agent-456"); + } + + #[tokio::test] + async fn ensure_registered_identity_uses_chatgpt_base_url() { + let server = MockServer::start().await; + let chatgpt_base_url = format!("{}/backend-api", server.uri()); + mount_human_biscuit(&server, &chatgpt_base_url).await; + Mock::given(method("POST")) + .and(path("/backend-api/v1/agent/register")) + .and(header("x-openai-authorization", "human-biscuit")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "agent_runtime_id": "agent_canonical", + }))) + .expect(1) + .mount(&server) + .await; + + let auth_manager = + AuthManager::from_auth_for_testing(make_chatgpt_auth("account-123", Some("user-123"))); + let manager = AgentIdentityManager::new_for_tests( + auth_manager, + /*feature_enabled*/ true, + chatgpt_base_url, + SessionSource::Cli, + ); + + let stored = manager + .ensure_registered_identity() + .await + .unwrap() + .expect("identity should be registered"); + assert_eq!(stored.agent_runtime_id, "agent_canonical"); + } + + async fn mount_human_biscuit(server: &MockServer, chatgpt_base_url: &str) { + let biscuit_url = agent_identity_biscuit_url(chatgpt_base_url); + let biscuit_path = reqwest::Url::parse(&biscuit_url) + .expect("biscuit URL parses") + .path() + .to_string(); + let target_url = agent_registration_url(chatgpt_base_url); + Mock::given(method("GET")) + .and(path(biscuit_path)) + .and(header("authorization", "Bearer access-token-account-123")) + .and(header("x-original-method", "GET")) + .and(header("x-original-url", target_url)) + .respond_with( + ResponseTemplate::new(200).insert_header("x-openai-authorization", "human-biscuit"), + ) + .expect(1) + .mount(server) + .await; + } + + #[test] + fn encode_ssh_ed25519_public_key_matches_expected_wire_shape() { + let key_material = generate_agent_key_material().expect("key material"); + let (_, encoded_blob) = key_material + .public_key_ssh + .split_once(' ') + .expect("public key contains scheme"); + let decoded = BASE64_STANDARD.decode(encoded_blob).expect("base64"); + + assert_eq!(&decoded[..4], 11u32.to_be_bytes().as_slice()); + assert_eq!(&decoded[4..15], b"ssh-ed25519"); + assert_eq!(&decoded[15..19], 32u32.to_be_bytes().as_slice()); + assert_eq!(decoded.len(), 51); + } + + fn make_chatgpt_auth(account_id: &str, user_id: Option<&str>) -> CodexAuth { + let tempdir = tempfile::tempdir().expect("tempdir"); + let auth_json = AuthDotJson { + auth_mode: Some(ApiAuthMode::Chatgpt), + openai_api_key: None, + tokens: Some(TokenData { + id_token: IdTokenInfo { + email: None, + chatgpt_plan_type: None, + chatgpt_user_id: user_id.map(ToOwned::to_owned), + chatgpt_account_id: Some(account_id.to_string()), + raw_jwt: fake_id_token(account_id, user_id), + }, + access_token: format!("access-token-{account_id}"), + refresh_token: "refresh-token".to_string(), + account_id: Some(account_id.to_string()), + }), + last_refresh: Some(Utc::now()), + agent_identity: None, + }; + save_auth(tempdir.path(), &auth_json, AuthCredentialsStoreMode::File).expect("save auth"); + CodexAuth::from_auth_storage(tempdir.path(), AuthCredentialsStoreMode::File) + .expect("load auth") + .expect("auth") + } + + fn fake_id_token(account_id: &str, user_id: Option<&str>) -> String { + let header = URL_SAFE_NO_PAD.encode(r#"{"alg":"none","typ":"JWT"}"#); + let payload = serde_json::json!({ + "https://api.openai.com/auth": { + "chatgpt_user_id": user_id, + "chatgpt_account_id": account_id, + } + }); + let payload = URL_SAFE_NO_PAD.encode(payload.to_string()); + format!("{header}.{payload}.signature") + } +} diff --git a/codex-rs/core/src/agent_identity/assertion.rs b/codex-rs/core/src/agent_identity/assertion.rs new file mode 100644 index 00000000000..6b8ead71d99 --- /dev/null +++ b/codex-rs/core/src/agent_identity/assertion.rs @@ -0,0 +1,270 @@ +use std::collections::BTreeMap; + +use anyhow::Context; +use anyhow::Result; +use base64::Engine as _; +use base64::engine::general_purpose::URL_SAFE_NO_PAD; +use ed25519_dalek::Signer as _; +use serde::Deserialize; +use serde::Serialize; +use thiserror::Error; +use tracing::debug; + +use super::*; + +#[derive(Debug, Error)] +#[error( + "agent task runtime {agent_runtime_id} does not match stored agent identity {stored_agent_runtime_id}" +)] +pub(crate) struct AgentTaskRuntimeMismatch { + pub(crate) agent_runtime_id: String, + pub(crate) task_id: String, + pub(crate) stored_agent_runtime_id: String, +} + +#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)] +pub(crate) struct AgentAssertionEnvelope { + pub(crate) agent_runtime_id: String, + pub(crate) task_id: String, + pub(crate) timestamp: String, + pub(crate) signature: String, +} + +impl AgentIdentityManager { + pub(crate) async fn authorization_header_for_task( + &self, + agent_task: &RegisteredAgentTask, + ) -> Result> { + if !self.feature_enabled { + return Ok(None); + } + + let Some(stored_identity) = self.ensure_registered_identity().await? else { + return Ok(None); + }; + if stored_identity.agent_runtime_id != agent_task.agent_runtime_id { + return Err(AgentTaskRuntimeMismatch { + agent_runtime_id: agent_task.agent_runtime_id.clone(), + task_id: agent_task.task_id.clone(), + stored_agent_runtime_id: stored_identity.agent_runtime_id, + } + .into()); + } + + let timestamp = Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true); + let envelope = AgentAssertionEnvelope { + agent_runtime_id: agent_task.agent_runtime_id.clone(), + task_id: agent_task.task_id.clone(), + timestamp: timestamp.clone(), + signature: sign_agent_assertion_payload(&stored_identity, agent_task, ×tamp)?, + }; + let serialized_assertion = serialize_agent_assertion(&envelope)?; + debug!( + agent_runtime_id = %envelope.agent_runtime_id, + task_id = %envelope.task_id, + "attaching agent assertion authorization to downstream request" + ); + Ok(Some(format!("AgentAssertion {serialized_assertion}"))) + } +} + +fn sign_agent_assertion_payload( + stored_identity: &StoredAgentIdentity, + agent_task: &RegisteredAgentTask, + timestamp: &str, +) -> Result { + let signing_key = stored_identity.signing_key()?; + let payload = format!( + "{}:{}:{timestamp}", + agent_task.agent_runtime_id, agent_task.task_id + ); + Ok(BASE64_STANDARD.encode(signing_key.sign(payload.as_bytes()).to_bytes())) +} + +fn serialize_agent_assertion(envelope: &AgentAssertionEnvelope) -> Result { + let payload = serde_json::to_vec(&BTreeMap::from([ + ("agent_runtime_id", envelope.agent_runtime_id.as_str()), + ("signature", envelope.signature.as_str()), + ("task_id", envelope.task_id.as_str()), + ("timestamp", envelope.timestamp.as_str()), + ])) + .context("failed to serialize agent assertion envelope")?; + Ok(URL_SAFE_NO_PAD.encode(payload)) +} + +#[cfg(test)] +mod tests { + use base64::engine::general_purpose::URL_SAFE_NO_PAD; + use ed25519_dalek::Signature; + use ed25519_dalek::Verifier as _; + use pretty_assertions::assert_eq; + + use super::*; + + #[tokio::test] + async fn authorization_header_for_task_skips_when_feature_is_disabled() { + let codex_home = tempfile::tempdir().expect("tempdir"); + let auth = make_chatgpt_auth(codex_home.path(), "account-123", Some("user-123")); + let auth_manager = AuthManager::from_auth_for_testing(auth); + let manager = AgentIdentityManager::new_for_tests( + auth_manager, + /*feature_enabled*/ false, + "https://chatgpt.com/backend-api/".to_string(), + SessionSource::Cli, + ); + let agent_task = RegisteredAgentTask { + agent_runtime_id: "agent-123".to_string(), + task_id: "task-123".to_string(), + registered_at: "2026-03-23T12:00:00Z".to_string(), + }; + + assert_eq!( + manager + .authorization_header_for_task(&agent_task) + .await + .unwrap(), + None + ); + } + + #[tokio::test] + async fn authorization_header_for_task_serializes_signed_agent_assertion() { + let codex_home = tempfile::tempdir().expect("tempdir"); + let auth = make_chatgpt_auth(codex_home.path(), "account-123", Some("user-123")); + let auth_manager = AuthManager::from_auth_for_testing(auth); + let manager = AgentIdentityManager::new_for_tests( + auth_manager, + /*feature_enabled*/ true, + "https://chatgpt.com/backend-api/".to_string(), + SessionSource::Cli, + ); + let stored_identity = manager + .seed_generated_identity_for_tests("agent-123") + .await + .expect("seed test identity"); + let agent_task = RegisteredAgentTask { + agent_runtime_id: "agent-123".to_string(), + task_id: "task-123".to_string(), + registered_at: "2026-03-23T12:00:00Z".to_string(), + }; + + let header = manager + .authorization_header_for_task(&agent_task) + .await + .expect("build agent assertion") + .expect("header should exist"); + let token = header + .strip_prefix("AgentAssertion ") + .expect("agent assertion scheme"); + let payload = URL_SAFE_NO_PAD + .decode(token) + .expect("valid base64url payload"); + let envelope: AgentAssertionEnvelope = + serde_json::from_slice(&payload).expect("valid assertion envelope"); + + assert_eq!( + envelope, + AgentAssertionEnvelope { + agent_runtime_id: "agent-123".to_string(), + task_id: "task-123".to_string(), + timestamp: envelope.timestamp.clone(), + signature: envelope.signature.clone(), + } + ); + let signature_bytes = BASE64_STANDARD + .decode(&envelope.signature) + .expect("valid base64 signature"); + let signature = Signature::from_slice(&signature_bytes).expect("valid signature bytes"); + let signing_key = stored_identity.signing_key().expect("signing key"); + signing_key + .verifying_key() + .verify( + format!( + "{}:{}:{}", + envelope.agent_runtime_id, envelope.task_id, envelope.timestamp + ) + .as_bytes(), + &signature, + ) + .expect("signature should verify"); + } + + #[tokio::test] + async fn authorization_header_for_task_reports_runtime_mismatch() { + let codex_home = tempfile::tempdir().expect("tempdir"); + let auth = make_chatgpt_auth(codex_home.path(), "account-123", Some("user-123")); + let auth_manager = AuthManager::from_auth_for_testing(auth); + let manager = AgentIdentityManager::new_for_tests( + auth_manager, + /*feature_enabled*/ true, + "https://chatgpt.com/backend-api/".to_string(), + SessionSource::Cli, + ); + manager + .seed_generated_identity_for_tests("agent-current") + .await + .expect("seed test identity"); + let agent_task = RegisteredAgentTask { + agent_runtime_id: "agent-stale".to_string(), + task_id: "task-123".to_string(), + registered_at: "2026-03-23T12:00:00Z".to_string(), + }; + + let error = manager + .authorization_header_for_task(&agent_task) + .await + .expect_err("stale task should be reported"); + let mismatch = error + .downcast_ref::() + .expect("runtime mismatch error"); + assert_eq!(mismatch.agent_runtime_id, "agent-stale"); + assert_eq!(mismatch.task_id, "task-123"); + assert_eq!(mismatch.stored_agent_runtime_id, "agent-current"); + } + + fn make_chatgpt_auth( + codex_home: &std::path::Path, + account_id: &str, + user_id: Option<&str>, + ) -> CodexAuth { + let auth_json = codex_login::AuthDotJson { + auth_mode: Some(codex_app_server_protocol::AuthMode::Chatgpt), + openai_api_key: None, + tokens: Some(codex_login::token_data::TokenData { + id_token: codex_login::token_data::IdTokenInfo { + email: None, + chatgpt_plan_type: None, + chatgpt_user_id: user_id.map(ToOwned::to_owned), + chatgpt_account_id: Some(account_id.to_string()), + raw_jwt: fake_id_token(account_id, user_id), + }, + access_token: format!("access-token-{account_id}"), + refresh_token: "refresh-token".to_string(), + account_id: Some(account_id.to_string()), + }), + last_refresh: Some(chrono::Utc::now()), + agent_identity: None, + }; + codex_login::save_auth( + codex_home, + &auth_json, + codex_login::AuthCredentialsStoreMode::File, + ) + .expect("save auth"); + CodexAuth::from_auth_storage(codex_home, codex_login::AuthCredentialsStoreMode::File) + .expect("load auth") + .expect("auth") + } + + fn fake_id_token(account_id: &str, user_id: Option<&str>) -> String { + let header = URL_SAFE_NO_PAD.encode(r#"{"alg":"none","typ":"JWT"}"#); + let payload = serde_json::json!({ + "https://api.openai.com/auth": { + "chatgpt_user_id": user_id, + "chatgpt_account_id": account_id, + } + }); + let payload = URL_SAFE_NO_PAD.encode(payload.to_string()); + format!("{header}.{payload}.signature") + } +} diff --git a/codex-rs/core/src/agent_identity/task_registration.rs b/codex-rs/core/src/agent_identity/task_registration.rs new file mode 100644 index 00000000000..9bbc364db3b --- /dev/null +++ b/codex-rs/core/src/agent_identity/task_registration.rs @@ -0,0 +1,456 @@ +use std::time::Duration; + +use anyhow::Context; +use anyhow::Result; +use codex_protocol::protocol::SessionAgentTask; +use crypto_box::SecretKey as Curve25519SecretKey; +use ed25519_dalek::Signer as _; +use serde::Deserialize; +use serde::Serialize; +use sha2::Digest as _; +use sha2::Sha512; +use tracing::info; + +use super::*; + +const AGENT_TASK_REGISTRATION_TIMEOUT: Duration = Duration::from_secs(15); + +#[derive(Clone, Debug, PartialEq, Eq)] +pub(crate) struct RegisteredAgentTask { + pub(crate) agent_runtime_id: String, + pub(crate) task_id: String, + pub(crate) registered_at: String, +} + +#[derive(Debug, Serialize)] +struct RegisterTaskRequest { + signature: String, + timestamp: String, +} + +#[derive(Debug, Deserialize)] +struct RegisterTaskResponse { + encrypted_task_id: String, +} + +impl AgentIdentityManager { + pub(crate) async fn register_task(&self) -> Result> { + if !self.feature_enabled { + return Ok(None); + } + + let Some((auth, binding)) = self.current_auth_binding().await else { + return Ok(None); + }; + + self.register_task_for_binding(auth, binding).await + } + + async fn register_task_for_binding( + &self, + auth: CodexAuth, + binding: AgentIdentityBinding, + ) -> Result> { + let stored_identity = self + .ensure_registered_identity_for_binding(&auth, &binding) + .await?; + + let timestamp = Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true); + let request_body = RegisterTaskRequest { + signature: sign_task_registration_payload(&stored_identity, ×tamp)?, + timestamp, + }; + + let client = create_client(); + let url = + agent_task_registration_url(&self.chatgpt_base_url, &stored_identity.agent_runtime_id); + let human_biscuit = self.mint_human_biscuit(&binding, &url).await?; + let response = client + .post(&url) + .header("X-OpenAI-Authorization", human_biscuit) + .json(&request_body) + .timeout(AGENT_TASK_REGISTRATION_TIMEOUT) + .send() + .await + .with_context(|| format!("failed to send agent task registration request to {url}"))?; + + if response.status().is_success() { + let response_body = response + .json::() + .await + .with_context(|| format!("failed to parse agent task response from {url}"))?; + let registered_task = RegisteredAgentTask { + agent_runtime_id: stored_identity.agent_runtime_id.clone(), + task_id: decrypt_task_id_response( + &stored_identity, + &response_body.encrypted_task_id, + )?, + registered_at: Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true), + }; + info!( + agent_runtime_id = %registered_task.agent_runtime_id, + task_id = %registered_task.task_id, + "registered agent task" + ); + return Ok(Some(registered_task)); + } + + let status = response.status(); + let body = response.text().await.unwrap_or_default(); + anyhow::bail!("agent task registration failed with status {status} from {url}: {body}") + } +} + +impl RegisteredAgentTask { + pub(crate) fn to_session_agent_task(&self) -> SessionAgentTask { + SessionAgentTask { + agent_runtime_id: self.agent_runtime_id.clone(), + task_id: self.task_id.clone(), + registered_at: self.registered_at.clone(), + } + } + + pub(crate) fn from_session_agent_task(task: SessionAgentTask) -> Self { + Self { + agent_runtime_id: task.agent_runtime_id, + task_id: task.task_id, + registered_at: task.registered_at, + } + } +} + +fn sign_task_registration_payload( + stored_identity: &StoredAgentIdentity, + timestamp: &str, +) -> Result { + let signing_key = stored_identity.signing_key()?; + let payload = format!("{}:{timestamp}", stored_identity.agent_runtime_id); + Ok(BASE64_STANDARD.encode(signing_key.sign(payload.as_bytes()).to_bytes())) +} + +fn decrypt_task_id_response( + stored_identity: &StoredAgentIdentity, + encrypted_task_id: &str, +) -> Result { + let signing_key = stored_identity.signing_key()?; + let ciphertext = BASE64_STANDARD + .decode(encrypted_task_id) + .context("encrypted task id is not valid base64")?; + let plaintext = curve25519_secret_key_from_signing_key(&signing_key) + .unseal(&ciphertext) + .map_err(|_| anyhow::anyhow!("failed to decrypt encrypted task id"))?; + String::from_utf8(plaintext).context("decrypted task id is not valid UTF-8") +} + +fn curve25519_secret_key_from_signing_key(signing_key: &SigningKey) -> Curve25519SecretKey { + let digest = Sha512::digest(signing_key.to_bytes()); + let mut secret_key = [0u8; 32]; + secret_key.copy_from_slice(&digest[..32]); + secret_key[0] &= 248; + secret_key[31] &= 127; + secret_key[31] |= 64; + Curve25519SecretKey::from(secret_key) +} + +fn agent_task_registration_url(chatgpt_base_url: &str, agent_runtime_id: &str) -> String { + let trimmed = chatgpt_base_url.trim_end_matches('/'); + format!("{trimmed}/v1/agent/{agent_runtime_id}/task/register") +} + +#[cfg(test)] +mod tests { + use base64::engine::general_purpose::URL_SAFE_NO_PAD; + use codex_app_server_protocol::AuthMode as ApiAuthMode; + use codex_login::AuthCredentialsStoreMode; + use codex_login::AuthDotJson; + use codex_login::save_auth; + use codex_login::token_data::IdTokenInfo; + use codex_login::token_data::TokenData; + use pretty_assertions::assert_eq; + use wiremock::Mock; + use wiremock::MockServer; + use wiremock::ResponseTemplate; + use wiremock::matchers::header; + use wiremock::matchers::method; + use wiremock::matchers::path; + + use super::*; + + #[tokio::test] + async fn register_task_skips_when_feature_is_disabled() { + let auth_manager = + AuthManager::from_auth_for_testing(make_chatgpt_auth("account-123", Some("user-123"))); + let manager = AgentIdentityManager::new_for_tests( + auth_manager, + /*feature_enabled*/ false, + "https://chatgpt.com/backend-api/".to_string(), + SessionSource::Cli, + ); + + assert_eq!(manager.register_task().await.unwrap(), None); + } + + #[tokio::test] + async fn register_task_skips_for_api_key_auth() { + let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("test-key")); + let manager = AgentIdentityManager::new_for_tests( + auth_manager, + /*feature_enabled*/ true, + "https://chatgpt.com/backend-api/".to_string(), + SessionSource::Cli, + ); + + assert_eq!(manager.register_task().await.unwrap(), None); + } + + #[tokio::test] + async fn register_task_registers_and_decrypts_plaintext_task_id() { + let server = MockServer::start().await; + let chatgpt_base_url = server.uri(); + mount_human_biscuit(&server, &chatgpt_base_url, "agent-123").await; + let auth = make_chatgpt_auth("account-123", Some("user-123")); + let auth_manager = AuthManager::from_auth_for_testing(auth.clone()); + let manager = AgentIdentityManager::new_for_tests( + auth_manager, + /*feature_enabled*/ true, + chatgpt_base_url, + SessionSource::Cli, + ); + let stored_identity = seed_stored_identity(&manager, &auth, "agent-123", "account-123"); + let encrypted_task_id = + encrypt_task_id_for_identity(&stored_identity, "task_123").expect("task ciphertext"); + + Mock::given(method("POST")) + .and(path("/v1/agent/agent-123/task/register")) + .and(header("x-openai-authorization", "human-biscuit")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "encrypted_task_id": encrypted_task_id, + }))) + .expect(1) + .mount(&server) + .await; + + let task = manager + .register_task() + .await + .unwrap() + .expect("task should be registered"); + + assert_eq!( + task, + RegisteredAgentTask { + agent_runtime_id: "agent-123".to_string(), + task_id: "task_123".to_string(), + registered_at: task.registered_at.clone(), + } + ); + } + + #[tokio::test] + async fn register_task_uses_chatgpt_base_url() { + let server = MockServer::start().await; + let chatgpt_base_url = format!("{}/backend-api", server.uri()); + mount_human_biscuit(&server, &chatgpt_base_url, "agent-fallback").await; + let auth = make_chatgpt_auth("account-123", Some("user-123")); + let auth_manager = AuthManager::from_auth_for_testing(auth.clone()); + let manager = AgentIdentityManager::new_for_tests( + auth_manager, + /*feature_enabled*/ true, + chatgpt_base_url, + SessionSource::Cli, + ); + let stored_identity = + seed_stored_identity(&manager, &auth, "agent-fallback", "account-123"); + let encrypted_task_id = encrypt_task_id_for_identity(&stored_identity, "task_fallback") + .expect("task ciphertext"); + + Mock::given(method("POST")) + .and(path("/backend-api/v1/agent/agent-fallback/task/register")) + .and(header("x-openai-authorization", "human-biscuit")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "encrypted_task_id": encrypted_task_id, + }))) + .expect(1) + .mount(&server) + .await; + + let task = manager + .register_task() + .await + .unwrap() + .expect("task should be registered"); + + assert_eq!(task.agent_runtime_id, "agent-fallback"); + assert_eq!(task.task_id, "task_fallback"); + } + + #[tokio::test] + async fn register_task_for_binding_keeps_one_auth_snapshot() { + let server = MockServer::start().await; + let chatgpt_base_url = server.uri(); + mount_human_biscuit(&server, &chatgpt_base_url, "agent-123").await; + let binding_auth = make_chatgpt_auth("account-123", Some("user-123")); + let auth_manager = + AuthManager::from_auth_for_testing(make_chatgpt_auth("account-456", Some("user-456"))); + let manager = AgentIdentityManager::new_for_tests( + auth_manager, + /*feature_enabled*/ true, + chatgpt_base_url, + SessionSource::Cli, + ); + let stored_identity = + seed_stored_identity(&manager, &binding_auth, "agent-123", "account-123"); + let encrypted_task_id = + encrypt_task_id_for_identity(&stored_identity, "task_123").expect("task ciphertext"); + let binding = + AgentIdentityBinding::from_auth(&binding_auth, /*forced_workspace_id*/ None) + .expect("binding"); + + Mock::given(method("POST")) + .and(path("/v1/agent/agent-123/task/register")) + .and(header("x-openai-authorization", "human-biscuit")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "encrypted_task_id": encrypted_task_id, + }))) + .expect(1) + .mount(&server) + .await; + + let task = manager + .register_task_for_binding(binding_auth, binding) + .await + .unwrap() + .expect("task should be registered"); + + assert_eq!( + task, + RegisteredAgentTask { + agent_runtime_id: "agent-123".to_string(), + task_id: "task_123".to_string(), + registered_at: task.registered_at.clone(), + } + ); + } + + #[tokio::test] + async fn task_matches_current_identity_rejects_stale_registered_identity() { + let auth_manager = + AuthManager::from_auth_for_testing(make_chatgpt_auth("account-456", Some("user-456"))); + let manager = AgentIdentityManager::new_for_tests( + auth_manager, + /*feature_enabled*/ true, + "https://chatgpt.com/backend-api/".to_string(), + SessionSource::Cli, + ); + let task = RegisteredAgentTask { + agent_runtime_id: "agent-123".to_string(), + task_id: "task_123".to_string(), + registered_at: "2026-03-23T12:00:00Z".to_string(), + }; + + assert!(!manager.task_matches_current_identity(&task).await); + } + + async fn mount_human_biscuit( + server: &MockServer, + chatgpt_base_url: &str, + agent_runtime_id: &str, + ) { + let biscuit_url = agent_identity_biscuit_url(chatgpt_base_url); + let biscuit_path = reqwest::Url::parse(&biscuit_url) + .expect("biscuit URL parses") + .path() + .to_string(); + let target_url = agent_task_registration_url(chatgpt_base_url, agent_runtime_id); + Mock::given(method("GET")) + .and(path(biscuit_path)) + .and(header("authorization", "Bearer access-token-account-123")) + .and(header("x-original-method", "GET")) + .and(header("x-original-url", target_url)) + .respond_with( + ResponseTemplate::new(200).insert_header("x-openai-authorization", "human-biscuit"), + ) + .expect(1) + .mount(server) + .await; + } + + fn seed_stored_identity( + manager: &AgentIdentityManager, + auth: &CodexAuth, + agent_runtime_id: &str, + account_id: &str, + ) -> StoredAgentIdentity { + let key_material = generate_agent_key_material().expect("key material"); + let binding = AgentIdentityBinding::from_auth(auth, None).expect("binding"); + let stored_identity = StoredAgentIdentity { + binding_id: binding.binding_id, + chatgpt_account_id: account_id.to_string(), + chatgpt_user_id: Some("user-123".to_string()), + agent_runtime_id: agent_runtime_id.to_string(), + private_key_pkcs8_base64: key_material.private_key_pkcs8_base64, + public_key_ssh: key_material.public_key_ssh, + registered_at: Utc::now().to_rfc3339_opts(SecondsFormat::Secs, true), + abom: manager.abom.clone(), + }; + manager + .store_identity(auth, &stored_identity) + .expect("store identity"); + let persisted = auth + .get_agent_identity(account_id) + .expect("persisted identity"); + assert_eq!(persisted.agent_runtime_id, agent_runtime_id); + stored_identity + } + + fn encrypt_task_id_for_identity( + stored_identity: &StoredAgentIdentity, + task_id: &str, + ) -> Result { + let mut rng = crypto_box::aead::OsRng; + let public_key = + curve25519_secret_key_from_signing_key(&stored_identity.signing_key()?).public_key(); + let ciphertext = public_key + .seal(&mut rng, task_id.as_bytes()) + .map_err(|_| anyhow::anyhow!("failed to encrypt test task id"))?; + Ok(BASE64_STANDARD.encode(ciphertext)) + } + + fn make_chatgpt_auth(account_id: &str, user_id: Option<&str>) -> CodexAuth { + let tempdir = tempfile::tempdir().expect("tempdir"); + let auth_json = AuthDotJson { + auth_mode: Some(ApiAuthMode::Chatgpt), + openai_api_key: None, + tokens: Some(TokenData { + id_token: IdTokenInfo { + email: None, + chatgpt_plan_type: None, + chatgpt_user_id: user_id.map(ToOwned::to_owned), + chatgpt_account_id: Some(account_id.to_string()), + raw_jwt: fake_id_token(account_id, user_id), + }, + access_token: format!("access-token-{account_id}"), + refresh_token: "refresh-token".to_string(), + account_id: Some(account_id.to_string()), + }), + last_refresh: Some(Utc::now()), + agent_identity: None, + }; + save_auth(tempdir.path(), &auth_json, AuthCredentialsStoreMode::File).expect("save auth"); + CodexAuth::from_auth_storage(tempdir.path(), AuthCredentialsStoreMode::File) + .expect("load auth") + .expect("auth") + } + + fn fake_id_token(account_id: &str, user_id: Option<&str>) -> String { + let header = URL_SAFE_NO_PAD.encode(r#"{"alg":"none","typ":"JWT"}"#); + let payload = serde_json::json!({ + "https://api.openai.com/auth": { + "chatgpt_user_id": user_id, + "chatgpt_account_id": account_id, + } + }); + let payload = URL_SAFE_NO_PAD.encode(payload.to_string()); + format!("{header}.{payload}.signature") + } +} diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 16f743943a0..7a24f437656 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -31,6 +31,9 @@ use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicU64; use std::sync::atomic::Ordering; +use crate::agent_identity::AgentIdentityManager; +use crate::agent_identity::AgentTaskRuntimeMismatch; +use crate::agent_identity::RegisteredAgentTask; use codex_api::ApiError; use codex_api::CompactClient as ApiCompactClient; use codex_api::CompactionInput as ApiCompactionInput; @@ -92,6 +95,7 @@ use tokio::sync::oneshot; use tokio::sync::oneshot::error::TryRecvError; use tokio_tungstenite::tungstenite::Error; use tokio_tungstenite::tungstenite::Message; +use tracing::debug; use tracing::instrument; use tracing::trace; use tracing::warn; @@ -144,6 +148,7 @@ pub(crate) const WEBSOCKET_CONNECT_TIMEOUT: Duration = #[derive(Debug)] struct ModelClientState { auth_manager: Option>, + agent_identity_manager: Option>, conversation_id: ThreadId, window_generation: AtomicU64, installation_id: String, @@ -155,7 +160,7 @@ struct ModelClientState { include_timing_metrics: bool, beta_features_header: Option, disable_websockets: AtomicBool, - cached_websocket_session: StdMutex, + cached_websocket_session: StdMutex, } /// Resolved API client setup for a single request attempt. @@ -211,6 +216,8 @@ pub struct ModelClient { pub struct ModelClientSession { client: ModelClient, websocket_session: WebsocketSession, + agent_task: Option, + cache_websocket_session_on_drop: bool, /// Turn state for sticky routing. /// /// This is an `OnceLock` that stores the turn state value received from the server @@ -238,6 +245,12 @@ struct WebsocketSession { connection_reused: StdMutex, } +#[derive(Debug, Default)] +struct CachedWebsocketSession { + agent_task: Option, + websocket_session: WebsocketSession, +} + impl WebsocketSession { fn set_connection_reused(&self, connection_reused: bool) { *self @@ -306,6 +319,33 @@ impl ModelClient { enable_request_compression: bool, include_timing_metrics: bool, beta_features_header: Option, + ) -> Self { + Self::new_with_agent_identity_manager( + auth_manager, + /*agent_identity_manager*/ None, + conversation_id, + installation_id, + provider, + session_source, + model_verbosity, + enable_request_compression, + include_timing_metrics, + beta_features_header, + ) + } + + #[allow(clippy::too_many_arguments)] + pub(crate) fn new_with_agent_identity_manager( + auth_manager: Option>, + agent_identity_manager: Option>, + conversation_id: ThreadId, + installation_id: String, + provider: ModelProviderInfo, + session_source: SessionSource, + model_verbosity: Option, + enable_request_compression: bool, + include_timing_metrics: bool, + beta_features_header: Option, ) -> Self { let auth_manager = auth_manager_for_provider(auth_manager, &provider); let codex_api_key_env_enabled = auth_manager @@ -315,6 +355,7 @@ impl ModelClient { Self { state: Arc::new(ModelClientState { auth_manager, + agent_identity_manager, conversation_id, window_generation: AtomicU64::new(0), installation_id, @@ -326,7 +367,7 @@ impl ModelClient { include_timing_metrics, beta_features_header, disable_websockets: AtomicBool::new(false), - cached_websocket_session: StdMutex::new(WebsocketSession::default()), + cached_websocket_session: StdMutex::new(CachedWebsocketSession::default()), }), } } @@ -336,9 +377,22 @@ impl ModelClient { /// This constructor does not perform network I/O itself; the session opens a websocket lazily /// when the first stream request is issued. pub fn new_session(&self) -> ModelClientSession { + self.new_session_with_agent_task(/*agent_task*/ None) + } + + pub(crate) fn new_session_with_agent_task( + &self, + agent_task: Option, + ) -> ModelClientSession { + // WebSocket auth is bound to the task that opened the connection. Reuse only when the + // cached connection was created for the same task, and drop mismatched taskless/task-scoped + // sessions rather than mixing auth contexts. + let websocket_session = self.take_cached_websocket_session(agent_task.as_ref()); ModelClientSession { client: self.clone(), - websocket_session: self.take_cached_websocket_session(), + websocket_session, + agent_task, + cache_websocket_session_on_drop: true, turn_state: Arc::new(OnceLock::new()), } } @@ -351,12 +405,12 @@ impl ModelClient { self.state .window_generation .store(window_generation, Ordering::Relaxed); - self.store_cached_websocket_session(WebsocketSession::default()); + self.clear_cached_websocket_session(); } pub(crate) fn advance_window_generation(&self) { self.state.window_generation.fetch_add(1, Ordering::Relaxed); - self.store_cached_websocket_session(WebsocketSession::default()); + self.clear_cached_websocket_session(); } fn current_window_id(&self) -> String { @@ -365,21 +419,44 @@ impl ModelClient { format!("{conversation_id}:{window_generation}") } - fn take_cached_websocket_session(&self) -> WebsocketSession { + fn take_cached_websocket_session( + &self, + agent_task: Option<&RegisteredAgentTask>, + ) -> WebsocketSession { let mut cached_websocket_session = self .state .cached_websocket_session .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); - std::mem::take(&mut *cached_websocket_session) + if cached_websocket_session.agent_task.as_ref() == agent_task { + return std::mem::take(&mut *cached_websocket_session).websocket_session; + } + + *cached_websocket_session = CachedWebsocketSession::default(); + WebsocketSession::default() } - fn store_cached_websocket_session(&self, websocket_session: WebsocketSession) { + fn store_cached_websocket_session( + &self, + agent_task: Option, + websocket_session: WebsocketSession, + ) { *self .state .cached_websocket_session .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) = websocket_session; + .unwrap_or_else(std::sync::PoisonError::into_inner) = CachedWebsocketSession { + agent_task, + websocket_session, + }; + } + + fn clear_cached_websocket_session(&self) { + *self + .state + .cached_websocket_session + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = CachedWebsocketSession::default(); } pub(crate) fn force_http_fallback( @@ -399,7 +476,7 @@ impl ModelClient { ); } - self.store_cached_websocket_session(WebsocketSession::default()); + self.clear_cached_websocket_session(); activated } @@ -421,7 +498,7 @@ impl ModelClient { if prompt.input.is_empty() { return Ok(Vec::new()); } - let client_setup = self.current_client_setup().await?; + let client_setup = self.current_client_setup(/*agent_task*/ None).await?; let transport = ReqwestTransport::new(build_reqwest_client()); let request_telemetry = Self::build_request_telemetry( session_telemetry, @@ -485,7 +562,7 @@ impl ModelClient { ) -> Result { // Create the media call over HTTP first, then retain matching auth so realtime can attach // the server-side control WebSocket to the call id from that HTTP response. - let client_setup = self.current_client_setup().await?; + let client_setup = self.current_client_setup(/*agent_task*/ None).await?; let mut sideband_headers = extra_headers.clone(); sideband_headers.extend(sideband_websocket_auth_headers(&client_setup.api_auth)); let transport = ReqwestTransport::new(build_reqwest_client()); @@ -518,7 +595,7 @@ impl ModelClient { return Ok(Vec::new()); } - let client_setup = self.current_client_setup().await?; + let client_setup = self.current_client_setup(/*agent_task*/ None).await?; let transport = ReqwestTransport::new(build_reqwest_client()); let request_telemetry = Self::build_request_telemetry( session_telemetry, @@ -659,7 +736,10 @@ impl ModelClient { /// /// This centralizes setup used by both prewarm and normal request paths so they stay in /// lockstep when auth/provider resolution changes. - async fn current_client_setup(&self) -> Result { + async fn current_client_setup( + &self, + agent_task: Option<&RegisteredAgentTask>, + ) -> Result { let auth = match self.state.auth_manager.as_ref() { Some(manager) => manager.auth().await, None => None, @@ -668,7 +748,42 @@ impl ModelClient { .state .provider .to_api_provider(auth.as_ref().map(CodexAuth::auth_mode))?; - let api_auth = auth_provider_from_auth(auth.clone(), &self.state.provider)?; + let api_auth = match (agent_task, self.state.agent_identity_manager.as_ref()) { + (Some(agent_task), Some(agent_identity_manager)) => { + if let Some(authorization_header_value) = agent_identity_manager + .authorization_header_for_task(agent_task) + .await + .map_err(|err| { + if let Some(mismatch) = err.downcast_ref::() { + debug!( + agent_runtime_id = %mismatch.agent_runtime_id, + task_id = %mismatch.task_id, + stored_agent_runtime_id = %mismatch.stored_agent_runtime_id, + "agent task no longer matches stored identity" + ); + return CodexErr::AgentTaskStale; + } + CodexErr::Stream( + format!("failed to build agent assertion authorization: {err}"), + None, + ) + })? + { + debug!( + agent_runtime_id = %agent_task.agent_runtime_id, + task_id = %agent_task.task_id, + "using agent assertion authorization for downstream request" + ); + CoreAuthProvider::from_authorization_header_value( + Some(authorization_header_value), + /*account_id*/ None, + ) + } else { + auth_provider_from_auth(auth.clone(), &self.state.provider)? + } + } + _ => auth_provider_from_auth(auth.clone(), &self.state.provider)?, + }; Ok(CurrentClientSetup { auth, api_provider, @@ -802,12 +917,22 @@ impl ModelClient { impl Drop for ModelClientSession { fn drop(&mut self) { let websocket_session = std::mem::take(&mut self.websocket_session); - self.client - .store_cached_websocket_session(websocket_session); + if self.cache_websocket_session_on_drop { + self.client + .store_cached_websocket_session(self.agent_task.clone(), websocket_session); + } } } impl ModelClientSession { + pub(crate) fn agent_task(&self) -> Option<&RegisteredAgentTask> { + self.agent_task.as_ref() + } + + pub(crate) fn disable_cached_websocket_session_on_drop(&mut self) { + self.cache_websocket_session_on_drop = false; + } + pub(crate) fn reset_websocket_session(&mut self) { self.websocket_session.connection = None; self.websocket_session.last_request = None; @@ -1009,11 +1134,15 @@ impl ModelClientSession { return Ok(()); } - let client_setup = self.client.current_client_setup().await.map_err(|err| { - ApiError::Stream(format!( - "failed to build websocket prewarm client setup: {err}" - )) - })?; + let client_setup = self + .client + .current_client_setup(self.agent_task.as_ref()) + .await + .map_err(|err| { + ApiError::Stream(format!( + "failed to build websocket prewarm client setup: {err}" + )) + })?; let auth_context = AuthRequestTelemetryContext::new( client_setup.auth.as_ref().map(CodexAuth::auth_mode), &client_setup.api_auth, @@ -1167,7 +1296,10 @@ impl ModelClientSession { .map(AuthManager::unauthorized_recovery); let mut pending_retry = PendingUnauthorizedRetry::default(); loop { - let client_setup = self.client.current_client_setup().await?; + let client_setup = self + .client + .current_client_setup(self.agent_task.as_ref()) + .await?; let transport = ReqwestTransport::new(build_reqwest_client()); let request_auth_context = AuthRequestTelemetryContext::new( client_setup.auth.as_ref().map(CodexAuth::auth_mode), @@ -1256,7 +1388,10 @@ impl ModelClientSession { .map(AuthManager::unauthorized_recovery); let mut pending_retry = PendingUnauthorizedRetry::default(); loop { - let client_setup = self.client.current_client_setup().await?; + let client_setup = self + .client + .current_client_setup(self.agent_task.as_ref()) + .await?; let request_auth_context = AuthRequestTelemetryContext::new( client_setup.auth.as_ref().map(CodexAuth::auth_mode), &client_setup.api_auth, diff --git a/codex-rs/core/src/client_tests.rs b/codex-rs/core/src/client_tests.rs index 2fd5f04f951..9f83df60d1c 100644 --- a/codex-rs/core/src/client_tests.rs +++ b/codex-rs/core/src/client_tests.rs @@ -1,3 +1,5 @@ +use std::sync::Arc; + use super::AuthRequestTelemetryContext; use super::ModelClient; use super::PendingUnauthorizedRetry; @@ -7,17 +9,43 @@ use super::X_CODEX_PARENT_THREAD_ID_HEADER; use super::X_CODEX_TURN_METADATA_HEADER; use super::X_CODEX_WINDOW_ID_HEADER; use super::X_OPENAI_SUBAGENT_HEADER; +use crate::Prompt; +use crate::ResponseEvent; +use crate::agent_identity::AgentAssertionEnvelope; +use crate::agent_identity::AgentIdentityManager; +use crate::agent_identity::RegisteredAgentTask; +use crate::agent_identity::StoredAgentIdentity; +use base64::Engine as _; +use base64::engine::general_purpose::URL_SAFE_NO_PAD; +use chrono::Utc; use codex_api::CoreAuthProvider; use codex_app_server_protocol::AuthMode; +use codex_login::AuthCredentialsStoreMode; +use codex_login::AuthDotJson; +use codex_login::AuthManager; +use codex_login::CodexAuth; +use codex_login::save_auth; +use codex_login::token_data::IdTokenInfo; +use codex_login::token_data::TokenData; +use codex_model_provider_info::ModelProviderInfo; use codex_model_provider_info::WireApi; use codex_model_provider_info::create_oss_provider_with_base_url; use codex_otel::SessionTelemetry; use codex_protocol::ThreadId; +use codex_protocol::config_types::ReasoningSummary; +use codex_protocol::error::CodexErr; +use codex_protocol::models::ContentItem; +use codex_protocol::models::ResponseItem; use codex_protocol::openai_models::ModelInfo; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::SubAgentSource; +use core_test_support::responses; +use ed25519_dalek::Signature; +use ed25519_dalek::Verifier as _; +use futures::StreamExt; use pretty_assertions::assert_eq; use serde_json::json; +use tempfile::TempDir; fn test_model_client(session_source: SessionSource) -> ModelClient { let provider = create_oss_provider_with_base_url("https://example.com/v1", WireApi::Responses); @@ -79,6 +107,152 @@ fn test_session_telemetry() -> SessionTelemetry { ) } +fn test_prompt(text: &str) -> Prompt { + Prompt { + input: vec![ResponseItem::Message { + id: None, + role: "user".into(), + content: vec![ContentItem::InputText { + text: text.to_string(), + }], + end_turn: None, + phase: None, + }], + ..Prompt::default() + } +} + +async fn drain_stream_to_completion(stream: &mut crate::ResponseStream) -> anyhow::Result<()> { + while let Some(event) = stream.next().await { + if matches!(event?, ResponseEvent::Completed { .. }) { + break; + } + } + Ok(()) +} + +async fn model_client_with_agent_task( + provider: ModelProviderInfo, +) -> ( + TempDir, + ModelClient, + RegisteredAgentTask, + StoredAgentIdentity, +) { + let codex_home = tempfile::tempdir().expect("tempdir"); + let auth = make_chatgpt_auth(codex_home.path(), "account-123", Some("user-123")); + 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, + "https://chatgpt.com/backend-api/".to_string(), + SessionSource::Cli, + )); + let stored_identity = agent_identity_manager + .seed_generated_identity_for_tests("agent-123") + .await + .expect("seed test identity"); + let agent_task = RegisteredAgentTask { + agent_runtime_id: stored_identity.agent_runtime_id.clone(), + task_id: "task-123".to_string(), + registered_at: "2026-03-23T12:00:00Z".to_string(), + }; + let client = ModelClient::new_with_agent_identity_manager( + Some(auth_manager), + Some(agent_identity_manager), + ThreadId::new(), + /*installation_id*/ "11111111-1111-4111-8111-111111111111".to_string(), + provider, + SessionSource::Cli, + /*model_verbosity*/ None, + /*enable_request_compression*/ false, + /*include_timing_metrics*/ false, + /*beta_features_header*/ None, + ); + (codex_home, client, agent_task, stored_identity) +} + +fn make_chatgpt_auth( + codex_home: &std::path::Path, + account_id: &str, + user_id: Option<&str>, +) -> CodexAuth { + let auth_json = AuthDotJson { + auth_mode: Some(AuthMode::Chatgpt), + openai_api_key: None, + tokens: Some(TokenData { + id_token: IdTokenInfo { + email: None, + chatgpt_plan_type: None, + chatgpt_user_id: user_id.map(ToOwned::to_owned), + chatgpt_account_id: Some(account_id.to_string()), + raw_jwt: fake_id_token(account_id, user_id), + }, + access_token: format!("access-token-{account_id}"), + refresh_token: "refresh-token".to_string(), + account_id: Some(account_id.to_string()), + }), + last_refresh: Some(Utc::now()), + agent_identity: None, + }; + save_auth(codex_home, &auth_json, AuthCredentialsStoreMode::File).expect("save auth"); + CodexAuth::from_auth_storage(codex_home, AuthCredentialsStoreMode::File) + .expect("load auth") + .expect("auth") +} + +fn fake_id_token(account_id: &str, user_id: Option<&str>) -> String { + let header = URL_SAFE_NO_PAD.encode(r#"{"alg":"none","typ":"JWT"}"#); + let payload = serde_json::json!({ + "https://api.openai.com/auth": { + "chatgpt_user_id": user_id, + "chatgpt_account_id": account_id, + } + }); + let payload = URL_SAFE_NO_PAD.encode(payload.to_string()); + format!("{header}.{payload}.signature") +} + +fn assert_agent_assertion_header( + authorization_header: &str, + stored_identity: &StoredAgentIdentity, + expected_agent_runtime_id: &str, + expected_task_id: &str, +) { + let token = authorization_header + .strip_prefix("AgentAssertion ") + .expect("agent assertion authorization scheme"); + let envelope: AgentAssertionEnvelope = serde_json::from_slice( + &URL_SAFE_NO_PAD + .decode(token) + .expect("base64url-encoded agent assertion"), + ) + .expect("valid agent assertion envelope"); + + assert_eq!(envelope.agent_runtime_id, expected_agent_runtime_id); + assert_eq!(envelope.task_id, expected_task_id); + + let signature = Signature::from_slice( + &base64::engine::general_purpose::STANDARD + .decode(&envelope.signature) + .expect("base64 signature"), + ) + .expect("signature bytes"); + stored_identity + .signing_key() + .expect("signing key") + .verifying_key() + .verify( + format!( + "{}:{}:{}", + envelope.agent_runtime_id, envelope.task_id, envelope.timestamp + ) + .as_bytes(), + &signature, + ) + .expect("signature should verify"); +} + #[test] fn build_subagent_headers_sets_other_subagent_label() { let client = test_model_client(SessionSource::SubAgent(SubAgentSource::Other( @@ -169,3 +343,236 @@ fn auth_request_telemetry_context_tracks_attached_auth_and_retry_phase() { assert_eq!(auth_context.recovery_mode, Some("managed")); assert_eq!(auth_context.recovery_phase, Some("refresh_token")); } + +#[tokio::test] +async fn responses_http_uses_agent_assertion_when_agent_task_is_present() { + core_test_support::skip_if_no_network!(); + + let server = responses::start_mock_server().await; + let request_recorder = responses::mount_sse_once( + &server, + responses::sse(vec![ + responses::ev_response_created("resp-1"), + responses::ev_completed("resp-1"), + ]), + ) + .await; + let provider = + create_oss_provider_with_base_url(&format!("{}/v1", server.uri()), WireApi::Responses); + let (_codex_home, client, agent_task, stored_identity) = + model_client_with_agent_task(provider).await; + let model_info = test_model_info(); + let session_telemetry = test_session_telemetry(); + let mut client_session = client.new_session_with_agent_task(Some(agent_task.clone())); + + let mut stream = client_session + .stream( + &test_prompt("hello"), + &model_info, + &session_telemetry, + /*effort*/ None, + ReasoningSummary::Auto, + /*service_tier*/ None, + /*turn_metadata_header*/ None, + ) + .await + .expect("stream request should succeed"); + drain_stream_to_completion(&mut stream) + .await + .expect("stream should complete"); + + let request = request_recorder.single_request(); + let authorization = request + .header("authorization") + .expect("authorization header should be present"); + assert_agent_assertion_header( + &authorization, + &stored_identity, + &agent_task.agent_runtime_id, + &agent_task.task_id, + ); + assert_eq!(request.header("chatgpt-account-id"), None); +} + +#[tokio::test] +async fn responses_http_reports_stale_agent_task_when_identity_changed() { + let provider = create_oss_provider_with_base_url("https://example.com/v1", WireApi::Responses); + let (_codex_home, client, mut agent_task, _stored_identity) = + model_client_with_agent_task(provider).await; + agent_task.agent_runtime_id = "agent-stale".to_string(); + let model_info = test_model_info(); + let session_telemetry = test_session_telemetry(); + let mut client_session = client.new_session_with_agent_task(Some(agent_task)); + + let error = match client_session + .stream( + &test_prompt("hello"), + &model_info, + &session_telemetry, + /*effort*/ None, + ReasoningSummary::Auto, + /*service_tier*/ None, + /*turn_metadata_header*/ None, + ) + .await + { + Ok(_) => panic!("stale task should be reported before sending a request"), + Err(error) => error, + }; + + assert!(matches!(error, CodexErr::AgentTaskStale)); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn websocket_agent_task_bypasses_cached_bearer_prewarm() { + core_test_support::skip_if_no_network!(); + + let server = responses::start_websocket_server(vec![ + vec![vec![ + responses::ev_response_created("resp-prewarm"), + responses::ev_completed("resp-prewarm"), + ]], + vec![vec![ + responses::ev_response_created("resp-1"), + responses::ev_completed("resp-1"), + ]], + ]) + .await; + let mut provider = + create_oss_provider_with_base_url(&format!("{}/v1", server.uri()), WireApi::Responses); + provider.supports_websockets = true; + provider.websocket_connect_timeout_ms = Some(5_000); + let (_codex_home, client, agent_task, stored_identity) = + model_client_with_agent_task(provider).await; + let model_info = test_model_info(); + let session_telemetry = test_session_telemetry(); + let prompt = test_prompt("hello"); + + let mut prewarm_session = client.new_session(); + prewarm_session + .prewarm_websocket( + &prompt, + &model_info, + &session_telemetry, + /*effort*/ None, + ReasoningSummary::Auto, + /*service_tier*/ None, + /*turn_metadata_header*/ None, + ) + .await + .expect("bearer prewarm should succeed"); + drop(prewarm_session); + + let mut agent_task_session = client.new_session_with_agent_task(Some(agent_task.clone())); + let mut stream = agent_task_session + .stream( + &prompt, + &model_info, + &session_telemetry, + /*effort*/ None, + ReasoningSummary::Auto, + /*service_tier*/ None, + /*turn_metadata_header*/ None, + ) + .await + .expect("agent task stream should succeed"); + drain_stream_to_completion(&mut stream) + .await + .expect("agent task websocket stream should complete"); + + let handshakes = server.handshakes(); + assert_eq!(handshakes.len(), 2); + assert_eq!( + handshakes[0].header("authorization"), + Some("Bearer access-token-account-123".to_string()) + ); + let agent_authorization = handshakes[1] + .header("authorization") + .expect("agent handshake should include authorization"); + assert_agent_assertion_header( + &agent_authorization, + &stored_identity, + &agent_task.agent_runtime_id, + &agent_task.task_id, + ); + assert_eq!(handshakes[1].header("chatgpt-account-id"), None); + + server.shutdown().await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn websocket_agent_task_reuses_cached_connection_for_same_task() { + core_test_support::skip_if_no_network!(); + + let server = responses::start_websocket_server(vec![vec![ + vec![ + responses::ev_response_created("resp-1"), + responses::ev_completed("resp-1"), + ], + vec![ + responses::ev_response_created("resp-2"), + responses::ev_completed("resp-2"), + ], + ]]) + .await; + let mut provider = + create_oss_provider_with_base_url(&format!("{}/v1", server.uri()), WireApi::Responses); + provider.supports_websockets = true; + provider.websocket_connect_timeout_ms = Some(5_000); + let (_codex_home, client, agent_task, stored_identity) = + model_client_with_agent_task(provider).await; + let model_info = test_model_info(); + let session_telemetry = test_session_telemetry(); + let prompt = test_prompt("hello"); + + { + let mut first_session = client.new_session_with_agent_task(Some(agent_task.clone())); + let mut stream = first_session + .stream( + &prompt, + &model_info, + &session_telemetry, + /*effort*/ None, + ReasoningSummary::Auto, + /*service_tier*/ None, + /*turn_metadata_header*/ None, + ) + .await + .expect("first agent task stream should succeed"); + drain_stream_to_completion(&mut stream) + .await + .expect("first agent task websocket stream should complete"); + } + + let mut second_session = client.new_session_with_agent_task(Some(agent_task.clone())); + let mut stream = second_session + .stream( + &prompt, + &model_info, + &session_telemetry, + /*effort*/ None, + ReasoningSummary::Auto, + /*service_tier*/ None, + /*turn_metadata_header*/ None, + ) + .await + .expect("second agent task stream should succeed"); + drain_stream_to_completion(&mut stream) + .await + .expect("second agent task websocket stream should complete"); + + let handshakes = server.handshakes(); + assert_eq!(handshakes.len(), 1); + let agent_authorization = handshakes[0] + .header("authorization") + .expect("agent handshake should include authorization"); + assert_agent_assertion_header( + &agent_authorization, + &stored_identity, + &agent_task.agent_runtime_id, + &agent_task.task_id, + ); + assert_eq!(server.single_connection().len(), 2); + + server.shutdown().await; +} diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index cd6f1e84fdd..21bf19971d0 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -14,6 +14,8 @@ use crate::agent::Mailbox; use crate::agent::MailboxReceiver; use crate::agent::agent_status_from_event; use crate::agent::status::is_final; +use crate::agent_identity::AgentIdentityManager; +use crate::agent_identity::RegisteredAgentTask; use crate::apps::render_apps_section; use crate::commit_attribution::commit_message_trailer_instruction; use crate::compact; @@ -119,6 +121,7 @@ use codex_protocol::protocol::RawResponseItemEvent; use codex_protocol::protocol::ReviewRequest; use codex_protocol::protocol::RolloutItem; use codex_protocol::protocol::SessionSource; +use codex_protocol::protocol::SessionStateUpdate; use codex_protocol::protocol::SubAgentSource; use codex_protocol::protocol::TurnAbortReason; use codex_protocol::protocol::TurnContextItem; @@ -843,6 +846,7 @@ pub(crate) struct Session { pub(crate) services: SessionServices, js_repl: Arc, next_internal_sub_id: AtomicU64, + agent_task_registration_lock: Mutex<()>, } #[derive(Clone, Debug)] @@ -1477,6 +1481,190 @@ impl Session { }); } + fn start_agent_identity_registration(self: &Arc) { + if !self.services.agent_identity_manager.is_enabled() { + return; + } + + let weak_sess = Arc::downgrade(self); + let mut auth_state_rx = self.services.auth_manager.subscribe_auth_state(); + tokio::spawn(async move { + loop { + let Some(sess) = weak_sess.upgrade() else { + return; + }; + match sess + .services + .agent_identity_manager + .ensure_registered_identity() + .await + { + Ok(Some(_)) => return, + Ok(None) => { + drop(sess); + if auth_state_rx.changed().await.is_err() { + return; + } + } + Err(error) => { + sess.fail_agent_identity_registration(error).await; + return; + } + } + } + }); + } + + async fn fail_agent_identity_registration(self: &Arc, error: anyhow::Error) { + warn!(error = %error, "agent identity registration failed"); + let message = format!( + "Agent identity registration failed. Codex cannot continue while `features.use_agent_identity` is enabled: {error}" + ); + self.send_event_raw(Event { + id: self.next_internal_sub_id(), + msg: EventMsg::Error(ErrorEvent { + message, + codex_error_info: Some(CodexErrorInfo::Other), + }), + }) + .await; + handlers::shutdown(self, self.next_internal_sub_id()).await; + } + + fn latest_persisted_agent_task( + rollout_items: &[RolloutItem], + ) -> Option> { + rollout_items.iter().rev().find_map(|item| match item { + RolloutItem::SessionState(update) => Some( + update + .agent_task + .clone() + .map(RegisteredAgentTask::from_session_agent_task), + ), + _ => None, + }) + } + + async fn restore_persisted_agent_task(&self, rollout_items: &[RolloutItem]) { + let Some(agent_task) = Self::latest_persisted_agent_task(rollout_items).flatten() else { + return; + }; + + let mut state = self.state.lock().await; + state.set_agent_task(agent_task); + } + + async fn persist_agent_task_update(&self, agent_task: Option<&RegisteredAgentTask>) { + self.persist_rollout_items(&[RolloutItem::SessionState(SessionStateUpdate { + agent_task: agent_task.map(RegisteredAgentTask::to_session_agent_task), + })]) + .await; + } + + async fn clear_cached_agent_task(&self, agent_task: &RegisteredAgentTask) { + let cleared = { + let mut state = self.state.lock().await; + if state.agent_task().as_ref() == Some(agent_task) { + state.clear_agent_task(); + true + } else { + false + } + }; + if cleared { + self.persist_agent_task_update(/*agent_task*/ None).await; + } + } + + async fn cache_agent_task(&self, agent_task: RegisteredAgentTask) -> RegisteredAgentTask { + let changed = { + let mut state = self.state.lock().await; + if state.agent_task().as_ref() == Some(&agent_task) { + false + } else { + state.set_agent_task(agent_task.clone()); + true + } + }; + if changed { + self.persist_agent_task_update(Some(&agent_task)).await; + } + agent_task + } + + async fn cached_agent_task_for_current_identity(&self) -> Option { + let agent_task = { + let state = self.state.lock().await; + state.agent_task() + }?; + + if self + .services + .agent_identity_manager + .task_matches_current_identity(&agent_task) + .await + { + debug!( + agent_runtime_id = %agent_task.agent_runtime_id, + task_id = %agent_task.task_id, + "reusing cached agent task" + ); + return Some(agent_task); + } + + debug!( + agent_runtime_id = %agent_task.agent_runtime_id, + task_id = %agent_task.task_id, + "discarding cached agent task because the registered agent identity changed" + ); + self.clear_cached_agent_task(&agent_task).await; + None + } + + async fn ensure_agent_task_registered(&self) -> anyhow::Result> { + if let Some(agent_task) = self.cached_agent_task_for_current_identity().await { + return Ok(Some(agent_task)); + } + + let _guard = self.agent_task_registration_lock.lock().await; + if let Some(agent_task) = self.cached_agent_task_for_current_identity().await { + return Ok(Some(agent_task)); + } + + for _ in 0..2 { + let Some(agent_task) = self.services.agent_identity_manager.register_task().await? + else { + return Ok(None); + }; + + if !self + .services + .agent_identity_manager + .task_matches_current_identity(&agent_task) + .await + { + debug!( + agent_runtime_id = %agent_task.agent_runtime_id, + task_id = %agent_task.task_id, + "discarding newly registered agent task because the registered agent identity changed" + ); + continue; + } + + let agent_task = self.cache_agent_task(agent_task).await; + + info!( + thread_id = %self.conversation_id, + agent_runtime_id = %agent_task.agent_runtime_id, + task_id = %agent_task.task_id, + "registered agent task for thread" + ); + return Ok(Some(agent_task)); + } + + Ok(None) + } + #[allow(clippy::too_many_arguments)] fn make_turn_context( conversation_id: ThreadId, @@ -1997,6 +2185,11 @@ impl Session { config.analytics_enabled, ) }); + let agent_identity_manager = Arc::new(AgentIdentityManager::new( + config.as_ref(), + Arc::clone(&auth_manager), + session_configuration.session_source.clone(), + )); let services = SessionServices { // Initialize the MCP connection manager with an uninitialized // instance. It will be replaced with one created via @@ -2019,6 +2212,7 @@ impl Session { hooks, rollout: Mutex::new(rollout_recorder), user_shell: Arc::new(default_shell), + agent_identity_manager: Arc::clone(&agent_identity_manager), shell_snapshot_tx, show_raw_agent_reasoning: config.show_raw_agent_reasoning, exec_policy, @@ -2035,8 +2229,9 @@ impl Session { network_proxy, network_approval: Arc::clone(&network_approval), state_db: state_db_ctx.clone(), - model_client: ModelClient::new( + model_client: ModelClient::new_with_agent_identity_manager( Some(Arc::clone(&auth_manager)), + Some(agent_identity_manager), conversation_id, installation_id, session_configuration.provider.clone(), @@ -2080,6 +2275,7 @@ impl Session { services, js_repl, next_internal_sub_id: AtomicU64::new(0), + agent_task_registration_lock: Mutex::new(()), }); if let Some(network_policy_decider_session) = network_policy_decider_session { let mut guard = network_policy_decider_session.write().await; @@ -2116,6 +2312,7 @@ impl Session { // Start the watcher after SessionConfigured so it cannot emit earlier events. sess.start_skills_watcher_listener(); + sess.start_agent_identity_registration(); // Construct sandbox_state before MCP startup so it can be sent to each // MCP server immediately after it becomes ready (avoiding blocking). let sandbox_state = SandboxState { @@ -2348,6 +2545,7 @@ impl Session { } InitialHistory::Resumed(resumed_history) => { let rollout_items = resumed_history.history; + self.restore_persisted_agent_task(&rollout_items).await; let previous_turn_settings = self .apply_rollout_reconstruction(&turn_context, &rollout_items) .await; @@ -6156,6 +6354,14 @@ pub(crate) async fn run_turn( .await; user_prompt_submit_outcome.additional_contexts }; + let agent_task_registration = if sess.services.agent_identity_manager.is_enabled() { + let sess = Arc::clone(&sess); + Some(tokio::spawn(async move { + sess.ensure_agent_task_registered().await + })) + } else { + None + }; sess.services .analytics_events_client .track_app_mentioned(tracking.clone(), mentioned_app_invocations); @@ -6177,6 +6383,22 @@ pub(crate) async fn run_turn( })) .await; } + let agent_task_result = match agent_task_registration { + Some(registration) => registration.await.unwrap_or_else(|error| { + Err(anyhow::anyhow!( + "agent task registration task failed: {error}" + )) + }), + None => sess.ensure_agent_task_registered().await, + }; + let agent_task = match agent_task_result { + Ok(agent_task) => agent_task, + Err(error) => { + warn!(error = %error, "agent task registration failed"); + sess.fail_agent_identity_registration(error).await; + return None; + } + }; if !skill_items.is_empty() { sess.record_conversation_items(&turn_context, &skill_items) @@ -6199,8 +6421,21 @@ pub(crate) async fn run_turn( // `ModelClientSession` is turn-scoped and caches WebSocket + sticky routing state, so we reuse // one instance across retries within this turn. - let mut client_session = - prewarmed_client_session.unwrap_or_else(|| sess.services.model_client.new_session()); + let mut prewarmed_client_session = prewarmed_client_session; + if agent_task.is_some() + && let Some(prewarmed_client_session) = prewarmed_client_session.as_mut() + { + prewarmed_client_session.disable_cached_websocket_session_on_drop(); + } + let mut client_session = if let Some(agent_task) = agent_task { + sess.services + .model_client + .new_session_with_agent_task(Some(agent_task)) + } else if let Some(prewarmed_client_session) = prewarmed_client_session.take() { + prewarmed_client_session + } else { + sess.services.model_client.new_session() + }; // Pending input is drained into history before building the next model request. // However, we defer that drain until after sampling in two cases: // 1. At the start of a turn, so the fresh user prompt in `input` gets sampled first. @@ -6812,6 +7047,7 @@ async fn run_sampling_request( ) .await; let mut retries = 0; + let mut stale_agent_task_refreshed = false; loop { let err = match try_run_sampling_request( tool_runtime.clone(), @@ -6833,6 +7069,40 @@ async fn run_sampling_request( sess.set_total_tokens_full(&turn_context).await; return Err(CodexErr::ContextWindowExceeded); } + Err(CodexErr::AgentTaskStale) => { + if stale_agent_task_refreshed { + return Err(CodexErr::AgentTaskStale); + } + stale_agent_task_refreshed = true; + let stale_agent_task = client_session.agent_task().cloned(); + client_session.disable_cached_websocket_session_on_drop(); + if let Some(stale_agent_task) = stale_agent_task.as_ref() { + sess.clear_cached_agent_task(stale_agent_task).await; + } + match sess.ensure_agent_task_registered().await { + Ok(Some(agent_task)) => { + *client_session = sess + .services + .model_client + .new_session_with_agent_task(Some(agent_task)); + retries = 0; + continue; + } + Ok(None) => { + return Err(CodexErr::Stream( + "agent assertion task became unavailable after identity changed" + .to_string(), + None, + )); + } + Err(error) => { + return Err(CodexErr::Stream( + format!("failed to refresh stale agent task: {error}"), + None, + )); + } + } + } Err(CodexErr::UsageLimitReached(e)) => { let rate_limits = e.rate_limits.clone(); if let Some(rate_limits) = rate_limits { diff --git a/codex-rs/core/src/codex/rollout_reconstruction.rs b/codex-rs/core/src/codex/rollout_reconstruction.rs index a4c042af0c8..3e407c4cd79 100644 --- a/codex-rs/core/src/codex/rollout_reconstruction.rs +++ b/codex-rs/core/src/codex/rollout_reconstruction.rs @@ -207,7 +207,9 @@ impl Session { active_segment.get_or_insert_with(ActiveReplaySegment::default); active_segment.counts_as_user_turn |= is_user_turn_boundary(response_item); } - RolloutItem::EventMsg(_) | RolloutItem::SessionMeta(_) => {} + RolloutItem::EventMsg(_) + | RolloutItem::SessionMeta(_) + | RolloutItem::SessionState(_) => {} } if base_replacement_history.is_some() @@ -275,6 +277,7 @@ impl Session { history.drop_last_n_user_turns(rollback.num_turns); } RolloutItem::EventMsg(_) + | RolloutItem::SessionState(_) | RolloutItem::TurnContext(_) | RolloutItem::SessionMeta(_) => {} } diff --git a/codex-rs/core/src/codex_tests.rs b/codex-rs/core/src/codex_tests.rs index 6e901a03467..9e5957b003d 100644 --- a/codex-rs/core/src/codex_tests.rs +++ b/codex-rs/core/src/codex_tests.rs @@ -110,6 +110,7 @@ use opentelemetry::trace::TraceId; use std::path::Path; use std::time::Duration; use tokio::time::sleep; +use tokio::time::timeout; use tracing_opentelemetry::OpenTelemetrySpanExt; use codex_protocol::mcp::CallToolResult as McpCallToolResult; @@ -1076,6 +1077,67 @@ async fn record_initial_history_reconstructs_resumed_transcript() { assert_eq!(expected, history.raw_items()); } +#[tokio::test] +async fn record_initial_history_restores_latest_persisted_agent_task() { + let (session, _turn_context) = make_session_and_context().await; + let expected = RegisteredAgentTask { + agent_runtime_id: "agent-123".to_string(), + task_id: "task-123".to_string(), + registered_at: "2026-03-23T12:00:00Z".to_string(), + }; + let rollout_items = vec![ + RolloutItem::SessionState(codex_protocol::protocol::SessionStateUpdate { + agent_task: Some(expected.to_session_agent_task()), + }), + RolloutItem::SessionState(codex_protocol::protocol::SessionStateUpdate { + agent_task: None, + }), + RolloutItem::SessionState(codex_protocol::protocol::SessionStateUpdate { + agent_task: Some(expected.to_session_agent_task()), + }), + ]; + + session + .record_initial_history(InitialHistory::Resumed(ResumedHistory { + conversation_id: ThreadId::default(), + history: rollout_items, + rollout_path: PathBuf::from("/tmp/resume.jsonl"), + })) + .await; + + assert_eq!(session.state.lock().await.agent_task(), Some(expected)); +} + +#[tokio::test] +async fn record_initial_history_honors_cleared_persisted_agent_task() { + let (session, _turn_context) = make_session_and_context().await; + let rollout_items = vec![ + RolloutItem::SessionState(codex_protocol::protocol::SessionStateUpdate { + agent_task: Some( + RegisteredAgentTask { + agent_runtime_id: "agent-123".to_string(), + task_id: "task-123".to_string(), + registered_at: "2026-03-23T12:00:00Z".to_string(), + } + .to_session_agent_task(), + ), + }), + RolloutItem::SessionState(codex_protocol::protocol::SessionStateUpdate { + agent_task: None, + }), + ]; + + session + .record_initial_history(InitialHistory::Resumed(ResumedHistory { + conversation_id: ThreadId::default(), + history: rollout_items, + rollout_path: PathBuf::from("/tmp/resume.jsonl"), + })) + .await; + + assert_eq!(session.state.lock().await.agent_task(), None); +} + #[tokio::test] async fn record_initial_history_new_defers_initial_context_until_first_turn() { let (session, _turn_context) = make_session_and_context().await; @@ -2866,6 +2928,11 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { }), rollout: Mutex::new(None), user_shell: Arc::new(default_user_shell()), + agent_identity_manager: Arc::new(crate::agent_identity::AgentIdentityManager::new( + config.as_ref(), + Arc::clone(&auth_manager), + session_configuration.session_source.clone(), + )), shell_snapshot_tx: watch::channel(None).0, show_raw_agent_reasoning: config.show_raw_agent_reasoning, exec_policy, @@ -2948,6 +3015,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { services, js_repl, next_internal_sub_id: AtomicU64::new(0), + agent_task_registration_lock: Mutex::new(()), }; (session, turn_context) @@ -3711,6 +3779,11 @@ pub(crate) async fn make_session_and_context_with_dynamic_tools_and_rx( }), rollout: Mutex::new(None), user_shell: Arc::new(default_user_shell()), + agent_identity_manager: Arc::new(crate::agent_identity::AgentIdentityManager::new( + config.as_ref(), + Arc::clone(&auth_manager), + session_configuration.session_source.clone(), + )), shell_snapshot_tx: watch::channel(None).0, show_raw_agent_reasoning: config.show_raw_agent_reasoning, exec_policy, @@ -3793,6 +3866,7 @@ pub(crate) async fn make_session_and_context_with_dynamic_tools_and_rx( services, js_repl, next_internal_sub_id: AtomicU64::new(0), + agent_task_registration_lock: Mutex::new(()), }); (session, turn_context, rx_event) @@ -3808,6 +3882,42 @@ pub(crate) async fn make_session_and_context_with_rx() -> ( make_session_and_context_with_dynamic_tools_and_rx(Vec::new()).await } +#[tokio::test] +async fn fail_agent_identity_registration_emits_error_and_shutdown() { + let (session, _turn_context, rx_event) = make_session_and_context_with_rx().await; + + session + .fail_agent_identity_registration(anyhow::anyhow!("registration exploded")) + .await; + + let error_event = timeout(Duration::from_secs(1), rx_event.recv()) + .await + .expect("error event should arrive") + .expect("error event should be readable"); + match error_event.msg { + EventMsg::Error(ErrorEvent { + message, + codex_error_info, + }) => { + assert_eq!( + message, + "Agent identity registration failed. Codex cannot continue while `features.use_agent_identity` is enabled: registration exploded".to_string() + ); + assert_eq!(codex_error_info, Some(CodexErrorInfo::Other)); + } + other => panic!("expected error event, got {other:?}"), + } + + let shutdown_event = timeout(Duration::from_secs(1), rx_event.recv()) + .await + .expect("shutdown event should arrive") + .expect("shutdown event should be readable"); + match shutdown_event.msg { + EventMsg::ShutdownComplete => {} + other => panic!("expected shutdown event, got {other:?}"), + } +} + #[tokio::test] async fn refresh_mcp_servers_is_deferred_until_next_turn() { let (session, turn_context) = make_session_and_context().await; diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 5e62199d2d9..2af3fd900a9 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -5,6 +5,7 @@ // the TUI or the tracing stack). #![deny(clippy::print_stdout, clippy::print_stderr)] +mod agent_identity; mod apply_patch; mod apps; mod arc_monitor; diff --git a/codex-rs/core/src/mcp_openai_file.rs b/codex-rs/core/src/mcp_openai_file.rs index 587dd4b7701..8ee0700f2ef 100644 --- a/codex-rs/core/src/mcp_openai_file.rs +++ b/codex-rs/core/src/mcp_openai_file.rs @@ -112,10 +112,8 @@ async fn build_uploaded_local_argument_value( let token_data = auth .get_token_data() .map_err(|error| format!("failed to read ChatGPT auth for file upload: {error}"))?; - let upload_auth = CoreAuthProvider { - token: Some(token_data.access_token), - account_id: token_data.account_id, - }; + let upload_auth = + CoreAuthProvider::from_bearer_token(Some(token_data.access_token), token_data.account_id); let uploaded = upload_local_file( turn_context.config.chatgpt_base_url.trim_end_matches('/'), &upload_auth, diff --git a/codex-rs/core/src/session_startup_prewarm.rs b/codex-rs/core/src/session_startup_prewarm.rs index acd9232c09f..cbce8e2baa9 100644 --- a/codex-rs/core/src/session_startup_prewarm.rs +++ b/codex-rs/core/src/session_startup_prewarm.rs @@ -13,6 +13,7 @@ use crate::codex::INITIAL_SUBMIT_ID; use crate::codex::Session; use crate::codex::build_prompt; use crate::codex::built_tools; +use codex_app_server_protocol::AuthMode; use codex_otel::STARTUP_PREWARM_AGE_AT_FIRST_TURN_METRIC; use codex_otel::STARTUP_PREWARM_DURATION_METRIC; use codex_otel::SessionTelemetry; @@ -157,6 +158,15 @@ impl SessionStartupPrewarmHandle { impl Session { pub(crate) async fn schedule_startup_prewarm(self: &Arc, base_instructions: String) { + if self.services.agent_identity_manager.is_enabled() + && self.services.auth_manager.auth_mode() != Some(AuthMode::ApiKey) + { + info!( + "skipping startup websocket prewarm because agent identity requires task-scoped auth" + ); + return; + } + let session_telemetry = self.services.session_telemetry.clone(); let websocket_connect_timeout = self.provider().await.websocket_connect_timeout(); let started_at = Instant::now(); diff --git a/codex-rs/core/src/state/service.rs b/codex-rs/core/src/state/service.rs index 2e89585f99f..4b6c5b30f51 100644 --- a/codex-rs/core/src/state/service.rs +++ b/codex-rs/core/src/state/service.rs @@ -4,6 +4,7 @@ use std::sync::Arc; use crate::RolloutRecorder; use crate::SkillsManager; use crate::agent::AgentControl; +use crate::agent_identity::AgentIdentityManager; use crate::client::ModelClient; use crate::config::StartedNetworkProxy; use crate::exec_policy::ExecPolicyManager; @@ -41,6 +42,7 @@ pub(crate) struct SessionServices { pub(crate) hooks: Hooks, pub(crate) rollout: Mutex>, pub(crate) user_shell: Arc, + pub(crate) agent_identity_manager: Arc, pub(crate) shell_snapshot_tx: watch::Sender>>, pub(crate) show_raw_agent_reasoning: bool, pub(crate) exec_policy: Arc, diff --git a/codex-rs/core/src/state/session.rs b/codex-rs/core/src/state/session.rs index 206f75060c7..cb3bfe3a545 100644 --- a/codex-rs/core/src/state/session.rs +++ b/codex-rs/core/src/state/session.rs @@ -6,6 +6,7 @@ use codex_sandboxing::policy_transforms::merge_permission_profiles; use std::collections::HashMap; use std::collections::HashSet; +use crate::agent_identity::RegisteredAgentTask; use crate::codex::PreviousTurnSettings; use crate::codex::SessionConfiguration; use crate::context_manager::ContextManager; @@ -30,6 +31,7 @@ pub(crate) struct SessionState { previous_turn_settings: Option, /// Startup prewarmed session prepared during session initialization. pub(crate) startup_prewarm: Option, + pub(crate) agent_task: Option, pub(crate) active_connector_selection: HashSet, pub(crate) pending_session_start_source: Option, granted_permissions: Option, @@ -48,6 +50,7 @@ impl SessionState { mcp_dependency_prompted: HashSet::new(), previous_turn_settings: None, startup_prewarm: None, + agent_task: None, active_connector_selection: HashSet::new(), pending_session_start_source: None, granted_permissions: None, @@ -174,6 +177,18 @@ impl SessionState { self.startup_prewarm.take() } + pub(crate) fn agent_task(&self) -> Option { + self.agent_task.clone() + } + + pub(crate) fn set_agent_task(&mut self, agent_task: RegisteredAgentTask) { + self.agent_task = Some(agent_task); + } + + pub(crate) fn clear_agent_task(&mut self) { + self.agent_task = None; + } + // Adds connector IDs to the active set and returns the merged selection. pub(crate) fn merge_connector_selection(&mut self, connector_ids: I) -> HashSet where diff --git a/codex-rs/core/src/state/session_tests.rs b/codex-rs/core/src/state/session_tests.rs index 1af7ccc8f60..d5ae7fcde68 100644 --- a/codex-rs/core/src/state/session_tests.rs +++ b/codex-rs/core/src/state/session_tests.rs @@ -1,4 +1,5 @@ use super::*; +use crate::agent_identity::RegisteredAgentTask; use crate::codex::make_session_configuration_for_tests; use codex_protocol::protocol::CreditsSnapshot; use codex_protocol::protocol::RateLimitWindow; @@ -33,6 +34,37 @@ async fn clear_connector_selection_removes_entries() { assert_eq!(state.get_connector_selection(), HashSet::new()); } +#[tokio::test] +async fn set_agent_task_persists_plaintext_task_for_session_reuse() { + let session_configuration = make_session_configuration_for_tests().await; + let mut state = SessionState::new(session_configuration); + let agent_task = RegisteredAgentTask { + agent_runtime_id: "agent-123".to_string(), + task_id: "task_123".to_string(), + registered_at: "2026-03-23T12:00:00Z".to_string(), + }; + + state.set_agent_task(agent_task.clone()); + + assert_eq!(state.agent_task(), Some(agent_task)); +} + +#[tokio::test] +async fn clear_agent_task_removes_cached_task() { + let session_configuration = make_session_configuration_for_tests().await; + let mut state = SessionState::new(session_configuration); + let agent_task = RegisteredAgentTask { + agent_runtime_id: "agent_123".to_string(), + task_id: "task_123".to_string(), + registered_at: "2026-03-23T12:00:00Z".to_string(), + }; + + state.set_agent_task(agent_task); + state.clear_agent_task(); + + assert_eq!(state.agent_task(), None); +} + #[tokio::test] async fn set_rate_limits_defaults_limit_id_to_codex_when_missing() { let session_configuration = make_session_configuration_for_tests().await; diff --git a/codex-rs/core/tests/suite/responses_api_proxy_headers.rs b/codex-rs/core/tests/suite/responses_api_proxy_headers.rs index 7dc888ffce1..a4d3d8d6ee0 100644 --- a/codex-rs/core/tests/suite/responses_api_proxy_headers.rs +++ b/codex-rs/core/tests/suite/responses_api_proxy_headers.rs @@ -138,6 +138,10 @@ async fn responses_api_proxy_dumps_parent_and_subagent_identity_headers() -> Res let proxy_base_url = proxy.base_url(); let mut builder = test_codex().with_config(move |config| { config.model_provider.base_url = Some(proxy_base_url); + config + .features + .enable(Feature::Collab) + .expect("test config should allow feature update"); config .features .disable(Feature::EnableRequestCompression) @@ -179,7 +183,23 @@ async fn responses_api_proxy_dumps_parent_and_subagent_identity_headers() -> Res } fn request_body_contains(req: &wiremock::Request, text: &str) -> bool { - std::str::from_utf8(&req.body).is_ok_and(|body| body.contains(text)) + let is_zstd = req + .headers + .get("content-encoding") + .and_then(|value| value.to_str().ok()) + .is_some_and(|value| { + value + .split(',') + .any(|entry| entry.trim().eq_ignore_ascii_case("zstd")) + }); + let bytes = if is_zstd { + zstd::stream::decode_all(std::io::Cursor::new(&req.body)).ok() + } else { + Some(req.body.clone()) + }; + bytes + .and_then(|body| String::from_utf8(body).ok()) + .is_some_and(|body| body.contains(text)) } fn wait_for_proxy_request_dumps(dump_dir: &Path) -> Result> { @@ -212,12 +232,39 @@ fn read_proxy_request_dumps(dump_dir: &Path) -> Result> { .and_then(|name| name.to_str()) .is_some_and(|name| name.ends_with("-request.json")) { - dumps.push(serde_json::from_str(&std::fs::read_to_string(&path)?)?); + let contents = std::fs::read_to_string(&path)?; + if contents.trim().is_empty() { + continue; + } + + match serde_json::from_str(&contents) { + Ok(dump) => dumps.push(dump), + Err(err) if err.is_eof() => continue, + Err(err) => return Err(err.into()), + } } } Ok(dumps) } +#[test] +fn read_proxy_request_dumps_ignores_in_progress_files() -> Result<()> { + let dump_dir = TempDir::new()?; + std::fs::write(dump_dir.path().join("empty-request.json"), "")?; + std::fs::write(dump_dir.path().join("partial-request.json"), "{\"body\"")?; + std::fs::write( + dump_dir.path().join("complete-request.json"), + serde_json::to_string(&json!({ "body": "ready" }))?, + )?; + + assert_eq!( + read_proxy_request_dumps(dump_dir.path())?, + vec![json!({ "body": "ready" })] + ); + + Ok(()) +} + fn dump_body_contains(dump: &Value, text: &str) -> bool { dump.get("body") .is_some_and(|body| body.to_string().contains(text)) diff --git a/codex-rs/login/src/api_bridge.rs b/codex-rs/login/src/api_bridge.rs index d8b9dbb77cb..85b02f6f169 100644 --- a/codex-rs/login/src/api_bridge.rs +++ b/codex-rs/login/src/api_bridge.rs @@ -8,29 +8,28 @@ pub fn auth_provider_from_auth( provider: &ModelProviderInfo, ) -> codex_protocol::error::Result { if let Some(api_key) = provider.api_key()? { - return Ok(CoreAuthProvider { - token: Some(api_key), - account_id: None, - }); + return Ok(CoreAuthProvider::from_bearer_token( + Some(api_key), + /*account_id*/ None, + )); } if let Some(token) = provider.experimental_bearer_token.clone() { - return Ok(CoreAuthProvider { - token: Some(token), - account_id: None, - }); + return Ok(CoreAuthProvider::from_bearer_token( + Some(token), + /*account_id*/ None, + )); } if let Some(auth) = auth { let token = auth.get_token()?; - Ok(CoreAuthProvider { - token: Some(token), - account_id: auth.get_account_id(), - }) + Ok(CoreAuthProvider::from_bearer_token( + Some(token), + auth.get_account_id(), + )) } else { - Ok(CoreAuthProvider { - token: None, - account_id: None, - }) + Ok(CoreAuthProvider::from_bearer_token( + /*token*/ None, /*account_id*/ None, + )) } } diff --git a/codex-rs/login/src/auth/auth_tests.rs b/codex-rs/login/src/auth/auth_tests.rs index 27064b68312..cbe15f59616 100644 --- a/codex-rs/login/src/auth/auth_tests.rs +++ b/codex-rs/login/src/auth/auth_tests.rs @@ -16,6 +16,8 @@ use serde_json::json; use std::sync::Arc; use tempfile::TempDir; use tempfile::tempdir; +use tokio::time::Duration; +use tokio::time::timeout; #[tokio::test] async fn refresh_without_id_token() { @@ -135,6 +137,7 @@ async fn pro_account_with_no_api_key_uses_chatgpt_auth() { account_id: None, }), last_refresh: Some(last_refresh), + agent_identity: None, }, auth_dot_json ); @@ -172,6 +175,7 @@ fn logout_removes_auth_file() -> Result<(), std::io::Error> { openai_api_key: Some("sk-test-key".to_string()), tokens: None, last_refresh: None, + agent_identity: None, }; super::save_auth(dir.path(), &auth_dot_json, AuthCredentialsStoreMode::File)?; let auth_file = get_auth_file(dir.path()); @@ -181,6 +185,48 @@ fn logout_removes_auth_file() -> Result<(), std::io::Error> { Ok(()) } +#[test] +fn chatgpt_auth_persists_agent_identity_for_workspace() { + let codex_home = tempdir().unwrap(); + write_auth_file( + AuthFileParams { + openai_api_key: None, + chatgpt_plan_type: Some("pro".to_string()), + chatgpt_account_id: Some("account-123".to_string()), + }, + codex_home.path(), + ) + .expect("failed to write auth file"); + let auth = super::load_auth( + codex_home.path(), + /*enable_codex_api_key_env*/ false, + AuthCredentialsStoreMode::File, + ) + .expect("load auth") + .expect("auth available"); + let record = AgentIdentityAuthRecord { + workspace_id: "account-123".to_string(), + agent_runtime_id: "agent_123".to_string(), + agent_private_key: "pkcs8-base64".to_string(), + registered_at: "2026-04-13T12:00:00Z".to_string(), + }; + + auth.set_agent_identity(record.clone()) + .expect("set agent identity"); + + assert_eq!(auth.get_agent_identity("account-123"), Some(record.clone())); + assert_eq!(auth.get_agent_identity("other-account"), None); + let storage = FileAuthStorage::new(codex_home.path().to_path_buf()); + let persisted = storage + .load() + .expect("load auth") + .expect("auth should exist"); + assert_eq!(persisted.agent_identity, Some(record)); + + assert!(auth.remove_agent_identity().expect("remove agent identity")); + assert_eq!(auth.get_agent_identity("account-123"), None); +} + #[test] fn unauthorized_recovery_reports_mode_and_step_names() { let dir = tempdir().unwrap(); @@ -474,6 +520,67 @@ exit 1 } } +#[tokio::test] +async fn auth_manager_notifies_when_auth_state_changes() { + let dir = tempdir().unwrap(); + let manager = AuthManager::shared( + dir.path().to_path_buf(), + /*enable_codex_api_key_env*/ false, + AuthCredentialsStoreMode::File, + ); + let mut auth_state_rx = manager.subscribe_auth_state(); + + save_auth( + dir.path(), + &AuthDotJson { + auth_mode: Some(ApiAuthMode::ApiKey), + openai_api_key: Some("sk-test-key".to_string()), + tokens: None, + last_refresh: None, + agent_identity: None, + }, + AuthCredentialsStoreMode::File, + ) + .expect("save auth"); + + assert!( + manager.reload(), + "reload should report a changed auth state" + ); + timeout(Duration::from_secs(1), auth_state_rx.changed()) + .await + .expect("auth change notification should arrive") + .expect("auth state watch should remain open"); + + save_auth( + dir.path(), + &AuthDotJson { + auth_mode: Some(ApiAuthMode::ApiKey), + openai_api_key: Some("sk-updated-key".to_string()), + tokens: None, + last_refresh: None, + agent_identity: None, + }, + AuthCredentialsStoreMode::File, + ) + .expect("save updated auth"); + + assert!( + !manager.reload(), + "reload remains mode-stable even when the underlying credentials change" + ); + timeout(Duration::from_secs(1), auth_state_rx.changed()) + .await + .expect("auth reload notification should still arrive") + .expect("auth state watch should remain open"); + + manager.set_forced_chatgpt_workspace_id(Some("workspace-123".to_string())); + timeout(Duration::from_secs(1), auth_state_rx.changed()) + .await + .expect("workspace change notification should arrive") + .expect("auth state watch should remain open"); +} + struct AuthFileParams { openai_api_key: Option, chatgpt_plan_type: Option, diff --git a/codex-rs/login/src/auth/manager.rs b/codex-rs/login/src/auth/manager.rs index 71857c9700a..41b097b70ff 100644 --- a/codex-rs/login/src/auth/manager.rs +++ b/codex-rs/login/src/auth/manager.rs @@ -13,6 +13,7 @@ use std::sync::Arc; use std::sync::Mutex; use std::sync::RwLock; use tokio::sync::Mutex as AsyncMutex; +use tokio::sync::watch; use codex_app_server_protocol::AuthMode; use codex_app_server_protocol::AuthMode as ApiAuthMode; @@ -20,6 +21,7 @@ use codex_protocol::config_types::ForcedLoginMethod; use codex_protocol::config_types::ModelProviderAuthInfo; use super::external_bearer::BearerTokenRefresher; +pub use crate::auth::storage::AgentIdentityAuthRecord; pub use crate::auth::storage::AuthDotJson; use crate::auth::storage::AuthStorageBackend; use crate::auth::storage::create_auth_storage; @@ -60,6 +62,7 @@ pub struct ChatgptAuth { #[derive(Debug, Clone)] pub struct ChatgptAuthTokens { state: ChatgptAuthState, + storage: Arc, } #[derive(Debug, Clone)] @@ -204,14 +207,13 @@ impl CodexAuth { client, }; + let storage = create_auth_storage(codex_home.to_path_buf(), storage_mode); match auth_mode { - ApiAuthMode::Chatgpt => { - let storage = create_auth_storage(codex_home.to_path_buf(), storage_mode); - Ok(Self::Chatgpt(ChatgptAuth { state, storage })) - } - ApiAuthMode::ChatgptAuthTokens => { - Ok(Self::ChatgptAuthTokens(ChatgptAuthTokens { state })) - } + ApiAuthMode::Chatgpt => Ok(Self::Chatgpt(ChatgptAuth { state, storage })), + ApiAuthMode::ChatgptAuthTokens => Ok(Self::ChatgptAuthTokens(ChatgptAuthTokens { + state, + storage, + })), ApiAuthMode::ApiKey => unreachable!("api key mode is handled above"), } } @@ -351,6 +353,52 @@ impl CodexAuth { self.get_current_auth_json().and_then(|t| t.tokens) } + pub fn get_agent_identity(&self, workspace_id: &str) -> Option { + self.get_current_auth_json() + .and_then(|auth| auth.agent_identity) + .filter(|identity| identity.workspace_id == workspace_id) + } + + pub fn set_agent_identity(&self, record: AgentIdentityAuthRecord) -> std::io::Result<()> { + let (state, storage) = match self { + Self::Chatgpt(auth) => (&auth.state, &auth.storage), + Self::ChatgptAuthTokens(auth) => (&auth.state, &auth.storage), + Self::ApiKey(_) => return Ok(()), + }; + let mut guard = state + .auth_dot_json + .lock() + .map_err(|_| std::io::Error::other("failed to lock auth state"))?; + let mut auth = guard + .clone() + .ok_or_else(|| std::io::Error::other("auth data is not available"))?; + auth.agent_identity = Some(record); + storage.save(&auth)?; + *guard = Some(auth); + Ok(()) + } + + pub fn remove_agent_identity(&self) -> std::io::Result { + let (state, storage) = match self { + Self::Chatgpt(auth) => (&auth.state, &auth.storage), + Self::ChatgptAuthTokens(auth) => (&auth.state, &auth.storage), + Self::ApiKey(_) => return Ok(false), + }; + let mut guard = state + .auth_dot_json + .lock() + .map_err(|_| std::io::Error::other("failed to lock auth state"))?; + let Some(mut auth) = guard.clone() else { + return Ok(false); + }; + let removed = auth.agent_identity.take().is_some(); + if removed { + storage.save(&auth)?; + *guard = Some(auth); + } + Ok(removed) + } + /// Consider this private to integration tests. pub fn create_dummy_chatgpt_auth_for_testing() -> Self { let auth_dot_json = AuthDotJson { @@ -363,6 +411,7 @@ impl CodexAuth { account_id: Some("account_id".to_string()), }), last_refresh: Some(Utc::now()), + agent_identity: None, }; let client = create_client(); @@ -438,6 +487,7 @@ pub fn login_with_api_key( openai_api_key: Some(api_key.to_string()), tokens: None, last_refresh: None, + agent_identity: None, }; save_auth(codex_home, &auth_dot_json, auth_credentials_store_mode) } @@ -811,6 +861,7 @@ impl AuthDotJson { openai_api_key: None, tokens: Some(tokens), last_refresh: Some(Utc::now()), + agent_identity: None, }) } @@ -1106,6 +1157,7 @@ pub struct AuthManager { forced_chatgpt_workspace_id: RwLock>, refresh_lock: AsyncMutex<()>, external_auth: RwLock>>, + auth_state_tx: watch::Sender<()>, } /// Configuration view required to construct a shared [`AuthManager`]. @@ -1154,6 +1206,7 @@ impl AuthManager { enable_codex_api_key_env: bool, auth_credentials_store_mode: AuthCredentialsStoreMode, ) -> Self { + let (auth_state_tx, _) = watch::channel(()); let managed_auth = load_auth( &codex_home, enable_codex_api_key_env, @@ -1172,11 +1225,13 @@ impl AuthManager { forced_chatgpt_workspace_id: RwLock::new(None), refresh_lock: AsyncMutex::new(()), external_auth: RwLock::new(None), + auth_state_tx, } } /// Create an AuthManager with a specific CodexAuth, for testing only. pub fn from_auth_for_testing(auth: CodexAuth) -> Arc { + let (auth_state_tx, _) = watch::channel(()); let cached = CachedAuth { auth: Some(auth), permanent_refresh_failure: None, @@ -1190,11 +1245,13 @@ impl AuthManager { forced_chatgpt_workspace_id: RwLock::new(None), refresh_lock: AsyncMutex::new(()), external_auth: RwLock::new(None), + auth_state_tx, }) } /// Create an AuthManager with a specific CodexAuth and codex home, for testing only. pub fn from_auth_for_testing_with_home(auth: CodexAuth, codex_home: PathBuf) -> Arc { + let (auth_state_tx, _) = watch::channel(()); let cached = CachedAuth { auth: Some(auth), permanent_refresh_failure: None, @@ -1207,10 +1264,12 @@ impl AuthManager { forced_chatgpt_workspace_id: RwLock::new(None), refresh_lock: AsyncMutex::new(()), external_auth: RwLock::new(None), + auth_state_tx, }) } pub fn external_bearer_only(config: ModelProviderAuthInfo) -> Arc { + let (auth_state_tx, _) = watch::channel(()); Arc::new(Self { codex_home: PathBuf::from("non-existent"), inner: RwLock::new(CachedAuth { @@ -1224,6 +1283,7 @@ impl AuthManager { external_auth: RwLock::new(Some( Arc::new(BearerTokenRefresher::new(config)) as Arc )), + auth_state_tx, }) } @@ -1363,6 +1423,7 @@ impl AuthManager { } tracing::info!("Reloaded auth, changed: {changed}"); guard.auth = new_auth; + self.auth_state_tx.send_replace(()); changed } else { false @@ -1372,18 +1433,23 @@ impl AuthManager { pub fn set_external_auth(&self, external_auth: Arc) { if let Ok(mut guard) = self.external_auth.write() { *guard = Some(external_auth); + self.auth_state_tx.send_replace(()); } } pub fn clear_external_auth(&self) { if let Ok(mut guard) = self.external_auth.write() { *guard = None; + self.auth_state_tx.send_replace(()); } } pub fn set_forced_chatgpt_workspace_id(&self, workspace_id: Option) { - if let Ok(mut guard) = self.forced_chatgpt_workspace_id.write() { + if let Ok(mut guard) = self.forced_chatgpt_workspace_id.write() + && *guard != workspace_id + { *guard = workspace_id; + self.auth_state_tx.send_replace(()); } } @@ -1394,6 +1460,10 @@ impl AuthManager { .and_then(|guard| guard.clone()) } + pub fn subscribe_auth_state(&self) -> watch::Receiver<()> { + self.auth_state_tx.subscribe() + } + pub fn has_external_auth(&self) -> bool { self.external_auth().is_some() } @@ -1637,8 +1707,14 @@ impl AuthManager { ), ))); } - let auth_dot_json = + let mut auth_dot_json = AuthDotJson::from_external_tokens(&refreshed).map_err(RefreshTokenError::Transient)?; + if let Some(previous_auth) = self + .auth_cached() + .and_then(|auth| auth.get_current_auth_json()) + { + auth_dot_json.agent_identity = previous_auth.agent_identity; + } save_auth( &self.codex_home, &auth_dot_json, diff --git a/codex-rs/login/src/auth/storage.rs b/codex-rs/login/src/auth/storage.rs index 97e801415ca..b6794c96198 100644 --- a/codex-rs/login/src/auth/storage.rs +++ b/codex-rs/login/src/auth/storage.rs @@ -39,6 +39,17 @@ pub struct AuthDotJson { #[serde(default, skip_serializing_if = "Option::is_none")] pub last_refresh: Option>, + + #[serde(default, skip_serializing_if = "Option::is_none")] + pub agent_identity: Option, +} + +#[derive(Deserialize, Serialize, Clone, Debug, PartialEq, Eq)] +pub struct AgentIdentityAuthRecord { + pub workspace_id: String, + pub agent_runtime_id: String, + pub agent_private_key: String, + pub registered_at: String, } pub(super) fn get_auth_file(codex_home: &Path) -> PathBuf { diff --git a/codex-rs/login/src/auth/storage_tests.rs b/codex-rs/login/src/auth/storage_tests.rs index 4bf72c11b94..8b408c56bbb 100644 --- a/codex-rs/login/src/auth/storage_tests.rs +++ b/codex-rs/login/src/auth/storage_tests.rs @@ -18,6 +18,7 @@ async fn file_storage_load_returns_auth_dot_json() -> anyhow::Result<()> { openai_api_key: Some("test-key".to_string()), tokens: None, last_refresh: Some(Utc::now()), + agent_identity: None, }; storage @@ -38,6 +39,7 @@ async fn file_storage_save_persists_auth_dot_json() -> anyhow::Result<()> { openai_api_key: Some("test-key".to_string()), tokens: None, last_refresh: Some(Utc::now()), + agent_identity: None, }; let file = get_auth_file(codex_home.path()); @@ -52,6 +54,29 @@ async fn file_storage_save_persists_auth_dot_json() -> anyhow::Result<()> { Ok(()) } +#[tokio::test] +async fn file_storage_persists_agent_identity() -> anyhow::Result<()> { + let codex_home = tempdir()?; + let storage = FileAuthStorage::new(codex_home.path().to_path_buf()); + let auth_dot_json = AuthDotJson { + auth_mode: Some(AuthMode::Chatgpt), + openai_api_key: None, + tokens: None, + last_refresh: Some(Utc::now()), + agent_identity: Some(AgentIdentityAuthRecord { + workspace_id: "account-123".to_string(), + agent_runtime_id: "agent_123".to_string(), + agent_private_key: "pkcs8-base64".to_string(), + registered_at: "2026-04-13T12:00:00Z".to_string(), + }), + }; + + storage.save(&auth_dot_json)?; + + assert_eq!(storage.load()?, Some(auth_dot_json)); + Ok(()) +} + #[test] fn file_storage_delete_removes_auth_file() -> anyhow::Result<()> { let dir = tempdir()?; @@ -60,6 +85,7 @@ fn file_storage_delete_removes_auth_file() -> anyhow::Result<()> { openai_api_key: Some("sk-test-key".to_string()), tokens: None, last_refresh: None, + agent_identity: None, }; let storage = create_auth_storage(dir.path().to_path_buf(), AuthCredentialsStoreMode::File); storage.save(&auth_dot_json)?; @@ -83,6 +109,7 @@ fn ephemeral_storage_save_load_delete_is_in_memory_only() -> anyhow::Result<()> openai_api_key: Some("sk-ephemeral".to_string()), tokens: None, last_refresh: Some(Utc::now()), + agent_identity: None, }; storage.save(&auth_dot_json)?; @@ -181,6 +208,7 @@ fn auth_with_prefix(prefix: &str) -> AuthDotJson { account_id: Some(format!("{prefix}-account-id")), }), last_refresh: None, + agent_identity: None, } } @@ -197,6 +225,7 @@ fn keyring_auth_storage_load_returns_deserialized_auth() -> anyhow::Result<()> { openai_api_key: Some("sk-test".to_string()), tokens: None, last_refresh: None, + agent_identity: None, }; seed_keyring_with_auth( &mock_keyring, @@ -239,6 +268,7 @@ fn keyring_auth_storage_save_persists_and_removes_fallback_file() -> anyhow::Res account_id: Some("account".to_string()), }), last_refresh: Some(Utc::now()), + agent_identity: None, }; storage.save(&auth)?; diff --git a/codex-rs/login/src/lib.rs b/codex-rs/login/src/lib.rs index 1b52223bc33..4880e431061 100644 --- a/codex-rs/login/src/lib.rs +++ b/codex-rs/login/src/lib.rs @@ -20,6 +20,7 @@ pub use server::ShutdownHandle; pub use server::run_login_server; pub use api_bridge::auth_provider_from_auth; +pub use auth::AgentIdentityAuthRecord; pub use auth::AuthConfig; pub use auth::AuthDotJson; pub use auth::AuthManager; diff --git a/codex-rs/login/src/server.rs b/codex-rs/login/src/server.rs index 169a8a3091b..0c7b8101843 100644 --- a/codex-rs/login/src/server.rs +++ b/codex-rs/login/src/server.rs @@ -781,6 +781,7 @@ pub(crate) async fn persist_tokens_async( openai_api_key: api_key, tokens: Some(tokens), last_refresh: Some(Utc::now()), + agent_identity: None, }; save_auth(&codex_home, &auth, auth_credentials_store_mode) }) diff --git a/codex-rs/login/tests/suite/auth_refresh.rs b/codex-rs/login/tests/suite/auth_refresh.rs index bf9e03bc263..d754c9589d9 100644 --- a/codex-rs/login/tests/suite/auth_refresh.rs +++ b/codex-rs/login/tests/suite/auth_refresh.rs @@ -54,6 +54,7 @@ async fn refresh_token_succeeds_updates_storage() -> Result<()> { openai_api_key: None, tokens: Some(initial_tokens.clone()), last_refresh: Some(initial_last_refresh), + agent_identity: None, }; ctx.write_auth(&initial_auth)?; @@ -117,6 +118,7 @@ async fn refresh_token_refreshes_when_auth_is_unchanged() -> Result<()> { openai_api_key: None, tokens: Some(initial_tokens.clone()), last_refresh: Some(initial_last_refresh), + agent_identity: None, }; ctx.write_auth(&initial_auth)?; @@ -171,6 +173,7 @@ async fn refresh_token_skips_refresh_when_auth_changed() -> Result<()> { openai_api_key: None, tokens: Some(initial_tokens), last_refresh: Some(initial_last_refresh), + agent_identity: None, }; ctx.write_auth(&initial_auth)?; @@ -180,6 +183,7 @@ async fn refresh_token_skips_refresh_when_auth_changed() -> Result<()> { openai_api_key: None, tokens: Some(disk_tokens.clone()), last_refresh: Some(initial_last_refresh), + agent_identity: None, }; save_auth( ctx.codex_home.path(), @@ -234,6 +238,7 @@ async fn refresh_token_errors_on_account_mismatch() -> Result<()> { openai_api_key: None, tokens: Some(initial_tokens.clone()), last_refresh: Some(initial_last_refresh), + agent_identity: None, }; ctx.write_auth(&initial_auth)?; @@ -244,6 +249,7 @@ async fn refresh_token_errors_on_account_mismatch() -> Result<()> { openai_api_key: None, tokens: Some(disk_tokens), last_refresh: Some(initial_last_refresh), + agent_identity: None, }; save_auth( ctx.codex_home.path(), @@ -302,6 +308,7 @@ async fn returns_fresh_tokens_as_is() -> Result<()> { openai_api_key: None, tokens: Some(initial_tokens.clone()), last_refresh: Some(stale_refresh), + agent_identity: None, }; ctx.write_auth(&initial_auth)?; @@ -349,6 +356,7 @@ async fn refreshes_token_when_access_token_is_expired() -> Result<()> { openai_api_key: None, tokens: Some(initial_tokens.clone()), last_refresh: Some(fresh_refresh), + agent_identity: None, }; ctx.write_auth(&initial_auth)?; @@ -398,6 +406,7 @@ async fn auth_reloads_disk_auth_when_cached_auth_is_stale() -> Result<()> { openai_api_key: None, tokens: Some(initial_tokens), last_refresh: Some(stale_refresh), + agent_identity: None, }; ctx.write_auth(&initial_auth)?; @@ -408,6 +417,7 @@ async fn auth_reloads_disk_auth_when_cached_auth_is_stale() -> Result<()> { openai_api_key: None, tokens: Some(disk_tokens.clone()), last_refresh: Some(fresh_refresh), + agent_identity: None, }; save_auth( ctx.codex_home.path(), @@ -459,6 +469,7 @@ async fn auth_reloads_disk_auth_without_calling_expired_refresh_token() -> Resul openai_api_key: None, tokens: Some(initial_tokens), last_refresh: Some(stale_refresh), + agent_identity: None, }; ctx.write_auth(&initial_auth)?; @@ -469,6 +480,7 @@ async fn auth_reloads_disk_auth_without_calling_expired_refresh_token() -> Resul openai_api_key: None, tokens: Some(disk_tokens.clone()), last_refresh: Some(fresh_refresh), + agent_identity: None, }; save_auth( ctx.codex_home.path(), @@ -518,6 +530,7 @@ async fn refresh_token_returns_permanent_error_for_expired_refresh_token() -> Re openai_api_key: None, tokens: Some(initial_tokens.clone()), last_refresh: Some(initial_last_refresh), + agent_identity: None, }; ctx.write_auth(&initial_auth)?; @@ -570,6 +583,7 @@ async fn refresh_token_does_not_retry_after_permanent_failure() -> Result<()> { openai_api_key: None, tokens: Some(initial_tokens.clone()), last_refresh: Some(initial_last_refresh), + agent_identity: None, }; ctx.write_auth(&initial_auth)?; @@ -636,6 +650,7 @@ async fn refresh_token_reloads_changed_auth_after_permanent_failure() -> Result< openai_api_key: None, tokens: Some(initial_tokens.clone()), last_refresh: Some(initial_last_refresh), + agent_identity: None, }; ctx.write_auth(&initial_auth)?; @@ -657,6 +672,7 @@ async fn refresh_token_reloads_changed_auth_after_permanent_failure() -> Result< openai_api_key: None, tokens: Some(disk_tokens.clone()), last_refresh: Some(fresh_refresh), + agent_identity: None, }; save_auth( ctx.codex_home.path(), @@ -715,6 +731,7 @@ async fn refresh_token_returns_transient_error_on_server_failure() -> Result<()> openai_api_key: None, tokens: Some(initial_tokens.clone()), last_refresh: Some(initial_last_refresh), + agent_identity: None, }; ctx.write_auth(&initial_auth)?; @@ -767,6 +784,7 @@ async fn unauthorized_recovery_reloads_then_refreshes_tokens() -> Result<()> { openai_api_key: None, tokens: Some(initial_tokens.clone()), last_refresh: Some(initial_last_refresh), + agent_identity: None, }; ctx.write_auth(&initial_auth)?; @@ -776,6 +794,7 @@ async fn unauthorized_recovery_reloads_then_refreshes_tokens() -> Result<()> { openai_api_key: None, tokens: Some(disk_tokens.clone()), last_refresh: Some(initial_last_refresh), + agent_identity: None, }; save_auth( ctx.codex_home.path(), @@ -859,6 +878,7 @@ async fn unauthorized_recovery_errors_on_account_mismatch() -> Result<()> { openai_api_key: None, tokens: Some(initial_tokens.clone()), last_refresh: Some(initial_last_refresh), + agent_identity: None, }; ctx.write_auth(&initial_auth)?; @@ -869,6 +889,7 @@ async fn unauthorized_recovery_errors_on_account_mismatch() -> Result<()> { openai_api_key: None, tokens: Some(disk_tokens), last_refresh: Some(initial_last_refresh), + agent_identity: None, }; save_auth( ctx.codex_home.path(), @@ -926,6 +947,7 @@ async fn unauthorized_recovery_requires_chatgpt_auth() -> Result<()> { openai_api_key: Some("sk-test".to_string()), tokens: None, last_refresh: None, + agent_identity: None, }; ctx.write_auth(&auth)?; diff --git a/codex-rs/protocol/src/error.rs b/codex-rs/protocol/src/error.rs index 3ffdeb48e04..1cb13c9b147 100644 --- a/codex-rs/protocol/src/error.rs +++ b/codex-rs/protocol/src/error.rs @@ -76,6 +76,8 @@ pub enum CodexErr { /// Optionally includes the requested delay before retrying the turn. #[error("stream disconnected before completion: {0}")] Stream(String, Option), + #[error("agent task no longer matches the current agent identity")] + AgentTaskStale, #[error( "Codex ran out of room in the model's context window. Start a new thread or clear earlier history before retrying." )] @@ -183,6 +185,7 @@ impl CodexErr { | CodexErr::ContextWindowExceeded | CodexErr::ThreadNotFound(_) | CodexErr::AgentLimitReached { .. } + | CodexErr::AgentTaskStale | CodexErr::Spawn | CodexErr::SessionConfiguredNotFirstEvent | CodexErr::UsageLimitReached(_) diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index 65d6dab04ce..8fb4dc6138e 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -2666,6 +2666,21 @@ impl fmt::Display for SubAgentSource { } } +/// Persisted agent-task details that let a resumed thread keep using the same backend task. +#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, JsonSchema, TS)] +pub struct SessionAgentTask { + pub agent_runtime_id: String, + pub task_id: String, + pub registered_at: String, +} + +/// Session-scoped state updates that can be appended after the canonical SessionMeta line. +#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, JsonSchema, TS, Default)] +pub struct SessionStateUpdate { + #[serde(default)] + pub agent_task: Option, +} + /// SessionMeta contains session-level data that doesn't correspond to a specific turn. /// /// NOTE: There used to be an `instructions` field here, which stored user_instructions, but we @@ -2735,6 +2750,7 @@ pub struct SessionMetaLine { #[serde(tag = "type", content = "payload", rename_all = "snake_case")] pub enum RolloutItem { SessionMeta(SessionMetaLine), + SessionState(SessionStateUpdate), ResponseItem(ResponseItem), Compacted(CompactedItem), TurnContext(TurnContextItem), diff --git a/codex-rs/rollout/src/list.rs b/codex-rs/rollout/src/list.rs index e7d3dae5de9..78510b8cacb 100644 --- a/codex-rs/rollout/src/list.rs +++ b/codex-rs/rollout/src/list.rs @@ -1063,6 +1063,9 @@ async fn read_head_summary(path: &Path, head_limit: usize) -> io::Result { // Not included in `head`; skip. } + RolloutItem::SessionState(_) => { + // Not included in `head`; skip. + } RolloutItem::EventMsg(ev) => { if let EventMsg::UserMessage(user) = ev { summary.saw_user_event = true; @@ -1115,6 +1118,7 @@ pub async fn read_head_for_summary(path: &Path) -> io::Result {} } diff --git a/codex-rs/rollout/src/metadata.rs b/codex-rs/rollout/src/metadata.rs index e8a95839d70..be920b79e28 100644 --- a/codex-rs/rollout/src/metadata.rs +++ b/codex-rs/rollout/src/metadata.rs @@ -70,7 +70,8 @@ pub fn builder_from_items( ) -> Option { if let Some(session_meta) = items.iter().find_map(|item| match item { RolloutItem::SessionMeta(meta_line) => Some(meta_line), - RolloutItem::ResponseItem(_) + RolloutItem::SessionState(_) + | RolloutItem::ResponseItem(_) | RolloutItem::Compacted(_) | RolloutItem::TurnContext(_) | RolloutItem::EventMsg(_) => None, @@ -124,7 +125,8 @@ pub async fn extract_metadata_from_rollout( metadata, memory_mode: items.iter().rev().find_map(|item| match item { RolloutItem::SessionMeta(meta_line) => meta_line.meta.memory_mode.clone(), - RolloutItem::ResponseItem(_) + RolloutItem::SessionState(_) + | RolloutItem::ResponseItem(_) | RolloutItem::Compacted(_) | RolloutItem::TurnContext(_) | RolloutItem::EventMsg(_) => None, diff --git a/codex-rs/rollout/src/policy.rs b/codex-rs/rollout/src/policy.rs index 4b50781e76e..6e6fb8b7a26 100644 --- a/codex-rs/rollout/src/policy.rs +++ b/codex-rs/rollout/src/policy.rs @@ -16,9 +16,10 @@ pub fn is_persisted_response_item(item: &RolloutItem, mode: EventPersistenceMode RolloutItem::ResponseItem(item) => should_persist_response_item(item), RolloutItem::EventMsg(ev) => should_persist_event_msg(ev, mode), // Persist Codex executive markers so we can analyze flows (e.g., compaction, API turns). - RolloutItem::Compacted(_) | RolloutItem::TurnContext(_) | RolloutItem::SessionMeta(_) => { - true - } + RolloutItem::Compacted(_) + | RolloutItem::TurnContext(_) + | RolloutItem::SessionMeta(_) + | RolloutItem::SessionState(_) => true, } } diff --git a/codex-rs/rollout/src/recorder.rs b/codex-rs/rollout/src/recorder.rs index 01ea10cb958..309865613c0 100644 --- a/codex-rs/rollout/src/recorder.rs +++ b/codex-rs/rollout/src/recorder.rs @@ -683,6 +683,9 @@ impl RolloutRecorder { RolloutItem::Compacted(item) => { items.push(RolloutItem::Compacted(item)); } + RolloutItem::SessionState(update) => { + items.push(RolloutItem::SessionState(update)); + } RolloutItem::TurnContext(item) => { items.push(RolloutItem::TurnContext(item)); } @@ -1303,6 +1306,7 @@ async fn resume_candidate_matches_cwd( && let Some(latest_turn_context_cwd) = items.iter().rev().find_map(|item| match item { RolloutItem::TurnContext(turn_context) => Some(turn_context.cwd.as_path()), RolloutItem::SessionMeta(_) + | RolloutItem::SessionState(_) | RolloutItem::ResponseItem(_) | RolloutItem::Compacted(_) | RolloutItem::EventMsg(_) => None, diff --git a/codex-rs/state/src/extract.rs b/codex-rs/state/src/extract.rs index 91c6755ffe9..92946373615 100644 --- a/codex-rs/state/src/extract.rs +++ b/codex-rs/state/src/extract.rs @@ -19,6 +19,7 @@ pub fn apply_rollout_item( ) { match item { RolloutItem::SessionMeta(meta_line) => apply_session_meta_from_item(metadata, meta_line), + RolloutItem::SessionState(_) => {} RolloutItem::TurnContext(turn_ctx) => apply_turn_context(metadata, turn_ctx), RolloutItem::EventMsg(event) => apply_event_msg(metadata, event), RolloutItem::ResponseItem(item) => apply_response_item(metadata, item), @@ -36,9 +37,10 @@ pub fn rollout_item_affects_thread_metadata(item: &RolloutItem) -> bool { RolloutItem::EventMsg( EventMsg::TokenCount(_) | EventMsg::UserMessage(_) | EventMsg::ThreadNameUpdated(_), ) => true, - RolloutItem::EventMsg(_) | RolloutItem::ResponseItem(_) | RolloutItem::Compacted(_) => { - false - } + RolloutItem::SessionState(_) + | RolloutItem::EventMsg(_) + | RolloutItem::ResponseItem(_) + | RolloutItem::Compacted(_) => false, } } diff --git a/codex-rs/state/src/runtime/threads.rs b/codex-rs/state/src/runtime/threads.rs index e56ca3f386e..0ae73bcc141 100644 --- a/codex-rs/state/src/runtime/threads.rs +++ b/codex-rs/state/src/runtime/threads.rs @@ -912,7 +912,8 @@ fn one_thread_id_from_rows( pub(super) fn extract_dynamic_tools(items: &[RolloutItem]) -> Option>> { items.iter().find_map(|item| match item { RolloutItem::SessionMeta(meta_line) => Some(meta_line.meta.dynamic_tools.clone()), - RolloutItem::ResponseItem(_) + RolloutItem::SessionState(_) + | RolloutItem::ResponseItem(_) | RolloutItem::Compacted(_) | RolloutItem::TurnContext(_) | RolloutItem::EventMsg(_) => None, @@ -922,7 +923,8 @@ pub(super) fn extract_dynamic_tools(items: &[RolloutItem]) -> Option Option { items.iter().rev().find_map(|item| match item { RolloutItem::SessionMeta(meta_line) => meta_line.meta.memory_mode.clone(), - RolloutItem::ResponseItem(_) + RolloutItem::SessionState(_) + | RolloutItem::ResponseItem(_) | RolloutItem::Compacted(_) | RolloutItem::TurnContext(_) | RolloutItem::EventMsg(_) => None, diff --git a/codex-rs/tui/src/local_chatgpt_auth.rs b/codex-rs/tui/src/local_chatgpt_auth.rs index e888c0387c3..1f84b289a78 100644 --- a/codex-rs/tui/src/local_chatgpt_auth.rs +++ b/codex-rs/tui/src/local_chatgpt_auth.rs @@ -108,6 +108,7 @@ mod tests { account_id: Some("workspace-1".to_string()), }), last_refresh: Some(Utc::now()), + agent_identity: None, }; save_auth(codex_home, &auth, AuthCredentialsStoreMode::File) .expect("chatgpt auth should save"); @@ -154,6 +155,7 @@ mod tests { openai_api_key: Some("sk-test".to_string()), tokens: None, last_refresh: None, + agent_identity: None, }, AuthCredentialsStoreMode::File, ) diff --git a/notes/pr17387-agent-task-persistence-assumption.md b/notes/pr17387-agent-task-persistence-assumption.md new file mode 100644 index 00000000000..9b1059ebbc2 --- /dev/null +++ b/notes/pr17387-agent-task-persistence-assumption.md @@ -0,0 +1,15 @@ +# PR 17387 Agent Task Persistence Assumption + +The review request said to persist the registered agent task_id "in the same place that other +session details are stored" without specifying the exact persistence shape. + +Assumption used in this patch: + +- Persist session-scoped agent task state in the thread rollout file, alongside other resumable + thread/session data. +- Keep agent identity registration on startup, but register the backend task lazily when a session + first needs task-scoped auth. +- Model task updates as an append-only RolloutItem::SessionState record instead of mutating the + canonical first SessionMeta line, because resumed threads need later updates and clears to win. +- Do not carry auth-binding fields on the persisted session task; only persist the task fields + needed to resume the same backend task for the same stored registered agent identity.