Skip to content

Commit 84aa70a

Browse files
hengyunabchengyunabc
authored
Fix MCP automatic polling timeout (#3205)
1 parent fc55581 commit 84aa70a

4 files changed

Lines changed: 164 additions & 13 deletions

File tree

‎arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/task/ServerTaskToolHandler.java‎

Lines changed: 47 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,8 @@ public class ServerTaskToolHandler extends AbstractTaskHandler<McpSchema.ServerT
3535

3636
private static final Logger logger = LoggerFactory.getLogger(ServerTaskToolHandler.class);
3737

38+
private static final Duration AUTOMATIC_POLLING_GRACE = Duration.ofSeconds(10);
39+
3840
private final ObjectMapper objectMapper;
3941
private final TaskManagerOptions taskOptions;
4042
private final Duration automaticPollingTimeout;
@@ -305,7 +307,7 @@ private CompletableFuture<McpSchema.CallToolResult> handleAutomaticTaskPolling(
305307
String taskId = task.getTaskId();
306308
String sessionId = extractSessionId(exchange);
307309

308-
return pollTaskUntilTerminal(taskId, sessionId, task, taskTool);
310+
return pollTaskUntilTerminal(taskId, sessionId, task, taskTool, request);
309311
});
310312
}
311313

@@ -314,15 +316,17 @@ private CompletableFuture<McpSchema.CallToolResult> pollTaskUntilTerminal(
314316
String taskId,
315317
String sessionId,
316318
McpSchema.Task initialTask,
317-
TaskAwareToolSpecification taskTool) {
319+
TaskAwareToolSpecification taskTool,
320+
McpSchema.CallToolRequest request) {
318321

319322
long pollInterval = initialTask.getPollInterval() != null
320323
? initialTask.getPollInterval()
321324
: TaskDefaults.DEFAULT_POLL_INTERVAL_MS;
322325

323-
Duration timeout = this.automaticPollingTimeout != null
324-
? this.automaticPollingTimeout
326+
Duration defaultTimeout = this.automaticPollingTimeout != null
327+
? this.automaticPollingTimeout
325328
: Duration.ofMillis(TaskDefaults.DEFAULT_AUTOMATIC_POLLING_TIMEOUT_MS);
329+
Duration timeout = resolveAutomaticPollingTimeout(request, defaultTimeout);
326330

327331
CompletableFuture<List<McpSchema.Task>> watchFuture = taskStore.watchTaskUntilTerminal(
328332
taskId,
@@ -392,6 +396,45 @@ private CompletableFuture<McpSchema.CallToolResult> pollTaskUntilTerminal(
392396
});
393397
}
394398

399+
Duration resolveAutomaticPollingTimeout(
400+
McpSchema.CallToolRequest request,
401+
Duration defaultTimeout) {
402+
Duration base = defaultTimeout != null
403+
? defaultTimeout
404+
: Duration.ofMillis(TaskDefaults.DEFAULT_AUTOMATIC_POLLING_TIMEOUT_MS);
405+
try {
406+
Long timeoutSeconds = readPositiveLongArgument(request, "timeout");
407+
Duration commandTimeout = timeoutSeconds != null
408+
? Duration.ofSeconds(timeoutSeconds)
409+
: base;
410+
Duration derived = commandTimeout.plus(AUTOMATIC_POLLING_GRACE);
411+
return base.compareTo(derived) >= 0 ? base : derived;
412+
} catch (RuntimeException e) {
413+
logger.warn("Failed to resolve automatic polling timeout, fallback to default with grace", e);
414+
return base.plus(AUTOMATIC_POLLING_GRACE);
415+
}
416+
}
417+
418+
private Long readPositiveLongArgument(McpSchema.CallToolRequest request, String name) {
419+
if (request == null || request.getArguments() == null) {
420+
return null;
421+
}
422+
Object value = request.getArguments().get(name);
423+
if (value instanceof Number) {
424+
long n = ((Number) value).longValue();
425+
return n > 0 ? n : null;
426+
}
427+
if (value instanceof String) {
428+
try {
429+
long n = Long.parseLong(((String) value).trim());
430+
return n > 0 ? n : null;
431+
} catch (NumberFormatException ignored) {
432+
return null;
433+
}
434+
}
435+
return null;
436+
}
437+
395438
private String extractSessionId(McpNettyServerExchange exchange) {
396439
return exchange.sessionId();
397440
}

