From 9ed44ae8ebe2f476693908833d790f0a3590f9e1 Mon Sep 17 00:00:00 2001 From: dj-oyu <68707227+dj-oyu@users.noreply.github.com> Date: Wed, 4 Mar 2026 17:37:39 +0900 Subject: [PATCH] Isolate task draft/status tracking by chat thread --- pkg/channels/manager.go | 39 +++++++++-- pkg/channels/manager_test.go | 123 ++++++++++++++++++++++++++++++++++- 2 files changed, 156 insertions(+), 6 deletions(-) diff --git a/pkg/channels/manager.go b/pkg/channels/manager.go index eac7b872a..6037b29b6 100644 --- a/pkg/channels/manager.go +++ b/pkg/channels/manager.go @@ -12,6 +12,7 @@ import ( "fmt" "hash/fnv" "math" + "strings" "sync" "time" @@ -100,7 +101,7 @@ type Manager struct { typingStops sync.Map // "channel:chatID" → typingEntry reactionUndos sync.Map // "channel:chatID" → reactionEntry statusMsgIDs sync.Map // "channel:chatID" → statusMsgEntry (streaming preview) - taskMsgIDs sync.Map // taskID → statusMsgEntry (background task status) + taskMsgIDs sync.Map // "channel:chatID:taskID" → statusMsgEntry (background task status) statusEditTimes sync.Map // key → time.Time — last EditMessage time for throttling } @@ -594,6 +595,16 @@ func (m *Manager) handleStatusSend(ctx context.Context, name string, w *channelW // 4. Channel doesn't support SendWithID or editing — drop silently } +func taskStatusKey(channel, chatID, taskID string) string { + if taskID == "" { + return "" + } + if channel == "" || chatID == "" { + return taskID + } + return channel + ":" + chatID + ":" + taskID +} + // handleTaskStatusSend processes IsTaskStatus messages (background task status). // It reuses a previously tracked task message, or sends a new one via SendWithID. // For channels implementing DraftSender, sendMessageDraft is used to avoid "(edited)". @@ -603,14 +614,27 @@ func (m *Manager) handleTaskStatusSend(ctx context.Context, name string, w *chan return } - taskKey := msg.TaskID + taskKey := taskStatusKey(name, msg.ChatID, msg.TaskID) // Final message: send as permanent (non-draft) message so it persists. // Drafts are ephemeral and disappear after a short time; the completion // message must survive. Clear the draft tracking and send via SendWithID // or regular Send, which creates a permanent Telegram message. if msg.Final { - m.taskMsgIDs.Delete(taskKey) + if v, loaded := m.taskMsgIDs.LoadAndDelete(taskKey); loaded { + if entry, ok := v.(statusMsgEntry); ok && entry.draftID != 0 { + if drafter, ok := w.ch.(DraftSender); ok { + if err := drafter.SendDraft(ctx, msg.ChatID, entry.draftID, ""); err != nil { + logger.WarnCF("channels", "Failed to dismiss task draft before final message", map[string]any{ + "task_id": taskKey, + "chat_id": msg.ChatID, + "draft_id": entry.draftID, + "error": err.Error(), + }) + } + } + } + } m.statusEditTimes.Delete(taskKey) if sender, ok := w.ch.(MessageSenderWithID); ok { if msgID, err := sender.SendWithID(ctx, msg.ChatID, msg.Content); err == nil && msgID != "" { @@ -992,7 +1016,7 @@ func (m *Manager) runTTLJanitor(ctx context.Context) { } // PromoteStatusToTask moves the tracked streaming status message for the given -// channel:chatID key into the task message map under taskID. This allows the +// channel:chatID key into the task message map under channel:chatID:taskID. This allows the // next IsTaskStatus publish to edit the streaming bubble instead of creating a // new message. Returns true if a status message was found and promoted. func (m *Manager) PromoteStatusToTask(statusKey, taskID string) bool { @@ -1000,6 +1024,13 @@ func (m *Manager) PromoteStatusToTask(statusKey, taskID string) bool { if !loaded { return false } + + parts := strings.SplitN(statusKey, ":", 2) + if len(parts) == 2 { + m.taskMsgIDs.Store(taskStatusKey(parts[0], parts[1], taskID), v) + return true + } + m.taskMsgIDs.Store(taskID, v) return true } diff --git a/pkg/channels/manager_test.go b/pkg/channels/manager_test.go index 3b8d28ac3..2afb16a35 100644 --- a/pkg/channels/manager_test.go +++ b/pkg/channels/manager_test.go @@ -1019,7 +1019,7 @@ func TestHandleTaskStatusSend_EditsExisting(t *testing.T) { w := &channelWorker{ch: ch, limiter: rate.NewLimiter(rate.Inf, 1)} // Pre-store task message - m.taskMsgIDs.Store("task-abc", statusMsgEntry{messageID: "task-msg-1", createdAt: time.Now()}) + m.taskMsgIDs.Store(taskStatusKey("test", "123", "task-abc"), statusMsgEntry{messageID: "task-msg-1", createdAt: time.Now()}) msg := bus.OutboundMessage{ Channel: "test", @@ -1065,7 +1065,7 @@ func TestHandleTaskStatusSend_SendsNewAndTracks(t *testing.T) { t.Fatal("expected SendWithID to be called") } - v, ok := m.taskMsgIDs.Load("task-xyz") + v, ok := m.taskMsgIDs.Load(taskStatusKey("test", "123", "task-xyz")) if !ok { t.Fatal("expected taskMsgIDs to contain tracked entry") } @@ -1436,6 +1436,125 @@ func TestHandleTaskStatusSend_UsesDraftSender(t *testing.T) { } } +func TestHandleTaskStatusSend_Final_DismissesDraftBeforePermanentMessage(t *testing.T) { + m := newTestManager() + var dismissCalled bool + var dismissDraftID int + var dismissContent string + var finalSendWithIDCalled bool + + ch := &mockDraftSender{ + mockChannel: mockChannel{ + sendFn: func(_ context.Context, _ bus.OutboundMessage) error { + t.Fatal("Send should not be called when SendWithID succeeds") + return nil + }, + }, + draftFn: func(_ context.Context, chatID string, draftID int, content string) error { + dismissCalled = true + dismissDraftID = draftID + dismissContent = content + if chatID != "123" { + t.Fatalf("expected dismiss chatID 123, got %s", chatID) + } + return nil + }, + editFn: func(_ context.Context, _, _, _ string) error { return nil }, + sendWithID: func(_ context.Context, chatID, content string) (string, error) { + finalSendWithIDCalled = true + if chatID != "123" { + t.Fatalf("expected final chatID 123, got %s", chatID) + } + if content != "task completed" { + t.Fatalf("expected final content 'task completed', got %s", content) + } + return "task-final-1", nil + }, + } + + w := &channelWorker{ch: ch, limiter: rate.NewLimiter(rate.Inf, 1)} + + m.taskMsgIDs.Store(taskStatusKey("test", "123", "task-final"), statusMsgEntry{draftID: 42, createdAt: time.Now()}) + m.statusEditTimes.Store(taskStatusKey("test", "123", "task-final"), time.Now()) + + msg := bus.OutboundMessage{ + Channel: "test", + ChatID: "123", + Content: "task completed", + IsTaskStatus: true, + TaskID: "task-final", + Final: true, + } + m.handleTaskStatusSend(context.Background(), "test", w, msg) + + if !dismissCalled { + t.Fatal("expected SendDraft dismiss call for final task status") + } + if dismissDraftID != 42 { + t.Fatalf("expected dismiss draftID 42, got %d", dismissDraftID) + } + if dismissContent != "" { + t.Fatalf("expected empty dismiss content, got %q", dismissContent) + } + if !finalSendWithIDCalled { + t.Fatal("expected final SendWithID to be called") + } + if _, loaded := m.taskMsgIDs.Load(taskStatusKey("test", "123", "task-final")); loaded { + t.Fatal("expected taskMsgIDs entry to be deleted for final task status") + } + if _, loaded := m.statusEditTimes.Load(taskStatusKey("test", "123", "task-final")); loaded { + t.Fatal("expected statusEditTimes entry to be deleted for final task status") + } +} + +func TestHandleTaskStatusSend_DraftStreaming_IsolatedByChatThread(t *testing.T) { + m := newTestManager() + + type draftCall struct { + chatID string + draftID int + content string + } + calls := make([]draftCall, 0, 2) + + ch := &mockDraftSender{ + mockChannel: mockChannel{ + sendFn: func(_ context.Context, _ bus.OutboundMessage) error { return nil }, + }, + draftFn: func(_ context.Context, chatID string, draftID int, content string) error { + calls = append(calls, draftCall{chatID: chatID, draftID: draftID, content: content}) + return nil + }, + editFn: func(_ context.Context, _, _, _ string) error { return nil }, + sendWithID: func(_ context.Context, _, _ string) (string, error) { return "", nil }, + } + + w := &channelWorker{ch: ch, limiter: rate.NewLimiter(rate.Inf, 1)} + + msgA := bus.OutboundMessage{Channel: "test", ChatID: "-100/10", Content: "A:10%", IsTaskStatus: true, TaskID: "shared-task"} + msgB := bus.OutboundMessage{Channel: "test", ChatID: "-100/20", Content: "B:10%", IsTaskStatus: true, TaskID: "shared-task"} + + m.handleTaskStatusSend(context.Background(), "test", w, msgA) + m.handleTaskStatusSend(context.Background(), "test", w, msgB) + + if len(calls) != 2 { + t.Fatalf("expected 2 SendDraft calls, got %d", len(calls)) + } + if calls[0].chatID == calls[1].chatID { + t.Fatalf("expected different chat threads, got %q and %q", calls[0].chatID, calls[1].chatID) + } + if calls[0].draftID == calls[1].draftID { + t.Fatalf("expected distinct draft IDs per thread key, both got %d", calls[0].draftID) + } + + if _, loaded := m.taskMsgIDs.Load(taskStatusKey("test", "-100/10", "shared-task")); !loaded { + t.Fatal("expected taskMsgIDs entry for thread A") + } + if _, loaded := m.taskMsgIDs.Load(taskStatusKey("test", "-100/20", "shared-task")); !loaded { + t.Fatal("expected taskMsgIDs entry for thread B") + } +} + func TestHandleTaskStatusSend_DraftFailure_DoesNotClobberTrackedMessageID(t *testing.T) { m := newTestManager() var sendWithIDCount int