mirror of
https://github.com/memohai/Memoh.git
synced 2026-04-27 07:16:19 +09:00
feat: remove extra tools and add exec tool
This commit is contained in:
Acbox
+47
-249
@@ -3,10 +3,8 @@ package mcp
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"fmt"
|
||||
"io/fs"
|
||||
"mime"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
@@ -17,14 +15,6 @@ import (
|
||||
sdkmcp "github.com/modelcontextprotocol/go-sdk/mcp"
|
||||
)
|
||||
|
||||
type EchoInput struct {
|
||||
Text string `json:"text" jsonschema:"text to echo"`
|
||||
}
|
||||
|
||||
type EchoOutput struct {
|
||||
Text string `json:"text" jsonschema:"echoed text"`
|
||||
}
|
||||
|
||||
type FSReadInput struct {
|
||||
Path string `json:"path" jsonschema:"relative file path"`
|
||||
}
|
||||
@@ -60,88 +50,33 @@ type FSListOutput struct {
|
||||
Entries []FSFileEntry `json:"entries" jsonschema:"entries"`
|
||||
}
|
||||
|
||||
type FSStatInput struct {
|
||||
Path string `json:"path" jsonschema:"relative path"`
|
||||
}
|
||||
|
||||
type FSStatOutput struct {
|
||||
Entry FSFileEntry `json:"entry" jsonschema:"entry"`
|
||||
}
|
||||
|
||||
type FSDeleteInput struct {
|
||||
Path string `json:"path" jsonschema:"relative path"`
|
||||
}
|
||||
|
||||
type FSDeleteOutput struct {
|
||||
OK bool `json:"ok" jsonschema:"delete result"`
|
||||
}
|
||||
|
||||
type FSApplyPatchInput struct {
|
||||
type FSEditInput struct {
|
||||
Path string `json:"path" jsonschema:"relative file path"`
|
||||
Patch string `json:"patch" jsonschema:"unified diff patch"`
|
||||
}
|
||||
|
||||
type FSApplyPatchOutput struct {
|
||||
type FSEditOutput struct {
|
||||
OK bool `json:"ok" jsonschema:"apply result"`
|
||||
}
|
||||
|
||||
type FSMkdirInput struct {
|
||||
Path string `json:"path" jsonschema:"relative directory path"`
|
||||
type ExecInput struct {
|
||||
Command string `json:"command" jsonschema:"command to run"`
|
||||
Args []string `json:"args" jsonschema:"command arguments"`
|
||||
}
|
||||
|
||||
type FSMkdirOutput struct {
|
||||
OK bool `json:"ok" jsonschema:"mkdir result"`
|
||||
}
|
||||
|
||||
type FSRenameInput struct {
|
||||
Source string `json:"source" jsonschema:"relative source path"`
|
||||
Destination string `json:"destination" jsonschema:"relative destination path"`
|
||||
}
|
||||
|
||||
type FSRenameOutput struct {
|
||||
OK bool `json:"ok" jsonschema:"rename result"`
|
||||
}
|
||||
|
||||
type FSReadBase64Input struct {
|
||||
Path string `json:"path" jsonschema:"relative file path"`
|
||||
}
|
||||
|
||||
type FSReadBase64Output struct {
|
||||
Data string `json:"data" jsonschema:"base64-encoded file bytes"`
|
||||
MimeType string `json:"mime_type" jsonschema:"detected mime type"`
|
||||
}
|
||||
|
||||
type GrepInput struct {
|
||||
Pattern string `json:"pattern" jsonschema:"grep pattern"`
|
||||
Args []string `json:"args" jsonschema:"grep options (flags only)"`
|
||||
}
|
||||
|
||||
type GrepOutput struct {
|
||||
Stdout string `json:"stdout" jsonschema:"grep standard output"`
|
||||
Stderr string `json:"stderr" jsonschema:"grep standard error"`
|
||||
ExitCode int `json:"exit_code" jsonschema:"grep exit code"`
|
||||
type ExecOutput struct {
|
||||
OK bool `json:"ok" jsonschema:"execution success"`
|
||||
ExitCode int `json:"exit_code" jsonschema:"process exit code"`
|
||||
Stdout string `json:"stdout" jsonschema:"standard output"`
|
||||
Stderr string `json:"stderr" jsonschema:"standard error"`
|
||||
}
|
||||
|
||||
func RegisterTools(server *sdkmcp.Server) {
|
||||
sdkmcp.AddTool(server, &sdkmcp.Tool{Name: "echo", Description: "echo input text"}, echoTool)
|
||||
sdkmcp.AddTool(server, &sdkmcp.Tool{Name: "fs.read", Description: "read file content"}, fsReadTool)
|
||||
sdkmcp.AddTool(server, &sdkmcp.Tool{Name: "fs.read_base64", Description: "read file bytes as base64"}, fsReadBase64Tool)
|
||||
sdkmcp.AddTool(server, &sdkmcp.Tool{Name: "fs.write", Description: "write file content"}, fsWriteTool)
|
||||
sdkmcp.AddTool(server, &sdkmcp.Tool{Name: "fs.list", Description: "list directory entries"}, fsListTool)
|
||||
sdkmcp.AddTool(server, &sdkmcp.Tool{Name: "fs.stat", Description: "stat file or directory"}, fsStatTool)
|
||||
sdkmcp.AddTool(server, &sdkmcp.Tool{Name: "fs.delete", Description: "delete file or directory"}, fsDeleteTool)
|
||||
sdkmcp.AddTool(server, &sdkmcp.Tool{Name: "fs.apply_patch", Description: "apply unified diff patch"}, fsApplyPatchTool)
|
||||
sdkmcp.AddTool(server, &sdkmcp.Tool{Name: "fs.mkdir", Description: "create directory (mkdir -p)"}, fsMkdirTool)
|
||||
sdkmcp.AddTool(server, &sdkmcp.Tool{Name: "fs.rename", Description: "rename/move file or directory"}, fsRenameTool)
|
||||
sdkmcp.AddTool(server, &sdkmcp.Tool{Name: "grep", Description: "grep within /data using GNU grep"}, grepTool)
|
||||
}
|
||||
|
||||
func echoTool(ctx context.Context, req *sdkmcp.CallToolRequest, input EchoInput) (
|
||||
*sdkmcp.CallToolResult,
|
||||
EchoOutput,
|
||||
error,
|
||||
) {
|
||||
return nil, EchoOutput{Text: input.Text}, nil
|
||||
sdkmcp.AddTool(server, &sdkmcp.Tool{Name: "read", Description: "read file content"}, fsReadTool)
|
||||
sdkmcp.AddTool(server, &sdkmcp.Tool{Name: "write", Description: "write file content"}, fsWriteTool)
|
||||
sdkmcp.AddTool(server, &sdkmcp.Tool{Name: "list", Description: "list directory entries"}, fsListTool)
|
||||
sdkmcp.AddTool(server, &sdkmcp.Tool{Name: "edit", Description: "apply unified diff patch"}, fsEditTool)
|
||||
sdkmcp.AddTool(server, &sdkmcp.Tool{Name: "exec", Description: "execute command"}, execTool)
|
||||
}
|
||||
|
||||
func fsReadTool(ctx context.Context, req *sdkmcp.CallToolRequest, input FSReadInput) (
|
||||
@@ -161,97 +96,6 @@ func fsReadTool(ctx context.Context, req *sdkmcp.CallToolRequest, input FSReadIn
|
||||
return nil, FSReadOutput{Content: string(data)}, nil
|
||||
}
|
||||
|
||||
func fsReadBase64Tool(ctx context.Context, req *sdkmcp.CallToolRequest, input FSReadBase64Input) (
|
||||
*sdkmcp.CallToolResult,
|
||||
FSReadBase64Output,
|
||||
error,
|
||||
) {
|
||||
root := dataRoot()
|
||||
target, err := resolvePath(root, input.Path)
|
||||
if err != nil {
|
||||
return nil, FSReadBase64Output{}, err
|
||||
}
|
||||
data, err := os.ReadFile(target)
|
||||
if err != nil {
|
||||
return nil, FSReadBase64Output{}, err
|
||||
}
|
||||
ext := strings.ToLower(filepath.Ext(target))
|
||||
mimeType := mime.TypeByExtension(ext)
|
||||
if strings.TrimSpace(mimeType) == "" {
|
||||
// Fallback mapping for common image/audio extensions.
|
||||
switch ext {
|
||||
case ".png":
|
||||
mimeType = "image/png"
|
||||
case ".jpg", ".jpeg":
|
||||
mimeType = "image/jpeg"
|
||||
case ".gif":
|
||||
mimeType = "image/gif"
|
||||
case ".webp":
|
||||
mimeType = "image/webp"
|
||||
case ".bmp":
|
||||
mimeType = "image/bmp"
|
||||
case ".svg":
|
||||
mimeType = "image/svg+xml"
|
||||
case ".mp3":
|
||||
mimeType = "audio/mpeg"
|
||||
case ".wav":
|
||||
mimeType = "audio/wav"
|
||||
case ".ogg":
|
||||
mimeType = "audio/ogg"
|
||||
case ".flac":
|
||||
mimeType = "audio/flac"
|
||||
default:
|
||||
mimeType = "application/octet-stream"
|
||||
}
|
||||
}
|
||||
return nil, FSReadBase64Output{
|
||||
Data: base64.StdEncoding.EncodeToString(data),
|
||||
MimeType: mimeType,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func grepTool(ctx context.Context, req *sdkmcp.CallToolRequest, input GrepInput) (
|
||||
*sdkmcp.CallToolResult,
|
||||
GrepOutput,
|
||||
error,
|
||||
) {
|
||||
if strings.TrimSpace(input.Pattern) == "" {
|
||||
return nil, GrepOutput{}, fmt.Errorf("pattern is required")
|
||||
}
|
||||
if stat, err := os.Stat("/data"); err != nil || !stat.IsDir() {
|
||||
return nil, GrepOutput{}, fmt.Errorf("/data is not available")
|
||||
}
|
||||
|
||||
args := append([]string{}, input.Args...)
|
||||
args = append(args, input.Pattern, ".")
|
||||
|
||||
cmd := exec.CommandContext(ctx, "grep", args...)
|
||||
cmd.Dir = "/data"
|
||||
|
||||
var stdout bytes.Buffer
|
||||
var stderr bytes.Buffer
|
||||
cmd.Stdout = &stdout
|
||||
cmd.Stderr = &stderr
|
||||
|
||||
exitCode := 0
|
||||
if err := cmd.Run(); err != nil {
|
||||
if exitErr, ok := err.(*exec.ExitError); ok {
|
||||
exitCode = exitErr.ExitCode()
|
||||
if exitCode != 1 {
|
||||
return nil, GrepOutput{}, fmt.Errorf("grep failed: %s", strings.TrimSpace(stderr.String()))
|
||||
}
|
||||
} else {
|
||||
return nil, GrepOutput{}, err
|
||||
}
|
||||
}
|
||||
|
||||
return nil, GrepOutput{
|
||||
Stdout: stdout.String(),
|
||||
Stderr: stderr.String(),
|
||||
ExitCode: exitCode,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func fsWriteTool(ctx context.Context, req *sdkmcp.CallToolRequest, input FSWriteInput) (
|
||||
*sdkmcp.CallToolResult,
|
||||
FSWriteOutput,
|
||||
@@ -338,112 +182,66 @@ func fsListTool(ctx context.Context, req *sdkmcp.CallToolRequest, input FSListIn
|
||||
return nil, FSListOutput{Path: listedPath, Entries: entries}, nil
|
||||
}
|
||||
|
||||
func fsStatTool(ctx context.Context, req *sdkmcp.CallToolRequest, input FSStatInput) (
|
||||
func fsEditTool(ctx context.Context, req *sdkmcp.CallToolRequest, input FSEditInput) (
|
||||
*sdkmcp.CallToolResult,
|
||||
FSStatOutput,
|
||||
error,
|
||||
) {
|
||||
root := dataRoot()
|
||||
target, err := resolvePathAllowRoot(root, input.Path)
|
||||
if err != nil {
|
||||
return nil, FSStatOutput{}, err
|
||||
}
|
||||
info, err := os.Stat(target)
|
||||
if err != nil {
|
||||
return nil, FSStatOutput{}, err
|
||||
}
|
||||
entry, err := entryForPath(root, target, info)
|
||||
if err != nil {
|
||||
return nil, FSStatOutput{}, err
|
||||
}
|
||||
return nil, FSStatOutput{Entry: entry}, nil
|
||||
}
|
||||
|
||||
func fsDeleteTool(ctx context.Context, req *sdkmcp.CallToolRequest, input FSDeleteInput) (
|
||||
*sdkmcp.CallToolResult,
|
||||
FSDeleteOutput,
|
||||
FSEditOutput,
|
||||
error,
|
||||
) {
|
||||
root := dataRoot()
|
||||
target, err := resolvePath(root, input.Path)
|
||||
if err != nil {
|
||||
return nil, FSDeleteOutput{}, err
|
||||
}
|
||||
if err := os.RemoveAll(target); err != nil {
|
||||
return nil, FSDeleteOutput{}, err
|
||||
}
|
||||
return nil, FSDeleteOutput{OK: true}, nil
|
||||
}
|
||||
|
||||
func fsApplyPatchTool(ctx context.Context, req *sdkmcp.CallToolRequest, input FSApplyPatchInput) (
|
||||
*sdkmcp.CallToolResult,
|
||||
FSApplyPatchOutput,
|
||||
error,
|
||||
) {
|
||||
root := dataRoot()
|
||||
target, err := resolvePath(root, input.Path)
|
||||
if err != nil {
|
||||
return nil, FSApplyPatchOutput{}, err
|
||||
return nil, FSEditOutput{}, err
|
||||
}
|
||||
orig, err := os.ReadFile(target)
|
||||
if err != nil {
|
||||
return nil, FSApplyPatchOutput{}, err
|
||||
return nil, FSEditOutput{}, err
|
||||
}
|
||||
updated, err := applyUnifiedPatch(string(orig), input.Patch)
|
||||
if err != nil {
|
||||
return nil, FSApplyPatchOutput{}, err
|
||||
return nil, FSEditOutput{}, err
|
||||
}
|
||||
info, err := os.Stat(target)
|
||||
if err != nil {
|
||||
return nil, FSApplyPatchOutput{}, err
|
||||
return nil, FSEditOutput{}, err
|
||||
}
|
||||
if err := os.WriteFile(target, []byte(updated), info.Mode().Perm()); err != nil {
|
||||
return nil, FSApplyPatchOutput{}, err
|
||||
return nil, FSEditOutput{}, err
|
||||
}
|
||||
return nil, FSApplyPatchOutput{OK: true}, nil
|
||||
return nil, FSEditOutput{OK: true}, nil
|
||||
}
|
||||
|
||||
func fsMkdirTool(ctx context.Context, req *sdkmcp.CallToolRequest, input FSMkdirInput) (
|
||||
func execTool(ctx context.Context, req *sdkmcp.CallToolRequest, input ExecInput) (
|
||||
*sdkmcp.CallToolResult,
|
||||
FSMkdirOutput,
|
||||
ExecOutput,
|
||||
error,
|
||||
) {
|
||||
root := dataRoot()
|
||||
target, err := resolvePath(root, input.Path)
|
||||
if err != nil {
|
||||
return nil, FSMkdirOutput{}, err
|
||||
if strings.TrimSpace(input.Command) == "" {
|
||||
return nil, ExecOutput{}, fmt.Errorf("command is required")
|
||||
}
|
||||
if err := os.MkdirAll(target, 0o755); err != nil {
|
||||
return nil, FSMkdirOutput{}, err
|
||||
}
|
||||
return nil, FSMkdirOutput{OK: true}, nil
|
||||
}
|
||||
cmd := exec.CommandContext(ctx, input.Command, input.Args...)
|
||||
var stdout bytes.Buffer
|
||||
var stderr bytes.Buffer
|
||||
cmd.Stdout = &stdout
|
||||
cmd.Stderr = &stderr
|
||||
|
||||
func fsRenameTool(ctx context.Context, req *sdkmcp.CallToolRequest, input FSRenameInput) (
|
||||
*sdkmcp.CallToolResult,
|
||||
FSRenameOutput,
|
||||
error,
|
||||
) {
|
||||
root := dataRoot()
|
||||
source, err := resolvePath(root, input.Source)
|
||||
if err != nil {
|
||||
return nil, FSRenameOutput{}, err
|
||||
if err := cmd.Run(); err != nil {
|
||||
if exitErr, ok := err.(*exec.ExitError); ok {
|
||||
return nil, ExecOutput{
|
||||
OK: false,
|
||||
ExitCode: exitErr.ExitCode(),
|
||||
Stdout: stdout.String(),
|
||||
Stderr: stderr.String(),
|
||||
}, nil
|
||||
}
|
||||
destination, err := resolvePath(root, input.Destination)
|
||||
if err != nil {
|
||||
return nil, FSRenameOutput{}, err
|
||||
return nil, ExecOutput{}, err
|
||||
}
|
||||
|
||||
if _, err := os.Lstat(destination); err == nil {
|
||||
return nil, FSRenameOutput{}, fmt.Errorf("destination already exists")
|
||||
} else if !os.IsNotExist(err) {
|
||||
return nil, FSRenameOutput{}, err
|
||||
}
|
||||
|
||||
if err := os.Rename(source, destination); err != nil {
|
||||
return nil, FSRenameOutput{}, err
|
||||
}
|
||||
return nil, FSRenameOutput{OK: true}, nil
|
||||
return nil, ExecOutput{
|
||||
OK: true,
|
||||
ExitCode: 0,
|
||||
Stdout: stdout.String(),
|
||||
Stderr: stderr.String(),
|
||||
}, nil
|
||||
}
|
||||
|
||||
func dataRoot() string {
|
||||
|
||||
Reference in New Issue
Block a user