Hare/yoi
1
1
Fork 0
Code Actions Packages Releases Activity

feat: capture bounded Bash output over replayable byte transport

Browse Source
This commit is contained in:
Hare
2026-10-20 03:01:22 +09:00
parent e5749b9be0
commit 4555aace28
46 changed files with 4183 additions and 1914 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/workdir/Cargo.toml
+1 -1
View File
@@ -20,7 +20,7 @@ serde_json.workspace = true
sha2.workspace = true
tempfile.workspace = true
thiserror.workspace = true
tokio = { workspace = true, features = ["process", "rt", "sync", "time"] }
tokio = { workspace = true, features = ["io-util", "process", "rt", "sync", "time"] }
[dev-dependencies]
serde_json.workspace = true
crates/workdir/src/command_test_support.rs
+43
View File
@@ -0,0 +1,43 @@
//! Bounded collection for tests whose subject is command routing/lifecycle,
//! rather than the byte-page protocol itself.
use crate::*;
#[derive(Debug)]
pub(crate) struct CollectedCommand {
pub content: String,
pub status: CommandStatus,
pub exit_code: Option<i32>,
pub timed_out: bool,
}
#[async_trait::async_trait]
pub(crate) trait CollectCommand: WorkdirSession {
async fn collect_command_output(
&self,
mut request: CommandOutputRequest,
) -> Result<CollectedCommand, WorkdirError> {
let mut bytes = Vec::new();
request.wait = true;
loop {
let page = self.command_output(request.clone()).await?;
for chunk in page.chunks {
bytes.extend(chunk.bytes);
assert!(
bytes.len() <= 16 * 1024 * 1024,
"fixture output exceeded bounded test collection"
);
}
request.cursor = page.next_cursor;
if let Some(terminal) = page.terminal {
return Ok(CollectedCommand {
content: String::from_utf8_lossy(&bytes).into_owned(),
status: terminal.status,
exit_code: terminal.exit_code,
timed_out: terminal.status == CommandStatus::TimedOut,
});
}
}
}
}
impl<T: WorkdirSession + ?Sized> CollectCommand for T {}
crates/workdir/src/external.rs
+136 -38
View File
@@ -9,7 +9,8 @@
//! COMMAND remains separate explicit authority. General session capability intersections may form
//! subsets independently of this grant hierarchy; the codec never implicitly expands nonhierarchical
//! grants. Commands run with the shared host's OS permissions, not a sandbox confining processes
//! to a directory. Output is bounded and inline; provider-local spill paths are not exposed.
//! to a directory. Output uses bounded byte pages with cursor acknowledgements;
//! providers never receive caller-selected output paths.
use std::path::{Component, Path};
@@ -18,7 +19,7 @@ use serde::{Deserialize, Serialize};
use crate::http::{WorkdirSessionOperation, WorkdirSessionOperationResult, WorkdirTransportError};
use crate::{BoundedReadLimits, WorkdirId, WorkdirPath, WorkdirSessionCapabilities};
pub const EXTERNAL_WORKDIR_PROVIDER_PROTOCOL_VERSION: u16 = 2;
pub const EXTERNAL_WORKDIR_PROVIDER_PROTOCOL_VERSION: u16 = 3;
/// Maximum serialized Backend-to-provider command frame size.
pub const MAX_EXTERNAL_SERVER_FRAME_BYTES: usize = 1024 * 1024;
/// Maximum serialized provider-to-Backend result frame size. Read bytes are
@@ -33,9 +34,9 @@ const MAX_EXTERNAL_GREP_CONTEXT: usize = 100;
const MAX_EXTERNAL_SCOPE_RULES: usize = 64;
pub const MAX_EXTERNAL_WRITE_BYTES: usize = 512 * 1024;
pub const MAX_EXTERNAL_COMMAND_BYTES: usize = 64 * 1024;
pub const MAX_EXTERNAL_COMMAND_OUTPUT_BYTES: usize = 512 * 1024;
pub const MAX_EXTERNAL_COMMAND_OUTPUT_BYTES: usize = crate::MAX_COMMAND_READ_BYTES;
pub const MAX_EXTERNAL_COMMAND_TIMEOUT_SECS: u64 = 600;
const MAX_EXTERNAL_COMMAND_CURSOR: usize = 1_000_000;
const MAX_EXTERNAL_COMMAND_CURSOR: u64 = crate::MAX_COMMAND_TRANSPORT_BYTES;
fn valid_command_handle(handle: &crate::CommandHandle) -> bool {
!handle.0.is_empty()
@@ -267,11 +268,10 @@ fn validate_external_operation(operation: &WorkdirSessionOperation) -> Result<()
}
WorkdirSessionOperation::CommandStart(request) => {
!request.command.is_empty()
&& request.validate().is_ok()
&& request.command.len() <= MAX_EXTERNAL_COMMAND_BYTES
&& (1..=MAX_EXTERNAL_COMMAND_TIMEOUT_SECS).contains(&request.timeout_secs)
&& (1..=MAX_EXTERNAL_COMMAND_OUTPUT_BYTES).contains(&request.output_limit)
&& is_root_relative(&request.cwd)
&& request.spill_dir.is_none()
&& request.tool_call_id.as_ref().is_none_or(|tool_call_id| {
!tool_call_id.is_empty()
&& tool_call_id.len() <= MAX_PROTOCOL_ID_BYTES
@@ -450,11 +450,19 @@ fn validate_external_result(result: &WorkdirSessionOperationResult) -> Result<()
WorkdirSessionOperationResult::CommandStart(handle) => valid_command_handle(handle),
WorkdirSessionOperationResult::CommandStatus(_) => true,
WorkdirSessionOperationResult::CommandOutput(output) => {
output.content.len() <= MAX_EXTERNAL_COMMAND_OUTPUT_BYTES
output.transport_bounds_valid()
&& output.chunks.len() <= MAX_EXTERNAL_COMMAND_OUTPUT_BYTES
&& output
.next_cursor
.is_none_or(|cursor| cursor <= MAX_EXTERNAL_COMMAND_CURSOR)
&& output.output_path.is_none()
.chunks
.iter()
.map(|chunk| chunk.bytes.len())
.sum::<usize>()
<= MAX_EXTERNAL_COMMAND_OUTPUT_BYTES
&& output.next_cursor <= MAX_EXTERNAL_COMMAND_CURSOR
&& output
.terminal
.as_ref()
.is_none_or(|terminal| terminal.status != crate::CommandStatus::Running)
}
WorkdirSessionOperationResult::CommandCancel => true,
};
@@ -801,11 +809,16 @@ mod tests {
#[test]
fn unsupported_protocol_versions_and_invalid_read_limits_fail_closed() {
let future = r#"{
"version": 3,
"message": {"kind": "heartbeat", "generation": 1, "sequence": 1}
}"#;
assert!(serde_json::from_str::<ExternalWorkdirServerFrame>(future).is_err());
for version in [
EXTERNAL_WORKDIR_PROVIDER_PROTOCOL_VERSION - 1,
EXTERNAL_WORKDIR_PROVIDER_PROTOCOL_VERSION + 1,
] {
let unsupported = serde_json::json!({
"version": version,
"message": {"kind": "heartbeat", "generation": 1, "sequence": 1}
});
assert!(serde_json::from_value::<ExternalWorkdirServerFrame>(unsupported).is_err());
}
let invalid_limits = r#"{
"max_source_bytes": 1024,
@@ -825,32 +838,35 @@ mod tests {
assert!(ExternalWorkdirOperation::try_from(absolute).is_err());
let command = WorkdirSessionOperation::CommandStart(CommandRequest {
command: "cat /home/operator/private/session.log".to_string(),
transport_byte_budget: crate::MAX_COMMAND_TRANSPORT_BYTES,
command: "printf unsafe".to_string(),
timeout_secs: 1,
output_limit: 1024,
cwd: WorkdirPath::root(),
spill_dir: Some("/home/operator/private/output".into()),
cwd: serde_json::from_value(serde_json::json!("/absolute")).unwrap(),
tool_call_id: None,
});
assert!(ExternalWorkdirOperation::try_from(command).is_err());
let command = WorkdirSessionOperation::CommandStart(CommandRequest {
transport_byte_budget: crate::MAX_COMMAND_TRANSPORT_BYTES,
command: "printf bounded".to_string(),
timeout_secs: 30,
output_limit: 12 * 1024,
cwd: WorkdirPath::root(),
spill_dir: None,
tool_call_id: Some("tool-call-1".to_string()),
});
assert!(ExternalWorkdirOperation::try_from(command).is_ok());
assert!(
ExternalWorkdirOperation::try_from(WorkdirSessionOperation::CommandStart(
CommandRequest {
transport_byte_budget: crate::MAX_COMMAND_TRANSPORT_BYTES,
command: "x".repeat(MAX_EXTERNAL_COMMAND_BYTES + 1),
timeout_secs: 30,
output_limit: 12 * 1024,
cwd: WorkdirPath::root(),
spill_dir: None,
tool_call_id: None,
}
))
@@ -887,13 +903,17 @@ mod tests {
assert!(
ExternalWorkdirOperationResult::try_from(WorkdirSessionOperationResult::CommandOutput(
crate::CommandOutput {
status: crate::CommandStatus::Completed,
exit_code: Some(0),
timed_out: false,
content: "bounded".to_string(),
next_cursor: None,
truncated: false,
output_path: None,
chunks: vec![],
next_cursor: 0,
terminal: Some(crate::CommandTerminal {
status: crate::CommandStatus::Completed,
exit_code: Some(0),
stdout_bytes: 0,
stderr_bytes: 0,
stdout_delivered_bytes: 0,
stderr_delivered_bytes: 0,
truncation_reason: None,
})
}
))
.is_ok()
@@ -901,13 +921,14 @@ mod tests {
assert!(
ExternalWorkdirOperationResult::try_from(WorkdirSessionOperationResult::CommandOutput(
crate::CommandOutput {
status: crate::CommandStatus::Completed,
exit_code: Some(0),
timed_out: false,
content: "bounded".to_string(),
next_cursor: None,
truncated: true,
output_path: Some("/provider/private/output".into()),
chunks: vec![crate::CommandByteChunk {
sequence: 0,
stream: crate::CommandStream::Stdout,
offset: 0,
bytes: vec![0; MAX_EXTERNAL_COMMAND_OUTPUT_BYTES + 1]
}],
next_cursor: 0,
terminal: None
}
))
.is_err()
@@ -958,7 +979,7 @@ mod tests {
#[test]
fn protocol_rejects_unknown_nested_fields_and_unbounded_requests() {
let unknown = r#"{
"version": 2,
"version": 3,
"message": {
"kind": "heartbeat",
"generation": 1,
@@ -1012,4 +1033,81 @@ mod tests {
assert!(ExternalWorkdirOperationId::new("x".repeat(129)).is_err());
assert!(ExternalWorkdirOperationId::new("operation-1").is_ok());
}
#[test]
fn command_budget_and_page_metadata_obey_external_hard_bounds() {
fn request(budget: u64) -> WorkdirSessionOperation {
WorkdirSessionOperation::CommandStart(CommandRequest {
command: "true".into(),
timeout_secs: 1,
cwd: WorkdirPath::root(),
transport_byte_budget: budget,
tool_call_id: None,
})
}
for budget in [0, crate::MAX_COMMAND_TRANSPORT_BYTES] {
let operation = ExternalWorkdirOperation::try_from(request(budget)).unwrap();
let wire = serde_json::to_value(operation).unwrap();
assert!(serde_json::from_value::<ExternalWorkdirOperation>(wire).is_ok());
}
assert!(
ExternalWorkdirOperation::try_from(request(crate::MAX_COMMAND_TRANSPORT_BYTES + 1))
.is_err()
);
fn page() -> crate::CommandOutput {
crate::CommandOutput {
chunks: vec![crate::CommandByteChunk {
sequence: 0,
stream: crate::CommandStream::Stdout,
offset: 0,
bytes: b"abc".to_vec(),
}],
next_cursor: 3,
terminal: Some(crate::CommandTerminal {
status: crate::CommandStatus::Completed,
exit_code: Some(0),
stdout_bytes: 100 * 1024 * 1024,
stderr_bytes: 7,
stdout_delivered_bytes: 3,
stderr_delivered_bytes: 0,
truncation_reason: Some(crate::CommandOutputTruncation::ByteBudget),
}),
}
}
let accepted = ExternalWorkdirOperationResult::try_from(
WorkdirSessionOperationResult::CommandOutput(page()),
)
.unwrap();
assert!(
serde_json::from_value::<ExternalWorkdirOperationResult>(
serde_json::to_value(accepted).unwrap()
)
.is_ok()
);
for mutate in [
|p: &mut crate::CommandOutput| p.next_cursor = crate::MAX_COMMAND_TRANSPORT_BYTES + 1,
|p: &mut crate::CommandOutput| p.chunks[0].offset = crate::MAX_COMMAND_TRANSPORT_BYTES,
|p: &mut crate::CommandOutput| p.chunks[0].sequence = u64::MAX,
|p: &mut crate::CommandOutput| p.terminal.as_mut().unwrap().stdout_delivered_bytes = 4,
|p: &mut crate::CommandOutput| p.terminal.as_mut().unwrap().stdout_bytes = 2,
|p: &mut crate::CommandOutput| p.terminal.as_mut().unwrap().truncation_reason = None,
] {
let mut invalid = page();
mutate(&mut invalid);
let wire = serde_json::to_value(WorkdirSessionOperationResult::CommandOutput(invalid))
.unwrap();
assert!(serde_json::from_value::<ExternalWorkdirOperationResult>(wire).is_err());
}
assert!(
ExternalWorkdirOperation::try_from(WorkdirSessionOperation::CommandOutput(
crate::CommandOutputRequest {
handle: crate::CommandHandle("command-1".into()),
cursor: crate::MAX_COMMAND_TRANSPORT_BYTES + 1,
limit: 1,
wait: false,
}
))
.is_err()
);
}
}
crates/workdir/src/external_local_tests.rs
+13 -7
View File
@@ -1,6 +1,9 @@
#[cfg(test)]
use crate::command_test_support::CollectCommand as _;
use crate::http::{
WorkdirSessionOperation, WorkdirSessionOperationResult, dispatch_workdir_session_operation,
};
use crate::{
BoundedReadLimits, CommandOutputRequest, CommandRequest, EditRequest, ExternalWorkdirRoot,
GlobRequest, GrepOutputMode, GrepRequest, LocalWorkdirSession, ReadRequest, Workdir,
@@ -59,11 +62,12 @@ async fn external_local_provider_is_strictly_read_only() {
let command = WorkdirSession::start_command(
&session,
CommandRequest {
transport_byte_budget: crate::MAX_COMMAND_TRANSPORT_BYTES,
command: "echo changed > item.txt".to_string(),
timeout_secs: 1,
output_limit: 1024,
cwd: WorkdirPath::root(),
spill_dir: None,
tool_call_id: None,
},
)
@@ -145,11 +149,12 @@ async fn external_local_provider_read_write_is_filesystem_only() {
let command = WorkdirSession::start_command(
&session,
CommandRequest {
transport_byte_budget: crate::MAX_COMMAND_TRANSPORT_BYTES,
command: "touch command-ran".to_string(),
timeout_secs: 1,
output_limit: 1024,
cwd: WorkdirPath::root(),
spill_dir: None,
tool_call_id: None,
},
)
@@ -193,17 +198,18 @@ async fn external_command_permission_is_explicit_and_independent_from_file_acces
let handle = session
.start_command(CommandRequest {
transport_byte_budget: crate::MAX_COMMAND_TRANSPORT_BYTES,
command: "printf command-ok".to_string(),
timeout_secs: 5,
output_limit: 1024,
cwd: WorkdirPath::root(),
spill_dir: None,
tool_call_id: None,
})
.await
.unwrap();
let output = session
.command_output(CommandOutputRequest {
.collect_command_output(CommandOutputRequest {
handle,
cursor: 0,
limit: 1024,
crates/workdir/src/http.rs
+67 -17
View File
@@ -822,6 +822,7 @@ mod client {
&self,
request: CommandRequest,
) -> Result<CommandHandle, WorkdirError> {
request.validate()?;
match self
.operate(WorkdirSessionOperation::CommandStart(request))
.await?
@@ -848,17 +849,39 @@ mod client {
&self,
request: CommandOutputRequest,
) -> Result<CommandOutput, WorkdirError> {
if !(1..=crate::MAX_COMMAND_READ_BYTES).contains(&request.limit) {
return Err(WorkdirError::InvalidArgument(
"command byte limit is outside transport bounds".into(),
));
}
let wait = request.wait;
loop {
match self
.operate(WorkdirSessionOperation::CommandOutput(request.clone()))
.await?
{
WorkdirSessionOperationResult::CommandOutput(result)
if wait && result.status == CommandStatus::Running => {}
WorkdirSessionOperationResult::CommandOutput(result) => return Ok(result),
_ => return Err(Self::mismatch("command_output")),
match self
.operate(WorkdirSessionOperation::CommandOutput(request.clone()))
.await?
{
WorkdirSessionOperationResult::CommandOutput(result) => {
let size = result
.chunks
.iter()
.try_fold(0usize, |size, chunk| size.checked_add(chunk.bytes.len()));
if !result.transport_bounds_valid()
|| size.is_none_or(|size| {
size > request.limit
|| request.cursor.checked_add(size as u64)
!= Some(result.next_cursor)
})
{
return Err(Self::mismatch("bounded command byte page"));
}
if wait && result.chunks.is_empty() && result.terminal.is_none() {
// A waiting request must block at the provider. Never spin HTTP
// requests against an immediate empty response. A caller may retry
// this exact cursor with its own bounded transport retry policy.
return Err(Self::mismatch("waiting command byte page"));
}
return Ok(result);
}
_ => return Err(Self::mismatch("command_output")),
}
}
@@ -917,20 +940,47 @@ mod client {
response: reqwest::Response,
) -> Result<T, WorkdirError> {
if response.status().is_success() {
response.json().await.map_err(http_unavailable)
let bytes = bounded_response_bytes(response).await?;
serde_json::from_slice(&bytes)
.map_err(|_| WorkdirError::Transport("Runtime Workdir response was invalid".into()))
} else {
Err(decode_error(response).await)
}
}
async fn bounded_response_bytes(
mut response: reqwest::Response,
) -> Result<Vec<u8>, WorkdirError> {
const MAX_RESPONSE_BYTES: usize = crate::external::MAX_EXTERNAL_PROVIDER_FRAME_BYTES;
if response
.content_length()
.is_some_and(|size| size > MAX_RESPONSE_BYTES as u64)
{
return Err(WorkdirError::Transport(
"Runtime Workdir response exceeded transport bound".into(),
));
}
let mut bytes = Vec::new();
while let Some(chunk) = response.chunk().await.map_err(http_unavailable)? {
if chunk.len() > MAX_RESPONSE_BYTES.saturating_sub(bytes.len()) {
return Err(WorkdirError::Transport(
"Runtime Workdir response exceeded transport bound".into(),
));
}
bytes.extend_from_slice(&chunk);
}
Ok(bytes)
}
async fn decode_error(response: reqwest::Response) -> WorkdirError {
response
.json::<WorkdirTransportError>()
.await
.map(WorkdirTransportError::into_workdir_error)
.unwrap_or_else(|_| {
WorkdirError::Transport("Runtime Workdir error response was invalid".to_string())
})
match bounded_response_bytes(response).await {
Ok(bytes) => serde_json::from_slice::<WorkdirTransportError>(&bytes)
.map(WorkdirTransportError::into_workdir_error)
.unwrap_or_else(|_| {
WorkdirError::Transport("Runtime Workdir error response was invalid".into())
}),
Err(error) => error,
}
}
fn http_unavailable(_error: reqwest::Error) -> WorkdirError {
crates/workdir/src/lib.rs
+3
View File
@@ -618,3 +618,6 @@ mod capability_tests {
assert_eq!(bits(WorkdirSessionCapabilities::ALL), 63);
}
}
#[cfg(test)]
mod command_test_support;
crates/workdir/src/local.rs
+100 -875
View File
File diff suppressed because it is too large Load Diff
crates/workdir/src/local/command_transport.rs
+377
View File
@@ -0,0 +1,377 @@
//! Captured process bytes use a bounded queue, never the lossy Console channel.
//! Pipe backpressure bounds unread process output without temporary-file growth.
//! A replay page survives a lost HTTP response until its next cursor is sent.
use super::*;
use tokio::io::{AsyncRead, AsyncReadExt};
use tokio::sync::mpsc;
const QUEUE_CHUNKS: usize = 8;
pub(super) const MAX_COMMANDS: usize = 64;
#[derive(Debug)]
enum Message {
Bytes(crate::CommandByteChunk),
}
#[derive(Debug)]
pub(super) struct LocalCommand {
pub task: JoinHandle<()>,
pub cancel: watch::Sender<bool>,
pub status: watch::Receiver<Option<crate::CommandTerminal>>,
pub transport: Arc<Mutex<Transport>>,
}
#[derive(Debug)]
pub(super) struct Transport {
receiver: mpsc::Receiver<Message>,
pending: Option<crate::CommandByteChunk>,
cursor: u64,
replay: Option<(u64, CommandOutput)>,
terminal: Option<crate::CommandTerminal>,
terminal_status: watch::Receiver<Option<crate::CommandTerminal>>,
}
impl Transport {
pub async fn read(
&mut self,
request: &CommandOutputRequest,
) -> Result<CommandOutput, WorkdirError> {
if !(1..=crate::MAX_COMMAND_READ_BYTES).contains(&request.limit) {
return Err(WorkdirError::Conflict(
"command byte limit is outside transport bounds".into(),
));
}
if let Some((cursor, page)) = &self.replay {
if request.cursor == *cursor {
if page
.chunks
.iter()
.map(|chunk| chunk.bytes.len())
.sum::<usize>()
> request.limit
{
return Err(WorkdirError::Conflict(
"retry byte limit is smaller than retained page".into(),
));
}
return Ok(page.clone());
}
if request.cursor != page.next_cursor {
return Err(WorkdirError::Conflict(
"command cursor is not the current page acknowledgement".into(),
));
}
self.cursor = page.next_cursor;
self.replay = None;
}
if request.cursor != self.cursor {
return Err(WorkdirError::Conflict(
"command cursor was acknowledged or is ahead of output".into(),
));
}
let mut page = CommandOutput {
chunks: Vec::new(),
next_cursor: self.cursor,
terminal: self.terminal.clone(),
};
let mut remaining = request.limit;
loop {
let message = if let Some(chunk) = self.pending.take() {
Some(Message::Bytes(chunk))
} else if page.chunks.is_empty() && page.terminal.is_none() && request.wait {
self.receiver.recv().await
} else {
match self.receiver.try_recv() {
Ok(message) => Some(message),
Err(mpsc::error::TryRecvError::Empty) => break,
Err(mpsc::error::TryRecvError::Disconnected) => None,
}
};
match message {
Some(Message::Bytes(mut chunk)) => {
let count = remaining.min(chunk.bytes.len());
if count < chunk.bytes.len() {
let tail = chunk.bytes.split_off(count);
self.pending = Some(crate::CommandByteChunk {
sequence: chunk.sequence + count as u64,
stream: chunk.stream,
offset: chunk.offset + count as u64,
bytes: tail,
});
}
page.next_cursor += count as u64;
remaining -= count;
page.chunks.push(chunk);
if remaining == 0 {
break;
}
}
None if page.terminal.is_some() => break,
None => {
let terminal = self.terminal_status.borrow().clone().ok_or_else(|| {
WorkdirError::Unavailable(
"command transport ended before terminal fence".into(),
)
})?;
self.terminal = Some(terminal.clone());
page.terminal = Some(terminal);
break;
}
}
}
// An empty poll must not become a permanently replayed empty page.
if !page.chunks.is_empty() || page.terminal.is_some() {
self.replay = Some((self.cursor, page.clone()));
}
Ok(page)
}
pub fn terminal_delivered(&self) -> bool {
self.replay
.as_ref()
.is_some_and(|(_, page)| page.terminal.is_some())
}
pub fn close(&mut self) {
self.receiver.close();
}
}
pub(super) fn start(
cwd: PathBuf,
request: CommandRequest,
command_id: String,
telemetry: CommandTelemetry,
environment: BTreeMap<String, String>,
resources: Vec<Arc<dyn WorkdirSessionResource>>,
) -> LocalCommand {
let (sender, receiver) = mpsc::channel(QUEUE_CHUNKS);
let (cancel, cancel_rx) = watch::channel(false);
let (status_tx, status) = watch::channel(None);
let transport = Arc::new(Mutex::new(Transport {
receiver,
pending: None,
cursor: 0,
replay: None,
terminal: None,
terminal_status: status.clone(),
}));
let task = tokio::spawn(async move {
let _resources = resources;
let terminal = execute(
cwd,
request,
&command_id,
&telemetry,
environment,
cancel_rx,
&sender,
)
.await;
telemetry.terminal(&command_id, terminal.status, terminal.exit_code);
// Closing the byte queue fences this already-published terminal, without
// requiring spare queue capacity or a receiver to release task resources.
let _ = status_tx.send(Some(terminal));
});
LocalCommand {
task,
cancel,
status,
transport,
}
}
#[derive(Default)]
struct ByteAccounting {
stdout: u64,
stderr: u64,
stdout_delivered: u64,
stderr_delivered: u64,
}
impl ByteAccounting {
fn terminal(&self, status: CommandStatus, exit_code: Option<i32>) -> crate::CommandTerminal {
crate::CommandTerminal {
status,
exit_code,
stdout_bytes: self.stdout,
stderr_bytes: self.stderr,
stdout_delivered_bytes: self.stdout_delivered,
stderr_delivered_bytes: self.stderr_delivered,
truncation_reason: (self.stdout > self.stdout_delivered
|| self.stderr > self.stderr_delivered)
.then_some(crate::CommandOutputTruncation::ByteBudget),
}
}
}
async fn pump<R: AsyncRead + Unpin>(
mut stdout: R,
mut stderr: impl AsyncRead + Unpin,
sender: &mpsc::Sender<Message>,
telemetry: &CommandTelemetry,
command_id: &str,
mut remaining: u64,
accounting: &mut ByteAccounting,
) -> std::io::Result<()> {
let mut stdout_done = false;
let mut stderr_done = false;
let mut stdout_decoder = CommandOutputDecoder::default();
let mut stderr_decoder = CommandOutputDecoder::default();
let mut sequence = 0;
let mut stdout_buffer = [0; COMMAND_EVENT_CHUNK_BYTES];
let mut stderr_buffer = [0; COMMAND_EVENT_CHUNK_BYTES];
while !stdout_done || !stderr_done {
// After quota, drain locally even if nobody acknowledges the queue.
// Before quota, reserve first to retain bounded pipe backpressure.
let permit = if remaining > 0 {
let Ok(permit) = sender.reserve().await else {
return Ok(());
};
Some(permit)
} else {
None
};
let (stream, result) = tokio::select! {
result = stdout.read(&mut stdout_buffer), if !stdout_done => (CommandStream::Stdout, result),
result = stderr.read(&mut stderr_buffer), if !stderr_done => (CommandStream::Stderr, result),
};
let buffer = match stream {
CommandStream::Stdout => &stdout_buffer,
CommandStream::Stderr => &stderr_buffer,
};
let count = result?;
let decoder = match stream {
CommandStream::Stdout => &mut stdout_decoder,
CommandStream::Stderr => &mut stderr_decoder,
};
if count == 0 {
match stream {
CommandStream::Stdout => stdout_done = true,
CommandStream::Stderr => stderr_done = true,
}
publish_decoded_output(decoder, telemetry, command_id, stream, true);
continue;
}
let (seen, delivered) = match stream {
CommandStream::Stdout => (&mut accounting.stdout, &mut accounting.stdout_delivered),
CommandStream::Stderr => (&mut accounting.stderr, &mut accounting.stderr_delivered),
};
*seen = seen.saturating_add(count as u64);
let captured = remaining.min(count as u64) as usize;
if captured > 0 {
let offset = *delivered;
decoder.read_offset += captured as u64;
decoder.pending.extend_from_slice(&buffer[..captured]);
publish_decoded_output(decoder, telemetry, command_id, stream, false);
permit
.expect("reserved while budget remains")
.send(Message::Bytes(crate::CommandByteChunk {
sequence,
stream,
offset,
bytes: buffer[..captured].to_vec(),
}));
*delivered += captured as u64;
sequence += captured as u64;
remaining -= captured as u64;
}
}
Ok(())
}
async fn execute(
cwd: PathBuf,
request: CommandRequest,
command_id: &str,
telemetry: &CommandTelemetry,
environment: BTreeMap<String, String>,
mut cancel: watch::Receiver<bool>,
sender: &mpsc::Sender<Message>,
) -> crate::CommandTerminal {
telemetry.started(command_id, request.tool_call_id);
let mut command = Command::new("bash");
command
.arg("-c")
.arg(request.command)
.current_dir(cwd)
.envs(environment)
.stdin(Stdio::null())
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.kill_on_drop(true);
#[cfg(unix)]
command.as_std_mut().process_group(0);
let mut accounting = ByteAccounting::default();
let Ok(mut child) = command.spawn() else {
return accounting.terminal(CommandStatus::Failed, None);
};
let mut process_group = CommandProcessGroupGuard::new(child.id());
let stdout = child.stdout.take().expect("piped stdout");
let stderr = child.stderr.take().expect("piped stderr");
let (status, exit_code) = {
let output = pump(
stdout,
stderr,
sender,
telemetry,
command_id,
request.transport_byte_budget,
&mut accounting,
);
tokio::pin!(output);
let mut drained = false;
let timeout = tokio::time::sleep(Duration::from_secs(request.timeout_secs.max(1)));
tokio::pin!(timeout);
let terminal = loop {
tokio::select! {
result = &mut output, if !drained => {
drained = true;
if result.is_err() || sender.is_closed() {
process_group.terminate();
let _ = child.start_kill();
let _ = child.wait().await;
break (CommandStatus::Failed, None);
}
}
exit = child.wait() => {
process_group.terminate();
break match exit {
Ok(exit) => (if exit.success() { CommandStatus::Completed } else { CommandStatus::Failed }, exit.code()),
Err(_) => (CommandStatus::Failed, None),
};
}
_ = &mut timeout => {
process_group.terminate();
let _ = child.start_kill();
let code = child.wait().await.ok().and_then(|exit| exit.code());
break (CommandStatus::TimedOut, code);
}
_ = sender.closed() => {
process_group.terminate();
let _ = child.start_kill();
let _ = child.wait().await;
break (CommandStatus::Cancelled, None);
}
changed = cancel.changed() => {
if changed.is_err() || *cancel.borrow() {
process_group.terminate();
let _ = child.start_kill();
let code = child.wait().await.ok().and_then(|exit| exit.code());
break (CommandStatus::Cancelled, code);
}
}
}
};
// Even after kill, preserve the bounded pipe tail before publishing terminal.
if !drained {
let _ = output.await;
}
terminal
};
accounting.terminal(status, exit_code)
}
#[cfg(test)]
mod tests;
crates/workdir/src/local/command_transport/tests.rs
+490
View File
@@ -0,0 +1,490 @@
//! Process/pipe integration through the production local provider. Temporary roots
//! and provider tasks are RAII-owned; byte pulls and bounded timeouts synchronize tests.
use super::*;
use tempfile::TempDir;
struct Harness {
session: LocalWorkdirSession,
_root: TempDir,
}
impl Harness {
fn new() -> Self {
let root = TempDir::new().unwrap();
let session = LocalWorkdirSession::new(
Scope::writable(root.path()).unwrap(),
root.path().to_path_buf(),
);
Self {
session,
_root: root,
}
}
async fn start(&self, command: &str, timeout_secs: u64) -> CommandHandle {
self.start_budget(command, timeout_secs, crate::MAX_COMMAND_TRANSPORT_BYTES)
.await
}
fn request(command: &str, timeout_secs: u64, budget: u64) -> CommandRequest {
CommandRequest {
transport_byte_budget: budget,
command: command.into(),
timeout_secs,
cwd: WorkdirPath::root(),
tool_call_id: None,
}
}
async fn start_budget(&self, command: &str, timeout_secs: u64, budget: u64) -> CommandHandle {
self.session
.start_command(Self::request(command, timeout_secs, budget))
.await
.unwrap()
}
async fn page(&self, handle: &CommandHandle, cursor: u64, limit: usize) -> CommandOutput {
tokio::time::timeout(
Duration::from_secs(10),
self.session.command_output(CommandOutputRequest {
handle: handle.clone(),
cursor,
limit,
wait: true,
}),
)
.await
.expect("bounded command pull")
.unwrap()
}
async fn collect(
&self,
handle: &CommandHandle,
mut cursor: u64,
) -> (Vec<crate::CommandByteChunk>, crate::CommandTerminal) {
let mut chunks = Vec::new();
loop {
let page = self
.page(handle, cursor, crate::MAX_COMMAND_READ_BYTES)
.await;
assert!(
page.chunks
.iter()
.map(|chunk| chunk.bytes.len())
.sum::<usize>()
<= crate::MAX_COMMAND_READ_BYTES
);
cursor = page.next_cursor;
chunks.extend(page.chunks);
assert!(chunks.iter().map(|chunk| chunk.bytes.len()).sum::<usize>() <= 8 * 1024 * 1024);
if let Some(terminal) = page.terminal {
return (chunks, terminal);
}
}
}
}
fn stream_bytes(chunks: &[crate::CommandByteChunk], stream: CommandStream) -> Vec<u8> {
let mut bytes = Vec::new();
for chunk in chunks.iter().filter(|chunk| chunk.stream == stream) {
assert_eq!(chunk.offset, bytes.len() as u64);
bytes.extend_from_slice(&chunk.bytes);
}
bytes
}
#[tokio::test]
async fn high_single_line_output_survives_absent_and_lagged_console_subscribers() {
for lagged in [false, true] {
let harness = Harness::new();
let _lagged_receiver = lagged.then(|| harness.session.subscribe_command_events().unwrap());
let handle = harness.start("head -c 3145728 /dev/zero | tr '\\0' x; head -c 65536 /dev/zero | tr '\\0' e >&2; exit 7", 20).await;
let (chunks, terminal) = harness.collect(&handle, 0).await;
assert_eq!(terminal.status, CommandStatus::Failed);
assert_eq!(terminal.exit_code, Some(7));
assert_eq!(
stream_bytes(&chunks, CommandStream::Stdout),
vec![b'x'; 3145728]
);
assert_eq!(
stream_bytes(&chunks, CommandStream::Stderr),
vec![b'e'; 65536]
);
let mut cursor = 0;
for chunk in chunks {
assert_eq!(chunk.sequence, cursor);
cursor += chunk.bytes.len() as u64;
}
harness.session.close().await.unwrap();
}
}
#[tokio::test]
async fn one_byte_pages_preserve_raw_utf8_invalid_bytes_and_replay_until_acknowledged() {
let harness = Harness::new();
let handle = harness
.start(
"printf '\\342\\202\\254\\377\\000'; printf '\\346\\227\\245' >&2",
5,
)
.await;
let mut chunks = Vec::new();
let mut cursor = 0;
loop {
let page = harness.page(&handle, cursor, 1).await;
assert_eq!(harness.page(&handle, cursor, 1).await, page);
cursor = page.next_cursor;
chunks.extend(page.chunks);
if let Some(terminal) = page.terminal {
assert_eq!(terminal.status, CommandStatus::Completed);
break;
}
}
assert_eq!(
stream_bytes(&chunks, CommandStream::Stdout),
b"\xe2\x82\xac\xff\0"
);
assert_eq!(
stream_bytes(&chunks, CommandStream::Stderr),
b"\xe6\x97\xa5"
);
assert!(
harness
.session
.command_output(CommandOutputRequest {
handle,
cursor: 0,
limit: 1,
wait: false
})
.await
.is_err()
);
}
#[tokio::test]
async fn cancellation_while_backpressured_drains_bytes_before_terminal_fence() {
let harness = Harness::new();
let handle = harness.start("printf ready; exec yes x", 30).await;
let first = harness.page(&handle, 0, 1).await;
harness
.session
.cancel_command(handle.clone())
.await
.unwrap();
let (mut chunks, terminal) = harness.collect(&handle, first.next_cursor).await;
let mut all = first.chunks;
all.append(&mut chunks);
assert!(stream_bytes(&all, CommandStream::Stdout).starts_with(b"ready"));
assert_eq!(terminal.status, CommandStatus::Cancelled);
let final_page = harness
.page(
&handle,
all.iter().map(|chunk| chunk.bytes.len() as u64).sum(),
1,
)
.await;
assert!(final_page.chunks.is_empty());
assert_eq!(final_page.terminal, Some(terminal));
}
#[tokio::test]
async fn timeout_and_session_close_finalize_without_an_output_subscriber() {
let harness = Harness::new();
let handle = harness
.start("printf ready; exec tail -f /dev/null", 1)
.await;
let (chunks, terminal) = harness.collect(&handle, 0).await;
assert_eq!(stream_bytes(&chunks, CommandStream::Stdout), b"ready");
assert_eq!(terminal.status, CommandStatus::TimedOut);
let handle = harness.start("exec yes x", 30).await;
// Receiving one byte establishes actual process output, without a timing sleep.
harness.page(&handle, 0, 1).await;
tokio::time::timeout(Duration::from_secs(5), harness.session.close())
.await
.unwrap()
.unwrap();
assert!(harness.session.command_status(handle).await.is_err());
}
#[tokio::test]
async fn interleaved_streams_keep_independent_offsets_and_bounded_read_validation() {
let harness = Harness::new();
let handle = harness
.start(
"for ((i=0;i<100;i++)); do printf 'out%s;' \"$i\"; printf 'err%s;' \"$i\" >&2; done",
5,
)
.await;
for limit in [0, crate::MAX_COMMAND_READ_BYTES + 1] {
assert!(
harness
.session
.command_output(CommandOutputRequest {
handle: handle.clone(),
cursor: 0,
limit,
wait: false
})
.await
.is_err()
);
}
let (chunks, terminal) = harness.collect(&handle, 0).await;
let expected = |prefix| {
(0..100)
.map(|i| format!("{prefix}{i};"))
.collect::<String>()
.into_bytes()
};
assert_eq!(
stream_bytes(&chunks, CommandStream::Stdout),
expected("out")
);
assert_eq!(
stream_bytes(&chunks, CommandStream::Stderr),
expected("err")
);
assert_eq!(terminal.status, CommandStatus::Completed);
}
#[tokio::test]
async fn byte_budget_zero_exact_partial_and_shared_keep_seen_and_delivered_distinct() {
let harness = Harness::new();
for (command, budget, stdout_seen, stderr_seen, exit_code) in [
("printf abc; printf defg >&2", 0, 3, 4, 0),
("printf abc", 3, 3, 0, 0),
("printf abcdefg", 3, 7, 0, 0),
("printf abcdefg; exit 7", 3, 7, 0, 7),
(
"head -c 32768 /dev/zero; head -c 32768 /dev/zero >&2",
12345,
32768,
32768,
0,
),
] {
let handle = harness.start_budget(command, 5, budget).await;
let (chunks, terminal) = harness.collect(&handle, 0).await;
let stdout = stream_bytes(&chunks, CommandStream::Stdout);
let stderr = stream_bytes(&chunks, CommandStream::Stderr);
let mut cursor = 0;
for chunk in &chunks {
assert_eq!(chunk.sequence, cursor);
cursor += chunk.bytes.len() as u64;
}
assert_eq!(cursor, budget.min(stdout_seen + stderr_seen));
assert_eq!(terminal.stdout_bytes, stdout_seen);
assert_eq!(terminal.stderr_bytes, stderr_seen);
assert_eq!(terminal.stdout_delivered_bytes, stdout.len() as u64);
assert_eq!(terminal.stderr_delivered_bytes, stderr.len() as u64);
assert_eq!(terminal.exit_code, Some(exit_code));
assert_eq!(
terminal.status,
if exit_code == 0 {
CommandStatus::Completed
} else {
CommandStatus::Failed
}
);
assert_eq!(
terminal.truncation_reason,
(cursor < stdout_seen + stderr_seen)
.then_some(crate::CommandOutputTruncation::ByteBudget)
);
}
}
#[tokio::test]
async fn exhausted_budget_finishes_without_receiver_and_keeps_timeout_cancel_distinct() {
let harness = Harness::new();
for budget in [0, 1, (QUEUE_CHUNKS * COMMAND_EVENT_CHUNK_BYTES) as u64] {
let handle = harness
.start_budget("head -c 1048576 /dev/zero; printf error >&2", 5, budget)
.await;
let mut status = harness.session.inner.commands.lock().await[&handle.0]
.status
.clone();
let terminal = tokio::time::timeout(
Duration::from_secs(5),
status.wait_for(|value| value.is_some()),
)
.await
.unwrap()
.unwrap()
.clone()
.unwrap();
assert_eq!(terminal.status, CommandStatus::Completed);
assert_eq!(terminal.stdout_bytes, 1048576);
assert_eq!(terminal.stderr_bytes, 5);
assert_eq!(
terminal.stdout_delivered_bytes + terminal.stderr_delivered_bytes,
budget
);
// Full capture queue cannot retain the process task or its resources.
{
let mut commands = harness.session.inner.commands.lock().await;
tokio::time::timeout(
Duration::from_secs(5),
&mut commands.get_mut(&handle.0).unwrap().task,
)
.await
.unwrap()
.unwrap();
}
// Only now acknowledge the captured queue and terminal fence.
let (_, delivered_terminal) = harness.collect(&handle, 0).await;
assert_eq!(delivered_terminal, terminal);
}
for (cancel, expected) in [
(false, CommandStatus::TimedOut),
(true, CommandStatus::Cancelled),
] {
let handle = harness
.start_budget(
"printf ready-extra; head -c 32768 /dev/zero; exec tail -f /dev/null",
1,
5,
)
.await;
let first = harness.page(&handle, 0, 5).await;
assert_eq!(first.next_cursor, 5);
if cancel {
harness
.session
.cancel_command(handle.clone())
.await
.unwrap();
}
let (rest, terminal) = harness.collect(&handle, first.next_cursor).await;
assert!(rest.is_empty());
assert_eq!(terminal.status, expected);
assert_eq!(terminal.stdout_delivered_bytes, 5);
assert!(terminal.stdout_bytes >= 11);
assert_eq!(
terminal.truncation_reason,
Some(crate::CommandOutputTruncation::ByteBudget)
);
if !cancel {
assert_eq!(terminal.stdout_bytes, 32779);
assert_eq!(
terminal.truncation_reason,
Some(crate::CommandOutputTruncation::ByteBudget)
);
}
}
}
#[tokio::test]
async fn budget_replay_preserves_unique_bytes_and_rejects_future_and_old_cursors() {
let harness = Harness::new();
let handle = harness.start_budget("printf abcdefghij", 5, 7).await;
let first = harness.page(&handle, 0, 2).await;
for _ in 0..3 {
assert_eq!(harness.page(&handle, 0, 2).await, first);
}
assert!(
harness
.session
.command_output(CommandOutputRequest {
handle: handle.clone(),
cursor: 8,
limit: 2,
wait: false,
})
.await
.is_err()
);
let (rest, terminal) = harness.collect(&handle, first.next_cursor).await;
let mut all = first.chunks;
all.extend(rest);
assert_eq!(stream_bytes(&all, CommandStream::Stdout), b"abcdefg");
assert_eq!(terminal.stdout_bytes, 10);
assert_eq!(terminal.stdout_delivered_bytes, 7);
assert!(
harness
.session
.command_output(CommandOutputRequest {
handle,
cursor: 0,
limit: 2,
wait: false,
})
.await
.is_err()
);
}
#[tokio::test]
async fn invalid_budget_is_rejected_before_process_start() {
let harness = Harness::new();
assert!(matches!(
harness
.session
.start_command(Harness::request(
"touch must-not-exist",
5,
crate::MAX_COMMAND_TRANSPORT_BYTES + 1,
))
.await,
Err(WorkdirError::InvalidArgument(_))
));
assert!(!harness._root.path().join("must-not-exist").exists());
assert!(harness.session.inner.commands.lock().await.is_empty());
}
#[tokio::test]
async fn slow_receiver_backpressures_before_quota_but_full_queue_does_not_block_discard() {
use std::future::Future;
use std::task::Poll;
let telemetry = CommandTelemetry::new();
let input = vec![b'x'; 1024 * 1024];
for budget in [
crate::MAX_COMMAND_TRANSPORT_BYTES,
(QUEUE_CHUNKS * COMMAND_EVENT_CHUNK_BYTES) as u64,
] {
let (sender, mut receiver) = mpsc::channel(QUEUE_CHUNKS);
let mut accounting = ByteAccounting::default();
{
let output = pump(
input.as_slice(),
&b""[..],
&sender,
&telemetry,
"test",
budget,
&mut accounting,
);
tokio::pin!(output);
let polled = std::future::poll_fn(|cx| Poll::Ready(output.as_mut().poll(cx))).await;
if budget == crate::MAX_COMMAND_TRANSPORT_BYTES {
assert!(polled.is_pending(), "a slow receiver must stop pipe reads");
} else {
assert!(
matches!(polled, Poll::Ready(Ok(()))),
"discard must not reserve queue capacity"
);
}
}
assert_eq!(receiver.len(), QUEUE_CHUNKS);
assert_eq!(
accounting.stdout_delivered,
(QUEUE_CHUNKS * COMMAND_EVENT_CHUNK_BYTES) as u64
);
assert_eq!(
accounting.stdout,
if budget == crate::MAX_COMMAND_TRANSPORT_BYTES {
accounting.stdout_delivered
} else {
input.len() as u64
}
);
let mut queued_bytes = 0;
while let Ok(Message::Bytes(chunk)) = receiver.try_recv() {
queued_bytes += chunk.bytes.len();
}
assert_eq!(queued_bytes, QUEUE_CHUNKS * COMMAND_EVENT_CHUNK_BYTES);
}
}
crates/workdir/src/operation.rs
+162 -19
View File
@@ -1,5 +1,3 @@
use std::path::PathBuf;
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
@@ -8,18 +6,18 @@ use serde::{Deserialize, Serialize};
pub struct CommandHandle(pub String);
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
#[serde(deny_unknown_fields)]
pub struct CommandRequest {
pub command: String,
#[schemars(range(min = 0, max = 9_007_199_254_740_991_u64))]
pub timeout_secs: u64,
#[schemars(range(min = 0, max = 9_007_199_254_740_991_usize))]
pub output_limit: usize,
/// Aggregate unique-byte capture allowance, independent of consumer storage.
/// Providers drain and count excess bytes without delivering their payload.
#[schemars(range(min = 0, max = 67_108_864_u64))]
pub transport_byte_budget: u64,
/// Workdir-relative command directory. Providers validate it against the
/// active session before process start.
pub cwd: fs_operation::FsPath,
/// Provider-local directory where complete output is retained when the
/// inline result exceeds `output_limit`.
pub spill_dir: Option<PathBuf>,
/// Optional caller-owned correlation id. Bash supplies its tool-call id so
/// user-facing command telemetry can update the corresponding Console row
/// without exposing provider/session handles.
@@ -27,12 +25,23 @@ pub struct CommandRequest {
pub tool_call_id: Option<String>,
}
impl CommandRequest {
pub fn validate(&self) -> Result<(), crate::WorkdirError> {
if self.transport_byte_budget > MAX_COMMAND_TRANSPORT_BYTES {
return Err(crate::WorkdirError::InvalidArgument(
"command transport byte budget exceeds provider maximum".into(),
));
}
Ok(())
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
pub struct CommandOutputRequest {
pub handle: CommandHandle,
#[schemars(range(min = 0, max = 9_007_199_254_740_991_usize))]
pub cursor: usize,
#[schemars(range(min = 0, max = 9_007_199_254_740_991_usize))]
#[schemars(range(min = 0, max = 67_108_864_u64))]
pub cursor: u64,
#[schemars(range(min = 1, max = 65_536_usize))]
pub limit: usize,
pub wait: bool,
}
@@ -107,16 +116,150 @@ pub enum CommandEvent {
},
}
/// Sequence is the merged transport byte position. Offset addresses raw bytes
/// within the named stream; neither ordering nor offsets depend on UTF-8 decoding.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
pub struct CommandOutput {
pub struct CommandByteChunk {
pub sequence: u64,
pub stream: CommandStream,
pub offset: u64,
pub bytes: Vec<u8>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
#[serde(rename_all = "snake_case")]
pub enum CommandOutputTruncation {
ByteBudget,
}
/// Raw totals include discarded pipe bytes. Delivered totals count only prefixes
/// admitted to the transport queue, not page reads or retries by a consumer.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
pub struct CommandTerminal {
pub status: CommandStatus,
#[schemars(range(min = -2147483648_i32, max = 2147483647_i32))]
pub exit_code: Option<i32>,
pub timed_out: bool,
pub content: String,
#[schemars(range(min = 0, max = 9_007_199_254_740_991_usize))]
pub next_cursor: Option<usize>,
pub truncated: bool,
/// Complete output retained by the provider when `truncated` is true.
pub output_path: Option<PathBuf>,
pub stdout_bytes: u64,
pub stderr_bytes: u64,
pub stdout_delivered_bytes: u64,
pub stderr_delivered_bytes: u64,
pub truncation_reason: Option<CommandOutputTruncation>,
}
/// Sending next_cursor acknowledges this page. Only the current page may be
/// retried; acknowledged bytes expire. Terminal fences all captured bytes,
/// including on timeout/cancel. Consumers must pull while the process runs.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
#[serde(deny_unknown_fields)]
pub struct CommandOutput {
pub chunks: Vec<CommandByteChunk>,
pub next_cursor: u64,
pub terminal: Option<CommandTerminal>,
}
impl CommandOutput {
pub(crate) fn transport_bounds_valid(&self) -> bool {
let mut end = None;
let mut stream_ends = [None, None];
for chunk in &self.chunks {
let stream_index = match chunk.stream {
CommandStream::Stdout => 0,
CommandStream::Stderr => 1,
};
if chunk.bytes.is_empty()
|| chunk.offset > chunk.sequence
|| stream_ends[stream_index].is_some_and(|end| chunk.offset != end)
|| end.is_some_and(|end| chunk.sequence != end)
|| chunk
.sequence
.checked_add(chunk.bytes.len() as u64)
.is_none_or(|end| end > MAX_COMMAND_TRANSPORT_BYTES)
|| chunk
.offset
.checked_add(chunk.bytes.len() as u64)
.is_none_or(|end| end > MAX_COMMAND_TRANSPORT_BYTES)
{
return false;
}
end = chunk.sequence.checked_add(chunk.bytes.len() as u64);
stream_ends[stream_index] = chunk.offset.checked_add(chunk.bytes.len() as u64);
}
self.next_cursor <= MAX_COMMAND_TRANSPORT_BYTES
&& end.is_none_or(|end| end == self.next_cursor)
&& self.terminal.as_ref().is_none_or(|terminal| {
let truncated = terminal.stdout_bytes > terminal.stdout_delivered_bytes
|| terminal.stderr_bytes > terminal.stderr_delivered_bytes;
stream_ends[0].is_none_or(|end| end == terminal.stdout_delivered_bytes)
&& stream_ends[1].is_none_or(|end| end == terminal.stderr_delivered_bytes)
&& terminal.status != CommandStatus::Running
&& terminal.stdout_delivered_bytes <= terminal.stdout_bytes
&& terminal.stderr_delivered_bytes <= terminal.stderr_bytes
&& terminal
.stdout_delivered_bytes
.checked_add(terminal.stderr_delivered_bytes)
== Some(self.next_cursor)
&& terminal.truncation_reason
== truncated.then_some(CommandOutputTruncation::ByteBudget)
})
}
}
pub const MAX_COMMAND_TRANSPORT_BYTES: u64 = 64 * 1024 * 1024;
pub const MAX_COMMAND_READ_BYTES: usize = 64 * 1024;
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn command_wire_rejects_legacy_output_host_paths() {
let command = serde_json::json!({
"command": "printf bytes", "timeout_secs": 1, "cwd": "", "transport_byte_budget": 0
});
assert!(serde_json::from_value::<CommandRequest>(command.clone()).is_ok());
for field in ["spill_dir", "output_limit"] {
let mut legacy = command.clone();
legacy[field] = serde_json::json!("/host/private");
assert!(serde_json::from_value::<CommandRequest>(legacy).is_err());
}
let output = serde_json::json!({"chunks": [], "next_cursor": 0, "terminal": null, "output_path": "/host/private"});
assert!(serde_json::from_value::<CommandOutput>(output).is_err());
}
#[test]
fn command_budget_wire_requires_field_and_checks_provider_hard_bound() {
let wire = serde_json::json!({"command": "true", "timeout_secs": 1, "cwd": ""});
assert!(serde_json::from_value::<CommandRequest>(wire.clone()).is_err());
for budget in [
0,
MAX_COMMAND_TRANSPORT_BYTES,
MAX_COMMAND_TRANSPORT_BYTES + 1,
] {
let mut value = wire.clone();
value["transport_byte_budget"] = serde_json::json!(budget);
let request: CommandRequest = serde_json::from_value(value).unwrap();
assert_eq!(
request.validate().is_ok(),
budget <= MAX_COMMAND_TRANSPORT_BYTES
);
}
let mut negative = wire;
negative["transport_byte_budget"] = serde_json::json!(-1);
assert!(serde_json::from_value::<CommandRequest>(negative).is_err());
assert_eq!(
serde_json::to_value(CommandOutputTruncation::ByteBudget).unwrap(),
"byte_budget"
);
let schema = serde_json::to_value(schemars::schema_for!(CommandRequest)).unwrap();
assert_eq!(
schema["properties"]["transport_byte_budget"]["maximum"],
MAX_COMMAND_TRANSPORT_BYTES
);
assert!(
schema["required"]
.as_array()
.unwrap()
.contains(&serde_json::json!("transport_byte_budget"))
);
}
}
crates/workdir/src/router.rs
+12 -7
View File
@@ -341,7 +341,9 @@ impl WorkdirSession for RoutedWorkdirSession {
let _active = self.state.enter()?;
let handle = request.handle.clone();
let output = self.state.session.command_output(request).await?;
self.state.mark_terminal(&handle, output.status);
if let Some(terminal) = &output.terminal {
self.state.mark_terminal(&handle, terminal.status);
}
Ok(output)
}
@@ -789,11 +791,12 @@ mod tests {
.unwrap();
let handle = routed
.start_command(CommandRequest {
transport_byte_budget: crate::MAX_COMMAND_TRANSPORT_BYTES,
command: "printf done".to_string(),
timeout_secs: 5,
output_limit: 1024,
cwd: crate::WorkdirPath::root(),
spill_dir: None,
tool_call_id: None,
})
.await
@@ -823,11 +826,12 @@ mod tests {
let routed = router.attach(alias, Arc::clone(&underlying)).unwrap();
let handle = routed
.start_command(CommandRequest {
transport_byte_budget: crate::MAX_COMMAND_TRANSPORT_BYTES,
command: "sleep 30".to_string(),
timeout_secs: 60,
output_limit: 1024,
cwd: crate::WorkdirPath::root(),
spill_dir: None,
tool_call_id: None,
})
.await
@@ -850,11 +854,12 @@ mod tests {
let routed = router.attach(alias.clone(), session("wd", &dir)).unwrap();
let handle = routed
.start_command(CommandRequest {
transport_byte_budget: crate::MAX_COMMAND_TRANSPORT_BYTES,
command: "sleep 30".to_string(),
timeout_secs: 60,
output_limit: 1024,
cwd: crate::WorkdirPath::root(),
spill_dir: None,
tool_call_id: None,
})
.await
crates/workdir/src/scope.rs
+291 -42
View File
@@ -10,6 +10,9 @@
//! provider-side before name enumeration, descent, or content reads. Returned coordinates are not
//! prefixed with cwd again.
#[cfg(test)]
use crate::command_test_support::CollectCommand as _;
use std::collections::{BTreeMap, HashMap, HashSet};
use std::path::Path;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
@@ -124,6 +127,7 @@ impl WorkdirToolBroker {
next_lease_id: AtomicU64::new(1),
close_lock: Arc::new(tokio::sync::Mutex::new(())),
owned_commands: Arc::new(Mutex::new(HashSet::new())),
command_reads: Mutex::new(HashMap::new()),
pending_command_events: Arc::new(Mutex::new(HashMap::new())),
starting_tool_calls: Arc::new(Mutex::new(HashSet::new())),
forwarded_starts: Arc::new(Mutex::new(HashSet::new())),
@@ -230,25 +234,48 @@ impl WorkdirScopeLease {
.source
.cancel_command(handle.clone())
.await;
let terminal = self
let read_cursor = self
.broker
.authority
.source
.command_output(CommandOutputRequest {
handle,
cursor: 0,
limit: 1,
wait: true,
})
.await;
.command_reads
.lock()
.expect("command cursor mutex poisoned")
.get(&command_id)
.cloned()
.expect("owned command cursor");
let mut read_cursor = read_cursor.lock().await;
let mut cursor = *read_cursor;
let terminal = loop {
let output = self
.broker
.authority
.source
.command_output(CommandOutputRequest {
handle: handle.clone(),
cursor,
limit: crate::MAX_COMMAND_READ_BYTES,
wait: true,
})
.await;
match output {
Ok(output) => {
cursor = output.next_cursor;
*read_cursor = cursor;
if let Some(terminal) = output.terminal {
break Ok(terminal);
}
}
Err(error) => break Err(error),
}
};
match (cancel, terminal) {
(_, Ok(output)) => {
self.broker.authority.publish_terminal_if_missing(
&command_id,
output.status,
output.exit_code,
output.next_cursor.unwrap_or(output.content.len()) as u64,
&output.content,
cursor,
"",
);
self.broker
.authority
@@ -582,7 +609,12 @@ struct ScopedWorkdirSession {
scope_lock: tokio::sync::Mutex<()>,
next_lease_id: AtomicU64,
close_lock: Arc<tokio::sync::Mutex<()>>,
// Live cleanup responsibility ends at terminal delivery. Read ownership is
// retained separately so a lost terminal response can still be replayed.
owned_commands: Arc<Mutex<HashSet<String>>>,
// Starting commands prunes completed entries at the bounded replay ceiling;
// current scope validity is still checked before every read or retry.
command_reads: Mutex<HashMap<String, Arc<tokio::sync::Mutex<u64>>>>,
pending_command_events: Arc<Mutex<HashMap<String, Vec<CommandEvent>>>>,
starting_tool_calls: Arc<Mutex<HashSet<String>>>,
forwarded_starts: Arc<Mutex<HashSet<String>>>,
@@ -700,6 +732,11 @@ impl ScopedWorkdirSession {
.lock()
.expect("scoped command set mutex poisoned")
.contains(&handle.0)
|| self
.command_reads
.lock()
.expect("command cursor mutex poisoned")
.contains_key(&handle.0)
{
Ok(())
} else {
@@ -1013,6 +1050,7 @@ impl ScopedWorkdirSession {
next_lease_id: AtomicU64::new(1),
close_lock: close_lock.clone(),
owned_commands,
command_reads: Mutex::new(HashMap::new()),
pending_command_events,
starting_tool_calls,
forwarded_starts,
@@ -1407,6 +1445,16 @@ impl WorkdirSession for ScopedWorkdirSession {
.lock()
.expect("scoped command set mutex poisoned");
owned.insert(handle.0.clone());
{
let mut reads = self
.command_reads
.lock()
.expect("command cursor mutex poisoned");
if reads.len() >= 64 {
reads.retain(|id, _| owned.contains(id));
}
reads.insert(handle.0.clone(), Arc::new(tokio::sync::Mutex::new(0)));
}
if let Some(tool_call_id) = &tool_call_id {
self.starting_tool_calls
.lock()
@@ -1457,14 +1505,28 @@ impl WorkdirSession for ScopedWorkdirSession {
) -> Result<CommandOutput, WorkdirError> {
self.ensure_owned_command(&request.handle)?;
let command_id = request.handle.0.clone();
let read_cursor = self
.command_reads
.lock()
.expect("command cursor mutex poisoned")
.get(&command_id)
.cloned();
let mut read_cursor = match &read_cursor {
Some(cursor) => Some(cursor.lock().await),
None => None,
};
self.ensure_owned_command(&request.handle)?;
let output = self.source.command_output(request).await?;
if !matches!(output.status, CommandStatus::Running) {
if let Some(cursor) = &mut read_cursor {
**cursor = output.next_cursor;
}
if let Some(terminal) = &output.terminal {
self.publish_terminal_if_missing(
&command_id,
output.status,
output.exit_code,
output.next_cursor.unwrap_or(output.content.len()) as u64,
&output.content,
terminal.status,
terminal.exit_code,
output.next_cursor,
"",
);
self.owned_commands
.lock()
@@ -2080,20 +2142,21 @@ mod tests {
session: &WorkdirSessionHandle,
command: impl Into<String>,
tool_call_id: impl Into<String>,
) -> CommandOutput {
) -> crate::command_test_support::CollectedCommand {
let handle = session
.start_command(CommandRequest {
transport_byte_budget: crate::MAX_COMMAND_TRANSPORT_BYTES,
command: command.into(),
timeout_secs: 5,
output_limit: 1024,
cwd: FsPath::root(),
spill_dir: None,
tool_call_id: Some(tool_call_id.into()),
})
.await
.unwrap();
session
.command_output(CommandOutputRequest {
.collect_command_output(CommandOutputRequest {
handle,
cursor: 0,
limit: 1024,
@@ -2112,11 +2175,12 @@ mod tests {
.expect("delegation wrapper must preserve command observation");
let handle = parent
.start_command(CommandRequest {
transport_byte_budget: crate::MAX_COMMAND_TRANSPORT_BYTES,
command: "printf ready; sleep 0.2; printf done".into(),
timeout_secs: 5,
output_limit: 1024,
cwd: FsPath::root(),
spill_dir: None,
tool_call_id: Some("tool-delegated".into()),
})
.await
@@ -2139,7 +2203,7 @@ mod tests {
assert_eq!(snapshots[0].stdout.content, "ready");
let output = parent
.command_output(CommandOutputRequest {
.collect_command_output(CommandOutputRequest {
handle,
cursor: 0,
limit: 1024,
@@ -2179,11 +2243,12 @@ mod tests {
);
let error = child
.start_command(CommandRequest {
transport_byte_budget: crate::MAX_COMMAND_TRANSPORT_BYTES,
command: "pwd".into(),
timeout_secs: 5,
output_limit: 1024,
cwd: FsPath::root(),
spill_dir: None,
tool_call_id: None,
})
.await
@@ -2211,11 +2276,12 @@ mod tests {
let handle = first
.start_command(CommandRequest {
transport_byte_budget: crate::MAX_COMMAND_TRANSPORT_BYTES,
command: "pwd; sleep 0.2".into(),
timeout_secs: 5,
output_limit: 4096,
cwd: FsPath::root(),
spill_dir: None,
tool_call_id: Some("first-command".into()),
})
.await
@@ -2239,7 +2305,7 @@ mod tests {
));
let output = first
.command_output(CommandOutputRequest {
.collect_command_output(CommandOutputRequest {
handle,
cursor: 0,
limit: 4096,
@@ -2316,11 +2382,12 @@ mod tests {
assert!(matches!(
child
.start_command(CommandRequest {
transport_byte_budget: crate::MAX_COMMAND_TRANSPORT_BYTES,
command: "printf denied".into(),
timeout_secs: 5,
output_limit: 1024,
cwd: FsPath::root(),
spill_dir: None,
tool_call_id: Some("read-only-command".into()),
})
.await,
@@ -2622,11 +2689,12 @@ mod tests {
assert!(matches!(
child
.start_command(CommandRequest {
transport_byte_budget: crate::MAX_COMMAND_TRANSPORT_BYTES,
command: "printf revoked".into(),
timeout_secs: 5,
output_limit: 1024,
cwd: FsPath::root(),
spill_dir: None,
tool_call_id: Some("revoked-child-command".into()),
})
.await,
@@ -2724,11 +2792,12 @@ mod tests {
for index in 0..MAX_SCOPED_COMMANDS {
child
.start_command(CommandRequest {
transport_byte_budget: crate::MAX_COMMAND_TRANSPORT_BYTES,
command: "sleep 30".into(),
timeout_secs: 60,
output_limit: 1024,
cwd: FsPath::root(),
spill_dir: None,
tool_call_id: Some(format!("command-{index}")),
})
.await
@@ -2737,11 +2806,12 @@ mod tests {
let error = child
.start_command(CommandRequest {
transport_byte_budget: crate::MAX_COMMAND_TRANSPORT_BYTES,
command: "sleep 30".into(),
timeout_secs: 60,
output_limit: 1024,
cwd: FsPath::root(),
spill_dir: None,
tool_call_id: Some("command-over-limit".into()),
})
.await
@@ -2771,11 +2841,12 @@ mod tests {
let command = tokio::spawn(async move {
command_child
.start_command(CommandRequest {
transport_byte_budget: crate::MAX_COMMAND_TRANSPORT_BYTES,
command: "sleep 30".into(),
timeout_secs: 60,
output_limit: 1024,
cwd: FsPath::root(),
spill_dir: None,
tool_call_id: Some("racing-command".into()),
})
.await
@@ -2817,11 +2888,12 @@ mod tests {
let mut events = child.subscribe_command_events().unwrap();
let handle = child
.start_command(CommandRequest {
transport_byte_budget: crate::MAX_COMMAND_TRANSPORT_BYTES,
command: "sleep 30; printf leaked > marker".into(),
timeout_secs: 60,
output_limit: 1024,
cwd: FsPath::root(),
spill_dir: None,
tool_call_id: Some("owned-command".into()),
})
.await
@@ -3012,11 +3084,12 @@ mod tests {
assert!(matches!(
parent
.start_command(CommandRequest {
transport_byte_budget: crate::MAX_COMMAND_TRANSPORT_BYTES,
command: "printf closed".into(),
timeout_secs: 5,
output_limit: 1024,
cwd: FsPath::root(),
spill_dir: None,
tool_call_id: Some("closed-parent-command".into()),
})
.await,
@@ -3028,4 +3101,180 @@ mod tests {
"child result after parent close: {child_result:?}"
);
}
#[tokio::test]
async fn scoped_byte_cursor_survives_partial_ack_then_revocation_and_terminal_retry() {
let root = TempDir::new().unwrap();
fs::create_dir(root.path().join("child")).unwrap();
let parent = session(root.path());
let child = parent
.scope(request("child", WorkdirToolScopePermission::Write))
.await
.unwrap();
let tool = child.tool_session();
let handle = tool
.start_command(CommandRequest {
transport_byte_budget: crate::MAX_COMMAND_TRANSPORT_BYTES,
command: "printf ready; exec yes x".into(),
timeout_secs: 10,
cwd: FsPath::root(),
tool_call_id: None,
})
.await
.unwrap();
let mut cursor = 0;
for _ in 0..2 {
let page = tool
.command_output(CommandOutputRequest {
handle: handle.clone(),
cursor,
limit: 1,
wait: true,
})
.await
.unwrap();
cursor = page.next_cursor;
}
tokio::time::timeout(std::time::Duration::from_secs(5), child.close())
.await
.unwrap()
.unwrap();
assert!(
tool.command_output(CommandOutputRequest {
handle,
cursor,
limit: 1,
wait: false
})
.await
.is_err()
);
let handle = parent
.start_command(CommandRequest {
transport_byte_budget: crate::MAX_COMMAND_TRANSPORT_BYTES,
command: "printf parent".into(),
timeout_secs: 5,
cwd: FsPath::root(),
tool_call_id: None,
})
.await
.unwrap();
let mut cursor = 0;
loop {
let request = CommandOutputRequest {
handle: handle.clone(),
cursor,
limit: 1024,
wait: true,
};
let page = parent.command_output(request.clone()).await.unwrap();
assert_eq!(parent.command_output(request).await.unwrap(), page);
cursor = page.next_cursor;
if page.terminal.is_some() {
break;
}
}
}
#[tokio::test]
async fn scoped_terminal_response_loss_replays_without_duplicate_capture_and_rejects_invalid_cursors()
{
tokio::time::timeout(std::time::Duration::from_secs(5), async {
let root = TempDir::new().unwrap();
fs::create_dir(root.path().join("child")).unwrap();
let parent = session(root.path());
let child = parent
.scope(request("child", WorkdirToolScopePermission::Write))
.await
.unwrap();
let tool = child.tool_session();
let handle = tool
.start_command(CommandRequest {
transport_byte_budget: crate::MAX_COMMAND_TRANSPORT_BYTES,
command: "printf abcdef".into(),
timeout_secs: 5,
cwd: FsPath::root(),
tool_call_id: None,
})
.await
.unwrap();
let mut cursor = 0;
let mut captured = Vec::new();
let terminal_request = loop {
let request = CommandOutputRequest {
handle: handle.clone(),
cursor,
limit: 2,
wait: true,
};
// Simulate losing each produced HTTP response, including terminal:
// append only the retried page, never both copies.
let lost = tool.command_output(request.clone()).await.unwrap();
let replay = tool.command_output(request.clone()).await.unwrap();
assert_eq!(replay, lost);
for chunk in &replay.chunks {
captured.extend_from_slice(&chunk.bytes);
}
cursor = replay.next_cursor;
if let Some(terminal) = replay.terminal {
assert_eq!(terminal.status, CommandStatus::Completed);
break request;
}
};
assert_eq!(captured, b"abcdef");
assert!(
!child
.broker
.authority
.owned_commands
.lock()
.unwrap()
.contains(&handle.0)
);
assert!(
child
.broker
.authority
.command_reads
.lock()
.unwrap()
.contains_key(&handle.0)
);
assert_eq!(
tool.command_status(handle.clone()).await.unwrap(),
CommandStatus::Completed
);
for invalid in [0, cursor + 1] {
assert!(
tool.command_output(CommandOutputRequest {
handle: handle.clone(),
cursor: invalid,
limit: 2,
wait: false
})
.await
.is_err()
);
}
// Invalid cursor requests do not destroy the retained terminal page.
assert!(
tool.command_output(terminal_request)
.await
.unwrap()
.terminal
.is_some()
);
child.close().await.unwrap();
assert!(
tool.command_output(CommandOutputRequest {
handle,
cursor,
limit: 2,
wait: false
})
.await
.is_err()
);
})
.await
.expect("bounded scoped terminal replay");
}
}
crates/workdir/tests/command_denial_diagnostics.rs
+3 -2
View File
@@ -103,11 +103,12 @@ impl WorkdirSession for CommandFailureSession {
}
fn command() -> CommandRequest {
CommandRequest {
transport_byte_budget: workdir::MAX_COMMAND_TRANSPORT_BYTES,
command: "fake command, never executed".into(),
timeout_secs: 1,
output_limit: 100,
cwd: WorkdirPath::root(),
spill_dir: None,
tool_call_id: None,
}
}
crates/workdir/tests/denial_reasons.rs
+3 -2
View File
@@ -167,11 +167,12 @@ async fn read_only_and_scoped_capability_refusals_do_not_collapse() {
assert_reason(
child
.start_command(workdir::CommandRequest {
transport_byte_budget: workdir::MAX_COMMAND_TRANSPORT_BYTES,
command: "must-not-run".into(),
timeout_secs: 1,
output_limit: 1,
cwd: WorkdirPath::root(),
spill_dir: None,
tool_call_id: None,
})
.await