codex: wrap app-server attestation transport

This commit is contained in:
Jiaming Zhang
2026-05-07 08:58:12 -07:00
parent 5fdd555edc
commit 8ee7bf6abc
8 changed files with 108 additions and 11 deletions

View File

@@ -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"
}
},

View File

@@ -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"
}
},

View File

@@ -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, };

View File

@@ -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,
}

View File

@@ -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.<opaque>" }`, 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.<opaque>" }`, 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.<opaque>" }`, 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

View File

@@ -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<OutgoingMessageSender>,
@@ -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::<AttestationGenerateResponse>(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;

View File

@@ -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}"#
);
}

View File

@@ -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;