hermes_state_local.py

"""Private local creation receipts in the canonical runtime transaction owner."""
import json
import threading
import time
from collections import OrderedDict

from hermes_state_runtime import RuntimeStoreError, _epoch, _json

POLICY_PREFIX = 'gateway.local_policy.v1:'
API_BINDING_PREFIX = 'gateway.api.binding.v1.'
API_DECLARED_PREFIX = 'gateway.api.conversation.v1.'


def commit_local_session(db, *, epoch, receipt):
    """Reserve the row, route and immutable policy together, before publication. A staged receipt's
    ``on_create`` (owner memory, never stored) is published by the transaction that creates the row
    and withdrawn if that transaction fails: an exact retry, a refused or a losing attempt adds none."""
    on_create = getattr(receipt, 'on_create', None)
    request = getattr(receipt, 'client_request', None)  # the staged request's own client_env, never stored
    if request is None:
        if receipt['policy'].get('client_env_json') is not None:
            # Lost (a copy): never judge a retry by this attempt's own filtered map instead.
            raise RuntimeStoreError('invalid_params')
        request = {}
    if on_create is not None:
        on_create.read()  # its file read: in the caller's thread, outside the write transaction
    encoded = _json(receipt)
    receipt = json.loads(encoded)
    sid, route = receipt['session_id'], receipt['route']
    key = POLICY_PREFIX + sid
    def write(conn):
        _epoch(conn, epoch)
        old = conn.execute('SELECT value FROM state_meta WHERE key=?', (key,)).fetchone()
        if old is not None:
            try:
                saved = json.loads(old[0])
                same = all(saved[k] == receipt[k] for k in (
                    'profile_id', 'principal_id', 'request_id', 'session_id', 'route'))
                same = same and same_create_retry(saved['policy'], receipt['policy']['request_json'], request)
            except (ValueError, TypeError, KeyError) as exc:
                raise RuntimeStoreError('storage_unavailable') from exc
            if not same:
                raise RuntimeStoreError('invalid_params')
            return saved
        from hermes_state_mutation_retirement import RETIRED_PREFIX
        if conn.execute('SELECT 1 FROM state_meta WHERE key=?', (RETIRED_PREFIX + sid,)).fetchone():
            # Deleted: its receipt was retired with the transcript; a create/cron retry must not
            # resurrect the row (and re-fire its prompt) under the same deterministic id.
            raise RuntimeStoreError('not_found')
        if ('legacy_session_id' not in receipt and
                conn.execute('SELECT 1 FROM sessions WHERE id=?', (sid,)).fetchone()):
            raise RuntimeStoreError('storage_unavailable')
        if conn.execute("SELECT 1 FROM gateway_routing WHERE scope='' AND session_key=?", (route,)).fetchone():
            raise RuntimeStoreError('admission_conflict')
        if 'legacy_session_id' in receipt:
            from hermes_state_local_migration import bind_legacy_target
            bind_legacy_target(db, conn, receipt)
        policy, entry = receipt['policy'], receipt['entry']
        from datetime import datetime
        started = datetime.fromisoformat(entry['created_at']).timestamp()
        # Auxiliary usage inserts created_source='unknown' when it creates the row
        # (record_auxiliary_usage carries no surface). The creator upsert repairs
        # 'unknown', but this DO NOTHING insert does not, so the surface is stamped here.
        # The agent's own accounting passes its surface (turn_usage._agent_session_source)
        # and is not that writer. gui's in-process agent is desktop; a managed worker
        # creates with the policy label gui.
        from gateway.session_policy import created_surface
        created = created_surface(policy['source'])
        conn.execute('''INSERT INTO sessions(id,source,created_source,user_id,session_key,chat_id,chat_type,
            model,cwd,profile_name,origin_json,started_at) VALUES(?,?,?,?,?,?,'dm',?,?,?,?,?)
            ON CONFLICT(id) DO NOTHING''',
            (sid, policy['source'], created, receipt['principal_id'], route,
             entry['origin']['chat_id'], policy['model'], _row_cwd(policy), db._own_profile_name(),
             _json(entry['origin']), started))
        conn.execute("INSERT INTO gateway_routing(scope,session_key,entry_json,updated_at) VALUES('',?,?,?)",
                     (route, _json(entry), started))
        conn.execute('INSERT INTO state_meta(key,value) VALUES(?,?)', (key, encoded))
        if on_create is not None:
            on_create.publish()  # before the row is visible, so no reader lazy-loads around it
        return receipt
    try:
        return db._execute_write(write)
    except BaseException:
        if on_create is not None:
            on_create.withdraw()  # only this attempt's own snapshot, never another winner's
        raise


