feat: McpToken — HTTP-mode mcplocal, CLI verbs, audit plumbing #50

Merged
michal merged 12 commits from feat/mcptoken into main 2026-04-18 16:37:53 +00:00
7 changed files with 191 additions and 14 deletions
Showing only changes of commit 39df459bb1 - Show all commits

View File

@@ -41,6 +41,11 @@ export class AuditCollector {
this.sessionPrincipals.set(sessionId, { ...existing, tokenName: token.tokenName, tokenSha: token.tokenSha }); this.sessionPrincipals.set(sessionId, { ...existing, tokenName: token.tokenName, tokenSha: token.tokenSha });
} }
/** Look up the McpToken SHA for a session. Returns undefined for non-HTTP-mode sessions. */
getSessionMcpTokenSha(sessionId: string): string | undefined {
return this.sessionPrincipals.get(sessionId)?.tokenSha;
}
/** Queue an audit event. Auto-fills projectName, userName, tokenName, and tokenSha. */ /** Queue an audit event. Auto-fills projectName, userName, tokenName, and tokenSha. */
emit(event: Omit<AuditEvent, 'projectName'>): void { emit(event: Omit<AuditEvent, 'projectName'>): void {
const enriched: AuditEvent = { ...event, projectName: this.projectName }; const enriched: AuditEvent = { ...event, projectName: this.projectName };

View File

@@ -3,6 +3,21 @@
* *
* Tracks whether a session has gone through the prompt selection flow. * Tracks whether a session has gone through the prompt selection flow.
* When gated, only begin_session is accessible. After ungating, all tools work. * When gated, only begin_session is accessible. After ungating, all tools work.
*
* Per-token ungate cache:
* When the caller authenticated via an `McpToken` (HTTP-mode service agent),
* we also remember the ungate keyed on the token's SHA. Subsequent sessions
* from the same token automatically start ungated for a TTL window.
*
* Why: LiteLLM and similar MCP-proxying clients don't preserve the
* `mcp-session-id` header across chat completion calls, so every tool call
* lands on a fresh upstream session — which would otherwise be gated anew,
* forcing the agent into a begin_session loop. Keying on the token (which IS
* preserved, because it's in the Authorization header) gives us a stable
* identity that survives stateless proxies.
*
* Claude Code's stdio path keeps its session-id, so this code is a no-op for
* that case (session-id ungate still applies, token ungate is purely additive).
*/ */
import type { PromptIndexEntry, TagMatchResult } from './tag-matcher.js'; import type { PromptIndexEntry, TagMatchResult } from './tag-matcher.js';
@@ -14,15 +29,37 @@ export interface SessionState {
briefing: string | null; briefing: string | null;
} }
interface TokenUngateEntry {
tokenSha: string;
tags: string[];
ungatedAt: number;
retrievedPrompts: Set<string>;
}
/** Default TTL for per-token ungate cache (1 hour). Tunable via env for testing. */
const DEFAULT_TOKEN_UNGATE_TTL_MS = Number(process.env['MCPLOCAL_TOKEN_UNGATE_TTL_MS']) || 60 * 60 * 1000;
export class SessionGate { export class SessionGate {
private sessions = new Map<string, SessionState>(); private sessions = new Map<string, SessionState>();
private tokenUngates = new Map<string, TokenUngateEntry>();
private readonly tokenUngateTtlMs: number;
/** Create a new session. Starts gated if the project is gated. */ constructor(tokenUngateTtlMs = DEFAULT_TOKEN_UNGATE_TTL_MS) {
createSession(sessionId: string, projectGated: boolean): void { this.tokenUngateTtlMs = tokenUngateTtlMs;
}
/**
* Create a new session. Starts gated if the project is gated, UNLESS the
* caller's McpToken already ungated within the last TTL window — in which
* case the session inherits the previous tags + retrievedPrompts so the
* agent doesn't get the full gated greeting on every fresh session.
*/
createSession(sessionId: string, projectGated: boolean, tokenSha?: string): void {
const priorEntry = tokenSha ? this.getActiveTokenEntry(tokenSha) : null;
this.sessions.set(sessionId, { this.sessions.set(sessionId, {
gated: projectGated, gated: projectGated && priorEntry === null,
tags: [], tags: priorEntry ? [...priorEntry.tags] : [],
retrievedPrompts: new Set(), retrievedPrompts: priorEntry ? new Set(priorEntry.retrievedPrompts) : new Set(),
briefing: null, briefing: null,
}); });
} }
@@ -37,18 +74,37 @@ export class SessionGate {
return this.sessions.get(sessionId)?.gated ?? false; return this.sessions.get(sessionId)?.gated ?? false;
} }
/** Ungate a session after prompt selection is complete. */ /** True when a token has an active (non-expired) ungate entry. */
ungate(sessionId: string, tags: string[], matchResult: TagMatchResult): void { isTokenUngated(tokenSha: string): boolean {
return this.getActiveTokenEntry(tokenSha) !== null;
}
/**
* Ungate a session after prompt selection is complete.
*
* When `tokenSha` is supplied, also remember the ungate keyed on the token
* so future sessions from the same token start ungated (survives proxies
* that drop `mcp-session-id`).
*/
ungate(sessionId: string, tags: string[], matchResult: TagMatchResult, tokenSha?: string): void {
const session = this.sessions.get(sessionId); const session = this.sessions.get(sessionId);
if (!session) return; if (!session) return;
session.gated = false; session.gated = false;
session.tags = [...session.tags, ...tags]; session.tags = [...session.tags, ...tags];
// Track which prompts have been sent
for (const p of matchResult.fullContent) { for (const p of matchResult.fullContent) {
session.retrievedPrompts.add(p.name); session.retrievedPrompts.add(p.name);
} }
if (tokenSha !== undefined && tokenSha !== '') {
this.tokenUngates.set(tokenSha, {
tokenSha,
tags: [...session.tags],
ungatedAt: Date.now(),
retrievedPrompts: new Set(session.retrievedPrompts),
});
}
} }
/** Record additional prompts retrieved via read_prompts. */ /** Record additional prompts retrieved via read_prompts. */
@@ -73,4 +129,19 @@ export class SessionGate {
removeSession(sessionId: string): void { removeSession(sessionId: string): void {
this.sessions.delete(sessionId); this.sessions.delete(sessionId);
} }
/** Forget a token's ungate entry (e.g. on revocation signal). */
revokeToken(tokenSha: string): void {
this.tokenUngates.delete(tokenSha);
}
private getActiveTokenEntry(tokenSha: string): TokenUngateEntry | null {
const entry = this.tokenUngates.get(tokenSha);
if (!entry) return null;
if (Date.now() - entry.ungatedAt > this.tokenUngateTtlMs) {
this.tokenUngates.delete(tokenSha);
return null;
}
return entry;
}
} }

View File

@@ -25,6 +25,13 @@ export interface PluginContextDeps {
queueNotification: (notification: JsonRpcNotification) => void; queueNotification: (notification: JsonRpcNotification) => void;
postToMcpd: (path: string, body: Record<string, unknown>) => Promise<unknown>; postToMcpd: (path: string, body: Record<string, unknown>) => Promise<unknown>;
auditCollector?: AuditCollector; auditCollector?: AuditCollector;
/**
* Resolves the principal's McpToken SHA for this session, if the caller
* authenticated via an McpToken. Called lazily so the value reflects the
* session's current state even when the token is attached after the plugin
* context is created.
*/
getMcpTokenSha?: () => string | undefined;
} }
/** /**
@@ -55,6 +62,11 @@ export class PluginContextImpl implements PluginSessionContext {
this.deps = deps; this.deps = deps;
} }
/** McpToken SHA for the current caller, or undefined for STDIO/session-auth callers. */
getMcpTokenSha(): string | undefined {
return this.deps.getMcpTokenSha?.();
}
registerTool(tool: ToolDefinition, handler: VirtualToolHandler): void { registerTool(tool: ToolDefinition, handler: VirtualToolHandler): void {
this.virtualTools.set(tool.name, { definition: tool, handler }); this.virtualTools.set(tool.name, { definition: tool, handler });
} }

View File

@@ -50,6 +50,14 @@ export interface PluginSessionContext {
// Audit event emission (auto-fills sessionId and projectName) // Audit event emission (auto-fills sessionId and projectName)
emitAuditEvent(event: Omit<AuditEvent, 'sessionId' | 'projectName'>): void; emitAuditEvent(event: Omit<AuditEvent, 'sessionId' | 'projectName'>): void;
/**
* McpToken SHA for the current caller, or undefined if the session was
* authenticated via a User session (STDIO/Claude Code path). Plugins can use
* this to key state on the token principal rather than the session-id —
* useful when the session-id doesn't survive a proxy (e.g. LiteLLM).
*/
getMcpTokenSha(): string | undefined;
} }
// ── Virtual Server ────────────────────────────────────────────────── // ── Virtual Server ──────────────────────────────────────────────────

