fix: harden auto tune scheduler throttling

Junjie

2026-04-27 e4e91b46d0ce781e7dc87dcdf0d2909b01911d4b

fix: harden auto tune scheduler throttling

 src/main/java/com/zy/ai/service/impl/AutoTuneAgentServiceImpl.java

@@ -2,6 +2,7 @@

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.zy.ai.domain.autotune.AutoTuneTriggerType;
import com.zy.ai.entity.AiPromptTemplate;
import com.zy.ai.entity.ChatCompletionRequest;
import com.zy.ai.entity.ChatCompletionResponse;
@@ -73,7 +74,7 @@
                }

                for (ChatCompletionRequest.ToolCall toolCall : toolCalls) {
                    Object toolOutput = callMountedTool(toolCall, runState);
                    Object toolOutput = callMountedTool(toolCall, runState, normalizedTriggerType);
                    messages.add(buildToolMessage(toolCall, toolOutput));
                }
            }
@@ -116,12 +117,13 @@
        return message;
    }

    private Object callMountedTool(ChatCompletionRequest.ToolCall toolCall, RunState runState) {
    private Object callMountedTool(ChatCompletionRequest.ToolCall toolCall, RunState runState, String triggerType) {
        String toolName = resolveToolName(toolCall);
        if (!ALLOWED_TOOL_NAMES.contains(toolName)) {
            throw new IllegalArgumentException("Disallowed auto-tune MCP tool: " + toolName);
        }
        JSONObject arguments = parseArguments(toolCall);
        applySchedulerTriggerType(toolName, triggerType, arguments);
        try {
            Object output = mcpToolManager.callTool(toolName, arguments);
            runState.markToolSuccess(toolName);
@@ -132,6 +134,16 @@
        }
    }

    private void applySchedulerTriggerType(String toolName, String triggerType, JSONObject arguments) {
        if (!TOOL_APPLY_CHANGES.equals(toolName)) {
            return;
        }
        if (!AutoTuneTriggerType.AUTO.getCode().equals(triggerType)) {
            return;
        }
        arguments.put("triggerType", AutoTuneTriggerType.AUTO.getCode());
    }

    private ChatCompletionRequest.Message buildToolMessage(ChatCompletionRequest.ToolCall toolCall, Object toolOutput) {
        ChatCompletionRequest.Message toolMessage = new ChatCompletionRequest.Message();
        toolMessage.setRole("tool");

 src/main/java/com/zy/ai/service/impl/AutoTuneCoordinatorServiceImpl.java

@@ -82,15 +82,15 @@
        }

        AutoTuneAgentService.AutoTuneAgentResult agentResult = null;
        markLastTriggerGuard(intervalMinutes);
        try {
            agentResult = autoTuneAgentService.runAutoTune(AutoTuneTriggerType.AUTO.getCode());
            writeOperateLog(agentResult);
            markNoChangeGuardIfNeeded(latestSuccessfulJob, agentResult, intervalMinutes);
            safeWriteOperateLog(agentResult);
            return AutoTuneCoordinatorResult.triggered(agentResult);
        } catch (Exception exception) {
            log.error("Auto tune coordinator failed to run agent", exception);
            agentResult = failedAgentResult(exception);
            writeOperateLog(agentResult);
            safeWriteOperateLog(agentResult);
            return AutoTuneCoordinatorResult.triggered(agentResult);
        } finally {
            redisUtil.compareAndDelete(lockKey, lockToken);
@@ -159,26 +159,10 @@
        return System.currentTimeMillis() - latestFinishTime.getTime() >= intervalMillis;
    }

    private void markNoChangeGuardIfNeeded(AiAutoTuneJob beforeJob,
                                           AutoTuneAgentService.AutoTuneAgentResult agentResult,
                                           int intervalMinutes) {
        if (agentResult == null || !Boolean.TRUE.equals(agentResult.getSuccess())) {
            return;
        }
        AiAutoTuneJob afterJob = latestSuccessfulAutoJob();
        if (!isSameJob(beforeJob, afterJob)) {
            return;
        }
    private void markLastTriggerGuard(int intervalMinutes) {
        long expireSeconds = intervalMinutes * 60L;
        redisUtil.set(RedisKeyType.AI_AUTO_TUNE_LAST_TRIGGER_GUARD.key,
                String.valueOf(System.currentTimeMillis()), expireSeconds);
    }

    private boolean isSameJob(AiAutoTuneJob beforeJob, AiAutoTuneJob afterJob) {
        if (beforeJob == null || afterJob == null) {
            return beforeJob == afterJob;
        }
        return beforeJob.getId() != null && beforeJob.getId().equals(afterJob.getId());
    }

    private AutoTuneAgentService.AutoTuneAgentResult failedAgentResult(Exception exception) {
@@ -195,6 +179,14 @@
        return result;
    }

    private void safeWriteOperateLog(AutoTuneAgentService.AutoTuneAgentResult agentResult) {
        try {
            writeOperateLog(agentResult);
        } catch (Exception exception) {
            log.warn("Auto tune coordinator failed to write operate log", exception);
        }
    }

    private void writeOperateLog(AutoTuneAgentService.AutoTuneAgentResult agentResult) {
        if (agentResult == null) {
            return;

 src/test/java/com/zy/ai/service/AutoTuneCoordinatorServiceImplTest.java

@@ -49,6 +49,8 @@
import static org.mockito.ArgumentMatchers.anyInt;
import static org.mockito.ArgumentMatchers.anyLong;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.doThrow;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
@@ -308,7 +310,7 @@
        when(configService.getConfigValue("aiAutoTuneIntervalMinutes", "10")).thenReturn("10");
        when(wrkMastService.count(any(Wrapper.class))).thenReturn(1L);
        when(redisUtil.get(RedisKeyType.AI_AUTO_TUNE_LAST_TRIGGER_GUARD.key)).thenReturn(null);
        when(aiAutoTuneJobService.list(any(Wrapper.class))).thenReturn(Collections.emptyList(), Collections.emptyList());
        when(aiAutoTuneJobService.list(any(Wrapper.class))).thenReturn(Collections.emptyList());
        when(redisUtil.trySetStringIfAbsent(anyString(), anyString(), anyLong())).thenReturn(true);
        when(autoTuneAgentService.runAutoTune(AutoTuneTriggerType.AUTO.getCode())).thenReturn(agentResult);

@@ -320,6 +322,62 @@
        verify(autoTuneAgentService).runAutoTune(AutoTuneTriggerType.AUTO.getCode());
        verify(operateLogService).save(any());
        verify(redisUtil).set(anyString(), any(), anyLong());
        verify(redisUtil).compareAndDelete(anyString(), anyString());
    }

    @Test
    void coordinatorKeepsAgentResultWhenOperateLogFails() {
        AutoTuneAgentService.AutoTuneAgentResult agentResult = successfulAgentResult();
        when(configService.getConfigValue("aiAutoTuneEnabled", "N")).thenReturn("Y");
        when(configService.getConfigValue("aiAutoTuneIntervalMinutes", "10")).thenReturn("10");
        when(wrkMastService.count(any(Wrapper.class))).thenReturn(1L);
        when(redisUtil.get(RedisKeyType.AI_AUTO_TUNE_LAST_TRIGGER_GUARD.key)).thenReturn(null);
        when(aiAutoTuneJobService.list(any(Wrapper.class))).thenReturn(Collections.emptyList());
        when(redisUtil.trySetStringIfAbsent(anyString(), anyString(), anyLong())).thenReturn(true);
        when(autoTuneAgentService.runAutoTune(AutoTuneTriggerType.AUTO.getCode())).thenReturn(agentResult);
        doThrow(new RuntimeException("log failed")).when(operateLogService).save(any());

        AutoTuneCoordinatorService.AutoTuneCoordinatorResult result = coordinatorService().runAutoTuneIfEligible();

        assertFalse(result.getSkipped());
        assertTrue(result.getTriggered());
        assertSame(agentResult, result.getAgentResult());
        verify(redisUtil).compareAndDelete(anyString(), anyString());
    }

    @Test
    void coordinatorSetsGuardWhenAgentReturnsFailure() {
        AutoTuneAgentService.AutoTuneAgentResult agentResult = failedAgentResult();
        when(configService.getConfigValue("aiAutoTuneEnabled", "N")).thenReturn("Y");
        when(configService.getConfigValue("aiAutoTuneIntervalMinutes", "10")).thenReturn("10");
        when(wrkMastService.count(any(Wrapper.class))).thenReturn(1L);
        when(redisUtil.get(RedisKeyType.AI_AUTO_TUNE_LAST_TRIGGER_GUARD.key)).thenReturn(null);
        when(aiAutoTuneJobService.list(any(Wrapper.class))).thenReturn(Collections.emptyList());
        when(redisUtil.trySetStringIfAbsent(anyString(), anyString(), anyLong())).thenReturn(true);
        when(autoTuneAgentService.runAutoTune(AutoTuneTriggerType.AUTO.getCode())).thenReturn(agentResult);

        AutoTuneCoordinatorService.AutoTuneCoordinatorResult result = coordinatorService().runAutoTuneIfEligible();

        assertFalse(result.getSkipped());
        assertSame(agentResult, result.getAgentResult());
        verify(redisUtil).set(eq(RedisKeyType.AI_AUTO_TUNE_LAST_TRIGGER_GUARD.key), any(), eq(600L));
    }

    @Test
    void coordinatorSetsGuardWhenAgentThrows() {
        when(configService.getConfigValue("aiAutoTuneEnabled", "N")).thenReturn("Y");
        when(configService.getConfigValue("aiAutoTuneIntervalMinutes", "10")).thenReturn("10");
        when(wrkMastService.count(any(Wrapper.class))).thenReturn(1L);
        when(redisUtil.get(RedisKeyType.AI_AUTO_TUNE_LAST_TRIGGER_GUARD.key)).thenReturn(null);
        when(aiAutoTuneJobService.list(any(Wrapper.class))).thenReturn(Collections.emptyList());
        when(redisUtil.trySetStringIfAbsent(anyString(), anyString(), anyLong())).thenReturn(true);
        when(autoTuneAgentService.runAutoTune(AutoTuneTriggerType.AUTO.getCode())).thenThrow(new RuntimeException("agent failed"));

        AutoTuneCoordinatorService.AutoTuneCoordinatorResult result = coordinatorService().runAutoTuneIfEligible();

        assertFalse(result.getSkipped());
        assertFalse(result.getAgentResult().getSuccess());
        verify(redisUtil).set(eq(RedisKeyType.AI_AUTO_TUNE_LAST_TRIGGER_GUARD.key), any(), eq(600L));
        verify(redisUtil).compareAndDelete(anyString(), anyString());
    }

@@ -402,6 +460,28 @@
        assertTrue(visibleToolNames.contains("wcs_local_dispatch_apply_auto_tune_changes"));
        assertTrue(visibleToolNames.contains("wcs_local_dispatch_revert_last_auto_tune_job"));
        assertFalse(visibleToolNames.contains("wcs_local_device_get_crn_status"));
    }

    @Test
    void agentForcesAutoTriggerTypeOnApplyTools() {
        AutoTuneAgentServiceImpl service = agentService();
        when(mcpToolManager.buildOpenAiTools()).thenReturn(allowedOpenAiTools());
        when(mcpToolManager.callTool(any(), any(JSONObject.class))).thenReturn(Collections.singletonMap("ok", true));
        when(llmChatService.chatCompletion(any(), anyDouble(), anyInt(), any()))
                .thenReturn(
                        response("snapshot", toolCall("call_1", "wcs_local_dispatch_get_auto_tune_snapshot",
                                "{}"), 10, 5),
                        response("dry run", toolCall("call_2", "wcs_local_dispatch_apply_auto_tune_changes",
                                "{\"dryRun\":true,\"changes\":[]}"), 10, 5),
                        response("done", null, 10, 5)
                );

        AutoTuneAgentService.AutoTuneAgentResult result = service.runAutoTune("auto");

        assertTrue(result.getSuccess());
        ArgumentCaptor<JSONObject> argumentCaptor = ArgumentCaptor.forClass(JSONObject.class);
        verify(mcpToolManager).callTool(eq("wcs_local_dispatch_apply_auto_tune_changes"), argumentCaptor.capture());
        assertEquals("auto", argumentCaptor.getValue().getString("triggerType"));
    }

    @Test
@@ -496,6 +576,13 @@
        return result;
    }

    private AutoTuneAgentService.AutoTuneAgentResult failedAgentResult() {
        AutoTuneAgentService.AutoTuneAgentResult result = successfulAgentResult();
        result.setSuccess(false);
        result.setSummary("failed");
        return result;
    }

    private AutoTuneChangeCommand change(String targetType, String targetId, String targetKey, String newValue) {
        AutoTuneChangeCommand command = new AutoTuneChangeCommand();
        command.setTargetType(targetType);