index.ts 54 KB

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