From 07ddaeba0bf2cd14f31224d0ab584cfe98f4de45 Mon Sep 17 00:00:00 2001 From: "ark-hand[bot]" <315378070+ark-hand[bot]@users.noreply.github.com> Date: Tue, 1 Sep 2026 03:42:35 +0000 Subject: [PATCH] fix(selfHosted): align worker lifecycle behavior MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## 简述 对齐 Java self-hosted worker 的 SSE 生命周期和工具执行超时语义。 ## 修改前 - self-hosted SSE 复用普通 OkHttp client，整次 call timeout 可能终止长时间 session。 - EnvironmentWorker.Options 无法直接配置单次工具执行超时。 - timeout 只进入 ToolContext，非协作 custom tool 如果不主动检查 isCancelled，仍会阻塞 SessionToolRunner。 - 初版 SSE 修复同时关闭 readTimeout，会让静默网络断流永久挂起。 ## 修改后 - 为 self-hosted SSE 创建独立 OkHttp 和 Retrofit client，保留认证、代理、拦截器及调用方 read timeout，仅将整次 callTimeout 设为 0。 - EnvironmentWorker.Options 增加 toolTimeoutMillis；正值覆盖 ToolContext，非正值保留原 120s 默认值或用户配置。 - SessionToolRunner 使用 daemon executor 执行工具，并在外层强制 deadline；以 50ms 间隔观察 worker 和父 context 取消，超时后立即返回 error result，不等待忽略中断的工具。 - timeout 和取消会同步到执行用 ToolContext；配合型工具可及时释放进程和资源。 - WorkPoller 文档明确 AutoStop 是串行 iterator 语义；并发 dispatch 必须关闭并自行维护 heartbeat 和 stop。 ## 边界 - 非协作工具线程可能继续运行到自身返回，但结果会被丢弃；该行为语义对齐 Anthropic Go/Python 的 abandon-on-cancel 实现，工具仍应响应 isCancelled 或线程中断。 - 本 MR 不修改 SIGTERM 时 force-stop 行为，不实现 work 迁移、release 或 requeue。 ## 验证 - Java 全量单测：47 passed - Checkstyle 通过 - Maven package 通过 - test/run.sh --sdk：Go、Python、Java 真实 STG worker 3/3 通过 - test/run.sh --all：真实 STG 全量默认套件通过，Docker 和直接 work-contract 用例按独立模式开关跳过 See merge request: !88 Sync-Source-Commit: f328f80ce14608dfc1e0b62436ccbbfd0727e3bf Hand-Written-Reason: Hand-written self-hosted worker lifecycle alignment; not produced by ark-apis generation. Release-Version: 0.4.0 --- .../runtime/selfhosted/EnvironmentWorker.java | 10 ++- .../runtime/selfhosted/SelfHostedClient.java | 11 ++- .../runtime/selfhosted/SessionToolRunner.java | 84 ++++++++++++++++++- .../ark/runtime/selfhosted/Tool.java | 6 ++ .../ark/runtime/selfhosted/WorkPoller.java | 6 ++ .../selfhosted/EnvironmentWorkerTest.java | 54 ++++++++++++ .../selfhosted/SelfHostedClientTest.java | 27 ++++++ .../selfhosted/SessionToolRunnerTest.java | 66 +++++++++++++++ 8 files changed, 258 insertions(+), 6 deletions(-) diff --git a/src/main/java/com/volcengine/ark/runtime/selfhosted/EnvironmentWorker.java b/src/main/java/com/volcengine/ark/runtime/selfhosted/EnvironmentWorker.java index fa7b997..381d89c 100644 --- a/src/main/java/com/volcengine/ark/runtime/selfhosted/EnvironmentWorker.java +++ b/src/main/java/com/volcengine/ark/runtime/selfhosted/EnvironmentWorker.java @@ -236,7 +236,9 @@ private ToolContext toolContext(String workdir, AtomicBoolean workStop) { copy.setEnv(new LinkedHashMap<>(context.getEnv())); } copy.setUnrestrictedPaths(options.unrestrictedPaths || context.isUnrestrictedPaths()); - copy.setToolTimeoutMillis(context.getToolTimeoutMillis()); + copy.setToolTimeoutMillis(options.toolTimeoutMillis > 0L + ? options.toolTimeoutMillis + : context.getToolTimeoutMillis()); copy.setCancelled(() -> closed.get() || workStop.get()); return copy; } @@ -408,6 +410,7 @@ public static class Options { private String workdir = "."; private boolean unrestrictedPaths; private ToolContext toolContext; + private long toolTimeoutMillis; private ToolSet tools; private long maxIdleMillis = SelfHostedConstants.DEFAULT_MAX_IDLE_MILLIS; private Map customTools = new LinkedHashMap<>(); @@ -438,6 +441,11 @@ public Options toolContext(ToolContext toolContext) { return this; } + public Options toolTimeoutMillis(long toolTimeoutMillis) { + this.toolTimeoutMillis = toolTimeoutMillis; + return this; + } + public Options tools(ToolSet tools) { this.tools = tools; return this; diff --git a/src/main/java/com/volcengine/ark/runtime/selfhosted/SelfHostedClient.java b/src/main/java/com/volcengine/ark/runtime/selfhosted/SelfHostedClient.java index 617a3f3..40bc205 100644 --- a/src/main/java/com/volcengine/ark/runtime/selfhosted/SelfHostedClient.java +++ b/src/main/java/com/volcengine/ark/runtime/selfhosted/SelfHostedClient.java @@ -47,6 +47,7 @@ public class SelfHostedClient { private static final long LIFECYCLE_TIMEOUT_SECONDS = 10L; private final ArkApi api; + private final ArkApi streamApi; private final ArkApi heartbeatApi; private final ArkApi lifecycleApi; private final OkHttpClient httpClient; @@ -70,6 +71,12 @@ private SelfHostedClient(Builder builder) { this.skillHubBaseUrl = trimTrailingSlash(builder.skillHubBaseUrl); Retrofit retrofit = ArkService.defaultRetrofit(this.httpClient, this.mapper, normalizeBaseUrl(builder.baseUrl), null); this.api = retrofit.create(ArkApi.class); + OkHttpClient streamClient = this.httpClient.newBuilder() + .callTimeout(0L, TimeUnit.MILLISECONDS) + .build(); + Retrofit streamRetrofit = ArkService.defaultRetrofit( + streamClient, this.mapper, normalizeBaseUrl(builder.baseUrl), null); + this.streamApi = streamRetrofit.create(ArkApi.class); OkHttpClient.Builder heartbeatClientBuilder = this.httpClient.newBuilder() .retryOnConnectionFailure(false) .connectTimeout(HEARTBEAT_TIMEOUT_SECONDS, TimeUnit.SECONDS) @@ -190,12 +197,12 @@ public SkillRef resolveSkill(SkillRef ref) { public Call streamEvents(String sessionId) { require(sessionId, "sessionId"); - return api.streamSessionEvents(sessionId, Collections.emptyMap()); + return streamApi.streamSessionEvents(sessionId, Collections.emptyMap()); } public EventStream openEventStream(String sessionId) { require(sessionId, "sessionId"); - Call call = api.streamSessionEvents(sessionId, Collections.emptyMap()); + Call call = streamApi.streamSessionEvents(sessionId, Collections.emptyMap()); try { return new EventStream(call, call.execute()); } catch (IOException e) { 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 f0c58c0..eb1b535 100644 --- a/src/main/java/com/volcengine/ark/runtime/selfhosted/SessionToolRunner.java +++ b/src/main/java/com/volcengine/ark/runtime/selfhosted/SessionToolRunner.java @@ -9,8 +9,17 @@ import java.util.List; import java.util.Map; import java.util.Random; +import java.util.concurrent.CancellationException; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.RejectedExecutionException; import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; import java.util.function.BooleanSupplier; import java.util.logging.Level; @@ -18,13 +27,20 @@ public class SessionToolRunner { private static final int STREAM_QUEUE_SIZE = 256; + private static final long TOOL_WAIT_SLICE_MILLIS = 50L; private static final Logger LOGGER = Logger.getLogger(SessionToolRunner.class.getName()); + private static final AtomicInteger TOOL_THREAD_ID = new AtomicInteger(); private final SelfHostedClient api; private final String sessionId; private final Options options; private final State state = new State(); private final List results = new ArrayList<>(); private final Random random = new Random(); + private final ExecutorService toolExecutor = Executors.newCachedThreadPool(runnable -> { + Thread thread = new Thread(runnable, "ma-self-host-tool-" + TOOL_THREAD_ID.incrementAndGet()); + thread.setDaemon(true); + return thread; + }); private volatile boolean closed; private volatile EventStream activeStream; @@ -198,6 +214,7 @@ private void consumeList() throws IOException { public void close() { closed = true; closeStream(activeStream); + toolExecutor.shutdownNow(); } public List getResults() { @@ -341,14 +358,75 @@ private void handleToolUse(Event event, boolean custom) throws IOException { } private ToolResult executeTool(Event event, boolean custom) { + 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 { + future = toolExecutor.submit(() -> executeTool(event, custom, context)); + } catch (RejectedExecutionException error) { + return ToolResult.error("tool execution canceled"); + } + long deadline = System.nanoTime() + TimeUnit.MILLISECONDS.toNanos(timeoutMillis); + while (true) { + if (context.isCancelled()) { + executionCancelled.set(true); + future.cancel(true); + return ToolResult.error("tool execution canceled"); + } + long remaining = deadline - System.nanoTime(); + if (remaining <= 0L) { + executionCancelled.set(true); + future.cancel(true); + return ToolResult.error("tool execution timed out after " + timeoutMillis + "ms"); + } + long waitNanos = Math.min(remaining, TimeUnit.MILLISECONDS.toNanos(TOOL_WAIT_SLICE_MILLIS)); + try { + return future.get(waitNanos, TimeUnit.NANOSECONDS); + } catch (TimeoutException ignored) { + // Continue so cancellation is observed without waiting for the full tool timeout. + } catch (CancellationException error) { + return ToolResult.error("tool execution canceled"); + } catch (InterruptedException error) { + Thread.currentThread().interrupt(); + executionCancelled.set(true); + future.cancel(true); + return ToolResult.error("tool execution canceled"); + } catch (ExecutionException error) { + Throwable cause = error.getCause(); + return ToolResult.error(cause == null ? error.toString() : errorText(cause)); + } + } + } + + private ToolResult executeTool(Event event, boolean custom, ToolContext context) { if (custom) { try { - return options.customTools.get(event.getName()).execute(event.getInput(), options.toolContext); + return options.customTools.get(event.getName()).execute(event.getInput(), context); } catch (RuntimeException error) { - return ToolResult.error(error.getMessage()); + return ToolResult.error(errorText(error)); } } - return options.tools.execute(event.getName(), event.getInput(), options.toolContext); + return options.tools.execute(event.getName(), event.getInput(), context); + } + + private ToolContext toolContextForExecution(long timeoutMillis, AtomicBoolean executionCancelled) { + ToolContext source = options.toolContext; + ToolContext context = new ToolContext(source.getWorkdir()); + if (source.hasExplicitEnv()) { + context.setEnv(new LinkedHashMap<>(source.getEnv())); + } + context.setUnrestrictedPaths(source.isUnrestrictedPaths()); + context.setToolTimeoutMillis(timeoutMillis); + context.setCancelled(() -> executionCancelled.get() || isClosed() || source.isCancelled()); + return context; + } + + 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 { diff --git a/src/main/java/com/volcengine/ark/runtime/selfhosted/Tool.java b/src/main/java/com/volcengine/ark/runtime/selfhosted/Tool.java index 6438504..9cdbc6a 100644 --- a/src/main/java/com/volcengine/ark/runtime/selfhosted/Tool.java +++ b/src/main/java/com/volcengine/ark/runtime/selfhosted/Tool.java @@ -3,6 +3,12 @@ package com.volcengine.ark.runtime.selfhosted; +/** + * Tool execution contract. + * + *

