Add pagination of batch requests

packages/typescript/src/api/async/client.ts

@@ -49,6 +49,7 @@ export class Client {
     private connection: MessageConnection | undefined;
     private options: ClientOptions;
     private connected = false;
+    private connecting: Promise<void> | undefined;
     private timing: TimingCollector | undefined;
     private batchedRequests: { method: APIRequest["method"]; params: APIRequest["params"]; resolve: (value: unknown) => void; reject: (reason?: any) => void; }[] = [];
     private nextBatch: NodeJS.Immediate | "manual" | undefined;
@@ -60,9 +61,14 @@ export class Client {
         }
     }
 
-    async connect(): Promise<void> {
-        if (this.connected) return;
+    connect(): Promise<void> {
+        if (this.connected) return Promise.resolve();
+        return this.connecting ??= this.connectWorker().finally(() => {
+            this.connecting = undefined;
+        });
+    }
 
+    private async connectWorker(): Promise<void> {
         if (isSpawnOptions(this.options)) {
             await this.connectViaSpawn(this.options);
         }
@@ -214,10 +220,27 @@ export class Client {
 
             const requestType = new RequestType<unknown, BatchRequestsResponse, void>("batchRequests");
             const params: BatchRequestsParams = { requests: requests.map(request => ({ method: request.method, params: request.params })) };
+            if (this.options.maxResponseBytesPerPage !== undefined) {
+                params.maxResponseBytesPerPage = this.options.maxResponseBytesPerPage;
+            }
             const response = await this.sendRequestWithTiming(requestType, params);
+            let responses = response.responses;
+            let continuationToken = response.continuationToken;
+            while (continuationToken) {
+                const pageParams: BatchRequestsParams = {
+                    requests: [],
+                    continuationToken,
+                };
+                if (this.options.maxResponseBytesPerPage !== undefined) {
+                    pageParams.maxResponseBytesPerPage = this.options.maxResponseBytesPerPage;
+                }
+                const page = await this.sendRequestWithTiming(requestType, pageParams);
+                responses = responses.concat(page.responses);
+                continuationToken = page.continuationToken;
+            }
             for (let i = 0; i < requests.length; i++) {
                 const { resolve, reject } = requests[i];
-                const item = response.responses[i];
+                const item = responses[i];
                 if (item.error !== undefined) {
                     reject(new Error(item.error));
                 }
@@ -253,7 +276,7 @@ export class Client {
         };
     }
 
-    async apiRequest<K extends keyof APIMethodInfo>(method: K, params: APIMethodInfo[K]["params"]): Promise<APIMethodInfo[K]["result"]> {
+    async apiRequest<K extends APIRequest["method"]>(method: K, params: APIMethodInfo[K]["params"]): Promise<APIMethodInfo[K]["result"]> {
         if (!this.connected) {
             await this.connect();
         }

packages/typescript/src/api/options.ts

@@ -8,6 +8,8 @@ import type { FileSystem } from "./fs.ts";
 export interface ClientSocketOptions {
     /** Path to the Unix domain socket or Windows named pipe for API communication */
     pipe: string;
+    /** Maximum encoded byte size of each batch response page. Defaults to 300 million bytes. */
+    maxResponseBytesPerPage?: number;
 }
 
 export interface ClientSpawnOptions {
@@ -19,6 +21,8 @@ export interface ClientSpawnOptions {
     fs?: FileSystem;
     /** Allow trusted projects to execute configured external content mapper processes. */
     runExternalCode?: boolean;
+    /** Maximum encoded byte size of each batch response page. Defaults to 300 million bytes. */
+    maxResponseBytesPerPage?: number;
     /**
      * When true, collect timing information for each request. The client
      * measures round-trip latency and bytes sent/received, and the server

packages/typescript/src/api/proto.generated.ts

@@ -163,10 +163,13 @@ export interface ReleaseParams {
 
 export interface BatchRequestsParams {
     requests: readonly BatchRequest[] | null;
+    continuationToken?: string;
+    maxResponseBytesPerPage?: number;
 }
 
 export interface BatchRequestsResponse {
     responses: BatchResponse[];
+    continuationToken?: string;
 }
 
 /** InitializeResponse is returned by the initialize method. */

packages/typescript/src/api/proto.ts

@@ -25,7 +25,8 @@ export type TypePropertyMethod = Exclude<APIMethodsReturning<TypeResponse>, Intr
 export type TypesPropertyMethod = APIMethodsReturning<TypeResponse[]>;
 export type IntrinsicTypeMethod = "getAnyType" | "getBigIntType" | "getBooleanType" | "getESSymbolType" | "getNeverType" | "getNonPrimitiveType" | "getNullType" | "getNumberType" | "getStringType" | "getUndefinedType" | "getUnknownType" | "getVoidType";
 
-export type APIRequest = { [K in keyof APIMethodInfo]: { method: K; params: APIMethodInfo[K]["params"]; }; }[keyof APIMethodInfo];
+type BatchableAPIMethod = Exclude<keyof APIMethodInfo, "batchRequests">;
+export type APIRequest = { [K in BatchableAPIMethod]: { method: K; params: APIMethodInfo[K]["params"]; }; }[BatchableAPIMethod];
 export type APIResponse<Request extends APIRequest = APIRequest> = Request extends APIRequest ?
         & {
             method: Request["method"];

packages/typescript/src/api/sync/client.ts

@@ -10,6 +10,7 @@ import {
 import type {
     APIMethodInfo,
     APIRequest,
+    BatchRequestsParams,
     BatchRequestsResponse,
     SourceFileResponseMethod,
 } from "../proto.ts";
@@ -28,13 +29,15 @@ export class Client {
     private channel: SyncRpcChannel;
     private encoder = new TextEncoder();
     private timing: TimingCollector | undefined;
+    private maxResponseBytesPerPage: number | undefined;
 
     constructor(options: ClientOptions) {
         if (!isSpawnOptions(options)) {
             throw new Error("Socket connections are not yet supported in the sync client");
         }
 
         const args = getAPIProcessArgs(options, false);
+        this.maxResponseBytesPerPage = options.maxResponseBytesPerPage;
 
         // Enable virtual FS callbacks for each provided FS function
         const enabledCallbacks: (typeof fsCallbackNames[number])[] = [];
@@ -99,7 +102,26 @@ export class Client {
     }
 
     batchRequests(requests: readonly APIRequest[]): BatchRequestsResponse {
-        return this.apiRequest("batchRequests", { requests });
+        const params: BatchRequestsParams = { requests };
+        if (this.maxResponseBytesPerPage !== undefined) {
+            params.maxResponseBytesPerPage = this.maxResponseBytesPerPage;
+        }
+        const response = this.apiRequest("batchRequests", params);
+        let responses = response.responses;
+        let continuationToken = response.continuationToken;
+        while (continuationToken) {
+            const pageParams: BatchRequestsParams = {
+                requests: [],
+                continuationToken,
+            };
+            if (this.maxResponseBytesPerPage !== undefined) {
+                pageParams.maxResponseBytesPerPage = this.maxResponseBytesPerPage;
+            }
+            const page = this.apiRequest("batchRequests", pageParams);
+            responses = responses.concat(page.responses);
+            continuationToken = page.continuationToken;
+        }
+        return { responses };
     }
 
     apiRequestBinary<K extends SourceFileResponseMethod>(method: K, params?: APIMethodInfo[K]["params"]): Uint8Array | undefined {

packages/typescript/test/async/api.test.ts

@@ -421,6 +421,28 @@ describe("API - automatic batching", () => {
 });
 
 describe("API - batchContext", () => {
+    test("transparently paginates batch responses", async () => {
+        const api = spawnAPI({ ...defaultFiles }, { maxResponseBytesPerPage: 1 });
+        try {
+            const requests = await (async () => {
+                using _ = api.batchContext();
+                return [
+                    api.parseCommandLine(["--strict"]),
+                    api.readConfigFile("/tsconfig.json"),
+                    api.parseCommandLine(["--noImplicitAny"]),
+                ] as const;
+            })();
+
+            const [strict, config, noImplicitAny] = await Promise.all(requests);
+            assert.equal(strict.options.strict, true);
+            assert.deepEqual(config.config, {});
+            assert.equal(noImplicitAny.options.noImplicitAny, true);
+        }
+        finally {
+            await api.close();
+        }
+    });
+
     test("holds requests until disposal", async () => {
         const api = spawnAPI();
         try {

packages/typescript/test/sync/api-generators.test.ts

@@ -481,6 +481,47 @@ describe("API - generator batching", () => {
         }
     });
 
+    test("transparently paginates batch responses", () => {
+        const api = spawnAPI(undefined, { maxResponseBytesPerPage: 1 });
+        try {
+            const [strict, config, noImplicitAny] = api.batch(
+                api.parseCommandLine.gen(["--strict"]),
+                api.readConfigFile.gen("/tsconfig.json"),
+                api.parseCommandLine.gen(["--noImplicitAny"]),
+            );
+
+            assert.equal(strict.options.strict, true);
+            assert.deepEqual(config.config, {});
+            assert.equal(noImplicitAny.options.noImplicitAny, true);
+        }
+        finally {
+            api.close();
+        }
+    });
+
+    test("transparently paginates responses at the default batch size limit", () => {
+        const largeConfigValue = "x".repeat(5_000_000);
+        const requestCount = 64;
+        const api = spawnAPI({ "/large.json": JSON.stringify({ largeConfigValue }) }, { collectTiming: true });
+        try {
+            api.parseCommandLine([]);
+            api.resetTimingInfo();
+
+            const configs = api.batch(...Array.from({ length: requestCount }, () => api.readConfigFile.gen("/large.json")));
+            assert.equal(configs.length, requestCount);
+            for (const config of configs) {
+                assert.deepEqual(config.config, { largeConfigValue });
+            }
+
+            const timing = api.getTimingInfo();
+            assert.equal(timing.totals.requestCount, 2);
+            assert.deepEqual(timing.recentRequests.map(request => request.method), ["batchRequests", "batchRequests"]);
+        }
+        finally {
+            api.close();
+        }
+    });
+
     test("all deduplicates only initialize requests within a batch round", () => {
         const api = spawnAPI();
         const requestBatches: string[][] = [];

tsc/internal/api/proto.go

@@ -621,7 +621,9 @@ type TranspileOutputResponse struct {
 }
 
 type BatchRequestsParams struct {
-	Requests []BatchRequest `json:"requests"`
+	Requests                []BatchRequest `json:"requests"`
+	ContinuationToken       string         `json:"continuationToken,omitempty"`
+	MaxResponseBytesPerPage int            `json:"maxResponseBytesPerPage,omitempty"`
 }
 
 type BatchRequest struct {
@@ -630,7 +632,48 @@ type BatchRequest struct {
 }
 
 type BatchRequestsResponse struct {
-	Responses []BatchResponse `json:"responses" nonnil:"true"`
+	Responses         []BatchResponse `json:"responses" nonnil:"true"`
+	ContinuationToken string          `json:"continuationToken,omitempty"`
+	encodedResponses  []json.Value
+}
+
+var _ json.MarshalerTo = (*BatchRequestsResponse)(nil)
+
+func (r *BatchRequestsResponse) MarshalJSONTo(enc *json.Encoder) error {
+	if err := enc.WriteToken(json.BeginObject); err != nil {
+		return err
+	}
+	if err := enc.WriteValue(json.Value(`"responses"`)); err != nil {
+		return err
+	}
+	if err := enc.WriteToken(json.BeginArray); err != nil {
+		return err
+	}
+	if r.encodedResponses != nil {
+		for _, response := range r.encodedResponses {
+			if err := enc.WriteValue(response); err != nil {
+				return err
+			}
+		}
+	} else {
+		for i := range r.Responses {
+			if err := json.MarshalEncode(enc, &r.Responses[i]); err != nil {
+				return err
+			}
+		}
+	}
+	if err := enc.WriteToken(json.EndArray); err != nil {
+		return err
+	}
+	if r.ContinuationToken != "" {
+		if err := enc.WriteValue(json.Value(`"continuationToken"`)); err != nil {
+			return err
+		}
+		if err := json.MarshalEncode(enc, r.ContinuationToken); err != nil {
+			return err
+		}
+	}
+	return enc.WriteToken(json.EndObject)
 }
 
 type BatchResponse struct {