Teknium
364 lines
14 KiB
Python
364 lines
14 KiB
Python
"""On-disk pet store — install / list / resolve pets.
|
|
|
|
Pets live under ``get_hermes_home()/pets/<slug>/`` (profile-scoped; NOT petdex's
|
|
``~/.codex/pets``, which its CLI owns): ``pet.json`` ({id, displayName,
|
|
description, spritesheetPath}) plus ``spritesheet.webp`` (or .png). The active
|
|
pet comes from the caller-supplied ``display.pet.slug`` (no config loader here).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import contextlib
|
|
import io
|
|
import itertools
|
|
import json
|
|
import logging
|
|
import re
|
|
import shutil
|
|
import zipfile
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from urllib.parse import urlparse
|
|
|
|
from hermes_constants import get_hermes_home
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_DOWNLOAD_TIMEOUT = 60.0
|
|
_HTTP_HEADERS = {"User-Agent": "hermes-agent-petdex"}
|
|
_THUMB_FRAME_W = 192
|
|
_THUMB_FRAME_H = 208
|
|
_THUMB_W = 96 # rendered ~40px; 2x+ keeps it crisp on HiDPI
|
|
|
|
|
|
class PetStoreError(RuntimeError):
|
|
"""Raised on install/IO failures."""
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class InstalledPet:
|
|
"""A pet present on disk."""
|
|
|
|
slug: str
|
|
display_name: str
|
|
description: str
|
|
directory: Path
|
|
spritesheet: Path
|
|
created_by: str = "" # "generator" for pets hatched locally; "" for petdex installs
|
|
|
|
@property
|
|
def exists(self) -> bool:
|
|
return self.spritesheet.is_file()
|
|
|
|
@property
|
|
def generated(self) -> bool:
|
|
return self.created_by == "generator"
|
|
|
|
|
|
def pets_dir() -> Path:
|
|
"""Return the profile-scoped pets directory (created on demand)."""
|
|
path = get_hermes_home() / "pets"
|
|
path.mkdir(parents=True, exist_ok=True)
|
|
return path
|
|
|
|
|
|
def _thumb_path(slug: str) -> Path:
|
|
"""Cached thumbnail for *slug* (lives OUTSIDE the pet dir, under ``pets/.thumbs/``)."""
|
|
path = pets_dir() / ".thumbs"
|
|
path.mkdir(parents=True, exist_ok=True)
|
|
return path / f"{slug}.png"
|
|
|
|
|
|
def _read_pet_json(directory: Path) -> dict:
|
|
pet_json = directory / "pet.json"
|
|
try:
|
|
return json.loads(pet_json.read_text(encoding="utf-8")) if pet_json.is_file() else {}
|
|
except (OSError, ValueError) as exc:
|
|
logger.debug("unreadable pet.json in %s: %s", directory, exc)
|
|
return {}
|
|
|
|
|
|
def _write_pet_json(directory: Path, meta: dict) -> None:
|
|
(directory / "pet.json").write_text(json.dumps(meta, indent=2), encoding="utf-8")
|
|
|
|
|
|
def _resolve_spritesheet(directory: Path, meta: dict) -> Path:
|
|
"""Honor ``spritesheetPath``, else probe conventional names (incl. petdex R2's ``sprite.webp``)."""
|
|
declared = str(meta.get("spritesheetPath", "") or "").strip()
|
|
names = ([declared] if declared else []) + ["spritesheet.webp", "spritesheet.png", "sprite.webp", "sprite.png"]
|
|
return next((directory / n for n in names if (directory / n).is_file()), directory / "spritesheet.webp") # stable default
|
|
|
|
|
|
def _safe_slug(slug: str) -> str:
|
|
"""Normalize a slug to one bare path segment so ``pets_dir()/<slug>`` can never escape the pets directory."""
|
|
segment = Path(str(slug).strip()).name
|
|
return "" if segment in ("", ".", "..") else segment
|
|
|
|
|
|
def load_pet(slug: str) -> InstalledPet | None:
|
|
"""Return the :class:`InstalledPet` for *slug*, or ``None`` if absent."""
|
|
slug = _safe_slug(slug)
|
|
directory = pets_dir() / slug
|
|
if not slug or not directory.is_dir():
|
|
return None
|
|
meta = _read_pet_json(directory)
|
|
return InstalledPet(
|
|
slug=slug,
|
|
display_name=str(meta.get("displayName", "") or slug),
|
|
description=str(meta.get("description", "") or ""),
|
|
directory=directory,
|
|
spritesheet=_resolve_spritesheet(directory, meta),
|
|
created_by=str(meta.get("createdBy", "") or ""),
|
|
)
|
|
|
|
|
|
def _usable_pet(slug: str, error: str = "") -> InstalledPet | None:
|
|
""":func:`load_pet` result only when its spritesheet is on disk; raises :class:`PetStoreError` *error* if given."""
|
|
pet = load_pet(slug)
|
|
if pet and pet.exists:
|
|
return pet
|
|
if error:
|
|
raise PetStoreError(error)
|
|
return None
|
|
|
|
|
|
def installed_pets() -> list[InstalledPet]:
|
|
"""Return every installed pet (dirs containing a usable spritesheet)."""
|
|
pets = (_usable_pet(child.name) for child in sorted(pets_dir().iterdir()) if child.is_dir())
|
|
return [pet for pet in pets if pet]
|
|
|
|
|
|
def resolve_active_pet(configured_slug: str | None = None) -> InstalledPet | None:
|
|
"""The configured slug (``display.pet.slug``) if installed, else the first pet alphabetically."""
|
|
if configured_slug and (pet := _usable_pet(configured_slug.strip())):
|
|
return pet
|
|
pets = installed_pets()
|
|
return pets[0] if pets else None
|
|
|
|
|
|
def install_pet(slug: str, *, force: bool = False, timeout: float = _DOWNLOAD_TIMEOUT) -> InstalledPet:
|
|
"""Download *slug* from the manifest; idempotent unless *force*. Raises :class:`PetStoreError` / ``ManifestError``."""
|
|
from agent.pet.manifest import find_entry
|
|
|
|
slug = _safe_slug(slug)
|
|
if not slug:
|
|
raise PetStoreError("invalid pet slug")
|
|
if not force and (existing := _usable_pet(slug)):
|
|
return existing
|
|
entry = find_entry(slug, timeout=timeout)
|
|
if entry is None:
|
|
raise PetStoreError(f"pet '{slug}' is not in the petdex manifest")
|
|
# Host-pin asset URLs so a compromised/spoofed manifest can't redirect the
|
|
# download to an arbitrary host (matches thumbnail_png).
|
|
if not _is_petdex_host(entry.spritesheet_url):
|
|
raise PetStoreError(f"refusing non-petdex spritesheet host for '{slug}'")
|
|
|
|
directory = pets_dir() / slug
|
|
directory.mkdir(parents=True, exist_ok=True)
|
|
sprite_ext = ".png" if entry.spritesheet_url.lower().split("?")[0].endswith(".png") else ".webp"
|
|
sprite_path = directory / f"spritesheet{sprite_ext}"
|
|
_download(entry.spritesheet_url, sprite_path, timeout=timeout)
|
|
|
|
# Prefer the upstream pet.json; else synthesize one so the layout is self-describing.
|
|
meta: dict = {}
|
|
if entry.pet_json_url and _is_petdex_host(entry.pet_json_url):
|
|
try:
|
|
data = _http_get(entry.pet_json_url, timeout).json()
|
|
meta = data if isinstance(data, dict) else {}
|
|
except Exception as exc: # noqa: BLE001 - non-fatal, fall back below
|
|
logger.debug("pet.json fetch failed for %s: %s", slug, exc)
|
|
meta = meta or {"id": slug, "displayName": entry.display_name, "description": ""}
|
|
meta["spritesheetPath"] = sprite_path.name # key order matters: pet.json is written verbatim
|
|
meta.setdefault("id", slug)
|
|
meta.setdefault("displayName", entry.display_name)
|
|
_write_pet_json(directory, meta)
|
|
return _usable_pet(slug, f"install of '{slug}' did not produce a spritesheet")
|
|
|
|
|
|
def slugify(name: str) -> str:
|
|
"""Lowercase, hyphenate, and strip a display name into a filesystem slug."""
|
|
slug = re.sub(r"[^a-z0-9]+", "-", (name or "").strip().lower()).strip("-")
|
|
return slug or "pet"
|
|
|
|
|
|
def unique_slug(name: str) -> str:
|
|
"""A :func:`slugify` result that doesn't collide with an existing pet dir."""
|
|
base = slugify(name)
|
|
candidates = (base if i == 1 else f"{base}-{i}" for i in itertools.count(1))
|
|
return next(slug for slug in candidates if not (pets_dir() / slug).exists())
|
|
|
|
|
|
def _write_spritesheet(source, dest: Path) -> None:
|
|
"""Write *source* (PIL image, bytes, or path) as a lossless WebP at *dest*."""
|
|
if isinstance(source, (bytes, bytearray)):
|
|
dest.write_bytes(bytes(source))
|
|
return
|
|
from agent.pet.generate.atlas import _load_rgba
|
|
|
|
_load_rgba(source).save(dest, format="WEBP", lossless=True, quality=100, method=6, exact=True)
|
|
|
|
|
|
def register_local_pet(spritesheet, *, slug: str, display_name: str = "", description: str = "") -> InstalledPet:
|
|
"""Write a locally-generated pet (PIL image, WebP/PNG bytes, or path) into the store.
|
|
|
|
Appears in :func:`installed_pets` immediately; :func:`install_pet` returns on-disk
|
|
pets before consulting the manifest, so no manifest entry is needed.
|
|
"""
|
|
slug = slugify(slug)
|
|
directory = pets_dir() / slug
|
|
directory.mkdir(parents=True, exist_ok=True)
|
|
sprite_path = directory / "spritesheet.webp"
|
|
try:
|
|
_write_spritesheet(spritesheet, sprite_path)
|
|
except Exception as exc: # noqa: BLE001 - normalize to one error type
|
|
raise PetStoreError(f"could not write spritesheet for '{slug}': {exc}") from exc
|
|
meta = {"id": slug, "displayName": display_name or slug, "description": description or "", "spritesheetPath": sprite_path.name}
|
|
_write_pet_json(directory, {**meta, "createdBy": "generator"})
|
|
return _usable_pet(slug, f"register of generated pet '{slug}' did not produce a spritesheet")
|
|
|
|
|
|
def export_pet(slug: str) -> tuple[str, bytes]:
|
|
"""Zip an installed pet's folder → ``(filename, bytes)``; dotfiles (thumbs, backups) skipped."""
|
|
root = pets_dir()
|
|
directory = root / slug.strip()
|
|
# Traversal guard: the target must be a direct child of pets_dir.
|
|
if directory.resolve().parent != root.resolve() or not directory.is_dir():
|
|
raise PetStoreError(f"pet '{slug}' is not installed")
|
|
buf = io.BytesIO()
|
|
with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as archive:
|
|
for path in sorted(directory.iterdir()):
|
|
if path.is_file() and not path.name.startswith("."):
|
|
archive.write(path, f"{directory.name}/{path.name}")
|
|
return f"{directory.name}.zip", buf.getvalue()
|
|
|
|
|
|
def _is_petdex_host(url: str) -> bool:
|
|
"""True only for petdex.dev hosts — bounds server-side fetch (anti-SSRF)."""
|
|
try:
|
|
host = (urlparse(url).hostname or "").lower()
|
|
except ValueError:
|
|
return False
|
|
return host == "petdex.dev" or host.endswith(".petdex.dev")
|
|
|
|
|
|
def _thumb_source_bytes(slug: str, source_url: str, timeout: float) -> bytes | None:
|
|
"""Spritesheet bytes for a thumbnail: the installed sheet, else a petdex-hosted *source_url*."""
|
|
pet = _usable_pet(slug)
|
|
if pet:
|
|
with contextlib.suppress(OSError):
|
|
return pet.spritesheet.read_bytes()
|
|
if source_url and _is_petdex_host(source_url):
|
|
try:
|
|
return _http_get(source_url, timeout).content
|
|
except Exception as exc: # noqa: BLE001 - cosmetic, degrade to placeholder
|
|
logger.debug("thumb fetch failed for %s: %s", slug, exc)
|
|
return None
|
|
|
|
|
|
def thumbnail_png(slug: str, *, source_url: str = "", timeout: float = 30.0) -> bytes | None:
|
|
"""Small idle-frame (top-left cell) PNG for *slug*, cached on disk; ``None`` on any failure.
|
|
|
|
Source: the installed sheet, else *source_url* only when petdex-hosted (the gateway
|
|
never fetches arbitrary client URLs). Server-side so it rides the authenticated
|
|
gateway as a same-origin data URL, sidestepping CSP/hotlink limits.
|
|
"""
|
|
slug = slug.strip()
|
|
if not slug:
|
|
return None
|
|
cache = _thumb_path(slug)
|
|
if cache.is_file():
|
|
with contextlib.suppress(OSError):
|
|
return cache.read_bytes()
|
|
sheet_bytes = _thumb_source_bytes(slug, source_url, timeout)
|
|
if not sheet_bytes:
|
|
return None
|
|
try:
|
|
from PIL import Image
|
|
|
|
with Image.open(io.BytesIO(sheet_bytes)) as im:
|
|
frame = im.convert("RGBA").crop((0, 0, min(_THUMB_FRAME_W, im.width), min(_THUMB_FRAME_H, im.height)))
|
|
frame = frame.resize((_THUMB_W, round(_THUMB_W * _THUMB_FRAME_H / _THUMB_FRAME_W)), Image.NEAREST)
|
|
buf = io.BytesIO()
|
|
frame.save(buf, format="PNG")
|
|
except Exception as exc: # noqa: BLE001
|
|
logger.debug("thumb crop failed for %s: %s", slug, exc)
|
|
return None
|
|
with contextlib.suppress(OSError):
|
|
cache.write_bytes(buf.getvalue())
|
|
return buf.getvalue()
|
|
|
|
|
|
def remove_pet(slug: str) -> bool:
|
|
"""Delete an installed pet directory. Returns True if anything was removed."""
|
|
slug = _safe_slug(slug)
|
|
if not slug:
|
|
return False
|
|
# Drop the cached thumb too or a later pet reusing this slug shows the stale one.
|
|
with contextlib.suppress(OSError):
|
|
_thumb_path(slug).unlink(missing_ok=True)
|
|
directory = pets_dir() / slug
|
|
if not directory.is_dir():
|
|
return False
|
|
shutil.rmtree(directory, ignore_errors=True)
|
|
return not directory.exists()
|
|
|
|
|
|
def rename_pet(slug: str, display_name: str) -> str | None:
|
|
"""Rename a pet's ``displayName`` AND realign its slug/dir to match.
|
|
|
|
Generated pets hatch under a provisional slug; naming makes that the identity.
|
|
The dir (and thumb) moves to ``slugify(name)`` when that's a free, different
|
|
slug. Returns the resulting slug, or ``None`` on failure.
|
|
"""
|
|
slug = _safe_slug(slug)
|
|
display_name = (display_name or "").strip()
|
|
directory = pets_dir() / slug
|
|
if not slug or not display_name or not (directory / "pet.json").is_file():
|
|
return None
|
|
meta = _read_pet_json(directory)
|
|
meta = meta if isinstance(meta, dict) else {}
|
|
meta["displayName"] = display_name
|
|
new_slug = slug
|
|
desired = slugify(display_name)
|
|
if desired and desired != slug and not (pets_dir() / desired).exists():
|
|
try:
|
|
directory.rename(pets_dir() / desired)
|
|
except OSError:
|
|
pass # keep the provisional slug if the move fails
|
|
else:
|
|
with contextlib.suppress(OSError):
|
|
_thumb_path(slug).rename(_thumb_path(desired))
|
|
directory = pets_dir() / desired
|
|
new_slug = meta["id"] = desired
|
|
|
|
try:
|
|
_write_pet_json(directory, meta)
|
|
except OSError:
|
|
return None
|
|
return new_slug
|
|
|
|
|
|
def _http_get(url: str, timeout: float):
|
|
"""GET *url* with the petdex UA, following redirects; raises on HTTP errors."""
|
|
import httpx
|
|
|
|
resp = httpx.get(url, timeout=timeout, follow_redirects=True, headers=_HTTP_HEADERS)
|
|
resp.raise_for_status()
|
|
return resp
|
|
|
|
|
|
def _download(url: str, dest: Path, *, timeout: float) -> None:
|
|
"""Stream *url* to *dest* via a ``.part`` temp file so a failed download never leaves a truncated sheet."""
|
|
import httpx
|
|
|
|
try:
|
|
with httpx.stream("GET", url, timeout=timeout, follow_redirects=True, headers=_HTTP_HEADERS) as resp:
|
|
resp.raise_for_status()
|
|
tmp = dest.with_suffix(dest.suffix + ".part")
|
|
with tmp.open("wb") as fh:
|
|
for chunk in resp.iter_bytes():
|
|
fh.write(chunk)
|
|
tmp.replace(dest)
|
|
except Exception as exc: # noqa: BLE001
|
|
raise PetStoreError(f"download failed for {url}: {exc}") from exc
|