next-ai-draw-io/tests/unit/ai-providers.test.ts

import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"
import {
    getAIModel,
    isAihubmixStandardBaseURL,
    resolveBaseURL,
    supportsPromptCaching,
} from "@/lib/ai-providers"
import { extractAihubmixModelIds } from "@/lib/aihubmix-models"

describe("extractAihubmixModelIds", () => {
    it("extracts unique chat model IDs from the AIHubMix model list payload", () => {
        const models = extractAihubmixModelIds({
            data: [
                { model_id: "claude-sonnet-4-5-20250929", types: "llm" },
                { model_id: "gpt-5.1", types: "llm" },
                { model_id: "gpt-5.1", types: "llm" },
                { model_id: "gpt-image-2", types: "image_generation,llm" },
                { model_id: "cohere-rerank-v4.0", types: "rerank" },
                { model_id: "", types: "llm" },
                { types: "llm" },
            ],
        })

        expect(models).toEqual(["claude-sonnet-4-5-20250929", "gpt-5.1"])
    })

    it("returns an empty list for malformed payloads", () => {
        expect(extractAihubmixModelIds({ data: null })).toEqual([])
        expect(extractAihubmixModelIds({})).toEqual([])
        expect(extractAihubmixModelIds(null)).toEqual([])
    })
})

describe("resolveBaseURL", () => {
    const SERVER_BASE_URL = "https://server-proxy.example.com"
    const USER_BASE_URL = "https://user-proxy.example.com"
    const DEFAULT_BASE_URL = "https://api.provider.com/v1"
    const USER_API_KEY = "user-api-key-123"

    describe("when user provides their own API key", () => {
        it("uses user's baseUrl when provided", () => {
            const result = resolveBaseURL(
                USER_API_KEY,
                USER_BASE_URL,
                SERVER_BASE_URL,
                DEFAULT_BASE_URL,
            )
            expect(result).toBe(USER_BASE_URL)
        })

        it("uses default baseUrl when user provides no baseUrl", () => {
            const result = resolveBaseURL(
                USER_API_KEY,
                null,
                SERVER_BASE_URL,
                DEFAULT_BASE_URL,
            )
            expect(result).toBe(DEFAULT_BASE_URL)
        })

        it("returns undefined when user provides no baseUrl and no default exists", () => {
            const result = resolveBaseURL(
                USER_API_KEY,
                null,
                SERVER_BASE_URL,
                undefined,
            )
            expect(result).toBeUndefined()
        })

        it("does NOT use server's baseUrl even when available", () => {
            const result = resolveBaseURL(
                USER_API_KEY,
                undefined,
                SERVER_BASE_URL,
                undefined,
            )
            // Should NOT return SERVER_BASE_URL
            expect(result).not.toBe(SERVER_BASE_URL)
            expect(result).toBeUndefined()
        })

        it("prefers user's baseUrl over default", () => {
            const result = resolveBaseURL(
                USER_API_KEY,
                USER_BASE_URL,
                SERVER_BASE_URL,
                DEFAULT_BASE_URL,
            )
            expect(result).toBe(USER_BASE_URL)
        })
    })

    describe("when using server credentials (no user API key)", () => {
        it("uses user's baseUrl when provided (overrides server)", () => {
            const result = resolveBaseURL(
                null,
                USER_BASE_URL,
                SERVER_BASE_URL,
                DEFAULT_BASE_URL,
            )
            expect(result).toBe(USER_BASE_URL)
        })

        it("falls back to server's baseUrl when no user baseUrl", () => {
            const result = resolveBaseURL(
                null,
                null,
                SERVER_BASE_URL,
                DEFAULT_BASE_URL,
            )
            expect(result).toBe(SERVER_BASE_URL)
        })

        it("falls back to default when no user or server baseUrl", () => {
            const result = resolveBaseURL(
                null,
                null,
                undefined,
                DEFAULT_BASE_URL,
            )
            expect(result).toBe(DEFAULT_BASE_URL)
        })

        it("returns undefined when no baseUrl available anywhere", () => {
            const result = resolveBaseURL(null, null, undefined, undefined)
            expect(result).toBeUndefined()
        })

        it("handles undefined apiKey same as null", () => {
            const result = resolveBaseURL(
                undefined,
                null,
                SERVER_BASE_URL,
                DEFAULT_BASE_URL,
            )
            expect(result).toBe(SERVER_BASE_URL)
        })
    })

    describe("edge cases", () => {
        it("handles empty string apiKey as falsy (uses server config)", () => {
            const result = resolveBaseURL(
                "",
                null,
                SERVER_BASE_URL,
                DEFAULT_BASE_URL,
            )
            // Empty string is falsy, so should use server config
            expect(result).toBe(SERVER_BASE_URL)
        })

        it("handles empty string baseUrl as falsy", () => {
            const result = resolveBaseURL(
                USER_API_KEY,
                "",
                SERVER_BASE_URL,
                DEFAULT_BASE_URL,
            )
            // Empty string baseUrl is falsy, should fall back to default
            expect(result).toBe(DEFAULT_BASE_URL)
        })
    })
})

