mirror of
https://github.com/openai/codex.git
synced 2026-10-19 17:28:46 +00:00
Add WebSocket transport to the code-mode host (#35078)
## What changed - Add a `--listen` option that accepts `stdio`, `stdio://`, or a `ws://IP:PORT` endpoint, while retaining stdio as the default. - Serve the existing length-prefixed protocol in binary WebSocket messages, with isolated connections, shared host limits, and a `/readyz` endpoint. - Reject browser-origin handshakes and contain malformed frames to the affected connection. ## Testing - Cover listen URL parsing and complete-frame encoding and decoding. - Exercise readiness, cell execution, tool callbacks, large frames, concurrent connections, malformed frames, and origin rejection through the WebSocket listener. GitOrigin-RevId: 01c8be4c6256b8ce4a3a0002440dcb3294e5f887
This commit is contained in:
Channing Conger
committed by
copyberry
copyberry
parent
f47f28cd0d
commit
0dfa778dae
@@ -0,0 +1,405 @@
|
||||
use std::process::Stdio;
|
||||
use std::time::Duration;
|
||||
|
||||
use anyhow::Context;
|
||||
use anyhow::Result;
|
||||
use codex_code_mode_protocol::host::Capability;
|
||||
use codex_code_mode_protocol::host::CapabilitySet;
|
||||
use codex_code_mode_protocol::host::ClientHello;
|
||||
use codex_code_mode_protocol::host::ClientToHost;
|
||||
use codex_code_mode_protocol::host::DelegateRequest;
|
||||
use codex_code_mode_protocol::host::DelegateResponse;
|
||||
use codex_code_mode_protocol::host::EncodedFrame;
|
||||
use codex_code_mode_protocol::host::HostHello;
|
||||
use codex_code_mode_protocol::host::HostRequest;
|
||||
use codex_code_mode_protocol::host::HostResponse;
|
||||
use codex_code_mode_protocol::host::HostToClient;
|
||||
use codex_code_mode_protocol::host::MAX_FRAME_BYTES;
|
||||
use codex_code_mode_protocol::host::ProtocolVersion;
|
||||
use codex_code_mode_protocol::host::RequestId;
|
||||
use codex_code_mode_protocol::host::SessionId;
|
||||
use codex_code_mode_protocol::host::SupportedProtocolVersions;
|
||||
use codex_code_mode_protocol::host::WireContentItem;
|
||||
use codex_code_mode_protocol::host::WireExecuteRequest;
|
||||
use codex_code_mode_protocol::host::WireResult;
|
||||
use codex_code_mode_protocol::host::WireRuntimeResponse;
|
||||
use codex_code_mode_protocol::host::WireToolDefinition;
|
||||
use codex_code_mode_protocol::host::WireToolKind;
|
||||
use codex_code_mode_protocol::host::WireToolName;
|
||||
use futures::SinkExt;
|
||||
use futures::StreamExt;
|
||||
use pretty_assertions::assert_eq;
|
||||
use serde_json::json;
|
||||
use tokio::io::AsyncBufReadExt;
|
||||
use tokio::io::AsyncReadExt;
|
||||
use tokio::io::AsyncWriteExt;
|
||||
use tokio::io::BufReader;
|
||||
use tokio::net::TcpStream;
|
||||
use tokio::process::Child;
|
||||
use tokio::process::Command;
|
||||
use tokio::time::timeout;
|
||||
use tokio_tungstenite::MaybeTlsStream;
|
||||
use tokio_tungstenite::WebSocketStream;
|
||||
use tokio_tungstenite::connect_async;
|
||||
use tokio_tungstenite::connect_async_with_config;
|
||||
use tokio_tungstenite::tungstenite::Error as WebSocketError;
|
||||
use tokio_tungstenite::tungstenite::Message;
|
||||
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
|
||||
use tokio_tungstenite::tungstenite::http::HeaderValue;
|
||||
use tokio_tungstenite::tungstenite::http::StatusCode;
|
||||
use tokio_tungstenite::tungstenite::http::header::ORIGIN;
|
||||
use tokio_tungstenite::tungstenite::protocol::WebSocketConfig;
|
||||
|
||||
const TEST_TIMEOUT: Duration = Duration::from_secs(10);
|
||||
const MAX_WEBSOCKET_FRAME_BYTES: usize = MAX_FRAME_BYTES + std::mem::size_of::<u32>();
|
||||
|
||||
struct HostHarness {
|
||||
child: Child,
|
||||
websocket_url: String,
|
||||
}
|
||||
|
||||
struct HostClient {
|
||||
websocket: WebSocketStream<MaybeTlsStream<TcpStream>>,
|
||||
}
|
||||
|
||||
impl HostHarness {
|
||||
async fn start() -> Result<Self> {
|
||||
let host_program = codex_utils_cargo_bin::cargo_bin("codex-code-mode-host")?;
|
||||
let mut command = Command::new(host_program);
|
||||
command
|
||||
.args(["--listen", "ws://127.0.0.1:0"])
|
||||
.stdin(Stdio::null())
|
||||
.stdout(Stdio::piped())
|
||||
.stderr(Stdio::piped())
|
||||
.kill_on_drop(true);
|
||||
let mut child = command.spawn().context("failed to start code-mode host")?;
|
||||
let stdout = child
|
||||
.stdout
|
||||
.take()
|
||||
.context("code-mode host stdout was not captured")?;
|
||||
let mut lines = BufReader::new(stdout).lines();
|
||||
let websocket_url = timeout(TEST_TIMEOUT, lines.next_line())
|
||||
.await
|
||||
.context("timed out waiting for code-mode host websocket URL")??
|
||||
.context("code-mode host exited before publishing its websocket URL")?;
|
||||
if !websocket_url.starts_with("ws://127.0.0.1:") {
|
||||
anyhow::bail!("unexpected code-mode host websocket URL `{websocket_url}`");
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
child,
|
||||
websocket_url,
|
||||
})
|
||||
}
|
||||
|
||||
async fn connect(&self) -> Result<HostClient> {
|
||||
let config = WebSocketConfig::default()
|
||||
.max_frame_size(Some(MAX_WEBSOCKET_FRAME_BYTES))
|
||||
.max_message_size(Some(MAX_WEBSOCKET_FRAME_BYTES));
|
||||
let (websocket, _) = timeout(
|
||||
TEST_TIMEOUT,
|
||||
connect_async_with_config(
|
||||
self.websocket_url.as_str(),
|
||||
Some(config),
|
||||
/*disable_nagle*/ false,
|
||||
),
|
||||
)
|
||||
.await
|
||||
.context("timed out connecting to code-mode host websocket")??;
|
||||
Ok(HostClient { websocket })
|
||||
}
|
||||
}
|
||||
|
||||
impl HostClient {
|
||||
async fn send(&mut self, message: &ClientToHost) -> Result<()> {
|
||||
let frame = EncodedFrame::encode(message)?;
|
||||
self.send_binary(frame.into_framed_bytes()).await
|
||||
}
|
||||
|
||||
async fn send_binary(&mut self, bytes: Vec<u8>) -> Result<()> {
|
||||
timeout(
|
||||
TEST_TIMEOUT,
|
||||
self.websocket.send(Message::Binary(bytes.into())),
|
||||
)
|
||||
.await
|
||||
.context("timed out writing code-mode websocket message")?
|
||||
.context("failed to write code-mode websocket message")
|
||||
}
|
||||
|
||||
async fn read(&mut self) -> Result<HostToClient> {
|
||||
loop {
|
||||
let message = timeout(TEST_TIMEOUT, self.websocket.next())
|
||||
.await
|
||||
.context("timed out waiting for code-mode websocket message")?
|
||||
.context("code-mode websocket closed before returning a message")?
|
||||
.context("failed to read code-mode websocket message")?;
|
||||
match message {
|
||||
Message::Binary(bytes) => {
|
||||
return EncodedFrame::decode_framed(&bytes)
|
||||
.context("failed to decode code-mode websocket frame");
|
||||
}
|
||||
Message::Ping(_) | Message::Pong(_) => {}
|
||||
Message::Close(frame) => {
|
||||
anyhow::bail!("code-mode websocket closed unexpectedly: {frame:?}");
|
||||
}
|
||||
Message::Text(text) => {
|
||||
anyhow::bail!("code-mode host returned a text websocket message: {text}");
|
||||
}
|
||||
Message::Frame(_) => {
|
||||
anyhow::bail!("code-mode host returned an unexpected raw websocket frame");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn negotiate(&mut self, optional_capabilities: CapabilitySet) -> Result<()> {
|
||||
let hello = ClientHello::new(
|
||||
SupportedProtocolVersions::try_new([ProtocolVersion::V1])?,
|
||||
CapabilitySet::empty(),
|
||||
optional_capabilities,
|
||||
)?;
|
||||
self.send(&ClientToHost::ClientHello(hello)).await?;
|
||||
assert_eq!(
|
||||
self.read().await?,
|
||||
HostToClient::HostHello(HostHello::new(ProtocolVersion::V1, CapabilitySet::empty()))
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn open_session(&mut self, session_id: SessionId) -> Result<()> {
|
||||
let id = RequestId::new(/*value*/ 1);
|
||||
self.send(&ClientToHost::Request {
|
||||
id,
|
||||
request: HostRequest::OpenSession {
|
||||
session_id: session_id.clone(),
|
||||
},
|
||||
})
|
||||
.await?;
|
||||
assert_eq!(
|
||||
self.read().await?,
|
||||
HostToClient::Response {
|
||||
id,
|
||||
result: WireResult::Ok {
|
||||
value: HostResponse::SessionReady { session_id },
|
||||
},
|
||||
}
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn websocket_listener_serves_readiness_endpoint() -> Result<()> {
|
||||
let host = HostHarness::start().await?;
|
||||
let address = host
|
||||
.websocket_url
|
||||
.strip_prefix("ws://")
|
||||
.context("code-mode host websocket URL should use ws://")?;
|
||||
|
||||
let response = timeout(TEST_TIMEOUT, async {
|
||||
let mut stream = TcpStream::connect(address)
|
||||
.await
|
||||
.context("failed to connect to code-mode host readiness endpoint")?;
|
||||
let request =
|
||||
format!("GET /readyz HTTP/1.1\r\nHost: {address}\r\nConnection: close\r\n\r\n");
|
||||
stream
|
||||
.write_all(request.as_bytes())
|
||||
.await
|
||||
.context("failed to request code-mode host readiness")?;
|
||||
|
||||
let mut response = String::new();
|
||||
stream
|
||||
.read_to_string(&mut response)
|
||||
.await
|
||||
.context("failed to read code-mode host readiness response")?;
|
||||
Ok::<_, anyhow::Error>(response)
|
||||
})
|
||||
.await
|
||||
.context("timed out requesting code-mode host readiness")??;
|
||||
|
||||
let status_line = response
|
||||
.lines()
|
||||
.next()
|
||||
.context("code-mode host readiness response is missing a status line")?;
|
||||
assert_eq!(status_line, "HTTP/1.1 200 OK");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn websocket_listener_executes_cells_and_forwards_tool_callbacks() -> Result<()> {
|
||||
let host = HostHarness::start().await?;
|
||||
let mut client = host.connect().await?;
|
||||
client.negotiate(CapabilitySet::empty()).await?;
|
||||
|
||||
let session_id = SessionId::new("websocket-session")?;
|
||||
client.open_session(session_id.clone()).await?;
|
||||
|
||||
let execute_id = RequestId::new(/*value*/ 2);
|
||||
client
|
||||
.send(&ClientToHost::Request {
|
||||
id: execute_id,
|
||||
request: HostRequest::Execute {
|
||||
session_id: session_id.clone(),
|
||||
request: WireExecuteRequest {
|
||||
tool_call_id: "websocket-call".to_string(),
|
||||
enabled_tools: vec![WireToolDefinition {
|
||||
name: "echo".to_string(),
|
||||
tool_name: WireToolName {
|
||||
name: "echo".to_string(),
|
||||
namespace: None,
|
||||
},
|
||||
description: String::new(),
|
||||
kind: WireToolKind::Function,
|
||||
input_schema: None,
|
||||
output_schema: None,
|
||||
}],
|
||||
source:
|
||||
r#"const result = await tools.echo({ value: "ping" }); text(result.value);"#
|
||||
.to_string(),
|
||||
yield_time_ms: Some(5_000),
|
||||
max_output_tokens: Some(1_000),
|
||||
},
|
||||
},
|
||||
})
|
||||
.await?;
|
||||
|
||||
let started = client.read().await?;
|
||||
let HostToClient::Response {
|
||||
id,
|
||||
result:
|
||||
WireResult::Ok {
|
||||
value: HostResponse::ExecutionStarted { cell_id },
|
||||
},
|
||||
} = started
|
||||
else {
|
||||
anyhow::bail!("expected execution-started response, got {started:?}");
|
||||
};
|
||||
assert_eq!(id, execute_id);
|
||||
|
||||
let callback = client.read().await?;
|
||||
let HostToClient::DelegateRequest {
|
||||
id: delegate_id,
|
||||
session_id: callback_session_id,
|
||||
request: DelegateRequest::InvokeTool { invocation },
|
||||
} = callback
|
||||
else {
|
||||
anyhow::bail!("expected tool callback, got {callback:?}");
|
||||
};
|
||||
assert_eq!(callback_session_id, session_id);
|
||||
assert_eq!(invocation.input, Some(json!({ "value": "ping" })));
|
||||
|
||||
client
|
||||
.send(&ClientToHost::DelegateResponse {
|
||||
id: delegate_id,
|
||||
result: WireResult::Ok {
|
||||
value: DelegateResponse::ToolResult {
|
||||
result: json!({ "value": "pong" }),
|
||||
},
|
||||
},
|
||||
})
|
||||
.await?;
|
||||
|
||||
assert_eq!(
|
||||
client.read().await?,
|
||||
HostToClient::InitialResponse {
|
||||
id: execute_id,
|
||||
result: WireResult::Ok {
|
||||
value: WireRuntimeResponse::Result {
|
||||
cell_id,
|
||||
content_items: vec![WireContentItem::InputText {
|
||||
text: "pong".to_string(),
|
||||
}],
|
||||
error_text: None,
|
||||
},
|
||||
},
|
||||
}
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn websocket_listener_accepts_frames_larger_than_default_websocket_limit() -> Result<()> {
|
||||
let host = HostHarness::start().await?;
|
||||
let mut client = host.connect().await?;
|
||||
let capability = Capability::new("x".repeat((16 * 1024 * 1024) + 1))?;
|
||||
|
||||
client
|
||||
.negotiate(CapabilitySet::try_new([capability])?)
|
||||
.await
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn websocket_listener_keeps_connections_and_session_ids_isolated() -> Result<()> {
|
||||
let host = HostHarness::start().await?;
|
||||
let mut first = host.connect().await?;
|
||||
let mut second = host.connect().await?;
|
||||
first.negotiate(CapabilitySet::empty()).await?;
|
||||
second.negotiate(CapabilitySet::empty()).await?;
|
||||
|
||||
let session_id = SessionId::new("shared-session-name")?;
|
||||
first.open_session(session_id.clone()).await?;
|
||||
second.open_session(session_id).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn malformed_websocket_frame_does_not_stop_the_listener() -> Result<()> {
|
||||
let mut host = HostHarness::start().await?;
|
||||
let stderr = host
|
||||
.child
|
||||
.stderr
|
||||
.take()
|
||||
.context("code-mode host stderr was not captured")?;
|
||||
let mut stderr_lines = BufReader::new(stderr).lines();
|
||||
let mut malformed = host.connect().await?;
|
||||
malformed.send_binary(vec![1, 0, 0, 0, b'{']).await?;
|
||||
|
||||
let close = timeout(TEST_TIMEOUT, malformed.websocket.next())
|
||||
.await
|
||||
.context("timed out waiting for malformed websocket connection to close")?;
|
||||
if let Some(Ok(message)) = close
|
||||
&& !matches!(message, Message::Close(_))
|
||||
{
|
||||
anyhow::bail!("malformed websocket returned an unexpected message: {message:?}");
|
||||
}
|
||||
|
||||
let diagnostic = timeout(TEST_TIMEOUT, async {
|
||||
loop {
|
||||
let line = stderr_lines
|
||||
.next_line()
|
||||
.await?
|
||||
.context("code-mode host exited before reporting the malformed frame")?;
|
||||
if line.contains("code-mode host websocket connection failed") {
|
||||
return Ok::<_, anyhow::Error>(line);
|
||||
}
|
||||
}
|
||||
})
|
||||
.await
|
||||
.context("timed out waiting for the malformed websocket diagnostic")??;
|
||||
assert!(
|
||||
diagnostic.contains("failed to read code-mode client hello"),
|
||||
"unexpected malformed websocket diagnostic: {diagnostic}",
|
||||
);
|
||||
|
||||
let mut recovered = host.connect().await?;
|
||||
recovered.negotiate(CapabilitySet::empty()).await
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn websocket_listener_rejects_browser_origin_handshakes() -> Result<()> {
|
||||
let host = HostHarness::start().await?;
|
||||
let mut request = host.websocket_url.as_str().into_client_request()?;
|
||||
request
|
||||
.headers_mut()
|
||||
.insert(ORIGIN, HeaderValue::from_static("https://evil.example"));
|
||||
|
||||
let error = match connect_async(request).await {
|
||||
Ok(_) => anyhow::bail!("browser-origin websocket handshake should be rejected"),
|
||||
Err(error) => error,
|
||||
};
|
||||
let WebSocketError::Http(response) = error else {
|
||||
anyhow::bail!("browser-origin websocket handshake failed unexpectedly: {error}");
|
||||
};
|
||||
assert_eq!(response.status(), StatusCode::FORBIDDEN);
|
||||
Ok(())
|
||||
}
|
||||
Reference in New Issue
Block a user