feat: implement Model Context Protocol (MCP) server and download manager capabilities

This commit is contained in:
ashish6298 committed 2026-09-22 16:16:01 +05:30
1 parent 88d150b933
commit e93437c95f
14 files changed
+976 -23

No files matched your search

+7
View File
@@ -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) {
+1
View File
@@ -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
}
+1
View File
@@ -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
+43
View File
@@ -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
View File
@@ -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)
+1
View File
@@ -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",
+64
View File
@@ -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")
}
}