plan.ts 4.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130
  1. import z from "zod"
  2. import path from "path"
  3. import { Tool } from "./tool"
  4. import { Question } from "../question"
  5. import { Session } from "../session"
  6. import { MessageV2 } from "../session/message-v2"
  7. import { Identifier } from "../id/id"
  8. import { Provider } from "../provider/provider"
  9. import { Instance } from "../project/instance"
  10. import EXIT_DESCRIPTION from "./plan-exit.txt"
  11. import ENTER_DESCRIPTION from "./plan-enter.txt"
  12. async function getLastModel(sessionID: string) {
  13. for await (const item of MessageV2.stream(sessionID)) {
  14. if (item.info.role === "user" && item.info.model) return item.info.model
  15. }
  16. return Provider.defaultModel()
  17. }
  18. export const PlanExitTool = Tool.define("plan_exit", {
  19. description: EXIT_DESCRIPTION,
  20. parameters: z.object({}),
  21. async execute(_params, ctx) {
  22. const session = await Session.get(ctx.sessionID)
  23. const plan = path.relative(Instance.worktree, Session.plan(session))
  24. const answers = await Question.ask({
  25. sessionID: ctx.sessionID,
  26. questions: [
  27. {
  28. question: `Plan at ${plan} is complete. Would you like to switch to the build agent and start implementing?`,
  29. header: "Build Agent",
  30. custom: false,
  31. options: [
  32. { label: "Yes", description: "Switch to build agent and start implementing the plan" },
  33. { label: "No", description: "Stay with plan agent to continue refining the plan" },
  34. ],
  35. },
  36. ],
  37. tool: ctx.callID ? { messageID: ctx.messageID, callID: ctx.callID } : undefined,
  38. })
  39. const answer = answers[0]?.[0]
  40. if (answer === "No") throw new Question.RejectedError()
  41. const model = await getLastModel(ctx.sessionID)
  42. const userMsg: MessageV2.User = {
  43. id: Identifier.ascending("message"),
  44. sessionID: ctx.sessionID,
  45. role: "user",
  46. time: {
  47. created: Date.now(),
  48. },
  49. agent: "build",
  50. model,
  51. }
  52. await Session.updateMessage(userMsg)
  53. await Session.updatePart({
  54. id: Identifier.ascending("part"),
  55. messageID: userMsg.id,
  56. sessionID: ctx.sessionID,
  57. type: "text",
  58. text: `The plan at ${plan} has been approved, you can now edit files. Execute the plan`,
  59. synthetic: true,
  60. } satisfies MessageV2.TextPart)
  61. return {
  62. title: "Switching to build agent",
  63. output: "User approved switching to build agent. Wait for further instructions.",
  64. metadata: {},
  65. }
  66. },
  67. })
  68. export const PlanEnterTool = Tool.define("plan_enter", {
  69. description: ENTER_DESCRIPTION,
  70. parameters: z.object({}),
  71. async execute(_params, ctx) {
  72. const session = await Session.get(ctx.sessionID)
  73. const plan = path.relative(Instance.worktree, Session.plan(session))
  74. const answers = await Question.ask({
  75. sessionID: ctx.sessionID,
  76. questions: [
  77. {
  78. question: `Would you like to switch to the plan agent and create a plan saved to ${plan}?`,
  79. header: "Plan Mode",
  80. custom: false,
  81. options: [
  82. { label: "Yes", description: "Switch to plan agent for research and planning" },
  83. { label: "No", description: "Stay with build agent to continue making changes" },
  84. ],
  85. },
  86. ],
  87. tool: ctx.callID ? { messageID: ctx.messageID, callID: ctx.callID } : undefined,
  88. })
  89. const answer = answers[0]?.[0]
  90. if (answer === "No") throw new Question.RejectedError()
  91. const model = await getLastModel(ctx.sessionID)
  92. const userMsg: MessageV2.User = {
  93. id: Identifier.ascending("message"),
  94. sessionID: ctx.sessionID,
  95. role: "user",
  96. time: {
  97. created: Date.now(),
  98. },
  99. agent: "plan",
  100. model,
  101. }
  102. await Session.updateMessage(userMsg)
  103. await Session.updatePart({
  104. id: Identifier.ascending("part"),
  105. messageID: userMsg.id,
  106. sessionID: ctx.sessionID,
  107. type: "text",
  108. text: "User has requested to enter plan mode. Switch to plan mode and begin planning.",
  109. synthetic: true,
  110. } satisfies MessageV2.TextPart)
  111. return {
  112. title: "Switching to plan agent",
  113. output: `User confirmed to switch to plan mode. A new message has been created to switch you to plan mode. The plan file will be at ${plan}. Begin planning.`,
  114. metadata: {},
  115. }
  116. },
  117. })