fix: PR branch names on filtered lists

src/features/pr-branch-names.test.ts

@@ -41,6 +41,20 @@ describe("injectPRBranchNames", () => {
     expect(document.querySelectorAll(".bg-skeleton-pill--branch")).toHaveLength(0);
   });
 
+  it("requests the visible PR numbers on filtered search pages", async () => {
+    setUrl(`${GH}/owner/repo/pulls?q=pre+bid+sort%3Aupdated-desc+`);
+    twoPRRows();
+    vi.mocked(fetchPRBranches).mockResolvedValue([
+      { number: 7, headRef: "feature/a" },
+      { number: 8, headRef: "fix/b" },
+    ]);
+
+    await injectPRBranchNames();
+
+    expect(fetchPRBranches).toHaveBeenCalledWith("owner", "repo", [7, 8], "open", 1);
+    expect(document.querySelectorAll(".better-github-branch-badge")).toHaveLength(2);
+  });
+
   it("ignores PR numbers the API did not return", async () => {
     twoPRRows();
     vi.mocked(fetchPRBranches).mockResolvedValue([{ number: 7, headRef: "feature/a" }]);

src/features/pr-branch-names.ts

@@ -50,14 +50,19 @@ export async function injectPRBranchNames(): Promise<void> {
   const existing = document.querySelectorAll(`.${BADGE_CLASS}`);
   if (existing.length > 0) return;
 
+  const prRows = document.querySelectorAll("[id^='issue_']:not([id$='_link'])");
+  const prNumbers = [...prRows]
+    .map((row) => Number(row.id.replace("issue_", "")))
+    .filter(Number.isInteger);
+  if (prNumbers.length === 0) return;
+
   try {
     const { state, page } = getPRListParams();
-    const branches = await fetchPRBranches(info.owner, info.repo, state, page);
+    const branches = await fetchPRBranches(info.owner, info.repo, prNumbers, state, page);
 
     if (branches.length === 0) return;
 
     const branchMap = new Map(branches.map((b) => [b.number, b.headRef]));
-    const prRows = document.querySelectorAll("[id^='issue_']");
 
     for (const row of prRows) {
       const id = row.getAttribute("id");

src/lib/github-api.test.ts

@@ -45,28 +45,28 @@ describe("github-api bridge", () => {
   it("forwards a typed request and resolves the worker's data on success", async () => {
     const runtime = mockRuntime({ response: { ok: true, data: [{ number: 7, headRef: "feature/a" }] } });
 
-    const result = await fetchPRBranches("owner", "repo", "open", 2);
+    const result = await fetchPRBranches("owner", "repo", [7], "open", 2);
 
     expect(result).toEqual([{ number: 7, headRef: "feature/a" }]);
     expect(runtime.sendMessage).toHaveBeenCalledWith(
-      { type: "FETCH_PR_BRANCHES", owner: "owner", repo: "repo", state: "open", page: 2 },
+      { type: "FETCH_PR_BRANCHES", owner: "owner", repo: "repo", prNumbers: [7], state: "open", page: 2 },
       expect.any(Function),
     );
   });
 
   it("swallows an ok:false response and returns the empty default", async () => {
     mockRuntime({ response: { ok: false, error: "boom" } });
-    expect(await fetchPRBranches("owner", "repo")).toEqual([]);
+    expect(await fetchPRBranches("owner", "repo", [1])).toEqual([]);
   });
 
   it("treats chrome.runtime.lastError as a failure", async () => {
     mockRuntime({ response: { ok: true, data: [] }, lastError: { message: "port closed" } });
-    expect(await fetchPRBranches("owner", "repo")).toEqual([]);
+    expect(await fetchPRBranches("owner", "repo", [1])).toEqual([]);
   });
 
   it("rejects without messaging when the extension context is invalidated", async () => {
     const runtime = mockRuntime({ id: undefined });
-    expect(await fetchPRBranches("owner", "repo")).toEqual([]);
+    expect(await fetchPRBranches("owner", "repo", [1])).toEqual([]);
     expect(runtime.sendMessage).not.toHaveBeenCalled();
   });
 

src/lib/github-api.ts

@@ -79,6 +79,7 @@ export async function fetchContributorInfo(
 export async function fetchPRBranches(
   owner: string,
   repo: string,
+  prNumbers: number[],
   state: string = "open",
   page: number = 1,
 ): Promise<PRBranchInfo[]> {
@@ -87,6 +88,7 @@ export async function fetchPRBranches(
       type: "FETCH_PR_BRANCHES",
       owner,
       repo,
+      prNumbers,
       state,
       page,
     });

src/lib/messages.ts

@@ -103,7 +103,7 @@ export interface ContributorInfo {
 }
 
 export type ServiceWorkerRequest =
-  | { type: "FETCH_PR_BRANCHES"; owner: string; repo: string; state: string; page: number }
+  | { type: "FETCH_PR_BRANCHES"; owner: string; repo: string; prNumbers: number[]; state: string; page: number }
   | { type: "FETCH_PR_CONFLICT_STATUSES"; owner: string; repo: string; prNumbers: number[] }
   | { type: "FETCH_PR_REVIEW_STATUSES"; owner: string; repo: string; prNumbers: number[] }
   | { type: "FETCH_PR_REVIEW_THREAD_DETAILS"; owner: string; repo: string; prNumber: number }

src/service-worker.test.ts

@@ -130,7 +130,7 @@ describe("service worker", () => {
   });
 
   it("coalesces concurrent requests and caches successful PR branch responses", async () => {
-    const state = await loadWorker("token");
+    const state = await loadWorker();
     const fetchDeferred = deferred<Response>();
     const fetchMock = vi.mocked(fetch);
     fetchMock.mockReturnValue(fetchDeferred.promise);
@@ -139,6 +139,7 @@ describe("service worker", () => {
       type: "FETCH_PR_BRANCHES",
       owner: "owner",
       repo: "repo",
+      prNumbers: [7],
       state: "open",
       page: 1,
     };
@@ -157,7 +158,7 @@ describe("service worker", () => {
 
   it("returns fresh cached data without fetching", async () => {
     vi.spyOn(Date, "now").mockReturnValue(1_000);
-    const state = await loadWorker("token");
+    const state = await loadWorker();
     state.sessionStore["cache:branches:owner/repo:open:1"] = {
       data: [{ number: 1, headRef: "cached" }],
       timestamp: 900,
@@ -167,6 +168,7 @@ describe("service worker", () => {
       type: "FETCH_PR_BRANCHES",
       owner: "owner",
       repo: "repo",
+      prNumbers: [1],
       state: "open",
       page: 1,
     });
@@ -177,7 +179,7 @@ describe("service worker", () => {
 
   it("expires stale cache entries before fetching", async () => {
     vi.spyOn(Date, "now").mockReturnValue(10 * 60 * 1000);
-    const state = await loadWorker("token");
+    const state = await loadWorker();
     state.sessionStore["cache:branches:owner/repo:open:1"] = {
       data: [{ number: 1, headRef: "stale" }],
       timestamp: 0,
@@ -188,6 +190,7 @@ describe("service worker", () => {
       type: "FETCH_PR_BRANCHES",
       owner: "owner",
       repo: "repo",
+      prNumbers: [2],
       state: "open",
       page: 1,
     });
@@ -198,6 +201,63 @@ describe("service worker", () => {
     });
   });
 
+  it("fetches branches by visible PR number when the list page response does not match", async () => {
+    const state = await loadWorker("token");
+    vi.mocked(fetch).mockResolvedValueOnce(
+      jsonResponse({
+        data: {
+          repository: {
+            pr_7: { headRefName: "feature/a" },
+            pr_8: { headRefName: "fix/b" },
+          },
+        },
+      }),
+    );
+
+    const response = await sendMessage(state.messageListeners[0], {
+      type: "FETCH_PR_BRANCHES",
+      owner: "owner",
+      repo: "repo",
+      prNumbers: [8, 7],
+      state: "open",
+      page: 1,
+    });
+
+    expect(response).toEqual({
+      ok: true,
+      data: [
+        { number: 7, headRef: "feature/a" },
+        { number: 8, headRef: "fix/b" },
+      ],
+    });
+    expect(fetch).toHaveBeenCalledTimes(1);
+    expect(vi.mocked(fetch).mock.calls[0][0]).toBe("https://api.github.com/graphql");
+  });
+
+  it("falls back to exact REST requests for public filtered lists", async () => {
+    const state = await loadWorker();
+    vi.mocked(fetch)
+      .mockResolvedValueOnce(jsonResponse([]))
+      .mockResolvedValueOnce(jsonResponse({ number: 7, head: { ref: "feature/a" } }));
+
+    const response = await sendMessage(state.messageListeners[0], {
+      type: "FETCH_PR_BRANCHES",
+      owner: "owner",
+      repo: "repo",
+      prNumbers: [7],
+      state: "open",
+      page: 1,
+    });
+
+    expect(response).toEqual({
+      ok: true,
+      data: [{ number: 7, headRef: "feature/a" }],
+    });
+    expect(vi.mocked(fetch).mock.calls[1][0]).toBe(
+      "https://api.github.com/repos/owner/repo/pulls/7",
+    );
+  });
+
   it("does not call GraphQL-backed endpoints without a token", async () => {
     const state = await loadWorker();
 

src/service-worker.ts

@@ -86,14 +86,33 @@ async function fetchPRBranches(
   repo: string,
   state: string,
   page: number,
+  prNumbers: number[],
 ): Promise<PRBranchInfo[]> {
+  const requested = [...new Set(prNumbers)].sort((a, b) => a - b);
+  if (requested.length === 0) return [];
+
+  const token = await getToken();
+  if (token) {
+    return fetchGraphQLBatch<number, PRBranchInfo>({
+      cachePrefix: "branches",
+      owner,
+      repo,
+      keys: requested,
+      aliasFor: (number) => `pr_${number}`,
+      buildNodeQuery: (number) => `pullRequest(number: ${number}) {
+        headRefName
+      }`,
+      parseNode: (number, pr) =>
+        typeof pr.headRefName === "string" ? { number, headRef: pr.headRefName } : null,
+    });
+  }
+
   const cacheKey = `cache:branches:${owner}/${repo}:${state}:${page}`;
-  return cachedFetch<PRBranchInfo[]>(cacheKey, async () => {
+  const pageBranches = await cachedFetch<PRBranchInfo[]>(cacheKey, async () => {
     const perPage = 30;
     const url = `https://api.github.com/repos/${owner}/${repo}/pulls?state=${state}&sort=updated&direction=desc&page=${page}&per_page=${perPage}`;
 
-    const headers = restHeaders(await getToken());
-    const response = await fetch(url, { headers });
+    const response = await fetch(url, { headers: restHeaders(token) });
 
     if (!response.ok) {
       console.error(`[Better GitHub] API error: ${response.status} ${response.statusText}`);
@@ -106,6 +125,27 @@ async function fetchPRBranches(
       headRef: pr.head.ref,
     }));
   });
+
+  const missing = requested.filter(
+    (number) => !pageBranches.some((branch) => branch.number === number),
+  );
+  const exactBranches = await Promise.all(
+    missing.map((number) =>
+      cachedFetch<PRBranchInfo[]>(`cache:branches:${owner}/${repo}:pr:${number}`, async () => {
+        const response = await fetch(
+          `https://api.github.com/repos/${owner}/${repo}/pulls/${number}`,
+          { headers: restHeaders(token) },
+        );
+        if (!response.ok) {
+          console.error(`[Better GitHub] API error: ${response.status} ${response.statusText}`);
+          return [];
+        }
+        const pull = (await response.json()) as { number: number; head: { ref: string } };
+        return [{ number: pull.number, headRef: pull.head.ref }];
+      }),
+    ),
+  );
+  return pageBranches.concat(...exactBranches);
 }
 
 interface GraphQLBatchSpec<K extends string | number, V> {
@@ -772,7 +812,7 @@ async function handleMessage(
     case "FETCH_PR_BRANCHES":
       return {
         ok: true,
-        data: await fetchPRBranches(request.owner, request.repo, request.state, request.page),
+        data: await fetchPRBranches(request.owner, request.repo, request.state, request.page, request.prNumbers),
       };
     case "FETCH_PR_CONFLICT_STATUSES":
       return {