task.ts 6.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174
  1. import * as Tool from "./tool"
  2. import DESCRIPTION from "./task.txt"
  3. import { Session } from "@/session/session"
  4. import { SessionID, MessageID } from "../session/schema"
  5. import { MessageV2 } from "../session/message-v2"
  6. import { Agent } from "../agent/agent"
  7. import { deriveSubagentSessionPermission } from "../agent/subagent-permissions"
  8. import type { SessionPrompt } from "../session/prompt"
  9. import { Config } from "@/config/config"
  10. import { Effect, Exit, Schema } from "effect"
  11. import { EffectBridge } from "@/effect/bridge"
  12. export interface TaskPromptOps {
  13. cancel(sessionID: SessionID): Effect.Effect<void>
  14. resolvePromptParts(template: string): Effect.Effect<SessionPrompt.PromptInput["parts"]>
  15. prompt(input: SessionPrompt.PromptInput): Effect.Effect<MessageV2.WithParts>
  16. }
  17. const id = "task"
  18. export const Parameters = Schema.Struct({
  19. description: Schema.String.annotate({ description: "A short (3-5 words) description of the task" }),
  20. prompt: Schema.String.annotate({ description: "The task for the agent to perform" }),
  21. subagent_type: Schema.String.annotate({ description: "The type of specialized agent to use for this task" }),
  22. task_id: Schema.optional(Schema.String).annotate({
  23. description:
  24. "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)",
  25. }),
  26. command: Schema.optional(Schema.String).annotate({ description: "The command that triggered this task" }),
  27. })
  28. export const TaskTool = Tool.define(
  29. id,
  30. Effect.gen(function* () {
  31. const agent = yield* Agent.Service
  32. const config = yield* Config.Service
  33. const sessions = yield* Session.Service
  34. const run = Effect.fn("TaskTool.execute")(function* (
  35. params: Schema.Schema.Type<typeof Parameters>,
  36. ctx: Tool.Context,
  37. ) {
  38. const cfg = yield* config.get()
  39. if (!ctx.extra?.bypassAgentCheck) {
  40. yield* ctx.ask({
  41. permission: id,
  42. patterns: [params.subagent_type],
  43. always: ["*"],
  44. metadata: {
  45. description: params.description,
  46. subagent_type: params.subagent_type,
  47. },
  48. })
  49. }
  50. const next = yield* agent.get(params.subagent_type)
  51. if (!next) {
  52. return yield* Effect.fail(new Error(`Unknown agent type: ${params.subagent_type} is not a valid agent type`))
  53. }
  54. const taskID = params.task_id
  55. const session = taskID
  56. ? yield* sessions.get(SessionID.make(taskID)).pipe(Effect.catchCause(() => Effect.succeed(undefined)))
  57. : undefined
  58. const parent = yield* sessions.get(ctx.sessionID)
  59. const parentAgent = parent.agent
  60. ? yield* agent.get(parent.agent).pipe(Effect.catchCause(() => Effect.succeed(undefined)))
  61. : undefined
  62. const nextSession =
  63. session ??
  64. (yield* sessions.create({
  65. parentID: ctx.sessionID,
  66. title: params.description + ` (@${next.name} subagent)`,
  67. permission: [
  68. ...deriveSubagentSessionPermission({
  69. parentSessionPermission: parent.permission ?? [],
  70. parentAgent,
  71. subagent: next,
  72. }),
  73. ...(cfg.experimental?.primary_tools?.map((item) => ({
  74. pattern: "*",
  75. action: "allow" as const,
  76. permission: item,
  77. })) ?? []),
  78. ],
  79. }))
  80. const msg = yield* Effect.sync(() => MessageV2.get({ sessionID: ctx.sessionID, messageID: ctx.messageID }))
  81. if (msg.info.role !== "assistant") return yield* Effect.fail(new Error("Not an assistant message"))
  82. const model = next.model ?? {
  83. modelID: msg.info.modelID,
  84. providerID: msg.info.providerID,
  85. }
  86. yield* ctx.metadata({
  87. title: params.description,
  88. metadata: {
  89. sessionId: nextSession.id,
  90. model,
  91. },
  92. })
  93. const ops = ctx.extra?.promptOps as TaskPromptOps
  94. if (!ops) return yield* Effect.fail(new Error("TaskTool requires promptOps in ctx.extra"))
  95. const runCancel = yield* EffectBridge.make()
  96. const messageID = MessageID.ascending()
  97. const cancel = ops.cancel(nextSession.id)
  98. function onAbort() {
  99. runCancel.fork(cancel)
  100. }
  101. return yield* Effect.acquireUseRelease(
  102. Effect.sync(() => {
  103. ctx.abort.addEventListener("abort", onAbort)
  104. }),
  105. () =>
  106. Effect.gen(function* () {
  107. const parts = yield* ops.resolvePromptParts(params.prompt)
  108. const result = yield* ops.prompt({
  109. messageID,
  110. sessionID: nextSession.id,
  111. model: {
  112. modelID: model.modelID,
  113. providerID: model.providerID,
  114. },
  115. agent: next.name,
  116. tools: {
  117. ...(next.permission.some((rule) => rule.permission === "todowrite") ? {} : { todowrite: false }),
  118. ...(next.permission.some((rule) => rule.permission === id) ? {} : { task: false }),
  119. ...Object.fromEntries((cfg.experimental?.primary_tools ?? []).map((item) => [item, false])),
  120. },
  121. parts,
  122. })
  123. return {
  124. title: params.description,
  125. metadata: {
  126. sessionId: nextSession.id,
  127. model,
  128. },
  129. output: [
  130. `task_id: ${nextSession.id} (for resuming to continue this task if needed)`,
  131. "",
  132. "<task_result>",
  133. result.parts.findLast((item) => item.type === "text")?.text ?? "",
  134. "</task_result>",
  135. ].join("\n"),
  136. }
  137. }),
  138. (_, exit) =>
  139. Effect.gen(function* () {
  140. if (Exit.hasInterrupts(exit)) yield* cancel
  141. }).pipe(
  142. Effect.ensuring(
  143. Effect.sync(() => {
  144. ctx.abort.removeEventListener("abort", onAbort)
  145. }),
  146. ),
  147. ),
  148. )
  149. })
  150. return {
  151. description: DESCRIPTION,
  152. parameters: Parameters,
  153. execute: (params: Schema.Schema.Type<typeof Parameters>, ctx: Tool.Context) =>
  154. run(params, ctx).pipe(Effect.orDie),
  155. }
  156. }),
  157. )