From 0f139f31206dd53ec16566b78c8a8216a6383e54 Mon Sep 17 00:00:00 2001 From: Pedro Heyerdahl Date: Wed, 18 Feb 2026 01:40:07 -0300 Subject: [PATCH 1/2] feat(usage): add feature attribution header to API requests Introduce X-KILOCODE-FEATURE header for tracking request origin across different entry points (CLI, VSCode extension, autocomplete). - Add HEADER_FEATURE, DEFAULT_FEATURE, and ENV_FEATURE constants - Add getFeatureHeader() helper reading from KILOCODE_FEATURE env var - Include feature header in buildKiloHeaders() when env var is set - Set feature to "autocomplete" for FIM/completion requests - Set KILOCODE_FEATURE="vscode-extension" when spawning CLI from VSCode - Default to "cli" feature in opencode entry point - Add unit tests for feature header logic --- packages/kilo-gateway/src/api/constants.ts | 11 +++ packages/kilo-gateway/src/headers.ts | 18 +++++ packages/kilo-gateway/src/index.ts | 4 +- packages/kilo-gateway/src/server/routes.ts | 7 +- .../services/cli-backend/server-manager.ts | 1 + packages/opencode/src/index.ts | 11 ++- .../kilocode/kilo-gateway-headers.test.ts | 67 +++++++++++++++++++ 7 files changed, 116 insertions(+), 3 deletions(-) create mode 100644 packages/opencode/test/kilocode/kilo-gateway-headers.test.ts diff --git a/packages/kilo-gateway/src/api/constants.ts b/packages/kilo-gateway/src/api/constants.ts index 50b677d5cd..c46260ca41 100644 --- a/packages/kilo-gateway/src/api/constants.ts +++ b/packages/kilo-gateway/src/api/constants.ts @@ -59,3 +59,14 @@ 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" + +/** Default feature value (CLI usage) */ +export const DEFAULT_FEATURE = "cli" + +/** Environment variable name for feature override */ +export const ENV_FEATURE = "KILOCODE_FEATURE" +// kilocode_change end diff --git a/packages/kilo-gateway/src/headers.ts b/packages/kilo-gateway/src/headers.ts index 5882cf4a00..cb0f167c96 100644 --- a/packages/kilo-gateway/src/headers.ts +++ b/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 { + // kilocode_change start + const feature = getFeatureHeader() const headers: Record = { [X_KILOCODE_EDITORNAME]: getEditorNameHeader(), + ...(feature ? { [X_KILOCODE_FEATURE]: feature } : {}), } + // kilocode_change end if (metadata?.taskId) { headers[X_KILOCODE_TASKID] = metadata.taskId diff --git a/packages/kilo-gateway/src/index.ts b/packages/kilo-gateway/src/index.ts index 3b23e57ace..a654a419e8 100644 --- a/packages/kilo-gateway/src/index.ts +++ b/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" diff --git a/packages/kilo-gateway/src/server/routes.ts b/packages/kilo-gateway/src/server/routes.ts index 0d5b7d371d..7393122c71 100644 --- a/packages/kilo-gateway/src/server/routes.ts +++ b/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, diff --git a/packages/kilo-vscode/src/services/cli-backend/server-manager.ts b/packages/kilo-vscode/src/services/cli-backend/server-manager.ts index 20e9dbafa1..679df2045b 100644 --- a/packages/kilo-vscode/src/services/cli-backend/server-manager.ts +++ b/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, diff --git a/packages/opencode/src/index.ts b/packages/opencode/src/index.ts index dd7caebab9..cc79306bc2 100644 --- a/packages/opencode/src/index.ts +++ b/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" diff --git a/packages/opencode/test/kilocode/kilo-gateway-headers.test.ts b/packages/opencode/test/kilocode/kilo-gateway-headers.test.ts new file mode 100644 index 0000000000..32050c74bf --- /dev/null +++ b/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") + }) +}) From 1be28bb24e0872b35656c54ccc4946be72fc5777 Mon Sep 17 00:00:00 2001 From: Pedro Heyerdahl Date: Wed, 18 Feb 2026 14:07:00 -0300 Subject: [PATCH 2/2] refactor(usage): remove DEFAULT_FEATURE constant and derive feature from command context Remove the hardcoded DEFAULT_FEATURE ("cli") constant from kilo-gateway and instead derive the feature value at runtime in opencode. When the KILOCODE_FEATURE env var is not set, the "serve" command is now tagged as "unknown" (to surface misconfiguration when a caller forgets to set the env var), while all other commands default to "cli". --- packages/kilo-gateway/src/api/constants.ts | 3 --- 1 file changed, 3 deletions(-) diff --git a/packages/kilo-gateway/src/api/constants.ts b/packages/kilo-gateway/src/api/constants.ts index c46260ca41..c5cf44a906 100644 --- a/packages/kilo-gateway/src/api/constants.ts +++ b/packages/kilo-gateway/src/api/constants.ts @@ -64,9 +64,6 @@ export const TESTER_SUPPRESS_VALUE = "SUPPRESS" /** Header name for feature tracking */ export const HEADER_FEATURE = "X-KILOCODE-FEATURE" -/** Default feature value (CLI usage) */ -export const DEFAULT_FEATURE = "cli" - /** Environment variable name for feature override */ export const ENV_FEATURE = "KILOCODE_FEATURE" // kilocode_change end