| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387 |
- 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"
- import type { SessionPrompt } from "../../src/session"
- import { MessageID, PartID } from "../../src/session/schema"
- import { ModelID, ProviderID } from "../../src/provider/schema"
- import { TaskTool, type TaskPromptOps } from "../../src/tool/task"
- import { Truncate } from "../../src/tool/truncate"
- import { ToolRegistry } from "../../src/tool/registry"
- 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 }
- })
- function stubOps(opts?: { onPrompt?: (input: SessionPrompt.PromptInput) => void; text?: string }): TaskPromptOps {
- return {
- cancel() {},
- resolvePromptParts: (template) => Effect.succeed([{ type: "text" as const, text: template }]),
- prompt: (input) =>
- Effect.sync(() => {
- opts?.onPrompt?.(input)
- return reply(input, opts?.text ?? "done")
- }),
- }
- }
- 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)
- expect(child.permission).toEqual([
- {
- permission: "todowrite",
- pattern: "*",
- action: "deny",
- },
- {
- permission: "bash",
- pattern: "*",
- action: "allow",
- },
- {
- permission: "read",
- pattern: "*",
- action: "allow",
- },
- ])
- expect(seen?.tools).toEqual({
- todowrite: false,
- bash: false,
- read: false,
- })
- }),
- {
- config: {
- agent: {
- reviewer: {
- mode: "subagent",
- permission: {
- task: "allow",
- },
- },
- },
- experimental: {
- primary_tools: ["bash", "read"],
- },
- },
- },
- ),
- )
- })
|