feat(mcp-task): cap concurrent task sessions (#3182)

arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/session/ArthasCommandSessionManager.java

@@ -1,6 +1,9 @@
 package com.taobao.arthas.mcp.server.session;
 
 import com.taobao.arthas.mcp.server.CommandExecutor;
+import com.taobao.arthas.mcp.server.protocol.spec.McpError;
+import com.taobao.arthas.mcp.server.protocol.spec.McpSchema;
+import com.taobao.arthas.mcp.server.task.TaskDefaults;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -14,19 +17,25 @@ import java.util.concurrent.ConcurrentHashMap;
  */
 public class ArthasCommandSessionManager {
     private static final Logger logger = LoggerFactory.getLogger(ArthasCommandSessionManager.class);
-    
+
     // Arthas 默认 session 超时时间是 30 分钟，这里设置一个稍短的时间作为预判断
     // 如果距离上次访问超过这个时间，认为 session 可能已过期，主动重建
     private static final long SESSION_EXPIRY_THRESHOLD_MS = 25 * 60 * 1000; // 25 分钟
-    
+
     private final CommandExecutor commandExecutor;
+    private final int maxConcurrentTaskSessions;
     private final ConcurrentHashMap<String, CommandSessionBinding> sessionBindings = new ConcurrentHashMap<>();
-    
+
     // 独立管理 task session
     private final ConcurrentHashMap<String, CommandSessionBinding> taskSessionBindings = new ConcurrentHashMap<>();
 
     public ArthasCommandSessionManager(CommandExecutor commandExecutor) {
+        this(commandExecutor, TaskDefaults.DEFAULT_MAX_CONCURRENT_TASK_SESSIONS);
+    }
+
+    public ArthasCommandSessionManager(CommandExecutor commandExecutor, int maxConcurrentTaskSessions) {
         this.commandExecutor = commandExecutor;
+        this.maxConcurrentTaskSessions = maxConcurrentTaskSessions;
     }
 
     public static class CommandSessionBinding {
@@ -161,20 +170,28 @@ public class ArthasCommandSessionManager {
      * 为 task 创建独立的 Arthas Session。
      */
     public CommandSessionBinding createIsolatedTaskSession(String taskId) {
-        Map<String, Object> result = commandExecutor.createSession();
-        
-        CommandSessionBinding binding = new CommandSessionBinding(
-            "task-" + taskId,  // 使用 task ID 作为 MCP session ID
-            (String) result.get("sessionId"),
-            (String) result.get("consumerId")
-        );
-        
-        // 注册到独立的 map，方便追踪和清理
-        taskSessionBindings.put(taskId, binding);
-        
-        logger.info("Created isolated task session: taskId={}, arthasSessionId={}", 
-                   taskId, binding.getArthasSessionId());
-        return binding;
+        synchronized (taskSessionBindings) {
+            if (taskSessionBindings.size() >= maxConcurrentTaskSessions) {
+                throw McpError.builder(McpSchema.ErrorCodes.INVALID_PARAMS)
+                    .message("concurrent task session limit reached: " + maxConcurrentTaskSessions)
+                    .build();
+            }
+
+            Map<String, Object> result = commandExecutor.createSession();
+
+            CommandSessionBinding binding = new CommandSessionBinding(
+                "task-" + taskId,  // 使用 task ID 作为 MCP session ID
+                (String) result.get("sessionId"),
+                (String) result.get("consumerId")
+            );
+
+            // 注册到独立的 map，方便追踪和清理
+            taskSessionBindings.put(taskId, binding);
+
+            logger.info("Created isolated task session: taskId={}, arthasSessionId={}",
+                       taskId, binding.getArthasSessionId());
+            return binding;
+        }
     }
 
     public void closeTaskSession(String taskId) {
@@ -194,4 +211,12 @@ public class ArthasCommandSessionManager {
     public int getActiveTaskSessionCount() {
         return taskSessionBindings.size();
     }
+
+    /**
+     * 检查当前活跃 task session 数是否已达并发上限。
+     * 供外层在创建 Task 前进行前置判断，避免产生孤儿 Task。
+     */
+    public boolean isAtConcurrencyLimit() {
+        return taskSessionBindings.size() >= maxConcurrentTaskSessions;
+    }
 }

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

@@ -46,4 +46,6 @@ public interface CreateTaskContext {
     ArthasCommandContext createIsolatedTaskSession(String taskId);
 
     void cleanupTaskSession(String taskId);
+
+    boolean isAtConcurrencyLimit();
 }

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

@@ -143,4 +143,12 @@ public class DefaultCreateTaskContext implements CreateTaskContext {
             sessionManager.closeTaskSession(taskId);
         }
     }
+
+    @Override
+    public boolean isAtConcurrencyLimit() {
+        if (sessionManager == null) {
+            return false;
+        }
+        return sessionManager.isAtConcurrencyLimit();
+    }
 }

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

@@ -44,12 +44,16 @@ public class InMemoryTaskStore<R extends McpSchema.Result> implements TaskStore<
 
     private final Set<String> cancellationRequests = ConcurrentHashMap.newKeySet();
 
+    private final Map<String, Long> expiredCancellationDeadlines = new ConcurrentHashMap<>();
+
     private final ScheduledExecutorService cleanupExecutor;
 
     private final long defaultTtl;
 
     private final long defaultPollInterval;
 
+    private final long expiredCancellationRetentionMs;
+
     private static final AtomicLong INSTANCE_COUNTER = new AtomicLong(0);
 
     private final long instanceId;
@@ -58,6 +62,8 @@ public class InMemoryTaskStore<R extends McpSchema.Result> implements TaskStore<
 
     private final int maxTasks;
 
+    private static final long DEFAULT_EXPIRED_CANCELLATION_RETENTION_MS = 30 * 60 * 1000L;
+
     public InMemoryTaskStore() {
         this(DEFAULT_TTL_MS, DEFAULT_POLL_INTERVAL_MS, null, DEFAULT_MAX_TASKS);
     }
@@ -72,15 +78,26 @@ public class InMemoryTaskStore<R extends McpSchema.Result> implements TaskStore<
 
     public InMemoryTaskStore(long defaultTtl, long defaultPollInterval,
                              TaskMessageQueue messageQueue, int maxTasks) {
+        this(defaultTtl, defaultPollInterval, messageQueue, maxTasks,
+                DEFAULT_EXPIRED_CANCELLATION_RETENTION_MS);
+    }
+
+    InMemoryTaskStore(long defaultTtl, long defaultPollInterval,
+                      TaskMessageQueue messageQueue, int maxTasks,
+                      long expiredCancellationRetentionMs) {
         if (defaultTtl <= 0) throw new IllegalArgumentException("defaultTtl must be positive");
         if (defaultPollInterval <= 0) throw new IllegalArgumentException("defaultPollInterval must be positive");
         if (maxTasks <= 0) throw new IllegalArgumentException("maxTasks must be positive");
+        if (expiredCancellationRetentionMs <= 0) {
+            throw new IllegalArgumentException("expiredCancellationRetentionMs must be positive");
+        }
 
         this.instanceId = INSTANCE_COUNTER.incrementAndGet();
         this.defaultTtl = defaultTtl;
         this.defaultPollInterval = defaultPollInterval;
         this.messageQueue = messageQueue;
         this.maxTasks = maxTasks;
+        this.expiredCancellationRetentionMs = expiredCancellationRetentionMs;
 
         this.cleanupExecutor = Executors.newSingleThreadScheduledExecutor(r -> {
             Thread t = new Thread(r, "mcp-task-cleanup-" + instanceId);
@@ -373,7 +390,10 @@ public class InMemoryTaskStore<R extends McpSchema.Result> implements TaskStore<
     public CompletableFuture<Boolean> isCancellationRequested(String taskId, String sessionId) {
         return CompletableFuture.supplyAsync(() -> {
             TaskEntry entry = tasks.get(taskId);
-            if (entry == null || !isSessionValid(entry, sessionId)) return false;
+            // Bug fix: entry == null means TTL cleanup already removed the task.
+            // We must NOT suppress the cancellation signal — if cancellationRequests still
+            // contains the taskId, the background thread must see it and exit cleanly.
+            if (entry != null && !isSessionValid(entry, sessionId)) return false;
             return cancellationRequests.contains(taskId);
         });
     }
@@ -452,8 +472,11 @@ public class InMemoryTaskStore<R extends McpSchema.Result> implements TaskStore<
     /** Package-visible for testing. */
     void cleanupExpiredTasks() {
         Instant now = Instant.now();
+        long nowMillis = now.toEpochMilli();
         List<String> expiredTaskIds = new ArrayList<>();
 
+        cleanupExpiredCancellationRequests(nowMillis);
+
         tasks.entrySet().removeIf(entry -> {
             McpSchema.Task task = entry.getValue().task();
             if (task.getTtl() == null) return false;
@@ -461,7 +484,17 @@ public class InMemoryTaskStore<R extends McpSchema.Result> implements TaskStore<
             if (now.isAfter(expiresAt)) {
                 String taskId = entry.getKey();
                 results.remove(taskId);
-                cancellationRequests.remove(taskId);
+                if (!TaskHelper.isTerminal(task.getStatus())
+                        || task.getStatus() == McpSchema.TaskStatus.CANCELLED) {
+                    // 对仍可能有后台执行线程的任务，移除 task 记录前先保留协作取消信号。
+                    // CANCELLED 对客户端是终态，但不代表后台 worker 已退出。
+                    // 保留取消信号一段时间，避免后台线程尚未观察到信号时被提前清理。
+                    markExpiredTaskCancellation(taskId, nowMillis);
+                    logger.debug("Retained cancellation signal for expired active task: {}", taskId);
+                } else {
+                    // COMPLETED/FAILED 已不需要协作取消信号，可以完整清理。
+                    clearCancellationRequest(taskId);
+                }
                 expiredTaskIds.add(taskId);
                 logger.debug("Removed expired task: {}", taskId);
                 return true;
@@ -480,6 +513,27 @@ public class InMemoryTaskStore<R extends McpSchema.Result> implements TaskStore<
         }
     }
 
+    private void markExpiredTaskCancellation(String taskId, long nowMillis) {
+        cancellationRequests.add(taskId);
+        expiredCancellationDeadlines.put(taskId, nowMillis + expiredCancellationRetentionMs);
+    }
+
+    private void clearCancellationRequest(String taskId) {
+        cancellationRequests.remove(taskId);
+        expiredCancellationDeadlines.remove(taskId);
+    }
+
+    private void cleanupExpiredCancellationRequests(long nowMillis) {
+        expiredCancellationDeadlines.entrySet().removeIf(entry -> {
+            if (nowMillis < entry.getValue()) {
+                return false;
+            }
+            cancellationRequests.remove(entry.getKey());
+            logger.debug("Removed expired cancellation signal: {}", entry.getKey());
+            return true;
+        });
+    }
+
     @Override
     public CompletableFuture<Void> shutdown() {
         return CompletableFuture.runAsync(() -> {

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

@@ -15,14 +15,24 @@ import java.time.Duration;
  */
 public final class TaskDefaults {
 
-    public static final long DEFAULT_TTL_MS = 60_000L;
+    public static final long DEFAULT_TTL_MS = 10 * 60 * 1000L;
 
     public static final long DEFAULT_POLL_INTERVAL_MS = 1000L;
 
     public static final int DEFAULT_PAGE_SIZE = 100;
 
+    /**
+     * task 记录存储容量上限（内存兜底）。
+     * 统计的是 tasks Map 中所有状态的 entry 总数（包含已完成但未被 TTL 清理的），
+     */
     public static final int DEFAULT_MAX_TASKS = 10_000;
 
+    /**
+     * 并发 task session 上限（资源保护）。
+     * 统计的是正在执行中的 Arthas session 数量，保护目标 JVM 不被过多 session 拖垮。
+     */
+    public static final int DEFAULT_MAX_CONCURRENT_TASK_SESSIONS = 10;
+
     public static final long DEFAULT_AUTOMATIC_POLLING_TIMEOUT_MS = 600000L;
 
     public static final int DEFAULT_SIDE_CHANNEL_TIMEOUT_MINUTES = 5;

arthas-mcp-server/src/main/java/com/taobao/arthas/mcp/server/tool/ToolCallbackCreateTaskHandler.java

@@ -1,7 +1,8 @@
 package com.taobao.arthas.mcp.server.tool;
 
-import com.taobao.arthas.mcp.server.protocol.server.McpTransportContext;
+import com.taobao.arthas.mcp.server.protocol.spec.McpError;
 import com.taobao.arthas.mcp.server.protocol.spec.McpSchema;
+import com.taobao.arthas.mcp.server.protocol.server.McpTransportContext;
 import com.taobao.arthas.mcp.server.session.ArthasCommandContext;
 import com.taobao.arthas.mcp.server.task.CreateTaskContext;
 import com.taobao.arthas.mcp.server.task.CreateTaskHandler;
@@ -12,6 +13,8 @@ import org.slf4j.LoggerFactory;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.Executor;
+import java.util.concurrent.RejectedExecutionException;
 
 import static com.taobao.arthas.mcp.server.tool.ToolContextKeys.*;
 
@@ -26,27 +29,62 @@ public class ToolCallbackCreateTaskHandler implements CreateTaskHandler {
     
     private final ToolCallback toolCallback;
 
-    public ToolCallbackCreateTaskHandler(ToolCallback toolCallback) {
+    // 专用 executor，避免 I/O 密集型任务污染 ForkJoinPool.commonPool
+    private final Executor taskExecutor;
+
+    public ToolCallbackCreateTaskHandler(ToolCallback toolCallback, Executor taskExecutor) {
         this.toolCallback = toolCallback;
+        this.taskExecutor = taskExecutor;
     }
 
     @Override
     public CompletableFuture<McpSchema.CreateTaskResult> createTask(
             Map<String, Object> args,
             CreateTaskContext context) {
-        
+
         logger.debug("Creating task for tool: {}", toolCallback.getToolDefinition().getName());
 
+        // 前置检查：在创建 Task 之前判断并发限制，避免产生孤儿 Task。
+        // 此时请求已经被应程层的并发限制，客户端可立刻感知并重试。
+        if (context.isAtConcurrencyLimit()) {
+            logger.warn("Concurrent task session limit reached, rejecting tool: {}",
+                    toolCallback.getToolDefinition().getName());
+            CompletableFuture<McpSchema.CreateTaskResult> rejected = new CompletableFuture<>();
+            rejected.completeExceptionally(
+                    McpError.builder(McpSchema.ErrorCodes.INVALID_PARAMS)
+                            .message("concurrent task session limit reached")
+                            .data("Server is at max concurrent task capacity. Retry after an existing task completes.")
+                            .build()
+            );
+            return rejected;
+        }
+
         return context.createTask(opts -> {
             // 使用默认配置，工具可以通过注解自定义 pollInterval 等
         }).thenCompose(task -> {
             String taskId = task.getTaskId();
-            
+
             logger.info("Task created: {}, starting async tool execution", taskId);
 
-            CompletableFuture.runAsync(() -> {
-                executeToolAndUpdateTaskStatus(taskId, args, context);
-            });
+            // 将 Task 提交到专用线程池（SynchronousQueue）。
+            // 若 executor 拒绝（理论上不应发生，前置检查已拦截），
+            // 作为安全底存捕获并标记任务失败，避免孤児 Task。
+            try {
+                CompletableFuture.runAsync(() -> {
+                    executeToolAndUpdateTaskStatus(taskId, args, context);
+                }, taskExecutor);
+            } catch (RejectedExecutionException e) {
+                logger.error("Task executor rejected task: {} (should not happen after pre-check)", taskId, e);
+                context.failTask(taskId, new McpSchema.CallToolResult(
+                        "Task rejected: executor at capacity", true, null))
+                        .exceptionally(ex -> {
+                            logger.error("Failed to mark rejected task as failed: {}", taskId, ex);
+                            return null;
+                        });
+                // 返回已创建但将就失败的 Task，让客户端能感知
+                return CompletableFuture.completedFuture(
+                        new McpSchema.CreateTaskResult(task, null));
+            }
 
             return CompletableFuture.completedFuture(
                 new McpSchema.CreateTaskResult(task, null)

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

@@ -0,0 +1,72 @@
+/*
+ * Copyright 2024-2024 the original author or authors.
+ */
+
+package com.taobao.arthas.mcp.server.task;
+
+import com.taobao.arthas.mcp.server.protocol.spec.McpSchema;
+import org.junit.jupiter.api.Test;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+class InMemoryTaskStoreTest {
+
+    @Test
+    void cleanupExpiredNonTerminalTaskCancellationSignalAfterRetention() throws Exception {
+        InMemoryTaskStore<McpSchema.CallToolResult> store = new InMemoryTaskStore<>(
+                1L, 1000L, null, 100, 1L);
+        String taskId = "expired-working-task";
+        String sessionId = "session-1";
+
+        try {
+            store.createTask(CreateTaskOptions.builder()
+                    .taskId(taskId)
+                    .sessionId(sessionId)
+                    .ttl(1L)
+                    .build()).join();
+
+            Thread.sleep(20L);
+            store.cleanupExpiredTasks();
+
+            assertThat(store.getTask(taskId, sessionId).join()).isNull();
+            assertThat(store.isCancellationRequested(taskId, sessionId).join()).isTrue();
+
+            Thread.sleep(20L);
+            store.cleanupExpiredTasks();
+
+            assertThat(store.isCancellationRequested(taskId, sessionId).join()).isFalse();
+        } finally {
+            store.shutdown().join();
+        }
+    }
+
+    @Test
+    void cleanupExpiredCancelledTaskRetainsCancellationSignalAfterTtl() throws Exception {
+        InMemoryTaskStore<McpSchema.CallToolResult> store = new InMemoryTaskStore<>(
+                1L, 1000L, null, 100, 1L);
+        String taskId = "expired-cancelled-task";
+        String sessionId = "session-1";
+
+        try {
+            store.createTask(CreateTaskOptions.builder()
+                    .taskId(taskId)
+                    .sessionId(sessionId)
+                    .ttl(1L)
+                    .build()).join();
+            store.requestCancellation(taskId, sessionId).join();
+
+            Thread.sleep(20L);
+            store.cleanupExpiredTasks();
+
+            assertThat(store.getTask(taskId, sessionId).join()).isNull();
+            assertThat(store.isCancellationRequested(taskId, sessionId).join()).isTrue();
+
+            Thread.sleep(20L);
+            store.cleanupExpiredTasks();
+
+            assertThat(store.isCancellationRequested(taskId, sessionId).join()).isFalse();
+        } finally {
+            store.shutdown().join();
+        }
+    }
+}

core/src/main/java/com/taobao/arthas/core/mcp/ArthasMcpServer.java

@@ -34,6 +34,12 @@ import org.slf4j.LoggerFactory;
 
 import java.time.Duration;
 import java.util.*;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.SynchronousQueue;
+import java.util.concurrent.ThreadFactory;
+import java.util.concurrent.ThreadPoolExecutor;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.stream.Collectors;
 
 /**
@@ -63,6 +69,10 @@ public class ArthasMcpServer {
 
     private McpStatelessHttpRequestHandler statelessHandler;
 
+    // MCP Task 专用线程池，大小与 maxConcurrentTaskSessions 对齐，
+    // 避免 I/O 密集型任务污染 ForkJoinPool.commonPool
+    private ExecutorService taskExecutor;
+
     public static final String DEFAULT_MCP_ENDPOINT = "/mcp";
     
     public ArthasMcpServer(String mcpEndpoint, CommandExecutor commandExecutor, String protocol) {
@@ -176,6 +186,18 @@ public class ArthasMcpServer {
         // 初始化 SessionManager
         this.sessionManager = new ArthasCommandSessionManager(commandExecutor);
         logger.info("Initialized ArthasCommandSessionManager for MCP server");
+
+        int maxSessions = com.taobao.arthas.mcp.server.task.TaskDefaults.DEFAULT_MAX_CONCURRENT_TASK_SESSIONS;
+        this.taskExecutor = new ThreadPoolExecutor(
+                maxSessions,                       // corePoolSize: 与 session 上限一致
+                maxSessions,                       // maxPoolSize: 固定大小，避免动态扩缩
+                0L, TimeUnit.MILLISECONDS,
+                new SynchronousQueue<>(),           // 无缓冲：直接交付或拒绝，无隐式排队
+                new McpTaskThreadFactory(),         // 具名线程，便于 thread dump 诊断
+                new ThreadPoolExecutor.AbortPolicy() // 兜底：超出直接抛 RejectedExecutionException
+        );
+        logger.info("Created MCP task executor: fixedSize={}, queue=SynchronousQueue (no buffering)",
+                maxSessions);
         
         McpStreamableServerTransportProvider transportProvider = createStreamableHttpTransportProvider(properties);
         streamableHandler = transportProvider.getMcpRequestHandler();
@@ -229,7 +251,7 @@ public class ArthasMcpServer {
         for (ToolCallback callback : taskAwareTools) {
             ToolDefinition def = callback.getToolDefinition();
 
-            ToolCallbackCreateTaskHandler createTaskHandler = new ToolCallbackCreateTaskHandler(callback);
+            ToolCallbackCreateTaskHandler createTaskHandler = new ToolCallbackCreateTaskHandler(callback, taskExecutor);
 
             TaskAwareToolSpecification spec = TaskAwareToolSpecification.builder()
                     .name(def.getName())
@@ -369,10 +391,32 @@ public class ArthasMcpServer {
                 statelessServer.closeGracefully().get();
                 logger.info("Stateless server stopped successfully");
             }
-            
+
+            if (taskExecutor != null) {
+                taskExecutor.shutdown();
+                if (!taskExecutor.awaitTermination(5, TimeUnit.SECONDS)) {
+                    taskExecutor.shutdownNow();
+                    logger.warn("MCP task executor did not terminate within 5s, forced shutdown");
+                } else {
+                    logger.info("MCP task executor stopped successfully");
+                }
+            }
+
             logger.info("Arthas MCP server stopped completely");
         } catch (Exception e) {
             logger.error("Failed to stop Arthas MCP server gracefully", e);
         }
     }
+
+
+    private static class McpTaskThreadFactory implements ThreadFactory {
+        private final AtomicInteger counter = new AtomicInteger(0);
+
+        @Override
+        public Thread newThread(Runnable r) {
+            Thread t = new Thread(r, "mcp-task-" + counter.getAndIncrement());
+            t.setDaemon(true);
+            return t;
+        }
+    }
 }