openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-29 19:08:43 +00:00
Code Issues Packages Projects Releases Wiki Activity

Prototype remote-control local thread stdio subscription

Browse Source
This commit is contained in:
Khoi
2026-05-18 16:57:24 -07:00
parent a66712c95d
commit df8f6372dc
9 changed files with 263 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/app-server/src/in_process.rs
+1
View File
@@ -449,6 +449,7 @@ async fn start_uninitialized(args: InProcessStartArgs) -> IoResult<InProcessClie
processor
.process_client_request(
IN_PROCESS_CONNECTION_ID,
crate::transport::ConnectionOrigin::InProcess,
*request,
Arc::clone(&session),
&outbound_initialized,
codex-rs/app-server/src/lib.rs
+2
View File
@@ -916,6 +916,7 @@ pub async fn run_main_with_transport_options(
processor
.process_request(
connection_id,
connection_state.origin,
request,
&transport,
Arc::clone(&connection_state.session),
@@ -961,6 +962,7 @@ pub async fn run_main_with_transport_options(
processor
.connection_initialized(
connection_id,
connection_state.origin,
connection_state
.session
.request_attestation(),
codex-rs/app-server/src/message_processor.rs
+9
View File
@@ -43,6 +43,7 @@ use crate::skills_watcher::SkillsWatcher;
use crate::thread_state::ConnectionCapabilities;
use crate::thread_state::ThreadStateManager;
use crate::transport::AppServerTransport;
use crate::transport::ConnectionOrigin;
use crate::transport::RemoteControlHandle;
use async_trait::async_trait;
use codex_analytics::AnalyticsEventsClient;
@@ -508,6 +509,7 @@ impl MessageProcessor {
pub(crate) async fn process_request(
self: &Arc<Self>,
connection_id: ConnectionId,
origin: ConnectionOrigin,
request: JSONRPCRequest,
transport: &AppServerTransport,
session: Arc<ConnectionSessionState>,
@@ -548,6 +550,7 @@ impl MessageProcessor {
self.handle_client_request(
request_id.clone(),
codex_request,
origin,
Arc::clone(&session),
/*outbound_initialized*/ None,
request_context.clone(),
@@ -571,6 +574,7 @@ impl MessageProcessor {
pub(crate) async fn process_client_request(
self: &Arc<Self>,
connection_id: ConnectionId,
origin: ConnectionOrigin,
request: ClientRequest,
session: Arc<ConnectionSessionState>,
outbound_initialized: &AtomicBool,
@@ -599,6 +603,7 @@ impl MessageProcessor {
.handle_client_request(
request_id.clone(),
request,
origin,
Arc::clone(&session),
Some(outbound_initialized),
request_context.clone(),
@@ -654,6 +659,7 @@ impl MessageProcessor {
pub(crate) async fn connection_initialized(
&self,
connection_id: ConnectionId,
origin: ConnectionOrigin,
request_attestation: bool,
) {
self.thread_processor
@@ -662,6 +668,7 @@ impl MessageProcessor {
ConnectionCapabilities {
request_attestation,
},
origin,
)
.await;
}
@@ -737,6 +744,7 @@ impl MessageProcessor {
self: &Arc<Self>,
connection_request_id: ConnectionRequestId,
codex_request: ClientRequest,
origin: ConnectionOrigin,
session: Arc<ConnectionSessionState>,
// `Some(...)` means the caller wants initialize to immediately mark the
// connection outbound-ready. Websocket JSON-RPC calls pass `None` so
@@ -763,6 +771,7 @@ impl MessageProcessor {
ConnectionCapabilities {
request_attestation: session.request_attestation(),
},
origin,
)
.await;
}
codex-rs/app-server/src/message_processor_tracing_tests.rs
+1
View File
@@ -179,6 +179,7 @@ impl TracingHarness {
self.processor
.process_request(
TEST_CONNECTION_ID,
crate::transport::ConnectionOrigin::Stdio,
request,
&AppServerTransport::Stdio,
Arc::clone(&self.session),
codex-rs/app-server/src/outgoing_message.rs
+46
View File
@@ -1319,4 +1319,50 @@ mod tests {
.is_empty()
);
}
#[tokio::test]
async fn thread_scoped_notifications_fan_out_to_all_subscribed_connections() {
let (tx, mut rx) = mpsc::channel::<OutgoingEnvelope>(8);
let outgoing = Arc::new(OutgoingMessageSender::new(
tx,
codex_analytics::AnalyticsEventsClient::disabled(),
));
let thread_id = ThreadId::new();
let notification = ServerNotification::GuardianWarning(GuardianWarningNotification {
thread_id: thread_id.to_string(),
message: "warn".to_string(),
});
let thread_outgoing = ThreadScopedOutgoingMessageSender::new(
outgoing,
vec![ConnectionId(2), ConnectionId(1)],
thread_id,
);
thread_outgoing
.send_server_notification(notification.clone())
.await;
let mut deliveries = Vec::new();
for _ in 0..2 {
let envelope = rx.recv().await.expect("notification should be delivered");
let OutgoingEnvelope::ToConnection {
connection_id,
message: OutgoingMessage::AppServerNotification(sent_notification),
..
} = envelope
else {
panic!("expected thread-scoped notification delivery");
};
match sent_notification {
ServerNotification::GuardianWarning(sent) => {
assert_eq!(sent.thread_id, thread_id.to_string());
assert_eq!(sent.message, "warn");
}
other => panic!("unexpected notification: {other:?}"),
}
deliveries.push(connection_id);
}
deliveries.sort_by_key(|connection_id| connection_id.0);
assert_eq!(deliveries, vec![ConnectionId(1), ConnectionId(2)]);
}
}
codex-rs/app-server/src/request_processors/thread_processor.rs
+25 -1
View File
@@ -1,5 +1,6 @@
use super::*;
use crate::error_code::method_not_found;
use crate::transport::ConnectionOrigin;
use codex_protocol::models::BUILT_IN_PERMISSION_PROFILE_DANGER_FULL_ACCESS;
use codex_protocol::models::BUILT_IN_PERMISSION_PROFILE_WORKSPACE;
@@ -1159,6 +1160,28 @@ impl ThreadRequestProcessor {
request_id.connection_id,
"thread",
);
for connection_id in listener_task_context
.thread_state_manager
.local_stdio_connection_ids_for_remote_control_thread_start(request_id.connection_id)
.await
{
log_listener_attach_result(
super::thread_lifecycle::ensure_conversation_listener(
listener_task_context.clone(),
thread_id,
connection_id,
/*raw_events_enabled*/ false,
)
.instrument(tracing::info_span!(
"app_server.thread_start.attach_stdio_listener",
otel.name = "app_server.thread_start.attach_stdio_listener",
))
.await,
thread_id,
connection_id,
"thread",
);
}
listener_task_context
.thread_watch_manager
@@ -2244,9 +2267,10 @@ impl ThreadRequestProcessor {
&self,
connection_id: ConnectionId,
capabilities: ConnectionCapabilities,
origin: ConnectionOrigin,
) {
self.thread_state_manager
.connection_initialized(connection_id, capabilities)
.connection_initialized(connection_id, capabilities, origin)
.await;
}
codex-rs/app-server/src/request_processors/thread_processor_tests.rs
+136 -7
View File
@@ -48,6 +48,7 @@ mod thread_processor_behavior_tests {
use super::super::*;
use crate::outgoing_message::OutgoingEnvelope;
use crate::outgoing_message::OutgoingMessage;
use crate::transport::ConnectionOrigin;
use anyhow::Result;
use chrono::DateTime;
use chrono::Utc;
@@ -1127,7 +1128,11 @@ mod thread_processor_behavior_tests {
let (cancel_tx, cancel_rx) = oneshot::channel();
manager
.connection_initialized(connection, ConnectionCapabilities::default())
.connection_initialized(
connection,
ConnectionCapabilities::default(),
ConnectionOrigin::Stdio,
)
.await;
manager
.try_ensure_connection_subscribed(
@@ -1172,10 +1177,18 @@ mod thread_processor_behavior_tests {
let (cancel_tx, mut cancel_rx) = oneshot::channel();
manager
.connection_initialized(connection_a, ConnectionCapabilities::default())
.connection_initialized(
connection_a,
ConnectionCapabilities::default(),
ConnectionOrigin::Stdio,
)
.await;
manager
.connection_initialized(connection_b, ConnectionCapabilities::default())
.connection_initialized(
connection_b,
ConnectionCapabilities::default(),
ConnectionOrigin::Stdio,
)
.await;
manager
.try_ensure_connection_subscribed(
@@ -1221,10 +1234,18 @@ mod thread_processor_behavior_tests {
let connection_b = ConnectionId(2);
manager
.connection_initialized(connection_a, ConnectionCapabilities::default())
.connection_initialized(
connection_a,
ConnectionCapabilities::default(),
ConnectionOrigin::Stdio,
)
.await;
manager
.connection_initialized(connection_b, ConnectionCapabilities::default())
.connection_initialized(
connection_b,
ConnectionCapabilities::default(),
ConnectionOrigin::Stdio,
)
.await;
manager
.try_ensure_connection_subscribed(
@@ -1271,7 +1292,11 @@ mod thread_processor_behavior_tests {
let connection = ConnectionId(1);
manager
.connection_initialized(connection, ConnectionCapabilities::default())
.connection_initialized(
connection,
ConnectionCapabilities::default(),
ConnectionOrigin::Stdio,
)
.await;
let threads_to_unload = manager.remove_connection(connection).await;
assert_eq!(threads_to_unload, Vec::<ThreadId>::new());
@@ -1305,6 +1330,7 @@ mod thread_processor_behavior_tests {
ConnectionCapabilities {
request_attestation: true,
},
ConnectionOrigin::Stdio,
)
.await;
manager
@@ -1313,6 +1339,7 @@ mod thread_processor_behavior_tests {
ConnectionCapabilities {
request_attestation: true,
},
ConnectionOrigin::Stdio,
)
.await;
manager
@@ -1321,10 +1348,15 @@ mod thread_processor_behavior_tests {
ConnectionCapabilities {
request_attestation: true,
},
ConnectionOrigin::Stdio,
)
.await;
manager
.connection_initialized(unsupported_connection, ConnectionCapabilities::default())
.connection_initialized(
unsupported_connection,
ConnectionCapabilities::default(),
ConnectionOrigin::Stdio,
)
.await;
assert!(
@@ -1362,4 +1394,101 @@ mod thread_processor_behavior_tests {
);
Ok(())
}
#[tokio::test]
async fn remote_control_thread_start_also_subscribes_local_stdio_connections() -> Result<()> {
let manager = ThreadStateManager::new();
let thread_id = ThreadId::from_string("4f1ddc9b-e4f5-4c59-b0d7-1ea3db8c9e4a")?;
let remote_control_connection = ConnectionId(1);
let stdio_connection = ConnectionId(2);
let websocket_connection = ConnectionId(3);
manager
.connection_initialized(
remote_control_connection,
ConnectionCapabilities::default(),
ConnectionOrigin::RemoteControl,
)
.await;
manager
.connection_initialized(
stdio_connection,
ConnectionCapabilities::default(),
ConnectionOrigin::Stdio,
)
.await;
manager
.connection_initialized(
websocket_connection,
ConnectionCapabilities::default(),
ConnectionOrigin::WebSocket,
)
.await;
manager
.try_ensure_connection_subscribed(
thread_id,
remote_control_connection,
/*experimental_raw_events*/ false,
)
.await
.expect("remote-control connection should be live");
for connection_id in manager
.local_stdio_connection_ids_for_remote_control_thread_start(remote_control_connection)
.await
{
manager
.try_ensure_connection_subscribed(
thread_id,
connection_id,
/*experimental_raw_events*/ false,
)
.await
.expect("stdio companion should be live");
}
let mut subscribed_connection_ids = manager.subscribed_connection_ids(thread_id).await;
subscribed_connection_ids.sort_by_key(|connection_id| connection_id.0);
assert_eq!(
subscribed_connection_ids,
vec![remote_control_connection, stdio_connection]
);
Ok(())
}
#[tokio::test]
async fn non_remote_control_thread_start_does_not_add_stdio_companions() -> Result<()> {
let manager = ThreadStateManager::new();
let stdio_connection = ConnectionId(1);
let websocket_connection = ConnectionId(2);
manager
.connection_initialized(
stdio_connection,
ConnectionCapabilities::default(),
ConnectionOrigin::Stdio,
)
.await;
manager
.connection_initialized(
websocket_connection,
ConnectionCapabilities::default(),
ConnectionOrigin::WebSocket,
)
.await;
assert_eq!(
manager
.local_stdio_connection_ids_for_remote_control_thread_start(stdio_connection)
.await,
Vec::<ConnectionId>::new()
);
assert_eq!(
manager
.local_stdio_connection_ids_for_remote_control_thread_start(websocket_connection)
.await,
Vec::<ConnectionId>::new()
);
Ok(())
}
}
codex-rs/app-server/src/thread_state.rs
+40 -6
View File
@@ -1,5 +1,6 @@
use crate::outgoing_message::ConnectionId;
use crate::outgoing_message::ConnectionRequestId;
use crate::transport::ConnectionOrigin;
use codex_app_server_protocol::RequestId;
use codex_app_server_protocol::ThreadGoal;
use codex_app_server_protocol::ThreadHistoryBuilder;
@@ -205,7 +206,7 @@ impl ThreadEntry {
#[derive(Default)]
struct ThreadStateManagerInner {
live_connections: HashMap<ConnectionId, ConnectionCapabilities>,
live_connections: HashMap<ConnectionId, LiveConnection>,
threads: HashMap<ThreadId, ThreadEntry>,
thread_ids_by_connection: HashMap<ConnectionId, HashSet<ThreadId>>,
}
@@ -215,6 +216,12 @@ pub(crate) struct ConnectionCapabilities {
pub(crate) request_attestation: bool,
}
#[derive(Clone, Copy)]
struct LiveConnection {
capabilities: ConnectionCapabilities,
origin: ConnectionOrigin,
}
#[derive(Clone, Default)]
pub(crate) struct ThreadStateManager {
state: Arc<Mutex<ThreadStateManagerInner>>,
@@ -229,12 +236,15 @@ impl ThreadStateManager {
&self,
connection_id: ConnectionId,
capabilities: ConnectionCapabilities,
origin: ConnectionOrigin,
) {
self.state
.lock()
.await
.live_connections
.insert(connection_id, capabilities);
self.state.lock().await.live_connections.insert(
connection_id,
LiveConnection {
capabilities,
origin,
},
);
}
pub(crate) async fn first_attestation_capable_connection_for_thread(
@@ -251,12 +261,36 @@ impl ThreadStateManager {
state
.live_connections
.get(connection_id)?
.capabilities
.request_attestation
.then_some(*connection_id)
})
.min_by_key(|connection_id| connection_id.0)
}
pub(crate) async fn local_stdio_connection_ids_for_remote_control_thread_start(
&self,
request_connection_id: ConnectionId,
) -> Vec<ConnectionId> {
let state = self.state.lock().await;
let Some(request_connection) = state.live_connections.get(&request_connection_id) else {
return Vec::new();
};
if request_connection.origin != ConnectionOrigin::RemoteControl {
return Vec::new();
}
let mut connection_ids = state
.live_connections
.iter()
.filter_map(|(connection_id, live_connection)| {
(live_connection.origin == ConnectionOrigin::Stdio).then_some(*connection_id)
})
.collect::<Vec<_>>();
connection_ids.sort_by_key(|connection_id| connection_id.0);
connection_ids
}
pub(crate) async fn subscribed_connection_ids(&self, thread_id: ThreadId) -> Vec<ConnectionId> {
let state = self.state.lock().await;
state
codex-rs/app-server/src/transport.rs
+3 -1
View File
@@ -30,6 +30,7 @@ pub(crate) use codex_app_server_transport::start_stdio_connection;
pub(crate) use codex_app_server_transport::start_websocket_acceptor;
pub(crate) struct ConnectionState {
pub(crate) origin: ConnectionOrigin,
pub(crate) outbound_initialized: Arc<AtomicBool>,
pub(crate) outbound_experimental_api_enabled: Arc<AtomicBool>,
pub(crate) outbound_opted_out_notification_methods: Arc<RwLock<HashSet<String>>>,
@@ -38,12 +39,13 @@ pub(crate) struct ConnectionState {
impl ConnectionState {
pub(crate) fn new(
_origin: ConnectionOrigin,
origin: ConnectionOrigin,
outbound_initialized: Arc<AtomicBool>,
outbound_experimental_api_enabled: Arc<AtomicBool>,
outbound_opted_out_notification_methods: Arc<RwLock<HashSet<String>>>,
) -> Self {
Self {
origin,
outbound_initialized,
outbound_experimental_api_enabled,
outbound_opted_out_notification_methods,