fix(vscode): attribute Agent Manager sessions by root context

packages/kilo-gateway/src/provider-debug.ts

@@ -5,6 +5,7 @@ import { getApiKey } from "./auth/token.js"
 import { buildKiloHeaders, getDefaultHeaders } from "./headers.js"
 import { ANONYMOUS_API_KEY } from "./api/constants.js"
 import { resolveKiloOpenRouterBaseUrl } from "./api/url.js"
+import { buildRequestHeaders } from "./provider.js"
 
 /**
  * Debug version of createKilo with extensive logging
@@ -40,12 +41,7 @@ export function createKiloDebug(options: KiloProviderOptions = {}): SDK {
     console.log("  - URL:", String(input))
     console.log("  - Method:", init?.method || "GET")
 
-    const headers = new Headers(init?.headers)
-
-    // Add custom headers
-    Object.entries(customHeaders).forEach(([key, value]) => {
-      headers.set(key, value)
-    })
+    const headers = buildRequestHeaders(customHeaders, init?.headers)
 
     // Add authorization if API key exists
     if (apiKey) {

packages/kilo-gateway/src/provider.ts

@@ -10,6 +10,14 @@ import { ANONYMOUS_API_KEY } from "./api/constants.js"
 import { resolveKiloOpenRouterBaseUrl } from "./api/url.js"
 import { sanitizeResponsesBody } from "./responses.js"
 
+export function buildRequestHeaders(defaultHeaders: Record<string, string>, requestHeaders?: HeadersInit): Headers {
+  const headers = new Headers(defaultHeaders)
+  new Headers(requestHeaders).forEach((value, key) => {
+    headers.set(key, value)
+  })
+  return headers
+}
+
 /**
  * Create a KiloCode provider instance
  *
@@ -45,14 +53,9 @@ export function createKilo(options: KiloProviderOptions = {}): KiloProvider {
   // Create custom fetch wrapper to add dynamic headers
   const originalFetch = options.fetch ?? fetch
   const wrappedFetch = async (input: string | URL | Request, init?: RequestInit) => {
-    const headers = new Headers(init?.headers)
+    const headers = buildRequestHeaders(customHeaders, init?.headers)
     const body = sanitizeResponsesBody(input, init?.body)
 
-    // Add custom headers
-    Object.entries(customHeaders).forEach(([key, value]) => {
-      headers.set(key, value)
-    })
-
     // Add authorization if API key exists
     if (apiKey) {
       headers.set("Authorization", `Bearer ${apiKey}`)

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

@@ -0,0 +1,23 @@
+import { describe, expect, test } from "bun:test"
+import { buildRequestHeaders } from "../src/provider"
+
+describe("Kilo provider request headers", () => {
+  test("request headers override provider defaults", () => {
+    const headers = buildRequestHeaders(
+      {
+        "content-type": "application/json",
+        "x-kilocode-feature": "vscode-extension",
+        "x-default-only": "kept",
+      },
+      {
+        "x-kilocode-feature": "agent-manager",
+        "x-request-only": "kept-too",
+      },
+    )
+
+    expect(headers.get("content-type")).toBe("application/json")
+    expect(headers.get("x-kilocode-feature")).toBe("agent-manager")
+    expect(headers.get("x-default-only")).toBe("kept")
+    expect(headers.get("x-request-only")).toBe("kept-too")
+  })
+})

packages/kilo-vscode/src/KiloProvider.ts

@@ -1336,7 +1336,10 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper
 
     try {
       const workspaceDir = this.getContextDirectory()
-      const { data: session } = await this.client.session.create({ directory: workspaceDir }, { throwOnError: true })
+      const { data: session } = await this.client.session.create(
+        { directory: workspaceDir, platform: this.opts.platform },
+        { throwOnError: true },
+      )
       this.setCurrentSession(session)
       this.contextSessionID = session.id
       this.trackDirectory(session.id, workspaceDir)
@@ -2438,7 +2441,10 @@ export class KiloProvider implements vscode.WebviewViewProvider, TelemetryProper
     })
 
     if (!sessionID && !this.currentSession) {
-      const { data: session } = await this.client.session.create({ directory: dir }, { throwOnError: true })
+      const { data: session } = await this.client.session.create(
+        { directory: dir, platform: this.opts.platform },
+        { throwOnError: true },
+      )
       this.setCurrentSession(session)
       this.contextSessionID = session.id
       this.trackDirectory(session.id, dir)

packages/kilo-vscode/src/agent-manager/vscode-host.ts

@@ -9,6 +9,7 @@ import * as vscode from "vscode"
 import type { Host, PanelContext, OutputHandle, SessionProvider, Disposable } from "./host"
 import type { KiloConnectionService } from "../services/cli-backend"
 import { KiloProvider } from "../KiloProvider"
+import { PLATFORM } from "./constants"
 import { DiffVirtualProvider } from "../DiffVirtualProvider"
 import { buildWebviewHtml } from "../utils"
 import { openFileInEditor, getWorkspaceRoot } from "../review-utils"
@@ -85,6 +86,7 @@ export class VscodeHost implements Host {
     })
 
     const provider = new KiloProvider(this.extensionUri, this.connectionService, this.context, {
+      platform: PLATFORM,
       slimEditMetadata: true,
     })
     if (this.diffVirtual) {

packages/kilo-vscode/src/kilo-provider/options.ts

@@ -1,5 +1,6 @@
 export type KiloProviderOptions = {
   projectDirectory?: string | null
+  platform?: string
   slimEditMetadata?: boolean
   tabTitle?: (title: string) => void
 }

packages/opencode/src/kilo-sessions/kilo-sessions.ts

@@ -467,11 +467,25 @@ export namespace KiloSessions {
       return
     }
 
+    const session = await Session.get(SessionID.make(sessionId)).catch(() => undefined)
+    if (session?.parentID) {
+      const parent = await get(session.parentID).catch(() => undefined)
+      if (!parent) await bootstrap(session.parentID)
+    }
+
     log.info("creating session", { sessionId })
 
+    const metadata = await meta(sessionId)
+
     const response = await client.fetch(`${client.url}/api/session`, {
       method: "POST",
-      body: JSON.stringify({ sessionId }),
+      body: JSON.stringify({
+        sessionId,
+        ...(session?.parentID ? { parentSessionId: session.parentID } : {}),
+        ...(session?.title ? { title: session.title } : {}),
+        platform: metadata.platform,
+        metadata,
+      }),
     })
 
     if (!response.ok) {
@@ -691,7 +705,7 @@ export namespace KiloSessions {
   }
 
   async function meta(sessionId?: string) {
-    const override = sessionId ? KiloSession.getPlatformOverride(sessionId) : undefined
+    const override = sessionId ? KiloSession.resolvePlatform(sessionId) : undefined
     const platform = override || process.env["KILO_PLATFORM"] || "cli"
     const orgId = await getOrgId()
     const gitBranch = await Vcs.branch().catch(() => undefined)

packages/opencode/src/kilocode/session/index.ts

@@ -58,6 +58,60 @@ export namespace KiloSession {
     return overrides.get(id)
   }
 
+  function getParentID(id: string): string | undefined {
+    const row = Database.use((db) =>
+      db
+        .select({ parentID: SessionTable.parent_id })
+        .from(SessionTable)
+        .where(eq(SessionTable.id, SessionID.make(id)))
+        .get(),
+    )
+    return row?.parentID ?? undefined
+  }
+
+  export function resolvePlatform(id: string): string | undefined {
+    const override = overrides.get(id)
+    if (override) return override
+
+    let current: string | undefined = id
+    const seen = new Set<string>()
+    while (current && !seen.has(current)) {
+      seen.add(current)
+      const parentID = getParentID(current)
+      if (!parentID) break
+      const parentOverride = overrides.get(parentID)
+      if (parentOverride) return parentOverride
+      current = parentID
+    }
+
+    return undefined
+  }
+
+  export function resolveRoot(id: string): string {
+    let root = id
+    const seen = new Set<string>()
+    while (!seen.has(root)) {
+      seen.add(root)
+      const parentID = getParentID(root)
+      if (!parentID) return root
+      root = parentID
+    }
+    return root
+  }
+
+  export function featureForPlatform(platform: string | undefined): string | undefined {
+    switch (platform) {
+      case "agent-manager":
+        return "agent-manager"
+      case "vscode":
+        return "vscode-extension"
+      case "cli":
+        return "cli"
+      default:
+        return undefined
+    }
+  }
+
   export function clearPlatformOverride(id: string) {
     overrides.delete(id)
   }

packages/opencode/src/session/llm.ts

@@ -22,9 +22,10 @@ import { Auth } from "@/auth"
 // kilocode_change start
 import { DEFAULT_HEADERS } from "@/kilocode/const"
 import { getKiloProjectId } from "@/kilocode/project-id"
-import { HEADER_PROJECTID, HEADER_MACHINEID, HEADER_TASKID } from "@kilocode/kilo-gateway"
+import { ENV_FEATURE, HEADER_FEATURE, HEADER_PROJECTID, HEADER_MACHINEID, HEADER_TASKID } from "@kilocode/kilo-gateway"
 import { Identity } from "@kilocode/kilo-telemetry"
 import { makeRuntime } from "@/effect/run-service"
+import { KiloSession } from "@/kilocode/session"
 // kilocode_change end
 import { Installation } from "@/installation"
 import { InstallationVersion } from "@opencode-ai/core/installation/version"
@@ -104,6 +105,9 @@ const live: Layer.Layer<
         ],
         { concurrency: "unbounded" },
       )
+      const rootSessionID = KiloSession.resolveRoot(input.sessionID)
+      const platform = KiloSession.resolvePlatform(rootSessionID) ?? process.env["KILO_PLATFORM"]
+      const feature = KiloSession.featureForPlatform(platform) ?? process.env[ENV_FEATURE]
 
       // TODO: move this to a proper hook
       const isOpenaiOauth = item.id === "openai" && info?.type === "oauth"
@@ -417,7 +421,8 @@ const live: Layer.Layer<
           ...(isKilo && input.agent.name ? { "x-kilocode-mode": input.agent.name.toLowerCase() } : {}),
           ...(isKilo && kiloProjectId ? { [HEADER_PROJECTID]: kiloProjectId } : {}),
           ...(isKilo && machineId ? { [HEADER_MACHINEID]: machineId } : {}),
-          ...(isKilo ? { [HEADER_TASKID]: input.sessionID } : {}),
+          ...(isKilo ? { [HEADER_TASKID]: rootSessionID } : {}),
+          ...(isKilo && feature ? { [HEADER_FEATURE]: feature } : {}),
           // kilocode_change end
           ...input.model.headers,
           ...headers,

packages/opencode/src/session/session.ts

@@ -489,6 +489,7 @@ export const layer: Layer.Layer<Service, never, Bus.Service | Storage.Service |
       directory: string
       path?: string
       permission?: Permission.Ruleset
+      platform?: string
     }) {
       const ctx = yield* InstanceState.context
       const result: Info = {
@@ -509,6 +510,10 @@ export const layer: Layer.Layer<Service, never, Bus.Service | Storage.Service |
       }
       log.info("created", result)
 
+      if (input.platform) {
+        KiloSession.setPlatformOverride(result.id, input.platform)
+      }
+
       yield* sync.run(Event.Created, { sessionID: result.id, info: result })
 
       if (!Flag.KILO_EXPERIMENTAL_WORKSPACES) {
@@ -652,13 +657,9 @@ export const layer: Layer.Layer<Service, never, Bus.Service | Storage.Service |
         path: sessionPath(ctx.worktree, ctx.directory),
         title: input?.title,
         permission: input?.permission,
+        platform: input?.platform ?? (input?.parentID ? KiloSession.resolvePlatform(input.parentID) : undefined),
         workspaceID: input?.workspaceID ?? workspace, // kilocode_change - allow explicit override
       })
-      // kilocode_change start - store platform override for session ingest
-      if (input?.platform) {
-        KiloSession.setPlatformOverride(session.id, input.platform)
-      }
-      // kilocode_change end
       return session
     })
 

packages/opencode/test/session/session.test.ts

@@ -8,6 +8,7 @@ import { MessageV2 } from "../../src/session/message-v2"
 import { MessageID, PartID, type SessionID } from "../../src/session/schema"
 import { AppRuntime } from "../../src/effect/app-runtime"
 import { tmpdir } from "../fixture/fixture"
+import { KiloSession } from "../../src/kilocode/session"
 
 const projectRoot = path.join(__dirname, "../..")
 void Log.init({ print: false })
@@ -91,6 +92,26 @@ describe("session.created event", () => {
   })
 })
 
+describe("session platform attribution", () => {
+  test("child sessions inherit the root platform override", async () => {
+    await Instance.provide({
+      directory: projectRoot,
+      fn: async () => {
+        const root = await create({ platform: "agent-manager" })
+        const child = await create({ parentID: root.id, title: "child" })
+
+        expect(KiloSession.getPlatformOverride(root.id)).toBe("agent-manager")
+        expect(KiloSession.getPlatformOverride(child.id)).toBe("agent-manager")
+        expect(KiloSession.resolvePlatform(child.id)).toBe("agent-manager")
+        expect(KiloSession.resolveRoot(child.id)).toBe(root.id)
+        expect(KiloSession.featureForPlatform(KiloSession.resolvePlatform(child.id))).toBe("agent-manager")
+
+        await remove(root.id)
+      },
+    })
+  })
+})
+
 describe("step-finish token propagation via Bus event", () => {
   test(
     "non-zero tokens propagate through PartUpdated event",