path: root/lua/99/providers.lua

--- @class _99.Providers.Observer
--- @field on_stdout fun(line: string): nil
--- @field on_stderr fun(line: string): nil
--- @field on_complete fun(status: _99.Prompt.EndingState, res: string|table): nil
--- @field on_start fun(): nil

--- @param fn fun(...: any): nil
--- @return fun(...: any): nil
local function once(fn)
    local called = false
    return function(...)
        if called then
            return
        end
        called = true
        fn(...)
    end
end

--- @class _A4.OpenAIResponse
--- @field completion _A4.OpenAIResponse.Completion
--- @field metadata _A4.OpenAIResponse.MetaData
--- @field id _A4.OpenAIResponse.Id
--- @field timings _A4.OpenAIResponse.Timings

--- @class _A4.OpenAIResponse.Completion
--- @field finishReason string
--- @field index integer
--- @field message _A4.OpenAIResponse.Completion.Message

--- @class _A4.OpenAIResponse.Completion.Message
--- @field role string
--- @field content string

--- @class _A4.OpenAIResponse.MetaData
--- @field created integer
--- @field model string
--- @field systemFingerPrint string
--- @field object string
--- @field usage _A4.OpenAIResponse.MetaData.Usage

--- @class _A4.OpenAIResponse.MetaData.Usage
--- @field completionTks integer
--- @field promptTks integer
--- @field totalTks integer
--- @field promptTksDetails _A4.OpenAIResponse.MetaData.Usage.PromptTksDetails

--- @class _A4.OpenAIResponse.MetaData.Usage.PromptTksDetails
--- @field cachedTks integer

--- @class _A4.OpenAIResponse.Id
--- @field id string

--- @class _A4.OpenAIResponse.Timings
--- @field cacheN integer
--- @field promptN integer
--- @field promptMs number
--- @field promptPerTkMs number
--- @field promptPerSec number
--- @field predictedN integer
--- @field predictedMs number
--- @field predictedPerTkMs number
--- @field predictedPerSec number

--- @class _99.Providers.BaseProvider
--- @field _build_command fun(self: _99.Providers.BaseProvider, query: string, context: _99.Prompt): string[]
--- @field _get_provider_name fun(self: _99.Providers.BaseProvider): string
--- @field _get_default_model fun(): string
local BaseProvider = {}

--- @param callback fun(models: string[]|nil, err: string|nil): nil
function BaseProvider.fetch_models(callback)
    callback(nil, "This provider does not support listing models")
end

--- @param data string
--- @return string
function BaseProvider._clean_response(data)
    local _data, _ = data:gsub("```[a-z]*\n", ""):gsub("```", "")
    return _data
end

-- TODO: Rewrite this...
--- @param context _99.Prompt
function BaseProvider:_retrieve_response(context)
    local logger = context.logger:set_area(self:_get_provider_name())

    if self._get_provider_name ~= "BareMetalProvider" then
        local tmp = context.tmp_file
        local success, result = pcall(function()
            return vim.fn.readfile(tmp)
        end)

        if not success then
            logger:error(
                "retrieve_results: failed to read file",
                "tmp_name",
                tmp,
                "error",
                result
            )
            return false, ""
        end

        local str = table.concat(result, "\n")
        logger:debug("retrieve_results", "results", str)

        return true, str
    else
        assert(
            context.req_response,
            "_retrieve_response: BareMetalProvider has no response"
        )
        logger:debug("retrieve_results", "results", context.req_response)
        return true, context.req_response
    end
end

--- @param query string
--- @param context _99.Prompt
--- @param observer _99.Providers.Observer
function BaseProvider:make_request(query, context, observer)
    observer.on_start()

    local logger = context.logger:set_area(self:_get_provider_name())
    logger:debug("make_request", "reason", context.operation)
    logger:debug("make_request", "endpoint", context.endpoint)

    local once_complete = once(
        --- @param status "success" | "failed" | "cancelled"
        --- @param resp string
        function(status, resp)
            observer.on_complete(status, resp)
        end
    )

    local command = self:_build_command(query, context)
    logger:debug("make_request", "command", command)

    local proc = vim.system(
        command,
        { text = true },
        vim.schedule_wrap(function(obj)
            logger:debug("exit callback fired", "code", tostring(obj.code))
            logger:debug("schedule_wrap", "stdout", tostring(obj.stdout))
            logger:debug("schedule_wrap", "stderr", tostring(obj.stderr))

            if context:is_cancelled() then
                once_complete("cancelled", "")
                return
            end

            if obj.code ~= 0 then
                once_complete("failed", obj.stderr or "")
                return
            end

            once_complete("success", obj.stdout)
        end)
    )

    logger:debug("proc spawned", "proc_id", tostring(proc))
    context:_set_process(proc)
