diff --git a/controller/mcp_call_latest.go b/controller/mcp_call_latest.go new file mode 100644 index 0000000000..21359a965f --- /dev/null +++ b/controller/mcp_call_latest.go @@ -0,0 +1,276 @@ +package controller + +import ( + "context" + "net/http" + "strings" + "time" + + "github.com/Laisky/errors/v2" + gmw "github.com/Laisky/gin-middlewares/v7" + "github.com/gin-gonic/gin" + + "github.com/Laisky/one-api/common/config" + "github.com/Laisky/one-api/model" + "github.com/Laisky/one-api/relay/mcp" +) + +type modernMCPToolCallPlan struct { + userID int + serverByID map[int]*model.MCPServer + candidates []mcp.ToolCandidate +} + +// executeModernMCPTool validates one modern call against its lossless descriptor and executes it. +// +// Parameters: +// - ctx: The request context controls database and upstream work. +// - c: The Gin context contains the authenticated user and request headers. +// - params: The call parameters contain the exact qualified tool name, arguments, signature, and optional MRTR state. +// +// Return values: +// - *mcp.CallToolResult: The normalized upstream result is returned on success. +// - error: A modern validation error or a wrapped routing, execution, or billing error is returned on failure. +func executeModernMCPTool(ctx context.Context, c *gin.Context, params modernMCPCallParams) (*mcp.CallToolResult, error) { + if params.Arguments == nil { + params.Arguments = map[string]any{} + } + + plan, err := prepareModernMCPToolCall(c, params) + if err != nil { + return nil, &modernMCPValidationError{ + Status: http.StatusOK, + Code: mcpErrInvalidParams, + Err: errors.Wrap(err, "resolve modern mcp tool descriptor"), + } + } + descriptor, err := resolveModernMCPToolDescriptor(ctx, c, params.Name, plan) + if err != nil { + return nil, &modernMCPValidationError{ + Status: http.StatusOK, + Code: mcpErrInvalidParams, + Err: errors.Wrap(err, "resolve modern mcp tool descriptor"), + } + } + if err := mcp.ValidateToolArgumentHeaders(c.Request.Header, descriptor.InputSchema, params.Arguments); err != nil { + return nil, &modernMCPValidationError{ + Status: http.StatusBadRequest, + Code: mcp.ErrorCodeHeaderMismatch, + Err: err, + } + } + + result, err := executeModernMCPToolCall(ctx, c, params, plan) + if err != nil { + return nil, err + } + return mcp.NormalizeCallToolResult(result), nil +} + +// prepareModernMCPToolCall loads and resolves one immutable candidate snapshot for validation and execution. +// +// Parameters: +// - c: The Gin context contains the authenticated user. +// - params: The call parameters contain the exact qualified tool name and optional signature. +// +// Return values: +// - *modernMCPToolCallPlan: The prepared user, server, and candidate snapshot. +// - error: A wrapped authentication, catalog, policy, or candidate error. +func prepareModernMCPToolCall(c *gin.Context, params modernMCPCallParams) (*modernMCPToolCallPlan, error) { + user, err := getUserFromContext(c) + if err != nil { + return nil, errors.Wrap(err, "get user from context") + } + + serverLabel, toolName := splitToolName(params.Name) + if toolName == "" { + toolName = strings.TrimSpace(params.Name) + } + if toolName == "" { + return nil, errors.WithStack(errors.New("tool name is required")) + } + + servers, serverByID, err := loadMCPCallServers(serverLabel) + if err != nil { + return nil, err + } + toolsByServer, err := loadMCPToolsByServer(servers) + if err != nil { + return nil, err + } + + candidates, err := mcp.BuildToolCandidates( + servers, + toolsByServer, + nil, + user.MCPToolBlacklist, + []string{toolName}, + toolName, + params.Signature, + ) + if err != nil { + return nil, errors.Wrapf(err, "build mcp tool candidates for %q", toolName) + } + candidates = filterExactMCPToolCandidates(candidates, toolName) + if len(candidates) == 0 { + return nil, errors.Errorf("no eligible MCP tool found for exact name %q", toolName) + } + + return &modernMCPToolCallPlan{ + userID: user.Id, + serverByID: serverByID, + candidates: candidates, + }, nil +} + +// resolveModernMCPToolDescriptor restores the descriptor from the prepared candidate snapshot. +// +// Parameters: +// - ctx: The request context is used only by the compatibility fallback. +// - c: The Gin context is used only by the compatibility fallback. +// - name: The exact qualified wire name returned by tools/list. +// - plan: The prepared candidate snapshot; nil retains the catalog lookup fallback for isolated callers. +// +// Return values: +// - mcp.ToolDescriptor: The exact policy-filtered descriptor used by execution. +// - error: A wrapped descriptor or unknown-tool error. +func resolveModernMCPToolDescriptor( + ctx context.Context, + c *gin.Context, + name string, + plan *modernMCPToolCallPlan, +) (mcp.ToolDescriptor, error) { + if plan == nil { + return findModernMCPToolDescriptor(ctx, c, name) + } + for _, candidate := range plan.candidates { + if candidate.Tool == nil { + continue + } + qualifiedName := candidate.ServerLabel + "." + strings.TrimSpace(candidate.Tool.Name) + if qualifiedName != name { + continue + } + descriptor, err := descriptorForMCPTool(candidate.Tool) + if err != nil { + return mcp.ToolDescriptor{}, errors.Wrapf( + err, + "restore mcp descriptor for server %d tool %q", + candidate.ServerID, + candidate.Tool.Name, + ) + } + descriptor.Name = qualifiedName + return descriptor, nil + } + return mcp.ToolDescriptor{}, errors.Errorf("no eligible MCP tool found for %q", name) +} + +// callMCPToolForUserLatest routes one exact tool name across eligible servers and applies final-result billing. +// +// Parameters: +// - ctx: The request context controls database and upstream work. +// - c: The Gin context contains the authenticated user and request-scoped logger. +// - params: The call parameters contain the exact qualified tool name, arguments, signature, and optional MRTR state. +// +// Return values: +// - *mcp.CallToolResult: The normalized upstream result, including input_required intermediates, is returned on success. +// - error: A wrapped authentication, catalog, routing, execution, or billing error is returned on failure. +func callMCPToolForUserLatest(ctx context.Context, c *gin.Context, params modernMCPCallParams) (*mcp.CallToolResult, error) { + plan, err := prepareModernMCPToolCall(c, params) + if err != nil { + return nil, err + } + return executeModernMCPToolCall(ctx, c, params, plan) +} + +// executeModernMCPToolCall executes one prepared modern tool call with fallback and final-result billing. +// +// Parameters: +// - ctx: The request context controls upstream work. +// - c: The Gin context contains the request-scoped logger. +// - params: The call parameters contain arguments and optional MRTR state. +// - plan: The immutable user, server, and candidate snapshot prepared for this request. +// +// Return values: +// - *mcp.CallToolResult: The normalized upstream result, including input_required intermediates, is returned on success. +// - error: A wrapped routing, execution, or billing error is returned on failure. +func executeModernMCPToolCall( + ctx context.Context, + c *gin.Context, + params modernMCPCallParams, + plan *modernMCPToolCallPlan, +) (*mcp.CallToolResult, error) { + if plan == nil { + return nil, errors.WithStack(errors.New("modern mcp tool call plan is nil")) + } + + logger := gmw.GetLogger(c) + startedAt := time.Now() // Preserve the monotonic component for elapsed-time measurement. + selected, result, err := mcp.CallWithFallback(ctx, plan.candidates, func(ctx context.Context, candidate mcp.ToolCandidate) (*mcp.CallToolResult, error) { + server := plan.serverByID[candidate.ServerID] + if server == nil { + return nil, errors.WithStack(errors.New("mcp server not loaded")) + } + descriptor, err := descriptorForMCPTool(candidate.Tool) + if err != nil { + return nil, errors.Wrapf(err, "build descriptor for tool %q", candidate.Tool.Name) + } + client := mcp.NewStreamableHTTPClientWithLogger( + server, + nil, + time.Duration(config.MCPToolCallTimeoutSec)*time.Second, + logger, + ) + callResult, err := client.CallToolLatestWithOptions(ctx, descriptor, params.Arguments, mcp.CallToolRequestOptions{ + InputResponses: params.InputResponses, + RequestState: params.RequestState, + }) + if err != nil { + return nil, errors.Wrapf(err, "call mcp tool %q on server %d", candidate.Tool.Name, candidate.ServerID) + } + return callResult, nil + }) + if err != nil { + return nil, errors.Wrap(err, "call mcp tool with fallback") + } + + result = mcp.NormalizeCallToolResult(result) + if !shouldBillMCPToolResult(result) { + return result, nil + } + if err := chargeAndRecordMCPToolCall(ctx, c, plan.userID, plan.serverByID, selected, startedAt); err != nil { + return nil, err + } + return result, nil +} + +// filterExactMCPToolCandidates removes candidates that only matched case-insensitive policy normalization. +// +// Parameters: +// - candidates: The policy-filtered candidates come from the shared registry. +// - exactName: The exactName value is the case-sensitive upstream wire name requested by the client. +// +// Return values: +// - []mcp.ToolCandidate: The returned candidates have stored wire names that exactly equal exactName. +func filterExactMCPToolCandidates(candidates []mcp.ToolCandidate, exactName string) []mcp.ToolCandidate { + filtered := make([]mcp.ToolCandidate, 0, len(candidates)) + for _, candidate := range candidates { + if candidate.Tool == nil || strings.TrimSpace(candidate.Tool.Name) != exactName { + continue + } + filtered = append(filtered, candidate) + } + return filtered +} + +// shouldBillMCPToolResult reports whether one result represents a successful completed logical call. +// +// Parameters: +// - result: The normalized upstream result is inspected for final completion. +// +// Return values: +// - bool: True is returned only for non-error results whose resultType is complete. +func shouldBillMCPToolResult(result *mcp.CallToolResult) bool { + return result != nil && !result.IsError && result.ResultType == mcp.ResultTypeComplete +} diff --git a/controller/mcp_call_latest_catalog_test.go b/controller/mcp_call_latest_catalog_test.go new file mode 100644 index 0000000000..0e5205d4cf --- /dev/null +++ b/controller/mcp_call_latest_catalog_test.go @@ -0,0 +1,38 @@ +package controller + +import ( + "context" + "testing" + + "github.com/stretchr/testify/require" + "gorm.io/gorm" + + "github.com/Laisky/one-api/model" +) + +func TestExecuteModernMCPToolReusesPreparedCatalog(t *testing.T) { + cleanup, fixture := setupMCPProxyTest(t) + defer cleanup() + + toolQueries := 0 + const callbackName = "test:modern-mcp-tool-call-reuses-prepared-catalog" + err := model.DB.Callback().Query().Before("gorm:query").Register(callbackName, func(tx *gorm.DB) { + if tx != nil && tx.Statement != nil && tx.Statement.Table == "mcp_tools" { + toolQueries++ + } + }) + require.NoError(t, err) + defer func() { + require.NoError(t, model.DB.Callback().Query().Remove(callbackName)) + }() + + c, _ := newMCPCallContext(t, fixture.user.Id, "modern-catalog-reuse") + result, err := executeModernMCPTool(context.Background(), c, modernMCPCallParams{ + Name: "fake-mcp.echo", + Arguments: map[string]any{"message": "hello"}, + }) + require.NoError(t, err) + require.NotNil(t, result) + require.Equal(t, 1, fixture.upstreamHits) + require.Equal(t, 1, toolQueries, "modern tools/call must reuse one prepared tool-catalog snapshot") +} diff --git a/controller/mcp_catalog_latest.go b/controller/mcp_catalog_latest.go new file mode 100644 index 0000000000..86507482a4 --- /dev/null +++ b/controller/mcp_catalog_latest.go @@ -0,0 +1,265 @@ +package controller + +import ( + "context" + "crypto/sha256" + "encoding/base64" + "encoding/hex" + "encoding/json" + "sort" + "strings" + + "github.com/Laisky/errors/v2" + "github.com/gin-gonic/gin" + + "github.com/Laisky/one-api/model" + "github.com/Laisky/one-api/relay/mcp" +) + +const modernMCPToolsPageSize = 128 + +type modernMCPToolsListParams struct { + Cursor string `json:"cursor,omitempty"` +} + +type modernMCPToolsCursor struct { + Offset int `json:"offset"` + Fingerprint string `json:"fingerprint"` +} + +// listModernMCPToolsPage returns one deterministic page from the authenticated aggregate tool catalog. +// +// Parameters: +// - ctx: the request context controlling database and policy work. +// - c: the Gin context containing the authenticated user. +// - rawParams: the encoded tools/list parameters containing an optional opaque cursor. +// +// Return values: +// - gin.H: a current MCP tools/list result with caching and optional nextCursor fields. +// - error: a wrapped authentication, database, descriptor, cursor, or encoding error. +func listModernMCPToolsPage(ctx context.Context, c *gin.Context, rawParams json.RawMessage) (gin.H, error) { + var params modernMCPToolsListParams + if len(rawParams) != 0 { + if err := json.Unmarshal(rawParams, ¶ms); err != nil { + return nil, &modernMCPValidationError{Status: 400, Code: mcpErrInvalidParams, Err: errors.Wrap(err, "decode mcp tools/list params")} + } + } + if len(params.Cursor) > 2048 { + return nil, &modernMCPValidationError{Status: 400, Code: mcpErrInvalidParams, Err: errors.New("mcp tools/list cursor is too large")} + } + + tools, err := listMCPToolDescriptorsForUser(ctx, c) + if err != nil { + return nil, err + } + fingerprint, err := fingerprintMCPToolDescriptors(tools) + if err != nil { + return nil, errors.Wrap(err, "fingerprint mcp tool catalog") + } + offset := 0 + if params.Cursor != "" { + cursor, err := decodeModernMCPToolsCursor(params.Cursor) + if err != nil { + return nil, &modernMCPValidationError{Status: 400, Code: mcpErrInvalidParams, Err: err} + } + if cursor.Fingerprint != fingerprint { + return nil, &modernMCPValidationError{Status: 400, Code: mcpErrInvalidParams, Err: errors.New("mcp tools/list cursor is stale")} + } + if cursor.Offset < 0 || cursor.Offset > len(tools) { + return nil, &modernMCPValidationError{Status: 400, Code: mcpErrInvalidParams, Err: errors.New("mcp tools/list cursor offset is invalid")} + } + offset = cursor.Offset + } + + end := offset + modernMCPToolsPageSize + if end > len(tools) { + end = len(tools) + } + page := append([]mcp.ToolDescriptor(nil), tools[offset:end]...) + result := gin.H{ + "resultType": mcp.ResultTypeComplete, + "tools": page, + "ttlMs": int64(60000), + "cacheScope": mcp.CacheScopePrivate, + } + if end < len(tools) { + nextCursor, err := encodeModernMCPToolsCursor(modernMCPToolsCursor{Offset: end, Fingerprint: fingerprint}) + if err != nil { + return nil, errors.Wrap(err, "encode next mcp tools/list cursor") + } + result["nextCursor"] = nextCursor + } + return result, nil +} + +// listMCPToolDescriptorsForUser builds the lossless policy-filtered aggregate catalog for one user. +// +// Parameters: +// - ctx: the request context controlling database and policy work. +// - c: the Gin context containing the authenticated user. +// +// Return values: +// - []mcp.ToolDescriptor: qualified descriptors sorted deterministically by exact wire name. +// - error: a wrapped authentication, database, policy, or stored-descriptor error. +func listMCPToolDescriptorsForUser(ctx context.Context, c *gin.Context) ([]mcp.ToolDescriptor, error) { + _ = ctx + user, err := getUserFromContext(c) + if err != nil { + return nil, errors.Wrap(err, "get user from context") + } + servers, err := model.ListEnabledMCPServers() + if err != nil { + return nil, errors.Wrap(err, "list enabled mcp servers") + } + sort.SliceStable(servers, func(left, right int) bool { + if servers[left].GetPriority() == servers[right].GetPriority() { + return servers[left].Id < servers[right].Id + } + return servers[left].GetPriority() > servers[right].GetPriority() + }) + + descriptors := make([]mcp.ToolDescriptor, 0) + for _, server := range servers { + if server == nil { + continue + } + tools, err := model.GetMCPToolsByServerID(server.Id) + if err != nil { + return nil, errors.Wrapf(err, "get mcp tools for server %d", server.Id) + } + resolved, err := mcp.ResolveTools(server, tools, nil, user.MCPToolBlacklist, nil) + if err != nil { + return nil, errors.Wrapf(err, "resolve mcp tools for server %d", server.Id) + } + for _, entry := range resolved { + if !entry.Policy.Allowed || entry.Tool == nil { + continue + } + descriptor, err := descriptorForMCPTool(entry.Tool) + if err != nil { + return nil, errors.Wrapf(err, "restore mcp descriptor for server %d tool %q", server.Id, entry.Tool.Name) + } + descriptor.Name = server.Name + "." + descriptor.Name + descriptors = append(descriptors, descriptor) + } + } + sort.SliceStable(descriptors, func(left, right int) bool { + return descriptors[left].Name < descriptors[right].Name + }) + return descriptors, nil +} + +// descriptorForMCPTool restores the complete wire descriptor with compatibility fallbacks for old rows. +// +// Parameters: +// - tool: the synchronized database row. +// +// Return values: +// - mcp.ToolDescriptor: the lossless descriptor with exact wire name and a non-nil input schema. +// - error: a wrapped stored JSON or legacy input-schema decoding error. +func descriptorForMCPTool(tool *model.MCPTool) (mcp.ToolDescriptor, error) { + if tool == nil { + return mcp.ToolDescriptor{}, errors.New("mcp tool is nil") + } + descriptor := mcp.ToolDescriptor{} + if strings.TrimSpace(tool.DescriptorJSON) != "" { + if err := json.Unmarshal([]byte(tool.DescriptorJSON), &descriptor); err != nil { + return mcp.ToolDescriptor{}, errors.Wrap(err, "decode stored mcp descriptor") + } + } + descriptor.Name = tool.Name + if descriptor.Title == "" { + descriptor.Title = tool.DisplayName + } + if descriptor.Description == "" { + descriptor.Description = tool.Description + } + if descriptor.InputSchema == nil && strings.TrimSpace(tool.InputSchema) != "" { + if err := json.Unmarshal([]byte(tool.InputSchema), &descriptor.InputSchema); err != nil { + return mcp.ToolDescriptor{}, errors.Wrap(err, "decode stored mcp input schema") + } + } + if descriptor.InputSchema == nil { + descriptor.InputSchema = map[string]any{"type": "object"} + } + return descriptor, nil +} + +// findModernMCPToolDescriptor resolves one exact qualified tool name for header validation. +// +// Parameters: +// - ctx: the request context controlling database and policy work. +// - c: the Gin context containing the authenticated user. +// - name: the exact qualified wire name returned by tools/list. +// +// Return values: +// - mcp.ToolDescriptor: the matching descriptor. +// - error: a wrapped catalog error or an unknown-tool error. +func findModernMCPToolDescriptor(ctx context.Context, c *gin.Context, name string) (mcp.ToolDescriptor, error) { + tools, err := listMCPToolDescriptorsForUser(ctx, c) + if err != nil { + return mcp.ToolDescriptor{}, errors.Wrap(err, "list mcp tools for header validation") + } + for _, tool := range tools { + if tool.Name == name { + return tool, nil + } + } + return mcp.ToolDescriptor{}, errors.Errorf("no eligible MCP tool found for %q", name) +} + +// fingerprintMCPToolDescriptors returns a stable digest of one deterministic catalog. +// +// Parameters: +// - tools: the sorted aggregate descriptors. +// +// Return values: +// - string: a lowercase hexadecimal SHA-256 digest. +// - error: a wrapped encoding error. +func fingerprintMCPToolDescriptors(tools []mcp.ToolDescriptor) (string, error) { + encoded, err := json.Marshal(tools) + if err != nil { + return "", errors.Wrap(err, "marshal mcp tool catalog") + } + digest := sha256.Sum256(encoded) + return hex.EncodeToString(digest[:]), nil +} + +// encodeModernMCPToolsCursor encodes one cursor as opaque URL-safe JSON. +// +// Parameters: +// - cursor: the validated offset and catalog fingerprint. +// +// Return values: +// - string: the URL-safe opaque cursor. +// - error: a wrapped JSON encoding error. +func encodeModernMCPToolsCursor(cursor modernMCPToolsCursor) (string, error) { + encoded, err := json.Marshal(cursor) + if err != nil { + return "", errors.Wrap(err, "marshal mcp tools/list cursor") + } + return base64.RawURLEncoding.EncodeToString(encoded), nil +} + +// decodeModernMCPToolsCursor decodes one opaque cursor and validates required fields. +// +// Parameters: +// - value: the URL-safe cursor supplied by a client. +// +// Return values: +// - modernMCPToolsCursor: the decoded cursor. +// - error: a wrapped Base64, JSON, or required-field error. +func decodeModernMCPToolsCursor(value string) (modernMCPToolsCursor, error) { + encoded, err := base64.RawURLEncoding.DecodeString(value) + if err != nil { + return modernMCPToolsCursor{}, errors.Wrap(err, "decode mcp tools/list cursor") + } + var cursor modernMCPToolsCursor + if err := json.Unmarshal(encoded, &cursor); err != nil { + return modernMCPToolsCursor{}, errors.Wrap(err, "unmarshal mcp tools/list cursor") + } + if cursor.Fingerprint == "" { + return modernMCPToolsCursor{}, errors.New("mcp tools/list cursor fingerprint is required") + } + return cursor, nil +} diff --git a/controller/mcp_proxy.go b/controller/mcp_proxy.go index 1bbe21350b..3c1a421b7c 100644 --- a/controller/mcp_proxy.go +++ b/controller/mcp_proxy.go @@ -5,7 +5,6 @@ import ( "encoding/json" "fmt" "net/http" - "sort" "strings" "time" @@ -30,22 +29,21 @@ type mcpRPCRequest struct { Params json.RawMessage `json:"params"` } +type mcpInitializeParams struct { + ProtocolVersion string `json:"protocolVersion"` +} + type mcpCallParams struct { Name string `json:"name"` Arguments map[string]any `json:"arguments"` Signature string `json:"signature,omitempty"` } -// MCP Streamable HTTP transport constants. The protocol version advertised here -// matches what the upstream client in relay/mcp/client.go negotiates by default -// and is supported by current MCP Inspector / SDK releases. const ( - mcpProtocolVersion = "2025-06-18" - mcpServerName = "one-api-mcp-proxy" - mcpServerVersion = "1.0.0" + mcpServerName = "one-api-mcp-proxy" + mcpServerVersion = "1.1.0" ) -// JSON-RPC 2.0 error codes (https://www.jsonrpc.org/specification#error_object). const ( mcpErrParseError = -32700 mcpErrInvalidRequest = -32600 @@ -54,10 +52,12 @@ const ( mcpErrInternal = -32603 ) -// MCPProxy handles MCP Streamable HTTP requests backed by configured MCP servers. -// Implements the single-endpoint Streamable HTTP transport: POST for JSON-RPC -// messages, GET for optional server-to-client SSE (not supported here, so 405), -// DELETE for session termination (stateless proxy, also 405). +// MCPProxy handles initialization-based MCP requests backed by the aggregate tool catalog. +// +// Parameters: +// - c: the Gin context containing the authenticated Streamable HTTP request. +// +// Return values: none; the function writes the complete HTTP response. func MCPProxy(c *gin.Context) { switch c.Request.Method { case http.MethodPost: @@ -69,23 +69,40 @@ func MCPProxy(c *gin.Context) { } } +// handleMCPPost dispatches one initialization-based JSON-RPC request or notification. +// +// Parameters: +// - c: the Gin context containing the authenticated request and request-scoped logger. +// +// Return values: none; the function writes a JSON-RPC response or HTTP 202 for notifications. func handleMCPPost(c *gin.Context) { ctx := gmw.Ctx(c) - var req mcpRPCRequest - if err := json.NewDecoder(c.Request.Body).Decode(&req); err != nil { + var request mcpRPCRequest + if err := json.NewDecoder(c.Request.Body).Decode(&request); err != nil { respondMCPError(c, nil, mcpErrParseError, errors.Wrap(err, "decode mcp request")) return } + if request.JSONRPC != "2.0" || strings.TrimSpace(request.Method) == "" { + respondMCPError(c, request.ID, mcpErrInvalidRequest, errors.New("jsonrpc must be 2.0 and method is required")) + return + } + isNotification := request.ID == nil - // JSON-RPC notifications carry no `id`. The Streamable HTTP transport - // requires the server to reply with HTTP 202 and an empty body — never a - // JSON-RPC envelope — so SDK clients don't try to correlate a response. - isNotification := req.ID == nil - - switch strings.ToLower(strings.TrimSpace(req.Method)) { + switch strings.ToLower(strings.TrimSpace(request.Method)) { case "initialize": - respondMCPResult(c, req.ID, gin.H{ - "protocolVersion": mcpProtocolVersion, + if isNotification { + respondMCPError(c, nil, mcpErrInvalidRequest, errors.New("initialize requires a request id")) + return + } + var params mcpInitializeParams + if len(request.Params) != 0 { + if err := json.Unmarshal(request.Params, ¶ms); err != nil { + respondMCPError(c, request.ID, mcpErrInvalidParams, errors.Wrap(err, "decode mcp initialize params")) + return + } + } + respondMCPResult(c, request.ID, gin.H{ + "protocolVersion": mcp.NegotiateLegacyProtocolVersion(params.ProtocolVersion), "capabilities": gin.H{ "tools": gin.H{"listChanged": false}, }, @@ -97,88 +114,70 @@ func handleMCPPost(c *gin.Context) { case "notifications/initialized", "notifications/cancelled", "notifications/progress", "notifications/roots/list_changed": c.AbortWithStatus(http.StatusAccepted) case "ping": - respondMCPResult(c, req.ID, gin.H{}) + if isNotification { + c.AbortWithStatus(http.StatusAccepted) + return + } + respondMCPResult(c, request.ID, gin.H{}) case "tools/list": + if isNotification { + c.AbortWithStatus(http.StatusAccepted) + return + } tools, err := listMCPToolsForUser(ctx, c) if err != nil { - respondMCPError(c, req.ID, mcpErrInternal, err) + respondMCPError(c, request.ID, mcpErrInternal, err) return } - respondMCPResult(c, req.ID, gin.H{"tools": tools}) + respondMCPResult(c, request.ID, gin.H{"tools": tools}) case "tools/call": + if isNotification { + c.AbortWithStatus(http.StatusAccepted) + return + } var params mcpCallParams - if err := json.Unmarshal(req.Params, ¶ms); err != nil { - respondMCPError(c, req.ID, mcpErrInvalidParams, errors.Wrap(err, "decode mcp call params")) + if err := json.Unmarshal(request.Params, ¶ms); err != nil { + respondMCPError(c, request.ID, mcpErrInvalidParams, errors.Wrap(err, "decode mcp call params")) return } result, err := callMCPToolForUser(ctx, c, params) if err != nil { - respondMCPError(c, req.ID, mcpErrInternal, err) + respondMCPError(c, request.ID, mcpErrInternal, err) return } - respondMCPResult(c, req.ID, result) + respondMCPResult(c, request.ID, result) default: if isNotification { c.AbortWithStatus(http.StatusAccepted) return } - respondMCPError(c, req.ID, mcpErrMethodNotFound, errors.Errorf("unsupported method %s", req.Method)) + respondMCPError(c, request.ID, mcpErrMethodNotFound, errors.Errorf("unsupported method %s", request.Method)) } } -// listMCPToolsForUser returns the allowed MCP tools for the authenticated user. +// listMCPToolsForUser returns lossless, policy-filtered, qualified descriptors for the authenticated user. +// +// Parameters: +// - ctx: the request context controlling database and policy work. +// - c: the Gin context containing the authenticated user. +// +// Return values: +// - []mcp.ToolDescriptor: the deterministic aggregate tool catalog. +// - error: a wrapped authentication, database, policy, or descriptor error. func listMCPToolsForUser(ctx context.Context, c *gin.Context) ([]mcp.ToolDescriptor, error) { - user, err := getUserFromContext(c) - if err != nil { - return nil, errors.Wrap(err, "get user from context") - } - - servers, err := model.ListEnabledMCPServers() - if err != nil { - return nil, errors.Wrap(err, "list enabled mcp servers") - } - - sort.SliceStable(servers, func(i, j int) bool { - if servers[i].GetPriority() == servers[j].GetPriority() { - return servers[i].Id < servers[j].Id - } - return servers[i].GetPriority() > servers[j].GetPriority() - }) - - // Initialize as a non-nil empty slice so the JSON-RPC `tools/list` - // response marshals to `[]` instead of `null` when no servers/tools - // resolve. Spec-compliant MCP clients (e.g. MCP Inspector with Zod - // schemas) reject `null` for the required `tools` array. See issue #340. - descriptors := make([]mcp.ToolDescriptor, 0) - for _, server := range servers { - tools, err := model.GetMCPToolsByServerID(server.Id) - if err != nil { - return nil, errors.Wrapf(err, "get mcp tools for server %d", server.Id) - } - resolved, err := mcp.ResolveTools(server, tools, nil, user.MCPToolBlacklist, nil) - if err != nil { - return nil, errors.Wrapf(err, "resolve mcp tools for server %d", server.Id) - } - for _, entry := range resolved { - if !entry.Policy.Allowed { - continue - } - var schema map[string]any - if entry.Tool.InputSchema != "" { - _ = json.Unmarshal([]byte(entry.Tool.InputSchema), &schema) - } - name := server.Name + "." + entry.Tool.Name - descriptors = append(descriptors, mcp.ToolDescriptor{ - Name: name, - Description: entry.Tool.Description, - InputSchema: schema, - }) - } - } - return descriptors, nil + return listMCPToolDescriptorsForUser(ctx, c) } -// callMCPToolForUser invokes a MCP tool and applies billing/logging. +// callMCPToolForUser routes one legacy downstream request through the modern-first upstream client. +// +// Parameters: +// - ctx: the request context controlling database and upstream work. +// - c: the Gin context containing the authenticated user and request-scoped logger. +// - params: the exact qualified tool name, arguments, and optional candidate signature. +// +// Return values: +// - *mcp.CallToolResult: the normalized upstream result. +// - error: a wrapped authentication, catalog, routing, execution, or billing error. func callMCPToolForUser(ctx context.Context, c *gin.Context, params mcpCallParams) (*mcp.CallToolResult, error) { logger := gmw.GetLogger(c) user, err := getUserFromContext(c) @@ -191,64 +190,56 @@ func callMCPToolForUser(ctx context.Context, c *gin.Context, params mcpCallParam toolName = strings.TrimSpace(params.Name) } if toolName == "" { - return nil, errors.New("tool name is required") + return nil, errors.WithStack(errors.New("tool name is required")) } - - var servers []*model.MCPServer - serverByID := make(map[int]*model.MCPServer) - if serverLabel != "" { - server, err := model.GetMCPServerByName(serverLabel) - if err != nil { - return nil, errors.Wrapf(err, "get mcp server by name %q", serverLabel) - } - servers = []*model.MCPServer{server} - serverByID[server.Id] = server - } else { - servers, err = model.ListEnabledMCPServers() - if err != nil { - return nil, errors.Wrap(err, "list enabled mcp servers") - } - for _, server := range servers { - if server == nil { - continue - } - serverByID[server.Id] = server - } + if params.Arguments == nil { + params.Arguments = map[string]any{} } - toolsByServer := make(map[int][]*model.MCPTool, len(servers)) - for _, server := range servers { - if server == nil { - continue - } - tools, err := model.GetMCPToolsByServerID(server.Id) - if err != nil { - return nil, errors.Wrapf(err, "get mcp tools for server %d", server.Id) - } - toolsByServer[server.Id] = tools + servers, serverByID, err := loadMCPCallServers(serverLabel) + if err != nil { + return nil, err + } + toolsByServer, err := loadMCPToolsByServer(servers) + if err != nil { + return nil, err } - candidates, err := mcp.BuildToolCandidates(servers, toolsByServer, nil, user.MCPToolBlacklist, []string{toolName}, toolName, params.Signature) + candidates, err := mcp.BuildToolCandidates( + servers, + toolsByServer, + nil, + user.MCPToolBlacklist, + []string{toolName}, + toolName, + params.Signature, + ) if err != nil { return nil, errors.Wrapf(err, "build mcp tool candidates for %q", toolName) } + candidates = filterExactMCPToolCandidates(candidates, toolName) if len(candidates) == 0 { - return nil, errors.New("no eligible MCP tool found") + return nil, errors.Errorf("no eligible MCP tool found for exact name %q", toolName) } - startedAt := time.Now() + startedAt := time.Now() // Preserve the monotonic component for elapsed-time measurement. selected, result, err := mcp.CallWithFallback(ctx, candidates, func(ctx context.Context, candidate mcp.ToolCandidate) (*mcp.CallToolResult, error) { server := serverByID[candidate.ServerID] if server == nil { - return nil, errors.New("mcp server not loaded") + return nil, errors.WithStack(errors.New("mcp server not loaded")) + } + descriptor, err := descriptorForMCPTool(candidate.Tool) + if err != nil { + return nil, errors.Wrapf(err, "build descriptor for tool %q", candidate.Tool.Name) } - client := mcp.NewStreamableHTTPClientWithLogger(server, nil, time.Duration(config.MCPToolCallTimeoutSec)*time.Second, logger) - callResult, err := client.CallTool(ctx, candidate.Tool.Name, params.Arguments) + client := mcp.NewStreamableHTTPClientWithLogger( + server, + nil, + time.Duration(config.MCPToolCallTimeoutSec)*time.Second, + logger, + ) + callResult, err := client.CallToolLatestWithDescriptor(ctx, descriptor, params.Arguments) if err != nil { - logger.Warn("mcp tool call failed", server.Ref().AppendZap([]zap.Field{ - zap.Error(err), - zap.String("tool", candidate.Tool.Name), - })...) return nil, errors.Wrapf(err, "call mcp tool %q on server %d", candidate.Tool.Name, candidate.ServerID) } return callResult, nil @@ -257,32 +248,116 @@ func callMCPToolForUser(ctx context.Context, c *gin.Context, params mcpCallParam return nil, errors.Wrap(err, "call mcp tool with fallback") } - if result.IsError { + result = mcp.NormalizeCallToolResult(result) + if !shouldBillMCPToolResult(result) { return result, nil } + if err := chargeAndRecordMCPToolCall(ctx, c, user.Id, serverByID, selected, startedAt); err != nil { + return nil, err + } + return result, nil +} + +// loadMCPCallServers loads either one explicitly selected server or every enabled server. +// +// Parameters: +// - serverLabel: an optional configured MCP server name. +// +// Return values: +// - []*model.MCPServer: candidate servers in repository-defined order. +// - map[int]*model.MCPServer: candidate servers indexed by internal id. +// - error: a wrapped server lookup error. +func loadMCPCallServers(serverLabel string) ([]*model.MCPServer, map[int]*model.MCPServer, error) { + serverByID := make(map[int]*model.MCPServer) + if serverLabel != "" { + server, err := model.GetMCPServerByName(serverLabel) + if err != nil { + return nil, nil, errors.Wrapf(err, "get mcp server by name %q", serverLabel) + } + serverByID[server.Id] = server + return []*model.MCPServer{server}, serverByID, nil + } + servers, err := model.ListEnabledMCPServers() + if err != nil { + return nil, nil, errors.Wrap(err, "list enabled mcp servers") + } + for _, server := range servers { + if server != nil { + serverByID[server.Id] = server + } + } + return servers, serverByID, nil +} + +// loadMCPToolsByServer loads synchronized tool rows for each candidate MCP server. +// +// Parameters: +// - servers: candidate MCP servers. +// +// Return values: +// - map[int][]*model.MCPTool: synchronized tools grouped by owning server id. +// - error: a wrapped database error. +func loadMCPToolsByServer(servers []*model.MCPServer) (map[int][]*model.MCPTool, error) { + toolsByServer := make(map[int][]*model.MCPTool, len(servers)) + for _, server := range servers { + if server == nil { + continue + } + tools, err := model.GetMCPToolsByServerID(server.Id) + if err != nil { + return nil, errors.Wrapf(err, "get mcp tools for server %d", server.Id) + } + toolsByServer[server.Id] = tools + } + return toolsByServer, nil +} + +// chargeAndRecordMCPToolCall applies quota and writes one finalized tool-call audit log. +// +// Parameters: +// - ctx: the request context controlling quota persistence. +// - c: the Gin context carrying request identity and tracing metadata. +// - userID: the authenticated user's internal id. +// - serverByID: loaded server configurations indexed by internal id. +// - selected: the successful tool candidate. +// - startedAt: the beginning of the logical tool call. +// +// Return values: +// - error: a wrapped server lookup or quota update error. +func chargeAndRecordMCPToolCall(ctx context.Context, c *gin.Context, userID int, serverByID map[int]*model.MCPServer, selected mcp.ToolCandidate, startedAt time.Time) error { server := serverByID[selected.ServerID] if server == nil { - return nil, errors.New("mcp server not loaded") + return errors.WithStack(errors.New("mcp server not loaded")) } - cost := resolveToolCost(server, selected.Tool.Name) if cost > 0 { - if err := model.DecreaseUserQuota(ctx, user.Id, cost); err != nil { - return nil, errors.Wrap(err, "decrease user quota for mcp tool call") + if err := model.DecreaseUserQuota(ctx, userID, cost); err != nil { + return errors.Wrap(err, "decrease user quota for mcp tool call") } - model.UpdateUserUsedQuotaAndRequestCountWithContext(ctx, user.Id, cost) + model.UpdateUserUsedQuotaAndRequestCountWithContext(ctx, userID, cost) } - qualifiedName := server.Name + "." + selected.Tool.Name - recordMCPToolLog(ctx, c, user.Id, server.Id, qualifiedName, cost, helper.CalcElapsedTime(startedAt)) - - return result, nil + recordMCPToolLog(ctx, c, userID, server.Id, qualifiedName, cost, helper.CalcElapsedTime(startedAt)) + return nil } -// resolveToolCost determines the quota cost for a MCP tool invocation. +// resolveToolCost determines the quota charge for one exact MCP tool name. +// +// Parameters: +// - server: the owning MCP server configuration. +// - toolName: the exact upstream tool name. +// +// Return values: +// - int64: the configured non-negative quota cost. func resolveToolCost(server *model.MCPServer, toolName string) int64 { - pricing := server.ToolPricing[strings.ToLower(toolName)] + if server == nil { + return 0 + } + pricing, exists := server.ToolPricing[toolName] + if !exists { + pricing = server.ToolPricing[strings.ToLower(toolName)] + } if pricing.QuotaPerCall > 0 { return pricing.QuotaPerCall } @@ -292,20 +367,21 @@ func resolveToolCost(server *model.MCPServer, toolName string) int64 { return 0 } -// mcpServerLabel renders an MCP server for log content by name and external -// UUID, never by internal integer id. +// mcpServerLabel renders an MCP server using its public name and UUID. +// // Parameters: -// - ctx: request context (reserved for logging; the store lookup is context-free). -// - serverId: internal MCP server id. +// - ctx: the request context reserved for future context-aware store access. +// - serverID: the internal MCP server id. // // Return values: // - string: " " when resolvable, otherwise "unknown". -func mcpServerLabel(ctx context.Context, serverId int) string { - server, err := model.GetMCPServerByID(serverId) +func mcpServerLabel(ctx context.Context, serverID int) string { + _ = ctx + server, err := model.GetMCPServerByID(serverID) if err != nil || server == nil { return "unknown" } - parts := []string{} + parts := make([]string, 0, 2) if name := strings.TrimSpace(server.Name); name != "" { parts = append(parts, name) } @@ -318,18 +394,26 @@ func mcpServerLabel(ctx context.Context, serverId int) string { return strings.Join(parts, " ") } -// recordMCPToolLog records an MCP tool invocation as a single LogTypeTool row. -// The dashboard tool charts aggregate strictly on type, so this becomes one -// row per invocation with ModelName=toolName and Quota=cost. Free invocations -// (cost == 0) still emit a row so every MCP call has a unified audit trail. -func recordMCPToolLog(ctx context.Context, c *gin.Context, userId int, serverId int, toolName string, cost int64, elapsedMs int64) { +// recordMCPToolLog records one finalized MCP tool invocation in the tool audit stream. +// +// Parameters: +// - ctx: the request context carrying cancellation and trace state. +// - c: the Gin context carrying authenticated UUIDs and request identifiers. +// - userID: the authenticated user id. +// - serverID: the selected MCP server id. +// - toolName: the qualified tool name exposed by one-api. +// - cost: the charged quota units. +// - elapsedMs: total logical tool-call latency in milliseconds. +// +// Return values: none; the shared model logging path owns persistence handling. +func recordMCPToolLog(ctx context.Context, c *gin.Context, userID int, serverID int, toolName string, cost int64, elapsedMs int64) { model.RecordToolLog(ctx, &model.Log{ - UserId: userId, + UserId: userID, UserUUID: model.StringPtrIfNotEmpty(c.GetString(ctxkey.UserUUID)), TokenUUID: model.StringPtrIfNotEmpty(c.GetString(ctxkey.TokenUUID)), ModelName: toolName, Quota: int(cost), - Content: fmt.Sprintf("MCP tool call: %s (server %s)", toolName, mcpServerLabel(ctx, serverId)), + Content: fmt.Sprintf("MCP tool call: %s (server %s)", toolName, mcpServerLabel(ctx, serverID)), RequestId: c.GetString(ctxkey.RequestId), TraceId: tracing.GetTraceID(c), IsStream: false, @@ -338,17 +422,22 @@ func recordMCPToolLog(ctx context.Context, c *gin.Context, userId int, serverId } // getUserFromContext loads the authenticated user from request context. -// It first checks for the cached UserObj set by auth middleware, -// falling back to a database lookup if not present. +// +// Parameters: +// - c: the Gin context populated by token authentication middleware. +// +// Return values: +// - *model.User: the authenticated user. +// - error: a wrapped lookup error or missing-identity error. func getUserFromContext(c *gin.Context) (*model.User, error) { - if userObj, exists := c.Get(ctxkey.UserObj); exists { - if u, ok := userObj.(*model.User); ok { - return u, nil + if userObject, exists := c.Get(ctxkey.UserObj); exists { + if user, ok := userObject.(*model.User); ok && user != nil { + return user, nil } } userID := c.GetInt(ctxkey.Id) if userID == 0 { - return nil, errors.New("user id missing") + return nil, errors.WithStack(errors.New("user id missing")) } user, err := model.GetUserById(userID, true) if err != nil { @@ -357,44 +446,61 @@ func getUserFromContext(c *gin.Context) (*model.User, error) { return user, nil } -// splitToolName splits server-qualified tool names. -func splitToolName(name string) (string, string) { - parts := strings.SplitN(name, ".", 2) +// splitToolName separates an optional server qualifier from the exact upstream tool name. +// +// Parameters: +// - value: a qualified name in server.tool form or an unqualified tool name. +// +// Return values: +// - string: the optional server name before the first dot. +// - string: the remaining exact tool name, which may itself contain dots. +func splitToolName(value string) (string, string) { + value = strings.TrimSpace(value) + parts := strings.SplitN(value, ".", 2) if len(parts) != 2 { - return "", "" + return "", value } return strings.TrimSpace(parts[0]), strings.TrimSpace(parts[1]) } -// isToolAllowed checks if a tool is permitted by the resolved policy. -func isToolAllowed(resolved []mcp.ResolvedTool, name string) bool { - for _, entry := range resolved { - if strings.EqualFold(entry.Tool.Name, name) { - return entry.Policy.Allowed - } - } - return false -} - -// respondMCPResult writes a JSON-RPC result payload. +// respondMCPResult writes one successful initialization-based JSON-RPC response. +// +// Parameters: +// - c: the Gin request context receiving the response. +// - id: the decoded JSON-RPC request identifier. +// - result: the result payload to encode. +// +// Return values: none; the function writes HTTP 200 JSON. func respondMCPResult(c *gin.Context, id any, result any) { - c.JSON(http.StatusOK, gin.H{ - "jsonrpc": "2.0", - "id": id, - "result": result, - }) + c.JSON(http.StatusOK, gin.H{"jsonrpc": "2.0", "id": id, "result": result}) } -// respondMCPError writes a JSON-RPC 2.0 error payload. The HTTP status stays -// 200 because JSON-RPC errors are envelope-level — clients parse the body to -// distinguish protocol errors from transport failures. +// respondMCPError writes one legacy JSON-RPC error and redacts internal implementation details. +// +// Parameters: +// - c: the Gin request context receiving the response and providing the request-scoped logger. +// - id: the decoded JSON-RPC request identifier. +// - code: the JSON-RPC error code. +// - err: the underlying validation or internal error. +// +// Return values: none; the function writes HTTP 200 JSON. func respondMCPError(c *gin.Context, id any, code int, err error) { + message := "mcp request failed" + if code == mcpErrInternal { + logger := gmw.GetLogger(c) + if err != nil { + logger.Error("mcp internal request failure", zap.Error(err)) + } + message = "internal MCP error" + } else if err != nil { + message = err.Error() + } c.JSON(http.StatusOK, gin.H{ "jsonrpc": "2.0", "id": id, "error": gin.H{ "code": code, - "message": err.Error(), + "message": message, }, }) } diff --git a/controller/mcp_proxy_latest.go b/controller/mcp_proxy_latest.go new file mode 100644 index 0000000000..8f8219b9f8 --- /dev/null +++ b/controller/mcp_proxy_latest.go @@ -0,0 +1,573 @@ +package controller + +import ( + "bytes" + "encoding/json" + stderrors "errors" + "fmt" + "io" + "math" + "net/http" + "net/url" + "strings" + + "github.com/Laisky/errors/v2" + gmw "github.com/Laisky/gin-middlewares/v7" + "github.com/Laisky/zap" + "github.com/gin-gonic/gin" + + "github.com/Laisky/one-api/relay/mcp" +) + +const ( + modernMCPMaxRequestBytes int64 = 4 << 20 + legacyMCPMaxRequestBytes int64 = 32 << 20 +) + +// modernMCPRequestMeta extracts the required per-request metadata for MCP 2026-07-28. +type modernMCPRequestMeta struct { + ProtocolVersion string `json:"io.modelcontextprotocol/protocolVersion"` + ClientInfo *mcp.ImplementationInfo `json:"io.modelcontextprotocol/clientInfo,omitempty"` + ClientCapabilities map[string]any `json:"io.modelcontextprotocol/clientCapabilities"` +} + +// modernMCPParamsEnvelope extracts modern metadata without constraining method-specific parameters. +type modernMCPParamsEnvelope struct { + Meta modernMCPRequestMeta `json:"_meta"` +} + +// modernMCPCallParams contains one tools/call request and optional multi-round-trip state. +type modernMCPCallParams struct { + Name string `json:"name"` + Arguments map[string]any `json:"arguments"` + Signature string `json:"signature,omitempty"` + InputResponses map[string]any `json:"inputResponses,omitempty"` + RequestState string `json:"requestState,omitempty"` +} + +// modernMCPValidationError carries HTTP and JSON-RPC details for one rejected modern request. +type modernMCPValidationError struct { + Status int + Code int + Err error + Data any +} + +// replayReadCloser replays bytes already inspected while retaining ownership of the original request body. +type replayReadCloser struct { + io.Reader + closer io.Closer +} + +// Error returns the underlying modern request validation message. +// +// Parameters: none. +// +// Return values: +// - string: the validation message or a stable fallback when the receiver is incomplete. +func (e *modernMCPValidationError) Error() string { + if e == nil || e.Err == nil { + return "invalid modern mcp request" + } + return e.Err.Error() +} + +// Unwrap returns the underlying validation error for errors.Is and errors.As. +// +// Parameters: none. +// +// Return values: +// - error: The underlying validation error is returned, or nil for an empty receiver. +func (e *modernMCPValidationError) Unwrap() error { + if e == nil { + return nil + } + return e.Err +} + +// Close closes the original HTTP request body retained by replayReadCloser. +// +// Parameters: none. +// +// Return values: +// - error: the original request body close error, or nil when no closer is present. +func (r *replayReadCloser) Close() error { + if r == nil || r.closer == nil { + return nil + } + return r.closer.Close() +} + +// MCPProxyLatest dispatches modern MCP requests and delegates recognized legacy traffic unchanged. +// +// Parameters: +// - c: the Gin request context containing the authenticated Streamable HTTP request. +// +// Return values: none; the function writes the complete HTTP response or delegates to MCPProxy. +func MCPProxyLatest(c *gin.Context) { + if err := validateModernMCPOrigin(c.Request); err != nil { + respondMCPModernError(c, nil, http.StatusForbidden, mcpErrInvalidRequest, err, nil) + return + } + + versionValues := c.Request.Header.Values(mcp.ProtocolVersionHeader) + if len(versionValues) > 1 { + respondMCPModernError(c, nil, http.StatusBadRequest, mcp.ErrorCodeHeaderMismatch, errors.Errorf("%s must occur at most once", mcp.ProtocolVersionHeader), nil) + return + } + modernTransport := len(versionValues) == 1 && !mcp.IsLegacyProtocolVersion(versionValues[0]) + + switch c.Request.Method { + case http.MethodGet, http.MethodDelete: + if modernTransport { + c.Header("Allow", http.MethodPost) + c.AbortWithStatus(http.StatusMethodNotAllowed) + return + } + MCPProxy(c) + return + case http.MethodPost: + // Continue below. + default: + c.Header("Allow", http.MethodPost) + c.AbortWithStatus(http.StatusMethodNotAllowed) + return + } + + requestLimit := legacyMCPMaxRequestBytes + if modernTransport { + requestLimit = modernMCPMaxRequestBytes + } + body, err := readBoundedMCPRequestBody(c.Request.Body, requestLimit) + if err != nil { + var tooLarge *mcpRequestTooLargeError + if stderrors.As(err, &tooLarge) { + respondMCPModernError(c, nil, http.StatusRequestEntityTooLarge, mcpErrInvalidRequest, err, nil) + return + } + respondMCPModernError(c, nil, http.StatusBadRequest, mcpErrParseError, errors.Wrap(err, "read mcp request"), nil) + return + } + originalBody := c.Request.Body + c.Request.Body = &replayReadCloser{Reader: bytes.NewReader(body), closer: originalBody} + + if len(versionValues) == 1 && mcp.IsLegacyProtocolVersion(versionValues[0]) { + MCPProxy(c) + return + } + + var request mcpRPCRequest + if err := json.Unmarshal(body, &request); err != nil { + if len(versionValues) == 0 { + MCPProxy(c) + return + } + respondMCPModernError(c, nil, http.StatusBadRequest, mcpErrParseError, errors.Wrap(err, "decode modern mcp request"), nil) + return + } + if !isModernMCPRequest(c, request) { + MCPProxy(c) + return + } + if int64(len(body)) > modernMCPMaxRequestBytes { + respondMCPModernError(c, request.ID, http.StatusRequestEntityTooLarge, mcpErrInvalidRequest, errors.Errorf("modern mcp request body exceeds %d bytes", modernMCPMaxRequestBytes), nil) + return + } + if err := validateModernMCPRequest(c, request); err != nil { + respondModernValidationError(c, request.ID, err) + return + } + handleModernMCPPost(c, request) +} + +// mcpRequestTooLargeError records the configured request-body limit that was exceeded. +type mcpRequestTooLargeError struct { + Limit int64 +} + +// Error returns a stable request-body limit message. +// +// Parameters: none. +// +// Return values: +// - string: The configured byte limit is included in the message. +func (e *mcpRequestTooLargeError) Error() string { + if e == nil { + return "mcp request body is too large" + } + return fmt.Sprintf("mcp request body exceeds %d bytes", e.Limit) +} + +// readBoundedMCPRequestBody reads one complete request without exceeding a fixed allocation boundary. +// +// Parameters: +// - reader: The inbound request body supplies the bytes to consume. +// - limit: The maximum accepted body size is expressed in bytes. +// +// Return values: +// - []byte: The complete request body is returned when it is within the limit. +// - error: A wrapped read error or mcpRequestTooLargeError is returned on failure. +func readBoundedMCPRequestBody(reader io.Reader, limit int64) ([]byte, error) { + if reader == nil { + return nil, errors.WithStack(errors.New("mcp request body is nil")) + } + if limit < 0 { + return nil, errors.WithStack(errors.New("mcp request body limit is negative")) + } + limited := &io.LimitedReader{R: reader, N: limit} + body, err := io.ReadAll(limited) + if err != nil { + return nil, errors.Wrap(err, "read bounded mcp request body") + } + if limited.N == 0 { + var extra [1]byte + _, probeErr := io.ReadFull(reader, extra[:]) + if probeErr == nil { + return nil, &mcpRequestTooLargeError{Limit: limit} + } + if !stderrors.Is(probeErr, io.EOF) { + return nil, errors.Wrap(probeErr, "probe bounded mcp request body") + } + } + return body, nil +} + +// isModernMCPRequest reports whether a request selects the 2026-07-28 stateless protocol profile. +// +// Parameters: +// - c: the Gin request context containing transport headers. +// - request: the parsed JSON-RPC request. +// +// Return values: +// - bool: true for modern metadata, the modern version header, discovery, or an unknown non-legacy version. +func isModernMCPRequest(c *gin.Context, request mcpRPCRequest) bool { + if strings.TrimSpace(request.Method) == "server/discover" { + return true + } + var params modernMCPParamsEnvelope + if json.Unmarshal(request.Params, ¶ms) == nil && strings.TrimSpace(params.Meta.ProtocolVersion) != "" { + return true + } + versions := c.Request.Header.Values(mcp.ProtocolVersionHeader) + if len(versions) != 1 { + return len(versions) > 1 + } + version := strings.TrimSpace(versions[0]) + return version != "" && !mcp.IsLegacyProtocolVersion(version) +} + +// validateModernMCPRequest validates JSON-RPC identity, metadata, and mirrored transport headers. +// +// Parameters: +// - c: the Gin request context containing transport headers. +// - request: the parsed modern JSON-RPC request. +// +// Return values: +// - error: a modernMCPValidationError describing the required HTTP status and JSON-RPC code. +func validateModernMCPRequest(c *gin.Context, request mcpRPCRequest) error { + if request.JSONRPC != "2.0" || strings.TrimSpace(request.Method) == "" { + return &modernMCPValidationError{Status: http.StatusBadRequest, Code: mcpErrInvalidRequest, Err: errors.New("jsonrpc must be 2.0 and method is required")} + } + if request.ID == nil && isModernMCPRequestMethod(request.Method) { + return &modernMCPValidationError{Status: http.StatusBadRequest, Code: mcpErrInvalidRequest, Err: errors.New("modern mcp requests require a non-null id")} + } + if request.ID != nil && !isValidModernMCPRequestID(request.ID) { + return &modernMCPValidationError{Status: http.StatusBadRequest, Code: mcpErrInvalidRequest, Err: errors.New("modern mcp request id must be a string or integer")} + } + + var params modernMCPParamsEnvelope + if err := json.Unmarshal(request.Params, ¶ms); err != nil { + return &modernMCPValidationError{Status: http.StatusBadRequest, Code: mcpErrInvalidRequest, Err: errors.Wrap(err, "decode modern mcp metadata")} + } + bodyVersion := strings.TrimSpace(params.Meta.ProtocolVersion) + headerVersion, err := singleMCPHeaderValue(c.Request.Header, mcp.ProtocolVersionHeader) + if err != nil || bodyVersion == "" { + if err == nil { + err = errors.New("modern mcp requests require protocol version metadata") + } + return &modernMCPValidationError{Status: http.StatusBadRequest, Code: mcp.ErrorCodeHeaderMismatch, Err: err} + } + if bodyVersion != headerVersion { + return &modernMCPValidationError{Status: http.StatusBadRequest, Code: mcp.ErrorCodeHeaderMismatch, Err: errors.New("MCP-Protocol-Version does not match request _meta")} + } + if bodyVersion != mcp.ProtocolVersion { + return &modernMCPValidationError{ + Status: http.StatusBadRequest, + Code: mcp.ErrorCodeUnsupportedProtocolVersion, + Err: errors.Errorf("unsupported mcp protocol version %q", bodyVersion), + Data: gin.H{ + "supported": mcp.SupportedProtocolVersions(), + "requested": bodyVersion, + }, + } + } + if params.Meta.ClientCapabilities == nil { + return &modernMCPValidationError{Status: http.StatusBadRequest, Code: mcpErrInvalidRequest, Err: errors.New("modern mcp requests require client capabilities in _meta")} + } + headerMethod, err := singleMCPHeaderValue(c.Request.Header, mcp.MethodHeader) + if err != nil || headerMethod != request.Method { + if err == nil { + err = errors.New("MCP-Method does not match the JSON-RPC method") + } + return &modernMCPValidationError{Status: http.StatusBadRequest, Code: mcp.ErrorCodeHeaderMismatch, Err: err} + } + if request.Method != "tools/call" { + return nil + } + + var callParams modernMCPCallParams + if err := json.Unmarshal(request.Params, &callParams); err != nil { + return &modernMCPValidationError{Status: http.StatusBadRequest, Code: mcpErrInvalidParams, Err: errors.Wrap(err, "decode mcp call params")} + } + headerName, err := singleMCPHeaderValue(c.Request.Header, mcp.NameHeader) + if err != nil { + return &modernMCPValidationError{Status: http.StatusBadRequest, Code: mcp.ErrorCodeHeaderMismatch, Err: err} + } + decodedName, err := mcp.DecodeMCPHeaderValue(headerName) + if err != nil || decodedName != callParams.Name { + if err == nil { + err = errors.New("MCP-Name does not match tools/call params.name") + } + return &modernMCPValidationError{Status: http.StatusBadRequest, Code: mcp.ErrorCodeHeaderMismatch, Err: err} + } + return nil +} + +// isModernMCPRequestMethod reports whether the tools-only modern endpoint expects a response identifier. +// +// Parameters: +// - method: the JSON-RPC method to classify. +// +// Return values: +// - bool: true for discovery, tool listing, and tool execution requests. +func isModernMCPRequestMethod(method string) bool { + switch method { + case "server/discover", "tools/list", "tools/call": + return true + default: + return false + } +} + +// isValidModernMCPRequestID reports whether a decoded JSON-RPC identifier is a string or integer. +// +// Parameters: +// - id: the JSON-decoded request identifier. +// +// Return values: +// - bool: true for strings and finite integral JSON numbers. +func isValidModernMCPRequestID(id any) bool { + switch typed := id.(type) { + case string: + return true + case float64: + return !math.IsNaN(typed) && !math.IsInf(typed, 0) && math.Trunc(typed) == typed + default: + return false + } +} + +// singleMCPHeaderValue returns one required non-empty header value. +// +// Parameters: +// - headers: the HTTP request headers. +// - name: the case-insensitive header field name. +// +// Return values: +// - string: the sole non-empty value. +// - error: a cardinality or empty-value error. +func singleMCPHeaderValue(headers http.Header, name string) (string, error) { + values := headers.Values(name) + if len(values) != 1 || strings.TrimSpace(values[0]) == "" { + return "", errors.Errorf("%s must occur exactly once", name) + } + return values[0], nil +} + +// validateModernMCPOrigin protects browser-accessible HTTP endpoints from DNS rebinding. +// +// Parameters: +// - request: the inbound HTTP request. +// +// Return values: +// - error: a validation error when a present Origin is malformed or targets another host. +func validateModernMCPOrigin(request *http.Request) error { + origin := strings.TrimSpace(request.Header.Get("Origin")) + if origin == "" { + return nil + } + parsed, err := url.Parse(origin) + if err != nil || (parsed.Scheme != "http" && parsed.Scheme != "https") || parsed.Host == "" || parsed.User != nil || parsed.RawQuery != "" || parsed.Fragment != "" || (parsed.Path != "" && parsed.Path != "/") { + return errors.New("invalid Origin header") + } + if !strings.EqualFold(parsed.Host, strings.TrimSpace(request.Host)) { + return errors.Errorf("Origin host %q does not match the MCP endpoint host", parsed.Host) + } + return nil +} + +// respondModernValidationError writes one modern protocol validation failure. +// +// Parameters: +// - c: the Gin request context receiving the response. +// - id: the decoded JSON-RPC request identifier. +// - err: a modernMCPValidationError or an unexpected validation error. +// +// Return values: none; the function writes a complete JSON-RPC error response. +func respondModernValidationError(c *gin.Context, id any, err error) { + var validationErr *modernMCPValidationError + if !stderrors.As(err, &validationErr) || validationErr == nil { + respondMCPModernError(c, id, http.StatusBadRequest, mcpErrInvalidRequest, err, nil) + return + } + respondMCPModernError(c, id, validationErr.Status, validationErr.Code, validationErr.Err, validationErr.Data) +} + +// handleModernMCPPost serves the tools-only MCP 2026-07-28 method surface. +// +// Parameters: +// - c: the Gin request context containing authentication and request-scoped logging. +// - request: the validated modern JSON-RPC request. +// +// Return values: none; the function writes a complete JSON-RPC response. +func handleModernMCPPost(c *gin.Context, request mcpRPCRequest) { + switch strings.TrimSpace(request.Method) { + case "server/discover": + respondMCPModernResult(c, request.ID, mcp.DiscoveryResult{ + ResultType: mcp.ResultTypeComplete, + SupportedVersions: mcp.SupportedProtocolVersions(), + Capabilities: gin.H{ + "tools": gin.H{"listChanged": false}, + }, + TTLMS: 3600000, + CacheScope: mcp.CacheScopePrivate, + Meta: mcp.ServerResponseMeta(mcpServerName, mcpServerVersion), + }) + case "tools/list": + result, err := listModernMCPToolsPage(gmw.Ctx(c), c, request.Params) + if err != nil { + respondModernDispatchError(c, request.ID, err) + return + } + respondMCPModernResult(c, request.ID, result) + case "tools/call": + var params modernMCPCallParams + if err := json.Unmarshal(request.Params, ¶ms); err != nil { + respondMCPModernError(c, request.ID, http.StatusBadRequest, mcpErrInvalidParams, errors.Wrap(err, "decode mcp call params"), nil) + return + } + result, err := executeModernMCPTool(gmw.Ctx(c), c, params) + if err != nil { + respondModernDispatchError(c, request.ID, err) + return + } + respondMCPModernResult(c, request.ID, result) + default: + if request.ID == nil { + c.AbortWithStatus(http.StatusAccepted) + return + } + respondMCPModernError(c, request.ID, http.StatusNotFound, mcpErrMethodNotFound, errors.Errorf("unsupported method %s", request.Method), nil) + } +} + +// respondModernDispatchError preserves modern validation status while redacting unexpected internal failures. +// +// Parameters: +// - c: The Gin request context receives the JSON-RPC response. +// - id: The JSON-RPC request identifier is reflected in the response. +// - err: The method failure is classified as a validation or internal error. +// +// Return values: none; the function writes the complete JSON-RPC error response. +func respondModernDispatchError(c *gin.Context, id any, err error) { + var validationErr *modernMCPValidationError + if stderrors.As(err, &validationErr) && validationErr != nil { + respondModernValidationError(c, id, validationErr) + return + } + respondMCPModernError(c, id, http.StatusOK, mcpErrInternal, err, nil) +} + +// respondMCPModernResult writes a successful JSON-RPC result with required defaults and server identity. +// +// Parameters: +// - c: the Gin request context receiving the response. +// - id: the decoded JSON-RPC request identifier. +// - result: the result object to normalize and encode. +// +// Return values: none; the function writes a complete JSON-RPC response. +func respondMCPModernResult(c *gin.Context, id any, result any) { + encoded, err := json.Marshal(result) + if err != nil { + respondMCPModernError(c, id, http.StatusOK, mcpErrInternal, errors.Wrap(err, "marshal modern mcp result"), nil) + return + } + var normalized map[string]any + if err := json.Unmarshal(encoded, &normalized); err != nil { + respondMCPModernError(c, id, http.StatusOK, mcpErrInternal, errors.Wrap(err, "normalize modern mcp result"), nil) + return + } + if normalized["resultType"] == nil || normalized["resultType"] == "" { + normalized["resultType"] = mcp.ResultTypeComplete + } + if content, exists := normalized["content"]; exists && content == nil { + delete(normalized, "content") + } + promoteModernResultAlias(normalized, "structured_content", "structuredContent") + promoteModernResultAlias(normalized, "is_error", "isError") + promoteModernResultAlias(normalized, "input_requests", "inputRequests") + promoteModernResultAlias(normalized, "request_state", "requestState") + meta, _ := normalized["_meta"].(map[string]any) + if meta == nil { + meta = make(map[string]any) + } + meta[mcp.MetaServerInfoKey] = gin.H{"name": mcpServerName, "version": mcpServerVersion} + normalized["_meta"] = meta + c.JSON(http.StatusOK, gin.H{"jsonrpc": "2.0", "id": id, "result": normalized}) +} + +// promoteModernResultAlias moves one legacy result field to its current camelCase name. +// +// Parameters: +// - result: the mutable result object. +// - legacyName: the legacy snake_case field name. +// - modernName: the current camelCase field name. +// +// Return values: none; result is updated in place. +func promoteModernResultAlias(result map[string]any, legacyName, modernName string) { + if value, exists := result[legacyName]; exists { + if _, modernExists := result[modernName]; !modernExists { + result[modernName] = value + } + delete(result, legacyName) + } +} + +// respondMCPModernError writes one JSON-RPC error while redacting internal implementation details. +// +// Parameters: +// - c: the Gin request context receiving the response and providing the request-scoped logger. +// - id: the decoded JSON-RPC request identifier, which may be nil for parse failures. +// - status: the HTTP status required by the transport profile. +// - code: the JSON-RPC or MCP error code. +// - err: the underlying validation or internal error. +// - data: optional protocol-safe structured error details. +// +// Return values: none; the function writes a complete JSON-RPC error response. +func respondMCPModernError(c *gin.Context, id any, status int, code int, err error, data any) { + message := "mcp request failed" + if code == mcpErrInternal { + logger := gmw.GetLogger(c) + if err != nil { + logger.Error("mcp internal request failure", zap.Error(err)) + } + message = "internal MCP error" + } else if err != nil { + message = err.Error() + } + errorObject := gin.H{"code": code, "message": message} + if data != nil { + errorObject["data"] = data + } + c.JSON(status, gin.H{"jsonrpc": "2.0", "id": id, "error": errorObject}) +} diff --git a/controller/mcp_proxy_latest_test.go b/controller/mcp_proxy_latest_test.go new file mode 100644 index 0000000000..15c9a50f92 --- /dev/null +++ b/controller/mcp_proxy_latest_test.go @@ -0,0 +1,204 @@ +package controller + +import ( + "bytes" + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + + "github.com/Laisky/one-api/model" + "github.com/Laisky/one-api/relay/mcp" +) + +// TestMCPProxyLatestDiscover verifies modern clients can discover protocol support without initialize. +func TestMCPProxyLatestDiscover(t *testing.T) { + gin.SetMode(gin.TestMode) + router := gin.New() + router.POST("/mcp", MCPProxyLatest) + + body := modernMCPRequestBody(t, "discover-1", "server/discover", map[string]any{}) + request := httptest.NewRequest(http.MethodPost, "/mcp", bytes.NewReader(body)) + request.Header.Set("Content-Type", "application/json") + request.Header.Set(mcp.ProtocolVersionHeader, mcp.ProtocolVersion) + request.Header.Set(mcp.MethodHeader, "server/discover") + response := httptest.NewRecorder() + + router.ServeHTTP(response, request) + require.Equal(t, http.StatusOK, response.Code) + var envelope struct { + Result map[string]any `json:"result"` + } + require.NoError(t, json.Unmarshal(response.Body.Bytes(), &envelope)) + require.Equal(t, mcp.ResultTypeComplete, envelope.Result["resultType"]) + require.Equal(t, []any{mcp.ProtocolVersion}, envelope.Result["supportedVersions"]) + meta := envelope.Result["_meta"].(map[string]any) + serverInfo := meta[mcp.MetaServerInfoKey].(map[string]any) + require.Equal(t, mcpServerName, serverInfo["name"]) + require.Equal(t, mcp.CacheScopePrivate, envelope.Result["cacheScope"]) + require.NotZero(t, envelope.Result["ttlMs"]) +} + +// TestMCPProxyLatestRejectsHeaderMismatch verifies modern mirrored protocol fields are enforced. +func TestMCPProxyLatestRejectsHeaderMismatch(t *testing.T) { + gin.SetMode(gin.TestMode) + router := gin.New() + router.POST("/mcp", MCPProxyLatest) + + body := modernMCPRequestBody(t, "mismatch-1", "server/discover", map[string]any{}) + request := httptest.NewRequest(http.MethodPost, "/mcp", bytes.NewReader(body)) + request.Header.Set("Content-Type", "application/json") + request.Header.Set(mcp.ProtocolVersionHeader, "2026-01-01") + request.Header.Set(mcp.MethodHeader, "server/discover") + response := httptest.NewRecorder() + + router.ServeHTTP(response, request) + require.Equal(t, http.StatusBadRequest, response.Code) + var envelope struct { + Error struct { + Code int `json:"code"` + } `json:"error"` + } + require.NoError(t, json.Unmarshal(response.Body.Bytes(), &envelope)) + require.Equal(t, mcp.ErrorCodeHeaderMismatch, envelope.Error.Code) +} + +// TestMCPProxyLatestRejectsCrossOrigin verifies modern browser requests cannot target a mismatched host. +func TestMCPProxyLatestRejectsCrossOrigin(t *testing.T) { + gin.SetMode(gin.TestMode) + router := gin.New() + router.POST("/mcp", MCPProxyLatest) + + body := modernMCPRequestBody(t, "origin-1", "server/discover", map[string]any{}) + request := httptest.NewRequest(http.MethodPost, "https://gateway.example/mcp", bytes.NewReader(body)) + request.Header.Set("Content-Type", "application/json") + request.Header.Set("Origin", "https://attacker.example") + request.Header.Set(mcp.ProtocolVersionHeader, mcp.ProtocolVersion) + request.Header.Set(mcp.MethodHeader, "server/discover") + response := httptest.NewRecorder() + + router.ServeHTTP(response, request) + require.Equal(t, http.StatusForbidden, response.Code) +} + +// TestMCPProxyLatestDelegatesLegacyInitialize verifies pre-2026 clients retain the existing lifecycle. +func TestMCPProxyLatestDelegatesLegacyInitialize(t *testing.T) { + gin.SetMode(gin.TestMode) + router := gin.New() + router.POST("/mcp", MCPProxyLatest) + + body, err := json.Marshal(map[string]any{ + "jsonrpc": "2.0", + "id": "legacy-1", + "method": "initialize", + "params": map[string]any{ + "protocolVersion": mcp.LegacyProtocolVersionFallback, + "capabilities": map[string]any{}, + "clientInfo": map[string]any{"name": "legacy", "version": "1"}, + }, + }) + require.NoError(t, err) + request := httptest.NewRequest(http.MethodPost, "/mcp", bytes.NewReader(body)) + request.Header.Set("Content-Type", "application/json") + response := httptest.NewRecorder() + + router.ServeHTTP(response, request) + require.Equal(t, http.StatusOK, response.Code) + var envelope struct { + Result map[string]any `json:"result"` + } + require.NoError(t, json.Unmarshal(response.Body.Bytes(), &envelope)) + require.Equal(t, mcp.LegacyProtocolVersionFallback, envelope.Result["protocolVersion"]) +} + +// TestMCPProxyLatestClientServerContract verifies one-api's modern client can drive its own modern server. +func TestMCPProxyLatestClientServerContract(t *testing.T) { + gin.SetMode(gin.TestMode) + router := gin.New() + router.POST("/mcp", MCPProxyLatest) + server := httptest.NewServer(router) + defer server.Close() + + client := mcp.NewStreamableHTTPClient(&model.MCPServer{BaseURL: server.URL + "/mcp"}, nil, 5*time.Second) + result, err := client.DiscoverLatest(context.Background()) + require.NoError(t, err) + require.Equal(t, mcp.ResultTypeComplete, result.ResultType) + require.Equal(t, []string{mcp.ProtocolVersion}, result.SupportedVersions) + require.Equal(t, mcpServerName, result.Meta.ServerInfo.Name) +} + +// modernMCPRequestBody builds a protocol-complete request body for controller tests. +func modernMCPRequestBody(t *testing.T, id, method string, params map[string]any) []byte { + t.Helper() + params = mcp.WithModernMeta(params) + body, err := json.Marshal(map[string]any{ + "jsonrpc": "2.0", + "id": id, + "method": method, + "params": params, + }) + require.NoError(t, err) + return body +} + +// TestValidateModernMCPRequestAcceptsEncodedName verifies mirrored tool names are decoded before comparison. +func TestValidateModernMCPRequestAcceptsEncodedName(t *testing.T) { + gin.SetMode(gin.TestMode) + context, _ := gin.CreateTestContext(httptest.NewRecorder()) + params := mcp.WithModernMeta(map[string]any{ + "name": "天气", + "arguments": map[string]any{}, + }) + encoded, err := json.Marshal(params) + require.NoError(t, err) + context.Request = httptest.NewRequest(http.MethodPost, "https://gateway.example/mcp", nil) + context.Request.Header.Set(mcp.ProtocolVersionHeader, mcp.ProtocolVersion) + context.Request.Header.Set(mcp.MethodHeader, "tools/call") + context.Request.Header.Set(mcp.NameHeader, mcp.EncodeMCPHeaderValue("天气")) + + err = validateModernMCPRequest(context, mcpRPCRequest{ + JSONRPC: "2.0", + ID: "call-1", + Method: "tools/call", + Params: encoded, + }) + require.NoError(t, err) +} + +// TestMCPProxyLatestRejectsMissingClientCapabilities verifies every modern request is self-contained. +func TestMCPProxyLatestRejectsMissingClientCapabilities(t *testing.T) { + gin.SetMode(gin.TestMode) + router := gin.New() + router.POST("/mcp", MCPProxyLatest) + + body, err := json.Marshal(map[string]any{ + "jsonrpc": "2.0", + "id": "missing-capabilities", + "method": "server/discover", + "params": map[string]any{ + "_meta": map[string]any{ + mcp.MetaProtocolVersionKey: mcp.ProtocolVersion, + }, + }, + }) + require.NoError(t, err) + request := httptest.NewRequest(http.MethodPost, "/mcp", bytes.NewReader(body)) + request.Header.Set(mcp.ProtocolVersionHeader, mcp.ProtocolVersion) + request.Header.Set(mcp.MethodHeader, "server/discover") + response := httptest.NewRecorder() + + router.ServeHTTP(response, request) + require.Equal(t, http.StatusBadRequest, response.Code) + var envelope struct { + Error struct { + Code int `json:"code"` + } `json:"error"` + } + require.NoError(t, json.Unmarshal(response.Body.Bytes(), &envelope)) + require.Equal(t, mcpErrInvalidRequest, envelope.Error.Code) +} diff --git a/controller/mcp_proxy_test.go b/controller/mcp_proxy_test.go index 5ec6e67801..99a4b76f9a 100644 --- a/controller/mcp_proxy_test.go +++ b/controller/mcp_proxy_test.go @@ -739,7 +739,7 @@ func TestMCPProxy_FullInspectorHandshake(t *testing.T) { require.Nil(t, initResp["error"]) result, ok := initResp["result"].(map[string]any) require.True(t, ok) - require.Equal(t, mcpProtocolVersion, result["protocolVersion"]) + require.Equal(t, mcp.LegacyProtocolVersionFallback, result["protocolVersion"]) info := result["serverInfo"].(map[string]any) require.Equal(t, mcpServerName, info["name"]) diff --git a/controller/mcp_reviewer_regression_test.go b/controller/mcp_reviewer_regression_test.go new file mode 100644 index 0000000000..b153a27d46 --- /dev/null +++ b/controller/mcp_reviewer_regression_test.go @@ -0,0 +1,182 @@ +package controller + +import ( + "bytes" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/Laisky/one-api/common/ctxkey" + "github.com/Laisky/one-api/model" + "github.com/Laisky/one-api/relay/mcp" +) + +// TestMCPProxyLatestRoutesExplicitLegacyVersionHeaders verifies versioned legacy clients bypass modern metadata validation. +// +// Parameters: +// - t: The test owns the in-memory request and compatibility assertions. +// +// Return values: none; failures are reported through t. +func TestMCPProxyLatestRoutesExplicitLegacyVersionHeaders(t *testing.T) { + body := []byte(`{"jsonrpc":"2.0","id":"legacy-header","method":"initialize","params":{"protocolVersion":"2025-06-18"}}`) + request := httptest.NewRequest(http.MethodPost, "/mcp", bytes.NewReader(body)) + request.Header.Set(mcp.ProtocolVersionHeader, mcp.LegacyProtocolVersionFallback) + context, response := newMCPCallContext(t, 1, "legacy-header") + context.Request = request + MCPProxyLatest(context) + require.Equal(t, http.StatusOK, response.Code) + require.Contains(t, response.Body.String(), `"protocolVersion":"2025-06-18"`) +} + +// TestMCPProxyLatestAllowsLargeLegacyBody verifies legacy compatibility is preserved above the modern four-megabyte limit. +// +// Parameters: +// - t: The test owns the large initialization request and compatibility assertion. +// +// Return values: none; failures are reported through t. +func TestMCPProxyLatestAllowsLargeLegacyBody(t *testing.T) { + padding := strings.Repeat("x", int(modernMCPMaxRequestBytes)+1024) + body, err := json.Marshal(map[string]any{ + "jsonrpc": "2.0", + "id": "large-legacy", + "method": "initialize", + "params": map[string]any{ + "protocolVersion": mcp.LegacyProtocolVersionFallback, + "padding": padding, + }, + }) + require.NoError(t, err) + request := httptest.NewRequest(http.MethodPost, "/mcp", bytes.NewReader(body)) + context, response := newMCPCallContext(t, 1, "large-legacy") + context.Request = request + MCPProxyLatest(context) + require.Equal(t, http.StatusOK, response.Code) + require.Contains(t, response.Body.String(), `"protocolVersion":"2025-06-18"`) +} + +// TestReadBoundedMCPRequestBodyRejectsOverflow verifies unversioned legacy fallback has a hard allocation boundary. +// +// Parameters: +// - t: The test owns the bounded reader and overflow assertion. +// +// Return values: none; failures are reported through t. +func TestReadBoundedMCPRequestBodyRejectsOverflow(t *testing.T) { + _, err := readBoundedMCPRequestBody(strings.NewReader("12345"), 4) + require.Error(t, err) + var tooLarge *mcpRequestTooLargeError + require.ErrorAs(t, err, &tooLarge) + require.Equal(t, int64(4), tooLarge.Limit) +} + +// TestMCPProxyLatestAcceptsLegacyInitializeWithoutParams verifies omitted optional initialize params retain compatibility. +// +// Parameters: +// - t: The test owns the in-memory legacy handshake and response assertions. +// +// Return values: none; failures are reported through t. +func TestMCPProxyLatestAcceptsLegacyInitializeWithoutParams(t *testing.T) { + body := []byte(`{"jsonrpc":"2.0","id":"legacy-no-params","method":"initialize"}`) + request := httptest.NewRequest(http.MethodPost, "/mcp", bytes.NewReader(body)) + context, response := newMCPCallContext(t, 1, "legacy-no-params") + context.Request = request + MCPProxyLatest(context) + require.Equal(t, http.StatusOK, response.Code) + require.Contains(t, response.Body.String(), `"protocolVersion":"2025-11-25"`) +} + +// TestMCPProxyLatestPreservesInvalidCursorError verifies method validation errors are not rewritten as internal failures. +// +// Parameters: +// - t: The test owns the SQLite fixture and modern tools/list request. +// +// Return values: none; failures are reported through t. +func TestMCPProxyLatestPreservesInvalidCursorError(t *testing.T) { + cleanup, fixture := setupMCPProxyTest(t) + defer cleanup() + + body := modernMCPRequestBody(t, "bad-cursor", "tools/list", map[string]any{"cursor": "%%%"}) + response := invokeModernMCPProxy(t, fixture, "bad-cursor", body, func(request *http.Request) { + request.Header.Set(mcp.ProtocolVersionHeader, mcp.ProtocolVersion) + request.Header.Set(mcp.MethodHeader, "tools/list") + }) + require.Equal(t, http.StatusBadRequest, response.Code) + require.Equal(t, mcpErrInvalidParams, decodeMCPErrorCode(t, response)) +} + +// TestMCPProxyLatestPreservesToolHeaderMismatch verifies schema-driven validation errors keep the MCP header code. +// +// Parameters: +// - t: The test owns the SQLite fixture and modern tools/call request. +// +// Return values: none; failures are reported through t. +func TestMCPProxyLatestPreservesToolHeaderMismatch(t *testing.T) { + cleanup, fixture := setupMCPProxyTest(t) + defer cleanup() + + schema := `{"type":"object","properties":{"tenant":{"type":"string","x-mcp-header":"Tenant"}},"required":["tenant"]}` + require.NoError(t, model.DB.Model(&model.MCPTool{}).Where("id = ?", fixture.tool.Id).Updates(map[string]any{ + "input_schema": schema, + "descriptor_json": "", + }).Error) + + body := modernMCPRequestBody(t, "header-mismatch", "tools/call", map[string]any{ + "name": "fake-mcp.echo", + "arguments": map[string]any{"tenant": "acme"}, + }) + response := invokeModernMCPProxy(t, fixture, "header-mismatch", body, func(request *http.Request) { + request.Header.Set(mcp.ProtocolVersionHeader, mcp.ProtocolVersion) + request.Header.Set(mcp.MethodHeader, "tools/call") + request.Header.Set(mcp.NameHeader, "fake-mcp.echo") + }) + require.Equal(t, http.StatusBadRequest, response.Code) + require.Equal(t, mcp.ErrorCodeHeaderMismatch, decodeMCPErrorCode(t, response)) + require.Zero(t, fixture.upstreamHits) +} + +// invokeModernMCPProxy invokes the authenticated modern endpoint with caller-provided headers. +// +// Parameters: +// - t: The test receives fixture and decoding failures. +// - fixture: The fixture supplies the authenticated user and database state. +// - requestID: The request identifier is attached to tracing context. +// - body: The complete JSON-RPC request body is submitted. +// - configure: The optional callback adds method-specific request headers. +// +// Return values: +// - *httptest.ResponseRecorder: The captured HTTP response is returned. +func invokeModernMCPProxy(t *testing.T, fixture *mcpFixture, requestID string, body []byte, configure func(*http.Request)) *httptest.ResponseRecorder { + t.Helper() + context, response := newMCPCallContext(t, fixture.user.Id, requestID) + context.Set(ctxkey.UserObj, fixture.user) + request := httptest.NewRequest(http.MethodPost, "/mcp", bytes.NewReader(body)) + request.Header.Set("Content-Type", "application/json") + if configure != nil { + configure(request) + } + context.Request = request + MCPProxyLatest(context) + return response +} + +// decodeMCPErrorCode extracts one JSON-RPC error code from an HTTP response. +// +// Parameters: +// - t: The test receives JSON decoding failures. +// - response: The recorded response contains a JSON-RPC error envelope. +// +// Return values: +// - int: The decoded error code is returned. +func decodeMCPErrorCode(t *testing.T, response *httptest.ResponseRecorder) int { + t.Helper() + var envelope struct { + Error struct { + Code int `json:"code"` + } `json:"error"` + } + require.NoError(t, json.Unmarshal(response.Body.Bytes(), &envelope)) + return envelope.Error.Code +} diff --git a/docs/manuals/mcp_protocol_2026_07_28.md b/docs/manuals/mcp_protocol_2026_07_28.md new file mode 100644 index 0000000000..9ded7ebb96 --- /dev/null +++ b/docs/manuals/mcp_protocol_2026_07_28.md @@ -0,0 +1,135 @@ +--- +title: MCP 2026-07-28 Protocol Compatibility +version: 1.0 +last_updated: 2026-08-31 +--- + +# MCP 2026-07-28 Protocol Compatibility + +one-api supports the stable MCP `2026-07-28` protocol on both sides of the gateway while retaining the existing initialization-based Streamable HTTP lifecycle for legacy clients and upstream servers. + +## Modern server behavior + +The authenticated `/mcp` endpoint accepts modern requests without an `initialize` exchange. + +Every modern request must include these values in `params._meta`: + +- `io.modelcontextprotocol/protocolVersion`; +- `io.modelcontextprotocol/clientCapabilities`; +- `io.modelcontextprotocol/clientInfo` should also identify the client. + +Every HTTP POST must also include: + +- `MCP-Protocol-Version`, matching `io.modelcontextprotocol/protocolVersion`; +- `Mcp-Method`, matching the JSON-RPC method; +- `Mcp-Name` for `tools/call`, matching `params.name` after protocol decoding; +- any schema-driven `Mcp-Param-*` headers declared through `x-mcp-header`. + +`server/discover` reports supported protocol versions, capabilities, cache metadata, and server identity. Successful results include `resultType` and `result._meta["io.modelcontextprotocol/serverInfo"]`. `tools/list` responses use deterministic ordering, a private cache scope, and a cache TTL. + +The HTTP endpoint validates `Origin` whenever it is present. The Origin host must match the MCP endpoint host, preventing DNS-rebinding access from an unrelated browser origin. + +## Modern client behavior + +The production synchronization and tool-call paths send `2026-07-28` requests directly. They do not initialize a protocol session before a modern request. + +The client: + +1. attaches namespaced modern `_meta` fields to every request; +2. mirrors the protocol method and tool name into HTTP headers; +3. derives `Mcp-Param-*` values from statically reachable `x-mcp-header` annotations; +4. supports JSON and Server-Sent Events responses; +5. accepts `resultType`, `structuredContent`, `isError`, `inputRequests`, and `requestState`; +6. preserves `inputResponses` and `requestState` when retrying a multi-round-trip tool call; +7. excludes malformed `x-mcp-header` tool definitions without hiding valid tools; +8. retries through the legacy initialize/session lifecycle only when the remote endpoint does not return a recognized modern protocol error. + +Authentication failures and recognized modern errors such as `HeaderMismatch`, `MissingRequiredClientCapability`, and `UnsupportedProtocolVersion` are returned to the caller rather than being misclassified as legacy-server failures. + +## Credential transport policy + +Configured remote MCP endpoints that receive an API key, an authorization or cookie header, custom authentication headers, or URL user information must use HTTPS. The same rule is enforced both when server configuration is validated and immediately before outbound network I/O, so previously persisted or directly constructed clients cannot bypass it. + +Plaintext HTTP remains available for unauthenticated endpoints. Credentialed HTTP is allowed only for `localhost`, `127.0.0.0/8`, and `::1` loopback endpoints used by local development and integration tests. The exception is host-bound: a credentialed redirect must remain on the exact original origin. HTTPS-to-HTTP redirects are always rejected, and credentialed redirects must preserve the original scheme, hostname, and effective port. + +## Schema-driven parameter headers + +An input-schema property may define an `x-mcp-header` annotation when its type is `string`, `integer`, or `boolean`. + +```json +{ + "type": "object", + "properties": { + "tenant": { + "type": "string", + "x-mcp-header": "Tenant-ID" + } + } +} +``` + +For `{"tenant":"acme"}`, the client sends: + +```text +Mcp-Param-Tenant-ID: acme +``` + +Header annotations must be unique case-insensitively and reachable from the schema root through `properties` only. Annotated integers must remain within the JavaScript-safe integer range. A missing or `null` parameter produces no header. + +Values that are not safe plain HTTP field values, or that already resemble the sentinel, use this exact encoding: + +```text +=?base64??= +``` + +The same encoding applies to `Mcp-Name`. The server decodes mirrored values, independently derives the expected parameter values from the JSON body, and rejects missing, repeated, malformed, or mismatched headers with error code `-32020`. + +## Legacy compatibility + +Requests without modern protocol metadata continue through the original handler. Existing legacy behavior remains available, including: + +- `initialize` and `notifications/initialized`; +- `Mcp-Session-Id`; +- the existing `2025-06-18` Streamable HTTP compatibility path; +- legacy tool-result aliases such as `is_error` and `structured_content`. + +The original `ListTools` and `CallTool` methods remain available for code that deliberately requires the legacy lifecycle. one-api's active synchronization and proxy execution paths use the modern-first methods. + +## Error and status behavior + +Modern transport-level validation uses HTTP status codes in addition to JSON-RPC errors: + +| Condition | HTTP status | JSON-RPC code | +| --- | ---: | ---: | +| Header/body mismatch | 400 | `-32020` | +| Missing required client capability | 400 | `-32021` | +| Unsupported protocol version | 400 | `-32022` | +| Invalid Origin | 403 | `-32600` | +| Unknown modern method | 404 | `-32601` | + +Tool execution failures that occur after a valid request remain JSON-RPC errors with HTTP 200, preserving normal RPC semantics. + +## Validation coverage + +Regression tests cover: + +- handshake-free modern tool listing; +- modern-to-legacy client fallback; +- namespaced request and result metadata; +- protocol, method, encoded tool-name, and schema-driven parameter headers; +- nested extraction, null omission, safe-integer enforcement, and exact Base64 sentinel encoding; +- exclusion of invalid tool header schemas; +- modern and legacy result-field aliases; +- multi-round-trip request fields; +- `server/discover`, cache metadata, Origin validation, and header mismatch rejection; +- legacy `initialize` delegation through the same `/mcp` endpoint; +- configuration-time and runtime rejection of credentialed remote plaintext HTTP; +- loopback-only credentialed HTTP compatibility and redirect downgrade protection. + +## Specification references + +- MCP changelog: +- Versioning: +- Streamable HTTP: +- Server discovery: +- Tools: diff --git a/model/mcp_server.go b/model/mcp_server.go index 6fd3be50b0..1c7f122c34 100644 --- a/model/mcp_server.go +++ b/model/mcp_server.go @@ -1,6 +1,7 @@ package model import ( + "net" "net/url" "strings" "time" @@ -109,6 +110,9 @@ func (s *MCPServer) NormalizeAndValidate() error { if s.AuthType == "" { s.AuthType = MCPAuthTypeNone } + if parsedURL.Scheme == "http" && s.HasSensitiveCredentials() && !isLoopbackMCPServerHost(parsedURL.Hostname()) { + return errkind.InvalidRequestErr(errors.New("credentialed mcp server base_url must use https unless it targets a loopback host")) + } if s.AutoSyncIntervalMinutes == 0 { s.AutoSyncIntervalMinutes = 60 @@ -125,6 +129,79 @@ func (s *MCPServer) NormalizeAndValidate() error { return nil } +// HasSensitiveCredentials reports whether the server configuration carries credentials that must not traverse remote plaintext HTTP. +// +// Parameters: none. +// +// Return values: +// - bool: True is returned when the API key, URL user information, or configured authentication headers contain sensitive data. +func (s *MCPServer) HasSensitiveCredentials() bool { + if s == nil { + return false + } + if strings.TrimSpace(s.APIKey) != "" { + return true + } + if parsedURL, err := url.Parse(strings.TrimSpace(s.BaseURL)); err == nil && parsedURL.User != nil && parsedURL.User.String() != "" { + return true + } + + for key, value := range s.Headers { + if strings.TrimSpace(value) == "" { + continue + } + normalizedKey := strings.ToLower(strings.TrimSpace(key)) + if isSensitiveMCPServerHeaderName(normalizedKey) { + return true + } + switch normalizedKey { + case "accept", "accept-encoding", "content-type", "user-agent", "mcp-protocol-version", "mcp-session-id": + continue + default: + // Arbitrary configured headers can implement custom authentication even + // when their names do not contain a conventional credential token. + return true + } + } + return false +} + +// isSensitiveMCPServerHeaderName reports whether a configured header name conventionally carries credentials. +// +// Parameters: +// - name: The configured HTTP header name is inspected case-insensitively. +// +// Return values: +// - bool: True is returned when the header name contains a credential-bearing token. +func isSensitiveMCPServerHeaderName(name string) bool { + lower := strings.ToLower(strings.TrimSpace(name)) + if lower == "" { + return false + } + for _, token := range []string{"authorization", "proxy-authorization", "api_key", "apikey", "token", "secret", "password", "passwd", "x-api-key", "cookie"} { + if strings.Contains(lower, token) { + return true + } + } + return false +} + +// isLoopbackMCPServerHost reports whether a hostname is restricted to the local machine. +// +// Parameters: +// - hostname: The URL hostname is checked as localhost or a loopback IP address. +// +// Return values: +// - bool: True is returned only for localhost and IP loopback addresses. +func isLoopbackMCPServerHost(hostname string) bool { + hostname = strings.TrimSpace(hostname) + if strings.EqualFold(hostname, "localhost") { + return true + } + address := net.ParseIP(hostname) + return address != nil && address.IsLoopback() +} + // ValidateToolPricing ensures per-tool pricing values are non-negative. func (s *MCPServer) ValidateToolPricing() error { for name, pricing := range s.ToolPricing { diff --git a/model/mcp_server_security_test.go b/model/mcp_server_security_test.go new file mode 100644 index 0000000000..133203ced1 --- /dev/null +++ b/model/mcp_server_security_test.go @@ -0,0 +1,105 @@ +package model + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +// TestMCPServerNormalizeAndValidateCredentialTransport verifies remote credentials require HTTPS while loopback development endpoints remain available. +// +// Parameters: +// - t: The test owns table-driven server configurations and policy assertions. +// +// Return values: none; failures are reported through t. +func TestMCPServerNormalizeAndValidateCredentialTransport(t *testing.T) { + testCases := []struct { + name string + server MCPServer + wantError string + }{ + { + name: "remote bearer credentials require HTTPS", + server: MCPServer{ + Name: "remote-bearer", + BaseURL: "http://mcp.example.com/mcp", + AuthType: MCPAuthTypeBearer, + APIKey: "secret", + }, + wantError: "must use https", + }, + { + name: "remote sensitive header requires HTTPS", + server: MCPServer{ + Name: "remote-header", + BaseURL: "http://mcp.example.com/mcp", + Headers: JSONStringMap{"Authorization": "Bearer secret"}, + }, + wantError: "must use https", + }, + { + name: "remote custom authentication headers require HTTPS", + server: MCPServer{ + Name: "remote-custom", + BaseURL: "http://mcp.example.com/mcp", + AuthType: MCPAuthTypeCustomHeaders, + Headers: JSONStringMap{"X-Tenant-Identity": "secret"}, + }, + wantError: "must use https", + }, + { + name: "remote URL user information requires HTTPS", + server: MCPServer{ + Name: "remote-userinfo", + BaseURL: "http://user:secret@mcp.example.com/mcp", + }, + wantError: "must use https", + }, + { + name: "credentialed HTTPS endpoint is accepted", + server: MCPServer{ + Name: "secure", + BaseURL: "https://mcp.example.com/mcp", + AuthType: MCPAuthTypeAPIKey, + APIKey: "secret", + }, + }, + { + name: "credentialed IPv4 loopback HTTP endpoint is accepted", + server: MCPServer{ + Name: "loopback-v4", + BaseURL: "http://127.0.0.1:8080/mcp", + AuthType: MCPAuthTypeAPIKey, + APIKey: "secret", + }, + }, + { + name: "credentialed IPv6 loopback HTTP endpoint is accepted", + server: MCPServer{ + Name: "loopback-v6", + BaseURL: "http://[::1]:8080/mcp", + AuthType: MCPAuthTypeAPIKey, + APIKey: "secret", + }, + }, + { + name: "unauthenticated remote HTTP endpoint remains compatible", + server: MCPServer{ + Name: "public-http", + BaseURL: "http://mcp.example.com/mcp", + }, + }, + } + + for _, testCase := range testCases { + t.Run(testCase.name, func(t *testing.T) { + testCase.server.AutoSyncIntervalMinutes = 60 + err := testCase.server.NormalizeAndValidate() + if testCase.wantError == "" { + require.NoError(t, err) + return + } + require.ErrorContains(t, err, testCase.wantError) + }) + } +} diff --git a/model/mcp_tool.go b/model/mcp_tool.go index a3766e494f..c85f55eb13 100644 --- a/model/mcp_tool.go +++ b/model/mcp_tool.go @@ -2,7 +2,7 @@ package model import "strings" -// MCPTool stores tool metadata synchronized from MCP servers. +// MCPTool stores searchable MCP tool metadata and the complete upstream wire descriptor. type MCPTool struct { Id int `json:"-"` UUID string `json:"uuid" gorm:"type:char(36);column:uuid"` @@ -12,16 +12,21 @@ type MCPTool struct { DisplayName string `json:"display_name" gorm:"type:varchar(128)"` Description string `json:"description" gorm:"type:text"` InputSchema string `json:"input_schema" gorm:"type:text"` + DescriptorJSON string `json:"-" gorm:"type:text;column:descriptor_json"` DefaultPricing ToolPricingLocalJSON `json:"default_pricing" gorm:"type:text"` Status int `json:"status" gorm:"type:int;default:1"` CreatedAt int64 `json:"created_at" gorm:"bigint;autoCreateTime:milli"` UpdatedAt int64 `json:"updated_at" gorm:"bigint;autoUpdateTime:milli"` } -// NormalizeName ensures tool names are normalized consistently. +// NormalizeName trims transport-insignificant whitespace while preserving the case-sensitive wire name. +// +// Parameters: none. +// +// Return values: none; the receiver is updated in place when it is non-nil. func (t *MCPTool) NormalizeName() { if t == nil { return } - t.Name = strings.TrimSpace(strings.ToLower(t.Name)) + t.Name = strings.TrimSpace(t.Name) } diff --git a/model/mcp_tool_store.go b/model/mcp_tool_store.go index 149c271db6..271ea65504 100644 --- a/model/mcp_tool_store.go +++ b/model/mcp_tool_store.go @@ -41,7 +41,19 @@ var MCPToolSortFields = map[string]string{ "updated_at": "updated_at", } -// ListMCPTools returns MCP tools filtered by server id and status. +// ListMCPTools returns MCP tools filtered by server, status, pagination, and ordering. +// +// Parameters: +// - serverID: owning server id; non-positive means every server. +// - status: optional status filter. +// - offset: pagination offset. +// - limit: page size; non-positive means unlimited. +// - sortBy: whitelisted sort column. +// - sortOrder: ascending or descending order. +// +// Return values: +// - []*MCPTool: matching tools. +// - error: a wrapped database error when the query fails. func ListMCPTools(serverID int, status *int, offset int, limit int, sortBy string, sortOrder string) ([]*MCPTool, error) { return SearchMCPTools(serverID, status, "", offset, limit, sortBy, sortOrder) } @@ -86,7 +98,15 @@ func SearchMCPTools(serverID int, status *int, keyword string, offset int, limit return tools, nil } -// CountMCPTools returns the total number of MCP tools matching filters. +// CountMCPTools returns the total number of MCP tools matching server and status filters. +// +// Parameters: +// - serverID: owning server id; non-positive means every server. +// - status: optional status filter. +// +// Return values: +// - int64: the number of matching tools. +// - error: a wrapped database error when the count fails. func CountMCPTools(serverID int, status *int) (int64, error) { return CountSearchedMCPTools(serverID, status, "") } @@ -121,8 +141,12 @@ func CountSearchedMCPTools(serverID int, status *int, keyword string) (int64, er // GetMCPToolsByServerID fetches tools for a specific server. // -// The empty case returns a non-nil zero-length slice so HTTP handlers that -// pass the result straight to c.JSON marshal "data" as [] rather than null. +// Parameters: +// - serverID: the positive internal id of the owning MCP server. +// +// Return values: +// - []*MCPTool: a non-nil slice containing every stored tool for the server. +// - error: a validation or wrapped database error. func GetMCPToolsByServerID(serverID int) ([]*MCPTool, error) { if serverID <= 0 { return nil, errors.New("server id is invalid") @@ -134,26 +158,45 @@ func GetMCPToolsByServerID(serverID int) ([]*MCPTool, error) { return tools, nil } -// UpsertMCPTools replaces tools for a server with the provided list. +// UpsertMCPTools atomically replaces a server catalog while preserving exact wire names and descriptors. +// +// Parameters: +// - serverID: the positive internal id of the owning MCP server. +// - serverUUID: the stable owning server UUID copied onto every tool row. +// - tools: the complete replacement catalog; nil entries are ignored. +// +// Return values: +// - error: a validation or wrapped transactional database error. func UpsertMCPTools(serverID int, serverUUID string, tools []*MCPTool) error { if serverID <= 0 { return errors.New("server id is invalid") } - if err := DB.Where("server_id = ?", serverID).Delete(&MCPTool{}).Error; err != nil { - return errors.Wrap(err, "clear mcp tools") + if DB == nil { + return errors.New("database is not initialized") } - for _, tool := range tools { - if tool == nil { - continue - } - tool.ServerId = serverID - if serverUUID != "" { - tool.ServerUUID = &serverUUID + if err := DB.Transaction(func(tx *gorm.DB) error { + if err := tx.Where("server_id = ?", serverID).Delete(&MCPTool{}).Error; err != nil { + return errors.Wrap(err, "clear mcp tools") } - tool.NormalizeName() - if err := DB.Create(tool).Error; err != nil { - return errors.Wrap(err, "create mcp tool") + for _, tool := range tools { + if tool == nil { + continue + } + tool.ServerId = serverID + if serverUUID != "" { + tool.ServerUUID = &serverUUID + } + tool.NormalizeName() + if tool.Name == "" { + return errors.New("mcp tool name is required") + } + if err := tx.Create(tool).Error; err != nil { + return errors.Wrapf(err, "create mcp tool %q", tool.Name) + } } + return nil + }); err != nil { + return errors.Wrap(err, "replace mcp tool catalog") } return nil } diff --git a/relay/mcp/client.go b/relay/mcp/client.go index 7de825b2d0..07ece29742 100644 --- a/relay/mcp/client.go +++ b/relay/mcp/client.go @@ -4,7 +4,6 @@ import ( "bytes" "context" "encoding/json" - "io" "net/http" "strings" "sync" @@ -20,29 +19,22 @@ import ( "github.com/Laisky/one-api/model" ) -// Client defines MCP operations required by the aggregator. +// Client defines the MCP tool-listing and tool-call operations required by the aggregator. type Client interface { ListTools(ctx context.Context) ([]ToolDescriptor, error) CallTool(ctx context.Context, name string, arguments any) (*CallToolResult, error) } -// StreamableHTTPClient implements MCP client calls over the Streamable HTTP -// transport. The client performs the protocol handshake (initialize + -// notifications/initialized) lazily on first use, captures any -// server-issued Mcp-Session-Id, and supports both JSON and SSE response -// content types. +// StreamableHTTPClient implements modern and legacy MCP over the Streamable HTTP transport. type StreamableHTTPClient struct { BaseURL string Headers map[string]string Timeout time.Duration Logger glog.Logger - // serverRef identifies the MCP server this client talks to. It is logged on - // every client log line because the caller's logger, even when - // request-bound, carries the user/token/channel identity but not the MCP - // server's. serverRef identity.MCPServerRef + headerMu sync.RWMutex initMu sync.Mutex initialized bool sessionID string @@ -50,70 +42,99 @@ type StreamableHTTPClient struct { } const ( - mcpProtocolVersionHeader = "Mcp-Protocol-Version" - mcpSessionIDHeader = "Mcp-Session-Id" - mcpDefaultProtocolVersion = "2025-06-18" + mcpProtocolVersionHeader = ProtocolVersionHeader + mcpSessionIDHeader = SessionIDHeader + mcpDefaultProtocolVersion = LegacyProtocolVersion mcpAcceptHeaderValue = "application/json, text/event-stream" mcpClientName = "one-api-mcp-client" mcpClientVersion = "1.0.0" ) // NewStreamableHTTPClient constructs a StreamableHTTPClient from MCP server metadata. +// +// Parameters: +// - server: the configured upstream MCP server. +// - headers: request headers that override server-level configured headers. +// - timeout: the per-request HTTP timeout. +// +// Return values: +// - *StreamableHTTPClient: a client configured for modern-first requests and legacy fallback. func NewStreamableHTTPClient(server *model.MCPServer, headers map[string]string, timeout time.Duration) *StreamableHTTPClient { return newStreamableHTTPClient(server, headers, timeout, nil) } -// NewStreamableHTTPClientWithLogger constructs a StreamableHTTPClient with logging enabled. +// NewStreamableHTTPClientWithLogger constructs a StreamableHTTPClient with request-aware logging. +// +// Parameters: +// - server: the configured upstream MCP server. +// - headers: request headers that override server-level configured headers. +// - timeout: the per-request HTTP timeout. +// - logger: the logger used for sanitized transport diagnostics. +// +// Return values: +// - *StreamableHTTPClient: a client configured for modern-first requests and legacy fallback. func NewStreamableHTTPClientWithLogger(server *model.MCPServer, headers map[string]string, timeout time.Duration, logger glog.Logger) *StreamableHTTPClient { return newStreamableHTTPClient(server, headers, timeout, logger) } -// newStreamableHTTPClient constructs a StreamableHTTPClient from MCP server metadata. -// The Mcp-Session-Id header is intentionally NOT pre-populated — per the -// Streamable HTTP transport spec, the session id is issued by the server in -// the initialize response and only then attached to subsequent requests. +// newStreamableHTTPClient merges configuration without preselecting a legacy protocol session. +// +// Parameters: +// - server: the configured upstream MCP server. +// - headers: request headers that override server-level configured headers. +// - timeout: the per-request HTTP timeout. +// - logger: the optional logger used for sanitized transport diagnostics. +// +// Return values: +// - *StreamableHTTPClient: a new client whose legacy lifecycle is initialized lazily. func newStreamableHTTPClient(server *model.MCPServer, headers map[string]string, timeout time.Duration, logger glog.Logger) *StreamableHTTPClient { merged := make(map[string]string) - for k, v := range server.Headers { - merged[k] = v - } - for k, v := range headers { - merged[k] = v + if server != nil { + for key, value := range server.Headers { + merged[key] = value + } } - if _, ok := merged[mcpProtocolVersionHeader]; !ok { - merged[mcpProtocolVersionHeader] = mcpDefaultProtocolVersion + for key, value := range headers { + merged[key] = value } - if _, ok := merged["Accept"]; !ok { + delete(merged, mcpProtocolVersionHeader) + delete(merged, mcpSessionIDHeader) + if _, exists := merged["Accept"]; !exists { merged["Accept"] = mcpAcceptHeaderValue } - switch strings.ToLower(server.AuthType) { - case model.MCPAuthTypeBearer: - if server.APIKey != "" { - merged["Authorization"] = "Bearer " + server.APIKey - } - case model.MCPAuthTypeAPIKey: - if server.APIKey != "" { - merged["X-API-Key"] = server.APIKey + if server != nil { + switch strings.ToLower(server.AuthType) { + case model.MCPAuthTypeBearer: + if server.APIKey != "" { + merged["Authorization"] = "Bearer " + server.APIKey + } + case model.MCPAuthTypeAPIKey: + if server.APIKey != "" { + merged["X-API-Key"] = server.APIKey + } } } - return &StreamableHTTPClient{ - BaseURL: strings.TrimSpace(server.BaseURL), - Headers: merged, - Timeout: timeout, - Logger: logger, - serverRef: server.Ref(), + client := &StreamableHTTPClient{ + Headers: merged, + Timeout: timeout, + Logger: logger, } + if server != nil { + client.BaseURL = strings.TrimSpace(server.BaseURL) + client.serverRef = server.Ref() + } + return client } -// Initialize performs the MCP protocol handshake: sends an `initialize` -// request and the corresponding `notifications/initialized` notification. -// Captures the server-issued Mcp-Session-Id (if any) and the negotiated -// protocol version, then attaches both to subsequent requests. +// Initialize performs the preferred 2025-11-25 initialize exchange and records the negotiated legacy state. +// +// Parameters: +// - ctx: the request context controlling cancellation and deadlines. // -// Safe to call concurrently and idempotent — the handshake runs at most -// once per client instance. +// Return values: +// - error: a wrapped transport, negotiation, or notification error when initialization cannot complete. func (c *StreamableHTTPClient) Initialize(ctx context.Context) error { if c == nil { return errors.New("mcp client is nil") @@ -132,62 +153,102 @@ func (c *StreamableHTTPClient) Initialize(ctx context.Context) error { "version": mcpClientVersion, }, } - var initResult struct { ProtocolVersion string `json:"protocolVersion"` Capabilities map[string]any `json:"capabilities"` ServerInfo map[string]any `json:"serverInfo"` } - respHeaders, err := c.doRPCRaw(ctx, "initialize", initParams, &initResult) + responseHeaders, err := c.doRPCRaw(ctx, "initialize", initParams, &initResult) if err != nil { return errors.Wrap(err, "mcp initialize") } - if sid := respHeaders.Get(mcpSessionIDHeader); sid != "" { - c.sessionID = sid - c.Headers[mcpSessionIDHeader] = sid + negotiatedVersion := strings.TrimSpace(initResult.ProtocolVersion) + if negotiatedVersion == "" { + negotiatedVersion = mcpDefaultProtocolVersion } - if initResult.ProtocolVersion != "" { - c.protocolVersion = initResult.ProtocolVersion - c.Headers[mcpProtocolVersionHeader] = initResult.ProtocolVersion + if !IsLegacyProtocolVersion(negotiatedVersion) { + return errors.Errorf("mcp initialize negotiated unsupported legacy version %q", negotiatedVersion) + } + c.protocolVersion = negotiatedVersion + c.setClientHeader(mcpProtocolVersionHeader, negotiatedVersion) + if sessionID := strings.TrimSpace(responseHeaders.Get(mcpSessionIDHeader)); sessionID != "" { + c.sessionID = sessionID + c.setClientHeader(mcpSessionIDHeader, sessionID) } if err := c.sendNotification(ctx, "notifications/initialized", nil); err != nil { - // Notification failure is non-fatal — log and proceed so a server - // that diverges on this notification does not block tool calls. if c.Logger != nil { c.Logger.Warn("mcp notifications/initialized failed", append(c.serverRef.Zap(), zap.Error(err))...) } } - c.initialized = true return nil } -// ListTools calls the MCP tools/list method. +// ListTools lists every tool through the negotiated legacy lifecycle and follows pagination cursors. +// +// Parameters: +// - ctx: the request context controlling cancellation and deadlines. +// +// Return values: +// - []ToolDescriptor: tools collected across every legacy tools/list page. +// - error: a wrapped initialization, transport, pagination, or decoding error. func (c *StreamableHTTPClient) ListTools(ctx context.Context) ([]ToolDescriptor, error) { if err := c.Initialize(ctx); err != nil { return nil, err } - var result struct { - Tools []ToolDescriptor `json:"tools"` - } - if err := c.doRPC(ctx, "tools/list", nil, &result); err != nil { - return nil, errors.Wrap(err, "mcp rpc tools/list") + tools := make([]ToolDescriptor, 0) + seenCursors := make(map[string]struct{}) + cursor := "" + for page := 0; ; page++ { + if page >= 1000 { + return nil, errors.New("legacy mcp tools/list exceeded 1000 pages") + } + var params map[string]any + if cursor != "" { + params = map[string]any{"cursor": cursor} + } + var result ListToolsResult + if err := c.doRPC(ctx, "tools/list", params, &result); err != nil { + return nil, errors.Wrap(err, "mcp rpc tools/list") + } + tools = append(tools, result.Tools...) + nextCursor := strings.TrimSpace(result.NextCursor) + if nextCursor == "" { + return tools, nil + } + if _, exists := seenCursors[nextCursor]; exists { + return nil, errors.Errorf("legacy mcp tools/list repeated cursor %q", nextCursor) + } + seenCursors[nextCursor] = struct{}{} + cursor = nextCursor } - return result.Tools, nil } -// CallTool invokes a MCP tool by name. +// CallTool invokes one exact case-sensitive tool through the negotiated legacy lifecycle. +// +// Parameters: +// - ctx: the request context controlling cancellation and deadlines. +// - name: the exact upstream tool name. +// - arguments: a JSON-compatible argument object; nil becomes an empty object. +// +// Return values: +// - *CallToolResult: the decoded upstream result. +// - error: a wrapped initialization, validation, transport, or decoding error. func (c *StreamableHTTPClient) CallTool(ctx context.Context, name string, arguments any) (*CallToolResult, error) { if err := c.Initialize(ctx); err != nil { return nil, err } + argumentMap, err := normalizeToolArguments(arguments) + if err != nil { + return nil, errors.Wrap(err, "normalize legacy mcp tool arguments") + } params := map[string]any{ "name": name, - "arguments": arguments, + "arguments": argumentMap, } var result CallToolResult if err := c.doRPC(ctx, "tools/call", params, &result); err != nil { @@ -196,28 +257,40 @@ func (c *StreamableHTTPClient) CallTool(ctx context.Context, name string, argume return &result, nil } -// doRPC performs a JSON-RPC call and discards the response headers. +// doRPC performs one legacy JSON-RPC request and discards the response headers. +// +// Parameters: +// - ctx: the request context controlling cancellation and deadlines. +// - method: the legacy JSON-RPC method. +// - params: optional structured request parameters. +// - out: the destination for a successful result, or nil to discard it. +// +// Return values: +// - error: a wrapped transport, correlation, protocol, or decoding error. func (c *StreamableHTTPClient) doRPC(ctx context.Context, method string, params any, out any) error { _, err := c.doRPCRaw(ctx, method, params, out) return err } -// doRPCRaw performs a JSON-RPC call and returns the response headers, which -// the initialize handshake needs to read the Mcp-Session-Id assigned by the -// server. Handles both `application/json` and `text/event-stream` response -// content types per the Streamable HTTP transport spec. +// doRPCRaw performs one correlated legacy JSON-RPC request and returns the response headers. +// +// Parameters: +// - ctx: the request context controlling cancellation and deadlines. +// - method: the legacy JSON-RPC method. +// - params: optional structured request parameters. +// - out: the destination for a successful result, or nil to discard it. +// +// Return values: +// - http.Header: response headers used by initialize to capture session state. +// - error: a wrapped transport, size, correlation, protocol, or decoding error. func (c *StreamableHTTPClient) doRPCRaw(ctx context.Context, method string, params any, out any) (http.Header, error) { if c == nil { return nil, errors.New("mcp client is nil") } - // Per JSON-RPC 2.0, the params member MUST be a structured value (Object - // or Array) when present. Strict validators (e.g. the TypeScript MCP - // SDK's Zod schema used by aas-ee/open-web-search) reject `"params":null` - // with -32700 "Parse error: Invalid JSON-RPC message". Omit the field - // entirely when no params were supplied. + requestID := random.GetUUID() payload := map[string]any{ "jsonrpc": "2.0", - "id": random.GetUUID(), + "id": requestID, "method": method, } if params != nil { @@ -228,52 +301,51 @@ func (c *StreamableHTTPClient) doRPCRaw(ctx context.Context, method string, para return nil, errors.Wrap(err, "marshal mcp request") } - client := &http.Client{Timeout: c.Timeout} req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.BaseURL, bytes.NewReader(data)) if err != nil { return nil, errors.Wrap(err, "create mcp request") } req.Header.Set("Content-Type", "application/json") - for key, value := range c.Headers { + for key, value := range c.headerSnapshot() { req.Header.Set(key, value) } - c.debugLogRequest(method, req.Header, data) + client := c.httpClient() resp, err := client.Do(req) if err != nil { return nil, errors.Wrap(err, "send mcp request") } defer resp.Body.Close() - body, readErr := io.ReadAll(resp.Body) - if readErr != nil { - return resp.Header, errors.Wrap(readErr, "read mcp response body") + body, err := readMCPResponseBody(resp.Body) + if err != nil { + return resp.Header, errors.Wrap(err, "read mcp response body") } c.debugLogResponse(method, resp, body) - if resp.StatusCode < 200 || resp.StatusCode >= 300 { - return resp.Header, errors.Errorf("mcp request failed with status %d: %s", resp.StatusCode, strings.TrimSpace(string(body))) - } - - contentType := strings.ToLower(resp.Header.Get("Content-Type")) - if strings.Contains(contentType, "text/event-stream") { - jsonBody, perr := parseSSEResponse(body) - if perr != nil { - return resp.Header, errors.Wrap(perr, "parse mcp sse response") + if strings.Contains(strings.ToLower(resp.Header.Get("Content-Type")), "text/event-stream") { + body, err = extractMCPResponseEnvelope(body, requestID) + if err != nil { + return resp.Header, errors.Wrap(err, "parse mcp SSE response") } - body = jsonBody } - - var envelope struct { - Result json.RawMessage `json:"result"` - Error map[string]any `json:"error"` + if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices { + return resp.Header, decodeModernProtocolError(resp.StatusCode, body) } - if err := json.Unmarshal(body, &envelope); err != nil { - return resp.Header, errors.Wrap(err, "decode mcp response") + + envelope, err := parseMCPResponseEnvelope(body, requestID) + if err != nil { + return resp.Header, err } if envelope.Error != nil { - return resp.Header, errors.Errorf("mcp error: %v", envelope.Error) + return resp.Header, &ProtocolError{ + HTTPStatus: resp.StatusCode, + Code: envelope.Error.Code, + Message: envelope.Error.Message, + Data: envelope.Error.Data, + Body: strings.TrimSpace(string(body)), + } } if out == nil { return resp.Header, nil @@ -284,9 +356,15 @@ func (c *StreamableHTTPClient) doRPCRaw(ctx context.Context, method string, para return resp.Header, nil } -// sendNotification sends a JSON-RPC notification (no `id` field). Per the -// Streamable HTTP transport spec, the server replies with HTTP 202 and an -// empty body — there is no JSON-RPC envelope to parse. +// sendNotification sends one legacy JSON-RPC notification without a request identifier. +// +// Parameters: +// - ctx: the request context controlling cancellation and deadlines. +// - method: the notification method. +// - params: optional structured notification parameters. +// +// Return values: +// - error: a wrapped transport, size, or HTTP status error. func (c *StreamableHTTPClient) sendNotification(ctx context.Context, method string, params any) error { if c == nil { return errors.New("mcp client is nil") @@ -303,39 +381,82 @@ func (c *StreamableHTTPClient) sendNotification(ctx context.Context, method stri return errors.Wrap(err, "marshal mcp notification") } - client := &http.Client{Timeout: c.Timeout} req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.BaseURL, bytes.NewReader(data)) if err != nil { return errors.Wrap(err, "create mcp notification request") } req.Header.Set("Content-Type", "application/json") - for key, value := range c.Headers { + for key, value := range c.headerSnapshot() { req.Header.Set(key, value) } - c.debugLogRequest(method, req.Header, data) + client := c.httpClient() resp, err := client.Do(req) if err != nil { return errors.Wrap(err, "send mcp notification") } defer resp.Body.Close() - body, _ := io.ReadAll(resp.Body) + body, err := readMCPResponseBody(resp.Body) + if err != nil { + return errors.Wrap(err, "read mcp notification response body") + } c.debugLogResponse(method, resp, body) - - if resp.StatusCode < 200 || resp.StatusCode >= 300 { + if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices { return errors.Errorf("mcp notification failed with status %d: %s", resp.StatusCode, strings.TrimSpace(string(body))) } return nil } -// parseSSEResponse extracts the JSON payload from a Server-Sent Events body. -// The MCP Streamable HTTP transport allows a server to reply to a single -// request/response with one SSE event whose `data:` field contains the -// JSON-RPC envelope. Multi-line `data:` fields are concatenated with `\n` -// per the SSE spec. +// headerSnapshot returns a copy of mutable client headers for one HTTP request. +// +// Parameters: none. +// +// Return values: +// - map[string]string: an isolated header map safe for request construction. +func (c *StreamableHTTPClient) headerSnapshot() map[string]string { + if c == nil { + return nil + } + c.headerMu.RLock() + defer c.headerMu.RUnlock() + snapshot := make(map[string]string, len(c.Headers)) + for key, value := range c.Headers { + snapshot[key] = value + } + return snapshot +} + +// setClientHeader updates one internally managed header under the client header lock. +// +// Parameters: +// - key: the HTTP header name. +// - value: the replacement value; an empty value removes the header. +// +// Return values: none. +func (c *StreamableHTTPClient) setClientHeader(key, value string) { + if c == nil { + return + } + c.headerMu.Lock() + defer c.headerMu.Unlock() + if value == "" { + delete(c.Headers, key) + return + } + c.Headers[key] = value +} + +// parseSSEResponse extracts data fields from a single finite SSE message for compatibility tests. +// +// Parameters: +// - body: a finite Server-Sent Events response body. +// +// Return values: +// - []byte: concatenated data fields. +// - error: an error when the body contains no data field. func parseSSEResponse(body []byte) ([]byte, error) { - var dataLines []string + dataLines := make([]string, 0) for _, raw := range strings.Split(string(body), "\n") { line := strings.TrimRight(raw, "\r") if !strings.HasPrefix(line, "data:") { @@ -350,49 +471,64 @@ func parseSSEResponse(body []byte) ([]byte, error) { } // debugLogRequest records sanitized outbound MCP request metadata and payload. +// +// Parameters: +// - method: the JSON-RPC method. +// - headers: the outbound HTTP headers. +// - body: the encoded request body. +// +// Return values: none. func (c *StreamableHTTPClient) debugLogRequest(method string, headers http.Header, body []byte) { if c == nil || c.Logger == nil { return } - sanitizedHeaders := sanitizeHeadersForLog(headers) - sanitizedBody := sanitizeBodyForLog(body) c.Logger.Debug("mcp outbound request", append(c.serverRef.Zap(), zap.String("method", method), zap.String("url", c.BaseURL), - zap.Any("headers", sanitizedHeaders), + zap.Any("headers", sanitizeHeadersForLog(headers)), zap.Int("body_bytes", len(body)), - zap.String("body", sanitizedBody), + zap.String("body", sanitizeBodyForLog(body)), )...) } // debugLogResponse records sanitized inbound MCP response metadata and payload. +// +// Parameters: +// - method: the JSON-RPC method. +// - resp: the HTTP response metadata. +// - body: the bounded response body. +// +// Return values: none. func (c *StreamableHTTPClient) debugLogResponse(method string, resp *http.Response, body []byte) { if c == nil || c.Logger == nil || resp == nil { return } - sanitizedHeaders := sanitizeHeadersForLog(resp.Header) - sanitizedBody := sanitizeBodyForLog(body) c.Logger.Debug("mcp inbound response", append(c.serverRef.Zap(), zap.String("method", method), zap.String("url", c.BaseURL), zap.Int("status_code", resp.StatusCode), - zap.Any("headers", sanitizedHeaders), + zap.Any("headers", sanitizeHeadersForLog(resp.Header)), zap.Int("body_bytes", len(body)), - zap.String("body", sanitizedBody), + zap.String("body", sanitizeBodyForLog(body)), )...) } -// sanitizeHeadersForLog redacts sensitive header values for logging. +// sanitizeHeadersForLog redacts sensitive header values before structured logging. +// +// Parameters: +// - headers: the HTTP headers to sanitize. +// +// Return values: +// - map[string]string: flattened headers with sensitive values redacted. func sanitizeHeadersForLog(headers http.Header) map[string]string { if headers == nil { return nil } sanitized := make(map[string]string, len(headers)) for key, values := range headers { - lower := strings.ToLower(strings.TrimSpace(key)) - if isSensitiveKey(lower) { + if isSensitiveKey(strings.ToLower(strings.TrimSpace(key))) { sanitized[key] = "" continue } @@ -401,7 +537,13 @@ func sanitizeHeadersForLog(headers http.Header) map[string]string { return sanitized } -// sanitizeBodyForLog returns a sanitized body string for logging. +// sanitizeBodyForLog redacts secrets and binary-like fields from a request or response body. +// +// Parameters: +// - body: the raw request or response body. +// +// Return values: +// - string: a sanitized JSON string or a safe textual placeholder. func sanitizeBodyForLog(body []byte) string { if len(body) == 0 { return "" @@ -428,7 +570,14 @@ func sanitizeBodyForLog(body []byte) string { return string(encoded) } -// scrubJSONValue redacts sensitive or binary-like data from JSON values. +// scrubJSONValue recursively redacts sensitive and binary-like JSON values. +// +// Parameters: +// - value: the decoded JSON value to sanitize. +// - keyHint: the parent field name used for redaction heuristics. +// +// Return values: +// - any: the sanitized JSON-compatible value. func scrubJSONValue(value any, keyHint string) any { if value == nil { return nil @@ -449,8 +598,8 @@ func scrubJSONValue(value any, keyHint string) any { } return typed case []any: - for idx, inner := range typed { - typed[idx] = scrubJSONValue(inner, keyHint) + for index, inner := range typed { + typed[index] = scrubJSONValue(inner, keyHint) } return typed case string: @@ -467,13 +616,18 @@ func scrubJSONValue(value any, keyHint string) any { } } -// isSensitiveKey reports whether a key is likely to contain secrets. +// isSensitiveKey reports whether a field or header name is likely to contain a secret. +// +// Parameters: +// - key: a normalized field or header name. +// +// Return values: +// - bool: true when the value must be redacted from logs. func isSensitiveKey(key string) bool { if key == "" { return false } - sensitive := []string{"authorization", "proxy-authorization", "api_key", "apikey", "token", "secret", "password", "passwd", "x-api-key"} - for _, token := range sensitive { + for _, token := range []string{"authorization", "proxy-authorization", "api_key", "apikey", "token", "secret", "password", "passwd", "x-api-key", "cookie"} { if strings.Contains(key, token) { return true } @@ -481,13 +635,18 @@ func isSensitiveKey(key string) bool { return false } -// isBinaryKey reports whether a key is likely to contain binary payloads. +// isBinaryKey reports whether a JSON field name is likely to contain binary payload data. +// +// Parameters: +// - key: a normalized JSON field name. +// +// Return values: +// - bool: true when the value should be omitted from logs. func isBinaryKey(key string) bool { if key == "" { return false } - tokens := []string{"image", "audio", "video", "binary", "base64", "bytes", "file", "blob"} - for _, token := range tokens { + for _, token := range []string{"image", "audio", "video", "binary", "base64", "bytes", "file", "blob"} { if strings.Contains(key, token) { return true } @@ -495,7 +654,13 @@ func isBinaryKey(key string) bool { return false } -// isLikelyBinary performs a heuristic check for binary payloads. +// isLikelyBinary reports whether a response body contains invalid UTF-8 or control-heavy data. +// +// Parameters: +// - body: the raw body to inspect. +// +// Return values: +// - bool: true when logging the body as text would be unsafe or unhelpful. func isLikelyBinary(body []byte) bool { if len(body) == 0 { return false @@ -504,18 +669,24 @@ func isLikelyBinary(body []byte) bool { return true } nonPrintable := 0 - for _, r := range body { - if r == '\n' || r == '\r' || r == '\t' { + for _, value := range body { + if value == '\n' || value == '\r' || value == '\t' { continue } - if r < 0x20 || r == 0x7f { + if value < 0x20 || value == 0x7f { nonPrintable++ } } return nonPrintable > len(body)/20 } -// isLikelyBase64 checks whether a string looks like base64 data. +// isLikelyBase64 reports whether a long string resembles encoded binary data. +// +// Parameters: +// - value: the string to inspect. +// +// Return values: +// - bool: true when the value should be omitted from logs. func isLikelyBase64(value string) bool { trimmed := strings.TrimSpace(value) if len(trimmed) < 128 { @@ -524,8 +695,8 @@ func isLikelyBase64(value string) bool { if strings.HasPrefix(trimmed, "data:") { return true } - for _, r := range trimmed { - if r == '=' || r == '+' || r == '/' || r == '-' || r == '_' || (r >= '0' && r <= '9') || (r >= 'A' && r <= 'Z') || (r >= 'a' && r <= 'z') { + for _, character := range trimmed { + if character == '=' || character == '+' || character == '/' || character == '-' || character == '_' || (character >= '0' && character <= '9') || (character >= 'A' && character <= 'Z') || (character >= 'a' && character <= 'z') { continue } return false diff --git a/relay/mcp/client_latest.go b/relay/mcp/client_latest.go new file mode 100644 index 0000000000..df68556432 --- /dev/null +++ b/relay/mcp/client_latest.go @@ -0,0 +1,388 @@ +package mcp + +import ( + "bytes" + "context" + "encoding/json" + "net/http" + "strings" + + "github.com/Laisky/errors/v2" + "github.com/Laisky/zap" + + "github.com/Laisky/one-api/common/random" +) + +// DiscoverLatest calls server/discover using the MCP 2026-07-28 request model. +// +// Parameters: +// - ctx: the request context controlling cancellation and deadlines. +// +// Return values: +// - *DiscoveryResult: the server's modern discovery result. +// - error: a wrapped transport, protocol, or decoding error. +func (c *StreamableHTTPClient) DiscoverLatest(ctx context.Context) (*DiscoveryResult, error) { + if c == nil { + return nil, errors.New("mcp client is nil") + } + var result DiscoveryResult + if err := c.doModernRPC(ctx, "server/discover", nil, "", nil, &result); err != nil { + return nil, errors.Wrap(err, "mcp server/discover") + } + if result.ResultType == "" { + result.ResultType = ResultTypeComplete + } + return &result, nil +} + +// ListToolsLatest lists every tool through modern pagination and falls back to the legacy lifecycle when required. +// +// Parameters: +// - ctx: the request context controlling cancellation and deadlines. +// +// Return values: +// - []ToolDescriptor: validated tools collected across all result pages. +// - error: a wrapped transport, protocol, pagination, or descriptor-validation error. +func (c *StreamableHTTPClient) ListToolsLatest(ctx context.Context) ([]ToolDescriptor, error) { + if c == nil { + return nil, errors.New("mcp client is nil") + } + if c.legacyInitialized() { + tools, err := c.ListTools(ctx) + if err != nil { + return nil, err + } + return c.normalizeAndFilterToolDescriptors(tools), nil + } + + tools := make([]ToolDescriptor, 0) + seenCursors := make(map[string]struct{}) + cursor := "" + for page := 0; ; page++ { + if page >= 1000 { + return nil, errors.New("mcp tools/list exceeded 1000 pages") + } + var params map[string]any + if cursor != "" { + params = map[string]any{"cursor": cursor} + } + var result ListToolsResult + err := c.doModernRPC(ctx, "tools/list", params, "", nil, &result) + if err != nil { + if page == 0 && IsModernFallbackCandidate(err) { + legacyTools, legacyErr := c.ListTools(ctx) + if legacyErr != nil { + return nil, legacyErr + } + return c.normalizeAndFilterToolDescriptors(legacyTools), nil + } + return nil, errors.Wrap(err, "mcp modern tools/list") + } + tools = append(tools, result.Tools...) + nextCursor := strings.TrimSpace(result.NextCursor) + if nextCursor == "" { + break + } + if _, exists := seenCursors[nextCursor]; exists { + return nil, errors.Errorf("mcp tools/list repeated cursor %q", nextCursor) + } + seenCursors[nextCursor] = struct{}{} + cursor = nextCursor + } + return c.normalizeAndFilterToolDescriptors(tools), nil +} + +// CallToolLatest invokes one tool with MCP 2026-07-28 and no multi-round-trip state. +// +// Parameters: +// - ctx: the request context controlling cancellation and deadlines. +// - name: the exact case-sensitive upstream tool name. +// - arguments: a JSON-compatible argument object; nil becomes an empty object. +// +// Return values: +// - *CallToolResult: the normalized upstream result. +// - error: a wrapped transport, protocol, validation, or decoding error. +func (c *StreamableHTTPClient) CallToolLatest(ctx context.Context, name string, arguments any) (*CallToolResult, error) { + return c.CallToolLatestWithDescriptor(ctx, ToolDescriptor{Name: name}, arguments) +} + +// CallToolLatestWithDescriptor invokes one tool and derives schema-driven request headers. +// +// Parameters: +// - ctx: the request context controlling cancellation and deadlines. +// - tool: the upstream descriptor containing the exact name and input schema. +// - arguments: a JSON-compatible argument object; nil becomes an empty object. +// +// Return values: +// - *CallToolResult: the normalized upstream result. +// - error: a wrapped transport, protocol, validation, or decoding error. +func (c *StreamableHTTPClient) CallToolLatestWithDescriptor(ctx context.Context, tool ToolDescriptor, arguments any) (*CallToolResult, error) { + return c.CallToolLatestWithOptions(ctx, tool, arguments, CallToolRequestOptions{}) +} + +// CallToolLatestWithOptions invokes one tool and preserves multi-round-trip retry fields. +// +// Parameters: +// - ctx: the request context controlling cancellation and deadlines. +// - tool: the upstream descriptor containing the exact name and input schema. +// - arguments: a JSON-compatible argument object; nil becomes an empty object. +// - options: optional input responses and opaque request state from an input_required result. +// +// Return values: +// - *CallToolResult: the normalized upstream result. +// - error: a wrapped transport, protocol, validation, or decoding error. +func (c *StreamableHTTPClient) CallToolLatestWithOptions(ctx context.Context, tool ToolDescriptor, arguments any, options CallToolRequestOptions) (*CallToolResult, error) { + if c == nil { + return nil, errors.New("mcp client is nil") + } + name := strings.TrimSpace(tool.Name) + if name == "" { + return nil, errors.New("mcp tool name is required") + } + if c.legacyInitialized() { + if hasCallToolRequestOptions(options) { + return nil, errors.New("legacy mcp servers cannot accept multi-round-trip tool retry fields") + } + return c.CallTool(ctx, name, arguments) + } + + argumentMap, err := normalizeToolArguments(arguments) + if err != nil { + return nil, errors.Wrap(err, "normalize mcp tool arguments") + } + parameterHeaders, err := ToolArgumentHeaders(tool.InputSchema, argumentMap) + if err != nil { + return nil, errors.Wrapf(err, "derive mcp headers for tool %q", name) + } + params := map[string]any{ + "name": name, + "arguments": argumentMap, + } + if options.InputResponses != nil { + params["inputResponses"] = options.InputResponses + } + if options.RequestState != "" { + params["requestState"] = options.RequestState + } + var result CallToolResult + err = c.doModernRPC(ctx, "tools/call", params, name, parameterHeaders, &result) + if err != nil { + if IsModernFallbackCandidate(err) && !hasCallToolRequestOptions(options) { + return c.CallTool(ctx, name, argumentMap) + } + return nil, errors.Wrapf(err, "mcp modern tools/call %s", name) + } + return NormalizeCallToolResult(&result), nil +} + +// normalizeAndFilterToolDescriptors supplies legacy defaults and excludes invalid HTTP tool definitions. +// +// Parameters: +// - tools: descriptors collected from one or more tools/list pages. +// +// Return values: +// - []ToolDescriptor: valid descriptors with a non-nil object input schema. +func (c *StreamableHTTPClient) normalizeAndFilterToolDescriptors(tools []ToolDescriptor) []ToolDescriptor { + for index := range tools { + if tools[index].InputSchema == nil { + tools[index].InputSchema = map[string]any{"type": "object"} + } + } + validTools, rejected := FilterValidToolDescriptors(tools) + if c != nil && c.Logger != nil { + for _, rejection := range rejected { + c.Logger.Warn("excluding invalid mcp tool descriptor", + append(c.serverRef.Zap(), zap.String("tool", rejection.Name), zap.Error(rejection.Err))...) + } + } + return validTools +} + +// hasCallToolRequestOptions reports whether a call carries modern multi-round-trip fields. +// +// Parameters: +// - options: the optional retry fields supplied by the caller. +// +// Return values: +// - bool: true when legacy fallback cannot represent the request. +func hasCallToolRequestOptions(options CallToolRequestOptions) bool { + return options.InputResponses != nil || options.RequestState != "" +} + +// legacyInitialized reports whether the client has committed to the legacy session lifecycle. +// +// Parameters: none. +// +// Return values: +// - bool: true after a successful initialize exchange. +func (c *StreamableHTTPClient) legacyInitialized() bool { + if c == nil { + return false + } + c.initMu.Lock() + defer c.initMu.Unlock() + return c.initialized +} + +// doModernRPC performs one stateless MCP 2026-07-28 JSON-RPC request. +// +// Parameters: +// - ctx: the request context controlling cancellation and deadlines. +// - method: the exact JSON-RPC method mirrored into the HTTP request. +// - params: method-specific parameters before required modern metadata is attached. +// - name: an optional tool or resource name mirrored into the HTTP request. +// - parameterHeaders: schema-derived MCP parameter headers for this request only. +// - out: the destination for a successful result, or nil to discard it. +// +// Return values: +// - error: a wrapped transport, size, correlation, protocol, or decoding error. +func (c *StreamableHTTPClient) doModernRPC(ctx context.Context, method string, params map[string]any, name string, parameterHeaders http.Header, out any) error { + if c == nil { + return errors.New("mcp client is nil") + } + requestID := random.GetUUID() + payload := map[string]any{ + "jsonrpc": "2.0", + "id": requestID, + "method": method, + "params": WithModernMeta(params), + } + data, err := json.Marshal(payload) + if err != nil { + return errors.Wrap(err, "marshal modern mcp request") + } + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.BaseURL, bytes.NewReader(data)) + if err != nil { + return errors.Wrap(err, "create modern mcp request") + } + req.Header.Set("Content-Type", "application/json") + for key, value := range c.headerSnapshot() { + req.Header.Set(key, value) + } + for key := range req.Header { + if strings.HasPrefix(strings.ToLower(key), strings.ToLower(ParameterHeaderPrefix)) { + req.Header.Del(key) + } + } + requestHeaderSet := func(key, value string) { + if value == "" { + req.Header.Del(key) + return + } + req.Header.Set(key, value) + } + requestHeaderSet("Accept", mcpAcceptHeaderValue) + req.Header.Del(SessionIDHeader) + req.Header.Del("Last-Event-ID") + requestHeaderSet(ProtocolVersionHeader, ProtocolVersion) + requestHeaderSet(MethodHeader, method) + if name != "" { + requestHeaderSet(NameHeader, EncodeMCPHeaderValue(name)) + } else { + req.Header.Del(NameHeader) + } + for key, values := range parameterHeaders { + if len(values) != 1 { + return errors.Errorf("mcp parameter header %s must contain exactly one value", key) + } + req.Header.Set(key, values[0]) + } + + c.debugLogRequest(method, req.Header, data) + client := c.httpClient() + resp, err := client.Do(req) + if err != nil { + return errors.Wrap(err, "send modern mcp request") + } + defer resp.Body.Close() + + body, err := readMCPResponseBody(resp.Body) + if err != nil { + return errors.Wrap(err, "read modern mcp response body") + } + c.debugLogResponse(method, resp, body) + + if strings.Contains(strings.ToLower(resp.Header.Get("Content-Type")), "text/event-stream") { + body, err = extractMCPResponseEnvelope(body, requestID) + if err != nil { + return errors.Wrap(err, "parse modern mcp SSE response") + } + } + + if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices { + return decodeModernProtocolError(resp.StatusCode, body) + } + + envelope, err := parseMCPResponseEnvelope(body, requestID) + if err != nil { + return err + } + if envelope.Error != nil { + return &ProtocolError{ + HTTPStatus: resp.StatusCode, + Code: envelope.Error.Code, + Message: envelope.Error.Message, + Data: envelope.Error.Data, + Body: strings.TrimSpace(string(body)), + } + } + if out == nil { + return nil + } + if err := json.Unmarshal(envelope.Result, out); err != nil { + return errors.Wrap(err, "unmarshal modern mcp result") + } + return nil +} + +// decodeModernProtocolError parses an HTTP-level modern MCP error response. +// +// Parameters: +// - status: the non-success HTTP status returned by the peer. +// - body: the bounded response body. +// +// Return values: +// - error: a ProtocolError retaining the HTTP and JSON-RPC details available to fallback policy. +func decodeModernProtocolError(status int, body []byte) error { + protocolErr := &ProtocolError{HTTPStatus: status, Body: strings.TrimSpace(string(body))} + var envelope mcpJSONRPCEnvelope + if err := json.Unmarshal(body, &envelope); err == nil && envelope.Error != nil { + protocolErr.Code = envelope.Error.Code + protocolErr.Message = envelope.Error.Message + protocolErr.Data = envelope.Error.Data + } + return protocolErr +} + +// normalizeToolArguments converts JSON-compatible arguments into a non-nil object. +// +// Parameters: +// - arguments: nil, a map, or another value that JSON can decode as an object. +// +// Return values: +// - map[string]any: the normalized argument object. +// - error: a wrapped encoding or type error when arguments are not a JSON object. +func normalizeToolArguments(arguments any) (map[string]any, error) { + if arguments == nil { + return map[string]any{}, nil + } + if object, ok := arguments.(map[string]any); ok { + if object == nil { + return map[string]any{}, nil + } + return object, nil + } + encoded, err := json.Marshal(arguments) + if err != nil { + return nil, errors.Wrap(err, "marshal mcp tool arguments") + } + var object map[string]any + if err := json.Unmarshal(encoded, &object); err != nil { + return nil, errors.Wrap(err, "decode mcp tool arguments object") + } + if object == nil { + return nil, errors.New("mcp tool arguments must be an object") + } + return object, nil +} diff --git a/relay/mcp/client_latest_test.go b/relay/mcp/client_latest_test.go new file mode 100644 index 0000000000..9b9c19f1e4 --- /dev/null +++ b/relay/mcp/client_latest_test.go @@ -0,0 +1,298 @@ +package mcp + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/require" + + "github.com/Laisky/one-api/model" +) + +// TestListToolsLatestUsesModernRequestModel verifies tool discovery no longer requires initialize on modern servers. +func TestListToolsLatestUsesModernRequestModel(t *testing.T) { + var requests int + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + requests++ + require.Equal(t, ProtocolVersion, r.Header.Get(ProtocolVersionHeader)) + require.Equal(t, "tools/list", r.Header.Get(MethodHeader)) + require.Empty(t, r.Header.Get(SessionIDHeader)) + var request struct { + ID any `json:"id"` + Method string `json:"method"` + Params map[string]any `json:"params"` + } + require.NoError(t, json.NewDecoder(r.Body).Decode(&request)) + require.Equal(t, "tools/list", request.Method) + meta := request.Params["_meta"].(map[string]any) + require.Equal(t, ProtocolVersion, meta[MetaProtocolVersionKey]) + require.NotNil(t, meta[MetaClientCapabilitiesKey]) + w.Header().Set("Content-Type", "application/json") + writeMCPTestJSON(t, w, map[string]any{ + "jsonrpc": "2.0", + "id": request.ID, + "result": map[string]any{ + "resultType": "complete", + "tools": []any{map[string]any{ + "name": "echo", + "title": "Echo", + "description": "echo input", + "inputSchema": map[string]any{"type": "object"}, + }}, + }, + }) + })) + defer server.Close() + + client := NewStreamableHTTPClient(&model.MCPServer{BaseURL: server.URL}, nil, 5*time.Second) + tools, err := client.ListToolsLatest(context.Background()) + require.NoError(t, err) + require.Equal(t, 1, requests) + require.Equal(t, "Echo", tools[0].Title) +} + +// TestCallToolLatestSendsSchemaDrivenHeaders verifies named and parameter headers mirror the JSON body. +func TestCallToolLatestSendsSchemaDrivenHeaders(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + require.Equal(t, "tools/call", r.Header.Get(MethodHeader)) + require.Equal(t, "echo", r.Header.Get(NameHeader)) + require.Equal(t, "acme", r.Header.Get("Mcp-Param-Tenant")) + var request struct { + ID any `json:"id"` + } + require.NoError(t, json.NewDecoder(r.Body).Decode(&request)) + w.Header().Set("Content-Type", "application/json") + writeMCPTestJSON(t, w, map[string]any{ + "jsonrpc": "2.0", + "id": request.ID, + "result": map[string]any{ + "resultType": "complete", + "content": []any{map[string]any{"type": "text", "text": "ok"}}, + "structuredContent": map[string]any{"tenant": "acme"}, + "isError": false, + }, + }) + })) + defer server.Close() + + client := NewStreamableHTTPClient(&model.MCPServer{BaseURL: server.URL}, nil, 5*time.Second) + result, err := client.CallToolLatestWithDescriptor(context.Background(), ToolDescriptor{ + Name: "echo", + InputSchema: map[string]any{ + "properties": map[string]any{ + "tenant": map[string]any{"type": "string", "x-mcp-header": "Tenant"}, + }, + }, + }, map[string]any{"tenant": "acme"}) + require.NoError(t, err) + require.Equal(t, ResultTypeComplete, result.ResultType) + require.Equal(t, map[string]any{"tenant": "acme"}, result.StructuredContent) +} + +// TestListToolsLatestFallsBackToLegacy verifies old Streamable HTTP servers retain initialize compatibility. +func TestListToolsLatestFallsBackToLegacy(t *testing.T) { + var mu sync.Mutex + methods := make([]string, 0, 4) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var request struct { + ID any `json:"id"` + Method string `json:"method"` + } + require.NoError(t, json.NewDecoder(r.Body).Decode(&request)) + mu.Lock() + methods = append(methods, request.Method) + mu.Unlock() + + w.Header().Set("Content-Type", "application/json") + if request.Method == "tools/list" && r.Header.Get(MethodHeader) != "" { + w.WriteHeader(http.StatusBadRequest) + writeMCPTestText(t, w, `{"jsonrpc":"2.0","id":null,"error":{"code":-32600,"message":"initialize required"}}`) + return + } + switch request.Method { + case "initialize": + writeMCPTestJSON(t, w, map[string]any{ + "jsonrpc": "2.0", + "id": request.ID, + "result": map[string]any{ + "protocolVersion": LegacyProtocolVersionFallback, + "capabilities": map[string]any{"tools": map[string]any{}}, + "serverInfo": map[string]any{"name": "legacy", "version": "1"}, + }, + }) + case "notifications/initialized": + w.WriteHeader(http.StatusAccepted) + case "tools/list": + writeMCPTestJSON(t, w, map[string]any{ + "jsonrpc": "2.0", + "id": request.ID, + "result": map[string]any{"tools": []any{map[string]any{"name": "legacy.echo"}}}, + }) + } + })) + defer server.Close() + + client := NewStreamableHTTPClient(&model.MCPServer{BaseURL: server.URL}, nil, 5*time.Second) + tools, err := client.ListToolsLatest(context.Background()) + require.NoError(t, err) + require.Equal(t, "legacy.echo", tools[0].Name) + mu.Lock() + defer mu.Unlock() + require.Equal(t, []string{"tools/list", "initialize", "notifications/initialized", "tools/list"}, methods) +} + +// TestCallToolResultAcceptsModernAndLegacyAliases verifies the wire model preserves both protocol eras. +func TestCallToolResultAcceptsModernAndLegacyAliases(t *testing.T) { + var modern CallToolResult + require.NoError(t, json.Unmarshal([]byte(`{"resultType":"input_required","inputRequests":{"confirm":{"method":"elicitation/create"}},"requestState":"opaque"}`), &modern)) + require.Equal(t, ResultTypeInputRequired, modern.ResultType) + require.Equal(t, map[string]any{"method": "elicitation/create"}, modern.InputRequests["confirm"]) + require.Equal(t, "opaque", modern.RequestState) + + var legacy CallToolResult + require.NoError(t, json.Unmarshal([]byte(`{"is_error":true,"structured_content":{"step":2}}`), &legacy)) + require.Empty(t, legacy.ResultType) + require.True(t, legacy.IsError) +} + +// TestCallToolLatestEncodesNameAndMRTRFields verifies non-ASCII tool names and retry state follow the modern wire format. +func TestCallToolLatestEncodesNameAndMRTRFields(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + require.Equal(t, "=?base64?5aSp5rCU?=", r.Header.Get(NameHeader)) + var request struct { + ID any `json:"id"` + Params map[string]any `json:"params"` + } + require.NoError(t, json.NewDecoder(r.Body).Decode(&request)) + require.Equal(t, "state", request.Params["requestState"]) + require.Equal(t, map[string]any{"answer": "yes"}, request.Params["inputResponses"]) + w.Header().Set("Content-Type", "application/json") + writeMCPTestJSON(t, w, map[string]any{ + "jsonrpc": "2.0", + "id": request.ID, + "result": map[string]any{ + "resultType": "complete", + "content": []any{}, + }, + }) + })) + defer server.Close() + + client := NewStreamableHTTPClient(&model.MCPServer{BaseURL: server.URL}, nil, 5*time.Second) + _, err := client.CallToolLatestWithOptions(context.Background(), ToolDescriptor{ + Name: "天气", + InputSchema: map[string]any{"type": "object"}, + }, map[string]any{}, CallToolRequestOptions{ + InputResponses: map[string]any{"answer": "yes"}, + RequestState: "state", + }) + require.NoError(t, err) +} + +// TestListToolsLatestExcludesInvalidHeaderSchemas verifies malformed tools are filtered instead of failing the entire catalog. +func TestListToolsLatestExcludesInvalidHeaderSchemas(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var request struct { + ID any `json:"id"` + } + require.NoError(t, json.NewDecoder(r.Body).Decode(&request)) + w.Header().Set("Content-Type", "application/json") + writeMCPTestJSON(t, w, map[string]any{ + "jsonrpc": "2.0", + "id": request.ID, + "result": map[string]any{ + "resultType": "complete", + "tools": []any{ + map[string]any{"name": "valid", "inputSchema": map[string]any{"type": "object"}}, + map[string]any{"name": "invalid", "inputSchema": map[string]any{"items": map[string]any{"x-mcp-header": "Bad"}}}, + }, + }, + }) + })) + defer server.Close() + + client := NewStreamableHTTPClient(&model.MCPServer{BaseURL: server.URL}, nil, 5*time.Second) + tools, err := client.ListToolsLatest(context.Background()) + require.NoError(t, err) + require.Len(t, tools, 1) + require.Equal(t, "valid", tools[0].Name) +} + +// TestRecognizedModernErrorsDoNotFallBack verifies validation errors are not misclassified as legacy servers. +func TestRecognizedModernErrorsDoNotFallBack(t *testing.T) { + for _, code := range []int{ErrorCodeHeaderMismatch, ErrorCodeMissingRequiredClientCapability, ErrorCodeUnsupportedProtocolVersion} { + require.False(t, IsModernFallbackCandidate(&ProtocolError{HTTPStatus: http.StatusBadRequest, Code: code})) + } + require.False(t, IsModernFallbackCandidate(&ProtocolError{HTTPStatus: http.StatusNotFound, Code: -32601})) + require.True(t, IsModernFallbackCandidate(&ProtocolError{HTTPStatus: http.StatusBadRequest, Code: -32600, Message: "initialize required"})) +} + +// TestListToolsLatestFollowsPagination verifies modern tool catalogs are collected across cursors. +func TestListToolsLatestFollowsPagination(t *testing.T) { + var calls int + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls++ + var request struct { + ID any `json:"id"` + Params map[string]any `json:"params"` + } + require.NoError(t, json.NewDecoder(r.Body).Decode(&request)) + result := map[string]any{ + "resultType": "complete", + "tools": []any{map[string]any{ + "name": fmt.Sprintf("tool-%d", calls), + "inputSchema": map[string]any{"type": "object"}, + }}, + "ttlMs": 1000, + "cacheScope": "private", + } + if calls == 1 { + require.Empty(t, request.Params["cursor"]) + result["nextCursor"] = "page-2" + } else { + require.Equal(t, "page-2", request.Params["cursor"]) + } + w.Header().Set("Content-Type", "application/json") + writeMCPTestJSON(t, w, map[string]any{"jsonrpc": "2.0", "id": request.ID, "result": result}) + })) + defer server.Close() + + client := NewStreamableHTTPClient(&model.MCPServer{BaseURL: server.URL}, nil, 5*time.Second) + tools, err := client.ListToolsLatest(context.Background()) + require.NoError(t, err) + require.Equal(t, []string{"tool-1", "tool-2"}, []string{tools[0].Name, tools[1].Name}) + require.Equal(t, 2, calls) +} + +// writeMCPTestJSON encodes one mock HTTP response and reports failures at the write site. +// +// Parameters: +// - t: The owning test receives any encoding failure. +// - writer: The mock HTTP response writer receives the JSON payload. +// - value: The JSON-compatible response value is encoded. +// +// Return values: none; failures are reported through t. +func writeMCPTestJSON(t *testing.T, writer http.ResponseWriter, value any) { + t.Helper() + require.NoError(t, json.NewEncoder(writer).Encode(value)) +} + +// writeMCPTestText writes one mock HTTP response and reports failures at the write site. +// +// Parameters: +// - t: The owning test receives any write failure. +// - writer: The mock HTTP response writer receives the text payload. +// - value: The response text is written exactly once. +// +// Return values: none; failures are reported through t. +func writeMCPTestText(t *testing.T, writer http.ResponseWriter, value string) { + t.Helper() + _, err := fmt.Fprint(writer, value) + require.NoError(t, err) +} diff --git a/relay/mcp/headers.go b/relay/mcp/headers.go new file mode 100644 index 0000000000..152b723387 --- /dev/null +++ b/relay/mcp/headers.go @@ -0,0 +1,551 @@ +package mcp + +import ( + "encoding/base64" + "encoding/json" + "math" + "math/big" + "net/http" + "sort" + "strconv" + "strings" + "unicode/utf8" + + "github.com/Laisky/errors/v2" +) + +const ( + mcpBase64SentinelPrefix = "=?base64?" + mcpBase64SentinelSuffix = "?=" + maxMCPHeaderInteger = int64(1<<53 - 1) + mcpHeaderNumberEpsilon = 1e-9 +) + +type toolHeaderBinding struct { + HeaderName string + Path []string + ValueType string +} + +// ToolDescriptorRejection describes one tool excluded because its x-mcp-header schema is invalid. +type ToolDescriptorRejection struct { + Name string + Err error +} + +// ToolArgumentHeaders derives MCP-Param-* headers from reachable x-mcp-header annotations. +// +// Parameters: +// - schema: the tool input schema containing optional x-mcp-header annotations. +// - arguments: the normalized JSON object supplied to tools/call. +// +// Return values: +// - http.Header: exactly one encoded header value for every present annotated argument. +// - error: a wrapped schema, lookup, type, range, or encoding error. +func ToolArgumentHeaders(schema map[string]any, arguments map[string]any) (http.Header, error) { + bindings, err := collectToolHeaderBindings(schema) + if err != nil { + return nil, errors.Wrap(err, "collect mcp tool header bindings") + } + + headers := make(http.Header, len(bindings)) + for _, binding := range bindings { + value, exists := lookupToolArgument(arguments, binding.Path) + if !exists || value == nil { + continue + } + encoded, err := formatToolHeaderValue(value, binding.ValueType) + if err != nil { + return nil, errors.Wrapf(err, "format mcp tool header %s", binding.HeaderName) + } + headers.Set(ParameterHeaderPrefix+binding.HeaderName, encoded) + } + return headers, nil +} + +// ValidateToolArgumentHeaders verifies mirrored parameter headers against the schema and JSON arguments. +// +// Parameters: +// - requestHeaders: the complete inbound HTTP request headers. +// - schema: the selected tool input schema. +// - arguments: the normalized tools/call argument object. +// +// Return values: +// - error: a wrapped schema, cardinality, decoding, numeric, or value-mismatch error. +func ValidateToolArgumentHeaders(requestHeaders http.Header, schema map[string]any, arguments map[string]any) error { + bindings, err := collectToolHeaderBindings(schema) + if err != nil { + return errors.Wrap(err, "collect expected mcp tool header bindings") + } + expected, err := ToolArgumentHeaders(schema, arguments) + if err != nil { + return errors.Wrap(err, "derive expected mcp tool headers") + } + + valueTypes := make(map[string]string, len(bindings)) + for _, binding := range bindings { + valueTypes[strings.ToLower(ParameterHeaderPrefix+binding.HeaderName)] = binding.ValueType + } + + actual := make(http.Header) + for key, values := range requestHeaders { + if !strings.HasPrefix(strings.ToLower(key), strings.ToLower(ParameterHeaderPrefix)) { + continue + } + actual[key] = append([]string(nil), values...) + } + if len(actual) != len(expected) { + return errors.Errorf("mcp parameter header count mismatch: expected %d, got %d", len(expected), len(actual)) + } + + for key, expectedValues := range expected { + if len(expectedValues) != 1 { + return errors.Errorf("mcp parameter header %s has invalid expected cardinality", key) + } + actualValues := actual.Values(key) + if len(actualValues) != 1 { + return errors.Errorf("mcp parameter header %s must occur exactly once", key) + } + + want, err := DecodeMCPHeaderValue(expectedValues[0]) + if err != nil { + return errors.Wrapf(err, "decode expected mcp parameter header %s", key) + } + got, err := DecodeMCPHeaderValue(actualValues[0]) + if err != nil { + return errors.Wrapf(err, "decode mcp parameter header %s", key) + } + if valueTypes[strings.ToLower(key)] == "integer" { + if !equalMCPHeaderNumbers(got, want) { + return errors.Errorf("mcp parameter header %s mismatch", key) + } + continue + } + if got != want { + return errors.Errorf("mcp parameter header %s mismatch", key) + } + } + return nil +} + +// ValidateToolSchemaHeaders validates x-mcp-header annotations without requiring argument values. +// +// Parameters: +// - schema: the tool input schema to validate. +// +// Return values: +// - error: a wrapped annotation placement, token, type, or uniqueness error. +func ValidateToolSchemaHeaders(schema map[string]any) error { + if _, err := collectToolHeaderBindings(schema); err != nil { + return errors.Wrap(err, "validate mcp tool schema headers") + } + return nil +} + +// ValidateToolDescriptor validates required tool fields and HTTP header annotations. +// +// Parameters: +// - tool: the MCP tool descriptor to validate for Streamable HTTP use. +// +// Return values: +// - error: a wrapped required-field or input-schema annotation error. +func ValidateToolDescriptor(tool ToolDescriptor) error { + if strings.TrimSpace(tool.Name) == "" { + return errors.New("mcp tool name is required") + } + if tool.InputSchema == nil { + return errors.New("mcp tool inputSchema is required") + } + if err := ValidateToolSchemaHeaders(tool.InputSchema); err != nil { + return errors.Wrap(err, "validate mcp tool inputSchema") + } + return nil +} + +// FilterValidToolDescriptors separates valid tools from descriptors rejected by HTTP rules. +// +// Parameters: +// - tools: descriptors returned by one or more tools/list pages. +// +// Return values: +// - []ToolDescriptor: valid descriptors in their original order. +// - []ToolDescriptorRejection: names and reasons for excluded descriptors. +func FilterValidToolDescriptors(tools []ToolDescriptor) ([]ToolDescriptor, []ToolDescriptorRejection) { + valid := make([]ToolDescriptor, 0, len(tools)) + rejected := make([]ToolDescriptorRejection, 0) + for _, tool := range tools { + if err := ValidateToolDescriptor(tool); err != nil { + rejected = append(rejected, ToolDescriptorRejection{Name: tool.Name, Err: err}) + continue + } + valid = append(valid, tool) + } + return valid, rejected +} + +// EncodeMCPHeaderValue applies the protocol Base64 sentinel to unsafe or sentinel-like values. +// +// Parameters: +// - value: the decoded UTF-8 value to represent in one HTTP header field. +// +// Return values: +// - string: the original safe value or its exact sentinel-encoded form. +func EncodeMCPHeaderValue(value string) string { + if headerValueRequiresEncoding(value) { + return mcpBase64SentinelPrefix + base64.StdEncoding.EncodeToString([]byte(value)) + mcpBase64SentinelSuffix + } + return value +} + +// DecodeMCPHeaderValue decodes a sentinel value or validates one safe plain HTTP field value. +// +// Parameters: +// - value: the HTTP field value received from an MCP peer. +// +// Return values: +// - string: the decoded UTF-8 protocol value. +// - error: a wrapped sentinel, Base64, UTF-8, or plain-field validation error. +func DecodeMCPHeaderValue(value string) (string, error) { + if strings.HasPrefix(value, mcpBase64SentinelPrefix) || strings.HasSuffix(value, mcpBase64SentinelSuffix) { + if !strings.HasPrefix(value, mcpBase64SentinelPrefix) || !strings.HasSuffix(value, mcpBase64SentinelSuffix) { + return "", errors.New("malformed mcp base64 header sentinel") + } + encoded := strings.TrimSuffix(strings.TrimPrefix(value, mcpBase64SentinelPrefix), mcpBase64SentinelSuffix) + decoded, err := base64.StdEncoding.DecodeString(encoded) + if err != nil { + return "", errors.Wrap(err, "decode mcp base64 header value") + } + if !utf8.Valid(decoded) { + return "", errors.New("decoded mcp header value is not valid UTF-8") + } + return string(decoded), nil + } + if headerValueRequiresEncoding(value) { + return "", errors.New("unsafe plain mcp header value") + } + return value, nil +} + +// collectToolHeaderBindings returns deterministic bindings reachable through object properties only. +// +// Parameters: +// - schema: the tool input schema to inspect. +// +// Return values: +// - []toolHeaderBinding: validated header names, argument paths, and primitive types. +// - error: a wrapped placement, token, type, or uniqueness error. +func collectToolHeaderBindings(schema map[string]any) ([]toolHeaderBinding, error) { + if len(schema) == 0 { + return nil, nil + } + if err := validateToolHeaderAnnotationPlacement(schema, true, false, "$"); err != nil { + return nil, errors.Wrap(err, "validate x-mcp-header placement") + } + bindings := make([]toolHeaderBinding, 0) + seen := make(map[string]string) + if err := walkToolSchemaProperties(schema, nil, &bindings, seen); err != nil { + return nil, errors.Wrap(err, "walk mcp tool schema properties") + } + return bindings, nil +} + +// validateToolHeaderAnnotationPlacement rejects annotations outside a properties-only root path. +// +// Parameters: +// - value: the current schema fragment. +// - reachable: whether properties-only traversal can reach the fragment. +// - isProperty: whether the fragment is itself a property schema. +// - location: the diagnostic JSON path. +// +// Return values: +// - error: an annotation-placement error, or nil when every descendant is valid. +func validateToolHeaderAnnotationPlacement(value any, reachable bool, isProperty bool, location string) error { + switch typed := value.(type) { + case map[string]any: + if _, exists := typed["x-mcp-header"]; exists && (!reachable || !isProperty) { + return errors.Errorf("x-mcp-header at %s is not statically reachable through properties", location) + } + for key, child := range typed { + if key == "x-mcp-header" { + continue + } + if key == "properties" && reachable { + properties, ok := child.(map[string]any) + if !ok { + continue + } + propertyNames := make([]string, 0, len(properties)) + for name := range properties { + propertyNames = append(propertyNames, name) + } + sort.Strings(propertyNames) + for _, name := range propertyNames { + if err := validateToolHeaderAnnotationPlacement(properties[name], true, true, location+".properties."+name); err != nil { + return err + } + } + continue + } + if err := validateToolHeaderAnnotationPlacement(child, false, false, location+"."+key); err != nil { + return err + } + } + case []any: + for index, child := range typed { + if err := validateToolHeaderAnnotationPlacement(child, false, false, location+"["+strconv.Itoa(index)+"]"); err != nil { + return err + } + } + } + return nil +} + +// walkToolSchemaProperties recursively visits object properties with deterministic argument paths. +// +// Parameters: +// - schema: the current object schema. +// - path: the property path from the input root. +// - bindings: the destination slice for validated bindings. +// - seen: case-insensitive header names and their first locations. +// +// Return values: +// - error: a token, type, or duplicate-name error. +func walkToolSchemaProperties(schema map[string]any, path []string, bindings *[]toolHeaderBinding, seen map[string]string) error { + properties, ok := schema["properties"].(map[string]any) + if !ok { + return nil + } + propertyNames := make([]string, 0, len(properties)) + for name := range properties { + propertyNames = append(propertyNames, name) + } + sort.Strings(propertyNames) + + for _, propertyName := range propertyNames { + propertySchema, ok := properties[propertyName].(map[string]any) + if !ok { + continue + } + propertyPath := append(append([]string(nil), path...), propertyName) + if annotation, exists := propertySchema["x-mcp-header"]; exists { + headerName, ok := annotation.(string) + if !ok || headerName == "" { + return errors.Errorf("x-mcp-header at %s must be a non-empty string", strings.Join(propertyPath, ".")) + } + if !isValidMCPHeaderToken(headerName) { + return errors.Errorf("x-mcp-header %q at %s is not a valid header token", headerName, strings.Join(propertyPath, ".")) + } + valueType, _ := propertySchema["type"].(string) + if valueType != "string" && valueType != "integer" && valueType != "boolean" { + return errors.Errorf("x-mcp-header %q at %s requires string, integer, or boolean type", headerName, strings.Join(propertyPath, ".")) + } + canonical := strings.ToLower(headerName) + if previous, exists := seen[canonical]; exists { + return errors.Errorf("x-mcp-header %q is duplicated at %s and %s", headerName, previous, strings.Join(propertyPath, ".")) + } + seen[canonical] = strings.Join(propertyPath, ".") + *bindings = append(*bindings, toolHeaderBinding{HeaderName: headerName, Path: propertyPath, ValueType: valueType}) + } + if _, hasNestedProperties := propertySchema["properties"]; hasNestedProperties { + if err := walkToolSchemaProperties(propertySchema, propertyPath, bindings, seen); err != nil { + return err + } + } + } + return nil +} + +// lookupToolArgument resolves one deterministic property path from JSON-like arguments. +// +// Parameters: +// - arguments: the normalized tools/call argument object. +// - path: the statically reachable property path. +// +// Return values: +// - any: the resolved value when present. +// - bool: true when every path segment exists. +func lookupToolArgument(arguments map[string]any, path []string) (any, bool) { + if len(path) == 0 { + return nil, false + } + var current any = arguments + for _, segment := range path { + object, ok := current.(map[string]any) + if !ok { + return nil, false + } + current, ok = object[segment] + if !ok { + return nil, false + } + } + return current, true +} + +// formatToolHeaderValue converts one primitive JSON value into its MCP header representation. +// +// Parameters: +// - value: the argument value selected by a schema binding. +// - valueType: string, boolean, or integer from the schema. +// +// Return values: +// - string: the safe plain or sentinel-encoded HTTP value. +// - error: a type, range, or formatting error. +func formatToolHeaderValue(value any, valueType string) (string, error) { + var rendered string + switch valueType { + case "string": + text, ok := value.(string) + if !ok { + return "", errors.Errorf("expected string, got %T", value) + } + rendered = text + case "boolean": + boolean, ok := value.(bool) + if !ok { + return "", errors.Errorf("expected boolean, got %T", value) + } + rendered = strconv.FormatBool(boolean) + case "integer": + integer, err := renderInteger(value) + if err != nil { + return "", err + } + rendered = integer + default: + return "", errors.Errorf("unsupported mcp header value type %q", valueType) + } + return EncodeMCPHeaderValue(rendered), nil +} + +// renderInteger formats a JavaScript-safe integer without losing precision. +// +// Parameters: +// - value: an integer-compatible Go or JSON number. +// +// Return values: +// - string: the base-10 integer representation. +// - error: a type, fractional, non-finite, parse, or safe-range error. +func renderInteger(value any) (string, error) { + var signed int64 + var unsigned uint64 + var isUnsigned bool + + switch typed := value.(type) { + case int: + signed = int64(typed) + case int8: + signed = int64(typed) + case int16: + signed = int64(typed) + case int32: + signed = int64(typed) + case int64: + signed = typed + case uint: + unsigned, isUnsigned = uint64(typed), true + case uint8: + unsigned, isUnsigned = uint64(typed), true + case uint16: + unsigned, isUnsigned = uint64(typed), true + case uint32: + unsigned, isUnsigned = uint64(typed), true + case uint64: + unsigned, isUnsigned = typed, true + case float64: + if math.IsNaN(typed) || math.IsInf(typed, 0) || math.Trunc(typed) != typed || math.Abs(typed) > float64(maxMCPHeaderInteger) { + return "", errors.Errorf("expected JavaScript-safe integer, got %v", typed) + } + return strconv.FormatFloat(typed, 'f', -1, 64), nil + case json.Number: + rational, ok := new(big.Rat).SetString(string(typed)) + if !ok || !rational.IsInt() { + return "", errors.Errorf("expected exact JSON integer, got %q", typed) + } + integer := rational.Num() + if !integer.IsInt64() { + return "", errors.Errorf("integer %q exceeds the supported range", typed) + } + signed = integer.Int64() + default: + return "", errors.Errorf("expected integer, got %T", value) + } + + if isUnsigned { + if unsigned > uint64(maxMCPHeaderInteger) { + return "", errors.Errorf("integer %d exceeds the JavaScript-safe range", unsigned) + } + return strconv.FormatUint(unsigned, 10), nil + } + if signed < -maxMCPHeaderInteger || signed > maxMCPHeaderInteger { + return "", errors.Errorf("integer %d exceeds the JavaScript-safe range", signed) + } + return strconv.FormatInt(signed, 10), nil +} + +// equalMCPHeaderNumbers compares two numeric header representations using SEP-2243 precision semantics. +// +// Parameters: +// - left: the decoded inbound header value. +// - right: the canonical value derived from the JSON body. +// +// Return values: +// - bool: true when both values are finite and differ by at most 1E-9 relative precision. +func equalMCPHeaderNumbers(left, right string) bool { + leftNumber, leftErr := strconv.ParseFloat(left, 64) + rightNumber, rightErr := strconv.ParseFloat(right, 64) + if leftErr != nil || rightErr != nil || math.IsNaN(leftNumber) || math.IsNaN(rightNumber) || math.IsInf(leftNumber, 0) || math.IsInf(rightNumber, 0) { + return false + } + difference := math.Abs(leftNumber - rightNumber) + scale := math.Max(1, math.Max(math.Abs(leftNumber), math.Abs(rightNumber))) + return difference <= mcpHeaderNumberEpsilon*scale +} + +// headerValueRequiresEncoding reports whether a value requires the MCP Base64 sentinel. +// +// Parameters: +// - value: the decoded UTF-8 field value. +// +// Return values: +// - bool: true for whitespace-sensitive, non-ASCII/control, or sentinel-like values. +func headerValueRequiresEncoding(value string) bool { + if !utf8.ValidString(value) || strings.TrimSpace(value) != value { + return true + } + if strings.HasPrefix(value, mcpBase64SentinelPrefix) || strings.HasSuffix(value, mcpBase64SentinelSuffix) { + return true + } + for _, character := range value { + if character == '\t' || character == ' ' || (character >= 0x21 && character <= 0x7e) { + continue + } + return true + } + return false +} + +// isValidMCPHeaderToken reports whether an annotation can safely extend an HTTP header name. +// +// Parameters: +// - value: the x-mcp-header annotation value. +// +// Return values: +// - bool: true when every character satisfies HTTP token syntax. +func isValidMCPHeaderToken(value string) bool { + if value == "" { + return false + } + for _, character := range value { + if (character >= 'a' && character <= 'z') || (character >= 'A' && character <= 'Z') || (character >= '0' && character <= '9') { + continue + } + switch character { + case '!', '#', '$', '%', '&', '\'', '*', '+', '-', '.', '^', '_', '`', '|', '~': + continue + default: + return false + } + } + return true +} diff --git a/relay/mcp/headers_test.go b/relay/mcp/headers_test.go new file mode 100644 index 0000000000..55bef04f11 --- /dev/null +++ b/relay/mcp/headers_test.go @@ -0,0 +1,180 @@ +package mcp + +import ( + "encoding/json" + "net/http" + "testing" + + "github.com/stretchr/testify/require" +) + +// TestToolArgumentHeadersDerivesNestedValues verifies primitive schema annotations become mirrored headers. +func TestToolArgumentHeadersDerivesNestedValues(t *testing.T) { + schema := map[string]any{ + "type": "object", + "properties": map[string]any{ + "tenant": map[string]any{ + "type": "string", + "x-mcp-header": "Tenant-ID", + }, + "options": map[string]any{ + "type": "object", + "properties": map[string]any{ + "attempt": map[string]any{ + "type": "integer", + "x-mcp-header": "Attempt", + }, + "enabled": map[string]any{ + "type": "boolean", + "x-mcp-header": "Enabled", + }, + }, + }, + }, + } + arguments := map[string]any{ + "tenant": "acme", + "options": map[string]any{"attempt": float64(3), "enabled": true}, + } + + headers, err := ToolArgumentHeaders(schema, arguments) + require.NoError(t, err) + require.Equal(t, "acme", headers.Get("Mcp-Param-Tenant-ID")) + require.Equal(t, "3", headers.Get("Mcp-Param-Attempt")) + require.Equal(t, "true", headers.Get("Mcp-Param-Enabled")) +} + +// TestToolArgumentHeadersEncodesUnsafeValues verifies the exact Base64 sentinel prevents invalid HTTP header values. +func TestToolArgumentHeadersEncodesUnsafeValues(t *testing.T) { + schema := map[string]any{ + "properties": map[string]any{ + "value": map[string]any{"type": "string", "x-mcp-header": "Value"}, + }, + } + headers, err := ToolArgumentHeaders(schema, map[string]any{"value": "=?base64?already-prefixed?="}) + require.NoError(t, err) + require.Equal(t, "=?base64?PT9iYXNlNjQ/YWxyZWFkeS1wcmVmaXhlZD89?=", headers.Get("Mcp-Param-Value")) + + decoded, err := DecodeMCPHeaderValue("=?base64?SGVsbG8sIOS4lueVjA==?=") + require.NoError(t, err) + require.Equal(t, "Hello, 世界", decoded) +} + +// TestToolArgumentHeadersOmitsNullValues verifies null and absent parameters do not produce headers. +func TestToolArgumentHeadersOmitsNullValues(t *testing.T) { + schema := map[string]any{ + "properties": map[string]any{ + "value": map[string]any{"type": "string", "x-mcp-header": "Value"}, + }, + } + headers, err := ToolArgumentHeaders(schema, map[string]any{"value": nil}) + require.NoError(t, err) + require.Empty(t, headers) +} + +// TestToolArgumentHeadersRejectsUnsafeIntegerRange verifies integer headers remain exactly representable in JavaScript. +func TestToolArgumentHeadersRejectsUnsafeIntegerRange(t *testing.T) { + schema := map[string]any{ + "properties": map[string]any{ + "value": map[string]any{"type": "integer", "x-mcp-header": "Value"}, + }, + } + _, err := ToolArgumentHeaders(schema, map[string]any{"value": float64(1 << 53)}) + require.ErrorContains(t, err, "JavaScript-safe") +} + +// TestValidateToolSchemaHeadersRejectsDuplicateNames verifies header annotations are unique case-insensitively. +func TestValidateToolSchemaHeadersRejectsDuplicateNames(t *testing.T) { + schema := map[string]any{ + "properties": map[string]any{ + "first": map[string]any{"type": "string", "x-mcp-header": "Tenant"}, + "second": map[string]any{"type": "string", "x-mcp-header": "tenant"}, + }, + } + require.ErrorContains(t, ValidateToolSchemaHeaders(schema), "duplicated") +} + +// TestValidateToolSchemaHeadersRejectsUnreachableAnnotations verifies annotations cannot hide behind arrays or composition keywords. +func TestValidateToolSchemaHeadersRejectsUnreachableAnnotations(t *testing.T) { + schema := map[string]any{ + "properties": map[string]any{ + "items": map[string]any{ + "type": "array", + "items": map[string]any{ + "type": "string", + "x-mcp-header": "Hidden", + }, + }, + }, + } + require.ErrorContains(t, ValidateToolSchemaHeaders(schema), "not statically reachable") +} + +// TestValidateToolArgumentHeadersRejectsMismatch verifies server-side mirrored header checks reject tampering. +func TestValidateToolArgumentHeadersRejectsMismatch(t *testing.T) { + schema := map[string]any{ + "properties": map[string]any{ + "tenant": map[string]any{"type": "string", "x-mcp-header": "Tenant"}, + }, + } + headers := make(http.Header) + headers.Set("Mcp-Param-Tenant", "other") + require.ErrorContains(t, ValidateToolArgumentHeaders(headers, schema, map[string]any{"tenant": "acme"}), "mismatch") +} + +// TestFilterValidToolDescriptorsExcludesInvalidTools verifies one malformed tool does not hide valid tools. +func TestFilterValidToolDescriptorsExcludesInvalidTools(t *testing.T) { + tools := []ToolDescriptor{ + {Name: "valid", InputSchema: map[string]any{"type": "object"}}, + {Name: "invalid", InputSchema: map[string]any{"items": map[string]any{"x-mcp-header": "Hidden"}}}, + } + valid, rejected := FilterValidToolDescriptors(tools) + require.Equal(t, []ToolDescriptor{tools[0]}, valid) + require.Len(t, rejected, 1) + require.Equal(t, "invalid", rejected[0].Name) +} + +// TestToolArgumentHeadersEncodeEitherSentinelBoundary verifies a single reserved boundary is never emitted as plain text. +// +// Parameters: +// - t: The test owns reserved-sentinel encoding assertions. +// +// Return values: none; failures are reported through t. +func TestToolArgumentHeadersEncodeEitherSentinelBoundary(t *testing.T) { + schema := map[string]any{"properties": map[string]any{"value": map[string]any{"type": "string", "x-mcp-header": "Value"}}} + for _, value := range []string{"=?base64?literal", "literal?="} { + headers, err := ToolArgumentHeaders(schema, map[string]any{"value": value}) + require.NoError(t, err) + require.NotEqual(t, value, headers.Get("Mcp-Param-Value")) + decoded, err := DecodeMCPHeaderValue(headers.Get("Mcp-Param-Value")) + require.NoError(t, err) + require.Equal(t, value, decoded) + } +} + +// TestValidateToolArgumentHeadersComparesIntegersNumerically verifies equivalent decimal representations are accepted. +// +// Parameters: +// - t: The test owns numeric header validation assertions. +// +// Return values: none; failures are reported through t. +func TestValidateToolArgumentHeadersComparesIntegersNumerically(t *testing.T) { + schema := map[string]any{"properties": map[string]any{"value": map[string]any{"type": "integer", "x-mcp-header": "Value"}}} + headers := make(http.Header) + headers.Set("Mcp-Param-Value", "42.0") + require.NoError(t, ValidateToolArgumentHeaders(headers, schema, map[string]any{"value": float64(42)})) +} + +// TestRenderIntegerRejectsRoundedJSONNumbers verifies fractional JSON numbers cannot round into accepted integers. +// +// Parameters: +// - t: The test owns exact-number parsing assertions. +// +// Return values: none; failures are reported through t. +func TestRenderIntegerRejectsRoundedJSONNumbers(t *testing.T) { + _, err := renderInteger(json.Number("1.0000000000000001")) + require.ErrorContains(t, err, "exact JSON integer") + value, err := renderInteger(json.Number("1e3")) + require.NoError(t, err) + require.Equal(t, "1000", value) +} diff --git a/relay/mcp/protocol.go b/relay/mcp/protocol.go new file mode 100644 index 0000000000..0d1553bef8 --- /dev/null +++ b/relay/mcp/protocol.go @@ -0,0 +1,269 @@ +package mcp + +import ( + stderrors "errors" + "fmt" + "net/http" + "strings" +) + +const ( + // ProtocolVersion is the latest stable MCP protocol version supported by one-api. + ProtocolVersion = "2026-07-28" + // LegacyProtocolVersion is the preferred initialization-based MCP protocol version. + LegacyProtocolVersion = "2025-11-25" + // LegacyProtocolVersionFallback keeps compatibility with older Streamable HTTP servers. + LegacyProtocolVersionFallback = "2025-06-18" + + // ProtocolVersionHeader carries the protocol version for every modern request. + ProtocolVersionHeader = "Mcp-Protocol-Version" + // MethodHeader mirrors the JSON-RPC method for modern request validation. + MethodHeader = "Mcp-Method" + // NameHeader mirrors the named resource or tool for modern request validation. + NameHeader = "Mcp-Name" + // ParameterHeaderPrefix prefixes schema-driven argument headers. + ParameterHeaderPrefix = "Mcp-Param-" + // SessionIDHeader carries legacy Streamable HTTP session identifiers. + SessionIDHeader = "Mcp-Session-Id" + + // MetaProtocolVersionKey identifies the request protocol version in _meta. + MetaProtocolVersionKey = "io.modelcontextprotocol/protocolVersion" + // MetaClientInfoKey identifies the client implementation in _meta. + MetaClientInfoKey = "io.modelcontextprotocol/clientInfo" + // MetaClientCapabilitiesKey identifies per-request client capabilities in _meta. + MetaClientCapabilitiesKey = "io.modelcontextprotocol/clientCapabilities" + // MetaServerInfoKey identifies the server implementation in result _meta. + MetaServerInfoKey = "io.modelcontextprotocol/serverInfo" + + // ResultTypeComplete marks a final successful result. + ResultTypeComplete = "complete" + // ResultTypeInputRequired marks a result that requires additional client input. + ResultTypeInputRequired = "input_required" + + // CacheScopePublic allows a cacheable result to be reused between users. + CacheScopePublic = "public" + // CacheScopePrivate restricts a cacheable result to the authenticated user. + CacheScopePrivate = "private" + + // ErrorCodeHeaderMismatch reports disagreement between JSON fields and mirrored headers. + ErrorCodeHeaderMismatch = -32020 + // ErrorCodeMissingRequiredClientCapability reports an undeclared required capability. + ErrorCodeMissingRequiredClientCapability = -32021 + // ErrorCodeUnsupportedProtocolVersion reports an unsupported modern protocol version. + ErrorCodeUnsupportedProtocolVersion = -32022 +) + +// ImplementationInfo identifies an MCP client or server implementation. +type ImplementationInfo struct { + Name string `json:"name"` + Version string `json:"version"` +} + +// RequestMeta contains the namespaced per-request metadata required by MCP 2026-07-28. +type RequestMeta struct { + ProtocolVersion string `json:"io.modelcontextprotocol/protocolVersion"` + ClientInfo ImplementationInfo `json:"io.modelcontextprotocol/clientInfo"` + ClientCapabilities map[string]any `json:"io.modelcontextprotocol/clientCapabilities"` +} + +// ResponseMeta contains the namespaced server identity returned with modern MCP results. +type ResponseMeta struct { + ServerInfo ImplementationInfo `json:"io.modelcontextprotocol/serverInfo"` +} + +// DiscoveryResult describes protocol versions and capabilities exposed by an MCP server. +type DiscoveryResult struct { + ResultType string `json:"resultType"` + SupportedVersions []string `json:"supportedVersions"` + Capabilities map[string]any `json:"capabilities"` + Instructions string `json:"instructions,omitempty"` + TTLMS int64 `json:"ttlMs"` + CacheScope string `json:"cacheScope"` + Meta ResponseMeta `json:"_meta"` +} + +// ProtocolError preserves HTTP and JSON-RPC error details for negotiation decisions. +type ProtocolError struct { + HTTPStatus int + Code int + Message string + Data any + Body string +} + +// Error renders ProtocolError without discarding transport or JSON-RPC context. +// +// Parameters: none. +// +// Return values: +// - string: a stable diagnostic string containing the available status, code, and message. +func (e *ProtocolError) Error() string { + if e == nil { + return "mcp protocol error" + } + parts := make([]string, 0, 3) + if e.HTTPStatus != 0 { + parts = append(parts, fmt.Sprintf("status %d", e.HTTPStatus)) + } + if e.Code != 0 { + parts = append(parts, fmt.Sprintf("code %d", e.Code)) + } + message := strings.TrimSpace(e.Message) + if message == "" { + message = strings.TrimSpace(e.Body) + } + if message != "" { + parts = append(parts, message) + } + if len(parts) == 0 { + return "mcp protocol error" + } + return "mcp protocol error: " + strings.Join(parts, ": ") +} + +// SupportedProtocolVersions returns protocol versions that the gateway can serve. +// +// Parameters: none. +// +// Return values: +// - []string: a new slice ordered from the preferred modern version to legacy compatibility versions. +func SupportedProtocolVersions() []string { + return []string{ProtocolVersion, LegacyProtocolVersion, LegacyProtocolVersionFallback} +} + +// IsLegacyProtocolVersion reports whether version selects the initialization-based protocol era. +// +// Parameters: +// - version: the protocol version supplied by an MCP peer or HTTP header. +// +// Return values: +// - bool: true only for legacy versions that one-api explicitly supports. +func IsLegacyProtocolVersion(version string) bool { + switch strings.TrimSpace(version) { + case LegacyProtocolVersion, LegacyProtocolVersionFallback: + return true + default: + return false + } +} + +// IsSupportedProtocolVersion reports whether version is supported by either MCP protocol era. +// +// Parameters: +// - version: the protocol version supplied by an MCP peer. +// +// Return values: +// - bool: true when one-api can process the version through a modern or legacy path. +func IsSupportedProtocolVersion(version string) bool { + return strings.TrimSpace(version) == ProtocolVersion || IsLegacyProtocolVersion(version) +} + +// NegotiateLegacyProtocolVersion chooses the legacy protocol version returned by initialize. +// +// Parameters: +// - requested: the version requested in legacy initialize parameters. +// +// Return values: +// - string: the requested version when supported, otherwise the preferred legacy version. +func NegotiateLegacyProtocolVersion(requested string) string { + requested = strings.TrimSpace(requested) + if IsLegacyProtocolVersion(requested) { + return requested + } + return LegacyProtocolVersion +} + +// ModernRequestMeta returns the metadata attached to each MCP 2026-07-28 request. +// +// Parameters: none. +// +// Return values: +// - RequestMeta: fresh request metadata with an explicit empty optional-capability set. +func ModernRequestMeta() RequestMeta { + return RequestMeta{ + ProtocolVersion: ProtocolVersion, + ClientInfo: ImplementationInfo{ + Name: mcpClientName, + Version: mcpClientVersion, + }, + ClientCapabilities: map[string]any{}, + } +} + +// ServerResponseMeta returns one-api's server identity for modern results. +// +// Parameters: +// - name: the public server implementation name. +// - version: the public server implementation version. +// +// Return values: +// - ResponseMeta: metadata containing the supplied server identity. +func ServerResponseMeta(name, version string) ResponseMeta { + return ResponseMeta{ServerInfo: ImplementationInfo{Name: name, Version: version}} +} + +// WithModernMeta returns object parameters containing the required modern request metadata. +// +// Parameters: +// - params: method-specific parameters; nil is treated as an empty object. +// +// Return values: +// - map[string]any: a new parameters object that does not mutate the caller's map. +func WithModernMeta(params map[string]any) map[string]any { + out := make(map[string]any) + for key, value := range params { + out[key] = value + } + out["_meta"] = ModernRequestMeta() + return out +} + +// IsRecognizedModernError reports whether an error proves that the peer speaks modern MCP. +// +// Parameters: +// - err: the transport or JSON-RPC error returned by a modern request. +// +// Return values: +// - bool: true when retrying through the legacy handshake would be incorrect. +func IsRecognizedModernError(err error) bool { + var protocolErr *ProtocolError + if !stderrors.As(err, &protocolErr) || protocolErr == nil { + return false + } + switch protocolErr.Code { + case ErrorCodeHeaderMismatch, ErrorCodeMissingRequiredClientCapability, ErrorCodeUnsupportedProtocolVersion: + return true + case -32601: + return protocolErr.HTTPStatus == http.StatusNotFound + default: + return false + } +} + +// IsModernFallbackCandidate reports whether a failed modern request should retry through the legacy handshake. +// +// Parameters: +// - err: the error returned by a modern request attempt. +// +// Return values: +// - bool: true only for failures that plausibly indicate a legacy endpoint. +func IsModernFallbackCandidate(err error) bool { + var protocolErr *ProtocolError + if !stderrors.As(err, &protocolErr) || protocolErr == nil { + return false + } + if protocolErr.HTTPStatus == http.StatusUnauthorized || protocolErr.HTTPStatus == http.StatusForbidden { + return false + } + if IsRecognizedModernError(err) { + return false + } + if protocolErr.HTTPStatus == http.StatusBadRequest || protocolErr.HTTPStatus == http.StatusNotFound || protocolErr.HTTPStatus == http.StatusMethodNotAllowed { + return true + } + message := strings.ToLower(protocolErr.Message + " " + protocolErr.Body) + if protocolErr.Code == -32002 || protocolErr.Code == -32600 || protocolErr.Code == -32601 { + return strings.Contains(message, "initial") || strings.Contains(message, "session") || strings.Contains(message, "protocol") + } + return false +} diff --git a/relay/mcp/protocol_test.go b/relay/mcp/protocol_test.go new file mode 100644 index 0000000000..d6dc84ce4e --- /dev/null +++ b/relay/mcp/protocol_test.go @@ -0,0 +1,26 @@ +package mcp + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +// TestWithModernMetaCopiesParametersAndSupportsNil verifies metadata injection is allocation-safe and non-mutating. +// +// Parameters: +// - t: The test owns map-copy and nil-input assertions. +// +// Return values: none; failures are reported through t. +func TestWithModernMetaCopiesParametersAndSupportsNil(t *testing.T) { + input := map[string]any{"cursor": "page-2"} + output := WithModernMeta(input) + require.Equal(t, "page-2", output["cursor"]) + require.NotNil(t, output["_meta"]) + output["cursor"] = "changed" + require.Equal(t, "page-2", input["cursor"]) + + nilOutput := WithModernMeta(nil) + require.Len(t, nilOutput, 1) + require.NotNil(t, nilOutput["_meta"]) +} diff --git a/relay/mcp/reviewer_regression_test.go b/relay/mcp/reviewer_regression_test.go new file mode 100644 index 0000000000..8ada610469 --- /dev/null +++ b/relay/mcp/reviewer_regression_test.go @@ -0,0 +1,210 @@ +package mcp + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "net/url" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/require" + + "github.com/Laisky/one-api/model" +) + +// TestCallToolLatestNormalizesNilArguments verifies zero-argument calls transmit an empty object instead of null. +// +// Parameters: +// - t: The test owns the mock MCP server and assertions. +// +// Return values: none; failures are reported through t. +func TestCallToolLatestNormalizesNilArguments(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + var envelope struct { + ID any `json:"id"` + Params struct { + Arguments map[string]any `json:"arguments"` + } `json:"params"` + } + require.NoError(t, json.NewDecoder(request.Body).Decode(&envelope)) + require.NotNil(t, envelope.Params.Arguments) + require.Empty(t, envelope.Params.Arguments) + writer.Header().Set("Content-Type", "application/json") + require.NoError(t, json.NewEncoder(writer).Encode(map[string]any{ + "jsonrpc": "2.0", + "id": envelope.ID, + "result": map[string]any{"resultType": ResultTypeComplete, "content": []any{}}, + })) + })) + defer server.Close() + + client := NewStreamableHTTPClient(&model.MCPServer{BaseURL: server.URL}, nil, 5*time.Second) + result, err := client.CallToolLatestWithDescriptor(context.Background(), ToolDescriptor{ + Name: "health", + InputSchema: map[string]any{"type": "object", "additionalProperties": false}, + }, nil) + require.NoError(t, err) + require.Equal(t, ResultTypeComplete, result.ResultType) +} + +// TestModernMCPClientRejectsUncorrelatedSSE verifies an unrelated SSE response cannot satisfy the active request. +// +// Parameters: +// - t: The test owns the mock MCP server and correlation assertion. +// +// Return values: none; failures are reported through t. +func TestModernMCPClientRejectsUncorrelatedSSE(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + writer.Header().Set("Content-Type", "text/event-stream") + _, err := fmt.Fprint(writer, "data: {\"jsonrpc\":\"2.0\",\"id\":\"other-request\",\"result\":{\"resultType\":\"complete\"}}\n\n") + require.NoError(t, err) + })) + defer server.Close() + + client := NewStreamableHTTPClient(&model.MCPServer{BaseURL: server.URL}, nil, 5*time.Second) + _, err := client.DiscoverLatest(context.Background()) + require.ErrorContains(t, err, "no event for request id") +} + +// TestModernMCPClientHeaderSnapshotIsRaceSafe verifies internal header updates cannot race request construction. +// +// Parameters: +// - t: The test owns concurrent requests, header mutations, and assertions. +// +// Return values: none; failures are reported through t. +func TestModernMCPClientHeaderSnapshotIsRaceSafe(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + var envelope struct { + ID any `json:"id"` + } + require.NoError(t, json.NewDecoder(request.Body).Decode(&envelope)) + writer.Header().Set("Content-Type", "application/json") + require.NoError(t, json.NewEncoder(writer).Encode(map[string]any{ + "jsonrpc": "2.0", + "id": envelope.ID, + "result": map[string]any{ + "resultType": ResultTypeComplete, + "supportedVersions": []string{ProtocolVersion}, + "capabilities": map[string]any{}, + }, + })) + })) + defer server.Close() + + client := NewStreamableHTTPClient(&model.MCPServer{BaseURL: server.URL}, nil, 5*time.Second) + const operations = 64 + start := make(chan struct{}) + errorsChannel := make(chan error, operations) + var group sync.WaitGroup + group.Add(operations + 1) + go func() { + defer group.Done() + <-start + for index := 0; index < operations*4; index++ { + client.setClientHeader("X-Regression-Header", fmt.Sprintf("value-%d", index)) + } + }() + for index := 0; index < operations; index++ { + go func() { + defer group.Done() + <-start + _, err := client.DiscoverLatest(context.Background()) + errorsChannel <- err + }() + } + close(start) + group.Wait() + close(errorsChannel) + for err := range errorsChannel { + require.NoError(t, err) + } +} + +// TestMCPClientRedirectPolicyPreventsCredentialLeakage verifies credentials never cross an origin or HTTPS downgrade. +// +// Parameters: +// - t: The test owns redirect targets and transport policy assertions. +// +// Return values: none; failures are reported through t. +func TestMCPClientRedirectPolicyPreventsCredentialLeakage(t *testing.T) { + t.Run("cross-origin API key redirect is blocked before target delivery", func(t *testing.T) { + var targetHits atomic.Int32 + target := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + targetHits.Add(1) + writer.WriteHeader(http.StatusNoContent) + })) + defer target.Close() + + source := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + http.Redirect(writer, request, target.URL, http.StatusTemporaryRedirect) + })) + defer source.Close() + + client := NewStreamableHTTPClient( + &model.MCPServer{BaseURL: source.URL}, + map[string]string{"X-API-Key": "secret"}, + 5*time.Second, + ) + _, err := client.DiscoverLatest(context.Background()) + require.ErrorContains(t, err, "preserve the endpoint origin") + require.Zero(t, targetHits.Load()) + }) + + t.Run("HTTPS downgrade is blocked", func(t *testing.T) { + client := NewStreamableHTTPClient(&model.MCPServer{BaseURL: "https://example.com/mcp"}, nil, time.Second) + initial := &http.Request{URL: mustMCPReviewerURL(t, "https://example.com/mcp")} + next := &http.Request{URL: mustMCPReviewerURL(t, "http://example.com/mcp")} + require.ErrorContains(t, client.httpClient().CheckRedirect(next, []*http.Request{initial}), "downgrade") + }) +} + +// TestMCPClientKeepsExplicitCredentialedHTTPCompatibility verifies direct operator-configured HTTP endpoints remain supported. +// +// Parameters: +// - t: The test owns the explicit HTTP endpoint and credential assertion. +// +// Return values: none; failures are reported through t. +func TestMCPClientKeepsExplicitCredentialedHTTPCompatibility(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + require.Equal(t, "secret", request.Header.Get("X-API-Key")) + var envelope struct { + ID any `json:"id"` + } + require.NoError(t, json.NewDecoder(request.Body).Decode(&envelope)) + writer.Header().Set("Content-Type", "application/json") + require.NoError(t, json.NewEncoder(writer).Encode(map[string]any{ + "jsonrpc": "2.0", + "id": envelope.ID, + "result": map[string]any{ + "resultType": ResultTypeComplete, + "supportedVersions": []string{ProtocolVersion}, + "capabilities": map[string]any{}, + }, + })) + })) + defer server.Close() + + client := NewStreamableHTTPClient(&model.MCPServer{BaseURL: server.URL}, map[string]string{"X-API-Key": "secret"}, 5*time.Second) + _, err := client.DiscoverLatest(context.Background()) + require.NoError(t, err) +} + +// mustMCPReviewerURL parses a fixed URL fixture and reports impossible errors through the test. +// +// Parameters: +// - t: The test receives URL parsing failures. +// - raw: The fixed absolute URL is parsed for redirect construction. +// +// Return values: +// - *url.URL: The parsed URL is returned. +func mustMCPReviewerURL(t *testing.T, raw string) *url.URL { + t.Helper() + parsed, err := url.Parse(raw) + require.NoError(t, err) + return parsed +} diff --git a/relay/mcp/sync.go b/relay/mcp/sync.go index 0f6eabac9a..71d41e3c11 100644 --- a/relay/mcp/sync.go +++ b/relay/mcp/sync.go @@ -15,14 +15,22 @@ import ( const defaultSyncTimeout = 20 * time.Second -// SyncServerTools fetches tools from the MCP server and stores them locally. +// SyncServerTools fetches a complete upstream catalog and atomically stores lossless descriptors. +// +// Parameters: +// - ctx: the request context controlling cancellation and deadlines. +// - server: the configured upstream MCP server to synchronize. +// +// Return values: +// - int: the number of valid tools stored in the replacement catalog. +// - error: a wrapped client, encoding, or database error. func SyncServerTools(ctx context.Context, server *model.MCPServer) (int, error) { if server == nil { return 0, errors.New("mcp server is nil") } client := NewStreamableHTTPClient(server, nil, defaultSyncTimeout) - tools, err := client.ListTools(ctx) + tools, err := client.ListToolsLatest(ctx) if err != nil { return 0, errors.Wrap(err, "list mcp tools from server") } @@ -36,18 +44,25 @@ func SyncServerTools(ctx context.Context, server *model.MCPServer) (int, error) if tool.InputSchema != nil { schemaBytes, err := json.Marshal(tool.InputSchema) if err != nil { - return 0, errors.Wrap(err, "marshal mcp tool schema") - } - if string(schemaBytes) != "null" { - inputSchema = string(schemaBytes) + return 0, errors.Wrapf(err, "marshal input schema for mcp tool %q", tool.Name) } + inputSchema = string(schemaBytes) + } + descriptorBytes, err := json.Marshal(tool) + if err != nil { + return 0, errors.Wrapf(err, "marshal complete descriptor for mcp tool %q", tool.Name) + } + displayName := tool.Title + if displayName == "" { + displayName = tool.Name } stored = append(stored, &model.MCPTool{ - Name: tool.Name, - DisplayName: tool.Name, - Description: tool.Description, - InputSchema: inputSchema, - Status: 1, + Name: tool.Name, + DisplayName: displayName, + Description: tool.Description, + InputSchema: inputSchema, + DescriptorJSON: string(descriptorBytes), + Status: 1, }) } @@ -56,11 +71,15 @@ func SyncServerTools(ctx context.Context, server *model.MCPServer) (int, error) errors.Wrapf(err, "upsert mcp tools for server %d", server.Id), server.Ref()) } - return len(stored), nil } -// StartAutoSync triggers MCP server tool syncs on a periodic schedule. +// StartAutoSync starts the periodic MCP catalog synchronization loop for enabled servers. +// +// Parameters: +// - ctx: the process context controlling worker shutdown. +// +// Return values: none; the worker logs each background result with the server identity. func StartAutoSync(ctx context.Context) { log := logger.FromContext(ctx) if log == nil { @@ -95,8 +114,6 @@ func StartAutoSync(ctx context.Context) { syncCtx, cancel := context.WithTimeout(ctx, defaultSyncTimeout) count, err := SyncServerTools(syncCtx, server) cancel() - // This is a background job: the logger is not request-bound, - // so every line must carry the MCP server identity explicitly. serverRef := server.Ref() if err != nil { server.MarkSyncResult(false, err.Error()) diff --git a/relay/mcp/transport.go b/relay/mcp/transport.go new file mode 100644 index 0000000000..c7d0b0a5ec --- /dev/null +++ b/relay/mcp/transport.go @@ -0,0 +1,134 @@ +package mcp + +import ( + "bytes" + "encoding/json" + "io" + "strings" + + "github.com/Laisky/errors/v2" +) + +const maxMCPResponseBodyBytes int64 = 32 << 20 + +type mcpJSONRPCError struct { + Code int `json:"code"` + Message string `json:"message"` + Data any `json:"data,omitempty"` +} + +type mcpJSONRPCEnvelope struct { + JSONRPC string `json:"jsonrpc"` + ID json.RawMessage `json:"id"` + Result json.RawMessage `json:"result"` + Error *mcpJSONRPCError `json:"error,omitempty"` +} + +// readMCPResponseBody reads an MCP response through a fixed allocation boundary. +// +// Parameters: +// - reader: the remote HTTP response body. +// +// Return values: +// - []byte: the response body when it does not exceed maxMCPResponseBodyBytes. +// - error: a wrapped read or size error when the response cannot be consumed safely. +func readMCPResponseBody(reader io.Reader) ([]byte, error) { + if reader == nil { + return nil, errors.New("mcp response body is nil") + } + body, err := io.ReadAll(io.LimitReader(reader, maxMCPResponseBodyBytes+1)) + if err != nil { + return nil, errors.Wrap(err, "read bounded mcp response body") + } + if int64(len(body)) > maxMCPResponseBodyBytes { + return nil, errors.Errorf("mcp response body exceeds %d bytes", maxMCPResponseBodyBytes) + } + return body, nil +} + +// parseMCPResponseEnvelope validates a JSON-RPC response and correlates it with one request. +// +// Parameters: +// - body: the JSON response envelope. +// - expectedID: the string request identifier generated by the client. +// +// Return values: +// - *mcpJSONRPCEnvelope: the validated response envelope. +// - error: a wrapped JSON, version, or response-correlation error. +func parseMCPResponseEnvelope(body []byte, expectedID string) (*mcpJSONRPCEnvelope, error) { + var envelope mcpJSONRPCEnvelope + if err := json.Unmarshal(body, &envelope); err != nil { + return nil, errors.Wrap(err, "decode mcp JSON-RPC response") + } + if envelope.JSONRPC != "2.0" { + return nil, errors.Errorf("mcp response jsonrpc must be 2.0, got %q", envelope.JSONRPC) + } + if err := validateMCPResponseID(envelope.ID, expectedID); err != nil { + return nil, err + } + return &envelope, nil +} + +// extractMCPResponseEnvelope finds the SSE data event correlated with one request identifier. +// +// Parameters: +// - body: a finite Server-Sent Events response body. +// - expectedID: the string request identifier generated by the client. +// +// Return values: +// - []byte: the matching JSON-RPC response envelope. +// - error: a wrapped error when no valid event matches expectedID. +func extractMCPResponseEnvelope(body []byte, expectedID string) ([]byte, error) { + normalized := strings.ReplaceAll(string(body), "\r\n", "\n") + normalized = strings.ReplaceAll(normalized, "\r", "\n") + foundData := false + for _, block := range strings.Split(normalized, "\n\n") { + dataLines := make([]string, 0) + for _, line := range strings.Split(block, "\n") { + if !strings.HasPrefix(line, "data:") { + continue + } + foundData = true + dataLines = append(dataLines, strings.TrimPrefix(strings.TrimPrefix(line, "data:"), " ")) + } + if len(dataLines) == 0 { + continue + } + candidate := []byte(strings.Join(dataLines, "\n")) + var envelope struct { + ID json.RawMessage `json:"id"` + } + if err := json.Unmarshal(candidate, &envelope); err != nil { + continue + } + if validateMCPResponseID(envelope.ID, expectedID) == nil { + return candidate, nil + } + } + if !foundData { + return nil, errors.New("mcp SSE response has no data fields") + } + return nil, errors.Errorf("mcp SSE response has no event for request id %q", expectedID) +} + +// validateMCPResponseID verifies that a JSON-RPC response ID is the expected string identifier. +// +// Parameters: +// - rawID: the encoded JSON-RPC response ID. +// - expectedID: the string request identifier generated by the client. +// +// Return values: +// - error: a correlation error when the ID is missing, has another type, or does not match. +func validateMCPResponseID(rawID json.RawMessage, expectedID string) error { + if len(bytes.TrimSpace(rawID)) == 0 { + return errors.New("mcp response id is missing") + } + var responseID string + if err := json.Unmarshal(rawID, &responseID); err != nil { + return errors.Wrap(err, "decode mcp response id") + } + if responseID != expectedID { + return errors.Errorf("mcp response id mismatch: expected %q, got %q", expectedID, responseID) + } + return nil +} diff --git a/relay/mcp/transport_security.go b/relay/mcp/transport_security.go new file mode 100644 index 0000000000..242d279897 --- /dev/null +++ b/relay/mcp/transport_security.go @@ -0,0 +1,184 @@ +package mcp + +import ( + "net" + "net/http" + "net/url" + "strings" + + "github.com/Laisky/errors/v2" +) + +const maxMCPRedirects = 10 + +// mcpSecurityRoundTripper rejects credential-bearing requests that would leave the process over remote plaintext HTTP. +type mcpSecurityRoundTripper struct { + base http.RoundTripper + credentialed bool +} + +// RoundTrip validates the outbound request transport before delegating to the configured HTTP transport. +// +// Parameters: +// - request: The outbound MCP HTTP request is checked before network I/O begins. +// +// Return values: +// - *http.Response: The delegated transport response is returned for an allowed request. +// - error: A transport-policy or delegated round-trip error is returned on failure. +func (transport mcpSecurityRoundTripper) RoundTrip(request *http.Request) (*http.Response, error) { + if err := validateMCPOutboundTransport(request, transport.credentialed); err != nil { + return nil, err + } + base := transport.base + if base == nil { + base = http.DefaultTransport + } + return base.RoundTrip(request) +} + +// httpClient returns an HTTP client that enforces MCP request and redirect transport boundaries. +// +// Parameters: none. +// +// Return values: +// - *http.Client: The client uses the configured timeout and rejects unsafe plaintext or redirected credential delivery. +func (c *StreamableHTTPClient) httpClient() *http.Client { + sensitive := hasSensitiveMCPClientState(c) + client := &http.Client{ + Timeout: c.Timeout, + Transport: mcpSecurityRoundTripper{ + base: http.DefaultTransport, + credentialed: sensitive, + }, + } + client.CheckRedirect = func(request *http.Request, via []*http.Request) error { + if len(via) >= maxMCPRedirects { + return errors.WithStack(errors.New("stopped after 10 MCP redirects")) + } + if request == nil || request.URL == nil || len(via) == 0 || via[0] == nil || via[0].URL == nil { + return nil + } + initial := via[0].URL + if strings.EqualFold(initial.Scheme, "https") && !strings.EqualFold(request.URL.Scheme, "https") { + return errors.WithStack(errors.New("MCP redirect would downgrade HTTPS to plaintext HTTP")) + } + if sensitive && !sameMCPOrigin(initial, request.URL) { + return errors.WithStack(errors.New("credentialed MCP redirect must preserve the endpoint origin")) + } + return nil + } + return client +} + +// validateMCPOutboundTransport enforces HTTPS for credentialed remote MCP requests while retaining loopback development compatibility. +// +// Parameters: +// - request: The outbound HTTP request supplies the destination URL. +// - credentialed: The flag indicates whether the client carries API keys, authorization values, cookies, URL user information, or similar secrets. +// +// Return values: +// - error: A policy error is returned when a credentialed request targets remote plaintext HTTP; otherwise nil is returned. +func validateMCPOutboundTransport(request *http.Request, credentialed bool) error { + if !credentialed { + return nil + } + if request == nil || request.URL == nil { + return errors.WithStack(errors.New("credentialed MCP request URL is missing")) + } + if strings.EqualFold(request.URL.Scheme, "https") { + return nil + } + if strings.EqualFold(request.URL.Scheme, "http") && isLoopbackMCPHostname(request.URL.Hostname()) { + return nil + } + return errors.WithStack(errors.New("credentialed MCP requests require HTTPS unless the endpoint is a loopback host")) +} + +// hasSensitiveMCPClientState reports whether an MCP client sends credentials with requests. +// +// Parameters: +// - client: The client supplies configured headers and URL user information. +// +// Return values: +// - bool: True is returned when redirects or plaintext transport could expose credentials. +func hasSensitiveMCPClientState(client *StreamableHTTPClient) bool { + if client == nil { + return false + } + for key, value := range client.headerSnapshot() { + if strings.TrimSpace(value) == "" { + continue + } + normalizedKey := http.CanonicalHeaderKey(strings.TrimSpace(key)) + if isSensitiveKey(strings.ToLower(normalizedKey)) { + return true + } + switch normalizedKey { + case "Accept", "Accept-Encoding", "Content-Type", "User-Agent", http.CanonicalHeaderKey(ProtocolVersionHeader), http.CanonicalHeaderKey(SessionIDHeader): + continue + default: + // Arbitrary configured headers can implement custom authentication even + // when their names do not contain a conventional credential token. + return true + } + } + parsed, err := url.Parse(client.BaseURL) + return err == nil && parsed.User != nil +} + +// isLoopbackMCPHostname reports whether a hostname is restricted to the local machine. +// +// Parameters: +// - hostname: The URL hostname is checked as localhost or a loopback IP address. +// +// Return values: +// - bool: True is returned only for localhost and IP loopback addresses. +func isLoopbackMCPHostname(hostname string) bool { + hostname = strings.TrimSpace(hostname) + if strings.EqualFold(hostname, "localhost") { + return true + } + address := net.ParseIP(hostname) + return address != nil && address.IsLoopback() +} + +// sameMCPOrigin reports whether two endpoint URLs share scheme, hostname, and effective port. +// +// Parameters: +// - left: The original MCP endpoint supplies the expected origin. +// - right: The redirect destination is compared with the original origin. +// +// Return values: +// - bool: True is returned only when the complete origins match. +func sameMCPOrigin(left, right *url.URL) bool { + if left == nil || right == nil { + return false + } + return strings.EqualFold(left.Scheme, right.Scheme) && + strings.EqualFold(left.Hostname(), right.Hostname()) && + effectiveMCPPort(left) == effectiveMCPPort(right) +} + +// effectiveMCPPort returns the explicit port or the scheme's standard port. +// +// Parameters: +// - endpoint: The URL supplies an optional explicit port and scheme. +// +// Return values: +// - string: The effective port is returned, or an empty string for an unknown scheme. +func effectiveMCPPort(endpoint *url.URL) string { + if endpoint == nil { + return "" + } + if port := endpoint.Port(); port != "" { + return port + } + switch strings.ToLower(endpoint.Scheme) { + case "http": + return "80" + case "https": + return "443" + default: + return "" + } +} diff --git a/relay/mcp/transport_security_test.go b/relay/mcp/transport_security_test.go new file mode 100644 index 0000000000..194b9d093a --- /dev/null +++ b/relay/mcp/transport_security_test.go @@ -0,0 +1,80 @@ +package mcp + +import ( + "context" + "net/http" + "testing" + "time" + + "github.com/stretchr/testify/require" + + "github.com/Laisky/one-api/model" +) + +// TestMCPClientRejectsCredentialedRemotePlaintextHTTP verifies the runtime transport blocks secrets before any remote plaintext network request begins. +// +// Parameters: +// - t: The test owns client configurations and transport-policy assertions. +// +// Return values: none; failures are reported through t. +func TestMCPClientRejectsCredentialedRemotePlaintextHTTP(t *testing.T) { + testCases := []struct { + name string + server model.MCPServer + headers map[string]string + }{ + { + name: "API key header", + server: model.MCPServer{BaseURL: "http://example.invalid/mcp"}, + headers: map[string]string{ + "X-API-Key": "secret", + }, + }, + { + name: "authorization header", + server: model.MCPServer{BaseURL: "http://example.invalid/mcp"}, + headers: map[string]string{ + "Authorization": "Bearer secret", + }, + }, + { + name: "custom authentication header", + server: model.MCPServer{BaseURL: "http://example.invalid/mcp"}, + headers: map[string]string{ + "X-Tenant-Identity": "secret", + }, + }, + { + name: "URL user information", + server: model.MCPServer{BaseURL: "http://user:secret@example.invalid/mcp"}, + }, + } + + for _, testCase := range testCases { + t.Run(testCase.name, func(t *testing.T) { + client := NewStreamableHTTPClient(&testCase.server, testCase.headers, time.Second) + _, err := client.DiscoverLatest(context.Background()) + require.ErrorContains(t, err, "require HTTPS") + }) + } +} + +// TestValidateMCPOutboundTransportPreservesExplicitCompatibility verifies plaintext remains available only when no credentials leave the host or the endpoint is loopback. +// +// Parameters: +// - t: The test owns request fixtures and transport-policy assertions. +// +// Return values: none; failures are reported through t. +func TestValidateMCPOutboundTransportPreservesExplicitCompatibility(t *testing.T) { + remoteRequest, err := http.NewRequest(http.MethodPost, "http://mcp.example.com/mcp", nil) + require.NoError(t, err) + require.NoError(t, validateMCPOutboundTransport(remoteRequest, false)) + + loopbackRequest, err := http.NewRequest(http.MethodPost, "http://127.0.0.1:8080/mcp", nil) + require.NoError(t, err) + require.NoError(t, validateMCPOutboundTransport(loopbackRequest, true)) + + secureRequest, err := http.NewRequest(http.MethodPost, "https://mcp.example.com/mcp", nil) + require.NoError(t, err) + require.NoError(t, validateMCPOutboundTransport(secureRequest, true)) +} diff --git a/relay/mcp/types.go b/relay/mcp/types.go index 41330c2f8b..d089e59420 100644 --- a/relay/mcp/types.go +++ b/relay/mcp/types.go @@ -1,76 +1,433 @@ package mcp import ( + "bytes" "encoding/json" "github.com/Laisky/errors/v2" ) -// ToolDescriptor describes a tool returned by MCP servers. +// ToolDescriptor describes a tool returned by MCP servers and preserves extension fields. type ToolDescriptor struct { - Name string `json:"name"` - Description string `json:"description,omitempty"` - InputSchema map[string]any `json:"inputSchema,omitempty"` + Name string `json:"-"` + Title string `json:"-"` + Description string `json:"-"` + InputSchema map[string]any `json:"-"` + OutputSchema map[string]any `json:"-"` + Annotations map[string]any `json:"-"` + Icons []map[string]any `json:"-"` + Meta map[string]any `json:"-"` + AdditionalFields map[string]any `json:"-"` } -// UnmarshalJSON decodes MCP tool descriptors while supporting multiple schema field names. +// MarshalJSON encodes a ToolDescriptor with current MCP field names and preserved extensions. +// +// Parameters: none. +// +// Return values: +// - []byte: the encoded MCP tool descriptor. +// - error: a wrapped encoding error when a field cannot be represented as JSON. +func (t ToolDescriptor) MarshalJSON() ([]byte, error) { + payload := make(map[string]any, len(t.AdditionalFields)+8) + for key, value := range t.AdditionalFields { + payload[key] = value + } + payload["name"] = t.Name + if t.Title != "" { + payload["title"] = t.Title + } + if t.Description != "" { + payload["description"] = t.Description + } + if t.InputSchema != nil { + payload["inputSchema"] = t.InputSchema + } + if t.OutputSchema != nil { + payload["outputSchema"] = t.OutputSchema + } + if t.Annotations != nil { + payload["annotations"] = t.Annotations + } + if t.Icons != nil { + payload["icons"] = t.Icons + } + if t.Meta != nil { + payload["_meta"] = t.Meta + } + encoded, err := json.Marshal(payload) + if err != nil { + return nil, errors.Wrap(err, "marshal mcp tool descriptor") + } + return encoded, nil +} + +// UnmarshalJSON decodes an MCP tool descriptor, accepts legacy aliases, and preserves extensions. +// +// Parameters: +// - data: a complete JSON object containing one MCP tool descriptor. +// +// Return values: +// - error: a wrapped decoding error when required field types or schemas are malformed. func (t *ToolDescriptor) UnmarshalJSON(data []byte) error { if t == nil { return errors.New("mcp tool descriptor is nil") } - var raw map[string]any + var raw map[string]json.RawMessage if err := json.Unmarshal(data, &raw); err != nil { return errors.Wrap(err, "unmarshal mcp tool descriptor") } - if name, ok := raw["name"].(string); ok { - t.Name = name + + var decoded ToolDescriptor + if err := decodeOptionalString(raw, "name", &decoded.Name); err != nil { + return errors.Wrap(err, "decode mcp tool name") } - if description, ok := raw["description"].(string); ok { - t.Description = description + if err := decodeOptionalString(raw, "title", &decoded.Title); err != nil { + return errors.Wrap(err, "decode mcp tool title") } - schema := raw["input_schema"] - if schema == nil { - schema = raw["inputSchema"] + if err := decodeOptionalString(raw, "description", &decoded.Description); err != nil { + return errors.Wrap(err, "decode mcp tool description") } - if schemaMap, ok := schema.(map[string]any); ok { - t.InputSchema = schemaMap - return nil + var err error + decoded.InputSchema, err = decodeOptionalObject(raw, "inputSchema", "input_schema") + if err != nil { + return errors.Wrap(err, "decode mcp tool input schema") + } + decoded.OutputSchema, err = decodeNullableOptionalObject(raw, "outputSchema", "output_schema") + if err != nil { + return errors.Wrap(err, "decode mcp tool output schema") + } + decoded.Annotations, err = decodeNullableOptionalObject(raw, "annotations") + if err != nil { + return errors.Wrap(err, "decode mcp tool annotations") } - if schema != nil { - encoded, err := json.Marshal(schema) - if err != nil { - return errors.Wrap(err, "marshal mcp tool schema") + decoded.Meta, err = decodeNullableOptionalObject(raw, "_meta") + if err != nil { + return errors.Wrap(err, "decode mcp tool metadata") + } + decoded.Icons, err = decodeNullableOptionalObjectSlice(raw, "icons") + if err != nil { + return errors.Wrap(err, "decode mcp tool icons") + } + + known := map[string]struct{}{ + "name": {}, "title": {}, "description": {}, + "inputSchema": {}, "input_schema": {}, + "outputSchema": {}, "output_schema": {}, + "annotations": {}, "icons": {}, "_meta": {}, + } + for key, value := range raw { + if _, exists := known[key]; exists { + continue + } + var extension any + if err := json.Unmarshal(value, &extension); err != nil { + return errors.Wrapf(err, "decode mcp tool extension %s", key) } - var parsed map[string]any - if err := json.Unmarshal(encoded, &parsed); err != nil { - return errors.Wrap(err, "decode mcp tool schema") + if decoded.AdditionalFields == nil { + decoded.AdditionalFields = make(map[string]any) } - t.InputSchema = parsed + decoded.AdditionalFields[key] = extension } + + *t = decoded return nil } -// CallToolResult represents a MCP tool call response. +// ListToolsResult represents a tools/list response across current and legacy MCP versions. +type ListToolsResult struct { + ResultType string `json:"resultType,omitempty"` + Tools []ToolDescriptor `json:"tools"` + NextCursor string `json:"nextCursor,omitempty"` + TTLMS int64 `json:"ttlMs,omitempty"` + CacheScope string `json:"cacheScope,omitempty"` + Meta map[string]any `json:"_meta,omitempty"` +} + +// CallToolRequestOptions carries MCP 2026-07-28 multi-round-trip retry fields. +type CallToolRequestOptions struct { + InputResponses map[string]any `json:"inputResponses,omitempty"` + RequestState string `json:"requestState,omitempty"` +} + +// CallToolResult represents an MCP tool result while preserving extension fields. type CallToolResult struct { - Content any `json:"content"` - IsError bool `json:"is_error,omitempty"` - Raw json.RawMessage `json:"-"` + ResultType string + Content any + StructuredContent any + IsError bool + InputRequests map[string]any + RequestState string + Meta map[string]any + AdditionalFields map[string]any + Raw json.RawMessage } -// UnmarshalJSON parses the MCP tool call result while keeping the raw payload. +// MarshalJSON encodes a CallToolResult with current camelCase field names and preserved extensions. +// +// Parameters: none. +// +// Return values: +// - []byte: the encoded MCP tool result. +// - error: a wrapped encoding error when a field cannot be represented as JSON. +func (c CallToolResult) MarshalJSON() ([]byte, error) { + payload := make(map[string]any, len(c.AdditionalFields)+7) + for key, value := range c.AdditionalFields { + payload[key] = value + } + if c.ResultType != "" { + payload["resultType"] = c.ResultType + } + if c.Content != nil { + payload["content"] = c.Content + } + if c.StructuredContent != nil { + payload["structuredContent"] = c.StructuredContent + } + if c.IsError { + payload["isError"] = true + } + if c.InputRequests != nil { + payload["inputRequests"] = c.InputRequests + } + if c.RequestState != "" { + payload["requestState"] = c.RequestState + } + if c.Meta != nil { + payload["_meta"] = c.Meta + } + encoded, err := json.Marshal(payload) + if err != nil { + return nil, errors.Wrap(err, "marshal mcp tool result") + } + return encoded, nil +} + +// UnmarshalJSON decodes a tool result, accepts legacy aliases, and preserves extension fields. +// +// Parameters: +// - data: a complete JSON object containing one MCP tool result. +// +// Return values: +// - error: a wrapped decoding error when a known result field has an invalid type. func (c *CallToolResult) UnmarshalJSON(data []byte) error { if c == nil { return errors.New("mcp tool result is nil") } - c.Raw = append(c.Raw[:0], data...) - var parsed struct { - Content any `json:"content"` - IsError bool `json:"is_error,omitempty"` - } - if err := json.Unmarshal(data, &parsed); err != nil { + var raw map[string]json.RawMessage + if err := json.Unmarshal(data, &raw); err != nil { return errors.Wrap(err, "unmarshal mcp tool result") } - c.Content = parsed.Content - c.IsError = parsed.IsError + + var decoded CallToolResult + decoded.Raw = append(decoded.Raw, data...) + if err := decodeOptionalString(raw, "resultType", &decoded.ResultType); err != nil { + return errors.Wrap(err, "decode mcp result type") + } + if encoded, exists := raw["content"]; exists { + if err := json.Unmarshal(encoded, &decoded.Content); err != nil { + return errors.Wrap(err, "decode mcp result content") + } + } + if encoded, exists := firstRawMessage(raw, "structuredContent", "structured_content"); exists { + if err := json.Unmarshal(encoded, &decoded.StructuredContent); err != nil { + return errors.Wrap(err, "decode mcp structured content") + } + } + if encoded, exists := firstRawMessage(raw, "isError", "is_error"); exists { + if err := json.Unmarshal(encoded, &decoded.IsError); err != nil { + return errors.Wrap(err, "decode mcp result error flag") + } + } + var err error + decoded.InputRequests, err = decodeNullableOptionalObject(raw, "inputRequests", "input_requests") + if err != nil { + return errors.Wrap(err, "decode mcp input requests") + } + if encoded, exists := firstRawMessage(raw, "requestState", "request_state"); exists { + if err := json.Unmarshal(encoded, &decoded.RequestState); err != nil { + return errors.Wrap(err, "decode mcp request state") + } + } + decoded.Meta, err = decodeNullableOptionalObject(raw, "_meta") + if err != nil { + return errors.Wrap(err, "decode mcp result metadata") + } + + known := map[string]struct{}{ + "resultType": {}, "content": {}, + "structuredContent": {}, "structured_content": {}, + "isError": {}, "is_error": {}, + "inputRequests": {}, "input_requests": {}, + "requestState": {}, "request_state": {}, "_meta": {}, + } + for key, value := range raw { + if _, exists := known[key]; exists { + continue + } + var extension any + if err := json.Unmarshal(value, &extension); err != nil { + return errors.Wrapf(err, "decode mcp result extension %s", key) + } + if decoded.AdditionalFields == nil { + decoded.AdditionalFields = make(map[string]any) + } + decoded.AdditionalFields[key] = extension + } + + *c = decoded return nil } + +// NormalizeCallToolResult fills protocol-required defaults without changing tool payloads. +// +// Parameters: +// - result: the tool result returned by an upstream MCP server; nil is allowed. +// +// Return values: +// - *CallToolResult: a non-nil result whose resultType is populated. +func NormalizeCallToolResult(result *CallToolResult) *CallToolResult { + if result == nil { + return &CallToolResult{ResultType: ResultTypeComplete} + } + if result.ResultType == "" { + result.ResultType = ResultTypeComplete + } + return result +} + +// decodeOptionalString decodes a string field when it is present in raw. +// +// Parameters: +// - raw: the source JSON object. +// - name: the field name to inspect. +// - destination: the string receiving a present value. +// +// Return values: +// - error: a wrapped type error for a present non-string value. +func decodeOptionalString(raw map[string]json.RawMessage, name string, destination *string) error { + encoded, exists := raw[name] + if !exists { + return nil + } + if err := json.Unmarshal(encoded, destination); err != nil { + return errors.Wrapf(err, "decode string field %s", name) + } + return nil +} + +// decodeOptionalObject decodes the first present object alias and rejects null or non-object values. +// +// Parameters: +// - raw: the source JSON object. +// - names: modern and legacy aliases in precedence order. +// +// Return values: +// - map[string]any: the decoded object, or nil when every alias is absent. +// - error: a wrapped type error for a present null or non-object value. +func decodeOptionalObject(raw map[string]json.RawMessage, names ...string) (map[string]any, error) { + encoded, exists := firstRawMessage(raw, names...) + if !exists { + return nil, nil + } + if bytes.Equal(bytes.TrimSpace(encoded), []byte("null")) { + return nil, errors.Errorf("field %s must be an object", names[0]) + } + var object map[string]any + if err := json.Unmarshal(encoded, &object); err != nil { + return nil, errors.Wrapf(err, "decode object field %s", names[0]) + } + if object == nil { + return nil, errors.Errorf("field %s must be an object", names[0]) + } + return object, nil +} + +// decodeNullableOptionalObject decodes an optional object alias and treats explicit null as absence. +// +// Parameters: +// - raw: The source JSON object supplies the candidate fields. +// - names: The modern and legacy aliases are checked in precedence order. +// +// Return values: +// - map[string]any: The decoded object is returned, or nil when the field is absent or null. +// - error: A wrapped type error is returned for a present non-object, non-null value. +func decodeNullableOptionalObject(raw map[string]json.RawMessage, names ...string) (map[string]any, error) { + encoded, exists := firstRawMessage(raw, names...) + if !exists || bytes.Equal(bytes.TrimSpace(encoded), []byte("null")) { + return nil, nil + } + return decodeOptionalObject(raw, names...) +} + +// decodeOptionalObjectSlice decodes an optional array whose elements must be JSON objects. +// +// Parameters: +// - raw: the source JSON object. +// - name: the array field name. +// +// Return values: +// - []map[string]any: the decoded object array, or nil when the field is absent. +// - error: a wrapped type error for null, non-array, or non-object elements. +func decodeOptionalObjectSlice(raw map[string]json.RawMessage, name string) ([]map[string]any, error) { + encoded, exists := raw[name] + if !exists { + return nil, nil + } + if bytes.Equal(bytes.TrimSpace(encoded), []byte("null")) { + return nil, errors.Errorf("field %s must be an array", name) + } + var values []json.RawMessage + if err := json.Unmarshal(encoded, &values); err != nil { + return nil, errors.Wrapf(err, "decode object array field %s", name) + } + objects := make([]map[string]any, 0, len(values)) + for index, value := range values { + if bytes.Equal(bytes.TrimSpace(value), []byte("null")) { + return nil, errors.Errorf("field %s element %d must be an object", name, index) + } + var object map[string]any + if err := json.Unmarshal(value, &object); err != nil { + return nil, errors.Wrapf(err, "decode object array field %s element %d", name, index) + } + if object == nil { + return nil, errors.Errorf("field %s element %d must be an object", name, index) + } + objects = append(objects, object) + } + return objects, nil +} + +// decodeNullableOptionalObjectSlice decodes an optional object array and treats field-level null as absence. +// +// Parameters: +// - raw: The source JSON object supplies the candidate array field. +// - name: The field name identifies the optional object array. +// +// Return values: +// - []map[string]any: The decoded objects are returned, or nil when the field is absent or null. +// - error: A wrapped type error is returned for a non-array value or a null/non-object element. +func decodeNullableOptionalObjectSlice(raw map[string]json.RawMessage, name string) ([]map[string]any, error) { + encoded, exists := raw[name] + if !exists || bytes.Equal(bytes.TrimSpace(encoded), []byte("null")) { + return nil, nil + } + return decodeOptionalObjectSlice(raw, name) +} + +// firstRawMessage returns the first present alias from a JSON object. +// +// Parameters: +// - raw: the source JSON object. +// - names: aliases in precedence order. +// +// Return values: +// - json.RawMessage: the encoded value for the first present alias. +// - bool: true when an alias was present, including an explicit null value. +func firstRawMessage(raw map[string]json.RawMessage, names ...string) (json.RawMessage, bool) { + for _, name := range names { + if encoded, exists := raw[name]; exists { + return encoded, true + } + } + return nil, false +} diff --git a/relay/mcp/types_test.go b/relay/mcp/types_test.go index 81d5b1449f..172c8d39f6 100644 --- a/relay/mcp/types_test.go +++ b/relay/mcp/types_test.go @@ -60,3 +60,47 @@ func TestToolDescriptor_UnmarshalJSON_HandlesSchemaFields(t *testing.T) { require.NotNil(t, underscore.InputSchema) require.Equal(t, "object", underscore.InputSchema["type"]) } + +// TestToolDescriptorOptionalNullFieldsAreAbsent verifies nullable optional fields do not reject an entire tools/list page. +// +// Parameters: +// - t: The test owns JSON compatibility assertions. +// +// Return values: none; failures are reported through t. +func TestToolDescriptorOptionalNullFieldsAreAbsent(t *testing.T) { + var descriptor ToolDescriptor + err := json.Unmarshal([]byte(`{"name":"echo","inputSchema":{"type":"object"},"outputSchema":null,"annotations":null,"icons":null,"_meta":null}`), &descriptor) + require.NoError(t, err) + require.Nil(t, descriptor.OutputSchema) + require.Nil(t, descriptor.Annotations) + require.Nil(t, descriptor.Icons) + require.Nil(t, descriptor.Meta) +} + +// TestToolDescriptorRequiredAndArrayFieldsRemainStrict verifies null compatibility does not weaken required schema validation. +// +// Parameters: +// - t: The test owns malformed descriptor assertions. +// +// Return values: none; failures are reported through t. +func TestToolDescriptorRequiredAndArrayFieldsRemainStrict(t *testing.T) { + var descriptor ToolDescriptor + require.Error(t, json.Unmarshal([]byte(`{"name":"echo","inputSchema":null}`), &descriptor)) + require.Error(t, json.Unmarshal([]byte(`{"name":"echo","inputSchema":{"type":"object"},"annotations":42}`), &descriptor)) + require.Error(t, json.Unmarshal([]byte(`{"name":"echo","inputSchema":{"type":"object"},"icons":[null]}`), &descriptor)) +} + +// TestCallToolResultOptionalNullObjectsAreAbsent verifies nullable result objects remain backward compatible. +// +// Parameters: +// - t: The test owns JSON compatibility assertions. +// +// Return values: none; failures are reported through t. +func TestCallToolResultOptionalNullObjectsAreAbsent(t *testing.T) { + var result CallToolResult + err := json.Unmarshal([]byte(`{"resultType":"complete","content":[],"structuredContent":null,"inputRequests":null,"_meta":null}`), &result) + require.NoError(t, err) + require.Nil(t, result.StructuredContent) + require.Nil(t, result.InputRequests) + require.Nil(t, result.Meta) +} diff --git a/router/relay.go b/router/relay.go index f743d670a9..d99e427ae0 100644 --- a/router/relay.go +++ b/router/relay.go @@ -8,6 +8,7 @@ import ( "github.com/Laisky/one-api/middleware" ) +// SetRelayRouter registers the public inference and MCP relay endpoints. func SetRelayRouter(router *gin.Engine) { // Rewrite various Claude Code prefixes to the canonical /v1/messages path. // Put this before other middlewares to avoid double-running them on redispatch. @@ -42,10 +43,9 @@ func SetRelayRouter(router *gin.Engine) { modelsRouter.GET("/:model", controller.RetrieveModel) } - // MCP Streamable HTTP transport: a single endpoint serves POST (JSON-RPC - // requests/notifications), GET (optional server-initiated SSE), and - // DELETE (session termination). The handler dispatches by method. - router.Any("/mcp", middleware.TokenAuth(), controller.MCPProxy) + // MCP Streamable HTTP transport: a single endpoint serves MCP 2026-07-28 + // requests and transparently delegates legacy initialize/session clients. + router.Any("/mcp", middleware.TokenAuth(), controller.MCPProxyLatest) relayMws := []gin.HandlerFunc{ // Track in-flight requests for graceful shutdown/drain