fix(mcp): atomically replace pending OAuth connections

packages/core/src/mcp/connection-api.test.ts

@@ -36,6 +36,10 @@ import {
 } from "./connection-api.js";
 import { MCP_REFRESH_TTL_MS, MCP_TOKEN_TTL_MS } from "./token-expiry.js";
 import { ensureMcpGranteeRegistered } from "./builder-registration.js";
+import type {
+  McpConnectionStore,
+  McpOAuthAuthorizationStore,
+} from "./types.js";
 
 const PUBLIC_ORIGIN = "https://example-session.relay.test";
 const REDIRECT_URI = "https://claude.ai/api/mcp/auth_callback";
@@ -1212,10 +1216,8 @@ describe("mcp/connection-api superseded OAuth authorizations", () => {
 
   async function authorize(
     stores: {
-      connectionStore: ReturnType<typeof createInMemoryMcpConnectionStore>;
-      authorizationStore: ReturnType<
-        typeof createInMemoryMcpOAuthAuthorizationStore
-      >;
+      connectionStore: McpConnectionStore;
+      authorizationStore: McpOAuthAuthorizationStore;
     },
     clientId: string,
     now = START,
@@ -1258,6 +1260,87 @@ describe("mcp/connection-api superseded OAuth authorizations", () => {
     ).toBe("pending");
   });
 
+  it("keeps exactly one pending connection after concurrent registrations", async () => {
+    const stores = newStores();
+    const registrations = await Promise.all(
+      Array.from({ length: 100 }, () => authorize(stores, "same-client")),
+    );
+
+    const connections = await stores.connectionStore.list();
+    expect(
+      connections.filter(
+        (connection) =>
+          connection.clientId === "same-client" &&
+          connection.status === "pending",
+      ),
+    ).toHaveLength(1);
+    expect(connections).toHaveLength(100);
+    expect(registrations).toHaveLength(100);
+  });
+
+  it("preserves an approval that races with a replacement", async () => {
+    const stores = newStores();
+    const first = await authorize(stores, "client-a");
+    const originalList = stores.connectionStore.list.bind(
+      stores.connectionStore,
+    );
+    let replacementStarted!: () => void;
+    let continueReplacement!: () => void;
+    const started = new Promise<void>((resolve) => {
+      replacementStarted = resolve;
+    });
+    const gate = new Promise<void>((resolve) => {
+      continueReplacement = resolve;
+    });
+    const fallbackStore: McpConnectionStore = {
+      ...stores.connectionStore,
+      replacePendingOAuthConnections: undefined,
+      list: async () => {
+        const snapshot = await originalList();
+        replacementStarted();
+        await gate;
+        return snapshot;
+      },
+    };
+    const raceStores = { ...stores, connectionStore: fallbackStore };
+
+    const replacing = authorize(raceStores, "client-a");
+    await started;
+    await approveMcpOAuthAuthorization(
+      {
+        authorizationId: first.authorizationId,
+        grants: [{ grantId: "grant-1", scopes: ["chatgpt.history"] }],
+      },
+      { ...raceStores, now: () => START },
+    );
+    continueReplacement();
+    await replacing;
+
+    expect(
+      await stores.connectionStore.getById(first.connectionId),
+    ).toMatchObject({
+      status: "approved",
+      grants: [{ grantId: "grant-1", scopes: ["chatgpt.history"] }],
+    });
+  });
+
+  it("creates one pending connection and authorization for a single registration", async () => {
+    const stores = newStores();
+    const created = await authorize(stores, "client-a");
+
+    expect(await stores.connectionStore.list()).toMatchObject([
+      { id: created.connectionId, clientId: "client-a", status: "pending" },
+    ]);
+    expect(
+      await stores.authorizationStore.getById(created.authorizationId),
+    ).toMatchObject({
+      id: created.authorizationId,
+      connectionId: created.connectionId,
+      clientId: "client-a",
+      status: "pending",
+    });
+  });
+
   it("signs no grant for a superseded authorization", async () => {
     const stores = newStores();
     const first = await authorize(stores, "client-a");

packages/core/src/mcp/connection-api.ts

@@ -58,6 +58,10 @@ import {
 
 const TOKEN_BYTES = 32;
 const OAUTH_AUTHORIZATION_TTL_MS = 10 * 60 * 1000;
+const pendingOAuthReplacements = new WeakMap<
+  McpConnectionStore,
+  Promise<void>
+>();
 
 /**
  * Default lifetime of a grant the server signs itself for an MCP connection
@@ -188,6 +192,15 @@ export async function createMcpConnection(
   input: CreateMcpConnectionInput,
   options: CreateMcpConnectionOptions,
 ): Promise<CreateMcpConnectionOutput> {
+  const { record, output } = await buildMcpConnection(input, options);
+  await options.store.create(record);
+  return output;
+}
+
+async function buildMcpConnection(
+  input: CreateMcpConnectionInput,
+  options: Pick<CreateMcpConnectionOptions, "publicOrigin" | "now">,
+): Promise<{ record: McpConnectionRecord; output: CreateMcpConnectionOutput }> {
   const grantee = generateMcpGrantee();
   const token = randomToken();
   const tokenHash = await hashConnectionToken(token);
@@ -209,14 +222,15 @@ export async function createMcpConnection(
     ...(input.clientId ? { clientId: input.clientId } : {}),
   };
 
-  await options.store.create(record);
-
   return {
-    connectionId: id,
-    granteeAddress: grantee.key.address,
-    connectionToken: token,
-    mcpUrl: buildMcpUrl(options.publicOrigin, token),
-    createdAt,
+    record,
+    output: {
+      connectionId: id,
+      granteeAddress: grantee.key.address,
+      connectionToken: token,
+      mcpUrl: buildMcpUrl(options.publicOrigin, token),
+      createdAt,
+    },
   };
 }
 
@@ -954,19 +968,23 @@ export async function createMcpOAuthAuthorization(
   const expiresAt = new Date(
     now.getTime() + OAUTH_AUTHORIZATION_TTL_MS,
   ).toISOString();
-  await revokeStalePending(input.clientId, now, options.connectionStore);
-  const created = await createMcpConnection(
+  const { record: connection, output: created } = await buildMcpConnection(
     {
       displayName: input.clientName ?? DEFAULT_MCP_CLIENT_NAME,
       ingress: input.ingress,
       clientId: input.clientId,
     },
     {
-      store: options.connectionStore,
       publicOrigin: options.publicOrigin,
       now: () => now,
     },
   );
+  await replacePendingOAuthConnections(
+    options.connectionStore,
+    input.clientId,
+    now.getTime() - OAUTH_AUTHORIZATION_TTL_MS,
+    connection,
+  );
   const authorizationId = randomId();
   const record: McpOAuthAuthorizationRecord = {
     id: authorizationId,
@@ -994,27 +1012,52 @@ export async function createMcpOAuthAuthorization(
   };
 }
 
-/**
- * Revoke never-approved OAuth connections a new authorize makes moot: the
- * same client's earlier tries (a retry or restart replaces them) and any
- * whose authorization expired. Owner-made connections carry no clientId and
- * are kept. Runs before the new connection exists, so two racing tries leave
- * an extra row at worst, never none.
- */
-async function revokeStalePending(
-  clientId: string,
-  now: Date,
+async function replacePendingOAuthConnections(
   store: McpConnectionStore,
+  clientId: string,
+  expiredBefore: number,
+  replacement: McpConnectionRecord,
 ): Promise<void> {
-  const cutoff = now.getTime() - OAUTH_AUTHORIZATION_TTL_MS;
-  for (const connection of await store.list()) {
-    if (connection.status !== "pending" || !connection.clientId) continue;
-
-    const expired = Date.parse(connection.createdAt) <= cutoff;
-    if (connection.clientId !== clientId && !expired) continue;
-
-    await revokeMcpConnection(connection.id, { store, now: () => now });
+  if (store.replacePendingOAuthConnections) {
+    await store.replacePendingOAuthConnections(
+      clientId,
+      expiredBefore,
+      replacement,
+    );
+    return;
   }
+
+  const previous = pendingOAuthReplacements.get(store) ?? Promise.resolve();
+  const operation = previous.then(async () => {
+    for (const connection of await store.list()) {
+      const expired = Date.parse(connection.createdAt) <= expiredBefore;
+      if (
+        connection.status !== "pending" ||
+        !connection.clientId ||
+        (connection.clientId !== clientId && !expired)
+      ) {
+        continue;
+      }
+      await store.mutate(connection.id, (latest) =>
+        latest.status === "pending"
+          ? {
+              ...latest,
+              status: "revoked",
+              revokedAt: replacement.createdAt,
+            }
+          : latest,
+      );
+    }
+    await store.create(replacement);
+  });
+  pendingOAuthReplacements.set(
+    store,
+    operation.then(
+      () => undefined,
+      () => undefined,
+    ),
+  );
+  await operation;
 }
 
 /** A superseded or removed authorization has a revoked connection; refuse it. */

packages/core/src/mcp/store.ts

@@ -31,6 +31,35 @@ export function createInMemoryMcpConnectionStore(): McpConnectionStore {
       byTokenHash.set(record.tokenHash, record.id);
     },
 
+    async replacePendingOAuthConnections(clientId, expiredBefore, replacement) {
+      if (
+        replacement.status !== "pending" ||
+        replacement.clientId !== clientId
+      ) {
+        throw new Error("OAuth replacement must be pending for its client");
+      }
+      if (byId.has(replacement.id)) {
+        throw new Error(`mcp connection ${replacement.id} already exists`);
+      }
+      for (const [id, record] of byId) {
+        const expired = Date.parse(record.createdAt) <= expiredBefore;
+        if (
+          record.status !== "pending" ||
+          !record.clientId ||
+          (record.clientId !== clientId && !expired)
+        ) {
+          continue;
+        }
+        byId.set(id, {
+          ...record,
+          status: "revoked",
+          revokedAt: replacement.createdAt,
+        });
+      }
+      byId.set(replacement.id, { ...replacement });
+      byTokenHash.set(replacement.tokenHash, replacement.id);
+    },
+
     async list() {
       return Array.from(byId.values()).map((r) => ({ ...r }));
     },

packages/core/src/mcp/types.ts

@@ -142,6 +142,12 @@ export interface McpConnectionRecord {
 
 export interface McpConnectionStore {
   create(record: McpConnectionRecord): Promise<void>;
+  /** Atomically replace pending OAuth connections when the store can transact. */
+  replacePendingOAuthConnections?(
+    clientId: string,
+    expiredBefore: number,
+    replacement: McpConnectionRecord,
+  ): Promise<void>;
   list(): Promise<McpConnectionRecord[]>;
   getById(id: string): Promise<McpConnectionRecord | null>;
   /**