sql.ts 6.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178
  1. import { sqliteTable, text, integer, index, primaryKey, real, uniqueIndex } from "drizzle-orm/sqlite-core"
  2. import * as DatabasePath from "../database/path"
  3. import { ProjectTable } from "../project/sql"
  4. import type { SessionMessage } from "./message"
  5. import type { Prompt } from "./prompt"
  6. import type { SessionInput } from "./input"
  7. import type { Snapshot } from "../snapshot"
  8. import { PermissionV1 } from "../v1/permission"
  9. import { ProjectV2 } from "../project"
  10. import type { SessionSchema } from "./schema"
  11. import type { MessageID, PartID, SessionV1 } from "../v1/session"
  12. import { WorkspaceV2 } from "../workspace"
  13. import { Timestamps } from "../database/schema.sql"
  14. import type { SystemContext } from "../system-context/index"
  15. import { AgentV2 } from "../agent"
  16. type SessionMessageData = Omit<(typeof SessionMessage.Message)["Encoded"], "type" | "id">
  17. type V1MessageData = Omit<SessionV1.Info, "id" | "sessionID">
  18. type V1PartData = Omit<SessionV1.Part, "id" | "sessionID" | "messageID">
  19. export const SessionTable = sqliteTable(
  20. "session",
  21. {
  22. id: text().$type<SessionSchema.ID>().primaryKey(),
  23. project_id: text()
  24. .$type<ProjectV2.ID>()
  25. .notNull()
  26. .references(() => ProjectTable.id, { onDelete: "cascade" }),
  27. workspace_id: text().$type<WorkspaceV2.ID>(),
  28. parent_id: text().$type<SessionSchema.ID>(),
  29. slug: text().notNull(),
  30. directory: DatabasePath.directoryColumn().notNull(),
  31. path: DatabasePath.pathColumn(),
  32. title: text().notNull(),
  33. version: text().notNull(),
  34. share_url: text(),
  35. summary_additions: integer(),
  36. summary_deletions: integer(),
  37. summary_files: integer(),
  38. summary_diffs: text({ mode: "json" }).$type<Snapshot.FileDiff[]>(),
  39. metadata: text({ mode: "json" }).$type<Record<string, unknown>>(),
  40. cost: real().notNull().default(0),
  41. tokens_input: integer().notNull().default(0),
  42. tokens_output: integer().notNull().default(0),
  43. tokens_reasoning: integer().notNull().default(0),
  44. tokens_cache_read: integer().notNull().default(0),
  45. tokens_cache_write: integer().notNull().default(0),
  46. revert: text({ mode: "json" }).$type<{ messageID: MessageID; partID?: PartID; snapshot?: string; diff?: string }>(),
  47. permission: text({ mode: "json" }).$type<PermissionV1.Ruleset>(),
  48. agent: text(),
  49. model: text({ mode: "json" }).$type<{
  50. id: string
  51. providerID: string
  52. variant?: string
  53. }>(),
  54. ...Timestamps,
  55. time_compacting: integer(),
  56. time_archived: integer(),
  57. },
  58. (table) => [
  59. index("session_project_idx").on(table.project_id),
  60. index("session_workspace_idx").on(table.workspace_id),
  61. index("session_parent_idx").on(table.parent_id),
  62. ],
  63. )
  64. export const MessageTable = sqliteTable(
  65. "message",
  66. {
  67. id: text().$type<MessageID>().primaryKey(),
  68. session_id: text()
  69. .$type<SessionSchema.ID>()
  70. .notNull()
  71. .references(() => SessionTable.id, { onDelete: "cascade" }),
  72. ...Timestamps,
  73. data: text({ mode: "json" }).notNull().$type<V1MessageData>(),
  74. },
  75. (table) => [index("message_session_time_created_id_idx").on(table.session_id, table.time_created, table.id)],
  76. )
  77. export const PartTable = sqliteTable(
  78. "part",
  79. {
  80. id: text().$type<PartID>().primaryKey(),
  81. message_id: text()
  82. .$type<MessageID>()
  83. .notNull()
  84. .references(() => MessageTable.id, { onDelete: "cascade" }),
  85. session_id: text().$type<SessionSchema.ID>().notNull(),
  86. ...Timestamps,
  87. data: text({ mode: "json" }).notNull().$type<V1PartData>(),
  88. },
  89. (table) => [
  90. index("part_message_id_id_idx").on(table.message_id, table.id),
  91. index("part_session_idx").on(table.session_id),
  92. ],
  93. )
  94. export const TodoTable = sqliteTable(
  95. "todo",
  96. {
  97. session_id: text()
  98. .$type<SessionSchema.ID>()
  99. .notNull()
  100. .references(() => SessionTable.id, { onDelete: "cascade" }),
  101. content: text().notNull(),
  102. status: text().notNull(),
  103. priority: text().notNull(),
  104. position: integer().notNull(),
  105. ...Timestamps,
  106. },
  107. (table) => [
  108. primaryKey({ columns: [table.session_id, table.position] }),
  109. index("todo_session_idx").on(table.session_id),
  110. ],
  111. )
  112. export const SessionMessageTable = sqliteTable(
  113. "session_message",
  114. {
  115. id: text().$type<SessionMessage.ID>().primaryKey(),
  116. session_id: text()
  117. .$type<SessionSchema.ID>()
  118. .notNull()
  119. .references(() => SessionTable.id, { onDelete: "cascade" }),
  120. type: text().$type<SessionMessage.Type>().notNull(),
  121. seq: integer().notNull(),
  122. ...Timestamps,
  123. data: text({ mode: "json" }).notNull().$type<SessionMessageData>(),
  124. },
  125. (table) => [
  126. uniqueIndex("session_message_session_seq_idx").on(table.session_id, table.seq),
  127. index("session_message_session_type_seq_idx").on(table.session_id, table.type, table.seq),
  128. index("session_message_session_time_created_id_idx").on(table.session_id, table.time_created, table.id),
  129. index("session_message_time_created_idx").on(table.time_created),
  130. ],
  131. )
  132. export const SessionInputTable = sqliteTable(
  133. "session_input",
  134. {
  135. id: text().$type<SessionMessage.ID>().primaryKey(),
  136. session_id: text()
  137. .$type<SessionSchema.ID>()
  138. .notNull()
  139. .references(() => SessionTable.id, { onDelete: "cascade" }),
  140. prompt: text({ mode: "json" }).notNull().$type<Prompt>(),
  141. delivery: text().$type<SessionInput.Delivery>().notNull(),
  142. admitted_seq: integer().notNull(),
  143. promoted_seq: integer(),
  144. time_created: integer()
  145. .notNull()
  146. .$default(() => Date.now()),
  147. },
  148. (table) => [
  149. index("session_input_session_pending_delivery_seq_idx").on(
  150. table.session_id,
  151. table.promoted_seq,
  152. table.delivery,
  153. table.admitted_seq,
  154. ),
  155. uniqueIndex("session_input_session_admitted_seq_idx").on(table.session_id, table.admitted_seq),
  156. uniqueIndex("session_input_session_promoted_seq_idx").on(table.session_id, table.promoted_seq),
  157. ],
  158. )
  159. export const SessionContextEpochTable = sqliteTable("session_context_epoch", {
  160. session_id: text()
  161. .$type<SessionSchema.ID>()
  162. .primaryKey()
  163. .references(() => SessionTable.id, { onDelete: "cascade" }),
  164. baseline: text().notNull(),
  165. agent: text().$type<AgentV2.ID>().notNull().default(AgentV2.defaultID),
  166. snapshot: text({ mode: "json" }).notNull().$type<SystemContext.Snapshot>(),
  167. baseline_seq: integer().notNull(),
  168. replacement_seq: integer(),
  169. revision: integer().notNull().default(0),
  170. })