Skip to content

Commit 4471db4

Browse files
fix(selfHosted): align worker lifecycle behavior
## 简述 对齐 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
1 parent a8cfe7a commit 4471db4

8 files changed

Lines changed: 258 additions & 6 deletions

File tree

src/main/java/com/volcengine/ark/runtime/selfhosted/EnvironmentWorker.java

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -236,7 +236,9 @@ private ToolContext toolContext(String workdir, AtomicBoolean workStop) {
236236
copy.setEnv(new LinkedHashMap<>(context.getEnv()));
237237
}
238238
copy.setUnrestrictedPaths(options.unrestrictedPaths || context.isUnrestrictedPaths());
239-
copy.setToolTimeoutMillis(context.getToolTimeoutMillis());
239+
copy.setToolTimeoutMillis(options.toolTimeoutMillis > 0L
240+
? options.toolTimeoutMillis
241+
: context.getToolTimeoutMillis());
240242
copy.setCancelled(() -> closed.get() || workStop.get());
241243
return copy;
242244
}
@@ -408,6 +410,7 @@ public static class Options {
408410
private String workdir = ".";
409411
private boolean unrestrictedPaths;
410412
private ToolContext toolContext;
413+
private long toolTimeoutMillis;
411414
private ToolSet tools;
412415
private long maxIdleMillis = SelfHostedConstants.DEFAULT_MAX_IDLE_MILLIS;
413416
private Map<String, Tool> customTools = new LinkedHashMap<>();
@@ -438,6 +441,11 @@ public Options toolContext(ToolContext toolContext) {
438441
return this;
439442
}
440443

444+
public Options toolTimeoutMillis(long toolTimeoutMillis) {
445+
this.toolTimeoutMillis = toolTimeoutMillis;
446+
return this;
447+
}
448+
441449
public Options tools(ToolSet tools) {
442450
this.tools = tools;
443451
return this;

src/main/java/com/volcengine/ark/runtime/selfhosted/SelfHostedClient.java

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@ public class SelfHostedClient {
4747
private static final long LIFECYCLE_TIMEOUT_SECONDS = 10L;
4848

4949
private final ArkApi api;
50+
private final ArkApi streamApi;
5051
private final ArkApi heartbeatApi;
5152
private final ArkApi lifecycleApi;
5253
private final OkHttpClient httpClient;
@@ -70,6 +71,12 @@ private SelfHostedClient(Builder builder) {
7071
this.skillHubBaseUrl = trimTrailingSlash(builder.skillHubBaseUrl);
7172
Retrofit retrofit = ArkService.defaultRetrofit(this.httpClient, this.mapper, normalizeBaseUrl(builder.baseUrl), null);
7273
this.api = retrofit.create(ArkApi.class);
74+
OkHttpClient streamClient = this.httpClient.newBuilder()
75+
.callTimeout(0L, TimeUnit.MILLISECONDS)
76+
.build();
77+
Retrofit streamRetrofit = ArkService.defaultRetrofit(
78+
streamClient, this.mapper, normalizeBaseUrl(builder.baseUrl), null);
79+
this.streamApi = streamRetrofit.create(ArkApi.class);
7380
OkHttpClient.Builder heartbeatClientBuilder = this.httpClient.newBuilder()
7481
.retryOnConnectionFailure(false)
7582
.connectTimeout(HEARTBEAT_TIMEOUT_SECONDS, TimeUnit.SECONDS)
@@ -190,12 +197,12 @@ public SkillRef resolveSkill(SkillRef ref) {
190197

191198
public Call<ResponseBody> streamEvents(String sessionId) {
192199
require(sessionId, "sessionId");
193-
return api.streamSessionEvents(sessionId, Collections.<String, String>emptyMap());
200+
return streamApi.streamSessionEvents(sessionId, Collections.<String, String>emptyMap());
194201
}
195202

196203
public EventStream openEventStream(String sessionId) {
197204
require(sessionId, "sessionId");
198-
Call<ResponseBody> call = api.streamSessionEvents(sessionId, Collections.<String, String>emptyMap());
205+
Call<ResponseBody> call = streamApi.streamSessionEvents(sessionId, Collections.<String, String>emptyMap());
199206
try {
200207
return new EventStream(call, call.execute());
201208
} catch (IOException e) {

src/main/java/com/volcengine/ark/runtime/selfhosted/SessionToolRunner.java

Lines changed: 81 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -9,22 +9,38 @@
99
import java.util.List;
1010
import java.util.Map;
1111
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;
1217
import java.util.concurrent.LinkedBlockingQueue;
18+
import java.util.concurrent.RejectedExecutionException;
1319
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;
1423
import java.util.concurrent.atomic.AtomicReference;
1524
import java.util.function.BooleanSupplier;
1625
import java.util.logging.Level;
1726
import java.util.logging.Logger;
1827

1928
public class SessionToolRunner {
2029
private static final int STREAM_QUEUE_SIZE = 256;
30+
private static final long TOOL_WAIT_SLICE_MILLIS = 50L;
2131
private static final Logger LOGGER = Logger.getLogger(SessionToolRunner.class.getName());
32+
private static final AtomicInteger TOOL_THREAD_ID = new AtomicInteger();
2233
private final SelfHostedClient api;
2334
private final String sessionId;
2435
private final Options options;
2536
private final State state = new State();
2637
private final List<ToolCallResult> results = new ArrayList<>();
2738
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+
});
2844
private volatile boolean closed;
2945
private volatile EventStream activeStream;
3046

@@ -198,6 +214,7 @@ private void consumeList() throws IOException {
198214
public void close() {
199215
closed = true;
200216
closeStream(activeStream);
217+
toolExecutor.shutdownNow();
201218
}
202219

203220
public List<ToolCallResult> getResults() {
@@ -341,14 +358,75 @@ private void handleToolUse(Event event, boolean custom) throws IOException {
341358
}
342359

343360
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) {
344406
if (custom) {
345407
try {
346-
return options.customTools.get(event.getName()).execute(event.getInput(), options.toolContext);
408+
return options.customTools.get(event.getName()).execute(event.getInput(), context);
347409
} catch (RuntimeException error) {
348-
return ToolResult.error(error.getMessage());
410+
return ToolResult.error(errorText(error));
349411
}
350412
}
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();
352430
}
353431

354432
private void sendResult(String callId, Event source, boolean custom, String confirmation, Event out) throws IOException {

src/main/java/com/volcengine/ark/runtime/selfhosted/Tool.java

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,12 @@
33

44
package com.volcengine.ark.runtime.selfhosted;
55

6+
/**
7+
* Tool execution contract.
8+
*
9+
* <p>The runner enforces {@link ToolContext#getToolTimeoutMillis()}. Custom tools should also
10+
* poll {@link ToolContext#isCancelled()} so timed-out work releases resources promptly.
11+
*/
612
public interface Tool {
713
String name();
814

src/main/java/com/volcengine/ark/runtime/selfhosted/WorkPoller.java

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,12 @@
88
import java.util.logging.Level;
99
import java.util.logging.Logger;
1010

11+
/**
12+
* Serial work poller with optional ownership cleanup.
13+
*
14+
* <p>The default auto-stop behavior is intended for iterator-style serial processing. Callers
15+
* dispatching work concurrently must disable it and own heartbeat and stop.
16+
*/
1117
public class WorkPoller implements AutoCloseable {
1218
private static final long POLL_BACKOFF_CAP_MILLIS = 60000L;
1319

src/test/java/com/volcengine/ark/runtime/selfhosted/EnvironmentWorkerTest.java

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,9 +9,11 @@
99
import com.volcengine.ark.runtime.models.environment.HeartbeatWorkResponse;
1010
import com.volcengine.ark.runtime.models.environment.WorkState;
1111
import java.io.IOException;
12+
import java.lang.reflect.Method;
1213
import java.nio.file.Files;
1314
import java.util.concurrent.CountDownLatch;
1415
import java.util.concurrent.TimeUnit;
16+
import java.util.concurrent.atomic.AtomicBoolean;
1517
import java.util.concurrent.atomic.AtomicInteger;
1618
import okhttp3.MediaType;
1719
import okhttp3.OkHttpClient;
@@ -123,6 +125,58 @@ public void leaseNotExtendedWithoutTtlDoesNotStopWorkOwnedByAnotherWorker() thro
123125
assertEquals(0, client.stops.get());
124126
}
125127

128+
@Test
129+
public void workerToolTimeoutOverridesToolContext() throws Exception {
130+
ToolContext baseContext = new ToolContext(".");
131+
baseContext.setToolTimeoutMillis(5000L);
132+
EnvironmentWorker worker = new EnvironmentWorker(
133+
new SelfHostedClient("test-key"),
134+
new EnvironmentWorker.Options()
135+
.toolContext(baseContext)
136+
.toolTimeoutMillis(20L));
137+
Method method = EnvironmentWorker.class.getDeclaredMethod(
138+
"toolContext", String.class, AtomicBoolean.class);
139+
method.setAccessible(true);
140+
141+
ToolContext context = (ToolContext) method.invoke(worker, ".", new AtomicBoolean(false));
142+
143+
assertEquals(20L, context.getToolTimeoutMillis());
144+
}
145+
146+
@Test
147+
public void workerKeepsToolContextTimeoutWhenUnset() throws Exception {
148+
ToolContext baseContext = new ToolContext(".");
149+
baseContext.setToolTimeoutMillis(5000L);
150+
EnvironmentWorker worker = new EnvironmentWorker(
151+
new SelfHostedClient("test-key"),
152+
new EnvironmentWorker.Options().toolContext(baseContext));
153+
Method method = EnvironmentWorker.class.getDeclaredMethod(
154+
"toolContext", String.class, AtomicBoolean.class);
155+
method.setAccessible(true);
156+
157+
ToolContext context = (ToolContext) method.invoke(worker, ".", new AtomicBoolean(false));
158+
159+
assertEquals(5000L, context.getToolTimeoutMillis());
160+
}
161+
162+
@Test
163+
public void workerKeepsToolContextTimeoutWhenNonpositive() throws Exception {
164+
ToolContext baseContext = new ToolContext(".");
165+
baseContext.setToolTimeoutMillis(5000L);
166+
EnvironmentWorker worker = new EnvironmentWorker(
167+
new SelfHostedClient("test-key"),
168+
new EnvironmentWorker.Options()
169+
.toolContext(baseContext)
170+
.toolTimeoutMillis(-1L));
171+
Method method = EnvironmentWorker.class.getDeclaredMethod(
172+
"toolContext", String.class, AtomicBoolean.class);
173+
method.setAccessible(true);
174+
175+
ToolContext context = (ToolContext) method.invoke(worker, ".", new AtomicBoolean(false));
176+
177+
assertEquals(5000L, context.getToolTimeoutMillis());
178+
}
179+
126180
private static Response response(Request request, String body) throws IOException {
127181
return new Response.Builder()
128182
.request(request)

src/test/java/com/volcengine/ark/runtime/selfhosted/SelfHostedClientTest.java

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -198,6 +198,33 @@ public void heartbeatUsesOneLeaseBoundedAttempt() {
198198
assertEquals(15L, timeoutSeconds.get());
199199
}
200200

201+
@Test
202+
public void eventStreamPreservesReadTimeoutWithoutTotalTimeout() throws Exception {
203+
AtomicLong readTimeoutMillis = new AtomicLong(-1L);
204+
AtomicLong callTimeoutNanos = new AtomicLong(-1L);
205+
OkHttpClient httpClient = new OkHttpClient.Builder()
206+
.readTimeout(25L, TimeUnit.MILLISECONDS)
207+
.callTimeout(25L, TimeUnit.MILLISECONDS)
208+
.addInterceptor(chain -> {
209+
readTimeoutMillis.set(chain.readTimeoutMillis());
210+
callTimeoutNanos.set(chain.call().timeout().timeoutNanos());
211+
return response(chain.request(), new AtomicReference<>(), "");
212+
})
213+
.build();
214+
SelfHostedClient client = new SelfHostedClient.Builder()
215+
.apiKey("test-api-key")
216+
.httpClient(httpClient)
217+
.build();
218+
219+
try (ResponseBody ignored = client.streamEvents("session-1").execute().body()) {
220+
assertEquals(25L, readTimeoutMillis.get());
221+
assertEquals(0L, callTimeoutNanos.get());
222+
}
223+
224+
assertEquals(25, httpClient.readTimeoutMillis());
225+
assertEquals(25, httpClient.callTimeoutMillis());
226+
}
227+
201228
@Test
202229
public void atomicEnvironmentWorkAPIMatchesOpenAPIContract() throws Exception {
203230
AtomicInteger calls = new AtomicInteger();

src/test/java/com/volcengine/ark/runtime/selfhosted/SessionToolRunnerTest.java

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -193,6 +193,72 @@ public void markSent(String callId) throws IOException {
193193
assertTrue(runner.getResults().get(0).isPosted());
194194
}
195195

196+
@Test
197+
public void builtinToolTimeoutAbandonsNoncooperativeTool() throws Exception {
198+
assertToolTimeoutAbandonsNoncooperativeTool(false);
199+
}
200+
201+
@Test
202+
public void customToolTimeoutAbandonsNoncooperativeTool() throws Exception {
203+
assertToolTimeoutAbandonsNoncooperativeTool(true);
204+
}
205+
206+
private static void assertToolTimeoutAbandonsNoncooperativeTool(boolean custom) throws Exception {
207+
CountDownLatch started = new CountDownLatch(1);
208+
CountDownLatch release = new CountDownLatch(1);
209+
Tool tool = new Tool() {
210+
@Override
211+
public String name() {
212+
return "blocking";
213+
}
214+
215+
@Override
216+
public ToolResult execute(Object input, ToolContext context) {
217+
started.countDown();
218+
while (true) {
219+
try {
220+
release.await();
221+
return ToolResult.text("late");
222+
} catch (InterruptedException ignored) {
223+
// Deliberately ignore cancellation to verify the runner's outer deadline.
224+
}
225+
}
226+
}
227+
};
228+
ToolContext context = new ToolContext(Files.createTempDirectory("ark-java-timeout-").toString());
229+
context.setToolTimeoutMillis(20L);
230+
SessionToolRunner.Options options = new SessionToolRunner.Options()
231+
.tools(custom ? new ToolSet() : new ToolSet().add(tool))
232+
.toolContext(context);
233+
if (custom) {
234+
options.customTools(Collections.singletonMap(tool.name(), tool));
235+
}
236+
SessionToolRunner runner = new SessionToolRunner(new SelfHostedClient("test-key"), "session-1", options);
237+
Map<String, Object> raw = new LinkedHashMap<>();
238+
raw.put("id", "tool-1");
239+
raw.put("type", custom ? "agent.custom_tool_use" : "agent.tool_use");
240+
raw.put("name", tool.name());
241+
raw.put(custom ? "custom_tool_use_id" : "tool_use_id", "call-1");
242+
raw.put("input", Collections.emptyMap());
243+
Method execute = SessionToolRunner.class.getDeclaredMethod("executeTool", Event.class, boolean.class);
244+
execute.setAccessible(true);
245+
246+
long startedAt = System.nanoTime();
247+
ToolResult result;
248+
try {
249+
result = (ToolResult) execute.invoke(runner, Event.fromMap(raw), custom);
250+
assertTrue(started.await(1L, TimeUnit.SECONDS));
251+
} finally {
252+
release.countDown();
253+
runner.close();
254+
}
255+
256+
long elapsedMillis = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startedAt);
257+
assertTrue("elapsed=" + elapsedMillis, elapsedMillis < 500L);
258+
assertTrue(result.isError());
259+
assertEquals("tool execution timed out after 20ms", result.getContent().get(0).getText());
260+
}
261+
196262
private static SessionToolRunner idleRunner() throws IOException {
197263
return new SessionToolRunner(
198264
new SelfHostedClient("test-key"),

0 commit comments

Comments
 (0)