openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-09-03 14:59:03 +00:00
Code Issues Packages Projects Releases Wiki Activity

Bound remote HTTP response buffering

Browse Source
This commit is contained in:
jif-oai
2026-07-09 13:03:45 +01:00
parent 13ba8058f2
commit 939bff96b6
5 changed files with 258 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
codex-rs/exec-server-protocol/src/protocol.rs
View File

@@ -40,6 +40,8 @@ pub const FS_COPY_METHOD: &str = "fs/copy";
pub const HTTP_REQUEST_METHOD: &str = "http/request";
/// JSON-RPC notification method for streamed executor HTTP response bodies.
pub const HTTP_REQUEST_BODY_DELTA_METHOD: &str = "http/request/bodyDelta";
/// Maximum decoded response-body bytes carried by one streamed HTTP notification.
pub const MAX_HTTP_BODY_DELTA_BYTES: usize = 1024 * 1024;
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(transparent)]

8
codex-rs/exec-server/src/client.rs
View File

@@ -15,6 +15,7 @@ use futures::future::BoxFuture;
use serde_json::Value;
use tokio::sync::Mutex;
use tokio::sync::OnceCell;
use tokio::sync::Semaphore;
use tokio::sync::mpsc;
use tokio::sync::watch;
use tokio_util::task::AbortOnDropHandle;
@@ -24,6 +25,8 @@ use tracing::Instrument;
use tracing::debug;
use crate::ProcessId;
use crate::client::http_client::response_body_stream::MAX_QUEUED_HTTP_BODY_BYTES;
use crate::client::http_client::response_body_stream::QueuedHttpBodyDelta;
use crate::client_api::ExecServerClientConnectOptions;
use crate::client_api::ExecServerTransportParams;
use crate::client_api::HttpClient;
@@ -86,7 +89,6 @@ use crate::protocol::FsWalkResponse;
use crate::protocol::FsWriteFileParams;
use crate::protocol::FsWriteFileResponse;
use crate::protocol::HTTP_REQUEST_BODY_DELTA_METHOD;
use crate::protocol::HttpRequestBodyDeltaNotification;
use crate::protocol::INITIALIZE_METHOD;
use crate::protocol::INITIALIZED_METHOD;
use crate::protocol::InitializeParams;
@@ -202,9 +204,10 @@ struct Inner {
// because they share the same connection-global notification channel as
// process output. Keep the routing table local to the client so higher
// layers can consume body chunks like a normal byte stream.
http_body_streams: ArcSwap<HashMap<String, mpsc::Sender<HttpRequestBodyDeltaNotification>>>,
http_body_streams: ArcSwap<HashMap<String, mpsc::Sender<QueuedHttpBodyDelta>>>,
http_body_stream_failures: ArcSwap<HashMap<String, String>>,
http_body_streams_write_lock: Mutex<()>,
http_body_stream_byte_budget: Arc<Semaphore>,
http_body_stream_next_id: AtomicU64,
session_id: OnceLock<String>,
reconnect_strategy: Option<ExecServerReconnectStrategy>,
@@ -851,6 +854,7 @@ impl ExecServerClient {
http_body_streams: ArcSwap::from_pointee(HashMap::new()),
http_body_stream_failures: ArcSwap::from_pointee(HashMap::new()),
http_body_streams_write_lock: Mutex::new(()),
http_body_stream_byte_budget: Arc::new(Semaphore::new(MAX_QUEUED_HTTP_BODY_BYTES)),
http_body_stream_next_id: AtomicU64::new(1),
session_id,
reconnect_strategy,

98
codex-rs/exec-server/src/client/http_response_body_stream.rs
View File

@@ -15,6 +15,7 @@ use reqwest::Response;
use serde_json::Value;
use serde_json::from_value;
use tokio::runtime::Handle;
use tokio::sync::OwnedSemaphorePermit;
use tokio::sync::mpsc;
use tokio::sync::mpsc::error::TrySendError;
use tracing::debug;
@@ -23,8 +24,29 @@ use crate::client::ExecServerError;
use crate::client::Inner;
use crate::protocol::HTTP_REQUEST_BODY_DELTA_METHOD;
use crate::protocol::HttpRequestBodyDeltaNotification;
use crate::protocol::MAX_HTTP_BODY_DELTA_BYTES;
use crate::rpc::RpcNotificationSender;
pub(crate) const MAX_QUEUED_HTTP_BODY_BYTES: usize = 16 * 1024 * 1024;
const MAX_ENCODED_HTTP_BODY_DELTA_BYTES: usize = MAX_HTTP_BODY_DELTA_BYTES.div_ceil(3) * 4;
pub(crate) struct QueuedHttpBodyDelta {
notification: HttpRequestBodyDeltaNotification,
_byte_permit: Option<OwnedSemaphorePermit>,
}
impl QueuedHttpBodyDelta {
pub(crate) fn new(
notification: HttpRequestBodyDeltaNotification,
byte_permit: Option<OwnedSemaphorePermit>,
) -> Self {
Self {
notification,
_byte_permit: byte_permit,
}
}
}
pub(super) struct HttpBodyStreamRegistration {
inner: Arc<Inner>,
request_id: String,
@@ -39,7 +61,7 @@ enum HttpResponseBodyStreamInner {
inner: Arc<Inner>,
request_id: String,
next_seq: u64,
rx: mpsc::Receiver<HttpRequestBodyDeltaNotification>,
rx: mpsc::Receiver<QueuedHttpBodyDelta>,
pending_eof: bool,
closed: bool,
},
@@ -66,7 +88,7 @@ impl HttpResponseBodyStream {
pub(super) fn remote(
inner: Arc<Inner>,
request_id: String,
rx: mpsc::Receiver<HttpRequestBodyDeltaNotification>,
rx: mpsc::Receiver<QueuedHttpBodyDelta>,
) -> Self {
Self {
inner: HttpResponseBodyStreamInner::Remote {
@@ -107,7 +129,11 @@ impl HttpResponseBodyStream {
return Ok(None);
}
let Some(delta) = rx.recv().await else {
let Some(QueuedHttpBodyDelta {
notification: delta,
..
}) = rx.recv().await
else {
finish_remote_stream(inner, request_id, closed).await;
if let Some(error) = inner.take_http_body_stream_failure(request_id).await {
return Err(ExecServerError::Protocol(format!(
@@ -220,7 +246,22 @@ impl Inner {
&self,
params: Option<Value>,
) -> Result<(), ExecServerError> {
let params: HttpRequestBodyDeltaNotification = from_value(params.unwrap_or(Value::Null))?;
let params = params.unwrap_or(Value::Null);
if params
.get("deltaBase64")
.and_then(Value::as_str)
.is_some_and(|delta| delta.len() > MAX_ENCODED_HTTP_BODY_DELTA_BYTES)
{
return Err(ExecServerError::Protocol(format!(
"http response body delta exceeds {MAX_HTTP_BODY_DELTA_BYTES} bytes"
)));
}
let params: HttpRequestBodyDeltaNotification = from_value(params)?;
if params.delta.0.len() > MAX_HTTP_BODY_DELTA_BYTES {
return Err(ExecServerError::Protocol(format!(
"http response body delta exceeds {MAX_HTTP_BODY_DELTA_BYTES} bytes"
)));
}
// Unknown request ids are ignored intentionally: a stream may have already
// reached EOF and released its route.
if let Some(tx) = self
@@ -231,7 +272,33 @@ impl Inner {
{
let request_id = params.request_id.clone();
let terminal_delta = params.done || params.error.is_some();
match tx.try_send(params) {
let queued_bytes = params
.delta
.0
.len()
.saturating_add(params.error.as_deref().map_or(0, str::len));
let byte_permit = if queued_bytes == 0 {
None
} else {
u32::try_from(queued_bytes).ok().and_then(|queued_bytes| {
Arc::clone(&self.http_body_stream_byte_budget)
.try_acquire_many_owned(queued_bytes)
.ok()
})
};
if queued_bytes > 0 && byte_permit.is_none() {
self.record_http_body_stream_failure(
&request_id,
format!("queued body deltas exceed {MAX_QUEUED_HTTP_BODY_BYTES} bytes"),
)
.await;
self.remove_http_body_stream(&request_id).await;
debug!(
"closing http response stream `{request_id}` after exhausting the queued byte budget"
);
return Ok(());
}
match tx.try_send(QueuedHttpBodyDelta::new(params, byte_permit)) {
Ok(()) => {
if terminal_delta {
self.remove_http_body_stream(&request_id).await;
@@ -266,13 +333,16 @@ impl Inner {
self.http_body_streams.store(Arc::new(HashMap::new()));
for (request_id, tx) in streams {
if tx
.try_send(HttpRequestBodyDeltaNotification {
request_id: request_id.clone(),
seq: 1,
delta: Vec::new().into(),
done: true,
error: Some(message.clone()),
})
.try_send(QueuedHttpBodyDelta::new(
HttpRequestBodyDeltaNotification {
request_id: request_id.clone(),
seq: 1,
delta: Vec::new().into(),
done: true,
error: Some(message.clone()),
},
/*byte_permit*/ None,
))
.is_err()
{
let mut next_failures = self.http_body_stream_failures.load().as_ref().clone();
@@ -295,7 +365,7 @@ impl Inner {
pub(super) async fn insert_http_body_stream(
&self,
request_id: String,
tx: mpsc::Sender<HttpRequestBodyDeltaNotification>,
tx: mpsc::Sender<QueuedHttpBodyDelta>,
) -> Result<(), ExecServerError> {
let _streams_write_guard = self.http_body_streams_write_lock.lock().await;
let streams = self.http_body_streams.load();
@@ -321,7 +391,7 @@ impl Inner {
pub(super) async fn remove_http_body_stream(
&self,
request_id: &str,
) -> Option<mpsc::Sender<HttpRequestBodyDeltaNotification>> {
) -> Option<mpsc::Sender<QueuedHttpBodyDelta>> {
let _streams_write_guard = self.http_body_streams_write_lock.lock().await;
let streams = self.http_body_streams.load();
let stream = streams.get(request_id).cloned();

31
codex-rs/exec-server/src/client/reqwest_http_client.rs
View File

@@ -30,6 +30,7 @@ use crate::protocol::HttpRedirectPolicy;
use crate::protocol::HttpRequestBodyDeltaNotification;
use crate::protocol::HttpRequestParams;
use crate::protocol::HttpRequestResponse;
use crate::protocol::MAX_HTTP_BODY_DELTA_BYTES;
use crate::rpc::RpcNotificationSender;
use crate::rpc::internal_error;
use crate::rpc::invalid_params;
@@ -222,21 +223,23 @@ impl ReqwestHttpRequestRunner {
while let Some(chunk) = body.next().await {
match chunk {
Ok(bytes) => {
if !send_body_delta(
&notifications,
HttpRequestBodyDeltaNotification {
request_id: request_id.clone(),
seq,
delta: bytes.to_vec().into(),
done: false,
error: None,
},
)
.await
{
return;
for chunk in bytes.chunks(MAX_HTTP_BODY_DELTA_BYTES) {
if !send_body_delta(
&notifications,
HttpRequestBodyDeltaNotification {
request_id: request_id.clone(),
seq,
delta: chunk.to_vec().into(),
done: false,
error: None,
},
)
.await
{
return;
}
seq += 1;
}
seq += 1;
}
Err(error) => {
let _ = send_body_delta(

149
codex-rs/exec-server/tests/http_client.rs
View File

@@ -17,6 +17,7 @@ use codex_exec_server_protocol::JSONRPCMessage;
use codex_exec_server_protocol::JSONRPCNotification;
use codex_exec_server_protocol::JSONRPCRequest;
use codex_exec_server_protocol::JSONRPCResponse;
use codex_exec_server_protocol::MAX_HTTP_BODY_DELTA_BYTES;
use codex_exec_server_protocol::RequestId;
use futures::SinkExt;
use futures::StreamExt;
@@ -44,6 +45,7 @@ const INITIALIZE_METHOD: &str = "initialize";
const INITIALIZED_METHOD: &str = "initialized";
const TEST_TIMEOUT: Duration = Duration::from_secs(5);
const HTTP_BODY_DELTA_CHANNEL_CAPACITY: u64 = 256;
const HTTP_BODY_DELTA_BYTE_BUDGET: usize = 16 * 1024 * 1024;
const OVERFLOWING_BODY_DELTA_FRAMES: u64 = 1_024;
/// What this tests: the buffered HTTP helper always sends a buffered
@@ -764,6 +766,153 @@ async fn http_response_body_stream_fails_when_transport_disconnects() -> Result<
Ok(())
}
/// What this tests: an executor cannot make the orchestrator decode and retain
/// a body frame larger than the response-stream wire contract allows.
#[tokio::test]
async fn http_response_body_stream_rejects_oversized_delta() -> Result<()> {
let (finish_tx, finish_rx) = oneshot::channel();
let server = spawn_scripted_exec_server(|mut peer| async move {
let (_request_id, params) = peer.read_http_request().await?;
assert_eq!(
params,
HttpRequestParams {
method: "GET".to_string(),
url: "https://example.test/mcp/oversized-delta".to_string(),
headers: Vec::new(),
body: None,
timeout_ms: None,
redirect_policy: HttpRedirectPolicy::Follow,
request_id: "http-1".to_string(),
stream_response: true,
}
);
peer.write_body_delta(HttpRequestBodyDeltaNotification {
request_id: "http-1".to_string(),
seq: 1,
delta: vec![0; MAX_HTTP_BODY_DELTA_BYTES + 1].into(),
done: false,
error: None,
})
.await?;
finish_rx.await.expect("test should finish server task");
Ok(())
})
.await?;
let client = server.connect_client().await?;
let request = HttpRequestParams {
method: "GET".to_string(),
url: "https://example.test/mcp/oversized-delta".to_string(),
headers: Vec::new(),
body: None,
timeout_ms: None,
redirect_policy: HttpRedirectPolicy::Follow,
request_id: "caller-stream-id".to_string(),
stream_response: false,
};
let result = timeout(TEST_TIMEOUT, client.http_request_stream(request))
.await
.context("oversized body delta should close the executor transport")?;
let error = match result {
Ok(_) => bail!("oversized body delta should fail the request"),
Err(error) => error,
};
let error = error.to_string();
assert_eq!(error, "exec-server transport disconnected");
finish_tx.send(()).expect("server task should stay active");
drop(client);
server.finish().await?;
Ok(())
}
/// What this tests: frame-count backpressure cannot hide an unbounded amount
/// of executor-controlled body bytes across the orchestrator's stream queues.
#[tokio::test]
async fn http_response_body_stream_enforces_queued_byte_budget() -> Result<()> {
let (finish_tx, finish_rx) = oneshot::channel();
let server = spawn_scripted_exec_server(|mut peer| async move {
let (request_id, params) = peer.read_http_request().await?;
assert_eq!(
params,
HttpRequestParams {
method: "GET".to_string(),
url: "https://example.test/mcp/byte-budget".to_string(),
headers: Vec::new(),
body: None,
timeout_ms: None,
redirect_policy: HttpRedirectPolicy::Follow,
request_id: "http-1".to_string(),
stream_response: true,
}
);
let frame_count = HTTP_BODY_DELTA_BYTE_BUDGET / MAX_HTTP_BODY_DELTA_BYTES + 1;
for seq in 1..=frame_count as u64 {
peer.write_body_delta(HttpRequestBodyDeltaNotification {
request_id: "http-1".to_string(),
seq,
delta: vec![0; MAX_HTTP_BODY_DELTA_BYTES].into(),
done: false,
error: None,
})
.await?;
}
tokio::time::sleep(Duration::from_millis(100)).await;
peer.write_response(
request_id,
HttpRequestResponse {
status: 200,
headers: Vec::new(),
body: Vec::new().into(),
},
)
.await?;
finish_rx.await.expect("test should finish server task");
Ok(())
})
.await?;
let client = server.connect_client().await?;
let (_response, mut body_stream) = timeout(
TEST_TIMEOUT,
client.http_request_stream(HttpRequestParams {
method: "GET".to_string(),
url: "https://example.test/mcp/byte-budget".to_string(),
headers: Vec::new(),
body: None,
timeout_ms: None,
redirect_policy: HttpRedirectPolicy::Follow,
request_id: "caller-stream-id".to_string(),
stream_response: false,
}),
)
.await
.context("streamed http/request should return headers")??;
let mut delivered_bytes = 0;
let error = loop {
match timeout(TEST_TIMEOUT, body_stream.recv())
.await
.context("queued body stream should finish")?
{
Ok(Some(chunk)) => delivered_bytes += chunk.len(),
Ok(None) => bail!("byte-budget exhaustion should not look like clean EOF"),
Err(error) => break error,
}
};
assert_eq!(delivered_bytes, HTTP_BODY_DELTA_BYTE_BUDGET);
assert!(
error
.to_string()
.contains("queued body deltas exceed 16777216 bytes")
);
finish_tx.send(()).expect("server task should stay active");
drop(client);
server.finish().await?;
Ok(())
}
/// What this tests: transport disconnect still records a terminal stream
/// failure even when the client-side body-delta queue is already full.
#[tokio::test]