"""HF Dataset persistent storage + OAuth for the loom critique service. Persists the SQLite DB to a private HF Dataset so data survives Space rebuilds. Also provides HF OAuth helpers so users can authorize dataset access. Tier 3 extension: also persists the ``.brain/`` folder for each active Conscious workspace as ``brains/.tar.gz`` (one tarball per workspace, uploaded every 120s alongside ``loom.db``; re-downloaded on boot). See TIER3_PLAN.md §7.2 for the spec. Env vars required: HF_CLIENT_ID — HF OAuth App client ID HF_CLIENT_SECRET — HF OAuth App client secret ENCRYPTION_KEY — Fernet key for token encryption (auto-generated if absent) Env vars optional: LOOM_DATASET_SUFFIX — suffix for the dataset name (default: main) The full name is {hf_username}/loom-priv-{suffix} """ from __future__ import annotations import io import json import os import tarfile import threading import time import urllib.error import urllib.request from datetime import datetime, timezone, timedelta from pathlib import Path from urllib.parse import urlencode import crypto import db import jwt_auth # --------------------------------------------------------------------------- # Config # --------------------------------------------------------------------------- HF_AUTHORIZE_URL = "https://huggingface.co/oauth/authorize" HF_TOKEN_URL = "https://huggingface.co/oauth/token" HF_API_BASE = "https://huggingface.co" HF_SCOPES = "openid profile write-repos" # --------------------------------------------------------------------------- # HF API helpers # --------------------------------------------------------------------------- def _hf_api(path: str, *, token: str = "", method: str = "GET", body: dict | None = None) -> dict: """Raw HF Hub API call via stdlib urllib.""" url = f"{HF_API_BASE}{path}" data = json.dumps(body).encode() if body is not None else None req = urllib.request.Request(url, data=data, method=method) req.add_header("Content-Type", "application/json") req.add_header("Accept", "application/json") if token: req.add_header("Authorization", f"Bearer {token}") try: with urllib.request.urlopen(req, timeout=30) as resp: raw = resp.read().decode() return json.loads(raw) if raw.strip() else {} except urllib.error.HTTPError as exc: detail = exc.read().decode(errors="replace")[:400] raise RuntimeError(f"HF API {path} failed (HTTP {exc.code}): {detail}") from exc def _hf_configured() -> bool: return bool( os.environ.get("HF_CLIENT_ID", "").strip() and os.environ.get("HF_CLIENT_SECRET", "").strip() ) # --------------------------------------------------------------------------- # Space name detection # --------------------------------------------------------------------------- def _space_name() -> str: """Return the Space repo name for dataset naming.""" name = os.environ.get("SPACE_NAME", "").strip() if name: return name sid = os.environ.get("SPACE_ID", "").strip() if sid and "/" in sid: return sid.split("/")[-1] return "main" # --------------------------------------------------------------------------- # HF OAuth flow # --------------------------------------------------------------------------- def make_hf_authorize_url(state: str, redirect_uri: str = "") -> str: """Build the HF OAuth authorize URL with state token. ``redirect_uri`` must match the registered callback URL exactly. """ params = urlencode({ "client_id": os.environ.get("HF_CLIENT_ID", ""), "redirect_uri": redirect_uri, "scope": HF_SCOPES, "response_type": "code", "state": state, }) return f"{HF_AUTHORIZE_URL}?{params}" def exchange_hf_code(code: str, redirect_uri: str = "") -> dict: """Exchange OAuth code for tokens. Returns dict with access_token, refresh_token, expires_in. ``redirect_uri`` must match the value used in the authorize request. Raises RuntimeError on failure. """ body = urlencode({ "client_id": os.environ.get("HF_CLIENT_ID", ""), "client_secret": os.environ.get("HF_CLIENT_SECRET", ""), "code": code, "grant_type": "authorization_code", "redirect_uri": redirect_uri, }).encode() req = urllib.request.Request(HF_TOKEN_URL, data=body, method="POST") req.add_header("Accept", "application/json") req.add_header("Content-Type", "application/x-www-form-urlencoded") try: with urllib.request.urlopen(req, timeout=30) as resp: data = json.loads(resp.read().decode()) except urllib.error.HTTPError as exc: detail = exc.read().decode(errors="replace")[:200] # don't include the full response — it may contain tokens (audit H4) raise RuntimeError(f"HF token exchange failed (HTTP {exc.code})") from exc token = data.get("access_token", "").strip() if not token: # log only the keys present, not the values (which may contain tokens) raise RuntimeError(f"HF token exchange returned no access_token (keys: {list(data.keys())})") return data def get_hf_user(token: str) -> dict: """Fetch authenticated user info from HF API.""" return _hf_api("/api/whoami-v2", token=token) def _valid_user_id(user_id: str | None) -> bool: """Return True if user_id looks like a valid 16-char hex ID. Rejects JWT strings (contain dots) and other garbage that might be accidentally passed by the frontend. """ if not user_id or not isinstance(user_id, str): return False if len(user_id) != 16: return False if "." in user_id: return False try: int(user_id, 16) return True except ValueError: return False def upsert_user_from_hf(token_data: dict, user_id: str | None = None) -> dict: """Fetch HF user info, upsert into DB, return the user row. ``token_data`` must contain ``access_token``, and may contain ``refresh_token`` and ``expires_in`` from the token exchange response. If ``user_id`` is provided, look up existing user by that internal UUID first. If found, use their github_id for the upsert to ensure merge works. If ``user_id`` is None, derives a deterministic ID from the HF username. """ hf_token = token_data["access_token"] info = get_hf_user(hf_token) encrypted = crypto.encrypt_token(hf_token) refresh_encrypted = crypto.encrypt_token(token_data["refresh_token"]) if token_data.get("refresh_token") else None expires_at = "" if token_data.get("expires_in"): expires_at = (datetime.now(timezone.utc) + timedelta(seconds=int(token_data["expires_in"]))).isoformat() # Derive a deterministic ID if none provided if user_id is None or not _valid_user_id(user_id): user_id = jwt_auth.derive_user_id_from_hf(info["name"]) github_id = None if user_id: existing = db.get_user(user_id) if existing: github_id = existing.get("github_id") return db.upsert_user( user_id=user_id, github_id=github_id, hf_id=info["name"], hf_username=info["name"], hf_token_encrypted=encrypted, hf_refresh_token_encrypted=refresh_encrypted, hf_token_expires_at=expires_at, ) def _hf_token_for_user(user_id: str) -> str: """Retrieve and decrypt a user's stored HF token. Auto-refreshes the token if it expires within 7 days and a refresh token is available. Raises RuntimeError if no token is stored. """ user = db.get_user(user_id) if not user or not user.get("hf_token_encrypted"): raise RuntimeError("HF not connected — please link your Hugging Face account") expires_at = user.get("hf_token_expires_at") or "" if _hf_token_expires_soon(expires_at) and user.get("hf_refresh_token_encrypted"): try: _refresh_hf_token(user_id, user) except Exception as exc: print(f"[dataset_persistence] token refresh failed: {type(exc).__name__}", flush=True) user = db.get_user(user_id) return crypto.decrypt_token(user["hf_token_encrypted"]) def _hf_token_expires_soon(expires_at: str, within_days: int = 7) -> bool: """Return True if the token expires within ``within_days`` or is already expired.""" if not expires_at: return True try: expiry = datetime.fromisoformat(expires_at) return expiry - timedelta(days=within_days) < datetime.now(timezone.utc) except (ValueError, TypeError): return True def _refresh_hf_token(user_id: str, user: dict) -> None: """Refresh a user's HF token using the stored refresh token. Updates the DB with new access/refresh tokens and expiry. """ refresh_encrypted = user.get("hf_refresh_token_encrypted") if not refresh_encrypted: raise RuntimeError("no refresh token stored") old_refresh = crypto.decrypt_token(refresh_encrypted) body = urlencode({ "client_id": os.environ.get("HF_CLIENT_ID", ""), "client_secret": os.environ.get("HF_CLIENT_SECRET", ""), "grant_type": "refresh_token", "refresh_token": old_refresh, }).encode() req = urllib.request.Request(HF_TOKEN_URL, data=body, method="POST") req.add_header("Accept", "application/json") req.add_header("Content-Type", "application/x-www-form-urlencoded") try: with urllib.request.urlopen(req, timeout=30) as resp: data = json.loads(resp.read().decode()) except urllib.error.HTTPError as exc: detail = exc.read().decode(errors="replace")[:400] raise RuntimeError(f"HF token refresh failed (HTTP {exc.code}): {detail}") from exc new_token = data.get("access_token", "").strip() if not new_token: raise RuntimeError(f"HF token refresh returned no access_token: {data}") encrypted = crypto.encrypt_token(new_token) new_refresh_encrypted = None if data.get("refresh_token"): new_refresh_encrypted = crypto.encrypt_token(data["refresh_token"]) expires_at = "" if data.get("expires_in"): expires_at = (datetime.now(timezone.utc) + timedelta(seconds=int(data["expires_in"]))).isoformat() db.upsert_user( hf_id=user["hf_id"], hf_token_encrypted=encrypted, hf_refresh_token_encrypted=new_refresh_encrypted, hf_token_expires_at=expires_at, ) # --------------------------------------------------------------------------- # Dataset persistence # --------------------------------------------------------------------------- _upload_scheduler: threading.Thread | None = None _shutdown_event = threading.Event() _current_hf_token: str = "" def dataset_name(hf_username: str) -> str: """Return the full dataset repo name for a user. Format: {hf_username}/loom-priv-{suffix} """ suffix = os.environ.get("LOOM_DATASET_SUFFIX", "").strip() or _space_name() return f"{hf_username}/loom-priv-{suffix}" def init_persistence(user_id: str) -> None: """Initialize dataset persistence for the current user. Ensures the private dataset repo exists, downloads the latest DB, and starts the periodic upload scheduler. Tier 3 extension: after the DB download, also restores every active Conscious's ``.brain/`` folder (§7.4). If the workspace sandbox doesn't exist (Space was rebuilt), it's recreated from the workspace's ``source_repo`` via the existing token-safe ``_clone_repo`` (Tier 1). """ global _current_hf_token token = _hf_token_for_user(user_id) _current_hf_token = token user = db.get_user(user_id) hf_user = user.get("hf_username", "") if user else "" if not hf_user: info = get_hf_user(token) hf_user = info.get("name", "") ds_name = dataset_name(hf_user) try: _hf_api(f"/api/repos/{ds_name}", token=token) except RuntimeError as exc: if "404" in str(exc): _hf_api("/api/repos", method="POST", token=token, body={ "name": ds_name, "private": True, "type": "dataset", }) else: raise _download_db(token, ds_name) # Tier 3 — restore .brain/ folders for all active Conscious workspaces. _restore_all_brains(token, ds_name) _start_upload_scheduler(token, ds_name) def _download_db(token: str, ds_name: str) -> None: """Download loom.db from the dataset repo to /data/loom.db if it exists.""" from huggingface_hub import hf_hub_download from huggingface_hub.utils import RepositoryNotFoundError, RevisionNotFoundError try: path = hf_hub_download( repo_id=ds_name, filename="loom.db", token=token, repo_type="dataset", ) import shutil shutil.copy2(path, db.DB_PATH) except (RepositoryNotFoundError, RevisionNotFoundError, OSError): pass def upload_db(token: str, ds_name: str) -> None: """Upload the current loom.db to the dataset repo.""" if not db.DB_PATH.exists(): return from huggingface_hub import upload_file from huggingface_hub.utils import HfHubHTTPError try: upload_file( path_or_fileobj=str(db.DB_PATH), path_in_repo="loom.db", repo_id=ds_name, repo_type="dataset", token=token, commit_message="auto-sync loom.db", ) except HfHubHTTPError as exc: print(f"[dataset_persistence] upload failed: {type(exc).__name__}", flush=True) def _start_upload_scheduler(token: str, ds_name: str, interval: int = 120) -> None: """Start a daemon thread that uploads the DB every ``interval`` seconds. Tier 3 extension: also uploads every active Conscious's ``.brain/`` folder (one tarball per workspace). Single-threaded; no new locks. Failures in one workspace don't block others (each wrapped in try/except, logged). """ global _upload_scheduler if _upload_scheduler is not None: return def loop(): while not _shutdown_event.is_set(): if _shutdown_event.wait(interval): break try: upload_db(token, ds_name) except Exception: pass # Tier 3 — also sync brains try: _upload_all_brains(token, ds_name) except Exception as exc: print(f"[dataset_persistence] brain sync failed: {type(exc).__name__}", flush=True) _upload_scheduler = threading.Thread(target=loop, daemon=True) _upload_scheduler.start() # --------------------------------------------------------------------------- # Tier 3 — .brain/ folder sync (§7.2, §7.4) # --------------------------------------------------------------------------- BRAIN_SYNC_INTERVAL_S = 120 # same as DB; shares the scheduler def _tar_brain(brain_path: Path) -> bytes: """Pack a .brain/ folder into a tar.gz bytes blob. Skips if path missing.""" buf = io.BytesIO() with tarfile.open(fileobj=buf, mode="w:gz") as tar: tar.add(str(brain_path), arcname=".brain") return buf.getvalue() def _untar_brain(blob: bytes, dest_workspace: Path) -> None: """Extract a tar.gz brain blob into ``dest_workspace`` (.brain/ lands there). Phase 5 security: uses ``filter='data'`` (Python 3.12+) to prevent path-traversal via malicious tarball entries (audit L4). On older Pythons where the filter param doesn't exist, falls back to manual member validation. """ buf = io.BytesIO(blob) with tarfile.open(fileobj=buf, mode="r:gz") as tar: try: # Python 3.12+ — built-in data filter blocks path traversal tar.extractall(str(dest_workspace), filter="data") except TypeError: # older Python — manual validation: reject members with .. or abs paths dest = dest_workspace.resolve() for member in tar.getmembers(): member_path = (dest_workspace / member.name).resolve() if not str(member_path).startswith(str(dest)): raise RuntimeError(f"tar member escapes dest: {member.name}") tar.extractall(str(dest_workspace)) def upload_brain(token: str, ds_name: str, workspace_id: str, brain_path: Path) -> None: """Upload the .brain/ folder for one workspace to the HF Dataset. Packs .brain/ as a tar.gz, uploads as ``brains/.tar.gz``. Idempotent: overwrites the previous upload. Skips if brain_path missing. Per TIER3_PLAN §7.2: one tarball per workspace per 120s tick (not per-file) to keep HF Dataset upload calls bounded. Phase 3: after a successful upload, stamps ``conscious.last_synced_at`` for every active conscious in this workspace so the API can surface stale syncs (>10min) — TIER3_PLAN §17 risk mitigation. Phase 5 security: scans .brain/ for known secret patterns BEFORE upload. If secrets are found, the upload is REFUSED (TIER3_PLAN §17: "refuse to commit if any file in .brain/ matches a secret pattern"). This is the second defense layer — commit_brain scans first, this scans before HF. """ if not brain_path.is_dir(): return # Security: scan for secrets before uploading to HF (audit C1) try: import brain as _brain_mod findings = _brain_mod.scan_for_secrets(brain_path) if findings: print(f"[dataset_persistence] BRAIN UPLOAD BLOCKED for {workspace_id}: " f"{len(findings)} secret(s) detected in .brain/", flush=True) return # refuse to upload except Exception: pass # scan failure must never block — but log it try: blob = _tar_brain(brain_path) except Exception as exc: print(f"[dataset_persistence] tar brain failed for {workspace_id}: " f"{type(exc).__name__}", # don't leak details (audit H4) flush=True) return from huggingface_hub import upload_file from huggingface_hub.utils import HfHubHTTPError try: upload_file( path_or_fileobj=blob, path_in_repo=f"brains/{workspace_id}.tar.gz", repo_id=ds_name, repo_type="dataset", token=token, commit_message=f"auto-sync brain {workspace_id}", ) except HfHubHTTPError as exc: print(f"[dataset_persistence] brain upload failed for {workspace_id}: {type(exc).__name__}", flush=True) return # Phase 3 — stamp last_synced_at on every active conscious in this workspace try: import conscious_db now = conscious_db._iso_now() with _db_write_lock_for_sync(): _db_for_sync().execute( "UPDATE conscious SET last_synced_at = ?, updated_at = ? " "WHERE workspace_id = ? AND status = 'active'", (now, now, workspace_id)) _db_for_sync().commit() except Exception as exc: print(f"[dataset_persistence] stamp last_synced_at failed for {workspace_id}: {type(exc).__name__}", flush=True) # Phase 3 — local helpers to avoid circular import of db._write_lock at module # load time. conscious_db imports db, and dataset_persistence imports # conscious_db lazily inside functions; these helpers mirror that pattern. def _db_for_sync(): import db as _db return _db._db() def _db_write_lock_for_sync(): import db as _db return _db._write_lock def download_brain(token: str, ds_name: str, workspace_id: str, dest_workspace: Path) -> None: """Download and extract ``brains/.tar.gz`` into dest_workspace. No-op if the file doesn't exist in the dataset (first run for that workspace). """ from huggingface_hub import hf_hub_download from huggingface_hub.utils import RepositoryNotFoundError, RevisionNotFoundError, EntryNotFoundError try: path = hf_hub_download( repo_id=ds_name, filename=f"brains/{workspace_id}.tar.gz", token=token, repo_type="dataset", ) except (RepositoryNotFoundError, RevisionNotFoundError, EntryNotFoundError, OSError): return # first run — no brain yet try: with open(path, "rb") as f: blob = f.read() dest_workspace.mkdir(parents=True, exist_ok=True) _untar_brain(blob, dest_workspace) except Exception as exc: print(f"[dataset_persistence] brain extract failed for {workspace_id}: {type(exc).__name__}", flush=True) def _upload_all_brains(token: str, ds_name: str) -> None: """Walk all active Conscious workspaces and upload each .brain/ folder.""" try: import conscious_db items = conscious_db.list_active_conscious_workspaces() except Exception: return # conscious tables not initialized yet for item in items: ws_id = item["workspace_id"] sandbox = item.get("sandbox_path") or "" if not sandbox: continue brain_path = Path(sandbox) / ".brain" try: upload_brain(token, ds_name, ws_id, brain_path) except Exception as exc: print(f"[dataset_persistence] brain upload error for {ws_id}: {type(exc).__name__}", flush=True) continue def _restore_all_brains(token: str, ds_name: str) -> None: """On boot, restore .brain/ for every Conscious in the restored DB. For each Conscious row: if the workspace sandbox doesn't exist, recreate it from the workspace's source_repo via Tier 1's token-safe ``_clone_repo`` (in github_integration); then extract the brain tarball into it. """ try: import conscious_db items = conscious_db.list_active_conscious_workspaces() except Exception: return # conscious tables not initialized for item in items: ws_id = item["workspace_id"] sandbox = item.get("sandbox_path") or "" if not sandbox: continue sandbox_path = Path(sandbox) try: if not sandbox_path.is_dir(): # recreate the workspace from source_repo (Tier 1 token-safe clone) ws = db.get_workspace(ws_id) if ws and ws.get("source_repo") and ws.get("user_id"): try: import github_integration github_integration._clone_repo( ws["user_id"], ws["source_repo"], branch=ws.get("source_branch"), dest=str(sandbox_path), branches=None) except Exception as exc: print(f"[dataset_persistence] workspace re-clone failed " f"for {ws_id}: {type(exc).__name__}", flush=True) continue else: continue download_brain(token, ds_name, ws_id, sandbox_path) except Exception as exc: print(f"[dataset_persistence] brain restore error for {ws_id}: {type(exc).__name__}", flush=True) continue def shutdown_upload() -> None: """Signal the upload scheduler to stop and do one final upload. Tier 3: also does a final brain sync so the latest brain state survives a Space rebuild. """ global _upload_scheduler _shutdown_event.set() if _upload_scheduler is not None: _upload_scheduler.join(timeout=10) _upload_scheduler = None if _current_hf_token: try: info = get_hf_user(_current_hf_token) ds_name = dataset_name(info.get("name", "")) upload_db(_current_hf_token, ds_name) _upload_all_brains(_current_hf_token, ds_name) except Exception: pass