fix(bigquery): carry job location into page 2+ (getQueryResults)

google-big-query/server/lib/bigquery-client.ts

@@ -191,10 +191,16 @@ export class BigQueryClient {
       pageToken?: string;
       startIndex?: string;
       timeoutMs?: number;
+      /** Job location — required for jobs outside the US/EU multi-regions. */
+      location?: string;
     },
   ): Promise<GetQueryResultsResponse> {
     const url = new URL(ENDPOINTS.QUERY_RESULTS(projectId, jobId));
 
+    if (options?.location) {
+      url.searchParams.set("location", options.location);
+    }
+
     if (options?.maxResults) {
       url.searchParams.set("maxResults", String(options.maxResults));
     }
@@ -258,6 +264,7 @@ export class BigQueryClient {
         projectId,
         response.jobReference.jobId,
         {
+          location: response.jobReference.location,
           timeoutMs: options.timeoutMs ?? DEFAULTS.QUERY_TIMEOUT_MS,
         },
       );
@@ -308,6 +315,7 @@ export class BigQueryClient {
     totalBytesProcessed: string;
     jobId?: string;
     pageToken?: string;
+    location?: string;
   }> {
     let response = await this.query(projectId, options);
 
@@ -325,6 +333,7 @@ export class BigQueryClient {
         projectId,
         response.jobReference.jobId,
         {
+          location: response.jobReference.location,
           timeoutMs: options.timeoutMs ?? DEFAULTS.QUERY_TIMEOUT_MS,
           maxResults: options.maxResults ?? DEFAULTS.MAX_RESULTS,
         },
@@ -351,6 +360,7 @@ export class BigQueryClient {
       totalBytesProcessed: response.totalBytesProcessed || "0",
       jobId: response.pageToken ? response.jobReference?.jobId : undefined,
       pageToken: response.pageToken,
+      location: response.jobReference?.location,
     };
   }
 

google-big-query/server/lib/pagination-token.ts

@@ -1,20 +1,31 @@
 // server/lib/pagination-token.ts
 
 /**
- * Encodes a BigQuery jobId and API pageToken into a single opaque string.
- * Format: base64(jobId + "\n" + apiPageToken)
- * The newline is the separator — jobIds and pageTokens never contain newlines.
+ * Encodes a BigQuery jobId, API pageToken and job location into one opaque
+ * string. Format: base64(jobId + "\n" + apiPageToken [+ "\n" + location])
+ * The newline is the separator — none of the three ever contains one.
+ *
+ * The location is required to fetch page 2+ of a job outside the US/EU
+ * multi-regions (e.g. southamerica-east1): jobs.getQueryResults without it
+ * looks the job up in US and 404s "Not found: Job".
+ * Tokens without a location (minted before this change) still decode.
  */
-export function encodePageToken(jobId: string, apiToken: string): string {
-  return btoa(`${jobId}\n${apiToken}`);
+export function encodePageToken(
+  jobId: string,
+  apiToken: string,
+  location?: string,
+): string {
+  return btoa(
+    location ? `${jobId}\n${apiToken}\n${location}` : `${jobId}\n${apiToken}`,
+  );
 }
 
 export function decodePageToken(token: string): {
   jobId: string;
   apiToken: string;
+  location?: string;
 } {
-  const decoded = atob(token);
-  const nl = decoded.indexOf("\n");
-  if (nl === -1) throw new Error("Invalid pageToken format");
-  return { jobId: decoded.slice(0, nl), apiToken: decoded.slice(nl + 1) };
+  const [jobId, apiToken, location] = atob(token).split("\n");
+  if (apiToken === undefined) throw new Error("Invalid pageToken format");
+  return { jobId, apiToken, ...(location ? { location } : {}) };
 }

google-big-query/server/tools/bigquery.ts

@@ -148,14 +148,18 @@ export const createQueryTool = (env: Env) =>
         totalBytesProcessed: string;
         jobId?: string;
         pageToken?: string;
+        location?: string;
       };
 
       if (context.pageToken) {
         // Subsequent page: decode opaque token and call getQueryResults
-        const { jobId, apiToken } = decodePageToken(context.pageToken);
+        const { jobId, apiToken, location } = decodePageToken(
+          context.pageToken,
+        );
         const page = await client.getQueryResults(context.projectId, jobId, {
           maxResults: context.maxResults,
           pageToken: apiToken,
+          location,
         });
         if (!page.jobComplete) {
           throw new Error("Unexpected: job not complete on page fetch");
@@ -168,6 +172,7 @@ export const createQueryTool = (env: Env) =>
           totalBytesProcessed: "0",
           jobId: page.pageToken ? jobId : undefined,
           pageToken: page.pageToken,
+          location,
         };
       } else {
         // First page: execute query
@@ -203,7 +208,7 @@ export const createQueryTool = (env: Env) =>
 
       const nextPageToken =
         result.jobId && result.pageToken
-          ? encodePageToken(result.jobId, result.pageToken)
+          ? encodePageToken(result.jobId, result.pageToken, result.location)
           : undefined;
 
       return {

google-big-query/tests/tools/bigquery-query.test.ts

@@ -149,3 +149,54 @@ test("throws when neither query nor pageToken is provided", async () => {
     }),
   ).rejects.toThrow("query is required when pageToken is not provided");
 });
+
+test("job location round-trips through nextPageToken into page 2's getQueryResults", async () => {
+  // Non-US jobs (southamerica-east1) 404 on getQueryResults without location.
+  let capturedUrl = "";
+  globalThis.fetch = mock((url: string | URL) => {
+    capturedUrl = url.toString();
+    return Promise.resolve(
+      makeResponse({
+        jobComplete: true,
+        jobReference: {
+          jobId: "jobSA",
+          projectId: "myproject",
+          location: "southamerica-east1",
+        },
+        schema: SCHEMA,
+        rows: ROWS,
+        totalRows: "500",
+        totalBytesProcessed: "100",
+        cacheHit: false,
+        pageToken: "BQ_P2",
+      }),
+    );
+  }) as unknown as typeof fetch;
+
+  const { createQueryTool } = await import("../../server/tools/bigquery.ts");
+  const tool = createQueryTool({
+    MESH_REQUEST_CONTEXT: { authorization: "Bearer fake" },
+  } as never);
+
+  // @ts-ignore
+  const first = await tool.execute({
+    context: { projectId: "myproject", query: "SELECT 1", maxResults: 1 },
+    runtimeContext: {} as never,
+  });
+  expect(decodePageToken(first.nextPageToken!).location).toBe(
+    "southamerica-east1",
+  );
+
+  // @ts-ignore
+  await tool.execute({
+    context: { projectId: "myproject", pageToken: first.nextPageToken },
+    runtimeContext: {} as never,
+  });
+  expect(capturedUrl).toContain("/queries/jobSA");
+  expect(capturedUrl).toContain("location=southamerica-east1");
+});
+
+test("legacy token without location still decodes", () => {
+  const legacy = btoa("job1\ntok1");
+  expect(decodePageToken(legacy)).toEqual({ jobId: "job1", apiToken: "tok1" });
+});