app-server: serialize stateful rpcs by object key

Motivation:
Concurrent initialized request handling lets later RPCs for the same high-level object run before earlier RPCs publish their side effects. A bare per-key mutex would not guarantee arrival order once tasks are spawned, so same-key requests need an explicit FIFO queue.

Summary:
Add Rust-only serialization metadata to client_request_definitions!, expose ClientRequest::serialization_scope(), and implement a private app-server per-key FIFO dispatcher. The server maps logical protocol scopes to runtime queue keys, including connection-scoped command/exec process IDs and fs watch IDs, while leaving unkeyed requests concurrent.

Testing:
- cargo test -p codex-app-server-protocol client_request_serialization_scope
- cargo test -p codex-app-server request_serialization --lib
- cargo test -p codex-app-server-protocol
- cargo test -p codex-app-server
This commit is contained in:
Ruslan Nigmatullin
2026-04-09 23:11:28 -07:00
parent 1385a44433
commit 3281a1cc5b
4 changed files with 610 additions and 1 deletions
+1
View File
@@ -79,6 +79,7 @@ pub mod in_process;
mod message_processor;
mod models;
mod outgoing_message;
mod request_serialization;
mod server_request_error;
mod thread_state;
mod thread_status;
+19 -1
View File
@@ -19,6 +19,9 @@ use crate::outgoing_message::ConnectionId;
use crate::outgoing_message::ConnectionRequestId;
use crate::outgoing_message::OutgoingMessageSender;
use crate::outgoing_message::RequestContext;
use crate::request_serialization::QueuedInitializedRequest;
use crate::request_serialization::RequestSerializationQueueKey;
use crate::request_serialization::RequestSerializationQueues;
use crate::transport::AppServerTransport;
use crate::transport::RemoteControlHandle;
use async_trait::async_trait;
@@ -173,6 +176,7 @@ pub(crate) struct MessageProcessor {
config_warnings: Arc<Vec<ConfigWarningNotification>>,
rpc_transport: AppServerRpcTransport,
remote_control_handle: Option<RemoteControlHandle>,
request_serialization_queues: RequestSerializationQueues,
}
#[derive(Debug, Default)]
@@ -330,6 +334,7 @@ impl MessageProcessor {
config_warnings: Arc::new(config_warnings),
rpc_transport,
remote_control_handle,
request_serialization_queues: RequestSerializationQueues::default(),
}
}
@@ -745,11 +750,13 @@ impl MessageProcessor {
return;
}
let serialization_scope = codex_request.serialization_scope();
let app_server_client_name = session.app_server_client_name().map(str::to_string);
let client_version = session.client_version().map(str::to_string);
let connection_id = connection_request_id.connection_id;
let processor = Arc::clone(self);
let span = request_context.span();
tokio::spawn(
let request = QueuedInitializedRequest::new(
async move {
processor
.handle_initialized_client_request(
@@ -763,6 +770,17 @@ impl MessageProcessor {
}
.instrument(span),
);
if let Some(scope) = serialization_scope {
let key = RequestSerializationQueueKey::from_scope(connection_id, scope);
self.request_serialization_queues
.enqueue(key, request)
.await;
} else {
tokio::spawn(async move {
request.run().await;
});
}
}
async fn handle_initialized_client_request(
@@ -0,0 +1,227 @@
use std::collections::HashMap;
use std::collections::VecDeque;
use std::future::Future;
use std::path::PathBuf;
use std::pin::Pin;
use std::sync::Arc;
use codex_app_server_protocol::ClientRequestSerializationScope;
use codex_utils_absolute_path::AbsolutePathBuf;
use tokio::sync::Mutex;
use tracing::Instrument;
use crate::outgoing_message::ConnectionId;
type BoxFutureUnit = Pin<Box<dyn Future<Output = ()> + Send + 'static>>;
#[derive(Clone, Debug, Eq, Hash, PartialEq)]
pub(crate) enum RequestSerializationQueueKey {
Global(&'static str),
Thread {
thread_id: String,
},
ThreadPath {
path: PathBuf,
},
CommandExecProcess {
connection_id: ConnectionId,
process_id: String,
},
FuzzyFileSearchSession {
session_id: String,
},
FsWatch {
connection_id: ConnectionId,
watch_id: String,
},
Plugin {
marketplace_path: AbsolutePathBuf,
plugin_name: String,
},
PluginId {
plugin_id: String,
},
McpOauth {
server_name: String,
},
}
impl RequestSerializationQueueKey {
pub(crate) fn from_scope(
connection_id: ConnectionId,
scope: ClientRequestSerializationScope,
) -> Self {
match scope {
ClientRequestSerializationScope::Global(name) => Self::Global(name),
ClientRequestSerializationScope::Thread { thread_id } => Self::Thread { thread_id },
ClientRequestSerializationScope::ThreadPath { path } => Self::ThreadPath { path },
ClientRequestSerializationScope::CommandExecProcess { process_id } => {
Self::CommandExecProcess {
connection_id,
process_id,
}
}
ClientRequestSerializationScope::FuzzyFileSearchSession { session_id } => {
Self::FuzzyFileSearchSession { session_id }
}
ClientRequestSerializationScope::FsWatch { watch_id } => Self::FsWatch {
connection_id,
watch_id,
},
ClientRequestSerializationScope::Plugin {
marketplace_path,
plugin_name,
} => Self::Plugin {
marketplace_path,
plugin_name,
},
ClientRequestSerializationScope::PluginId { plugin_id } => Self::PluginId { plugin_id },
ClientRequestSerializationScope::McpOauth { server_name } => {
Self::McpOauth { server_name }
}
}
}
}
pub(crate) struct QueuedInitializedRequest {
future: BoxFutureUnit,
}
impl QueuedInitializedRequest {
pub(crate) fn new(future: impl Future<Output = ()> + Send + 'static) -> Self {
Self {
future: Box::pin(future),
}
}
pub(crate) async fn run(self) {
self.future.await;
}
}
#[derive(Clone, Default)]
pub(crate) struct RequestSerializationQueues {
inner: Arc<Mutex<HashMap<RequestSerializationQueueKey, VecDeque<QueuedInitializedRequest>>>>,
}
impl RequestSerializationQueues {
pub(crate) async fn enqueue(
&self,
key: RequestSerializationQueueKey,
request: QueuedInitializedRequest,
) {
let should_spawn = {
let mut queues = self.inner.lock().await;
match queues.get_mut(&key) {
Some(queue) => {
queue.push_back(request);
false
}
None => {
let mut queue = VecDeque::new();
queue.push_back(request);
queues.insert(key.clone(), queue);
true
}
}
};
if should_spawn {
let queues = self.clone();
let span = tracing::debug_span!("app_server.serialized_request_queue", ?key);
tokio::spawn(async move { queues.drain(key).await }.instrument(span));
}
}
async fn drain(self, key: RequestSerializationQueueKey) {
loop {
let request = {
let mut queues = self.inner.lock().await;
let Some(queue) = queues.get_mut(&key) else {
return;
};
match queue.pop_front() {
Some(request) => request,
None => {
queues.remove(&key);
return;
}
}
};
request.run().await;
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use pretty_assertions::assert_eq;
use tokio::sync::mpsc;
use tokio::sync::oneshot;
use tokio::time::Duration;
use tokio::time::timeout;
#[tokio::test]
async fn same_key_requests_run_fifo() {
let queues = RequestSerializationQueues::default();
let key = RequestSerializationQueueKey::Global("test");
let (tx, mut rx) = mpsc::unbounded_channel();
for value in [1, 2, 3] {
let tx = tx.clone();
queues
.enqueue(
key.clone(),
QueuedInitializedRequest::new(async move {
tx.send(value).expect("receiver should be open");
}),
)
.await;
}
drop(tx);
let mut values = Vec::new();
while let Some(value) = timeout(Duration::from_secs(1), rx.recv())
.await
.expect("timed out waiting for queued request")
{
values.push(value);
}
assert_eq!(values, vec![1, 2, 3]);
}
#[tokio::test]
async fn different_keys_run_concurrently() {
let queues = RequestSerializationQueues::default();
let (blocked_tx, blocked_rx) = oneshot::channel::<()>();
let (ran_tx, ran_rx) = oneshot::channel::<()>();
queues
.enqueue(
RequestSerializationQueueKey::Global("blocked"),
QueuedInitializedRequest::new(async move {
let _ = blocked_rx.await;
}),
)
.await;
queues
.enqueue(
RequestSerializationQueueKey::Global("other"),
QueuedInitializedRequest::new(async move {
ran_tx.send(()).expect("receiver should be open");
}),
)
.await;
timeout(Duration::from_secs(1), ran_rx)
.await
.expect("other key should not be blocked")
.expect("sender should be open");
blocked_tx
.send(())
.expect("blocked request should be waiting");
}
}