299 lines
9.3 KiB
TypeScript
299 lines
9.3 KiB
TypeScript
import { Tool } from "./tool"
|
|
import DESCRIPTION from "./task.txt"
|
|
import z from "zod"
|
|
import { Session } from "../session"
|
|
import { MessageV2 } from "../session/message-v2"
|
|
import { Identifier } from "../id/id"
|
|
import { Agent } from "../agent/agent"
|
|
import { SessionPrompt } from "../session/prompt"
|
|
import { SessionStatus } from "../session/status"
|
|
import { iife } from "@/util/iife"
|
|
import { defer } from "@/util/defer"
|
|
import { Config } from "../config/config"
|
|
import { PermissionNext } from "@/permission/next"
|
|
|
|
const parameters = z.object({
|
|
description: z.string().describe("A short (3-5 words) description of the task"),
|
|
prompt: z.string().describe("The task for the agent to perform"),
|
|
subagent_type: z.string().describe("The type of specialized agent to use for this task"),
|
|
task_id: z
|
|
.string()
|
|
.describe(
|
|
"This should only be set if you mean to resume a previous task (you can pass a prior task_id and the task will continue the same subagent session as before instead of creating a fresh one)",
|
|
)
|
|
.optional(),
|
|
command: z.string().describe("The command that triggered this task").optional(),
|
|
background: z
|
|
.boolean()
|
|
.optional()
|
|
.describe("When true, launch the subagent in the background and return immediately"),
|
|
})
|
|
|
|
function output(sessionID: string, text: string) {
|
|
return [
|
|
`task_id: ${sessionID} (for resuming to continue this task if needed)`,
|
|
"",
|
|
"<task_result>",
|
|
text,
|
|
"</task_result>",
|
|
].join("\n")
|
|
}
|
|
|
|
function backgroundOutput(sessionID: string) {
|
|
return [
|
|
`task_id: ${sessionID} (for polling this task with task_status)`,
|
|
"state: running",
|
|
"",
|
|
"<task_result>",
|
|
"Background task started. Continue your current work and call task_status when you need the result.",
|
|
"</task_result>",
|
|
].join("\n")
|
|
}
|
|
|
|
function backgroundMessage(input: {
|
|
sessionID: string
|
|
description: string
|
|
state: "completed" | "error"
|
|
text: string
|
|
}) {
|
|
const tag = input.state === "completed" ? "task_result" : "task_error"
|
|
const title =
|
|
input.state === "completed"
|
|
? `Background task completed: ${input.description}`
|
|
: `Background task failed: ${input.description}`
|
|
return [title, `task_id: ${input.sessionID}`, `state: ${input.state}`, `<${tag}>`, input.text, `</${tag}>`].join("\n")
|
|
}
|
|
|
|
function errorText(error: unknown) {
|
|
if (error instanceof Error) return error.message
|
|
return String(error)
|
|
}
|
|
|
|
function resultTaskID(input: unknown) {
|
|
if (!input || typeof input !== "object") return
|
|
const taskID = Reflect.get(input, "task_id")
|
|
if (typeof taskID === "string") return taskID
|
|
}
|
|
|
|
function polled(input: { message: MessageV2.WithParts; taskID: string }) {
|
|
if (input.message.info.role !== "assistant") return false
|
|
return input.message.parts.some((part) => {
|
|
if (part.type !== "tool") return false
|
|
if (part.tool !== "task_status") return false
|
|
if (part.state.status !== "completed") return false
|
|
return resultTaskID(part.state.input) === input.taskID
|
|
})
|
|
}
|
|
|
|
async function latestUser(sessionID: string) {
|
|
const [message] = await Session.messages({
|
|
sessionID,
|
|
limit: 1,
|
|
})
|
|
if (!message) return
|
|
if (message.info.role !== "user") return
|
|
return message.info.id
|
|
}
|
|
|
|
async function continueParent(input: { parentID: string; userID: string; taskID: string }) {
|
|
const message =
|
|
SessionStatus.get(input.parentID).type === "idle"
|
|
? undefined
|
|
: await SessionPrompt.loop({
|
|
sessionID: input.parentID,
|
|
}).catch(() => undefined)
|
|
if (message && polled({ message, taskID: input.taskID })) return
|
|
if (SessionStatus.get(input.parentID).type !== "idle") return
|
|
if ((await latestUser(input.parentID)) !== input.userID) return
|
|
await SessionPrompt.loop({
|
|
sessionID: input.parentID,
|
|
})
|
|
}
|
|
|
|
export const TaskTool = Tool.define("task", async (ctx) => {
|
|
const agents = await Agent.list().then((x) => x.filter((a) => a.mode !== "primary"))
|
|
|
|
// Filter agents by permissions if agent provided
|
|
const caller = ctx?.agent
|
|
const accessibleAgents = caller
|
|
? agents.filter((a) => PermissionNext.evaluate("task", a.name, caller.permission).action !== "deny")
|
|
: agents
|
|
|
|
const description = DESCRIPTION.replace(
|
|
"{agents}",
|
|
accessibleAgents
|
|
.map((a) => `- ${a.name}: ${a.description ?? "This subagent should only be called manually by the user."}`)
|
|
.join("\n"),
|
|
)
|
|
return {
|
|
description,
|
|
parameters,
|
|
async execute(params: z.infer<typeof parameters>, ctx) {
|
|
const config = await Config.get()
|
|
|
|
// Skip permission check when user explicitly invoked via @ or command subtask
|
|
if (!ctx.extra?.bypassAgentCheck) {
|
|
await ctx.ask({
|
|
permission: "task",
|
|
patterns: [params.subagent_type],
|
|
always: ["*"],
|
|
metadata: {
|
|
description: params.description,
|
|
subagent_type: params.subagent_type,
|
|
},
|
|
})
|
|
}
|
|
|
|
const agent = await Agent.get(params.subagent_type)
|
|
if (!agent) 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 session = await iife(async () => {
|
|
if (params.task_id) {
|
|
const found = await Session.get(params.task_id).catch(() => {})
|
|
if (found) return found
|
|
}
|
|
|
|
return await Session.create({
|
|
parentID: ctx.sessionID,
|
|
title: params.description + ` (@${agent.name} subagent)`,
|
|
permission: [
|
|
{
|
|
permission: "todowrite",
|
|
pattern: "*",
|
|
action: "deny",
|
|
},
|
|
{
|
|
permission: "todoread",
|
|
pattern: "*",
|
|
action: "deny",
|
|
},
|
|
...(hasTaskPermission
|
|
? []
|
|
: [
|
|
{
|
|
permission: "task" as const,
|
|
pattern: "*" as const,
|
|
action: "deny" as const,
|
|
},
|
|
]),
|
|
...(config.experimental?.primary_tools?.map((t) => ({
|
|
pattern: "*",
|
|
action: "allow" as const,
|
|
permission: t,
|
|
})) ?? []),
|
|
],
|
|
})
|
|
})
|
|
const msg = await MessageV2.get({ sessionID: ctx.sessionID, messageID: ctx.messageID })
|
|
if (msg.info.role !== "assistant") throw new Error("Not an assistant message")
|
|
|
|
const parentModel = {
|
|
modelID: msg.info.modelID,
|
|
providerID: msg.info.providerID,
|
|
}
|
|
const model = agent.model ?? parentModel
|
|
const background = params.background === true
|
|
const metadata = {
|
|
sessionId: session.id,
|
|
model,
|
|
...(background ? { background: true } : {}),
|
|
}
|
|
|
|
ctx.metadata({
|
|
title: params.description,
|
|
metadata,
|
|
})
|
|
|
|
const run = async () => {
|
|
const promptParts = await SessionPrompt.resolvePromptParts(params.prompt)
|
|
const result = await SessionPrompt.prompt({
|
|
messageID: Identifier.ascending("message"),
|
|
sessionID: session.id,
|
|
model: {
|
|
modelID: model.modelID,
|
|
providerID: model.providerID,
|
|
},
|
|
agent: agent.name,
|
|
tools: {
|
|
todowrite: false,
|
|
todoread: false,
|
|
...(hasTaskPermission ? {} : { task: false }),
|
|
...Object.fromEntries((config.experimental?.primary_tools ?? []).map((t) => [t, false])),
|
|
},
|
|
parts: promptParts,
|
|
})
|
|
return result.parts.findLast((x) => x.type === "text")?.text ?? ""
|
|
}
|
|
|
|
if (background) {
|
|
const inject = (state: "completed" | "error", text: string) =>
|
|
SessionPrompt.prompt({
|
|
sessionID: ctx.sessionID,
|
|
noReply: true,
|
|
model: {
|
|
modelID: parentModel.modelID,
|
|
providerID: parentModel.providerID,
|
|
},
|
|
agent: ctx.agent,
|
|
parts: [
|
|
{
|
|
type: "text",
|
|
synthetic: true,
|
|
text: backgroundMessage({
|
|
sessionID: session.id,
|
|
description: params.description,
|
|
state,
|
|
text,
|
|
}),
|
|
},
|
|
],
|
|
})
|
|
|
|
void run()
|
|
.then((text) =>
|
|
inject("completed", text)
|
|
.then((message) =>
|
|
continueParent({
|
|
parentID: ctx.sessionID,
|
|
userID: message.info.id,
|
|
taskID: session.id,
|
|
}),
|
|
)
|
|
.catch(() => {}),
|
|
)
|
|
.catch((error) =>
|
|
inject("error", errorText(error))
|
|
.then((message) =>
|
|
continueParent({
|
|
parentID: ctx.sessionID,
|
|
userID: message.info.id,
|
|
taskID: session.id,
|
|
}),
|
|
)
|
|
.catch(() => {}),
|
|
)
|
|
|
|
return {
|
|
title: params.description,
|
|
metadata,
|
|
output: backgroundOutput(session.id),
|
|
}
|
|
}
|
|
|
|
function cancel() {
|
|
SessionPrompt.cancel(session.id)
|
|
}
|
|
ctx.abort.addEventListener("abort", cancel)
|
|
using _ = defer(() => ctx.abort.removeEventListener("abort", cancel))
|
|
const text = await run()
|
|
|
|
return {
|
|
title: params.description,
|
|
metadata,
|
|
output: output(session.id, text),
|
|
}
|
|
},
|
|
}
|
|
})
|