mirror of
https://github.com/openai/codex.git
synced 2026-11-07 20:30:33 +00:00
Support caller-owned sideband channels for realtime attachments (#58953)
## What changed Add `RealtimeSidebandTransport` and `Op::RealtimeConversationAttachStream` so callers can attach Core to an existing realtime call through bounded JSON channels without opening or reconnecting a socket. Read events independently of outbound backpressure, enforce a 1 MiB message limit, and drain accepted input on detach or channel closure before flushing the transcript. Pending writes are never replayed. ## Testing Add transport and Core integration tests covering delegation without a callback socket, stalled output, attachment replacement and detach, channel closure, oversized messages, and preservation of queued transcripts in history. GitOrigin-RevId: 06b7d1dfe2388054438d72d925e510a959dfe106
This commit is contained in:
Jiahao
@@ -63,11 +63,13 @@ use tungstenite::protocol::frame::coding::CloseCode;
|
||||
use url::Url;
|
||||
|
||||
const REALTIME_WIRE_LOG_TARGET: &str = "codex_api::realtime_websocket::wire";
|
||||
mod sideband;
|
||||
const OPENAI_REALTIME_API_BASE_URL: &str = "https://api.openai.com/v1";
|
||||
|
||||
struct WsStream {
|
||||
tx_command: mpsc::Sender<WsCommand>,
|
||||
pump_task: tokio::task::JoinHandle<()>,
|
||||
transport_task: tokio::task::JoinHandle<()>,
|
||||
sideband_input_stop: Option<tokio::sync::watch::Sender<bool>>,
|
||||
terminal_policy_error: tokio::sync::watch::Receiver<Option<NetworkPolicyDenied>>,
|
||||
}
|
||||
|
||||
@@ -90,7 +92,7 @@ impl WsStream {
|
||||
let (tx_message, rx_message) = async_channel::unbounded::<Result<Message, WsError>>();
|
||||
|
||||
let (terminal_policy_tx, terminal_policy_error) = tokio::sync::watch::channel(None);
|
||||
let pump_task = tokio::spawn(async move {
|
||||
let transport_task = tokio::spawn(async move {
|
||||
let mut inner = inner;
|
||||
let report_write_policy_error = |error: &WsError| {
|
||||
if let Some(denied) = codex_websocket_client::network_policy_denial(error) {
|
||||
@@ -191,7 +193,8 @@ impl WsStream {
|
||||
(
|
||||
Self {
|
||||
tx_command,
|
||||
pump_task,
|
||||
transport_task,
|
||||
sideband_input_stop: None,
|
||||
terminal_policy_error,
|
||||
},
|
||||
rx_message,
|
||||
@@ -238,7 +241,7 @@ impl WsStream {
|
||||
|
||||
impl Drop for WsStream {
|
||||
fn drop(&mut self) {
|
||||
self.pump_task.abort();
|
||||
self.transport_task.abort();
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,158 @@
|
||||
//! Adapts caller-owned, bounded JSON channels to the existing realtime parser and writer.
|
||||
//!
|
||||
//! ```text
|
||||
//! Caller -- incoming JSON --> read loop -- Message::Text --> Core event parser
|
||||
//! Caller <-- outgoing JSON -- write loop <-- WsCommand ----- Core writer
|
||||
//! ```
|
||||
//!
|
||||
//! Each channel item carries one JSON message of up to 1 MiB. The loops run
|
||||
//! independently so Core can still receive events while a command waits for
|
||||
//! space in the caller's queue. Shutdown seals incoming messages and drains accepted
|
||||
//! input before closing both directions, including a frame awaiting parser capacity.
|
||||
|
||||
use super::*;
|
||||
use codex_protocol::protocol::RealtimeSidebandTransport;
|
||||
|
||||
const MAX_MESSAGE_BYTES: usize = 1024 * 1024;
|
||||
|
||||
impl RealtimeWebsocketConnection {
|
||||
/// Uses an already authenticated, caller-owned stream without opening or reconnecting a socket.
|
||||
pub fn from_sideband(
|
||||
transport: RealtimeSidebandTransport,
|
||||
event_parser: RealtimeEventParser,
|
||||
) -> Self {
|
||||
let RealtimeSidebandTransport {
|
||||
mut incoming,
|
||||
outgoing,
|
||||
} = transport;
|
||||
|
||||
let (tx_command, mut rx_command) = mpsc::channel::<WsCommand>(32);
|
||||
let (tx_message, rx_message) = async_channel::bounded(32);
|
||||
let (_, terminal_policy_error) = tokio::sync::watch::channel(None);
|
||||
|
||||
let (input_stop, mut input_stopped) = tokio::sync::watch::channel(false);
|
||||
let transport_input_stop = input_stop.clone();
|
||||
let transport_task = tokio::spawn(async move {
|
||||
let read = async {
|
||||
let mut is_draining = false;
|
||||
loop {
|
||||
let text = if is_draining {
|
||||
// A reserved but unsent permit must not hold shutdown open.
|
||||
// try_recv() will drain messages already queued and exit when emptied.
|
||||
incoming.try_recv().ok()
|
||||
} else {
|
||||
tokio::select! {
|
||||
biased;
|
||||
// Core requested shutdown, or the transport's write loop ended.
|
||||
_ = input_stopped.changed() => {
|
||||
incoming.close();
|
||||
is_draining = true;
|
||||
continue;
|
||||
}
|
||||
// The caller dropped its receiver for Core commands.
|
||||
// Detect this even when the writer is idle, then drain accepted input.
|
||||
_ = outgoing.closed() => {
|
||||
incoming.close();
|
||||
is_draining = true;
|
||||
continue;
|
||||
}
|
||||
text = incoming.recv() => text,
|
||||
}
|
||||
};
|
||||
let Some(text) = text else {
|
||||
break;
|
||||
};
|
||||
if text.len() > MAX_MESSAGE_BYTES {
|
||||
break;
|
||||
}
|
||||
if tx_message
|
||||
.send(Ok(Message::Text(text.into())))
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
break;
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
let write = async {
|
||||
while let Some(command) = rx_command.recv().await {
|
||||
// These commands are still called WsCommand for legacy reasons.
|
||||
match command {
|
||||
WsCommand::Send {
|
||||
message: Message::Text(text),
|
||||
tx_result,
|
||||
} => {
|
||||
let result = if text.len() > MAX_MESSAGE_BYTES {
|
||||
Err(WsError::Io(std::io::Error::other(
|
||||
"sideband message too large",
|
||||
)))
|
||||
} else {
|
||||
outgoing
|
||||
.send(text.to_string())
|
||||
.await
|
||||
.map_err(|_| WsError::ConnectionClosed)
|
||||
};
|
||||
let failed = result.is_err();
|
||||
// Tell Core's writer whether the caller's channel accepted the command.
|
||||
let _ = tx_result.send(result);
|
||||
if failed {
|
||||
break;
|
||||
}
|
||||
}
|
||||
WsCommand::Close {
|
||||
frame: _,
|
||||
tx_result,
|
||||
} => {
|
||||
let _ = tx_result.send(Ok(()));
|
||||
break;
|
||||
}
|
||||
WsCommand::Send { tx_result, .. } => {
|
||||
let _ = tx_result.send(Err(WsError::Io(std::io::Error::other(
|
||||
"sideband requires JSON text",
|
||||
))));
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
tokio::pin!(read);
|
||||
tokio::select! {
|
||||
_ = &mut read => {},
|
||||
_ = write => {
|
||||
transport_input_stop.send_replace(true);
|
||||
// Preserve the read future: it may already own a frame that
|
||||
// has not yet fit into the parser queue.
|
||||
read.await;
|
||||
},
|
||||
}
|
||||
});
|
||||
|
||||
Self::new(
|
||||
WsStream {
|
||||
tx_command,
|
||||
transport_task,
|
||||
sideband_input_stop: Some(input_stop),
|
||||
terminal_policy_error,
|
||||
},
|
||||
rx_message,
|
||||
event_parser,
|
||||
RealtimeTranscriptState::default(),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
impl RealtimeWebsocketWriter {
|
||||
/// Seal caller-owned input so its accepted events can drain without waiting for producers.
|
||||
/// Socket-backed connections are unaffected.
|
||||
pub fn close_sideband_input(&self) {
|
||||
if let Some(input_stop) = &self.stream.sideband_input_stop {
|
||||
input_stop.send_replace(true);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "sideband_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,210 @@
|
||||
#![cfg(test)]
|
||||
|
||||
use super::*;
|
||||
use pretty_assertions::assert_eq;
|
||||
use std::time::Duration;
|
||||
use tokio::time::timeout;
|
||||
|
||||
fn connection() -> (
|
||||
RealtimeWebsocketConnection,
|
||||
mpsc::Sender<String>,
|
||||
mpsc::Receiver<String>,
|
||||
) {
|
||||
let (events, incoming) = mpsc::channel(1);
|
||||
let (outgoing, commands) = mpsc::channel(1);
|
||||
(
|
||||
RealtimeWebsocketConnection::from_sideband(
|
||||
RealtimeSidebandTransport { incoming, outgoing },
|
||||
RealtimeEventParser::FramelessBidi,
|
||||
),
|
||||
events,
|
||||
commands,
|
||||
)
|
||||
}
|
||||
|
||||
/// A blocked command receiver must not prevent receiving a realtime event on the other half.
|
||||
#[tokio::test]
|
||||
async fn stalled_output_does_not_block_input() {
|
||||
let (connection, events, commands) = connection();
|
||||
let writer = connection.writer();
|
||||
writer.send_payload("first".into()).await.unwrap();
|
||||
let blocked = tokio::spawn(async move { writer.send_payload("second".into()).await });
|
||||
events
|
||||
.send(r#"{"type":"session.started","session":{"id":"call"}}"#.into())
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(
|
||||
timeout(Duration::from_secs(1), connection.next_event())
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap()
|
||||
.is_some()
|
||||
);
|
||||
assert!(!blocked.is_finished());
|
||||
drop(commands);
|
||||
assert!(
|
||||
timeout(Duration::from_secs(1), blocked)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap()
|
||||
.is_err()
|
||||
);
|
||||
assert!(
|
||||
timeout(Duration::from_secs(1), connection.next_event())
|
||||
.await
|
||||
.unwrap()
|
||||
.is_err()
|
||||
);
|
||||
}
|
||||
|
||||
/// Ending the input stream closes the output too and never attempts reconnection.
|
||||
#[tokio::test]
|
||||
async fn input_half_close_ends_attachment() {
|
||||
let (connection, events, mut commands) = connection();
|
||||
drop(events);
|
||||
assert!(
|
||||
timeout(Duration::from_secs(1), connection.next_event())
|
||||
.await
|
||||
.unwrap()
|
||||
.is_err()
|
||||
);
|
||||
assert_eq!(
|
||||
timeout(Duration::from_secs(1), commands.recv())
|
||||
.await
|
||||
.unwrap(),
|
||||
None
|
||||
);
|
||||
assert!(
|
||||
connection
|
||||
.writer()
|
||||
.send_payload("late".into())
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
}
|
||||
|
||||
/// A caller dropping its response half releases an idle transport without waiting for another command.
|
||||
#[tokio::test]
|
||||
async fn response_half_close_ends_idle_attachment() {
|
||||
let (_connection, events, commands) = connection();
|
||||
drop(commands);
|
||||
timeout(Duration::from_secs(1), events.closed())
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
/// Oversized payloads fail the transport before entering the parser or caller's outgoing queue.
|
||||
#[tokio::test]
|
||||
async fn oversized_payloads_close_both_directions() {
|
||||
let (input_connection, events, mut commands) = connection();
|
||||
events
|
||||
.send("x".repeat(MAX_MESSAGE_BYTES + 1))
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(
|
||||
timeout(Duration::from_secs(1), input_connection.next_event())
|
||||
.await
|
||||
.unwrap()
|
||||
.is_err()
|
||||
);
|
||||
assert_eq!(
|
||||
timeout(Duration::from_secs(1), commands.recv())
|
||||
.await
|
||||
.unwrap(),
|
||||
None
|
||||
);
|
||||
let (output_connection, _events, mut commands) = connection();
|
||||
assert!(
|
||||
output_connection
|
||||
.writer()
|
||||
.send_payload("x".repeat(MAX_MESSAGE_BYTES + 1))
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
assert_eq!(
|
||||
timeout(Duration::from_secs(1), commands.recv())
|
||||
.await
|
||||
.unwrap(),
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
/// A final transcript accepted before peer hangup must reach both the parser and transcript tail.
|
||||
#[tokio::test(flavor = "current_thread")]
|
||||
async fn queued_transcript_survives_peer_hangup() {
|
||||
for _ in 0..64 {
|
||||
let (connection, events, commands) = connection();
|
||||
events
|
||||
.try_send(r#"{"type":"input_transcript.added","item":{"text":"last words"}}"#.into())
|
||||
.unwrap();
|
||||
drop(events);
|
||||
drop(commands);
|
||||
assert_eq!(
|
||||
timeout(Duration::from_secs(1), connection.next_event())
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap(),
|
||||
Some(
|
||||
codex_protocol::protocol::RealtimeEvent::InputTranscriptDelta(
|
||||
codex_protocol::protocol::RealtimeTranscriptDelta {
|
||||
delta: "last words".into()
|
||||
}
|
||||
)
|
||||
)
|
||||
);
|
||||
assert_eq!(
|
||||
connection.events().take_transcript_tail().await,
|
||||
vec![codex_protocol::protocol::RealtimeTranscriptEntry {
|
||||
role: "user".into(),
|
||||
text: "last words".into()
|
||||
}]
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// Explicit shutdown preserves a frame blocked on parser capacity and does not wait for an unsent permit.
|
||||
#[tokio::test]
|
||||
async fn input_stop_drains_backpressure_without_waiting_for_reserved_permit() {
|
||||
let (events, incoming) = mpsc::channel(64);
|
||||
let (outgoing, _commands) = mpsc::channel(1);
|
||||
let connection = RealtimeWebsocketConnection::from_sideband(
|
||||
RealtimeSidebandTransport { incoming, outgoing },
|
||||
RealtimeEventParser::FramelessBidi,
|
||||
);
|
||||
let _permit = events.reserve().await.unwrap();
|
||||
for _ in 0..63 {
|
||||
events
|
||||
.try_send(r#"{"type":"input_transcript.added","item":{"text":"tail"}}"#.into())
|
||||
.unwrap();
|
||||
}
|
||||
timeout(Duration::from_secs(1), async {
|
||||
while events.capacity() < 33 {
|
||||
tokio::task::yield_now().await;
|
||||
}
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
connection.writer().close_sideband_input();
|
||||
for _ in 0..63 {
|
||||
assert_eq!(
|
||||
timeout(Duration::from_secs(1), connection.next_event())
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap(),
|
||||
Some(
|
||||
codex_protocol::protocol::RealtimeEvent::InputTranscriptDelta(
|
||||
codex_protocol::protocol::RealtimeTranscriptDelta {
|
||||
delta: "tail".into()
|
||||
}
|
||||
)
|
||||
)
|
||||
);
|
||||
}
|
||||
assert!(
|
||||
timeout(Duration::from_secs(1), connection.next_event())
|
||||
.await
|
||||
.unwrap()
|
||||
.is_err()
|
||||
);
|
||||
assert!(events.is_closed());
|
||||
}
|
||||
Reference in New Issue
Block a user