def client_request(request_json, derived=None):
    """A frozen create request as the client sent it: an owner-resolved ``-m`` route (``model_requested``)
    compares by the name the client typed, not by the model and the *derived* fields (a receipt from
    before ``route_derived``: the endpoint) its ``${VAR}`` refs expanded to at the time. ``desktop_local``
    (``gateway.session_policy.LOCAL_DESKTOP_CHAT``), which an earlier owner stamped into the request, was
    never sent."""
    request = json.loads(request_json)
    request.pop('desktop_local', None)
    if derived is None:
        request.pop('route_derived', None)
    if 'model_requested' in request:
        request.pop('model', None)
        for key in ('base_url',) if derived is None else derived:
            request.pop(key, None)
    return request


def same_create_request(saved_json, retry_json):
    """Whether a create retry carries the frozen request it resumes: one check for a retry that finds the
    session committed and one staged before it was. The saved receipt's provenance names the owner's
    fields; one from before a same-name alias got ``model_requested`` compares exactly, as it did."""
    derived = json.loads(saved_json).get('route_derived')
    if client_request(saved_json, derived) == client_request(retry_json, derived):
        return True
    retry = json.loads(retry_json)
    if retry.get('model_requested') == retry.get('model'):
        retry.pop('model_requested', None)
        retry.pop('route_derived', None)
    return json.loads(saved_json) == retry


def same_create_retry(saved_policy, retry_json, client_request):
    """The one exact-retry identity of a create: ``same_create_request`` on what the client sent, and its
    ``client_env`` (``client_request``: the retry's own, never stored) judged by the rules frozen with
    *saved_policy*. Used by a retry that finds the session committed, at commit, and by one overlapping
    the attempt that holds the session's key slot."""
    from agent.session_client_env import retry_matches
    return same_create_request(saved_policy['request_json'], retry_json) and retry_matches(saved_policy, client_request)


def _row_cwd(policy):
    """The sessions-row ``cwd``: the create request's ``row_cwd`` (``""`` = none) when it sent one,
    else the execution cwd. Row metadata only (workspace grouping, ``-c``, cwd restore); the
    frozen policy's ``cwd`` / ``TERMINAL_CWD`` stay where the session's tools run."""
    try:
        request = json.loads(policy.get('request_json') or '{}')
    except (TypeError, ValueError):
        request = {}
    value = request.get('row_cwd') if isinstance(request, dict) else None
    return policy['cwd'] if not isinstance(value, str) else (value or None)


def _policy_target(raw):
    """Current physical target a creation receipt routes to; None for a malformed receipt, which
    then stays (its own reads already refuse it as ``storage_unavailable``) and cannot abort
    the delete of an unrelated session."""
    try:
        value = json.loads(raw)
    except (TypeError, ValueError):
        return None
    entry = value.get('entry') if isinstance(value, dict) else None
    return entry.get('session_id') if isinstance(entry, dict) else None


def retire_local_receipts(conn, session_ids):
    """Deletion fence half for owner receipts. A local creation policy carries the launch request
    (``request_json``: a cron fire's job and raw prompt), so it goes exactly when its route does:
    when the receipt's current physical target is deleted (a pruned reset ancestor whose logical
    session lives on in a child keeps it). An API binding and the declared conversation key that
    resolves to it go with their id, so that key starts a fresh session."""
    targets = set(session_ids)
    if not targets:
        return
    # Prefix scans (GLOB uses the key index), not one LIKE per id: a prune retires thousands.
    gone = [key for key, value in conn.execute('SELECT key,value FROM state_meta WHERE key GLOB ?',
                                               (POLICY_PREFIX + '*',))
            if _policy_target(value) in targets]
    gone += [key for key, value in conn.execute('SELECT key,value FROM state_meta WHERE key GLOB ?',
                                                (API_DECLARED_PREFIX + '*',)) if value in targets]
    gone += [API_BINDING_PREFIX + sid for sid in targets]
    conn.executemany('DELETE FROM state_meta WHERE key=?', [(key,) for key in gone])


