fix(channels): preserve tool feedback progress state

fix(pico): preserve context usage when finalizing tool feedback

chore: record branch review pass

fix: preserve tool feedback finalization state

fix(web): handle pico history update fallback
This commit is contained in:
lxowalle 2026-04-22 18:02:25 +08:00
parent 669bf10cb5
commit 4886f28a3c
8 changed files with 460 additions and 27 deletions

View file

@ -892,6 +892,77 @@ func TestAgentLoop_HookRespond_BusFallback(t *testing.T) {
}
}
func TestAgentLoop_HookRespond_ResponseHandledMediaPreservesOutboundContext(t *testing.T) {
provider := &multiToolProvider{
toolCalls: []providers.ToolCall{
{ID: "call-1", Name: "media_tool", Arguments: map[string]any{}},
},
finalContent: "done",
}
al, agent, cleanup := newHookTestLoop(t, provider)
defer cleanup()
hook := &respondWithMediaHook{
respondTools: map[string]bool{"media_tool": true},
media: []string{"media://test/image.png"},
responseHandled: true,
forLLM: "media sent successfully",
}
if err := al.MountHook(NamedHook("media-hook", hook)); err != nil {
t.Fatalf("MountHook failed: %v", err)
}
telegramChannel := &fakeMediaChannel{fakeChannel: fakeChannel{id: "rid-telegram"}}
al.channelManager = newStartedTestChannelManager(t,
al.bus.(*bus.MessageBus), al.mediaStore, "telegram", telegramChannel)
_, err := al.runAgentLoop(context.Background(), agent, processOptions{
Dispatch: DispatchRequest{
SessionKey: "session-topic-media",
SessionScope: &session.SessionScope{
Version: session.ScopeVersionV1,
AgentID: agent.ID,
Channel: "telegram",
Dimensions: []string{"chat"},
Values: map[string]string{
"chat": "forum:-100123/42",
},
},
InboundContext: &bus.InboundContext{
Channel: "telegram",
ChatID: "-100123",
TopicID: "42",
ChatType: "group",
SenderID: "user1",
},
UserMessage: "send media",
},
DefaultResponse: defaultResponse,
EnableSummary: false,
SendResponse: false,
})
if err != nil {
t.Fatalf("runAgentLoop failed: %v", err)
}
if len(telegramChannel.sentMedia) != 1 {
t.Fatalf("expected exactly 1 sent media message, got %d", len(telegramChannel.sentMedia))
}
sent := telegramChannel.sentMedia[0]
if sent.Context.Channel != "telegram" || sent.Context.ChatID != "-100123" || sent.Context.TopicID != "42" {
t.Fatalf("unexpected media context: %+v", sent.Context)
}
if sent.AgentID != agent.ID {
t.Fatalf("sent media agent_id = %q, want %q", sent.AgentID, agent.ID)
}
if sent.SessionKey != "session-topic-media" {
t.Fatalf("sent media session_key = %q, want session-topic-media", sent.SessionKey)
}
if sent.Scope == nil || sent.Scope.Values["chat"] != "forum:-100123/42" {
t.Fatalf("unexpected sent media scope: %+v", sent.Scope)
}
}
type multiToolProvider struct {
mu sync.Mutex
callCount int

View file

@ -126,6 +126,15 @@ toolLoop:
outboundMedia := bus.OutboundMediaMessage{
Channel: ts.channel,
ChatID: ts.chatID,
Context: outboundContextFromInbound(
ts.opts.Dispatch.InboundContext,
ts.channel,
ts.chatID,
ts.opts.Dispatch.ReplyToMessageID(),
),
AgentID: ts.agent.ID,
SessionKey: ts.sessionKey,
Scope: outboundScopeFromSessionScope(ts.opts.Dispatch.SessionScope),
Parts: parts,
}
if al.channelManager != nil && ts.channel != "" && !constants.IsInternalChannel(ts.channel) {

View file

@ -111,6 +111,10 @@ type toolFeedbackMessageTargetResolver interface {
ToolFeedbackMessageChatID(chatID string, outboundCtx *bus.InboundContext) string
}
type toolFeedbackMessageContentPreparer interface {
PrepareToolFeedbackMessageContent(content string) string
}
type asyncTask struct {
cancel context.CancelFunc
}
@ -166,6 +170,19 @@ func dismissTrackedToolFeedbackMessage(
}
}
func prepareToolFeedbackMessageContent(ch Channel, content string) string {
prepared := strings.TrimSpace(content)
if prepared == "" {
return ""
}
if preparer, ok := ch.(toolFeedbackMessageContentPreparer); ok {
if candidate := strings.TrimSpace(preparer.PrepareToolFeedbackMessageContent(prepared)); candidate != "" {
return candidate
}
}
return prepared
}
// RecordPlaceholder registers a placeholder message for later editing.
// Implements PlaceholderRecorder.
func (m *Manager) RecordPlaceholder(channel, chatID, placeholderID string) {
@ -281,13 +298,15 @@ func (m *Manager) preSend(ctx context.Context, name string, msg bus.OutboundMess
if entry, ok := v.(placeholderEntry); ok && entry.id != "" {
if editor, ok := ch.(MessageEditor); ok {
content := msg.Content
trackedContent := msg.Content
if isToolFeedback {
content = InitialAnimatedToolFeedbackContent(msg.Content)
trackedContent = prepareToolFeedbackMessageContent(ch, msg.Content)
content = InitialAnimatedToolFeedbackContent(trackedContent)
}
if err := editor.EditMessage(ctx, chatID, entry.id, content); err == nil {
trackedChatID := trackedToolFeedbackMessageChatID(ch, chatID, &msg.Context)
if tracker, ok := ch.(toolFeedbackMessageTracker); ok && isToolFeedback {
tracker.RecordToolFeedbackMessage(trackedChatID, entry.id, msg.Content)
tracker.RecordToolFeedbackMessage(trackedChatID, entry.id, trackedContent)
} else if !isToolFeedback {
dismissTrackedToolFeedbackMessage(ctx, ch, chatID, &msg.Context)
}
@ -389,7 +408,15 @@ func (m *Manager) GetStreamer(ctx context.Context, channelName, chatID string) (
return &finalizeHookStreamer{
Streamer: streamer,
onFinalize: func(finalizeCtx context.Context) {
dismissTrackedToolFeedbackMessage(finalizeCtx, ch, chatID, nil)
dismissTrackedToolFeedbackMessage(
finalizeCtx,
ch,
chatID,
&bus.InboundContext{
Channel: channelName,
ChatID: chatID,
},
)
m.streamActive.Store(key, true)
},
}, true

View file

@ -119,6 +119,7 @@ func (m *mockStreamer) Cancel(context.Context) {}
type mockStreamingChannel struct {
mockMessageEditor
streamer Streamer
resolveChatIDFn func(chatID string, outboundCtx *bus.InboundContext) string
}
func (m *mockStreamingChannel) BeginStream(context.Context, string) (Streamer, error) {
@ -128,6 +129,16 @@ func (m *mockStreamingChannel) BeginStream(context.Context, string) (Streamer, e
return m.streamer, nil
}
func (m *mockStreamingChannel) ToolFeedbackMessageChatID(
chatID string,
outboundCtx *bus.InboundContext,
) string {
if m.resolveChatIDFn != nil {
return m.resolveChatIDFn(chatID, outboundCtx)
}
return chatID
}
// newTestManager creates a minimal Manager suitable for unit tests.
func newTestManager() *Manager {
return &Manager{
@ -754,6 +765,7 @@ type mockMessageEditor struct {
finalizeCalled bool
recordedChatID string
recordedMessageID string
recordedContent string
clearedChatID string
dismissedChatID string
}
@ -762,9 +774,10 @@ func (m *mockMessageEditor) EditMessage(ctx context.Context, chatID, messageID,
return m.editFn(ctx, chatID, messageID, content)
}
func (m *mockMessageEditor) RecordToolFeedbackMessage(chatID, messageID, _ string) {
func (m *mockMessageEditor) RecordToolFeedbackMessage(chatID, messageID, content string) {
m.recordedChatID = chatID
m.recordedMessageID = messageID
m.recordedContent = content
}
func (m *mockMessageEditor) ClearToolFeedbackMessage(chatID string) {
@ -801,6 +814,18 @@ func (m *mockResolvedToolFeedbackEditor) ToolFeedbackMessageChatID(
return chatID
}
type mockPreparedToolFeedbackEditor struct {
mockMessageEditor
prepareFn func(content string) string
}
func (m *mockPreparedToolFeedbackEditor) PrepareToolFeedbackMessageContent(content string) string {
if m.prepareFn != nil {
return m.prepareFn(content)
}
return content
}
func TestPreSend_PlaceholderEditSuccess(t *testing.T) {
m := newTestManager()
var sendCalled bool
@ -928,6 +953,56 @@ func TestPreSend_ToolFeedbackPlaceholderEditUsesResolvedTrackedChatID(t *testing
}
}
func TestPreSend_ToolFeedbackPlaceholderEditUsesPreparedContent(t *testing.T) {
m := newTestManager()
const rawContent = "🔧 `read_file`\n" + "<raw>"
const preparedContent = "🔧 `read_file`\n&lt;raw&gt;"
ch := &mockPreparedToolFeedbackEditor{
mockMessageEditor: mockMessageEditor{
editFn: func(_ context.Context, chatID, messageID, content string) error {
if chatID != "123" || messageID != "456" {
t.Fatalf("unexpected edit target: %s/%s", chatID, messageID)
}
if content != InitialAnimatedToolFeedbackContent(preparedContent) {
t.Fatalf("unexpected prepared content: %q", content)
}
return nil
},
},
prepareFn: func(content string) string {
if content != rawContent {
t.Fatalf("unexpected raw tool feedback: %q", content)
}
return preparedContent
},
}
m.RecordPlaceholder("test", "123", "456")
msg := testOutboundMessage(bus.OutboundMessage{
Channel: "test",
ChatID: "123",
Content: rawContent,
Context: bus.InboundContext{
Channel: "test",
ChatID: "123",
Raw: map[string]string{
"message_kind": "tool_feedback",
},
},
})
_, edited := m.preSend(context.Background(), "test", msg, ch)
if !edited {
t.Fatal("expected preSend to edit placeholder")
}
if ch.recordedContent != preparedContent {
t.Fatalf("expected tracked content %q, got %q", preparedContent, ch.recordedContent)
}
}
func TestPreSend_NonToolFeedbackLeavesTrackedMessageForChannelSend(t *testing.T) {
m := newTestManager()
ch := &mockMessageEditor{}
@ -1157,6 +1232,45 @@ func TestGetStreamer_FinalizeDismissesTrackedToolFeedback(t *testing.T) {
}
}
func TestGetStreamer_FinalizeDismissesResolvedTrackedToolFeedback(t *testing.T) {
m := newTestManager()
ch := &mockStreamingChannel{
mockMessageEditor: mockMessageEditor{},
streamer: &mockStreamer{
finalizeFn: func(_ context.Context, content string) error {
if content != "final reply" {
t.Fatalf("unexpected finalize content: %q", content)
}
return nil
},
},
resolveChatIDFn: func(chatID string, outboundCtx *bus.InboundContext) string {
if outboundCtx == nil {
t.Fatal("expected outbound context during stream finalize")
}
if outboundCtx.ChatID != "-100123/42" {
t.Fatalf("unexpected outbound context: %+v", outboundCtx)
}
return outboundCtx.ChatID
},
}
m.channels["test"] = ch
streamer, ok := m.GetStreamer(context.Background(), "test", "-100123/42")
if !ok {
t.Fatal("expected streamer to be available")
}
if err := streamer.Finalize(context.Background(), "final reply"); err != nil {
t.Fatalf("Finalize() error = %v", err)
}
if ch.dismissedChatID != "-100123/42" {
t.Fatalf("expected resolved tracked tool feedback dismissal, got %q", ch.dismissedChatID)
}
if _, ok := m.streamActive.Load("test:-100123/42"); !ok {
t.Fatal("expected streamActive marker to be recorded after finalize")
}
}
func TestPreSend_PlaceholderEditSuccessDismissesResolvedTrackedToolFeedback(t *testing.T) {
m := newTestManager()

View file

@ -57,6 +57,10 @@ func outboundMessageIsToolFeedback(msg bus.OutboundMessage) bool {
return strings.EqualFold(strings.TrimSpace(msg.Context.Raw["message_kind"]), "tool_feedback")
}
func outboundMessageFinalizesTrackedToolFeedback(msg bus.OutboundMessage) bool {
return !outboundMessageIsToolFeedback(msg) && !outboundMessageIsThought(msg)
}
// writeJSON sends a JSON message to the connection with write locking.
func (pc *picoConn) writeJSON(v any) error {
if pc.closed.Load() {
@ -294,7 +298,7 @@ func (c *PicoChannel) Send(ctx context.Context, msg bus.OutboundMessage) ([]stri
}
}
trackedMsgID, hasTrackedMsg := c.currentToolFeedbackMessage(msg.ChatID)
if !isToolFeedback {
if outboundMessageFinalizesTrackedToolFeedback(msg) {
if msgIDs, handled := c.FinalizeToolFeedbackMessage(ctx, msg); handled {
return msgIDs, nil
}
@ -319,7 +323,7 @@ func (c *PicoChannel) Send(ctx context.Context, msg bus.OutboundMessage) ([]stri
}
if isToolFeedback {
c.RecordToolFeedbackMessage(msg.ChatID, msgID, msg.Content)
} else if hasTrackedMsg {
} else if hasTrackedMsg && outboundMessageFinalizesTrackedToolFeedback(msg) {
c.dismissTrackedToolFeedbackMessage(ctx, msg.ChatID, trackedMsgID)
}
return []string{msgID}, nil
@ -327,11 +331,7 @@ func (c *PicoChannel) Send(ctx context.Context, msg bus.OutboundMessage) ([]stri
// EditMessage implements channels.MessageEditor.
func (c *PicoChannel) EditMessage(ctx context.Context, chatID string, messageID string, content string) error {
outMsg := newMessage(TypeMessageUpdate, map[string]any{
"message_id": messageID,
"content": content,
})
return c.broadcastToSession(chatID, outMsg)
return c.editMessage(ctx, chatID, messageID, content, nil)
}
// DeleteMessage implements channels.MessageDeleter.
@ -394,13 +394,14 @@ func (c *PicoChannel) finalizeTrackedToolFeedbackMessage(
ctx context.Context,
chatID string,
content string,
editFn func(context.Context, string, string, string) error,
editFn func(context.Context, string, string, string, *bus.ContextUsage) error,
contextUsage *bus.ContextUsage,
) ([]string, bool) {
msgID, baseContent, ok := c.takeToolFeedbackMessage(chatID)
if !ok || editFn == nil {
return nil, false
}
if err := editFn(ctx, chatID, msgID, content); err != nil {
if err := editFn(ctx, chatID, msgID, content, contextUsage); err != nil {
c.RecordToolFeedbackMessage(chatID, msgID, baseContent)
return nil, false
}
@ -408,10 +409,10 @@ func (c *PicoChannel) finalizeTrackedToolFeedbackMessage(
}
func (c *PicoChannel) FinalizeToolFeedbackMessage(ctx context.Context, msg bus.OutboundMessage) ([]string, bool) {
if outboundMessageIsToolFeedback(msg) {
if !outboundMessageFinalizesTrackedToolFeedback(msg) {
return nil, false
}
return c.finalizeTrackedToolFeedbackMessage(ctx, msg.ChatID, msg.Content, c.EditMessage)
return c.finalizeTrackedToolFeedbackMessage(ctx, msg.ChatID, msg.Content, c.editMessage, msg.ContextUsage)
}
// StartTyping implements channels.TypingCapable.
@ -1068,3 +1069,19 @@ func setContextUsagePayload(payload map[string]any, u *bus.ContextUsage) {
"used_percent": u.UsedPercent,
}
}
func (c *PicoChannel) editMessage(
ctx context.Context,
chatID string,
messageID string,
content string,
contextUsage *bus.ContextUsage,
) error {
payload := map[string]any{
"message_id": messageID,
"content": content,
}
setContextUsagePayload(payload, contextUsage)
outMsg := newMessage(TypeMessageUpdate, payload)
return c.broadcastToSession(chatID, outMsg)
}

View file

@ -46,15 +46,19 @@ func TestFinalizeTrackedToolFeedbackMessage_StopsTrackingBeforeEdit(t *testing.T
context.Background(),
"pico:chat-1",
"final reply",
func(_ context.Context, chatID, messageID, content string) error {
func(_ context.Context, chatID, messageID, content string, contextUsage *bus.ContextUsage) error {
if _, ok := ch.currentToolFeedbackMessage(chatID); ok {
t.Fatal("expected tracked tool feedback to be stopped before edit")
}
if chatID != "pico:chat-1" || messageID != "msg-1" || content != "final reply" {
t.Fatalf("unexpected edit args: %s %s %s", chatID, messageID, content)
}
if contextUsage != nil {
t.Fatalf("unexpected context usage: %+v", contextUsage)
}
return nil
},
nil,
)
if !handled {
t.Fatal("expected finalizeTrackedToolFeedbackMessage to handle tracked message")
@ -90,6 +94,105 @@ func TestDismissTrackedToolFeedbackMessage_DeletesProgressMessage(t *testing.T)
}
}
func TestSend_ThoughtMessageDoesNotFinalizeTrackedToolFeedback(t *testing.T) {
ch := newTestPicoChannel(t)
if err := ch.Start(context.Background()); err != nil {
t.Fatalf("Start() error = %v", err)
}
defer ch.Stop(context.Background())
clientConn, received, cleanup := newTestPicoWebSocket(t)
defer cleanup()
ch.addConnForTest(&picoConn{id: "conn-1", conn: clientConn, sessionID: "sess-1"})
ch.RecordToolFeedbackMessage("pico:sess-1", "msg-progress", "🔧 `read_file`\nReading config")
if _, err := ch.Send(context.Background(), bus.OutboundMessage{
ChatID: "pico:sess-1",
Content: "thinking trace",
Context: bus.InboundContext{
Channel: "pico",
ChatID: "pico:sess-1",
Raw: map[string]string{
"message_kind": MessageKindThought,
},
},
}); err != nil {
t.Fatalf("Send(thought) error = %v", err)
}
select {
case msg := <-received:
if msg.Type != TypeMessageCreate {
t.Fatalf("thought message type = %q, want %q", msg.Type, TypeMessageCreate)
}
payload := msg.Payload
if got := payload[PayloadKeyContent]; got != "thinking trace" {
t.Fatalf("thought content = %#v, want %q", got, "thinking trace")
}
if got := payload[PayloadKeyThought]; got != true {
t.Fatalf("thought flag = %#v, want true", got)
}
if got := payload["message_id"]; got == "msg-progress" || got == nil || got == "" {
t.Fatalf("thought message_id = %#v, want new non-progress id", got)
}
case <-time.After(time.Second):
t.Fatal("expected thought message to be delivered")
}
if msgID, ok := ch.currentToolFeedbackMessage("pico:sess-1"); !ok || msgID != "msg-progress" {
t.Fatalf("tracked tool feedback = (%q, %v), want (msg-progress, true)", msgID, ok)
}
if _, err := ch.Send(context.Background(), bus.OutboundMessage{
ChatID: "pico:sess-1",
Content: "final reply",
Context: bus.InboundContext{
Channel: "pico",
ChatID: "pico:sess-1",
},
ContextUsage: &bus.ContextUsage{
UsedTokens: 321,
TotalTokens: 4096,
CompressAtTokens: 3072,
UsedPercent: 8,
},
}); err != nil {
t.Fatalf("Send(final) error = %v", err)
}
select {
case msg := <-received:
if msg.Type != TypeMessageUpdate {
t.Fatalf("final message type = %q, want %q", msg.Type, TypeMessageUpdate)
}
payload := msg.Payload
if got := payload["message_id"]; got != "msg-progress" {
t.Fatalf("final message_id = %#v, want %q", got, "msg-progress")
}
if got := payload[PayloadKeyContent]; got != "final reply" {
t.Fatalf("final content = %#v, want %q", got, "final reply")
}
rawUsage, ok := payload["context_usage"].(map[string]any)
if !ok {
t.Fatalf("final context_usage = %#v, want map payload", payload["context_usage"])
}
if got, ok := rawUsage["used_tokens"].(float64); !ok || got != 321 {
t.Fatalf("used_tokens = %#v, want 321", rawUsage["used_tokens"])
}
if got, ok := rawUsage["total_tokens"].(float64); !ok || got != 4096 {
t.Fatalf("total_tokens = %#v, want 4096", rawUsage["total_tokens"])
}
case <-time.After(time.Second):
t.Fatal("expected final reply to finalize tracked tool feedback")
}
if _, ok := ch.currentToolFeedbackMessage("pico:sess-1"); ok {
t.Fatal("expected tracked tool feedback to be cleared after final reply")
}
}
func TestCreateAndAddConnection_RespectsMaxConnectionsConcurrently(t *testing.T) {
ch := newTestPicoChannel(t)

View file

@ -1115,6 +1115,13 @@ func fitToolFeedbackForTelegram(content string, useMarkdownV2 bool, maxParsedLen
return best
}
func (c *TelegramChannel) PrepareToolFeedbackMessageContent(content string) string {
if c == nil || c.tgCfg == nil {
return strings.TrimSpace(content)
}
return fitToolFeedbackForTelegram(content, c.tgCfg.UseMarkdownV2, 4096)
}
func telegramToolFeedbackChatKey(chatID string, outboundCtx *bus.InboundContext) string {
resolvedChatID, threadID, err := resolveTelegramOutboundTarget(chatID, outboundCtx)
if err != nil || threadID == 0 {

View file

@ -4,6 +4,7 @@ import { normalizeUnixTimestamp } from "@/features/chat/state"
import {
type AssistantMessageKind,
type ChatAttachment,
type ChatMessage,
type ContextUsage,
updateChatStore,
} from "@/store/chat"
@ -90,6 +91,35 @@ function parseContextUsage(
}
}
function isToolFeedbackMessage(message: ChatMessage): boolean {
if (message.role !== "assistant") {
return false
}
const firstLine = message.content.split("\n", 1)[0]?.trim() ?? ""
return /^🔧\s+`[^`]+`/.test(firstLine)
}
function findToolFeedbackMessageIndex(messages: ChatMessage[]): number {
let lastUserIndex = -1
for (let i = messages.length - 1; i >= 0; i -= 1) {
if (messages[i].role === "user") {
lastUserIndex = i
break
}
}
for (let i = messages.length - 1; i >= 0; i -= 1) {
if (i <= lastUserIndex) {
break
}
if (isToolFeedbackMessage(messages[i])) {
return i
}
}
return -1
}
export function handlePicoMessage(
message: PicoMessage,
expectedSessionId: string,
@ -138,21 +168,64 @@ export function handlePicoMessage(
const hasKind = hasAssistantKindPayload(payload)
const kind = parseAssistantMessageKind(payload)
const attachments = parseAttachments(payload)
const contextUsage = parseContextUsage(payload)
const timestamp =
message.timestamp !== undefined &&
Number.isFinite(Number(message.timestamp))
? normalizeUnixTimestamp(Number(message.timestamp))
: Date.now()
if (!messageId) {
break
}
updateChatStore((prev) => ({
messages: prev.messages.map((msg) =>
msg.id === messageId
? {
...msg,
content,
...(hasKind ? { kind } : {}),
...(attachments ? { attachments } : {}),
}
: msg,
),
messages: (() => {
let found = false
const messages = prev.messages.map((msg) => {
if (msg.id !== messageId) {
return msg
}
found = true
return {
...msg,
id: messageId,
content,
...(hasKind ? { kind } : {}),
...(attachments ? { attachments } : {}),
}
})
if (found) {
return messages
}
const fallbackIndex = findToolFeedbackMessageIndex(messages)
if (fallbackIndex >= 0) {
return messages.map((msg, index) =>
index === fallbackIndex
? {
...msg,
id: messageId,
content,
...(hasKind ? { kind } : {}),
...(attachments ? { attachments } : {}),
}
: msg,
)
}
return [
...messages,
{
id: messageId,
role: "assistant" as const,
content,
...(hasKind ? { kind } : {}),
...(attachments ? { attachments } : {}),
timestamp,
},
]
})(),
...(contextUsage ? { contextUsage } : {}),
}))
break
}
@ -164,7 +237,19 @@ export function handlePicoMessage(
}
updateChatStore((prev) => ({
messages: prev.messages.filter((msg) => msg.id !== messageId),
messages: (() => {
const exactMessages = prev.messages.filter((msg) => msg.id !== messageId)
if (exactMessages.length !== prev.messages.length) {
return exactMessages
}
const fallbackIndex = findToolFeedbackMessageIndex(prev.messages)
if (fallbackIndex < 0) {
return prev.messages
}
return prev.messages.filter((_, index) => index !== fallbackIndex)
})(),
}))
break
}