Junjie
2 天以前 63b01db83d9aad8a15276b4236a9a22e4aeef065
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;
@@ -7,19 +8,29 @@
import com.zy.ai.domain.autotune.AutoTuneSnapshot;
import com.zy.ai.entity.AiAutoTuneChange;
import com.zy.ai.entity.AiAutoTuneJob;
import com.zy.ai.entity.AiAutoTuneMcpCall;
import com.zy.ai.enums.AiPromptScene;
import com.zy.ai.service.AiAutoTuneChangeService;
import com.zy.ai.service.AiAutoTuneJobService;
import com.zy.ai.service.AiAutoTuneMcpCallService;
import com.zy.ai.service.AutoTuneApplyService;
import com.zy.ai.service.AutoTuneSnapshotService;
import com.zy.ai.utils.AutoTuneWriteBehaviorUtils;
import lombok.RequiredArgsConstructor;
import org.springframework.ai.tool.annotation.Tool;
import org.springframework.ai.tool.annotation.ToolParam;
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;
import java.util.function.LongSupplier;
@Component
@RequiredArgsConstructor
@@ -27,13 +38,17 @@
    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 AiAutoTuneMcpCallService aiAutoTuneMcpCallService;
    private final ConcurrentMap<String, DryRunPreview> dryRunPreviews = new ConcurrentHashMap<>();
    private LongSupplier currentTimeMillisSupplier = System::currentTimeMillis;
    @Tool(name = "dispatch_get_auto_tune_snapshot", description = "获取WCS自动调参所需的调度快照、站点运行态、拓扑容量和当前可写参数")
    @Tool(name = "dispatch_get_auto_tune_snapshot", description = "获取WCS自动调参所需的调度快照、站点运行态、拓扑容量、当前可写参数和调参规则约束")
    public AutoTuneSnapshot getAutoTuneSnapshot() {
        return autoTuneSnapshotService.buildSnapshot();
    }
@@ -43,6 +58,7 @@
            @ToolParam(description = "返回任务数量上限,默认5,最大20", required = false) Integer limit) {
        int safeLimit = boundLimit(limit);
        List<AiAutoTuneJob> jobs = aiAutoTuneJobService.list(new QueryWrapper<AiAutoTuneJob>()
                .eq("prompt_scene_code", AiPromptScene.AUTO_TUNE_DISPATCH.getCode())
                .orderByDesc("start_time")
                .orderByDesc("id")
                .last("limit " + safeLimit));
@@ -63,14 +79,29 @@
            @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 成功后返回的预览令牌。dryRun=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.");
        }
        AutoTuneApplyRequest request = new AutoTuneApplyRequest();
        request.setReason(reason);
        request.setAnalysisIntervalMinutes(analysisIntervalMinutes);
        request.setTriggerType(triggerType);
        request.setDryRun(dryRun);
        request.setChanges(changes);
        return autoTuneApplyService.apply(request);
        String fingerprint = buildChangeFingerprint(changes);
        if (Boolean.FALSE.equals(dryRun)) {
            requireMatchingDryRunToken(dryRunToken, fingerprint);
        }
        AutoTuneApplyResult result = autoTuneApplyService.apply(request);
        if (Boolean.TRUE.equals(dryRun) && hasApplicableDryRunChanges(result)) {
            result.setDryRunToken(createDryRunToken(fingerprint));
        }
        return result;
    }
    @Tool(name = "dispatch_revert_last_auto_tune_job", description = "回滚最近一次成功的自动调参任务")
@@ -90,16 +121,70 @@
        item.put("successCount", job.getSuccessCount());
        item.put("rejectCount", job.getRejectCount());
        item.put("errorMessage", job.getErrorMessage());
        item.put("changes", listChangeSummaries(job.getId()));
        List<AiAutoTuneMcpCall> mcpCalls = listMcpCalls(job.getId());
        List<Map<String, Object>> mcpCallSummaries = toMcpCallSummaries(mcpCalls);
        List<Map<String, Object>> changeSummaries = listChangeSummaries(job, mcpCalls);
        AutoTuneWriteBehaviorUtils.addWriteBehavior(item,
                AutoTuneWriteBehaviorUtils.resolveJobWriteBehavior(job, mcpCallSummaries, changeSummaries));
        item.put("mcpCallCount", mcpCalls.size());
        item.put("mcpCalls", mcpCallSummaries);
        item.put("changes", changeSummaries);
        return item;
    }
    private List<Map<String, Object>> listChangeSummaries(Long jobId) {
        if (jobId == null) {
    private List<AiAutoTuneMcpCall> listMcpCalls(Long agentJobId) {
        if (agentJobId == null) {
            return new ArrayList<>();
        }
        List<AiAutoTuneMcpCall> mcpCalls = aiAutoTuneMcpCallService.list(new QueryWrapper<AiAutoTuneMcpCall>()
                .eq("agent_job_id", agentJobId)
                .orderByAsc("call_seq")
                .orderByAsc("id"));
        return mcpCalls == null ? new ArrayList<>() : mcpCalls;
    }
    private List<Map<String, Object>> toMcpCallSummaries(List<AiAutoTuneMcpCall> mcpCalls) {
        List<Map<String, Object>> result = new ArrayList<>();
        if (mcpCalls == null || mcpCalls.isEmpty()) {
            return result;
        }
        for (AiAutoTuneMcpCall mcpCall : mcpCalls) {
            result.add(toMcpCallSummary(mcpCall));
        }
        return result;
    }
    private Map<String, Object> toMcpCallSummary(AiAutoTuneMcpCall mcpCall) {
        LinkedHashMap<String, Object> item = new LinkedHashMap<>();
        item.put("callSeq", mcpCall.getCallSeq());
        item.put("toolName", mcpCall.getToolName());
        item.put("status", mcpCall.getStatus());
        item.put("dryRun", toBoolean(mcpCall.getDryRun()));
        item.put("applyJobId", mcpCall.getApplyJobId());
        item.put("successCount", mcpCall.getSuccessCount());
        item.put("rejectCount", mcpCall.getRejectCount());
        item.put("errorMessage", mcpCall.getErrorMessage());
        AutoTuneWriteBehaviorUtils.addWriteBehavior(item,
                AutoTuneWriteBehaviorUtils.resolveMcpWriteBehavior(mcpCall));
        return item;
    }
    private Boolean toBoolean(Integer value) {
        if (value == null) {
            return null;
        }
        return value == 1;
    }
    private List<Map<String, Object>> listChangeSummaries(AiAutoTuneJob job, List<AiAutoTuneMcpCall> mcpCalls) {
        Map<Long, String> ownerTriggerTypes = collectChangeOwnerTriggerTypes(job, mcpCalls);
        List<Long> applyJobIds = new ArrayList<>(ownerTriggerTypes.keySet());
        if (applyJobIds.isEmpty()) {
            return new ArrayList<>();
        }
        List<AiAutoTuneChange> changes = aiAutoTuneChangeService.list(new QueryWrapper<AiAutoTuneChange>()
                .eq("job_id", jobId)
                .in("job_id", applyJobIds)
                .orderByAsc("job_id")
                .orderByAsc("id"));
        if (changes == null || changes.isEmpty()) {
            return new ArrayList<>();
@@ -107,13 +192,40 @@
        List<Map<String, Object>> result = new ArrayList<>();
        for (AiAutoTuneChange change : changes) {
            result.add(toChangeSummary(change));
            result.add(toChangeSummary(change, ownerTriggerTypes.get(change.getJobId())));
        }
        return result;
    }
    private Map<String, Object> toChangeSummary(AiAutoTuneChange change) {
    private Map<Long, String> collectChangeOwnerTriggerTypes(AiAutoTuneJob job, List<AiAutoTuneMcpCall> mcpCalls) {
        LinkedHashMap<Long, String> result = new LinkedHashMap<>();
        if (job != null && job.getId() != null) {
            result.put(job.getId(), job.getTriggerType());
        }
        if (mcpCalls == null || mcpCalls.isEmpty()) {
            return result;
        }
        for (AiAutoTuneMcpCall mcpCall : mcpCalls) {
            Long applyJobId = mcpCall.getApplyJobId();
            if (applyJobId == null || result.containsKey(applyJobId)) {
                continue;
            }
            result.put(applyJobId, resolveMcpApplyJobTriggerType(mcpCall));
        }
        return result;
    }
    private String resolveMcpApplyJobTriggerType(AiAutoTuneMcpCall mcpCall) {
        if (mcpCall == null || mcpCall.getToolName() == null) {
            return null;
        }
        String toolName = mcpCall.getToolName().toLowerCase(Locale.ROOT);
        return toolName.contains("revert_last_auto_tune_job") || toolName.contains("rollback") ? "rollback" : null;
    }
    private Map<String, Object> toChangeSummary(AiAutoTuneChange change, String ownerTriggerType) {
        LinkedHashMap<String, Object> item = new LinkedHashMap<>();
        item.put("jobId", change.getJobId());
        item.put("targetType", change.getTargetType());
        item.put("targetId", change.getTargetId());
        item.put("targetKey", change.getTargetKey());
@@ -124,6 +236,8 @@
        item.put("rejectReason", change.getRejectReason());
        item.put("cooldownExpireTime", change.getCooldownExpireTime());
        item.put("createTime", change.getCreateTime());
        AutoTuneWriteBehaviorUtils.addWriteBehavior(item,
                AutoTuneWriteBehaviorUtils.resolveChangeWriteBehavior(change, ownerTriggerType));
        return item;
    }
@@ -133,4 +247,139 @@
        }
        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(currentTimeMillis())) {
            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, currentTimeMillis() + DRY_RUN_TOKEN_TTL_MILLIS));
        return token;
    }
    private void cleanExpiredDryRunPreviews() {
        long currentTimeMillis = currentTimeMillis();
        for (Map.Entry<String, DryRunPreview> entry : dryRunPreviews.entrySet()) {
            if (entry.getValue() == null || entry.getValue().isExpired(currentTimeMillis)) {
                dryRunPreviews.remove(entry.getKey());
            }
        }
    }
    private boolean isSuccessful(AutoTuneApplyResult result) {
        return result != null && Boolean.TRUE.equals(result.getSuccess());
    }
    private boolean hasApplicableDryRunChanges(AutoTuneApplyResult result) {
        if (!isSuccessful(result) || result.getChanges() == null || result.getChanges().isEmpty()) {
            return false;
        }
        for (AiAutoTuneChange change : result.getChanges()) {
            if (change != null && "dry_run".equals(normalizeLower(change.getResultStatus()))) {
                return true;
            }
        }
        return false;
    }
    private String buildChangeFingerprint(List<AutoTuneChangeCommand> changes) {
        List<Map<String, String>> normalizedChanges = new ArrayList<>();
        if (changes != null) {
            for (AutoTuneChangeCommand change : changes) {
                normalizedChanges.add(toNormalizedChange(change));
            }
        }
        validateUniqueChangeTargets(normalizedChanges);
        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 void validateUniqueChangeTargets(List<Map<String, String>> normalizedChanges) {
        Map<String, Map<String, String>> uniqueTargets = new LinkedHashMap<>();
        for (Map<String, String> change : normalizedChanges) {
            String targetSignature = buildTargetSignature(change);
            if (uniqueTargets.containsKey(targetSignature)) {
                throw new IllegalArgumentException("Duplicate auto-tune change target in same request: "
                        + "targetType=" + change.get("targetType")
                        + ", targetId=" + change.get("targetId")
                        + ", targetKey=" + change.get("targetKey"));
            }
            uniqueTargets.put(targetSignature, change);
        }
    }
    private String buildTargetSignature(Map<String, String> change) {
        return change.get("targetType") + "\n"
                + change.get("targetId") + "\n"
                + change.get("targetKey");
    }
    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 long currentTimeMillis() {
        return currentTimeMillisSupplier.getAsLong();
    }
    void setCurrentTimeMillisSupplier(LongSupplier currentTimeMillisSupplier) {
        this.currentTimeMillisSupplier = currentTimeMillisSupplier == null
                ? System::currentTimeMillis
                : currentTimeMillisSupplier;
    }
    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(long currentTimeMillis) {
            return currentTimeMillis > expireAtMillis;
        }
    }
}