feat(proxy): favourite-index tool presentation #86

Merged
michal merged 4 commits from feat/favourite-index-presentation into main 2026-07-23 21:28:23 +00:00
3 changed files with 95 additions and 44 deletions
Showing only changes of commit c9364a2f2a - Show all commits

View File

@@ -52,6 +52,64 @@ function splitServer(canonical: string): { server: string; short: string } {
: { server: canonical.slice(0, i), short: canonical.slice(i + 1) };
}
interface Presentation {
tools: ToolDefinition[];
/** presented name → canonical `server/tool` */
resolver: Map<string, string>;
}
/**
* Reshape an upstream catalog into favourites-first + full all/ catalog, and
* the presented→canonical routing map. Pure — shared by onToolsList and the
* lazy rebuild in onToolCallBefore.
*/
function buildPresentation(
upstreamTools: ToolDefinition[],
favourites: string[],
cap?: number,
): Presentation {
const byName = new Map(upstreamTools.map((t) => [t.name, t]));
const resolver = new Map<string, string>();
const usedPresented = new Set<string>();
const seenCanonical = new Set<string>();
// Favourites first: configured order, only those present, capped, with
// short-name collision disambiguation.
const limit = cap ?? favourites.length;
const favTools: ToolDefinition[] = [];
for (const canonical of favourites) {
if (favTools.length >= limit) break;
if (seenCanonical.has(canonical)) continue;
const t = byName.get(canonical);
if (!t) continue; // stale/missing pin — skip (all/ still exposes it)
const { server, short } = splitServer(canonical);
let presented = `favourite/${short}`;
if (usedPresented.has(presented)) presented = `favourite/${server}-${short}`;
if (usedPresented.has(presented)) continue; // still collides (rare) — skip
seenCanonical.add(canonical);
usedPresented.add(presented);
resolver.set(presented, canonical);
favTools.push({
name: presented,
description: `${humanize(short)} — common (${server})`,
...(t.inputSchema !== undefined ? { inputSchema: t.inputSchema } : {}),
});
}
// Then the full catalog under all/<server>/<tool>.
const allTools: ToolDefinition[] = upstreamTools.map((t) => {
const presented = `all/${t.name}`;
resolver.set(presented, t.name);
return {
name: presented,
...(t.description !== undefined ? { description: t.description } : {}),
...(t.inputSchema !== undefined ? { inputSchema: t.inputSchema } : {}),
};
});
return { tools: [...favTools, ...allTools], resolver };
}
export function createFavouriteIndexPlugin(config: FavouriteIndexConfig): ProxyModelPlugin {
const favourites = config.favourites ?? [];
const cap = config.maxFavourites && config.maxFavourites > 0 ? config.maxFavourites : undefined;
@@ -79,54 +137,33 @@ export function createFavouriteIndexPlugin(config: FavouriteIndexConfig): ProxyM
return tools;
}
const byName = new Map(upstreamTools.map((t) => [t.name, t]));
const resolver = new Map<string, string>(); // presented → canonical
const usedPresented = new Set<string>();
const seenCanonical = new Set<string>();
// Favourites first: configured order, only those present, capped, with
// short-name collision disambiguation.
const limit = cap ?? favourites.length;
const favTools: ToolDefinition[] = [];
for (const canonical of favourites) {
if (favTools.length >= limit) break;
if (seenCanonical.has(canonical)) continue;
const t = byName.get(canonical);
if (!t) continue; // stale/missing pin — skip (all/ still exposes it)
const { server, short } = splitServer(canonical);
let presented = `favourite/${short}`;
if (usedPresented.has(presented)) presented = `favourite/${server}-${short}`;
if (usedPresented.has(presented)) continue; // still collides (rare) — skip
seenCanonical.add(canonical);
usedPresented.add(presented);
resolver.set(presented, canonical);
favTools.push({
name: presented,
description: `${humanize(short)} — common (${server})`,
...(t.inputSchema !== undefined ? { inputSchema: t.inputSchema } : {}),
});
}
// Then the full catalog under all/<server>/<tool>.
const allTools: ToolDefinition[] = upstreamTools.map((t) => {
const presented = `all/${t.name}`;
resolver.set(presented, t.name);
return {
name: presented,
...(t.description !== undefined ? { description: t.description } : {}),
...(t.inputSchema !== undefined ? { inputSchema: t.inputSchema } : {}),
};
});
const { tools: presented, resolver } = buildPresentation(upstreamTools, favourites, cap);
ctx.state.set(RESOLVER_KEY, resolver);
return [...favTools, ...allTools, ...passthrough];
return [...presented, ...passthrough];
},
async onToolCallBefore(toolName, _args, request, ctx) {
const resolver = ctx.state.get(RESOLVER_KEY) as Map<string, string> | undefined;
const canonical =
resolver?.get(toolName) ??
(toolName.startsWith('all/') ? toolName.slice('all/'.length) : undefined);
if (!toolName.startsWith('favourite/') && !toolName.startsWith('all/')) return null;
let resolver = ctx.state.get(RESOLVER_KEY) as Map<string, string> | undefined;
let canonical = resolver?.get(toolName);
// Resolver missing (a call arrived before tools/list this session): rebuild
// it so favourite/ names — which, unlike all/, can't be prefix-stripped —
// still route.
if (canonical === undefined) {
const upstream = await ctx.discoverTools();
if (upstream.length > 0) {
resolver = buildPresentation(upstream, favourites, cap).resolver;
ctx.state.set(RESOLVER_KEY, resolver);
canonical = resolver.get(toolName);
}
// all/ is always derivable by stripping the prefix.
if (canonical === undefined && toolName.startsWith('all/')) {
canonical = toolName.slice('all/'.length);
}
}
if (canonical && canonical !== toolName) {
// Rewrite presented name → canonical `server/tool`, then continue the
// chain (return null) so normal routing + content-pipeline still run.

View File

@@ -134,6 +134,19 @@ describe('favourite-index plugin', () => {
expect(text).toBe('Called scale_deployment');
});
it('routes a favourite/ call even before tools/list (lazy resolver rebuild)', async () => {
const router = setup({ favourites: ['vault/read_secret'] });
router.addUpstream(mockUpstream('vault', CATALOG.vault));
// initialize only — deliberately NO tools/list, so the resolver is unset.
await router.route({ jsonrpc: '2.0', id: 1, method: 'initialize' }, { sessionId: 's1' });
const res = await router.route(
{ jsonrpc: '2.0', id: 2, method: 'tools/call', params: { name: 'favourite/read_secret', arguments: {} } },
{ sessionId: 's1' },
);
expect((res.result as { content: Array<{ text: string }> }).content[0]!.text).toBe('Called read_secret');
});
it('no-ops while the session is gated (only begin_session visible)', async () => {
const router = setup({ gated: true, favourites: ['k8s/get_pods'] });
router.addUpstream(mockUpstream('k8s', CATALOG.k8s));

View File

@@ -104,6 +104,7 @@ describe('Smoke: favourite-index presentation', () => {
chat.section('favourite-index routing');
try {
await chat.initialize();
await chat.listTools(); // real clients list before calling — populates the resolver
// Calling with no args → the UPSTREAM tool's arg validation fires, proving
// the presented name routed to the real server (not a -32601 unknown tool).