insights/services/mcp/tests/unit/tool-schema-snapshots.test.ts

import { mkdir, readdir, rm } from 'node:fs/promises'
import path from 'node:path'
import { fileURLToPath } from 'node:url'
import { format } from 'oxfmt'
import { describe, expect, it } from 'vitest'
import { z } from 'zod'

import { OAUTH_SCOPES_HIDDEN } from '@/lib/constants'
import { SessionManager } from '@/lib/SessionManager'
import { getToolsFromContext } from '@/tools'
import type { Context } from '@/tools/types'

function createMockContext(): Context {
    return {
        api: {} as any,
        cache: {} as any,
        env: {
            MCP_APPS_BASE_URL: undefined,
            INSIGHTS_ANALYTICS_API_KEY: undefined,
            INSIGHTS_ANALYTICS_HOST: undefined,
            INSIGHTS_API_BASE_URL: undefined,
            INSIGHTS_PUBLIC_URL: undefined,
            INSIGHTS_MCP_APPS_ANALYTICS_BASE_URL: undefined,
            INSIGHTS_UI_APPS_TOKEN: undefined,
        },
        stateManager: {
            // Staff-only tools require their OAuth-hidden scope explicitly on the key
            // (`*` alone does not match) plus a staff user, so grant both here to keep
            // their schemas in the snapshot surface.
            getApiKey: async () => ({ scopes: ['*', ...OAUTH_SCOPES_HIDDEN] }),
            getAiConsentGiven: async () => true,
            getUser: async () => ({ is_staff: true }),
        } as any,
        sessionManager: new SessionManager({} as any),
        getDistinctId: async () => 'test-distinct-id',
        trackEvent: async () => {},
    }
}

function deepSortKeys(value: unknown): unknown {
    if (Array.isArray(value)) {
        return value.map((item) => deepSortKeys(item))
    }

    if (value && typeof value === 'object') {
        const obj = value as Record<string, unknown>
        const sortedEntries = Object.entries(obj)
            .sort(([a], [b]) => a.localeCompare(b))
            .map(([key, nestedValue]) => [key, deepSortKeys(nestedValue)] as const)
        return Object.fromEntries(sortedEntries)
    }

    return value
}

async function listSnapshotFiles(root: string): Promise<string[]> {
    try {
        const files = await readdir(root)
        return files.filter((file) => file.endsWith('.json')).map((file) => path.join(root, file))
    } catch {
        // Directory might not exist yet.
        return []
    }
}

function isSnapshotUpdateAll(): boolean {
    // Vitest strips CLI flags from process.argv in worker threads, so we read
    // the snapshot-update mode from vitest's internal state instead.
    const state = expect.getState() as unknown as { snapshotState?: { _updateSnapshot?: string } }
    return state.snapshotState?._updateSnapshot === 'all'
}

async function formatSnapshotJson(snapshotPath: string, schema: unknown): Promise<string> {
    const content = `${JSON.stringify(schema, null, 4)}\n`
    const result = await format(snapshotPath, content, { tabWidth: 4, printWidth: 120 })

    if (result.errors.length > 0) {
        const errorMessage = result.errors.map((error) => error.message ?? 'unknown formatting error').join('; ')
        throw new Error(`Failed formatting snapshot ${snapshotPath}: ${errorMessage}`)
    }

    return result.code
}

describe('Tool schema snapshots', () => {
    const __filename = fileURLToPath(import.meta.url)
    const __dirname = path.dirname(__filename)
    const context = createMockContext()
    it('snapshots runtime tool schemas', async () => {
        const shouldUpdateSnapshots = isSnapshotUpdateAll()
        const root = path.resolve(__dirname, '__snapshots__', 'tool-schemas')
        // Enable flag-gated tools we snapshot here: tracing (APM spans), tasks, loops,
        // dashboard-widgets. Other flag-gated tools (logs-alerts, visual-review, etc.) stay off to keep the surface stable.
        // agent-feedback is always_available and no longer flag-gated, so it appears regardless.
        const featureFlags = {
            tracing: true,
            tasks: true,
            loops: true,
            'dashboard-widgets': true,
        }
        const tools = [...(await getToolsFromContext(context, { featureFlags }))].sort((a, b) =>
            a.name.localeCompare(b.name)
        )

        expect(tools.length).toBeGreaterThan(0)

        await mkdir(root, { recursive: true })

        const expectedPaths = new Set<string>()

        for (const tool of tools) {
            const jsonSchema = deepSortKeys(z.toJSONSchema(tool.schema, { io: 'input', reused: 'inline' }))
            const snapshotPath = path.join(root, `${tool.name}.json`)
            expectedPaths.add(snapshotPath)
            const content = await formatSnapshotJson(snapshotPath, jsonSchema)
            await expect(content).toMatchFileSnapshot(snapshotPath)
        }

        const existingPaths = await listSnapshotFiles(root)
        const stalePaths = existingPaths.filter((existingPath) => !expectedPaths.has(existingPath)).sort()

        if (shouldUpdateSnapshots) {
            for (const stalePath of stalePaths) {
                await rm(stalePath)
            }
        } else if (stalePaths.length > 0) {
            throw new Error(
                `Found stale snapshot files (run vitest -u to clean them):\n${stalePaths
                    .map((file) => `- ${file}`)
                    .join('\n')}`
            )
        }
    })
})