Merge pull request #407 from Kilo-Org/feat/usage-feature-attribution

feat(usage): add feature attribution header to API requests

Author: pedroheyerdahl

packages/kilo-gateway/src/api/constants.ts

@@ -59,3 +59,11 @@ export const ENV_EDITOR_NAME = "KILOCODE_EDITOR_NAME"
 
 /** Tester header value for suppressing warnings */
 export const TESTER_SUPPRESS_VALUE = "SUPPRESS"
+
+// kilocode_change start
+/** Header name for feature tracking */
+export const HEADER_FEATURE = "X-KILOCODE-FEATURE"
+
+/** Environment variable name for feature override */
+export const ENV_FEATURE = "KILOCODE_FEATURE"
+// kilocode_change end

packages/kilo-gateway/src/headers.ts

@@ -4,11 +4,13 @@ import {
   HEADER_PROJECTID,
   HEADER_TESTER,
   HEADER_EDITORNAME,
+  HEADER_FEATURE, // kilocode_change
   USER_AGENT,
   CONTENT_TYPE,
   DEFAULT_EDITOR_NAME,
   ENV_EDITOR_NAME,
   TESTER_SUPPRESS_VALUE,
+  ENV_FEATURE, // kilocode_change
 } from "./api/constants.js"
 
 /**
@@ -20,6 +22,18 @@ export const X_KILOCODE_TASKID = HEADER_TASKID
 export const X_KILOCODE_PROJECTID = HEADER_PROJECTID
 export const X_KILOCODE_TESTER = HEADER_TESTER
 export const X_KILOCODE_EDITORNAME = HEADER_EDITORNAME
+export const X_KILOCODE_FEATURE = HEADER_FEATURE // kilocode_change
+
+/**
+ * Get feature header value from KILOCODE_FEATURE env var.
+ * Returns undefined when not set — the gateway stores NULL (unattributed).
+ * Callers must explicitly set the env var to get attribution.
+ */
+// kilocode_change start
+export function getFeatureHeader(): string | undefined {
+  return process.env[ENV_FEATURE] || undefined
+}
+// kilocode_change end
 
 /**
  * Default headers for KiloCode requests
@@ -47,9 +61,13 @@ export function buildKiloHeaders(
     kilocodeTesterWarningsDisabledUntil?: number
   },
 ): Record<string, string> {
+  // kilocode_change start
+  const feature = getFeatureHeader()
   const headers: Record<string, string> = {
     [X_KILOCODE_EDITORNAME]: getEditorNameHeader(),
+    ...(feature ? { [X_KILOCODE_FEATURE]: feature } : {}),
   }
+  // kilocode_change end
 
   if (metadata?.taskId) {
     headers[X_KILOCODE_TASKID] = metadata.taskId

packages/kilo-gateway/src/index.ts

@@ -9,7 +9,7 @@ export { KiloAuthPlugin, default } from "./plugin.js"
 export { createKilo } from "./provider.js"
 export { createKiloDebug } from "./provider-debug.js"
 export { kiloCustomLoader } from "./loader.js"
-export { buildKiloHeaders, getEditorNameHeader, DEFAULT_HEADERS } from "./headers.js"
+export { buildKiloHeaders, getEditorNameHeader, getFeatureHeader, DEFAULT_HEADERS } from "./headers.js" // kilocode_change - added getFeatureHeader
 
 // ============================================================================
 // Auth
@@ -89,7 +89,9 @@ export {
   HEADER_PROJECTID,
   HEADER_TESTER,
   HEADER_EDITORNAME,
+  HEADER_FEATURE, // kilocode_change
   DEFAULT_EDITOR_NAME,
   ENV_EDITOR_NAME,
   TESTER_SUPPRESS_VALUE,
+  ENV_FEATURE, // kilocode_change
 } from "./api/constants.js"

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

@@ -8,7 +8,8 @@
 
 import { fetchProfile, fetchBalance } from "../api/profile.js"
 import { fetchKilocodeNotifications, KilocodeNotificationSchema } from "../api/notifications.js"
-import { KILO_API_BASE } from "../api/constants.js"
+import { KILO_API_BASE, HEADER_FEATURE } from "../api/constants.js" // kilocode_change - added HEADER_FEATURE
+import { buildKiloHeaders } from "../headers.js" // kilocode_change
 
 // Type definitions for OpenCode dependencies (injected at runtime)
 type Hono = any
@@ -216,6 +217,10 @@ export function createKiloRoutes(deps: KiloRoutesDeps) {
           headers: {
             "Content-Type": "application/json",
             Authorization: `Bearer ${token}`,
+            // kilocode_change start - include kilo headers with autocomplete feature override
+            ...buildKiloHeaders(),
+            [HEADER_FEATURE]: "autocomplete",
+            // kilocode_change end
           },
           body: JSON.stringify({
             model: fimModel,

packages/kilo-vscode/src/services/cli-backend/server-manager.ts

@@ -66,6 +66,7 @@ export class ServerManager {
           ...process.env,
           KILO_SERVER_PASSWORD: password,
           KILO_CLIENT: "vscode",
+          KILOCODE_FEATURE: "vscode-extension", // kilocode_change - feature tracking
           KILO_TELEMETRY_LEVEL: vscode.env.isTelemetryEnabled ? "all" : "off",
           KILO_APP_NAME: "kilo-code",
           KILO_EDITOR_NAME: vscode.env.appName,

packages/opencode/src/index.ts

@@ -28,7 +28,16 @@ import { PrCommand } from "./cli/cmd/pr"
 import { SessionCommand } from "./cli/cmd/session"
 // kilocode_change start - Import telemetry and legacy migration
 import { Telemetry } from "@kilocode/kilo-telemetry"
-import { migrateLegacyKiloAuth } from "@kilocode/kilo-gateway"
+import { migrateLegacyKiloAuth, ENV_FEATURE } from "@kilocode/kilo-gateway"
+
+// kilocode_change - set feature for tracking. 'serve' is spawned by other services
+// (extension, cloud) which set their own KILOCODE_FEATURE env var. Direct CLI use
+// (any command other than 'serve') is tagged as 'cli'. If 'serve' is spawned without
+// the env var, it gets 'unknown' so the misconfiguration is visible in data.
+if (!process.env[ENV_FEATURE]) {
+  const isServe = process.argv.includes("serve")
+  process.env[ENV_FEATURE] = isServe ? "unknown" : "cli"
+}
 import { Global } from "./global"
 import { Config } from "./config/config"
 import { Auth } from "./auth"

packages/opencode/test/kilocode/kilo-gateway-headers.test.ts

@@ -0,0 +1,67 @@
+// kilocode_change - new file
+import { describe, it, expect, afterEach } from "bun:test"
+import { buildKiloHeaders, getFeatureHeader, getEditorNameHeader } from "@kilocode/kilo-gateway"
+import { HEADER_FEATURE, ENV_FEATURE } from "@kilocode/kilo-gateway"
+
+describe("getFeatureHeader", () => {
+  const original = process.env[ENV_FEATURE]
+
+  afterEach(() => {
+    if (original === undefined) {
+      delete process.env[ENV_FEATURE]
+    } else {
+      process.env[ENV_FEATURE] = original
+    }
+  })
+
+  it("returns undefined when env var is not set", () => {
+    delete process.env[ENV_FEATURE]
+    expect(getFeatureHeader()).toBeUndefined()
+  })
+
+  it("returns the env var value when set", () => {
+    process.env[ENV_FEATURE] = "cloud-agent"
+    expect(getFeatureHeader()).toBe("cloud-agent")
+  })
+
+  it("returns undefined for empty string", () => {
+    process.env[ENV_FEATURE] = ""
+    expect(getFeatureHeader()).toBeUndefined()
+  })
+})
+
+describe("buildKiloHeaders", () => {
+  const original = process.env[ENV_FEATURE]
+
+  afterEach(() => {
+    if (original === undefined) {
+      delete process.env[ENV_FEATURE]
+    } else {
+      process.env[ENV_FEATURE] = original
+    }
+  })
+
+  it("includes feature header when env var is set", () => {
+    process.env[ENV_FEATURE] = "vscode-extension"
+    const headers = buildKiloHeaders()
+    expect(headers[HEADER_FEATURE]).toBe("vscode-extension")
+  })
+
+  it("omits feature header when env var is not set", () => {
+    delete process.env[ENV_FEATURE]
+    const headers = buildKiloHeaders()
+    expect(headers[HEADER_FEATURE]).toBeUndefined()
+  })
+
+  it("always includes editor name header", () => {
+    delete process.env[ENV_FEATURE]
+    const headers = buildKiloHeaders()
+    expect(headers["X-KILOCODE-EDITORNAME"]).toBe(getEditorNameHeader())
+  })
+
+  it("passes through any feature value from env", () => {
+    process.env[ENV_FEATURE] = "custom-feature"
+    const headers = buildKiloHeaders()
+    expect(headers[HEADER_FEATURE]).toBe("custom-feature")
+  })
+})