The runner enforces {@link ToolContext#getToolTimeoutMillis()}. Custom tools should also + * poll {@link ToolContext#isCancelled()} so timed-out work releases resources promptly. + */ public interface Tool { String name(); diff --git a/src/main/java/com/volcengine/ark/runtime/selfhosted/WorkPoller.java b/src/main/java/com/volcengine/ark/runtime/selfhosted/WorkPoller.java index 857a830..18526ac 100644 --- a/src/main/java/com/volcengine/ark/runtime/selfhosted/WorkPoller.java +++ b/src/main/java/com/volcengine/ark/runtime/selfhosted/WorkPoller.java @@ -8,6 +8,12 @@ import java.util.logging.Level; import java.util.logging.Logger; +/** + * Serial work poller with optional ownership cleanup. + * + *

The default auto-stop behavior is intended for iterator-style serial processing. Callers + * dispatching work concurrently must disable it and own heartbeat and stop. + */ public class WorkPoller implements AutoCloseable { private static final long POLL_BACKOFF_CAP_MILLIS = 60000L; diff --git a/src/test/java/com/volcengine/ark/runtime/selfhosted/EnvironmentWorkerTest.java b/src/test/java/com/volcengine/ark/runtime/selfhosted/EnvironmentWorkerTest.java index e4dc1b7..71715c2 100644 --- a/src/test/java/com/volcengine/ark/runtime/selfhosted/EnvironmentWorkerTest.java +++ b/src/test/java/com/volcengine/ark/runtime/selfhosted/EnvironmentWorkerTest.java @@ -9,9 +9,11 @@ import com.volcengine.ark.runtime.models.environment.HeartbeatWorkResponse; import com.volcengine.ark.runtime.models.environment.WorkState; import java.io.IOException; +import java.lang.reflect.Method; import java.nio.file.Files; import java.util.concurrent.CountDownLatch; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; import okhttp3.MediaType; import okhttp3.OkHttpClient; @@ -123,6 +125,58 @@ public void leaseNotExtendedWithoutTtlDoesNotStopWorkOwnedByAnotherWorker() thro assertEquals(0, client.stops.get()); } + @Test + public void workerToolTimeoutOverridesToolContext() throws Exception { + ToolContext baseContext = new ToolContext("."); + baseContext.setToolTimeoutMillis(5000L); + EnvironmentWorker worker = new EnvironmentWorker( + new SelfHostedClient("test-key"), + new EnvironmentWorker.Options() + .toolContext(baseContext) + .toolTimeoutMillis(20L)); + Method method = EnvironmentWorker.class.getDeclaredMethod( + "toolContext", String.class, AtomicBoolean.class); + method.setAccessible(true); + + ToolContext context = (ToolContext) method.invoke(worker, ".", new AtomicBoolean(false)); + + assertEquals(20L, context.getToolTimeoutMillis()); + } + + @Test + public void workerKeepsToolContextTimeoutWhenUnset() throws Exception { + ToolContext baseContext = new ToolContext("."); + baseContext.setToolTimeoutMillis(5000L); + EnvironmentWorker worker = new EnvironmentWorker( + new SelfHostedClient("test-key"), + new EnvironmentWorker.Options().toolContext(baseContext)); + Method method = EnvironmentWorker.class.getDeclaredMethod( + "toolContext", String.class, AtomicBoolean.class); + method.setAccessible(true); + + ToolContext context = (ToolContext) method.invoke(worker, ".", new AtomicBoolean(false)); + + assertEquals(5000L, context.getToolTimeoutMillis()); + } + + @Test + public void workerKeepsToolContextTimeoutWhenNonpositive() throws Exception { + ToolContext baseContext = new ToolContext("."); + baseContext.setToolTimeoutMillis(5000L); + EnvironmentWorker worker = new EnvironmentWorker( + new SelfHostedClient("test-key"), + new EnvironmentWorker.Options() + .toolContext(baseContext) + .toolTimeoutMillis(-1L)); + Method method = EnvironmentWorker.class.getDeclaredMethod( + "toolContext", String.class, AtomicBoolean.class); + method.setAccessible(true); + + ToolContext context = (ToolContext) method.invoke(worker, ".", new AtomicBoolean(false)); + + assertEquals(5000L, context.getToolTimeoutMillis()); + } + private static Response response(Request request, String body) throws IOException { return new Response.Builder() .request(request) diff --git a/src/test/java/com/volcengine/ark/runtime/selfhosted/SelfHostedClientTest.java b/src/test/java/com/volcengine/ark/runtime/selfhosted/SelfHostedClientTest.java index a081f6e..128fde4 100644 --- a/src/test/java/com/volcengine/ark/runtime/selfhosted/SelfHostedClientTest.java +++ b/src/test/java/com/volcengine/ark/runtime/selfhosted/SelfHostedClientTest.java @@ -198,6 +198,33 @@ public void heartbeatUsesOneLeaseBoundedAttempt() { assertEquals(15L, timeoutSeconds.get()); } + @Test + public void eventStreamPreservesReadTimeoutWithoutTotalTimeout() throws Exception { + AtomicLong readTimeoutMillis = new AtomicLong(-1L); + AtomicLong callTimeoutNanos = new AtomicLong(-1L); + OkHttpClient httpClient = new OkHttpClient.Builder() + .readTimeout(25L, TimeUnit.MILLISECONDS) + .callTimeout(25L, TimeUnit.MILLISECONDS) + .addInterceptor(chain -> { + readTimeoutMillis.set(chain.readTimeoutMillis()); + callTimeoutNanos.set(chain.call().timeout().timeoutNanos()); + return response(chain.request(), new AtomicReference<>(), ""); + }) + .build(); + SelfHostedClient client = new SelfHostedClient.Builder() + .apiKey("test-api-key") + .httpClient(httpClient) + .build(); + + try (ResponseBody ignored = client.streamEvents("session-1").execute().body()) { + assertEquals(25L, readTimeoutMillis.get()); + assertEquals(0L, callTimeoutNanos.get()); + } + + assertEquals(25, httpClient.readTimeoutMillis()); + assertEquals(25, httpClient.callTimeoutMillis()); + } + @Test public void atomicEnvironmentWorkAPIMatchesOpenAPIContract() throws Exception { AtomicInteger calls = new AtomicInteger(); 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 5694d2a..c199e98 100644 --- a/src/test/java/com/volcengine/ark/runtime/selfhosted/SessionToolRunnerTest.java +++ b/src/test/java/com/volcengine/ark/runtime/selfhosted/SessionToolRunnerTest.java @@ -193,6 +193,72 @@ public void markSent(String callId) throws IOException { assertTrue(runner.getResults().get(0).isPosted()); } + @Test + public void builtinToolTimeoutAbandonsNoncooperativeTool() throws Exception { + assertToolTimeoutAbandonsNoncooperativeTool(false); + } + + @Test + public void customToolTimeoutAbandonsNoncooperativeTool() throws Exception { + assertToolTimeoutAbandonsNoncooperativeTool(true); + } + + private static void assertToolTimeoutAbandonsNoncooperativeTool(boolean custom) throws Exception { + CountDownLatch started = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + Tool tool = new Tool() { + @Override + public String name() { + return "blocking"; + } + + @Override + public ToolResult execute(Object input, ToolContext context) { + started.countDown(); + while (true) { + try { + release.await(); + return ToolResult.text("late"); + } catch (InterruptedException ignored) { + // Deliberately ignore cancellation to verify the runner's outer deadline. + } + } + } + }; + ToolContext context = new ToolContext(Files.createTempDirectory("ark-java-timeout-").toString()); + context.setToolTimeoutMillis(20L); + SessionToolRunner.Options options = new SessionToolRunner.Options() + .tools(custom ? new ToolSet() : new ToolSet().add(tool)) + .toolContext(context); + if (custom) { + options.customTools(Collections.singletonMap(tool.name(), tool)); + } + SessionToolRunner runner = new SessionToolRunner(new SelfHostedClient("test-key"), "session-1", options); + Map raw = new LinkedHashMap<>(); + raw.put("id", "tool-1"); + raw.put("type", custom ? "agent.custom_tool_use" : "agent.tool_use"); + 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); + execute.setAccessible(true); + + long startedAt = System.nanoTime(); + ToolResult result; + try { + result = (ToolResult) execute.invoke(runner, Event.fromMap(raw), custom); + assertTrue(started.await(1L, TimeUnit.SECONDS)); + } finally { + release.countDown(); + runner.close(); + } + + long elapsedMillis = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startedAt); + assertTrue("elapsed=" + elapsedMillis, elapsedMillis < 500L); + assertTrue(result.isError()); + assertEquals("tool execution timed out after 20ms", result.getContent().get(0).getText()); + } + private static SessionToolRunner idleRunner() throws IOException { return new SessionToolRunner( new SelfHostedClient("test-key"),