From 4f314bd37fe2558b860030cdcd0b2e57175634ab Mon Sep 17 00:00:00 2001 From: krau <71133316+krau@users.noreply.github.com> Date: Sun, 16 Nov 2025 21:44:51 +0800 Subject: [PATCH] feat: implement configurable media group handling timeout, close #137 --- client/bot/handlers/media.go | 101 ------------------------ client/bot/handlers/media_group.go | 122 +++++++++++++++++++++++++++++ config/tg.go | 13 +-- 3 files changed, 129 insertions(+), 107 deletions(-) create mode 100644 client/bot/handlers/media_group.go diff --git a/client/bot/handlers/media.go b/client/bot/handlers/media.go index 9891af5..bbd9d71 100644 --- a/client/bot/handlers/media.go +++ b/client/bot/handlers/media.go @@ -1,22 +1,14 @@ package handlers import ( - "fmt" - "sync" - "time" - "github.com/celestix/gotgproto/dispatcher" "github.com/celestix/gotgproto/ext" "github.com/charmbracelet/log" - "github.com/gotd/td/tg" "github.com/krau/SaveAny-Bot/client/bot/handlers/utils/dirutil" "github.com/krau/SaveAny-Bot/client/bot/handlers/utils/mediautil" "github.com/krau/SaveAny-Bot/client/bot/handlers/utils/msgelem" "github.com/krau/SaveAny-Bot/client/bot/handlers/utils/shortcut" - "github.com/krau/SaveAny-Bot/common/utils/tgutil" "github.com/krau/SaveAny-Bot/database" - "github.com/krau/SaveAny-Bot/pkg/tcbdata" - "github.com/krau/SaveAny-Bot/pkg/tfile" "github.com/krau/SaveAny-Bot/storage" ) @@ -76,96 +68,3 @@ func handleSilentSaveMedia(ctx *ext.Context, update *ext.Update) error { } return shortcut.CreateAndAddTGFileTaskWithEdit(ctx, userID, stor, dirutil.PathFromContext(ctx), file, msg.ID) } - -type MediaGroupHandler struct { - groups map[int64][]tfile.TGFileMessage - timers map[int64]*time.Timer - mu sync.Mutex - timeout time.Duration -} - -var mediaGroupHandler = &MediaGroupHandler{ - groups: make(map[int64][]tfile.TGFileMessage), - timers: make(map[int64]*time.Timer), - timeout: 1 * time.Second, -} - -func handleGroupMediaMessage(ctx *ext.Context, update *ext.Update, message *tg.Message, groupID int64) error { - logger := log.FromContext(ctx) - media := message.Media - supported := mediautil.IsSupported(media) - if !supported { - return dispatcher.EndGroups - } - file, err := tfile.FromMediaMessage(media, ctx.Raw, message, tfile.WithNameIfEmpty( - tgutil.GenFileNameFromMessage(*message), - )) - if err != nil { - logger.Errorf("Failed to get file from media: %s", err) - return dispatcher.EndGroups - } - mediaGroupHandler.mu.Lock() - defer mediaGroupHandler.mu.Unlock() - if mediaGroupHandler.groups[groupID] == nil { - mediaGroupHandler.groups[groupID] = make([]tfile.TGFileMessage, 0) - } - mediaGroupHandler.groups[groupID] = append(mediaGroupHandler.groups[groupID], file) - - if timer, exists := mediaGroupHandler.timers[groupID]; exists { - timer.Stop() - } - mediaGroupHandler.timers[groupID] = time.AfterFunc(mediaGroupHandler.timeout, func() { - processMediaGroup(ctx, update, groupID) - }) - return dispatcher.EndGroups -} - -func processMediaGroup(ctx *ext.Context, update *ext.Update, groupID int64) { - logger := log.FromContext(ctx) - mediaGroupHandler.mu.Lock() - items := mediaGroupHandler.groups[groupID] - delete(mediaGroupHandler.groups, groupID) - delete(mediaGroupHandler.timers, groupID) - mediaGroupHandler.mu.Unlock() - if len(items) == 0 { - logger.Warn("No media items to process for group", "groupID", groupID) - return - } - logger.Debugf("Processing media group %d with %d items", groupID, len(items)) - - userId := update.GetUserChat().GetID() - msg, err := ctx.Reply(update, ext.ReplyTextString("正在保存文件..."), nil) - if err != nil { - logger.Errorf("Failed to reply: %s", err) - return - } - stor := storage.FromContext(ctx) - if stor != nil { - // In silent mode - if len(items) == 1 { - shortcut.CreateAndAddTGFileTaskWithEdit(ctx, userId, stor, "", items[0], msg.ID) - return - } - shortcut.CreateAndAddBatchTGFileTaskWithEdit(ctx, userId, stor, "", items, msg.ID) - return - } - - stors := storage.GetUserStorages(ctx, userId) - markup, err := msgelem.BuildAddSelectStorageKeyboard(stors, tcbdata.Add{ - Files: items, - AsBatch: len(items) > 1, - }) - if err != nil { - logger.Errorf("构建存储选择键盘失败: %s", err) - ctx.EditMessage(userId, &tg.MessagesEditMessageRequest{ - ID: msg.ID, - Message: "构建存储选择键盘失败: " + err.Error(), - }) - return - } - ctx.EditMessage(userId, &tg.MessagesEditMessageRequest{ - ID: msg.ID, - Message: fmt.Sprintf("共 %d 个文件, 请选择存储位置", len(items)), - ReplyMarkup: markup, - }) -} diff --git a/client/bot/handlers/media_group.go b/client/bot/handlers/media_group.go new file mode 100644 index 0000000..f71ac27 --- /dev/null +++ b/client/bot/handlers/media_group.go @@ -0,0 +1,122 @@ +package handlers + +import ( + "fmt" + "sync" + "time" + + "github.com/celestix/gotgproto/dispatcher" + "github.com/celestix/gotgproto/ext" + "github.com/charmbracelet/log" + "github.com/gotd/td/tg" + "github.com/krau/SaveAny-Bot/client/bot/handlers/utils/mediautil" + "github.com/krau/SaveAny-Bot/client/bot/handlers/utils/msgelem" + "github.com/krau/SaveAny-Bot/client/bot/handlers/utils/shortcut" + "github.com/krau/SaveAny-Bot/common/utils/tgutil" + "github.com/krau/SaveAny-Bot/config" + "github.com/krau/SaveAny-Bot/pkg/tcbdata" + "github.com/krau/SaveAny-Bot/pkg/tfile" + "github.com/krau/SaveAny-Bot/storage" +) + +type MediaGroupHandler struct { + groups map[int64][]tfile.TGFileMessage + timers map[int64]*time.Timer + mu sync.Mutex + timeout time.Duration +} + +var ( + mediaGroupHandler *MediaGroupHandler + onceMediaGroup sync.Once + setupMediaGroupHandler = func() { + onceMediaGroup.Do(func() { + mediaGroupHandler = &MediaGroupHandler{ + groups: make(map[int64][]tfile.TGFileMessage), + timers: make(map[int64]*time.Timer), + timeout: time.Duration(min(config.C().Telegram.MediaGroupTimeout, 1)) * time.Second, + } + }) + } +) + +func handleGroupMediaMessage(ctx *ext.Context, update *ext.Update, message *tg.Message, groupID int64) error { + onceMediaGroup.Do(setupMediaGroupHandler) + logger := log.FromContext(ctx) + media := message.Media + supported := mediautil.IsSupported(media) + if !supported { + return dispatcher.EndGroups + } + file, err := tfile.FromMediaMessage(media, ctx.Raw, message, tfile.WithNameIfEmpty( + tgutil.GenFileNameFromMessage(*message), + )) + if err != nil { + logger.Errorf("Failed to get file from media: %s", err) + return dispatcher.EndGroups + } + mediaGroupHandler.mu.Lock() + defer mediaGroupHandler.mu.Unlock() + if mediaGroupHandler.groups[groupID] == nil { + mediaGroupHandler.groups[groupID] = make([]tfile.TGFileMessage, 0) + } + mediaGroupHandler.groups[groupID] = append(mediaGroupHandler.groups[groupID], file) + + if timer, exists := mediaGroupHandler.timers[groupID]; exists { + timer.Stop() + } + mediaGroupHandler.timers[groupID] = time.AfterFunc(mediaGroupHandler.timeout, func() { + processMediaGroup(ctx, update, groupID) + }) + return dispatcher.EndGroups +} + +func processMediaGroup(ctx *ext.Context, update *ext.Update, groupID int64) { + logger := log.FromContext(ctx) + mediaGroupHandler.mu.Lock() + items := mediaGroupHandler.groups[groupID] + delete(mediaGroupHandler.groups, groupID) + delete(mediaGroupHandler.timers, groupID) + mediaGroupHandler.mu.Unlock() + if len(items) == 0 { + logger.Warn("No media items to process for group", "groupID", groupID) + return + } + logger.Debugf("Processing media group %d with %d items", groupID, len(items)) + + userId := update.GetUserChat().GetID() + msg, err := ctx.Reply(update, ext.ReplyTextString("正在保存文件..."), nil) + if err != nil { + logger.Errorf("Failed to reply: %s", err) + return + } + stor := storage.FromContext(ctx) + if stor != nil { + // In silent mode + if len(items) == 1 { + shortcut.CreateAndAddTGFileTaskWithEdit(ctx, userId, stor, "", items[0], msg.ID) + return + } + shortcut.CreateAndAddBatchTGFileTaskWithEdit(ctx, userId, stor, "", items, msg.ID) + return + } + + stors := storage.GetUserStorages(ctx, userId) + markup, err := msgelem.BuildAddSelectStorageKeyboard(stors, tcbdata.Add{ + Files: items, + AsBatch: len(items) > 1, + }) + if err != nil { + logger.Errorf("构建存储选择键盘失败: %s", err) + ctx.EditMessage(userId, &tg.MessagesEditMessageRequest{ + ID: msg.ID, + Message: "构建存储选择键盘失败: " + err.Error(), + }) + return + } + ctx.EditMessage(userId, &tg.MessagesEditMessageRequest{ + ID: msg.ID, + Message: fmt.Sprintf("共 %d 个文件, 请选择存储位置", len(items)), + ReplyMarkup: markup, + }) +} diff --git a/config/tg.go b/config/tg.go index 2f15790..9b26e40 100644 --- a/config/tg.go +++ b/config/tg.go @@ -1,12 +1,13 @@ package config type telegramConfig struct { - Token string `toml:"token" mapstructure:"token"` - AppID int `toml:"app_id" mapstructure:"app_id" json:"app_id"` - AppHash string `toml:"app_hash" mapstructure:"app_hash" json:"app_hash"` - Proxy tgProxyConfig `toml:"proxy" mapstructure:"proxy"` - RpcRetry int `toml:"rpc_retry" mapstructure:"rpc_retry" json:"rpc_retry"` - Userbot userbotConfig `toml:"userbot" mapstructure:"userbot" json:"userbot"` + Token string `toml:"token" mapstructure:"token"` + AppID int `toml:"app_id" mapstructure:"app_id" json:"app_id"` + AppHash string `toml:"app_hash" mapstructure:"app_hash" json:"app_hash"` + Proxy tgProxyConfig `toml:"proxy" mapstructure:"proxy"` + RpcRetry int `toml:"rpc_retry" mapstructure:"rpc_retry" json:"rpc_retry"` + Userbot userbotConfig `toml:"userbot" mapstructure:"userbot" json:"userbot"` + MediaGroupTimeout int `toml:"media_group_timeout" mapstructure:"media_group_timeout" json:"media_group_timeout"` } type userbotConfig struct {