Add ML Intern Hub artifact metadata (#225)

* Add ML Intern Hub artifact metadata

Co-authored-by: OpenAI Codex <codex@openai.com>

* Extend Hub artifact hooks to sandbox bash

Co-authored-by: OpenAI Codex <codex@openai.com>

* Create Hub artifact collections at session start

Co-authored-by: OpenAI Codex <codex@openai.com>

* Address Hub artifact PR review comments

Co-authored-by: OpenAI Codex <codex@openai.com>

* Respect Hub collection title limits

Co-authored-by: OpenAI Codex <codex@openai.com>

* Shorten collection session UUID fragments

Co-authored-by: OpenAI Codex <codex@openai.com>

---------

Co-authored-by: OpenAI Codex <codex@openai.com>

agent/core/agent_loop.py

@@ -26,6 +26,7 @@ from agent.core.cost_estimation import CostEstimate, estimate_tool_cost
 from agent.messaging.gateway import NotificationGateway
 from agent.core import telemetry
 from agent.core.doom_loop import check_for_doom_loop
+from agent.core.hub_artifacts import start_session_artifact_collection_task
 from agent.core.llm_params import _resolve_llm_params
 from agent.core.prompt_caching import with_prompt_caching
 from agent.core.session import Event, OpType, Session
@@ -1998,6 +1999,7 @@ async def submission_loop(
     )
     if session_holder is not None:
         session_holder[0] = session
+    start_session_artifact_collection_task(session, token=hf_token)
     logger.info("Agent loop started")
 
     # Retry any failed uploads from previous sessions (fire-and-forget).

agent/core/hub_artifacts.py

@@ -0,0 +1,765 @@
+"""Best-effort Hub metadata for artifacts generated by ML Intern sessions."""
+
+import asyncio
+import base64
+import logging
+import re
+import shlex
+import tempfile
+import textwrap
+from datetime import datetime
+from pathlib import Path
+from typing import Any
+
+from huggingface_hub import HfApi, hf_hub_download
+from huggingface_hub.repocard import metadata_load, metadata_save
+from huggingface_hub.utils import EntryNotFoundError, RepositoryNotFoundError
+
+logger = logging.getLogger(__name__)
+
+ML_INTERN_TAG = "ml-intern"
+SUPPORTED_REPO_TYPES = {"model", "dataset", "space"}
+PROVENANCE_MARKER = "<!-- ml-intern-provenance -->"
+_COLLECTION_TITLE_PREFIX = "ml-intern-artifacts"
+_COLLECTION_TITLE_MAX_LENGTH = 59
+_UUID_SESSION_ID_RE = re.compile(
+    r"^[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-"
+    r"[0-9a-fA-F]{4}-[0-9a-fA-F]{12}$"
+)
+_KNOWN_ARTIFACTS_ATTR = "_ml_intern_known_hub_artifacts"
+_REGISTERED_ARTIFACTS_ATTR = "_ml_intern_registered_hub_artifacts"
+_COLLECTION_SLUG_ATTR = "_ml_intern_artifact_collection_slug"
+_COLLECTION_TASK_ATTR = "_ml_intern_artifact_collection_task"
+_SESSION_ARTIFACT_SET_FALLBACK: dict[tuple[int, str], set[str]] = {}
+_USAGE_HEADING_RE = re.compile(
+    r"^#{2,6}\s+(usage|how to use|using this (model|dataset)|use this (model|dataset))\b",
+    re.IGNORECASE | re.MULTILINE,
+)
+_FRONT_MATTER_RE = re.compile(r"\A---\s*\n.*?\n---\s*\n?", re.DOTALL)
+
+
+def _safe_session_id(session: Any) -> str:
+    raw = str(getattr(session, "session_id", "") or "unknown-session")
+    safe = re.sub(r"[^A-Za-z0-9._-]+", "-", raw).strip("-")
+    return safe or "unknown-session"
+
+
+def session_artifact_date(session: Any) -> str:
+    """Return the YYYY-MM-DD partition date for a session."""
+    raw = getattr(session, "session_start_time", None)
+    if raw:
+        try:
+            return datetime.fromisoformat(str(raw).replace("Z", "+00:00")).strftime(
+                "%Y-%m-%d"
+            )
+        except ValueError:
+            logger.debug("Could not parse session_start_time=%r", raw)
+    return datetime.utcnow().strftime("%Y-%m-%d")
+
+
+def _collection_session_id_fragment(session: Any) -> str:
+    safe_id = _safe_session_id(session)
+    if _UUID_SESSION_ID_RE.match(safe_id):
+        return safe_id[:8]
+    stem = f"{_COLLECTION_TITLE_PREFIX}-{session_artifact_date(session)}-"
+    max_id_length = max(1, _COLLECTION_TITLE_MAX_LENGTH - len(stem))
+    if len(safe_id) <= max_id_length:
+        return safe_id
+    return safe_id[:max_id_length].rstrip("-._") or safe_id[:max_id_length]
+
+
+def artifact_collection_title(session: Any) -> str:
+    return (
+        f"{_COLLECTION_TITLE_PREFIX}-{session_artifact_date(session)}-"
+        f"{_collection_session_id_fragment(session)}"
+    )
+
+
+def _artifact_key(repo_id: str, repo_type: str | None) -> str:
+    return f"{repo_type or 'model'}:{repo_id}"
+
+
+def _session_artifact_set(session: Any, attr: str) -> set[str]:
+    current = getattr(session, attr, None)
+    if isinstance(current, set):
+        return current
+    current = set()
+    try:
+        setattr(session, attr, current)
+    except Exception:
+        logger.warning(
+            "Could not attach %s to session; using process-local fallback state",
+            attr,
+        )
+        return _SESSION_ARTIFACT_SET_FALLBACK.setdefault((id(session), attr), set())
+    return current
+
+
+def remember_hub_artifact(session: Any, repo_id: str, repo_type: str | None) -> None:
+    if session is None or not repo_id:
+        return
+    _session_artifact_set(session, _KNOWN_ARTIFACTS_ATTR).add(
+        _artifact_key(repo_id, repo_type)
+    )
+
+
+def is_known_hub_artifact(session: Any, repo_id: str, repo_type: str | None) -> bool:
+    if session is None or not repo_id:
+        return False
+    return _artifact_key(repo_id, repo_type) in _session_artifact_set(
+        session, _KNOWN_ARTIFACTS_ATTR
+    )
+
+
+def _merge_tags(metadata: dict[str, Any], tag: str = ML_INTERN_TAG) -> dict[str, Any]:
+    merged = dict(metadata)
+    raw_tags = merged.get("tags")
+    if raw_tags is None:
+        tags: list[str] = []
+    elif isinstance(raw_tags, str):
+        tags = [raw_tags]
+    elif isinstance(raw_tags, list):
+        tags = [str(item) for item in raw_tags]
+    else:
+        tags = [str(raw_tags)]
+
+    if tag not in tags:
+        tags.append(tag)
+    merged["tags"] = tags
+    return merged
+
+
+def _metadata_from_content(content: str) -> dict[str, Any]:
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        path = Path(tmp_dir) / "README.md"
+        path.write_text(content, encoding="utf-8")
+        return metadata_load(path) or {}
+
+
+def _content_with_metadata(content: str, metadata: dict[str, Any]) -> str:
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        path = Path(tmp_dir) / "README.md"
+        path.write_text(content, encoding="utf-8")
+        metadata_save(path, metadata)
+        return path.read_text(encoding="utf-8")
+
+
+def _body_without_metadata(content: str) -> str:
+    return _FRONT_MATTER_RE.sub("", content, count=1).strip()
+
+
+def _append_section(content: str, section: str) -> str:
+    base = content.rstrip()
+    if base:
+        return f"{base}\n\n{section.strip()}\n"
+    return f"{section.strip()}\n"
+
+
+def _provenance_section(repo_type: str) -> str:
+    label = {"model": "model", "dataset": "dataset"}.get(repo_type, "Hub")
+    return f"""{PROVENANCE_MARKER}
+## Generated by ML Intern
+
+This {label} repository was generated by [ML Intern](https://github.com/huggingface/ml-intern), an agent for machine learning research and development on the Hugging Face Hub.
+
+- Try ML Intern: https://smolagents-ml-intern.hf.space
+- Source code: https://github.com/huggingface/ml-intern
+"""
+
+
+def _usage_section(repo_id: str, repo_type: str) -> str:
+    if repo_type == "dataset":
+        return f"""## Usage
+
+```python
+from datasets import load_dataset
+
+dataset = load_dataset("{repo_id}")
+```
+"""
+
+    return f"""## Usage
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model_id = "{repo_id}"
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(model_id)
+```
+
+For non-causal architectures, replace `AutoModelForCausalLM` with the appropriate `AutoModel` class.
+"""
+
+
+def augment_repo_card_content(
+    content: str | None,
+    repo_id: str,
+    repo_type: str = "model",
+    *,
+    extra_metadata: dict[str, Any] | None = None,
+) -> str:
+    """Return README content with ML Intern metadata and provenance added."""
+    repo_type = repo_type or "model"
+    content = content or ""
+    metadata = _metadata_from_content(content)
+    if extra_metadata:
+        metadata = {**extra_metadata, **metadata}
+    metadata = _merge_tags(metadata)
+    updated = _content_with_metadata(content, metadata)
+
+    if not _body_without_metadata(updated):
+        updated = _append_section(updated, f"# {repo_id}")
+
+    if repo_type in {"model", "dataset"} and PROVENANCE_MARKER not in updated:
+        updated = _append_section(updated, _provenance_section(repo_type))
+        if not _USAGE_HEADING_RE.search(content):
+            updated = _append_section(updated, _usage_section(repo_id, repo_type))
+
+    return updated
+
+
+def _read_remote_readme(
+    api: Any,
+    repo_id: str,
+    repo_type: str,
+    *,
+    token: str | bool | None = None,
+) -> str:
+    token_value = token if token is not None else getattr(api, "token", None)
+    try:
+        readme_path = hf_hub_download(
+            repo_id=repo_id,
+            filename="README.md",
+            repo_type=repo_type,
+            token=token_value,
+        )
+    except (EntryNotFoundError, RepositoryNotFoundError):
+        return ""
+    return Path(readme_path).read_text(encoding="utf-8")
+
+
+def _update_repo_card(
+    api: Any,
+    repo_id: str,
+    repo_type: str,
+    *,
+    token: str | bool | None = None,
+    extra_metadata: dict[str, Any] | None = None,
+) -> None:
+    current = _read_remote_readme(api, repo_id, repo_type, token=token)
+    updated = augment_repo_card_content(
+        current,
+        repo_id,
+        repo_type,
+        extra_metadata=extra_metadata,
+    )
+    if updated == current:
+        return
+    api.upload_file(
+        path_or_fileobj=updated.encode("utf-8"),
+        path_in_repo="README.md",
+        repo_id=repo_id,
+        repo_type=repo_type,
+        token=token,
+        commit_message="Update ML Intern artifact metadata",
+    )
+
+
+def _ensure_collection_slug(
+    api: Any,
+    session: Any,
+    *,
+    token: str | bool | None = None,
+) -> str | None:
+    slug = getattr(session, _COLLECTION_SLUG_ATTR, None)
+    if slug:
+        return slug
+
+    title = artifact_collection_title(session)
+    collection = api.create_collection(
+        title=title,
+        description=(
+            f"Artifacts generated by ML Intern session {_safe_session_id(session)} "
+            f"on {session_artifact_date(session)}."
+        ),
+        private=True,
+        exists_ok=True,
+        token=token,
+    )
+    slug = getattr(collection, "slug", None)
+    if slug:
+        setattr(session, _COLLECTION_SLUG_ATTR, slug)
+    return slug
+
+
+async def ensure_session_artifact_collection(
+    session: Any,
+    *,
+    token: str | bool | None = None,
+) -> str | None:
+    """Create/cache the per-session artifact collection without raising."""
+    if session is None or not getattr(session, "session_id", None):
+        return None
+    token_value = token if token is not None else getattr(session, "hf_token", None)
+    if not token_value:
+        return None
+
+    try:
+        api = HfApi(token=token_value)
+        return await asyncio.to_thread(
+            _ensure_collection_slug,
+            api,
+            session,
+            token=token_value,
+        )
+    except Exception as e:
+        logger.warning(
+            "ML Intern session collection creation failed for %s: %s",
+            _safe_session_id(session),
+            e,
+        )
+        return None
+
+
+def start_session_artifact_collection_task(
+    session: Any,
+    *,
+    token: str | bool | None = None,
+) -> asyncio.Task | None:
+    """Schedule best-effort collection creation for a newly started session."""
+    if session is None or not getattr(session, "session_id", None):
+        return None
+    if getattr(session, _COLLECTION_SLUG_ATTR, None):
+        return None
+
+    token_value = token if token is not None else getattr(session, "hf_token", None)
+    if not token_value:
+        return None
+
+    existing = getattr(session, _COLLECTION_TASK_ATTR, None)
+    if isinstance(existing, asyncio.Task) and not existing.done():
+        return existing
+
+    try:
+        loop = asyncio.get_running_loop()
+    except RuntimeError:
+        return None
+
+    async def _run() -> None:
+        await ensure_session_artifact_collection(session, token=token_value)
+
+    task = loop.create_task(_run())
+    try:
+        setattr(session, _COLLECTION_TASK_ATTR, task)
+    except Exception:
+        logger.debug("Could not attach ML Intern collection task to session")
+    return task
+
+
+def _add_to_collection(
+    api: Any,
+    session: Any,
+    repo_id: str,
+    repo_type: str,
+    *,
+    token: str | bool | None = None,
+) -> None:
+    slug = _ensure_collection_slug(api, session, token=token)
+    if not slug:
+        return
+    api.add_collection_item(
+        collection_slug=slug,
+        item_id=repo_id,
+        item_type=repo_type,
+        note=(
+            f"Generated by ML Intern session {_safe_session_id(session)} "
+            f"on {session_artifact_date(session)}."
+        ),
+        exists_ok=True,
+        token=token,
+    )
+
+
+def register_hub_artifact(
+    api: Any,
+    repo_id: str,
+    repo_type: str = "model",
+    *,
+    session: Any = None,
+    token: str | bool | None = None,
+    extra_metadata: dict[str, Any] | None = None,
+    force: bool = False,
+) -> bool:
+    """Tag, card, and collection-register a Hub artifact without raising."""
+    if session is None or not repo_id:
+        return False
+    repo_type = repo_type or "model"
+    if repo_type not in SUPPORTED_REPO_TYPES:
+        return False
+
+    key = _artifact_key(repo_id, repo_type)
+    remember_hub_artifact(session, repo_id, repo_type)
+    registered = _session_artifact_set(session, _REGISTERED_ARTIFACTS_ATTR)
+    if key in registered and not force:
+        return True
+
+    token_value = token if token is not None else getattr(api, "token", None)
+    card_updated = False
+    collection_updated = False
+    try:
+        _update_repo_card(
+            api,
+            repo_id,
+            repo_type,
+            token=token_value,
+            extra_metadata=extra_metadata,
+        )
+        card_updated = True
+    except Exception as e:
+        logger.debug("ML Intern repo-card update failed for %s: %s", repo_id, e)
+
+    try:
+        _add_to_collection(api, session, repo_id, repo_type, token=token_value)
+        collection_updated = True
+    except Exception as e:
+        logger.debug("ML Intern collection update failed for %s: %s", repo_id, e)
+
+    if card_updated and collection_updated:
+        registered.add(key)
+        return True
+    return False
+
+
+def build_hub_artifact_sitecustomize(session: Any) -> str:
+    """Build standalone sitecustomize.py code for HF Jobs Python processes."""
+    if session is None or not getattr(session, "session_id", None):
+        return ""
+
+    session_id = _safe_session_id(session)
+    session_date = session_artifact_date(session)
+    collection_title = artifact_collection_title(session)
+    collection_slug = getattr(session, _COLLECTION_SLUG_ATTR, None)
+
+    return (
+        textwrap.dedent(
+            f"""
+        # Auto-generated by ML Intern. Best-effort Hub artifact metadata only.
+        def _install_ml_intern_artifact_hooks():
+            import os
+            import re
+            import tempfile
+            from pathlib import Path
+
+            try:
+                import huggingface_hub as _hub
+                from huggingface_hub import HfApi, hf_hub_download
+                from huggingface_hub.repocard import metadata_load, metadata_save
+                from huggingface_hub.utils import EntryNotFoundError, RepositoryNotFoundError
+            except Exception:
+                return
+
+            session_id = {session_id!r}
+            session_date = {session_date!r}
+            collection_title = {collection_title!r}
+            tag = {ML_INTERN_TAG!r}
+            marker = {PROVENANCE_MARKER!r}
+            supported = {sorted(SUPPORTED_REPO_TYPES)!r}
+            registering = False
+            collection_slug = {collection_slug!r}
+            registered = set()
+            usage_re = re.compile(
+                r"^#{{2,6}}\\s+(usage|how to use|using this (model|dataset)|use this (model|dataset))\\b",
+                re.IGNORECASE | re.MULTILINE,
+            )
+            front_matter_re = re.compile(r"\\A---\\s*\\n.*?\\n---\\s*\\n?", re.DOTALL)
+
+            def _token(value=None, api=None):
+                if isinstance(value, str) and value:
+                    return value
+                api_token = getattr(api, "token", None)
+                if isinstance(api_token, str) and api_token:
+                    return api_token
+                return (
+                    os.environ.get("HF_TOKEN")
+                    or os.environ.get("HUGGINGFACE_HUB_TOKEN")
+                    or None
+                )
+
+            def _merge_tags(metadata):
+                metadata = dict(metadata or {{}})
+                raw_tags = metadata.get("tags")
+                if raw_tags is None:
+                    tags = []
+                elif isinstance(raw_tags, str):
+                    tags = [raw_tags]
+                elif isinstance(raw_tags, list):
+                    tags = [str(item) for item in raw_tags]
+                else:
+                    tags = [str(raw_tags)]
+                if tag not in tags:
+                    tags.append(tag)
+                metadata["tags"] = tags
+                return metadata
+
+            def _metadata_from_content(content):
+                with tempfile.TemporaryDirectory() as tmp_dir:
+                    path = Path(tmp_dir) / "README.md"
+                    path.write_text(content or "", encoding="utf-8")
+                    return metadata_load(path) or {{}}
+
+            def _content_with_metadata(content, metadata):
+                with tempfile.TemporaryDirectory() as tmp_dir:
+                    path = Path(tmp_dir) / "README.md"
+                    path.write_text(content or "", encoding="utf-8")
+                    metadata_save(path, metadata)
+                    return path.read_text(encoding="utf-8")
+
+            def _body_without_metadata(content):
+                return front_matter_re.sub("", content or "", count=1).strip()
+
+            def _append_section(content, section):
+                base = (content or "").rstrip()
+                if base:
+                    return base + "\\n\\n" + section.strip() + "\\n"
+                return section.strip() + "\\n"
+
+            def _provenance(repo_type):
+                label = {{"model": "model", "dataset": "dataset"}}.get(
+                    repo_type, "Hub"
+                )
+                return (
+                    marker
+                    + "\\n## Generated by ML Intern\\n\\n"
+                    + f"This {{label}} repository was generated by [ML Intern](https://github.com/huggingface/ml-intern), an agent for machine learning research and development on the Hugging Face Hub.\\n\\n"
+                    + "- Try ML Intern: https://smolagents-ml-intern.hf.space\\n"
+                    + "- Source code: https://github.com/huggingface/ml-intern\\n"
+                )
+
+            def _usage(repo_id, repo_type):
+                if repo_type == "dataset":
+                    return (
+                        "## Usage\\n\\n"
+                        "```python\\n"
+                        "from datasets import load_dataset\\n\\n"
+                        f"dataset = load_dataset({{repo_id!r}})\\n"
+                        "```\\n"
+                    )
+                return (
+                    "## Usage\\n\\n"
+                    "```python\\n"
+                    "from transformers import AutoModelForCausalLM, AutoTokenizer\\n\\n"
+                    f"model_id = {{repo_id!r}}\\n"
+                    "tokenizer = AutoTokenizer.from_pretrained(model_id)\\n"
+                    "model = AutoModelForCausalLM.from_pretrained(model_id)\\n"
+                    "```\\n\\n"
+                    "For non-causal architectures, replace `AutoModelForCausalLM` with the appropriate `AutoModel` class.\\n"
+                )
+
+            def _augment(content, repo_id, repo_type, extra_metadata=None):
+                metadata = _metadata_from_content(content or "")
+                if extra_metadata:
+                    metadata = {{**extra_metadata, **metadata}}
+                updated = _content_with_metadata(content or "", _merge_tags(metadata))
+                if not _body_without_metadata(updated):
+                    updated = _append_section(updated, f"# {{repo_id}}")
+                if repo_type in {{"model", "dataset"}} and marker not in updated:
+                    updated = _append_section(updated, _provenance(repo_type))
+                    if not usage_re.search(content or ""):
+                        updated = _append_section(updated, _usage(repo_id, repo_type))
+                return updated
+
+            def _readme(api, repo_id, repo_type, token_value):
+                try:
+                    path = hf_hub_download(
+                        repo_id=repo_id,
+                        filename="README.md",
+                        repo_type=repo_type,
+                        token=token_value,
+                    )
+                except (EntryNotFoundError, RepositoryNotFoundError):
+                    return ""
+                return Path(path).read_text(encoding="utf-8")
+
+            def _ensure_collection(api, token_value):
+                nonlocal collection_slug
+                if collection_slug:
+                    return collection_slug
+                collection = api.create_collection(
+                    title=collection_title,
+                    description=(
+                        f"Artifacts generated by ML Intern session {{session_id}} "
+                        f"on {{session_date}}."
+                    ),
+                    private=True,
+                    exists_ok=True,
+                    token=token_value,
+                )
+                collection_slug = getattr(collection, "slug", None)
+                return collection_slug
+
+            def _register(
+                repo_id,
+                repo_type="model",
+                token_value=None,
+                extra_metadata=None,
+                force=False,
+            ):
+                nonlocal registering
+                if registering or not repo_id:
+                    return
+                repo_type = repo_type or "model"
+                if repo_type not in supported:
+                    return
+                key = f"{{repo_type}}:{{repo_id}}"
+                if key in registered and not force:
+                    return
+                registering = True
+                try:
+                    token_value = _token(token_value)
+                    api = HfApi(token=token_value)
+                    try:
+                        current = _readme(api, repo_id, repo_type, token_value)
+                        updated = _augment(
+                            current, repo_id, repo_type, extra_metadata=extra_metadata
+                        )
+                        if updated != current:
+                            _original_upload_file(
+                                api,
+                                path_or_fileobj=updated.encode("utf-8"),
+                                path_in_repo="README.md",
+                                repo_id=repo_id,
+                                repo_type=repo_type,
+                                token=token_value,
+                                commit_message="Update ML Intern artifact metadata",
+                            )
+                    except Exception:
+                        pass
+                    try:
+                        slug = _ensure_collection(api, token_value)
+                        if slug:
+                            api.add_collection_item(
+                                collection_slug=slug,
+                                item_id=repo_id,
+                                item_type=repo_type,
+                                note=(
+                                    f"Generated by ML Intern session {{session_id}} "
+                                    f"on {{session_date}}."
+                                ),
+                                exists_ok=True,
+                                token=token_value,
+                            )
+                    except Exception:
+                        pass
+                    registered.add(key)
+                finally:
+                    registering = False
+
+            _original_create_repo = HfApi.create_repo
+            _original_upload_file = HfApi.upload_file
+            _original_upload_folder = getattr(HfApi, "upload_folder", None)
+            _original_create_commit = getattr(HfApi, "create_commit", None)
+
+            def _repo_id(args, kwargs):
+                return kwargs.get("repo_id") or (args[0] if args else None)
+
+            def _repo_type(kwargs):
+                return kwargs.get("repo_type") or "model"
+
+            def _patched_create_repo(self, *args, **kwargs):
+                result = _original_create_repo(self, *args, **kwargs)
+                repo_id = _repo_id(args, kwargs)
+                repo_type = _repo_type(kwargs)
+                extra = None
+                if repo_type == "space" and kwargs.get("space_sdk"):
+                    extra = {{"sdk": kwargs.get("space_sdk")}}
+                _register(repo_id, repo_type, _token(kwargs.get("token"), self), extra)
+                return result
+
+            def _patched_upload_file(self, *args, **kwargs):
+                result = _original_upload_file(self, *args, **kwargs)
+                if not kwargs.get("create_pr"):
+                    force = kwargs.get("path_in_repo") == "README.md"
+                    _register(
+                        kwargs.get("repo_id"),
+                        _repo_type(kwargs),
+                        _token(kwargs.get("token"), self),
+                        force=force,
+                    )
+                return result
+
+            def _patched_upload_folder(self, *args, **kwargs):
+                result = _original_upload_folder(self, *args, **kwargs)
+                if not kwargs.get("create_pr"):
+                    _register(
+                        kwargs.get("repo_id"),
+                        _repo_type(kwargs),
+                        _token(kwargs.get("token"), self),
+                        force=True,
+                    )
+                return result
+
+            def _patched_create_commit(self, *args, **kwargs):
+                result = _original_create_commit(self, *args, **kwargs)
+                if not kwargs.get("create_pr"):
+                    _register(
+                        _repo_id(args, kwargs),
+                        _repo_type(kwargs),
+                        _token(kwargs.get("token"), self),
+                        force=True,
+                    )
+                return result
+
+            HfApi.create_repo = _patched_create_repo
+            HfApi.upload_file = _patched_upload_file
+            if _original_upload_folder is not None:
+                HfApi.upload_folder = _patched_upload_folder
+            if _original_create_commit is not None:
+                HfApi.create_commit = _patched_create_commit
+
+            def _patch_module_func(name, method_name):
+                original = getattr(_hub, name, None)
+                if original is None:
+                    return
+                method = getattr(HfApi, method_name)
+
+                def _patched(*args, **kwargs):
+                    api = HfApi(token=_token(kwargs.get("token")))
+                    return method(api, *args, **kwargs)
+
+                setattr(_hub, name, _patched)
+
+            _patch_module_func("create_repo", "create_repo")
+            _patch_module_func("upload_file", "upload_file")
+            if _original_upload_folder is not None:
+                _patch_module_func("upload_folder", "upload_folder")
+            if _original_create_commit is not None:
+                _patch_module_func("create_commit", "create_commit")
+
+        try:
+            _install_ml_intern_artifact_hooks()
+        except Exception:
+            pass
+        """
+        ).strip()
+        + "\n"
+    )
+
+
+def wrap_shell_command_with_hub_artifact_bootstrap(
+    command: str,
+    session: Any,
+) -> str:
+    """Prefix a shell command so child Python processes load Hub hooks."""
+    sitecustomize = build_hub_artifact_sitecustomize(session)
+    if not sitecustomize or not command:
+        return command
+
+    encoded = base64.b64encode(sitecustomize.encode("utf-8")).decode("ascii")
+    bootstrap = (
+        '_ml_intern_artifacts_dir="$(mktemp -d 2>/dev/null)" '
+        f"&& printf %s {shlex.quote(encoded)} | base64 -d "
+        '> "$_ml_intern_artifacts_dir/sitecustomize.py" '
+        '&& export PYTHONPATH="$_ml_intern_artifacts_dir${PYTHONPATH:+:$PYTHONPATH}"'
+    )
+    return f"{bootstrap}; {command}"

agent/tools/hf_repo_files_tool.py

@@ -10,6 +10,7 @@ from typing import Any, Dict, Literal, Optional
 from huggingface_hub import HfApi, hf_hub_download
 from huggingface_hub.utils import EntryNotFoundError, RepositoryNotFoundError
 
+from agent.core.hub_artifacts import is_known_hub_artifact, register_hub_artifact
 from agent.tools.types import ToolResult
 
 OperationType = Literal["list", "read", "upload", "delete"]
@@ -39,8 +40,9 @@ def _format_size(size_bytes: int) -> str:
 class HfRepoFilesTool:
     """Tool for file operations on HF repos."""
 
-    def __init__(self, hf_token: Optional[str] = None):
+    def __init__(self, hf_token: Optional[str] = None, session: Any = None):
         self.api = HfApi(token=hf_token)
+        self.session = session
 
     async def execute(self, args: Dict[str, Any]) -> ToolResult:
         """Execute the specified operation."""
@@ -214,6 +216,16 @@ class HfRepoFilesTool:
             create_pr=create_pr,
         )
 
