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.system.service.SystemSettingService;
|
||||||
import vip.mate.tool.ToolRegistry;
|
import vip.mate.tool.ToolRegistry;
|
||||||
import vip.mate.tool.disclosure.ToolUsageRecencyTracker;
|
import vip.mate.tool.disclosure.ToolUsageRecencyTracker;
|
||||||
|
import vip.mate.tool.mcp.runtime.McpProgressContext;
|
||||||
import vip.mate.memory.spi.MemoryManager;
|
import vip.mate.memory.spi.MemoryManager;
|
||||||
import vip.mate.workspace.document.WorkspaceFileService;
|
import vip.mate.workspace.document.WorkspaceFileService;
|
||||||
import vip.mate.tool.guard.service.ToolGuardService;
|
import vip.mate.tool.guard.service.ToolGuardService;
|
||||||
@ -105,6 +106,7 @@ public class AgentGraphBuilder {
|
|||||||
private final ModelContextWindowResolver contextWindowResolver;
|
private final ModelContextWindowResolver contextWindowResolver;
|
||||||
private final PrefixBudgetPlanner prefixBudgetPlanner;
|
private final PrefixBudgetPlanner prefixBudgetPlanner;
|
||||||
private final ToolUsageRecencyTracker toolUsageRecencyTracker;
|
private final ToolUsageRecencyTracker toolUsageRecencyTracker;
|
||||||
|
private final McpProgressContext progressContext;
|
||||||
private final vip.mate.llm.service.ModelCapabilityService modelCapabilityService;
|
private final vip.mate.llm.service.ModelCapabilityService modelCapabilityService;
|
||||||
private final ProviderRouter providerRouter;
|
private final ProviderRouter providerRouter;
|
||||||
private final PlanningService planningService;
|
private final PlanningService planningService;
|
||||||
@ -632,6 +634,7 @@ public class AgentGraphBuilder {
|
|||||||
// the right invocation pattern instead of a dead-end error.
|
// the right invocation pattern instead of a dead-end error.
|
||||||
executor.setSkillRuntimeService(skillRuntimeService);
|
executor.setSkillRuntimeService(skillRuntimeService);
|
||||||
executor.setUsageRecencyTracker(toolUsageRecencyTracker);
|
executor.setUsageRecencyTracker(toolUsageRecencyTracker);
|
||||||
|
executor.setProgressContext(progressContext);
|
||||||
// Optional: route child-agent denied-tool audit events through
|
// Optional: route child-agent denied-tool audit events through
|
||||||
// the audit pipeline. Null when audit is not wired (legacy / test).
|
// the audit pipeline. Null when audit is not wired (legacy / test).
|
||||||
if (auditEventService != null) {
|
if (auditEventService != null) {
|
||||||
@ -928,6 +931,7 @@ public class AgentGraphBuilder {
|
|||||||
// the right invocation pattern instead of a dead-end error.
|
// the right invocation pattern instead of a dead-end error.
|
||||||
executor.setSkillRuntimeService(skillRuntimeService);
|
executor.setSkillRuntimeService(skillRuntimeService);
|
||||||
executor.setUsageRecencyTracker(toolUsageRecencyTracker);
|
executor.setUsageRecencyTracker(toolUsageRecencyTracker);
|
||||||
|
executor.setProgressContext(progressContext);
|
||||||
// Optional: route child-agent denied-tool audit events through
|
// Optional: route child-agent denied-tool audit events through
|
||||||
// the audit pipeline. Null when audit is not wired (legacy / test).
|
// the audit pipeline. Null when audit is not wired (legacy / test).
|
||||||
if (auditEventService != null) {
|
if (auditEventService != null) {
|
||||||
|
|||||||
@ -8,7 +8,9 @@ import org.springframework.ai.chat.model.ToolContext;
|
|||||||
import org.springframework.ai.tool.ToolCallback;
|
import org.springframework.ai.tool.ToolCallback;
|
||||||
import vip.mate.tool.builtin.ToolExecutionContext;
|
import vip.mate.tool.builtin.ToolExecutionContext;
|
||||||
import vip.mate.tool.disclosure.ToolUsageRecencyTracker;
|
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.McpToolNameResolver;
|
||||||
|
import vip.mate.tool.mcp.runtime.ProgressAwareMcpToolCallback;
|
||||||
import vip.mate.agent.AgentToolSet;
|
import vip.mate.agent.AgentToolSet;
|
||||||
import vip.mate.agent.GraphEventPublisher;
|
import vip.mate.agent.GraphEventPublisher;
|
||||||
import vip.mate.agent.context.ChatOrigin;
|
import vip.mate.agent.context.ChatOrigin;
|
||||||
@ -256,10 +258,17 @@ public class ToolExecutionExecutor {
|
|||||||
/** Optional recency feed for budget-driven tool-disclosure demotion. */
|
/** Optional recency feed for budget-driven tool-disclosure demotion. */
|
||||||
private ToolUsageRecencyTracker usageRecencyTracker;
|
private ToolUsageRecencyTracker usageRecencyTracker;
|
||||||
|
|
||||||
|
/** Optional MCP progress context for long-running tool progress relay. */
|
||||||
|
private McpProgressContext progressContext;
|
||||||
|
|
||||||
public void setUsageRecencyTracker(ToolUsageRecencyTracker tracker) {
|
public void setUsageRecencyTracker(ToolUsageRecencyTracker tracker) {
|
||||||
this.usageRecencyTracker = tracker;
|
this.usageRecencyTracker = tracker;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
public void setProgressContext(McpProgressContext ctx) {
|
||||||
|
this.progressContext = ctx;
|
||||||
|
}
|
||||||
|
|
||||||
public void setSkillRuntimeService(vip.mate.skill.runtime.SkillRuntimeService s) {
|
public void setSkillRuntimeService(vip.mate.skill.runtime.SkillRuntimeService s) {
|
||||||
this.skillRuntimeService = s;
|
this.skillRuntimeService = s;
|
||||||
}
|
}
|
||||||
@ -883,14 +892,31 @@ public class ToolExecutionExecutor {
|
|||||||
// not yet migrated to ToolContext keep working unchanged.
|
// not yet migrated to ToolContext keep working unchanged.
|
||||||
ToolExecutionContext.set(pc.conversationId, pc.requesterId, pc.workspaceBasePath);
|
ToolExecutionContext.set(pc.conversationId, pc.requesterId, pc.workspaceBasePath);
|
||||||
String result;
|
String result;
|
||||||
|
String progressToken = null;
|
||||||
try {
|
try {
|
||||||
ChatOrigin runtimeOrigin = pc.origin != null ? pc.origin : ChatOrigin.EMPTY;
|
ChatOrigin runtimeOrigin = pc.origin != null ? pc.origin : ChatOrigin.EMPTY;
|
||||||
runtimeOrigin = runtimeOrigin
|
runtimeOrigin = runtimeOrigin
|
||||||
.withConversationId(pc.conversationId)
|
.withConversationId(pc.conversationId)
|
||||||
.withWorkspace(runtimeOrigin.workspaceId(), pc.workspaceBasePath);
|
.withWorkspace(runtimeOrigin.workspaceId(), pc.workspaceBasePath);
|
||||||
ToolContext toolContext = runtimeOrigin.toToolContext();
|
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);
|
result = pc.callback.call(pc.arguments, toolContext);
|
||||||
} finally {
|
} finally {
|
||||||
|
if (progressToken != null) {
|
||||||
|
progressContext.remove(progressToken);
|
||||||
|
progressContext.removeSnapshot(pc.conversationId, pc.toolCall.id());
|
||||||
|
}
|
||||||
ToolExecutionContext.clear();
|
ToolExecutionContext.clear();
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@ -3,10 +3,13 @@ package vip.mate.channel.web;
|
|||||||
import com.fasterxml.jackson.databind.ObjectMapper;
|
import com.fasterxml.jackson.databind.ObjectMapper;
|
||||||
import jakarta.annotation.PreDestroy;
|
import jakarta.annotation.PreDestroy;
|
||||||
import lombok.extern.slf4j.Slf4j;
|
import lombok.extern.slf4j.Slf4j;
|
||||||
|
import org.springframework.beans.factory.annotation.Autowired;
|
||||||
import org.springframework.beans.factory.annotation.Value;
|
import org.springframework.beans.factory.annotation.Value;
|
||||||
|
import org.springframework.context.ApplicationContext;
|
||||||
import org.springframework.stereotype.Component;
|
import org.springframework.stereotype.Component;
|
||||||
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
|
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
|
||||||
import reactor.core.Disposable;
|
import reactor.core.Disposable;
|
||||||
|
import vip.mate.tool.mcp.runtime.McpProgressContext;
|
||||||
import vip.mate.workspace.conversation.model.MessageContentPart;
|
import vip.mate.workspace.conversation.model.MessageContentPart;
|
||||||
|
|
||||||
import java.io.IOException;
|
import java.io.IOException;
|
||||||
@ -106,6 +109,9 @@ public class ChatStreamTracker {
|
|||||||
@Value("${mateclaw.stream.heartbeat.tool-sec:5}")
|
@Value("${mateclaw.stream.heartbeat.tool-sec:5}")
|
||||||
private int heartbeatToolSec = 5;
|
private int heartbeatToolSec = 5;
|
||||||
|
|
||||||
|
@Autowired
|
||||||
|
private ApplicationContext applicationContext;
|
||||||
|
|
||||||
public ChatStreamTracker(ObjectMapper objectMapper) {
|
public ChatStreamTracker(ObjectMapper objectMapper) {
|
||||||
this.objectMapper = objectMapper;
|
this.objectMapper = objectMapper;
|
||||||
}
|
}
|
||||||
@ -582,6 +588,15 @@ public class ChatStreamTracker {
|
|||||||
* early-return remains.
|
* early-return remains.
|
||||||
*/
|
*/
|
||||||
public void broadcast(String conversationId, String eventName, String jsonData) {
|
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);
|
RunState state = runs.get(conversationId);
|
||||||
|
|
||||||
boolean isDone = "done".equals(eventName);
|
boolean isDone = "done".equals(eventName);
|
||||||
@ -652,18 +667,23 @@ public class ChatStreamTracker {
|
|||||||
}
|
}
|
||||||
|
|
||||||
synchronized (state.lock) {
|
synchronized (state.lock) {
|
||||||
long id = ++state.nextEventId;
|
if (!skipBuffer) {
|
||||||
SseEvent event = new SseEvent(id, eventName, jsonData);
|
long id = ++state.nextEventId;
|
||||||
state.buffer.add(event);
|
SseEvent event = new SseEvent(id, eventName, jsonData);
|
||||||
// buffer 容量保护:超出上限时优先丢弃 thinking_delta(占比最大且非关键)
|
state.buffer.add(event);
|
||||||
if (state.buffer.size() > MAX_BUFFER_SIZE) {
|
if (state.buffer.size() > MAX_BUFFER_SIZE) {
|
||||||
trimBuffer(state.buffer);
|
trimBuffer(state.buffer);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
Iterator<SseEmitter> it = state.subscribers.iterator();
|
Iterator<SseEmitter> it = state.subscribers.iterator();
|
||||||
while (it.hasNext()) {
|
while (it.hasNext()) {
|
||||||
SseEmitter emitter = it.next();
|
SseEmitter emitter = it.next();
|
||||||
try {
|
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) {
|
} catch (IOException | IllegalStateException e) {
|
||||||
log.debug("Removing dead subscriber for {}: {}", conversationId, e.getMessage());
|
log.debug("Removing dead subscriber for {}: {}", conversationId, e.getMessage());
|
||||||
it.remove();
|
it.remove();
|
||||||
@ -695,6 +715,13 @@ public class ChatStreamTracker {
|
|||||||
* @param data 事件载荷,将被 Jackson 序列化为 JSON
|
* @param data 事件载荷,将被 Jackson 序列化为 JSON
|
||||||
*/
|
*/
|
||||||
public void broadcastObject(String conversationId, String eventName, Object data) {
|
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;
|
String json;
|
||||||
try {
|
try {
|
||||||
json = objectMapper.writeValueAsString(data);
|
json = objectMapper.writeValueAsString(data);
|
||||||
@ -702,7 +729,32 @@ public class ChatStreamTracker {
|
|||||||
log.warn("Failed to serialize broadcast data for event {}: {}", eventName, e.getMessage());
|
log.warn("Failed to serialize broadcast data for event {}: {}", eventName, e.getMessage());
|
||||||
json = "{\"error\":\"serialization_failed\"}";
|
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
|
// Without this, async_task_completed fired after `done` would be silently
|
||||||
// dropped, leaving the chat UI stuck on the "正在生成中" placeholder.
|
// dropped, leaving the chat UI stuck on the "正在生成中" placeholder.
|
||||||
state.subscribers.add(emitter);
|
state.subscribers.add(emitter);
|
||||||
|
|
||||||
|
// Deliver MCP progress snapshots on reconnect (progress events skip buffer replay)
|
||||||
|
sendProgressSnapshots(conversationId, emitter);
|
||||||
|
|
||||||
if (state.done) {
|
if (state.done) {
|
||||||
log.info("[SSE] Replayed {} buffered events; emitter stays subscribed for late async events: {}",
|
log.info("[SSE] Replayed {} buffered events; emitter stays subscribed for late async events: {}",
|
||||||
state.buffer.size(), conversationId);
|
state.buffer.size(), conversationId);
|
||||||
|
|||||||
@ -77,7 +77,12 @@ public final class IdentityForwardingToolCallback implements ToolCallback {
|
|||||||
return delegate;
|
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)
|
return identityService.resolve(toolContext, audience)
|
||||||
.map(i -> withClaim(toolInput, i.key(), i.value()))
|
.map(i -> withClaim(toolInput, i.key(), i.value()))
|
||||||
.orElse(toolInput);
|
.orElse(toolInput);
|
||||||
|
|||||||
@ -1,6 +1,7 @@
|
|||||||
package vip.mate.tool.mcp.runtime;
|
package vip.mate.tool.mcp.runtime;
|
||||||
|
|
||||||
import cn.hutool.json.JSONUtil;
|
import cn.hutool.json.JSONUtil;
|
||||||
|
import com.fasterxml.jackson.databind.ObjectMapper;
|
||||||
import io.modelcontextprotocol.client.McpClient;
|
import io.modelcontextprotocol.client.McpClient;
|
||||||
import io.modelcontextprotocol.client.McpSyncClient;
|
import io.modelcontextprotocol.client.McpSyncClient;
|
||||||
import io.modelcontextprotocol.client.transport.ServerParameters;
|
import io.modelcontextprotocol.client.transport.ServerParameters;
|
||||||
@ -67,13 +68,21 @@ public class McpClientManager {
|
|||||||
|
|
||||||
private final McpIdentityForwardService identityForwardService;
|
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. */
|
/** serverId -> server name, captured at build time for identity-forward opt-in matching. */
|
||||||
private final ConcurrentHashMap<Long, String> serverNames = new ConcurrentHashMap<>();
|
private final ConcurrentHashMap<Long, String> serverNames = new ConcurrentHashMap<>();
|
||||||
|
|
||||||
public McpClientManager(ApplicationEventPublisher eventPublisher,
|
public McpClientManager(ApplicationEventPublisher eventPublisher,
|
||||||
McpIdentityForwardService identityForwardService) {
|
McpIdentityForwardService identityForwardService,
|
||||||
|
McpProgressContext progressContext,
|
||||||
|
ObjectMapper objectMapper) {
|
||||||
this.eventPublisher = eventPublisher;
|
this.eventPublisher = eventPublisher;
|
||||||
this.identityForwardService = identityForwardService;
|
this.identityForwardService = identityForwardService;
|
||||||
|
this.progressContext = progressContext;
|
||||||
|
this.objectMapper = objectMapper;
|
||||||
}
|
}
|
||||||
|
|
||||||
/** serverId -> connection result info */
|
/** serverId -> connection result info */
|
||||||
@ -201,7 +210,8 @@ public class McpClientManager {
|
|||||||
McpIdentityForwardService idSvc =
|
McpIdentityForwardService idSvc =
|
||||||
identityForwardService.forwardsTo(serverId, serverName) ? identityForwardService : null;
|
identityForwardService.forwardsTo(serverId, serverName) ? identityForwardService : null;
|
||||||
String audience = idSvc != null ? identityForwardService.audienceFor(serverId, serverName) : 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);
|
lastGoodCallbacks.put(serverId, wrapped);
|
||||||
allCallbacks.addAll(wrapped);
|
allCallbacks.addAll(wrapped);
|
||||||
continue;
|
continue;
|
||||||
@ -253,7 +263,7 @@ public class McpClientManager {
|
|||||||
* real {@link McpSyncClient}.
|
* real {@link McpSyncClient}.
|
||||||
*/
|
*/
|
||||||
static List<ToolCallback> wrapServerCallbacks(long serverId, ToolCallback[] cbs) {
|
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
|
* @param serverName human-readable MCP server name; forwarded into each
|
||||||
* {@link PrefixedNameToolCallback} so the tool description is tagged
|
* {@link PrefixedNameToolCallback} so the tool description is tagged
|
||||||
* {@code [MCP server: <name>]}. May be {@code null} when unknown.
|
* {@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,
|
static List<ToolCallback> wrapServerCallbacks(long serverId, ToolCallback[] cbs,
|
||||||
McpIdentityForwardService identitySvc, String audience,
|
McpIdentityForwardService identitySvc, String audience,
|
||||||
String serverName) {
|
String serverName,
|
||||||
|
McpSyncClient mcpClient, ObjectMapper objectMapper) {
|
||||||
List<String> rawNames = new ArrayList<>(cbs.length);
|
List<String> rawNames = new ArrayList<>(cbs.length);
|
||||||
for (ToolCallback cb : cbs) {
|
for (ToolCallback cb : cbs) {
|
||||||
rawNames.add(cb.getToolDefinition() != null ? cb.getToolDefinition().name() : null);
|
rawNames.add(cb.getToolDefinition() != null ? cb.getToolDefinition().name() : null);
|
||||||
@ -298,6 +313,13 @@ public class McpClientManager {
|
|||||||
ToolCallback inner = identitySvc != null
|
ToolCallback inner = identitySvc != null
|
||||||
? new IdentityForwardingToolCallback(cb, identitySvc, audience)
|
? new IdentityForwardingToolCallback(cb, identitySvc, audience)
|
||||||
: cb;
|
: 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));
|
out.add(new PrefixedNameToolCallback(d.prefixedName(), inner, serverName));
|
||||||
}
|
}
|
||||||
return out;
|
return out;
|
||||||
@ -418,6 +440,29 @@ public class McpClientManager {
|
|||||||
Long serverId = server.getId();
|
Long serverId = server.getId();
|
||||||
spec.toolsChangeConsumer(tools ->
|
spec.toolsChangeConsumer(tools ->
|
||||||
eventPublisher.publishEvent(new McpServerChangedEvent("mcp-tools-changed:" + serverId)));
|
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();
|
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 {
|
void staleListToolsServesSnapshotAndRequestsReconnect() throws Exception {
|
||||||
ApplicationEventPublisher publisher = mock(ApplicationEventPublisher.class);
|
ApplicationEventPublisher publisher = mock(ApplicationEventPublisher.class);
|
||||||
McpClientManager manager = new McpClientManager(publisher,
|
McpClientManager manager = new McpClientManager(publisher,
|
||||||
new McpIdentityForwardService(new McpIdentityForwardProperties()));
|
new McpIdentityForwardService(new McpIdentityForwardProperties()),
|
||||||
|
null, null);
|
||||||
|
|
||||||
// A client whose connection went stale: every listTools() throws.
|
// A client whose connection went stale: every listTools() throws.
|
||||||
McpSyncClient deadClient = mock(McpSyncClient.class);
|
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 isSuccess = computed(() => props.segment.status === 'completed' && props.segment.toolSuccess !== false)
|
||||||
const isError = computed(() => props.segment.status === 'error' || props.segment.toolSuccess === false)
|
const isError = computed(() => props.segment.status === 'error' || props.segment.toolSuccess === false)
|
||||||
const isRunning = computed(() => props.segment.status === 'running')
|
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.
|
// A delegation flagged by the heartbeat watchdog as making no progress.
|
||||||
const isStalled = computed(() => isDelegation.value && isRunning.value && !!props.segment.delegationStale)
|
const isStalled = computed(() => isDelegation.value && isRunning.value && !!props.segment.delegationStale)
|
||||||
// Fire-and-forget delegation: runs detached, result comes via task_output later.
|
// 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">
|
<div class="seg-tool__header" @click="hasBody ? (expanded = !expanded) : null">
|
||||||
<span class="seg-tool__status">
|
<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-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="isRunning" class="is-loading" :size="13"><Loading /></el-icon>
|
||||||
<el-icon v-else-if="isSuccess" :size="13"><Select /></el-icon>
|
<el-icon v-else-if="isSuccess" :size="13"><Select /></el-icon>
|
||||||
<el-icon v-else :size="13"><CloseBold /></el-icon>
|
<el-icon v-else :size="13"><CloseBold /></el-icon>
|
||||||
@ -147,6 +150,14 @@ const detailStatus = computed<'running' | 'completed' | 'error'>(() => {
|
|||||||
><ArrowDown /></el-icon>
|
><ArrowDown /></el-icon>
|
||||||
</span>
|
</span>
|
||||||
</div>
|
</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">
|
<Transition name="seg-slide">
|
||||||
<div v-if="expanded && hasBody" class="seg-tool__body">
|
<div v-if="expanded && hasBody" class="seg-tool__body">
|
||||||
<!-- Nested subagent timeline (delegation segments) -->
|
<!-- Nested subagent timeline (delegation segments) -->
|
||||||
@ -385,4 +396,39 @@ const detailStatus = computed<'running' | 'completed' | 'error'>(() => {
|
|||||||
opacity: 0;
|
opacity: 0;
|
||||||
transform: translateY(-4px);
|
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>
|
</style>
|
||||||
|
|||||||
@ -958,6 +958,21 @@ export function useChat(options: UseChatOptions): UseChatReturn {
|
|||||||
stream.on('tool_call_started', handleToolCallStarted)
|
stream.on('tool_call_started', handleToolCallStarted)
|
||||||
stream.on('tool_call_completed', handleToolCallCompleted)
|
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 =====
|
// ===== Browser action events =====
|
||||||
|
|
||||||
stream.on('browser_action', (data) => {
|
stream.on('browser_action', (data) => {
|
||||||
|
|||||||
@ -268,6 +268,12 @@ export interface MessageSegment {
|
|||||||
supersededBySegmentId?: string
|
supersededBySegmentId?: string
|
||||||
/** Machine-readable reason for superseding this segment. */
|
/** Machine-readable reason for superseding this segment. */
|
||||||
supersededReason?: string
|
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. */
|
/** A file artifact generated by a tool during the turn, surfaced in the run-overview rail. */
|
||||||
|
|||||||
Loading…
Reference in New Issue
Block a user