feat(gateway): add per-request Kilo headers, body sanitization, and error contract

packages/core/src/kilocode/gateway-rpc.ts

@@ -46,8 +46,15 @@ export const ModelEntry = Schema.Struct({
 
 export type ModelEntry = typeof ModelEntry.Type
 
+export const SessionMetadata = Schema.Struct({
+  routedModelID: Schema.optional(Schema.String),
+  vercelID: Schema.optional(Schema.String),
+})
+
+export type SessionMetadata = typeof SessionMetadata.Type
+
 // The TUI and IDE clients use this contract for the account picker after login, /teams,
-// and Kilo-specific model catalog metadata.
+// Kilo-specific model catalog metadata, and per-session routed model metadata.
 export const Definition = Rpc.define({
   id: "kilocode.gateway",
   methods: {
@@ -66,6 +73,11 @@ export const Definition = Rpc.define({
       output: portable(Schema.Array(ModelEntry)),
       errors: Failure,
     },
+    "session.metadata": {
+      input: portable(Schema.Struct({ sessionID: Schema.String })),
+      output: portable(SessionMetadata),
+      errors: Failure,
+    },
   },
   events: {
     organization: { schema: portable(Schema.Struct({ organizationID: OrganizationID })) },

packages/core/src/kilocode/gateway.ts

@@ -2,10 +2,15 @@ import {
   KiloAccount,
   KiloDefaults,
   KiloDevice,
+  KiloErrors,
   KiloHeaders,
+  KiloMachine,
   KiloModels,
   KiloProfile,
+  KiloProject,
+  KiloRequests,
   KiloServer,
+  KiloStream,
 } from "@kilocode/gateway-core"
 import type { ConnectionInfo } from "@opencode/client"
 import type { IntegrationOAuthMethodRegistration } from "@opencode/plugin/effect/integration"
@@ -15,13 +20,14 @@ import { IntegrationID, IntegrationMethodID } from "@opencode/schema/integration
 import { Model } from "@opencode/schema/model"
 import { Money } from "@opencode/schema/money"
 import { Provider } from "@opencode/schema/provider"
+import { Global } from "@opencode/util/global"
 import { createHash, randomUUID } from "node:crypto"
 import { Clock, Duration, Effect, Fiber, Option, Schedule, Schema, Scope, Semaphore, Stream } from "effect"
 import { Bus } from "../bus.js"
 import { Variant } from "../variant.js"
 import { KiloGatewayRpc } from "./gateway-rpc.js"
 
-// v2 adapter for the Kilo Gateway (issue #14991, PR 1 & PR 2: login, teams, and catalog authority).
+// v2 adapter for the Kilo Gateway (issue #14991, PRs 1-3: login, teams, catalog authority, requests, and errors).
 // Kilo logic lives in @kilocode/gateway-core; this file maps it onto the v2 plugin API.
 const id = Provider.ID.make("kilo")
 const integrationID = IntegrationID.make("kilo")
@@ -76,7 +82,25 @@ export const KiloGatewayPlugin = define<Bus.Service | Scope.Scope>({
       anonymous: true,
     }
     const lastGood = new Map<string, readonly KiloModels.Model[]>()
+    const sessionMetadata = new Map<string, KiloGatewayRpc.SessionMetadata>()
     let retryFiber: Fiber.Fiber<void, never> | undefined
+    let cachedProjectID: string | null | undefined
+    let cachedMachineID: string | null | undefined
+
+    const getProjectID = Effect.gen(function* () {
+      if (cachedProjectID !== undefined) return cachedProjectID ?? undefined
+      const resolved = yield* KiloProject.resolve(ctx.location.directory)
+      cachedProjectID = resolved ?? null
+      return resolved
+    })
+
+    const getMachineID = Effect.gen(function* () {
+      if (process.env.KILO_MACHINE_ID?.trim()) return process.env.KILO_MACHINE_ID.trim()
+      if (cachedMachineID !== undefined) return cachedMachineID ?? undefined
+      const resolved = yield* KiloMachine.get(Global.Path.data)
+      cachedMachineID = resolved ?? null
+      return resolved
+    })
 
     const cancelRetry = Effect.gen(function* () {
       if (!retryFiber) return
@@ -403,6 +427,80 @@ export const KiloGatewayPlugin = define<Bus.Service | Scope.Scope>({
       }
     })
 
+    // Per-request attribution headers (`X-KILOCODE-TASKID`, `X-KILOCODE-PARENT-TASKID`, `X-KILOCODE-PROJECTID`,
+    // `X-KILOCODE-MACHINEID`, `x-kilocode-mode`, `User-Agent`, plus provider headers).
+    yield* ctx.session.hook(
+      "model.request",
+      (event) =>
+        Effect.gen(function* () {
+          const [session, projectID, machineID] = yield* Effect.all(
+            [
+              ctx.session.get({ sessionID: event.sessionID }).pipe(Effect.option),
+              getProjectID,
+              getMachineID,
+            ],
+            { concurrency: 3 },
+          )
+          const parentSessionID = Option.isSome(session) ? session.value.parentID : undefined
+          const requestHeaders = KiloHeaders.request({
+            version: ctx.app.version,
+            organizationID: catalog.account?.organizationID,
+            sessionID: event.sessionID,
+            parentSessionID,
+            projectID,
+            machineID,
+            mode: event.agent,
+          })
+          Object.assign(event.headers, requestHeaders)
+        }),
+      { providerID: id },
+    )
+
+    // Request body and auth normalization (`x-api-key` -> `Bearer`, credential metadata stripping,
+    // stateless `/responses` item stripping, and `provider.data_collection`).
+    yield* ctx.session.hook(
+      "http.request",
+      (event) =>
+        Effect.promise(async () => {
+          const dataCollection =
+            process.env.KILO_DATA_COLLECTION === "deny"
+              ? ("deny" as const)
+              : process.env.KILO_DATA_COLLECTION === "allow"
+                ? ("allow" as const)
+                : undefined
+          event.request = await KiloRequests.rewriteRequest(event.request, { dataCollection })
+        }),
+      { providerID: id },
+    )
+
+    // Error contract and SSE stream metadata (`routedModelID`, `x-vercel-id`, and mid-stream SSE errors).
+    yield* ctx.session.hook(
+      "http.response",
+      (event) =>
+        Effect.promise(async () => {
+          if (event.response.status >= 400) {
+            event.response = await KiloErrors.rewriteHttpResponse(event.response)
+            return
+          }
+          event.response = KiloStream.transformResponse(event.response, (meta) => {
+            if (meta.routedModelID || meta.vercelID) {
+              const prev = sessionMetadata.get(event.sessionID)
+              if (!prev && sessionMetadata.size >= 500) {
+                const oldest = sessionMetadata.keys().next().value
+                if (oldest) sessionMetadata.delete(oldest)
+              }
+              sessionMetadata.set(event.sessionID, {
+                ...(prev?.routedModelID ? { routedModelID: prev.routedModelID } : {}),
+                ...(prev?.vercelID ? { vercelID: prev.vercelID } : {}),
+                ...(meta.routedModelID ? { routedModelID: meta.routedModelID } : {}),
+                ...(meta.vercelID ? { vercelID: meta.vercelID } : {}),
+              })
+            }
+          })
+        }),
+      { providerID: id },
+    )
+
     yield* ctx.rpc
       .register(KiloGatewayRpc.Definition, {
         profile: (_input, context) =>
@@ -442,6 +540,8 @@ export const KiloGatewayPlugin = define<Bus.Service | Scope.Scope>({
             }
             return KiloModels.entries(catalog.models)
           }),
+        "session.metadata": (input) =>
+          Effect.sync(() => sessionMetadata.get(input.sessionID) ?? {}),
       })
       .pipe(Effect.orDie)
 

packages/core/test/kilocode/gateway.test.ts

@@ -8,10 +8,13 @@ import { KiloGatewayPlugin } from "@opencode/core/kilocode/gateway"
 import { KiloGatewayRpc } from "@opencode/core/kilocode/gateway-rpc"
 import { Model } from "@opencode/core/model"
 import { Plugin } from "@opencode/core/plugin"
+import { PluginHooks } from "@opencode/core/plugin/hooks"
 import { PluginHost } from "@opencode/core/plugin/host"
 import { ProviderPlugins } from "@opencode/core/plugin/provider"
 import { Provider } from "@opencode/core/provider"
 import { Rpc } from "@opencode/core/rpc"
+import { Agent } from "@opencode/schema/agent"
+import { Session } from "@opencode/schema/session"
 import { withEnv } from "../fixture/env"
 import { testEffect } from "../lib/effect"
 import { PluginTestLayer } from "../plugin/fixture"
@@ -28,6 +31,8 @@ const env = {
   KILO_API_URL: undefined,
   KILO_AI_GATEWAY_URL: undefined,
   KILO_ORG_ID: undefined,
+  KILO_MACHINE_ID: undefined,
+  KILO_DATA_COLLECTION: undefined,
   KILOCODE_EDITOR_NAME: undefined,
   KILOCODE_FEATURE: undefined,
   KILOCODE_VERSION: undefined,
@@ -408,4 +413,137 @@ describe("KiloGatewayPlugin", () => {
       (backend) => Effect.sync(() => backend[Symbol.dispose]()),
     ),
   )
+
+  it.live(
+    "injects per-request headers, sanitizes request bodies, rewrites Gateway errors, and records routedModelID",
+    () =>
+      Effect.acquireUseRelease(
+        Effect.sync(() => fixture()),
+        (backend) =>
+          withEnv(
+            {
+              ...env,
+              KILO_API_URL: backend.url,
+              KILO_API_KEY: "env-token",
+              KILO_ORG_ID: "selected",
+              KILO_MACHINE_ID: "mach-test",
+              KILO_DATA_COLLECTION: "deny",
+            },
+            () =>
+              Effect.gen(function* () {
+                const host = yield* addPlugin()
+                const hooks = yield* PluginHooks.Service
+                const rpc = yield* Rpc.Service
+                const parent = yield* host.session.create({ title: "parent" })
+                const child = yield* host.session.create({ title: "child", parentID: parent.id })
+                const modelRef = {
+                  providerID: kilo,
+                  id: Model.ID.make("kilo-auto/frontier"),
+                }
+
+                const modelEvt = yield* hooks.trigger("session", "model.request", {
+                  sessionID: child.id,
+                  agent: Agent.ID.make("code"),
+                  model: modelRef,
+                  kind: "primary",
+                  headers: {},
+                })
+                expect(modelEvt.headers).toMatchObject({
+                  "HTTP-Referer": "https://kilocode.ai",
+                  "X-Title": "Kilo Code",
+                  "X-KILOCODE-ORGANIZATIONID": "selected",
+                  "X-KILOCODE-TASKID": child.id,
+                  "X-KILOCODE-PARENT-TASKID": parent.id,
+                  "X-KILOCODE-MACHINEID": "mach-test",
+                  "x-kilocode-mode": "code",
+                  "User-Agent": "opencode-kilo-provider/unknown",
+                })
+
+                const httpReq = yield* hooks.trigger("session", "http.request", {
+                  sessionID: child.id,
+                  agent: Agent.ID.make("code"),
+                  model: modelRef,
+                  kind: "primary",
+                  request: new Request(`${backend.url}/api/gateway/responses`, {
+                    method: "POST",
+                    headers: { "content-type": "application/json", "x-api-key": "env-token" },
+                    body: JSON.stringify({
+                      model: "openai/gpt-5.4",
+                      server: backend.url,
+                      organizationID: "selected",
+                      input: [
+                        { type: "item_reference", id: "item_1" },
+                        { type: "message", id: "msg_1", role: "user", content: "hello" },
+                      ],
+                    }),
+                  }),
+                })
+                expect(httpReq.request.headers.get("authorization")).toBe("Bearer env-token")
+                expect(httpReq.request.headers.get("x-api-key")).toBeNull()
+                expect(yield* Effect.promise(() => httpReq.request.json())).toEqual({
+                  model: "openai/gpt-5.4",
+                  input: [{ type: "message", role: "user", content: "hello" }],
+                  provider: { data_collection: "deny" },
+                })
+
+                const errorRes = yield* hooks.trigger("session", "http.response", {
+                  sessionID: child.id,
+                  agent: Agent.ID.make("code"),
+                  model: modelRef,
+                  kind: "primary",
+                  request: httpReq.request,
+                  response: Response.json(
+                    {
+                      error: {
+                        title: "Low balance",
+                        message: "Balance is depleted.",
+                        balance: 0,
+                        buyCreditsUrl: "https://app.kilo.ai/billing",
+                      },
+                      error_type: "usage_limit_exceeded",
+                    },
+                    { status: 402 },
+                  ),
+                })
+                expect(yield* Effect.promise(() => errorRes.response.json())).toMatchObject({
+                  error: {
+                    code: "usage_limit_exceeded",
+                    message:
+                      "Low balance: Balance is depleted. (Balance: $0.00) Buy credits at https://app.kilo.ai/billing",
+                  },
+                })
+
+                const sseBody = [
+                  `data: ${JSON.stringify({
+                    model: "anthropic/claude-sonnet-4.6",
+                    choices: [{ delta: { content: "PONG" }, finish_reason: "stop" }],
+                    usage: { prompt_tokens: 4, completion_tokens: 1, total_tokens: 5 },
+                  })}`,
+                  "data: [DONE]",
+                  "",
+                ].join("\n")
+                const streamRes = yield* hooks.trigger("session", "http.response", {
+                  sessionID: Session.ID.make(child.id),
+                  agent: Agent.ID.make("code"),
+                  model: modelRef,
+                  kind: "primary",
+                  request: httpReq.request,
+                  response: new Response(sseBody, {
+                    status: 200,
+                    headers: { "content-type": "text/event-stream", "x-vercel-id": "iad1::abc" },
+                  }),
+                })
+                const streamText = yield* Effect.promise(() => streamRes.response.text())
+                expect(streamText).toContain('"routedModelID":"anthropic/claude-sonnet-4.6"')
+                expect(
+                  yield* rpc.call("kilocode.gateway", "session.metadata", { sessionID: child.id }),
+                ).toEqual({
+                  routedModelID: "anthropic/claude-sonnet-4.6",
+                  vercelID: "iad1::abc",
+                })
+              }),
+          ),
+        (backend) => Effect.sync(() => backend[Symbol.dispose]()),
+      ),
+  )
 })