end

--- @class OpenCodeProvider : _99.Providers.BaseProvider
local OpenCodeProvider = setmetatable({}, { __index = BaseProvider })

--- @param query string
--- @param context _99.Prompt
--- @return string[]
function OpenCodeProvider._build_command(_, query, context)
    return {
        "opencode",
        "run",
        "--agent",
        "build",
        "-m",
        context.model,
        query,
    }
end

--- @return string
function OpenCodeProvider._get_provider_name()
    return "OpenCodeProvider"
end

--- @return string
function OpenCodeProvider._get_default_model()
    return "opencode/claude-sonnet-4-5"
end

function OpenCodeProvider.fetch_models(callback)
    vim.system({ "opencode", "models" }, { text = true }, function(obj)
        vim.schedule(function()
            if obj.code ~= 0 then
                callback(nil, "Failed to fetch models from opencode")
                return
            end
            local models = vim.split(obj.stdout, "\n", { trimempty = true })
            callback(models, nil)
        end)
    end)
end

--- @class ClaudeCodeProvider : _99.Providers.BaseProvider
local ClaudeCodeProvider = setmetatable({}, { __index = BaseProvider })

--- @param query string
--- @param context _99.Prompt
--- @return string[]
function ClaudeCodeProvider._build_command(_, query, context)
    return {
        "claude",
        "--dangerously-skip-permissions",
        "--model",
        context.model,
        "--print",
        query,
    }
end

--- @return string
function ClaudeCodeProvider._get_provider_name()
    return "ClaudeCodeProvider"
end

--- @return string
function ClaudeCodeProvider._get_default_model()
    return "claude-sonnet-4-5"
end

