Alex Alecu
aaaa191b76
When a session spawned subagents via the task tool, the displayed cost reflected only the parent's own LLM steps. Propagate each subagent's total up to the invoking assistant message so every UI sums the full tree.
620 lines
20 KiB
TypeScript
620 lines
20 KiB
TypeScript
import { afterEach, describe, expect } from "bun:test"
|
|
import { Effect, Layer } from "effect"
|
|
import { Agent } from "../../src/agent/agent"
|
|
import { Config } from "../../src/config"
|
|
import * as CrossSpawnSpawner from "../../src/effect/cross-spawn-spawner"
|
|
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, 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"
|
|
import { ToolRegistry } from "../../src/tool"
|
|
import { provideTmpdirInstance } from "../fixture/fixture"
|
|
import { testEffect } from "../lib/effect"
|
|
|
|
afterEach(async () => {
|
|
await Instance.disposeAll()
|
|
})
|
|
|
|
const ref = {
|
|
providerID: ProviderID.make("test"),
|
|
modelID: ModelID.make("test-model"),
|
|
}
|
|
|
|
const it = testEffect(
|
|
Layer.mergeAll(
|
|
Agent.defaultLayer,
|
|
Config.defaultLayer,
|
|
CrossSpawnSpawner.defaultLayer,
|
|
Session.defaultLayer,
|
|
Truncate.defaultLayer,
|
|
ToolRegistry.defaultLayer,
|
|
),
|
|
)
|
|
|
|
const seed = Effect.fn("TaskToolTest.seed")(function* (title = "Pinned") {
|
|
const session = yield* Session.Service
|
|
const chat = yield* session.create({ title })
|
|
const user = yield* session.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* session.updateMessage(assistant)
|
|
return { chat, assistant }
|
|
})
|
|
|
|
// 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.gen(function* () {
|
|
opts?.onPrompt?.(input)
|
|
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()
|
|
return {
|
|
info: {
|
|
id,
|
|
role: "assistant",
|
|
parentID: input.messageID ?? MessageID.ascending(),
|
|
sessionID: input.sessionID,
|
|
mode: input.agent ?? "general",
|
|
agent: input.agent ?? "general",
|
|
cost: 0,
|
|
path: { cwd: "/tmp", root: "/tmp" },
|
|
tokens: { input: 0, output: 0, reasoning: 0, cache: { read: 0, write: 0 } },
|
|
modelID: input.model?.modelID ?? ref.modelID,
|
|
providerID: input.model?.providerID ?? ref.providerID,
|
|
time: { created: Date.now() },
|
|
finish: "stop",
|
|
},
|
|
parts: [
|
|
{
|
|
id: PartID.ascending(),
|
|
messageID: id,
|
|
sessionID: input.sessionID,
|
|
type: "text",
|
|
text,
|
|
},
|
|
],
|
|
}
|
|
}
|
|
|
|
describe("tool.task", () => {
|
|
it.live("description sorts subagents by name and is stable across calls", () =>
|
|
provideTmpdirInstance(
|
|
() =>
|
|
Effect.gen(function* () {
|
|
const agent = yield* Agent.Service
|
|
const build = yield* agent.get("build")
|
|
const registry = yield* ToolRegistry.Service
|
|
const get = Effect.fnUntraced(function* () {
|
|
const tools = yield* registry.tools({ ...ref, agent: build })
|
|
return tools.find((tool) => tool.id === TaskTool.id)?.description ?? ""
|
|
})
|
|
const first = yield* get()
|
|
const second = yield* get()
|
|
|
|
expect(first).toBe(second)
|
|
|
|
const alpha = first.indexOf("- alpha: Alpha agent")
|
|
const explore = first.indexOf("- explore:")
|
|
const general = first.indexOf("- general:")
|
|
const zebra = first.indexOf("- zebra: Zebra agent")
|
|
|
|
expect(alpha).toBeGreaterThan(-1)
|
|
expect(explore).toBeGreaterThan(alpha)
|
|
expect(general).toBeGreaterThan(explore)
|
|
expect(zebra).toBeGreaterThan(general)
|
|
}),
|
|
{
|
|
config: {
|
|
agent: {
|
|
zebra: {
|
|
description: "Zebra agent",
|
|
mode: "subagent",
|
|
},
|
|
alpha: {
|
|
description: "Alpha agent",
|
|
mode: "subagent",
|
|
},
|
|
},
|
|
},
|
|
},
|
|
),
|
|
)
|
|
|
|
it.live("description hides denied subagents for the caller", () =>
|
|
provideTmpdirInstance(
|
|
() =>
|
|
Effect.gen(function* () {
|
|
const agent = yield* Agent.Service
|
|
const build = yield* agent.get("build")
|
|
const registry = yield* ToolRegistry.Service
|
|
const description =
|
|
(yield* registry.tools({ ...ref, agent: build })).find((tool) => tool.id === TaskTool.id)?.description ?? ""
|
|
|
|
expect(description).toContain("- alpha: Alpha agent")
|
|
expect(description).not.toContain("- zebra: Zebra agent")
|
|
}),
|
|
{
|
|
config: {
|
|
permission: {
|
|
task: {
|
|
"*": "allow",
|
|
zebra: "deny",
|
|
},
|
|
},
|
|
agent: {
|
|
zebra: {
|
|
description: "Zebra agent",
|
|
mode: "subagent",
|
|
},
|
|
alpha: {
|
|
description: "Alpha agent",
|
|
mode: "subagent",
|
|
},
|
|
},
|
|
},
|
|
},
|
|
),
|
|
)
|
|
|
|
it.live("execute resumes an existing task session from task_id", () =>
|
|
provideTmpdirInstance(() =>
|
|
Effect.gen(function* () {
|
|
const sessions = yield* Session.Service
|
|
const { chat, assistant } = yield* seed()
|
|
const child = yield* sessions.create({ parentID: chat.id, title: "Existing child" })
|
|
const tool = yield* TaskTool
|
|
const def = yield* tool.init()
|
|
let seen: SessionPrompt.PromptInput | undefined
|
|
const promptOps = stubOps({ text: "resumed", onPrompt: (input) => (seen = input) })
|
|
|
|
const result = 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 kids = yield* sessions.children(chat.id)
|
|
expect(kids).toHaveLength(1)
|
|
expect(kids[0]?.id).toBe(child.id)
|
|
expect(result.metadata.sessionId).toBe(child.id)
|
|
expect(result.output).toContain(`task_id: ${child.id}`)
|
|
expect(seen?.sessionID).toBe(child.id)
|
|
}),
|
|
),
|
|
)
|
|
|
|
it.live("execute asks by default and skips checks when bypassed", () =>
|
|
provideTmpdirInstance(() =>
|
|
Effect.gen(function* () {
|
|
const { chat, assistant } = yield* seed()
|
|
const tool = yield* TaskTool
|
|
const def = yield* tool.init()
|
|
const calls: unknown[] = []
|
|
const promptOps = stubOps()
|
|
|
|
const exec = (extra?: Record<string, any>) =>
|
|
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, ...extra },
|
|
messages: [],
|
|
metadata: () => Effect.void,
|
|
ask: (input) =>
|
|
Effect.sync(() => {
|
|
calls.push(input)
|
|
}),
|
|
},
|
|
)
|
|
|
|
yield* exec()
|
|
yield* exec({ bypassAgentCheck: true })
|
|
|
|
expect(calls).toHaveLength(1)
|
|
expect(calls[0]).toEqual({
|
|
permission: "task",
|
|
patterns: ["general"],
|
|
always: ["*"],
|
|
metadata: {
|
|
description: "inspect bug",
|
|
subagent_type: "general",
|
|
},
|
|
})
|
|
}),
|
|
),
|
|
)
|
|
|
|
it.live("execute creates a child when task_id does not exist", () =>
|
|
provideTmpdirInstance(() =>
|
|
Effect.gen(function* () {
|
|
const sessions = yield* Session.Service
|
|
const { chat, assistant } = yield* seed()
|
|
const tool = yield* TaskTool
|
|
const def = yield* tool.init()
|
|
let seen: SessionPrompt.PromptInput | undefined
|
|
const promptOps = stubOps({ text: "created", onPrompt: (input) => (seen = input) })
|
|
|
|
const result = yield* def.execute(
|
|
{
|
|
description: "inspect bug",
|
|
prompt: "look into the cache key path",
|
|
subagent_type: "general",
|
|
task_id: "ses_missing",
|
|
},
|
|
{
|
|
sessionID: chat.id,
|
|
messageID: assistant.id,
|
|
agent: "build",
|
|
abort: new AbortController().signal,
|
|
extra: { promptOps },
|
|
messages: [],
|
|
metadata: () => Effect.void,
|
|
ask: () => Effect.void,
|
|
},
|
|
)
|
|
|
|
const kids = yield* sessions.children(chat.id)
|
|
expect(kids).toHaveLength(1)
|
|
expect(kids[0]?.id).toBe(result.metadata.sessionId)
|
|
expect(result.metadata.sessionId).not.toBe("ses_missing")
|
|
expect(result.output).toContain(`task_id: ${result.metadata.sessionId}`)
|
|
expect(seen?.sessionID).toBe(result.metadata.sessionId)
|
|
}),
|
|
),
|
|
)
|
|
|
|
it.live("execute shapes child permissions for task, todowrite, and primary tools", () =>
|
|
provideTmpdirInstance(
|
|
() =>
|
|
Effect.gen(function* () {
|
|
const sessions = yield* Session.Service
|
|
const { chat, assistant } = yield* seed()
|
|
const tool = yield* TaskTool
|
|
const def = yield* tool.init()
|
|
let seen: SessionPrompt.PromptInput | undefined
|
|
const promptOps = stubOps({ onPrompt: (input) => (seen = input) })
|
|
|
|
const result = yield* def.execute(
|
|
{
|
|
description: "inspect bug",
|
|
prompt: "look into the cache key path",
|
|
subagent_type: "reviewer",
|
|
},
|
|
{
|
|
sessionID: chat.id,
|
|
messageID: assistant.id,
|
|
agent: "build",
|
|
abort: new AbortController().signal,
|
|
extra: { promptOps },
|
|
messages: [],
|
|
metadata: () => Effect.void,
|
|
ask: () => Effect.void,
|
|
},
|
|
)
|
|
|
|
const child = yield* sessions.get(result.metadata.sessionId)
|
|
expect(child.parentID).toBe(chat.id)
|
|
// kilocode_change start — use arrayContaining: Kilo appends inherited caller restrictions
|
|
expect(child.permission).toEqual(
|
|
expect.arrayContaining([
|
|
{
|
|
permission: "todowrite",
|
|
pattern: "*",
|
|
action: "deny",
|
|
},
|
|
{
|
|
permission: "bash",
|
|
pattern: "*",
|
|
action: "allow",
|
|
},
|
|
{
|
|
permission: "read",
|
|
pattern: "*",
|
|
action: "allow",
|
|
},
|
|
{
|
|
permission: "task",
|
|
pattern: "*",
|
|
action: "deny",
|
|
},
|
|
]),
|
|
)
|
|
// kilocode_change end
|
|
expect(seen?.tools).toEqual({
|
|
todowrite: false,
|
|
bash: false,
|
|
read: false,
|
|
})
|
|
}),
|
|
{
|
|
config: {
|
|
agent: {
|
|
reviewer: {
|
|
mode: "subagent",
|
|
permission: {
|
|
task: "allow",
|
|
},
|
|
},
|
|
},
|
|
experimental: {
|
|
primary_tools: ["bash", "read"],
|
|
openTelemetry: true, // kilocode_change
|
|
},
|
|
},
|
|
},
|
|
),
|
|
)
|
|
})
|
|
|
|
// 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
|