fix(copy-message): missing citations

Author: pnd280

src/hooks/useCopyPplxThread.ts

@@ -2,9 +2,10 @@ import { useQuery } from "@tanstack/react-query";
 import { sendMessage } from "webext-bridge/content-script";
 
 import { toast } from "@/components/ui/use-toast";
+import { threadMessageBlocksDomObserverStore } from "@/plugins/_core/dom-observers/thread/message-blocks/store";
 import { ThreadMessageApiResponse } from "@/services/pplx-api/pplx-api.types";
 import { pplxApiQueries } from "@/services/pplx-api/query-keys";
-import { INTERNAL_ATTRIBUTES, DOM_SELECTORS } from "@/utils/dom-selectors";
+import { DOM_SELECTORS } from "@/utils/dom-selectors";
 import { errorWrapper } from "@/utils/error-wrapper";
 import { ThreadExport } from "@/utils/thread-export";
 import { parseUrl } from "@/utils/utils";
@@ -91,22 +92,40 @@ async function copyMessageWithCitations({
 }: {
   messageBlockIndex: number;
 }) {
-  const message = await sendMessage(
+  const content = await sendMessage(
     "reactVdom:getMessageContent",
     {
       index: messageBlockIndex,
     },
     "window",
   );
 
-  if (message == null) {
-    $(
-      `[data-cplx-component="${INTERNAL_ATTRIBUTES.THREAD.MESSAGE.BLOCK}"][data-index="${messageBlockIndex}"] [data-cplx-component="${INTERNAL_ATTRIBUTES.THREAD.MESSAGE.TEXT_COL_CHILD.BOTTOM_BAR}"] ${DOM_SELECTORS.THREAD.MESSAGE.TEXT_COL_CHILD.BOTTOM_BAR_CHILD.COPY_BUTTON}`,
-    ).trigger("click");
+  if (content == null) {
+    const $bottomBar =
+      threadMessageBlocksDomObserverStore.getState().messageBlocks?.[
+        messageBlockIndex
+      ]?.nodes.$bottomBar;
+
+    if (!$bottomBar || !$bottomBar.length) return;
+
+    const $copyButton = $bottomBar.find(
+      DOM_SELECTORS.THREAD.MESSAGE.TEXT_COL_CHILD.BOTTOM_BAR_CHILD.COPY_BUTTON,
+    );
+
+    if (!$copyButton.length) return;
+
+    $copyButton.trigger("click");
+
     return;
   }
 
-  navigator.clipboard.writeText(message);
+  if (content.webResults && content.webResults.length) {
+    navigator.clipboard.writeText(
+      `${content.answer}\n\nCitations:\n${ThreadExport.formatWebResults(content.webResults)}`,
+    );
+  } else {
+    navigator.clipboard.writeText(content.answer);
+  }
 }
 
 async function copyMessageWithoutCitations({

src/plugins/_core/react-vdom/listeners.ts

@@ -3,6 +3,7 @@ import { onMessage } from "webext-bridge/window";
 import { LanguageModelCode } from "@/data/plugins/query-box/language-model-selector/language-models.types";
 import { INTERNAL_ATTRIBUTES, DOM_SELECTORS } from "@/utils/dom-selectors";
 import { errorWrapper } from "@/utils/error-wrapper";
+import { PplxWebResult } from "@/utils/thread-export";
 import { getCookie, getReactFiberKey } from "@/utils/utils";
 
 export type ReactVdomEvents = {
@@ -12,7 +13,10 @@ export type ReactVdomEvents = {
   "reactVdom:getMessageDisplayModelCode": (params: {
     index: number;
   }) => string | null;
-  "reactVdom:getMessageContent": (params: { index: number }) => string | null;
+  "reactVdom:getMessageContent": (params: { index: number }) => {
+    answer: string;
+    webResults: PplxWebResult[] | undefined;
+  } | null;
   "reactVdom:getCodeBlockContent": (params: {
     messageBlockIndex: number;
     codeBlockIndex: number;
@@ -94,21 +98,23 @@ export function setupReactVdomListeners() {
 
     if (!$el.length) return null;
 
-    const [answer, error] = errorWrapper(() =>
+    const [result, error] = errorWrapper(() =>
       findReactFiberNodeValue({
         fiberNode: ($el[0] as any)[getReactFiberKey($el[0])],
         condition: (node) =>
-          node.memoizedProps.children[3].props.response.answer != null,
-        select: (node) =>
-          node.memoizedProps.children[3].props.response.answer as string,
+          node.memoizedProps.children[3].props.response != null,
+        select: (node) => ({
+          answer: node.memoizedProps.children[3].props.response
+            .answer as string,
+          webResults: node.memoizedProps.children[3].props.response
+            ?.web_results as PplxWebResult[] | undefined,
+        }),
       }),
     )();
 
     if (error) console.warn("[VDOM Plugin] getMessageContent", error);
 
-    if (error || answer == null) return null;
-
-    return answer;
+    return result;
   });
 
   onMessage(

src/plugins/thread-better-message-toolbars/message-words-and-characters-count.ts

@@ -62,19 +62,19 @@ csLoaderRegistry.register({
               return;
             }
 
-            const answer = await sendMessage(
+            const content = await sendMessage(
               "reactVdom:getMessageContent",
               { index },
               "window",
             );
 
-            if (answer == null) {
+            if (content == null) {
               return;
             }
 
-            const answerWordsCount = answer.split(" ").length;
-            const answerCharactersCount = answer.length;
-            const answerTokensCount = Math.ceil(answer.length / 4);
+            const answerWordsCount = content.answer.split(" ").length;
+            const answerCharactersCount = content.answer.length;
+            const answerTokensCount = Math.ceil(content.answer.length / 4);
 
             const answerWordsAndCharactersCountContainer =
               createAnswerHeadingContainer(

src/utils/dom-selectors.ts

@@ -32,8 +32,7 @@ export const DOM_SELECTORS = {
       VISUAL_COL: ".col-span-4",
       TEXT_COL_CHILD: {
         /** The query box */
-        QUERY_WRAPPER:
-          "div.mt-md.md\\:mt-lg.mb-3.md\\:mb-6:has(> .group\\/title)",
+        QUERY_WRAPPER: `div.mt-md.md\\:mt-lg.mb-3.md\\:mb-6:has(> .group\\/title), div.mt-md.md\\:mt-lg.mb-3.md\\:mb-6:has(div[style*="height"])`,
         QUERY: ".group\\/query",
         QUERY_HOVER_CONTAINER: ".pointer-events-none.absolute.bottom-0.right-0",
         QUERY_HOVER_CONTAINER_CHILD: {
@@ -46,7 +45,7 @@ export const DOM_SELECTORS = {
         /** The sources heading */
         SOURCES: ".mb-lg.border-borderMain\\/50.ring-borderMain\\/50",
         /** The answer wrapper */
-        ANSWER: ".relative.default.font-sans.text-base",
+        ANSWER: ".mb-md > .relative.default.font-sans.text-base",
         /** The bottom toolbar of the message (share, rewrite, model name, etc.) */
         BOTTOM_BAR: ".mt-sm.flex.items-center.justify-between",
         BOTTOM_BAR_CHILD: {

src/utils/thread-export.ts

@@ -4,10 +4,10 @@ import { jsonUtils } from "@/utils/utils";
 
 type ThreadAnswer = {
   answer: string;
-  web_results: WebResult[];
+  web_results: PplxWebResult[];
 };
 
-type WebResult = {
+export type PplxWebResult = {
   name: string;
   url: string;
   snippet: string;
@@ -33,10 +33,10 @@ export class ThreadExport {
 
   private static extractWebResults(
     message: ThreadMessageApiResponse,
-  ): WebResult[] {
+  ): PplxWebResult[] {
     const text = jsonUtils.safeParse(message.text);
 
-    const webResults = text.web_results as WebResult[] | undefined;
+    const webResults = text.web_results as PplxWebResult[] | undefined;
 
     if (webResults != null) {
       return webResults;
@@ -54,7 +54,7 @@ export class ThreadExport {
     );
   }
 
-  private static formatWebResults(webResults: WebResult[]) {
+  static formatWebResults(webResults: PplxWebResult[]) {
     return webResults
       .map(
         (webResult, index) =>
@@ -86,7 +86,7 @@ export class ThreadExport {
       .join("  \n");
   }
 
-  private static trimReferences(answer: string, webResults: WebResult[]) {
+  private static trimReferences(answer: string, webResults: PplxWebResult[]) {
     webResults.forEach((_, index) => {
       const findText = `\\[${index + 1}\\]`;
       answer = answer.replace(new RegExp(findText, "g"), "");