-
Notifications
You must be signed in to change notification settings - Fork 10.7k
[3/8] Add pushed exec process events #18020
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: dev/remote-mcp-exec-stdin
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -6,6 +6,7 @@ use arc_swap::ArcSwap; | |
| use codex_app_server_protocol::JSONRPCNotification; | ||
| use serde_json::Value; | ||
| use tokio::sync::Mutex; | ||
| use tokio::sync::broadcast; | ||
| use tokio::sync::watch; | ||
|
|
||
| use tokio::time::timeout; | ||
|
|
@@ -16,6 +17,7 @@ use crate::ProcessId; | |
| use crate::client_api::ExecServerClientConnectOptions; | ||
| use crate::client_api::RemoteExecServerConnectArgs; | ||
| use crate::connection::JsonRpcConnection; | ||
| use crate::process::ExecProcessEvent; | ||
| use crate::protocol::EXEC_CLOSED_METHOD; | ||
| use crate::protocol::EXEC_EXITED_METHOD; | ||
| use crate::protocol::EXEC_METHOD; | ||
|
|
@@ -53,6 +55,7 @@ use crate::protocol::INITIALIZE_METHOD; | |
| use crate::protocol::INITIALIZED_METHOD; | ||
| use crate::protocol::InitializeParams; | ||
| use crate::protocol::InitializeResponse; | ||
| use crate::protocol::ProcessOutputChunk; | ||
| use crate::protocol::ReadParams; | ||
| use crate::protocol::ReadResponse; | ||
| use crate::protocol::TerminateParams; | ||
|
|
@@ -65,6 +68,7 @@ use crate::rpc::RpcClientEvent; | |
|
|
||
| const CONNECT_TIMEOUT: Duration = Duration::from_secs(10); | ||
| const INITIALIZE_TIMEOUT: Duration = Duration::from_secs(10); | ||
| const PROCESS_EVENT_CHANNEL_CAPACITY: usize = 256; | ||
|
|
||
| impl Default for ExecServerClientConnectOptions { | ||
| fn default() -> Self { | ||
|
|
@@ -100,6 +104,7 @@ impl RemoteExecServerConnectArgs { | |
|
|
||
| pub(crate) struct SessionState { | ||
| wake_tx: watch::Sender<u64>, | ||
| event_tx: broadcast::Sender<ExecProcessEvent>, | ||
| failure: Mutex<Option<String>>, | ||
| } | ||
|
|
||
|
|
@@ -450,8 +455,10 @@ impl From<RpcCallError> for ExecServerError { | |
| impl SessionState { | ||
| fn new() -> Self { | ||
| let (wake_tx, _wake_rx) = watch::channel(0); | ||
| let (event_tx, _event_rx) = broadcast::channel(PROCESS_EVENT_CHANNEL_CAPACITY); | ||
| Self { | ||
| wake_tx, | ||
| event_tx, | ||
| failure: Mutex::new(None), | ||
| } | ||
| } | ||
|
|
@@ -460,19 +467,31 @@ impl SessionState { | |
| self.wake_tx.subscribe() | ||
| } | ||
|
|
||
| pub(crate) fn subscribe_events(&self) -> broadcast::Receiver<ExecProcessEvent> { | ||
| self.event_tx.subscribe() | ||
| } | ||
|
|
||
| fn note_change(&self, seq: u64) { | ||
| let next = (*self.wake_tx.borrow()).max(seq); | ||
| let _ = self.wake_tx.send(next); | ||
| } | ||
|
|
||
| fn publish_event(&self, event: ExecProcessEvent) { | ||
| let _ = self.event_tx.send(event); | ||
| } | ||
|
|
||
| async fn set_failure(&self, message: String) { | ||
| let mut failure = self.failure.lock().await; | ||
| if failure.is_none() { | ||
| *failure = Some(message); | ||
| let should_publish = failure.is_none(); | ||
| if should_publish { | ||
| *failure = Some(message.clone()); | ||
| } | ||
| drop(failure); | ||
| let next = (*self.wake_tx.borrow()).saturating_add(1); | ||
| let _ = self.wake_tx.send(next); | ||
| if should_publish { | ||
| self.publish_event(ExecProcessEvent::Failed(message)); | ||
| } | ||
| } | ||
|
|
||
| async fn failed_response(&self) -> Option<ReadResponse> { | ||
|
|
@@ -505,6 +524,10 @@ impl Session { | |
| self.state.subscribe() | ||
| } | ||
|
|
||
| pub(crate) fn subscribe_events(&self) -> broadcast::Receiver<ExecProcessEvent> { | ||
| self.state.subscribe_events() | ||
| } | ||
|
|
||
| pub(crate) async fn read( | ||
| &self, | ||
| after_seq: Option<u64>, | ||
|
|
@@ -628,13 +651,22 @@ async fn handle_server_notification( | |
| serde_json::from_value(notification.params.unwrap_or(Value::Null))?; | ||
| if let Some(session) = inner.get_session(¶ms.process_id) { | ||
| session.note_change(params.seq); | ||
| session.publish_event(ExecProcessEvent::Output(ProcessOutputChunk { | ||
| seq: params.seq, | ||
| stream: params.stream, | ||
| chunk: params.chunk, | ||
| })); | ||
| } | ||
| } | ||
| EXEC_EXITED_METHOD => { | ||
| let params: ExecExitedNotification = | ||
| serde_json::from_value(notification.params.unwrap_or(Value::Null))?; | ||
| if let Some(session) = inner.get_session(¶ms.process_id) { | ||
| session.note_change(params.seq); | ||
| session.publish_event(ExecProcessEvent::Exited { | ||
| seq: params.seq, | ||
| exit_code: params.exit_code, | ||
| }); | ||
| } | ||
| } | ||
| EXEC_CLOSED_METHOD => { | ||
|
|
@@ -645,6 +677,7 @@ async fn handle_server_notification( | |
| let session = inner.remove_session(¶ms.process_id).await; | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
On Useful? React with 👍 / 👎. |
||
| if let Some(session) = session { | ||
| session.note_change(params.seq); | ||
| session.publish_event(ExecProcessEvent::Closed { seq: params.seq }); | ||
| } | ||
| } | ||
| other => { | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
SessionState::newcreates a broadcast channel and immediately drops the initial receiver. If a short-lived process outputs/exits beforesubscribe_events()is called,publish_eventdrops those events (no receivers), yet the sender remains alive on the process handle. An event-only consumer can then wait forever forExited/Closed, which is a functional hang.Useful? React with 👍 / 👎.