diff --git a/src/mcplocal/src/gate/llm-selector.ts b/src/mcplocal/src/gate/llm-selector.ts index 628fb73..c857351 100644 --- a/src/mcplocal/src/gate/llm-selector.ts +++ b/src/mcplocal/src/gate/llm-selector.ts @@ -1,12 +1,16 @@ /** * LLM-based prompt selection for the gating flow. * - * Sends tags + prompt index to the heavy LLM, which returns - * a ranked list of relevant prompt names. + * Sends tags + prompt index to a heavy LLM, which returns a ranked list of + * relevant prompt names. Credential tiering (user rule): prefer mcpd's server + * `Llm` (cloud/server keys live at the k8s level) via the inference proxy, and + * fall back to the local (personal-token) provider registry only when the + * server path is unavailable. Cloud keys never live in mcplocal config. */ import type { ProviderRegistry } from '../providers/registry.js'; import type { SystemPromptFetcher } from '../proxymodel/types.js'; +import type { ChatMessage, CompletionOptions } from '../providers/types.js'; export interface PromptIndexForLlm { name: string; @@ -20,6 +24,23 @@ export interface LlmSelectionResult { reasoning: string; } +/** + * Route a selection inference through mcpd's server `Llm` (typically a + * `POST /api/v1/llms/:name/infer`). Returns the completion text. Must throw on + * transport/HTTP failure so the selector can fall back to the local provider. + */ +export type ServerInfer = ( + messages: ChatMessage[], + opts: { maxTokens: number; temperature: number; signal?: AbortSignal }, +) => Promise; + +export interface SelectPromptsOptions { + getSystemPromptFn?: SystemPromptFetcher; + signal?: AbortSignal; + /** Preferred inference path: mcpd server Llm (cloud/server keys at k8s). */ + serverInfer?: ServerInfer; +} + /** * Token budget for the selection call. A reasoning model spends most of its * budget inside reasoning_content before emitting the {selectedNames} JSON, so @@ -31,17 +52,49 @@ const SELECT_MAX_TOKENS = (() => { return Number.isFinite(v) && v > 0 ? v : 4000; })(); +/** + * Pull the assistant text from an OpenAI-shaped completion, falling back to + * reasoning_content — thinking models emit their answer there with content null. + */ +export function pickCompletionText(resp: unknown): string { + const m = (resp as { + choices?: Array<{ + message?: { + content?: string | null; + reasoning_content?: string | null; + provider_specific_fields?: { reasoning_content?: string | null }; + }; + }>; + }).choices?.[0]?.message; + const primary = m?.content ?? ''; + if (primary !== '') return primary; + return m?.reasoning_content ?? m?.provider_specific_fields?.reasoning_content ?? ''; +} + +/** Extract the `{ "selectedNames": [...], "reasoning": "..." }` object, or null. */ +export function extractSelection(text: string): LlmSelectionResult | null { + const jsonMatch = text.match(/\{[\s\S]*"selectedNames"[\s\S]*\}/); + if (!jsonMatch) return null; + try { + const parsed = JSON.parse(jsonMatch[0]) as { selectedNames?: string[]; reasoning?: string }; + return { selectedNames: parsed.selectedNames ?? [], reasoning: parsed.reasoning ?? '' }; + } catch { + return null; + } +} + export class LlmPromptSelector { constructor( - private readonly providerRegistry: ProviderRegistry, + private readonly providerRegistry: ProviderRegistry | null, ) {} async selectPrompts( tags: string[], promptIndex: PromptIndexForLlm[], - getSystemPromptFn?: SystemPromptFetcher, - signal?: AbortSignal, + opts: SelectPromptsOptions = {}, ): Promise { + const { getSystemPromptFn, signal, serverInfer } = opts; + const DEFAULT_SYSTEM_PROMPT = `You are a context selection assistant. Given a developer's task keywords and a list of available project prompts, select which prompts are relevant to their work. Return a JSON object with "selectedNames" (array of prompt names) and "reasoning" (brief explanation). Priority 10 prompts must always be included.`; const systemPrompt = getSystemPromptFn ? await getSystemPromptFn('llm-gate-context-selector', DEFAULT_SYSTEM_PROMPT) @@ -54,45 +107,62 @@ ${promptIndex.map((p) => `- ${p.name} (priority: ${p.priority}): ${p.summary ?? Select the relevant prompts. Return JSON: { "selectedNames": [...], "reasoning": "..." }`; - const provider = this.providerRegistry.getProvider('heavy'); - if (!provider) { - throw new Error('No heavy LLM provider available'); + const messages: ChatMessage[] = [ + { role: 'system', content: systemPrompt }, + { role: 'user', content: userPrompt }, + ]; + + // Sources in priority order: mcpd server Llm first (cloud/server keys at + // k8s), then the local personal-token provider as fallback. + const sources: Array<{ label: string; run: () => Promise }> = []; + if (serverInfer) { + sources.push({ + label: 'mcpd server Llm', + run: () => serverInfer(messages, { maxTokens: SELECT_MAX_TOKENS, temperature: 0, ...(signal ? { signal } : {}) }), + }); + } + if (this.providerRegistry) { + sources.push({ + label: 'local provider', + run: async () => { + const provider = this.providerRegistry!.getProvider('heavy'); + if (!provider) throw new Error('No heavy LLM provider available'); + const completionOptions: CompletionOptions = { + messages, + temperature: 0, + maxTokens: SELECT_MAX_TOKENS, + ...(signal ? { signal } : {}), + }; + return (await provider.complete(completionOptions)).content; + }, + }); + } + if (sources.length === 0) { + throw new Error('No LLM provider available for prompt selection'); } - // Deliberately NOT setting `model` here. Prompt-ranking only needs a - // working LLM; forcing the project's vLLM model onto the (anthropic) heavy - // provider made every selection fail silently. Use the provider's own model. - const completionOptions: import('../providers/types.js').CompletionOptions = { - messages: [ - { role: 'system', content: systemPrompt }, - { role: 'user', content: userPrompt }, - ], - temperature: 0, - maxTokens: SELECT_MAX_TOKENS, - ...(signal ? { signal } : {}), - }; - - const result = await provider.complete(completionOptions); - - const response = result.content; - - // Parse JSON from response (may be wrapped in markdown code blocks) - const jsonMatch = response.match(/\{[\s\S]*"selectedNames"[\s\S]*\}/); - if (!jsonMatch) { - throw new Error('LLM response did not contain valid selection JSON'); - } - - const parsed = JSON.parse(jsonMatch[0]) as { selectedNames?: string[]; reasoning?: string }; - const selectedNames = parsed.selectedNames ?? []; - const reasoning = parsed.reasoning ?? ''; - - // Always include priority 10 prompts - for (const p of promptIndex) { - if (p.priority === 10 && !selectedNames.includes(p.name)) { - selectedNames.push(p.name); + let lastErr: Error | null = null; + for (const src of sources) { + let text: string; + try { + text = await src.run(); + } catch (err) { + lastErr = err instanceof Error ? err : new Error(String(err)); + continue; // transport/HTTP failure → try the next source } + const sel = extractSelection(text); + if (!sel) { + lastErr = new Error(`${src.label}: LLM response did not contain valid selection JSON`); + continue; + } + // Always include priority 10 prompts. + for (const p of promptIndex) { + if (p.priority === 10 && !sel.selectedNames.includes(p.name)) { + sel.selectedNames.push(p.name); + } + } + return sel; } - - return { selectedNames, reasoning }; + throw lastErr ?? new Error('LLM prompt-selection failed'); } } diff --git a/src/mcplocal/src/http/project-mcp-endpoint.ts b/src/mcplocal/src/http/project-mcp-endpoint.ts index 94b8b91..4a38d37 100644 --- a/src/mcplocal/src/http/project-mcp-endpoint.ts +++ b/src/mcplocal/src/http/project-mcp-endpoint.ts @@ -147,6 +147,10 @@ export function registerProjectMcpEndpoint(app: FastifyInstance, mcpdClient: Mcp providerRegistry: effectiveRegistry, }; if (resolvedModel) pluginConfig.modelOverride = resolvedModel; + // Route gate prompt-selection through this project's server Llm (mcpd + // inference proxy) so cloud/server keys stay at the k8s level; the local + // personal-token provider is the fallback. See credential-tiering. + if (mcpdConfig.llmProvider) pluginConfig.llmProvider = mcpdConfig.llmProvider; const basePlugin = createDefaultPlugin(pluginConfig); // Optional favourite-index presentation: curated favourite/ + full // all// + a "prefer favourite/" instruction. Composed AFTER diff --git a/src/mcplocal/src/proxymodel/plugins/gate.ts b/src/mcplocal/src/proxymodel/plugins/gate.ts index 24f057d..78ef5d9 100644 --- a/src/mcplocal/src/proxymodel/plugins/gate.ts +++ b/src/mcplocal/src/proxymodel/plugins/gate.ts @@ -13,7 +13,7 @@ import type { ProxyModelPlugin, PluginSessionContext } from '../plugin.js'; import { SessionGate } from '../../gate/session-gate.js'; import { TagMatcher, extractKeywordsFromToolCall, tokenizeDescription } from '../../gate/tag-matcher.js'; import type { TagMatchResult } from '../../gate/tag-matcher.js'; -import { LlmPromptSelector } from '../../gate/llm-selector.js'; +import { LlmPromptSelector, pickCompletionText, type ServerInfer } from '../../gate/llm-selector.js'; import type { ProviderRegistry } from '../../providers/registry.js'; import { withTimeout, TimeoutError } from '../../util/with-timeout.js'; @@ -26,6 +26,13 @@ export interface GatePluginConfig { providerRegistry?: ProviderRegistry | null; modelOverride?: string; byteBudget?: number; + /** + * Name of the project's server `Llm` (mcpd). When set (and not 'none'), gate + * prompt-selection routes through mcpd's inference proxy first — keeping + * cloud/server keys at the k8s level — and only falls back to the local + * personal-token provider registry. See credential-tiering principle. + */ + llmProvider?: string; } const MAX_RESPONSE_CHARS = 24_000; @@ -33,8 +40,12 @@ const MAX_RESPONSE_CHARS = 24_000; export function createGatePlugin(config: GatePluginConfig = {}): ProxyModelPlugin { const isGated = config.gated !== false; const tagMatcher = new TagMatcher(config.byteBudget); - const llmSelector = config.providerRegistry - ? new LlmPromptSelector(config.providerRegistry) + // Prefer routing selection through mcpd's server Llm (cloud/server keys at + // k8s); the local provider registry (personal tokens) is the fallback. A gate + // with a server Llm but no local providers still selects (server-only). + const serverLlm = config.llmProvider && config.llmProvider !== 'none' ? config.llmProvider : undefined; + const llmSelector = (config.providerRegistry || serverLlm) + ? new LlmPromptSelector(config.providerRegistry ?? null) : null; // Per-session state tracking (plugin-scoped, not global SessionGate) @@ -49,7 +60,7 @@ export function createGatePlugin(config: GatePluginConfig = {}): ProxyModelPlugi // Register begin_session virtual tool ctx.registerTool(getBeginSessionTool(llmSelector), async (args, callCtx) => { - return handleBeginSession(args, callCtx, sessionGate, tagMatcher, llmSelector); + return handleBeginSession(args, callCtx, sessionGate, tagMatcher, llmSelector, serverLlm); }); // Register read_prompts virtual tool (available even when ungated) @@ -250,6 +261,7 @@ async function handleBeginSession( sessionGate: SessionGate, tagMatcher: TagMatcher, llmSelector: LlmPromptSelector | null, + serverLlm?: string, ): Promise { const rawTags = args['tags'] as string[] | undefined; const description = args['description'] as string | undefined; @@ -287,8 +299,25 @@ async function handleBeginSession( chapters: p.chapters, })); const getSystemPromptFn = ctx.getSystemPrompt.bind(ctx); + // Prefer mcpd's server Llm (keeps cloud/server keys at k8s); the selector + // falls back to the local personal-token provider on failure. + const serverInfer: ServerInfer | undefined = serverLlm + ? async (messages, o) => { + const resp = await ctx.postToMcpd(`/api/v1/llms/${encodeURIComponent(serverLlm)}/infer`, { + messages, + temperature: o.temperature, + max_tokens: o.maxTokens, + stream: false, + }); + return pickCompletionText(resp); + } + : undefined; const llmResult = await withTimeout( - (signal) => llmSelector.selectPrompts(tags, llmIndex, getSystemPromptFn, signal), + (signal) => llmSelector.selectPrompts(tags, llmIndex, { + getSystemPromptFn, + signal, + ...(serverInfer ? { serverInfer } : {}), + }), GATE_LLM_TIMEOUT_MS, 'gate LLM prompt-selection', ); diff --git a/src/mcplocal/tests/llm-selector.test.ts b/src/mcplocal/tests/llm-selector.test.ts index 91aa890..612aa2b 100644 --- a/src/mcplocal/tests/llm-selector.test.ts +++ b/src/mcplocal/tests/llm-selector.test.ts @@ -109,12 +109,49 @@ describe('LlmPromptSelector', () => { const selector = new LlmPromptSelector(registry); const ac = new AbortController(); - await selector.selectPrompts(['test'], sampleIndex, undefined, ac.signal); + await selector.selectPrompts(['test'], sampleIndex, { signal: ac.signal }); const call = (provider.complete as ReturnType).mock.calls[0]![0] as CompletionOptions; expect(call.signal).toBe(ac.signal); }); + it('prefers the mcpd server Llm (serverInfer) over the local provider', async () => { + const provider = makeMockProvider('{ "selectedNames": ["mqtt-config"], "reasoning": "local" }'); + const registry = makeRegistry(provider); + const selector = new LlmPromptSelector(registry); + const serverInfer = vi.fn().mockResolvedValue('{ "selectedNames": ["zigbee-pairing"], "reasoning": "server" }'); + + const result = await selector.selectPrompts(['x'], sampleIndex, { serverInfer }); + + expect(serverInfer).toHaveBeenCalledOnce(); + expect(provider.complete).not.toHaveBeenCalled(); // server succeeded → no local call + expect(result.selectedNames).toContain('zigbee-pairing'); + expect(result.reasoning).toBe('server'); + }); + + it('falls back to the local provider when serverInfer throws', async () => { + const provider = makeMockProvider('{ "selectedNames": ["mqtt-config"], "reasoning": "local" }'); + const registry = makeRegistry(provider); + const selector = new LlmPromptSelector(registry); + const serverInfer = vi.fn().mockRejectedValue(new Error('mcpd HTTP 503')); + + const result = await selector.selectPrompts(['x'], sampleIndex, { serverInfer }); + + expect(serverInfer).toHaveBeenCalledOnce(); + expect(provider.complete).toHaveBeenCalledOnce(); + expect(result.selectedNames).toContain('mqtt-config'); + expect(result.reasoning).toBe('local'); + }); + + it('works server-only (null registry) when a serverInfer is supplied', async () => { + const selector = new LlmPromptSelector(null); + const serverInfer = vi.fn().mockResolvedValue('{ "selectedNames": ["mqtt-config"], "reasoning": "s" }'); + + const result = await selector.selectPrompts(['x'], sampleIndex, { serverInfer }); + + expect(result.selectedNames).toContain('mqtt-config'); + }); + it('throws when no heavy provider is available', async () => { const registry = new ProviderRegistry(); // Empty registry const selector = new LlmPromptSelector(registry);