task.ts 6.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191
  1. import { Tool } from "./tool"
  2. import DESCRIPTION from "./task.txt"
  3. import z from "zod"
  4. import { Session } from "../session"
  5. import { Bus } from "../bus"
  6. import { MessageV2 } from "../session/message-v2"
  7. import { Identifier } from "../id/id"
  8. import { Agent } from "../agent/agent"
  9. import { SessionPrompt } from "../session/prompt"
  10. import { iife } from "@/util/iife"
  11. import { defer } from "@/util/defer"
  12. import { Config } from "../config/config"
  13. import { PermissionNext } from "@/permission/next"
  14. const parameters = z.object({
  15. description: z.string().describe("A short (3-5 words) description of the task"),
  16. prompt: z.string().describe("The task for the agent to perform"),
  17. subagent_type: z.string().describe("The type of specialized agent to use for this task"),
  18. session_id: z.string().describe("Existing Task session to continue").optional(),
  19. command: z.string().describe("The command that triggered this task").optional(),
  20. })
  21. export const TaskTool = Tool.define("task", async (ctx) => {
  22. const agents = await Agent.list().then((x) => x.filter((a) => a.mode !== "primary"))
  23. // Filter agents by permissions if agent provided
  24. const caller = ctx?.agent
  25. const accessibleAgents = caller
  26. ? agents.filter((a) => PermissionNext.evaluate("task", a.name, caller.permission).action !== "deny")
  27. : agents
  28. const description = DESCRIPTION.replace(
  29. "{agents}",
  30. accessibleAgents
  31. .map((a) => `- ${a.name}: ${a.description ?? "This subagent should only be called manually by the user."}`)
  32. .join("\n"),
  33. )
  34. return {
  35. description,
  36. parameters,
  37. async execute(params: z.infer<typeof parameters>, ctx) {
  38. const config = await Config.get()
  39. // Skip permission check when user explicitly invoked via @ or command subtask
  40. if (!ctx.extra?.bypassAgentCheck) {
  41. await ctx.ask({
  42. permission: "task",
  43. patterns: [params.subagent_type],
  44. always: ["*"],
  45. metadata: {
  46. description: params.description,
  47. subagent_type: params.subagent_type,
  48. },
  49. })
  50. }
  51. const agent = await Agent.get(params.subagent_type)
  52. if (!agent) throw new Error(`Unknown agent type: ${params.subagent_type} is not a valid agent type`)
  53. const hasTaskPermission = agent.permission.some((rule) => rule.permission === "task")
  54. const session = await iife(async () => {
  55. if (params.session_id) {
  56. const found = await Session.get(params.session_id).catch(() => {})
  57. if (found) return found
  58. }
  59. return await Session.create({
  60. parentID: ctx.sessionID,
  61. title: params.description + ` (@${agent.name} subagent)`,
  62. permission: [
  63. {
  64. permission: "todowrite",
  65. pattern: "*",
  66. action: "deny",
  67. },
  68. {
  69. permission: "todoread",
  70. pattern: "*",
  71. action: "deny",
  72. },
  73. ...(hasTaskPermission
  74. ? []
  75. : [
  76. {
  77. permission: "task" as const,
  78. pattern: "*" as const,
  79. action: "deny" as const,
  80. },
  81. ]),
  82. ...(config.experimental?.primary_tools?.map((t) => ({
  83. pattern: "*",
  84. action: "allow" as const,
  85. permission: t,
  86. })) ?? []),
  87. ],
  88. })
  89. })
  90. const msg = await MessageV2.get({ sessionID: ctx.sessionID, messageID: ctx.messageID })
  91. if (msg.info.role !== "assistant") throw new Error("Not an assistant message")
  92. const model = agent.model ?? {
  93. modelID: msg.info.modelID,
  94. providerID: msg.info.providerID,
  95. }
  96. ctx.metadata({
  97. title: params.description,
  98. metadata: {
  99. sessionId: session.id,
  100. model,
  101. },
  102. })
  103. const messageID = Identifier.ascending("message")
  104. const parts: Record<string, { id: string; tool: string; state: { status: string; title?: string } }> = {}
  105. const unsub = Bus.subscribe(MessageV2.Event.PartUpdated, async (evt) => {
  106. if (evt.properties.part.sessionID !== session.id) return
  107. if (evt.properties.part.messageID === messageID) return
  108. if (evt.properties.part.type !== "tool") return
  109. const part = evt.properties.part
  110. parts[part.id] = {
  111. id: part.id,
  112. tool: part.tool,
  113. state: {
  114. status: part.state.status,
  115. title: part.state.status === "completed" ? part.state.title : undefined,
  116. },
  117. }
  118. ctx.metadata({
  119. title: params.description,
  120. metadata: {
  121. summary: Object.values(parts).sort((a, b) => a.id.localeCompare(b.id)),
  122. sessionId: session.id,
  123. model,
  124. },
  125. })
  126. })
  127. function cancel() {
  128. SessionPrompt.cancel(session.id)
  129. }
  130. ctx.abort.addEventListener("abort", cancel)
  131. using _ = defer(() => ctx.abort.removeEventListener("abort", cancel))
  132. const promptParts = await SessionPrompt.resolvePromptParts(params.prompt)
  133. const result = await SessionPrompt.prompt({
  134. messageID,
  135. sessionID: session.id,
  136. model: {
  137. modelID: model.modelID,
  138. providerID: model.providerID,
  139. },
  140. agent: agent.name,
  141. tools: {
  142. todowrite: false,
  143. todoread: false,
  144. ...(hasTaskPermission ? {} : { task: false }),
  145. ...Object.fromEntries((config.experimental?.primary_tools ?? []).map((t) => [t, false])),
  146. },
  147. parts: promptParts,
  148. })
  149. unsub()
  150. const messages = await Session.messages({ sessionID: session.id })
  151. const summary = messages
  152. .filter((x) => x.info.role === "assistant")
  153. .flatMap((msg) => msg.parts.filter((x: any) => x.type === "tool") as MessageV2.ToolPart[])
  154. .map((part) => ({
  155. id: part.id,
  156. tool: part.tool,
  157. state: {
  158. status: part.state.status,
  159. title: part.state.status === "completed" ? part.state.title : undefined,
  160. },
  161. }))
  162. const text = result.parts.findLast((x) => x.type === "text")?.text ?? ""
  163. const output = text + "\n\n" + ["<task_metadata>", `session_id: ${session.id}`, "</task_metadata>"].join("\n")
  164. return {
  165. title: params.description,
  166. metadata: {
  167. summary,
  168. sessionId: session.id,
  169. model,
  170. },
  171. output,
  172. }
  173. },
  174. }
  175. })