From fb8f8271279b1774e830b4ea5053f7db3c42fd07 Mon Sep 17 00:00:00 2001 From: Elliott de Launay Date: Fri, 6 Mar 2026 14:17:18 +0000 Subject: [PATCH] feat(SecretStorageService): service for mcp hub to store access and refresh tokens --- src/services/mcp/McpServerManager.ts | 4 + src/services/mcp/SecretStorageService.ts | 44 ++++++++ .../__tests__/SecretStorageService.spec.ts | 102 ++++++++++++++++++ 3 files changed, 150 insertions(+) create mode 100644 src/services/mcp/SecretStorageService.ts create mode 100644 src/services/mcp/__tests__/SecretStorageService.spec.ts diff --git a/src/services/mcp/McpServerManager.ts b/src/services/mcp/McpServerManager.ts index 3fd7146d9f..cc456d7e89 100644 --- a/src/services/mcp/McpServerManager.ts +++ b/src/services/mcp/McpServerManager.ts @@ -1,5 +1,6 @@ import * as vscode from "vscode" import { McpHub } from "./McpHub" +import { SecretStorageService } from "./SecretStorageService" import { ClineProvider } from "../../core/webview/ClineProvider" /** @@ -37,6 +38,9 @@ export class McpServerManager { // Double-check instance in case it was created while we were waiting if (!this.instance) { const hub = new McpHub(provider) + // Set the secret storage service for OAuth operations + const secretStorage = new SecretStorageService(context) + hub.setSecretStorage(secretStorage) // Wait for all MCP servers to finish connecting (or timing out) await hub.waitUntilReady() this.instance = hub diff --git a/src/services/mcp/SecretStorageService.ts b/src/services/mcp/SecretStorageService.ts new file mode 100644 index 0000000000..d4d1684e54 --- /dev/null +++ b/src/services/mcp/SecretStorageService.ts @@ -0,0 +1,44 @@ +import * as vscode from "vscode" +import type { OAuthTokens } from "@modelcontextprotocol/sdk/shared/auth.js" + +export interface StoredMcpOAuthData { + tokens: OAuthTokens + /** Unix ms timestamp after which the access token should be considered expired. */ + expires_at: number +} + +/** + * Thin wrapper around VS Code SecretStorage for persisting MCP OAuth tokens. + * Tokens are stored per-server (keyed by host) so different servers on the + * same host share credentials, which is the common case for multi-path APIs. + */ +export class SecretStorageService { + private readonly _storage: vscode.SecretStorage + private readonly _namespace = "mcp.oauth." + + constructor(context: vscode.ExtensionContext) { + this._storage = context.secrets + } + + private _key(serverUrl: string): string { + return `${this._namespace}${new URL(serverUrl).host}.data` + } + + async getOAuthData(serverUrl: string): Promise { + const raw = await this._storage.get(this._key(serverUrl)) + if (!raw) return undefined + try { + return JSON.parse(raw) as StoredMcpOAuthData + } catch { + return undefined + } + } + + async saveOAuthData(serverUrl: string, data: StoredMcpOAuthData): Promise { + await this._storage.store(this._key(serverUrl), JSON.stringify(data)) + } + + async deleteOAuthData(serverUrl: string): Promise { + await this._storage.delete(this._key(serverUrl)) + } +} diff --git a/src/services/mcp/__tests__/SecretStorageService.spec.ts b/src/services/mcp/__tests__/SecretStorageService.spec.ts new file mode 100644 index 0000000000..12709347d0 --- /dev/null +++ b/src/services/mcp/__tests__/SecretStorageService.spec.ts @@ -0,0 +1,102 @@ +import { describe, it, expect, vi, beforeEach } from "vitest" + +vi.mock("vscode", () => ({})) + +import { SecretStorageService, StoredMcpOAuthData } from "../SecretStorageService" + +function createMockContext() { + const store = new Map() + return { + secrets: { + get: vi.fn(async (key: string) => store.get(key)), + store: vi.fn(async (key: string, value: string) => { + store.set(key, value) + }), + delete: vi.fn(async (key: string) => { + store.delete(key) + }), + }, + } as any +} + +describe("SecretStorageService", () => { + let service: SecretStorageService + let context: ReturnType + + beforeEach(() => { + context = createMockContext() + service = new SecretStorageService(context) + }) + + describe("getOAuthData", () => { + it("should return undefined when no data stored", async () => { + const result = await service.getOAuthData("https://example.com/mcp") + expect(result).toBeUndefined() + }) + + it("should return stored data", async () => { + const data: StoredMcpOAuthData = { + tokens: { access_token: "tok", token_type: "Bearer" }, + expires_at: Date.now() + 3600_000, + } + await service.saveOAuthData("https://example.com/mcp", data) + + const result = await service.getOAuthData("https://example.com/mcp") + expect(result).toEqual(data) + }) + + it("should return undefined for malformed JSON", async () => { + // Manually store garbage via the underlying mock + context.secrets.store("mcp.oauth.example.com.data", "not-json") + + const result = await service.getOAuthData("https://example.com/mcp") + expect(result).toBeUndefined() + }) + }) + + describe("saveOAuthData", () => { + it("should persist data under host-based key", async () => { + const data: StoredMcpOAuthData = { + tokens: { access_token: "abc", token_type: "Bearer" }, + expires_at: 12345, + } + await service.saveOAuthData("https://example.com/mcp", data) + + expect(context.secrets.store).toHaveBeenCalledWith("mcp.oauth.example.com.data", JSON.stringify(data)) + }) + }) + + describe("deleteOAuthData", () => { + it("should delete stored data", async () => { + const data: StoredMcpOAuthData = { + tokens: { access_token: "tok", token_type: "Bearer" }, + expires_at: Date.now() + 3600_000, + } + await service.saveOAuthData("https://example.com/mcp", data) + + await service.deleteOAuthData("https://example.com/mcp") + + expect(context.secrets.delete).toHaveBeenCalledWith("mcp.oauth.example.com.data") + const result = await service.getOAuthData("https://example.com/mcp") + expect(result).toBeUndefined() + }) + }) + + describe("key isolation", () => { + it("should isolate data by host", async () => { + const data1: StoredMcpOAuthData = { + tokens: { access_token: "a", token_type: "Bearer" }, + expires_at: 1, + } + const data2: StoredMcpOAuthData = { + tokens: { access_token: "b", token_type: "Bearer" }, + expires_at: 2, + } + await service.saveOAuthData("https://host1.com/mcp", data1) + await service.saveOAuthData("https://host2.com/mcp", data2) + + expect((await service.getOAuthData("https://host1.com/mcp"))?.tokens.access_token).toBe("a") + expect((await service.getOAuthData("https://host2.com/mcp"))?.tokens.access_token).toBe("b") + }) + }) +})