app-server: implement device key v2 methods (#18430)

## Why

The device-key protocol needs an app-server implementation that keeps
local key operations behind the same request-processing boundary as
other v2 APIs.

app-server owns request dispatch, transport policy, documentation, and
JSON-RPC error shaping. `codex-device-key` owns key binding, validation,
platform provider selection, and signing mechanics. Keeping the adapter
thin makes the boundary easier to review and avoids moving local
key-management details into thread orchestration code.

## What changed

- Added `DeviceKeyApi` as the app-server adapter around
`DeviceKeyStore`.
- Converted protocol protection policies, payload variants, algorithms,
and protection classes to and from the device-key crate types.
- Encoded SPKI public keys and DER signatures as base64 protocol fields.
- Routed `device/key/create`, `device/key/public`, and `device/key/sign`
through `MessageProcessor`.
- Rejected remote transports before provider access while allowing local
`stdio` and in-process callers to reach the device-key API.
- Added stdio, in-process, and websocket tests for device-key validation
and transport policy.
- Documented the device-key methods in the app-server v2 method list.

## Test coverage

- `device_key_create_rejects_empty_account_user_id`
- `in_process_allows_device_key_requests_to_reach_device_key_api`
- `device_key_methods_are_rejected_over_websocket`

## Stack

This is PR 3 of 4 in the device-key app-server stack. It is stacked on
#18429.

## Validation

- `cargo test -p codex-app-server device_key`
- `just fix -p codex-app-server`
This commit is contained in:
Ruslan Nigmatullin
2026-04-21 14:07:08 -07:00
committed by GitHub
Unverified
parent e502f0b52d
commit 69c3d12274
17 changed files with 734 additions and 6 deletions
+1
View File
@@ -1472,6 +1472,7 @@ dependencies = [
"codex-config",
"codex-core",
"codex-core-plugins",
"codex-device-key",
"codex-exec-server",
"codex-features",
"codex-feedback",
+1
View File
@@ -133,6 +133,7 @@ codex-connectors = { path = "connectors" }
codex-core = { path = "core" }
codex-core-plugins = { path = "core-plugins" }
codex-core-skills = { path = "core-skills" }
codex-device-key = { path = "device-key" }
codex-exec = { path = "exec" }
codex-exec-server = { path = "exec-server" }
codex-execpolicy = { path = "execpolicy" }
+1
View File
@@ -35,6 +35,7 @@ codex-cloud-requirements = { workspace = true }
codex-config = { workspace = true }
codex-core = { workspace = true }
codex-core-plugins = { workspace = true }
codex-device-key = { workspace = true }
codex-exec-server = { workspace = true }
codex-features = { workspace = true }
codex-git-utils = { workspace = true }
+3
View File
@@ -190,6 +190,9 @@ Example with notification opt-out:
- `plugin/read` — read one plugin by `marketplacePath` plus `pluginName`, returning marketplace info, a list-style `summary`, manifest descriptions/interface metadata, and bundled skills/apps/MCP server names. Returned plugin skills include their current `enabled` state after local config filtering. Plugin app summaries also include `needsAuth` when the server can determine connector accessibility (**under development; do not call from production clients yet**).
- `skills/changed` — notification emitted when watched local skill files change.
- `app/list` — list available apps.
- `device/key/create` — create or load a controller-local device signing key for an account/client binding. This local-key API is available only over local transports such as stdio and in-process; remote transports reject it. Hardware-backed providers are the target protection class; an OS-protected non-extractable fallback is allowed only with `protectionPolicy: "allow_os_protected_nonextractable"` and returns the reported `protectionClass`.
- `device/key/public` — return a device key's SPKI DER public key as base64 plus its `algorithm` and `protectionClass`.
- `device/key/sign` — sign one of the accepted structured payload variants with a controller-local device key. The only accepted payload today is `remoteControlClientConnection`, which binds a server-issued `/client` websocket challenge to the enrolled controller device without signing the bearer token itself; this is intentionally not an arbitrary-byte signing API.
- `skills/config/write` — write user-level skill config by name or absolute path.
- `plugin/install` — install a plugin from a discovered marketplace entry, rejecting marketplace entries marked unavailable for install, install MCPs if any, and return the effective plugin auth policy plus any apps that still need auth (**under development; do not call from production clients yet**).
- `plugin/uninstall` — uninstall a plugin by id by removing its cached files and clearing its user-level config entry (**under development; do not call from production clients yet**).
+226
View File
@@ -0,0 +1,226 @@
use crate::error_code::INTERNAL_ERROR_CODE;
use crate::error_code::INVALID_REQUEST_ERROR_CODE;
use base64::Engine;
use base64::engine::general_purpose::STANDARD;
use codex_app_server_protocol::DeviceKeyAlgorithm;
use codex_app_server_protocol::DeviceKeyCreateParams;
use codex_app_server_protocol::DeviceKeyCreateResponse;
use codex_app_server_protocol::DeviceKeyProtectionClass;
use codex_app_server_protocol::DeviceKeyPublicParams;
use codex_app_server_protocol::DeviceKeyPublicResponse;
use codex_app_server_protocol::DeviceKeySignParams;
use codex_app_server_protocol::DeviceKeySignPayload;
use codex_app_server_protocol::DeviceKeySignResponse;
use codex_app_server_protocol::JSONRPCErrorError;
use codex_device_key::DeviceKeyBinding;
use codex_device_key::DeviceKeyCreateRequest;
use codex_device_key::DeviceKeyError;
use codex_device_key::DeviceKeyGetPublicRequest;
use codex_device_key::DeviceKeyInfo;
use codex_device_key::DeviceKeyProtectionPolicy;
use codex_device_key::DeviceKeySignRequest;
use codex_device_key::DeviceKeyStore;
use codex_device_key::RemoteControlClientConnectionAudience;
use codex_device_key::RemoteControlClientConnectionSignPayload;
use codex_device_key::RemoteControlClientEnrollmentAudience;
use codex_device_key::RemoteControlClientEnrollmentSignPayload;
#[derive(Clone, Default)]
pub(crate) struct DeviceKeyApi {
store: DeviceKeyStore,
}
impl DeviceKeyApi {
pub(crate) fn create(
&self,
params: DeviceKeyCreateParams,
) -> Result<DeviceKeyCreateResponse, JSONRPCErrorError> {
let info = self
.store
.create(DeviceKeyCreateRequest {
protection_policy: protection_policy_from_params(params.protection_policy),
binding: DeviceKeyBinding {
account_user_id: params.account_user_id,
client_id: params.client_id,
},
})
.map_err(map_device_key_error)?;
Ok(create_response_from_info(info))
}
pub(crate) fn public(
&self,
params: DeviceKeyPublicParams,
) -> Result<DeviceKeyPublicResponse, JSONRPCErrorError> {
let info = self
.store
.get_public(DeviceKeyGetPublicRequest {
key_id: params.key_id,
})
.map_err(map_device_key_error)?;
Ok(public_response_from_info(info))
}
pub(crate) fn sign(
&self,
params: DeviceKeySignParams,
) -> Result<DeviceKeySignResponse, JSONRPCErrorError> {
let signature = self
.store
.sign(DeviceKeySignRequest {
key_id: params.key_id,
payload: payload_from_params(params.payload),
})
.map_err(map_device_key_error)?;
Ok(DeviceKeySignResponse {
signature_der_base64: STANDARD.encode(signature.signature_der),
signed_payload_base64: STANDARD.encode(signature.signed_payload),
algorithm: algorithm_from_store(signature.algorithm),
})
}
}
fn create_response_from_info(info: DeviceKeyInfo) -> DeviceKeyCreateResponse {
DeviceKeyCreateResponse {
key_id: info.key_id,
public_key_spki_der_base64: STANDARD.encode(info.public_key_spki_der),
algorithm: algorithm_from_store(info.algorithm),
protection_class: protection_class_from_store(info.protection_class),
}
}
fn public_response_from_info(info: DeviceKeyInfo) -> DeviceKeyPublicResponse {
DeviceKeyPublicResponse {
key_id: info.key_id,
public_key_spki_der_base64: STANDARD.encode(info.public_key_spki_der),
algorithm: algorithm_from_store(info.algorithm),
protection_class: protection_class_from_store(info.protection_class),
}
}
fn protection_policy_from_params(
protection_policy: Option<codex_app_server_protocol::DeviceKeyProtectionPolicy>,
) -> DeviceKeyProtectionPolicy {
match protection_policy
.unwrap_or(codex_app_server_protocol::DeviceKeyProtectionPolicy::HardwareOnly)
{
codex_app_server_protocol::DeviceKeyProtectionPolicy::HardwareOnly => {
DeviceKeyProtectionPolicy::HardwareOnly
}
codex_app_server_protocol::DeviceKeyProtectionPolicy::AllowOsProtectedNonextractable => {
DeviceKeyProtectionPolicy::AllowOsProtectedNonextractable
}
}
}
fn payload_from_params(payload: DeviceKeySignPayload) -> codex_device_key::DeviceKeySignPayload {
match payload {
DeviceKeySignPayload::RemoteControlClientConnection {
nonce,
audience,
session_id,
target_origin,
target_path,
account_user_id,
client_id,
token_sha256_base64url,
token_expires_at,
scopes,
} => codex_device_key::DeviceKeySignPayload::RemoteControlClientConnection(
RemoteControlClientConnectionSignPayload {
nonce,
audience: remote_control_client_connection_audience_from_protocol(audience),
session_id,
target_origin,
target_path,
account_user_id,
client_id,
token_sha256_base64url,
token_expires_at,
scopes,
},
),
DeviceKeySignPayload::RemoteControlClientEnrollment {
nonce,
audience,
challenge_id,
target_origin,
target_path,
account_user_id,
client_id,
device_identity_sha256_base64url,
challenge_expires_at,
} => codex_device_key::DeviceKeySignPayload::RemoteControlClientEnrollment(
RemoteControlClientEnrollmentSignPayload {
nonce,
audience: remote_control_client_enrollment_audience_from_protocol(audience),
challenge_id,
target_origin,
target_path,
account_user_id,
client_id,
device_identity_sha256_base64url,
challenge_expires_at,
},
),
}
}
fn remote_control_client_connection_audience_from_protocol(
audience: codex_app_server_protocol::RemoteControlClientConnectionAudience,
) -> RemoteControlClientConnectionAudience {
match audience {
codex_app_server_protocol::RemoteControlClientConnectionAudience::RemoteControlClientWebsocket => {
RemoteControlClientConnectionAudience::RemoteControlClientWebsocket
}
}
}
fn remote_control_client_enrollment_audience_from_protocol(
audience: codex_app_server_protocol::RemoteControlClientEnrollmentAudience,
) -> RemoteControlClientEnrollmentAudience {
match audience {
codex_app_server_protocol::RemoteControlClientEnrollmentAudience::RemoteControlClientEnrollment => {
RemoteControlClientEnrollmentAudience::RemoteControlClientEnrollment
}
}
}
fn algorithm_from_store(algorithm: codex_device_key::DeviceKeyAlgorithm) -> DeviceKeyAlgorithm {
match algorithm {
codex_device_key::DeviceKeyAlgorithm::EcdsaP256Sha256 => {
DeviceKeyAlgorithm::EcdsaP256Sha256
}
}
}
fn protection_class_from_store(
protection_class: codex_device_key::DeviceKeyProtectionClass,
) -> DeviceKeyProtectionClass {
match protection_class {
codex_device_key::DeviceKeyProtectionClass::HardwareSecureEnclave => {
DeviceKeyProtectionClass::HardwareSecureEnclave
}
codex_device_key::DeviceKeyProtectionClass::HardwareTpm => {
DeviceKeyProtectionClass::HardwareTpm
}
codex_device_key::DeviceKeyProtectionClass::OsProtectedNonextractable => {
DeviceKeyProtectionClass::OsProtectedNonextractable
}
}
}
fn map_device_key_error(error: DeviceKeyError) -> JSONRPCErrorError {
let code = match error {
DeviceKeyError::DegradedProtectionNotAllowed { .. }
| DeviceKeyError::HardwareBackedKeysUnavailable
| DeviceKeyError::KeyNotFound
| DeviceKeyError::InvalidPayload(_) => INVALID_REQUEST_ERROR_CODE,
DeviceKeyError::Platform(_) | DeviceKeyError::Crypto(_) => INTERNAL_ERROR_CODE,
};
JSONRPCErrorError {
code,
message: error.to_string(),
data: None,
}
}
+88 -1
View File
@@ -63,6 +63,7 @@ use crate::outgoing_message::OutgoingMessage;
use crate::outgoing_message::OutgoingMessageSender;
use crate::outgoing_message::QueuedOutgoingMessage;
use crate::transport::CHANNEL_CAPACITY;
use crate::transport::ConnectionOrigin;
use crate::transport::OutboundConnectionState;
use crate::transport::route_outgoing_envelope;
use codex_analytics::AppServerRpcTransport;
@@ -416,7 +417,7 @@ fn start_uninitialized(args: InProcessStartArgs) -> InProcessClientHandle {
remote_control_handle: None,
}));
let mut thread_created_rx = processor.thread_created_receiver();
let session = Arc::new(ConnectionSessionState::default());
let session = Arc::new(ConnectionSessionState::new(ConnectionOrigin::InProcess));
let mut listen_for_threads = true;
loop {
@@ -713,6 +714,11 @@ mod tests {
use super::*;
use codex_app_server_protocol::ClientInfo;
use codex_app_server_protocol::ConfigRequirementsReadResponse;
use codex_app_server_protocol::DeviceKeyPublicParams;
use codex_app_server_protocol::DeviceKeySignParams;
use codex_app_server_protocol::DeviceKeySignPayload;
use codex_app_server_protocol::RemoteControlClientConnectionAudience;
use codex_app_server_protocol::RemoteControlClientEnrollmentAudience;
use codex_app_server_protocol::SessionSource as ApiSessionSource;
use codex_app_server_protocol::ThreadStartParams;
use codex_app_server_protocol::ThreadStartResponse;
@@ -786,6 +792,87 @@ mod tests {
.expect("in-process runtime should shutdown cleanly");
}
#[tokio::test]
async fn in_process_allows_device_key_requests_to_reach_device_key_api() {
let client = start_test_client(SessionSource::Cli).await;
const MALFORMED_KEY_ID_MESSAGE: &str = concat!(
"invalid device key payload: keyId must be dk_hse_, dk_tpm_, or dk_osn_ ",
"followed by unpadded base64url-encoded 32 bytes"
);
let requests = [
(
ClientRequest::DeviceKeyPublic {
request_id: RequestId::Integer(11),
params: DeviceKeyPublicParams {
key_id: String::new(),
},
},
MALFORMED_KEY_ID_MESSAGE,
),
(
ClientRequest::DeviceKeySign {
request_id: RequestId::Integer(12),
params: DeviceKeySignParams {
key_id: String::new(),
payload: DeviceKeySignPayload::RemoteControlClientConnection {
nonce: "nonce-123".to_string(),
audience:
RemoteControlClientConnectionAudience::RemoteControlClientWebsocket,
session_id: "wssess_123".to_string(),
target_origin: "https://chatgpt.com".to_string(),
target_path: "/api/codex/remote/control/client".to_string(),
account_user_id: "acct_123".to_string(),
client_id: "cli_123".to_string(),
token_expires_at: 4_102_444_800,
token_sha256_base64url: "47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU"
.to_string(),
scopes: vec!["remote_control_controller_websocket".to_string()],
},
},
},
MALFORMED_KEY_ID_MESSAGE,
),
(
ClientRequest::DeviceKeySign {
request_id: RequestId::Integer(13),
params: DeviceKeySignParams {
key_id: String::new(),
payload: DeviceKeySignPayload::RemoteControlClientEnrollment {
nonce: "nonce-123".to_string(),
audience:
RemoteControlClientEnrollmentAudience::RemoteControlClientEnrollment,
challenge_id: "rch_123".to_string(),
target_origin: "https://chatgpt.com".to_string(),
target_path: "/wham/remote/control/client/enroll".to_string(),
account_user_id: "acct_123".to_string(),
client_id: "cli_123".to_string(),
device_identity_sha256_base64url:
"47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU".to_string(),
challenge_expires_at: 4_102_444_800,
},
},
},
MALFORMED_KEY_ID_MESSAGE,
),
];
for (request, expected_message) in requests {
let error = client
.request(request)
.await
.expect("request transport should work")
.expect_err("request should be rejected");
assert_eq!(error.code, INVALID_REQUEST_ERROR_CODE);
assert_eq!(error.message, expected_message);
}
client
.shutdown()
.await
.expect("in-process runtime should shutdown cleanly");
}
#[tokio::test]
async fn in_process_start_uses_requested_session_source_for_thread_start() {
for (requested_source, expected_source) in [
+3
View File
@@ -73,6 +73,7 @@ mod config;
mod config_api;
mod config_manager;
mod config_manager_service;
mod device_key_api;
mod dynamic_tools;
mod error_code;
mod external_agent_config_api;
@@ -706,6 +707,7 @@ pub async fn run_main_with_transport(
match event {
TransportEvent::ConnectionOpened {
connection_id,
origin,
writer,
disconnect_sender,
} => {
@@ -735,6 +737,7 @@ pub async fn run_main_with_transport(
connections.insert(
connection_id,
ConnectionState::new(
origin,
outbound_initialized,
outbound_experimental_api_enabled,
outbound_opted_out_notification_methods,
+155 -1
View File
@@ -9,6 +9,7 @@ use crate::codex_message_processor::CodexMessageProcessor;
use crate::codex_message_processor::CodexMessageProcessorArgs;
use crate::config_api::ConfigApi;
use crate::config_manager::ConfigManager;
use crate::device_key_api::DeviceKeyApi;
use crate::error_code::INVALID_REQUEST_ERROR_CODE;
use crate::external_agent_config_api::ExternalAgentConfigApi;
use crate::fs_api::FsApi;
@@ -18,6 +19,7 @@ use crate::outgoing_message::ConnectionRequestId;
use crate::outgoing_message::OutgoingMessageSender;
use crate::outgoing_message::RequestContext;
use crate::transport::AppServerTransport;
use crate::transport::ConnectionOrigin;
use crate::transport::RemoteControlHandle;
use async_trait::async_trait;
use axum::http::HeaderValue;
@@ -35,6 +37,9 @@ use codex_app_server_protocol::ConfigBatchWriteParams;
use codex_app_server_protocol::ConfigReadParams;
use codex_app_server_protocol::ConfigValueWriteParams;
use codex_app_server_protocol::ConfigWarningNotification;
use codex_app_server_protocol::DeviceKeyCreateParams;
use codex_app_server_protocol::DeviceKeyPublicParams;
use codex_app_server_protocol::DeviceKeySignParams;
use codex_app_server_protocol::ExperimentalApi;
use codex_app_server_protocol::ExperimentalFeatureEnablementSetParams;
use codex_app_server_protocol::ExternalAgentConfigDetectParams;
@@ -164,6 +169,7 @@ pub(crate) struct MessageProcessor {
codex_message_processor: CodexMessageProcessor,
thread_manager: Arc<ThreadManager>,
config_api: ConfigApi,
device_key_api: DeviceKeyApi,
external_agent_config_api: ExternalAgentConfigApi,
fs_api: FsApi,
auth_manager: Arc<AuthManager>,
@@ -175,8 +181,9 @@ pub(crate) struct MessageProcessor {
remote_control_handle: Option<RemoteControlHandle>,
}
#[derive(Debug, Default)]
#[derive(Debug)]
pub(crate) struct ConnectionSessionState {
origin: ConnectionOrigin,
initialized: OnceLock<InitializedConnectionSessionState>,
}
@@ -188,11 +195,28 @@ struct InitializedConnectionSessionState {
client_version: String,
}
impl Default for ConnectionSessionState {
fn default() -> Self {
Self::new(ConnectionOrigin::WebSocket)
}
}
impl ConnectionSessionState {
pub(crate) fn new(origin: ConnectionOrigin) -> Self {
Self {
origin,
initialized: OnceLock::new(),
}
}
pub(crate) fn initialized(&self) -> bool {
self.initialized.get().is_some()
}
fn allows_device_key_requests(&self) -> bool {
self.origin.allows_device_key_requests()
}
pub(crate) fn experimental_api_enabled(&self) -> bool {
self.initialized
.get()
@@ -301,6 +325,7 @@ impl MessageProcessor {
thread_manager.clone(),
analytics_events_client.clone(),
);
let device_key_api = DeviceKeyApi::default();
let external_agent_config_api =
ExternalAgentConfigApi::new(config.codex_home.to_path_buf());
let fs_api = FsApi::default();
@@ -311,6 +336,7 @@ impl MessageProcessor {
codex_message_processor,
thread_manager: Arc::clone(&thread_manager),
config_api,
device_key_api,
external_agent_config_api,
fs_api,
auth_manager,
@@ -748,6 +774,7 @@ impl MessageProcessor {
let app_server_client_name = session.app_server_client_name().map(str::to_string);
let client_version = session.client_version().map(str::to_string);
let device_key_requests_allowed = session.allows_device_key_requests();
Arc::clone(self)
.handle_initialized_client_request(
connection_request_id,
@@ -755,6 +782,7 @@ impl MessageProcessor {
request_context,
app_server_client_name,
client_version,
device_key_requests_allowed,
)
.await;
}
@@ -766,6 +794,7 @@ impl MessageProcessor {
request_context: RequestContext,
app_server_client_name: Option<String>,
client_version: Option<String>,
device_key_requests_allowed: bool,
) {
let connection_id = connection_request_id.connection_id;
@@ -840,6 +869,39 @@ impl MessageProcessor {
})
.await;
}
ClientRequest::DeviceKeyCreate { request_id, params } => {
self.handle_device_key_create(
ConnectionRequestId {
connection_id,
request_id,
},
params,
device_key_requests_allowed,
)
.await;
}
ClientRequest::DeviceKeyPublic { request_id, params } => {
self.handle_device_key_public(
ConnectionRequestId {
connection_id,
request_id,
},
params,
device_key_requests_allowed,
)
.await;
}
ClientRequest::DeviceKeySign { request_id, params } => {
self.handle_device_key_sign(
ConnectionRequestId {
connection_id,
request_id,
},
params,
device_key_requests_allowed,
)
.await;
}
ClientRequest::FsReadFile { request_id, params } => {
self.handle_fs_read_file(
ConnectionRequestId {
@@ -1103,6 +1165,98 @@ impl MessageProcessor {
}
}
async fn handle_device_key_create(
&self,
request_id: ConnectionRequestId,
params: DeviceKeyCreateParams,
device_key_requests_allowed: bool,
) {
if self
.reject_device_key_request_over_remote_transport(
request_id.clone(),
"device/key/create",
device_key_requests_allowed,
)
.await
{
return;
}
match self.device_key_api.create(params) {
Ok(response) => self.outgoing.send_response(request_id, response).await,
Err(error) => self.outgoing.send_error(request_id, error).await,
}
}
async fn handle_device_key_public(
&self,
request_id: ConnectionRequestId,
params: DeviceKeyPublicParams,
device_key_requests_allowed: bool,
) {
if self
.reject_device_key_request_over_remote_transport(
request_id.clone(),
"device/key/public",
device_key_requests_allowed,
)
.await
{
return;
}
match self.device_key_api.public(params) {
Ok(response) => self.outgoing.send_response(request_id, response).await,
Err(error) => self.outgoing.send_error(request_id, error).await,
}
}
async fn handle_device_key_sign(
&self,
request_id: ConnectionRequestId,
params: DeviceKeySignParams,
device_key_requests_allowed: bool,
) {
if self
.reject_device_key_request_over_remote_transport(
request_id.clone(),
"device/key/sign",
device_key_requests_allowed,
)
.await
{
return;
}
match self.device_key_api.sign(params) {
Ok(response) => self.outgoing.send_response(request_id, response).await,
Err(error) => self.outgoing.send_error(request_id, error).await,
}
}
async fn reject_device_key_request_over_remote_transport(
&self,
request_id: ConnectionRequestId,
method: &str,
device_key_requests_allowed: bool,
) -> bool {
if device_key_requests_allowed {
return false;
}
self.outgoing
.send_error(
request_id,
JSONRPCErrorError {
code: INVALID_REQUEST_ERROR_CODE,
message: format!("{method} is not available over remote transports"),
data: None,
},
)
.await;
true
}
async fn handle_external_agent_config_detect(
&self,
request_id: ConnectionRequestId,
@@ -5,16 +5,21 @@ use crate::config_manager::ConfigManager;
use crate::outgoing_message::ConnectionId;
use crate::outgoing_message::OutgoingMessageSender;
use crate::transport::AppServerTransport;
use crate::transport::ConnectionOrigin;
use anyhow::Result;
use app_test_support::create_mock_responses_server_repeating_assistant;
use app_test_support::write_mock_responses_config_toml;
use codex_analytics::AppServerRpcTransport;
use codex_app_server_protocol::ClientInfo;
use codex_app_server_protocol::ClientRequest;
use codex_app_server_protocol::DeviceKeySignParams;
use codex_app_server_protocol::DeviceKeySignPayload;
use codex_app_server_protocol::InitializeCapabilities;
use codex_app_server_protocol::InitializeParams;
use codex_app_server_protocol::InitializeResponse;
use codex_app_server_protocol::JSONRPCErrorError;
use codex_app_server_protocol::JSONRPCRequest;
use codex_app_server_protocol::RemoteControlClientConnectionAudience;
use codex_app_server_protocol::RequestId;
use codex_app_server_protocol::ThreadStartParams;
use codex_app_server_protocol::ThreadStartResponse;
@@ -114,6 +119,10 @@ struct TracingHarness {
impl TracingHarness {
async fn new() -> Result<Self> {
Self::new_with_origin(ConnectionOrigin::WebSocket).await
}
async fn new_with_origin(origin: ConnectionOrigin) -> Result<Self> {
let server = create_mock_responses_server_repeating_assistant("Done").await;
let codex_home = TempDir::new()?;
let config = Arc::new(build_test_config(codex_home.path(), &server.uri()).await?);
@@ -126,7 +135,7 @@ impl TracingHarness {
_codex_home: codex_home,
processor,
outgoing_rx,
session: Arc::new(ConnectionSessionState::default()),
session: Arc::new(ConnectionSessionState::new(origin)),
tracing,
};
@@ -185,6 +194,29 @@ impl TracingHarness {
read_response(&mut self.outgoing_rx, request_id).await
}
async fn request_error(
&mut self,
request: ClientRequest,
trace: Option<W3cTraceContext>,
) -> JSONRPCErrorError {
let request_id = match request.id() {
RequestId::Integer(request_id) => *request_id,
request_id => panic!("expected integer request id in test harness, got {request_id:?}"),
};
let mut request = request_from_client_request(request);
request.trace = trace;
self.processor
.process_request(
TEST_CONNECTION_ID,
request,
AppServerTransport::Stdio,
Arc::clone(&self.session),
)
.await;
read_error(&mut self.outgoing_rx, request_id).await
}
async fn start_thread(
&mut self,
request_id: i64,
@@ -420,6 +452,36 @@ async fn read_response<T: serde::de::DeserializeOwned>(
}
}
async fn read_error(
outgoing_rx: &mut mpsc::Receiver<crate::outgoing_message::OutgoingEnvelope>,
request_id: i64,
) -> JSONRPCErrorError {
loop {
let envelope = tokio::time::timeout(std::time::Duration::from_secs(5), outgoing_rx.recv())
.await
.expect("timed out waiting for error")
.expect("outgoing channel closed");
let crate::outgoing_message::OutgoingEnvelope::ToConnection {
connection_id,
message,
..
} = envelope
else {
continue;
};
if connection_id != TEST_CONNECTION_ID {
continue;
}
let crate::outgoing_message::OutgoingMessage::Error(error) = message else {
continue;
};
if error.id != RequestId::Integer(request_id) {
continue;
}
return error.error;
}
}
async fn read_thread_started_notification(
outgoing_rx: &mut mpsc::Receiver<crate::outgoing_message::OutgoingEnvelope>,
) {
@@ -585,6 +647,47 @@ async fn thread_start_jsonrpc_span_exports_server_span_and_parents_children() ->
Ok(())
}
#[tokio::test(flavor = "current_thread")]
#[serial(app_server_tracing)]
async fn remote_control_origin_rejects_device_key_requests() -> Result<()> {
let mut harness = TracingHarness::new_with_origin(ConnectionOrigin::RemoteControl).await?;
let error = harness
.request_error(
ClientRequest::DeviceKeySign {
request_id: RequestId::Integer(20_004),
params: DeviceKeySignParams {
key_id: "dk_123".to_string(),
payload: DeviceKeySignPayload::RemoteControlClientConnection {
nonce: "nonce-123".to_string(),
audience:
RemoteControlClientConnectionAudience::RemoteControlClientWebsocket,
session_id: "wssess_123".to_string(),
target_origin: "https://chatgpt.com".to_string(),
target_path: "/api/codex/remote/control/client".to_string(),
account_user_id: "acct_123".to_string(),
client_id: "cli_123".to_string(),
token_expires_at: 4_102_444_800,
token_sha256_base64url: "47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU"
.to_string(),
scopes: vec!["remote_control_controller_websocket".to_string()],
},
},
},
/*trace*/ None,
)
.await;
assert_eq!(error.code, crate::error_code::INVALID_REQUEST_ERROR_CODE);
assert_eq!(
error.message,
"device/key/sign is not available over remote transports"
);
harness.shutdown().await;
Ok(())
}
#[tokio::test(flavor = "current_thread")]
#[serial(app_server_tracing)]
async fn turn_start_jsonrpc_span_parents_core_turn_spans() -> Result<()> {
+19 -1
View File
@@ -106,6 +106,7 @@ impl FromStr for AppServerTransport {
pub(crate) enum TransportEvent {
ConnectionOpened {
connection_id: ConnectionId,
origin: ConnectionOrigin,
writer: mpsc::Sender<QueuedOutgoingMessage>,
disconnect_sender: Option<CancellationToken>,
},
@@ -118,6 +119,22 @@ pub(crate) enum TransportEvent {
},
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum ConnectionOrigin {
Stdio,
InProcess,
WebSocket,
RemoteControl,
}
impl ConnectionOrigin {
pub(crate) fn allows_device_key_requests(self) -> bool {
// Device-key endpoints are only for local connections that own the app-server instance.
// Do not include remote transports such as SSH or remote-control websocket connections.
matches!(self, Self::Stdio | Self::InProcess)
}
}
pub(crate) struct ConnectionState {
pub(crate) outbound_initialized: Arc<AtomicBool>,
pub(crate) outbound_experimental_api_enabled: Arc<AtomicBool>,
@@ -127,6 +144,7 @@ pub(crate) struct ConnectionState {
impl ConnectionState {
pub(crate) fn new(
origin: ConnectionOrigin,
outbound_initialized: Arc<AtomicBool>,
outbound_experimental_api_enabled: Arc<AtomicBool>,
outbound_opted_out_notification_methods: Arc<RwLock<HashSet<String>>>,
@@ -135,7 +153,7 @@ impl ConnectionState {
outbound_initialized,
outbound_experimental_api_enabled,
outbound_opted_out_notification_methods,
session: Arc::new(ConnectionSessionState::default()),
session: Arc::new(ConnectionSessionState::new(origin)),
}
}
}
@@ -9,6 +9,7 @@ use super::protocol::ServerEvent;
use super::protocol::StreamId;
use crate::outgoing_message::ConnectionId;
use crate::outgoing_message::QueuedOutgoingMessage;
use crate::transport::ConnectionOrigin;
use crate::transport::remote_control::QueuedServerEnvelope;
use codex_app_server_protocol::JSONRPCMessage;
use std::collections::HashMap;
@@ -160,6 +161,7 @@ impl ClientTracker {
let disconnect_token = self.shutdown_token.child_token();
self.send_transport_event(TransportEvent::ConnectionOpened {
connection_id,
origin: ConnectionOrigin::RemoteControl,
writer: writer_tx,
disconnect_sender: Some(disconnect_token.clone()),
})
@@ -12,6 +12,7 @@ use super::*;
use crate::outgoing_message::OutgoingMessage;
use crate::outgoing_message::QueuedOutgoingMessage;
use crate::transport::CHANNEL_CAPACITY;
use crate::transport::ConnectionOrigin;
use crate::transport::TransportEvent;
use base64::Engine;
use codex_app_server_protocol::AuthMode;
@@ -226,9 +227,13 @@ async fn remote_control_transport_manages_virtual_clients_and_routes_messages()
{
TransportEvent::ConnectionOpened {
connection_id,
origin,
writer,
..
} => (connection_id, writer),
} => {
assert_eq!(origin, ConnectionOrigin::RemoteControl);
(connection_id, writer)
}
other => panic!("expected connection open event, got {other:?}"),
};
@@ -1,4 +1,5 @@
use super::CHANNEL_CAPACITY;
use super::ConnectionOrigin;
use super::TransportEvent;
use super::forward_incoming_message;
use super::next_connection_id;
@@ -31,6 +32,7 @@ pub(crate) async fn start_stdio_connection(
transport_event_tx
.send(TransportEvent::ConnectionOpened {
connection_id,
origin: ConnectionOrigin::Stdio,
writer: writer_tx,
disconnect_sender: None,
})
@@ -1,4 +1,5 @@
use super::CHANNEL_CAPACITY;
use super::ConnectionOrigin;
use super::TransportEvent;
use super::auth::WebsocketAuthPolicy;
use super::auth::authorize_upgrade;
@@ -172,6 +173,7 @@ async fn run_websocket_connection(
if transport_event_tx
.send(TransportEvent::ConnectionOpened {
connection_id,
origin: ConnectionOrigin::WebSocket,
writer: writer_tx,
disconnect_sender: Some(disconnect_token.clone()),
})
@@ -777,7 +777,7 @@ pub(super) async fn read_response_and_notification_for_method(
Ok((response, notification))
}
async fn read_error_for_id(stream: &mut WsClient, id: i64) -> Result<JSONRPCError> {
pub(super) async fn read_error_for_id(stream: &mut WsClient, id: i64) -> Result<JSONRPCError> {
let target_id = RequestId::Integer(id);
loop {
let message = read_jsonrpc_message(stream).await?;
@@ -0,0 +1,119 @@
use super::connection_handling_websocket::connect_websocket;
use super::connection_handling_websocket::create_config_toml;
use super::connection_handling_websocket::read_error_for_id;
use super::connection_handling_websocket::read_response_for_id;
use super::connection_handling_websocket::send_initialize_request;
use super::connection_handling_websocket::send_request;
use super::connection_handling_websocket::spawn_websocket_server;
use anyhow::Result;
use app_test_support::McpProcess;
use app_test_support::create_mock_responses_server_sequence_unchecked;
use codex_app_server_protocol::RequestId;
use pretty_assertions::assert_eq;
use serde_json::json;
use tempfile::TempDir;
use tokio::time::Duration;
use tokio::time::timeout;
#[cfg(any(target_os = "macos", windows))]
const DEFAULT_READ_TIMEOUT: Duration = Duration::from_secs(60);
#[cfg(not(any(target_os = "macos", windows)))]
const DEFAULT_READ_TIMEOUT: Duration = Duration::from_secs(10);
async fn initialized_mcp(codex_home: &TempDir) -> Result<McpProcess> {
let mut mcp = McpProcess::new(codex_home.path()).await?;
timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??;
Ok(mcp)
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn device_key_create_rejects_empty_account_user_id() -> Result<()> {
let codex_home = TempDir::new()?;
let mut mcp = initialized_mcp(&codex_home).await?;
let request_id = mcp
.send_raw_request(
"device/key/create",
Some(json!({
"accountUserId": "",
"clientId": "cli_123",
})),
)
.await?;
let error = timeout(
DEFAULT_READ_TIMEOUT,
mcp.read_stream_until_error_message(RequestId::Integer(request_id)),
)
.await??;
assert_eq!(error.error.code, -32600);
assert_eq!(
error.error.message,
"invalid device key payload: accountUserId must not be empty"
);
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn device_key_methods_are_rejected_over_websocket() -> Result<()> {
let server = create_mock_responses_server_sequence_unchecked(Vec::new()).await;
let codex_home = TempDir::new()?;
create_config_toml(codex_home.path(), &server.uri(), "never")?;
let (mut process, bind_addr) = spawn_websocket_server(codex_home.path()).await?;
let mut ws = connect_websocket(bind_addr).await?;
send_initialize_request(&mut ws, /*id*/ 1, "device_key_ws_test").await?;
let initialize_response = read_response_for_id(&mut ws, /*id*/ 1).await?;
assert_eq!(initialize_response.id, RequestId::Integer(1));
let cases = [
(
"device/key/create",
json!({
"accountUserId": "acct_123",
"clientId": "cli_123",
}),
),
(
"device/key/public",
json!({
"keyId": "device-key-123",
}),
),
(
"device/key/sign",
json!({
"keyId": "device-key-123",
"payload": {
"type": "remoteControlClientConnection",
"nonce": "nonce-123",
"audience": "remote_control_client_websocket",
"sessionId": "wssess_123",
"targetOrigin": "https://chatgpt.com",
"targetPath": "/api/codex/remote/control/client",
"accountUserId": "acct_123",
"clientId": "cli_123",
"tokenExpiresAt": 4_102_444_800i64,
"tokenSha256Base64url": "47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU",
"scopes": ["remote_control_controller_websocket"],
},
}),
),
];
for (index, (method, params)) in cases.into_iter().enumerate() {
let id = 2 + index as i64;
send_request(&mut ws, method, id, Some(params)).await?;
let error = read_error_for_id(&mut ws, id).await?;
assert_eq!(error.error.code, -32600);
assert_eq!(
error.error.message,
format!("{method} is not available over remote transports")
);
}
process.kill().await?;
Ok(())
}
@@ -10,6 +10,7 @@ mod config_rpc;
mod connection_handling_websocket;
#[cfg(unix)]
mod connection_handling_websocket_unix;
mod device_key;
mod dynamic_tools;
mod experimental_api;
mod experimental_feature_list;