feat: auto-create private HF dataset using OAuth token

- No more need for HF_TOKEN secret; the logged-in user's OAuth token
(with write-repos scope) is used for all HF API calls
- Dataset ID is derived from SPACE_ID (<space>-data) when not explicit
- Dataset is created automatically (private) on first write operation
- Last valid OAuth token is cached in memory for background operations
(debounced document push)

Made-with: Cursor

README.md

@@ -7,6 +7,8 @@ sdk: docker
 app_port: 8080
 pinned: false
 hf_oauth: true
+hf_oauth_scopes:
+  - write-repos
 ---
 
 # Collab Editor

backend/src/hf-storage.ts

@@ -1,12 +1,60 @@
-import { uploadFile, downloadFile, type RepoDesignation } from "@huggingface/hub";
+import { uploadFile, downloadFile, createRepo, type RepoDesignation } from "@huggingface/hub";
 
-const HF_TOKEN = process.env.HF_TOKEN || "";
-const HF_DATASET_ID = process.env.HF_DATASET_ID || "";
+const SPACE_ID = process.env.SPACE_ID || "";
+
+// Derive dataset ID: explicit env > "<space-id>-data"
+const HF_DATASET_ID =
+  process.env.HF_DATASET_ID ||
+  (SPACE_ID ? `${SPACE_ID}-data` : "");
 
 const repo: RepoDesignation = { type: "dataset", name: HF_DATASET_ID };
 
