diff --git a/src/agents/extensions/memory/sqlalchemy_session.py b/src/agents/extensions/memory/sqlalchemy_session.py index 81f7dcdae8..6be385e854 100644 --- a/src/agents/extensions/memory/sqlalchemy_session.py +++ b/src/agents/extensions/memory/sqlalchemy_session.py @@ -48,8 +48,14 @@ text as sql_text, update, ) -from sqlalchemy.exc import IntegrityError, OperationalError -from sqlalchemy.ext.asyncio import AsyncEngine, async_sessionmaker, create_async_engine +from sqlalchemy.dialects import mysql as mysql_dialect +from sqlalchemy.exc import IntegrityError, OperationalError, SQLAlchemyError +from sqlalchemy.ext.asyncio import ( + AsyncConnection, + AsyncEngine, + async_sessionmaker, + create_async_engine, +) from ...items import TResponseInputItem from ...memory.session import SessionABC @@ -62,6 +68,24 @@ _T = TypeVar("_T") +# MySQL-family dialects require a bounded VARCHAR for indexed string columns. +_MYSQL_SESSION_ID_MAX_LENGTH = 190 +# ``CHARACTER SET`` is declared alongside the collation: a column given only a +# collation inherits the database character set, and the server rejects +# ``utf8mb4_bin`` against a non-utf8mb4 inherited set with +# "ERROR 1253 COLLATION 'utf8mb4_bin' is not valid for CHARACTER SET ''". +# A MySQL 5.7 install defaulting to latin1 would otherwise fail in +# ``create_all()`` before either table exists. +_SESSION_ID_TYPE = String().with_variant( + mysql_dialect.VARCHAR( + _MYSQL_SESSION_ID_MAX_LENGTH, + charset="utf8mb4", + collation="utf8mb4_bin", + ), + "mysql", + "mariadb", +) + class SQLAlchemySession(SessionABC): """SQLAlchemy implementation of [`Session`][agents.memory.session.Session].""" @@ -163,7 +187,9 @@ def __init__( '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. + development and testing when migrations aren't used. Automatically created + MySQL and MariaDB schemas store session IDs in VARCHAR(190) columns, and + session IDs longer than that are rejected only for those schemas. 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 @@ -189,7 +215,7 @@ def __init__( self._sessions = Table( sessions_table, self._metadata, - Column("session_id", String, primary_key=True), + Column("session_id", _SESSION_ID_TYPE, primary_key=True), Column( "created_at", TIMESTAMP(timezone=False), @@ -211,7 +237,7 @@ def __init__( Column("id", Integer, primary_key=True, autoincrement=True), Column( "session_id", - String, + _SESSION_ID_TYPE, ForeignKey(f"{sessions_table}.session_id", ondelete="CASCADE"), nullable=False, ), @@ -234,6 +260,7 @@ def __init__( self._session_factory = async_sessionmaker(self._engine, expire_on_commit=False) self._create_tables = create_tables + self._session_id_collation_validated = False # --------------------------------------------------------------------- # Convenience constructors @@ -278,9 +305,70 @@ async def _deserialize_item(self, item: str) -> TResponseInputItem: # ------------------------------------------------------------------ # Session protocol implementation # ------------------------------------------------------------------ + async def _validate_session_id_collation(self, conn: AsyncConnection) -> None: + """Reject trailing-space IDs only when the actual MySQL collation pads spaces.""" + if self._engine.dialect.name not in {"mysql", "mariadb"}: + return + if not self.session_id.endswith(" "): + return + + try: + collation_result = await conn.execute( + sql_text( + "SELECT COLLATION_NAME FROM information_schema.COLUMNS " + "WHERE TABLE_SCHEMA = DATABASE() AND TABLE_NAME = :table_name " + "AND COLUMN_NAME = 'session_id'" + ), + {"table_name": self._sessions.name}, + ) + collation = collation_result.scalar_one_or_none() + if not collation: + return + + pad_attribute: str | None + if getattr(self._engine.dialect, "is_mariadb", False): + pad_attribute = "NO PAD" if "_nopad_" in collation.casefold() else "PAD SPACE" + else: + try: + pad_result = await conn.execute( + sql_text( + "SELECT PAD_ATTRIBUTE FROM information_schema.COLLATIONS " + "WHERE COLLATION_NAME = :collation" + ), + {"collation": collation}, + ) + pad_attribute = pad_result.scalar_one_or_none() + except SQLAlchemyError: + version_result = await conn.execute(sql_text("SELECT VERSION()")) + version = version_result.scalar_one_or_none() + pad_attribute = ( + "PAD SPACE" + if version + and version.partition(".")[0].isdigit() + and int(version.partition(".")[0]) < 8 + else None + ) + except SQLAlchemyError: + return + + if pad_attribute == "PAD SPACE": + raise ValueError( + f"session_id {self.session_id!r} ends with a space, which is not distinct " + f"under the column's PAD SPACE collation {collation!r}; two sessions would " + "silently share one history" + ) + async def _ensure_tables(self) -> None: """Ensure tables are created before any database operations.""" if not self._create_tables: + if ( + not self._session_id_collation_validated + and self._engine.dialect.name in {"mysql", "mariadb"} + and self.session_id.endswith(" ") + ): + async with self._engine.connect() as conn: + await self._validate_session_id_collation(conn) + self._session_id_collation_validated = True return assert self._init_lock is not None @@ -294,6 +382,8 @@ async def _ensure_tables(self) -> None: async with self._engine.begin() as conn: await conn.run_sync(self._metadata.create_all) + await self._validate_session_id_collation(conn) + self._session_id_collation_validated = True self._create_tables = False # Only create once finally: self._init_lock.release() diff --git a/tests/extensions/memory/test_sqlalchemy_session.py b/tests/extensions/memory/test_sqlalchemy_session.py index b985d0a7e9..eb7faf3932 100644 --- a/tests/extensions/memory/test_sqlalchemy_session.py +++ b/tests/extensions/memory/test_sqlalchemy_session.py @@ -8,7 +8,9 @@ from contextlib import asynccontextmanager from datetime import datetime, timedelta from pathlib import Path +from types import SimpleNamespace from typing import Any, cast +from unittest.mock import AsyncMock, MagicMock import pytest from openai.types.responses.response_output_message_param import ResponseOutputMessageParam @@ -17,7 +19,9 @@ ResponseReasoningItemParam, Summary, ) -from sqlalchemy import event, insert, select, text, update +from sqlalchemy import create_mock_engine, event, insert, select, text, update +from sqlalchemy.dialects import postgresql, sqlite +from sqlalchemy.exc import SQLAlchemyError from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine from sqlalchemy.sql import Select @@ -35,6 +39,60 @@ DB_URL = "sqlite+aiosqlite:///:memory:" +@pytest.mark.parametrize("dialect_url", ["mysql://", "mariadb://"]) +async def test_schema_create_all_compiles_for_mysql_family(dialect_url: str): + """MySQL-family schema creation includes both tables and the session-time index.""" + session = SQLAlchemySession.from_url("schema_compile", url=DB_URL) + tables = (session._sessions, session._messages) + statements: list[str] = [] + + def record(statement: Any, *args: Any, **kwargs: Any) -> None: + statements.append(str(statement.compile(dialect=engine.dialect))) + + engine = create_mock_engine(dialect_url, record) + + try: + session._metadata.create_all(engine) + for table in tables: + # CHARACTER SET must be emitted with the collation: a column given + # only a collation inherits the database character set, and the + # server rejects utf8mb4_bin against a non-utf8mb4 set with + # ERROR 1253, failing create_all() on e.g. a latin1 MySQL 5.7. + assert ( + table.c.session_id.type.compile(dialect=engine.dialect) + == "VARCHAR(190) CHARACTER SET utf8mb4 COLLATE utf8mb4_bin" + ) + finally: + await session.engine.dispose() + + assert any("CREATE TABLE agent_sessions" in statement for statement in statements) + messages_ddl = next( + statement for statement in statements if "CREATE TABLE agent_messages" in statement + ) + assert ( + "FOREIGN KEY(session_id) REFERENCES agent_sessions (session_id) ON DELETE CASCADE" + in messages_ddl + ) + assert any( + "CREATE INDEX idx_agent_messages_session_time " + "ON agent_messages (session_id, created_at)" in statement + for statement in statements + ) + + +async def test_schema_keeps_unbounded_session_ids_for_sqlite_and_postgresql(): + """SQLite and PostgreSQL retain the pre-existing unbounded string type.""" + session = SQLAlchemySession.from_url("schema_compile", url=DB_URL) + + try: + for table in (session._sessions, session._messages): + session_id_type = table.c.session_id.type + assert session_id_type.compile(dialect=postgresql.dialect()) == "VARCHAR" + assert session_id_type.compile(dialect=sqlite.dialect()) == "VARCHAR" + finally: + await session.engine.dispose() + + def _make_message_item(item_id: str, text_value: str) -> TResponseInputItem: content: ResponseOutputTextParam = { "type": "output_text", @@ -203,6 +261,193 @@ async def test_session_isolation(agent: Agent): assert "dogs" not in result.final_output.lower() +async def test_session_ids_are_case_sensitive(): + """Session IDs that differ only by case retain separate histories.""" + engine = create_async_engine(DB_URL) + upper = SQLAlchemySession("Foo", engine=engine, create_tables=True) + lower = SQLAlchemySession("foo", engine=engine, create_tables=True) + + try: + await upper.add_items([{"role": "user", "content": "upper"}]) + await lower.add_items([{"role": "user", "content": "lower"}]) + + assert await upper.get_items() == [{"role": "user", "content": "upper"}] + assert await lower.get_items() == [{"role": "user", "content": "lower"}] + finally: + await engine.dispose() + + +@pytest.mark.parametrize("dialect_name", ["mysql", "mariadb"]) +@pytest.mark.parametrize("create_tables", [True, False]) +async def test_constructor_does_not_impose_a_session_id_length_bound( + dialect_name: str, create_tables: bool +): + """The constructor leaves the actual schema authoritative for session ID length.""" + engine = MagicMock(spec=AsyncEngine) + engine.dialect = SimpleNamespace(name=dialect_name) + long_id = "a" * 191 + + session = SQLAlchemySession(long_id, engine=engine, create_tables=create_tables) + + assert session.session_id == long_id + + +class _ScalarResult: + def __init__(self, value: str | None) -> None: + self._value = value + + def scalar_one_or_none(self) -> str | None: + return self._value + + +@pytest.mark.parametrize("dialect_name", ["mysql", "mariadb"]) +async def test_validate_session_id_collation_rejects_pad_space( + dialect_name: str, +) -> None: + engine = MagicMock(spec=AsyncEngine) + engine.dialect = SimpleNamespace(name=dialect_name, is_mariadb=dialect_name == "mariadb") + session = SQLAlchemySession( + "tenant ", + engine=engine, + create_tables=True, + sessions_table="custom_sessions", + ) + conn = MagicMock() + conn.execute = AsyncMock(side_effect=[_ScalarResult("utf8mb4_bin"), _ScalarResult("PAD SPACE")]) + + with pytest.raises( + ValueError, + match=( + r"session_id 'tenant ' ends with a space, which is not distinct under the " + r"column's PAD SPACE collation 'utf8mb4_bin'; two sessions would silently " + r"share one history" + ), + ): + await session._validate_session_id_collation(conn) + + expected_calls = 1 if dialect_name == "mariadb" else 2 + assert conn.execute.await_count == expected_calls + assert conn.execute.await_args_list[0].args[1] == {"table_name": "custom_sessions"} + if dialect_name == "mysql": + assert conn.execute.await_args_list[1].args[1] == {"collation": "utf8mb4_bin"} + + +async def test_validate_session_id_collation_rejects_pad_space_on_mysql_57() -> None: + engine = MagicMock(spec=AsyncEngine) + engine.dialect = SimpleNamespace(name="mysql", is_mariadb=False) + session = SQLAlchemySession("tenant ", engine=engine, create_tables=True) + conn = MagicMock() + conn.execute = AsyncMock( + side_effect=[ + _ScalarResult("utf8mb4_bin"), + SQLAlchemyError("Unknown column PAD_ATTRIBUTE"), + _ScalarResult("5.7.44"), + ] + ) + + with pytest.raises(ValueError, match="PAD SPACE collation"): + await session._validate_session_id_collation(conn) + + assert conn.execute.await_count == 3 + + +async def test_validate_session_id_collation_allows_mariadb_nopad() -> None: + engine = MagicMock(spec=AsyncEngine) + engine.dialect = SimpleNamespace(name="mysql", is_mariadb=True) + session = SQLAlchemySession("tenant ", engine=engine, create_tables=True) + conn = MagicMock() + conn.execute = AsyncMock(return_value=_ScalarResult("utf8mb4_nopad_bin")) + + await session._validate_session_id_collation(conn) + + assert conn.execute.await_count == 1 + + +@pytest.mark.parametrize("pad_attribute", ["NO PAD", None]) +async def test_validate_session_id_collation_allows_non_pad_or_unknown( + pad_attribute: str | None, +) -> None: + engine = MagicMock(spec=AsyncEngine) + engine.dialect = SimpleNamespace(name="mysql") + session = SQLAlchemySession("tenant ", engine=engine, create_tables=True) + conn = MagicMock() + if pad_attribute is None: + conn.execute = AsyncMock(return_value=_ScalarResult(None)) + else: + conn.execute = AsyncMock( + side_effect=[_ScalarResult("utf8mb4_0900_bin"), _ScalarResult(pad_attribute)] + ) + + await session._validate_session_id_collation(conn) + + +@pytest.mark.parametrize("dialect_name", ["sqlite", "postgresql"]) +async def test_validate_session_id_collation_skips_non_mysql_dialects( + dialect_name: str, + monkeypatch: pytest.MonkeyPatch, +) -> None: + engine = MagicMock(spec=AsyncEngine) + engine.dialect = SimpleNamespace(name=dialect_name) + monkeypatch.setattr(SQLAlchemySession, "_configure_sqlite_engine", MagicMock()) + session = SQLAlchemySession("tenant ", engine=engine, create_tables=True) + conn = MagicMock() + conn.execute = AsyncMock() + + await session._validate_session_id_collation(conn) + + conn.execute.assert_not_awaited() + + +async def test_create_tables_false_skips_session_id_collation_validation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + session = SQLAlchemySession.from_url("tenant ", url=DB_URL, create_tables=False) + validate = AsyncMock() + monkeypatch.setattr(session, "_validate_session_id_collation", validate) + + try: + await session._ensure_tables() + finally: + await session.engine.dispose() + + validate.assert_not_awaited() + + +async def test_existing_mysql_schema_validates_trailing_space_session_id( + monkeypatch: pytest.MonkeyPatch, +) -> None: + engine = MagicMock(spec=AsyncEngine) + engine.dialect = SimpleNamespace(name="mysql", is_mariadb=False) + conn = MagicMock() + connection_context = MagicMock() + connection_context.__aenter__ = AsyncMock(return_value=conn) + connection_context.__aexit__ = AsyncMock(return_value=None) + engine.connect.return_value = connection_context + session = SQLAlchemySession("tenant ", engine=engine, create_tables=False) + validate = AsyncMock() + monkeypatch.setattr(session, "_validate_session_id_collation", validate) + + await session._ensure_tables() + + validate.assert_awaited_once_with(conn) + + +async def test_session_ids_keep_trailing_spaces_on_sqlite(): + """SQLite stores trailing-space IDs as distinct values.""" + engine = create_async_engine(DB_URL) + bare = SQLAlchemySession("tenant", engine=engine, create_tables=True) + padded = SQLAlchemySession("tenant ", engine=engine, create_tables=True) + + try: + await bare.add_items([{"role": "user", "content": "bare"}]) + await padded.add_items([{"role": "user", "content": "padded"}]) + + assert await bare.get_items() == [{"role": "user", "content": "bare"}] + assert await padded.get_items() == [{"role": "user", "content": "padded"}] + finally: + await engine.dispose() + + async def test_get_items_with_limit(agent: Agent): """Test the limit parameter in get_items.""" session_id = "limit_test"