diff --git a/src/main/java/com/volcengine/ark/runtime/selfhosted/DefaultTools.java b/src/main/java/com/volcengine/ark/runtime/selfhosted/DefaultTools.java index dfd2b47..ce3ca82 100644 --- a/src/main/java/com/volcengine/ark/runtime/selfhosted/DefaultTools.java +++ b/src/main/java/com/volcengine/ark/runtime/selfhosted/DefaultTools.java @@ -81,10 +81,15 @@ public ToolResult execute(Object input, ToolContext context) { : new LinkedHashMap<>(System.getenv()); builder.environment().putAll(scrubbedEnv(environment)); builder.redirectErrorStream(true); + Process process = null; + Thread reader = null; try { - Process process = builder.start(); + process = builder.start(); BoundedOutput output = new BoundedOutput(MAX_OUTPUT_BYTES); - Thread reader = new Thread(() -> drainOutput(process.getInputStream(), output), "ma-self-host-bash-output"); + Process runningProcess = process; + reader = new Thread( + () -> drainOutput(runningProcess.getInputStream(), output), + "ma-self-host-bash-output"); reader.setDaemon(true); reader.start(); long deadline = System.nanoTime() + TimeUnit.MILLISECONDS.toNanos(Math.max(0L, context.getToolTimeoutMillis())); @@ -117,12 +122,46 @@ public ToolResult execute(Object input, ToolContext context) { return ToolResult.error("exit code " + process.exitValue() + "\n" + text); } return ToolResult.text(text); + } catch (InterruptedException e) { + // SessionToolRunner interrupts the execution future when user.interrupt + // arrives. Do not let that interruption orphan the shell: the shell may + // otherwise resume the remaining compound command and mutate the workdir + // after the call has already been marked interrupted. + terminateProcess(process); + joinReader(reader); + Thread.currentThread().interrupt(); + return ToolResult.error("tool execution canceled"); } catch (Exception e) { + terminateProcess(process); + joinReader(reader); return ToolResult.error(e.getMessage()); } } } + private static void terminateProcess(Process process) { + if (process == null || !process.isAlive()) { + return; + } + process.destroyForcibly(); + try { + process.waitFor(PROCESS_TERMINATION_GRACE_MILLIS, TimeUnit.MILLISECONDS); + } catch (InterruptedException ignored) { + Thread.currentThread().interrupt(); + } + } + + private static void joinReader(Thread reader) { + if (reader == null) { + return; + } + try { + reader.join(PROCESS_TERMINATION_GRACE_MILLIS); + } catch (InterruptedException ignored) { + Thread.currentThread().interrupt(); + } + } + static class ReadTool implements Tool { @Override public String name() { diff --git a/src/test/java/com/volcengine/ark/runtime/selfhosted/DefaultToolsTest.java b/src/test/java/com/volcengine/ark/runtime/selfhosted/DefaultToolsTest.java index 69aa31f..063f01c 100644 --- a/src/test/java/com/volcengine/ark/runtime/selfhosted/DefaultToolsTest.java +++ b/src/test/java/com/volcengine/ark/runtime/selfhosted/DefaultToolsTest.java @@ -14,6 +14,8 @@ import java.util.LinkedHashMap; import java.util.Map; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicReference; import org.junit.Test; public class DefaultToolsTest { @@ -214,6 +216,38 @@ public void bashTimeoutReturnsPromptly() throws Exception { assertTrue(result.getContent().get(0).getText().contains("timed out")); } + @Test + public void bashInterruptionTerminatesCompoundCommand() throws Exception { + Path workdir = Files.createTempDirectory("ark-java-cancel-"); + Path started = workdir.resolve("started"); + Path completed = workdir.resolve("completed"); + AtomicBoolean canceled = new AtomicBoolean(); + ToolContext context = new ToolContext(workdir.toString()); + context.setCancelled(canceled::get); + Map input = new LinkedHashMap<>(); + input.put("command", "printf started > started; sleep 1; printf completed > completed"); + AtomicReference result = new AtomicReference<>(); + Thread execution = new Thread( + () -> result.set(new DefaultTools.BashTool().execute(input, context)), + "bash-cancellation-test"); + + execution.start(); + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(2L); + while (!Files.exists(started) && System.nanoTime() < deadline) { + Thread.sleep(10L); + } + assertTrue("bash command did not start", Files.exists(started)); + canceled.set(true); + execution.interrupt(); + execution.join(3000L); + + assertFalse("bash execution did not return after cancellation", execution.isAlive()); + assertTrue(result.get().isError()); + assertTrue(result.get().getContent().get(0).getText().contains("canceled")); + Thread.sleep(1200L); + assertFalse("compound command continued after interruption", Files.exists(completed)); + } + @Test public void fileToolRejectsSymlinkEscape() throws Exception { Path workdir = Files.createTempDirectory("ark-java-root-");