simstudioai--sim/apps/sim/app/api/chat/utils.test.ts

/**
 * Tests for chat API utils
 *
 * @vitest-environment node
 */
import {
  dbChainMock,
  dbChainMockFns,
  encryptionMock,
  encryptionMockFns,
  loggingSessionMock,
  workflowsUtilsMock,
} from '@sim/testing'
import type { NextResponse } from 'next/server'
import { beforeEach, describe, expect, it, vi } from 'vitest'

const {
  mockMergeSubblockStateWithValues,
  mockMergeSubBlockValues,
  mockValidateAuthToken,
  mockSetDeploymentAuthCookie,
  mockIsEmailAllowed,
  mockGetSession,
  mockCheckRateLimitDirect,
  mockIsWorkspaceApiExecutionEntitled,
  flagState,
} = vi.hoisted(() => ({
  mockMergeSubblockStateWithValues: vi.fn().mockReturnValue({}),
  mockMergeSubBlockValues: vi.fn().mockReturnValue({}),
  mockValidateAuthToken: vi.fn().mockReturnValue(false),
  mockSetDeploymentAuthCookie: vi.fn(),
  mockIsEmailAllowed: vi.fn(),
  mockGetSession: vi.fn(),
  mockCheckRateLimitDirect: vi.fn().mockResolvedValue({ allowed: true }),
  mockIsWorkspaceApiExecutionEntitled: vi.fn().mockResolvedValue(true),
  flagState: { isBillingEnabled: false, isFreeApiDeploymentGateEnabled: true },
}))

vi.mock('@sim/db', () => dbChainMock)

vi.mock('@/lib/billing/core/api-access', () => ({
  isWorkspaceApiExecutionEntitled: mockIsWorkspaceApiExecutionEntitled,
}))

vi.mock('@/lib/core/rate-limiter', () => ({
  RateLimiter: class {
    checkRateLimitDirect = mockCheckRateLimitDirect
  },
}))

vi.mock('@/lib/auth', () => ({
  auth: { api: { getSession: vi.fn() } },
  getSession: mockGetSession,
}))

const mockDecryptSecret = encryptionMockFns.mockDecryptSecret

vi.mock('@/lib/logs/execution/logging-session', () => loggingSessionMock)

vi.mock('@/executor', () => ({
  Executor: vi.fn(),
}))

vi.mock('@/serializer', () => ({
  Serializer: vi.fn(),
}))

vi.mock('@sim/workflow-persistence/subblocks', () => ({
  mergeSubblockStateWithValues: mockMergeSubblockStateWithValues,
  mergeSubBlockValues: mockMergeSubBlockValues,
}))

vi.mock('@/lib/core/security/encryption', () => encryptionMock)

vi.mock('@/lib/core/security/deployment', () => ({
  validateAuthToken: mockValidateAuthToken,
  setDeploymentAuthCookie: mockSetDeploymentAuthCookie,
  isEmailAllowed: mockIsEmailAllowed,
  deploymentAuthCookieName: (prefix: string, id: string) => `${prefix}_auth_${id}`,
}))

vi.mock('@/lib/core/config/env-flags', () => ({
  isDev: true,
  isProd: false,
  get isBillingEnabled() {
    return flagState.isBillingEnabled
  },
  get isFreeApiDeploymentGateEnabled() {
    return flagState.isFreeApiDeploymentGateEnabled
  },
}))

vi.mock('@/lib/workflows/utils', () => workflowsUtilsMock)

import { NextRequest } from 'next/server'
import { decryptSecret } from '@/lib/core/security/encryption'
import { assertChatEmbedAllowed, setChatAuthCookie, validateChatAuth } from '@/app/api/chat/utils'

function chatRequest(origin?: string): NextRequest {
  return new NextRequest('https://www.sim.ai/api/chat/abc', {
    headers: origin ? { origin } : undefined,
  })
}

