feat: ServerContext lifecycle — AsyncDisposable + stale-handler guards (#148)

* feat: ServerContext lifecycle — AsyncDisposable + stale-handler guards

* fixed format error

* addressed ai reviews

---------

Co-authored-by: zho <jornathanm910923@gmail.com>

Author: jonathanMLDev

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<McpServer> {
+export async function setupAllianceServer(config?: ServerConfig): Promise<ServerHandle> {
   const server = await setupCoreServer(config ?? resolveAllianceConfig({}), {
     instructions: ALLIANCE_SERVER_INSTRUCTIONS,
   });

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,

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'));

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';

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();
+  });
+});

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<string, UrlGeneratorFn>();
@@ -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;
+    }
+  }
 }
 
 let defaultContext: ServerContext | null = null;

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'),

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<typeof toolErrorCodeSchema>;
 
 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<typeof toolErrorSchema>;
@@ -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);
 }

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) {

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) {