research/research_workspace.py

"""Resolve and prepare per-user Hugging Face bucket workspaces."""

from __future__ import annotations

import hashlib
import json
import os
import re
from contextvars import ContextVar
from dataclasses import dataclass
from datetime import UTC, datetime
from typing import Any, Mapping
from uuid import uuid4

from huggingface_hub import HfApi, get_token
from huggingface_hub.errors import BucketNotFoundError

from fast_agent import AgentAuth
from fast_agent.mcp.server.common import normalize_serve_oauth_provider


_SAFE_SEGMENT = re.compile(r"[^A-Za-z0-9._-]+")


@dataclass(frozen=True, slots=True)
class ResearchWorkspace:
    username: str
    session_id: str
    bucket_id: str
    root: str
    scratch: str
    output: str
    bucket_created: bool
    marker_paths: tuple[str, ...]
    bearer_token: str | None
    archive_space_id: str | None = None
    archive_space_url: str | None = None
    archive_app_url: str | None = None
    archive_status: str | None = None
    archive_template_version: str | None = None
    archive_installed_version: str | None = None
    archive_error: str | None = None


current_research_workspace: ContextVar[ResearchWorkspace | None] = ContextVar(
    "current_research_workspace",
    default=None,
)


def ensure_workspace(
    *,
    auth: AgentAuth | None,
    request_metadata: Mapping[str, Any],
    open_metadata: Mapping[str, object],
    create_bucket: bool = True,
    write_markers: bool = True,
    api: HfApi | None = None,
) -> ResearchWorkspace:
    """Resolve identity/session, ensure the bucket exists, and write markers."""
    token = _token(auth)
    whoami = _whoami(auth, token)
    username = _username(whoami)
    session_id = _session_id(request_metadata, open_metadata)
    bucket_id = f"{username}/research-agent"
    root = f"hf://buckets/{bucket_id}/{session_id}/"

    api = api or HfApi()
    bucket_created = False
    try:
        api.bucket_info(bucket_id, token=token)
    except BucketNotFoundError as exc:
        if not create_bucket:
            raise RuntimeError(
                f"Bucket {bucket_id!r} is not accessible: {exc}"
            ) from exc
        try:
            api.create_bucket(bucket_id, private=True, exist_ok=True, token=token)
            bucket_created = True
        except Exception as create_exc:
            raise RuntimeError(
                f"Could not create/access bucket {bucket_id!r}: {create_exc}"
            ) from create_exc

    marker_paths: tuple[str, ...] = ()
    if write_markers:
        marker = {
            "server": "research-agent",
            "username": username,
            "session_id": session_id,
            "bucket_id": bucket_id,
            "checked_at": datetime.now(UTC).isoformat(),
        }
        try:
            api.batch_bucket_files(
                bucket_id,
                add=[
                    (
                        json.dumps(marker, indent=2).encode("utf-8"),
                        f"{session_id}/scratch/.workspace.json",
                    ),
                    (b"", f"{session_id}/output/.keep"),
                ],
                token=token,
            )
        except Exception as exc:
            raise RuntimeError(
                f"Bucket {bucket_id!r} is accessible but marker write failed: {exc}"
            ) from exc
        marker_paths = (
            f"{root}scratch/.workspace.json",
            f"{root}output/.keep",
        )

    return ResearchWorkspace(
        username=username,
        session_id=session_id,
        bucket_id=bucket_id,
        root=root,
        scratch=f"{root}scratch/",
        output=f"{root}output/",
        bucket_created=bucket_created,
        marker_paths=marker_paths,
        bearer_token=token,
    )


def _token(auth: AgentAuth | None) -> str | None:
    if auth is not None and auth.token:
        return auth.token
    oauth_provider = normalize_serve_oauth_provider(os.getenv("FAST_AGENT_SERVE_OAUTH"))
    if oauth_provider == "huggingface":
        raise RuntimeError(
            "Hugging Face OAuth is enabled, but this request has no caller token."
        )
    env_token = os.getenv("HF_TOKEN")
    if env_token:
        return env_token
    return get_token()


def _whoami(auth: AgentAuth | None, token: str | bool | None) -> Mapping[str, Any]:
    """Return the authoritative Hugging Face whoami payload for this caller."""
    claims = dict(auth.claims) if auth is not None else {}
    whoami = claims.get("huggingface_whoami")
    if isinstance(whoami, dict) and whoami:
        return whoami

    try:
        return HfApi().whoami(token=token)
    except Exception as exc:
        raise RuntimeError(
            "Could not determine the Hugging Face user. Provide a bearer token, "
            "enable Hugging Face OAuth, set HF_TOKEN, or run `hf auth login`."
        ) from exc


def _username(whoami: Mapping[str, Any]) -> str:
    username = safe_segment(whoami.get("name"))
    if username:
        return username
    raise RuntimeError(
        f"Hugging Face whoami response did not include a usable name: {dict(whoami)!r}."
    )


def _session_id(
    request_metadata: Mapping[str, Any],
    open_metadata: Mapping[str, object],
) -> str:
    candidates = [
        request_metadata.get("research_workspace_id"),
        open_metadata.get("research_workspace_id"),
        request_metadata.get("request_session_id"),
        request_metadata.get("harness_session_id"),
        request_metadata.get("requested_session_id"),
        request_metadata.get("mcp_session_id"),
        open_metadata.get("harness_session_id"),
        open_metadata.get("requested_session_id"),
        open_metadata.get("mcp_session_id"),
    ]
    for candidate in candidates:
        value = _safe_session_segment(candidate)
        if value:
            return value
    # No usable session identity was supplied. Never fall back to a shared
    # constant ("default") — concurrent runs would collide on one bucket path
    # and leak one run's report into another's UI. Mint a unique id instead.
    return f"session-{uuid4().hex}"


def _safe_session_segment(value: object) -> str | None:
    """Sanitize a session id, keeping distinct inputs on distinct segments.

    ``safe_segment`` truncates to 96 chars and maps disallowed characters to
    ``-``, so two different client-supplied ids can collapse to the same
    segment. When sanitization loses information, append a short stable hash of
    the original so the mapping stays collision-resistant (and deterministic, so
    the same input still resolves to the same workspace across requests).
    """
    if value is None:
        return None
    raw = str(value).strip().strip("/")
    if not raw:
        return None
    safe = safe_segment(raw)
    if safe is None:
        return None
    if safe != raw:
        digest = hashlib.sha256(raw.encode("utf-8")).hexdigest()[:8]
        safe = f"{safe[:87].rstrip('.-_')}-{digest}"
    return safe


def safe_segment(value: object) -> str | None:
    if value is None:
        return None
    text = str(value).strip().strip("/")
    if not text:
        return None
    return _SAFE_SEGMENT.sub("-", text)[:96].strip(".-_") or None