From 10bdc4b6e9701befd1a83bccd2998dcc96cb2c43 Mon Sep 17 00:00:00 2001
From: Junjie <fallin.jie@qq.com>
Date: 星期一, 27 四月 2026 12:07:38 +0800
Subject: [PATCH] fix: enforce auto tune agent tool safety
---
src/main/java/com/zy/ai/mcp/tool/AutoTuneMcpTools.java | 114 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++
1 files changed, 113 insertions(+), 1 deletions(-)
diff --git a/src/main/java/com/zy/ai/mcp/tool/AutoTuneMcpTools.java b/src/main/java/com/zy/ai/mcp/tool/AutoTuneMcpTools.java
index db0e6d0..e822cf3 100644
--- a/src/main/java/com/zy/ai/mcp/tool/AutoTuneMcpTools.java
+++ b/src/main/java/com/zy/ai/mcp/tool/AutoTuneMcpTools.java
@@ -1,5 +1,6 @@
package com.zy.ai.mcp.tool;
+import com.alibaba.fastjson.JSON;
import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import com.zy.ai.domain.autotune.AutoTuneApplyRequest;
import com.zy.ai.domain.autotune.AutoTuneApplyResult;
@@ -17,9 +18,14 @@
import org.springframework.stereotype.Component;
import java.util.ArrayList;
+import java.util.Comparator;
import java.util.LinkedHashMap;
import java.util.List;
+import java.util.Locale;
import java.util.Map;
+import java.util.UUID;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
@Component
@RequiredArgsConstructor
@@ -27,11 +33,13 @@
private static final int DEFAULT_RECENT_JOB_LIMIT = 5;
private static final int MAX_RECENT_JOB_LIMIT = 20;
+ private static final long DRY_RUN_TOKEN_TTL_MILLIS = 10L * 60L * 1000L;
private final AutoTuneSnapshotService autoTuneSnapshotService;
private final AutoTuneApplyService autoTuneApplyService;
private final AiAutoTuneJobService aiAutoTuneJobService;
private final AiAutoTuneChangeService aiAutoTuneChangeService;
+ private final ConcurrentMap<String, DryRunPreview> dryRunPreviews = new ConcurrentHashMap<>();
@Tool(name = "dispatch_get_auto_tune_snapshot", description = "鑾峰彇WCS鑷姩璋冨弬鎵�闇�鐨勮皟搴﹀揩鐓с�佺珯鐐硅繍琛屾�併�佹嫇鎵戝閲忓拰褰撳墠鍙啓鍙傛暟")
public AutoTuneSnapshot getAutoTuneSnapshot() {
@@ -63,14 +71,27 @@
@ToolParam(description = "寤鸿鑷姩璋冨弬鍒嗘瀽闂撮殧鍒嗛挓", required = false) Integer analysisIntervalMinutes,
@ToolParam(description = "瑙﹀彂绫诲瀷锛屼緥濡� scheduler/manual/agent", required = false) String triggerType,
@ToolParam(description = "鏄惁浠呰瘯绠楋紝瀹為檯搴旂敤鍓嶅繀椤诲厛浼� true", required = false) Boolean dryRun,
+ @ToolParam(description = "dry-run 鎴愬姛鍚庤繑鍥炵殑棰勮浠ょ墝銆俤ryRun=false 鏃跺繀椤绘彁渚涳紝涓斿彉鏇撮泦蹇呴』瀹屽叏涓�鑷�", required = false) String dryRunToken,
@ToolParam(description = "璋冨弬鍙樻洿鍒楄〃") List<AutoTuneChangeCommand> changes) {
+ if (dryRun == null) {
+ throw new IllegalArgumentException("dryRun is required. Use dryRun=true first to create a preview token.");
+ }
+ String fingerprint = buildChangeFingerprint(changes);
+ if (Boolean.FALSE.equals(dryRun)) {
+ requireMatchingDryRunToken(dryRunToken, fingerprint);
+ }
+
AutoTuneApplyRequest request = new AutoTuneApplyRequest();
request.setReason(reason);
request.setAnalysisIntervalMinutes(analysisIntervalMinutes);
request.setTriggerType(triggerType);
request.setDryRun(dryRun);
request.setChanges(changes);
- return autoTuneApplyService.apply(request);
+ AutoTuneApplyResult result = autoTuneApplyService.apply(request);
+ if (Boolean.TRUE.equals(dryRun) && isSuccessful(result)) {
+ result.setDryRunToken(createDryRunToken(fingerprint));
+ }
+ return result;
}
@Tool(name = "dispatch_revert_last_auto_tune_job", description = "鍥炴粴鏈�杩戜竴娆℃垚鍔熺殑鑷姩璋冨弬浠诲姟")
@@ -133,4 +154,95 @@
}
return Math.min(limit, MAX_RECENT_JOB_LIMIT);
}
+
+ private void requireMatchingDryRunToken(String dryRunToken, String fingerprint) {
+ cleanExpiredDryRunPreviews();
+ if (isBlank(dryRunToken)) {
+ throw new IllegalArgumentException("dryRunToken is required when dryRun=false. Run dryRun=true first.");
+ }
+ DryRunPreview preview = dryRunPreviews.remove(dryRunToken.trim());
+ if (preview == null) {
+ throw new IllegalArgumentException("dryRunToken is missing, expired, or already used.");
+ }
+ if (preview.isExpired()) {
+ throw new IllegalArgumentException("dryRunToken is expired. Run dryRun=true again.");
+ }
+ if (!preview.getFingerprint().equals(fingerprint)) {
+ throw new IllegalArgumentException("dryRunToken does not match the requested change set.");
+ }
+ }
+
+ private String createDryRunToken(String fingerprint) {
+ cleanExpiredDryRunPreviews();
+ String token = UUID.randomUUID().toString();
+ dryRunPreviews.put(token, new DryRunPreview(fingerprint, System.currentTimeMillis() + DRY_RUN_TOKEN_TTL_MILLIS));
+ return token;
+ }
+
+ private void cleanExpiredDryRunPreviews() {
+ for (Map.Entry<String, DryRunPreview> entry : dryRunPreviews.entrySet()) {
+ if (entry.getValue() == null || entry.getValue().isExpired()) {
+ dryRunPreviews.remove(entry.getKey());
+ }
+ }
+ }
+
+ private boolean isSuccessful(AutoTuneApplyResult result) {
+ return result != null && Boolean.TRUE.equals(result.getSuccess());
+ }
+
+ private String buildChangeFingerprint(List<AutoTuneChangeCommand> changes) {
+ List<Map<String, String>> normalizedChanges = new ArrayList<>();
+ if (changes != null) {
+ for (AutoTuneChangeCommand change : changes) {
+ normalizedChanges.add(toNormalizedChange(change));
+ }
+ }
+ normalizedChanges.sort(Comparator
+ .comparing((Map<String, String> item) -> item.get("targetType"))
+ .thenComparing(item -> item.get("targetId"))
+ .thenComparing(item -> item.get("targetKey"))
+ .thenComparing(item -> item.get("newValue")));
+ return JSON.toJSONString(normalizedChanges);
+ }
+
+ private Map<String, String> toNormalizedChange(AutoTuneChangeCommand change) {
+ LinkedHashMap<String, String> item = new LinkedHashMap<>();
+ String targetType = normalizeLower(change == null ? null : change.getTargetType());
+ item.put("targetType", targetType);
+ item.put("targetId", "sys_config".equals(targetType) ? "" : normalizeText(change == null ? null : change.getTargetId()));
+ item.put("targetKey", normalizeText(change == null ? null : change.getTargetKey()));
+ item.put("newValue", normalizeText(change == null ? null : change.getNewValue()));
+ return item;
+ }
+
+ private String normalizeLower(String value) {
+ return normalizeText(value).toLowerCase(Locale.ROOT);
+ }
+
+ private String normalizeText(String value) {
+ return value == null ? "" : value.trim();
+ }
+
+ private boolean isBlank(String value) {
+ return value == null || value.trim().isEmpty();
+ }
+
+ private static class DryRunPreview {
+ private final String fingerprint;
+ private final long expireAtMillis;
+
+ DryRunPreview(String fingerprint, long expireAtMillis) {
+ this.fingerprint = fingerprint;
+ this.expireAtMillis = expireAtMillis;
+ }
+
+ String getFingerprint() {
+ return fingerprint;
+ }
+
+ boolean isExpired() {
+ return System.currentTimeMillis() > expireAtMillis;
+ }
+ }
}
--
Gitblit v1.9.1