Visitar URL original
refactor(core): share OAuth credential refresh between provider logins and MCP by rekram1-node · Pull Request #63022 · anomalyco/opencode · GitHub
Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
75 changes: 57 additions & 18 deletions packages/core/src/credential.ts
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
export * as Credential from "./credential.js"

import { asc, desc, eq, type SQL } from "drizzle-orm"
import { Cause, Context, Effect, Fiber, Layer, Schema } from "effect"
import { Cause, Clock, Context, Duration, Effect, Fiber, Layer, Schema } from "effect"
import { Credential } from "@opencode/schema/credential"
import { CredentialLegacy } from "@opencode/schema/credential-legacy"
import { Integration } from "@opencode/schema/integration"
Expand Down Expand Up @@ -29,6 +29,10 @@ export type Value = Credential.Value

export const Event = Credential.Event

/** Whether an access token is due for a refresh. `expires` 0 means the issuer gave no lifetime. */
export const expiring = (value: OAuth, now: number) =>
value.expires !== 0 && value.expires <= now + Duration.toMillis(Duration.minutes(5))

export class Info extends Schema.Class<Info>("Credential.Info")({
id: ID,
integrationID: Integration.ID,
Expand All @@ -43,10 +47,19 @@ export interface Interface {
readonly list: (integrationID: Integration.ID) => Effect.Effect<Info[]>
/** Returns one stored credential by ID. */
readonly get: (id: ID) => Effect.Effect<Info | undefined>
/** Concurrent callers for the same credential share the running operation's result, including failures. */
readonly lock: (
/**
* Refreshes a stored OAuth credential and saves the result. Every process and Location sharing the database
* refreshes a credential one at a time, and concurrent callers share the running refresh, including its
* failure. The row is read again under the lock: if `stale` is false, another caller already refreshed it and
* that row is returned. If the row changes while `refresh` runs, the stored row wins. Resolves undefined when
* the credential is gone.
*/
readonly refresh: (
id: ID,
operation: Effect.Effect<OAuth | undefined, unknown>,
input: {
readonly stale: (current: OAuth, now: number) => boolean
readonly refresh: (current: OAuth) => Effect.Effect<OAuth, unknown>
},
) => Effect.Effect<OAuth | undefined, unknown>
/**
* Creates a credential for an integration and returns the new record. The new credential becomes the
Expand Down Expand Up @@ -126,21 +139,46 @@ const layer = Layer.effect(
Effect.orDie,
)

return Service.of({
lock: Effect.fn("Credential.lock")((id, operation) =>
Effect.uninterruptibleMask((restore) =>
const lock = (id: ID, operation: Effect.Effect<OAuth | undefined, unknown>) =>
Effect.uninterruptibleMask((restore) =>
Effect.gen(function* () {
const current = pending.get(id)
if (current) return yield* restore(Fiber.join(current))
// The global scope owns the operation, not the first Session or Location waiting for it.
// The file lock extends this to other processes sharing the database.
const fiber = yield* operation.pipe(
flock.withLock(`credential-refresh:${id}`),
Effect.ensuring(Effect.sync(() => pending.delete(id))),
Effect.forkIn(scope),
)
pending.set(id, fiber)
return yield* restore(Fiber.join(fiber))
}),
)

const service: Interface = Service.of({
refresh: Effect.fn("Credential.refresh")((id, input) =>
lock(
id,
Effect.gen(function* () {
const current = pending.get(id)
if (current) return yield* restore(Fiber.join(current))
// The global scope owns the operation, not the first Session or Location waiting for it.
// The file lock extends this to other processes sharing the database.
const fiber = yield* operation.pipe(
flock.withLock(`credential-refresh:${id}`),
Effect.ensuring(Effect.sync(() => pending.delete(id))),
Effect.forkIn(scope),
)
pending.set(id, fiber)
return yield* restore(Fiber.join(fiber))
const current = yield* service.get(id)
if (current?.value.type !== "oauth") return undefined
if (!input.stale(current.value, yield* Clock.currentTimeMillis)) {
yield* Effect.logInfo("credential refresh skipped", { credentialID: id, reason: "already_refreshed" })
return current.value
}
const value = yield* input.refresh(current.value).pipe(Effect.timeout("30 seconds"))
const latest = yield* service.get(id)
if (latest?.value.type !== "oauth") return undefined
if (latest.value.access !== current.value.access || latest.value.refresh !== current.value.refresh) {
yield* Effect.logWarning("credential refresh discarded", {
credentialID: id,
reason: "changed_meanwhile",
})
return latest.value
}
yield* service.update(id, { value })
return value
}),
),
),
Expand Down Expand Up @@ -305,6 +343,7 @@ const layer = Layer.effect(
)
}),
})
return service
}),
)

Expand Down
22 changes: 3 additions & 19 deletions packages/core/src/integration.ts
Original file line number Diff line number Diff line change
Expand Up @@ -708,25 +708,9 @@ const layer = Layer.effect(
.get()
.integrations.get(credential.integrationID)
?.implementations.get(credential.value.methodID)?.refresh
const expiring = (value: Credential.OAuth) =>
Clock.currentTimeMillis.pipe(
Effect.map((now) => value.expires <= now + Duration.toMillis(Duration.minutes(5))),
)
if (!refresh || !(yield* expiring(credential.value))) return credential.value
return yield* authorize(
credentials.lock(
credential.id,
Effect.gen(function* () {
// Another Location may have refreshed while this caller waited.
const current = yield* credentials.get(credential.id)
if (!current || current.value.type !== "oauth") return undefined
if (!(yield* expiring(current.value))) return current.value
const value = yield* refresh(current.value).pipe(Effect.timeout("30 seconds"))
yield* credentials.update(current.id, { value })
return value
}),
),
)
if (!refresh || !Credential.expiring(credential.value, yield* Clock.currentTimeMillis))
return credential.value
return yield* authorize(credentials.refresh(credential.id, { stale: Credential.expiring, refresh }))
}),
key: Effect.fn("Integration.connection.key")(function* (input) {
const method = state
Expand Down
77 changes: 34 additions & 43 deletions packages/core/src/mcp/oauth.ts
Original file line number Diff line number Diff line change
Expand Up @@ -337,52 +337,43 @@ export const connectProvider = Effect.fnUntraced(function* (input: {
}
// Refresh tokens rotate and the row is shared across connections: only drop it while it still holds ours.
let presented = found.value.refresh
// Every process and Location sharing the database refreshes this row under one lock, and the lock is held
// until the rotated tokens are saved. A caller holding a refresh token the row no longer has gets the
// stored tokens instead of spending it: servers that rotate refresh tokens revoke the grant on reuse.
// Every process and Location sharing the database refreshes this row one at a time. A caller holding a
// refresh token the row no longer has gets the stored tokens instead of spending it: servers that rotate
// refresh tokens revoke the grant on reuse.
const refresh: Refresh = (send, url, init, token) =>
run(
credentials
.lock(
id,
Effect.gen(function* () {
const current = yield* credentials.get(id)
if (current?.value.type !== "oauth") return undefined
if (current.value.refresh !== token) {
yield* Effect.logInfo("mcp oauth refresh skipped", { credentialID: id, reason: "already_rotated" })
return current.value
}
// Bounds reading the body too, so a server that stalls after its headers cannot hold the lock.
const { response, body } = yield* Effect.tryPromise({
try: async (signal) => {
const response = await send(url, { ...init, signal })
return { response, body: await response.text() }
},
catch: (error) => error,
}).pipe(Effect.timeout("30 seconds"))
const parsed = OAuthTokensSchema.safeParse(
response.ok ? Option.getOrUndefined(decodeJson(body)) : undefined,
)
if (!parsed.success)
return yield* new RefreshRejected({ status: response.status, headers: response.headers, body })
const latest = yield* credentials.get(id)
if (latest?.value.type !== "oauth") return undefined
if (latest.value.refresh !== token) {
yield* Effect.logWarning("mcp oauth refresh discarded", { credentialID: id, reason: "rotated_meanwhile" })
return latest.value
}
// The SDK ignores an issuer in the token response; keep the one this credential was stamped with.
const issuer = latest.value.metadata?.issuer
const value = toCredential({
methodID,
serverUrl: input.config.url,
tokens: { refresh_token: token, ...parsed.data, issuer: typeof issuer === "string" ? issuer : undefined },
client: clientFromCredential(latest.value),
})
yield* credentials.update(id, { value })
return value
}),
)
.refresh(id, {
stale: (current) => current.refresh === token,
refresh: (current) =>
Effect.gen(function* () {
// The body is read here so the refresh time limit also covers a server that stalls after its headers.
const { response, body } = yield* Effect.tryPromise({
try: async (signal) => {
const response = await send(url, { ...init, signal })
return { response, body: await response.text() }
},
catch: (error) => error,
})
const parsed = OAuthTokensSchema.safeParse(
response.ok ? Option.getOrUndefined(decodeJson(body)) : undefined,
)
if (!parsed.success)
return yield* new RefreshRejected({ status: response.status, headers: response.headers, body })
// The SDK ignores an issuer in the token response; keep the one this credential was stamped with.
const issuer = current.metadata?.issuer
return toCredential({
methodID,
serverUrl: input.config.url,
tokens: {
refresh_token: token,
...parsed.data,
issuer: typeof issuer === "string" ? issuer : undefined,
},
client: clientFromCredential(current),
})
}),
})
.pipe(
Effect.map((value) =>
value
Expand Down
4 changes: 2 additions & 2 deletions packages/core/test/config/config.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -910,7 +910,7 @@ describe("Config", () => {
layer: Layer.succeed(
Credential.Service,
Credential.Service.of({
lock: () => Effect.die("unused Credential.lock"),
refresh: () => Effect.die("unused Credential.refresh"),
all: () => Effect.die("unused Credential.all"),
list: () =>
Effect.succeed([
Expand Down Expand Up @@ -1153,7 +1153,7 @@ describe("Config", () => {
layer: Layer.succeed(
Credential.Service,
Credential.Service.of({
lock: () => Effect.die("unused Credential.lock"),
refresh: () => Effect.die("unused Credential.refresh"),
all: () => Effect.die("unused Credential.all"),
list: () =>
Effect.succeed([
Expand Down
41 changes: 0 additions & 41 deletions packages/core/test/credential-lock.test.ts

This file was deleted.

Loading
Loading