Skip to content

fix(sandbox-e2b): reliable exec timeout with server-side kill #2976

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
chcodex wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix(sandbox-e2b): reliable exec timeout with server-side kill #2976

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 24 additions & 6 deletions ...-sandbox-e2b/src/main/java/io/agentscope/extensions/sandbox/e2b/E2bEnvdProcessClient.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -119,10 +119,19 @@ byte[] runShellBinaryStdout(
private ShellCapture runShellCapture(
E2bSandboxState state, String cwd, String shellCommand, int timeoutSeconds)
throws Exception {
if (timeoutSeconds <= 0) {
throw new IllegalArgumentException(
"timeoutSeconds must be positive: " + timeoutSeconds);
}
OkHttpClient callClient =
timeoutSeconds > 0
? http.newBuilder().callTimeout(timeoutSeconds, TimeUnit.SECONDS).build()
: http;
http.newBuilder()
.callTimeout(timeoutSeconds, TimeUnit.SECONDS)
// Disable the idle read timeout: it fires on gaps between bytes and
// would preempt callTimeout (both surface as InterruptedIOException),
// misreporting a short idle stall as a full exec timeout (#2974).
// Total-duration semantics is owned solely by callTimeout.
.readTimeout(0, TimeUnit.SECONDS)
.build();
String host = envdHost(state);
String url = host + "/process.Process/Start";
byte[] envelope = encodeStartRequestEnvelope(shellCommand, cwd);
Expand All @@ -134,7 +143,10 @@ private ShellCapture runShellCapture(
.addHeader("User-Agent", "agentscope-java-e2b")
.addHeader("E2b-Sandbox-Id", state.getSandboxId())
.addHeader("E2b-Sandbox-Port", Integer.toString(ENVD_PORT))
.addHeader("Authorization", basicAuthUser(opt.getRunUser()));
.addHeader("Authorization", basicAuthUser(opt.getRunUser()))
// Ask envd to kill the remote process when this lapses; otherwise a
// client-side timeout only drops our HTTP stream and leaks the process.
.addHeader("Connect-Timeout-Ms", String.valueOf(timeoutSeconds * 1000L));
if (state.getEnvdAccessToken() != null && !state.getEnvdAccessToken().isBlank()) {
rb.addHeader("X-Access-Token", state.getEnvdAccessToken());
}
Expand All @@ -145,7 +157,7 @@ private ShellCapture runShellCapture(
int exit;
try (Response res = callClient.newCall(req).execute()) {
if (!res.isSuccessful()) {
String err = res.body() != null ? res.body().string() : "";
String err = res.body().string();
throw new SandboxException.SandboxRuntimeException(
SandboxErrorCode.WORKSPACE_START_ERROR,
"envd Start failed HTTP " + res.code() + ": " + err);
Expand All @@ -154,6 +166,12 @@ private ShellCapture runShellCapture(
exit = drainStartStream(in, stdout, stderr);
}
} catch (InterruptedIOException e) {
// External cancellation surfaces here too; don't misreport it as a timeout
// and don't swallow the interrupt bit.
if (Thread.currentThread().isInterrupted()) {
Thread.currentThread().interrupt();
throw e;
}
throw new SandboxException.ExecTimeoutException(shellCommand, timeoutSeconds);
}
return new ShellCapture(exit, stdout, stderr);
Expand All @@ -179,7 +197,7 @@ private int drainStartStream(
break;
}
int len = ByteBuffer.wrap(lenB).order(ByteOrder.BIG_ENDIAN).getInt() & 0x7FFFFFFF;
if (len < 0 || len > 64 * 1024 * 1024) {
if (len > 64 * 1024 * 1024) {
throw new IOException("Invalid connect frame length: " + len);
}
byte[] data = in.readNBytes(len);
Expand Down
77 changes: 77 additions & 0 deletions ...dbox-e2b/src/test/java/io/agentscope/extensions/sandbox/e2b/E2bEnvdProcessClientTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,23 +17,30 @@

import static org.junit.jupiter.api.Assertions.assertArrayEquals;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertNull;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;

import com.google.protobuf.ByteString;
import com.google.protobuf.Descriptors;
import com.google.protobuf.DynamicMessage;
import io.agentscope.harness.agent.sandbox.SandboxException;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.net.SocketTimeoutException;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.nio.charset.StandardCharsets;
import java.util.Base64;
import java.util.concurrent.atomic.AtomicReference;
import okhttp3.Interceptor;
import okhttp3.MediaType;
import okhttp3.OkHttpClient;
import org.junit.jupiter.api.Test;

class E2bEnvdProcessClientTest {
Expand Down Expand Up @@ -307,4 +314,74 @@ private static E2bSandboxClientOptions options(E2bCodec codec) {
options.setCodec(codec);
return options;
}

@Test
void requestCarriesConnectTimeoutMsHeader() throws Exception {
AtomicReference<String> header = new AtomicReference<>();
Interceptor capture =
chain -> {
header.set(chain.request().header("Connect-Timeout-Ms"));
throw new SocketTimeoutException("timeout");
};
E2bEnvdProcessClient client = clientWithInterceptor(capture);

assertThrows(
SandboxException.ExecTimeoutException.class,
() -> client.runShell(state(), "/workspace", "sleep 1000", 3));
assertEquals("3000", header.get());
}

@Test
void interruptionIsRethrownNotWrappedAsTimeout() throws Exception {
Interceptor interrupting =
chain -> {
Thread.currentThread().interrupt();
throw new SocketTimeoutException("read timed out");
};
E2bEnvdProcessClient client = clientWithInterceptor(interrupting);
try {
SocketTimeoutException thrown =
assertThrows(
SocketTimeoutException.class,
() -> client.runShell(state(), "/workspace", "sleep 1000", 3));
assertEquals("read timed out", thrown.getMessage());
assertTrue(Thread.currentThread().isInterrupted(), "interrupt bit must be restored");
} finally {
// Do not leak the interrupt bit into other tests.
Thread.interrupted();
}
assertFalse(Thread.currentThread().isInterrupted());
}

@Test
void nonPositiveTimeoutFailsFastWithoutRequest() throws Exception {
AtomicReference<String> header = new AtomicReference<>();
Interceptor capture =
chain -> {
header.set(chain.request().header("Connect-Timeout-Ms"));
throw new SocketTimeoutException("must not be called");
};
E2bEnvdProcessClient client = clientWithInterceptor(capture);

assertThrows(
IllegalArgumentException.class,
() -> client.runShell(state(), "/workspace", "echo hi", 0));
assertThrows(
IllegalArgumentException.class,
() -> client.runShell(state(), "/workspace", "echo hi", -5));
assertNull(header.get(), "no HTTP request must be issued for invalid timeout");
}

private static E2bEnvdProcessClient clientWithInterceptor(Interceptor interceptor)
throws Exception {
E2bSandboxClientOptions opt = options(E2bCodec.PROTO);
opt.setHttpClient(new OkHttpClient.Builder().addInterceptor(interceptor).build());
return new E2bEnvdProcessClient(opt);
}

private static E2bSandboxState state() {
E2bSandboxState state = new E2bSandboxState();
state.setSandboxId("test-sandbox");
return state;
}
}
Loading