Add a withTimeout function to set a timeout for a request

src/extension/__tests__/rpc.test.ts

@@ -409,9 +409,9 @@ test("times out if no message received within configured timeout", async () => {
   const adapter = createTestAdapter();
   const client = createRpcClient<Message>(adapter);
 
-  const promise = client.request("add", { x: 1, y: 2 }, { timeoutMs: 1000 });
+  const promise = client.withTimeout(1000).request("add", { x: 1, y: 2 });
 
-  jest.advanceTimersByTime(1_000);
+  jest.advanceTimersByTime(1000);
 
   await expect(promise).rejects.toEqual(
     new Error("Timeout waiting for message")

src/extension/rpc.ts

@@ -15,10 +15,10 @@ import {
 type MessageCollection = Record<string, (...parameters: [SafeAny]) => SafeAny>;
 
 export interface RpcClient<Messages extends MessageCollection> {
+  withTimeout: (timeoutMs: number) => RpcClient<Messages>;
   request: <TName extends keyof Messages & string>(
     name: TName,
-    params: Parameters<Messages[TName]>[0],
-    options?: { timeoutMs?: number }
+    params: Parameters<Messages[TName]>[0]
   ) => Promise<Awaited<ReturnType<Messages[TName]>>>;
 }
 
@@ -27,15 +27,24 @@ const DEFAULT_TIMEOUT = 30_000;
 export function createRpcClient<Messages extends MessageCollection>(
   adapter: MessageAdapter<RPCRequestMessage>
 ): RpcClient<Messages> {
-  return {
-    request: (name, params, options) => {
+  let timeout = DEFAULT_TIMEOUT;
+
+  const client: RpcClient<Messages> = {
+    withTimeout: (timeoutMs) => {
+      timeout = timeoutMs;
+      return client;
+    },
+    request: (name, params) => {
+      const configuredTimeout = timeout;
+      timeout = DEFAULT_TIMEOUT;
+
       return new Promise<SafeAny>((resolve, reject) => {
         const id = createId();
 
         const timeout = setTimeout(() => {
           removeListener();
           reject(new Error("Timeout waiting for message"));
-        }, options?.timeoutMs ?? DEFAULT_TIMEOUT);
+        }, configuredTimeout);
 
         const removeListener = adapter.addListener((message) => {
           if (!isRPCResponseMessage(message) || message.sourceId !== id) {
@@ -62,6 +71,8 @@ export function createRpcClient<Messages extends MessageCollection>(
       });
     },
   };
+
+  return client;
 }
 
 export function createRpcHandler<Messages extends MessageCollection>(