def moved_receipts(conn, target_ids=None):
    """``{owner id: (moved row id, its source)}`` of local receipts whose current physical target row
    now carries another lane's route. An explicit repoint (``/handoff``, a platform ``/resume``)
    rewrote that row's routing identity, so the receipt can never restore the conversation again
    (the stored identity check in ``gateway.session_local_recovery._checked_entry``). Only the
    current target counts: an earlier ``/new`` segment that moved leaves the conversation the
    receipt guards local, and releasing it would strand that transcript. *target_ids* narrows the
    scan to the rows one repoint rewrote."""
    rows = conn.execute(
        "SELECT substr(m.key, ?), s.id, s.source FROM state_meta m JOIN sessions s "
        "ON s.id = CASE WHEN json_valid(m.value) THEN json_extract(m.value, '$.entry.session_id') END "
        "WHERE m.key GLOB ? AND s.session_key IS NOT NULL "
        "AND s.session_key IS NOT CASE WHEN json_valid(m.value) THEN json_extract(m.value, '$.route') END",
        (len(POLICY_PREFIX) + 1, POLICY_PREFIX + '*')).fetchall()
    return {owner: (target, source) for owner, target, source in rows
            if target_ids is None or target in target_ids}


def release_receipts(conn, owners):
    """Retire *owners*' local receipts (creation policy and legacy binding): the row they guarded now
    belongs to a messaging lane, which restores it from its route like any other platform session."""
    from hermes_state_local_migration import LEGACY_PREFIX
    conn.executemany('DELETE FROM state_meta WHERE key IN (?,?)',
                     [(POLICY_PREFIX + owner, LEGACY_PREFIX + owner) for owner in owners])


def handoff_head(db, session_id):
    """``(local owner or None, row)`` a ``/handoff`` of *session_id* moves: the conversation the user
    is looking at. A local conversation continues from its receipt's current target (its latest
    ``/new`` segment, then that segment's compression tip), never from an earlier segment; any other
    row from its compression tip."""
    with db._read_ctx() as conn:
        owner = local_lineage_owner(conn, session_id)
        saved = conn.execute('SELECT value FROM state_meta WHERE key=?', (POLICY_PREFIX + owner,)).fetchone()
    target = _policy_target(saved[0]) if saved else None
    start = target or session_id
    return (owner if target else None), db.get_compression_tip(start) or start


_LANE_IDENTITY = 'session_key, source, user_id, chat_id, chat_type, thread_id, display_name, origin_json, transport_profile'


def hand_off_conversations(conn, moved):
    """Retire the receipts in *moved* (``moved_receipts``: a messaging lane took over each one's
    current target) and move the rest of each conversation onto that lane with it. Its earlier
    ``/new`` segments take the lane's identity, so they open there as read-only history (a lane
    moved past them) instead of staying local rows no receipt restores; the conversation's name
    follows its head, which the lane continues as its own conversation. Returns the receipts it
    handed off: one whose lineage is malformed stays as it is (its own reads refuse it, logged)."""
    handed = {}
    for owner, (target, source) in moved.items():
        saved = conn.execute('SELECT value FROM state_meta WHERE key=?', (POLICY_PREFIX + owner,)).fetchone()
        lineage = (_receipt(saved[0]) or {}).get('lineage', [owner]) if saved else [owner]
        if not isinstance(lineage, list) or not all(isinstance(sid, str) for sid in lineage):
            continue
        segments = list(dict.fromkeys([owner, *lineage]))
        conn.execute(f'UPDATE sessions SET ({_LANE_IDENTITY}) = (SELECT {_LANE_IDENTITY} FROM sessions WHERE id=?) '
                     'WHERE id IN (SELECT value FROM json_each(?)) AND id != ?', (target, json.dumps(segments), target))
        _title_follows_head(conn, owner, _segment_root(conn, segments, target))
        handed[owner] = (target, source)
    release_receipts(conn, handed)
    return handed


def _receipt(raw):
    try:
        value = json.loads(raw)
    except (TypeError, ValueError):
        return None
    return value if isinstance(value, dict) else None


def _segment_root(conn, segments, target):
    """The first row of *target*'s ``/new`` segment: back through the compressions behind it."""
    if target not in segments:
        return target
    index = segments.index(target)
    while index > 0 and (row := conn.execute('SELECT end_reason FROM sessions WHERE id=?',
                                               (segments[index - 1],)).fetchone()) and row[0] == 'compression':
        index -= 1
    return segments[index]


