From c4a198323f1bf3be80c56047f4eb1784b7f00a86 Mon Sep 17 00:00:00 2001 From: Jiaming Zhang Date: Wed, 6 May 2026 16:48:52 -0700 Subject: [PATCH 01/23] feat(attestation): request device attestation from desktop app Co-authored-by: Codex --- .../analytics/src/analytics_client_tests.rs | 3 + codex-rs/app-server-client/src/lib.rs | 1 + codex-rs/app-server-client/src/remote.rs | 1 + .../json/AttestationGenerateParams.json | 5 + .../json/AttestationGenerateResponse.json | 14 ++ .../schema/json/ClientRequest.json | 5 + .../schema/json/ServerRequest.json | 28 +++ .../codex_app_server_protocol.schemas.json | 49 ++++ .../codex_app_server_protocol.v2.schemas.json | 5 + .../schema/json/v1/InitializeParams.json | 5 + .../typescript/InitializeCapabilities.ts | 4 + .../schema/typescript/ServerRequest.ts | 3 +- .../v2/AttestationGenerateParams.ts | 5 + .../v2/AttestationGenerateResponse.ts | 9 + .../schema/typescript/v2/index.ts | 2 + .../src/protocol/common.rs | 32 +++ .../app-server-protocol/src/protocol/v1.rs | 3 + .../src/protocol/v2/attestation.rs | 17 ++ .../src/protocol/v2/mod.rs | 2 + codex-rs/app-server-test-client/src/lib.rs | 1 + codex-rs/app-server/README.md | 4 + codex-rs/app-server/src/message_processor.rs | 111 +++++++- codex-rs/app-server/src/outgoing_message.rs | 2 +- .../initialize_processor.rs | 21 +- .../tests/suite/v2/experimental_api.rs | 7 + .../app-server/tests/suite/v2/initialize.rs | 1 + .../tests/suite/v2/thread_status.rs | 1 + codex-rs/core/src/attestation.rs | 39 +++ codex-rs/core/src/client.rs | 145 ++++++++++- codex-rs/core/src/client_tests.rs | 238 ++++++++++++++++++ codex-rs/core/src/codex_delegate.rs | 1 + codex-rs/core/src/lib.rs | 2 + codex-rs/core/src/session/mod.rs | 4 + codex-rs/core/src/session/session.rs | 5 +- codex-rs/core/src/session/tests.rs | 5 + .../core/src/session/tests/guardian_tests.rs | 1 + codex-rs/core/src/state/service.rs | 2 + codex-rs/core/src/thread_manager.rs | 30 +++ codex-rs/debug-client/src/client.rs | 1 + codex-rs/exec/src/lib.rs | 9 + codex-rs/model-provider-info/src/lib.rs | 3 +- .../tui/src/app/app_server_event_targets.rs | 1 + codex-rs/tui/src/app/app_server_requests.rs | 7 + codex-rs/tui/src/app/side.rs | 1 + codex-rs/tui/src/chatwidget.rs | 1 + 45 files changed, 813 insertions(+), 23 deletions(-) create mode 100644 codex-rs/app-server-protocol/schema/json/AttestationGenerateParams.json create mode 100644 codex-rs/app-server-protocol/schema/json/AttestationGenerateResponse.json create mode 100644 codex-rs/app-server-protocol/schema/typescript/v2/AttestationGenerateParams.ts create mode 100644 codex-rs/app-server-protocol/schema/typescript/v2/AttestationGenerateResponse.ts create mode 100644 codex-rs/app-server-protocol/src/protocol/v2/attestation.rs create mode 100644 codex-rs/core/src/attestation.rs diff --git a/codex-rs/analytics/src/analytics_client_tests.rs b/codex-rs/analytics/src/analytics_client_tests.rs index 880adfc254fc..9bcc30bfaef0 100644 --- a/codex-rs/analytics/src/analytics_client_tests.rs +++ b/codex-rs/analytics/src/analytics_client_tests.rs @@ -1122,6 +1122,7 @@ async fn initialize_caches_client_and_thread_lifecycle_publishes_once_initialize }, capabilities: Some(InitializeCapabilities { experimental_api: false, + request_attestation: false, opt_out_notification_methods: None, }), }, @@ -1269,6 +1270,7 @@ async fn compaction_event_ingests_custom_fact() { }, capabilities: Some(InitializeCapabilities { experimental_api: false, + request_attestation: false, opt_out_notification_methods: None, }), }, @@ -1382,6 +1384,7 @@ async fn guardian_review_event_ingests_custom_fact_with_optional_target_item() { }, capabilities: Some(InitializeCapabilities { experimental_api: false, + request_attestation: false, opt_out_notification_methods: None, }), }, diff --git a/codex-rs/app-server-client/src/lib.rs b/codex-rs/app-server-client/src/lib.rs index ebafe351af2f..e167bdb7dbe8 100644 --- a/codex-rs/app-server-client/src/lib.rs +++ b/codex-rs/app-server-client/src/lib.rs @@ -375,6 +375,7 @@ impl InProcessClientStartArgs { pub fn initialize_params(&self) -> InitializeParams { let capabilities = InitializeCapabilities { experimental_api: self.experimental_api, + request_attestation: false, opt_out_notification_methods: if self.opt_out_notification_methods.is_empty() { None } else { diff --git a/codex-rs/app-server-client/src/remote.rs b/codex-rs/app-server-client/src/remote.rs index d75534c16045..4a4426a26069 100644 --- a/codex-rs/app-server-client/src/remote.rs +++ b/codex-rs/app-server-client/src/remote.rs @@ -73,6 +73,7 @@ impl RemoteAppServerConnectArgs { fn initialize_params(&self) -> InitializeParams { let capabilities = InitializeCapabilities { experimental_api: self.experimental_api, + request_attestation: false, opt_out_notification_methods: if self.opt_out_notification_methods.is_empty() { None } else { diff --git a/codex-rs/app-server-protocol/schema/json/AttestationGenerateParams.json b/codex-rs/app-server-protocol/schema/json/AttestationGenerateParams.json new file mode 100644 index 000000000000..310552bb7d6b --- /dev/null +++ b/codex-rs/app-server-protocol/schema/json/AttestationGenerateParams.json @@ -0,0 +1,5 @@ +{ + "$schema": "http://json-schema.org/draft-07/schema#", + "title": "AttestationGenerateParams", + "type": "object" +} \ No newline at end of file diff --git a/codex-rs/app-server-protocol/schema/json/AttestationGenerateResponse.json b/codex-rs/app-server-protocol/schema/json/AttestationGenerateResponse.json new file mode 100644 index 000000000000..b7b7f8c474f0 --- /dev/null +++ b/codex-rs/app-server-protocol/schema/json/AttestationGenerateResponse.json @@ -0,0 +1,14 @@ +{ + "$schema": "http://json-schema.org/draft-07/schema#", + "properties": { + "headerValue": { + "description": "Opaque upstream `x-oai-attestation` header value.", + "type": "string" + } + }, + "required": [ + "headerValue" + ], + "title": "AttestationGenerateResponse", + "type": "object" +} \ No newline at end of file diff --git a/codex-rs/app-server-protocol/schema/json/ClientRequest.json b/codex-rs/app-server-protocol/schema/json/ClientRequest.json index cac1c33d9bb3..4fe50a4d4638 100644 --- a/codex-rs/app-server-protocol/schema/json/ClientRequest.json +++ b/codex-rs/app-server-protocol/schema/json/ClientRequest.json @@ -1259,6 +1259,11 @@ "array", "null" ] + }, + "requestAttestation": { + "default": false, + "description": "Opt into `attestation/generate` requests for upstream `x-oai-attestation`.", + "type": "boolean" } }, "type": "object" diff --git a/codex-rs/app-server-protocol/schema/json/ServerRequest.json b/codex-rs/app-server-protocol/schema/json/ServerRequest.json index 51cab50810fd..31e3651951b1 100644 --- a/codex-rs/app-server-protocol/schema/json/ServerRequest.json +++ b/codex-rs/app-server-protocol/schema/json/ServerRequest.json @@ -121,6 +121,9 @@ ], "type": "object" }, + "AttestationGenerateParams": { + "type": "object" + }, "ChatgptAuthTokensRefreshParams": { "properties": { "previousAccountId": { @@ -1900,6 +1903,31 @@ "title": "Account/chatgptAuthTokens/refreshRequest", "type": "object" }, + { + "description": "Generate a fresh upstream attestation result on demand.", + "properties": { + "id": { + "$ref": "#/definitions/RequestId" + }, + "method": { + "enum": [ + "attestation/generate" + ], + "title": "Attestation/generateRequestMethod", + "type": "string" + }, + "params": { + "$ref": "#/definitions/AttestationGenerateParams" + } + }, + "required": [ + "id", + "method", + "params" + ], + "title": "Attestation/generateRequest", + "type": "object" + }, { "description": "DEPRECATED APIs below Request to approve a patch. This request is used for Turns started via the legacy APIs (i.e. SendUserTurn, SendUserMessage).", "properties": { diff --git a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json index 5e248457291d..71d9f81bc5d5 100644 --- a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json +++ b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json @@ -83,6 +83,25 @@ "title": "ApplyPatchApprovalResponse", "type": "object" }, + "AttestationGenerateParams": { + "$schema": "http://json-schema.org/draft-07/schema#", + "title": "AttestationGenerateParams", + "type": "object" + }, + "AttestationGenerateResponse": { + "$schema": "http://json-schema.org/draft-07/schema#", + "properties": { + "headerValue": { + "description": "Opaque upstream `x-oai-attestation` header value.", + "type": "string" + } + }, + "required": [ + "headerValue" + ], + "title": "AttestationGenerateResponse", + "type": "object" + }, "ChatgptAuthTokensRefreshParams": { "$schema": "http://json-schema.org/draft-07/schema#", "properties": { @@ -2608,6 +2627,11 @@ "array", "null" ] + }, + "requestAttestation": { + "default": false, + "description": "Opt into `attestation/generate` requests for upstream `x-oai-attestation`.", + "type": "boolean" } }, "type": "object" @@ -5189,6 +5213,31 @@ "title": "Account/chatgptAuthTokens/refreshRequest", "type": "object" }, + { + "description": "Generate a fresh upstream attestation result on demand.", + "properties": { + "id": { + "$ref": "#/definitions/v2/RequestId" + }, + "method": { + "enum": [ + "attestation/generate" + ], + "title": "Attestation/generateRequestMethod", + "type": "string" + }, + "params": { + "$ref": "#/definitions/AttestationGenerateParams" + } + }, + "required": [ + "id", + "method", + "params" + ], + "title": "Attestation/generateRequest", + "type": "object" + }, { "description": "DEPRECATED APIs below Request to approve a patch. This request is used for Turns started via the legacy APIs (i.e. SendUserTurn, SendUserMessage).", "properties": { diff --git a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json index 6153a54eb9f5..854083ceb2d9 100644 --- a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json +++ b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json @@ -6409,6 +6409,11 @@ "array", "null" ] + }, + "requestAttestation": { + "default": false, + "description": "Opt into `attestation/generate` requests for upstream `x-oai-attestation`.", + "type": "boolean" } }, "type": "object" diff --git a/codex-rs/app-server-protocol/schema/json/v1/InitializeParams.json b/codex-rs/app-server-protocol/schema/json/v1/InitializeParams.json index 6048b8224268..af5c509249a2 100644 --- a/codex-rs/app-server-protocol/schema/json/v1/InitializeParams.json +++ b/codex-rs/app-server-protocol/schema/json/v1/InitializeParams.json @@ -39,6 +39,11 @@ "array", "null" ] + }, + "requestAttestation": { + "default": false, + "description": "Opt into `attestation/generate` requests for upstream `x-oai-attestation`.", + "type": "boolean" } }, "type": "object" diff --git a/codex-rs/app-server-protocol/schema/typescript/InitializeCapabilities.ts b/codex-rs/app-server-protocol/schema/typescript/InitializeCapabilities.ts index 5d42cc4852db..c5043e3b64fc 100644 --- a/codex-rs/app-server-protocol/schema/typescript/InitializeCapabilities.ts +++ b/codex-rs/app-server-protocol/schema/typescript/InitializeCapabilities.ts @@ -10,6 +10,10 @@ export type InitializeCapabilities = { * Opt into receiving experimental API methods and fields. */ experimentalApi: boolean, +/** + * Opt into `attestation/generate` requests for upstream `x-oai-attestation`. + */ +requestAttestation: boolean, /** * Exact notification method names that should be suppressed for this * connection (for example `thread/started`). diff --git a/codex-rs/app-server-protocol/schema/typescript/ServerRequest.ts b/codex-rs/app-server-protocol/schema/typescript/ServerRequest.ts index 13d04b0be701..80e9ffc1162c 100644 --- a/codex-rs/app-server-protocol/schema/typescript/ServerRequest.ts +++ b/codex-rs/app-server-protocol/schema/typescript/ServerRequest.ts @@ -4,6 +4,7 @@ import type { ApplyPatchApprovalParams } from "./ApplyPatchApprovalParams"; import type { ExecCommandApprovalParams } from "./ExecCommandApprovalParams"; import type { RequestId } from "./RequestId"; +import type { AttestationGenerateParams } from "./v2/AttestationGenerateParams"; import type { ChatgptAuthTokensRefreshParams } from "./v2/ChatgptAuthTokensRefreshParams"; import type { CommandExecutionRequestApprovalParams } from "./v2/CommandExecutionRequestApprovalParams"; import type { DynamicToolCallParams } from "./v2/DynamicToolCallParams"; @@ -15,4 +16,4 @@ import type { ToolRequestUserInputParams } from "./v2/ToolRequestUserInputParams /** * Request initiated from the server and sent to the client. */ -export type ServerRequest = { "method": "item/commandExecution/requestApproval", id: RequestId, params: CommandExecutionRequestApprovalParams, } | { "method": "item/fileChange/requestApproval", id: RequestId, params: FileChangeRequestApprovalParams, } | { "method": "item/tool/requestUserInput", id: RequestId, params: ToolRequestUserInputParams, } | { "method": "mcpServer/elicitation/request", id: RequestId, params: McpServerElicitationRequestParams, } | { "method": "item/permissions/requestApproval", id: RequestId, params: PermissionsRequestApprovalParams, } | { "method": "item/tool/call", id: RequestId, params: DynamicToolCallParams, } | { "method": "account/chatgptAuthTokens/refresh", id: RequestId, params: ChatgptAuthTokensRefreshParams, } | { "method": "applyPatchApproval", id: RequestId, params: ApplyPatchApprovalParams, } | { "method": "execCommandApproval", id: RequestId, params: ExecCommandApprovalParams, }; +export type ServerRequest = { "method": "item/commandExecution/requestApproval", id: RequestId, params: CommandExecutionRequestApprovalParams, } | { "method": "item/fileChange/requestApproval", id: RequestId, params: FileChangeRequestApprovalParams, } | { "method": "item/tool/requestUserInput", id: RequestId, params: ToolRequestUserInputParams, } | { "method": "mcpServer/elicitation/request", id: RequestId, params: McpServerElicitationRequestParams, } | { "method": "item/permissions/requestApproval", id: RequestId, params: PermissionsRequestApprovalParams, } | { "method": "item/tool/call", id: RequestId, params: DynamicToolCallParams, } | { "method": "account/chatgptAuthTokens/refresh", id: RequestId, params: ChatgptAuthTokensRefreshParams, } | { "method": "attestation/generate", id: RequestId, params: AttestationGenerateParams, } | { "method": "applyPatchApproval", id: RequestId, params: ApplyPatchApprovalParams, } | { "method": "execCommandApproval", id: RequestId, params: ExecCommandApprovalParams, }; diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/AttestationGenerateParams.ts b/codex-rs/app-server-protocol/schema/typescript/v2/AttestationGenerateParams.ts new file mode 100644 index 000000000000..0e87e7d3e4a4 --- /dev/null +++ b/codex-rs/app-server-protocol/schema/typescript/v2/AttestationGenerateParams.ts @@ -0,0 +1,5 @@ +// GENERATED CODE! DO NOT MODIFY BY HAND! + +// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. + +export type AttestationGenerateParams = Record; diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/AttestationGenerateResponse.ts b/codex-rs/app-server-protocol/schema/typescript/v2/AttestationGenerateResponse.ts new file mode 100644 index 000000000000..48eef943fc54 --- /dev/null +++ b/codex-rs/app-server-protocol/schema/typescript/v2/AttestationGenerateResponse.ts @@ -0,0 +1,9 @@ +// GENERATED CODE! DO NOT MODIFY BY HAND! + +// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. + +export type AttestationGenerateResponse = { +/** + * Opaque upstream `x-oai-attestation` header value. + */ +headerValue: string, }; diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/index.ts b/codex-rs/app-server-protocol/schema/typescript/v2/index.ts index 950dd9839a55..4517e2adb770 100644 --- a/codex-rs/app-server-protocol/schema/typescript/v2/index.ts +++ b/codex-rs/app-server-protocol/schema/typescript/v2/index.ts @@ -28,6 +28,8 @@ export type { AppsDefaultConfig } from "./AppsDefaultConfig"; export type { AppsListParams } from "./AppsListParams"; export type { AppsListResponse } from "./AppsListResponse"; export type { AskForApproval } from "./AskForApproval"; +export type { AttestationGenerateParams } from "./AttestationGenerateParams"; +export type { AttestationGenerateResponse } from "./AttestationGenerateResponse"; export type { AutoReviewDecisionSource } from "./AutoReviewDecisionSource"; export type { ByteRange } from "./ByteRange"; export type { CancelLoginAccountParams } from "./CancelLoginAccountParams"; diff --git a/codex-rs/app-server-protocol/src/protocol/common.rs b/codex-rs/app-server-protocol/src/protocol/common.rs index e79c99a9c971..eab1e6235ab9 100644 --- a/codex-rs/app-server-protocol/src/protocol/common.rs +++ b/codex-rs/app-server-protocol/src/protocol/common.rs @@ -1305,6 +1305,12 @@ server_request_definitions! { response: v2::ChatgptAuthTokensRefreshResponse, }, + /// Generate a fresh upstream attestation result on demand. + AttestationGenerate => "attestation/generate" { + params: v2::AttestationGenerateParams, + response: v2::AttestationGenerateResponse, + }, + /// DEPRECATED APIs below /// Request to approve a patch. /// This request is used for Turns started via the legacy APIs (i.e. SendUserTurn, SendUserMessage). @@ -1891,6 +1897,7 @@ mod tests { }, capabilities: Some(v1::InitializeCapabilities { experimental_api: true, + request_attestation: true, opt_out_notification_methods: Some(vec![ "thread/started".to_string(), "item/agentMessage/delta".to_string(), @@ -1911,6 +1918,7 @@ mod tests { }, "capabilities": { "experimentalApi": true, + "requestAttestation": true, "optOutNotificationMethods": [ "thread/started", "item/agentMessage/delta" @@ -1936,6 +1944,7 @@ mod tests { }, "capabilities": { "experimentalApi": true, + "requestAttestation": true, "optOutNotificationMethods": [ "thread/started", "item/agentMessage/delta" @@ -1956,6 +1965,7 @@ mod tests { }, capabilities: Some(v1::InitializeCapabilities { experimental_api: true, + request_attestation: true, opt_out_notification_methods: Some(vec![ "thread/started".to_string(), "item/agentMessage/delta".to_string(), @@ -2072,6 +2082,28 @@ mod tests { Ok(()) } + #[test] + fn serialize_attestation_generate_request() -> Result<()> { + let params = v2::AttestationGenerateParams {}; + let request = ServerRequest::AttestationGenerate { + request_id: RequestId::Integer(9), + params: params.clone(), + }; + assert_eq!( + json!({ + "method": "attestation/generate", + "id": 9, + "params": {} + }), + serde_json::to_value(&request)?, + ); + + let payload = ServerRequestPayload::AttestationGenerate(params); + assert_eq!(request.id(), &RequestId::Integer(9)); + assert_eq!(payload.request_with_id(RequestId::Integer(9)), request); + Ok(()) + } + #[test] fn serialize_server_response() -> Result<()> { let response = ServerResponse::CommandExecutionRequestApproval { diff --git a/codex-rs/app-server-protocol/src/protocol/v1.rs b/codex-rs/app-server-protocol/src/protocol/v1.rs index d642e7fab954..95ab710a6bd6 100644 --- a/codex-rs/app-server-protocol/src/protocol/v1.rs +++ b/codex-rs/app-server-protocol/src/protocol/v1.rs @@ -46,6 +46,9 @@ pub struct InitializeCapabilities { /// Opt into receiving experimental API methods and fields. #[serde(default)] pub experimental_api: bool, + /// Opt into `attestation/generate` requests for upstream `x-oai-attestation`. + #[serde(default)] + pub request_attestation: bool, /// Exact notification method names that should be suppressed for this /// connection (for example `thread/started`). #[ts(optional = nullable)] diff --git a/codex-rs/app-server-protocol/src/protocol/v2/attestation.rs b/codex-rs/app-server-protocol/src/protocol/v2/attestation.rs new file mode 100644 index 000000000000..36173b63609f --- /dev/null +++ b/codex-rs/app-server-protocol/src/protocol/v2/attestation.rs @@ -0,0 +1,17 @@ +use schemars::JsonSchema; +use serde::Deserialize; +use serde::Serialize; +use ts_rs::TS; + +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS, Default)] +#[serde(rename_all = "camelCase")] +#[ts(export_to = "v2/")] +pub struct AttestationGenerateParams {} + +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)] +#[serde(rename_all = "camelCase")] +#[ts(export_to = "v2/")] +pub struct AttestationGenerateResponse { + /// Opaque upstream `x-oai-attestation` header value. + pub header_value: String, +} diff --git a/codex-rs/app-server-protocol/src/protocol/v2/mod.rs b/codex-rs/app-server-protocol/src/protocol/v2/mod.rs index 275e7ca45b4f..32c24bff1d26 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2/mod.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2/mod.rs @@ -2,6 +2,7 @@ mod shared; mod account; mod apps; +mod attestation; mod collaboration_mode; mod command_exec; mod config; @@ -26,6 +27,7 @@ mod windows_sandbox; pub use account::*; pub use apps::*; +pub use attestation::*; pub use collaboration_mode::*; pub use command_exec::*; pub use config::*; diff --git a/codex-rs/app-server-test-client/src/lib.rs b/codex-rs/app-server-test-client/src/lib.rs index edea431c61f8..ff3a181f2e47 100644 --- a/codex-rs/app-server-test-client/src/lib.rs +++ b/codex-rs/app-server-test-client/src/lib.rs @@ -1551,6 +1551,7 @@ impl CodexClient { }, capabilities: Some(InitializeCapabilities { experimental_api, + request_attestation: false, opt_out_notification_methods: Some( NOTIFICATIONS_TO_OPT_OUT .iter() diff --git a/codex-rs/app-server/README.md b/codex-rs/app-server/README.md index babfac99ba36..c55f84ab276a 100644 --- a/codex-rs/app-server/README.md +++ b/codex-rs/app-server/README.md @@ -1322,6 +1322,10 @@ UI guidance for IDEs: surface an approval dialog as soon as the request arrives. When the client responds to `item/tool/requestUserInput`, the server emits `serverRequest/resolved` with `{ threadId, requestId }`. If the pending request is cleared by turn start, turn completion, or turn interruption before the client answers, the server emits the same notification for that cleanup. +### Attestation generation + +Desktop hosts that provide upstream attestation should set `capabilities.requestAttestation` during `initialize` and handle the server-initiated `attestation/generate` request. App-server issues it just in time before ChatGPT Codex requests that forward `x-oai-attestation`; the client responds with `{ "headerValue": "v1." }`, where `headerValue` is the complete upstream header value. App-server treats that value as opaque and forwards it unchanged. If no initialized client opted into attestation, or if the opted-in client is unavailable, times out, or returns invalid data, app-server omits `x-oai-attestation` for that upstream request. + ### MCP server elicitations MCP servers can interrupt a turn and ask the client for structured input via `mcpServer/elicitation/request`. diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index 2ca0a87d84fc..81dcaf02a6ab 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -40,6 +40,8 @@ use crate::transport::RemoteControlHandle; use async_trait::async_trait; use codex_analytics::AnalyticsEventsClient; use codex_analytics::AppServerRpcTransport; +use codex_app_server_protocol::AttestationGenerateParams; +use codex_app_server_protocol::AttestationGenerateResponse; use codex_app_server_protocol::AuthMode as LoginAuthMode; use codex_app_server_protocol::ChatgptAuthTokensRefreshParams; use codex_app_server_protocol::ChatgptAuthTokensRefreshReason; @@ -58,6 +60,7 @@ use codex_app_server_protocol::ServerRequestPayload; use codex_app_server_protocol::experimental_required_message; use codex_arg0::Arg0DispatchPaths; use codex_chatgpt::workspace_settings; +use codex_core::AttestationProvider; use codex_core::ThreadManager; use codex_core::config::Config; use codex_core::thread_store_from_config; @@ -80,7 +83,9 @@ use tokio::sync::watch; use tokio::time::Duration; use tokio::time::timeout; use tracing::Instrument; +use tracing::warn; +const ATTESTATION_GENERATE_TIMEOUT: Duration = Duration::from_secs(5); const EXTERNAL_AUTH_REFRESH_TIMEOUT: Duration = Duration::from_secs(10); #[derive(Clone)] struct ExternalAuthRefreshBridge { @@ -150,6 +155,86 @@ impl ExternalAuth for ExternalAuthRefreshBridge { } } +fn app_server_attestation_provider( + outgoing: Arc, + attestation_connection_ids: Arc>>, +) -> AttestationProvider { + AttestationProvider::new(move || { + let outgoing = outgoing.clone(); + let attestation_connection_ids = attestation_connection_ids.clone(); + Box::pin(request_attestation_header_value( + outgoing, + attestation_connection_ids, + )) + }) +} + +async fn request_attestation_header_value( + outgoing: Arc, + attestation_connection_ids: Arc>>, +) -> Option { + request_attestation_header_value_with_timeout( + outgoing, + attestation_connection_ids, + ATTESTATION_GENERATE_TIMEOUT, + ) + .await +} + +async fn request_attestation_header_value_with_timeout( + outgoing: Arc, + attestation_connection_ids: Arc>>, + timeout_duration: Duration, +) -> Option { + let connection_id = attestation_connection_ids + .lock() + .await + .iter() + .min_by_key(|connection_id| connection_id.0) + .copied()?; + + let connection_ids = [connection_id]; + let (request_id, rx) = outgoing + .send_request_to_connections( + Some(&connection_ids), + ServerRequestPayload::AttestationGenerate(AttestationGenerateParams {}), + /*thread_id*/ None, + ) + .await; + + let result = match timeout(timeout_duration, rx).await { + Ok(Ok(Ok(result))) => result, + Ok(Ok(Err(err))) => { + warn!( + code = err.code, + message = %err.message, + "attestation generation request failed" + ); + return None; + } + Ok(Err(err)) => { + warn!("attestation generation request canceled: {err}"); + return None; + } + Err(_) => { + let _canceled = outgoing.cancel_request(&request_id).await; + warn!( + timeout_seconds = timeout_duration.as_secs(), + "attestation generation request timed out" + ); + return None; + } + }; + + match serde_json::from_value::(result) { + Ok(response) => Some(response.header_value), + Err(err) => { + warn!("failed to deserialize attestation generation response: {err}"); + None + } + } +} + pub(crate) struct MessageProcessor { outgoing: Arc, account_processor: AccountRequestProcessor, @@ -172,6 +257,7 @@ pub(crate) struct MessageProcessor { turn_processor: TurnRequestProcessor, windows_sandbox_processor: WindowsSandboxRequestProcessor, request_serialization_queues: RequestSerializationQueues, + attestation_connection_ids: Arc>>, } #[derive(Debug)] @@ -186,6 +272,7 @@ pub(crate) struct InitializedConnectionSessionState { pub(crate) opted_out_notification_methods: HashSet, pub(crate) app_server_client_name: String, pub(crate) client_version: String, + pub(crate) request_attestation: bool, } impl Default for ConnectionSessionState { @@ -231,6 +318,12 @@ impl ConnectionSessionState { .map(|session| session.client_version.as_str()) } + pub(crate) fn request_attestation(&self) -> bool { + self.initialized + .get() + .is_some_and(|session| session.request_attestation) + } + pub(crate) fn initialize(&self, session: InitializedConnectionSessionState) -> Result<(), ()> { self.initialized.set(session).map_err(|_| ()) } @@ -280,11 +373,12 @@ impl MessageProcessor { auth_manager.set_external_auth(Arc::new(ExternalAuthRefreshBridge { outgoing: outgoing.clone(), })); + let attestation_connection_ids = Arc::new(Mutex::new(HashSet::new())); // The thread store is intentionally process-scoped. Config reloads can // affect per-thread behavior, but they must not move newly started, // resumed, or forked threads to a different persistence backend/root. let thread_store = thread_store_from_config(config.as_ref(), state_db.clone()); - let thread_manager = Arc::new(ThreadManager::new( + let thread_manager = Arc::new(ThreadManager::new_with_attestation_provider( config.as_ref(), auth_manager.clone(), session_source, @@ -293,6 +387,10 @@ impl MessageProcessor { Arc::clone(&thread_store), state_db.clone(), installation_id, + Some(app_server_attestation_provider( + outgoing.clone(), + attestation_connection_ids.clone(), + )), )); thread_manager .plugins_manager() @@ -467,6 +565,7 @@ impl MessageProcessor { turn_processor, windows_sandbox_processor, request_serialization_queues: RequestSerializationQueues::default(), + attestation_connection_ids, } } @@ -664,6 +763,10 @@ impl MessageProcessor { session_state: &ConnectionSessionState, ) { session_state.rpc_gate.shutdown().await; + self.attestation_connection_ids + .lock() + .await + .remove(&connection_id); self.outgoing.connection_closed(connection_id).await; self.fs_processor.connection_closed(connection_id).await; self.command_exec_processor @@ -721,6 +824,12 @@ impl MessageProcessor { .connection_initialized(connection_id) .await; } + if session.request_attestation() { + self.attestation_connection_ids + .lock() + .await + .insert(connection_id); + } return Ok(()); } diff --git a/codex-rs/app-server/src/outgoing_message.rs b/codex-rs/app-server/src/outgoing_message.rs index a7420f8c7814..1d2035305485 100644 --- a/codex-rs/app-server/src/outgoing_message.rs +++ b/codex-rs/app-server/src/outgoing_message.rs @@ -265,7 +265,7 @@ impl OutgoingMessageSender { RequestId::Integer(self.next_server_request_id.fetch_add(1, Ordering::Relaxed)) } - async fn send_request_to_connections( + pub(crate) async fn send_request_to_connections( &self, connection_ids: Option<&[ConnectionId]>, request: ServerRequestPayload, diff --git a/codex-rs/app-server/src/request_processors/initialize_processor.rs b/codex-rs/app-server/src/request_processors/initialize_processor.rs index a206b2faa02a..c13ce4340f20 100644 --- a/codex-rs/app-server/src/request_processors/initialize_processor.rs +++ b/codex-rs/app-server/src/request_processors/initialize_processor.rs @@ -65,15 +65,17 @@ impl InitializeRequestProcessor { // experimental API). Proposed direction is instance-global first-write-wins // with initialize-time mismatch rejection. let analytics_initialize_params = params.clone(); - let (experimental_api_enabled, opt_out_notification_methods) = match params.capabilities { - Some(capabilities) => ( - capabilities.experimental_api, - capabilities - .opt_out_notification_methods - .unwrap_or_default(), - ), - None => (false, Vec::new()), - }; + let (experimental_api_enabled, request_attestation, opt_out_notification_methods) = + match params.capabilities { + Some(capabilities) => ( + capabilities.experimental_api, + capabilities.request_attestation, + capabilities + .opt_out_notification_methods + .unwrap_or_default(), + ), + None => (false, false, Vec::new()), + }; let ClientInfo { name, title: _title, @@ -95,6 +97,7 @@ impl InitializeRequestProcessor { opted_out_notification_methods: opt_out_notification_methods.into_iter().collect(), app_server_client_name: name.clone(), client_version: version, + request_attestation, }) .is_err() { diff --git a/codex-rs/app-server/tests/suite/v2/experimental_api.rs b/codex-rs/app-server/tests/suite/v2/experimental_api.rs index 9ac0dc3e21f1..4096e3d96fd4 100644 --- a/codex-rs/app-server/tests/suite/v2/experimental_api.rs +++ b/codex-rs/app-server/tests/suite/v2/experimental_api.rs @@ -36,6 +36,7 @@ async fn mock_experimental_method_requires_experimental_api_capability() -> Resu default_client_info(), Some(InitializeCapabilities { experimental_api: false, + request_attestation: false, opt_out_notification_methods: None, }), ) @@ -66,6 +67,7 @@ async fn realtime_conversation_start_requires_experimental_api_capability() -> R default_client_info(), Some(InitializeCapabilities { experimental_api: false, + request_attestation: false, opt_out_notification_methods: None, }), ) @@ -103,6 +105,7 @@ async fn thread_memory_mode_set_requires_experimental_api_capability() -> Result default_client_info(), Some(InitializeCapabilities { experimental_api: false, + request_attestation: false, opt_out_notification_methods: None, }), ) @@ -136,6 +139,7 @@ async fn realtime_webrtc_start_requires_experimental_api_capability() -> Result< default_client_info(), Some(InitializeCapabilities { experimental_api: false, + request_attestation: false, opt_out_notification_methods: None, }), ) @@ -177,6 +181,7 @@ async fn thread_start_mock_field_requires_experimental_api_capability() -> Resul default_client_info(), Some(InitializeCapabilities { experimental_api: false, + request_attestation: false, opt_out_notification_methods: None, }), ) @@ -214,6 +219,7 @@ async fn thread_start_without_dynamic_tools_allows_without_experimental_api_capa default_client_info(), Some(InitializeCapabilities { experimental_api: false, + request_attestation: false, opt_out_notification_methods: None, }), ) @@ -250,6 +256,7 @@ async fn thread_start_granular_approval_policy_requires_experimental_api_capabil default_client_info(), Some(InitializeCapabilities { experimental_api: false, + request_attestation: false, opt_out_notification_methods: None, }), ) diff --git a/codex-rs/app-server/tests/suite/v2/initialize.rs b/codex-rs/app-server/tests/suite/v2/initialize.rs index 165160468f78..dcfd4e54996b 100644 --- a/codex-rs/app-server/tests/suite/v2/initialize.rs +++ b/codex-rs/app-server/tests/suite/v2/initialize.rs @@ -158,6 +158,7 @@ async fn initialize_opt_out_notification_methods_filters_notifications() -> Resu }, Some(InitializeCapabilities { experimental_api: true, + request_attestation: false, opt_out_notification_methods: Some(vec!["thread/started".to_string()]), }), ), diff --git a/codex-rs/app-server/tests/suite/v2/thread_status.rs b/codex-rs/app-server/tests/suite/v2/thread_status.rs index ad90e4900afd..957969c3eaf1 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_status.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_status.rs @@ -145,6 +145,7 @@ async fn thread_status_changed_can_be_opted_out() -> Result<()> { }, Some(InitializeCapabilities { experimental_api: true, + request_attestation: false, opt_out_notification_methods: Some(vec!["thread/status/changed".to_string()]), }), ), diff --git a/codex-rs/core/src/attestation.rs b/codex-rs/core/src/attestation.rs new file mode 100644 index 000000000000..aa637304be26 --- /dev/null +++ b/codex-rs/core/src/attestation.rs @@ -0,0 +1,39 @@ +use std::fmt; +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; + +use http::HeaderValue; + +pub(crate) const X_OAI_ATTESTATION_HEADER: &str = "x-oai-attestation"; + +type GenerateAttestationFuture = Pin> + Send>>; +type GenerateAttestationCallback = dyn Fn() -> GenerateAttestationFuture + Send + Sync + 'static; + +/// Session-scoped source for just-in-time attestation header values. +/// +/// Host integrations provide the opaque string expected by the upstream +/// `x-oai-attestation` header. Core validates only that it is legal as an HTTP +/// header value before forwarding it. +#[derive(Clone)] +pub struct AttestationProvider { + generate: Arc, +} + +impl fmt::Debug for AttestationProvider { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.debug_struct("AttestationProvider").finish() + } +} + +impl AttestationProvider { + pub fn new(generate: impl Fn() -> GenerateAttestationFuture + Send + Sync + 'static) -> Self { + Self { + generate: Arc::new(generate), + } + } + + pub(crate) async fn generate_header(&self) -> Option { + HeaderValue::from_str(&(self.generate)().await?).ok() + } +} diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 39e6e85e2020..1f72ca93ea93 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -105,6 +105,8 @@ use tracing::instrument; use tracing::trace; use tracing::warn; +use crate::attestation::AttestationProvider; +use crate::attestation::X_OAI_ATTESTATION_HEADER; use crate::client_common::Prompt; use crate::client_common::ResponseEvent; use crate::client_common::ResponseStream; @@ -118,6 +120,7 @@ use codex_login::auth_env_telemetry::AuthEnvTelemetry; use codex_login::auth_env_telemetry::collect_auth_env_telemetry; use codex_model_provider::SharedModelProvider; use codex_model_provider::create_model_provider; +use codex_model_provider_info::CHATGPT_CODEX_BASE_URL; #[cfg(test)] use codex_model_provider_info::DEFAULT_WEBSOCKET_CONNECT_TIMEOUT_MS; use codex_model_provider_info::ModelProviderInfo; @@ -170,6 +173,7 @@ struct ModelClientState { enable_request_compression: bool, include_timing_metrics: bool, beta_features_header: Option, + attestation_provider: Option, disable_websockets: AtomicBool, cached_websocket_session: StdMutex, } @@ -314,6 +318,35 @@ impl ModelClient { enable_request_compression: bool, include_timing_metrics: bool, beta_features_header: Option, + ) -> Self { + Self::new_with_attestation_provider( + auth_manager, + session_id, + thread_id, + installation_id, + provider_info, + session_source, + model_verbosity, + enable_request_compression, + include_timing_metrics, + beta_features_header, + /*attestation_provider*/ None, + ) + } + + #[allow(clippy::too_many_arguments)] + pub(crate) fn new_with_attestation_provider( + auth_manager: Option>, + session_id: SessionId, + thread_id: ThreadId, + installation_id: String, + provider_info: ModelProviderInfo, + session_source: SessionSource, + model_verbosity: Option, + enable_request_compression: bool, + include_timing_metrics: bool, + beta_features_header: Option, + attestation_provider: Option, ) -> Self { let model_provider = create_model_provider(provider_info, auth_manager); let codex_api_key_env_enabled = model_provider @@ -335,6 +368,7 @@ impl ModelClient { enable_request_compression, include_timing_metrics, beta_features_header, + attestation_provider, disable_websockets: AtomicBool::new(false), cached_websocket_session: StdMutex::new(WebsocketSession::default()), }), @@ -463,9 +497,6 @@ impl ModelClient { text, .. } = request; - let client = - ApiCompactClient::new(transport, client_setup.api_provider, client_setup.api_auth) - .with_telemetry(Some(request_telemetry)); let payload = ApiCompactionInput { model: &model, input: &input, @@ -492,6 +523,15 @@ impl ModelClient { Some(self.state.session_id.to_string()), Some(self.state.thread_id.to_string()), )); + self.extend_attestation_header_for( + &mut extra_headers, + &client_setup.api_provider, + AttestationPurpose::Compaction, + ) + .await; + let client = + ApiCompactClient::new(transport, client_setup.api_provider, client_setup.api_auth) + .with_telemetry(Some(request_telemetry)); let trace_attempt = compaction_trace.start_attempt(&payload); let result = client .compact_input(&payload, extra_headers) @@ -505,11 +545,17 @@ impl ModelClient { &self, sdp: String, session_config: ApiRealtimeSessionConfig, - extra_headers: ApiHeaderMap, + mut extra_headers: ApiHeaderMap, ) -> 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?; + self.extend_attestation_header_for( + &mut extra_headers, + &client_setup.api_provider, + AttestationPurpose::RealtimeWebrtcCallSetup, + ) + .await; let mut sideband_headers = extra_headers.clone(); sideband_headers.extend(sideband_websocket_auth_headers( client_setup.api_auth.as_ref(), @@ -640,6 +686,25 @@ impl ModelClient { client_metadata } + async fn generate_attestation_header(&self) -> Option { + self.state + .attestation_provider + .as_ref()? + .generate_header() + .await + } + + async fn generate_attestation_header_for( + &self, + provider: &codex_api::Provider, + purpose: AttestationPurpose, + ) -> Option { + if !should_send_attestation(provider, purpose) { + return None; + } + self.generate_attestation_header().await + } + /// Builds request telemetry for unary API calls (e.g., Compact endpoint). fn build_request_telemetry( session_telemetry: &SessionTelemetry, @@ -777,7 +842,9 @@ impl ModelClient { auth_context: AuthRequestTelemetryContext, request_route_telemetry: RequestRouteTelemetry, ) -> std::result::Result { - let headers = self.build_websocket_headers(turn_state.as_ref(), turn_metadata_header); + let headers = self + .build_websocket_headers(&api_provider, turn_state.as_ref(), turn_metadata_header) + .await; let websocket_telemetry = ModelClientSession::build_websocket_telemetry( session_telemetry, auth_context, @@ -854,8 +921,9 @@ impl ModelClient { /// /// Callers should pass the current turn-state lock when available so sticky-routing state is /// replayed on reconnect within the same turn. - fn build_websocket_headers( + async fn build_websocket_headers( &self, + provider: &codex_api::Provider, turn_state: Option<&Arc>>, turn_metadata_header: Option<&str>, ) -> ApiHeaderMap { @@ -872,6 +940,8 @@ impl ModelClient { } headers.extend(build_session_headers(Some(session_id), Some(thread_id))); headers.extend(self.build_responses_identity_headers()); + self.extend_attestation_header_for(&mut headers, provider, AttestationPurpose::Response) + .await; headers.insert( OPENAI_BETA_HEADER, HeaderValue::from_static(RESPONSES_WEBSOCKETS_V2_BETA_HEADER_VALUE), @@ -920,8 +990,9 @@ impl ModelClientSession { /// /// Keeping option construction in one place ensures request-scoped headers are consistent /// regardless of transport choice. - fn build_responses_options( + async fn build_responses_options( &self, + provider: &codex_api::Provider, turn_metadata_header: Option<&str>, compression: Compression, ) -> ApiResponsesOptions { @@ -939,6 +1010,13 @@ impl ModelClientSession { turn_metadata_header.as_ref(), ); headers.extend(self.client.build_responses_identity_headers()); + self.client + .extend_attestation_header_for( + &mut headers, + provider, + AttestationPurpose::Response, + ) + .await; headers }, compression, @@ -1215,7 +1293,13 @@ impl ModelClientSession { self.client.state.auth_env_telemetry.clone(), ); let compression = self.responses_request_compression(client_setup.auth.as_ref()); - let options = self.build_responses_options(turn_metadata_header, compression); + let options = self + .build_responses_options( + &client_setup.api_provider, + turn_metadata_header, + compression, + ) + .await; let request = self.client.build_responses_request( &client_setup.api_provider, @@ -1322,7 +1406,13 @@ impl ModelClientSession { ); let compression = self.responses_request_compression(client_setup.auth.as_ref()); - let options = self.build_responses_options(turn_metadata_header, compression); + let options = self + .build_responses_options( + &client_setup.api_provider, + turn_metadata_header, + compression, + ) + .await; let request = self.client.build_responses_request( &client_setup.api_provider, prompt, @@ -1626,6 +1716,43 @@ fn build_responses_headers( headers } +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum AttestationPurpose { + Response, + Compaction, + RealtimeWebrtcCallSetup, +} + +fn should_send_attestation(provider: &codex_api::Provider, purpose: AttestationPurpose) -> bool { + let provider_is_chatgpt_codex = provider + .base_url + .trim_end_matches('/') + .eq_ignore_ascii_case(CHATGPT_CODEX_BASE_URL); + provider_is_chatgpt_codex + && matches!( + purpose, + AttestationPurpose::Response + | AttestationPurpose::Compaction + | AttestationPurpose::RealtimeWebrtcCallSetup + ) +} + +impl ModelClient { + async fn extend_attestation_header_for( + &self, + headers: &mut ApiHeaderMap, + provider: &codex_api::Provider, + purpose: AttestationPurpose, + ) { + if let Some(header_value) = self + .generate_attestation_header_for(provider, purpose) + .await + { + headers.insert(X_OAI_ATTESTATION_HEADER, header_value); + } + } +} + fn subagent_header_value(session_source: &SessionSource) -> Option { match session_source { SessionSource::SubAgent(subagent_source) => match subagent_source { diff --git a/codex-rs/core/src/client_tests.rs b/codex-rs/core/src/client_tests.rs index 2ba65d7c453d..622bddb77aed 100644 --- a/codex-rs/core/src/client_tests.rs +++ b/codex-rs/core/src/client_tests.rs @@ -1,3 +1,4 @@ +use super::AttestationPurpose; use super::AuthRequestTelemetryContext; use super::ModelClient; use super::PendingUnauthorizedRetry; @@ -7,6 +8,7 @@ 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::AttestationProvider; use codex_api::ApiError; use codex_api::ResponseEvent; use codex_app_server_protocol::AuthMode; @@ -14,6 +16,7 @@ use codex_model_provider::BearerAuthProvider; use codex_model_provider_info::WireApi; use codex_model_provider_info::create_oss_provider_with_base_url; use codex_otel::SessionTelemetry; +use codex_protocol::SessionId; use codex_protocol::ThreadId; use codex_protocol::models::ContentItem; use codex_protocol::models::ResponseItem; @@ -36,6 +39,8 @@ use std::collections::VecDeque; use std::pin::Pin; use std::sync::Arc; use std::sync::Mutex; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; use std::task::Context; use std::task::Poll; use std::time::Duration; @@ -67,6 +72,23 @@ fn test_model_client(session_source: SessionSource) -> ModelClient { ) } +fn api_provider(base_url: &str) -> codex_api::Provider { + codex_api::Provider { + name: "test".to_string(), + base_url: base_url.to_string(), + query_params: None, + headers: http::HeaderMap::new(), + retry: codex_api::RetryConfig { + max_attempts: 1, + base_delay: Duration::from_millis(1), + retry_429: false, + retry_5xx: true, + retry_transport: true, + }, + stream_idle_timeout: Duration::from_secs(1), + } +} + fn test_model_info() -> ModelInfo { serde_json::from_value(json!({ "slug": "gpt-test", @@ -466,3 +488,219 @@ 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")); } + +fn model_client_with_counting_attestation() -> (ModelClient, Arc) { + let attestation_calls = Arc::new(AtomicUsize::new(0)); + let calls = attestation_calls.clone(); + let model_client = ModelClient::new_with_attestation_provider( + /*auth_manager*/ None, + SessionId::new(), + ThreadId::new(), + /*installation_id*/ "11111111-1111-4111-8111-111111111111".to_string(), + create_oss_provider_with_base_url("https://example.com/v1", WireApi::Responses), + SessionSource::Exec, + /*model_verbosity*/ None, + /*enable_request_compression*/ false, + /*include_timing_metrics*/ false, + /*beta_features_header*/ None, + Some(AttestationProvider::new(move || { + let calls = calls.clone(); + Box::pin(async move { + let call = calls.fetch_add(1, Ordering::Relaxed) + 1; + Some(format!("v1.header-{call}")) + }) + })), + ); + (model_client, attestation_calls) +} + +#[test] +fn should_send_attestation_for_allowed_chatgpt_codex_purposes() { + let provider = api_provider("https://chatgpt.com/backend-api/codex/"); + + for purpose in [ + AttestationPurpose::Response, + AttestationPurpose::Compaction, + AttestationPurpose::RealtimeWebrtcCallSetup, + ] { + assert!(super::should_send_attestation(&provider, purpose)); + } +} + +#[test] +fn should_not_send_attestation_for_non_chatgpt_codex_provider() { + let provider = api_provider("https://api.openai.com/v1"); + + assert!(!super::should_send_attestation( + &provider, + AttestationPurpose::Response, + )); +} + +#[tokio::test] +async fn responses_generate_fresh_attestation_headers_for_chatgpt_codex() { + let provider = api_provider("https://chatgpt.com/backend-api/codex/"); + let (model_client, attestation_calls) = model_client_with_counting_attestation(); + let mut first_headers = http::HeaderMap::new(); + let mut second_headers = http::HeaderMap::new(); + + model_client + .extend_attestation_header_for(&mut first_headers, &provider, AttestationPurpose::Response) + .await; + model_client + .extend_attestation_header_for(&mut second_headers, &provider, AttestationPurpose::Response) + .await; + + assert_eq!( + first_headers + .get(crate::attestation::X_OAI_ATTESTATION_HEADER) + .and_then(|value| value.to_str().ok()), + Some("v1.header-1"), + ); + assert_eq!( + second_headers + .get(crate::attestation::X_OAI_ATTESTATION_HEADER) + .and_then(|value| value.to_str().ok()), + Some("v1.header-2"), + ); + assert_eq!(attestation_calls.load(Ordering::Relaxed), 2); +} + +#[tokio::test] +async fn websocket_handshake_includes_attestation_for_chatgpt_codex_responses() { + let provider = api_provider("https://chatgpt.com/backend-api/codex/"); + let (model_client, attestation_calls) = model_client_with_counting_attestation(); + + let headers = model_client + .build_websocket_headers( + &provider, /*turn_state*/ None, /*turn_metadata_header*/ None, + ) + .await; + + assert_eq!( + headers + .get(crate::attestation::X_OAI_ATTESTATION_HEADER) + .and_then(|value| value.to_str().ok()), + Some("v1.header-1"), + ); + assert_eq!(attestation_calls.load(Ordering::Relaxed), 1); +} + +#[tokio::test] +async fn compact_generate_fresh_attestation_headers_for_chatgpt_codex() { + let provider = api_provider("https://chatgpt.com/backend-api/codex/"); + let (model_client, attestation_calls) = model_client_with_counting_attestation(); + let mut first_headers = http::HeaderMap::new(); + let mut second_headers = http::HeaderMap::new(); + + model_client + .extend_attestation_header_for( + &mut first_headers, + &provider, + AttestationPurpose::Compaction, + ) + .await; + model_client + .extend_attestation_header_for( + &mut second_headers, + &provider, + AttestationPurpose::Compaction, + ) + .await; + + assert_eq!( + first_headers + .get(crate::attestation::X_OAI_ATTESTATION_HEADER) + .and_then(|value| value.to_str().ok()), + Some("v1.header-1"), + ); + assert_eq!( + second_headers + .get(crate::attestation::X_OAI_ATTESTATION_HEADER) + .and_then(|value| value.to_str().ok()), + Some("v1.header-2"), + ); + assert_eq!(attestation_calls.load(Ordering::Relaxed), 2); +} + +#[tokio::test] +async fn realtime_setup_generate_fresh_attestation_headers_for_chatgpt_codex() { + let provider = api_provider("https://chatgpt.com/backend-api/codex/"); + let (model_client, attestation_calls) = model_client_with_counting_attestation(); + let mut first_headers = http::HeaderMap::new(); + let mut second_headers = http::HeaderMap::new(); + + model_client + .extend_attestation_header_for( + &mut first_headers, + &provider, + AttestationPurpose::RealtimeWebrtcCallSetup, + ) + .await; + model_client + .extend_attestation_header_for( + &mut second_headers, + &provider, + AttestationPurpose::RealtimeWebrtcCallSetup, + ) + .await; + + assert_eq!( + first_headers + .get(crate::attestation::X_OAI_ATTESTATION_HEADER) + .and_then(|value| value.to_str().ok()), + Some("v1.header-1"), + ); + assert_eq!( + second_headers + .get(crate::attestation::X_OAI_ATTESTATION_HEADER) + .and_then(|value| value.to_str().ok()), + Some("v1.header-2"), + ); + assert_eq!(attestation_calls.load(Ordering::Relaxed), 2); +} + +#[tokio::test] +async fn non_chatgpt_codex_endpoints_omit_attestation_generation() { + let provider = api_provider("https://api.openai.com/v1"); + let (model_client, attestation_calls) = model_client_with_counting_attestation(); + let mut response_headers = http::HeaderMap::new(); + + model_client + .extend_attestation_header_for( + &mut response_headers, + &provider, + AttestationPurpose::Response, + ) + .await; + let mut compaction_headers = http::HeaderMap::new(); + model_client + .extend_attestation_header_for( + &mut compaction_headers, + &provider, + AttestationPurpose::Compaction, + ) + .await; + let mut realtime_headers = http::HeaderMap::new(); + model_client + .extend_attestation_header_for( + &mut realtime_headers, + &provider, + AttestationPurpose::RealtimeWebrtcCallSetup, + ) + .await; + + assert_eq!( + response_headers.get(crate::attestation::X_OAI_ATTESTATION_HEADER), + None, + ); + assert_eq!( + compaction_headers.get(crate::attestation::X_OAI_ATTESTATION_HEADER), + None, + ); + assert_eq!( + realtime_headers.get(crate::attestation::X_OAI_ATTESTATION_HEADER), + None, + ); + assert_eq!(attestation_calls.load(Ordering::Relaxed), 0); +} diff --git a/codex-rs/core/src/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index a89d8fc9737c..25037649046f 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/codex-rs/core/src/codex_delegate.rs @@ -99,6 +99,7 @@ pub(crate) async fn run_codex_thread_interactive( environment_selections: parent_ctx.environments.clone(), analytics_events_client: Some(parent_session.services.analytics_events_client.clone()), thread_store: Arc::clone(&parent_session.services.thread_store), + attestation_provider: parent_session.services.attestation_provider.clone(), })) .or_cancel(&cancel_token) .await??; diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 0cdf0e2d4669..c178b5d4e8bd 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -23,6 +23,7 @@ pub use codex_thread::CodexThread; pub use codex_thread::CodexThreadTurnContextOverrides; pub use codex_thread::ThreadConfigSnapshot; mod agent; +mod attestation; mod codex_delegate; mod command_canonicalization; mod commit_attribution; @@ -177,6 +178,7 @@ mod tasks; mod user_shell_command; pub mod util; +pub use attestation::AttestationProvider; pub use client::ModelClient; pub use client::ModelClientSession; pub use client::X_CODEX_INSTALLATION_ID_HEADER; diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 6d4b27542182..07ef550f5179 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -14,6 +14,7 @@ use crate::agent::Mailbox; use crate::agent::MailboxReceiver; use crate::agent::agent_status_from_event; use crate::agent::status::is_final; +use crate::attestation::AttestationProvider; use crate::build_available_skills; use crate::commit_attribution::commit_message_trailer_instruction; use crate::compact; @@ -412,6 +413,7 @@ pub(crate) struct CodexSpawnArgs { pub(crate) environment_selections: ResolvedTurnEnvironments, pub(crate) analytics_events_client: Option, pub(crate) thread_store: Arc, + pub(crate) attestation_provider: Option, } pub(crate) const INITIAL_SUBMIT_ID: &str = ""; @@ -471,6 +473,7 @@ impl Codex { environment_selections, analytics_events_client, thread_store, + attestation_provider, } = args; let (tx_sub, rx_sub) = async_channel::bounded(SUBMISSION_CHANNEL_CAPACITY); let (tx_event, rx_event) = async_channel::unbounded(); @@ -656,6 +659,7 @@ impl Codex { analytics_events_client, thread_store, parent_rollout_thread_trace, + attestation_provider, ) .await .map_err(|e| { diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index f72a173c80fb..828ecca4c2c4 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -370,6 +370,7 @@ impl Session { analytics_events_client: Option, thread_store: Arc, parent_rollout_thread_trace: ThreadTraceContext, + attestation_provider: Option, ) -> anyhow::Result> { debug!( "Configuring session: model={}; provider={:?}", @@ -852,7 +853,8 @@ impl Session { state_db: state_db_ctx.clone(), live_thread: live_thread_init.as_ref().cloned(), thread_store: Arc::clone(&thread_store), - model_client: ModelClient::new( + attestation_provider: attestation_provider.clone(), + model_client: ModelClient::new_with_attestation_provider( Some(Arc::clone(&auth_manager)), session_id, thread_id, @@ -863,6 +865,7 @@ impl Session { config.features.enabled(Feature::EnableRequestCompression), config.features.enabled(Feature::RuntimeMetrics), Self::build_model_client_beta_features_header(config.as_ref()), + attestation_provider, ), code_mode_service: crate::tools::code_mode::CodeModeService::new(), environment_manager, diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index b63b16cbf4f7..67fe370ac7ba 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -3733,6 +3733,7 @@ async fn session_new_fails_when_zsh_fork_enabled_without_zsh_path() { /*state_db*/ None, )), codex_rollout_trace::ThreadTraceContext::disabled(), + /*attestation_provider*/ None, ) .await; @@ -3881,6 +3882,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { codex_thread_store::LocalThreadStoreConfig::from_config(config.as_ref()), /*state_db*/ None, )), + attestation_provider: None, model_client: ModelClient::new( Some(auth_manager.clone()), thread_id.into(), @@ -4069,6 +4071,7 @@ async fn make_session_with_config_and_rx( /*state_db*/ None, )), codex_rollout_trace::ThreadTraceContext::disabled(), + /*attestation_provider*/ None, ) .await?; @@ -4178,6 +4181,7 @@ async fn make_session_with_history_source_and_agent_control_and_rx( ), )), codex_rollout_trace::ThreadTraceContext::disabled(), + /*attestation_provider*/ None, ) .await?; @@ -5596,6 +5600,7 @@ where codex_thread_store::LocalThreadStoreConfig::from_config(config.as_ref()), state_db, )), + attestation_provider: None, model_client: ModelClient::new( Some(Arc::clone(&auth_manager)), thread_id.into(), diff --git a/codex-rs/core/src/session/tests/guardian_tests.rs b/codex-rs/core/src/session/tests/guardian_tests.rs index 5c473ef1f9d4..37a5e6bb09f2 100644 --- a/codex-rs/core/src/session/tests/guardian_tests.rs +++ b/codex-rs/core/src/session/tests/guardian_tests.rs @@ -763,6 +763,7 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() { }, analytics_events_client: None, thread_store, + attestation_provider: None, }) .await .expect("spawn guardian subagent"); diff --git a/codex-rs/core/src/state/service.rs b/codex-rs/core/src/state/service.rs index 9cd9e97fbba7..4506c0054cf7 100644 --- a/codex-rs/core/src/state/service.rs +++ b/codex-rs/core/src/state/service.rs @@ -3,6 +3,7 @@ use std::sync::Arc; use crate::SkillsManager; use crate::agent::AgentControl; +use crate::attestation::AttestationProvider; use crate::client::ModelClient; use crate::config::StartedNetworkProxy; use crate::exec_policy::ExecPolicyManager; @@ -66,6 +67,7 @@ pub(crate) struct SessionServices { pub(crate) state_db: Option, pub(crate) live_thread: Option, pub(crate) thread_store: Arc, + pub(crate) attestation_provider: Option, /// Session-scoped model client shared across turns. pub(crate) model_client: ModelClient, pub(crate) code_mode_service: CodeModeService, diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index a19832717d3b..b1bfdc47b0d0 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -1,5 +1,6 @@ use crate::SkillsManager; use crate::agent::AgentControl; +use crate::attestation::AttestationProvider; use crate::codex_thread::CodexThread; use crate::config::Config; use crate::config::ThreadStoreConfig; @@ -249,6 +250,7 @@ pub(crate) struct ThreadManagerState { mcp_manager: Arc, skills_watcher: Arc, thread_store: Arc, + attestation_provider: Option, session_source: SessionSource, installation_id: String, analytics_events_client: Option, @@ -293,6 +295,31 @@ impl ThreadManager { thread_store: Arc, state_db: Option, installation_id: String, + ) -> Self { + Self::new_with_attestation_provider( + config, + auth_manager, + session_source, + environment_manager, + analytics_events_client, + thread_store, + state_db, + installation_id, + /*attestation_provider*/ None, + ) + } + + #[allow(clippy::too_many_arguments)] + pub fn new_with_attestation_provider( + config: &Config, + auth_manager: Arc, + session_source: SessionSource, + environment_manager: Arc, + analytics_events_client: Option, + thread_store: Arc, + state_db: Option, + installation_id: String, + attestation_provider: Option, ) -> Self { let codex_home = config.codex_home.clone(); let restriction_product = session_source.restriction_product(); @@ -319,6 +346,7 @@ impl ThreadManager { mcp_manager, skills_watcher, thread_store, + attestation_provider, auth_manager, session_source, installation_id, @@ -420,6 +448,7 @@ impl ThreadManager { mcp_manager, skills_watcher, thread_store, + attestation_provider: None, auth_manager, session_source: SessionSource::Exec, installation_id, @@ -1206,6 +1235,7 @@ impl ThreadManagerState { environment_selections, analytics_events_client: self.analytics_events_client.clone(), thread_store: Arc::clone(&self.thread_store), + attestation_provider: self.attestation_provider.clone(), }) .await?; let new_thread = self diff --git a/codex-rs/debug-client/src/client.rs b/codex-rs/debug-client/src/client.rs index 2edabfac0084..69eb474aaf69 100644 --- a/codex-rs/debug-client/src/client.rs +++ b/codex-rs/debug-client/src/client.rs @@ -103,6 +103,7 @@ impl AppServerClient { }, capabilities: Some(InitializeCapabilities { experimental_api: true, + request_attestation: false, opt_out_notification_methods: None, }), }, diff --git a/codex-rs/exec/src/lib.rs b/codex-rs/exec/src/lib.rs index b035a195172b..6e79b8bb3dfe 100644 --- a/codex-rs/exec/src/lib.rs +++ b/codex-rs/exec/src/lib.rs @@ -1604,6 +1604,15 @@ async fn handle_server_request( ) .await } + ServerRequest::AttestationGenerate { request_id, .. } => { + reject_server_request( + client, + request_id, + &method, + "attestation generation is not supported in exec mode".to_string(), + ) + .await + } ServerRequest::ApplyPatchApproval { request_id, params } => { reject_server_request( client, diff --git a/codex-rs/model-provider-info/src/lib.rs b/codex-rs/model-provider-info/src/lib.rs index 0fb8be474690..6fca7e6a1f59 100644 --- a/codex-rs/model-provider-info/src/lib.rs +++ b/codex-rs/model-provider-info/src/lib.rs @@ -34,6 +34,7 @@ const MAX_REQUEST_MAX_RETRIES: u64 = 100; const OPENAI_PROVIDER_NAME: &str = "OpenAI"; pub const OPENAI_PROVIDER_ID: &str = "openai"; +pub const CHATGPT_CODEX_BASE_URL: &str = "https://chatgpt.com/backend-api/codex"; const AMAZON_BEDROCK_PROVIDER_NAME: &str = "Amazon Bedrock"; pub const AMAZON_BEDROCK_PROVIDER_ID: &str = "amazon-bedrock"; pub const AMAZON_BEDROCK_DEFAULT_BASE_URL: &str = @@ -234,7 +235,7 @@ impl ModelProviderInfo { auth_mode, Some(AuthMode::Chatgpt | AuthMode::ChatgptAuthTokens | AuthMode::AgentIdentity) ) { - "https://chatgpt.com/backend-api/codex" + CHATGPT_CODEX_BASE_URL } else { "https://api.openai.com/v1" }; diff --git a/codex-rs/tui/src/app/app_server_event_targets.rs b/codex-rs/tui/src/app/app_server_event_targets.rs index 382a82a19f5b..d535bf8e3d83 100644 --- a/codex-rs/tui/src/app/app_server_event_targets.rs +++ b/codex-rs/tui/src/app/app_server_event_targets.rs @@ -25,6 +25,7 @@ pub(super) fn server_request_thread_id(request: &ServerRequest) -> Option None, } diff --git a/codex-rs/tui/src/app/app_server_requests.rs b/codex-rs/tui/src/app/app_server_requests.rs index 4b587b0fc894..ff3c755b6ee8 100644 --- a/codex-rs/tui/src/app/app_server_requests.rs +++ b/codex-rs/tui/src/app/app_server_requests.rs @@ -134,6 +134,12 @@ impl PendingAppServerRequests { }) } ServerRequest::ChatgptAuthTokensRefresh { .. } => None, + ServerRequest::AttestationGenerate { request_id, .. } => { + Some(UnsupportedAppServerRequest { + request_id: request_id.clone(), + message: "Attestation generation is not available in TUI.".to_string(), + }) + } ServerRequest::ApplyPatchApproval { request_id, .. } => { Some(UnsupportedAppServerRequest { request_id: request_id.clone(), @@ -332,6 +338,7 @@ impl PendingAppServerRequests { .any(|pending_request_id| pending_request_id == request_id), ServerRequest::DynamicToolCall { .. } | ServerRequest::ChatgptAuthTokensRefresh { .. } + | ServerRequest::AttestationGenerate { .. } | ServerRequest::ApplyPatchApproval { .. } | ServerRequest::ExecCommandApproval { .. } => true, } diff --git a/codex-rs/tui/src/app/side.rs b/codex-rs/tui/src/app/side.rs index 59f3d71991fb..d3ea62da70bb 100644 --- a/codex-rs/tui/src/app/side.rs +++ b/codex-rs/tui/src/app/side.rs @@ -92,6 +92,7 @@ impl SideParentStatus { | ServerRequest::ApplyPatchApproval { .. } | ServerRequest::ExecCommandApproval { .. } => Some(SideParentStatus::NeedsApproval), ServerRequest::DynamicToolCall { .. } + | ServerRequest::AttestationGenerate { .. } | ServerRequest::ChatgptAuthTokensRefresh { .. } => None, } } diff --git a/codex-rs/tui/src/chatwidget.rs b/codex-rs/tui/src/chatwidget.rs index adaba76b24fa..1f6774da152f 100644 --- a/codex-rs/tui/src/chatwidget.rs +++ b/codex-rs/tui/src/chatwidget.rs @@ -6204,6 +6204,7 @@ impl ChatWidget { self.on_request_user_input(params); } ServerRequest::DynamicToolCall { .. } + | ServerRequest::AttestationGenerate { .. } | ServerRequest::ChatgptAuthTokensRefresh { .. } | ServerRequest::ApplyPatchApproval { .. } | ServerRequest::ExecCommandApproval { .. } => { From bc141c4f277cd06579ef7017af6e0fa9022ec7ae Mon Sep 17 00:00:00 2001 From: Jiaming Zhang Date: Wed, 6 May 2026 16:52:46 -0700 Subject: [PATCH 02/23] test(app-server): cover attestation websocket flow Co-authored-by: Codex --- .../app-server/tests/suite/v2/attestation.rs | 193 ++++++++++++++++++ codex-rs/app-server/tests/suite/v2/mod.rs | 1 + codex-rs/codex-api/src/provider.rs | 1 + codex-rs/core/src/client.rs | 7 +- codex-rs/core/src/client_tests.rs | 4 + codex-rs/model-provider-info/src/lib.rs | 4 + 6 files changed, 204 insertions(+), 6 deletions(-) create mode 100644 codex-rs/app-server/tests/suite/v2/attestation.rs diff --git a/codex-rs/app-server/tests/suite/v2/attestation.rs b/codex-rs/app-server/tests/suite/v2/attestation.rs new file mode 100644 index 000000000000..5030b61f2ea7 --- /dev/null +++ b/codex-rs/app-server/tests/suite/v2/attestation.rs @@ -0,0 +1,193 @@ +use anyhow::Result; +use anyhow::bail; +use app_test_support::ChatGptAuthFixture; +use app_test_support::McpProcess; +use app_test_support::to_response; +use app_test_support::write_chatgpt_auth; +use codex_app_server_protocol::AttestationGenerateResponse; +use codex_app_server_protocol::ClientInfo; +use codex_app_server_protocol::InitializeCapabilities; +use codex_app_server_protocol::JSONRPCMessage; +use codex_app_server_protocol::JSONRPCResponse; +use codex_app_server_protocol::RequestId; +use codex_app_server_protocol::ServerRequest; +use codex_app_server_protocol::ThreadStartParams; +use codex_app_server_protocol::ThreadStartResponse; +use codex_app_server_protocol::TurnStartParams; +use codex_app_server_protocol::TurnStartResponse; +use codex_app_server_protocol::UserInput as V2UserInput; +use codex_config::types::AuthCredentialsStoreMode; +use core_test_support::responses; +use core_test_support::responses::WebSocketConnectionConfig; +use core_test_support::responses::start_websocket_server_with_headers; +use core_test_support::skip_if_no_network; +use pretty_assertions::assert_eq; +use std::path::Path; +use tempfile::TempDir; +use tokio::time::Duration; +use tokio::time::timeout; + +const DEFAULT_READ_TIMEOUT: Duration = Duration::from_secs(60); +const ATTESTATION_HEADER: &str = "v1.integration-test"; + +#[tokio::test] +async fn attestation_generate_round_trip_adds_header_to_responses_websocket_handshake() -> Result<()> +{ + skip_if_no_network!(Ok(())); + + let websocket_server = start_websocket_server_with_headers(vec![ + // App-server refreshes `/models` over HTTP during thread startup. It points at the same + // local test base URL, so let that non-websocket probe consume one connection before the + // websocket handshake under test arrives. + WebSocketConnectionConfig { + requests: Vec::new(), + response_headers: Vec::new(), + accept_delay: None, + close_after_requests: true, + }, + WebSocketConnectionConfig { + requests: vec![ + vec![ + responses::ev_response_created("warm-1"), + responses::ev_completed("warm-1"), + ], + vec![ + responses::ev_response_created("resp-1"), + responses::ev_assistant_message("msg-1", "Done"), + responses::ev_completed("resp-1"), + ], + ], + response_headers: Vec::new(), + accept_delay: None, + close_after_requests: true, + }, + ]) + .await; + + let codex_home = TempDir::new()?; + create_chatgpt_websocket_config( + codex_home.path(), + &websocket_server.uri().replacen("ws://", "http://", 1), + )?; + write_chatgpt_auth( + codex_home.path(), + ChatGptAuthFixture::new("access-chatgpt").plan_type("pro"), + AuthCredentialsStoreMode::File, + )?; + + let mut mcp = McpProcess::new_with_env(codex_home.path(), &[("OPENAI_API_KEY", None)]).await?; + let initialized = timeout( + DEFAULT_READ_TIMEOUT, + mcp.initialize_with_capabilities( + ClientInfo { + name: "codex_desktop".to_string(), + title: Some("Codex Desktop".to_string()), + version: "0.1.0".to_string(), + }, + Some(InitializeCapabilities { + experimental_api: true, + request_attestation: true, + opt_out_notification_methods: None, + }), + ), + ) + .await??; + let JSONRPCMessage::Response(_) = initialized else { + bail!("expected initialize response, got {initialized:?}"); + }; + + let thread_request_id = mcp + .send_thread_start_request(ThreadStartParams::default()) + .await?; + let thread_response: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(thread_request_id)), + ) + .await??; + let ThreadStartResponse { thread, .. } = to_response(thread_response)?; + + let turn_request_id = mcp + .send_turn_start_request(TurnStartParams { + thread_id: thread.id, + input: vec![V2UserInput::Text { + text: "Hello".to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + }) + .await?; + let turn_response: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(turn_request_id)), + ) + .await??; + let _: TurnStartResponse = to_response(turn_response)?; + + let mut attestation_requests = 0; + timeout(DEFAULT_READ_TIMEOUT, async { + loop { + match mcp.read_next_message().await? { + JSONRPCMessage::Request(request) => { + let request = ServerRequest::try_from(request)?; + let ServerRequest::AttestationGenerate { request_id, .. } = request else { + bail!("expected attestation/generate request, got {request:?}"); + }; + attestation_requests += 1; + mcp.send_response( + request_id, + serde_json::to_value(AttestationGenerateResponse { + header_value: ATTESTATION_HEADER.to_string(), + })?, + ) + .await?; + } + JSONRPCMessage::Notification(notification) + if notification.method == "turn/completed" => + { + break Ok(()); + } + _ => {} + } + } + }) + .await??; + assert!(attestation_requests > 0); + + assert!( + websocket_server + .wait_for_handshakes(/*expected*/ 1, DEFAULT_READ_TIMEOUT) + .await + ); + let handshake = websocket_server.single_handshake(); + assert_eq!( + handshake.header("x-oai-attestation").as_deref(), + Some(ATTESTATION_HEADER) + ); + + websocket_server.shutdown().await; + Ok(()) +} + +fn create_chatgpt_websocket_config(codex_home: &Path, server_uri: &str) -> std::io::Result<()> { + std::fs::write( + codex_home.join("config.toml"), + format!( + r#" +model = "mock-model" +approval_policy = "never" +sandbox_mode = "read-only" + +model_provider = "mock_provider" + +[model_providers.mock_provider] +name = "Mock ChatGPT provider for test" +base_url = "{server_uri}/v1" +wire_api = "responses" +request_max_retries = 0 +stream_max_retries = 0 +requires_openai_auth = true +supports_websockets = true +"# + ), + ) +} diff --git a/codex-rs/app-server/tests/suite/v2/mod.rs b/codex-rs/app-server/tests/suite/v2/mod.rs index 8e13df7825f4..642be8ad4ab2 100644 --- a/codex-rs/app-server/tests/suite/v2/mod.rs +++ b/codex-rs/app-server/tests/suite/v2/mod.rs @@ -1,6 +1,7 @@ mod account; mod analytics; mod app_list; +mod attestation; mod client_metadata; mod collaboration_mode_list; #[cfg(unix)] diff --git a/codex-rs/codex-api/src/provider.rs b/codex-rs/codex-api/src/provider.rs index 45f2512dc39a..75a1062f65e9 100644 --- a/codex-rs/codex-api/src/provider.rs +++ b/codex-rs/codex-api/src/provider.rs @@ -43,6 +43,7 @@ impl RetryConfig { pub struct Provider { pub name: String, pub base_url: String, + pub uses_chatgpt_auth: bool, pub query_params: Option>, pub headers: HeaderMap, pub retry: RetryConfig, diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 1f72ca93ea93..e9ebf12406a0 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -120,7 +120,6 @@ use codex_login::auth_env_telemetry::AuthEnvTelemetry; use codex_login::auth_env_telemetry::collect_auth_env_telemetry; use codex_model_provider::SharedModelProvider; use codex_model_provider::create_model_provider; -use codex_model_provider_info::CHATGPT_CODEX_BASE_URL; #[cfg(test)] use codex_model_provider_info::DEFAULT_WEBSOCKET_CONNECT_TIMEOUT_MS; use codex_model_provider_info::ModelProviderInfo; @@ -1724,11 +1723,7 @@ enum AttestationPurpose { } fn should_send_attestation(provider: &codex_api::Provider, purpose: AttestationPurpose) -> bool { - let provider_is_chatgpt_codex = provider - .base_url - .trim_end_matches('/') - .eq_ignore_ascii_case(CHATGPT_CODEX_BASE_URL); - provider_is_chatgpt_codex + provider.uses_chatgpt_auth && matches!( purpose, AttestationPurpose::Response diff --git a/codex-rs/core/src/client_tests.rs b/codex-rs/core/src/client_tests.rs index 622bddb77aed..40509ae9d81c 100644 --- a/codex-rs/core/src/client_tests.rs +++ b/codex-rs/core/src/client_tests.rs @@ -13,6 +13,7 @@ use codex_api::ApiError; use codex_api::ResponseEvent; use codex_app_server_protocol::AuthMode; use codex_model_provider::BearerAuthProvider; +use codex_model_provider_info::CHATGPT_CODEX_BASE_URL; use codex_model_provider_info::WireApi; use codex_model_provider_info::create_oss_provider_with_base_url; use codex_otel::SessionTelemetry; @@ -76,6 +77,9 @@ fn api_provider(base_url: &str) -> codex_api::Provider { codex_api::Provider { name: "test".to_string(), base_url: base_url.to_string(), + uses_chatgpt_auth: base_url + .trim_end_matches('/') + .eq_ignore_ascii_case(CHATGPT_CODEX_BASE_URL), query_params: None, headers: http::HeaderMap::new(), retry: codex_api::RetryConfig { diff --git a/codex-rs/model-provider-info/src/lib.rs b/codex-rs/model-provider-info/src/lib.rs index 6fca7e6a1f59..f8bee9ff34a7 100644 --- a/codex-rs/model-provider-info/src/lib.rs +++ b/codex-rs/model-provider-info/src/lib.rs @@ -256,6 +256,10 @@ impl ModelProviderInfo { Ok(ApiProvider { name: self.name.clone(), base_url, + uses_chatgpt_auth: matches!( + auth_mode, + Some(AuthMode::Chatgpt | AuthMode::ChatgptAuthTokens | AuthMode::AgentIdentity) + ), query_params: self.query_params.clone(), headers, retry, From a96ba757081716b7e3f7a4be1e2707afe0e0a368 Mon Sep 17 00:00:00 2001 From: Jiaming Zhang Date: Wed, 6 May 2026 17:24:02 -0700 Subject: [PATCH 03/23] fix(attestation): tighten timeout behavior Co-authored-by: Codex --- codex-rs/app-server/src/message_processor.rs | 10 +++++----- codex-rs/model-provider-info/src/lib.rs | 2 +- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index 81dcaf02a6ab..d35118e0cc07 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -85,7 +85,7 @@ use tokio::time::timeout; use tracing::Instrument; use tracing::warn; -const ATTESTATION_GENERATE_TIMEOUT: Duration = Duration::from_secs(5); +const ATTESTATION_GENERATE_TIMEOUT: Duration = Duration::from_millis(100); const EXTERNAL_AUTH_REFRESH_TIMEOUT: Duration = Duration::from_secs(10); #[derive(Clone)] struct ExternalAuthRefreshBridge { @@ -210,11 +210,11 @@ async fn request_attestation_header_value_with_timeout( message = %err.message, "attestation generation request failed" ); - return None; + return Some(String::new()); } Ok(Err(err)) => { warn!("attestation generation request canceled: {err}"); - return None; + return Some(String::new()); } Err(_) => { let _canceled = outgoing.cancel_request(&request_id).await; @@ -222,7 +222,7 @@ async fn request_attestation_header_value_with_timeout( timeout_seconds = timeout_duration.as_secs(), "attestation generation request timed out" ); - return None; + return Some(String::new()); } }; @@ -230,7 +230,7 @@ async fn request_attestation_header_value_with_timeout( Ok(response) => Some(response.header_value), Err(err) => { warn!("failed to deserialize attestation generation response: {err}"); - None + Some(String::new()) } } } diff --git a/codex-rs/model-provider-info/src/lib.rs b/codex-rs/model-provider-info/src/lib.rs index f8bee9ff34a7..d59a7503683e 100644 --- a/codex-rs/model-provider-info/src/lib.rs +++ b/codex-rs/model-provider-info/src/lib.rs @@ -258,7 +258,7 @@ impl ModelProviderInfo { base_url, uses_chatgpt_auth: matches!( auth_mode, - Some(AuthMode::Chatgpt | AuthMode::ChatgptAuthTokens | AuthMode::AgentIdentity) + Some(AuthMode::Chatgpt | AuthMode::ChatgptAuthTokens) ), query_params: self.query_params.clone(), headers, From e11df6aa3ef4039dc00e5bae98c0bae54c23fd54 Mon Sep 17 00:00:00 2001 From: Jiaming Zhang Date: Wed, 6 May 2026 19:05:21 -0700 Subject: [PATCH 04/23] refactor(attestation): reuse existing constructors Co-authored-by: Codex --- codex-rs/app-server/src/mcp_refresh.rs | 1 + codex-rs/app-server/src/message_processor.rs | 2 +- codex-rs/core/src/client.rs | 28 ------------------- codex-rs/core/src/client_tests.rs | 3 +- codex-rs/core/src/prompt_debug.rs | 1 + codex-rs/core/src/session/session.rs | 2 +- codex-rs/core/src/session/tests.rs | 3 ++ codex-rs/core/src/thread_manager.rs | 24 ---------------- codex-rs/core/src/thread_manager_tests.rs | 10 +++++++ .../src/tools/handlers/multi_agents_tests.rs | 1 + codex-rs/core/tests/common/test_codex.rs | 1 + codex-rs/core/tests/responses_headers.rs | 3 ++ codex-rs/core/tests/suite/client.rs | 3 ++ .../core/tests/suite/client_websockets.rs | 1 + codex-rs/mcp-server/src/message_processor.rs | 1 + codex-rs/memories/write/src/runtime.rs | 1 + codex-rs/thread-manager-sample/src/main.rs | 1 + 17 files changed, 31 insertions(+), 55 deletions(-) diff --git a/codex-rs/app-server/src/mcp_refresh.rs b/codex-rs/app-server/src/mcp_refresh.rs index 8e1ccd3c0aaf..a327b4b125b9 100644 --- a/codex-rs/app-server/src/mcp_refresh.rs +++ b/codex-rs/app-server/src/mcp_refresh.rs @@ -187,6 +187,7 @@ mod tests { thread_store, Some(state_db.clone()), "11111111-1111-4111-8111-111111111111".to_string(), + /*attestation_provider*/ None, )); thread_manager.start_thread(good_config).await?; thread_manager.start_thread(bad_config).await?; diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index d35118e0cc07..be59de451ea9 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -378,7 +378,7 @@ impl MessageProcessor { // affect per-thread behavior, but they must not move newly started, // resumed, or forked threads to a different persistence backend/root. let thread_store = thread_store_from_config(config.as_ref(), state_db.clone()); - let thread_manager = Arc::new(ThreadManager::new_with_attestation_provider( + let thread_manager = Arc::new(ThreadManager::new( config.as_ref(), auth_manager.clone(), session_source, diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index e9ebf12406a0..cb7502b1ba3d 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -317,34 +317,6 @@ impl ModelClient { enable_request_compression: bool, include_timing_metrics: bool, beta_features_header: Option, - ) -> Self { - Self::new_with_attestation_provider( - auth_manager, - session_id, - thread_id, - installation_id, - provider_info, - session_source, - model_verbosity, - enable_request_compression, - include_timing_metrics, - beta_features_header, - /*attestation_provider*/ None, - ) - } - - #[allow(clippy::too_many_arguments)] - pub(crate) fn new_with_attestation_provider( - auth_manager: Option>, - session_id: SessionId, - thread_id: ThreadId, - installation_id: String, - provider_info: ModelProviderInfo, - session_source: SessionSource, - model_verbosity: Option, - enable_request_compression: bool, - include_timing_metrics: bool, - beta_features_header: Option, attestation_provider: Option, ) -> Self { let model_provider = create_model_provider(provider_info, auth_manager); diff --git a/codex-rs/core/src/client_tests.rs b/codex-rs/core/src/client_tests.rs index 40509ae9d81c..bfdc679b698f 100644 --- a/codex-rs/core/src/client_tests.rs +++ b/codex-rs/core/src/client_tests.rs @@ -70,6 +70,7 @@ fn test_model_client(session_source: SessionSource) -> ModelClient { /*enable_request_compression*/ false, /*include_timing_metrics*/ false, /*beta_features_header*/ None, + /*attestation_provider*/ None, ) } @@ -496,7 +497,7 @@ fn auth_request_telemetry_context_tracks_attached_auth_and_retry_phase() { fn model_client_with_counting_attestation() -> (ModelClient, Arc) { let attestation_calls = Arc::new(AtomicUsize::new(0)); let calls = attestation_calls.clone(); - let model_client = ModelClient::new_with_attestation_provider( + let model_client = ModelClient::new( /*auth_manager*/ None, SessionId::new(), ThreadId::new(), diff --git a/codex-rs/core/src/prompt_debug.rs b/codex-rs/core/src/prompt_debug.rs index 8717427afeb5..7a82bb07d9e7 100644 --- a/codex-rs/core/src/prompt_debug.rs +++ b/codex-rs/core/src/prompt_debug.rs @@ -49,6 +49,7 @@ pub async fn build_prompt_input( thread_store, state_db.clone(), installation_id, + /*attestation_provider*/ None, ); let thread = thread_manager.start_thread(config).await?; diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 828ecca4c2c4..f5faccf0be45 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -854,7 +854,7 @@ impl Session { live_thread: live_thread_init.as_ref().cloned(), thread_store: Arc::clone(&thread_store), attestation_provider: attestation_provider.clone(), - model_client: ModelClient::new_with_attestation_provider( + model_client: ModelClient::new( Some(Arc::clone(&auth_manager)), session_id, thread_id, diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 67fe370ac7ba..c083f6aad7d1 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -407,6 +407,7 @@ fn test_model_client_session() -> crate::client::ModelClientSession { /*enable_request_compression*/ false, /*include_timing_metrics*/ false, /*beta_features_header*/ None, + /*attestation_provider*/ None, ) .new_session() } @@ -3894,6 +3895,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { config.features.enabled(Feature::EnableRequestCompression), config.features.enabled(Feature::RuntimeMetrics), Session::build_model_client_beta_features_header(config.as_ref()), + /*attestation_provider*/ None, ), code_mode_service: crate::tools::code_mode::CodeModeService::new(), environment_manager: Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), @@ -5612,6 +5614,7 @@ where config.features.enabled(Feature::EnableRequestCompression), config.features.enabled(Feature::RuntimeMetrics), Session::build_model_client_beta_features_header(config.as_ref()), + /*attestation_provider*/ None, ), code_mode_service: crate::tools::code_mode::CodeModeService::new(), environment_manager: Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index b1bfdc47b0d0..2b8cfa8b39dd 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -295,30 +295,6 @@ impl ThreadManager { thread_store: Arc, state_db: Option, installation_id: String, - ) -> Self { - Self::new_with_attestation_provider( - config, - auth_manager, - session_source, - environment_manager, - analytics_events_client, - thread_store, - state_db, - installation_id, - /*attestation_provider*/ None, - ) - } - - #[allow(clippy::too_many_arguments)] - pub fn new_with_attestation_provider( - config: &Config, - auth_manager: Arc, - session_source: SessionSource, - environment_manager: Arc, - analytics_events_client: Option, - thread_store: Arc, - state_db: Option, - installation_id: String, attestation_provider: Option, ) -> Self { let codex_home = config.codex_home.clone(); diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index 0834c18e21b9..1cf0df95f2ea 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -401,6 +401,7 @@ async fn resume_and_fork_do_not_restore_thread_environments_from_rollout() { thread_store_from_config(&config, /*state_db*/ None), /*state_db*/ None, TEST_INSTALLATION_ID.to_string(), + /*attestation_provider*/ None, ); let selected_cwd = AbsolutePathBuf::try_from(config.cwd.as_path().join("selected")).expect("absolute path"); @@ -517,6 +518,7 @@ async fn explicit_installation_id_skips_codex_home_file() { thread_store, state_db.clone(), installation_id.clone(), + /*attestation_provider*/ None, ); let thread = manager @@ -554,6 +556,7 @@ async fn resume_active_thread_from_rollout_returns_running_thread() { thread_store_from_config(&config, /*state_db*/ None), /*state_db*/ None, TEST_INSTALLATION_ID.to_string(), + /*attestation_provider*/ None, ); let source = manager @@ -609,6 +612,7 @@ async fn resume_stopped_thread_from_rollout_spawns_new_thread() { thread_store_from_config(&config, /*state_db*/ None), /*state_db*/ None, TEST_INSTALLATION_ID.to_string(), + /*attestation_provider*/ None, ); let source = manager @@ -671,6 +675,7 @@ async fn resume_stopped_thread_from_rollout_preserves_thread_source() { thread_store, state_db.clone(), TEST_INSTALLATION_ID.to_string(), + /*attestation_provider*/ None, ); let source = manager @@ -860,6 +865,7 @@ async fn new_uses_active_provider_for_model_refresh() { thread_store_from_config(&config, /*state_db*/ None), /*state_db*/ None, TEST_INSTALLATION_ID.to_string(), + /*attestation_provider*/ None, ); let _ = manager.list_models(RefreshStrategy::Online).await; @@ -1074,6 +1080,7 @@ async fn interrupted_fork_snapshot_does_not_synthesize_turn_id_for_legacy_histor thread_store_from_config(&config, state_db.clone()), state_db.clone(), TEST_INSTALLATION_ID.to_string(), + /*attestation_provider*/ None, ); let source = manager @@ -1180,6 +1187,7 @@ async fn interrupted_fork_snapshot_preserves_explicit_turn_id() { thread_store_from_config(&config, state_db.clone()), state_db.clone(), TEST_INSTALLATION_ID.to_string(), + /*attestation_provider*/ None, ); let source = manager @@ -1275,6 +1283,7 @@ async fn interrupted_fork_snapshot_uses_persisted_mid_turn_history_without_live_ thread_store_from_config(&config, state_db.clone()), state_db.clone(), TEST_INSTALLATION_ID.to_string(), + /*attestation_provider*/ None, ); let source = manager @@ -1416,6 +1425,7 @@ async fn resumed_thread_keeps_paused_goal_paused() -> anyhow::Result<()> { thread_store_from_config(&config, state_db.clone()), state_db.clone(), TEST_INSTALLATION_ID.to_string(), + /*attestation_provider*/ None, ); let source = manager diff --git a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs index 43503be8c170..e7625fa44ffe 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs @@ -3167,6 +3167,7 @@ async fn tool_handlers_cascade_close_and_resume_and_keep_explicitly_closed_subtr thread_store_from_config(&config, state_db.clone()), state_db.clone(), "11111111-1111-4111-8111-111111111111".to_string(), + /*attestation_provider*/ None, ); let parent = manager diff --git a/codex-rs/core/tests/common/test_codex.rs b/codex-rs/core/tests/common/test_codex.rs index 7082d2608905..0a40076c6813 100644 --- a/codex-rs/core/tests/common/test_codex.rs +++ b/codex-rs/core/tests/common/test_codex.rs @@ -436,6 +436,7 @@ impl TestCodexBuilder { thread_store, state_db.clone(), installation_id, + /*attestation_provider*/ None, ); let thread_manager = Arc::new(thread_manager); let user_shell_override = self.user_shell_override.clone(); diff --git a/codex-rs/core/tests/responses_headers.rs b/codex-rs/core/tests/responses_headers.rs index af99790a1fec..6f0429e64499 100644 --- a/codex-rs/core/tests/responses_headers.rs +++ b/codex-rs/core/tests/responses_headers.rs @@ -109,6 +109,7 @@ async fn responses_stream_includes_subagent_header_on_review() { /*enable_request_compression*/ false, /*include_timing_metrics*/ false, /*beta_features_header*/ None, + /*attestation_provider*/ None, ); let mut client_session = client.new_session(); @@ -236,6 +237,7 @@ async fn responses_stream_includes_subagent_header_on_other() { /*enable_request_compression*/ false, /*include_timing_metrics*/ false, /*beta_features_header*/ None, + /*attestation_provider*/ None, ); let mut client_session = client.new_session(); @@ -352,6 +354,7 @@ async fn responses_respects_model_info_overrides_from_config() { /*enable_request_compression*/ false, /*include_timing_metrics*/ false, /*beta_features_header*/ None, + /*attestation_provider*/ None, ); let mut client_session = client.new_session(); diff --git a/codex-rs/core/tests/suite/client.rs b/codex-rs/core/tests/suite/client.rs index 432b57de9f8b..b273c6d5a276 100644 --- a/codex-rs/core/tests/suite/client.rs +++ b/codex-rs/core/tests/suite/client.rs @@ -900,6 +900,7 @@ async fn send_provider_auth_request(server: &MockServer, auth: ModelProviderAuth /*enable_request_compression*/ false, /*include_timing_metrics*/ false, /*beta_features_header*/ None, + /*attestation_provider*/ None, ); let mut client_session = client.new_session(); let mut prompt = Prompt::default(); @@ -1126,6 +1127,7 @@ async fn prefers_apikey_when_config_prefers_apikey_even_with_chatgpt_tokens() { thread_store_from_config(&config, /*state_db*/ None), /*state_db*/ None, installation_id, + /*attestation_provider*/ None, ); let NewThread { thread: codex, .. } = thread_manager .start_thread(config.clone()) @@ -2313,6 +2315,7 @@ async fn azure_responses_request_includes_store_and_reasoning_ids() { /*enable_request_compression*/ false, /*include_timing_metrics*/ false, /*beta_features_header*/ None, + /*attestation_provider*/ None, ); let mut client_session = client.new_session(); diff --git a/codex-rs/core/tests/suite/client_websockets.rs b/codex-rs/core/tests/suite/client_websockets.rs index 796b85ca5ed7..0c9bad38787c 100755 --- a/codex-rs/core/tests/suite/client_websockets.rs +++ b/codex-rs/core/tests/suite/client_websockets.rs @@ -1960,6 +1960,7 @@ async fn websocket_harness_with_provider_options( /*enable_request_compression*/ false, runtime_metrics_enabled, /*beta_features_header*/ None, + /*attestation_provider*/ None, ); WebsocketTestHarness { diff --git a/codex-rs/mcp-server/src/message_processor.rs b/codex-rs/mcp-server/src/message_processor.rs index d64fc43b1b81..f963747e8774 100644 --- a/codex-rs/mcp-server/src/message_processor.rs +++ b/codex-rs/mcp-server/src/message_processor.rs @@ -72,6 +72,7 @@ impl MessageProcessor { thread_store_from_config(config.as_ref(), state_db.clone()), state_db.clone(), installation_id, + /*attestation_provider*/ None, )); Self { outgoing, diff --git a/codex-rs/memories/write/src/runtime.rs b/codex-rs/memories/write/src/runtime.rs index 53a10934a594..2b39805f4e72 100644 --- a/codex-rs/memories/write/src/runtime.rs +++ b/codex-rs/memories/write/src/runtime.rs @@ -183,6 +183,7 @@ impl MemoryStartupContext { config.features.enabled(Feature::EnableRequestCompression), config.features.enabled(Feature::RuntimeMetrics), /*beta_features_header*/ None, + /*attestation_provider*/ None, ); let mut client_session = model_client.new_session(); diff --git a/codex-rs/thread-manager-sample/src/main.rs b/codex-rs/thread-manager-sample/src/main.rs index 6817f677e6b6..dc49771af819 100644 --- a/codex-rs/thread-manager-sample/src/main.rs +++ b/codex-rs/thread-manager-sample/src/main.rs @@ -126,6 +126,7 @@ async fn run_main(arg0_paths: Arg0DispatchPaths) -> anyhow::Result<()> { Arc::clone(&thread_store), state_db, installation_id, + /*attestation_provider*/ None, ); let NewThread { From d6e3cb76ecf0984614c1ae1213d14adb5b2e103e Mon Sep 17 00:00:00 2001 From: Jiaming Zhang Date: Wed, 6 May 2026 19:12:53 -0700 Subject: [PATCH 05/23] refactor(attestation): remove request purpose enum Co-authored-by: Codex --- codex-rs/core/src/client.rs | 50 ++++------------------ codex-rs/core/src/client_tests.rs | 70 ++++--------------------------- 2 files changed, 17 insertions(+), 103 deletions(-) diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index cb7502b1ba3d..9a0ec4e2ca6b 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -494,12 +494,8 @@ impl ModelClient { Some(self.state.session_id.to_string()), Some(self.state.thread_id.to_string()), )); - self.extend_attestation_header_for( - &mut extra_headers, - &client_setup.api_provider, - AttestationPurpose::Compaction, - ) - .await; + self.extend_attestation_header_for(&mut extra_headers, &client_setup.api_provider) + .await; let client = ApiCompactClient::new(transport, client_setup.api_provider, client_setup.api_auth) .with_telemetry(Some(request_telemetry)); @@ -521,12 +517,8 @@ impl ModelClient { // 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?; - self.extend_attestation_header_for( - &mut extra_headers, - &client_setup.api_provider, - AttestationPurpose::RealtimeWebrtcCallSetup, - ) - .await; + self.extend_attestation_header_for(&mut extra_headers, &client_setup.api_provider) + .await; let mut sideband_headers = extra_headers.clone(); sideband_headers.extend(sideband_websocket_auth_headers( client_setup.api_auth.as_ref(), @@ -668,9 +660,8 @@ impl ModelClient { async fn generate_attestation_header_for( &self, provider: &codex_api::Provider, - purpose: AttestationPurpose, ) -> Option { - if !should_send_attestation(provider, purpose) { + if !provider.uses_chatgpt_auth { return None; } self.generate_attestation_header().await @@ -911,7 +902,7 @@ impl ModelClient { } headers.extend(build_session_headers(Some(session_id), Some(thread_id))); headers.extend(self.build_responses_identity_headers()); - self.extend_attestation_header_for(&mut headers, provider, AttestationPurpose::Response) + self.extend_attestation_header_for(&mut headers, provider) .await; headers.insert( OPENAI_BETA_HEADER, @@ -982,11 +973,7 @@ impl ModelClientSession { ); headers.extend(self.client.build_responses_identity_headers()); self.client - .extend_attestation_header_for( - &mut headers, - provider, - AttestationPurpose::Response, - ) + .extend_attestation_header_for(&mut headers, provider) .await; headers }, @@ -1687,34 +1674,13 @@ fn build_responses_headers( headers } -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -enum AttestationPurpose { - Response, - Compaction, - RealtimeWebrtcCallSetup, -} - -fn should_send_attestation(provider: &codex_api::Provider, purpose: AttestationPurpose) -> bool { - provider.uses_chatgpt_auth - && matches!( - purpose, - AttestationPurpose::Response - | AttestationPurpose::Compaction - | AttestationPurpose::RealtimeWebrtcCallSetup - ) -} - impl ModelClient { async fn extend_attestation_header_for( &self, headers: &mut ApiHeaderMap, provider: &codex_api::Provider, - purpose: AttestationPurpose, ) { - if let Some(header_value) = self - .generate_attestation_header_for(provider, purpose) - .await - { + if let Some(header_value) = self.generate_attestation_header_for(provider).await { headers.insert(X_OAI_ATTESTATION_HEADER, header_value); } } diff --git a/codex-rs/core/src/client_tests.rs b/codex-rs/core/src/client_tests.rs index bfdc679b698f..bc278b449040 100644 --- a/codex-rs/core/src/client_tests.rs +++ b/codex-rs/core/src/client_tests.rs @@ -1,4 +1,3 @@ -use super::AttestationPurpose; use super::AuthRequestTelemetryContext; use super::ModelClient; use super::PendingUnauthorizedRetry; @@ -519,29 +518,6 @@ fn model_client_with_counting_attestation() -> (ModelClient, Arc) { (model_client, attestation_calls) } -#[test] -fn should_send_attestation_for_allowed_chatgpt_codex_purposes() { - let provider = api_provider("https://chatgpt.com/backend-api/codex/"); - - for purpose in [ - AttestationPurpose::Response, - AttestationPurpose::Compaction, - AttestationPurpose::RealtimeWebrtcCallSetup, - ] { - assert!(super::should_send_attestation(&provider, purpose)); - } -} - -#[test] -fn should_not_send_attestation_for_non_chatgpt_codex_provider() { - let provider = api_provider("https://api.openai.com/v1"); - - assert!(!super::should_send_attestation( - &provider, - AttestationPurpose::Response, - )); -} - #[tokio::test] async fn responses_generate_fresh_attestation_headers_for_chatgpt_codex() { let provider = api_provider("https://chatgpt.com/backend-api/codex/"); @@ -550,10 +526,10 @@ async fn responses_generate_fresh_attestation_headers_for_chatgpt_codex() { let mut second_headers = http::HeaderMap::new(); model_client - .extend_attestation_header_for(&mut first_headers, &provider, AttestationPurpose::Response) + .extend_attestation_header_for(&mut first_headers, &provider) .await; model_client - .extend_attestation_header_for(&mut second_headers, &provider, AttestationPurpose::Response) + .extend_attestation_header_for(&mut second_headers, &provider) .await; assert_eq!( @@ -599,18 +575,10 @@ async fn compact_generate_fresh_attestation_headers_for_chatgpt_codex() { let mut second_headers = http::HeaderMap::new(); model_client - .extend_attestation_header_for( - &mut first_headers, - &provider, - AttestationPurpose::Compaction, - ) + .extend_attestation_header_for(&mut first_headers, &provider) .await; model_client - .extend_attestation_header_for( - &mut second_headers, - &provider, - AttestationPurpose::Compaction, - ) + .extend_attestation_header_for(&mut second_headers, &provider) .await; assert_eq!( @@ -636,18 +604,10 @@ async fn realtime_setup_generate_fresh_attestation_headers_for_chatgpt_codex() { let mut second_headers = http::HeaderMap::new(); model_client - .extend_attestation_header_for( - &mut first_headers, - &provider, - AttestationPurpose::RealtimeWebrtcCallSetup, - ) + .extend_attestation_header_for(&mut first_headers, &provider) .await; model_client - .extend_attestation_header_for( - &mut second_headers, - &provider, - AttestationPurpose::RealtimeWebrtcCallSetup, - ) + .extend_attestation_header_for(&mut second_headers, &provider) .await; assert_eq!( @@ -672,27 +632,15 @@ async fn non_chatgpt_codex_endpoints_omit_attestation_generation() { let mut response_headers = http::HeaderMap::new(); model_client - .extend_attestation_header_for( - &mut response_headers, - &provider, - AttestationPurpose::Response, - ) + .extend_attestation_header_for(&mut response_headers, &provider) .await; let mut compaction_headers = http::HeaderMap::new(); model_client - .extend_attestation_header_for( - &mut compaction_headers, - &provider, - AttestationPurpose::Compaction, - ) + .extend_attestation_header_for(&mut compaction_headers, &provider) .await; let mut realtime_headers = http::HeaderMap::new(); model_client - .extend_attestation_header_for( - &mut realtime_headers, - &provider, - AttestationPurpose::RealtimeWebrtcCallSetup, - ) + .extend_attestation_header_for(&mut realtime_headers, &provider) .await; assert_eq!( From 92a0b57d3668962516890def093f7f2408827733 Mon Sep 17 00:00:00 2001 From: Jiaming Zhang Date: Wed, 6 May 2026 19:24:27 -0700 Subject: [PATCH 06/23] test(attestation): trim duplicate unit coverage Co-authored-by: Codex --- codex-rs/app-server/src/message_processor.rs | 15 +--- codex-rs/core/src/client_tests.rs | 87 -------------------- 2 files changed, 2 insertions(+), 100 deletions(-) diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index be59de451ea9..e4860238068e 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -162,25 +162,14 @@ fn app_server_attestation_provider( AttestationProvider::new(move || { let outgoing = outgoing.clone(); let attestation_connection_ids = attestation_connection_ids.clone(); - Box::pin(request_attestation_header_value( + Box::pin(request_attestation_header_value_with_timeout( outgoing, attestation_connection_ids, + ATTESTATION_GENERATE_TIMEOUT, )) }) } -async fn request_attestation_header_value( - outgoing: Arc, - attestation_connection_ids: Arc>>, -) -> Option { - request_attestation_header_value_with_timeout( - outgoing, - attestation_connection_ids, - ATTESTATION_GENERATE_TIMEOUT, - ) - .await -} - async fn request_attestation_header_value_with_timeout( outgoing: Arc, attestation_connection_ids: Arc>>, diff --git a/codex-rs/core/src/client_tests.rs b/codex-rs/core/src/client_tests.rs index bc278b449040..4bfa5b8a5a91 100644 --- a/codex-rs/core/src/client_tests.rs +++ b/codex-rs/core/src/client_tests.rs @@ -518,35 +518,6 @@ fn model_client_with_counting_attestation() -> (ModelClient, Arc) { (model_client, attestation_calls) } -#[tokio::test] -async fn responses_generate_fresh_attestation_headers_for_chatgpt_codex() { - let provider = api_provider("https://chatgpt.com/backend-api/codex/"); - let (model_client, attestation_calls) = model_client_with_counting_attestation(); - let mut first_headers = http::HeaderMap::new(); - let mut second_headers = http::HeaderMap::new(); - - model_client - .extend_attestation_header_for(&mut first_headers, &provider) - .await; - model_client - .extend_attestation_header_for(&mut second_headers, &provider) - .await; - - assert_eq!( - first_headers - .get(crate::attestation::X_OAI_ATTESTATION_HEADER) - .and_then(|value| value.to_str().ok()), - Some("v1.header-1"), - ); - assert_eq!( - second_headers - .get(crate::attestation::X_OAI_ATTESTATION_HEADER) - .and_then(|value| value.to_str().ok()), - Some("v1.header-2"), - ); - assert_eq!(attestation_calls.load(Ordering::Relaxed), 2); -} - #[tokio::test] async fn websocket_handshake_includes_attestation_for_chatgpt_codex_responses() { let provider = api_provider("https://chatgpt.com/backend-api/codex/"); @@ -567,64 +538,6 @@ async fn websocket_handshake_includes_attestation_for_chatgpt_codex_responses() assert_eq!(attestation_calls.load(Ordering::Relaxed), 1); } -#[tokio::test] -async fn compact_generate_fresh_attestation_headers_for_chatgpt_codex() { - let provider = api_provider("https://chatgpt.com/backend-api/codex/"); - let (model_client, attestation_calls) = model_client_with_counting_attestation(); - let mut first_headers = http::HeaderMap::new(); - let mut second_headers = http::HeaderMap::new(); - - model_client - .extend_attestation_header_for(&mut first_headers, &provider) - .await; - model_client - .extend_attestation_header_for(&mut second_headers, &provider) - .await; - - assert_eq!( - first_headers - .get(crate::attestation::X_OAI_ATTESTATION_HEADER) - .and_then(|value| value.to_str().ok()), - Some("v1.header-1"), - ); - assert_eq!( - second_headers - .get(crate::attestation::X_OAI_ATTESTATION_HEADER) - .and_then(|value| value.to_str().ok()), - Some("v1.header-2"), - ); - assert_eq!(attestation_calls.load(Ordering::Relaxed), 2); -} - -#[tokio::test] -async fn realtime_setup_generate_fresh_attestation_headers_for_chatgpt_codex() { - let provider = api_provider("https://chatgpt.com/backend-api/codex/"); - let (model_client, attestation_calls) = model_client_with_counting_attestation(); - let mut first_headers = http::HeaderMap::new(); - let mut second_headers = http::HeaderMap::new(); - - model_client - .extend_attestation_header_for(&mut first_headers, &provider) - .await; - model_client - .extend_attestation_header_for(&mut second_headers, &provider) - .await; - - assert_eq!( - first_headers - .get(crate::attestation::X_OAI_ATTESTATION_HEADER) - .and_then(|value| value.to_str().ok()), - Some("v1.header-1"), - ); - assert_eq!( - second_headers - .get(crate::attestation::X_OAI_ATTESTATION_HEADER) - .and_then(|value| value.to_str().ok()), - Some("v1.header-2"), - ); - assert_eq!(attestation_calls.load(Ordering::Relaxed), 2); -} - #[tokio::test] async fn non_chatgpt_codex_endpoints_omit_attestation_generation() { let provider = api_provider("https://api.openai.com/v1"); From bb55e050d75c7b52af5b0977bb660cfa0ee14f70 Mon Sep 17 00:00:00 2001 From: Jiaming Zhang Date: Wed, 6 May 2026 21:43:17 -0700 Subject: [PATCH 07/23] refactor(attestation): move policy behind provider Co-authored-by: Codex --- codex-rs/app-server/src/message_processor.rs | 47 ++++++++++++++++---- codex-rs/core/src/attestation.rs | 43 ++++++------------ codex-rs/core/src/client.rs | 25 +++++------ codex-rs/core/src/client_tests.rs | 33 +++++++++++--- codex-rs/core/src/lib.rs | 2 + codex-rs/core/src/session/mod.rs | 2 +- codex-rs/core/src/session/session.rs | 2 +- codex-rs/core/src/state/service.rs | 2 +- codex-rs/core/src/thread_manager.rs | 4 +- 9 files changed, 95 insertions(+), 65 deletions(-) diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index e4860238068e..f9506611802d 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -60,7 +60,9 @@ use codex_app_server_protocol::ServerRequestPayload; use codex_app_server_protocol::experimental_required_message; use codex_arg0::Arg0DispatchPaths; use codex_chatgpt::workspace_settings; +use codex_core::AttestationContext; use codex_core::AttestationProvider; +use codex_core::GenerateAttestationFuture; use codex_core::ThreadManager; use codex_core::config::Config; use codex_core::thread_store_from_config; @@ -158,18 +160,45 @@ impl ExternalAuth for ExternalAuthRefreshBridge { fn app_server_attestation_provider( outgoing: Arc, attestation_connection_ids: Arc>>, -) -> AttestationProvider { - AttestationProvider::new(move || { - let outgoing = outgoing.clone(); - let attestation_connection_ids = attestation_connection_ids.clone(); - Box::pin(request_attestation_header_value_with_timeout( - outgoing, - attestation_connection_ids, - ATTESTATION_GENERATE_TIMEOUT, - )) +) -> Arc { + Arc::new(AppServerAttestationProvider { + outgoing, + attestation_connection_ids, }) } +struct AppServerAttestationProvider { + outgoing: Arc, + attestation_connection_ids: Arc>>, +} + +impl std::fmt::Debug for AppServerAttestationProvider { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("AppServerAttestationProvider") + .finish() + } +} + +impl AttestationProvider for AppServerAttestationProvider { + fn generate_header_value(&self, context: AttestationContext) -> GenerateAttestationFuture<'_> { + let outgoing = self.outgoing.clone(); + let attestation_connection_ids = self.attestation_connection_ids.clone(); + Box::pin(async move { + if !context.uses_chatgpt_auth { + return None; + } + + request_attestation_header_value_with_timeout( + outgoing, + attestation_connection_ids, + ATTESTATION_GENERATE_TIMEOUT, + ) + .await + }) + } +} + async fn request_attestation_header_value_with_timeout( outgoing: Arc, attestation_connection_ids: Arc>>, diff --git a/codex-rs/core/src/attestation.rs b/codex-rs/core/src/attestation.rs index aa637304be26..8e5bdcf5388a 100644 --- a/codex-rs/core/src/attestation.rs +++ b/codex-rs/core/src/attestation.rs @@ -1,39 +1,22 @@ -use std::fmt; use std::future::Future; use std::pin::Pin; -use std::sync::Arc; - -use http::HeaderValue; pub(crate) const X_OAI_ATTESTATION_HEADER: &str = "x-oai-attestation"; -type GenerateAttestationFuture = Pin> + Send>>; -type GenerateAttestationCallback = dyn Fn() -> GenerateAttestationFuture + Send + Sync + 'static; - -/// Session-scoped source for just-in-time attestation header values. -/// -/// Host integrations provide the opaque string expected by the upstream -/// `x-oai-attestation` header. Core validates only that it is legal as an HTTP -/// header value before forwarding it. -#[derive(Clone)] -pub struct AttestationProvider { - generate: Arc, -} +pub type GenerateAttestationFuture<'a> = Pin> + Send + 'a>>; -impl fmt::Debug for AttestationProvider { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - formatter.debug_struct("AttestationProvider").finish() - } +/// Request context that host integrations can use when deciding whether to +/// generate an attestation header value. +#[derive(Clone, Copy, Debug)] +pub struct AttestationContext { + pub uses_chatgpt_auth: bool, } -impl AttestationProvider { - pub fn new(generate: impl Fn() -> GenerateAttestationFuture + Send + Sync + 'static) -> Self { - Self { - generate: Arc::new(generate), - } - } - - pub(crate) async fn generate_header(&self) -> Option { - HeaderValue::from_str(&(self.generate)().await?).ok() - } +/// Host integration boundary for just-in-time attestation header values. +/// +/// Implementations own the policy for when attestation should be attempted and +/// return the opaque string expected by the upstream `x-oai-attestation` +/// header. Core only forwards valid HTTP header values returned by the host. +pub trait AttestationProvider: std::fmt::Debug + Send + Sync { + fn generate_header_value(&self, context: AttestationContext) -> GenerateAttestationFuture<'_>; } diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 9a0ec4e2ca6b..21b69064a713 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -105,6 +105,7 @@ use tracing::instrument; use tracing::trace; use tracing::warn; +use crate::attestation::AttestationContext; use crate::attestation::AttestationProvider; use crate::attestation::X_OAI_ATTESTATION_HEADER; use crate::client_common::Prompt; @@ -172,7 +173,7 @@ struct ModelClientState { enable_request_compression: bool, include_timing_metrics: bool, beta_features_header: Option, - attestation_provider: Option, + attestation_provider: Option>, disable_websockets: AtomicBool, cached_websocket_session: StdMutex, } @@ -317,7 +318,7 @@ impl ModelClient { enable_request_compression: bool, include_timing_metrics: bool, beta_features_header: Option, - attestation_provider: Option, + attestation_provider: Option>, ) -> Self { let model_provider = create_model_provider(provider_info, auth_manager); let codex_api_key_env_enabled = model_provider @@ -649,22 +650,18 @@ impl ModelClient { client_metadata } - async fn generate_attestation_header(&self) -> Option { - self.state - .attestation_provider - .as_ref()? - .generate_header() - .await - } - async fn generate_attestation_header_for( &self, provider: &codex_api::Provider, ) -> Option { - if !provider.uses_chatgpt_auth { - return None; - } - self.generate_attestation_header().await + self.state + .attestation_provider + .as_ref()? + .generate_header_value(AttestationContext { + uses_chatgpt_auth: provider.uses_chatgpt_auth, + }) + .await + .and_then(|value| HeaderValue::from_str(&value).ok()) } /// Builds request telemetry for unary API calls (e.g., Compact endpoint). diff --git a/codex-rs/core/src/client_tests.rs b/codex-rs/core/src/client_tests.rs index 4bfa5b8a5a91..38058e2baa43 100644 --- a/codex-rs/core/src/client_tests.rs +++ b/codex-rs/core/src/client_tests.rs @@ -7,7 +7,9 @@ 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::AttestationContext; use crate::AttestationProvider; +use crate::GenerateAttestationFuture; use codex_api::ApiError; use codex_api::ResponseEvent; use codex_app_server_protocol::AuthMode; @@ -494,8 +496,29 @@ fn auth_request_telemetry_context_tracks_attached_auth_and_retry_phase() { } fn model_client_with_counting_attestation() -> (ModelClient, Arc) { + #[derive(Debug)] + struct CountingAttestationProvider { + calls: Arc, + } + + impl AttestationProvider for CountingAttestationProvider { + fn generate_header_value( + &self, + context: AttestationContext, + ) -> GenerateAttestationFuture<'_> { + let calls = self.calls.clone(); + Box::pin(async move { + if !context.uses_chatgpt_auth { + return None; + } + + let call = calls.fetch_add(1, Ordering::Relaxed) + 1; + Some(format!("v1.header-{call}")) + }) + } + } + let attestation_calls = Arc::new(AtomicUsize::new(0)); - let calls = attestation_calls.clone(); let model_client = ModelClient::new( /*auth_manager*/ None, SessionId::new(), @@ -507,12 +530,8 @@ fn model_client_with_counting_attestation() -> (ModelClient, Arc) { /*enable_request_compression*/ false, /*include_timing_metrics*/ false, /*beta_features_header*/ None, - Some(AttestationProvider::new(move || { - let calls = calls.clone(); - Box::pin(async move { - let call = calls.fetch_add(1, Ordering::Relaxed) + 1; - Some(format!("v1.header-{call}")) - }) + Some(Arc::new(CountingAttestationProvider { + calls: attestation_calls.clone(), })), ); (model_client, attestation_calls) diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index c178b5d4e8bd..57ac9b4a5944 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -178,7 +178,9 @@ mod tasks; mod user_shell_command; pub mod util; +pub use attestation::AttestationContext; pub use attestation::AttestationProvider; +pub use attestation::GenerateAttestationFuture; pub use client::ModelClient; pub use client::ModelClientSession; pub use client::X_CODEX_INSTALLATION_ID_HEADER; diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 07ef550f5179..d3a00b3108eb 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -413,7 +413,7 @@ pub(crate) struct CodexSpawnArgs { pub(crate) environment_selections: ResolvedTurnEnvironments, pub(crate) analytics_events_client: Option, pub(crate) thread_store: Arc, - pub(crate) attestation_provider: Option, + pub(crate) attestation_provider: Option>, } pub(crate) const INITIAL_SUBMIT_ID: &str = ""; diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index f5faccf0be45..1a790314d589 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -370,7 +370,7 @@ impl Session { analytics_events_client: Option, thread_store: Arc, parent_rollout_thread_trace: ThreadTraceContext, - attestation_provider: Option, + attestation_provider: Option>, ) -> anyhow::Result> { debug!( "Configuring session: model={}; provider={:?}", diff --git a/codex-rs/core/src/state/service.rs b/codex-rs/core/src/state/service.rs index 4506c0054cf7..0dba931296c1 100644 --- a/codex-rs/core/src/state/service.rs +++ b/codex-rs/core/src/state/service.rs @@ -67,7 +67,7 @@ pub(crate) struct SessionServices { pub(crate) state_db: Option, pub(crate) live_thread: Option, pub(crate) thread_store: Arc, - pub(crate) attestation_provider: Option, + pub(crate) attestation_provider: Option>, /// Session-scoped model client shared across turns. pub(crate) model_client: ModelClient, pub(crate) code_mode_service: CodeModeService, diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 2b8cfa8b39dd..d2c433d44e6d 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -250,7 +250,7 @@ pub(crate) struct ThreadManagerState { mcp_manager: Arc, skills_watcher: Arc, thread_store: Arc, - attestation_provider: Option, + attestation_provider: Option>, session_source: SessionSource, installation_id: String, analytics_events_client: Option, @@ -295,7 +295,7 @@ impl ThreadManager { thread_store: Arc, state_db: Option, installation_id: String, - attestation_provider: Option, + attestation_provider: Option>, ) -> Self { let codex_home = config.codex_home.clone(); let restriction_product = session_source.restriction_product(); From 4e4e6b6114f3f1f2f0c98e402cede80ba64509e8 Mon Sep 17 00:00:00 2001 From: Jiaming Zhang Date: Wed, 6 May 2026 22:00:04 -0700 Subject: [PATCH 08/23] refactor(attestation): return header values from provider Co-authored-by: Codex --- codex-rs/app-server/src/message_processor.rs | 4 +++- codex-rs/core/src/attestation.rs | 10 ++++++---- codex-rs/core/src/client.rs | 3 +-- codex-rs/core/src/client_tests.rs | 7 ++----- 4 files changed, 12 insertions(+), 12 deletions(-) diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index f9506611802d..f45f9d554ec8 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -38,6 +38,7 @@ use crate::thread_state::ThreadStateManager; use crate::transport::AppServerTransport; use crate::transport::RemoteControlHandle; use async_trait::async_trait; +use axum::http::HeaderValue; use codex_analytics::AnalyticsEventsClient; use codex_analytics::AppServerRpcTransport; use codex_app_server_protocol::AttestationGenerateParams; @@ -181,7 +182,7 @@ impl std::fmt::Debug for AppServerAttestationProvider { } impl AttestationProvider for AppServerAttestationProvider { - fn generate_header_value(&self, context: AttestationContext) -> GenerateAttestationFuture<'_> { + fn header_for_request(&self, context: AttestationContext) -> GenerateAttestationFuture<'_> { let outgoing = self.outgoing.clone(); let attestation_connection_ids = self.attestation_connection_ids.clone(); Box::pin(async move { @@ -195,6 +196,7 @@ impl AttestationProvider for AppServerAttestationProvider { ATTESTATION_GENERATE_TIMEOUT, ) .await + .and_then(|value| HeaderValue::from_bytes(value.as_bytes()).ok()) }) } } diff --git a/codex-rs/core/src/attestation.rs b/codex-rs/core/src/attestation.rs index 8e5bdcf5388a..01d4e72e4ea8 100644 --- a/codex-rs/core/src/attestation.rs +++ b/codex-rs/core/src/attestation.rs @@ -1,9 +1,12 @@ use std::future::Future; use std::pin::Pin; +use http::HeaderValue; + pub(crate) const X_OAI_ATTESTATION_HEADER: &str = "x-oai-attestation"; -pub type GenerateAttestationFuture<'a> = Pin> + Send + 'a>>; +pub type GenerateAttestationFuture<'a> = + Pin> + Send + 'a>>; /// Request context that host integrations can use when deciding whether to /// generate an attestation header value. @@ -15,8 +18,7 @@ pub struct AttestationContext { /// Host integration boundary for just-in-time attestation header values. /// /// Implementations own the policy for when attestation should be attempted and -/// return the opaque string expected by the upstream `x-oai-attestation` -/// header. Core only forwards valid HTTP header values returned by the host. +/// return the upstream `x-oai-attestation` header value when one should be sent. pub trait AttestationProvider: std::fmt::Debug + Send + Sync { - fn generate_header_value(&self, context: AttestationContext) -> GenerateAttestationFuture<'_>; + fn header_for_request(&self, context: AttestationContext) -> GenerateAttestationFuture<'_>; } diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 21b69064a713..ab850770ca72 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -657,11 +657,10 @@ impl ModelClient { self.state .attestation_provider .as_ref()? - .generate_header_value(AttestationContext { + .header_for_request(AttestationContext { uses_chatgpt_auth: provider.uses_chatgpt_auth, }) .await - .and_then(|value| HeaderValue::from_str(&value).ok()) } /// Builds request telemetry for unary API calls (e.g., Compact endpoint). diff --git a/codex-rs/core/src/client_tests.rs b/codex-rs/core/src/client_tests.rs index 38058e2baa43..55b0c227c660 100644 --- a/codex-rs/core/src/client_tests.rs +++ b/codex-rs/core/src/client_tests.rs @@ -502,10 +502,7 @@ fn model_client_with_counting_attestation() -> (ModelClient, Arc) { } impl AttestationProvider for CountingAttestationProvider { - fn generate_header_value( - &self, - context: AttestationContext, - ) -> GenerateAttestationFuture<'_> { + fn header_for_request(&self, context: AttestationContext) -> GenerateAttestationFuture<'_> { let calls = self.calls.clone(); Box::pin(async move { if !context.uses_chatgpt_auth { @@ -513,7 +510,7 @@ fn model_client_with_counting_attestation() -> (ModelClient, Arc) { } let call = calls.fetch_add(1, Ordering::Relaxed) + 1; - Some(format!("v1.header-{call}")) + Some(http::HeaderValue::from_bytes(format!("v1.header-{call}").as_bytes()).unwrap()) }) } } From 441a97bc5d7837b7d6eea2c7c5fe39f96fb6e4ad Mon Sep 17 00:00:00 2001 From: Jiaming Zhang Date: Wed, 6 May 2026 22:35:24 -0700 Subject: [PATCH 09/23] refactor(app-server): track attestation on live connections Co-authored-by: Codex --- codex-rs/app-server/src/lib.rs | 9 ++- codex-rs/app-server/src/message_processor.rs | 59 +++++++++---------- codex-rs/app-server/src/request_processors.rs | 1 + .../request_processors/thread_processor.rs | 8 ++- .../thread_processor_tests.rs | 57 ++++++++++++++++-- codex-rs/app-server/src/thread_state.rs | 31 ++++++++-- 6 files changed, 121 insertions(+), 44 deletions(-) diff --git a/codex-rs/app-server/src/lib.rs b/codex-rs/app-server/src/lib.rs index 08aab99f6549..9bb7937e19db 100644 --- a/codex-rs/app-server/src/lib.rs +++ b/codex-rs/app-server/src/lib.rs @@ -933,7 +933,14 @@ pub async fn run_main_with_transport_options( ), ) .await; - processor.connection_initialized(connection_id).await; + processor + .connection_initialized( + connection_id, + connection_state + .session + .request_attestation(), + ) + .await; connection_state .outbound_initialized .store(true, std::sync::atomic::Ordering::Release); diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index f45f9d554ec8..9bf24c4f163f 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -34,6 +34,7 @@ use crate::request_processors::WindowsSandboxRequestProcessor; use crate::request_serialization::QueuedInitializedRequest; use crate::request_serialization::RequestSerializationQueueKey; use crate::request_serialization::RequestSerializationQueues; +use crate::thread_state::ConnectionCapabilities; use crate::thread_state::ThreadStateManager; use crate::transport::AppServerTransport; use crate::transport::RemoteControlHandle; @@ -160,17 +161,17 @@ impl ExternalAuth for ExternalAuthRefreshBridge { fn app_server_attestation_provider( outgoing: Arc, - attestation_connection_ids: Arc>>, + thread_state_manager: ThreadStateManager, ) -> Arc { Arc::new(AppServerAttestationProvider { outgoing, - attestation_connection_ids, + thread_state_manager, }) } struct AppServerAttestationProvider { outgoing: Arc, - attestation_connection_ids: Arc>>, + thread_state_manager: ThreadStateManager, } impl std::fmt::Debug for AppServerAttestationProvider { @@ -184,7 +185,7 @@ impl std::fmt::Debug for AppServerAttestationProvider { impl AttestationProvider for AppServerAttestationProvider { fn header_for_request(&self, context: AttestationContext) -> GenerateAttestationFuture<'_> { let outgoing = self.outgoing.clone(); - let attestation_connection_ids = self.attestation_connection_ids.clone(); + let thread_state_manager = self.thread_state_manager.clone(); Box::pin(async move { if !context.uses_chatgpt_auth { return None; @@ -192,7 +193,7 @@ impl AttestationProvider for AppServerAttestationProvider { request_attestation_header_value_with_timeout( outgoing, - attestation_connection_ids, + thread_state_manager, ATTESTATION_GENERATE_TIMEOUT, ) .await @@ -203,15 +204,12 @@ impl AttestationProvider for AppServerAttestationProvider { async fn request_attestation_header_value_with_timeout( outgoing: Arc, - attestation_connection_ids: Arc>>, + thread_state_manager: ThreadStateManager, timeout_duration: Duration, ) -> Option { - let connection_id = attestation_connection_ids - .lock() - .await - .iter() - .min_by_key(|connection_id| connection_id.0) - .copied()?; + let connection_id = thread_state_manager + .first_attestation_capable_connection() + .await?; let connection_ids = [connection_id]; let (request_id, rx) = outgoing @@ -277,7 +275,6 @@ pub(crate) struct MessageProcessor { turn_processor: TurnRequestProcessor, windows_sandbox_processor: WindowsSandboxRequestProcessor, request_serialization_queues: RequestSerializationQueues, - attestation_connection_ids: Arc>>, } #[derive(Debug)] @@ -393,7 +390,7 @@ impl MessageProcessor { auth_manager.set_external_auth(Arc::new(ExternalAuthRefreshBridge { outgoing: outgoing.clone(), })); - let attestation_connection_ids = Arc::new(Mutex::new(HashSet::new())); + let thread_state_manager = ThreadStateManager::new(); // The thread store is intentionally process-scoped. Config reloads can // affect per-thread behavior, but they must not move newly started, // resumed, or forked threads to a different persistence backend/root. @@ -409,7 +406,7 @@ impl MessageProcessor { installation_id, Some(app_server_attestation_provider( outgoing.clone(), - attestation_connection_ids.clone(), + thread_state_manager.clone(), )), )); thread_manager @@ -417,7 +414,6 @@ impl MessageProcessor { .set_analytics_events_client(analytics_events_client.clone()); let pending_thread_unloads = Arc::new(Mutex::new(HashSet::new())); - let thread_state_manager = ThreadStateManager::new(); let thread_watch_manager = crate::thread_status::ThreadWatchManager::new_with_outgoing(outgoing.clone()); let thread_list_state_permit = Arc::new(Semaphore::new(/*permits*/ 1)); @@ -585,7 +581,6 @@ impl MessageProcessor { turn_processor, windows_sandbox_processor, request_serialization_queues: RequestSerializationQueues::default(), - attestation_connection_ids, } } @@ -739,9 +734,18 @@ impl MessageProcessor { .await; } - pub(crate) async fn connection_initialized(&self, connection_id: ConnectionId) { + pub(crate) async fn connection_initialized( + &self, + connection_id: ConnectionId, + request_attestation: bool, + ) { self.thread_processor - .connection_initialized(connection_id) + .connection_initialized( + connection_id, + ConnectionCapabilities { + request_attestation, + }, + ) .await; } @@ -783,10 +787,6 @@ impl MessageProcessor { session_state: &ConnectionSessionState, ) { session_state.rpc_gate.shutdown().await; - self.attestation_connection_ids - .lock() - .await - .remove(&connection_id); self.outgoing.connection_closed(connection_id).await; self.fs_processor.connection_closed(connection_id).await; self.command_exec_processor @@ -841,15 +841,14 @@ impl MessageProcessor { .await?; if connection_initialized { self.thread_processor - .connection_initialized(connection_id) + .connection_initialized( + connection_id, + ConnectionCapabilities { + request_attestation: session.request_attestation(), + }, + ) .await; } - if session.request_attestation() { - self.attestation_connection_ids - .lock() - .await - .insert(connection_id); - } return Ok(()); } diff --git a/codex-rs/app-server/src/request_processors.rs b/codex-rs/app-server/src/request_processors.rs index 795b5d4c3235..da3bdefc6674 100644 --- a/codex-rs/app-server/src/request_processors.rs +++ b/codex-rs/app-server/src/request_processors.rs @@ -471,6 +471,7 @@ use crate::error_code::internal_error; use crate::error_code::invalid_request; use crate::filters::compute_source_filters; use crate::filters::source_kind_matches; +use crate::thread_state::ConnectionCapabilities; use crate::thread_state::ThreadListenerCommand; use crate::thread_state::ThreadState; use crate::thread_state::ThreadStateManager; diff --git a/codex-rs/app-server/src/request_processors/thread_processor.rs b/codex-rs/app-server/src/request_processors/thread_processor.rs index d83f5e631f1c..554b1480826b 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor.rs @@ -2183,9 +2183,13 @@ impl ThreadRequestProcessor { self.thread_manager.subscribe_thread_created() } - pub(crate) async fn connection_initialized(&self, connection_id: ConnectionId) { + pub(crate) async fn connection_initialized( + &self, + connection_id: ConnectionId, + capabilities: ConnectionCapabilities, + ) { self.thread_state_manager - .connection_initialized(connection_id) + .connection_initialized(connection_id, capabilities) .await; } diff --git a/codex-rs/app-server/src/request_processors/thread_processor_tests.rs b/codex-rs/app-server/src/request_processors/thread_processor_tests.rs index 5642dbbe81bf..3cfde4ff4324 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor_tests.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor_tests.rs @@ -1115,7 +1115,9 @@ mod thread_processor_behavior_tests { let connection = ConnectionId(1); let (cancel_tx, cancel_rx) = oneshot::channel(); - manager.connection_initialized(connection).await; + manager + .connection_initialized(connection, ConnectionCapabilities::default()) + .await; manager .try_ensure_connection_subscribed( thread_id, connection, /*experimental_raw_events*/ false, @@ -1158,8 +1160,12 @@ mod thread_processor_behavior_tests { let connection_b = ConnectionId(2); let (cancel_tx, mut cancel_rx) = oneshot::channel(); - manager.connection_initialized(connection_a).await; - manager.connection_initialized(connection_b).await; + manager + .connection_initialized(connection_a, ConnectionCapabilities::default()) + .await; + manager + .connection_initialized(connection_b, ConnectionCapabilities::default()) + .await; manager .try_ensure_connection_subscribed( thread_id, @@ -1203,8 +1209,12 @@ mod thread_processor_behavior_tests { let connection_a = ConnectionId(1); let connection_b = ConnectionId(2); - manager.connection_initialized(connection_a).await; - manager.connection_initialized(connection_b).await; + manager + .connection_initialized(connection_a, ConnectionCapabilities::default()) + .await; + manager + .connection_initialized(connection_b, ConnectionCapabilities::default()) + .await; manager .try_ensure_connection_subscribed( thread_id, @@ -1249,7 +1259,9 @@ mod thread_processor_behavior_tests { let thread_id = ThreadId::from_string("ad7f0408-99b8-4f6e-a46f-bd0eec433370")?; let connection = ConnectionId(1); - manager.connection_initialized(connection).await; + manager + .connection_initialized(connection, ConnectionCapabilities::default()) + .await; let threads_to_unload = manager.remove_connection(connection).await; assert_eq!(threads_to_unload, Vec::::new()); @@ -1264,4 +1276,37 @@ mod thread_processor_behavior_tests { assert!(!manager.has_subscribers(thread_id).await); Ok(()) } + + #[tokio::test] + async fn first_attestation_capable_connection_returns_lowest_live_connection_id() { + let manager = ThreadStateManager::new(); + let unsupported_connection = ConnectionId(1); + let later_supported_connection = ConnectionId(3); + let earlier_supported_connection = ConnectionId(2); + + manager + .connection_initialized(unsupported_connection, ConnectionCapabilities::default()) + .await; + manager + .connection_initialized( + later_supported_connection, + ConnectionCapabilities { + request_attestation: true, + }, + ) + .await; + manager + .connection_initialized( + earlier_supported_connection, + ConnectionCapabilities { + request_attestation: true, + }, + ) + .await; + + assert_eq!( + manager.first_attestation_capable_connection().await, + Some(earlier_supported_connection) + ); + } } diff --git a/codex-rs/app-server/src/thread_state.rs b/codex-rs/app-server/src/thread_state.rs index dddbcf483b09..8552d21f710e 100644 --- a/codex-rs/app-server/src/thread_state.rs +++ b/codex-rs/app-server/src/thread_state.rs @@ -199,11 +199,16 @@ impl ThreadEntry { #[derive(Default)] struct ThreadStateManagerInner { - live_connections: HashSet, + live_connections: HashMap, threads: HashMap, thread_ids_by_connection: HashMap>, } +#[derive(Clone, Copy, Default)] +pub(crate) struct ConnectionCapabilities { + pub(crate) request_attestation: bool, +} + #[derive(Clone, Default)] pub(crate) struct ThreadStateManager { state: Arc>, @@ -214,12 +219,28 @@ impl ThreadStateManager { Self::default() } - pub(crate) async fn connection_initialized(&self, connection_id: ConnectionId) { + pub(crate) async fn connection_initialized( + &self, + connection_id: ConnectionId, + capabilities: ConnectionCapabilities, + ) { + self.state + .lock() + .await + .live_connections + .insert(connection_id, capabilities); + } + + pub(crate) async fn first_attestation_capable_connection(&self) -> Option { self.state .lock() .await .live_connections - .insert(connection_id); + .iter() + .filter_map(|(connection_id, capabilities)| { + capabilities.request_attestation.then_some(*connection_id) + }) + .min_by_key(|connection_id| connection_id.0) } pub(crate) async fn subscribed_connection_ids(&self, thread_id: ThreadId) -> Vec { @@ -338,7 +359,7 @@ impl ThreadStateManager { ) -> Option>> { let thread_state = { let mut state = self.state.lock().await; - if !state.live_connections.contains(&connection_id) { + if !state.live_connections.contains_key(&connection_id) { return None; } state @@ -366,7 +387,7 @@ impl ThreadStateManager { connection_id: ConnectionId, ) -> bool { let mut state = self.state.lock().await; - if !state.live_connections.contains(&connection_id) { + if !state.live_connections.contains_key(&connection_id) { return false; } state From 5fdd555edcae0eb15f3338cefa94400a7f18fc26 Mon Sep 17 00:00:00 2001 From: Jiaming Zhang Date: Wed, 6 May 2026 23:35:08 -0700 Subject: [PATCH 10/23] codex: fix CI failure on PR #20619 --- codex-rs/analytics/src/analytics_client_tests.rs | 1 + codex-rs/codex-api/src/endpoint/memories.rs | 1 + codex-rs/codex-api/src/endpoint/models.rs | 1 + codex-rs/codex-api/src/endpoint/realtime_call.rs | 1 + .../codex-api/src/endpoint/realtime_websocket/methods.rs | 5 +++++ codex-rs/codex-api/tests/clients.rs | 1 + codex-rs/codex-api/tests/models_integration.rs | 1 + codex-rs/codex-api/tests/realtime_websocket_e2e.rs | 1 + codex-rs/codex-api/tests/sse_end_to_end.rs | 1 + 9 files changed, 13 insertions(+) diff --git a/codex-rs/analytics/src/analytics_client_tests.rs b/codex-rs/analytics/src/analytics_client_tests.rs index 9bcc30bfaef0..6b3808e707d0 100644 --- a/codex-rs/analytics/src/analytics_client_tests.rs +++ b/codex-rs/analytics/src/analytics_client_tests.rs @@ -636,6 +636,7 @@ fn sample_initialize_fact(connection_id: u64) -> AnalyticsFact { }, capabilities: Some(InitializeCapabilities { experimental_api: false, + request_attestation: false, opt_out_notification_methods: None, }), }, diff --git a/codex-rs/codex-api/src/endpoint/memories.rs b/codex-rs/codex-api/src/endpoint/memories.rs index a6c25641f25d..e9060ccb9e9a 100644 --- a/codex-rs/codex-api/src/endpoint/memories.rs +++ b/codex-rs/codex-api/src/endpoint/memories.rs @@ -142,6 +142,7 @@ mod tests { Provider { name: "test".to_string(), base_url: base_url.to_string(), + uses_chatgpt_auth: false, query_params: None, headers: HeaderMap::new(), retry: RetryConfig { diff --git a/codex-rs/codex-api/src/endpoint/models.rs b/codex-rs/codex-api/src/endpoint/models.rs index ec9ee7aac6d3..c491f610b801 100644 --- a/codex-rs/codex-api/src/endpoint/models.rs +++ b/codex-rs/codex-api/src/endpoint/models.rs @@ -140,6 +140,7 @@ mod tests { Provider { name: "test".to_string(), base_url: base_url.to_string(), + uses_chatgpt_auth: false, query_params: None, headers: HeaderMap::new(), retry: RetryConfig { diff --git a/codex-rs/codex-api/src/endpoint/realtime_call.rs b/codex-rs/codex-api/src/endpoint/realtime_call.rs index b0342c53498d..ba5b88782481 100644 --- a/codex-rs/codex-api/src/endpoint/realtime_call.rs +++ b/codex-rs/codex-api/src/endpoint/realtime_call.rs @@ -297,6 +297,7 @@ mod tests { Provider { name: "test".to_string(), base_url: base_url.to_string(), + uses_chatgpt_auth: false, query_params: None, headers: HeaderMap::new(), retry: RetryConfig { diff --git a/codex-rs/codex-api/src/endpoint/realtime_websocket/methods.rs b/codex-rs/codex-api/src/endpoint/realtime_websocket/methods.rs index 9fcca1c3e318..41f9b92fca2c 100644 --- a/codex-rs/codex-api/src/endpoint/realtime_websocket/methods.rs +++ b/codex-rs/codex-api/src/endpoint/realtime_websocket/methods.rs @@ -1661,6 +1661,7 @@ mod tests { let provider = Provider { name: "test".to_string(), base_url: format!("http://{addr}"), + uses_chatgpt_auth: false, query_params: Some(HashMap::new()), headers: HeaderMap::new(), retry: crate::provider::RetryConfig { @@ -1955,6 +1956,7 @@ mod tests { let provider = Provider { name: "test".to_string(), base_url: format!("http://{addr}"), + uses_chatgpt_auth: false, query_params: Some(HashMap::new()), headers: HeaderMap::new(), retry: crate::provider::RetryConfig { @@ -2070,6 +2072,7 @@ mod tests { let provider = Provider { name: "test".to_string(), base_url: format!("http://{addr}"), + uses_chatgpt_auth: false, query_params: Some(HashMap::new()), headers: HeaderMap::new(), retry: crate::provider::RetryConfig { @@ -2174,6 +2177,7 @@ mod tests { let provider = Provider { name: "test".to_string(), base_url: format!("http://{addr}"), + uses_chatgpt_auth: false, query_params: Some(HashMap::new()), headers: HeaderMap::new(), retry: crate::provider::RetryConfig { @@ -2264,6 +2268,7 @@ mod tests { let provider = Provider { name: "test".to_string(), base_url: format!("http://{addr}"), + uses_chatgpt_auth: false, query_params: Some(HashMap::new()), headers: HeaderMap::new(), retry: crate::provider::RetryConfig { diff --git a/codex-rs/codex-api/tests/clients.rs b/codex-rs/codex-api/tests/clients.rs index a2a29ba16d37..6cdc9a09b9b8 100644 --- a/codex-rs/codex-api/tests/clients.rs +++ b/codex-rs/codex-api/tests/clients.rs @@ -127,6 +127,7 @@ fn provider(name: &str) -> Provider { Provider { name: name.to_string(), base_url: "https://example.com/v1".to_string(), + uses_chatgpt_auth: false, query_params: None, headers: HeaderMap::new(), retry: codex_api::RetryConfig { diff --git a/codex-rs/codex-api/tests/models_integration.rs b/codex-rs/codex-api/tests/models_integration.rs index d2b31180b907..e3653d82ec5f 100644 --- a/codex-rs/codex-api/tests/models_integration.rs +++ b/codex-rs/codex-api/tests/models_integration.rs @@ -32,6 +32,7 @@ fn provider(base_url: &str) -> Provider { Provider { name: "test".to_string(), base_url: base_url.to_string(), + uses_chatgpt_auth: false, query_params: None, headers: HeaderMap::new(), retry: RetryConfig { diff --git a/codex-rs/codex-api/tests/realtime_websocket_e2e.rs b/codex-rs/codex-api/tests/realtime_websocket_e2e.rs index cb9d7122f4b0..f8c59fa32e26 100644 --- a/codex-rs/codex-api/tests/realtime_websocket_e2e.rs +++ b/codex-rs/codex-api/tests/realtime_websocket_e2e.rs @@ -62,6 +62,7 @@ fn test_provider(base_url: String) -> Provider { Provider { name: "test".to_string(), base_url, + uses_chatgpt_auth: false, query_params: Some(HashMap::new()), headers: HeaderMap::new(), retry: RetryConfig { diff --git a/codex-rs/codex-api/tests/sse_end_to_end.rs b/codex-rs/codex-api/tests/sse_end_to_end.rs index bf880fefcf9f..bf31f87f77d3 100644 --- a/codex-rs/codex-api/tests/sse_end_to_end.rs +++ b/codex-rs/codex-api/tests/sse_end_to_end.rs @@ -61,6 +61,7 @@ fn provider(name: &str) -> Provider { Provider { name: name.to_string(), base_url: "https://example.com/v1".to_string(), + uses_chatgpt_auth: false, query_params: None, headers: HeaderMap::new(), retry: codex_api::RetryConfig { From 8ee7bf6abcdaaee2f057cb57f7da00a365f7df26 Mon Sep 17 00:00:00 2001 From: Jiaming Zhang Date: Thu, 7 May 2026 08:58:12 -0700 Subject: [PATCH 11/23] codex: wrap app-server attestation transport --- .../json/AttestationGenerateResponse.json | 2 +- .../codex_app_server_protocol.schemas.json | 2 +- .../v2/AttestationGenerateResponse.ts | 2 +- .../src/protocol/v2/attestation.rs | 2 +- codex-rs/app-server/README.md | 2 +- codex-rs/app-server/src/message_processor.rs | 72 +++++++++++++++++-- .../message_processor_attestation_tests.rs | 34 +++++++++ .../app-server/tests/suite/v2/attestation.rs | 3 +- 8 files changed, 108 insertions(+), 11 deletions(-) create mode 100644 codex-rs/app-server/src/message_processor_attestation_tests.rs diff --git a/codex-rs/app-server-protocol/schema/json/AttestationGenerateResponse.json b/codex-rs/app-server-protocol/schema/json/AttestationGenerateResponse.json index b7b7f8c474f0..921cc6f4947e 100644 --- a/codex-rs/app-server-protocol/schema/json/AttestationGenerateResponse.json +++ b/codex-rs/app-server-protocol/schema/json/AttestationGenerateResponse.json @@ -2,7 +2,7 @@ "$schema": "http://json-schema.org/draft-07/schema#", "properties": { "headerValue": { - "description": "Opaque upstream `x-oai-attestation` header value.", + "description": "Opaque client attestation payload to embed in the upstream header envelope.", "type": "string" } }, diff --git a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json index 71d9f81bc5d5..9cbb3e8c37b7 100644 --- a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json +++ b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json @@ -92,7 +92,7 @@ "$schema": "http://json-schema.org/draft-07/schema#", "properties": { "headerValue": { - "description": "Opaque upstream `x-oai-attestation` header value.", + "description": "Opaque client attestation payload to embed in the upstream header envelope.", "type": "string" } }, diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/AttestationGenerateResponse.ts b/codex-rs/app-server-protocol/schema/typescript/v2/AttestationGenerateResponse.ts index 48eef943fc54..7f11e8364f00 100644 --- a/codex-rs/app-server-protocol/schema/typescript/v2/AttestationGenerateResponse.ts +++ b/codex-rs/app-server-protocol/schema/typescript/v2/AttestationGenerateResponse.ts @@ -4,6 +4,6 @@ export type AttestationGenerateResponse = { /** - * Opaque upstream `x-oai-attestation` header value. + * Opaque client attestation payload to embed in the upstream header envelope. */ headerValue: string, }; diff --git a/codex-rs/app-server-protocol/src/protocol/v2/attestation.rs b/codex-rs/app-server-protocol/src/protocol/v2/attestation.rs index 36173b63609f..d2ad34b87980 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2/attestation.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2/attestation.rs @@ -12,6 +12,6 @@ pub struct AttestationGenerateParams {} #[serde(rename_all = "camelCase")] #[ts(export_to = "v2/")] pub struct AttestationGenerateResponse { - /// Opaque upstream `x-oai-attestation` header value. + /// Opaque client attestation payload to embed in the upstream header envelope. pub header_value: String, } diff --git a/codex-rs/app-server/README.md b/codex-rs/app-server/README.md index c55f84ab276a..6c48ac9b8a56 100644 --- a/codex-rs/app-server/README.md +++ b/codex-rs/app-server/README.md @@ -1324,7 +1324,7 @@ When the client responds to `item/tool/requestUserInput`, the server emits `serv ### Attestation generation -Desktop hosts that provide upstream attestation should set `capabilities.requestAttestation` during `initialize` and handle the server-initiated `attestation/generate` request. App-server issues it just in time before ChatGPT Codex requests that forward `x-oai-attestation`; the client responds with `{ "headerValue": "v1." }`, where `headerValue` is the complete upstream header value. App-server treats that value as opaque and forwards it unchanged. If no initialized client opted into attestation, or if the opted-in client is unavailable, times out, or returns invalid data, app-server omits `x-oai-attestation` for that upstream request. +Desktop hosts that provide upstream attestation should set `capabilities.requestAttestation` during `initialize` and handle the server-initiated `attestation/generate` request. App-server issues it just in time before ChatGPT Codex requests that forward `x-oai-attestation`; the client responds with `{ "headerValue": "v1." }`, where `headerValue` is an opaque client-owned payload. When app-server receives a client response, it forwards a consistent outer envelope such as `{ "v": 1, "s": 0, "t": "v1." }`, where `t` contains the client payload unchanged. If app-server attempts attestation but fails within its own boundary, it sends the same envelope shape with an app-server status code and without `t` (`1 = timeout`, `2 = request failed`, `3 = request canceled`, `4 = malformed response`). If no initialized client opted into attestation, app-server omits `x-oai-attestation` for that upstream request. ### MCP server elicitations diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index 9bf24c4f163f..68aca7ee75ae 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -80,6 +80,7 @@ use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::W3cTraceContext; use codex_rollout::StateDbHandle; use codex_state::log_db::LogDbLayer; +use serde::Serialize; use tokio::sync::Mutex; use tokio::sync::Semaphore; use tokio::sync::broadcast; @@ -91,6 +92,48 @@ use tracing::warn; const ATTESTATION_GENERATE_TIMEOUT: Duration = Duration::from_millis(100); const EXTERNAL_AUTH_REFRESH_TIMEOUT: Duration = Duration::from_secs(10); + +#[derive(Clone, Copy)] +enum AppServerAttestationStatus { + Ok, + Timeout, + RequestFailed, + RequestCanceled, + MalformedResponse, +} + +impl AppServerAttestationStatus { + const fn code(self) -> u8 { + match self { + Self::Ok => 0, + Self::Timeout => 1, + Self::RequestFailed => 2, + Self::RequestCanceled => 3, + Self::MalformedResponse => 4, + } + } +} + +#[derive(Serialize)] +struct AppServerAttestationEnvelope<'a> { + v: u8, + s: u8, + #[serde(skip_serializing_if = "Option::is_none")] + t: Option<&'a str>, +} + +fn app_server_attestation_header_value( + status: AppServerAttestationStatus, + token: Option<&str>, +) -> String { + serde_json::to_string(&AppServerAttestationEnvelope { + v: 1, + s: status.code(), + t: token, + }) + .expect("app-server attestation envelope should serialize") +} + #[derive(Clone)] struct ExternalAuthRefreshBridge { outgoing: Arc, @@ -228,11 +271,17 @@ async fn request_attestation_header_value_with_timeout( message = %err.message, "attestation generation request failed" ); - return Some(String::new()); + return Some(app_server_attestation_header_value( + AppServerAttestationStatus::RequestFailed, + None, + )); } Ok(Err(err)) => { warn!("attestation generation request canceled: {err}"); - return Some(String::new()); + return Some(app_server_attestation_header_value( + AppServerAttestationStatus::RequestCanceled, + None, + )); } Err(_) => { let _canceled = outgoing.cancel_request(&request_id).await; @@ -240,15 +289,24 @@ async fn request_attestation_header_value_with_timeout( timeout_seconds = timeout_duration.as_secs(), "attestation generation request timed out" ); - return Some(String::new()); + return Some(app_server_attestation_header_value( + AppServerAttestationStatus::Timeout, + None, + )); } }; match serde_json::from_value::(result) { - Ok(response) => Some(response.header_value), + Ok(response) => Some(app_server_attestation_header_value( + AppServerAttestationStatus::Ok, + Some(&response.header_value), + )), Err(err) => { warn!("failed to deserialize attestation generation response: {err}"); - Some(String::new()) + Some(app_server_attestation_header_value( + AppServerAttestationStatus::MalformedResponse, + None, + )) } } } @@ -1398,6 +1456,10 @@ impl MessageProcessor { } } +#[cfg(test)] +#[path = "message_processor_attestation_tests.rs"] +mod message_processor_attestation_tests; + #[cfg(test)] #[path = "message_processor_tracing_tests.rs"] mod message_processor_tracing_tests; diff --git a/codex-rs/app-server/src/message_processor_attestation_tests.rs b/codex-rs/app-server/src/message_processor_attestation_tests.rs new file mode 100644 index 000000000000..86b0b706ae5c --- /dev/null +++ b/codex-rs/app-server/src/message_processor_attestation_tests.rs @@ -0,0 +1,34 @@ +use super::AppServerAttestationStatus; +use super::app_server_attestation_header_value; +use pretty_assertions::assert_eq; + +#[test] +fn app_server_attestation_header_value_wraps_opaque_client_payloads() { + assert_eq!( + app_server_attestation_header_value( + AppServerAttestationStatus::Ok, + Some("v1.opaque-client-payload"), + ), + r#"{"v":1,"s":0,"t":"v1.opaque-client-payload"}"# + ); +} + +#[test] +fn app_server_attestation_header_value_reports_app_server_failures() { + assert_eq!( + app_server_attestation_header_value(AppServerAttestationStatus::Timeout, None), + r#"{"v":1,"s":1}"# + ); + assert_eq!( + app_server_attestation_header_value(AppServerAttestationStatus::RequestFailed, None), + r#"{"v":1,"s":2}"# + ); + assert_eq!( + app_server_attestation_header_value(AppServerAttestationStatus::RequestCanceled, None), + r#"{"v":1,"s":3}"# + ); + assert_eq!( + app_server_attestation_header_value(AppServerAttestationStatus::MalformedResponse, None), + r#"{"v":1,"s":4}"# + ); +} diff --git a/codex-rs/app-server/tests/suite/v2/attestation.rs b/codex-rs/app-server/tests/suite/v2/attestation.rs index 5030b61f2ea7..101d7192a92e 100644 --- a/codex-rs/app-server/tests/suite/v2/attestation.rs +++ b/codex-rs/app-server/tests/suite/v2/attestation.rs @@ -29,6 +29,7 @@ use tokio::time::timeout; const DEFAULT_READ_TIMEOUT: Duration = Duration::from_secs(60); const ATTESTATION_HEADER: &str = "v1.integration-test"; +const APP_SERVER_ATTESTATION_HEADER: &str = r#"{"v":1,"s":0,"t":"v1.integration-test"}"#; #[tokio::test] async fn attestation_generate_round_trip_adds_header_to_responses_websocket_handshake() -> Result<()> @@ -161,7 +162,7 @@ async fn attestation_generate_round_trip_adds_header_to_responses_websocket_hand let handshake = websocket_server.single_handshake(); assert_eq!( handshake.header("x-oai-attestation").as_deref(), - Some(ATTESTATION_HEADER) + Some(APP_SERVER_ATTESTATION_HEADER) ); websocket_server.shutdown().await; From e77b359ade63b4f0f2320784285062a003ab2a78 Mon Sep 17 00:00:00 2001 From: Jiaming Zhang Date: Thu, 7 May 2026 09:14:43 -0700 Subject: [PATCH 12/23] codex: move app-server attestation into module --- codex-rs/app-server/src/attestation.rs | 204 ++++++++++++++++++ codex-rs/app-server/src/lib.rs | 1 + codex-rs/app-server/src/message_processor.rs | 164 +------------- .../message_processor_attestation_tests.rs | 34 --- 4 files changed, 206 insertions(+), 197 deletions(-) create mode 100644 codex-rs/app-server/src/attestation.rs delete mode 100644 codex-rs/app-server/src/message_processor_attestation_tests.rs diff --git a/codex-rs/app-server/src/attestation.rs b/codex-rs/app-server/src/attestation.rs new file mode 100644 index 000000000000..adfb415babbb --- /dev/null +++ b/codex-rs/app-server/src/attestation.rs @@ -0,0 +1,204 @@ +use std::sync::Arc; + +use axum::http::HeaderValue; +use codex_app_server_protocol::AttestationGenerateParams; +use codex_app_server_protocol::AttestationGenerateResponse; +use codex_app_server_protocol::ServerRequestPayload; +use codex_core::AttestationContext; +use codex_core::AttestationProvider; +use codex_core::GenerateAttestationFuture; +use serde::Serialize; +use tokio::time::Duration; +use tokio::time::timeout; +use tracing::warn; + +use crate::outgoing_message::OutgoingMessageSender; +use crate::thread_state::ThreadStateManager; + +const ATTESTATION_GENERATE_TIMEOUT: Duration = Duration::from_millis(100); + +pub(crate) fn app_server_attestation_provider( + outgoing: Arc, + thread_state_manager: ThreadStateManager, +) -> Arc { + Arc::new(AppServerAttestationProvider { + outgoing, + thread_state_manager, + }) +} + +struct AppServerAttestationProvider { + outgoing: Arc, + thread_state_manager: ThreadStateManager, +} + +impl std::fmt::Debug for AppServerAttestationProvider { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("AppServerAttestationProvider") + .finish() + } +} + +impl AttestationProvider for AppServerAttestationProvider { + fn header_for_request(&self, context: AttestationContext) -> GenerateAttestationFuture<'_> { + let outgoing = self.outgoing.clone(); + let thread_state_manager = self.thread_state_manager.clone(); + Box::pin(async move { + if !context.uses_chatgpt_auth { + return None; + } + + request_attestation_header_value_with_timeout( + outgoing, + thread_state_manager, + ATTESTATION_GENERATE_TIMEOUT, + ) + .await + .and_then(|value| HeaderValue::from_bytes(value.as_bytes()).ok()) + }) + } +} + +async fn request_attestation_header_value_with_timeout( + outgoing: Arc, + thread_state_manager: ThreadStateManager, + timeout_duration: Duration, +) -> Option { + let connection_id = thread_state_manager + .first_attestation_capable_connection() + .await?; + + let connection_ids = [connection_id]; + let (request_id, rx) = outgoing + .send_request_to_connections( + Some(&connection_ids), + ServerRequestPayload::AttestationGenerate(AttestationGenerateParams {}), + /*thread_id*/ None, + ) + .await; + + let result = match timeout(timeout_duration, rx).await { + Ok(Ok(Ok(result))) => result, + Ok(Ok(Err(err))) => { + warn!( + code = err.code, + message = %err.message, + "attestation generation request failed" + ); + return app_server_attestation_header_value( + AppServerAttestationStatus::RequestFailed, + None, + ); + } + Ok(Err(err)) => { + warn!("attestation generation request canceled: {err}"); + return app_server_attestation_header_value( + AppServerAttestationStatus::RequestCanceled, + None, + ); + } + Err(_) => { + let _canceled = outgoing.cancel_request(&request_id).await; + warn!( + timeout_seconds = timeout_duration.as_secs(), + "attestation generation request timed out" + ); + return app_server_attestation_header_value(AppServerAttestationStatus::Timeout, None); + } + }; + + match serde_json::from_value::(result) { + Ok(response) => app_server_attestation_header_value( + AppServerAttestationStatus::Ok, + Some(&response.header_value), + ), + Err(err) => { + warn!("failed to deserialize attestation generation response: {err}"); + app_server_attestation_header_value(AppServerAttestationStatus::MalformedResponse, None) + } + } +} + +#[derive(Clone, Copy)] +enum AppServerAttestationStatus { + Ok, + Timeout, + RequestFailed, + RequestCanceled, + MalformedResponse, +} + +impl AppServerAttestationStatus { + const fn code(self) -> u8 { + match self { + Self::Ok => 0, + Self::Timeout => 1, + Self::RequestFailed => 2, + Self::RequestCanceled => 3, + Self::MalformedResponse => 4, + } + } +} + +#[derive(Serialize)] +struct AppServerAttestationEnvelope<'a> { + v: u8, + s: u8, + #[serde(skip_serializing_if = "Option::is_none")] + t: Option<&'a str>, +} + +fn app_server_attestation_header_value( + status: AppServerAttestationStatus, + token: Option<&str>, +) -> Option { + serde_json::to_string(&AppServerAttestationEnvelope { + v: 1, + s: status.code(), + t: token, + }) + .map_err(|err| warn!("failed to serialize app-server attestation envelope: {err}")) + .ok() +} + +#[cfg(test)] +mod tests { + use super::AppServerAttestationStatus; + use super::app_server_attestation_header_value; + use pretty_assertions::assert_eq; + + #[test] + fn app_server_attestation_header_value_wraps_opaque_client_payloads() { + assert_eq!( + app_server_attestation_header_value( + AppServerAttestationStatus::Ok, + Some("v1.opaque-client-payload"), + ), + Some(r#"{"v":1,"s":0,"t":"v1.opaque-client-payload"}"#.to_string()) + ); + } + + #[test] + fn app_server_attestation_header_value_reports_app_server_failures() { + assert_eq!( + app_server_attestation_header_value(AppServerAttestationStatus::Timeout, None), + Some(r#"{"v":1,"s":1}"#.to_string()) + ); + assert_eq!( + app_server_attestation_header_value(AppServerAttestationStatus::RequestFailed, None), + Some(r#"{"v":1,"s":2}"#.to_string()) + ); + assert_eq!( + app_server_attestation_header_value(AppServerAttestationStatus::RequestCanceled, None), + Some(r#"{"v":1,"s":3}"#.to_string()) + ); + assert_eq!( + app_server_attestation_header_value( + AppServerAttestationStatus::MalformedResponse, + None + ), + Some(r#"{"v":1,"s":4}"#.to_string()) + ); + } +} diff --git a/codex-rs/app-server/src/lib.rs b/codex-rs/app-server/src/lib.rs index 9bb7937e19db..dd6397170506 100644 --- a/codex-rs/app-server/src/lib.rs +++ b/codex-rs/app-server/src/lib.rs @@ -74,6 +74,7 @@ use tracing_subscriber::util::SubscriberInitExt; mod analytics_utils; mod app_server_tracing; +mod attestation; mod bespoke_event_handling; mod command_exec; mod config; diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index 68aca7ee75ae..ea89f0b03667 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -4,6 +4,7 @@ use std::sync::Arc; use std::sync::OnceLock; use std::sync::atomic::AtomicBool; +use crate::attestation::app_server_attestation_provider; use crate::config_manager::ConfigManager; use crate::connection_rpc_gate::ConnectionRpcGate; use crate::error_code::invalid_request; @@ -39,11 +40,8 @@ use crate::thread_state::ThreadStateManager; use crate::transport::AppServerTransport; use crate::transport::RemoteControlHandle; use async_trait::async_trait; -use axum::http::HeaderValue; use codex_analytics::AnalyticsEventsClient; use codex_analytics::AppServerRpcTransport; -use codex_app_server_protocol::AttestationGenerateParams; -use codex_app_server_protocol::AttestationGenerateResponse; use codex_app_server_protocol::AuthMode as LoginAuthMode; use codex_app_server_protocol::ChatgptAuthTokensRefreshParams; use codex_app_server_protocol::ChatgptAuthTokensRefreshReason; @@ -62,9 +60,6 @@ use codex_app_server_protocol::ServerRequestPayload; use codex_app_server_protocol::experimental_required_message; use codex_arg0::Arg0DispatchPaths; use codex_chatgpt::workspace_settings; -use codex_core::AttestationContext; -use codex_core::AttestationProvider; -use codex_core::GenerateAttestationFuture; use codex_core::ThreadManager; use codex_core::config::Config; use codex_core::thread_store_from_config; @@ -80,7 +75,6 @@ use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::W3cTraceContext; use codex_rollout::StateDbHandle; use codex_state::log_db::LogDbLayer; -use serde::Serialize; use tokio::sync::Mutex; use tokio::sync::Semaphore; use tokio::sync::broadcast; @@ -88,52 +82,9 @@ use tokio::sync::watch; use tokio::time::Duration; use tokio::time::timeout; use tracing::Instrument; -use tracing::warn; -const ATTESTATION_GENERATE_TIMEOUT: Duration = Duration::from_millis(100); const EXTERNAL_AUTH_REFRESH_TIMEOUT: Duration = Duration::from_secs(10); -#[derive(Clone, Copy)] -enum AppServerAttestationStatus { - Ok, - Timeout, - RequestFailed, - RequestCanceled, - MalformedResponse, -} - -impl AppServerAttestationStatus { - const fn code(self) -> u8 { - match self { - Self::Ok => 0, - Self::Timeout => 1, - Self::RequestFailed => 2, - Self::RequestCanceled => 3, - Self::MalformedResponse => 4, - } - } -} - -#[derive(Serialize)] -struct AppServerAttestationEnvelope<'a> { - v: u8, - s: u8, - #[serde(skip_serializing_if = "Option::is_none")] - t: Option<&'a str>, -} - -fn app_server_attestation_header_value( - status: AppServerAttestationStatus, - token: Option<&str>, -) -> String { - serde_json::to_string(&AppServerAttestationEnvelope { - v: 1, - s: status.code(), - t: token, - }) - .expect("app-server attestation envelope should serialize") -} - #[derive(Clone)] struct ExternalAuthRefreshBridge { outgoing: Arc, @@ -202,115 +153,6 @@ impl ExternalAuth for ExternalAuthRefreshBridge { } } -fn app_server_attestation_provider( - outgoing: Arc, - thread_state_manager: ThreadStateManager, -) -> Arc { - Arc::new(AppServerAttestationProvider { - outgoing, - thread_state_manager, - }) -} - -struct AppServerAttestationProvider { - outgoing: Arc, - thread_state_manager: ThreadStateManager, -} - -impl std::fmt::Debug for AppServerAttestationProvider { - fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - formatter - .debug_struct("AppServerAttestationProvider") - .finish() - } -} - -impl AttestationProvider for AppServerAttestationProvider { - fn header_for_request(&self, context: AttestationContext) -> GenerateAttestationFuture<'_> { - let outgoing = self.outgoing.clone(); - let thread_state_manager = self.thread_state_manager.clone(); - Box::pin(async move { - if !context.uses_chatgpt_auth { - return None; - } - - request_attestation_header_value_with_timeout( - outgoing, - thread_state_manager, - ATTESTATION_GENERATE_TIMEOUT, - ) - .await - .and_then(|value| HeaderValue::from_bytes(value.as_bytes()).ok()) - }) - } -} - -async fn request_attestation_header_value_with_timeout( - outgoing: Arc, - thread_state_manager: ThreadStateManager, - timeout_duration: Duration, -) -> Option { - let connection_id = thread_state_manager - .first_attestation_capable_connection() - .await?; - - let connection_ids = [connection_id]; - let (request_id, rx) = outgoing - .send_request_to_connections( - Some(&connection_ids), - ServerRequestPayload::AttestationGenerate(AttestationGenerateParams {}), - /*thread_id*/ None, - ) - .await; - - let result = match timeout(timeout_duration, rx).await { - Ok(Ok(Ok(result))) => result, - Ok(Ok(Err(err))) => { - warn!( - code = err.code, - message = %err.message, - "attestation generation request failed" - ); - return Some(app_server_attestation_header_value( - AppServerAttestationStatus::RequestFailed, - None, - )); - } - Ok(Err(err)) => { - warn!("attestation generation request canceled: {err}"); - return Some(app_server_attestation_header_value( - AppServerAttestationStatus::RequestCanceled, - None, - )); - } - Err(_) => { - let _canceled = outgoing.cancel_request(&request_id).await; - warn!( - timeout_seconds = timeout_duration.as_secs(), - "attestation generation request timed out" - ); - return Some(app_server_attestation_header_value( - AppServerAttestationStatus::Timeout, - None, - )); - } - }; - - match serde_json::from_value::(result) { - Ok(response) => Some(app_server_attestation_header_value( - AppServerAttestationStatus::Ok, - Some(&response.header_value), - )), - Err(err) => { - warn!("failed to deserialize attestation generation response: {err}"); - Some(app_server_attestation_header_value( - AppServerAttestationStatus::MalformedResponse, - None, - )) - } - } -} - pub(crate) struct MessageProcessor { outgoing: Arc, account_processor: AccountRequestProcessor, @@ -1456,10 +1298,6 @@ impl MessageProcessor { } } -#[cfg(test)] -#[path = "message_processor_attestation_tests.rs"] -mod message_processor_attestation_tests; - #[cfg(test)] #[path = "message_processor_tracing_tests.rs"] mod message_processor_tracing_tests; diff --git a/codex-rs/app-server/src/message_processor_attestation_tests.rs b/codex-rs/app-server/src/message_processor_attestation_tests.rs deleted file mode 100644 index 86b0b706ae5c..000000000000 --- a/codex-rs/app-server/src/message_processor_attestation_tests.rs +++ /dev/null @@ -1,34 +0,0 @@ -use super::AppServerAttestationStatus; -use super::app_server_attestation_header_value; -use pretty_assertions::assert_eq; - -#[test] -fn app_server_attestation_header_value_wraps_opaque_client_payloads() { - assert_eq!( - app_server_attestation_header_value( - AppServerAttestationStatus::Ok, - Some("v1.opaque-client-payload"), - ), - r#"{"v":1,"s":0,"t":"v1.opaque-client-payload"}"# - ); -} - -#[test] -fn app_server_attestation_header_value_reports_app_server_failures() { - assert_eq!( - app_server_attestation_header_value(AppServerAttestationStatus::Timeout, None), - r#"{"v":1,"s":1}"# - ); - assert_eq!( - app_server_attestation_header_value(AppServerAttestationStatus::RequestFailed, None), - r#"{"v":1,"s":2}"# - ); - assert_eq!( - app_server_attestation_header_value(AppServerAttestationStatus::RequestCanceled, None), - r#"{"v":1,"s":3}"# - ); - assert_eq!( - app_server_attestation_header_value(AppServerAttestationStatus::MalformedResponse, None), - r#"{"v":1,"s":4}"# - ); -} From bdb777233c78cbb9de7c5adbbca053364ab4170b Mon Sep 17 00:00:00 2001 From: Jiaming Zhang Date: Thu, 7 May 2026 11:21:56 -0700 Subject: [PATCH 13/23] fix(app-server): scope attestation to thread connections --- codex-rs/app-server/src/attestation.rs | 4 +- .../thread_processor_tests.rs | 56 ++++++++++++++++--- codex-rs/app-server/src/thread_state.rs | 22 +++++--- codex-rs/core/src/attestation.rs | 3 + codex-rs/core/src/client.rs | 1 + 5 files changed, 71 insertions(+), 15 deletions(-) diff --git a/codex-rs/app-server/src/attestation.rs b/codex-rs/app-server/src/attestation.rs index adfb415babbb..35cc82a18e0a 100644 --- a/codex-rs/app-server/src/attestation.rs +++ b/codex-rs/app-server/src/attestation.rs @@ -52,6 +52,7 @@ impl AttestationProvider for AppServerAttestationProvider { request_attestation_header_value_with_timeout( outgoing, thread_state_manager, + context.thread_id, ATTESTATION_GENERATE_TIMEOUT, ) .await @@ -63,10 +64,11 @@ impl AttestationProvider for AppServerAttestationProvider { async fn request_attestation_header_value_with_timeout( outgoing: Arc, thread_state_manager: ThreadStateManager, + thread_id: codex_protocol::ThreadId, timeout_duration: Duration, ) -> Option { let connection_id = thread_state_manager - .first_attestation_capable_connection() + .first_attestation_capable_connection_for_thread(thread_id) .await?; let connection_ids = [connection_id]; diff --git a/codex-rs/app-server/src/request_processors/thread_processor_tests.rs b/codex-rs/app-server/src/request_processors/thread_processor_tests.rs index 3cfde4ff4324..5068f199635b 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor_tests.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor_tests.rs @@ -1278,18 +1278,27 @@ mod thread_processor_behavior_tests { } #[tokio::test] - async fn first_attestation_capable_connection_returns_lowest_live_connection_id() { + async fn first_attestation_capable_connection_for_thread_only_uses_thread_subscribers() + -> Result<()> { let manager = ThreadStateManager::new(); - let unsupported_connection = ConnectionId(1); - let later_supported_connection = ConnectionId(3); + let thread_id = ThreadId::from_string("dfbd9a95-2f44-470a-8bd8-1cfc04efc243")?; + let other_thread_id = ThreadId::from_string("6c9a74e4-5e59-479e-90bf-5c5798bb50aa")?; + let unrelated_supported_connection = ConnectionId(1); let earlier_supported_connection = ConnectionId(2); + let later_supported_connection = ConnectionId(3); + let unsupported_connection = ConnectionId(4); manager - .connection_initialized(unsupported_connection, ConnectionCapabilities::default()) + .connection_initialized( + unrelated_supported_connection, + ConnectionCapabilities { + request_attestation: true, + }, + ) .await; manager .connection_initialized( - later_supported_connection, + earlier_supported_connection, ConnectionCapabilities { request_attestation: true, }, @@ -1297,16 +1306,49 @@ mod thread_processor_behavior_tests { .await; manager .connection_initialized( - earlier_supported_connection, + later_supported_connection, ConnectionCapabilities { request_attestation: true, }, ) .await; + manager + .connection_initialized(unsupported_connection, ConnectionCapabilities::default()) + .await; + + assert!( + manager + .try_add_connection_to_thread(other_thread_id, unrelated_supported_connection) + .await + ); + assert!( + manager + .try_add_connection_to_thread(thread_id, later_supported_connection) + .await + ); + assert!( + manager + .try_add_connection_to_thread(thread_id, earlier_supported_connection) + .await + ); + assert!( + manager + .try_add_connection_to_thread(thread_id, unsupported_connection) + .await + ); assert_eq!( - manager.first_attestation_capable_connection().await, + manager + .first_attestation_capable_connection_for_thread(thread_id) + .await, Some(earlier_supported_connection) ); + assert_eq!( + manager + .first_attestation_capable_connection_for_thread(other_thread_id) + .await, + Some(unrelated_supported_connection) + ); + Ok(()) } } diff --git a/codex-rs/app-server/src/thread_state.rs b/codex-rs/app-server/src/thread_state.rs index 8552d21f710e..82871fca8b40 100644 --- a/codex-rs/app-server/src/thread_state.rs +++ b/codex-rs/app-server/src/thread_state.rs @@ -231,14 +231,22 @@ impl ThreadStateManager { .insert(connection_id, capabilities); } - pub(crate) async fn first_attestation_capable_connection(&self) -> Option { - self.state - .lock() - .await - .live_connections + pub(crate) async fn first_attestation_capable_connection_for_thread( + &self, + thread_id: ThreadId, + ) -> Option { + let state = self.state.lock().await; + state + .threads + .get(&thread_id)? + .connection_ids .iter() - .filter_map(|(connection_id, capabilities)| { - capabilities.request_attestation.then_some(*connection_id) + .filter_map(|connection_id| { + state + .live_connections + .get(connection_id)? + .request_attestation + .then_some(*connection_id) }) .min_by_key(|connection_id| connection_id.0) } diff --git a/codex-rs/core/src/attestation.rs b/codex-rs/core/src/attestation.rs index 01d4e72e4ea8..c985f91b29c5 100644 --- a/codex-rs/core/src/attestation.rs +++ b/codex-rs/core/src/attestation.rs @@ -1,6 +1,7 @@ use std::future::Future; use std::pin::Pin; +use codex_protocol::ThreadId; use http::HeaderValue; pub(crate) const X_OAI_ATTESTATION_HEADER: &str = "x-oai-attestation"; @@ -13,6 +14,8 @@ pub type GenerateAttestationFuture<'a> = #[derive(Clone, Copy, Debug)] pub struct AttestationContext { pub uses_chatgpt_auth: bool, + /// Thread whose upstream request is being prepared. + pub thread_id: ThreadId, } /// Host integration boundary for just-in-time attestation header values. diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index ab850770ca72..34cdcde8189f 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -659,6 +659,7 @@ impl ModelClient { .as_ref()? .header_for_request(AttestationContext { uses_chatgpt_auth: provider.uses_chatgpt_auth, + thread_id: self.state.thread_id, }) .await } From 5daabda6a4443b861a15acc0866189489a9b53f8 Mon Sep 17 00:00:00 2001 From: Jiaming Zhang Date: Thu, 7 May 2026 11:52:05 -0700 Subject: [PATCH 14/23] codex: fix CI failure on PR #20619 --- codex-rs/app-server/src/attestation.rs | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/codex-rs/app-server/src/attestation.rs b/codex-rs/app-server/src/attestation.rs index 35cc82a18e0a..66eedbbad6c2 100644 --- a/codex-rs/app-server/src/attestation.rs +++ b/codex-rs/app-server/src/attestation.rs @@ -90,14 +90,14 @@ async fn request_attestation_header_value_with_timeout( ); return app_server_attestation_header_value( AppServerAttestationStatus::RequestFailed, - None, + /*token*/ None, ); } Ok(Err(err)) => { warn!("attestation generation request canceled: {err}"); return app_server_attestation_header_value( AppServerAttestationStatus::RequestCanceled, - None, + /*token*/ None, ); } Err(_) => { @@ -106,7 +106,10 @@ async fn request_attestation_header_value_with_timeout( timeout_seconds = timeout_duration.as_secs(), "attestation generation request timed out" ); - return app_server_attestation_header_value(AppServerAttestationStatus::Timeout, None); + return app_server_attestation_header_value( + AppServerAttestationStatus::Timeout, + /*token*/ None, + ); } }; @@ -117,7 +120,10 @@ async fn request_attestation_header_value_with_timeout( ), Err(err) => { warn!("failed to deserialize attestation generation response: {err}"); - app_server_attestation_header_value(AppServerAttestationStatus::MalformedResponse, None) + app_server_attestation_header_value( + AppServerAttestationStatus::MalformedResponse, + /*token*/ None, + ) } } } From 5da2adef8ec1b258cf57c7df8866dbaacee3222e Mon Sep 17 00:00:00 2001 From: Jiaming Zhang Date: Thu, 7 May 2026 11:59:50 -0700 Subject: [PATCH 15/23] codex: fix CI failure on PR #20619 --- codex-rs/app-server/src/attestation.rs | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/codex-rs/app-server/src/attestation.rs b/codex-rs/app-server/src/attestation.rs index 66eedbbad6c2..75c57107c2ba 100644 --- a/codex-rs/app-server/src/attestation.rs +++ b/codex-rs/app-server/src/attestation.rs @@ -190,21 +190,30 @@ mod tests { #[test] fn app_server_attestation_header_value_reports_app_server_failures() { assert_eq!( - app_server_attestation_header_value(AppServerAttestationStatus::Timeout, None), + app_server_attestation_header_value( + AppServerAttestationStatus::Timeout, + /*token*/ None, + ), Some(r#"{"v":1,"s":1}"#.to_string()) ); assert_eq!( - app_server_attestation_header_value(AppServerAttestationStatus::RequestFailed, None), + app_server_attestation_header_value( + AppServerAttestationStatus::RequestFailed, + /*token*/ None, + ), Some(r#"{"v":1,"s":2}"#.to_string()) ); assert_eq!( - app_server_attestation_header_value(AppServerAttestationStatus::RequestCanceled, None), + app_server_attestation_header_value( + AppServerAttestationStatus::RequestCanceled, + /*token*/ None, + ), Some(r#"{"v":1,"s":3}"#.to_string()) ); assert_eq!( app_server_attestation_header_value( AppServerAttestationStatus::MalformedResponse, - None + /*token*/ None ), Some(r#"{"v":1,"s":4}"#.to_string()) ); From a99661637c10aab2b0c259f1b6a2c2bddc4e4e6e Mon Sep 17 00:00:00 2001 From: Jiaming Zhang Date: Thu, 7 May 2026 12:13:51 -0700 Subject: [PATCH 16/23] codex: fix CI failure on PR #20619 --- codex-rs/core/src/thread_manager_tests.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index 1cf0df95f2ea..2a7bf7d2c3a4 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -764,6 +764,7 @@ async fn rollout_path_resume_and_fork_read_history_through_thread_store() { thread_store.clone(), state_db, TEST_INSTALLATION_ID.to_string(), + /*attestation_provider*/ None, ); let source = manager From eac9cff3665e34fdf7c508d80de5e9e6339948fc Mon Sep 17 00:00:00 2001 From: Jiaming Zhang Date: Thu, 7 May 2026 17:27:56 -0700 Subject: [PATCH 17/23] codex: address PR review feedback (#20619) --- .../schema/json/AttestationGenerateResponse.json | 6 +++--- .../schema/json/codex_app_server_protocol.schemas.json | 8 ++++---- .../schema/typescript/v2/AttestationGenerateResponse.ts | 4 ++-- .../app-server-protocol/src/protocol/v2/attestation.rs | 4 ++-- codex-rs/app-server/README.md | 2 +- codex-rs/app-server/src/attestation.rs | 2 +- codex-rs/app-server/tests/suite/v2/attestation.rs | 2 +- 7 files changed, 14 insertions(+), 14 deletions(-) diff --git a/codex-rs/app-server-protocol/schema/json/AttestationGenerateResponse.json b/codex-rs/app-server-protocol/schema/json/AttestationGenerateResponse.json index 921cc6f4947e..e6bd59ec250c 100644 --- a/codex-rs/app-server-protocol/schema/json/AttestationGenerateResponse.json +++ b/codex-rs/app-server-protocol/schema/json/AttestationGenerateResponse.json @@ -1,13 +1,13 @@ { "$schema": "http://json-schema.org/draft-07/schema#", "properties": { - "headerValue": { - "description": "Opaque client attestation payload to embed in the upstream header envelope.", + "token": { + "description": "Opaque client attestation token.", "type": "string" } }, "required": [ - "headerValue" + "token" ], "title": "AttestationGenerateResponse", "type": "object" diff --git a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json index 9cbb3e8c37b7..c68795c015de 100644 --- a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json +++ b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json @@ -91,13 +91,13 @@ "AttestationGenerateResponse": { "$schema": "http://json-schema.org/draft-07/schema#", "properties": { - "headerValue": { - "description": "Opaque client attestation payload to embed in the upstream header envelope.", + "token": { + "description": "Opaque client attestation token.", "type": "string" } }, "required": [ - "headerValue" + "token" ], "title": "AttestationGenerateResponse", "type": "object" @@ -18438,4 +18438,4 @@ }, "title": "CodexAppServerProtocol", "type": "object" -} +} \ No newline at end of file diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/AttestationGenerateResponse.ts b/codex-rs/app-server-protocol/schema/typescript/v2/AttestationGenerateResponse.ts index 7f11e8364f00..6821c898ece7 100644 --- a/codex-rs/app-server-protocol/schema/typescript/v2/AttestationGenerateResponse.ts +++ b/codex-rs/app-server-protocol/schema/typescript/v2/AttestationGenerateResponse.ts @@ -4,6 +4,6 @@ export type AttestationGenerateResponse = { /** - * Opaque client attestation payload to embed in the upstream header envelope. + * Opaque client attestation token. */ -headerValue: string, }; +token: string, }; diff --git a/codex-rs/app-server-protocol/src/protocol/v2/attestation.rs b/codex-rs/app-server-protocol/src/protocol/v2/attestation.rs index d2ad34b87980..ef8828b58067 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2/attestation.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2/attestation.rs @@ -12,6 +12,6 @@ pub struct AttestationGenerateParams {} #[serde(rename_all = "camelCase")] #[ts(export_to = "v2/")] pub struct AttestationGenerateResponse { - /// Opaque client attestation payload to embed in the upstream header envelope. - pub header_value: String, + /// Opaque client attestation token. + pub token: String, } diff --git a/codex-rs/app-server/README.md b/codex-rs/app-server/README.md index 6c48ac9b8a56..6a3e6f4b3ad6 100644 --- a/codex-rs/app-server/README.md +++ b/codex-rs/app-server/README.md @@ -1324,7 +1324,7 @@ When the client responds to `item/tool/requestUserInput`, the server emits `serv ### Attestation generation -Desktop hosts that provide upstream attestation should set `capabilities.requestAttestation` during `initialize` and handle the server-initiated `attestation/generate` request. App-server issues it just in time before ChatGPT Codex requests that forward `x-oai-attestation`; the client responds with `{ "headerValue": "v1." }`, where `headerValue` is an opaque client-owned payload. When app-server receives a client response, it forwards a consistent outer envelope such as `{ "v": 1, "s": 0, "t": "v1." }`, where `t` contains the client payload unchanged. If app-server attempts attestation but fails within its own boundary, it sends the same envelope shape with an app-server status code and without `t` (`1 = timeout`, `2 = request failed`, `3 = request canceled`, `4 = malformed response`). If no initialized client opted into attestation, app-server omits `x-oai-attestation` for that upstream request. +Desktop hosts that provide upstream attestation should set `capabilities.requestAttestation` during `initialize` and handle the server-initiated `attestation/generate` request. App-server issues it just in time before ChatGPT Codex requests that forward `x-oai-attestation`; the client responds with `{ "token": "v1." }`, where `token` is an opaque client-owned value. When app-server receives a client response, it forwards a consistent outer envelope such as `{ "v": 1, "s": 0, "t": "v1." }`, where `t` contains the client token unchanged. If app-server attempts attestation but fails within its own boundary, it sends the same envelope shape with an app-server status code and without `t` (`1 = timeout`, `2 = request failed`, `3 = request canceled`, `4 = malformed response`). If no initialized client opted into attestation, app-server omits `x-oai-attestation` for that upstream request. ### MCP server elicitations diff --git a/codex-rs/app-server/src/attestation.rs b/codex-rs/app-server/src/attestation.rs index 75c57107c2ba..b1fbf6945a2e 100644 --- a/codex-rs/app-server/src/attestation.rs +++ b/codex-rs/app-server/src/attestation.rs @@ -116,7 +116,7 @@ async fn request_attestation_header_value_with_timeout( match serde_json::from_value::(result) { Ok(response) => app_server_attestation_header_value( AppServerAttestationStatus::Ok, - Some(&response.header_value), + Some(&response.token), ), Err(err) => { warn!("failed to deserialize attestation generation response: {err}"); diff --git a/codex-rs/app-server/tests/suite/v2/attestation.rs b/codex-rs/app-server/tests/suite/v2/attestation.rs index 101d7192a92e..d0565e2571d8 100644 --- a/codex-rs/app-server/tests/suite/v2/attestation.rs +++ b/codex-rs/app-server/tests/suite/v2/attestation.rs @@ -137,7 +137,7 @@ async fn attestation_generate_round_trip_adds_header_to_responses_websocket_hand mcp.send_response( request_id, serde_json::to_value(AttestationGenerateResponse { - header_value: ATTESTATION_HEADER.to_string(), + token: ATTESTATION_HEADER.to_string(), })?, ) .await?; From 0712d3b9210d4c81eea931fbcfc18260d79f6f1e Mon Sep 17 00:00:00 2001 From: Jiaming Zhang Date: Thu, 7 May 2026 18:02:04 -0700 Subject: [PATCH 18/23] codex: move attestation support to model provider --- codex-rs/Cargo.lock | 118 ------------------ codex-rs/app-server/src/attestation.rs | 4 - codex-rs/codex-api/src/endpoint/memories.rs | 1 - codex-rs/codex-api/src/endpoint/models.rs | 1 - .../codex-api/src/endpoint/realtime_call.rs | 1 - .../endpoint/realtime_websocket/methods.rs | 5 - codex-rs/codex-api/src/provider.rs | 1 - codex-rs/codex-api/tests/clients.rs | 1 - .../codex-api/tests/models_integration.rs | 1 - .../codex-api/tests/realtime_websocket_e2e.rs | 1 - codex-rs/codex-api/tests/sse_end_to_end.rs | 1 - codex-rs/core/src/attestation.rs | 1 - codex-rs/core/src/client.rs | 36 +++--- codex-rs/core/src/client_tests.rs | 47 ++++--- codex-rs/model-provider-info/src/lib.rs | 4 - codex-rs/model-provider/src/provider.rs | 12 ++ 16 files changed, 58 insertions(+), 177 deletions(-) diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index c656a791a60c..e544fdb5f3bc 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -1180,12 +1180,6 @@ dependencies = [ "windows-link", ] -[[package]] -name = "base16ct" -version = "0.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4c7f02d4ea65f2c1853089ffd8d2787bdbc63de2f0d29dedbcf8ccdfa0ccd4cf" - [[package]] name = "base64" version = "0.21.7" @@ -4425,18 +4419,6 @@ version = "0.2.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" -[[package]] -name = "crypto-bigint" -version = "0.5.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0dc92fb57ca44df6db8059111ab3af99a63d5d0f8375d9972e319a379c6bab76" -dependencies = [ - "generic-array", - "rand_core 0.6.4", - "subtle", - "zeroize", -] - [[package]] name = "crypto-common" version = "0.1.7" @@ -5150,20 +5132,6 @@ version = "1.0.20" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d0881ea181b1df73ff77ffaaf9c7544ecc11e82fba9b5f27b262a3c73a332555" -[[package]] -name = "ecdsa" -version = "0.16.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ee27f32b5c5292967d2d4a9d7f1e0b0aed2c15daded5a60300e4abb9d8020bca" -dependencies = [ - "der", - "digest", - "elliptic-curve", - "rfc6979", - "signature", - "spki", -] - [[package]] name = "ed25519" version = "2.2.3" @@ -5197,26 +5165,6 @@ dependencies = [ "serde", ] -[[package]] -name = "elliptic-curve" -version = "0.13.8" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b5e6043086bf7973472e0c7dff2142ea0b680d30e18d9cc40f267efbf222bd47" -dependencies = [ - "base16ct", - "crypto-bigint", - "digest", - "ff", - "generic-array", - "group", - "pem-rfc7468", - "pkcs8", - "rand_core 0.6.4", - "sec1", - "subtle", - "zeroize", -] - [[package]] name = "ena" version = "0.14.3" @@ -5470,16 +5418,6 @@ dependencies = [ "simd-adler32", ] -[[package]] -name = "ff" -version = "0.13.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c0b50bfb653653f9ca9095b427bed08ab8d75a137839d9ad64eb11810d5b6393" -dependencies = [ - "rand_core 0.6.4", - "subtle", -] - [[package]] name = "fiat-crypto" version = "0.2.9" @@ -6874,17 +6812,6 @@ dependencies = [ "system-deps", ] -[[package]] -name = "group" -version = "0.13.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f0f9ef7462f7c099f518d754361858f86d8a07af53ba9af0fe635bbccb151a63" -dependencies = [ - "ff", - "rand_core 0.6.4", - "subtle", -] - [[package]] name = "gzip-header" version = "1.0.0" @@ -9384,18 +9311,6 @@ dependencies = [ "supports-color 3.0.2", ] -[[package]] -name = "p256" -version = "0.13.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c9863ad85fa8f4460f9c48cb909d38a0d689dba1f6f6988a5e3e0d31071bcd4b" -dependencies = [ - "ecdsa", - "elliptic-curve", - "primeorder", - "sha2", -] - [[package]] name = "parking" version = "2.2.1" @@ -9825,15 +9740,6 @@ dependencies = [ "syn 2.0.114", ] -[[package]] -name = "primeorder" -version = "0.13.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "353e1ca18966c16d9deb1c69278edbc5f194139612772bd9537af60ac231e1e6" -dependencies = [ - "elliptic-curve", -] - [[package]] name = "proc-macro-crate" version = "3.4.0" @@ -10784,16 +10690,6 @@ version = "0.7.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1e061d1b48cb8d38042de4ae0a7a6401009d6143dc80d2e2d6f31f0bdd6470c7" -[[package]] -name = "rfc6979" -version = "0.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f8dd2a808d456c4a54e300a23e9f5a67e122c3024119acbfd73e3bf664491cb2" -dependencies = [ - "hmac", - "subtle", -] - [[package]] name = "ring" version = "0.17.14" @@ -11253,20 +11149,6 @@ version = "3.0.10" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "490dcfcbfef26be6800d11870ff2df8774fa6e86d047e3e8c8a76b25655e41ca" -[[package]] -name = "sec1" -version = "0.7.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d3e97a565f76233a6003f9f5c54be1d9c5bdfa3eccfb189469f11ec4901c47dc" -dependencies = [ - "base16ct", - "der", - "generic-array", - "pkcs8", - "subtle", - "zeroize", -] - [[package]] name = "seccompiler" version = "0.5.0" diff --git a/codex-rs/app-server/src/attestation.rs b/codex-rs/app-server/src/attestation.rs index b1fbf6945a2e..17bb10c38c76 100644 --- a/codex-rs/app-server/src/attestation.rs +++ b/codex-rs/app-server/src/attestation.rs @@ -45,10 +45,6 @@ impl AttestationProvider for AppServerAttestationProvider { let outgoing = self.outgoing.clone(); let thread_state_manager = self.thread_state_manager.clone(); Box::pin(async move { - if !context.uses_chatgpt_auth { - return None; - } - request_attestation_header_value_with_timeout( outgoing, thread_state_manager, diff --git a/codex-rs/codex-api/src/endpoint/memories.rs b/codex-rs/codex-api/src/endpoint/memories.rs index e9060ccb9e9a..a6c25641f25d 100644 --- a/codex-rs/codex-api/src/endpoint/memories.rs +++ b/codex-rs/codex-api/src/endpoint/memories.rs @@ -142,7 +142,6 @@ mod tests { Provider { name: "test".to_string(), base_url: base_url.to_string(), - uses_chatgpt_auth: false, query_params: None, headers: HeaderMap::new(), retry: RetryConfig { diff --git a/codex-rs/codex-api/src/endpoint/models.rs b/codex-rs/codex-api/src/endpoint/models.rs index c491f610b801..ec9ee7aac6d3 100644 --- a/codex-rs/codex-api/src/endpoint/models.rs +++ b/codex-rs/codex-api/src/endpoint/models.rs @@ -140,7 +140,6 @@ mod tests { Provider { name: "test".to_string(), base_url: base_url.to_string(), - uses_chatgpt_auth: false, query_params: None, headers: HeaderMap::new(), retry: RetryConfig { diff --git a/codex-rs/codex-api/src/endpoint/realtime_call.rs b/codex-rs/codex-api/src/endpoint/realtime_call.rs index ba5b88782481..b0342c53498d 100644 --- a/codex-rs/codex-api/src/endpoint/realtime_call.rs +++ b/codex-rs/codex-api/src/endpoint/realtime_call.rs @@ -297,7 +297,6 @@ mod tests { Provider { name: "test".to_string(), base_url: base_url.to_string(), - uses_chatgpt_auth: false, query_params: None, headers: HeaderMap::new(), retry: RetryConfig { diff --git a/codex-rs/codex-api/src/endpoint/realtime_websocket/methods.rs b/codex-rs/codex-api/src/endpoint/realtime_websocket/methods.rs index 41f9b92fca2c..9fcca1c3e318 100644 --- a/codex-rs/codex-api/src/endpoint/realtime_websocket/methods.rs +++ b/codex-rs/codex-api/src/endpoint/realtime_websocket/methods.rs @@ -1661,7 +1661,6 @@ mod tests { let provider = Provider { name: "test".to_string(), base_url: format!("http://{addr}"), - uses_chatgpt_auth: false, query_params: Some(HashMap::new()), headers: HeaderMap::new(), retry: crate::provider::RetryConfig { @@ -1956,7 +1955,6 @@ mod tests { let provider = Provider { name: "test".to_string(), base_url: format!("http://{addr}"), - uses_chatgpt_auth: false, query_params: Some(HashMap::new()), headers: HeaderMap::new(), retry: crate::provider::RetryConfig { @@ -2072,7 +2070,6 @@ mod tests { let provider = Provider { name: "test".to_string(), base_url: format!("http://{addr}"), - uses_chatgpt_auth: false, query_params: Some(HashMap::new()), headers: HeaderMap::new(), retry: crate::provider::RetryConfig { @@ -2177,7 +2174,6 @@ mod tests { let provider = Provider { name: "test".to_string(), base_url: format!("http://{addr}"), - uses_chatgpt_auth: false, query_params: Some(HashMap::new()), headers: HeaderMap::new(), retry: crate::provider::RetryConfig { @@ -2268,7 +2264,6 @@ mod tests { let provider = Provider { name: "test".to_string(), base_url: format!("http://{addr}"), - uses_chatgpt_auth: false, query_params: Some(HashMap::new()), headers: HeaderMap::new(), retry: crate::provider::RetryConfig { diff --git a/codex-rs/codex-api/src/provider.rs b/codex-rs/codex-api/src/provider.rs index 75a1062f65e9..45f2512dc39a 100644 --- a/codex-rs/codex-api/src/provider.rs +++ b/codex-rs/codex-api/src/provider.rs @@ -43,7 +43,6 @@ impl RetryConfig { pub struct Provider { pub name: String, pub base_url: String, - pub uses_chatgpt_auth: bool, pub query_params: Option>, pub headers: HeaderMap, pub retry: RetryConfig, diff --git a/codex-rs/codex-api/tests/clients.rs b/codex-rs/codex-api/tests/clients.rs index 6cdc9a09b9b8..a2a29ba16d37 100644 --- a/codex-rs/codex-api/tests/clients.rs +++ b/codex-rs/codex-api/tests/clients.rs @@ -127,7 +127,6 @@ fn provider(name: &str) -> Provider { Provider { name: name.to_string(), base_url: "https://example.com/v1".to_string(), - uses_chatgpt_auth: false, query_params: None, headers: HeaderMap::new(), retry: codex_api::RetryConfig { diff --git a/codex-rs/codex-api/tests/models_integration.rs b/codex-rs/codex-api/tests/models_integration.rs index e3653d82ec5f..d2b31180b907 100644 --- a/codex-rs/codex-api/tests/models_integration.rs +++ b/codex-rs/codex-api/tests/models_integration.rs @@ -32,7 +32,6 @@ fn provider(base_url: &str) -> Provider { Provider { name: "test".to_string(), base_url: base_url.to_string(), - uses_chatgpt_auth: false, query_params: None, headers: HeaderMap::new(), retry: RetryConfig { diff --git a/codex-rs/codex-api/tests/realtime_websocket_e2e.rs b/codex-rs/codex-api/tests/realtime_websocket_e2e.rs index f8c59fa32e26..cb9d7122f4b0 100644 --- a/codex-rs/codex-api/tests/realtime_websocket_e2e.rs +++ b/codex-rs/codex-api/tests/realtime_websocket_e2e.rs @@ -62,7 +62,6 @@ fn test_provider(base_url: String) -> Provider { Provider { name: "test".to_string(), base_url, - uses_chatgpt_auth: false, query_params: Some(HashMap::new()), headers: HeaderMap::new(), retry: RetryConfig { diff --git a/codex-rs/codex-api/tests/sse_end_to_end.rs b/codex-rs/codex-api/tests/sse_end_to_end.rs index bf31f87f77d3..bf880fefcf9f 100644 --- a/codex-rs/codex-api/tests/sse_end_to_end.rs +++ b/codex-rs/codex-api/tests/sse_end_to_end.rs @@ -61,7 +61,6 @@ fn provider(name: &str) -> Provider { Provider { name: name.to_string(), base_url: "https://example.com/v1".to_string(), - uses_chatgpt_auth: false, query_params: None, headers: HeaderMap::new(), retry: codex_api::RetryConfig { diff --git a/codex-rs/core/src/attestation.rs b/codex-rs/core/src/attestation.rs index c985f91b29c5..e2ec309cb733 100644 --- a/codex-rs/core/src/attestation.rs +++ b/codex-rs/core/src/attestation.rs @@ -13,7 +13,6 @@ pub type GenerateAttestationFuture<'a> = /// generate an attestation header value. #[derive(Clone, Copy, Debug)] pub struct AttestationContext { - pub uses_chatgpt_auth: bool, /// Thread whose upstream request is being prepared. pub thread_id: ThreadId, } diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 34cdcde8189f..d328e143de99 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -173,6 +173,7 @@ struct ModelClientState { enable_request_compression: bool, include_timing_metrics: bool, beta_features_header: Option, + include_attestation: bool, attestation_provider: Option>, disable_websockets: AtomicBool, cached_websocket_session: StdMutex, @@ -327,6 +328,7 @@ impl ModelClient { .is_some_and(|manager| manager.codex_api_key_env_enabled()); let auth_env_telemetry = collect_auth_env_telemetry(model_provider.info(), codex_api_key_env_enabled); + let include_attestation = model_provider.supports_attestation(); Self { state: Arc::new(ModelClientState { session_id, @@ -340,6 +342,7 @@ impl ModelClient { enable_request_compression, include_timing_metrics, beta_features_header, + include_attestation, attestation_provider, disable_websockets: AtomicBool::new(false), cached_websocket_session: StdMutex::new(WebsocketSession::default()), @@ -495,8 +498,7 @@ impl ModelClient { Some(self.state.session_id.to_string()), Some(self.state.thread_id.to_string()), )); - self.extend_attestation_header_for(&mut extra_headers, &client_setup.api_provider) - .await; + self.extend_attestation_header_for(&mut extra_headers).await; let client = ApiCompactClient::new(transport, client_setup.api_provider, client_setup.api_auth) .with_telemetry(Some(request_telemetry)); @@ -518,8 +520,7 @@ impl ModelClient { // 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?; - self.extend_attestation_header_for(&mut extra_headers, &client_setup.api_provider) - .await; + self.extend_attestation_header_for(&mut extra_headers).await; let mut sideband_headers = extra_headers.clone(); sideband_headers.extend(sideband_websocket_auth_headers( client_setup.api_auth.as_ref(), @@ -650,15 +651,15 @@ impl ModelClient { client_metadata } - async fn generate_attestation_header_for( - &self, - provider: &codex_api::Provider, - ) -> Option { + async fn generate_attestation_header_for(&self) -> Option { + if !self.state.include_attestation { + return None; + } + self.state .attestation_provider .as_ref()? .header_for_request(AttestationContext { - uses_chatgpt_auth: provider.uses_chatgpt_auth, thread_id: self.state.thread_id, }) .await @@ -882,7 +883,7 @@ impl ModelClient { /// replayed on reconnect within the same turn. async fn build_websocket_headers( &self, - provider: &codex_api::Provider, + _provider: &codex_api::Provider, turn_state: Option<&Arc>>, turn_metadata_header: Option<&str>, ) -> ApiHeaderMap { @@ -899,8 +900,7 @@ impl ModelClient { } headers.extend(build_session_headers(Some(session_id), Some(thread_id))); headers.extend(self.build_responses_identity_headers()); - self.extend_attestation_header_for(&mut headers, provider) - .await; + self.extend_attestation_header_for(&mut headers).await; headers.insert( OPENAI_BETA_HEADER, HeaderValue::from_static(RESPONSES_WEBSOCKETS_V2_BETA_HEADER_VALUE), @@ -951,7 +951,7 @@ impl ModelClientSession { /// regardless of transport choice. async fn build_responses_options( &self, - provider: &codex_api::Provider, + _provider: &codex_api::Provider, turn_metadata_header: Option<&str>, compression: Compression, ) -> ApiResponsesOptions { @@ -970,7 +970,7 @@ impl ModelClientSession { ); headers.extend(self.client.build_responses_identity_headers()); self.client - .extend_attestation_header_for(&mut headers, provider) + .extend_attestation_header_for(&mut headers) .await; headers }, @@ -1672,12 +1672,8 @@ fn build_responses_headers( } impl ModelClient { - async fn extend_attestation_header_for( - &self, - headers: &mut ApiHeaderMap, - provider: &codex_api::Provider, - ) { - if let Some(header_value) = self.generate_attestation_header_for(provider).await { + async fn extend_attestation_header_for(&self, headers: &mut ApiHeaderMap) { + if let Some(header_value) = self.generate_attestation_header_for().await { headers.insert(X_OAI_ATTESTATION_HEADER, header_value); } } diff --git a/codex-rs/core/src/client_tests.rs b/codex-rs/core/src/client_tests.rs index 55b0c227c660..4370bd009af1 100644 --- a/codex-rs/core/src/client_tests.rs +++ b/codex-rs/core/src/client_tests.rs @@ -13,8 +13,11 @@ use crate::GenerateAttestationFuture; use codex_api::ApiError; use codex_api::ResponseEvent; use codex_app_server_protocol::AuthMode; +use codex_login::AuthManager; +use codex_login::CodexAuth; use codex_model_provider::BearerAuthProvider; use codex_model_provider_info::CHATGPT_CODEX_BASE_URL; +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; @@ -79,9 +82,6 @@ fn api_provider(base_url: &str) -> codex_api::Provider { codex_api::Provider { name: "test".to_string(), base_url: base_url.to_string(), - uses_chatgpt_auth: base_url - .trim_end_matches('/') - .eq_ignore_ascii_case(CHATGPT_CODEX_BASE_URL), query_params: None, headers: http::HeaderMap::new(), retry: codex_api::RetryConfig { @@ -495,20 +495,21 @@ fn auth_request_telemetry_context_tracks_attached_auth_and_retry_phase() { assert_eq!(auth_context.recovery_phase, Some("refresh_token")); } -fn model_client_with_counting_attestation() -> (ModelClient, Arc) { +fn model_client_with_counting_attestation( + include_attestation: bool, +) -> (ModelClient, Arc) { #[derive(Debug)] struct CountingAttestationProvider { calls: Arc, } impl AttestationProvider for CountingAttestationProvider { - fn header_for_request(&self, context: AttestationContext) -> GenerateAttestationFuture<'_> { + fn header_for_request( + &self, + _context: AttestationContext, + ) -> GenerateAttestationFuture<'_> { let calls = self.calls.clone(); Box::pin(async move { - if !context.uses_chatgpt_auth { - return None; - } - let call = calls.fetch_add(1, Ordering::Relaxed) + 1; Some(http::HeaderValue::from_bytes(format!("v1.header-{call}").as_bytes()).unwrap()) }) @@ -516,12 +517,25 @@ fn model_client_with_counting_attestation() -> (ModelClient, Arc) { } let attestation_calls = Arc::new(AtomicUsize::new(0)); + let (auth_manager, provider) = if include_attestation { + ( + Some(AuthManager::from_auth_for_testing( + CodexAuth::create_dummy_chatgpt_auth_for_testing(), + )), + ModelProviderInfo::create_openai_provider(Some(CHATGPT_CODEX_BASE_URL.to_string())), + ) + } else { + ( + None, + create_oss_provider_with_base_url("https://example.com/v1", WireApi::Responses), + ) + }; let model_client = ModelClient::new( - /*auth_manager*/ None, + auth_manager, SessionId::new(), ThreadId::new(), /*installation_id*/ "11111111-1111-4111-8111-111111111111".to_string(), - create_oss_provider_with_base_url("https://example.com/v1", WireApi::Responses), + provider, SessionSource::Exec, /*model_verbosity*/ None, /*enable_request_compression*/ false, @@ -537,7 +551,7 @@ fn model_client_with_counting_attestation() -> (ModelClient, Arc) { #[tokio::test] async fn websocket_handshake_includes_attestation_for_chatgpt_codex_responses() { let provider = api_provider("https://chatgpt.com/backend-api/codex/"); - let (model_client, attestation_calls) = model_client_with_counting_attestation(); + let (model_client, attestation_calls) = model_client_with_counting_attestation(true); let headers = model_client .build_websocket_headers( @@ -556,20 +570,19 @@ async fn websocket_handshake_includes_attestation_for_chatgpt_codex_responses() #[tokio::test] async fn non_chatgpt_codex_endpoints_omit_attestation_generation() { - let provider = api_provider("https://api.openai.com/v1"); - let (model_client, attestation_calls) = model_client_with_counting_attestation(); + let (model_client, attestation_calls) = model_client_with_counting_attestation(false); let mut response_headers = http::HeaderMap::new(); model_client - .extend_attestation_header_for(&mut response_headers, &provider) + .extend_attestation_header_for(&mut response_headers) .await; let mut compaction_headers = http::HeaderMap::new(); model_client - .extend_attestation_header_for(&mut compaction_headers, &provider) + .extend_attestation_header_for(&mut compaction_headers) .await; let mut realtime_headers = http::HeaderMap::new(); model_client - .extend_attestation_header_for(&mut realtime_headers, &provider) + .extend_attestation_header_for(&mut realtime_headers) .await; assert_eq!( diff --git a/codex-rs/model-provider-info/src/lib.rs b/codex-rs/model-provider-info/src/lib.rs index d59a7503683e..6fca7e6a1f59 100644 --- a/codex-rs/model-provider-info/src/lib.rs +++ b/codex-rs/model-provider-info/src/lib.rs @@ -256,10 +256,6 @@ impl ModelProviderInfo { Ok(ApiProvider { name: self.name.clone(), base_url, - uses_chatgpt_auth: matches!( - auth_mode, - Some(AuthMode::Chatgpt | AuthMode::ChatgptAuthTokens) - ), query_params: self.query_params.clone(), headers, retry, diff --git a/codex-rs/model-provider/src/provider.rs b/codex-rs/model-provider/src/provider.rs index 0c5e8e0ffea7..8e1d37b29de4 100644 --- a/codex-rs/model-provider/src/provider.rs +++ b/codex-rs/model-provider/src/provider.rs @@ -85,6 +85,11 @@ pub trait ModelProvider: fmt::Debug + Send + Sync { ProviderCapabilities::default() } + /// Returns whether requests made through this provider should include attestation. + fn supports_attestation(&self) -> bool { + false + } + /// Returns the provider-scoped auth manager, when this provider uses one. /// /// TODO(celia-oai): Make auth manager access internal to this crate so callers @@ -167,6 +172,13 @@ impl ModelProvider for ConfiguredModelProvider { self.auth_manager.clone() } + fn supports_attestation(&self) -> bool { + self.auth_manager + .as_ref() + .and_then(|auth_manager| auth_manager.auth_cached()) + .is_some_and(|auth| auth.is_chatgpt_auth()) + } + async fn auth(&self) -> Option { match self.auth_manager.as_ref() { Some(auth_manager) => auth_manager.auth().await, From 10877f09fbc7793d64ce830674d8969f647eb08e Mon Sep 17 00:00:00 2001 From: Jiaming Zhang Date: Thu, 7 May 2026 18:04:26 -0700 Subject: [PATCH 19/23] codex: inline attestation header insertion --- codex-rs/core/src/client.rs | 26 ++++++++++++-------------- codex-rs/core/src/client_tests.rs | 18 +++++++++--------- 2 files changed, 21 insertions(+), 23 deletions(-) diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index d328e143de99..789515a35572 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -498,7 +498,9 @@ impl ModelClient { Some(self.state.session_id.to_string()), Some(self.state.thread_id.to_string()), )); - self.extend_attestation_header_for(&mut extra_headers).await; + if let Some(header_value) = self.generate_attestation_header_for().await { + extra_headers.insert(X_OAI_ATTESTATION_HEADER, header_value); + } let client = ApiCompactClient::new(transport, client_setup.api_provider, client_setup.api_auth) .with_telemetry(Some(request_telemetry)); @@ -520,7 +522,9 @@ impl ModelClient { // 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?; - self.extend_attestation_header_for(&mut extra_headers).await; + if let Some(header_value) = self.generate_attestation_header_for().await { + extra_headers.insert(X_OAI_ATTESTATION_HEADER, header_value); + } let mut sideband_headers = extra_headers.clone(); sideband_headers.extend(sideband_websocket_auth_headers( client_setup.api_auth.as_ref(), @@ -900,7 +904,9 @@ impl ModelClient { } headers.extend(build_session_headers(Some(session_id), Some(thread_id))); headers.extend(self.build_responses_identity_headers()); - self.extend_attestation_header_for(&mut headers).await; + if let Some(header_value) = self.generate_attestation_header_for().await { + headers.insert(X_OAI_ATTESTATION_HEADER, header_value); + } headers.insert( OPENAI_BETA_HEADER, HeaderValue::from_static(RESPONSES_WEBSOCKETS_V2_BETA_HEADER_VALUE), @@ -969,9 +975,9 @@ impl ModelClientSession { turn_metadata_header.as_ref(), ); headers.extend(self.client.build_responses_identity_headers()); - self.client - .extend_attestation_header_for(&mut headers) - .await; + if let Some(header_value) = self.client.generate_attestation_header_for().await { + headers.insert(X_OAI_ATTESTATION_HEADER, header_value); + } headers }, compression, @@ -1671,14 +1677,6 @@ fn build_responses_headers( headers } -impl ModelClient { - async fn extend_attestation_header_for(&self, headers: &mut ApiHeaderMap) { - if let Some(header_value) = self.generate_attestation_header_for().await { - headers.insert(X_OAI_ATTESTATION_HEADER, header_value); - } - } -} - fn subagent_header_value(session_source: &SessionSource) -> Option { match session_source { SessionSource::SubAgent(subagent_source) => match subagent_source { diff --git a/codex-rs/core/src/client_tests.rs b/codex-rs/core/src/client_tests.rs index 4370bd009af1..901eba32bb67 100644 --- a/codex-rs/core/src/client_tests.rs +++ b/codex-rs/core/src/client_tests.rs @@ -573,17 +573,17 @@ async fn non_chatgpt_codex_endpoints_omit_attestation_generation() { let (model_client, attestation_calls) = model_client_with_counting_attestation(false); let mut response_headers = http::HeaderMap::new(); - model_client - .extend_attestation_header_for(&mut response_headers) - .await; + if let Some(header_value) = model_client.generate_attestation_header_for().await { + response_headers.insert(crate::attestation::X_OAI_ATTESTATION_HEADER, header_value); + } let mut compaction_headers = http::HeaderMap::new(); - model_client - .extend_attestation_header_for(&mut compaction_headers) - .await; + if let Some(header_value) = model_client.generate_attestation_header_for().await { + compaction_headers.insert(crate::attestation::X_OAI_ATTESTATION_HEADER, header_value); + } let mut realtime_headers = http::HeaderMap::new(); - model_client - .extend_attestation_header_for(&mut realtime_headers) - .await; + if let Some(header_value) = model_client.generate_attestation_header_for().await { + realtime_headers.insert(crate::attestation::X_OAI_ATTESTATION_HEADER, header_value); + } assert_eq!( response_headers.get(crate::attestation::X_OAI_ATTESTATION_HEADER), From ed193932bdd8f06fb259ca871294562ef41d6e8a Mon Sep 17 00:00:00 2001 From: Jiaming Zhang Date: Thu, 7 May 2026 18:33:36 -0700 Subject: [PATCH 20/23] codex: fix CI failure on PR #20619 --- codex-rs/core/src/client_tests.rs | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/codex-rs/core/src/client_tests.rs b/codex-rs/core/src/client_tests.rs index 901eba32bb67..17aef00136d9 100644 --- a/codex-rs/core/src/client_tests.rs +++ b/codex-rs/core/src/client_tests.rs @@ -551,7 +551,8 @@ fn model_client_with_counting_attestation( #[tokio::test] async fn websocket_handshake_includes_attestation_for_chatgpt_codex_responses() { let provider = api_provider("https://chatgpt.com/backend-api/codex/"); - let (model_client, attestation_calls) = model_client_with_counting_attestation(true); + let (model_client, attestation_calls) = + model_client_with_counting_attestation(/*include_attestation*/ true); let headers = model_client .build_websocket_headers( @@ -570,7 +571,8 @@ async fn websocket_handshake_includes_attestation_for_chatgpt_codex_responses() #[tokio::test] async fn non_chatgpt_codex_endpoints_omit_attestation_generation() { - let (model_client, attestation_calls) = model_client_with_counting_attestation(false); + let (model_client, attestation_calls) = + model_client_with_counting_attestation(/*include_attestation*/ false); let mut response_headers = http::HeaderMap::new(); if let Some(header_value) = model_client.generate_attestation_header_for().await { From bd1da5f3ca2b31de7635ccd07cc57f7753a03431 Mon Sep 17 00:00:00 2001 From: Jiaming Zhang Date: Fri, 8 May 2026 09:52:46 -0700 Subject: [PATCH 21/23] codex: remove unused provider parameters --- codex-rs/core/src/client.rs | 16 +++------------- codex-rs/core/src/client_tests.rs | 5 +---- 2 files changed, 4 insertions(+), 17 deletions(-) diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 789515a35572..94f88ce01096 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -807,7 +807,7 @@ impl ModelClient { request_route_telemetry: RequestRouteTelemetry, ) -> std::result::Result { let headers = self - .build_websocket_headers(&api_provider, turn_state.as_ref(), turn_metadata_header) + .build_websocket_headers(turn_state.as_ref(), turn_metadata_header) .await; let websocket_telemetry = ModelClientSession::build_websocket_telemetry( session_telemetry, @@ -887,7 +887,6 @@ impl ModelClient { /// replayed on reconnect within the same turn. async fn build_websocket_headers( &self, - _provider: &codex_api::Provider, turn_state: Option<&Arc>>, turn_metadata_header: Option<&str>, ) -> ApiHeaderMap { @@ -957,7 +956,6 @@ impl ModelClientSession { /// regardless of transport choice. async fn build_responses_options( &self, - _provider: &codex_api::Provider, turn_metadata_header: Option<&str>, compression: Compression, ) -> ApiResponsesOptions { @@ -1255,11 +1253,7 @@ impl ModelClientSession { ); let compression = self.responses_request_compression(client_setup.auth.as_ref()); let options = self - .build_responses_options( - &client_setup.api_provider, - turn_metadata_header, - compression, - ) + .build_responses_options(turn_metadata_header, compression) .await; let request = self.client.build_responses_request( @@ -1368,11 +1362,7 @@ impl ModelClientSession { let compression = self.responses_request_compression(client_setup.auth.as_ref()); let options = self - .build_responses_options( - &client_setup.api_provider, - turn_metadata_header, - compression, - ) + .build_responses_options(turn_metadata_header, compression) .await; let request = self.client.build_responses_request( &client_setup.api_provider, diff --git a/codex-rs/core/src/client_tests.rs b/codex-rs/core/src/client_tests.rs index 17aef00136d9..a2d08e01d928 100644 --- a/codex-rs/core/src/client_tests.rs +++ b/codex-rs/core/src/client_tests.rs @@ -550,14 +550,11 @@ fn model_client_with_counting_attestation( #[tokio::test] async fn websocket_handshake_includes_attestation_for_chatgpt_codex_responses() { - let provider = api_provider("https://chatgpt.com/backend-api/codex/"); let (model_client, attestation_calls) = model_client_with_counting_attestation(/*include_attestation*/ true); let headers = model_client - .build_websocket_headers( - &provider, /*turn_state*/ None, /*turn_metadata_header*/ None, - ) + .build_websocket_headers(/*turn_state*/ None, /*turn_metadata_header*/ None) .await; assert_eq!( From f3cda57d7baa9a8d7a6625dae4fd413c1cc05fbe Mon Sep 17 00:00:00 2001 From: Jiaming Zhang Date: Fri, 8 May 2026 11:04:27 -0700 Subject: [PATCH 22/23] codex: fix CI failure on PR #20619 --- codex-rs/Cargo.lock | 118 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 118 insertions(+) diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index e544fdb5f3bc..c656a791a60c 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -1180,6 +1180,12 @@ dependencies = [ "windows-link", ] +[[package]] +name = "base16ct" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4c7f02d4ea65f2c1853089ffd8d2787bdbc63de2f0d29dedbcf8ccdfa0ccd4cf" + [[package]] name = "base64" version = "0.21.7" @@ -4419,6 +4425,18 @@ version = "0.2.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" +[[package]] +name = "crypto-bigint" +version = "0.5.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0dc92fb57ca44df6db8059111ab3af99a63d5d0f8375d9972e319a379c6bab76" +dependencies = [ + "generic-array", + "rand_core 0.6.4", + "subtle", + "zeroize", +] + [[package]] name = "crypto-common" version = "0.1.7" @@ -5132,6 +5150,20 @@ version = "1.0.20" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d0881ea181b1df73ff77ffaaf9c7544ecc11e82fba9b5f27b262a3c73a332555" +[[package]] +name = "ecdsa" +version = "0.16.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ee27f32b5c5292967d2d4a9d7f1e0b0aed2c15daded5a60300e4abb9d8020bca" +dependencies = [ + "der", + "digest", + "elliptic-curve", + "rfc6979", + "signature", + "spki", +] + [[package]] name = "ed25519" version = "2.2.3" @@ -5165,6 +5197,26 @@ dependencies = [ "serde", ] +[[package]] +name = "elliptic-curve" +version = "0.13.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b5e6043086bf7973472e0c7dff2142ea0b680d30e18d9cc40f267efbf222bd47" +dependencies = [ + "base16ct", + "crypto-bigint", + "digest", + "ff", + "generic-array", + "group", + "pem-rfc7468", + "pkcs8", + "rand_core 0.6.4", + "sec1", + "subtle", + "zeroize", +] + [[package]] name = "ena" version = "0.14.3" @@ -5418,6 +5470,16 @@ dependencies = [ "simd-adler32", ] +[[package]] +name = "ff" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c0b50bfb653653f9ca9095b427bed08ab8d75a137839d9ad64eb11810d5b6393" +dependencies = [ + "rand_core 0.6.4", + "subtle", +] + [[package]] name = "fiat-crypto" version = "0.2.9" @@ -6812,6 +6874,17 @@ dependencies = [ "system-deps", ] +[[package]] +name = "group" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0f9ef7462f7c099f518d754361858f86d8a07af53ba9af0fe635bbccb151a63" +dependencies = [ + "ff", + "rand_core 0.6.4", + "subtle", +] + [[package]] name = "gzip-header" version = "1.0.0" @@ -9311,6 +9384,18 @@ dependencies = [ "supports-color 3.0.2", ] +[[package]] +name = "p256" +version = "0.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c9863ad85fa8f4460f9c48cb909d38a0d689dba1f6f6988a5e3e0d31071bcd4b" +dependencies = [ + "ecdsa", + "elliptic-curve", + "primeorder", + "sha2", +] + [[package]] name = "parking" version = "2.2.1" @@ -9740,6 +9825,15 @@ dependencies = [ "syn 2.0.114", ] +[[package]] +name = "primeorder" +version = "0.13.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "353e1ca18966c16d9deb1c69278edbc5f194139612772bd9537af60ac231e1e6" +dependencies = [ + "elliptic-curve", +] + [[package]] name = "proc-macro-crate" version = "3.4.0" @@ -10690,6 +10784,16 @@ version = "0.7.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1e061d1b48cb8d38042de4ae0a7a6401009d6143dc80d2e2d6f31f0bdd6470c7" +[[package]] +name = "rfc6979" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8dd2a808d456c4a54e300a23e9f5a67e122c3024119acbfd73e3bf664491cb2" +dependencies = [ + "hmac", + "subtle", +] + [[package]] name = "ring" version = "0.17.14" @@ -11149,6 +11253,20 @@ version = "3.0.10" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "490dcfcbfef26be6800d11870ff2df8774fa6e86d047e3e8c8a76b25655e41ca" +[[package]] +name = "sec1" +version = "0.7.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3e97a565f76233a6003f9f5c54be1d9c5bdfa3eccfb189469f11ec4901c47dc" +dependencies = [ + "base16ct", + "der", + "generic-array", + "pkcs8", + "subtle", + "zeroize", +] + [[package]] name = "seccompiler" version = "0.5.0" From 34bf9f5c97444aa1b2dd00f5e67cc05a9f27212b Mon Sep 17 00:00:00 2001 From: Jiaming Zhang Date: Fri, 8 May 2026 11:14:04 -0700 Subject: [PATCH 23/23] codex: fix CI failure on PR #20619 --- codex-rs/core/src/client_tests.rs | 17 ----------------- 1 file changed, 17 deletions(-) diff --git a/codex-rs/core/src/client_tests.rs b/codex-rs/core/src/client_tests.rs index a2d08e01d928..b9d9172c8398 100644 --- a/codex-rs/core/src/client_tests.rs +++ b/codex-rs/core/src/client_tests.rs @@ -78,23 +78,6 @@ fn test_model_client(session_source: SessionSource) -> ModelClient { ) } -fn api_provider(base_url: &str) -> codex_api::Provider { - codex_api::Provider { - name: "test".to_string(), - base_url: base_url.to_string(), - query_params: None, - headers: http::HeaderMap::new(), - retry: codex_api::RetryConfig { - max_attempts: 1, - base_delay: Duration::from_millis(1), - retry_429: false, - retry_5xx: true, - retry_transport: true, - }, - stream_idle_timeout: Duration::from_secs(1), - } -} - fn test_model_info() -> ModelInfo { serde_json::from_value(json!({ "slug": "gpt-test",