strix/strix/core/sessions.py

"""SDK session helpers for Strix agents."""

from __future__ import annotations

import asyncio
import logging
import sqlite3
from contextlib import contextmanager
from typing import TYPE_CHECKING, Any, cast
from weakref import WeakKeyDictionary

from agents.items import ItemHelpers
from agents.memory import SQLiteSession


if TYPE_CHECKING:
    from collections.abc import Callable, Iterator
    from pathlib import Path

    from agents.items import TResponseInputItem
    from agents.memory import Session


logger = logging.getLogger(__name__)


class _PooledConnectionSession(SQLiteSession):
    @contextmanager
    def _locked_connection(self) -> Iterator[sqlite3.Connection]:
        with self._lock:
            if self._closed:
                raise RuntimeError("SQLiteSession is closed")
            if self._is_memory_db:
                yield self._shared_connection
                return
            connection = sqlite3.connect(str(self.db_path), check_same_thread=False)
            try:
                yield connection
            finally:
                connection.close()


def open_agent_session(agent_id: str, path: Path) -> SQLiteSession:
    path.parent.mkdir(parents=True, exist_ok=True)
    return _PooledConnectionSession(session_id=agent_id, db_path=path)


async def seed_initial_input(session: Session, initial_input: Any) -> bool:
    """Commit an agent's opening identity/task input before its first run cycle."""
    items = ItemHelpers.input_to_new_input_list(initial_input)
    if not items:
        return False
    async with session_write_lock(session):
        if await session.get_items():
            return False
        await session.add_items(items)
    return True


_IMAGE_REJECTED_TEXT = "[image rejected by the model]"
_IMAGE_ELIDED_TEXT = "[older screenshot elided to bound context memory]"
_INHERITED_IMAGE_TEXT = "[screenshot omitted from inherited context]"


def _output_has_image(item_dict: dict[str, Any]) -> bool:
    return (
        item_dict.get("type") == "function_call_output"
        and isinstance(item_dict.get("output"), list)
        and any(isinstance(b, dict) and b.get("type") == "input_image" for b in item_dict["output"])
    )


def _elided_output(item_dict: dict[str, Any], text: str) -> dict[str, Any]:
    # Replace only image blocks; sibling text blocks are preserved.
    output = item_dict.get("output")
    blocks = output if isinstance(output, list) else []
    return {
        "type": "function_call_output",
        "call_id": item_dict.get("call_id"),
        "output": [
            {"type": "input_text", "text": text}
            if isinstance(block, dict) and block.get("type") == "input_image"
            else block
            for block in blocks
        ],
    }


_session_write_locks: WeakKeyDictionary[Session, asyncio.Lock] = WeakKeyDictionary()


def session_write_lock(session: Session) -> asyncio.Lock:
    """Lock serialising all out-of-band writes to ``session``."""
    lock = _session_write_locks.get(session)
    if lock is None:
        lock = asyncio.Lock()
        _session_write_locks[session] = lock
    return lock


async def _rewrite_session(
    session: Session,
    transform: Callable[[list[Any]], tuple[list[Any], bool]],
) -> bool:
    """Read-modify-write a session under its write lock, restoring on failure."""
    async with session_write_lock(session):
        items = await session.get_items()
        if not items:
            return False
        rebuilt, changed = transform(list(items))
        if not changed:
            return False
        rebuilt_items = cast("list[TResponseInputItem]", rebuilt)
        original_items = cast("list[TResponseInputItem]", list(items))
        await session.clear_session()
        try:
            await session.add_items(rebuilt_items)
        except Exception:
            logger.exception("session rewrite failed; restoring original items")
            await session.clear_session()
            await session.add_items(original_items)
            raise
        return True


async def replace_session_items(
    session: Session,
    new_items: list[Any],
    *,
    expected_len: int | None = None,
) -> bool:
    """Overwrite the session's items, restoring the originals on failure.

    When ``expected_len`` is given, the rewrite is skipped if the session no
    longer has that many items (a concurrent writer changed it), so a slow
    compaction summary can't clobber newer turns.
    """
    async with session_write_lock(session):
        original = list(await session.get_items())
        if expected_len is not None and len(original) != expected_len:
            logger.warning(
                "skipping session rewrite: expected %d items, found %d",
                expected_len,
                len(original),
            )
            return False
        rebuilt = cast("list[TResponseInputItem]", new_items)
        await session.clear_session()
        try:
            await session.add_items(rebuilt)
        except Exception:
            logger.exception("session rewrite failed; restoring original items")
            await session.clear_session()
            await session.add_items(original)
            raise
        return True


async def strip_all_images_from_session(session: Session) -> bool:
    """Replace every image tool output with a text placeholder (rejection recovery)."""

    def _transform(items: list[Any]) -> tuple[list[Any], bool]:
        rebuilt: list[Any] = []
        changed = False
        for item in items:
            item_dict = cast("dict[str, Any]", item) if isinstance(item, dict) else None
            if item_dict is not None and _output_has_image(item_dict):
                rebuilt.append(_elided_output(item_dict, _IMAGE_REJECTED_TEXT))
                changed = True
            else:
                rebuilt.append(item)
        return rebuilt, changed

    return await _rewrite_session(session, _transform)


async def enforce_image_budget(session: Session, max_images: int) -> bool:
    """Keep only the most recent ``max_images`` image outputs; elide older ones."""
    if max_images < 0:
        return False

    def _transform(items: list[Any]) -> tuple[list[Any], bool]:
        image_indices = [
            i
            for i, item in enumerate(items)
            if isinstance(item, dict) and _output_has_image(cast("dict[str, Any]", item))
        ]
        if len(image_indices) <= max_images:
            return items, False
        to_elide = set(image_indices[: len(image_indices) - max_images])
        rebuilt = [
            _elided_output(cast("dict[str, Any]", item), _IMAGE_ELIDED_TEXT)
            if i in to_elide
            else item
            for i, item in enumerate(items)
        ]
        return rebuilt, True

    return await _rewrite_session(session, _transform)


def scrub_images_from_items(items: list[Any]) -> list[Any]:
    """Return a copy of ``items`` with every image block replaced by text."""

    def _scrub(obj: Any) -> Any:
        if isinstance(obj, dict):
            if obj.get("type") == "input_image":
                return {"type": "input_text", "text": _INHERITED_IMAGE_TEXT}
            return {k: _scrub(v) for k, v in obj.items()}
        if isinstance(obj, list):
            return [_scrub(v) for v in obj]
        return obj

    return [_scrub(item) for item in items]