diff --git a/src/alliance/setup.ts b/src/alliance/setup.ts index e72bc29..f1495bc 100644 --- a/src/alliance/setup.ts +++ b/src/alliance/setup.ts @@ -1,9 +1,8 @@ -import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; import { ALLIANCE_SERVER_INSTRUCTIONS } from '../constants.js'; import type { ServerConfig } from '../core/config.js'; import { getDefaultServerContext } from '../core/server/server-context.js'; import { resolveAllianceConfig } from './config.js'; -import { setupCoreServer } from '../core/setup.js'; +import { setupCoreServer, type ServerHandle } from '../core/setup.js'; import { registerBuiltinUrlGenerators } from './url-builtins.js'; import { registerGuidedQueryTool } from './tools/guided-query-tool.js'; import { registerSuggestQueryParamsTool } from './tools/suggest-query-params-tool.js'; @@ -14,7 +13,7 @@ import { registerSuggestQueryParamsTool } from './tools/suggest-query-params-too * * When `config` is omitted, resolves env via {@link resolveAllianceConfig} (Alliance index/rerank defaults when unset). */ -export async function setupAllianceServer(config?: ServerConfig): Promise { +export async function setupAllianceServer(config?: ServerConfig): Promise { const server = await setupCoreServer(config ?? resolveAllianceConfig({}), { instructions: ALLIANCE_SERVER_INSTRUCTIONS, }); diff --git a/src/alliance/tools/guided-query-tool.ts b/src/alliance/tools/guided-query-tool.ts index 82f47a6..0471db1 100644 --- a/src/alliance/tools/guided-query-tool.ts +++ b/src/alliance/tools/guided-query-tool.ts @@ -17,6 +17,7 @@ import { suggestQueryParams } from '../../core/server/query-suggestion.js'; import { markSuggested } from '../../core/server/suggestion-flow.js'; import { classifyToolCatchError, + lifecycleToolError, logToolError, pineconeToolError, validationToolError, @@ -82,6 +83,9 @@ export function registerGuidedQueryTool(server: McpServer, ctx?: ServerContext): }, async (params) => { try { + if (ctx?.disposed) { + return jsonErrorResponse(lifecycleToolError('ServerContext has been disposed')); + } const { user_query, namespace: inputNamespace, diff --git a/src/alliance/tools/suggest-query-params-tool.ts b/src/alliance/tools/suggest-query-params-tool.ts index d059309..83c36c4 100644 --- a/src/alliance/tools/suggest-query-params-tool.ts +++ b/src/alliance/tools/suggest-query-params-tool.ts @@ -7,6 +7,7 @@ import type { ServerContext } from '../../core/server/server-context.js'; import { markSuggested } from '../../core/server/suggestion-flow.js'; import { classifyToolCatchError, + lifecycleToolError, logToolError, validationToolError, } from '../../core/server/tool-error.js'; @@ -37,6 +38,9 @@ export function registerSuggestQueryParamsTool(server: McpServer, ctx?: ServerCo }, async (params) => { try { + if (ctx?.disposed) { + return jsonErrorResponse(lifecycleToolError('ServerContext has been disposed')); + } const { namespace, user_query } = params; if (!user_query?.trim()) { return jsonErrorResponse(validationToolError('user_query cannot be empty', 'user_query')); diff --git a/src/core/index.ts b/src/core/index.ts index 7d96c75..5611f89 100644 --- a/src/core/index.ts +++ b/src/core/index.ts @@ -41,4 +41,4 @@ export type { HybridQueryResult, HybridLegFailed, } from '../types.js'; -export { setupCoreServer, teardownServer } from './setup.js'; +export { setupCoreServer, teardownServer, type ServerHandle } from './setup.js'; diff --git a/src/core/server/server-context.lifecycle.test.ts b/src/core/server/server-context.lifecycle.test.ts new file mode 100644 index 0000000..8b582ac --- /dev/null +++ b/src/core/server/server-context.lifecycle.test.ts @@ -0,0 +1,55 @@ +import { afterEach, describe, expect, it } from 'vitest'; +import { PineconeClient } from '../pinecone-client.js'; +import { resolveConfig } from '../config.js'; +import { setPineconeClient } from '../server/client-context.js'; +import { setupCoreServer, teardownServer } from '../setup.js'; +import { resolveTestConfig } from './tools/test-helpers.js'; +import { ServerContext, createServer, teardownDefaultServerContext } from './server-context.js'; + +describe('ServerContext lifecycle', () => { + afterEach(() => { + teardownServer(); + }); + + it('sets disposed after teardown()', () => { + const ctx = new ServerContext(resolveTestConfig()); + ctx.teardown(); + expect(ctx.disposed).toBe(true); + }); + + it('sets disposed after teardownDefaultServerContext()', () => { + const ctx = createServer(resolveTestConfig()); + teardownDefaultServerContext(); + expect(ctx.disposed).toBe(true); + }); + + it('await using disposes ServerContext on scope exit', async () => { + const config = resolveTestConfig(); + let ctx!: ServerContext; + await (async () => { + await using scoped = new ServerContext(config); + ctx = scoped; + expect(scoped.disposed).toBe(false); + })(); + expect(ctx.disposed).toBe(true); + }); + + it('await using on setupCoreServer return value tears down and allows re-setup', async () => { + const cfg = resolveConfig({ apiKey: 'lifecycle-await-key', indexName: 'test-index' }); + setPineconeClient( + new PineconeClient({ + apiKey: cfg.apiKey, + indexName: cfg.indexName, + rerankModel: cfg.rerankModel, + defaultTopK: cfg.defaultTopK, + }) + ); + + await (async () => { + await using _server = await setupCoreServer(cfg); + })(); + + await expect(setupCoreServer(cfg)).resolves.toBeDefined(); + teardownServer(); + }); +}); diff --git a/src/core/server/server-context.ts b/src/core/server/server-context.ts index 91d1c63..36f60b5 100644 --- a/src/core/server/server-context.ts +++ b/src/core/server/server-context.ts @@ -43,7 +43,8 @@ function buildPineconeClient(config: ServerConfig): PineconeClient { * Encapsulates per-server state: Pinecone client, config, URL registry, * suggest-flow gate, and namespaces cache. */ -export class ServerContext { +export class ServerContext implements AsyncDisposable { + disposed = false; private client: PineconeClient | null = null; private configValue: ServerConfig | null = null; private readonly urlGenerators = new Map(); @@ -257,12 +258,21 @@ export class ServerContext { /** Clear all encapsulated state (client handle, caches, registries). */ teardown(): void { + this.disposed = true; this.client = null; this.configValue = null; this.urlGenerators.clear(); this.suggestionFlow.clear(); this.namespacesCache = null; } + + async [Symbol.asyncDispose](): Promise { + this.teardown(); + if (defaultContext === this) { + defaultContext = null; + pendingConfig = null; + } + } } let defaultContext: ServerContext | null = null; diff --git a/src/core/server/tool-error.test.ts b/src/core/server/tool-error.test.ts index 87b8000..73821bb 100644 --- a/src/core/server/tool-error.test.ts +++ b/src/core/server/tool-error.test.ts @@ -2,6 +2,7 @@ import { describe, expect, it } from 'vitest'; import { classifyToolCatchError, flowGateToolError, + lifecycleToolError, toolErrorSchema, validationToolError, } from './tool-error.js'; @@ -37,6 +38,14 @@ describe('ToolError schema and builders', () => { expect(toolErrorSchema.parse(err).code).toBe('PINECONE_ERROR'); }); + it('LIFECYCLE: not recoverable and parses', () => { + const err = lifecycleToolError('ServerContext has been disposed'); + const parsed = toolErrorSchema.parse(err); + expect(parsed.code).toBe('LIFECYCLE'); + expect(parsed.recoverable).toBe(false); + expect(parsed.message).toContain('disposed'); + }); + it('TIMEOUT: classifyToolCatchError matches withTimeout message prefix', () => { const err = classifyToolCatchError( new Error('Timeout after 100ms while waiting for query'), diff --git a/src/core/server/tool-error.ts b/src/core/server/tool-error.ts index dffe42d..b241071 100644 --- a/src/core/server/tool-error.ts +++ b/src/core/server/tool-error.ts @@ -17,7 +17,13 @@ export function logToolError(toolName: string, error: unknown): void { logError(`Error in ${toolName} tool`, error); } -export const toolErrorCodeSchema = z.enum(['FLOW_GATE', 'VALIDATION', 'PINECONE_ERROR', 'TIMEOUT']); +export const toolErrorCodeSchema = z.enum([ + 'FLOW_GATE', + 'VALIDATION', + 'PINECONE_ERROR', + 'TIMEOUT', + 'LIFECYCLE', +]); export type ToolErrorCode = z.infer; const flowGateToolErrorSchema = z.object({ @@ -49,11 +55,19 @@ const timeoutToolErrorSchema = z.object({ suggestion: z.string().optional(), }); +const lifecycleToolErrorSchema = z.object({ + code: z.literal('LIFECYCLE'), + message: z.string(), + recoverable: z.literal(false), + suggestion: z.string().optional(), +}); + export const toolErrorSchema = z.discriminatedUnion('code', [ flowGateToolErrorSchema, validationToolErrorSchema, pineconeToolErrorSchema, timeoutToolErrorSchema, + lifecycleToolErrorSchema, ]); export type ToolError = z.infer; @@ -107,6 +121,14 @@ export function timeoutToolError(message: string, options?: { suggestion?: strin }; } +export function lifecycleToolError(message: string): ToolError { + return { + code: 'LIFECYCLE', + message, + recoverable: false, + }; +} + function rawErrorMessage(error: unknown): string { return error instanceof Error ? error.message : String(error); } diff --git a/src/core/server/tools/count-tool.ts b/src/core/server/tools/count-tool.ts index a31809f..50e6fc6 100644 --- a/src/core/server/tools/count-tool.ts +++ b/src/core/server/tools/count-tool.ts @@ -8,6 +8,7 @@ import { requireSuggested } from '../suggestion-flow.js'; import { classifyToolCatchError, flowGateToolError, + lifecycleToolError, logToolError, validationToolError, } from '../tool-error.js'; @@ -30,6 +31,9 @@ type CountExecParams = { async function executeCount(params: CountExecParams, ctx?: ServerContext) { try { + if (ctx?.disposed) { + return jsonErrorResponse(lifecycleToolError('ServerContext has been disposed')); + } const { namespace, query_text, metadata_filter } = params; const nsNorm = normalizeNamespace(namespace); if (!nsNorm) { diff --git a/src/core/server/tools/generate-urls-tool.ts b/src/core/server/tools/generate-urls-tool.ts index 43a6b4b..8ad19ae 100644 --- a/src/core/server/tools/generate-urls-tool.ts +++ b/src/core/server/tools/generate-urls-tool.ts @@ -3,7 +3,12 @@ import { z } from 'zod'; import { normalizeNamespace } from '../namespace-utils.js'; import type { ServerContext } from '../server-context.js'; import { generateUrlForNamespace } from '../url-registry.js'; -import { classifyToolCatchError, logToolError, validationToolError } from '../tool-error.js'; +import { + classifyToolCatchError, + lifecycleToolError, + 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). */ @@ -39,6 +44,9 @@ export function registerGenerateUrlsTool(server: McpServer, ctx?: ServerContext) }, async (params) => { try { + if (ctx?.disposed) { + return jsonErrorResponse(lifecycleToolError('ServerContext has been disposed')); + } const { namespace, records } = params; const nsNorm = normalizeNamespace(namespace); if (!nsNorm) { diff --git a/src/core/server/tools/keyword-search-tool.ts b/src/core/server/tools/keyword-search-tool.ts index b4e226d..e594dc6 100644 --- a/src/core/server/tools/keyword-search-tool.ts +++ b/src/core/server/tools/keyword-search-tool.ts @@ -6,7 +6,12 @@ import { formatQueryResultRows } from '../format-query-result.js'; import type { ServerContext } from '../server-context.js'; import { metadataFilterSchema, validateMetadataFilterDetailed } from '../metadata-filter.js'; import type { ToolError } from '../tool-error.js'; -import { classifyToolCatchError, logToolError, validationToolError } from '../tool-error.js'; +import { + classifyToolCatchError, + lifecycleToolError, + logToolError, + validationToolError, +} from '../tool-error.js'; import { jsonErrorResponse, jsonResponse } from '../tool-response.js'; /** Success response shape for keyword_search (aligned with query tool fields). */ @@ -47,6 +52,9 @@ async function executeKeywordSearch( }, ctx?: ServerContext ): Promise { + if (ctx?.disposed) { + return { ok: false, error: lifecycleToolError('ServerContext has been disposed') }; + } const { query_text, namespace, top_k, metadata_filter, fields } = params; const normalizedQuery = query_text.trim(); diff --git a/src/core/server/tools/lifecycle.context.test.ts b/src/core/server/tools/lifecycle.context.test.ts new file mode 100644 index 0000000..9f6b9cf --- /dev/null +++ b/src/core/server/tools/lifecycle.context.test.ts @@ -0,0 +1,115 @@ +import { afterEach, describe, expect, it, vi } from 'vitest'; +import { PineconeClient } from '../../pinecone-client.js'; +import { resolveConfig } from '../../config.js'; +import { + getDefaultServerContext, + setPineconeClient, + setupCoreServer, + teardownServer, +} from '../../index.js'; +import { registerQueryTool } from './query-tool.js'; +import { + assertToolErrorCode, + createMockServer, + createTestServerContext, + makeHybridQueryResult, + parseToolJson, + resolveTestConfig, +} from './test-helpers.js'; + +describe('tool handler lifecycle guards', () => { + afterEach(() => { + teardownServer(); + }); + + it('returns LIFECYCLE when a stale handler invokes a disposed context', async () => { + const ctx = createTestServerContext({ + client: { query: vi.fn() } as never, + }); + const server = createMockServer(); + registerQueryTool(server as never, ctx); + const handler = server.getHandler('query')!; + + ctx.teardown(); + + const raw = await handler({ + query_text: 'contracts', + namespace: 'wg21', + preset: 'fast', + }); + const err = assertToolErrorCode(raw, 'LIFECYCLE'); + expect(err.message).toContain('disposed'); + }); + + it('fresh context after teardown produces a working handler', async () => { + const query = vi.fn().mockResolvedValue(makeHybridQueryResult()); + const ctx1 = createTestServerContext({ client: { query } as never }); + const server1 = createMockServer(); + registerQueryTool(server1 as never, ctx1); + ctx1.teardown(); + + const ctx2 = createTestServerContext({ + config: { disableSuggestFlow: true }, + client: { query } as never, + }); + const server2 = createMockServer(); + registerQueryTool(server2 as never, ctx2); + + const body = parseToolJson( + await server2.getHandler('query')!({ + query_text: 'contracts', + namespace: 'wg21', + preset: 'fast', + }) + ); + expect(body['status']).toBe('success'); + }); + + it('setup cycle: teardownServer invalidates prior context handlers; re-setup works', async () => { + const cfg = resolveConfig({ apiKey: 'lifecycle-cycle-key', indexName: 'test-index' }); + setPineconeClient( + new PineconeClient({ + apiKey: cfg.apiKey, + indexName: cfg.indexName, + rerankModel: cfg.rerankModel, + defaultTopK: cfg.defaultTopK, + }) + ); + + await setupCoreServer(cfg); + const staleCtx = getDefaultServerContext(); + const mockServer = createMockServer(); + registerQueryTool(mockServer as never, staleCtx); + const staleHandler = mockServer.getHandler('query')!; + + teardownServer(); + + const staleRaw = await staleHandler({ + query_text: 'contracts', + namespace: 'wg21', + preset: 'fast', + }); + assertToolErrorCode(staleRaw, 'LIFECYCLE'); + + await expect(setupCoreServer(cfg)).resolves.toBeDefined(); + + const freshCtx = createTestServerContext({ + config: resolveTestConfig({ disableSuggestFlow: true }), + client: { + query: vi.fn().mockResolvedValue(makeHybridQueryResult()), + } as never, + }); + const freshServer = createMockServer(); + registerQueryTool(freshServer as never, freshCtx); + const body = parseToolJson( + await freshServer.getHandler('query')!({ + query_text: 'contracts', + namespace: 'wg21', + preset: 'fast', + }) + ); + expect(body['status']).toBe('success'); + + teardownServer(); + }); +}); diff --git a/src/core/server/tools/list-namespaces-tool.ts b/src/core/server/tools/list-namespaces-tool.ts index 5bc094a..0d841d8 100644 --- a/src/core/server/tools/list-namespaces-tool.ts +++ b/src/core/server/tools/list-namespaces-tool.ts @@ -1,11 +1,14 @@ import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; import { getNamespacesWithCache } from '../namespaces-cache.js'; import type { ServerContext } from '../server-context.js'; -import { classifyToolCatchError, logToolError } from '../tool-error.js'; +import { classifyToolCatchError, lifecycleToolError, logToolError } from '../tool-error.js'; import { jsonErrorResponse, jsonResponse } from '../tool-response.js'; async function executeListNamespaces(ctx?: ServerContext) { try { + if (ctx?.disposed) { + return jsonErrorResponse(lifecycleToolError('ServerContext has been disposed')); + } const { data: namespacesInfo, cache_hit, diff --git a/src/core/server/tools/namespace-router-tool.ts b/src/core/server/tools/namespace-router-tool.ts index e602e35..092ea52 100644 --- a/src/core/server/tools/namespace-router-tool.ts +++ b/src/core/server/tools/namespace-router-tool.ts @@ -3,7 +3,12 @@ import { z } from 'zod'; import { getNamespacesWithCache } from '../namespaces-cache.js'; import { rankNamespacesByQuery } from '../namespace-router.js'; import type { ServerContext } from '../server-context.js'; -import { classifyToolCatchError, logToolError, validationToolError } from '../tool-error.js'; +import { + classifyToolCatchError, + lifecycleToolError, + logToolError, + validationToolError, +} from '../tool-error.js'; import { jsonErrorResponse, jsonResponse } from '../tool-response.js'; /** Register the namespace_router tool on the MCP server. */ @@ -29,6 +34,9 @@ export function registerNamespaceRouterTool(server: McpServer, ctx?: ServerConte }, async (params) => { try { + if (ctx?.disposed) { + return jsonErrorResponse(lifecycleToolError('ServerContext has been disposed')); + } const { user_query, top_n } = params; if (!user_query?.trim()) { return jsonErrorResponse(validationToolError('user_query cannot be empty', 'user_query')); diff --git a/src/core/server/tools/query-documents-tool.ts b/src/core/server/tools/query-documents-tool.ts index 7b9db14..c0e9c16 100644 --- a/src/core/server/tools/query-documents-tool.ts +++ b/src/core/server/tools/query-documents-tool.ts @@ -14,6 +14,7 @@ import { requireSuggested } from '../suggestion-flow.js'; import { classifyToolCatchError, flowGateToolError, + lifecycleToolError, logToolError, validationToolError, } from '../tool-error.js'; @@ -74,6 +75,9 @@ export function registerQueryDocumentsTool(server: McpServer, ctx?: ServerContex }, async (params) => { try { + if (ctx?.disposed) { + return jsonErrorResponse(lifecycleToolError('ServerContext has been disposed')); + } const { query_text, namespace, diff --git a/src/core/server/tools/query-tool.ts b/src/core/server/tools/query-tool.ts index 2020666..d3b2695 100644 --- a/src/core/server/tools/query-tool.ts +++ b/src/core/server/tools/query-tool.ts @@ -11,6 +11,7 @@ import { requireSuggested } from '../suggestion-flow.js'; import { classifyToolCatchError, flowGateToolError, + lifecycleToolError, logToolError, validationToolError, } from '../tool-error.js'; @@ -32,6 +33,9 @@ type QueryExecParams = { async function executeQuery(params: QueryExecParams, ctx?: ServerContext) { const { query_text, namespace, top_k, use_reranking, metadata_filter, fields, mode } = params; try { + if (ctx?.disposed) { + return jsonErrorResponse(lifecycleToolError('ServerContext has been disposed')); + } if (!query_text.trim()) { return jsonErrorResponse(validationToolError('Query text cannot be empty', 'query_text')); } diff --git a/src/core/setup.ts b/src/core/setup.ts index 5ed279a..47a5925 100644 --- a/src/core/setup.ts +++ b/src/core/setup.ts @@ -18,6 +18,9 @@ import { registerQueryTool } from './server/tools/query-tool.js'; let mcpServerInitialized = false; +/** MCP server handle with automatic teardown via `await using`. */ +export type ServerHandle = McpServer & AsyncDisposable; + /** * Reset process-global MCP server state (suggest-flow, namespace cache, active config, * Pinecone client handle, URL generator registry). Call before a second {@link setupCoreServer}. @@ -42,7 +45,7 @@ export type SetupCoreServerOptions = { export async function setupCoreServer( config?: ServerConfig, options?: SetupCoreServerOptions -): Promise { +): Promise { if (mcpServerInitialized) { throw new Error( 'setupCoreServer() 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.' @@ -83,6 +86,14 @@ export async function setupCoreServer( registerQueryDocumentsTool(server, ctx); registerGenerateUrlsTool(server, ctx); + const handle = server as ServerHandle; + handle[Symbol.asyncDispose] = async () => { + try { + await ctx[Symbol.asyncDispose](); + } finally { + mcpServerInitialized = false; + } + }; mcpServerInitialized = true; - return server; + return handle; } diff --git a/tsconfig.json b/tsconfig.json index feb39f1..eee93ee 100644 --- a/tsconfig.json +++ b/tsconfig.json @@ -1,6 +1,7 @@ { "compilerOptions": { "target": "ES2022", + "lib": ["ES2022", "ESNext.Disposable"], "module": "Node16", "moduleResolution": "Node16", "outDir": "./dist",