|
1 | 1 | export * as SessionBlob from "./blob.js" |
2 | 2 |
|
3 | | -import { inArray } from "drizzle-orm" |
4 | | -import { Effect, Schema } from "effect" |
| 3 | +import { createHash } from "node:crypto" |
| 4 | +import { and, eq, gte, inArray, notExists, notInArray, type SQL, sql } from "drizzle-orm" |
| 5 | +import { Effect, Option, Schema } from "effect" |
5 | 6 | import { map } from "effect/Array" |
6 | 7 | import type { FileAttachment } from "@opencode/schema/prompt" |
7 | 8 | import type { Tool } from "@opencode/schema/tool" |
| 9 | +import { Config } from "../config.js" |
8 | 10 | import type { Database } from "../database/database.js" |
9 | | -import { BlobTable, SessionInboxTable, SessionMessageTable } from "./sql.js" |
| 11 | +import type { SessionMessage } from "./message.js" |
| 12 | +import type { SessionSchema } from "./schema.js" |
| 13 | +import { BlobRefTable, BlobTable, SessionInboxTable, SessionMessageTable } from "./sql.js" |
10 | 14 | import { SessionStored } from "./stored.js" |
11 | 15 |
|
12 | 16 | type DatabaseService = Database.Interface["db"] |
@@ -105,3 +109,172 @@ export const hydrateInbox = Effect.fnUntraced(function* (db: DatabaseService, ro |
105 | 109 | return { ...row, payload: { ...payload, files: payload.files.map((file) => inlineFile(file, blobs)) } } |
106 | 110 | }) |
107 | 111 | }) |
| 112 | + |
| 113 | +/** A blob's hash is the lowercase hex SHA-256 of its raw bytes. */ |
| 114 | +export const hashBytes = (raw: Uint8Array) => createHash("sha256").update(raw).digest("hex") |
| 115 | + |
| 116 | +/** The row a ref keeps its blob alive for: a message at its `seq`, or a pending inbox row at its `enqueued_seq`. */ |
| 117 | +export interface Owner { |
| 118 | + readonly sessionID: SessionSchema.ID |
| 119 | + readonly messageID: SessionMessage.ID |
| 120 | + readonly seq: number |
| 121 | +} |
| 122 | + |
| 123 | +/** Blobs taken out of one row, by hash. */ |
| 124 | +type Pending = Map<string, typeof BlobTable.$inferInsert> |
| 125 | +type AssistantPart = SessionMessage.Assistant["content"][number] |
| 126 | + |
| 127 | +/** Reads `experimental.media_blob` from the Config in scope on every write; without one, media stays inline. */ |
| 128 | +const writeEnabled = Effect.gen(function* () { |
| 129 | + const config = yield* Effect.serviceOption(Config.Service) |
| 130 | + if (Option.isNone(config)) return false |
| 131 | + return Config.latest(yield* config.value.entries(), "experimental")?.media_blob === true |
| 132 | +}) |
| 133 | + |
| 134 | +const DATA_URI = /^data:[^,]*;base64,(.*)$/s |
| 135 | + |
| 136 | +function take(pending: Pending, raw: Buffer, mime: string) { |
| 137 | + const hash = hashBytes(raw) |
| 138 | + if (!pending.has(hash)) pending.set(hash, { hash, mime, bytes: raw.byteLength, data: raw }) |
| 139 | + return SessionStored.ref(hash) |
| 140 | +} |
| 141 | + |
| 142 | +function fileRef(pending: Pending, file: SessionStored.File): SessionStored.File { |
| 143 | + if (isFileRef(file)) return file |
| 144 | + const { data, ...rest } = file |
| 145 | + return { ...rest, uri: take(pending, Buffer.from(data, "base64"), file.mime) } |
| 146 | +} |
| 147 | + |
| 148 | +function toolRef(pending: Pending, item: Tool.Content): Tool.Content { |
| 149 | + if (item.type !== "file") return item |
| 150 | + const data = DATA_URI.exec(item.uri)?.[1] |
| 151 | + return data === undefined ? item : { ...item, uri: take(pending, Buffer.from(data, "base64"), item.mime) } |
| 152 | +} |
| 153 | + |
| 154 | +function toolPartRefs(pending: Pending, part: AssistantPart): AssistantPart { |
| 155 | + if (part.type !== "tool" || (part.state.status !== "completed" && part.state.status !== "error")) return part |
| 156 | + const content = part.state.content |
| 157 | + if (!content) return part |
| 158 | + return { ...part, state: { ...part.state, content: map(content, (item) => toolRef(pending, item)) } } |
| 159 | +} |
| 160 | + |
| 161 | +const save = Effect.fnUntraced(function* (db: DatabaseService, pending: Pending) { |
| 162 | + if (pending.size === 0) return |
| 163 | + yield* db |
| 164 | + .insert(BlobTable) |
| 165 | + .values([...pending.values()]) |
| 166 | + .onConflictDoNothing() |
| 167 | + .run() |
| 168 | + .pipe(Effect.orDie) |
| 169 | +}) |
| 170 | + |
| 171 | +/** {@link externalize} for a user message or pending inbox user payload. */ |
| 172 | +export const externalizePayload = Effect.fnUntraced(function* < |
| 173 | + T extends { readonly files?: ReadonlyArray<SessionStored.File> }, |
| 174 | +>(db: DatabaseService, input: T): Effect.fn.Return<T> { |
| 175 | + if (input.files === undefined || !(yield* writeEnabled)) return input |
| 176 | + const pending: Pending = new Map() |
| 177 | + const files = input.files.map((file) => fileRef(pending, file)) |
| 178 | + if (pending.size === 0) return input |
| 179 | + yield* save(db, pending) |
| 180 | + return { ...input, files } |
| 181 | +}) |
| 182 | + |
| 183 | +/** |
| 184 | + * With `experimental.media_blob` on, moves inline user files and tool `data:` files into `blob` and returns the |
| 185 | + * message to store with refs in their place. Existing refs pass through. |
| 186 | + */ |
| 187 | +export const externalize = Effect.fnUntraced(function* ( |
| 188 | + db: DatabaseService, |
| 189 | + message: SessionMessage.Info | SessionStored.User, |
| 190 | +) { |
| 191 | + if (message.type === "user") return yield* externalizePayload(db, message) |
| 192 | + if (message.type !== "assistant" || !(yield* writeEnabled)) return message |
| 193 | + const pending: Pending = new Map() |
| 194 | + const content = message.content.map((part) => toolPartRefs(pending, part)) |
| 195 | + if (pending.size === 0) return message |
| 196 | + yield* save(db, pending) |
| 197 | + return { ...message, content } |
| 198 | +}) |
| 199 | + |
| 200 | +const hashesOf = (data: MessageRow["data"] | InboxRow["payload"]) => { |
| 201 | + const hashes = new Set<string>() |
| 202 | + collect(data, hashes) |
| 203 | + return hashes |
| 204 | +} |
| 205 | + |
| 206 | +/** Records that `owner` references each blob its stored row names, so release keeps them until it lets go. */ |
| 207 | +export const recordRefs = Effect.fnUntraced(function* ( |
| 208 | + db: DatabaseService, |
| 209 | + owner: Owner, |
| 210 | + data: MessageRow["data"] | InboxRow["payload"], |
| 211 | +) { |
| 212 | + const hashes = hashesOf(data) |
| 213 | + if (hashes.size === 0) return |
| 214 | + // A ref whose blob is already gone keeps nothing alive; reads render it as unavailable. |
| 215 | + const found = yield* db |
| 216 | + .select({ hash: BlobTable.hash }) |
| 217 | + .from(BlobTable) |
| 218 | + .where(inArray(BlobTable.hash, [...hashes])) |
| 219 | + .all() |
| 220 | + .pipe(Effect.orDie) |
| 221 | + if (found.length === 0) return |
| 222 | + yield* db |
| 223 | + .insert(BlobRefTable) |
| 224 | + .values( |
| 225 | + found.map(({ hash }) => ({ session_id: owner.sessionID, message_id: owner.messageID, hash, seq: owner.seq })), |
| 226 | + ) |
| 227 | + .onConflictDoUpdate({ |
| 228 | + target: [BlobRefTable.session_id, BlobRefTable.message_id, BlobRefTable.hash], |
| 229 | + set: { seq: owner.seq }, |
| 230 | + }) |
| 231 | + .run() |
| 232 | + .pipe(Effect.orDie) |
| 233 | +}) |
| 234 | + |
| 235 | +/** Records the refs a rewritten message names and releases the ones it no longer does. */ |
| 236 | +export const syncRefs = Effect.fnUntraced(function* (db: DatabaseService, owner: Owner, data: MessageRow["data"]) { |
| 237 | + yield* recordRefs(db, owner, data) |
| 238 | + const present = [...hashesOf(data)] |
| 239 | + const own = [eq(BlobRefTable.session_id, owner.sessionID), eq(BlobRefTable.message_id, owner.messageID)] as const |
| 240 | + yield* present.length === 0 ? release(db, ...own) : release(db, ...own, notInArray(BlobRefTable.hash, present)) |
| 241 | +}) |
| 242 | + |
| 243 | +/** Deletes the matching refs, then every blob among theirs that no ref names anymore. */ |
| 244 | +const release = Effect.fnUntraced(function* (db: DatabaseService, ...where: readonly [SQL, ...SQL[]]) { |
| 245 | + const released = yield* db |
| 246 | + .delete(BlobRefTable) |
| 247 | + .where(and(...where)) |
| 248 | + .returning({ hash: BlobRefTable.hash }) |
| 249 | + .all() |
| 250 | + .pipe(Effect.orDie) |
| 251 | + const hashes = [...new Set(released.map((ref) => ref.hash))] |
| 252 | + for (let offset = 0; offset < hashes.length; offset += 500) |
| 253 | + yield* db |
| 254 | + .delete(BlobTable) |
| 255 | + .where( |
| 256 | + and( |
| 257 | + inArray(BlobTable.hash, hashes.slice(offset, offset + 500)), |
| 258 | + notExists( |
| 259 | + db |
| 260 | + .select({ one: sql`1` }) |
| 261 | + .from(BlobRefTable) |
| 262 | + .where(eq(BlobRefTable.hash, BlobTable.hash)), |
| 263 | + ), |
| 264 | + ), |
| 265 | + ) |
| 266 | + .run() |
| 267 | + .pipe(Effect.orDie) |
| 268 | +}) |
| 269 | + |
| 270 | +/** Releases the blobs a removed message or cancelled inbox row referenced. */ |
| 271 | +export const releaseMessage = (db: DatabaseService, sessionID: SessionSchema.ID, messageID: SessionMessage.ID) => |
| 272 | + release(db, eq(BlobRefTable.session_id, sessionID), eq(BlobRefTable.message_id, messageID)) |
| 273 | + |
| 274 | +/** Releases the blobs referenced at or after `seq`, for a revert. */ |
| 275 | +export const releaseFrom = (db: DatabaseService, sessionID: SessionSchema.ID, seq: number) => |
| 276 | + release(db, eq(BlobRefTable.session_id, sessionID), gte(BlobRefTable.seq, seq)) |
| 277 | + |
| 278 | +/** Releases every blob a deleted session referenced. */ |
| 279 | +export const releaseSession = (db: DatabaseService, sessionID: SessionSchema.ID) => |
| 280 | + release(db, eq(BlobRefTable.session_id, sessionID)) |
0 commit comments