From 2bf49bddb6ac20b87a518a9487362f10de5ec06a Mon Sep 17 00:00:00 2001 From: Mark IJbema Date: Mon, 25 May 2026 17:13:01 +0200 Subject: [PATCH] refactor(gateway): extract FIM routing --- packages/kilo-gateway/src/server/fim.ts | 140 +++++++++++++++++++++ packages/kilo-gateway/src/server/routes.ts | 126 +------------------ packages/kilo-gateway/test/fim.test.ts | 2 +- 3 files changed, 143 insertions(+), 125 deletions(-) create mode 100644 packages/kilo-gateway/src/server/fim.ts diff --git a/packages/kilo-gateway/src/server/fim.ts b/packages/kilo-gateway/src/server/fim.ts new file mode 100644 index 00000000000..b668646197d --- /dev/null +++ b/packages/kilo-gateway/src/server/fim.ts @@ -0,0 +1,140 @@ +import { HEADER_FEATURE, KILO_API_BASE } from "../api/constants.js" +import { buildKiloHeaders } from "../headers.js" + +type Auth = any + +type FimProvider = "kilo" | "mistral" | "inception" + +interface FimTarget { + provider: FimProvider + model: string + urls: string[] +} + +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] } +} + +async function getProxyAuth(Auth: Auth) { + const auth = await Auth.get("kilo") + const token = auth?.type === "api" ? auth.key : auth?.type === "oauth" ? auth.access : undefined + return { + auth, + token, + organizationId: auth?.type === "oauth" ? auth.accountId : undefined, + } +} + +async function getProviderKey(Auth: Auth, provider: FimProvider) { + const auth = await Auth.get(provider) + return auth?.type === "api" ? auth.key : undefined +} + +async function fetchFim( + target: FimTarget, + url: string, + fallbacks: string[], + key: string, + input: { + prefix: string + suffix: string + maxTokens: number + temperature: number + signal: AbortSignal + 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 [next] = fallbacks + if (response.status === 401 && next) return fetchFim(target, next, fallbacks.slice(1), key, input) + return response +} + +export function createFimHandler(Auth: Auth) { + return async (c: any) => { + 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(Auth) : undefined + const token = target.provider === "kilo" ? proxy?.token : await getProviderKey(Auth, target.provider) + + if (target.provider === "kilo" && !proxy?.auth) { + return c.json({ error: "Not authenticated with Kilo Gateway" }, 401) + } + + if (target.provider === "kilo" && !token) { + return c.json({ error: "No valid token found" }, 401) + } + + 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 { + 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, + organizationId: proxy?.organizationId, + }) + } catch (err) { + if (err instanceof DOMException && err.name === "TimeoutError") { + return c.json({ error: "FIM request timed out" }, 504 as any) + } + if (signal.aborted) return c.json({ error: "FIM request canceled" }, 499 as any) + throw err + } + + if (!response.ok) { + const text = await response.text() + return c.json({ error: `FIM request failed: ${response.status} ${text}` }, response.status as any) + } + + return new Response(response.body, { + headers: { + "Content-Type": "text/event-stream", + "Cache-Control": "no-cache", + Connection: "keep-alive", + }, + }) + } +} diff --git a/packages/kilo-gateway/src/server/routes.ts b/packages/kilo-gateway/src/server/routes.ts index 381b2b184ef..f1544ade69d 100644 --- a/packages/kilo-gateway/src/server/routes.ts +++ b/packages/kilo-gateway/src/server/routes.ts @@ -18,6 +18,7 @@ import { import { buildKiloHeaders } from "../headers.js" import type { ImportDeps, DrizzleDb } from "../cloud-sessions.js" import { fetchCloudSession, fetchCloudSessionForImport, importSessionToDb } from "../cloud-sessions.js" +import { createFimHandler } from "./fim.js" // Type definitions for OpenCode dependencies (injected at runtime) type Hono = any @@ -29,14 +30,6 @@ 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 @@ -49,22 +42,6 @@ interface KiloRoutesDeps extends ImportDeps { InstanceStore: { disposeAllInstances(): Promise } } -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 * @@ -169,52 +146,6 @@ 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 => { - 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", @@ -406,60 +337,7 @@ export function createKiloRoutes(deps: KiloRoutesDeps) { temperature: z.number().optional(), }), ), - async (c: any) => { - 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 (target.provider === "kilo" && !proxy?.auth) { - return c.json({ error: "Not authenticated with Kilo Gateway" }, 401) - } - - if (target.provider === "kilo" && !token) { - return c.json({ error: "No valid token found" }, 401) - } - - 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 { - 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, - organizationId: proxy?.organizationId, - }) - } catch (err) { - if (err instanceof DOMException && err.name === "TimeoutError") - return c.json({ error: "FIM request timed out" }, 504 as any) - if (signal.aborted) return c.json({ error: "FIM request canceled" }, 499 as any) - throw err - } - - if (!response.ok) { - const text = await response.text() - return c.json({ error: `FIM request failed: ${response.status} ${text}` }, response.status as any) - } - - return new Response(response.body, { - headers: { - "Content-Type": "text/event-stream", - "Cache-Control": "no-cache", - Connection: "keep-alive", - }, - }) - }, + createFimHandler(Auth), ) .post( "/audio/transcriptions", diff --git a/packages/kilo-gateway/test/fim.test.ts b/packages/kilo-gateway/test/fim.test.ts index 9118506ad6e..e90315311a5 100644 --- a/packages/kilo-gateway/test/fim.test.ts +++ b/packages/kilo-gateway/test/fim.test.ts @@ -1,5 +1,5 @@ import { describe, expect, test } from "bun:test" -import { resolveFimTarget } from "../src/server/routes" +import { resolveFimTarget } from "../src/server/fim" describe("FIM target resolution", () => { test("keeps gateway autocomplete models on Kilo Gateway", () => {