feat(vscode,gateway): route Mercury Next Edit through Kilo Gateway

.changeset/nextedit-via-kilo-gateway.md

@@ -0,0 +1,6 @@
+---
+"@kilocode/kilo-gateway": minor
+"kilo-code": minor
+---
+
+Support Mercury Next Edit through the Kilo Gateway. The new "Mercury Next Edit via Kilo Gateway" autocomplete model routes Next Edit predictions through your Kilo account (no separate Inception API key required).

packages/kilo-gateway/src/autocomplete.ts

@@ -45,6 +45,19 @@ const models: AutocompleteModelDef[] = [
     requestModel: "inception/mercury-edit-2",
     temperature: 0,
   },
+  {
+    // Same wire-level model as `kilo/inception/mercury-edit-2`, but routed
+    // through the Kilo Gateway's Next Edit endpoint instead of FIM. Picked by
+    // users who want multi-line next-edit predictions with the jump-to-edit UX.
+    id: "kilo/inception/mercury-next-edit",
+    modelID: "inception/mercury-next-edit",
+    label: "Mercury Next Edit",
+    providerID: "kilo",
+    provider: "Kilo Gateway",
+    requestModel: "inception/mercury-edit-2",
+    temperature: 0,
+    kind: "edit",
+  },
   {
     id: "mistral/codestral-2508",
     modelID: "codestral-2508",

packages/kilo-gateway/src/edit.ts

@@ -1,3 +1,4 @@
+import { KILO_API_BASE } from "./api/constants.js"
 import { getAutocompleteModel, type DirectAutocompleteProviderID } from "./autocomplete.js"
 
 /**
@@ -20,20 +21,30 @@ export interface EditUpstreamResponse {
 }
 
 const INCEPTION_EDIT_URL = "https://api.inceptionlabs.ai/v1/edit/completions"
+const KILO_NEXTEDIT_URL = KILO_API_BASE + "/api/edit/completions"
 
 /**
- * Pick the upstream edit endpoint for a (provider, model) pair. Only Inception
- * is wired up today — Mercury is the only model family with a documented
- * /v1/edit/completions endpoint. Mistral does not expose a comparable surface.
+ * Pick the upstream edit endpoint for a (provider, model) pair. Today this is
+ * either Inception's `/v1/edit/completions` (direct BYOK) or the Kilo Gateway's
+ * `/api/edit/completions` proxy, which forwards to Inception server-side.
+ * Mistral does not expose a comparable surface.
  */
 export function resolveEditTarget(provider?: string, model?: string): EditTarget {
   const info = getAutocompleteModel(provider, model)
-  if (info.kind === "edit" && info.directProvider === "inception") {
-    return { provider: "inception", model: info.requestModel, url: INCEPTION_EDIT_URL }
+  if (info.kind === "edit") {
+    if (info.providerID === "kilo") {
+      // The gateway expects the upstream model id with the `inception/` prefix
+      // (it strips it before forwarding to Inception). The kilo entry's
+      // `requestModel` already carries the prefix.
+      const m = info.requestModel.includes("/") ? info.requestModel : `inception/${info.requestModel}`
+      return { provider: "kilo", model: m, url: KILO_NEXTEDIT_URL }
+    }
+    if (info.directProvider === "inception") {
+      return { provider: "inception", model: info.requestModel, url: INCEPTION_EDIT_URL }
+    }
   }
-  // Kilo Gateway does not currently proxy an edit endpoint; callers should
-  // fall back to FIM. We still return a kilo target so the handler can surface
-  // a 400 rather than silently routing somewhere unexpected.
+  // Non-edit models fall through to a kilo placeholder with no URL so the
+  // handler can surface a 400 rather than silently routing somewhere unexpected.
   return { provider: "kilo", model: info.requestModel, url: "" }
 }
 

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

@@ -1,6 +1,8 @@
+import { HEADER_FEATURE } from "../api/constants.js"
 import { DIRECT_EDIT_ENV, extractFencedBody, resolveEditTarget, type EditTarget, type EditUpstreamResponse } from "../edit.js"
 import { buildMercuryEditPrompt, type MercuryEditContext } from "../edit-prompt.js"
 import type { DirectAutocompleteProviderID } from "../autocomplete.js"
+import { buildKiloHeaders } from "../headers.js"
 import type { AuthStore } from "./handlers.js"
 
 type Auth = Pick<AuthStore, "get">
@@ -14,16 +16,33 @@ async function getProviderKey(Auth: Auth, provider: DirectAutocompleteProviderID
   return DIRECT_EDIT_ENV[provider].map((key) => process.env[key]).find(Boolean)
 }
 
+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,
+  }
+}
+
 export function createEditHandler(Auth: Auth) {
   return async (c: any) => {
     const { provider, model, maxTokens, ...context } = c.req.valid("json")
     const target = resolveEditTarget(provider, model)
 
-    if (target.provider !== "inception") {
+    if (target.provider === "kilo" && !target.url) {
       return c.json({ error: "Next Edit currently requires the Inception provider (mercury-edit-2)." }, 400 as any)
     }
 
-    const token = await getProviderKey(Auth, target.provider)
+    const proxy = target.provider === "kilo" ? await getProxyAuth(Auth) : undefined
+    const token =
+      target.provider === "kilo" ? proxy?.token : await getProviderKey(Auth, target.provider as DirectAutocompleteProviderID)
+
+    if (target.provider === "kilo" && !proxy?.auth) {
+      return c.json({ error: "Not authenticated with Kilo Gateway" }, 401 as any)
+    }
+
     if (!token) {
       return c.json({ error: `Missing ${target.provider} provider API key` }, 401 as any)
     }
@@ -40,6 +59,10 @@ export function createEditHandler(Auth: Auth) {
         headers: {
           "Content-Type": "application/json",
           Authorization: `Bearer ${token}`,
+          ...(target.provider === "kilo"
+            ? buildKiloHeaders(undefined, { kilocodeOrganizationId: proxy?.organizationId })
+            : {}),
+          ...(target.provider === "kilo" ? { [HEADER_FEATURE]: "autocomplete" } : {}),
         },
         signal,
         body: JSON.stringify({

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

@@ -16,6 +16,13 @@ describe("Edit target resolution", () => {
     expect(resolveEditTarget("inception", "mercury-edit-2").provider).toBe("kilo")
   })
 
+  test("routes the Kilo Gateway next-edit model to the gateway proxy", () => {
+    const target = resolveEditTarget("kilo", "inception/mercury-next-edit")
+    expect(target.provider).toBe("kilo")
+    expect(target.model).toBe("inception/mercury-edit-2")
+    expect(target.url).toMatch(/\/api\/edit\/completions$/)
+  })
+
   test("falls back to a kilo placeholder (no upstream) for non-edit models", () => {
     expect(resolveEditTarget("kilo", "mistralai/codestral-2508")).toEqual({
       provider: "kilo",

packages/kilo-vscode/eslint.config.mjs

@@ -38,7 +38,7 @@ export default [
   // New code must stay ≤ 20. Do not raise these caps; refactor instead.
   {
     files: ["src/KiloProvider.ts"],
-    rules: { complexity: ["error", 150], "max-lines": ["error", 3600] },
+    rules: { complexity: ["error", 150], "max-lines": ["error", 3700] },
   },
   {
     files: ["webview-ui/agent-manager/AgentManagerApp.tsx"],

packages/kilo-vscode/package.json

@@ -845,13 +845,15 @@
           "enum": [
             "mistralai/codestral-2508",
             "inception/mercury-edit-2",
+            "inception/mercury-next-edit",
             "codestral-2508",
             "mercury-edit-2",
             "mercury-next-edit"
           ],
           "enumDescriptions": [
             "Codestral via Kilo Gateway",
             "Mercury Edit 2 via Kilo Gateway",
+            "Mercury Next Edit (multi-line edit predictions with jump-to-edit UX) via Kilo Gateway",
             "Codestral via your connected Mistral provider API key",
             "Mercury Edit 2 (FIM) via your connected Inception provider API key",
             "Mercury Next Edit (multi-line edit predictions with jump-to-edit UX) via your connected Inception provider API key"

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

@@ -109,6 +109,10 @@ export class AutocompleteServiceManager {
     this.nextEditProvider = new NextEditInlineCompletionProvider({
       connectionService,
       suggestionManager: this.nextEditSuggestionManager,
+      getModelSelection: () => {
+        const info = getAutocompleteModel(this.settings?.provider, this.settings?.model)
+        return { providerId: info.providerID, modelId: info.modelID }
+      },
       isFileAllowed: async (fsPath) => {
         const ignore = await this.inlineCompletionProvider.ignoreController
         return ignore.validateAccess(fsPath)

packages/kilo-vscode/src/services/autocomplete/next-edit/MercuryEditProvider.ts

@@ -3,13 +3,17 @@ import { nesLog, nesWarn } from "./log"
 import type { MercuryEditRequestContext, MercuryEditSuggestion } from "./types"
 
 const MERCURY_MAX_TOKENS = 512
-const PROVIDER_ID = "inception"
-const MODEL_ID = "mercury-next-edit"
+const DEFAULT_PROVIDER_ID = "inception"
+const DEFAULT_MODEL_ID = "mercury-next-edit"
 
 type EditResponseData = { content?: string; usage?: { prompt_tokens?: number; completion_tokens?: number } }
 
 export interface MercuryEditProviderOptions {
   connectionService: KiloConnectionService
+  /** Provider id to send to the gateway (e.g. `"kilo"` or `"inception"`). */
+  providerId?: string
+  /** Model id to send to the gateway (e.g. `"inception/mercury-next-edit"`). */
+  modelId?: string
   /** AbortSignal for cancellation (cursor moves, escape, etc.). */
   signal?: AbortSignal
 }
@@ -25,17 +29,19 @@ export class MercuryEditProvider {
 
   async suggest(ctx: MercuryEditRequestContext): Promise<MercuryEditSuggestion | null> {
     const start = Date.now()
+    const provider = this.options.providerId ?? DEFAULT_PROVIDER_ID
+    const model = this.options.modelId ?? DEFAULT_MODEL_ID
     nesLog(
-      `-> /kilo/edit model=${MODEL_ID} region=[${ctx.editableRegionStartLine},${ctx.editableRegionEndLine}] diffs=${ctx.editDiffHistory.length} snippets=${ctx.recentlyViewedSnippets.length}`,
+      `-> /kilo/edit provider=${provider} model=${model} region=[${ctx.editableRegionStartLine},${ctx.editableRegionEndLine}] diffs=${ctx.editDiffHistory.length} snippets=${ctx.recentlyViewedSnippets.length}`,
     )
 
     const client = await this.options.connectionService.getClientAsync()
     try {
       // Send structured editor context; the gateway assembles the Mercury prompt.
       const { data, error, response } = await client.kilo.edit(
         {
-          provider: PROVIDER_ID,
-          model: MODEL_ID,
+          provider,
+          model,
           maxTokens: MERCURY_MAX_TOKENS,
           currentFilePath: ctx.currentFilePath,
           currentFileContent: ctx.currentFileContent,

packages/kilo-vscode/src/services/autocomplete/next-edit/NextEditInlineCompletionProvider.ts

@@ -22,6 +22,8 @@ export interface NextEditProviderDeps {
   onFatalError?: (status: number | null) => void
   /** Stash for diffs that don't land on the cursor's line — rendered as a jump affordance. */
   suggestionManager?: NextEditSuggestionManager
+  /** Resolves the currently selected (provider, model) at request time. */
+  getModelSelection?: () => { providerId: string; modelId: string }
 }
 
 export interface NextEditSuggestionEvent {
@@ -78,8 +80,11 @@ export class NextEditInlineCompletionProvider implements vscode.InlineCompletion
 
     const abort = this.swapAbortController(token)
     const ctx = await this.buildRequestContext(document, position)
+    const sel = this.deps.getModelSelection?.()
     const provider = new MercuryEditProvider({
       connectionService: this.deps.connectionService,
+      providerId: sel?.providerId,
+      modelId: sel?.modelId,
       signal: abort.signal,
     })