diff --git a/mateclaw-server/src/main/java/vip/mate/agent/graph/executor/ToolCallDeadline.java b/mateclaw-server/src/main/java/vip/mate/agent/graph/executor/ToolCallDeadline.java new file mode 100644 index 00000000..c78f6800 --- /dev/null +++ b/mateclaw-server/src/main/java/vip/mate/agent/graph/executor/ToolCallDeadline.java @@ -0,0 +1,60 @@ +package vip.mate.agent.graph.executor; + +import java.util.concurrent.ScheduledFuture; +import java.util.concurrent.ScheduledThreadPoolExecutor; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import java.util.function.Supplier; + +/** Interrupts cooperative callbacks in place, preserving their thread-local context. */ +final class ToolCallDeadline { + private static final ScheduledThreadPoolExecutor TIMER = new ScheduledThreadPoolExecutor( + 1, Thread.ofPlatform().daemon().name("tool-deadline-", 0).factory()); + + static { + TIMER.setRemoveOnCancelPolicy(true); + } + + private final Thread owner = Thread.currentThread(); + private boolean active = true; + private boolean expired; + + private synchronized void expire() { + if (active) { + expired = true; + owner.interrupt(); + } + } + + private synchronized boolean finish() { + active = false; + return expired; + } + + static T call(String toolName, long timeoutMs, Supplier callback) throws TimeoutException { + ToolCallDeadline deadline = new ToolCallDeadline(); + ScheduledFuture timer = TIMER.schedule(deadline::expire, Math.max(1L, timeoutMs), TimeUnit.MILLISECONDS); + try { + T result = callback.get(); + if (deadline.finish()) throw timeout(toolName, timeoutMs); + return result; + } catch (RuntimeException failure) { + if (deadline.finish()) { + TimeoutException timeout = timeout(toolName, timeoutMs); + timeout.initCause(failure); + throw timeout; + } + throw failure; + } finally { + // Synchronize with the timer before this thread can execute another + // tool. Never leave a late watchdog interrupt on a reused worker. + boolean expired = deadline.finish(); + timer.cancel(false); + if (expired) Thread.interrupted(); + } + } + + private static TimeoutException timeout(String name, long timeoutMs) { + return new TimeoutException("Tool " + name + " timed out after " + timeoutMs + "ms"); + } +} diff --git a/mateclaw-server/src/main/java/vip/mate/agent/graph/executor/ToolExecutionExecutor.java b/mateclaw-server/src/main/java/vip/mate/agent/graph/executor/ToolExecutionExecutor.java index 5c43307a..99a5e3fe 100644 --- a/mateclaw-server/src/main/java/vip/mate/agent/graph/executor/ToolExecutionExecutor.java +++ b/mateclaw-server/src/main/java/vip/mate/agent/graph/executor/ToolExecutionExecutor.java @@ -1793,9 +1793,11 @@ public class ToolExecutionExecutor { } private String invokeObserved(ToolCallback callback, String arguments, ToolContext context, - String invocationKey, String providerCallId) { - return executionEvidenceRecorder == null ? callback.call(arguments, context) - : executionEvidenceRecorder.invoke(callback, arguments, context, invocationKey, providerCallId); + String invocationKey, String providerCallId) throws TimeoutException { + String toolName = callback.getToolDefinition().name(); + return ToolCallDeadline.call(toolName, getToolTimeoutMs(toolName), + () -> executionEvidenceRecorder == null ? callback.call(arguments, context) + : executionEvidenceRecorder.invoke(callback, arguments, context, invocationKey, providerCallId)); } // ==================== 内部数据类 ==================== diff --git a/mateclaw-server/src/main/java/vip/mate/tool/builtin/DocumentExtractTool.java b/mateclaw-server/src/main/java/vip/mate/tool/builtin/DocumentExtractTool.java index aa37e770..85341498 100644 --- a/mateclaw-server/src/main/java/vip/mate/tool/builtin/DocumentExtractTool.java +++ b/mateclaw-server/src/main/java/vip/mate/tool/builtin/DocumentExtractTool.java @@ -149,7 +149,7 @@ public class DocumentExtractTool { String forcedMethod = extractOption(options, "method"); if ("tika".equalsIgnoreCase(forcedMethod)) { long t = System.currentTimeMillis(); - String text = TikaExtractor.extract(path); + String text = TikaExtractor.extract(path, MAX_OUTPUT_LENGTH + 1); attempts.add("user-forced method=tika: skipped automatic fallback chain"); if (text == null || text.isBlank()) { attempts.add("tika: 失败或不可用 (" + (System.currentTimeMillis() - t) + "ms)"); @@ -161,7 +161,7 @@ public class DocumentExtractTool { boolean trunc = false; if (capped.length() > MAX_OUTPUT_LENGTH) { capped = capped.substring(0, MAX_OUTPUT_LENGTH) - + "\n\n... [内容已截断，总长度: " + text.length() + " 字符]"; + + "\n\n... [内容已截断，总长度至少: " + text.length() + " 字符]"; trunc = true; } result.set("text", capped); @@ -195,7 +195,7 @@ public class DocumentExtractTool { String text = content.text(); boolean truncated = false; if (text.length() > MAX_OUTPUT_LENGTH) { - text = text.substring(0, MAX_OUTPUT_LENGTH) + "\n\n... [内容已截断，总长度: " + content.text().length() + " 字符]"; + text = text.substring(0, MAX_OUTPUT_LENGTH) + "\n\n... [内容已截断，总长度至少: " + content.text().length() + " 字符]"; truncated = true; } @@ -878,7 +878,9 @@ public class DocumentExtractTool { private ExtractedContent extractXlsx(Path path, String options, List attempts) throws Exception { long t = System.currentTimeMillis(); - String text = TikaExtractor.extract(path); + // Stop at the response budget instead of parsing millions of unused + // characters. The extra character preserves the truncation marker. + String text = TikaExtractor.extract(path, MAX_OUTPUT_LENGTH + 1); long elapsed = System.currentTimeMillis() - t; if (text != null && !text.isBlank()) { attempts.add("tika: 成功 (" + elapsed + "ms)"); diff --git a/mateclaw-server/src/main/java/vip/mate/tool/builtin/TikaExtractor.java b/mateclaw-server/src/main/java/vip/mate/tool/builtin/TikaExtractor.java index 1965adc0..6747e12e 100644 --- a/mateclaw-server/src/main/java/vip/mate/tool/builtin/TikaExtractor.java +++ b/mateclaw-server/src/main/java/vip/mate/tool/builtin/TikaExtractor.java @@ -6,8 +6,13 @@ import org.apache.tika.metadata.Metadata; import org.apache.tika.parser.AutoDetectParser; import org.apache.tika.parser.ParseContext; import org.apache.tika.sax.BodyContentHandler; +import org.xml.sax.Attributes; +import org.xml.sax.SAXException; +import java.io.FilterInputStream; +import java.io.IOException; import java.io.InputStream; +import java.io.InterruptedIOException; import java.nio.file.Files; import java.nio.file.Path; @@ -64,12 +69,39 @@ public final class TikaExtractor { } int cap = maxChars <= 0 ? DEFAULT_MAX_CHARS : maxChars; - BodyContentHandler handler = new BodyContentHandler(cap); + BodyContentHandler handler = new BodyContentHandler(cap) { + @Override + public void characters(char[] chars, int start, int length) throws SAXException { + checkParseInterrupted(); + super.characters(chars, start, length); + } + + @Override + public void startElement(String uri, String localName, String name, Attributes attributes) + throws SAXException { + checkParseInterrupted(); + super.startElement(uri, localName, name, attributes); + } + }; AutoDetectParser parser = new AutoDetectParser(); Metadata metadata = new Metadata(); ParseContext context = new ParseContext(); - try (InputStream is = Files.newInputStream(path)) { + try (InputStream is = new FilterInputStream(Files.newInputStream(path)) { + @Override public int read() throws IOException { + checkInterrupted(); + return super.read(); + } + @Override public int read(byte[] bytes, int offset, int length) throws IOException { + checkInterrupted(); + return super.read(bytes, offset, length); + } + @Override public long skip(long count) throws IOException { + checkInterrupted(); + return super.skip(count); + } + }) { + checkInterrupted(); parser.parse(is, handler, metadata, context); return handler.toString(); } catch (WriteLimitReachedException truncated) { @@ -81,6 +113,12 @@ public final class TikaExtractor { partial.length(), path.getFileName()); return partial.isBlank() ? null : partial; } catch (Throwable t) { + // Office parsers may wrap the SAX write-limit exception. A bounded + // spreadsheet preview is still a successful extraction in that case. + if (!Thread.currentThread().isInterrupted() && WriteLimitReachedException.isWriteLimitReached(t)) { + String partial = handler.toString(); + return partial.isBlank() ? null : partial; + } // Catching Throwable on purpose: Tika can throw NoClassDefFoundError / // LinkageError when an obscure transitive parser is missing on a // minimal classpath, and that should not crash the extract chain. @@ -88,4 +126,18 @@ public final class TikaExtractor { return null; } } + + private static void checkInterrupted() throws InterruptedIOException { + if (Thread.currentThread().isInterrupted()) { + throw new InterruptedIOException("Document extraction interrupted"); + } + } + + private static void checkParseInterrupted() throws SAXException { + try { + checkInterrupted(); + } catch (InterruptedIOException e) { + throw new SAXException(e); + } + } } diff --git a/mateclaw-server/src/test/java/vip/mate/agent/AgentServiceToolTimeoutTest.java b/mateclaw-server/src/test/java/vip/mate/agent/AgentServiceToolTimeoutTest.java new file mode 100644 index 00000000..e772fe26 --- /dev/null +++ b/mateclaw-server/src/test/java/vip/mate/agent/AgentServiceToolTimeoutTest.java @@ -0,0 +1,66 @@ +package vip.mate.agent; + +import org.junit.jupiter.api.Test; +import org.springframework.ai.chat.messages.AssistantMessage; +import org.springframework.ai.tool.ToolCallback; +import org.springframework.ai.tool.definition.ToolDefinition; +import org.springframework.test.util.ReflectionTestUtils; +import reactor.core.publisher.Flux; +import vip.mate.agent.graph.executor.ToolExecutionExecutor; +import vip.mate.config.ToolTimeoutProperties; +import vip.mate.memory.MemoryProperties; + +import java.util.List; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.function.BiFunction; +import java.util.function.Function; + +import static org.junit.jupiter.api.Assertions.*; + +class AgentServiceToolTimeoutTest { + @Test + void singleToolTimeoutFinishesTurnAndSameConversationCanRunAgain() { + AtomicBoolean firstCall = new AtomicBoolean(true); + AtomicBoolean interrupted = new AtomicBoolean(); + ToolCallback tool = new ToolCallback() { + @Override public ToolDefinition getToolDefinition() { + return ToolDefinition.builder().name("extract_document_text") + .description("document extraction").inputSchema("{\"type\":\"object\"}").build(); + } + @Override public String call(String arguments) { + if (firstCall.getAndSet(false)) { + try { + Thread.sleep(3_000); + } catch (InterruptedException e) { + interrupted.set(true); + Thread.currentThread().interrupt(); + } + } + return "ok"; + } + }; + ToolTimeoutProperties timeouts = new ToolTimeoutProperties(); + timeouts.setDefaultTimeoutSeconds(1); + var executor = new ToolExecutionExecutor(AgentToolSet.fromCallbacks(List.of(), List.of(tool)), + null, null, null, timeouts); + var call = new AssistantMessage.ToolCall("call", "function", "extract_document_text", "{}"); + MemoryProperties memory = new MemoryProperties(); + memory.setLifecycleMediatorEnabled(false); + var service = new AgentService(null, null, null, null, memory, null, null); + BiFunction> invoke = (message, conversation) -> Flux.defer(() -> + Flux.just(executor.execute(List.of(call), conversation, "1", false) + .responses().getFirst().responseData())); + Function content = Function.identity(); + Flux first = ReflectionTestUtils.invokeMethod(service, "withLifecycleFlux", + 1L, "read spreadsheet", "same-conversation", invoke, content); + assertNotNull(first); + assertTrue(first.blockLast().contains("timed out")); + assertTrue(interrupted.get()); + assertFalse(Thread.currentThread().isInterrupted()); + + Flux second = ReflectionTestUtils.invokeMethod(service, "withLifecycleFlux", + 1L, "retry", "same-conversation", invoke, content); + assertNotNull(second); + assertEquals("ok", second.blockLast(), "the previous turn must release admission"); + } +} diff --git a/mateclaw-server/src/test/java/vip/mate/agent/graph/executor/ToolCallDeadlineTest.java b/mateclaw-server/src/test/java/vip/mate/agent/graph/executor/ToolCallDeadlineTest.java new file mode 100644 index 00000000..8a00d892 --- /dev/null +++ b/mateclaw-server/src/test/java/vip/mate/agent/graph/executor/ToolCallDeadlineTest.java @@ -0,0 +1,37 @@ +package vip.mate.agent.graph.executor; + +import org.junit.jupiter.api.Test; + +import java.util.concurrent.TimeoutException; + +import static org.junit.jupiter.api.Assertions.*; + +class ToolCallDeadlineTest { + @Test + void timeoutInterruptsCallbackAndDoesNotPoisonNextCall() throws Exception { + ThreadLocal context = new ThreadLocal<>(); + context.set("conversation"); + try { + assertThrows(TimeoutException.class, () -> ToolCallDeadline.call("extract_document_text", 30, () -> { + assertEquals("conversation", context.get()); + try { + Thread.sleep(2_000); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + return "partial result must not be treated as success"; + })); + assertFalse(Thread.currentThread().isInterrupted()); + assertEquals("next", ToolCallDeadline.call("next", 1000, () -> "next")); + } finally { + context.remove(); + } + } + + @Test + void successfulCallbackCancelsItsWatchdog() throws Exception { + assertEquals("ok", ToolCallDeadline.call("fast", 30, () -> "ok")); + Thread.sleep(80); + assertFalse(Thread.currentThread().isInterrupted()); + } +} diff --git a/mateclaw-server/src/test/java/vip/mate/tool/builtin/DocumentExtractToolSpreadsheetTest.java b/mateclaw-server/src/test/java/vip/mate/tool/builtin/DocumentExtractToolSpreadsheetTest.java new file mode 100644 index 00000000..a0cb64d0 --- /dev/null +++ b/mateclaw-server/src/test/java/vip/mate/tool/builtin/DocumentExtractToolSpreadsheetTest.java @@ -0,0 +1,44 @@ +package vip.mate.tool.builtin; + +import cn.hutool.json.JSONUtil; +import org.apache.poi.xssf.streaming.SXSSFWorkbook; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.Base64; +import java.util.Random; + +import static org.junit.jupiter.api.Assertions.*; + +class DocumentExtractToolSpreadsheetTest { + @Test + void largeSpreadsheetReturnsBoundedPreviewWithTruncation(@TempDir Path dir) throws Exception { + Path file = dir.resolve("large.xlsx"); + try (SXSSFWorkbook workbook = new SXSSFWorkbook(100)) { + var sheet = workbook.createSheet("Data"); + Random random = new Random(635); + byte[] bytes = new byte[750]; + for (int row = 0; row < 10_000; row++) { + random.nextBytes(bytes); + sheet.createRow(row).createCell(0).setCellValue( + "row-" + row + "-" + Base64.getEncoder().encodeToString(bytes)); + } + try (var out = Files.newOutputStream(file)) { + workbook.write(out); + } + } + assertTrue(Files.size(file) > 7_000_000, "exercise a real 7 MB+ XLSX upload"); + for (String options : new String[]{null, "{\"method\":\"tika\"}"}) { + var result = JSONUtil.parseObj(new DocumentExtractTool().extractTrustedDocument(file.toString(), options)); + assertTrue(result.getBool("success"), result.toString()); + assertTrue(result.getBool("truncated")); + String text = result.getStr("text"); + assertTrue(text.contains("row-0-")); + assertFalse(text.contains("row-9999-")); + assertTrue(text.length() < 501_000); + assertTrue(text.contains("总长度至少: 500001"), "parse must stop at the response budget"); + } + } +} diff --git a/mateclaw-server/src/test/java/vip/mate/tool/builtin/TikaExtractorTest.java b/mateclaw-server/src/test/java/vip/mate/tool/builtin/TikaExtractorTest.java index e8582808..f113e6e4 100644 --- a/mateclaw-server/src/test/java/vip/mate/tool/builtin/TikaExtractorTest.java +++ b/mateclaw-server/src/test/java/vip/mate/tool/builtin/TikaExtractorTest.java @@ -1,14 +1,19 @@ package vip.mate.tool.builtin; +import org.apache.tika.parser.AutoDetectParser; import org.junit.jupiter.api.DisplayName; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.io.TempDir; +import org.xml.sax.ContentHandler; import java.io.IOException; import java.nio.file.Files; import java.nio.file.Path; import static org.junit.jupiter.api.Assertions.*; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.mockConstruction; /** * RFC-051 §5.2: pin TikaExtractor's safety guarantees. @@ -20,6 +25,42 @@ import static org.junit.jupiter.api.Assertions.*; */ class TikaExtractorTest { + @Test + void stopsWhenCancellationArrivesAfterInputWasBuffered(@TempDir Path tmp) throws Exception { + Path file = tmp.resolve("buffered.txt"); + Files.writeString(file, "buffered document"); + try (var ignored = mockConstruction(AutoDetectParser.class, (parser, context) -> { + doAnswer(invocation -> { + ContentHandler handler = invocation.getArgument(1); + handler.startDocument(); + Thread.currentThread().interrupt(); + // Office parsers may already have buffered the input. The SAX + // callback must still observe cancellation without another read. + handler.characters("text".toCharArray(), 0, 4); + fail("cancelled parsing must not continue"); + return null; + }).when(parser).parse(any(), any(), any(), any()); + })) { + assertNull(TikaExtractor.extract(file)); + assertTrue(Thread.currentThread().isInterrupted()); + } finally { + Thread.interrupted(); + } + } + + @Test + void interruptedExtractionStopsAndPreservesCancellation(@TempDir Path tmp) throws IOException { + Path file = tmp.resolve("cancelled.txt"); + Files.writeString(file, "Do not parse after cancellation"); + Thread.currentThread().interrupt(); + try { + assertNull(TikaExtractor.extract(file)); + assertTrue(Thread.currentThread().isInterrupted()); + } finally { + Thread.interrupted(); + } + } + @Test @DisplayName("null path returns null without throwing") void nullPath() {