66 lines
1.7 KiB
Go
66 lines
1.7 KiB
Go
package server
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"net/http"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/modelcontextprotocol/go-sdk/mcp"
|
|
log "github.com/nacorid/logger"
|
|
)
|
|
|
|
type LoggingMiddleware struct{}
|
|
|
|
func NewLoggingMiddleware() *LoggingMiddleware {
|
|
return &LoggingMiddleware{}
|
|
}
|
|
|
|
// OnCall logs native MCP tool requests streamed over SSE/WebSockets
|
|
func (lm *LoggingMiddleware) OnCall(next mcp.MethodHandler) mcp.MethodHandler {
|
|
return func(ctx context.Context, method string, request mcp.Request) (mcp.Result, error) {
|
|
if method == "tools/call" {
|
|
start := time.Now()
|
|
|
|
result, err := next(ctx, method, request)
|
|
|
|
duration := fmt.Sprintf("%.2fms", float64(time.Since(start))/float64(time.Millisecond))
|
|
|
|
toolName := "unknown"
|
|
if callReq, ok := request.(*mcp.CallToolRequest); ok && callReq.Params != nil {
|
|
toolName = callReq.Params.Name
|
|
}
|
|
|
|
log.InfoWithContext(ctx, "MCP Tool Processed",
|
|
"transport", "mcp_native",
|
|
"tool", toolName,
|
|
"duration", duration,
|
|
"success", err == nil,
|
|
)
|
|
return result, err
|
|
}
|
|
return next(ctx, method, request)
|
|
}
|
|
}
|
|
|
|
func (lm *LoggingMiddleware) OnCallHTTP(next http.Handler) http.Handler {
|
|
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
start := time.Now()
|
|
|
|
ipAddr := r.RemoteAddr
|
|
if forwarded := r.Header.Get("X-Forwarded-For"); forwarded != "" {
|
|
ipAddr = strings.TrimSpace(strings.Split(forwarded, ",")[0])
|
|
}
|
|
|
|
next.ServeHTTP(w, r)
|
|
|
|
duration := fmt.Sprintf("%.2fms", float64(time.Since(start))/float64(time.Millisecond))
|
|
log.InfoWithContext(r.Context(), "REST Tool Processed",
|
|
"transport", "rest",
|
|
"path", r.URL.Path,
|
|
"duration", duration,
|
|
"remote_ip", ipAddr,
|
|
)
|
|
})
|
|
}
|