P0-R2 HF Bucket root-path correction

hf_bucket_store.py

@@ -0,0 +1,380 @@
+# -*- coding: utf-8 -*-
+"""Best-effort private Hugging Face Bucket storage for completed ZIP delivery.
+
+This module is deliberately outside the paper-generation core. Bucket storage is
+an optional recovery layer only: a Bucket error must never turn an already valid
+local paper package into a failed generation or alter quota/accounting state.
+"""
+from __future__ import annotations
+
+import hashlib
+import json
+import logging
+import os
+import re
+import tempfile
+import zipfile
+from dataclasses import dataclass
+from datetime import datetime, timedelta, timezone
+from pathlib import Path
+from typing import Any
+
+_TRUE = {"1", "true", "yes", "on", "enabled"}
+_DEFAULT_TTL_HOURS = 168
+_DEFAULT_WARNING_GIB = 70.0
+_DEFAULT_HARD_STOP_GIB = 80.0
+_MAX_HARD_STOP_GIB = 80.0  # Free-only safety ceiling approved for Article Agent.
+_GIB = 1024 ** 3
+_TASK_RE = re.compile(r"^[A-Za-z0-9_.-]{1,160}$")
+_BUCKET_RE = re.compile(r"^[A-Za-z0-9_.-]+/[A-Za-z0-9_.-]+$")
+
+
+@dataclass(frozen=True)
+class BucketConfig:
+    enabled: bool
+    bucket_id: str
+    token: str
+    ttl_hours: int
+    warning_bytes: int
+    hard_stop_bytes: int
+
+
+def _truthy(value: Any) -> bool:
+    return str(value or "").strip().casefold() in _TRUE
+
+
+def _bounded_float(name: str, default: float, low: float, high: float) -> float:
+    try:
+        value = float(os.environ.get(name, str(default)))
+    except (TypeError, ValueError):
+        value = default
+    return max(low, min(high, value))
+
+
+def _bounded_int(name: str, default: int, low: int, high: int) -> int:
+    try:
+        value = int(os.environ.get(name, str(default)))
+    except (TypeError, ValueError):
+        value = default
+    return max(low, min(high, value))
+
+
+def config() -> BucketConfig:
+    enabled = _truthy(os.environ.get("SMARTPAPER_HF_ARTIFACTS_ENABLED"))
+    bucket_id = str(os.environ.get("SMARTPAPER_HF_ARTIFACT_BUCKET") or "").strip()
+    token = str(os.environ.get("SMARTPAPER_HF_BUCKET_TOKEN") or "").strip()
+    ttl_hours = _bounded_int("SMARTPAPER_HF_ARTIFACT_TTL_HOURS", _DEFAULT_TTL_HOURS, 1, _DEFAULT_TTL_HOURS)
+    requested_hard = _bounded_float(
+        "SMARTPAPER_HF_ARTIFACT_HARD_STOP_GIB", _DEFAULT_HARD_STOP_GIB, 1.0, _MAX_HARD_STOP_GIB
+    )
+    requested_warning = _bounded_float(
+        "SMARTPAPER_HF_ARTIFACT_WARNING_GIB", _DEFAULT_WARNING_GIB, 1.0, requested_hard
+    )
+    return BucketConfig(
+        enabled=enabled,
+        bucket_id=bucket_id,
+        token=token,
+        ttl_hours=ttl_hours,
+        warning_bytes=int(requested_warning * _GIB),
+        hard_stop_bytes=int(requested_hard * _GIB),
+    )
+
+
+def _new_api(token: str):
+    # huggingface_hub reads these at import time. Keep recovery bounded without
+    # importing the optional dependency when the feature is disabled.
+    os.environ.setdefault("HF_HUB_ETAG_TIMEOUT", "8")
+    os.environ.setdefault("HF_HUB_DOWNLOAD_TIMEOUT", "8")
+    from huggingface_hub import HfApi  # lazy: absence must not break normal generation
+
+    return HfApi(token=token)
+
+
+def _validate_config(cfg: BucketConfig) -> None:
+    if not cfg.enabled:
+        raise RuntimeError("hf_bucket_disabled")
+    if not cfg.bucket_id or not _BUCKET_RE.fullmatch(cfg.bucket_id):
+        raise RuntimeError("hf_bucket_id_invalid")
+    if not cfg.token:
+        raise RuntimeError("hf_bucket_token_missing")
+    if cfg.hard_stop_bytes > int(_MAX_HARD_STOP_GIB * _GIB):
+        raise RuntimeError("hf_bucket_hard_stop_exceeds_free_safety_ceiling")
+
+
+def _trusted_key_parts(user_id: int | str, task_id: str) -> tuple[str, str]:
+    try:
+        trusted_user_id = str(int(user_id))
+    except (TypeError, ValueError) as exc:
+        raise RuntimeError("hf_bucket_user_id_invalid") from exc
+    trusted_task_id = str(task_id or "").strip()
+    if not _TASK_RE.fullmatch(trusted_task_id):
+        raise RuntimeError("hf_bucket_task_id_invalid")
+    return trusted_user_id, trusted_task_id
+
+
+def object_paths(user_id: int | str, task_id: str) -> tuple[str, str]:
+    trusted_user_id, trusted_task_id = _trusted_key_parts(user_id, task_id)
+    prefix = f"v1/users/{trusted_user_id}/tasks/{trusted_task_id}"
+    return f"{prefix}/final.zip", f"{prefix}/final.json"
+
+
+def _safe_output_filename(value: str) -> str:
+    text = str(value or "").strip()
+    name = os.path.basename(text)
+    if not name or name != text or not name.lower().endswith(".zip") or len(name) > 240:
+        raise RuntimeError("hf_bucket_output_filename_invalid")
+    return name
+
+
+def _sha256_file(path: Path) -> tuple[str, int]:
+    digest = hashlib.sha256()
+    size = 0
+    with path.open("rb") as handle:
+        for block in iter(lambda: handle.read(1024 * 1024), b""):
+            size += len(block)
+            digest.update(block)
+    return digest.hexdigest(), size
+
+
+def _private_bucket_info(api, cfg: BucketConfig):
+    info = api.bucket_info(cfg.bucket_id)
+    if getattr(info, "private", None) is not True:
+        raise RuntimeError("hf_bucket_not_private")
+    return info
+
+
+def _file_map(api, bucket_id: str, paths: list[str]) -> dict[str, Any]:
+    return {str(item.path): item for item in api.get_bucket_paths_info(bucket_id, paths)}
+
+
+def _metadata_bytes(*, user_id: str, task_id: str, filename: str, sha256: str, size_bytes: int, uploaded_at: datetime) -> bytes:
+    payload = {
+        "schema_version": "articleagent-hf-bucket-v1",
+        "user_id": user_id,
+        "task_id": task_id,
+        "output_filename": filename,
+        "sha256": sha256,
+        "size_bytes": int(size_bytes),
+        "uploaded_at": uploaded_at.astimezone(timezone.utc).isoformat(),
+    }
+    return json.dumps(payload, sort_keys=True, separators=(",", ":"), ensure_ascii=False).encode("utf-8")
+
+
+def _parse_uploaded_at(value: Any) -> datetime:
+    text = str(value or "").strip()
+    if not text:
+        raise RuntimeError("hf_bucket_metadata_time_missing")
+    try:
+        observed = datetime.fromisoformat(text.replace("Z", "+00:00"))
+    except ValueError as exc:
+        raise RuntimeError("hf_bucket_metadata_time_invalid") from exc
+    if observed.tzinfo is None:
+        observed = observed.replace(tzinfo=timezone.utc)
+    return observed.astimezone(timezone.utc)
+
+
+def _download_metadata(api, bucket_id: str, remote_path: str) -> dict[str, Any]:
+    with tempfile.TemporaryDirectory(prefix="sp_hf_meta_") as temp_dir:
+        target = Path(temp_dir) / "final.json"
+        api.download_bucket_files(bucket_id, [(remote_path, target)], raise_on_missing_files=True)
+        raw = target.read_bytes()
+    if len(raw) > 64 * 1024:
+        raise RuntimeError("hf_bucket_metadata_too_large")
+    value = json.loads(raw.decode("utf-8"))
+    if not isinstance(value, dict):
+        raise RuntimeError("hf_bucket_metadata_invalid")
+    return value
+
+
+def cleanup_expired(*, api, cfg: BucketConfig, now: datetime | None = None) -> dict[str, Any]:
+    """Delete only files older than the approved 7-day retention window.
+
+    Cleanup itself is best-effort at the call site. This function raises on API
+    failure so callers can log it, while never guessing that a younger object is
+    safe to delete.
+    """
+    current = (now or datetime.now(timezone.utc)).astimezone(timezone.utc)
+    cutoff = current - timedelta(hours=cfg.ttl_hours)
+    expired: list[str] = []
+    for item in api.list_bucket_tree(cfg.bucket_id, prefix="v1/users/", recursive=True):
+        if str(getattr(item, "type", "")) != "file":
+            continue
+        mtime = getattr(item, "mtime", None)
+        if not isinstance(mtime, datetime):
+            continue
+        if mtime.tzinfo is None:
+            mtime = mtime.replace(tzinfo=timezone.utc)
+        if mtime.astimezone(timezone.utc) < cutoff:
+            expired.append(str(item.path))
+    if expired:
+        api.batch_bucket_files(cfg.bucket_id, delete=expired)
+    return {"deleted": len(expired), "paths": expired}
+
+
+def archive_completed_zip(
+    *,
+    user_id: int | str,
+    task_id: str,
+    output_filename: str,
+    output_path: str | Path,
+    now: datetime | None = None,
+) -> dict[str, Any]:
+    """Best-effort archive primitive. Caller must catch all errors.
+
+    This function never modifies task/quota state and should only be called after
+    the authoritative Worker has accepted the completed task transition.
+    """
+    cfg = config()
+    if not cfg.enabled:
+        return {"ok": True, "stored": False, "reason": "disabled"}
+    _validate_config(cfg)
+    trusted_user_id, trusted_task_id = _trusted_key_parts(user_id, task_id)
+    filename = _safe_output_filename(output_filename)
+    local = Path(output_path).resolve()
+    if not local.is_file() or local.name != filename:
+        raise RuntimeError("hf_bucket_local_zip_missing")
+    sha256, size_bytes = _sha256_file(local)
+    if size_bytes <= 0:
+        raise RuntimeError("hf_bucket_local_zip_empty")
+
+    api = _new_api(cfg.token)
+    _private_bucket_info(api, cfg)
+
+    cleanup_warning = ""
+    try:
+        cleanup_expired(api=api, cfg=cfg, now=now)
+    except Exception as exc:  # cleanup must never destroy local delivery
+        cleanup_warning = type(exc).__name__
+        logging.warning("[SP_BUCKET] cleanup_failed task_id=%s error=%s", trusted_task_id, cleanup_warning)
+
+    zip_path, metadata_path = object_paths(trusted_user_id, trusted_task_id)
+    remote = _file_map(api, cfg.bucket_id, [zip_path, metadata_path])
+    if zip_path in remote or metadata_path in remote:
+        if zip_path in remote and metadata_path in remote:
+            existing = _download_metadata(api, cfg.bucket_id, metadata_path)
+            same = (
+                str(existing.get("user_id") or "") == trusted_user_id
+                and str(existing.get("task_id") or "") == trusted_task_id
+                and str(existing.get("output_filename") or "") == filename
+                and str(existing.get("sha256") or "") == sha256
+                and int(existing.get("size_bytes") or -1) == size_bytes
+                and int(getattr(remote[zip_path], "size", -1) or -1) == size_bytes
+            )
+            if same:
+                return {
+                    "ok": True,
+                    "stored": True,
+                    "idempotent": True,
+                    "sha256": sha256,
+                    "size_bytes": size_bytes,
+                    "cleanup_warning": cleanup_warning,
+                }
+        raise RuntimeError("hf_bucket_artifact_conflict")
+
+    info = _private_bucket_info(api, cfg)
+    current_size = int(getattr(info, "size", 0) or 0)
+    if current_size >= cfg.warning_bytes:
+        logging.warning("[SP_BUCKET] capacity_warning bytes=%s", current_size)
+    uploaded_at = (now or datetime.now(timezone.utc)).astimezone(timezone.utc)
+    metadata = _metadata_bytes(
+        user_id=trusted_user_id,
+        task_id=trusted_task_id,
+        filename=filename,
+        sha256=sha256,
+        size_bytes=size_bytes,
+        uploaded_at=uploaded_at,
+    )
+    projected = current_size + size_bytes + len(metadata)
+    if projected >= cfg.hard_stop_bytes:
+        raise RuntimeError("hf_bucket_hard_stop_reached")
+
+    api.batch_bucket_files(
+        cfg.bucket_id,
+        add=[(local, zip_path), (metadata, metadata_path)],
+    )
+    verified = _file_map(api, cfg.bucket_id, [zip_path, metadata_path])
+    if zip_path not in verified or metadata_path not in verified:
+        raise RuntimeError("hf_bucket_upload_incomplete")
+    if int(getattr(verified[zip_path], "size", -1) or -1) != size_bytes:
+        raise RuntimeError("hf_bucket_upload_size_mismatch")
+    existing = _download_metadata(api, cfg.bucket_id, metadata_path)
+    if (
+        str(existing.get("sha256") or "") != sha256
+        or int(existing.get("size_bytes") or -1) != size_bytes
+        or str(existing.get("task_id") or "") != trusted_task_id
+        or str(existing.get("user_id") or "") != trusted_user_id
+    ):
+        raise RuntimeError("hf_bucket_upload_metadata_mismatch")
+    return {
+        "ok": True,
+        "stored": True,
+        "idempotent": False,
+        "sha256": sha256,
+        "size_bytes": size_bytes,
+        "cleanup_warning": cleanup_warning,
+    }
+
+
+def restore_completed_zip(
+    *,
+    user_id: int | str,
+    task_id: str,
+    output_filename: str,
+    output_path: str | Path,
+    now: datetime | None = None,
+) -> bool:
+    """Restore a completed ZIP from a private bucket, with ownership/integrity checks."""
+    cfg = config()
+    if not cfg.enabled:
+        return False
+    _validate_config(cfg)
+    trusted_user_id, trusted_task_id = _trusted_key_parts(user_id, task_id)
+    filename = _safe_output_filename(output_filename)
+    target = Path(output_path).resolve()
+    if target.is_file():
+        return True
+    if target.name != filename:
+        raise RuntimeError("hf_bucket_restore_target_invalid")
+
+    api = _new_api(cfg.token)
+    _private_bucket_info(api, cfg)
+    zip_path, metadata_path = object_paths(trusted_user_id, trusted_task_id)
+    remote = _file_map(api, cfg.bucket_id, [zip_path, metadata_path])
+    if zip_path not in remote or metadata_path not in remote:
+        return False
+
+    metadata = _download_metadata(api, cfg.bucket_id, metadata_path)
+    uploaded_at = _parse_uploaded_at(metadata.get("uploaded_at"))
+    current = (now or datetime.now(timezone.utc)).astimezone(timezone.utc)
+    if uploaded_at > current + timedelta(minutes=5):
+        raise RuntimeError("hf_bucket_metadata_time_future")
+    if current - uploaded_at > timedelta(hours=cfg.ttl_hours):
+        return False
+    expected_size = int(metadata.get("size_bytes") or -1)
+    expected_sha = str(metadata.get("sha256") or "")
+    if (
+        str(metadata.get("schema_version") or "") != "articleagent-hf-bucket-v1"
+        or str(metadata.get("user_id") or "") != trusted_user_id
+        or str(metadata.get("task_id") or "") != trusted_task_id
+        or str(metadata.get("output_filename") or "") != filename
+        or expected_size <= 0
+        or not re.fullmatch(r"[0-9a-f]{64}", expected_sha)
+        or int(getattr(remote[zip_path], "size", -1) or -1) != expected_size
+    ):
+        raise RuntimeError("hf_bucket_restore_metadata_mismatch")
+
+    target.parent.mkdir(parents=True, exist_ok=True)
+    temporary = target.with_suffix(target.suffix + ".bucket.tmp")
+    try:
+        temporary.unlink(missing_ok=True)
+        api.download_bucket_files(cfg.bucket_id, [(zip_path, temporary)], raise_on_missing_files=True)
+        observed_sha, observed_size = _sha256_file(temporary)
+        if observed_size != expected_size or observed_sha != expected_sha:
+            raise RuntimeError("hf_bucket_restore_integrity_mismatch")
+        with zipfile.ZipFile(temporary, "r") as archive:
+            if archive.testzip() is not None:
+                raise RuntimeError("hf_bucket_restore_zip_invalid")
+        os.replace(temporary, target)
+    finally:
+        temporary.unlink(missing_ok=True)
+    return True

