mirror of
https://github.com/openai/codex.git
synced 2026-11-10 21:00:47 +00:00
Fragment large Responses WebSocket messages into bounded frames (#60522)
## What changed Send Responses text messages larger than 1 MiB as a single WebSocket message split into frames of at most 1 MiB, awaiting each frame send to apply backpressure. Keep smaller text messages and non-text messages on the existing send path. Update the `tokio-tungstenite` and `tungstenite` fork revisions and associated Cargo and Bazel lockfiles. ## Testing Add a unit test covering the 1 MiB boundary, continuation opcodes, the final-frame flag, and preservation of payload bytes when a multibyte UTF-8 character crosses a frame boundary. GitOrigin-RevId: 372ee31b77e185b963d3bdca11847e50e9018421
This commit is contained in:
Generated
+2
-2
@@ -16043,7 +16043,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "tokio-tungstenite"
|
||||
version = "0.28.0"
|
||||
source = "git+https://github.com/openai-oss-forks/tokio-tungstenite?rev=02cb49fa827411e9279e3688d19ab9e82661b0e8#02cb49fa827411e9279e3688d19ab9e82661b0e8"
|
||||
source = "git+https://github.com/openai-oss-forks/tokio-tungstenite?rev=7447c427513cfe685f236e1ac0513c4ea3dd9c0b#7447c427513cfe685f236e1ac0513c4ea3dd9c0b"
|
||||
dependencies = [
|
||||
"futures-util",
|
||||
"log",
|
||||
@@ -16531,7 +16531,7 @@ dependencies = [
|
||||
[[package]]
|
||||
name = "tungstenite"
|
||||
version = "0.27.0"
|
||||
source = "git+https://github.com/openai-oss-forks/tungstenite-rs?rev=db4be399bc1b8b5de2edb46c7187834edabf7e30#db4be399bc1b8b5de2edb46c7187834edabf7e30"
|
||||
source = "git+https://github.com/openai-oss-forks/tungstenite-rs?rev=bc589d684755a8c02da0fc7e58d0861cbfc306f1#bc589d684755a8c02da0fc7e58d0861cbfc306f1"
|
||||
dependencies = [
|
||||
"bytes",
|
||||
"data-encoding",
|
||||
|
||||
+3
-3
@@ -685,12 +685,12 @@ crossterm = { git = "https://github.com/openai-oss-forks/crossterm", rev = "ed1c
|
||||
# the same reviewed source until a release includes transact-rs/sqlx#4502.
|
||||
sqlx-core = { git = "https://github.com/transact-rs/sqlx", rev = "10462fa4414c6c7f1ed1d39bb31678e3b726878c" }
|
||||
sqlx-sqlite = { git = "https://github.com/transact-rs/sqlx", rev = "10462fa4414c6c7f1ed1d39bb31678e3b726878c" }
|
||||
tokio-tungstenite = { git = "https://github.com/openai-oss-forks/tokio-tungstenite", rev = "02cb49fa827411e9279e3688d19ab9e82661b0e8" }
|
||||
tungstenite = { git = "https://github.com/openai-oss-forks/tungstenite-rs", rev = "db4be399bc1b8b5de2edb46c7187834edabf7e30" }
|
||||
tokio-tungstenite = { git = "https://github.com/openai-oss-forks/tokio-tungstenite", rev = "7447c427513cfe685f236e1ac0513c4ea3dd9c0b" }
|
||||
tungstenite = { git = "https://github.com/openai-oss-forks/tungstenite-rs", rev = "bc589d684755a8c02da0fc7e58d0861cbfc306f1" }
|
||||
v8 = { git = "https://github.com/openai-oss-forks/rusty_v8.git", rev = "90a65f887a0c0f048f3784afbcb320b373c8a2e0" }
|
||||
|
||||
# Uncomment to debug local changes.
|
||||
# rmcp = { path = "../../rust-sdk/crates/rmcp" }
|
||||
|
||||
[patch."ssh://git@github.com/openai-oss-forks/tungstenite-rs.git"]
|
||||
tungstenite = { git = "https://github.com/openai-oss-forks/tungstenite-rs", rev = "db4be399bc1b8b5de2edb46c7187834edabf7e30" }
|
||||
tungstenite = { git = "https://github.com/openai-oss-forks/tungstenite-rs", rev = "bc589d684755a8c02da0fc7e58d0861cbfc306f1" }
|
||||
|
||||
@@ -18,6 +18,7 @@ use crate::safety_buffering::treatment_from_headers;
|
||||
use crate::sse::ResponsesStreamEvent;
|
||||
use crate::sse::process_responses_event;
|
||||
use crate::telemetry::WebsocketTelemetry;
|
||||
use bytes::Bytes;
|
||||
use codex_client::TransportError;
|
||||
use codex_http_client::HttpClientFactory;
|
||||
use codex_http_client::RetryAfter;
|
||||
@@ -44,6 +45,9 @@ use tokio_tungstenite::tungstenite::Message;
|
||||
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
|
||||
use tokio_tungstenite::tungstenite::handshake::client::Response;
|
||||
use tokio_tungstenite::tungstenite::protocol::CloseFrame;
|
||||
use tokio_tungstenite::tungstenite::protocol::frame::Fragment;
|
||||
use tokio_tungstenite::tungstenite::protocol::frame::Frame;
|
||||
use tokio_tungstenite::tungstenite::protocol::frame::coding::Data;
|
||||
use tracing::Instrument;
|
||||
use tracing::Span;
|
||||
use tracing::debug;
|
||||
@@ -69,6 +73,36 @@ enum WsCommand {
|
||||
},
|
||||
}
|
||||
|
||||
/// Send large Responses turns as one text message with backpressure between frames.
|
||||
async fn send_message(
|
||||
socket: &mut (impl futures::Sink<Message, Error = WsError> + Unpin),
|
||||
message: Message,
|
||||
) -> Result<(), WsError> {
|
||||
const FRAGMENT_SIZE: usize = 1024 * 1024;
|
||||
|
||||
let Message::Text(text) = message else {
|
||||
return socket.send(message).await;
|
||||
};
|
||||
if text.len() <= FRAGMENT_SIZE {
|
||||
return socket.send(Message::Text(text)).await;
|
||||
}
|
||||
|
||||
let mut data = Bytes::from(text);
|
||||
let mut opcode = Data::Text;
|
||||
while !data.is_empty() {
|
||||
let payload = data.split_to(data.len().min(FRAGMENT_SIZE));
|
||||
let fragment = if data.is_empty() {
|
||||
Fragment::Final
|
||||
} else {
|
||||
Fragment::More
|
||||
};
|
||||
let frame = Frame::message_fragment(payload, opcode, fragment);
|
||||
socket.send(Message::Frame(frame)).await?;
|
||||
opcode = Data::Continue;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
impl WsStream {
|
||||
fn new(inner: WebSocketConnection) -> Self {
|
||||
let routing_affinity = inner.routing_affinity().cloned();
|
||||
@@ -85,7 +119,7 @@ impl WsStream {
|
||||
};
|
||||
match command {
|
||||
WsCommand::Send { message, tx_result } => {
|
||||
let result = inner.send(message).await;
|
||||
let result = send_message(&mut inner, message).await;
|
||||
let should_break = result.is_err();
|
||||
let _ = tx_result.send(result);
|
||||
if should_break {
|
||||
@@ -1173,6 +1207,38 @@ mod tests {
|
||||
assert_eq!(wire_payload, expected_payload);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn large_turn_is_one_text_message_in_bounded_frames() {
|
||||
const MIB: usize = 1024 * 1024;
|
||||
let (tx, mut frames) = futures::channel::mpsc::unbounded();
|
||||
let mut socket = tx.sink_map_err(|_| WsError::ConnectionClosed);
|
||||
|
||||
let at_limit = Message::text("a".repeat(MIB));
|
||||
send_message(&mut socket, at_limit.clone()).await.unwrap();
|
||||
assert_eq!(frames.next().await, Some(at_limit));
|
||||
|
||||
let mut text = "a".repeat(MIB - 1);
|
||||
text.push('🌱');
|
||||
text.push_str(&"b".repeat(MIB));
|
||||
send_message(&mut socket, Message::text(text.clone()))
|
||||
.await
|
||||
.unwrap();
|
||||
drop(socket);
|
||||
for (i, expected) in text.as_bytes().chunks(MIB).enumerate() {
|
||||
let Some(Message::Frame(frame)) = frames.next().await else {
|
||||
panic!("expected a data frame");
|
||||
};
|
||||
let opcode = if i == 0 { Data::Text } else { Data::Continue };
|
||||
assert_eq!(
|
||||
frame.header().opcode,
|
||||
tokio_tungstenite::tungstenite::protocol::frame::coding::OpCode::Data(opcode)
|
||||
);
|
||||
assert_eq!(frame.header().is_final, (i + 1) * MIB >= text.len());
|
||||
assert_eq!(frame.payload(), expected);
|
||||
}
|
||||
assert!(frames.next().await.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn websocket_config_enables_permessage_deflate() {
|
||||
let config = websocket_config();
|
||||
|
||||
Reference in New Issue
Block a user