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 27ab32f..dfd2b47 100644 --- a/src/main/java/com/volcengine/ark/runtime/selfhosted/DefaultTools.java +++ b/src/main/java/com/volcengine/ark/runtime/selfhosted/DefaultTools.java @@ -11,6 +11,9 @@ import java.io.File; import java.io.IOException; import java.io.InputStream; +import java.nio.ByteBuffer; +import java.nio.charset.CharacterCodingException; +import java.nio.charset.CodingErrorAction; import java.nio.charset.StandardCharsets; import java.nio.file.AtomicMoveNotSupportedException; import java.nio.file.Files; @@ -19,9 +22,12 @@ import java.nio.file.Paths; import java.nio.file.StandardCopyOption; import java.nio.file.attribute.PosixFilePermission; +import java.util.ArrayList; +import java.util.Base64; import java.util.Collections; import java.util.Iterator; import java.util.LinkedHashMap; +import java.util.List; import java.util.Locale; import java.util.Map; import java.util.Set; @@ -33,6 +39,8 @@ public final class DefaultTools { private static final ObjectMapper MAPPER = ArkService.defaultObjectMapper(); private static final int MAX_OUTPUT_BYTES = 100000; + private static final int DEFAULT_READ_LINES = 2000; + private static final int MAX_READ_LINE_CHARS = 2000; private static final int MAX_SEARCH_MATCHES = 1000; private static final long PROCESS_TERMINATION_GRACE_MILLIS = 1000L; @@ -125,9 +133,68 @@ public String name() { public ToolResult execute(Object input, ToolContext context) { try { Map args = asMap(input); - Path path = safePath(context, firstNonEmpty(stringValue(args.get("path")), stringValue(args.get("file")))); - byte[] data = readBounded(path, MAX_OUTPUT_BYTES); - return ToolResult.text(new String(data, StandardCharsets.UTF_8)); + Path path = safePath(context, readFilePath(args)); + List viewRange = readViewRange(args.get("view_range")); + boolean hasViewRange = viewRange != null && !viewRange.isEmpty(); + Long offset = optionalLong(args, "offset"); + Long lineLimit = optionalLong(args, "limit"); + if (hasViewRange && (offset != null || lineLimit != null)) { + return ToolResult.error("view_range cannot be combined with offset or limit"); + } + MediaInfo media = detectReadMedia(path); + if (media != null) { + if (hasViewRange || offset != null || lineLimit != null) { + return ToolResult.error( + "view_range, offset, and limit are only supported for text files"); + } + long size = Files.size(path); + long mediaLimit = context.getMaxMediaFileBytes(); + if (mediaLimit == 0L) { + mediaLimit = context.getMaxInputFileBytes(); + } + if (mediaLimit > 0L && size > mediaLimit) { + return ToolResult.error("media file too large: " + size + " bytes"); + } + int readLimit = mediaLimit > 0L + ? (int) Math.min(mediaLimit + 1L, Integer.MAX_VALUE) + : Integer.MAX_VALUE; + byte[] data = mediaLimit > 0L + ? readBounded(path, readLimit) + : Files.readAllBytes(path); + if (mediaLimit > 0L && data.length > mediaLimit) { + return ToolResult.error("media file too large: " + Math.max(size, data.length) + " bytes"); + } + ContentBlock block = new ContentBlock(); + block.setType(media.blockType); + Map source = new LinkedHashMap<>(); + source.put("type", "base64"); + source.put("media_type", media.mediaType); + source.put("data", Base64.getEncoder().encodeToString(data)); + block.setSource(source); + return new ToolResult(Collections.singletonList(block), false); + } + long configuredLimit = context.getMaxInputFileBytes(); + int byteLimit = configuredLimit > 0L + ? (int) Math.min(configuredLimit, Integer.MAX_VALUE) + : MAX_OUTPUT_BYTES; + long size = Files.size(path); + if (byteLimit > 0 && size > byteLimit) { + return ToolResult.error("file too large: " + size + " bytes"); + } + String text; + try { + text = decodeUTF8(readBounded(path, byteLimit)); + } catch (CharacterCodingException e) { + return ToolResult.error("binary file cannot be read directly"); + } + if (hasViewRange) { + return ToolResult.text(renderViewRange(text, viewRange)); + } + if (offset != null && offset < 1L) { + return ToolResult.error( + "offset is the 1-based start line and must be >= 1, got " + offset); + } + return ToolResult.text(renderReadLines(text, offset, lineLimit)); } catch (Exception e) { return ToolResult.error(e.getMessage()); } @@ -298,6 +365,218 @@ static Map asMap(Object input) { return Collections.emptyMap(); } + private static MediaInfo detectReadMedia(Path path) throws IOException { + byte[] header = readBounded(path, 512); + if (startsWith(header, new byte[] {(byte) 0xff, (byte) 0xd8, (byte) 0xff})) { + return new MediaInfo("image", "image/jpeg"); + } + if (startsWith(header, new byte[] {(byte) 0x89, 'P', 'N', 'G', '\r', '\n', (byte) 0x1a, '\n'})) { + return new MediaInfo("image", "image/png"); + } + if (startsWith(header, "GIF87a".getBytes(StandardCharsets.US_ASCII)) + || startsWith(header, "GIF89a".getBytes(StandardCharsets.US_ASCII))) { + return new MediaInfo("image", "image/gif"); + } + if (header.length >= 12 + && startsWith(header, "RIFF".getBytes(StandardCharsets.US_ASCII)) + && matchesAt(header, 8, "WEBP".getBytes(StandardCharsets.US_ASCII))) { + return new MediaInfo("image", "image/webp"); + } + if (startsWith(header, "%PDF-".getBytes(StandardCharsets.US_ASCII))) { + return new MediaInfo("document", "application/pdf"); + } + if (!looksBinary(header)) { + return null; + } + String name = path.getFileName().toString().toLowerCase(Locale.ROOT); + if (name.endsWith(".jpg") || name.endsWith(".jpeg")) { + return new MediaInfo("image", "image/jpeg"); + } + if (name.endsWith(".png")) { + return new MediaInfo("image", "image/png"); + } + if (name.endsWith(".gif")) { + return new MediaInfo("image", "image/gif"); + } + if (name.endsWith(".webp")) { + return new MediaInfo("image", "image/webp"); + } + if (name.endsWith(".pdf")) { + return new MediaInfo("document", "application/pdf"); + } + return null; + } + + private static boolean startsWith(byte[] value, byte[] prefix) { + return matchesAt(value, 0, prefix); + } + + private static boolean matchesAt(byte[] value, int offset, byte[] expected) { + if (value.length - offset < expected.length) { + return false; + } + for (int index = 0; index < expected.length; index++) { + if (value[offset + index] != expected[index]) { + return false; + } + } + return true; + } + + private static boolean looksBinary(byte[] value) { + for (byte item : value) { + if (item == 0) { + return true; + } + } + try { + StandardCharsets.UTF_8 + .newDecoder() + .onMalformedInput(CodingErrorAction.REPORT) + .onUnmappableCharacter(CodingErrorAction.REPORT) + .decode(ByteBuffer.wrap(value)); + return false; + } catch (CharacterCodingException error) { + return true; + } + } + + private static long longValue(Object value, String name) { + if (value instanceof Byte || value instanceof Short || value instanceof Integer || value instanceof Long) { + return ((Number) value).longValue(); + } + if (value == null || value.toString().isEmpty()) { + return 0L; + } + try { + return Long.parseLong(value.toString()); + } catch (NumberFormatException error) { + throw new IllegalArgumentException(name + " must be an integer", error); + } + } + + private static String readFilePath(Map args) { + String path = ""; + for (String name : new String[] {"file_path", "path", "file"}) { + Object value = args.get(name); + if (value != null && !(value instanceof String)) { + throw new IllegalArgumentException(name + " must be a string"); + } + String candidate = stringValue(value); + if (candidate.isEmpty()) { + continue; + } + if (!path.isEmpty() && !path.equals(candidate)) { + throw new IllegalArgumentException("file_path, path, and file must not conflict"); + } + path = candidate; + } + if (path.isEmpty()) { + throw new IllegalArgumentException("file_path is required"); + } + return path; + } + + private static Long optionalLong(Map args, String name) { + return args.containsKey(name) && args.get(name) != null ? longValue(args.get(name), name) : null; + } + + private static List readViewRange(Object value) { + if (value == null) { + return null; + } + if (!(value instanceof List)) { + throw new IllegalArgumentException("view_range must be [start_line, end_line]"); + } + List raw = (List) value; + if (raw.isEmpty()) { + return Collections.emptyList(); + } + if (raw.size() != 2) { + throw new IllegalArgumentException("view_range must be [start_line, end_line]"); + } + List result = new ArrayList<>(2); + result.add(longValue(raw.get(0), "view_range")); + result.add(longValue(raw.get(1), "view_range")); + return result; + } + + private static String decodeUTF8(byte[] data) throws CharacterCodingException { + return StandardCharsets.UTF_8 + .newDecoder() + .onMalformedInput(CodingErrorAction.REPORT) + .onUnmappableCharacter(CodingErrorAction.REPORT) + .decode(ByteBuffer.wrap(data)) + .toString(); + } + + private static String renderViewRange(String text, List viewRange) { + String[] lines = text.split("\n", -1); + long startLine = viewRange.get(0); + long zeroBasedStart = startLine <= 1L ? 0L : startLine - 1L; + int start = (int) Math.min(zeroBasedStart, lines.length); + int end = lines.length; + if (viewRange.get(1) > 0L) { + end = (int) Math.min(viewRange.get(1), lines.length); + } + if (end < start) { + throw new IllegalArgumentException( + "view_range end line " + viewRange.get(1) + + " is before start line " + viewRange.get(0)); + } + StringBuilder output = new StringBuilder(); + for (int index = start; index < end; index++) { + if (index > start) { + output.append('\n'); + } + output.append(lines[index]); + } + return output.toString(); + } + + private static String renderReadLines(String text, Long offset, Long limit) { + String[] lines = text.split("\n", -1); + int total = lines.length; + if (total > 0 && lines[total - 1].isEmpty() && text.endsWith("\n")) { + total--; + } + long startValue = offset == null ? 0L : offset - 1L; + int start = (int) Math.min(startValue, total); + long count = limit != null && limit > 0L ? limit : DEFAULT_READ_LINES; + int end = count >= total - start ? total : start + (int) count; + StringBuilder output = new StringBuilder(); + for (int index = start; index < end; index++) { + output.append(String.format(Locale.ROOT, "%6d\t%s\n", index + 1, truncateLine(lines[index]))); + } + if (end < total) { + output.append(String.format( + Locale.ROOT, + "\n[truncated: showing lines %d-%d of %d]\n", + start + 1, + end, + total)); + } + return output.toString(); + } + + private static String truncateLine(String line) { + if (line.codePointCount(0, line.length()) <= MAX_READ_LINE_CHARS) { + return line; + } + return line.substring(0, line.offsetByCodePoints(0, MAX_READ_LINE_CHARS)) + + " [line truncated]"; + } + + private static class MediaInfo { + final String blockType; + final String mediaType; + + MediaInfo(String blockType, String mediaType) { + this.blockType = blockType; + this.mediaType = mediaType; + } + } + private static int appendMatches( Pattern regex, ToolContext context, diff --git a/src/main/java/com/volcengine/ark/runtime/selfhosted/Event.java b/src/main/java/com/volcengine/ark/runtime/selfhosted/Event.java index 0a97e89..931127a 100644 --- a/src/main/java/com/volcengine/ark/runtime/selfhosted/Event.java +++ b/src/main/java/com/volcengine/ark/runtime/selfhosted/Event.java @@ -221,6 +221,10 @@ public Object getInput() { return input; } + public String getProcessedAt() { + return processedAt; + } + public String getEvaluatedPermission() { return evaluatedPermission; } diff --git a/src/main/java/com/volcengine/ark/runtime/selfhosted/SelfHostedConstants.java b/src/main/java/com/volcengine/ark/runtime/selfhosted/SelfHostedConstants.java index 8f1c8f2..c454858 100644 --- a/src/main/java/com/volcengine/ark/runtime/selfhosted/SelfHostedConstants.java +++ b/src/main/java/com/volcengine/ark/runtime/selfhosted/SelfHostedConstants.java @@ -10,6 +10,7 @@ public final class SelfHostedConstants { public static final String EVENT_TYPE_AGENT_TOOL_USE = "agent.tool_use"; public static final String EVENT_TYPE_AGENT_CUSTOM_TOOL_USE = "agent.custom_tool_use"; public static final String EVENT_TYPE_USER_TOOL_CONFIRMATION = "user.tool_confirmation"; + public static final String EVENT_TYPE_USER_INTERRUPT = "user.interrupt"; public static final String EVENT_TYPE_USER_TOOL_RESULT = "user.tool_result"; public static final String EVENT_TYPE_USER_CUSTOM_TOOL_RESULT = "user.custom_tool_result"; public static final String EVENT_TYPE_SESSION_STATUS_IDLE = "session.status_idle"; @@ -31,6 +32,8 @@ public final class SelfHostedConstants { public static final long DEFAULT_MAX_IDLE_MILLIS = 60000L; public static final long DEFAULT_TOOL_TIMEOUT_MILLIS = 120000L; + public static final long DEFAULT_MAX_INPUT_FILE_BYTES = 100000L; + public static final long DEFAULT_MAX_MEDIA_FILE_BYTES = 7L << 20; public static final long DEFAULT_HEARTBEAT_MILLIS = 30000L; public static final int DEFAULT_POLL_BLOCK_MILLIS = 999; diff --git a/src/main/java/com/volcengine/ark/runtime/selfhosted/SessionToolRunner.java b/src/main/java/com/volcengine/ark/runtime/selfhosted/SessionToolRunner.java index 70faeab..8694c16 100644 --- a/src/main/java/com/volcengine/ark/runtime/selfhosted/SessionToolRunner.java +++ b/src/main/java/com/volcengine/ark/runtime/selfhosted/SessionToolRunner.java @@ -63,24 +63,29 @@ public SessionToolRunner(SelfHostedClient api, String sessionId, Options options } public List run() throws IOException { - if (options.resultStore != null) { - FileToolResultStore.RecoverResult recovered = options.resultStore.recover(); - state.pendingResults.putAll(recovered.getPending()); - for (String callId : recovered.getPending().keySet()) { - state.recoveredResults.put(callId, Boolean.TRUE); + try { + if (options.resultStore != null) { + FileToolResultStore.RecoverResult recovered = options.resultStore.recover(); + state.pendingResults.putAll(recovered.getPending()); + for (String callId : recovered.getPending().keySet()) { + state.recoveredResults.put(callId, Boolean.TRUE); + } + state.processed.putAll(recovered.getProcessed()); + state.answered.putAll(recovered.getProcessed()); } - state.processed.putAll(recovered.getProcessed()); - state.answered.putAll(recovered.getProcessed()); - } - if (options.preferStream) { - try { - consumeStreamLoop(); - return results; - } catch (StreamUnsupportedException ignored) { + if (options.preferStream) { + try { + consumeStreamLoop(); + return results; + } catch (StreamUnsupportedException ignored) { + } } + consumeList(); + return results; + } finally { + cancelActiveExecution(); + toolExecutor.shutdownNow(); } - consumeList(); - return results; } private void consumeStreamLoop() throws IOException { @@ -94,6 +99,7 @@ private void consumeStreamLoop() throws IOException { try { reconcile(true); while (!isClosed() && (pump.isAlive() || !events.isEmpty())) { + drainExecutionDone(); flushResults(); if (idleExpired()) { throw new IdleTimeoutException(); @@ -205,6 +211,7 @@ private void reconcileOnce(boolean reconcile) throws IOException { private void consumeList() throws IOException { while (!isClosed()) { + drainExecutionDone(); reconcile(false); flushResults(); if (idleExpired()) { @@ -216,6 +223,7 @@ private void consumeList() throws IOException { public void close() { closed = true; + cancelActiveExecution(); closeStream(activeStream); toolExecutor.shutdownNow(); } @@ -227,6 +235,7 @@ public List getResults() { private void processListedEvents(List events, boolean reconcile) throws IOException { List pending = new ArrayList<>(); Map pendingIds = new LinkedHashMap<>(); + Map replayedToolUses = new LinkedHashMap<>(); boolean touchedIdle = false; boolean lastWasEndTurn = false; for (Event event : events) { @@ -243,12 +252,15 @@ private void processListedEvents(List events, boolean reconcile) throws I String type = event.getType(); if (SelfHostedConstants.EVENT_TYPE_USER_TOOL_CONFIRMATION.equals(type)) { recordConfirmation(event); + } else if (SelfHostedConstants.EVENT_TYPE_USER_INTERRUPT.equals(type)) { + handleInterrupt(event, reconcile ? replayedToolUses : null); } else if (SelfHostedConstants.EVENT_TYPE_USER_TOOL_RESULT.equals(type) || SelfHostedConstants.EVENT_TYPE_USER_CUSTOM_TOOL_RESULT.equals(type)) { markAnswered(event.resultCallId()); } else if (SelfHostedConstants.EVENT_TYPE_AGENT_TOOL_USE.equals(type) || SelfHostedConstants.EVENT_TYPE_AGENT_CUSTOM_TOOL_USE.equals(type)) { String callId = event.callId(); + replayedToolUses.put(callId, Boolean.TRUE); if (!callId.isEmpty() && !pendingIds.containsKey(callId)) { pending.add(event); pendingIds.put(callId, Boolean.TRUE); @@ -269,11 +281,7 @@ private void processListedEvents(List events, boolean reconcile) throws I } releaseConfirmedToolUses(); if (touchedIdle && lastWasEndTurn) { - if (hasUnblockedOutstandingTool(pending)) { - disarmIdle(); - } else { - armIdle(); - } + armIdle(); } } @@ -304,6 +312,8 @@ private void handleEvent(Event event) throws IOException { if (SelfHostedConstants.EVENT_TYPE_USER_TOOL_CONFIRMATION.equals(type)) { recordConfirmation(event); releaseConfirmedToolUses(); + } else if (SelfHostedConstants.EVENT_TYPE_USER_INTERRUPT.equals(type)) { + handleInterrupt(event, null); } else if (SelfHostedConstants.EVENT_TYPE_USER_TOOL_RESULT.equals(type) || SelfHostedConstants.EVENT_TYPE_USER_CUSTOM_TOOL_RESULT.equals(type)) { markAnswered(event.resultCallId()); @@ -318,7 +328,7 @@ private void handleEvent(Event event) throws IOException { private void handleToolUse(Event event, boolean custom) throws IOException { String callId = event.callId(); - if (callId.isEmpty() || isAnswered(callId)) { + if (callId.isEmpty() || isAnswered(callId) || state.scheduled.containsKey(callId)) { return; } Event pending = state.pendingResults.get(callId); @@ -352,8 +362,59 @@ private void handleToolUse(Event event, boolean custom) throws IOException { return; } } - ToolResult result = executeTool(event, custom); - Event out = custom + state.scheduled.put(callId, Boolean.TRUE); + state.executionQueue.add(new PendingToolEvent(event, custom, decision.confirmation)); + startNextToolExecution(); + } + + private void startNextToolExecution() { + if (state.activeExecution != null) { + return; + } + while (!state.executionQueue.isEmpty()) { + PendingToolEvent pending = state.executionQueue.remove(0); + String callId = pending.event.callId(); + if (isAnswered(callId)) { + state.scheduled.remove(callId); + continue; + } + AtomicBoolean cancelled = new AtomicBoolean(); + state.activeExecution = new ActiveToolExecution(pending, cancelled); + try { + toolExecutor.submit(() -> { + ToolResult result = executeTool(pending.event, pending.custom, cancelled); + state.executionDone.offer(new ToolExecutionResult(pending, result)); + }); + } catch (RejectedExecutionException error) { + state.executionDone.offer(new ToolExecutionResult(pending, ToolResult.error("tool execution canceled"))); + } + return; + } + } + + private void drainExecutionDone() { + ToolExecutionResult completed; + while ((completed = state.executionDone.poll()) != null) { + finishToolExecution(completed); + } + } + + private void finishToolExecution(ToolExecutionResult completed) { + String callId = completed.pending.event.callId(); + if (state.activeExecution != null && state.activeExecution.pending.event.callId().equals(callId)) { + state.activeExecution = null; + } + state.scheduled.remove(callId); + if (!isAnswered(callId)) { + postResult(completed.pending, completed.result); + } + startNextToolExecution(); + } + + private void postResult(PendingToolEvent pending, ToolResult result) { + Event event = pending.event; + String callId = event.callId(); + Event out = pending.custom ? Event.newUserCustomToolResultEvent(callId, result.getContent(), result.isError(), event.getSessionThreadId()) : Event.newUserToolResultEvent(callId, result.getContent(), result.isError(), event.getSessionThreadId()); if (options.resultStore != null) { @@ -364,15 +425,67 @@ private void handleToolUse(Event event, boolean custom) throws IOException { } state.pendingResults.put(callId, out); } - sendResult(callId, event, custom, decision.confirmation, out); + sendResult(callId, event, pending.custom, pending.confirmation, out); + } + + private void handleInterrupt(Event event, Map eligible) { + String threadId = event.getSessionThreadId(); + ActiveToolExecution active = state.activeExecution; + if (active != null && interruptMatches(active.pending.event, threadId, eligible)) { + active.cancelled.set(true); + settleInterruptedToolUse(active.pending.event.callId()); + } + List retained = new ArrayList<>(); + for (PendingToolEvent pending : state.executionQueue) { + if (interruptMatches(pending.event, threadId, eligible)) { + settleInterruptedToolUse(pending.event.callId()); + } else { + retained.add(pending); + } + } + state.executionQueue.clear(); + state.executionQueue.addAll(retained); + for (Map.Entry entry : new ArrayList<>(state.toolUseEvents.entrySet())) { + if (!isAnswered(entry.getKey()) && interruptMatches(entry.getValue(), threadId, eligible)) { + settleInterruptedToolUse(entry.getKey()); + } + } + } + + private boolean interruptMatches(Event toolEvent, String threadId, Map eligible) { + String callId = toolEvent.callId(); + if (eligible != null && !eligible.containsKey(callId)) { + return false; + } + return threadId == null || threadId.isEmpty() || threadId.equals(toolEvent.getSessionThreadId()); } - private ToolResult executeTool(Event event, boolean custom) { + private void settleInterruptedToolUse(String callId) { + if (callId == null || callId.isEmpty() || isAnswered(callId)) { + return; + } + state.scheduled.remove(callId); + markAnswered(callId); + if (options.resultStore != null) { + try { + options.resultStore.discard(callId); + } catch (IOException error) { + LOGGER.log(Level.WARNING, "discard interrupted tool result failed tool_use_id=" + callId, error); + } + } + } + + private void cancelActiveExecution() { + if (state.activeExecution != null) { + state.activeExecution.cancelled.set(true); + } + } + + private ToolResult executeTool(Event event, boolean custom, AtomicBoolean executionCancelled) { long timeoutMillis = options.toolContext.getToolTimeoutMillis(); if (timeoutMillis <= 0L) { timeoutMillis = SelfHostedConstants.DEFAULT_TOOL_TIMEOUT_MILLIS; } - AtomicBoolean executionCancelled = new AtomicBoolean(); ToolContext context = toolContextForExecution(timeoutMillis, executionCancelled); Future future; try { @@ -431,6 +544,8 @@ private ToolContext toolContextForExecution(long timeoutMillis, AtomicBoolean ex } context.setUnrestrictedPaths(source.isUnrestrictedPaths()); context.setToolTimeoutMillis(timeoutMillis); + context.setMaxInputFileBytes(source.getMaxInputFileBytes()); + context.setMaxMediaFileBytes(source.getMaxMediaFileBytes()); context.setCancelled(() -> executionCancelled.get() || isClosed() || source.isCancelled()); return context; } @@ -439,7 +554,7 @@ private static String errorText(Throwable error) { return error.getMessage() == null ? error.toString() : error.getMessage(); } - private void sendResult(String callId, Event source, boolean custom, String confirmation, Event out) throws IOException { + private void sendResult(String callId, Event source, boolean custom, String confirmation, Event out) { boolean posted = retrySendEvent(out); if (posted) { markAnswered(callId); @@ -506,6 +621,9 @@ private void observeSessionState(Event event) { String callId = event.callId(); if (!callId.isEmpty()) { state.sessionToolUses.put(callId, Boolean.TRUE); + if (!isAnswered(callId)) { + state.toolUseEvents.put(callId, event); + } state.toolUsesSinceStatus.put(callId, Boolean.TRUE); } return; @@ -596,6 +714,9 @@ private PermissionDecision permissionAllows(Event event, boolean custom, String private boolean markEventSeen(Event event) { String key = event.getId().isEmpty() ? event.callId() : event.getId(); + if (key.isEmpty() && SelfHostedConstants.EVENT_TYPE_USER_INTERRUPT.equals(event.getType())) { + key = "interrupt:" + event.getProcessedAt() + ":" + event.getSessionThreadId(); + } if (key.isEmpty()) { return true; } @@ -616,6 +737,7 @@ private void markAnswered(String callId) { state.recoveredResults.remove(callId); state.pendingAsk.remove(callId); state.externalTools.remove(callId); + state.toolUseEvents.remove(callId); maybeArmPendingIdle(); } @@ -639,20 +761,6 @@ private void releaseConfirmedToolUses() throws IOException { } } - private boolean hasUnblockedOutstandingTool(List pending) { - for (Event event : pending) { - String callId = event.callId(); - if (callId.isEmpty() || isAnswered(callId) || !shouldHandleToolUse(callId)) { - continue; - } - if (state.pendingAsk.containsKey(callId) || state.pendingResults.containsKey(callId)) { - continue; - } - return true; - } - return false; - } - private void armIdle() { if (options.maxIdleMillis <= 0) { return; @@ -679,7 +787,10 @@ private void maybeArmPendingIdle() { } private boolean hasIdleBlockers() { - return !state.pendingAsk.isEmpty() || !state.pendingResults.isEmpty() || !state.externalTools.isEmpty(); + return !state.pendingAsk.isEmpty() + || !state.pendingResults.isEmpty() + || !state.externalTools.isEmpty() + || !state.scheduled.isEmpty(); } private boolean idleExpired() { @@ -691,6 +802,7 @@ private boolean idleExpired() { private void sleepOrIdle(long millis) { long deadline = System.currentTimeMillis() + Math.max(millis, 0L); while (!isClosed()) { + drainExecutionDone(); if (idleExpired()) { throw new IdleTimeoutException(); } @@ -742,6 +854,11 @@ private static class State { Map pendingAsk = new LinkedHashMap<>(); Map confirmations = new LinkedHashMap<>(); Map externalTools = new LinkedHashMap<>(); + Map toolUseEvents = new LinkedHashMap<>(); + Map scheduled = new LinkedHashMap<>(); + List executionQueue = new ArrayList<>(); + volatile ActiveToolExecution activeExecution; + LinkedBlockingQueue executionDone = new LinkedBlockingQueue<>(); Map sessionToolUses = new LinkedHashMap<>(); Map toolUsesSinceStatus = new LinkedHashMap<>(); Map blockingEventIds = new LinkedHashMap<>(); @@ -750,6 +867,38 @@ private static class State { boolean idleArmPending; } + private static class PendingToolEvent { + final Event event; + final boolean custom; + final String confirmation; + + PendingToolEvent(Event event, boolean custom, String confirmation) { + this.event = event; + this.custom = custom; + this.confirmation = confirmation; + } + } + + private static class ActiveToolExecution { + final PendingToolEvent pending; + final AtomicBoolean cancelled; + + ActiveToolExecution(PendingToolEvent pending, AtomicBoolean cancelled) { + this.pending = pending; + this.cancelled = cancelled; + } + } + + private static class ToolExecutionResult { + final PendingToolEvent pending; + final ToolResult result; + + ToolExecutionResult(PendingToolEvent pending, ToolResult result) { + this.pending = pending; + this.result = result; + } + } + private static class PermissionDecision { final String confirmation; final boolean allowed; diff --git a/src/main/java/com/volcengine/ark/runtime/selfhosted/ToolContext.java b/src/main/java/com/volcengine/ark/runtime/selfhosted/ToolContext.java index d9a5c46..cb12051 100644 --- a/src/main/java/com/volcengine/ark/runtime/selfhosted/ToolContext.java +++ b/src/main/java/com/volcengine/ark/runtime/selfhosted/ToolContext.java @@ -13,6 +13,8 @@ public class ToolContext { private boolean explicitEnv; private boolean unrestrictedPaths; private long toolTimeoutMillis = SelfHostedConstants.DEFAULT_TOOL_TIMEOUT_MILLIS; + private long maxInputFileBytes = SelfHostedConstants.DEFAULT_MAX_INPUT_FILE_BYTES; + private long maxMediaFileBytes = SelfHostedConstants.DEFAULT_MAX_MEDIA_FILE_BYTES; private BooleanSupplier cancelled = () -> false; public ToolContext(String workdir) { @@ -56,6 +58,30 @@ public void setToolTimeoutMillis(long toolTimeoutMillis) { this.toolTimeoutMillis = toolTimeoutMillis; } + public long getMaxInputFileBytes() { + return maxInputFileBytes; + } + + /** + * Sets the text read limit. Non-positive values retain the legacy text cap; when used as the + * media fallback, zero combines with a zero media limit to disable the media size limit. + */ + public void setMaxInputFileBytes(long maxInputFileBytes) { + this.maxInputFileBytes = maxInputFileBytes; + } + + public long getMaxMediaFileBytes() { + return maxMediaFileBytes; + } + + /** + * Sets the media read limit. Zero follows the input limit, so both limits set to zero disable + * the media size limit; a negative value also disables it. + */ + public void setMaxMediaFileBytes(long maxMediaFileBytes) { + this.maxMediaFileBytes = maxMediaFileBytes; + } + public boolean isCancelled() { return cancelled != null && cancelled.getAsBoolean(); } 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 de93212..69aa31f 100644 --- a/src/test/java/com/volcengine/ark/runtime/selfhosted/DefaultToolsTest.java +++ b/src/test/java/com/volcengine/ark/runtime/selfhosted/DefaultToolsTest.java @@ -9,12 +9,167 @@ import java.nio.file.Files; import java.nio.file.Path; +import java.util.Base64; +import java.util.Collections; import java.util.LinkedHashMap; import java.util.Map; import java.util.concurrent.TimeUnit; import org.junit.Test; public class DefaultToolsTest { + @Test + public void readReturnsNativeImageAndDocumentBlocks() throws Exception { + Path workdir = Files.createTempDirectory("ark-java-media-"); + byte[] image = new byte[] {(byte) 0x89, 'P', 'N', 'G', '\r', '\n', (byte) 0x1a, '\n', 1}; + byte[] document = "%PDF-1.7\ndata".getBytes(java.nio.charset.StandardCharsets.US_ASCII); + Files.write(workdir.resolve("image.png"), image); + Files.write(workdir.resolve("document.pdf"), document); + ToolContext context = new ToolContext(workdir.toString()); + + ToolResult imageResult = new DefaultTools.ReadTool().execute( + Collections.singletonMap("file_path", "image.png"), context); + ToolResult documentResult = new DefaultTools.ReadTool().execute( + Collections.singletonMap("file_path", "document.pdf"), context); + + assertMediaBlock(imageResult, "image", "image/png", image); + assertMediaBlock(documentResult, "document", "application/pdf", document); + } + + @Test + public void readRejectsRangesAndOversizedMedia() throws Exception { + Path workdir = Files.createTempDirectory("ark-java-media-limit-"); + byte[] image = new byte[] {(byte) 0x89, 'P', 'N', 'G', '\r', '\n', (byte) 0x1a, '\n', 1}; + Files.write(workdir.resolve("image.png"), image); + ToolContext context = new ToolContext(workdir.toString()); + Map rangedInput = new LinkedHashMap<>(); + rangedInput.put("path", "image.png"); + rangedInput.put("limit", 1); + + ToolResult ranged = new DefaultTools.ReadTool().execute(rangedInput, context); + context.setMaxMediaFileBytes(image.length - 1L); + ToolResult oversized = new DefaultTools.ReadTool().execute( + Collections.singletonMap("path", "image.png"), context); + + assertTrue(ranged.isError()); + assertEquals( + "view_range, offset, and limit are only supported for text files", + ranged.getContent().get(0).getText()); + assertTrue(oversized.isError()); + assertEquals("media file too large: " + image.length + " bytes", oversized.getContent().get(0).getText()); + } + + @Test + public void readReportsInvalidMediaRangeField() throws Exception { + Path workdir = Files.createTempDirectory("ark-java-media-range-"); + Files.write(workdir.resolve("image.png"), new byte[] { + (byte) 0x89, 'P', 'N', 'G', '\r', '\n', (byte) 0x1a, '\n', 1 + }); + Map input = new LinkedHashMap<>(); + input.put("path", "image.png"); + input.put("offset", "abc"); + + ToolResult result = new DefaultTools.ReadTool().execute(input, new ToolContext(workdir.toString())); + + assertTrue(result.isError()); + assertEquals("offset must be an integer", result.getContent().get(0).getText()); + } + + @Test + public void readUsesManagedAgentLineRange() throws Exception { + Path workdir = Files.createTempDirectory("ark-java-read-lines-"); + Files.write( + workdir.resolve("example.txt"), + "a\nb\nc\nd\n".getBytes(java.nio.charset.StandardCharsets.UTF_8)); + ToolContext context = new ToolContext(workdir.toString()); + DefaultTools.ReadTool tool = new DefaultTools.ReadTool(); + Map rangedInput = new LinkedHashMap<>(); + rangedInput.put("file_path", "example.txt"); + rangedInput.put("offset", 2); + rangedInput.put("limit", 2); + + ToolResult whole = tool.execute(Collections.singletonMap("file_path", "example.txt"), context); + ToolResult ranged = tool.execute(rangedInput, context); + Map invalidInput = new LinkedHashMap<>(); + invalidInput.put("file_path", "example.txt"); + invalidInput.put("offset", 0); + ToolResult invalid = tool.execute(invalidInput, context); + + assertFalse(whole.isError()); + assertEquals(" 1\ta\n 2\tb\n 3\tc\n 4\td\n", whole.getContent().get(0).getText()); + assertFalse(ranged.isError()); + assertEquals( + " 2\tb\n 3\tc\n\n[truncated: showing lines 2-3 of 4]\n", + ranged.getContent().get(0).getText()); + assertTrue(invalid.isError()); + assertEquals( + "offset is the 1-based start line and must be >= 1, got 0", + invalid.getContent().get(0).getText()); + } + + @Test + public void readMarksTruncatedLinesAndRejectsInvalidUTF8() throws Exception { + Path workdir = Files.createTempDirectory("ark-java-read-truncated-"); + Files.write( + workdir.resolve("long.txt"), + repeat('a', 2001).getBytes(java.nio.charset.StandardCharsets.UTF_8)); + Files.write(workdir.resolve("binary.txt"), new byte[] {(byte) 0xff}); + ToolContext context = new ToolContext(workdir.toString()); + DefaultTools.ReadTool tool = new DefaultTools.ReadTool(); + + ToolResult truncated = tool.execute(Collections.singletonMap("file_path", "long.txt"), context); + ToolResult invalidUTF8 = tool.execute(Collections.singletonMap("file_path", "binary.txt"), context); + + assertFalse(truncated.isError()); + assertEquals( + " 1\t" + repeat('a', 2000) + " [line truncated]\n", + truncated.getContent().get(0).getText()); + assertTrue(invalidUTF8.isError()); + assertEquals("binary file cannot be read directly", invalidUTF8.getContent().get(0).getText()); + } + + @Test + public void readKeepsLegacyInputsWithoutAmbiguity() throws Exception { + Path workdir = Files.createTempDirectory("ark-java-read-compat-"); + Files.write( + workdir.resolve("example.txt"), + "a\nb\nc\n".getBytes(java.nio.charset.StandardCharsets.UTF_8)); + ToolContext context = new ToolContext(workdir.toString()); + DefaultTools.ReadTool tool = new DefaultTools.ReadTool(); + + assertFalse(tool.execute(Collections.singletonMap("path", "example.txt"), context).isError()); + assertFalse(tool.execute(Collections.singletonMap("file", "example.txt"), context).isError()); + Map samePath = new LinkedHashMap<>(); + samePath.put("file_path", "example.txt"); + samePath.put("path", "example.txt"); + assertFalse(tool.execute(samePath, context).isError()); + + Map legacyRangeInput = new LinkedHashMap<>(); + legacyRangeInput.put("file_path", "example.txt"); + legacyRangeInput.put("view_range", java.util.Arrays.asList(2, 3)); + ToolResult legacyRange = tool.execute(legacyRangeInput, context); + assertFalse(legacyRange.isError()); + assertEquals("b\nc", legacyRange.getContent().get(0).getText()); + + Map pathConflictInput = new LinkedHashMap<>(); + pathConflictInput.put("file_path", "example.txt"); + pathConflictInput.put("path", "other.txt"); + ToolResult pathConflict = tool.execute(pathConflictInput, context); + assertTrue(pathConflict.isError()); + assertEquals( + "file_path, path, and file must not conflict", + pathConflict.getContent().get(0).getText()); + + Map rangeConflictInput = new LinkedHashMap<>(); + rangeConflictInput.put("file_path", "example.txt"); + rangeConflictInput.put("view_range", java.util.Arrays.asList(1, 1)); + rangeConflictInput.put("offset", 1); + ToolResult rangeConflict = tool.execute(rangeConflictInput, context); + assertTrue(rangeConflict.isError()); + assertEquals( + "view_range cannot be combined with offset or limit", + rangeConflict.getContent().get(0).getText()); + } + @Test public void bashScrubsSensitiveExplicitEnvironment() throws Exception { Path workdir = Files.createTempDirectory("ark-java-tools-"); @@ -149,4 +304,24 @@ public void bashDrainsAndBoundsLargeOutput() throws Exception { private static String envName(String... parts) { return String.join("_", parts); } + + private static String repeat(char value, int count) { + StringBuilder result = new StringBuilder(count); + for (int index = 0; index < count; index++) { + result.append(value); + } + return result.toString(); + } + + @SuppressWarnings("unchecked") + private static void assertMediaBlock( + ToolResult result, String blockType, String mediaType, byte[] expected) { + assertFalse(result.isError()); + ContentBlock block = result.getContent().get(0); + assertEquals(blockType, block.getType()); + Map source = (Map) block.getSource(); + assertEquals("base64", source.get("type")); + assertEquals(mediaType, source.get("media_type")); + assertEquals(Base64.getEncoder().encodeToString(expected), source.get("data")); + } } diff --git a/src/test/java/com/volcengine/ark/runtime/selfhosted/SessionToolRunnerTest.java b/src/test/java/com/volcengine/ark/runtime/selfhosted/SessionToolRunnerTest.java index 6a56549..3dbe324 100644 --- a/src/test/java/com/volcengine/ark/runtime/selfhosted/SessionToolRunnerTest.java +++ b/src/test/java/com/volcengine/ark/runtime/selfhosted/SessionToolRunnerTest.java @@ -18,7 +18,9 @@ import java.util.List; import java.util.Map; import java.util.concurrent.CountDownLatch; +import java.util.concurrent.BlockingQueue; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; import okhttp3.MediaType; @@ -225,6 +227,7 @@ public void recoveredResultsAreFilteredAgainstCurrentBlockers() throws Exception toolUse("stale-call"), toolUse("current-call"), requiresAction("current-call"))); + finishExecution(runner); assertEquals(1, executions.get()); assertEquals(1, sent.size()); @@ -236,6 +239,157 @@ public void recoveredResultsAreFilteredAgainstCurrentBlockers() throws Exception runner.close(); } + @Test + public void interruptCancelsActiveToolWithoutPostingResult() throws Exception { + CountDownLatch started = new CountDownLatch(1); + CountDownLatch canceled = new CountDownLatch(1); + Tool tool = blockingTool(started, canceled); + List discarded = new ArrayList<>(); + List sent = new ArrayList<>(); + SessionToolRunner runner = runnerWithStore(recordingStore(discarded, new ArrayList<>()), tool, sent); + + handleStreamEvent(runner, toolUse("call-1", "thread-1", "blocking")); + assertTrue(started.await(1L, TimeUnit.SECONDS)); + handleStreamEvent(runner, interrupt("interrupt-1", "thread-1")); + assertTrue(canceled.await(1L, TimeUnit.SECONDS)); + finishExecution(runner); + + assertTrue(sent.isEmpty()); + assertTrue(answered(runner).containsKey("call-1")); + assertEquals(Collections.singletonList("call-1"), discarded); + runner.close(); + } + + @Test + public void interruptOnlyCancelsTargetThread() throws Exception { + CountDownLatch started = new CountDownLatch(1); + CountDownLatch canceled = new CountDownLatch(1); + Tool tool = blockingTool(started, canceled); + SessionToolRunner runner = runnerWithStore(null, tool, new ArrayList<>()); + + handleStreamEvent(runner, toolUse("call-1", "thread-a", "blocking")); + assertTrue(started.await(1L, TimeUnit.SECONDS)); + handleStreamEvent(runner, interrupt("interrupt-other", "thread-b")); + + assertFalse(answered(runner).containsKey("call-1")); + assertFalse(canceled.await(50L, TimeUnit.MILLISECONDS)); + + handleStreamEvent(runner, interrupt("interrupt-target", "thread-a")); + assertTrue(canceled.await(1L, TimeUnit.SECONDS)); + finishExecution(runner); + runner.close(); + } + + @Test + public void listReplayOfAnonymousInterruptDoesNotRedispatchOrCancelLaterTool() throws Exception { + CountDownLatch started = new CountDownLatch(1); + CountDownLatch canceled = new CountDownLatch(1); + AtomicInteger executions = new AtomicInteger(); + Tool tool = new Tool() { + @Override + public String name() { + return "blocking"; + } + + @Override + public ToolResult execute(Object input, ToolContext context) { + executions.incrementAndGet(); + started.countDown(); + while (!context.isCancelled()) { + try { + Thread.sleep(5L); + } catch (InterruptedException ignored) { + } + } + canceled.countDown(); + return ToolResult.text("late"); + } + }; + SessionToolRunner runner = runnerWithStore(null, tool, new ArrayList<>()); + List events = java.util.Arrays.asList( + toolUse("old-call", "thread-1", "blocking"), + interrupt("", "thread-1"), + toolUse("new-call", "thread-1", "blocking")); + + processListedEvents(runner, events, false); + assertTrue(started.await(1L, TimeUnit.SECONDS)); + processListedEvents(runner, events, false); + processListedEvents(runner, events, true); + + assertTrue(answered(runner).containsKey("old-call")); + assertFalse(answered(runner).containsKey("new-call")); + assertEquals(1, executions.get()); + assertFalse(canceled.await(50L, TimeUnit.MILLISECONDS)); + assertFalse(stateMap(runner, "toolUseEvents").containsKey("old-call")); + + handleStreamEvent(runner, interrupt("interrupt-all", "")); + assertTrue(canceled.await(1L, TimeUnit.SECONDS)); + finishExecution(runner); + runner.close(); + } + + @Test + public void listInterruptCancelsToolFromEarlierPoll() throws Exception { + CountDownLatch started = new CountDownLatch(1); + CountDownLatch canceled = new CountDownLatch(1); + SessionToolRunner runner = runnerWithStore(null, blockingTool(started, canceled), new ArrayList<>()); + Event toolUse = toolUse("cross-poll-call", "thread-1", "blocking"); + + processListedEvents(runner, Collections.singletonList(toolUse), false); + assertTrue(started.await(1L, TimeUnit.SECONDS)); + processListedEvents( + runner, + java.util.Arrays.asList( + toolUse, + interruptAt("", "2026-09-20T00:00:00Z", "thread-1")), + false); + + assertTrue(canceled.await(1L, TimeUnit.SECONDS)); + finishExecution(runner); + assertTrue(answered(runner).containsKey("cross-poll-call")); + runner.close(); + } + + @Test + public void interruptRemovesMatchingQueuedToolOnly() throws Exception { + CountDownLatch started = new CountDownLatch(1); + CountDownLatch canceled = new CountDownLatch(1); + SessionToolRunner runner = runnerWithStore(null, blockingTool(started, canceled), new ArrayList<>()); + Event active = toolUse("active-call", "thread-a", "blocking"); + Event queued = toolUse("queued-call", "thread-b", "blocking"); + + processListedEvents(runner, java.util.Arrays.asList(active, queued), false); + assertTrue(started.await(1L, TimeUnit.SECONDS)); + handleStreamEvent(runner, interrupt("interrupt-b", "thread-b")); + + assertTrue(answered(runner).containsKey("queued-call")); + assertFalse(answered(runner).containsKey("active-call")); + assertFalse(canceled.await(50L, TimeUnit.MILLISECONDS)); + + handleStreamEvent(runner, interrupt("interrupt-all", "")); + assertTrue(canceled.await(1L, TimeUnit.SECONDS)); + finishExecution(runner); + runner.close(); + } + + @Test + public void listEndTurnArmsIdleAfterToolCompletes() throws Exception { + SessionToolRunner runner = runnerWithStore(null, countingTool(new AtomicInteger()), new ArrayList<>()); + Event toolUse = toolUse("idle-call"); + processListedEvents(runner, Collections.singletonList(toolUse), false); + processListedEvents( + runner, + java.util.Arrays.asList(toolUse, idleEvent("idle-end-turn")), + false); + + assertTrue(stateBoolean(runner, "idleArmPending")); + assertEquals(0L, idleArmedAt(runner)); + finishExecution(runner); + assertFalse(stateBoolean(runner, "idleArmPending")); + assertTrue(idleArmedAt(runner) > 0L); + runner.close(); + } + @Test public void currentBlockerReusesRecoveredResultWithoutReexecution() throws Exception { List discarded = new ArrayList<>(); @@ -311,13 +465,14 @@ public ToolResult execute(Object input, ToolContext context) { raw.put("name", tool.name()); raw.put(custom ? "custom_tool_use_id" : "tool_use_id", "call-1"); raw.put("input", Collections.emptyMap()); - Method execute = SessionToolRunner.class.getDeclaredMethod("executeTool", Event.class, boolean.class); + Method execute = SessionToolRunner.class.getDeclaredMethod( + "executeTool", Event.class, boolean.class, AtomicBoolean.class); execute.setAccessible(true); long startedAt = System.nanoTime(); ToolResult result; try { - result = (ToolResult) execute.invoke(runner, Event.fromMap(raw), custom); + result = (ToolResult) execute.invoke(runner, Event.fromMap(raw), custom, new AtomicBoolean()); assertTrue(started.await(1L, TimeUnit.SECONDS)); } finally { release.countDown(); @@ -379,6 +534,28 @@ public ToolResult execute(Object input, ToolContext context) { }; } + private static Tool blockingTool(CountDownLatch started, CountDownLatch canceled) { + return new Tool() { + @Override + public String name() { + return "blocking"; + } + + @Override + public ToolResult execute(Object input, ToolContext context) { + started.countDown(); + while (!context.isCancelled()) { + try { + Thread.sleep(5L); + } catch (InterruptedException ignored) { + } + } + canceled.countDown(); + return ToolResult.text("late"); + } + }; + } + private static SessionToolRunner runnerWithStore(FileToolResultStore store, Tool tool, List sent) throws IOException { SelfHostedClient client = new SelfHostedClient("test-key") { @@ -394,19 +571,38 @@ public void sendEvent(String sessionId, Event event) { .tools(new ToolSet()) .toolContext(new ToolContext( Files.createTempDirectory("ark-java-recovery-runner-").toString())) - .customTools(Collections.singletonMap("custom", tool)) + .customTools(Collections.singletonMap(tool.name(), tool)) .resultStore(store)); } private static Event toolUse(String callId) { + return toolUse(callId, "", "custom"); + } + + private static Event toolUse(String callId, String threadId, String name) { Map raw = new LinkedHashMap<>(); raw.put("id", callId); raw.put("type", "agent.custom_tool_use"); - raw.put("name", "custom"); + raw.put("name", name); + raw.put("custom_tool_use_id", callId); + raw.put("session_thread_id", threadId); raw.put("input", Collections.emptyMap()); return Event.fromMap(raw); } + private static Event interrupt(String eventId, String threadId) { + return interruptAt(eventId, "", threadId); + } + + private static Event interruptAt(String eventId, String processedAt, String threadId) { + Map raw = new LinkedHashMap<>(); + raw.put("id", eventId); + raw.put("type", "user.interrupt"); + raw.put("processed_at", processedAt); + raw.put("session_thread_id", threadId); + return Event.fromMap(raw); + } + private static Event requiresAction(String callId) { Map stopReason = new LinkedHashMap<>(); stopReason.put("type", "requires_action"); @@ -419,16 +615,44 @@ private static Event requiresAction(String callId) { } private static void processListedEvents(SessionToolRunner runner, List events) throws Exception { + processListedEvents(runner, events, true); + } + + private static void processListedEvents(SessionToolRunner runner, List events, boolean reconcile) + throws Exception { Method process = SessionToolRunner.class.getDeclaredMethod("processListedEvents", List.class, boolean.class); process.setAccessible(true); - process.invoke(runner, events, true); + process.invoke(runner, events, reconcile); + } + + private static void handleStreamEvent(SessionToolRunner runner, Event event) throws Exception { + Method handle = SessionToolRunner.class.getDeclaredMethod("handleStreamEvent", Event.class); + handle.setAccessible(true); + handle.invoke(runner, event); + } + + private static void finishExecution(SessionToolRunner runner) throws Exception { + Field stateField = SessionToolRunner.class.getDeclaredField("state"); + stateField.setAccessible(true); + Object state = stateField.get(runner); + Field doneField = state.getClass().getDeclaredField("executionDone"); + doneField.setAccessible(true); + Object completed = ((BlockingQueue) doneField.get(state)).poll(1L, TimeUnit.SECONDS); + assertTrue("tool execution did not finish", completed != null); + Method finish = SessionToolRunner.class.getDeclaredMethod("finishToolExecution", completed.getClass()); + finish.setAccessible(true); + finish.invoke(runner, completed); } private static Event idleEvent() { + return idleEvent("idle-1"); + } + + private static Event idleEvent(String eventId) { Map stopReason = new LinkedHashMap<>(); stopReason.put("type", "end_turn"); Map raw = new LinkedHashMap<>(); - raw.put("id", "idle-1"); + raw.put("id", eventId); raw.put("type", "session.status_idle"); raw.put("stop_reason", stopReason); return Event.fromMap(raw); @@ -443,6 +667,15 @@ private static long idleArmedAt(SessionToolRunner runner) throws Exception { return idleField.getLong(state); } + private static boolean stateBoolean(SessionToolRunner runner, String fieldName) throws Exception { + Field stateField = SessionToolRunner.class.getDeclaredField("state"); + stateField.setAccessible(true); + Object state = stateField.get(runner); + Field field = state.getClass().getDeclaredField(fieldName); + field.setAccessible(true); + return field.getBoolean(state); + } + @SuppressWarnings("unchecked") private static Map answered(SessionToolRunner runner) throws Exception { Field stateField = SessionToolRunner.class.getDeclaredField("state");