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

@@ -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

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) {

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

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" + "<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()
 

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{
+	return c.editMessage(ctx, chatID, messageID, content, nil)
-		"message_id": messageID,
-		"content":    content,
-	})
-	return c.broadcastToSession(chatID, outMsg)
 }
 
 // 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)
+}

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)
 

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 {

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) =>
+        messages: (() => {
-          msg.id === messageId
+          let found = false
-            ? {
+          const messages = prev.messages.map((msg) => {
-                ...msg,
+            if (msg.id !== messageId) {
-                content,
+              return msg
-                ...(hasKind ? { kind } : {}),
+            }
-                ...(attachments ? { attachments } : {}),
+            found = true
-              }
+            return {
-            : msg,
+              ...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
     }