Add Mercury Next Edit (NES) as an opt-in autocomplete mode

packages/kilo-gateway/package.json

@@ -19,6 +19,8 @@
     ".": "./src/index.ts",
     "./autocomplete": "./src/autocomplete.ts",
     "./fim": "./src/fim.ts",
+    "./edit": "./src/edit.ts",
+    "./edit-prompt": "./src/edit-prompt.ts",
     "./tui": "./src/tui.ts"
   },
   "files": [

packages/kilo-gateway/src/autocomplete.ts

@@ -4,20 +4,26 @@ export type DirectAutocompleteProviderID = Exclude<AutocompleteProviderID, "kilo
 export interface AutocompleteModelDef {
   /** Stable combined value for internal comparisons. */
   readonly id: string
-  /** Model value stored in settings and sent to the FIM API. */
+  /** Model value stored in settings and sent to the autocomplete API. */
   readonly modelID: string
   /** Human-readable label shown in settings. */
   readonly label: string
   /** Provider value stored in settings and used by the selector group. */
   readonly providerID: AutocompleteProviderID
   /** Provider display name for status bar / telemetry. */
   readonly provider: string
-  /** Full model ID sent upstream by the FIM route. */
+  /** Full model ID sent upstream by the autocomplete route. */
   readonly requestModel: string
-  /** Provider key to use for direct BYOK FIM. Empty means Kilo Gateway. */
+  /** Provider key to use for direct BYOK. Empty means Kilo Gateway. */
   readonly directProvider?: DirectAutocompleteProviderID
-  /** FIM request temperature. */
+  /** Request temperature. */
   readonly temperature: number
+  /**
+   * Which gateway endpoint this model targets. Defaults to "fim" if omitted
+   * (back-compat with existing entries). Models with `kind: "edit"` route
+   * through `/kilo/edit` and use Mercury's Next Edit pipeline.
+   */
+  readonly kind?: "fim" | "edit"
 }
 
 const models: AutocompleteModelDef[] = [
@@ -59,6 +65,20 @@ const models: AutocompleteModelDef[] = [
     directProvider: "inception",
     temperature: 0,
   },
+  {
+    // Same wire-level model as `mercury-edit-2`, but routed through the
+    // Mercury Next Edit endpoint instead of FIM. Picked by users who want
+    // multi-line next-edit predictions with the jump-to-edit UX.
+    id: "inception/mercury-next-edit",
+    modelID: "mercury-next-edit",
+    label: "Mercury Next Edit",
+    providerID: "inception",
+    provider: "Inception",
+    requestModel: "mercury-edit-2",
+    directProvider: "inception",
+    temperature: 0,
+    kind: "edit",
+  },
 ]
 
 export const AUTOCOMPLETE_MODELS: readonly AutocompleteModelDef[] = models

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

@@ -0,0 +1,107 @@
+/**
+ * Mercury Next Edit prompt assembly. Lives in the gateway so every client
+ * (VS Code, JetBrains, TUI) sends the same structured editor context and the
+ * Mercury-specific sentinel format is defined in exactly one place.
+ *
+ * Tag set is defined by the model and must be reproduced verbatim — see
+ * https://docs.inceptionlabs.ai/capabilities/next-edit
+ */
+
+const RECENTLY_VIEWED_SNIPPETS_OPEN = "<|recently_viewed_code_snippets|>"
+const RECENTLY_VIEWED_SNIPPETS_CLOSE = "<|/recently_viewed_code_snippets|>"
+const RECENTLY_VIEWED_SNIPPET_OPEN = "<|recently_viewed_code_snippet|>"
+const RECENTLY_VIEWED_SNIPPET_CLOSE = "<|/recently_viewed_code_snippet|>"
+const CURRENT_FILE_CONTENT_OPEN = "<|current_file_content|>"
+const CURRENT_FILE_CONTENT_CLOSE = "<|/current_file_content|>"
+const CODE_TO_EDIT_OPEN = "<|code_to_edit|>"
+const CODE_TO_EDIT_CLOSE = "<|/code_to_edit|>"
+const EDIT_DIFF_HISTORY_OPEN = "<|edit_diff_history|>"
+const EDIT_DIFF_HISTORY_CLOSE = "<|/edit_diff_history|>"
+const CURSOR = "<|cursor|>"
+/** Trailing token that tells the model this is a next-edit (not chat) request. */
+const UNIQUE_TOKEN = "<|!@#IS_NEXT_EDIT!@#|>"
+
+export interface MercuryRecentSnippet {
+  filepath: string
+  content: string
+}
+
+/** Editor-derived context a client sends; the gateway turns it into a prompt. */
+export interface MercuryEditContext {
+  currentFilePath: string
+  currentFileContent: string
+  cursorLine: number
+  cursorCharacter: number
+  editableRegionStartLine: number
+  editableRegionEndLine: number
+  recentlyViewedSnippets: MercuryRecentSnippet[]
+  editDiffHistory: string[]
+}
+
+function insertCursorToken(lines: string[], cursorLine: number, cursorCharacter: number): string[] {
+  if (cursorLine < 0 || cursorLine >= lines.length) return lines
+  const line = lines[cursorLine]
+  const safeChar = Math.min(Math.max(cursorCharacter, 0), line.length)
+  const next = line.slice(0, safeChar) + CURSOR + line.slice(safeChar)
+  return [...lines.slice(0, cursorLine), next, ...lines.slice(cursorLine + 1)]
+}
+
+export function recentlyViewedSnippetsBlock(snippets: MercuryRecentSnippet[]): string {
+  const inner = snippets
+    .map((s) =>
+      [RECENTLY_VIEWED_SNIPPET_OPEN, `code_snippet_file_path: ${s.filepath}`, s.content, RECENTLY_VIEWED_SNIPPET_CLOSE].join("\n"),
+    )
+    .join("\n")
+  return [RECENTLY_VIEWED_SNIPPETS_OPEN, inner, RECENTLY_VIEWED_SNIPPETS_CLOSE].join("\n")
+}
+
+export function currentFileContentBlock(
+  currentFilePath: string,
+  currentFileContent: string,
+  editableRegionStartLine: number,
+  editableRegionEndLine: number,
+  cursorLine: number,
+  cursorCharacter: number,
+): string {
+  const rawLines = currentFileContent.split("\n")
+  const withCursor = insertCursorToken(rawLines, cursorLine, cursorCharacter)
+  const start = Math.max(0, Math.min(editableRegionStartLine, withCursor.length))
+  const end = Math.max(start, Math.min(editableRegionEndLine, withCursor.length - 1))
+  const instrumented = [
+    ...withCursor.slice(0, start),
+    CODE_TO_EDIT_OPEN,
+    ...withCursor.slice(start, end + 1),
+    CODE_TO_EDIT_CLOSE,
+    ...withCursor.slice(end + 1),
+  ]
+  return [CURRENT_FILE_CONTENT_OPEN, `current_file_path: ${currentFilePath}`, instrumented.join("\n"), CURRENT_FILE_CONTENT_CLOSE].join("\n")
+}
+
+export function editDiffHistoryBlock(diffs: string[]): string {
+  // Each unidiff from `diff.createPatch` opens with an Index line + separator we
+  // strip. Diffs are blank-line separated so the model reads them as distinct hunks.
+  const trimmed = diffs.map((d) => {
+    const lines = d.split("\n")
+    return lines.length > 2 ? lines.slice(2).join("\n") : d
+  })
+  return [EDIT_DIFF_HISTORY_OPEN, trimmed.join("\n\n"), EDIT_DIFF_HISTORY_CLOSE].join("\n")
+}
+
+export function buildMercuryEditPrompt(ctx: MercuryEditContext): string {
+  return [
+    recentlyViewedSnippetsBlock(ctx.recentlyViewedSnippets),
+    "",
+    currentFileContentBlock(
+      ctx.currentFilePath,
+      ctx.currentFileContent,
+      ctx.editableRegionStartLine,
+      ctx.editableRegionEndLine,
+      ctx.cursorLine,
+      ctx.cursorCharacter,
+    ),
+    "",
+    editDiffHistoryBlock(ctx.editDiffHistory),
+    "",
+    UNIQUE_TOKEN,
+  ].join("\n")
+}

packages/kilo-gateway/src/edit.ts

@@ -0,0 +1,60 @@
+import { getAutocompleteModel, type DirectAutocompleteProviderID } from "./autocomplete.js"
+
+/**
+ * Env var(s) consulted as a fallback for BYOK keys when the provider hasn't
+ * been authenticated via the gateway's Auth store. Mirrors `DIRECT_FIM_ENV`.
+ */
+export const DIRECT_EDIT_ENV: Record<DirectAutocompleteProviderID, string[]> = {
+  mistral: ["MISTRAL_API_KEY"],
+  inception: ["INCEPTION_API_KEY"],
+}
+
+export type EditTarget =
+  | { provider: "inception"; model: string; url: string }
+  | { provider: "kilo"; model: string; url: string }
+
+/** Shape of the upstream (Mercury) chat/edit completion response we read from. */
+export interface EditUpstreamResponse {
+  choices?: Array<{ message?: { content?: string } }>
+  usage?: { prompt_tokens?: number; completion_tokens?: number }
+}
+
+const INCEPTION_EDIT_URL = "https://api.inceptionlabs.ai/v1/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.
+ */
+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 }
+  }
+  // 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.
+  return { provider: "kilo", model: info.requestModel, url: "" }
+}
+
+/**
+ * Mercury wraps the rewritten editable region in a triple-backtick fence,
+ * sometimes with a language tag and sometimes with `<|code_to_edit|>` sentinels
+ * inside. Strip all of that down to the bare code. Shared by both the hono and
+ * the Effect HttpApi edit handlers so the parsing can't drift between them.
+ */
+export function extractFencedBody(message: string): string {
+  if (!message) return ""
+  const fenceOpen = message.indexOf("```")
+  if (fenceOpen === -1) return message
+  const afterFenceOpen = message.indexOf("\n", fenceOpen + 3)
+  if (afterFenceOpen === -1) return ""
+  // A missing closing fence means the response was truncated (max_tokens hit).
+  // Take everything after the opening fence rather than dropping the suggestion.
+  const fenceClose = message.indexOf("```", afterFenceOpen + 1)
+  let body = fenceClose === -1 ? message.slice(afterFenceOpen + 1) : message.slice(afterFenceOpen + 1, fenceClose)
+  if (body.endsWith("\n")) body = body.slice(0, -1)
+  body = body.replace(/^<\|code_to_edit\|>\n?/, "")
+  body = body.replace(/\n?<\|\/code_to_edit\|>$/, "")
+  return body
+}

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

