mirror of
https://gitee.com/mateos/mateclaw.git
synced 2026-09-13 19:23:42 +08:00
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:
parent
4ae4731d54
commit
e35c07f742
520
MCP长时任务进度推送-落地方案.md
Normal file
520
MCP长时任务进度推送-落地方案.md
Normal 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` 机制,在 MateClaw(MCP 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` 请求发给 server,server 在 `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);
|
||||
|
||||
// 广播到 SSE(skipBuffer = 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 Server(SSE 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 均支持
|
||||
@ -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) {
|
||||
|
||||
@ -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();
|
||||
}
|
||||
|
||||
|
||||
@ -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);
|
||||
|
||||
@ -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);
|
||||
|
||||
@ -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();
|
||||
}
|
||||
|
||||
@ -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);
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -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; }
|
||||
}
|
||||
@ -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";
|
||||
}
|
||||
}
|
||||
@ -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;
|
||||
}
|
||||
}
|
||||
@ -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));
|
||||
}
|
||||
}
|
||||
@ -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);
|
||||
|
||||
@ -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());
|
||||
}
|
||||
}
|
||||
@ -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"));
|
||||
}
|
||||
}
|
||||
@ -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));
|
||||
}
|
||||
}
|
||||
@ -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>
|
||||
|
||||
@ -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) => {
|
||||
|
||||
@ -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. */
|
||||
|
||||
Loading…
Reference in New Issue
Block a user