Reject pending async API requests when the connection closes

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

@@ -70,6 +70,7 @@ class RpcClient {
     private timing: TimingCollector | undefined;
     private timingDump: TimingDump | undefined;
     private closing: Promise<void> | undefined;
+    private draining: Promise<void> | undefined;
     private batchedRequests: { method: APIRequest["method"]; params: APIRequest["params"]; resolve: (value: unknown) => void; reject: (reason?: any) => void; }[] = [];
     private nextBatch: NodeJS.Immediate | "manual" | undefined;
 
@@ -129,9 +130,9 @@ class RpcClient {
                 resolve();
             });
 
-            this.connection = this.jsonrpc!.createStreamConnection(this.process.stdout!, this.process.stdin!);
-            this.registerFSCallbacks(this.connection, options.fs, fsConfiguration);
-            this.connection.listen();
+            const connection = this.setConnection(this.jsonrpc!.createStreamConnection(this.process.stdout!, this.process.stdin!));
+            this.registerFSCallbacks(connection, options.fs, fsConfiguration);
+            connection.listen();
         });
     }
 
@@ -140,8 +141,8 @@ class RpcClient {
 
         return new Promise((resolve, reject) => {
             this.socket = createConnection(options.pipe, () => {
-                this.connection = this.jsonrpc!.createStreamConnection(this.socket!, this.socket!);
-                this.connection.listen();
+                const connection = this.setConnection(this.jsonrpc!.createStreamConnection(this.socket!, this.socket!));
+                connection.listen();
                 this.connected = true;
                 resolve();
             });
@@ -152,6 +153,23 @@ class RpcClient {
         });
     }
 
+    private setConnection({ connection, drain }: ReturnType<typeof JsonRpc.createStreamConnection>): JsonRpc.MessageConnection {
+        this.connection = connection;
+        this.draining = undefined;
+        connection.onClose(() => {
+            // The JSON-RPC connection emits close without rejecting pending requests.
+            this.connectionClosing(connection);
+            const draining = drain().then(() => connection.dispose());
+            if (this.connection === connection) this.draining = draining;
+        });
+        return connection;
+    }
+
+    private connectionClosing(connection: JsonRpc.MessageConnection): void {
+        // A failed spawn may close its streams after a later connection is established.
+        if (this.connected && this.connection === connection) this.rejectQueuedRequests();
+    }
+
     private registerFSCallbacks(
         connection: JsonRpc.MessageConnection,
         fs: FileSystemCallbacks | undefined,
@@ -423,7 +441,7 @@ class RpcClient {
         return this.closing ??= this.closeWorker();
     }
 
-    private async closeWorker(): Promise<void> {
+    private rejectQueuedRequests(): void {
         this.closed = true;
         if (this.nextBatch && this.nextBatch !== "manual") {
             clearImmediate(this.nextBatch);
@@ -434,29 +452,28 @@ class RpcClient {
         for (const { reject } of requests) {
             reject(new Error("Client is closed"));
         }
+    }
+
+    private async closeWorker(): Promise<void> {
+        this.rejectQueuedRequests();
         await this.connecting?.catch(() => {}); // if connection is still in-progress, wait for it to finish before closing the connection
         try {
             if (this.timingDump && this.timing) {
-                const local = this.timing.getInfo();
                 let server: ServerTimingInfo | undefined;
                 let serverTimingError: string | undefined;
-                const connection = this.connection;
-                // Closing alone does not reject pending JSON-RPC requests.
-                const onClose = connection?.onClose(() => connection.dispose());
                 try {
                     if (this.connected) server = await this.fetchServerTiming();
                 }
                 catch (error) {
                     serverTimingError = String(error);
                 }
-                finally {
-                    onClose?.dispose();
-                }
-                this.timingDump.write(local, server, serverTimingError);
+                if (this.draining) await this.draining;
+                this.timingDump.write(this.timing.getInfo(), server, serverTimingError);
             }
         }
         finally {
             if (this.connection) {
+                if (this.draining) await this.draining;
                 this.connection.dispose();
                 this.connection = undefined;
             }

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

@@ -6,6 +6,7 @@ import {
     AbstractMessageBuffer,
     createMessageConnection,
     Disposable,
+    Message,
     type MessageConnection,
     NullLogger,
     RAL,
@@ -97,16 +98,74 @@ function installRuntime(): void {
 }
 
 /** Creates a JSON-RPC connection over Node.js streams. */
-export function createStreamConnection(readable: EventSource, writable: WritableStream): MessageConnection {
+export function createStreamConnection(readable: EventSource, writable: WritableStream): { connection: MessageConnection; drain(): Promise<void>; } {
     installRuntime();
+    const frames = new IncomingFrames();
+    const responses = new Map<string, Message>();
+    let closing = false;
+    let resolveDrain: (() => void) | undefined;
+
+    function finishDrain(): void {
+        if (frames.pending === 0 && responses.size === 0 && resolveDrain) {
+            resolveDrain();
+            resolveDrain = undefined;
+        }
+    }
+
+    function onClose(stream: EventSource, listener: () => void): Disposable {
+        return listen(stream, "close", () => {
+            // A write-side close can precede EOF on the read side. Only frames
+            // received before close belong to this drain; do not wait for later data.
+            closing = true;
+            listener();
+        });
+    }
+
+    const decoder = RAL().applicationJson.decoder;
     const reader = new ReadableStreamMessageReader({
-        onClose: listener => listen(readable, "close", listener),
+        onClose: listener => onClose(readable, listener),
         onError: listener => listen(readable, "error", listener),
         onEnd: listener => listen(readable, "end", listener),
-        onData: listener => listen(readable, "data", listener),
+        onData: listener =>
+            listen(readable, "data", (data: Uint8Array | string) => {
+                if (closing) return;
+                const bytes = typeof data === "string" ? Buffer.from(data, "utf8") : data;
+                frames.append(bytes);
+                listener(bytes);
+            }),
+    }, {
+        contentTypeDecoder: {
+            name: decoder.name,
+            async decode(data, options) {
+                try {
+                    return await decoder.decode(data, options);
+                }
+                catch (error) {
+                    frames.finish();
+                    finishDrain();
+                    throw error;
+                }
+            },
+        },
     });
+    // This connection does not consume partial-message events. In particular,
+    // avoid leaving their repeating timer alive after a truncated frame closes.
+    reader.partialMessageTimeout = 0;
+    const listenForMessages = reader.listen.bind(reader);
+    reader.listen = callback =>
+        listenForMessages(message => {
+            const key = responseKey(message);
+            if (key !== undefined) responses.set(key, message);
+            try {
+                callback(message);
+            }
+            finally {
+                frames.finish();
+                finishDrain();
+            }
+        });
     const writer = new WriteableStreamMessageWriter({
-        onClose: listener => listen(writable, "close", listener),
+        onClose: listener => onClose(writable, listener),
         onError: listener => listen(writable, "error", listener),
         onEnd: listener => listen(writable, "end", listener),
         write: (data: Uint8Array | string, encoding?: RAL.MessageBufferEncoding) =>
@@ -121,5 +180,110 @@ export function createStreamConnection(readable: EventSource, writable: Writable
             }),
         end: () => writable.end(),
     });
-    return createMessageConnection(reader, writer, NullLogger);
+    const connection = createMessageConnection(reader, writer, NullLogger, {
+        messageStrategy: {
+            handleMessage(message, next) {
+                try {
+                    return next(message);
+                }
+                finally {
+                    const key = responseKey(message);
+                    if (key !== undefined && responses.get(key) === message) responses.delete(key);
+                    finishDrain();
+                }
+            },
+        },
+    });
+    return {
+        connection,
+        drain: () =>
+            frames.pending === 0 && responses.size === 0
+                ? Promise.resolve()
+                : new Promise<void>(resolve => {
+                    resolveDrain = resolve;
+                }),
+    };
+}
+
+function responseKey(message: Message): string | undefined {
+    // Match the vendored connection's request-first classification and its
+    // response queue key, which stringifies both numeric and string IDs.
+    return !Message.isRequest(message) && Message.isResponse(message) && message.id !== null
+        ? String(message.id)
+        : undefined;
+}
+
+/** Counts complete Content-Length frames, retaining only bytes deferred after malformed headers. */
+class IncomingFrames {
+    private header = "";
+    private headerState = 0;
+    private remaining: number | undefined;
+    private deferred: Uint8Array | undefined;
+    private valid = true;
+    pending = 0;
+
+    append(data: Uint8Array): void {
+        if (!this.valid) return;
+        if (this.deferred) {
+            data = Buffer.concat([this.deferred, data]);
+            this.deferred = undefined;
+        }
+        let offset = 0;
+        while (offset < data.length) {
+            if (this.remaining !== undefined) {
+                const consumed = Math.min(this.remaining, data.length - offset);
+                this.remaining -= consumed;
+                offset += consumed;
+                if (this.remaining === 0) {
+                    this.pending++;
+                    this.remaining = undefined;
+                }
+                continue;
+            }
+            const value = data[offset++];
+            this.header += String.fromCharCode(value & 0x7f);
+            // Match AbstractMessageBuffer.tryReadHeaders on raw bytes, before
+            // ASCII decoding, including its reset on overlapping CR sequences.
+            switch (value) {
+                case 13:
+                    this.headerState = this.headerState === 0 ? 1 : this.headerState === 2 ? 3 : 0;
+                    break;
+                case 10:
+                    this.headerState = this.headerState === 1 ? 2 : this.headerState === 3 ? 4 : 0;
+                    break;
+                default:
+                    this.headerState = 0;
+            }
+            if (this.headerState !== 4) continue;
+            this.headerState = 0;
+            const header = this.header;
+            this.header = "";
+            // The reader consumes malformed headers, then retries the remaining
+            // buffered bytes on the next data event, including an empty event.
+            let contentLength: string | undefined;
+            for (const line of header.slice(0, -4).split("\r\n")) {
+                const separator = line.indexOf(":");
+                if (separator < 0) {
+                    this.deferred = offset < data.length ? data.subarray(offset) : undefined;
+                    return;
+                }
+                if (line.slice(0, separator).toLowerCase() === "content-length") contentLength = line.slice(separator + 1).trim();
+            }
+            const length = contentLength ? parseInt(contentLength) : NaN;
+            if (Number.isNaN(length)) {
+                this.deferred = offset < data.length ? data.subarray(offset) : undefined;
+                return;
+            }
+            if (!Number.isSafeInteger(length) || length < 0) {
+                this.valid = false;
+                return;
+            }
+            if (length === 0) this.pending++;
+            else this.remaining = length;
+        }
+    }
+
+    finish(): void {
+        if (this.pending > 0) this.pending--;
+    }
 }