feat(memory): P1 长期记忆升级 —— 异步攒批 Consolidate + 软删 + importance/last_seen #1
@@ -47,6 +47,8 @@ class McpGateway:
|
|||||||
self.interpreter = CodeInterpreter() # Docker 隔离沙箱
|
self.interpreter = CodeInterpreter() # Docker 隔离沙箱
|
||||||
self._nc: nats.NATS | None = None
|
self._nc: nats.NATS | None = None
|
||||||
self._tools: dict[str, callable] = {}
|
self._tools: dict[str, callable] = {}
|
||||||
|
# 单实例工具并发上限(垂直扩);队列组多副本再叠加水平扩。
|
||||||
|
self._sem = asyncio.Semaphore(int(os.getenv("MCP_TOOL_CONCURRENCY", "16")))
|
||||||
|
|
||||||
async def register_tools(self) -> None:
|
async def register_tools(self) -> None:
|
||||||
"""注册工具名 → 处理协程。"""
|
"""注册工具名 → 处理协程。"""
|
||||||
@@ -77,25 +79,30 @@ class McpGateway:
|
|||||||
await asyncio.Event().wait()
|
await asyncio.Event().wait()
|
||||||
|
|
||||||
async def _on_call(self, msg) -> None:
|
async def _on_call(self, msg) -> None:
|
||||||
"""解析 ToolCall → 路由到工具 → Respond ToolResult。"""
|
"""立即派发到独立 task 并发处理,不阻塞订阅消息循环(nats-py 单订阅回调本是串行 await)。"""
|
||||||
try:
|
asyncio.create_task(self._handle(msg))
|
||||||
req = json.loads(msg.data)
|
|
||||||
except Exception as e: # noqa: BLE001
|
async def _handle(self, msg) -> None:
|
||||||
await self._reply(msg, ok=False, error=f"bad tool call: {e}")
|
"""解析 ToolCall → 路由到工具 → Respond ToolResult(限并发)。"""
|
||||||
return
|
async with self._sem:
|
||||||
tool = req.get("tool", "")
|
try:
|
||||||
task_id = req.get("task_id", "")
|
req = json.loads(msg.data)
|
||||||
args = req.get("args") or {}
|
except Exception as e: # noqa: BLE001
|
||||||
log.info("[mcp_py] tool=%s task=%s args=%s", tool, task_id, args)
|
await self._reply(msg, ok=False, error=f"bad tool call: {e}")
|
||||||
fn = self._tools.get(tool)
|
return
|
||||||
if fn is None:
|
tool = req.get("tool", "")
|
||||||
await self._reply(msg, ok=False, error=f"unknown tool: {tool}")
|
task_id = req.get("task_id", "")
|
||||||
return
|
args = req.get("args") or {}
|
||||||
try:
|
log.info("[mcp_py] tool=%s task=%s args=%s", tool, task_id, args)
|
||||||
content = await fn(args)
|
fn = self._tools.get(tool)
|
||||||
await self._reply(msg, ok=True, content=content)
|
if fn is None:
|
||||||
except Exception as e: # noqa: BLE001
|
await self._reply(msg, ok=False, error=f"unknown tool: {tool}")
|
||||||
await self._reply(msg, ok=False, error=str(e))
|
return
|
||||||
|
try:
|
||||||
|
content = await fn(args)
|
||||||
|
await self._reply(msg, ok=True, content=content)
|
||||||
|
except Exception as e: # noqa: BLE001
|
||||||
|
await self._reply(msg, ok=False, error=str(e))
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
async def _reply(msg, *, ok: bool, content: str = "", error: str = "") -> None:
|
async def _reply(msg, *, ok: bool, content: str = "", error: str = "") -> None:
|
||||||
|
|||||||
+40
-19
@@ -233,27 +233,48 @@ func (b *Bus) CallTool(ctx context.Context, subject string, call *contract.ToolC
|
|||||||
// ToolHandler 处理一次工具调用并返回结果。
|
// ToolHandler 处理一次工具调用并返回结果。
|
||||||
type ToolHandler func(ctx context.Context, call *contract.ToolCall) *contract.ToolResult
|
type ToolHandler func(ctx context.Context, call *contract.ToolCall) *contract.ToolResult
|
||||||
|
|
||||||
// ServeTool 以队列组订阅工具主题(可用通配 sundynix.tools.go.>),
|
// toolConcurrency 返回单实例工具并发处理上限(env MCP_TOOL_CONCURRENCY,默认 16)。
|
||||||
// 对每个请求调用 h 并 Respond,队列组内多副本自动负载均衡。
|
func toolConcurrency() int {
|
||||||
// 返回的 unsub 用于退订。
|
if v := os.Getenv("MCP_TOOL_CONCURRENCY"); v != "" {
|
||||||
|
if n, err := strconv.Atoi(v); err == nil && n > 0 {
|
||||||
|
return n
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return 16
|
||||||
|
}
|
||||||
|
|
||||||
|
// ServeTool 以队列组订阅工具主题(可用通配 sundynix.tools.go.>),队列组内多副本自动负载均衡(水平扩)。
|
||||||
|
// 单实例内每个请求分发到独立 goroutine 并发处理(垂直扩,上限 MCP_TOOL_CONCURRENCY)——
|
||||||
|
// 回调立即返回不阻塞 NATS 投递;信号量约束同时执行的 handler 数。返回的 unsub 用于退订。
|
||||||
func (b *Bus) ServeTool(subject, queue string, h ToolHandler) (unsub func() error, err error) {
|
func (b *Bus) ServeTool(subject, queue string, h ToolHandler) (unsub func() error, err error) {
|
||||||
|
sem := make(chan struct{}, toolConcurrency())
|
||||||
sub, err := b.nc.QueueSubscribe(subject, queue, func(m *nats.Msg) {
|
sub, err := b.nc.QueueSubscribe(subject, queue, func(m *nats.Msg) {
|
||||||
var call contract.ToolCall
|
// 立即起 goroutine 并返回,让 NATS 继续投递下一条(核心 NATS 单订阅回调本是串行)。
|
||||||
if err := json.Unmarshal(m.Data, &call); err != nil {
|
go func() {
|
||||||
respond(m, &contract.ToolResult{OK: false, Error: "bad tool call: " + err.Error()})
|
sem <- struct{}{} // 限并发:超额则在此短暂排队
|
||||||
return
|
defer func() {
|
||||||
}
|
<-sem
|
||||||
// 还原上游链路并开服务端 span(成为 dispatcher tool.call span 的子节点)。
|
if r := recover(); r != nil { // handler panic → 回错误结果,调用方不必干等超时
|
||||||
mctx := extractTrace(context.Background(), m.Header)
|
respond(m, &contract.ToolResult{OK: false, Error: fmt.Sprintf("tool panic: %v", r)})
|
||||||
mctx, span := tracer().Start(mctx, "tool.serve "+call.Tool,
|
}
|
||||||
trace.WithSpanKind(trace.SpanKindServer),
|
}()
|
||||||
trace.WithAttributes(attribute.String("sundynix.tool", call.Tool)))
|
var call contract.ToolCall
|
||||||
res := h(mctx, &call)
|
if err := json.Unmarshal(m.Data, &call); err != nil {
|
||||||
if res != nil {
|
respond(m, &contract.ToolResult{OK: false, Error: "bad tool call: " + err.Error()})
|
||||||
span.SetAttributes(attribute.Bool("sundynix.tool.ok", res.OK))
|
return
|
||||||
}
|
}
|
||||||
span.End()
|
// 还原上游链路并开服务端 span(成为 dispatcher tool.call span 的子节点)。
|
||||||
respond(m, res)
|
mctx := extractTrace(context.Background(), m.Header)
|
||||||
|
mctx, span := tracer().Start(mctx, "tool.serve "+call.Tool,
|
||||||
|
trace.WithSpanKind(trace.SpanKindServer),
|
||||||
|
trace.WithAttributes(attribute.String("sundynix.tool", call.Tool)))
|
||||||
|
defer span.End()
|
||||||
|
res := h(mctx, &call)
|
||||||
|
if res != nil {
|
||||||
|
span.SetAttributes(attribute.Bool("sundynix.tool.ok", res.OK))
|
||||||
|
}
|
||||||
|
respond(m, res)
|
||||||
|
}()
|
||||||
})
|
})
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, fmt.Errorf("serve tool %s: %w", subject, err)
|
return nil, fmt.Errorf("serve tool %s: %w", subject, err)
|
||||||
|
|||||||
@@ -254,3 +254,55 @@ func TestConcurrentConsume(t *testing.T) {
|
|||||||
t.Fatal("B 被阻塞的 A 卡住——并发消费未生效(仍是串行)")
|
t.Fatal("B 被阻塞的 A 卡住——并发消费未生效(仍是串行)")
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// TestConcurrentToolServe 验证单实例工具服务并发:一个慢工具调用不阻塞另一个快调用。
|
||||||
|
// 旧的串行 ServeTool 下本测试会超时失败。
|
||||||
|
func TestConcurrentToolServe(t *testing.T) {
|
||||||
|
url := startEmbeddedNATS(t)
|
||||||
|
|
||||||
|
srv, err := bus.Connect(url)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("server connect: %v", err)
|
||||||
|
}
|
||||||
|
defer srv.Close()
|
||||||
|
|
||||||
|
blockSlow := make(chan struct{})
|
||||||
|
unsub, err := srv.ServeTool(contract.SubjectToolsGoAll, contract.QueueToolsGo,
|
||||||
|
func(_ context.Context, call *contract.ToolCall) *contract.ToolResult {
|
||||||
|
if call.Tool == "slow" {
|
||||||
|
<-blockSlow // 模拟慢工具(如 RAG 嵌入)
|
||||||
|
}
|
||||||
|
return &contract.ToolResult{OK: true, Content: call.Tool}
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("serve tool: %v", err)
|
||||||
|
}
|
||||||
|
defer unsub()
|
||||||
|
defer close(blockSlow)
|
||||||
|
|
||||||
|
cli, err := bus.Connect(url)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("client connect: %v", err)
|
||||||
|
}
|
||||||
|
defer cli.Close()
|
||||||
|
|
||||||
|
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||||
|
defer cancel()
|
||||||
|
|
||||||
|
// 先发 slow(会阻塞在 handler 里),再发 fast——fast 应在 slow 仍阻塞时返回。
|
||||||
|
go func() {
|
||||||
|
_, _ = cli.CallTool(ctx, contract.ToolSubjectGo("slow"), &contract.ToolCall{Tool: "slow"})
|
||||||
|
}()
|
||||||
|
time.Sleep(300 * time.Millisecond)
|
||||||
|
|
||||||
|
fctx, fcancel := context.WithTimeout(ctx, 3*time.Second)
|
||||||
|
defer fcancel()
|
||||||
|
res, err := cli.CallTool(fctx, contract.ToolSubjectGo("fast"), &contract.ToolCall{Tool: "fast"})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("fast 工具被慢工具阻塞——单实例并发未生效: %v", err)
|
||||||
|
}
|
||||||
|
if res.Content != "fast" {
|
||||||
|
t.Fatalf("unexpected result: %q", res.Content)
|
||||||
|
}
|
||||||
|
t.Log("✓ 单实例工具并发生效:slow 仍阻塞时 fast 已返回")
|
||||||
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user