simstudioai--sim/apps/sim/lib/mcp/oauth/storage.test.ts

/**
 * @vitest-environment node
 */
import {
  dbChainMock,
  dbChainMockFns,
  encryptionMock,
  encryptionMockFns,
  resetDbChainMock,
  schemaMock,
} from '@sim/testing'
import { beforeEach, describe, expect, it, vi } from 'vitest'

const { mockAcquireLock, mockReleaseLock, mockExtendLock } = vi.hoisted(() => ({
  mockAcquireLock: vi.fn(),
  mockReleaseLock: vi.fn(),
  mockExtendLock: vi.fn(),
}))

vi.mock('@sim/db', () => dbChainMock)
vi.mock('@sim/db/schema', () => schemaMock)
vi.mock('@/lib/core/security/encryption', () => encryptionMock)
vi.mock('@/lib/core/config/redis', () => ({
  acquireLock: mockAcquireLock,
  releaseLock: mockReleaseLock,
  extendLock: mockExtendLock,
}))

import {
  getOrCreateOauthRow,
  loadOauthRow,
  setOauthRowUser,
  withMcpOauthRefreshLock,
} from './storage'

describe('MCP OAuth storage', () => {
  beforeEach(() => {
    vi.clearAllMocks()
    resetDbChainMock()
    encryptionMockFns.mockDecryptSecret.mockResolvedValue({ decrypted: '{}' })
    encryptionMockFns.mockEncryptSecret.mockResolvedValue({
      encrypted: 'encrypted',
      iv: 'iv',
    })
  })

  it('loads OAuth state by MCP server, independent of the requesting user', async () => {
    dbChainMockFns.limit.mockResolvedValueOnce([
      {
        id: 'oauth-row-1',
        mcpServerId: 'server-1',
        userId: 'authorizer-1',
        workspaceId: 'workspace-1',
        clientInformation: null,
        tokens: null,
        codeVerifier: null,
        state: null,
        updatedAt: new Date(),
      },
    ])

    const row = await loadOauthRow({ mcpServerId: 'server-1' })

    expect(row).toMatchObject({
      id: 'oauth-row-1',
      mcpServerId: 'server-1',
      userId: 'authorizer-1',
      workspaceId: 'workspace-1',
    })
    expect(dbChainMockFns.limit).toHaveBeenCalledTimes(1)
  })

  it('reuses the existing workspace OAuth row instead of creating a per-user row', async () => {
    dbChainMockFns.limit.mockResolvedValueOnce([
      {
        id: 'oauth-row-1',
        mcpServerId: 'server-1',
        userId: 'authorizer-1',
        workspaceId: 'workspace-1',
        clientInformation: null,
        tokens: null,
        codeVerifier: null,
        state: null,
        updatedAt: new Date(),
      },
    ])

    const row = await getOrCreateOauthRow({
      mcpServerId: 'server-1',
      userId: 'different-user',
      workspaceId: 'workspace-1',
    })

    expect(row.id).toBe('oauth-row-1')
    expect(row.userId).toBe('authorizer-1')
    expect(dbChainMockFns.insert).not.toHaveBeenCalled()
  })

  it('records the latest authorizing user without changing row ownership', async () => {
    await setOauthRowUser('oauth-row-1', 'user-2')

    expect(dbChainMockFns.update).toHaveBeenCalledTimes(1)
    expect(dbChainMockFns.set).toHaveBeenCalledWith(
      expect.objectContaining({
        userId: 'user-2',
        updatedAt: expect.any(Date),
      })
    )
  })
})

