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 7ef2ec7e1 fix(ai): bound request and stream payloads (#2539)
7ef2ec7e1 is described below
commit 7ef2ec7e14b406a630faf43a3f89f3f0c49ab8b5
Author: xdz997 <[email protected]>
AuthorDate: Tue Aug 25 17:28:06 2026 +0800
fix(ai): bound request and stream payloads (#2539)
---
.../rocketmq/studio/ops/ai/AiPayloadGuard.java | 157 +++++++++++++++++++++
.../apache/rocketmq/studio/ops/ai/AiService.java | 16 ++-
.../rocketmq/studio/ops/ai/CliAgentProvider.java | 4 +
.../studio/ops/ai/OpenAiCompatibleLlmClient.java | 65 ++++++++-
.../rocketmq/studio/ops/ai/AiPayloadGuardTest.java | 77 ++++++++++
.../rocketmq/studio/ops/ai/AiServiceTest.java | 46 +++++-
.../studio/ops/ai/CliAgentProviderTest.java | 12 ++
.../ops/ai/OpenAiCompatibleLlmClientTest.java | 39 +++++
8 files changed, 409 insertions(+), 7 deletions(-)
diff --git
a/server/src/main/java/org/apache/rocketmq/studio/ops/ai/AiPayloadGuard.java
b/server/src/main/java/org/apache/rocketmq/studio/ops/ai/AiPayloadGuard.java
new file mode 100644
index 000000000..b8bc68e07
--- /dev/null
+++ b/server/src/main/java/org/apache/rocketmq/studio/ops/ai/AiPayloadGuard.java
@@ -0,0 +1,157 @@
+/*
+ * 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 com.fasterxml.jackson.core.JsonProcessingException;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import org.apache.rocketmq.studio.common.exception.BusinessException;
+import org.springframework.util.StringUtils;
+
+import java.util.Map;
+
+/**
+ * Central byte budgets for user-controlled AI payloads. Character-count
validation is not
+ * sufficient here because provider requests, CLI arguments and tool payloads
are encoded as
+ * UTF-8 before leaving Studio.
+ */
+final class AiPayloadGuard {
+
+ static final int MAX_MESSAGE_BYTES = 64 * 1024;
+ static final int MAX_CONTEXT_BYTES = 256 * 1024;
+ static final int MAX_TOOL_INPUT_BYTES = 256 * 1024;
+ static final int MAX_OUTBOUND_PROMPT_BYTES = MAX_MESSAGE_BYTES +
MAX_CONTEXT_BYTES + 1024;
+ static final int MAX_MODEL_BYTES = 512;
+ static final int MAX_CONVERSATION_ID_BYTES = 256;
+ static final int MAX_SELECTOR_BYTES = 64;
+ static final int MAX_TOOL_NAME_BYTES = 256;
+
+ private AiPayloadGuard() {
+ }
+
+ static void validateChat(ChatDTO request) {
+ if (request == null) {
+ throw new BusinessException(400, "Chat request is required");
+ }
+ requireText(request.getMessage(), "Chat message is required");
+ requireWithin(request.getMessage(), MAX_MESSAGE_BYTES, "Chat message");
+ requireOptionalWithin(request.getModel(), MAX_MODEL_BYTES, "Chat
model");
+ requireOptionalWithin(request.getConversationId(),
MAX_CONVERSATION_ID_BYTES,
+ "Conversation ID");
+ requireOptionalWithin(request.getMode(), MAX_SELECTOR_BYTES, "Chat
mode");
+ requireOptionalWithin(request.getEngine(), MAX_SELECTOR_BYTES, "Chat
engine");
+ }
+
+ static void validateCommand(AiCommandDTO command, ObjectMapper
objectMapper) {
+ if (command == null) {
+ throw new BusinessException(400, "Command request is required");
+ }
+ if (!StringUtils.hasText(command.getPrompt()) &&
!StringUtils.hasText(command.getCommand())) {
+ throw new BusinessException(400, "Command or prompt is required");
+ }
+ requireOptionalWithin(command.getPrompt(), MAX_MESSAGE_BYTES, "Command
prompt");
+ requireOptionalWithin(command.getCommand(), MAX_MESSAGE_BYTES,
"Command text");
+ requireOptionalWithin(command.getModel(), MAX_MODEL_BYTES, "Command
model");
+ requireOptionalWithin(command.getConversationId(),
MAX_CONVERSATION_ID_BYTES,
+ "Conversation ID");
+ requireOptionalWithin(command.getMode(), MAX_SELECTOR_BYTES, "Command
mode");
+ requireOptionalWithin(command.getEngine(), MAX_SELECTOR_BYTES,
"Command engine");
+ requireJsonWithin(command.getContext(), MAX_CONTEXT_BYTES, "Command
context", objectMapper);
+ }
+
+ static void validateToolInvocation(String name, Map<String, Object> input,
ObjectMapper objectMapper) {
+ requireText(name, "Tool name is required");
+ requireWithin(name, MAX_TOOL_NAME_BYTES, "Tool name");
+ requireJsonWithin(input, MAX_TOOL_INPUT_BYTES, "Tool input",
objectMapper);
+ }
+
+ static void validateOutboundPrompt(String prompt, String model) {
+ if (exceedsUtf8Limit(prompt, MAX_OUTBOUND_PROMPT_BYTES)) {
+ throw requestTooLarge("LLM prompt", MAX_OUTBOUND_PROMPT_BYTES);
+ }
+ if (exceedsUtf8Limit(model, MAX_MODEL_BYTES)) {
+ throw requestTooLarge("LLM model", MAX_MODEL_BYTES);
+ }
+ }
+
+ private static void requireJsonWithin(Object value, int limitBytes, String
field, ObjectMapper objectMapper) {
+ if (value == null) {
+ return;
+ }
+ final int size;
+ try {
+ size = objectMapper.writeValueAsBytes(value).length;
+ } catch (JsonProcessingException exception) {
+ throw new BusinessException(400, field + " must be valid JSON");
+ }
+ if (size > limitBytes) {
+ throw new BusinessException(400, field + " must not exceed " +
limitBytes + " UTF-8 bytes");
+ }
+ }
+
+ private static void requireText(String value, String message) {
+ if (!StringUtils.hasText(value)) {
+ throw new BusinessException(400, message);
+ }
+ }
+
+ private static void requireOptionalWithin(String value, int limitBytes,
String field) {
+ if (value != null) {
+ requireWithin(value, limitBytes, field);
+ }
+ }
+
+ private static void requireWithin(String value, int limitBytes, String
field) {
+ if (exceedsUtf8Limit(value, limitBytes)) {
+ throw new BusinessException(400, field + " must not exceed " +
limitBytes + " UTF-8 bytes");
+ }
+ }
+
+ private static boolean exceedsUtf8Limit(String value, int limitBytes) {
+ if (value == null) {
+ return false;
+ }
+ int bytes = 0;
+ for (int index = 0; index < value.length(); index++) {
+ char current = value.charAt(index);
+ if (current <= 0x7f) {
+ bytes++;
+ } else if (current <= 0x7ff) {
+ bytes += 2;
+ } else if (Character.isHighSurrogate(current)
+ && index + 1 < value.length()
+ && Character.isLowSurrogate(value.charAt(index + 1))) {
+ bytes += 4;
+ index++;
+ } else if (Character.isSurrogate(current)) {
+ // The JDK UTF-8 encoder replaces an unpaired surrogate with a
one-byte '?'.
+ bytes++;
+ } else {
+ bytes += 3;
+ }
+ if (bytes > limitBytes) {
+ return true;
+ }
+ }
+ return false;
+ }
+
+ private static LlmGatewayException requestTooLarge(String field, int
limitBytes) {
+ return new LlmGatewayException(400, "llm.request.payload_too_large",
+ field + " exceeded the maximum of " + limitBytes + " UTF-8
bytes",
+ "Reduce the request size and retry.");
+ }
+}
diff --git
a/server/src/main/java/org/apache/rocketmq/studio/ops/ai/AiService.java
b/server/src/main/java/org/apache/rocketmq/studio/ops/ai/AiService.java
index 000331437..62ede2ee0 100644
--- a/server/src/main/java/org/apache/rocketmq/studio/ops/ai/AiService.java
+++ b/server/src/main/java/org/apache/rocketmq/studio/ops/ai/AiService.java
@@ -16,6 +16,7 @@
*/
package org.apache.rocketmq.studio.ops.ai;
+import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.apache.rocketmq.studio.common.exception.BusinessException;
@@ -32,13 +33,11 @@ public class AiService {
private final LlmGateway llmGateway;
private final McpServerRegistry mcpServerRegistry;
+ private final ObjectMapper objectMapper;
public SseEmitter chat(ChatDTO request) {
- if (request == null) {
- log.warn("Chat request body is missing");
- throw new BusinessException(400, "Chat request is required");
- }
+ AiPayloadGuard.validateChat(request);
log.info("Chat request received: mode={}, conversationId={}",
request.getMode(), request.getConversationId());
return llmGateway.chat(request);
}
@@ -52,6 +51,14 @@ public class AiService {
.result("Command request is required")
.build();
}
+ try {
+ AiPayloadGuard.validateCommand(command, objectMapper);
+ } catch (BusinessException exception) {
+ return AiExecuteResultVO.builder()
+ .success(false)
+ .result(exception.getMessage())
+ .build();
+ }
log.info("Executing AI command: {}", command.getCommand());
try {
String result = llmGateway.execute(command);
@@ -80,6 +87,7 @@ public class AiService {
}
public Object executeTool(String name, Map<String, Object> input) {
+ AiPayloadGuard.validateToolInvocation(name, input, objectMapper);
log.info("Executing registered AI tool: {}", name);
return mcpServerRegistry.execute(name, input);
}
diff --git
a/server/src/main/java/org/apache/rocketmq/studio/ops/ai/CliAgentProvider.java
b/server/src/main/java/org/apache/rocketmq/studio/ops/ai/CliAgentProvider.java
index d2db6d3f1..7e937d720 100644
---
a/server/src/main/java/org/apache/rocketmq/studio/ops/ai/CliAgentProvider.java
+++
b/server/src/main/java/org/apache/rocketmq/studio/ops/ai/CliAgentProvider.java
@@ -77,6 +77,10 @@ public abstract class CliAgentProvider implements
AgentProvider {
@Override
public String complete(LlmConfigVO config, String prompt, String
modelOverride) {
+ String effectiveModel = StringUtils.hasText(modelOverride)
+ ? modelOverride
+ : config == null ? null : config.getModel();
+ AiPayloadGuard.validateOutboundPrompt(prompt, effectiveModel);
if (!available()) {
throw new LlmGatewayException(503, "llm.provider.cli_missing",
binaryName() + " CLI is not installed in the server
runtime",
diff --git
a/server/src/main/java/org/apache/rocketmq/studio/ops/ai/OpenAiCompatibleLlmClient.java
b/server/src/main/java/org/apache/rocketmq/studio/ops/ai/OpenAiCompatibleLlmClient.java
index 1b1b6d228..054a98765 100644
---
a/server/src/main/java/org/apache/rocketmq/studio/ops/ai/OpenAiCompatibleLlmClient.java
+++
b/server/src/main/java/org/apache/rocketmq/studio/ops/ai/OpenAiCompatibleLlmClient.java
@@ -25,6 +25,7 @@ import org.springframework.util.StringUtils;
import java.io.BufferedReader;
import java.io.ByteArrayOutputStream;
+import java.io.FilterInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
@@ -84,6 +85,7 @@ public class OpenAiCompatibleLlmClient {
public String complete(LlmConfigVO config, String prompt, String
modelOverride) {
validate(config);
+ AiPayloadGuard.validateOutboundPrompt(prompt, effectiveModel(config,
modelOverride));
Map<String, Object> requestBody = requestBody(config, prompt,
modelOverride, false);
HttpRequest request = request(config, "application/json", requestBody);
try {
@@ -135,6 +137,7 @@ public class OpenAiCompatibleLlmClient {
public void stream(LlmConfigVO config, String prompt, String
modelOverride, Consumer<String> tokenConsumer) {
validate(config);
+ AiPayloadGuard.validateOutboundPrompt(prompt, effectiveModel(config,
modelOverride));
Map<String, Object> requestBody = requestBody(config, prompt,
modelOverride, true);
HttpRequest request = request(config, "text/event-stream",
requestBody);
try {
@@ -143,6 +146,10 @@ public class OpenAiCompatibleLlmClient {
throw upstreamException(response.statusCode(),
checkedBody(response.body().errorBody()));
}
parseStreamWithTimeout(response.body().stream(), tokenConsumer);
+ } catch (ResponseLimitException exception) {
+ throw new LlmGatewayException(502,
"llm.provider.response_too_large",
+ "LLM provider stream exceeded the maximum of " +
exception.limitBytes() + " bytes",
+ "Reduce the provider response size and retry.", exception);
} catch (HttpTimeoutException exception) {
throw new LlmGatewayException(504, "llm.provider.timeout",
"LLM provider stream timed out",
@@ -231,7 +238,8 @@ public class OpenAiCompatibleLlmClient {
new LimitedBodySubscriber(responseBodyLimitBytes()),
StreamBody::error);
}
return HttpResponse.BodySubscribers.mapping(
- HttpResponse.BodySubscribers.ofInputStream(),
StreamBody::stream);
+ HttpResponse.BodySubscribers.ofInputStream(),
+ input -> StreamBody.stream(new LimitedInputStream(input,
responseBodyLimitBytes())));
};
}
@@ -301,6 +309,10 @@ public class OpenAiCompatibleLlmClient {
return body;
}
+ private String effectiveModel(LlmConfigVO config, String modelOverride) {
+ return StringUtils.hasText(modelOverride) ? modelOverride :
config.getModel();
+ }
+
private URI chatCompletionsUri(LlmConfigVO config) {
return providerUri(config, CHAT_COMPLETIONS_PATH);
}
@@ -552,4 +564,55 @@ public class OpenAiCompatibleLlmClient {
}
}
}
+
+ private static final class LimitedInputStream extends FilterInputStream {
+
+ private final int limitBytes;
+ private int bytesRead;
+
+ private LimitedInputStream(InputStream input, int limitBytes) {
+ super(input);
+ this.limitBytes = limitBytes;
+ }
+
+ @Override
+ public int read() throws IOException {
+ int value = super.read();
+ if (value != -1) {
+ addBytes(1);
+ }
+ return value;
+ }
+
+ @Override
+ public int read(byte[] buffer, int offset, int length) throws
IOException {
+ int allowed = Math.min(length, limitBytes - bytesRead + 1);
+ int read = super.read(buffer, offset, allowed);
+ if (read > 0) {
+ addBytes(read);
+ }
+ return read;
+ }
+
+ private void addBytes(int count) throws ResponseLimitException {
+ bytesRead += count;
+ if (bytesRead > limitBytes) {
+ throw new ResponseLimitException(limitBytes);
+ }
+ }
+ }
+
+ private static final class ResponseLimitException extends IOException {
+
+ private final int limitBytes;
+
+ private ResponseLimitException(int limitBytes) {
+ super("LLM provider stream exceeds " + limitBytes + " bytes");
+ this.limitBytes = limitBytes;
+ }
+
+ private int limitBytes() {
+ return limitBytes;
+ }
+ }
}
diff --git
a/server/src/test/java/org/apache/rocketmq/studio/ops/ai/AiPayloadGuardTest.java
b/server/src/test/java/org/apache/rocketmq/studio/ops/ai/AiPayloadGuardTest.java
new file mode 100644
index 000000000..f4e6b026e
--- /dev/null
+++
b/server/src/test/java/org/apache/rocketmq/studio/ops/ai/AiPayloadGuardTest.java
@@ -0,0 +1,77 @@
+/*
+ * 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 com.fasterxml.jackson.databind.ObjectMapper;
+import org.apache.rocketmq.studio.common.exception.BusinessException;
+import org.junit.jupiter.api.Test;
+
+import java.util.Map;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatCode;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+class AiPayloadGuardTest {
+
+ private final ObjectMapper objectMapper = new ObjectMapper();
+
+ @Test
+ void chatMeasuresUtf8BytesAndAcceptsTheExactBoundary() {
+ ChatDTO boundary = ChatDTO.builder()
+ .message("\u754c".repeat(AiPayloadGuard.MAX_MESSAGE_BYTES / 3)
+ "x")
+ .build();
+ ChatDTO oversized = ChatDTO.builder()
+ .message(boundary.getMessage() + "\u754c")
+ .build();
+
+ assertThatCode(() ->
AiPayloadGuard.validateChat(boundary)).doesNotThrowAnyException();
+ assertThatThrownBy(() -> AiPayloadGuard.validateChat(oversized))
+ .isInstanceOf(BusinessException.class)
+ .hasMessageContaining("UTF-8 bytes");
+ }
+
+ @Test
+ void commandRequiresPromptOrCommand() {
+ AiCommandDTO request =
AiCommandDTO.builder().context(Map.of("cluster", "main")).build();
+
+ assertThatThrownBy(() -> AiPayloadGuard.validateCommand(request,
objectMapper))
+ .isInstanceOf(BusinessException.class)
+ .hasMessage("Command or prompt is required");
+ }
+
+ @Test
+ void toolInputUsesItsSerializedJsonSize() {
+ Map<String, Object> input = Map.of(
+ "payload", "x".repeat(AiPayloadGuard.MAX_TOOL_INPUT_BYTES));
+
+ assertThatThrownBy(() ->
AiPayloadGuard.validateToolInvocation("rmq.query", input, objectMapper))
+ .isInstanceOf(BusinessException.class)
+ .hasMessageContaining("Tool input must not exceed");
+ }
+
+ @Test
+ void outboundModelUsesTheProviderErrorContract() {
+ assertThatThrownBy(() -> AiPayloadGuard.validateOutboundPrompt(
+ "hello", "\u6a21".repeat(AiPayloadGuard.MAX_MODEL_BYTES / 3 +
1)))
+ .isInstanceOfSatisfying(LlmGatewayException.class, exception
-> {
+ assertThat(exception.getStatusCode()).isEqualTo(400);
+
assertThat(exception.getCode()).isEqualTo("llm.request.payload_too_large");
+ assertThat(exception.getMessage()).contains("LLM model");
+ });
+ }
+}
diff --git
a/server/src/test/java/org/apache/rocketmq/studio/ops/ai/AiServiceTest.java
b/server/src/test/java/org/apache/rocketmq/studio/ops/ai/AiServiceTest.java
index d98a751f3..d50f9c46c 100644
--- a/server/src/test/java/org/apache/rocketmq/studio/ops/ai/AiServiceTest.java
+++ b/server/src/test/java/org/apache/rocketmq/studio/ops/ai/AiServiceTest.java
@@ -16,10 +16,11 @@
*/
package org.apache.rocketmq.studio.ops.ai;
+import com.fasterxml.jackson.databind.ObjectMapper;
import org.apache.rocketmq.studio.common.exception.BusinessException;
+import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
-import org.mockito.InjectMocks;
import org.mockito.Mock;
import org.mockito.junit.jupiter.MockitoExtension;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
@@ -46,9 +47,13 @@ class AiServiceTest {
@Mock
private McpServerRegistry mcpServerRegistry;
- @InjectMocks
private AiService aiService;
+ @BeforeEach
+ void setUp() {
+ aiService = new AiService(llmGateway, mcpServerRegistry, new
ObjectMapper());
+ }
+
@Test
void chatShouldReturnSseEmitterFromGateway() {
ChatDTO request = ChatDTO.builder()
@@ -215,6 +220,43 @@ class AiServiceTest {
verifyNoInteractions(llmGateway);
}
+ @Test
+ void chatRejectsOversizedMessageBeforeCallingGateway() {
+ ChatDTO request = ChatDTO.builder()
+ .message("\u754c".repeat(AiPayloadGuard.MAX_MESSAGE_BYTES / 3
+ 1))
+ .build();
+
+ assertThatThrownBy(() -> aiService.chat(request))
+ .isInstanceOf(BusinessException.class)
+ .hasMessageContaining("Chat message must not exceed");
+ verifyNoInteractions(llmGateway);
+ }
+
+ @Test
+ void executeRejectsOversizedContextBeforeCallingGateway() {
+ AiCommandDTO command = AiCommandDTO.builder()
+ .command("query_metrics")
+ .context(Map.of("payload",
"x".repeat(AiPayloadGuard.MAX_CONTEXT_BYTES)))
+ .build();
+
+ AiExecuteResultVO result = aiService.execute(command);
+
+ assertThat(result.isSuccess()).isFalse();
+ assertThat(result.getResult()).contains("Command context must not
exceed");
+ verifyNoInteractions(llmGateway);
+ }
+
+ @Test
+ void executeToolRejectsOversizedInputBeforeCallingRegistry() {
+ Map<String, Object> input = Map.of(
+ "payload", "x".repeat(AiPayloadGuard.MAX_TOOL_INPUT_BYTES));
+
+ assertThatThrownBy(() -> aiService.executeTool("rmq.capabilities",
input))
+ .isInstanceOf(BusinessException.class)
+ .hasMessageContaining("Tool input must not exceed");
+ verifyNoInteractions(mcpServerRegistry);
+ }
+
@Test
void executeHandlesNullCommand() {
AiExecuteResultVO result = aiService.execute(null);
diff --git
a/server/src/test/java/org/apache/rocketmq/studio/ops/ai/CliAgentProviderTest.java
b/server/src/test/java/org/apache/rocketmq/studio/ops/ai/CliAgentProviderTest.java
index 19fc6aac7..12f97eff4 100644
---
a/server/src/test/java/org/apache/rocketmq/studio/ops/ai/CliAgentProviderTest.java
+++
b/server/src/test/java/org/apache/rocketmq/studio/ops/ai/CliAgentProviderTest.java
@@ -125,6 +125,18 @@ class CliAgentProviderTest {
assertThat(cli.complete(null, "prompt", null)).isNotEmpty();
}
+ @Test
+ void completeRejectsOversizedPromptBeforeStartingCli() {
+ FakeCli cli = new FakeCli("echo should-not-run");
+
+ assertThatThrownBy(() -> cli.complete(
+ null, "x".repeat(AiPayloadGuard.MAX_OUTBOUND_PROMPT_BYTES +
1), null))
+ .isInstanceOfSatisfying(LlmGatewayException.class, exception
-> {
+ assertThat(exception.getStatusCode()).isEqualTo(400);
+
assertThat(exception.getCode()).isEqualTo("llm.request.payload_too_large");
+ });
+ }
+
@Test
void availabilityAndCompletionUseTheIsolatedEnvironment() {
RecordingEnvironment processEnvironment = new RecordingEnvironment();
diff --git
a/server/src/test/java/org/apache/rocketmq/studio/ops/ai/OpenAiCompatibleLlmClientTest.java
b/server/src/test/java/org/apache/rocketmq/studio/ops/ai/OpenAiCompatibleLlmClientTest.java
index eb07c2067..f10a908de 100644
---
a/server/src/test/java/org/apache/rocketmq/studio/ops/ai/OpenAiCompatibleLlmClientTest.java
+++
b/server/src/test/java/org/apache/rocketmq/studio/ops/ai/OpenAiCompatibleLlmClientTest.java
@@ -370,6 +370,45 @@ class OpenAiCompatibleLlmClientTest {
assertThat(exception.getCode()).isEqualTo("llm.provider.response_too_large"));
}
+ @Test
+ void streamShouldRejectOversizedSuccessfulResponse() {
+ client = clientWithLimit(1024);
+ String body = "data: {\"choices\":[{\"delta\":{\"content\":\""
+ + "x".repeat(1024)
+ + "\"}}]}\n\n";
+ server.createContext("/v1/chat/completions",
+ exchange -> respond(exchange, 200, body, "text/event-stream"));
+
+ assertThatThrownBy(() -> client.stream(config("openai", "sk-test"),
"hello", null, token -> { }))
+ .isInstanceOfSatisfying(LlmGatewayException.class, exception
-> {
+ assertThat(exception.getStatusCode()).isEqualTo(502);
+
assertThat(exception.getCode()).isEqualTo("llm.provider.response_too_large");
+ assertThat(exception.getMessage()).contains("1024 bytes");
+ });
+ }
+
+ @Test
+ void completeShouldRejectOversizedPromptBeforeCallingUpstream() {
+ assertThatThrownBy(() -> client.complete(
+ config("openai", "sk-test"),
+ "x".repeat(AiPayloadGuard.MAX_OUTBOUND_PROMPT_BYTES + 1),
+ null))
+ .isInstanceOfSatisfying(LlmGatewayException.class, exception
-> {
+ assertThat(exception.getStatusCode()).isEqualTo(400);
+
assertThat(exception.getCode()).isEqualTo("llm.request.payload_too_large");
+ });
+ }
+
+ @Test
+ void completeShouldRejectOversizedConfiguredModel() {
+ LlmConfigVO config = config("openai", "sk-test");
+ config.setModel("x".repeat(AiPayloadGuard.MAX_MODEL_BYTES + 1));
+
+ assertThatThrownBy(() -> client.complete(config, "hello", null))
+ .isInstanceOfSatisfying(LlmGatewayException.class, exception ->
+
assertThat(exception.getCode()).isEqualTo("llm.request.payload_too_large"));
+ }
+
@Test
void completeShouldExposeTimeoutAsGatewayTimeout() {
OpenAiCompatibleLlmClient timeoutClient = new
OpenAiCompatibleLlmClient(