+        if not create_pr and is_known_hub_artifact(self.session, repo_id, repo_type):
+            await _async_call(
+                register_hub_artifact,
+                self.api,
+                repo_id,
+                repo_type,
+                session=self.session,
+                force=path == "README.md",
+            )
+
         url = _build_repo_url(repo_id, repo_type)
         if create_pr and hasattr(result, "pr_url"):
             response = f"**Uploaded as PR**\n{result.pr_url}"
@@ -343,7 +355,7 @@ async def hf_repo_files_handler(
     """Handler for agent tool router."""
     try:
         hf_token = session.hf_token if session else None
-        tool = HfRepoFilesTool(hf_token=hf_token)
+        tool = HfRepoFilesTool(hf_token=hf_token, session=session)
         result = await tool.execute(arguments)
         return result["formatted"], not result.get("isError", False)
     except Exception as e:

agent/tools/hf_repo_git_tool.py

@@ -10,6 +10,7 @@ from typing import Any, Dict, Literal, Optional
 from huggingface_hub import HfApi
 from huggingface_hub.utils import RepositoryNotFoundError
 
+from agent.core.hub_artifacts import register_hub_artifact
 from agent.tools.types import ToolResult
 
 OperationType = Literal[
@@ -45,8 +46,9 @@ def _build_repo_url(repo_id: str, repo_type: str = "model") -> str:
 class HfRepoGitTool:
     """Tool for git-like operations on HF repos."""
 
-    def __init__(self, hf_token: Optional[str] = None):
+    def __init__(self, hf_token: Optional[str] = None, session: Any = None):
         self.api = HfApi(token=hf_token)
+        self.session = session
 
     async def execute(self, args: Dict[str, Any]) -> ToolResult:
         """Execute the specified operation."""
@@ -552,6 +554,17 @@ class HfRepoGitTool:
             kwargs["space_sdk"] = space_sdk
 
         result = await _async_call(self.api.create_repo, **kwargs)
+        extra_metadata = None
+        if repo_type == "space" and space_sdk:
+            extra_metadata = {"sdk": space_sdk}
+        await _async_call(
+            register_hub_artifact,
+            self.api,
+            repo_id,
+            repo_type,
+            session=self.session,
+            extra_metadata=extra_metadata,
+        )
 
         return {
             "formatted": f"**Repository created:** {repo_id}\n**Private:** {private}\n{result}",
@@ -747,7 +760,7 @@ async def hf_repo_git_handler(
     """Handler for agent tool router."""
     try:
         hf_token = session.hf_token if session else None
-        tool = HfRepoGitTool(hf_token=hf_token)
+        tool = HfRepoGitTool(hf_token=hf_token, session=session)
         result = await tool.execute(arguments)
         return result["formatted"], not result.get("isError", False)
     except Exception as e:

agent/tools/jobs_tool.py

@@ -9,6 +9,7 @@ import base64
 import http.client
 import logging
 import re
+import shlex
 from typing import Any, Awaitable, Callable, Dict, Literal, Optional
 
 import httpx
@@ -20,6 +21,7 @@ from agent.core.hf_access import (
     is_billing_error,
     resolve_jobs_namespace,
 )
+from agent.core.hub_artifacts import build_hub_artifact_sitecustomize
 from agent.core.session import Event
 from agent.tools.trackio_seed import ensure_trackio_dashboard
 from agent.tools.types import ToolResult
@@ -237,6 +239,26 @@ def _resolve_uv_command(
     return _build_uv_command(script, with_deps, python, script_args)
 
 
+def _wrap_command_with_artifact_bootstrap(
+    command: list[str], session: Any = None
+) -> list[str]:
+    """Install sitecustomize hooks before the user command runs in HF Jobs."""
+    sitecustomize = build_hub_artifact_sitecustomize(session)
+    if not sitecustomize:
+        return command
+
+    encoded = base64.b64encode(sitecustomize.encode("utf-8")).decode("ascii")
+    original_command = shlex.join(command)
+    shell = (
+        'set -e; _ml_intern_artifacts_dir="$(mktemp -d)"; '
+        f"printf %s {shlex.quote(encoded)} | base64 -d "
+        '> "$_ml_intern_artifacts_dir/sitecustomize.py"; '
+        'export PYTHONPATH="$_ml_intern_artifacts_dir${PYTHONPATH:+:$PYTHONPATH}"; '
+        f"exec {original_command}"
+    )
+    return ["/bin/sh", "-lc", shell]
+
+
 async def _async_call(func, *args, **kwargs):
     """Wrap synchronous HfApi calls for async context"""
     return await asyncio.to_thread(func, *args, **kwargs)
@@ -560,6 +582,8 @@ class HfJobsTool:
                 image = args.get("image", "python:3.12")
                 job_type = "Docker"
 
+            command = _wrap_command_with_artifact_bootstrap(command, self.session)
+
             # Run the job
             flavor = args.get("hardware_flavor", "cpu-basic")
             timeout_str = args.get("timeout", "30m")
@@ -912,6 +936,8 @@ To verify, call this tool with `{{"operation": "inspect", "job_id": "{job_id}"}}
                 image = args.get("image", "python:3.12")
                 job_type = "Docker"
 
+            command = _wrap_command_with_artifact_bootstrap(command, self.session)
+
             # Create scheduled job
             scheduled_job = await _async_call(
                 self.api.create_scheduled_job,

agent/tools/local_tools.py

@@ -15,6 +15,8 @@ import tempfile
 from pathlib import Path
 from typing import Any
 
+from agent.core.hub_artifacts import wrap_shell_command_with_hub_artifact_bootstrap
+
 
 MAX_OUTPUT_CHARS = 25_000
 MAX_LINE_LENGTH = 4000
@@ -98,10 +100,13 @@ def _truncate_output(
 # ── Handlers ────────────────────────────────────────────────────────────
 
 
-async def _bash_handler(args: dict[str, Any], **_kw) -> tuple[str, bool]:
+async def _bash_handler(
+    args: dict[str, Any], session: Any = None, **_kw
+) -> tuple[str, bool]:
     command = args.get("command", "")
     if not command:
         return "No command provided.", False
+    command = wrap_shell_command_with_hub_artifact_bootstrap(command, session)
     work_dir = args.get("work_dir", ".")
     timeout = min(args.get("timeout") or DEFAULT_TIMEOUT, MAX_TIMEOUT)
     try:

agent/tools/sandbox_tool.py

@@ -21,6 +21,7 @@ from typing import Any
 
 from huggingface_hub import HfApi, SpaceHardware
 
+from agent.core.hub_artifacts import wrap_shell_command_with_hub_artifact_bootstrap
 from agent.core.session import Event
 from agent.tools.sandbox_client import Sandbox
 from agent.tools.trackio_seed import ensure_trackio_dashboard
@@ -729,6 +730,14 @@ def _make_tool_handler(sandbox_tool_name: str):
             return "Sandbox is still starting. Please retry shortly.", False
 
         try:
+            if sandbox_tool_name == "bash" and args.get("command"):
+                args = {
+                    **args,
+                    "command": wrap_shell_command_with_hub_artifact_bootstrap(
+                        args["command"],
+                        session,
+                    ),
+                }
             result = await asyncio.to_thread(sb.call_tool, sandbox_tool_name, args)
             if result.success:
                 output = result.output or "(no output)"

backend/session_manager.py

@@ -12,10 +12,11 @@ from typing import Any, Optional
 
 from agent.config import load_config
 from agent.core.agent_loop import process_submission
-from agent.messaging.gateway import NotificationGateway
+from agent.core.hub_artifacts import start_session_artifact_collection_task
 from agent.core.session import Event, OpType, Session
 from agent.core.session_persistence import get_session_store
 from agent.core.tools import ToolRouter
+from agent.messaging.gateway import NotificationGateway
 
 # Get project root (parent of backend directory)
 PROJECT_ROOT = Path(__file__).parent.parent
@@ -135,6 +136,7 @@ class SessionManager:
         self.sessions: dict[str, AgentSession] = {}
         self._lock = asyncio.Lock()
         self.persistence_store = None
+        self.enable_hub_artifact_collections = True
 
     async def start(self) -> None:
         """Start shared background resources."""
@@ -411,6 +413,28 @@ class SessionManager:
         session.sandbox_preload_cancel_event = None
         self._start_cpu_sandbox_preload(agent_session)
 
+    def _start_hub_artifact_collection(self, agent_session: AgentSession) -> None:
+        """Kick off best-effort Hub collection creation for the session."""
+        if not getattr(self, "enable_hub_artifact_collections", False):
+            return
+        session = agent_session.session
+        if not getattr(session, "session_id", None):
+            try:
+                session.session_id = agent_session.session_id
+            except Exception:
+                logger.debug("Could not attach session id for Hub artifact collection")
+        token = agent_session.hf_token or getattr(session, "hf_token", None)
+        if not token:
+            return
+        try:
+            start_session_artifact_collection_task(session, token=token)
+        except Exception as e:
+            logger.debug(
+                "Failed to schedule Hub artifact collection for %s: %s",
+                agent_session.session_id,
+                e,
+            )
+
     async def _clear_persisted_sandbox_metadata(self, session_id: str) -> None:
         try:
             await self._store().update_session_fields(
@@ -567,6 +591,7 @@ class SessionManager:
                     existing,
                     preload_sandbox=preload_sandbox,
                 )
+                self._start_hub_artifact_collection(existing)
                 return existing
             return None
 
@@ -588,6 +613,7 @@ class SessionManager:
                     existing,
                     preload_sandbox=preload_sandbox,
                 )
+                self._start_hub_artifact_collection(existing)
                 return existing
             return None
 
@@ -674,7 +700,9 @@ class SessionManager:
                 hf_token=hf_token,
                 hf_username=hf_username,
             )
+            self._start_hub_artifact_collection(started)
             return started
+        self._start_hub_artifact_collection(agent_session)
         if preload_sandbox:
             self._start_cpu_sandbox_preload(agent_session)
         logger.info("Restored session %s for user %s", session_id, owner or user_id)
@@ -757,6 +785,7 @@ class SessionManager:
             event_queue=event_queue,
             tool_router=tool_router,
         )
+        self._start_hub_artifact_collection(agent_session)
         await self.persist_session_snapshot(agent_session, runtime_state="idle")
         self._start_cpu_sandbox_preload(agent_session)
 

tests/unit/test_hub_artifacts.py

@@ -0,0 +1,505 @@
+import asyncio
+import logging
+from types import SimpleNamespace
+
+import pytest
+
+from agent.core import hub_artifacts
+from agent.core.hub_artifacts import (
+    ML_INTERN_TAG,
+    PROVENANCE_MARKER,
+    artifact_collection_title,
+    augment_repo_card_content,
+    build_hub_artifact_sitecustomize,
+    ensure_session_artifact_collection,
+    is_known_hub_artifact,
+    register_hub_artifact,
+    remember_hub_artifact,
+    start_session_artifact_collection_task,
+    wrap_shell_command_with_hub_artifact_bootstrap,
+)
+from agent.tools import local_tools, sandbox_tool
+from agent.tools.hf_repo_files_tool import HfRepoFilesTool
+from agent.tools.hf_repo_git_tool import HfRepoGitTool
+from agent.tools.jobs_tool import _wrap_command_with_artifact_bootstrap
+
+
+def _session() -> SimpleNamespace:
+    return SimpleNamespace(
+        session_id="session-123",
+        session_start_time="2026-05-05T10:20:30",
+    )
+
+
+def test_artifact_collection_title_uses_session_date_and_id():
+    assert (
+        artifact_collection_title(_session())
+        == "ml-intern-artifacts-2026-05-05-session-123"
+    )
+
+
+def test_artifact_collection_title_uses_short_uuid_fragment():
+    session = SimpleNamespace(
+        session_id="fadcbc77-3439-4c2b-bc52-50d7f6353af3",
+        session_start_time="2026-05-05T10:20:30",
+    )
+
+    title = artifact_collection_title(session)
+
+    assert title == "ml-intern-artifacts-2026-05-05-fadcbc77"
+    assert len(title) < 60
+
+
+def test_artifact_collection_title_still_truncates_long_non_uuid_ids():
+    session = SimpleNamespace(
+        session_id="custom-session-id-that-is-longer-than-the-hub-title-limit",
+        session_start_time="2026-05-05T10:20:30",
+    )
+
+    title = artifact_collection_title(session)
+
+    assert title.startswith("ml-intern-artifacts-2026-05-05-custom-session-id")
+    assert len(title) < 60
+
+
+def test_model_card_merges_tags_and_appends_provenance_and_usage():
+    content = """---
+license: apache-2.0
+tags:
+- text-generation
+---
+# Existing Model
+
+Existing details stay here.
+"""
+
+    updated = augment_repo_card_content(content, "alice/model", "model")
+    second_pass = augment_repo_card_content(updated, "alice/model", "model")
+
+    assert "license: apache-2.0" in updated
+    assert "- text-generation" in updated
+    assert f"- {ML_INTERN_TAG}" in updated
+    assert "# Existing Model" in updated
+    assert "Existing details stay here." in updated
+    assert PROVENANCE_MARKER in updated
+    assert "AutoModelForCausalLM" in updated
+    assert second_pass.count(PROVENANCE_MARKER) == 1
+    assert second_pass.count("AutoModelForCausalLM") == updated.count(
+        "AutoModelForCausalLM"
+    )
+
+
+def test_dataset_card_adds_load_dataset_usage():
+    updated = augment_repo_card_content("", "alice/dataset", "dataset")
+
+    assert f"- {ML_INTERN_TAG}" in updated
+    assert "# alice/dataset" in updated
+    assert "from datasets import load_dataset" in updated
+    assert 'load_dataset("alice/dataset")' in updated
+
+
+def test_existing_usage_section_is_preserved_without_duplicate_usage():
+    content = """# Existing Dataset
+
+## Usage
+
+Use the custom loader in this repository.
+"""
+
+    updated = augment_repo_card_content(content, "alice/dataset", "dataset")
+
+    assert "Use the custom loader in this repository." in updated
+    assert "from datasets import load_dataset" not in updated
+    assert PROVENANCE_MARKER in updated
+
+
+def test_space_card_gets_metadata_without_provenance_body():
+    updated = augment_repo_card_content("# Existing Space\n", "alice/space", "space")
+
+    assert f"- {ML_INTERN_TAG}" in updated
+    assert "# Existing Space" in updated
+    assert PROVENANCE_MARKER not in updated
+
+
+def test_register_hub_artifact_creates_private_collection_and_adds_item_once(
+    monkeypatch,
+):
+    session = _session()
+
+    class FakeApi:
+        token = "hf-token"
+
+        def __init__(self):
+            self.created_collections = []
+            self.collection_items = []
+            self.uploads = []
+
+        def create_collection(self, **kwargs):
+            self.created_collections.append(kwargs)
+            return SimpleNamespace(slug="alice/ml-intern-artifacts")
+
+        def add_collection_item(self, **kwargs):
+            self.collection_items.append(kwargs)
+
+        def upload_file(self, **kwargs):
+            self.uploads.append(kwargs)
+
+    api = FakeApi()
+    monkeypatch.setattr(hub_artifacts, "_read_remote_readme", lambda *_, **__: "")
+
+    assert register_hub_artifact(api, "alice/model", "model", session=session)
+    assert register_hub_artifact(api, "alice/model", "model", session=session)
+
+    assert is_known_hub_artifact(session, "alice/model", "model")
+    assert len(api.created_collections) == 1
+    assert api.created_collections[0]["title"] == artifact_collection_title(session)
+    assert api.created_collections[0]["private"] is True
+    assert len(api.collection_items) == 1
+    assert api.collection_items[0]["item_id"] == "alice/model"
+    assert api.collection_items[0]["item_type"] == "model"
+    assert api.collection_items[0]["exists_ok"] is True
+    assert len(api.uploads) == 1
+    assert b"ml-intern" in api.uploads[0]["path_or_fileobj"]
+
+
+def test_register_hub_artifact_retries_after_partial_failure(monkeypatch):
+    session = _session()
+    api = SimpleNamespace(token="hf-token")
+    card_attempts = 0
+    collection_attempts = 0
+
+    def flaky_update_repo_card(*args, **kwargs):
+        nonlocal card_attempts
+        card_attempts += 1
+        if card_attempts == 1:
+            raise RuntimeError("temporary card failure")
+
+    def add_to_collection(*args, **kwargs):
+        nonlocal collection_attempts
+        collection_attempts += 1
+
+    monkeypatch.setattr(
+        hub_artifacts,
+        "_update_repo_card",
+        flaky_update_repo_card,
+    )
+    monkeypatch.setattr(hub_artifacts, "_add_to_collection", add_to_collection)
+
+    assert not register_hub_artifact(api, "alice/model", "model", session=session)
+    assert register_hub_artifact(api, "alice/model", "model", session=session)
+    assert register_hub_artifact(api, "alice/model", "model", session=session)
+
+    assert card_attempts == 2
+    assert collection_attempts == 2
+
+
+def test_register_hub_artifact_retries_after_collection_failure(monkeypatch):
+    session = _session()
+    api = SimpleNamespace(token="hf-token")
+    card_attempts = 0
+    collection_attempts = 0
+
+    def update_repo_card(*args, **kwargs):
+        nonlocal card_attempts
+        card_attempts += 1
+
+    def flaky_add_to_collection(*args, **kwargs):
+        nonlocal collection_attempts
+        collection_attempts += 1
+        if collection_attempts == 1:
+            raise RuntimeError("temporary collection failure")
+
+    monkeypatch.setattr(hub_artifacts, "_update_repo_card", update_repo_card)
+    monkeypatch.setattr(
+        hub_artifacts,
+        "_add_to_collection",
+        flaky_add_to_collection,
+    )
+
+    assert not register_hub_artifact(api, "alice/model", "model", session=session)
+    assert register_hub_artifact(api, "alice/model", "model", session=session)
+    assert register_hub_artifact(api, "alice/model", "model", session=session)
+
+    assert card_attempts == 2
+    assert collection_attempts == 2
+
+
+def test_session_artifact_set_falls_back_when_session_rejects_attrs(caplog):
+    class SlottedSession:
+        __slots__ = ("session_id", "session_start_time")
+
+        def __init__(self):
+            self.session_id = "session-123"
+            self.session_start_time = "2026-05-05T10:20:30"
+
+    session = SlottedSession()
+
+    with caplog.at_level(logging.WARNING):
+        remember_hub_artifact(session, "alice/model", "model")
+
+    assert is_known_hub_artifact(session, "alice/model", "model")
+    assert "using process-local fallback state" in caplog.text
+
+
+@pytest.mark.asyncio
+async def test_ensure_session_artifact_collection_uses_user_token(monkeypatch):
+    session = _session()
+    calls = []
+
+    class FakeApi:
+        def __init__(self, token):
+            self.token = token
+
+    def fake_ensure_collection_slug(api, seen_session, **kwargs):
+        calls.append((api.token, seen_session, kwargs))
+        return "alice/ml-intern-artifacts"
+
+    monkeypatch.setattr(hub_artifacts, "HfApi", FakeApi)
+    monkeypatch.setattr(
+        hub_artifacts,
+        "_ensure_collection_slug",
+        fake_ensure_collection_slug,
+    )
+
+    slug = await ensure_session_artifact_collection(session, token="hf-token")
+
+    assert slug == "alice/ml-intern-artifacts"
+    assert calls == [
+        ("hf-token", session, {"token": "hf-token"}),
+    ]
+
+
+@pytest.mark.asyncio
+async def test_start_session_artifact_collection_task_dedupes(monkeypatch):
+    session = _session()
+    calls = []
+
+    async def fake_ensure_session_artifact_collection(seen_session, **kwargs):
+        calls.append((seen_session, kwargs))
+        await asyncio.sleep(0)
+        return "alice/ml-intern-artifacts"
+
+    monkeypatch.setattr(
+        hub_artifacts,
+        "ensure_session_artifact_collection",
+        fake_ensure_session_artifact_collection,
+    )
+
+    task = start_session_artifact_collection_task(session, token="hf-token")
+    second = start_session_artifact_collection_task(session, token="hf-token")
+
+    assert task is not None
+    assert second is task
+    await task
+    assert calls == [(session, {"token": "hf-token"})]
+
+
+def test_start_session_artifact_collection_task_skips_without_token():
+    assert start_session_artifact_collection_task(_session()) is None
+
+
+@pytest.mark.asyncio
+async def test_hf_repo_git_create_repo_registers_artifact(monkeypatch):
+    session = _session()
+    calls = []
+
+    class FakeApi:
+        token = "hf-token"
+
+        def create_repo(self, **kwargs):
+            self.create_kwargs = kwargs
+            return "https://huggingface.co/spaces/alice/demo"
+
+    def fake_register(api, repo_id, repo_type, **kwargs):
+        calls.append((api, repo_id, repo_type, kwargs))
+        return True
+
+    monkeypatch.setattr(
+        "agent.tools.hf_repo_git_tool.register_hub_artifact",
+        fake_register,
+    )
+    tool = HfRepoGitTool(hf_token="hf-token", session=session)
+    tool.api = FakeApi()
+
+    result = await tool._create_repo(
+        {
+            "repo_id": "alice/demo",
+            "repo_type": "space",
+            "space_sdk": "gradio",
+            "private": True,
+        }
+    )
+
+    assert result["totalResults"] == 1
+    assert calls == [
+        (
+            tool.api,
+            "alice/demo",
+            "space",
+            {"session": session, "extra_metadata": {"sdk": "gradio"}},
+        )
+    ]
+
+
+@pytest.mark.asyncio
+async def test_hf_repo_files_upload_registers_known_artifact_with_force(monkeypatch):
+    session = _session()
+    calls = []
+    uploads = []
+
+    class FakeApi:
+        token = "hf-token"
+
+        def upload_file(self, **kwargs):
+            uploads.append(kwargs)
+            return SimpleNamespace()
+
+    def fake_register(api, repo_id, repo_type, **kwargs):
+        calls.append((api, repo_id, repo_type, kwargs))
+        return True
+
+    monkeypatch.setattr(
+        "agent.tools.hf_repo_files_tool.register_hub_artifact",
+        fake_register,
+    )
+    remember_hub_artifact(session, "alice/model", "model")
+
+    tool = HfRepoFilesTool(hf_token="hf-token", session=session)
+    tool.api = FakeApi()
+
+    result = await tool._upload(
+        {
+            "repo_id": "alice/model",
+            "repo_type": "model",
+            "path": "weights.bin",
+            "content": b"weights",
+        }
+    )
+    readme_result = await tool._upload(
+        {
+            "repo_id": "alice/model",
+            "repo_type": "model",
+            "path": "README.md",
+            "content": "# Model",
+        }
+    )
+
+    assert result["totalResults"] == 1
+    assert readme_result["totalResults"] == 1
+    assert [upload["path_in_repo"] for upload in uploads] == [
+        "weights.bin",
+        "README.md",
+    ]
+    assert calls == [
+        (
+            tool.api,
+            "alice/model",
+            "model",
+            {"session": session, "force": False},
+        ),
+        (
+            tool.api,
+            "alice/model",
+            "model",
+            {"session": session, "force": True},
+        ),
+    ]
+
+
+def test_hf_jobs_artifact_bootstrap_wraps_command_without_changing_exec_target():
+    command = ["uv", "run", "train.py"]
+    wrapped = _wrap_command_with_artifact_bootstrap(command, _session())
+
+    assert wrapped[0:2] == ["/bin/sh", "-lc"]
+    assert "sitecustomize.py" in wrapped[2]
+    assert "PYTHONPATH" in wrapped[2]
+    assert "exec uv run train.py" in wrapped[2]
+    assert _wrap_command_with_artifact_bootstrap(command, None) == command
+
+
+def test_shell_bootstrap_wraps_capybara_push_to_hub_pattern():
+    command = (
+        "pip install -q datasets huggingface_hub && python -c "
+        "\"subset.push_to_hub('lewtun/Capybara-100', private=False)\""
+    )
+
+    wrapped = wrap_shell_command_with_hub_artifact_bootstrap(command, _session())
+
+    assert "sitecustomize.py" in wrapped
+    assert "PYTHONPATH" in wrapped
+    assert command in wrapped
+    assert wrap_shell_command_with_hub_artifact_bootstrap(command, None) == command
+    assert (
+        wrap_shell_command_with_hub_artifact_bootstrap(
+            command,
+            SimpleNamespace(session_start_time="2026-05-05T10:20:30"),
+        )
+        == command
+    )
+
+
+@pytest.mark.asyncio
+async def test_sandbox_bash_wraps_command_for_session_artifact_hooks():
+    calls = []
+
+    class FakeSandbox:
+        def call_tool(self, name, args):
+            calls.append((name, args))
+            return SimpleNamespace(success=True, output="ok", error="")
+
+    session = _session()
+    session.sandbox = FakeSandbox()
+
+    handler = sandbox_tool._make_tool_handler("bash")
+    output, ok = await handler({"command": "python make_dataset.py"}, session=session)
+
+    assert ok is True
+    assert output == "ok"
+    assert calls[0][0] == "bash"
+    assert "sitecustomize.py" in calls[0][1]["command"]
+    assert "python make_dataset.py" in calls[0][1]["command"]
+
+
+@pytest.mark.asyncio
+async def test_local_bash_wraps_command_for_session_artifact_hooks(monkeypatch):
+    seen = {}
+
+    def fake_run(command, **kwargs):
+        seen["command"] = command
+        seen["kwargs"] = kwargs
+        return SimpleNamespace(stdout="ok", stderr="", returncode=0)
+
+    monkeypatch.setattr(local_tools.subprocess, "run", fake_run)
+
+    output, ok = await local_tools._bash_handler(
+        {"command": "python make_dataset.py"},
+        session=_session(),
+    )
+
+    assert ok is True
+    assert output == "ok"
+    assert "sitecustomize.py" in seen["command"]
+    assert "python make_dataset.py" in seen["command"]
+
+
+def test_sitecustomize_bootstrap_is_valid_python():
+    code = build_hub_artifact_sitecustomize(_session())
+
+    compile(code, "sitecustomize.py", "exec")
+    assert "ml-intern-artifacts-2026-05-05-session-123" in code
+
+
+def test_sitecustomize_bootstrap_reuses_existing_collection_slug():
+    session = _session()
+    setattr(
+        session,
+        hub_artifacts._COLLECTION_SLUG_ATTR,
+        "alice/ml-intern-artifacts-2026-05-05-session-123",
+    )
+
+    code = build_hub_artifact_sitecustomize(session)
+
+    compile(code, "sitecustomize.py", "exec")
+    assert (
+        "collection_slug = 'alice/ml-intern-artifacts-2026-05-05-session-123'" in code
+    )

tests/unit/test_session_manager_persistence.py

@@ -430,6 +430,32 @@ async def test_create_session_schedules_cpu_sandbox_preload():
         await _cancel_runtime_tasks(manager)
 
 
+@pytest.mark.asyncio
+async def test_create_session_starts_hub_artifact_collection(monkeypatch):
+    manager = _manager_with_store(NoopSessionStore())
+    manager.enable_hub_artifact_collections = True
+    stop = _install_fake_runtime(manager)
+    started: list[tuple[str, str]] = []
+
+    def fake_start_session_artifact_collection_task(session, **kwargs):
+        started.append((session.session_id, kwargs["token"]))
+        return None
+
+    monkeypatch.setattr(
+        "session_manager.start_session_artifact_collection_task",
+        fake_start_session_artifact_collection_task,
+    )
+    manager._start_cpu_sandbox_preload = lambda _: None  # type: ignore[method-assign]
+
+    try:
+        session_id = await manager.create_session(user_id="owner", hf_token="token")
+
+        assert started == [(session_id, "token")]
+    finally:
+        stop.set()
+        await _cancel_runtime_tasks(manager)
+
+
 @pytest.mark.asyncio
 async def test_lazy_restore_schedules_cpu_sandbox_preload():
     manager = _manager_with_store(RestoreStore())
@@ -454,6 +480,37 @@ async def test_lazy_restore_schedules_cpu_sandbox_preload():
         await _cancel_runtime_tasks(manager)
 
 
+@pytest.mark.asyncio
+async def test_lazy_restore_starts_hub_artifact_collection(monkeypatch):
+    manager = _manager_with_store(RestoreStore())
+    manager.enable_hub_artifact_collections = True
+    stop = _install_fake_runtime(manager)
+    started: list[tuple[str, str]] = []
+
+    def fake_start_session_artifact_collection_task(session, **kwargs):
+        started.append((session.session_id, kwargs["token"]))
+        return None
+
+    monkeypatch.setattr(
+        "session_manager.start_session_artifact_collection_task",
+        fake_start_session_artifact_collection_task,
+    )
+    manager._start_cpu_sandbox_preload = lambda _: None  # type: ignore[method-assign]
+
+    try:
+        restored = await manager.ensure_session_loaded(
+            "persisted-session",
+            user_id="owner",
+            hf_token="token",
+        )
+
+        assert restored is not None
+        assert started == [("persisted-session", "token")]
+    finally:
+        stop.set()
+        await _cancel_runtime_tasks(manager)
+
+
 @pytest.mark.asyncio
 async def test_lazy_restore_deletes_persisted_sandbox_before_preload(monkeypatch):
     deleted: list[tuple[str, str, str]] = []