diff --git a/tests/tui_gateway/test_pet_payload_extract.py b/tests/tui_gateway/test_pet_payload_extract.py new file mode 100644 index 000000000000..9191ea5e97f9 --- /dev/null +++ b/tests/tui_gateway/test_pet_payload_extract.py @@ -0,0 +1,435 @@ +"""R3-S1 extraction regression: pet payload moved to tui_gateway.pet_payload. + +Covers the consensus seam contract (epic #78647, target #78630): +- identity-preserving re-export of all 23 members (server namespace binds the + exact same objects the new module owns), +- handler + watcher liveness through the re-export binding, +- no import cycle in any order, +- byte-verbatim golden-sha regression for the moved span, +- aggressive unit coverage of the moved cluster (payload shape, cache cap, + clone-on-read, cancel lifecycle, reference-image validation, downscale, + gen sweep staleness, config-scale fail-open). +""" + +from __future__ import annotations + +import base64 +import hashlib +import pathlib +import subprocess +import sys +import time +from types import SimpleNamespace + +import pytest + +pytest.importorskip("PIL") +from PIL import Image # noqa: E402 + +from tui_gateway import pet_payload, server # noqa: E402 + +GOLDEN_SHA = "99120f354c2675612f796c4d4fb19477f7a9983b1135d5dd0c17471ea3aa59aa" + +MEMBERS = [ + "_pet_frame_counts", "_pet_payload_cache_lock", "_pet_payload_cache", + "_pet_sheet_revision", "_pet_payload_cache_key", "_clone_pet_payload", + "_pet_row_frame_counts", "_pet_config_scale", "_pet_sprite_payload", + "_pet_active_selection", "_pet_state_rows", "_pet_gen_root", "_pet_gen_sweep", + "_pet_png_data_uri", "_pet_cancel_lock", "_pet_cancelled", + "_PET_REFERENCE_MIME_EXT", "_PET_REFERENCE_MAX_BYTES", + "_pet_reference_images_from_data_url", "_pet_cancel_arm", "_pet_cancel_request", + "_pet_is_cancelled", "_pet_cancel_release", +] + +WIRE_KEYS = { + "slug", "displayName", "mime", "spritesheetBase64", "spritesheetRevision", + "frameW", "frameH", "framesPerState", "framesByState", "framesByRow", + "loopMs", "scale", "stateRows", +} + + +@pytest.fixture(autouse=True) +def _reset_pet_state(): + """Pet caches/cancel set now live in pet_payload (post-extraction).""" + pet_payload._pet_payload_cache.clear() + pet_payload._pet_cancelled.clear() + yield + + +def _png(path, size=(64, 64), color=(200, 80, 80, 255)): + Image.new("RGBA", size, color).save(path) + return pathlib.Path(path) + + +def _fake_pet(path, slug="test-pet", display_name="Test Pet"): + return SimpleNamespace( + slug=slug, + display_name=display_name, + spritesheet=pathlib.Path(path), + exists=True, + ) + + +# --------------------------------------------------------------------------- +# Seam identity: re-export binds the exact same objects (23 members) +# --------------------------------------------------------------------------- + +def test_reexport_identity_all_23_members(): + for name in MEMBERS: + assert name in pet_payload.__dict__, f"{name} missing from pet_payload" + assert name in vars(server), f"{name} missing from server namespace" + assert getattr(server, name) is getattr(pet_payload, name), ( + f"{name} re-export is not identity-preserving" + ) + + +def test_pet_payload_defines_exactly_the_23_members(): + import ast + + src = pathlib.Path(pet_payload.__file__).read_text(encoding="utf-8") + tree = ast.parse(src) + defined = set() + for node in ast.walk(tree): + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): + defined.add(node.name) + elif isinstance(node, (ast.Assign, ast.AnnAssign)): + targets = node.targets if isinstance(node, ast.Assign) else [node.target] + for t in targets: + if isinstance(t, ast.Name): + defined.add(t.id) + assert set(MEMBERS) <= defined, f"missing definitions: {set(MEMBERS) - defined}" + # no stray pet-cluster names beyond the adjudicated 23 + stray = {n for n in defined if n.startswith(("_pet", "_PET", "_clone"))} - set(MEMBERS) + assert not stray, f"unexpected pet-cluster definitions: {stray}" + + +def test_golden_sha_span_still_verbatim_in_module(): + src = pathlib.Path(pet_payload.__file__).read_text(encoding="utf-8") + marker = "def _pet_frame_counts" + assert marker in src + span = src[src.index(marker):] + digest = hashlib.sha256(span.encode("utf-8")).hexdigest() + assert digest == GOLDEN_SHA, f"moved span drifted: {digest}" + + +def test_lock_and_cancel_state_same_object_across_boundary(): + # The lock objects must be the SAME objects across the module boundary so + # concurrency semantics survive (consensus section 4.3). + assert server._pet_payload_cache_lock is pet_payload._pet_payload_cache_lock + assert server._pet_cancel_lock is pet_payload._pet_cancel_lock + assert server._pet_payload_cache is pet_payload._pet_payload_cache + assert server._pet_cancelled is pet_payload._pet_cancelled + + +# --------------------------------------------------------------------------- +# Import cycle: all orders clean +# --------------------------------------------------------------------------- + +@pytest.mark.parametrize("order", [ + "import tui_gateway.pet_payload; import tui_gateway.server; import tui_gateway.methods_session", + "import tui_gateway.methods_session; import tui_gateway.server; import tui_gateway.pet_payload", + "import tui_gateway.server; import tui_gateway.pet_payload; import tui_gateway.methods_session", +]) +def test_no_import_cycle_any_order(order): + code = f"{order}; print('ok')" + proc = subprocess.run( + [sys.executable, "-c", code], + capture_output=True, text=True, cwd=str(pathlib.Path(__file__).parents[2]), + timeout=120, + ) + combined = proc.stdout + proc.stderr + assert proc.returncode == 0, combined + assert "ok" in combined + + +# --------------------------------------------------------------------------- +# Handler + watcher liveness through the re-export binding +# --------------------------------------------------------------------------- + +def test_pet_info_handler_live_through_reexport(): + # Rebound handler's bare _pet_active_selection/_pet_sprite_payload resolve + # through server.py's namespace (HandlerRegistry.install rebinds __globals__). + resp = server._methods["pet.info"]("r1", {}) + assert resp["id"] == "r1" + assert resp["jsonrpc"] == "2.0" + assert isinstance(resp["result"], dict) + assert resp["result"].get("enabled") is False # fail-open, no pet configured + + +def test_pet_changed_watcher_patch_liveness(monkeypatch, tmp_path): + # test_change_watcher monkeypatches server._pet_active_selection; the + # re-export binding must keep intercepting (consensus section 3 seam #3). + sheet = _png(tmp_path / "sheet.png") + fake = lambda: (True, _fake_pet(sheet, slug="patch-pet", display_name="Patch Pet"), 0.5) + monkeypatch.setattr(server, "_pet_active_selection", fake) + payload = server._pet_changed_payload() + assert payload == { + "enabled": True, + "slug": "patch-pet", + "displayName": "Patch Pet", + "scale": 0.5, + "spritesheetRevision": f"{sheet.stat().st_mtime_ns}:{sheet.stat().st_size}", + } + + +def test_pet_sig_watcher_second_hop(monkeypatch, tmp_path): + # _pet_sig (defined in server.py, R1) reads _pet_active_selection + + # _pet_sheet_revision through the server namespace. + sheet = _png(tmp_path / "sheet.png") + monkeypatch.setattr(server, "_load_cfg", lambda: {"display": {"pet": {"enabled": True}}}) + monkeypatch.setattr( + server, "_pet_active_selection", + lambda: (True, _fake_pet(sheet, slug="sig-pet"), 0.75), + ) + assert server._pet_sig() == ( + "sig-pet", + f"{sheet.stat().st_mtime_ns}:{sheet.stat().st_size}", + 0.75, + ) + + +# --------------------------------------------------------------------------- +# Payload shape, cache, clone-on-read (wire contract, consensus section 2.5) +# --------------------------------------------------------------------------- + +def test_sprite_payload_wire_shape(tmp_path): + sheet = _png(tmp_path / "sheet.png") + payload = pet_payload._pet_sprite_payload(_fake_pet(sheet), scale=1.0) + assert set(payload) == WIRE_KEYS + assert payload["slug"] == "test-pet" + assert payload["displayName"] == "Test Pet" + assert payload["mime"] == "image/png" + assert base64.b64decode(payload["spritesheetBase64"]) == sheet.read_bytes() + assert payload["scale"] == 1.0 + assert isinstance(payload["framesByState"], dict) + assert isinstance(payload["stateRows"], list) + + +def test_cache_hit_and_clone_on_read_no_alias(tmp_path): + sheet = _png(tmp_path / "sheet.png") + pet = _fake_pet(sheet) + first = pet_payload._pet_sprite_payload(pet, scale=1.0) + assert len(pet_payload._pet_payload_cache) == 1 + + # Mutating the returned clone must not corrupt the cache entry. + first["slug"] = "MUTATED" + first["framesByState"]["hacked"] = 999 + cached = next(iter(pet_payload._pet_payload_cache.values())) + assert cached["slug"] == "test-pet" + assert "hacked" not in cached["framesByState"] + + # Second call: distinct object, same wire values (served from cache). + second = pet_payload._pet_sprite_payload(pet, scale=1.0) + assert second is not first + assert second["slug"] == "test-pet" + assert second["spritesheetBase64"] == base64.standard_b64encode( + sheet.read_bytes() + ).decode("ascii") + assert len(pet_payload._pet_payload_cache) == 1 + + +def test_cache_key_change_busts_cache(tmp_path): + sheet = _png(tmp_path / "sheet.png") + pet = _fake_pet(sheet) + pet_payload._pet_sprite_payload(pet, scale=1.0) + assert len(pet_payload._pet_payload_cache) == 1 + # mtime bump -> new key -> new entry (cap still enforced) + time.sleep(0.02) + os_utime_backdate = time.time() + 5 # future mtime, definitely different + import os + os.utime(sheet, (os_utime_backdate, os_utime_backdate)) + pet_payload._pet_sprite_payload(pet, scale=1.0) + assert len(pet_payload._pet_payload_cache) == 2 + + +def test_cache_cap_stays_at_8(tmp_path): + for i in range(10): + sheet = _png(tmp_path / f"sheet{i}.png", size=(32 + i, 32 + i)) + pet_payload._pet_sprite_payload(_fake_pet(sheet, slug=f"pet-{i}"), scale=1.0) + assert len(pet_payload._pet_payload_cache) <= 8 + + +def test_sprite_payload_scale_is_part_of_key(tmp_path): + sheet = _png(tmp_path / "sheet.png") + pet = _fake_pet(sheet) + pet_payload._pet_sprite_payload(pet, scale=1.0) + pet_payload._pet_sprite_payload(pet, scale=2.0) + assert len(pet_payload._pet_payload_cache) == 2 + + +# --------------------------------------------------------------------------- +# Cache key / revision helpers +# --------------------------------------------------------------------------- + +def test_payload_cache_key_shape(tmp_path): + sheet = _png(tmp_path / "sheet.png") + key = pet_payload._pet_payload_cache_key(_fake_pet(sheet, slug="k", display_name="K"), scale=1.25) + assert key == ( + str(sheet), sheet.stat().st_mtime_ns, sheet.stat().st_size, + "k", "K", round(1.25, 4), + ) + + +def test_payload_cache_key_missing_file_returns_none(tmp_path): + missing = tmp_path / "nope.png" + assert pet_payload._pet_payload_cache_key(_fake_pet(missing), scale=1.0) is None + + +def test_sheet_revision_shape(tmp_path): + sheet = _png(tmp_path / "sheet.png") + assert pet_payload._pet_sheet_revision(sheet) == f"{sheet.stat().st_mtime_ns}:{sheet.stat().st_size}" + + +def test_sheet_revision_fail_open(tmp_path): + assert pet_payload._pet_sheet_revision(tmp_path / "missing.png") == "0:0" + + +# --------------------------------------------------------------------------- +# Reference-image data URL validation +# --------------------------------------------------------------------------- + +def _data_url(mime, raw): + return f"data:image/{mime};base64," + base64.b64encode(raw).decode("ascii") + + +def test_reference_images_valid_png(tmp_path): + raw = _png(tmp_path / "src.png").read_bytes() + out = pet_payload._pet_reference_images_from_data_url(_data_url("png", raw), tmp_path) + assert out == [tmp_path / "reference.png"] + assert (tmp_path / "reference.png").read_bytes() == raw + + +def test_reference_images_valid_jpeg(tmp_path): + raw = b"\xff\xd8\xff\xe0fakejpeg" + out = pet_payload._pet_reference_images_from_data_url(_data_url("jpeg", raw), tmp_path) + assert out == [tmp_path / "reference.jpg"] + + +def test_reference_images_mime_whitelist(tmp_path): + with pytest.raises(ValueError, match="unsupported reference image type"): + pet_payload._pet_reference_images_from_data_url(_data_url("bmp", b"x" * 16), tmp_path) + + +def test_reference_images_invalid_format(tmp_path): + with pytest.raises(ValueError, match="invalid reference image format"): + pet_payload._pet_reference_images_from_data_url("not-a-data-url", tmp_path) + + +def test_reference_images_size_cap(monkeypatch, tmp_path): + monkeypatch.setattr(pet_payload, "_PET_REFERENCE_MAX_BYTES", 10) + with pytest.raises(ValueError, match="reference image too large"): + pet_payload._pet_reference_images_from_data_url(_data_url("png", b"x" * 64), tmp_path) + + +def test_reference_images_invalid_base64(tmp_path): + with pytest.raises(ValueError, match="invalid reference image data"): + pet_payload._pet_reference_images_from_data_url( + "data:image/png;base64,!!!not-base64!!!", tmp_path + ) + + +# --------------------------------------------------------------------------- +# Cancel token lifecycle +# --------------------------------------------------------------------------- + +def test_cancel_lifecycle(): + assert pet_payload._pet_is_cancelled("tok") is False + pet_payload._pet_cancel_request("tok") + assert pet_payload._pet_is_cancelled("tok") is True + pet_payload._pet_cancel_arm("tok") + assert pet_payload._pet_is_cancelled("tok") is False + pet_payload._pet_cancel_request("tok") + pet_payload._pet_cancel_release("tok") + assert pet_payload._pet_is_cancelled("tok") is False + + +def test_cancel_tokens_isolated(): + pet_payload._pet_cancel_request("a") + assert pet_payload._pet_is_cancelled("a") is True + assert pet_payload._pet_is_cancelled("b") is False + pet_payload._pet_cancel_release("b") # release of unset token is a no-op + assert pet_payload._pet_is_cancelled("a") is True + + +# --------------------------------------------------------------------------- +# PNG data-URI downscale +# --------------------------------------------------------------------------- + +def test_png_data_uri_downscales(tmp_path): + src = _png(tmp_path / "big.png", size=(400, 400)) + uri = pet_payload._pet_png_data_uri(src, max_px=64) + assert uri.startswith("data:image/png;base64,") + decoded = base64.b64decode(uri.split(",", 1)[1]) + assert decoded[:8] == b"\x89PNG\r\n\x1a\n" + with Image.open(__import__("io").BytesIO(decoded)) as img: + assert max(img.size) <= 64 + + +# --------------------------------------------------------------------------- +# Gen sweep staleness +# --------------------------------------------------------------------------- + +def test_gen_sweep_removes_only_stale(tmp_path): + root = tmp_path / "pet-gen" + root.mkdir() + old = root / "old" + fresh = root / "fresh" + old.mkdir() + fresh.mkdir() + backdate = time.time() - 7200 + import os + os.utime(old, (backdate, backdate)) + pet_payload._pet_gen_sweep(root, max_age_s=3600.0) + assert not old.exists() + assert fresh.exists() + + +def test_gen_sweep_missing_root_no_raise(tmp_path): + pet_payload._pet_gen_sweep(tmp_path / "missing", max_age_s=1.0) # no exception + + +# --------------------------------------------------------------------------- +# Config scale fail-open +# --------------------------------------------------------------------------- + +def test_config_scale_reads_display_pet_scale(monkeypatch): + import hermes_cli.config + monkeypatch.setattr( + hermes_cli.config, "load_config", + lambda: {"display": {"pet": {"scale": 2.5}}}, + ) + assert pet_payload._pet_config_scale() == 2.5 + + +def test_config_scale_fail_open(monkeypatch): + import hermes_cli.config + from agent.pet import constants + + def boom(): + raise RuntimeError("config broken") + + monkeypatch.setattr(hermes_cli.config, "load_config", boom) + assert pet_payload._pet_config_scale() == constants.DEFAULT_SCALE + + +def test_config_scale_missing_keys_fail_open(monkeypatch): + import hermes_cli.config + from agent.pet import constants + monkeypatch.setattr(hermes_cli.config, "load_config", lambda: {}) + assert pet_payload._pet_config_scale() == constants.DEFAULT_SCALE + + +# --------------------------------------------------------------------------- +# Fail-open decode helpers +# --------------------------------------------------------------------------- + +def test_frame_counts_fail_open_on_garbage(tmp_path): + garbage = tmp_path / "sheet.png" + garbage.write_bytes(b"not an image") + # Fail-open contract: decode hiccups degrade to a dict (never raise). + counts = pet_payload._pet_frame_counts(garbage) + assert isinstance(counts, dict) + rows = pet_payload._pet_row_frame_counts(garbage) + assert isinstance(rows, dict) + states = pet_payload._pet_state_rows(garbage) + assert isinstance(states, list) diff --git a/tui_gateway/pet_payload.py b/tui_gateway/pet_payload.py new file mode 100644 index 000000000000..6c1d670b4725 --- /dev/null +++ b/tui_gateway/pet_payload.py @@ -0,0 +1,312 @@ +"""Pet payload / spritesheet build helpers (extracted from tui_gateway.server). + +Slice R3-S1 of the tui_gateway/server.py god-file extraction - epic #78647, +target #78630. Byte-verbatim move of the pet payload/spritesheet cluster +(window 7989-8282); server.py re-exports every name so bare-name consumers +(HandlerRegistry-installed pet handlers, the R1 ``pet.changed`` watcher, and +test monkeypatches) keep resolving through the server namespace unchanged. +""" + +from __future__ import annotations + +import logging +import os +import threading + +logger = logging.getLogger(__name__) + + +def _pet_frame_counts(spritesheet) -> dict: + """Real (padding-trimmed) frame count per state, for the desktop canvas. + + Fail-open: a decode hiccup returns ``{}`` and the canvas falls back to its + static ``framesPerState`` rather than breaking the (cosmetic) pet. + """ + try: + from agent.pet import render + + return render.state_frame_counts(str(spritesheet)) + except Exception: # noqa: BLE001 - cosmetic, never break the surface + return {} + + +_pet_payload_cache_lock = threading.Lock() +_pet_payload_cache: dict[tuple, dict] = {} + + +def _pet_sheet_revision(spritesheet) -> str: + """Stable revision id for one spritesheet file.""" + try: + stat = spritesheet.stat() + return f"{stat.st_mtime_ns}:{stat.st_size}" + except Exception: # noqa: BLE001 - cosmetic, never break the surface + return "0:0" + + +def _pet_payload_cache_key(pet, *, scale: float) -> tuple | None: + """Cache key for the expensive sprite payload build.""" + try: + stat = pet.spritesheet.stat() + except Exception: # noqa: BLE001 + return None + return ( + str(pet.spritesheet), + stat.st_mtime_ns, + stat.st_size, + pet.slug, + pet.display_name, + round(scale, 4), + ) + + +def _clone_pet_payload(payload: dict) -> dict: + """Shallow-clone cached payloads so callers can't mutate shared state.""" + out = dict(payload) + if isinstance(payload.get("framesByState"), dict): + out["framesByState"] = dict(payload["framesByState"]) + if isinstance(payload.get("framesByRow"), dict): + out["framesByRow"] = dict(payload["framesByRow"]) + if isinstance(payload.get("stateRows"), list): + out["stateRows"] = list(payload["stateRows"]) + return out + + +def _pet_row_frame_counts(spritesheet) -> dict: + """Real frame count per concrete spritesheet row name.""" + try: + from PIL import Image + + from agent.pet import constants, render + + with Image.open(spritesheet) as opened: + image = opened.convert("RGBA") + cols = max(1, image.width // constants.FRAME_W) + row_count = max(1, image.height // constants.FRAME_H) + rows = constants.state_rows_for_grid(row_count) + out: dict[str, int] = {} + for row_idx, name in enumerate(rows[:row_count]): + top = row_idx * constants.FRAME_H + count = 0 + for col in range(cols): + left = col * constants.FRAME_W + frame = image.crop((left, top, left + constants.FRAME_W, top + constants.FRAME_H)) + if render._frame_is_blank(frame): + break + count += 1 + out[name] = count + return out + except Exception: # noqa: BLE001 - cosmetic, never break the surface + return {} + + +def _pet_config_scale() -> float: + """Configured ``display.pet.scale`` (or the engine default), never raises.""" + from agent.pet import constants + + try: + from hermes_cli.config import load_config + + cfg = load_config() + display = cfg.get("display", {}) if isinstance(cfg.get("display"), dict) else {} + pet_cfg = display.get("pet", {}) if isinstance(display.get("pet"), dict) else {} + return float(pet_cfg.get("scale", constants.DEFAULT_SCALE) or constants.DEFAULT_SCALE) + except Exception: # noqa: BLE001 + return constants.DEFAULT_SCALE + + +def _pet_sprite_payload(pet, *, scale: float) -> dict: + """Build the renderer payload (spritesheet bytes + geometry) for *pet*. + + Shared by ``pet.info`` (the active mascot) and ``pet.hatch`` (the unadopted + preview) so both feed the desktop canvas / TUI from one shape. + """ + import base64 + + from agent.pet import constants + + cache_key = _pet_payload_cache_key(pet, scale=scale) + if cache_key is not None: + with _pet_payload_cache_lock: + cached = _pet_payload_cache.get(cache_key) + if cached is not None: + return _clone_pet_payload(cached) + + raw = pet.spritesheet.read_bytes() + suffix = pet.spritesheet.suffix.lower() + mime = "image/png" if suffix == ".png" else "image/webp" + payload = { + "slug": pet.slug, + "displayName": pet.display_name, + "mime": mime, + "spritesheetBase64": base64.standard_b64encode(raw).decode("ascii"), + "spritesheetRevision": _pet_sheet_revision(pet.spritesheet), + "frameW": constants.FRAME_W, + "frameH": constants.FRAME_H, + "framesPerState": constants.FRAMES_PER_STATE, + "framesByState": _pet_frame_counts(pet.spritesheet), + "framesByRow": _pet_row_frame_counts(pet.spritesheet), + "loopMs": constants.LOOP_MS, + "scale": scale, + "stateRows": _pet_state_rows(pet.spritesheet), + } + if cache_key is not None: + with _pet_payload_cache_lock: + _pet_payload_cache[cache_key] = payload + while len(_pet_payload_cache) > 8: + _pet_payload_cache.pop(next(iter(_pet_payload_cache))) + return _clone_pet_payload(payload) + + +def _pet_active_selection(): + """Resolve configured active pet + scale from config.""" + from agent.pet import constants, store + + try: + from hermes_cli.config import load_config + + cfg = load_config() + display = cfg.get("display", {}) if isinstance(cfg.get("display"), dict) else {} + pet_cfg = display.get("pet", {}) if isinstance(display.get("pet"), dict) else {} + except Exception: + pet_cfg = {} + + enabled = bool(pet_cfg.get("enabled")) + configured_slug = str(pet_cfg.get("slug", "") or "") + pet = store.resolve_active_pet(configured_slug) if enabled else None + scale = float(pet_cfg.get("scale", constants.DEFAULT_SCALE) or constants.DEFAULT_SCALE) + return enabled, pet, scale + + +def _pet_state_rows(spritesheet) -> list[str]: + """Row taxonomy for the concrete active pet sheet. + + Hermes has to support both the legacy 8-row petdex atlas and the current + Codex/petdex 9-row atlas. The desktop canvas gets this list and indexes it + with the same `PetState` names the Python renderer uses. + """ + try: + from PIL import Image + + from agent.pet import constants + + with Image.open(spritesheet) as image: + row_count = max(1, image.height // constants.FRAME_H) + return list(constants.state_rows_for_grid(row_count)) + except Exception: # noqa: BLE001 - cosmetic, never break the surface + from agent.pet import constants + + return list(constants.STATE_ROWS) + + +def _pet_gen_root(): + """Profile-scoped staging dir for in-progress generation drafts.""" + from hermes_constants import get_hermes_home + + root = get_hermes_home() / "cache" / "pet-gen" + root.mkdir(parents=True, exist_ok=True) + return root + + +def _pet_gen_sweep(root, *, max_age_s: float = 3600.0) -> None: + """Drop stale draft staging dirs so cache never grows unbounded.""" + import shutil + import time + + try: + now = time.time() + for child in root.iterdir(): + if child.is_dir() and now - child.stat().st_mtime > max_age_s: + shutil.rmtree(child, ignore_errors=True) + except Exception as exc: # noqa: BLE001 - cleanup is best-effort + logger.debug("pet-gen sweep failed: %s", exc) + + +def _pet_png_data_uri(path, *, max_px: int = 160) -> str: + """Downscaled PNG data URI for a draft image (small preview payload).""" + import base64 + import io + + from PIL import Image + + with Image.open(path) as opened: + img = opened.convert("RGBA") + img.thumbnail((max_px, max_px), Image.LANCZOS) + buf = io.BytesIO() + img.save(buf, format="PNG") + return "data:image/png;base64," + base64.standard_b64encode(buf.getvalue()).decode("ascii") + + +# Cooperative cancellation for the heavy pet generation paths. The client's Stop +# aborts its RPC immediately, but the worker-pool generation keeps running unless +# told to stop — pet.cancel flips a token's flag, which generate_base_drafts / +# hatch_pet poll between provider calls to skip work they haven't started. +_pet_cancel_lock = threading.Lock() +_pet_cancelled: set[str] = set() +_PET_REFERENCE_MIME_EXT = { + "png": "png", + "jpeg": "jpg", + "jpg": "jpg", + "webp": "webp", + "gif": "gif", +} +try: + _PET_REFERENCE_MAX_BYTES = max( + 1, + int(os.environ.get("HERMES_PET_REFERENCE_MAX_BYTES") or str(16 * 1024 * 1024)), + ) +except (TypeError, ValueError): + _PET_REFERENCE_MAX_BYTES = 16 * 1024 * 1024 + + +def _pet_reference_images_from_data_url(ref_raw: str, stage) -> list: + """Decode + validate a reference-image data URL into the stage dir.""" + import base64 + import binascii + import re as _re + + match = _re.match(r"^data:image/([a-zA-Z0-9.+-]+);base64,(.*)$", ref_raw, _re.DOTALL) + if not match: + raise ValueError("invalid reference image format") + + mime = match.group(1).lower() + ext = _PET_REFERENCE_MIME_EXT.get(mime) + if ext is None: + raise ValueError("unsupported reference image type") + + payload = "".join(match.group(2).split()) + approx = (len(payload) * 3) // 4 + if approx > _PET_REFERENCE_MAX_BYTES: + raise ValueError("reference image too large") + + try: + raw = base64.b64decode(payload, validate=True) + except (binascii.Error, ValueError) as exc: + raise ValueError("invalid reference image data") from exc + + if len(raw) > _PET_REFERENCE_MAX_BYTES: + raise ValueError("reference image too large") + + ref_path = stage / f"reference.{ext}" + ref_path.write_bytes(raw) + return [ref_path] + + +def _pet_cancel_arm(token: str) -> None: + """Clear a stale cancel flag at the start of a generate/hatch run.""" + with _pet_cancel_lock: + _pet_cancelled.discard(token) + + +def _pet_cancel_request(token: str) -> None: + with _pet_cancel_lock: + _pet_cancelled.add(token) + + +def _pet_is_cancelled(token: str) -> bool: + with _pet_cancel_lock: + return token in _pet_cancelled + + +def _pet_cancel_release(token: str) -> None: + with _pet_cancel_lock: + _pet_cancelled.discard(token) diff --git a/tui_gateway/server.py b/tui_gateway/server.py index a36a539408b1..c2ab298bb333 100644 --- a/tui_gateway/server.py +++ b/tui_gateway/server.py @@ -7986,300 +7986,16 @@ def _main_runtime_from_agent(agent) -> dict | None: return runtime or None -def _pet_frame_counts(spritesheet) -> dict: - """Real (padding-trimmed) frame count per state, for the desktop canvas. - - Fail-open: a decode hiccup returns ``{}`` and the canvas falls back to its - static ``framesPerState`` rather than breaking the (cosmetic) pet. - """ - try: - from agent.pet import render - - return render.state_frame_counts(str(spritesheet)) - except Exception: # noqa: BLE001 - cosmetic, never break the surface - return {} - - -_pet_payload_cache_lock = threading.Lock() -_pet_payload_cache: dict[tuple, dict] = {} - - -def _pet_sheet_revision(spritesheet) -> str: - """Stable revision id for one spritesheet file.""" - try: - stat = spritesheet.stat() - return f"{stat.st_mtime_ns}:{stat.st_size}" - except Exception: # noqa: BLE001 - cosmetic, never break the surface - return "0:0" - - -def _pet_payload_cache_key(pet, *, scale: float) -> tuple | None: - """Cache key for the expensive sprite payload build.""" - try: - stat = pet.spritesheet.stat() - except Exception: # noqa: BLE001 - return None - return ( - str(pet.spritesheet), - stat.st_mtime_ns, - stat.st_size, - pet.slug, - pet.display_name, - round(scale, 4), - ) - - -def _clone_pet_payload(payload: dict) -> dict: - """Shallow-clone cached payloads so callers can't mutate shared state.""" - out = dict(payload) - if isinstance(payload.get("framesByState"), dict): - out["framesByState"] = dict(payload["framesByState"]) - if isinstance(payload.get("framesByRow"), dict): - out["framesByRow"] = dict(payload["framesByRow"]) - if isinstance(payload.get("stateRows"), list): - out["stateRows"] = list(payload["stateRows"]) - return out - - -def _pet_row_frame_counts(spritesheet) -> dict: - """Real frame count per concrete spritesheet row name.""" - try: - from PIL import Image - - from agent.pet import constants, render - - with Image.open(spritesheet) as opened: - image = opened.convert("RGBA") - cols = max(1, image.width // constants.FRAME_W) - row_count = max(1, image.height // constants.FRAME_H) - rows = constants.state_rows_for_grid(row_count) - out: dict[str, int] = {} - for row_idx, name in enumerate(rows[:row_count]): - top = row_idx * constants.FRAME_H - count = 0 - for col in range(cols): - left = col * constants.FRAME_W - frame = image.crop((left, top, left + constants.FRAME_W, top + constants.FRAME_H)) - if render._frame_is_blank(frame): - break - count += 1 - out[name] = count - return out - except Exception: # noqa: BLE001 - cosmetic, never break the surface - return {} - - -def _pet_config_scale() -> float: - """Configured ``display.pet.scale`` (or the engine default), never raises.""" - from agent.pet import constants - - try: - from hermes_cli.config import load_config - - cfg = load_config() - display = cfg.get("display", {}) if isinstance(cfg.get("display"), dict) else {} - pet_cfg = display.get("pet", {}) if isinstance(display.get("pet"), dict) else {} - return float(pet_cfg.get("scale", constants.DEFAULT_SCALE) or constants.DEFAULT_SCALE) - except Exception: # noqa: BLE001 - return constants.DEFAULT_SCALE - - -def _pet_sprite_payload(pet, *, scale: float) -> dict: - """Build the renderer payload (spritesheet bytes + geometry) for *pet*. - - Shared by ``pet.info`` (the active mascot) and ``pet.hatch`` (the unadopted - preview) so both feed the desktop canvas / TUI from one shape. - """ - import base64 - - from agent.pet import constants - - cache_key = _pet_payload_cache_key(pet, scale=scale) - if cache_key is not None: - with _pet_payload_cache_lock: - cached = _pet_payload_cache.get(cache_key) - if cached is not None: - return _clone_pet_payload(cached) - - raw = pet.spritesheet.read_bytes() - suffix = pet.spritesheet.suffix.lower() - mime = "image/png" if suffix == ".png" else "image/webp" - payload = { - "slug": pet.slug, - "displayName": pet.display_name, - "mime": mime, - "spritesheetBase64": base64.standard_b64encode(raw).decode("ascii"), - "spritesheetRevision": _pet_sheet_revision(pet.spritesheet), - "frameW": constants.FRAME_W, - "frameH": constants.FRAME_H, - "framesPerState": constants.FRAMES_PER_STATE, - "framesByState": _pet_frame_counts(pet.spritesheet), - "framesByRow": _pet_row_frame_counts(pet.spritesheet), - "loopMs": constants.LOOP_MS, - "scale": scale, - "stateRows": _pet_state_rows(pet.spritesheet), - } - if cache_key is not None: - with _pet_payload_cache_lock: - _pet_payload_cache[cache_key] = payload - while len(_pet_payload_cache) > 8: - _pet_payload_cache.pop(next(iter(_pet_payload_cache))) - return _clone_pet_payload(payload) - - -def _pet_active_selection(): - """Resolve configured active pet + scale from config.""" - from agent.pet import constants, store - - try: - from hermes_cli.config import load_config - - cfg = load_config() - display = cfg.get("display", {}) if isinstance(cfg.get("display"), dict) else {} - pet_cfg = display.get("pet", {}) if isinstance(display.get("pet"), dict) else {} - except Exception: - pet_cfg = {} - - enabled = bool(pet_cfg.get("enabled")) - configured_slug = str(pet_cfg.get("slug", "") or "") - pet = store.resolve_active_pet(configured_slug) if enabled else None - scale = float(pet_cfg.get("scale", constants.DEFAULT_SCALE) or constants.DEFAULT_SCALE) - return enabled, pet, scale - - -def _pet_state_rows(spritesheet) -> list[str]: - """Row taxonomy for the concrete active pet sheet. - - Hermes has to support both the legacy 8-row petdex atlas and the current - Codex/petdex 9-row atlas. The desktop canvas gets this list and indexes it - with the same `PetState` names the Python renderer uses. - """ - try: - from PIL import Image - - from agent.pet import constants - - with Image.open(spritesheet) as image: - row_count = max(1, image.height // constants.FRAME_H) - return list(constants.state_rows_for_grid(row_count)) - except Exception: # noqa: BLE001 - cosmetic, never break the surface - from agent.pet import constants - - return list(constants.STATE_ROWS) - - -def _pet_gen_root(): - """Profile-scoped staging dir for in-progress generation drafts.""" - from hermes_constants import get_hermes_home - - root = get_hermes_home() / "cache" / "pet-gen" - root.mkdir(parents=True, exist_ok=True) - return root - - -def _pet_gen_sweep(root, *, max_age_s: float = 3600.0) -> None: - """Drop stale draft staging dirs so cache never grows unbounded.""" - import shutil - import time - - try: - now = time.time() - for child in root.iterdir(): - if child.is_dir() and now - child.stat().st_mtime > max_age_s: - shutil.rmtree(child, ignore_errors=True) - except Exception as exc: # noqa: BLE001 - cleanup is best-effort - logger.debug("pet-gen sweep failed: %s", exc) - - -def _pet_png_data_uri(path, *, max_px: int = 160) -> str: - """Downscaled PNG data URI for a draft image (small preview payload).""" - import base64 - import io - - from PIL import Image - - with Image.open(path) as opened: - img = opened.convert("RGBA") - img.thumbnail((max_px, max_px), Image.LANCZOS) - buf = io.BytesIO() - img.save(buf, format="PNG") - return "data:image/png;base64," + base64.standard_b64encode(buf.getvalue()).decode("ascii") - - -# Cooperative cancellation for the heavy pet generation paths. The client's Stop -# aborts its RPC immediately, but the worker-pool generation keeps running unless -# told to stop — pet.cancel flips a token's flag, which generate_base_drafts / -# hatch_pet poll between provider calls to skip work they haven't started. -_pet_cancel_lock = threading.Lock() -_pet_cancelled: set[str] = set() -_PET_REFERENCE_MIME_EXT = { - "png": "png", - "jpeg": "jpg", - "jpg": "jpg", - "webp": "webp", - "gif": "gif", -} -try: - _PET_REFERENCE_MAX_BYTES = max( - 1, - int(os.environ.get("HERMES_PET_REFERENCE_MAX_BYTES") or str(16 * 1024 * 1024)), - ) -except (TypeError, ValueError): - _PET_REFERENCE_MAX_BYTES = 16 * 1024 * 1024 - - -def _pet_reference_images_from_data_url(ref_raw: str, stage) -> list: - """Decode + validate a reference-image data URL into the stage dir.""" - import base64 - import binascii - import re as _re - - match = _re.match(r"^data:image/([a-zA-Z0-9.+-]+);base64,(.*)$", ref_raw, _re.DOTALL) - if not match: - raise ValueError("invalid reference image format") - - mime = match.group(1).lower() - ext = _PET_REFERENCE_MIME_EXT.get(mime) - if ext is None: - raise ValueError("unsupported reference image type") - - payload = "".join(match.group(2).split()) - approx = (len(payload) * 3) // 4 - if approx > _PET_REFERENCE_MAX_BYTES: - raise ValueError("reference image too large") - - try: - raw = base64.b64decode(payload, validate=True) - except (binascii.Error, ValueError) as exc: - raise ValueError("invalid reference image data") from exc - - if len(raw) > _PET_REFERENCE_MAX_BYTES: - raise ValueError("reference image too large") - - ref_path = stage / f"reference.{ext}" - ref_path.write_bytes(raw) - return [ref_path] - - -def _pet_cancel_arm(token: str) -> None: - """Clear a stale cancel flag at the start of a generate/hatch run.""" - with _pet_cancel_lock: - _pet_cancelled.discard(token) - - -def _pet_cancel_request(token: str) -> None: - with _pet_cancel_lock: - _pet_cancelled.add(token) - - -def _pet_is_cancelled(token: str) -> bool: - with _pet_cancel_lock: - return token in _pet_cancelled - - -def _pet_cancel_release(token: str) -> None: - with _pet_cancel_lock: - _pet_cancelled.discard(token) +from tui_gateway.pet_payload import ( # noqa: E402,F401 - legacy re-exports + _pet_frame_counts, _pet_payload_cache_lock, _pet_payload_cache, + _pet_sheet_revision, _pet_payload_cache_key, _clone_pet_payload, + _pet_row_frame_counts, _pet_config_scale, _pet_sprite_payload, + _pet_active_selection, _pet_state_rows, _pet_gen_root, _pet_gen_sweep, + _pet_png_data_uri, _pet_cancel_lock, _pet_cancelled, + _PET_REFERENCE_MIME_EXT, _PET_REFERENCE_MAX_BYTES, + _pet_reference_images_from_data_url, _pet_cancel_arm, _pet_cancel_request, + _pet_is_cancelled, _pet_cancel_release, +) # ===========================================================================