diff --git a/agentscope-extensions/agentscope-extensions-sandbox/agentscope-extensions-sandbox-e2b/src/main/java/io/agentscope/extensions/sandbox/e2b/E2bEnvdProcessClient.java b/agentscope-extensions/agentscope-extensions-sandbox/agentscope-extensions-sandbox-e2b/src/main/java/io/agentscope/extensions/sandbox/e2b/E2bEnvdProcessClient.java index d8fb0b9dc0..849056983a 100644 --- a/agentscope-extensions/agentscope-extensions-sandbox/agentscope-extensions-sandbox-e2b/src/main/java/io/agentscope/extensions/sandbox/e2b/E2bEnvdProcessClient.java +++ b/agentscope-extensions/agentscope-extensions-sandbox/agentscope-extensions-sandbox-e2b/src/main/java/io/agentscope/extensions/sandbox/e2b/E2bEnvdProcessClient.java @@ -142,7 +142,7 @@ private ShellCapture runShellCapture( ByteArrayOutputStream stdout = new ByteArrayOutputStream(); ByteArrayOutputStream stderr = new ByteArrayOutputStream(); - int exit = Integer.MIN_VALUE; + int exit; try (Response res = callClient.newCall(req).execute()) { if (!res.isSuccessful()) { String err = res.body() != null ? res.body().string() : ""; @@ -165,7 +165,7 @@ private record ShellCapture( private int drainStartStream( InputStream in, ByteArrayOutputStream stdout, ByteArrayOutputStream stderr) throws IOException { - int exit = Integer.MIN_VALUE; + Integer exit = null; Descriptors.FieldDescriptor srEventF = startResponseDesc.findFieldByName("event"); Descriptors.FieldDescriptor peDataF = processEventDesc.findFieldByName("data"); Descriptors.FieldDescriptor peEndF = processEventDesc.findFieldByName("end"); @@ -212,6 +212,9 @@ private int drainStartStream( continue; } } + if (exit == null) { + throw new IOException("envd process stream ended before receiving a process exit code"); + } return exit; } diff --git a/agentscope-extensions/agentscope-extensions-sandbox/agentscope-extensions-sandbox-e2b/src/test/java/io/agentscope/extensions/sandbox/e2b/E2bEnvdProcessClientTest.java b/agentscope-extensions/agentscope-extensions-sandbox/agentscope-extensions-sandbox-e2b/src/test/java/io/agentscope/extensions/sandbox/e2b/E2bEnvdProcessClientTest.java index 95d37d8dd5..3d2c19dfb7 100644 --- a/agentscope-extensions/agentscope-extensions-sandbox/agentscope-extensions-sandbox-e2b/src/test/java/io/agentscope/extensions/sandbox/e2b/E2bEnvdProcessClientTest.java +++ b/agentscope-extensions/agentscope-extensions-sandbox/agentscope-extensions-sandbox-e2b/src/test/java/io/agentscope/extensions/sandbox/e2b/E2bEnvdProcessClientTest.java @@ -17,6 +17,7 @@ import static org.junit.jupiter.api.Assertions.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import com.google.protobuf.ByteString; @@ -24,7 +25,9 @@ import com.google.protobuf.DynamicMessage; 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.nio.ByteBuffer; import java.nio.ByteOrder; @@ -116,24 +119,67 @@ void jsonCodecSkipsMalformedBase64AndKeepsStreaming() throws Exception { } @Test - void jsonCodecReturnsSentinelWhenEndMissing() throws Exception { + void jsonCodecRejectsEofBeforeEnd() throws Exception { E2bEnvdProcessClient client = new E2bEnvdProcessClient(options(E2bCodec.JSON)); ByteArrayOutputStream stdout = new ByteArrayOutputStream(); ByteArrayOutputStream stderr = new ByteArrayOutputStream(); - int exit = - drainStartStream( - client, - connectFrames( - responseJson(base64("hello\n"), null, null), - responseJson(null, base64("warn\n"), null)), - stdout, - stderr); + IOException exception = + assertThrows( + IOException.class, + () -> + drainStartStream( + client, + connectFrames( + responseJson(base64("hello\n"), null, null), + responseJson(null, base64("warn\n"), null)), + stdout, + stderr)); - assertEquals(Integer.MIN_VALUE, exit); + assertTrue(exception.getMessage().contains("before receiving a process exit code")); assertEquals("hello\n", stdout.toString(StandardCharsets.UTF_8)); assertEquals("warn\n", stderr.toString(StandardCharsets.UTF_8)); } + @Test + void jsonCodecRejectsTruncatedFrameLengthBeforeEnd() throws Exception { + E2bEnvdProcessClient client = new E2bEnvdProcessClient(options(E2bCodec.JSON)); + byte[] truncatedLength = new byte[] {0x00, 0x00, 0x00}; + + IOException exception = + assertThrows( + IOException.class, + () -> + drainStartStream( + client, + truncatedLength, + new ByteArrayOutputStream(), + new ByteArrayOutputStream())); + + assertTrue(exception.getMessage().contains("before receiving a process exit code")); + } + + @Test + void jsonCodecRejectsTruncatedFramePayloadBeforeEnd() throws Exception { + E2bEnvdProcessClient client = new E2bEnvdProcessClient(options(E2bCodec.JSON)); + ByteBuffer truncatedPayload = ByteBuffer.allocate(7).order(ByteOrder.BIG_ENDIAN); + truncatedPayload.put((byte) 0x00); + truncatedPayload.putInt(4); + truncatedPayload.put((byte) '{'); + truncatedPayload.put((byte) '}'); + + IOException exception = + assertThrows( + IOException.class, + () -> + drainStartStream( + client, + truncatedPayload.array(), + new ByteArrayOutputStream(), + new ByteArrayOutputStream())); + + assertTrue(exception.getMessage().contains("before receiving a process exit code")); + } + @Test void jsonCodecReturnsEmptyOutputsWhenOnlyEndPresent() throws Exception { E2bEnvdProcessClient client = new E2bEnvdProcessClient(options(E2bCodec.JSON)); @@ -195,7 +241,19 @@ private static int drainStartStream( ByteArrayOutputStream.class, ByteArrayOutputStream.class); method.setAccessible(true); - return (int) method.invoke(client, new ByteArrayInputStream(connectFrame), stdout, stderr); + try { + return (int) + method.invoke(client, new ByteArrayInputStream(connectFrame), stdout, stderr); + } catch (InvocationTargetException e) { + Throwable cause = e.getCause(); + if (cause instanceof Exception exception) { + throw exception; + } + if (cause instanceof Error error) { + throw error; + } + throw e; + } } private static byte[] connectFrame(String json) { diff --git a/agentscope-extensions/agentscope-spring-boot-starters/agentscope-agui-spring-boot-starter/src/test/java/io/agentscope/spring/boot/agui/mvc/AguiMvcControllerTest.java b/agentscope-extensions/agentscope-spring-boot-starters/agentscope-agui-spring-boot-starter/src/test/java/io/agentscope/spring/boot/agui/mvc/AguiMvcControllerTest.java index fafcadfa04..be5cb5d1e9 100644 --- a/agentscope-extensions/agentscope-spring-boot-starters/agentscope-agui-spring-boot-starter/src/test/java/io/agentscope/spring/boot/agui/mvc/AguiMvcControllerTest.java +++ b/agentscope-extensions/agentscope-spring-boot-starters/agentscope-agui-spring-boot-starter/src/test/java/io/agentscope/spring/boot/agui/mvc/AguiMvcControllerTest.java @@ -85,6 +85,7 @@ void sseErrorCancelsSubscriptionAndReleasesThread() throws Exception { assertTrue(fixture.firstRunSubscribed.await(5, TimeUnit.SECONDS)); invokeError(emitter); + assertTrue(fixture.firstRunTerminated.await(5, TimeUnit.SECONDS)); awaitSecondRunAccepted(fixture); @@ -132,6 +133,7 @@ void sseTimeoutCancelsSubscriptionAndReleasesThread() throws Exception { Object timeoutCallback = ReflectionTestUtils.getField(emitter, "timeoutCallback"); ReflectionTestUtils.invokeMethod(timeoutCallback, "run"); + assertTrue(fixture.firstRunTerminated.await(5, TimeUnit.SECONDS)); awaitSecondRunAccepted(fixture); @@ -250,8 +252,9 @@ private TestAdapter( @Override public Flux run(RunAgentInput input, RuntimeContext runtimeContext) { if (runCount.incrementAndGet() == 1) { - firstRunSubscribed.countDown(); - return firstRunEvents.doFinally(signalType -> firstRunTerminated.countDown()); + return firstRunEvents + .doOnRequest(ignored -> firstRunSubscribed.countDown()) + .doFinally(signalType -> firstRunTerminated.countDown()); } return Flux.empty(); }