Staging

e2e/citation-dedupe.e2e.ts

@@ -276,3 +276,103 @@ test("draws citation regions instead of a full-page highlight", async ({
     )
   }
 })
+
+test("moves a deduplicated cited page to the top from the same source", async ({
+  context,
+  page,
+}) => {
+  await context.addCookies([
+    {
+      name: "better-auth.session_token",
+      value: "playwright",
+      url: "http://localhost:3000",
+    },
+  ])
+  await page.setViewportSize({ width: 1280, height: 832 })
+  await page.route("**/api/sources/source_spacex/chunks**", async (route) => {
+    const pageAsset = {
+      pageNumber: 26,
+      assetUrl: "/images/knowhere/logo-icon.png",
+      contentType: "image/png",
+      width: 1000,
+      height: 1400,
+    }
+    await route.fulfill({
+      status: 200,
+      contentType: "application/json",
+      body: JSON.stringify({
+        chunks: [
+          {
+            chunkId: "page_1",
+            documentId: "doc_spacex",
+            sectionPath: "Page 1",
+            type: "page",
+            content: "Content from page 1.",
+            readableContent: "Content from page 1.",
+            pageNums: [1],
+            pageAssets: [{ ...pageAsset, pageNumber: 1 }],
+            sourceTitle: "spacex-s1.pdf",
+          },
+          {
+            chunkId: "page_26_first_section",
+            documentId: "doc_spacex",
+            sectionPath: "Page 26 / First section",
+            type: "page",
+            content: "Other content from page 26.",
+            readableContent: "Other content from page 26.",
+            pageNums: [26],
+            pageAssets: [pageAsset],
+            sourceTitle: "spacex-s1.pdf",
+          },
+          {
+            chunkId: "page_26_cited_section",
+            documentId: "doc_spacex",
+            sectionPath: "Page 26",
+            type: "page",
+            content: "Revenue evidence on page 26.",
+            readableContent: "Revenue evidence on page 26.",
+            pageNums: [26],
+            pageAssets: [pageAsset],
+            sourceTitle: "spacex-s1.pdf",
+          },
+        ],
+        pagination: {
+          page: 1,
+          pageSize: 50,
+          total: 3,
+          totalPages: 1,
+        },
+      }),
+    })
+  })
+
+  await page.goto("/e2e/citation-same-page")
+  await page
+    .getByTestId("desktop-sources-panel")
+    .getByRole("button", { name: "Open spacex-s1.pdf parsed chunks" })
+    .click()
+  const chunksPanel = page.getByTestId("desktop-chunks-panel")
+  await chunksPanel.getByRole("button", { name: "List" }).click()
+  await expect
+    .poll(() =>
+      chunksPanel
+        .locator('[data-index="0"]')
+        .getAttribute("data-chunk-id"),
+    )
+    .toBe("page_1")
+
+  await page
+    .getByTestId("desktop-chat-panel")
+    .getByTestId("citation-chip")
+    .first()
+    .click()
+
+  await expect
+    .poll(() =>
+      chunksPanel
+        .locator('[data-index="0"]')
+        .getAttribute("data-chunk-id"),
+    )
+    .toBe("page_26_first_section")
+  await expect(page.getByTestId("citation-region-highlight")).toHaveCount(1)
+})

src/app/e2e/citation-same-page/page.tsx