describe("supportsPromptCaching", () => {
    it("returns true for Claude models", () => {
        expect(supportsPromptCaching("claude-sonnet-4-5")).toBe(true)
        expect(supportsPromptCaching("anthropic.claude-3-5-sonnet")).toBe(true)
        expect(supportsPromptCaching("us.anthropic.claude-3-5-sonnet")).toBe(
            true,
        )
        expect(supportsPromptCaching("eu.anthropic.claude-3-5-sonnet")).toBe(
            true,
        )
    })

    it("returns false for non-Claude models", () => {
        expect(supportsPromptCaching("gpt-4o")).toBe(false)
        expect(supportsPromptCaching("gemini-pro")).toBe(false)
        expect(supportsPromptCaching("deepseek-chat")).toBe(false)
    })
})

vi.mock("ollama-ai-provider-v2", () => {
    const mockModel = { modelId: "test-model" }
    const mockProviderFn = vi.fn(() => mockModel)
    const mockCreateOllama = vi.fn(() => mockProviderFn)
    const mockOllama = vi.fn(() => mockModel)
    return { createOllama: mockCreateOllama, ollama: mockOllama }
})

vi.mock("@ai-sdk/deepseek", () => {
    const mockModel = { modelId: "test-model" }
    const mockProviderFn = vi.fn(() => mockModel)
    const mockCreateDeepSeek = vi.fn(() => mockProviderFn)
    const mockDeepseek = vi.fn(() => mockModel)
    return { createDeepSeek: mockCreateDeepSeek, deepseek: mockDeepseek }
})

vi.mock("@aihubmix/ai-sdk-provider", () => {
    const mockModel = { modelId: "test-model" }
    const mockProviderFn = vi.fn(() => mockModel)
    const mockCreateAihubmix = vi.fn(() => mockProviderFn)
    const mockAihubmix = vi.fn(() => mockModel)
    return { aihubmix: mockAihubmix, createAihubmix: mockCreateAihubmix }
})

describe("AIHubMix provider", () => {
    let createAihubmixMock: ReturnType<typeof vi.fn>
    const savedEnv: Record<string, string | undefined> = {}

    beforeEach(async () => {
        savedEnv.AIHUBMIX_API_KEY = process.env.AIHUBMIX_API_KEY
        savedEnv.AIHUBMIX_BASE_URL = process.env.AIHUBMIX_BASE_URL
        delete process.env.AIHUBMIX_BASE_URL

        const mod = await import("@aihubmix/ai-sdk-provider")
        createAihubmixMock = mod.createAihubmix as ReturnType<typeof vi.fn>
        createAihubmixMock.mockClear()
    })

    afterEach(() => {
        process.env.AIHUBMIX_API_KEY = savedEnv.AIHUBMIX_API_KEY
        process.env.AIHUBMIX_BASE_URL = savedEnv.AIHUBMIX_BASE_URL
    })

    it("uses AIHUBMIX_API_KEY for server configured AIHubMix", () => {
        process.env.AIHUBMIX_API_KEY = "server-aihubmix-key"

        getAIModel({
            provider: "aihubmix",
            modelId: "claude-sonnet-4-5-20250929",
        })

        expect(createAihubmixMock).toHaveBeenCalledWith({
            apiKey: "server-aihubmix-key",
            appCode: "MSBS9675",
        })
    })

    it("uses client BYOK API key for AIHubMix", () => {
        getAIModel({
            provider: "aihubmix",
            apiKey: "client-aihubmix-key",
            modelId: "gpt-5.1",
        })

        expect(createAihubmixMock).toHaveBeenCalledWith({
            apiKey: "client-aihubmix-key",
            appCode: "MSBS9675",
        })
    })

    it("recognizes AIHubMix standard endpoints", () => {
        expect(isAihubmixStandardBaseURL(undefined)).toBe(true)
        expect(isAihubmixStandardBaseURL("https://aihubmix.com")).toBe(true)
        expect(isAihubmixStandardBaseURL("https://aihubmix.com/v1/")).toBe(true)
        expect(isAihubmixStandardBaseURL("https://proxy.example.com/v1")).toBe(
            false,
        )
    })
})

