From d02145a549ab263f28b4c00f3fddde239dbd6475 Mon Sep 17 00:00:00 2001 From: Mark IJbema Date: Tue, 26 May 2026 11:44:13 +0200 Subject: [PATCH] refactor(gateway): simplify Mistral FIM endpoint cache --- .../kilo-gateway/src/mistral-fim-endpoint.ts | 29 +++++++---------- packages/kilo-gateway/src/server/fim.ts | 2 +- .../test/mistral-fim-endpoint.test.ts | 31 +++++++++---------- .../server/httpapi/handlers/kilo-gateway.ts | 2 +- 4 files changed, 28 insertions(+), 36 deletions(-) diff --git a/packages/kilo-gateway/src/mistral-fim-endpoint.ts b/packages/kilo-gateway/src/mistral-fim-endpoint.ts index deb2630faf..ebd5b70fb4 100644 --- a/packages/kilo-gateway/src/mistral-fim-endpoint.ts +++ b/packages/kilo-gateway/src/mistral-fim-endpoint.ts @@ -1,37 +1,30 @@ -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() - -function fingerprint(key: string) { - return createHash("sha256").update(key).digest("hex").slice(0, 16) -} +let preferred: string | undefined export function isMistralEndpointMismatch(response: Response) { return response.status === 401 || response.status === 403 } export function clearMistralFimEndpointCache() { - cache.clear() + preferred = undefined } -export function getCachedMistralFimEndpoint(key: string) { - return cache.get(fingerprint(key)) +export function getCachedMistralFimEndpoint() { + return preferred } -export async function requestMistralFim(key: string, request: (url: string) => Promise) { - 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) +export async function requestMistralFim(request: (url: string) => Promise) { + const firstUrl = preferred ?? MISTRAL_FIM_URL + const secondUrl = firstUrl === MISTRAL_FIM_URL ? CODESTRAL_FIM_URL : MISTRAL_FIM_URL + const first = await request(firstUrl) 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) + preferred = undefined + const second = await request(secondUrl) + if (second.ok) preferred = secondUrl return second } diff --git a/packages/kilo-gateway/src/server/fim.ts b/packages/kilo-gateway/src/server/fim.ts index 5c0f6428fa..b9628378b4 100644 --- a/packages/kilo-gateway/src/server/fim.ts +++ b/packages/kilo-gateway/src/server/fim.ts @@ -90,7 +90,7 @@ async function fetchFim( }) } - if (target.provider === "mistral") return requestMistralFim(key, run) + if (target.provider === "mistral") return requestMistralFim(run) const [url] = target.urls if (!url) throw new Error("No FIM endpoint configured") diff --git a/packages/kilo-gateway/test/mistral-fim-endpoint.test.ts b/packages/kilo-gateway/test/mistral-fim-endpoint.test.ts index 48c4a58d37..d10c617728 100644 --- a/packages/kilo-gateway/test/mistral-fim-endpoint.test.ts +++ b/packages/kilo-gateway/test/mistral-fim-endpoint.test.ts @@ -12,14 +12,14 @@ function response(status: number) { } describe("Mistral FIM endpoint cache", () => { - test("caches Codestral endpoint after successful fallback", async () => { + test("remembers Codestral endpoint after successful fallback", async () => { clearMistralFimEndpointCache() const urls: string[] = [] - const first = await requestMistralFim("key-a", async (url) => { + const first = await requestMistralFim(async (url) => { urls.push(url) return response(url === MISTRAL_FIM_URL ? 401 : 200) }) - const second = await requestMistralFim("key-a", async (url) => { + const second = await requestMistralFim(async (url) => { urls.push(url) return response(200) }) @@ -27,50 +27,49 @@ describe("Mistral FIM endpoint cache", () => { 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) + expect(getCachedMistralFimEndpoint()).toBe(CODESTRAL_FIM_URL) }) - test("does not cache fallback for invalid credentials", async () => { + test("does not remember fallback for invalid credentials", async () => { clearMistralFimEndpointCache() const urls: string[] = [] - const res = await requestMistralFim("key-b", async (url) => { + const res = await requestMistralFim(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() + expect(getCachedMistralFimEndpoint()).toBeUndefined() }) - test("keeps endpoint preference scoped to credential fingerprint", async () => { + test("uses one process-local endpoint preference", async () => { clearMistralFimEndpointCache() const urls: string[] = [] - await requestMistralFim("key-c", async (url) => { + await requestMistralFim(async (url) => { urls.push(url) return response(url === MISTRAL_FIM_URL ? 403 : 200) }) - await requestMistralFim("key-d", async (url) => { + await requestMistralFim(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() + expect(urls).toEqual([MISTRAL_FIM_URL, CODESTRAL_FIM_URL, CODESTRAL_FIM_URL]) + expect(getCachedMistralFimEndpoint()).toBe(CODESTRAL_FIM_URL) }) 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) => { + await requestMistralFim(async (url) => response(url === MISTRAL_FIM_URL ? 401 : 200)) + const res = await requestMistralFim(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) + expect(getCachedMistralFimEndpoint()).toBe(MISTRAL_FIM_URL) }) }) diff --git a/packages/opencode/src/kilocode/server/httpapi/handlers/kilo-gateway.ts b/packages/opencode/src/kilocode/server/httpapi/handlers/kilo-gateway.ts index 7502716474..a024add296 100644 --- a/packages/opencode/src/kilocode/server/httpapi/handlers/kilo-gateway.ts +++ b/packages/opencode/src/kilocode/server/httpapi/handlers/kilo-gateway.ts @@ -145,7 +145,7 @@ export const kiloGatewayHandlers = HttpApiBuilder.group(InstanceHttpApi, "kilo", }), }) } - if (target.provider === "mistral") return requestMistralFim(token, run) + if (target.provider === "mistral") return requestMistralFim(run) return run(target.urls[0]!) } catch (err) { if (err instanceof DOMException && err.name === "TimeoutError")