clone/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-04-30 01:16:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

app-server: thread resume subscriptions (#11474)

Browse Source 
This stack layer makes app-server thread event delivery connection-aware
so resumed/attached threads only emit notifications and approval prompts
to subscribed connections.

- Added per-thread subscription tracking in `ThreadState`
(`subscribed_connections`) and mapped subscription ids to `(thread_id,
connection_id)`.
- Updated listener lifecycle so removing a subscription or closing a
connection only removes that connection from the thread’s subscriber
set; listener shutdown now happens when the last subscriber is gone.
- Added `connection_closed(connection_id)` plumbing (`lib.rs` ->
`message_processor.rs` -> `codex_message_processor.rs`) so disconnect
cleanup happens immediately.
- Scoped bespoke event handling outputs through `TargetedOutgoing` to
send requests/notifications only to subscribed connections.
- Kept existing threadresume behavior while aligning with the latest
split-loop transport structure.
This commit is contained in:
Max Johnson
2026-02-11 16:21:13 -08:00
committed by GitHub
parent 703fb38d2a
commit c0ecc2e1e1
7 changed files with 648 additions and 150 deletions
Download Patch File Download Diff File Expand all files Collapse all files

302
codex-rs/app-server/src/codex_message_processor.rs
View File

@@ -7,6 +7,7 @@ use crate::outgoing_message::ConnectionId;
use crate::outgoing_message::ConnectionRequestId;
use crate::outgoing_message::OutgoingMessageSender;
use crate::outgoing_message::OutgoingNotification;
use crate::outgoing_message::ThreadScopedOutgoingMessageSender;
use chrono::DateTime;
use chrono::SecondsFormat;
use chrono::Utc;
@@ -251,6 +252,7 @@ use uuid::Uuid;
use crate::filters::compute_source_filters;
use crate::filters::source_kind_matches;
use crate::thread_state::ThreadState;
use crate::thread_state::ThreadStateManager;
const THREAD_LIST_DEFAULT_LIMIT: usize = 25;
@@ -1961,8 +1963,9 @@ impl CodexMessageProcessor {
// Auto-attach a thread listener when starting a thread.
// Use the same behavior as the v1 API, with opt-in support for raw item events.
if let Err(err) = self
.attach_conversation_listener(
.ensure_conversation_listener(
thread_id,
request_id.connection_id,
experimental_raw_events,
ApiVersion::V2,
)
@@ -2628,20 +2631,28 @@ impl CodexMessageProcessor {
self.thread_manager.subscribe_thread_created()
}
/// Best-effort: attach a listener for thread_id if missing.
pub(crate) async fn try_attach_thread_listener(&mut self, thread_id: ThreadId) {
if self.thread_state_manager.has_listener_for_thread(thread_id) {
return;
}
pub(crate) async fn connection_closed(&mut self, connection_id: ConnectionId) {
self.thread_state_manager
.remove_connection(connection_id)
.await;
}
if let Err(err) = self
.attach_conversation_listener(thread_id, false, ApiVersion::V2)
.await
{
warn!(
"failed to attach listener for thread {thread_id}: {message}",
message = err.message
);
/// Best-effort: ensure initialized connections are subscribed to this thread.
pub(crate) async fn try_attach_thread_listener(
&mut self,
thread_id: ThreadId,
connection_ids: Vec<ConnectionId>,
) {
for connection_id in connection_ids {
if let Err(err) = self
.ensure_conversation_listener(thread_id, connection_id, false, ApiVersion::V2)
.await
{
warn!(
"failed to auto-attach listener for thread {thread_id}: {message}",
message = err.message
);
}
}
}
@@ -2793,7 +2804,12 @@ impl CodexMessageProcessor {
};
// Auto-attach a thread listener when resuming a thread.
if let Err(err) = self
.attach_conversation_listener(thread_id, false, ApiVersion::V2)
.ensure_conversation_listener(
thread_id,
request_id.connection_id,
false,
ApiVersion::V2,
)
.await
{
tracing::warn!(
@@ -3019,7 +3035,12 @@ impl CodexMessageProcessor {
};
// Auto-attach a conversation listener when forking a thread.
if let Err(err) = self
.attach_conversation_listener(thread_id, false, ApiVersion::V2)
.ensure_conversation_listener(
thread_id,
request_id.connection_id,
false,
ApiVersion::V2,
)
.await
{
tracing::warn!(
@@ -5136,7 +5157,12 @@ impl CodexMessageProcessor {
})?;
if let Err(err) = self
.attach_conversation_listener(thread_id, false, ApiVersion::V2)
.ensure_conversation_listener(
thread_id,
request_id.connection_id,
false,
ApiVersion::V2,
)
.await
{
tracing::warn!(
@@ -5281,18 +5307,38 @@ impl CodexMessageProcessor {
conversation_id,
experimental_raw_events,
} = params;
match self
.attach_conversation_listener(conversation_id, experimental_raw_events, ApiVersion::V1)
.await
{
Ok(subscription_id) => {
let response = AddConversationSubscriptionResponse { subscription_id };
self.outgoing.send_response(request_id, response).await;
let conversation = match self.thread_manager.get_thread(conversation_id).await {
Ok(conv) => conv,
Err(_) => {
let error = JSONRPCErrorError {
code: INVALID_REQUEST_ERROR_CODE,
message: format!("thread not found: {conversation_id}"),
data: None,
};
self.outgoing.send_error(request_id, error).await;
return;
}
Err(err) => {
self.outgoing.send_error(request_id, err).await;
}
}
};
let subscription_id = Uuid::new_v4();
let thread_state = self
.thread_state_manager
.set_listener(
subscription_id,
conversation_id,
request_id.connection_id,
experimental_raw_events,
)
.await;
self.ensure_listener_task_running(
conversation_id,
conversation,
thread_state,
ApiVersion::V1,
)
.await;
let response = AddConversationSubscriptionResponse { subscription_id };
self.outgoing.send_response(request_id, response).await;
}
async fn remove_thread_listener(
@@ -5322,12 +5368,13 @@ impl CodexMessageProcessor {
}
}
async fn attach_conversation_listener(
async fn ensure_conversation_listener(
&mut self,
conversation_id: ThreadId,
connection_id: ConnectionId,
raw_events_enabled: bool,
api_version: ApiVersion,
) -> Result<Uuid, JSONRPCErrorError> {
) -> Result<(), JSONRPCErrorError> {
let conversation = match self.thread_manager.get_thread(conversation_id).await {
Ok(conv) => conv,
Err(_) => {
@@ -5338,13 +5385,30 @@ impl CodexMessageProcessor {
});
}
};
let subscription_id = Uuid::new_v4();
let (cancel_tx, mut cancel_rx) = oneshot::channel();
let thread_state = self
.thread_state_manager
.set_listener(subscription_id, conversation_id, cancel_tx)
.ensure_connection_subscribed(conversation_id, connection_id, raw_events_enabled)
.await;
self.ensure_listener_task_running(conversation_id, conversation, thread_state, api_version)
.await;
Ok(())
}
async fn ensure_listener_task_running(
&self,
conversation_id: ThreadId,
conversation: Arc<CodexThread>,
thread_state: Arc<Mutex<ThreadState>>,
api_version: ApiVersion,
) {
let (cancel_tx, mut cancel_rx) = oneshot::channel();
{
let mut thread_state = thread_state.lock().await;
if thread_state.listener_matches(&conversation) {
return;
}
thread_state.set_listener(cancel_tx, &conversation);
}
let outgoing_for_task = self.outgoing.clone();
let fallback_model_provider = self.config.model_provider_id.clone();
tokio::spawn(async move {
@@ -5363,10 +5427,6 @@ impl CodexMessageProcessor {
}
};
if let EventMsg::RawResponseItem(_) = &event.msg && !raw_events_enabled {
continue;
}
// For now, we send a notification for every event,
// JSON-serializing the `Event` as-is, but these should
// be migrated to be variants of `ServerNotification`
@@ -5391,19 +5451,38 @@ impl CodexMessageProcessor {
"conversationId".to_string(),
conversation_id.to_string().into(),
);
let (subscribed_connection_ids, raw_events_enabled) = {
let thread_state = thread_state.lock().await;
(
thread_state.subscribed_connection_ids(),
thread_state.experimental_raw_events,
)
};
if let EventMsg::RawResponseItem(_) = &event.msg && !raw_events_enabled {
continue;
}
outgoing_for_task
.send_notification(OutgoingNotification {
method: format!("codex/event/{event_formatted}"),
params: Some(params.into()),
})
.await;
if !subscribed_connection_ids.is_empty() {
outgoing_for_task
.send_notification_to_connections(
&subscribed_connection_ids,
OutgoingNotification {
method: format!("codex/event/{event_formatted}"),
params: Some(params.into()),
},
)
.await;
}
let thread_outgoing = ThreadScopedOutgoingMessageSender::new(
outgoing_for_task.clone(),
subscribed_connection_ids,
);
apply_bespoke_event_handling(
event.clone(),
conversation_id,
conversation.clone(),
outgoing_for_task.clone(),
thread_outgoing,
thread_state.clone(),
api_version,
fallback_model_provider.clone(),
@@ -5413,9 +5492,7 @@ impl CodexMessageProcessor {
}
}
});
Ok(subscription_id)
}
async fn git_diff_to_origin(&self, request_id: ConnectionRequestId, cwd: PathBuf) {
let diff = git_diff_to_remote(&cwd).await;
match diff {
@@ -6299,22 +6376,137 @@ mod tests {
let thread_id = ThreadId::from_string("ad7f0408-99b8-4f6e-a46f-bd0eec433370")?;
let listener_a = Uuid::new_v4();
let listener_b = Uuid::new_v4();
let (cancel_a, cancel_rx_a) = oneshot::channel();
let (cancel_b, mut cancel_rx_b) = oneshot::channel();
let connection_a = ConnectionId(1);
let connection_b = ConnectionId(2);
let (cancel_tx, mut cancel_rx) = oneshot::channel();
manager.set_listener(listener_a, thread_id, cancel_a).await;
manager.set_listener(listener_b, thread_id, cancel_b).await;
manager
.set_listener(listener_a, thread_id, connection_a, false)
.await;
manager
.set_listener(listener_b, thread_id, connection_b, false)
.await;
{
let state = manager.thread_state(thread_id);
state.lock().await.cancel_tx = Some(cancel_tx);
}
assert_eq!(manager.remove_listener(listener_a).await, Some(thread_id));
assert_eq!(cancel_rx_a.await, Ok(()));
assert!(
tokio::time::timeout(Duration::from_millis(20), &mut cancel_rx_b)
tokio::time::timeout(Duration::from_millis(20), &mut cancel_rx)
.await
.is_err()
);
assert_eq!(manager.remove_listener(listener_b).await, Some(thread_id));
assert_eq!(cancel_rx.await, Ok(()));
Ok(())
}
#[tokio::test]
async fn removing_listener_unsubscribes_its_connection() -> Result<()> {
let mut manager = ThreadStateManager::new();
let thread_id = ThreadId::from_string("ad7f0408-99b8-4f6e-a46f-bd0eec433370")?;
let listener_a = Uuid::new_v4();
let listener_b = Uuid::new_v4();
let connection_a = ConnectionId(1);
let connection_b = ConnectionId(2);
manager
.set_listener(listener_a, thread_id, connection_a, false)
.await;
manager
.set_listener(listener_b, thread_id, connection_b, false)
.await;
assert_eq!(manager.remove_listener(listener_a).await, Some(thread_id));
let state = manager.thread_state(thread_id);
let subscribed_connection_ids = state.lock().await.subscribed_connection_ids();
assert_eq!(subscribed_connection_ids, vec![connection_b]);
Ok(())
}
#[tokio::test]
async fn set_listener_uses_last_write_for_raw_events() -> Result<()> {
let mut manager = ThreadStateManager::new();
let thread_id = ThreadId::from_string("ad7f0408-99b8-4f6e-a46f-bd0eec433370")?;
let listener_a = Uuid::new_v4();
let listener_b = Uuid::new_v4();
let connection_a = ConnectionId(1);
let connection_b = ConnectionId(2);
manager
.set_listener(listener_a, thread_id, connection_a, true)
.await;
{
let state = manager.thread_state(thread_id);
assert!(state.lock().await.experimental_raw_events);
}
manager
.set_listener(listener_b, thread_id, connection_b, false)
.await;
let state = manager.thread_state(thread_id);
assert!(!state.lock().await.experimental_raw_events);
Ok(())
}
#[tokio::test]
async fn removing_connection_clears_subscription_and_listener_when_last_subscriber()
-> Result<()> {
let mut manager = ThreadStateManager::new();
let thread_id = ThreadId::from_string("ad7f0408-99b8-4f6e-a46f-bd0eec433370")?;
let listener = Uuid::new_v4();
let connection = ConnectionId(1);
let (cancel_tx, cancel_rx) = oneshot::channel();
manager
.set_listener(listener, thread_id, connection, false)
.await;
{
let state = manager.thread_state(thread_id);
state.lock().await.cancel_tx = Some(cancel_tx);
}
manager.remove_connection(connection).await;
assert_eq!(cancel_rx.await, Ok(()));
assert_eq!(manager.remove_listener(listener).await, None);
let state = manager.thread_state(thread_id);
assert!(state.lock().await.subscribed_connection_ids().is_empty());
Ok(())
}
#[tokio::test]
async fn removing_auto_attached_connection_preserves_listener_for_other_connections()
-> Result<()> {
let mut manager = ThreadStateManager::new();
let thread_id = ThreadId::from_string("ad7f0408-99b8-4f6e-a46f-bd0eec433370")?;
let connection_a = ConnectionId(1);
let connection_b = ConnectionId(2);
let (cancel_tx, mut cancel_rx) = oneshot::channel();
manager
.ensure_connection_subscribed(thread_id, connection_a, false)
.await;
manager
.ensure_connection_subscribed(thread_id, connection_b, false)
.await;
{
let state = manager.thread_state(thread_id);
state.lock().await.cancel_tx = Some(cancel_tx);
}
manager.remove_connection(connection_a).await;
assert!(
tokio::time::timeout(Duration::from_millis(20), &mut cancel_rx)
.await
.is_err()
);
assert_eq!(manager.remove_listener(listener_b).await, Some(thread_id));
assert_eq!(cancel_rx_b.await, Ok(()));
let state = manager.thread_state(thread_id);
assert_eq!(
state.lock().await.subscribed_connection_ids(),
vec![connection_b]
);
Ok(())
}
}