FranklinWH-Toolbox/service/agentmcp/kit/middleware.go

package kit

import (
	"context"
	"encoding/json"
	"fmt"
	"runtime/debug"
	"strings"
	"sync/atomic"
	"time"

	"github.com/mark3labs/mcp-go/mcp"
	"github.com/mark3labs/mcp-go/server"
	"go.uber.org/zap"
)

const (
	// DefaultTimeout 普通工具单次调用超时
	DefaultTimeout = 60 * time.Second
	// ExportTimeout 导出工具单次调用超时
	ExportTimeout = 300 * time.Second
	// QueueWait 并发已满时的最长排队时间
	QueueWait = 30 * time.Second

	maxParamsLogBytes = 2 * 1024
)

var callSeq int64

// Middleware 统一的工具调用中间件：并发上限、超时、panic 恢复、调用日志（不含调用人）。
//
// 数据访问函数大多不接收 context，超时后无法真正中止上游请求：处理函数在后台继续执行，
// 直到返回才释放并发名额，从而保证对 FWH 平台和数据库的并发压力不超过上限。
func Middleware(log *zap.Logger, maxConcurrency int) server.ToolHandlerMiddleware {
	if maxConcurrency <= 0 {
		maxConcurrency = 8
	}
	sem := make(chan struct{}, maxConcurrency)

	return func(next server.ToolHandlerFunc) server.ToolHandlerFunc {
		return func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
			name := req.Params.Name
			seq := atomic.AddInt64(&callSeq, 1)
			started := time.Now()
			fields := []zap.Field{zap.Int64("seq", seq), zap.String("tool", name)}

			// 1. 并发名额
			waitCtx, cancelWait := context.WithTimeout(ctx, QueueWait)
			select {
			case sem <- struct{}{}:
				cancelWait()
			case <-waitCtx.Done():
				cancelWait()
				log.Warn("[agent-mcp] 排队超时", fields...)
				return mcp.NewToolResultError("服务繁忙，请稍后重试。"), nil
			}

			timeout := DefaultTimeout
			if strings.HasSuffix(name, "_export") {
				timeout = ExportTimeout
			}
			callCtx, cancel := context.WithTimeout(ctx, timeout)
			defer cancel()

			log.Info("[agent-mcp] 调用", append(fields, zap.String("params", paramsSummary(req)))...)

			type outcome struct {
				res *mcp.CallToolResult
				err error
			}
			done := make(chan outcome, 1)
			go func() {
				defer func() { <-sem }()
				defer func() {
					if r := recover(); r != nil {
						log.Error("[agent-mcp] panic", append(fields, zap.Any("panic", r), zap.String("stack", string(debug.Stack())))...)
						done <- outcome{res: mcp.NewToolResultError("内部错误，请稍后重试。")}
					}
				}()
				res, err := next(callCtx, req)
				done <- outcome{res: res, err: err}
			}()

			select {
			case o := <-done:
				size, isErr := resultInfo(o.res)
				l := log.Info
				if o.err != nil {
					l = log.Error
				} else if isErr {
					l = log.Warn
				}
				l("[agent-mcp] 完成", append(fields,
					zap.Int64("duration_ms", time.Since(started).Milliseconds()),
					zap.Int("result_size", size),
					zap.Bool("is_error", isErr),
					zap.Error(o.err),
				)...)
				if o.err != nil {
					// 不把 Go error 透传为协议错误，统一给模型一条可读信息
					return mcp.NewToolResultError(fmt.Sprintf("工具执行失败：%v", o.err)), nil
				}
				return o.res, nil
			case <-callCtx.Done():
				log.Warn("[agent-mcp] 超时", append(fields, zap.Duration("timeout", timeout))...)
				return mcp.NewToolResultError(fmt.Sprintf("查询超时（%d 秒），请缩小时间范围或稍后重试。", int(timeout.Seconds()))), nil
			}
		}
	}
}

func paramsSummary(req mcp.CallToolRequest) string {
	b, err := json.Marshal(req.GetArguments())
	if err != nil {
		return ""
	}
	if len(b) > maxParamsLogBytes {
		return string(b[:maxParamsLogBytes]) + "...(truncated)"
	}
	return string(b)
}

func resultInfo(r *mcp.CallToolResult) (size int, isErr bool) {
	if r == nil {
		return 0, false
	}
	for _, c := range r.Content {
		switch t := c.(type) {
		case mcp.TextContent:
			size += len(t.Text)
		case mcp.EmbeddedResource:
			if tr, ok := t.Resource.(mcp.TextResourceContents); ok {
				size += len(tr.Text)
			}
		}
	}
	return size, r.IsError
}