yaoyao/integrations/hermes-bots-bridge/context_projection.py

"""Request-only Bot context selection; the durable Hermes transcript stays intact."""
from __future__ import annotations

import hashlib
import json
import os
import tempfile
import threading
from pathlib import Path
from types import MethodType


def _hash(text):
    return hashlib.sha256(text.encode('utf-8')).hexdigest()


def _atomic(path, text):
    fd, tmp = tempfile.mkstemp(dir=path.parent, prefix='.context-')
    try:
        with os.fdopen(fd, 'w', encoding='utf-8') as out:
            os.fchmod(out.fileno(), 0o600)
            out.write(text)
        os.replace(tmp, path)
    finally:
        Path(tmp).unlink(missing_ok=True)


class BotContext:
    def __init__(self, home, session_id):
        self.root = Path(home) / 'cache' / 'yaoyao-bot-context' / _hash(session_id)
        self.root.mkdir(parents=True, exist_ok=True, mode=0o700)
        self.lock = threading.RLock()
        self.prompts = {}
        path = self.root / 'prompts.json'
        if path.exists():
            data = json.loads(path.read_text(encoding='utf-8'))
            if not isinstance(data, dict) or any(not isinstance(k, str) or not isinstance(v, str) for k, v in data.items()):
                raise ValueError('Invalid Bot context provenance')
            self.prompts = data
        self.saved = set()
        self.last_stats = {}

    def register(self, prompt, history_text):
        # Only the authenticated bridge may declare the historical form. Never
        # parse section titles/markers out of user-authored or legacy messages.
        if not isinstance(prompt, str) or not isinstance(history_text, str) or not prompt or not history_text:
            raise ValueError('Invalid Bot context envelope')
        if len(prompt) > 2_000_000 or len(history_text) > 2_000_000:
            raise ValueError('Bot context envelope too large')
        key = _hash(prompt)
        with self.lock:
            if key in self.prompts and self.prompts[key] != history_text:
                raise ValueError('Conflicting Bot context provenance')
            updated = {**self.prompts, key: history_text}
            # The retained Hermes tail is bounded; old unknown rows fail open.
            updated = dict(list(updated.items())[-2048:])
            _atomic(self.root / 'prompts.json', json.dumps(updated, ensure_ascii=False))
            self.prompts = updated

    def _archive(self, text):
        key = _hash(text)
        if key not in self.saved:
            path = self.root / (key + '.txt')
            if not path.exists():
                _atomic(path, text)
            self.saved.add(key)
        return key

    def read(self, ref, offset=0, limit=6000):
        if not isinstance(ref, str) or len(ref) != 64 or any(c not in '0123456789abcdef' for c in ref):
            raise ValueError('Invalid history reference')
        if type(offset) is not int or offset < 0 or type(limit) is not int or not 1 <= limit <= 12000:
            raise ValueError('Invalid history range')
        text = (self.root / (ref + '.txt')).read_text(encoding='utf-8')
        return {'ref': ref, 'text': text[offset:offset + limit], 'totalChars': len(text),
                'nextOffset': offset + limit if offset + limit < len(text) else None}

    def select(self, messages, incoming_message=None):
        with self.lock:
            # Leave the latest full Bot prompt in place, including its current
            # authorization, environment, task, facts and attachment boundaries.
            matched = [(i, _hash(m['content'])) for i, m in enumerate(messages)
                       if isinstance(m, dict) and m.get('role') == 'user'
                       and isinstance(m.get('content'), str) and _hash(m['content']) in self.prompts]
            latest = matched[-1][0] if matched else -1
            incoming = incoming_message.get('content') if isinstance(incoming_message, dict) else None
            incoming_key = _hash(incoming) if isinstance(incoming, str) else None
            if incoming_key in self.prompts and matched and matched[-1][1] != incoming_key:
                # Native ephemeral context may extend the current wire content.
                # Its original incoming message still proves that a newer full
                # policy is present; do not retain the previous policy as well.
                latest = -1
            changes = {}
            for i, key in matched:
                if i != latest:
                    changes[i] = self.prompts[key]
            users = [i for i, m in enumerate(messages) if isinstance(m, dict) and m.get('role') == 'user']
            # Every result in the two newest user turns stays byte-identical.
            cutoff = users[-2] if len(users) >= 2 else 0
            seen = set()
            for i in range(len(messages) - 1, -1, -1):
                m = messages[i]
                if not isinstance(m, dict) or m.get('role') != 'tool' or not isinstance(m.get('content'), str):
                    continue
                text = m['content']
                identity = (m.get('name') or m.get('tool_name'), _hash(text))
                duplicate = identity in seen and len(text) > 1500
                seen.add(identity)
                if i >= cutoff or (len(text) <= 8000 and not duplicate):
                    continue
                ref = self._archive(text)
                notice = (f'\n[Earlier tool result: {len(text)} characters. '
                          f'Full original: yaoyao_history_read(ref="{ref}", offset=0).]\n')
                changes[i] = notice if duplicate else text[:2500] + notice + text[-1000:]
            result = [{**m, 'content': changes[i]} if i in changes else m for i, m in enumerate(messages)]
            prompt_replacements = len([i for i, _ in matched if i != latest])
            self.last_stats = {'promptReplacements': prompt_replacements,
                               'toolReplacements': len(changes) - prompt_replacements,
                               'savedChars': sum(len(messages[i]['content']) - len(text) for i, text in changes.items())}
            return result


def install(agent, home, session_id, active):
    existing = getattr(agent, '_yaoyao_bot_context', None)
    if existing is not None and getattr(agent, '_yaoyao_bot_context_engine', None) is getattr(agent, 'context_compressor', None):
        return existing
    engine = getattr(agent, 'context_compressor', None)
    original = getattr(engine, 'select_context', None)
    if not callable(original):
        raise ValueError('Hermes context selection is unavailable')
    context = existing or BotContext(home, session_id)

    def select_context(self, request_messages, **kwargs):
        selected = original(request_messages, **kwargs)
        source = selected if isinstance(selected, list) else request_messages
        return context.select(source, kwargs.get('incoming_message')) if active() else selected

    engine.select_context = MethodType(select_context, engine)
    agent._yaoyao_bot_context = context
    agent._yaoyao_bot_context_engine = engine
    return context