Open WebUI persisted-history vision replay canary toolkit

vision_replay_canary.py

"""
title: Vision Replay Canary
author: Adam Smith
author_url: https://github.com/rndmcnlly
version: 0.1.0
license: MIT
description: Tests whether a model can recover a Python toolkit image from persisted chat history.
requirements: pillow
"""

import base64
import io
import secrets
from dataclasses import dataclass

from PIL import Image, ImageDraw, ImageFont


@dataclass
class Challenge:
    code: str
    attempted: bool = False


class Tools:
    """Two-turn manual test for model-visible toolkit images in stored history.

    Turn 1 prompt:
      Start a vision replay canary. Call issue_vision_challenge exactly once.
      Inspect the image silently. Do not transcribe, mention, or verify its code.
      Reply exactly READY.

    Turn 2 prompt (send as a fresh user message after Turn 1 finishes):
      Do not generate a new challenge. Recover the code from the earlier tool
      image and call verify_recalled_code exactly once. Report its result.

    Before the replay fix, Turn 2 loses the image and cannot reliably pass.
    After the fix, persisted history supplies the image and the verifier passes.
    """

    def __init__(self):
        self._challenges: dict[tuple[str, str], Challenge] = {}

    @staticmethod
    def _conversation_key(__metadata__: dict | None, __user__: dict | None) -> tuple[str, str]:
        metadata = __metadata__ or {}
        user = __user__ or {}
        return str(user.get('id', 'unknown-user')), str(metadata.get('chat_id', 'unknown-chat'))

    @staticmethod
    def _render_challenge(code: str) -> str:
        # Render small, then scale with nearest-neighbor sampling for large, crisp glyphs.
        image = Image.new('RGB', (96, 34), '#fff7df')
        draw = ImageDraw.Draw(image)
        font = ImageFont.load_default()
        draw.rectangle((1, 1, 94, 32), outline='#b42318', width=2)
        draw.text((7, 5), 'VISION CANARY', fill='#344054', font=font)
        draw.text((29, 18), code, fill='#101828', font=font)
        image = image.resize((768, 272), Image.Resampling.NEAREST)

        buffer = io.BytesIO()
        image.save(buffer, format='PNG', optimize=True)
        encoded = base64.b64encode(buffer.getvalue()).decode('ascii')
        return f'data:image/png;base64,{encoded}'

    def issue_vision_challenge(
        self,
        __metadata__: dict | None = None,
        __user__: dict | None = None,
    ) -> str:
        """Create a one-time visual challenge for a history-replay test.

        The challenge code exists only in the returned image. Call this exactly
        once in Turn 1, inspect the image silently, and follow the user's response
        instructions. Do not call the verifier in the same turn.
        """
        alphabet = '23456789'
        code = ''.join(secrets.choice(alphabet) for _ in range(6))
        key = self._conversation_key(__metadata__, __user__)
        self._challenges[key] = Challenge(code=code)
        return self._render_challenge(code)

    def verify_recalled_code(
        self,
        code: str,
        __metadata__: dict | None = None,
        __user__: dict | None = None,
    ) -> str:
        """Check one recalled vision-canary code without revealing the answer.

        Call this exactly once in Turn 2 with the code recovered from the earlier
        tool image. Never call issue_vision_challenge again before verification.

        :param code: The six-character code read from the earlier image.
        """
        key = self._conversation_key(__metadata__, __user__)
        challenge = self._challenges.get(key)
        if challenge is None:
            return 'CANARY ERROR: No challenge exists for this conversation.'
        if challenge.attempted:
            return 'CANARY ERROR: This challenge has already been attempted.'

        challenge.attempted = True
        normalized = ''.join(character for character in code.upper() if character.isalnum())
        if secrets.compare_digest(normalized, challenge.code):
            return 'CANARY PASS: The model recovered the image code from conversation history.'
        return 'CANARY FAIL: The submitted code did not match the image.'