Wire request timeout and 429/5xx retry onto Pinecone I/O hot path (#225)

* Wire request timeout and 429/5xx retry onto Pinecone I/O hot path

* Propagated app timeout errors from hybrid search.

* fixed some nit findings

* fixed format error

* add mapPc function

* - keep the SDK methods inside the guard
- added some tests

---------

Co-authored-by: zho <jornathanm910923@gmail.com>

Author: jonathanMLDev

src/core/pinecone-client.test.ts

@@ -1,4 +1,4 @@
-import { describe, it, expect, beforeEach, vi } from 'vitest';
+import { describe, it, expect, beforeEach, vi, afterEach } from 'vitest';
 import { PineconeClient } from './pinecone-client.js';
 import { resolveConfig } from './config.js';
 import type { SearchableIndex, PineconeHit } from '../types.js';
@@ -21,6 +21,16 @@ function stubPineconeClient(client: PineconeClient): PineconeClientMethodStubs {
   return client as unknown as PineconeClientMethodStubs;
 }
 
+function stubDualLegSearchFailure(testClient: PineconeClientMethodStubs, searchError: Error): void {
+  testClient.ensureIndexes = async () => ({
+    denseIndex: {} as SearchableIndex,
+    sparseIndex: {} as SearchableIndex,
+  });
+  testClient.searchIndex = async () => {
+    throw searchError;
+  };
+}
+
 describe('PineconeClient', () => {
   let client: PineconeClient;
 
@@ -32,6 +42,10 @@ describe('PineconeClient', () => {
     });
   });
 
+  afterEach(() => {
+    vi.useRealTimers();
+  });
+
   describe('constructor', () => {
     it('should initialize with provided config', () => {
       expect(client).toBeDefined();
@@ -143,14 +157,7 @@ describe('PineconeClient', () => {
 
     it('should throw when both dense and sparse searches fail', async () => {
       const testClient = stubPineconeClient(client);
-
-      testClient.ensureIndexes = async () => ({
-        denseIndex: {} as SearchableIndex,
-        sparseIndex: {} as SearchableIndex,
-      });
-      testClient.searchIndex = async () => {
-        throw new Error('index failure');
-      };
+      stubDualLegSearchFailure(testClient, new Error('index failure'));
 
       await expect(
         client.query({
@@ -161,6 +168,23 @@ describe('PineconeClient', () => {
         })
       ).rejects.toThrow('Hybrid search failed: both dense and sparse index searches failed.');
     });
+
+    it('propagates app timeout when both dense and sparse searches fail', async () => {
+      const testClient = stubPineconeClient(client);
+      stubDualLegSearchFailure(
+        testClient,
+        new Error('Timeout after 50ms while waiting for search')
+      );
+
+      await expect(
+        client.query({
+          query: 'hybrid search',
+          namespace: 'test',
+          topK: 5,
+          useReranking: false,
+        })
+      ).rejects.toThrow('Timeout after 50ms while waiting for search');
+    });
   });
 
   describe('count', () => {
@@ -443,6 +467,27 @@ describe('PineconeClient', () => {
       );
     });
 
+    it('passes configured requestTimeoutMs to search call sites', async () => {
+      vi.useFakeTimers();
+      const timeoutClient = new PineconeClient({
+        apiKey: 'test-api-key',
+        indexName: 'test-index',
+        requestTimeoutMs: 50,
+      });
+      const testClient = stubPineconeClient(timeoutClient);
+      const search = vi.fn(() => new Promise(() => {}));
+      const sparseRef = { search } as SearchableIndex;
+      testClient.ensureIndexes = async () => ({
+        denseIndex: {} as SearchableIndex,
+        sparseIndex: sparseRef,
+      });
+
+      const p = timeoutClient.keywordSearch({ query: 'q', namespace: 'n' });
+      const assertion = expect(p).rejects.toThrow(/Timeout after 50ms while waiting for search/);
+      await vi.advanceTimersByTimeAsync(50);
+      await assertion;
+    });
+
     it('searches sparse index only and maps hits', async () => {
       const testClient = stubPineconeClient(client);
       const denseRef = {} as SearchableIndex;

src/core/pinecone-client.ts

@@ -25,6 +25,7 @@ import {
   sliceMergedHitsToSearchResults,
 } from './pinecone/search.js';
 import { rerankResults as rerankResultsImpl } from './pinecone/rerank.js';
+import { isAppTimeoutError } from './server/retry.js';
 
 export class PineconeClient {
   private readonly rerankModel: string | undefined;
@@ -105,7 +106,15 @@ export class PineconeClient {
     metadataFilter?: Record<string, unknown>,
     options?: { fields?: string[] }
   ): Promise<PineconeHit[]> {
-    return searchIndexImpl(index, query, topK, namespace, metadataFilter, options);
+    return searchIndexImpl(
+      index,
+      query,
+      topK,
+      namespace,
+      metadataFilter,
+      options,
+      this.indexSession.getRequestTimeoutMs()
+    );
   }
 
   async query(params: QueryParams): Promise<HybridQueryResult> {
@@ -149,6 +158,8 @@ export class PineconeClient {
       logError('Sparse index search failed', sparseResult.reason);
     }
     if (denseResult.status === 'rejected' && sparseResult.status === 'rejected') {
+      if (isAppTimeoutError(denseResult.reason)) throw denseResult.reason;
+      if (isAppTimeoutError(sparseResult.reason)) throw sparseResult.reason;
       throw new Error('Hybrid search failed: both dense and sparse index searches failed.');
     }
 
@@ -179,7 +190,8 @@ export class PineconeClient {
         this.rerankModel,
         query,
         mergedResults,
-        topK
+        topK,
+        this.indexSession.getRequestTimeoutMs()
       );
       documents = rerankOut.results;
       degraded = rerankOut.degraded;

src/core/pinecone/indexes.test.ts

@@ -1,4 +1,4 @@
-import { describe, it, expect, vi } from 'vitest';
+import { describe, it, expect, vi, afterEach } from 'vitest';
 import { PineconeIndexSession } from './indexes.js';
 import type { SearchableIndex } from '../../types.js';
 
@@ -29,6 +29,60 @@ class ThrowingEnsureSession extends PineconeIndexSession {
   }
 }
 
+/** Session with a short timeout for fake-timer I/O deadline tests. */
+class TimeoutIndexSession extends PineconeIndexSession {
+  constructor(
+    private readonly pair: { dense: SearchableIndex; sparse: SearchableIndex },
+    requestTimeoutMs = 50
+  ) {
+    super('test-api-key', 'test-index', undefined, requestTimeoutMs);
+  }
+
+  override async ensureIndexes(): Promise<{
+    denseIndex: SearchableIndex;
+    sparseIndex: SearchableIndex;
+  }> {
+    return { denseIndex: this.pair.dense, sparseIndex: this.pair.sparse };
+  }
+}
+
+afterEach(() => {
+  vi.useRealTimers();
+});
+
+/** Mock describeIndexStats that fails when the SDK method is invoked detached. */
+function makeThisSensitiveIndex(
+  result: {
+    namespaces?: Record<string, { recordCount?: number }>;
+    dimension?: number;
+  } = {}
+): SearchableIndex {
+  const index = {
+    describeIndexStats() {
+      if (this !== index) {
+        throw new TypeError('detached SDK method: wrong this binding');
+      }
+      return Promise.resolve(result);
+    },
+  };
+  return index as unknown as SearchableIndex;
+}
+
+/** Mock namespace().query that fails when the SDK method is invoked detached. */
+function makeThisSensitiveNamespaceHandle(
+  matches: Array<{ metadata?: Record<string, unknown> }> = []
+) {
+  const handle = {
+    query() {
+      if (this !== handle) {
+        throw new TypeError('detached SDK method: wrong this binding');
+      }
+      return Promise.resolve({ matches });
+    },
+  };
+  return handle;
+}
+
 describe('PineconeIndexSession', () => {
   describe('listNamespacesFromKeywordIndex', () => {
     it('returns namespace rows when describeIndexStats succeeds', async () => {
@@ -66,6 +120,63 @@ describe('PineconeIndexSession', () => {
         expect(result.error).toContain('stats unavailable');
       }
     });
+
+    it('retries describeIndexStats on 503 then succeeds', async () => {
+      let n = 0;
+      const sparse = {
+        describeIndexStats: vi.fn().mockImplementation(async () => {
+          n++;
+          if (n < 2) throw new Error('HTTP 503');
+          return { namespaces: { papers: { recordCount: 7 } } };
+        }),
+      } as unknown as SearchableIndex;
+      const session = new PineconeIndexSessionTestDouble({
+        dense: {} as SearchableIndex,
+        sparse,
+      });
+
+      const result = await session.listNamespacesFromKeywordIndex();
+
+      expect(result.ok).toBe(true);
+      if (result.ok) {
+        expect(result.namespaces).toEqual([{ namespace: 'papers', recordCount: 7 }]);
+      }
+      expect(sparse.describeIndexStats).toHaveBeenCalledTimes(2);
+    });
+
+    it('invokes describeIndexStats on the sparse index receiver through runIo', async () => {
+      const sparse = makeThisSensitiveIndex({
+        namespaces: { papers: { recordCount: 42 } },
+      });
+      const session = new PineconeIndexSessionTestDouble({
+        dense: {} as SearchableIndex,
+        sparse,
+      });
+
+      const result = await session.listNamespacesFromKeywordIndex();
+
+      expect(result.ok).toBe(true);
+      if (result.ok) {
+        expect(result.namespaces).toEqual([{ namespace: 'papers', recordCount: 42 }]);
+      }
+    });
+
+    it('times out describeIndexStats at requestTimeoutMs', async () => {
+      vi.useFakeTimers();
+      const sparse = {
+        describeIndexStats: vi.fn(() => new Promise(() => {})),
+      } as unknown as SearchableIndex;
+      const session = new TimeoutIndexSession({ dense: {} as SearchableIndex, sparse });
+      const p = session.listNamespacesFromKeywordIndex();
+      const assertion = expect(p).resolves.toMatchObject({
+        ok: false,
+        error: expect.stringMatching(
+          /Timeout after 50ms while waiting for describeIndexStats-sparse/
+        ),
+      });
+      await vi.advanceTimersByTimeAsync(50);
+      await assertion;
+    });
   });
 
   describe('listNamespacesWithMetadata', () => {
@@ -105,6 +216,28 @@ describe('PineconeIndexSession', () => {
       });
     });
 
+    it('invokes namespace().query on the namespace receiver through runIo when sampling', async () => {
+      const nsHandle = makeThisSensitiveNamespaceHandle([
+        { metadata: { title: 'T', tags: ['a'] } },
+      ]);
+      const dense = makeThisSensitiveIndex({
+        namespaces: { ns1: { recordCount: 2 } },
+        dimension: 4,
+      }) as SearchableIndex & { namespace: (name: string) => typeof nsHandle };
+      dense.namespace = () => nsHandle;
+      const session = new PineconeIndexSessionTestDouble({
+        dense,
+        sparse: {} as SearchableIndex,
+      });
+
+      const rows = await session.listNamespacesWithMetadata();
+      expect(rows.namespaces).toHaveLength(1);
+      expect(rows.namespaces[0]?.metadata).toMatchObject({
+        title: 'string',
+        tags: 'string[]',
+      });
+    });
+
     it('samples metadata when records exist and namespace.query returns matches', async () => {
       const dense = {
         describeIndexStats: vi.fn().mockResolvedValue({
@@ -141,6 +274,33 @@ describe('PineconeIndexSession', () => {
       expect(rows.namespaces[0]?.schema_source).toBe('sampled');
     });
 
+    it('times out metadata sampling at requestTimeoutMs', async () => {
+      vi.useFakeTimers();
+      const dense = {
+        describeIndexStats: vi.fn().mockResolvedValue({
+          namespaces: { ns1: { recordCount: 2 } },
+          dimension: 4,
+        }),
+        namespace: vi.fn().mockReturnValue({
+          query: vi.fn(() => new Promise(() => {})),
+        }),
+      } as unknown as SearchableIndex;
+      const session = new TimeoutIndexSession({ dense, sparse: {} as SearchableIndex });
+      const p = session.listNamespacesWithMetadata();
+      const assertion = expect(p).resolves.toMatchObject({
+        namespaces: [
+          {
+            namespace: 'ns1',
+            recordCount: 2,
+            metadata: {},
+            schema_source: 'sampled',
+          },
+        ],
+      });
+      await vi.advanceTimersByTimeAsync(50);
+      await assertion;
+    });
+
     it('uses declared schema and skips sampling when declaredSchemas provides one', async () => {
       const query = vi.fn();
       const dense = {
@@ -235,6 +395,16 @@ describe('PineconeIndexSession', () => {
   });
 
   describe('checkIndexes', () => {
+    it('invokes describeIndexStats on dense and sparse receivers through runIo', async () => {
+      const dense = makeThisSensitiveIndex();
+      const sparse = makeThisSensitiveIndex();
+      const session = new PineconeIndexSessionTestDouble({ dense, sparse });
+
+      const result = await session.checkIndexes();
+      expect(result.ok).toBe(true);
+      expect(result.errors).toHaveLength(0);
+    });
+
     it('returns ok when describeIndexStats succeeds for dense and sparse', async () => {
       const dense = {
         describeIndexStats: vi.fn().mockResolvedValue({}),
@@ -348,5 +518,31 @@ describe('PineconeIndexSession', () => {
         /Timeout after 50ms while waiting for fetchRecordFields/
       );
     });
+
+    it('retries fetch on 503 then succeeds', async () => {
+      let n = 0;
+      class RetryFetchSession extends PineconeIndexSession {
+        override ensureClient() {
+          return {
+            index: () => ({
+              fetch: vi.fn().mockImplementation(async () => {
+                n++;
+                if (n < 2) throw new Error('HTTP 503');
+                return {
+                  records: {
+                    schema_manifest: { metadata: { chunk_text: '{"ok":true}' } },
+                  },
+                };
+              }),
+            }),
+          } as never;
+        }
+      }
+
+      const session = new RetryFetchSession();
+      const fields = await session.fetchRecordFields('_mcp_config', 'schema_manifest');
+      expect(fields?.chunk_text).toBe('{"ok":true}');
+      expect(n).toBe(2);
+    });
   });
 });