Add assistant message copy action and harden related test/storage fallbacks (#1211)

Co-authored-by: Julius Marminge <julius0216@outlook.com>

Author: shivamhwp

apps/web/src/components/ChatView.tsx

@@ -3353,6 +3353,7 @@ export default function ChatView(props: ChatViewProps) {
                 hasMessages={timelineEntries.length > 0}
                 isWorking={isWorking}
                 activeTurnInProgress={isWorking || !latestTurnSettled}
+                activeTurnId={activeLatestTurn?.turnId ?? null}
                 activeTurnStartedAt={activeWorkStartedAt}
                 scrollContainer={messagesScrollElement}
                 timelineEntries={timelineEntries}

apps/web/src/components/chat/MessageCopyButton.tsx

@@ -1,20 +1,82 @@
-import { memo } from "react";
+import { memo, useRef } from "react";
 import { CopyIcon, CheckIcon } from "lucide-react";
 import { Button } from "../ui/button";
 import { useCopyToClipboard } from "~/hooks/useCopyToClipboard";
+import { cn } from "~/lib/utils";
+import { anchoredToastManager } from "../ui/toast";
+import { Tooltip, TooltipPopup, TooltipTrigger } from "../ui/tooltip";
 
-export const MessageCopyButton = memo(function MessageCopyButton({ text }: { text: string }) {
-  const { copyToClipboard, isCopied } = useCopyToClipboard();
+const ANCHORED_TOAST_TIMEOUT_MS = 1000;
+const onCopy = (ref: React.RefObject<HTMLButtonElement | null>) => {
+  if (ref.current) {
+    anchoredToastManager.add({
+      data: {
+        tooltipStyle: true,
+      },
+      positionerProps: {
+        anchor: ref.current,
+      },
+      timeout: ANCHORED_TOAST_TIMEOUT_MS,
+      title: "Copied!",
+    });
+  }
+};
+
+const onCopyError = (ref: React.RefObject<HTMLButtonElement | null>, error: Error) => {
+  if (ref.current) {
+    anchoredToastManager.add({
+      data: {
+        tooltipStyle: true,
+      },
+      positionerProps: {
+        anchor: ref.current,
+      },
+      timeout: ANCHORED_TOAST_TIMEOUT_MS,
+      title: "Failed to copy",
+      description: error.message,
+    });
+  }
+};
+
+export const MessageCopyButton = memo(function MessageCopyButton({
+  text,
+  size = "xs",
+  variant = "outline",
+  className,
+}: {
+  text: string;
+  size?: "xs" | "icon-xs";
+  variant?: "outline" | "ghost";
+  className?: string;
+}) {
+  const ref = useRef<HTMLButtonElement>(null);
+  const { copyToClipboard, isCopied } = useCopyToClipboard<void>({
+    onCopy: () => onCopy(ref),
+    onError: (error: Error) => onCopyError(ref, error),
+    timeout: ANCHORED_TOAST_TIMEOUT_MS,
+  });
 
   return (
-    <Button
-      type="button"
-      size="xs"
-      variant="outline"
-      onClick={() => copyToClipboard(text)}
-      title="Copy message"
-    >
-      {isCopied ? <CheckIcon className="size-3 text-success" /> : <CopyIcon className="size-3" />}
-    </Button>
+    <Tooltip>
+      <TooltipTrigger
+        render={
+          <Button
+            aria-label="Copy link"
+            disabled={isCopied}
+            onClick={() => copyToClipboard(text)}
+            ref={ref}
+            type="button"
+            size={size}
+            variant={variant}
+            className={cn(className)}
+          />
+        }
+      >
+        {isCopied ? <CheckIcon className="size-3 text-success" /> : <CopyIcon className="size-3" />}
+      </TooltipTrigger>
+      <TooltipPopup>
+        <p>Copy to clipboard</p>
+      </TooltipPopup>
+    </Tooltip>
   );
 });

apps/web/src/components/chat/MessagesTimeline.logic.test.ts

@@ -1,5 +1,10 @@
 import { describe, expect, it } from "vitest";
-import { computeMessageDurationStart, normalizeCompactToolLabel } from "./MessagesTimeline.logic";
+import {
+  computeMessageDurationStart,
+  deriveMessagesTimelineRows,
+  normalizeCompactToolLabel,
+  resolveAssistantMessageCopyState,
+} from "./MessagesTimeline.logic";
 
 describe("computeMessageDurationStart", () => {
   it("returns message createdAt when there is no preceding user message", () => {
@@ -143,3 +148,120 @@ describe("normalizeCompactToolLabel", () => {
     expect(normalizeCompactToolLabel("Read file completed")).toBe("Read file");
   });
 });
+
+describe("resolveAssistantMessageCopyState", () => {
+  it("returns enabled copy state for completed assistant messages", () => {
+    expect(
+      resolveAssistantMessageCopyState({
+        showCopyButton: true,
+        text: "Ship it",
+        streaming: false,
+      }),
+    ).toEqual({
+      text: "Ship it",
+      visible: true,
+    });
+  });
+
+  it("hides copy while an assistant message is still streaming", () => {
+    expect(
+      resolveAssistantMessageCopyState({
+        showCopyButton: true,
+        text: "Still streaming",
+        streaming: true,
+      }),
+    ).toEqual({
+      text: "Still streaming",
+      visible: false,
+    });
+  });
+
+  it("hides copy for empty completed assistant messages", () => {
+    expect(
+      resolveAssistantMessageCopyState({
+        showCopyButton: true,
+        text: "   ",
+        streaming: false,
+      }),
+    ).toEqual({
+      text: null,
+      visible: false,
+    });
+  });
+
+  it("hides copy for non-terminal assistant messages", () => {
+    expect(
+      resolveAssistantMessageCopyState({
+        showCopyButton: false,
+        text: "Interim thought",
+        streaming: false,
+      }),
+    ).toEqual({
+      text: "Interim thought",
+      visible: false,
+    });
+  });
+});
+
+describe("deriveMessagesTimelineRows", () => {
+  it("only enables assistant copy for the terminal assistant message in a turn", () => {
+    const rows = deriveMessagesTimelineRows({
+      timelineEntries: [
+        {
+          id: "user-1-entry",
+          kind: "message",
+          createdAt: "2026-01-01T00:00:00Z",
+          message: {
+            id: "user-1" as never,
+            role: "user",
+            text: "Write a poem",
+            turnId: null,
+            createdAt: "2026-01-01T00:00:00Z",
+            streaming: false,
+          },
+        },
+        {
+          id: "assistant-thought-entry",
+          kind: "message",
+          createdAt: "2026-01-01T00:00:10Z",
+          message: {
+            id: "assistant-thought" as never,
+            role: "assistant",
+            text: "I should ground this first.",
+            turnId: "turn-1" as never,
+            createdAt: "2026-01-01T00:00:10Z",
+            completedAt: "2026-01-01T00:00:11Z",
+            streaming: false,
+          },
+        },
+        {
+          id: "assistant-final-entry",
+          kind: "message",
+          createdAt: "2026-01-01T00:00:20Z",
+          message: {
+            id: "assistant-final" as never,
+            role: "assistant",
+            text: "Here is the poem.",
+            turnId: "turn-1" as never,
+            createdAt: "2026-01-01T00:00:20Z",
+            completedAt: "2026-01-01T00:00:30Z",
+            streaming: false,
+          },
+        },
+      ],
+      completionDividerBeforeEntryId: "assistant-final-entry",
+      isWorking: false,
+      activeTurnStartedAt: null,
+    });
+
+    const assistantRows = rows.filter(
+      (row): row is Extract<(typeof rows)[number], { kind: "message" }> =>
+        row.kind === "message" && row.message.role === "assistant",
+    );
+
+    expect(assistantRows).toHaveLength(2);
+    expect(assistantRows[0]?.showAssistantCopyButton).toBe(false);
+    expect(assistantRows[1]?.showAssistantCopyButton).toBe(true);
+    expect(assistantRows[1]?.showCompletionDivider).toBe(true);
+  });
+});

apps/web/src/components/chat/MessagesTimeline.logic.ts

@@ -27,6 +27,7 @@ export type MessagesTimelineRow =
       message: ChatMessage;
       durationStart: string;
       showCompletionDivider: boolean;
+      showAssistantCopyButton: boolean;
     }
   | {
       kind: "proposed-plan";
@@ -59,6 +60,48 @@ export function normalizeCompactToolLabel(value: string): string {
   return value.replace(/\s+(?:complete|completed)\s*$/i, "").trim();
 }
 
+export function resolveAssistantMessageCopyState({
+  text,
+  showCopyButton,
+  streaming,
+}: {
+  text: string | null;
+  showCopyButton: boolean;
+  streaming: boolean;
+}) {
+  const hasText = text !== null && text.trim().length > 0;
+  return {
+    text: hasText ? text : null,
+    visible: showCopyButton && hasText && !streaming,
+  };
+}
+
+function deriveTerminalAssistantMessageIds(timelineEntries: ReadonlyArray<TimelineEntry>) {
+  const lastAssistantMessageIdByResponseKey = new Map<string, string>();
+  let nullTurnResponseIndex = 0;
+
+  for (const timelineEntry of timelineEntries) {
+    if (timelineEntry.kind !== "message") {
+      continue;
+    }
+    const { message } = timelineEntry;
+    if (message.role === "user") {
+      nullTurnResponseIndex += 1;
+      continue;
+    }
+    if (message.role !== "assistant") {
+      continue;
+    }
+
+    const responseKey = message.turnId
+      ? `turn:${message.turnId}`
+      : `unkeyed:${nullTurnResponseIndex}`;
+    lastAssistantMessageIdByResponseKey.set(responseKey, message.id);
+  }
+
+  return new Set(lastAssistantMessageIdByResponseKey.values());
+}
+
 export function deriveMessagesTimelineRows(input: {
   timelineEntries: ReadonlyArray<TimelineEntry>;
   completionDividerBeforeEntryId: string | null;
@@ -69,6 +112,7 @@ export function deriveMessagesTimelineRows(input: {
   const durationStartByMessageId = computeMessageDurationStart(
     input.timelineEntries.flatMap((entry) => (entry.kind === "message" ? [entry.message] : [])),
   );
+  const terminalAssistantMessageIds = deriveTerminalAssistantMessageIds(input.timelineEntries);
 
   for (let index = 0; index < input.timelineEntries.length; index += 1) {
     const timelineEntry = input.timelineEntries[index];
@@ -115,6 +159,9 @@ export function deriveMessagesTimelineRows(input: {
       showCompletionDivider:
         timelineEntry.message.role === "assistant" &&
         input.completionDividerBeforeEntryId === timelineEntry.id,
+      showAssistantCopyButton:
+        timelineEntry.message.role === "assistant" &&
+        terminalAssistantMessageIds.has(timelineEntry.message.id),
     });
   }