From 4886f28a3cf5dc8fad9e8de3e41cf1b08ed1277e Mon Sep 17 00:00:00 2001 From: lxowalle Date: Wed, 22 Apr 2026 18:02:25 +0800 Subject: [PATCH] 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 --- pkg/agent/hooks_test.go | 71 +++++++++++++ pkg/agent/pipeline_execute.go | 9 ++ pkg/channels/manager.go | 33 +++++- pkg/channels/manager_test.go | 116 ++++++++++++++++++++- pkg/channels/pico/pico.go | 39 +++++-- pkg/channels/pico/pico_test.go | 105 ++++++++++++++++++- pkg/channels/telegram/telegram.go | 7 ++ web/frontend/src/features/chat/protocol.ts | 107 +++++++++++++++++-- 8 files changed, 460 insertions(+), 27 deletions(-) diff --git a/pkg/agent/hooks_test.go b/pkg/agent/hooks_test.go index bccc6bba4..1cfa341a7 100644 --- a/pkg/agent/hooks_test.go +++ b/pkg/agent/hooks_test.go @@ -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 diff --git a/pkg/agent/pipeline_execute.go b/pkg/agent/pipeline_execute.go index 1043d06a1..e1e4ef04a 100644 --- a/pkg/agent/pipeline_execute.go +++ b/pkg/agent/pipeline_execute.go @@ -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) { diff --git a/pkg/channels/manager.go b/pkg/channels/manager.go index b2df35414..2ffb1bb10 100644 --- a/pkg/channels/manager.go +++ b/pkg/channels/manager.go @@ -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 diff --git a/pkg/channels/manager_test.go b/pkg/channels/manager_test.go index 31e77704d..09f589183 100644 --- a/pkg/channels/manager_test.go +++ b/pkg/channels/manager_test.go @@ -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" + "" + const preparedContent = "🔧 `read_file`\n<raw>" + + 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() diff --git a/pkg/channels/pico/pico.go b/pkg/channels/pico/pico.go index b7361af4c..31360b3de 100644 --- a/pkg/channels/pico/pico.go +++ b/pkg/channels/pico/pico.go @@ -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) +} diff --git a/pkg/channels/pico/pico_test.go b/pkg/channels/pico/pico_test.go index 05bdfb2be..22ed5451a 100644 --- a/pkg/channels/pico/pico_test.go +++ b/pkg/channels/pico/pico_test.go @@ -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) diff --git a/pkg/channels/telegram/telegram.go b/pkg/channels/telegram/telegram.go index cc148fd5b..cebebfed6 100644 --- a/pkg/channels/telegram/telegram.go +++ b/pkg/channels/telegram/telegram.go @@ -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 { diff --git a/web/frontend/src/features/chat/protocol.ts b/web/frontend/src/features/chat/protocol.ts index de74b332a..3c4259014 100644 --- a/web/frontend/src/features/chat/protocol.ts +++ b/web/frontend/src/features/chat/protocol.ts @@ -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 }