mirror of
https://github.com/GopeedLab/gopeed.git
synced 2026-11-08 17:40:28 +00:00
feat(protocol/http, mcpserver, download): add checksum verification, MCP speed tool, pause webhook, and CLI flags (#1542)
This commit is contained in:
1 parent
fa622a4063
commit
6f20790b17
41 files changed
+1852
-25
No files matched your search
+67
-7
@@ -12,6 +12,7 @@ import (
|
||||
|
||||
"github.com/GopeedLab/gopeed/pkg/base"
|
||||
"github.com/GopeedLab/gopeed/pkg/download"
|
||||
phttp "github.com/GopeedLab/gopeed/pkg/protocol/http"
|
||||
"github.com/google/jsonschema-go/jsonschema"
|
||||
"github.com/modelcontextprotocol/go-sdk/mcp"
|
||||
)
|
||||
@@ -65,12 +66,13 @@ type btRequestExtra struct {
|
||||
}
|
||||
|
||||
type httpTaskExtra struct {
|
||||
Connections int `json:"connections,omitempty" jsonschema:"Number of HTTP download connections"`
|
||||
AutoTorrent *bool `json:"autoTorrent,omitempty" jsonschema:"Automatically create a BitTorrent task after downloading a torrent file"`
|
||||
DeleteTorrentAfterDownload *bool `json:"deleteTorrentAfterDownload,omitempty" jsonschema:"Delete the torrent file after automatically creating its BitTorrent task"`
|
||||
AutoExtract *bool `json:"autoExtract,omitempty" jsonschema:"Automatically extract the downloaded archive"`
|
||||
ArchivePassword string `json:"archivePassword,omitempty" jsonschema:"Password for an encrypted archive"`
|
||||
DeleteAfterExtract bool `json:"deleteAfterExtract,omitempty" jsonschema:"Delete archive files after successful extraction"`
|
||||
Connections int `json:"connections,omitempty" jsonschema:"Number of HTTP download connections"`
|
||||
AutoTorrent *bool `json:"autoTorrent,omitempty" jsonschema:"Automatically create a BitTorrent task after downloading a torrent file"`
|
||||
DeleteTorrentAfterDownload *bool `json:"deleteTorrentAfterDownload,omitempty" jsonschema:"Delete the torrent file after automatically creating its BitTorrent task"`
|
||||
AutoExtract *bool `json:"autoExtract,omitempty" jsonschema:"Automatically extract the downloaded archive"`
|
||||
ArchivePassword string `json:"archivePassword,omitempty" jsonschema:"Password for an encrypted archive"`
|
||||
DeleteAfterExtract bool `json:"deleteAfterExtract,omitempty" jsonschema:"Delete archive files after successful extraction"`
|
||||
Checksum *checksumOption `json:"checksum,omitempty" jsonschema:"Optional checksum verification settings"`
|
||||
}
|
||||
|
||||
// downloadRequest is the MCP-facing form of base.Request. The REST model uses
|
||||
@@ -129,6 +131,11 @@ 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 {
|
||||
@@ -186,12 +193,31 @@ func (o *downloadOptions) baseOptions() *base.Options {
|
||||
if o == nil {
|
||||
return nil
|
||||
}
|
||||
var extra any
|
||||
if o.Extra != nil {
|
||||
var checksum *phttp.ChecksumOption
|
||||
if o.Extra.Checksum != nil {
|
||||
checksum = &phttp.ChecksumOption{
|
||||
Algorithm: o.Extra.Checksum.Algorithm,
|
||||
Expected: o.Extra.Checksum.Expected,
|
||||
}
|
||||
}
|
||||
extra = &phttp.OptsExtra{
|
||||
Connections: o.Extra.Connections,
|
||||
AutoTorrent: o.Extra.AutoTorrent,
|
||||
DeleteTorrentAfterDownload: o.Extra.DeleteTorrentAfterDownload,
|
||||
AutoExtract: o.Extra.AutoExtract,
|
||||
ArchivePassword: o.Extra.ArchivePassword,
|
||||
DeleteAfterExtract: o.Extra.DeleteAfterExtract,
|
||||
Checksum: checksum,
|
||||
}
|
||||
}
|
||||
return &base.Options{
|
||||
Name: o.Name,
|
||||
Path: o.Path,
|
||||
AsDefaultPath: o.AsDefaultPath,
|
||||
SelectFiles: o.SelectFiles,
|
||||
Extra: o.Extra,
|
||||
Extra: extra,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -270,6 +296,13 @@ type getTaskStatsOutput struct {
|
||||
Stats any `json:"stats"`
|
||||
}
|
||||
|
||||
type getDownloadSpeedInput struct{}
|
||||
|
||||
type getDownloadSpeedOutput struct {
|
||||
TotalSpeed int64 `json:"totalSpeed"`
|
||||
ActiveTaskCount int `json:"activeTaskCount"`
|
||||
}
|
||||
|
||||
type taskActionOutput struct {
|
||||
ID string `json:"id"`
|
||||
Success bool `json:"success"`
|
||||
@@ -416,6 +449,33 @@ 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: "Returns the combined download speed across all currently active tasks and how many tasks are actively downloading. Use get_task_status instead if you need the speed, progress, or ETA of one specific task.",
|
||||
Annotations: &mcp.ToolAnnotations{ReadOnlyHint: true, OpenWorldHint: &closedWorld},
|
||||
}, func(_ context.Context, _ *mcp.CallToolRequest, _ *getDownloadSpeedInput) (*mcp.CallToolResult, *getDownloadSpeedOutput, error) {
|
||||
allTasks := downloader.GetTasks()
|
||||
var totalSpeed int64
|
||||
activeCount := 0
|
||||
|
||||
for _, task := range allTasks {
|
||||
status, err := downloader.RuntimeStatus(task.ID)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
if status.Status == base.DownloadStatusRunning {
|
||||
activeCount++
|
||||
totalSpeed += status.Speed
|
||||
}
|
||||
}
|
||||
|
||||
return nil, &getDownloadSpeedOutput{
|
||||
TotalSpeed: totalSpeed,
|
||||
ActiveTaskCount: activeCount,
|
||||
}, 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,103 @@
|
||||
package mcpserver
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/GopeedLab/gopeed/pkg/base"
|
||||
"github.com/GopeedLab/gopeed/pkg/download"
|
||||
"github.com/modelcontextprotocol/go-sdk/mcp"
|
||||
)
|
||||
|
||||
func TestGetDownloadSpeedTool(t *testing.T) {
|
||||
// Serve a test payload that takes a little bit of time or streams bytes
|
||||
testServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Length", "1048576")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
chunk := make([]byte, 8192)
|
||||
for i := 0; i < 128; i++ {
|
||||
_, _ = w.Write(chunk)
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
}
|
||||
}))
|
||||
defer testServer.Close()
|
||||
|
||||
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. Zero tasks test (returns 0 totalSpeed and 0 activeTaskCount)
|
||||
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"])
|
||||
}
|
||||
if totalSpeed, ok := structured["totalSpeed"].(float64); !ok || totalSpeed != 0 {
|
||||
t.Fatalf("totalSpeed = %v, want 0", structured["totalSpeed"])
|
||||
}
|
||||
|
||||
// 2. Start a real download task
|
||||
taskID, err := downloader.CreateDirect(&base.Request{
|
||||
URL: testServer.URL,
|
||||
}, &base.Options{
|
||||
Path: t.TempDir(),
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer func() {
|
||||
_ = downloader.Delete(&download.TaskFilter{IDs: []string{taskID}}, true)
|
||||
}()
|
||||
|
||||
// Query get_download_speed while downloading
|
||||
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 _, ok := structured["activeTaskCount"].(float64); !ok {
|
||||
t.Fatalf("activeTaskCount missing or wrong type: %v", structured["activeTaskCount"])
|
||||
}
|
||||
if _, ok := structured["totalSpeed"].(float64); !ok {
|
||||
t.Fatalf("totalSpeed missing or wrong type: %v", structured["totalSpeed"])
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user