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
6 changes: 4 additions & 2 deletions src/alliance/setup.ts
Original file line number Diff line number Diff line change
@@ -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';
Expand All @@ -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;
}
56 changes: 56 additions & 0 deletions src/alliance/tools/guided-query-tool.context.test.ts
Original file line number Diff line number Diff line change
@@ -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();
});
});
27 changes: 19 additions & 8 deletions src/alliance/tools/guided-query-tool.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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',
{
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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,
Expand Down
44 changes: 44 additions & 0 deletions src/alliance/tools/suggest-query-params-tool.context.test.ts
Original file line number Diff line number Diff line change
@@ -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);
});
});
25 changes: 18 additions & 7 deletions src/alliance/tools/suggest-query-params-tool.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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',
{
Expand Down Expand Up @@ -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,
Expand Down
31 changes: 31 additions & 0 deletions src/core/server/tools/generate-urls-tool.context.test.ts
Original file line number Diff line number Diff line change
@@ -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',
});
});
});
9 changes: 6 additions & 3 deletions src/core/server/tools/generate-urls-tool.ts
Original file line number Diff line number Diff line change
@@ -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';

Expand All @@ -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',
{
Expand Down Expand Up @@ -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,
Expand Down
37 changes: 37 additions & 0 deletions src/core/server/tools/keyword-search-tool.context.test.ts
Original file line number Diff line number Diff line change
@@ -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();
});
});
Loading
Loading