feat: heartbeat message dedup + Telegram sendMessageDraft streaming
Task 1 - Heartbeat message dedup: - Include LLM final response in task completion status bubble instead of sending separate messages via sendResponse - Split long responses (>4096 chars) into status header + regular message (auto-split by SplitMessage infrastructure) - Always return SilentResult from heartbeat handler to prevent duplicate bubbles - Remove unused sendResponse method from HeartbeatService Task 2 - Telegram sendMessageDraft streaming: - Add DraftSender interface for channels supporting progressive drafts - Implement SendDraft on TelegramChannel using telego SendMessageDraft - Prefer draft-based streaming in handleStatusSend/handleTaskStatusSend with automatic fallback to EditMessage - Replace fixed-interval streaming throttle with Go channel + goroutine consumer for natural backpressure - Add 500ms EditMessage throttle for non-draft channels to avoid API rate limits and "(edited)" flicker - Clean up draft state in preSend so final sendMessage replaces draft - Add statusEditTimes cleanup to TTL janitor Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
parent
3706b2d5f0
commit
bdf0cc1ea2
8 changed files with 474 additions and 86 deletions
|
|
@ -125,14 +125,11 @@ func gatewayCmd(debug bool, orchestration bool, enableStats bool) error {
|
|||
if err != nil {
|
||||
return tools.ErrorResult(fmt.Sprintf("Heartbeat error: %v", err))
|
||||
}
|
||||
// Always return SilentResult — the task completion message in runAgentLoop
|
||||
// already includes the LLM response in the same status bubble.
|
||||
if response == "HEARTBEAT_OK" {
|
||||
return tools.SilentResult("Heartbeat OK")
|
||||
}
|
||||
// Deliver response to user when a plan interview/review needs resuming.
|
||||
// For async tasks (spawn), results are delivered separately via processSystemMessage.
|
||||
if status := agentLoop.GetPlanStatus(); status == "interviewing" || status == "review" {
|
||||
return tools.UserResult(response)
|
||||
}
|
||||
return tools.SilentResult(response)
|
||||
})
|
||||
|
||||
|
|
|
|||
|
|
@ -1028,6 +1028,10 @@ func (al *AgentLoop) runAgentLoop(ctx context.Context, agent *AgentInstance, opt
|
|||
}
|
||||
}
|
||||
|
||||
// Shared variable for capturing LLM's final response. The defer below reads it
|
||||
// to include the response in the task completion message.
|
||||
var finalContent string
|
||||
|
||||
// Use TaskID as key if available (for background tasks), else sessionKey
|
||||
taskKey := opts.SessionKey
|
||||
if opts.TaskID != "" {
|
||||
|
|
@ -1037,14 +1041,43 @@ func (al *AgentLoop) runAgentLoop(ctx context.Context, agent *AgentInstance, opt
|
|||
defer func() {
|
||||
al.activeTasks.Delete(taskKey)
|
||||
|
||||
// Publish final task status on completion for background tasks
|
||||
// Publish final task status on completion for background tasks.
|
||||
// Include finalContent so the LLM response appears in the same bubble
|
||||
// as the completion status, avoiding duplicate messages.
|
||||
if opts.TaskID != "" {
|
||||
elapsed := time.Since(task.StartedAt)
|
||||
completionMsg := fmt.Sprintf("\u2705 Task completed (%.1fs)", elapsed.Seconds())
|
||||
if finalContent != "" && finalContent != defaultResponse {
|
||||
// Keep completion + response in one bubble if short enough (4096 = Telegram limit).
|
||||
// If too long, edit the status bubble with the header, then send the
|
||||
// full response as a regular message — the channel worker's SplitMessage
|
||||
// will automatically chunk it for channels with MaxMessageLength.
|
||||
combined := completionMsg + "\n\n" + finalContent
|
||||
if len([]rune(combined)) <= 4096 {
|
||||
completionMsg = combined
|
||||
} else {
|
||||
doneCtx, doneCancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
_ = al.bus.PublishOutbound(doneCtx, bus.OutboundMessage{
|
||||
Channel: opts.Channel,
|
||||
ChatID: opts.ChatID,
|
||||
Content: completionMsg,
|
||||
IsTaskStatus: true,
|
||||
TaskID: opts.TaskID,
|
||||
})
|
||||
_ = al.bus.PublishOutbound(doneCtx, bus.OutboundMessage{
|
||||
Channel: opts.Channel,
|
||||
ChatID: opts.ChatID,
|
||||
Content: finalContent,
|
||||
})
|
||||
doneCancel()
|
||||
return
|
||||
}
|
||||
}
|
||||
doneCtx, doneCancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
_ = al.bus.PublishOutbound(doneCtx, bus.OutboundMessage{
|
||||
Channel: opts.Channel,
|
||||
ChatID: opts.ChatID,
|
||||
Content: fmt.Sprintf("\u2705 Task completed (%.1fs)\n%s", elapsed.Seconds(), task.Description),
|
||||
Content: completionMsg,
|
||||
IsTaskStatus: true,
|
||||
TaskID: opts.TaskID,
|
||||
})
|
||||
|
|
@ -1182,7 +1215,6 @@ func (al *AgentLoop) runAgentLoop(ctx context.Context, agent *AgentInstance, opt
|
|||
}
|
||||
|
||||
// 5. Run LLM iteration loop (with automatic phase transitions)
|
||||
var finalContent string
|
||||
var iteration int
|
||||
const maxPhaseTransitions = 10
|
||||
|
||||
|
|
@ -2155,23 +2187,47 @@ func (al *AgentLoop) runLLMIteration(
|
|||
var err error
|
||||
|
||||
// Build onChunk callback for streaming preview.
|
||||
// When sending responses to a real (non-internal) channel, publish
|
||||
// throttled status updates so the user sees LLM output in real time.
|
||||
// Instead of a fixed-interval throttle, use a Go channel with
|
||||
// latest-value semantics: a consumer goroutine publishes status
|
||||
// updates as fast as the bus → manager → channel pipeline allows.
|
||||
// Backpressure is provided naturally by the per-channel rate limiter
|
||||
// (e.g. 20 msg/s for Telegram's SendDraft, 1 msg/s for Discord's EditMessage).
|
||||
type streamUpdate struct{ accumulated, reasoning string }
|
||||
var onChunk func(string, string)
|
||||
var streamCh chan streamUpdate
|
||||
var streamDone chan struct{}
|
||||
if !constants.IsInternalChannel(opts.Channel) {
|
||||
lastPublish := time.Time{}
|
||||
onChunk = func(accumulated, reasoning string) {
|
||||
if time.Since(lastPublish) < 500*time.Millisecond {
|
||||
return
|
||||
streamCh = make(chan streamUpdate, 1)
|
||||
streamDone = make(chan struct{})
|
||||
go func() {
|
||||
defer close(streamDone)
|
||||
for up := range streamCh {
|
||||
display := buildStreamingDisplay(up.accumulated, up.reasoning)
|
||||
_ = al.bus.PublishOutbound(ctx, bus.OutboundMessage{
|
||||
Channel: opts.Channel,
|
||||
ChatID: opts.ChatID,
|
||||
Content: display,
|
||||
IsStatus: true,
|
||||
})
|
||||
}
|
||||
}()
|
||||
onChunk = func(accumulated, reasoning string) {
|
||||
up := streamUpdate{accumulated, reasoning}
|
||||
// Non-blocking latest-value send: if the consumer hasn't
|
||||
// drained the previous update, replace it with the latest.
|
||||
select {
|
||||
case streamCh <- up:
|
||||
default:
|
||||
// Channel full — drain stale value, then send latest.
|
||||
select {
|
||||
case <-streamCh:
|
||||
default:
|
||||
}
|
||||
select {
|
||||
case streamCh <- up:
|
||||
default:
|
||||
}
|
||||
}
|
||||
lastPublish = time.Now()
|
||||
display := buildStreamingDisplay(accumulated, reasoning)
|
||||
_ = al.bus.PublishOutbound(ctx, bus.OutboundMessage{
|
||||
Channel: opts.Channel,
|
||||
ChatID: opts.ChatID,
|
||||
Content: display,
|
||||
IsStatus: true,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -2306,6 +2362,17 @@ func (al *AgentLoop) runLLMIteration(
|
|||
break
|
||||
}
|
||||
|
||||
// Streaming finished — close the stream goroutine so it flushes
|
||||
// the last update and exits cleanly before we process the response.
|
||||
if streamDone != nil {
|
||||
// onChunk is captured by doCall closures; nil it to avoid
|
||||
// writes after the channel is closed during retries.
|
||||
onChunk = nil
|
||||
close(streamCh)
|
||||
<-streamDone
|
||||
streamDone = nil
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
logger.ErrorCF("agent", "LLM call failed",
|
||||
map[string]any{
|
||||
|
|
|
|||
|
|
@ -37,6 +37,13 @@ type PlaceholderCapable interface {
|
|||
SendPlaceholder(ctx context.Context, chatID string) (messageID string, err error)
|
||||
}
|
||||
|
||||
// DraftSender — channels that can send progressive draft messages.
|
||||
// Used for streaming LLM output without the "edited" indicator.
|
||||
// draftID must be non-zero and consistent across updates for the same draft.
|
||||
type DraftSender interface {
|
||||
SendDraft(ctx context.Context, chatID string, draftID int, content string) error
|
||||
}
|
||||
|
||||
// PlaceholderRecorder is injected into channels by Manager.
|
||||
// Channels call these methods on inbound to register typing/placeholder state.
|
||||
// Manager uses the registered state on outbound to stop typing and edit placeholders.
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ import (
|
|||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"hash/fnv"
|
||||
"math"
|
||||
"sync"
|
||||
"time"
|
||||
|
|
@ -37,6 +38,12 @@ const (
|
|||
placeholderTTL = 10 * time.Minute
|
||||
statusMsgTTL = 5 * time.Minute
|
||||
taskMsgTTL = 30 * time.Minute
|
||||
|
||||
// statusEditInterval is the minimum interval between EditMessage calls
|
||||
// for the same status/task bubble. EditMessage APIs are more rate-sensitive
|
||||
// than SendMessageDraft, so we throttle edits to avoid "(edited)" flicker
|
||||
// and API rate limit errors. Draft-based channels bypass this throttle.
|
||||
statusEditInterval = 500 * time.Millisecond
|
||||
)
|
||||
|
||||
// typingEntry wraps a typing stop function with a creation timestamp for TTL eviction.
|
||||
|
|
@ -60,6 +67,7 @@ type placeholderEntry struct {
|
|||
// statusMsgEntry tracks a status or task message ID for later editing.
|
||||
type statusMsgEntry struct {
|
||||
messageID string
|
||||
draftID int // non-zero when using draft-based streaming
|
||||
createdAt time.Time
|
||||
}
|
||||
|
||||
|
|
@ -81,18 +89,19 @@ type channelWorker struct {
|
|||
}
|
||||
|
||||
type Manager struct {
|
||||
channels map[string]Channel
|
||||
workers map[string]*channelWorker
|
||||
bus *bus.MessageBus
|
||||
config *config.Config
|
||||
mediaStore media.MediaStore
|
||||
dispatchTask *asyncTask
|
||||
mu sync.RWMutex
|
||||
placeholders sync.Map // "channel:chatID" → placeholderEntry
|
||||
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)
|
||||
channels map[string]Channel
|
||||
workers map[string]*channelWorker
|
||||
bus *bus.MessageBus
|
||||
config *config.Config
|
||||
mediaStore media.MediaStore
|
||||
dispatchTask *asyncTask
|
||||
mu sync.RWMutex
|
||||
placeholders sync.Map // "channel:chatID" → placeholderEntry
|
||||
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)
|
||||
statusEditTimes sync.Map // key → time.Time — last EditMessage time for throttling
|
||||
}
|
||||
|
||||
type asyncTask struct {
|
||||
|
|
@ -139,12 +148,18 @@ func (m *Manager) preSend(ctx context.Context, name string, msg bus.OutboundMess
|
|||
}
|
||||
}
|
||||
|
||||
// 3. Try editing a tracked status message (from streaming preview)
|
||||
// 3. Try editing a tracked status message (from streaming preview).
|
||||
// If the status was draft-based (draftID != 0), just clear the entry —
|
||||
// the final sendMessage will automatically replace the draft bubble.
|
||||
if v, loaded := m.statusMsgIDs.LoadAndDelete(key); loaded {
|
||||
if entry, ok := v.(statusMsgEntry); ok && entry.messageID != "" {
|
||||
if editor, ok := ch.(MessageEditor); ok {
|
||||
if err := editor.EditMessage(ctx, msg.ChatID, entry.messageID, msg.Content); err == nil {
|
||||
return true // edited successfully, skip Send
|
||||
if entry, ok := v.(statusMsgEntry); ok {
|
||||
if entry.draftID != 0 {
|
||||
// Draft-based: sendMessage replaces the draft, no edit needed
|
||||
} else if entry.messageID != "" {
|
||||
if editor, ok := ch.(MessageEditor); ok {
|
||||
if err := editor.EditMessage(ctx, msg.ChatID, entry.messageID, msg.Content); err == nil {
|
||||
return true // edited successfully, skip Send
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -474,6 +489,8 @@ func (m *Manager) runWorker(ctx context.Context, name string, w *channelWorker)
|
|||
// handleStatusSend processes IsStatus messages (streaming previews).
|
||||
// It reuses an existing placeholder or tracked status message, or sends a new
|
||||
// one via SendWithID so subsequent status updates edit the same bubble.
|
||||
// For channels implementing DraftSender (e.g. Telegram private chats),
|
||||
// sendMessageDraft is preferred as it avoids the "(edited)" indicator.
|
||||
// If the channel doesn't support editing, the message is silently dropped.
|
||||
func (m *Manager) handleStatusSend(ctx context.Context, name string, w *channelWorker, msg bus.OutboundMessage) {
|
||||
if err := w.limiter.Wait(ctx); err != nil {
|
||||
|
|
@ -482,11 +499,41 @@ func (m *Manager) handleStatusSend(ctx context.Context, name string, w *channelW
|
|||
|
||||
key := name + ":" + msg.ChatID
|
||||
|
||||
// 0. Draft-based streaming (preferred for supported channels)
|
||||
if drafter, ok := w.ch.(DraftSender); ok {
|
||||
var did int
|
||||
if v, loaded := m.statusMsgIDs.Load(key); loaded {
|
||||
if entry, ok := v.(statusMsgEntry); ok && entry.draftID != 0 {
|
||||
did = entry.draftID
|
||||
}
|
||||
}
|
||||
if did == 0 {
|
||||
did = generateDraftID(key)
|
||||
m.statusMsgIDs.Store(key, statusMsgEntry{
|
||||
draftID: did,
|
||||
createdAt: time.Now(),
|
||||
})
|
||||
}
|
||||
if err := drafter.SendDraft(ctx, msg.ChatID, did, msg.Content); err == nil {
|
||||
return
|
||||
}
|
||||
// Draft failed — fall through to edit-based approach
|
||||
}
|
||||
|
||||
// Edit-based path: throttle to statusEditInterval per key to avoid
|
||||
// API rate limit errors and "(edited)" flicker.
|
||||
if v, loaded := m.statusEditTimes.Load(key); loaded {
|
||||
if t, ok := v.(time.Time); ok && time.Since(t) < statusEditInterval {
|
||||
return // too recent, skip this update
|
||||
}
|
||||
}
|
||||
|
||||
// 1. Try editing an existing placeholder
|
||||
if v, loaded := m.placeholders.Load(key); loaded {
|
||||
if entry, ok := v.(placeholderEntry); ok && entry.id != "" {
|
||||
if editor, ok := w.ch.(MessageEditor); ok {
|
||||
if err := editor.EditMessage(ctx, msg.ChatID, entry.id, msg.Content); err == nil {
|
||||
m.statusEditTimes.Store(key, time.Now())
|
||||
return
|
||||
}
|
||||
}
|
||||
|
|
@ -498,6 +545,7 @@ func (m *Manager) handleStatusSend(ctx context.Context, name string, w *channelW
|
|||
if entry, ok := v.(statusMsgEntry); ok && entry.messageID != "" {
|
||||
if editor, ok := w.ch.(MessageEditor); ok {
|
||||
if err := editor.EditMessage(ctx, msg.ChatID, entry.messageID, msg.Content); err == nil {
|
||||
m.statusEditTimes.Store(key, time.Now())
|
||||
return
|
||||
}
|
||||
}
|
||||
|
|
@ -520,6 +568,7 @@ func (m *Manager) handleStatusSend(ctx context.Context, name string, w *channelW
|
|||
|
||||
// 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)".
|
||||
// If the channel doesn't support editing, falls back to regular Send.
|
||||
func (m *Manager) handleTaskStatusSend(ctx context.Context, name string, w *channelWorker, msg bus.OutboundMessage) {
|
||||
if err := w.limiter.Wait(ctx); err != nil {
|
||||
|
|
@ -528,12 +577,43 @@ func (m *Manager) handleTaskStatusSend(ctx context.Context, name string, w *chan
|
|||
|
||||
taskKey := msg.TaskID
|
||||
|
||||
// 0. Draft-based streaming (preferred for supported channels)
|
||||
if drafter, ok := w.ch.(DraftSender); ok && taskKey != "" {
|
||||
var did int
|
||||
if v, loaded := m.taskMsgIDs.Load(taskKey); loaded {
|
||||
if entry, ok := v.(statusMsgEntry); ok && entry.draftID != 0 {
|
||||
did = entry.draftID
|
||||
}
|
||||
}
|
||||
if did == 0 {
|
||||
did = generateDraftID(taskKey)
|
||||
m.taskMsgIDs.Store(taskKey, statusMsgEntry{
|
||||
draftID: did,
|
||||
createdAt: time.Now(),
|
||||
})
|
||||
}
|
||||
if err := drafter.SendDraft(ctx, msg.ChatID, did, msg.Content); err == nil {
|
||||
return
|
||||
}
|
||||
// Draft failed — fall through to edit-based approach
|
||||
}
|
||||
|
||||
// Edit-based path: throttle to statusEditInterval per task key.
|
||||
if taskKey != "" {
|
||||
if v, loaded := m.statusEditTimes.Load(taskKey); loaded {
|
||||
if t, ok := v.(time.Time); ok && time.Since(t) < statusEditInterval {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 1. Try editing an existing task message
|
||||
if taskKey != "" {
|
||||
if v, loaded := m.taskMsgIDs.Load(taskKey); loaded {
|
||||
if entry, ok := v.(statusMsgEntry); ok && entry.messageID != "" {
|
||||
if editor, ok := w.ch.(MessageEditor); ok {
|
||||
if err := editor.EditMessage(ctx, msg.ChatID, entry.messageID, msg.Content); err == nil {
|
||||
m.statusEditTimes.Store(taskKey, time.Now())
|
||||
return
|
||||
}
|
||||
}
|
||||
|
|
@ -558,6 +638,22 @@ func (m *Manager) handleTaskStatusSend(ctx context.Context, name string, w *chan
|
|||
_ = w.ch.Send(ctx, msg)
|
||||
}
|
||||
|
||||
// generateDraftID produces a stable non-zero int from a key string.
|
||||
// The same key always maps to the same draft ID so successive calls
|
||||
// animate the same Telegram draft bubble.
|
||||
func generateDraftID(key string) int {
|
||||
h := fnv.New32a()
|
||||
h.Write([]byte(key))
|
||||
v := int(h.Sum32())
|
||||
if v == 0 {
|
||||
v = 1 // draftID must be non-zero
|
||||
}
|
||||
if v < 0 {
|
||||
v = -v
|
||||
}
|
||||
return v
|
||||
}
|
||||
|
||||
// sendWithRetry sends a message through the channel with rate limiting and
|
||||
// retry logic. It classifies errors to determine the retry strategy:
|
||||
// - ErrNotRunning / ErrSendFailed: permanent, no retry
|
||||
|
|
@ -835,6 +931,16 @@ func (m *Manager) runTTLJanitor(ctx context.Context) {
|
|||
}
|
||||
return true
|
||||
})
|
||||
// Clean up stale edit-time entries (only needed for a few seconds,
|
||||
// but janitor runs infrequently so use a generous TTL).
|
||||
m.statusEditTimes.Range(func(key, value any) bool {
|
||||
if t, ok := value.(time.Time); ok {
|
||||
if now.Sub(t) > statusMsgTTL {
|
||||
m.statusEditTimes.Delete(key)
|
||||
}
|
||||
}
|
||||
return true
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1242,3 +1242,196 @@ func TestStatusMsgTTLJanitor(t *testing.T) {
|
|||
t.Fatal("expected fresh status entry to survive")
|
||||
}
|
||||
}
|
||||
|
||||
// --- DraftSender tests ---
|
||||
|
||||
// mockDraftSender implements DraftSender + MessageSenderWithID + MessageEditor.
|
||||
type mockDraftSender struct {
|
||||
mockChannel
|
||||
draftFn func(ctx context.Context, chatID string, draftID int, content string) error
|
||||
editFn func(ctx context.Context, chatID, messageID, content string) error
|
||||
sendWithID func(ctx context.Context, chatID, content string) (string, error)
|
||||
}
|
||||
|
||||
func (m *mockDraftSender) SendDraft(ctx context.Context, chatID string, draftID int, content string) error {
|
||||
return m.draftFn(ctx, chatID, draftID, content)
|
||||
}
|
||||
|
||||
func (m *mockDraftSender) EditMessage(ctx context.Context, chatID, messageID, content string) error {
|
||||
return m.editFn(ctx, chatID, messageID, content)
|
||||
}
|
||||
|
||||
func (m *mockDraftSender) SendWithID(ctx context.Context, chatID, content string) (string, error) {
|
||||
return m.sendWithID(ctx, chatID, content)
|
||||
}
|
||||
|
||||
func TestHandleStatusSend_UsesDraftSender(t *testing.T) {
|
||||
m := newTestManager()
|
||||
var draftCalled bool
|
||||
var draftContent string
|
||||
var draftDID int
|
||||
|
||||
ch := &mockDraftSender{
|
||||
mockChannel: mockChannel{
|
||||
sendFn: func(_ context.Context, _ bus.OutboundMessage) error { return nil },
|
||||
},
|
||||
draftFn: func(_ context.Context, chatID string, draftID int, content string) error {
|
||||
draftCalled = true
|
||||
draftContent = content
|
||||
draftDID = draftID
|
||||
return nil
|
||||
},
|
||||
editFn: func(_ context.Context, _, _, _ string) error {
|
||||
t.Fatal("EditMessage should not be called when draft succeeds")
|
||||
return nil
|
||||
},
|
||||
sendWithID: func(_ context.Context, _, _ string) (string, error) {
|
||||
t.Fatal("SendWithID should not be called when draft succeeds")
|
||||
return "", nil
|
||||
},
|
||||
}
|
||||
|
||||
w := &channelWorker{ch: ch, limiter: rate.NewLimiter(rate.Inf, 1)}
|
||||
|
||||
msg := bus.OutboundMessage{Channel: "test", ChatID: "123", Content: "streaming preview", IsStatus: true}
|
||||
m.handleStatusSend(context.Background(), "test", w, msg)
|
||||
|
||||
if !draftCalled {
|
||||
t.Fatal("expected SendDraft to be called")
|
||||
}
|
||||
if draftContent != "streaming preview" {
|
||||
t.Fatalf("expected draft content 'streaming preview', got %s", draftContent)
|
||||
}
|
||||
if draftDID == 0 {
|
||||
t.Fatal("expected non-zero draftID")
|
||||
}
|
||||
|
||||
// Second call should reuse the same draftID
|
||||
draftCalled = false
|
||||
var secondDID int
|
||||
ch.draftFn = func(_ context.Context, _ string, draftID int, _ string) error {
|
||||
draftCalled = true
|
||||
secondDID = draftID
|
||||
return nil
|
||||
}
|
||||
msg.Content = "streaming preview updated"
|
||||
m.handleStatusSend(context.Background(), "test", w, msg)
|
||||
|
||||
if !draftCalled {
|
||||
t.Fatal("expected SendDraft to be called again")
|
||||
}
|
||||
if secondDID != draftDID {
|
||||
t.Fatalf("expected same draftID %d, got %d", draftDID, secondDID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandleStatusSend_DraftFails_FallsToEdit(t *testing.T) {
|
||||
m := newTestManager()
|
||||
var editCalled bool
|
||||
|
||||
ch := &mockDraftSender{
|
||||
mockChannel: mockChannel{
|
||||
sendFn: func(_ context.Context, _ bus.OutboundMessage) error { return nil },
|
||||
},
|
||||
draftFn: func(_ context.Context, _ string, _ int, _ string) error {
|
||||
return fmt.Errorf("draft not supported in group")
|
||||
},
|
||||
editFn: func(_ context.Context, _, _, _ string) error {
|
||||
editCalled = true
|
||||
return nil
|
||||
},
|
||||
sendWithID: func(_ context.Context, _, _ string) (string, error) {
|
||||
return "msg-1", nil
|
||||
},
|
||||
}
|
||||
|
||||
w := &channelWorker{ch: ch, limiter: rate.NewLimiter(rate.Inf, 1)}
|
||||
|
||||
// No existing placeholder/status — draft fails, then SendWithID
|
||||
msg := bus.OutboundMessage{Channel: "test", ChatID: "123", Content: "preview", IsStatus: true}
|
||||
m.handleStatusSend(context.Background(), "test", w, msg)
|
||||
|
||||
// Draft failed, so it should fall through; no placeholder → no edit → SendWithID
|
||||
if editCalled {
|
||||
t.Fatal("expected EditMessage NOT to be called (no placeholder)")
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandleTaskStatusSend_UsesDraftSender(t *testing.T) {
|
||||
m := newTestManager()
|
||||
var draftCalled bool
|
||||
|
||||
ch := &mockDraftSender{
|
||||
mockChannel: mockChannel{
|
||||
sendFn: func(_ context.Context, _ bus.OutboundMessage) error { return nil },
|
||||
},
|
||||
draftFn: func(_ context.Context, _ string, _ int, _ string) error {
|
||||
draftCalled = true
|
||||
return nil
|
||||
},
|
||||
editFn: func(_ context.Context, _, _, _ string) error {
|
||||
t.Fatal("EditMessage should not be called when draft succeeds")
|
||||
return nil
|
||||
},
|
||||
sendWithID: func(_ context.Context, _, _ string) (string, error) {
|
||||
t.Fatal("SendWithID should not be called when draft succeeds")
|
||||
return "", nil
|
||||
},
|
||||
}
|
||||
|
||||
w := &channelWorker{ch: ch, limiter: rate.NewLimiter(rate.Inf, 1)}
|
||||
|
||||
msg := bus.OutboundMessage{
|
||||
Channel: "test",
|
||||
ChatID: "123",
|
||||
Content: "task progress 50%",
|
||||
IsTaskStatus: true,
|
||||
TaskID: "task-draft",
|
||||
}
|
||||
m.handleTaskStatusSend(context.Background(), "test", w, msg)
|
||||
|
||||
if !draftCalled {
|
||||
t.Fatal("expected SendDraft to be called for task status")
|
||||
}
|
||||
}
|
||||
|
||||
func TestPreSend_ClearsDraftState(t *testing.T) {
|
||||
m := newTestManager()
|
||||
|
||||
ch := &mockChannel{
|
||||
sendFn: func(_ context.Context, _ bus.OutboundMessage) error { return nil },
|
||||
}
|
||||
|
||||
// Store a draft-based status entry (draftID != 0, messageID empty)
|
||||
m.statusMsgIDs.Store("test:123", statusMsgEntry{draftID: 42, createdAt: time.Now()})
|
||||
|
||||
msg := bus.OutboundMessage{Channel: "test", ChatID: "123", Content: "final response"}
|
||||
edited := m.preSend(context.Background(), "test", msg, ch)
|
||||
|
||||
// Draft-based entries don't trigger edit; the final sendMessage replaces the draft
|
||||
if edited {
|
||||
t.Fatal("expected preSend to return false for draft-based status (sendMessage replaces draft)")
|
||||
}
|
||||
|
||||
// Verify draft state was consumed
|
||||
if _, loaded := m.statusMsgIDs.Load("test:123"); loaded {
|
||||
t.Fatal("expected draft status entry to be deleted after preSend")
|
||||
}
|
||||
}
|
||||
|
||||
func TestGenerateDraftID_Stable(t *testing.T) {
|
||||
id1 := generateDraftID("telegram:123")
|
||||
id2 := generateDraftID("telegram:123")
|
||||
if id1 != id2 {
|
||||
t.Fatalf("expected stable draft ID, got %d vs %d", id1, id2)
|
||||
}
|
||||
if id1 == 0 {
|
||||
t.Fatal("expected non-zero draft ID")
|
||||
}
|
||||
|
||||
// Different key should produce different ID
|
||||
id3 := generateDraftID("telegram:456")
|
||||
if id1 == id3 {
|
||||
t.Fatalf("expected different draft IDs for different keys, both got %d", id1)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -295,6 +295,33 @@ func (c *TelegramChannel) SendPlaceholder(ctx context.Context, chatID string) (s
|
|||
return fmt.Sprintf("%d", pMsg.MessageID), nil
|
||||
}
|
||||
|
||||
// SendDraft implements channels.DraftSender.
|
||||
// It uses Telegram Bot API's sendMessageDraft for progressive message streaming
|
||||
// without the "edited" indicator. Only works in private chats.
|
||||
func (c *TelegramChannel) SendDraft(ctx context.Context, chatID string, draftID int, content string) error {
|
||||
if !c.IsRunning() {
|
||||
return channels.ErrNotRunning
|
||||
}
|
||||
cid, err := parseChatID(chatID)
|
||||
if err != nil {
|
||||
return fmt.Errorf("invalid chat ID %s: %w", chatID, channels.ErrSendFailed)
|
||||
}
|
||||
htmlContent := markdownToTelegramHTML(content)
|
||||
params := &telego.SendMessageDraftParams{
|
||||
ChatID: cid,
|
||||
DraftID: draftID,
|
||||
Text: htmlContent,
|
||||
ParseMode: telego.ModeHTML,
|
||||
}
|
||||
if err = c.bot.SendMessageDraft(ctx, params); err != nil {
|
||||
// HTML parse failure — retry as plain text
|
||||
params.ParseMode = ""
|
||||
params.Text = content
|
||||
return c.bot.SendMessageDraft(ctx, params)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// SendMedia implements the channels.MediaSender interface.
|
||||
func (c *TelegramChannel) SendMedia(ctx context.Context, msg bus.OutboundMediaMessage) error {
|
||||
if !c.IsRunning() {
|
||||
|
|
|
|||
|
|
@ -7,7 +7,6 @@
|
|||
package heartbeat
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
|
|
@ -21,7 +20,6 @@ import (
|
|||
"github.com/sipeed/picoclaw/pkg/logger"
|
||||
"github.com/sipeed/picoclaw/pkg/state"
|
||||
"github.com/sipeed/picoclaw/pkg/tools"
|
||||
"github.com/sipeed/picoclaw/pkg/utils"
|
||||
)
|
||||
|
||||
const (
|
||||
|
|
@ -228,12 +226,8 @@ func (hs *HeartbeatService) executeHeartbeat() {
|
|||
return
|
||||
}
|
||||
|
||||
// Send result to user
|
||||
if result.ForUser != "" {
|
||||
hs.sendResponse(result.ForUser)
|
||||
} else if result.ForLLM != "" {
|
||||
hs.sendResponse(result.ForLLM)
|
||||
}
|
||||
// Skip sendResponse — the task completion message in runAgentLoop already
|
||||
// includes the LLM response, so sending here would create a duplicate bubble.
|
||||
|
||||
hs.mu.Lock()
|
||||
hs.lastNotifiedAt = time.Now()
|
||||
|
|
@ -309,45 +303,6 @@ Add your heartbeat tasks below this line:
|
|||
}
|
||||
}
|
||||
|
||||
// sendResponse sends the heartbeat response to the last channel.
|
||||
// Think blocks are stripped as a safety net to prevent LLM reasoning
|
||||
// artifacts from leaking into user-facing messages.
|
||||
func (hs *HeartbeatService) sendResponse(response string) {
|
||||
response = utils.StripThinkBlocks(response)
|
||||
hs.mu.RLock()
|
||||
msgBus := hs.bus
|
||||
hs.mu.RUnlock()
|
||||
|
||||
if msgBus == nil {
|
||||
hs.logInfof("No message bus configured, heartbeat result not sent")
|
||||
return
|
||||
}
|
||||
|
||||
// Get last channel from state
|
||||
lastChannel := hs.state.GetLastChannel()
|
||||
if lastChannel == "" {
|
||||
hs.logInfof("No last channel recorded, heartbeat result not sent")
|
||||
return
|
||||
}
|
||||
|
||||
platform, userID := hs.parseLastChannel(lastChannel)
|
||||
|
||||
// Skip internal channels that can't receive messages
|
||||
if platform == "" || userID == "" {
|
||||
return
|
||||
}
|
||||
|
||||
pubCtx, pubCancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer pubCancel()
|
||||
msgBus.PublishOutbound(pubCtx, bus.OutboundMessage{
|
||||
Channel: platform,
|
||||
ChatID: userID,
|
||||
Content: response,
|
||||
})
|
||||
|
||||
hs.logInfof("Heartbeat result sent to %s", platform)
|
||||
}
|
||||
|
||||
// parseLastChannel parses the last channel string into platform and userID.
|
||||
// Returns empty strings for invalid or internal channels.
|
||||
func (hs *HeartbeatService) parseLastChannel(lastChannel string) (platform, userID string) {
|
||||
|
|
|
|||
|
|
@ -200,6 +200,42 @@ func TestLogPath(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
// TestExecuteHeartbeat_NoSendResponse verifies that heartbeat results
|
||||
// do not trigger sendResponse (dedup: response is included in task status instead).
|
||||
func TestExecuteHeartbeat_NoSendResponse(t *testing.T) {
|
||||
tmpDir, err := os.MkdirTemp("", "heartbeat-test-*")
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create temp dir: %v", err)
|
||||
}
|
||||
defer os.RemoveAll(tmpDir)
|
||||
|
||||
hs := NewHeartbeatService(tmpDir, 30, true)
|
||||
hs.stopChan = make(chan struct{})
|
||||
|
||||
hs.SetHandler(func(prompt, channel, chatID string) *tools.ToolResult {
|
||||
return &tools.ToolResult{
|
||||
ForUser: "Task result for user",
|
||||
ForLLM: "Task result for LLM",
|
||||
Silent: false,
|
||||
IsError: false,
|
||||
Async: false,
|
||||
}
|
||||
})
|
||||
|
||||
os.WriteFile(filepath.Join(tmpDir, "HEARTBEAT.md"), []byte("Test task"), 0o644)
|
||||
|
||||
// Execute heartbeat — since bus is nil, sendResponse would log but not crash.
|
||||
// The key assertion is that lastNotifiedAt is still updated (flow reaches end).
|
||||
hs.executeHeartbeat()
|
||||
|
||||
hs.mu.RLock()
|
||||
notified := !hs.lastNotifiedAt.IsZero()
|
||||
hs.mu.RUnlock()
|
||||
if !notified {
|
||||
t.Error("Expected lastNotifiedAt to be set after heartbeat completion")
|
||||
}
|
||||
}
|
||||
|
||||
// TestHeartbeatFilePath verifies HEARTBEAT.md is at workspace root
|
||||
func TestHeartbeatFilePath(t *testing.T) {
|
||||
tmpDir, err := os.MkdirTemp("", "heartbeat-test-*")
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue