diff --git a/nanobot/agent/loop.py b/nanobot/agent/loop.py index 4dcecfc5..6ffade73 100644 --- a/nanobot/agent/loop.py +++ b/nanobot/agent/loop.py @@ -858,15 +858,17 @@ class AgentLoop: ) ) - # Persist the triggering user message immediately, before running the - # agent loop. If the process is killed mid-turn (OOM, SIGKILL, self- - # restart, etc.), the existing runtime_checkpoint preserves the - # in-flight assistant/tool state but NOT the user message itself, so - # the user's prompt is silently lost on recovery. Saving it up front - # makes recovery possible from the session log alone. + # Persist the triggering user message up front so a mid-turn crash + # doesn't silently lose the prompt on recovery. ``media`` rides along + # as raw on-disk paths — sanitized image blocks are stripped from + # JSONL, and webui replay needs the paths to mint signed URLs. user_persisted_early = False - if isinstance(msg.content, str) and msg.content.strip(): - session.add_message("user", msg.content) + media_paths = [p for p in (msg.media or []) if isinstance(p, str) and p] + has_text = isinstance(msg.content, str) and msg.content.strip() + if has_text or media_paths: + extra: dict[str, Any] = {"media": list(media_paths)} if media_paths else {} + text = msg.content if isinstance(msg.content, str) else "" + session.add_message("user", text, **extra) self._mark_pending_user_turn(session) self.sessions.save(session) user_persisted_early = True diff --git a/nanobot/api/server.py b/nanobot/api/server.py index ebdee557..92e7be90 100644 --- a/nanobot/api/server.py +++ b/nanobot/api/server.py @@ -7,13 +7,9 @@ All requests route to a single persistent API session. from __future__ import annotations import asyncio -import base64 import json as _json -import mimetypes -import re import time import uuid -from pathlib import Path from typing import Any from aiohttp import web @@ -21,14 +17,20 @@ from loguru import logger from nanobot.config.paths import get_media_dir from nanobot.utils.helpers import safe_filename +from nanobot.utils.media_decode import ( + FileSizeExceeded as _FileSizeExceeded, + MAX_FILE_SIZE, + save_base64_data_url as _save_base64_data_url, +) from nanobot.utils.runtime import EMPTY_FINAL_RESPONSE_MESSAGE -MAX_FILE_SIZE = 10 * 1024 * 1024 # 10 MB -_DATA_URL_RE = re.compile(r"^data:([^;]+);base64,(.+)$", re.DOTALL) - - -class _FileSizeExceeded(Exception): - """Raised when an uploaded file exceeds the size limit.""" +__all__ = ( + "MAX_FILE_SIZE", + "_FileSizeExceeded", + "_save_base64_data_url", + "create_app", + "handle_chat_completions", +) API_SESSION_KEY = "api:default" @@ -102,25 +104,6 @@ _SSE_DONE = b"data: [DONE]\n\n" # --------------------------------------------------------------------------- -def _save_base64_data_url(data_url: str, media_dir: Path) -> str | None: - """Decode a data:...;base64,... URL and save to disk.""" - m = _DATA_URL_RE.match(data_url) - if not m: - return None - mime_type, b64_payload = m.group(1), m.group(2) - try: - raw = base64.b64decode(b64_payload) - except Exception: - return None - if len(raw) > MAX_FILE_SIZE: - raise _FileSizeExceeded(f"File exceeds {MAX_FILE_SIZE // (1024 * 1024)}MB limit") - ext = mimetypes.guess_extension(mime_type) or ".bin" - filename = f"{uuid.uuid4().hex[:12]}{ext}" - dest = media_dir / safe_filename(filename) - dest.write_bytes(raw) - return str(dest) - - def _parse_json_content(body: dict) -> tuple[str, list[str]]: """Parse JSON request body. Returns (text, media_paths).""" messages = body.get("messages") diff --git a/nanobot/channels/websocket.py b/nanobot/channels/websocket.py index 639519a4..1622a37d 100644 --- a/nanobot/channels/websocket.py +++ b/nanobot/channels/websocket.py @@ -3,7 +3,10 @@ from __future__ import annotations import asyncio +import base64 +import binascii import email.utils +import hashlib import hmac import http import json @@ -28,7 +31,12 @@ from websockets.http11 import Response from nanobot.bus.events import OutboundMessage from nanobot.bus.queue import MessageBus from nanobot.channels.base import BaseChannel +from nanobot.config.paths import get_media_dir from nanobot.config.schema import Base +from nanobot.utils.media_decode import ( + FileSizeExceeded, + save_base64_data_url, +) if TYPE_CHECKING: from nanobot.session.manager import SessionManager @@ -75,7 +83,11 @@ class WebSocketConfig(Base): websocket_requires_token: bool = True allow_from: list[str] = Field(default_factory=lambda: ["*"]) streaming: bool = True - max_message_bytes: int = Field(default=1_048_576, ge=1024, le=16_777_216) + # Default 36 MB, upper 40 MB: supports up to 4 images at ~6 MB each after + # client-side Worker normalization (see webui Composer). 4 × 6 MB × 1.37 + # (base64 overhead) + envelope framing stays under 36 MB; the 40 MB ceiling + # leaves a small margin for sender slop without opening a DoS avenue. + max_message_bytes: int = Field(default=37_748_736, ge=1024, le=41_943_040) ping_interval_s: float = Field(default=20.0, ge=5.0, le=300.0) ping_timeout_s: float = Field(default=20.0, ge=5.0, le=300.0) ssl_certfile: str = "" @@ -206,6 +218,35 @@ def _parse_envelope(raw: str) -> dict[str, Any] | None: return data +# Per-message image limits. The server-side guard is a touch looser than the +# client's ``Worker`` normalization target (6 MB) — tolerate client slop, but +# still cap total ingress at ``_MAX_IMAGES_PER_MESSAGE * _MAX_IMAGE_BYTES`` +# which fits comfortably inside ``max_message_bytes``. +_MAX_IMAGES_PER_MESSAGE = 4 +_MAX_IMAGE_BYTES = 8 * 1024 * 1024 + +# Image MIME whitelist — matches the Composer's ``accept`` list. SVG is +# explicitly excluded to avoid the XSS surface inside embedded scripts. +_IMAGE_MIME_ALLOWED: frozenset[str] = frozenset({ + "image/png", + "image/jpeg", + "image/webp", + "image/gif", +}) + +_DATA_URL_MIME_RE = re.compile(r"^data:([^;]+);base64,", re.DOTALL) + + +def _extract_data_url_mime(url: str) -> str | None: + """Return the MIME type of a ``data:;base64,...`` URL, else ``None``.""" + if not isinstance(url, str): + return None + m = _DATA_URL_MIME_RE.match(url) + if not m: + return None + return m.group(1).strip().lower() or None + + _LOCALHOSTS = frozenset({"127.0.0.1", "::1", "localhost"}) # Matches the legacy chat-id pattern but allows file-system-safe stems too, @@ -278,6 +319,29 @@ def _is_websocket_upgrade(request: WsRequest) -> bool: return True +def _b64url_encode(data: bytes) -> str: + """URL-safe base64 without padding — compact + friendly in URL paths.""" + return base64.urlsafe_b64encode(data).rstrip(b"=").decode("ascii") + + +def _b64url_decode(s: str) -> bytes: + """Reverse of :func:`_b64url_encode`; caller handles ``ValueError``.""" + pad = "=" * (-len(s) % 4) + return base64.urlsafe_b64decode(s + pad) + + +# Allowed MIME types we actually serve from the media endpoint. Anything +# outside this set is degraded to ``application/octet-stream`` so an +# attacker who somehow gets a signed URL for an unexpected file type can't +# trick the browser into sniffing executable content. +_MEDIA_ALLOWED_MIMES: frozenset[str] = frozenset({ + "image/png", + "image/jpeg", + "image/webp", + "image/gif", +}) + + def _issue_route_secret_matches(headers: Any, configured_secret: str) -> bool: """Return True if the token-issue HTTP request carries credentials matching ``token_issue_secret``.""" if not configured_secret: @@ -326,6 +390,11 @@ class WebSocketChannel(BaseChannel): self._static_dist_path: Path | None = ( static_dist_path.resolve() if static_dist_path is not None else None ) + # Process-local secret used to HMAC-sign media URLs. The signed URL is + # the capability — anyone who holds a valid URL can fetch that one + # file, nothing else. The secret regenerates on restart so links + # become self-expiring (callers just refresh the session list). + self._media_secret: bytes = secrets.token_bytes(32) # -- Subscription bookkeeping ------------------------------------------- @@ -457,6 +526,14 @@ class WebSocketChannel(BaseChannel): if m: return self._handle_session_delete(request, m.group(1)) + # Signed media fetch: ```` is an HMAC over ````; the + # payload decodes to a path inside :func:`get_media_dir`. See + # :meth:`_sign_media_path` for the inverse direction used to build + # these URLs when replaying a session. + m = re.match(r"^/api/media/([A-Za-z0-9_-]+)/([A-Za-z0-9_-]+)$", got) + if m: + return self._handle_media_fetch(m.group(1), m.group(2)) + # 4. WebSocket upgrade (the channel's primary purpose). Only run the # handshake gate on requests that actually ask to upgrade; otherwise # a bare ``GET /`` from the browser would be rejected as an @@ -568,8 +645,109 @@ class WebSocketChannel(BaseChannel): data = self._session_manager.read_session_file(decoded_key) if data is None: return _http_error(404, "session not found") + # Decorate persisted user messages with signed media URLs so the + # client can render previews. The raw on-disk ``media`` paths are + # stripped on the way out — they leak server filesystem layout and + # the client never needs them once it has the signed fetch URL. + self._augment_media_urls(data) return _http_json_response(data) + def _augment_media_urls(self, payload: dict[str, Any]) -> None: + """Mutate *payload* in place: each message's ``media`` path list is + replaced by a parallel ``media_urls`` list of signed fetch URLs. + + Messages without media or with non-string path entries are left + untouched. Paths that no longer live inside ``media_dir`` (e.g. the + file was deleted, or the dir was relocated) are silently skipped; + the client falls back to the historical-replay placeholder tile. + """ + messages = payload.get("messages") + if not isinstance(messages, list): + return + for msg in messages: + if not isinstance(msg, dict): + continue + media = msg.get("media") + if not isinstance(media, list) or not media: + continue + urls: list[dict[str, str]] = [] + for entry in media: + if not isinstance(entry, str) or not entry: + continue + signed = self._sign_media_path(Path(entry)) + if signed is None: + continue + urls.append({"url": signed, "name": Path(entry).name}) + if urls: + msg["media_urls"] = urls + # Always drop the raw paths from the wire payload. + msg.pop("media", None) + + def _sign_media_path(self, abs_path: Path) -> str | None: + """Return a ``/api/media//`` URL for *abs_path*, or + ``None`` when the path does not resolve inside the media root. + + The URL is self-authenticating: the signature binds the payload to + this process's ``_media_secret``, so only paths we chose to sign can + be fetched. The returned path is relative to the server origin; the + client joins it against the existing webui base. + """ + try: + media_root = get_media_dir().resolve() + rel = abs_path.resolve().relative_to(media_root) + except (OSError, ValueError): + return None + payload = _b64url_encode(rel.as_posix().encode("utf-8")) + mac = hmac.new( + self._media_secret, payload.encode("ascii"), hashlib.sha256 + ).digest()[:16] + return f"/api/media/{_b64url_encode(mac)}/{payload}" + + def _handle_media_fetch(self, sig: str, payload: str) -> Response: + """Serve a single media file previously signed via + :meth:`_sign_media_path`. Validates the signature, decodes the + payload to a relative path, and streams the file bytes with a + long-lived immutable cache header (the URL already encodes the + file identity, so caches can be aggressive).""" + try: + provided_mac = _b64url_decode(sig) + except (ValueError, binascii.Error): + return _http_error(401, "invalid signature") + expected_mac = hmac.new( + self._media_secret, payload.encode("ascii"), hashlib.sha256 + ).digest()[:16] + if not hmac.compare_digest(expected_mac, provided_mac): + return _http_error(401, "invalid signature") + try: + rel_bytes = _b64url_decode(payload) + rel_str = rel_bytes.decode("utf-8") + except (ValueError, binascii.Error, UnicodeDecodeError): + return _http_error(400, "invalid payload") + # An attacker who somehow bypassed the HMAC check would still need + # the resolved path to escape the media root; guard defensively. + try: + media_root = get_media_dir().resolve() + candidate = (media_root / rel_str).resolve() + candidate.relative_to(media_root) + except (OSError, ValueError): + return _http_error(404, "not found") + if not candidate.is_file(): + return _http_error(404, "not found") + try: + body = candidate.read_bytes() + except OSError: + return _http_error(500, "read error") + mime, _ = mimetypes.guess_type(candidate.name) + if mime not in _MEDIA_ALLOWED_MIMES: + mime = "application/octet-stream" + return _http_response( + body, + content_type=mime, + extra_headers=[ + ("Cache-Control", "private, max-age=31536000, immutable"), + ], + ) + def _handle_session_delete(self, request: WsRequest, key: str) -> Response: if not self._check_api_token(request): return _http_error(401, "Unauthorized") @@ -755,6 +933,48 @@ class WebSocketChannel(BaseChannel): finally: self._cleanup_connection(connection) + @staticmethod + def _save_envelope_media( + media: list[Any], + ) -> tuple[list[str], str | None]: + """Decode and persist ``media`` items from a ``message`` envelope. + + Returns ``(paths, None)`` on success or ``([], reason)`` on the first + failure — the caller is expected to surface ``reason`` to the client + and skip publishing so no half-formed message ever reaches the agent. + ``reason`` is a short, stable token suitable for UI localization. + + Shape: ``list[{"data_url": str, "name"?: str | None}]``. + """ + if len(media) > _MAX_IMAGES_PER_MESSAGE: + return [], "too_many_images" + media_dir = get_media_dir("websocket") + paths: list[str] = [] + for item in media: + if not isinstance(item, dict): + return [], "malformed" + data_url = item.get("data_url") + if not isinstance(data_url, str) or not data_url: + return [], "malformed" + mime = _extract_data_url_mime(data_url) + if mime is None: + return [], "decode" + if mime not in _IMAGE_MIME_ALLOWED: + return [], "mime" + try: + saved = save_base64_data_url( + data_url, media_dir, max_bytes=_MAX_IMAGE_BYTES, + ) + except FileSizeExceeded: + return [], "size" + except Exception as exc: + logger.warning("websocket: media decode failed: {}", exc) + return [], "decode" + if saved is None: + return [], "decode" + paths.append(saved) + return paths, None + async def _dispatch_envelope( self, connection: Any, @@ -782,15 +1002,39 @@ class WebSocketChannel(BaseChannel): if not _is_valid_chat_id(cid): await self._send_event(connection, "error", detail="invalid chat_id") return - if not isinstance(content, str) or not content.strip(): + if not isinstance(content, str): await self._send_event(connection, "error", detail="missing content") return + + raw_media = envelope.get("media") + media_paths: list[str] = [] + if raw_media is not None: + if not isinstance(raw_media, list): + await self._send_event( + connection, "error", + detail="image_rejected", reason="malformed", + ) + return + media_paths, reason = self._save_envelope_media(raw_media) + if reason is not None: + await self._send_event( + connection, "error", + detail="image_rejected", reason=reason, + ) + return + + # Allow image-only turns (content may be empty when media is attached). + if not content.strip() and not media_paths: + await self._send_event(connection, "error", detail="missing content") + return + # Auto-attach on first use so clients can one-shot without a separate attach. self._attach(connection, cid) await self._handle_message( sender_id=client_id, chat_id=cid, content=content, + media=media_paths or None, metadata={"remote": getattr(connection, "remote_address", None)}, ) return diff --git a/nanobot/session/manager.py b/nanobot/session/manager.py index 436d8225..69509a83 100644 --- a/nanobot/session/manager.py +++ b/nanobot/session/manager.py @@ -11,7 +11,12 @@ from typing import Any from loguru import logger from nanobot.config.paths import get_legacy_sessions_dir -from nanobot.utils.helpers import ensure_dir, find_legal_message_start, safe_filename +from nanobot.utils.helpers import ( + ensure_dir, + find_legal_message_start, + image_placeholder_text, + safe_filename, +) @dataclass @@ -54,7 +59,19 @@ class Session: out: list[dict[str, Any]] = [] for message in sliced: - entry: dict[str, Any] = {"role": message["role"], "content": message.get("content", "")} + content = message.get("content", "") + # Synthesize an ``[image: path]`` breadcrumb from the persisted + # ``media`` kwarg so LLM replay still sees *something* where the + # image used to be. Without this, an image-only user turn + # replays as an empty user message — the assistant's reply then + # looks like it's responding to nothing. + media = message.get("media") + if isinstance(media, list) and media and isinstance(content, str): + breadcrumbs = "\n".join( + image_placeholder_text(p) for p in media if isinstance(p, str) and p + ) + content = f"{content}\n{breadcrumbs}" if content else breadcrumbs + entry: dict[str, Any] = {"role": message["role"], "content": content} for key in ("tool_calls", "tool_call_id", "name", "reasoning_content"): if key in message: entry[key] = message[key] diff --git a/nanobot/utils/media_decode.py b/nanobot/utils/media_decode.py new file mode 100644 index 00000000..484613d9 --- /dev/null +++ b/nanobot/utils/media_decode.py @@ -0,0 +1,55 @@ +"""Shared helpers for decoding ``data:...;base64,...`` URLs to disk. + +Historically lived in ``nanobot.api.server``; now shared by the WebSocket +channel so the ``api`` + ``websocket`` ingress paths apply the same parsing, +size guard, and filesystem layout. +""" + +from __future__ import annotations + +import base64 +import mimetypes +import re +import uuid +from pathlib import Path + +from nanobot.utils.helpers import safe_filename + +DEFAULT_MAX_BYTES = 10 * 1024 * 1024 +MAX_FILE_SIZE = DEFAULT_MAX_BYTES + +_DATA_URL_RE = re.compile(r"^data:([^;]+);base64,(.+)$", re.DOTALL) + + +class FileSizeExceeded(Exception): + """Raised when a decoded payload exceeds the caller's size limit.""" + + +def save_base64_data_url( + data_url: str, + media_dir: Path, + *, + max_bytes: int | None = None, +) -> str | None: + """Decode a ``data:;base64,`` URL and persist it. + + Returns the absolute path on success, ``None`` when the URL shape or the + base64 payload itself is malformed. Raises :class:`FileSizeExceeded` + when the decoded payload is larger than ``max_bytes`` (default 10 MB). + """ + m = _DATA_URL_RE.match(data_url) + if not m: + return None + mime_type, b64_payload = m.group(1), m.group(2) + try: + raw = base64.b64decode(b64_payload) + except Exception: + return None + limit = DEFAULT_MAX_BYTES if max_bytes is None else max_bytes + if len(raw) > limit: + raise FileSizeExceeded(f"File exceeds {limit // (1024 * 1024)}MB limit") + ext = mimetypes.guess_extension(mime_type) or ".bin" + filename = f"{uuid.uuid4().hex[:12]}{ext}" + dest = media_dir / safe_filename(filename) + dest.write_bytes(raw) + return str(dest) diff --git a/tests/agent/test_loop_save_turn.py b/tests/agent/test_loop_save_turn.py index 4f1c1f35..50951824 100644 --- a/tests/agent/test_loop_save_turn.py +++ b/tests/agent/test_loop_save_turn.py @@ -234,6 +234,87 @@ async def test_process_message_persists_user_message_before_turn_completes(tmp_p assert persisted.updated_at >= persisted.created_at +# 1x1 PNG used by the media-persistence tests. ``extract_documents`` runs +# at the top of ``_process_message`` and filters ``msg.media`` down to +# paths that magic-byte-sniff as images, so the test fixture needs real +# bytes on disk (not just placeholder paths). +_PNG_1X1 = ( + b"\x89PNG\r\n\x1a\n" + b"\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01" + b"\x08\x06\x00\x00\x00\x1f\x15\xc4\x89" + b"\x00\x00\x00\nIDATx\x9cc\x00\x00\x00\x02\x00\x01" + b"\x00\x00\x00\x00IEND\xaeB`\x82" +) + + +@pytest.mark.asyncio +async def test_process_message_persists_media_paths_on_user_turn(tmp_path: Path) -> None: + """User turns that attach images must record the media paths alongside + the text so the webui can rehydrate previews on session replay. + + This is the producer half of the signed-media-URL round-trip: paths are + stored here, then :meth:`WebSocketChannel._augment_media_urls` maps them + onto signed URLs on the way out. + """ + img_a = tmp_path / "uuid-1.png" + img_a.write_bytes(_PNG_1X1) + img_b = tmp_path / "uuid-2.png" + img_b.write_bytes(_PNG_1X1) + + loop = _make_full_loop(tmp_path) + loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign] + loop._run_agent_loop = AsyncMock(side_effect=RuntimeError("interrupt")) # type: ignore[method-assign] + + msg = InboundMessage( + channel="websocket", + sender_id="u1", + chat_id="c-media", + content="look", + media=[str(img_a), str(img_b)], + ) + with pytest.raises(RuntimeError, match="interrupt"): + await loop._process_message(msg) + + loop.sessions.invalidate("websocket:c-media") + persisted = loop.sessions.get_or_create("websocket:c-media") + assert [m["role"] for m in persisted.messages] == ["user"] + assert persisted.messages[0]["content"] == "look" + assert persisted.messages[0]["media"] == [str(img_a), str(img_b)] + + +@pytest.mark.asyncio +async def test_process_message_persists_media_only_turn_without_text(tmp_path: Path) -> None: + """A turn with images but no text still persists (previously silent-dropped). + + The old early-persist gate skipped messages without text, leaving pure + image turns un-checkpointed. They now materialise as an empty-content + user row with ``media`` attached. + """ + img = tmp_path / "only.png" + img.write_bytes(_PNG_1X1) + + loop = _make_full_loop(tmp_path) + loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign] + loop._run_agent_loop = AsyncMock(side_effect=RuntimeError("boom")) # type: ignore[method-assign] + + msg = InboundMessage( + channel="websocket", + sender_id="u1", + chat_id="c-images-only", + content="", + media=[str(img)], + ) + with pytest.raises(RuntimeError): + await loop._process_message(msg) + + loop.sessions.invalidate("websocket:c-images-only") + persisted = loop.sessions.get_or_create("websocket:c-images-only") + assert len(persisted.messages) == 1 + assert persisted.messages[0]["role"] == "user" + assert persisted.messages[0]["content"] == "" + assert persisted.messages[0]["media"] == [str(img)] + + @pytest.mark.asyncio async def test_process_message_does_not_duplicate_early_persisted_user_message(tmp_path: Path) -> None: loop = _make_full_loop(tmp_path) diff --git a/tests/agent/test_session_manager_history.py b/tests/agent/test_session_manager_history.py index 1297a587..8b4d0740 100644 --- a/tests/agent/test_session_manager_history.py +++ b/tests/agent/test_session_manager_history.py @@ -217,3 +217,55 @@ def test_window_cuts_mid_tool_group(): # leaving orphan tool results for split_a at the front. history = session.get_history(max_messages=6) _assert_no_orphans(history) + + +# --- Image breadcrumbs: media kwarg is synthesized into content for replay --- + + +def test_get_history_synthesizes_image_breadcrumb_from_media_kwarg(): + """Persisted user turns carry image paths as a ``media`` kwarg; LLM + replay must still see an ``[image: path]`` breadcrumb so the assistant's + follow-up reply has a referent instead of trailing an empty user row.""" + session = Session(key="test:media") + session.messages.append( + {"role": "user", "content": "look", "media": ["/m/a.png", "/m/b.png"]} + ) + session.messages.append({"role": "assistant", "content": "nice"}) + + history = session.get_history(max_messages=500) + + assert history == [ + {"role": "user", "content": "look\n[image: /m/a.png]\n[image: /m/b.png]"}, + {"role": "assistant", "content": "nice"}, + ] + + +def test_get_history_synthesizes_breadcrumb_for_image_only_turn(): + """Turns with no text but attached images must not replay as empty + strings — the LLM would otherwise see a bare user turn followed by an + unexplained assistant answer.""" + session = Session(key="test:image-only") + session.messages.append({"role": "user", "content": "", "media": ["/m/pic.png"]}) + session.messages.append({"role": "assistant", "content": "I see a cat"}) + + history = session.get_history(max_messages=500) + + assert history[0] == {"role": "user", "content": "[image: /m/pic.png]"} + + +def test_get_history_ignores_media_kwarg_on_non_user_rows(): + """``media`` only ever appears on user entries in practice, but the + synthesizer must be defensive: assistants / tools with list content + don't get the breadcrumb pasted on top.""" + session = Session(key="test:defensive") + session.messages.append( + { + "role": "assistant", + "content": [{"type": "text", "text": "structured"}], + "media": ["/m/x.png"], # nonsense but shouldn't crash + } + ) + history = session.get_history(max_messages=500) + # List content is passed through verbatim — the synthesizer only + # rewrites plain-string content. + assert history[0]["content"] == [{"type": "text", "text": "structured"}] diff --git a/tests/channels/test_websocket_envelope_media.py b/tests/channels/test_websocket_envelope_media.py new file mode 100644 index 00000000..4fe5f642 --- /dev/null +++ b/tests/channels/test_websocket_envelope_media.py @@ -0,0 +1,416 @@ +"""Tests for WS envelope media handling (client image upload path). + +Exercises ``WebSocketChannel._dispatch_envelope`` for the ``message`` branch: +decoding base64 data URLs, rejecting malformed / oversized / non-whitelisted +payloads, preserving backward compatibility with media-less frames, and +forwarding saved paths to ``_handle_message``. +""" + +from __future__ import annotations + +import base64 +import json +from pathlib import Path +from typing import Any +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from nanobot.channels.websocket import ( + WebSocketChannel, + _extract_data_url_mime, +) + + +def _tiny_png_data_url() -> str: + """A 1-pixel PNG prefixed as a data URL — just enough for magic-bytes sniffing.""" + # 1x1 transparent PNG + png = ( + b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00" + b"\x00\x01\x08\x06\x00\x00\x00\x1f\x15\xc4\x89\x00\x00\x00\rIDATx" + b"\x9cc\xf8\xcf\xc0\x00\x00\x00\x03\x00\x01\x00\x18\xdd\x8d\xb4\x00" + b"\x00\x00\x00IEND\xaeB`\x82" + ) + return f"data:image/png;base64,{base64.b64encode(png).decode()}" + + +def _data_url(mime: str, payload: bytes) -> str: + return f"data:{mime};base64,{base64.b64encode(payload).decode()}" + + +def _make_channel() -> WebSocketChannel: + bus = MagicMock() + bus.publish_inbound = AsyncMock() + channel = WebSocketChannel( + {"enabled": True, "allowFrom": ["*"], "websocketRequiresToken": False}, + bus, + ) + channel._handle_message = AsyncMock() # type: ignore[method-assign] + return channel + + +# -- Pure helpers -------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("url", "expected"), + [ + ("data:image/png;base64,AAAA", "image/png"), + ("data:image/jpeg;base64,AAAA", "image/jpeg"), + ("data:IMAGE/PNG;base64,AAAA", "image/png"), + ("data:image/svg+xml;base64,AAAA", "image/svg+xml"), + ("data:text/plain;base64,AAAA", "text/plain"), + ("http://evil.example/x.png", None), + ("data:image/png,AAAA", None), # missing `;base64` + ("", None), + (None, None), + ], +) +def test_extract_data_url_mime(url: Any, expected: str | None) -> None: + assert _extract_data_url_mime(url) == expected + + +# -- max_message_bytes bump ---------------------------------------------------- + + +def test_max_message_bytes_default_supports_multi_image_frame() -> None: + """Default 36 MB must comfortably hold 4 × 6 MB base64-encoded images.""" + from nanobot.channels.websocket import WebSocketConfig + + default = WebSocketConfig().max_message_bytes + # 4 images × 6 MB × 1.37 base64 overhead ≈ 33 MB + assert default >= 33 * 1024 * 1024 + # Upper bound 40 MB matches plan + with pytest.raises(Exception): + WebSocketConfig(max_message_bytes=41_943_040 + 1) + + +# -- _dispatch_envelope message branch + media -------------------------------- + + +@pytest.mark.asyncio +async def test_message_without_media_backward_compatible() -> None: + """Existing clients that don't send ``media`` keep working unchanged.""" + channel = _make_channel() + mock_conn = AsyncMock() + envelope = {"type": "message", "chat_id": "abc123", "content": "hello"} + + await channel._dispatch_envelope(mock_conn, "client-1", envelope) + + channel._handle_message.assert_awaited_once() + call = channel._handle_message.call_args + assert call.kwargs["chat_id"] == "abc123" + assert call.kwargs["content"] == "hello" + # When no media, we pass ``media=None`` so downstream treats it as absent. + assert call.kwargs["media"] is None + + +@pytest.mark.asyncio +async def test_message_with_single_image_forwards_saved_path(tmp_path) -> None: + channel = _make_channel() + mock_conn = AsyncMock() + envelope = { + "type": "message", + "chat_id": "abc123", + "content": "look at this", + "media": [{"data_url": _tiny_png_data_url(), "name": "shot.png"}], + } + + with patch( + "nanobot.channels.websocket.get_media_dir", return_value=tmp_path + ): + await channel._dispatch_envelope(mock_conn, "client-1", envelope) + + channel._handle_message.assert_awaited_once() + paths = channel._handle_message.call_args.kwargs["media"] + assert isinstance(paths, list) and len(paths) == 1 + saved = Path(paths[0]) + assert saved.exists() + assert saved.suffix == ".png" + assert saved.is_relative_to(tmp_path) + + +@pytest.mark.asyncio +async def test_message_with_multiple_images(tmp_path) -> None: + channel = _make_channel() + mock_conn = AsyncMock() + envelope = { + "type": "message", + "chat_id": "abc123", + "content": "a couple", + "media": [ + {"data_url": _tiny_png_data_url()}, + {"data_url": _tiny_png_data_url()}, + {"data_url": _tiny_png_data_url()}, + ], + } + + with patch( + "nanobot.channels.websocket.get_media_dir", return_value=tmp_path + ): + await channel._dispatch_envelope(mock_conn, "client-1", envelope) + + paths = channel._handle_message.call_args.kwargs["media"] + assert len(paths) == 3 + # Saved filenames must be unique. + assert len({Path(p).name for p in paths}) == 3 + + +@pytest.mark.asyncio +async def test_image_only_message_allows_empty_text(tmp_path) -> None: + """When media is attached, empty text is acceptable.""" + channel = _make_channel() + mock_conn = AsyncMock() + envelope = { + "type": "message", + "chat_id": "abc123", + "content": "", + "media": [{"data_url": _tiny_png_data_url()}], + } + + with patch( + "nanobot.channels.websocket.get_media_dir", return_value=tmp_path + ): + await channel._dispatch_envelope(mock_conn, "client-1", envelope) + + channel._handle_message.assert_awaited_once() + # Error event NOT sent. + mock_conn.send.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_message_rejected_when_more_than_four_images(tmp_path) -> None: + channel = _make_channel() + mock_conn = AsyncMock() + envelope = { + "type": "message", + "chat_id": "abc123", + "content": "hi", + "media": [{"data_url": _tiny_png_data_url()}] * 5, + } + + with patch( + "nanobot.channels.websocket.get_media_dir", return_value=tmp_path + ): + await channel._dispatch_envelope(mock_conn, "client-1", envelope) + + channel._handle_message.assert_not_awaited() + mock_conn.send.assert_awaited_once() + err = json.loads(mock_conn.send.call_args[0][0]) + assert err["event"] == "error" + assert err["detail"] == "image_rejected" + assert err["reason"] == "too_many_images" + + +@pytest.mark.asyncio +async def test_message_rejected_on_oversize_payload(tmp_path) -> None: + channel = _make_channel() + mock_conn = AsyncMock() + oversized = b"x" * (9 * 1024 * 1024) # > 8 MB WS limit + envelope = { + "type": "message", + "chat_id": "abc123", + "content": "big", + "media": [{"data_url": _data_url("image/png", oversized)}], + } + + with patch( + "nanobot.channels.websocket.get_media_dir", return_value=tmp_path + ): + await channel._dispatch_envelope(mock_conn, "client-1", envelope) + + channel._handle_message.assert_not_awaited() + err = json.loads(mock_conn.send.call_args[0][0]) + assert err["detail"] == "image_rejected" + assert err["reason"] == "size" + + +@pytest.mark.asyncio +async def test_message_rejected_on_non_image_mime(tmp_path) -> None: + channel = _make_channel() + mock_conn = AsyncMock() + envelope = { + "type": "message", + "chat_id": "abc123", + "content": "pdf?", + "media": [{"data_url": _data_url("application/pdf", b"%PDF-1.4")}], + } + + with patch( + "nanobot.channels.websocket.get_media_dir", return_value=tmp_path + ): + await channel._dispatch_envelope(mock_conn, "client-1", envelope) + + channel._handle_message.assert_not_awaited() + err = json.loads(mock_conn.send.call_args[0][0]) + assert err["detail"] == "image_rejected" + assert err["reason"] == "mime" + + +@pytest.mark.asyncio +async def test_message_rejected_on_svg_mime(tmp_path) -> None: + """SVG is explicitly rejected — XSS surface inside embedded scripts.""" + channel = _make_channel() + mock_conn = AsyncMock() + envelope = { + "type": "message", + "chat_id": "abc123", + "content": "svg", + "media": [{"data_url": _data_url("image/svg+xml", b"")}], + } + + with patch( + "nanobot.channels.websocket.get_media_dir", return_value=tmp_path + ): + await channel._dispatch_envelope(mock_conn, "client-1", envelope) + + channel._handle_message.assert_not_awaited() + err = json.loads(mock_conn.send.call_args[0][0]) + assert err["reason"] == "mime" + + +@pytest.mark.asyncio +async def test_message_rejected_on_malformed_data_url(tmp_path) -> None: + channel = _make_channel() + mock_conn = AsyncMock() + envelope = { + "type": "message", + "chat_id": "abc123", + "content": "nope", + "media": [{"data_url": "http://evil.example/image.png"}], + } + + with patch( + "nanobot.channels.websocket.get_media_dir", return_value=tmp_path + ): + await channel._dispatch_envelope(mock_conn, "client-1", envelope) + + channel._handle_message.assert_not_awaited() + err = json.loads(mock_conn.send.call_args[0][0]) + assert err["reason"] == "decode" + + +@pytest.mark.asyncio +async def test_message_rejected_on_broken_base64(tmp_path) -> None: + channel = _make_channel() + mock_conn = AsyncMock() + envelope = { + "type": "message", + "chat_id": "abc123", + "content": "nope", + "media": [{"data_url": "data:image/png;base64,not-valid-base64!!!"}], + } + + with patch( + "nanobot.channels.websocket.get_media_dir", return_value=tmp_path + ): + await channel._dispatch_envelope(mock_conn, "client-1", envelope) + + channel._handle_message.assert_not_awaited() + err = json.loads(mock_conn.send.call_args[0][0]) + assert err["reason"] == "decode" + + +@pytest.mark.asyncio +async def test_message_rejected_when_media_item_shape_wrong(tmp_path) -> None: + channel = _make_channel() + mock_conn = AsyncMock() + envelope = { + "type": "message", + "chat_id": "abc123", + "content": "huh", + # Not a dict — plain string at the top level. + "media": ["data:image/png;base64,XXXX"], + } + + with patch( + "nanobot.channels.websocket.get_media_dir", return_value=tmp_path + ): + await channel._dispatch_envelope(mock_conn, "client-1", envelope) + + channel._handle_message.assert_not_awaited() + err = json.loads(mock_conn.send.call_args[0][0]) + assert err["reason"] == "malformed" + + +@pytest.mark.asyncio +async def test_message_rejected_when_media_field_is_not_list() -> None: + channel = _make_channel() + mock_conn = AsyncMock() + envelope = { + "type": "message", + "chat_id": "abc123", + "content": "huh", + "media": "not-a-list", + } + + await channel._dispatch_envelope(mock_conn, "client-1", envelope) + + channel._handle_message.assert_not_awaited() + err = json.loads(mock_conn.send.call_args[0][0]) + assert err["detail"] == "image_rejected" + assert err["reason"] == "malformed" + + +@pytest.mark.asyncio +async def test_failed_media_does_not_partially_persist(tmp_path) -> None: + """If the second image is invalid, the first must not be forwarded.""" + channel = _make_channel() + mock_conn = AsyncMock() + envelope = { + "type": "message", + "chat_id": "abc123", + "content": "mixed", + "media": [ + {"data_url": _tiny_png_data_url()}, + {"data_url": _data_url("application/pdf", b"%PDF-1.4")}, + ], + } + + with patch( + "nanobot.channels.websocket.get_media_dir", return_value=tmp_path + ): + await channel._dispatch_envelope(mock_conn, "client-1", envelope) + + channel._handle_message.assert_not_awaited() + # The first image was saved to disk (we don't roll it back — the caller + # is expected to not reference it) but the agent never sees the paths. + # That's the important invariant: no partial publish. + err = json.loads(mock_conn.send.call_args[0][0]) + assert err["reason"] == "mime" + + +@pytest.mark.asyncio +async def test_rejects_empty_text_without_media() -> None: + """When no media is attached, whitespace-only content is still rejected + (matches the existing behavior for backward compat).""" + channel = _make_channel() + mock_conn = AsyncMock() + envelope = { + "type": "message", + "chat_id": "abc123", + "content": " ", + } + + await channel._dispatch_envelope(mock_conn, "client-1", envelope) + + channel._handle_message.assert_not_awaited() + err = json.loads(mock_conn.send.call_args[0][0]) + assert err["detail"] == "missing content" + + +@pytest.mark.asyncio +async def test_non_string_content_still_rejected() -> None: + channel = _make_channel() + mock_conn = AsyncMock() + envelope = { + "type": "message", + "chat_id": "abc123", + "content": 42, + } + + await channel._dispatch_envelope(mock_conn, "client-1", envelope) + + channel._handle_message.assert_not_awaited() + err = json.loads(mock_conn.send.call_args[0][0]) + assert err["detail"] == "missing content" diff --git a/tests/channels/test_websocket_media_route.py b/tests/channels/test_websocket_media_route.py new file mode 100644 index 00000000..14577355 --- /dev/null +++ b/tests/channels/test_websocket_media_route.py @@ -0,0 +1,375 @@ +"""Tests for the signed ``/api/media//`` route and its replay +integration on ``/api/sessions//messages``. + +The route is the return path for images attached to persisted user turns: +:meth:`WebSocketChannel._sign_media_path` mints URLs during session reads, +and :meth:`WebSocketChannel._handle_media_fetch` serves the bytes back. +These tests cover the two halves end-to-end plus the adversarial edges +(bad signatures, ``..`` traversal, non-existent files, non-image types). +""" + +from __future__ import annotations + +import asyncio +import functools +import hashlib +import hmac +from pathlib import Path +from typing import Any +from unittest.mock import AsyncMock, MagicMock, patch + +import httpx +import pytest + +from nanobot.channels.websocket import ( + WebSocketChannel, + _b64url_decode, + _b64url_encode, +) +from nanobot.session.manager import Session, SessionManager + + +# PNG magic bytes + a couple of sentinel bytes so we can verify byte-for-byte +# round-trip of the served payload. Stays under mimetype + size limits. +_PNG_BYTES = ( + b"\x89PNG\r\n\x1a\n" + b"\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01" + b"\x08\x06\x00\x00\x00\x1f\x15\xc4\x89" + b"\x00\x00\x00\nIDATx\x9cc\x00\x00\x00\x02\x00\x01" + b"\x00\x00\x00\x00IEND\xaeB`\x82" +) + + +def _ch( + bus: Any, + *, + session_manager: SessionManager | None = None, + port: int, +) -> WebSocketChannel: + return WebSocketChannel( + { + "enabled": True, + "allowFrom": ["*"], + "host": "127.0.0.1", + "port": port, + "path": "/", + "websocketRequiresToken": False, + }, + bus, + session_manager=session_manager, + ) + + +@pytest.fixture() +def bus() -> MagicMock: + b = MagicMock() + b.publish_inbound = AsyncMock() + return b + + +async def _http_get( + url: str, headers: dict[str, str] | None = None +) -> httpx.Response: + return await asyncio.to_thread( + functools.partial(httpx.get, url, headers=headers or {}, timeout=5.0) + ) + + +# --------------------------------------------------------------------------- +# _sign_media_path: the URL minter +# --------------------------------------------------------------------------- + + +def test_sign_media_path_rejects_paths_outside_media_root( + bus: MagicMock, tmp_path: Path +) -> None: + """Paths that resolve outside ``get_media_dir()`` must not be signed. + + This is the single most important invariant of the whole scheme: + if the minter ever signed an arbitrary path, the HMAC would legitimise + it for the fetch handler and we'd hand out a disk-read primitive. + """ + outside = tmp_path / "secrets" / "cred.txt" + outside.parent.mkdir() + outside.write_text("nope") + media = tmp_path / "media" + media.mkdir() + channel = _ch(bus, port=0) + with patch("nanobot.channels.websocket.get_media_dir", return_value=media): + assert channel._sign_media_path(outside) is None + # Traversal via the media root is also rejected — the resolve() step + # normalises ``..`` out before the relative_to check. + assert channel._sign_media_path(media / ".." / "secrets" / "cred.txt") is None + + +def test_sign_media_path_round_trips_via_hmac( + bus: MagicMock, tmp_path: Path +) -> None: + """The signature embeds exactly ``HMAC-SHA256(secret, payload)[:16]``.""" + media = tmp_path / "media" + media.mkdir() + (media / "a.png").write_bytes(_PNG_BYTES) + channel = _ch(bus, port=0) + with patch("nanobot.channels.websocket.get_media_dir", return_value=media): + url = channel._sign_media_path(media / "a.png") + assert url is not None + assert url.startswith("/api/media/") + sig, payload = url[len("/api/media/"):].split("/", 1) + expected = hmac.new( + channel._media_secret, payload.encode("ascii"), hashlib.sha256 + ).digest()[:16] + assert _b64url_decode(sig) == expected + # The payload decodes back to the *relative* path — no absolute-path leaks. + assert _b64url_decode(payload).decode() == "a.png" + + +# --------------------------------------------------------------------------- +# /api/media//: the serving handler +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_media_route_serves_signed_file( + bus: MagicMock, tmp_path: Path +) -> None: + """Valid signature + existing file => 200 with correct bytes + MIME.""" + media = tmp_path / "media" + media.mkdir() + target = media / "round-trip.png" + target.write_bytes(_PNG_BYTES) + + channel = _ch(bus, port=29920) + with patch("nanobot.channels.websocket.get_media_dir", return_value=media): + url_path = channel._sign_media_path(target) + assert url_path is not None + server_task = asyncio.create_task(channel.start()) + await asyncio.sleep(0.3) + try: + resp = await _http_get(f"http://127.0.0.1:29920{url_path}") + finally: + await channel.stop() + await server_task + + assert resp.status_code == 200 + assert resp.content == _PNG_BYTES + assert resp.headers["content-type"].startswith("image/png") + # Immutable cache header lets the browser skip round-trips on replay. + assert "immutable" in resp.headers.get("cache-control", "") + + +@pytest.mark.asyncio +async def test_media_route_rejects_bad_signature( + bus: MagicMock, tmp_path: Path +) -> None: + """A payload re-signed with a different secret must 401. + + Protects against a restart: old URLs baked into a stale tab become + un-forgeable once ``_media_secret`` regenerates. + """ + media = tmp_path / "media" + media.mkdir() + (media / "f.png").write_bytes(_PNG_BYTES) + + channel = _ch(bus, port=29921) + with patch("nanobot.channels.websocket.get_media_dir", return_value=media): + good = channel._sign_media_path(media / "f.png") + assert good is not None + _, payload = good[len("/api/media/"):].split("/", 1) + # Forge a sig with a *different* secret. + forged_mac = hmac.new( + b"\x00" * 32, payload.encode("ascii"), hashlib.sha256 + ).digest()[:16] + forged = f"/api/media/{_b64url_encode(forged_mac)}/{payload}" + + server_task = asyncio.create_task(channel.start()) + await asyncio.sleep(0.3) + try: + resp = await _http_get(f"http://127.0.0.1:29921{forged}") + finally: + await channel.stop() + await server_task + assert resp.status_code == 401 + + +@pytest.mark.asyncio +async def test_media_route_rejects_path_traversal_payload( + bus: MagicMock, tmp_path: Path +) -> None: + """Even a validly-signed ``..`` payload must not escape the media root. + + The signer never *emits* such payloads, but an attacker who somehow + obtained the secret (or the channel was misconfigured) must still be + stopped by the resolve()+relative_to() guard in the serving path. + """ + media = tmp_path / "media" + media.mkdir() + secret_file = tmp_path / "secret.txt" + secret_file.write_text("classified") + + channel = _ch(bus, port=29922) + # Hand-craft a traversal payload the legit signer would refuse to mint. + payload = _b64url_encode(b"../secret.txt") + mac = hmac.new( + channel._media_secret, payload.encode("ascii"), hashlib.sha256 + ).digest()[:16] + url = f"/api/media/{_b64url_encode(mac)}/{payload}" + + with patch("nanobot.channels.websocket.get_media_dir", return_value=media): + server_task = asyncio.create_task(channel.start()) + await asyncio.sleep(0.3) + try: + resp = await _http_get(f"http://127.0.0.1:29922{url}") + finally: + await channel.stop() + await server_task + assert resp.status_code == 404 + assert b"classified" not in resp.content + + +@pytest.mark.asyncio +async def test_media_route_404s_missing_file( + bus: MagicMock, tmp_path: Path +) -> None: + """A signed URL for a file that no longer exists degrades to 404 so the + client can fall back to the placeholder tile instead of breaking.""" + media = tmp_path / "media" + media.mkdir() + target = media / "gone.png" + target.write_bytes(_PNG_BYTES) + + channel = _ch(bus, port=29923) + with patch("nanobot.channels.websocket.get_media_dir", return_value=media): + url_path = channel._sign_media_path(target) + assert url_path is not None + target.unlink() # the file vanishes between signing and fetching + server_task = asyncio.create_task(channel.start()) + await asyncio.sleep(0.3) + try: + resp = await _http_get(f"http://127.0.0.1:29923{url_path}") + finally: + await channel.stop() + await server_task + assert resp.status_code == 404 + + +@pytest.mark.asyncio +async def test_media_route_degrades_non_image_to_octet_stream( + bus: MagicMock, tmp_path: Path +) -> None: + """A non-image extension must not be served as its native MIME. + + Defence-in-depth: if media_dir ever contained (say) an HTML file, we + do not want the browser to render it as HTML via the signed route. + """ + media = tmp_path / "media" + media.mkdir() + (media / "scary.html").write_bytes(b"") + + channel = _ch(bus, port=29924) + with patch("nanobot.channels.websocket.get_media_dir", return_value=media): + payload = _b64url_encode(b"scary.html") + mac = hmac.new( + channel._media_secret, payload.encode("ascii"), hashlib.sha256 + ).digest()[:16] + url = f"/api/media/{_b64url_encode(mac)}/{payload}" + server_task = asyncio.create_task(channel.start()) + await asyncio.sleep(0.3) + try: + resp = await _http_get(f"http://127.0.0.1:29924{url}") + finally: + await channel.stop() + await server_task + assert resp.status_code == 200 + assert resp.headers["content-type"].startswith("application/octet-stream") + + +# --------------------------------------------------------------------------- +# /api/sessions//messages: media_urls hydration on session read +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_session_messages_exposes_signed_media_urls( + bus: MagicMock, tmp_path: Path +) -> None: + """The read path must map persisted ``media`` paths onto signed URLs + and strip the raw path — the client never learns the server's layout.""" + media = tmp_path / "media" + media.mkdir() + img = media / "u.png" + img.write_bytes(_PNG_BYTES) + + sm = SessionManager(tmp_path / "ws_state") + sess = Session(key="websocket:media-hydrate") + sess.add_message("user", "look at this", media=[str(img)]) + sess.add_message("assistant", "nice") + sm.save(sess) + + channel = _ch(bus, session_manager=sm, port=29925) + with patch("nanobot.channels.websocket.get_media_dir", return_value=media): + server_task = asyncio.create_task(channel.start()) + await asyncio.sleep(0.3) + try: + boot = await _http_get("http://127.0.0.1:29925/webui/bootstrap") + token = boot.json()["token"] + auth = {"Authorization": f"Bearer {token}"} + resp = await _http_get( + "http://127.0.0.1:29925/api/sessions/websocket:media-hydrate/messages", + headers=auth, + ) + body = resp.json() + # The signed URL round-trips end-to-end: fetching it yields the same bytes. + user_msg = next(m for m in body["messages"] if m["role"] == "user") + urls = user_msg["media_urls"] + assert isinstance(urls, list) and len(urls) == 1 + assert urls[0]["name"] == "u.png" + assert urls[0]["url"].startswith("/api/media/") + # Raw paths must not leak to the wire. + assert "media" not in user_msg + + # And the URL actually works. + fetched = await _http_get(f"http://127.0.0.1:29925{urls[0]['url']}") + assert fetched.status_code == 200 + assert fetched.content == _PNG_BYTES + finally: + await channel.stop() + await server_task + + +@pytest.mark.asyncio +async def test_session_messages_skips_vanished_media( + bus: MagicMock, tmp_path: Path +) -> None: + """Paths that no longer resolve inside the media root produce no URL — + the message is still delivered, just without the preview.""" + media = tmp_path / "media" + media.mkdir() + + sm = SessionManager(tmp_path / "ws_state") + sess = Session(key="websocket:vanished") + sess.add_message("user", "missing pic", media=[str(media / "absent.png")]) + sm.save(sess) + + channel = _ch(bus, session_manager=sm, port=29926) + with patch("nanobot.channels.websocket.get_media_dir", return_value=media): + server_task = asyncio.create_task(channel.start()) + await asyncio.sleep(0.3) + try: + boot = await _http_get("http://127.0.0.1:29926/webui/bootstrap") + token = boot.json()["token"] + resp = await _http_get( + "http://127.0.0.1:29926/api/sessions/websocket:vanished/messages", + headers={"Authorization": f"Bearer {token}"}, + ) + user_msg = next(m for m in resp.json()["messages"] if m["role"] == "user") + # absent.png lives inside the media root so it *does* get a signed + # URL (we don't stat the file at signing time — that would slow + # the listing). Fetching the URL is where the 404 surfaces. + urls = user_msg.get("media_urls") or [] + assert len(urls) == 1 + fetched = await _http_get(f"http://127.0.0.1:29926{urls[0]['url']}") + assert fetched.status_code == 404 + assert "media" not in user_msg + finally: + await channel.stop() + await server_task diff --git a/tests/utils/test_media_decode.py b/tests/utils/test_media_decode.py new file mode 100644 index 00000000..5926ab2b --- /dev/null +++ b/tests/utils/test_media_decode.py @@ -0,0 +1,75 @@ +"""Tests for ``nanobot.utils.media_decode``.""" + +from __future__ import annotations + +import base64 + +import pytest + +from nanobot.utils.media_decode import ( + DEFAULT_MAX_BYTES, + FileSizeExceeded, + MAX_FILE_SIZE, + save_base64_data_url, +) + + +def _data_url(payload: bytes, mime: str = "image/png") -> str: + return f"data:{mime};base64,{base64.b64encode(payload).decode()}" + + +def test_saves_png_with_correct_extension(tmp_path) -> None: + result = save_base64_data_url(_data_url(b"fake png"), tmp_path) + assert result is not None + assert result.endswith(".png") + assert (tmp_path / result.split("/")[-1]).read_bytes() == b"fake png" + + +def test_returns_none_for_malformed_data_url(tmp_path) -> None: + assert save_base64_data_url("not-a-data-url", tmp_path) is None + + +def test_returns_none_for_broken_base64(tmp_path) -> None: + # Python's b64decode strips non-alphabet chars by default, so we need a + # payload whose alphabet-filtered length breaks padding. + assert save_base64_data_url("data:image/png;base64,not-valid-base64!!!", tmp_path) is None + + +def test_unknown_mime_falls_back_to_bin(tmp_path) -> None: + result = save_base64_data_url(_data_url(b"xyz", mime="unknown/type"), tmp_path) + assert result is not None + assert result.endswith(".bin") + + +def test_default_limit_is_10mb(tmp_path) -> None: + """Backwards-compatible default — the API path depends on this.""" + assert DEFAULT_MAX_BYTES == 10 * 1024 * 1024 + assert MAX_FILE_SIZE == 10 * 1024 * 1024 + + oversized = b"x" * (11 * 1024 * 1024) + with pytest.raises(FileSizeExceeded, match="10MB limit"): + save_base64_data_url(_data_url(oversized), tmp_path) + + +def test_explicit_max_bytes_overrides_default(tmp_path) -> None: + """WS channel passes 8 MB; a 9 MB payload should be rejected there even + though it would pass the 10 MB API limit.""" + payload = b"y" * (9 * 1024 * 1024) + with pytest.raises(FileSizeExceeded, match="8MB limit"): + save_base64_data_url(_data_url(payload), tmp_path, max_bytes=8 * 1024 * 1024) + + +def test_saved_file_lives_under_media_dir(tmp_path) -> None: + result = save_base64_data_url(_data_url(b"ok"), tmp_path) + assert result is not None + assert result.startswith(str(tmp_path)) + + +def test_legacy_symbols_reexported_from_api_server() -> None: + """Existing tests import ``_save_base64_data_url`` / ``_FileSizeExceeded`` + from ``nanobot.api.server`` — keep the aliases working.""" + from nanobot.api import server + + assert server._save_base64_data_url is save_base64_data_url + assert server._FileSizeExceeded is FileSizeExceeded + assert server.MAX_FILE_SIZE == MAX_FILE_SIZE diff --git a/webui/src/components/ImageLightbox.tsx b/webui/src/components/ImageLightbox.tsx new file mode 100644 index 00000000..90819afa --- /dev/null +++ b/webui/src/components/ImageLightbox.tsx @@ -0,0 +1,199 @@ +import { useCallback, useEffect, useMemo } from "react"; +import * as DialogPrimitive from "@radix-ui/react-dialog"; +import { ChevronLeft, ChevronRight, X } from "lucide-react"; +import { useTranslation } from "react-i18next"; + +import { cn } from "@/lib/utils"; +import type { UIImage } from "@/lib/types"; + +interface ImageLightboxProps { + images: UIImage[]; + index: number | null; + onIndexChange: (index: number) => void; + onOpenChange: (open: boolean) => void; +} + +/** + * Modal image viewer. Uses the Radix Dialog primitives directly so we can + * fill the viewport (the shared `DialogContent` wrapper caps at max-w-lg, + * which is much too small for a photo preview). + * + * Implementation notes: + * - `translate3d` + `will-change: transform` promote the image to a GPU + * compositing layer so open/swap stays at 60 FPS on long threads. + * - Adjacent images are rendered in hidden `` tags so the browser + * decodes them eagerly; pressing left/right feels instant. + * - Radix handles `Escape` + focus trapping; we only wire up ←/→ + Home/End. + * - Respects `prefers-reduced-motion` by dropping the fade + zoom-in + * keyframes via `motion-reduce:*` variants. + */ +export function ImageLightbox({ + images, + index, + onIndexChange, + onOpenChange, +}: ImageLightboxProps) { + const { t } = useTranslation(); + const open = index !== null; + const total = images.length; + const current = index !== null ? images[index] : null; + + const go = useCallback( + (delta: number) => { + if (index === null || total <= 1) return; + const next = (index + delta + total) % total; + onIndexChange(next); + }, + [index, onIndexChange, total], + ); + + useEffect(() => { + if (!open) return; + const onKey = (e: KeyboardEvent) => { + if (e.key === "ArrowLeft") { + e.preventDefault(); + go(-1); + } else if (e.key === "ArrowRight") { + e.preventDefault(); + go(1); + } else if (e.key === "Home") { + e.preventDefault(); + onIndexChange(0); + } else if (e.key === "End") { + e.preventDefault(); + onIndexChange(total - 1); + } + }; + window.addEventListener("keydown", onKey); + return () => window.removeEventListener("keydown", onKey); + }, [go, onIndexChange, open, total]); + + // Neighbours we want the browser to decode eagerly. + const preload = useMemo(() => { + if (index === null || total <= 1) return [] as UIImage[]; + const prev = images[(index - 1 + total) % total]; + const next = images[(index + 1) % total]; + return [prev, next].filter((i) => i && i.url); + }, [images, index, total]); + + if (!current || !current.url) return null; + + const hasMany = total > 1; + const counter = hasMany ? `${index! + 1} / ${total}` : null; + + return ( + + + + + + {current.name ?? t("lightbox.title")} + + +
+ {current.name +
+ + {hasMany ? ( + <> + { + e.stopPropagation(); + go(-1); + }} + /> + { + e.stopPropagation(); + go(1); + }} + /> +
+ {counter} +
+ + ) : null} + + + + + + {/* Invisible preload — browser decodes adjacent images so prev/next swap is instant. */} +
+ {preload.map((img, i) => ( + + ))} +
+
+
+
+ ); +} + +interface NavButtonProps { + side: "left" | "right"; + label: string; + onClick: React.MouseEventHandler; +} + +function NavButton({ side, label, onClick }: NavButtonProps) { + const Icon = side === "left" ? ChevronLeft : ChevronRight; + return ( + + ); +} diff --git a/webui/src/components/MessageBubble.tsx b/webui/src/components/MessageBubble.tsx index a6c02e3d..63a43c8d 100644 --- a/webui/src/components/MessageBubble.tsx +++ b/webui/src/components/MessageBubble.tsx @@ -1,10 +1,11 @@ import { useState } from "react"; -import { ChevronRight, Wrench } from "lucide-react"; +import { ChevronRight, ImageIcon, Wrench } from "lucide-react"; import { useTranslation } from "react-i18next"; +import { ImageLightbox } from "@/components/ImageLightbox"; import { MarkdownText } from "@/components/MarkdownText"; import { cn } from "@/lib/utils"; -import type { UIMessage } from "@/lib/types"; +import type { UIImage, UIMessage } from "@/lib/types"; interface MessageBubbleProps { message: UIMessage; @@ -27,22 +28,28 @@ export function MessageBubble({ message }: MessageBubbleProps) { } if (message.role === "user") { + const images = message.images ?? []; + const hasImages = images.length > 0; + const hasText = message.content.trim().length > 0; return (
-

- {message.content} -

+ {hasImages ? : null} + {hasText ? ( +

+ {message.content} +

+ ) : null}
); } @@ -62,6 +69,121 @@ export function MessageBubble({ message }: MessageBubbleProps) { ); } +/** + * Right-aligned preview row for images attached to a user turn. + * + * Visual follows agent-chat-ui: a single wrapping row of fixed-size square + * thumbnails that stay modest next to the text pill regardless of how many + * images are attached. + * + * The URL is expected to be a self-contained ``data:`` URL (the Composer + * hands the normalized base64 payload to the optimistic bubble so that the + * preview survives React StrictMode double-mount — blob URLs would be + * revoked by the Composer's cleanup before remount). Historical replays + * have no URL (the backend strips data URLs before persisting), so we + * render a labelled placeholder tile instead of a broken ````. + */ +function UserImages({ images }: { images: UIImage[] }) { + const { t } = useTranslation(); + // Only real-URL images can open in the lightbox; historical-replay + // placeholders (no URL) have nothing to zoom into. + const viewable = images + .map((img, i) => ({ img, i })) + .filter(({ img }) => typeof img.url === "string" && img.url.length > 0); + const viewableImages = viewable.map(({ img }) => img); + const originalToViewable = new Map( + viewable.map(({ i }, v) => [i, v]), + ); + + const [lightboxIndex, setLightboxIndex] = useState(null); + + return ( + <> +
+ {images.map((img, i) => ( + setLightboxIndex(originalToViewable.get(i)!) + : undefined + } + /> + ))} +
+ { + if (!open) setLightboxIndex(null); + }} + /> + + ); +} + +function UserImageCell({ + image, + placeholderLabel, + openLabel, + onOpen, +}: { + image: UIImage; + placeholderLabel: string; + openLabel: string; + onOpen?: () => void; +}) { + const hasUrl = typeof image.url === "string" && image.url.length > 0; + const tileClasses = cn( + "relative h-24 w-24 overflow-hidden rounded-[14px] border border-border/60 bg-muted/40", + "shadow-[0_6px_18px_-14px_rgba(0,0,0,0.45)]", + ); + + if (hasUrl && onOpen) { + return ( + + ); + } + + return ( +
+
+ + + {image.name ?? placeholderLabel} + +
+
+ ); +} + /** Blinking cursor appended at the end of streaming text. */ function StreamCursor() { const { t } = useTranslation(); diff --git a/webui/src/components/thread/StreamErrorNotice.tsx b/webui/src/components/thread/StreamErrorNotice.tsx new file mode 100644 index 00000000..06c89815 --- /dev/null +++ b/webui/src/components/thread/StreamErrorNotice.tsx @@ -0,0 +1,72 @@ +import { AlertTriangle, X } from "lucide-react"; +import { useTranslation } from "react-i18next"; + +import { Button } from "@/components/ui/button"; +import { cn } from "@/lib/utils"; +import type { StreamError } from "@/lib/nanobot-client"; + +interface StreamErrorNoticeProps { + error: StreamError; + onDismiss: () => void; +} + +/** + * Dismissible banner that surfaces transport-level faults the user needs to + * know about. Rendered above the composer so the message the fault referred + * to remains in view just above. ``role="alert"`` + ``aria-live="assertive"`` + * ensures screen readers announce the failure. + */ +export function StreamErrorNotice({ error, onDismiss }: StreamErrorNoticeProps) { + const { t } = useTranslation(); + + const { title, body } = resolveCopy(error, t); + + return ( +
+ +
+

{title}

+

{body}

+
+ +
+ ); +} + +function resolveCopy( + error: StreamError, + t: (key: string) => string, +): { title: string; body: string } { + switch (error.kind) { + case "message_too_big": + return { + title: t("errors.messageTooBig.title"), + body: t("errors.messageTooBig.body"), + }; + default: { + // Exhaustiveness guard: if a new StreamError kind is added, TS will + // complain here until we add a corresponding i18n branch. + const _exhaustive: never = error.kind; + return { title: String(_exhaustive), body: "" }; + } + } +} diff --git a/webui/src/components/thread/ThreadComposer.tsx b/webui/src/components/thread/ThreadComposer.tsx index 35350fef..9b38d903 100644 --- a/webui/src/components/thread/ThreadComposer.tsx +++ b/webui/src/components/thread/ThreadComposer.tsx @@ -1,12 +1,43 @@ -import { useCallback, useEffect, useRef, useState } from "react"; -import { ArrowUp } from "lucide-react"; +import { + useCallback, + useEffect, + useMemo, + useRef, + useState, + type KeyboardEvent as ReactKeyboardEvent, +} from "react"; +import { + ArrowUp, + ImageIcon, + Loader2, + Paperclip, + X, +} from "lucide-react"; import { useTranslation } from "react-i18next"; import { Button } from "@/components/ui/button"; +import { + useAttachedImages, + type AttachedImage, + type AttachmentError, + MAX_IMAGES_PER_MESSAGE, +} from "@/hooks/useAttachedImages"; +import { useClipboardAndDrop } from "@/hooks/useClipboardAndDrop"; +import type { SendImage } from "@/hooks/useNanobotStream"; import { cn } from "@/lib/utils"; +/** ````: aligned with the server's MIME whitelist. SVG is + * deliberately excluded to avoid an embedded-script XSS surface. */ +const ACCEPT_ATTR = "image/png,image/jpeg,image/webp,image/gif"; + +function formatBytes(n: number): string { + if (n < 1024) return `${n} B`; + if (n < 1024 * 1024) return `${(n / 1024).toFixed(1)} KB`; + return `${(n / (1024 * 1024)).toFixed(1)} MB`; +} + interface ThreadComposerProps { - onSend: (content: string) => void; + onSend: (content: string, images?: SendImage[]) => void; disabled?: boolean; placeholder?: string; modelLabel?: string | null; @@ -22,11 +53,47 @@ export function ThreadComposer({ }: ThreadComposerProps) { const { t } = useTranslation(); const [value, setValue] = useState(""); + const [inlineError, setInlineError] = useState(null); const textareaRef = useRef(null); + const fileInputRef = useRef(null); + const chipRefs = useRef(new Map()); const isHero = variant === "hero"; const resolvedPlaceholder = placeholder ?? t("thread.composer.placeholderThread"); + const { images, enqueue, remove, clear, encoding, full } = + useAttachedImages(); + + const formatRejection = useCallback( + (reason: AttachmentError): string => { + const key = `thread.composer.imageRejected.${reason}`; + return t(key, { max: MAX_IMAGES_PER_MESSAGE }); + }, + [t], + ); + + const addFiles = useCallback( + (files: File[]) => { + if (files.length === 0) return; + const { rejected } = enqueue(files); + if (rejected.length > 0) { + setInlineError(formatRejection(rejected[0].reason)); + } else { + setInlineError(null); + } + }, + [enqueue, formatRejection], + ); + + const { + isDragging, + onPaste, + onDragEnter, + onDragOver, + onDragLeave, + onDrop, + } = useClipboardAndDrop(addFiles); + useEffect(() => { if (disabled) return; const el = textareaRef.current; @@ -35,11 +102,43 @@ export function ThreadComposer({ return () => cancelAnimationFrame(id); }, [disabled]); + const readyImages = useMemo( + () => images.filter((img): img is AttachedImage & { dataUrl: string } => + img.status === "ready" && typeof img.dataUrl === "string", + ), + [images], + ); + const hasErrors = images.some((img) => img.status === "error"); + + const canSend = + !disabled + && !encoding + && !hasErrors + && (value.trim().length > 0 || readyImages.length > 0); + const submit = useCallback(() => { + if (!canSend) return; const trimmed = value.trim(); - if (!trimmed || disabled) return; - onSend(trimmed); + // Share the same normalized ``data:`` URL with both the wire payload and + // the optimistic bubble preview: data URLs are self-contained (no blob + // lifetime, safe under React StrictMode double-mount) and keep the + // bubble in sync with whatever the backend actually sees. + const payload: SendImage[] | undefined = + readyImages.length > 0 + ? readyImages.map((img) => ({ + media: { + data_url: img.dataUrl, + name: img.file.name, + }, + preview: { url: img.dataUrl, name: img.file.name }, + })) + : undefined; + onSend(trimmed, payload); setValue(""); + setInlineError(null); + // Bubble owns the data URL copy; safe to revoke every staged blob + // preview here without affecting the rendered message. + clear(); requestAnimationFrame(() => { const el = textareaRef.current; if (el) { @@ -47,9 +146,9 @@ export function ThreadComposer({ el.focus(); } }); - }, [disabled, onSend, value]); + }, [canSend, clear, onSend, readyImages, value]); - const onKeyDown: React.KeyboardEventHandler = (e) => { + const onKeyDown = (e: ReactKeyboardEvent) => { if (e.key === "Enter" && !e.shiftKey && !e.nativeEvent.isComposing) { e.preventDefault(); submit(); @@ -62,12 +161,55 @@ export function ThreadComposer({ el.style.height = `${Math.min(el.scrollHeight, 260)}px`; }; + const onFilePick: React.ChangeEventHandler = (e) => { + const files = Array.from(e.target.files ?? []); + e.target.value = ""; + addFiles(files); + }; + + const removeChip = useCallback( + (id: string) => { + const { nextFocusId } = remove(id); + setInlineError(null); + requestAnimationFrame(() => { + const el = nextFocusId ? chipRefs.current.get(nextFocusId) : null; + if (el) { + el.focus(); + } else { + textareaRef.current?.focus(); + } + }); + }, + [remove], + ); + + const onChipKey = useCallback( + (id: string) => (e: ReactKeyboardEvent) => { + if ( + e.key === "Delete" || + e.key === "Backspace" || + e.key === "Enter" || + e.key === " " + ) { + e.preventDefault(); + removeChip(id); + } + }, + [removeChip], + ); + + const attachButtonDisabled = disabled || full; + return (
{ e.preventDefault(); submit(); }} + onDragEnter={onDragEnter} + onDragOver={onDragOver} + onDragLeave={onDragLeave} + onDrop={onDrop} className={cn("w-full", isHero ? "px-0" : "px-1 pb-1.5 pt-1 sm:px-0")} >
+ {images.length > 0 ? ( +
+ {images.map((img) => ( + + t("thread.composer.normalizedSizeHint", { + orig: formatBytes(orig), + current: formatBytes(current), + }) + } + formatError={formatRejection} + onRemove={() => removeChip(img.id)} + onKeyDown={onChipKey(img.id)} + registerRef={(el) => { + if (el) chipRefs.current.set(img.id, el); + else chipRefs.current.delete(img.id); + }} + /> + ))} +
+ ) : null}