mirror of
https://github.com/openai/codex.git
synced 2026-09-05 15:18:41 +00:00
[codex-app-server-transport] cover remote control HTTP state routes [ci changed_files]
This commit is contained in:
1
codex-rs/Cargo.lock
generated
1
codex-rs/Cargo.lock
generated
@@ -2094,6 +2094,7 @@ dependencies = [
|
||||
"codex-app-server-protocol",
|
||||
"codex-config",
|
||||
"codex-core",
|
||||
"codex-http-state",
|
||||
"codex-login",
|
||||
"codex-model-provider",
|
||||
"codex-state",
|
||||
|
||||
@@ -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 }
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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,
|
||||
};
|
||||
|
||||
@@ -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");
|
||||
|
||||
@@ -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))
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user