managed_auth.py

@@ -14,15 +14,19 @@ import hmac
 import html
 import io
 import json
+import logging
 import os
+import re
 import secrets
 import time
+from pathlib import Path
 from typing import Any, Callable
 from urllib.parse import unquote, urlsplit
 
 import qrcode
 import streamlit as st
 
+from config import APP_CONFIG
 from i18n import tr
 from smartpaper_restore_bridge import edge_bridge
 
@@ -388,6 +392,48 @@ def start_generation_task(
     return dict(result)
 
 
+def _completed_output_path(task: dict[str, Any]) -> Path | None:
+    user = dict(st.session_state.get("sp_user") or {})
+    user_key = re.sub(r"[^A-Za-z0-9_.-]", "_", str(user.get("id") or "").strip())[:100]
+    task_key = re.sub(r"[^A-Za-z0-9_.-]", "_", str((task or {}).get("task_id") or "").strip())[:100]
+    filename = os.path.basename(str((task or {}).get("output_filename") or "").strip())
+    if not user_key or not task_key or not filename or not filename.lower().endswith(".zip"):
+        return None
+    workspace = Path(str(APP_CONFIG.get("workspace_dir", ".workspace")))
+    return workspace / "users" / user_key / "tasks" / task_key / "outputs" / filename
+
+
+def _restore_completed_artifact_if_needed(task: dict[str, Any]) -> None:
+    """Best-effort restore for a completed task when this runtime lost its local ZIP."""
+    if str((task or {}).get("status") or "").strip().lower() != "completed":
+        return
+    output_path = _completed_output_path(task)
+    if output_path is None or output_path.is_file():
+        return
+    user = dict(st.session_state.get("sp_user") or {})
+    numeric_user_id = user.get("numeric_id")
+    if numeric_user_id in (None, ""):
+        logging.warning("[SP_AUTH] bucket_restore_skipped task_id=%s reason=numeric_user_missing", str(task.get("task_id") or ""))
+        return
+    try:
+        from hf_bucket_store import restore_completed_zip
+
+        restored = restore_completed_zip(
+            user_id=numeric_user_id,
+            task_id=str(task.get("task_id") or ""),
+            output_filename=str(task.get("output_filename") or ""),
+            output_path=output_path,
+        )
+        if restored:
+            logging.info("[SP_AUTH] bucket_restore_ok task_id=%s", str(task.get("task_id") or ""))
+    except Exception as exc:
+        # Recovery storage is never authoritative for task state or quota.
+        logging.warning(
+            "[SP_AUTH] bucket_restore_nonfatal task_id=%s error=%s",
+            str((task or {}).get("task_id") or ""), type(exc).__name__,
+        )
+
+
 def recover_generation_task() -> dict[str, Any]:
     """Initialize current task state once from the authoritative Worker.
 
@@ -404,6 +450,7 @@ def recover_generation_task() -> dict[str, Any]:
     _merge_quota_into_session(result)
     task = result.get("task")
     if isinstance(task, dict) and task.get("task_id"):
+        _restore_completed_artifact_if_needed(task)
         st.session_state["sp_generation_task"] = dict(task)
         st.session_state["sp_tracked_task_id"] = str(task["task_id"])
         st.session_state["sp_last_ui_task_status"] = str(task.get("status") or "")

requirements.txt

@@ -12,3 +12,5 @@ PyMuPDF>=1.24,<2
 openpyxl>=3.1,<4
 tabulate>=0.9,<1
 qrcode[pil]>=8,<9
+
+huggingface_hub>=1.16,<2

task_executor.py

@@ -755,6 +755,7 @@ def main() -> int:
             logging.warning("[SP_EXECUTOR] final_heartbeat_error task_id=%s error=%s", args.task_id, _safe(exc, 180))
         if state.cancelled():
             raise RuntimeError("task_cancel_requested")
+        output_filename = str(result.get("output_filename") or "")[:240]
         api.finish(
             {
                 "task_id": args.task_id,
@@ -763,7 +764,7 @@ def main() -> int:
                 "executor_instance_id": instance_id,
                 "status": "completed",
                 "current_node": "package_validated",
-                "output_filename": str(result.get("output_filename") or "")[:240],
+                "output_filename": output_filename,
                 # R14-TQ2-PQ1-R3-R1: completed tasks may carry a nonfatal,
                 # backend-trusted quality warning. Worker/D1 already persists
                 # safe_error_code/detail for completed terminal transitions;
@@ -772,6 +773,28 @@ def main() -> int:
                 "safe_error_detail": str(result.get("safe_error_detail") or "")[:800],
             }
         )
+        # P0-R2: private HF Bucket storage is a best-effort recovery layer only.
+        # The authoritative completed/quota transition above must succeed first;
+        # any storage failure below is nonfatal and cannot alter task accounting.
+        try:
+            from hf_bucket_store import archive_completed_zip
+
+            archive_result = archive_completed_zip(
+                user_id=args.user_id,
+                task_id=args.task_id,
+                output_filename=output_filename,
+                output_path=str(result.get("output_path") or ""),
+            )
+            if archive_result.get("stored"):
+                logging.info(
+                    "[SP_EXECUTOR] bucket_backup_ok task_id=%s bytes=%s idempotent=%s",
+                    args.task_id, archive_result.get("size_bytes"), bool(archive_result.get("idempotent")),
+                )
+        except Exception as exc:
+            logging.warning(
+                "[SP_EXECUTOR] bucket_backup_nonfatal task_id=%s error=%s",
+                args.task_id, type(exc).__name__,
+            )
         logging.info("[SP_EXECUTOR] completed task_id=%s epoch=%s", args.task_id, epoch)
         return 0
     except Exception as exc: