2
0

task.ts 4.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136
  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. export const TaskTool = Tool.define("task", async () => {
  14. const agents = await Agent.list().then((x) => x.filter((a) => a.mode !== "primary"))
  15. const description = DESCRIPTION.replace(
  16. "{agents}",
  17. agents
  18. .map((a) => `- ${a.name}: ${a.description ?? "This subagent should only be called manually by the user."}`)
  19. .join("\n"),
  20. )
  21. return {
  22. description,
  23. parameters: z.object({
  24. description: z.string().describe("A short (3-5 words) description of the task"),
  25. prompt: z.string().describe("The task for the agent to perform"),
  26. subagent_type: z.string().describe("The type of specialized agent to use for this task"),
  27. session_id: z.string().describe("Existing Task session to continue").optional(),
  28. command: z.string().describe("The command that triggered this task").optional(),
  29. }),
  30. async execute(params, ctx) {
  31. const agent = await Agent.get(params.subagent_type)
  32. if (!agent) throw new Error(`Unknown agent type: ${params.subagent_type} is not a valid agent type`)
  33. const session = await iife(async () => {
  34. if (params.session_id) {
  35. const found = await Session.get(params.session_id).catch(() => {})
  36. if (found) return found
  37. }
  38. return await Session.create({
  39. parentID: ctx.sessionID,
  40. title: params.description + ` (@${agent.name} subagent)`,
  41. })
  42. })
  43. const msg = await MessageV2.get({ sessionID: ctx.sessionID, messageID: ctx.messageID })
  44. if (msg.info.role !== "assistant") throw new Error("Not an assistant message")
  45. ctx.metadata({
  46. title: params.description,
  47. metadata: {
  48. sessionId: session.id,
  49. },
  50. })
  51. const messageID = Identifier.ascending("message")
  52. const parts: Record<string, { id: string; tool: string; state: { status: string; title?: string } }> = {}
  53. const unsub = Bus.subscribe(MessageV2.Event.PartUpdated, async (evt) => {
  54. if (evt.properties.part.sessionID !== session.id) return
  55. if (evt.properties.part.messageID === messageID) return
  56. if (evt.properties.part.type !== "tool") return
  57. const part = evt.properties.part
  58. parts[part.id] = {
  59. id: part.id,
  60. tool: part.tool,
  61. state: {
  62. status: part.state.status,
  63. title: part.state.status === "completed" ? part.state.title : undefined,
  64. },
  65. }
  66. ctx.metadata({
  67. title: params.description,
  68. metadata: {
  69. summary: Object.values(parts).sort((a, b) => a.id.localeCompare(b.id)),
  70. sessionId: session.id,
  71. },
  72. })
  73. })
  74. const model = agent.model ?? {
  75. modelID: msg.info.modelID,
  76. providerID: msg.info.providerID,
  77. }
  78. function cancel() {
  79. SessionPrompt.cancel(session.id)
  80. }
  81. ctx.abort.addEventListener("abort", cancel)
  82. using _ = defer(() => ctx.abort.removeEventListener("abort", cancel))
  83. const promptParts = await SessionPrompt.resolvePromptParts(params.prompt)
  84. const config = await Config.get()
  85. const result = await SessionPrompt.prompt({
  86. messageID,
  87. sessionID: session.id,
  88. model: {
  89. modelID: model.modelID,
  90. providerID: model.providerID,
  91. },
  92. agent: agent.name,
  93. tools: {
  94. todowrite: false,
  95. todoread: false,
  96. task: false,
  97. ...Object.fromEntries((config.experimental?.primary_tools ?? []).map((t) => [t, false])),
  98. ...agent.tools,
  99. },
  100. parts: promptParts,
  101. })
  102. unsub()
  103. const messages = await Session.messages({ sessionID: session.id })
  104. const summary = messages
  105. .filter((x) => x.info.role === "assistant")
  106. .flatMap((msg) => msg.parts.filter((x: any) => x.type === "tool") as MessageV2.ToolPart[])
  107. .map((part) => ({
  108. id: part.id,
  109. tool: part.tool,
  110. state: {
  111. status: part.state.status,
  112. title: part.state.status === "completed" ? part.state.title : undefined,
  113. },
  114. }))
  115. const text = result.parts.findLast((x) => x.type === "text")?.text ?? ""
  116. const output = text + "\n\n" + ["<task_metadata>", `session_id: ${session.id}`, "</task_metadata>"].join("\n")
  117. return {
  118. title: params.description,
  119. metadata: {
  120. summary,
  121. sessionId: session.id,
  122. },
  123. output,
  124. }
  125. },
  126. }
  127. })