@@ -0,0 +1,91 @@
+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 type { AuthStore } from "./handlers.js"
+
+type Auth = Pick<AuthStore, "get">
+
+const EDIT_TIMEOUT_MS = 30_000
+const MAX_TOKENS_DEFAULT = 512
+
+async function getProviderKey(Auth: Auth, provider: DirectAutocompleteProviderID): Promise<string | undefined> {
+  const auth = await Auth.get(provider)
+  if (auth?.type === "api") return auth.key
+  return DIRECT_EDIT_ENV[provider].map((key) => process.env[key]).find(Boolean)
+}
+
+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") {
+      return c.json({ error: "Next Edit currently requires the Inception provider (mercury-edit-2)." }, 400 as any)
+    }
+
+    const token = await getProviderKey(Auth, target.provider)
+    if (!token) {
+      return c.json({ error: `Missing ${target.provider} provider API key` }, 401 as any)
+    }
+
+    // Build the Mercury sentinel prompt here so every client only sends
+    // structured editor context.
+    const content = buildMercuryEditPrompt(context as MercuryEditContext)
+    const signal = AbortSignal.any([c.req.raw.signal, AbortSignal.timeout(EDIT_TIMEOUT_MS)])
+    console.info(`[EDIT] request provider=${target.provider} model=${target.model} url=${target.url} chars=${content.length}`)
+
+    let response: Response
+    try {
+      response = await fetch(target.url, {
+        method: "POST",
+        headers: {
+          "Content-Type": "application/json",
+          Authorization: `Bearer ${token}`,
+        },
+        signal,
+        body: JSON.stringify({
+          model: target.model,
+          max_tokens: maxTokens ?? MAX_TOKENS_DEFAULT,
+          // Mercury rejects role:"system" on this endpoint — must be a single
+          // user message. See the integration's constants.ts for context.
+          messages: [{ role: "user", content }],
+        }),
+      })
+    } catch (err) {
+      if (err instanceof DOMException && err.name === "TimeoutError") {
+        return c.json({ error: "Edit request timed out" }, 504 as any)
+      }
+      if (signal.aborted) return c.json({ error: "Edit request canceled" }, 499 as any)
+      throw err
+    }
+
+    if (!response.ok) {
+      const text = await safeText(response)
+      return c.json({ error: `Edit request failed: ${response.status} ${text}` }, response.status as any)
+    }
+
+    const json = (await response.json()) as EditUpstreamResponse
+    const replyContent = json.choices?.[0]?.message?.content ?? ""
+    const body = extractFencedBody(replyContent)
+    return c.json({
+      content: body,
+      usage: json.usage
+        ? {
+            prompt_tokens: json.usage.prompt_tokens,
+            completion_tokens: json.usage.completion_tokens,
+          }
+        : undefined,
+    })
+  }
+}
+
+async function safeText(res: Response): Promise<string> {
+  try {
+    return await res.text()
+  } catch {
+    return "<unreadable>"
+  }
+}
+
+// Re-export the target type for tests + the opencode handler
+export type { EditTarget }

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

