feat(mcp): progress notifications for long-running MCP tools

Wire MCP standard notifications/progress into the existing SSE stream so long-running MCP tool calls surface live progress instead of a bare spinner. A per-call progressToken maps back to (conversationId, toolCallId); ProgressAwareMcpToolCallback injects it into tools/call _meta and calls McpSyncClient directly (falling back to the delegate on error, and applying identity forwarding first). Progress events skip the ring buffer and are replayed from a latest-value snapshot on SSE reconnect. Frontend renders a gradient progress bar in ToolCallSegment when a running tool reports progress.
This commit is contained in:
MIST 2026-07-09 18:05:14 +08:00 committed by GitHub
parent 4ae4731d54
commit e35c07f742
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
18 changed files with 1574 additions and 14 deletions

View File

@ -0,0 +1,520 @@
# MCP 长时任务进度推送落地方案(复用现有 SSE 通道)
## 一、目标
让 MateClaw 支持 MCP 长时任务(最大 360 分钟)的**实时进度展示**,兼容 **SSE 和 streamable_http** 两种 MCP transport全程**零新增 HTTP 接口、零轮询、零中间件**。
---
## 二、背景
### 2.1 当前问题
- MCP 工具调用为同步阻塞模式(`McpSyncClient.callTool()`),最长 60 秒超时
- `tool_call_started``tool_call_completed` 之间前端只显示旋转加载器,用户无感知
- 长时任务(如 Linux 源码编译安装,最长 360 分钟)缺乏进度反馈,用户体验差
### 2.2 为什么不用原有方案文档stderr / 自定义 WS
| 原有方案 | 问题 |
|---------|------|
| stderr 管道 | 绕过 MCP 协议标准;仅 stdio transport 可用AI 侧需额外解析 |
| 自定义 WebSocket | MCP Server 需自建 WS 服务AI 侧需额外建立 WS 连接;非 MCP 标准 |
### 2.3 本方案的核心思路
利用 MCP 协议标准 `notifications/progress` 机制,在 MateClawMCP Client 侧)接收进度通知后,**直接注入现有 SSE 推送通道**传到前端浏览器——数据流完全复用已有基础设施。
---
## 三、架构与数据流
### 3.1 全链路数据流
```
MCP Server任意 transport: SSE / streamable_http / stdio
│ notifications/progress {progressToken, progress, total, message}
McpClientManager.progressConsumer ← 新增注册
│ 根据 progressToken 查表得到 (conversationId, toolCallId)
Spring McpProgressEvent ← 新增事件类型
McpProgressRelay.onMcpProgress() ← 新增监听器
│ 调用 ChatStreamTracker.broadcastObject()
ChatStreamTracker ← 已有,纯内存广播
│ SSE: event=tool_call_progress
│ data={toolCallId, toolName, percent, stage, message}
浏览器 ToolCallSegment.vue ← 已有组件,加进度条渲染
```
### 3.2 progressToken 映射机制
MCP 协议要求 client 生成唯一的 `progressToken``tools/call` 请求发给 serverserver 在 `notifications/progress` 中**原样回传**——这是天然的请求-响应绑定。
```
工具调用前:
progressToken = UUID.randomUUID()
progressTokenMap.put(progressToken, ProgressContext(conversationId, toolCallId, serverId, toolName))
MCP tools/call 请求:
{ name: "linux_source_install", _meta: { progressToken: "xxx-uuid" }, ... }
MCP 服务端推送:
{ method: "notifications/progress", params: { progressToken: "xxx-uuid", progress: 0.5, ... } }
MateClaw 收到:
context = progressTokenMap.get("xxx-uuid")
→ ChatStreamTracker.broadcastObject(context.conversationId, "tool_call_progress", {...})
工具调用完成后:
progressTokenMap.remove("xxx-uuid")
```
### 3.3 360 分钟超长任务处理
`ChatStreamTracker` 的环形缓冲区上限 16000 条事件360 分钟 × 每 2 秒一次 = 10800 条 progress会挤占 content/thinking delta 空间。
**策略**
- progress 事件**不缓存**到 event buffer`skipBuffer = true`
- 维护独立内存快照:`Map<conversationId, Map<toolCallId, ProgressSnapshot>>`
- SSE 重连时不做全量 progress 回放,只下发**一条最新进度快照**
- 快照仅存最新值,内存恒定 O(1) per tool call
---
## 四、涉及文件与改动说明
| # | 文件路径 | 改动类型 | 说明 |
|---|---------|---------|------|
| 1 | `mateclaw-server/.../mcp/runtime/McpClientManager.java` | 修改 | 注册 progressConsumer |
| 2 | `mateclaw-server/.../mcp/runtime/McpProgressContext.java` | 新建 | progressToken 映射表 + 线程安全存取 |
| 3 | `mateclaw-server/.../mcp/runtime/McpProgressEvent.java` | 新建 | Spring Event 定义 |
| 4 | `mateclaw-server/.../mcp/runtime/McpProgressRelay.java` | 新建 | Event Listener → ChatStreamTracker |
| 5 | `mateclaw-server/.../agent/ToolExecutionExecutor.java` | 修改 | 调用前注册映射,完成后清理 |
| 6 | `mateclaw-server/.../mcp/runtime/SyncMcpToolCallbackProvider.java` | 修改 | 向 tools/call 请求注入 progressToken |
| 7 | `mateclaw-server/.../channel/web/ChatStreamTracker.java` | 修改 | 支持 skipBuffer + 重连下发进度快照 |
| 8 | `mateclaw-ui/.../chat/ToolCallSegment.vue` | 修改 | 渲染进度条 |
| 9 | `mateclaw-ui/.../chat/useChat.ts` | 修改 | 监听 tool_call_progress 事件 |
---
## 五、逐文件实现规格
### 5.1 McpProgressEvent.java新建
位置:`mateclaw-server/src/main/java/vip/mate/tool/mcp/runtime/McpProgressEvent.java`
```java
package vip.mate.tool.mcp.runtime;
import org.springframework.context.ApplicationEvent;
import java.util.Map;
/**
* MCP 工具调用进度事件。
* 由 McpClientManager.progressConsumer 发布,
* 由 McpProgressRelay 消费并转发到 ChatStreamTracker。
*/
public class McpProgressEvent extends ApplicationEvent {
private final String conversationId;
private final String toolCallId;
private final String toolName;
private final double progress; // 0.0 ~ 1.0
private final Double total; // 可为 null
private final String message; // 当前阶段描述
public McpProgressEvent(Object source, String conversationId, String toolCallId,
String toolName, double progress, Double total, String message) {
super(source);
this.conversationId = conversationId;
this.toolCallId = toolCallId;
this.toolName = toolName;
this.progress = progress;
this.total = total;
this.message = message;
}
// getters...
}
```
### 5.2 McpProgressContext.java新建
位置:`mateclaw-server/src/main/java/vip/mate/tool/mcp/runtime/McpProgressContext.java`
职责:
- `Map<String, ProgressEntry>` — progressToken → {conversationId, toolCallId, serverId, toolName}
- 线程安全(`ConcurrentHashMap`
- 提供 `register(token, entry)` / `lookup(token)` / `remove(token)`
- 提供 `getLatestSnapshot(conversationId, toolCallId)` — 用于 SSE 重连时下发进度快照
```java
package vip.mate.tool.mcp.runtime;
import org.springframework.stereotype.Component;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
@Component
public class McpProgressContext {
private final Map<String, ProgressEntry> tokenMap = new ConcurrentHashMap<>();
// 进度快照conversationId -> toolCallId -> 最新进度 JSON
private final Map<String, Map<String, String>> snapshotMap = new ConcurrentHashMap<>();
public record ProgressEntry(String conversationId, String toolCallId,
String serverId, String toolName) {}
public void register(String progressToken, ProgressEntry entry) {
tokenMap.put(progressToken, entry);
}
public ProgressEntry lookup(String progressToken) {
return tokenMap.get(progressToken);
}
public void remove(String progressToken) {
tokenMap.remove(progressToken);
}
/** 更新进度快照(每次收到 progress 时调用) */
public void updateSnapshot(String conversationId, String toolCallId, String progressJson) {
snapshotMap.computeIfAbsent(conversationId, k -> new ConcurrentHashMap<>())
.put(toolCallId, progressJson);
}
/** SSE 重连时获取进度快照 */
public String getSnapshot(String conversationId, String toolCallId) {
Map<String, String> tools = snapshotMap.get(conversationId);
return tools != null ? tools.get(toolCallId) : null;
}
/** 工具完成后清理快照 */
public void removeSnapshot(String conversationId, String toolCallId) {
Map<String, String> tools = snapshotMap.get(conversationId);
if (tools != null) {
tools.remove(toolCallId);
}
}
}
```
### 5.3 McpClientManager.java修改
位置:`mateclaw-server/src/main/java/vip/mate/tool/mcp/runtime/McpClientManager.java`
`buildClient()` 方法中(约第 413 行,`spec.toolsChangeConsumer(...)` 之后)新增:
```java
import vip.mate.tool.mcp.runtime.McpProgressContext;
import vip.mate.tool.mcp.runtime.McpProgressEvent;
// 字段注入
private final McpProgressContext progressContext;
// buildClient() 中toolsChangeConsumer 之后:
spec.progressConsumer(progressNotification -> {
if (progressNotification == null || progressNotification.progressToken() == null) return;
McpProgressContext.ProgressEntry entry = progressContext.lookup(progressNotification.progressToken());
if (entry == null) return;
try {
McpProgressEvent event = new McpProgressEvent(
this,
entry.conversationId(),
entry.toolCallId(),
entry.toolName(),
progressNotification.progress(),
progressNotification.total(),
progressNotification.message()
);
eventPublisher.publishEvent(event);
} catch (Exception e) {
log.warn("Failed to publish McpProgressEvent: {}", e.getMessage());
}
});
```
### 5.4 SyncMcpToolCallbackProvider.java修改
位置:`mateclaw-server/src/main/java/vip/mate/tool/mcp/runtime/SyncMcpToolCallbackProvider.java`
> 注意:如果 `SyncMcpToolCallbackProvider` 来自 Spring AI SDK 且无法直接修改,则需要创建一个 **wrapper**`ProgressAwareSyncMcpToolCallback`,在 `call()` 方法中:
> 1. 生成 `progressToken = UUID.randomUUID().toString()`
> 2. 将 `progressToken` 设置到 `CallToolRequest._meta`
> 3. 委托给原始 `SyncMcpToolCallback.call()` 或直接调 `McpSyncClient.callTool(request)`
核心逻辑:
```java
public String call(String toolInput, ToolContext toolContext) {
// 仅对 MCP 工具生效
if (!isMcpTool) return delegate.call(toolInput, toolContext);
String progressToken = UUID.randomUUID().toString();
// 注册映射
progressContext.register(progressToken,
new McpProgressContext.ProgressEntry(conversationId, toolCallId, serverId, toolName));
try {
// 构造带 progressToken 的 CallToolRequest
McpSchema.CallToolRequest request = McpSchema.CallToolRequest.builder()
.name(toolName)
.arguments(arguments)
.meta(Map.of("progressToken", progressToken))
.build();
return mcpSyncClient.callTool(request).content().toString();
} finally {
progressContext.remove(progressToken);
}
}
```
### 5.5 ToolExecutionExecutor.java修改
位置:`mateclaw-server/src/main/java/vip/mate/tool/agent/ToolExecutionExecutor.java`
`executeSingleTool()` 方法中(约第 892 行,`callback.call()` 调用前后):
```java
// 调用前:对于 MCP 工具,注册 progressToken 映射
// (这部分逻辑实际在 SyncMcpToolCallbackProvider wrapper 中完成)
// ToolExecutionExecutor 此处主要负责调用完成后通知清理
```
> 实际需要改动的地方较少——progressToken 的注册和清理已由 wrapper 负责ToolExecutionExecutor 的改动主要是确保 `toolCallId``conversationId` 能传递到 wrapper 中。
### 5.6 McpProgressRelay.java新建
位置:`mateclaw-server/src/main/java/vip/mate/tool/mcp/runtime/McpProgressRelay.java`
```java
package vip.mate.tool.mcp.runtime;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.context.event.EventListener;
import org.springframework.stereotype.Component;
import vip.mate.channel.web.ChatStreamTracker;
import java.util.Map;
@Slf4j
@Component
@RequiredArgsConstructor
public class McpProgressRelay {
private final ChatStreamTracker streamTracker;
private final McpProgressContext progressContext;
private final ObjectMapper objectMapper;
@EventListener
public void onMcpProgress(McpProgressEvent event) {
try {
Map<String, Object> data = Map.of(
"toolCallId", event.getToolCallId(),
"toolName", event.getToolName(),
"percent", Math.round(event.getProgress() * 10000.0) / 100.0, // 保留两位小数
"total", event.getTotal() != null ? event.getTotal() : 1.0,
"message", event.getMessage() != null ? event.getMessage() : "",
"stage", inferStage(event.getProgress()) // 根据百分比推断阶段
);
String jsonData = objectMapper.writeValueAsString(data);
// 更新进度快照(用于重连)
progressContext.updateSnapshot(event.getConversationId(), event.getToolCallId(), jsonData);
// 广播到 SSEskipBuffer = true不缓存到环形缓冲区
streamTracker.broadcast(event.getConversationId(), "tool_call_progress", jsonData, true);
} catch (Exception e) {
log.warn("Failed to relay MCP progress: {}", e.getMessage());
}
}
/** 根据百分比推断阶段名 */
private String inferStage(double progress) {
if (progress <= 0.05) return "prepare";
if (progress <= 0.95) return "execute";
return "finalize";
}
}
```
### 5.7 ChatStreamTracker.java修改
位置:`mateclaw-server/src/main/java/vip/mate/channel/web/ChatStreamTracker.java`
改动 1`broadcast()` 方法新增 `skipBuffer` 参数重载
```java
/**
* 广播事件到所有 SSE 订阅者(可选是否缓存)。
* @param skipBuffer true 时不写入环形缓冲区,用于高频 transient 事件(如 progress
*/
public void broadcast(String conversationId, String eventName, String jsonData, boolean skipBuffer) {
// 现有 broadcast 逻辑 + skipBuffer 判断
}
```
改动 2`attach()` 重连时下发进度快照
```java
// 在 attach() 方法的 buffer 回放完成后:
McpProgressContext progressCtx = springContext.getBean(McpProgressContext.class);
Map<String, String> toolSnapshots = progressCtx.getSnapshots(conversationId);
if (toolSnapshots != null) {
for (Map.Entry<String, String> entry : toolSnapshots.entrySet()) {
sendToEmitter(emitter, "tool_call_progress", entry.getValue());
}
}
```
### 5.8 useChat.ts修改
位置:`mateclaw-ui/src/composables/chat/useChat.ts`
在 SSE 事件处理注册中添加:
```typescript
stream.on('tool_call_progress', (event: SSEEvent) => {
const data = parseSSEData(event.data)
if (!data?.toolCallId) return
const msgIdx = messages.value.findIndex(m =>
m.segments?.some(s => s.toolCallId === data.toolCallId))
if (msgIdx < 0) return
const msg = messages.value[msgIdx]
const segIdx = msg.segments!.findIndex(s => s.toolCallId === data.toolCallId)
if (segIdx < 0) return
// 更新 segment 的 progress 字段
msg.segments![segIdx] = {
...msg.segments![segIdx],
progress: data.percent,
progressMessage: data.message,
progressStage: data.stage
}
})
```
### 5.9 ToolCallSegment.vue修改
位置:`mateclaw-ui/src/components/chat/ToolCallSegment.vue`
在运行状态(`status === 'running'`)时,如果有 progress 数据,渲染进度条替代纯旋转加载器:
```vue
<!-- 运行中 + 有 progress 数据 → 显示进度条 -->
<div v-if="segment.status === 'running' && segment.progress != null" class="progress-bar-wrapper">
<div class="progress-label">{{ segment.progress }}%</div>
<div class="progress-bar">
<div class="progress-fill" :style="{ width: segment.progress + '%' }"></div>
</div>
<div class="progress-message">{{ segment.progressMessage }}</div>
</div>
<!-- 运行中 + 无 progress 数据 → 显示原有旋转加载器 -->
<div v-else-if="segment.status === 'running'" class="loading-spinner">...</div>
```
---
## 六、边界情况处理
| 场景 | 处理方式 |
|------|---------|
| MCP Server 不支持 progress | `progressConsumer` 收不到回调,路径完全不变,前端展示旋转加载器 |
| progressConsumer 内部异常 | try-catch 包围log.warn不传播异常 |
| progressTokenMap 内存泄漏 | `finally` 块保证清理;工具超时后通过定时任务扫描清理超过 400 分钟的陈旧 entry |
| SSE 断开重连5 分钟内) | progress 不参与 buffer 回放attach 后从快照下发最新进度 |
| SSE 断开超过 5 分钟 | RunState 已销毁attach 失败,前端重新发起请求 |
| progress 推送频率过高 | 接收端不节流(交给 MCP Server 侧控制),前端直接渲染,无性能问题 |
| 多个 MCP Server 同时运行 | progressToken 全局唯一UUID不同 server 的 token 不会冲突 |
| conversationId 找不到 | `ChatStreamTracker.broadcast()` 内部 state 为 null 时静默丢弃,不报错 |
---
## 七、验证方法
### 7.1 后端验证
**步骤 1**:启动一个 MCP ServerSSE transport 或 streamable_http实现 `@McpProgressToken` 推送进度。
用 [Spring AI MCP Server Boot Starter](https://docs.spring.io/spring-ai/reference/api/mcp/mcp-server-boot-starter.html) 写一个简单的测试工具:
```java
@McpTool(name = "long_running_test", description = "模拟长时任务")
public String longRunning(@McpProgressToken String progressToken,
McpSyncServerExchange exchange) throws Exception {
for (int i = 0; i <= 10; i++) {
Thread.sleep(2000); // 每 2 秒推进 10%
exchange.progressNotification(p -> p
.progressToken(progressToken)
.progress(i * 0.1)
.total(1.0)
.message("Step " + i + "/10"));
}
return "done";
}
```
**步骤 2**:在 MateClaw 中注册该 MCP Server通过聊天界面触发 `long_running_test` 工具。
**预期结果**
- MateClaw 后端日志输出:`McpProgressRelay` 收到 progress 事件并广播
- `ChatStreamTracker` 广播 `tool_call_progress` 事件(`skipBuffer=true`
### 7.2 前端验证
**步骤 1**:调用 `long_running_test` 后,打开浏览器 DevTools → Network → 找到 `/api/v1/chat/stream` 的 SSE 响应。
**预期结果**
- SSE 流中出现 `event: tool_call_progress` 事件data 包含 `toolCallId`、`percent`、`message`
- `tool_call_started` 之后,`ToolCallSegment` 不再只显示旋转加载器,而是显示进度条和百分比
**步骤 2**:在任务运行过程中,刷新浏览器页面(模拟重连)。
**预期结果**
- SSE 重连成功(`Last-Event-ID` 回放)
- progress 事件不会批量回放(因为 `skipBuffer=true`
- 连接恢复后立即下发一条最新的 progress 快照
- 后续 progress 照常实时推送
### 7.3 兼容性验证
| 测试项 | 方法 | 预期 |
|--------|------|------|
| SSE transport MCP Server | 用 SSE transport 注册 MCP Server触发 progress 工具 | 正常展示进度 |
| streamable_http transport | 用 streamable_http transport 注册(需客户端支持),触发 progress 工具 | 正常展示进度 |
| 非 MCP 工具(内置工具) | 调用 read_file / write_file 等 | 不受影响,仍展示旋转加载器 |
| 无 progress 的 MCP 工具 | 调用不发送 progress 的 MCP 工具 | 不受影响,仍展示旋转加载器 |
| 360 分钟长任务 | 模拟推送 360 分钟的 progress 事件 | 进度持续更新event buffer 未被挤占,内存不增长 |
---
## 八、实施顺序(推荐)
1. **`McpProgressContext.java`** — 先建映射表
2. **`McpProgressEvent.java`** — 事件定义
3. **`McpClientManager.java`** — 注册 progressConsumer
4. **`SyncMcpToolCallbackProvider.java`** — 注入 progressToken
5. **`McpProgressRelay.java`** — 转发到 SSE
6. **`ChatStreamTracker.java`** — skipBuffer + 重连快照
7. **`useChat.ts` + `ToolCallSegment.vue`** — 前端渲染
8. **集成测试** — 用测试 MCP Server 端到端验证
---
## 九、注意事项
- `progressConsumer` 注册在 `McpClientManager.buildClient()` 中,每次 MCP 连接建立/重建时生效
- `progressToken` 的生命周期必须与工具调用严格绑定调用前注册、finally 清理
- progress 事件不走 `StreamAccumulator``GraphEvent` 管道(因为图节点在工具执行期间阻塞),直接由 `McpProgressRelay` 注入 `ChatStreamTracker`
- transport 类型对方案无影响——`progressConsumer` 是 SDK 层面抽象stdio/SSE/streamable_http 均支持

View File

@ -53,6 +53,7 @@ import vip.mate.skill.service.SkillService;
import vip.mate.system.service.SystemSettingService;
import vip.mate.tool.ToolRegistry;
import vip.mate.tool.disclosure.ToolUsageRecencyTracker;
import vip.mate.tool.mcp.runtime.McpProgressContext;
import vip.mate.memory.spi.MemoryManager;
import vip.mate.workspace.document.WorkspaceFileService;
import vip.mate.tool.guard.service.ToolGuardService;
@ -105,6 +106,7 @@ public class AgentGraphBuilder {
private final ModelContextWindowResolver contextWindowResolver;
private final PrefixBudgetPlanner prefixBudgetPlanner;
private final ToolUsageRecencyTracker toolUsageRecencyTracker;
private final McpProgressContext progressContext;
private final vip.mate.llm.service.ModelCapabilityService modelCapabilityService;
private final ProviderRouter providerRouter;
private final PlanningService planningService;
@ -632,6 +634,7 @@ public class AgentGraphBuilder {
// the right invocation pattern instead of a dead-end error.
executor.setSkillRuntimeService(skillRuntimeService);
executor.setUsageRecencyTracker(toolUsageRecencyTracker);
executor.setProgressContext(progressContext);
// Optional: route child-agent denied-tool audit events through
// the audit pipeline. Null when audit is not wired (legacy / test).
if (auditEventService != null) {
@ -928,6 +931,7 @@ public class AgentGraphBuilder {
// the right invocation pattern instead of a dead-end error.
executor.setSkillRuntimeService(skillRuntimeService);
executor.setUsageRecencyTracker(toolUsageRecencyTracker);
executor.setProgressContext(progressContext);
// Optional: route child-agent denied-tool audit events through
// the audit pipeline. Null when audit is not wired (legacy / test).
if (auditEventService != null) {

View File

@ -8,7 +8,9 @@ import org.springframework.ai.chat.model.ToolContext;
import org.springframework.ai.tool.ToolCallback;
import vip.mate.tool.builtin.ToolExecutionContext;
import vip.mate.tool.disclosure.ToolUsageRecencyTracker;
import vip.mate.tool.mcp.runtime.McpProgressContext;
import vip.mate.tool.mcp.runtime.McpToolNameResolver;
import vip.mate.tool.mcp.runtime.ProgressAwareMcpToolCallback;
import vip.mate.agent.AgentToolSet;
import vip.mate.agent.GraphEventPublisher;
import vip.mate.agent.context.ChatOrigin;
@ -256,10 +258,17 @@ public class ToolExecutionExecutor {
/** Optional recency feed for budget-driven tool-disclosure demotion. */
private ToolUsageRecencyTracker usageRecencyTracker;
/** Optional MCP progress context for long-running tool progress relay. */
private McpProgressContext progressContext;
public void setUsageRecencyTracker(ToolUsageRecencyTracker tracker) {
this.usageRecencyTracker = tracker;
}
public void setProgressContext(McpProgressContext ctx) {
this.progressContext = ctx;
}
public void setSkillRuntimeService(vip.mate.skill.runtime.SkillRuntimeService s) {
this.skillRuntimeService = s;
}
@ -883,14 +892,31 @@ public class ToolExecutionExecutor {
// not yet migrated to ToolContext keep working unchanged.
ToolExecutionContext.set(pc.conversationId, pc.requesterId, pc.workspaceBasePath);
String result;
String progressToken = null;
try {
ChatOrigin runtimeOrigin = pc.origin != null ? pc.origin : ChatOrigin.EMPTY;
runtimeOrigin = runtimeOrigin
.withConversationId(pc.conversationId)
.withWorkspace(runtimeOrigin.workspaceId(), pc.workspaceBasePath);
ToolContext toolContext = runtimeOrigin.toToolContext();
// MCP progress: generate progressToken and inject into ToolContext
// so ProgressAwareMcpToolCallback can include it in tools/call _meta.
if (progressContext != null) {
progressToken = UUID.randomUUID().toString();
progressContext.register(progressToken,
new McpProgressContext.ProgressEntry(pc.conversationId, pc.toolCall.id(), toolName));
Map<String, Object> ctxMap = new HashMap<>(toolContext.getContext());
ctxMap.put(ProgressAwareMcpToolCallback.MCP_PROGRESS_TOKEN_KEY, progressToken);
toolContext = new ToolContext(ctxMap);
}
result = pc.callback.call(pc.arguments, toolContext);
} finally {
if (progressToken != null) {
progressContext.remove(progressToken);
progressContext.removeSnapshot(pc.conversationId, pc.toolCall.id());
}
ToolExecutionContext.clear();
}

View File

@ -3,10 +3,13 @@ package vip.mate.channel.web;
import com.fasterxml.jackson.databind.ObjectMapper;
import jakarta.annotation.PreDestroy;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.ApplicationContext;
import org.springframework.stereotype.Component;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import reactor.core.Disposable;
import vip.mate.tool.mcp.runtime.McpProgressContext;
import vip.mate.workspace.conversation.model.MessageContentPart;
import java.io.IOException;
@ -106,6 +109,9 @@ public class ChatStreamTracker {
@Value("${mateclaw.stream.heartbeat.tool-sec:5}")
private int heartbeatToolSec = 5;
@Autowired
private ApplicationContext applicationContext;
public ChatStreamTracker(ObjectMapper objectMapper) {
this.objectMapper = objectMapper;
}
@ -582,6 +588,15 @@ public class ChatStreamTracker {
* early-return remains.
*/
public void broadcast(String conversationId, String eventName, String jsonData) {
broadcast(conversationId, eventName, jsonData, false);
}
/**
* Broadcast an event to all subscribers (optionally skip buffer).
* @param skipBuffer if true, do not write to the ring buffer used for
* high-frequency transient events (e.g. progress).
*/
public void broadcast(String conversationId, String eventName, String jsonData, boolean skipBuffer) {
RunState state = runs.get(conversationId);
boolean isDone = "done".equals(eventName);
@ -652,18 +667,23 @@ public class ChatStreamTracker {
}
synchronized (state.lock) {
if (!skipBuffer) {
long id = ++state.nextEventId;
SseEvent event = new SseEvent(id, eventName, jsonData);
state.buffer.add(event);
// buffer 容量保护超出上限时优先丢弃 thinking_delta占比最大且非关键
if (state.buffer.size() > MAX_BUFFER_SIZE) {
trimBuffer(state.buffer);
}
}
Iterator<SseEmitter> it = state.subscribers.iterator();
while (it.hasNext()) {
SseEmitter emitter = it.next();
try {
emitter.send(SseEmitter.event().id(String.valueOf(id)).name(eventName).data(jsonData));
if (skipBuffer) {
emitter.send(SseEmitter.event().name(eventName).data(jsonData));
} else {
emitter.send(SseEmitter.event().id(String.valueOf(state.nextEventId)).name(eventName).data(jsonData));
}
} catch (IOException | IllegalStateException e) {
log.debug("Removing dead subscriber for {}: {}", conversationId, e.getMessage());
it.remove();
@ -695,6 +715,13 @@ public class ChatStreamTracker {
* @param data 事件载荷将被 Jackson 序列化为 JSON
*/
public void broadcastObject(String conversationId, String eventName, Object data) {
broadcastObject(conversationId, eventName, data, false);
}
/**
* Broadcast an Object directly (auto-serialized to JSON), optionally skipping the buffer.
*/
public void broadcastObject(String conversationId, String eventName, Object data, boolean skipBuffer) {
String json;
try {
json = objectMapper.writeValueAsString(data);
@ -702,7 +729,32 @@ public class ChatStreamTracker {
log.warn("Failed to serialize broadcast data for event {}: {}", eventName, e.getMessage());
json = "{\"error\":\"serialization_failed\"}";
}
broadcast(conversationId, eventName, json);
broadcast(conversationId, eventName, json, skipBuffer);
}
/**
* Deliver MCP progress snapshots on SSE reconnect. Progress events do not
* participate in buffer replay, so the latest snapshot is read from
* {@link McpProgressContext} and delivered separately on attach.
*/
private void sendProgressSnapshots(String conversationId, SseEmitter emitter) {
try {
McpProgressContext progressCtx = applicationContext.getBean(McpProgressContext.class);
Map<String, String> snapshots = progressCtx.getSnapshots(conversationId);
if (snapshots != null && !snapshots.isEmpty()) {
for (Map.Entry<String, String> entry : snapshots.entrySet()) {
try {
emitter.send(SseEmitter.event()
.name("tool_call_progress")
.data(entry.getValue()));
} catch (IOException e) {
log.debug("Failed to send progress snapshot for {}: {}", conversationId, e.getMessage());
}
}
}
} catch (Exception e) {
log.debug("Failed to send progress snapshots for {}: {}", conversationId, e.getMessage());
}
}
/**
@ -921,6 +973,10 @@ public class ChatStreamTracker {
// Without this, async_task_completed fired after `done` would be silently
// dropped, leaving the chat UI stuck on the "正在生成中" placeholder.
state.subscribers.add(emitter);
// Deliver MCP progress snapshots on reconnect (progress events skip buffer replay)
sendProgressSnapshots(conversationId, emitter);
if (state.done) {
log.info("[SSE] Replayed {} buffered events; emitter stays subscribed for late async events: {}",
state.buffer.size(), conversationId);

View File

@ -77,7 +77,12 @@ public final class IdentityForwardingToolCallback implements ToolCallback {
return delegate;
}
private String inject(String toolInput, ToolContext toolContext) {
/**
* Inject identity claim into the toolInput JSON.
* Package-private so {@link ProgressAwareMcpToolCallback} can apply identity
* forwarding before calling mcpClient directly (progress path).
*/
String inject(String toolInput, ToolContext toolContext) {
return identityService.resolve(toolContext, audience)
.map(i -> withClaim(toolInput, i.key(), i.value()))
.orElse(toolInput);

View File

@ -1,6 +1,7 @@
package vip.mate.tool.mcp.runtime;
import cn.hutool.json.JSONUtil;
import com.fasterxml.jackson.databind.ObjectMapper;
import io.modelcontextprotocol.client.McpClient;
import io.modelcontextprotocol.client.McpSyncClient;
import io.modelcontextprotocol.client.transport.ServerParameters;
@ -67,13 +68,21 @@ public class McpClientManager {
private final McpIdentityForwardService identityForwardService;
private final McpProgressContext progressContext;
private final ObjectMapper objectMapper;
/** serverId -> server name, captured at build time for identity-forward opt-in matching. */
private final ConcurrentHashMap<Long, String> serverNames = new ConcurrentHashMap<>();
public McpClientManager(ApplicationEventPublisher eventPublisher,
McpIdentityForwardService identityForwardService) {
McpIdentityForwardService identityForwardService,
McpProgressContext progressContext,
ObjectMapper objectMapper) {
this.eventPublisher = eventPublisher;
this.identityForwardService = identityForwardService;
this.progressContext = progressContext;
this.objectMapper = objectMapper;
}
/** serverId -> connection result info */
@ -201,7 +210,8 @@ public class McpClientManager {
McpIdentityForwardService idSvc =
identityForwardService.forwardsTo(serverId, serverName) ? identityForwardService : null;
String audience = idSvc != null ? identityForwardService.audienceFor(serverId, serverName) : null;
List<ToolCallback> wrapped = wrapServerCallbacks(serverId, cbs, idSvc, audience, serverName);
List<ToolCallback> wrapped = wrapServerCallbacks(serverId, cbs, idSvc, audience, serverName,
entry.getValue(), objectMapper);
lastGoodCallbacks.put(serverId, wrapped);
allCallbacks.addAll(wrapped);
continue;
@ -253,7 +263,7 @@ public class McpClientManager {
* real {@link McpSyncClient}.
*/
static List<ToolCallback> wrapServerCallbacks(long serverId, ToolCallback[] cbs) {
return wrapServerCallbacks(serverId, cbs, null, null, null);
return wrapServerCallbacks(serverId, cbs, null, null, null, null, null);
}
/**
@ -266,10 +276,15 @@ public class McpClientManager {
* @param serverName human-readable MCP server name; forwarded into each
* {@link PrefixedNameToolCallback} so the tool description is tagged
* {@code [MCP server: <name>]}. May be {@code null} when unknown.
* @param mcpClient the active {@link McpSyncClient} for this server; when
* non-null each callback is wrapped in {@link ProgressAwareMcpToolCallback}
* so {@code _meta.progressToken} can be injected into tools/call requests.
* @param objectMapper JSON mapper for argument serialization inside the wrapper.
*/
static List<ToolCallback> wrapServerCallbacks(long serverId, ToolCallback[] cbs,
McpIdentityForwardService identitySvc, String audience,
String serverName) {
String serverName,
McpSyncClient mcpClient, ObjectMapper objectMapper) {
List<String> rawNames = new ArrayList<>(cbs.length);
for (ToolCallback cb : cbs) {
rawNames.add(cb.getToolDefinition() != null ? cb.getToolDefinition().name() : null);
@ -298,6 +313,13 @@ public class McpClientManager {
ToolCallback inner = identitySvc != null
? new IdentityForwardingToolCallback(cb, identitySvc, audience)
: cb;
// Wrap with progress-aware callback so _meta.progressToken is
// injected when ToolContext carries a progress token. Must sit
// inside PrefixedNameToolCallback so both the prefixed name and
// the call-path see the same chain.
if (mcpClient != null) {
inner = new ProgressAwareMcpToolCallback(inner, mcpClient, raw, objectMapper);
}
out.add(new PrefixedNameToolCallback(d.prefixedName(), inner, serverName));
}
return out;
@ -418,6 +440,29 @@ public class McpClientManager {
Long serverId = server.getId();
spec.toolsChangeConsumer(tools ->
eventPublisher.publishEvent(new McpServerChangedEvent("mcp-tools-changed:" + serverId)));
// MCP progress notifications: the server pushes progress for
// long-running tool calls. progressToken context lookup + event
// publish so McpProgressRelay can forward to the SSE stream.
spec.progressConsumer(progressNotification -> {
if (progressNotification == null || progressNotification.progressToken() == null) return;
McpProgressContext.ProgressEntry entry = progressContext.lookup(progressNotification.progressToken());
if (entry == null) return;
try {
McpProgressEvent event = new McpProgressEvent(
this,
entry.conversationId(),
entry.toolCallId(),
entry.toolName(),
progressNotification.progress(),
progressNotification.total(),
progressNotification.message()
);
eventPublisher.publishEvent(event);
} catch (Exception e) {
log.warn("Failed to publish McpProgressEvent: {}", e.getMessage());
}
});
}
return spec.build();
}

View File

@ -0,0 +1,54 @@
package vip.mate.tool.mcp.runtime;
import org.springframework.stereotype.Component;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
/**
* Thread-safe progressToken mapping table.
* Maintains {@code progressToken (conversationId, toolCallId, toolName)} mappings,
* and progress snapshots for SSE reconnect delivery.
*/
@Component
public class McpProgressContext {
private final Map<String, ProgressEntry> tokenMap = new ConcurrentHashMap<>();
/** Progress snapshots: conversationId → (toolCallId → latest progress JSON) */
private final Map<String, Map<String, String>> snapshotMap = new ConcurrentHashMap<>();
public record ProgressEntry(String conversationId, String toolCallId, String toolName) {}
public void register(String progressToken, ProgressEntry entry) {
tokenMap.put(progressToken, entry);
}
public ProgressEntry lookup(String progressToken) {
return tokenMap.get(progressToken);
}
public void remove(String progressToken) {
tokenMap.remove(progressToken);
}
/** Update progress snapshot (called on each progress notification). */
public void updateSnapshot(String conversationId, String toolCallId, String progressJson) {
snapshotMap.computeIfAbsent(conversationId, k -> new ConcurrentHashMap<>())
.put(toolCallId, progressJson);
}
/** SSE 重连时获取某个 conversation 下所有进行中的进度快照 */
public Map<String, String> getSnapshots(String conversationId) {
Map<String, String> tools = snapshotMap.get(conversationId);
return tools != null ? Map.copyOf(tools) : Map.of();
}
/** Remove snapshot after tool completion. */
public void removeSnapshot(String conversationId, String toolCallId) {
Map<String, String> tools = snapshotMap.get(conversationId);
if (tools != null) {
tools.remove(toolCallId);
}
}
}

View File

@ -0,0 +1,36 @@
package vip.mate.tool.mcp.runtime;
import org.springframework.context.ApplicationEvent;
/**
* MCP tool-call progress event.
* Published by {@code McpClientManager.progressConsumer} and consumed
* by {@link McpProgressRelay} for forwarding to {@code ChatStreamTracker}.
*/
public class McpProgressEvent extends ApplicationEvent {
private final String conversationId;
private final String toolCallId;
private final String toolName;
private final double progress; // 0.0 ~ 1.0
private final Double total; // may be null
private final String message; // current stage description
public McpProgressEvent(Object source, String conversationId, String toolCallId,
String toolName, double progress, Double total, String message) {
super(source);
this.conversationId = conversationId;
this.toolCallId = toolCallId;
this.toolName = toolName;
this.progress = progress;
this.total = total;
this.message = message;
}
public String getConversationId() { return conversationId; }
public String getToolCallId() { return toolCallId; }
public String getToolName() { return toolName; }
public double getProgress() { return progress; }
public Double getTotal() { return total; }
public String getMessage() { return message; }
}

View File

@ -0,0 +1,61 @@
package vip.mate.tool.mcp.runtime;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.context.event.EventListener;
import org.springframework.stereotype.Component;
import vip.mate.channel.web.ChatStreamTracker;
import java.util.Map;
/**
* MCP progress event relay listens for {@link McpProgressEvent} and forwards to
* {@link ChatStreamTracker}. Progress events skip the event buffer ({@code skipBuffer=true})
* and do not participate in SSE reconnect replay. On reconnect, the latest snapshot is
* read from {@link McpProgressContext} by {@code ChatStreamTracker.attach()}.
*/
@Slf4j
@Component
@RequiredArgsConstructor
public class McpProgressRelay {
/**
* SSE event name constant, agreed upon between frontend and backend.
*/
public static final String EVENT_TOOL_PROGRESS = "tool_call_progress";
private final ChatStreamTracker streamTracker;
private final McpProgressContext progressContext;
private final ObjectMapper objectMapper;
@EventListener
public void onMcpProgress(McpProgressEvent event) {
try {
Map<String, Object> data = Map.of(
"toolCallId", event.getToolCallId(),
"toolName", event.getToolName(),
"percent", Math.round(event.getProgress() * 10000.0) / 100.0,
"total", event.getTotal() != null ? event.getTotal() : 1.0,
"message", event.getMessage() != null ? event.getMessage() : "",
"stage", inferStage(event.getProgress())
);
String jsonData = objectMapper.writeValueAsString(data);
// Update snapshot for SSE reconnect
progressContext.updateSnapshot(event.getConversationId(), event.getToolCallId(), jsonData);
// Broadcast to SSE (skipBuffer=true, not cached in ring buffer)
streamTracker.broadcastObject(event.getConversationId(), EVENT_TOOL_PROGRESS, data, true);
} catch (Exception e) {
log.warn("Failed to relay MCP progress: {}", e.getMessage());
}
}
/** Infer stage name from progress percentage. */
private String inferStage(double progress) {
if (progress <= 0.05) return "prepare";
if (progress <= 0.95) return "execute";
return "finalize";
}
}

View File

@ -0,0 +1,122 @@
package vip.mate.tool.mcp.runtime;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import io.modelcontextprotocol.client.McpSyncClient;
import io.modelcontextprotocol.spec.McpSchema;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.model.ToolContext;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.ai.tool.definition.ToolDefinition;
import org.springframework.ai.tool.metadata.ToolMetadata;
import java.util.Map;
import java.util.UUID;
/**
* MCP tool callback wrapper that injects {@code _meta.progressToken} into
* {@code tools/call} requests, enabling MCP Servers to push progress via
* {@code notifications/progress}.
*
* <p>When {@code MCP_PROGRESS_TOKEN} is present in {@link ToolContext}, this wrapper
* calls {@link McpSyncClient#callTool(McpSchema.CallToolRequest)} directly with the
* progressToken injected. Otherwise it delegates to the original callback
* (compatible with MCP Servers that do not support progress, and with built-in tools).
*/
@Slf4j
public final class ProgressAwareMcpToolCallback implements ToolCallback {
/** Key in ToolContext where the progressToken is stored. */
public static final String MCP_PROGRESS_TOKEN_KEY = "_mcp_progress_token";
private final ToolCallback delegate;
private final McpSyncClient mcpClient;
private final String rawToolName;
private final ObjectMapper objectMapper;
public ProgressAwareMcpToolCallback(ToolCallback delegate, McpSyncClient mcpClient,
String rawToolName, ObjectMapper objectMapper) {
this.delegate = delegate;
this.mcpClient = mcpClient;
this.rawToolName = rawToolName;
this.objectMapper = objectMapper;
}
@Override
public ToolDefinition getToolDefinition() {
return delegate.getToolDefinition();
}
@Override
public ToolMetadata getToolMetadata() {
return delegate.getToolMetadata();
}
@Override
public String call(String toolInput) {
return delegate.call(toolInput);
}
@Override
public String call(String toolInput, ToolContext toolContext) {
String progressToken = null;
if (toolContext != null && toolContext.getContext() != null) {
Object token = toolContext.getContext().get(MCP_PROGRESS_TOKEN_KEY);
if (token instanceof String s && !s.isBlank()) {
progressToken = s;
}
}
if (progressToken == null) {
return delegate.call(toolInput, toolContext);
}
try {
// Apply identity forwarding BEFORE building CallToolRequest
// otherwise the progress path would silently bypass identity injection.
String effectiveInput = toolInput;
if (delegate instanceof IdentityForwardingToolCallback idFwd) {
effectiveInput = idFwd.inject(toolInput, toolContext);
}
Map<String, Object> arguments = parseArguments(effectiveInput);
McpSchema.CallToolRequest request = McpSchema.CallToolRequest.builder()
.name(rawToolName)
.arguments(arguments != null ? arguments : Map.of())
.meta(Map.of("progressToken", progressToken))
.build();
McpSchema.CallToolResult result = mcpClient.callTool(request);
return serializeResult(result);
} catch (Exception e) {
log.warn("Progress-aware MCP call failed for tool '{}', falling back to delegate: {}",
rawToolName, e.getMessage());
return delegate.call(toolInput, toolContext);
}
}
private Map<String, Object> parseArguments(String toolInput) {
if (toolInput == null || toolInput.isBlank()) return Map.of();
try {
return objectMapper.readValue(toolInput, new TypeReference<Map<String, Object>>() {});
} catch (Exception e) {
log.debug("Failed to parse MCP tool arguments as JSON, using raw string: {}", e.getMessage());
return Map.of("input", toolInput);
}
}
private String serializeResult(McpSchema.CallToolResult result) {
if (result == null) return "";
if (result.content() == null || result.content().isEmpty()) return "";
StringBuilder sb = new StringBuilder();
for (var content : result.content()) {
if (content instanceof McpSchema.TextContent tc) {
sb.append(tc.text());
} else {
sb.append(content.toString());
}
}
return sb.toString();
}
/** Return the underlying delegate (for ReturnDirect / IdentityForward detection). */
public ToolCallback getDelegate() {
return delegate;
}
}

View File

@ -0,0 +1,110 @@
package vip.mate.tool.mcp.runtime;
import com.fasterxml.jackson.databind.ObjectMapper;
import io.modelcontextprotocol.client.McpSyncClient;
import org.junit.jupiter.api.DisplayName;
import org.junit.jupiter.api.Test;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.ai.tool.definition.DefaultToolDefinition;
import org.springframework.ai.tool.definition.ToolDefinition;
import org.springframework.ai.tool.metadata.ToolMetadata;
import org.springframework.ai.chat.model.ToolContext;
import java.util.List;
import static org.junit.jupiter.api.Assertions.*;
import static org.mockito.Mockito.mock;
/**
* Black-box regression suite verifying that the new
* {@link McpClientManager#wrapServerCallbacks(long, ToolCallback[], McpIdentityForwardService,
* String, String, McpSyncClient, ObjectMapper)} overload:
* <ol>
* <li>does <b>not</b> break the existing null-mcpClient path</li>
* <li>wraps with {@link ProgressAwareMcpToolCallback} when mcpClient is provided</li>
* </ol>
*/
class McpClientManagerProgressWrapTest {
private static ToolCallback stub(String name) {
ToolDefinition def = DefaultToolDefinition.builder()
.name(name).description("").inputSchema("{}").build();
return new ToolCallback() {
@Override public ToolDefinition getToolDefinition() { return def; }
@Override public ToolMetadata getToolMetadata() { return ToolCallback.super.getToolMetadata(); }
@Override public String call(String toolInput) { return name + ":" + toolInput; }
@Override public String call(String toolInput, ToolContext toolContext) { return call(toolInput); }
};
}
// Backward-compat: null McpSyncClient (same as before)
@Test
@DisplayName("null McpSyncClient → no ProgressAwareMcpToolCallback wrapping")
void nullMcpClientNoProgressWrap() {
ToolCallback cb = stub("search");
List<ToolCallback> wrapped = McpClientManager.wrapServerCallbacks(99L,
new ToolCallback[]{cb}, null, null, null, null, null);
assertEquals(1, wrapped.size());
assertInstanceOf(PrefixedNameToolCallback.class, wrapped.get(0));
PrefixedNameToolCallback p = (PrefixedNameToolCallback) wrapped.get(0);
// Inner should be the original stub, NOT a ProgressAwareMcpToolCallback
assertFalse(p.getDelegate() instanceof ProgressAwareMcpToolCallback,
"should NOT wrap when mcpClient is null");
}
// With McpSyncClient wraps
@Test
@DisplayName("non-null McpSyncClient wraps with ProgressAwareMcpToolCallback")
void withMcpClientWrapsProgress() {
McpSyncClient client = mock(McpSyncClient.class);
ToolCallback cb = stub("long_task");
ObjectMapper mapper = new ObjectMapper();
List<ToolCallback> wrapped = McpClientManager.wrapServerCallbacks(88L,
new ToolCallback[]{cb}, null, null, null, client, mapper);
assertEquals(1, wrapped.size());
assertInstanceOf(PrefixedNameToolCallback.class, wrapped.get(0));
PrefixedNameToolCallback p = (PrefixedNameToolCallback) wrapped.get(0);
assertInstanceOf(ProgressAwareMcpToolCallback.class, p.getDelegate(),
"should wrap with ProgressAwareMcpToolCallback when mcpClient is provided");
ProgressAwareMcpToolCallback prog = (ProgressAwareMcpToolCallback) p.getDelegate();
assertEquals(cb, prog.getDelegate(), "original callback preserved as delegate");
}
@Test
@DisplayName("ProgressAwareMcpToolCallback sits outside IdentityForward but inside PrefixedName (correct chain)")
void chainOrder() {
McpSyncClient client = mock(McpSyncClient.class);
ToolCallback cb = stub("private_data");
ObjectMapper mapper = new ObjectMapper();
// identitySvc=null no identity wrapping
List<ToolCallback> wrapped = McpClientManager.wrapServerCallbacks(77L,
new ToolCallback[]{cb}, null, null, "my-server", client, mapper);
assertEquals(1, wrapped.size());
assertInstanceOf(PrefixedNameToolCallback.class, wrapped.get(0));
PrefixedNameToolCallback p = (PrefixedNameToolCallback) wrapped.get(0);
assertInstanceOf(ProgressAwareMcpToolCallback.class, p.getDelegate());
ProgressAwareMcpToolCallback prog = (ProgressAwareMcpToolCallback) p.getDelegate();
// IdentityForwarding is NOT wrapped because identitySvc=null; the raw stub IS the delegate
assertSame(cb, prog.getDelegate());
}
@Test
@DisplayName("existing two-arg wrapServerCallbacks overload still compiles and works")
void twoArgOverloadStillWorks() {
ToolCallback cb = stub("read_file");
// This is the original API used by McpClientManagerWrapTest must still work
List<ToolCallback> wrapped = McpClientManager.wrapServerCallbacks(66L,
new ToolCallback[]{cb});
assertEquals(1, wrapped.size());
assertInstanceOf(PrefixedNameToolCallback.class, wrapped.get(0));
}
}

View File

@ -38,7 +38,8 @@ class McpClientManagerSnapshotTest {
void staleListToolsServesSnapshotAndRequestsReconnect() throws Exception {
ApplicationEventPublisher publisher = mock(ApplicationEventPublisher.class);
McpClientManager manager = new McpClientManager(publisher,
new McpIdentityForwardService(new McpIdentityForwardProperties()));
new McpIdentityForwardService(new McpIdentityForwardProperties()),
null, null);
// A client whose connection went stale: every listTools() throws.
McpSyncClient deadClient = mock(McpSyncClient.class);

View File

@ -0,0 +1,148 @@
package vip.mate.tool.mcp.runtime;
import org.junit.jupiter.api.DisplayName;
import org.junit.jupiter.api.Test;
import static org.junit.jupiter.api.Assertions.*;
/**
* White-box unit tests for {@link McpProgressContext}.
* Covers register lookup remove lifecycle, snapshot persistence,
* multi-conversation isolation, and concurrent safety.
*/
class McpProgressContextTest {
private McpProgressContext ctx() {
return new McpProgressContext();
}
// Token Map
@Test
@DisplayName("register → lookup returns same entry")
void registerAndLookup() {
McpProgressContext ctx = ctx();
var entry = new McpProgressContext.ProgressEntry("conv_1", "call_1", "search");
ctx.register("token-1", entry);
assertSame(entry, ctx.lookup("token-1"));
}
@Test
@DisplayName("lookup for unregistered token returns null")
void lookupMissingReturnsNull() {
assertNull(ctx().lookup("nonexistent"));
}
@Test
@DisplayName("remove makes subsequent lookup return null")
void removeThenLookupReturnsNull() {
McpProgressContext ctx = ctx();
ctx.register("tok", new McpProgressContext.ProgressEntry("c", "t", "n"));
ctx.remove("tok");
assertNull(ctx.lookup("tok"));
}
@Test
@DisplayName("register overwrites existing entry for same token")
void registerOverwrites() {
McpProgressContext ctx = ctx();
var first = new McpProgressContext.ProgressEntry("c1", "t1", "n1");
var second = new McpProgressContext.ProgressEntry("c2", "t2", "n2");
ctx.register("tok", first);
ctx.register("tok", second);
assertSame(second, ctx.lookup("tok"));
}
@Test
@DisplayName("remove of non-existent token is no-op")
void removeNonexistentIsNoop() {
McpProgressContext ctx = ctx();
assertDoesNotThrow(() -> ctx.remove("ghost"));
}
// Snapshot Map
@Test
@DisplayName("updateSnapshot stores and getSnapshots returns latest")
void snapshotStoreAndRetrieve() {
McpProgressContext ctx = ctx();
ctx.updateSnapshot("conv_1", "call_a", "{\"percent\":30}");
ctx.updateSnapshot("conv_1", "call_a", "{\"percent\":70}");
var snapshots = ctx.getSnapshots("conv_1");
assertEquals(1, snapshots.size());
assertEquals("{\"percent\":70}", snapshots.get("call_a"));
}
@Test
@DisplayName("getSnapshots for unknown conversation returns empty map")
void snapshotsForUnknownConversation() {
assertTrue(ctx().getSnapshots("no_such_conv").isEmpty());
}
@Test
@DisplayName("removeSnapshot cleans up individual tool snapshot")
void removeSnapshot() {
McpProgressContext ctx = ctx();
ctx.updateSnapshot("conv_1", "call_a", "{\"p\":50}");
ctx.updateSnapshot("conv_1", "call_b", "{\"p\":80}");
ctx.removeSnapshot("conv_1", "call_a");
var snapshots = ctx.getSnapshots("conv_1");
assertEquals(1, snapshots.size());
assertNull(snapshots.get("call_a"));
assertEquals("{\"p\":80}", snapshots.get("call_b"));
}
@Test
@DisplayName("removeSnapshot for unknown keys is no-op")
void removeSnapshotNoop() {
assertDoesNotThrow(() -> {
McpProgressContext ctx = ctx();
ctx.removeSnapshot("no_conv", "no_call");
ctx.updateSnapshot("cv", "cl", "{}");
ctx.removeSnapshot("cv", "other");
assertEquals(1, ctx.getSnapshots("cv").size());
});
}
@Test
@DisplayName("getSnapshots returns immutable copy")
void snapshotsImmutable() {
McpProgressContext ctx = ctx();
ctx.updateSnapshot("c", "t", "{}");
var snap = ctx.getSnapshots("c");
assertThrows(UnsupportedOperationException.class, () -> snap.put("x", "y"));
}
@Test
@DisplayName("multiple conversations isolated")
void multiConversationIsolation() {
McpProgressContext ctx = ctx();
ctx.updateSnapshot("c1", "t1", "A");
ctx.updateSnapshot("c2", "t2", "B");
assertEquals(1, ctx.getSnapshots("c1").size());
assertEquals(1, ctx.getSnapshots("c2").size());
assertEquals("A", ctx.getSnapshots("c1").get("t1"));
assertEquals("B", ctx.getSnapshots("c2").get("t2"));
}
@Test
@DisplayName("register + snapshot lifecycle: full round-trip")
void fullRoundtrip() {
McpProgressContext ctx = ctx();
var entry = new McpProgressContext.ProgressEntry("conv_x", "call_x", "long_task");
ctx.register("pt-1", entry);
assertEquals(entry, ctx.lookup("pt-1"));
ctx.updateSnapshot("conv_x", "call_x", "{\"percent\":33}");
ctx.updateSnapshot("conv_x", "call_x", "{\"percent\":99}");
assertEquals("{\"percent\":99}", ctx.getSnapshots("conv_x").get("call_x"));
ctx.remove("pt-1");
ctx.removeSnapshot("conv_x", "call_x");
assertNull(ctx.lookup("pt-1"));
assertTrue(ctx.getSnapshots("conv_x").isEmpty());
}
}

View File

@ -0,0 +1,134 @@
package vip.mate.tool.mcp.runtime;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.DisplayName;
import org.junit.jupiter.api.Test;
import vip.mate.channel.web.ChatStreamTracker;
import java.util.Map;
import static org.junit.jupiter.api.Assertions.*;
import static org.mockito.ArgumentMatchers.*;
import static org.mockito.Mockito.*;
/**
* Integration test for the {@link McpProgressRelay} event listener.
* Verifies that {@link McpProgressEvent} {@link ChatStreamTracker#broadcastObject}
* forwarding works correctly, including snapshot updates and skipBuffer=true.
*/
class McpProgressRelayTest {
private ChatStreamTracker streamTracker;
private McpProgressContext progressContext;
private McpProgressRelay relay;
@BeforeEach
void setUp() {
streamTracker = mock(ChatStreamTracker.class);
progressContext = new McpProgressContext();
relay = new McpProgressRelay(streamTracker, progressContext, new ObjectMapper());
}
@Test
@DisplayName("relay forwards event to ChatStreamTracker with skipBuffer=true")
void forwardsEvent() {
McpProgressEvent event = new McpProgressEvent(
this, "conv_1", "call_abc", "long_task", 0.5, 1.0, "Processing...");
relay.onMcpProgress(event);
verify(streamTracker).broadcastObject(
eq("conv_1"),
eq(McpProgressRelay.EVENT_TOOL_PROGRESS),
any(Object.class),
eq(true));
}
@Test
@DisplayName("relay updates progress snapshot")
void updatesSnapshot() {
McpProgressEvent event = new McpProgressEvent(
this, "conv_1", "call_abc", "task", 0.75, 1.0, "Almost done");
relay.onMcpProgress(event);
var snapshots = progressContext.getSnapshots("conv_1");
assertEquals(1, snapshots.size());
String json = snapshots.get("call_abc");
assertNotNull(json);
assertTrue(json.contains("\"percent\":75"));
assertTrue(json.contains("\"call_abc\""));
}
@Test
@DisplayName("streamTracker throws → relay logs warning, does not propagate")
void streamTrackerThrowsDoesNotPropagate() {
doThrow(new RuntimeException("SSE dead")).when(streamTracker)
.broadcastObject(any(), any(), any(), anyBoolean());
McpProgressEvent event = new McpProgressEvent(
this, "conv", "call", "tool", 0.0, null, "init");
// Should not throw
assertDoesNotThrow(() -> relay.onMcpProgress(event));
}
@Test
@DisplayName("null progress → broadcast still succeeds with 0.0")
void nullProgress() {
// This would be an edge case from MCP SDK; not expected but guarded
McpProgressEvent event = new McpProgressEvent(
this, "conv_2", "call_2", "task", 0.0, null, null);
relay.onMcpProgress(event);
verify(streamTracker).broadcastObject(
eq("conv_2"),
eq(McpProgressRelay.EVENT_TOOL_PROGRESS),
any(Object.class),
eq(true));
}
@Test
@DisplayName("stage inference: 0-5% → prepare, 5-95% → execute, 95%+ → finalize")
void stageInference() {
// Test via the relay that stage reflects in the broadcast data
McpProgressEvent event = new McpProgressEvent(
this, "c", "t", "task", 0.97, 1.0, "Finishing");
relay.onMcpProgress(event);
verify(streamTracker).broadcastObject(
eq("c"), eq("tool_call_progress"),
argThat((Object data) -> {
if (data instanceof Map<?, ?> m) {
return "finalize".equals(m.get("stage"));
}
return false;
}),
eq(true));
}
@Test
@DisplayName("event constant matches frontend expectation")
void eventConstantCorrect() {
assertEquals("tool_call_progress", McpProgressRelay.EVENT_TOOL_PROGRESS,
"must match the SSE event name used in useChat.ts and ChatStreamTracker");
}
@Test
@DisplayName("multiple events for same tool call update snapshot idempotently")
void multipleEventsUpdateSameSnapshot() {
relay.onMcpProgress(new McpProgressEvent(this, "c", "t", "n", 0.2, 1.0, "A"));
relay.onMcpProgress(new McpProgressEvent(this, "c", "t", "n", 0.6, 1.0, "B"));
relay.onMcpProgress(new McpProgressEvent(this, "c", "t", "n", 0.99, 1.0, "C"));
// Only 1 snapshot (latest)
var snapshots = progressContext.getSnapshots("c");
assertEquals(1, snapshots.size());
String json = snapshots.get("t");
assertTrue(json.contains("\"percent\":99"));
assertTrue(json.contains("C"));
}
}

View File

@ -0,0 +1,171 @@
package vip.mate.tool.mcp.runtime;
import com.fasterxml.jackson.databind.ObjectMapper;
import io.modelcontextprotocol.client.McpSyncClient;
import io.modelcontextprotocol.spec.McpSchema;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.DisplayName;
import org.junit.jupiter.api.Test;
import org.springframework.ai.chat.model.ToolContext;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.ai.tool.definition.DefaultToolDefinition;
import org.springframework.ai.tool.definition.ToolDefinition;
import org.springframework.ai.tool.metadata.ToolMetadata;
import java.util.List;
import java.util.Map;
import static org.junit.jupiter.api.Assertions.*;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.*;
/**
* White-box tests for {@link ProgressAwareMcpToolCallback}.
*
* <p>Covers the two code paths:
* <ol>
* <li>progressToken present in ToolContext direct McpSyncClient.callTool() with injected meta</li>
* <li>progressToken absent delegates to inner callback (backward-compatible)</li>
* </ol>
*/
class ProgressAwareMcpToolCallbackTest {
private static final ObjectMapper MAPPER = new ObjectMapper();
private ToolCallback delegate;
private McpSyncClient mcpClient;
private ProgressAwareMcpToolCallback wrapper;
@BeforeEach
void setUp() {
delegate = mock(ToolCallback.class);
mcpClient = mock(McpSyncClient.class);
when(delegate.getToolDefinition()).thenReturn(
DefaultToolDefinition.builder().name("search").description("desc").inputSchema("{}").build());
when(delegate.getToolMetadata()).thenReturn(ToolMetadata.builder().build());
wrapper = new ProgressAwareMcpToolCallback(delegate, mcpClient, "search", MAPPER);
}
@Test
@DisplayName("getToolDefinition delegates to inner callback")
void delegatesGetToolDefinition() {
assertEquals("search", wrapper.getToolDefinition().name());
verify(delegate).getToolDefinition();
}
@Test
@DisplayName("getToolMetadata delegates to inner callback")
void delegatesGetToolMetadata() {
assertNotNull(wrapper.getToolMetadata());
verify(delegate).getToolMetadata();
}
@Test
@DisplayName("call(toolInput) without ToolContext delegates to inner")
void callWithoutToolContextDelegates() {
when(delegate.call("{}")).thenReturn("result");
assertEquals("result", wrapper.call("{}"));
verify(delegate).call("{}");
verifyNoInteractions(mcpClient);
}
@Test
@DisplayName("call with ToolContext but without progressToken delegates to inner")
void callWithoutProgressTokenDelegates() {
ToolContext ctx = new ToolContext(Map.of());
when(delegate.call("{}", ctx)).thenReturn("delegated");
assertEquals("delegated", wrapper.call("{}", ctx));
verify(delegate).call("{}", ctx);
verifyNoInteractions(mcpClient);
}
@Test
@DisplayName("call with null ToolContext delegates to inner")
void callWithNullToolContextDelegates() {
when(delegate.call("{}", null)).thenReturn("null_ctx");
assertEquals("null_ctx", wrapper.call("{}", (ToolContext) null));
verify(delegate).call("{}", (ToolContext) null);
verifyNoInteractions(mcpClient);
}
@Test
@DisplayName("call with progressToken in ToolContext calls McpSyncClient directly with meta injected")
void callWithProgressTokenUsesMcpClient() {
ToolContext ctx = new ToolContext(Map.of(
ProgressAwareMcpToolCallback.MCP_PROGRESS_TOKEN_KEY, "pt-uuid-123"));
McpSchema.TextContent textContent = new McpSchema.TextContent("mcp result");
McpSchema.CallToolResult result = new McpSchema.CallToolResult(List.of(textContent), false);
when(mcpClient.callTool(any())).thenReturn(result);
String output = wrapper.call("{\"q\":\"hello\"}", ctx);
assertEquals("mcp result", output);
verify(mcpClient).callTool(any(McpSchema.CallToolRequest.class));
verify(delegate, never()).call(any(), any());
}
@Test
@DisplayName("progressToken present but blank → delegates (edge case)")
void blankProgressTokenDelegates() {
ToolContext ctx = new ToolContext(Map.of(
ProgressAwareMcpToolCallback.MCP_PROGRESS_TOKEN_KEY, " "));
when(delegate.call("{}", ctx)).thenReturn("fallback");
assertEquals("fallback", wrapper.call("{}", ctx));
verify(delegate).call("{}", ctx);
verifyNoInteractions(mcpClient);
}
@Test
@DisplayName("McpSyncClient throws → falls back to delegate")
void mcpClientThrowsFallsBackToDelegate() {
ToolContext ctx = new ToolContext(Map.of(
ProgressAwareMcpToolCallback.MCP_PROGRESS_TOKEN_KEY, "tok"));
when(mcpClient.callTool(any())).thenThrow(new RuntimeException("connection lost"));
when(delegate.call(eq("{}"), any(ToolContext.class))).thenReturn("fallback result");
String output = wrapper.call("{}", ctx);
assertEquals("fallback result", output);
verify(mcpClient).callTool(any());
verify(delegate).call(eq("{}"), any());
}
@Test
@DisplayName("callTool succeeds with multi-text content concatenated")
void multiTextContentConcatenated() {
ToolContext ctx = new ToolContext(Map.of(
ProgressAwareMcpToolCallback.MCP_PROGRESS_TOKEN_KEY, "tok"));
McpSchema.CallToolResult result = new McpSchema.CallToolResult(List.of(
new McpSchema.TextContent("part1"),
new McpSchema.TextContent("part2")), false);
when(mcpClient.callTool(any())).thenReturn(result);
assertEquals("part1part2", wrapper.call("{}", ctx));
}
@Test
@DisplayName("getDelegate returns inner callback (for ReturnDirect / IdentityForward detection)")
void getDelegateReturnsInner() {
assertSame(delegate, wrapper.getDelegate());
}
@Test
@DisplayName("parseArguments handles null input")
void parseArgumentsHandlesNull() {
ToolContext ctx = new ToolContext(Map.of(
ProgressAwareMcpToolCallback.MCP_PROGRESS_TOKEN_KEY, "tok"));
when(mcpClient.callTool(any())).thenReturn(
new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("ok")), false));
assertEquals("ok", wrapper.call(null, ctx));
}
@Test
@DisplayName("parseArguments handles blank input")
void parseArgumentsHandlesBlank() {
ToolContext ctx = new ToolContext(Map.of(
ProgressAwareMcpToolCallback.MCP_PROGRESS_TOKEN_KEY, "tok"));
when(mcpClient.callTool(any())).thenReturn(
new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("ok")), false));
assertEquals("ok", wrapper.call(" ", ctx));
}
}

View File

@ -58,6 +58,8 @@ const isRead = computed(() => {
const isSuccess = computed(() => props.segment.status === 'completed' && props.segment.toolSuccess !== false)
const isError = computed(() => props.segment.status === 'error' || props.segment.toolSuccess === false)
const isRunning = computed(() => props.segment.status === 'running')
// MCP progress bar: show when running AND progress data is available
const hasProgress = computed(() => isRunning.value && props.segment.progress != null)
// A delegation flagged by the heartbeat watchdog as making no progress.
const isStalled = computed(() => isDelegation.value && isRunning.value && !!props.segment.delegationStale)
// Fire-and-forget delegation: runs detached, result comes via task_output later.
@ -118,6 +120,7 @@ const detailStatus = computed<'running' | 'completed' | 'error'>(() => {
<div class="seg-tool__header" @click="hasBody ? (expanded = !expanded) : null">
<span class="seg-tool__status">
<el-icon v-if="isAsync" class="seg-tool__async" :title="$t('chat.subagentAsync')" :size="13"><Clock /></el-icon>
<el-icon v-else-if="hasProgress" :size="13"><Loading /></el-icon>
<el-icon v-else-if="isRunning" class="is-loading" :size="13"><Loading /></el-icon>
<el-icon v-else-if="isSuccess" :size="13"><Select /></el-icon>
<el-icon v-else :size="13"><CloseBold /></el-icon>
@ -147,6 +150,14 @@ const detailStatus = computed<'running' | 'completed' | 'error'>(() => {
><ArrowDown /></el-icon>
</span>
</div>
<!-- MCP progress bar: shown when running and progress data is available -->
<div v-if="hasProgress" class="seg-tool__progress">
<div class="seg-tool__progress-bar">
<div class="seg-tool__progress-fill" :style="{ width: (segment.progress || 0) + '%' }"></div>
</div>
<div class="seg-tool__progress-label">{{ segment.progress }}%</div>
<div v-if="segment.progressMessage" class="seg-tool__progress-msg">{{ segment.progressMessage }}</div>
</div>
<Transition name="seg-slide">
<div v-if="expanded && hasBody" class="seg-tool__body">
<!-- Nested subagent timeline (delegation segments) -->
@ -385,4 +396,39 @@ const detailStatus = computed<'running' | 'completed' | 'error'>(() => {
opacity: 0;
transform: translateY(-4px);
}
/* MCP progress bar */
.seg-tool__progress {
padding: 0 10px 6px 22px;
display: flex;
flex-wrap: wrap;
align-items: center;
gap: 8px;
}
.seg-tool__progress-bar {
flex: 1;
min-width: 80px;
height: 6px;
background: var(--mc-bg-muted);
border-radius: 3px;
overflow: hidden;
}
.seg-tool__progress-fill {
height: 100%;
background: linear-gradient(90deg, var(--mc-primary), var(--mc-primary-light, #f0a070));
border-radius: 3px;
transition: width 0.3s ease;
}
.seg-tool__progress-label {
font-size: 12px;
font-weight: 500;
color: var(--mc-primary);
white-space: nowrap;
}
.seg-tool__progress-msg {
width: 100%;
font-size: 11px;
color: var(--mc-text-tertiary);
line-height: 1.3;
}
</style>

View File

@ -958,6 +958,21 @@ export function useChat(options: UseChatOptions): UseChatReturn {
stream.on('tool_call_started', handleToolCallStarted)
stream.on('tool_call_completed', handleToolCallCompleted)
// MCP long-running tool progress: update the matching tool_call segment's
// progress field so ToolCallSegment can render a progress bar.
stream.on('tool_call_progress', (data: any) => {
if (isStaleEvent(data)) return
if (!data?.toolCallId) return
const segs = currentSegments.value
const toolSeg = segs.find((s: MessageSegment) =>
s.type === 'tool_call' && s.status === 'running' && s.toolCallId === data.toolCallId)
if (toolSeg) {
toolSeg.progress = data.percent
toolSeg.progressMessage = data.message
toolSeg.progressStage = data.stage
}
})
// ===== Browser action events =====
stream.on('browser_action', (data) => {

View File

@ -268,6 +268,12 @@ export interface MessageSegment {
supersededBySegmentId?: string
/** Machine-readable reason for superseding this segment. */
supersededReason?: string
/** MCP progress: 0-100 percentage */
progress?: number
/** MCP progress: human-readable stage message */
progressMessage?: string
/** MCP progress: current stage (prepare/execute/finalize) */
progressStage?: string
}
/** A file artifact generated by a tool during the turn, surfaced in the run-overview rail. */