Merge pull request #9448 from Kilo-Org/fix/costs-of-subagents

fix(cli): include subagent costs in session total
2026-08-28 19:11:03 +08:00 · 2026-04-27 13:30:17 +03:00
parent b344ac97c6 e101781bfc
commit 5be897cec3
11 changed files with 706 additions and 9 deletions
@@ -0,0 +1,5 @@
+---
+"@kilocode/cli": patch
+---
+
+Fix session cost display missing subagent costs. The TUI footer, sidebar, web context panel, and ACP usage reports now include the cost of every subagent the session spawned, including nested ones.
@@ -193,7 +193,11 @@ export async function aggregateSessionStats(days?: number, projectFilter?: strin

      for (const message of messages) {
        if (message.info.role === "assistant") {
-          sessionCost += message.info.cost || 0
+          // kilocode_change start - count propagated subagent cost once but keep child model stats (#6321)
+          const parts = message.parts.filter((part) => part.type === "step-finish")
+          const cost = parts.length ? parts.reduce((sum, part) => sum + part.cost, 0) : message.info.cost || 0
+          if (!session.parentID) sessionCost += message.info.cost || 0
+          // kilocode_change end

          const modelKey = `${message.info.providerID}/${message.info.modelID}`
          if (!sessionModelUsage[modelKey]) {
@@ -204,7 +208,7 @@ export async function aggregateSessionStats(days?: number, projectFilter?: strin
            }
          }
          sessionModelUsage[modelKey].messages++
-          sessionModelUsage[modelKey].cost += message.info.cost || 0
+          sessionModelUsage[modelKey].cost += cost // kilocode_change

          if (message.info.tokens) {
            sessionTokens.input += message.info.tokens.input || 0
@@ -0,0 +1,69 @@
+// kilocode_change - new file
+import { Effect } from "effect"
+import { Session } from "@/session"
+import { MessageV2 } from "@/session/message-v2"
+import { SessionID, MessageID } from "@/session/schema"
+
+export namespace KiloCostPropagation {
+  /**
+   * Per-key promise chain that serializes concurrent `propagate` calls against
+   * the same parent message. Prevents lost updates when the LLM launches
+   * several `task` tool calls in parallel (each release stage races to
+   * read-modify-write the same parent cost field).
+   */
+  const locks = new Map<string, Promise<void>>()
+
+  function acquire(key: string): Promise<() => void> {
+    const prev = locks.get(key) ?? Promise.resolve()
+    let release!: () => void
+    const current = new Promise<void>((r) => (release = r))
+    const chain = prev.catch(() => {}).then(() => current)
+    locks.set(key, chain)
+    return prev
+      .catch(() => {})
+      .then(() => () => {
+        release()
+        if (locks.get(key) === chain) locks.delete(key)
+      })
+  }
+
+  /**
+   * Total assistant-message cost in a session. Because each subagent propagates
+   * its own total into the parent assistant message when it finishes, this sum
+   * already reflects descendant sessions recursively — no tree walk needed.
+   */
+  export const childCost = Effect.fn("KiloCostPropagation.childCost")(function* (
+    sessions: Session.Interface,
+    id: SessionID,
+  ) {
+    const msgs = yield* sessions.messages({ sessionID: id })
+    return msgs.reduce((sum, m) => sum + (m.info.role === "assistant" ? m.info.cost : 0), 0)
+  })
+
+  /**
+   * Add `amount` to the given parent assistant message's cost. No-op when
+   * `amount` is non-positive or the target is not an assistant message.
+   *
+   * Concurrent calls against the same parent are serialized internally so the
+   * read-modify-write cannot lose updates when subagents complete in parallel.
+   */
+  export const propagate = Effect.fn("KiloCostPropagation.propagate")(function* (
+    sessions: Session.Interface,
+    sid: SessionID,
+    mid: MessageID,
+    amount: number,
+  ) {
+    if (!(amount > 0)) return
+    yield* Effect.acquireUseRelease(
+      Effect.promise(() => acquire(`${sid}:${mid}`)),
+      () =>
+        Effect.gen(function* () {
+          const parent = yield* Effect.sync(() => MessageV2.get({ sessionID: sid, messageID: mid }))
+          if (parent.info.role !== "assistant") return
+          parent.info.cost += amount
+          yield* sessions.updateMessage(parent.info)
+        }),
+      (release) => Effect.sync(() => release()),
+    )
+  })
+}
@@ -19,6 +19,7 @@ import type { Provider } from "@/provider"
 import { Question } from "@/question"
 import { KiloSessionProcessor } from "@/kilocode/session/processor" // kilocode_change
 import { Suggestion } from "@/kilocode/suggestion" // kilocode_change
+import { NotFoundError } from "@/storage" // kilocode_change
 import { errorMessage } from "@/util/error"
 import { Log } from "@/util"
 import { isRecord } from "@/util/record"
@@ -157,6 +158,22 @@ export const layer: Layer.Layer<
        return { call, part }
      })

+      // kilocode_change start - tolerate deleted sessions during subagent cost reconciliation (#6321)
+      const reconcile = Effect.fn("SessionProcessor.reconcileCost")(function* () {
+        const fresh = yield* Effect.sync(() => {
+          try {
+            return MessageV2.get({ sessionID: ctx.assistantMessage.sessionID, messageID: ctx.assistantMessage.id })
+          } catch (err) {
+            if (NotFoundError.isInstance(err)) return
+            throw err
+          }
+        })
+        if (fresh?.info.role !== "assistant") return
+        if (fresh.info.cost <= ctx.assistantMessage.cost) return
+        ctx.assistantMessage.cost = fresh.info.cost
+      })
+      // kilocode_change end
+
      const updateToolCall = Effect.fn("SessionProcessor.updateToolCall")(function* (
        toolCallID: string,
        update: (part: MessageV2.ToolPart) => MessageV2.ToolPart,
@@ -410,6 +427,9 @@ export const layer: Layer.Layer<
            })
            // kilocode_change end
            ctx.assistantMessage.finish = value.finishReason
+            // kilocode_change start - capture any subagent cost propagated by tool calls during this step (#6321)
+            yield* reconcile()
+            // kilocode_change end
            ctx.assistantMessage.cost += usage.cost
            ctx.assistantMessage.tokens = usage.tokens
            yield* session.updatePart({
@@ -567,6 +587,9 @@ export const layer: Layer.Layer<
        ctx.toolcalls = {}
        KiloSessionProcessor.guardEmptyToolCalls(ctx.assistantMessage, MessageV2.parts(ctx.assistantMessage.id)) // kilocode_change
        ctx.assistantMessage.time.completed = Date.now()
+        // kilocode_change start - reconcile cost with any subagent propagation written during tool calls (#6321)
+        yield* reconcile()
+        // kilocode_change end
        yield* session.updateMessage(ctx.assistantMessage)
      })

@@ -4,6 +4,7 @@ import fs from "fs/promises"
 import { KiloSessionPrompt } from "@/kilocode/session/prompt" // kilocode_change
 import { KiloSessionPromptQueue } from "@/kilocode/session/prompt-queue" // kilocode_change
 import { KiloSession } from "@/kilocode/session" // kilocode_change
+import { KiloCostPropagation } from "@/kilocode/session/cost-propagation" // kilocode_change
 import { Suggestion } from "@/kilocode/suggestion" // kilocode_change
 import { Question } from "@/question" // kilocode_change
 import z from "zod"
@@ -598,6 +599,12 @@ NOTE: At any point in time through this workflow you should feel free to ask the

      let error: Error | undefined
      const taskAbort = new AbortController()
+      // kilocode_change start - shared reader for the child session id written by task.ts ctx.metadata (#6321)
+      const childID = () => {
+        const meta = part.state.status !== "pending" ? part.state.metadata : undefined
+        return (meta as { sessionId?: string } | undefined)?.sessionId
+      }
+      // kilocode_change end
      const result = yield* taskTool
        .execute(taskArgs, {
          agent: task.agent,
@@ -636,6 +643,12 @@ NOTE: At any point in time through this workflow you should feel free to ask the
              taskAbort.abort()
              assistantMessage.finish = "tool-calls"
              assistantMessage.time.completed = Date.now()
+              // kilocode_change start - propagate partial subagent cost on cancel (#6321)
+              const cid = childID()
+              if (cid) {
+                assistantMessage.cost = yield* KiloCostPropagation.childCost(sessions, SessionID.make(cid))
+              }
+              // kilocode_change end
              yield* sessions.updateMessage(assistantMessage)
              if (part.state.status === "running") {
                yield* sessions.updatePart({
@@ -668,6 +681,12 @@ NOTE: At any point in time through this workflow you should feel free to ask the

      assistantMessage.finish = "tool-calls"
      assistantMessage.time.completed = Date.now()
+      // kilocode_change start - include subagent total cost on the wrapper message (#6321)
+      const cid = result?.metadata?.sessionId ?? childID()
+      if (cid) {
+        assistantMessage.cost = yield* KiloCostPropagation.childCost(sessions, SessionID.make(cid))
+      }
+      // kilocode_change end
      yield* sessions.updateMessage(assistantMessage)

      if (result && part.state.status === "running") {
@@ -9,6 +9,7 @@ import type { SessionPrompt } from "../session/prompt"
 import { Config } from "../config"
 import { Effect } from "effect"
 import { KiloTask } from "../kilocode/tool/task" // kilocode_change
+import { KiloCostPropagation } from "../kilocode/session/cost-propagation" // kilocode_change

 export interface TaskPromptOps {
  cancel(sessionID: SessionID): void
@@ -141,9 +142,12 @@ export const TaskTool = Tool.define(
      }

      return yield* Effect.acquireUseRelease(
-        Effect.sync(() => {
+        // kilocode_change start - snapshot child cost so we propagate only the delta on resume (#6321)
+        Effect.gen(function* () {
          ctx.abort.addEventListener("abort", cancel)
+          return yield* KiloCostPropagation.childCost(sessions, nextSession.id)
        }),
+        // kilocode_change end
        () =>
          Effect.gen(function* () {
            const parts = yield* ops.resolvePromptParts(params.prompt)
@@ -180,10 +184,14 @@ export const TaskTool = Tool.define(
              ].join("\n"),
            }
          }),
-        () =>
-          Effect.sync(() => {
+        // kilocode_change start - propagate subagent cost delta to parent on every exit path (#6321)
+        (costBefore) =>
+          Effect.gen(function* () {
            ctx.abort.removeEventListener("abort", cancel)
+            const costAfter = yield* KiloCostPropagation.childCost(sessions, nextSession.id)
+            yield* KiloCostPropagation.propagate(sessions, ctx.sessionID, ctx.messageID, costAfter - costBefore)
          }),
+        // kilocode_change end
      )
    })

@@ -0,0 +1,94 @@
+// Verifies KiloCostPropagation.propagate() serializes concurrent writes to
+// the same parent assistant message. Without the internal lock, parallel
+// subagent completions race on read-modify-write and lose deltas (#6321).
+
+import { afterEach, describe, expect } from "bun:test"
+import { Effect, Layer } from "effect"
+import { Bus } from "../../src/bus"
+import * as CrossSpawnSpawner from "../../src/effect/cross-spawn-spawner"
+import { KiloCostPropagation } from "../../src/kilocode/session/cost-propagation"
+import { Instance } from "../../src/project/instance"
+import { ProviderID, ModelID } from "../../src/provider/schema"
+import { Session } from "../../src/session"
+import { MessageV2 } from "../../src/session/message-v2"
+import { MessageID } from "../../src/session/schema"
+import { Log } from "../../src/util"
+import { provideTmpdirInstance } from "../fixture/fixture"
+import { testEffect } from "../lib/effect"
+
+Log.init({ print: false })
+
+afterEach(async () => {
+  await Instance.disposeAll()
+})
+
+const ref = {
+  providerID: ProviderID.make("test"),
+  modelID: ModelID.make("test-model"),
+}
+
+const it = testEffect(Layer.mergeAll(Session.defaultLayer, Bus.layer, CrossSpawnSpawner.defaultLayer))
+
+const seed = Effect.fn("CostPropagationTest.seed")(function* () {
+  const sessions = yield* Session.Service
+  const chat = yield* sessions.create({ title: "parent" })
+  const user = yield* sessions.updateMessage({
+    id: MessageID.ascending(),
+    role: "user",
+    sessionID: chat.id,
+    agent: "build",
+    model: ref,
+    time: { created: Date.now() },
+  })
+  const assistant: MessageV2.Assistant = {
+    id: MessageID.ascending(),
+    role: "assistant",
+    parentID: user.id,
+    sessionID: chat.id,
+    mode: "build",
+    agent: "build",
+    cost: 0,
+    path: { cwd: "/tmp", root: "/tmp" },
+    tokens: { input: 0, output: 0, reasoning: 0, cache: { read: 0, write: 0 } },
+    modelID: ref.modelID,
+    providerID: ref.providerID,
+    time: { created: Date.now() },
+  }
+  yield* sessions.updateMessage(assistant)
+  return { chat, assistant }
+})
+
+describe("KiloCostPropagation.propagate", () => {
+  it.live("sums deltas correctly under parallel execution", () =>
+    provideTmpdirInstance(() =>
+      Effect.gen(function* () {
+        const sessions = yield* Session.Service
+        const { chat, assistant } = yield* seed()
+        const deltas = [0.01, 0.02, 0.04, 0.08, 0.16, 0.32, 0.64, 1.28]
+        yield* Effect.all(
+          deltas.map((d) => KiloCostPropagation.propagate(sessions, chat.id, assistant.id, d)),
+          { concurrency: "unbounded" },
+        )
+        const parent = yield* Effect.sync(() => MessageV2.get({ sessionID: chat.id, messageID: assistant.id }))
+        expect(parent.info.role).toBe("assistant")
+        if (parent.info.role !== "assistant") return
+        const total = deltas.reduce((a, b) => a + b, 0)
+        expect(parent.info.cost).toBeCloseTo(total, 6)
+      }),
+    ),
+  )
+
+  it.live("is a no-op when amount is non-positive", () =>
+    provideTmpdirInstance(() =>
+      Effect.gen(function* () {
+        const sessions = yield* Session.Service
+        const { chat, assistant } = yield* seed()
+        yield* KiloCostPropagation.propagate(sessions, chat.id, assistant.id, 0)
+        yield* KiloCostPropagation.propagate(sessions, chat.id, assistant.id, -1.5)
+        const parent = yield* Effect.sync(() => MessageV2.get({ sessionID: chat.id, messageID: assistant.id }))
+        if (parent.info.role !== "assistant") return
+        expect(parent.info.cost).toBe(0)
+      }),
+    ),
+  )
+})
@@ -183,6 +183,77 @@ describe("session processor empty tool-calls", () => {
    ),
  )

+  it.live("ignores deleted session during cost reconciliation", () =>
+    provideTmpdirInstance(
+      (dir) =>
+        Effect.gen(function* () {
+          const test = yield* TestLLM
+          const processors = yield* SessionProcessor.Service
+          const session = yield* Session.Service
+
+          yield* test.reply(
+            { type: "start" },
+            { type: "start-step" } as LLM.Event,
+            {
+              type: "finish-step",
+              finishReason: "stop",
+              usage: usage(),
+              providerMetadata: undefined,
+            } as LLM.Event,
+            { type: "finish" } as LLM.Event,
+          )
+
+          const chat = yield* session.create({})
+          const parent = yield* session.updateMessage({
+            id: MessageID.ascending(),
+            role: "user",
+            sessionID: chat.id,
+            agent: "code",
+            model: ref,
+            time: { created: Date.now() },
+          })
+          const msg: MessageV2.Assistant = {
+            id: MessageID.ascending(),
+            role: "assistant",
+            sessionID: chat.id,
+            parentID: parent.id,
+            mode: "code",
+            agent: "code",
+            path: { cwd: path.resolve(dir), root: path.resolve(dir) },
+            cost: 0,
+            tokens: { input: 0, output: 0, reasoning: 0, cache: { read: 0, write: 0 } },
+            modelID: ref.modelID,
+            providerID: ref.providerID,
+            time: { created: Date.now() },
+          }
+          yield* session.updateMessage(msg)
+
+          const mdl = model()
+          const handle = yield* processors.create({
+            assistantMessage: msg,
+            sessionID: chat.id,
+            model: mdl,
+          })
+          yield* session.remove(chat.id)
+
+          const input: LLM.StreamInput = {
+            user: parent as MessageV2.User,
+            sessionID: chat.id,
+            model: mdl,
+            agent: { name: "code", mode: "primary", permission: [], options: {} } as any,
+            system: [],
+            messages: [],
+            tools: {},
+          }
+
+          const result = yield* handle.process(input)
+          expect(result).toBe("continue")
+          expect(handle.message.error).toBeUndefined()
+        }),
+      { git: true },
+    ),
+  )
+
  it.live("preserves tool-calls finish when tool parts exist", () =>
    provideTmpdirInstance(
      (dir) =>
@@ -0,0 +1,123 @@
+// Verifies `kilo stats` does not double-count subagent cost while still
+// including child-session messages, tokens, tools, and model usage. The task
+// tool propagates each child session's total cost up to the parent's
+// tool-wrapper assistant message (#6321).
+
+import { afterEach, describe, expect, test } from "bun:test"
+import { aggregateSessionStats } from "../../src/cli/cmd/stats"
+import { MessageV2 } from "../../src/session/message-v2"
+import { Instance } from "../../src/project/instance"
+import { ProviderID, ModelID } from "../../src/provider/schema"
+import { Session } from "../../src/session"
+import { MessageID, PartID } from "../../src/session/schema"
+import { Log } from "../../src/util"
+import { tmpdir } from "../fixture/fixture"
+
+Log.init({ print: false })
+
+afterEach(async () => {
+  await Instance.disposeAll()
+})
+
+const ref = {
+  providerID: ProviderID.make("test"),
+  modelID: ModelID.make("test-model"),
+}
+
+function assistant(sessionID: string, parentID: string, cost: number): MessageV2.Assistant {
+  return {
+    id: MessageID.ascending(),
+    role: "assistant",
+    parentID: parentID as any,
+    sessionID: sessionID as any,
+    mode: "build",
+    agent: "build",
+    cost,
+    path: { cwd: "/tmp", root: "/tmp" },
+    tokens: { input: 10, output: 5, reasoning: 0, cache: { read: 0, write: 0 } },
+    modelID: ref.modelID,
+    providerID: ref.providerID,
+    time: { created: Date.now() },
+  }
+}
+
+async function step(sessionID: string, messageID: string, cost: number) {
+  await Session.updatePart({
+    id: PartID.ascending(),
+    messageID: messageID as any,
+    sessionID: sessionID as any,
+    type: "step-finish",
+    reason: "stop",
+    cost,
+    tokens: { total: 15, input: 10, output: 5, reasoning: 0, cache: { read: 0, write: 0 } },
+  })
+}
+
+async function tool(sessionID: string, messageID: string) {
+  const time = Date.now()
+  await Session.updatePart({
+    id: PartID.ascending(),
+    messageID: messageID as any,
+    sessionID: sessionID as any,
+    type: "tool",
+    callID: "call_1",
+    tool: "bash",
+    state: {
+      status: "completed",
+      input: {},
+      output: "ok",
+      title: "bash",
+      metadata: {},
+      time: { start: time, end: time },
+    },
+  })
+}
+
+describe("stats subagent cost", () => {
+  test("counts child usage without double-counting propagated cost", async () => {
+    await using tmp = await tmpdir({ git: true })
+    await Instance.provide({
+      directory: tmp.path,
+      fn: async () => {
+        const parent = await Session.create({ title: "root" })
+        const child = await Session.create({ parentID: parent.id, title: "subagent" })
+
+        const userMsg = await Session.updateMessage({
+          id: MessageID.ascending(),
+          role: "user",
+          sessionID: parent.id,
+          agent: "build",
+          model: ref,
+          time: { created: Date.now() },
+        } as any)
+        const parentMsg = await Session.updateMessage(assistant(parent.id, userMsg.id, 1.5))
+        await step(parent.id, parentMsg.id, 1)
+
+        const childUser = await Session.updateMessage({
+          id: MessageID.ascending(),
+          role: "user",
+          sessionID: child.id,
+          agent: "general",
+          model: ref,
+          time: { created: Date.now() },
+        } as any)
+        const childMsg = await Session.updateMessage(assistant(child.id, childUser.id, 0.5))
+        await step(child.id, childMsg.id, 0.5)
+        await tool(child.id, childMsg.id)
+
+        const stats = await aggregateSessionStats()
+        const model = stats.modelUsage["test/test-model"]!
+        expect(stats.totalCost).toBeCloseTo(1.5, 6)
+        expect(stats.totalSessions).toBe(2)
+        expect(stats.totalMessages).toBe(4)
+        expect(stats.totalTokens.input).toBe(20)
+        expect(stats.totalTokens.output).toBe(10)
+        expect(stats.toolUsage.bash).toBe(1)
+        expect(model.messages).toBe(2)
+        expect(model.tokens.input).toBe(20)
+        expect(model.tokens.output).toBe(10)
+        expect(model.cost).toBeCloseTo(1.5, 6)
+      },
+    })
+  })
+})
@@ -812,6 +812,65 @@ it.live(
  30_000,
 )

+// kilocode_change start - handleSubtask propagates child session cost to wrapper (#6321)
+it.live(
+  "handleSubtask propagates subagent cost to wrapper message",
+  () =>
+    provideTmpdirServer(
+      Effect.fnUntraced(function* ({ llm }) {
+        const prompt = yield* SessionPrompt.Service
+        const sessions = yield* Session.Service
+        const registry = yield* ToolRegistry.Service
+        const { task } = yield* registry.named()
+        const original = task.execute
+        // Simulate task tool: create a child session, persist an assistant with cost, return metadata.
+        task.execute = (_args, ctx) =>
+          Effect.gen(function* () {
+            const child = yield* sessions.create({ parentID: ctx.sessionID, title: "subagent" })
+            const childAssistant: MessageV2.Assistant = {
+              id: MessageID.ascending(),
+              role: "assistant",
+              parentID: ctx.messageID,
+              sessionID: child.id,
+              mode: "general",
+              agent: "general",
+              cost: 0.42,
+              path: { cwd: "/tmp", root: "/tmp" },
+              tokens: { input: 0, output: 0, reasoning: 0, cache: { read: 0, write: 0 } },
+              modelID: ref.modelID,
+              providerID: ref.providerID,
+              time: { created: Date.now(), completed: Date.now() },
+              finish: "stop",
+            }
+            yield* sessions.updateMessage(childAssistant)
+            yield* ctx.metadata({
+              title: "done",
+              metadata: { sessionId: child.id, model: ref, variant: undefined },
+            })
+            return { title: "done", metadata: { sessionId: child.id, model: ref, variant: undefined }, output: "done" }
+          })
+        yield* Effect.addFinalizer(() => Effect.sync(() => void (task.execute = original)))
+
+        const chat = yield* sessions.create({ title: "Pinned" })
+        const msg = yield* user(chat.id, "hello")
+        yield* addSubtask(chat.id, msg.id)
+        // The loop continues past handleSubtask into a normal LLM step; provide one response to exit.
+        yield* llm.text("wrapped")
+
+        yield* prompt.loop({ sessionID: chat.id })
+
+        const msgs = yield* MessageV2.filterCompactedEffect(chat.id)
+        const wrapper = msgs.find((item) => item.info.role === "assistant" && item.info.agent === "general")
+        expect(wrapper?.info.role).toBe("assistant")
+        if (!wrapper || wrapper.info.role !== "assistant") return
+        expect(wrapper.info.cost).toBeCloseTo(0.42, 6)
+      }),
+      { git: true, config: providerCfg },
+    ),
+  30_000,
+)
+// kilocode_change end
+
 it.live(
  "cancel with queued callers resolves all cleanly",
  () =>
@@ -7,7 +7,7 @@ import { Instance } from "../../src/project/instance"
 import { Session } from "../../src/session"
 import { MessageV2 } from "../../src/session/message-v2"
 import type { SessionPrompt } from "../../src/session/prompt"
-import { MessageID, PartID } from "../../src/session/schema"
+import { MessageID, PartID, SessionID } from "../../src/session/schema" // kilocode_change - SessionID used by cost propagation tests
 import { ModelID, ProviderID } from "../../src/provider/schema"
 import { TaskTool, type TaskPromptOps } from "../../src/tool/task"
 import { Truncate } from "../../src/tool"
@@ -64,17 +64,28 @@ const seed = Effect.fn("TaskToolTest.seed")(function* (title = "Pinned") {
  return { chat, assistant }
 })

-function stubOps(opts?: { onPrompt?: (input: SessionPrompt.PromptInput) => void; text?: string }): TaskPromptOps {
+// kilocode_change start - stub signature + prompt body extended to persist assistant cost for propagation tests
+function stubOps(opts?: {
+  onPrompt?: (input: SessionPrompt.PromptInput) => void
+  text?: string
+  sessions?: Session.Interface
+  childCost?: number
+}): TaskPromptOps {
  return {
    cancel() {},
    resolvePromptParts: (template) => Effect.succeed([{ type: "text" as const, text: template }]),
    prompt: (input) =>
-      Effect.sync(() => {
+      Effect.gen(function* () {
        opts?.onPrompt?.(input)
-        return reply(input, opts?.text ?? "done")
+        const rep = reply(input, opts?.text ?? "done")
+        if (opts?.sessions && opts?.childCost != null) {
+          yield* opts.sessions.updateMessage({ ...rep.info, cost: opts.childCost })
+        }
+        return rep
      }),
  }
 }
+// kilocode_change end

 function reply(input: SessionPrompt.PromptInput, text: string): MessageV2.WithParts {
  const id = MessageID.ascending()
@@ -395,3 +406,214 @@ describe("tool.task", () => {
    ),
  )
 })
+
+// kilocode_change start - subagent cost propagation coverage (#6321)
+const assistantCost = Effect.fn("TaskToolTest.assistantCost")(function* (sessionID: string) {
+  const sessions = yield* Session.Service
+  const msgs = yield* sessions.messages({ sessionID: SessionID.make(sessionID) })
+  return msgs.reduce((sum, m) => sum + (m.info.role === "assistant" ? m.info.cost : 0), 0)
+})
+
+describe("tool.task cost propagation", () => {
+  it.live("propagates subagent cost to parent assistant message", () =>
+    provideTmpdirInstance(() =>
+      Effect.gen(function* () {
+        const sessions = yield* Session.Service
+        const { chat, assistant } = yield* seed()
+        const tool = yield* TaskTool
+        const def = yield* tool.init()
+        const promptOps = stubOps({ sessions, childCost: 0.25 })
+
+        yield* def.execute(
+          {
+            description: "inspect bug",
+            prompt: "look into the cache key path",
+            subagent_type: "general",
+          },
+          {
+            sessionID: chat.id,
+            messageID: assistant.id,
+            agent: "build",
+            abort: new AbortController().signal,
+            extra: { promptOps },
+            messages: [],
+            metadata: () => Effect.void,
+            ask: () => Effect.void,
+          },
+        )
+        const parent = yield* Effect.sync(() => MessageV2.get({ sessionID: chat.id, messageID: assistant.id }))
+        expect(parent.info.role).toBe("assistant")
+        if (parent.info.role !== "assistant") return
+        expect(parent.info.cost).toBeCloseTo(0.25, 6)
+      }),
+    ),
+  )
+
+  it.live("propagates recursively through nested subagent costs", () =>
+    provideTmpdirInstance(() =>
+      Effect.gen(function* () {
+        const sessions = yield* Session.Service
+        const { chat, assistant } = yield* seed()
+        // Pre-create a child with its own assistant already bearing a grandchild cost.
+        const child = yield* sessions.create({ parentID: chat.id, title: "grandchild-accumulated" })
+        yield* sessions.updateMessage({
+          id: MessageID.ascending(),
+          role: "assistant",
+          parentID: assistant.id,
+          sessionID: child.id,
+          mode: "build",
+          agent: "general",
+          cost: 0.4,
+          path: { cwd: "/tmp", root: "/tmp" },
+          tokens: { input: 0, output: 0, reasoning: 0, cache: { read: 0, write: 0 } },
+          modelID: ref.modelID,
+          providerID: ref.providerID,
+          time: { created: Date.now() },
+        })
+
+        const tool = yield* TaskTool
+        const def = yield* tool.init()
+        // Resuming into the same child via task_id and the stub tacks on another 0.15.
+        const promptOps = stubOps({ sessions, childCost: 0.15 })
+
+        yield* def.execute(
+          {
+            description: "inspect bug",
+            prompt: "look into the cache key path",
+            subagent_type: "general",
+            task_id: child.id,
+          },
+          {
+            sessionID: chat.id,
+            messageID: assistant.id,
+            agent: "build",
+            abort: new AbortController().signal,
+            extra: { promptOps },
+            messages: [],
+            metadata: () => Effect.void,
+            ask: () => Effect.void,
+          },
+        )
+        const parent = yield* Effect.sync(() => MessageV2.get({ sessionID: chat.id, messageID: assistant.id }))
+        if (parent.info.role !== "assistant") return
+        // Only the delta since the start of this invocation propagates.
+        expect(parent.info.cost).toBeCloseTo(0.15, 6)
+        // Child session keeps the full cumulative total (0.4 pre-existing + 0.15 this run).
+        expect(yield* assistantCost(child.id)).toBeCloseTo(0.55, 6)
+      }),
+    ),
+  )
+
+  it.live("resumed task_id only propagates the delta", () =>
+    provideTmpdirInstance(() =>
+      Effect.gen(function* () {
+        const sessions = yield* Session.Service
+        const { chat, assistant } = yield* seed()
+        const child = yield* sessions.create({ parentID: chat.id, title: "resume target" })
+        yield* sessions.updateMessage({
+          id: MessageID.ascending(),
+          role: "assistant",
+          parentID: assistant.id,
+          sessionID: child.id,
+          mode: "build",
+          agent: "general",
+          cost: 0.1,
+          path: { cwd: "/tmp", root: "/tmp" },
+          tokens: { input: 0, output: 0, reasoning: 0, cache: { read: 0, write: 0 } },
+          modelID: ref.modelID,
+          providerID: ref.providerID,
+          time: { created: Date.now() },
+        })
+
+        const tool = yield* TaskTool
+        const def = yield* tool.init()
+        const promptOps = stubOps({ sessions, childCost: 0.05 })
+
+        yield* def.execute(
+          {
+            description: "inspect bug",
+            prompt: "continue investigation",
+            subagent_type: "general",
+            task_id: child.id,
+          },
+          {
+            sessionID: chat.id,
+            messageID: assistant.id,
+            agent: "build",
+            abort: new AbortController().signal,
+            extra: { promptOps },
+            messages: [],
+            metadata: () => Effect.void,
+            ask: () => Effect.void,
+          },
+        )
+        const parent = yield* Effect.sync(() => MessageV2.get({ sessionID: chat.id, messageID: assistant.id }))
+        if (parent.info.role !== "assistant") return
+        // Delta-only: only the 0.05 from this run, not 0.15 including the pre-existing 0.10.
+        expect(parent.info.cost).toBeCloseTo(0.05, 6)
+      }),
+    ),
+  )
+
+  it.live("propagates partial cost on abort", () =>
+    provideTmpdirInstance(() =>
+      Effect.gen(function* () {
+        const sessions = yield* Session.Service
+        const { chat, assistant } = yield* seed()
+        const tool = yield* TaskTool
+        const def = yield* tool.init()
+        const abort = new AbortController()
+        // Stub that persists a partial cost, then aborts — mimics interrupted run after tokens billed.
+        const ops: TaskPromptOps = {
+          cancel() {},
+          resolvePromptParts: (template) => Effect.succeed([{ type: "text" as const, text: template }]),
+          prompt: (input) =>
+            Effect.gen(function* () {
+              const info: MessageV2.Assistant = {
+                id: MessageID.ascending(),
+                role: "assistant",
+                parentID: input.messageID ?? MessageID.ascending(),
+                sessionID: input.sessionID,
+                mode: "general",
+                agent: "general",
+                cost: 0.07,
+                path: { cwd: "/tmp", root: "/tmp" },
+                tokens: { input: 0, output: 0, reasoning: 0, cache: { read: 0, write: 0 } },
+                modelID: ref.modelID,
+                providerID: ref.providerID,
+                time: { created: Date.now() },
+              }
+              yield* sessions.updateMessage(info)
+              abort.abort()
+              return yield* Effect.interrupt
+            }),
+        }
+
+        yield* def
+          .execute(
+            {
+              description: "partial",
+              prompt: "will abort",
+              subagent_type: "general",
+            },
+            {
+              sessionID: chat.id,
+              messageID: assistant.id,
+              agent: "build",
+              abort: abort.signal,
+              extra: { promptOps: ops },
+              messages: [],
+              metadata: () => Effect.void,
+              ask: () => Effect.void,
+            },
+          )
+          .pipe(Effect.exit)
+
+        const parent = yield* Effect.sync(() => MessageV2.get({ sessionID: chat.id, messageID: assistant.id }))
+        if (parent.info.role !== "assistant") return
+        expect(parent.info.cost).toBeCloseTo(0.07, 6)
+      }),
+    ),
+  )
+})
+// kilocode_change end