fix(mcplocal): rebind upstream auth token on every project refresh #133

Merged
michal merged 1 commits from fix/stale-upstream-token into main 2026-09-22 22:46:02 +00:00
6 changed files with 111 additions and 10 deletions

View File

@@ -167,12 +167,16 @@ function syncUpstreams(router: McpRouter, mcpdClient: McpdClient, servers: McpdS
}
}
// Add/update upstreams for each server
// Add/update upstreams for each server. Always rebuild, even when an
// upstream with this name already exists: McpdUpstream binds the caller's
// bearer token (and the server id) at construction, and a cached project
// router outlives both. Reusing the old object meant a token rotation left
// every pre-rotation server silently 401-ing on tools/list — attached,
// healthy, zero tools, no error — while servers attached after the rotation
// worked, because only THEY got an upstream carrying the fresh token.
for (const server of servers) {
if (!currentNames.has(server.name)) {
const upstream = new McpdUpstream(server.id, server.name, toolClient, server.description, discoveryClient);
router.addUpstream(upstream);
}
const upstream = new McpdUpstream(server.id, server.name, toolClient, server.description, discoveryClient);
router.addUpstream(upstream);
// Applied on every sync, not just on first registration: raising a
// server's timeout should take effect at the next refresh rather than
// requiring the upstream to be dropped and rebuilt.

View File

@@ -19,6 +19,7 @@ export interface PluginContextDeps {
cache: CacheProvider;
log: StageLogger;
discoverTools: () => Promise<ToolDefinition[]>;
getToolDiscoveryFailures?: () => Array<{ server: string; message: string }>;
routeToUpstream: (request: JsonRpcRequest) => Promise<JsonRpcResponse>;
fetchPromptIndex: () => Promise<PromptIndexEntry[]>;
getSystemPrompt: (name: string, fallback: string) => Promise<string>;
@@ -96,6 +97,10 @@ export class PluginContextImpl implements PluginSessionContext {
return this.deps.discoverTools();
}
getToolDiscoveryFailures(): Array<{ server: string; message: string }> {
return this.deps.getToolDiscoveryFailures?.() ?? [];
}
routeToUpstream(request: JsonRpcRequest): Promise<JsonRpcResponse> {
return this.deps.routeToUpstream(request);
}

View File

@@ -36,6 +36,8 @@ export interface PluginSessionContext {
// Upstream access
discoverTools(): Promise<ToolDefinition[]>;
/** Servers that contributed no tools on the last discoverTools run, with the error. */
getToolDiscoveryFailures(): Array<{ server: string; message: string }>;
routeToUpstream(request: JsonRpcRequest): Promise<JsonRpcResponse>;
// Prompt access

View File

@@ -127,6 +127,7 @@ export function createGatePlugin(config: GatePluginConfig = {}): ProxyModelPlugi
parts.push(` ${sanitizeWireName(t.name)}`);
}
}
parts.push(...formatDiscoveryFailures(ctx.getToolDiscoveryFailures()));
} catch {
// Tool discovery is optional
}
@@ -177,6 +178,22 @@ export function createGatePlugin(config: GatePluginConfig = {}): ProxyModelPlugi
};
}
/**
* A partial catalog must never look like a healthy one. When attached servers
* contributed no tools, say so in-band — the client has no other channel to
* learn that six servers silently vanished from its tool list.
*/
function formatDiscoveryFailures(failures: Array<{ server: string; message: string }>): string[] {
if (failures.length === 0) return [];
const parts: string[] = [];
parts.push(`\n⚠ WARNING: ${failures.length} attached server(s) contributed NO tools to this catalog:`);
for (const f of failures) {
parts.push(` - ${f.server}: ${f.message}`);
}
parts.push('The tool list above is INCOMPLETE. Report this to the user instead of treating a missing tool as nonexistent.');
return parts;
}
// ── begin_session tool definition ──
function getBeginSessionTool(llmSelector: LlmPromptSelector | null): ToolDefinition {
@@ -442,6 +459,7 @@ async function handleBeginSession(
responseParts.push(` ${sanitizeWireName(t.name)}`);
}
}
responseParts.push(...formatDiscoveryFailures(ctx.getToolDiscoveryFailures()));
} catch {
// Tool discovery is optional
}

View File

