From fb2c2cab97422498223feeac8195c564a774482a Mon Sep 17 00:00:00 2001 From: zho Date: Tue, 19 May 2026 01:18:41 +0800 Subject: [PATCH 1/3] resolved lifecycle issues --- benchmarks/latency.ts | 6 +- src/config.ts | 10 +- src/pinecone-client.test.ts | 143 +++++++++++++----- src/pinecone-client.ts | 41 +++-- src/pinecone/rerank.test.ts | 26 ++-- src/pinecone/rerank.ts | 34 +++-- src/server.test.ts | 59 +++++++- src/server.ts | 39 ++++- src/server/client-context.ts | 5 + src/server/config-context.ts | 5 + src/server/namespace-utils.ts | 12 ++ src/server/suggestion-flow.ts | 23 ++- src/server/tools/count-tool.test.ts | 11 ++ src/server/tools/count-tool.ts | 17 ++- src/server/tools/generate-urls-tool.test.ts | 11 ++ src/server/tools/generate-urls-tool.ts | 15 +- src/server/tools/guided-query-tool.test.ts | 38 ++++- src/server/tools/guided-query-tool.ts | 44 ++++-- src/server/tools/query-documents-tool.test.ts | 4 +- src/server/tools/query-documents-tool.ts | 27 +++- src/server/tools/query-tool.test.ts | 46 +++++- src/server/tools/query-tool.ts | 27 +++- .../tools/suggest-query-params-tool.test.ts | 44 ++++++ src/server/tools/suggest-query-params-tool.ts | 15 +- src/server/tools/test-helpers.ts | 14 +- src/server/url-generation.ts | 9 ++ src/types.ts | 27 +++- 27 files changed, 626 insertions(+), 126 deletions(-) create mode 100644 src/server/namespace-utils.ts diff --git a/benchmarks/latency.ts b/benchmarks/latency.ts index b04cee2..56a57c8 100644 --- a/benchmarks/latency.ts +++ b/benchmarks/latency.ts @@ -222,7 +222,11 @@ function createBenchPineconeMock(): PineconeClient { return { async query() { - return mockQueryResults; + return { + results: mockQueryResults, + degraded: false, + hybrid_leg_failed: null, + }; }, async count() { return { count: 42, truncated: false }; diff --git a/src/config.ts b/src/config.ts index 11b2fc4..9872b2c 100644 --- a/src/config.ts +++ b/src/config.ts @@ -6,7 +6,12 @@ * `process.env` directly anymore — they receive their slice of the config. */ -import { DEFAULT_INDEX_NAME, DEFAULT_RERANK_MODEL, FLOW_CACHE_TTL_MS } from './constants.js'; +import { + DEFAULT_INDEX_NAME, + DEFAULT_RERANK_MODEL, + DEFAULT_TOP_K, + FLOW_CACHE_TTL_MS, +} from './constants.js'; /** Allowed log levels, in ascending severity. */ export type LogLevel = 'DEBUG' | 'INFO' | 'WARN' | 'ERROR'; @@ -49,9 +54,6 @@ export interface ServerConfig { /** Default per-call timeout for Pinecone requests, in milliseconds. */ export const DEFAULT_REQUEST_TIMEOUT_MS = 15_000; -/** Default top-k mirrors constants.DEFAULT_TOP_K but is duplicated here to avoid a cycle. */ -const DEFAULT_TOP_K = 10; - function asLogLevel(value: string | undefined, fallback: LogLevel): LogLevel { const allowed: LogLevel[] = ['DEBUG', 'INFO', 'WARN', 'ERROR']; return allowed.includes(value as LogLevel) ? (value as LogLevel) : fallback; diff --git a/src/pinecone-client.test.ts b/src/pinecone-client.test.ts index 98f1523..a101cb0 100644 --- a/src/pinecone-client.test.ts +++ b/src/pinecone-client.test.ts @@ -1,5 +1,6 @@ -import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest'; +import { describe, it, expect, beforeEach, vi } from 'vitest'; import { PineconeClient } from './pinecone-client.js'; +import { resolveConfig } from './config.js'; import type { SearchableIndex, PineconeHit } from './types.js'; import * as rerankModule from './pinecone/rerank.js'; @@ -31,26 +32,40 @@ describe('PineconeClient', () => { }); }); - afterEach(() => { - delete process.env['PINECONE_INDEX_NAME']; - delete process.env['PINECONE_RERANK_MODEL']; - delete process.env['PINECONE_TOP_K']; - }); - describe('constructor', () => { it('should initialize with provided config', () => { expect(client).toBeDefined(); }); - it('should use environment variables as fallbacks', () => { - process.env['PINECONE_INDEX_NAME'] = 'env-index'; - process.env['PINECONE_RERANK_MODEL'] = 'env-model'; - - const envClient = new PineconeClient({ - apiKey: 'test-api-key', - }); - - expect(envClient).toBeDefined(); + it('honors resolveConfig overrides without PINECONE_* env on the client path', () => { + const prevIndex = process.env['PINECONE_INDEX_NAME']; + const prevModel = process.env['PINECONE_RERANK_MODEL']; + const prevTopK = process.env['PINECONE_TOP_K']; + delete process.env['PINECONE_INDEX_NAME']; + delete process.env['PINECONE_RERANK_MODEL']; + delete process.env['PINECONE_TOP_K']; + try { + const resolved = resolveConfig({ + apiKey: 'test-api-key', + indexName: 'resolved-index', + rerankModel: 'resolved-model', + defaultTopK: 42, + }); + const c = new PineconeClient({ + apiKey: resolved.apiKey, + indexName: resolved.indexName, + rerankModel: resolved.rerankModel, + defaultTopK: resolved.defaultTopK, + }); + expect(c.getSparseIndexName()).toBe('resolved-index-sparse'); + } finally { + if (prevIndex !== undefined) process.env['PINECONE_INDEX_NAME'] = prevIndex; + else delete process.env['PINECONE_INDEX_NAME']; + if (prevModel !== undefined) process.env['PINECONE_RERANK_MODEL'] = prevModel; + else delete process.env['PINECONE_RERANK_MODEL']; + if (prevTopK !== undefined) process.env['PINECONE_TOP_K'] = prevTopK; + else delete process.env['PINECONE_TOP_K']; + } }); }); @@ -76,16 +91,16 @@ describe('PineconeClient', () => { it('should continue hybrid search when one index fails', async () => { const testClient = stubPineconeClient(client); + const denseRef = {} as SearchableIndex; + const sparseRef = {} as SearchableIndex; testClient.ensureIndexes = async () => ({ - denseIndex: {} as SearchableIndex, - sparseIndex: {} as SearchableIndex, + denseIndex: denseRef, + sparseIndex: sparseRef, }); - let searchCall = 0; - testClient.searchIndex = async () => { - searchCall += 1; - if (searchCall === 1) { + testClient.searchIndex = async (index) => { + if (index === denseRef) { throw new Error('dense failure'); } return [ @@ -97,16 +112,18 @@ describe('PineconeClient', () => { ]; }; - const results = await client.query({ + const out = await client.query({ query: 'hybrid search', namespace: 'test', topK: 5, useReranking: false, }); - expect(results).toHaveLength(1); - expect(results[0].content).toBe('hybrid content'); - expect(results[0].metadata.author).toBe('tester'); + expect(out.results).toHaveLength(1); + expect(out.results[0]?.content).toBe('hybrid content'); + expect(out.results[0]?.metadata.author).toBe('tester'); + expect(out.hybrid_leg_failed).toBe('dense'); + expect(out.degraded).toBe(false); }); it('should throw when both dense and sparse searches fail', async () => { @@ -249,15 +266,18 @@ describe('PineconeClient', () => { }); it('uses rerankResults from pinecone/rerank when useReranking is true', async () => { - const spy = vi.spyOn(rerankModule, 'rerankResults').mockResolvedValue([ - { - id: 'd1', - content: 'from dense', - score: 0.9, - metadata: {}, - reranked: true, - }, - ]); + const spy = vi.spyOn(rerankModule, 'rerankResults').mockResolvedValue({ + results: [ + { + id: 'd1', + content: 'from dense', + score: 0.9, + metadata: {}, + reranked: true, + }, + ], + degraded: false, + }); try { const testClient = stubPineconeClient(client); const denseRef = {} as SearchableIndex; @@ -280,9 +300,54 @@ describe('PineconeClient', () => { useReranking: true, }); - expect(results).toHaveLength(1); - expect(results[0].reranked).toBe(true); - expect(results[0].content).toBe('from dense'); + expect(results.results).toHaveLength(1); + expect(results.results[0]?.reranked).toBe(true); + expect(results.results[0]?.content).toBe('from dense'); + expect(spy).toHaveBeenCalled(); + } finally { + spy.mockRestore(); + } + }); + + it('propagates rerank degradation to hybrid query outcome', async () => { + const spy = vi.spyOn(rerankModule, 'rerankResults').mockResolvedValue({ + results: [ + { + id: 'd1', + content: 'from dense', + score: 0.9, + metadata: {}, + reranked: false, + }, + ], + degraded: true, + degradation_reason: 'rerank_failed: timeout', + }); + try { + const testClient = stubPineconeClient(client); + const denseRef = {} as SearchableIndex; + const sparseRef = {} as SearchableIndex; + testClient.ensureIndexes = async () => ({ + denseIndex: denseRef, + sparseIndex: sparseRef, + }); + testClient.searchIndex = async (index) => { + if (index === denseRef) { + return [{ _id: 'd1', _score: 0.9, fields: { chunk_text: 'from dense' } }]; + } + return []; + }; + + const out = await client.query({ + query: 'q', + namespace: 'n', + topK: 5, + useReranking: true, + }); + + expect(out.degraded).toBe(true); + expect(out.degradation_reason).toBe('rerank_failed: timeout'); + expect(out.results[0]?.reranked).toBe(false); expect(spy).toHaveBeenCalled(); } finally { spy.mockRestore(); @@ -314,7 +379,7 @@ describe('PineconeClient', () => { useReranking: false, }); - expect(results.length).toBe(2); + expect(results.results.length).toBe(2); }); }); diff --git a/src/pinecone-client.ts b/src/pinecone-client.ts index 37c5f0b..3606072 100644 --- a/src/pinecone-client.ts +++ b/src/pinecone-client.ts @@ -11,6 +11,8 @@ import type { KeywordSearchParams, KeywordIndexNamespacesResult, SearchableIndex, + HybridQueryResult, + HybridLegFailed, } from './types.js'; import { DEFAULT_INDEX_NAME, @@ -35,16 +37,16 @@ export class PineconeClient { private defaultTopK: number; private readonly indexSession: PineconeIndexSession; - /** Create a client with the given config; env vars override index name, rerank model, and top-k. */ + /** + * Create a client from a resolved {@link PineconeClientConfig}. + * Index name, rerank model, and default top-k come only from this object (typically + * built via {@link resolveConfig} / CLI); this class does not read `process.env`. + */ constructor(config: PineconeClientConfig) { - const indexName = config.indexName || process.env['PINECONE_INDEX_NAME'] || DEFAULT_INDEX_NAME; + const indexName = config.indexName ?? DEFAULT_INDEX_NAME; this.indexSession = new PineconeIndexSession(config.apiKey, indexName); - this.rerankModel = - config.rerankModel || process.env['PINECONE_RERANK_MODEL'] || DEFAULT_RERANK_MODEL; - const envTopK = process.env['PINECONE_TOP_K']; - const parsedEnvTopK = envTopK !== undefined ? parseInt(envTopK, 10) : NaN; - this.defaultTopK = - config.defaultTopK ?? (Number.isFinite(parsedEnvTopK) ? parsedEnvTopK : DEFAULT_TOP_K); + this.rerankModel = config.rerankModel ?? DEFAULT_RERANK_MODEL; + this.defaultTopK = config.defaultTopK ?? DEFAULT_TOP_K; } /** Sparse index name `{indexName}-sparse` (keyword / hybrid sparse). */ @@ -105,7 +107,7 @@ export class PineconeClient { return searchIndexImpl(index, query, topK, namespace, metadataFilter, options); } - async query(params: QueryParams): Promise { + async query(params: QueryParams): Promise { const { query, topK: requestedTopK, @@ -148,17 +150,29 @@ export class PineconeClient { throw new Error('Hybrid search failed: both dense and sparse index searches failed.'); } + let hybridLegFailed: HybridLegFailed = null; + if (denseResult.status === 'rejected' && sparseResult.status === 'fulfilled') { + hybridLegFailed = 'dense'; + } else if (sparseResult.status === 'rejected' && denseResult.status === 'fulfilled') { + hybridLegFailed = 'sparse'; + } + const mergedResults = mergeResults(denseHits, sparseHits); + let degraded = false; + let degradation_reason: string | undefined; let documents: SearchResult[]; if (useReranking) { - documents = await rerankResultsImpl( + const rerankOut = await rerankResultsImpl( this.indexSession.ensureClient(), this.rerankModel, query, mergedResults, topK ); + documents = rerankOut.results; + degraded = rerankOut.degraded; + degradation_reason = rerankOut.degradation_reason; } else { documents = sliceMergedHitsToSearchResults(mergedResults, topK); } @@ -167,7 +181,12 @@ export class PineconeClient { `Retrieved ${documents.length} documents from hybrid search (dense: ${denseHits.length}, sparse: ${sparseHits.length})` ); - return documents; + return { + results: documents, + degraded, + ...(degradation_reason !== undefined ? { degradation_reason } : {}), + hybrid_leg_failed: hybridLegFailed, + }; } async keywordSearch(params: KeywordSearchParams): Promise { diff --git a/src/pinecone/rerank.test.ts b/src/pinecone/rerank.test.ts index 17c66a5..518f5ad 100644 --- a/src/pinecone/rerank.test.ts +++ b/src/pinecone/rerank.test.ts @@ -7,10 +7,11 @@ const sampleMerged: MergedHit[] = [ ]; describe('rerankResults', () => { - it('returns empty array when there are no merged hits', async () => { + it('returns empty outcome when there are no merged hits', async () => { const pc = {} as Parameters[0]; const out = await rerankResults(pc, 'any-model', 'q', [], 5); - expect(out).toEqual([]); + expect(out.results).toEqual([]); + expect(out.degraded).toBe(false); }); it('maps successful inference.rerank response', async () => { @@ -26,21 +27,24 @@ describe('rerankResults', () => { const out = await rerankResults(pc, 'm', 'q', sampleMerged, 5); - expect(out).toHaveLength(1); - expect(out[0]?.reranked).toBe(true); - expect(out[0]?.id).toBe('1'); - expect(out[0]?.content).toBe('hello'); - expect(out[0]?.score).toBeCloseTo(0.99); + expect(out.results).toHaveLength(1); + expect(out.degraded).toBe(false); + expect(out.results[0]?.reranked).toBe(true); + expect(out.results[0]?.id).toBe('1'); + expect(out.results[0]?.content).toBe('hello'); + expect(out.results[0]?.score).toBeCloseTo(0.99); }); - it('returns unreranked slice when rerank throws', async () => { + it('returns unreranked slice with degraded when rerank throws', async () => { const rerank = vi.fn().mockRejectedValue(new Error('rerank unavailable')); const pc = { inference: { rerank } } as Parameters[0]; const out = await rerankResults(pc, 'm', 'q', sampleMerged, 5); - expect(out).toHaveLength(1); - expect(out[0]?.reranked).toBe(false); - expect(out[0]?.content).toBe('hello'); + expect(out.results).toHaveLength(1); + expect(out.degraded).toBe(true); + expect(out.degradation_reason).toMatch(/^rerank_failed:/); + expect(out.results[0]?.reranked).toBe(false); + expect(out.results[0]?.content).toBe('hello'); }); }); diff --git a/src/pinecone/rerank.ts b/src/pinecone/rerank.ts index 90b6c85..a712691 100644 --- a/src/pinecone/rerank.ts +++ b/src/pinecone/rerank.ts @@ -6,6 +6,14 @@ import type { Pinecone } from '@pinecone-database/pinecone'; import { error as logError } from '../logger.js'; import type { MergedHit, SearchResult } from '../types.js'; +export type RerankOutcome = { + results: SearchResult[]; + /** True when rerank was attempted and failed; results are unreranked slice. */ + degraded: boolean; + /** Human-readable reason for MCP clients when {@link degraded} is true. */ + degradation_reason?: string; +}; + /** * Rerank merged hits using Pinecone's reranking model; on failure returns unreranked slice. */ @@ -15,9 +23,9 @@ export async function rerankResults( query: string, results: MergedHit[], topN: number -): Promise { +): Promise { if (!results || results.length === 0) { - return []; + return { results: [], degraded: false }; } try { @@ -48,16 +56,20 @@ export async function rerankResults( reranked: true, }); } - return reranked; + return { results: reranked, degraded: false }; } catch (error) { logError('Error reranking results', error); - // Fall back to returning unreranked results - return results.slice(0, topN).map((result) => ({ - id: result._id || '', - content: result.chunk_text || '', - score: result._score || 0, - metadata: result.metadata || {}, - reranked: false, - })); + const msg = error instanceof Error ? error.message : String(error); + return { + results: results.slice(0, topN).map((result) => ({ + id: result._id || '', + content: result.chunk_text || '', + score: result._score || 0, + metadata: result.metadata || {}, + reranked: false, + })), + degraded: true, + degradation_reason: `rerank_failed: ${msg}`, + }; } } diff --git a/src/server.test.ts b/src/server.test.ts index 602e98b..3da3982 100644 --- a/src/server.test.ts +++ b/src/server.test.ts @@ -1,5 +1,14 @@ -import { describe, expect, it } from 'vitest'; -import { validateMetadataFilter, suggestQueryParams } from './server.js'; +import { describe, expect, it, afterEach } from 'vitest'; +import { + validateMetadataFilter, + suggestQueryParams, + setupServer, + teardownServer, + setPineconeClient, + resolveConfig, + PineconeClient, + hasUrlGenerator, +} from './server.js'; describe('suggestQueryParams', () => { const wg21Fields = { @@ -96,3 +105,49 @@ describe('validateMetadataFilter', () => { expect(result).toContain('must use an array of primitive values'); }); }); + +describe('setupServer lifecycle', () => { + afterEach(() => { + teardownServer(); + }); + + it('throws on second setupServer without teardown', async () => { + const cfg = resolveConfig({ apiKey: 'lifecycle-test-key' }); + setPineconeClient( + new PineconeClient({ + apiKey: cfg.apiKey, + indexName: cfg.indexName, + rerankModel: cfg.rerankModel, + defaultTopK: cfg.defaultTopK, + }) + ); + await setupServer(cfg); + await expect(setupServer(cfg)).rejects.toThrow(/teardownServer/); + }); + + it('allows setup after teardown and reinstalls built-in URL generators', async () => { + const cfg = resolveConfig({ apiKey: 'lifecycle-test-key-2' }); + setPineconeClient( + new PineconeClient({ + apiKey: cfg.apiKey, + indexName: cfg.indexName, + rerankModel: cfg.rerankModel, + defaultTopK: cfg.defaultTopK, + }) + ); + await setupServer(cfg); + expect(hasUrlGenerator('mailing')).toBe(true); + teardownServer(); + expect(hasUrlGenerator('mailing')).toBe(false); + setPineconeClient( + new PineconeClient({ + apiKey: cfg.apiKey, + indexName: cfg.indexName, + rerankModel: cfg.rerankModel, + defaultTopK: cfg.defaultTopK, + }) + ); + await setupServer(cfg); + expect(hasUrlGenerator('mailing')).toBe(true); + }); +}); diff --git a/src/server.ts b/src/server.ts index 6ed356c..05fcad2 100644 --- a/src/server.ts +++ b/src/server.ts @@ -5,7 +5,8 @@ * * Import from the package root: * - * - {@link setupServer} — build an `McpServer` with all tools registered. + * - {@link setupServer} — build an `McpServer` with all tools registered (at most once per process unless {@link teardownServer} runs). + * - {@link teardownServer} — clear process-global server state so {@link setupServer} can run again (tests, re-embedding). * - {@link PineconeClient} — hybrid search, count, namespace listing, etc. * - {@link resolveConfig} — merge CLI-style overrides with `process.env`. * - {@link setPineconeClient} — inject a client instance before `setupServer()`. @@ -21,8 +22,14 @@ import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; import { SERVER_INSTRUCTIONS, SERVER_NAME, SERVER_VERSION } from './constants.js'; import type { ServerConfig } from './config.js'; -import { setServerConfig } from './server/config-context.js'; -import { registerBuiltinUrlGenerators } from './server/url-generation.js'; +import { setServerConfig, resetServerConfig } from './server/config-context.js'; +import { clearPineconeClient } from './server/client-context.js'; +import { + registerBuiltinUrlGenerators, + resetUrlGenerationRegistry, +} from './server/url-generation.js'; +import { invalidateNamespacesCache } from './server/namespaces-cache.js'; +import { resetSuggestionFlow } from './server/suggestion-flow.js'; import { registerCountTool } from './server/tools/count-tool.js'; import { registerGuidedQueryTool } from './server/tools/guided-query-tool.js'; import { registerGenerateUrlsTool } from './server/tools/generate-urls-tool.js'; @@ -75,8 +82,25 @@ export type { QueryResponse, QueryResultRowShape, KeywordIndexNamespacesResult, + HybridQueryResult, + HybridLegFailed, } from './types.js'; +let mcpServerInitialized = false; + +/** + * Reset process-global MCP server state (suggest-flow, namespace cache, active config, + * Pinecone client handle, URL generator registry). Call before a second {@link setupServer}. + */ +export function teardownServer(): void { + resetSuggestionFlow(); + invalidateNamespacesCache(); + resetServerConfig(); + clearPineconeClient(); + resetUrlGenerationRegistry(); + mcpServerInitialized = false; +} + /** * Create and configure the MCP server with all tools. * @@ -85,9 +109,17 @@ export type { * and {@link setServerConfig} — see README “Deployment model”. Multi-tenant HTTP * multiplexing can violate the suggest-flow guarantee unless you isolate by session. * + * A second call in the same process throws unless {@link teardownServer} runs first. + * * @returns the configured `McpServer` instance, ready to connect to a transport. */ export async function setupServer(config?: ServerConfig): Promise { + if (mcpServerInitialized) { + throw new Error( + 'setupServer() already called in this process. The MCP server uses process-global state (suggest-flow, namespace cache, URL generators, config). Call teardownServer() first if you need to re-initialize.' + ); + } + if (config) { setServerConfig(config); } @@ -114,5 +146,6 @@ export async function setupServer(config?: ServerConfig): Promise { registerGuidedQueryTool(server); registerGenerateUrlsTool(server); + mcpServerInitialized = true; return server; } diff --git a/src/server/client-context.ts b/src/server/client-context.ts index f2f759f..e894fc0 100644 --- a/src/server/client-context.ts +++ b/src/server/client-context.ts @@ -15,3 +15,8 @@ export function getPineconeClient(): PineconeClient { export function setPineconeClient(client: PineconeClient): void { pineconeClient = client; } + +/** Clear the shared client (used by {@link teardownServer} and tests). */ +export function clearPineconeClient(): void { + pineconeClient = null; +} diff --git a/src/server/config-context.ts b/src/server/config-context.ts index 763ac94..42ed419 100644 --- a/src/server/config-context.ts +++ b/src/server/config-context.ts @@ -8,6 +8,11 @@ export function setServerConfig(config: ServerConfig): void { activeConfig = config; } +/** Clear active config so the next `getServerConfig()` resolves again (used by {@link teardownServer}). */ +export function resetServerConfig(): void { + activeConfig = null; +} + /** * Active server config for modules that cannot receive `ServerConfig` through parameters * (namespace cache TTL, suggest-flow gate, etc.). diff --git a/src/server/namespace-utils.ts b/src/server/namespace-utils.ts new file mode 100644 index 0000000..3b4b440 --- /dev/null +++ b/src/server/namespace-utils.ts @@ -0,0 +1,12 @@ +/** + * Shared namespace string handling for the suggest-flow gate and Pinecone tool calls. + */ + +/** + * Trim surrounding whitespace. Returns `null` if the result is empty + * (callers should map this to a VALIDATION tool error). + */ +export function normalizeNamespace(input: string): string | null { + const trimmed = input.trim(); + return trimmed.length === 0 ? null : trimmed; +} diff --git a/src/server/suggestion-flow.ts b/src/server/suggestion-flow.ts index 2f77d13..d4aa312 100644 --- a/src/server/suggestion-flow.ts +++ b/src/server/suggestion-flow.ts @@ -1,5 +1,6 @@ import { getServerConfig } from './config-context.js'; import type { RecommendedTool } from './query-suggestion.js'; +import { normalizeNamespace } from './namespace-utils.js'; type FlowState = { updatedAt: number; @@ -26,8 +27,12 @@ function sweepExpired(): void { /** Record that suggest_query_params was called for this namespace (enables query/count for the flow). */ export function markSuggested(namespace: string, state: Omit): void { + const key = normalizeNamespace(namespace); + if (!key) { + throw new Error('markSuggested: namespace must not be empty after trim'); + } sweepExpired(); - stateByNamespace.set(namespace, { + stateByNamespace.set(key, { ...state, updatedAt: Date.now(), }); @@ -50,6 +55,14 @@ export function requireSuggested(namespace: string): ok: false; message: string; } { + const key = normalizeNamespace(namespace); + if (!key) { + return { + ok: false, + message: 'namespace cannot be empty after trimming whitespace.', + }; + } + const cfg = getServerConfig(); if (cfg.disableSuggestFlow) { return { @@ -63,7 +76,7 @@ export function requireSuggested(namespace: string): }; } - const state = stateByNamespace.get(namespace); + const state = stateByNamespace.get(key); if (!state) { return { ok: false, @@ -74,7 +87,7 @@ export function requireSuggested(namespace: string): const now = Date.now(); if (now - state.updatedAt > cfg.cacheTtlMs) { - stateByNamespace.delete(namespace); + stateByNamespace.delete(key); return { ok: false, message: @@ -85,7 +98,7 @@ export function requireSuggested(namespace: string): return { ok: true, flow: state }; } -/** Test-only: clear the in-memory flow state so each test starts clean. */ -export function resetSuggestionFlowForTests(): void { +/** Clear suggest-flow gate state (used by {@link teardownServer} and tests). */ +export function resetSuggestionFlow(): void { stateByNamespace.clear(); } diff --git a/src/server/tools/count-tool.test.ts b/src/server/tools/count-tool.test.ts index c899d2b..aa32003 100644 --- a/src/server/tools/count-tool.test.ts +++ b/src/server/tools/count-tool.test.ts @@ -33,6 +33,17 @@ describe('count tool handler', () => { vi.restoreAllMocks(); }); + it('returns VALIDATION when namespace is whitespace-only', async () => { + const server = createMockServer(); + registerCountTool(server as never); + const raw = await server.getHandler('count')!({ + namespace: ' ', + query_text: 'doc', + }); + const err = assertToolErrorCode(raw, 'VALIDATION'); + expect(err.field).toBe('namespace'); + }); + it('returns VALIDATION when query_text is empty', async () => { const server = createMockServer(); registerCountTool(server as never); diff --git a/src/server/tools/count-tool.ts b/src/server/tools/count-tool.ts index f9214ff..b582e8a 100644 --- a/src/server/tools/count-tool.ts +++ b/src/server/tools/count-tool.ts @@ -2,6 +2,7 @@ import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; import { z } from 'zod'; import { getPineconeClient } from '../client-context.js'; import { metadataFilterSchema, validateMetadataFilterDetailed } from '../metadata-filter.js'; +import { normalizeNamespace } from '../namespace-utils.js'; import { requireSuggested } from '../suggestion-flow.js'; import { classifyToolCatchError, @@ -53,6 +54,14 @@ export function registerCountTool(server: McpServer): void { async (params) => { try { const { namespace, query_text, metadata_filter } = params; + const nsNorm = normalizeNamespace(namespace); + if (!nsNorm) { + return jsonErrorResponse( + validationToolError('namespace cannot be empty', 'namespace', { + suggestion: 'Use a namespace name from list_namespaces (trimmed).', + }) + ); + } if (!query_text.trim()) { return jsonErrorResponse(validationToolError('query_text cannot be empty', 'query_text')); } @@ -62,21 +71,21 @@ export function registerCountTool(server: McpServer): void { return jsonErrorResponse(validationToolError(err.message, err.field)); } } - const flowCheck = requireSuggested(namespace); + const flowCheck = requireSuggested(nsNorm); if (!flowCheck.ok) { - return jsonErrorResponse(flowGateToolError(namespace, flowCheck.message)); + return jsonErrorResponse(flowGateToolError(nsNorm, flowCheck.message)); } const client = getPineconeClient(); const { count, truncated } = await client.count({ query: query_text.trim(), - namespace, + namespace: nsNorm, metadataFilter: metadata_filter, }); const response: CountResponse = { status: COUNT_RESPONSE_STATUS, count, truncated, - namespace, + namespace: nsNorm, metadata_filter, }; return jsonResponse(response); diff --git a/src/server/tools/generate-urls-tool.test.ts b/src/server/tools/generate-urls-tool.test.ts index 96a57a8..d08e640 100644 --- a/src/server/tools/generate-urls-tool.test.ts +++ b/src/server/tools/generate-urls-tool.test.ts @@ -8,6 +8,17 @@ describe('generate_urls tool handler', () => { vi.restoreAllMocks(); }); + it('returns VALIDATION when namespace is whitespace-only', async () => { + const server = createMockServer(); + registerGenerateUrlsTool(server as never); + const raw = await server.getHandler('generate_urls')!({ + namespace: ' ', + records: [{ document_number: 'P1234' }], + }); + const err = assertToolErrorCode(raw, 'VALIDATION'); + expect(err.field).toBe('namespace'); + }); + it('returns PINECONE_ERROR when generateUrlForNamespace throws', async () => { vi.spyOn(urlGeneration, 'generateUrlForNamespace').mockImplementation(() => { throw new Error('generator boom'); diff --git a/src/server/tools/generate-urls-tool.ts b/src/server/tools/generate-urls-tool.ts index c8b6941..34006da 100644 --- a/src/server/tools/generate-urls-tool.ts +++ b/src/server/tools/generate-urls-tool.ts @@ -1,7 +1,8 @@ import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; import { z } from 'zod'; import { generateUrlForNamespace } from '../url-generation.js'; -import { classifyToolCatchError, logToolError } from '../tool-error.js'; +import { normalizeNamespace } from '../namespace-utils.js'; +import { classifyToolCatchError, logToolError, validationToolError } from '../tool-error.js'; import { jsonErrorResponse, jsonResponse } from '../tool-response.js'; /** Get metadata from a record (either record.metadata or the record itself). */ @@ -38,9 +39,17 @@ export function registerGenerateUrlsTool(server: McpServer): void { async (params) => { try { const { namespace, records } = params; + const nsNorm = normalizeNamespace(namespace); + if (!nsNorm) { + return jsonErrorResponse( + validationToolError('namespace cannot be empty', 'namespace', { + suggestion: 'Use a namespace name from list_namespaces (trimmed).', + }) + ); + } const results = records.map((record, index) => { const metadata = extractMetadata(record); - const generated = generateUrlForNamespace(namespace, metadata); + const generated = generateUrlForNamespace(nsNorm, metadata); return { index, url: generated.url, @@ -52,7 +61,7 @@ export function registerGenerateUrlsTool(server: McpServer): void { return jsonResponse({ status: 'success', - namespace, + namespace: nsNorm, count: results.length, results, }); diff --git a/src/server/tools/guided-query-tool.test.ts b/src/server/tools/guided-query-tool.test.ts index a748f74..b1e8073 100644 --- a/src/server/tools/guided-query-tool.test.ts +++ b/src/server/tools/guided-query-tool.test.ts @@ -5,6 +5,7 @@ import { registerGuidedQueryTool } from './guided-query-tool.js'; import { assertToolErrorCode, createMockServer, + makeHybridQueryResult, makeNamespaceCacheEntry, makeSearchResult, parseToolJson, @@ -43,11 +44,43 @@ describe('guided_query tool handler', () => { expires_at: Date.now() + 1_800_000, }); mockedGetClient.mockReturnValue({ - query: vi.fn().mockResolvedValue([makeSearchResult()]), + query: vi.fn().mockResolvedValue(makeHybridQueryResult()), count: vi.fn().mockResolvedValue({ count: 7, truncated: false }), } as never); }); + it('guided_query: surfaces rerank failure in decision_trace', async () => { + mockedGetClient.mockReturnValue({ + query: vi.fn().mockResolvedValue( + makeHybridQueryResult({ + degraded: true, + degradation_reason: 'rerank_failed: timeout', + results: [makeSearchResult({ reranked: false })], + }) + ), + count: vi.fn().mockResolvedValue({ count: 7, truncated: false }), + } as never); + + const server = createMockServer(); + registerGuidedQueryTool(server as never); + + const body = parseToolJson( + await server.getHandler('guided_query')!({ + user_query: 'What does the paper say about contracts?', + namespace: 'papers', + top_k: 8, + preferred_tool: 'auto', + enrich_urls: false, + }) + ); + + const trace = body.decision_trace as Record; + expect(trace.rerank_status).toBe('failed'); + const result = body.result as Record; + expect(result.degraded).toBe(true); + expect(result.degradation_reason).toBe('rerank_failed: timeout'); + }); + it('runs query_detailed path on auto when user asks for content', async () => { const server = createMockServer(); registerGuidedQueryTool(server as never); @@ -67,6 +100,7 @@ describe('guided_query tool handler', () => { const trace = body.decision_trace as Record; expect(trace.selected_namespace).toBe('papers'); expect(trace.selected_tool).toBe('detailed'); + expect(trace.rerank_status).toBe('success'); expect(query).toHaveBeenCalledWith( expect.objectContaining({ namespace: 'papers', @@ -100,6 +134,8 @@ describe('guided_query tool handler', () => { const result = body.result as Record; expect(result.tool).toBe('count'); expect(result.count).toBe(7); + const trace = body.decision_trace as Record; + expect(trace.rerank_status).toBe('skipped'); }); it('returns error when user_query is empty', async () => { diff --git a/src/server/tools/guided-query-tool.ts b/src/server/tools/guided-query-tool.ts index f2abf72..a42d9c2 100644 --- a/src/server/tools/guided-query-tool.ts +++ b/src/server/tools/guided-query-tool.ts @@ -7,6 +7,7 @@ import { formatQueryResultRows } from '../format-query-result.js'; import { metadataFilterSchema, validateMetadataFilterDetailed } from '../metadata-filter.js'; import { rankNamespacesByQuery } from '../namespace-router.js'; import { getNamespacesWithCache } from '../namespaces-cache.js'; +import { normalizeNamespace } from '../namespace-utils.js'; import { suggestQueryParams } from '../query-suggestion.js'; import { markSuggested } from '../suggestion-flow.js'; import { @@ -100,11 +101,16 @@ export function registerGuidedQueryTool(server: McpServer): void { const { data: namespaces, cache_hit } = await getNamespacesWithCache(); const ranked = rankNamespacesByQuery(queryText, namespaces, 3); - const explicitNamespace = inputNamespace?.trim(); - if (inputNamespace !== undefined && !explicitNamespace) { - return jsonErrorResponse(validationToolError('namespace cannot be empty', 'namespace')); + let namespace: string | null = null; + if (inputNamespace !== undefined) { + namespace = normalizeNamespace(inputNamespace); + if (!namespace) { + return jsonErrorResponse(validationToolError('namespace cannot be empty', 'namespace')); + } + } else { + const top = ranked[0]?.namespace; + namespace = top ? normalizeNamespace(top) : null; } - const namespace = explicitNamespace ?? ranked[0]?.namespace; /* * ToolError mapping: empty index / no routable namespace is backend/data state * (PINECONE_ERROR, recoverable). Explicit namespace missing from cache is user/input @@ -122,7 +128,9 @@ export function registerGuidedQueryTool(server: McpServer): void { ); } - const ns = namespaces.find((n) => n.namespace === namespace); + const ns = namespaces.find( + (n) => n.namespace === namespace || normalizeNamespace(n.namespace) === namespace + ); const suggestion = suggestQueryParams(ns?.metadata ?? null, queryText); if (!suggestion.namespace_found) { return jsonErrorResponse( @@ -145,7 +153,7 @@ export function registerGuidedQueryTool(server: McpServer): void { }); const client = getPineconeClient(); - const decision_trace = { + const baseTrace = { cache_hit, input_namespace: inputNamespace ?? null, routed_namespace: ranked[0]?.namespace ?? null, @@ -166,7 +174,10 @@ export function registerGuidedQueryTool(server: McpServer): void { }); return jsonResponse({ status: 'success', - decision_trace, + decision_trace: { + ...baseTrace, + rerank_status: 'skipped' as const, + }, result: { tool: 'count', namespace, @@ -191,7 +202,7 @@ export function registerGuidedQueryTool(server: McpServer): void { : isFast ? [...FAST_QUERY_FIELDS] : undefined; - const queryResults = await client.query({ + const queryOutcome = await client.query({ query: queryText, namespace, topK: top_k, @@ -199,7 +210,12 @@ export function registerGuidedQueryTool(server: McpServer): void { useReranking: !isFast, fields: fields?.length ? fields : undefined, }); - const formattedResults = formatQueryResultRows(queryResults, { + const rerank_status: 'success' | 'skipped' | 'failed' = isFast + ? 'skipped' + : queryOutcome.degraded + ? 'failed' + : 'success'; + const formattedResults = formatQueryResultRows(queryOutcome.results, { namespace, enrichUrls: enrich_urls, }); @@ -212,10 +228,18 @@ export function registerGuidedQueryTool(server: McpServer): void { result_count: formattedResults.length, ...(fields?.length ? { fields } : {}), results: formattedResults, + degraded: queryOutcome.degraded, + ...(queryOutcome.degradation_reason !== undefined + ? { degradation_reason: queryOutcome.degradation_reason } + : {}), + hybrid_leg_failed: queryOutcome.hybrid_leg_failed, }; return jsonResponse({ status: 'success', - decision_trace, + decision_trace: { + ...baseTrace, + rerank_status, + }, result, }); } catch (error) { diff --git a/src/server/tools/query-documents-tool.test.ts b/src/server/tools/query-documents-tool.test.ts index b41768e..1d0a713 100644 --- a/src/server/tools/query-documents-tool.test.ts +++ b/src/server/tools/query-documents-tool.test.ts @@ -7,7 +7,7 @@ import { registerQueryDocumentsTool } from './query-documents-tool.js'; import { assertToolErrorCode, createMockServer, - makeSearchResult, + makeHybridQueryResult, parseToolJson, } from './test-helpers.js'; @@ -46,7 +46,7 @@ describe('query_documents tool handler', () => { }, ]); mockedGetClient.mockReturnValue({ - query: vi.fn().mockResolvedValue([makeSearchResult()]), + query: vi.fn().mockResolvedValue(makeHybridQueryResult()), } as never); }); diff --git a/src/server/tools/query-documents-tool.ts b/src/server/tools/query-documents-tool.ts index f59fea8..adc1c69 100644 --- a/src/server/tools/query-documents-tool.ts +++ b/src/server/tools/query-documents-tool.ts @@ -7,6 +7,7 @@ import { } from '../../constants.js'; import { getPineconeClient } from '../client-context.js'; import { metadataFilterSchema, validateMetadataFilterDetailed } from '../metadata-filter.js'; +import { normalizeNamespace } from '../namespace-utils.js'; import { reassembleByDocument } from '../reassemble-documents.js'; import { requireSuggested } from '../suggestion-flow.js'; import { @@ -91,23 +92,32 @@ export function registerQueryDocumentsTool(server: McpServer): void { } } - const flowCheck = requireSuggested(namespace); + const nsNorm = normalizeNamespace(namespace); + if (!nsNorm) { + return jsonErrorResponse( + validationToolError('namespace cannot be empty', 'namespace', { + suggestion: 'Use a namespace name from list_namespaces (trimmed).', + }) + ); + } + + const flowCheck = requireSuggested(nsNorm); if (!flowCheck.ok) { - return jsonErrorResponse(flowGateToolError(namespace, flowCheck.message)); + return jsonErrorResponse(flowGateToolError(nsNorm, flowCheck.message)); } const chunkLimit = Math.min(QUERY_DOCUMENTS_MAX_CHUNKS, top_k * CHUNKS_PER_DOCUMENT); const client = getPineconeClient(); - const results = await client.query({ + const queryOutcome = await client.query({ query: query_text.trim(), topK: chunkLimit, - namespace, + namespace: nsNorm, useReranking: true, metadataFilter: metadata_filter, fields: undefined, }); - const reassembled = reassembleByDocument(results, { + const reassembled = reassembleByDocument(queryOutcome.results, { maxChunksPerDocument: max_chunks_per_document ?? 200, }); @@ -118,9 +128,14 @@ export function registerQueryDocumentsTool(server: McpServer): void { return jsonResponse({ status: 'success', query: query_text.trim(), - namespace, + namespace: nsNorm, metadata_filter, result_count: topDocuments.length, + degraded: queryOutcome.degraded, + ...(queryOutcome.degradation_reason !== undefined + ? { degradation_reason: queryOutcome.degradation_reason } + : {}), + hybrid_leg_failed: queryOutcome.hybrid_leg_failed, documents: topDocuments.map((doc) => ({ document_id: doc.document_id, merged_content: doc.merged_content, diff --git a/src/server/tools/query-tool.test.ts b/src/server/tools/query-tool.test.ts index 1b004c0..365bc39 100644 --- a/src/server/tools/query-tool.test.ts +++ b/src/server/tools/query-tool.test.ts @@ -6,6 +6,7 @@ import { registerQueryTool } from './query-tool.js'; import { assertToolErrorCode, createMockServer, + makeHybridQueryResult, makeSearchResult, parseToolJson, } from './test-helpers.js'; @@ -31,7 +32,7 @@ describe('query tool handler (preset-driven)', () => { vi.clearAllMocks(); vi.spyOn(suggestionFlow, 'requireSuggested').mockReturnValue(flowOk); mockedGetClient.mockReturnValue({ - query: vi.fn().mockResolvedValue([makeSearchResult()]), + query: vi.fn().mockResolvedValue(makeHybridQueryResult()), count: vi.fn(), } as never); }); @@ -40,6 +41,43 @@ describe('query tool handler (preset-driven)', () => { vi.restoreAllMocks(); }); + it('query: trims namespace for flow gate and Pinecone', async () => { + const requireSpy = vi.spyOn(suggestionFlow, 'requireSuggested').mockReturnValue(flowOk); + const server = createMockServer(); + registerQueryTool(server as never); + const query = mockedGetClient().query as ReturnType; + + await server.getHandler('query')!({ + query_text: 'contracts', + namespace: 'wg21 ', + top_k: 5, + preset: 'full', + use_reranking: true, + }); + + expect(requireSpy).toHaveBeenCalledWith('wg21'); + expect(query).toHaveBeenCalledWith( + expect.objectContaining({ + namespace: 'wg21', + }) + ); + }); + + it('query: returns VALIDATION when namespace is whitespace-only', async () => { + const server = createMockServer(); + registerQueryTool(server as never); + const query = mockedGetClient().query as ReturnType; + const raw = await server.getHandler('query')!({ + query_text: 'hello', + namespace: ' ', + top_k: 5, + preset: 'full', + }); + expect((raw as { isError?: boolean }).isError).toBe(true); + assertToolErrorCode(raw, 'VALIDATION'); + expect(query).not.toHaveBeenCalled(); + }); + it('query (preset=full): happy path calls client.query and returns formatted rows', async () => { const server = createMockServer(); registerQueryTool(server as never); @@ -183,7 +221,11 @@ describe('query tool handler (preset-driven)', () => { mockedGetClient.mockReturnValue({ query: vi .fn() - .mockResolvedValue([makeSearchResult({ reranked: false, score: 0.5, content: 'x' })]), + .mockResolvedValue( + makeHybridQueryResult({ + results: [makeSearchResult({ reranked: false, score: 0.5, content: 'x' })], + }) + ), count: vi.fn(), } as never); diff --git a/src/server/tools/query-tool.ts b/src/server/tools/query-tool.ts index d23882d..ddbbc77 100644 --- a/src/server/tools/query-tool.ts +++ b/src/server/tools/query-tool.ts @@ -5,6 +5,7 @@ import type { QueryResponse } from '../../types.js'; import { getPineconeClient } from '../client-context.js'; import { formatQueryResultRows } from '../format-query-result.js'; import { metadataFilterSchema, validateMetadataFilterDetailed } from '../metadata-filter.js'; +import { normalizeNamespace } from '../namespace-utils.js'; import { requireSuggested } from '../suggestion-flow.js'; import { classifyToolCatchError, @@ -43,32 +44,46 @@ async function executeQuery(params: QueryExecParams) { } } - const flowCheck = requireSuggested(namespace); + const nsNorm = normalizeNamespace(namespace); + if (!nsNorm) { + return jsonErrorResponse( + validationToolError('namespace cannot be empty', 'namespace', { + suggestion: 'Use a namespace name from list_namespaces (trimmed).', + }) + ); + } + + const flowCheck = requireSuggested(nsNorm); if (!flowCheck.ok) { - return jsonErrorResponse(flowGateToolError(namespace, flowCheck.message)); + return jsonErrorResponse(flowGateToolError(nsNorm, flowCheck.message)); } const client = getPineconeClient(); - const results = await client.query({ + const queryOutcome = await client.query({ query: query_text.trim(), topK: top_k, - namespace, + namespace: nsNorm, useReranking: use_reranking, metadataFilter: metadata_filter, fields: fields?.length ? fields : undefined, }); - const formattedResults = formatQueryResultRows(results); + const formattedResults = formatQueryResultRows(queryOutcome.results); const response: QueryResponse = { status: 'success', mode, query: query_text, - namespace, + namespace: nsNorm, metadata_filter: metadata_filter, result_count: formattedResults.length, results: formattedResults, ...(fields?.length ? { fields } : {}), + degraded: queryOutcome.degraded, + ...(queryOutcome.degradation_reason !== undefined + ? { degradation_reason: queryOutcome.degradation_reason } + : {}), + hybrid_leg_failed: queryOutcome.hybrid_leg_failed, }; return jsonResponse(response); } catch (error) { diff --git a/src/server/tools/suggest-query-params-tool.test.ts b/src/server/tools/suggest-query-params-tool.test.ts index 472a573..5882a55 100644 --- a/src/server/tools/suggest-query-params-tool.test.ts +++ b/src/server/tools/suggest-query-params-tool.test.ts @@ -25,6 +25,50 @@ describe('suggest_query_params tool handler', () => { vi.clearAllMocks(); }); + it('marks suggestion flow with trimmed namespace when cache key has no padding', async () => { + mockedGetNamespaces.mockResolvedValue({ + data: [makeNamespaceCacheEntry('wg21')], + cache_hit: false, + expires_at: Date.now() + 1_800_000, + }); + + const server = createMockServer(); + registerSuggestQueryParamsTool(server as never); + const body = parseToolJson( + await server.getHandler('suggest_query_params')!({ + namespace: 'wg21 ', + user_query: 'List papers with titles', + }) + ); + + expect(body.status).toBe('success'); + expect(body.namespace_found).toBe(true); + expect(mockedMarkSuggested).toHaveBeenCalledWith( + 'wg21', + expect.objectContaining({ + user_query: 'List papers with titles', + }) + ); + }); + + it('returns VALIDATION when namespace is whitespace-only', async () => { + mockedGetNamespaces.mockResolvedValue({ + data: [makeNamespaceCacheEntry('wg21')], + cache_hit: false, + expires_at: Date.now() + 1_800_000, + }); + + const server = createMockServer(); + registerSuggestQueryParamsTool(server as never); + const raw = await server.getHandler('suggest_query_params')!({ + namespace: ' ', + user_query: 'hello', + }); + const err = assertToolErrorCode(raw, 'VALIDATION'); + expect(err.field).toBe('namespace'); + expect(mockedMarkSuggested).not.toHaveBeenCalled(); + }); + it('marks suggestion flow and returns success when namespace exists in cache', async () => { mockedGetNamespaces.mockResolvedValue({ data: [makeNamespaceCacheEntry('wg21')], diff --git a/src/server/tools/suggest-query-params-tool.ts b/src/server/tools/suggest-query-params-tool.ts index 10f929b..461a40e 100644 --- a/src/server/tools/suggest-query-params-tool.ts +++ b/src/server/tools/suggest-query-params-tool.ts @@ -1,5 +1,6 @@ import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; import { z } from 'zod'; +import { normalizeNamespace } from '../namespace-utils.js'; import { getNamespacesWithCache } from '../namespaces-cache.js'; import { suggestQueryParams } from '../query-suggestion.js'; import { markSuggested } from '../suggestion-flow.js'; @@ -35,12 +36,22 @@ export function registerSuggestQueryParamsTool(server: McpServer): void { if (!user_query?.trim()) { return jsonErrorResponse(validationToolError('user_query cannot be empty', 'user_query')); } + const nsNorm = normalizeNamespace(namespace); + if (!nsNorm) { + return jsonErrorResponse( + validationToolError('namespace cannot be empty', 'namespace', { + suggestion: 'Use a namespace name from list_namespaces (trimmed).', + }) + ); + } const { data: namespacesInfo, cache_hit } = await getNamespacesWithCache(); - const ns = namespacesInfo.find((n) => n.namespace === namespace); + const ns = namespacesInfo.find( + (n) => n.namespace === nsNorm || normalizeNamespace(n.namespace) === nsNorm + ); const metadataFields = ns?.metadata ?? null; const result = suggestQueryParams(metadataFields, user_query.trim()); if (result.namespace_found) { - markSuggested(namespace, { + markSuggested(nsNorm, { recommended_tool: result.recommended_tool, suggested_fields: result.suggested_fields, user_query: user_query.trim(), diff --git a/src/server/tools/test-helpers.ts b/src/server/tools/test-helpers.ts index 7dcd118..b0e4541 100644 --- a/src/server/tools/test-helpers.ts +++ b/src/server/tools/test-helpers.ts @@ -1,4 +1,4 @@ -import type { SearchResult } from '../../types.js'; +import type { HybridQueryResult, SearchResult } from '../../types.js'; import type { ToolError, ToolErrorCode } from '../tool-error.js'; import { toolErrorSchema } from '../tool-error.js'; @@ -76,6 +76,18 @@ export function makeSearchResult(overrides?: Partial): SearchResul }; } +/** Default hybrid query outcome for mocked {@link PineconeClient.query}. */ +export function makeHybridQueryResult(overrides?: Partial): HybridQueryResult { + return { + results: overrides?.results ?? [makeSearchResult()], + degraded: overrides?.degraded ?? false, + ...(overrides?.degradation_reason !== undefined + ? { degradation_reason: overrides.degradation_reason } + : {}), + hybrid_leg_failed: overrides?.hybrid_leg_failed ?? null, + }; +} + /** Shape returned by {@link getNamespacesWithCache} `data` entries. */ export function makeNamespaceCacheEntry( namespace: string, diff --git a/src/server/url-generation.ts b/src/server/url-generation.ts index 8acbd11..e64c4c9 100644 --- a/src/server/url-generation.ts +++ b/src/server/url-generation.ts @@ -135,6 +135,15 @@ export function registerBuiltinUrlGenerators(options?: RegisterBuiltinUrlGenerat builtinGeneratorsRegistered = true; } +/** + * Clear all URL generators and the built-in registration latch. + * Used by {@link teardownServer} so a subsequent `setupServer()` can reinstall built-ins. + */ +export function resetUrlGenerationRegistry(): void { + urlGenerators.clear(); + builtinGeneratorsRegistered = false; +} + /** * Register a URL generator for a namespace, replacing any existing entry. * diff --git a/src/types.ts b/src/types.ts index 965da6d..b99e7c2 100644 --- a/src/types.ts +++ b/src/types.ts @@ -9,8 +9,9 @@ export type PineconeMetadataValue = string | number | boolean | string[]; * Configuration for `new PineconeClient(config)`. * * `apiKey` is required. All other fields are optional and fall back to - * sensible defaults; library consumers usually pass these through from - * `ServerConfig`. + * sensible defaults from `src/constants.ts`. Values are expected to come from + * {@link resolveConfig} / CLI or an equivalent resolved object — `PineconeClient` + * does not read `process.env` directly. */ export interface PineconeClientConfig { apiKey: string; @@ -34,6 +35,22 @@ export interface SearchResult { reranked: boolean; } +/** Which hybrid leg failed when the other produced hits (partial hybrid success). */ +export type HybridLegFailed = 'dense' | 'sparse' | null; + +/** + * Outcome of {@link PineconeClient.query}: result rows plus degradation signals for MCP clients. + */ +export interface HybridQueryResult { + results: SearchResult[]; + /** True when reranking was attempted and failed (rows may have `reranked: false`). */ + degraded: boolean; + /** Present when {@link degraded} is true; suitable for LLM-facing tool output. */ + degradation_reason?: string; + /** Set when exactly one of dense/sparse search failed but the other succeeded. */ + hybrid_leg_failed: HybridLegFailed; +} + export interface PineconeHit { _id: string; _score: number; @@ -132,6 +149,12 @@ export interface QueryResponse { /** Present when the query requested specific fields. */ fields?: string[]; results?: QueryResultRowShape[]; + /** True when reranking was attempted but failed; see `degradation_reason`. */ + degraded?: boolean; + /** Human-readable explanation when {@link degraded} is true. */ + degradation_reason?: string; + /** Partial hybrid failure: one leg failed while the other returned hits. */ + hybrid_leg_failed?: HybridLegFailed; } /** Internal merged hit shape before rerank (dense + sparse deduped). */ From 82057de2a85082783f2bdf124cc119e9ffcc3b60 Mon Sep 17 00:00:00 2001 From: zho Date: Tue, 19 May 2026 01:27:34 +0800 Subject: [PATCH 2/3] fixed type check errors --- src/server/tools/query-tool.test.ts | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/src/server/tools/query-tool.test.ts b/src/server/tools/query-tool.test.ts index 365bc39..167199d 100644 --- a/src/server/tools/query-tool.test.ts +++ b/src/server/tools/query-tool.test.ts @@ -219,13 +219,11 @@ describe('query tool handler (preset-driven)', () => { it('query: surfaces unreranked hits when client returns reranked:false (rerank fallback shape)', async () => { mockedGetClient.mockReturnValue({ - query: vi - .fn() - .mockResolvedValue( - makeHybridQueryResult({ - results: [makeSearchResult({ reranked: false, score: 0.5, content: 'x' })], - }) - ), + query: vi.fn().mockResolvedValue( + makeHybridQueryResult({ + results: [makeSearchResult({ reranked: false, score: 0.5, content: 'x' })], + }) + ), count: vi.fn(), } as never); From d102381818bb5b44a9e096333fa1282b3f8edf99 Mon Sep 17 00:00:00 2001 From: zho Date: Tue, 19 May 2026 03:53:51 +0800 Subject: [PATCH 3/3] addressed ai review results --- src/pinecone-client.ts | 12 ++++++++++-- src/server/suggestion-flow.ts | 16 ++++++++-------- src/server/tools/guided-query-tool.ts | 2 +- 3 files changed, 19 insertions(+), 11 deletions(-) diff --git a/src/pinecone-client.ts b/src/pinecone-client.ts index 3606072..3efde17 100644 --- a/src/pinecone-client.ts +++ b/src/pinecone-client.ts @@ -151,9 +151,17 @@ export class PineconeClient { } let hybridLegFailed: HybridLegFailed = null; - if (denseResult.status === 'rejected' && sparseResult.status === 'fulfilled') { + if ( + denseResult.status === 'rejected' && + sparseResult.status === 'fulfilled' && + sparseHits.length > 0 + ) { hybridLegFailed = 'dense'; - } else if (sparseResult.status === 'rejected' && denseResult.status === 'fulfilled') { + } else if ( + sparseResult.status === 'rejected' && + denseResult.status === 'fulfilled' && + denseHits.length > 0 + ) { hybridLegFailed = 'sparse'; } diff --git a/src/server/suggestion-flow.ts b/src/server/suggestion-flow.ts index d4aa312..2e34cde 100644 --- a/src/server/suggestion-flow.ts +++ b/src/server/suggestion-flow.ts @@ -55,14 +55,6 @@ export function requireSuggested(namespace: string): ok: false; message: string; } { - const key = normalizeNamespace(namespace); - if (!key) { - return { - ok: false, - message: 'namespace cannot be empty after trimming whitespace.', - }; - } - const cfg = getServerConfig(); if (cfg.disableSuggestFlow) { return { @@ -76,6 +68,14 @@ export function requireSuggested(namespace: string): }; } + const key = normalizeNamespace(namespace); + if (!key) { + return { + ok: false, + message: 'namespace cannot be empty after trimming whitespace.', + }; + } + const state = stateByNamespace.get(key); if (!state) { return { diff --git a/src/server/tools/guided-query-tool.ts b/src/server/tools/guided-query-tool.ts index a42d9c2..a127a04 100644 --- a/src/server/tools/guided-query-tool.ts +++ b/src/server/tools/guided-query-tool.ts @@ -212,7 +212,7 @@ export function registerGuidedQueryTool(server: McpServer): void { }); const rerank_status: 'success' | 'skipped' | 'failed' = isFast ? 'skipped' - : queryOutcome.degraded + : queryOutcome.degraded && queryOutcome.hybrid_leg_failed === null ? 'failed' : 'success'; const formattedResults = formatQueryResultRows(queryOutcome.results, {