fix(mcp): harden tool schemas and request routing

This commit is contained in:
johnnyjoygh
2026-07-25 01:12:47 +08:00
parent 03e34bd0da
commit f43cfc3cfc
8 changed files with 598 additions and 40 deletions
+7 -4
View File
@@ -17,7 +17,10 @@ import (
"github.com/usememos/memos/store"
)
const maxAPIRequestBytes = 256 << 20
// MaxAPIRequestBytes caps the size of a request body accepted by the API. The
// in-process MCP endpoint forwards every tool call through these same routes, so
// it derives its own limit from this constant to keep the two gates in lockstep.
const MaxAPIRequestBytes = 256 << 20
type APIV1Service struct {
v1pb.UnimplementedInstanceServiceServer
@@ -123,7 +126,7 @@ func (s *APIV1Service) RegisterGateway(ctx context.Context, echoServer *echo.Ech
gwGroup := echoServer.Group("")
// Register SSE endpoint with same CORS as rest of /api/v1.
RegisterSSERoutes(gwGroup, s.SSEHub, s.Store, s.Secret)
handler := echo.WrapHandler(http.MaxBytesHandler(gwMux, maxAPIRequestBytes))
handler := echo.WrapHandler(http.MaxBytesHandler(gwMux, MaxAPIRequestBytes))
gwGroup.Any("/api/v1/*", handler)
gwGroup.Any("/file/*", handler)
@@ -138,10 +141,10 @@ func (s *APIV1Service) RegisterGateway(ctx context.Context, echoServer *echo.Ech
)
connectMux := http.NewServeMux()
connectHandler := NewConnectServiceHandler(s)
connectHandler.RegisterConnectHandlers(connectMux, connectInterceptors, connect.WithReadMaxBytes(maxAPIRequestBytes))
connectHandler.RegisterConnectHandlers(connectMux, connectInterceptors, connect.WithReadMaxBytes(MaxAPIRequestBytes))
connectGroup := echoServer.Group("")
connectGroup.Any("/memos.api.v1.*", echo.WrapHandler(http.MaxBytesHandler(connectMux, maxAPIRequestBytes)))
connectGroup.Any("/memos.api.v1.*", echo.WrapHandler(http.MaxBytesHandler(connectMux, MaxAPIRequestBytes)))
return nil
}
+21 -13
View File
@@ -49,17 +49,18 @@ fast on any inconsistency:
`RegisterRoutes` binds `echoServer.Any("/mcp", ...)`. Each request:
1. `isAllowedMCPOrigin` (`origin.go`) rejects disallowed cross-origin browser requests with `403`.
2. The SDK streamable handler dispatches the MCP message.
3. On a `tools/call` request, `newMCPToolHandler` (`service.go`) decodes the JSON
2. The request body is capped at 256 MiB before the SDK reads it.
3. The SDK streamable handler dispatches the MCP message.
4. On a `tools/call` request, `newMCPToolHandler` (`service.go`) decodes the JSON
arguments into a map.
4. `validateToolArguments` (`validation.go`) checks them against the tool's
5. `validateToolArguments` (`validation.go`) checks them against the tool's
input schema.
5. The caller's `Authorization` header is read from the request (`request.Extra.Header` on the SDK's `*sdkmcp.CallToolRequest`).
6. `apiAdapter.execute` (`adapter.go`) builds the API request
6. The caller's `Authorization` header is read from the request (`request.Extra.Header` on the SDK's `*sdkmcp.CallToolRequest`).
7. `apiAdapter.execute` (`adapter.go`) builds the API request
(`buildAPIRequest`: path-parameter substitution, query encoding, JSON body),
forwards the bearer token, and runs it against the Echo server through an
`httptest.ResponseRecorder`.
7. The recorder body is decoded; a non-2xx status becomes a tool error
8. The recorder body is decoded; a non-2xx status becomes a tool error
(`newToolErrorResult`), otherwise the value is wrapped by
`newStructuredToolResult`.
@@ -86,6 +87,10 @@ use `$ref`. `openapi.go` resolves these into local definitions:
parameter stays in `required`.
- A request body becomes a single `body` property; a required body adds `body`
to `required`. Body `$defs` are lifted to the schema's top-level `$defs`.
- Per-operation overrides relax resource-level requirements for create and
partial-update bodies and remove fields already supplied by a path binding
from `body: "*"` schemas. Memo updates may omit `updateMask` so the REST
gateway can infer it from the fields present in the request body.
- The schema sets `"additionalProperties": false`.
The output schema is the operation's 200 `application/json` schema. When a 200
@@ -100,6 +105,7 @@ response has no JSON body, the fallback is:
- **Endpoint:** `POST /mcp` (the SDK may also use `GET`/`DELETE` on the same
path for the Streamable HTTP transport).
- **Transport:** Streamable HTTP, **stateless**, JSON responses.
- **Request size:** request bodies are limited to 256 MiB before SDK dispatch.
- **Auth:** the caller's `Authorization: Bearer <token>` header is forwarded to
the in-process API request. Mutating tools therefore require a valid token
(personal access token or access token); public reads may work without one,
@@ -172,15 +178,16 @@ by `_`. So `MemoService_ListMemos → memo_list_memos`.
| DELETE | false | true | true |
| other (POST, PATCH, …) | false | false | false |
A per-operation override (`idempotentOperationIDs`) then corrects cases the
method heuristic gets wrong: `MemoService_SetMemoAttachments` and
`MemoService_SetMemoRelations` are PATCH but declaratively replace the full set
on a memo, so they report `IdempotentHint: true`.
Per-operation overrides then correct cases the method heuristic gets wrong.
`MemoService_SetMemoAttachments` and `MemoService_SetMemoRelations` are PATCH
but declaratively replace the full set on a memo, so they report both
`IdempotentHint: true` and `DestructiveHint: true`. `MemoService_UpdateMemo`
also reports `DestructiveHint: true` because it can overwrite existing fields.
`OpenWorldHint` is `false` for all tools. Annotations are client hints; they do
not replace API authorization.
**Result shape.** Every result carries object-shaped `structuredContent`
**Result shape.** Every successful result carries object-shaped `structuredContent`
(`normalizeStructuredContent` in `result.go`):
- a JSON object is returned unchanged;
@@ -194,8 +201,9 @@ where collection tools returned a bare array that strict MCP clients reject.
## Error handling
Failures are returned as MCP tool errors (`CallToolResult` with `IsError: true`
and an `error.message` in `structuredContent`), not JSON-RPC protocol errors —
the handler returns `(result, nil)`:
and a text content block), not JSON-RPC protocol errors — the handler returns
`(result, nil)`. Error results omit `structuredContent` so strict clients do not
validate an error payload against the tool's success-only output schema:
| Failure | Result |
| --- | --- |
+79 -14
View File
@@ -92,46 +92,111 @@ func buildAPIRequest(ctx context.Context, operation *openAPIOperation, arguments
}
func substitutePathParameters(operation *openAPIOperation, arguments map[string]any) (string, error) {
path := operation.Path
for _, parameter := range operation.Parameters {
if parameter.In != "path" {
continue
}
value, ok := arguments[parameter.Name]
if !ok || value == nil || valueToString(value) == "" {
return "", errors.Errorf(`missing required path parameter "%s"`, parameter.Name)
}
id := trimResourceNamePrefix(operation.Path, parameter.Name, valueToString(value))
path = strings.ReplaceAll(path, "{"+parameter.Name+"}", url.PathEscape(id))
}
// Resolve placeholders in the order they appear in the path so a nested
// resource name (e.g. "memos/abc123/reactions/reaction456") can be matched
// against its already-resolved parent segments. Each placeholder is resolved
// exactly once from the argument map and cached in resolved, so a value that
// itself contains a "{" can never be re-expanded into a longer prefix.
path := operation.Path
resolved := map[string]string{}
for _, name := range pathPlaceholderNames(operation.Path) {
value, ok := arguments[name]
if !ok {
continue
}
id := trimResourceNamePrefix(operation.Path, name, valueToString(value), resolved)
resolved[name] = id
path = strings.ReplaceAll(path, "{"+name+"}", url.PathEscape(id))
}
return path, nil
}
// pathPlaceholderNames returns the "{name}" placeholder names in the order they
// appear in path.
func pathPlaceholderNames(path string) []string {
var names []string
for {
start := strings.Index(path, "{")
if start < 0 {
break
}
endOffset := strings.Index(path[start:], "}")
if endOffset < 0 {
break
}
end := start + endOffset
names = append(names, path[start+1:end])
path = path[end+1:]
}
return names
}
// trimResourceNamePrefix accepts canonical resource names for path parameters.
// The API returns names like "memos/abc123", but the REST paths take the bare
// ID ("/api/v1/memos/{memo}"), so clients that round-trip a returned name
// would otherwise request "/api/v1/memos/memos/abc123" and get a 404. When the
// placeholder directly follows its collection segment and the value carries
// that collection prefix, strip the prefix; bare IDs pass through unchanged.
func trimResourceNamePrefix(path, parameterName, value string) string {
// It uses the already-resolved parent segments so nested names such as
// "memos/abc123/reactions/reaction456" are accepted only when their parent
// segments match the other arguments. Bare IDs pass through unchanged.
func trimResourceNamePrefix(path, parameterName, value string, resolved map[string]string) string {
placeholder := "/{" + parameterName + "}"
index := strings.Index(path, placeholder)
if index < 0 {
return value
}
head := path[:index]
collection := head[strings.LastIndex(head, "/")+1:]
if collection == "" {
prefix, ok := resolvedResourceNamePrefix(path[:index], resolved)
if !ok || prefix == "" {
return value
}
id, ok := strings.CutPrefix(value, collection+"/")
id, ok := strings.CutPrefix(value, prefix+"/")
if !ok || id == "" || strings.Contains(id, "/") {
return value
}
return id
}
// resolvedResourceNamePrefix rebuilds the collection prefix preceding a
// placeholder by substituting each earlier placeholder with its already-resolved
// bare id. It reads resolved ids from the map instead of recursing, and rejects
// any id that is not a single bare segment, so every iteration removes one "{"
// and the loop always terminates.
func resolvedResourceNamePrefix(prefix string, resolved map[string]string) (string, bool) {
const apiPrefix = "/api/v1/"
prefix, ok := strings.CutPrefix(prefix, apiPrefix)
if !ok {
return "", false
}
for {
start := strings.Index(prefix, "{")
if start < 0 {
break
}
endOffset := strings.Index(prefix[start:], "}")
if endOffset < 0 {
return "", false
}
end := start + endOffset
parameterName := prefix[start+1 : end]
id, ok := resolved[parameterName]
if !ok || id == "" || strings.ContainsAny(id, "/{}") {
return "", false
}
prefix = prefix[:start] + id + prefix[end+1:]
}
return strings.Trim(prefix, "/"), true
}
func valueToString(value any) string {
switch typed := value.(type) {
case string:
+104
View File
@@ -147,6 +147,110 @@ func TestBuildAPIRequestAcceptsResourceNamesForPathParameters(t *testing.T) {
}
}
func TestBuildAPIRequestAcceptsHierarchicalResourceNamesForPathParameters(t *testing.T) {
tests := []struct {
name string
path string
parameters []string
arguments map[string]any
wantPath string
}{
{
name: "canonical reaction name",
path: "/api/v1/memos/{memo}/reactions/{reaction}",
parameters: []string{"memo", "reaction"},
arguments: map[string]any{
"memo": "memos/abc123",
"reaction": "memos/abc123/reactions/reaction456",
},
wantPath: "/api/v1/memos/abc123/reactions/reaction456",
},
{
name: "canonical nested name with bare parent id",
path: "/api/v1/memos/{memo}/reactions/{reaction}",
parameters: []string{"memo", "reaction"},
arguments: map[string]any{
"memo": "abc123",
"reaction": "memos/abc123/reactions/reaction456",
},
wantPath: "/api/v1/memos/abc123/reactions/reaction456",
},
{
name: "parameter declaration order does not matter",
path: "/api/v1/users/{user}/shortcuts/{shortcut}",
parameters: []string{"shortcut", "user"},
arguments: map[string]any{
"user": "users/user123",
"shortcut": "users/user123/shortcuts/shortcut456",
},
wantPath: "/api/v1/users/user123/shortcuts/shortcut456",
},
{
name: "canonical nested name on action route",
path: "/api/v1/users/{user}/webhooks/{webhook}:getSigningSecret",
parameters: []string{"user", "webhook"},
arguments: map[string]any{
"user": "users/user123",
"webhook": "users/user123/webhooks/webhook456",
},
wantPath: "/api/v1/users/user123/webhooks/webhook456:getSigningSecret",
},
{
name: "mismatched parent is left untouched",
path: "/api/v1/memos/{memo}/reactions/{reaction}",
parameters: []string{"memo", "reaction"},
arguments: map[string]any{
"memo": "memos/abc123",
"reaction": "memos/other/reactions/reaction456",
},
wantPath: "/api/v1/memos/abc123/reactions/memos%2Fother%2Freactions%2Freaction456",
},
{
name: "extra nested segment is left untouched",
path: "/api/v1/memos/{memo}/reactions/{reaction}",
parameters: []string{"memo", "reaction"},
arguments: map[string]any{
"memo": "memos/abc123",
"reaction": "memos/abc123/reactions/reaction456/extra",
},
wantPath: "/api/v1/memos/abc123/reactions/memos%2Fabc123%2Freactions%2Freaction456%2Fextra",
},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
parameters := make([]openAPIParameter, 0, len(test.parameters))
for _, name := range test.parameters {
parameters = append(parameters, openAPIParameter{Name: name, In: "path", Required: true, Schema: jsonSchema{"type": "string"}})
}
operation := &openAPIOperation{
Method: "GET",
Path: test.path,
Parameters: parameters,
}
req, err := buildAPIRequest(context.Background(), operation, test.arguments, "")
require.NoError(t, err)
require.Equal(t, test.wantPath, req.URL.EscapedPath())
})
}
}
func TestBuildAPIRequestAcceptsCanonicalReactionNameForCuratedOperation(t *testing.T) {
spec, err := loadMCPServiceOpenAPISpec()
require.NoError(t, err)
registry, err := buildOperationRegistry(spec)
require.NoError(t, err)
operation := registry["MemoService_DeleteMemoReaction"]
require.NotNil(t, operation)
req, err := buildAPIRequest(context.Background(), operation, map[string]any{
"memo": "memos/abc123",
"reaction": "memos/abc123/reactions/reaction456",
}, "")
require.NoError(t, err)
require.Equal(t, "/api/v1/memos/abc123/reactions/reaction456", req.URL.EscapedPath())
}
func TestBuildAPIRequestRequiresPathParameters(t *testing.T) {
operation := &openAPIOperation{
Method: "GET",
+111 -1
View File
@@ -1,6 +1,7 @@
package mcp
import (
"maps"
"regexp"
"strings"
@@ -42,8 +43,78 @@ type registeredOperation struct {
InputSchema jsonSchema
}
type requestBodySchemaOverride struct {
required []string
omittedProperties []string
// minProperties, when > 0, requires the body to carry at least that many
// properties. It replaces a cleared required list for partial updates so an
// empty body is rejected up front instead of failing later at the API.
minProperties int
}
// requestBodySchemaOverrides adjusts resource schemas to match how each HTTP
// binding consumes its request body. Resource-level required fields are too
// strict for partial updates, while body: "*" schemas include fields already
// supplied by the path.
var requestBodySchemaOverrides = map[string]requestBodySchemaOverride{
"MemoService_CreateMemo": {
required: []string{"content"},
},
"MemoService_UpdateMemo": {
required: []string{},
omittedProperties: []string{"name"},
minProperties: 1,
},
"MemoService_CreateMemoComment": {
required: []string{"content"},
},
"MemoService_SetMemoAttachments": {
required: []string{"attachments"},
omittedProperties: []string{"name"},
},
"MemoService_SetMemoRelations": {
required: []string{"relations"},
omittedProperties: []string{"name"},
},
"MemoService_UpsertMemoReaction": {
required: []string{"reaction"},
omittedProperties: []string{"name"},
},
}
var wordBoundary = regexp.MustCompile(`([a-z0-9])([A-Z])`)
// validateOperationOverrides fails fast when a per-operation override table
// references an operation that is not in the registry (e.g. after a proto RPC
// rename). Without this, a stale key would silently miss and the renamed
// operation would lose its schema/annotation override with no error.
func validateOperationOverrides(registry map[string]*openAPIOperation) error {
tables := []struct {
name string
ids []string
}{
{"requestBodySchemaOverrides", mapKeys(requestBodySchemaOverrides)},
{"idempotentOperationIDs", mapKeys(idempotentOperationIDs)},
{"destructiveOperationIDs", mapKeys(destructiveOperationIDs)},
}
for _, table := range tables {
for _, operationID := range table.ids {
if _, ok := registry[operationID]; !ok {
return errors.Errorf("%s references unknown operation %q", table.name, operationID)
}
}
}
return nil
}
func mapKeys[V any](m map[string]V) []string {
keys := make([]string, 0, len(m))
for key := range m {
keys = append(keys, key)
}
return keys
}
func buildCuratedTools(registry map[string]*openAPIOperation) ([]*sdkmcp.Tool, map[string]*registeredOperation, error) {
tools := make([]*sdkmcp.Tool, 0, len(curatedOperationIDs))
operations := map[string]*registeredOperation{}
@@ -61,6 +132,10 @@ func buildCuratedTools(registry map[string]*openAPIOperation) ([]*sdkmcp.Tool, m
tools = append(tools, tool)
operations[tool.Name] = registered
}
if err := validateOperationOverrides(registry); err != nil {
return nil, nil, err
}
return tools, operations, nil
}
@@ -160,7 +235,30 @@ func requestBodySchema(operation *openAPIOperation) jsonSchema {
if operation.RequestBodySchema == nil {
return jsonSchema{"type": "object"}
}
return cloneSchema(operation.RequestBodySchema)
schema := cloneSchema(operation.RequestBodySchema)
override, ok := requestBodySchemaOverrides[operation.OperationID]
if !ok {
return schema
}
if len(override.required) == 0 {
delete(schema, "required")
} else {
schema["required"] = override.required
}
if len(override.omittedProperties) > 0 {
properties := maps.Clone(schemaProperties(schema["properties"]))
for _, name := range override.omittedProperties {
delete(properties, name)
}
schema["properties"] = properties
}
if override.minProperties > 0 {
schema["minProperties"] = override.minProperties
}
return schema
}
func outputSchemaForOperation(operation *openAPIOperation) jsonSchema {
@@ -197,6 +295,14 @@ var idempotentOperationIDs = map[string]bool{
"MemoService_SetMemoRelations": true,
}
// destructiveOperationIDs lists mutating operations that can overwrite or
// remove existing user data despite not using DELETE.
var destructiveOperationIDs = map[string]bool{
"MemoService_UpdateMemo": true,
"MemoService_SetMemoAttachments": true,
"MemoService_SetMemoRelations": true,
}
// annotationsForOperation derives the method-based annotations and then applies
// per-operation overrides that the HTTP method alone cannot express.
func annotationsForOperation(operation *openAPIOperation, title string) *sdkmcp.ToolAnnotations {
@@ -204,6 +310,10 @@ func annotationsForOperation(operation *openAPIOperation, title string) *sdkmcp.
if idempotentOperationIDs[operation.OperationID] {
annotations.IdempotentHint = true
}
if destructiveOperationIDs[operation.OperationID] {
destructive := true
annotations.DestructiveHint = &destructive
}
return annotations
}
+130 -4
View File
@@ -96,14 +96,114 @@ func TestBuildToolFromOperationIncludesRequestBodySchema(t *testing.T) {
err = validateToolArguments(input, map[string]any{
"body": map[string]any{
"state": "NORMAL",
"content": "hello",
"visibility": "PRIVATE",
"content": "hello",
},
})
require.NoError(t, err)
}
func TestBuildToolFromOperationTailorsRequestBodySchemas(t *testing.T) {
spec, err := loadOpenAPISpec("../../../proto/gen/openapi.yaml")
require.NoError(t, err)
registry, err := buildOperationRegistry(spec)
require.NoError(t, err)
tests := []struct {
name string
operationID string
arguments map[string]any
omittedProperties []string
}{
{
name: "partial memo update",
operationID: "MemoService_UpdateMemo",
arguments: map[string]any{
"memo": "memos/abc123",
"body": map[string]any{"content": "updated"},
},
omittedProperties: []string{"name"},
},
{
name: "comment defaults state and visibility",
operationID: "MemoService_CreateMemoComment",
arguments: map[string]any{
"memo": "memos/abc123",
"body": map[string]any{"content": "comment"},
},
},
{
name: "set attachments gets name from path",
operationID: "MemoService_SetMemoAttachments",
arguments: map[string]any{
"memo": "memos/abc123",
"body": map[string]any{"attachments": []any{}},
},
omittedProperties: []string{"name"},
},
{
name: "set relations gets name from path",
operationID: "MemoService_SetMemoRelations",
arguments: map[string]any{
"memo": "memos/abc123",
"body": map[string]any{"relations": []any{}},
},
omittedProperties: []string{"name"},
},
{
name: "upsert reaction gets name from path",
operationID: "MemoService_UpsertMemoReaction",
arguments: map[string]any{
"memo": "memos/abc123",
"body": map[string]any{
"reaction": map[string]any{
"contentId": "memos/abc123",
"reactionType": "👍",
},
},
},
omittedProperties: []string{"name"},
},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
tool, _ := buildToolFromOperation(registry[test.operationID])
input, ok := tool.InputSchema.(jsonSchema)
require.True(t, ok)
require.NoError(t, validateToolArguments(input, test.arguments))
properties := schemaProperties(input["properties"])
body := schemaProperties(properties["body"])
bodyProperties := schemaProperties(body["properties"])
for _, property := range test.omittedProperties {
require.NotContains(t, bodyProperties, property)
}
})
}
}
func TestBuildToolFromOperationRejectsEmptyMemoUpdateBody(t *testing.T) {
spec, err := loadOpenAPISpec("../../../proto/gen/openapi.yaml")
require.NoError(t, err)
registry, err := buildOperationRegistry(spec)
require.NoError(t, err)
tool, _ := buildToolFromOperation(registry["MemoService_UpdateMemo"])
input, ok := tool.InputSchema.(jsonSchema)
require.True(t, ok)
// An empty body carries no fields to update; reject it at the schema instead of
// letting the gateway infer an empty field mask and fail late.
require.Error(t, validateToolArguments(input, map[string]any{
"memo": "memos/abc123",
"body": map[string]any{},
}))
require.NoError(t, validateToolArguments(input, map[string]any{
"memo": "memos/abc123",
"body": map[string]any{"content": "updated"},
}))
}
func TestBuildToolFromOperationExposesCreateAttachment(t *testing.T) {
spec, err := loadOpenAPISpec("../../../proto/gen/openapi.yaml")
require.NoError(t, err)
@@ -183,10 +283,23 @@ func TestBuildToolFromOperationMarksSetOperationsIdempotent(t *testing.T) {
// override restores the declarative "set" semantics.
require.True(t, tool.Annotations.IdempotentHint, operationID)
require.False(t, tool.Annotations.ReadOnlyHint, operationID)
require.False(t, *tool.Annotations.DestructiveHint, operationID)
require.True(t, *tool.Annotations.DestructiveHint, operationID)
}
}
func TestBuildToolFromOperationMarksUpdateMemoDestructive(t *testing.T) {
spec, err := loadOpenAPISpec("../../../proto/gen/openapi.yaml")
require.NoError(t, err)
registry, err := buildOperationRegistry(spec)
require.NoError(t, err)
tool, operation := buildToolFromOperation(registry["MemoService_UpdateMemo"])
require.Equal(t, "PATCH", operation.Method)
require.False(t, tool.Annotations.ReadOnlyHint)
require.True(t, *tool.Annotations.DestructiveHint)
require.False(t, tool.Annotations.IdempotentHint)
}
func TestBuildCuratedToolsHasUniqueNames(t *testing.T) {
spec, err := loadOpenAPISpec("../../../proto/gen/openapi.yaml")
require.NoError(t, err)
@@ -221,6 +334,19 @@ func TestBuildCuratedToolsRejectsMissingOperation(t *testing.T) {
require.ErrorContains(t, err, "not found")
}
func TestValidateOperationOverridesRejectsStaleKey(t *testing.T) {
spec, err := loadOpenAPISpec("../../../proto/gen/openapi.yaml")
require.NoError(t, err)
registry, err := buildOperationRegistry(spec)
require.NoError(t, err)
require.NoError(t, validateOperationOverrides(registry))
// A renamed/removed operation must be reported instead of silently losing its
// override.
delete(registry, "MemoService_UpdateMemo")
require.ErrorContains(t, validateOperationOverrides(registry), "MemoService_UpdateMemo")
}
func TestBuildCuratedToolsRejectsDuplicateToolNames(t *testing.T) {
registry := make(map[string]*openAPIOperation, len(curatedOperationIDs))
for _, operationID := range curatedOperationIDs {
+9 -4
View File
@@ -6,14 +6,20 @@ import (
"net/http"
"github.com/labstack/echo/v5"
"github.com/labstack/echo/v5/middleware"
sdkmcp "github.com/modelcontextprotocol/go-sdk/mcp"
"github.com/pkg/errors"
"gopkg.in/yaml.v3"
"github.com/usememos/memos/internal/profile"
memosproto "github.com/usememos/memos/proto"
apiv1 "github.com/usememos/memos/server/router/api/v1"
)
// maxMCPRequestBytes caps the /mcp request body. It tracks the API limit because
// every tool call is forwarded in-process through the API routes.
const maxMCPRequestBytes int64 = apiv1.MaxAPIRequestBytes
// MCPService serves the OpenAPI-driven MCP endpoint.
type MCPService struct {
profile *profile.Profile
@@ -52,7 +58,7 @@ func NewMCPService(profile *profile.Profile, echoServer *echo.Echo) (*MCPService
server.AddTool(tool, newMCPToolHandler(adapter, operation))
}
handler := sdkmcp.NewStreamableHTTPHandler(func(*http.Request) *sdkmcp.Server {
streamableHandler := sdkmcp.NewStreamableHTTPHandler(func(*http.Request) *sdkmcp.Server {
return server
}, &sdkmcp.StreamableHTTPOptions{
Stateless: true,
@@ -65,11 +71,10 @@ func NewMCPService(profile *profile.Profile, echoServer *echo.Echo) (*MCPService
// CSRF / DNS-rebinding protection instead.
DisableLocalhostProtection: true,
})
return &MCPService{
profile: profile,
operationsByTool: operationsByTool,
handler: handler,
handler: streamableHandler,
}, nil
}
@@ -113,5 +118,5 @@ func (s *MCPService) RegisterRoutes(echoServer *echo.Echo) {
}
s.handler.ServeHTTP(c.Response(), request)
return nil
})
}, middleware.BodyLimit(maxMCPRequestBytes))
}
+137
View File
@@ -178,6 +178,143 @@ func TestMCPToolCallReturnsObjectStructuredContent(t *testing.T) {
}, result["structuredContent"])
}
func TestMCPToolCallAllowsGatewayToInferMemoUpdateMask(t *testing.T) {
echoServer := echo.New()
routeHits := 0
echoServer.PATCH("/api/v1/memos/:memo", func(c *echo.Context) error {
routeHits++
require.Equal(t, "abc123", c.Param("memo"))
require.Empty(t, c.QueryParam("updateMask"))
body := map[string]any{}
require.NoError(t, json.NewDecoder(c.Request().Body).Decode(&body))
require.Equal(t, map[string]any{"content": "updated"}, body)
return c.JSON(http.StatusOK, map[string]any{
"name": "memos/abc123",
"content": "updated",
})
})
service, err := NewMCPService(&profile.Profile{Version: "test-version"}, echoServer)
require.NoError(t, err)
service.RegisterRoutes(echoServer)
initializeMCP(t, echoServer)
response := postMCP(t, echoServer, map[string]any{
"jsonrpc": "2.0",
"id": 2,
"method": "tools/call",
"params": map[string]any{
"name": "memo_update_memo",
"arguments": map[string]any{
"memo": "memos/abc123",
"body": map[string]any{"content": "updated"},
},
},
})
result, ok := response["result"].(map[string]any)
require.True(t, ok)
require.NotEqual(t, true, result["isError"])
require.Equal(t, map[string]any{
"name": "memos/abc123",
"content": "updated",
}, result["structuredContent"])
require.Equal(t, 1, routeHits)
}
func TestMCPToolCallBindsMemoFromPathForBodyStarOperations(t *testing.T) {
register := func(e *echo.Echo, method, path string, handler func(*echo.Context) error) {
switch method {
case http.MethodPatch:
e.PATCH(path, handler)
case http.MethodPost:
e.POST(path, handler)
default:
t.Fatalf("unsupported method %q", method)
}
}
tests := []struct {
name string
method string
path string
toolName string
body map[string]any
response map[string]any
}{
{
name: "set attachments",
method: http.MethodPatch,
path: "/api/v1/memos/:memo/attachments",
toolName: "memo_set_memo_attachments",
body: map[string]any{"attachments": []any{}},
response: map[string]any{},
},
{
name: "set relations",
method: http.MethodPatch,
path: "/api/v1/memos/:memo/relations",
toolName: "memo_set_memo_relations",
body: map[string]any{"relations": []any{}},
response: map[string]any{},
},
{
name: "upsert reaction",
method: http.MethodPost,
path: "/api/v1/memos/:memo/reactions",
toolName: "memo_upsert_memo_reaction",
body: map[string]any{
"reaction": map[string]any{
"contentId": "memos/abc123",
"reactionType": "👍",
},
},
response: map[string]any{"reactionType": "👍"},
},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
echoServer := echo.New()
routeHits := 0
register(echoServer, test.method, test.path, func(c *echo.Context) error {
routeHits++
// The memo must be bound from the path, and the omitted "name"
// property must not reappear in the forwarded body.
require.Equal(t, "abc123", c.Param("memo"))
body := map[string]any{}
require.NoError(t, json.NewDecoder(c.Request().Body).Decode(&body))
require.NotContains(t, body, "name")
return c.JSON(http.StatusOK, test.response)
})
service, err := NewMCPService(&profile.Profile{Version: "test-version"}, echoServer)
require.NoError(t, err)
service.RegisterRoutes(echoServer)
initializeMCP(t, echoServer)
response := postMCP(t, echoServer, map[string]any{
"jsonrpc": "2.0",
"id": 2,
"method": "tools/call",
"params": map[string]any{
"name": test.toolName,
"arguments": map[string]any{
"memo": "memos/abc123",
"body": test.body,
},
},
})
result, ok := response["result"].(map[string]any)
require.True(t, ok)
require.NotEqual(t, true, result["isError"], result)
require.Equal(t, 1, routeHits)
})
}
}
func TestMCPToolCallRejectsInvalidArguments(t *testing.T) {
echoServer := echo.New()
routeHits := 0