agent/codex
1
0
Fork 0
mirror of https://github.com/pchuan98/codex.git synced 2026-07-01 00:31:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

app-server: track remote-control seq IDs per stream (#17902)

Browse Source 
## Summary

- Track outbound remote-control sequence IDs independently for each
client stream.
- Retain unacked outbound messages per stream using FIFO buffers.
- Require stream-scoped acks and update tests for contiguous per-stream
sequencing.

## Why

The remote-control peer uses outbound sequence gaps to detect lost
messages and re-initialize. A single global outbound sequence counter
can create apparent gaps on an individual stream when another stream
receives an interleaved message.

## Validation

- `just fmt`
- `cargo test -p codex-app-server remote_control`
- `just fix -p codex-app-server`
- `git diff --check`
This commit is contained in:
Ruslan Nigmatullin
2026-04-15 11:52:53 -07:00
committed by GitHub
Unverified
parent f5e8eac2ae
commit 83abf67d20
3 changed files with 216 additions and 98 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/app-server/src/transport/remote_control/protocol.rs
+3 -4
View File
@@ -48,9 +48,8 @@ pub enum ClientEvent {
message: JSONRPCMessage,
},
/// Backend-generated acknowledgement for all server envelopes addressed to
/// `client_id` whose envelope `seq_id` is less than or equal to this ack's
/// `seq_id`. This cursor is client-scoped, not stream-scoped, so receivers
/// must not use `stream_id` to partition acks.
/// `client_id` and `stream_id` whose envelope `seq_id` is less than or equal
/// to this ack's `seq_id`. This cursor is stream-scoped.
Ack,
Ping,
ClientClosed,
@@ -65,7 +64,7 @@ pub(crate) struct ClientEnvelope {
pub(crate) client_id: ClientId,
#[serde(rename = "stream_id", skip_serializing_if = "Option::is_none")]
pub(crate) stream_id: Option<StreamId>,
/// For `Ack`, this is the backend-generated per-client cursor over
/// For `Ack`, this is the backend-generated per-stream cursor over
/// `ServerEnvelope.seq_id`.
#[serde(rename = "seq_id", skip_serializing_if = "Option::is_none")]
pub(crate) seq_id: Option<u64>,
codex-rs/app-server/src/transport/remote_control/tests.rs
+23 -13
View File
@@ -5,6 +5,7 @@ use super::enroll::update_persisted_remote_control_enrollment;
use super::protocol::ClientEnvelope;
use super::protocol::ClientEvent;
use super::protocol::ClientId;
use super::protocol::StreamId;
use super::protocol::normalize_remote_control_url;
use super::websocket::REMOTE_CONTROL_PROTOCOL_VERSION;
use super::*;
@@ -163,7 +164,7 @@ async fn remote_control_transport_manages_virtual_clients_and_routes_messages()
json!({
"type": "pong",
"client_id": "client-1",
"seq_id": 0,
"seq_id": 1,
"status": "unknown",
})
);
@@ -369,7 +370,7 @@ async fn remote_control_transport_manages_virtual_clients_and_routes_messages()
json!({
"type": "pong",
"client_id": "client-1",
"seq_id": 3,
"seq_id": 1,
"status": "unknown",
})
);
@@ -618,12 +619,13 @@ async fn remote_control_transport_clears_outgoing_buffer_when_backend_acks() {
))
.await
.expect("remote writer should accept outgoing message");
let (server_event, stream_id) = read_server_event_with_stream_id(&mut first_websocket).await;
assert_eq!(
read_server_event(&mut first_websocket).await,
server_event,
json!({
"type": "server_message",
"client_id": "client-1",
"seq_id": 0,
"seq_id": 1,
"message": {
"method": "configWarning",
"params": {
@@ -639,8 +641,8 @@ async fn remote_control_transport_clears_outgoing_buffer_when_backend_acks() {
ClientEnvelope {
event: ClientEvent::Ack,
client_id: client_id.clone(),
stream_id: None,
seq_id: Some(0),
stream_id: Some(stream_id),
seq_id: Some(1),
cursor: None,
},
)
@@ -854,7 +856,7 @@ async fn remote_control_http_mode_enrolls_before_connecting() {
json!({
"type": "server_message",
"client_id": backend_client_id.0.clone(),
"seq_id": 0,
"seq_id": 1,
"message": {
"id": 11,
"result": {
@@ -882,7 +884,7 @@ async fn remote_control_http_mode_enrolls_before_connecting() {
json!({
"type": "server_message",
"client_id": backend_client_id.0.clone(),
"seq_id": 1,
"seq_id": 2,
"message": {
"method": "configWarning",
"params": {
@@ -1356,6 +1358,12 @@ async fn send_client_event(
}
async fn read_server_event(websocket: &mut WebSocketStream<TcpStream>) -> serde_json::Value {
read_server_event_with_stream_id(websocket).await.0
}
async fn read_server_event_with_stream_id(
websocket: &mut WebSocketStream<TcpStream>,
) -> (serde_json::Value, StreamId) {
loop {
let frame = timeout(Duration::from_secs(5), websocket.next())
.await
@@ -1366,13 +1374,15 @@ async fn read_server_event(websocket: &mut WebSocketStream<TcpStream>) -> serde_
tungstenite::Message::Text(text) => {
let mut event: serde_json::Value =
serde_json::from_str(text.as_ref()).expect("server event should deserialize");
if let Some(stream_id) = event
let stream_id = event
.as_object_mut()
.and_then(|event| event.remove("stream_id"))
{
assert!(stream_id.is_string(), "stream_id should be a string");
}
return event;
.expect("stream_id should be present");
let stream_id = stream_id
.as_str()
.expect("stream_id should be a string")
.to_string();
return (event, StreamId(stream_id));
}
tungstenite::Message::Ping(payload) => {
websocket
codex-rs/app-server/src/transport/remote_control/websocket.rs
+190 -81
View File
@@ -14,6 +14,7 @@ use super::protocol::ClientEvent;
use super::protocol::ClientId;
use super::protocol::RemoteControlTarget;
use super::protocol::ServerEnvelope;
use super::protocol::StreamId;
use axum::http::HeaderValue;
use base64::Engine;
use codex_core::util::backoff;
@@ -25,8 +26,8 @@ use futures::SinkExt;
use futures::StreamExt;
use futures::stream::SplitSink;
use futures::stream::SplitStream;
use std::collections::BTreeMap;
use std::collections::HashMap;
use std::collections::VecDeque;
use std::io;
use std::io::ErrorKind;
use std::sync::Arc;
@@ -57,10 +58,7 @@ const REMOTE_CONTROL_ACCOUNT_ID_RETRY_INTERVAL: std::time::Duration =
std::time::Duration::from_secs(1);
struct BoundedOutboundBuffer {
// Remote-control acks are generated by the backend at client scope, so
// retransmit retention is keyed by client_id only. stream_id stays on each
// envelope for routing, but it is not part of the ack cursor.
buffer_by_client: HashMap<ClientId, BTreeMap<u64, ServerEnvelope>>,
buffer_by_stream: HashMap<(ClientId, StreamId), VecDeque<ServerEnvelope>>,
used_tx: watch::Sender<usize>,
}
@@ -68,46 +66,50 @@ impl BoundedOutboundBuffer {
fn new() -> (Self, watch::Receiver<usize>) {
let (used_tx, used_rx) = watch::channel(0);
let buffer = Self {
buffer_by_client: HashMap::new(),
buffer_by_stream: HashMap::new(),
used_tx,
};
(buffer, used_rx)
}
fn insert(&mut self, server_envelope: &ServerEnvelope) {
self.buffer_by_client
.entry(server_envelope.client_id.clone())
self.buffer_by_stream
.entry((
server_envelope.client_id.clone(),
server_envelope.stream_id.clone(),
))
.or_default()
.insert(server_envelope.seq_id, server_envelope.clone());
.push_back(server_envelope.clone());
self.used_tx.send_modify(|used| *used += 1);
}
fn ack(&mut self, client_id: &ClientId, acked_seq_id: u64) {
let Some(buffer) = self.buffer_by_client.get_mut(client_id) else {
fn ack(&mut self, client_id: &ClientId, stream_id: &StreamId, acked_seq_id: u64) {
let key = (client_id.clone(), stream_id.clone());
let Some(buffer) = self.buffer_by_stream.get_mut(&key) else {
return;
};
while let Some(seq_id) = buffer.first_key_value().map(|(seq_id, _)| seq_id)
&& *seq_id <= acked_seq_id
while let Some(server_envelope) = buffer.front()
&& server_envelope.seq_id <= acked_seq_id
{
buffer.pop_first();
buffer.pop_front();
self.used_tx.send_modify(|used| *used -= 1);
}
if buffer.is_empty() {
self.buffer_by_client.remove(client_id);
self.buffer_by_stream.remove(&key);
}
}
fn server_envelopes(&self) -> impl Iterator<Item = &ServerEnvelope> {
self.buffer_by_client
self.buffer_by_stream
.values()
.flat_map(|buffer| buffer.values())
.flat_map(|buffer| buffer.iter())
}
}
struct WebsocketState {
outbound_buffer: BoundedOutboundBuffer,
subscribe_cursor: Option<String>,
next_seq_id: u64,
next_seq_id_by_stream: HashMap<(ClientId, StreamId), u64>,
}
pub(crate) struct RemoteControlWebsocket {
@@ -162,7 +164,7 @@ impl RemoteControlWebsocket {
state: Arc::new(Mutex::new(WebsocketState {
outbound_buffer,
subscribe_cursor: None,
next_seq_id: 0,
next_seq_id_by_stream: HashMap::new(),
})),
server_event_rx: Arc::new(Mutex::new(server_event_rx)),
used_rx,
@@ -445,10 +447,16 @@ impl RemoteControlWebsocket {
}
}
};
let (server_envelope, write_complete_tx) = {
let (payload, write_complete_tx) = {
let mut state = state.lock().await;
let seq_id = state.next_seq_id;
state.next_seq_id = state.next_seq_id.saturating_add(1);
let seq_key = (
queued_server_envelope.client_id.clone(),
queued_server_envelope.stream_id.clone(),
);
let seq_id = *state
.next_seq_id_by_stream
.entry(seq_key.clone())
.or_insert(1);
let server_envelope = ServerEnvelope {
event: queued_server_envelope.event,
@@ -456,17 +464,19 @@ impl RemoteControlWebsocket {
seq_id,
stream_id: queued_server_envelope.stream_id,
};
let payload = match serde_json::to_string(&server_envelope) {
Ok(payload) => payload,
Err(err) => {
error!("failed to serialize remote-control server event: {err}");
continue;
}
};
state
.next_seq_id_by_stream
.insert(seq_key, seq_id.saturating_add(1));
state.outbound_buffer.insert(&server_envelope);
(server_envelope, queued_server_envelope.write_complete_tx)
};
let payload = match serde_json::to_string(&server_envelope) {
Ok(payload) => payload,
Err(err) => {
error!("failed to serialize remote-control server event: {err}");
continue;
}
(payload, queued_server_envelope.write_complete_tx)
};
tokio::select! {
@@ -590,10 +600,13 @@ impl RemoteControlWebsocket {
}
if let ClientEvent::Ack = &client_envelope.event
&& let Some(acked_seq_id) = client_envelope.seq_id
&& let Some(stream_id) = client_envelope.stream_id.as_ref()
{
websocket_state
.outbound_buffer
.ack(&client_envelope.client_id, acked_seq_id);
websocket_state.outbound_buffer.ack(
&client_envelope.client_id,
stream_id,
acked_seq_id,
);
}
drop(websocket_state);
@@ -1220,7 +1233,7 @@ mod tests {
let state = Arc::new(Mutex::new(WebsocketState {
outbound_buffer,
subscribe_cursor: None,
next_seq_id: 0,
next_seq_id_by_stream: HashMap::new(),
}));
let (_server_event_tx, server_event_rx) = mpsc::channel(super::super::CHANNEL_CAPACITY);
let server_event_rx = Arc::new(Mutex::new(server_event_rx));
@@ -1248,6 +1261,83 @@ mod tests {
.expect("writer should stop cleanly");
}
#[tokio::test]
async fn run_server_writer_inner_assigns_contiguous_seq_ids_per_stream() {
let (client_stream, mut server_stream) = connected_websocket_pair().await;
let (websocket_writer, _websocket_reader) = client_stream.split();
let (outbound_buffer, used_rx) = BoundedOutboundBuffer::new();
let state = Arc::new(Mutex::new(WebsocketState {
outbound_buffer,
subscribe_cursor: None,
next_seq_id_by_stream: HashMap::new(),
}));
let (server_event_tx, server_event_rx) = mpsc::channel(super::super::CHANNEL_CAPACITY);
let server_event_rx = Arc::new(Mutex::new(server_event_rx));
let shutdown_token = CancellationToken::new();
let writer_task = tokio::spawn(RemoteControlWebsocket::run_server_writer_inner(
state,
server_event_rx,
used_rx,
websocket_writer,
Duration::from_secs(60),
shutdown_token.clone(),
));
let client_id = ClientId("client-1".to_string());
let first_stream = StreamId("stream-1".to_string());
let second_stream = StreamId("stream-2".to_string());
for stream_id in [&first_stream, &second_stream, &first_stream] {
server_event_tx
.send(super::super::QueuedServerEnvelope {
event: ServerEvent::Pong {
status: crate::transport::remote_control::protocol::PongStatus::Active,
},
client_id: client_id.clone(),
stream_id: stream_id.clone(),
write_complete_tx: None,
})
.await
.expect("server event should queue");
}
assert_eq!(
read_server_text_event(&mut server_stream).await,
serde_json::json!({
"type": "pong",
"client_id": "client-1",
"stream_id": "stream-1",
"seq_id": 1,
"status": "active",
})
);
assert_eq!(
read_server_text_event(&mut server_stream).await,
serde_json::json!({
"type": "pong",
"client_id": "client-1",
"stream_id": "stream-2",
"seq_id": 1,
"status": "active",
})
);
assert_eq!(
read_server_text_event(&mut server_stream).await,
serde_json::json!({
"type": "pong",
"client_id": "client-1",
"stream_id": "stream-1",
"seq_id": 2,
"status": "active",
})
);
shutdown_token.cancel();
writer_task
.await
.expect("writer task should join")
.expect("writer should stop cleanly");
}
#[tokio::test]
async fn run_websocket_reader_inner_times_out_without_pong_frames() {
let (client_stream, _server_stream) = connected_websocket_pair().await;
@@ -1256,7 +1346,7 @@ mod tests {
let state = Arc::new(Mutex::new(WebsocketState {
outbound_buffer,
subscribe_cursor: None,
next_seq_id: 0,
next_seq_id_by_stream: HashMap::new(),
}));
let (server_event_tx, _server_event_rx) = mpsc::channel(super::super::CHANNEL_CAPACITY);
let (transport_event_tx, _transport_event_rx) =
@@ -1287,70 +1377,32 @@ mod tests {
}
#[test]
fn outbound_buffer_acks_by_client_id_across_stream_ids() {
fn outbound_buffer_acks_by_stream_id() {
let (mut outbound_buffer, used_rx) = BoundedOutboundBuffer::new();
let client_1 = ClientId("client-1".to_string());
let client_2 = ClientId("client-2".to_string());
let stream_1 = StreamId("stream-1".to_string());
outbound_buffer.insert(&server_envelope(
&client_1,
"stream-1",
/*seq_id*/ 0,
/*seq_id*/ 1,
"first-client-old-stream",
));
outbound_buffer.insert(&server_envelope(
&client_2,
"stream-1",
/*seq_id*/ 1,
/*seq_id*/ 2,
"second-client",
));
outbound_buffer.insert(&server_envelope(
&client_1,
"stream-2",
/*seq_id*/ 2,
/*seq_id*/ 3,
"first-client-new-stream",
));
outbound_buffer.ack(&client_1, /*acked_seq_id*/ 2);
let mut retained = outbound_buffer
.server_envelopes()
.map(|server_envelope| {
(
server_envelope.client_id.0.as_str(),
server_envelope.stream_id.0.as_str(),
server_envelope.seq_id,
)
})
.collect::<Vec<_>>();
retained.sort_unstable();
assert_eq!(retained, vec![("client-2", "stream-1", 1)]);
assert_eq!(*used_rx.borrow(), 1);
}
#[test]
fn outbound_buffer_retains_unacked_messages_until_ack_advances() {
let (mut outbound_buffer, used_rx) = BoundedOutboundBuffer::new();
let client_1 = ClientId("client-1".to_string());
let client_2 = ClientId("client-2".to_string());
outbound_buffer.insert(&server_envelope(
&client_1,
"stream-1",
/*seq_id*/ 0,
"first-old",
));
outbound_buffer.insert(&server_envelope(
&client_1,
"stream-2",
/*seq_id*/ 1,
"first-new",
));
outbound_buffer.insert(&server_envelope(
&client_2, "stream-1", /*seq_id*/ 2, "second",
));
outbound_buffer.ack(&client_1, /*acked_seq_id*/ 0);
outbound_buffer.ack(&client_1, &stream_1, /*acked_seq_id*/ 3);
let mut retained = outbound_buffer
.server_envelopes()
@@ -1365,7 +1417,50 @@ mod tests {
retained.sort_unstable();
assert_eq!(
retained,
vec![("client-1", "stream-2", 1), ("client-2", "stream-1", 2)]
vec![("client-1", "stream-2", 3), ("client-2", "stream-1", 2)]
);
assert_eq!(*used_rx.borrow(), 2);
}
#[test]
fn outbound_buffer_retains_unacked_messages_until_ack_advances() {
let (mut outbound_buffer, used_rx) = BoundedOutboundBuffer::new();
let client_1 = ClientId("client-1".to_string());
let client_2 = ClientId("client-2".to_string());
let stream_1 = StreamId("stream-1".to_string());
outbound_buffer.insert(&server_envelope(
&client_1,
"stream-1",
/*seq_id*/ 1,
"first-old",
));
outbound_buffer.insert(&server_envelope(
&client_1,
"stream-2",
/*seq_id*/ 2,
"first-new",
));
outbound_buffer.insert(&server_envelope(
&client_2, "stream-1", /*seq_id*/ 3, "second",
));
outbound_buffer.ack(&client_1, &stream_1, /*acked_seq_id*/ 1);
let mut retained = outbound_buffer
.server_envelopes()
.map(|server_envelope| {
(
server_envelope.client_id.0.as_str(),
server_envelope.stream_id.0.as_str(),
server_envelope.seq_id,
)
})
.collect::<Vec<_>>();
retained.sort_unstable();
assert_eq!(
retained,
vec![("client-1", "stream-2", 2), ("client-2", "stream-1", 3)]
);
assert_eq!(*used_rx.borrow(), 2);
}
@@ -1450,6 +1545,20 @@ mod tests {
(client_stream, server_stream)
}
async fn read_server_text_event(
server_stream: &mut WebSocketStream<TcpStream>,
) -> serde_json::Value {
let message = timeout(Duration::from_secs(5), server_stream.next())
.await
.expect("server event should arrive in time")
.expect("server websocket should stay open")
.expect("server event should read");
let tungstenite::Message::Text(text) = message else {
panic!("expected text event, got {message:?}");
};
serde_json::from_str(text.as_ref()).expect("server event should deserialize")
}
async fn respond_with_status_and_headers(
mut stream: TcpStream,
status: &str,