fix(gateway): cache working Mistral FIM endpoint

packages/kilo-docs/source-links.md

@@ -11,8 +11,6 @@
   <!-- packages/opencode/src/kilocode/server/httpapi/handlers/kilo-gateway.ts -->
 - <https://api.kilo.ai>
   <!-- packages/opencode/src/cli/cmd/github.ts -->
-- <https://api.mistral.ai/v1/fim/completions>
-  <!-- packages/opencode/src/kilocode/server/httpapi/handlers/kilo-gateway.ts -->
 - <https://api.x.ai/v1>
   <!-- packages/opencode/src/plugin/xai.ts -->
 - <https://app.kilo.ai>
@@ -47,8 +45,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/opencode/src/kilocode/server/httpapi/handlers/kilo-gateway.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/package.json

@@ -18,6 +18,7 @@
   "exports": {
     ".": "./src/index.ts",
     "./autocomplete": "./src/autocomplete.ts",
+    "./fim-endpoint": "./src/fim-endpoint.ts",
     "./tui": "./src/tui.ts"
   },
   "files": [

packages/kilo-gateway/src/fim-endpoint.ts

@@ -0,0 +1,37 @@
+import { createHash } from "node:crypto"
+
+export const MISTRAL_FIM_URL = "https://api.mistral.ai/v1/fim/completions"
+export const CODESTRAL_FIM_URL = "https://codestral.mistral.ai/v1/fim/completions"
+
+const cache = new Map<string, string>()
+
+function fingerprint(key: string) {
+  return createHash("sha256").update(key).digest("hex").slice(0, 16)
+}
+
+export function isMistralEndpointMismatch(response: Response) {
+  return response.status === 401 || response.status === 403
+}
+
+export function clearMistralFimEndpointCache() {
+  cache.clear()
+}
+
+export function getCachedMistralFimEndpoint(key: string) {
+  return cache.get(fingerprint(key))
+}
+
+export async function requestMistralFim(key: string, request: (url: string) => Promise<Response>) {
+  const id = fingerprint(key)
+  const preferred = cache.get(id) ?? MISTRAL_FIM_URL
+  const alternate = preferred === MISTRAL_FIM_URL ? CODESTRAL_FIM_URL : MISTRAL_FIM_URL
+  const first = await request(preferred)
+
+  if (first.ok) return first
+  if (!isMistralEndpointMismatch(first)) return first
+
+  cache.delete(id)
+  const second = await request(alternate)
+  if (second.ok) cache.set(id, alternate)
+  return second
+}

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

@@ -1,5 +1,6 @@
 import { HEADER_FEATURE, KILO_API_BASE } from "../api/constants.js"
 import { getAutocompleteModel } from "../autocomplete.js"
+import { CODESTRAL_FIM_URL, MISTRAL_FIM_URL, requestMistralFim } from "../fim-endpoint.js"
 import { buildKiloHeaders } from "../headers.js"
 
 type Auth = any
@@ -19,8 +20,6 @@ interface FimTarget {
 
 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(provider?: string, model?: string): FimTarget {
@@ -57,8 +56,6 @@ async function getProviderKey(Auth: Auth, provider: FimProvider) {
 
 async function fetchFim(
   target: FimTarget,
-  url: string,
-  fallbacks: string[],
   key: string,
   input: {
     prefix: string
@@ -69,31 +66,35 @@ async function fetchFim(
     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 run = async (url: string) => {
+    console.info(`[FIM] request provider=${target.provider} model=${target.model} url=${url}`)
+    return 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
+  if (target.provider === "mistral") return requestMistralFim(key, run)
+
+  const [url] = target.urls
+  if (!url) throw new Error("No FIM endpoint configured")
+  return run(url)
 }
 
 export function createFimHandler(Auth: Auth) {
@@ -121,9 +122,7 @@ export function createFimHandler(Auth: Auth) {
 
     let response: Response
     try {
-      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, {
+      response = await fetchFim(target, token, {
         prefix,
         suffix,
         maxTokens: fimMaxTokens,

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

@@ -0,0 +1,76 @@
+import { describe, expect, test } from "bun:test"
+import {
+  CODESTRAL_FIM_URL,
+  MISTRAL_FIM_URL,
+  clearMistralFimEndpointCache,
+  getCachedMistralFimEndpoint,
+  requestMistralFim,
+} from "../src/fim-endpoint"
+
+function response(status: number) {
+  return new Response(null, { status })
+}
+
+describe("Mistral FIM endpoint cache", () => {
+  test("caches Codestral endpoint after successful fallback", async () => {
+    clearMistralFimEndpointCache()
+    const urls: string[] = []
+    const first = await requestMistralFim("key-a", async (url) => {
+      urls.push(url)
+      return response(url === MISTRAL_FIM_URL ? 401 : 200)
+    })
+    const second = await requestMistralFim("key-a", async (url) => {
+      urls.push(url)
+      return response(200)
+    })
+
+    expect(first.ok).toBe(true)
+    expect(second.ok).toBe(true)
+    expect(urls).toEqual([MISTRAL_FIM_URL, CODESTRAL_FIM_URL, CODESTRAL_FIM_URL])
+    expect(getCachedMistralFimEndpoint("key-a")).toBe(CODESTRAL_FIM_URL)
+  })
+
+  test("does not cache fallback for invalid credentials", async () => {
+    clearMistralFimEndpointCache()
+    const urls: string[] = []
+    const res = await requestMistralFim("key-b", async (url) => {
+      urls.push(url)
+      return response(401)
+    })
+
+    expect(res.status).toBe(401)
+    expect(urls).toEqual([MISTRAL_FIM_URL, CODESTRAL_FIM_URL])
+    expect(getCachedMistralFimEndpoint("key-b")).toBeUndefined()
+  })
+
+  test("keeps endpoint preference scoped to credential fingerprint", async () => {
+    clearMistralFimEndpointCache()
+    const urls: string[] = []
+    await requestMistralFim("key-c", async (url) => {
+      urls.push(url)
+      return response(url === MISTRAL_FIM_URL ? 403 : 200)
+    })
+    await requestMistralFim("key-d", async (url) => {
+      urls.push(url)
+      return response(200)
+    })
+
+    expect(urls).toEqual([MISTRAL_FIM_URL, CODESTRAL_FIM_URL, MISTRAL_FIM_URL])
+    expect(getCachedMistralFimEndpoint("key-c")).toBe(CODESTRAL_FIM_URL)
+    expect(getCachedMistralFimEndpoint("key-d")).toBeUndefined()
+  })
+
+  test("clears stale preference and probes alternate endpoint", async () => {
+    clearMistralFimEndpointCache()
+    const urls: string[] = []
+    await requestMistralFim("key-e", async (url) => response(url === MISTRAL_FIM_URL ? 401 : 200))
+    const res = await requestMistralFim("key-e", async (url) => {
+      urls.push(url)
+      return response(url === CODESTRAL_FIM_URL ? 401 : 200)
+    })
+
+    expect(res.ok).toBe(true)
+    expect(urls).toEqual([CODESTRAL_FIM_URL, MISTRAL_FIM_URL])
+    expect(getCachedMistralFimEndpoint("key-e")).toBe(MISTRAL_FIM_URL)
+  })
+})

packages/opencode/src/kilocode/server/httpapi/handlers/kilo-gateway.ts

@@ -20,6 +20,7 @@ import {
   fetchProfile,
 } from "@kilocode/kilo-gateway"
 import { getAutocompleteModel } from "@kilocode/kilo-gateway/autocomplete"
+import { CODESTRAL_FIM_URL, MISTRAL_FIM_URL, requestMistralFim } from "@kilocode/kilo-gateway/fim-endpoint"
 import { buildKiloHeaders } from "@kilocode/kilo-gateway"
 import { Effect } from "effect"
 import * as Stream from "effect/Stream"
@@ -40,8 +41,6 @@ import { AudioTranscriptionsBody, FimBody } from "../groups/kilo-gateway"
 
 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"
 
 type FimProvider = "kilo" | "mistral" | "inception"
@@ -123,9 +122,9 @@ export const kiloGatewayHandlers = HttpApiBuilder.group(InstanceHttpApi, "kilo",
           : AbortSignal.timeout(FIM_TIMEOUT_MS)
       const response = yield* Effect.promise(async () => {
         try {
-          const run = async (url: string, fallbacks: string[]): Promise<Response> => {
+          const run = async (url: string): Promise<Response> => {
             console.info(`[FIM] request provider=${target.provider} model=${target.model} url=${url}`)
-            const response = await fetch(url, {
+            return fetch(url, {
               method: "POST",
               headers: {
                 "Content-Type": "application/json",
@@ -145,11 +144,9 @@ export const kiloGatewayHandlers = HttpApiBuilder.group(InstanceHttpApi, "kilo",
                 stream: true,
               }),
             })
-            const [next] = fallbacks
-            if (response.status === 401 && next) return run(next, fallbacks.slice(1))
-            return response
           }
-          return run(target.urls[0]!, target.urls.slice(1))
+          if (target.provider === "mistral") return requestMistralFim(token, run)
+          return run(target.urls[0]!)
         } catch (err) {
           if (err instanceof DOMException && err.name === "TimeoutError")
             return Response.json({ error: "FIM request timed out" }, { status: 504 })