fix(mcp): refresh credentials on reauthentication (#33717)
This commit is contained in:
@@ -22,7 +22,7 @@ import { NamedError } from "@opencode-ai/core/util/error"
|
||||
import { InstallationVersion } from "@opencode-ai/core/installation/version"
|
||||
import { withTimeout } from "@/util/timeout"
|
||||
import { FSUtil } from "@opencode-ai/core/fs-util"
|
||||
import { McpOAuthProvider, OAUTH_CALLBACK_PATH } from "./oauth-provider"
|
||||
import { McpOAuthPendingProvider, McpOAuthProvider, OAUTH_CALLBACK_PATH } from "./oauth-provider"
|
||||
import { McpOAuthCallback } from "./oauth-callback"
|
||||
import { McpAuth } from "./auth"
|
||||
import { EventV2Bridge } from "@/event-v2-bridge"
|
||||
@@ -109,7 +109,7 @@ export type Status = Schema.Schema.Type<typeof Status>
|
||||
|
||||
// Store transports for OAuth servers to allow finishing auth
|
||||
type TransportWithAuth = StreamableHTTPClientTransport | SSEClientTransport
|
||||
const pendingOAuthTransports = new Map<string, TransportWithAuth>()
|
||||
const pendingOAuthTransports = new Map<string, { transport: TransportWithAuth; provider?: McpOAuthPendingProvider }>()
|
||||
|
||||
// Prompt cache types
|
||||
type PromptInfo = Awaited<ReturnType<MCPClient["listPrompts"]>>["prompts"][number]
|
||||
@@ -301,7 +301,7 @@ export const layer = Layer.effect(
|
||||
})
|
||||
.pipe(Effect.ignore, Effect.as(undefined))
|
||||
} else {
|
||||
pendingOAuthTransports.set(key, transport)
|
||||
pendingOAuthTransports.set(key, { transport })
|
||||
lastStatus = { status: "needs_auth" as const }
|
||||
return events
|
||||
.publish(TuiEvent.ToastShow, {
|
||||
@@ -819,7 +819,7 @@ export const layer = Layer.effect(
|
||||
.join("")
|
||||
yield* auth.updateOAuthState(mcpName, oauthState)
|
||||
let capturedUrl: URL | undefined
|
||||
const authProvider = new McpOAuthProvider(
|
||||
const authProvider = new McpOAuthPendingProvider(
|
||||
mcpName,
|
||||
mcpConfig.url,
|
||||
{
|
||||
@@ -845,15 +845,16 @@ export const layer = Layer.effect(
|
||||
return yield* Effect.tryPromise({
|
||||
try: () => {
|
||||
const client = createClient(directory)
|
||||
return client
|
||||
.connect(transport)
|
||||
.then(() => ({ authorizationUrl: "", oauthState, client }) satisfies AuthResult)
|
||||
return client.connect(transport).then(async () => {
|
||||
await authProvider.commit()
|
||||
return { authorizationUrl: "", oauthState, client } satisfies AuthResult
|
||||
})
|
||||
},
|
||||
catch: (error) => error,
|
||||
}).pipe(
|
||||
Effect.catch((error) => {
|
||||
if (error instanceof UnauthorizedError && capturedUrl) {
|
||||
pendingOAuthTransports.set(mcpName, transport)
|
||||
pendingOAuthTransports.set(mcpName, { transport, provider: authProvider })
|
||||
return Effect.succeed({ authorizationUrl: capturedUrl.toString(), oauthState } satisfies AuthResult)
|
||||
}
|
||||
return Effect.die(error)
|
||||
@@ -924,11 +925,11 @@ export const layer = Layer.effect(
|
||||
|
||||
const finishAuth = Effect.fn("MCP.finishAuth")(function* (mcpName: string, authorizationCode: string) {
|
||||
yield* requireMcpConfig(mcpName)
|
||||
const transport = pendingOAuthTransports.get(mcpName)
|
||||
if (!transport) throw new Error(`No pending OAuth flow for MCP server: ${mcpName}`)
|
||||
const pending = pendingOAuthTransports.get(mcpName)
|
||||
if (!pending) throw new Error(`No pending OAuth flow for MCP server: ${mcpName}`)
|
||||
|
||||
const result = yield* Effect.tryPromise({
|
||||
try: () => transport.finishAuth(authorizationCode).then(() => true as const),
|
||||
try: () => pending.transport.finishAuth(authorizationCode).then(() => true as const),
|
||||
catch: (error) => {
|
||||
return error
|
||||
},
|
||||
@@ -938,6 +939,7 @@ export const layer = Layer.effect(
|
||||
return { status: "failed", error: "OAuth completion failed" } satisfies Status
|
||||
}
|
||||
|
||||
yield* Effect.promise(() => pending.provider?.commit() ?? Promise.resolve())
|
||||
yield* auth.clearCodeVerifier(mcpName)
|
||||
pendingOAuthTransports.delete(mcpName)
|
||||
|
||||
|
||||
@@ -25,11 +25,11 @@ export interface McpOAuthCallbacks {
|
||||
|
||||
export class McpOAuthProvider implements OAuthClientProvider {
|
||||
constructor(
|
||||
private mcpName: string,
|
||||
private serverUrl: string,
|
||||
private config: McpOAuthConfig,
|
||||
protected mcpName: string,
|
||||
protected serverUrl: string,
|
||||
protected config: McpOAuthConfig,
|
||||
private callbacks: McpOAuthCallbacks,
|
||||
private auth: McpAuth.Interface,
|
||||
protected auth: McpAuth.Interface,
|
||||
) {}
|
||||
|
||||
get redirectUrl(): string {
|
||||
@@ -53,7 +53,6 @@ export class McpOAuthProvider implements OAuthClientProvider {
|
||||
}
|
||||
|
||||
async clientInformation(): Promise<OAuthClientInformation | undefined> {
|
||||
// Check config first (pre-registered client)
|
||||
if (this.config.clientId) {
|
||||
return {
|
||||
client_id: this.config.clientId,
|
||||
@@ -164,10 +163,7 @@ export class McpOAuthProvider implements OAuthClientProvider {
|
||||
|
||||
async invalidateCredentials(type: "all" | "client" | "tokens"): Promise<void> {
|
||||
const entry = await Effect.runPromise(this.auth.get(this.mcpName))
|
||||
if (!entry) {
|
||||
return
|
||||
}
|
||||
|
||||
if (!entry) return
|
||||
switch (type) {
|
||||
case "all":
|
||||
await Effect.runPromise(this.auth.remove(this.mcpName))
|
||||
@@ -184,6 +180,63 @@ export class McpOAuthProvider implements OAuthClientProvider {
|
||||
}
|
||||
}
|
||||
|
||||
export class McpOAuthPendingProvider extends McpOAuthProvider {
|
||||
private pendingClientInfo?: OAuthClientInformationFull
|
||||
private pendingTokens?: OAuthTokens
|
||||
|
||||
override async clientInformation(): Promise<OAuthClientInformation | undefined> {
|
||||
if (!this.config.clientId) return this.pendingClientInfo
|
||||
return {
|
||||
client_id: this.config.clientId,
|
||||
client_secret: this.config.clientSecret,
|
||||
}
|
||||
}
|
||||
|
||||
override async saveClientInformation(info: OAuthClientInformationFull): Promise<void> {
|
||||
this.pendingClientInfo = info
|
||||
}
|
||||
|
||||
override async tokens(): Promise<OAuthTokens | undefined> {
|
||||
return this.pendingTokens
|
||||
}
|
||||
|
||||
override async saveTokens(tokens: OAuthTokens): Promise<void> {
|
||||
this.pendingTokens = tokens
|
||||
}
|
||||
|
||||
override async invalidateCredentials(type: "all" | "client" | "tokens"): Promise<void> {
|
||||
if (type === "all" || type === "client") this.pendingClientInfo = undefined
|
||||
if (type === "all" || type === "tokens") this.pendingTokens = undefined
|
||||
}
|
||||
|
||||
async commit(): Promise<void> {
|
||||
if (!this.pendingTokens) return
|
||||
await Effect.runPromise(
|
||||
this.auth.set(
|
||||
this.mcpName,
|
||||
{
|
||||
tokens: {
|
||||
accessToken: this.pendingTokens.access_token,
|
||||
refreshToken: this.pendingTokens.refresh_token,
|
||||
expiresAt: this.pendingTokens.expires_in ? Date.now() / 1000 + this.pendingTokens.expires_in : undefined,
|
||||
scope: this.pendingTokens.scope,
|
||||
},
|
||||
clientInfo:
|
||||
this.pendingClientInfo && !this.config.clientId
|
||||
? {
|
||||
clientId: this.pendingClientInfo.client_id,
|
||||
clientSecret: this.pendingClientInfo.client_secret,
|
||||
clientIdIssuedAt: this.pendingClientInfo.client_id_issued_at,
|
||||
clientSecretExpiresAt: this.pendingClientInfo.client_secret_expires_at,
|
||||
}
|
||||
: undefined,
|
||||
},
|
||||
this.serverUrl,
|
||||
),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
export { OAUTH_CALLBACK_PORT, OAUTH_CALLBACK_PATH }
|
||||
|
||||
/**
|
||||
|
||||
Reference in New Issue
Block a user