"""SQLAlchemy-powered Session backend. Usage:: from agents.extensions.memory import SQLAlchemySession # Create from SQLAlchemy URL (uses asyncpg driver under the hood for Postgres) session = SQLAlchemySession.from_url( session_id="user-123", url="postgresql+asyncpg://app:secret@db.example.com/agents", create_tables=True, # If you want to auto-create tables, set to True. ) # Or pass an existing AsyncEngine that your application already manages session = SQLAlchemySession( session_id="user-123", engine=my_async_engine, create_tables=True, # If you want to auto-create tables, set to True. ) await Runner.run(agent, "Hello", session=session) """ from __future__ import annotations import asyncio import json import threading from typing import Any, ClassVar from sqlalchemy import ( TIMESTAMP, Column, ForeignKey, Index, Integer, MetaData, String, Table, Text, delete, event, insert, select, text as sql_text, update, ) from sqlalchemy.exc import IntegrityError, OperationalError from sqlalchemy.ext.asyncio import AsyncEngine, async_sessionmaker, create_async_engine from ...items import TResponseInputItem from ...memory.session import SessionABC from ...memory.session_settings import SessionSettings, resolve_session_limit class SQLAlchemySession(SessionABC): """SQLAlchemy implementation of :pyclass:`agents.memory.session.Session`.""" _table_init_locks: ClassVar[dict[tuple[str, str, str], threading.Lock]] = {} _table_init_locks_guard: ClassVar[threading.Lock] = threading.Lock() _sqlite_configured_engines: ClassVar[set[int]] = set() _sqlite_configured_engines_guard: ClassVar[threading.Lock] = threading.Lock() _SQLITE_BUSY_TIMEOUT_MS: ClassVar[int] = 5000 _SQLITE_LOCK_RETRY_DELAYS: ClassVar[tuple[float, ...]] = (0.05, 0.1, 0.2, 0.4, 0.8) _metadata: MetaData _sessions: Table _messages: Table session_settings: SessionSettings | None = None @classmethod def _get_table_init_lock( cls, engine: AsyncEngine, sessions_table: str, messages_table: str ) -> threading.Lock: lock_key = ( engine.url.render_as_string(hide_password=True), sessions_table, messages_table, ) with cls._table_init_locks_guard: lock = cls._table_init_locks.get(lock_key) if lock is None: lock = threading.Lock() cls._table_init_locks[lock_key] = lock return lock @classmethod def _configure_sqlite_engine(cls, engine: AsyncEngine) -> None: """Apply SQLite settings that reduce transient lock failures.""" if engine.dialect.name != "sqlite": return engine_key = id(engine.sync_engine) with cls._sqlite_configured_engines_guard: if engine_key in cls._sqlite_configured_engines: return @event.listens_for(engine.sync_engine, "connect") def _configure_sqlite_connection(dbapi_connection: Any, _: Any) -> None: cursor = dbapi_connection.cursor() try: cursor.execute(f"PRAGMA busy_timeout = {cls._SQLITE_BUSY_TIMEOUT_MS}") cursor.execute("PRAGMA journal_mode = WAL") finally: cursor.close() cls._sqlite_configured_engines.add(engine_key) @staticmethod def _is_sqlite_lock_error(exc: OperationalError) -> bool: return "database is locked" in str(exc).lower() async def _run_sqlite_write_with_retry(self, operation: Any) -> None: """Retry transient SQLite write lock failures with bounded backoff.""" if self._engine.dialect.name != "sqlite": await operation() return for attempt, delay in enumerate((0.0, *self._SQLITE_LOCK_RETRY_DELAYS)): if delay: await asyncio.sleep(delay) try: await operation() return except OperationalError as exc: if not self._is_sqlite_lock_error(exc): raise if attempt == len(self._SQLITE_LOCK_RETRY_DELAYS): raise def __init__( self, session_id: str, *, engine: AsyncEngine, create_tables: bool = False, sessions_table: str = "agent_sessions", messages_table: str = "agent_messages", session_settings: SessionSettings | None = None, ): """Initializes a new SQLAlchemySession. Args: session_id (str): Unique identifier for the conversation. engine (AsyncEngine): A pre-configured SQLAlchemy async engine. The engine must be created with an async driver (e.g., 'postgresql+asyncpg://', 'mysql+aiomysql://', or 'sqlite+aiosqlite://'). create_tables (bool, optional): Whether to automatically create the required tables and indexes. Defaults to False for production use. Set to True for development and testing when migrations aren't used. sessions_table (str, optional): Override the default table name for sessions if needed. messages_table (str, optional): Override the default table name for messages if needed. session_settings (SessionSettings | None, optional): Session configuration settings """ self.session_id = session_id self.session_settings = session_settings or SessionSettings() self._engine = engine self._configure_sqlite_engine(engine) self._init_lock = ( self._get_table_init_lock(engine, sessions_table, messages_table) if create_tables else None ) self._metadata = MetaData() self._sessions = Table( sessions_table, self._metadata, Column("session_id", String, primary_key=True), Column( "created_at", TIMESTAMP(timezone=False), server_default=sql_text("CURRENT_TIMESTAMP"), nullable=False, ), Column( "updated_at", TIMESTAMP(timezone=False), server_default=sql_text("CURRENT_TIMESTAMP"), onupdate=sql_text("CURRENT_TIMESTAMP"), nullable=False, ), ) self._messages = Table( messages_table, self._metadata, Column("id", Integer, primary_key=True, autoincrement=True), Column( "session_id", String, ForeignKey(f"{sessions_table}.session_id", ondelete="CASCADE"), nullable=False, ), Column("message_data", Text, nullable=False), Column( "created_at", TIMESTAMP(timezone=False), server_default=sql_text("CURRENT_TIMESTAMP"), nullable=False, ), Index( f"idx_{messages_table}_session_time", "session_id", "created_at", ), sqlite_autoincrement=True, ) # Async session factory self._session_factory = async_sessionmaker(self._engine, expire_on_commit=False) self._create_tables = create_tables # --------------------------------------------------------------------- # Convenience constructors # --------------------------------------------------------------------- @classmethod def from_url( cls, session_id: str, *, url: str, engine_kwargs: dict[str, Any] | None = None, session_settings: SessionSettings | None = None, **kwargs: Any, ) -> SQLAlchemySession: """Create a session from a database URL string. Args: session_id (str): Conversation ID. url (str): Any SQLAlchemy async URL, e.g. "postgresql+asyncpg://user:pass@host/db". engine_kwargs (dict[str, Any] | None): Additional keyword arguments forwarded to sqlalchemy.ext.asyncio.create_async_engine. session_settings (SessionSettings | None): Session configuration settings including default limit for retrieving items. If None, uses default SessionSettings(). **kwargs: Additional keyword arguments forwarded to the main constructor (e.g., create_tables, custom table names, etc.). Returns: SQLAlchemySession: An instance of SQLAlchemySession connected to the specified database. """ engine_kwargs = engine_kwargs or {} engine = create_async_engine(url, **engine_kwargs) return cls(session_id, engine=engine, session_settings=session_settings, **kwargs) async def _serialize_item(self, item: TResponseInputItem) -> str: """Serialize an item to JSON string. Can be overridden by subclasses.""" return json.dumps(item, separators=(",", ":")) async def _deserialize_item(self, item: str) -> TResponseInputItem: """Deserialize a JSON string to an item. Can be overridden by subclasses.""" return json.loads(item) # type: ignore[no-any-return] # ------------------------------------------------------------------ # Session protocol implementation # ------------------------------------------------------------------ async def _ensure_tables(self) -> None: """Ensure tables are created before any database operations.""" if not self._create_tables: return assert self._init_lock is not None while not self._init_lock.acquire(blocking=False): # Poll without handing lock acquisition to a background thread so # cancellation cannot strand the shared init lock in the acquired state. await asyncio.sleep(0.01) try: if not self._create_tables: return async with self._engine.begin() as conn: await conn.run_sync(self._metadata.create_all) self._create_tables = False # Only create once finally: self._init_lock.release() async def get_items(self, limit: int | None = None) -> list[TResponseInputItem]: """Retrieve the conversation history for this session. Args: limit: Maximum number of items to retrieve. If None, uses session_settings.limit. When specified, returns the latest N items in chronological order. Returns: List of input items representing the conversation history """ await self._ensure_tables() session_limit = resolve_session_limit(limit, self.session_settings) async with self._session_factory() as sess: if session_limit is None: stmt = ( select(self._messages.c.message_data) .where(self._messages.c.session_id == self.session_id) .order_by( self._messages.c.created_at.asc(), self._messages.c.id.asc(), ) ) else: stmt = ( select(self._messages.c.message_data) .where(self._messages.c.session_id == self.session_id) # Use DESC + LIMIT to get the latest N # then reverse later for chronological order. .order_by( self._messages.c.created_at.desc(), self._messages.c.id.desc(), ) .limit(session_limit) ) result = await sess.execute(stmt) rows: list[str] = [row[0] for row in result.all()] if session_limit is not None: rows.reverse() items: list[TResponseInputItem] = [] for raw in rows: try: items.append(await self._deserialize_item(raw)) except json.JSONDecodeError: # Skip corrupted rows continue return items async def add_items(self, items: list[TResponseInputItem]) -> None: """Add new items to the conversation history. Args: items: List of input items to add to the history """ if not items: return await self._ensure_tables() payload = [ { "session_id": self.session_id, "message_data": await self._serialize_item(item), } for item in items ] async def _write_items() -> None: async with self._session_factory() as sess: async with sess.begin(): # Avoid check-then-insert races on the first write while keeping # the common path free of avoidable integrity exceptions. existing = await sess.execute( select(self._sessions.c.session_id).where( self._sessions.c.session_id == self.session_id ) ) if not existing.scalar_one_or_none(): try: async with sess.begin_nested(): await sess.execute( insert(self._sessions).values({"session_id": self.session_id}) ) except IntegrityError: # Another concurrent writer created the parent row first. pass # Insert messages in bulk await sess.execute(insert(self._messages), payload) # Touch updated_at column await sess.execute( update(self._sessions) .where(self._sessions.c.session_id == self.session_id) .values(updated_at=sql_text("CURRENT_TIMESTAMP")) ) await self._run_sqlite_write_with_retry(_write_items) async def pop_item(self) -> TResponseInputItem | None: """Remove and return the most recent item from the session. Returns: The most recent item if it exists, None if the session is empty """ await self._ensure_tables() async with self._session_factory() as sess: async with sess.begin(): # Fallback for all dialects - get ID first, then delete subq = ( select(self._messages.c.id) .where(self._messages.c.session_id == self.session_id) .order_by( self._messages.c.created_at.desc(), self._messages.c.id.desc(), ) .limit(1) ) res = await sess.execute(subq) row_id = res.scalar_one_or_none() if row_id is None: return None # Fetch data before deleting res_data = await sess.execute( select(self._messages.c.message_data).where(self._messages.c.id == row_id) ) row = res_data.scalar_one_or_none() await sess.execute(delete(self._messages).where(self._messages.c.id == row_id)) if row is None: return None try: return await self._deserialize_item(row) except json.JSONDecodeError: return None async def clear_session(self) -> None: """Clear all items for this session.""" await self._ensure_tables() async with self._session_factory() as sess: async with sess.begin(): await sess.execute( delete(self._messages).where(self._messages.c.session_id == self.session_id) ) await sess.execute( delete(self._sessions).where(self._sessions.c.session_id == self.session_id) ) @property def engine(self) -> AsyncEngine: """Access the underlying SQLAlchemy AsyncEngine. This property provides direct access to the engine for advanced use cases, such as checking connection pool status, configuring engine settings, or manually disposing the engine when needed. Returns: AsyncEngine: The SQLAlchemy async engine instance. """ return self._engine