View File

@@ -40,7 +40,11 @@ export function createGatePlugin(config: GatePluginConfig = {}): ProxyModelPlugi
description: 'Gated session flow: begin_session → prompt selection → ungate.', description: 'Gated session flow: begin_session → prompt selection → ungate.',
async onSessionCreate(ctx) { async onSessionCreate(ctx) {
sessionGate.createSession(ctx.sessionId, isGated); // Pass the caller's McpToken SHA so the gate can honor a cross-session
// ungate cache keyed on the token principal. Fixes the LiteLLM case where
// each tool call lands on a fresh mcp-session-id → would otherwise loop
// on begin_session forever.
sessionGate.createSession(ctx.sessionId, isGated, ctx.getMcpTokenSha());
// Register begin_session virtual tool // Register begin_session virtual tool
ctx.registerTool(getBeginSessionTool(llmSelector), async (args, callCtx) => { ctx.registerTool(getBeginSessionTool(llmSelector), async (args, callCtx) => {
@@ -264,8 +268,9 @@ async function handleBeginSession(
matchResult = tagMatcher.match(tags, promptIndex); matchResult = tagMatcher.match(tags, promptIndex);
} }
// Ungate the session // Ungate the session (and remember the ungate per McpToken if this is a
sessionGate.ungate(ctx.sessionId, tags, matchResult); // service-token request, so the next session from the same token skips the gate).
sessionGate.ungate(ctx.sessionId, tags, matchResult, ctx.getMcpTokenSha());
ctx.queueNotification('notifications/tools/list_changed'); ctx.queueNotification('notifications/tools/list_changed');
// Audit: gate_decision for begin_session // Audit: gate_decision for begin_session
@@ -451,8 +456,8 @@ async function handleGatedIntercept(
const promptIndex = await ctx.fetchPromptIndex(); const promptIndex = await ctx.fetchPromptIndex();
const matchResult = tagMatcher.match(tags, promptIndex); const matchResult = tagMatcher.match(tags, promptIndex);
// Ungate the session // Ungate the session (and remember per-token if the caller is a McpToken).
sessionGate.ungate(ctx.sessionId, tags, matchResult); sessionGate.ungate(ctx.sessionId, tags, matchResult, ctx.getMcpTokenSha());
ctx.queueNotification('notifications/tools/list_changed'); ctx.queueNotification('notifications/tools/list_changed');
// Audit: gate_decision for auto-intercept // Audit: gate_decision for auto-intercept
@@ -522,7 +527,7 @@ async function handleGatedIntercept(
return response; return response;
} catch { } catch {
// If prompt retrieval fails, just ungate and route normally // If prompt retrieval fails, just ungate and route normally
sessionGate.ungate(ctx.sessionId, tags, { fullContent: [], indexOnly: [], remaining: [] }); sessionGate.ungate(ctx.sessionId, tags, { fullContent: [], indexOnly: [], remaining: [] }, ctx.getMcpTokenSha());
ctx.queueNotification('notifications/tools/list_changed'); ctx.queueNotification('notifications/tools/list_changed');
return ctx.routeToUpstream(request); return ctx.routeToUpstream(request);
} }

View File

@@ -198,6 +198,10 @@ export class McpRouter {
return this.mcpdClient.post(path, body); return this.mcpdClient.post(path, body);
}, },
...(this.auditCollector ? { auditCollector: this.auditCollector } : {}), ...(this.auditCollector ? { auditCollector: this.auditCollector } : {}),
// Lazily resolve the caller's McpToken SHA via the audit collector's
// session principal map. The token is attached in onsessioninitialized,
// which runs before any plugin context is created, so this is stable.
getMcpTokenSha: () => this.auditCollector?.getSessionMcpTokenSha(sessionId),
}; };
ctx = new PluginContextImpl(deps); ctx = new PluginContextImpl(deps);

View File

@@ -152,4 +152,76 @@ describe('SessionGate', () => {
expect(gate.isGated('s1')).toBe(false); expect(gate.isGated('s1')).toBe(false);
expect(gate.getSession('s2')!.tags).toEqual([]); // s2 untouched expect(gate.getSession('s2')!.tags).toEqual([]); // s2 untouched
}); });
describe('per-McpToken ungate cache', () => {
it('new session from an already-ungated token starts ungated, with prior tags + prompts', () => {
const gate = new SessionGate();
gate.createSession('session-1', true, 'tokA');
expect(gate.isGated('session-1')).toBe(true);
gate.ungate('session-1', ['ops'], makeMatchResult(['runbook']), 'tokA');
expect(gate.isTokenUngated('tokA')).toBe(true);
// LiteLLM semantics: same token, brand-new session-id.
gate.createSession('session-2', true, 'tokA');
expect(gate.isGated('session-2')).toBe(false);
const s2 = gate.getSession('session-2')!;
expect(s2.tags).toContain('ops');
expect(s2.retrievedPrompts.has('runbook')).toBe(true);
});
it('does not persist across tokens', () => {
const gate = new SessionGate();
gate.createSession('s1', true, 'tokA');
gate.ungate('s1', ['ops'], makeMatchResult(['p']), 'tokA');
// Different token → fresh gated session.
gate.createSession('s2', true, 'tokB');
expect(gate.isGated('s2')).toBe(true);
expect(gate.isTokenUngated('tokB')).toBe(false);
});
it('is not triggered when no tokenSha is supplied (STDIO path)', () => {
const gate = new SessionGate();
gate.createSession('s1', true);
gate.ungate('s1', ['ops'], makeMatchResult(['p']));
// A second session with no token starts gated — STDIO semantics preserved.
gate.createSession('s2', true);
expect(gate.isGated('s2')).toBe(true);
});
it('honors the TTL window and expires', () => {
const gate = new SessionGate(50); // 50ms TTL for the test
gate.createSession('s1', true, 'tokA');
gate.ungate('s1', ['ops'], makeMatchResult(['p']), 'tokA');
expect(gate.isTokenUngated('tokA')).toBe(true);
return new Promise<void>((resolve) => setTimeout(() => {
expect(gate.isTokenUngated('tokA')).toBe(false);
gate.createSession('s2', true, 'tokA');
expect(gate.isGated('s2')).toBe(true);
resolve();
}, 70));
});
it('revokeToken clears the ungate entry immediately', () => {
const gate = new SessionGate();
gate.createSession('s1', true, 'tokA');
gate.ungate('s1', ['ops'], makeMatchResult(['p']), 'tokA');
expect(gate.isTokenUngated('tokA')).toBe(true);
gate.revokeToken('tokA');
expect(gate.isTokenUngated('tokA')).toBe(false);
gate.createSession('s2', true, 'tokA');
expect(gate.isGated('s2')).toBe(true);
});
it('empty-string tokenSha does not register an ungate entry', () => {
const gate = new SessionGate();
gate.createSession('s1', true, '');
gate.ungate('s1', ['ops'], makeMatchResult(['p']), '');
expect(gate.isTokenUngated('')).toBe(false);
});
});
}); });