@@ -34,6 +34,7 @@ const chatMessages: ChatMessageView[] = [
       {
         chunkType: "page",
         score: 0.91,
+        content: "Revenue evidence on page 26.",
         pageCitationPageNumber: 26,
         highlightRegions: [{ x: 0.12, y: 0.18, w: 0.46, h: 0.08 }],
         source: {

src/components/chunks-panel-state.test.ts

@@ -208,6 +208,44 @@ describe("chunksPanelState", () => {
     ).toEqual(["page_4_first", "page_5"])
   })
 
+  it("focuses the retained page card when the cited logical chunk was deduplicated", () => {
+    const chunks: ParsedChunkView[] = [
+      {
+        chunkId: "page_1",
+        type: "page",
+        content: "Page 1.",
+        sourceTitle: "cdn.pdf",
+        pageNums: [1],
+      },
+      {
+        chunkId: "page_4_first_section",
+        type: "page",
+        content: "First section on page 4.",
+        sourceTitle: "cdn.pdf",
+        pageNums: [4],
+      },
+      {
+        chunkId: "page_4_cited_section",
+        type: "page",
+        content: "Cited section on page 4.",
+        sourceTitle: "cdn.pdf",
+        pageNums: [4],
+      },
+    ]
+    const displayedChunks =
+      chunksPanelState.getPageAssetChunksWithoutDuplicatePages(chunks)
+
+    expect(
+      chunksPanelState
+        .getChunksWithFocusedFirst(
+          displayedChunks,
+          "page_4_cited_section",
+          4,
+        )
+        .map((chunk) => chunk.chunkId),
+    ).toEqual(["page_4_first_section", "page_1"])
+  })
+
   it("keeps overlapping page-memory section chunks that share a boundary page", () => {
     const chunks: ParsedChunkView[] = [
       {

src/components/chunks-panel-state.ts

@@ -63,8 +63,12 @@ function getChunksWithFocusedFirst(
   const orderedChunks = getChunksOrderedByPageNumber(
     dedupeChunksById(chunks),
   )
-  const targetChunkId =
-    focusedChunkId ?? findChunkIdForPageNumber(orderedChunks, focusedPageNumber)
+  const hasVisibleFocusedChunk =
+    focusedChunkId !== null &&
+    orderedChunks.some((chunk) => chunk.chunkId === focusedChunkId)
+  const targetChunkId = hasVisibleFocusedChunk
+    ? focusedChunkId
+    : findChunkIdForPageNumber(orderedChunks, focusedPageNumber)
   if (!targetChunkId) return orderedChunks
 
   const focusedIndex = orderedChunks.findIndex(

src/domains/chat/index.test.ts

@@ -4,6 +4,8 @@ import type {
   KnowledgeGrepResponse,
   KnowledgeOutline,
   KnowledgeReadResponse,
+  RetrievalQueryParams,
+  RetrievalQueryResponse,
   RetrievalResult,
 } from "@ontos-ai/knowhere-sdk"
 import { Effect } from "effect"
@@ -362,6 +364,155 @@ describe("answerQuestionWithRetrieval", () => {
     });
   });
 
+  it("queries at most four namespaces concurrently", async () => {
+    const namespaces = [
+      "default",
+      "notebook-1",
+      "notebook-2",
+      "notebook-3",
+      "notebook-4",
+      "notebook-5",
+    ];
+    const releases: Array<() => void> = [];
+    let activeCount = 0;
+    let maxActiveCount = 0;
+    const retrieval = {
+      query: vi.fn(async (params: RetrievalQueryParams) => {
+        const namespace = requireRetrievalQueryText(
+          params.namespace,
+          "namespace",
+        );
+        const query = requireRetrievalQueryText(params.query, "query");
+        activeCount += 1;
+        maxActiveCount = Math.max(maxActiveCount, activeCount);
+        await new Promise<void>((resolve) => {
+          releases.push(() => {
+            activeCount -= 1;
+            resolve();
+          });
+        });
+        return makeRetrievalQueryResponse(namespace, query);
+      }),
+    };
+    const generateAnswer = vi.fn(async ({ searchSources }) => {
+      await searchSources({ query: "parallel evidence" });
+      return makeHarnessRunResult("The answer is grounded.");
+    });
+
+    const answerPromise = Effect.runPromise(
+      answerQuestionWithRetrieval({
+        question: "What do the documents say?",
+        namespace: "notebook-1",
+        namespaces,
+        sources: [makeSource()],
+        excludedSourceIds: [],
+        retrieval,
+        generateAnswer,
+        messages: [],
+      }),
+    );
+
+    await vi.waitFor(() => expect(retrieval.query).toHaveBeenCalledTimes(4));
+    expect(maxActiveCount).toBe(4);
+    releases.splice(0).forEach((release) => release());
+
+    await vi.waitFor(() => expect(retrieval.query).toHaveBeenCalledTimes(6));
+    expect(maxActiveCount).toBe(4);
+    releases.splice(0).forEach((release) => release());
+
+    await answerPromise;
+    expect(maxActiveCount).toBe(4);
+  });
+
+  it("retries rate-limited namespace queries sequentially", async () => {
+    const namespaces = ["default", "notebook-1", "notebook-2"];
+    const attempts = new Map<string, number>();
+    const fallbackReleases: Array<() => void> = [];
+    let activeFallbackCount = 0;
+    let maxActiveFallbackCount = 0;
+    const retrieval = {
+      query: vi.fn(async (params: RetrievalQueryParams) => {
+        const namespace = requireRetrievalQueryText(
+          params.namespace,
+          "namespace",
+        );
+        const query = requireRetrievalQueryText(params.query, "query");
+        const attempt = (attempts.get(namespace) ?? 0) + 1;
+        attempts.set(namespace, attempt);
+
+        if (namespace !== "default" && attempt === 1) {
+          throw Object.assign(
+            new Error(
+              "Too many concurrent requests. Please retry after 30 seconds.",
+            ),
+            {
+              statusCode: 429,
+              code: "RESOURCE_EXHAUSTED",
+            },
+          );
+        }
+
+        if (namespace !== "default") {
+          activeFallbackCount += 1;
+          maxActiveFallbackCount = Math.max(
+            maxActiveFallbackCount,
+            activeFallbackCount,
+          );
+          await new Promise<void>((resolve) => {
+            fallbackReleases.push(() => {
+              activeFallbackCount -= 1;
+              resolve();
+            });
+          });
+        }
+
+        return makeRetrievalQueryResponse(namespace, query);
+      }),
+    };
+    const generateAnswer = vi.fn(async ({ searchSources }) => {
+      await searchSources({ query: "rate-limited evidence" });
+      return makeHarnessRunResult("The answer is grounded.");
+    });
+
+    const answerPromise = Effect.runPromise(
+      answerQuestionWithRetrieval({
+        question: "What do the documents say?",
+        namespace: "notebook-1",
+        namespaces,
+        sources: [makeSource()],
+        excludedSourceIds: [],
+        retrieval,
+        generateAnswer,
+        messages: [],
+      }),
+    );
+
+    await vi.waitFor(() => expect(retrieval.query).toHaveBeenCalledTimes(4));
+    expect(maxActiveFallbackCount).toBe(1);
+    expect(attempts.get("notebook-2")).toBe(1);
+    fallbackReleases.shift()?.();
+
+    await vi.waitFor(() => expect(retrieval.query).toHaveBeenCalledTimes(5));
+    expect(maxActiveFallbackCount).toBe(1);
+    fallbackReleases.shift()?.();
+
+    await answerPromise;
+    expect(attempts).toEqual(
+      new Map([
+        ["default", 1],
+        ["notebook-1", 2],
+        ["notebook-2", 2],
+      ]),
+    );
+    expect(loggerMock.warn).toHaveBeenCalledWith(
+      "chat-agent: searchSources rate limited; retrying failed namespaces sequentially",
+      {
+        namespaceCount: 2,
+        namespaces: ["notebook-1", "notebook-2"],
+      },
+    );
+  });
+
   it("bounds merged retrieval evidence before passing it to the answer agent", async () => {
     const defaultResults = Array.from({ length: 40 }, (_, index) =>
       makeRetrievalResult({
@@ -3115,6 +3266,33 @@ function makeRetrievalResult(
   };
 }
 
+function makeRetrievalQueryResponse(
+  namespace: string,
+  query: string,
+): RetrievalQueryResponse {
+  return {
+    results: [makeRetrievalResult()],
+    evidenceText: `Evidence from ${namespace}`,
+    referencedChunks: [],
+    namespace,
+    query,
+    routerUsed: "workflow_single_step",
+    answerText: null,
+    stopReason: "answer_done",
+    failureReason: null,
+  };
+}
+
+function requireRetrievalQueryText(
+  value: string | undefined,
+  name: string,
+): string {
+  if (typeof value !== "string" || value.length === 0) {
+    throw new Error(`Expected retrieval ${name}.`);
+  }
+  return value;
+}
+
 function makeSource(overrides: Partial<Source> = {}): Source {
   return {
     id: "source_1",