Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
63 changes: 54 additions & 9 deletions src/core/pinecone-client.test.ts
Original file line number Diff line number Diff line change
@@ -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';
Expand All @@ -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;

Expand All @@ -32,6 +42,10 @@ describe('PineconeClient', () => {
});
});

afterEach(() => {
vi.useRealTimers();
});

describe('constructor', () => {
it('should initialize with provided config', () => {
expect(client).toBeDefined();
Expand Down Expand Up @@ -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({
Expand All @@ -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', () => {
Expand Down Expand Up @@ -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;
Expand Down
16 changes: 14 additions & 2 deletions src/core/pinecone-client.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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()
);
Comment thread
coderabbitai[bot] marked this conversation as resolved.
}

async query(params: QueryParams): Promise<HybridQueryResult> {
Expand Down Expand Up @@ -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.');
}

Expand Down Expand Up @@ -179,7 +190,8 @@ export class PineconeClient {
this.rerankModel,
query,
mergedResults,
topK
topK,
this.indexSession.getRequestTimeoutMs()
);
documents = rerankOut.results;
degraded = rerankOut.degraded;
Expand Down
198 changes: 197 additions & 1 deletion src/core/pinecone/indexes.test.ts
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
import { describe, it, expect, vi } from 'vitest';
import { describe, it, expect, vi, afterEach } from 'vitest';
Comment thread
jonathanMLDev marked this conversation as resolved.
import { PineconeIndexSession } from './indexes.js';
import type { SearchableIndex } from '../../types.js';

Expand Down Expand Up @@ -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 () => {
Expand Down Expand Up @@ -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', () => {
Expand Down Expand Up @@ -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({
Expand Down Expand Up @@ -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 = {
Expand Down Expand Up @@ -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({}),
Expand Down Expand Up @@ -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);
});
});
});
Loading