openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-09-13 11:47:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

Improve app-server streaming event throughput

Browse Source
This commit is contained in:
Fabian Ponce
2026-06-23 03:39:05 +00:00
parent 27f22b54ae
commit b78df5cce7
11 changed files with 421 additions and 80 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
codex-rs/Cargo.lock generated
View File

@@ -2162,11 +2162,13 @@ dependencies = [
"codex-core",
"codex-login",
"codex-model-provider",
"codex-protocol",
"codex-state",
"codex-uds",
"codex-utils-absolute-path",
"codex-utils-rustls-provider",
"constant_time_eq 0.3.1",
"divan",
"futures",
"gethostname",
"hmac 0.12.1",

4
codex-rs/analytics/src/client.rs
View File

@@ -478,7 +478,7 @@ impl AnalyticsEventsClient {
});
}
pub fn track_notification(&self, notification: ServerNotification) {
pub fn track_notification(&self, notification: &ServerNotification) {
if !matches!(
notification,
ServerNotification::TurnStarted(_)
@@ -491,7 +491,7 @@ impl AnalyticsEventsClient {
) {
return;
}
self.record_fact(AnalyticsFact::Notification(Box::new(notification)));
self.record_fact(AnalyticsFact::Notification(Box::new(notification.clone())));
}
}

28
codex-rs/app-server-protocol/src/protocol/event_mapping.rs
View File

@@ -432,7 +432,10 @@ pub fn item_event_to_server_notification(
}
EventMsg::ExecCommandOutputDelta(exec_command_output_delta_event) => {
let item_id = exec_command_output_delta_event.call_id;
let delta = String::from_utf8_lossy(&exec_command_output_delta_event.chunk).to_string();
let delta = match String::from_utf8(exec_command_output_delta_event.chunk) {
Ok(delta) => delta,
Err(error) => String::from_utf8_lossy(error.as_bytes()).into_owned(),
};
ServerNotification::CommandExecutionOutputDelta(
CommandExecutionOutputDeltaNotification {
thread_id,
@@ -608,4 +611,27 @@ mod tests {
},
);
}
#[test]
fn exec_command_output_delta_preserves_lossy_utf8_mapping() {
let notification = item_event_to_server_notification(
EventMsg::ExecCommandOutputDelta(ExecCommandOutputDeltaEvent {
call_id: "call-1".to_string(),
stream: ExecOutputStream::Stdout,
chunk: b"hello\xFFworld".to_vec(),
}),
"thread-1",
"turn-1",
);
assert_command_execution_output_delta_server_notification(
notification,
CommandExecutionOutputDeltaNotification {
thread_id: "thread-1".to_string(),
turn_id: "turn-1".to_string(),
item_id: "call-1".to_string(),
delta: "hello\u{FFFD}world".to_string(),
},
);
}
}

6
codex-rs/app-server-transport/Cargo.toml
View File

@@ -55,6 +55,12 @@ uuid = { workspace = true, features = ["serde", "v7"] }
[dev-dependencies]
chrono = { workspace = true }
codex-config = { workspace = true }
codex-protocol = { workspace = true }
divan = { workspace = true }
pretty_assertions = { workspace = true }
tempfile = { workspace = true }
tokio = { workspace = true, features = ["test-util"] }
[[bench]]
name = "streaming_events"
harness = false

102
codex-rs/app-server-transport/benches/streaming_events.rs Normal file
View File

@@ -0,0 +1,102 @@
use codex_app_server_protocol::item_event_to_server_notification;
use codex_app_server_transport::OutgoingMessage;
use codex_protocol::protocol::AgentMessageContentDeltaEvent;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::ExecCommandOutputDeltaEvent;
use codex_protocol::protocol::ExecOutputStream;
use divan::Bencher;
use divan::counter::BytesCount;
use divan::counter::ItemsCount;
const EVENTS_PER_BATCH: usize = 256;
const THREAD_ID: &str = "019ef47c-0000-7000-8000-000000000001";
const TURN_ID: &str = "019ef47c-0000-7000-8000-000000000002";
const ITEM_ID: &str = "item_019ef47c000070008000000000000003";
fn main() {
divan::main();
}
#[divan::bench(args = [16, 256, 4096])]
fn project_and_serialize_agent_message_delta(bencher: Bencher, delta_bytes: usize) {
// Input generation is excluded because the event already exists when the
// app-server receives it from core.
bencher
.counter(ItemsCount::new(1_usize))
.counter(BytesCount::new(delta_bytes))
.with_inputs(|| {
EventMsg::AgentMessageContentDelta(AgentMessageContentDeltaEvent {
thread_id: THREAD_ID.to_string(),
turn_id: TURN_ID.to_string(),
item_id: ITEM_ID.to_string(),
delta: "x".repeat(delta_bytes),
})
})
.bench_local_values(|event| {
let message = project_event(event);
serialize_outgoing_message(&message).len()
});
}
#[divan::bench(args = [64, 4096, 65536])]
fn project_and_serialize_command_output_delta(bencher: Bencher, chunk_bytes: usize) {
// Input generation is excluded because the event already exists when the
// app-server receives it from core.
bencher
.counter(ItemsCount::new(1_usize))
.counter(BytesCount::new(chunk_bytes))
.with_inputs(|| {
EventMsg::ExecCommandOutputDelta(ExecCommandOutputDeltaEvent {
call_id: ITEM_ID.to_string(),
stream: ExecOutputStream::Stdout,
chunk: vec![b'x'; chunk_bytes],
})
})
.bench_local_values(|event| {
let message = project_event(event);
serialize_outgoing_message(&message).len()
});
}
/// Measures dispatch preparation and encoding after a streaming event has
/// already been projected. The dispatcher clones for all but the final
/// subscriber and moves the original message to the final queue.
#[divan::bench(args = [1, 4, 16])]
fn serialize_agent_message_delta_fanout(bencher: Bencher, subscriber_count: usize) {
let message = project_event(EventMsg::AgentMessageContentDelta(
AgentMessageContentDeltaEvent {
thread_id: THREAD_ID.to_string(),
turn_id: TURN_ID.to_string(),
item_id: ITEM_ID.to_string(),
delta: "streaming assistant output".to_string(),
},
));
let serialized_messages = EVENTS_PER_BATCH * subscriber_count;
bencher
.counter(ItemsCount::new(serialized_messages))
.with_inputs(|| vec![message.clone(); EVENTS_PER_BATCH])
.bench_local_values(|messages| {
let mut encoded_bytes = 0;
for message in messages {
for _ in 1..subscriber_count {
let subscriber_message = message.clone();
encoded_bytes += serialize_outgoing_message(&subscriber_message).len();
}
encoded_bytes += serialize_outgoing_message(&message).len();
}
encoded_bytes
});
}
fn project_event(event: EventMsg) -> OutgoingMessage {
let notification = item_event_to_server_notification(event, THREAD_ID, TURN_ID);
OutgoingMessage::AppServerNotification(notification)
}
fn serialize_outgoing_message(message: &OutgoingMessage) -> String {
match serde_json::to_string(message) {
Ok(json) => json,
Err(error) => panic!("outgoing message should serialize: {error}"),
}
}

55
codex-rs/app-server-transport/src/transport/mod.rs
View File

@@ -256,17 +256,10 @@ async fn enqueue_incoming_message(
}
fn serialize_outgoing_message(outgoing_message: OutgoingMessage) -> Option<String> {
let value = match serde_json::to_value(outgoing_message) {
Ok(value) => value,
Err(err) => {
error!("Failed to convert OutgoingMessage to JSON value: {err}");
return None;
}
};
match serde_json::to_string(&value) {
match serde_json::to_string(&outgoing_message) {
Ok(json) => Some(json),
Err(err) => {
error!("Failed to serialize JSONRPCMessage: {err}");
error!("Failed to serialize OutgoingMessage: {err}");
None
}
}
@@ -275,12 +268,15 @@ fn serialize_outgoing_message(outgoing_message: OutgoingMessage) -> Option<Strin
#[cfg(test)]
mod tests {
use super::*;
use crate::outgoing_message::OutgoingResponse;
use codex_app_server_protocol::ConfigWarningNotification;
use codex_app_server_protocol::CurrentTimeReadParams;
use codex_app_server_protocol::JSONRPCNotification;
use codex_app_server_protocol::JSONRPCRequest;
use codex_app_server_protocol::JSONRPCResponse;
use codex_app_server_protocol::RequestId;
use codex_app_server_protocol::ServerNotification;
use codex_app_server_protocol::ServerRequest;
use pretty_assertions::assert_eq;
use serde_json::json;
use tokio::time::Duration;
@@ -294,6 +290,47 @@ mod tests {
);
}
#[test]
fn direct_outgoing_serialization_preserves_json_values() {
let messages = [
OutgoingMessage::Request(ServerRequest::CurrentTimeRead {
request_id: RequestId::Integer(1),
params: CurrentTimeReadParams {
thread_id: "thread-1".to_string(),
},
}),
OutgoingMessage::AppServerNotification(ServerNotification::ConfigWarning(
ConfigWarningNotification {
summary: "warning".to_string(),
details: Some("details".to_string()),
path: None,
range: None,
},
)),
OutgoingMessage::Response(OutgoingResponse {
id: RequestId::Integer(2),
result: json!({"ok": true}),
}),
OutgoingMessage::Error(OutgoingError {
id: RequestId::Integer(3),
error: JSONRPCErrorError {
code: -32000,
message: "error".to_string(),
data: Some(json!({"retryable": false})),
},
}),
];
for message in messages {
let expected = serde_json::to_value(&message).expect("message should serialize");
let serialized =
serialize_outgoing_message(message).expect("message should serialize directly");
let actual: serde_json::Value =
serde_json::from_str(&serialized).expect("message should be valid JSON");
assert_eq!(actual, expected);
}
}
#[tokio::test]
async fn enqueue_incoming_request_returns_overload_error_when_queue_is_full() {
let connection_id = ConnectionId(42);

88
codex-rs/app-server/src/bespoke_event_handling.rs
View File

@@ -139,13 +139,13 @@ struct CommandExecutionCompletionItem {
pub(crate) async fn apply_bespoke_event_handling(
event: Event,
conversation_id: ThreadId,
conversation: Arc<CodexThread>,
thread_manager: Arc<ThreadManager>,
conversation: &Arc<CodexThread>,
thread_manager: &Arc<ThreadManager>,
outgoing: ThreadScopedOutgoingMessageSender,
thread_state: Arc<tokio::sync::Mutex<ThreadState>>,
thread_watch_manager: ThreadWatchManager,
thread_list_state_permit: Arc<tokio::sync::Semaphore>,
fallback_model_provider: String,
thread_state: &Arc<tokio::sync::Mutex<ThreadState>>,
thread_watch_manager: &ThreadWatchManager,
thread_list_state_permit: &Arc<tokio::sync::Semaphore>,
fallback_model_provider: &str,
) {
let Event {
id: event_turn_id,
@@ -185,7 +185,7 @@ pub(crate) async fn apply_bespoke_event_handling(
EventMsg::TurnComplete(turn_complete_event) => {
// All per-thread requests are bound to a turn, so abort them.
outgoing.abort_pending_server_requests().await;
respond_to_pending_interrupts(&thread_state, &outgoing).await;
respond_to_pending_interrupts(thread_state, &outgoing).await;
let turn_failed = thread_state.lock().await.turn_summary.last_error.is_some();
thread_watch_manager
.note_turn_completed(&conversation_id.to_string(), turn_failed)
@@ -195,7 +195,7 @@ pub(crate) async fn apply_bespoke_event_handling(
event_turn_id,
turn_complete_event,
&outgoing,
&thread_state,
thread_state,
)
.await;
}
@@ -280,7 +280,7 @@ pub(crate) async fn apply_bespoke_event_handling(
completion_item.command_actions.clone(),
CommandExecutionSource::Agent,
&outgoing,
&thread_state,
thread_state,
)
.await;
}
@@ -315,7 +315,7 @@ pub(crate) async fn apply_bespoke_event_handling(
completion_item.command_actions,
completion_status,
&outgoing,
&thread_state,
thread_state,
)
.await;
}
@@ -543,13 +543,15 @@ pub(crate) async fn apply_bespoke_event_handling(
let (pending_request_id, rx) = outgoing
.send_request(ServerRequestPayload::FileChangeRequestApproval(params))
.await;
let conversation = Arc::clone(conversation);
let thread_state = Arc::clone(thread_state);
tokio::spawn(async move {
on_file_change_request_approval_response(
item_id,
pending_request_id,
rx,
conversation,
thread_state.clone(),
thread_state,
permission_guard,
)
.await;
@@ -623,7 +625,7 @@ pub(crate) async fn apply_bespoke_event_handling(
completion_item.command_actions.clone(),
CommandExecutionSource::Agent,
&outgoing,
&thread_state,
thread_state,
)
.await;
}
@@ -661,6 +663,8 @@ pub(crate) async fn apply_bespoke_event_handling(
params,
))
.await;
let conversation = Arc::clone(conversation);
let thread_state = Arc::clone(thread_state);
tokio::spawn(async move {
on_command_execution_request_approval_response(
event_turn_id,
@@ -672,7 +676,7 @@ pub(crate) async fn apply_bespoke_event_handling(
rx,
conversation,
outgoing,
thread_state.clone(),
thread_state,
permission_guard,
)
.await;
@@ -712,6 +716,8 @@ pub(crate) async fn apply_bespoke_event_handling(
let (pending_request_id, rx) = outgoing
.send_request(ServerRequestPayload::ToolRequestUserInput(params))
.await;
let conversation = Arc::clone(conversation);
let thread_state = Arc::clone(thread_state);
tokio::spawn(async move {
on_request_user_input_response(
event_turn_id,
@@ -769,6 +775,8 @@ pub(crate) async fn apply_bespoke_event_handling(
let (pending_request_id, rx) = outgoing
.send_request(ServerRequestPayload::McpServerElicitationRequest(params))
.await;
let conversation = Arc::clone(conversation);
let thread_state = Arc::clone(thread_state);
tokio::spawn(async move {
on_mcp_server_elicitation_response(
request.server_name,
@@ -815,6 +823,8 @@ pub(crate) async fn apply_bespoke_event_handling(
receiver: rx,
request_permissions_guard: permission_guard,
};
let conversation = Arc::clone(conversation);
let thread_state = Arc::clone(thread_state);
tokio::spawn(async move {
on_request_permissions_response(pending_response, conversation, thread_state).await;
});
@@ -855,6 +865,7 @@ pub(crate) async fn apply_bespoke_event_handling(
let (_pending_request_id, rx) = outgoing
.send_request(ServerRequestPayload::DynamicToolCall(params))
.await;
let conversation = Arc::clone(conversation);
tokio::spawn(async move {
crate::dynamic_tools::on_call_response(call_id, rx, conversation).await;
});
@@ -955,7 +966,7 @@ pub(crate) async fn apply_bespoke_event_handling(
return handle_thread_rollback_failed(
conversation_id,
message,
&thread_state,
thread_state,
&outgoing,
)
.await;
@@ -975,7 +986,7 @@ pub(crate) async fn apply_bespoke_event_handling(
&event_turn_id,
turn_error,
&outgoing,
&thread_state,
thread_state,
)
.await;
}
@@ -1167,7 +1178,7 @@ pub(crate) async fn apply_bespoke_event_handling(
EventMsg::TurnAborted(turn_aborted_event) => {
// All per-thread requests are bound to a turn, so abort them.
outgoing.abort_pending_server_requests().await;
respond_to_pending_interrupts(&thread_state, &outgoing).await;
respond_to_pending_interrupts(thread_state, &outgoing).await;
thread_watch_manager
.note_turn_interrupted(&conversation_id.to_string())
@@ -1177,7 +1188,7 @@ pub(crate) async fn apply_bespoke_event_handling(
event_turn_id,
turn_aborted_event,
&outgoing,
&thread_state,
thread_state,
)
.await;
}
@@ -1228,7 +1239,7 @@ pub(crate) async fn apply_bespoke_event_handling(
let response = match thread_rollback_response_from_stored_thread(
stored_thread,
conversation.session_configured().session_id.to_string(),
fallback_model_provider.as_str(),
fallback_model_provider,
&fallback_cwd,
loaded_status,
) {
@@ -2408,19 +2419,21 @@ mod tests {
impl GuardianAssessmentTestContext {
async fn apply_guardian_assessment_event(&self, assessment: GuardianAssessmentEvent) {
let event_turn_id = assessment.turn_id.clone();
let thread_list_state_permit =
Arc::new(tokio::sync::Semaphore::new(/*permits*/ 1));
apply_bespoke_event_handling(
Event {
id: event_turn_id,
msg: EventMsg::GuardianAssessment(assessment),
},
self.conversation_id,
self.conversation.clone(),
self.thread_manager.clone(),
&self.conversation,
&self.thread_manager,
self.outgoing.clone(),
self.thread_state.clone(),
self.thread_watch_manager.clone(),
Arc::new(tokio::sync::Semaphore::new(/*permits*/ 1)),
"test-provider".to_string(),
&self.thread_state,
&self.thread_watch_manager,
&thread_list_state_permit,
"test-provider",
)
.await;
}
@@ -3360,6 +3373,7 @@ mod tests {
vec![ConnectionId(1)],
conversation_id,
);
let thread_list_state_permit = Arc::new(tokio::sync::Semaphore::new(/*permits*/ 1));
apply_bespoke_event_handling(
Event {
@@ -3373,13 +3387,13 @@ mod tests {
}),
},
conversation_id,
conversation,
thread_manager,
&conversation,
&thread_manager,
outgoing,
thread_state,
thread_watch_manager,
Arc::new(tokio::sync::Semaphore::new(/*permits*/ 1)),
"test-provider".to_string(),
&thread_state,
&thread_watch_manager,
&thread_list_state_permit,
"test-provider",
)
.await;
@@ -3429,6 +3443,8 @@ mod tests {
vec![ConnectionId(1)],
conversation_id,
);
let thread_state = new_thread_state();
let thread_list_state_permit = Arc::new(tokio::sync::Semaphore::new(/*permits*/ 1));
apply_bespoke_event_handling(
Event {
@@ -3443,13 +3459,13 @@ mod tests {
}),
},
conversation_id,
conversation,
thread_manager,
&conversation,
&thread_manager,
outgoing,
new_thread_state(),
thread_watch_manager.clone(),
Arc::new(tokio::sync::Semaphore::new(/*permits*/ 1)),
"test-provider".to_string(),
&thread_state,
&thread_watch_manager,
&thread_list_state_permit,
"test-provider",
)
.await;

69
codex-rs/app-server/src/outgoing_message.rs
View File

@@ -107,7 +107,7 @@ pub(crate) struct OutgoingMessageSender {
#[derive(Clone)]
pub(crate) struct ThreadScopedOutgoingMessageSender {
outgoing: Arc<OutgoingMessageSender>,
connection_ids: Arc<Vec<ConnectionId>>,
connection_ids: Arc<[ConnectionId]>,
thread_id: ThreadId,
}
@@ -120,12 +120,12 @@ struct PendingCallbackEntry {
impl ThreadScopedOutgoingMessageSender {
pub(crate) fn new(
outgoing: Arc<OutgoingMessageSender>,
connection_ids: Vec<ConnectionId>,
connection_ids: impl Into<Arc<[ConnectionId]>>,
thread_id: ThreadId,
) -> Self {
Self {
outgoing,
connection_ids: Arc::new(connection_ids),
connection_ids: connection_ids.into(),
thread_id,
}
}
@@ -136,7 +136,7 @@ impl ThreadScopedOutgoingMessageSender {
) -> (RequestId, oneshot::Receiver<ClientRequestResult>) {
self.outgoing
.send_request_to_connections(
Some(self.connection_ids.as_slice()),
Some(self.connection_ids.as_ref()),
payload,
Some(self.thread_id),
)
@@ -160,12 +160,12 @@ impl ThreadScopedOutgoingMessageSender {
pub(crate) async fn send_server_notification(&self, notification: ServerNotification) {
self.outgoing
.analytics_events_client
.track_notification(notification.clone());
.track_notification(&notification);
if self.connection_ids.is_empty() {
return;
}
self.outgoing
.send_server_notification_to_connections(self.connection_ids.as_slice(), notification)
.send_server_notification_to_connections(self.connection_ids.as_ref(), notification)
.await;
}
@@ -564,7 +564,7 @@ impl OutgoingMessageSender {
targeted_connections = connection_ids.len(),
"app-server event: {notification}"
);
let outgoing_message = OutgoingMessage::AppServerNotification(notification.clone());
let outgoing_message = OutgoingMessage::AppServerNotification(notification);
if connection_ids.is_empty() {
if let Err(err) = self
.sender
@@ -577,7 +577,10 @@ impl OutgoingMessageSender {
}
return;
}
for connection_id in connection_ids {
let Some((last_connection_id, other_connection_ids)) = connection_ids.split_last() else {
return;
};
for connection_id in other_connection_ids {
if let Err(err) = self
.sender
.send(OutgoingEnvelope::ToConnection {
@@ -590,6 +593,17 @@ impl OutgoingMessageSender {
warn!("failed to send server notification to client: {err:?}");
}
}
if let Err(err) = self
.sender
.send(OutgoingEnvelope::ToConnection {
connection_id: *last_connection_id,
message: outgoing_message,
write_complete_tx: None,
})
.await
{
warn!("failed to send server notification to client: {err:?}");
}
}
pub(crate) async fn send_server_notification_to_connection_and_wait(
@@ -598,7 +612,7 @@ impl OutgoingMessageSender {
notification: ServerNotification,
) {
tracing::trace!("app-server event: {notification}");
let outgoing_message = OutgoingMessage::AppServerNotification(notification.clone());
let outgoing_message = OutgoingMessage::AppServerNotification(notification);
let (write_complete_tx, write_complete_rx) = oneshot::channel();
if let Err(err) = self
.sender
@@ -1116,6 +1130,43 @@ mod tests {
}
}
#[tokio::test]
async fn send_server_notification_fanout_preserves_connection_order() {
let (tx, mut rx) = mpsc::channel::<OutgoingEnvelope>(4);
let outgoing =
OutgoingMessageSender::new(tx, codex_analytics::AnalyticsEventsClient::disabled());
outgoing
.send_server_notification_to_connections(
&[ConnectionId(3), ConnectionId(5), ConnectionId(8)],
ServerNotification::ModelRerouted(ModelReroutedNotification {
thread_id: "thread-1".to_string(),
turn_id: "turn-1".to_string(),
from_model: "gpt-5.3-codex".to_string(),
to_model: "gpt-5.2".to_string(),
reason: ModelRerouteReason::HighRiskCyberActivity,
}),
)
.await;
for expected_connection_id in [ConnectionId(3), ConnectionId(5), ConnectionId(8)] {
let envelope = rx.recv().await.expect("fanout should queue one envelope");
let OutgoingEnvelope::ToConnection {
connection_id,
message,
..
} = envelope
else {
panic!("expected targeted server notification envelope");
};
assert_eq!(connection_id, expected_connection_id);
assert!(matches!(
message,
OutgoingMessage::AppServerNotification(ServerNotification::ModelRerouted(_))
));
}
}
#[tokio::test]
async fn send_server_notification_to_connection_and_wait_tracks_write_completion() {
let (tx, mut rx) = mpsc::channel::<OutgoingEnvelope>(4);

27
codex-rs/app-server/src/request_processors/thread_lifecycle.rs
View File

@@ -227,6 +227,15 @@ pub(super) async fn ensure_listener_task_running(
"thread {conversation_id} is closing; retry after the thread is closed"
)));
};
let Some(connection_ids_rx) = listener_task_context
.thread_state_manager
.subscribe_to_connection_ids(conversation_id)
.await
else {
return Err(invalid_request(format!(
"thread {conversation_id} is closing; retry after the thread is closed"
)));
};
let config = conversation.config().await;
let environments = conversation.environment_selections().await;
let watch_registration = listener_task_context
@@ -315,9 +324,7 @@ pub(super) async fn ensure_listener_task_running(
thread_state.track_current_turn_event(&event.id, &event.msg);
thread_state.experimental_raw_events
};
let subscribed_connection_ids = thread_state_manager
.subscribed_connection_ids(conversation_id)
.await;
let subscribed_connection_ids = connection_ids_rx.borrow().clone();
let thread_outgoing = ThreadScopedOutgoingMessageSender::new(
outgoing_for_task.clone(),
subscribed_connection_ids,
@@ -338,15 +345,15 @@ pub(super) async fn ensure_listener_task_running(
}
apply_bespoke_event_handling(
event.clone(),
event,
conversation_id,
conversation.clone(),
thread_manager.clone(),
&conversation,
&thread_manager,
thread_outgoing,
thread_state.clone(),
thread_watch_manager.clone(),
thread_list_state_permit.clone(),
fallback_model_provider.clone(),
&thread_state,
&thread_watch_manager,
&thread_list_state_permit,
&fallback_model_provider,
)
.await;
}

67
codex-rs/app-server/src/request_processors/thread_processor_tests.rs
View File

@@ -1328,8 +1328,8 @@ mod thread_processor_behavior_tests {
);
assert_eq!(
manager.subscribed_connection_ids(thread_id).await,
vec![connection_b]
manager.subscribed_connection_ids(thread_id).await.as_ref(),
&[connection_b]
);
Ok(())
}
@@ -1385,6 +1385,69 @@ mod thread_processor_behavior_tests {
Ok(())
}
#[tokio::test]
async fn connection_ids_watcher_replaces_immutable_snapshots() -> Result<()> {
let manager = ThreadStateManager::new();
let thread_id = ThreadId::from_string("ad7f0408-99b8-4f6e-a46f-bd0eec433370")?;
let connection = ConnectionId(1);
let _thread_state = manager.thread_state(thread_id).await;
let mut connection_ids = manager
.subscribe_to_connection_ids(thread_id)
.await
.expect("thread should have a connection-ids watcher");
let empty_snapshot = connection_ids.borrow().clone();
assert!(empty_snapshot.is_empty());
manager
.connection_initialized(connection, ConnectionCapabilities::default())
.await;
assert!(
manager
.try_add_connection_to_thread(thread_id, connection)
.await
);
tokio::time::timeout(Duration::from_secs(1), connection_ids.changed())
.await
.expect("timed out waiting for subscriber snapshot")
.expect("connection-ids watcher should remain open");
let subscribed_snapshot = connection_ids.borrow().clone();
assert_eq!(subscribed_snapshot.as_ref(), &[connection]);
assert!(empty_snapshot.is_empty());
assert!(
manager
.unsubscribe_connection_from_thread(thread_id, connection)
.await
);
tokio::time::timeout(Duration::from_secs(1), connection_ids.changed())
.await
.expect("timed out waiting for empty subscriber snapshot")
.expect("connection-ids watcher should remain open");
assert!(connection_ids.borrow().is_empty());
assert_eq!(subscribed_snapshot.as_ref(), &[connection]);
Ok(())
}
#[tokio::test]
async fn connection_ids_watcher_does_not_recreate_missing_thread_state() -> Result<()> {
let manager = ThreadStateManager::new();
let thread_id = ThreadId::from_string("ad7f0408-99b8-4f6e-a46f-bd0eec433370")?;
assert!(
manager
.subscribe_to_connection_ids(thread_id)
.await
.is_none()
);
assert!(
manager
.subscribe_to_has_connections(thread_id)
.await
.is_none()
);
Ok(())
}
#[tokio::test]
async fn wait_for_thread_subscriber_unblocks_after_connection_attaches() -> Result<()> {
let manager = ThreadStateManager::new();

53
codex-rs/app-server/src/thread_state.rs
View File

@@ -250,6 +250,7 @@ mod tests {
struct ThreadEntry {
state: Arc<Mutex<ThreadState>>,
connection_ids: HashSet<ConnectionId>,
connection_ids_watcher: watch::Sender<Arc<[ConnectionId]>>,
has_connections_watcher: watch::Sender<bool>,
}
@@ -258,13 +259,33 @@ impl Default for ThreadEntry {
Self {
state: Arc::new(Mutex::new(ThreadState::default())),
connection_ids: HashSet::new(),
connection_ids_watcher: watch::channel(Arc::from([])).0,
has_connections_watcher: watch::channel(false).0,
}
}
}
impl ThreadEntry {
fn update_has_connections(&self) {
fn insert_connection(&mut self, connection_id: ConnectionId) {
if self.connection_ids.insert(connection_id) {
self.publish_connections();
}
}
fn remove_connection(&mut self, connection_id: ConnectionId) {
if self.connection_ids.remove(&connection_id) {
self.publish_connections();
}
}
fn publish_connections(&self) {
self.connection_ids_watcher.send_replace(
self.connection_ids
.iter()
.copied()
.collect::<Vec<_>>()
.into(),
);
let _ = self.has_connections_watcher.send_if_modified(|current| {
let prev = *current;
*current = !self.connection_ids.is_empty();
@@ -348,15 +369,29 @@ impl ThreadStateManager {
}
}
pub(crate) async fn subscribed_connection_ids(&self, thread_id: ThreadId) -> Vec<ConnectionId> {
pub(crate) async fn subscribed_connection_ids(
&self,
thread_id: ThreadId,
) -> Arc<[ConnectionId]> {
let state = self.state.lock().await;
state
.threads
.get(&thread_id)
.map(|thread_entry| thread_entry.connection_ids.iter().copied().collect())
.map(|thread_entry| thread_entry.connection_ids_watcher.borrow().clone())
.unwrap_or_default()
}
pub(crate) async fn subscribe_to_connection_ids(
&self,
thread_id: ThreadId,
) -> Option<watch::Receiver<Arc<[ConnectionId]>>> {
let state = self.state.lock().await;
state
.threads
.get(&thread_id)
.map(|thread_entry| thread_entry.connection_ids_watcher.subscribe())
}
pub(crate) async fn thread_state(&self, thread_id: ThreadId) -> Arc<Mutex<ThreadState>> {
let mut state = self.state.lock().await;
state.threads.entry(thread_id).or_default().state.clone()
@@ -469,8 +504,7 @@ impl ThreadStateManager {
}
}
if let Some(thread_entry) = state.threads.get_mut(&thread_id) {
thread_entry.connection_ids.remove(&connection_id);
thread_entry.update_has_connections();
thread_entry.remove_connection(connection_id);
}
};
@@ -504,8 +538,7 @@ impl ThreadStateManager {
.or_default()
.insert(thread_id);
let thread_entry = state.threads.entry(thread_id).or_default();
thread_entry.connection_ids.insert(connection_id);
thread_entry.update_has_connections();
thread_entry.insert_connection(connection_id);
thread_entry.state.clone()
};
{
@@ -532,8 +565,7 @@ impl ThreadStateManager {
.or_default()
.insert(thread_id);
let thread_entry = state.threads.entry(thread_id).or_default();
thread_entry.connection_ids.insert(connection_id);
thread_entry.update_has_connections();
thread_entry.insert_connection(connection_id);
true
}
@@ -547,8 +579,7 @@ impl ThreadStateManager {
.unwrap_or_default();
for thread_id in &thread_ids {
if let Some(thread_entry) = state.threads.get_mut(thread_id) {
thread_entry.connection_ids.remove(&connection_id);
thread_entry.update_has_connections();
thread_entry.remove_connection(connection_id);
}
}
thread_ids