@@ -11,6 +11,7 @@ import { KILO_API_BASE, HEADER_FEATURE, HEADER_ORGANIZATIONID } from "../api/con
 import { buildKiloHeaders } from "../headers.js"
 import type { ImportDeps, DrizzleDb } from "../cloud-sessions.js"
 import { fetchCloudSession, fetchCloudSessionForImport, importSessionToDb } from "../cloud-sessions.js"
+import { createEditHandler } from "./edit.js"
 import { createFimHandler } from "./fim.js"
 import {
   GatewayError,
@@ -112,6 +113,16 @@ export function createKiloRoutes(deps: KiloRoutesDeps) {
     currentOrgId: z.string().nullable(),
   })
 
+  const EditCompletionResponse = z.object({
+    content: z.string(),
+    usage: z
+      .object({
+        prompt_tokens: z.number().optional(),
+        completion_tokens: z.number().optional(),
+      })
+      .optional(),
+  })
+
   const FimStreamChunk = z.object({
     choices: z
       .array(
@@ -325,6 +336,44 @@ export function createKiloRoutes(deps: KiloRoutesDeps) {
       ),
       createFimHandler(Auth),
     )
+    .post(
+      "/edit",
+      describeRoute({
+        summary: "Next Edit completion",
+        description:
+          "Proxy a Mercury-style Next Edit request. The client supplies structured editor " +
+          "context; the gateway assembles the sentinel-tagged prompt and forwards to the upstream edit endpoint.",
+        operationId: "kilo.edit",
+        responses: {
+          200: {
+            description: "Next Edit completion",
+            content: {
+              "application/json": {
+                schema: resolver(EditCompletionResponse),
+              },
+            },
+          },
+          ...errors(400, 401),
+        },
+      }),
+      validator(
+        "json",
+        z.object({
+          provider: z.string().optional(),
+          model: z.string().optional(),
+          maxTokens: z.number().optional(),
+          currentFilePath: z.string(),
+          currentFileContent: z.string(),
+          cursorLine: z.number(),
+          cursorCharacter: z.number(),
+          editableRegionStartLine: z.number(),
+          editableRegionEndLine: z.number(),
+          recentlyViewedSnippets: z.array(z.object({ filepath: z.string(), content: z.string() })),
+          editDiffHistory: z.array(z.string()),
+        }),
+      ),
+      createEditHandler(Auth),
+    )
     .post(
       "/audio/transcriptions",
       describeRoute({