use std::collections::BTreeMap;
use std::time::Duration;
use async_trait::async_trait;
use base64::engine::general_purpose::STANDARD as BASE64;
use base64::Engine as _;
use futures::stream::BoxStream;
use futures::StreamExt;
use serde::Deserialize;
use serde_json::json;
use crate::error::{ErrorData, Result};
use crate::traits::{CommandOutput, JobError, JobExit, JobPoll, JobStart, RunCommandRequest};
use alien_error::{AlienError, Context, ContextError, IntoAlienError};
pub use alien_core::sandbox_image::AGENT_PORT;
const RUN_COMMAND: &str = "sandbox.runCommand";
const JOB_START: &str = "sandbox.jobStart";
const JOB_POLL: &str = "sandbox.jobPoll";
const JOB_CANCEL: &str = "sandbox.jobCancel";
#[cfg(not(test))]
const AGENT_RESPONSE_TIMEOUT: Duration = Duration::from_secs(60);
#[cfg(test)]
const AGENT_RESPONSE_TIMEOUT: Duration = Duration::from_millis(200);
#[cfg(not(test))]
const TERMINAL_FRAME_GRACE: Duration = Duration::from_secs(60);
#[cfg(test)]
const TERMINAL_FRAME_GRACE: Duration = Duration::from_millis(200);
#[async_trait]
pub trait AgentTransport: Send + Sync + std::fmt::Debug {
async fn request(
&self,
sandbox_id: &str,
method: reqwest::Method,
path: &str,
) -> Result<reqwest::RequestBuilder>;
fn provider(&self) -> &'static str;
}
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase", tag = "t")]
enum AgentFrame {
Stdout {
seq: u64,
data: String,
},
Stderr {
seq: u64,
data: String,
},
Exit {
code: i32,
#[serde(default)]
truncated: bool,
},
Error {
code: String,
message: String,
},
}
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
struct ReadFileResponse {
contents_base64: String,
}
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
struct JobStartResponse {
job_id: String,
}
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
struct JobPollResponse {
running: bool,
#[serde(default)]
frames: Vec<AgentFrame>,
#[serde(default)]
exit_code: Option<i32>,
#[serde(default)]
truncated: Option<bool>,
#[serde(default)]
error: Option<JobErrorResponse>,
}
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
struct JobErrorResponse {
code: String,
message: String,
}
pub async fn run_command<T: AgentTransport + ?Sized>(
transport: &T,
sandbox_id: &str,
request: RunCommandRequest,
) -> Result<BoxStream<'static, Result<CommandOutput>>> {
let body = exec_body(&request)?;
let response = send(
transport
.request(sandbox_id, reqwest::Method::POST, "/v1/exec")
.await?
.json(&body),
RUN_COMMAND,
)
.await?;
Ok(frame_stream(
response,
transport.provider(),
request.timeout,
))
}
pub async fn start_job<T: AgentTransport + ?Sized>(
transport: &T,
sandbox_id: &str,
request: RunCommandRequest,
) -> Result<JobStart> {
let body = exec_body(&request)?;
let response = send(
transport
.request(sandbox_id, reqwest::Method::POST, "/v1/jobs/start")
.await?
.json(&body),
JOB_START,
)
.await?;
let started: JobStartResponse = response
.json()
.await
.into_alien_error()
.context(ErrorData::UnexpectedResponseFormat {
provider: transport.provider().to_string(),
binding_name: JOB_START.to_string(),
field: "jobId".to_string(),
response_json: "the agent returned a body this provider cannot parse".to_string(),
})
.context(unanswered(
JOB_START,
"the job started and its id could not be read",
))?;
Ok(JobStart {
job_id: started.job_id,
})
}
pub async fn poll_job<T: AgentTransport + ?Sized>(
transport: &T,
sandbox_id: &str,
job_id: &str,
since_seq: Option<u64>,
) -> Result<JobPoll> {
let response = send(
transport
.request(sandbox_id, reqwest::Method::POST, "/v1/jobs/poll")
.await?
.json(&json!({ "jobId": job_id, "sinceSeq": since_seq })),
JOB_POLL,
)
.await?;
let JobPollResponse {
running,
frames,
exit_code,
truncated,
error,
} = response
.json()
.await
.into_alien_error()
.context(unanswered(JOB_POLL, "the poll's body could not be read"))?;
Ok(JobPoll {
running,
frames: frames
.into_iter()
.map(|frame| frame.into_output(transport.provider()))
.collect::<Result<Vec<_>>>()?,
exit: exit_code.map(|code| JobExit {
code,
truncated: truncated.unwrap_or(false),
}),
error: error.map(|error| JobError {
code: error.code,
message: error.message,
}),
})
}
pub async fn cancel_job<T: AgentTransport + ?Sized>(
transport: &T,
sandbox_id: &str,
job_id: &str,
) -> Result<()> {
send(
transport
.request(sandbox_id, reqwest::Method::POST, "/v1/jobs/cancel")
.await?
.json(&json!({ "jobId": job_id })),
JOB_CANCEL,
)
.await?;
Ok(())
}
fn exec_body(request: &RunCommandRequest) -> Result<serde_json::Value> {
if timeout_millis(request.timeout) == 0 {
return Err(AlienError::new(ErrorData::SandboxCommandFailed {
failure: "invalidRequest".to_string(),
reason: "a command must carry a non-zero timeout".to_string(),
}));
}
Ok(json!({
"command": request.argv(),
"timeoutMs": timeout_millis(request.timeout),
"cwd": request.cwd,
"env": request.env,
}))
}
pub async fn read_file<T: AgentTransport + ?Sized>(
transport: &T,
sandbox_id: &str,
path: &str,
) -> Result<Vec<u8>> {
let response = send(
transport
.request(sandbox_id, reqwest::Method::GET, "/v1/files")
.await?
.query(&[("path", path)]),
"sandbox.readFile",
)
.await?;
let body: ReadFileResponse =
response
.json()
.await
.into_alien_error()
.context(ErrorData::UnexpectedResponseFormat {
provider: transport.provider().to_string(),
binding_name: "sandbox.readFile".to_string(),
field: "body".to_string(),
response_json: "the agent returned a body this provider cannot parse".to_string(),
})?;
decode(
&body.contents_base64,
transport.provider(),
"sandbox.readFile",
"contentsBase64",
)
}
pub async fn write_files<T: AgentTransport + ?Sized>(
transport: &T,
sandbox_id: &str,
files: BTreeMap<String, Vec<u8>>,
) -> Result<()> {
for (path, contents) in files {
send(
transport
.request(sandbox_id, reqwest::Method::PUT, "/v1/files")
.await?
.json(&json!({
"path": path,
"contentsBase64": BASE64.encode(contents),
})),
"sandbox.writeFiles",
)
.await?;
}
Ok(())
}
fn timeout_millis(timeout: Duration) -> u64 {
u64::try_from(timeout.as_millis()).unwrap_or(u64::MAX)
}
fn unanswered(operation: &str, reason: &str) -> ErrorData {
if operation == RUN_COMMAND || operation == JOB_START {
return ErrorData::SandboxOutcomeUnknown {
operation: operation.to_string(),
reason: reason.to_string(),
};
}
ErrorData::SandboxUnreachable {
operation: operation.to_string(),
reason: reason.to_string(),
}
}
pub async fn send(request: reqwest::RequestBuilder, operation: &str) -> Result<reqwest::Response> {
let response = match tokio::time::timeout(AGENT_RESPONSE_TIMEOUT, request.send()).await {
Ok(sent) => sent
.into_alien_error()
.context(unanswered(operation, "the request never reached the agent"))?,
Err(_) => {
return Err(AlienError::new(unanswered(
operation,
&format!(
"the agent did not answer within {}s",
AGENT_RESPONSE_TIMEOUT.as_secs()
),
)));
}
};
if response.status().is_success() {
return Ok(response);
}
let status = response.status();
let body = match response.text().await {
Ok(body) => body,
Err(error) => {
return Err(error).into_alien_error().context(unanswered(
operation,
&format!("{operation} returned {status} and its body could not be read"),
))
}
};
if status.is_server_error() {
return Err(AlienError::new(unanswered(
operation,
&format!("the sandbox host returned {status}: {body}"),
)));
}
Err(AlienError::new(ErrorData::SandboxCommandFailed {
failure: "agentRefused".to_string(),
reason: format!("{operation} returned {status}: {body}"),
}))
}
fn frame_stream(
response: reqwest::Response,
provider: &'static str,
timeout: Duration,
) -> BoxStream<'static, Result<CommandOutput>> {
struct State {
bytes: BoxStream<'static, reqwest::Result<bytes::Bytes>>,
buffer: Vec<u8>,
finished: bool,
saw_terminal: bool,
provider: &'static str,
stall: std::pin::Pin<Box<tokio::time::Sleep>>,
}
let state = State {
bytes: response.bytes_stream().boxed(),
buffer: Vec::new(),
finished: false,
saw_terminal: false,
provider,
stall: Box::pin(tokio::time::sleep(
timeout.saturating_add(TERMINAL_FRAME_GRACE),
)),
};
futures::stream::unfold(state, |mut state| async move {
loop {
if let Some(index) = state.buffer.iter().position(|byte| *byte == b'\n') {
let line: Vec<u8> = state.buffer.drain(..=index).collect();
let line = &line[..line.len() - 1];
if line.is_empty() {
continue;
}
let frame = match serde_json::from_slice::<AgentFrame>(line) {
Ok(frame) => frame,
Err(error) => {
state.finished = true;
state.buffer.clear();
let failure = malformed(&error.to_string(), state.provider)
.context(unanswered(RUN_COMMAND, "an output frame did not parse"));
return Some((Err(failure), state));
}
};
if matches!(frame, AgentFrame::Exit { .. } | AgentFrame::Error { .. }) {
state.saw_terminal = true;
}
let output = frame.into_output(state.provider);
if output.is_err() {
state.finished = true;
state.buffer.clear();
}
return Some((output, state));
}
if state.finished {
return None;
}
let chunk = tokio::select! {
biased;
chunk = state.bytes.next() => chunk,
() = &mut state.stall => {
state.finished = true;
state.bytes = futures::stream::empty().boxed();
return Some((
Err(AlienError::new(unanswered(
RUN_COMMAND,
&format!(
"no terminal frame arrived within {}s of the command's timeout",
TERMINAL_FRAME_GRACE.as_secs()
),
))),
state,
));
}
};
match chunk {
Some(Ok(chunk)) => state.buffer.extend_from_slice(&chunk),
Some(Err(error)) => {
state.finished = true;
return Some((
Err(AlienError::new(unanswered(
RUN_COMMAND,
&format!("the output stream failed: {error}"),
))),
state,
));
}
None => {
state.finished = true;
if !state.saw_terminal {
return Some((
Err(AlienError::new(unanswered(
RUN_COMMAND,
"the output stream ended without a terminal frame",
))),
state,
));
}
return None;
}
}
}
})
.boxed()
}
impl AgentFrame {
fn into_output(self, provider: &'static str) -> Result<CommandOutput> {
match self {
Self::Stdout { seq, data } => Ok(CommandOutput::Stdout {
seq,
data: decode(&data, provider, RUN_COMMAND, "data")
.context(unanswered(RUN_COMMAND, "an output frame did not decode"))?,
}),
Self::Stderr { seq, data } => Ok(CommandOutput::Stderr {
seq,
data: decode(&data, provider, RUN_COMMAND, "data")
.context(unanswered(RUN_COMMAND, "an output frame did not decode"))?,
}),
Self::Exit { code, truncated } => Ok(CommandOutput::Exit { code, truncated }),
Self::Error { code, message } => {
Err(AlienError::new(ErrorData::SandboxCommandFailed {
failure: code,
reason: message,
}))
}
}
}
}
fn decode(data: &str, provider: &'static str, binding_name: &str, field: &str) -> Result<Vec<u8>> {
BASE64
.decode(data)
.into_alien_error()
.context(ErrorData::UnexpectedResponseFormat {
provider: provider.to_string(),
binding_name: binding_name.to_string(),
field: field.to_string(),
response_json: format!("{field} was not valid base64"),
})
}
fn malformed(reason: &str, provider: &'static str) -> AlienError<ErrorData> {
AlienError::new(ErrorData::UnexpectedResponseFormat {
provider: provider.to_string(),
binding_name: RUN_COMMAND.to_string(),
field: "frame".to_string(),
response_json: format!("an output frame did not parse: {reason}"),
})
}
#[cfg(test)]
mod tests {
use super::*;
use crate::traits::CommandOutput;
use axum::extract::State;
use axum::http::StatusCode;
use axum::response::IntoResponse;
use axum::routing::post;
use axum::{Json, Router};
use std::net::SocketAddr;
use std::sync::{Arc, Mutex};
async fn serve_frames(chunks: Vec<&'static str>) -> String {
let handler = move || {
let chunks = chunks.clone();
async move {
let stream = futures::stream::iter(
chunks
.into_iter()
.map(|chunk| Ok::<_, std::io::Error>(bytes::Bytes::from(chunk))),
);
axum::body::Body::from_stream(stream).into_response()
}
};
let router = Router::new().route("/v1/exec", post(handler));
let listener = tokio::net::TcpListener::bind::<SocketAddr>("127.0.0.1:0".parse().unwrap())
.await
.expect("bind");
let address = listener.local_addr().expect("address");
tokio::spawn(async move {
axum::serve(listener, router).await.expect("serve");
});
format!("http://{address}")
}
async fn frames_from(chunks: Vec<&'static str>) -> Vec<Result<CommandOutput>> {
let base = serve_frames(chunks).await;
let response = reqwest::Client::new()
.post(format!("{base}/v1/exec"))
.send()
.await
.expect("responds");
frame_stream(response, "test-sandbox", Duration::from_secs(30))
.collect::<Vec<_>>()
.await
}
async fn send_status(status: StatusCode, body: &'static str) -> AlienError<ErrorData> {
let handler = move || async move { (status, body).into_response() };
let router = Router::new().route("/v1/exec", post(handler));
let listener = tokio::net::TcpListener::bind::<SocketAddr>("127.0.0.1:0".parse().unwrap())
.await
.expect("bind");
let address = listener.local_addr().expect("address");
tokio::spawn(async move { axum::serve(listener, router).await.expect("serve") });
send(
reqwest::Client::new().post(format!("http://{address}/v1/exec")),
RUN_COMMAND,
)
.await
.expect_err("a non-success must be an error")
}
#[tokio::test]
async fn a_bodyless_server_error_is_not_reported_as_the_agent_refusing() {
let error = send_status(StatusCode::BAD_GATEWAY, "").await;
let rendered = error.to_string();
assert!(
!rendered.contains("agentRefused"),
"a proxy 502 is not the agent refusing: {rendered}"
);
assert_eq!(
error.code, "SANDBOX_OUTCOME_UNKNOWN",
"a proxy can synthesize a 502 after the agent accepted the request, so the caller has \
to be told the outcome is unknown rather than that it is safe to repeat: {rendered}"
);
}
#[tokio::test]
async fn a_server_error_is_unknown_however_much_body_it_carries() {
let error = send_status(StatusCode::INTERNAL_SERVER_ERROR, "spawn failed: ENOMEM").await;
let rendered = error.to_string();
assert!(
rendered.contains("spawn failed: ENOMEM"),
"the cause has to survive: {rendered}"
);
assert_eq!(error.code, "SANDBOX_OUTCOME_UNKNOWN", "got: {rendered}");
}
#[tokio::test]
async fn a_gateway_timeout_with_an_error_page_does_not_read_as_the_agent_refusing() {
let error = send_status(
StatusCode::GATEWAY_TIMEOUT,
"<html><body>504 Gateway Time-out</body></html>",
)
.await;
assert_eq!(error.code, "SANDBOX_OUTCOME_UNKNOWN", "got: {error}");
assert!(
!error.retryable,
"the command may still be running behind the gateway: {error}"
);
}
#[tokio::test]
async fn a_client_error_is_the_agent_answering() {
let error = send_status(StatusCode::BAD_REQUEST, "the command was empty").await;
assert_ne!(
error.code, "SANDBOX_OUTCOME_UNKNOWN",
"a refusal before dispatch establishes that nothing ran: {error}"
);
}
#[tokio::test]
async fn an_agent_that_never_answers_is_refused_within_the_bound() {
let handler = || async {
tokio::time::sleep(AGENT_RESPONSE_TIMEOUT * 20).await;
"late".into_response()
};
let router = Router::new().route("/v1/exec", post(handler));
let listener = tokio::net::TcpListener::bind::<SocketAddr>("127.0.0.1:0".parse().unwrap())
.await
.expect("bind");
let address = listener.local_addr().expect("address");
tokio::spawn(async move { axum::serve(listener, router).await.expect("serve") });
let started = std::time::Instant::now();
let error = send(
reqwest::Client::new().post(format!("http://{address}/v1/exec")),
RUN_COMMAND,
)
.await
.expect_err("a stalled agent must be refused, not waited on");
assert!(
started.elapsed() < AGENT_RESPONSE_TIMEOUT * 10,
"refused at the bound, not at the agent's leisure: {:?}",
started.elapsed()
);
assert!(
error.to_string().contains("did not answer"),
"the refusal says the agent stalled: {error}"
);
}
#[tokio::test]
async fn a_slow_body_after_prompt_headers_is_not_cut_off() {
let handler = || async {
let frames = async_stream_frames(vec![
"{\"t\":\"stdout\",\"seq\":0,\"data\":\"aGk=\"}\n",
"{\"t\":\"exit\",\"code\":0,\"truncated\":false}\n",
]);
axum::body::Body::from_stream(frames).into_response()
};
let router = Router::new().route("/v1/exec", post(handler));
let listener = tokio::net::TcpListener::bind::<SocketAddr>("127.0.0.1:0".parse().unwrap())
.await
.expect("bind");
let address = listener.local_addr().expect("address");
tokio::spawn(async move { axum::serve(listener, router).await.expect("serve") });
let response = send(
reqwest::Client::new().post(format!("http://{address}/v1/exec")),
RUN_COMMAND,
)
.await
.expect("headers arrive at once");
let outputs = frame_stream(response, "test-sandbox", Duration::from_secs(30))
.collect::<Vec<_>>()
.await;
assert_eq!(outputs.len(), 2, "every frame arrived: {outputs:?}");
assert_eq!(
outputs[1].as_ref().expect("exit"),
&CommandOutput::Exit {
code: 0,
truncated: false
}
);
}
async fn frozen_after_one_frame() -> TestTransport {
let handler = || async {
let frames = futures::stream::iter(vec![Ok::<bytes::Bytes, std::io::Error>(
bytes::Bytes::from("{\"t\":\"stdout\",\"seq\":0,\"data\":\"aGk=\"}\n"),
)])
.chain(futures::stream::pending());
axum::body::Body::from_stream(frames).into_response()
};
serve(Router::new().route("/v1/exec", post(handler))).await
}
#[tokio::test]
async fn a_frozen_sandbox_refuses_an_open_stream_rather_than_parking_its_reader() {
let transport = frozen_after_one_frame().await;
let stream = run_command(
&transport,
"sandbox-1",
RunCommandRequest {
command: "sleep".to_string(),
args: vec!["1".to_string()],
cwd: None,
env: BTreeMap::new(),
timeout: TERMINAL_FRAME_GRACE,
},
)
.await
.expect("the agent answers with headers");
let outputs = tokio::time::timeout(TERMINAL_FRAME_GRACE * 20, stream.collect::<Vec<_>>())
.await
.expect("a frozen sandbox must end the stream, not park its reader");
assert_eq!(
outputs.len(),
2,
"the frame that did arrive is kept: {outputs:?}"
);
let error = outputs[1]
.as_ref()
.expect_err("a body that never ends is a failure, not an end");
assert_eq!(error.code, "SANDBOX_OUTCOME_UNKNOWN", "got: {error}");
assert!(
!error.retryable,
"the command may still be there when the sandbox resumes, so a repeat could run it \
twice: {error}"
);
assert!(
error.to_string().contains("no terminal frame"),
"the refusal says what was missing: {error}"
);
}
#[tokio::test]
async fn a_refused_stream_releases_the_transport_when_it_reports() {
use tokio::io::{AsyncReadExt, AsyncWriteExt};
let listener = tokio::net::TcpListener::bind::<SocketAddr>("127.0.0.1:0".parse().unwrap())
.await
.expect("bind");
let address = listener.local_addr().expect("address");
let (release, released) = tokio::sync::oneshot::channel::<()>();
tokio::spawn(async move {
let (mut socket, _) = listener.accept().await.expect("accept");
let mut seen = Vec::new();
let mut chunk = [0u8; 4096];
while !seen.windows(4).any(|window| window == b"\r\n\r\n") {
let read = socket.read(&mut chunk).await.expect("read the request");
assert!(read > 0, "the client left before finishing its request");
seen.extend_from_slice(&chunk[..read]);
}
let frame = "{\"t\":\"stdout\",\"seq\":0,\"data\":\"aGk=\"}\n";
socket
.write_all(
format!(
"HTTP/1.1 200 OK\r\nContent-Type: application/x-ndjson\r\n\
Transfer-Encoding: chunked\r\n\r\n{:x}\r\n{frame}\r\n",
frame.len()
)
.as_bytes(),
)
.await
.expect("write one frame");
socket.flush().await.expect("flush");
while socket.read(&mut chunk).await.unwrap_or(0) > 0 {}
release.send(()).expect("the test is still waiting");
});
let client = reqwest::Client::new();
let response = send(
client.post(format!("http://{address}/v1/exec")),
RUN_COMMAND,
)
.await
.expect("headers arrive at once");
let mut stream = frame_stream(response, "test-sandbox", TERMINAL_FRAME_GRACE);
stream
.next()
.await
.expect("the one frame the agent wrote")
.expect("a stdout frame");
let error = stream
.next()
.await
.expect("the stall is reported rather than parking the reader")
.expect_err("a body that never ends is a failure, not an end");
assert_eq!(error.code, "SANDBOX_OUTCOME_UNKNOWN", "got: {error}");
tokio::time::timeout(TERMINAL_FRAME_GRACE * 20, released)
.await
.expect("the refusal closes the connection to the agent")
.expect("the agent reports the release");
drop(stream);
drop(client);
}
fn async_stream_frames(
chunks: Vec<&'static str>,
) -> impl futures::Stream<Item = std::result::Result<&'static str, std::io::Error>> {
futures::stream::iter(chunks).then(|chunk| async move {
tokio::time::sleep(AGENT_RESPONSE_TIMEOUT * 2).await;
Ok(chunk)
})
}
async fn stalled_agent() -> String {
let handler = || async {
tokio::time::sleep(AGENT_RESPONSE_TIMEOUT * 20).await;
"late".into_response()
};
let router = Router::new().fallback(handler);
let listener = tokio::net::TcpListener::bind::<SocketAddr>("127.0.0.1:0".parse().unwrap())
.await
.expect("bind");
let address = listener.local_addr().expect("address");
tokio::spawn(async move { axum::serve(listener, router).await.expect("serve") });
format!("http://{address}")
}
#[tokio::test]
async fn a_dropped_run_command_connection_is_not_retryable() {
let listener = tokio::net::TcpListener::bind::<SocketAddr>("127.0.0.1:0".parse().unwrap())
.await
.expect("bind");
let address = listener.local_addr().expect("address");
tokio::spawn(async move {
loop {
if let Ok((socket, _)) = listener.accept().await {
drop(socket);
}
}
});
let error = send(
reqwest::Client::new().post(format!("http://{address}/v1/exec")),
RUN_COMMAND,
)
.await
.expect_err("a dropped connection is a refusal, not a response");
assert_eq!(error.code, "SANDBOX_OUTCOME_UNKNOWN", "got: {error}");
assert!(
!error.retryable,
"a command that may have started must not be retried: {error}"
);
assert!(
error.to_string().contains("may have taken effect"),
"the refusal says the outcome is unknown: {error}"
);
}
#[tokio::test]
async fn a_stalled_run_command_is_not_retryable() {
let base = stalled_agent().await;
let error = send(
reqwest::Client::new().post(format!("{base}/v1/exec")),
RUN_COMMAND,
)
.await
.expect_err("a stalled agent must be refused, not waited on");
assert_eq!(error.code, "SANDBOX_OUTCOME_UNKNOWN", "got: {error}");
assert!(
!error.retryable,
"a command with an unknown outcome must not be retried: {error}"
);
assert!(
error.to_string().contains("did not answer")
&& error.to_string().contains("may have taken effect"),
"the refusal says the agent stalled and the outcome is unknown: {error}"
);
}
#[tokio::test]
async fn a_stalled_file_operation_stays_retryable() {
let base = stalled_agent().await;
let error = send(
reqwest::Client::new().get(format!("{base}/v1/files")),
"sandbox.readFile",
)
.await
.expect_err("a stalled agent must be refused, not waited on");
assert_eq!(error.code, "SANDBOX_UNREACHABLE", "got: {error}");
assert!(
error.retryable,
"a stalled file read is safe to repeat: {error}"
);
assert!(
error.to_string().contains("did not answer"),
"the refusal says the agent stalled: {error}"
);
}
#[tokio::test]
async fn frames_decode_in_order_with_a_real_exit_code() {
let outputs = frames_from(vec![
"{\"t\":\"stdout\",\"seq\":0,\"data\":\"aGk=\"}\n",
"{\"t\":\"stderr\",\"seq\":1,\"data\":\"b29wcw==\"}\n",
"{\"t\":\"exit\",\"code\":7,\"truncated\":false}\n",
])
.await;
assert_eq!(outputs.len(), 3);
assert_eq!(
outputs[0].as_ref().expect("stdout"),
&CommandOutput::Stdout {
seq: 0,
data: b"hi".to_vec()
}
);
assert_eq!(
outputs[1].as_ref().expect("stderr"),
&CommandOutput::Stderr {
seq: 1,
data: b"oops".to_vec()
}
);
assert_eq!(
outputs[2].as_ref().expect("exit"),
&CommandOutput::Exit {
code: 7,
truncated: false
}
);
}
#[tokio::test]
async fn a_frame_split_across_chunks_is_reassembled() {
let outputs = frames_from(vec![
"{\"t\":\"stdo",
"ut\",\"seq\":0,\"data\":\"aGk=\"}\n{\"t\":\"ex",
"it\",\"code\":0,\"truncated\":false}\n",
])
.await;
assert_eq!(
outputs.len(),
2,
"a split frame must not become two frames or an error"
);
assert_eq!(
outputs[0].as_ref().expect("stdout"),
&CommandOutput::Stdout {
seq: 0,
data: b"hi".to_vec()
}
);
assert_eq!(
outputs[1].as_ref().expect("exit"),
&CommandOutput::Exit {
code: 0,
truncated: false
}
);
}
#[tokio::test]
async fn a_frame_that_arrives_but_does_not_decode_leaves_the_outcome_unknown() {
let outputs = frames_from(vec![
"{\"t\":\"stdout\",\"seq\":0,\"data\":\"!!not base64!!\"}\n",
])
.await;
let error = outputs[0]
.as_ref()
.expect_err("a payload that does not decode is not output");
assert_eq!(error.code, "SANDBOX_OUTCOME_UNKNOWN", "got: {error}");
assert!(
!error.retryable,
"the command ran; a repeat would run it twice: {error}"
);
assert!(
error.to_string().contains("base64"),
"the decode failure must stay in the chain: {error}"
);
}
#[tokio::test]
async fn a_frame_that_does_not_parse_leaves_the_outcome_unknown() {
let outputs = frames_from(vec!["{not json at all}\n"]).await;
let error = outputs[0]
.as_ref()
.expect_err("a malformed frame is not output");
assert_eq!(error.code, "SANDBOX_OUTCOME_UNKNOWN", "got: {error}");
assert!(!error.retryable, "the command ran: {error}");
assert!(
error.to_string().contains("did not parse"),
"the parse failure must stay in the chain: {error}"
);
}
#[tokio::test]
async fn an_unestablished_outcome_ends_the_stream_mid_chunk() {
let outputs = frames_from(vec![
"{not json at all}\n{\"t\":\"exit\",\"code\":0,\"truncated\":false}\n",
])
.await;
assert_eq!(
outputs.len(),
1,
"the exit frame must not follow the failure"
);
let error = outputs[0]
.as_ref()
.expect_err("a malformed frame is not output");
assert_eq!(error.code, "SANDBOX_OUTCOME_UNKNOWN", "got: {error}");
}
#[tokio::test]
async fn a_decode_failure_ends_the_stream_mid_chunk() {
let outputs = frames_from(vec![
"{\"t\":\"stdout\",\"seq\":0,\"data\":\"!!\"}\n{\"t\":\"exit\",\"code\":0,\"truncated\":false}\n",
])
.await;
assert_eq!(
outputs.len(),
1,
"the exit frame must not follow the failure"
);
assert_eq!(
outputs[0]
.as_ref()
.expect_err("a bad payload is not output")
.code,
"SANDBOX_OUTCOME_UNKNOWN"
);
}
#[tokio::test]
async fn a_stream_without_a_terminal_frame_is_an_unknown_outcome() {
let outputs = frames_from(vec!["{\"t\":\"stdout\",\"seq\":0,\"data\":\"aGk=\"}\n"]).await;
assert_eq!(outputs.len(), 2);
outputs[0].as_ref().expect("the stdout frame still arrives");
let error = outputs[1]
.as_ref()
.expect_err("a truncated stream must not read as success");
assert!(
error.to_string().contains("without a terminal frame"),
"the failure must name the cause: {error}"
);
assert_eq!(error.code, "SANDBOX_OUTCOME_UNKNOWN", "got: {error}");
assert!(
!error.retryable,
"a command that started and whose end was lost must not be retried: {error}"
);
}
#[tokio::test]
async fn an_error_frame_surfaces_as_an_error_not_a_silent_end() {
let outputs = frames_from(vec![
"{\"t\":\"error\",\"code\":\"timeoutExceeded\",\"message\":\"exceeded its 300ms timeout\"}\n",
])
.await;
assert_eq!(outputs.len(), 1);
let error = outputs[0]
.as_ref()
.expect_err("an error frame is a failure");
assert!(error.to_string().contains("timeoutExceeded"), "{error}");
}
#[derive(Debug)]
struct TestTransport(String);
#[async_trait]
impl AgentTransport for TestTransport {
async fn request(
&self,
_sandbox_id: &str,
method: reqwest::Method,
path: &str,
) -> Result<reqwest::RequestBuilder> {
Ok(reqwest::Client::new().request(method, format!("{}{path}", self.0)))
}
fn provider(&self) -> &'static str {
"test-sandbox"
}
}
async fn serve(router: Router) -> TestTransport {
let listener = tokio::net::TcpListener::bind::<SocketAddr>("127.0.0.1:0".parse().unwrap())
.await
.expect("bind");
let address = listener.local_addr().expect("address");
tokio::spawn(async move { axum::serve(listener, router).await.expect("serve") });
TestTransport(format!("http://{address}"))
}
type Cursors = Arc<Mutex<Vec<Option<u64>>>>;
async fn job_agent(cursors: Cursors) -> TestTransport {
async fn poll(
State(cursors): State<Cursors>,
Json(body): Json<serde_json::Value>,
) -> impl IntoResponse {
let since = body.get("sinceSeq").and_then(serde_json::Value::as_u64);
cursors.lock().expect("cursors").push(since);
match since {
None => Json(json!({
"running": true,
"frames": [{ "t": "stdout", "seq": 0, "data": "aGk=" }],
})),
Some(_) => Json(json!({
"running": false,
"frames": [],
"exitCode": 7,
"truncated": false,
})),
}
}
serve(
Router::new()
.route(
"/v1/jobs/start",
post(|| async { Json(json!({"jobId": "j1"})) }),
)
.route("/v1/jobs/poll", post(poll))
.route("/v1/jobs/cancel", post(|| async { Json(json!({})) }))
.with_state(cursors),
)
.await
}
#[tokio::test]
async fn a_job_starts_polls_from_its_cursor_and_cancels() {
let cursors: Cursors = Arc::new(Mutex::new(Vec::new()));
let transport = job_agent(Arc::clone(&cursors)).await;
let started = start_job(&transport, "s1", command(Duration::from_secs(600)))
.await
.expect("the job starts");
assert_eq!(started.job_id, "j1");
let first = poll_job(&transport, "s1", &started.job_id, None)
.await
.expect("the first poll answers");
assert!(first.running, "the job is still running: {first:?}");
assert_eq!(
first.frames,
vec![CommandOutput::Stdout {
seq: 0,
data: b"hi".to_vec()
}],
"the agent's base64 frame is decoded"
);
assert!(first.exit.is_none() && first.error.is_none());
let second = poll_job(&transport, "s1", &started.job_id, Some(0))
.await
.expect("the second poll answers");
assert!(!second.running);
assert!(second.frames.is_empty(), "nothing follows the last frame");
assert_eq!(
second.exit,
Some(crate::traits::JobExit {
code: 7,
truncated: false
}),
"the ending is the envelope's, not a frame's"
);
cancel_job(&transport, "s1", &started.job_id)
.await
.expect("the cancel is accepted");
assert_eq!(
*cursors.lock().expect("cursors"),
vec![None, Some(0)],
"the cursor a caller passes has to reach the agent, or every poll replays the whole \
output and a caller sees each frame twice"
);
}
#[tokio::test]
async fn a_server_error_on_a_start_leaves_the_outcome_unknown() {
let transport = serve(Router::new().route(
"/v1/jobs/start",
post(|| async { (StatusCode::INTERNAL_SERVER_ERROR, "spawn failed").into_response() }),
))
.await;
let error = start_job(&transport, "s1", command(Duration::from_secs(600)))
.await
.expect_err("a 5xx is not a job that started");
assert_eq!(error.code, "SANDBOX_OUTCOME_UNKNOWN", "got: {error}");
assert!(
!error.retryable,
"the agent may have taken the command, and a repeat would run it twice: {error}"
);
assert!(
error.to_string().contains("sandbox.jobStart"),
"the operation reaches callers and telemetry, so a failed start has to be tellable \
apart from a failed streaming command: {error}"
);
}
#[tokio::test]
async fn a_server_error_on_a_poll_stays_retryable() {
let transport = serve(Router::new().route(
"/v1/jobs/poll",
post(|| async { (StatusCode::SERVICE_UNAVAILABLE, "").into_response() }),
))
.await;
let error = poll_job(&transport, "s1", "j1", Some(4))
.await
.expect_err("a 5xx is not a poll that answered");
assert_eq!(error.code, "SANDBOX_UNREACHABLE", "got: {error}");
assert!(error.retryable, "the job is untouched: {error}");
}
#[tokio::test]
async fn a_poll_body_that_ends_early_is_retryable() {
let transport = serve(Router::new().route(
"/v1/jobs/poll",
post(|| async {
axum::body::Body::from_stream(futures::stream::iter(vec![
Ok::<_, std::io::Error>(bytes::Bytes::from_static(b"{\"running\":tr")),
Err(std::io::Error::other("the connection went")),
]))
.into_response()
}),
))
.await;
let error = poll_job(&transport, "s1", "j1", None)
.await
.expect_err("a truncated body is not a poll");
assert_eq!(error.code, "SANDBOX_UNREACHABLE", "got: {error}");
assert!(error.retryable, "polling again costs nothing: {error}");
}
#[tokio::test]
async fn a_refused_cancel_is_the_agent_answering() {
let transport = serve(Router::new().route(
"/v1/jobs/cancel",
post(|| async { (StatusCode::NOT_FOUND, "JOB_NOT_FOUND").into_response() }),
))
.await;
let error = cancel_job(&transport, "s1", "j1")
.await
.expect_err("a 404 is not a cancel that landed");
assert_eq!(error.code, "SANDBOX_COMMAND_FAILED", "got: {error}");
assert!(
error.to_string().contains("JOB_NOT_FOUND"),
"the agent's own reason has to survive: {error}"
);
}
#[test]
fn the_program_leads_its_arguments_in_the_body_the_agent_receives() {
let body = exec_body(&RunCommandRequest {
command: "python".to_string(),
args: vec!["-u".to_string(), "main.py".to_string()],
cwd: Some("/work".to_string()),
env: BTreeMap::from([("TOKEN".to_string(), "t".to_string())]),
timeout: Duration::from_millis(5_000),
})
.expect("a command with a non-zero timeout builds a body");
assert_eq!(
body["command"],
serde_json::json!(["python", "-u", "main.py"]),
"the program leads, and its arguments follow in order"
);
assert_eq!(body["cwd"], serde_json::json!("/work"));
assert_eq!(body["timeoutMs"], serde_json::json!(5_000));
assert_eq!(body["env"], serde_json::json!({ "TOKEN": "t" }));
}
fn command(timeout: Duration) -> RunCommandRequest {
RunCommandRequest {
command: "/bin/sleep".to_string(),
args: vec!["600".to_string()],
cwd: None,
env: BTreeMap::new(),
timeout,
}
}
}