refactor(vscode): route direct FIM through backend

packages/kilo-docs/source-links.md

@@ -1,17 +1,13 @@
 # Source Code Links
 
 <!-- Auto-generated by script/extract-source-links.ts — DO NOT EDIT -->
-<!-- 89 unique URLs extracted from extension and CLI source -->
+<!-- 86 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>
@@ -36,8 +32,6 @@
   <!-- 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>

packages/kilo-gateway/src/server/routes.ts

@@ -29,6 +29,14 @@ type Auth = any
 type ModelCache = { clear: (providerID: string) => void }
 type Z = any
 
+type FimProvider = "kilo" | "mistral" | "inception"
+
+interface FimTarget {
+  provider: FimProvider
+  model: string
+  urls: string[]
+}
+
 interface KiloRoutesDeps extends ImportDeps {
   Hono: new () => Hono
   describeRoute: DescribeRoute
@@ -42,6 +50,20 @@ interface KiloRoutesDeps extends ImportDeps {
 }
 
 const FIM_TIMEOUT_MS = 30_000
+const KILO_FIM_URL = KILO_API_BASE + "/api/fim/completions"
+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"
+
+export function resolveFimTarget(model?: string): FimTarget {
+  if (model === "mistral/codestral-2508") {
+    return { provider: "mistral", model: "codestral-2508", urls: [MISTRAL_FIM_URL, CODESTRAL_FIM_URL] }
+  }
+  if (model === "inception-direct/mercury-edit-2") {
+    return { provider: "inception", model: "mercury-edit-2", urls: [INCEPTION_FIM_URL] }
+  }
+  return { provider: "kilo", model: model ?? "mistralai/codestral-2501", urls: [KILO_FIM_URL] }
+}
 
 /**
  * Create Kilo Gateway routes with OpenCode dependencies injected
@@ -147,6 +169,52 @@ export function createKiloRoutes(deps: KiloRoutesDeps) {
     }
   }
 
+  const getProviderKey = async (provider: FimProvider) => {
+    const auth = await Auth.get(provider)
+    return auth?.type === "api" ? auth.key : undefined
+  }
+
+  const fetchFim = async (
+    target: FimTarget,
+    url: string,
+    fallbacks: string[],
+    key: string,
+    input: {
+      prefix: string
+      suffix: string
+      maxTokens: number
+      temperature: number
+      signal: AbortSignal
+      organizationId?: string
+    },
+  ): Promise<Response> => {
+    console.info(`[FIM] request provider=${target.provider} model=${target.model} url=${url}`)
+    const response = await fetch(url, {
+      method: "POST",
+      headers: {
+        "Content-Type": "application/json",
+        Authorization: `Bearer ${key}`,
+        ...(target.provider === "kilo"
+          ? buildKiloHeaders(undefined, { kilocodeOrganizationId: input.organizationId })
+          : {}),
+        ...(target.provider === "kilo" ? { [HEADER_FEATURE]: "autocomplete" } : {}),
+      },
+      signal: input.signal,
+      body: JSON.stringify({
+        model: target.model,
+        prompt: input.prefix,
+        suffix: input.suffix,
+        max_tokens: input.maxTokens,
+        temperature: input.temperature,
+        stream: true,
+      }),
+    })
+
+    const [next] = fallbacks
+    if (response.status === 401 && next) return fetchFim(target, next, fallbacks.slice(1), key, input)
+    return response
+  }
+
   return new Hono()
     .get(
       "/profile",
@@ -339,47 +407,38 @@ export function createKiloRoutes(deps: KiloRoutesDeps) {
         }),
       ),
       async (c: any) => {
-        const proxy = await getProxyAuth()
+        const { prefix, suffix, model, maxTokens, temperature } = c.req.valid("json")
+        const target = resolveFimTarget(model)
+        const fimMaxTokens = maxTokens ?? 256
+        const fimTemperature = temperature ?? 0.2
+        const proxy = target.provider === "kilo" ? await getProxyAuth() : undefined
+        const token = target.provider === "kilo" ? proxy?.token : await getProviderKey(target.provider)
 
-        if (!proxy.auth) {
+        if (target.provider === "kilo" && !proxy?.auth) {
           return c.json({ error: "Not authenticated with Kilo Gateway" }, 401)
         }
 
-        if (!proxy.token) {
+        if (target.provider === "kilo" && !token) {
           return c.json({ error: "No valid token found" }, 401)
         }
 
-        const { prefix, suffix, model, maxTokens, temperature } = c.req.valid("json")
-        const fimModel = model ?? "mistralai/codestral-2501"
-        const fimMaxTokens = maxTokens ?? 256
-        const fimTemperature = temperature ?? 0.2
-
-        const baseApiUrl = KILO_API_BASE + "/api/"
-        const endpoint = new URL("fim/completions", baseApiUrl)
-
-        const headers = {
-          "Content-Type": "application/json",
-          Authorization: `Bearer ${proxy.token}`,
-          ...buildKiloHeaders(undefined, { kilocodeOrganizationId: proxy.organizationId }),
-          [HEADER_FEATURE]: "autocomplete",
+        if (!token) {
+          return c.json({ error: `Missing ${target.provider} provider API key` }, 401)
         }
 
         const signal = AbortSignal.any([c.req.raw.signal, AbortSignal.timeout(FIM_TIMEOUT_MS)])
 
         let response: Response
         try {
-          response = await fetch(endpoint, {
-            method: "POST",
-            headers,
+          const [url] = target.urls
+          if (!url) return c.json({ error: "No FIM endpoint configured" }, 500 as any)
+          response = await fetchFim(target, url, target.urls.slice(1), token, {
+            prefix,
+            suffix,
+            maxTokens: fimMaxTokens,
+            temperature: fimTemperature,
             signal,
-            body: JSON.stringify({
-              model: fimModel,
-              prompt: prefix,
-              suffix,
-              max_tokens: fimMaxTokens,
-              temperature: fimTemperature,
-              stream: true,
-            }),
+            organizationId: proxy?.organizationId,
           })
         } catch (err) {
           if (err instanceof DOMException && err.name === "TimeoutError")

packages/kilo-gateway/test/fim.test.ts

@@ -0,0 +1,30 @@
+import { describe, expect, test } from "bun:test"
+import { resolveFimTarget } from "../src/server/routes"
+
+describe("FIM target resolution", () => {
+  test("keeps gateway autocomplete models on Kilo Gateway", () => {
+    expect(resolveFimTarget("mistralai/codestral-2508")).toEqual({
+      provider: "kilo",
+      model: "mistralai/codestral-2508",
+      urls: ["https://api.kilo.ai/api/fim/completions"],
+    })
+    expect(resolveFimTarget("inception/mercury-edit-2")).toEqual({
+      provider: "kilo",
+      model: "inception/mercury-edit-2",
+      urls: ["https://api.kilo.ai/api/fim/completions"],
+    })
+  })
+
+  test("routes explicit provider autocomplete models directly", () => {
+    expect(resolveFimTarget("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(resolveFimTarget("inception-direct/mercury-edit-2")).toEqual({
+      provider: "inception",
+      model: "mercury-edit-2",
+      urls: ["https://api.inceptionlabs.ai/v1/fim/completions"],
+    })
+  })
+})

packages/kilo-vscode/src/services/autocomplete/fim.ts

@@ -3,171 +3,6 @@ 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.
@@ -185,25 +20,6 @@ export async function generateFim(
 ): 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
   let outputTokens = 0
@@ -213,16 +29,15 @@ export async function generateFim(
   // ends the stream. Without this, errors never reach ErrorBackoff.
   let sseError: Error | undefined
 
-  const temp = info.temperature
-  console.info(`[FIM] request provider=kilo model=${info.requestModel} url=/kilo/fim`)
+  console.info(`[FIM] request provider=${info.providerID} model=${info.requestModel} url=/kilo/fim`)
 
   const { stream } = await client.kilo.fim(
     {
       prefix,
       suffix,
-      model: info.requestModel,
+      model: info.id,
       maxTokens: FIM_MAX_TOKENS,
-      temperature: temp,
+      temperature: info.temperature,
     },
     {
       signal,

packages/kilo-vscode/tests/unit/autocomplete-fim-direct.test.ts

@@ -1,94 +0,0 @@
-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"])
-  })
-})