Add arbitrary API request batching

Herebyfile.mjs

@@ -578,7 +578,10 @@ export const generateAST = task({
 export const generateAPI = task({
     name: "generate:api",
     description: "Generates API files from internal/api/proto.go and internal/api/session.go.",
-    run: () => $`go -C ./tools run ./gen-proto ../tsc/internal/api/proto.go ../packages/typescript/src/api/proto.generated.ts`,
+    run: async () => {
+        await $`go -C ./tools run ./gen-proto ../tsc/internal/api/proto.go ../packages/typescript/src/api/proto.generated.ts`;
+        await $`npx dprint fmt packages/typescript/src/api/proto.generated.ts`;
+    },
 });
 
 // ── Vendored npm dependencies ───────────────────────────────────

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

@@ -49,6 +49,8 @@ import {
     toPath,
 } from "../path.ts";
 import type {
+    APIRequest,
+    APIResponseTuple,
     CompilerOptions,
     Diagnostic,
     DocumentIdentifier,
@@ -230,6 +232,19 @@ export class API<FromLSP extends boolean = false> {
         return api;
     }
 
+    async batchRequests<const Requests extends readonly APIRequest[]>(requests: Requests): Promise<{ responses: APIResponseTuple<Requests>; }> {
+        const response = await this.client.apiRequest("batchRequests", { requests });
+        // we're replacing the `unknown`s in the autogenerated types with much more specific per-request types here, which creates some variance issues for the
+        // `batchRequests` method within the batch request object itself, so we cast.
+        return response as { responses: APIResponseTuple<Requests>; };
+    }
+
+    // @sync-skip-block-start
+    batchContext(): { [globalThis.Symbol.dispose](): void; } {
+        return this.client.batchContext();
+    }
+    // @sync-skip-block-end
+
     private async ensureInitialized(): Promise<void> {
         if (!this.initialized) {
             const response = await this.client.apiRequest("initialize", null);

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

@@ -23,6 +23,9 @@ import {
 } from "../options.ts";
 import type {
     APIMethodInfo,
+    APIRequest,
+    BatchRequestsParams,
+    BatchRequestsResponse,
     SourceFileResponseMethod,
 } from "../proto.ts";
 import {
@@ -47,6 +50,8 @@ export class Client {
     private options: ClientOptions;
     private connected = false;
     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;
 
     constructor(options: ClientOptions) {
         this.options = options;
@@ -158,15 +163,11 @@ export class Client {
         }
     }
 
-    async apiRequest<K extends keyof APIMethodInfo>(method: K, params: APIMethodInfo[K]["params"]): Promise<APIMethodInfo[K]["result"]> {
-        if (!this.connected) {
-            await this.connect();
-        }
+    private async sendRequestWithTiming<TResponse>(requestType: RequestType<unknown, TResponse, void>, params: unknown): Promise<TResponse> {
         if (!this.connection) {
             throw new Error("Connection not established");
         }
 
-        const requestType = new RequestType<unknown, APIMethodInfo[K]["result"], void>(method);
         if (!this.timing) {
             return this.connection.sendRequest(requestType, params);
         }
@@ -180,7 +181,7 @@ export class Client {
         const result = await this.connection.sendRequest(requestType, params);
         const roundTripMs = performance.now() - start;
         this.timing.record({
-            method,
+            method: requestType.method,
             roundTripMs,
             bytesSent,
             bytesReceived: result === undefined || result === null
@@ -190,6 +191,83 @@ export class Client {
         return result;
     }
 
+    private async doBatch(): Promise<void> {
+        this.nextBatch = undefined;
+        if (!this.batchedRequests.length) return;
+        const requests = this.batchedRequests;
+        this.batchedRequests = [];
+        try {
+            if (!this.connected) {
+                await this.connect();
+            }
+            if (!this.connection) {
+                throw new Error("Connection not established");
+            }
+
+            if (requests.length === 1) {
+                // send single queued requests directly instead of as a batched request
+                const requestType = new RequestType<unknown, unknown, void>(requests[0].method);
+                const response = await this.sendRequestWithTiming(requestType, requests[0].params);
+                requests[0].resolve(response);
+                return;
+            }
+
+            const requestType = new RequestType<unknown, BatchRequestsResponse, void>("batchRequests");
+            const params: BatchRequestsParams = { requests: requests.map(request => ({ method: request.method, params: request.params })) };
+            const response = await this.sendRequestWithTiming(requestType, params);
+            for (let i = 0; i < requests.length; i++) {
+                const { resolve, reject } = requests[i];
+                const item = response.responses[i];
+                if (item.error !== undefined) {
+                    reject(new Error(item.error));
+                }
+                else {
+                    resolve(item.result);
+                }
+            }
+        }
+        catch (error) {
+            for (const { reject } of requests) reject(error);
+        }
+    }
+
+    private scheduleImmediateBatch(): void {
+        if (this.nextBatch) return;
+        this.nextBatch = setImmediate(this.doBatch.bind(this));
+    }
+
+    batchContext(): { [Symbol.dispose](): void; } {
+        if (this.nextBatch === "manual") {
+            throw new Error("Already in a manual batch context");
+        }
+        if (this.nextBatch) {
+            clearImmediate(this.nextBatch);
+            this.doBatch(); // empty the queue before entering a manual batch context
+        }
+        this.nextBatch = "manual";
+        return {
+            [Symbol.dispose]: () => {
+                this.nextBatch = undefined;
+                this.scheduleImmediateBatch();
+            },
+        };
+    }
+
+    async apiRequest<K extends keyof APIMethodInfo>(method: K, params: APIMethodInfo[K]["params"]): Promise<APIMethodInfo[K]["result"]> {
+        if (!this.connected) {
+            await this.connect();
+        }
+        if (!this.connection) {
+            throw new Error("Connection not established");
+        }
+
+        const resultPromise = new Promise<APIMethodInfo[K]["result"]>((resolve, reject) => {
+            this.batchedRequests.push({ method, params, resolve, reject });
+            this.scheduleImmediateBatch();
+        });
+        return resultPromise;
+    }
+
     async apiRequestBinary<K extends SourceFileResponseMethod>(method: K, params: APIMethodInfo[K]["params"]): Promise<Uint8Array | undefined> {
         const response = await this.apiRequest(method, params);
         if (!response) return undefined;