| 1 | from __future__ import annotations |
| 2 | |
| 3 | import base64 |
| 4 | import threading |
| 5 | import time |
| 6 | import uuid |
| 7 | from dataclasses import dataclass |
| 8 | |
| 9 | |
| 10 | REF_PREFIX = "a0-ephemeral-image://" |
| 11 | DEFAULT_TTL_SECONDS = 15 * 60 |
| 12 | |
| 13 | |
| 14 | @dataclass(frozen=True) |
| 15 | class EphemeralImage: |
| 16 | ref: str |
| 17 | context_id: str |
| 18 | mime: str |
| 19 | data: str |
| 20 | name: str |
| 21 | created_at: float |
| 22 | expires_at: float |
| 23 | |
| 24 | @property |
| 25 | def data_url(self) -> str: |
| 26 | return f"data:{self.mime};base64,{self.data}" |
| 27 | |
| 28 | @property |
| 29 | def display_name(self) -> str: |
| 30 | return self.name or display_ref(self.ref) |
| 31 | |
| 32 | |
| 33 | _store: dict[str, EphemeralImage] = {} |
| 34 | _lock = threading.RLock() |
| 35 | |
| 36 | |
| 37 | def put_image_bytes( |
| 38 | *, |
| 39 | context_id: str, |
| 40 | mime: str, |
| 41 | payload: bytes, |
| 42 | name: str = "", |
| 43 | ttl_seconds: float = DEFAULT_TTL_SECONDS, |
| 44 | ) -> str: |
| 45 | data = base64.b64encode(bytes(payload or b"")).decode("ascii") |
| 46 | return put_image( |
| 47 | context_id=context_id, |
| 48 | mime=mime, |
| 49 | data=data, |
| 50 | name=name, |
| 51 | ttl_seconds=ttl_seconds, |
| 52 | ) |
| 53 | |
| 54 | |
| 55 | def put_image( |
| 56 | *, |
| 57 | context_id: str, |
| 58 | mime: str, |
| 59 | data: str, |
| 60 | name: str = "", |
| 61 | ttl_seconds: float = DEFAULT_TTL_SECONDS, |
| 62 | ) -> str: |
| 63 | compact_data = _compact_base64(data) |
| 64 | if not compact_data: |
| 65 | raise ValueError("ephemeral image data is empty") |
| 66 | base64.b64decode(compact_data, validate=True) |
| 67 | |
| 68 | normalized_mime = _normalize_mime(mime) |
| 69 | now = time.time() |
| 70 | ref = f"{REF_PREFIX}{uuid.uuid4().hex}" |
| 71 | image = EphemeralImage( |
| 72 | ref=ref, |
| 73 | context_id=str(context_id or "").strip(), |
| 74 | mime=normalized_mime, |
| 75 | data=compact_data, |
| 76 | name=str(name or "").strip(), |
| 77 | created_at=now, |
| 78 | expires_at=now + max(1.0, float(ttl_seconds or DEFAULT_TTL_SECONDS)), |
| 79 | ) |
| 80 | with _lock: |
| 81 | _prune_expired_locked(now) |
| 82 | _store[ref] = image |
| 83 | return ref |
| 84 | |
| 85 | |
| 86 | def is_ref(value: object) -> bool: |
| 87 | return str(value or "").strip().startswith(REF_PREFIX) |
| 88 | |
| 89 | |
| 90 | def display_ref(ref: str) -> str: |
| 91 | value = str(ref or "").strip() |
| 92 | if not is_ref(value): |
| 93 | return value |
| 94 | return f"{REF_PREFIX}<ephemeral>" |
| 95 | |
| 96 | |
| 97 | def get_image(ref: str, *, context_id: str = "") -> EphemeralImage | None: |
| 98 | return _resolve_image(ref, context_id=context_id, consume=False) |
| 99 | |
| 100 | |
| 101 | def consume_image(ref: str, *, context_id: str = "") -> EphemeralImage | None: |
| 102 | return _resolve_image(ref, context_id=context_id, consume=True) |
| 103 | |
| 104 | |
| 105 | def delete_image(ref: str) -> None: |
| 106 | with _lock: |
| 107 | _store.pop(str(ref or "").strip(), None) |
| 108 | |
| 109 | |
| 110 | def clear_context(context_id: str) -> None: |
| 111 | normalized_context = str(context_id or "").strip() |
| 112 | with _lock: |
| 113 | for ref, image in list(_store.items()): |
| 114 | if image.context_id == normalized_context: |
| 115 | _store.pop(ref, None) |
| 116 | |
| 117 | |
| 118 | def _resolve_image(ref: str, *, context_id: str = "", consume: bool) -> EphemeralImage | None: |
| 119 | value = str(ref or "").strip() |
| 120 | if not is_ref(value): |
| 121 | return None |
| 122 | |
| 123 | now = time.time() |
| 124 | with _lock: |
| 125 | _prune_expired_locked(now) |
| 126 | image = _store.get(value) |
| 127 | if image is None: |
| 128 | return None |
| 129 | requested_context = str(context_id or "").strip() |
| 130 | if requested_context and image.context_id and image.context_id != requested_context: |
| 131 | return None |
| 132 | if consume: |
| 133 | _store.pop(value, None) |
| 134 | return image |
| 135 | |
| 136 | |
| 137 | def _compact_base64(data: str) -> str: |
| 138 | return "".join(char for char in str(data or "") if not char.isspace()) |
| 139 | |
| 140 | |
| 141 | def _normalize_mime(mime: str) -> str: |
| 142 | value = str(mime or "").strip().lower() |
| 143 | return value if value.startswith("image/") else "image/jpeg" |
| 144 | |
| 145 | |
| 146 | def _prune_expired_locked(now: float | None = None) -> None: |
| 147 | current = time.time() if now is None else float(now) |
| 148 | for ref, image in list(_store.items()): |
| 149 | if image.expires_at <= current: |
| 150 | _store.pop(ref, None) |