describe('withMcpOauthRefreshLock', () => {
  beforeEach(() => {
    vi.clearAllMocks()
    mockAcquireLock.mockReset()
    mockReleaseLock.mockReset()
    mockExtendLock.mockReset()
    mockReleaseLock.mockResolvedValue(true)
    mockExtendLock.mockResolvedValue(true)
  })

  it('serializes concurrent in-process callers, each running its own fn()', async () => {
    mockAcquireLock.mockResolvedValue(true)
    let active = 0
    let maxActive = 0
    const fn = vi.fn(async () => {
      active++
      maxActive = Math.max(maxActive, active)
      await new Promise((r) => setTimeout(r, 1))
      active--
      return 'tokens'
    })

    const results = await Promise.all([
      withMcpOauthRefreshLock('row-serial', fn),
      withMcpOauthRefreshLock('row-serial', fn),
      withMcpOauthRefreshLock('row-serial', fn),
    ])

    expect(results).toEqual(['tokens', 'tokens', 'tokens'])
    // Each caller gets its own fn() invocation — critical because fn() returns
    // a stateful McpClient that can't be shared across consumers.
    expect(fn).toHaveBeenCalledTimes(3)
    // But never two at the same time within a process.
    expect(maxActive).toBe(1)
    expect(mockAcquireLock).toHaveBeenCalledTimes(3)
    expect(mockReleaseLock).toHaveBeenCalledTimes(3)
  })

  it('serializes cross-process callers: follower polls until leader releases', async () => {
    // First acquire fails (another process holds it), second succeeds.
    mockAcquireLock.mockResolvedValueOnce(false).mockResolvedValueOnce(true)
    const fn = vi.fn(async () => 'fresh')

    const result = await withMcpOauthRefreshLock('row-mutex', fn)

    expect(result).toBe('fresh')
    expect(mockAcquireLock).toHaveBeenCalledTimes(2)
    expect(fn).toHaveBeenCalledTimes(1)
  })

  it('falls open when Redis is unavailable on acquire', async () => {
    mockAcquireLock.mockRejectedValueOnce(new Error('Redis connection refused'))
    const fn = vi.fn(async () => 'uncoordinated')

    const result = await withMcpOauthRefreshLock('row-redis-down', fn)

    expect(result).toBe('uncoordinated')
    expect(fn).toHaveBeenCalledTimes(1)
    expect(mockReleaseLock).not.toHaveBeenCalled()
  })

  it('releases the lock even when fn throws', async () => {
    mockAcquireLock.mockResolvedValue(true)
    const fn = vi.fn(async () => {
      throw new Error('refresh failed')
    })

    await expect(withMcpOauthRefreshLock('row-throws', fn)).rejects.toThrow('refresh failed')

    expect(mockReleaseLock).toHaveBeenCalledTimes(1)
  })

  it('does not surface releaseLock failures to the caller', async () => {
    mockAcquireLock.mockResolvedValue(true)
    mockReleaseLock.mockRejectedValueOnce(new Error('release failed'))
    const fn = vi.fn(async () => 'value')

    const result = await withMcpOauthRefreshLock('row-release-fail', fn)
    expect(result).toBe('value')
  })

  it('uses per-row lock keys so different rows do not serialize', async () => {
    mockAcquireLock.mockResolvedValue(true)
    const fn = vi.fn(async () => 'ok')

    await Promise.all([withMcpOauthRefreshLock('row-a', fn), withMcpOauthRefreshLock('row-b', fn)])

    expect(mockAcquireLock).toHaveBeenCalledTimes(2)
    const keys = mockAcquireLock.mock.calls.map((c) => c[0])
    expect(keys).toContain('mcp:oauth:refresh:row-a')
    expect(keys).toContain('mcp:oauth:refresh:row-b')
  })

  it('throws when the lock is held longer than the max wait (does not race)', async () => {
    vi.useFakeTimers()
    try {
      // Acquire always fails — another process holds the lock with watchdog extension.
      mockAcquireLock.mockResolvedValue(false)
      const fn = vi.fn(async () => 'should-not-run')

      const pending = withMcpOauthRefreshLock('row-deadline', fn)
      // Attach the rejection expectation before draining so Vitest doesn't see
      // an unhandled rejection while timers advance.
      const assertion = expect(pending).rejects.toThrow(/held longer than/)
      await vi.advanceTimersByTimeAsync(31_000)
      await assertion
      expect(fn).not.toHaveBeenCalled()
    } finally {
      vi.useRealTimers()
    }
  })

  it('bounds the queue wait: callers stalled behind a wedged link reject without running fn', async () => {
    vi.useFakeTimers()
    try {
      mockAcquireLock.mockResolvedValue(true)
      let resolveFirst: (value: string) => void
      const hungFn = vi.fn(
        () =>
          new Promise<string>((resolve) => {
            resolveFirst = resolve
          })
      )
      const queuedFn = vi.fn(async () => 'second')

      const first = withMcpOauthRefreshLock('row-stall', hungFn)
      const second = withMcpOauthRefreshLock('row-stall', queuedFn)
      const secondAssertion = expect(second).rejects.toThrow(/stalled for/)

      await vi.advanceTimersByTimeAsync(90_000)
      await secondAssertion
      expect(queuedFn).not.toHaveBeenCalled()

      // The wedged link is untouched by the queue deadline (its own fn keeps
      // the lock, protecting a possibly mid-rotation refresh) and the row
      // heals once it settles — including skipping the abandoned link's fn.
      resolveFirst!('first')
      await expect(first).resolves.toBe('first')

      await expect(withMcpOauthRefreshLock('row-stall', async () => 'healed')).resolves.toBe(
        'healed'
      )
      expect(queuedFn).not.toHaveBeenCalled()
    } finally {
      vi.useRealTimers()
    }
  })

  it('extends the lock TTL while fn() is running so long refreshes do not lose the lock', async () => {
    vi.useFakeTimers()
    try {
      mockAcquireLock.mockResolvedValue(true)
      let resolveFn: (v: string) => void
      const fn = vi.fn(
        () =>
          new Promise<string>((resolve) => {
            resolveFn = resolve
          })
      )

      const pending = withMcpOauthRefreshLock('row-watchdog', fn)

      // Advance time past two extend intervals (5s + 5s = 10s).
      await vi.advanceTimersByTimeAsync(11_000)
      expect(mockExtendLock.mock.calls.length).toBeGreaterThanOrEqual(2)
      for (const call of mockExtendLock.mock.calls) {
        expect(call[0]).toBe('mcp:oauth:refresh:row-watchdog')
      }

      resolveFn!('done')
      await expect(pending).resolves.toBe('done')

      // Watchdog must stop once fn() settles — no more extend calls.
      const extendCallsAtFinish = mockExtendLock.mock.calls.length
      await vi.advanceTimersByTimeAsync(20_000)
      expect(mockExtendLock.mock.calls.length).toBe(extendCallsAtFinish)
    } finally {
      vi.useRealTimers()
    }
  })
})