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 3bc0f7473 fix(ai): cancel abandoned streams and reject overload (#2209)
3bc0f7473 is described below
commit 3bc0f74739ad11de005ff03991ec2d40c59c850e
Author: xdz997 <[email protected]>
AuthorDate: Wed Aug 19 14:24:09 2026 +0800
fix(ai): cancel abandoned streams and reject overload (#2209)
---
.../studio/ops/ai/ClaudeCodeAgentProvider.java | 10 +-
.../rocketmq/studio/ops/ai/LlmSseSession.java | 129 ++++++++++++
.../studio/ops/ai/OpenAiCompatibleLlmGateway.java | 171 ++++++++++++----
.../studio/ops/ai/ClaudeCodeAgentProviderTest.java | 70 ++++++-
.../rocketmq/studio/ops/ai/LlmSseSessionTest.java | 130 ++++++++++++
.../ops/ai/OpenAiCompatibleLlmGatewayTest.java | 227 +++++++++++++++++++++
6 files changed, 691 insertions(+), 46 deletions(-)
diff --git
a/server/src/main/java/org/apache/rocketmq/studio/ops/ai/ClaudeCodeAgentProvider.java
b/server/src/main/java/org/apache/rocketmq/studio/ops/ai/ClaudeCodeAgentProvider.java
index 033eb6ff3..8646ca74d 100644
---
a/server/src/main/java/org/apache/rocketmq/studio/ops/ai/ClaudeCodeAgentProvider.java
+++
b/server/src/main/java/org/apache/rocketmq/studio/ops/ai/ClaudeCodeAgentProvider.java
@@ -98,8 +98,9 @@ public class ClaudeCodeAgentProvider extends CliAgentProvider
{
ProcessBuilder builder = new ProcessBuilder(command);
processEnvironment().apply(builder, childEnv(config));
builder.redirectErrorStream(false);
+ Process process = null;
try {
- Process process = builder.start();
+ process = startProcess(builder);
AtomicBoolean emitted = new AtomicBoolean(false);
StringBuilder resultText = new StringBuilder();
CompletableFuture<Void> stdoutFuture = drainStdout(
@@ -128,12 +129,19 @@ public class ClaudeCodeAgentProvider extends
CliAgentProvider {
"Failed to execute " + binaryName() + " CLI",
"Check that the CLI binary is installed and executable.",
exception);
} catch (InterruptedException exception) {
+ if (process != null) {
+ process.destroyForcibly();
+ }
Thread.currentThread().interrupt();
throw new LlmGatewayException(502, "llm.provider.interrupted",
binaryName() + " CLI execution was interrupted", "Retry
the request.", exception);
}
}
+ protected Process startProcess(ProcessBuilder builder) throws IOException {
+ return builder.start();
+ }
+
protected long streamTimeoutSeconds() {
return STREAM_TIMEOUT_SECONDS;
}
diff --git
a/server/src/main/java/org/apache/rocketmq/studio/ops/ai/LlmSseSession.java
b/server/src/main/java/org/apache/rocketmq/studio/ops/ai/LlmSseSession.java
new file mode 100644
index 000000000..0835c4024
--- /dev/null
+++ b/server/src/main/java/org/apache/rocketmq/studio/ops/ai/LlmSseSession.java
@@ -0,0 +1,129 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.rocketmq.studio.ops.ai;
+
+import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
+
+import java.io.IOException;
+import java.util.Objects;
+import java.util.concurrent.Future;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Consumer;
+
+/**
+ * Couples one downstream SSE response with its asynchronous provider task.
+ * Client completion, timeout, and transport errors cancel the task so an
+ * abandoned response does not retain a gateway worker until provider timeout.
+ */
+final class LlmSseSession {
+
+ private enum State {
+ ACTIVE,
+ TERMINATING,
+ TERMINATED,
+ CANCELLED
+ }
+
+ private final SseEmitter emitter;
+ private final Consumer<LlmSseSession> terminationListener;
+ private final AtomicReference<State> state = new
AtomicReference<>(State.ACTIVE);
+ private final AtomicReference<Future<?>> task = new AtomicReference<>();
+ private final AtomicBoolean terminationNotified = new AtomicBoolean();
+
+ LlmSseSession(SseEmitter emitter, Consumer<LlmSseSession>
terminationListener) {
+ this.emitter = Objects.requireNonNull(emitter, "emitter");
+ this.terminationListener = Objects.requireNonNull(terminationListener,
"terminationListener");
+ emitter.onCompletion(this::cancel);
+ emitter.onTimeout(this::cancel);
+ emitter.onError(ignored -> cancel());
+ }
+
+ SseEmitter emitter() {
+ return emitter;
+ }
+
+ void attach(Future<?> submittedTask) {
+ Objects.requireNonNull(submittedTask, "submittedTask");
+ if (!task.compareAndSet(null, submittedTask)) {
+ submittedTask.cancel(true);
+ throw new IllegalStateException("An SSE session can own only one
task");
+ }
+ if (state.get() == State.CANCELLED) {
+ submittedTask.cancel(true);
+ }
+ }
+
+ boolean beginTerminal() {
+ return state.compareAndSet(State.ACTIVE, State.TERMINATING);
+ }
+
+ void send(SseEmitter.SseEventBuilder event) throws IOException {
+ State current = state.get();
+ if (current == State.CANCELLED || current == State.TERMINATED) {
+ throw new IOException("SSE client is no longer connected");
+ }
+ try {
+ emitter.send(event);
+ } catch (IOException exception) {
+ cancel();
+ throw exception;
+ }
+ }
+
+ void complete() {
+ if (state.compareAndSet(State.TERMINATING, State.TERMINATED)) {
+ notifyTermination();
+ emitter.complete();
+ }
+ }
+
+ void completeWithError(Throwable throwable) {
+ State previous = state.getAndUpdate(current -> switch (current) {
+ case ACTIVE, TERMINATING -> State.TERMINATED;
+ case TERMINATED, CANCELLED -> current;
+ });
+ if (previous == State.ACTIVE || previous == State.TERMINATING) {
+ notifyTermination();
+ emitter.completeWithError(throwable);
+ }
+ }
+
+ void cancel() {
+ State previous = state.getAndUpdate(current -> switch (current) {
+ case ACTIVE, TERMINATING -> State.CANCELLED;
+ case TERMINATED, CANCELLED -> current;
+ });
+ if (previous == State.ACTIVE || previous == State.TERMINATING) {
+ Future<?> submittedTask = task.get();
+ if (submittedTask != null) {
+ submittedTask.cancel(true);
+ }
+ notifyTermination();
+ }
+ }
+
+ boolean isCancelled() {
+ return state.get() == State.CANCELLED;
+ }
+
+ private void notifyTermination() {
+ if (terminationNotified.compareAndSet(false, true)) {
+ terminationListener.accept(this);
+ }
+ }
+}
diff --git
a/server/src/main/java/org/apache/rocketmq/studio/ops/ai/OpenAiCompatibleLlmGateway.java
b/server/src/main/java/org/apache/rocketmq/studio/ops/ai/OpenAiCompatibleLlmGateway.java
index 8ec0f825d..efdb6d1d4 100644
---
a/server/src/main/java/org/apache/rocketmq/studio/ops/ai/OpenAiCompatibleLlmGateway.java
+++
b/server/src/main/java/org/apache/rocketmq/studio/ops/ai/OpenAiCompatibleLlmGateway.java
@@ -19,8 +19,8 @@ package org.apache.rocketmq.studio.ops.ai;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import jakarta.annotation.PreDestroy;
-import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
+import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Primary;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;
@@ -31,27 +31,62 @@ import java.io.InputStream;
import java.nio.charset.StandardCharsets;
import java.util.Locale;
import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Future;
+import java.util.concurrent.RejectedExecutionException;
import java.util.concurrent.SynchronousQueue;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
+import java.util.function.Consumer;
+import java.util.function.LongFunction;
@Slf4j
@Primary
@Component
-@RequiredArgsConstructor
public class OpenAiCompatibleLlmGateway implements LlmGateway {
+ private static final long HTTP_STREAM_TIMEOUT_MILLIS = 60_000L;
+ private static final long CLI_STREAM_TIMEOUT_MILLIS = 300_000L;
+ private static final int MAX_CONCURRENT_CHATS = 16;
+
private final LlmConfigService configService;
private final OpenAiCompatibleLlmClient llmClient;
private final AgentProviderRegistry agentProviders;
private final ObjectMapper objectMapper;
- // Bounded pool: cached threads grow without limit under load and,
combined with a hung CLI
- // child, can exhaust memory. CallerRunsPolicy keeps SSE work from being
dropped under load.
- private final ExecutorService executor = new ThreadPoolExecutor(
- 0, 16, 60L, TimeUnit.SECONDS,
- new SynchronousQueue<>(),
- new ThreadPoolExecutor.CallerRunsPolicy());
+ private final ExecutorService executor;
+ private final LongFunction<SseEmitter> emitterFactory;
+ private final Set<LlmSseSession> activeSessions =
ConcurrentHashMap.newKeySet();
+
+ @Autowired
+ public OpenAiCompatibleLlmGateway(LlmConfigService configService,
+ OpenAiCompatibleLlmClient llmClient,
+ AgentProviderRegistry agentProviders,
+ ObjectMapper objectMapper) {
+ this(configService, llmClient, agentProviders, objectMapper,
newChatExecutor(), SseEmitter::new);
+ }
+
+ OpenAiCompatibleLlmGateway(LlmConfigService configService,
+ OpenAiCompatibleLlmClient llmClient,
+ AgentProviderRegistry agentProviders,
+ ObjectMapper objectMapper,
+ ExecutorService executor,
+ LongFunction<SseEmitter> emitterFactory) {
+ this.configService = configService;
+ this.llmClient = llmClient;
+ this.agentProviders = agentProviders;
+ this.objectMapper = objectMapper;
+ this.executor = executor;
+ this.emitterFactory = emitterFactory;
+ }
+
+ private static ExecutorService newChatExecutor() {
+ return new ThreadPoolExecutor(
+ 0, MAX_CONCURRENT_CHATS, 60L, TimeUnit.SECONDS,
+ new SynchronousQueue<>(),
+ new ThreadPoolExecutor.AbortPolicy());
+ }
@Override
public SseEmitter chat(ChatDTO request) {
@@ -61,17 +96,15 @@ public class OpenAiCompatibleLlmGateway implements
LlmGateway {
}
String engine = resolveEngine(request == null ? null :
request.getEngine(), config);
if (isCliEngine(engine)) {
- SseEmitter emitter = new SseEmitter(300_000L);
- executor.execute(() -> runCliChat(request, config, engine,
emitter));
- return emitter;
+ return submitChat(CLI_STREAM_TIMEOUT_MILLIS,
+ session -> runCliChat(request, config, engine, session));
}
if (!llmClient.supports(config)) {
return errorEmitter(unsupportedProviderException());
}
- SseEmitter emitter = new SseEmitter(60_000L);
- executor.execute(() -> streamChat(request, config, emitter));
- return emitter;
+ return submitChat(HTTP_STREAM_TIMEOUT_MILLIS,
+ session -> streamChat(request, config, session));
}
@Override
@@ -106,23 +139,28 @@ public class OpenAiCompatibleLlmGateway implements
LlmGateway {
return !LlmConfigVO.ENGINE_HTTP.equalsIgnoreCase(engine);
}
- private void runCliChat(ChatDTO request, LlmConfigVO config, String
engine, SseEmitter emitter) {
+ private void runCliChat(ChatDTO request, LlmConfigVO config, String
engine, LlmSseSession session) {
try {
AgentProvider provider = agentProviders.forEngine(engine);
String prompt = request == null ? null : request.getMessage();
if (request != null && request.isEnhance() &&
StringUtils.hasText(prompt)) {
- prompt = enhanceAndEmit(config, provider, prompt, emitter);
+ prompt = enhanceAndEmit(config, provider, prompt, session);
}
String result = provider.complete(config, prompt, request == null
? null : request.getModel());
- sendMessage(emitter, result);
- emitter.send(SseEmitter.event().name("done").data("[DONE]"));
- emitter.complete();
+ sendMessage(session, result);
+ finishSuccess(session);
} catch (LlmGatewayException exception) {
+ if (session.isCancelled()) {
+ return;
+ }
log.warn("Agent CLI chat failed: {}", exception.getCode(),
exception);
- sendError(emitter, exception);
+ sendError(session, exception);
} catch (Exception exception) {
+ if (session.isCancelled()) {
+ return;
+ }
log.error("Failed to run agent CLI chat", exception);
- sendError(emitter, new LlmGatewayException(502,
"llm.gateway_error",
+ sendError(session, new LlmGatewayException(502,
"llm.gateway_error",
"Failed to run agent CLI chat", "Check the agent provider
configuration and retry.", exception));
}
}
@@ -142,36 +180,36 @@ public class OpenAiCompatibleLlmGateway implements
LlmGateway {
}
/** Streams the prompt rewrite via the given provider, emitting per-chunk
"enhance" SSE events. */
- private String enhanceAndEmit(LlmConfigVO config, AgentProvider provider,
String rawPrompt, SseEmitter emitter)
+ private String enhanceAndEmit(LlmConfigVO config, AgentProvider provider,
String rawPrompt, LlmSseSession session)
throws IOException {
String metaPrompt = ENHANCE_PROMPT_TEMPLATE.formatted(rawPrompt);
StringBuilder accumulated = new StringBuilder();
provider.stream(config, metaPrompt, null, chunk -> {
accumulated.append(chunk);
- emitEnhanceChunk(emitter, chunk);
+ emitEnhanceChunk(session, chunk);
});
String enhanced = cleanEnhancedPrompt(accumulated.toString());
return StringUtils.hasText(enhanced) ? enhanced : rawPrompt;
}
- private String enhanceAndEmitHttp(LlmConfigVO config, String rawPrompt,
SseEmitter emitter)
+ private String enhanceAndEmitHttp(LlmConfigVO config, String rawPrompt,
LlmSseSession session)
throws IOException {
String metaPrompt = ENHANCE_PROMPT_TEMPLATE.formatted(rawPrompt);
StringBuilder accumulated = new StringBuilder();
llmClient.stream(config, metaPrompt, null, chunk -> {
accumulated.append(chunk);
- emitEnhanceChunk(emitter, chunk);
+ emitEnhanceChunk(session, chunk);
});
String enhanced = cleanEnhancedPrompt(accumulated.toString());
return StringUtils.hasText(enhanced) ? enhanced : rawPrompt;
}
- private void emitEnhanceChunk(SseEmitter emitter, String chunk) {
+ private void emitEnhanceChunk(LlmSseSession session, String chunk) {
if (!StringUtils.hasText(chunk)) {
return;
}
try {
- emitter.send(SseEmitter.event()
+ session.send(SseEmitter.event()
.name("enhance")
.data(objectMapper.writeValueAsString(Map.of("delta",
chunk))));
} catch (IOException exception) {
@@ -193,33 +231,39 @@ public class OpenAiCompatibleLlmGateway implements
LlmGateway {
@PreDestroy
void destroy() {
+ activeSessions.forEach(LlmSseSession::cancel);
executor.shutdownNow();
}
- private void streamChat(ChatDTO request, LlmConfigVO config, SseEmitter
emitter) {
+ private void streamChat(ChatDTO request, LlmConfigVO config, LlmSseSession
session) {
try {
String prompt = request == null ? null : request.getMessage();
if (request != null && request.isEnhance() &&
StringUtils.hasText(prompt)) {
- prompt = enhanceAndEmitHttp(config, prompt, emitter);
+ prompt = enhanceAndEmitHttp(config, prompt, session);
}
llmClient.stream(config, prompt,
request == null ? null : request.getModel(),
- token -> sendMessage(emitter, token));
- emitter.send(SseEmitter.event().name("done").data("[DONE]"));
- emitter.complete();
+ token -> sendMessage(session, token));
+ finishSuccess(session);
} catch (LlmGatewayException exception) {
+ if (session.isCancelled()) {
+ return;
+ }
log.warn("LLM chat stream failed: {}", exception.getCode(),
exception);
- sendError(emitter, exception);
+ sendError(session, exception);
} catch (Exception exception) {
+ if (session.isCancelled()) {
+ return;
+ }
log.error("Failed to stream LLM chat response", exception);
- sendError(emitter, new LlmGatewayException(502,
"llm.gateway_error",
+ sendError(session, new LlmGatewayException(502,
"llm.gateway_error",
"Failed to stream LLM chat response", "Check the LLM
provider configuration and retry.", exception));
}
}
- private void sendMessage(SseEmitter emitter, String token) {
+ private void sendMessage(LlmSseSession session, String token) {
try {
- emitter.send(SseEmitter.event()
+ session.send(SseEmitter.event()
.name("message")
.data(objectMapper.writeValueAsString(Map.of("text",
token))));
} catch (IOException exception) {
@@ -228,25 +272,58 @@ public class OpenAiCompatibleLlmGateway implements
LlmGateway {
}
}
+ private SseEmitter submitChat(long timeoutMillis, Consumer<LlmSseSession>
work) {
+ LlmSseSession session = newSession(timeoutMillis);
+ try {
+ Future<?> task = executor.submit(() -> work.accept(session));
+ session.attach(task);
+ } catch (RejectedExecutionException exception) {
+ sendError(session, overloadedException());
+ }
+ return session.emitter();
+ }
+
+ private LlmSseSession newSession(long timeoutMillis) {
+ LlmSseSession session = new LlmSseSession(
+ emitterFactory.apply(timeoutMillis), activeSessions::remove);
+ activeSessions.add(session);
+ return session;
+ }
+
+ private void finishSuccess(LlmSseSession session) {
+ if (!session.beginTerminal()) {
+ return;
+ }
+ try {
+ session.send(SseEmitter.event().name("done").data("[DONE]"));
+ session.complete();
+ } catch (IOException exception) {
+ session.completeWithError(exception);
+ }
+ }
+
private SseEmitter errorEmitter(LlmGatewayException exception) {
- SseEmitter emitter = new SseEmitter(60_000L);
- executor.execute(() -> sendError(emitter, exception));
- return emitter;
+ LlmSseSession session = newSession(HTTP_STREAM_TIMEOUT_MILLIS);
+ sendError(session, exception);
+ return session.emitter();
}
- private void sendError(SseEmitter emitter, LlmGatewayException exception) {
+ private void sendError(LlmSseSession session, LlmGatewayException
exception) {
+ if (!session.beginTerminal()) {
+ return;
+ }
try {
- emitter.send(SseEmitter.event()
+ session.send(SseEmitter.event()
.name("error")
.data(objectMapper.writeValueAsString(Map.of(
"status", exception.getStatusCode(),
"code", exception.getCode(),
"message", exception.getMessage(),
"hint", exception.getHint() == null ? "" :
exception.getHint()))));
- emitter.send(SseEmitter.event().name("done").data("[DONE]"));
- emitter.complete();
+ session.send(SseEmitter.event().name("done").data("[DONE]"));
+ session.complete();
} catch (IOException ioException) {
- emitter.completeWithError(ioException);
+ session.completeWithError(ioException);
}
}
@@ -272,6 +349,12 @@ public class OpenAiCompatibleLlmGateway implements
LlmGateway {
"Use one of: openai, deepseek, tongyi, ollama.");
}
+ private LlmGatewayException overloadedException() {
+ return new LlmGatewayException(503, "llm.gateway.overloaded",
+ "AI chat capacity is temporarily exhausted",
+ "Wait for an active chat to finish, then retry.");
+ }
+
private String commandPrompt(AiCommandDTO command) {
if (command == null) {
return "";
diff --git
a/server/src/test/java/org/apache/rocketmq/studio/ops/ai/ClaudeCodeAgentProviderTest.java
b/server/src/test/java/org/apache/rocketmq/studio/ops/ai/ClaudeCodeAgentProviderTest.java
index d9767c31c..a0bfc00cb 100644
---
a/server/src/test/java/org/apache/rocketmq/studio/ops/ai/ClaudeCodeAgentProviderTest.java
+++
b/server/src/test/java/org/apache/rocketmq/studio/ops/ai/ClaudeCodeAgentProviderTest.java
@@ -18,12 +18,23 @@ package org.apache.rocketmq.studio.ops.ai;
import org.junit.jupiter.api.Test;
+import java.io.ByteArrayInputStream;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.mockito.ArgumentMatchers.anyLong;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
class ClaudeCodeAgentProviderTest {
@@ -52,6 +63,7 @@ class ClaudeCodeAgentProviderTest {
}
@Test
+
void streamUsesTheIsolatedEnvironment() {
RecordingEnvironment processEnvironment = new RecordingEnvironment();
TestClaudeCodeAgentProvider provider = new TestClaudeCodeAgentProvider(
@@ -69,6 +81,44 @@ class ClaudeCodeAgentProviderTest {
.doesNotContainKey("SERVER_SECRET"));
}
+ @Test
+ void streamInterruptionDestroysTheChildProcess() throws Exception {
+ Process process = mock(Process.class);
+ CountDownLatch waitStarted = new CountDownLatch(1);
+ when(process.getInputStream()).thenReturn(new ByteArrayInputStream(new
byte[0]));
+ when(process.getErrorStream()).thenReturn(new ByteArrayInputStream(new
byte[0]));
+ doAnswer(invocation -> {
+ waitStarted.countDown();
+ new CountDownLatch(1).await();
+ return false;
+ }).when(process).waitFor(anyLong(), eq(TimeUnit.SECONDS));
+ TestClaudeCodeAgentProvider provider = new TestClaudeCodeAgentProvider(
+ List.of("claude"), 300, process);
+ AtomicReference<LlmGatewayException> failure = new AtomicReference<>();
+ AtomicBoolean interruptPreserved = new AtomicBoolean();
+ CountDownLatch finished = new CountDownLatch(1);
+ Thread worker = new Thread(() -> {
+ try {
+ provider.stream(LlmConfigVO.builder().build(), "prompt", null,
ignored -> { });
+ } catch (LlmGatewayException exception) {
+ failure.set(exception);
+ interruptPreserved.set(Thread.currentThread().isInterrupted());
+ } finally {
+ finished.countDown();
+ }
+ });
+ worker.start();
+ assertThat(waitStarted.await(5, TimeUnit.SECONDS)).isTrue();
+
+ worker.interrupt();
+
+ assertThat(finished.await(5, TimeUnit.SECONDS)).isTrue();
+ assertThat(failure.get()).isNotNull();
+
assertThat(failure.get().getCode()).isEqualTo("llm.provider.interrupted");
+ assertThat(interruptPreserved).isTrue();
+ verify(process).destroyForcibly();
+ }
+
private static final class RecordingEnvironment extends
CliProcessEnvironment {
private final List<Map<String, String>> childEnvironments = new
ArrayList<>();
@@ -84,23 +134,36 @@ class ClaudeCodeAgentProviderTest {
}
}
+
private static class TestClaudeCodeAgentProvider extends
ClaudeCodeAgentProvider {
private final List<String> command;
private final long timeoutSeconds;
private final Map<String, String> environment;
+ private final Process process;
TestClaudeCodeAgentProvider(List<String> command, long timeoutSeconds)
{
- this(command, timeoutSeconds, new
CliProcessEnvironment(List.of()), Map.of());
+ this(command, timeoutSeconds, new
CliProcessEnvironment(List.of()), Map.of(), null);
}
TestClaudeCodeAgentProvider(List<String> command, long timeoutSeconds,
CliProcessEnvironment processEnvironment,
Map<String, String> environment) {
+ this(command, timeoutSeconds, processEnvironment, environment,
null);
+ }
+
+ TestClaudeCodeAgentProvider(List<String> command, long timeoutSeconds,
Process process) {
+ this(command, timeoutSeconds, new
CliProcessEnvironment(List.of()), Map.of(), process);
+ }
+
+ TestClaudeCodeAgentProvider(List<String> command, long timeoutSeconds,
+ CliProcessEnvironment processEnvironment,
+ Map<String, String> environment, Process
process) {
super(null, processEnvironment);
this.command = command;
this.timeoutSeconds = timeoutSeconds;
this.environment = environment;
+ this.process = process;
}
@Override
@@ -127,5 +190,10 @@ class ClaudeCodeAgentProviderTest {
protected long streamTimeoutSeconds() {
return timeoutSeconds;
}
+
+ @Override
+ protected Process startProcess(ProcessBuilder builder) throws
java.io.IOException {
+ return process == null ? super.startProcess(builder) : process;
+ }
}
}
diff --git
a/server/src/test/java/org/apache/rocketmq/studio/ops/ai/LlmSseSessionTest.java
b/server/src/test/java/org/apache/rocketmq/studio/ops/ai/LlmSseSessionTest.java
new file mode 100644
index 000000000..625d24b6b
--- /dev/null
+++
b/server/src/test/java/org/apache/rocketmq/studio/ops/ai/LlmSseSessionTest.java
@@ -0,0 +1,130 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.rocketmq.studio.ops.ai;
+
+import org.junit.jupiter.api.Test;
+import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
+
+import java.util.concurrent.Future;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Consumer;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
+
+class LlmSseSessionTest {
+
+ @Test
+ void completionTimeoutAndErrorCancelAttachedTasks() {
+ assertLifecycleCallbackCancels(TestEmitter::triggerCompletion);
+ assertLifecycleCallbackCancels(TestEmitter::triggerTimeout);
+ assertLifecycleCallbackCancels(emitter -> emitter.triggerError(new
IllegalStateException("closed")));
+ }
+
+ @Test
+ void terminationBeforeAttachmentCancelsTheLateTask() {
+ TestEmitter emitter = new TestEmitter();
+ LlmSseSession session = new LlmSseSession(emitter, ignored -> { });
+ Future<?> task = mock(Future.class);
+
+ emitter.triggerTimeout();
+ session.attach(task);
+
+ assertThat(session.isCancelled()).isTrue();
+ verify(task).cancel(true);
+ }
+
+ @Test
+ void normalTerminalCompletionDoesNotCancelTheTask() {
+ TestEmitter emitter = new TestEmitter();
+ AtomicInteger terminations = new AtomicInteger();
+ LlmSseSession session = new LlmSseSession(emitter, ignored ->
terminations.incrementAndGet());
+ Future<?> task = mock(Future.class);
+ session.attach(task);
+
+ assertThat(session.beginTerminal()).isTrue();
+ session.complete();
+
+ assertThat(emitter.completed).isTrue();
+ assertThat(terminations).hasValue(1);
+ assertThat(session.isCancelled()).isFalse();
+ verify(task, never()).cancel(true);
+ }
+
+ private void assertLifecycleCallbackCancels(Consumer<TestEmitter>
callback) {
+ TestEmitter emitter = new TestEmitter();
+ AtomicInteger terminations = new AtomicInteger();
+ LlmSseSession session = new LlmSseSession(emitter, ignored ->
terminations.incrementAndGet());
+ Future<?> task = mock(Future.class);
+ session.attach(task);
+
+ callback.accept(emitter);
+ callback.accept(emitter);
+
+ assertThat(session.isCancelled()).isTrue();
+ assertThat(terminations).hasValue(1);
+ verify(task).cancel(true);
+ }
+
+ private static final class TestEmitter extends SseEmitter {
+ private Runnable completionCallback;
+ private Runnable timeoutCallback;
+ private Consumer<Throwable> errorCallback;
+ private boolean completed;
+
+ @Override
+ public synchronized void onCompletion(Runnable callback) {
+ this.completionCallback = callback;
+ }
+
+ @Override
+ public synchronized void onTimeout(Runnable callback) {
+ this.timeoutCallback = callback;
+ }
+
+ @Override
+ public synchronized void onError(Consumer<Throwable> callback) {
+ this.errorCallback = callback;
+ }
+
+ @Override
+ public synchronized void complete() {
+ completed = true;
+ triggerCompletion();
+ }
+
+ void triggerCompletion() {
+ if (completionCallback != null) {
+ completionCallback.run();
+ }
+ }
+
+ void triggerTimeout() {
+ if (timeoutCallback != null) {
+ timeoutCallback.run();
+ }
+ }
+
+ void triggerError(Throwable throwable) {
+ if (errorCallback != null) {
+ errorCallback.accept(throwable);
+ }
+ }
+ }
+}
diff --git
a/server/src/test/java/org/apache/rocketmq/studio/ops/ai/OpenAiCompatibleLlmGatewayTest.java
b/server/src/test/java/org/apache/rocketmq/studio/ops/ai/OpenAiCompatibleLlmGatewayTest.java
index afff18ff1..23074ab12 100644
---
a/server/src/test/java/org/apache/rocketmq/studio/ops/ai/OpenAiCompatibleLlmGatewayTest.java
+++
b/server/src/test/java/org/apache/rocketmq/studio/ops/ai/OpenAiCompatibleLlmGatewayTest.java
@@ -19,14 +19,28 @@ package org.apache.rocketmq.studio.ops.ai;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.junit.jupiter.api.Test;
import org.mockito.ArgumentCaptor;
+import
org.springframework.web.servlet.mvc.method.annotation.ResponseBodyEmitter;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CopyOnWriteArrayList;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.SynchronousQueue;
+import java.util.concurrent.ThreadPoolExecutor;
+import java.util.concurrent.TimeUnit;
+import java.util.function.Consumer;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
import static org.assertj.core.api.Assertions.assertThat;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.verify;
@@ -105,6 +119,153 @@ class OpenAiCompatibleLlmGatewayTest {
});
}
+ @Test
+ void
saturatedGatewayReturnsStructuredOverloadWithoutRunningProviderOnCaller()
throws Exception {
+ ExecutorService executor = singleChatExecutor();
+ List<RecordingSseEmitter> emitters = new CopyOnWriteArrayList<>();
+ OpenAiCompatibleLlmGateway testedGateway = gateway(executor, emitters);
+ LlmConfigVO config = config("openai", "sk-test");
+ CountDownLatch started = new CountDownLatch(1);
+ CountDownLatch release = new CountDownLatch(1);
+ when(configService.getConfig()).thenReturn(config);
+ when(llmClient.supports(config)).thenReturn(true);
+ doAnswer(invocation -> {
+ started.countDown();
+ release.await();
+ return null;
+ }).when(llmClient).stream(any(), any(), any(), any());
+ try {
+ testedGateway.chat(ChatDTO.builder().message("first").build());
+ assertThat(started.await(5, TimeUnit.SECONDS)).isTrue();
+
+ CompletableFuture<SseEmitter> overloaded =
CompletableFuture.supplyAsync(
+ () ->
testedGateway.chat(ChatDTO.builder().message("second").build()));
+ SseEmitter result = overloaded.get(1, TimeUnit.SECONDS);
+
+ assertThat(result).isSameAs(emitters.get(1));
+ assertThat(emitters.get(1).eventText())
+ .contains("event:error", "llm.gateway.overloaded", "503",
"event:done", "[DONE]");
+ assertThat(emitters.get(1).completed).isTrue();
+ } finally {
+ release.countDown();
+ testedGateway.destroy();
+ }
+ }
+
+ @Test
+ void
downstreamTimeoutInterruptsTheRunningProviderWithoutSendingAnotherTerminalEvent()
throws Exception {
+ ExecutorService executor = singleChatExecutor();
+ List<RecordingSseEmitter> emitters = new CopyOnWriteArrayList<>();
+ OpenAiCompatibleLlmGateway testedGateway = gateway(executor, emitters);
+ LlmConfigVO config = config("openai", "sk-test");
+ CountDownLatch started = new CountDownLatch(1);
+ CountDownLatch interrupted = new CountDownLatch(1);
+ when(configService.getConfig()).thenReturn(config);
+ when(llmClient.supports(config)).thenReturn(true);
+ doAnswer(invocation -> {
+ started.countDown();
+ try {
+ new CountDownLatch(1).await();
+ } catch (InterruptedException exception) {
+ interrupted.countDown();
+ Thread.currentThread().interrupt();
+ }
+ return null;
+ }).when(llmClient).stream(any(), any(), any(), any());
+ try {
+ testedGateway.chat(ChatDTO.builder().message("hello").build());
+ assertThat(started.await(5, TimeUnit.SECONDS)).isTrue();
+
+ emitters.get(0).triggerTimeout();
+
+ assertThat(interrupted.await(5, TimeUnit.SECONDS)).isTrue();
+ assertThat(emitters.get(0).sentEvents).isEmpty();
+ } finally {
+ testedGateway.destroy();
+ }
+ }
+
+ @Test
+ void gatewayShutdownInterruptsActiveProviderWork() throws Exception {
+ ExecutorService executor = singleChatExecutor();
+ List<RecordingSseEmitter> emitters = new CopyOnWriteArrayList<>();
+ OpenAiCompatibleLlmGateway testedGateway = gateway(executor, emitters);
+ LlmConfigVO config = config("openai", "sk-test");
+ CountDownLatch started = new CountDownLatch(1);
+ CountDownLatch interrupted = new CountDownLatch(1);
+ when(configService.getConfig()).thenReturn(config);
+ when(llmClient.supports(config)).thenReturn(true);
+ doAnswer(invocation -> {
+ started.countDown();
+ try {
+ new CountDownLatch(1).await();
+ } catch (InterruptedException exception) {
+ interrupted.countDown();
+ Thread.currentThread().interrupt();
+ }
+ return null;
+ }).when(llmClient).stream(any(), any(), any(), any());
+
+ testedGateway.chat(ChatDTO.builder().message("hello").build());
+ assertThat(started.await(5, TimeUnit.SECONDS)).isTrue();
+
+ testedGateway.destroy();
+
+ assertThat(interrupted.await(5, TimeUnit.SECONDS)).isTrue();
+ assertThat(emitters.get(0).sentEvents).isEmpty();
+ }
+
+ @Test
+ void successfulAndFailedStreamsEmitOneTerminalSequence() throws Exception {
+ ExecutorService executor = singleChatExecutor();
+ List<RecordingSseEmitter> emitters = new CopyOnWriteArrayList<>();
+ OpenAiCompatibleLlmGateway testedGateway = gateway(executor, emitters);
+ LlmConfigVO config = config("openai", "sk-test");
+ when(configService.getConfig()).thenReturn(config);
+ when(llmClient.supports(config)).thenReturn(true);
+ doAnswer(invocation -> {
+ @SuppressWarnings("unchecked")
+ Consumer<String> consumer = invocation.getArgument(3,
Consumer.class);
+ consumer.accept("hello");
+ return null;
+ }).doThrow(new LlmGatewayException(502, "llm.provider.failed",
"provider failed", "retry"))
+ .when(llmClient).stream(any(), any(), any(), any());
+ try {
+ testedGateway.chat(ChatDTO.builder().message("success").build());
+ assertThat(emitters.get(0).completedLatch.await(5,
TimeUnit.SECONDS)).isTrue();
+ testedGateway.chat(ChatDTO.builder().message("failure").build());
+ assertThat(emitters.get(1).completedLatch.await(5,
TimeUnit.SECONDS)).isTrue();
+
+ assertThat(emitters.get(0).eventText())
+ .contains("event:message", "hello", "event:done", "[DONE]")
+ .doesNotContain("event:error");
+ assertThat(emitters.get(1).eventText())
+ .contains("event:error", "llm.provider.failed",
"event:done", "[DONE]");
+ assertThat(emitters.get(0).eventCount("event:done")).isEqualTo(1);
+ assertThat(emitters.get(1).eventCount("event:done")).isEqualTo(1);
+ } finally {
+ testedGateway.destroy();
+ }
+ }
+
+ private OpenAiCompatibleLlmGateway gateway(ExecutorService executor,
+ List<RecordingSseEmitter>
emitters) {
+ return new OpenAiCompatibleLlmGateway(
+ configService, llmClient, new
AgentProviderRegistry(List.of()), new ObjectMapper(),
+ executor, timeout -> {
+ RecordingSseEmitter emitter = new
RecordingSseEmitter(timeout);
+ emitters.add(emitter);
+ return emitter;
+ });
+ }
+
+ private ExecutorService singleChatExecutor() {
+ return new ThreadPoolExecutor(
+ 0, 1, 60L, TimeUnit.SECONDS,
+ new SynchronousQueue<>(),
+ new ThreadPoolExecutor.AbortPolicy());
+ }
+
private LlmConfigVO config(String provider, String apiKey) {
return LlmConfigVO.builder()
.provider(provider)
@@ -116,4 +277,70 @@ class OpenAiCompatibleLlmGatewayTest {
.enabled(true)
.build();
}
+
+ private static final class RecordingSseEmitter extends SseEmitter {
+ private final List<Set<ResponseBodyEmitter.DataWithMediaType>>
sentEvents = new CopyOnWriteArrayList<>();
+ private final CountDownLatch completedLatch = new CountDownLatch(1);
+ private Runnable completionCallback;
+ private Runnable timeoutCallback;
+ private Consumer<Throwable> errorCallback;
+ private boolean completed;
+
+ RecordingSseEmitter(long timeout) {
+ super(timeout);
+ }
+
+ @Override
+ public void send(SseEventBuilder builder) throws IOException {
+ sentEvents.add(builder.build());
+ }
+
+ @Override
+ public synchronized void onCompletion(Runnable callback) {
+ completionCallback = callback;
+ }
+
+ @Override
+ public synchronized void onTimeout(Runnable callback) {
+ timeoutCallback = callback;
+ }
+
+ @Override
+ public synchronized void onError(Consumer<Throwable> callback) {
+ errorCallback = callback;
+ }
+
+ @Override
+ public synchronized void complete() {
+ completed = true;
+ completedLatch.countDown();
+ if (completionCallback != null) {
+ completionCallback.run();
+ }
+ }
+
+ @Override
+ public synchronized void completeWithError(Throwable throwable) {
+ completedLatch.countDown();
+ if (errorCallback != null) {
+ errorCallback.accept(throwable);
+ }
+ }
+
+ void triggerTimeout() {
+ timeoutCallback.run();
+ }
+
+ String eventText() {
+ List<String> values = new ArrayList<>();
+ sentEvents.forEach(event -> event.forEach(item ->
values.add(String.valueOf(item.getData()))));
+ return String.join("", values);
+ }
+
+ long eventCount(String marker) {
+ return sentEvents.stream()
+ .filter(event -> event.stream().anyMatch(item ->
String.valueOf(item.getData()).contains(marker)))
+ .count();
+ }
+ }
}
