scriptoriumadmin/Roo-Code
1
0
Fork 0
mirror of https://github.com/RooVetGit/Roo-Code.git synced 2026-09-07 08:26:51 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

feat(McpHub): update OAuth Flow to be opt in and handle concurrent flows across instances

Browse source
This commit is contained in:
Elliott de Launay 2026-03-14 00:26:23 -04:00
parent 653b71ea11
commit 3d0eb2775e
No known key found for this signature in database
GPG key ID: BB899BED766D1806
5 changed files with 403 additions and 59 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

194
src/services/mcp/McpHub.ts
View file

@ -168,6 +168,7 @@ export class McpHub {
private initializationPromise: Promise<void>
private secretStorage?: SecretStorageService
private reauthPromises: Map<string, Promise<void>> = new Map()
private _oauthWatchers: Map<string, { unsubscribe: () => void; abortHandle: NodeJS.Timeout }> = new Map()
constructor(provider: ClineProvider) {
this.providerRef = new WeakRef(provider)
@ -825,6 +826,13 @@ export class McpHub {
const connection = this.findConnection(name, source)
if (connection && connection.type === "connected") {
if (error instanceof UnauthorizedError && authProvider) {
// If we're already in the OAuth / polling flow, ignore transport
// retries that surface another 401 — the poll or _completeOAuthFlow
// will reconnect from scratch once tokens arrive.
if (connection.server.status === "connecting") {
return
}
// Mid-session re-auth triggered by a tool call (401)
connection.server.status = "connecting"
@ -861,6 +869,12 @@ export class McpHub {
transport.onclose = async () => {
const connection = this.findConnection(name, source)
if (connection) {
// If OAuth is in progress, don't overwrite "connecting" with "disconnected".
// The transport will close/retry while we await the browser flow or poll;
// the reconnect path (deleteConnection + connectToServer) handles cleanup.
if (connection.server.status === "connecting") {
return
}
connection.server.status = "disconnected"
}
await this.notifyWebviewOfServerChanges()
@ -943,21 +957,69 @@ export class McpHub {
try {
await client.connect(transport)
} catch (connectError) {
if (connectError instanceof UnauthorizedError && streamableHttpAuthProvider) {
// The server requires OAuth. The SDK has already called
// authProvider.redirectToAuthorization() which started the local callback
// server (lazily) and opened the user's browser.
//
// We fire-and-forget the rest of the flow so the extension (chat window,
// other servers) is not blocked waiting for the user's browser session.
if (connectError instanceof UnauthorizedError && streamableHttpAuthProvider && configInjected.url) {
// The server requires OAuth. Mark this connection as "connecting" and
// detach the toast + browser flow from the initialization path so that
// waitUntilReady() resolves immediately and the MCP panel can load.
const serverUrl = configInjected.url
connection.server.status = "connecting"
void this._completeOAuthFlow(
streamableHttpAuthProvider,
transport as StreamableHTTPClientTransport,
connection,
name,
source,
)
void (async () => {
const TOKEN_EXPIRY_BUFFER_MS = 5 * 60 * 1000
// Check if another window already saved valid tokens
const existing = await this.secretStorage!.getOAuthData(serverUrl)
if (existing && Date.now() < existing.expires_at - TOKEN_EXPIRY_BUFFER_MS) {
await streamableHttpAuthProvider.close()
await this.deleteConnection(name, source)
await this.connectToServer(name, config, source)
await this.notifyWebviewOfServerChanges()
return
}
// Show a confirmation toast so the user can decide whether to authenticate.
// This resolves immediately when the user responds — but connectToServer
// has already returned so the panel is not blocked.
const choice = await vscode.window.showInformationMessage(
`MCP server "${name}" requires authentication.`,
"Authenticate",
)
if (choice === "Authenticate") {
// Check tokens again — another window may have authed while toast was showing
const tokens = await this.secretStorage!.getOAuthData(serverUrl)
if (tokens && Date.now() < tokens.expires_at - TOKEN_EXPIRY_BUFFER_MS) {
await streamableHttpAuthProvider.close()
await this.deleteConnection(name, source)
await this.connectToServer(name, config, source)
await this.notifyWebviewOfServerChanges()
return
}
void this._completeOAuthFlow(
streamableHttpAuthProvider,
transport as StreamableHTTPClientTransport,
connection,
name,
source,
)
} else {
// Toast was dismissed or auto-timed-out.
// First do an immediate check — another window may have already
// completed auth while the toast was showing.
const tokens = await this.secretStorage!.getOAuthData(serverUrl)
if (tokens && Date.now() < tokens.expires_at - TOKEN_EXPIRY_BUFFER_MS) {
await streamableHttpAuthProvider.close()
await this.deleteConnection(name, source)
await this.connectToServer(name, config, source)
await this.notifyWebviewOfServerChanges()
return
}
// Tokens not available yet — start watching for them
await streamableHttpAuthProvider.close()
this._watchForOAuthTokens(name, source, serverUrl, config)
}
})()
return
}
// Non-OAuth error — let the outer catch handle it.
@ -1010,19 +1072,24 @@ export class McpHub {
name: string,
source: "global" | "project",
): Promise<void> {
const config = JSON.parse(connection.server.config)
try {
// Open the browser now that the user has confirmed the toast.
// redirectToAuthorization() was already called by the SDK (which stored
// the URL in _pendingAuthorizationUrl), but deliberately did not open it.
await authProvider.openBrowser()
const code = await authProvider.waitForAuthCode()
// Exchange auth code for tokens using the pre-fetched token_endpoint
// directly. The SDK's transport.finishAuth() re-runs discovery internally
// and hits the same broken URL for path-prefixed issuers (see
// utils/oauth.ts for upstream issue links).
await authProvider.exchangeCodeForTokens(code)
authProvider.close().catch(console.error)
await authProvider.close()
// Recover the validated server config stored on the connection so we
// can pass it directly to connectToServer without re-reading the file.
const parsedConfig = JSON.parse(connection.server.config)
const validatedConfig = this.validateServerConfig(parsedConfig, name)
const validatedConfig = this.validateServerConfig(config, name)
// Remove the broken connection (closes the old transport/client),
// then reconnect. The new McpOAuthClientProvider will find the token
@ -1045,6 +1112,82 @@ export class McpHub {
}
}
private _watchForOAuthTokens(
name: string,
source: "global" | "project",
serverUrl: string,
config: z.infer<typeof ServerConfigSchema>,
): void {
if (!this.secretStorage) return
const watcherKey = `${name}:${source}`
// Cancel any existing watcher for this connection before starting a new one
const existing = this._oauthWatchers.get(watcherKey)
if (existing) {
existing.unsubscribe()
clearTimeout(existing.abortHandle)
this._oauthWatchers.delete(watcherKey)
}
const TOKEN_EXPIRY_BUFFER_MS = 5 * 60 * 1000
// Called when SecretStorage fires onDidChange for this server's key.
// Runs in all VS Code windows the instant tokens are saved — no polling delay.
const onTokensChanged = async () => {
try {
if (this.isDisposed) return
const conn = this.findConnection(name, source)
if (!conn || conn.server.status === "connected") {
cleanup()
return
}
const data = await this.secretStorage?.getOAuthData(serverUrl)
if (data && Date.now() < data.expires_at - TOKEN_EXPIRY_BUFFER_MS) {
cleanup()
await this.deleteConnection(name, source)
const validatedConfig = this.validateServerConfig(config, name)
await this.connectToServer(name, validatedConfig, source)
await this.notifyWebviewOfServerChanges()
}
// If tokens aren't valid yet (e.g. a delete event fired), keep listening.
} catch (err) {
console.error(`[McpHub] OAuth token watcher failed for "${name}":`, err)
}
}
const cleanup = () => {
const entry = this._oauthWatchers.get(watcherKey)
if (entry) {
entry.unsubscribe()
clearTimeout(entry.abortHandle)
this._oauthWatchers.delete(watcherKey)
}
}
const unsubscribe = this.secretStorage.onDidChange(serverUrl, () => {
void onTokensChanged()
})
// Give up after 6 minutes if no token ever arrives
const abortHandle = setTimeout(
() => {
cleanup()
const conn = this.findConnection(name, source)
if (conn && conn.server.status === "connecting") {
conn.server.status = "disconnected"
this.appendErrorMessage(conn, "OAuth authentication timed out waiting for another window")
void this.notifyWebviewOfServerChanges()
}
},
6 * 60 * 1000,
)
this._oauthWatchers.set(watcherKey, { unsubscribe, abortHandle })
}
private appendErrorMessage(connection: McpConnection, error: string, level: "error" | "warn" | "info" = "error") {
const MAX_ERROR_LENGTH = 1000
const truncatedError =
@ -1222,6 +1365,15 @@ export class McpHub {
}
async deleteConnection(name: string, source?: "global" | "project"): Promise<void> {
// Cancel any active OAuth token watchers for this connection
const watcherKey = `${name}:${source}`
const watcher = this._oauthWatchers.get(watcherKey)
if (watcher) {
watcher.unsubscribe()
clearTimeout(watcher.abortHandle)
this._oauthWatchers.delete(watcherKey)
}
// Clean up file watchers for this server
this.removeFileWatchersForServer(name)
@ -2181,6 +2333,14 @@ export class McpHub {
}
this.isProgrammaticUpdate = false
// Cancel all active OAuth token watchers
for (const { unsubscribe, abortHandle } of this._oauthWatchers.values()) {
unsubscribe()
clearTimeout(abortHandle)
}
this._oauthWatchers.clear()
this.removeAllFileWatchers()
for (const connection of this.connections) {

75
src/services/mcp/McpOAuthClientProvider.ts
View file

@ -37,6 +37,10 @@ export class McpOAuthClientProvider implements OAuthClientProvider {
private _clientInfo?: OAuthClientInformationFull
private _closed = false
private _refreshPromise: Promise<OAuthTokens> | null = null
/** Stored by redirectToAuthorization(); opened on-demand via openBrowser(). */
private _pendingAuthorizationUrl: URL | null = null
/** Deduplicates concurrent _ensureCallbackServer() calls. */
private _ensureServerPromise: Promise<void> | null = null
private constructor(
private readonly _serverUrl: string,
@ -88,23 +92,15 @@ export class McpOAuthClientProvider implements OAuthClientProvider {
.map((b) => b.toString(16).padStart(2, "0"))
.join("")
// Start the callback server now so the port is known and stable.
// The SDK reads `redirectUrl` synchronously when building the authorization
// URL, so the port must be available before any connect attempt.
const { server, port, result } = await startCallbackServer(undefined, state)
const authCodePromise = result.then((r) => {
if (r.error) throw new Error(`OAuth authorization failed: ${r.error}`)
if (!r.code) throw new Error("No authorization code received in callback")
return r.code
})
// We start the callback server lazily in `redirectToAuthorization()` or `waitForAuthCode()`.
// We use a default port (0) initially; it will be updated when the server starts.
return new McpOAuthClientProvider(
serverUrl,
secretStorage,
server,
port,
authCodePromise,
null,
0,
null,
tokenEndpointAuthMethod,
grantTypes,
scopes,
@ -121,9 +117,20 @@ export class McpOAuthClientProvider implements OAuthClientProvider {
return `http://localhost:${this._port}/callback`
}
private async _ensureCallbackServer(): Promise<void> {
if (this._server && !this._closed) return
private _ensureCallbackServer(): Promise<void> {
// Guard against concurrent callers (e.g. redirectToAuthorization + registerClientIfNeeded
// called in parallel) both passing the "server not yet started" check and each launching
// their own startCallbackServer(), which would bind two ports and lose one handle.
if (this._server && !this._closed) return Promise.resolve()
if (!this._ensureServerPromise) {
this._ensureServerPromise = this._doStartCallbackServer().finally(() => {
this._ensureServerPromise = null
})
}
return this._ensureServerPromise
}
private async _doStartCallbackServer(): Promise<void> {
this._closed = false
const { server, port, result } = await startCallbackServer(this._port, this._state)
this._server = server
@ -188,6 +195,10 @@ export class McpOAuthClientProvider implements OAuthClientProvider {
if (!this._authServerMeta?.registration_endpoint) return // DCR not supported
// For Dynamic Client Registration, we MUST have a stable redirect URI.
// Ensure the callback server is started so we have a real port.
await this._ensureCallbackServer()
const response = await fetch(this._authServerMeta.registration_endpoint as string, {
method: "POST",
headers: {
@ -254,8 +265,11 @@ export class McpOAuthClientProvider implements OAuthClientProvider {
}
async redirectToAuthorization(authorizationUrl: URL): Promise<void> {
// Ensure the callback server is running before opening the browser.
// This handles mid-session re-auth where the initial server was closed.
// Ensure the callback server is running so redirectUrl has a real port.
// The server must be started here because the SDK calls this method as
// part of its internal auth flow (before throwing UnauthorizedError back
// to our caller). We do NOT open the browser here — that is deferred to
// openBrowser(), which McpHub calls only after the user confirms the toast.
await this._ensureCallbackServer()
// Workaround for SDK metadata discovery bug (see utils/oauth.ts for issue links).
@ -290,13 +304,25 @@ export class McpOAuthClientProvider implements OAuthClientProvider {
}
}
void vscode.window.showInformationMessage("MCP server requires authentication. Opening browser for OAuth…")
// Store the (possibly corrected) URL; it will be opened by openBrowser()
// once the user confirms the "Authenticate" toast in McpHub.
this._pendingAuthorizationUrl = correctedUrl
}
/**
* Opens the pending OAuth authorization URL in the system browser.
* Must be called after `redirectToAuthorization()` has been invoked by the SDK.
* McpHub calls this only after the user confirms the authentication toast.
*/
async openBrowser(): Promise<void> {
const url = this._pendingAuthorizationUrl
if (!url) {
throw new Error("No pending authorization URL — redirectToAuthorization() was not called")
}
try {
await vscode.env.openExternal(vscode.Uri.parse(correctedUrl.toString()))
await vscode.env.openExternal(vscode.Uri.parse(url.toString()))
} catch {
void vscode.window.showInformationMessage(
`Please open this URL in your browser to authenticate: ${correctedUrl}`,
)
void vscode.window.showInformationMessage(`Please open this URL in your browser to authenticate: ${url}`)
}
}
@ -429,6 +455,11 @@ export class McpOAuthClientProvider implements OAuthClientProvider {
/** Close the local callback server. Always call this when done. */
async close(): Promise<void> {
// If a server startup is in flight, wait for it to finish so we don't
// close before _server is set (which would leave a dangling server).
if (this._ensureServerPromise) {
await this._ensureServerPromise.catch(() => {})
}
if (!this._closed && this._server) {
this._closed = true
await stopCallbackServer(this._server).catch(() => {})

17
src/services/mcp/SecretStorageService.ts
View file

@ -48,4 +48,21 @@ export class SecretStorageService {
async deleteOAuthData(serverUrl: string): Promise<void> {
await this._storage.delete(this._key(serverUrl))
}
/**
* Subscribe to changes for a specific server URL's OAuth data.
* The callback fires (in all VS Code windows) immediately when another
* window writes or deletes the token for this server.
*
* @returns A dispose function — call it to stop listening.
*/
onDidChange(serverUrl: string, callback: () => void): () => void {
const key = this._key(serverUrl)
const disposable = this._storage.onDidChange((e) => {
if (e.key === key) {
callback()
}
})
return () => disposable.dispose()
}
}

113
src/services/mcp/__tests__/McpOAuthClientProvider.spec.ts
View file

@ -85,18 +85,29 @@ describe("McpOAuthClientProvider", () => {
})
describe("create", () => {
it("should start a callback server and return a provider", async () => {
it("should return a provider without starting a callback server", async () => {
setupCallbackServerMock()
const secretStorage = createMockSecretStorage()
const provider = await McpOAuthClientProvider.create("https://example.com/mcp", secretStorage)
expect(startCallbackServer).toHaveBeenCalledWith(undefined, expect.any(String))
expect(provider.redirectUrl).toBe("http://localhost:12345/callback")
expect(startCallbackServer).not.toHaveBeenCalled()
expect(provider.redirectUrl).toBe("http://localhost:0/callback")
await provider.close()
})
})
describe("redirectUrl (pre-server-start)", () => {
it("should return localhost:0 before the callback server is started (intentional lazy-init behaviour)", async () => {
const secretStorage = createMockSecretStorage()
const provider = await McpOAuthClientProvider.create("https://example.com/mcp", secretStorage)
// Port 0 is intentional: the server starts lazily in _ensureCallbackServer().
// DCR and authorization flows always call _ensureCallbackServer() first to obtain
// a real port before using redirectUrl, so port 0 is never sent to an OAuth server.
expect(provider.redirectUrl).toBe("http://localhost:0/callback")
})
})
describe("clientMetadata", () => {
it("should return correct metadata with redirect URI", async () => {
setupCallbackServerMock()
@ -105,7 +116,7 @@ describe("McpOAuthClientProvider", () => {
const metadata = provider.clientMetadata
expect(metadata.client_name).toBe("Roo Code")
expect(metadata.redirect_uris).toEqual(["http://localhost:12345/callback"])
expect(metadata.redirect_uris).toEqual(["http://localhost:0/callback"])
expect(metadata.grant_types).toContain("authorization_code")
expect(metadata.response_types).toContain("code")
expect(metadata.token_endpoint_auth_method).toBe("none")
@ -253,17 +264,29 @@ describe("McpOAuthClientProvider", () => {
})
describe("redirectToAuthorization", () => {
it("should open browser with the authorization URL", async () => {
it("should store the authorization URL without opening the browser", async () => {
setupCallbackServerMock()
const provider = await McpOAuthClientProvider.create("https://example.com/mcp", createMockSecretStorage())
const authUrl = new URL("https://auth.example.com/authorize?client_id=test")
await provider.redirectToAuthorization(authUrl)
// Browser must NOT have been opened yet — it is deferred to openBrowser()
expect(vscode.env.openExternal).not.toHaveBeenCalled()
await provider.close()
})
})
describe("openBrowser", () => {
it("should open the pending authorization URL in the browser", async () => {
setupCallbackServerMock()
const provider = await McpOAuthClientProvider.create("https://example.com/mcp", createMockSecretStorage())
const authUrl = new URL("https://auth.example.com/authorize?client_id=test")
await provider.redirectToAuthorization(authUrl)
await provider.openBrowser()
expect(vscode.env.openExternal).toHaveBeenCalled()
expect(vscode.window.showInformationMessage).toHaveBeenCalledWith(
expect.stringContaining("Opening browser for OAuth"),
)
await provider.close()
})
@ -274,6 +297,7 @@ describe("McpOAuthClientProvider", () => {
const authUrl = new URL("https://auth.example.com/authorize?client_id=test")
await provider.redirectToAuthorization(authUrl)
await provider.openBrowser()
expect(vscode.window.showInformationMessage).toHaveBeenCalledWith(
expect.stringContaining("Please open this URL"),
@ -281,6 +305,13 @@ describe("McpOAuthClientProvider", () => {
await provider.close()
})
it("should throw if called before redirectToAuthorization", async () => {
setupCallbackServerMock()
const provider = await McpOAuthClientProvider.create("https://example.com/mcp", createMockSecretStorage())
await expect(provider.openBrowser()).rejects.toThrow("No pending authorization URL")
await provider.close()
})
it("should correct a wrong authorization URL using pre-fetched metadata", async () => {
// Mock discovery to return an issuer with a path component.
// The SDK's discoverOAuthMetadata builds the wrong URL for such issuers,
@ -309,6 +340,7 @@ describe("McpOAuthClientProvider", () => {
// Simulate the SDK building the wrong base URL (using bare /authorize) and omitting scope
const sdkWrongUrl = new URL("https://mcp.kapa.ai/authorize?client_id=abc&code_challenge=xyz&state=123")
await provider.redirectToAuthorization(sdkWrongUrl)
await provider.openBrowser()
// The provider should have corrected the URL to use the real authorization_endpoint
const openedUri = (vscode.env.openExternal as any).mock.calls[0][0].toString()
@ -351,6 +383,7 @@ describe("McpOAuthClientProvider", () => {
const sdkUrl = new URL("https://mcp.kapa.ai/authorize?client_id=abc&state=123")
await provider.redirectToAuthorization(sdkUrl)
await provider.openBrowser()
const openedUri = (vscode.env.openExternal as any).mock.calls[0][0].toString()
// The resource indicator from the protected resource metadata must appear
@ -368,6 +401,7 @@ describe("McpOAuthClientProvider", () => {
"https://auth.example.com/authorize?client_id=abc&resource=https%3A%2F%2Fexample.com%2F&state=123",
)
await provider.redirectToAuthorization(sdkUrl)
await provider.openBrowser()
const openedUri = (vscode.env.openExternal as any).mock.calls[0][0].toString()
const resourceMatches = (openedUri.match(/resource=/g) || []).length
@ -382,6 +416,7 @@ describe("McpOAuthClientProvider", () => {
// SDK URL already includes scope=openid
const sdkUrl = new URL("https://auth.example.com/authorize?client_id=abc&scope=openid&state=123")
await provider.redirectToAuthorization(sdkUrl)
await provider.openBrowser()
// scope should appear exactly once
const openedUri = (vscode.env.openExternal as any).mock.calls[0][0].toString()
@ -445,7 +480,7 @@ describe("McpOAuthClientProvider", () => {
expect(body.get("client_id")).toBe("client-id-123")
expect(body.get("client_secret")).toBe("client-secret-abc")
expect(body.get("code_verifier")).toBe("pkce-verifier-123")
expect(body.get("redirect_uri")).toBe("http://localhost:12345/callback")
expect(body.get("redirect_uri")).toBe("http://localhost:0/callback")
// Verify tokens were saved
const saved = await provider.tokens()
@ -564,10 +599,13 @@ describe("McpOAuthClientProvider", () => {
})
describe("close", () => {
it("should stop the callback server", async () => {
it("should stop the callback server if it was started", async () => {
const { mockServer } = setupCallbackServerMock()
const provider = await McpOAuthClientProvider.create("https://example.com/mcp", createMockSecretStorage())
// Start server lazily
await provider.waitForAuthCode().catch(() => {})
await provider.close()
expect(stopCallbackServer).toHaveBeenCalledWith(mockServer)
@ -577,11 +615,23 @@ describe("McpOAuthClientProvider", () => {
setupCallbackServerMock()
const provider = await McpOAuthClientProvider.create("https://example.com/mcp", createMockSecretStorage())
// Start server lazily
await provider.waitForAuthCode().catch(() => {})
await provider.close()
await provider.close()
expect(stopCallbackServer).toHaveBeenCalledTimes(1)
})
it("should not call stopCallbackServer if server was never started", async () => {
setupCallbackServerMock()
const provider = await McpOAuthClientProvider.create("https://example.com/mcp", createMockSecretStorage())
await provider.close()
expect(stopCallbackServer).not.toHaveBeenCalled()
})
})
describe("registerClientIfNeeded", () => {
@ -668,5 +718,48 @@ describe("McpOAuthClientProvider", () => {
expect((await provider.clientInformation())?.client_id).toBe("new-client-id")
await provider.close()
})
it("should use the same redirect URI in DCR and authorization flow", async () => {
setupCallbackServerMock()
const secretStorage = createMockSecretStorage()
mockFetch.mockClear()
// Auth server metadata
mockFetch.mockResolvedValueOnce({
ok: true,
json: () =>
Promise.resolve({
issuer: "https://auth.example.com",
authorization_endpoint: "https://auth.example.com/authorize",
token_endpoint: "https://auth.example.com/token",
registration_endpoint: "https://auth.example.com/register",
response_types_supported: ["code"],
token_endpoint_auth_methods_supported: ["none"],
grant_types_supported: ["authorization_code", "refresh_token"],
}),
})
// DCR response
mockFetch.mockResolvedValueOnce({
ok: true,
json: () =>
Promise.resolve({
client_id: "consistency-client-id",
redirect_uris: ["http://localhost:12345/callback"],
client_name: "Roo Code",
grant_types: ["authorization_code", "refresh_token"],
response_types: ["code"],
token_endpoint_auth_method: "none",
}),
})
const provider = await McpOAuthClientProvider.create("https://example.com/mcp", secretStorage)
await provider.registerClientIfNeeded()
// The redirect_uri sent in the DCR body must equal the redirectUrl property
// used during the authorization redirect so RFC 6749 §4.1.3 validation passes.
const dcrBody = JSON.parse(mockFetch.mock.calls[1][1]?.body as string)
expect(dcrBody.redirect_uris).toContain(provider.redirectUrl)
await provider.close()
})
})
})

63
src/services/mcp/__tests__/SecretStorageService.spec.ts
View file

@ -6,17 +6,30 @@ import { SecretStorageService, StoredMcpOAuthData } from "../SecretStorageServic
function createMockContext() {
const store = new Map<string, string>()
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)
}),
// Listeners registered via onDidChange; keyed by arbitrary id for disposal.
const listeners = new Map<number, (e: { key: string }) => void>()
let nextId = 0
const 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)
}),
onDidChange: vi.fn((handler: (e: { key: string }) => void) => {
const id = nextId++
listeners.set(id, handler)
return { dispose: () => listeners.delete(id) }
}),
/** Test helper: simulate a storage change event. */
_emit: (key: string) => {
for (const handler of listeners.values()) handler({ key })
},
} as any
}
return { secrets } as any
}
describe("SecretStorageService", () => {
@ -90,6 +103,36 @@ describe("SecretStorageService", () => {
const result = await service.getOAuthData("https://example.com/mcp")
expect(result).toBeUndefined()
})
describe("onDidChange", () => {
it("should call the callback when the key for the given URL changes", () => {
const cb = vi.fn()
service.onDidChange("https://example.com/mcp", cb)
context.secrets._emit("mcp.oauth.example.com.mcp.data")
expect(cb).toHaveBeenCalledTimes(1)
})
it("should not call the callback for a different URL's key", () => {
const cb = vi.fn()
service.onDidChange("https://example.com/mcp", cb)
context.secrets._emit("mcp.oauth.other.com.mcp.data")
expect(cb).not.toHaveBeenCalled()
})
it("should stop calling the callback after the returned dispose function is called", () => {
const cb = vi.fn()
const unsubscribe = service.onDidChange("https://example.com/mcp", cb)
unsubscribe()
context.secrets._emit("mcp.oauth.example.com.mcp.data")
expect(cb).not.toHaveBeenCalled()
})
})
})
describe("key isolation", () => {