index.ts 26 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877
  1. import { Slug } from "@opencode-ai/util/slug"
  2. import path from "path"
  3. import { BusEvent } from "@/bus/bus-event"
  4. import { Bus } from "@/bus"
  5. import { Decimal } from "decimal.js"
  6. import z from "zod"
  7. import { type ProviderMetadata } from "ai"
  8. import { Config } from "../config/config"
  9. import { Flag } from "../flag/flag"
  10. import { Identifier } from "../id/id"
  11. import { Installation } from "../installation"
  12. import { Database, NotFoundError, eq, and, or, gte, isNull, desc, like, inArray, lt } from "../storage/db"
  13. import type { SQL } from "../storage/db"
  14. import { SessionTable, MessageTable, PartTable } from "./session.sql"
  15. import { ProjectTable } from "../project/project.sql"
  16. import { Storage } from "@/storage/storage"
  17. import { Log } from "../util/log"
  18. import { MessageV2 } from "./message-v2"
  19. import { Instance } from "../project/instance"
  20. import { SessionPrompt } from "./prompt"
  21. import { fn } from "@/util/fn"
  22. import { Command } from "../command"
  23. import { Snapshot } from "@/snapshot"
  24. import type { Provider } from "@/provider/provider"
  25. import { PermissionNext } from "@/permission/next"
  26. import { Global } from "@/global"
  27. import type { LanguageModelV2Usage } from "@ai-sdk/provider"
  28. import { iife } from "@/util/iife"
  29. export namespace Session {
  30. const log = Log.create({ service: "session" })
  31. const parentTitlePrefix = "New session - "
  32. const childTitlePrefix = "Child session - "
  33. function createDefaultTitle(isChild = false) {
  34. return (isChild ? childTitlePrefix : parentTitlePrefix) + new Date().toISOString()
  35. }
  36. export function isDefaultTitle(title: string) {
  37. return new RegExp(
  38. `^(${parentTitlePrefix}|${childTitlePrefix})\\d{4}-\\d{2}-\\d{2}T\\d{2}:\\d{2}:\\d{2}\\.\\d{3}Z$`,
  39. ).test(title)
  40. }
  41. type SessionRow = typeof SessionTable.$inferSelect
  42. export function fromRow(row: SessionRow): Info {
  43. const summary =
  44. row.summary_additions !== null || row.summary_deletions !== null || row.summary_files !== null
  45. ? {
  46. additions: row.summary_additions ?? 0,
  47. deletions: row.summary_deletions ?? 0,
  48. files: row.summary_files ?? 0,
  49. diffs: row.summary_diffs ?? undefined,
  50. }
  51. : undefined
  52. const share = row.share_url ? { url: row.share_url } : undefined
  53. const revert = row.revert ?? undefined
  54. return {
  55. id: row.id,
  56. slug: row.slug,
  57. projectID: row.project_id,
  58. directory: row.directory,
  59. parentID: row.parent_id ?? undefined,
  60. title: row.title,
  61. version: row.version,
  62. summary,
  63. share,
  64. revert,
  65. permission: row.permission ?? undefined,
  66. time: {
  67. created: row.time_created,
  68. updated: row.time_updated,
  69. compacting: row.time_compacting ?? undefined,
  70. archived: row.time_archived ?? undefined,
  71. },
  72. }
  73. }
  74. export function toRow(info: Info) {
  75. return {
  76. id: info.id,
  77. project_id: info.projectID,
  78. parent_id: info.parentID,
  79. slug: info.slug,
  80. directory: info.directory,
  81. title: info.title,
  82. version: info.version,
  83. share_url: info.share?.url,
  84. summary_additions: info.summary?.additions,
  85. summary_deletions: info.summary?.deletions,
  86. summary_files: info.summary?.files,
  87. summary_diffs: info.summary?.diffs,
  88. revert: info.revert ?? null,
  89. permission: info.permission,
  90. time_created: info.time.created,
  91. time_updated: info.time.updated,
  92. time_compacting: info.time.compacting,
  93. time_archived: info.time.archived,
  94. }
  95. }
  96. function getForkedTitle(title: string): string {
  97. const match = title.match(/^(.+) \(fork #(\d+)\)$/)
  98. if (match) {
  99. const base = match[1]
  100. const num = parseInt(match[2], 10)
  101. return `${base} (fork #${num + 1})`
  102. }
  103. return `${title} (fork #1)`
  104. }
  105. export const Info = z
  106. .object({
  107. id: Identifier.schema("session"),
  108. slug: z.string(),
  109. projectID: z.string(),
  110. directory: z.string(),
  111. parentID: Identifier.schema("session").optional(),
  112. summary: z
  113. .object({
  114. additions: z.number(),
  115. deletions: z.number(),
  116. files: z.number(),
  117. diffs: Snapshot.FileDiff.array().optional(),
  118. })
  119. .optional(),
  120. share: z
  121. .object({
  122. url: z.string(),
  123. })
  124. .optional(),
  125. title: z.string(),
  126. version: z.string(),
  127. time: z.object({
  128. created: z.number(),
  129. updated: z.number(),
  130. compacting: z.number().optional(),
  131. archived: z.number().optional(),
  132. }),
  133. permission: PermissionNext.Ruleset.optional(),
  134. revert: z
  135. .object({
  136. messageID: z.string(),
  137. partID: z.string().optional(),
  138. snapshot: z.string().optional(),
  139. diff: z.string().optional(),
  140. })
  141. .optional(),
  142. })
  143. .meta({
  144. ref: "Session",
  145. })
  146. export type Info = z.output<typeof Info>
  147. export const ProjectInfo = z
  148. .object({
  149. id: z.string(),
  150. name: z.string().optional(),
  151. worktree: z.string(),
  152. })
  153. .meta({
  154. ref: "ProjectSummary",
  155. })
  156. export type ProjectInfo = z.output<typeof ProjectInfo>
  157. export const GlobalInfo = Info.extend({
  158. project: ProjectInfo.nullable(),
  159. }).meta({
  160. ref: "GlobalSession",
  161. })
  162. export type GlobalInfo = z.output<typeof GlobalInfo>
  163. export const Event = {
  164. Created: BusEvent.define(
  165. "session.created",
  166. z.object({
  167. info: Info,
  168. }),
  169. ),
  170. Updated: BusEvent.define(
  171. "session.updated",
  172. z.object({
  173. info: Info,
  174. }),
  175. ),
  176. Deleted: BusEvent.define(
  177. "session.deleted",
  178. z.object({
  179. info: Info,
  180. }),
  181. ),
  182. Diff: BusEvent.define(
  183. "session.diff",
  184. z.object({
  185. sessionID: z.string(),
  186. diff: Snapshot.FileDiff.array(),
  187. }),
  188. ),
  189. Error: BusEvent.define(
  190. "session.error",
  191. z.object({
  192. sessionID: z.string().optional(),
  193. error: MessageV2.Assistant.shape.error,
  194. }),
  195. ),
  196. }
  197. export const create = fn(
  198. z
  199. .object({
  200. parentID: Identifier.schema("session").optional(),
  201. title: z.string().optional(),
  202. permission: Info.shape.permission,
  203. })
  204. .optional(),
  205. async (input) => {
  206. return createNext({
  207. parentID: input?.parentID,
  208. directory: Instance.directory,
  209. title: input?.title,
  210. permission: input?.permission,
  211. })
  212. },
  213. )
  214. export const fork = fn(
  215. z.object({
  216. sessionID: Identifier.schema("session"),
  217. messageID: Identifier.schema("message").optional(),
  218. }),
  219. async (input) => {
  220. const original = await get(input.sessionID)
  221. if (!original) throw new Error("session not found")
  222. const title = getForkedTitle(original.title)
  223. const session = await createNext({
  224. directory: Instance.directory,
  225. title,
  226. })
  227. const msgs = await messages({ sessionID: input.sessionID })
  228. const idMap = new Map<string, string>()
  229. for (const msg of msgs) {
  230. if (input.messageID && msg.info.id >= input.messageID) break
  231. const newID = Identifier.ascending("message")
  232. idMap.set(msg.info.id, newID)
  233. const parentID = msg.info.role === "assistant" && msg.info.parentID ? idMap.get(msg.info.parentID) : undefined
  234. const cloned = await updateMessage({
  235. ...msg.info,
  236. sessionID: session.id,
  237. id: newID,
  238. ...(parentID && { parentID }),
  239. })
  240. for (const part of msg.parts) {
  241. await updatePart({
  242. ...part,
  243. id: Identifier.ascending("part"),
  244. messageID: cloned.id,
  245. sessionID: session.id,
  246. })
  247. }
  248. }
  249. return session
  250. },
  251. )
  252. export const touch = fn(Identifier.schema("session"), async (sessionID) => {
  253. const now = Date.now()
  254. Database.use((db) => {
  255. const row = db
  256. .update(SessionTable)
  257. .set({ time_updated: now })
  258. .where(eq(SessionTable.id, sessionID))
  259. .returning()
  260. .get()
  261. if (!row) throw new NotFoundError({ message: `Session not found: ${sessionID}` })
  262. const info = fromRow(row)
  263. Database.effect(() => Bus.publish(Event.Updated, { info }))
  264. })
  265. })
  266. export async function createNext(input: {
  267. id?: string
  268. title?: string
  269. parentID?: string
  270. directory: string
  271. permission?: PermissionNext.Ruleset
  272. }) {
  273. const result: Info = {
  274. id: Identifier.descending("session", input.id),
  275. slug: Slug.create(),
  276. version: Installation.VERSION,
  277. projectID: Instance.project.id,
  278. directory: input.directory,
  279. parentID: input.parentID,
  280. title: input.title ?? createDefaultTitle(!!input.parentID),
  281. permission: input.permission,
  282. time: {
  283. created: Date.now(),
  284. updated: Date.now(),
  285. },
  286. }
  287. log.info("created", result)
  288. Database.use((db) => {
  289. db.insert(SessionTable).values(toRow(result)).run()
  290. Database.effect(() =>
  291. Bus.publish(Event.Created, {
  292. info: result,
  293. }),
  294. )
  295. })
  296. const cfg = await Config.get()
  297. if (!result.parentID && (Flag.OPENCODE_AUTO_SHARE || cfg.share === "auto"))
  298. share(result.id).catch(() => {
  299. // Silently ignore sharing errors during session creation
  300. })
  301. Bus.publish(Event.Updated, {
  302. info: result,
  303. })
  304. return result
  305. }
  306. export function plan(input: { slug: string; time: { created: number } }) {
  307. const base = Instance.project.vcs
  308. ? path.join(Instance.worktree, ".opencode", "plans")
  309. : path.join(Global.Path.data, "plans")
  310. return path.join(base, [input.time.created, input.slug].join("-") + ".md")
  311. }
  312. export const get = fn(Identifier.schema("session"), async (id) => {
  313. const row = Database.use((db) => db.select().from(SessionTable).where(eq(SessionTable.id, id)).get())
  314. if (!row) throw new NotFoundError({ message: `Session not found: ${id}` })
  315. return fromRow(row)
  316. })
  317. export const share = fn(Identifier.schema("session"), async (id) => {
  318. const cfg = await Config.get()
  319. if (cfg.share === "disabled") {
  320. throw new Error("Sharing is disabled in configuration")
  321. }
  322. const { ShareNext } = await import("@/share/share-next")
  323. const share = await ShareNext.create(id)
  324. Database.use((db) => {
  325. const row = db.update(SessionTable).set({ share_url: share.url }).where(eq(SessionTable.id, id)).returning().get()
  326. if (!row) throw new NotFoundError({ message: `Session not found: ${id}` })
  327. const info = fromRow(row)
  328. Database.effect(() => Bus.publish(Event.Updated, { info }))
  329. })
  330. return share
  331. })
  332. export const unshare = fn(Identifier.schema("session"), async (id) => {
  333. // Use ShareNext to remove the share (same as share function uses ShareNext to create)
  334. const { ShareNext } = await import("@/share/share-next")
  335. await ShareNext.remove(id)
  336. Database.use((db) => {
  337. const row = db.update(SessionTable).set({ share_url: null }).where(eq(SessionTable.id, id)).returning().get()
  338. if (!row) throw new NotFoundError({ message: `Session not found: ${id}` })
  339. const info = fromRow(row)
  340. Database.effect(() => Bus.publish(Event.Updated, { info }))
  341. })
  342. })
  343. export const setTitle = fn(
  344. z.object({
  345. sessionID: Identifier.schema("session"),
  346. title: z.string(),
  347. }),
  348. async (input) => {
  349. return Database.use((db) => {
  350. const row = db
  351. .update(SessionTable)
  352. .set({ title: input.title })
  353. .where(eq(SessionTable.id, input.sessionID))
  354. .returning()
  355. .get()
  356. if (!row) throw new NotFoundError({ message: `Session not found: ${input.sessionID}` })
  357. const info = fromRow(row)
  358. Database.effect(() => Bus.publish(Event.Updated, { info }))
  359. return info
  360. })
  361. },
  362. )
  363. export const setArchived = fn(
  364. z.object({
  365. sessionID: Identifier.schema("session"),
  366. time: z.number().optional(),
  367. }),
  368. async (input) => {
  369. return Database.use((db) => {
  370. const row = db
  371. .update(SessionTable)
  372. .set({ time_archived: input.time })
  373. .where(eq(SessionTable.id, input.sessionID))
  374. .returning()
  375. .get()
  376. if (!row) throw new NotFoundError({ message: `Session not found: ${input.sessionID}` })
  377. const info = fromRow(row)
  378. Database.effect(() => Bus.publish(Event.Updated, { info }))
  379. return info
  380. })
  381. },
  382. )
  383. export const setPermission = fn(
  384. z.object({
  385. sessionID: Identifier.schema("session"),
  386. permission: PermissionNext.Ruleset,
  387. }),
  388. async (input) => {
  389. return Database.use((db) => {
  390. const row = db
  391. .update(SessionTable)
  392. .set({ permission: input.permission, time_updated: Date.now() })
  393. .where(eq(SessionTable.id, input.sessionID))
  394. .returning()
  395. .get()
  396. if (!row) throw new NotFoundError({ message: `Session not found: ${input.sessionID}` })
  397. const info = fromRow(row)
  398. Database.effect(() => Bus.publish(Event.Updated, { info }))
  399. return info
  400. })
  401. },
  402. )
  403. export const setRevert = fn(
  404. z.object({
  405. sessionID: Identifier.schema("session"),
  406. revert: Info.shape.revert,
  407. summary: Info.shape.summary,
  408. }),
  409. async (input) => {
  410. return Database.use((db) => {
  411. const row = db
  412. .update(SessionTable)
  413. .set({
  414. revert: input.revert ?? null,
  415. summary_additions: input.summary?.additions,
  416. summary_deletions: input.summary?.deletions,
  417. summary_files: input.summary?.files,
  418. time_updated: Date.now(),
  419. })
  420. .where(eq(SessionTable.id, input.sessionID))
  421. .returning()
  422. .get()
  423. if (!row) throw new NotFoundError({ message: `Session not found: ${input.sessionID}` })
  424. const info = fromRow(row)
  425. Database.effect(() => Bus.publish(Event.Updated, { info }))
  426. return info
  427. })
  428. },
  429. )
  430. export const clearRevert = fn(Identifier.schema("session"), async (sessionID) => {
  431. return Database.use((db) => {
  432. const row = db
  433. .update(SessionTable)
  434. .set({
  435. revert: null,
  436. time_updated: Date.now(),
  437. })
  438. .where(eq(SessionTable.id, sessionID))
  439. .returning()
  440. .get()
  441. if (!row) throw new NotFoundError({ message: `Session not found: ${sessionID}` })
  442. const info = fromRow(row)
  443. Database.effect(() => Bus.publish(Event.Updated, { info }))
  444. return info
  445. })
  446. })
  447. export const setSummary = fn(
  448. z.object({
  449. sessionID: Identifier.schema("session"),
  450. summary: Info.shape.summary,
  451. }),
  452. async (input) => {
  453. return Database.use((db) => {
  454. const row = db
  455. .update(SessionTable)
  456. .set({
  457. summary_additions: input.summary?.additions,
  458. summary_deletions: input.summary?.deletions,
  459. summary_files: input.summary?.files,
  460. time_updated: Date.now(),
  461. })
  462. .where(eq(SessionTable.id, input.sessionID))
  463. .returning()
  464. .get()
  465. if (!row) throw new NotFoundError({ message: `Session not found: ${input.sessionID}` })
  466. const info = fromRow(row)
  467. Database.effect(() => Bus.publish(Event.Updated, { info }))
  468. return info
  469. })
  470. },
  471. )
  472. export const diff = fn(Identifier.schema("session"), async (sessionID) => {
  473. try {
  474. return await Storage.read<Snapshot.FileDiff[]>(["session_diff", sessionID])
  475. } catch {
  476. return []
  477. }
  478. })
  479. export const messages = fn(
  480. z.object({
  481. sessionID: Identifier.schema("session"),
  482. limit: z.number().optional(),
  483. }),
  484. async (input) => {
  485. const result = [] as MessageV2.WithParts[]
  486. for await (const msg of MessageV2.stream(input.sessionID)) {
  487. if (input.limit && result.length >= input.limit) break
  488. result.push(msg)
  489. }
  490. result.reverse()
  491. return result
  492. },
  493. )
  494. export function* list(input?: {
  495. directory?: string
  496. roots?: boolean
  497. start?: number
  498. search?: string
  499. limit?: number
  500. }) {
  501. const project = Instance.project
  502. const conditions = [eq(SessionTable.project_id, project.id)]
  503. if (input?.directory) {
  504. conditions.push(eq(SessionTable.directory, input.directory))
  505. }
  506. if (input?.roots) {
  507. conditions.push(isNull(SessionTable.parent_id))
  508. }
  509. if (input?.start) {
  510. conditions.push(gte(SessionTable.time_updated, input.start))
  511. }
  512. if (input?.search) {
  513. conditions.push(like(SessionTable.title, `%${input.search}%`))
  514. }
  515. const limit = input?.limit ?? 100
  516. const rows = Database.use((db) =>
  517. db
  518. .select()
  519. .from(SessionTable)
  520. .where(and(...conditions))
  521. .orderBy(desc(SessionTable.time_updated))
  522. .limit(limit)
  523. .all(),
  524. )
  525. for (const row of rows) {
  526. yield fromRow(row)
  527. }
  528. }
  529. export function* listGlobal(input?: {
  530. directory?: string
  531. roots?: boolean
  532. start?: number
  533. cursor?: number
  534. search?: string
  535. limit?: number
  536. archived?: boolean
  537. }) {
  538. const conditions: SQL[] = []
  539. if (input?.directory) {
  540. conditions.push(eq(SessionTable.directory, input.directory))
  541. }
  542. if (input?.roots) {
  543. conditions.push(isNull(SessionTable.parent_id))
  544. }
  545. if (input?.start) {
  546. conditions.push(gte(SessionTable.time_updated, input.start))
  547. }
  548. if (input?.cursor) {
  549. conditions.push(lt(SessionTable.time_updated, input.cursor))
  550. }
  551. if (input?.search) {
  552. conditions.push(like(SessionTable.title, `%${input.search}%`))
  553. }
  554. if (!input?.archived) {
  555. conditions.push(isNull(SessionTable.time_archived))
  556. }
  557. const limit = input?.limit ?? 100
  558. const rows = Database.use((db) => {
  559. const query =
  560. conditions.length > 0
  561. ? db
  562. .select()
  563. .from(SessionTable)
  564. .where(and(...conditions))
  565. : db.select().from(SessionTable)
  566. return query.orderBy(desc(SessionTable.time_updated), desc(SessionTable.id)).limit(limit).all()
  567. })
  568. const ids = [...new Set(rows.map((row) => row.project_id))]
  569. const projects = new Map<string, ProjectInfo>()
  570. if (ids.length > 0) {
  571. const items = Database.use((db) =>
  572. db
  573. .select({ id: ProjectTable.id, name: ProjectTable.name, worktree: ProjectTable.worktree })
  574. .from(ProjectTable)
  575. .where(inArray(ProjectTable.id, ids))
  576. .all(),
  577. )
  578. for (const item of items) {
  579. projects.set(item.id, {
  580. id: item.id,
  581. name: item.name ?? undefined,
  582. worktree: item.worktree,
  583. })
  584. }
  585. }
  586. for (const row of rows) {
  587. const project = projects.get(row.project_id) ?? null
  588. yield { ...fromRow(row), project }
  589. }
  590. }
  591. export const children = fn(Identifier.schema("session"), async (parentID) => {
  592. const project = Instance.project
  593. const rows = Database.use((db) =>
  594. db
  595. .select()
  596. .from(SessionTable)
  597. .where(and(eq(SessionTable.project_id, project.id), eq(SessionTable.parent_id, parentID)))
  598. .all(),
  599. )
  600. return rows.map(fromRow)
  601. })
  602. export const remove = fn(Identifier.schema("session"), async (sessionID) => {
  603. const project = Instance.project
  604. try {
  605. const session = await get(sessionID)
  606. for (const child of await children(sessionID)) {
  607. await remove(child.id)
  608. }
  609. await unshare(sessionID).catch(() => {})
  610. // CASCADE delete handles messages and parts automatically
  611. Database.use((db) => {
  612. db.delete(SessionTable).where(eq(SessionTable.id, sessionID)).run()
  613. Database.effect(() =>
  614. Bus.publish(Event.Deleted, {
  615. info: session,
  616. }),
  617. )
  618. })
  619. } catch (e) {
  620. log.error(e)
  621. }
  622. })
  623. export const updateMessage = fn(MessageV2.Info, async (msg) => {
  624. const time_created = msg.time.created
  625. const { id, sessionID, ...data } = msg
  626. Database.use((db) => {
  627. db.insert(MessageTable)
  628. .values({
  629. id,
  630. session_id: sessionID,
  631. time_created,
  632. data,
  633. })
  634. .onConflictDoUpdate({ target: MessageTable.id, set: { data } })
  635. .run()
  636. Database.effect(() =>
  637. Bus.publish(MessageV2.Event.Updated, {
  638. info: msg,
  639. }),
  640. )
  641. })
  642. return msg
  643. })
  644. export const removeMessage = fn(
  645. z.object({
  646. sessionID: Identifier.schema("session"),
  647. messageID: Identifier.schema("message"),
  648. }),
  649. async (input) => {
  650. // CASCADE delete handles parts automatically
  651. Database.use((db) => {
  652. db.delete(MessageTable)
  653. .where(and(eq(MessageTable.id, input.messageID), eq(MessageTable.session_id, input.sessionID)))
  654. .run()
  655. Database.effect(() =>
  656. Bus.publish(MessageV2.Event.Removed, {
  657. sessionID: input.sessionID,
  658. messageID: input.messageID,
  659. }),
  660. )
  661. })
  662. return input.messageID
  663. },
  664. )
  665. export const removePart = fn(
  666. z.object({
  667. sessionID: Identifier.schema("session"),
  668. messageID: Identifier.schema("message"),
  669. partID: Identifier.schema("part"),
  670. }),
  671. async (input) => {
  672. Database.use((db) => {
  673. db.delete(PartTable)
  674. .where(and(eq(PartTable.id, input.partID), eq(PartTable.session_id, input.sessionID)))
  675. .run()
  676. Database.effect(() =>
  677. Bus.publish(MessageV2.Event.PartRemoved, {
  678. sessionID: input.sessionID,
  679. messageID: input.messageID,
  680. partID: input.partID,
  681. }),
  682. )
  683. })
  684. return input.partID
  685. },
  686. )
  687. const UpdatePartInput = MessageV2.Part
  688. export const updatePart = fn(UpdatePartInput, async (part) => {
  689. const { id, messageID, sessionID, ...data } = part
  690. const time = Date.now()
  691. Database.use((db) => {
  692. db.insert(PartTable)
  693. .values({
  694. id,
  695. message_id: messageID,
  696. session_id: sessionID,
  697. time_created: time,
  698. data,
  699. })
  700. .onConflictDoUpdate({ target: PartTable.id, set: { data } })
  701. .run()
  702. Database.effect(() =>
  703. Bus.publish(MessageV2.Event.PartUpdated, {
  704. part,
  705. }),
  706. )
  707. })
  708. return part
  709. })
  710. export const updatePartDelta = fn(
  711. z.object({
  712. sessionID: z.string(),
  713. messageID: z.string(),
  714. partID: z.string(),
  715. field: z.string(),
  716. delta: z.string(),
  717. }),
  718. async (input) => {
  719. Bus.publish(MessageV2.Event.PartDelta, input)
  720. },
  721. )
  722. export const getUsage = fn(
  723. z.object({
  724. model: z.custom<Provider.Model>(),
  725. usage: z.custom<LanguageModelV2Usage>(),
  726. metadata: z.custom<ProviderMetadata>().optional(),
  727. }),
  728. (input) => {
  729. const safe = (value: number) => {
  730. if (!Number.isFinite(value)) return 0
  731. return value
  732. }
  733. const inputTokens = safe(input.usage.inputTokens ?? 0)
  734. const outputTokens = safe(input.usage.outputTokens ?? 0)
  735. const reasoningTokens = safe(input.usage.reasoningTokens ?? 0)
  736. const cacheReadInputTokens = safe(input.usage.cachedInputTokens ?? 0)
  737. const cacheWriteInputTokens = safe(
  738. (input.metadata?.["anthropic"]?.["cacheCreationInputTokens"] ??
  739. // @ts-expect-error
  740. input.metadata?.["bedrock"]?.["usage"]?.["cacheWriteInputTokens"] ??
  741. // @ts-expect-error
  742. input.metadata?.["venice"]?.["usage"]?.["cacheCreationInputTokens"] ??
  743. 0) as number,
  744. )
  745. // OpenRouter provides inputTokens as the total count of input tokens (including cached).
  746. // AFAIK other providers (OpenRouter/OpenAI/Gemini etc.) do it the same way e.g. vercel/ai#8794 (comment)
  747. // Anthropic does it differently though - inputTokens doesn't include cached tokens.
  748. // It looks like OpenCode's cost calculation assumes all providers return inputTokens the same way Anthropic does (I'm guessing getUsage logic was originally implemented with anthropic), so it's causing incorrect cost calculation for OpenRouter and others.
  749. const excludesCachedTokens = !!(input.metadata?.["anthropic"] || input.metadata?.["bedrock"])
  750. const adjustedInputTokens = safe(
  751. excludesCachedTokens ? inputTokens : inputTokens - cacheReadInputTokens - cacheWriteInputTokens,
  752. )
  753. const total = iife(() => {
  754. // Anthropic doesn't provide total_tokens, also ai sdk will vastly undercount if we
  755. // don't compute from components
  756. if (
  757. input.model.api.npm === "@ai-sdk/anthropic" ||
  758. input.model.api.npm === "@ai-sdk/amazon-bedrock" ||
  759. input.model.api.npm === "@ai-sdk/google-vertex/anthropic"
  760. ) {
  761. return adjustedInputTokens + outputTokens + cacheReadInputTokens + cacheWriteInputTokens
  762. }
  763. return input.usage.totalTokens
  764. })
  765. const tokens = {
  766. total,
  767. input: adjustedInputTokens,
  768. output: outputTokens,
  769. reasoning: reasoningTokens,
  770. cache: {
  771. write: cacheWriteInputTokens,
  772. read: cacheReadInputTokens,
  773. },
  774. }
  775. const costInfo =
  776. input.model.cost?.experimentalOver200K && tokens.input + tokens.cache.read > 200_000
  777. ? input.model.cost.experimentalOver200K
  778. : input.model.cost
  779. return {
  780. cost: safe(
  781. new Decimal(0)
  782. .add(new Decimal(tokens.input).mul(costInfo?.input ?? 0).div(1_000_000))
  783. .add(new Decimal(tokens.output).mul(costInfo?.output ?? 0).div(1_000_000))
  784. .add(new Decimal(tokens.cache.read).mul(costInfo?.cache?.read ?? 0).div(1_000_000))
  785. .add(new Decimal(tokens.cache.write).mul(costInfo?.cache?.write ?? 0).div(1_000_000))
  786. // TODO: update models.dev to have better pricing model, for now:
  787. // charge reasoning tokens at the same rate as output tokens
  788. .add(new Decimal(tokens.reasoning).mul(costInfo?.output ?? 0).div(1_000_000))
  789. .toNumber(),
  790. ),
  791. tokens,
  792. }
  793. },
  794. )
  795. export class BusyError extends Error {
  796. constructor(public readonly sessionID: string) {
  797. super(`Session ${sessionID} is busy`)
  798. }
  799. }
  800. export const initialize = fn(
  801. z.object({
  802. sessionID: Identifier.schema("session"),
  803. modelID: z.string(),
  804. providerID: z.string(),
  805. messageID: Identifier.schema("message"),
  806. }),
  807. async (input) => {
  808. await SessionPrompt.command({
  809. sessionID: input.sessionID,
  810. messageID: input.messageID,
  811. model: input.providerID + "/" + input.modelID,
  812. command: Command.Default.INIT,
  813. arguments: "",
  814. })
  815. },
  816. )
  817. }