lukasa1993/tanstack-skills-extracted/tanstack-ai-persistence-build-cloudflare-artifact-store

import type {
  ArtifactRecord,
  BlobBody,
  BlobGetOptions,
  BlobListOptions,
  BlobListPage,
  BlobObject,
  BlobPutOptions,
  BlobRecord,
} from '@tanstack/ai-persistence'
// BlobStore — the byte layer. R2 backs it.
interface BlobStore {
  put: (
    key: string,
    body: BlobBody,
    options?: BlobPutOptions,
  ) => Promise<BlobRecord>
  // metadata + byte accessors; `options.range` reads one slice (for `206`s)
  get: (key: string, options?: BlobGetOptions) => Promise<BlobObject | null>
  head: (key: string) => Promise<BlobRecord | null> // metadata only
  delete: (key: string) => Promise<void> // no-op if absent
  list: (options?: BlobListOptions) => Promise<BlobListPage>
}
// ArtifactStore — the metadata layer. D1 backs it.
interface ArtifactStore {
  save: (record: ArtifactRecord) => Promise<void> // insert or overwrite
  get: (artifactId: string) => Promise<ArtifactRecord | null>
  list: (runId: string) => Promise<Array<ArtifactRecord>> // [] when none
  listForThread: (threadId: string) => Promise<Array<ArtifactRecord>>
  delete: (artifactId: string) => Promise<void>
  deleteForRun: (runId: string) => Promise<void>
}

import { defineBlobStore, resolveBlobRange } from '@tanstack/ai-persistence'
import type { BlobObject, BlobRecord } from '@tanstack/ai-persistence'
 
// R2 multipart parts must be ≥ 5 MiB and — except for the last — all exactly
// the SAME size, so a part reader has to cut on an exact boundary and carry the
// remainder. 8 MiB × the 10,000-part ceiling puts the multipart path's limit at
// ~80 GiB; raise this for larger objects, and check R2's current object-size
// limits before promising more.
const MULTIPART_PART_SIZE = 8 * 1024 * 1024
 
/**
 * Cut exactly `limit` bytes off the stream (fewer only at EOF), carrying any
 * overshoot into the next part.
 *
 * `carry` is the leftover from the previous call. Chunks are collected by
 * reference and copied once per part: growing a `Uint8Array` chunk-by-chunk
 * instead would re-copy the whole part on every read — quadratic work for a
 * part built from hundreds of small chunks.
 */
async function readPart(
  reader: ReadableStreamDefaultReader<Uint8Array>,
  limit: number,
  carry: Uint8Array,
): Promise<{ bytes: Uint8Array; carry: Uint8Array; eof: boolean }> {
  const chunks: Array<Uint8Array> = carry.byteLength > 0 ? [carry] : []
  let total = carry.byteLength
  let eof = false
  while (total < limit) {
    const { value, done } = await reader.read()
    if (done) {
      eof = true
      break
    }
    chunks.push(value)
    total += value.byteLength
  }
  const joined = new Uint8Array(total)
  let offset = 0
  for (const chunk of chunks) {
    joined.set(chunk, offset)
    offset += chunk.byteLength
  }
  // Equal-sized parts: hand back exactly `limit` and keep the rest for next
  // time. At EOF the last part is whatever is left, which R2 allows.
  if (!eof && total > limit) {
    return {
      bytes: joined.subarray(0, limit),
      carry: joined.subarray(limit),
      eof,
    }
  }
  return { bytes: joined, carry: new Uint8Array(0), eof }
}
 
// R2 takes a single-shot put up to 5 GiB. Above that the upload has to be
// multipart even when the length is known.
const MAX_SINGLE_SHOT_BYTES = 5 * 1024 * 1024 * 1024
 
/**
 * `R2Bucket.put` requires a stream with a known length; a pass-through
 * TransformStream (what the middleware sends when it had to cap the body) has
 * none. Re-declare the length when `expectedLength` provides it — the
 * middleware only sets it when it is the exact decoded byte count — and
 * otherwise stream through a multipart upload, one buffered part at a time.
 */
async function putStream(
  bucket: R2Bucket,
  key: string,
  body: ReadableStream<Uint8Array>,
  options: R2PutOptions,
  expectedLength: number | undefined,
): Promise<R2Object | null> {
  if (expectedLength !== undefined && expectedLength <= MAX_SINGLE_SHOT_BYTES) {
    return bucket.put(
      key,
      body.pipeThrough(new FixedLengthStream(expectedLength)),
      options,
    )
  }
  // Unknown length, or too big for one shot: multipart, one part at a time.
  const reader = body.getReader()
  const first = await readPart(reader, MULTIPART_PART_SIZE, new Uint8Array(0))
  if (first.eof) {
    // The whole body fit in one part — a plain put is enough.
    return bucket.put(key, first.bytes, options)
  }
  const upload = await bucket.createMultipartUpload(key, options)
  try {
    const parts: Array<R2UploadedPart> = [
      await upload.uploadPart(1, first.bytes),
    ]
    let carry = first.carry
    let partNumber = 2
    for (;;) {
      const part = await readPart(reader, MULTIPART_PART_SIZE, carry)
      carry = part.carry
      if (part.bytes.byteLength > 0) {
        // 10,000 parts is R2's ceiling; failing here beats a confusing error
        // from `complete` after uploading gigabytes.
        if (partNumber > 10_000) {
          throw new Error(
            `Artifact at ${key} exceeds the multipart part limit — raise MULTIPART_PART_SIZE.`,
          )
        }
        parts.push(await upload.uploadPart(partNumber, part.bytes))
        partNumber += 1
      }
      if (part.eof) break
    }
    return await upload.complete(parts)
  } catch (error) {
    await upload.abort().catch(() => undefined)
    throw error
  }
}
 
function toRecord(obj: R2Object): BlobRecord {
  const uploaded = obj.uploaded.getTime()
  return {
    key: obj.key,
    size: obj.size,
    etag: obj.etag,
    ...(obj.httpMetadata?.contentType
      ? { contentType: obj.httpMetadata.contentType }
      : {}),
    ...(obj.customMetadata ? { customMetadata: obj.customMetadata } : {}),
    createdAt: uploaded,
    updatedAt: uploaded,
  }
}
 
export function r2BlobStore(bucket: R2Bucket) {
  return defineBlobStore({
    async put(key, body, options) {
      const r2Options = {
        ...(options?.contentType
          ? { httpMetadata: { contentType: options.contentType } }
          : {}),
        ...(options?.customMetadata
          ? { customMetadata: options.customMetadata }
          : {}),
      }
      const obj =
        body instanceof ReadableStream
          ? await putStream(
              bucket,
              key,
              body,
              r2Options,
              options?.expectedLength,
            )
          : await bucket.put(key, body, r2Options)
      // R2.put returns null only when an onlyIf precondition fails — not used here.
      if (!obj) throw new Error(`R2 put failed for ${key}`)
      return toRecord(obj)
    },