describe('Chat API Utils', () => {
  beforeEach(() => {
    vi.clearAllMocks()
    vi.stubGlobal('process', {
      ...process,
      env: {
        ...process.env,
        NODE_ENV: 'development',
      },
    })
  })

  describe('Auth token utils', () => {
    it('should accept valid auth cookie via validateChatAuth', async () => {
      mockValidateAuthToken.mockReturnValue(true)

      const deployment = {
        id: 'chat-id',
        authType: 'password',
        password: 'encrypted-password',
      }

      const mockRequest = {
        method: 'POST',
        cookies: {
          get: vi.fn().mockReturnValue({ value: 'valid-token' }),
        },
      } as any

      const result = await validateChatAuth('request-id', deployment, mockRequest)
      expect(mockValidateAuthToken).toHaveBeenCalledWith(
        'valid-token',
        'chat-id',
        'password',
        'encrypted-password'
      )
      expect(result.authorized).toBe(true)
    })

    it('should reject invalid auth cookie via validateChatAuth', async () => {
      mockValidateAuthToken.mockReturnValue(false)

      const deployment = {
        id: 'chat-id',
        authType: 'password',
        password: 'encrypted-password',
      }

      const mockRequest = {
        method: 'GET',
        cookies: {
          get: vi.fn().mockReturnValue({ value: 'invalid-token' }),
        },
      } as any

      const result = await validateChatAuth('request-id', deployment, mockRequest)
      expect(result.authorized).toBe(false)
    })
  })

  describe('Cookie handling', () => {
    it('should delegate to setDeploymentAuthCookie', () => {
      const mockResponse = {
        cookies: { set: vi.fn() },
      } as unknown as NextResponse

      setChatAuthCookie(mockResponse, 'test-chat-id', 'password')

      expect(mockSetDeploymentAuthCookie).toHaveBeenCalledWith(
        mockResponse,
        'chat',
        'test-chat-id',
        'password',
        undefined
      )
    })
  })

  describe('Chat auth validation', () => {
    beforeEach(() => {
      mockDecryptSecret.mockResolvedValue({ decrypted: 'correct-password' })
      mockCheckRateLimitDirect.mockResolvedValue({ allowed: true })
    })

    it('should allow access to public chats', async () => {
      const deployment = {
        id: 'chat-id',
        authType: 'public',
      }

      const mockRequest = {
        cookies: {
          get: vi.fn().mockReturnValue(null),
        },
      } as any

      const result = await validateChatAuth('request-id', deployment, mockRequest)

      expect(result.authorized).toBe(true)
    })

    it('should request password auth for GET requests', async () => {
      const deployment = {
        id: 'chat-id',
        authType: 'password',
      }

      const mockRequest = {
        method: 'GET',
        cookies: {
          get: vi.fn().mockReturnValue(null),
        },
      } as any

      const result = await validateChatAuth('request-id', deployment, mockRequest)

      expect(result.authorized).toBe(false)
      expect(result.error).toBe('auth_required_password')
    })

    it('should validate password for POST requests', async () => {
      const deployment = {
        id: 'chat-id',
        authType: 'password',
        password: 'encrypted-password',
      }

      const mockRequest = {
        method: 'POST',
        cookies: {
          get: vi.fn().mockReturnValue(null),
        },
      } as any

      const parsedBody = {
        password: 'correct-password',
      }

      const result = await validateChatAuth('request-id', deployment, mockRequest, parsedBody)

      expect(decryptSecret).toHaveBeenCalledWith('encrypted-password')
      expect(result.authorized).toBe(true)
    })

    it('should reject incorrect password', async () => {
      const deployment = {
        id: 'chat-id',
        authType: 'password',
        password: 'encrypted-password',
      }

      const mockRequest = {
        method: 'POST',
        cookies: {
          get: vi.fn().mockReturnValue(null),
        },
      } as any

      const parsedBody = {
        password: 'wrong-password',
      }

      const result = await validateChatAuth('request-id', deployment, mockRequest, parsedBody)

      expect(result.authorized).toBe(false)
      expect(result.error).toBe('Invalid password')
    })

    it('should return 429 when the password attempt rate limit is exceeded', async () => {
      mockCheckRateLimitDirect.mockResolvedValueOnce({ allowed: false, retryAfterMs: 60_000 })

      const deployment = {
        id: 'chat-id',
        authType: 'password',
        password: 'encrypted-password',
      }

      const mockRequest = {
        method: 'POST',
        cookies: {
          get: vi.fn().mockReturnValue(null),
        },
      } as any

      const result = await validateChatAuth('request-id', deployment, mockRequest, {
        password: 'any-guess',
      })

      expect(result.authorized).toBe(false)
      expect(result.status).toBe(429)
      expect(result.retryAfterMs).toBe(60_000)
      expect(decryptSecret).not.toHaveBeenCalled()
    })

    it('should request email auth for email-protected chats', async () => {
      const deployment = {
        id: 'chat-id',
        authType: 'email',
        allowedEmails: ['user@example.com', '@company.com'],
      }

      const mockRequest = {
        method: 'GET',
        cookies: {
          get: vi.fn().mockReturnValue(null),
        },
      } as any

      const result = await validateChatAuth('request-id', deployment, mockRequest)

      expect(result.authorized).toBe(false)
      expect(result.error).toBe('auth_required_email')
    })

    it('should check allowed emails for email auth', async () => {
      const deployment = {
        id: 'chat-id',
        authType: 'email',
        allowedEmails: ['user@example.com', '@company.com'],
      }

      const mockRequest = {
        method: 'POST',
        cookies: {
          get: vi.fn().mockReturnValue(null),
        },
      } as any

      mockIsEmailAllowed.mockReturnValue(true)
      const result1 = await validateChatAuth('request-id', deployment, mockRequest, {
        email: 'user@example.com',
      })
      expect(result1.authorized).toBe(false)
      expect(result1.error).toBe('otp_required')

      const result2 = await validateChatAuth('request-id', deployment, mockRequest, {
        email: 'other@company.com',
      })
      expect(result2.authorized).toBe(false)
      expect(result2.error).toBe('otp_required')

      mockIsEmailAllowed.mockReturnValue(false)
      const result3 = await validateChatAuth('request-id', deployment, mockRequest, {
        email: 'user@unknown.com',
      })
      expect(result3.authorized).toBe(false)
      expect(result3.error).toBe('Email not authorized')
    })

    describe('SSO auth', () => {
      const ssoDeployment = {
        id: 'chat-id',
        authType: 'sso',
        allowedEmails: ['user@example.com', '@company.com'],
      }

      const postRequest = {
        method: 'POST',
        cookies: { get: vi.fn().mockReturnValue(null) },
      } as any

      it('rejects when no session is present', async () => {
        mockGetSession.mockResolvedValue(null)

        const result = await validateChatAuth('request-id', ssoDeployment, postRequest, {
          input: 'hello',
        })

        expect(result.authorized).toBe(false)
        expect(result.error).toBe('auth_required_sso')
      })

      it('ignores body-supplied email and uses the session email', async () => {
        mockGetSession.mockResolvedValue({ user: { email: 'session@example.com' } })
        mockIsEmailAllowed.mockReturnValue(true)

        await validateChatAuth('request-id', ssoDeployment, postRequest, {
          email: 'attacker@evil.com',
          input: 'hello',
        })

        expect(mockIsEmailAllowed).toHaveBeenCalledWith(
          'session@example.com',
          ssoDeployment.allowedEmails
        )
      })

      it('authorizes execution when session email is allowlisted', async () => {
        mockGetSession.mockResolvedValue({ user: { email: 'user@example.com' } })
        mockIsEmailAllowed.mockReturnValue(true)

        const result = await validateChatAuth('request-id', ssoDeployment, postRequest, {
          input: 'hello',
        })

        expect(result.authorized).toBe(true)
      })

      it('rejects execution when session email is not allowlisted', async () => {
        mockGetSession.mockResolvedValue({ user: { email: 'stranger@other.com' } })
        mockIsEmailAllowed.mockReturnValue(false)

        const result = await validateChatAuth('request-id', ssoDeployment, postRequest, {
          input: 'hello',
        })

        expect(result.authorized).toBe(false)
        expect(result.error).toBe('Your email is not authorized to access this resource')
      })
    })
  })

  describe('Execution Result Processing', () => {
    it.concurrent('should process logs regardless of overall success status', () => {
      const executionResult = {
        success: false,
        output: {},
        logs: [
          {
            blockId: 'agent1',
            startedAt: '2023-01-01T00:00:00Z',
            endedAt: '2023-01-01T00:00:01Z',
            durationMs: 1000,
            success: true,
            output: { content: 'Agent 1 succeeded' },
            error: undefined,
          },
          {
            blockId: 'agent2',
            startedAt: '2023-01-01T00:00:00Z',
            endedAt: '2023-01-01T00:00:01Z',
            durationMs: 500,
            success: false,
            output: null,
            error: 'Agent 2 failed',
          },
        ],
        metadata: { duration: 1000 },
      }

      expect(executionResult.success).toBe(false)
      expect(executionResult.logs).toBeDefined()
      expect(executionResult.logs).toHaveLength(2)

      expect(executionResult.logs[0].success).toBe(true)
      expect(executionResult.logs[0].output?.content).toBe('Agent 1 succeeded')

      expect(executionResult.logs[1].success).toBe(false)
      expect(executionResult.logs[1].error).toBe('Agent 2 failed')
    })

    it.concurrent('should handle ExecutionResult vs StreamingExecution types correctly', () => {
      const executionResult = {
        success: true,
        output: { content: 'test' },
        logs: [],
        metadata: { duration: 100 },
      }

      const directResult = executionResult
      const extractedDirect = directResult
      expect(extractedDirect).toBe(executionResult)

      const streamingResult = {
        stream: new ReadableStream(),
        execution: executionResult,
      }

      const extractedFromStreaming =
        streamingResult && typeof streamingResult === 'object' && 'execution' in streamingResult
          ? streamingResult.execution
          : streamingResult

      expect(extractedFromStreaming).toBe(executionResult)
    })
  })
})

