|
9 | 9 | import java.util.List; |
10 | 10 | import java.util.Map; |
11 | 11 | import java.util.Random; |
| 12 | +import java.util.concurrent.CancellationException; |
| 13 | +import java.util.concurrent.ExecutionException; |
| 14 | +import java.util.concurrent.ExecutorService; |
| 15 | +import java.util.concurrent.Executors; |
| 16 | +import java.util.concurrent.Future; |
12 | 17 | import java.util.concurrent.LinkedBlockingQueue; |
| 18 | +import java.util.concurrent.RejectedExecutionException; |
13 | 19 | import java.util.concurrent.TimeUnit; |
| 20 | +import java.util.concurrent.TimeoutException; |
| 21 | +import java.util.concurrent.atomic.AtomicBoolean; |
| 22 | +import java.util.concurrent.atomic.AtomicInteger; |
14 | 23 | import java.util.concurrent.atomic.AtomicReference; |
15 | 24 | import java.util.function.BooleanSupplier; |
16 | 25 | import java.util.logging.Level; |
17 | 26 | import java.util.logging.Logger; |
18 | 27 |
|
19 | 28 | public class SessionToolRunner { |
20 | 29 | private static final int STREAM_QUEUE_SIZE = 256; |
| 30 | + private static final long TOOL_WAIT_SLICE_MILLIS = 50L; |
21 | 31 | private static final Logger LOGGER = Logger.getLogger(SessionToolRunner.class.getName()); |
| 32 | + private static final AtomicInteger TOOL_THREAD_ID = new AtomicInteger(); |
22 | 33 | private final SelfHostedClient api; |
23 | 34 | private final String sessionId; |
24 | 35 | private final Options options; |
25 | 36 | private final State state = new State(); |
26 | 37 | private final List<ToolCallResult> results = new ArrayList<>(); |
27 | 38 | private final Random random = new Random(); |
| 39 | + private final ExecutorService toolExecutor = Executors.newCachedThreadPool(runnable -> { |
| 40 | + Thread thread = new Thread(runnable, "ma-self-host-tool-" + TOOL_THREAD_ID.incrementAndGet()); |
| 41 | + thread.setDaemon(true); |
| 42 | + return thread; |
| 43 | + }); |
28 | 44 | private volatile boolean closed; |
29 | 45 | private volatile EventStream activeStream; |
30 | 46 |
|
@@ -198,6 +214,7 @@ private void consumeList() throws IOException { |
198 | 214 | public void close() { |
199 | 215 | closed = true; |
200 | 216 | closeStream(activeStream); |
| 217 | + toolExecutor.shutdownNow(); |
201 | 218 | } |
202 | 219 |
|
203 | 220 | public List<ToolCallResult> getResults() { |
@@ -341,14 +358,75 @@ private void handleToolUse(Event event, boolean custom) throws IOException { |
341 | 358 | } |
342 | 359 |
|
343 | 360 | private ToolResult executeTool(Event event, boolean custom) { |
| 361 | + long timeoutMillis = options.toolContext.getToolTimeoutMillis(); |
| 362 | + if (timeoutMillis <= 0L) { |
| 363 | + timeoutMillis = SelfHostedConstants.DEFAULT_TOOL_TIMEOUT_MILLIS; |
| 364 | + } |
| 365 | + AtomicBoolean executionCancelled = new AtomicBoolean(); |
| 366 | + ToolContext context = toolContextForExecution(timeoutMillis, executionCancelled); |
| 367 | + Future<ToolResult> future; |
| 368 | + try { |
| 369 | + future = toolExecutor.submit(() -> executeTool(event, custom, context)); |
| 370 | + } catch (RejectedExecutionException error) { |
| 371 | + return ToolResult.error("tool execution canceled"); |
| 372 | + } |
| 373 | + long deadline = System.nanoTime() + TimeUnit.MILLISECONDS.toNanos(timeoutMillis); |
| 374 | + while (true) { |
| 375 | + if (context.isCancelled()) { |
| 376 | + executionCancelled.set(true); |
| 377 | + future.cancel(true); |
| 378 | + return ToolResult.error("tool execution canceled"); |
| 379 | + } |
| 380 | + long remaining = deadline - System.nanoTime(); |
| 381 | + if (remaining <= 0L) { |
| 382 | + executionCancelled.set(true); |
| 383 | + future.cancel(true); |
| 384 | + return ToolResult.error("tool execution timed out after " + timeoutMillis + "ms"); |
| 385 | + } |
| 386 | + long waitNanos = Math.min(remaining, TimeUnit.MILLISECONDS.toNanos(TOOL_WAIT_SLICE_MILLIS)); |
| 387 | + try { |
| 388 | + return future.get(waitNanos, TimeUnit.NANOSECONDS); |
| 389 | + } catch (TimeoutException ignored) { |
| 390 | + // Continue so cancellation is observed without waiting for the full tool timeout. |
| 391 | + } catch (CancellationException error) { |
| 392 | + return ToolResult.error("tool execution canceled"); |
| 393 | + } catch (InterruptedException error) { |
| 394 | + Thread.currentThread().interrupt(); |
| 395 | + executionCancelled.set(true); |
| 396 | + future.cancel(true); |
| 397 | + return ToolResult.error("tool execution canceled"); |
| 398 | + } catch (ExecutionException error) { |
| 399 | + Throwable cause = error.getCause(); |
| 400 | + return ToolResult.error(cause == null ? error.toString() : errorText(cause)); |
| 401 | + } |
| 402 | + } |
| 403 | + } |
| 404 | + |
| 405 | + private ToolResult executeTool(Event event, boolean custom, ToolContext context) { |
344 | 406 | if (custom) { |
345 | 407 | try { |
346 | | - return options.customTools.get(event.getName()).execute(event.getInput(), options.toolContext); |
| 408 | + return options.customTools.get(event.getName()).execute(event.getInput(), context); |
347 | 409 | } catch (RuntimeException error) { |
348 | | - return ToolResult.error(error.getMessage()); |
| 410 | + return ToolResult.error(errorText(error)); |
349 | 411 | } |
350 | 412 | } |
351 | | - return options.tools.execute(event.getName(), event.getInput(), options.toolContext); |
| 413 | + return options.tools.execute(event.getName(), event.getInput(), context); |
| 414 | + } |
| 415 | + |
| 416 | + private ToolContext toolContextForExecution(long timeoutMillis, AtomicBoolean executionCancelled) { |
| 417 | + ToolContext source = options.toolContext; |
| 418 | + ToolContext context = new ToolContext(source.getWorkdir()); |
| 419 | + if (source.hasExplicitEnv()) { |
| 420 | + context.setEnv(new LinkedHashMap<>(source.getEnv())); |
| 421 | + } |
| 422 | + context.setUnrestrictedPaths(source.isUnrestrictedPaths()); |
| 423 | + context.setToolTimeoutMillis(timeoutMillis); |
| 424 | + context.setCancelled(() -> executionCancelled.get() || isClosed() || source.isCancelled()); |
| 425 | + return context; |
| 426 | + } |
| 427 | + |
| 428 | + private static String errorText(Throwable error) { |
| 429 | + return error.getMessage() == null ? error.toString() : error.getMessage(); |
352 | 430 | } |
353 | 431 |
|
354 | 432 | private void sendResult(String callId, Event source, boolean custom, String confirmation, Event out) throws IOException { |
|
0 commit comments