mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
app-server: prepare to run initialized rpcs concurrently (#17372)
## Summary - Refactors `MessageProcessor` and per-connection session state so initialized service RPC handling can be moved into spawned tasks in a follow-up PR. - Shares the processor and initialized session data with `Arc`/`OnceLock` instead of mutable borrowed connection state. - Keeps initialized request handling synchronous in this PR; it does **not** call `tokio::spawn` for service RPCs yet. ## Testing - `just fmt` - `cargo test -p codex-app-server` *(fails on existing hardening gaps covered by #17375, #17376, and #17377; the pipelined config regression passed before the unrelated failures)* - `just fix -p codex-app-server`
This commit is contained in:
committed by
GitHub
Unverified
parent
769b1c3d7e
commit
23d4098c0f
@@ -72,10 +72,10 @@ pub(crate) fn typed_request_span(
|
||||
&span,
|
||||
client_info
|
||||
.map(|(client_name, _)| client_name)
|
||||
.or(session.app_server_client_name.as_deref()),
|
||||
.or(session.app_server_client_name()),
|
||||
client_info
|
||||
.map(|(_, client_version)| client_version)
|
||||
.or(session.client_version.as_deref()),
|
||||
.or(session.client_version()),
|
||||
);
|
||||
|
||||
attach_parent_context(&span, &method, request.id(), /*parent_trace*/ None);
|
||||
@@ -147,7 +147,7 @@ fn client_name<'a>(
|
||||
if let Some(params) = initialize_client_info {
|
||||
return Some(params.client_info.name.as_str());
|
||||
}
|
||||
session.app_server_client_name.as_deref()
|
||||
session.app_server_client_name()
|
||||
}
|
||||
|
||||
fn client_version<'a>(
|
||||
@@ -157,7 +157,7 @@ fn client_version<'a>(
|
||||
if let Some(params) = initialize_client_info {
|
||||
return Some(params.client_info.version.as_str());
|
||||
}
|
||||
session.client_version.as_deref()
|
||||
session.client_version()
|
||||
}
|
||||
|
||||
fn initialize_client_info(request: &JSONRPCRequest) -> Option<InitializeParams> {
|
||||
|
||||
@@ -386,7 +386,7 @@ fn start_uninitialized(args: InProcessStartArgs) -> InProcessClientHandle {
|
||||
AuthManager::shared_from_config(args.config.as_ref(), args.enable_codex_api_key_env);
|
||||
let (processor_tx, mut processor_rx) = mpsc::channel::<ProcessorCommand>(channel_capacity);
|
||||
let mut processor_handle = tokio::spawn(async move {
|
||||
let processor = MessageProcessor::new(MessageProcessorArgs {
|
||||
let processor = Arc::new(MessageProcessor::new(MessageProcessorArgs {
|
||||
outgoing: Arc::clone(&processor_outgoing),
|
||||
arg0_paths: args.arg0_paths,
|
||||
config: args.config,
|
||||
@@ -401,9 +401,9 @@ fn start_uninitialized(args: InProcessStartArgs) -> InProcessClientHandle {
|
||||
auth_manager,
|
||||
rpc_transport: AppServerRpcTransport::InProcess,
|
||||
remote_control_handle: None,
|
||||
});
|
||||
}));
|
||||
let mut thread_created_rx = processor.thread_created_receiver();
|
||||
let mut session = ConnectionSessionState::default();
|
||||
let session = Arc::new(ConnectionSessionState::default());
|
||||
let mut listen_for_threads = true;
|
||||
|
||||
loop {
|
||||
@@ -411,28 +411,33 @@ fn start_uninitialized(args: InProcessStartArgs) -> InProcessClientHandle {
|
||||
command = processor_rx.recv() => {
|
||||
match command {
|
||||
Some(ProcessorCommand::Request(request)) => {
|
||||
let was_initialized = session.initialized;
|
||||
let was_initialized = session.initialized();
|
||||
processor
|
||||
.process_client_request(
|
||||
IN_PROCESS_CONNECTION_ID,
|
||||
*request,
|
||||
&mut session,
|
||||
Arc::clone(&session),
|
||||
&outbound_initialized,
|
||||
)
|
||||
.await;
|
||||
let opted_out_notification_methods_snapshot =
|
||||
session.opted_out_notification_methods();
|
||||
let experimental_api_enabled =
|
||||
session.experimental_api_enabled();
|
||||
let is_initialized = session.initialized();
|
||||
if let Ok(mut opted_out_notification_methods) =
|
||||
outbound_opted_out_notification_methods.write()
|
||||
{
|
||||
*opted_out_notification_methods =
|
||||
session.opted_out_notification_methods.clone();
|
||||
opted_out_notification_methods_snapshot;
|
||||
} else {
|
||||
warn!("failed to update outbound opted-out notifications");
|
||||
}
|
||||
outbound_experimental_api_enabled.store(
|
||||
session.experimental_api_enabled,
|
||||
experimental_api_enabled,
|
||||
Ordering::Release,
|
||||
);
|
||||
if !was_initialized && session.initialized {
|
||||
if !was_initialized && is_initialized {
|
||||
processor.send_initialize_notifications().await;
|
||||
}
|
||||
}
|
||||
@@ -447,7 +452,7 @@ fn start_uninitialized(args: InProcessStartArgs) -> InProcessClientHandle {
|
||||
created = thread_created_rx.recv(), if listen_for_threads => {
|
||||
match created {
|
||||
Ok(thread_id) => {
|
||||
let connection_ids = if session.initialized {
|
||||
let connection_ids = if session.initialized() {
|
||||
vec![IN_PROCESS_CONNECTION_ID]
|
||||
} else {
|
||||
Vec::<ConnectionId>::new()
|
||||
|
||||
@@ -652,7 +652,7 @@ pub async fn run_main_with_transport(
|
||||
AuthManager::shared_from_config(&config, /*enable_codex_api_key_env*/ false);
|
||||
let cli_overrides: Vec<(String, TomlValue)> = cli_kv_overrides.clone();
|
||||
let loader_overrides = loader_overrides_for_config_api;
|
||||
let processor = MessageProcessor::new(MessageProcessorArgs {
|
||||
let processor = Arc::new(MessageProcessor::new(MessageProcessorArgs {
|
||||
outgoing: outgoing_message_sender,
|
||||
arg0_paths,
|
||||
config: Arc::new(config),
|
||||
@@ -667,7 +667,7 @@ pub async fn run_main_with_transport(
|
||||
auth_manager,
|
||||
rpc_transport: analytics_rpc_transport(transport),
|
||||
remote_control_handle: Some(remote_control_handle),
|
||||
});
|
||||
}));
|
||||
let mut thread_created_rx = processor.thread_created_receiver();
|
||||
let mut running_turn_count_rx = processor.subscribe_running_assistant_turn_count();
|
||||
let mut connections = HashMap::<ConnectionId, ConnectionState>::new();
|
||||
@@ -769,23 +769,28 @@ pub async fn run_main_with_transport(
|
||||
warn!("dropping request from unknown connection: {connection_id:?}");
|
||||
continue;
|
||||
};
|
||||
let was_initialized = connection_state.session.initialized;
|
||||
let was_initialized =
|
||||
connection_state.session.initialized();
|
||||
processor
|
||||
.process_request(
|
||||
connection_id,
|
||||
request,
|
||||
transport,
|
||||
&mut connection_state.session,
|
||||
Arc::clone(&connection_state.session),
|
||||
)
|
||||
.await;
|
||||
let opted_out_notification_methods_snapshot = connection_state
|
||||
.session
|
||||
.opted_out_notification_methods();
|
||||
let experimental_api_enabled =
|
||||
connection_state.session.experimental_api_enabled();
|
||||
let is_initialized = connection_state.session.initialized();
|
||||
if let Ok(mut opted_out_notification_methods) = connection_state
|
||||
.outbound_opted_out_notification_methods
|
||||
.write()
|
||||
{
|
||||
*opted_out_notification_methods = connection_state
|
||||
.session
|
||||
.opted_out_notification_methods
|
||||
.clone();
|
||||
*opted_out_notification_methods =
|
||||
opted_out_notification_methods_snapshot;
|
||||
} else {
|
||||
warn!(
|
||||
"failed to update outbound opted-out notifications"
|
||||
@@ -794,10 +799,10 @@ pub async fn run_main_with_transport(
|
||||
connection_state
|
||||
.outbound_experimental_api_enabled
|
||||
.store(
|
||||
connection_state.session.experimental_api_enabled,
|
||||
experimental_api_enabled,
|
||||
std::sync::atomic::Ordering::Release,
|
||||
);
|
||||
if !was_initialized && connection_state.session.initialized {
|
||||
if !was_initialized && is_initialized {
|
||||
processor
|
||||
.send_initialize_notifications_to_connection(
|
||||
connection_id,
|
||||
@@ -837,12 +842,12 @@ pub async fn run_main_with_transport(
|
||||
created = thread_created_rx.recv(), if listen_for_threads => {
|
||||
match created {
|
||||
Ok(thread_id) => {
|
||||
let initialized_connection_ids: Vec<ConnectionId> = connections
|
||||
.iter()
|
||||
.filter_map(|(connection_id, connection_state)| {
|
||||
connection_state.session.initialized.then_some(*connection_id)
|
||||
})
|
||||
.collect();
|
||||
let mut initialized_connection_ids = Vec::new();
|
||||
for (connection_id, connection_state) in &connections {
|
||||
if connection_state.session.initialized() {
|
||||
initialized_connection_ids.push(*connection_id);
|
||||
}
|
||||
}
|
||||
processor
|
||||
.try_attach_thread_listener(
|
||||
thread_id,
|
||||
|
||||
@@ -2,6 +2,7 @@ use std::collections::BTreeMap;
|
||||
use std::collections::HashSet;
|
||||
use std::future::Future;
|
||||
use std::sync::Arc;
|
||||
use std::sync::OnceLock;
|
||||
use std::sync::RwLock;
|
||||
use std::sync::atomic::AtomicBool;
|
||||
use std::sync::atomic::Ordering;
|
||||
@@ -20,6 +21,7 @@ use crate::outgoing_message::RequestContext;
|
||||
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::AppListUpdatedNotification;
|
||||
@@ -173,13 +175,52 @@ pub(crate) struct MessageProcessor {
|
||||
remote_control_handle: Option<RemoteControlHandle>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Default)]
|
||||
#[derive(Debug, Default)]
|
||||
pub(crate) struct ConnectionSessionState {
|
||||
pub(crate) initialized: bool,
|
||||
pub(crate) experimental_api_enabled: bool,
|
||||
pub(crate) opted_out_notification_methods: HashSet<String>,
|
||||
pub(crate) app_server_client_name: Option<String>,
|
||||
pub(crate) client_version: Option<String>,
|
||||
initialized: OnceLock<InitializedConnectionSessionState>,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct InitializedConnectionSessionState {
|
||||
experimental_api_enabled: bool,
|
||||
opted_out_notification_methods: HashSet<String>,
|
||||
app_server_client_name: String,
|
||||
client_version: String,
|
||||
}
|
||||
|
||||
impl ConnectionSessionState {
|
||||
pub(crate) fn initialized(&self) -> bool {
|
||||
self.initialized.get().is_some()
|
||||
}
|
||||
|
||||
pub(crate) fn experimental_api_enabled(&self) -> bool {
|
||||
self.initialized
|
||||
.get()
|
||||
.is_some_and(|session| session.experimental_api_enabled)
|
||||
}
|
||||
|
||||
pub(crate) fn opted_out_notification_methods(&self) -> HashSet<String> {
|
||||
self.initialized
|
||||
.get()
|
||||
.map(|session| session.opted_out_notification_methods.clone())
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
pub(crate) fn app_server_client_name(&self) -> Option<&str> {
|
||||
self.initialized
|
||||
.get()
|
||||
.map(|session| session.app_server_client_name.as_str())
|
||||
}
|
||||
|
||||
pub(crate) fn client_version(&self) -> Option<&str> {
|
||||
self.initialized
|
||||
.get()
|
||||
.map(|session| session.client_version.as_str())
|
||||
}
|
||||
|
||||
fn initialize(&self, session: InitializedConnectionSessionState) -> Result<(), ()> {
|
||||
self.initialized.set(session).map_err(|_| ())
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) struct MessageProcessorArgs {
|
||||
@@ -299,11 +340,11 @@ impl MessageProcessor {
|
||||
}
|
||||
|
||||
pub(crate) async fn process_request(
|
||||
&self,
|
||||
self: &Arc<Self>,
|
||||
connection_id: ConnectionId,
|
||||
request: JSONRPCRequest,
|
||||
transport: AppServerTransport,
|
||||
session: &mut ConnectionSessionState,
|
||||
session: Arc<ConnectionSessionState>,
|
||||
) {
|
||||
let request_method = request.method.as_str();
|
||||
tracing::trace!(
|
||||
@@ -316,7 +357,7 @@ impl MessageProcessor {
|
||||
request_id: request.id.clone(),
|
||||
};
|
||||
let request_span =
|
||||
crate::app_server_tracing::request_span(&request, transport, connection_id, session);
|
||||
crate::app_server_tracing::request_span(&request, transport, connection_id, &session);
|
||||
let request_trace = request.trace.as_ref().map(|trace| W3cTraceContext {
|
||||
traceparent: trace.traceparent.clone(),
|
||||
tracestate: trace.tracestate.clone(),
|
||||
@@ -358,7 +399,7 @@ impl MessageProcessor {
|
||||
self.handle_client_request(
|
||||
request_id.clone(),
|
||||
codex_request,
|
||||
session,
|
||||
Arc::clone(&session),
|
||||
/*outbound_initialized*/ None,
|
||||
request_context.clone(),
|
||||
)
|
||||
@@ -373,10 +414,10 @@ impl MessageProcessor {
|
||||
/// This bypasses JSON request deserialization but keeps identical request
|
||||
/// semantics by delegating to `handle_client_request`.
|
||||
pub(crate) async fn process_client_request(
|
||||
&self,
|
||||
self: &Arc<Self>,
|
||||
connection_id: ConnectionId,
|
||||
request: ClientRequest,
|
||||
session: &mut ConnectionSessionState,
|
||||
session: Arc<ConnectionSessionState>,
|
||||
outbound_initialized: &AtomicBool,
|
||||
) {
|
||||
let request_id = ConnectionRequestId {
|
||||
@@ -384,7 +425,7 @@ impl MessageProcessor {
|
||||
request_id: request.id().clone(),
|
||||
};
|
||||
let request_span =
|
||||
crate::app_server_tracing::typed_request_span(&request, connection_id, session);
|
||||
crate::app_server_tracing::typed_request_span(&request, connection_id, &session);
|
||||
let request_context =
|
||||
RequestContext::new(request_id.clone(), request_span, /*parent_trace*/ None);
|
||||
tracing::trace!(
|
||||
@@ -402,7 +443,7 @@ impl MessageProcessor {
|
||||
self.handle_client_request(
|
||||
request_id.clone(),
|
||||
request,
|
||||
session,
|
||||
Arc::clone(&session),
|
||||
Some(outbound_initialized),
|
||||
request_context.clone(),
|
||||
)
|
||||
@@ -525,10 +566,10 @@ impl MessageProcessor {
|
||||
}
|
||||
|
||||
async fn handle_client_request(
|
||||
&self,
|
||||
self: &Arc<Self>,
|
||||
connection_request_id: ConnectionRequestId,
|
||||
codex_request: ClientRequest,
|
||||
session: &mut ConnectionSessionState,
|
||||
session: Arc<ConnectionSessionState>,
|
||||
// `Some(...)` means the caller wants initialize to immediately mark the
|
||||
// connection outbound-ready. Websocket JSON-RPC calls pass `None` so
|
||||
// lib.rs can deliver connection-scoped initialize notifications first.
|
||||
@@ -536,126 +577,166 @@ impl MessageProcessor {
|
||||
request_context: RequestContext,
|
||||
) {
|
||||
let connection_id = connection_request_id.connection_id;
|
||||
match codex_request {
|
||||
if let ClientRequest::Initialize { request_id, params } = codex_request {
|
||||
// Handle Initialize internally so CodexMessageProcessor does not have to concern
|
||||
// itself with the `initialized` bool.
|
||||
ClientRequest::Initialize { request_id, params } => {
|
||||
let connection_request_id = ConnectionRequestId {
|
||||
connection_id,
|
||||
request_id,
|
||||
let connection_request_id = ConnectionRequestId {
|
||||
connection_id,
|
||||
request_id,
|
||||
};
|
||||
if session.initialized() {
|
||||
let error = JSONRPCErrorError {
|
||||
code: INVALID_REQUEST_ERROR_CODE,
|
||||
message: "Already initialized".to_string(),
|
||||
data: None,
|
||||
};
|
||||
if session.initialized {
|
||||
let error = JSONRPCErrorError {
|
||||
code: INVALID_REQUEST_ERROR_CODE,
|
||||
message: "Already initialized".to_string(),
|
||||
data: None,
|
||||
};
|
||||
self.outgoing.send_error(connection_request_id, error).await;
|
||||
return;
|
||||
}
|
||||
|
||||
// TODO(maxj): Revisit capability scoping for `experimental_api_enabled`.
|
||||
// Current behavior is per-connection. Reviewer feedback notes this can
|
||||
// create odd cross-client behavior (for example dynamic tool calls on a
|
||||
// shared thread when another connected client did not opt into
|
||||
// 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()),
|
||||
};
|
||||
session.experimental_api_enabled = experimental_api_enabled;
|
||||
session.opted_out_notification_methods =
|
||||
opt_out_notification_methods.into_iter().collect();
|
||||
let ClientInfo {
|
||||
name,
|
||||
title: _title,
|
||||
version,
|
||||
} = params.client_info;
|
||||
session.app_server_client_name = Some(name.clone());
|
||||
session.client_version = Some(version.clone());
|
||||
let originator = name.clone();
|
||||
if let Err(error) = set_default_originator(originator.clone()) {
|
||||
match error {
|
||||
SetOriginatorError::InvalidHeaderValue => {
|
||||
let error = JSONRPCErrorError {
|
||||
code: INVALID_REQUEST_ERROR_CODE,
|
||||
message: format!(
|
||||
"Invalid clientInfo.name: '{name}'. Must be a valid HTTP header value."
|
||||
),
|
||||
data: None,
|
||||
};
|
||||
self.outgoing
|
||||
.send_error(connection_request_id.clone(), error)
|
||||
.await;
|
||||
return;
|
||||
}
|
||||
SetOriginatorError::AlreadyInitialized => {
|
||||
// No-op. This is expected to happen if the originator is already set via env var.
|
||||
// TODO(owen): Once we remove support for CODEX_INTERNAL_ORIGINATOR_OVERRIDE,
|
||||
// this will be an unexpected state and we can return a JSON-RPC error indicating
|
||||
// internal server error.
|
||||
}
|
||||
}
|
||||
}
|
||||
if self.config.features.enabled(Feature::GeneralAnalytics) {
|
||||
self.analytics_events_client.track_initialize(
|
||||
connection_id.0,
|
||||
analytics_initialize_params,
|
||||
originator,
|
||||
self.rpc_transport,
|
||||
);
|
||||
}
|
||||
set_default_client_residency_requirement(self.config.enforce_residency.value());
|
||||
let user_agent_suffix = format!("{name}; {version}");
|
||||
if let Ok(mut suffix) = USER_AGENT_SUFFIX.lock() {
|
||||
*suffix = Some(user_agent_suffix);
|
||||
}
|
||||
|
||||
let user_agent = get_codex_user_agent();
|
||||
let response = InitializeResponse {
|
||||
user_agent,
|
||||
codex_home: self.config.codex_home.clone(),
|
||||
platform_family: std::env::consts::FAMILY.to_string(),
|
||||
platform_os: std::env::consts::OS.to_string(),
|
||||
};
|
||||
self.outgoing
|
||||
.send_response(connection_request_id, response)
|
||||
.await;
|
||||
|
||||
session.initialized = true;
|
||||
if let Some(outbound_initialized) = outbound_initialized {
|
||||
// In-process clients can complete readiness immediately here. The
|
||||
// websocket path defers this until lib.rs finishes transport-layer
|
||||
// initialize handling for the specific connection.
|
||||
outbound_initialized.store(true, Ordering::Release);
|
||||
self.codex_message_processor
|
||||
.connection_initialized(connection_id)
|
||||
.await;
|
||||
}
|
||||
self.outgoing.send_error(connection_request_id, error).await;
|
||||
return;
|
||||
}
|
||||
_ => {
|
||||
if !session.initialized {
|
||||
let error = JSONRPCErrorError {
|
||||
code: INVALID_REQUEST_ERROR_CODE,
|
||||
message: "Not initialized".to_string(),
|
||||
data: None,
|
||||
};
|
||||
self.outgoing.send_error(connection_request_id, error).await;
|
||||
return;
|
||||
|
||||
// TODO(maxj): Revisit capability scoping for `experimental_api_enabled`.
|
||||
// Current behavior is per-connection. Reviewer feedback notes this can
|
||||
// create odd cross-client behavior (for example dynamic tool calls on a
|
||||
// shared thread when another connected client did not opt into
|
||||
// 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 ClientInfo {
|
||||
name,
|
||||
title: _title,
|
||||
version,
|
||||
} = params.client_info;
|
||||
// Validate before committing; set_default_originator validates while
|
||||
// mutating process-global metadata.
|
||||
if HeaderValue::from_str(&name).is_err() {
|
||||
let error = JSONRPCErrorError {
|
||||
code: INVALID_REQUEST_ERROR_CODE,
|
||||
message: format!(
|
||||
"Invalid clientInfo.name: '{name}'. Must be a valid HTTP header value."
|
||||
),
|
||||
data: None,
|
||||
};
|
||||
self.outgoing
|
||||
.send_error(connection_request_id.clone(), error)
|
||||
.await;
|
||||
return;
|
||||
}
|
||||
let originator = name.clone();
|
||||
let user_agent_suffix = format!("{name}; {version}");
|
||||
let codex_home = self.config.codex_home.clone();
|
||||
if session
|
||||
.initialize(InitializedConnectionSessionState {
|
||||
experimental_api_enabled,
|
||||
opted_out_notification_methods: opt_out_notification_methods
|
||||
.into_iter()
|
||||
.collect(),
|
||||
app_server_client_name: name.clone(),
|
||||
client_version: version,
|
||||
})
|
||||
.is_err()
|
||||
{
|
||||
let error = JSONRPCErrorError {
|
||||
code: INVALID_REQUEST_ERROR_CODE,
|
||||
message: "Already initialized".to_string(),
|
||||
data: None,
|
||||
};
|
||||
self.outgoing.send_error(connection_request_id, error).await;
|
||||
return;
|
||||
}
|
||||
|
||||
// Only the request that wins session initialization may mutate
|
||||
// process-global client metadata.
|
||||
if let Err(error) = set_default_originator(originator.clone()) {
|
||||
match error {
|
||||
SetOriginatorError::InvalidHeaderValue => {
|
||||
tracing::warn!(
|
||||
client_info_name = %name,
|
||||
"validated clientInfo.name was rejected while setting originator"
|
||||
);
|
||||
}
|
||||
SetOriginatorError::AlreadyInitialized => {
|
||||
// No-op. This is expected to happen if the originator is already set via env var.
|
||||
// TODO(owen): Once we remove support for CODEX_INTERNAL_ORIGINATOR_OVERRIDE,
|
||||
// this will be an unexpected state and we can return a JSON-RPC error indicating
|
||||
// internal server error.
|
||||
}
|
||||
}
|
||||
}
|
||||
if self.config.features.enabled(Feature::GeneralAnalytics) {
|
||||
self.analytics_events_client.track_initialize(
|
||||
connection_id.0,
|
||||
analytics_initialize_params,
|
||||
originator,
|
||||
self.rpc_transport,
|
||||
);
|
||||
}
|
||||
set_default_client_residency_requirement(self.config.enforce_residency.value());
|
||||
if let Ok(mut suffix) = USER_AGENT_SUFFIX.lock() {
|
||||
*suffix = Some(user_agent_suffix);
|
||||
}
|
||||
|
||||
let user_agent = get_codex_user_agent();
|
||||
let response = InitializeResponse {
|
||||
user_agent,
|
||||
codex_home,
|
||||
platform_family: std::env::consts::FAMILY.to_string(),
|
||||
platform_os: std::env::consts::OS.to_string(),
|
||||
};
|
||||
|
||||
self.outgoing
|
||||
.send_response(connection_request_id, response)
|
||||
.await;
|
||||
|
||||
if let Some(outbound_initialized) = outbound_initialized {
|
||||
// In-process clients can complete readiness immediately here. The
|
||||
// websocket path defers this until lib.rs finishes transport-layer
|
||||
// initialize handling for the specific connection.
|
||||
outbound_initialized.store(true, Ordering::Release);
|
||||
self.codex_message_processor
|
||||
.connection_initialized(connection_id)
|
||||
.await;
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
self.dispatch_initialized_client_request(
|
||||
connection_request_id,
|
||||
codex_request,
|
||||
session,
|
||||
request_context,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
async fn dispatch_initialized_client_request(
|
||||
self: &Arc<Self>,
|
||||
connection_request_id: ConnectionRequestId,
|
||||
codex_request: ClientRequest,
|
||||
session: Arc<ConnectionSessionState>,
|
||||
request_context: RequestContext,
|
||||
) {
|
||||
if !session.initialized() {
|
||||
let error = JSONRPCErrorError {
|
||||
code: INVALID_REQUEST_ERROR_CODE,
|
||||
message: "Not initialized".to_string(),
|
||||
data: None,
|
||||
};
|
||||
self.outgoing.send_error(connection_request_id, error).await;
|
||||
return;
|
||||
}
|
||||
|
||||
if let Some(reason) = codex_request.experimental_reason()
|
||||
&& !session.experimental_api_enabled
|
||||
&& !session.experimental_api_enabled()
|
||||
{
|
||||
let error = JSONRPCErrorError {
|
||||
code: INVALID_REQUEST_ERROR_CODE,
|
||||
@@ -676,6 +757,29 @@ 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);
|
||||
Arc::clone(self)
|
||||
.handle_initialized_client_request(
|
||||
connection_request_id,
|
||||
codex_request,
|
||||
request_context,
|
||||
app_server_client_name,
|
||||
client_version,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
async fn handle_initialized_client_request(
|
||||
self: Arc<Self>,
|
||||
connection_request_id: ConnectionRequestId,
|
||||
codex_request: ClientRequest,
|
||||
request_context: RequestContext,
|
||||
app_server_client_name: Option<String>,
|
||||
client_version: Option<String>,
|
||||
) {
|
||||
let connection_id = connection_request_id.connection_id;
|
||||
|
||||
match codex_request {
|
||||
ClientRequest::ConfigRead { request_id, params } => {
|
||||
self.handle_config_read(
|
||||
@@ -847,8 +951,8 @@ impl MessageProcessor {
|
||||
.process_request(
|
||||
connection_id,
|
||||
other,
|
||||
session.app_server_client_name.clone(),
|
||||
session.client_version.clone(),
|
||||
app_server_client_name,
|
||||
client_version,
|
||||
request_context,
|
||||
)
|
||||
.boxed()
|
||||
|
||||
@@ -109,9 +109,9 @@ fn tracing_test_guard() -> &'static tokio::sync::Mutex<()> {
|
||||
struct TracingHarness {
|
||||
_server: MockServer,
|
||||
_codex_home: TempDir,
|
||||
processor: MessageProcessor,
|
||||
processor: Arc<MessageProcessor>,
|
||||
outgoing_rx: mpsc::Receiver<crate::outgoing_message::OutgoingEnvelope>,
|
||||
session: ConnectionSessionState,
|
||||
session: Arc<ConnectionSessionState>,
|
||||
tracing: &'static TestTracing,
|
||||
}
|
||||
|
||||
@@ -129,7 +129,7 @@ impl TracingHarness {
|
||||
_codex_home: codex_home,
|
||||
processor,
|
||||
outgoing_rx,
|
||||
session: ConnectionSessionState::default(),
|
||||
session: Arc::new(ConnectionSessionState::default()),
|
||||
tracing,
|
||||
};
|
||||
|
||||
@@ -152,7 +152,7 @@ impl TracingHarness {
|
||||
/*trace*/ None,
|
||||
)
|
||||
.await;
|
||||
assert!(harness.session.initialized);
|
||||
assert!(harness.session.initialized());
|
||||
|
||||
Ok(harness)
|
||||
}
|
||||
@@ -182,7 +182,7 @@ impl TracingHarness {
|
||||
TEST_CONNECTION_ID,
|
||||
request,
|
||||
AppServerTransport::Stdio,
|
||||
&mut self.session,
|
||||
Arc::clone(&self.session),
|
||||
)
|
||||
.await;
|
||||
read_response(&mut self.outgoing_rx, request_id).await
|
||||
@@ -230,14 +230,14 @@ async fn build_test_config(codex_home: &Path, server_uri: &str) -> Result<Config
|
||||
fn build_test_processor(
|
||||
config: Arc<Config>,
|
||||
) -> (
|
||||
MessageProcessor,
|
||||
Arc<MessageProcessor>,
|
||||
mpsc::Receiver<crate::outgoing_message::OutgoingEnvelope>,
|
||||
) {
|
||||
let (outgoing_tx, outgoing_rx) = mpsc::channel(16);
|
||||
let outgoing = Arc::new(OutgoingMessageSender::new(outgoing_tx));
|
||||
let auth_manager =
|
||||
AuthManager::shared_from_config(config.as_ref(), /*enable_codex_api_key_env*/ false);
|
||||
let processor = MessageProcessor::new(MessageProcessorArgs {
|
||||
let processor = Arc::new(MessageProcessor::new(MessageProcessorArgs {
|
||||
outgoing,
|
||||
arg0_paths: Arg0DispatchPaths::default(),
|
||||
config,
|
||||
@@ -252,7 +252,7 @@ fn build_test_processor(
|
||||
auth_manager,
|
||||
rpc_transport: AppServerRpcTransport::Stdio,
|
||||
remote_control_handle: None,
|
||||
});
|
||||
}));
|
||||
(processor, outgoing_rx)
|
||||
}
|
||||
|
||||
|
||||
@@ -121,7 +121,7 @@ pub(crate) struct ConnectionState {
|
||||
pub(crate) outbound_initialized: Arc<AtomicBool>,
|
||||
pub(crate) outbound_experimental_api_enabled: Arc<AtomicBool>,
|
||||
pub(crate) outbound_opted_out_notification_methods: Arc<RwLock<HashSet<String>>>,
|
||||
pub(crate) session: ConnectionSessionState,
|
||||
pub(crate) session: Arc<ConnectionSessionState>,
|
||||
}
|
||||
|
||||
impl ConnectionState {
|
||||
@@ -134,7 +134,7 @@ impl ConnectionState {
|
||||
outbound_initialized,
|
||||
outbound_experimental_api_enabled,
|
||||
outbound_opted_out_notification_methods,
|
||||
session: ConnectionSessionState::default(),
|
||||
session: Arc::new(ConnectionSessionState::default()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user