‎arthas-mcp-server/src/test/java/com/taobao/arthas/mcp/server/task/ServerTaskToolHandlerTest.java‎

Lines changed: 106 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,106 @@
1+
/*
2+
* Copyright 2024-2024 the original author or authors.
3+
*/
4+
5+
package com.taobao.arthas.mcp.server.task;
6+
7+
import com.fasterxml.jackson.databind.ObjectMapper;
8+
import com.taobao.arthas.mcp.server.protocol.spec.McpSchema;
9+
import org.junit.jupiter.api.Test;
10+
11+
import java.time.Duration;
12+
import java.util.Collections;
13+
import java.util.HashMap;
14+
import java.util.Map;
15+
import java.util.concurrent.CompletableFuture;
16+
17+
import static org.assertj.core.api.Assertions.assertThat;
18+
19+
class ServerTaskToolHandlerTest {
20+
21+
@Test
22+
void automaticPollingTimeoutShouldAddGraceWhenTimeoutArgumentIsMissing() {
23+
ServerTaskToolHandler handler = newHandler(Duration.ofSeconds(30));
24+
25+
Duration timeout = handler.resolveAutomaticPollingTimeout(
26+
new McpSchema.CallToolRequest("watch", null, null),
27+
Duration.ofSeconds(30));
28+
29+
assertThat(timeout).isEqualTo(Duration.ofSeconds(40));
30+
}
31+
32+
@Test
33+
void automaticPollingTimeoutShouldCoverLongCommandTimeout() {
34+
ServerTaskToolHandler handler = newHandler(Duration.ofSeconds(30));
35+
36+
Duration timeout = handler.resolveAutomaticPollingTimeout(
37+
requestWithTimeout(45),
38+
Duration.ofSeconds(30));
39+
40+
assertThat(timeout).isEqualTo(Duration.ofSeconds(55));
41+
}
42+
43+
@Test
44+
void automaticPollingTimeoutShouldNotShrinkBelowDefault() {
45+
ServerTaskToolHandler handler = newHandler(Duration.ofSeconds(30));
46+
47+
Duration timeout = handler.resolveAutomaticPollingTimeout(
48+
requestWithTimeout(5),
49+
Duration.ofSeconds(30));
50+
51+
assertThat(timeout).isEqualTo(Duration.ofSeconds(30));
52+
}
53+
54+
@Test
55+
void automaticPollingTimeoutShouldParseNumericStringTimeout() {
56+
ServerTaskToolHandler handler = newHandler(Duration.ofSeconds(30));
57+
58+
Duration timeout = handler.resolveAutomaticPollingTimeout(
59+
requestWithTimeout("45"),
60+
Duration.ofSeconds(30));
61+
62+
assertThat(timeout).isEqualTo(Duration.ofSeconds(55));
63+
}
64+
65+
@Test
66+
void automaticPollingTimeoutShouldFallbackWithGraceForInvalidTimeoutArguments() {
67+
ServerTaskToolHandler handler = newHandler(Duration.ofSeconds(30));
68+
Object[] invalidValues = new Object[] { "bad", "0", "-1", 0, -1, new Object() };
69+
70+
for (Object invalidValue : invalidValues) {
71+
Duration timeout = handler.resolveAutomaticPollingTimeout(
72+
requestWithTimeout(invalidValue),
73+
Duration.ofSeconds(30));
74+
75+
assertThat(timeout).isEqualTo(Duration.ofSeconds(40));
76+
}
77+
}
78+
79+
@Test
80+
void automaticPollingTimeoutShouldFallbackToTaskDefaultWhenDefaultIsNull() {
81+
ServerTaskToolHandler handler = newHandler(null);
82+
83+
Duration timeout = handler.resolveAutomaticPollingTimeout(
84+
new McpSchema.CallToolRequest("watch", null, null),
85+
null);
86+
87+
assertThat(timeout).isEqualTo(Duration.ofMillis(TaskDefaults.DEFAULT_AUTOMATIC_POLLING_TIMEOUT_MS)
88+
.plus(Duration.ofSeconds(10)));
89+
}
90+
91+
private static ServerTaskToolHandler newHandler(Duration automaticPollingTimeout) {
92+
return new ServerTaskToolHandler(
93+
Collections.emptyList(),
94+
null,
95+
new ObjectMapper(),
96+
(method, payload) -> CompletableFuture.completedFuture(null),
97+
automaticPollingTimeout,
98+
null);
99+
}
100+
101+
private static McpSchema.CallToolRequest requestWithTimeout(Object timeout) {
102+
Map<String, Object> arguments = new HashMap<String, Object>();
103+
arguments.put("timeout", timeout);
104+
return new McpSchema.CallToolRequest("watch", arguments, null);
105+
}
106+
}

