Files
codex/codex-rs/app-server/src/in_process.rs
T
ztanner-oai 6562117a81 Preserve process output and control replies under backpressure (#57569)
## Why

Slow consumers can lose process output or block control replies when the in-process event queue fills. Descendants can also keep output streams open after the child exits, leaving clients unable to tell whether captured output is complete.

## What changed

- Require delivery of `process/outputDelta` and `process/exited` notifications, while allowing process request replies to proceed when event delivery is blocked.
- Exclude time spent delivering output to slow clients from the post-exit drain budget, and drain queued bytes after the deadline closes a stream.
- Keep draining after an output byte cap is reached so the child can finish writing; abort each stream's reader when its drain deadline expires.
- Add optional `outputDrainTimedOut` to `process/exited` and generated bindings. Report incomplete output separately from the child's exit code, treating a fully captured requested byte cap as complete.
- Deliver the exit notification before moving Windows console cleanup to a blocking task.

## Testing

Add regression coverage for full event queues, stalled consumers issuing process control requests, slow output delivery, inherited pipes, output caps, Windows console descendants, and compatibility with exit notifications that omit `outputDrainTimedOut`.

GitOrigin-RevId: 04cfeea5daa7c2b833cdec3f1fff7efeb4d1a86e
2026-10-27 21:51:50 +00:00

1441 lines
61 KiB
Rust

//! In-process app-server runtime host for local embedders.
//!
//! This module runs the existing [`MessageProcessor`] and outbound routing logic
//! on Tokio tasks, but replaces socket/stdio transports with bounded in-memory
//! channels. The intent is to preserve app-server semantics while avoiding a
//! process boundary for CLI surfaces that run in the same process.
//!
//! # Lifecycle
//!
//! 1. Construct runtime state with [`InProcessStartArgs`].
//! 2. Call [`start`], which performs the `initialize` / `initialized` handshake
//! internally and returns a ready-to-use [`InProcessClientHandle`].
//! 3. Send requests via [`InProcessClientHandle::request`], notifications via
//! [`InProcessClientHandle::notify`], and consume events via
//! [`InProcessClientHandle::next_event`].
//! 4. Terminate with [`InProcessClientHandle::shutdown`].
//!
//! # Transport model
//!
//! The runtime is transport-local but not protocol-free. Incoming requests are
//! typed [`ClientRequest`] values, yet responses still come back through the
//! same JSON-RPC result envelope that `MessageProcessor` uses for stdio and
//! websocket transports. This keeps in-process behavior aligned with
//! app-server rather than creating a second execution contract.
//!
//! # Backpressure
//!
//! Command submission uses `try_send` and can return `WouldBlock`, while event
//! fanout may drop notifications under saturation. Server requests are never
//! silently abandoned: if they cannot be queued they are failed back into
//! `MessageProcessor` with overload or internal errors so approval flows do
//! not hang indefinitely.
//!
//! # Relationship to `codex-app-server-client`
//!
//! This module provides the low-level runtime handle ([`InProcessClientHandle`]).
//! Higher-level callers (TUI, exec) should go through `codex-app-server-client`,
//! which wraps this module behind a worker task with async request/response
//! helpers, surface-specific startup policy, and bounded shutdown.
use std::collections::HashMap;
use std::collections::HashSet;
use std::collections::hash_map::Entry;
use std::io::Error as IoError;
use std::io::ErrorKind;
use std::io::Result as IoResult;
use std::sync::Arc;
use std::sync::RwLock;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use std::time::Duration;
use crate::analytics_utils::analytics_events_client_from_config;
use crate::config_manager::ConfigManager;
use crate::error_code::OVERLOADED_ERROR_CODE;
use crate::error_code::internal_error;
use crate::error_code::invalid_request;
use crate::message_processor::ConnectionSessionState;
use crate::message_processor::MessageProcessor;
use crate::message_processor::MessageProcessorArgs;
use crate::outgoing_message::ConnectionId;
use crate::outgoing_message::OutgoingEnvelope;
use crate::outgoing_message::OutgoingMessage;
use crate::outgoing_message::OutgoingMessageSender;
use crate::outgoing_message::QueuedOutgoingMessage;
use crate::plugin_config_reload::PluginStartupConfig;
use crate::transport::CHANNEL_CAPACITY;
use crate::transport::OutboundConnectionState;
use crate::transport::route_outgoing_envelope;
pub use bootstrap::EmbeddedNetworkPolicy;
use codex_analytics::AppServerRpcTransport;
use codex_app_server_protocol::AgentMessageDelivery;
use codex_app_server_protocol::ClientNotification;
use codex_app_server_protocol::ClientRequest;
use codex_app_server_protocol::ConfigWarningNotification;
use codex_app_server_protocol::InitializeParams;
use codex_app_server_protocol::ItemCompletedNotification;
use codex_app_server_protocol::JSONRPCErrorError;
use codex_app_server_protocol::RequestId;
use codex_app_server_protocol::Result;
use codex_app_server_protocol::ServerNotification;
use codex_app_server_protocol::ServerRequest;
use codex_app_server_protocol::ThreadItemPayload;
use codex_arg0::Arg0DispatchPaths;
use codex_config::CloudConfigBundleLoader;
use codex_config::LoaderOverrides;
use codex_config::ThreadConfigLoader;
use codex_core::check_execpolicy_for_warnings;
use codex_core::config::Config;
use codex_core::resolve_installation_id;
use codex_exec_server::EnvironmentManager;
use codex_feedback::CodexFeedback;
use codex_protocol::protocol::SessionSource;
pub use codex_rollout::StateDbHandle;
pub use codex_state::log_db::LogDbLayer;
use tokio::sync::mpsc;
use tokio::sync::oneshot;
use tokio::time::timeout;
use toml::Value as TomlValue;
use tracing::warn;
#[path = "in_process_bootstrap.rs"]
mod bootstrap;
const IN_PROCESS_CONNECTION_ID: ConnectionId = ConnectionId(0);
const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(5);
// Covers both bounded runtime drains plus the analytics client's 25-second best-effort flush.
const SHUTDOWN_ACK_TIMEOUT: Duration = Duration::from_secs(35);
/// Default bounded channel capacity for in-process runtime queues.
pub const DEFAULT_IN_PROCESS_CHANNEL_CAPACITY: usize = CHANNEL_CAPACITY;
type PendingClientRequestResponse = std::result::Result<Result, JSONRPCErrorError>;
fn server_notification_requires_delivery(notification: &ServerNotification) -> bool {
matches!(
notification,
ServerNotification::ProcessOutputDelta(_)
| ServerNotification::ProcessExited(_)
| ServerNotification::TurnCompleted(_)
| ServerNotification::ThreadQueueChanged(_)
| ServerNotification::ThreadSettingsUpdated(_)
| ServerNotification::ThreadAttachmentUpdated(_)
| ServerNotification::ExternalAgentConfigImportCompleted(_)
| ServerNotification::ItemCompleted(ItemCompletedNotification {
item: codex_app_server_protocol::ThreadItem {
payload: ThreadItemPayload::AgentMessage {
delivery: Some(AgentMessageDelivery::Async),
..
},
..
},
..
})
)
}
/// Input needed to start an in-process app-server runtime.
///
/// These fields mirror the pieces of ambient process state that stdio and
/// websocket transports normally assemble before `MessageProcessor` starts.
#[derive(Clone)]
pub struct InProcessStartArgs {
/// Resolved argv0 dispatch paths used by command execution internals.
pub arg0_paths: Arg0DispatchPaths,
/// Shared base config used to initialize core components.
pub config: Arc<Config>,
/// CLI config overrides that are already parsed into TOML values.
pub cli_overrides: Vec<(String, TomlValue)>,
/// Loader override knobs used by config API paths.
pub loader_overrides: LoaderOverrides,
/// Whether config API paths should reject unknown config fields.
pub strict_config: bool,
/// Preloaded cloud config bundle provider.
pub cloud_config_bundle: CloudConfigBundleLoader,
/// Policy shared with transports created by the embedder before startup.
pub embedded_network_policy: EmbeddedNetworkPolicy,
/// Loader used to fetch typed thread config sources before a thread starts.
pub thread_config_loader: Arc<dyn ThreadConfigLoader>,
/// Feedback sink used by app-server/core telemetry and logs.
pub feedback: CodexFeedback,
/// SQLite tracing layer used to flush recently emitted logs before feedback upload.
pub log_db: Option<LogDbLayer>,
/// Process-wide SQLite state handle shared with embedded app-server consumers.
pub state_db: Option<StateDbHandle>,
/// Environment manager used by core execution and filesystem operations.
pub environment_manager: Arc<EnvironmentManager>,
/// Startup warnings emitted after initialize succeeds.
pub config_warnings: Vec<ConfigWarningNotification>,
/// Session source stamped into thread/session metadata.
pub session_source: SessionSource,
/// Whether serving auth should honor `CODEX_API_KEY`; workspace policy still uses stored auth.
pub enable_codex_api_key_env: bool,
/// Initialize params used for initial handshake.
pub initialize: InitializeParams,
/// Capacity used for all runtime queues (clamped to at least 1).
pub channel_capacity: usize,
}
/// Event emitted from the app-server to the in-process client.
///
/// [`Lagged`](Self::Lagged) is a transport health marker, not an application
/// event — it signals that the consumer fell behind and some events were dropped.
#[derive(Debug, Clone)]
pub enum InProcessServerEvent {
/// Server request that requires client response/rejection.
ServerRequest(Box<ServerRequest>),
/// App-server notification directed to the embedded client.
ServerNotification(Box<ServerNotification>),
/// Indicates one or more events were dropped due to backpressure.
Lagged { skipped: usize },
}
/// Internal message sent from [`InProcessClientHandle`] methods to the runtime task.
///
/// Requests carry a oneshot sender for the response; notifications and server-request
/// replies are fire-and-forget from the caller's perspective (transport errors are
/// caught by `try_send` on the outer channel).
enum InProcessClientMessage {
Request {
request: Box<ClientRequest>,
response_tx: oneshot::Sender<PendingClientRequestResponse>,
cancellation: tokio_util::sync::CancellationToken,
},
Notification {
notification: ClientNotification,
},
ServerRequestResponse {
request_id: RequestId,
result: Result,
},
ServerRequestError {
request_id: RequestId,
error: JSONRPCErrorError,
},
Shutdown {
done_tx: oneshot::Sender<()>,
},
}
enum ProcessorCommand {
Request(Box<ClientRequest>, tokio_util::sync::CancellationToken),
Notification(ClientNotification),
}
#[derive(Clone)]
pub struct InProcessClientSender {
client_tx: mpsc::Sender<InProcessClientMessage>,
}
impl InProcessClientSender {
pub async fn request(&self, request: ClientRequest) -> IoResult<PendingClientRequestResponse> {
let (response_tx, response_rx) = oneshot::channel();
let cancellation = tokio_util::sync::CancellationToken::new();
let _cancel_on_drop = cancellation.clone().drop_guard();
self.try_send_client_message(InProcessClientMessage::Request {
request: Box::new(request),
response_tx,
cancellation,
})?;
response_rx.await.map_err(|err| {
IoError::new(
ErrorKind::BrokenPipe,
format!("in-process request response channel closed: {err}"),
)
})
}
pub fn notify(&self, notification: ClientNotification) -> IoResult<()> {
self.try_send_client_message(InProcessClientMessage::Notification { notification })
}
pub fn respond_to_server_request(&self, request_id: RequestId, result: Result) -> IoResult<()> {
self.try_send_client_message(InProcessClientMessage::ServerRequestResponse {
request_id,
result,
})
}
pub fn fail_server_request(
&self,
request_id: RequestId,
error: JSONRPCErrorError,
) -> IoResult<()> {
self.try_send_client_message(InProcessClientMessage::ServerRequestError {
request_id,
error,
})
}
fn try_send_client_message(&self, message: InProcessClientMessage) -> IoResult<()> {
match self.client_tx.try_send(message) {
Ok(()) => Ok(()),
Err(mpsc::error::TrySendError::Full(_)) => Err(IoError::new(
ErrorKind::WouldBlock,
"in-process app-server client queue is full",
)),
Err(mpsc::error::TrySendError::Closed(_)) => Err(IoError::new(
ErrorKind::BrokenPipe,
"in-process app-server runtime is closed",
)),
}
}
}
/// Handle used by an in-process client to call app-server and consume events.
///
/// This is the low-level runtime handle. Higher-level callers should usually go
/// through `codex-app-server-client`, which adds worker-task buffering,
/// request/response helpers, and surface-specific startup policy.
pub struct InProcessClientHandle {
client: InProcessClientSender,
event_rx: mpsc::Receiver<InProcessServerEvent>,
runtime_handle: tokio::task::JoinHandle<()>,
#[cfg(test)]
_test_codex_home: Option<tempfile::TempDir>,
}
impl InProcessClientHandle {
/// Sends a typed client request into the in-process runtime.
///
/// The returned value is a transport-level `IoResult` containing either a
/// JSON-RPC success payload or JSON-RPC error payload. Callers must keep
/// request IDs unique among concurrent requests; reusing an in-flight ID
/// produces an `INVALID_REQUEST` response and can make request routing
/// ambiguous in the caller.
pub async fn request(&self, request: ClientRequest) -> IoResult<PendingClientRequestResponse> {
self.client.request(request).await
}
/// Sends a typed client notification into the in-process runtime.
///
/// Notifications do not have an application-level response. Transport
/// errors indicate queue saturation or closed runtime.
pub fn notify(&self, notification: ClientNotification) -> IoResult<()> {
self.client.notify(notification)
}
/// Resolves a pending [`ServerRequest`](InProcessServerEvent::ServerRequest).
///
/// This should be used only with request IDs received from the current
/// runtime event stream; sending arbitrary IDs has no effect on app-server
/// state and can mask a stuck approval flow in the caller.
pub fn respond_to_server_request(&self, request_id: RequestId, result: Result) -> IoResult<()> {
self.client.respond_to_server_request(request_id, result)
}
/// Rejects a pending [`ServerRequest`](InProcessServerEvent::ServerRequest).
///
/// Use this when the embedder cannot satisfy a server request; leaving
/// requests unanswered can stall turn progress.
pub fn fail_server_request(
&self,
request_id: RequestId,
error: JSONRPCErrorError,
) -> IoResult<()> {
self.client.fail_server_request(request_id, error)
}
/// Receives the next server event from the in-process runtime.
///
/// Returns `None` when the runtime task exits and no more events are
/// available.
pub async fn next_event(&mut self) -> Option<InProcessServerEvent> {
self.event_rx.recv().await
}
/// Requests runtime shutdown and waits for worker termination.
///
/// Shutdown is bounded by internal timeouts and may abort background tasks
/// if graceful drain does not complete in time.
pub async fn shutdown(self) -> IoResult<()> {
// Shutdown owns the only event receiver, so nobody can drain it anymore.
// Unblock notification delivery before waiting for space in the client queue.
drop(self.event_rx);
let mut runtime_handle = self.runtime_handle;
let (done_tx, done_rx) = oneshot::channel();
if self
.client
.client_tx
.send(InProcessClientMessage::Shutdown { done_tx })
.await
.is_ok()
{
let _ = timeout(SHUTDOWN_ACK_TIMEOUT, done_rx).await;
}
if let Err(_elapsed) = timeout(SHUTDOWN_TIMEOUT, &mut runtime_handle).await {
runtime_handle.abort();
let _ = runtime_handle.await;
}
Ok(())
}
pub fn sender(&self) -> InProcessClientSender {
self.client.clone()
}
}
/// Starts an in-process app-server runtime and performs initialize handshake.
///
/// This function sends `initialize` followed by `initialized` before returning
/// the handle, so callers receive a ready-to-use runtime. If initialize fails,
/// the runtime is shut down and an `InvalidData` error is returned.
pub async fn start(mut args: InProcessStartArgs) -> IoResult<InProcessClientHandle> {
if let Ok(Some(err)) = check_execpolicy_for_warnings(&args.config.config_layer_stack).await {
let (path, range) = crate::exec_policy_warning_location(&err);
args.config_warnings.push(ConfigWarningNotification {
summary: "Error parsing rules; custom rules not applied.".to_string(),
details: Some(err.to_string()),
path,
range,
});
}
let initialize = args.initialize.clone();
let client = Box::pin(start_uninitialized(args)).await?;
let initialize_response = client
.request(ClientRequest::Initialize {
request_id: RequestId::Integer(0),
params: initialize,
})
.await?;
if let Err(error) = initialize_response {
let _ = client.shutdown().await;
return Err(IoError::new(
ErrorKind::InvalidData,
format!("in-process initialize failed: {}", error.message),
));
}
client.notify(ClientNotification::Initialized)?;
Ok(client)
}
async fn run_outbound_router(
mut outgoing_rx: mpsc::Receiver<OutgoingEnvelope>,
mut outbound_connections: HashMap<ConnectionId, OutboundConnectionState>,
mut shutdown_rx: oneshot::Receiver<()>,
) {
loop {
tokio::select! {
biased;
_ = &mut shutdown_rx => break,
envelope = outgoing_rx.recv() => {
let Some(envelope) = envelope else {
break;
};
route_outgoing_envelope(&mut outbound_connections, envelope).await;
}
}
}
}
async fn start_uninitialized(mut args: InProcessStartArgs) -> IoResult<InProcessClientHandle> {
let config_manager = ConfigManager::new(
args.config.codex_home.to_path_buf(),
args.cli_overrides,
args.loader_overrides,
args.strict_config,
args.cloud_config_bundle,
args.arg0_paths.clone(),
args.thread_config_loader,
)
.with_embedded_network_policy(args.embedded_network_policy);
// Keep the caller's project context independent of any fallback policy config.
let startup_cwd =
(!args.config.config_layer_stack.is_projectless()).then(|| args.config.cwd.to_path_buf());
let auth_manager = bootstrap::configure(
&config_manager,
&mut args.config,
args.enable_codex_api_key_env,
)
.await?;
let channel_capacity = args.channel_capacity.max(1);
let installation_id = resolve_installation_id(&args.config.codex_home).await?;
let (client_tx, mut client_rx) = mpsc::channel::<InProcessClientMessage>(channel_capacity);
let (event_tx, event_rx) = mpsc::channel::<InProcessServerEvent>(channel_capacity);
let runtime_handle = tokio::spawn(async move {
let (outgoing_tx, outgoing_rx) = mpsc::channel::<OutgoingEnvelope>(channel_capacity);
let analytics_events_client =
analytics_events_client_from_config(Arc::clone(&auth_manager), args.config.as_ref());
let analytics_events_flush_client = analytics_events_client.clone();
// A client can be waiting for process/kill while nobody is polling next_event.
let (process_reply_tx, mut process_reply_rx) = mpsc::channel(channel_capacity);
let outgoing_message_sender = Arc::new(
OutgoingMessageSender::new(outgoing_tx, analytics_events_client.clone())
.with_process_reply_sender(process_reply_tx),
);
let log_write_warning = crate::log_write_warning::LogWriteWarningReporter::new(
args.feedback.clone(),
&outgoing_message_sender,
&args.config,
);
if let Some(log_db) = &args.log_db {
log_db.set_failure_reporter(log_write_warning.clone());
if log_db.has_write_failure() {
log_write_warning.notify_failure();
}
}
let (writer_tx, mut writer_rx) = mpsc::channel::<QueuedOutgoingMessage>(channel_capacity);
let outbound_initialized = Arc::new(AtomicBool::new(false));
let outbound_experimental_api_enabled = Arc::new(AtomicBool::new(false));
let outbound_opted_out_notification_methods = Arc::new(RwLock::new(HashSet::new()));
let mut outbound_connections = HashMap::<ConnectionId, OutboundConnectionState>::new();
outbound_connections.insert(
IN_PROCESS_CONNECTION_ID,
OutboundConnectionState::new(
writer_tx,
Arc::clone(&outbound_initialized),
Arc::clone(&outbound_experimental_api_enabled),
Arc::clone(&outbound_opted_out_notification_methods),
/*disconnect_sender*/ None,
),
);
let (outbound_shutdown_tx, outbound_shutdown_rx) = oneshot::channel();
let mut outbound_handle = tokio::spawn(run_outbound_router(
outgoing_rx,
outbound_connections,
outbound_shutdown_rx,
));
let processor_outgoing = Arc::clone(&outgoing_message_sender);
let (processor_tx, mut processor_rx) = mpsc::channel::<ProcessorCommand>(channel_capacity);
let mut processor_handle = tokio::spawn(async move {
let (startup_warnings_tx, mut startup_warnings_rx) = mpsc::unbounded_channel();
let processor = Arc::new(MessageProcessor::new(MessageProcessorArgs {
cloud_skill_provider: None,
outgoing: Arc::clone(&processor_outgoing),
analytics_events_client,
arg0_paths: args.arg0_paths,
config: args.config,
startup_cwd,
config_manager,
environment_manager: args.environment_manager,
feedback: args.feedback,
log_db: args.log_db,
state_db: args.state_db,
config_warnings: args.config_warnings,
startup_warnings_tx,
session_source: args.session_source,
user_verification: Arc::new(crate::user_verification::Service::new(Arc::clone(
&auth_manager,
))),
auth_manager,
installation_id,
code_mode_session_provider: None,
rpc_transport: AppServerRpcTransport::InProcess,
remote_control_handle: None,
plugin_startup_tasks: Some(PluginStartupConfig::Current),
}));
let mut thread_created_rx = processor.thread_created_receiver();
let session = Arc::new(ConnectionSessionState::new(
crate::transport::ConnectionOrigin::InProcess,
));
let mut listen_for_threads = true;
loop {
tokio::select! {
command = processor_rx.recv() => {
match command {
Some(ProcessorCommand::Request(request, cancellation)) => {
let was_initialized = session.initialized();
processor
.process_client_request(
IN_PROCESS_CONNECTION_ID,
*request,
Arc::clone(&session),
&outbound_initialized,
cancellation,
)
.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 =
opted_out_notification_methods_snapshot;
} else {
warn!("failed to update outbound opted-out notifications");
}
outbound_experimental_api_enabled.store(
experimental_api_enabled,
Ordering::Release,
);
if !was_initialized && is_initialized {
processor.send_initialize_notifications().await;
}
}
Some(ProcessorCommand::Notification(notification)) => {
processor.process_client_notification(notification).await;
}
None => {
break;
}
}
}
Some(message) = startup_warnings_rx.recv(), if session.initialized() => {
crate::startup_warnings::send_warning(&processor_outgoing, message).await;
}
created = thread_created_rx.recv(), if listen_for_threads => {
match created {
Ok(thread_id) => {
let connection_ids = if session.initialized() {
vec![IN_PROCESS_CONNECTION_ID]
} else {
Vec::<ConnectionId>::new()
};
processor
.try_attach_thread_listener(thread_id, connection_ids)
.await;
}
Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => {
warn!("thread_created receiver lagged; skipping resync");
}
Err(tokio::sync::broadcast::error::RecvError::Closed) => {
listen_for_threads = false;
}
}
}
}
}
processor.clear_runtime_references();
processor.cancel_active_login().await;
processor
.connection_closed(IN_PROCESS_CONNECTION_ID, &session)
.await;
processor.clear_all_thread_listeners().await;
processor.drain_background_tasks().await;
processor.shutdown_threads().await;
});
let mut pending_request_responses =
HashMap::<RequestId, oneshot::Sender<PendingClientRequestResponse>>::new();
let mut shutdown_ack = None;
let mut pending_event: Option<(InProcessServerEvent, Option<oneshot::Sender<()>>)> = None;
loop {
tokio::select! {
message = client_rx.recv() => {
match message {
Some(InProcessClientMessage::Request { request, response_tx, cancellation }) => {
let request = *request;
let request_id = request.id().clone();
match pending_request_responses.entry(request_id.clone()) {
Entry::Vacant(entry) => {
entry.insert(response_tx);
}
Entry::Occupied(_) => {
let _ = response_tx.send(Err(invalid_request(format!(
"duplicate request id: {request_id:?}"
))));
continue;
}
}
match processor_tx.try_send(ProcessorCommand::Request(Box::new(request), cancellation)) {
Ok(()) => {}
Err(mpsc::error::TrySendError::Full(_)) => {
if let Some(response_tx) =
pending_request_responses.remove(&request_id)
{
let _ = response_tx.send(Err(JSONRPCErrorError {
code: OVERLOADED_ERROR_CODE,
message: "in-process app-server request queue is full"
.to_string(),
data: None,
}));
}
}
Err(mpsc::error::TrySendError::Closed(_)) => {
if let Some(response_tx) =
pending_request_responses.remove(&request_id)
{
let _ = response_tx.send(Err(internal_error(
"in-process app-server request processor is closed",
)));
}
break;
}
}
}
Some(InProcessClientMessage::Notification { notification }) => {
match processor_tx.try_send(ProcessorCommand::Notification(notification)) {
Ok(()) => {}
Err(mpsc::error::TrySendError::Full(_)) => {
warn!("dropping in-process client notification (queue full)");
}
Err(mpsc::error::TrySendError::Closed(_)) => {
break;
}
}
}
Some(InProcessClientMessage::ServerRequestResponse { request_id, result }) => {
outgoing_message_sender
.notify_client_response(IN_PROCESS_CONNECTION_ID, request_id, result)
.await;
}
Some(InProcessClientMessage::ServerRequestError { request_id, error }) => {
outgoing_message_sender
.notify_client_error(IN_PROCESS_CONNECTION_ID, request_id, error)
.await;
}
Some(InProcessClientMessage::Shutdown { done_tx }) => {
shutdown_ack = Some(done_tx);
break;
}
None => {
break;
}
}
}
permit = event_tx.reserve(), if pending_event.is_some() => {
let Ok(permit) = permit else { break };
let Some((event, written)) = pending_event.take() else { continue };
permit.send(event);
if let Some(written) = written {
let _ = written.send(());
}
}
queued_message = async {
tokio::select! {
reply = process_reply_rx.recv() => match reply {
Some(OutgoingEnvelope::ToConnection { message, write_complete_tx, .. }) => {
Some(QueuedOutgoingMessage { message, write_complete_tx })
}
Some(OutgoingEnvelope::Broadcast { .. }) => unreachable!("replies are connection scoped"),
None => None,
},
// Leave ordinary replies and events queued until the pending event fits.
event = writer_rx.recv(), if pending_event.is_none() => event,
}
} => {
let Some(queued_message) = queued_message else {
break;
};
let outgoing_message = queued_message.message;
match outgoing_message {
OutgoingMessage::Response(response) => {
if let Some(response_tx) = pending_request_responses.remove(&response.id) {
let result = serde_json::to_value(response.result).map_err(|err| {
internal_error(format!("failed to serialize response: {err}"))
});
let _ = response_tx.send(result);
} else {
warn!(
request_id = ?response.id,
"dropping unmatched in-process response"
);
}
}
OutgoingMessage::Error(error) => {
if let Some(response_tx) = pending_request_responses.remove(&error.id) {
let _ = response_tx.send(Err(error.error));
} else {
warn!(
request_id = ?error.id,
"dropping unmatched in-process error response"
);
}
}
OutgoingMessage::Request(request) => {
// Send directly to avoid cloning; on failure the
// original value is returned inside the error.
if let Err(send_error) = event_tx
.try_send(InProcessServerEvent::ServerRequest(Box::new(request)))
{
let (error, inner) = match send_error {
mpsc::error::TrySendError::Full(inner) => (
JSONRPCErrorError {
code: OVERLOADED_ERROR_CODE,
message:
"in-process server request queue is full".to_string(),
data: None,
},
inner,
),
mpsc::error::TrySendError::Closed(inner) => (
internal_error(
"in-process server request consumer is closed",
),
inner,
),
};
let request_id = match inner {
InProcessServerEvent::ServerRequest(req) => req.id().clone(),
_ => unreachable!("we just sent a ServerRequest variant"),
};
outgoing_message_sender
.notify_client_error(IN_PROCESS_CONNECTION_ID, request_id, error)
.await;
}
}
OutgoingMessage::AppServerNotification(envelope) => {
let notification = envelope.notification;
if server_notification_requires_delivery(&notification) {
pending_event = Some((
InProcessServerEvent::ServerNotification(Box::new(notification)),
queued_message.write_complete_tx,
));
continue;
}
if let Err(send_error) =
event_tx.try_send(InProcessServerEvent::ServerNotification(
Box::new(notification),
))
{
match send_error {
mpsc::error::TrySendError::Full(_) => {
warn!("dropping in-process server notification (queue full)");
continue;
}
mpsc::error::TrySendError::Closed(_) => {
break;
}
}
}
}
}
if let Some(write_complete_tx) = queued_message.write_complete_tx {
let _ = write_complete_tx.send(());
}
}
}
}
drop(pending_event);
drop(writer_rx);
drop(process_reply_rx);
drop(processor_tx);
outgoing_message_sender
.cancel_all_requests(Some(internal_error(
"in-process app-server runtime is shutting down",
)))
.await;
// Detached processor work can retain outgoing senders, so channel
// closure alone cannot be used to shut down the outbound router.
drop(outgoing_message_sender);
for (_, response_tx) in pending_request_responses {
let _ = response_tx.send(Err(internal_error(
"in-process app-server runtime is shutting down",
)));
}
if let Err(_elapsed) = timeout(SHUTDOWN_TIMEOUT, &mut processor_handle).await {
processor_handle.abort();
let _ = processor_handle.await;
}
let _ = outbound_shutdown_tx.send(());
if let Err(_elapsed) = timeout(SHUTDOWN_TIMEOUT, &mut outbound_handle).await {
outbound_handle.abort();
let _ = outbound_handle.await;
}
analytics_events_flush_client.flush().await;
if let Some(done_tx) = shutdown_ack {
let _ = done_tx.send(());
}
});
Ok(InProcessClientHandle {
client: InProcessClientSender { client_tx },
event_rx,
runtime_handle,
#[cfg(test)]
_test_codex_home: None,
})
}
#[cfg(test)]
mod tests {
use super::*;
use codex_app_server_protocol::ClientInfo;
use codex_app_server_protocol::ConfigRequirementsReadResponse;
use codex_app_server_protocol::ExternalAgentConfigImportCompletedNotification;
use codex_app_server_protocol::SessionSource as ApiSessionSource;
use codex_app_server_protocol::ThreadAttachmentAddParams;
use codex_app_server_protocol::ThreadAttachmentOperation;
use codex_app_server_protocol::ThreadAttachmentRemoveParams;
use codex_app_server_protocol::ThreadAttachmentUpdatedNotification;
use codex_app_server_protocol::ThreadQueueChangedNotification;
use codex_app_server_protocol::ThreadSetNameParams;
use codex_app_server_protocol::ThreadStartParams;
use codex_app_server_protocol::ThreadStartResponse;
use codex_app_server_protocol::Turn;
use codex_app_server_protocol::TurnCompletedNotification;
use codex_app_server_protocol::TurnItemsView;
use codex_app_server_protocol::TurnStatus;
use codex_core::config::ConfigBuilder;
use futures::FutureExt;
use pretty_assertions::assert_eq;
use std::path::Path;
use tempfile::TempDir;
async fn build_test_config(codex_home: &Path) -> Config {
match ConfigBuilder::default()
.codex_home(codex_home.to_path_buf())
.build()
.await
{
Ok(config) => config,
Err(_) => Config::load_default_with_cli_overrides_for_codex_home(
codex_home.to_path_buf(),
Vec::new(),
)
.await
.expect("default config should load"),
}
}
async fn start_test_client_with_options(
session_source: SessionSource,
channel_capacity: usize,
capabilities: Option<codex_app_server_protocol::InitializeCapabilities>,
) -> InProcessClientHandle {
let codex_home = TempDir::new().expect("temp dir");
let config = Arc::new(build_test_config(codex_home.path()).await);
let state_db = codex_rollout::state_db::try_init(config.as_ref())
.await
.expect("state db should initialize for in-process test");
let args = InProcessStartArgs {
arg0_paths: Arg0DispatchPaths::default(),
config,
cli_overrides: Vec::new(),
loader_overrides: LoaderOverrides::default(),
strict_config: false,
cloud_config_bundle: CloudConfigBundleLoader::default(),
embedded_network_policy: Default::default(),
thread_config_loader: Arc::new(codex_config::NoopThreadConfigLoader),
feedback: CodexFeedback::new(),
log_db: None,
state_db: Some(state_db),
environment_manager: Arc::new(EnvironmentManager::default_for_tests()),
config_warnings: Vec::new(),
session_source,
enable_codex_api_key_env: false,
initialize: InitializeParams {
client_info: ClientInfo {
name: "codex-in-process-test".to_string(),
title: None,
version: "0.0.0".to_string(),
},
capabilities,
},
channel_capacity,
};
let mut client = start(args).await.expect("in-process runtime should start");
client._test_codex_home = Some(codex_home);
client
}
async fn start_test_client(session_source: SessionSource) -> InProcessClientHandle {
start_test_client_with_options(
session_source,
DEFAULT_IN_PROCESS_CHANNEL_CAPACITY,
/*capabilities*/ None,
)
.await
}
#[tokio::test]
async fn in_process_start_initializes_and_handles_typed_v2_request() {
let client = start_test_client(SessionSource::Cli).await;
let response = client
.request(ClientRequest::ConfigRequirementsRead {
request_id: RequestId::Integer(1),
params: None,
})
.await
.expect("request transport should work")
.expect("request should succeed");
assert!(response.is_object());
let _parsed: ConfigRequirementsReadResponse =
serde_json::from_value(response).expect("response should match v2 schema");
client
.shutdown()
.await
.expect("in-process runtime should shutdown cleanly");
}
/// Importing a project makes its change visible to an embedder before the import reply.
#[tokio::test]
async fn project_import_delivers_change_before_reply() {
let mut client = start_test_client_with_options(
SessionSource::Cli,
DEFAULT_IN_PROCESS_CHANNEL_CAPACITY,
Some(codex_app_server_protocol::InitializeCapabilities {
experimental_api: true,
..Default::default()
}),
)
.await;
let params = codex_app_server_protocol::ProjectImportParams {
name: "Imported".into(),
roots: Vec::new(),
metadata: None,
threads: None,
idempotency_key: "test-import".into(),
};
client
.request(ClientRequest::ProjectImport {
request_id: RequestId::Integer(1),
params,
})
.await
.unwrap()
.unwrap();
loop {
let event = client
.event_rx
.try_recv()
.expect("project/changed must already be queued");
if let InProcessServerEvent::ServerNotification(notification) = event
&& matches!(*notification, ServerNotification::ProjectChanged(_))
{
break;
}
}
client.shutdown().await.unwrap();
}
#[tokio::test]
async fn in_process_start_uses_requested_session_source_for_thread_start() {
for (requested_source, expected_source) in [
(SessionSource::Cli, ApiSessionSource::Cli),
(SessionSource::Exec, ApiSessionSource::Exec),
] {
let client = start_test_client(requested_source).await;
let response = client
.request(ClientRequest::ThreadStart {
request_id: RequestId::Integer(2),
params: ThreadStartParams {
ephemeral: Some(true),
..ThreadStartParams::default()
},
})
.await
.expect("request transport should work")
.expect("thread/start should succeed");
let parsed: ThreadStartResponse =
serde_json::from_value(response).expect("thread/start response should parse");
assert_eq!(parsed.thread.source, expected_source);
client
.shutdown()
.await
.expect("in-process runtime should shutdown cleanly");
}
}
#[tokio::test]
async fn in_process_start_clamps_zero_channel_capacity() {
let client = start_test_client_with_options(
SessionSource::Cli,
/*channel_capacity*/ 0,
/*capabilities*/ None,
)
.await;
let response = loop {
match client
.request(ClientRequest::ConfigRequirementsRead {
request_id: RequestId::Integer(4),
params: None,
})
.await
{
Ok(response) => break response.expect("request should succeed"),
Err(err) if err.kind() == std::io::ErrorKind::WouldBlock => {
tokio::task::yield_now().await;
}
Err(err) => panic!("request transport should work: {err}"),
}
};
let _parsed: ConfigRequirementsReadResponse =
serde_json::from_value(response).expect("response should match v2 schema");
client
.shutdown()
.await
.expect("in-process runtime should shutdown cleanly");
}
#[tokio::test(start_paused = true)]
async fn in_process_outbound_router_shutdown_does_not_wait_for_retained_sender() {
let (outgoing_tx, outgoing_rx) = mpsc::channel(/*buffer*/ 1);
let retained_outgoing_tx = outgoing_tx.clone();
drop(outgoing_tx);
let (shutdown_tx, shutdown_rx) = oneshot::channel();
let mut outbound_handle = tokio::spawn(run_outbound_router(
outgoing_rx,
HashMap::new(),
shutdown_rx,
));
assert!(!retained_outgoing_tx.is_closed());
shutdown_tx
.send(())
.expect("outbound router should accept explicit shutdown");
timeout(SHUTDOWN_TIMEOUT, &mut outbound_handle)
.await
.expect("outbound router should not wait for its retained sender")
.expect("outbound router should complete successfully");
assert!(retained_outgoing_tx.is_closed());
}
#[tokio::test(start_paused = true)]
async fn in_process_shutdown_waits_for_analytics_flush_budget() {
let (client_tx, mut client_rx) = mpsc::channel(/*buffer*/ 1);
let (_event_tx, event_rx) = mpsc::channel(/*buffer*/ 1);
let completed = Arc::new(AtomicBool::new(false));
let runtime_completed = Arc::clone(&completed);
let runtime_handle = tokio::spawn(async move {
let done_tx = match client_rx.recv().await {
Some(InProcessClientMessage::Shutdown { done_tx }) => done_tx,
_ => panic!("expected in-process shutdown request"),
};
tokio::time::sleep(SHUTDOWN_TIMEOUT + SHUTDOWN_TIMEOUT + Duration::from_secs(24)).await;
runtime_completed.store(true, Ordering::Release);
let _ = done_tx.send(());
});
let client = InProcessClientHandle {
client: InProcessClientSender { client_tx },
event_rx,
runtime_handle,
_test_codex_home: None,
};
client
.shutdown()
.await
.expect("in-process runtime should shutdown cleanly");
assert!(completed.load(Ordering::Acquire));
}
#[test]
fn guaranteed_delivery_helpers_cover_required_server_notifications() {
assert!(server_notification_requires_delivery(
&ServerNotification::TurnCompleted(TurnCompletedNotification {
thread_id: "thread-1".to_string(),
turn: Turn {
id: "turn-1".to_string(),
root_turn_id: None,
items: Vec::new(),
items_view: TurnItemsView::NotLoaded,
status: TurnStatus::Completed,
error: None,
started_at: None,
completed_at: Some(0),
duration_ms: None,
},
})
));
assert!(server_notification_requires_delivery(
&ServerNotification::ThreadQueueChanged(ThreadQueueChangedNotification {
thread_id: "thread-1".to_string(),
})
));
assert!(server_notification_requires_delivery(
&ServerNotification::ExternalAgentConfigImportCompleted(
ExternalAgentConfigImportCompletedNotification {
import_id: "import".to_string(),
item_type_results: Vec::new(),
},
)
));
assert!(server_notification_requires_delivery(
&ServerNotification::ThreadAttachmentUpdated(ThreadAttachmentUpdatedNotification {
thread_id: "thread-1".to_string(),
attachment_type: "pull_request".to_string(),
identity_key: r#"["github.com","openai","codex",123]"#.to_string(),
attachment_id: "attachment-1".to_string(),
operation: ThreadAttachmentOperation::Deleted,
})
));
assert!(server_notification_requires_delivery(
&ServerNotification::ItemCompleted(ItemCompletedNotification::new(
"thread-1".to_string(),
"turn-1".to_string(),
ThreadItemPayload::AgentMessage {
id: "item-1".to_string(),
text: "Still working".to_string(),
phase: None,
memory_citation: None,
delivery: Some(AgentMessageDelivery::Async),
questions: None,
},
/*started_at_ms*/ None,
/*completed_at_ms*/ 0
))
));
}
#[cfg(unix)]
async fn spawn_streaming_test_process(
cwd: &Path,
handle: &str,
script: &str,
) -> InProcessClientHandle {
use codex_app_server_protocol::InitializeCapabilities;
use codex_app_server_protocol::ProcessSpawnParams;
let client = start_test_client_with_options(
SessionSource::Cli,
/*channel_capacity*/ 2,
Some(InitializeCapabilities {
experimental_api: true,
..Default::default()
}),
)
.await;
client
.request(ClientRequest::ProcessSpawn {
request_id: RequestId::Integer(1),
params: ProcessSpawnParams {
process_handle: handle.into(),
command: vec!["/bin/sh".into(), "-c".into(), script.into()],
cwd: cwd.try_into().unwrap(),
tty: false,
stream_stdin: false,
stream_stdout_stderr: true,
output_bytes_cap: Some(None),
timeout_ms: None,
env: None,
size: None,
},
})
.await
.unwrap()
.unwrap();
client
}
/// A kill after child exit must promptly report it has exited while the client is still
/// waiting to consume the complete download and its successful exit notification.
#[cfg(unix)]
#[tokio::test]
async fn process_output_and_exit_survive_a_full_in_process_event_queue() {
use base64::Engine;
use base64::engine::general_purpose::STANDARD;
use codex_app_server_protocol::ProcessKillParams;
let dir = TempDir::new().unwrap();
let contents: Vec<u8> = (0..128 * 1024u32).flat_map(u32::to_le_bytes).collect();
std::fs::write(dir.path().join("input"), &contents).unwrap();
let mut client = spawn_streaming_test_process(
dir.path(),
"download",
concat!(
"p=$$; (while kill -0 $p 2>/dev/null; do sleep 0.05; done; touch exited) ",
"</dev/null >/dev/null 2>&1 & exec cat input",
),
)
.await;
timeout(Duration::from_secs(10), async {
// The pipe readers can hold this payload while delivery is stalled.
while client.event_rx.capacity() != 0 || !dir.path().join("exited").exists() {
tokio::task::yield_now().await;
}
let error = client
.request(ClientRequest::ProcessKill {
request_id: RequestId::Integer(2),
params: ProcessKillParams {
process_handle: "download".into(),
},
})
.await
.unwrap()
.unwrap_err();
assert_eq!(error.message, "process \"download\" is no longer running");
let mut received = Vec::new();
loop {
if let InProcessServerEvent::ServerNotification(notification) =
client.next_event().await.unwrap()
{
match *notification {
ServerNotification::ProcessOutputDelta(delta) => {
received.extend(STANDARD.decode(delta.delta_base64).unwrap());
}
ServerNotification::ProcessExited(exit) => {
assert_eq!(exit.exit_code, 0);
assert_eq!(received.len(), contents.len());
assert_eq!(received, contents);
break;
}
_ => {}
}
}
}
})
.await
.expect("process output and exit must both reach the consumer");
client.shutdown().await.unwrap();
}
/// Unread output must not block process control replies or a subsequent kill on the same handle.
#[cfg(unix)]
#[tokio::test]
async fn in_process_kill_completes_while_process_output_is_stalled() {
use codex_app_server_protocol::ProcessKillParams;
use codex_app_server_protocol::ProcessResizePtyParams;
use codex_app_server_protocol::ProcessTerminalSize;
let dir = TempDir::new().unwrap();
let client = spawn_streaming_test_process(
dir.path(),
"writer",
"head -c 524288 /dev/zero; head -c 524288 /dev/zero >&2; touch written; exec sleep 30",
)
.await;
timeout(Duration::from_secs(10), async {
while client.event_rx.capacity() != 0 || !dir.path().join("written").exists() {
tokio::task::yield_now().await;
}
// A resize on a pipe fails, but its error still has to release the handle for kill.
client
.request(ClientRequest::ProcessResizePty {
request_id: RequestId::Integer(2),
params: ProcessResizePtyParams {
process_handle: "writer".into(),
size: ProcessTerminalSize { rows: 24, cols: 80 },
},
})
.await
.unwrap()
.unwrap_err();
client
.request(ClientRequest::ProcessKill {
request_id: RequestId::Integer(3),
params: ProcessKillParams {
process_handle: "writer".into(),
},
})
.await
.unwrap()
.unwrap();
})
.await
.expect("kill must finish before the client resumes reading output");
client.shutdown().await.unwrap();
}
#[tokio::test]
async fn in_process_shutdown_completes_with_full_request_and_event_queues() {
let capacity = DEFAULT_IN_PROCESS_CHANNEL_CAPACITY;
let mut client = start_test_client(SessionSource::Cli).await;
let response = client
.request(ClientRequest::ThreadStart {
request_id: RequestId::Integer(1),
params: ThreadStartParams::default(),
})
.await
.expect("thread/start transport should work")
.expect("thread/start should succeed");
let thread: ThreadStartResponse =
serde_json::from_value(response).expect("thread/start response should parse");
let thread_id = thread.thread.id;
// Persist the thread metadata so attachment requests can find it.
client
.request(ClientRequest::ThreadSetName {
request_id: RequestId::Integer(2),
params: ThreadSetNameParams {
thread_id: thread_id.clone(),
name: "Shutdown backpressure".to_string(),
},
})
.await
.expect("thread/name/set transport should work")
.expect("thread/name/set should succeed");
// Observe both mutation replies at the point their updates become visible, then
// leave capacity + 1 updates unread to exercise shutdown with full queues.
let sender = client.sender();
for index in 0..(capacity + 3) {
let request_id = RequestId::Integer((index + 3) as i64);
let request = if index % 2 == 0 {
ClientRequest::ThreadAttachmentAdd {
request_id,
params: ThreadAttachmentAddParams {
thread_id: thread_id.clone(),
attachment_type: "pull_request".to_string(),
identity_key: "one-attachment".to_string(),
payload: serde_json::json!({}),
},
}
} else {
ClientRequest::ThreadAttachmentRemove {
request_id,
params: ThreadAttachmentRemoveParams {
thread_id: thread_id.clone(),
attachment_type: "pull_request".to_string(),
identity_key: "one-attachment".to_string(),
},
}
};
let mut reply = std::pin::pin!(sender.request(request));
if index < 2 {
assert!(reply.as_mut().now_or_never().is_none());
loop {
let event = client
.next_event()
.await
.expect("event stream should remain open");
if let InProcessServerEvent::ServerNotification(notification) = event
&& matches!(
*notification,
ServerNotification::ThreadAttachmentUpdated(_)
)
{
reply
.as_mut()
.now_or_never()
.expect("the attachment reply must be ready before its update")
.expect("attachment request transport should work")
.expect("attachment mutation should succeed");
break;
}
}
} else {
reply
.await
.expect("attachment request transport should work")
.expect("attachment mutation should succeed");
}
}
assert_eq!(client.event_rx.capacity(), 0);
let mut queued_requests = Vec::with_capacity(capacity);
for index in 0..capacity {
let mut request = Box::pin(sender.request(ClientRequest::ConfigRequirementsRead {
request_id: RequestId::Integer((1000 + index) as i64),
params: None,
}));
// Fill the command queue on this single-thread runtime without giving the
// dispatcher a turn between submissions.
assert!(request.as_mut().now_or_never().is_none());
queued_requests.push(request);
}
assert_eq!(sender.client_tx.capacity(), 0);
// Only advance time after the real requests have filled both queues.
tokio::time::pause();
timeout(
SHUTDOWN_ACK_TIMEOUT + SHUTDOWN_TIMEOUT + Duration::from_secs(1),
client.shutdown(),
)
.await
.expect("shutdown must complete with full request and event queues")
.expect("shutdown should succeed");
}
}