diff --git a/packages/kilo-docs/source-links.md b/packages/kilo-docs/source-links.md index 21e2a798364..d4161189873 100644 --- a/packages/kilo-docs/source-links.md +++ b/packages/kilo-docs/source-links.md @@ -11,8 +11,6 @@ - -- - - - @@ -47,8 +45,6 @@ - -- - - - diff --git a/packages/kilo-gateway/package.json b/packages/kilo-gateway/package.json index 96128d08f01..8bd9023e4e6 100644 --- a/packages/kilo-gateway/package.json +++ b/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": [ diff --git a/packages/kilo-gateway/src/fim-endpoint.ts b/packages/kilo-gateway/src/fim-endpoint.ts new file mode 100644 index 00000000000..deb2630fafb --- /dev/null +++ b/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() + +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) { + 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 +} diff --git a/packages/kilo-gateway/src/server/fim.ts b/packages/kilo-gateway/src/server/fim.ts index d344337941d..1777c7b835b 100644 --- a/packages/kilo-gateway/src/server/fim.ts +++ b/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 { - 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, diff --git a/packages/kilo-gateway/test/fim-endpoint.test.ts b/packages/kilo-gateway/test/fim-endpoint.test.ts new file mode 100644 index 00000000000..f6567510cfe --- /dev/null +++ b/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) + }) +}) 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 02535ee8a02..9646b1b06f9 100644 --- a/packages/opencode/src/kilocode/server/httpapi/handlers/kilo-gateway.ts +++ b/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 => { + const run = async (url: string): Promise => { 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 })