diff --git a/src/core/pinecone-client.test.ts b/src/core/pinecone-client.test.ts index 437a1ca..3cd6831 100644 --- a/src/core/pinecone-client.test.ts +++ b/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; diff --git a/src/core/pinecone-client.ts b/src/core/pinecone-client.ts index 6c67e40..8da14f8 100644 --- a/src/core/pinecone-client.ts +++ b/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, options?: { fields?: string[] } ): Promise { - return searchIndexImpl(index, query, topK, namespace, metadataFilter, options); + return searchIndexImpl( + index, + query, + topK, + namespace, + metadataFilter, + options, + this.indexSession.getRequestTimeoutMs() + ); } async query(params: QueryParams): Promise { @@ -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; diff --git a/src/core/pinecone/indexes.test.ts b/src/core/pinecone/indexes.test.ts index 1b77fd4..7b28919 100644 --- a/src/core/pinecone/indexes.test.ts +++ b/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; + 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 }> = [] +) { + 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); + }); }); }); diff --git a/src/core/pinecone/indexes.ts b/src/core/pinecone/indexes.ts index 65beb0c..a0ff39d 100644 --- a/src/core/pinecone/indexes.ts +++ b/src/core/pinecone/indexes.ts @@ -4,7 +4,7 @@ import { Pinecone } from '@pinecone-database/pinecone'; import { DEFAULT_REQUEST_TIMEOUT_MS } from '../config.js'; -import { withTimeout } from '../server/retry.js'; +import { runWithPolicy, type PolicyOptions } from '../server/retry.js'; import { error as logError, info as logInfo } from '../../logger.js'; import type { KeywordIndexNamespacesResult, @@ -12,6 +12,9 @@ import type { SearchableIndex, } from '../../types.js'; +/** Startup probe: fail fast on unreachable indexes instead of retrying. */ +const CHECK_INDEXES_IO_POLICY = { retries: 0 } as const satisfies Pick; + function inferMetadataFieldType(value: unknown): string { if (value === null || value === undefined) { return 'unknown'; @@ -57,6 +60,18 @@ export class PineconeIndexSession { return this.sparseIndexName ?? `${this.indexName}-sparse`; } + getRequestTimeoutMs(): number { + return this.requestTimeoutMs; + } + + private runIo( + label: string, + fn: () => Promise, + policy?: Pick + ): Promise { + return runWithPolicy(() => fn(), { timeoutMs: this.requestTimeoutMs, label, ...policy }); + } + /** Ensure Pinecone client is initialized */ ensureClient(): Pinecone { if (!this.pc) { @@ -103,9 +118,11 @@ export class PineconeIndexSession { async listNamespacesFromKeywordIndex(): Promise { try { const { sparseIndex } = await this.ensureIndexes(); - const stats = sparseIndex.describeIndexStats - ? await sparseIndex.describeIndexStats() - : undefined; + // SDK methods must be invoked on the index receiver inside runIo, not detached. + const stats = + typeof sparseIndex.describeIndexStats === 'function' + ? await this.runIo('describeIndexStats-sparse', () => sparseIndex.describeIndexStats!()) + : undefined; const namespaces = stats?.namespaces ?? {}; const rows = Object.entries(namespaces).map(([namespace, info]) => ({ namespace, @@ -134,9 +151,10 @@ export class PineconeIndexSession { const { denseIndex } = await this.ensureIndexes(); // Get index stats to find namespaces - const stats = denseIndex.describeIndexStats - ? await denseIndex.describeIndexStats() - : undefined; + const stats = + typeof denseIndex.describeIndexStats === 'function' + ? await this.runIo('describeIndexStats-dense', () => denseIndex.describeIndexStats!()) + : undefined; const namespaces = stats?.namespaces ? Object.keys(stats.namespaces) : []; const liveSet = new Set(namespaces); @@ -176,11 +194,13 @@ export class PineconeIndexSession { const nsObj: NamespaceHandle = denseIndex.namespace(ns); const sampleQuery = typeof nsObj.query === 'function' - ? await nsObj.query({ - topK: 5, - vector: Array(stats?.dimension ?? 1536).fill(0), - includeMetadata: true, - }) + ? await this.runIo('sampleNamespaceMetadata', () => + nsObj.query!({ + topK: 5, + vector: Array(stats?.dimension ?? 1536).fill(0), + includeMetadata: true, + }) + ) : { matches: undefined }; // Collect unique metadata fields and infer types (including string[]) @@ -238,26 +258,23 @@ export class PineconeIndexSession { * may appear on the record or inside metadata. */ async fetchRecordFields(namespace: string, id: string): Promise | null> { - return withTimeout( - async (_signal) => { - const pc = this.ensureClient(); - const response = await pc.index(this.indexName).fetch({ ids: [id], namespace }); - const record = response.records?.[id] as - (Record & { metadata?: Record }) | undefined; - if (!record) { - return null; - } - const metadata = record.metadata ?? {}; - const merged: Record = { ...metadata }; - for (const [k, v] of Object.entries(record)) { - if (k !== 'metadata' && k !== 'values' && k !== 'sparseValues' && !(k in merged)) { - merged[k] = v; - } + return this.runIo('fetchRecordFields', async () => { + const pc = this.ensureClient(); + const response = await pc.index(this.indexName).fetch({ ids: [id], namespace }); + const record = response.records?.[id] as + (Record & { metadata?: Record }) | undefined; + if (!record) { + return null; + } + const metadata = record.metadata ?? {}; + const merged: Record = { ...metadata }; + for (const [k, v] of Object.entries(record)) { + if (k !== 'metadata' && k !== 'values' && k !== 'sparseValues' && !(k in merged)) { + merged[k] = v; } - return merged; - }, - { timeoutMs: this.requestTimeoutMs, label: 'fetchRecordFields' } - ); + } + return merged; + }); } /** @@ -277,7 +294,11 @@ export class PineconeIndexSession { ); } else { try { - await denseIndex.describeIndexStats(); + await this.runIo( + 'describeIndexStats-dense', + () => denseIndex.describeIndexStats!(), + CHECK_INDEXES_IO_POLICY + ); } catch (err) { const msg = err instanceof Error ? err.message : String(err); errors.push(`Dense index "${denseName}": ${msg}`); @@ -290,7 +311,11 @@ export class PineconeIndexSession { ); } else { try { - await sparseIndex.describeIndexStats(); + await this.runIo( + 'describeIndexStats-sparse', + () => sparseIndex.describeIndexStats!(), + CHECK_INDEXES_IO_POLICY + ); } catch (err) { const msg = err instanceof Error ? err.message : String(err); errors.push(`Sparse index "${sparseName}": ${msg}`); diff --git a/src/core/pinecone/rerank.test.ts b/src/core/pinecone/rerank.test.ts index 6156d30..c5f94b0 100644 --- a/src/core/pinecone/rerank.test.ts +++ b/src/core/pinecone/rerank.test.ts @@ -7,7 +7,7 @@ const sampleMerged: MergedHit[] = [ ]; function makePc(rerank: ReturnType) { - return { inference: { rerank } } as Parameters[0]; + return { inference: { rerank } } as unknown as Parameters[0]; } describe('rerankResults', () => { @@ -52,6 +52,40 @@ describe('rerankResults', () => { expect(out.results[0]?.content).toBe('hello'); }); + it('retries on 429 then succeeds without degrading', async () => { + let n = 0; + const rerank = vi.fn().mockImplementation(async () => { + n++; + if (n < 2) throw new Error('HTTP 429'); + return { + data: [ + { + score: 0.99, + document: { _id: '1', chunk_text: 'hello', metadata: { k: 'v' } }, + }, + ], + }; + }); + const pc = makePc(rerank); + + const out = await rerankResults(pc, 'm', 'q', sampleMerged, 5); + + expect(out.degraded).toBe(false); + expect(out.results[0]?.reranked).toBe(true); + expect(rerank).toHaveBeenCalledTimes(2); + }); + + it('degrades after exhausting retries on persistent 503', async () => { + const rerank = vi.fn().mockRejectedValue(new Error('HTTP 503')); + const pc = makePc(rerank); + + const out = await rerankResults(pc, 'm', 'q', sampleMerged, 5, 5000); + + expect(out.degraded).toBe(true); + expect(out.degradation_reason).toMatch(/^rerank_failed:/); + expect(rerank).toHaveBeenCalledTimes(3); + }); + it('returns empty reranked results with degraded=false when data is an empty array', async () => { const rerank = vi.fn().mockResolvedValue({ data: [] }); const pc = makePc(rerank); diff --git a/src/core/pinecone/rerank.ts b/src/core/pinecone/rerank.ts index 4e72cc9..28d7624 100644 --- a/src/core/pinecone/rerank.ts +++ b/src/core/pinecone/rerank.ts @@ -4,6 +4,8 @@ import type { Pinecone } from '@pinecone-database/pinecone'; import { error as logError, redactApiKey } from '../../logger.js'; +import { DEFAULT_REQUEST_TIMEOUT_MS } from '../config.js'; +import { runWithPolicy } from '../server/retry.js'; import type { MergedHit, SearchResult } from '../../types.js'; export type RerankOutcome = { @@ -22,28 +24,33 @@ export async function rerankResults( rerankModel: string, query: string, results: MergedHit[], - topN: number + topN: number, + requestTimeoutMs: number = DEFAULT_REQUEST_TIMEOUT_MS ): Promise { if (!results || results.length === 0) { return { results: [], degraded: false }; } try { - const rerankResult = await pc.inference.rerank({ - model: rerankModel, - query, - // The Pinecone SDK types constrain document values to `Record`, - // but the underlying HTTP API accepts any JSON value. We pass MergedHit objects - // (metadata may contain number/boolean/string[]) and only `chunk_text` — which is - // always a string — is accessed via rankFields. The double cast via `as unknown` - // is intentional: it bypasses the SDK's over-narrow type without stringifying - // metadata values that we need to read back from the returned documents. - documents: results as unknown as (string | Record)[], - topN, - rankFields: ['chunk_text'], - returnDocuments: true, - parameters: { truncate: 'END' }, - }); + const rerankResult = await runWithPolicy( + async (_signal) => + pc.inference.rerank({ + model: rerankModel, + query, + // The Pinecone SDK types constrain document values to `Record`, + // but the underlying HTTP API accepts any JSON value. We pass MergedHit objects + // (metadata may contain number/boolean/string[]) and only `chunk_text` — which is + // always a string — is accessed via rankFields. The double cast via `as unknown` + // is intentional: it bypasses the SDK's over-narrow type without stringifying + // metadata values that we need to read back from the returned documents. + documents: results as unknown as (string | Record)[], + topN, + rankFields: ['chunk_text'], + returnDocuments: true, + parameters: { truncate: 'END' }, + }), + { timeoutMs: requestTimeoutMs, label: 'rerank' } + ); const reranked: SearchResult[] = []; for (const item of rerankResult.data || []) { diff --git a/src/core/pinecone/search.test.ts b/src/core/pinecone/search.test.ts index 31f322a..61f5e05 100644 --- a/src/core/pinecone/search.test.ts +++ b/src/core/pinecone/search.test.ts @@ -1,4 +1,4 @@ -import { describe, it, expect, vi } from 'vitest'; +import { describe, it, expect, vi, afterEach } from 'vitest'; import { searchIndex, mergeResults, @@ -8,6 +8,10 @@ import { } from './search.js'; import type { PineconeHit, SearchableIndex } from '../../types.js'; +afterEach(() => { + vi.useRealTimers(); +}); + describe('searchIndex', () => { it('uses index.search when available and passes fields', async () => { const search = vi.fn().mockResolvedValue({ @@ -59,6 +63,36 @@ describe('searchIndex', () => { 'Pinecone search failed for namespace "my-ns": boom' ); }); + + it('retries on 503 then succeeds', async () => { + let n = 0; + const search = vi.fn().mockImplementation(async () => { + n++; + if (n < 2) throw new Error('HTTP 503'); + return { result: { hits: [{ _id: '1', _score: 1, fields: {} }] } }; + }); + const index = { search } as unknown as SearchableIndex; + const hits = await searchIndex(index, 'hi', 5); + expect(hits).toHaveLength(1); + expect(search).toHaveBeenCalledTimes(2); + }); + + it('does not retry on 401', async () => { + const search = vi.fn().mockRejectedValue(new Error('HTTP 401')); + const index = { search } as unknown as SearchableIndex; + await expect(searchIndex(index, 'hi', 5)).rejects.toThrow(/401/); + expect(search).toHaveBeenCalledTimes(1); + }); + + it('times out at requestTimeoutMs and preserves the timeout error prefix', async () => { + vi.useFakeTimers(); + const search = vi.fn(() => new Promise(() => {})); + const index = { search } as unknown as SearchableIndex; + const p = searchIndex(index, 'hi', 5, undefined, undefined, undefined, 50); + const assertion = expect(p).rejects.toThrow(/^Timeout after 50ms while waiting for search/); + await vi.advanceTimersByTimeAsync(50); + await assertion; + }); }); describe('mergeResults', () => { diff --git a/src/core/pinecone/search.ts b/src/core/pinecone/search.ts index 44ebe0e..5e8d5f3 100644 --- a/src/core/pinecone/search.ts +++ b/src/core/pinecone/search.ts @@ -4,6 +4,8 @@ import { COUNT_FIELDS, COUNT_TOP_K } from '../../constants.js'; import { debug as logDebug, warn as logWarn } from '../../logger.js'; +import { DEFAULT_REQUEST_TIMEOUT_MS } from '../config.js'; +import { isAppTimeoutError, runWithPolicy } from '../server/retry.js'; import type { CountResult, MergedHit, @@ -23,7 +25,8 @@ export async function searchIndex( topK: number, namespace?: string, metadataFilter?: Record, - options?: { fields?: string[] } + options?: { fields?: string[] }, + requestTimeoutMs: number = DEFAULT_REQUEST_TIMEOUT_MS ): Promise { // Build query payload in the same shape as Python implementation. const queryPayload: Record = { @@ -38,42 +41,50 @@ export async function searchIndex( } try { - // Preferred path: Pinecone search API. - if (typeof index.search === 'function') { - const searchOpts: { - namespace?: string; - query: Record; - fields?: string[]; - } = { - namespace, - query: queryPayload, - }; - if (options?.fields?.length) { - searchOpts.fields = options.fields; - } - const result = await index.search(searchOpts); - return result?.result?.hits || []; - } + return await runWithPolicy( + async (_signal) => { + // Preferred path: Pinecone search API. + if (typeof index.search === 'function') { + const searchOpts: { + namespace?: string; + query: Record; + fields?: string[]; + } = { + namespace, + query: queryPayload, + }; + if (options?.fields?.length) { + searchOpts.fields = options.fields; + } + const result = await index.search(searchOpts); + return result?.result?.hits || []; + } - // Backward-compatible fallback for older API shapes. - const target = namespace && index.namespace ? index.namespace(namespace) : index; - const queryParams: { query: Record; fields?: string[] } = { - query: { - topK, - inputs: { text: query }, + // Backward-compatible fallback for older API shapes. + const target = namespace && index.namespace ? index.namespace(namespace) : index; + const queryParams: { query: Record; fields?: string[] } = { + query: { + topK, + inputs: { text: query }, + }, + }; + if (metadataFilter !== undefined) { + queryParams.query['filter'] = metadataFilter; + } + if (options?.fields?.length) { + queryParams.fields = options.fields; + } + const result = target.searchRecords + ? await target.searchRecords(queryParams) + : { result: { hits: [] as PineconeHit[] } }; + return result?.result?.hits || []; }, - }; - if (metadataFilter !== undefined) { - queryParams.query['filter'] = metadataFilter; - } - if (options?.fields?.length) { - queryParams.fields = options.fields; - } - const result = target.searchRecords - ? await target.searchRecords(queryParams) - : { result: { hits: [] as PineconeHit[] } }; - return result?.result?.hits || []; + { timeoutMs: requestTimeoutMs, label: 'search' } + ); } catch (error) { + if (isAppTimeoutError(error)) { + throw error; + } const errorMessage = error instanceof Error ? error.message : String(error); throw new Error( `Pinecone search failed for namespace "${namespace ?? 'default'}": ${errorMessage}` diff --git a/src/core/pinecone/searchable-index-conformance.test.ts b/src/core/pinecone/searchable-index-conformance.test.ts index a711300..e5c44ab 100644 --- a/src/core/pinecone/searchable-index-conformance.test.ts +++ b/src/core/pinecone/searchable-index-conformance.test.ts @@ -40,4 +40,24 @@ describe('SearchableIndex SDK conformance (#220)', () => { expect(typeof namespaceHandle[method]).toBe('function'); } ); + + // indexes.ts must call SDK methods on the index/namespace receiver inside runIo; + // detached calls throw against the real SDK (wpak-ai PR #225). + it('describeIndexStats throws when called detached from the index', () => { + const describeIndexStats = index.describeIndexStats as () => Promise; + expect(() => describeIndexStats()).toThrow( + /_describeIndexStats|Cannot read properties of undefined/ + ); + }); + + it('namespace().query rejects when called detached from the namespace handle', async () => { + const query = namespaceHandle.query as (opts: { + topK: number; + vector: number[]; + includeMetadata: boolean; + }) => Promise; + await expect(query({ topK: 1, vector: [0], includeMetadata: true })).rejects.toThrow( + /_queryCommand|Cannot read properties of undefined/ + ); + }); }); diff --git a/src/core/server/retry.test.ts b/src/core/server/retry.test.ts index 7850895..562b634 100644 --- a/src/core/server/retry.test.ts +++ b/src/core/server/retry.test.ts @@ -1,5 +1,12 @@ import { describe, expect, it, vi, afterEach } from 'vitest'; -import { withRetry, withTimeout, defaultShouldRetry } from './retry.js'; +import { + withRetry, + withTimeout, + defaultShouldRetry, + transientShouldRetry, + isAppTimeoutError, + runWithPolicy, +} from './retry.js'; afterEach(() => { vi.useRealTimers(); @@ -14,6 +21,74 @@ describe('defaultShouldRetry', () => { }); }); +describe('transientShouldRetry', () => { + it('retries on 503 in message', () => { + expect(transientShouldRetry(new Error('HTTP 503'))).toBe(true); + }); + + it('does not retry on 401', () => { + expect(transientShouldRetry(new Error('HTTP 401'))).toBe(false); + }); + + it('does not retry app-level withTimeout deadlines', () => { + expect(transientShouldRetry(new Error('Timeout after 50ms while waiting for search'))).toBe( + false + ); + }); +}); + +describe('isAppTimeoutError', () => { + it('matches withTimeout rejection messages', () => { + expect(isAppTimeoutError(new Error('Timeout after 50ms while waiting for search'))).toBe(true); + }); +}); + +describe('runWithPolicy', () => { + it('retries then succeeds on transient 503', async () => { + let n = 0; + const v = await runWithPolicy( + async () => { + n++; + if (n < 2) throw new Error('HTTP 503'); + return 'done'; + }, + { timeoutMs: 1000, label: 'test', retries: 2, backoffMs: 1 } + ); + expect(v).toBe('done'); + expect(n).toBe(2); + }); + + it('fails fast on non-retryable 401', async () => { + let n = 0; + await expect( + runWithPolicy( + async () => { + n++; + throw new Error('HTTP 401'); + }, + { timeoutMs: 1000, label: 'test', retries: 2, backoffMs: 1 } + ) + ).rejects.toThrow('HTTP 401'); + expect(n).toBe(1); + }); + + it('does not retry when the per-attempt timeout fires', async () => { + vi.useFakeTimers(); + let n = 0; + const p = runWithPolicy( + async () => { + n++; + return new Promise(() => {}); + }, + { timeoutMs: 50, label: 'test', retries: 2, backoffMs: 1 } + ); + const assertion = expect(p).rejects.toThrow(/Timeout after 50ms/); + await vi.advanceTimersByTimeAsync(50); + await assertion; + expect(n).toBe(1); + }); +}); + describe('withTimeout', () => { it('aborts signal when deadline passes', async () => { vi.useFakeTimers(); diff --git a/src/core/server/retry.ts b/src/core/server/retry.ts index 5173f7d..a8a84db 100644 --- a/src/core/server/retry.ts +++ b/src/core/server/retry.ts @@ -7,6 +7,16 @@ * waiter still rejects immediately on timeout. */ +import { warn, redactApiKey } from '../../logger.js'; + +/** Matches {@link withTimeout} rejection message prefix; used by tool-error and callers. */ +export const APP_TIMEOUT_PATTERN = /^Timeout after \d+ms while waiting for /i; + +export function isAppTimeoutError(error: unknown): boolean { + const msg = error instanceof Error ? error.message : String(error); + return APP_TIMEOUT_PATTERN.test(msg); +} + /** Retry policy. */ export interface RetryOptions { /** Total number of attempts after the first try. Default 2. */ @@ -27,6 +37,14 @@ export interface TimeoutOptions { label?: string; } +/** Per-call timeout + transient retry policy for outbound Pinecone I/O. */ +export interface PolicyOptions { + timeoutMs: number; + label: string; + retries?: number; + backoffMs?: number; +} + /** Default predicate: retry on common transient HTTP statuses (429/5xx) and network-ish messages. */ export function defaultShouldRetry(error: unknown): boolean { if (error instanceof Error) { @@ -43,6 +61,12 @@ export function defaultShouldRetry(error: unknown): boolean { return false; } +/** Retry 429/5xx + network errors; do NOT retry app-level {@link withTimeout} deadlines. */ +export function transientShouldRetry(error: unknown): boolean { + if (isAppTimeoutError(error)) return false; + return defaultShouldRetry(error); +} + /** * Run `fn` and retry on transient failures. * @@ -102,3 +126,22 @@ export async function withTimeout( void fnPromise.catch(() => {}); } } + +/** + * Compose per-attempt timeout with bounded transient retry for Pinecone I/O. + * Each retry gets a fresh timeout window. + */ +export function runWithPolicy( + fn: (signal: AbortSignal) => Promise, + options: PolicyOptions +): Promise { + return withRetry(() => withTimeout(fn, { timeoutMs: options.timeoutMs, label: options.label }), { + retries: options.retries, + backoffMs: options.backoffMs, + shouldRetry: transientShouldRetry, + onRetry: (attempt, error) => { + const msg = redactApiKey(error instanceof Error ? error.message : String(error)); + warn(`Retrying ${options.label} (attempt ${attempt})`, msg); + }, + }); +} diff --git a/src/core/server/tool-error.ts b/src/core/server/tool-error.ts index 0d3f034..f299ce4 100644 --- a/src/core/server/tool-error.ts +++ b/src/core/server/tool-error.ts @@ -5,6 +5,7 @@ import { z } from 'zod'; import { getLogLevel, error as logError, info } from '../../logger.js'; +import { isAppTimeoutError } from './retry.js'; /** User-facing error message: detailed in DEBUG, generic otherwise. */ export function getToolErrorMessage(error: unknown, fallbackMessage: string): string { @@ -81,9 +82,6 @@ export type ToolError = z.infer; const DEFAULT_TIMEOUT_SUGGESTION = 'Retry the request, or increase --request-timeout-ms.'; -/** Matches {@link withTimeout} rejection message prefix in `retry.ts`. */ -const TIMEOUT_MESSAGE_PATTERN = /^Timeout after \d+ms while waiting for /i; - export function flowGateToolError(namespace: string, message: string): ToolError { return { code: 'FLOW_GATE', @@ -136,18 +134,12 @@ export function lifecycleToolError(message: string): ToolError { }; } -function rawErrorMessage(error: unknown): string { - return error instanceof Error ? error.message : String(error); -} - /** * Map an unexpected thrown error to {@link ToolError} for MCP responses. - * Uses raw `Error#message` for timeout detection (DEBUG mode replaces the user message). */ export function classifyToolCatchError(error: unknown, fallbackMessage: string): ToolError { - const raw = rawErrorMessage(error); const message = getToolErrorMessage(error, fallbackMessage); - if (TIMEOUT_MESSAGE_PATTERN.test(raw)) { + if (isAppTimeoutError(error)) { return timeoutToolError(message); } return pineconeToolError(message);