update 9 tools (cppalliance#132)

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

Author: jonathanMLDev

src/alliance/setup.ts

@@ -1,6 +1,7 @@
 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 { registerBuiltinUrlGenerators } from './url-builtins.js';
@@ -17,8 +18,9 @@ export async function setupAllianceServer(config?: ServerConfig): Promise<McpSer
   const server = await setupCoreServer(config ?? resolveAllianceConfig({}), {
     instructions: ALLIANCE_SERVER_INSTRUCTIONS,
   });
+  const ctx = getDefaultServerContext();
   registerBuiltinUrlGenerators();
-  registerSuggestQueryParamsTool(server);
-  registerGuidedQueryTool(server);
+  registerSuggestQueryParamsTool(server, ctx);
+  registerGuidedQueryTool(server, ctx);
   return server;
 }

src/alliance/tools/guided-query-tool.context.test.ts

@@ -0,0 +1,56 @@
+import { describe, expect, it, vi } from 'vitest';
+import { registerGuidedQueryTool } from './guided-query-tool.js';
+import {
+  createMockServer,
+  createTestServerContext,
+  makeHybridQueryResult,
+  parseToolJson,
+} from '../../core/server/tools/test-helpers.js';
+
+describe('guided_query tool handler (ServerContext instance path)', () => {
+  it('returns success with decision_trace using injected context', async () => {
+    const listNamespacesWithMetadata = vi.fn().mockResolvedValue([
+      {
+        namespace: 'papers',
+        recordCount: 42,
+        metadata: {
+          document_number: 'string',
+          title: 'string',
+          url: 'string',
+          author: 'string',
+          chunk_text: 'string',
+        },
+      },
+    ]);
+    const query = vi.fn().mockResolvedValue(makeHybridQueryResult());
+    const ctx = createTestServerContext({
+      client: {
+        listNamespacesWithMetadata,
+        query,
+        count: vi.fn().mockResolvedValue({ count: 7, truncated: false }),
+      } as never,
+    });
+
+    const server = createMockServer();
+    registerGuidedQueryTool(server as never, ctx);
+    const raw = 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 body = parseToolJson(raw);
+    expect(body).toMatchObject({
+      status: 'success',
+    });
+    const trace = body['decision_trace'] as Record<string, unknown>;
+    expect(trace).toMatchObject({
+      cache_hit: false,
+      selected_namespace: 'papers',
+      enrich_urls: false,
+    });
+    expect(trace['rerank_status']).toBeDefined();
+    expect(query).toHaveBeenCalledOnce();
+  });
+});

src/alliance/tools/guided-query-tool.ts

