Ver código fonte

feat:增加mcp server日志打印

Jerry Liang 2 meses atrás
pai
commit
983dabaf0b

+ 2 - 0
service/cmd/mcp/main.go

@@ -17,6 +17,7 @@ import (
 	minioUtil "fwh-toolbox-server/utils/minio"
 
 	"fwh-toolbox-server/cmd/mcp/tools"
+	"fwh-toolbox-server/cmd/mcp/tools/middleware"
 	"fwh-toolbox-server/cmd/mcp/tools/storage"
 )
 
@@ -63,6 +64,7 @@ func main() {
 		serverVersion,
 		server.WithToolCapabilities(false),
 		server.WithRecovery(),
+		server.WithToolHandlerMiddleware(middleware.LoggingMiddleware()),
 		server.WithInstructions(serverInstructions),
 	)
 

+ 165 - 0
service/cmd/mcp/tools/middleware/logging.go

@@ -0,0 +1,165 @@
+// 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)"
+}

+ 140 - 0
service/cmd/mcp/tools/middleware/logging_test.go

@@ -0,0 +1,140 @@
+package middleware
+
+import (
+	"context"
+	"strings"
+	"testing"
+
+	"github.com/mark3labs/mcp-go/mcp"
+	"github.com/mark3labs/mcp-go/server"
+	"go.uber.org/zap"
+	"go.uber.org/zap/zapcore"
+	"go.uber.org/zap/zaptest/observer"
+)
+
+// setupObservedLogger 替换全局 logger 为可观测 logger,返回观测器与还原函数。
+func setupObservedLogger(t *testing.T) (*observer.ObservedLogs, func()) {
+	t.Helper()
+	core, obs := observer.New(zapcore.InfoLevel)
+	orig := zap.L()
+	zap.ReplaceGlobals(zap.New(core))
+	return obs, func() { zap.ReplaceGlobals(orig) }
+}
+
+func newReq(name string, args map[string]any) mcp.CallToolRequest {
+	var req mcp.CallToolRequest
+	req.Params.Name = name
+	req.Params.Arguments = args
+	return req
+}
+
+func wrap(next server.ToolHandlerFunc) server.ToolHandlerFunc {
+	return LoggingMiddleware()(next)
+}
+
+func TestLogging_Success(t *testing.T) {
+	obs, restore := setupObservedLogger(t)
+	defer restore()
+
+	h := wrap(func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+		return mcp.NewToolResultText(`{"ok":true}`), nil
+	})
+	_, err := h(context.Background(), newReq("service_site_query_info", map[string]any{"sn": "SN123"}))
+	if err != nil {
+		t.Fatalf("unexpected err: %v", err)
+	}
+
+	logs := obs.All()
+	if len(logs) != 2 {
+		t.Fatalf("want 2 log entries (start+success), got %d: %+v", len(logs), logs)
+	}
+	if logs[0].Message != "[tool] call start" {
+		t.Errorf("first log = %q, want start", logs[0].Message)
+	}
+	if logs[1].Message != "[tool] call success" || logs[1].Level != zapcore.InfoLevel {
+		t.Errorf("second log = %q level=%v, want success/info", logs[1].Message, logs[1].Level)
+	}
+	// status 字段
+	if got := logs[1].ContextMap()["status"]; got != "success" {
+		t.Errorf("status = %v, want success", got)
+	}
+}
+
+func TestLogging_BusinessError(t *testing.T) {
+	obs, restore := setupObservedLogger(t)
+	defer restore()
+
+	h := wrap(func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+		return mcp.NewToolResultError("查询站点信息失败: not found"), nil
+	})
+	if _, err := h(context.Background(), newReq("service_site_query_info", nil)); err != nil {
+		t.Fatalf("unexpected err: %v", err)
+	}
+
+	last := obs.All()[len(obs.All())-1]
+	if last.Message != "[tool] call business_error" || last.Level != zapcore.WarnLevel {
+		t.Fatalf("last log = %q level=%v, want business_error/warn", last.Message, last.Level)
+	}
+	if got := last.ContextMap()["status"]; got != "business_error" {
+		t.Errorf("status = %v, want business_error", got)
+	}
+}
+
+func TestLogging_GoError(t *testing.T) {
+	obs, restore := setupObservedLogger(t)
+	defer restore()
+
+	h := wrap(func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+		return nil, context.DeadlineExceeded
+	})
+	if _, err := h(context.Background(), newReq("data_mongo_query", nil)); err == nil {
+		t.Fatal("expected error to propagate")
+	}
+
+	last := obs.All()[len(obs.All())-1]
+	if last.Message != "[tool] call error" || last.Level != zapcore.ErrorLevel {
+		t.Fatalf("last log = %q level=%v, want error/error", last.Message, last.Level)
+	}
+}
+
+func TestLogging_Panic(t *testing.T) {
+	obs, restore := setupObservedLogger(t)
+	defer restore()
+
+	h := wrap(func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+		panic("boom")
+	})
+
+	defer func() {
+		if r := recover(); r == nil {
+			t.Fatal("expected panic to be re-thrown")
+		}
+		last := obs.All()[len(obs.All())-1]
+		if last.Message != "[tool] call panic" || last.Level != zapcore.ErrorLevel {
+			t.Fatalf("last log = %q level=%v, want panic/error", last.Message, last.Level)
+		}
+		if _, ok := last.ContextMap()["stack"]; !ok {
+			t.Error("panic log missing stack field")
+		}
+	}()
+
+	_, _ = h(context.Background(), newReq("service_site_query_info", nil))
+}
+
+// TestLogging_ParamMasking 验证参数中的敏感字段被脱敏后再写日志。
+func TestLogging_ParamMasking(t *testing.T) {
+	obs, restore := setupObservedLogger(t)
+	defer restore()
+
+	h := wrap(func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+		return mcp.NewToolResultText("ok"), nil
+	})
+	// phone 是常见 PII,MaskJSONBytes 应对其脱敏
+	_, _ = h(context.Background(), newReq("service_site_query_info", map[string]any{"phone": "13800138000"}))
+
+	start := obs.All()[0]
+	params, _ := start.ContextMap()["params"].(string)
+	if strings.Contains(params, "13800138000") {
+		t.Errorf("params log leaked raw phone: %q", params)
+	}
+}