FranklinWH-Toolbox/service/cmd/mcp/tools/middleware/logging.go

// Package middleware 提供 MCP 工具调用的横切关注点（日志等）。
package middleware

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

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

	"fwh-toolbox-server/cmd/mcp/tools/helper"
)

// maxParamsLogBytes 参数写入日志时的最大字节数，超过则截断，避免日志过大。
const maxParamsLogBytes = 2 * 1024

// maxErrLogBytes 错误结果内容写入日志时的最大字节数。
const maxErrLogBytes = 1 * 1024

// requestSeq 为每次调用生成一个进程内自增的请求序号，用于关联同一次调用的多条日志。
var requestSeq int64

// LoggingMiddleware 返回一个 MCP 工具处理中间件，统一记录工具调用日志：
//
// 第 1 类（调用入口/出口）：
//   - 入口：tool、request_id、脱敏后的 params（截断）
//   - 出口：duration_ms、status、result_size / spilled
//
// 第 2 类（错误与异常）：
//   - handler 返回 Go error         -> level=error
//   - 业务错误结果（result.IsError）-> level=warn，附错误内容摘要
//   - panic                         -> level=error，附堆栈，并重新抛出（保持 WithRecovery 行为）
func LoggingMiddleware() server.ToolHandlerMiddleware {
	return func(next server.ToolHandlerFunc) server.ToolHandlerFunc {
		return func(ctx context.Context, req mcp.CallToolRequest) (result *mcp.CallToolResult, err error) {
			tool := req.Params.Name
			requestID := nextRequestID(tool)
			start := time.Now()

			// 入口日志
			zap.L().Info("[tool] call start",
				zap.String("request_id", requestID),
				zap.String("tool", tool),
				zap.String("params", maskedParams(req)),
			)

			// panic 兜底：记录后重新抛出，交由上层 WithRecovery 处理
			defer func() {
				if r := recover(); r != nil {
					zap.L().Error("[tool] call panic",
						zap.String("request_id", requestID),
						zap.String("tool", tool),
						zap.Int64("duration_ms", time.Since(start).Milliseconds()),
						zap.Any("panic", r),
						zap.String("stack", string(debug.Stack())),
					)
					panic(r)
				}
			}()

			result, err = next(ctx, req)
			durationMS := time.Since(start).Milliseconds()

			switch {
			case err != nil:
				// handler 返回 Go error（框架级/传输级错误）
				zap.L().Error("[tool] call error",
					zap.String("request_id", requestID),
					zap.String("tool", tool),
					zap.Int64("duration_ms", durationMS),
					zap.String("status", "error"),
					zap.Error(err),
				)
			case result != nil && result.IsError:
				// 业务错误结果（handler 返回 mcp.NewToolResultError）
				zap.L().Warn("[tool] call business_error",
					zap.String("request_id", requestID),
					zap.String("tool", tool),
					zap.Int64("duration_ms", durationMS),
					zap.String("status", "business_error"),
					zap.String("err", truncate(resultText(result), maxErrLogBytes)),
				)
			default:
				// 成功
				size, spilled := resultMeta(result)
				zap.L().Info("[tool] call success",
					zap.String("request_id", requestID),
					zap.String("tool", tool),
					zap.Int64("duration_ms", durationMS),
					zap.String("status", "success"),
					zap.Int("result_size", size),
					zap.Bool("spilled", spilled),
				)
			}

			return result, err
		}
	}
}

// nextRequestID 生成一个进程内唯一的请求 ID，格式：<tool>-<纳秒时间戳>-<序号>。
func nextRequestID(tool string) string {
	requestSeq++
	return fmt.Sprintf("%s-%d-%d", tool, time.Now().UnixNano(), requestSeq)
}

// maskedParams 将入参序列化并做 PII 脱敏，超长则截断。脱敏失败时不泄露原始值。
func maskedParams(req mcp.CallToolRequest) string {
	args := req.GetArguments()
	if len(args) == 0 {
		return "{}"
	}
	b, err := json.Marshal(args)
	if err != nil {
		return "<marshal error>"
	}
	masked, err := helper.MaskJSONBytes(b)
	if err != nil {
		// fail-closed：脱敏失败则不输出原始参数
		return "<masked: sanitize failed>"
	}
	return truncate(string(masked), maxParamsLogBytes)
}

// resultText 拼接结果中的文本内容。
func resultText(result *mcp.CallToolResult) string {
	if result == nil {
		return ""
	}
	var sb strings.Builder
	for _, c := range result.Content {
		if tc, ok := c.(mcp.TextContent); ok {
			sb.WriteString(tc.Text)
		}
	}
	return sb.String()
}

// resultMeta 返回结果文本字节数，以及是否为落盘结果（spilled）。
func resultMeta(result *mcp.CallToolResult) (size int, spilled bool) {
	text := resultText(result)
	size = len(text)
	// 落盘结果的返回体是包含 "spilled": true 的 JSON 摘要
	if strings.Contains(text, `"spilled": true`) || strings.Contains(text, `"spilled":true`) {
		spilled = true
	}
	return size, spilled
}

// truncate 按字节截断日志字段，避免单条日志过大。
func truncate(s string, max int) string {
	if len(s) <= max {
		return s
	}
	end := max
	for end > 0 && s[end]&0xC0 == 0x80 { // 回退到 UTF-8 rune 边界
		end--
	}
	return s[:end] + "...(truncated)"
}