fix(vscode): support direct autocomplete FIM providers

2026-09-24 16:02:55 +08:00 · 2026-05-25 16:38:29 +02:00
parent 39affcc140
commit 33eb706c58
10 changed files with 363 additions and 44 deletions
@@ -0,0 +1,5 @@
+---
+"kilo-code": patch
+---
+
+Use connected Mistral and Inception provider API keys for autocomplete FIM requests before falling back to Kilo Gateway.
@@ -1,13 +1,17 @@
 # Source Code Links

 <!-- Auto-generated by script/extract-source-links.ts — DO NOT EDIT -->
-<!-- 86 unique URLs extracted from extension and CLI source -->
+<!-- 89 unique URLs extracted from extension and CLI source -->

 - <https://api.apertis.ai/v1>
  <!-- packages/opencode/src/provider/model-cache.ts -->
  <!-- packages/opencode/src/provider/models.ts -->
+- <https://api.inceptionlabs.ai/v1/fim/completions>
+  <!-- packages/kilo-vscode/src/services/autocomplete/fim.ts -->
 - <https://api.kilo.ai>
  <!-- packages/opencode/src/cli/cmd/github.ts -->
+- <https://api.mistral.ai/v1/fim/completions>
+  <!-- packages/kilo-vscode/src/services/autocomplete/fim.ts -->
 - <https://app.kilo.ai>
  <!-- packages/opencode/src/kilocode/kilo-commands.tsx -->
 - <https://app.kilo.ai/claw>
@@ -32,6 +36,8 @@
  <!-- packages/kilo-vscode/src/agent-manager/WorktreeManager.ts -->
 - <https://cloudflare.com/cdn-cgi/trace>
  <!-- packages/opencode/src/session/network.ts -->
+- <https://codestral.mistral.ai/v1/fim/completions>
+  <!-- packages/kilo-vscode/src/services/autocomplete/fim.ts -->
 - <https://cookbook.openai.com/examples/using_logprobs>
  <!-- packages/opencode/src/provider/sdk/copilot/responses/openai-responses-language-model.ts -->
 - <https://docs.github.com/en/actions/how-tos/security-for-github-actions/security-hardening-your-deployments/configuring-openid-connect-in-amazon-web-services>
@@ -753,11 +753,15 @@
          "type": "string",
          "enum": [
            "mistralai/codestral-2508",
-            "inception/mercury-edit-2"
+            "inception/mercury-edit-2",
+            "mistral/codestral-2508",
+            "inception-direct/mercury-edit-2"
          ],
          "enumDescriptions": [
-            "Codestral by Mistral AI (default)",
-            "Mercury Edit 2 by Inception"
+            "Codestral via Kilo Gateway (default)",
+            "Mercury Edit 2 via Kilo Gateway",
+            "Codestral via your connected Mistral provider API key",
+            "Mercury Edit 2 via your connected Inception provider API key"
          ],
          "description": "Model to use for inline autocomplete suggestions"
        },