‎site/docs/.vuepress/theme/components/SidebarItem.vue‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,7 @@ function isActiveSidebarItem(sidebarItem, currentRoute) {
4747
4848
if (sidebarItem.children) {
4949
return sidebarItem.children.some((child) =>
50-
isActiveSidebarItem(child, currentRoute)
50+
isActiveSidebarItem(child, currentRoute),
5151
);
5252
}
5353
@@ -58,10 +58,10 @@ const hasChildren = computed(() => Boolean(item.value.children?.length));
5858
const usesTreeRow = computed(() => Boolean(item.value.link) && depth.value > 0);
5959
const isActive = computed(() => isActiveSidebarItem(item.value, route));
6060
const isCollapsible = computed(
61-
() => Boolean(item.value.collapsible) && hasChildren.value
61+
() => Boolean(item.value.collapsible) && hasChildren.value,
6262
);
6363
const isOpenDefault = computed(() =>
64-
isCollapsible.value ? forceOpen.value || isActive.value : true
64+
isCollapsible.value ? forceOpen.value || isActive.value : true,
6565
);
6666
const isOpen = ref(isOpenDefault.value);
6767
@@ -181,7 +181,9 @@ onBeforeUnmount(() => {
181181
margin: 0.125rem 0.75rem;
182182
border-left: 0.25rem solid transparent;
183183
border-radius: 6px;
184-
transition: background-color var(--t-color), border-color var(--t-color);
184+
transition:
185+
background-color var(--t-color),
186+
border-color var(--t-color);
185187
186188
&.active,
187189
&:hover {

‎site/docs/.vuepress/theme/components/SidebarItems.vue‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -14,10 +14,10 @@ const normalizedQuery = computed(() => normalizeText(searchQuery.value.trim()));
1414
const isSearchActive = computed(() => normalizedQuery.value.length > 0);
1515
const searchPlaceholder = computed(() => (isEnglish.value ? "Search" : "搜索"));
1616
const clearSearchLabel = computed(() =>
17-
isEnglish.value ? "Clear search" : "清除搜索"
17+
isEnglish.value ? "Clear search" : "清除搜索",
1818
);
1919
const noResultsText = computed(() =>
20-
isEnglish.value ? "No matches" : "没有匹配结果"
20+
isEnglish.value ? "No matches" : "没有匹配结果",
2121
);
2222
2323
function normalizeText(value) {
@@ -26,7 +26,7 @@ function normalizeText(value) {
2626
2727
function sidebarItemMatches(item, query) {
2828
return [item.text, item.link].some((value) =>
29-
normalizeText(value).includes(query)
29+
normalizeText(value).includes(query),
3030
);
3131
}
3232
@@ -67,7 +67,7 @@ onMounted(() => {
6767
if (!sidebar) return;
6868
6969
const activeSidebarItem = document.querySelector(
70-
`.sidebar a.sidebar-item[href="${route.path}${hash}"]`
70+
`.sidebar a.sidebar-item[href="${route.path}${hash}"]`,
7171
);
7272
if (!activeSidebarItem) return;
7373
@@ -84,7 +84,7 @@ onMounted(() => {
8484
) {
8585
activeSidebarItem.scrollIntoView(false);
8686
}
87-
}
87+
},
8888
);
8989
});
9090
</script>

0 commit comments

Comments
 (0)