fix: add timeout getMultichainClient to prevent remove connections fr… (#73)

* fix: add timeout getMultichainClient to prevent remove connections from dropping

* Move timeout management to Transport layer

* Fix retry

* fix lint after rebase

* fix: cleanup pendingRequests on timeout to prevent memory leak + tests

---------

Co-authored-by: Edouard Bougon <[email protected]>

Author: elribonazo

README.md

@@ -32,6 +32,15 @@ const result = await client.invokeMethod({
 await client.revokeSession();
 ```
 
+### Configuring Transport Timeout
+
+You can configure a default timeout (in milliseconds) for all requests made through the transport by passing the `defaultTimeout` option to `getDefaultTransport`:
+
+```typescript
+const transport = getDefaultTransport({ defaultTimeout: 5000 }); // 5 seconds timeout for all requests
+const client = getMultichainClient({ transport });
+```
+
 ## Extending RPC Types
 
 The client's RPC requests are strongly typed, enforcing the RPC methods and params to be defined ahead of usage. The client supports extending
@@ -78,6 +87,16 @@ const result = await client.invokeMethod({
 
 Transports handle the communication layer between your application and the wallet. You can create custom transports for different environments or communication methods.
 
+### Timeout Responsibility
+
+It is recommended that each custom transport implements its own request timeout mechanism rather than relying on higher layers. This ensures:
+
+- Transport-specific optimizations (e.g., aborting underlying network channels, clearing listeners)
+- Consistent error semantics (e.g., always throwing a dedicated `TransportTimeoutError` or custom error type)
+- Better resource cleanup in environments like browsers or workers
+
+Your `request` implementation should accept an optional `{ timeout?: number }` argument.
+
 ### Transport Interface
 
 A transport must implement the following interface:
@@ -87,29 +106,47 @@ type Transport = {
   connect: () => Promise<void>;
   disconnect: () => Promise<void>;
   isConnected: () => boolean;
-  request: <TRequest, TResponse>(request: TRequest) => Promise<TResponse>;
+  request: <TRequest, TResponse>(
+    request: TRequest,
+    options?: { timeout?: number }
+  ) => Promise<TResponse>;
   onNotification: (callback: (data: unknown) => void) => () => void;
 };
 ```
 
 ### Example: Custom Transport
 
 ```typescript
+import { TransportError, TransportTimeoutError } from '@metamask/multichain-api-client';
 import type { Transport, TransportRequest, TransportResponse } from '@metamask/multichain-api-client';
 
-export function getCustomTransport(): Transport {
+type CustomTransportOptions = {
+  defaultTimeout?: number; // ms
+};
+
+export function getCustomTransport(options: CustomTransportOptions = {}): Transport {
+  const { defaultTimeout = 5000 } = options;
+
   return {
     connect: async () => { ... },
     disconnect: async () => { ... },
     isConnected: () => { ...},
-    request: async <TRequest extends TransportRequest, TResponse extends TransportResponse>( request: TRequest ): Promise<TResponse> => { ... },
+    request: async <TRequest extends TransportRequest, TResponse extends TransportResponse>( request: TRequest, { timeout }: { timeout?: number } = {}): Promise<TResponse> => { ... },
     onNotification: (callback: (data: unknown) => void) => { ... },
   };
 }
 
 // Usage
-const transport = getCustomTransport();
+const transport = getCustomTransport({ defaultTimeout: 8000 });
 const client = getMultichainClient({ transport });
+
+// Per-request override
+await client.invokeMethod({
+  scope: 'eip155:1',
+  request: { method: 'eth_chainId', params: [] },
+  // The transport's request implementation can expose a timeout override
+  { timeout: 10000  // 10 seconds timeout for this request only
+});
 ```
 
 ## Error Handling

src/helpers/utils.test.ts

@@ -1,65 +1,110 @@
 import { describe, expect, it } from 'vitest';
-import { withRetry } from './utils';
+import { withRetry, withTimeout } from './utils';
 
-/**
- * Mock function that returns a promise that resolves only when called again after a delay
- * This function mocks MetaMask Multichain API wallet_getSession method, where early calls may never resolve
- *
- * @returns A promise that resolves after a delay
- */
-function mockMultichainApiRequest() {
-  const startTime = Date.now();
+describe('utils', () => {
+  class CustomTimeoutError extends Error {}
+  class CustomError extends Error {}
 
-  // Delay for the first call to resolve
-  const successThresholdDelay = 300;
+  describe('withRetry', () => {
+    it('retries on thrown error until success', async () => {
+      let attempts = 0;
+      const fn = async () => {
+        attempts++;
+        if (attempts < 3) {
+          throw new Error('fail');
+        }
+        return 'ok';
+      };
+      const result = await withRetry(fn, { maxRetries: 5 });
+      expect(result).toBe('ok');
+      expect(attempts).toBe(3);
+    });
 
-  return async () => {
-    const callTime = Date.now();
-    if (callTime - startTime < successThresholdDelay) {
-      // Promise that never resolves
-      await new Promise(() => {});
-    }
-    return 'success';
-  };
-}
+    it('throws last error after exceeding maxRetries', async () => {
+      let attempts = 0;
+      const fn = async () => {
+        attempts++;
+        throw new Error('boom');
+      };
+      await expect(withRetry(fn, { maxRetries: 2 })).rejects.toThrow('boom');
+      // maxRetries=2 => attempts 0,1,2 (3 total)
+      expect(attempts).toBe(3);
+    });
 
-function mockThrowingFn() {
-  const startTime = Date.now();
+    it('retries only specific error class with delay', async () => {
+      let attempts = 0;
+      const fn = async () => {
+        attempts++;
+        if (attempts < 3) {
+          throw new CustomError('Custom Error');
+        }
+        return 'done';
+      };
+      const start = Date.now();
+      const result = await withRetry(fn, { maxRetries: 5, retryDelay: 20, timeoutErrorClass: CustomTimeoutError });
+      const elapsed = Date.now() - start;
+      expect(result).toBe('done');
+      expect(attempts).toBe(3);
+      // Two retries with ~20ms delay each (allow some tolerance)
+      expect(elapsed).toBeGreaterThanOrEqual(30);
+    });
 
-  // Delay for the first call to resolve
-  const successThresholdDelay = 300;
+    it('retries only TimeoutError class without delay', async () => {
+      let attempts = 0;
+      const fn = async () => {
+        attempts++;
+        if (attempts < 3) {
+          throw new CustomTimeoutError('Custom Error');
+        }
+        return 'done';
+      };
+      const start = Date.now();
+      const result = await withRetry(fn, { maxRetries: 5, retryDelay: 20, timeoutErrorClass: CustomTimeoutError });
+      const elapsed = Date.now() - start;
+      expect(result).toBe('done');
+      expect(attempts).toBe(3);
+      expect(elapsed).toBeLessThanOrEqual(20);
+    });
 
-  return async () => {
-    const callTime = Date.now();
-    if (callTime - startTime < successThresholdDelay) {
-      throw new Error('error');
-    }
-    return 'success';
-  };
-}
+    it('continues retrying even if non-timeout errors occur (no delay applied for them)', async () => {
+      const sequenceErrors = [new Error('other'), new CustomTimeoutError('timeout'), new CustomTimeoutError('timeout')];
+      let attempts = 0;
+      const fn = async () => {
+        if (attempts < sequenceErrors.length) {
+          const err = sequenceErrors[attempts];
+          attempts++;
+          throw err;
+        }
+        attempts++;
+        return 'final';
+      };
+      const result = await withRetry(fn, { maxRetries: 5, retryDelay: 10, timeoutErrorClass: CustomTimeoutError });
+      expect(result).toBe('final');
+      expect(attempts).toBe(4); // 3 fail + 1 success
+    });
+  });
 
-describe('utils', () => {
-  describe('withRetry', () => {
-    it('should retry a function until it succeeds', async () => {
-      const result = await withRetry(mockMultichainApiRequest(), { maxRetries: 4, requestTimeout: 100 });
-      expect(result).toBe('success');
+  describe('withTimeout', () => {
+    it('should resolve before timeout', async () => {
+      const result = await withTimeout(Promise.resolve('ok'), 1000);
+      expect(result).toBe('ok');
     });
 
-    it('should retry a function that never resolves until it succeeds', async () => {
-      await expect(
-        async () => await withRetry(mockMultichainApiRequest(), { maxRetries: 2, requestTimeout: 100 }),
-      ).rejects.toThrow('Timeout reached');
+    it('should reject after timeout', async () => {
+      await expect(withTimeout(new Promise(() => {}), 50)).rejects.toThrow('Timeout after 50ms');
     });
 
-    it('should retry a throwing function until it succeeds', async () => {
-      const result = await withRetry(mockThrowingFn(), { maxRetries: 4, requestTimeout: 100 });
-      expect(result).toBe('success');
+    it('should propagate rejection from promise', async () => {
+      await expect(withTimeout(Promise.reject(new Error('fail')), 1000)).rejects.toThrow('fail');
     });
 
-    it('should retry a throwing function until it succeeds', async () => {
-      await expect(
-        async () => await withRetry(mockThrowingFn(), { maxRetries: 2, requestTimeout: 100 }),
-      ).rejects.toThrow('error');
+    it('should use custom error from errorFactory', async () => {
+      await expect(withTimeout(new Promise(() => {}), 10, () => new CustomTimeoutError('custom'))).rejects.toThrow(
+        CustomTimeoutError,
+      );
+      await expect(withTimeout(new Promise(() => {}), 10, () => new CustomTimeoutError('custom'))).rejects.toThrow(
+        'custom',
+      );
     });
   });
 });

src/helpers/utils.ts

@@ -1,3 +1,6 @@
+// chrome is a global injected by browser extensions
+declare const chrome: any;
+
 /**
  * Detects if we're in a Chrome-like environment with extension support
  */
@@ -17,35 +20,57 @@ export async function withRetry<T>(
   fn: () => Promise<T>,
   options: {
     maxRetries?: number;
-    requestTimeout?: number;
     retryDelay?: number;
+    timeoutErrorClass?: new (...args: any[]) => Error;
   } = {},
 ): Promise<T> {
-  const { maxRetries = 10, requestTimeout = 200, retryDelay = requestTimeout } = options;
-  const errorMessage = 'Timeout reached';
+  const { maxRetries = 10, retryDelay = 200, timeoutErrorClass } = options;
 
   for (let attempt = 0; attempt <= maxRetries; attempt++) {
     try {
-      // Use Promise.race to implement timeout
-      const timeoutPromise = new Promise<never>((_, reject) => {
-        setTimeout(() => reject(new Error(errorMessage)), requestTimeout);
-      });
-
-      const result = await Promise.race([fn(), timeoutPromise]);
-      return result;
+      return await fn();
     } catch (error) {
       // If this was the last attempt, throw the error
       if (attempt >= maxRetries) {
         throw error;
       }
 
       // Wait before retrying (unless it was a timeout, then retry immediately)
-      if (error instanceof Error && error.message !== errorMessage) {
-        await new Promise((resolve) => setTimeout(resolve, retryDelay));
+      if (timeoutErrorClass && typeof timeoutErrorClass === 'function' && error instanceof timeoutErrorClass) {
+        continue;
       }
+
+      await new Promise((resolve) => setTimeout(resolve, retryDelay));
     }
   }
 
   // This should never be reached due to the throw in the loop
   throw new Error('Max retries exceeded');
 }
+
+/**
+ * Returns a promise that resolves or rejects like the given promise, but fails if the timeout is exceeded.
+ * @param promise - The promise to monitor
+ * @param timeoutMs - Maximum duration in ms
+ * @param errorFactory - Optional callback to generate a custom error on timeout
+ */
+export function withTimeout<T>(promise: Promise<T>, timeoutMs: number, errorFactory?: () => Error): Promise<T> {
+  return new Promise<T>((resolve, reject) => {
+    const timer = setTimeout(() => {
+      if (errorFactory) {
+        reject(errorFactory());
+      } else {
+        reject(new Error(`Timeout after ${timeoutMs}ms`));
+      }
+    }, timeoutMs);
+    promise
+      .then((value) => {
+        clearTimeout(timer);
+        resolve(value);
+      })
+      .catch((err) => {
+        clearTimeout(timer);
+        reject(err);
+      });
+  });
+}

src/index.ts

@@ -20,9 +20,14 @@ import type { Transport } from './types/transport';
  * const transport = getDefaultTransport({ extensionId: '...' });
  * ```
  */
-function getDefaultTransport(params: { extensionId?: string } = {}): Transport {
+function getDefaultTransport({
+  extensionId,
+  defaultTimeout,
+}: { extensionId?: string; defaultTimeout?: number } = {}): Transport {
   const isChrome = isChromeRuntime();
-  return isChrome ? getExternallyConnectableTransport(params) : getWindowPostMessageTransport();
+  return isChrome
+    ? getExternallyConnectableTransport({ extensionId, defaultTimeout })
+    : getWindowPostMessageTransport({ defaultTimeout });
 }
 
 export { getMultichainClient, getDefaultTransport, getExternallyConnectableTransport, getWindowPostMessageTransport };