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

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 2 additions & 3 deletions src/alliance/setup.ts
Original file line number Diff line number Diff line change
@@ -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';
Expand All @@ -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<McpServer> {
export async function setupAllianceServer(config?: ServerConfig): Promise<ServerHandle> {
const server = await setupCoreServer(config ?? resolveAllianceConfig({}), {
instructions: ALLIANCE_SERVER_INSTRUCTIONS,
});
Expand Down
4 changes: 4 additions & 0 deletions src/alliance/tools/guided-query-tool.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down
4 changes: 4 additions & 0 deletions src/alliance/tools/suggest-query-params-tool.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand Down Expand Up @@ -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'));
Expand Down
2 changes: 1 addition & 1 deletion src/core/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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';
55 changes: 55 additions & 0 deletions src/core/server/server-context.lifecycle.test.ts
Original file line number Diff line number Diff line change
@@ -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();
});
});
Comment thread
coderabbitai[bot] marked this conversation as resolved.
12 changes: 11 additions & 1 deletion src/core/server/server-context.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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<string, UrlGeneratorFn>();
Expand Down Expand Up @@ -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<void> {
this.teardown();
if (defaultContext === this) {
defaultContext = null;
pendingConfig = null;
}
}
Comment thread
coderabbitai[bot] marked this conversation as resolved.
}

let defaultContext: ServerContext | null = null;
Expand Down
9 changes: 9 additions & 0 deletions src/core/server/tool-error.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ import { describe, expect, it } from 'vitest';
import {
classifyToolCatchError,
flowGateToolError,
lifecycleToolError,
toolErrorSchema,
validationToolError,
} from './tool-error.js';
Expand Down Expand Up @@ -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'),
Expand Down
24 changes: 23 additions & 1 deletion src/core/server/tool-error.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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<typeof toolErrorCodeSchema>;

const flowGateToolErrorSchema = z.object({
Expand Down Expand Up @@ -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<typeof toolErrorSchema>;
Expand Down Expand Up @@ -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);
}
Expand Down
4 changes: 4 additions & 0 deletions src/core/server/tools/count-tool.ts
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ import { requireSuggested } from '../suggestion-flow.js';
import {
classifyToolCatchError,
flowGateToolError,
lifecycleToolError,
logToolError,
validationToolError,
} from '../tool-error.js';
Expand All @@ -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) {
Expand Down
10 changes: 9 additions & 1 deletion src/core/server/tools/generate-urls-tool.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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). */
Expand Down Expand Up @@ -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) {
Expand Down
10 changes: 9 additions & 1 deletion src/core/server/tools/keyword-search-tool.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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). */
Expand Down Expand Up @@ -47,6 +52,9 @@ async function executeKeywordSearch(
},
ctx?: ServerContext
): Promise<KeywordSearchExecResult> {
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();
Expand Down
Loading
Loading