describe("Kimi provider uses createDeepSeek for reasoning_content support", () => {
    let createDeepSeekMock: ReturnType<typeof vi.fn>
    const savedEnv: Record<string, string | undefined> = {}

    beforeEach(async () => {
        savedEnv.KIMI_API_KEY = process.env.KIMI_API_KEY
        savedEnv.KIMI_BASE_URL = process.env.KIMI_BASE_URL
        delete process.env.KIMI_BASE_URL

        const mod = await import("@ai-sdk/deepseek")
        createDeepSeekMock = mod.createDeepSeek as ReturnType<typeof vi.fn>
        createDeepSeekMock.mockClear()
    })

    afterEach(() => {
        process.env.KIMI_API_KEY = savedEnv.KIMI_API_KEY
        process.env.KIMI_BASE_URL = savedEnv.KIMI_BASE_URL
    })

    it("uses createDeepSeek with Kimi default base URL for reasoning_content support", () => {
        process.env.KIMI_API_KEY = "test-kimi-key"

        getAIModel({
            provider: "kimi",
            apiKey: "test-kimi-key",
            modelId: "moonshot-v1-8k",
        })

        expect(createDeepSeekMock).toHaveBeenCalledWith(
            expect.objectContaining({
                baseURL: "https://api.moonshot.cn/v1",
            }),
        )
    })

    it("uses custom base URL when provided for kimi provider", () => {
        process.env.KIMI_API_KEY = "test-kimi-key"

        getAIModel({
            provider: "kimi",
            apiKey: "test-kimi-key",
            baseUrl: "https://custom-kimi-endpoint.com/v1",
            modelId: "kimi-k2.6",
        })

        expect(createDeepSeekMock).toHaveBeenCalledWith(
            expect.objectContaining({
                baseURL: "https://custom-kimi-endpoint.com/v1",
            }),
        )
    })
})

describe("Ollama API key security", () => {
    let createOllamaMock: ReturnType<typeof vi.fn>
    const savedEnv: Record<string, string | undefined> = {}

    beforeEach(async () => {
        savedEnv.OLLAMA_API_KEY = process.env.OLLAMA_API_KEY
        savedEnv.OLLAMA_BASE_URL = process.env.OLLAMA_BASE_URL
        delete process.env.OLLAMA_BASE_URL

        const mod = await import("ollama-ai-provider-v2")
        createOllamaMock = mod.createOllama as ReturnType<typeof vi.fn>
        createOllamaMock.mockClear()
    })

    afterEach(() => {
        process.env.OLLAMA_API_KEY = savedEnv.OLLAMA_API_KEY
        process.env.OLLAMA_BASE_URL = savedEnv.OLLAMA_BASE_URL
    })

    it("applies server OLLAMA_API_KEY when no client baseUrl is provided", () => {
        process.env.OLLAMA_API_KEY = "server-secret-key"

        getAIModel({ provider: "ollama", modelId: "llama2" })

        expect(createOllamaMock).toHaveBeenCalledWith(
            expect.objectContaining({
                headers: { Authorization: "Bearer server-secret-key" },
            }),
        )
    })

    it("does NOT leak server OLLAMA_API_KEY when client provides a custom baseUrl", () => {
        process.env.OLLAMA_API_KEY = "server-secret-key"

        // When server has OLLAMA_API_KEY, the SSRF guard rejects
        // client-provided baseUrl without an apiKey outright
        expect(() =>
            getAIModel({
                provider: "ollama",
                baseUrl: "https://evil-server.com",
                modelId: "llama2",
            }),
        ).toThrow("API key is required")
    })

    it("uses client API key when client provides both baseUrl and apiKey", () => {
        process.env.OLLAMA_API_KEY = "server-secret-key"

        getAIModel({
            provider: "ollama",
            baseUrl: "https://my-ollama.com",
            apiKey: "client-key",
            modelId: "llama2",
        })

        expect(createOllamaMock).toHaveBeenCalledWith(
            expect.objectContaining({
                baseURL: "https://my-ollama.com",
                headers: { Authorization: "Bearer client-key" },
            }),
        )
    })

    it("applies both server OLLAMA_BASE_URL and OLLAMA_API_KEY when no client overrides", () => {
        process.env.OLLAMA_BASE_URL = "https://cloud.ollama.com"
        process.env.OLLAMA_API_KEY = "server-key"

        getAIModel({ provider: "ollama", modelId: "llama2" })

        expect(createOllamaMock).toHaveBeenCalledWith(
            expect.objectContaining({
                baseURL: "https://cloud.ollama.com",
                headers: { Authorization: "Bearer server-key" },
            }),
        )
    })

    it("works when OLLAMA_API_KEY is set but OLLAMA_BASE_URL is not", () => {
        process.env.OLLAMA_API_KEY = "server-key"
        delete process.env.OLLAMA_BASE_URL

        getAIModel({ provider: "ollama", modelId: "llama2" })

        expect(createOllamaMock).toHaveBeenCalledTimes(1)
        const callArgs = createOllamaMock.mock.calls[0][0]
        expect(callArgs).not.toHaveProperty("baseURL")
        expect(callArgs).toEqual(
            expect.objectContaining({
                headers: { Authorization: "Bearer server-key" },
            }),
        )
    })

    it("allows client custom baseUrl without apiKey when no server OLLAMA_API_KEY", () => {
        delete process.env.OLLAMA_API_KEY

        getAIModel({
            provider: "ollama",
            baseUrl: "https://my-ollama.com",
            modelId: "llama2",
        })

        expect(createOllamaMock).toHaveBeenCalledTimes(1)
        const callArgs = createOllamaMock.mock.calls[0][0]
        expect(callArgs.baseURL).toBe("https://my-ollama.com")
        expect(callArgs).not.toHaveProperty("headers")
    })
})