|
| 1 | +import { auth } from '@modelcontextprotocol/sdk/client/auth.js' |
| 2 | +import { dbChainMockFns, encryptionMock, resetDbChainMock } from '@sim/testing' |
| 3 | +import { beforeEach, describe, expect, it, vi } from 'vitest' |
| 4 | + |
| 5 | +vi.mock('@/lib/core/security/encryption', () => encryptionMock) |
| 6 | + |
| 7 | +import { SimMcpOauthProvider } from '@/lib/mcp/oauth/provider' |
| 8 | +import type { McpOauthRow } from '@/lib/mcp/oauth/storage' |
| 9 | + |
| 10 | +const SERVER_URL = 'https://mcp.example.com/mcp' |
| 11 | +const ISSUER = 'https://auth.example.com' |
| 12 | + |
| 13 | +function grantRow(): McpOauthRow { |
| 14 | + return { |
| 15 | + id: 'oauth-row-1', |
| 16 | + mcpServerId: 'server-1', |
| 17 | + userId: 'authorizer-1', |
| 18 | + workspaceId: 'workspace-1', |
| 19 | + organizationId: null, |
| 20 | + clientInformation: { client_id: 'client-1', issuer: ISSUER }, |
| 21 | + tokens: { access_token: 'expired', token_type: 'bearer', refresh_token: 'refresh-1' }, |
| 22 | + codeVerifier: null, |
| 23 | + state: null, |
| 24 | + stateCreatedAt: null, |
| 25 | + updatedAt: new Date(), |
| 26 | + } |
| 27 | +} |
| 28 | + |
| 29 | +/** An authorization server whose token endpoint answers every refresh with `token`. */ |
| 30 | +function authorizationServer(token: { status: number; body: Record<string, unknown> }) { |
| 31 | + return vi.fn(async (input: string | URL) => { |
| 32 | + const url = String(input) |
| 33 | + const json = (status: number, body: unknown) => |
| 34 | + new Response(JSON.stringify(body), { |
| 35 | + status, |
| 36 | + headers: { 'content-type': 'application/json' }, |
| 37 | + }) |
| 38 | + if (url.startsWith('https://mcp.example.com/.well-known/oauth-protected-resource')) |
| 39 | + return json(200, { resource: SERVER_URL, authorization_servers: [ISSUER] }) |
| 40 | + if (url.startsWith(`${ISSUER}/.well-known/`)) |
| 41 | + return json(200, { |
| 42 | + issuer: ISSUER, |
| 43 | + authorization_endpoint: `${ISSUER}/authorize`, |
| 44 | + token_endpoint: `${ISSUER}/token`, |
| 45 | + response_types_supported: ['code'], |
| 46 | + code_challenge_methods_supported: ['S256'], |
| 47 | + }) |
| 48 | + if (url === `${ISSUER}/token`) return json(token.status, token.body) |
| 49 | + return json(404, {}) |
| 50 | + }) |
| 51 | +} |
| 52 | + |
| 53 | +/** |
| 54 | + * Best-effort discovery connects with the workspace's one shared grant. A refresh that fails |
| 55 | + * there must leave the grant as it was: the SDK otherwise clears the tokens and starts a new |
| 56 | + * authorization flow, disconnecting the server for every member. |
| 57 | + */ |
| 58 | +describe('SimMcpOauthProvider preserving the shared grant', () => { |
| 59 | + beforeEach(() => { |
| 60 | + resetDbChainMock() |
| 61 | + }) |
| 62 | + |
| 63 | + it('writes nothing to the OAuth row when a refresh is rejected', async () => { |
| 64 | + const provider = new SimMcpOauthProvider({ row: grantRow(), preserveGrant: true }) |
| 65 | + |
| 66 | + await expect( |
| 67 | + auth(provider, { |
| 68 | + serverUrl: SERVER_URL, |
| 69 | + fetchFn: authorizationServer({ status: 400, body: { error: 'invalid_grant' } }), |
| 70 | + }) |
| 71 | + ).rejects.toThrow('MCP OAuth grant left unchanged') |
| 72 | + |
| 73 | + expect(dbChainMockFns.update).not.toHaveBeenCalled() |
| 74 | + expect(dbChainMockFns.insert).not.toHaveBeenCalled() |
| 75 | + expect(provider.tokens()?.refresh_token).toBe('refresh-1') |
| 76 | + }) |
| 77 | + |
| 78 | + it('persists a successful refresh, which keeps the grant valid for everyone', async () => { |
| 79 | + const provider = new SimMcpOauthProvider({ row: grantRow(), preserveGrant: true }) |
| 80 | + |
| 81 | + await expect( |
| 82 | + auth(provider, { |
| 83 | + serverUrl: SERVER_URL, |
| 84 | + fetchFn: authorizationServer({ |
| 85 | + status: 200, |
| 86 | + body: { access_token: 'fresh', token_type: 'bearer', refresh_token: 'refresh-2' }, |
| 87 | + }), |
| 88 | + }) |
| 89 | + ).resolves.toBe('AUTHORIZED') |
| 90 | + |
| 91 | + expect(dbChainMockFns.update).toHaveBeenCalledTimes(1) |
| 92 | + expect(provider.tokens()?.access_token).toBe('fresh') |
| 93 | + }) |
| 94 | +}) |
0 commit comments