145 lines · 4.6 KB · TypeScript
RawDownload
import path from "path"import z from "zod"import { Global } from "../global"import { JsonStore } from "../util/jsonstore" export namespace McpAuth {  export const Tokens = z.object({    accessToken: z.string(),    refreshToken: z.string().optional(),    expiresAt: z.number().optional(),    scope: z.string().optional(),  })  export type Tokens = z.infer<typeof Tokens>   export const ClientInfo = z.object({    clientId: z.string(),    clientSecret: z.string().optional(),    clientIdIssuedAt: z.number().optional(),    clientSecretExpiresAt: z.number().optional(),  })  export type ClientInfo = z.infer<typeof ClientInfo>   export const Entry = z.object({    tokens: Tokens.optional(),    clientInfo: ClientInfo.optional(),    codeVerifier: z.string().optional(),    oauthState: z.string().optional(),    serverUrl: z.string().optional(), // Track the URL these credentials are for  })  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]  }   /**   * 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     // If no serverUrl is stored, this is from an old version - consider it invalid    if (!entry.serverUrl) return undefined     // If URL has changed, credentials are invalid    if (entry.serverUrl !== serverUrl) return undefined     return entry  }   export async function all(): Promise<Record<string, Entry>> {    return (await JsonStore.read(filepath)) as Record<string, Entry>  }   /** Atomic read-modify-write of one server's entry. Serialized behind the   *  store lock so concurrent updates can't drop each other's writes. */  async function update(mcpName: string, fn: (entry: Entry) => void, serverUrl?: string): Promise<void> {    await JsonStore.update(filepath, (data) => {      const entry = (data[mcpName] ?? {}) as Entry      fn(entry)      // Always update serverUrl if provided      if (serverUrl) entry.serverUrl = serverUrl      data[mcpName] = entry    })  }   export async function set(mcpName: string, entry: Entry, serverUrl?: string): Promise<void> {    // Always update serverUrl if provided    if (serverUrl) entry.serverUrl = serverUrl    await JsonStore.update(filepath, (data) => ({ ...data, [mcpName]: entry }))  }   export async function remove(mcpName: string): Promise<void> {    await JsonStore.update(filepath, (data) => {      delete data[mcpName]    })  }   export async function updateTokens(mcpName: string, tokens: Tokens, serverUrl?: string): Promise<void> {    await update(      mcpName,      (entry) => {        entry.tokens = tokens      },      serverUrl,    )  }   export async function updateClientInfo(mcpName: string, clientInfo: ClientInfo, serverUrl?: string): Promise<void> {    await update(      mcpName,      (entry) => {        entry.clientInfo = clientInfo      },      serverUrl,    )  }   export async function updateCodeVerifier(mcpName: string, codeVerifier: string): Promise<void> {    await update(mcpName, (entry) => {      entry.codeVerifier = codeVerifier    })  }   export async function clearCodeVerifier(mcpName: string): Promise<void> {    await JsonStore.update(filepath, (data) => {      const entry = data[mcpName] as Entry | undefined      if (entry) delete entry.codeVerifier    })  }   export async function updateOAuthState(mcpName: string, oauthState: string): Promise<void> {    await update(mcpName, (entry) => {      entry.oauthState = oauthState    })  }   export async function getOAuthState(mcpName: string): Promise<string | undefined> {    const entry = await get(mcpName)    return entry?.oauthState  }   export async function clearOAuthState(mcpName: string): Promise<void> {    await JsonStore.update(filepath, (data) => {      const entry = data[mcpName] as Entry | undefined      if (entry) delete entry.oauthState    })  }   /**   * 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  }}