Skip to content

Commit 2e33986

Browse files
rekram1-noderekram1-node
committed
fix(mcp): guard credential invalidation races
1 parent b9fe4d6 commit 2e33986

3 files changed

Lines changed: 62 additions & 52 deletions

File tree

‎packages/opencode/src/mcp/auth.ts‎

Lines changed: 28 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import { Global } from "@opencode-ai/core/global"
55
import { Effect, Layer, Context, Option, Schema } from "effect"
66
import { FSUtil } from "@opencode-ai/core/fs-util"
77
import { EffectFlock } from "@opencode-ai/core/util/effect-flock"
8+
import { isDeepEqual } from "remeda"
89

910
export const Tokens = Schema.Struct({
1011
accessToken: Schema.mutableKey(Schema.String),
@@ -44,9 +45,12 @@ export interface Interface {
4445
readonly set: (mcpName: string, entry: Entry, serverUrl?: string) => Effect.Effect<void>
4546
readonly remove: (mcpName: string) => Effect.Effect<void>
4647
readonly updateTokens: (mcpName: string, tokens: Tokens, serverUrl?: string) => Effect.Effect<void>
47-
readonly clearTokens: (mcpName: string) => Effect.Effect<void>
4848
readonly updateClientInfo: (mcpName: string, clientInfo: ClientInfo, serverUrl?: string) => Effect.Effect<void>
49-
readonly clearClientInfo: (mcpName: string) => Effect.Effect<void>
49+
readonly invalidateCredentials: (
50+
mcpName: string,
51+
type: "all" | "client" | "tokens",
52+
expected: Pick<Entry, "tokens" | "clientInfo">,
53+
) => Effect.Effect<void>
5054
readonly updateCodeVerifier: (mcpName: string, codeVerifier: string) => Effect.Effect<void>
5155
readonly clearCodeVerifier: (mcpName: string) => Effect.Effect<void>
5256
readonly updateOAuthState: (mcpName: string, oauthState: string) => Effect.Effect<void>
@@ -136,11 +140,26 @@ export const layer = Layer.effect(
136140
const updateClientInfo = updateField("clientInfo", "updateClientInfo")
137141
const updateCodeVerifier = updateField("codeVerifier", "updateCodeVerifier")
138142
const updateOAuthState = updateField("oauthState", "updateOAuthState")
139-
const clearTokens = clearField("tokens", "clearTokens")
140-
const clearClientInfo = clearField("clientInfo", "clearClientInfo")
141143
const clearCodeVerifier = clearField("codeVerifier", "clearCodeVerifier")
142144
const clearOAuthState = clearField("oauthState", "clearOAuthState")
143145

146+
const invalidateCredentials = Effect.fn("McpAuth.invalidateCredentials")(function* (
147+
mcpName: string,
148+
type: "all" | "client" | "tokens",
149+
expected: Pick<Entry, "tokens" | "clientInfo">,
150+
) {
151+
yield* mutate((data) => {
152+
const current = data[mcpName]
153+
if (!current) return undefined
154+
const entry = { ...current }
155+
if ((type === "all" || type === "tokens") && credentialsEqual(current.tokens, expected.tokens))
156+
delete entry.tokens
157+
if ((type === "all" || type === "client") && credentialsEqual(current.clientInfo, expected.clientInfo))
158+
delete entry.clientInfo
159+
return { ...data, [mcpName]: entry }
160+
})
161+
})
162+
144163
const resetOAuthFlow = Effect.fn("McpAuth.resetOAuthFlow")(function* (mcpName: string) {
145164
yield* mutate((data) => {
146165
const current = data[mcpName]
@@ -164,9 +183,8 @@ export const layer = Layer.effect(
164183
set,
165184
remove,
166185
updateTokens,
167-
clearTokens,
168186
updateClientInfo,
169-
clearClientInfo,
187+
invalidateCredentials,
170188
updateCodeVerifier,
171189
clearCodeVerifier,
172190
updateOAuthState,
@@ -181,4 +199,8 @@ export const defaultLayer = layer.pipe(Layer.provide(EffectFlock.defaultLayer),
181199

182200
export const node = LayerNode.make({ service: Service, layer: layer, deps: [FSUtil.node, EffectFlock.node] })
183201

202+
function credentialsEqual(left: Tokens | ClientInfo | undefined, right: Tokens | ClientInfo | undefined) {
203+
return isDeepEqual(left, right)
204+
}
205+
184206
export * as McpAuth from "./auth"

‎packages/opencode/src/mcp/oauth-provider.ts‎

Lines changed: 24 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,8 @@ export interface McpOAuthCallbacks {
2727
export class McpOAuthProvider implements OAuthClientProvider {
2828
private pendingClientInfo?: OAuthClientInformationFull
2929
private pendingTokens?: OAuthTokens
30+
private observedClientInfo?: McpAuth.ClientInfo
31+
private observedTokens?: McpAuth.Tokens
3032

3133
constructor(
3234
private mcpName: string,
@@ -71,6 +73,7 @@ export class McpOAuthProvider implements OAuthClientProvider {
7173
// Use getForUrl to validate credentials are for the current server URL
7274
const entry = await Effect.runPromise(this.auth.getForUrl(this.mcpName, this.serverUrl))
7375
if (entry?.clientInfo) {
76+
this.observedClientInfo = entry.clientInfo
7477
// Check if client secret has expired
7578
if (entry.clientInfo.clientSecretExpiresAt && entry.clientInfo.clientSecretExpiresAt < Date.now() / 1000) {
7679
return undefined
@@ -90,25 +93,21 @@ export class McpOAuthProvider implements OAuthClientProvider {
9093
this.pendingClientInfo = info
9194
return
9295
}
93-
await Effect.runPromise(
94-
this.auth.updateClientInfo(
95-
this.mcpName,
96-
{
97-
clientId: info.client_id,
98-
clientSecret: info.client_secret,
99-
clientIdIssuedAt: info.client_id_issued_at,
100-
clientSecretExpiresAt: info.client_secret_expires_at,
101-
},
102-
this.serverUrl,
103-
),
104-
)
96+
this.observedClientInfo = {
97+
clientId: info.client_id,
98+
clientSecret: info.client_secret,
99+
clientIdIssuedAt: info.client_id_issued_at,
100+
clientSecretExpiresAt: info.client_secret_expires_at,
101+
}
102+
await Effect.runPromise(this.auth.updateClientInfo(this.mcpName, this.observedClientInfo, this.serverUrl))
105103
}
106104

107105
async tokens(): Promise<OAuthTokens | undefined> {
108106
if (this.config.deferPersistence) return this.pendingTokens
109107
// Use getForUrl to validate tokens are for the current server URL
110108
const entry = await Effect.runPromise(this.auth.getForUrl(this.mcpName, this.serverUrl))
111109
if (!entry?.tokens) return undefined
110+
this.observedTokens = entry.tokens
112111

113112
return {
114113
access_token: entry.tokens.accessToken,
@@ -126,18 +125,13 @@ export class McpOAuthProvider implements OAuthClientProvider {
126125
this.pendingTokens = tokens
127126
return
128127
}
129-
await Effect.runPromise(
130-
this.auth.updateTokens(
131-
this.mcpName,
132-
{
133-
accessToken: tokens.access_token,
134-
refreshToken: tokens.refresh_token,
135-
expiresAt: tokens.expires_in ? Date.now() / 1000 + tokens.expires_in : undefined,
136-
scope: tokens.scope,
137-
},
138-
this.serverUrl,
139-
),
140-
)
128+
this.observedTokens = {
129+
accessToken: tokens.access_token,
130+
refreshToken: tokens.refresh_token,
131+
expiresAt: tokens.expires_in ? Date.now() / 1000 + tokens.expires_in : undefined,
132+
scope: tokens.scope,
133+
}
134+
await Effect.runPromise(this.auth.updateTokens(this.mcpName, this.observedTokens, this.serverUrl))
141135
}
142136

143137
async redirectToAuthorization(authorizationUrl: URL): Promise<void> {
@@ -183,17 +177,12 @@ export class McpOAuthProvider implements OAuthClientProvider {
183177
if (type === "all" || type === "tokens") this.pendingTokens = undefined
184178
return
185179
}
186-
switch (type) {
187-
case "all":
188-
await Effect.runPromise(this.auth.remove(this.mcpName))
189-
break
190-
case "client":
191-
await Effect.runPromise(this.auth.clearClientInfo(this.mcpName))
192-
break
193-
case "tokens":
194-
await Effect.runPromise(this.auth.clearTokens(this.mcpName))
195-
break
196-
}
180+
await Effect.runPromise(
181+
this.auth.invalidateCredentials(this.mcpName, type, {
182+
clientInfo: this.observedClientInfo,
183+
tokens: this.observedTokens,
184+
}),
185+
)
197186
}
198187

199188
async commit(): Promise<void> {

‎packages/opencode/test/mcp/auth.test.ts‎

Lines changed: 10 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -61,9 +61,8 @@ function gateCredentialInvalidation(
6161
const gate = Deferred.succeed(ready, undefined).pipe(Effect.ignore, Effect.andThen(Deferred.await(release)))
6262
return McpAuth.Service.of({
6363
...auth,
64-
get: (mcpName) => auth.get(mcpName).pipe(Effect.tap(() => gate)),
65-
clearTokens: (mcpName) => gate.pipe(Effect.andThen(auth.clearTokens(mcpName))),
66-
clearClientInfo: (mcpName) => gate.pipe(Effect.andThen(auth.clearClientInfo(mcpName))),
64+
invalidateCredentials: (mcpName, type, expected) =>
65+
gate.pipe(Effect.andThen(auth.invalidateCredentials(mcpName, type, expected))),
6766
})
6867
}
6968

@@ -92,7 +91,7 @@ test("serializes concurrent auth file updates across service instances", async (
9291
)
9392
})
9493

95-
test("concurrent token invalidation does not overwrite newer client registration", async () => {
94+
test("stale token invalidation does not delete newer tokens", async () => {
9695
const name = `token-invalidation-${crypto.randomUUID()}`
9796
const url = "https://mcp.example.com/exact/path"
9897

@@ -111,22 +110,22 @@ test("concurrent token invalidation does not overwrite newer client registration
111110
)
112111

113112
yield* first.updateTokens(name, { accessToken: "old-token" }, url)
113+
yield* Effect.promise(() => provider.tokens())
114114
const invalidation = yield* Effect.promise(() => provider.invalidateCredentials("tokens")).pipe(Effect.forkChild)
115115
yield* Deferred.await(ready)
116-
yield* second.updateClientInfo(name, { clientId: "new-client" }, url)
116+
yield* second.updateTokens(name, { accessToken: "new-token" }, url)
117117
yield* Deferred.succeed(release, undefined).pipe(Effect.ignore)
118118
yield* Fiber.join(invalidation)
119119

120120
const entry = yield* first.get(name)
121-
expect(entry?.tokens).toBeUndefined()
122-
expect(entry?.clientInfo?.clientId).toBe("new-client")
121+
expect(entry?.tokens?.accessToken).toBe("new-token")
123122
expect(entry?.serverUrl).toBe(url)
124123
yield* first.remove(name)
125124
}).pipe(Effect.provide(McpAuth.defaultLayer)),
126125
)
127126
})
128127

129-
test("concurrent client invalidation does not overwrite newer tokens", async () => {
128+
test("stale client invalidation does not delete newer registration", async () => {
130129
const name = `client-invalidation-${crypto.randomUUID()}`
131130
const url = "https://mcp.example.com/exact/path"
132131

@@ -145,15 +144,15 @@ test("concurrent client invalidation does not overwrite newer tokens", async ()
145144
)
146145

147146
yield* first.updateClientInfo(name, { clientId: "old-client" }, url)
147+
yield* Effect.promise(() => provider.clientInformation())
148148
const invalidation = yield* Effect.promise(() => provider.invalidateCredentials("client")).pipe(Effect.forkChild)
149149
yield* Deferred.await(ready)
150-
yield* second.updateTokens(name, { accessToken: "new-token" }, url)
150+
yield* second.updateClientInfo(name, { clientId: "new-client" }, url)
151151
yield* Deferred.succeed(release, undefined).pipe(Effect.ignore)
152152
yield* Fiber.join(invalidation)
153153

154154
const entry = yield* first.get(name)
155-
expect(entry?.clientInfo).toBeUndefined()
156-
expect(entry?.tokens?.accessToken).toBe("new-token")
155+
expect(entry?.clientInfo?.clientId).toBe("new-client")
157156
expect(entry?.serverUrl).toBe(url)
158157
yield* first.remove(name)
159158
}).pipe(Effect.provide(McpAuth.defaultLayer)),

0 commit comments

Comments
 (0)