@@ -12,6 +12,7 @@ import {
 import { rankNamespacesByQuery } from '../../core/server/namespace-router.js';
 import { getNamespacesWithCache } from '../../core/server/namespaces-cache.js';
 import { normalizeNamespace } from '../../core/server/namespace-utils.js';
+import type { ServerContext } from '../../core/server/server-context.js';
 import { suggestQueryParams } from '../../core/server/query-suggestion.js';
 import { markSuggested } from '../../core/server/suggestion-flow.js';
 import {
@@ -39,7 +40,7 @@ function resolveGuidedToolName(
  * Registers `guided_query` (routing + suggestion + execution in one call).
  * See "Retrieval tool decision matrix" in README.md for tool-selection guidance.
  */
-export function registerGuidedQueryTool(server: McpServer): void {
+export function registerGuidedQueryTool(server: McpServer, ctx?: ServerContext): void {
   server.registerTool(
     'guided_query',
     {
@@ -102,7 +103,9 @@ export function registerGuidedQueryTool(server: McpServer): void {
         }
 
         const queryText = user_query.trim();
-        const { data: namespaces, cache_hit } = await getNamespacesWithCache();
+        const { data: namespaces, cache_hit } = ctx
+          ? await ctx.getNamespacesWithCache()
+          : await getNamespacesWithCache();
         const ranked = rankNamespacesByQuery(queryText, namespaces, 3);
 
         let namespace: string | null = null;
@@ -150,13 +153,21 @@ export function registerGuidedQueryTool(server: McpServer): void {
         }
 
         const selectedTool: GuidedToolName = resolveGuidedToolName(preferred_tool, suggestion);
-        markSuggested(namespace, {
-          recommended_tool: selectedTool,
-          suggested_fields: suggestion.suggested_fields,
-          user_query: queryText,
-        });
+        if (ctx) {
+          ctx.markSuggested(namespace, {
+            recommended_tool: selectedTool,
+            suggested_fields: suggestion.suggested_fields,
+            user_query: queryText,
+          });
+        } else {
+          markSuggested(namespace, {
+            recommended_tool: selectedTool,
+            suggested_fields: suggestion.suggested_fields,
+            user_query: queryText,
+          });
+        }
 
-        const client = getPineconeClient();
+        const client = ctx ? ctx.getClient() : getPineconeClient();
         const baseTrace = {
           cache_hit,
           input_namespace: inputNamespace ?? null,

src/alliance/tools/suggest-query-params-tool.context.test.ts

@@ -0,0 +1,44 @@
+import { describe, expect, it, vi } from 'vitest';
+import { registerSuggestQueryParamsTool } from './suggest-query-params-tool.js';
+import {
+  createMockServer,
+  createTestServerContext,
+  parseToolJson,
+} from '../../core/server/tools/test-helpers.js';
+
+describe('suggest_query_params tool handler (ServerContext instance path)', () => {
+  it('marks suggest-flow on injected context when namespace exists', async () => {
+    const listNamespacesWithMetadata = vi.fn().mockResolvedValue([
+      {
+        namespace: 'wg21',
+        recordCount: 42,
+        metadata: {
+          document_number: 'string',
+          title: 'string',
+          url: 'string',
+          author: 'string',
+          chunk_text: 'string',
+        },
+      },
+    ]);
+    const ctx = createTestServerContext({
+      client: { listNamespacesWithMetadata } as never,
+    });
+
+    const server = createMockServer();
+    registerSuggestQueryParamsTool(server as never, ctx);
+    const raw = await server.getHandler('suggest_query_params')!({
+      namespace: 'wg21',
+      user_query: 'List papers with titles',
+    });
+    const body = parseToolJson(raw);
+    expect(body).toMatchObject({
+      status: 'success',
+      namespace_found: true,
+      cache_hit: false,
+    });
+
+    const flowCheck = ctx.requireSuggested('wg21');
+    expect(flowCheck.ok).toBe(true);
+  });
+});

src/alliance/tools/suggest-query-params-tool.ts

@@ -3,6 +3,7 @@ import { z } from 'zod';
 import { normalizeNamespace } from '../../core/server/namespace-utils.js';
 import { getNamespacesWithCache } from '../../core/server/namespaces-cache.js';
 import { suggestQueryParams } from '../../core/server/query-suggestion.js';
+import type { ServerContext } from '../../core/server/server-context.js';
 import { markSuggested } from '../../core/server/suggestion-flow.js';
 import {
   classifyToolCatchError,
@@ -12,7 +13,7 @@ import {
 import { jsonErrorResponse, jsonResponse } from '../../core/server/tool-response.js';
 
 /** Register the suggest_query_params tool on the MCP server. */
-export function registerSuggestQueryParamsTool(server: McpServer): void {
+export function registerSuggestQueryParamsTool(server: McpServer, ctx?: ServerContext): void {
   server.registerTool(
     'suggest_query_params',
     {
@@ -48,18 +49,28 @@ export function registerSuggestQueryParamsTool(server: McpServer): void {
             })
           );
         }
-        const { data: namespacesInfo, cache_hit } = await getNamespacesWithCache();
+        const { data: namespacesInfo, cache_hit } = ctx
+          ? await ctx.getNamespacesWithCache()
+          : await getNamespacesWithCache();
         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(nsNorm, {
-            recommended_tool: result.recommended_tool,
-            suggested_fields: result.suggested_fields,
-            user_query: user_query.trim(),
-          });
+          if (ctx) {
+            ctx.markSuggested(nsNorm, {
+              recommended_tool: result.recommended_tool,
+              suggested_fields: result.suggested_fields,
+              user_query: user_query.trim(),
+            });
+          } else {
+            markSuggested(nsNorm, {
+              recommended_tool: result.recommended_tool,
+              suggested_fields: result.suggested_fields,
+              user_query: user_query.trim(),
+            });
+          }
         }
         const response = {
           ...result,

src/core/server/tools/generate-urls-tool.context.test.ts

@@ -0,0 +1,31 @@
+import { describe, expect, it } from 'vitest';
+import { registerGenerateUrlsTool } from './generate-urls-tool.js';
+import { createMockServer, createTestServerContext, parseToolJson } from './test-helpers.js';
+
+describe('generate_urls tool handler (ServerContext instance path)', () => {
+  it('uses URL generator registered on injected context', async () => {
+    const ctx = createTestServerContext();
+    ctx.registerUrlGenerator('mailing', () => ({
+      url: 'https://example.com/doc/P1234',
+      method: 'generator',
+    }));
+
+    const server = createMockServer();
+    registerGenerateUrlsTool(server as never, ctx);
+    const raw = await server.getHandler('generate_urls')!({
+      namespace: 'mailing',
+      records: [{ document_number: 'P1234' }],
+    });
+    const body = parseToolJson(raw);
+    expect(body).toMatchObject({
+      status: 'success',
+      namespace: 'mailing',
+      count: 1,
+    });
+    const results = body['results'] as Array<{ url: string; method: string }>;
+    expect(results[0]).toMatchObject({
+      url: 'https://example.com/doc/P1234',
+      method: 'generator',
+    });
+  });
+});

src/core/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-registry.js';
 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 { jsonErrorResponse, jsonResponse } from '../tool-response.js';
 
@@ -15,7 +16,7 @@ function extractMetadata(record: Record<string, unknown>): Record<string, unknow
 }
 
 /** Register the generate_urls tool on the MCP server. */
-export function registerGenerateUrlsTool(server: McpServer): void {
+export function registerGenerateUrlsTool(server: McpServer, ctx?: ServerContext): void {
   server.registerTool(
     'generate_urls',
     {
@@ -49,7 +50,9 @@ export function registerGenerateUrlsTool(server: McpServer): void {
         }
         const results = records.map((record, index) => {
           const metadata = extractMetadata(record);
-          const generated = generateUrlForNamespace(nsNorm, metadata);
+          const generated = ctx
+            ? ctx.generateUrlForNamespace(nsNorm, metadata)
+            : generateUrlForNamespace(nsNorm, metadata);
           return {
             index,
             url: generated.url,

src/core/server/tools/keyword-search-tool.context.test.ts

@@ -0,0 +1,37 @@
+import { describe, expect, it, vi } from 'vitest';
+import { registerKeywordSearchTool } from './keyword-search-tool.js';
+import {
+  createMockServer,
+  createTestServerContext,
+  makeSearchResult,
+  parseToolJson,
+} from './test-helpers.js';
+
+describe('keyword_search tool handler (ServerContext instance path)', () => {
+  it('returns success using injected client', async () => {
+    const keywordSearch = vi.fn().mockResolvedValue([makeSearchResult()]);
+    const ctx = createTestServerContext({
+      client: {
+        keywordSearch,
+        getSparseIndexName: () => 'test-index-sparse',
+      } as never,
+    });
+
+    const server = createMockServer();
+    registerKeywordSearchTool(server as never, ctx);
+    const raw = await server.getHandler('keyword_search')!({
+      query_text: 'contracts',
+      namespace: 'wg21',
+      top_k: 5,
+    });
+    const body = parseToolJson(raw);
+    expect(body).toMatchObject({
+      status: 'success',
+      query: 'contracts',
+      namespace: 'wg21',
+      index: 'test-index-sparse',
+      result_count: 1,
+    });
+    expect(keywordSearch).toHaveBeenCalledOnce();
+  });
+});