mirror of https://github.com/usememos/memos
feat(mcp): harden tool exposure and side effects (#5850)
boojack
1 month ago
committed by
GitHub
parent
0fc1dab28b
commit
583c3d24f4
@ -0,0 +1,74 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"github.com/pkg/errors"
|
||||
|
||||
v1pb "github.com/usememos/memos/proto/gen/api/v1"
|
||||
apiv1 "github.com/usememos/memos/server/router/api/v1"
|
||||
"github.com/usememos/memos/store"
|
||||
)
|
||||
|
||||
func visibilityToProto(visibility store.Visibility) v1pb.Visibility {
|
||||
switch visibility {
|
||||
case store.Protected:
|
||||
return v1pb.Visibility_PROTECTED
|
||||
case store.Public:
|
||||
return v1pb.Visibility_PUBLIC
|
||||
default:
|
||||
return v1pb.Visibility_PRIVATE
|
||||
}
|
||||
}
|
||||
|
||||
func rowStatusToProto(rowStatus store.RowStatus) v1pb.State {
|
||||
switch rowStatus {
|
||||
case store.Archived:
|
||||
return v1pb.State_ARCHIVED
|
||||
default:
|
||||
return v1pb.State_NORMAL
|
||||
}
|
||||
}
|
||||
|
||||
func (s *MCPService) loadMemoJSONByName(ctx context.Context, name string) (memoJSON, error) {
|
||||
uid, err := parseMemoUID(name)
|
||||
if err != nil {
|
||||
return memoJSON{}, err
|
||||
}
|
||||
memo, err := s.store.GetMemo(ctx, &store.FindMemo{UID: &uid})
|
||||
if err != nil {
|
||||
return memoJSON{}, errors.Wrap(err, "failed to get memo")
|
||||
}
|
||||
if memo == nil {
|
||||
return memoJSON{}, errors.New("memo not found")
|
||||
}
|
||||
return storeMemoToJSONWithStore(ctx, s.store, memo)
|
||||
}
|
||||
|
||||
func (s *MCPService) loadReactionJSONByID(ctx context.Context, reactionID int32) (reactionJSON, error) {
|
||||
reaction, err := s.store.GetReaction(ctx, &store.FindReaction{ID: &reactionID})
|
||||
if err != nil {
|
||||
return reactionJSON{}, errors.Wrap(err, "failed to get reaction")
|
||||
}
|
||||
if reaction == nil {
|
||||
return reactionJSON{}, errors.New("reaction not found")
|
||||
}
|
||||
creator, err := lookupUsername(ctx, s.store, reaction.CreatorID)
|
||||
if err != nil {
|
||||
return reactionJSON{}, errors.Wrap(err, "failed to resolve reaction creator")
|
||||
}
|
||||
return reactionJSON{
|
||||
ID: reaction.ID,
|
||||
Creator: creator,
|
||||
ReactionType: reaction.ReactionType,
|
||||
CreateTime: reaction.CreatedTs,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (s *MCPService) loadReactionJSONByName(ctx context.Context, name string) (reactionJSON, error) {
|
||||
_, reactionID, err := apiv1.ExtractMemoReactionIDFromName(name)
|
||||
if err != nil {
|
||||
return reactionJSON{}, err
|
||||
}
|
||||
return s.loadReactionJSONByID(ctx, reactionID)
|
||||
}
|
||||
@ -0,0 +1,102 @@
|
||||
package mcp
|
||||
|
||||
import "github.com/mark3labs/mcp-go/mcp"
|
||||
|
||||
var mcpToolsByToolset = map[string]map[string]struct{}{
|
||||
"memos": stringSet(
|
||||
"list_memos",
|
||||
"get_memo",
|
||||
"create_memo",
|
||||
"update_memo",
|
||||
"delete_memo",
|
||||
"search_memos",
|
||||
"list_memo_comments",
|
||||
"create_memo_comment",
|
||||
),
|
||||
"tags": stringSet(
|
||||
"list_tags",
|
||||
),
|
||||
"attachments": stringSet(
|
||||
"list_attachments",
|
||||
"get_attachment",
|
||||
"delete_attachment",
|
||||
"link_attachment_to_memo",
|
||||
),
|
||||
"relations": stringSet(
|
||||
"list_memo_relations",
|
||||
"create_memo_relation",
|
||||
"delete_memo_relation",
|
||||
),
|
||||
"reactions": stringSet(
|
||||
"list_reactions",
|
||||
"upsert_reaction",
|
||||
"delete_reaction",
|
||||
),
|
||||
}
|
||||
|
||||
var allMCPToolNames = func() map[string]struct{} {
|
||||
names := map[string]struct{}{}
|
||||
for _, tools := range mcpToolsByToolset {
|
||||
for name := range tools {
|
||||
names[name] = struct{}{}
|
||||
}
|
||||
}
|
||||
return names
|
||||
}()
|
||||
|
||||
var mcpMutationTools = stringSet(
|
||||
"create_memo",
|
||||
"update_memo",
|
||||
"delete_memo",
|
||||
"create_memo_comment",
|
||||
"delete_attachment",
|
||||
"link_attachment_to_memo",
|
||||
"create_memo_relation",
|
||||
"delete_memo_relation",
|
||||
"upsert_reaction",
|
||||
"delete_reaction",
|
||||
)
|
||||
|
||||
type deletedJSON struct {
|
||||
Deleted bool `json:"deleted"`
|
||||
}
|
||||
|
||||
func stringSet(values ...string) map[string]struct{} {
|
||||
result := make(map[string]struct{}, len(values))
|
||||
for _, value := range values {
|
||||
result[value] = struct{}{}
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
func readOnlyToolOptions(title string, description string, opts ...mcp.ToolOption) []mcp.ToolOption {
|
||||
return annotatedToolOptions(title, description, true, false, true, false, opts...)
|
||||
}
|
||||
|
||||
func createToolOptions(title string, description string, idempotent bool, opts ...mcp.ToolOption) []mcp.ToolOption {
|
||||
return annotatedToolOptions(title, description, false, false, idempotent, false, opts...)
|
||||
}
|
||||
|
||||
func updateToolOptions(title string, description string, opts ...mcp.ToolOption) []mcp.ToolOption {
|
||||
return annotatedToolOptions(title, description, false, true, false, false, opts...)
|
||||
}
|
||||
|
||||
func annotatedToolOptions(title string, description string, readOnly bool, destructive bool, idempotent bool, openWorld bool, opts ...mcp.ToolOption) []mcp.ToolOption {
|
||||
base := []mcp.ToolOption{
|
||||
mcp.WithTitleAnnotation(title),
|
||||
mcp.WithDescription(description),
|
||||
mcp.WithReadOnlyHintAnnotation(readOnly),
|
||||
mcp.WithDestructiveHintAnnotation(destructive),
|
||||
mcp.WithIdempotentHintAnnotation(idempotent),
|
||||
mcp.WithOpenWorldHintAnnotation(openWorld),
|
||||
}
|
||||
return append(base, opts...)
|
||||
}
|
||||
|
||||
func newToolResultJSON(v any) (*mcp.CallToolResult, error) {
|
||||
return mcp.NewToolResultJSON(v)
|
||||
}
|
||||
|
||||
func newDeletedToolResult() (*mcp.CallToolResult, error) {
|
||||
return newToolResultJSON(deletedJSON{Deleted: true})
|
||||
}
|
||||
Loading…
Reference in New Issue