refactor(SecretStorageService): caching full DCR repose

This commit is contained in:
Elliott de Launay 2026-03-26 23:23:46 -04:00
parent b1d44ce58a
commit f2a3cfba07
No known key found for this signature in database
GPG key ID: BB899BED766D1806
4 changed files with 62 additions and 16 deletions

View file

@ -227,14 +227,12 @@ export class McpOAuthClientProvider implements OAuthClientProvider {
// Check if we have a cached client_id from previous registration // Check if we have a cached client_id from previous registration
const cachedData = await this._secretStorage.getOAuthData(this._serverUrl) const cachedData = await this._secretStorage.getOAuthData(this._serverUrl)
if (cachedData?.client_id) { if (cachedData?.client_info) {
// Use the full DCR response, override redirect_uris with the
// current port (which may have changed between sessions).
this._clientInfo = { this._clientInfo = {
client_id: cachedData.client_id, ...cachedData.client_info,
redirect_uris: [this.redirectUrl], redirect_uris: [this.redirectUrl],
client_name: this._clientName,
grant_types: this._grantTypes,
response_types: ["code"],
token_endpoint_auth_method: this._tokenEndpointAuthMethod,
} }
return return
} }
@ -280,7 +278,7 @@ export class McpOAuthClientProvider implements OAuthClientProvider {
// auth server bound the refresh token to. `this._clientInfo.client_id` // auth server bound the refresh token to. `this._clientInfo.client_id`
// may differ if a fresh DCR was performed (e.g. after stale token // may differ if a fresh DCR was performed (e.g. after stale token
// cleanup removed the cached data). // cleanup removed the cached data).
const clientIdForRefresh = data.client_id ?? this._clientInfo?.client_id const clientIdForRefresh = data.client_info?.client_id ?? this._clientInfo?.client_id
this._refreshPromise = this.refreshAccessToken(data.tokens.refresh_token, clientIdForRefresh).finally( this._refreshPromise = this.refreshAccessToken(data.tokens.refresh_token, clientIdForRefresh).finally(
() => { () => {
@ -303,10 +301,14 @@ export class McpOAuthClientProvider implements OAuthClientProvider {
async saveTokens(tokens: OAuthTokens, clientIdOverride?: string): Promise<void> { async saveTokens(tokens: OAuthTokens, clientIdOverride?: string): Promise<void> {
const expires_at = tokens.expires_in ? Date.now() + tokens.expires_in * 1000 : Date.now() + 3600 * 1000 // default 1 hour when server omits expires_in const expires_at = tokens.expires_in ? Date.now() + tokens.expires_in * 1000 : Date.now() + 3600 * 1000 // default 1 hour when server omits expires_in
const clientInfo =
clientIdOverride && this._clientInfo
? { ...this._clientInfo, client_id: clientIdOverride }
: this._clientInfo
await this._secretStorage.saveOAuthData(this._serverUrl, { await this._secretStorage.saveOAuthData(this._serverUrl, {
tokens, tokens,
expires_at, expires_at,
client_id: clientIdOverride ?? this._clientInfo?.client_id, ...(clientInfo ? { client_info: clientInfo } : {}),
}) })
} }

View file

@ -1,12 +1,17 @@
import * as vscode from "vscode" import * as vscode from "vscode"
import type { OAuthTokens } from "@modelcontextprotocol/sdk/shared/auth.js" import type { OAuthClientInformationFull, OAuthTokens } from "@modelcontextprotocol/sdk/shared/auth.js"
export interface StoredMcpOAuthData { export interface StoredMcpOAuthData {
tokens: OAuthTokens tokens: OAuthTokens
/** Unix ms timestamp after which the access token should be considered expired. */ /** Unix ms timestamp after which the access token should be considered expired. */
expires_at: number expires_at: number
/** The client_id used to obtain these tokens (for token reuse without re-registration). */ /**
client_id?: string * Full DCR response from the auth server, persisted so that fields like
* client_secret, grant_types, and token_endpoint_auth_method survive restarts.
* Note: redirect_uris within this object may be stale (port changes between
* sessions); callers must override redirect_uris with the current value.
*/
client_info?: OAuthClientInformationFull
} }
/** /**

View file

@ -702,15 +702,22 @@ describe("McpOAuthClientProvider", () => {
}) })
describe("registerClientIfNeeded", () => { describe("registerClientIfNeeded", () => {
it("should reuse cached client_id from previous registration", async () => { it("should reuse cached client_info from previous registration", async () => {
setupCallbackServerMock() setupCallbackServerMock()
const secretStorage = createMockSecretStorage() const secretStorage = createMockSecretStorage()
// Pre-populate storage with cached data // Pre-populate storage with cached data including full client_info
await secretStorage.saveOAuthData("https://example.com/mcp", { await secretStorage.saveOAuthData("https://example.com/mcp", {
tokens: { access_token: "cached-token", token_type: "Bearer" }, tokens: { access_token: "cached-token", token_type: "Bearer" },
expires_at: Date.now() + 3600000, expires_at: Date.now() + 3600000,
client_info: {
client_id: "cached-client-id", client_id: "cached-client-id",
client_name: "Test Client",
redirect_uris: ["http://localhost:9999/callback"],
grant_types: ["authorization_code", "refresh_token"],
response_types: ["code"],
token_endpoint_auth_method: "none",
},
}) })
const provider = await McpOAuthClientProvider.create("https://example.com/mcp", secretStorage) const provider = await McpOAuthClientProvider.create("https://example.com/mcp", secretStorage)
@ -720,7 +727,7 @@ describe("McpOAuthClientProvider", () => {
await provider.close() await provider.close()
}) })
it("should reuse cached client_id even when callback server port has changed", async () => { it("should reuse cached client_info even when callback server port has changed", async () => {
setupCallbackServerMock() setupCallbackServerMock()
const secretStorage = createMockSecretStorage() const secretStorage = createMockSecretStorage()
@ -730,7 +737,14 @@ describe("McpOAuthClientProvider", () => {
await secretStorage.saveOAuthData("https://example.com/mcp", { await secretStorage.saveOAuthData("https://example.com/mcp", {
tokens: { access_token: "cached-token", token_type: "Bearer" }, tokens: { access_token: "cached-token", token_type: "Bearer" },
expires_at: Date.now() + 3600000, expires_at: Date.now() + 3600000,
client_info: {
client_id: "cached-client-id", client_id: "cached-client-id",
client_name: "Test Client",
redirect_uris: ["http://localhost:9999/callback"],
grant_types: ["authorization_code", "refresh_token"],
response_types: ["code"],
token_endpoint_auth_method: "none",
},
}) })
const provider = await McpOAuthClientProvider.create("https://example.com/mcp", secretStorage) const provider = await McpOAuthClientProvider.create("https://example.com/mcp", secretStorage)

View file

@ -160,6 +160,31 @@ describe("SecretStorageService", () => {
}) })
}) })
describe("client_info round-trip", () => {
it("should persist and retrieve full client_info", async () => {
const data: StoredMcpOAuthData = {
tokens: { access_token: "tok", token_type: "Bearer" },
expires_at: Date.now() + 3600_000,
client_info: {
client_id: "cid-123",
client_secret: "secret-456",
client_name: "Test Client",
redirect_uris: ["http://localhost:12345/callback"],
grant_types: ["authorization_code", "refresh_token"],
response_types: ["code"],
token_endpoint_auth_method: "client_secret_post",
},
}
await service.saveOAuthData("https://example.com/mcp", data)
const result = await service.getOAuthData("https://example.com/mcp")
expect(result).toEqual(data)
expect(result?.client_info?.client_id).toBe("cid-123")
expect(result?.client_info?.client_secret).toBe("secret-456")
expect(result?.client_info?.token_endpoint_auth_method).toBe("client_secret_post")
})
})
describe("key isolation", () => { describe("key isolation", () => {
it("should isolate data by host", async () => { it("should isolate data by host", async () => {
const data1: StoredMcpOAuthData = { const data1: StoredMcpOAuthData = {