Add reusable OAuth lifecycle custody

.changeset/calm-oauth-custody.md

@@ -0,0 +1,5 @@
+---
+"@agent-native/core": patch
+---
+
+Add reusable owner- and resource-bound OAuth credential lifecycle primitives with concurrency-safe refresh, revocation, and explicit connection states.

packages/core/src/extensions/url-safety.spec.ts

@@ -130,6 +130,21 @@ describe("ssrfSafeFetch per-hop policies", () => {
     expect(redirectResponse.bodyUsed).toBe(true);
   });
 
+  it("can return a validated redirect for a caller with its own redirect policy", async () => {
+    const redirectResponse = new Response("moved", {
+      status: 302,
+      headers: { location: httpOrigin },
+    });
+    const fetchMock = vi.fn(async () => redirectResponse);
+    vi.stubGlobal("fetch", fetchMock);
+
+    await expect(
+      ssrfSafeFetch(httpsOrigin, {}, { followRedirects: false }),
+    ).resolves.toBe(redirectResponse);
+    expect(fetchMock).toHaveBeenCalledTimes(1);
+    expect(redirectResponse.bodyUsed).toBe(false);
+  });
+
   it("allows configured loopback aliases without allowing an unconfigured port", async () => {
     const fetchMock = vi.fn(async () => new Response("ok", { status: 200 }));
     vi.stubGlobal("fetch", fetchMock);

packages/core/src/extensions/url-safety.ts

@@ -331,6 +331,7 @@ export async function ssrfSafeFetch(
   init: RequestInit = {},
   options: {
     maxRedirects?: number;
+    followRedirects?: boolean;
     httpsOnly?: boolean;
     assertUrlAllowed?: (url: string) => void | Promise<void>;
     /**
@@ -387,6 +388,7 @@ export async function ssrfSafeFetch(
 
     const response = await fetch(currentUrl, fetchOpts);
     if (response.status >= 300 && response.status < 400) {
+      if (options.followRedirects === false) return response;
       const location = response.headers.get("location");
       if (!location) return response;
       // Drain the redirect body so the hop's connection is released instead

packages/core/src/mcp-client/oauth-client.spec.ts

@@ -6,6 +6,9 @@ const validateAuthorizationResponseIssuerMock = vi.hoisted(() => vi.fn());
 const deleteOAuthTokensMock = vi.hoisted(() => vi.fn());
 const getOAuthTokensMock = vi.hoisted(() => vi.fn());
 const saveOAuthTokensMock = vi.hoisted(() => vi.fn());
+const replaceOAuthTokensIfRevisionMock = vi.hoisted(() => vi.fn());
+const deleteOAuthTokensIfRevisionMock = vi.hoisted(() => vi.fn());
+const ssrfSafeFetchMock = vi.hoisted(() => vi.fn());
 
 vi.mock("@modelcontextprotocol/client", () => ({
   auth: authMock,
@@ -17,6 +20,29 @@ vi.mock("../oauth-tokens/store.js", () => ({
   deleteOAuthTokens: deleteOAuthTokensMock,
   getOAuthTokens: getOAuthTokensMock,
   saveOAuthTokens: saveOAuthTokensMock,
+  getOAuthTokenSnapshot: vi.fn(
+    async (provider: string, accountId: string, owner: string) => {
+      const tokens = await getOAuthTokensMock(provider, accountId, owner);
+      return tokens ? { tokens, owner, revision: 1 } : null;
+    },
+  ),
+  replaceOAuthTokensIfRevision: replaceOAuthTokensIfRevisionMock,
+  deleteOAuthTokensIfRevision: deleteOAuthTokensIfRevisionMock,
+}));
+
+vi.mock("../settings/store.js", () => ({
+  mutateSetting: vi.fn(
+    async (
+      _key: string,
+      updater: (
+        current: Record<string, unknown> | null,
+      ) => Record<string, unknown> | Promise<Record<string, unknown>>,
+    ) => updater(null),
+  ),
+}));
+
+vi.mock("../extensions/url-safety.js", () => ({
+  ssrfSafeFetch: ssrfSafeFetchMock,
 }));
 
 import {
@@ -25,6 +51,7 @@ import {
   getMcpOAuthAccessToken,
   McpOAuthClientProvider,
   readMcpOAuthCredentials,
+  revokeMcpOAuthCredentials,
   saveMcpOAuthCredentials,
   startMcpOAuthAuthorization,
   tokenExpiresAt,
@@ -60,11 +87,37 @@ const credentials = {
 };
 
 beforeEach(() => {
+  vi.restoreAllMocks();
   authMock.mockReset();
   refreshAuthorizationMock.mockReset();
   deleteOAuthTokensMock.mockReset();
   getOAuthTokensMock.mockReset();
   saveOAuthTokensMock.mockReset();
+  replaceOAuthTokensIfRevisionMock.mockReset();
+  replaceOAuthTokensIfRevisionMock.mockImplementation(
+    async (
+      _provider: string,
+      _accountId: string,
+      _owner: string,
+      _revision: number,
+      tokens: Record<string, unknown>,
+    ) => {
+      getOAuthTokensMock.mockResolvedValue(tokens);
+      saveOAuthTokensMock(_provider, _accountId, tokens, _owner);
+      return true;
+    },
+  );
+  deleteOAuthTokensIfRevisionMock.mockReset();
+  deleteOAuthTokensIfRevisionMock.mockImplementation(
+    async (provider: string, accountId: string, owner: string) => {
+      deleteOAuthTokensMock(provider, accountId, owner);
+      getOAuthTokensMock.mockResolvedValue(null);
+      return true;
+    },
+  );
+  ssrfSafeFetchMock.mockImplementation((url: string, init?: RequestInit) =>
+    fetch(url, init),
+  );
   validateAuthorizationResponseIssuerMock.mockReset();
 });
 
@@ -142,6 +195,7 @@ describe("MCP OAuth client", () => {
     await expect(
       fetchFn!("https://127.0.0.1/.well-known/oauth-authorization-server"),
     ).rejects.toThrow(/private\/internal address/);
+    expect(ssrfSafeFetchMock).not.toHaveBeenCalled();
     const provider = new McpOAuthClientProvider({
       serverUrl: "https://mcp.example.com/mcp",
       redirectUrl: "https://app.example.com/callback",
@@ -154,6 +208,44 @@ describe("MCP OAuth client", () => {
     ).toThrow(/private\/internal address/);
   });
 
+  it("routes OAuth discovery through the DNS-aware SSRF guard", async () => {
+    let fetchFn:
+      | ((url: string | URL, init?: RequestInit) => Promise<Response>)
+      | undefined;
+    authMock.mockImplementationOnce(
+      async (
+        provider: McpOAuthClientProvider,
+        options: { fetchFn?: typeof fetchFn },
+      ) => {
+        fetchFn = options.fetchFn;
+        provider.saveClientInformation(clientInformation as any);
+        provider.saveCodeVerifier("<CODE_VERIFIER>");
+        provider.redirectToAuthorization(
+          new URL("https://auth.example.com/authorize"),
+        );
+        return "REDIRECT";
+      },
+    );
+    ssrfSafeFetchMock.mockResolvedValueOnce(new Response("ok"));
+
+    await startMcpOAuthAuthorization({
+      serverUrl: "https://mcp.example.com/mcp",
+      redirectUrl: "https://app.example.com/callback",
+      state: "<STATE>",
+    });
+    await fetchFn!("https://auth.example.com/discovery");
+
+    expect(ssrfSafeFetchMock).toHaveBeenCalledWith(
+      "https://auth.example.com/discovery",
+      expect.objectContaining({ redirect: "manual" }),
+      expect.objectContaining({
+        maxRedirects: 0,
+        followRedirects: false,
+        allowedPrivateOrigins: [],
+      }),
+    );
+  });
+
   it("validates every OAuth redirect hop and strips credentials across origins", async () => {
     let fetchFn:
       | ((url: string | URL, init?: RequestInit) => Promise<Response>)
@@ -302,7 +394,15 @@ describe("MCP OAuth client", () => {
     expect(saveOAuthTokensMock).toHaveBeenCalledWith(
       "mcp",
       "mcp_oauth:test",
-      credentials,
+      expect.objectContaining({
+        ...credentials,
+        oauthLifecycle: {
+          version: 1,
+          provider: "mcp",
+          resource: "https://mcp.example.com/mcp",
+          owner: "user:alice@example.com",
+        },
+      }),
       "user:alice@example.com",
     );
   });
@@ -312,7 +412,7 @@ describe("MCP OAuth client", () => {
       ...credentials,
       tokenExpiresAt: Date.now() - 1,
     };
-    getOAuthTokensMock.mockResolvedValueOnce(expiring);
+    getOAuthTokensMock.mockResolvedValue(expiring);
     refreshAuthorizationMock.mockResolvedValueOnce({
       access_token: "<NEW_ACCESS_TOKEN>",
       token_type: "bearer",
@@ -341,7 +441,7 @@ describe("MCP OAuth client", () => {
   });
 
   it("requires reauthorization for expiring legacy credentials without issuer binding", async () => {
-    getOAuthTokensMock.mockResolvedValueOnce({
+    getOAuthTokensMock.mockResolvedValue({
       ...credentials,
       clientInformation: {
         client_id: "legacy-client",
@@ -361,7 +461,7 @@ describe("MCP OAuth client", () => {
   });
 
   it("does not return an expired token when refresh fails", async () => {
-    getOAuthTokensMock.mockResolvedValueOnce({
+    getOAuthTokensMock.mockResolvedValue({
       ...credentials,
       tokenExpiresAt: Date.now() - 1,
     });
@@ -380,7 +480,7 @@ describe("MCP OAuth client", () => {
   });
 
   it("keeps a still-valid token when an early refresh fails", async () => {
-    getOAuthTokensMock.mockResolvedValueOnce({
+    getOAuthTokensMock.mockResolvedValue({
       ...credentials,
       tokenExpiresAt: Date.now() + 30_000,
     });
@@ -398,6 +498,93 @@ describe("MCP OAuth client", () => {
     ).resolves.toBe("<ACCESS_TOKEN>");
   });
 
+  it("does not return a legacy MCP token for a different resource", async () => {
+    getOAuthTokensMock.mockResolvedValue(credentials);
+
+    await expect(
+      getMcpOAuthAccessToken({
+        key: "mcp_oauth:test",
+        scope: "user",
+        scopeId: "alice@example.com",
+        serverUrl: "https://other.example.com/mcp",
+      }),
+    ).resolves.toBeNull();
+    expect(refreshAuthorizationMock).not.toHaveBeenCalled();
+  });
+
+  it("revokes the refresh token before deleting local MCP custody", async () => {
+    getOAuthTokensMock.mockResolvedValue({
+      ...credentials,
+      discoveryState: {
+        ...credentials.discoveryState,
+        authorizationServerMetadata: {
+          ...credentials.discoveryState.authorizationServerMetadata,
+          revocation_endpoint: "https://auth.example.com/revoke",
+        },
+      },
+    });
+    const fetchMock = vi
+      .spyOn(globalThis, "fetch")
+      .mockResolvedValueOnce(new Response(null, { status: 200 }));
+
+    await expect(
+      revokeMcpOAuthCredentials({
+        key: "mcp_oauth:test",
+        scope: "user",
+        scopeId: "alice@example.com",
+        serverUrl: "https://mcp.example.com/mcp",
+      }),
+    ).resolves.toEqual({ remote: "succeeded", local: "deleted" });
+
+    expect(String(fetchMock.mock.calls[0]?.[0])).toBe(
+      "https://auth.example.com/revoke",
+    );
+    const request = fetchMock.mock.calls[0]?.[1] as RequestInit;
+    const body = new URLSearchParams(String(request.body));
+    expect(body.get("token")).toBe("<REFRESH_TOKEN>");
+    expect(body.get("token_type_hint")).toBe("refresh_token");
+    expect(body.get("client_id")).toBe("mcp-client-test");
+    expect(deleteOAuthTokensIfRevisionMock).toHaveBeenCalledTimes(1);
+    expect(ssrfSafeFetchMock).toHaveBeenCalledWith(
+      "https://auth.example.com/revoke",
+      expect.any(Object),
+      expect.objectContaining({ maxRedirects: 0, httpsOnly: true }),
+    );
+  });
+
+  it("fails closed instead of posting a token to a loopback revocation endpoint", async () => {
+    getOAuthTokensMock.mockResolvedValue({
+      ...credentials,
+      discoveryState: {
+        ...credentials.discoveryState,
+        authorizationServerMetadata: {
+          ...credentials.discoveryState.authorizationServerMetadata,
+          revocation_endpoint: "http://127.0.0.1:9000/revoke",
+        },
+      },
+    });
+    ssrfSafeFetchMock.mockRejectedValueOnce(
+      new Error("SSRF blocked: refusing to fetch private/internal address"),
+    );
+    const fetchMock = vi.spyOn(globalThis, "fetch");
+
+    await expect(
+      revokeMcpOAuthCredentials({
+        key: "mcp_oauth:test",
+        scope: "user",
+        scopeId: "alice@example.com",
+        serverUrl: "https://mcp.example.com/mcp",
+      }),
+    ).resolves.toEqual({ remote: "failed", local: "deleted" });
+
+    expect(ssrfSafeFetchMock).toHaveBeenCalledWith(
+      "http://127.0.0.1:9000/revoke",
+      expect.any(Object),
+      expect.objectContaining({ maxRedirects: 0, httpsOnly: true }),
+    );
+    expect(fetchMock).not.toHaveBeenCalled();
+  });
+
   it("rejects malformed stored bundles", async () => {
     getOAuthTokensMock.mockResolvedValueOnce({ access_token: "<TOKEN>" });
 
@@ -406,13 +593,14 @@ describe("MCP OAuth client", () => {
         key: "mcp_oauth:test",
         scope: "user",
         scopeId: "alice@example.com",
+        serverUrl: "https://mcp.example.com/mcp",
       }),
     ).resolves.toBeNull();
   });
 
-  it("binds reads and deletes to the credential owner", async () => {
-    getOAuthTokensMock.mockResolvedValueOnce(null);
-    deleteOAuthTokensMock.mockResolvedValueOnce(1);
+  it("preserves owner-bound legacy reads and deletes without a serverUrl argument", async () => {
+    getOAuthTokensMock.mockResolvedValue(credentials);
+    deleteOAuthTokensMock.mockResolvedValue(1);
 
     await readMcpOAuthCredentials({
       key: "mcp_oauth:test",