feat: basic aria2 integration

This commit is contained in:
krau
2026-01-17 14:57:03 +08:00
parent f17a380579
commit b05d86509c
16 changed files with 809 additions and 15 deletions

View File

@@ -90,6 +90,15 @@ func handleAddCallback(ctx *ext.Context, update *ext.Update) error {
shortcut.CreateAndAddParsedTaskWithEdit(ctx, selectedStorage, dirPath, data.ParsedItem, msgID, userID)
case tasktype.TaskTypeDirectlinks:
shortcut.CreateAndAddDirectTaskWithEdit(ctx, selectedStorage, dirPath, data.DirectLinks, msgID, userID)
case tasktype.TaskTypeAria2:
client := GetAria2Client()
if client == nil {
ctx.AnswerCallback(msgelem.AlertCallbackAnswer(queryID, i18n.T(i18nk.BotMsgAria2ErrorAria2ClientInitFailed, map[string]any{
"Error": "aria2 client not initialized",
})))
return dispatcher.EndGroups
}
shortcut.CreateAndAddAria2TaskWithEdit(ctx, selectedStorage, dirPath, data.Aria2URIs, client, msgID, userID)
default:
return fmt.Errorf("unexcept task type: %s", data.TaskType)
}

View File

@@ -58,6 +58,11 @@ var aria2ClientInitOnce sync.Once
var aria2ClientInitErr error
var aria2Client *aria2.Client
// GetAria2Client returns the shared aria2 client instance
func GetAria2Client() *aria2.Client {
return aria2Client
}
func handleAria2DlCmd(ctx *ext.Context, update *ext.Update) error {
if !config.C().Aria2.Enable {
ctx.Reply(update, ext.ReplyTextString(i18n.T(i18nk.BotMsgAria2ErrorAria2NotEnabled)), nil)
@@ -78,7 +83,9 @@ func handleAria2DlCmd(ctx *ext.Context, update *ext.Update) error {
ctx.Reply(update, ext.ReplyTextString(i18n.T(i18nk.BotMsgDlErrorNoValidLinks)), nil)
return nil
}
logger.Debug("Adding aria2 download", "links", links)
logger.Debug("Preparing aria2 download", "links", links)
// Initialize aria2 client to check connection
aria2ClientInitOnce.Do(func() {
aria2Client, aria2ClientInitErr = aria2.NewClient(config.C().Aria2.Url, config.C().Aria2.Secret)
})
@@ -89,17 +96,18 @@ func handleAria2DlCmd(ctx *ext.Context, update *ext.Update) error {
})), nil)
return nil
}
gid, err := aria2Client.AddURI(ctx, links, nil)
// Build storage selection keyboard (don't add to aria2 yet)
markup, err := msgelem.BuildAddSelectStorageKeyboard(storage.GetUserStorages(ctx, update.GetUserChat().GetID()), tcbdata.Add{
TaskType: tasktype.TaskTypeAria2,
Aria2URIs: links,
})
if err != nil {
logger.Error("Failed to add aria2 download", "error", err)
ctx.Reply(update, ext.ReplyTextString(i18n.T(i18nk.BotMsgAria2ErrorAddingAria2Download, map[string]any{
"Error": err.Error(),
})), nil)
return nil
return err
}
logger.Info("Aria2 download added", "gid", gid)
ctx.Reply(update, ext.ReplyTextString(i18n.T(i18nk.BotMsgAria2InfoAria2DownloadAdded, map[string]any{
"GID": gid,
})), nil)
ctx.Reply(update, ext.ReplyTextString(i18n.T(i18nk.BotMsgAria2InfoSelectStorage)), &ext.ReplyOpts{
Markup: markup,
})
return nil
}

View File

@@ -49,6 +49,8 @@ func BuildAddSelectStorageKeyboard(stors []storage.Storage, adddata tcbdata.Add)
ParsedItem: adddata.ParsedItem,
DirectLinks: adddata.DirectLinks,
Aria2URIs: adddata.Aria2URIs,
}
dataid := xid.New().String()
err := cache.Set(dataid, data)

View File

@@ -0,0 +1,65 @@
package shortcut
import (
"github.com/celestix/gotgproto/dispatcher"
"github.com/celestix/gotgproto/ext"
"github.com/charmbracelet/log"
"github.com/gotd/td/tg"
"github.com/krau/SaveAny-Bot/common/i18n"
"github.com/krau/SaveAny-Bot/common/i18n/i18nk"
"github.com/krau/SaveAny-Bot/common/utils/tgutil"
"github.com/krau/SaveAny-Bot/core"
"github.com/krau/SaveAny-Bot/core/tasks/aria2dl"
"github.com/krau/SaveAny-Bot/pkg/aria2"
"github.com/krau/SaveAny-Bot/storage"
"github.com/rs/xid"
)
func CreateAndAddAria2TaskWithEdit(ctx *ext.Context, stor storage.Storage, dirPath string, uris []string, aria2Client *aria2.Client, msgID int, userID int64) error {
logger := log.FromContext(ctx)
injectCtx := tgutil.ExtWithContext(ctx.Context, ctx)
// Now add to aria2 after user selected storage
logger.Infof("Adding download to aria2, uris type: %T, value: %+v", uris, uris)
// Ensure uris is valid
if len(uris) == 0 {
logger.Error("URIs list is empty")
ctx.EditMessage(userID, &tg.MessagesEditMessageRequest{
ID: msgID,
Message: i18n.T(i18nk.BotMsgDlErrorNoValidLinks, nil),
})
return dispatcher.EndGroups
}
gid, err := aria2Client.AddURI(ctx, uris, nil)
if err != nil {
logger.Errorf("Failed to add aria2 download: %s", err)
ctx.EditMessage(userID, &tg.MessagesEditMessageRequest{
ID: msgID,
Message: i18n.T(i18nk.BotMsgAria2ErrorAddingAria2Download, map[string]any{
"Error": err.Error(),
}),
})
return dispatcher.EndGroups
}
logger.Infof("Aria2 download added with GID: %s", gid)
// Create task with the GID
task := aria2dl.NewTask(xid.New().String(), injectCtx, gid, uris, aria2Client, stor, stor.JoinStoragePath(dirPath), aria2dl.NewProgress(msgID, userID))
if err := core.AddTask(injectCtx, task); err != nil {
logger.Errorf("Failed to add task: %s", err)
ctx.EditMessage(userID, &tg.MessagesEditMessageRequest{
ID: msgID,
Message: i18n.T(i18nk.BotMsgCommonErrorTaskAddFailed, map[string]any{
"Error": err.Error(),
}),
})
return dispatcher.EndGroups
}
ctx.EditMessage(userID, &tg.MessagesEditMessageRequest{
ID: msgID,
Message: i18n.T(i18nk.BotMsgCommonInfoTaskAdded, nil),
})
return dispatcher.EndGroups
}