def _title_follows_head(conn, owner, head):
    """The conversation's name (a Bot Chat's registry title included) moves to its head, over a
    title of lower provenance there (a generated one); a name of equal rank the head carries stays
    (``SessionDB._title_rank``: a pre-provenance title ranks as the user's)."""
    from hermes_state import SessionDB
    row = conn.execute('SELECT title, title_source, hidden FROM sessions WHERE id=?', (owner,)).fetchone()
    if head == owner or row is None or not row[0]:
        return
    kept = conn.execute('SELECT title, title_source FROM sessions WHERE id=?', (head,)).fetchone()
    if kept is None or (kept[0] and SessionDB._title_rank(kept[1]) >= SessionDB._title_rank(row[1])):
        return
    conn.execute('UPDATE sessions SET title=NULL WHERE id IN (?,?)', (owner, head))
    conn.execute('UPDATE sessions SET title=?, title_source=?, hidden=? WHERE id=?', (*row, head))


def local_lineage_owner(conn, session_id):
    """The creation id whose receipt lineage holds physical segment *session_id* (a reset or
    compression child, which carries the owner as ``chat_id``); *session_id* itself otherwise."""
    row = conn.execute('SELECT chat_id FROM sessions WHERE id=?', (session_id,)).fetchone()
    owner = row[0] if row else None
    if not owner or owner == session_id:
        return session_id
    saved = conn.execute('SELECT value FROM state_meta WHERE key=?', (POLICY_PREFIX + owner,)).fetchone()
    try:
        lineage = json.loads(saved[0]).get('lineage') if saved else None
    except (TypeError, ValueError, AttributeError):
        return session_id
    return owner if isinstance(lineage, list) and session_id in lineage else session_id


def owned_lineage_ids(conn, session_id):
    """[session_id] for any row outside a local reset/compression lineage, else the whole
    conversation (*session_id* first). Any segment deletes as all of it: the owner's policy goes
    with its current target, and an earlier segment's gap breaks the lineage, so a lone segment
    delete strands the rest (and the owner's queued work) readable but never continuable."""
    return list(dict.fromkeys([session_id, *local_conversation_ids(conn, session_id)]))


def local_conversation_ids(conn, session_id):
    """Every row of the local conversation *session_id* belongs to: its creation id first, then
    each segment of the receipt lineage. A canonical reset never forks the logical session (one
    policy, FIFO and generation), so listing, delete and retirement treat these ids as one."""
    owner = local_lineage_owner(conn, session_id)
    saved = conn.execute('SELECT value FROM state_meta WHERE key=?', (POLICY_PREFIX + owner,)).fetchone()
    try:
        lineage = json.loads(saved[0]).get('lineage') if saved else None
    except (TypeError, ValueError, AttributeError):
        lineage = None
    if not isinstance(lineage, list):
        return [session_id]
    return list(dict.fromkeys([owner, *(sid for sid in lineage if isinstance(sid, str))]))


def local_lineage_index(conn):
    """``(superseded, heads)`` for listing a local conversation as one row, as
    ``acp_adapter.catalog`` does: a canonical reset keeps the creation id as the owner and moves the
    transcript to a child, so every segment is one conversation. ``superseded`` holds the listable
    roots of earlier reset segments (filtered in SQL, before paging); ``heads`` maps the latest
    reset segment to ``(owner, lineage)`` for the row that represents the conversation. A delete
    of that row removes the whole lineage (``hermes_state_mutation_guards.delete_targets``)."""
    superseded, heads = set(), {}
    rows = conn.execute("SELECT value FROM state_meta WHERE key GLOB ? AND json_valid(value) "
                        "AND json_array_length(value, '$.lineage') > 1", (POLICY_PREFIX + '*',)).fetchall()
    for (raw,) in rows:
        receipt = json.loads(raw)
        owner, lineage = receipt.get('session_id'), receipt.get('lineage')
        if not isinstance(owner, str) or not all(isinstance(sid, str) for sid in lineage):
            continue
        marks = ','.join('?' * len(lineage))
        ended = dict(conn.execute(f'SELECT id,end_reason FROM sessions WHERE id IN ({marks})', lineage).fetchall())
        starts = [sid for i, sid in enumerate(lineage) if i == 0 or ended.get(lineage[i - 1]) == 'session_reset']
        superseded.update(starts[:-1])
        heads[starts[-1]] = (owner, lineage)
    return superseded, heads


