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 305c95c0 feat(ai): persist LLM generation settings and add message
query/trace tools (#1931)
305c95c0 is described below
commit 305c95c0565e3ef2adf50c0fa62abc6ff1fb69a5
Author: yyqdbngt <[email protected]>
AuthorDate: Fri Aug 14 11:33:52 2026 +0800
feat(ai): persist LLM generation settings and add message query/trace tools
(#1931)
Consolidates #1915, #1931, #1939: persist LLM maxTokens/temperature in
general settings, add message query and trace AI tools, bound the LLM
gateway executor with a caller-runs pool, and normalize gateway engine
identifiers with Locale.ROOT.
Co-authored-by: yyqdbngt <[email protected]>
---
.../rocketmq/studio/ops/ai/LlmConfigService.java | 6 +-
.../studio/ops/ai/OpenAiCompatibleLlmGateway.java | 15 ++-
.../ops/ai/tool/MessageQueryToolHandler.java | 94 +++++++++++++++
.../ops/ai/tool/MessageTraceToolHandler.java | 89 ++++++++++++++
.../studio/settings/GeneralSettingsVO.java | 2 +
.../src/main/resources/tool-catalog/rmq-tools.yaml | 128 +++++++++++++++++++++
.../rocketmq/studio/ops/ai/AiEngineLocaleTest.java | 55 +++++++++
.../studio/ops/ai/LlmConfigServiceTest.java | 4 +
.../ops/ai/tool/MessageQueryToolHandlerTest.java | 88 ++++++++++++++
.../ops/ai/tool/MessageTraceToolHandlerTest.java | 86 ++++++++++++++
.../studio/ops/ai/tool/ToolCatalogTest.java | 2 +
.../studio/ops/ai/tool/ToolGatewayServiceTest.java | 29 ++++-
12 files changed, 588 insertions(+), 10 deletions(-)
diff --git
a/server/src/main/java/org/apache/rocketmq/studio/ops/ai/LlmConfigService.java
b/server/src/main/java/org/apache/rocketmq/studio/ops/ai/LlmConfigService.java
index e2b9ad66..d5666b63 100644
---
a/server/src/main/java/org/apache/rocketmq/studio/ops/ai/LlmConfigService.java
+++
b/server/src/main/java/org/apache/rocketmq/studio/ops/ai/LlmConfigService.java
@@ -117,6 +117,8 @@ public class LlmConfigService {
.deploymentName(normalized.getDeploymentName())
.apiVersion(normalized.getApiVersion())
.awsRegion(normalized.getAwsRegion())
+ .maxTokens(normalized.getMaxTokens())
+ .temperature(normalized.getTemperature())
.build();
LlmConfigVO nextOverrides = copy(normalized);
settingsService.saveGeneralSettings(updated);
@@ -239,8 +241,8 @@ public class LlmConfigService {
.apiKey(apiKey)
.apiBase(apiBase)
.model(model)
- .maxTokens(DEFAULT_MAX_TOKENS)
- .temperature(DEFAULT_TEMPERATURE)
+ .maxTokens(settings.getMaxTokens() != null ?
settings.getMaxTokens() : DEFAULT_MAX_TOKENS)
+ .temperature(settings.getTemperature() != null ?
settings.getTemperature() : DEFAULT_TEMPERATURE)
.enabled(!requiresApiKey(provider) ||
!!StringUtils.hasText(apiKey))
.deploymentName(defaultString(settings.getDeploymentName(),
""))
.apiVersion(defaultString(settings.getApiVersion(),
"2024-02-15-preview"))
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 2ea5eda4..8ec0f825 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
@@ -29,9 +29,12 @@ import
org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import java.io.IOException;
import java.io.InputStream;
import java.nio.charset.StandardCharsets;
+import java.util.Locale;
import java.util.Map;
import java.util.concurrent.ExecutorService;
-import java.util.concurrent.Executors;
+import java.util.concurrent.SynchronousQueue;
+import java.util.concurrent.ThreadPoolExecutor;
+import java.util.concurrent.TimeUnit;
@Slf4j
@Primary
@@ -43,7 +46,12 @@ public class OpenAiCompatibleLlmGateway implements
LlmGateway {
private final OpenAiCompatibleLlmClient llmClient;
private final AgentProviderRegistry agentProviders;
private final ObjectMapper objectMapper;
- private final ExecutorService executor = Executors.newCachedThreadPool();
+ // 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());
@Override
public SseEmitter chat(ChatDTO request) {
@@ -83,7 +91,8 @@ public class OpenAiCompatibleLlmGateway implements LlmGateway
{
/** Request-level engine (per-user preference) overrides the global
config. */
private String resolveEngine(String requestEngine, LlmConfigVO config) {
- String engine = StringUtils.hasText(requestEngine) ?
requestEngine.trim().toLowerCase() : null;
+ String engine = StringUtils.hasText(requestEngine)
+ ? requestEngine.trim().toLowerCase(Locale.ROOT) : null;
if (engine == null) {
return config.normalizeEngine();
}
diff --git
a/server/src/main/java/org/apache/rocketmq/studio/ops/ai/tool/MessageQueryToolHandler.java
b/server/src/main/java/org/apache/rocketmq/studio/ops/ai/tool/MessageQueryToolHandler.java
new file mode 100644
index 00000000..6ea69ca0
--- /dev/null
+++
b/server/src/main/java/org/apache/rocketmq/studio/ops/ai/tool/MessageQueryToolHandler.java
@@ -0,0 +1,94 @@
+/*
+ * 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.tool;
+
+import org.apache.rocketmq.studio.instance.message.MessageRecordVO;
+import org.apache.rocketmq.studio.instance.message.MessageService;
+import lombok.RequiredArgsConstructor;
+import org.springframework.stereotype.Component;
+
+import java.util.LinkedHashMap;
+import java.util.Map;
+
+/**
+ * Queries messages in a RocketMQ instance by topic, message id, business key
or time range. Read-only
+ * and safe for AI/MCP/CLI callers; delegates to the standard message service
used by the web UI.
+ */
+@Component
+@RequiredArgsConstructor
+public class MessageQueryToolHandler implements ToolHandler {
+
+ private static final String NAME = "rmq.message.query";
+
+ private final MessageService messageService;
+
+ @Override
+ public String name() {
+ return NAME;
+ }
+
+ @Override
+ public Object execute(Map<String, Object> input) {
+ String instanceId = (String) input.get("cluster");
+ String topic = (String) input.get("topic");
+ String msgId = (String) input.get("msgId");
+ String tag = (String) input.get("tag");
+ String key = (String) input.get("key");
+ Long startTime = asLong(input.get("startTime"));
+ Long endTime = asLong(input.get("endTime"));
+ return messageService.queryMessages(instanceId, topic, msgId, tag,
key, startTime, endTime).stream()
+ .map(MessageQueryToolHandler::safeProjection)
+ .toList();
+ }
+
+ private static Map<String, Object> safeProjection(MessageRecordVO message)
{
+ Map<String, Object> result = new LinkedHashMap<>();
+ result.put("msgId", require(message.getMsgId(), "msgId"));
+ result.put("topic", blankIfNull(message.getTopic()));
+ result.put("tag", blankIfNull(message.getTag()));
+ result.put("key", blankIfNull(message.getKey()));
+ result.put("storeTime", message.getStoreTime());
+ result.put("storeHost", blankIfNull(message.getStoreHost()));
+ result.put("bornHost", blankIfNull(message.getBornHost()));
+ result.put("body", blankIfNull(message.getBody()));
+ result.put("bodyEncoding", blankIfNull(message.getBodyEncoding()));
+ result.put("bodyTruncated", message.isBodyTruncated());
+ result.put("size", message.getSize());
+ return result;
+ }
+
+ private static Long asLong(Object value) {
+ if (value == null) {
+ return null;
+ }
+ if (value instanceof Number number) {
+ return number.longValue();
+ }
+ return Long.parseLong(value.toString());
+ }
+
+ private static String require(String value, String field) {
+ if (value == null || value.isBlank()) {
+ throw new IllegalStateException("Message " + field + " is
unavailable");
+ }
+ return value;
+ }
+
+ private static String blankIfNull(String value) {
+ return value == null ? "" : value;
+ }
+}
diff --git
a/server/src/main/java/org/apache/rocketmq/studio/ops/ai/tool/MessageTraceToolHandler.java
b/server/src/main/java/org/apache/rocketmq/studio/ops/ai/tool/MessageTraceToolHandler.java
new file mode 100644
index 00000000..ef30a5e4
--- /dev/null
+++
b/server/src/main/java/org/apache/rocketmq/studio/ops/ai/tool/MessageTraceToolHandler.java
@@ -0,0 +1,89 @@
+/*
+ * 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.tool;
+
+import org.apache.rocketmq.studio.instance.message.ConsumerStatusVO;
+import org.apache.rocketmq.studio.instance.message.MessageService;
+import org.apache.rocketmq.studio.instance.message.TraceNodeVO;
+import org.apache.rocketmq.studio.instance.message.TraceRecordVO;
+import lombok.RequiredArgsConstructor;
+import org.springframework.stereotype.Component;
+
+import java.util.LinkedHashMap;
+import java.util.Map;
+
+/**
+ * Returns the message trace timeline for one message id. Read-only and safe
for AI/MCP/CLI callers.
+ */
+@Component
+@RequiredArgsConstructor
+public class MessageTraceToolHandler implements ToolHandler {
+
+ private static final String NAME = "rmq.message.trace";
+
+ private final MessageService messageService;
+
+ @Override
+ public String name() {
+ return NAME;
+ }
+
+ @Override
+ public Object execute(Map<String, Object> input) {
+ String instanceId = (String) input.get("cluster");
+ String msgId = (String) input.get("msgId");
+ String topic = (String) input.get("topic");
+ TraceRecordVO trace = messageService.getMessageTrace(instanceId,
msgId, topic);
+ return project(msgId, trace);
+ }
+
+ private static Map<String, Object> project(String msgId, TraceRecordVO
trace) {
+ Map<String, Object> result = new LinkedHashMap<>();
+ result.put("msgId", msgId);
+ result.put("nodes", trace.getNodes().stream()
+ .map(MessageTraceToolHandler::projectNode)
+ .toList());
+ result.put("consumerStatus", trace.getConsumerStatus().stream()
+ .map(MessageTraceToolHandler::projectConsumerStatus)
+ .toList());
+ return result;
+ }
+
+ private static Map<String, Object> projectNode(TraceNodeVO node) {
+ Map<String, Object> result = new LinkedHashMap<>();
+ result.put("title", blankIfNull(node.getTitle()));
+ result.put("timestamp", node.getTimestamp());
+ result.put("status", blankIfNull(node.getStatus()));
+ result.put("costTime", node.getCostTime());
+ result.put("description", blankIfNull(node.getDescription()));
+ return result;
+ }
+
+ private static Map<String, Object> projectConsumerStatus(ConsumerStatusVO
status) {
+ Map<String, Object> result = new LinkedHashMap<>();
+ result.put("group", blankIfNull(status.getGroup()));
+ result.put("deliveryStatus",
+ status.getDeliveryStatus() == null ? "" :
status.getDeliveryStatus().name());
+ result.put("consumeTime", status.getConsumeTime());
+ result.put("retryCount", status.getRetryCount());
+ return result;
+ }
+
+ private static String blankIfNull(String value) {
+ return value == null ? "" : value;
+ }
+}
diff --git
a/server/src/main/java/org/apache/rocketmq/studio/settings/GeneralSettingsVO.java
b/server/src/main/java/org/apache/rocketmq/studio/settings/GeneralSettingsVO.java
index bdcd48d5..814205bf 100644
---
a/server/src/main/java/org/apache/rocketmq/studio/settings/GeneralSettingsVO.java
+++
b/server/src/main/java/org/apache/rocketmq/studio/settings/GeneralSettingsVO.java
@@ -47,6 +47,8 @@ public class GeneralSettingsVO {
private String deploymentName;
private String apiVersion;
private String awsRegion;
+ private Integer maxTokens;
+ private Double temperature;
@JsonProperty(value = "apiKeyConfigured", access =
JsonProperty.Access.READ_ONLY)
public boolean isApiKeyConfigured() {
diff --git a/server/src/main/resources/tool-catalog/rmq-tools.yaml
b/server/src/main/resources/tool-catalog/rmq-tools.yaml
index 42885442..867a7b11 100644
--- a/server/src/main/resources/tool-catalog/rmq-tools.yaml
+++ b/server/src/main/resources/tool-catalog/rmq-tools.yaml
@@ -303,6 +303,134 @@ tools:
type: integer
viewHint: table
deprecated: false
+ - name: rmq.message.query
+ cli:
+ resource: message
+ verb: query
+ description: Query messages in a RocketMQ instance by topic, message id,
key or time range.
+ riskLevel: L1
+ permission: message:read
+ requiredCapabilities: []
+ inputSchema:
+ type: object
+ required:
+ - cluster
+ additionalProperties: false
+ properties:
+ cluster:
+ type: string
+ minLength: 1
+ topic:
+ type: string
+ msgId:
+ type: string
+ key:
+ type: string
+ tag:
+ type: string
+ startTime:
+ type: integer
+ endTime:
+ type: integer
+ outputSchema:
+ type: array
+ items:
+ type: object
+ required:
+ - msgId
+ - topic
+ additionalProperties: false
+ properties:
+ msgId:
+ type: string
+ topic:
+ type: string
+ tag:
+ type: string
+ key:
+ type: string
+ storeTime:
+ type: integer
+ storeHost:
+ type: string
+ bornHost:
+ type: string
+ body:
+ type: string
+ bodyEncoding:
+ type: string
+ bodyTruncated:
+ type: boolean
+ size:
+ type: integer
+ viewHint: table
+ deprecated: false
+ - name: rmq.message.trace
+ cli:
+ resource: message
+ verb: trace
+ description: Return the message trace timeline for one message id.
+ riskLevel: L1
+ permission: message:read
+ requiredCapabilities: []
+ inputSchema:
+ type: object
+ required:
+ - cluster
+ - msgId
+ additionalProperties: false
+ properties:
+ cluster:
+ type: string
+ minLength: 1
+ msgId:
+ type: string
+ minLength: 1
+ topic:
+ type: string
+ outputSchema:
+ type: object
+ required:
+ - msgId
+ - nodes
+ additionalProperties: false
+ properties:
+ msgId:
+ type: string
+ nodes:
+ type: array
+ items:
+ type: object
+ required:
+ - title
+ additionalProperties: false
+ properties:
+ title:
+ type: string
+ timestamp:
+ type: integer
+ status:
+ type: string
+ costTime:
+ type: integer
+ description:
+ type: string
+ consumerStatus:
+ type: array
+ items:
+ type: object
+ additionalProperties: false
+ properties:
+ group:
+ type: string
+ deliveryStatus:
+ type: string
+ consumeTime:
+ type: integer
+ retryCount:
+ type: integer
+ viewHint: text
+ deprecated: false
- name: rmq.alert.rule.list
cli:
resource: alert-rule
diff --git
a/server/src/test/java/org/apache/rocketmq/studio/ops/ai/AiEngineLocaleTest.java
b/server/src/test/java/org/apache/rocketmq/studio/ops/ai/AiEngineLocaleTest.java
new file mode 100644
index 00000000..6df014e9
--- /dev/null
+++
b/server/src/test/java/org/apache/rocketmq/studio/ops/ai/AiEngineLocaleTest.java
@@ -0,0 +1,55 @@
+/*
+ * 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 java.util.List;
+import java.util.Locale;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+class AiEngineLocaleTest {
+
+ @Test
+ void engineIdentifiersShouldIgnoreTheJvmDefaultLocale() {
+ Locale previous = Locale.getDefault();
+ Locale.setDefault(Locale.forLanguageTag("tr-TR"));
+ try {
+ LlmConfigVO config = LlmConfigVO.builder().engine(" CUSTOM-CLI
").build();
+ AgentProvider provider = new StubProvider("custom-cli");
+ AgentProviderRegistry registry = new
AgentProviderRegistry(List.of(provider));
+
+ assertThat(config.normalizeEngine()).isEqualTo("custom-cli");
+ assertThat(registry.forEngine(" CUSTOM-CLI ")).isSameAs(provider);
+ } finally {
+ Locale.setDefault(previous);
+ }
+ }
+
+ private record StubProvider(String engine) implements AgentProvider {
+ @Override
+ public boolean available() {
+ return true;
+ }
+
+ @Override
+ public String complete(LlmConfigVO config, String prompt, String
modelOverride) {
+ return "ok";
+ }
+ }
+}
diff --git
a/server/src/test/java/org/apache/rocketmq/studio/ops/ai/LlmConfigServiceTest.java
b/server/src/test/java/org/apache/rocketmq/studio/ops/ai/LlmConfigServiceTest.java
index f942b2f4..77a774cd 100644
---
a/server/src/test/java/org/apache/rocketmq/studio/ops/ai/LlmConfigServiceTest.java
+++
b/server/src/test/java/org/apache/rocketmq/studio/ops/ai/LlmConfigServiceTest.java
@@ -185,7 +185,11 @@ class LlmConfigServiceTest {
assertThat(saved.getApiKey()).isEqualTo("sk-deepseek");
assertThat(saved.getModel()).isEqualTo("deepseek-chat");
assertThat(saved.getBaseUrl()).isEqualTo("https://api.deepseek.com/v1";);
+ assertThat(saved.getMaxTokens()).isEqualTo(8192);
+ assertThat(saved.getTemperature()).isEqualTo(0.2);
assertThat(llmConfigService.getConfig().getProvider()).isEqualTo("deepseek");
+
assertThat(llmConfigService.getConfig().getMaxTokens()).isEqualTo(8192);
+
assertThat(llmConfigService.getConfig().getTemperature()).isEqualTo(0.2);
}
@Test
diff --git
a/server/src/test/java/org/apache/rocketmq/studio/ops/ai/tool/MessageQueryToolHandlerTest.java
b/server/src/test/java/org/apache/rocketmq/studio/ops/ai/tool/MessageQueryToolHandlerTest.java
new file mode 100644
index 00000000..067ca8c8
--- /dev/null
+++
b/server/src/test/java/org/apache/rocketmq/studio/ops/ai/tool/MessageQueryToolHandlerTest.java
@@ -0,0 +1,88 @@
+/*
+ * 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.tool;
+
+import org.apache.rocketmq.studio.instance.message.MessageRecordVO;
+import org.apache.rocketmq.studio.instance.message.MessageService;
+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 java.util.List;
+import java.util.Map;
+
+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.verify;
+import static org.mockito.Mockito.when;
+
+@ExtendWith(MockitoExtension.class)
+class MessageQueryToolHandlerTest {
+
+ @Mock
+ private MessageService messageService;
+
+ @InjectMocks
+ private MessageQueryToolHandler handler;
+
+ @Test
+ void executeShouldDelegateToMessageServiceAndProject() {
+ MessageRecordVO message = MessageRecordVO.builder()
+ .msgId("msg-1")
+ .topic("TopicA")
+ .tag("tag1")
+ .key("key1")
+ .body("hello")
+ .bodyEncoding("UTF-8")
+ .bodyTruncated(false)
+ .storeTime(1000L)
+ .bornHost("10.0.0.1")
+ .storeHost("10.0.0.2")
+ .size(5)
+ .build();
+ when(messageService.queryMessages(eq("instance-a"), eq("TopicA"),
any(), any(), any(), any(), any()))
+ .thenReturn(List.of(message));
+
+ Object result = handler.execute(Map.of("cluster", "instance-a",
"topic", "TopicA"));
+
+ assertThat(result).isInstanceOf(List.class);
+ List<?> rows = (List<?>) result;
+ assertThat(rows).hasSize(1);
+ Map<?, ?> row = (Map<?, ?>) rows.get(0);
+ assertThat(row.get("msgId")).isEqualTo("msg-1");
+ assertThat(row.get("topic")).isEqualTo("TopicA");
+ assertThat(row.get("tag")).isEqualTo("tag1");
+ assertThat(row.get("storeTime")).isEqualTo(1000L);
+ assertThat(row.get("body")).isEqualTo("hello");
+ assertThat(row.get("size")).isEqualTo(5);
+ }
+
+ @Test
+ void executeShouldConvertNumericTimeArguments() {
+ when(messageService.queryMessages(any(), any(), any(), any(), any(),
any(), any()))
+ .thenReturn(List.of());
+
+ handler.execute(Map.of("cluster", "instance-a", "topic", "TopicA",
+ "startTime", 1000L, "endTime", 2000L));
+
+ verify(messageService)
+ .queryMessages(eq("instance-a"), any(), any(), any(), any(),
eq(1000L), eq(2000L));
+ }
+}
diff --git
a/server/src/test/java/org/apache/rocketmq/studio/ops/ai/tool/MessageTraceToolHandlerTest.java
b/server/src/test/java/org/apache/rocketmq/studio/ops/ai/tool/MessageTraceToolHandlerTest.java
new file mode 100644
index 00000000..e495a8ca
--- /dev/null
+++
b/server/src/test/java/org/apache/rocketmq/studio/ops/ai/tool/MessageTraceToolHandlerTest.java
@@ -0,0 +1,86 @@
+/*
+ * 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.tool;
+
+import org.apache.rocketmq.studio.instance.message.ConsumerStatusVO;
+import org.apache.rocketmq.studio.instance.message.MessageService;
+import org.apache.rocketmq.studio.instance.message.TraceNodeVO;
+import org.apache.rocketmq.studio.instance.message.TraceRecordVO;
+import org.apache.rocketmq.studio.common.domain.enums.DeliveryStatus;
+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 java.util.List;
+import java.util.Map;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+@ExtendWith(MockitoExtension.class)
+class MessageTraceToolHandlerTest {
+
+ @Mock
+ private MessageService messageService;
+
+ @InjectMocks
+ private MessageTraceToolHandler handler;
+
+ @Test
+ void executeShouldDelegateToMessageServiceAndProject() {
+ TraceRecordVO trace = TraceRecordVO.builder()
+ .nodes(List.of(TraceNodeVO.builder()
+ .title("Send message")
+ .timestamp(1000L)
+ .status("SUCCESS")
+ .costTime(5L)
+ .description("msg-1 sent")
+ .build()))
+ .consumerStatus(List.of(ConsumerStatusVO.builder()
+ .group("group-a")
+ .deliveryStatus(DeliveryStatus.success)
+ .consumeTime(2000L)
+ .retryCount(0)
+ .build()))
+ .build();
+ when(messageService.getMessageTrace(eq("instance-a"), eq("msg-1"),
eq("TopicA")))
+ .thenReturn(trace);
+
+ Object result = handler.execute(Map.of(
+ "cluster", "instance-a", "msgId", "msg-1", "topic", "TopicA"));
+
+ assertThat(result).isInstanceOf(Map.class);
+ Map<?, ?> row = (Map<?, ?>) result;
+ assertThat(row.get("msgId")).isEqualTo("msg-1");
+ List<?> nodes = (List<?>) row.get("nodes");
+ assertThat(nodes).hasSize(1);
+ Map<?, ?> node = (Map<?, ?>) nodes.get(0);
+ assertThat(node.get("title")).isEqualTo("Send message");
+ assertThat(node.get("status")).isEqualTo("SUCCESS");
+ List<?> statuses = (List<?>) row.get("consumerStatus");
+ assertThat(statuses).hasSize(1);
+ Map<?, ?> status = (Map<?, ?>) statuses.get(0);
+ assertThat(status.get("group")).isEqualTo("group-a");
+ assertThat(status.get("deliveryStatus")).isEqualTo("success");
+
+ verify(messageService).getMessageTrace("instance-a", "msg-1",
"TopicA");
+ }
+}
diff --git
a/server/src/test/java/org/apache/rocketmq/studio/ops/ai/tool/ToolCatalogTest.java
b/server/src/test/java/org/apache/rocketmq/studio/ops/ai/tool/ToolCatalogTest.java
index 0a46bb20..1389e979 100644
---
a/server/src/test/java/org/apache/rocketmq/studio/ops/ai/tool/ToolCatalogTest.java
+++
b/server/src/test/java/org/apache/rocketmq/studio/ops/ai/tool/ToolCatalogTest.java
@@ -44,6 +44,8 @@ class ToolCatalogTest {
"rmq.dashboard.summary",
"rmq.topic.list",
"rmq.group.list",
+ "rmq.message.query",
+ "rmq.message.trace",
"rmq.alert.rule.list",
"rmq.nameserver.config.diff");
assertThat(catalog.find("rmq.cluster.list")).isPresent();
diff --git
a/server/src/test/java/org/apache/rocketmq/studio/ops/ai/tool/ToolGatewayServiceTest.java
b/server/src/test/java/org/apache/rocketmq/studio/ops/ai/tool/ToolGatewayServiceTest.java
index e3fc309a..9876b4b7 100644
---
a/server/src/test/java/org/apache/rocketmq/studio/ops/ai/tool/ToolGatewayServiceTest.java
+++
b/server/src/test/java/org/apache/rocketmq/studio/ops/ai/tool/ToolGatewayServiceTest.java
@@ -29,6 +29,7 @@ import
org.apache.rocketmq.studio.common.domain.enums.TopicPerm;
import org.apache.rocketmq.studio.common.domain.enums.TopicType;
import org.apache.rocketmq.studio.common.exception.BusinessException;
import org.apache.rocketmq.studio.instance.group.ConsumerGroupVO;
+import org.apache.rocketmq.studio.instance.message.MessageService;
import org.apache.rocketmq.studio.instance.topic.MetadataService;
import org.apache.rocketmq.studio.instance.topic.TopicVO;
import org.apache.rocketmq.studio.ops.ai.AiToolVO;
@@ -60,6 +61,7 @@ class ToolGatewayServiceTest {
private ToolCatalog catalog;
private ClusterService clusterService;
private DashboardService dashboardService;
+ private MessageService messageService;
private MetadataService metadataService;
private AlertService alertService;
private NameServerConfigDiffService nameServerConfigDiffService;
@@ -71,6 +73,8 @@ class ToolGatewayServiceTest {
private ConsumerGroupListToolHandler consumerGroupListHandler;
private AlertRuleListToolHandler alertRuleListHandler;
private NameServerConfigDiffToolHandler nameServerConfigDiffHandler;
+ private MessageQueryToolHandler messageQueryHandler;
+ private MessageTraceToolHandler messageTraceHandler;
private ToolGatewayService gateway;
@BeforeEach
@@ -78,6 +82,7 @@ class ToolGatewayServiceTest {
catalog = canonicalCatalog();
clusterService = mock(ClusterService.class);
dashboardService = mock(DashboardService.class);
+ messageService = mock(MessageService.class);
metadataService = mock(MetadataService.class);
alertService = mock(AlertService.class);
nameServerConfigDiffService = mock(NameServerConfigDiffService.class);
@@ -90,6 +95,8 @@ class ToolGatewayServiceTest {
alertRuleListHandler = new AlertRuleListToolHandler(alertService);
nameServerConfigDiffHandler = new NameServerConfigDiffToolHandler(
nameServerConfigDiffService);
+ messageQueryHandler = new MessageQueryToolHandler(messageService);
+ messageTraceHandler = new MessageTraceToolHandler(messageService);
gateway = gateway(
catalog,
clusterListHandler,
@@ -98,7 +105,9 @@ class ToolGatewayServiceTest {
topicListHandler,
consumerGroupListHandler,
alertRuleListHandler,
- nameServerConfigDiffHandler);
+ nameServerConfigDiffHandler,
+ messageQueryHandler,
+ messageTraceHandler);
}
@Test
@@ -121,6 +130,8 @@ class ToolGatewayServiceTest {
"rmq.dashboard.summary",
"rmq.topic.list",
"rmq.group.list",
+ "rmq.message.query",
+ "rmq.message.trace",
"rmq.alert.rule.list",
"rmq.nameserver.config.diff");
}
@@ -481,7 +492,9 @@ class ToolGatewayServiceTest {
topicListHandler,
consumerGroupListHandler,
alertRuleListHandler,
- nameServerConfigDiffHandler);
+ nameServerConfigDiffHandler,
+ messageQueryHandler,
+ messageTraceHandler);
assertThatThrownBy(() -> l2Gateway.execute("rmq.cluster.list",
Map.of()))
.isInstanceOf(BusinessException.class)
@@ -533,7 +546,9 @@ class ToolGatewayServiceTest {
topicListHandler,
consumerGroupListHandler,
alertRuleListHandler,
- nameServerConfigDiffHandler))
+ nameServerConfigDiffHandler,
+ messageQueryHandler,
+ messageTraceHandler))
.isInstanceOf(IllegalStateException.class)
.hasMessageContaining("input schema")
.hasMessageContaining("rmq.cluster.list");
@@ -559,7 +574,9 @@ class ToolGatewayServiceTest {
topicListHandler,
consumerGroupListHandler,
alertRuleListHandler,
- nameServerConfigDiffHandler))
+ nameServerConfigDiffHandler,
+ messageQueryHandler,
+ messageTraceHandler))
.isInstanceOf(IllegalStateException.class)
.hasMessageContaining("input schema")
.hasMessageContaining("rmq.cluster.list");
@@ -586,7 +603,9 @@ class ToolGatewayServiceTest {
topicListHandler,
consumerGroupListHandler,
alertRuleListHandler,
- nameServerConfigDiffHandler);
+ nameServerConfigDiffHandler,
+ messageQueryHandler,
+ messageTraceHandler);
assertThatThrownBy(() -> invalidGateway.execute("rmq.cluster.list",
Map.of()))
.isInstanceOf(IllegalStateException.class)
