refactor(effect): build task tool from agent services (#21017)
This commit is contained in:
@@ -1,5 +1,5 @@
|
||||
import { NodeFileSystem } from "@effect/platform-node"
|
||||
import { expect, spyOn } from "bun:test"
|
||||
import { expect } from "bun:test"
|
||||
import { Cause, Effect, Exit, Fiber, Layer } from "effect"
|
||||
import path from "path"
|
||||
import z from "zod"
|
||||
@@ -29,7 +29,6 @@ import { MessageID, PartID, SessionID } from "../../src/session/schema"
|
||||
import { SessionStatus } from "../../src/session/status"
|
||||
import { Shell } from "../../src/shell/shell"
|
||||
import { Snapshot } from "../../src/snapshot"
|
||||
import { TaskTool } from "../../src/tool/task"
|
||||
import { ToolRegistry } from "../../src/tool/registry"
|
||||
import { Truncate } from "../../src/tool/truncate"
|
||||
import { Log } from "../../src/util/log"
|
||||
@@ -627,11 +626,13 @@ it.live(
|
||||
"cancel finalizes subtask tool state",
|
||||
() =>
|
||||
provideTmpdirInstance(
|
||||
(dir) =>
|
||||
() =>
|
||||
Effect.gen(function* () {
|
||||
const ready = defer<void>()
|
||||
const aborted = defer<void>()
|
||||
const init = spyOn(TaskTool, "init").mockImplementation(async () => ({
|
||||
const registry = yield* ToolRegistry.Service
|
||||
const init = registry.named.task.init
|
||||
registry.named.task.init = async () => ({
|
||||
description: "task",
|
||||
parameters: z.object({
|
||||
description: z.string(),
|
||||
@@ -653,8 +654,8 @@ it.live(
|
||||
output: "",
|
||||
}
|
||||
},
|
||||
}))
|
||||
yield* Effect.addFinalizer(() => Effect.sync(() => init.mockRestore()))
|
||||
})
|
||||
yield* Effect.addFinalizer(() => Effect.sync(() => void (registry.named.task.init = init)))
|
||||
|
||||
const { prompt, chat } = yield* boot()
|
||||
const msg = yield* user(chat.id, "hello")
|
||||
|
||||
@@ -1,50 +1,412 @@
|
||||
import { Effect } from "effect"
|
||||
import { afterEach, describe, expect, test } from "bun:test"
|
||||
import { afterEach, describe, expect } from "bun:test"
|
||||
import { Effect, Layer } from "effect"
|
||||
import { Agent } from "../../src/agent/agent"
|
||||
import { Config } from "../../src/config/config"
|
||||
import * as CrossSpawnSpawner from "../../src/effect/cross-spawn-spawner"
|
||||
import { Instance } from "../../src/project/instance"
|
||||
import { TaskDescription } from "../../src/tool/task"
|
||||
import { tmpdir } from "../fixture/fixture"
|
||||
import { Session } from "../../src/session"
|
||||
import { MessageV2 } from "../../src/session/message-v2"
|
||||
import { SessionPrompt } from "../../src/session/prompt"
|
||||
import { MessageID, PartID } from "../../src/session/schema"
|
||||
import { ModelID, ProviderID } from "../../src/provider/schema"
|
||||
import { TaskDescription, TaskTool } from "../../src/tool/task"
|
||||
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),
|
||||
)
|
||||
|
||||
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 reply(input: Parameters<typeof SessionPrompt.prompt>[0], 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", () => {
|
||||
test("description sorts subagents by name and is stable across calls", async () => {
|
||||
await using tmp = await tmpdir({
|
||||
config: {
|
||||
agent: {
|
||||
zebra: {
|
||||
description: "Zebra agent",
|
||||
mode: "subagent",
|
||||
},
|
||||
alpha: {
|
||||
description: "Alpha agent",
|
||||
mode: "subagent",
|
||||
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 first = yield* TaskDescription(build)
|
||||
const second = yield* TaskDescription(build)
|
||||
|
||||
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",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
})
|
||||
),
|
||||
)
|
||||
|
||||
await Instance.provide({
|
||||
directory: tmp.path,
|
||||
fn: async () => {
|
||||
const agent = { name: "build", mode: "primary" as const, permission: [], options: {} }
|
||||
const first = await Effect.runPromise(TaskDescription(agent))
|
||||
const second = await Effect.runPromise(TaskDescription(agent))
|
||||
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 description = yield* TaskDescription(build)
|
||||
|
||||
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)
|
||||
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* Effect.promise(() => tool.init())
|
||||
const resolve = SessionPrompt.resolvePromptParts
|
||||
const prompt = SessionPrompt.prompt
|
||||
let seen: Parameters<typeof SessionPrompt.prompt>[0] | undefined
|
||||
|
||||
SessionPrompt.resolvePromptParts = async (template) => [{ type: "text", text: template }]
|
||||
SessionPrompt.prompt = async (input) => {
|
||||
seen = input
|
||||
return reply(input, "resumed")
|
||||
}
|
||||
yield* Effect.addFinalizer(() =>
|
||||
Effect.sync(() => {
|
||||
SessionPrompt.resolvePromptParts = resolve
|
||||
SessionPrompt.prompt = prompt
|
||||
}),
|
||||
)
|
||||
|
||||
const result = yield* Effect.promise(() =>
|
||||
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,
|
||||
messages: [],
|
||||
metadata() {},
|
||||
ask: async () => {},
|
||||
},
|
||||
),
|
||||
)
|
||||
|
||||
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* Effect.promise(() => tool.init())
|
||||
const resolve = SessionPrompt.resolvePromptParts
|
||||
const prompt = SessionPrompt.prompt
|
||||
const calls: unknown[] = []
|
||||
|
||||
SessionPrompt.resolvePromptParts = async (template) => [{ type: "text", text: template }]
|
||||
SessionPrompt.prompt = async (input) => reply(input, "done")
|
||||
yield* Effect.addFinalizer(() =>
|
||||
Effect.sync(() => {
|
||||
SessionPrompt.resolvePromptParts = resolve
|
||||
SessionPrompt.prompt = prompt
|
||||
}),
|
||||
)
|
||||
|
||||
const exec = (extra?: { bypassAgentCheck?: boolean }) =>
|
||||
Effect.promise(() =>
|
||||
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,
|
||||
messages: [],
|
||||
metadata() {},
|
||||
ask: async (input) => {
|
||||
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* Effect.promise(() => tool.init())
|
||||
const resolve = SessionPrompt.resolvePromptParts
|
||||
const prompt = SessionPrompt.prompt
|
||||
let seen: Parameters<typeof SessionPrompt.prompt>[0] | undefined
|
||||
|
||||
SessionPrompt.resolvePromptParts = async (template) => [{ type: "text", text: template }]
|
||||
SessionPrompt.prompt = async (input) => {
|
||||
seen = input
|
||||
return reply(input, "created")
|
||||
}
|
||||
yield* Effect.addFinalizer(() =>
|
||||
Effect.sync(() => {
|
||||
SessionPrompt.resolvePromptParts = resolve
|
||||
SessionPrompt.prompt = prompt
|
||||
}),
|
||||
)
|
||||
|
||||
const result = yield* Effect.promise(() =>
|
||||
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,
|
||||
messages: [],
|
||||
metadata() {},
|
||||
ask: async () => {},
|
||||
},
|
||||
),
|
||||
)
|
||||
|
||||
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* Effect.promise(() => tool.init())
|
||||
const resolve = SessionPrompt.resolvePromptParts
|
||||
const prompt = SessionPrompt.prompt
|
||||
let seen: Parameters<typeof SessionPrompt.prompt>[0] | undefined
|
||||
|
||||
SessionPrompt.resolvePromptParts = async (template) => [{ type: "text", text: template }]
|
||||
SessionPrompt.prompt = async (input) => {
|
||||
seen = input
|
||||
return reply(input, "done")
|
||||
}
|
||||
yield* Effect.addFinalizer(() =>
|
||||
Effect.sync(() => {
|
||||
SessionPrompt.resolvePromptParts = resolve
|
||||
SessionPrompt.prompt = prompt
|
||||
}),
|
||||
)
|
||||
|
||||
const result = yield* Effect.promise(() =>
|
||||
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,
|
||||
messages: [],
|
||||
metadata() {},
|
||||
ask: async () => {},
|
||||
},
|
||||
),
|
||||
)
|
||||
|
||||
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"],
|
||||
},
|
||||
},
|
||||
},
|
||||
),
|
||||
)
|
||||
})
|
||||
|
||||
Reference in New Issue
Block a user