def _lineage_version(conn):
    # data_version sees OTHER connections' commits; total_changes sees this connection's writes.
    return conn.total_changes, conn.execute('PRAGMA main.data_version').fetchone()[0]


def _copy_lineage_index(value):
    superseded, heads = value
    return set(superseded), {key: (owner, list(lineage)) for key, (owner, lineage) in heads.items()}


def _cached_local_lineage_index(db, conn):
    # Owned by this SessionDB, not a process-global path/id cache. Bound retained connections
    # across read-pool churn and close/reopen; a closed/replaced connection cannot match a new one.
    cache, lock = db.__dict__.setdefault('_local_lineage_cache', (OrderedDict(), threading.Lock()))
    if conn.in_transaction:
        # A snapshot or uncommitted answer must never outlive its transaction (including rollback).
        with lock:
            cache.pop(conn, None)
        return local_lineage_index(conn)
    version = _lineage_version(conn)
    with lock:
        saved = cache.get(conn)
        if saved is not None and saved[0] == version:
            cache.move_to_end(conn)
            return _copy_lineage_index(saved[1])
    value = local_lineage_index(conn)
    # A concurrent writer may commit between the receipt scan and the end_reason reads. Never
    # publish that in-flight projection as a stable cache entry; the next call must rebuild it.
    if not conn.in_transaction and _lineage_version(conn) == version:
        with lock:
            cache[conn] = (version, value)
            cache.move_to_end(conn)
            if len(cache) > 8:
                cache.popitem(last=False)
    return _copy_lineage_index(value)


def exclude_superseded_segments(db, where_clauses, params):
    """Drop earlier reset segments from a listing or count query (one conversation, one row),
    in SQL before paging; returns ``heads`` for :func:`annotate_local_lineages`."""
    with db._read_ctx() as conn:
        superseded, heads = _cached_local_lineage_index(db, conn)
    if superseded:
        where_clauses.append("s.id NOT IN (SELECT value FROM json_each(?))")
        params.append(json.dumps(sorted(superseded)))
    return heads


def annotate_local_lineages(sessions, heads):
    """Name every segment on the conversation's representative row (``_lineage_ids``, which the
    Desktop matches deletes, tombstones and stored ids against). ``_lineage_root_id`` stays the
    compression root: pins and archive flags live on the listed segment's own chain."""
    for row in sessions:
        found = heads.get(row.get('_lineage_root_id') or row['id'])
        if found is not None:
            row['_lineage_ids'] = list(dict.fromkeys([*found[1], *(row.get('_lineage_ids') or ())]))
    return sessions


def local_receipt(db, session_id):
    with db._read_ctx() as conn:
        row = conn.execute('SELECT value FROM state_meta WHERE key=?', (POLICY_PREFIX + session_id,)).fetchone()
    if row is None:
        raise RuntimeStoreError('storage_unavailable')
    try:
        value = json.loads(row[0])
        if not isinstance(value, dict):
            raise ValueError('invalid receipt')
        return value
    except (TypeError, ValueError) as exc:
        raise RuntimeStoreError('storage_unavailable') from exc


def end_idle_local_session(db, *, epoch, session_id, target_id, reason):
    """Stamp ``ended_at`` on *target_id* only while the logical *session_id*'s FIFO is idle.

    The idle check and the stamp share one owner write txn, so an admission committed first keeps
    the row open; one committed after is reopened by the drain (``reopen_local_session``)."""
    def write(conn):
        _epoch(conn, epoch)
        if conn.execute("SELECT 1 FROM session_admissions WHERE target_session_id=? AND status IN ('queued','started','unknown')",
                        (session_id,)).fetchone():
            return 0
        if conn.execute("SELECT 1 FROM worker_executions WHERE session_id IN (?,?) AND status IN ('registered','running','unknown')",
                        (session_id, target_id)).fetchone():
            return 0
        return db._end_and_bump(conn, 'UPDATE sessions SET ended_at=?, end_reason=? WHERE id=? AND ended_at IS NULL',
                                (time.time(), reason, target_id), target_id, reason)
    return db._execute_write(write)