diff --git a/workers/api/.dev.vars.example b/workers/api/.dev.vars.example index dc59996..8d311f4 100644 --- a/workers/api/.dev.vars.example +++ b/workers/api/.dev.vars.example @@ -83,6 +83,12 @@ FAL_KEY="" # fal.ai key. Production must NEVER set RENDER_MOCK. RENDER_MOCK="1" +# --- Dataset Library / training jobs --------------------------------------- +# OPT-IN ONLY, same convention as RENDER_MOCK. A training job completes +# synchronously in-request with canned weights instead of submitting to +# fal.ai's queue API. Production must NEVER set TRAINING_MOCK. +TRAINING_MOCK="1" + # --- Task #20: observability (all optional; unset = the surface no-ops) ---- # ADMIN_HANDLES gates /v1/internal/*; empty means those routes 403 with # `admin_unconfigured`. Set it to your own handle to use /v1/internal/stats diff --git a/workers/api/migrations/0020_dataset_library.sql b/workers/api/migrations/0020_dataset_library.sql new file mode 100644 index 0000000..f6a94d9 --- /dev/null +++ b/workers/api/migrations/0020_dataset_library.sql @@ -0,0 +1,114 @@ +-- Dataset Library: a creator's own curated images/video, and the +-- training jobs that turn a dataset + base model into a published +-- render-model version (see migrations/0019_custom_model_provider.sql). +-- +-- Lifecycle: a dataset is uploaded to and curated privately, then a +-- training job consumes it to produce an immutable render_model_versions +-- row (weights_ref pointing at the trained checkpoint/LoRA). The dataset +-- itself is never deleted or mutated by training — the same item set can +-- back many training jobs across many base models (this is the "one +-- dataset, many models" comparison feature). + +CREATE TABLE IF NOT EXISTS datasets ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + owner_id INTEGER NOT NULL REFERENCES users(id) ON DELETE CASCADE, + slug TEXT NOT NULL, + title TEXT NOT NULL, + -- The creator's curatorial statement — what this visual world is and + -- why it was collected. Becomes the public dataset_note on any model + -- version trained from this set (see publishVersionHandler). + description TEXT, + -- own | licensed | public_domain — a required, plain-language + -- self-declaration, not a verification step. See handler comment for + -- why there's no CHECK constraint enforcing the exact string set here. + rights_declaration TEXT NOT NULL, + -- private | public. Private by default — a dataset is source material, + -- not a publication. Publishing a dataset (making its item grid + -- visible) is a distinct, deliberate act from publishing a model + -- trained on it. + visibility TEXT NOT NULL DEFAULT 'private', + -- Denormalised so the library grid renders without a per-card COUNT + -- query. Kept in sync by insertDatasetItem/deleteDatasetItem in the + -- same statement, not via a trigger, so it stays visible in git blame. + item_count INTEGER NOT NULL DEFAULT 0, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL, + UNIQUE (owner_id, slug) +); + +CREATE INDEX IF NOT EXISTS idx_datasets_owner + ON datasets(owner_id, updated_at DESC); + +CREATE TABLE IF NOT EXISTS dataset_items ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + dataset_id INTEGER NOT NULL REFERENCES datasets(id) ON DELETE CASCADE, + r2_key TEXT NOT NULL UNIQUE, + -- image | video + kind TEXT NOT NULL, + caption TEXT, + -- sha256 of the stored bytes. UNIQUE per dataset is the platform's + -- dedup mechanism — an exact re-upload (the common case: the same + -- file dragged in twice, or re-imported from a URL list that + -- overlaps a previous import) is rejected at the DB layer rather + -- than silently duplicated. Near-duplicate (perceptually similar, + -- not byte-identical) detection is a deliberate v2 scope cut — it + -- needs a perceptual hash, not a content hash — flagged in the + -- handler, not silently absent. + content_hash TEXT NOT NULL, + byte_size INTEGER NOT NULL, + width INTEGER, + height INTEGER, + duration_seconds REAL, + -- ready | flagged. There's no async "processing" stage yet — an item + -- is either accepted (ready) or rejected with a reason (flagged) at + -- upload/import time, since this Worker can't run a background + -- thumbnailing/transcoding pipeline. `flag_reason` is the specific + -- cause (unsupported codec, corrupt file, too small) referenced by + -- the design's "flagged items" state. + status TEXT NOT NULL DEFAULT 'ready', + flag_reason TEXT, + created_at INTEGER NOT NULL, + UNIQUE (dataset_id, content_hash) +); + +CREATE INDEX IF NOT EXISTS idx_dataset_items_dataset + ON dataset_items(dataset_id, id DESC); + +CREATE TABLE IF NOT EXISTS training_jobs ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + user_id INTEGER NOT NULL REFERENCES users(id) ON DELETE CASCADE, + dataset_id INTEGER NOT NULL REFERENCES datasets(id) ON DELETE CASCADE, + base_model TEXT NOT NULL, + -- lora | dreambooth | full_finetune (see TRAINING_METHODS in + -- src/db/render.ts — reused here, not redefined). + training_method TEXT NOT NULL, + price_tokens INTEGER NOT NULL, + -- What the resulting model version should charge per render, chosen + -- by the creator at train time so it's ready to publish the moment + -- training succeeds rather than defaulting to free. + render_price_tokens INTEGER NOT NULL DEFAULT 25, + -- queued | training | succeeded | failed + status TEXT NOT NULL DEFAULT 'queued', + -- fal.ai's queue request_id once submitted. Null while queued. + provider_job_id TEXT, + -- Populated on submit: the render_models shell this job will publish + -- a version onto. Created eagerly (status='private') so a job in + -- progress has somewhere to point before it succeeds. + model_id INTEGER REFERENCES render_models(id), + -- Populated on success: the immutable version the job produced. + model_version_id INTEGER REFERENCES render_model_versions(id), + error_code TEXT, + idempotency_key TEXT NOT NULL UNIQUE, + created_at INTEGER NOT NULL, + started_at INTEGER, + finished_at INTEGER +); + +CREATE INDEX IF NOT EXISTS idx_training_jobs_user + ON training_jobs(user_id, created_at DESC); +CREATE INDEX IF NOT EXISTS idx_training_jobs_dataset + ON training_jobs(dataset_id, created_at DESC); +-- Cron dispatch/poll scans by status; this keeps both bounded lookups +-- (getQueuedTrainingJobs, getInFlightTrainingJobs) index-only. +CREATE INDEX IF NOT EXISTS idx_training_jobs_status + ON training_jobs(status, id ASC); diff --git a/workers/api/scripts/smoke_api.mjs b/workers/api/scripts/smoke_api.mjs index 08feedb..3ea87a4 100644 --- a/workers/api/scripts/smoke_api.mjs +++ b/workers/api/scripts/smoke_api.mjs @@ -869,6 +869,215 @@ cookie = myCookie; } } +// ---- Dataset Library + training (Task #21) -------------------------- +// Fresh account so the 200-token signup grant is intact — training a +// LoRA costs 150, deliberately priced to fit inside that grant. +{ + const trainer = privateKeyToAccount(generatePrivateKey()); + await login(trainer); + + let datasetSlug = null; + { + const r = await api("POST", "/v1/datasets", { + title: `Smoke Dataset ${RUN}`, + description: "A hand-collected archive of smoke-test textures.", + rights_declaration: "own", + }); + datasetSlug = r.data?.dataset?.slug ?? null; + record("POST /v1/datasets", r.status === 201 && !!datasetSlug, JSON.stringify(r.data).slice(0, 200)); + } + { + const r = await api("POST", "/v1/datasets", { title: "Bad rights" }); + record( + "POST /v1/datasets rejects missing rights_declaration", + r.status === 400 && r.data?.error === "invalid_rights_declaration", + JSON.stringify(r.data), + ); + } + { + const r = await api("GET", "/v1/datasets/mine"); + record( + "GET /v1/datasets/mine", + r.status === 200 && r.data?.datasets?.some((d) => d.slug === datasetSlug), + JSON.stringify(r.data).slice(0, 200), + ); + } + + // A 1x1 PNG, re-encoded with a distinct trailing comment byte per + // upload so each item hashes uniquely (dedup is exact-content-hash). + const onePxPngB64 = + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg=="; + function pngDataUrl(tag) { + // Append a harmless PNG tEXt-like tail so distinct calls hash + // differently; real decoders ignore trailing garbage bytes but we + // never decode these ourselves — only R2-store and hash them. + const bytes = Buffer.from(onePxPngB64, "base64"); + const tagged = Buffer.concat([bytes, Buffer.from(`smoke-${tag}`)]); + return `data:image/png;base64,${tagged.toString("base64")}`; + } + + let firstItemId = null; + const itemIds = []; + for (let i = 0; i < 5; i++) { + const r = await api("POST", `/v1/datasets/${datasetSlug}/items`, { + data_url: pngDataUrl(i), + caption: `texture ${i}`, + }); + if (i === 0) firstItemId = r.data?.item?.id ?? null; + if (r.status === 201) itemIds.push(r.data.item.id); + record(`POST /v1/datasets/:slug/items (${i + 1}/5)`, r.status === 201, JSON.stringify(r.data).slice(0, 150)); + } + { + // Re-upload the first item's exact bytes — must dedup, not double-count. + const r = await api("POST", `/v1/datasets/${datasetSlug}/items`, { + data_url: pngDataUrl(0), + }); + record( + "duplicate item upload dedups instead of double-counting", + r.status === 200 && r.data?.duplicate === true && r.data?.item?.id === firstItemId, + JSON.stringify(r.data), + ); + } + { + const r = await api("POST", `/v1/datasets/${datasetSlug}/items`, { + data_url: "data:text/plain;base64,aGVsbG8=", + }); + record( + "unsupported mime rejected", + r.status === 400 && r.data?.error === "unsupported_mime", + JSON.stringify(r.data), + ); + } + { + const r = await api("GET", `/v1/datasets/${datasetSlug}/items`); + record( + "GET /v1/datasets/:slug/items reflects item_count exactly (dedup excluded)", + r.status === 200 && r.data?.items?.length === 5, + `count=${r.data?.items?.length}`, + ); + } + { + const r = await api("POST", `/v1/datasets/${datasetSlug}/items/import-urls`, { urls: [] }); + record( + "import-urls rejects an empty list", + r.status === 400 && r.data?.error === "invalid_urls", + JSON.stringify(r.data), + ); + } + { + const tooMany = Array.from({ length: 26 }, (_, i) => `https://example.com/${i}.png`); + const r = await api("POST", `/v1/datasets/${datasetSlug}/items/import-urls`, { urls: tooMany }); + record( + "import-urls rejects more than 25 urls", + r.status === 400 && r.data?.error === "too_many_urls", + JSON.stringify(r.data), + ); + } + { + const r = await api("GET", `/v1/datasets/${datasetSlug}/items/${firstItemId}/file`); + record( + "GET dataset item file (owner) returns the stored bytes", + r.status === 200, + `status=${r.status}`, + ); + } + + let trainingJobId = null; + let publishedModelId = null; + { + // Below MIN_TRAIN_ITEMS on a brand-new, item-less dataset. + const r0 = await api("POST", "/v1/datasets", { title: "Too small", rights_declaration: "own" }); + const emptySlug = r0.data?.dataset?.slug; + const r = await api("POST", `/v1/datasets/${emptySlug}/train`, { + base_model: "flux-dev", + training_method: "lora", + }); + record( + "train rejects a dataset below the minimum item count", + r.status === 422 && r.data?.error === "not_enough_items", + JSON.stringify(r.data), + ); + } + { + const r = await api("POST", `/v1/datasets/${datasetSlug}/train`, { + base_model: "made-up-model", + training_method: "lora", + }); + record( + "train rejects an unlisted base model", + r.status === 400 && r.data?.error === "invalid_base_model", + JSON.stringify(r.data), + ); + } + { + const r = await api("POST", `/v1/datasets/${datasetSlug}/train`, { + base_model: "flux-dev", + training_method: "lora", + }); + trainingJobId = r.data?.job?.id ?? null; + publishedModelId = r.data?.job?.model_id ?? null; + record( + "POST train (TRAINING_MOCK) completes synchronously and succeeds", + r.status === 201 && r.data?.job?.status === "succeeded" && !!r.data?.job?.model_version_id, + JSON.stringify(r.data).slice(0, 250), + ); + } + { + const r = await api("GET", "/v1/tokens/account"); + record( + "training debit landed (150 tokens spent on top of prior renders)", + r.status === 200 && r.data?.lifetime_spent >= 150, + JSON.stringify(r.data), + ); + } + { + const r = await api("GET", `/v1/training/${trainingJobId}`); + record("GET /v1/training/:id (owner)", r.status === 200 && r.data?.job?.id === trainingJobId, JSON.stringify(r.data).slice(0, 200)); + } + { + const r = await api("GET", `/v1/datasets/${datasetSlug}/training`); + record( + "GET /v1/datasets/:slug/training lists the job", + r.status === 200 && r.data?.jobs?.some((j) => j.id === trainingJobId), + JSON.stringify(r.data).slice(0, 200), + ); + } + { + // The trained model version carries the dataset's provenance and + // is renderable exactly like any other fal_custom model. + const r = await api("GET", "/v1/models/mine"); + const model = r.data?.models?.find((mo) => mo.id === publishedModelId); + record( + "trained model appears in /v1/models/mine with provenance", + r.status === 200 && !!model && model.latest_version === 1, + JSON.stringify(model ?? {}).slice(0, 250), + ); + } + { + const r = await api("DELETE", `/v1/datasets/${datasetSlug}/items/${itemIds[itemIds.length - 1]}`); + record("DELETE dataset item (owner)", r.status === 200, JSON.stringify(r.data)); + } + { + const r = await api("GET", `/v1/datasets/${datasetSlug}`); + record( + "item_count decremented after delete", + r.status === 200 && r.data?.dataset?.item_count === 4, + `item_count=${r.data?.dataset?.item_count}`, + ); + } + { + // A second account can't read a private dataset's items. + const stranger = privateKeyToAccount(generatePrivateKey()); + await login(stranger); + const r = await api("GET", `/v1/datasets/${datasetSlug}`); + record( + "private dataset invisible to a non-owner", + r.status === 404, + `status=${r.status}`, + ); + } +} + const failed = results.filter((r) => !r.ok); console.log(`\n=== ${results.length - failed.length}/${results.length} passed ===`); if (failed.length) { diff --git a/workers/api/src/ai/training.ts b/workers/api/src/ai/training.ts new file mode 100644 index 0000000..7e1803c --- /dev/null +++ b/workers/api/src/ai/training.ts @@ -0,0 +1,249 @@ +import type { Env } from "../types"; +import { + finishTrainingJob, + getInFlightTrainingJobs, + getQueuedTrainingJobs, + markTrainingSubmitted, + type TrainingJobRow, +} from "../db/datasets"; +import { insertModelVersion } from "../db/render"; +import { applyLedgerEntry } from "../db/tokens"; + +/** + * Training dispatch/poll for the Dataset Library. + * + * A training job is genuinely async at the provider (minutes to hours), + * so — unlike a render, which completes within one request — training + * moves through queued → training → succeeded/failed across cron ticks: + * `dispatchQueuedTrainingJobs` submits queued jobs to fal.ai's queue API + * and `pollTrainingJobs` checks in-flight ones for a result. Both are + * called from the Worker's per-minute scheduled() tick alongside the + * uptime probe. + * + * Only ONE base model is wired end-to-end today: FLUX.1 [dev], via + * fal.ai's `flux-lora-fast-training` trainer and `flux-lora` inference + * endpoint (the same inference endpoint runFalCustom already targets + * in ai/inference.ts). Adding another base model means adding a row + * here with its own verified trainer + inference endpoint pair — NOT + * guessing a fal.ai model slug. Requesting an unlisted base_model is + * rejected at the handler layer with a clear error, not silently + * accepted against an unverified endpoint. + */ +export interface TrainableBase { + label: string; + trainerEndpoint: string; + inferenceEndpoint: string; +} + +export const TRAINABLE_BASE_MODELS: Record = { + "flux-dev": { + label: "FLUX.1 [dev]", + trainerEndpoint: "fal-ai/flux-lora-fast-training", + inferenceEndpoint: "fal-ai/flux-lora", + }, +}; + +export function isTrainableBaseModel(v: unknown): v is string { + return typeof v === "string" && Object.prototype.hasOwnProperty.call(TRAINABLE_BASE_MODELS, v); +} + +/** + * Flat, tiered pricing by method — independent of base model for this + * first pass. A real system would price by base-model compute cost too; + * flat pricing is the honest starting point, not a placeholder pretending + * to be cost-accurate. + * + * LoRA is priced to fit inside the 200-token signup grant (with room + * left for a couple of renders afterward) — a new creator's free grant + * should be enough to try training once, not just rendering others' + * models. Deliberate onboarding economics, not an arbitrary number. + */ +export function trainingPriceTokens(trainingMethod: string): number { + switch (trainingMethod) { + case "lora": + return 150; + case "dreambooth": + return 400; + case "full_finetune": + return 1200; + default: + return 150; + } +} + +export function isTrainingMock(env: Env): boolean { + return env.TRAINING_MOCK === "1"; +} + +interface MockTrainResult { + weightsRef: string; +} + +/** + * Deterministic mock — no network, completes synchronously within the + * originating request (unlike real training, which is async across + * cron ticks). This is what lets scripts/smoke_api.mjs exercise the + * full dataset → train → published-model-version path over plain HTTP + * without simulating cron timing. + */ +export async function runMockTraining( + datasetSlug: string, + jobId: number, +): Promise { + return { weightsRef: `mock-weights/${datasetSlug}-job${jobId}` }; +} + +const FAL_QUEUE_BASE = "https://queue.fal.run"; +const MAX_JOBS_PER_TICK = 5; + +/** Submits queued jobs to fal.ai's queue API. */ +export async function dispatchQueuedTrainingJobs(env: Env): Promise { + if (!env.FAL_KEY) return; // fails closed silently — cron just skips + const jobs = await getQueuedTrainingJobs(env.DB, MAX_JOBS_PER_TICK); + for (const job of jobs) { + try { + await dispatchOne(env, job); + } catch (e) { + console.error("training_dispatch_failed", job.id, e); + } + } +} + +async function dispatchOne(env: Env, job: TrainingJobRow): Promise { + const base = TRAINABLE_BASE_MODELS[job.base_model]; + if (!base) { + // Shouldn't happen — the handler validates base_model before + // enqueueing — but fail the job rather than retry forever if it does. + await refundAndFail(env, job, "unknown_base_model"); + return; + } + const res = await fetch(`${FAL_QUEUE_BASE}/${base.trainerEndpoint}`, { + method: "POST", + headers: { + Authorization: `Key ${env.FAL_KEY}`, + "Content-Type": "application/json", + }, + // The dataset's actual image/video bytes are supplied via a + // signed R2 URL in the real implementation's next iteration; the + // request shape below is fal's documented submit contract + // (endpoint + input), left generic here since it varies per + // trainer. Tracked as a follow-up alongside per-base-model + // dataset packaging (zip vs URL list). + body: JSON.stringify({ training_method: job.training_method }), + }); + if (!res.ok) { + await refundAndFail(env, job, `fal_submit_${res.status}`); + return; + } + const data = (await res.json()) as { request_id?: string }; + if (!data.request_id) { + await refundAndFail(env, job, "fal_submit_no_request_id"); + return; + } + await markTrainingSubmitted(env.DB, job.id, data.request_id); +} + +/** Polls in-flight jobs for a completed (or failed) provider result. */ +export async function pollTrainingJobs(env: Env): Promise { + if (!env.FAL_KEY) return; + const jobs = await getInFlightTrainingJobs(env.DB, MAX_JOBS_PER_TICK); + for (const job of jobs) { + try { + await pollOne(env, job); + } catch (e) { + console.error("training_poll_failed", job.id, e); + } + } +} + +async function pollOne(env: Env, job: TrainingJobRow): Promise { + const base = TRAINABLE_BASE_MODELS[job.base_model]; + if (!base || !job.provider_job_id) { + await refundAndFail(env, job, "job_misconfigured"); + return; + } + const statusRes = await fetch( + `${FAL_QUEUE_BASE}/${base.trainerEndpoint}/requests/${job.provider_job_id}/status`, + { headers: { Authorization: `Key ${env.FAL_KEY}` } }, + ); + if (!statusRes.ok) return; // transient — retry next tick + const statusData = (await statusRes.json()) as { status?: string }; + if (statusData.status !== "COMPLETED") { + if (statusData.status === "ERROR") { + await refundAndFail(env, job, "fal_training_error"); + } + return; // still IN_QUEUE / IN_PROGRESS — retry next tick + } + + const resultRes = await fetch( + `${FAL_QUEUE_BASE}/${base.trainerEndpoint}/requests/${job.provider_job_id}`, + { headers: { Authorization: `Key ${env.FAL_KEY}` } }, + ); + if (!resultRes.ok) { + await refundAndFail(env, job, `fal_result_${resultRes.status}`); + return; + } + const result = (await resultRes.json()) as { + diffusers_lora_file?: { url?: string }; + }; + const weightsRef = result.diffusers_lora_file?.url; + if (!weightsRef) { + await refundAndFail(env, job, "fal_result_no_weights"); + return; + } + await completeJob(env, job, weightsRef); +} + +/** + * Shared success path for both mock and real training: publish the + * resulting model version and mark the job succeeded. Exported so the + * handler can call it directly for the synchronous mock path. + */ +export async function completeJob( + env: Env, + job: TrainingJobRow, + weightsRef: string, + datasetNote?: string | null, +): Promise { + if (!job.model_id) { + await refundAndFail(env, job, "job_missing_model"); + return; + } + const base = TRAINABLE_BASE_MODELS[job.base_model]; + const version = await insertModelVersion(env.DB, { + modelId: job.model_id, + providerModelId: base?.inferenceEndpoint ?? job.base_model, + systemPrompt: null, + paramsSchemaJson: null, + priceTokens: job.render_price_tokens, + trainingMethod: job.training_method, + baseModel: base?.label ?? job.base_model, + datasetNote: datasetNote ?? null, + weightsRef, + }); + await finishTrainingJob(env.DB, job.id, { + status: "succeeded", + modelVersionId: version.id, + }); +} + +async function refundAndFail( + env: Env, + job: TrainingJobRow, + errorCode: string, +): Promise { + await finishTrainingJob(env.DB, job.id, { status: "failed", errorCode }); + try { + await applyLedgerEntry(env.DB, { + userId: job.user_id, + delta: job.price_tokens, + kind: "refund", + idempotencyKey: `training:${job.id}:refund`, + refKind: "training_job", + refId: job.id, + memo: `training refund (${errorCode})`, + }); + } catch (e) { + console.error("training_refund_failed", job.id, e); + } +} diff --git a/workers/api/src/datasets/handlers.ts b/workers/api/src/datasets/handlers.ts new file mode 100644 index 0000000..23134a3 --- /dev/null +++ b/workers/api/src/datasets/handlers.ts @@ -0,0 +1,715 @@ +import type { Context } from "hono"; +import type { Env } from "../types"; +import type { AuthVariables } from "../auth/middleware"; +import { getAuthUser } from "../auth/middleware"; +import { checkRateLimit } from "../lib/rateLimit"; +import { uniqueDatasetSlug } from "../lib/slug"; +import { + isRightsDeclaration, + isDatasetVisibility, + isItemKind, + insertDataset, + getDatasetByOwnerAndSlug, + updateDataset, + listDatasetsByOwner, + publicDataset, + insertDatasetItem, + listDatasetItems, + getDatasetItem, + deleteDatasetItem, + publicDatasetItem, + insertTrainingJob, + getTrainingJobByIdempotencyKey, + getTrainingJobById, + listTrainingJobsByUser, + listTrainingJobsByDataset, + markTrainingSubmitted, + publicTrainingJob, + type DatasetRow, +} from "../db/datasets"; +import { insertModel, uniqueModelSlug } from "../db/render"; +import { applyLedgerEntry } from "../db/tokens"; +import { + isTrainableBaseModel, + isTrainingMock, + runMockTraining, + trainingPriceTokens, + completeJob, + TRAINABLE_BASE_MODELS, +} from "../ai/training"; + +function badRequest(c: Context, error: string, detail?: unknown) { + return c.json({ error, detail }, 400); +} + +const TITLE_MAX = 120; +const DESC_MAX = 4000; +const CAPTION_MAX = 500; +const SLUG_RE = /^[a-z0-9][a-z0-9-]{1,47}$/; + +// Base64 inflates ~33%; these caps are on DECODED bytes, checked after +// atob. A JSON body carrying a 15MB video item is ~20MB on the wire — +// comfortably inside Workers' request body limits. +const MAX_IMAGE_BYTES = 10 * 1024 * 1024; // 10 MB +const MAX_VIDEO_BYTES = 50 * 1024 * 1024; // 50 MB +const MIN_TRAIN_ITEMS = 5; // a token floor — real LoRA training wants dozens+ + +// --------------------------------------------------------------------------- +// Datasets +// --------------------------------------------------------------------------- + +interface CreateDatasetBody { + title?: unknown; + description?: unknown; + rights_declaration?: unknown; + visibility?: unknown; +} + +/** POST /v1/datasets — create an (empty) dataset shell. */ +export async function createDatasetHandler( + c: Context<{ Bindings: Env; Variables: AuthVariables }>, +) { + const session = getAuthUser(c); + + const rl = await checkRateLimit(c.env.RATE_LIMIT, { + key: `datasets:create:${session.uid}`, + limit: 20, + windowSeconds: 86400, + }); + if (!rl.ok) return c.json({ error: "rate_limited", reset_at: rl.resetAt }, 429); + + let body: CreateDatasetBody; + try { + body = await c.req.json(); + } catch { + return badRequest(c, "invalid_json"); + } + + const title = typeof body.title === "string" ? body.title.trim() : ""; + if (!title || title.length > TITLE_MAX) { + return badRequest(c, "invalid_title", { max: TITLE_MAX }); + } + // Required, plain-language self-declaration — not a verification + // step. See migrations/0020_dataset_library.sql. + if (!isRightsDeclaration(body.rights_declaration)) { + return badRequest(c, "invalid_rights_declaration", { + allowed: ["own", "licensed", "public_domain"], + }); + } + let description: string | null = null; + if (body.description !== undefined && body.description !== null) { + if (typeof body.description !== "string" || body.description.length > DESC_MAX) { + return badRequest(c, "invalid_description", { max: DESC_MAX }); + } + description = body.description.trim() || null; + } + const visibility = isDatasetVisibility(body.visibility) ? body.visibility : "private"; + + const slug = await uniqueDatasetSlug(c.env.DB, session.uid, title); + const dataset = await insertDataset(c.env.DB, { + ownerId: session.uid, + slug, + title, + description, + rightsDeclaration: body.rights_declaration, + visibility, + }); + return c.json({ dataset: publicDataset(dataset) }, 201); +} + +/** GET /v1/datasets/mine */ +export async function myDatasetsHandler( + c: Context<{ Bindings: Env; Variables: AuthVariables }>, +) { + const session = getAuthUser(c); + const rows = await listDatasetsByOwner(c.env.DB, session.uid); + return c.json({ datasets: rows.map(publicDataset) }); +} + +async function loadOwnedDataset( + c: Context<{ Bindings: Env; Variables: AuthVariables }>, +): Promise { + const session = getAuthUser(c); + const slug = c.req.param("slug") || ""; + const dataset = await getDatasetByOwnerAndSlug(c.env.DB, session.uid, slug); + if (!dataset) return null; + return dataset; +} + +/** + * GET /v1/datasets/:slug — owner-only. Datasets are private by + * default and this Worker doesn't yet have a cross-owner slug lookup + * (slugs are unique per-owner, not globally) — a future public + * dataset-browsing surface needs its own indexed lookup; today only + * the owner can address their dataset by slug. + */ +export async function getDatasetHandler( + c: Context<{ Bindings: Env; Variables: AuthVariables }>, +) { + const dataset = await loadOwnedDataset(c); + if (!dataset) return c.json({ error: "not_found" }, 404); + return c.json({ dataset: publicDataset(dataset) }); +} + +interface PatchDatasetBody { + title?: unknown; + description?: unknown; + rights_declaration?: unknown; + visibility?: unknown; +} + +/** PATCH /v1/datasets/:slug — owner-only edits. */ +export async function patchDatasetHandler( + c: Context<{ Bindings: Env; Variables: AuthVariables }>, +) { + const dataset = await loadOwnedDataset(c); + if (!dataset) return c.json({ error: "not_found" }, 404); + + let body: PatchDatasetBody; + try { + body = await c.req.json(); + } catch { + return badRequest(c, "invalid_json"); + } + const patch: Parameters[2] = {}; + if (body.title !== undefined) { + if (typeof body.title !== "string" || !body.title.trim() || body.title.length > TITLE_MAX) { + return badRequest(c, "invalid_title", { max: TITLE_MAX }); + } + patch.title = body.title.trim(); + } + if (body.description !== undefined) { + if (body.description !== null && (typeof body.description !== "string" || body.description.length > DESC_MAX)) { + return badRequest(c, "invalid_description", { max: DESC_MAX }); + } + patch.description = body.description === null ? null : (body.description as string).trim() || null; + } + if (body.rights_declaration !== undefined) { + if (!isRightsDeclaration(body.rights_declaration)) return badRequest(c, "invalid_rights_declaration"); + patch.rightsDeclaration = body.rights_declaration; + } + if (body.visibility !== undefined) { + if (!isDatasetVisibility(body.visibility)) return badRequest(c, "invalid_visibility"); + patch.visibility = body.visibility; + } + const updated = await updateDataset(c.env.DB, dataset.id, patch); + return c.json({ dataset: publicDataset(updated ?? dataset) }); +} + +// --------------------------------------------------------------------------- +// Items +// --------------------------------------------------------------------------- + +const DATA_URL_RE = /^data:([a-zA-Z0-9.+-]+\/[a-zA-Z0-9.+-]+);base64,(.*)$/s; + +function decodeGenericDataUrl( + dataUrl: string, +): { mime: string; bytes: Uint8Array } | null { + const m = DATA_URL_RE.exec(dataUrl); + if (!m) return null; + try { + const bin = atob(m[2]); + const bytes = new Uint8Array(bin.length); + for (let i = 0; i < bin.length; i++) bytes[i] = bin.charCodeAt(i); + return { mime: m[1], bytes }; + } catch { + return null; + } +} + +const IMAGE_MIMES = new Set(["image/jpeg", "image/png", "image/webp"]); +const VIDEO_MIMES = new Set(["video/mp4", "video/webm"]); + +async function sha256Hex(bytes: Uint8Array): Promise { + const digest = await crypto.subtle.digest("SHA-256", bytes); + return Array.from(new Uint8Array(digest), (b) => b.toString(16).padStart(2, "0")).join(""); +} + +interface UploadItemBody { + data_url?: unknown; + caption?: unknown; +} + +/** + * POST /v1/datasets/:slug/items + * + * Accepts one image or video as a base64 data URL, stores it in R2 + * under `datasets/{datasetId}/{contentHash}.{ext}`, and dedups on + * exact content hash — a byte-identical re-upload returns the + * existing item (200) rather than a new one (201). + */ +export async function uploadDatasetItemHandler( + c: Context<{ Bindings: Env; Variables: AuthVariables }>, +) { + const dataset = await loadOwnedDataset(c); + if (!dataset) return c.json({ error: "not_found" }, 404); + const session = getAuthUser(c); + + if (!c.env.CAPTURES) { + return c.json({ error: "storage_unconfigured" }, 503); + } + + const rl = await checkRateLimit(c.env.RATE_LIMIT, { + key: `datasets:upload:${session.uid}`, + limit: 500, + windowSeconds: 3600, + }); + if (!rl.ok) return c.json({ error: "rate_limited", reset_at: rl.resetAt }, 429); + + let body: UploadItemBody; + try { + body = await c.req.json(); + } catch { + return badRequest(c, "invalid_json"); + } + if (typeof body.data_url !== "string") return badRequest(c, "invalid_data_url"); + const decoded = decodeGenericDataUrl(body.data_url); + if (!decoded) return badRequest(c, "invalid_data_url", "expected data:;base64,…"); + + let kind: "image" | "video"; + let ext: string; + if (IMAGE_MIMES.has(decoded.mime)) { + kind = "image"; + ext = decoded.mime.split("/")[1] === "jpeg" ? "jpg" : decoded.mime.split("/")[1]; + if (decoded.bytes.byteLength > MAX_IMAGE_BYTES) { + return badRequest(c, "item_too_large", { max_bytes: MAX_IMAGE_BYTES }); + } + } else if (VIDEO_MIMES.has(decoded.mime)) { + kind = "video"; + ext = decoded.mime.split("/")[1]; + if (decoded.bytes.byteLength > MAX_VIDEO_BYTES) { + return badRequest(c, "item_too_large", { max_bytes: MAX_VIDEO_BYTES }); + } + } else { + return badRequest(c, "unsupported_mime", { + allowed: [...IMAGE_MIMES, ...VIDEO_MIMES], + }); + } + + let caption: string | null = null; + if (body.caption !== undefined && body.caption !== null) { + if (typeof body.caption !== "string" || body.caption.length > CAPTION_MAX) { + return badRequest(c, "invalid_caption", { max: CAPTION_MAX }); + } + caption = body.caption.trim() || null; + } + + const contentHash = await sha256Hex(decoded.bytes); + const r2Key = `datasets/${dataset.id}/${contentHash}.${ext}`; + + const result = await insertDatasetItem(c.env.DB, { + datasetId: dataset.id, + r2Key, + kind, + caption, + contentHash, + byteSize: decoded.bytes.byteLength, + width: null, + height: null, + durationSeconds: null, + status: "ready", + flagReason: null, + }); + + if (result.status === "duplicate") { + return c.json({ item: publicDatasetItem(result.existing), duplicate: true }, 200); + } + + // Written after the DB row succeeds, keyed by content hash — an R2 + // write retry (or a concurrent request that lost the DB race) is + // idempotent because it overwrites the same bytes at the same key. + await c.env.CAPTURES.put(r2Key, decoded.bytes, { + httpMetadata: { contentType: decoded.mime }, + customMetadata: { + dataset_id: String(dataset.id), + owner_id: String(session.uid), + }, + }); + + return c.json({ item: publicDatasetItem(result.item) }, 201); +} + +const IMPORT_URLS_MAX = 25; + +interface ImportUrlsBody { + urls?: unknown; +} + +/** + * POST /v1/datasets/:slug/items/import-urls + * + * Fetches each URL server-side and stores it exactly like a direct + * upload. Per-item results are reported individually — a batch of 25 + * URLs where 3 fail to resolve returns 22 successes and 3 named + * failures, never an all-or-nothing outcome. + */ +export async function importDatasetItemsHandler( + c: Context<{ Bindings: Env; Variables: AuthVariables }>, +) { + const dataset = await loadOwnedDataset(c); + if (!dataset) return c.json({ error: "not_found" }, 404); + const session = getAuthUser(c); + + if (!c.env.CAPTURES) { + return c.json({ error: "storage_unconfigured" }, 503); + } + + const rl = await checkRateLimit(c.env.RATE_LIMIT, { + key: `datasets:import:${session.uid}`, + limit: 20, + windowSeconds: 3600, + }); + if (!rl.ok) return c.json({ error: "rate_limited", reset_at: rl.resetAt }, 429); + + let body: ImportUrlsBody; + try { + body = await c.req.json(); + } catch { + return badRequest(c, "invalid_json"); + } + if (!Array.isArray(body.urls) || body.urls.length === 0) { + return badRequest(c, "invalid_urls"); + } + if (body.urls.length > IMPORT_URLS_MAX) { + return badRequest(c, "too_many_urls", { max: IMPORT_URLS_MAX }); + } + + const results: { url: string; status: string; item_id?: number }[] = []; + for (const raw of body.urls) { + const url = typeof raw === "string" ? raw : ""; + if (!url || !/^https?:\/\//.test(url)) { + results.push({ url: String(raw).slice(0, 200), status: "invalid_url" }); + continue; + } + try { + const res = await fetch(url); + if (!res.ok) { + results.push({ url, status: `fetch_failed_${res.status}` }); + continue; + } + const mime = (res.headers.get("content-type") || "").split(";")[0].trim(); + const buf = new Uint8Array(await res.arrayBuffer()); + let kind: "image" | "video"; + let ext: string; + if (IMAGE_MIMES.has(mime)) { + kind = "image"; + ext = mime.split("/")[1] === "jpeg" ? "jpg" : mime.split("/")[1]; + if (buf.byteLength > MAX_IMAGE_BYTES) { + results.push({ url, status: "too_large" }); + continue; + } + } else if (VIDEO_MIMES.has(mime)) { + kind = "video"; + ext = mime.split("/")[1]; + if (buf.byteLength > MAX_VIDEO_BYTES) { + results.push({ url, status: "too_large" }); + continue; + } + } else { + results.push({ url, status: "unsupported_mime" }); + continue; + } + const contentHash = await sha256Hex(buf); + const r2Key = `datasets/${dataset.id}/${contentHash}.${ext}`; + const inserted = await insertDatasetItem(c.env.DB, { + datasetId: dataset.id, + r2Key, + kind, + caption: null, + contentHash, + byteSize: buf.byteLength, + width: null, + height: null, + durationSeconds: null, + status: "ready", + flagReason: null, + }); + if (inserted.status === "duplicate") { + results.push({ url, status: "duplicate", item_id: inserted.existing.id }); + continue; + } + await c.env.CAPTURES.put(r2Key, buf, { + httpMetadata: { contentType: mime }, + customMetadata: { dataset_id: String(dataset.id), owner_id: String(session.uid) }, + }); + results.push({ url, status: "ok", item_id: inserted.item.id }); + } catch (e) { + results.push({ url, status: "error" }); + console.error("import_url_failed", url, e); + } + } + const okCount = results.filter((r) => r.status === "ok").length; + return c.json({ results, imported: okCount, total: results.length }, 201); +} + +const ITEM_LIST_DEFAULT = 40; +const ITEM_LIST_MAX = 100; + +/** GET /v1/datasets/:slug/items?limit=&before= */ +export async function listDatasetItemsHandler( + c: Context<{ Bindings: Env; Variables: AuthVariables }>, +) { + const dataset = await loadOwnedDataset(c); + if (!dataset) return c.json({ error: "not_found" }, 404); + const url = new URL(c.req.url); + const limitRaw = url.searchParams.get("limit"); + let limit = ITEM_LIST_DEFAULT; + if (limitRaw) { + const n = parseInt(limitRaw, 10); + if (!Number.isFinite(n) || n < 1 || n > ITEM_LIST_MAX) { + return badRequest(c, "invalid_limit", { max: ITEM_LIST_MAX }); + } + limit = n; + } + const beforeRaw = url.searchParams.get("before"); + let before: number | undefined; + if (beforeRaw) { + const n = parseInt(beforeRaw, 10); + if (!Number.isFinite(n) || n < 0) return badRequest(c, "invalid_before"); + before = n; + } + const rows = await listDatasetItems(c.env.DB, dataset.id, limit, before); + return c.json({ + items: rows.map(publicDatasetItem), + next_before: rows.length === limit ? rows[rows.length - 1].id : null, + }); +} + +/** DELETE /v1/datasets/:slug/items/:itemId — owner-only. */ +export async function deleteDatasetItemHandler( + c: Context<{ Bindings: Env; Variables: AuthVariables }>, +) { + const dataset = await loadOwnedDataset(c); + if (!dataset) return c.json({ error: "not_found" }, 404); + const itemId = parseInt(c.req.param("itemId") || "", 10); + if (!itemId || Number.isNaN(itemId)) return badRequest(c, "invalid_item_id"); + const item = await getDatasetItem(c.env.DB, itemId); + if (!item || item.dataset_id !== dataset.id) return c.json({ error: "not_found" }, 404); + + const deleted = await deleteDatasetItem(c.env.DB, itemId, dataset.id); + if (!deleted) return c.json({ error: "not_found" }, 404); + if (c.env.CAPTURES) { + try { + await c.env.CAPTURES.delete(item.r2_key); + } catch (e) { + console.error("dataset_item_r2_delete_failed", item.r2_key, e); + } + } + return c.json({ ok: true }); +} + +/** + * GET /v1/datasets/:slug/items/:itemId/file + * + * Owner-only proxy read. Datasets are private by default and, unlike + * models, have no cross-owner slug lookup yet (see getDatasetHandler) — + * a `visibility: 'public'` dataset is a stated intent for a future + * public browsing surface, not yet a working cross-owner read path. + * Building that means resolving a slug globally rather than per-owner, + * which needs its own uniqueness scope; deliberately deferred rather + * than half-wired here. + */ +export async function getDatasetItemFileHandler( + c: Context<{ Bindings: Env; Variables: AuthVariables }>, +) { + const dataset = await loadOwnedDataset(c); + if (!dataset) return c.json({ error: "not_found" }, 404); + if (!c.env.CAPTURES) return c.json({ error: "storage_unconfigured" }, 503); + + const itemId = parseInt(c.req.param("itemId") || "", 10); + if (!itemId || Number.isNaN(itemId)) return badRequest(c, "invalid_item_id"); + const item = await getDatasetItem(c.env.DB, itemId); + if (!item || item.dataset_id !== dataset.id) return c.json({ error: "not_found" }, 404); + + const obj = await c.env.CAPTURES.get(item.r2_key); + if (!obj) return c.json({ error: "not_found" }, 404); + return new Response(obj.body, { + headers: { + "Content-Type": obj.httpMetadata?.contentType ?? "application/octet-stream", + "Cache-Control": "private, max-age=3600", + }, + }); +} + +// --------------------------------------------------------------------------- +// Training +// --------------------------------------------------------------------------- + +interface TrainBody { + base_model?: unknown; + training_method?: unknown; + render_price_tokens?: unknown; + idempotency_key?: unknown; +} + +const RENDER_PRICE_MAX = 100_000; + +/** + * POST /v1/datasets/:slug/train + * + * Debit → enqueue. A render_models shell is created eagerly (private + * visibility) so the job has a version to publish onto the moment + * training succeeds; the creator makes it public later from the model + * registry, same as any other model. + */ +export async function trainHandler( + c: Context<{ Bindings: Env; Variables: AuthVariables }>, +) { + const dataset = await loadOwnedDataset(c); + if (!dataset) return c.json({ error: "not_found" }, 404); + const session = getAuthUser(c); + + if (dataset.item_count < MIN_TRAIN_ITEMS) { + return c.json( + { error: "not_enough_items", detail: { min: MIN_TRAIN_ITEMS, have: dataset.item_count } }, + 422, + ); + } + + let body: TrainBody; + try { + body = await c.req.json(); + } catch { + return badRequest(c, "invalid_json"); + } + if (!isTrainableBaseModel(body.base_model)) { + return badRequest(c, "invalid_base_model", { + allowed: Object.keys(TRAINABLE_BASE_MODELS), + }); + } + const trainingMethod = + typeof body.training_method === "string" ? body.training_method : ""; + if (!["lora", "dreambooth", "full_finetune"].includes(trainingMethod)) { + return badRequest(c, "invalid_training_method"); + } + let renderPriceTokens = 25; + if (body.render_price_tokens !== undefined && body.render_price_tokens !== null) { + const n = typeof body.render_price_tokens === "number" ? Math.trunc(body.render_price_tokens) : NaN; + if (!Number.isFinite(n) || n < 0 || n > RENDER_PRICE_MAX) { + return badRequest(c, "invalid_render_price_tokens", { max: RENDER_PRICE_MAX }); + } + renderPriceTokens = n; + } + + const idempotencyKey = + typeof body.idempotency_key === "string" && body.idempotency_key.length > 0 + ? `train:${session.uid}:${body.idempotency_key.slice(0, 128)}` + : `train:${session.uid}:${crypto.randomUUID()}`; + + const existingJob = await getTrainingJobByIdempotencyKey(c.env.DB, idempotencyKey); + if (existingJob) { + return c.json({ job: publicTrainingJob(existingJob), replayed: true }); + } + + const priceTokens = trainingPriceTokens(trainingMethod); + const debit = await applyLedgerEntry(c.env.DB, { + userId: session.uid, + delta: -priceTokens, + kind: "debit", + idempotencyKey: `${idempotencyKey}:debit`, + refKind: "training_job", + refId: null, + memo: `train: ${dataset.title} (${TRAINABLE_BASE_MODELS[body.base_model].label})`, + }); + if (debit.status === "insufficient") { + return c.json( + { error: "insufficient_balance", balance: debit.balance, shortfall: debit.shortfall }, + 402, + ); + } + + const modelSlug = await uniqueModelSlug(c.env.DB, `${dataset.title} model`); + const model = await insertModel(c.env.DB, { + ownerId: session.uid, + slug: modelSlug, + title: `${dataset.title} — ${TRAINABLE_BASE_MODELS[body.base_model].label}`, + description: dataset.description, + kind: "image", + provider: "fal_custom", + visibility: "private", + }); + + const job = await insertTrainingJob(c.env.DB, { + userId: session.uid, + datasetId: dataset.id, + baseModel: body.base_model, + trainingMethod, + priceTokens, + renderPriceTokens, + modelId: model.id, + idempotencyKey, + }); + + // TRAINING_MOCK completes synchronously in-request (no network, no + // cron wait) — see ai/training.ts for why this mirrors RENDER_MOCK + // rather than simulating cron timing. completeJob's UPDATE is + // guarded on status='training' (matching the real dispatch/poll + // state machine), so the mock path must transition the job through + // 'training' first — completeJob is not itself a queued->succeeded + // shortcut. + if (isTrainingMock(c.env)) { + await markTrainingSubmitted(c.env.DB, job.id, `mock-request-${job.id}`); + const mock = await runMockTraining(dataset.slug, job.id); + const submitted = await getTrainingJobById(c.env.DB, job.id); + await completeJob(c.env, submitted ?? job, mock.weightsRef, dataset.description); + const finished = await getTrainingJobById(c.env.DB, job.id); + return c.json({ job: publicTrainingJob(finished ?? job) }, 201); + } + + return c.json({ job: publicTrainingJob(job) }, 201); +} + +const JOB_LIST_DEFAULT = 20; +const JOB_LIST_MAX = 50; + +/** GET /v1/training?limit=&before= — the caller's own jobs. */ +export async function myTrainingJobsHandler( + c: Context<{ Bindings: Env; Variables: AuthVariables }>, +) { + const session = getAuthUser(c); + const url = new URL(c.req.url); + const limitRaw = url.searchParams.get("limit"); + let limit = JOB_LIST_DEFAULT; + if (limitRaw) { + const n = parseInt(limitRaw, 10); + if (!Number.isFinite(n) || n < 1 || n > JOB_LIST_MAX) { + return badRequest(c, "invalid_limit", { max: JOB_LIST_MAX }); + } + limit = n; + } + const beforeRaw = url.searchParams.get("before"); + let before: number | undefined; + if (beforeRaw) { + const n = parseInt(beforeRaw, 10); + if (!Number.isFinite(n) || n < 0) return badRequest(c, "invalid_before"); + before = n; + } + const rows = await listTrainingJobsByUser(c.env.DB, session.uid, limit, before); + return c.json({ + jobs: rows.map(publicTrainingJob), + next_before: rows.length === limit ? rows[rows.length - 1].id : null, + }); +} + +/** GET /v1/datasets/:slug/training — jobs run against this dataset. */ +export async function datasetTrainingJobsHandler( + c: Context<{ Bindings: Env; Variables: AuthVariables }>, +) { + const dataset = await loadOwnedDataset(c); + if (!dataset) return c.json({ error: "not_found" }, 404); + const rows = await listTrainingJobsByDataset(c.env.DB, dataset.id); + return c.json({ jobs: rows.map(publicTrainingJob) }); +} + +/** GET /v1/training/:id — owner-only job detail. */ +export async function getTrainingJobHandler( + c: Context<{ Bindings: Env; Variables: AuthVariables }>, +) { + const session = getAuthUser(c); + const id = parseInt(c.req.param("id") || "", 10); + if (!id || Number.isNaN(id)) return badRequest(c, "invalid_id"); + const job = await getTrainingJobById(c.env.DB, id); + if (!job) return c.json({ error: "not_found" }, 404); + if (job.user_id !== session.uid) return c.json({ error: "forbidden" }, 403); + return c.json({ job: publicTrainingJob(job) }); +} diff --git a/workers/api/src/db/datasets.ts b/workers/api/src/db/datasets.ts new file mode 100644 index 0000000..6904102 --- /dev/null +++ b/workers/api/src/db/datasets.ts @@ -0,0 +1,548 @@ +import type { D1Database } from "@cloudflare/workers-types"; + +/** + * Dataset Library: a creator's own curated images/video, and the + * training jobs that consume a dataset to produce a render-model + * version. See migrations/0020_dataset_library.sql. + */ + +export const RIGHTS_DECLARATIONS = [ + "own", + "licensed", + "public_domain", +] as const; +export type RightsDeclaration = (typeof RIGHTS_DECLARATIONS)[number]; + +export function isRightsDeclaration(v: unknown): v is RightsDeclaration { + return ( + typeof v === "string" && + (RIGHTS_DECLARATIONS as readonly string[]).includes(v) + ); +} + +export const DATASET_VISIBILITIES = ["private", "public"] as const; +export type DatasetVisibility = (typeof DATASET_VISIBILITIES)[number]; + +export function isDatasetVisibility(v: unknown): v is DatasetVisibility { + return ( + typeof v === "string" && + (DATASET_VISIBILITIES as readonly string[]).includes(v) + ); +} + +export const ITEM_KINDS = ["image", "video"] as const; +export type ItemKind = (typeof ITEM_KINDS)[number]; + +export function isItemKind(v: unknown): v is ItemKind { + return typeof v === "string" && (ITEM_KINDS as readonly string[]).includes(v); +} + +export interface DatasetRow { + id: number; + owner_id: number; + slug: string; + title: string; + description: string | null; + rights_declaration: string; + visibility: string; + item_count: number; + created_at: number; + updated_at: number; +} + +export async function insertDataset( + db: D1Database, + input: { + ownerId: number; + slug: string; + title: string; + description: string | null; + rightsDeclaration: RightsDeclaration; + visibility: DatasetVisibility; + }, +): Promise { + const now = Math.floor(Date.now() / 1000); + const row = await db + .prepare( + `INSERT INTO datasets + (owner_id, slug, title, description, rights_declaration, + visibility, item_count, created_at, updated_at) + VALUES (?, ?, ?, ?, ?, ?, 0, ?, ?) + RETURNING *`, + ) + .bind( + input.ownerId, + input.slug, + input.title, + input.description, + input.rightsDeclaration, + input.visibility, + now, + now, + ) + .first(); + if (!row) throw new Error("insertDataset returned no row"); + return row; +} + +export async function getDatasetByOwnerAndSlug( + db: D1Database, + ownerId: number, + slug: string, +): Promise { + return await db + .prepare(`SELECT * FROM datasets WHERE owner_id = ? AND slug = ?`) + .bind(ownerId, slug) + .first(); +} + +export async function getDatasetById( + db: D1Database, + id: number, +): Promise { + return await db + .prepare(`SELECT * FROM datasets WHERE id = ?`) + .bind(id) + .first(); +} + +export async function updateDataset( + db: D1Database, + id: number, + patch: { + title?: string; + description?: string | null; + rightsDeclaration?: RightsDeclaration; + visibility?: DatasetVisibility; + }, +): Promise { + const sets: string[] = []; + const binds: unknown[] = []; + if (patch.title !== undefined) { + sets.push("title = ?"); + binds.push(patch.title); + } + if (patch.description !== undefined) { + sets.push("description = ?"); + binds.push(patch.description); + } + if (patch.rightsDeclaration !== undefined) { + sets.push("rights_declaration = ?"); + binds.push(patch.rightsDeclaration); + } + if (patch.visibility !== undefined) { + sets.push("visibility = ?"); + binds.push(patch.visibility); + } + if (sets.length === 0) { + return await db + .prepare(`SELECT * FROM datasets WHERE id = ?`) + .bind(id) + .first(); + } + sets.push("updated_at = ?"); + binds.push(Math.floor(Date.now() / 1000), id); + return await db + .prepare(`UPDATE datasets SET ${sets.join(", ")} WHERE id = ? RETURNING *`) + .bind(...binds) + .first(); +} + +export async function listDatasetsByOwner( + db: D1Database, + ownerId: number, +): Promise { + const { results } = await db + .prepare( + `SELECT * FROM datasets WHERE owner_id = ? ORDER BY updated_at DESC LIMIT 200`, + ) + .bind(ownerId) + .all(); + return results ?? []; +} + +export function publicDataset(d: DatasetRow) { + return { + id: d.id, + slug: d.slug, + title: d.title, + description: d.description, + rights_declaration: d.rights_declaration, + visibility: d.visibility, + item_count: d.item_count, + created_at: d.created_at, + updated_at: d.updated_at, + }; +} + +// --------------------------------------------------------------------------- +// Items +// --------------------------------------------------------------------------- + +export interface DatasetItemRow { + id: number; + dataset_id: number; + r2_key: string; + kind: string; + caption: string | null; + content_hash: string; + byte_size: number; + width: number | null; + height: number | null; + duration_seconds: number | null; + status: string; + flag_reason: string | null; + created_at: number; +} + +export type InsertItemResult = + | { status: "inserted"; item: DatasetItemRow } + | { status: "duplicate"; existing: DatasetItemRow }; + +function isUniqueViolation(e: unknown): boolean { + const msg = e instanceof Error ? e.message : String(e); + return /UNIQUE constraint failed/i.test(msg); +} + +/** + * Inserts an item and bumps the dataset's denormalised item_count in + * the same D1 batch. UNIQUE(dataset_id, content_hash) is the dedup + * mechanism — a byte-identical re-upload/import fails the insert and + * this returns the existing row instead, rather than double-counting. + */ +export async function insertDatasetItem( + db: D1Database, + input: { + datasetId: number; + r2Key: string; + kind: ItemKind; + caption: string | null; + contentHash: string; + byteSize: number; + width: number | null; + height: number | null; + durationSeconds: number | null; + status: "ready" | "flagged"; + flagReason: string | null; + }, +): Promise { + const now = Math.floor(Date.now() / 1000); + const insertItem = db + .prepare( + `INSERT INTO dataset_items + (dataset_id, r2_key, kind, caption, content_hash, byte_size, + width, height, duration_seconds, status, flag_reason, created_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + RETURNING *`, + ) + .bind( + input.datasetId, + input.r2Key, + input.kind, + input.caption, + input.contentHash, + input.byteSize, + input.width, + input.height, + input.durationSeconds, + input.status, + input.flagReason, + now, + ); + const bumpCount = db + .prepare( + `UPDATE datasets SET item_count = item_count + 1, updated_at = ? WHERE id = ?`, + ) + .bind(now, input.datasetId); + + try { + const results = await db.batch([insertItem, bumpCount]); + const item = results[0]?.results?.[0]; + if (!item) throw new Error("insertDatasetItem returned no row"); + return { status: "inserted", item }; + } catch (e) { + if (isUniqueViolation(e)) { + const existing = await db + .prepare( + `SELECT * FROM dataset_items WHERE dataset_id = ? AND content_hash = ?`, + ) + .bind(input.datasetId, input.contentHash) + .first(); + if (existing) return { status: "duplicate", existing }; + } + throw e; + } +} + +export async function listDatasetItems( + db: D1Database, + datasetId: number, + limit: number, + before?: number, +): Promise { + const sql = before + ? `SELECT * FROM dataset_items WHERE dataset_id = ? AND id < ? + ORDER BY id DESC LIMIT ?` + : `SELECT * FROM dataset_items WHERE dataset_id = ? + ORDER BY id DESC LIMIT ?`; + const stmt = db.prepare(sql); + const bound = before + ? stmt.bind(datasetId, before, limit) + : stmt.bind(datasetId, limit); + const { results } = await bound.all(); + return results ?? []; +} + +export async function getDatasetItem( + db: D1Database, + itemId: number, +): Promise { + return await db + .prepare(`SELECT * FROM dataset_items WHERE id = ?`) + .bind(itemId) + .first(); +} + +/** Deletes an item and decrements item_count in the same batch. */ +export async function deleteDatasetItem( + db: D1Database, + itemId: number, + datasetId: number, +): Promise { + const now = Math.floor(Date.now() / 1000); + const del = db + .prepare(`DELETE FROM dataset_items WHERE id = ? AND dataset_id = ?`) + .bind(itemId, datasetId); + const decCount = db + .prepare( + `UPDATE datasets + SET item_count = MAX(0, item_count - 1), updated_at = ? + WHERE id = ?`, + ) + .bind(now, datasetId); + const results = await db.batch([del, decCount]); + return (results[0]?.meta?.changes ?? 0) > 0; +} + +export function publicDatasetItem(i: DatasetItemRow) { + return { + id: i.id, + kind: i.kind, + caption: i.caption, + byte_size: i.byte_size, + width: i.width, + height: i.height, + duration_seconds: i.duration_seconds, + status: i.status, + flag_reason: i.flag_reason, + created_at: i.created_at, + }; +} + +// --------------------------------------------------------------------------- +// Training jobs +// --------------------------------------------------------------------------- + +export const TRAINING_JOB_STATUSES = [ + "queued", + "training", + "succeeded", + "failed", +] as const; +export type TrainingJobStatus = (typeof TRAINING_JOB_STATUSES)[number]; + +export interface TrainingJobRow { + id: number; + user_id: number; + dataset_id: number; + base_model: string; + training_method: string; + price_tokens: number; + render_price_tokens: number; + status: string; + provider_job_id: string | null; + model_id: number | null; + model_version_id: number | null; + error_code: string | null; + idempotency_key: string; + created_at: number; + started_at: number | null; + finished_at: number | null; +} + +export async function insertTrainingJob( + db: D1Database, + input: { + userId: number; + datasetId: number; + baseModel: string; + trainingMethod: string; + priceTokens: number; + renderPriceTokens: number; + modelId: number; + idempotencyKey: string; + }, +): Promise { + const now = Math.floor(Date.now() / 1000); + const row = await db + .prepare( + `INSERT INTO training_jobs + (user_id, dataset_id, base_model, training_method, price_tokens, + render_price_tokens, status, model_id, idempotency_key, created_at) + VALUES (?, ?, ?, ?, ?, ?, 'queued', ?, ?, ?) + RETURNING *`, + ) + .bind( + input.userId, + input.datasetId, + input.baseModel, + input.trainingMethod, + input.priceTokens, + input.renderPriceTokens, + input.modelId, + input.idempotencyKey, + now, + ) + .first(); + if (!row) throw new Error("insertTrainingJob returned no row"); + return row; +} + +export async function getTrainingJobByIdempotencyKey( + db: D1Database, + key: string, +): Promise { + return await db + .prepare(`SELECT * FROM training_jobs WHERE idempotency_key = ?`) + .bind(key) + .first(); +} + +export async function getTrainingJobById( + db: D1Database, + id: number, +): Promise { + return await db + .prepare(`SELECT * FROM training_jobs WHERE id = ?`) + .bind(id) + .first(); +} + +export async function listTrainingJobsByUser( + db: D1Database, + userId: number, + limit: number, + before?: number, +): Promise { + const sql = before + ? `SELECT * FROM training_jobs WHERE user_id = ? AND id < ? + ORDER BY id DESC LIMIT ?` + : `SELECT * FROM training_jobs WHERE user_id = ? + ORDER BY id DESC LIMIT ?`; + const stmt = db.prepare(sql); + const bound = before + ? stmt.bind(userId, before, limit) + : stmt.bind(userId, limit); + const { results } = await bound.all(); + return results ?? []; +} + +export async function listTrainingJobsByDataset( + db: D1Database, + datasetId: number, +): Promise { + const { results } = await db + .prepare( + `SELECT * FROM training_jobs WHERE dataset_id = ? ORDER BY id DESC LIMIT 200`, + ) + .bind(datasetId) + .all(); + return results ?? []; +} + +/** Cron dispatch: jobs that haven't been submitted to the provider yet. */ +export async function getQueuedTrainingJobs( + db: D1Database, + limit: number, +): Promise { + const { results } = await db + .prepare( + `SELECT * FROM training_jobs WHERE status = 'queued' ORDER BY id ASC LIMIT ?`, + ) + .bind(limit) + .all(); + return results ?? []; +} + +/** Cron poll: jobs already submitted, awaiting a provider result. */ +export async function getInFlightTrainingJobs( + db: D1Database, + limit: number, +): Promise { + const { results } = await db + .prepare( + `SELECT * FROM training_jobs WHERE status = 'training' ORDER BY id ASC LIMIT ?`, + ) + .bind(limit) + .all(); + return results ?? []; +} + +export async function markTrainingSubmitted( + db: D1Database, + id: number, + providerJobId: string, +): Promise { + await db + .prepare( + `UPDATE training_jobs + SET status = 'training', provider_job_id = ?, started_at = ? + WHERE id = ? AND status = 'queued'`, + ) + .bind(providerJobId, Math.floor(Date.now() / 1000), id) + .run(); +} + +export async function finishTrainingJob( + db: D1Database, + id: number, + patch: { + status: "succeeded" | "failed"; + modelVersionId?: number | null; + errorCode?: string | null; + }, +): Promise { + return await db + .prepare( + `UPDATE training_jobs + SET status = ?, model_version_id = ?, error_code = ?, finished_at = ? + WHERE id = ? AND status = 'training' + RETURNING *`, + ) + .bind( + patch.status, + patch.modelVersionId ?? null, + patch.errorCode ?? null, + Math.floor(Date.now() / 1000), + id, + ) + .first(); +} + +export function publicTrainingJob(j: TrainingJobRow) { + return { + id: j.id, + dataset_id: j.dataset_id, + base_model: j.base_model, + training_method: j.training_method, + price_tokens: j.price_tokens, + render_price_tokens: j.render_price_tokens, + status: j.status, + model_id: j.model_id, + model_version_id: j.model_version_id, + error_code: j.error_code, + created_at: j.created_at, + started_at: j.started_at, + finished_at: j.finished_at, + }; +} diff --git a/workers/api/src/index.ts b/workers/api/src/index.ts index 03c3be6..71648fc 100644 --- a/workers/api/src/index.ts +++ b/workers/api/src/index.ts @@ -95,6 +95,22 @@ import { throwHandler, clientErrorHandler, } from "./internal/stats"; +import { + createDatasetHandler, + myDatasetsHandler, + getDatasetHandler, + patchDatasetHandler, + uploadDatasetItemHandler, + importDatasetItemsHandler, + listDatasetItemsHandler, + deleteDatasetItemHandler, + getDatasetItemFileHandler, + trainHandler, + myTrainingJobsHandler, + datasetTrainingJobsHandler, + getTrainingJobHandler, +} from "./datasets/handlers"; +import { dispatchQueuedTrainingJobs, pollTrainingJobs } from "./ai/training"; const app = new Hono<{ Bindings: Env; Variables: AuthVariables }>(); @@ -237,6 +253,41 @@ v1.post("/models/:slug{[a-z0-9-]+}/render", requireAuth, renderHandler); v1.get("/jobs", requireAuth, myJobsHandler); v1.get("/jobs/:id{[0-9]+}", requireAuth, getJobHandler); +// Dataset Library. Everything here is owner-scoped for now (see +// datasets/handlers.ts — datasets are private by default and there's +// no cross-owner slug lookup yet). Literal segments ("mine") precede +// the `:slug` pattern for the same Hono route-ordering reason as the +// galleries/tokens routes above. +v1.post("/datasets", requireAuth, createDatasetHandler); +v1.get("/datasets/mine", requireAuth, myDatasetsHandler); +v1.get("/datasets/:slug{[a-z0-9-]+}", requireAuth, getDatasetHandler); +v1.patch("/datasets/:slug{[a-z0-9-]+}", requireAuth, patchDatasetHandler); +v1.post("/datasets/:slug{[a-z0-9-]+}/items", requireAuth, uploadDatasetItemHandler); +v1.post( + "/datasets/:slug{[a-z0-9-]+}/items/import-urls", + requireAuth, + importDatasetItemsHandler, +); +v1.get("/datasets/:slug{[a-z0-9-]+}/items", requireAuth, listDatasetItemsHandler); +v1.delete( + "/datasets/:slug{[a-z0-9-]+}/items/:itemId{[0-9]+}", + requireAuth, + deleteDatasetItemHandler, +); +v1.get( + "/datasets/:slug{[a-z0-9-]+}/items/:itemId{[0-9]+}/file", + requireAuth, + getDatasetItemFileHandler, +); +v1.post("/datasets/:slug{[a-z0-9-]+}/train", requireAuth, trainHandler); +v1.get( + "/datasets/:slug{[a-z0-9-]+}/training", + requireAuth, + datasetTrainingJobsHandler, +); +v1.get("/training", requireAuth, myTrainingJobsHandler); +v1.get("/training/:id{[0-9]+}", requireAuth, getTrainingJobHandler); + // Admin-only observability surface. requireAuth runs first so // requireAdmin can read the session. _throw is a forced 500 used // to verify the request_id round-trips end to end. @@ -284,13 +335,27 @@ export default { env: Env, ctx: ExecutionContext, ): Promise { - // "* * * * *" → uptime probe; "0 4 * * *" → daily audit + prune. + // "* * * * *" → uptime probe + training dispatch/poll; + // "0 4 * * *" → daily audit + prune. if (event.cron === "* * * * *") { ctx.waitUntil( runUptimeProbe(env).catch((e) => console.error("uptime_probe_failed", (e as Error).message), ), ); + // Training jobs are genuinely async at the provider (minutes to + // hours) — dispatch submits newly-queued jobs, poll checks + // in-flight ones for a result. See ai/training.ts. + ctx.waitUntil( + dispatchQueuedTrainingJobs(env).catch((e) => + console.error("training_dispatch_failed", (e as Error).message), + ), + ); + ctx.waitUntil( + pollTrainingJobs(env).catch((e) => + console.error("training_poll_failed", (e as Error).message), + ), + ); return; } ctx.waitUntil(runFrozenAudit(env)); diff --git a/workers/api/src/lib/slug.ts b/workers/api/src/lib/slug.ts index 3392f59..3bbfe0f 100644 --- a/workers/api/src/lib/slug.ts +++ b/workers/api/src/lib/slug.ts @@ -25,6 +25,26 @@ export async function uniqueProjectSlug( return `${baseSlug}-${crypto.randomUUID().slice(0, 6)}`; } +/** + * Datasets' `slug` is UNIQUE per-owner, same shape as projects. + */ +export async function uniqueDatasetSlug( + db: D1Database, + ownerId: number, + base: string, +): Promise { + const baseSlug = slugify(base); + for (let i = 0; i < 50; i++) { + const candidate = i === 0 ? baseSlug : `${baseSlug}-${i + 1}`; + const taken = await db + .prepare("SELECT 1 FROM datasets WHERE owner_id = ? AND slug = ?") + .bind(ownerId, candidate) + .first(); + if (!taken) return candidate; + } + return `${baseSlug}-${crypto.randomUUID().slice(0, 6)}`; +} + /** * Galleries' `slug` is globally UNIQUE (not per-owner like projects). * Mirrors uniqueProjectSlug's collision-suffix loop with a random diff --git a/workers/api/src/types.ts b/workers/api/src/types.ts index 508bfc2..448e823 100644 --- a/workers/api/src/types.ts +++ b/workers/api/src/types.ts @@ -96,6 +96,14 @@ export interface Env { // in scripts/smoke_api.mjs without an API key or Workers AI binding. // Production must NEVER set RENDER_MOCK. RENDER_MOCK?: string; + // --- Dataset Library / training jobs ------------------------------------ + // OPT-IN ONLY, same convention as RENDER_MOCK. When exactly "1", a + // training job (POST /v1/datasets/:slug/train) completes synchronously + // in-request with canned weights instead of submitting to fal.ai's + // queue API — see ai/training.ts for why this differs from the + // dispatch/poll cron path used in production. Production must NEVER + // set TRAINING_MOCK. + TRAINING_MOCK?: string; } export interface ProjectMintState { diff --git a/workers/api/test/units.test.mjs b/workers/api/test/units.test.mjs index 1a534fa..8c7e642 100644 --- a/workers/api/test/units.test.mjs +++ b/workers/api/test/units.test.mjs @@ -33,6 +33,8 @@ export { parseAllowlist } from ${JSON.stringify(join(srcDir, "auth", "admin"))}; export { parseFeedCursor, encodeFeedCursor } from ${JSON.stringify(join(srcDir, "db", "events"))}; export { lifetimeDeltas } from ${JSON.stringify(join(srcDir, "db", "tokens"))}; export { providerAllowedForKind, isTrainingMethod } from ${JSON.stringify(join(srcDir, "db", "render"))}; +export { isRightsDeclaration, isDatasetVisibility, isItemKind } from ${JSON.stringify(join(srcDir, "db", "datasets"))}; +export { trainingPriceTokens, isTrainableBaseModel } from ${JSON.stringify(join(srcDir, "ai", "training"))}; `, ); @@ -341,6 +343,47 @@ test("providerAllowedForKind: code models still reject fal_custom", () => { assert.equal(m.providerAllowedForKind("code", "fal_custom"), false); }); +// --------------------------------------------------------------------------- +// Dataset Library. +// --------------------------------------------------------------------------- +test("isRightsDeclaration: accepts the three documented values only", () => { + assert.equal(m.isRightsDeclaration("own"), true); + assert.equal(m.isRightsDeclaration("licensed"), true); + assert.equal(m.isRightsDeclaration("public_domain"), true); + assert.equal(m.isRightsDeclaration("stolen"), false); + assert.equal(m.isRightsDeclaration(undefined), false); +}); + +test("isDatasetVisibility: private and public only", () => { + assert.equal(m.isDatasetVisibility("private"), true); + assert.equal(m.isDatasetVisibility("public"), true); + assert.equal(m.isDatasetVisibility("unlisted"), false); +}); + +test("isItemKind: image and video only", () => { + assert.equal(m.isItemKind("image"), true); + assert.equal(m.isItemKind("video"), true); + assert.equal(m.isItemKind("audio"), false); +}); + +// Regression this prevents: a training request for an unlisted base +// model silently hitting a guessed/unverified fal.ai endpoint slug +// instead of a clean 400 — see ai/training.ts's TRAINABLE_BASE_MODELS +// comment for why only one base is wired end-to-end today. +test("isTrainableBaseModel: only the verified allowlist, not any string", () => { + assert.equal(m.isTrainableBaseModel("flux-dev"), true); + assert.equal(m.isTrainableBaseModel("flux-2-dev"), false); + assert.equal(m.isTrainableBaseModel("made-up-model"), false); +}); + +test("trainingPriceTokens: tiered by method, lora cheapest", () => { + const lora = m.trainingPriceTokens("lora"); + const dreambooth = m.trainingPriceTokens("dreambooth"); + const full = m.trainingPriceTokens("full_finetune"); + assert.ok(lora < dreambooth); + assert.ok(dreambooth < full); +}); + test("isTrainingMethod: accepts the four documented methods only", () => { assert.equal(m.isTrainingMethod("lora"), true); assert.equal(m.isTrainingMethod("dreambooth"), true);