+// Fall-back token from env (optional); overridden by user OAuth tokens at runtime
+const ENV_TOKEN = process.env.HF_TOKEN || "";
+
+// Last known valid OAuth token, kept in memory for background operations
+let _cachedToken = ENV_TOKEN;
+
+export function getDatasetId(): string {
+  return HF_DATASET_ID;
+}
+
+export function setUserToken(token: string): void {
+  if (token) _cachedToken = token;
+}
+
+function getToken(explicit?: string): string {
+  return explicit || _cachedToken || ENV_TOKEN;
+}
+
 export function isHfStorageEnabled(): boolean {
-  return Boolean(HF_TOKEN && HF_DATASET_ID);
+  return Boolean(HF_DATASET_ID && (_cachedToken || ENV_TOKEN));
+}
+
+/**
+ * Ensure the backing dataset exists on HF (private).
+ * Safe to call multiple times; ignores "already exists".
+ */
+let _datasetReady = false;
+
+export async function ensureDatasetExists(token?: string): Promise<void> {
+  if (_datasetReady || !HF_DATASET_ID) return;
+
+  const accessToken = getToken(token);
+  if (!accessToken) return;
+
+  try {
+    await createRepo({ repo, private: true, accessToken });
+    console.log(`[hf-storage] created private dataset ${HF_DATASET_ID}`);
+  } catch (err: any) {
+    if (err?.statusCode === 409 || String(err?.message).includes("already")) {
+      console.log(`[hf-storage] dataset ${HF_DATASET_ID} already exists`);
+    } else {
+      console.error(`[hf-storage] failed to create dataset:`, err);
+      return;
+    }
+  }
+  _datasetReady = true;
 }
 
 function sanitizeName(name: string): string {
@@ -18,13 +66,15 @@ function sanitizeName(name: string): string {
 export async function uploadImageToHf(
   buffer: Buffer,
   filename: string,
+  token?: string,
 ): Promise<string> {
-  const path = `images/${filename}`;
+  await ensureDatasetExists(token);
 
+  const path = `images/${filename}`;
   await uploadFile({
     repo,
     file: { path, content: new Blob([new Uint8Array(buffer)]) },
-    accessToken: HF_TOKEN,
+    accessToken: getToken(token),
     commitTitle: `upload image ${filename}`,
   });
 
@@ -36,10 +86,6 @@ export async function uploadImageToHf(
 const dirtyDocs = new Map<string, { state: Buffer; timer: NodeJS.Timeout }>();
 const DEBOUNCE_MS = 10_000;
 
-/**
- * Mark a document as dirty. After DEBOUNCE_MS of inactivity,
- * push the latest state to HF.
- */
 export function schedulePush(docName: string, state: Buffer): void {
   const existing = dirtyDocs.get(docName);
   if (existing) clearTimeout(existing.timer);
@@ -53,6 +99,8 @@ async function pushDocument(docName: string): Promise<void> {
   if (!entry) return;
   dirtyDocs.delete(docName);
 
+  await ensureDatasetExists();
+
   const safeName = sanitizeName(docName);
   const path = `articles/${safeName}.yjs`;
 
@@ -60,7 +108,7 @@ async function pushDocument(docName: string): Promise<void> {
     await uploadFile({
       repo,
       file: { path, content: new Blob([new Uint8Array(entry.state)]) },
-      accessToken: HF_TOKEN,
+      accessToken: getToken(),
       commitTitle: `save ${safeName}`,
     });
     console.log(`[hf-storage] pushed ${path}`);
@@ -77,7 +125,7 @@ export async function pullDocument(docName: string): Promise<Uint8Array | null>
     const res = await downloadFile({
       repo,
       path,
-      accessToken: HF_TOKEN,
+      accessToken: getToken(),
     });
     if (!res) return null;
     const arrayBuf = await res.arrayBuffer();
@@ -87,10 +135,6 @@ export async function pullDocument(docName: string): Promise<Uint8Array | null>
   }
 }
 
-/**
- * Flush all pending documents to HF immediately.
- * Call this on graceful shutdown.
- */
 export async function flushAll(): Promise<void> {
   const names = [...dirtyDocs.keys()];
   await Promise.allSettled(names.map((n) => pushDocument(n)));
@@ -107,16 +151,19 @@ interface PublishedPayload {
 
 export async function uploadPublishedAssets(
   docName: string,
-  payload: PublishedPayload
+  payload: PublishedPayload,
+  token?: string,
 ): Promise<{ htmlUrl: string; pdfUrl: string | null; thumbUrl: string | null }> {
+  await ensureDatasetExists(token);
+
+  const accessToken = getToken(token);
   const safeName = sanitizeName(docName);
   const base = `published/${safeName}`;
 
-  // Upload HTML
   await uploadFile({
     repo,
     file: { path: `${base}/index.html`, content: new Blob([payload.html]) },
-    accessToken: HF_TOKEN,
+    accessToken,
     commitTitle: `publish ${safeName}: HTML`,
   });
 
@@ -124,36 +171,33 @@ export async function uploadPublishedAssets(
   let pdfUrl: string | null = null;
   let thumbUrl: string | null = null;
 
-  // Upload PDF
   if (payload.pdf) {
     await uploadFile({
       repo,
       file: { path: `${base}/article.pdf`, content: new Blob([new Uint8Array(payload.pdf)]) },
-      accessToken: HF_TOKEN,
+      accessToken,
       commitTitle: `publish ${safeName}: PDF`,
     });
     pdfUrl = `${baseUrl}/article.pdf`;
   }
 
-  // Upload thumbnail
   if (payload.thumbnail) {
     await uploadFile({
       repo,
       file: { path: `${base}/thumb.jpg`, content: new Blob([new Uint8Array(payload.thumbnail)]) },
-      accessToken: HF_TOKEN,
+      accessToken,
       commitTitle: `publish ${safeName}: thumbnail`,
     });
     thumbUrl = `${baseUrl}/thumb.jpg`;
   }
 
-  // Upload meta.json
   await uploadFile({
     repo,
     file: {
       path: `${base}/meta.json`,
       content: new Blob([JSON.stringify(payload.meta, null, 2)]),
     },
-    accessToken: HF_TOKEN,
+    accessToken,
     commitTitle: `publish ${safeName}: metadata`,
   });
 

backend/src/publisher/index.ts

@@ -91,7 +91,7 @@ export interface PublishResult {
 /**
  * Run the full publish pipeline for a document.
  */
-export async function publishDocument(docName: string): Promise<PublishResult> {
+export async function publishDocument(docName: string, token?: string): Promise<PublishResult> {
   const p = docPath(docName);
 
   if (!existsSync(p)) {
@@ -144,7 +144,7 @@ export async function publishDocument(docName: string): Promise<PublishResult> {
         pdf,
         thumbnail,
         meta: meta as any,
-      });
+      }, token);
       return { ...urls, success: true };
     } catch (err) {
       console.error("[publish] Upload failed:", err);

backend/src/server.ts

@@ -13,6 +13,9 @@ import { handleChat } from "./agent/chat.js";
 import { citationsRouter } from "./citations.js";
 import {
   isHfStorageEnabled,
+  ensureDatasetExists,
+  getDatasetId,
+  setUserToken,
   uploadImageToHf,
   pullDocument,
   schedulePush,
@@ -34,11 +37,12 @@ function docPath(name: string) {
   return join(DATA_DIR, `${name.replace(/[^a-zA-Z0-9_-]/g, "_")}.yjs`);
 }
 
-const hfEnabled = isHfStorageEnabled();
-if (hfEnabled) {
+if (isHfStorageEnabled()) {
   console.log("[server] HF Dataset persistence enabled");
+} else if (getDatasetId()) {
+  console.log("[server] HF Dataset ID ready, waiting for first user login to activate persistence");
 } else {
-  console.log("[server] HF Dataset persistence disabled (no HF_TOKEN / HF_DATASET_ID)");
+  console.log("[server] HF Dataset persistence disabled (no SPACE_ID or HF_DATASET_ID)");
   console.log("[server] falling back to local file persistence");
 }
 
@@ -62,12 +66,10 @@ const hocuspocus = new Hocuspocus({
   extensions: [
     new Database({
       fetch: async ({ documentName }: { documentName: string }) => {
-        // Try local cache first
         const p = docPath(documentName);
         if (existsSync(p)) return readFileSync(p);
 
-        // Then try HF dataset
-        if (hfEnabled) {
+        if (isHfStorageEnabled()) {
           const data = await pullDocument(documentName);
           if (data) {
             writeFileSync(p, data);
@@ -85,11 +87,9 @@ const hocuspocus = new Hocuspocus({
         documentName: string;
         state: Buffer;
       }) => {
-        // Always write locally for fast access
         writeFileSync(docPath(documentName), state);
 
-        // Schedule debounced push to HF
-        if (hfEnabled) {
+        if (isHfStorageEnabled()) {
           schedulePush(documentName, state);
         }
       },
@@ -124,6 +124,11 @@ app.get("/api/auth/status", async (req, res) => {
     return;
   }
 
+  // Cache the user's OAuth token for HF API calls (dataset storage)
+  if (user.canEdit && token) {
+    setUserToken(token);
+  }
+
   res.json({
     authenticated: true,
     canEdit: user.canEdit,
@@ -152,7 +157,8 @@ app.use("/api/citations", citationsRouter);
 // ---------- Publish ----------
 
 app.post("/api/publish", async (req, res) => {
-  // Auth check: only editors can publish
+  let userToken: string | undefined;
+
   if (oauthEnabled) {
     const token = extractToken(req.headers.cookie);
     const user = await resolveUser(token);
@@ -166,6 +172,8 @@ app.post("/api/publish", async (req, res) => {
       res.status(403).json({ error: "Unauthorized: write access required" });
       return;
     }
+    userToken = token ?? undefined;
+    if (userToken) setUserToken(userToken);
   }
 
   try {
@@ -179,7 +187,7 @@ app.post("/api/publish", async (req, res) => {
     }
     await conn.disconnect();
 
-    const result = await publishDocument(DEFAULT_DOC_NAME);
+    const result = await publishDocument(DEFAULT_DOC_NAME, userToken);
     if (!result.success) {
       res.status(500).json({ error: result.error || "Publish failed" });
       return;
@@ -206,8 +214,11 @@ app.post("/api/upload", upload.single("file"), async (req, res) => {
     const ext = file.originalname.split(".").pop() || "png";
     const filename = `${randomUUID()}.${ext}`;
 
-    if (hfEnabled) {
-      const url = await uploadImageToHf(file.buffer, filename);
+    const userToken = extractToken(req.headers.cookie) ?? undefined;
+    if (userToken) setUserToken(userToken);
+
+    if (isHfStorageEnabled()) {
+      const url = await uploadImageToHf(file.buffer, filename, userToken);
       res.json({ url });
       return;
     }
@@ -238,18 +249,18 @@ app.get("/api/images/*", async (req, res) => {
     return;
   }
 
-  const HF_DATASET_ID = process.env.HF_DATASET_ID || "";
-  const HF_TOKEN = process.env.HF_TOKEN || "";
-
-  if (!HF_DATASET_ID || !HF_TOKEN) {
+  const datasetId = getDatasetId();
+  if (!datasetId || !isHfStorageEnabled()) {
     res.status(404).send("Not configured");
     return;
   }
 
+  const userToken = extractToken(req.headers.cookie) ?? undefined;
+
   try {
-    const url = `https://huggingface.co/datasets/${HF_DATASET_ID}/resolve/main/images/${imagePath}`;
+    const url = `https://huggingface.co/datasets/${datasetId}/resolve/main/images/${imagePath}`;
     const response = await fetch(url, {
-      headers: { Authorization: `Bearer ${HF_TOKEN}` },
+      headers: { Authorization: `Bearer ${userToken || ""}` },
     });
 
     if (!response.ok) {