fix+refactor(mcp): lifecycle tests, cancelPending fix, Effect migration (#19042)
This commit is contained in:
Kit Langton
@@ -1,7 +1,9 @@
|
||||
import path from "path"
|
||||
import z from "zod"
|
||||
import { Global } from "../global"
|
||||
import { Filesystem } from "../util/filesystem"
|
||||
import { Effect, Layer, ServiceMap } from "effect"
|
||||
import { AppFileSystem } from "@/filesystem"
|
||||
import { makeRunPromise } from "@/effect/run-service"
|
||||
|
||||
export namespace McpAuth {
|
||||
export const Tokens = z.object({
|
||||
@@ -25,106 +27,155 @@ export namespace McpAuth {
|
||||
clientInfo: ClientInfo.optional(),
|
||||
codeVerifier: z.string().optional(),
|
||||
oauthState: z.string().optional(),
|
||||
serverUrl: z.string().optional(), // Track the URL these credentials are for
|
||||
serverUrl: z.string().optional(),
|
||||
})
|
||||
export type Entry = z.infer<typeof Entry>
|
||||
|
||||
const filepath = path.join(Global.Path.data, "mcp-auth.json")
|
||||
|
||||
export async function get(mcpName: string): Promise<Entry | undefined> {
|
||||
const data = await all()
|
||||
return data[mcpName]
|
||||
export interface Interface {
|
||||
readonly all: () => Effect.Effect<Record<string, Entry>>
|
||||
readonly get: (mcpName: string) => Effect.Effect<Entry | undefined>
|
||||
readonly getForUrl: (mcpName: string, serverUrl: string) => Effect.Effect<Entry | undefined>
|
||||
readonly set: (mcpName: string, entry: Entry, serverUrl?: string) => Effect.Effect<void>
|
||||
readonly remove: (mcpName: string) => Effect.Effect<void>
|
||||
readonly updateTokens: (mcpName: string, tokens: Tokens, serverUrl?: string) => Effect.Effect<void>
|
||||
readonly updateClientInfo: (mcpName: string, clientInfo: ClientInfo, serverUrl?: string) => Effect.Effect<void>
|
||||
readonly updateCodeVerifier: (mcpName: string, codeVerifier: string) => Effect.Effect<void>
|
||||
readonly clearCodeVerifier: (mcpName: string) => Effect.Effect<void>
|
||||
readonly updateOAuthState: (mcpName: string, oauthState: string) => Effect.Effect<void>
|
||||
readonly getOAuthState: (mcpName: string) => Effect.Effect<string | undefined>
|
||||
readonly clearOAuthState: (mcpName: string) => Effect.Effect<void>
|
||||
readonly isTokenExpired: (mcpName: string) => Effect.Effect<boolean | null>
|
||||
}
|
||||
|
||||
/**
|
||||
* Get auth entry and validate it's for the correct URL.
|
||||
* Returns undefined if URL has changed (credentials are invalid).
|
||||
*/
|
||||
export async function getForUrl(mcpName: string, serverUrl: string): Promise<Entry | undefined> {
|
||||
const entry = await get(mcpName)
|
||||
if (!entry) return undefined
|
||||
export class Service extends ServiceMap.Service<Service, Interface>()("@opencode/McpAuth") {}
|
||||
|
||||
// If no serverUrl is stored, this is from an old version - consider it invalid
|
||||
if (!entry.serverUrl) return undefined
|
||||
export const layer = Layer.effect(
|
||||
Service,
|
||||
Effect.gen(function* () {
|
||||
const fs = yield* AppFileSystem.Service
|
||||
|
||||
// If URL has changed, credentials are invalid
|
||||
if (entry.serverUrl !== serverUrl) return undefined
|
||||
const all = Effect.fn("McpAuth.all")(function* () {
|
||||
return yield* fs.readJson(filepath).pipe(
|
||||
Effect.map((data) => data as Record<string, Entry>),
|
||||
Effect.catch(() => Effect.succeed({} as Record<string, Entry>)),
|
||||
)
|
||||
})
|
||||
|
||||
return entry
|
||||
}
|
||||
const get = Effect.fn("McpAuth.get")(function* (mcpName: string) {
|
||||
const data = yield* all()
|
||||
return data[mcpName]
|
||||
})
|
||||
|
||||
export async function all(): Promise<Record<string, Entry>> {
|
||||
return Filesystem.readJson<Record<string, Entry>>(filepath).catch(() => ({}))
|
||||
}
|
||||
const getForUrl = Effect.fn("McpAuth.getForUrl")(function* (mcpName: string, serverUrl: string) {
|
||||
const entry = yield* get(mcpName)
|
||||
if (!entry) return undefined
|
||||
if (!entry.serverUrl) return undefined
|
||||
if (entry.serverUrl !== serverUrl) return undefined
|
||||
return entry
|
||||
})
|
||||
|
||||
export async function set(mcpName: string, entry: Entry, serverUrl?: string): Promise<void> {
|
||||
const data = await all()
|
||||
// Always update serverUrl if provided
|
||||
if (serverUrl) {
|
||||
entry.serverUrl = serverUrl
|
||||
}
|
||||
await Filesystem.writeJson(filepath, { ...data, [mcpName]: entry }, 0o600)
|
||||
}
|
||||
const set = Effect.fn("McpAuth.set")(function* (mcpName: string, entry: Entry, serverUrl?: string) {
|
||||
const data = yield* all()
|
||||
if (serverUrl) entry.serverUrl = serverUrl
|
||||
yield* fs.writeJson(filepath, { ...data, [mcpName]: entry }, 0o600).pipe(Effect.orDie)
|
||||
})
|
||||
|
||||
export async function remove(mcpName: string): Promise<void> {
|
||||
const data = await all()
|
||||
delete data[mcpName]
|
||||
await Filesystem.writeJson(filepath, data, 0o600)
|
||||
}
|
||||
const remove = Effect.fn("McpAuth.remove")(function* (mcpName: string) {
|
||||
const data = yield* all()
|
||||
delete data[mcpName]
|
||||
yield* fs.writeJson(filepath, data, 0o600).pipe(Effect.orDie)
|
||||
})
|
||||
|
||||
export async function updateTokens(mcpName: string, tokens: Tokens, serverUrl?: string): Promise<void> {
|
||||
const entry = (await get(mcpName)) ?? {}
|
||||
entry.tokens = tokens
|
||||
await set(mcpName, entry, serverUrl)
|
||||
}
|
||||
const updateField = <K extends keyof Entry>(field: K, spanName: string) =>
|
||||
Effect.fn(`McpAuth.${spanName}`)(function* (mcpName: string, value: NonNullable<Entry[K]>, serverUrl?: string) {
|
||||
const entry = (yield* get(mcpName)) ?? {}
|
||||
entry[field] = value
|
||||
yield* set(mcpName, entry, serverUrl)
|
||||
})
|
||||
|
||||
export async function updateClientInfo(mcpName: string, clientInfo: ClientInfo, serverUrl?: string): Promise<void> {
|
||||
const entry = (await get(mcpName)) ?? {}
|
||||
entry.clientInfo = clientInfo
|
||||
await set(mcpName, entry, serverUrl)
|
||||
}
|
||||
const clearField = <K extends keyof Entry>(field: K, spanName: string) =>
|
||||
Effect.fn(`McpAuth.${spanName}`)(function* (mcpName: string) {
|
||||
const entry = yield* get(mcpName)
|
||||
if (entry) {
|
||||
delete entry[field]
|
||||
yield* set(mcpName, entry)
|
||||
}
|
||||
})
|
||||
|
||||
export async function updateCodeVerifier(mcpName: string, codeVerifier: string): Promise<void> {
|
||||
const entry = (await get(mcpName)) ?? {}
|
||||
entry.codeVerifier = codeVerifier
|
||||
await set(mcpName, entry)
|
||||
}
|
||||
const updateTokens = updateField("tokens", "updateTokens")
|
||||
const updateClientInfo = updateField("clientInfo", "updateClientInfo")
|
||||
const updateCodeVerifier = updateField("codeVerifier", "updateCodeVerifier")
|
||||
const updateOAuthState = updateField("oauthState", "updateOAuthState")
|
||||
const clearCodeVerifier = clearField("codeVerifier", "clearCodeVerifier")
|
||||
const clearOAuthState = clearField("oauthState", "clearOAuthState")
|
||||
|
||||
export async function clearCodeVerifier(mcpName: string): Promise<void> {
|
||||
const entry = await get(mcpName)
|
||||
if (entry) {
|
||||
delete entry.codeVerifier
|
||||
await set(mcpName, entry)
|
||||
}
|
||||
}
|
||||
const getOAuthState = Effect.fn("McpAuth.getOAuthState")(function* (mcpName: string) {
|
||||
const entry = yield* get(mcpName)
|
||||
return entry?.oauthState
|
||||
})
|
||||
|
||||
export async function updateOAuthState(mcpName: string, oauthState: string): Promise<void> {
|
||||
const entry = (await get(mcpName)) ?? {}
|
||||
entry.oauthState = oauthState
|
||||
await set(mcpName, entry)
|
||||
}
|
||||
const isTokenExpired = Effect.fn("McpAuth.isTokenExpired")(function* (mcpName: string) {
|
||||
const entry = yield* get(mcpName)
|
||||
if (!entry?.tokens) return null
|
||||
if (!entry.tokens.expiresAt) return false
|
||||
return entry.tokens.expiresAt < Date.now() / 1000
|
||||
})
|
||||
|
||||
export async function getOAuthState(mcpName: string): Promise<string | undefined> {
|
||||
const entry = await get(mcpName)
|
||||
return entry?.oauthState
|
||||
}
|
||||
return Service.of({
|
||||
all,
|
||||
get,
|
||||
getForUrl,
|
||||
set,
|
||||
remove,
|
||||
updateTokens,
|
||||
updateClientInfo,
|
||||
updateCodeVerifier,
|
||||
clearCodeVerifier,
|
||||
updateOAuthState,
|
||||
getOAuthState,
|
||||
clearOAuthState,
|
||||
isTokenExpired,
|
||||
})
|
||||
}),
|
||||
)
|
||||
|
||||
export async function clearOAuthState(mcpName: string): Promise<void> {
|
||||
const entry = await get(mcpName)
|
||||
if (entry) {
|
||||
delete entry.oauthState
|
||||
await set(mcpName, entry)
|
||||
}
|
||||
}
|
||||
const defaultLayer = layer.pipe(Layer.provide(AppFileSystem.defaultLayer))
|
||||
|
||||
/**
|
||||
* Check if stored tokens are expired.
|
||||
* Returns null if no tokens exist, false if no expiry or not expired, true if expired.
|
||||
*/
|
||||
export async function isTokenExpired(mcpName: string): Promise<boolean | null> {
|
||||
const entry = await get(mcpName)
|
||||
if (!entry?.tokens) return null
|
||||
if (!entry.tokens.expiresAt) return false
|
||||
return entry.tokens.expiresAt < Date.now() / 1000
|
||||
}
|
||||
const runPromise = makeRunPromise(Service, defaultLayer)
|
||||
|
||||
// Async facades for backward compat (used by McpOAuthProvider, CLI)
|
||||
|
||||
export const get = async (mcpName: string) => runPromise((svc) => svc.get(mcpName))
|
||||
|
||||
export const getForUrl = async (mcpName: string, serverUrl: string) =>
|
||||
runPromise((svc) => svc.getForUrl(mcpName, serverUrl))
|
||||
|
||||
export const all = async () => runPromise((svc) => svc.all())
|
||||
|
||||
export const set = async (mcpName: string, entry: Entry, serverUrl?: string) =>
|
||||
runPromise((svc) => svc.set(mcpName, entry, serverUrl))
|
||||
|
||||
export const remove = async (mcpName: string) => runPromise((svc) => svc.remove(mcpName))
|
||||
|
||||
export const updateTokens = async (mcpName: string, tokens: Tokens, serverUrl?: string) =>
|
||||
runPromise((svc) => svc.updateTokens(mcpName, tokens, serverUrl))
|
||||
|
||||
export const updateClientInfo = async (mcpName: string, clientInfo: ClientInfo, serverUrl?: string) =>
|
||||
runPromise((svc) => svc.updateClientInfo(mcpName, clientInfo, serverUrl))
|
||||
|
||||
export const updateCodeVerifier = async (mcpName: string, codeVerifier: string) =>
|
||||
runPromise((svc) => svc.updateCodeVerifier(mcpName, codeVerifier))
|
||||
|
||||
export const clearCodeVerifier = async (mcpName: string) => runPromise((svc) => svc.clearCodeVerifier(mcpName))
|
||||
|
||||
export const updateOAuthState = async (mcpName: string, oauthState: string) =>
|
||||
runPromise((svc) => svc.updateOAuthState(mcpName, oauthState))
|
||||
|
||||
export const getOAuthState = async (mcpName: string) => runPromise((svc) => svc.getOAuthState(mcpName))
|
||||
|
||||
export const clearOAuthState = async (mcpName: string) => runPromise((svc) => svc.clearOAuthState(mcpName))
|
||||
|
||||
export const isTokenExpired = async (mcpName: string) => runPromise((svc) => svc.isTokenExpired(mcpName))
|
||||
}
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -54,6 +54,9 @@ interface PendingAuth {
|
||||
export namespace McpOAuthCallback {
|
||||
let server: ReturnType<typeof Bun.serve> | undefined
|
||||
const pendingAuths = new Map<string, PendingAuth>()
|
||||
// Reverse index: mcpName → oauthState, so cancelPending(mcpName) can
|
||||
// find the right entry in pendingAuths (which is keyed by oauthState).
|
||||
const mcpNameToState = new Map<string, string>()
|
||||
|
||||
const CALLBACK_TIMEOUT_MS = 5 * 60 * 1000 // 5 minutes
|
||||
|
||||
@@ -98,6 +101,12 @@ export namespace McpOAuthCallback {
|
||||
const pending = pendingAuths.get(state)!
|
||||
clearTimeout(pending.timeout)
|
||||
pendingAuths.delete(state)
|
||||
for (const [name, s] of mcpNameToState) {
|
||||
if (s === state) {
|
||||
mcpNameToState.delete(name)
|
||||
break
|
||||
}
|
||||
}
|
||||
pending.reject(new Error(errorMsg))
|
||||
}
|
||||
return new Response(HTML_ERROR(errorMsg), {
|
||||
@@ -126,6 +135,13 @@ export namespace McpOAuthCallback {
|
||||
|
||||
clearTimeout(pending.timeout)
|
||||
pendingAuths.delete(state)
|
||||
// Clean up reverse index
|
||||
for (const [name, s] of mcpNameToState) {
|
||||
if (s === state) {
|
||||
mcpNameToState.delete(name)
|
||||
break
|
||||
}
|
||||
}
|
||||
pending.resolve(code)
|
||||
|
||||
return new Response(HTML_SUCCESS, {
|
||||
@@ -137,11 +153,13 @@ export namespace McpOAuthCallback {
|
||||
log.info("oauth callback server started", { port: OAUTH_CALLBACK_PORT })
|
||||
}
|
||||
|
||||
export function waitForCallback(oauthState: string): Promise<string> {
|
||||
export function waitForCallback(oauthState: string, mcpName?: string): Promise<string> {
|
||||
if (mcpName) mcpNameToState.set(mcpName, oauthState)
|
||||
return new Promise((resolve, reject) => {
|
||||
const timeout = setTimeout(() => {
|
||||
if (pendingAuths.has(oauthState)) {
|
||||
pendingAuths.delete(oauthState)
|
||||
if (mcpName) mcpNameToState.delete(mcpName)
|
||||
reject(new Error("OAuth callback timeout - authorization took too long"))
|
||||
}
|
||||
}, CALLBACK_TIMEOUT_MS)
|
||||
@@ -151,10 +169,14 @@ export namespace McpOAuthCallback {
|
||||
}
|
||||
|
||||
export function cancelPending(mcpName: string): void {
|
||||
const pending = pendingAuths.get(mcpName)
|
||||
// Look up the oauthState for this mcpName via the reverse index
|
||||
const oauthState = mcpNameToState.get(mcpName)
|
||||
const key = oauthState ?? mcpName
|
||||
const pending = pendingAuths.get(key)
|
||||
if (pending) {
|
||||
clearTimeout(pending.timeout)
|
||||
pendingAuths.delete(mcpName)
|
||||
pendingAuths.delete(key)
|
||||
mcpNameToState.delete(mcpName)
|
||||
pending.reject(new Error("Authorization cancelled"))
|
||||
}
|
||||
}
|
||||
@@ -184,6 +206,7 @@ export namespace McpOAuthCallback {
|
||||
pending.reject(new Error("OAuth callback server stopped"))
|
||||
}
|
||||
pendingAuths.clear()
|
||||
mcpNameToState.clear()
|
||||
}
|
||||
|
||||
export function isRunning(): boolean {
|
||||
|
||||
Reference in New Issue
Block a user