diff --git a/src/mcplocal/src/audit/collector.ts b/src/mcplocal/src/audit/collector.ts index 0cd73c3..a3b8aa8 100644 --- a/src/mcplocal/src/audit/collector.ts +++ b/src/mcplocal/src/audit/collector.ts @@ -41,11 +41,6 @@ export class AuditCollector { 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. */ emit(event: Omit): void { const enriched: AuditEvent = { ...event, projectName: this.projectName }; diff --git a/src/mcplocal/src/gate/session-gate.ts b/src/mcplocal/src/gate/session-gate.ts index 137516f..920451a 100644 --- a/src/mcplocal/src/gate/session-gate.ts +++ b/src/mcplocal/src/gate/session-gate.ts @@ -3,21 +3,6 @@ * * Tracks whether a session has gone through the prompt selection flow. * 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'; @@ -29,37 +14,15 @@ export interface SessionState { briefing: string | null; } -interface TokenUngateEntry { - tokenSha: string; - tags: string[]; - ungatedAt: number; - retrievedPrompts: Set; -} - -/** 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 { private sessions = new Map(); - private tokenUngates = new Map(); - private readonly tokenUngateTtlMs: number; - constructor(tokenUngateTtlMs = DEFAULT_TOKEN_UNGATE_TTL_MS) { - 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; + /** Create a new session. Starts gated if the project is gated. */ + createSession(sessionId: string, projectGated: boolean): void { this.sessions.set(sessionId, { - gated: projectGated && priorEntry === null, - tags: priorEntry ? [...priorEntry.tags] : [], - retrievedPrompts: priorEntry ? new Set(priorEntry.retrievedPrompts) : new Set(), + gated: projectGated, + tags: [], + retrievedPrompts: new Set(), briefing: null, }); } @@ -74,37 +37,18 @@ export class SessionGate { return this.sessions.get(sessionId)?.gated ?? false; } - /** True when a token has an active (non-expired) ungate entry. */ - 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 { + /** Ungate a session after prompt selection is complete. */ + ungate(sessionId: string, tags: string[], matchResult: TagMatchResult): void { const session = this.sessions.get(sessionId); if (!session) return; session.gated = false; session.tags = [...session.tags, ...tags]; + // Track which prompts have been sent for (const p of matchResult.fullContent) { 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. */ @@ -129,19 +73,4 @@ export class SessionGate { removeSession(sessionId: string): void { 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; - } } diff --git a/src/mcplocal/src/proxymodel/plugin-context.ts b/src/mcplocal/src/proxymodel/plugin-context.ts index 21fbb6e..39a6fd7 100644 --- a/src/mcplocal/src/proxymodel/plugin-context.ts +++ b/src/mcplocal/src/proxymodel/plugin-context.ts @@ -25,13 +25,6 @@ export interface PluginContextDeps { queueNotification: (notification: JsonRpcNotification) => void; postToMcpd: (path: string, body: Record) => Promise; 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; } /** @@ -62,11 +55,6 @@ export class PluginContextImpl implements PluginSessionContext { 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 { this.virtualTools.set(tool.name, { definition: tool, handler }); } diff --git a/src/mcplocal/src/proxymodel/plugin.ts b/src/mcplocal/src/proxymodel/plugin.ts index b2d2f39..9c2fc6b 100644 --- a/src/mcplocal/src/proxymodel/plugin.ts +++ b/src/mcplocal/src/proxymodel/plugin.ts @@ -50,14 +50,6 @@ export interface PluginSessionContext { // Audit event emission (auto-fills sessionId and projectName) emitAuditEvent(event: Omit): 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 ────────────────────────────────────────────────── diff --git a/src/mcplocal/src/proxymodel/plugins/gate.ts b/src/mcplocal/src/proxymodel/plugins/gate.ts index bc51268..41e1af6 100644 --- a/src/mcplocal/src/proxymodel/plugins/gate.ts +++ b/src/mcplocal/src/proxymodel/plugins/gate.ts @@ -40,11 +40,7 @@ export function createGatePlugin(config: GatePluginConfig = {}): ProxyModelPlugi description: 'Gated session flow: begin_session → prompt selection → ungate.', async onSessionCreate(ctx) { - // 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()); + sessionGate.createSession(ctx.sessionId, isGated); // Register begin_session virtual tool ctx.registerTool(getBeginSessionTool(llmSelector), async (args, callCtx) => { @@ -268,9 +264,8 @@ async function handleBeginSession( matchResult = tagMatcher.match(tags, promptIndex); } - // Ungate the session (and remember the ungate per McpToken if this is a - // service-token request, so the next session from the same token skips the gate). - sessionGate.ungate(ctx.sessionId, tags, matchResult, ctx.getMcpTokenSha()); + // Ungate the session + sessionGate.ungate(ctx.sessionId, tags, matchResult); ctx.queueNotification('notifications/tools/list_changed'); // Audit: gate_decision for begin_session @@ -456,8 +451,8 @@ async function handleGatedIntercept( const promptIndex = await ctx.fetchPromptIndex(); const matchResult = tagMatcher.match(tags, promptIndex); - // Ungate the session (and remember per-token if the caller is a McpToken). - sessionGate.ungate(ctx.sessionId, tags, matchResult, ctx.getMcpTokenSha()); + // Ungate the session + sessionGate.ungate(ctx.sessionId, tags, matchResult); ctx.queueNotification('notifications/tools/list_changed'); // Audit: gate_decision for auto-intercept @@ -527,7 +522,7 @@ async function handleGatedIntercept( return response; } catch { // If prompt retrieval fails, just ungate and route normally - sessionGate.ungate(ctx.sessionId, tags, { fullContent: [], indexOnly: [], remaining: [] }, ctx.getMcpTokenSha()); + sessionGate.ungate(ctx.sessionId, tags, { fullContent: [], indexOnly: [], remaining: [] }); ctx.queueNotification('notifications/tools/list_changed'); return ctx.routeToUpstream(request); } diff --git a/src/mcplocal/src/router.ts b/src/mcplocal/src/router.ts index 6e8f365..e5c80d0 100644 --- a/src/mcplocal/src/router.ts +++ b/src/mcplocal/src/router.ts @@ -198,10 +198,6 @@ export class McpRouter { return this.mcpdClient.post(path, body); }, ...(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); diff --git a/src/mcplocal/tests/session-gate.test.ts b/src/mcplocal/tests/session-gate.test.ts index b604899..a086df1 100644 --- a/src/mcplocal/tests/session-gate.test.ts +++ b/src/mcplocal/tests/session-gate.test.ts @@ -152,76 +152,4 @@ describe('SessionGate', () => { expect(gate.isGated('s1')).toBe(false); 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((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); - }); - }); });