From e4e91b46d0ce781e7dc87dcdf0d2909b01911d4b Mon Sep 17 00:00:00 2001
From: Junjie <fallin.jie@qq.com>
Date: 星期一, 27 四月 2026 12:34:31 +0800
Subject: [PATCH] fix: harden auto tune scheduler throttling
---
src/main/java/com/zy/ai/service/impl/AutoTuneAgentServiceImpl.java | 16 ++++
src/main/java/com/zy/ai/service/impl/AutoTuneCoordinatorServiceImpl.java | 32 ++++------
src/test/java/com/zy/ai/service/AutoTuneCoordinatorServiceImplTest.java | 89 +++++++++++++++++++++++++++++
3 files changed, 114 insertions(+), 23 deletions(-)
diff --git a/src/main/java/com/zy/ai/service/impl/AutoTuneAgentServiceImpl.java b/src/main/java/com/zy/ai/service/impl/AutoTuneAgentServiceImpl.java
index 7af6821..fe0ca9b 100644
--- a/src/main/java/com/zy/ai/service/impl/AutoTuneAgentServiceImpl.java
+++ b/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");
diff --git a/src/main/java/com/zy/ai/service/impl/AutoTuneCoordinatorServiceImpl.java b/src/main/java/com/zy/ai/service/impl/AutoTuneCoordinatorServiceImpl.java
index 808895a..bbafe96 100644
--- a/src/main/java/com/zy/ai/service/impl/AutoTuneCoordinatorServiceImpl.java
+++ b/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;
diff --git a/src/test/java/com/zy/ai/service/AutoTuneCoordinatorServiceImplTest.java b/src/test/java/com/zy/ai/service/AutoTuneCoordinatorServiceImplTest.java
index 70bd28d..2f880c2 100644
--- a/src/test/java/com/zy/ai/service/AutoTuneCoordinatorServiceImplTest.java
+++ b/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);
--
Gitblit v1.9.1