diff --git a/src/alliance/setup.ts b/src/alliance/setup.ts index 44cfdf9..e72bc29 100644 --- a/src/alliance/setup.ts +++ b/src/alliance/setup.ts @@ -19,7 +19,7 @@ export async function setupAllianceServer(config?: ServerConfig): Promise }) { + return { + listNamespacesWithMetadata: vi + .fn() + .mockResolvedValue([{ namespace: 'papers', recordCount: 42, metadata: namespaceMetadata }]), + query: overrides?.query ?? vi.fn().mockResolvedValue(makeHybridQueryResult()), + count: vi.fn().mockResolvedValue({ count: 7, truncated: false }), + }; +} + describe('guided_query tool handler (ServerContext instance path)', () => { it('returns success with decision_trace using injected context', async () => { const listNamespacesWithMetadata = vi.fn().mockResolvedValue([ @@ -53,4 +76,137 @@ describe('guided_query tool handler (ServerContext instance path)', () => { expect(trace['rerank_status']).toBeDefined(); expect(query).toHaveBeenCalledOnce(); }); + + it('surfaces degraded and hybrid_leg_failed in result', async () => { + const query = vi.fn().mockResolvedValue( + makeHybridQueryResult({ + degraded: true, + degradation_reason: 'sparse_leg_empty', + hybrid_leg_failed: 'sparse', + }) + ); + const ctx = createTestServerContext({ + client: papersNamespaceClient({ query }) as never, + }); + const server = createMockServer(); + registerGuidedQueryTool(server as never, ctx); + const body = parseToolJson( + await server.getHandler('guided_query')!({ + user_query: 'contracts', + namespace: 'papers', + preferred_tool: 'fast', + enrich_urls: false, + }) + ); + const result = body['result'] as Record; + expect(result['degraded']).toBe(true); + expect(result['hybrid_leg_failed']).toBe('sparse'); + expect(result['degradation_reason']).toBe('sparse_leg_empty'); + }); + + it('enriches urls via ctx builtins when enrich_urls is true', async () => { + const mailingDocId = 'boost-announce@lists.boost.org/message/O5VYCDZADVDHK5Z5LAYJBHMDOAFQL7P6'; + const query = vi.fn().mockResolvedValue( + makeHybridQueryResult({ + results: [ + makeSearchResult({ + metadata: { + document_number: 'MSG-1', + title: 'T', + author: 'A', + doc_id: mailingDocId, + }, + }), + ], + }) + ); + const ctx = createTestServerContext({ + client: { + listNamespacesWithMetadata: vi.fn().mockResolvedValue([ + { + namespace: 'mailing', + recordCount: 42, + metadata: { + document_number: 'string', + title: 'string', + author: 'string', + chunk_text: 'string', + }, + }, + ]), + query, + count: vi.fn(), + } as never, + }); + registerBuiltinUrlGenerators(ctx); + const server = createMockServer(); + registerGuidedQueryTool(server as never, ctx); + const body = parseToolJson( + await server.getHandler('guided_query')!({ + user_query: 'announcement', + namespace: 'mailing', + preferred_tool: 'fast', + enrich_urls: true, + }) + ); + const result = body['result'] as Record; + const rows = result['results'] as Array<{ url: string }>; + expect(rows[0]?.url).toContain('lists.boost.org'); + }); + + it('returns TIMEOUT when orchestrator client throws timeout error', async () => { + const query = vi + .fn() + .mockRejectedValue(new Error('Timeout after 5000ms while waiting for query')); + const ctx = createTestServerContext({ + client: papersNamespaceClient({ query }) as never, + }); + const server = createMockServer(); + registerGuidedQueryTool(server as never, ctx); + const err = assertToolErrorCode( + await server.getHandler('guided_query')!({ + user_query: 'contracts', + namespace: 'papers', + preferred_tool: 'fast', + enrich_urls: false, + }), + 'TIMEOUT' + ); + expect(err.suggestion).toMatch(/retry|timeout/i); + }); + + it('does not block explicit suggest_query_params after internal suggest', async () => { + const ctx = createTestServerContext({ + client: papersNamespaceClient() as never, + }); + const guidedServer = createMockServer(); + registerGuidedQueryTool(guidedServer as never, ctx); + await guidedServer.getHandler('guided_query')!({ + user_query: 'What does the paper say?', + namespace: 'papers', + preferred_tool: 'fast', + enrich_urls: false, + }); + + const suggestServer = createMockServer(); + registerSuggestQueryParamsTool(suggestServer as never, ctx); + const suggestBody = parseToolJson( + await suggestServer.getHandler('suggest_query_params')!({ + namespace: 'papers', + user_query: 'List titles', + }) + ); + expect(suggestBody['status']).toBe('success'); + + const queryServer = createMockServer(); + registerQueryTool(queryServer as never, ctx); + const queryBody = parseToolJson( + await queryServer.getHandler('query')!({ + query_text: 'List titles', + namespace: 'papers', + preset: 'fast', + }) + ); + expect(queryBody['status']).toBe('success'); + }); }); diff --git a/src/alliance/tools/guided-query-tool.ts b/src/alliance/tools/guided-query-tool.ts index df6939a..82f47a6 100644 --- a/src/alliance/tools/guided-query-tool.ts +++ b/src/alliance/tools/guided-query-tool.ts @@ -229,6 +229,7 @@ export function registerGuidedQueryTool(server: McpServer, ctx?: ServerContext): const formattedResults = formatQueryResultRows(queryOutcome.results, { namespace, enrichUrls: enrich_urls, + ctx, }); const result: QueryResponse = { status: 'success', diff --git a/src/alliance/tools/isolated-context.context.test.ts b/src/alliance/tools/isolated-context.context.test.ts new file mode 100644 index 0000000..6c46137 --- /dev/null +++ b/src/alliance/tools/isolated-context.context.test.ts @@ -0,0 +1,79 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import { registerBuiltinUrlGenerators } from '../url-builtins.js'; +import { registerGuidedQueryTool } from './guided-query-tool.js'; +import { + createMockServer, + createTestServerContext, + isolateFromDefaultContext, + makeHybridQueryResult, + makeSearchResult, + parseToolJson, +} from '../../core/server/tools/test-helpers.js'; +import { + getDefaultServerContext, + teardownDefaultServerContext, +} from '../../core/server/server-context.js'; + +const MAILING_DOC_ID = 'boost-announce@lists.boost.org/message/O5VYCDZADVDHK5Z5LAYJBHMDOAFQL7P6'; +const EXPECTED_MAILING_URL = `https://lists.boost.org/archives/list/${MAILING_DOC_ID}/`; + +describe('isolated ServerContext with zero default context', () => { + beforeEach(() => { + isolateFromDefaultContext(); + }); + + afterEach(() => { + teardownDefaultServerContext(); + }); + + it('guided_query enrich_urls uses ctx builtins, not default registry', async () => { + const listNamespacesWithMetadata = vi.fn().mockResolvedValue([ + { + namespace: 'mailing', + recordCount: 42, + metadata: { + document_number: 'string', + title: 'string', + author: 'string', + chunk_text: 'string', + }, + }, + ]); + const query = vi.fn().mockResolvedValue( + makeHybridQueryResult({ + results: [ + makeSearchResult({ + metadata: { + document_number: 'MSG-1', + title: 'T', + author: 'A', + doc_id: MAILING_DOC_ID, + }, + }), + ], + }) + ); + const ctx = createTestServerContext({ + client: { + listNamespacesWithMetadata, + query, + count: vi.fn(), + } as never, + }); + registerBuiltinUrlGenerators(ctx); + + const server = createMockServer(); + registerGuidedQueryTool(server as never, ctx); + const raw = await server.getHandler('guided_query')!({ + user_query: 'What was announced?', + namespace: 'mailing', + preferred_tool: 'fast', + enrich_urls: true, + }); + const body = parseToolJson(raw); + const result = body['result'] as Record; + const rows = result['results'] as Array<{ url: string }>; + expect(rows[0]?.url).toBe(EXPECTED_MAILING_URL); + expect(getDefaultServerContext().hasUrlGenerator('mailing')).toBe(false); + }); +}); diff --git a/src/alliance/tools/suggest-query-params-tool.context.test.ts b/src/alliance/tools/suggest-query-params-tool.context.test.ts index 5539b63..7455d98 100644 --- a/src/alliance/tools/suggest-query-params-tool.context.test.ts +++ b/src/alliance/tools/suggest-query-params-tool.context.test.ts @@ -1,12 +1,34 @@ -import { describe, expect, it, vi } from 'vitest'; +import { afterEach, describe, expect, it, vi } from 'vitest'; +import { registerQueryTool } from '../../core/server/tools/query-tool.js'; import { registerSuggestQueryParamsTool } from './suggest-query-params-tool.js'; import { createMockServer, createTestServerContext, + makeHybridQueryResult, parseToolJson, } from '../../core/server/tools/test-helpers.js'; +const namespaceMetadata = { + document_number: 'string', + title: 'string', + url: 'string', + author: 'string', + chunk_text: 'string', +}; + +function mockNamespacesClient() { + return { + listNamespacesWithMetadata: vi + .fn() + .mockResolvedValue([{ namespace: 'wg21', recordCount: 42, metadata: namespaceMetadata }]), + }; +} + describe('suggest_query_params tool handler (ServerContext instance path)', () => { + afterEach(() => { + vi.useRealTimers(); + }); + it('marks suggest-flow on injected context when namespace exists', async () => { const listNamespacesWithMetadata = vi.fn().mockResolvedValue([ { @@ -41,4 +63,55 @@ describe('suggest_query_params tool handler (ServerContext instance path)', () = const flowCheck = ctx.requireSuggested('wg21'); expect(flowCheck.ok).toBe(true); }); + + it('updates suggestion state when called twice for the same namespace', async () => { + const ctx = createTestServerContext({ + client: mockNamespacesClient() as never, + }); + const server = createMockServer(); + registerSuggestQueryParamsTool(server as never, ctx); + const handler = server.getHandler('suggest_query_params')!; + + await handler({ namespace: 'wg21', user_query: 'List papers with titles' }); + await handler({ namespace: 'wg21', user_query: 'how many records match?' }); + + const flowCheck = ctx.requireSuggested('wg21'); + expect(flowCheck.ok).toBe(true); + if (flowCheck.ok) { + expect(flowCheck.flow.user_query).toBe('how many records match?'); + expect(flowCheck.flow.recommended_tool).toBe('count'); + } + }); + + it('re-suggest after expiry allows query without FLOW_GATE', async () => { + vi.useFakeTimers(); + const query = vi.fn().mockResolvedValue(makeHybridQueryResult()); + const ctx = createTestServerContext({ + config: { cacheTtlSeconds: 1 }, + client: { ...mockNamespacesClient(), query } as never, + }); + const suggestServer = createMockServer(); + registerSuggestQueryParamsTool(suggestServer as never, ctx); + const suggestHandler = suggestServer.getHandler('suggest_query_params')!; + + await suggestHandler({ namespace: 'wg21', user_query: 'contracts' }); + vi.advanceTimersByTime(2000); + + const reSuggestBody = parseToolJson( + await suggestHandler({ namespace: 'wg21', user_query: 'contracts again' }) + ); + expect(reSuggestBody['status']).toBe('success'); + + const queryServer = createMockServer(); + registerQueryTool(queryServer as never, ctx); + const queryBody = parseToolJson( + await queryServer.getHandler('query')!({ + query_text: 'contracts again', + namespace: 'wg21', + preset: 'fast', + }) + ); + expect(queryBody['status']).toBe('success'); + expect(query).toHaveBeenCalledOnce(); + }); }); diff --git a/src/alliance/url-builtins.context.test.ts b/src/alliance/url-builtins.context.test.ts new file mode 100644 index 0000000..1d763f3 --- /dev/null +++ b/src/alliance/url-builtins.context.test.ts @@ -0,0 +1,22 @@ +import { describe, expect, it } from 'vitest'; +import { registerBuiltinUrlGenerators } from './url-builtins.js'; +import { createTestServerContext } from '../core/server/tools/test-helpers.js'; + +const MAILING_DOC_ID = 'boost-announce@lists.boost.org/message/O5VYCDZADVDHK5Z5LAYJBHMDOAFQL7P6'; + +describe('registerBuiltinUrlGenerators (ServerContext instance path)', () => { + it('registers builtins only on the target context', () => { + const ctxA = createTestServerContext(); + const ctxB = createTestServerContext(); + registerBuiltinUrlGenerators(ctxA); + + const metadata = { doc_id: MAILING_DOC_ID }; + const fromA = ctxA.generateUrlForNamespace('mailing', metadata); + const fromB = ctxB.generateUrlForNamespace('mailing', metadata); + + expect(fromA.url).toContain('lists.boost.org'); + expect(fromA.method).toBe('generated.mailing'); + expect(fromB.method).toBe('unavailable'); + expect(fromB.url).toBeNull(); + }); +}); diff --git a/src/alliance/url-builtins.ts b/src/alliance/url-builtins.ts index ae2ea13..48f6f68 100644 --- a/src/alliance/url-builtins.ts +++ b/src/alliance/url-builtins.ts @@ -2,6 +2,7 @@ * C++ Alliance domain-specific URL generators (Boost mailing list, Slack). */ +import type { ServerContext } from '../core/server/server-context.js'; import { registerUrlGenerator } from '../core/server/url-registry.js'; import type { UrlGenerationResult } from '../core/server/url-registry.js'; @@ -62,7 +63,8 @@ export function generatorSlackCpplang(metadata: Record): UrlGen }; } -let builtinGeneratorsRegistered = false; +const builtinGeneratorsRegisteredContexts = new WeakSet(); +let defaultBuiltinGeneratorsRegistered = false; /** Options for {@link registerBuiltinUrlGenerators}. */ export type RegisterBuiltinUrlGeneratorsOptions = { @@ -73,18 +75,57 @@ export type RegisterBuiltinUrlGeneratorsOptions = { reinstallBuiltins?: boolean; }; -/** - * Register built-in Alliance generators (`mailing`, `slack-Cpplang`). - */ -export function registerBuiltinUrlGenerators(options?: RegisterBuiltinUrlGeneratorsOptions): void { +function registerBuiltinsOnContext(ctx: ServerContext): void { + ctx.registerUrlGenerator('mailing', generatorMailing); + ctx.registerUrlGenerator('slack-Cpplang', generatorSlackCpplang); +} + +function isServerContext(value: unknown): value is ServerContext { + return ( + typeof value === 'object' && + value !== null && + typeof (value as ServerContext).getConfig === 'function' && + typeof (value as ServerContext).generateUrlForNamespace === 'function' && + typeof (value as ServerContext).registerUrlGenerator === 'function' + ); +} + +function registerBuiltinsOnDefaultContext(options?: RegisterBuiltinUrlGeneratorsOptions): void { if (options?.reinstallBuiltins) { registerUrlGenerator('mailing', generatorMailing); registerUrlGenerator('slack-Cpplang', generatorSlackCpplang); - builtinGeneratorsRegistered = true; + defaultBuiltinGeneratorsRegistered = true; return; } - if (builtinGeneratorsRegistered) return; + if (defaultBuiltinGeneratorsRegistered) return; registerUrlGenerator('mailing', generatorMailing); registerUrlGenerator('slack-Cpplang', generatorSlackCpplang); - builtinGeneratorsRegistered = true; + defaultBuiltinGeneratorsRegistered = true; +} + +/** Register built-in Alliance generators on the process-default context (legacy). */ +export function registerBuiltinUrlGenerators(options?: RegisterBuiltinUrlGeneratorsOptions): void; +/** Register built-in Alliance generators on the given {@link ServerContext}. */ +export function registerBuiltinUrlGenerators( + ctx: ServerContext, + options?: RegisterBuiltinUrlGeneratorsOptions +): void; +export function registerBuiltinUrlGenerators( + ctxOrOptions?: ServerContext | RegisterBuiltinUrlGeneratorsOptions, + options?: RegisterBuiltinUrlGeneratorsOptions +): void { + if (isServerContext(ctxOrOptions)) { + const ctx = ctxOrOptions; + if (options?.reinstallBuiltins) { + registerBuiltinsOnContext(ctx); + builtinGeneratorsRegisteredContexts.add(ctx); + return; + } + if (builtinGeneratorsRegisteredContexts.has(ctx)) return; + registerBuiltinsOnContext(ctx); + builtinGeneratorsRegisteredContexts.add(ctx); + return; + } + + registerBuiltinsOnDefaultContext(ctxOrOptions); } diff --git a/src/core/server/format-query-result.context.test.ts b/src/core/server/format-query-result.context.test.ts new file mode 100644 index 0000000..af8f24e --- /dev/null +++ b/src/core/server/format-query-result.context.test.ts @@ -0,0 +1,36 @@ +import { afterEach, describe, expect, it } from 'vitest'; +import { formatSearchResultAsRow } from './format-query-result.js'; +import { ServerContext } from './server-context.js'; +import { isolateFromDefaultContext, resolveTestConfig } from './tools/test-helpers.js'; +import { teardownDefaultServerContext } from './server-context.js'; +import type { SearchResult } from '../../types.js'; + +describe('formatSearchResultAsRow (ServerContext instance path)', () => { + afterEach(() => { + teardownDefaultServerContext(); + }); + + it('enriches url from injected ctx registry when default context has no generator', () => { + isolateFromDefaultContext(); + const ctx = new ServerContext(resolveTestConfig()); + ctx.registerUrlGenerator('papers', () => ({ + url: 'https://ctx.example/papers/doc-1', + method: 'generated.custom', + })); + + const doc: SearchResult = { + id: 'v1', + content: 'body', + score: 0.9, + metadata: { document_number: 'DOC-1', title: 'T', author: 'A' }, + reranked: false, + }; + + const row = formatSearchResultAsRow(doc, { + namespace: 'papers', + enrichUrls: true, + ctx, + }); + expect(row.url).toBe('https://ctx.example/papers/doc-1'); + }); +}); diff --git a/src/core/server/format-query-result.ts b/src/core/server/format-query-result.ts index 700f5a5..9c6914a 100644 --- a/src/core/server/format-query-result.ts +++ b/src/core/server/format-query-result.ts @@ -5,8 +5,16 @@ import type { PineconeMetadataValue, SearchResult } from '../../types.js'; import { warn as logWarn } from '../../logger.js'; +import type { ServerContext } from './server-context.js'; import { generateUrlForNamespace } from './url-registry.js'; +export type FormatQueryResultOptions = { + namespace?: string; + enrichUrls?: boolean; + contentMaxLength?: number; + ctx?: ServerContext; +}; + const DEFAULT_CONTENT_MAX_LENGTH = 2000; /** @@ -53,17 +61,15 @@ export function resetPaperNumberDeprecationLatchForTests(): void { */ export function formatSearchResultAsRow( doc: SearchResult, - options?: { - namespace?: string; - enrichUrls?: boolean; - contentMaxLength?: number; - } + options?: FormatQueryResultOptions ): QueryResultRow { const contentMaxLength = options?.contentMaxLength ?? DEFAULT_CONTENT_MAX_LENGTH; const metadata = { ...doc.metadata } as Record; if (options?.enrichUrls && options?.namespace) { - const generated = generateUrlForNamespace(options.namespace, metadata); + const generated = options.ctx + ? options.ctx.generateUrlForNamespace(options.namespace, metadata) + : generateUrlForNamespace(options.namespace, metadata); const existingUrl = metadata['url']; const urlIsBlank = typeof existingUrl !== 'string' || existingUrl.trim() === ''; if (generated.url && urlIsBlank) { @@ -105,11 +111,7 @@ export function formatSearchResultAsRow( */ export function formatQueryResultRows( results: SearchResult[], - options?: { - namespace?: string; - enrichUrls?: boolean; - contentMaxLength?: number; - } + options?: FormatQueryResultOptions ): QueryResultRow[] { return results.map((doc) => formatSearchResultAsRow(doc, options)); } diff --git a/src/core/server/server-context.test.ts b/src/core/server/server-context.test.ts index 66d5314..9122a8d 100644 --- a/src/core/server/server-context.test.ts +++ b/src/core/server/server-context.test.ts @@ -5,13 +5,13 @@ import { ServerContext, createServer, getDefaultServerContext, - setDefaultServerContext, teardownDefaultServerContext, } from './server-context.js'; describe('ServerContext', () => { afterEach(() => { teardownDefaultServerContext(); + vi.useRealTimers(); }); const testConfig = () => resolveTestConfig(); @@ -142,8 +142,24 @@ describe('ServerContext', () => { it('teardownDefaultServerContext clears process default', () => { createServer(testConfig()); teardownDefaultServerContext(); - setDefaultServerContext(null); const fresh = getDefaultServerContext(); expect(fresh).not.toBeNull(); }); + + it('requireSuggested returns expiry message after TTL on instance context', () => { + vi.useFakeTimers(); + const ctx = new ServerContext(resolveTestConfig({ cacheTtlSeconds: 1 })); + ctx.markSuggested('wg21', { + recommended_tool: 'fast', + suggested_fields: ['title'], + user_query: 'contracts', + }); + expect(ctx.requireSuggested('wg21').ok).toBe(true); + vi.advanceTimersByTime(2000); + const expired = ctx.requireSuggested('wg21'); + expect(expired.ok).toBe(false); + if (!expired.ok) { + expect(expired.message).toMatch(/expired/); + } + }); }); diff --git a/src/core/server/suggestion-flow.test.ts b/src/core/server/suggestion-flow.test.ts new file mode 100644 index 0000000..1a4c265 --- /dev/null +++ b/src/core/server/suggestion-flow.test.ts @@ -0,0 +1,71 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import { markSuggested, requireSuggested, resetSuggestionFlow } from './suggestion-flow.js'; +import { createServer, teardownDefaultServerContext } from './server-context.js'; +import { resolveTestConfig } from './tools/test-helpers.js'; + +describe('suggestion-flow facade (default ServerContext)', () => { + beforeEach(() => { + teardownDefaultServerContext(); + createServer(resolveTestConfig({ disableSuggestFlow: false, cacheTtlSeconds: 1 })); + }); + + afterEach(() => { + teardownDefaultServerContext(); + vi.useRealTimers(); + }); + + it('markSuggested then requireSuggested succeeds', () => { + markSuggested('wg21', { + recommended_tool: 'fast', + suggested_fields: ['title'], + user_query: 'contracts', + }); + const check = requireSuggested('wg21'); + expect(check.ok).toBe(true); + if (check.ok) { + expect(check.flow.user_query).toBe('contracts'); + expect(check.flow.recommended_tool).toBe('fast'); + } + }); + + it('requireSuggested fails when namespace was never suggested', () => { + const check = requireSuggested('wg21'); + expect(check.ok).toBe(false); + if (!check.ok) { + expect(check.message).toMatch(/suggest_query_params first/); + } + }); + + it('requireSuggested fails after TTL expiry', () => { + vi.useFakeTimers(); + markSuggested('wg21', { + recommended_tool: 'detailed', + suggested_fields: ['chunk_text'], + user_query: 'q', + }); + vi.advanceTimersByTime(2000); + const check = requireSuggested('wg21'); + expect(check.ok).toBe(false); + if (!check.ok) { + expect(check.message).toMatch(/expired/); + } + }); + + it('requireSuggested bypasses gate when disableSuggestFlow is true', () => { + teardownDefaultServerContext(); + createServer(resolveTestConfig({ disableSuggestFlow: true, cacheTtlSeconds: 1 })); + const check = requireSuggested('wg21'); + expect(check.ok).toBe(true); + }); + + it('resetSuggestionFlow clears prior suggestion state', () => { + markSuggested('wg21', { + recommended_tool: 'count', + suggested_fields: [], + user_query: 'how many', + }); + expect(requireSuggested('wg21').ok).toBe(true); + resetSuggestionFlow(); + expect(requireSuggested('wg21').ok).toBe(false); + }); +}); diff --git a/src/core/server/tools/count-tool.context.test.ts b/src/core/server/tools/count-tool.context.test.ts index c5c108f..13c17f8 100644 --- a/src/core/server/tools/count-tool.context.test.ts +++ b/src/core/server/tools/count-tool.context.test.ts @@ -48,4 +48,22 @@ describe('count tool handler (ServerContext instance path)', () => { const err = assertToolErrorCode(raw, 'FLOW_GATE'); expect(err.suggestion).toBe("Call suggest_query_params for namespace 'wg21' first"); }); + + it('succeeds without prior suggest when disableSuggestFlow is true', async () => { + const count = vi.fn().mockResolvedValue({ count: 2, truncated: false }); + const ctx = createTestServerContext({ + config: { disableSuggestFlow: true }, + client: { count } as never, + }); + const server = createMockServer(); + registerCountTool(server as never, ctx); + const body = parseToolJson( + await server.getHandler('count')!({ + namespace: 'wg21', + query_text: 'how many', + }) + ); + expect(body['status']).toBe('success'); + expect(count).toHaveBeenCalledOnce(); + }); }); diff --git a/src/core/server/tools/keyword-search-tool.ts b/src/core/server/tools/keyword-search-tool.ts index 7b801f3..b4e226d 100644 --- a/src/core/server/tools/keyword-search-tool.ts +++ b/src/core/server/tools/keyword-search-tool.ts @@ -87,6 +87,7 @@ async function executeKeywordSearch( const formattedResults = formatQueryResultRows(results, { namespace: normalizedNamespace, + ctx, }); const response: KeywordSearchResponse = { diff --git a/src/core/server/tools/query-tool.context.test.ts b/src/core/server/tools/query-tool.context.test.ts index e8dd6a8..96221ab 100644 --- a/src/core/server/tools/query-tool.context.test.ts +++ b/src/core/server/tools/query-tool.context.test.ts @@ -51,4 +51,105 @@ describe('query tool handler (ServerContext instance path)', () => { const err = assertToolErrorCode(raw, 'FLOW_GATE'); expect(err.suggestion).toBe("Call suggest_query_params for namespace 'wg21' first"); }); + + it('succeeds without prior suggest when disableSuggestFlow is true', async () => { + const query = vi.fn().mockResolvedValue(makeHybridQueryResult()); + const ctx = createTestServerContext({ + config: { disableSuggestFlow: true }, + client: { query } as never, + }); + const server = createMockServer(); + registerQueryTool(server as never, ctx); + const body = parseToolJson( + await server.getHandler('query')!({ + query_text: 'hello', + namespace: 'wg21', + preset: 'fast', + }) + ); + expect(body['status']).toBe('success'); + expect(query).toHaveBeenCalledOnce(); + }); + + it('forwards rerank_skipped_reason and degradation_reason on injected context', async () => { + const query = vi.fn().mockResolvedValue( + makeHybridQueryResult({ + rerank_skipped_reason: 'no_model', + degradation_reason: 'rerank_skipped_no_model: set PINECONE_RERANK_MODEL', + }) + ); + const ctx = createTestServerContext({ + client: { query } as never, + }); + ctx.markSuggested('wg21', { + recommended_tool: 'detailed', + suggested_fields: ['chunk_text'], + user_query: 'q', + }); + const server = createMockServer(); + registerQueryTool(server as never, ctx); + const body = parseToolJson( + await server.getHandler('query')!({ + query_text: 'hello', + namespace: 'wg21', + preset: 'detailed', + }) + ); + expect(body['rerank_skipped_reason']).toBe('no_model'); + expect(body['degradation_reason']).toMatch(/rerank_skipped_no_model/); + }); + + it('forwards hybrid_leg_failed for dense and sparse partial hybrid', async () => { + const query = vi.fn(); + const ctx = createTestServerContext({ + client: { query } as never, + }); + ctx.markSuggested('wg21', { + recommended_tool: 'fast', + suggested_fields: ['title'], + user_query: 'q', + }); + const server = createMockServer(); + registerQueryTool(server as never, ctx); + const handler = server.getHandler('query')!; + + query.mockResolvedValue(makeHybridQueryResult({ hybrid_leg_failed: 'dense', degraded: false })); + const denseBody = parseToolJson( + await handler({ query_text: 'a', namespace: 'wg21', preset: 'fast' }) + ); + expect(denseBody['hybrid_leg_failed']).toBe('dense'); + + query.mockResolvedValue( + makeHybridQueryResult({ hybrid_leg_failed: 'sparse', degraded: false }) + ); + const sparseBody = parseToolJson( + await handler({ query_text: 'b', namespace: 'wg21', preset: 'fast' }) + ); + expect(sparseBody['hybrid_leg_failed']).toBe('sparse'); + }); + + it('returns TIMEOUT when client throws timeout error', async () => { + const query = vi + .fn() + .mockRejectedValue(new Error('Timeout after 5000ms while waiting for query')); + const ctx = createTestServerContext({ + client: { query } as never, + }); + ctx.markSuggested('wg21', { + recommended_tool: 'fast', + suggested_fields: ['title'], + user_query: 'q', + }); + const server = createMockServer(); + registerQueryTool(server as never, ctx); + const err = assertToolErrorCode( + await server.getHandler('query')!({ + query_text: 'hello', + namespace: 'wg21', + preset: 'fast', + }), + 'TIMEOUT' + ); + expect(err.suggestion).toMatch(/retry|timeout/i); + }); }); diff --git a/src/core/server/tools/query-tool.ts b/src/core/server/tools/query-tool.ts index 7f449eb..2020666 100644 --- a/src/core/server/tools/query-tool.ts +++ b/src/core/server/tools/query-tool.ts @@ -69,7 +69,7 @@ async function executeQuery(params: QueryExecParams, ctx?: ServerContext) { fields: fields?.length ? fields : undefined, }); - const formattedResults = formatQueryResultRows(queryOutcome.results); + const formattedResults = formatQueryResultRows(queryOutcome.results, { ctx }); const response: QueryResponse = { status: 'success', diff --git a/src/core/server/tools/test-helpers.ts b/src/core/server/tools/test-helpers.ts index 17412a0..770dc0b 100644 --- a/src/core/server/tools/test-helpers.ts +++ b/src/core/server/tools/test-helpers.ts @@ -2,7 +2,7 @@ import type { HybridQueryResult, SearchResult } from '../../../types.js'; import { resolveConfig } from '../../config.js'; import type { PineconeClient } from '../../pinecone-client.js'; import type { ConfigOverrides, ServerConfig } from '../../config.js'; -import { ServerContext } from '../server-context.js'; +import { ServerContext, teardownDefaultServerContext } from '../server-context.js'; import type { ToolError, ToolErrorCode } from '../tool-error.js'; import { toolErrorSchema } from '../tool-error.js'; @@ -127,6 +127,11 @@ export function resolveTestConfig(overrides: ConfigOverrides = {}): ServerConfig }); } +/** Clear process-default context so tests exercise an isolated instance only. */ +export function isolateFromDefaultContext(): void { + teardownDefaultServerContext(); +} + /** Build an isolated {@link ServerContext} for instance-path tool tests. */ export function createTestServerContext(options?: { config?: ConfigOverrides; diff --git a/vitest.config.ts b/vitest.config.ts index 95efacf..1ab26de 100644 --- a/vitest.config.ts +++ b/vitest.config.ts @@ -16,7 +16,7 @@ export default defineConfig({ thresholds: { lines: 73, statements: 72, - branches: 58, + branches: 65, functions: 76, }, },