refactor(effect): build task tool from agent services

pull/21017/head
Kit Langton 2026-04-02 22:42:11 -04:00
parent baff53b759
commit 78f7258c6d
4 changed files with 198 additions and 172 deletions

View File

@ -33,6 +33,7 @@ import { Effect, Layer, ServiceMap } from "effect"
import { InstanceState } from "@/effect/instance-state" import { InstanceState } from "@/effect/instance-state"
import { makeRuntime } from "@/effect/run-service" import { makeRuntime } from "@/effect/run-service"
import { Env } from "../env" import { Env } from "../env"
import { Agent as AgentSvc } from "../agent/agent"
import { Question } from "../question" import { Question } from "../question"
import { Todo } from "../session/todo" import { Todo } from "../session/todo"
import { LSP } from "../lsp" import { LSP } from "../lsp"
@ -68,6 +69,7 @@ export namespace ToolRegistry {
| Plugin.Service | Plugin.Service
| Question.Service | Question.Service
| Todo.Service | Todo.Service
| AgentSvc.Service
| LSP.Service | LSP.Service
| FileTime.Service | FileTime.Service
| Instruction.Service | Instruction.Service
@ -237,6 +239,7 @@ export namespace ToolRegistry {
layer.pipe( layer.pipe(
Layer.provide(Config.defaultLayer), Layer.provide(Config.defaultLayer),
Layer.provide(Plugin.defaultLayer), Layer.provide(Plugin.defaultLayer),
Layer.provide(AgentSvc.defaultLayer),
Layer.provide(Question.defaultLayer), Layer.provide(Question.defaultLayer),
Layer.provide(Todo.defaultLayer), Layer.provide(Todo.defaultLayer),
Layer.provide(LSP.defaultLayer), Layer.provide(LSP.defaultLayer),

View File

@ -1,6 +1,7 @@
import { Tool } from "./tool" import { Tool } from "./tool"
import DESCRIPTION from "./task.txt" import DESCRIPTION from "./task.txt"
import z from "zod" import z from "zod"
import { Effect } from "effect"
import { Session } from "../session" import { Session } from "../session"
import { SessionID, MessageID } from "../session/schema" import { SessionID, MessageID } from "../session/schema"
import { MessageV2 } from "../session/message-v2" import { MessageV2 } from "../session/message-v2"
@ -25,10 +26,18 @@ const parameters = z.object({
command: z.string().describe("The command that triggered this task").optional(), command: z.string().describe("The command that triggered this task").optional(),
}) })
export const TaskTool = Tool.define("task", async (ctx) => { export const TaskTool = Tool.defineEffect(
const agents = await Agent.list().then((x) => x.filter((a) => a.mode !== "primary")) "task",
Effect.gen(function* () {
const agent = yield* Agent.Service
const config = yield* Config.Service
return async (ctx) => {
const agents = await agent.list().pipe(
Effect.map((x) => x.filter((a) => a.mode !== "primary")),
Effect.runPromise,
)
// Filter agents by permissions if agent provided
const caller = ctx?.agent const caller = ctx?.agent
const accessibleAgents = caller const accessibleAgents = caller
? agents.filter((a) => Permission.evaluate("task", a.name, caller.permission).action !== "deny") ? agents.filter((a) => Permission.evaluate("task", a.name, caller.permission).action !== "deny")
@ -41,11 +50,12 @@ export const TaskTool = Tool.define("task", async (ctx) => {
.map((a) => `- ${a.name}: ${a.description ?? "This subagent should only be called manually by the user."}`) .map((a) => `- ${a.name}: ${a.description ?? "This subagent should only be called manually by the user."}`)
.join("\n"), .join("\n"),
) )
return { return {
description, description,
parameters, parameters,
async execute(params: z.infer<typeof parameters>, ctx) { async execute(params: z.infer<typeof parameters>, ctx) {
const config = await Config.get() const cfg = await config.get().pipe(Effect.runPromise)
// Skip permission check when user explicitly invoked via @ or command subtask // Skip permission check when user explicitly invoked via @ or command subtask
if (!ctx.extra?.bypassAgentCheck) { if (!ctx.extra?.bypassAgentCheck) {
@ -60,11 +70,14 @@ export const TaskTool = Tool.define("task", async (ctx) => {
}) })
} }
const agent = await Agent.get(params.subagent_type) const next = await agent
if (!agent) throw new Error(`Unknown agent type: ${params.subagent_type} is not a valid agent type`) .get(params.subagent_type)
.pipe(Effect.runPromise)
.catch(() => undefined)
if (!next) throw new Error(`Unknown agent type: ${params.subagent_type} is not a valid agent type`)
const hasTaskPermission = agent.permission.some((rule) => rule.permission === "task") const hasTaskPermission = next.permission.some((rule) => rule.permission === "task")
const hasTodoWritePermission = agent.permission.some((rule) => rule.permission === "todowrite") const hasTodoWritePermission = next.permission.some((rule) => rule.permission === "todowrite")
const session = await iife(async () => { const session = await iife(async () => {
if (params.task_id) { if (params.task_id) {
@ -74,7 +87,7 @@ export const TaskTool = Tool.define("task", async (ctx) => {
return await Session.create({ return await Session.create({
parentID: ctx.sessionID, parentID: ctx.sessionID,
title: params.description + ` (@${agent.name} subagent)`, title: params.description + ` (@${next.name} subagent)`,
permission: [ permission: [
...(hasTodoWritePermission ...(hasTodoWritePermission
? [] ? []
@ -94,7 +107,7 @@ export const TaskTool = Tool.define("task", async (ctx) => {
action: "deny" as const, action: "deny" as const,
}, },
]), ]),
...(config.experimental?.primary_tools?.map((t) => ({ ...(cfg.experimental?.primary_tools?.map((t) => ({
pattern: "*", pattern: "*",
action: "allow" as const, action: "allow" as const,
permission: t, permission: t,
@ -105,7 +118,7 @@ export const TaskTool = Tool.define("task", async (ctx) => {
const msg = await MessageV2.get({ sessionID: ctx.sessionID, messageID: ctx.messageID }) const msg = await MessageV2.get({ sessionID: ctx.sessionID, messageID: ctx.messageID })
if (msg.info.role !== "assistant") throw new Error("Not an assistant message") if (msg.info.role !== "assistant") throw new Error("Not an assistant message")
const model = agent.model ?? { const model = next.model ?? {
modelID: msg.info.modelID, modelID: msg.info.modelID,
providerID: msg.info.providerID, providerID: msg.info.providerID,
} }
@ -134,11 +147,11 @@ export const TaskTool = Tool.define("task", async (ctx) => {
modelID: model.modelID, modelID: model.modelID,
providerID: model.providerID, providerID: model.providerID,
}, },
agent: agent.name, agent: next.name,
tools: { tools: {
...(hasTodoWritePermission ? {} : { todowrite: false }), ...(hasTodoWritePermission ? {} : { todowrite: false }),
...(hasTaskPermission ? {} : { task: false }), ...(hasTaskPermission ? {} : { task: false }),
...Object.fromEntries((config.experimental?.primary_tools ?? []).map((t) => [t, false])), ...Object.fromEntries((cfg.experimental?.primary_tools ?? []).map((t) => [t, false])),
}, },
parts: promptParts, parts: promptParts,
}) })
@ -163,4 +176,6 @@ export const TaskTool = Tool.define("task", async (ctx) => {
} }
}, },
} }
}) }
}),
)

View File

@ -29,7 +29,6 @@ import { MessageID, PartID, SessionID } from "../../src/session/schema"
import { SessionStatus } from "../../src/session/status" import { SessionStatus } from "../../src/session/status"
import { Shell } from "../../src/shell/shell" import { Shell } from "../../src/shell/shell"
import { Snapshot } from "../../src/snapshot" import { Snapshot } from "../../src/snapshot"
import { TaskTool } from "../../src/tool/task"
import { ToolRegistry } from "../../src/tool/registry" import { ToolRegistry } from "../../src/tool/registry"
import { Truncate } from "../../src/tool/truncate" import { Truncate } from "../../src/tool/truncate"
import { Log } from "../../src/util/log" import { Log } from "../../src/util/log"
@ -631,7 +630,9 @@ it.live(
Effect.gen(function* () { Effect.gen(function* () {
const ready = defer<void>() const ready = defer<void>()
const aborted = 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", description: "task",
parameters: z.object({ parameters: z.object({
description: z.string(), description: z.string(),
@ -653,8 +654,8 @@ it.live(
output: "", 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 { prompt, chat } = yield* boot()
const msg = yield* user(chat.id, "hello") const msg = yield* user(chat.id, "hello")

View File

@ -1,36 +1,29 @@
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 { 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 { Instance } from "../../src/project/instance"
import { TaskTool } from "../../src/tool/task" import { TaskTool } from "../../src/tool/task"
import { tmpdir } from "../fixture/fixture" import { provideTmpdirInstance } from "../fixture/fixture"
import { testEffect } from "../lib/effect"
afterEach(async () => { afterEach(async () => {
await Instance.disposeAll() await Instance.disposeAll()
}) })
describe("tool.task", () => { const it = testEffect(Layer.mergeAll(Agent.defaultLayer, Config.defaultLayer, CrossSpawnSpawner.defaultLayer))
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",
},
},
},
})
await Instance.provide({ describe("tool.task", () => {
directory: tmp.path, it.live("description sorts subagents by name and is stable across calls", () =>
fn: async () => { provideTmpdirInstance(
const build = await Agent.get("build") () =>
const first = await TaskTool.init({ agent: build }) Effect.gen(function* () {
const second = await TaskTool.init({ agent: build }) const agent = yield* Agent.Service
const build = yield* agent.get("build")
const tool = yield* TaskTool
const first = yield* Effect.promise(() => tool.init({ agent: build }))
const second = yield* Effect.promise(() => tool.init({ agent: build }))
expect(first.description).toBe(second.description) expect(first.description).toBe(second.description)
@ -43,7 +36,21 @@ describe("tool.task", () => {
expect(explore).toBeGreaterThan(alpha) expect(explore).toBeGreaterThan(alpha)
expect(general).toBeGreaterThan(explore) expect(general).toBeGreaterThan(explore)
expect(zebra).toBeGreaterThan(general) expect(zebra).toBeGreaterThan(general)
}),
{
config: {
agent: {
zebra: {
description: "Zebra agent",
mode: "subagent",
}, },
}) alpha: {
}) description: "Alpha agent",
mode: "subagent",
},
},
},
},
),
)
}) })