github/codex
4
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-20 00:18:49 +00:00
Code Issues Packages Projects Releases Wiki Activity

Report selected environment connection transitions (#33251)

Browse Source 
## What changed

- Track connected and disconnected states across initial remote exec-server connections and reconnection attempts.
- Emit experimental `thread/environment/connected` and `thread/environment/disconnected` app-server notifications for each thread selecting the environment. Each payload identifies the thread and environment; current state is not replayed when a thread starts.
- Stop forwarding connection events when an environment selection is removed or replaced.

## Testing

- Cover connection, disconnection, reconnection, shared-environment notifications, and replacement of a selected environment.

GitOrigin-RevId: 5dd767372363c4a2a8319fc16164be117d5bd20c
This commit is contained in:
sayan-oai
2026-07-15 05:21:13 +00:00
committed by copyberry
parent 8aae858958
commit 3afbd8dd45
28 changed files with 672 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/app-server-protocol/schema/json/ServerNotification.json
Generated
+55
View File
@@ -1140,6 +1140,21 @@
],
"type": "string"
},
"EnvironmentConnectionNotification": {
"properties": {
"environmentId": {
"type": "string"
},
"threadId": {
"type": "string"
}
},
"required": [
"environmentId",
"threadId"
],
"type": "object"
},
"ErrorNotification": {
"properties": {
"error": {
@@ -5991,6 +6006,46 @@
"title": "Thread/goal/clearedNotification",
"type": "object"
},
{
"properties": {
"method": {
"enum": [
"thread/environment/connected"
],
"title": "Thread/environment/connectedNotificationMethod",
"type": "string"
},
"params": {
"$ref": "#/definitions/EnvironmentConnectionNotification"
}
},
"required": [
"method",
"params"
],
"title": "Thread/environment/connectedNotification",
"type": "object"
},
{
"properties": {
"method": {
"enum": [
"thread/environment/disconnected"
],
"title": "Thread/environment/disconnectedNotificationMethod",
"type": "string"
},
"params": {
"$ref": "#/definitions/EnvironmentConnectionNotification"
}
},
"required": [
"method",
"params"
],
"title": "Thread/environment/disconnectedNotification",
"type": "object"
},
{
"properties": {
"method": {
codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json
Generated
+57
View File
@@ -4328,6 +4328,46 @@
"title": "Thread/goal/clearedNotification",
"type": "object"
},
{
"properties": {
"method": {
"enum": [
"thread/environment/connected"
],
"title": "Thread/environment/connectedNotificationMethod",
"type": "string"
},
"params": {
"$ref": "#/definitions/v2/EnvironmentConnectionNotification"
}
},
"required": [
"method",
"params"
],
"title": "Thread/environment/connectedNotification",
"type": "object"
},
{
"properties": {
"method": {
"enum": [
"thread/environment/disconnected"
],
"title": "Thread/environment/disconnectedNotificationMethod",
"type": "string"
},
"params": {
"$ref": "#/definitions/v2/EnvironmentConnectionNotification"
}
},
"required": [
"method",
"params"
],
"title": "Thread/environment/disconnectedNotification",
"type": "object"
},
{
"properties": {
"method": {
@@ -9025,6 +9065,23 @@
}
]
},
"EnvironmentConnectionNotification": {
"$schema": "http://json-schema.org/draft-07/schema#",
"properties": {
"environmentId": {
"type": "string"
},
"threadId": {
"type": "string"
}
},
"required": [
"environmentId",
"threadId"
],
"title": "EnvironmentConnectionNotification",
"type": "object"
},
"ErrorNotification": {
"$schema": "http://json-schema.org/draft-07/schema#",
"properties": {
codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json
Generated
+57
View File
@@ -5258,6 +5258,23 @@
}
]
},
"EnvironmentConnectionNotification": {
"$schema": "http://json-schema.org/draft-07/schema#",
"properties": {
"environmentId": {
"type": "string"
},
"threadId": {
"type": "string"
}
},
"required": [
"environmentId",
"threadId"
],
"title": "EnvironmentConnectionNotification",
"type": "object"
},
"ErrorNotification": {
"$schema": "http://json-schema.org/draft-07/schema#",
"properties": {
@@ -13290,6 +13307,46 @@
"title": "Thread/goal/clearedNotification",
"type": "object"
},
{
"properties": {
"method": {
"enum": [
"thread/environment/connected"
],
"title": "Thread/environment/connectedNotificationMethod",
"type": "string"
},
"params": {
"$ref": "#/definitions/EnvironmentConnectionNotification"
}
},
"required": [
"method",
"params"
],
"title": "Thread/environment/connectedNotification",
"type": "object"
},
{
"properties": {
"method": {
"enum": [
"thread/environment/disconnected"
],
"title": "Thread/environment/disconnectedNotificationMethod",
"type": "string"
},
"params": {
"$ref": "#/definitions/EnvironmentConnectionNotification"
}
},
"required": [
"method",
"params"
],
"title": "Thread/environment/disconnectedNotification",
"type": "object"
},
{
"properties": {
"method": {
codex-rs/app-server-protocol/schema/json/v2/EnvironmentConnectionNotification.json
Generated
+17
View File
@@ -0,0 +1,17 @@
{
"$schema": "http://json-schema.org/draft-07/schema#",
"properties": {
"environmentId": {
"type": "string"
},
"threadId": {
"type": "string"
}
},
"required": [
"environmentId",
"threadId"
],
"title": "EnvironmentConnectionNotification",
"type": "object"
}
codex-rs/app-server-protocol/schema/typescript/ServerNotification.ts
Generated
+2 -1
View File
File diff suppressed because one or more lines are too long
codex-rs/app-server-protocol/schema/typescript/ServerNotificationEnvelope.ts
Generated
+2 -1
View File
File diff suppressed because one or more lines are too long
codex-rs/app-server-protocol/schema/typescript/v2/EnvironmentConnectionNotification.ts
Generated
+5
View File
@@ -0,0 +1,5 @@
// GENERATED CODE! DO NOT MODIFY BY HAND!
// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually.
export type EnvironmentConnectionNotification = { threadId: string, environmentId: string, };
codex-rs/app-server-protocol/schema/typescript/v2/index.ts
Generated
+1
View File
@@ -99,6 +99,7 @@ export type { DynamicToolFunctionSpec } from "./DynamicToolFunctionSpec";
export type { DynamicToolNamespaceSpec } from "./DynamicToolNamespaceSpec";
export type { DynamicToolNamespaceTool } from "./DynamicToolNamespaceTool";
export type { DynamicToolSpec } from "./DynamicToolSpec";
export type { EnvironmentConnectionNotification } from "./EnvironmentConnectionNotification";
export type { ErrorNotification } from "./ErrorNotification";
export type { ExecPolicyAmendment } from "./ExecPolicyAmendment";
export type { ExperimentalFeature } from "./ExperimentalFeature";
codex-rs/app-server-protocol/src/protocol/common.rs
+4
View File
@@ -1630,6 +1630,10 @@ server_notification_definitions! {
ThreadNameUpdated => "thread/name/updated" (v2::ThreadNameUpdatedNotification),
ThreadGoalUpdated => "thread/goal/updated" (v2::ThreadGoalUpdatedNotification),
ThreadGoalCleared => "thread/goal/cleared" (v2::ThreadGoalClearedNotification),
#[experimental("thread/environment/connected")]
EnvironmentConnected => "thread/environment/connected" (v2::EnvironmentConnectionNotification),
#[experimental("thread/environment/disconnected")]
EnvironmentDisconnected => "thread/environment/disconnected" (v2::EnvironmentConnectionNotification),
#[experimental("thread/settings/updated")]
ThreadSettingsUpdated => "thread/settings/updated" (v2::ThreadSettingsUpdatedNotification),
ThreadTokenUsageUpdated => "thread/tokenUsage/updated" (v2::ThreadTokenUsageUpdatedNotification),
codex-rs/app-server-protocol/src/protocol/v2/environment.rs
+8
View File
@@ -21,6 +21,14 @@ pub struct EnvironmentAddParams {
#[ts(export_to = "v2/")]
pub struct EnvironmentAddResponse {}
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)]
#[serde(rename_all = "camelCase")]
#[ts(rename_all = "camelCase", export_to = "v2/")]
pub struct EnvironmentConnectionNotification {
pub thread_id: String,
pub environment_id: String,
}
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, JsonSchema, TS)]
#[serde(rename_all = "camelCase")]
#[ts(export_to = "v2/")]
codex-rs/app-server/README.md
+1
View File
@@ -207,6 +207,7 @@ Example with notification opt-out:
- `environment/add` — experimental; add or replace a named remote environment by `environmentId` and `execServerUrl` for later selection by `thread/start` or `turn/start`; optional `connectTimeoutMs` overrides the WebSocket connection timeout; returns `{}` and does not change the default environment.
- `environment/info` — experimental; connect to a configured environment by `environmentId` and return its detected `shell` plus its default `cwd` as a canonical environment-native `file:` URI. Connection failures are returned as request errors.
- `environment/status` — experimental; read the current status for one configured `environmentId`. Ready remote environments are probed over their existing exec-server connection without starting or reconnecting environments; the response reports `ready`, `pending`, `disconnected`, or `unknown`.
- `thread/environment/connected` and `thread/environment/disconnected` — experimental; report exec-server connection transitions observed after thread startup for selected environments. Current connection state is not replayed.
- `collaborationMode/list` — list available collaboration mode presets (experimental, no pagination). Built-in presets do not select a model; the Plan preset selects medium reasoning effort. This response omits built-in developer instructions; clients should either pass `settings.developer_instructions: null` when setting a mode to use Codex's built-in instructions, or provide their own instructions explicitly.
- `skills/list` — list skills for one or more `cwd` values (optional `forceReload`).
- `skills/extraRoots/set` — replace the app-server process runtime extra standalone skill roots. The roots are not persisted; missing directories are accepted and simply load no skills.
codex-rs/app-server/src/bespoke_event_handling.rs
+21
View File
@@ -22,6 +22,7 @@ use codex_app_server_protocol::CommandExecutionSource;
use codex_app_server_protocol::CommandExecutionStatus;
use codex_app_server_protocol::DeprecationNoticeNotification;
use codex_app_server_protocol::DynamicToolCallParams;
use codex_app_server_protocol::EnvironmentConnectionNotification;
use codex_app_server_protocol::ErrorNotification;
use codex_app_server_protocol::ExecPolicyAmendment as V2ExecPolicyAmendment;
use codex_app_server_protocol::FileChangeApprovalDecision;
@@ -225,6 +226,26 @@ pub(crate) async fn apply_bespoke_event_handling(
.send_server_notification(ServerNotification::McpServerStatusUpdated(notification))
.await;
}
EventMsg::EnvironmentConnected(event) => {
outgoing
.send_server_notification(ServerNotification::EnvironmentConnected(
EnvironmentConnectionNotification {
thread_id: conversation_id.to_string(),
environment_id: event.environment_id,
},
))
.await;
}
EventMsg::EnvironmentDisconnected(event) => {
outgoing
.send_server_notification(ServerNotification::EnvironmentDisconnected(
EnvironmentConnectionNotification {
thread_id: conversation_id.to_string(),
environment_id: event.environment_id,
},
))
.await;
}
EventMsg::Warning(warning_event) => {
let notification = WarningNotification {
thread_id: Some(conversation_id.to_string()),
codex-rs/app-server/tests/suite/v2/environment_add.rs
+144
View File
@@ -4,14 +4,22 @@ use anyhow::Result;
use app_test_support::TestAppServer;
use app_test_support::to_response;
use codex_app_server_protocol::EnvironmentAddResponse;
use codex_app_server_protocol::EnvironmentConnectionNotification;
use codex_app_server_protocol::JSONRPCResponse;
use codex_app_server_protocol::RequestId;
use codex_app_server_protocol::ThreadStartParams;
use codex_app_server_protocol::ThreadStartResponse;
use codex_app_server_protocol::TurnEnvironmentParams;
use pretty_assertions::assert_eq;
use serde_json::json;
use tempfile::TempDir;
use tokio::io::AsyncReadExt;
use tokio::net::TcpListener;
use tokio::sync::oneshot;
use tokio::time::timeout;
use super::exec_server_test_support::accept_exec_server_environment;
const RPC_TIMEOUT: Duration = Duration::from_secs(10);
const CONNECTION_CLOSE_TIMEOUT: Duration = Duration::from_secs(5);
@@ -54,3 +62,139 @@ async fn environment_add_applies_connect_timeout() -> Result<()> {
timeout(CONNECTION_CLOSE_TIMEOUT, stalled_server).await???;
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn selected_environment_emits_connection_lifecycle_notifications() -> Result<()> {
let listener = TcpListener::bind("127.0.0.1:0").await?;
let exec_server_url = format!("ws://{}", listener.local_addr()?);
let codex_home = TempDir::new()?;
std::fs::write(
codex_home.path().join("config.toml"),
"[features]\ndeferred_executor = true\n",
)?;
let mut app_server = TestAppServer::builder()
.with_codex_home(codex_home.path())
.without_auto_env()
.build()
.await?;
timeout(RPC_TIMEOUT, app_server.initialize()).await??;
let request_id = app_server
.send_raw_request(
"environment/add",
Some(json!({
"environmentId": "remote-a",
"execServerUrl": exec_server_url,
})),
)
.await?;
let response: JSONRPCResponse = timeout(
RPC_TIMEOUT,
app_server.read_stream_until_response_message(RequestId::Integer(request_id)),
)
.await??;
let _: EnvironmentAddResponse = to_response(response)?;
let environment = TurnEnvironmentParams {
environment_id: "remote-a".to_string(),
cwd: codex_utils_absolute_path::AbsolutePathBuf::try_from(codex_home.path().to_path_buf())?
.into(),
runtime_workspace_roots: None,
};
let request_id = app_server
.send_thread_start_request(ThreadStartParams {
environments: Some(vec![environment.clone()]),
..Default::default()
})
.await?;
let response: JSONRPCResponse = timeout(
RPC_TIMEOUT,
app_server.read_stream_until_response_message(RequestId::Integer(request_id)),
)
.await??;
let ThreadStartResponse { thread, .. } = to_response(response)?;
let (disconnect_tx, disconnect_rx) = oneshot::channel();
let exec_server = tokio::spawn(async move {
let mut websocket = accept_exec_server_environment(
listener,
json!({"shell": {"name": "zsh", "path": "/bin/zsh"}}),
)
.await?;
disconnect_rx.await?;
websocket.close(None).await?;
Ok::<_, anyhow::Error>(())
});
let connected = timeout(
RPC_TIMEOUT,
app_server.read_stream_until_notification_message("thread/environment/connected"),
)
.await??;
assert_eq!(
serde_json::from_value::<EnvironmentConnectionNotification>(
connected.params.expect("connected notification params"),
)?,
EnvironmentConnectionNotification {
thread_id: thread.id.clone(),
environment_id: "remote-a".to_string(),
}
);
let request_id = app_server
.send_thread_start_request(ThreadStartParams {
environments: Some(vec![environment]),
..Default::default()
})
.await?;
let response: JSONRPCResponse = timeout(
RPC_TIMEOUT,
app_server.read_stream_until_response_message(RequestId::Integer(request_id)),
)
.await??;
let ThreadStartResponse {
thread: second_thread,
..
} = to_response(response)?;
disconnect_tx
.send(())
.map_err(|_| anyhow::anyhow!("exec-server disconnect receiver closed"))?;
let mut disconnected = Vec::new();
for _ in 0..2 {
let notification = timeout(
RPC_TIMEOUT,
app_server.read_stream_until_notification_message("thread/environment/disconnected"),
)
.await??;
disconnected.push(serde_json::from_value::<EnvironmentConnectionNotification>(
notification
.params
.expect("disconnected notification params"),
)?);
}
disconnected.sort_by(|left, right| left.thread_id.cmp(&right.thread_id));
let mut expected = vec![
EnvironmentConnectionNotification {
thread_id: thread.id,
environment_id: "remote-a".to_string(),
},
EnvironmentConnectionNotification {
thread_id: second_thread.id,
environment_id: "remote-a".to_string(),
},
];
expected.sort_by(|left, right| left.thread_id.cmp(&right.thread_id));
assert_eq!(disconnected, expected);
assert!(
!app_server
.pending_notification_methods()
.iter()
.any(|method| method == "thread/environment/connected"),
"connection state should not be replayed when a thread starts"
);
timeout(RPC_TIMEOUT, exec_server).await???;
Ok(())
}
codex-rs/core/src/environment_selection.rs
+144 -4
View File
@@ -2,18 +2,25 @@ use std::collections::HashMap;
use std::collections::HashSet;
use std::fmt;
use std::sync::Arc;
use std::sync::OnceLock;
use arc_swap::ArcSwap;
use async_channel::Sender;
use codex_exec_server::Environment;
use codex_exec_server::EnvironmentConnectionState;
use codex_exec_server::EnvironmentManager;
use codex_exec_server::ExecServerError;
use codex_exec_server::ExecutorFileSystem;
use codex_protocol::protocol::EnvironmentConnectionEvent;
use codex_protocol::protocol::Event;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::TurnEnvironmentSelection;
use codex_utils_absolute_path::AbsolutePathBuf;
use codex_utils_path_uri::PathUri;
use futures::FutureExt;
use futures::future::BoxFuture;
use futures::future::Shared;
use tokio_util::task::AbortOnDropHandle;
use crate::session::turn_context::TurnEnvironment;
use crate::shell::Shell;
@@ -41,6 +48,9 @@ type TurnEnvironmentResolution = Shared<BoxFuture<'static, TurnEnvironmentResult
#[derive(Clone)]
struct SelectedTurnEnvironment {
selection: TurnEnvironmentSelection,
environment: Arc<Environment>,
// Selection clones share one listener; the final handle drop aborts it.
connection_events_task: Option<Arc<AbortOnDropHandle<()>>>,
resolution: TurnEnvironmentResolution,
}
@@ -72,6 +82,7 @@ pub(crate) struct ThreadEnvironments {
shell_snapshot: ShellSnapshot,
non_blocking_snapshots: bool,
environments: ArcSwap<Vec<SelectedTurnEnvironment>>,
connection_event_tx: OnceLock<Sender<Event>>,
}
impl ThreadEnvironments {
@@ -88,10 +99,13 @@ impl ThreadEnvironments {
.into_iter()
.map(|environment| {
let selection = environment.selection();
let selected_environment = Arc::clone(&environment.environment);
let resolution: TurnEnvironmentResolution =
futures::future::ready(Ok(environment)).boxed().shared();
SelectedTurnEnvironment {
selection,
environment: selected_environment,
connection_events_task: None,
resolution,
}
})
@@ -102,6 +116,7 @@ impl ThreadEnvironments {
shell_snapshot,
non_blocking_snapshots,
environments: ArcSwap::from_pointee(environments),
connection_event_tx: OnceLock::new(),
}
}
@@ -127,9 +142,26 @@ impl ThreadEnvironments {
tracing::warn!("skipping unknown turn environment `{environment_id}`");
continue;
};
// Connection state belongs to the environment instance, not its cwd or roots.
let connection_events_task = previous
.iter()
.find(|previous| {
previous.selection.environment_id.as_str() == environment_id.as_str()
&& Arc::ptr_eq(&previous.environment, &environment)
})
.and_then(|previous| previous.connection_events_task.clone())
.or_else(|| {
self.connection_event_tx.get().and_then(|tx_event| {
Self::spawn_connection_event_listener(
environment.as_ref(),
environment_id.clone(),
tx_event.clone(),
)
})
});
let (resolution_task, resolution) = Self::resolve_environment(
selected_environment.clone(),
environment,
Arc::clone(&environment),
self.local_shell.clone(),
self.shell_snapshot.clone(),
)
@@ -138,10 +170,92 @@ impl ThreadEnvironments {
let resolution = resolution.boxed().shared();
next.push(SelectedTurnEnvironment {
selection: selected_environment.clone(),
environment,
connection_events_task,
resolution,
});
}
let removed_connection_tasks = previous
.iter()
.filter_map(|previous| {
let task = previous.connection_events_task.as_ref()?;
(!next.iter().any(|next| {
next.connection_events_task
.as_ref()
.is_some_and(|next_task| Arc::ptr_eq(task, next_task))
}))
.then(|| Arc::clone(task))
})
.collect::<Vec<_>>();
self.environments.store(Arc::new(next));
// ArcSwap readers may retain removed selections, so abort at logical removal.
for task in removed_connection_tasks {
task.abort();
}
}
fn spawn_connection_event_listener(
environment: &Environment,
environment_id: String,
tx_event: Sender<Event>,
) -> Option<Arc<AbortOnDropHandle<()>>> {
let mut connection_state = environment.subscribe_connection_state()?;
let task = tokio::spawn(async move {
loop {
let state = tokio::select! {
_ = tx_event.closed() => return,
changed = connection_state.changed() => {
if changed.is_err() {
return;
}
*connection_state.borrow_and_update()
}
};
let msg = match state {
EnvironmentConnectionState::Connected => {
EventMsg::EnvironmentConnected(EnvironmentConnectionEvent {
environment_id: environment_id.clone(),
})
}
EnvironmentConnectionState::Disconnected => {
EventMsg::EnvironmentDisconnected(EnvironmentConnectionEvent {
environment_id: environment_id.clone(),
})
}
};
if tx_event
.send(Event {
id: String::new(),
msg,
})
.await
.is_err()
{
return;
}
}
});
Some(Arc::new(AbortOnDropHandle::new(task)))
}
pub(crate) fn start_connection_event_forwarding(&self, tx_event: Sender<Event>) {
let tx_event = self.connection_event_tx.get_or_init(|| tx_event);
let current = self.environments.load_full();
let environments = current
.iter()
.map(|selected| {
let mut selected = selected.clone();
if selected.connection_events_task.is_none() {
selected.connection_events_task = Self::spawn_connection_event_listener(
selected.environment.as_ref(),
selected.selection.environment_id.clone(),
tx_event.clone(),
);
}
selected
})
.collect();
self.environments.store(Arc::new(environments));
}
fn resolve_environment(
@@ -760,7 +874,7 @@ url = "ws://127.0.0.1:8765"
}
#[tokio::test]
async fn matching_environment_id_and_cwd_reuse_resolution() {
async fn replacement_environment_events_follow_selected_environment() {
let cwd = AbsolutePathBuf::current_dir().expect("cwd");
let first_listener = TcpListener::bind("127.0.0.1:0")
.await
@@ -780,13 +894,15 @@ url = "ws://127.0.0.1:8765"
cwd: PathUri::from_abs_path(&cwd),
workspace_roots: Vec::new(),
};
let environments = ThreadEnvironments::new(
let (tx_event, rx_event) = async_channel::unbounded();
let environments = Arc::new(ThreadEnvironments::new(
Arc::clone(&manager),
crate::shell::default_user_shell(),
ShellSnapshot::disabled(),
TurnEnvironmentSnapshot::default(),
/*non_blocking_snapshots*/ true,
);
));
environments.start_connection_event_forwarding(tx_event);
environments.update_selections(std::slice::from_ref(&selection));
let initial_snapshot = environments.snapshot().await;
let second_listener = TcpListener::bind("127.0.0.1:0")
@@ -827,6 +943,30 @@ url = "ws://127.0.0.1:8765"
.expect("changed environment");
assert!(initial.resolution.ptr_eq(&reused.resolution));
assert!(!reused.resolution.ptr_eq(&changed.resolution));
serve_environment_info(first_listener).await;
assert!(
timeout(Duration::from_millis(250), rx_event.recv())
.await
.is_err(),
"old environment event should not be forwarded"
);
serve_environment_info(second_listener).await;
let event = timeout(Duration::from_secs(5), rx_event.recv())
.await
.expect("replacement environment event")
.expect("event channel");
let event = match event.msg {
EventMsg::EnvironmentConnected(event) => event,
other => panic!("expected connected event, got {other:?}"),
};
assert_eq!(
event,
EnvironmentConnectionEvent {
environment_id: REMOTE_ENVIRONMENT_ID.to_string(),
}
);
}
#[tokio::test]
codex-rs/core/src/session/session.rs
+1
View File
@@ -1183,6 +1183,7 @@ impl Session {
for event in events {
sess.send_event_raw(event).await;
}
turn_environments.start_connection_event_forwarding(tx_event.clone());
let mcp_startup_cancellation_token = {
let mut cancel_guard = sess.services.mcp_startup_cancellation_token.lock().await;
codex-rs/core/src/session/turn.rs
+2
View File
@@ -1552,6 +1552,8 @@ pub(super) fn realtime_text_for_event(msg: &EventMsg) -> Option<(String, Option<
| EventMsg::AgentReasoningRawContent(_)
| EventMsg::AgentReasoningSectionBreak(_)
| EventMsg::SessionConfigured(_)
| EventMsg::EnvironmentConnected(_)
| EventMsg::EnvironmentDisconnected(_)
| EventMsg::ThreadGoalUpdated(_)
| EventMsg::McpStartupUpdate(_)
| EventMsg::McpStartupComplete(_)
codex-rs/exec-server/src/client.rs
+68 -10
View File
@@ -35,6 +35,7 @@ use crate::client_api::RemoteExecServerConnectArgs;
use crate::client_api::StdioExecServerConnectArgs;
use crate::client_transport::ExecServerReconnectStrategy;
use crate::connection::JsonRpcConnection;
use crate::environment::EnvironmentConnectionState;
use crate::process::ExecProcessEvent;
use crate::process::ExecProcessEventLog;
use crate::process::ExecProcessEventReceiver;
@@ -220,6 +221,7 @@ struct Inner {
struct ConnectionState {
status: ConnectionStatus,
active_process_starts: usize,
environment_connection_state_tx: watch::Sender<EnvironmentConnectionState>,
}
enum ConnectionStatus {
@@ -228,6 +230,34 @@ enum ConnectionStatus {
Failed(String),
}
impl ConnectionState {
fn set_status(&mut self, status: ConnectionStatus) {
self.status = status;
self.publish_environment_connection_state();
}
fn publish_environment_connection_state(&self) {
let state = match &self.status {
ConnectionStatus::Connected(rpc_client) if !rpc_client.is_disconnected() => {
EnvironmentConnectionState::Connected
}
ConnectionStatus::Connected(_)
| ConnectionStatus::Recovering
| ConnectionStatus::Failed(_) => EnvironmentConnectionState::Disconnected,
};
let _ = self
.environment_connection_state_tx
.send_if_modified(|current| {
if *current == state {
false
} else {
*current = state;
true
}
});
}
}
#[derive(Clone, Copy)]
enum RecoveryPolicy {
Wait,
@@ -262,6 +292,7 @@ pub(crate) struct LazyRemoteExecServerClient {
// The latest successful client, replaced whenever reconnecting succeeds.
current_client: Arc<StdMutex<Option<ExecServerClient>>>,
reconnect: Arc<StdMutex<Option<Arc<ConnectionAttempt>>>>,
environment_connection_state_tx: watch::Sender<EnvironmentConnectionState>,
}
impl LazyRemoteExecServerClient {
@@ -272,9 +303,17 @@ impl LazyRemoteExecServerClient {
startup: Arc::new(ConnectionAttempt::new()),
current_client: Arc::new(StdMutex::new(None)),
reconnect: Arc::new(StdMutex::new(None)),
environment_connection_state_tx: watch::channel(
EnvironmentConnectionState::Disconnected,
)
.0,
}
}
pub(crate) fn subscribe_connection_state(&self) -> watch::Receiver<EnvironmentConnectionState> {
self.environment_connection_state_tx.subscribe()
}
pub(crate) fn start_connecting(&self) -> Option<AbortOnDropHandle<()>> {
// Stdio starts a process, so keep it lazy until the environment is used.
if matches!(
@@ -379,10 +418,7 @@ impl LazyRemoteExecServerClient {
async fn initial_client(&self) -> Result<ExecServerClient, ExecServerError> {
// The first caller starts the work; every other caller waits for that same result.
let result = self
.startup
.get_or_init(|| connect_once(self.transport_params.clone()))
.await;
let result = self.startup.get_or_init(|| self.connect_once()).await;
match result {
Ok(client) => {
let mut current_client = self
@@ -414,7 +450,7 @@ impl LazyRemoteExecServerClient {
};
let result = attempt
.get_or_init(|| async {
let result = connect_once(self.transport_params.clone()).await;
let result = self.connect_once().await;
if let Ok(client) = &result {
*self
.current_client
@@ -458,12 +494,17 @@ impl LazyRemoteExecServerClient {
| ExecServerTransportParams::NoiseRendezvous { .. }
)
}
}
async fn connect_once(transport_params: ExecServerTransportParams) -> ConnectionResult {
ExecServerClient::connect_for_transport(transport_params)
.await
.map_err(Arc::new)
async fn connect_once(&self) -> ConnectionResult {
let result = ExecServerClient::connect_for_transport(self.transport_params.clone())
.await
.map_err(Arc::new);
if let Ok(client) = &result {
client
.attach_environment_connection_state(self.environment_connection_state_tx.clone());
}
result
}
}
impl HttpClient for LazyRemoteExecServerClient {
@@ -534,6 +575,19 @@ pub enum ExecServerError {
}
impl ExecServerClient {
fn attach_environment_connection_state(
&self,
state_tx: watch::Sender<EnvironmentConnectionState>,
) {
let mut connection = self
.inner
.connection
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
connection.environment_connection_state_tx = state_tx;
connection.publish_environment_connection_state();
}
fn fail_fast(&self) -> Result<Self, ExecServerError> {
self.rpc_client_without_recovery()?;
Ok(Self {
@@ -890,6 +944,10 @@ impl ExecServerClient {
connection: StdMutex::new(ConnectionState {
status: ConnectionStatus::Connected(Arc::clone(&rpc_client)),
active_process_starts: 0,
environment_connection_state_tx: watch::channel(
EnvironmentConnectionState::Connected,
)
.0,
}),
connection_changed,
sessions: ArcSwap::from_pointee(HashMap::new()),
codex-rs/exec-server/src/client_recovery.rs
+3 -3
View File
@@ -298,7 +298,7 @@ impl Inner {
ConnectionStatus::Connected(current)
if Arc::ptr_eq(current, &failed_rpc_client) =>
{
connection.status = ConnectionStatus::Recovering;
connection.set_status(ConnectionStatus::Recovering);
true
}
ConnectionStatus::Connected(_)
@@ -411,7 +411,7 @@ impl Inner {
{
false
} else {
connection.status = ConnectionStatus::Connected(rpc_client);
connection.set_status(ConnectionStatus::Connected(rpc_client));
true
}
};
@@ -512,7 +512,7 @@ impl Inner {
match &connection.status {
ConnectionStatus::Failed(existing) => (existing.clone(), false),
ConnectionStatus::Connected(_) | ConnectionStatus::Recovering => {
connection.status = ConnectionStatus::Failed(message.clone());
connection.set_status(ConnectionStatus::Failed(message.clone()));
(message, true)
}
}
codex-rs/exec-server/src/environment.rs
+28 -13
View File
@@ -29,6 +29,7 @@ use crate::remote::NoiseRendezvousEnvironmentConfig;
use crate::remote_file_system::RemoteFileSystem;
use crate::remote_process::RemoteProcess;
use tokio::sync::oneshot;
use tokio::sync::watch;
use tokio_util::task::AbortOnDropHandle;
pub const CODEX_EXEC_SERVER_URL_ENV_VAR: &str = "CODEX_EXEC_SERVER_URL";
@@ -40,6 +41,15 @@ pub const CODEX_EXEC_SERVER_NOISE_AUTH_TOKEN_ENV_VAR: &str = "CODEX_EXEC_SERVER_
pub const CODEX_EXEC_SERVER_NOISE_CHATGPT_ACCOUNT_ID_ENV_VAR: &str =
"CODEX_EXEC_SERVER_NOISE_CHATGPT_ACCOUNT_ID";
/// The current connection state for one concrete environment.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum EnvironmentConnectionState {
/// An initialized exec-server connection is available.
Connected,
/// No initialized exec-server connection is currently available.
Disconnected,
}
/// Owns the execution/filesystem environments available to the Codex runtime.
///
/// `EnvironmentManager` is a shared registry for concrete environments. Its
@@ -344,11 +354,7 @@ impl EnvironmentManager {
),
self.local_runtime_paths.clone(),
));
environment.start_connecting();
self.environments
.write()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.insert(environment_id, environment);
self.insert_environment(environment_id, environment);
Ok(())
}
@@ -368,11 +374,7 @@ impl EnvironmentManager {
})),
self.local_runtime_paths.clone(),
));
environment.start_connecting();
self.environments
.write()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.insert(environment_id, environment);
self.insert_environment(environment_id, environment);
Ok(DeferredEnvironmentRegistration(completion))
}
@@ -392,12 +394,16 @@ impl EnvironmentManager {
ExecServerTransportParams::NoiseRendezvous { provider, identity },
self.local_runtime_paths.clone(),
));
environment.start_connecting();
self.insert_environment(environment_id, environment);
Ok(())
}
fn insert_environment(&self, environment_id: String, environment: Arc<Environment>) {
self.environments
.write()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.insert(environment_id, environment);
Ok(())
.insert(environment_id, Arc::clone(&environment));
environment.start_connecting();
}
}
@@ -608,6 +614,15 @@ impl Environment {
self.remote_client.is_some()
}
/// Subscribes to the current connection state for this remote environment.
pub fn subscribe_connection_state(
&self,
) -> Option<watch::Receiver<EnvironmentConnectionState>> {
self.remote_client
.as_ref()
.map(LazyRemoteExecServerClient::subscribe_connection_state)
}
pub fn local_runtime_paths(&self) -> Option<&ExecServerRuntimePaths> {
self.local_runtime_paths.as_ref()
}
codex-rs/exec-server/src/lib.rs
+1
View File
@@ -67,6 +67,7 @@ pub use environment::CODEX_EXEC_SERVER_NOISE_REGISTRY_URL_ENV_VAR;
pub use environment::CODEX_EXEC_SERVER_URL_ENV_VAR;
pub use environment::DeferredEnvironmentRegistration;
pub use environment::Environment;
pub use environment::EnvironmentConnectionState;
pub use environment::EnvironmentManager;
pub use environment::EnvironmentObservedStatus;
pub use environment::LOCAL_ENVIRONMENT_ID;
codex-rs/exec-server/tests/deferred_environment.rs
+4
View File
@@ -45,6 +45,9 @@ async fn deferred_environment_waits_before_connecting() -> anyhow::Result<()> {
let registration =
manager.register_deferred_noise_environment("tools".to_string(), provider.clone())?;
let environment = manager.get_environment("tools").expect("environment");
let connection_state = environment
.subscribe_connection_state()
.expect("remote environment connection state");
let mut readiness = Box::pin(environment.wait_until_ready());
assert!(poll!(&mut readiness).is_pending());
@@ -54,6 +57,7 @@ async fn deferred_environment_waits_before_connecting() -> anyhow::Result<()> {
let error = readiness.await.unwrap_err();
assert!(error.to_string().contains("test Noise provider called"));
assert_eq!(provider.calls(), 1);
assert!(!connection_state.has_changed()?);
Ok(())
}
codex-rs/exec-server/tests/relay.rs
+24
View File
@@ -15,6 +15,7 @@ use anyhow::Result;
use base64::Engine as _;
use base64::engine::general_purpose::STANDARD;
use codex_api::AuthProvider;
use codex_exec_server::EnvironmentConnectionState;
use codex_exec_server::EnvironmentManager;
use codex_exec_server::ExecParams;
use codex_exec_server::ExecResponse;
@@ -42,6 +43,7 @@ use relay_proto::relay_message_frame;
use tempfile::TempDir;
use tokio::net::TcpListener;
use tokio::net::TcpStream;
use tokio::sync::watch;
use tokio::time::timeout;
use tokio_tungstenite::WebSocketStream;
use tokio_tungstenite::accept_async;
@@ -157,6 +159,9 @@ async fn deferred_noise_environment_connects_and_reconnects_with_fresh_bundle()
let environment = manager
.get_environment(ENVIRONMENT_ID)
.context("deferred Noise environment")?;
let mut connection_state = environment
.subscribe_connection_state()
.context("remote environment connection state")?;
assert_eq!(provider.calls(), 0);
registration.complete(Ok(()))?;
@@ -170,9 +175,17 @@ async fn deferred_noise_environment_connects_and_reconnects_with_fresh_bundle()
.await
.context("deferred Noise environment should become ready")??;
assert_eq!(provider.calls(), 1);
assert_eq!(
next_connection_state(&mut connection_state).await?,
EnvironmentConnectionState::Connected
);
first_relay.abort();
let _ = first_relay.await;
assert_eq!(
next_connection_state(&mut connection_state).await?,
EnvironmentConnectionState::Disconnected
);
let first_reconnected_websocket = accept_websocket(&listener, "reconnected peer").await?;
let second_reconnected_websocket = accept_websocket(&listener, "reconnected peer").await?;
let second_relay = tokio::spawn(proxy_relay_frames(
@@ -186,6 +199,10 @@ async fn deferred_noise_environment_connects_and_reconnects_with_fresh_bundle()
assert_eq!(recovered_info, initial_info);
assert_eq!(provider.calls(), 2);
assert_eq!(
next_connection_state(&mut connection_state).await?,
EnvironmentConnectionState::Connected
);
registry.verify().await;
second_relay.abort();
@@ -195,6 +212,13 @@ async fn deferred_noise_environment_connects_and_reconnects_with_fresh_bundle()
Ok(())
}
async fn next_connection_state(
state: &mut watch::Receiver<EnvironmentConnectionState>,
) -> Result<EnvironmentConnectionState> {
timeout(TEST_TIMEOUT, state.changed()).await??;
Ok(*state.borrow_and_update())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn remote_environment_routes_encrypted_exec_server_rpc() -> Result<()> {
let listener = TcpListener::bind("127.0.0.1:0").await?;
codex-rs/mcp-server/src/codex_tool_runner.rs
+2
View File
@@ -333,6 +333,8 @@ async fn run_codex_tool_session_inner(
EventMsg::AgentReasoningRawContent(_)
| EventMsg::TurnStarted(_)
| EventMsg::ThreadSettingsApplied(_)
| EventMsg::EnvironmentConnected(_)
| EventMsg::EnvironmentDisconnected(_)
| EventMsg::TokenCount(_)
| EventMsg::AgentReasoning(_)
| EventMsg::AgentReasoningSectionBreak(_)
codex-rs/protocol/src/protocol.rs
+11
View File
@@ -1268,6 +1268,11 @@ pub struct Event {
pub msg: EventMsg,
}
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq, JsonSchema, TS)]
pub struct EnvironmentConnectionEvent {
pub environment_id: String,
}
/// Response event from the agent
/// NOTE: Make sure none of these values have optional types, as it will mess up the extension code-gen.
#[derive(Debug, Clone, Deserialize, Serialize, Display, JsonSchema, TS)]
@@ -1351,6 +1356,12 @@ pub enum EventMsg {
/// Ack the client's configure message.
SessionConfigured(SessionConfiguredEvent),
/// A selected environment completed its connection handshake.
EnvironmentConnected(EnvironmentConnectionEvent),
/// A selected environment lost its established connection.
EnvironmentDisconnected(EnvironmentConnectionEvent),
/// Updated long-running goal metadata for the thread.
ThreadGoalUpdated(ThreadGoalUpdatedEvent),
codex-rs/rollout-trace/src/protocol_event.rs
+4
View File
@@ -365,6 +365,8 @@ pub(crate) fn tool_runtime_trace_event(event: &EventMsg) -> Option<ToolRuntimeTr
| EventMsg::AgentReasoningRawContent(_)
| EventMsg::AgentReasoningSectionBreak(_)
| EventMsg::SessionConfigured(_)
| EventMsg::EnvironmentConnected(_)
| EventMsg::EnvironmentDisconnected(_)
| EventMsg::McpStartupUpdate(_)
| EventMsg::McpStartupComplete(_)
| EventMsg::WebSearchBegin(_)
@@ -430,6 +432,8 @@ pub(crate) fn wrapped_protocol_event_type(event: &EventMsg) -> Option<&'static s
| EventMsg::TurnModerationMetadata(_)
| EventMsg::ContextCompacted(_)
| EventMsg::ThreadSettingsApplied(_)
| EventMsg::EnvironmentConnected(_)
| EventMsg::EnvironmentDisconnected(_)
| EventMsg::TokenCount(_)
| EventMsg::AgentMessage(_)
| EventMsg::UserMessage(_)
codex-rs/rollout/src/policy.rs
+2
View File
@@ -144,6 +144,8 @@ pub fn should_persist_event_msg(ev: &EventMsg, history_mode: ThreadHistoryMode)
| EventMsg::RawResponseItem(_)
| EventMsg::RawResponseCompleted(_)
| EventMsg::SessionConfigured(_)
| EventMsg::EnvironmentConnected(_)
| EventMsg::EnvironmentDisconnected(_)
| EventMsg::McpToolCallBegin(_)
| EventMsg::ExecCommandBegin(_)
| EventMsg::TerminalInteraction(_)
codex-rs/tui/src/app/app_server_event_targets.rs
+2
View File
@@ -169,6 +169,8 @@ pub(super) fn server_notification_thread_target(
| ServerNotification::AccountUpdated(_)
| ServerNotification::AccountRateLimitsUpdated(_)
| ServerNotification::AppListUpdated(_)
| ServerNotification::EnvironmentConnected(_)
| ServerNotification::EnvironmentDisconnected(_)
| ServerNotification::RemoteControlStatusChanged(_)
| ServerNotification::ExternalAgentConfigImportProgress(_)
| ServerNotification::ExternalAgentConfigImportCompleted(_)
codex-rs/tui/src/chatwidget/protocol.rs
+2
View File
@@ -208,6 +208,8 @@ impl ChatWidget {
| ServerNotification::McpToolCallProgress(_)
| ServerNotification::McpServerOauthLoginCompleted(_)
| ServerNotification::AppListUpdated(_)
| ServerNotification::EnvironmentConnected(_)
| ServerNotification::EnvironmentDisconnected(_)
| ServerNotification::RemoteControlStatusChanged(_)
| ServerNotification::ExternalAgentConfigImportProgress(_)
| ServerNotification::ExternalAgentConfigImportCompleted(_)