@@ -75,6 +75,8 @@ export class McpRouter {
// `tools/list` from re-hitting slow/dead upstreams and absorbs negative results so one
// dead server only stalls the first POST, not every subsequent one.
private listCache = new Map<string, ListCacheEntry>();
/** Per-server failures from the most recent discoverTools run. */
private toolDiscoveryFailures: Array<{ server: string; message: string }> = [];
private readonly LIST_CACHE_POSITIVE_TTL_MS = 30_000;
private readonly LIST_CACHE_NEGATIVE_TTL_MS = 30_000;
@@ -176,6 +178,7 @@ export class McpRouter {
error: (msg: string) => console.error(`[plugin:${sessionId.slice(0, 8)}] ${msg}`),
},
discoverTools: () => this.discoverTools(),
getToolDiscoveryFailures: () => this.getToolDiscoveryFailures(),
routeToUpstream: (request) => this.routeNamespacedCall(request, 'name', this.toolToServer),
fetchPromptIndex: () => this.fetchPromptIndex(),
getSystemPrompt: (name, fallback) => this.getSystemPrompt(name, fallback),
@@ -329,7 +332,7 @@ export class McpRouter {
const started = Date.now();
let cachedCount = 0;
let freshCount = 0;
const failed: string[] = [];
const failed: Array<{ server: string; message: string }> = [];
// Discover tools from all servers in parallel so one slow server doesn't block the rest
const entries = [...this.upstreams.entries()];
@@ -369,7 +372,7 @@ export class McpRouter {
const cached = result.value.cached;
if (cached.kind === 'err') {
cachedCount++;
failed.push(serverName);
failed.push({ server: serverName, message: cached.message });
continue;
}
response = { jsonrpc: '2.0', id: `cached-${serverName}`, result: cached.result };
@@ -381,7 +384,7 @@ export class McpRouter {
const message = (response.error as { message?: string }).message ?? 'unknown error';
this.setListCacheEntry(serverName, 'tools/list', { kind: 'err', message, fetchedAt: Date.now() });
console.warn(`[discoverTools] ${serverName}: ${message}`);
failed.push(serverName);
failed.push({ server: serverName, message });
continue;
}
if (response.result !== undefined) {
@@ -410,16 +413,27 @@ export class McpRouter {
}
}
this.toolDiscoveryFailures = failed;
if (entries.length > 0) {
const elapsed = Date.now() - started;
const project = this.projectName ? ` project=${this.projectName}` : '';
const failedStr = failed.length > 0 ? ` failed=[${failed.join(',')}]` : '';
const failedStr = failed.length > 0 ? ` failed=[${failed.map((f) => f.server).join(',')}]` : '';
console.info(`[discoverTools]${project} fresh=${freshCount} cached=${cachedCount}${failedStr} elapsed=${elapsed}ms`);
}
return allTools;
}
/**
* Servers that contributed no tools on the most recent `discoverTools` run,
* with the error that silenced them. A partial catalog must never look like
* a healthy one: the gate reads this to warn the client in-band.
*/
getToolDiscoveryFailures(): Array<{ server: string; message: string }> {
return this.toolDiscoveryFailures;
}
/**
* Discover resources from all upstreams by calling resources/list on each.
* Shares the per-server list cache with `discoverTools`.

View File

@@ -1,5 +1,5 @@
import { describe, it, expect, vi } from 'vitest';
import { refreshUpstreams } from '../src/discovery.js';
import { refreshUpstreams, refreshProjectUpstreams } from '../src/discovery.js';
import { McpRouter } from '../src/router.js';
function mockMcpdClient(servers: Array<{ id: string; name: string; transport: string }>) {
@@ -69,3 +69,61 @@ describe('refreshUpstreams', () => {
expect(router.getUpstreamNames()).toHaveLength(0);
});
});
describe('refreshProjectUpstreams', () => {
/**
* Regression: a cached project router outlives the caller's bearer token.
* Upstreams used to be created only for NEW server names, so after a token
* rotation every pre-existing upstream kept posting the old token — six
* healthy servers 401'd on tools/list forever while newly attached ones
* worked. Every refresh must rebind existing upstreams to the fresh token.
*/
it('rebinds existing upstreams to the current auth token on every refresh', async () => {
const router = new McpRouter();
const servers = [{ id: 'srv-1', name: 'gitea', transport: 'stdio' }];
// Base client: withToken(t) mints a child client whose posts we can count.
const children = new Map<string, ReturnType<typeof mockMcpdClient>>();
const base = mockMcpdClient(servers);
(base as any).forward = vi.fn(async () => ({ status: 200, body: servers }));
(base as any).withToken = vi.fn((t: string) => {
if (!children.has(t)) {
const child = mockMcpdClient(servers);
child.post = vi.fn(async () => ({ result: { tools: [{ name: 'merge_pull_request' }] } }));
children.set(t, child);
}
return children.get(t);
});
await refreshProjectUpstreams(router, base as any, 'sre', 'old-token');
await refreshProjectUpstreams(router, base as any, 'sre', 'new-token');
const tools = await router.discoverTools();
expect(tools.map((t) => t.name)).toEqual(['gitea/merge_pull_request']);
// The proxy call must ride the NEW token's client, not the stale one.
expect(children.get('new-token')!.post).toHaveBeenCalledWith('/api/v1/mcp/proxy', expect.anything());
expect(children.get('old-token')!.post).not.toHaveBeenCalled();
});
it('updates the server id when a server was recreated under the same name', async () => {
const router = new McpRouter();
const base = mockMcpdClient([]);
const posts: unknown[] = [];
const child = mockMcpdClient([]);
child.post = vi.fn(async (_path: string, body: unknown) => {
posts.push(body);
return { result: { tools: [] } };
});
(base as any).withToken = vi.fn(() => child);
(base as any).forward = vi.fn()
.mockResolvedValueOnce({ status: 200, body: [{ id: 'srv-old', name: 'gitea', transport: 'stdio' }] })
.mockResolvedValueOnce({ status: 200, body: [{ id: 'srv-new', name: 'gitea', transport: 'stdio' }] });
await refreshProjectUpstreams(router, base as any, 'sre', 't');
await refreshProjectUpstreams(router, base as any, 'sre', 't');
await router.discoverTools();
expect(posts).toHaveLength(1);
expect((posts[0] as { serverId: string }).serverId).toBe('srv-new');
});
});