@@ -319,7 +319,7 @@ export class AutocompleteServiceManager {
  }

  private getCurrentModelName(): string {
-    return this.inlineCompletionProvider.getModelId()
+    return getAutocompleteModel(this.inlineCompletionProvider.getModelId()).label
  }

  private getCurrentProviderName(): string {
@@ -65,36 +65,14 @@ export class FimPromptBuilder {
    signal?: AbortSignal,
  ): Promise<FimCompletionResult> {
    const { formattedPrefix, prunedSuffix, autocompleteInput } = prompt
-    let perflog = ""
-    const logtime = (() => {
-      let timestamp = performance.now()
-      return (msg: string) => {
-        const baseline = timestamp
-        timestamp = performance.now()
-        perflog += `${msg}: ${timestamp - baseline}\n`
-      }
-    })()
-
-    logtime("snippets")
-
-    console.log("[FIM] formattedPrefix:", formattedPrefix)
-
    let response = ""
    const onChunk = (text: string) => {
      response += text
    }
-    logtime("prep fim")
    const usageInfo = await generateFim(connection, modelId, formattedPrefix, prunedSuffix, onChunk, signal)
-    logtime("fim network")
-    console.log("[FIM] response:", response)

    const fillInAtCursorSuggestion = processSuggestion(response)

-    if (fillInAtCursorSuggestion.text) {
-      console.info("Final FIM suggestion:", fillInAtCursorSuggestion)
-    }
-    logtime("processSuggestion")
-    console.log(perflog + `lengths: ${formattedPrefix.length + prunedSuffix.length}\n`)
    return {
      suggestion: fillInAtCursorSuggestion,
      cost: usageInfo.cost,
@@ -2,6 +2,173 @@ import { ResponseMetaData } from "./types"
 import type { KiloConnectionService } from "../cli-backend"
 import { getAutocompleteModel } from "../../shared/autocomplete-models"

+const FIM_MAX_TOKENS = 256
+const MISTRAL_FIM_URL = "https://api.mistral.ai/v1/fim/completions"
+const CODESTRAL_FIM_URL = "https://codestral.mistral.ai/v1/fim/completions"
+const INCEPTION_FIM_URL = "https://api.inceptionlabs.ai/v1/fim/completions"
+
+type FimProvider = "mistral" | "inception"
+
+interface DirectFimTarget {
+  provider: FimProvider
+  model: string
+  urls: string[]
+}
+
+interface ProviderItem {
+  id: string
+  key?: string
+  options?: Record<string, unknown>
+}
+
+interface ProviderListResponse {
+  all: ProviderItem[]
+}
+
+interface DirectFimChunk {
+  choices?: Array<{
+    delta?: { content?: string }
+    text?: string
+  }>
+  usage?: {
+    prompt_tokens?: number
+    completion_tokens?: number
+  }
+}
+
+interface DirectFimOptions {
+  apiKey: string
+  target: DirectFimTarget
+  prefix: string
+  suffix: string
+  temperature: number
+  onChunk: (text: string) => void
+  signal?: AbortSignal
+  fetchImpl?: typeof fetch
+}
+
+export function getDirectFimTarget(model: string): DirectFimTarget | null {
+  const info = getAutocompleteModel(model)
+  if (info.directProvider === "mistral") {
+    return { provider: "mistral", model: info.requestModel, urls: [MISTRAL_FIM_URL, CODESTRAL_FIM_URL] }
+  }
+  if (info.directProvider === "inception") {
+    return { provider: "inception", model: info.requestModel, urls: [INCEPTION_FIM_URL] }
+  }
+  return null
+}
+
+async function resolveProviderKey(connectionService: KiloConnectionService, provider: FimProvider): Promise<string | null> {
+  const client = await connectionService.getClientAsync()
+  const result = await client.provider.list({}, { throwOnError: true })
+  const data = result.data as ProviderListResponse | undefined
+  const item = data?.all.find((p) => p.id === provider)
+  const key = item?.options?.apiKey
+  return item?.key ?? (typeof key === "string" ? key : null)
+}
+
+function extractDirectFimContent(chunk: DirectFimChunk): string {
+  const choice = chunk.choices?.[0]
+  return choice?.delta?.content ?? choice?.text ?? ""
+}
+
+function parseDirectFimEvent(data: string): DirectFimChunk | null {
+  if (data === "[DONE]") return null
+  return JSON.parse(data) as DirectFimChunk
+}
+
+function parseDirectFimLine(line: string): string | null {
+  const trimmed = line.trim()
+  if (!trimmed.startsWith("data:")) return null
+  return trimmed.slice("data:".length).trim()
+}
+
+function handleDirectFimLine(
+  line: string,
+  onChunk: (text: string) => void,
+  usage: { inputTokens: number; outputTokens: number },
+) {
+  const data = parseDirectFimLine(line)
+  if (!data) return
+  const event = parseDirectFimEvent(data)
+  if (!event) return
+  const content = extractDirectFimContent(event)
+  if (content) onChunk(content)
+  usage.inputTokens = event.usage?.prompt_tokens ?? usage.inputTokens
+  usage.outputTokens = event.usage?.completion_tokens ?? usage.outputTokens
+}
+
+export async function generateDirectFim(options: DirectFimOptions): Promise<ResponseMetaData> {
+  const urls = [...options.target.urls]
+  const [url] = urls
+  if (!url) throw new Error("FIM request failed: 500 missing provider endpoint")
+  return generateDirectFimWithUrl(options, url, urls.slice(1))
+}
+
+async function generateDirectFimWithUrl(
+  options: DirectFimOptions,
+  url: string,
+  fallbacks: string[],
+): Promise<ResponseMetaData> {
+  const fetchImpl = options.fetchImpl ?? fetch
+  console.info(`[FIM] request provider=${options.target.provider} model=${options.target.model} url=${url}`)
+  const res = await fetchImpl(url, {
+    method: "POST",
+    headers: {
+      "Content-Type": "application/json",
+      Authorization: `Bearer ${options.apiKey}`,
+    },
+    body: JSON.stringify({
+      model: options.target.model,
+      prompt: options.prefix,
+      suffix: options.suffix,
+      max_tokens: FIM_MAX_TOKENS,
+      temperature: options.temperature,
+      stream: true,
+    }),
+    signal: options.signal,
+  })
+
+  if (!res.ok) {
+    const body = await res.text().catch(() => "")
+    const [next] = fallbacks
+    if (res.status === 401 && next) return generateDirectFimWithUrl(options, next, fallbacks.slice(1))
+    throw new Error(`FIM request failed: ${res.status} ${res.statusText}: ${body}`)
+  }
+
+  if (!res.body) throw new Error("FIM request failed: 500 empty response body")
+
+  const reader = res.body.getReader()
+  const decoder = new TextDecoder()
+  const usage = { inputTokens: 0, outputTokens: 0 }
+  let pending = ""
+
+  while (true) {
+    const chunk = await reader.read()
+    if (chunk.done) break
+    pending += decoder.decode(chunk.value, { stream: true })
+    const lines = pending.split("\n")
+    pending = lines.pop() ?? ""
+
+    for (const line of lines) {
+      handleDirectFimLine(line, options.onChunk, usage)
+    }
+  }
+
+  pending += decoder.decode()
+  handleDirectFimLine(pending, options.onChunk, usage)
+
+  reader.releaseLock()
+
+  return {
+    cost: 0,
+    inputTokens: usage.inputTokens,
+    outputTokens: usage.outputTokens,
+    cacheWriteTokens: 0,
+    cacheReadTokens: 0,
+  }
+}
+
 /**
 * Generate a FIM (Fill-in-the-Middle) completion via the CLI backend.
 * Uses the SDK's kilo.fim() SSE endpoint which handles auth and streaming.
@@ -17,6 +184,25 @@ export async function generateFim(
  signal?: AbortSignal,
 ): Promise<ResponseMetaData> {
  const client = await connectionService.getClientAsync()
+  const info = getAutocompleteModel(modelId)
+  const target = getDirectFimTarget(modelId)
+  const key = info.directProvider ? await resolveProviderKey(connectionService, info.directProvider).catch(() => null) : null
+
+  if (target && !key) {
+    throw new Error(`FIM request failed: 401 Missing ${target.provider} provider API key`)
+  }
+
+  if (target && key) {
+    return generateDirectFim({
+      apiKey: key,
+      target,
+      prefix,
+      suffix,
+      temperature: info.temperature,
+      onChunk,
+      signal,
+    })
+  }

  let cost = 0
  let inputTokens = 0
@@ -27,14 +213,15 @@ export async function generateFim(
  // ends the stream. Without this, errors never reach ErrorBackoff.
  let sseError: Error | undefined

-  const temp = getAutocompleteModel(modelId).temperature
+  const temp = info.temperature
+  console.info(`[FIM] request provider=kilo model=${info.requestModel} url=/kilo/fim`)

  const { stream } = await client.kilo.fim(
    {
      prefix,
      suffix,
-      model: modelId,
-      maxTokens: 256,
+      model: info.requestModel,
+      maxTokens: FIM_MAX_TOKENS,
      temperature: temp,
    },
    {
@@ -7,12 +7,20 @@
 */

 export interface AutocompleteModelDef {
-  /** Full model ID sent to the gateway, e.g. "mistralai/codestral-2508" */
+  /** Stable setting value. */
  readonly id: string
+  /** Model ID displayed under the selector provider group. */
+  readonly modelID: string
  /** Human-readable label shown in the settings dropdown */
  readonly label: string
+  /** Provider ID used by the selector group. */
+  readonly providerID: string
  /** Provider display name for status bar / telemetry */
  readonly provider: string
+  /** Full model ID sent to the FIM API. */
+  readonly requestModel: string
+  /** Provider key to use for direct BYOK FIM. Empty means Kilo Gateway. */
+  readonly directProvider?: "mistral" | "inception"
  /** FIM request temperature */
  readonly temperature: number
 }
@@ -20,14 +28,40 @@ export interface AutocompleteModelDef {
 const models: AutocompleteModelDef[] = [
  {
    id: "mistralai/codestral-2508",
-    label: "Codestral (Mistral AI)",
-    provider: "Mistral AI",
+    modelID: "mistralai/codestral-2508",
+    label: "Codestral",
+    providerID: "kilo",
+    provider: "Kilo Gateway",
+    requestModel: "mistralai/codestral-2508",
    temperature: 0.2,
  },
  {
    id: "inception/mercury-edit-2",
-    label: "Mercury Edit 2 (Inception)",
+    modelID: "inception/mercury-edit-2",
+    label: "Mercury Edit 2",
+    providerID: "kilo",
+    provider: "Kilo Gateway",
+    requestModel: "inception/mercury-edit-2",
+    temperature: 0,
+  },
+  {
+    id: "mistral/codestral-2508",
+    modelID: "codestral-2508",
+    label: "Codestral",
+    providerID: "mistral",
+    provider: "Mistral",
+    requestModel: "codestral-2508",
+    directProvider: "mistral",
+    temperature: 0.2,
+  },
+  {
+    id: "inception-direct/mercury-edit-2",
+    modelID: "mercury-edit-2",
+    label: "Mercury Edit 2",
+    providerID: "inception",
    provider: "Inception",
+    requestModel: "mercury-edit-2",
+    directProvider: "inception",
    temperature: 0,
  },
 ]
@@ -0,0 +1,94 @@
+import { describe, expect, it } from "bun:test"
+import { generateDirectFim, getDirectFimTarget } from "../../src/services/autocomplete/fim"
+
+function stream(text: string) {
+  return new ReadableStream<Uint8Array>({
+    start(controller) {
+      controller.enqueue(new TextEncoder().encode(text))
+      controller.close()
+    },
+  })
+}
+
+describe("direct autocomplete FIM", () => {
+  it("maps autocomplete models to direct provider endpoints", () => {
+    expect(getDirectFimTarget("mistralai/codestral-2508")).toBeNull()
+    expect(getDirectFimTarget("inception/mercury-edit-2")).toBeNull()
+    expect(getDirectFimTarget("mistral/codestral-2508")).toEqual({
+      provider: "mistral",
+      model: "codestral-2508",
+      urls: ["https://api.mistral.ai/v1/fim/completions", "https://codestral.mistral.ai/v1/fim/completions"],
+    })
+    expect(getDirectFimTarget("inception-direct/mercury-edit-2")).toEqual({
+      provider: "inception",
+      model: "mercury-edit-2",
+      urls: ["https://api.inceptionlabs.ai/v1/fim/completions"],
+    })
+    expect(getDirectFimTarget("openai/gpt-5")).toBeNull()
+  })
+
+  it("streams provider FIM chunks and returns usage", async () => {
+    const chunks = [
+      'data: {"choices":[{"delta":{"content":"hel"}}]}\n\n',
+      'data: {"choices":[{"delta":{"content":"lo"}}],"usage":{"prompt_tokens":3,"completion_tokens":2}}\n\n',
+      "data: [DONE]\n\n",
+    ].join("")
+    const calls: RequestInit[] = []
+    const fetchImpl: typeof fetch = async (_url, init) => {
+      calls.push(init ?? {})
+      return new Response(stream(chunks), { status: 200 })
+    }
+    const text: string[] = []
+    const usage = await generateDirectFim({
+      apiKey: "test-key",
+      target: getDirectFimTarget("inception-direct/mercury-edit-2")!,
+      prefix: "const value = ",
+      suffix: "\n",
+      temperature: 0,
+      onChunk: (chunk) => text.push(chunk),
+      fetchImpl,
+    })
+
+    expect(text.join("")).toBe("hello")
+    expect(usage).toEqual({
+      cost: 0,
+      inputTokens: 3,
+      outputTokens: 2,
+      cacheWriteTokens: 0,
+      cacheReadTokens: 0,
+    })
+    expect(calls[0]?.headers).toEqual({
+      "Content-Type": "application/json",
+      Authorization: "Bearer test-key",
+    })
+    expect(JSON.parse(String(calls[0]?.body))).toEqual({
+      model: "mercury-edit-2",
+      prompt: "const value = ",
+      suffix: "\n",
+      max_tokens: 256,
+      temperature: 0,
+      stream: true,
+    })
+  })
+
+  it("retries Codestral-specific endpoint when Mistral rejects a Codestral key", async () => {
+    const urls: string[] = []
+    const fetchImpl: typeof fetch = async (url) => {
+      urls.push(String(url))
+      if (urls.length === 1) return new Response("unauthorized", { status: 401, statusText: "Unauthorized" })
+      return new Response(stream("data: [DONE]\n\n"), { status: 200 })
+    }
+
+    await generateDirectFim({
+      apiKey: "codestral-key",
+      target: getDirectFimTarget("mistral/codestral-2508")!,
+      prefix: "",
+      suffix: "",
+      temperature: 0.2,
+      onChunk: () => {},
+      fetchImpl,
+    })
+
+    expect(urls).toEqual(["https://api.mistral.ai/v1/fim/completions", "https://codestral.mistral.ai/v1/fim/completions"])
+  })
+})
@@ -7,7 +7,11 @@ import { parseModelString } from "../../../../src/shared/provider-model"
 import { DEFAULT_AUTOCOMPLETE_MODEL } from "../../../../src/shared/autocomplete-models"
 import { ModelSelectorBase } from "../shared/ModelSelector"
 import SettingsRow from "./SettingsRow"
-import { AUTOCOMPLETE_PROVIDER_ID, AUTOCOMPLETE_SELECTOR_MODELS } from "./autocomplete-model-selector"
+import {
+  AUTOCOMPLETE_SELECTOR_MODELS,
+  getAutocompleteSelection,
+  getAutocompleteSettingID,
+} from "./autocomplete-model-selector"

 const ModelsTab: Component = () => {
  const { config, settings, updateConfig, updateSetting } = useConfig()
@@ -39,8 +43,9 @@ const ModelsTab: Component = () => {
  }

  function handleAutocompleteModelSelect(providerID: string, modelID: string) {
-    if (providerID !== AUTOCOMPLETE_PROVIDER_ID || !modelID) return
-    updateSetting("autocomplete.model", modelID)
+    const id = getAutocompleteSettingID(providerID, modelID)
+    if (!id) return
+    updateSetting("autocomplete.model", id)
  }

  return (
@@ -77,7 +82,7 @@ const ModelsTab: Component = () => {
          last
        >
          <ModelSelectorBase
-            value={{ providerID: AUTOCOMPLETE_PROVIDER_ID, modelID: autocompleteModel() }}
+            value={getAutocompleteSelection(autocompleteModel())}
            onSelect={handleAutocompleteModelSelect}
            placement="bottom-start"
            models={AUTOCOMPLETE_SELECTOR_MODELS}
@@ -1,12 +1,18 @@
-import { AUTOCOMPLETE_MODELS } from "../../../../src/shared/autocomplete-models"
+import { AUTOCOMPLETE_MODELS, getAutocompleteModel } from "../../../../src/shared/autocomplete-models"
 import type { EnrichedModel } from "../../context/provider"

-export const AUTOCOMPLETE_PROVIDER_ID = "kilo"
-export const AUTOCOMPLETE_PROVIDER_NAME = "Kilo Gateway"
+export function getAutocompleteSelection(id: string) {
+  const model = getAutocompleteModel(id)
+  return { providerID: model.providerID, modelID: model.modelID }
+}
+
+export function getAutocompleteSettingID(providerID: string, modelID: string) {
+  return AUTOCOMPLETE_MODELS.find((m) => m.providerID === providerID && m.modelID === modelID)?.id
+}

 export const AUTOCOMPLETE_SELECTOR_MODELS: EnrichedModel[] = AUTOCOMPLETE_MODELS.map((m) => ({
-  id: m.id,
+  id: m.modelID,
  name: m.label,
-  providerID: AUTOCOMPLETE_PROVIDER_ID,
-  providerName: AUTOCOMPLETE_PROVIDER_NAME,
+  providerID: m.providerID,
+  providerName: m.provider,
 }))