mirror of
https://github.com/GopeedLab/gopeed.git
synced 2026-11-07 17:10:29 +00:00
feat: implement Model Context Protocol (MCP) server and download manager capabilities
This commit is contained in:
1 parent
88d150b933
commit
e93437c95f
14 files changed
+976
-23
No files matched your search
@@ -124,6 +124,11 @@ type FileInfo struct {
|
||||
Req *Request `json:"req"`
|
||||
}
|
||||
|
||||
type ChecksumOption struct {
|
||||
Algorithm string `json:"algorithm"` // "md5" | "sha1" | "sha256"
|
||||
Expected string `json:"expected"` // hex-encoded expected hash
|
||||
}
|
||||
|
||||
// Options for download
|
||||
type Options struct {
|
||||
// Download file name
|
||||
@@ -136,6 +141,8 @@ type Options struct {
|
||||
SelectFiles []int `json:"selectFiles"`
|
||||
// Extra info for specific fetcher
|
||||
Extra any `json:"extra"`
|
||||
// Checksum verification option
|
||||
Checksum *ChecksumOption `json:"checksum,omitempty"`
|
||||
}
|
||||
|
||||
func (o *Options) InitSelectFiles(fileSize int) {
|
||||
|
||||
@@ -1912,6 +1912,7 @@ func (d *Downloader) doPause(task *Task) (err error) {
|
||||
}
|
||||
if handled {
|
||||
d.emit(EventKeyPause, task)
|
||||
d.triggerWebhooks(WebhookEventDownloadPause, task, nil)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -21,6 +21,7 @@ type WebhookEvent string
|
||||
const (
|
||||
WebhookEventDownloadDone WebhookEvent = "DOWNLOAD_DONE"
|
||||
WebhookEventDownloadError WebhookEvent = "DOWNLOAD_ERROR"
|
||||
WebhookEventDownloadPause WebhookEvent = "DOWNLOAD_PAUSE"
|
||||
)
|
||||
|
||||
// WebhookData is the data sent to webhook URLs
|
||||
|
||||
@@ -814,6 +814,49 @@ func TestWebhook_TestWebhookUrl_VerifyTestPayload(t *testing.T) {
|
||||
})
|
||||
}
|
||||
|
||||
func TestWebhook_TriggerOnPause(t *testing.T) {
|
||||
receivedData := make(chan *WebhookData, 1)
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
var data WebhookData
|
||||
if err := json.NewDecoder(r.Body).Decode(&data); err != nil {
|
||||
t.Errorf("Failed to decode data: %v", err)
|
||||
return
|
||||
}
|
||||
receivedData <- &data
|
||||
w.WriteHeader(http.StatusOK)
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
setupWebhookTest(t, func(downloader *Downloader) {
|
||||
cfg, _ := downloader.GetConfig()
|
||||
cfg.Webhook = &base.WebhookConfig{
|
||||
Enable: true,
|
||||
URLs: []string{server.URL},
|
||||
}
|
||||
downloader.PutConfig(cfg)
|
||||
|
||||
task := NewTask()
|
||||
task.Protocol = "http"
|
||||
task.Meta = &mockFetcherMeta
|
||||
|
||||
downloader.triggerWebhooks(WebhookEventDownloadPause, task, nil)
|
||||
|
||||
select {
|
||||
case data := <-receivedData:
|
||||
if data.Event != WebhookEventDownloadPause {
|
||||
t.Errorf("Expected event 'DOWNLOAD_PAUSE', got '%s'", data.Event)
|
||||
}
|
||||
if data.Payload == nil || data.Payload.Task == nil {
|
||||
t.Error("Expected payload.task to be present")
|
||||
} else if data.Payload.Task.ID != task.ID {
|
||||
t.Errorf("Expected task ID '%s', got '%s'", task.ID, data.Payload.Task.ID)
|
||||
}
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Error("Timeout waiting for webhook")
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func setupWebhookTest(t *testing.T, fn func(downloader *Downloader)) {
|
||||
defaultDownloader.Setup()
|
||||
defaultDownloader.cfg.StorageDir = ".test_storage"
|
||||
|
||||
+138
-5
@@ -129,14 +129,20 @@ func isBitTorrentRequest(rawURL string) bool {
|
||||
strings.HasPrefix(upperURL, "DATA:APPLICATION/X-BITTORRENT;BASE64,")
|
||||
}
|
||||
|
||||
type checksumOption struct {
|
||||
Algorithm string `json:"algorithm" jsonschema:"Checksum algorithm: md5, sha1, or sha256"`
|
||||
Expected string `json:"expected" jsonschema:"Expected hex-encoded checksum hash"`
|
||||
}
|
||||
|
||||
// downloadOptions is the MCP-facing form of base.Options. Only HTTP currently
|
||||
// defines protocol-specific task options, so Extra has one concrete schema.
|
||||
type downloadOptions struct {
|
||||
Name string `json:"name,omitempty" jsonschema:"Optional output file or directory name"`
|
||||
Path string `json:"path,omitempty" jsonschema:"Optional download directory"`
|
||||
AsDefaultPath bool `json:"asDefaultPath,omitempty" jsonschema:"Save path as the default download directory after task creation"`
|
||||
SelectFiles []int `json:"selectFiles,omitempty" jsonschema:"Optional zero-based file indexes to download from a multi-file resource"`
|
||||
Extra *httpTaskExtra `json:"extra,omitempty" jsonschema:"Optional HTTP task settings"`
|
||||
Name string `json:"name,omitempty" jsonschema:"Optional output file or directory name"`
|
||||
Path string `json:"path,omitempty" jsonschema:"Optional download directory"`
|
||||
AsDefaultPath bool `json:"asDefaultPath,omitempty" jsonschema:"Save path as the default download directory after task creation"`
|
||||
SelectFiles []int `json:"selectFiles,omitempty" jsonschema:"Optional zero-based file indexes to download from a multi-file resource"`
|
||||
Extra *httpTaskExtra `json:"extra,omitempty" jsonschema:"Optional HTTP task settings"`
|
||||
Checksum *checksumOption `json:"checksum,omitempty" jsonschema:"Optional checksum verification settings"`
|
||||
}
|
||||
|
||||
func requestExtraSchema() *jsonschema.Schema {
|
||||
@@ -186,12 +192,20 @@ func (o *downloadOptions) baseOptions() *base.Options {
|
||||
if o == nil {
|
||||
return nil
|
||||
}
|
||||
var checksum *base.ChecksumOption
|
||||
if o.Checksum != nil {
|
||||
checksum = &base.ChecksumOption{
|
||||
Algorithm: o.Checksum.Algorithm,
|
||||
Expected: o.Checksum.Expected,
|
||||
}
|
||||
}
|
||||
return &base.Options{
|
||||
Name: o.Name,
|
||||
Path: o.Path,
|
||||
AsDefaultPath: o.AsDefaultPath,
|
||||
SelectFiles: o.SelectFiles,
|
||||
Extra: o.Extra,
|
||||
Checksum: checksum,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -270,6 +284,29 @@ type getTaskStatsOutput struct {
|
||||
Stats any `json:"stats"`
|
||||
}
|
||||
|
||||
type getDownloadSpeedInput struct {
|
||||
ID string `json:"id,omitempty" jsonschema:"Optional Gopeed task ID. If omitted, returns global download and upload speed across all active tasks."`
|
||||
}
|
||||
|
||||
type TaskSpeedInfo struct {
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name,omitempty"`
|
||||
Status base.Status `json:"status"`
|
||||
Speed int64 `json:"speed"` // bytes/s
|
||||
UploadSpeed int64 `json:"uploadSpeed"` // bytes/s
|
||||
Downloaded int64 `json:"downloaded"`
|
||||
Total int64 `json:"total"`
|
||||
Progress float64 `json:"progress"` // percentage 0.0 - 100.0
|
||||
ETA *int64 `json:"eta,omitempty"` // estimated seconds remaining, if running and speed > 0
|
||||
}
|
||||
|
||||
type getDownloadSpeedOutput struct {
|
||||
TotalSpeed int64 `json:"totalSpeed"` // total download speed in bytes/s
|
||||
TotalUploadSpeed int64 `json:"totalUploadSpeed"` // total upload speed in bytes/s
|
||||
ActiveTaskCount int `json:"activeTaskCount"`
|
||||
Tasks []TaskSpeedInfo `json:"tasks,omitempty"`
|
||||
}
|
||||
|
||||
type taskActionOutput struct {
|
||||
ID string `json:"id"`
|
||||
Success bool `json:"success"`
|
||||
@@ -416,6 +453,102 @@ func registerTools(server *mcp.Server, downloader *download.Downloader) {
|
||||
return nil, &getTaskStatsOutput{Stats: stats}, nil
|
||||
})
|
||||
|
||||
mcp.AddTool(server, &mcp.Tool{
|
||||
Name: "get_download_speed",
|
||||
Title: "Get download speed",
|
||||
Description: "Get instantaneous download/upload speed and progress/ETA for a specific task or across all active tasks.",
|
||||
Annotations: &mcp.ToolAnnotations{ReadOnlyHint: true, OpenWorldHint: &closedWorld},
|
||||
}, func(_ context.Context, _ *mcp.CallToolRequest, input *getDownloadSpeedInput) (*mcp.CallToolResult, *getDownloadSpeedOutput, error) {
|
||||
if input != nil && strings.TrimSpace(input.ID) != "" {
|
||||
task, err := requireTask(downloader, input.ID)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
status, err := downloader.RuntimeStatus(task.ID)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
var progress float64
|
||||
if status.Total > 0 {
|
||||
progress = float64(status.Downloaded) / float64(status.Total) * 100.0
|
||||
}
|
||||
var eta *int64
|
||||
if status.Status == base.DownloadStatusRunning && status.Speed > 0 && status.Total > status.Downloaded {
|
||||
remainingBytes := status.Total - status.Downloaded
|
||||
etaSeconds := remainingBytes / status.Speed
|
||||
eta = &etaSeconds
|
||||
}
|
||||
taskInfo := TaskSpeedInfo{
|
||||
ID: task.ID,
|
||||
Name: task.Name(),
|
||||
Status: status.Status,
|
||||
Speed: status.Speed,
|
||||
UploadSpeed: status.UploadSpeed,
|
||||
Downloaded: status.Downloaded,
|
||||
Total: status.Total,
|
||||
Progress: progress,
|
||||
ETA: eta,
|
||||
}
|
||||
activeCount := 0
|
||||
if status.Status == base.DownloadStatusRunning {
|
||||
activeCount = 1
|
||||
}
|
||||
return nil, &getDownloadSpeedOutput{
|
||||
TotalSpeed: status.Speed,
|
||||
TotalUploadSpeed: status.UploadSpeed,
|
||||
ActiveTaskCount: activeCount,
|
||||
Tasks: []TaskSpeedInfo{taskInfo},
|
||||
}, nil
|
||||
}
|
||||
|
||||
// All tasks query
|
||||
allTasks := downloader.GetTasks()
|
||||
var totalSpeed int64
|
||||
var totalUploadSpeed int64
|
||||
activeCount := 0
|
||||
tasksInfo := make([]TaskSpeedInfo, 0)
|
||||
|
||||
for _, task := range allTasks {
|
||||
status, err := downloader.RuntimeStatus(task.ID)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
if status.Status == base.DownloadStatusRunning {
|
||||
activeCount++
|
||||
totalSpeed += status.Speed
|
||||
totalUploadSpeed += status.UploadSpeed
|
||||
}
|
||||
var progress float64
|
||||
if status.Total > 0 {
|
||||
progress = float64(status.Downloaded) / float64(status.Total) * 100.0
|
||||
}
|
||||
var eta *int64
|
||||
if status.Status == base.DownloadStatusRunning && status.Speed > 0 && status.Total > status.Downloaded {
|
||||
remainingBytes := status.Total - status.Downloaded
|
||||
etaSeconds := remainingBytes / status.Speed
|
||||
eta = &etaSeconds
|
||||
}
|
||||
tasksInfo = append(tasksInfo, TaskSpeedInfo{
|
||||
ID: task.ID,
|
||||
Name: task.Name(),
|
||||
Status: status.Status,
|
||||
Speed: status.Speed,
|
||||
UploadSpeed: status.UploadSpeed,
|
||||
Downloaded: status.Downloaded,
|
||||
Total: status.Total,
|
||||
Progress: progress,
|
||||
ETA: eta,
|
||||
})
|
||||
}
|
||||
|
||||
return nil, &getDownloadSpeedOutput{
|
||||
TotalSpeed: totalSpeed,
|
||||
TotalUploadSpeed: totalUploadSpeed,
|
||||
ActiveTaskCount: activeCount,
|
||||
Tasks: tasksInfo,
|
||||
}, nil
|
||||
})
|
||||
|
||||
registerTaskAction(server, downloader, "pause_task", "Pause task", "Pause one Gopeed task.", &nondestructive, &closedWorld, downloader.Pause)
|
||||
registerTaskAction(server, downloader, "continue_task", "Continue task", "Continue a paused or failed Gopeed task.", &nondestructive, &openWorld, downloader.Continue)
|
||||
|
||||
|
||||
@@ -64,6 +64,7 @@ func TestStreamableHTTPTools(t *testing.T) {
|
||||
"continue_task",
|
||||
"create_task",
|
||||
"delete_task",
|
||||
"get_download_speed",
|
||||
"get_task",
|
||||
"get_task_stats",
|
||||
"get_task_status",
|
||||
|
||||
@@ -0,0 +1,64 @@
|
||||
package mcpserver
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/GopeedLab/gopeed/pkg/download"
|
||||
"github.com/modelcontextprotocol/go-sdk/mcp"
|
||||
)
|
||||
|
||||
func TestGetDownloadSpeedTool(t *testing.T) {
|
||||
downloader := download.NewDownloader(&download.DownloaderConfig{
|
||||
Storage: download.NewMemStorage(),
|
||||
StorageDir: t.TempDir(),
|
||||
})
|
||||
if err := downloader.Setup(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer downloader.Close()
|
||||
|
||||
httpServer := httptest.NewServer(NewHandler(downloader))
|
||||
defer httpServer.Close()
|
||||
|
||||
client := mcp.NewClient(&mcp.Implementation{Name: "gopeed-test", Version: "1.0.0"}, nil)
|
||||
session, err := client.Connect(context.Background(), &mcp.StreamableClientTransport{Endpoint: httpServer.URL}, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer session.Close()
|
||||
|
||||
// 1. Query all tasks speed (empty list)
|
||||
result, err := session.CallTool(context.Background(), &mcp.CallToolParams{
|
||||
Name: "get_download_speed",
|
||||
Arguments: map[string]any{},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if result.IsError {
|
||||
t.Fatalf("get_download_speed returned error: %+v", result.Content)
|
||||
}
|
||||
structured, ok := result.StructuredContent.(map[string]any)
|
||||
if !ok {
|
||||
t.Fatalf("structured content type = %T", result.StructuredContent)
|
||||
}
|
||||
if activeCount, ok := structured["activeTaskCount"].(float64); !ok || activeCount != 0 {
|
||||
t.Fatalf("activeTaskCount = %v, want 0", structured["activeTaskCount"])
|
||||
}
|
||||
|
||||
// 2. Query non-existent task ID -> should return error
|
||||
notFoundResult, err := session.CallTool(context.Background(), &mcp.CallToolParams{
|
||||
Name: "get_download_speed",
|
||||
Arguments: map[string]any{
|
||||
"id": "non-existent-task-id",
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !notFoundResult.IsError {
|
||||
t.Fatal("get_download_speed should return error for non-existent task ID")
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user