feat(gate): route prompt-selection through mcpd's server Llm (credential tiering)
Some checks failed
CI/CD / typecheck (pull_request) Successful in 1m3s
CI/CD / lint (pull_request) Successful in 2m6s
CI/CD / test (pull_request) Successful in 1m18s
CI/CD / build (pull_request) Successful in 2m10s
CI/CD / smoke (pull_request) Failing after 3m9s
CI/CD / publish (pull_request) Has been skipped

Cloud/server keys belong at the k8s/mcpd level; mcplocal handles only the user's
personal tokens. The gate's prompt-selection used mcplocal's LOCAL provider
registry (heavy=anthropic = the user's OAuth token, which is API-gated and 404s),
so selection always degraded.

Now LlmPromptSelector tries sources in priority order: (1) the project's server
Llm via mcpd's inference proxy (POST /api/v1/llms/:name/infer) — cloud/server
keys stay at k8s — then (2) the local personal-token provider as fallback. First
source that returns valid selection JSON wins. Extracted pickCompletionText
(content ?? reasoning_content) + extractSelection helpers.

Wiring: GatePluginConfig.llmProvider (from the project), threaded via
createDefaultPlugin; handleBeginSession builds serverInfer from ctx.postToMcpd.
Verified live: POST .../llms/vllm-current/infer returns valid selection JSON
(qwen3-thinking, 4000-token budget). + selector tests (server preferred, local
fallback on error, server-only). mcplocal suite green (755).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Michal
2026-07-24 12:43:48 +01:00
parent e75d2ba296
commit 2729580974
4 changed files with 187 additions and 47 deletions

View File

@@ -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<string>;
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<LlmSelectionResult> {
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<string> }> = [];
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');
}
}

View File

@@ -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/<tool> + full
// all/<server>/<tool> + a "prefer favourite/" instruction. Composed AFTER

View File

@@ -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<unknown> {
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',
);

View File

@@ -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<typeof vi.fn>).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);