[codex-app-server-transport] cover remote control HTTP state routes [ci changed_files]

This commit is contained in:
Cooper Gamble
2026-06-03 04:48:21 +00:00
parent 4dd16dc182
commit 7236cfefa6
6 changed files with 130 additions and 9 deletions

1
codex-rs/Cargo.lock generated
View File

@@ -2094,6 +2094,7 @@ dependencies = [
"codex-app-server-protocol",
"codex-config",
"codex-core",
"codex-http-state",
"codex-login",
"codex-model-provider",
"codex-state",

View File

@@ -25,6 +25,7 @@ clap = { workspace = true, features = ["derive"] }
codex-api = { workspace = true }
codex-app-server-protocol = { workspace = true }
codex-core = { workspace = true }
codex-http-state = { workspace = true }
codex-login = { workspace = true }
codex-model-provider = { workspace = true }
codex-state = { workspace = true }

View File

@@ -6,6 +6,8 @@ use super::protocol::RemoteControlTarget;
use super::protocol::StartRemoteControlPairingRequest;
use super::protocol::StartRemoteControlPairingResponse;
use axum::http::HeaderMap;
use axum::http::HeaderValue;
use axum::http::header::AUTHORIZATION;
use codex_api::SharedAuthProvider;
use codex_app_server_protocol::RemoteControlPairingStartResponse;
use codex_login::default_client::build_reqwest_client;
@@ -45,6 +47,7 @@ pub(super) struct RemoteControlEnrollment {
impl RemoteControlEnrollment {
pub(super) async fn start_pairing(
&self,
auth: &RemoteControlConnectionAuth,
request: StartRemoteControlPairingRequest,
) -> io::Result<RemoteControlPairingStartResponse> {
if self.should_refresh_server_token() {
@@ -55,10 +58,22 @@ impl RemoteControlEnrollment {
.as_deref()
.ok_or_else(pairing_unavailable_error)?;
let mut auth_headers = HeaderMap::new();
auth.server_token_auth_provider
.add_auth_headers_for_url(&self.remote_control_target.pair_url, &mut auth_headers);
auth_headers.insert(
AUTHORIZATION,
HeaderValue::from_str(&format!("Bearer {remote_control_token}")).map_err(|err| {
io::Error::new(
ErrorKind::InvalidInput,
format!("invalid remote control pairing token: {err}"),
)
})?,
);
let response = build_reqwest_client()
.post(&self.remote_control_target.pair_url)
.timeout(REMOTE_CONTROL_PAIRING_TIMEOUT)
.bearer_auth(remote_control_token)
.headers(auth_headers.clone())
.json(&request)
.send()
.await
@@ -69,6 +84,11 @@ impl RemoteControlEnrollment {
))
})?;
let headers = response.headers().clone();
auth.server_token_auth_provider.observe_response_headers(
&self.remote_control_target.pair_url,
&auth_headers,
&headers,
);
let status = response.status();
let body = response.bytes().await.map_err(|err| {
io::Error::other(format!(
@@ -153,6 +173,7 @@ impl RemoteControlEnrollment {
pub(super) struct RemoteControlConnectionAuth {
pub(super) auth_provider: SharedAuthProvider,
pub(super) server_token_auth_provider: SharedAuthProvider,
pub(super) account_id: String,
}
@@ -430,11 +451,12 @@ where
{
let client = build_reqwest_client();
let mut auth_headers = HeaderMap::new();
auth.auth_provider.add_auth_headers(&mut auth_headers);
auth.auth_provider
.add_auth_headers_for_url(url, &mut auth_headers);
let response = client
.post(url)
.timeout(REMOTE_CONTROL_ENROLL_TIMEOUT)
.headers(auth_headers)
.headers(auth_headers.clone())
.header(REMOTE_CONTROL_ACCOUNT_ID_HEADER, &auth.account_id)
.header(REMOTE_CONTROL_INSTALLATION_ID_HEADER, installation_id)
.json(request)
@@ -446,6 +468,8 @@ where
))
})?;
let headers = response.headers().clone();
auth.auth_provider
.observe_response_headers(url, &auth_headers, &headers);
let status = response.status();
let body = response.bytes().await.map_err(|err| {
io::Error::other(format!(
@@ -745,6 +769,7 @@ mod tests {
&remote_control_target,
&RemoteControlConnectionAuth {
auth_provider: codex_model_provider::unauthenticated_auth_provider(),
server_token_auth_provider: codex_model_provider::unauthenticated_auth_provider(),
account_id: "account_id".to_string(),
},
"11111111-1111-4111-8111-111111111111",

View File

@@ -174,7 +174,7 @@ impl RemoteControlHandle {
let pairing_request = || protocol::StartRemoteControlPairingRequest {
manual_code: params.manual_code,
};
let pairing_response = match enrollment.start_pairing(pairing_request()).await {
let pairing_response = match enrollment.start_pairing(&auth, pairing_request()).await {
Err(err) if err.kind() == io::ErrorKind::PermissionDenied => {
clear_pairing_server_token(&self.current_enrollment, &mut enrollment)?;
refresh_pairing_enrollment(
@@ -185,7 +185,7 @@ impl RemoteControlHandle {
&mut enrollment,
)
.await?;
enrollment.start_pairing(pairing_request()).await
enrollment.start_pairing(&auth, pairing_request()).await
}
pairing_response => pairing_response,
};

View File

@@ -1,3 +1,4 @@
use super::super::enroll::RemoteControlConnectionAuth;
use super::super::protocol::StartRemoteControlPairingRequest;
use super::*;
use pretty_assertions::assert_eq;
@@ -21,6 +22,14 @@ fn remote_control_enrollment(
}
}
fn remote_control_connection_auth() -> RemoteControlConnectionAuth {
RemoteControlConnectionAuth {
auth_provider: codex_model_provider::unauthenticated_auth_provider(),
server_token_auth_provider: codex_model_provider::unauthenticated_auth_provider(),
account_id: "account-id".to_string(),
}
}
async fn pairing_error(status: &'static str, body: &'static str) -> (String, String) {
let listener = TcpListener::bind("127.0.0.1:0")
.await
@@ -41,7 +50,10 @@ async fn pairing_error(status: &'static str, body: &'static str) -> (String, Str
});
let err = remote_control_enrollment(&remote_control_url, "remote-control-token")
.start_pairing(StartRemoteControlPairingRequest { manual_code: false })
.start_pairing(
&remote_control_connection_auth(),
StartRemoteControlPairingRequest { manual_code: false },
)
.await
.expect_err("pairing should fail");
server_task.await.expect("server task should finish");
@@ -59,7 +71,10 @@ async fn pairing_response_error(body: serde_json::Value) -> String {
});
let err = remote_control_enrollment(&remote_control_url, "remote-control-token")
.start_pairing(StartRemoteControlPairingRequest { manual_code: false })
.start_pairing(
&remote_control_connection_auth(),
StartRemoteControlPairingRequest { manual_code: false },
)
.await
.expect_err("pairing should fail");
server_task.await.expect("server task should finish");

View File

@@ -28,6 +28,8 @@ use base64::Engine;
use codex_app_server_protocol::RemoteControlConnectionStatus;
use codex_app_server_protocol::RemoteControlStatusChangedNotification;
use codex_core::util::backoff;
use codex_http_state::HttpStateContext;
use codex_http_state::HttpStateSurface;
use codex_login::AuthManager;
use codex_login::UnauthorizedRecovery;
use codex_state::StateRuntime;
@@ -1124,6 +1126,7 @@ fn set_remote_control_header(
fn build_remote_control_websocket_request(
websocket_url: &str,
auth_provider: &codex_api::SharedAuthProvider,
enrollment: &RemoteControlEnrollment,
installation_id: &str,
subscribe_cursor: Option<&str>,
@@ -1135,6 +1138,7 @@ fn build_remote_control_websocket_request(
)
})?;
let headers = request.headers_mut();
auth_provider.add_auth_headers_for_url(websocket_url, headers);
set_remote_control_header(headers, "x-codex-server-id", &enrollment.server_id)?;
set_remote_control_header(
headers,
@@ -1206,8 +1210,21 @@ pub(crate) async fn load_remote_control_auth(
));
}
let http_state = HttpStateContext::new(
auth_manager.codex_home().to_path_buf(),
HttpStateSurface::CodexRemoteControl,
);
Ok(RemoteControlConnectionAuth {
auth_provider: codex_model_provider::auth_provider_from_auth(&auth),
auth_provider: codex_model_provider::with_native_integrity_state(
codex_model_provider::auth_provider_from_auth(&auth),
Some(&auth),
Some(http_state.clone()),
),
server_token_auth_provider: codex_model_provider::with_native_integrity_state(
codex_model_provider::unauthenticated_auth_provider(),
Some(&auth),
Some(http_state),
),
account_id: auth.get_account_id().ok_or_else(|| {
io::Error::new(
ErrorKind::WouldBlock,
@@ -1386,10 +1403,12 @@ pub(super) async fn connect_remote_control_websocket(
publish_current_enrollment(current_enrollment, enrollment_ref);
let request = build_remote_control_websocket_request(
&remote_control_target.websocket_url,
&auth.server_token_auth_provider,
enrollment_ref,
connect_options.installation_id,
connect_options.subscribe_cursor,
)?;
let request_headers = request.headers().clone();
let websocket_connect_result = tokio::time::timeout(
REMOTE_CONTROL_WEBSOCKET_CONNECT_TIMEOUT,
@@ -1407,8 +1426,22 @@ pub(super) async fn connect_remote_control_websocket(
})?;
match websocket_connect_result {
Ok((websocket_stream, response)) => Ok((websocket_stream, response.map(|_| ()))),
Ok((websocket_stream, response)) => {
auth.server_token_auth_provider.observe_response_headers(
&remote_control_target.websocket_url,
&request_headers,
response.headers(),
);
Ok((websocket_stream, response.map(|_| ())))
}
Err(err) => {
if let tungstenite::Error::Http(response) = &err {
auth.server_token_auth_provider.observe_response_headers(
&remote_control_target.websocket_url,
&request_headers,
response.headers(),
);
}
match &err {
tungstenite::Error::Http(response) if response.status().as_u16() == 404 => {
info!(
@@ -1630,6 +1663,7 @@ mod tests {
use codex_app_server_protocol::ServerNotification;
use codex_config::types::AuthCredentialsStoreMode;
use codex_core::test_support::auth_manager_from_auth;
use codex_http_state::HttpStateStore;
use codex_login::AuthDotJson;
use codex_login::CodexAuth;
use codex_login::save_auth;
@@ -1673,6 +1707,51 @@ mod tests {
}
}
#[test]
fn websocket_request_attaches_integrity_state_without_chatgpt_auth_headers() {
let codex_home = TempDir::new().expect("temp dir should create");
HttpStateStore::new(codex_home.path().to_path_buf())
.set(
HttpStateSurface::CodexRemoteControl,
"integrity-state".to_string(),
)
.expect("state should persist");
let auth = CodexAuth::create_dummy_chatgpt_auth_for_testing();
let auth_provider = codex_model_provider::with_native_integrity_state(
codex_model_provider::unauthenticated_auth_provider(),
Some(&auth),
Some(HttpStateContext::new(
codex_home.path().to_path_buf(),
HttpStateSurface::CodexRemoteControl,
)),
);
let request = build_remote_control_websocket_request(
"wss://chatgpt.com/backend-api/wham/remote/control/server",
&auth_provider,
&remote_control_enrollment(Some(TEST_REMOTE_CONTROL_SERVER_TOKEN)),
TEST_INSTALLATION_ID,
None,
)
.expect("request should build");
assert_eq!(
request
.headers()
.get("x-oai-is")
.and_then(|value| value.to_str().ok()),
Some("integrity-state")
);
assert_eq!(
request
.headers()
.get("authorization")
.and_then(|value| value.to_str().ok()),
Some("Bearer Remote Control Token")
);
assert!(!request.headers().contains_key("chatgpt-account-id"));
}
fn test_current_enrollment() -> CurrentRemoteControlEnrollment {
Arc::new(std::sync::Mutex::new(None))
}