mirror of
https://github.com/openai/codex.git
synced 2026-10-23 18:08:45 +00:00
feat(tui): stream staged uploads over websocket sftp
This commit is contained in:
@@ -16,6 +16,7 @@
|
||||
//! surfaces as channel-full errors rather than unbounded memory growth.
|
||||
|
||||
mod remote;
|
||||
mod sftp_upload;
|
||||
|
||||
use std::error::Error;
|
||||
use std::fmt;
|
||||
@@ -850,6 +851,16 @@ impl AppServerRequestHandle {
|
||||
Self::Remote(handle) => handle.request_typed(request).await,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn upload_file_over_sftp(&self, path: &str, bytes: &[u8]) -> IoResult<()> {
|
||||
match self {
|
||||
Self::InProcess(_) => Err(IoError::new(
|
||||
ErrorKind::Unsupported,
|
||||
"staged SFTP uploads require a remote app-server connection",
|
||||
)),
|
||||
Self::Remote(handle) => handle.upload_file_over_sftp(path, bytes).await,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl AppServerClient {
|
||||
@@ -1453,6 +1464,46 @@ mod tests {
|
||||
client.shutdown().await.expect("shutdown should complete");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn remote_binary_lane_roundtrip_works() {
|
||||
let websocket_url = start_test_remote_server(|mut websocket| async move {
|
||||
expect_remote_initialize(&mut websocket).await;
|
||||
let frame = websocket
|
||||
.next()
|
||||
.await
|
||||
.expect("binary frame should arrive")
|
||||
.expect("binary frame should decode");
|
||||
let Message::Binary(bytes) = frame else {
|
||||
panic!("expected binary frame");
|
||||
};
|
||||
assert_eq!(bytes.as_ref(), [1, 2, 3]);
|
||||
websocket
|
||||
.send(Message::Binary(vec![4, 5, 6].into()))
|
||||
.await
|
||||
.expect("binary response should send");
|
||||
websocket.close(None).await.expect("close should succeed");
|
||||
})
|
||||
.await;
|
||||
let client = RemoteAppServerClient::connect(test_remote_connect_args(websocket_url))
|
||||
.await
|
||||
.expect("remote client should connect");
|
||||
let handle = client.request_handle();
|
||||
|
||||
handle
|
||||
.send_binary(vec![1, 2, 3])
|
||||
.await
|
||||
.expect("binary request should send");
|
||||
assert_eq!(
|
||||
handle
|
||||
.next_binary()
|
||||
.await
|
||||
.expect("binary response should arrive"),
|
||||
vec![4, 5, 6]
|
||||
);
|
||||
|
||||
client.shutdown().await.expect("shutdown should complete");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn remote_typed_request_accepts_large_single_frame_response() {
|
||||
let padding = "x".repeat((17 << 20) + 1024);
|
||||
|
||||
@@ -40,6 +40,7 @@ use futures::SinkExt;
|
||||
use futures::StreamExt;
|
||||
use serde::de::DeserializeOwned;
|
||||
use tokio::net::TcpStream;
|
||||
use tokio::sync::Mutex;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio::sync::oneshot;
|
||||
use tokio::time::timeout;
|
||||
@@ -120,6 +121,13 @@ enum RemoteClientCommand {
|
||||
error: JSONRPCErrorError,
|
||||
response_tx: oneshot::Sender<IoResult<()>>,
|
||||
},
|
||||
SendBinary {
|
||||
bytes: Vec<u8>,
|
||||
response_tx: oneshot::Sender<IoResult<()>>,
|
||||
},
|
||||
NextBinary {
|
||||
response_tx: oneshot::Sender<IoResult<Vec<u8>>>,
|
||||
},
|
||||
Shutdown {
|
||||
response_tx: oneshot::Sender<IoResult<()>>,
|
||||
},
|
||||
@@ -129,12 +137,14 @@ pub struct RemoteAppServerClient {
|
||||
command_tx: mpsc::Sender<RemoteClientCommand>,
|
||||
event_rx: mpsc::UnboundedReceiver<AppServerEvent>,
|
||||
pending_events: VecDeque<AppServerEvent>,
|
||||
upload_lock: std::sync::Arc<Mutex<()>>,
|
||||
worker_handle: tokio::task::JoinHandle<()>,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct RemoteAppServerRequestHandle {
|
||||
command_tx: mpsc::Sender<RemoteClientCommand>,
|
||||
upload_lock: std::sync::Arc<Mutex<()>>,
|
||||
}
|
||||
|
||||
impl RemoteAppServerClient {
|
||||
@@ -212,6 +222,8 @@ impl RemoteAppServerClient {
|
||||
let worker_handle = tokio::spawn(async move {
|
||||
let mut pending_requests =
|
||||
HashMap::<RequestId, oneshot::Sender<IoResult<RequestResult>>>::new();
|
||||
let mut pending_binary_packets = VecDeque::<Vec<u8>>::new();
|
||||
let mut pending_binary_waiters = VecDeque::<oneshot::Sender<IoResult<Vec<u8>>>>::new();
|
||||
let mut worker_exit_error: Option<(ErrorKind, String)> = None;
|
||||
loop {
|
||||
tokio::select! {
|
||||
@@ -298,6 +310,38 @@ impl RemoteAppServerClient {
|
||||
.await;
|
||||
let _ = response_tx.send(result);
|
||||
}
|
||||
RemoteClientCommand::SendBinary { bytes, response_tx } => {
|
||||
let result = write_binary_message(
|
||||
&mut stream,
|
||||
bytes,
|
||||
&websocket_url,
|
||||
)
|
||||
.await;
|
||||
let should_exit = result.is_err();
|
||||
let err_message = result.as_ref().err().map(ToString::to_string);
|
||||
let _ = response_tx.send(result);
|
||||
if should_exit {
|
||||
let message = format!(
|
||||
"remote app server at `{websocket_url}` binary write failed: {}",
|
||||
err_message.unwrap_or_else(|| "unknown error".to_string())
|
||||
);
|
||||
let _ = deliver_event(
|
||||
&event_tx,
|
||||
AppServerEvent::Disconnected {
|
||||
message: message.clone(),
|
||||
},
|
||||
);
|
||||
worker_exit_error = Some((ErrorKind::BrokenPipe, message));
|
||||
break;
|
||||
}
|
||||
}
|
||||
RemoteClientCommand::NextBinary { response_tx } => {
|
||||
if let Some(bytes) = pending_binary_packets.pop_front() {
|
||||
let _ = response_tx.send(Ok(bytes));
|
||||
} else {
|
||||
pending_binary_waiters.push_back(response_tx);
|
||||
}
|
||||
}
|
||||
RemoteClientCommand::Shutdown { response_tx } => {
|
||||
let close_result = stream.close(None).await.map_err(|err| {
|
||||
IoError::other(format!(
|
||||
@@ -421,8 +465,15 @@ impl RemoteAppServerClient {
|
||||
));
|
||||
break;
|
||||
}
|
||||
Some(Ok(Message::Binary(_)))
|
||||
| Some(Ok(Message::Ping(_)))
|
||||
Some(Ok(Message::Binary(bytes))) => {
|
||||
let bytes = bytes.to_vec();
|
||||
if let Some(response_tx) = pending_binary_waiters.pop_front() {
|
||||
let _ = response_tx.send(Ok(bytes));
|
||||
} else {
|
||||
pending_binary_packets.push_back(bytes);
|
||||
}
|
||||
}
|
||||
Some(Ok(Message::Ping(_)))
|
||||
| Some(Ok(Message::Pong(_)))
|
||||
| Some(Ok(Message::Frame(_))) => {}
|
||||
Some(Err(err)) => {
|
||||
@@ -465,12 +516,16 @@ impl RemoteAppServerClient {
|
||||
for (_, response_tx) in pending_requests {
|
||||
let _ = response_tx.send(Err(IoError::new(err_kind, err_message.clone())));
|
||||
}
|
||||
for response_tx in pending_binary_waiters {
|
||||
let _ = response_tx.send(Err(IoError::new(err_kind, err_message.clone())));
|
||||
}
|
||||
});
|
||||
|
||||
Ok(Self {
|
||||
command_tx,
|
||||
event_rx,
|
||||
pending_events: pending_events.into(),
|
||||
upload_lock: std::sync::Arc::new(Mutex::new(())),
|
||||
worker_handle,
|
||||
})
|
||||
}
|
||||
@@ -478,6 +533,7 @@ impl RemoteAppServerClient {
|
||||
pub fn request_handle(&self) -> RemoteAppServerRequestHandle {
|
||||
RemoteAppServerRequestHandle {
|
||||
command_tx: self.command_tx.clone(),
|
||||
upload_lock: std::sync::Arc::clone(&self.upload_lock),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -611,6 +667,7 @@ impl RemoteAppServerClient {
|
||||
command_tx,
|
||||
event_rx,
|
||||
pending_events: _pending_events,
|
||||
upload_lock: _upload_lock,
|
||||
worker_handle,
|
||||
} = self;
|
||||
let mut worker_handle = worker_handle;
|
||||
@@ -675,6 +732,53 @@ impl RemoteAppServerRequestHandle {
|
||||
serde_json::from_value(result)
|
||||
.map_err(|source| TypedRequestError::Deserialize { method, source })
|
||||
}
|
||||
|
||||
pub(crate) async fn send_binary(&self, bytes: Vec<u8>) -> IoResult<()> {
|
||||
let (response_tx, response_rx) = oneshot::channel();
|
||||
self.command_tx
|
||||
.send(RemoteClientCommand::SendBinary { bytes, response_tx })
|
||||
.await
|
||||
.map_err(|_| {
|
||||
IoError::new(
|
||||
ErrorKind::BrokenPipe,
|
||||
"remote app-server worker channel is closed",
|
||||
)
|
||||
})?;
|
||||
response_rx.await.map_err(|_| {
|
||||
IoError::new(
|
||||
ErrorKind::BrokenPipe,
|
||||
"remote app-server binary send channel is closed",
|
||||
)
|
||||
})?
|
||||
}
|
||||
|
||||
pub(crate) async fn next_binary(&self) -> IoResult<Vec<u8>> {
|
||||
let (response_tx, response_rx) = oneshot::channel();
|
||||
self.command_tx
|
||||
.send(RemoteClientCommand::NextBinary { response_tx })
|
||||
.await
|
||||
.map_err(|_| {
|
||||
IoError::new(
|
||||
ErrorKind::BrokenPipe,
|
||||
"remote app-server worker channel is closed",
|
||||
)
|
||||
})?;
|
||||
response_rx.await.map_err(|_| {
|
||||
IoError::new(
|
||||
ErrorKind::BrokenPipe,
|
||||
"remote app-server binary receive channel is closed",
|
||||
)
|
||||
})?
|
||||
}
|
||||
|
||||
#[expect(
|
||||
clippy::await_holding_invalid_type,
|
||||
reason = "one SFTP upload must own the shared websocket binary lane at a time"
|
||||
)]
|
||||
pub async fn upload_file_over_sftp(&self, path: &str, bytes: &[u8]) -> IoResult<()> {
|
||||
let _guard = self.upload_lock.lock().await;
|
||||
crate::sftp_upload::upload_file_over_sftp(self, path, bytes).await
|
||||
}
|
||||
}
|
||||
|
||||
async fn initialize_remote_connection(
|
||||
@@ -864,6 +968,21 @@ async fn write_jsonrpc_message(
|
||||
))
|
||||
})
|
||||
}
|
||||
|
||||
async fn write_binary_message(
|
||||
stream: &mut WebSocketStream<MaybeTlsStream<TcpStream>>,
|
||||
bytes: Vec<u8>,
|
||||
websocket_url: &str,
|
||||
) -> IoResult<()> {
|
||||
stream
|
||||
.send(Message::Binary(bytes.into()))
|
||||
.await
|
||||
.map_err(|err| {
|
||||
IoError::other(format!(
|
||||
"failed to write binary websocket message to `{websocket_url}`: {err}"
|
||||
))
|
||||
})
|
||||
}
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
@@ -879,6 +998,7 @@ mod tests {
|
||||
command_tx,
|
||||
event_rx,
|
||||
pending_events: VecDeque::new(),
|
||||
upload_lock: std::sync::Arc::new(Mutex::new(())),
|
||||
worker_handle,
|
||||
};
|
||||
|
||||
|
||||
@@ -0,0 +1,290 @@
|
||||
use std::io::Error as IoError;
|
||||
use std::io::ErrorKind;
|
||||
use std::io::Result as IoResult;
|
||||
|
||||
use crate::remote::RemoteAppServerRequestHandle;
|
||||
|
||||
const SSH_FXP_INIT: u8 = 1;
|
||||
const SSH_FXP_VERSION: u8 = 2;
|
||||
const SSH_FXP_OPEN: u8 = 3;
|
||||
const SSH_FXP_CLOSE: u8 = 4;
|
||||
const SSH_FXP_WRITE: u8 = 6;
|
||||
const SSH_FXP_STATUS: u8 = 101;
|
||||
const SSH_FXP_HANDLE: u8 = 102;
|
||||
|
||||
const SSH_FX_OK: u32 = 0;
|
||||
const SSH_FXF_WRITE: u32 = 0x0000_0002;
|
||||
const SSH_FXF_CREAT: u32 = 0x0000_0008;
|
||||
const SSH_FXF_TRUNC: u32 = 0x0000_0010;
|
||||
const SFTP_VERSION: u32 = 3;
|
||||
const SFTP_WRITE_CHUNK_BYTES: usize = 64 * 1024;
|
||||
const MAX_SFTP_PACKET_BYTES: usize = 1024 * 1024;
|
||||
|
||||
pub(crate) async fn upload_file_over_sftp(
|
||||
handle: &RemoteAppServerRequestHandle,
|
||||
path: &str,
|
||||
bytes: &[u8],
|
||||
) -> IoResult<()> {
|
||||
let mut packets = PacketReader::default();
|
||||
|
||||
handle.send_binary(init_packet()).await?;
|
||||
expect_version(&packets.next_packet(handle).await?)?;
|
||||
|
||||
let mut request_id = 1;
|
||||
handle
|
||||
.send_binary(open_packet(
|
||||
request_id,
|
||||
path,
|
||||
SSH_FXF_WRITE | SSH_FXF_CREAT | SSH_FXF_TRUNC,
|
||||
))
|
||||
.await?;
|
||||
let file_handle = expect_handle(&packets.next_packet(handle).await?, request_id)?;
|
||||
|
||||
for (chunk_index, chunk) in bytes.chunks(SFTP_WRITE_CHUNK_BYTES).enumerate() {
|
||||
request_id += 1;
|
||||
let offset = chunk_index
|
||||
.checked_mul(SFTP_WRITE_CHUNK_BYTES)
|
||||
.and_then(|offset| u64::try_from(offset).ok())
|
||||
.ok_or_else(|| IoError::other("staged upload offset overflowed"))?;
|
||||
handle
|
||||
.send_binary(write_packet(request_id, &file_handle, offset, chunk))
|
||||
.await?;
|
||||
expect_ok_status(&packets.next_packet(handle).await?, request_id)?;
|
||||
}
|
||||
|
||||
request_id += 1;
|
||||
handle
|
||||
.send_binary(close_packet(request_id, &file_handle))
|
||||
.await?;
|
||||
expect_ok_status(&packets.next_packet(handle).await?, request_id)
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct PacketReader {
|
||||
pending_bytes: Vec<u8>,
|
||||
}
|
||||
|
||||
impl PacketReader {
|
||||
async fn next_packet(&mut self, handle: &RemoteAppServerRequestHandle) -> IoResult<Vec<u8>> {
|
||||
loop {
|
||||
if let Some(packet) = take_packet(&mut self.pending_bytes)? {
|
||||
return Ok(packet);
|
||||
}
|
||||
self.pending_bytes.extend(handle.next_binary().await?);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn take_packet(buffer: &mut Vec<u8>) -> IoResult<Option<Vec<u8>>> {
|
||||
if buffer.len() < 4 {
|
||||
return Ok(None);
|
||||
}
|
||||
let packet_len = u32::from_be_bytes(
|
||||
buffer[..4]
|
||||
.try_into()
|
||||
.map_err(|_| invalid_packet("invalid SFTP packet length prefix"))?,
|
||||
) as usize;
|
||||
if packet_len > MAX_SFTP_PACKET_BYTES {
|
||||
return Err(invalid_packet("SFTP packet exceeds 1 MiB"));
|
||||
}
|
||||
if buffer.len() < packet_len + 4 {
|
||||
return Ok(None);
|
||||
}
|
||||
let packet = buffer[4..packet_len + 4].to_vec();
|
||||
buffer.drain(..packet_len + 4);
|
||||
Ok(Some(packet))
|
||||
}
|
||||
|
||||
fn init_packet() -> Vec<u8> {
|
||||
packet(|payload| {
|
||||
payload.push(SSH_FXP_INIT);
|
||||
payload.extend_from_slice(&SFTP_VERSION.to_be_bytes());
|
||||
})
|
||||
}
|
||||
|
||||
fn open_packet(request_id: u32, path: &str, flags: u32) -> Vec<u8> {
|
||||
packet(|payload| {
|
||||
payload.push(SSH_FXP_OPEN);
|
||||
payload.extend_from_slice(&request_id.to_be_bytes());
|
||||
push_string(payload, path.as_bytes());
|
||||
payload.extend_from_slice(&flags.to_be_bytes());
|
||||
payload.extend_from_slice(&0_u32.to_be_bytes());
|
||||
})
|
||||
}
|
||||
|
||||
fn write_packet(request_id: u32, handle: &[u8], offset: u64, data: &[u8]) -> Vec<u8> {
|
||||
packet(|payload| {
|
||||
payload.push(SSH_FXP_WRITE);
|
||||
payload.extend_from_slice(&request_id.to_be_bytes());
|
||||
push_string(payload, handle);
|
||||
payload.extend_from_slice(&offset.to_be_bytes());
|
||||
push_string(payload, data);
|
||||
})
|
||||
}
|
||||
|
||||
fn close_packet(request_id: u32, handle: &[u8]) -> Vec<u8> {
|
||||
packet(|payload| {
|
||||
payload.push(SSH_FXP_CLOSE);
|
||||
payload.extend_from_slice(&request_id.to_be_bytes());
|
||||
push_string(payload, handle);
|
||||
})
|
||||
}
|
||||
|
||||
fn expect_version(packet: &[u8]) -> IoResult<()> {
|
||||
let mut cursor = PacketCursor::new(packet);
|
||||
if cursor.read_u8()? != SSH_FXP_VERSION {
|
||||
return Err(invalid_packet(
|
||||
"staged upload returned a non-version response",
|
||||
));
|
||||
}
|
||||
let version = cursor.read_u32()?;
|
||||
cursor.finish()?;
|
||||
if version != SFTP_VERSION {
|
||||
return Err(invalid_packet(format!(
|
||||
"staged upload returned SFTP version {version}"
|
||||
)));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn expect_handle(packet: &[u8], request_id: u32) -> IoResult<Vec<u8>> {
|
||||
let mut cursor = PacketCursor::new(packet);
|
||||
if cursor.read_u8()? != SSH_FXP_HANDLE {
|
||||
return Err(invalid_packet(
|
||||
"staged upload returned a non-handle response",
|
||||
));
|
||||
}
|
||||
expect_request_id(&mut cursor, request_id)?;
|
||||
let handle = cursor.read_string()?;
|
||||
cursor.finish()?;
|
||||
Ok(handle)
|
||||
}
|
||||
|
||||
fn expect_ok_status(packet: &[u8], request_id: u32) -> IoResult<()> {
|
||||
let mut cursor = PacketCursor::new(packet);
|
||||
if cursor.read_u8()? != SSH_FXP_STATUS {
|
||||
return Err(invalid_packet(
|
||||
"staged upload returned a non-status response",
|
||||
));
|
||||
}
|
||||
expect_request_id(&mut cursor, request_id)?;
|
||||
let status = cursor.read_u32()?;
|
||||
let message = cursor.read_string_utf8()?;
|
||||
let _language = cursor.read_string_utf8()?;
|
||||
cursor.finish()?;
|
||||
if status != SSH_FX_OK {
|
||||
return Err(invalid_packet(format!(
|
||||
"staged upload returned status {status}: {message}"
|
||||
)));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn expect_request_id(cursor: &mut PacketCursor<'_>, request_id: u32) -> IoResult<()> {
|
||||
let response_request_id = cursor.read_u32()?;
|
||||
if response_request_id != request_id {
|
||||
return Err(invalid_packet(format!(
|
||||
"staged upload returned response {response_request_id} for request {request_id}"
|
||||
)));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn packet(fill_payload: impl FnOnce(&mut Vec<u8>)) -> Vec<u8> {
|
||||
let mut payload = Vec::new();
|
||||
fill_payload(&mut payload);
|
||||
let mut packet = Vec::with_capacity(payload.len() + 4);
|
||||
packet.extend_from_slice(&(payload.len() as u32).to_be_bytes());
|
||||
packet.extend_from_slice(&payload);
|
||||
packet
|
||||
}
|
||||
|
||||
fn push_string(buffer: &mut Vec<u8>, value: &[u8]) {
|
||||
buffer.extend_from_slice(&(value.len() as u32).to_be_bytes());
|
||||
buffer.extend_from_slice(value);
|
||||
}
|
||||
|
||||
struct PacketCursor<'a> {
|
||||
bytes: &'a [u8],
|
||||
offset: usize,
|
||||
}
|
||||
|
||||
impl<'a> PacketCursor<'a> {
|
||||
fn new(bytes: &'a [u8]) -> Self {
|
||||
Self { bytes, offset: 0 }
|
||||
}
|
||||
|
||||
fn read_u8(&mut self) -> IoResult<u8> {
|
||||
let value = *self
|
||||
.bytes
|
||||
.get(self.offset)
|
||||
.ok_or_else(|| invalid_packet("unexpected end of SFTP packet"))?;
|
||||
self.offset += 1;
|
||||
Ok(value)
|
||||
}
|
||||
|
||||
fn read_u32(&mut self) -> IoResult<u32> {
|
||||
let bytes = self.read_bytes(/*len*/ 4)?;
|
||||
Ok(u32::from_be_bytes(
|
||||
bytes
|
||||
.try_into()
|
||||
.map_err(|_| invalid_packet("invalid u32 field"))?,
|
||||
))
|
||||
}
|
||||
|
||||
fn read_string(&mut self) -> IoResult<Vec<u8>> {
|
||||
let len = self.read_u32()? as usize;
|
||||
Ok(self.read_bytes(len)?.to_vec())
|
||||
}
|
||||
|
||||
fn read_string_utf8(&mut self) -> IoResult<String> {
|
||||
String::from_utf8(self.read_string()?)
|
||||
.map_err(|_| invalid_packet("invalid UTF-8 string in SFTP packet"))
|
||||
}
|
||||
|
||||
fn read_bytes(&mut self, len: usize) -> IoResult<&'a [u8]> {
|
||||
let end = self
|
||||
.offset
|
||||
.checked_add(len)
|
||||
.ok_or_else(|| invalid_packet("SFTP packet offset overflowed"))?;
|
||||
let bytes = self
|
||||
.bytes
|
||||
.get(self.offset..end)
|
||||
.ok_or_else(|| invalid_packet("unexpected end of SFTP packet"))?;
|
||||
self.offset = end;
|
||||
Ok(bytes)
|
||||
}
|
||||
|
||||
fn finish(&self) -> IoResult<()> {
|
||||
if self.offset == self.bytes.len() {
|
||||
Ok(())
|
||||
} else {
|
||||
Err(invalid_packet("trailing bytes in SFTP packet"))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn invalid_packet(message: impl Into<String>) -> IoError {
|
||||
IoError::new(ErrorKind::InvalidData, message.into())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn take_packet_waits_for_a_complete_binary_chunk() {
|
||||
let mut bytes = vec![0, 0, 0, 2, 7];
|
||||
assert_eq!(
|
||||
take_packet(&mut bytes).expect("partial packet should parse"),
|
||||
None
|
||||
);
|
||||
|
||||
bytes.push(8);
|
||||
assert_eq!(
|
||||
take_packet(&mut bytes).expect("complete packet should parse"),
|
||||
Some(vec![7, 8])
|
||||
);
|
||||
assert!(bytes.is_empty());
|
||||
}
|
||||
}
|
||||
@@ -5,6 +5,11 @@
|
||||
//! the main event loop remains single-threaded.
|
||||
|
||||
use super::*;
|
||||
use base64::Engine;
|
||||
use codex_app_server_protocol::FsCreateUploadParams;
|
||||
use codex_app_server_protocol::FsCreateUploadResponse;
|
||||
use codex_app_server_protocol::FsWriteFileParams;
|
||||
use codex_app_server_protocol::FsWriteFileResponse;
|
||||
use codex_app_server_protocol::HookTrustStatus;
|
||||
use codex_app_server_protocol::MarketplaceAddParams;
|
||||
use codex_app_server_protocol::MarketplaceAddResponse;
|
||||
@@ -16,6 +21,16 @@ use codex_app_server_protocol::MarketplaceUpgradeResponse;
|
||||
use codex_app_server_protocol::RequestId;
|
||||
|
||||
use codex_utils_absolute_path::AbsolutePathBuf;
|
||||
use serde::Deserialize;
|
||||
use tokio::io::AsyncReadExt;
|
||||
|
||||
const MAX_UPLOAD_FILE_BYTES: u64 = 45 * 1024 * 1024;
|
||||
|
||||
#[derive(Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct RemoteFsCreateUploadResponse {
|
||||
path: String,
|
||||
}
|
||||
|
||||
impl App {
|
||||
pub(super) fn fetch_mcp_inventory(
|
||||
@@ -141,6 +156,17 @@ impl App {
|
||||
});
|
||||
}
|
||||
|
||||
pub(super) fn upload_local_file(&mut self, app_server: &AppServerSession, local_path: PathBuf) {
|
||||
let request_handle = app_server.request_handle();
|
||||
let app_event_tx = self.app_event_tx.clone();
|
||||
tokio::spawn(async move {
|
||||
let result = upload_local_file_request(request_handle, &local_path)
|
||||
.await
|
||||
.map_err(|err| err.to_string());
|
||||
app_event_tx.send(AppEvent::LocalFileUploaded { local_path, result });
|
||||
});
|
||||
}
|
||||
|
||||
pub(super) fn fetch_hooks_list(&mut self, app_server: &AppServerSession, cwd: PathBuf) {
|
||||
let request_handle = app_server.request_handle();
|
||||
let app_event_tx = self.app_event_tx.clone();
|
||||
@@ -604,6 +630,74 @@ pub(super) async fn fetch_all_mcp_server_statuses(
|
||||
Ok(statuses)
|
||||
}
|
||||
|
||||
pub(super) async fn upload_local_file_request(
|
||||
request_handle: AppServerRequestHandle,
|
||||
local_path: &Path,
|
||||
) -> Result<PathBuf> {
|
||||
let file_name = local_path
|
||||
.file_name()
|
||||
.and_then(|name| name.to_str())
|
||||
.ok_or_else(|| color_eyre::eyre::eyre!("Upload path must name a file."))?
|
||||
.to_string();
|
||||
let metadata = tokio::fs::metadata(local_path).await?;
|
||||
if !metadata.is_file() {
|
||||
color_eyre::eyre::bail!("Upload path must name a regular file.");
|
||||
}
|
||||
ensure_upload_file_size(metadata.len())?;
|
||||
let file = tokio::fs::File::open(local_path).await?;
|
||||
let mut bytes = Vec::new();
|
||||
file.take(MAX_UPLOAD_FILE_BYTES + 1)
|
||||
.read_to_end(&mut bytes)
|
||||
.await?;
|
||||
ensure_upload_file_size(bytes.len() as u64)?;
|
||||
match &request_handle {
|
||||
AppServerRequestHandle::Remote(_) => {
|
||||
let request_id = RequestId::String(format!("upload-local-file-{}", Uuid::new_v4()));
|
||||
let response: RemoteFsCreateUploadResponse = request_handle
|
||||
.request_typed(ClientRequest::FsCreateUpload {
|
||||
request_id,
|
||||
params: FsCreateUploadParams { file_name },
|
||||
})
|
||||
.await
|
||||
.wrap_err("fs/createUpload failed in TUI")?;
|
||||
request_handle
|
||||
.upload_file_over_sftp(&response.path, &bytes)
|
||||
.await
|
||||
.wrap_err("staged upload stream failed in TUI")?;
|
||||
Ok(PathBuf::from(response.path))
|
||||
}
|
||||
AppServerRequestHandle::InProcess(_) => {
|
||||
let request_id = RequestId::String(format!("upload-local-file-{}", Uuid::new_v4()));
|
||||
let response: FsCreateUploadResponse = request_handle
|
||||
.request_typed(ClientRequest::FsCreateUpload {
|
||||
request_id,
|
||||
params: FsCreateUploadParams { file_name },
|
||||
})
|
||||
.await
|
||||
.wrap_err("fs/createUpload failed in TUI")?;
|
||||
let request_id = RequestId::String(format!("write-local-upload-{}", Uuid::new_v4()));
|
||||
request_handle
|
||||
.request_typed::<FsWriteFileResponse>(ClientRequest::FsWriteFile {
|
||||
request_id,
|
||||
params: FsWriteFileParams {
|
||||
path: response.path.clone(),
|
||||
data_base64: base64::engine::general_purpose::STANDARD.encode(bytes),
|
||||
},
|
||||
})
|
||||
.await
|
||||
.wrap_err("fs/writeFile failed for local staged upload in TUI")?;
|
||||
Ok(response.path.into())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn ensure_upload_file_size(file_size: u64) -> Result<()> {
|
||||
if file_size > MAX_UPLOAD_FILE_BYTES {
|
||||
color_eyre::eyre::bail!("Uploads accept files up to {MAX_UPLOAD_FILE_BYTES} bytes.");
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(super) async fn fetch_account_rate_limits(
|
||||
request_handle: AppServerRequestHandle,
|
||||
) -> Result<Vec<RateLimitSnapshot>> {
|
||||
@@ -1040,6 +1134,17 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn upload_rejects_files_larger_than_the_upload_limit() {
|
||||
let err = ensure_upload_file_size(MAX_UPLOAD_FILE_BYTES + 1)
|
||||
.expect_err("file should exceed upload limit");
|
||||
|
||||
assert_eq!(
|
||||
err.to_string(),
|
||||
format!("Uploads accept files up to {MAX_UPLOAD_FILE_BYTES} bytes.")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mcp_inventory_maps_prefix_tool_names_by_server() {
|
||||
let statuses = vec![
|
||||
|
||||
@@ -666,6 +666,21 @@ impl App {
|
||||
AppEvent::FileSearchResult { query, matches } => {
|
||||
self.chat_widget.apply_file_search_result(query, matches);
|
||||
}
|
||||
AppEvent::UploadLocalFile { path } => {
|
||||
self.upload_local_file(app_server, path);
|
||||
}
|
||||
AppEvent::LocalFileUploaded { local_path, result } => match result {
|
||||
Ok(remote_path) => {
|
||||
self.chat_widget.insert_uploaded_file_path(&remote_path);
|
||||
self.chat_widget.maybe_send_next_queued_input();
|
||||
}
|
||||
Err(err) => {
|
||||
self.chat_widget.add_error_message(format!(
|
||||
"Failed to upload '{}': {err}",
|
||||
local_path.display()
|
||||
));
|
||||
}
|
||||
},
|
||||
AppEvent::RefreshRateLimits { origin } => {
|
||||
self.refresh_rate_limits(app_server, origin);
|
||||
}
|
||||
|
||||
@@ -229,6 +229,17 @@ pub(crate) enum AppEvent {
|
||||
matches: Vec<FileMatch>,
|
||||
},
|
||||
|
||||
/// Upload a local client file into Codex-managed app-server storage.
|
||||
UploadLocalFile {
|
||||
path: PathBuf,
|
||||
},
|
||||
|
||||
/// Result of uploading a local client file into Codex-managed app-server storage.
|
||||
LocalFileUploaded {
|
||||
local_path: PathBuf,
|
||||
result: Result<PathBuf, String>,
|
||||
},
|
||||
|
||||
/// Refresh account rate limits in the background.
|
||||
RefreshRateLimits {
|
||||
origin: RateLimitRefreshOrigin,
|
||||
|
||||
@@ -10472,6 +10472,27 @@ impl ChatWidget {
|
||||
self.bottom_pane.insert_str(text);
|
||||
}
|
||||
|
||||
pub(crate) fn insert_uploaded_file_path(&mut self, path: &Path) {
|
||||
let path = path.to_string_lossy();
|
||||
let queued_message_index = self.queued_user_messages.iter().position(|queued_message| {
|
||||
self.queued_message_accepts_uploaded_file_path(queued_message)
|
||||
});
|
||||
if let Some(queued_message) =
|
||||
queued_message_index.and_then(|index| self.queued_user_messages.get_mut(index))
|
||||
{
|
||||
if !queued_message.user_message.text.is_empty() {
|
||||
queued_message.user_message.text.push(' ');
|
||||
}
|
||||
queued_message.user_message.text.push_str(&path);
|
||||
self.refresh_pending_input_preview();
|
||||
return;
|
||||
}
|
||||
if !self.bottom_pane.composer_text().is_empty() {
|
||||
self.bottom_pane.insert_str(" ");
|
||||
}
|
||||
self.bottom_pane.insert_str(&path);
|
||||
}
|
||||
|
||||
/// Replace the composer content with the provided text and reset cursor.
|
||||
pub(crate) fn set_composer_text(
|
||||
&mut self,
|
||||
|
||||
@@ -33,6 +33,7 @@ const SIDE_SLASH_COMMAND_UNAVAILABLE_HINT: &str = "Press Esc to return to the ma
|
||||
const GOAL_USAGE: &str = "Usage: /goal <objective>";
|
||||
const GOAL_USAGE_HINT: &str = "Example: /goal improve benchmark coverage";
|
||||
const RAW_USAGE: &str = "Usage: /raw [on|off]";
|
||||
const UPLOAD_USAGE: &str = "Usage: /upload <local-path>";
|
||||
|
||||
impl ChatWidget {
|
||||
/// Dispatch a bare slash command and record its staged local-history entry.
|
||||
@@ -393,6 +394,9 @@ impl ChatWidget {
|
||||
SlashCommand::Ps => {
|
||||
self.add_ps_output();
|
||||
}
|
||||
SlashCommand::Upload => {
|
||||
self.add_info_message(UPLOAD_USAGE.to_string(), /*hint*/ None);
|
||||
}
|
||||
SlashCommand::Stop => {
|
||||
self.clean_background_terminals();
|
||||
}
|
||||
@@ -776,6 +780,16 @@ impl ChatWidget {
|
||||
self.app_event_tx
|
||||
.send(AppEvent::BeginWindowsSandboxGrantReadRoot { path: args });
|
||||
}
|
||||
SlashCommand::Upload if !trimmed.is_empty() => {
|
||||
let path = PathBuf::from(trimmed);
|
||||
self.app_event_tx.send(AppEvent::UploadLocalFile {
|
||||
path: if path.is_absolute() {
|
||||
path
|
||||
} else {
|
||||
self.config.cwd.join(path).to_path_buf()
|
||||
},
|
||||
});
|
||||
}
|
||||
_ => self.dispatch_command(cmd),
|
||||
}
|
||||
if source == SlashCommandDispatchSource::Live && cmd != SlashCommand::Goal {
|
||||
@@ -826,6 +840,9 @@ impl ChatWidget {
|
||||
|
||||
if rest.is_empty() {
|
||||
self.dispatch_command(cmd);
|
||||
if cmd == SlashCommand::Upload {
|
||||
return QueueDrain::Continue;
|
||||
}
|
||||
return self.queued_command_drain_result(cmd);
|
||||
}
|
||||
|
||||
@@ -890,6 +907,35 @@ impl ChatWidget {
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn queued_message_accepts_uploaded_file_path(
|
||||
&self,
|
||||
queued_message: &QueuedUserMessage,
|
||||
) -> bool {
|
||||
match queued_message.action {
|
||||
QueuedInputAction::Plain => true,
|
||||
QueuedInputAction::RunShell => false,
|
||||
QueuedInputAction::ParseSlash => {
|
||||
let Some((name, rest, _rest_offset)) = parse_slash_name(&queued_message.text)
|
||||
else {
|
||||
return true;
|
||||
};
|
||||
if name.contains('/') {
|
||||
return true;
|
||||
}
|
||||
let Some(cmd) =
|
||||
slash_commands::find_builtin_command(name, self.builtin_command_flags())
|
||||
else {
|
||||
return false;
|
||||
};
|
||||
if rest.trim().is_empty() {
|
||||
return false;
|
||||
}
|
||||
matches!(cmd, SlashCommand::Plan | SlashCommand::Side)
|
||||
|| !cmd.supports_inline_args()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn queued_command_drain_result(&self, cmd: SlashCommand) -> QueueDrain {
|
||||
if self.is_user_turn_pending_or_running() || !self.bottom_pane.no_modal_or_popup_active() {
|
||||
return QueueDrain::Stop;
|
||||
@@ -938,6 +984,7 @@ impl ChatWidget {
|
||||
| SlashCommand::Experimental
|
||||
| SlashCommand::AutoReview
|
||||
| SlashCommand::Memories
|
||||
| SlashCommand::Upload
|
||||
| SlashCommand::Quit
|
||||
| SlashCommand::Exit
|
||||
| SlashCommand::Logout
|
||||
|
||||
@@ -1950,6 +1950,163 @@ async fn queued_fast_slash_applies_before_next_queued_message() {
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn queued_upload_stops_before_follow_up_prompt() {
|
||||
let (mut chat, mut rx, mut op_rx) = make_chatwidget_manual(/*model_override*/ None).await;
|
||||
chat.thread_id = Some(ThreadId::new());
|
||||
handle_turn_started(&mut chat, "turn-1");
|
||||
|
||||
queue_composer_text_with_tab(&mut chat, "/upload /tmp/demo.txt");
|
||||
queue_composer_text_with_tab(&mut chat, "use the upload");
|
||||
|
||||
complete_turn_with_message(&mut chat, "turn-1", Some("done"));
|
||||
|
||||
let events = std::iter::from_fn(|| rx.try_recv().ok()).collect::<Vec<_>>();
|
||||
assert!(
|
||||
events.iter().any(|event| matches!(
|
||||
event,
|
||||
AppEvent::UploadLocalFile { path } if path == &test_path_buf("/tmp/demo.txt")
|
||||
)),
|
||||
"expected queued /upload to start upload; events: {events:?}"
|
||||
);
|
||||
assert_eq!(chat.queued_user_messages.len(), 1);
|
||||
assert_eq!(
|
||||
chat.queued_user_messages.front().unwrap().text,
|
||||
"use the upload"
|
||||
);
|
||||
assert_matches!(op_rx.try_recv(), Err(TryRecvError::Empty));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn upload_resolves_relative_paths_against_session_cwd() {
|
||||
let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await;
|
||||
let tempdir = tempdir().unwrap();
|
||||
chat.config.cwd = tempdir.path().to_path_buf().abs();
|
||||
|
||||
submit_composer_text(&mut chat, "/upload demo.txt");
|
||||
|
||||
let events = std::iter::from_fn(|| rx.try_recv().ok()).collect::<Vec<_>>();
|
||||
assert!(
|
||||
events.iter().any(|event| matches!(
|
||||
event,
|
||||
AppEvent::UploadLocalFile { path }
|
||||
if path == &chat.config.cwd.join("demo.txt").to_path_buf()
|
||||
)),
|
||||
"expected relative upload to use the session cwd; events: {events:?}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn queued_upload_resumes_follow_up_prompt_after_upload_succeeds() {
|
||||
let (mut chat, mut rx, mut op_rx) = make_chatwidget_manual(/*model_override*/ None).await;
|
||||
chat.thread_id = Some(ThreadId::new());
|
||||
handle_turn_started(&mut chat, "turn-1");
|
||||
|
||||
queue_composer_text_with_tab(&mut chat, "/upload /tmp/demo.txt");
|
||||
queue_composer_text_with_tab(&mut chat, "use the upload");
|
||||
|
||||
complete_turn_with_message(&mut chat, "turn-1", Some("done"));
|
||||
let _events = std::iter::from_fn(|| rx.try_recv().ok()).collect::<Vec<_>>();
|
||||
|
||||
chat.insert_uploaded_file_path(Path::new("/tmp/.codex/uploads/demo.txt"));
|
||||
chat.maybe_send_next_queued_input();
|
||||
|
||||
match next_submit_op(&mut op_rx) {
|
||||
Op::UserTurn { items, .. } => assert_eq!(
|
||||
items,
|
||||
vec![UserInput::Text {
|
||||
text: "use the upload /tmp/.codex/uploads/demo.txt".to_string(),
|
||||
text_elements: Vec::new(),
|
||||
}]
|
||||
),
|
||||
other => panic!("expected queued follow-up prompt to submit, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn queued_upload_attaches_to_follow_up_plan_prompt() {
|
||||
let (mut chat, mut rx, mut op_rx) = make_chatwidget_manual(/*model_override*/ None).await;
|
||||
chat.thread_id = Some(ThreadId::new());
|
||||
chat.set_feature_enabled(Feature::CollaborationModes, /*enabled*/ true);
|
||||
handle_turn_started(&mut chat, "turn-1");
|
||||
|
||||
queue_composer_text_with_tab(&mut chat, "/upload /tmp/demo.txt");
|
||||
queue_composer_text_with_tab(&mut chat, "/plan use the upload");
|
||||
|
||||
complete_turn_with_message(&mut chat, "turn-1", Some("done"));
|
||||
let _events = std::iter::from_fn(|| rx.try_recv().ok()).collect::<Vec<_>>();
|
||||
|
||||
chat.insert_uploaded_file_path(Path::new("/tmp/.codex/uploads/demo.txt"));
|
||||
chat.maybe_send_next_queued_input();
|
||||
|
||||
match next_submit_op(&mut op_rx) {
|
||||
Op::UserTurn { items, .. } => assert_eq!(
|
||||
items,
|
||||
vec![UserInput::Text {
|
||||
text: "use the upload /tmp/.codex/uploads/demo.txt".to_string(),
|
||||
text_elements: Vec::new(),
|
||||
}]
|
||||
),
|
||||
other => panic!("expected queued /plan prompt to submit, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn queued_upload_skips_command_only_follow_ups_before_plain_prompt() {
|
||||
let (mut chat, mut rx, _op_rx) = make_chatwidget_manual(/*model_override*/ None).await;
|
||||
chat.thread_id = Some(ThreadId::new());
|
||||
handle_turn_started(&mut chat, "turn-1");
|
||||
|
||||
queue_composer_text_with_tab(&mut chat, "/upload /tmp/first.txt");
|
||||
queue_composer_text_with_tab(&mut chat, "/upload /tmp/second.txt");
|
||||
queue_composer_text_with_tab(&mut chat, "use both uploads");
|
||||
|
||||
complete_turn_with_message(&mut chat, "turn-1", Some("done"));
|
||||
let _events = std::iter::from_fn(|| rx.try_recv().ok()).collect::<Vec<_>>();
|
||||
|
||||
chat.insert_uploaded_file_path(Path::new("/tmp/.codex/uploads/first.txt"));
|
||||
|
||||
assert_eq!(chat.queued_user_messages.len(), 2);
|
||||
assert_eq!(
|
||||
chat.queued_user_messages.front().unwrap().text,
|
||||
"/upload /tmp/second.txt"
|
||||
);
|
||||
assert_eq!(
|
||||
chat.queued_user_messages.back().unwrap().text,
|
||||
"use both uploads /tmp/.codex/uploads/first.txt"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn queued_upload_without_args_keeps_draining() {
|
||||
let (mut chat, mut rx, mut op_rx) = make_chatwidget_manual(/*model_override*/ None).await;
|
||||
chat.thread_id = Some(ThreadId::new());
|
||||
handle_turn_started(&mut chat, "turn-1");
|
||||
|
||||
queue_composer_text_with_tab(&mut chat, "/upload");
|
||||
queue_composer_text_with_tab(&mut chat, "after upload help");
|
||||
|
||||
complete_turn_with_message(&mut chat, "turn-1", Some("done"));
|
||||
|
||||
let events = std::iter::from_fn(|| rx.try_recv().ok()).collect::<Vec<_>>();
|
||||
assert!(
|
||||
!events
|
||||
.iter()
|
||||
.any(|event| matches!(event, AppEvent::UploadLocalFile { .. })),
|
||||
"bare queued /upload should not start an upload; events: {events:?}"
|
||||
);
|
||||
match next_submit_op(&mut op_rx) {
|
||||
Op::UserTurn { items, .. } => assert_eq!(
|
||||
items,
|
||||
vec![UserInput::Text {
|
||||
text: "after upload help".to_string(),
|
||||
text_elements: Vec::new(),
|
||||
}]
|
||||
),
|
||||
other => panic!("expected follow-up prompt after bare /upload, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn user_turn_sends_standard_override_after_fast_is_turned_off() {
|
||||
let (mut chat, mut rx, mut op_rx) = make_chatwidget_manual(Some("gpt-5.3-codex")).await;
|
||||
|
||||
@@ -58,6 +58,7 @@ pub enum SlashCommand {
|
||||
Feedback,
|
||||
Rollout,
|
||||
Ps,
|
||||
Upload,
|
||||
#[strum(to_string = "stop", serialize = "clean")]
|
||||
Stop,
|
||||
Clear,
|
||||
@@ -100,6 +101,7 @@ impl SlashCommand {
|
||||
SlashCommand::Statusline => "configure which items appear in the status line",
|
||||
SlashCommand::Theme => "choose a syntax highlighting theme",
|
||||
SlashCommand::Ps => "list background terminals",
|
||||
SlashCommand::Upload => "upload a local file to the app-server host",
|
||||
SlashCommand::Stop => "stop all background terminals",
|
||||
SlashCommand::MemoryDrop => "DO NOT USE",
|
||||
SlashCommand::MemoryUpdate => "DO NOT USE",
|
||||
@@ -159,6 +161,7 @@ impl SlashCommand {
|
||||
| SlashCommand::Side
|
||||
| SlashCommand::Resume
|
||||
| SlashCommand::SandboxReadRoot
|
||||
| SlashCommand::Upload
|
||||
)
|
||||
}
|
||||
|
||||
@@ -209,6 +212,7 @@ impl SlashCommand {
|
||||
| SlashCommand::Status
|
||||
| SlashCommand::DebugConfig
|
||||
| SlashCommand::Ps
|
||||
| SlashCommand::Upload
|
||||
| SlashCommand::Stop
|
||||
| SlashCommand::Goal
|
||||
| SlashCommand::Mcp
|
||||
|
||||
Reference in New Issue
Block a user