-- TODO: the claude CLI has no way to list available models.
-- We could use the Anthropic API (https://docs.anthropic.com/en/api/models)
-- but that requires the user to have an ANTHROPIC_API_KEY set which isn't ideal.
-- Until Anthropic adds a CLI command for this, we have to hardcode the list here.
-- See https://github.com/anthropics/claude-code/issues/12612
function ClaudeCodeProvider.fetch_models(callback)
    callback({
        "claude-opus-4-6",
        "claude-sonnet-4-5",
        "claude-haiku-4-5",
        "claude-opus-4-5",
        "claude-opus-4-1",
        "claude-sonnet-4-0",
        "claude-opus-4-0",
        "claude-3-7-sonnet-latest",
    }, nil)
end

--- @class CursorAgentProvider : _99.Providers.BaseProvider
local CursorAgentProvider = setmetatable({}, { __index = BaseProvider })

--- @param query string
--- @param context _99.Prompt
--- @return string[]
function CursorAgentProvider._build_command(_, query, context)
    return { "cursor-agent", "--model", context.model, "--print", query }
end

--- @return string
function CursorAgentProvider._get_provider_name()
    return "CursorAgentProvider"
end

--- @return string
function CursorAgentProvider._get_default_model()
    return "sonnet-4.5"
end

function CursorAgentProvider.fetch_models(callback)
    vim.system({ "cursor-agent", "models" }, { text = true }, function(obj)
        vim.schedule(function()
            if obj.code ~= 0 then
                callback(nil, "Failed to fetch models from cursor-agent")
                return
            end
            local models = {}
            for _, line in
                ipairs(vim.split(obj.stdout, "\n", { trimempty = true }))
            do
                -- `cursor-agent models` outputs lines like "model-id - description",
                -- so we grab everything before the first " - " separator
                local id = line:match("^(%S+)%s+%-")
                if id then
                    table.insert(models, id)
                end
            end
            callback(models, nil)
        end)
    end)
end

--- @class KiroProvider : _99.Providers.BaseProvider
local KiroProvider = setmetatable({}, { __index = BaseProvider })

--- @param query string
--- @param context _99.Prompt
--- @return string[]
function KiroProvider._build_command(_, query, context)
    return {
        "kiro-cli",
        "chat",
        "--no-interactive",
        "--model",
        context.model,
        "--trust-all-tools",
        query,
    }
end

--- @return string
function KiroProvider._get_provider_name()
    return "KiroProvider"
end

--- @return string
function KiroProvider._get_default_model()
    return "claude-sonnet-4.5"
end

--- @class GeminiCLIProvider : _99.Providers.BaseProvider
local GeminiCLIProvider = setmetatable({}, { __index = BaseProvider })

--- @param query string
--- @param context _99.Prompt
--- @return string[]
function GeminiCLIProvider._build_command(_, query, context)
    return {
        "gemini",
        "--approval-mode",
        -- Allow writing to temp files by default. See:
        -- https://geminicli.com/docs/core/policy-engine/#default-policies
        "auto_edit",
        "--model",
        context.model,
        "--prompt",
        query,
    }
end

--- @return string
function GeminiCLIProvider._get_provider_name()
    return "GeminiCLIProvider"
end

--- @return string
function GeminiCLIProvider._get_default_model()
    -- Default to auto-routing between pro and flash. See:
    -- https://geminicli.com/docs/cli/model/
    return "auto"
end

--- @class BareMetalProvider : _99.Providers.BaseProvider
--- @field _clean_response fun(data : string): string
--- @field _parse_openai_response fun(raw: table): _A4.OpenAIResponse
local BareMetalProvider = setmetatable({}, { __index = BaseProvider })

--- @param query string
--- @param context _99.Prompt
--- @return string[]
function BareMetalProvider._build_command(_, query, context)
    return {
        "curl",
        "-s",
        context.endpoint,
        "-H",
        "Content-Type: application/json",
        "-d",
        vim.json.encode({
            messages = {
                {
                    role = "system",
                    content = "You are a strict code completion backend for Neovim. Your input will be a code snippet, function signature, or a comment requesting code. CRITICAL DIRECTIONS: 1. Output ONLY valid, executable programming code. 2. Do NOT wrap your response in markdown code blocks. 3. Do NOT include any conversational filler, explanations, greetings, or sign-offs. 4. If you cannot fulfill the request, output nothing or a code comment explaining why.",
                },
                {
                    role = "user",
                    content = query,
                },
            },
            temperature = 0.0,
            stream = false,
        }),
    }
end

--- @return string
function BareMetalProvider._get_provider_name()
    return "BareMetalProvider"
end

--- @return string
function BareMetalProvider._get_default_model()
    return "Qwen2.5-Coder-3B-Instruct"
end

function BareMetalProvider.fetch_models(callback)
    callback({
        "Qwen2.5-Coder-3B-Instruct",
    }, nil)
end

--- @param raw table
--- @return _A4.OpenAIResponse
function BareMetalProvider._parse_openai_response(raw)
    local choice = raw.choices[1]

    --- @type _A4.OpenAIResponse
    local response = {
        completion = {
            finishReason = choice.finish_reason,
            index = choice.index,
            message = {
                role = choice.message.role,
                content = choice.message.content,
            },
        },
        metadata = {
            created = raw.created,
            model = raw.model,
            systemFingerPrint = raw.system_fingerprint,
            object = raw.object,
            usage = {
                completionTks = raw.usage.completion_tokens,
                promptTks = raw.usage.prompt_tokens,
                totalTks = raw.usage.total_tokens,
                promptTksDetails = {
                    cachedTks = raw.usage.prompt_tokens_details.cached_tokens,
                },
            },
        },
        id = {
            id = raw.id,
        },
        timings = {
            cacheN = raw.timings.cache_n,
            promptN = raw.timings.prompt_n,
            promptMs = raw.timings.prompt_ms,
            promptPerTkMs = raw.timings.prompt_per_token_ms,
            promptPerSec = raw.timings.prompt_per_second,
            predictedN = raw.timings.predicted_n,
            predictedMs = raw.timings.predicted_ms,
            predictedPerTkMs = raw.timings.predicted_per_token_ms,
            predictedPerSec = raw.timings.predicted_per_second,
        },
    }

    return response
end

return {
    BaseProvider = BaseProvider,
    OpenCodeProvider = OpenCodeProvider,
    ClaudeCodeProvider = ClaudeCodeProvider,
    CursorAgentProvider = CursorAgentProvider,
    KiroProvider = KiroProvider,
    GeminiCLIProvider = GeminiCLIProvider,
    BareMetalProvider = BareMetalProvider,
}