@agent-native/core 0.105.0 → 0.106.1

package/corpus/core/src/mcp-client/oauth-client.ts

@@ -0,0 +1,462 @@
+/**
+ * OAuth 2.1 client support for remote MCP servers.
+ *
+ * MCP servers advertise their OAuth endpoints through the standard protected
+ * resource and authorization-server metadata documents. The SDK handles the
+ * protocol details; this module owns the framework-specific encrypted storage
+ * and refresh boundary.
+ */
+
+import {
+  auth,
+  refreshAuthorization,
+  type OAuthClientProvider,
+} from "@modelcontextprotocol/sdk/client/auth.js";
+import type {
+  AuthorizationServerMetadata,
+  OAuthClientInformationMixed,
+  OAuthClientMetadata,
+  OAuthProtectedResourceMetadata,
+  OAuthTokens,
+} from "@modelcontextprotocol/sdk/shared/auth.js";
+
+import {
+  deleteOAuthTokens,
+  getOAuthTokens,
+  saveOAuthTokens,
+} from "../oauth-tokens/store.js";
+import { validateRemoteUrl } from "./remote-url.js";
+
+const TOKEN_EXPIRY_SKEW_MS = 60_000;
+const MAX_OAUTH_REDIRECTS = 5;
+
+type GuardedFetch = (
+  url: string | URL,
+  init?: RequestInit,
+) => Promise<Response>;
+
+function checkedRemoteUrl(value: string | URL, label: string): URL {
+  const result = validateRemoteUrl(String(value));
+  if (!result.ok || !result.url) {
+    throw new Error(
+      `MCP OAuth ${label} is not allowed: ${result.error ?? "invalid URL"}`,
+    );
+  }
+  return result.url;
+}
+
+function guardedOAuthFetch(): GuardedFetch {
+  return async (url, init) => {
+    let currentUrl = checkedRemoteUrl(url, "request");
+    let currentInit: RequestInit = { ...init, redirect: "manual" };
+
+    for (let redirectCount = 0; ; redirectCount += 1) {
+      const response = await fetch(currentUrl, currentInit);
+      if (![301, 302, 303, 307, 308].includes(response.status)) {
+        return response;
+      }
+      if (redirectCount >= MAX_OAUTH_REDIRECTS) {
+        throw new Error("MCP OAuth redirect limit exceeded.");
+      }
+      const location = response.headers.get("location");
+      if (!location) throw new Error("MCP OAuth redirect is missing a target.");
+
+      let nextUrl: URL;
+      try {
+        nextUrl = checkedRemoteUrl(
+          new URL(location, currentUrl),
+          "redirect target",
+        );
+      } catch {
+        throw new Error("MCP OAuth redirect target is not allowed.");
+      }
+      await response.body?.cancel().catch(() => undefined);
+
+      const nextHeaders = new Headers(currentInit.headers);
+      if (nextUrl.origin !== currentUrl.origin) {
+        nextHeaders.delete("authorization");
+        nextHeaders.delete("cookie");
+        nextHeaders.delete("proxy-authorization");
+      }
+      const nextInit: RequestInit = {
+        ...currentInit,
+        headers: nextHeaders,
+      };
+      if (
+        response.status === 303 ||
+        ((response.status === 301 || response.status === 302) &&
+          (currentInit.method ?? "GET").toUpperCase() === "POST")
+      ) {
+        nextInit.method = "GET";
+        delete nextInit.body;
+        nextHeaders.delete("content-length");
+        nextHeaders.delete("content-type");
+      }
+      currentUrl = nextUrl;
+      currentInit = nextInit;
+    }
+  };
+}
+
+function validateDiscoveryUrls(state: {
+  authorizationServerUrl: string;
+  authorizationServerMetadata?: AuthorizationServerMetadata;
+  resourceMetadata?: OAuthProtectedResourceMetadata;
+  resourceMetadataUrl?: string;
+}): void {
+  checkedRemoteUrl(state.authorizationServerUrl, "authorization server");
+  if (state.resourceMetadataUrl) {
+    checkedRemoteUrl(state.resourceMetadataUrl, "resource metadata");
+  }
+  for (const url of state.resourceMetadata?.authorization_servers ?? []) {
+    checkedRemoteUrl(url, "discovered authorization server");
+  }
+  for (const field of [
+    "authorization_endpoint",
+    "token_endpoint",
+    "registration_endpoint",
+    "introspection_endpoint",
+    "revocation_endpoint",
+  ] as const) {
+    const url = (
+      state.authorizationServerMetadata as Record<string, unknown> | undefined
+    )?.[field];
+    if (typeof url === "string") {
+      checkedRemoteUrl(url, `OAuth ${field}`);
+    }
+  }
+}
+
+function credentialOwner(options: { scope: "user" | "org"; scopeId: string }) {
+  return `${options.scope}:${options.scopeId}`;
+}
+
+export interface McpOAuthDiscoveryState {
+  authorizationServerUrl: string;
+  authorizationServerMetadata?: AuthorizationServerMetadata;
+  resourceMetadata?: OAuthProtectedResourceMetadata;
+  resourceMetadataUrl?: string;
+}
+
+export interface McpOAuthCredentialBundle {
+  serverUrl: string;
+  clientInformation: OAuthClientInformationMixed;
+  discoveryState?: McpOAuthDiscoveryState;
+  tokens: OAuthTokens;
+  tokenExpiresAt?: number;
+}
+
+export interface McpOAuthStartResult {
+  authorizationUrl: URL;
+  codeVerifier: string;
+  state: string;
+  clientInformation: OAuthClientInformationMixed;
+  discoveryState?: McpOAuthDiscoveryState;
+}
+
+export interface McpOAuthCallbackResult {
+  credentials: McpOAuthCredentialBundle;
+}
+
+export interface McpOAuthProviderOptions {
+  serverUrl: string;
+  redirectUrl: string;
+  state: string;
+  clientInformation?: OAuthClientInformationMixed;
+  codeVerifier?: string;
+  discoveryState?: McpOAuthDiscoveryState;
+}
+
+/**
+ * A small adapter around the MCP SDK's OAuth provider interface. The route
+ * stores the adapter's state in an encrypted, short-lived browser cookie; the
+ * durable credential bundle is written only after the callback succeeds.
+ */
+export class McpOAuthClientProvider implements OAuthClientProvider {
+  private readonly redirectUrlValue: string;
+  private readonly stateValue: string;
+  private readonly metadata: OAuthClientMetadata;
+  private clientInfo?: OAuthClientInformationMixed;
+  private savedTokens?: OAuthTokens;
+  private savedCodeVerifier?: string;
+  private savedDiscovery?: McpOAuthDiscoveryState;
+  private authorizationUrl?: URL;
+
+  constructor(options: McpOAuthProviderOptions) {
+    this.redirectUrlValue = options.redirectUrl;
+    this.stateValue = options.state;
+    this.clientInfo = options.clientInformation;
+    this.savedCodeVerifier = options.codeVerifier;
+    this.savedDiscovery = options.discoveryState;
+    this.metadata = {
+      redirect_uris: [options.redirectUrl],
+      token_endpoint_auth_method: "none",
+      grant_types: ["authorization_code"],
+      response_types: ["code"],
+      client_name: "Agent Native MCP connector",
+    };
+  }
+
+  get redirectUrl(): string {
+    return this.redirectUrlValue;
+  }
+
+  get clientMetadata(): OAuthClientMetadata {
+    return this.metadata;
+  }
+
+  state(): string {
+    return this.stateValue;
+  }
+
+  clientInformation(): OAuthClientInformationMixed | undefined {
+    return this.clientInfo;
+  }
+
+  saveClientInformation(info: OAuthClientInformationMixed): void {
+    this.clientInfo = info;
+  }
+
+  tokens(): OAuthTokens | undefined {
+    return this.savedTokens;
+  }
+
+  saveTokens(tokens: OAuthTokens): void {
+    this.savedTokens = tokens;
+  }
+
+  redirectToAuthorization(authorizationUrl: URL): void {
+    this.authorizationUrl = checkedRemoteUrl(
+      authorizationUrl,
+      "authorization redirect",
+    );
+  }
+
+  saveCodeVerifier(codeVerifier: string): void {
+    this.savedCodeVerifier = codeVerifier;
+  }
+
+  codeVerifier(): string {
+    if (!this.savedCodeVerifier) {
+      throw new Error("MCP OAuth code verifier is missing");
+    }
+    return this.savedCodeVerifier;
+  }
+
+  saveDiscoveryState(state: {
+    authorizationServerUrl: string;
+    authorizationServerMetadata?: AuthorizationServerMetadata;
+    resourceMetadata?: OAuthProtectedResourceMetadata;
+    resourceMetadataUrl?: string;
+  }): void {
+    validateDiscoveryUrls(state);
+    this.savedDiscovery = state;
+  }
+
+  discoveryState(): McpOAuthDiscoveryState | undefined {
+    return this.savedDiscovery;
+  }
+
+  get authorizationRedirect(): URL | undefined {
+    return this.authorizationUrl;
+  }
+
+  get savedCodeVerifierValue(): string | undefined {
+    return this.savedCodeVerifier;
+  }
+
+  get savedTokensValue(): OAuthTokens | undefined {
+    return this.savedTokens;
+  }
+
+  get savedClientInformation(): OAuthClientInformationMixed | undefined {
+    return this.clientInfo;
+  }
+}
+
+export async function startMcpOAuthAuthorization(
+  options: McpOAuthProviderOptions & { scope?: string },
+): Promise<McpOAuthStartResult> {
+  checkedRemoteUrl(options.serverUrl, "server");
+  const provider = new McpOAuthClientProvider(options);
+  const result = await auth(provider, {
+    serverUrl: options.serverUrl,
+    scope: options.scope,
+    fetchFn: guardedOAuthFetch(),
+  });
+  if (result !== "REDIRECT" || !provider.authorizationRedirect) {
+    throw new Error("MCP server did not start an interactive OAuth flow");
+  }
+  const clientInformation = provider.savedClientInformation;
+  if (!clientInformation) {
+    throw new Error("MCP OAuth client registration did not complete");
+  }
+  return {
+    authorizationUrl: provider.authorizationRedirect,
+    codeVerifier: provider.savedCodeVerifierValue ?? provider.codeVerifier(),
+    state: options.state,
+    clientInformation,
+    discoveryState: provider.discoveryState(),
+  };
+}
+
+export async function finishMcpOAuthAuthorization(
+  options: McpOAuthProviderOptions & { authorizationCode: string },
+): Promise<McpOAuthCallbackResult> {
+  checkedRemoteUrl(options.serverUrl, "server");
+  const provider = new McpOAuthClientProvider(options);
+  const result = await auth(provider, {
+    serverUrl: options.serverUrl,
+    authorizationCode: options.authorizationCode,
+    fetchFn: guardedOAuthFetch(),
+  });
+  if (result !== "AUTHORIZED" || !provider.savedTokensValue) {
+    throw new Error("MCP OAuth token exchange did not complete");
+  }
+  const tokens = provider.savedTokensValue;
+  return {
+    credentials: {
+      serverUrl: options.serverUrl,
+      clientInformation: provider.savedClientInformation!,
+      discoveryState: provider.discoveryState(),
+      tokens,
+      tokenExpiresAt: tokenExpiresAt(tokens),
+    },
+  };
+}
+
+export function tokenExpiresAt(tokens: OAuthTokens): number | undefined {
+  const expiresIn = Number(tokens.expires_in);
+  if (!Number.isFinite(expiresIn) || expiresIn <= 0) return undefined;
+  return Date.now() + expiresIn * 1_000;
+}
+
+export async function saveMcpOAuthCredentials(options: {
+  key: string;
+  scope: "user" | "org";
+  scopeId: string;
+  credentials: McpOAuthCredentialBundle;
+}): Promise<void> {
+  checkedRemoteUrl(options.credentials.serverUrl, "server");
+  if (options.credentials.discoveryState) {
+    validateDiscoveryUrls(options.credentials.discoveryState);
+  }
+  await saveOAuthTokens(
+    "mcp",
+    options.key,
+    options.credentials as unknown as Record<string, unknown>,
+    `${options.scope}:${options.scopeId}`,
+  );
+}
+
+export async function readMcpOAuthCredentials(options: {
+  key: string;
+  scope: "user" | "org";
+  scopeId: string;
+}): Promise<McpOAuthCredentialBundle | null> {
+  const stored = await getOAuthTokens(
+    "mcp",
+    options.key,
+    credentialOwner(options),
+  );
+  if (!stored) return null;
+  const parsed = stored as Partial<McpOAuthCredentialBundle>;
+  if (
+    typeof parsed.serverUrl !== "string" ||
+    !parsed.clientInformation ||
+    !parsed.tokens ||
+    typeof parsed.tokens.access_token !== "string"
+  ) {
+    return null;
+  }
+  if (!validateRemoteUrl(parsed.serverUrl).ok) return null;
+  if (parsed.discoveryState) {
+    try {
+      validateDiscoveryUrls(parsed.discoveryState);
+    } catch {
+      return null;
+    }
+  }
+  return parsed as McpOAuthCredentialBundle;
+}
+
+export async function deleteMcpOAuthCredentials(options: {
+  key: string;
+  scope: "user" | "org";
+  scopeId: string;
+}): Promise<boolean> {
+  return (
+    (await deleteOAuthTokens("mcp", options.key, credentialOwner(options))) > 0
+  );
+}
+
+/**
+ * Resolve an access token for the MCP manager. Refreshing happens only when a
+ * token is near expiry, so ordinary manager reconfiguration does not perform
+ * a network request for every connector.
+ */
+export async function getMcpOAuthAccessToken(options: {
+  key: string;
+  scope: "user" | "org";
+  scopeId: string;
+  serverUrl: string;
+}): Promise<string | null> {
+  if (!validateRemoteUrl(options.serverUrl).ok) return null;
+  const credentials = await readMcpOAuthCredentials(options);
+  if (!credentials || credentials.serverUrl !== options.serverUrl) return null;
+
+  const accessToken = credentials.tokens.access_token;
+  if (
+    typeof credentials.tokenExpiresAt !== "number" ||
+    credentials.tokenExpiresAt - Date.now() > TOKEN_EXPIRY_SKEW_MS
+  ) {
+    return accessToken;
+  }
+
+  const refreshToken = credentials.tokens.refresh_token;
+  const discovery = credentials.discoveryState;
+  if (!refreshToken || !discovery?.authorizationServerUrl) {
+    return accessToken;
+  }
+
+  try {
+    const resource = discovery.resourceMetadata?.resource
+      ? checkedRemoteUrl(discovery.resourceMetadata.resource, "resource")
+      : undefined;
+    const authorizationServerUrl = checkedRemoteUrl(
+      discovery.authorizationServerUrl,
+      "authorization server",
+    );
+    const refreshed = await refreshAuthorization(authorizationServerUrl, {
+      metadata: discovery.authorizationServerMetadata,
+      clientInformation: credentials.clientInformation,
+      refreshToken,
+      resource,
+      fetchFn: guardedOAuthFetch(),
+    });
+    const nextTokens: OAuthTokens = {
+      ...credentials.tokens,
+      ...refreshed,
+      ...(refreshed.refresh_token
+        ? { refresh_token: refreshed.refresh_token }
+        : credentials.tokens.refresh_token
+          ? { refresh_token: credentials.tokens.refresh_token }
+          : {}),
+    };
+    const next: McpOAuthCredentialBundle = {
+      ...credentials,
+      tokens: nextTokens,
+      tokenExpiresAt: tokenExpiresAt(nextTokens),
+    };
+    await saveMcpOAuthCredentials({
+      ...options,
+      credentials: next,
+    });
+    return nextTokens.access_token;
+  } catch {
+    // Keep the old token so the MCP request can return the provider's normal
+    // auth error. A single expired connector must not remove every server from
+    // the process-wide manager during a config refresh.
+    return accessToken;
+  }
+}