index.ts 55 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576157715781579158015811582158315841585158615871588158915901591159215931594159515961597159815991600160116021603160416051606160716081609161016111612161316141615161616171618161916201621162216231624162516261627162816291630163116321633163416351636163716381639164016411642164316441645164616471648164916501651165216531654165516561657165816591660166116621663166416651666166716681669167016711672167316741675167616771678167916801681168216831684168516861687168816891690169116921693169416951696169716981699170017011702170317041705170617071708170917101711171217131714171517161717171817191720172117221723172417251726172717281729173017311732173317341735173617371738173917401741174217431744174517461747174817491750175117521753175417551756175717581759176017611762176317641765176617671768176917701771177217731774177517761777177817791780178117821783178417851786178717881789179017911792179317941795179617971798179918001801180218031804180518061807180818091810181118121813181418151816181718181819182018211822
  1. import os from "os"
  2. import path from "path"
  3. import fs from "fs/promises"
  4. import { spawn } from "child_process"
  5. import { Decimal } from "decimal.js"
  6. import { z, ZodSchema } from "zod"
  7. import {
  8. generateText,
  9. LoadAPIKeyError,
  10. streamText,
  11. tool,
  12. wrapLanguageModel,
  13. type Tool as AITool,
  14. type LanguageModelUsage,
  15. type ProviderMetadata,
  16. type ModelMessage,
  17. type StreamTextResult,
  18. } from "ai"
  19. import PROMPT_INITIALIZE from "../session/prompt/initialize.txt"
  20. import PROMPT_PLAN from "../session/prompt/plan.txt"
  21. import BUILD_SWITCH from "../session/prompt/build-switch.txt"
  22. import { Bus } from "../bus"
  23. import { Config } from "../config/config"
  24. import { Flag } from "../flag/flag"
  25. import { Identifier } from "../id/id"
  26. import { Installation } from "../installation"
  27. import { MCP } from "../mcp"
  28. import { Provider } from "../provider/provider"
  29. import { ProviderTransform } from "../provider/transform"
  30. import type { ModelsDev } from "../provider/models"
  31. import { Share } from "../share/share"
  32. import { Snapshot } from "../snapshot"
  33. import { Storage } from "../storage/storage"
  34. import { Log } from "../util/log"
  35. import { NamedError } from "../util/error"
  36. import { SystemPrompt } from "./system"
  37. import { FileTime } from "../file/time"
  38. import { MessageV2 } from "./message-v2"
  39. import { LSP } from "../lsp"
  40. import { ReadTool } from "../tool/read"
  41. import { mergeDeep, pipe, splitWhen } from "remeda"
  42. import { ToolRegistry } from "../tool/registry"
  43. import { Plugin } from "../plugin"
  44. import { Project } from "../project/project"
  45. import { Instance } from "../project/instance"
  46. import { Agent } from "../agent/agent"
  47. import { Permission } from "../permission"
  48. import { Wildcard } from "../util/wildcard"
  49. import { ulid } from "ulid"
  50. import { defer } from "../util/defer"
  51. import { Command } from "../command"
  52. import { $ } from "bun"
  53. export namespace Session {
  54. const log = Log.create({ service: "session" })
  55. const OUTPUT_TOKEN_MAX = 32_000
  56. const parentSessionTitlePrefix = "New session - "
  57. const childSessionTitlePrefix = "Child session - "
  58. function createDefaultTitle(isChild = false) {
  59. return (isChild ? childSessionTitlePrefix : parentSessionTitlePrefix) + new Date().toISOString()
  60. }
  61. function isDefaultTitle(title: string) {
  62. return title.startsWith(parentSessionTitlePrefix)
  63. }
  64. export const Info = z
  65. .object({
  66. id: Identifier.schema("session"),
  67. projectID: z.string(),
  68. directory: z.string(),
  69. parentID: Identifier.schema("session").optional(),
  70. share: z
  71. .object({
  72. url: z.string(),
  73. })
  74. .optional(),
  75. title: z.string(),
  76. version: z.string(),
  77. time: z.object({
  78. created: z.number(),
  79. updated: z.number(),
  80. }),
  81. revert: z
  82. .object({
  83. messageID: z.string(),
  84. partID: z.string().optional(),
  85. snapshot: z.string().optional(),
  86. diff: z.string().optional(),
  87. })
  88. .optional(),
  89. })
  90. .openapi({
  91. ref: "Session",
  92. })
  93. export type Info = z.output<typeof Info>
  94. export const ShareInfo = z
  95. .object({
  96. secret: z.string(),
  97. url: z.string(),
  98. })
  99. .openapi({
  100. ref: "SessionShare",
  101. })
  102. export type ShareInfo = z.output<typeof ShareInfo>
  103. export const Event = {
  104. Updated: Bus.event(
  105. "session.updated",
  106. z.object({
  107. info: Info,
  108. }),
  109. ),
  110. Deleted: Bus.event(
  111. "session.deleted",
  112. z.object({
  113. info: Info,
  114. }),
  115. ),
  116. Idle: Bus.event(
  117. "session.idle",
  118. z.object({
  119. sessionID: z.string(),
  120. }),
  121. ),
  122. Error: Bus.event(
  123. "session.error",
  124. z.object({
  125. sessionID: z.string().optional(),
  126. error: MessageV2.Assistant.shape.error,
  127. }),
  128. ),
  129. }
  130. const state = Instance.state(
  131. () => {
  132. const pending = new Map<string, AbortController>()
  133. const autoCompacting = new Map<string, boolean>()
  134. const queued = new Map<
  135. string,
  136. {
  137. input: ChatInput
  138. message: MessageV2.User
  139. parts: MessageV2.Part[]
  140. processed: boolean
  141. callback: (input: { info: MessageV2.Assistant; parts: MessageV2.Part[] }) => void
  142. }[]
  143. >()
  144. return {
  145. pending,
  146. autoCompacting,
  147. queued,
  148. }
  149. },
  150. async (state) => {
  151. for (const [_, controller] of state.pending) {
  152. controller.abort()
  153. }
  154. },
  155. )
  156. export async function create(parentID?: string, title?: string) {
  157. return createNext({
  158. parentID,
  159. directory: Instance.directory,
  160. title,
  161. })
  162. }
  163. export async function createNext(input: { id?: string; title?: string; parentID?: string; directory: string }) {
  164. const result: Info = {
  165. id: Identifier.descending("session", input.id),
  166. version: Installation.VERSION,
  167. projectID: Instance.project.id,
  168. directory: input.directory,
  169. parentID: input.parentID,
  170. title: input.title ?? createDefaultTitle(!!input.parentID),
  171. time: {
  172. created: Date.now(),
  173. updated: Date.now(),
  174. },
  175. }
  176. log.info("created", result)
  177. await Storage.write(["session", Instance.project.id, result.id], result)
  178. const cfg = await Config.get()
  179. if (!result.parentID && (Flag.OPENCODE_AUTO_SHARE || cfg.share === "auto"))
  180. share(result.id)
  181. .then((share) => {
  182. update(result.id, (draft) => {
  183. draft.share = share
  184. })
  185. })
  186. .catch(() => {
  187. // Silently ignore sharing errors during session creation
  188. })
  189. Bus.publish(Event.Updated, {
  190. info: result,
  191. })
  192. return result
  193. }
  194. export async function get(id: string) {
  195. const read = await Storage.read<Info>(["session", Instance.project.id, id])
  196. return read as Info
  197. }
  198. export async function getShare(id: string) {
  199. return Storage.read<ShareInfo>(["share", id])
  200. }
  201. export async function share(id: string) {
  202. const cfg = await Config.get()
  203. if (cfg.share === "disabled") {
  204. throw new Error("Sharing is disabled in configuration")
  205. }
  206. const session = await get(id)
  207. if (session.share) return session.share
  208. const share = await Share.create(id)
  209. await update(id, (draft) => {
  210. draft.share = {
  211. url: share.url,
  212. }
  213. })
  214. await Storage.write(["share", id], share)
  215. await Share.sync("session/info/" + id, session)
  216. for (const msg of await messages(id)) {
  217. await Share.sync("session/message/" + id + "/" + msg.info.id, msg.info)
  218. for (const part of msg.parts) {
  219. await Share.sync("session/part/" + id + "/" + msg.info.id + "/" + part.id, part)
  220. }
  221. }
  222. return share
  223. }
  224. export async function unshare(id: string) {
  225. const share = await getShare(id)
  226. if (!share) return
  227. await Storage.remove(["share", id])
  228. await update(id, (draft) => {
  229. draft.share = undefined
  230. })
  231. await Share.remove(id, share.secret)
  232. }
  233. export async function update(id: string, editor: (session: Info) => void) {
  234. const project = Instance.project
  235. const result = await Storage.update<Info>(["session", project.id, id], (draft) => {
  236. editor(draft)
  237. draft.time.updated = Date.now()
  238. })
  239. Bus.publish(Event.Updated, {
  240. info: result,
  241. })
  242. return result
  243. }
  244. export async function messages(sessionID: string) {
  245. const result = [] as {
  246. info: MessageV2.Info
  247. parts: MessageV2.Part[]
  248. }[]
  249. for (const p of await Storage.list(["message", sessionID])) {
  250. const read = await Storage.read<MessageV2.Info>(p)
  251. result.push({
  252. info: read,
  253. parts: await getParts(read.id),
  254. })
  255. }
  256. result.sort((a, b) => (a.info.id > b.info.id ? 1 : -1))
  257. return result
  258. }
  259. export async function getMessage(sessionID: string, messageID: string) {
  260. return {
  261. info: await Storage.read<MessageV2.Info>(["message", sessionID, messageID]),
  262. parts: await getParts(messageID),
  263. }
  264. }
  265. export async function getParts(messageID: string) {
  266. const result = [] as MessageV2.Part[]
  267. for (const item of await Storage.list(["part", messageID])) {
  268. const read = await Storage.read<MessageV2.Part>(item)
  269. result.push(read)
  270. }
  271. result.sort((a, b) => (a.id > b.id ? 1 : -1))
  272. return result
  273. }
  274. export async function* list() {
  275. const project = Instance.project
  276. for (const item of await Storage.list(["session", project.id])) {
  277. yield Storage.read<Info>(item)
  278. }
  279. }
  280. export async function children(parentID: string) {
  281. const project = Instance.project
  282. const result = [] as Session.Info[]
  283. for (const item of await Storage.list(["session", project.id])) {
  284. const session = await Storage.read<Info>(item)
  285. if (session.parentID !== parentID) continue
  286. result.push(session)
  287. }
  288. return result
  289. }
  290. export function abort(sessionID: string) {
  291. const controller = state().pending.get(sessionID)
  292. if (!controller) return false
  293. log.info("aborting", {
  294. sessionID,
  295. })
  296. controller.abort()
  297. state().pending.delete(sessionID)
  298. return true
  299. }
  300. export async function remove(sessionID: string, emitEvent = true) {
  301. const project = Instance.project
  302. try {
  303. abort(sessionID)
  304. const session = await get(sessionID)
  305. for (const child of await children(sessionID)) {
  306. await remove(child.id, false)
  307. }
  308. await unshare(sessionID).catch(() => {})
  309. for (const msg of await Storage.list(["message", sessionID])) {
  310. for (const part of await Storage.list(["part", msg.at(-1)!])) {
  311. await Storage.remove(part)
  312. }
  313. await Storage.remove(msg)
  314. }
  315. await Storage.remove(["session", project.id, sessionID])
  316. if (emitEvent) {
  317. Bus.publish(Event.Deleted, {
  318. info: session,
  319. })
  320. }
  321. } catch (e) {
  322. log.error(e)
  323. }
  324. }
  325. async function updateMessage(msg: MessageV2.Info) {
  326. await Storage.write(["message", msg.sessionID, msg.id], msg)
  327. Bus.publish(MessageV2.Event.Updated, {
  328. info: msg,
  329. })
  330. }
  331. async function updatePart(part: MessageV2.Part) {
  332. await Storage.write(["part", part.messageID, part.id], part)
  333. Bus.publish(MessageV2.Event.PartUpdated, {
  334. part,
  335. })
  336. return part
  337. }
  338. export const PromptInput = z.object({
  339. sessionID: Identifier.schema("session"),
  340. messageID: Identifier.schema("message").optional(),
  341. model: z
  342. .object({
  343. providerID: z.string(),
  344. modelID: z.string(),
  345. })
  346. .optional(),
  347. agent: z.string().optional(),
  348. system: z.string().optional(),
  349. tools: z.record(z.boolean()).optional(),
  350. parts: z.array(
  351. z.discriminatedUnion("type", [
  352. MessageV2.TextPart.omit({
  353. messageID: true,
  354. sessionID: true,
  355. })
  356. .partial({
  357. id: true,
  358. })
  359. .openapi({
  360. ref: "TextPartInput",
  361. }),
  362. MessageV2.FilePart.omit({
  363. messageID: true,
  364. sessionID: true,
  365. })
  366. .partial({
  367. id: true,
  368. })
  369. .openapi({
  370. ref: "FilePartInput",
  371. }),
  372. MessageV2.AgentPart.omit({
  373. messageID: true,
  374. sessionID: true,
  375. })
  376. .partial({
  377. id: true,
  378. })
  379. .openapi({
  380. ref: "AgentPartInput",
  381. }),
  382. ]),
  383. ),
  384. })
  385. export type ChatInput = z.infer<typeof PromptInput>
  386. export async function prompt(
  387. input: z.infer<typeof PromptInput>,
  388. ): Promise<{ info: MessageV2.Assistant; parts: MessageV2.Part[] }> {
  389. const l = log.clone().tag("session", input.sessionID)
  390. l.info("chatting")
  391. const inputAgent = input.agent ?? "build"
  392. // Process revert cleanup first, before creating new messages
  393. const session = await get(input.sessionID)
  394. if (session.revert) {
  395. let msgs = await messages(input.sessionID)
  396. const messageID = session.revert.messageID
  397. const [preserve, remove] = splitWhen(msgs, (x) => x.info.id === messageID)
  398. msgs = preserve
  399. for (const msg of remove) {
  400. await Storage.remove(["message", input.sessionID, msg.info.id])
  401. await Bus.publish(MessageV2.Event.Removed, { sessionID: input.sessionID, messageID: msg.info.id })
  402. }
  403. const last = preserve.at(-1)
  404. if (session.revert.partID && last) {
  405. const partID = session.revert.partID
  406. const [preserveParts, removeParts] = splitWhen(last.parts, (x) => x.id === partID)
  407. last.parts = preserveParts
  408. for (const part of removeParts) {
  409. await Storage.remove(["part", last.info.id, part.id])
  410. await Bus.publish(MessageV2.Event.PartRemoved, {
  411. sessionID: input.sessionID,
  412. messageID: last.info.id,
  413. partID: part.id,
  414. })
  415. }
  416. }
  417. await update(input.sessionID, (draft) => {
  418. draft.revert = undefined
  419. })
  420. }
  421. const userMsg: MessageV2.Info = {
  422. id: input.messageID ?? Identifier.ascending("message"),
  423. role: "user",
  424. sessionID: input.sessionID,
  425. time: {
  426. created: Date.now(),
  427. },
  428. }
  429. const userParts = await Promise.all(
  430. input.parts.map(async (part): Promise<MessageV2.Part[]> => {
  431. if (part.type === "file") {
  432. const url = new URL(part.url)
  433. switch (url.protocol) {
  434. case "data:":
  435. if (part.mime === "text/plain") {
  436. return [
  437. {
  438. id: Identifier.ascending("part"),
  439. messageID: userMsg.id,
  440. sessionID: input.sessionID,
  441. type: "text",
  442. synthetic: true,
  443. text: `Called the Read tool with the following input: ${JSON.stringify({ filePath: part.filename })}`,
  444. },
  445. {
  446. id: Identifier.ascending("part"),
  447. messageID: userMsg.id,
  448. sessionID: input.sessionID,
  449. type: "text",
  450. synthetic: true,
  451. text: Buffer.from(part.url, "base64url").toString(),
  452. },
  453. {
  454. ...part,
  455. id: part.id ?? Identifier.ascending("part"),
  456. messageID: userMsg.id,
  457. sessionID: input.sessionID,
  458. },
  459. ]
  460. }
  461. break
  462. case "file:":
  463. // have to normalize, symbol search returns absolute paths
  464. // Decode the pathname since URL constructor doesn't automatically decode it
  465. const filePath = decodeURIComponent(url.pathname)
  466. if (part.mime === "text/plain") {
  467. let offset: number | undefined = undefined
  468. let limit: number | undefined = undefined
  469. const range = {
  470. start: url.searchParams.get("start"),
  471. end: url.searchParams.get("end"),
  472. }
  473. if (range.start != null) {
  474. const filePath = part.url.split("?")[0]
  475. let start = parseInt(range.start)
  476. let end = range.end ? parseInt(range.end) : undefined
  477. // some LSP servers (eg, gopls) don't give full range in
  478. // workspace/symbol searches, so we'll try to find the
  479. // symbol in the document to get the full range
  480. if (start === end) {
  481. const symbols = await LSP.documentSymbol(filePath)
  482. for (const symbol of symbols) {
  483. let range: LSP.Range | undefined
  484. if ("range" in symbol) {
  485. range = symbol.range
  486. } else if ("location" in symbol) {
  487. range = symbol.location.range
  488. }
  489. if (range?.start?.line && range?.start?.line === start) {
  490. start = range.start.line
  491. end = range?.end?.line ?? start
  492. break
  493. }
  494. }
  495. }
  496. offset = Math.max(start - 1, 0)
  497. if (end) {
  498. limit = end - offset
  499. }
  500. }
  501. const args = { filePath, offset, limit }
  502. const result = await ReadTool.init().then((t) =>
  503. t.execute(args, {
  504. sessionID: input.sessionID,
  505. abort: new AbortController().signal,
  506. agent: input.agent!,
  507. messageID: userMsg.id,
  508. extra: { bypassCwdCheck: true },
  509. metadata: async () => {},
  510. }),
  511. )
  512. return [
  513. {
  514. id: Identifier.ascending("part"),
  515. messageID: userMsg.id,
  516. sessionID: input.sessionID,
  517. type: "text",
  518. synthetic: true,
  519. text: `Called the Read tool with the following input: ${JSON.stringify(args)}`,
  520. },
  521. {
  522. id: Identifier.ascending("part"),
  523. messageID: userMsg.id,
  524. sessionID: input.sessionID,
  525. type: "text",
  526. synthetic: true,
  527. text: result.output,
  528. },
  529. {
  530. ...part,
  531. id: part.id ?? Identifier.ascending("part"),
  532. messageID: userMsg.id,
  533. sessionID: input.sessionID,
  534. },
  535. ]
  536. }
  537. let file = Bun.file(filePath)
  538. FileTime.read(input.sessionID, filePath)
  539. return [
  540. {
  541. id: Identifier.ascending("part"),
  542. messageID: userMsg.id,
  543. sessionID: input.sessionID,
  544. type: "text",
  545. text: `Called the Read tool with the following input: {\"filePath\":\"${filePath}\"}`,
  546. synthetic: true,
  547. },
  548. {
  549. id: part.id ?? Identifier.ascending("part"),
  550. messageID: userMsg.id,
  551. sessionID: input.sessionID,
  552. type: "file",
  553. url: `data:${part.mime};base64,` + Buffer.from(await file.bytes()).toString("base64"),
  554. mime: part.mime,
  555. filename: part.filename!,
  556. source: part.source,
  557. },
  558. ]
  559. }
  560. }
  561. if (part.type === "agent") {
  562. return [
  563. {
  564. id: Identifier.ascending("part"),
  565. ...part,
  566. messageID: userMsg.id,
  567. sessionID: input.sessionID,
  568. },
  569. {
  570. id: Identifier.ascending("part"),
  571. messageID: userMsg.id,
  572. sessionID: input.sessionID,
  573. type: "text",
  574. synthetic: true,
  575. text:
  576. "Use the above message and context to generate a prompt and call the task tool with subagent: " +
  577. part.name,
  578. },
  579. ]
  580. }
  581. return [
  582. {
  583. id: Identifier.ascending("part"),
  584. ...part,
  585. messageID: userMsg.id,
  586. sessionID: input.sessionID,
  587. },
  588. ]
  589. }),
  590. ).then((x) => x.flat())
  591. await Plugin.trigger(
  592. "chat.message",
  593. {},
  594. {
  595. message: userMsg,
  596. parts: userParts,
  597. },
  598. )
  599. await updateMessage(userMsg)
  600. for (const part of userParts) {
  601. await updatePart(part)
  602. }
  603. // mark session as updated
  604. // used for session list sorting (indicates when session was most recently interacted with)
  605. await update(input.sessionID, (_draft) => {})
  606. if (isLocked(input.sessionID)) {
  607. return new Promise((resolve) => {
  608. const queue = state().queued.get(input.sessionID) ?? []
  609. queue.push({
  610. input: input,
  611. message: userMsg,
  612. parts: userParts,
  613. processed: false,
  614. callback: resolve,
  615. })
  616. state().queued.set(input.sessionID, queue)
  617. })
  618. }
  619. const agent = await Agent.get(inputAgent)
  620. const model = await (async () => {
  621. if (input.model) {
  622. return input.model
  623. }
  624. if (agent.model) {
  625. return agent.model
  626. }
  627. return Provider.defaultModel()
  628. })().then((x) => Provider.getModel(x.providerID, x.modelID))
  629. let msgs = await messages(input.sessionID)
  630. const previous = msgs.filter((x) => x.info.role === "assistant").at(-1)?.info as MessageV2.Assistant
  631. const outputLimit = Math.min(model.info.limit.output, OUTPUT_TOKEN_MAX) || OUTPUT_TOKEN_MAX
  632. // auto summarize if too long
  633. if (previous && previous.tokens) {
  634. const tokens =
  635. previous.tokens.input + previous.tokens.cache.read + previous.tokens.cache.write + previous.tokens.output
  636. if (model.info.limit.context && tokens > Math.max((model.info.limit.context - outputLimit) * 0.9, 0)) {
  637. state().autoCompacting.set(input.sessionID, true)
  638. await summarize({
  639. sessionID: input.sessionID,
  640. providerID: model.providerID,
  641. modelID: model.info.id,
  642. })
  643. return prompt(input)
  644. }
  645. }
  646. using abort = lock(input.sessionID)
  647. const lastSummary = msgs.findLast((msg) => msg.info.role === "assistant" && msg.info.summary === true)
  648. if (lastSummary) msgs = msgs.filter((msg) => msg.info.id >= lastSummary.info.id)
  649. if (msgs.filter((m) => m.info.role === "user").length === 1 && !session.parentID && isDefaultTitle(session.title)) {
  650. const small = (await Provider.getSmallModel(model.providerID)) ?? model
  651. generateText({
  652. maxOutputTokens: small.info.reasoning ? 1024 : 20,
  653. providerOptions: {
  654. [model.providerID]: {
  655. ...small.info.options,
  656. ...ProviderTransform.options(small.providerID, small.modelID, input.sessionID),
  657. },
  658. },
  659. messages: [
  660. ...SystemPrompt.title(model.providerID).map(
  661. (x): ModelMessage => ({
  662. role: "system",
  663. content: x,
  664. }),
  665. ),
  666. ...MessageV2.toModelMessage([
  667. {
  668. info: {
  669. id: Identifier.ascending("message"),
  670. role: "user",
  671. sessionID: input.sessionID,
  672. time: {
  673. created: Date.now(),
  674. },
  675. },
  676. parts: userParts,
  677. },
  678. ]),
  679. ],
  680. model: small.language,
  681. })
  682. .then((result) => {
  683. if (result.text)
  684. return Session.update(input.sessionID, (draft) => {
  685. const cleaned = result.text.replace(/<think>[\s\S]*?<\/think>\s*/g, "")
  686. const title = cleaned.length > 100 ? cleaned.substring(0, 97) + "..." : cleaned
  687. draft.title = title.trim()
  688. })
  689. })
  690. .catch((error) => {
  691. log.error("failed to generate title", { error, model: small.info.id })
  692. })
  693. }
  694. if (agent.name === "plan") {
  695. msgs.at(-1)?.parts.push({
  696. id: Identifier.ascending("part"),
  697. messageID: userMsg.id,
  698. sessionID: input.sessionID,
  699. type: "text",
  700. text: PROMPT_PLAN,
  701. synthetic: true,
  702. })
  703. }
  704. const lastAssistantMsg = msgs.filter((x) => x.info.role === "assistant").at(-1)?.info as MessageV2.Assistant
  705. if (lastAssistantMsg?.mode === "plan" && agent.name === "build") {
  706. msgs.at(-1)?.parts.push({
  707. id: Identifier.ascending("part"),
  708. messageID: userMsg.id,
  709. sessionID: input.sessionID,
  710. type: "text",
  711. text: BUILD_SWITCH,
  712. synthetic: true,
  713. })
  714. }
  715. let system = SystemPrompt.header(model.providerID)
  716. system.push(
  717. ...(() => {
  718. if (input.system) return [input.system]
  719. if (agent.prompt) return [agent.prompt]
  720. return SystemPrompt.provider(model.modelID)
  721. })(),
  722. )
  723. system.push(...(await SystemPrompt.environment()))
  724. system.push(...(await SystemPrompt.custom()))
  725. // max 2 system prompt messages for caching purposes
  726. const [first, ...rest] = system
  727. system = [first, rest.join("\n")]
  728. const assistantMsg: MessageV2.Info = {
  729. id: Identifier.ascending("message"),
  730. role: "assistant",
  731. system,
  732. mode: inputAgent,
  733. path: {
  734. cwd: Instance.directory,
  735. root: Instance.worktree,
  736. },
  737. cost: 0,
  738. tokens: {
  739. input: 0,
  740. output: 0,
  741. reasoning: 0,
  742. cache: { read: 0, write: 0 },
  743. },
  744. modelID: model.modelID,
  745. providerID: model.providerID,
  746. time: {
  747. created: Date.now(),
  748. },
  749. sessionID: input.sessionID,
  750. }
  751. await updateMessage(assistantMsg)
  752. await using _ = defer(async () => {
  753. if (assistantMsg.time.completed) return
  754. await Storage.remove(["session", "message", input.sessionID, assistantMsg.id])
  755. await Bus.publish(MessageV2.Event.Removed, { sessionID: input.sessionID, messageID: assistantMsg.id })
  756. })
  757. const tools: Record<string, AITool> = {}
  758. const processor = createProcessor(assistantMsg, model.info)
  759. const enabledTools = pipe(
  760. agent.tools,
  761. mergeDeep(await ToolRegistry.enabled(model.providerID, model.modelID, agent)),
  762. mergeDeep(input.tools ?? {}),
  763. )
  764. for (const item of await ToolRegistry.tools(model.providerID, model.modelID)) {
  765. if (Wildcard.all(item.id, enabledTools) === false) continue
  766. tools[item.id] = tool({
  767. id: item.id as any,
  768. description: item.description,
  769. inputSchema: item.parameters as ZodSchema,
  770. async execute(args, options) {
  771. await Plugin.trigger(
  772. "tool.execute.before",
  773. {
  774. tool: item.id,
  775. sessionID: input.sessionID,
  776. callID: options.toolCallId,
  777. },
  778. {
  779. args,
  780. },
  781. )
  782. const result = await item.execute(args, {
  783. sessionID: input.sessionID,
  784. abort: options.abortSignal!,
  785. messageID: assistantMsg.id,
  786. callID: options.toolCallId,
  787. agent: agent.name,
  788. metadata: async (val) => {
  789. const match = processor.partFromToolCall(options.toolCallId)
  790. if (match && match.state.status === "running") {
  791. await updatePart({
  792. ...match,
  793. state: {
  794. title: val.title,
  795. metadata: val.metadata,
  796. status: "running",
  797. input: args,
  798. time: {
  799. start: Date.now(),
  800. },
  801. },
  802. })
  803. }
  804. },
  805. })
  806. await Plugin.trigger(
  807. "tool.execute.after",
  808. {
  809. tool: item.id,
  810. sessionID: input.sessionID,
  811. callID: options.toolCallId,
  812. },
  813. result,
  814. )
  815. return result
  816. },
  817. toModelOutput(result) {
  818. return {
  819. type: "text",
  820. value: result.output,
  821. }
  822. },
  823. })
  824. }
  825. for (const [key, item] of Object.entries(await MCP.tools())) {
  826. if (Wildcard.all(key, enabledTools) === false) continue
  827. const execute = item.execute
  828. if (!execute) continue
  829. item.execute = async (args, opts) => {
  830. await Plugin.trigger(
  831. "tool.execute.before",
  832. {
  833. tool: key,
  834. sessionID: input.sessionID,
  835. callID: opts.toolCallId,
  836. },
  837. {
  838. args,
  839. },
  840. )
  841. const result = await execute(args, opts)
  842. const output = result.content
  843. .filter((x: any) => x.type === "text")
  844. .map((x: any) => x.text)
  845. .join("\n\n")
  846. await Plugin.trigger(
  847. "tool.execute.after",
  848. {
  849. tool: key,
  850. sessionID: input.sessionID,
  851. callID: opts.toolCallId,
  852. },
  853. result,
  854. )
  855. return {
  856. output,
  857. }
  858. }
  859. item.toModelOutput = (result) => {
  860. return {
  861. type: "text",
  862. value: result.output,
  863. }
  864. }
  865. tools[key] = item
  866. }
  867. const params = await Plugin.trigger(
  868. "chat.params",
  869. {
  870. model: model.info,
  871. provider: await Provider.getProvider(model.providerID),
  872. message: userMsg,
  873. },
  874. {
  875. temperature: model.info.temperature
  876. ? (agent.temperature ?? ProviderTransform.temperature(model.providerID, model.modelID))
  877. : undefined,
  878. topP: agent.topP ?? ProviderTransform.topP(model.providerID, model.modelID),
  879. options: {
  880. ...ProviderTransform.options(model.providerID, model.modelID, input.sessionID),
  881. ...model.info.options,
  882. ...agent.options,
  883. },
  884. },
  885. )
  886. const stream = streamText({
  887. onError(e) {
  888. log.error("streamText error", {
  889. error: e,
  890. })
  891. },
  892. async prepareStep({ messages }) {
  893. const queue = (state().queued.get(input.sessionID) ?? []).filter((x) => !x.processed)
  894. if (queue.length) {
  895. for (const item of queue) {
  896. if (item.processed) continue
  897. messages.push(
  898. ...MessageV2.toModelMessage([
  899. {
  900. info: item.message,
  901. parts: item.parts,
  902. },
  903. ]),
  904. )
  905. item.processed = true
  906. }
  907. assistantMsg.time.completed = Date.now()
  908. await updateMessage(assistantMsg)
  909. Object.assign(assistantMsg, {
  910. id: Identifier.ascending("message"),
  911. role: "assistant",
  912. system,
  913. path: {
  914. cwd: Instance.directory,
  915. root: Instance.worktree,
  916. },
  917. cost: 0,
  918. tokens: {
  919. input: 0,
  920. output: 0,
  921. reasoning: 0,
  922. cache: { read: 0, write: 0 },
  923. },
  924. modelID: model.modelID,
  925. providerID: model.providerID,
  926. mode: inputAgent,
  927. time: {
  928. created: Date.now(),
  929. },
  930. sessionID: input.sessionID,
  931. })
  932. await updateMessage(assistantMsg)
  933. }
  934. return {
  935. messages,
  936. }
  937. },
  938. async experimental_repairToolCall(input) {
  939. return {
  940. ...input.toolCall,
  941. input: JSON.stringify({
  942. tool: input.toolCall.toolName,
  943. error: input.error.message,
  944. }),
  945. toolName: "invalid",
  946. }
  947. },
  948. headers:
  949. model.providerID === "opencode"
  950. ? {
  951. "x-opencode-session": input.sessionID,
  952. "x-opencode-request": userMsg.id,
  953. }
  954. : undefined,
  955. maxRetries: 3,
  956. activeTools: Object.keys(tools).filter((x) => x !== "invalid"),
  957. maxOutputTokens: outputLimit,
  958. abortSignal: abort.signal,
  959. stopWhen: async ({ steps }) => {
  960. if (steps.length >= 1000) {
  961. return true
  962. }
  963. // Check if processor flagged that we should stop
  964. if (processor.getShouldStop()) {
  965. return true
  966. }
  967. return false
  968. },
  969. providerOptions: {
  970. [model.providerID]: params.options,
  971. },
  972. temperature: params.temperature,
  973. topP: params.topP,
  974. messages: [
  975. ...system.map(
  976. (x): ModelMessage => ({
  977. role: "system",
  978. content: x,
  979. }),
  980. ),
  981. ...MessageV2.toModelMessage(msgs.filter((m) => !(m.info.role === "assistant" && m.info.error))),
  982. ],
  983. tools: model.info.tool_call === false ? undefined : tools,
  984. model: wrapLanguageModel({
  985. model: model.language,
  986. middleware: [
  987. {
  988. async transformParams(args) {
  989. if (args.type === "stream") {
  990. // @ts-expect-error
  991. args.params.prompt = ProviderTransform.message(args.params.prompt, model.providerID, model.modelID)
  992. }
  993. return args.params
  994. },
  995. },
  996. ],
  997. }),
  998. })
  999. const result = await processor.process(stream)
  1000. const queued = state().queued.get(input.sessionID) ?? []
  1001. const unprocessed = queued.find((x) => !x.processed)
  1002. if (unprocessed) {
  1003. unprocessed.processed = true
  1004. return prompt(unprocessed.input)
  1005. }
  1006. for (const item of queued) {
  1007. item.callback(result)
  1008. }
  1009. state().queued.delete(input.sessionID)
  1010. return result
  1011. }
  1012. export const ShellInput = z.object({
  1013. sessionID: Identifier.schema("session"),
  1014. agent: z.string(),
  1015. command: z.string(),
  1016. })
  1017. export type ShellInput = z.infer<typeof ShellInput>
  1018. export async function shell(input: ShellInput) {
  1019. using abort = lock(input.sessionID)
  1020. const userMsg: MessageV2.User = {
  1021. id: Identifier.ascending("message"),
  1022. sessionID: input.sessionID,
  1023. time: {
  1024. created: Date.now(),
  1025. },
  1026. role: "user",
  1027. }
  1028. await updateMessage(userMsg)
  1029. const userPart: MessageV2.Part = {
  1030. type: "text",
  1031. id: Identifier.ascending("part"),
  1032. messageID: userMsg.id,
  1033. sessionID: input.sessionID,
  1034. text: "The following tool was executed by the user",
  1035. synthetic: true,
  1036. }
  1037. await updatePart(userPart)
  1038. const msg: MessageV2.Assistant = {
  1039. id: Identifier.ascending("message"),
  1040. sessionID: input.sessionID,
  1041. system: [],
  1042. mode: input.agent,
  1043. cost: 0,
  1044. path: {
  1045. cwd: Instance.directory,
  1046. root: Instance.worktree,
  1047. },
  1048. time: {
  1049. created: Date.now(),
  1050. },
  1051. role: "assistant",
  1052. tokens: {
  1053. input: 0,
  1054. output: 0,
  1055. reasoning: 0,
  1056. cache: { read: 0, write: 0 },
  1057. },
  1058. modelID: "",
  1059. providerID: "",
  1060. }
  1061. await updateMessage(msg)
  1062. const part: MessageV2.Part = {
  1063. type: "tool",
  1064. id: Identifier.ascending("part"),
  1065. messageID: msg.id,
  1066. sessionID: input.sessionID,
  1067. tool: "bash",
  1068. callID: ulid(),
  1069. state: {
  1070. status: "running",
  1071. time: {
  1072. start: Date.now(),
  1073. },
  1074. input: {
  1075. command: input.command,
  1076. },
  1077. },
  1078. }
  1079. await updatePart(part)
  1080. const shell = process.env["SHELL"] ?? "bash"
  1081. const shellName = path.basename(shell)
  1082. const scripts: Record<string, string> = {
  1083. nu: input.command,
  1084. fish: `eval "${input.command}"`,
  1085. }
  1086. const script =
  1087. scripts[shellName] ??
  1088. `[[ -f ~/.zshenv ]] && source ~/.zshenv >/dev/null 2>&1 || true
  1089. [[ -f "\${ZDOTDIR:-$HOME}/.zshrc" ]] && source "\${ZDOTDIR:-$HOME}/.zshrc" >/dev/null 2>&1 || true
  1090. [[ -f ~/.bashrc ]] && source ~/.bashrc >/dev/null 2>&1 || true
  1091. eval "${input.command}"`
  1092. const isFishOrNu = shellName === "fish" || shellName === "nu"
  1093. const args = isFishOrNu ? ["-c", script] : ["-c", "-l", script]
  1094. const proc = spawn(shell, args, {
  1095. cwd: Instance.directory,
  1096. signal: abort.signal,
  1097. detached: true,
  1098. stdio: ["ignore", "pipe", "pipe"],
  1099. env: {
  1100. ...process.env,
  1101. TERM: "dumb",
  1102. },
  1103. })
  1104. abort.signal.addEventListener("abort", () => {
  1105. if (!proc.pid) return
  1106. process.kill(-proc.pid)
  1107. })
  1108. let output = ""
  1109. proc.stdout?.on("data", (chunk) => {
  1110. output += chunk.toString()
  1111. if (part.state.status === "running") {
  1112. part.state.metadata = {
  1113. output: output,
  1114. description: "",
  1115. }
  1116. updatePart(part)
  1117. }
  1118. })
  1119. proc.stderr?.on("data", (chunk) => {
  1120. output += chunk.toString()
  1121. if (part.state.status === "running") {
  1122. part.state.metadata = {
  1123. output: output,
  1124. description: "",
  1125. }
  1126. updatePart(part)
  1127. }
  1128. })
  1129. await new Promise<void>((resolve) => {
  1130. proc.on("close", () => {
  1131. resolve()
  1132. })
  1133. })
  1134. msg.time.completed = Date.now()
  1135. await updateMessage(msg)
  1136. if (part.state.status === "running") {
  1137. part.state = {
  1138. status: "completed",
  1139. time: {
  1140. ...part.state.time,
  1141. end: Date.now(),
  1142. },
  1143. input: part.state.input,
  1144. title: "",
  1145. metadata: {
  1146. output,
  1147. description: "",
  1148. },
  1149. output,
  1150. }
  1151. await updatePart(part)
  1152. }
  1153. return { info: msg, parts: [part] }
  1154. }
  1155. export const CommandInput = z.object({
  1156. messageID: Identifier.schema("message").optional(),
  1157. sessionID: Identifier.schema("session"),
  1158. agent: z.string().optional(),
  1159. model: z.string().optional(),
  1160. arguments: z.string(),
  1161. command: z.string(),
  1162. })
  1163. export type CommandInput = z.infer<typeof CommandInput>
  1164. const bashRegex = /!`([^`]+)`/g
  1165. /**
  1166. * Regular expression to match @ file references in text
  1167. * Matches @ followed by file paths, excluding commas, periods at end of sentences, and backticks
  1168. * Does not match when preceded by word characters or backticks (to avoid email addresses and quoted references)
  1169. */
  1170. export const fileRegex = /(?<![\w`])@(\.?[^\s`,.]*(?:\.[^\s`,.]+)*)/g
  1171. export async function command(input: CommandInput) {
  1172. log.info("command", input)
  1173. const command = await Command.get(input.command)
  1174. const agent = command.agent ?? input.agent ?? "build"
  1175. let template = command.template.replace("$ARGUMENTS", input.arguments)
  1176. const bash = Array.from(template.matchAll(bashRegex))
  1177. if (bash.length > 0) {
  1178. const results = await Promise.all(
  1179. bash.map(async ([, cmd]) => {
  1180. try {
  1181. return await $`${{ raw: cmd }}`.nothrow().text()
  1182. } catch (error) {
  1183. return `Error executing command: ${error instanceof Error ? error.message : String(error)}`
  1184. }
  1185. }),
  1186. )
  1187. let index = 0
  1188. template = template.replace(bashRegex, () => results[index++])
  1189. }
  1190. const parts = [
  1191. {
  1192. type: "text",
  1193. text: template,
  1194. },
  1195. ] as ChatInput["parts"]
  1196. const matches = Array.from(template.matchAll(fileRegex))
  1197. await Promise.all(
  1198. matches.map(async (match) => {
  1199. const name = match[1]
  1200. const filepath = name.startsWith("~/")
  1201. ? path.join(os.homedir(), name.slice(2))
  1202. : path.resolve(Instance.worktree, name)
  1203. const stats = await fs.stat(filepath).catch(() => undefined)
  1204. if (!stats) {
  1205. const agent = await Agent.get(name)
  1206. if (agent) {
  1207. parts.push({
  1208. type: "agent",
  1209. name: agent.name,
  1210. })
  1211. }
  1212. return
  1213. }
  1214. if (stats.isDirectory()) return
  1215. parts.push({
  1216. type: "file",
  1217. url: `file://${filepath}`,
  1218. filename: name,
  1219. mime: "text/plain",
  1220. })
  1221. }),
  1222. )
  1223. return prompt({
  1224. sessionID: input.sessionID,
  1225. messageID: input.messageID,
  1226. model: (() => {
  1227. if (input.model) {
  1228. return Provider.parseModel(input.model)
  1229. }
  1230. if (command.model) {
  1231. return Provider.parseModel(command.model)
  1232. }
  1233. return undefined
  1234. })(),
  1235. agent,
  1236. parts,
  1237. })
  1238. }
  1239. function createProcessor(assistantMsg: MessageV2.Assistant, model: ModelsDev.Model) {
  1240. const toolcalls: Record<string, MessageV2.ToolPart> = {}
  1241. let snapshot: string | undefined
  1242. let shouldStop = false
  1243. return {
  1244. partFromToolCall(toolCallID: string) {
  1245. return toolcalls[toolCallID]
  1246. },
  1247. getShouldStop() {
  1248. return shouldStop
  1249. },
  1250. async process(stream: StreamTextResult<Record<string, AITool>, never>) {
  1251. try {
  1252. let currentText: MessageV2.TextPart | undefined
  1253. let reasoningMap: Record<string, MessageV2.ReasoningPart> = {}
  1254. for await (const value of stream.fullStream) {
  1255. log.info("part", {
  1256. type: value.type,
  1257. })
  1258. switch (value.type) {
  1259. case "start":
  1260. break
  1261. case "reasoning-start":
  1262. if (value.id in reasoningMap) {
  1263. continue
  1264. }
  1265. reasoningMap[value.id] = {
  1266. id: Identifier.ascending("part"),
  1267. messageID: assistantMsg.id,
  1268. sessionID: assistantMsg.sessionID,
  1269. type: "reasoning",
  1270. text: "",
  1271. time: {
  1272. start: Date.now(),
  1273. },
  1274. }
  1275. break
  1276. case "reasoning-delta":
  1277. if (value.id in reasoningMap) {
  1278. const part = reasoningMap[value.id]
  1279. part.text += value.text
  1280. if (part.text) await updatePart(part)
  1281. }
  1282. break
  1283. case "reasoning-end":
  1284. if (value.id in reasoningMap) {
  1285. const part = reasoningMap[value.id]
  1286. part.text = part.text.trimEnd()
  1287. part.metadata = value.providerMetadata
  1288. part.time = {
  1289. ...part.time,
  1290. end: Date.now(),
  1291. }
  1292. await updatePart(part)
  1293. delete reasoningMap[value.id]
  1294. }
  1295. break
  1296. case "tool-input-start":
  1297. const part = await updatePart({
  1298. id: toolcalls[value.id]?.id ?? Identifier.ascending("part"),
  1299. messageID: assistantMsg.id,
  1300. sessionID: assistantMsg.sessionID,
  1301. type: "tool",
  1302. tool: value.toolName,
  1303. callID: value.id,
  1304. state: {
  1305. status: "pending",
  1306. },
  1307. })
  1308. toolcalls[value.id] = part as MessageV2.ToolPart
  1309. break
  1310. case "tool-input-delta":
  1311. break
  1312. case "tool-input-end":
  1313. break
  1314. case "tool-call": {
  1315. const match = toolcalls[value.toolCallId]
  1316. if (match) {
  1317. const part = await updatePart({
  1318. ...match,
  1319. tool: value.toolName,
  1320. state: {
  1321. status: "running",
  1322. input: value.input,
  1323. time: {
  1324. start: Date.now(),
  1325. },
  1326. },
  1327. })
  1328. toolcalls[value.toolCallId] = part as MessageV2.ToolPart
  1329. }
  1330. break
  1331. }
  1332. case "tool-result": {
  1333. const match = toolcalls[value.toolCallId]
  1334. if (match && match.state.status === "running") {
  1335. await updatePart({
  1336. ...match,
  1337. state: {
  1338. status: "completed",
  1339. input: value.input,
  1340. output: value.output.output,
  1341. metadata: value.output.metadata,
  1342. title: value.output.title,
  1343. time: {
  1344. start: match.state.time.start,
  1345. end: Date.now(),
  1346. },
  1347. },
  1348. })
  1349. delete toolcalls[value.toolCallId]
  1350. }
  1351. break
  1352. }
  1353. case "tool-error": {
  1354. const match = toolcalls[value.toolCallId]
  1355. if (match && match.state.status === "running") {
  1356. if (value.error instanceof Permission.RejectedError) {
  1357. shouldStop = true
  1358. }
  1359. await updatePart({
  1360. ...match,
  1361. state: {
  1362. status: "error",
  1363. input: value.input,
  1364. error: (value.error as any).toString(),
  1365. metadata: value.error instanceof Permission.RejectedError ? value.error.metadata : undefined,
  1366. time: {
  1367. start: match.state.time.start,
  1368. end: Date.now(),
  1369. },
  1370. },
  1371. })
  1372. delete toolcalls[value.toolCallId]
  1373. }
  1374. break
  1375. }
  1376. case "error":
  1377. throw value.error
  1378. case "start-step":
  1379. await updatePart({
  1380. id: Identifier.ascending("part"),
  1381. messageID: assistantMsg.id,
  1382. sessionID: assistantMsg.sessionID,
  1383. type: "step-start",
  1384. })
  1385. snapshot = await Snapshot.track()
  1386. break
  1387. case "finish-step":
  1388. const usage = getUsage(model, value.usage, value.providerMetadata)
  1389. assistantMsg.cost += usage.cost
  1390. assistantMsg.tokens = usage.tokens
  1391. await updatePart({
  1392. id: Identifier.ascending("part"),
  1393. messageID: assistantMsg.id,
  1394. sessionID: assistantMsg.sessionID,
  1395. type: "step-finish",
  1396. tokens: usage.tokens,
  1397. cost: usage.cost,
  1398. })
  1399. await updateMessage(assistantMsg)
  1400. if (snapshot) {
  1401. const patch = await Snapshot.patch(snapshot)
  1402. if (patch.files.length) {
  1403. await updatePart({
  1404. id: Identifier.ascending("part"),
  1405. messageID: assistantMsg.id,
  1406. sessionID: assistantMsg.sessionID,
  1407. type: "patch",
  1408. hash: patch.hash,
  1409. files: patch.files,
  1410. })
  1411. }
  1412. snapshot = undefined
  1413. }
  1414. break
  1415. case "text-start":
  1416. currentText = {
  1417. id: Identifier.ascending("part"),
  1418. messageID: assistantMsg.id,
  1419. sessionID: assistantMsg.sessionID,
  1420. type: "text",
  1421. text: "",
  1422. time: {
  1423. start: Date.now(),
  1424. },
  1425. }
  1426. break
  1427. case "text-delta":
  1428. if (currentText) {
  1429. currentText.text += value.text
  1430. if (currentText.text) await updatePart(currentText)
  1431. }
  1432. break
  1433. case "text-end":
  1434. if (currentText) {
  1435. currentText.text = currentText.text.trimEnd()
  1436. currentText.time = {
  1437. start: Date.now(),
  1438. end: Date.now(),
  1439. }
  1440. await updatePart(currentText)
  1441. }
  1442. currentText = undefined
  1443. break
  1444. case "finish":
  1445. assistantMsg.time.completed = Date.now()
  1446. await updateMessage(assistantMsg)
  1447. break
  1448. default:
  1449. log.info("unhandled", {
  1450. ...value,
  1451. })
  1452. continue
  1453. }
  1454. }
  1455. } catch (e) {
  1456. log.error("", {
  1457. error: e,
  1458. })
  1459. switch (true) {
  1460. case e instanceof DOMException && e.name === "AbortError":
  1461. assistantMsg.error = new MessageV2.AbortedError(
  1462. { message: e.message },
  1463. {
  1464. cause: e,
  1465. },
  1466. ).toObject()
  1467. break
  1468. case MessageV2.OutputLengthError.isInstance(e):
  1469. assistantMsg.error = e
  1470. break
  1471. case LoadAPIKeyError.isInstance(e):
  1472. assistantMsg.error = new MessageV2.AuthError(
  1473. {
  1474. providerID: model.id,
  1475. message: e.message,
  1476. },
  1477. { cause: e },
  1478. ).toObject()
  1479. break
  1480. case e instanceof Error:
  1481. assistantMsg.error = new NamedError.Unknown({ message: e.toString() }, { cause: e }).toObject()
  1482. break
  1483. default:
  1484. assistantMsg.error = new NamedError.Unknown({ message: JSON.stringify(e) }, { cause: e })
  1485. }
  1486. Bus.publish(Event.Error, {
  1487. sessionID: assistantMsg.sessionID,
  1488. error: assistantMsg.error,
  1489. })
  1490. }
  1491. const p = await getParts(assistantMsg.id)
  1492. for (const part of p) {
  1493. if (part.type === "tool" && part.state.status !== "completed" && part.state.status !== "error") {
  1494. updatePart({
  1495. ...part,
  1496. state: {
  1497. status: "error",
  1498. error: "Tool execution aborted",
  1499. time: {
  1500. start: Date.now(),
  1501. end: Date.now(),
  1502. },
  1503. input: {},
  1504. },
  1505. })
  1506. }
  1507. }
  1508. assistantMsg.time.completed = Date.now()
  1509. await updateMessage(assistantMsg)
  1510. return { info: assistantMsg, parts: p }
  1511. },
  1512. }
  1513. }
  1514. export const RevertInput = z.object({
  1515. sessionID: Identifier.schema("session"),
  1516. messageID: Identifier.schema("message"),
  1517. partID: Identifier.schema("part").optional(),
  1518. })
  1519. export type RevertInput = z.infer<typeof RevertInput>
  1520. export async function revert(input: RevertInput) {
  1521. const all = await messages(input.sessionID)
  1522. let lastUser: MessageV2.User | undefined
  1523. const session = await get(input.sessionID)
  1524. let revert: Info["revert"]
  1525. const patches: Snapshot.Patch[] = []
  1526. for (const msg of all) {
  1527. if (msg.info.role === "user") lastUser = msg.info
  1528. const remaining = []
  1529. for (const part of msg.parts) {
  1530. if (revert) {
  1531. if (part.type === "patch") {
  1532. patches.push(part)
  1533. }
  1534. continue
  1535. }
  1536. if (!revert) {
  1537. if ((msg.info.id === input.messageID && !input.partID) || part.id === input.partID) {
  1538. // if no useful parts left in message, same as reverting whole message
  1539. const partID = remaining.some((item) => ["text", "tool"].includes(item.type)) ? input.partID : undefined
  1540. revert = {
  1541. messageID: !partID && lastUser ? lastUser.id : msg.info.id,
  1542. partID,
  1543. }
  1544. }
  1545. remaining.push(part)
  1546. }
  1547. }
  1548. }
  1549. if (revert) {
  1550. const session = await get(input.sessionID)
  1551. revert.snapshot = session.revert?.snapshot ?? (await Snapshot.track())
  1552. await Snapshot.revert(patches)
  1553. if (revert.snapshot) revert.diff = await Snapshot.diff(revert.snapshot)
  1554. return update(input.sessionID, (draft) => {
  1555. draft.revert = revert
  1556. })
  1557. }
  1558. return session
  1559. }
  1560. export async function unrevert(input: { sessionID: string }) {
  1561. log.info("unreverting", input)
  1562. const session = await get(input.sessionID)
  1563. if (!session.revert) return session
  1564. if (session.revert.snapshot) await Snapshot.restore(session.revert.snapshot)
  1565. const next = await update(input.sessionID, (draft) => {
  1566. draft.revert = undefined
  1567. })
  1568. return next
  1569. }
  1570. export async function summarize(input: { sessionID: string; providerID: string; modelID: string }) {
  1571. using abort = lock(input.sessionID)
  1572. const msgs = await messages(input.sessionID)
  1573. const lastSummary = msgs.findLast((msg) => msg.info.role === "assistant" && msg.info.summary === true)
  1574. const filtered = msgs.filter((msg) => !lastSummary || msg.info.id >= lastSummary.info.id)
  1575. const model = await Provider.getModel(input.providerID, input.modelID)
  1576. const system = [
  1577. ...SystemPrompt.summarize(model.providerID),
  1578. ...(await SystemPrompt.environment()),
  1579. ...(await SystemPrompt.custom()),
  1580. ]
  1581. const next: MessageV2.Info = {
  1582. id: Identifier.ascending("message"),
  1583. role: "assistant",
  1584. sessionID: input.sessionID,
  1585. system,
  1586. mode: "build",
  1587. path: {
  1588. cwd: Instance.directory,
  1589. root: Instance.worktree,
  1590. },
  1591. summary: true,
  1592. cost: 0,
  1593. modelID: input.modelID,
  1594. providerID: model.providerID,
  1595. tokens: {
  1596. input: 0,
  1597. output: 0,
  1598. reasoning: 0,
  1599. cache: { read: 0, write: 0 },
  1600. },
  1601. time: {
  1602. created: Date.now(),
  1603. },
  1604. }
  1605. await updateMessage(next)
  1606. const processor = createProcessor(next, model.info)
  1607. const stream = streamText({
  1608. maxRetries: 10,
  1609. abortSignal: abort.signal,
  1610. model: model.language,
  1611. messages: [
  1612. ...system.map(
  1613. (x): ModelMessage => ({
  1614. role: "system",
  1615. content: x,
  1616. }),
  1617. ),
  1618. ...MessageV2.toModelMessage(filtered),
  1619. {
  1620. role: "user",
  1621. content: [
  1622. {
  1623. type: "text",
  1624. text: "Provide a detailed but concise summary of our conversation above. Focus on information that would be helpful for continuing the conversation, including what we did, what we're doing, which files we're working on, and what we're going to do next.",
  1625. },
  1626. ],
  1627. },
  1628. ],
  1629. })
  1630. const result = await processor.process(stream)
  1631. return result
  1632. }
  1633. function isLocked(sessionID: string) {
  1634. return state().pending.has(sessionID)
  1635. }
  1636. function lock(sessionID: string) {
  1637. log.info("locking", { sessionID })
  1638. if (state().pending.has(sessionID)) throw new BusyError(sessionID)
  1639. const controller = new AbortController()
  1640. state().pending.set(sessionID, controller)
  1641. return {
  1642. signal: controller.signal,
  1643. async [Symbol.dispose]() {
  1644. log.info("unlocking", { sessionID })
  1645. state().pending.delete(sessionID)
  1646. const isAutoCompacting = state().autoCompacting.get(sessionID) ?? false
  1647. if (isAutoCompacting) {
  1648. state().autoCompacting.delete(sessionID)
  1649. return
  1650. }
  1651. const session = await get(sessionID)
  1652. if (session.parentID) return
  1653. Bus.publish(Event.Idle, {
  1654. sessionID,
  1655. })
  1656. },
  1657. }
  1658. }
  1659. function getUsage(model: ModelsDev.Model, usage: LanguageModelUsage, metadata?: ProviderMetadata) {
  1660. const tokens = {
  1661. input: usage.inputTokens ?? 0,
  1662. output: usage.outputTokens ?? 0,
  1663. reasoning: usage?.reasoningTokens ?? 0,
  1664. cache: {
  1665. write: (metadata?.["anthropic"]?.["cacheCreationInputTokens"] ??
  1666. // @ts-expect-error
  1667. metadata?.["bedrock"]?.["usage"]?.["cacheWriteInputTokens"] ??
  1668. 0) as number,
  1669. read: usage.cachedInputTokens ?? 0,
  1670. },
  1671. }
  1672. return {
  1673. cost: new Decimal(0)
  1674. .add(new Decimal(tokens.input).mul(model.cost?.input ?? 0).div(1_000_000))
  1675. .add(new Decimal(tokens.output).mul(model.cost?.output ?? 0).div(1_000_000))
  1676. .add(new Decimal(tokens.cache.read).mul(model.cost?.cache_read ?? 0).div(1_000_000))
  1677. .add(new Decimal(tokens.cache.write).mul(model.cost?.cache_write ?? 0).div(1_000_000))
  1678. .toNumber(),
  1679. tokens,
  1680. }
  1681. }
  1682. export class BusyError extends Error {
  1683. constructor(public readonly sessionID: string) {
  1684. super(`Session ${sessionID} is busy`)
  1685. }
  1686. }
  1687. export async function initialize(input: {
  1688. sessionID: string
  1689. modelID: string
  1690. providerID: string
  1691. messageID: string
  1692. }) {
  1693. await Session.prompt({
  1694. sessionID: input.sessionID,
  1695. messageID: input.messageID,
  1696. model: {
  1697. providerID: input.providerID,
  1698. modelID: input.modelID,
  1699. },
  1700. parts: [
  1701. {
  1702. id: Identifier.ascending("part"),
  1703. type: "text",
  1704. text: PROMPT_INITIALIZE.replace("${path}", Instance.worktree),
  1705. },
  1706. ],
  1707. })
  1708. await Project.setInitialized(Instance.project.id)
  1709. }
  1710. }