This is an automated email from the ASF dual-hosted git repository.
lizhimins pushed a commit to branch rocketmq-studio
in repository https://gitbox.apache.org/repos/asf/rocketmq-dashboard.git
The following commit(s) were added to refs/heads/rocketmq-studio by this push:
new d712772bc fix(ai): enforce a shared run deadline across provider
stages (#7044)
d712772bc is described below
commit d712772bc08c21facd01d5c3436fa7a2dd922d83
Author: aias00 <[email protected]>
AuthorDate: Mon Oct 26 16:33:37 2026 +0800
fix(ai): enforce a shared run deadline across provider stages (#7044)
* fix(ai): enforce a shared run deadline across provider stages
---
docs/ai-agent-architecture.md | 6 ++
.../studio/ops/ai/conversation/AiRunExecutor.java | 12 ++++
.../ops/ai/conversation/AiRunExecutorTest.java | 83 ++++++++++++++++++++++
3 files changed, 101 insertions(+)
diff --git a/docs/ai-agent-architecture.md b/docs/ai-agent-architecture.md
index caa28fa17..3d42f78a2 100644
--- a/docs/ai-agent-architecture.md
+++ b/docs/ai-agent-architecture.md
@@ -166,6 +166,12 @@ onError` 接到 `cancel()`，TCP 一断就杀 worker。改造后：事件先落
代价要说清楚：关掉标签页后 Agent 会跑到完，由 CLI 超时与孤儿 run 清扫兜住。换来的是刷新不丢、
崩溃可恢复、以及「run 在没有客户端时也能完成」这个可验证性质。
+整轮 worker 有一个共享墙钟期限，包含提示增强、正式生成和丢失 resume 后的重试，而不是每个阶段
+重新获得一整份预算。期限从 worker 开始执行时计算，预留最多 5 秒（短预算预留 10%）供停止和
+终态写入；到期通过现有 handle 记录 `TIMEOUT`、停止当前执行，下一次 provider 调用会被拒绝。
+完成时取消定时任务。它不把用户断线改成取消；底层 I/O 的中断响应、可配置的停止宽限和数据库
+延迟仍可能延后最终退出，因此预留窗口不是持久化一定在连接关闭前完成的保证。
+
## 8. 取消
三层，全部落在 `status=STOPPED/FAILED`：
diff --git
a/server/src/main/java/org/apache/rocketmq/studio/ops/ai/conversation/AiRunExecutor.java
b/server/src/main/java/org/apache/rocketmq/studio/ops/ai/conversation/AiRunExecutor.java
index 400c703b3..f5e85e8e6 100644
---
a/server/src/main/java/org/apache/rocketmq/studio/ops/ai/conversation/AiRunExecutor.java
+++
b/server/src/main/java/org/apache/rocketmq/studio/ops/ai/conversation/AiRunExecutor.java
@@ -53,6 +53,7 @@ import java.util.concurrent.ExecutorService;
import java.util.concurrent.Future;
import java.util.concurrent.RejectedExecutionException;
import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.ScheduledThreadPoolExecutor;
import java.util.concurrent.SynchronousQueue;
import java.util.concurrent.ThreadPoolExecutor;
@@ -276,7 +277,13 @@ public class AiRunExecutor {
private void execute(RunContext context) {
RmqAiRun run = context.getRun();
Outcome outcome = new Outcome();
+ ScheduledFuture<?> deadline = null;
try {
+ long timeoutMillis = context.getTimeout().toMillis();
+ long reserve = Math.min(5_000L, timeoutMillis / 10);
+ deadline = scheduler.schedule(() -> {
+ if (!context.isTerminal())
context.getHandle().stop(AbortReason.TIMEOUT);
+ }, Math.max(1L, timeoutMillis - reserve), TimeUnit.MILLISECONDS);
markRunning(context);
streamWithLostResumeRetry(context, outcome);
} catch (LlmGatewayException exception) {
@@ -294,6 +301,7 @@ public class AiRunExecutor {
outcome.unexpected = error;
throw error;
} finally {
+ if (deadline != null) deadline.cancel(false);
try {
finalizeRun(context, decide(context, outcome));
} catch (RuntimeException exception) {
@@ -359,6 +367,10 @@ public class AiRunExecutor {
}
private void stream(RunContext context, String prompt, Outcome outcome,
boolean allowResume) {
+ if (context.getHandle().isStopRequested()) {
+ throw new LlmGatewayException(504, PROVIDER_TIMEOUT_CODE,
+ "Agent run was cancelled before the next provider call",
"Retry the request.");
+ }
Consumer<AgentEvent> events = event -> onAgentEvent(context, outcome,
event);
if (isHttpEngine(context.getEngine())) {
streamHttp(context, prompt, events);
diff --git
a/server/src/test/java/org/apache/rocketmq/studio/ops/ai/conversation/AiRunExecutorTest.java
b/server/src/test/java/org/apache/rocketmq/studio/ops/ai/conversation/AiRunExecutorTest.java
index c423a7d51..f8e33c0bc 100644
---
a/server/src/test/java/org/apache/rocketmq/studio/ops/ai/conversation/AiRunExecutorTest.java
+++
b/server/src/test/java/org/apache/rocketmq/studio/ops/ai/conversation/AiRunExecutorTest.java
@@ -46,8 +46,11 @@ import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CopyOnWriteArrayList;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.TimeUnit;
import java.util.function.Consumer;
@@ -162,6 +165,86 @@ class AiRunExecutorTest {
registry.liveRunIds().forEach(registry::unregister);
}
+ @Test
+ void realDeadlineShouldInterruptEnhancementAndPreventGenerationTest()
throws Exception {
+ ExecutorService worker = Executors.newSingleThreadExecutor();
+ ScheduledExecutorService timers = Executors.newScheduledThreadPool(2);
+ CountDownLatch enhancing = new CountDownLatch(1);
+ CountDownLatch terminal = new CountDownLatch(1);
+ try {
+ executor = new AiRunExecutor(new
AgentProviderRegistry(List.of(provider)), llmClient,
+ promptEnhancer, registry, runRepository,
conversationRepository, eventRepository,
+ properties(), objectMapper, worker,
AiRunTestSupport.recordingInto(emitters),
+ timers, Clock.systemUTC());
+ when(promptEnhancer.enhance(any(), any(), any(),
any())).thenAnswer(invocation -> {
+ enhancing.countDown();
+ try {
+ new CountDownLatch(1).await();
+ } catch (InterruptedException interrupted) {
+ Thread.currentThread().interrupt();
+ }
+ return "raw prompt";
+ });
+ doAnswer(invocation -> {
+ RmqAiRun updated = invocation.getArgument(0);
+ runUpdates.add(AiRunTestSupport.copyOf(updated));
+ if (RunStatus.FAILED.name().equals(updated.getStatus()))
terminal.countDown();
+ return null;
+ }).when(runRepository).update(any(RmqAiRun.class));
+ AiRunExecutor.RunContext base = context(ENGINE, true);
+ AiRunExecutor.RunContext shortRun = new
AiRunExecutor.RunContext(conversation, run,
+ base.getSink(), base.getHandle(), base.getConfig(),
ENGINE, "hello", null,
+ true, Duration.ofMillis(500), null);
+ executor.submit(shortRun);
+ assertThat(enhancing.await(3, TimeUnit.SECONDS)).isTrue();
+ assertThat(terminal.await(3, TimeUnit.SECONDS)).isTrue();
+
assertThat(runRow().getStopReason()).isEqualTo(StopReason.TIMEOUT.name());
+ assertThat(provider.wasCalled()).isFalse();
+ } finally {
+ timers.shutdownNow();
+ worker.shutdownNow();
+ worker.awaitTermination(3, TimeUnit.SECONDS);
+ }
+ }
+
+ @Test
+ void deadlineDuringEnhancementShouldNotStartAnotherProviderCallTest() {
+ Runnable[] deadline = new Runnable[1];
+ ScheduledFuture<?> scheduled = mock(ScheduledFuture.class);
+ when(scheduler.schedule(any(Runnable.class), anyLong(),
eq(TimeUnit.MILLISECONDS)))
+ .thenAnswer(invocation -> {
+ deadline[0] = invocation.getArgument(0);
+ return scheduled;
+ });
+ when(promptEnhancer.enhance(any(), any(), any(),
any())).thenAnswer(invocation -> {
+ if (deadline[0] != null) deadline[0].run();
+ return "rewritten";
+ });
+
+ startAndRun(true);
+
+ assertThat(provider.wasCalled()).isFalse();
+ assertThat(runRow().getStatus()).isEqualTo(RunStatus.FAILED.name());
+
assertThat(runRow().getStopReason()).isEqualTo(StopReason.TIMEOUT.name());
+ verify(scheduled).cancel(false);
+
assertThat(emitters.get(0).eventCount("\"type\":\"run_finished\"")).isEqualTo(1);
+ }
+
+ @Test
+ void successfulCompletionShouldCancelTheRunDeadlineTest() {
+ ScheduledFuture<?> scheduled = mock(ScheduledFuture.class);
+ when(scheduler.schedule(any(Runnable.class), anyLong(),
eq(TimeUnit.MILLISECONDS)))
+ .thenAnswer(invocation -> scheduled);
+ provider.emit(new AgentEvent.ResultMeta("session-9", 12L, 3, 9,
+ AgentEventProjector.SUCCESS_SUBTYPE));
+
+ startAndRun();
+
+ assertThat(runRow().getStatus()).isEqualTo(RunStatus.COMPLETED.name());
+ verify(scheduler).schedule(any(Runnable.class), eq(295_000L),
eq(TimeUnit.MILLISECONDS));
+ verify(scheduled).cancel(false);
+ }
+
@Test
void
textDeltasShouldBeCoalescedIntoBlocksWhileTheStreamStaysPerTokenTest() {
// 5000 characters in 500 deltas: 2048-character cuts leave two rows,
and the finalisation flush
