| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| 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; |
| 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::ThreadItem; |
| 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_login::AuthManager; |
| 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; |
|
|
| const IN_PROCESS_CONNECTION_ID: ConnectionId = ConnectionId(0); |
| const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(5); |
| |
| const SHUTDOWN_ACK_TIMEOUT: Duration = Duration::from_secs(35); |
| |
| 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::TurnCompleted(_) |
| | ServerNotification::ThreadQueueChanged(_) |
| | ServerNotification::ThreadSettingsUpdated(_) |
| | ServerNotification::ThreadAttachmentUpdated(_) |
| | ServerNotification::ExternalAgentConfigImportCompleted(_) |
| | ServerNotification::ItemCompleted(ItemCompletedNotification { |
| item: ThreadItem::AgentMessage { |
| delivery: Some(AgentMessageDelivery::Async), |
| .. |
| }, |
| .. |
| }) |
| ) |
| } |
|
|
| |
| |
| |
| |
| #[derive(Clone)] |
| pub struct InProcessStartArgs { |
| |
| pub arg0_paths: Arg0DispatchPaths, |
| |
| pub config: Arc<Config>, |
| |
| pub cli_overrides: Vec<(String, TomlValue)>, |
| |
| pub loader_overrides: LoaderOverrides, |
| |
| pub strict_config: bool, |
| |
| pub cloud_config_bundle: CloudConfigBundleLoader, |
| |
| pub thread_config_loader: Arc<dyn ThreadConfigLoader>, |
| |
| pub feedback: CodexFeedback, |
| |
| pub log_db: Option<LogDbLayer>, |
| |
| pub state_db: Option<StateDbHandle>, |
| |
| pub environment_manager: Arc<EnvironmentManager>, |
| |
| pub config_warnings: Vec<ConfigWarningNotification>, |
| |
| pub session_source: SessionSource, |
| |
| pub enable_codex_api_key_env: bool, |
| |
| pub initialize: InitializeParams, |
| |
| pub channel_capacity: usize, |
| } |
|
|
| |
| |
| |
| |
| #[derive(Debug, Clone)] |
| pub enum InProcessServerEvent { |
| |
| ServerRequest(Box<ServerRequest>), |
| |
| ServerNotification(Box<ServerNotification>), |
| |
| Lagged { skipped: usize }, |
| } |
|
|
| |
| |
| |
| |
| |
| 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", |
| )), |
| } |
| } |
| } |
|
|
| |
| |
| |
| |
| |
| 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 { |
| |
| |
| |
| |
| |
| |
| |
| pub async fn request(&self, request: ClientRequest) -> IoResult<PendingClientRequestResponse> { |
| self.client.request(request).await |
| } |
|
|
| |
| |
| |
| |
| pub fn notify(&self, notification: ClientNotification) -> IoResult<()> { |
| self.client.notify(notification) |
| } |
|
|
| |
| |
| |
| |
| |
| pub fn respond_to_server_request(&self, request_id: RequestId, result: Result) -> IoResult<()> { |
| self.client.respond_to_server_request(request_id, result) |
| } |
|
|
| |
| |
| |
| |
| pub fn fail_server_request( |
| &self, |
| request_id: RequestId, |
| error: JSONRPCErrorError, |
| ) -> IoResult<()> { |
| self.client.fail_server_request(request_id, error) |
| } |
|
|
| |
| |
| |
| |
| pub async fn next_event(&mut self) -> Option<InProcessServerEvent> { |
| self.event_rx.recv().await |
| } |
|
|
| |
| |
| |
| |
| pub async fn shutdown(self) -> IoResult<()> { |
| 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() |
| } |
| } |
|
|
| |
| |
| |
| |
| |
| 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 = 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(args: InProcessStartArgs) -> IoResult<InProcessClientHandle> { |
| args.config.auth_config().validate()?; |
| let channel_capacity = args.channel_capacity.max(1); |
| let installation_id = resolve_installation_id(&args.config.codex_home).await?; |
| let auth_manager = |
| AuthManager::shared_from_config(args.config.as_ref(), args.enable_codex_api_key_env) |
| .await |
| .map_err(IoError::other)?; |
| 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(); |
| let outgoing_message_sender = Arc::new(OutgoingMessageSender::new( |
| outgoing_tx, |
| analytics_events_client.clone(), |
| )); |
|
|
| 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), |
| 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 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, |
| ); |
| let (processor_tx, mut processor_rx) = mpsc::channel::<ProcessorCommand>(channel_capacity); |
| let mut processor_handle = tokio::spawn(async move { |
| let processor = Arc::new(MessageProcessor::new(MessageProcessorArgs { |
| outgoing: Arc::clone(&processor_outgoing), |
| analytics_events_client, |
| arg0_paths: args.arg0_paths, |
| config: args.config, |
| 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, |
| 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; |
| } |
| } |
| } |
| 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; |
|
|
| 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; |
| } |
| } |
| } |
| queued_message = writer_rx.recv() => { |
| 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) => { |
| |
| |
| 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(¬ification) { |
| if event_tx |
| .send(InProcessServerEvent::ServerNotification(Box::new( |
| notification, |
| ))) |
| .await |
| .is_err() |
| { |
| break; |
| } |
| } else 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(writer_rx); |
| drop(processor_tx); |
| outgoing_message_sender |
| .cancel_all_requests(Some(internal_error( |
| "in-process app-server runtime is shutting down", |
| ))) |
| .await; |
| |
| |
| 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::ThreadAttachmentOperation; |
| use codex_app_server_protocol::ThreadAttachmentUpdatedNotification; |
| use codex_app_server_protocol::ThreadQueueChangedNotification; |
| 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 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_capacity( |
| session_source: SessionSource, |
| channel_capacity: usize, |
| ) -> 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(), |
| 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: None, |
| }, |
| 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_capacity(session_source, DEFAULT_IN_PROCESS_CHANNEL_CAPACITY).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"); |
| } |
|
|
| #[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_capacity(SessionSource::Cli, 0).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( 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( 1); |
| let (_event_tx, event_rx) = mpsc::channel( 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(), |
| 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 { |
| item: ThreadItem::AgentMessage { |
| id: "item-1".to_string(), |
| text: "Still working".to_string(), |
| phase: None, |
| memory_citation: None, |
| delivery: Some(AgentMessageDelivery::Async), |
| questions: None, |
| }, |
| thread_id: "thread-1".to_string(), |
| turn_id: "turn-1".to_string(), |
| completed_at_ms: 0, |
| }) |
| )); |
| } |
| } |
|
|