describe('assertChatEmbedAllowed', () => {
  beforeEach(() => {
    vi.clearAllMocks()
    flagState.isBillingEnabled = true
    flagState.isFreeApiDeploymentGateEnabled = true
    mockIsWorkspaceApiExecutionEntitled.mockResolvedValue(true)
    dbChainMockFns.limit.mockResolvedValue([{ workspaceId: 'ws-1' }])
  })

  it('returns 403 for a cross-site origin when the owner is on the free plan', async () => {
    mockIsWorkspaceApiExecutionEntitled.mockResolvedValueOnce(false)
    const res = await assertChatEmbedAllowed(
      chatRequest('https://evil.example.com'),
      'wf-1',
      'req-1'
    )
    expect(res?.status).toBe(403)
  })

  it('allows a cross-site origin when the owner is on a paid plan', async () => {
    const res = await assertChatEmbedAllowed(
      chatRequest('https://evil.example.com'),
      'wf-1',
      'req-1'
    )
    expect(res).toBeNull()
  })

  it('returns 403 for a cross-site origin when the workflow has no active workspace', async () => {
    dbChainMockFns.limit.mockResolvedValueOnce([])
    const res = await assertChatEmbedAllowed(
      chatRequest('https://evil.example.com'),
      'wf-1',
      'req-1'
    )
    expect(res?.status).toBe(403)
    expect(mockIsWorkspaceApiExecutionEntitled).not.toHaveBeenCalled()
  })

  it('allows a first-party *.sim.ai origin without gating', async () => {
    const res = await assertChatEmbedAllowed(chatRequest('https://chat.sim.ai'), 'wf-1', 'req-1')
    expect(res).toBeNull()
    expect(mockIsWorkspaceApiExecutionEntitled).not.toHaveBeenCalled()
  })

  it('allows requests with no Origin header', async () => {
    const res = await assertChatEmbedAllowed(chatRequest(), 'wf-1', 'req-1')
    expect(res).toBeNull()
    expect(mockIsWorkspaceApiExecutionEntitled).not.toHaveBeenCalled()
  })

  it('is a no-op when billing is disabled', async () => {
    flagState.isBillingEnabled = false
    const res = await assertChatEmbedAllowed(
      chatRequest('https://evil.example.com'),
      'wf-1',
      'req-1'
    )
    expect(res).toBeNull()
    expect(mockIsWorkspaceApiExecutionEntitled).not.toHaveBeenCalled()
  })

  it('is a no-op when the gate feature flag is disabled', async () => {
    flagState.isFreeApiDeploymentGateEnabled = false
    const res = await assertChatEmbedAllowed(
      chatRequest('https://evil.example.com'),
      'wf-1',
      'req-1'
    )
    expect(res).toBeNull()
    expect(mockIsWorkspaceApiExecutionEntitled).not.toHaveBeenCalled()
  })
})