From 0cea7ed32f028d1745b81aa164025e8ea0beef96 Mon Sep 17 00:00:00 2001
From: dj-oyu <68707227+dj-oyu@users.noreply.github.com>
Date: Tue, 3 Mar 2026 11:40:06 +0900
Subject: [PATCH] fix: redirect message tool to task status bubble for
 background tasks

Root cause: LLM calls the `message` tool during heartbeat execution,
which publishes directly to the bus via sendCallback, bypassing the
SendResponse=false guard. This creates a second bubble separate from
the task status.

Fix: For background tasks with SendResponse=false (e.g. heartbeat),
replace the message tool's sendCallback to publish as IsTaskStatus
with the same TaskID. This routes message tool output into the same
bubble as streaming preview and task completion.

Also removes the now-unnecessary PromoteStatusToTask call since
streaming preview already uses IsTaskStatus directly.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 pkg/agent/loop.go | 34 +++++++++++++++++++++++-----------
 1 file changed, 23 insertions(+), 11 deletions(-)

diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go
index 62012b461..56773f3dd 100644
--- a/pkg/agent/loop.go
+++ b/pkg/agent/loop.go
@@ -1127,6 +1127,28 @@ func (al *AgentLoop) runAgentLoop(ctx context.Context, agent *AgentInstance, opt
 	// 1. Update tool contexts
 	al.updateToolContexts(agent, opts.Channel, opts.ChatID)
 
+	// 1-bis. For background tasks that don't send a final response (e.g. heartbeat),
+	// redirect the message tool to publish as IsTaskStatus so its output lands in
+	// the same bubble as the task status instead of creating a separate message.
+	if opts.Background && !opts.SendResponse && opts.TaskID != "" {
+		if tool, ok := agent.Tools.Get("message"); ok {
+			if mt, ok := tool.(*tools.MessageTool); ok {
+				taskID := opts.TaskID
+				mt.SetSendCallback(func(channel, chatID, content string) error {
+					pubCtx, pubCancel := context.WithTimeout(context.Background(), 5*time.Second)
+					defer pubCancel()
+					return al.bus.PublishOutbound(pubCtx, bus.OutboundMessage{
+						Channel:      channel,
+						ChatID:       chatID,
+						Content:      content,
+						IsTaskStatus: true,
+						TaskID:       taskID,
+					})
+				})
+			}
+		}
+	}
+
 	// 1a. Set session-specific working directory for bootstrap file lookup.
 	// Prefer the tool-detected project directory (touch_dir) from the session tracker,
 	// resolved as an absolute path under workspace. Fall back to worktree or workspace.
@@ -1391,17 +1413,7 @@ func (al *AgentLoop) runAgentLoop(ctx context.Context, agent *AgentInstance, opt
 		task.Result = utils.Truncate(finalContent, 280)
 	}
 
-	// 5e. Promote streaming status bubble to task message for background tasks.
-	// When SendResponse is false (e.g. heartbeat), no final non-status message
-	// triggers preSend's statusMsgIDs.LoadAndDelete cleanup, so the last
-	// streaming chunk persists on Telegram. Move the tracked status message ID
-	// into taskMsgIDs so the defer's IsTaskStatus completion notification edits
-	// the existing bubble instead of creating a duplicate message.
-	if opts.Background && !opts.SendResponse && !constants.IsInternalChannel(opts.Channel) && task != nil &&
-		task.streamedChunks && opts.TaskID != "" {
-		statusKey := opts.Channel + ":" + opts.ChatID
-		al.channelManager.PromoteStatusToTask(statusKey, opts.TaskID)
-	}
+
 
 	// 6. Save final assistant message to session (deferred write-behind)
 	agent.Sessions.AddMessage(opts.SessionKey, "assistant", finalContent)