#

Junjie

9 天以前 dc3f9cc91759823ce59486f19b138be4b296a0f1

 src/main/java/com/zy/ai/mcp/tool/AutoTuneMcpTools.java

@@ -26,6 +26,7 @@
import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.function.LongSupplier;

@Component
@RequiredArgsConstructor
@@ -40,8 +41,9 @@
    private final AiAutoTuneJobService aiAutoTuneJobService;
    private final AiAutoTuneChangeService aiAutoTuneChangeService;
    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();
    }
@@ -164,7 +166,7 @@
        if (preview == null) {
            throw new IllegalArgumentException("dryRunToken is missing, expired, or already used.");
        }
        if (preview.isExpired()) {
        if (preview.isExpired(currentTimeMillis())) {
            throw new IllegalArgumentException("dryRunToken is expired. Run dryRun=true again.");
        }
        if (!preview.getFingerprint().equals(fingerprint)) {
@@ -175,13 +177,14 @@
    private String createDryRunToken(String fingerprint) {
        cleanExpiredDryRunPreviews();
        String token = UUID.randomUUID().toString();
        dryRunPreviews.put(token, new DryRunPreview(fingerprint, System.currentTimeMillis() + DRY_RUN_TOKEN_TTL_MILLIS));
        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()) {
            if (entry.getValue() == null || entry.getValue().isExpired(currentTimeMillis)) {
                dryRunPreviews.remove(entry.getKey());
            }
        }
@@ -228,6 +231,16 @@
        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;
@@ -241,8 +254,8 @@
            return fingerprint;
        }

        boolean isExpired() {
            return System.currentTimeMillis() > expireAtMillis;
        boolean isExpired(long currentTimeMillis) {
            return currentTimeMillis > expireAtMillis;
        }
    }
}