Skip to content

Commit 6517db9

Browse files
committed
fix(memory): require a complete structure-table layout owned by the configured pair
Address review feedback. Validation is now one contract. create_tables=True creates and claims the layout inside a single BEGIN IMMEDIATE transaction, and create_tables=False accepts only an already-existing layout owned by the configured pair. An incomplete or unclaimed layout is rejected, so a no-create session can no longer open a file before any pair owns it and then have another pair claim the structure tables underneath it. Identifier equality is delegated to SQLite's ASCII-only NOCASE collation instead of Python casefold(), which equated names SQLite keeps distinct such as 'sssessions' and the sharp-s spelling.
1 parent f68b409 commit 6517db9

2 files changed

Lines changed: 189 additions & 33 deletions

File tree

‎src/agents/extensions/memory/advanced_sqlite_session.py‎

Lines changed: 47 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
import time
88
from contextlib import closing
99
from pathlib import Path
10-
from typing import Any, cast
10+
from typing import Any, ClassVar, cast
1111

1212
from agents.result import RunResult
1313
from agents.usage import Usage
@@ -72,7 +72,7 @@ def __init__(
7272
self._init_structure_tables()
7373
else:
7474
with self._locked_connection() as conn:
75-
self._verify_structure_tables_scope(conn)
75+
self._claim_structure_tables(conn)
7676
except BaseException:
7777
try:
7878
self.close()
@@ -111,47 +111,61 @@ def _commit_branch_pointer(self, branch_id: str, generation: int) -> bool:
111111
self._current_branch_id = branch_id
112112
return True
113113

114-
def _verify_structure_tables_scope(self, conn: sqlite3.Connection) -> None:
115-
"""Reject a database file whose structure tables belong to other base tables.
114+
# The structure tables that record which base-table pair owns a database file, and the
115+
# foreign keys that record it. `branch_reservations` and `session_clear_generations` carry no
116+
# foreign keys, but enforcing one pair per file keeps them unambiguous too.
117+
_STRUCTURE_TABLE_OWNERS: ClassVar[dict[str, tuple[str, ...]]] = {
118+
"message_structure": ("session_id", "message_id"),
119+
"turn_usage": ("session_id",),
120+
}
121+
122+
def _claim_structure_tables(self, conn: sqlite3.Connection) -> None:
123+
"""Require a complete structure-table layout owned by the configured base-table pair.
116124
117125
The structure tables are not named after ``sessions_table``/``messages_table``, so a
118-
database file can only hold the structure rows of a single base-table pair. ``CREATE
119-
TABLE IF NOT EXISTS`` keeps the first pair's foreign keys, so a second session configured
120-
with different base table names would join ``message_structure`` against the wrong
121-
messages table and read back rows that belong to the other pair.
122-
123-
Both owner-bearing structure tables are checked, so a file whose two structure tables were
124-
created for different pairs is rejected instead of being accepted by whichever one happens
125-
to match.
126+
database file can only hold the structure rows of a single pair. ``CREATE TABLE IF NOT
127+
EXISTS`` keeps the first pair's foreign keys, so a second session configured with
128+
different base table names would join ``message_structure`` against the wrong messages
129+
table and read back rows that belong to the other pair.
130+
131+
An incomplete layout is rejected rather than accepted, so a ``create_tables=False``
132+
session cannot open a file before any pair has claimed it and then have another pair
133+
claim it underneath.
126134
"""
127-
expected = {
128-
"message_structure": {
129-
"session_id": self.sessions_table,
130-
"message_id": self.messages_table,
131-
},
132-
"turn_usage": {"session_id": self.sessions_table},
133-
}
134-
for table, expected_owners in expected.items():
135-
# SQLite resolves table names case-insensitively but `PRAGMA foreign_key_list` reports
136-
# the spelling used at CREATE time, so compare case-folded.
137-
owners = {
138-
row[3].casefold(): row[2]
139-
for row in conn.execute(f"PRAGMA foreign_key_list({table})")
140-
}
141-
if not owners:
142-
continue
135+
owned_by = {"session_id": self.sessions_table, "message_id": self.messages_table}
136+
for table, columns in self._STRUCTURE_TABLE_OWNERS.items():
137+
owners = {row[3]: row[2] for row in conn.execute(f"PRAGMA foreign_key_list({table})")}
138+
missing = [column for column in columns if column not in owners]
139+
if missing:
140+
raise ValueError(
141+
f"The `{table}` table in {self.db_path} does not exist yet or does not record "
142+
"which sessions_table/messages_table pair owns it. Construct an "
143+
"AdvancedSQLiteSession with create_tables=True to create and claim the "
144+
"structure tables before opening the database without them."
145+
)
143146
if any(
144-
owners.get(column, "").casefold() != owner.casefold()
145-
for column, owner in expected_owners.items()
147+
not self._identifiers_equal(conn, owners[column], owned_by[column])
148+
for column in columns
146149
):
147-
found = "/".join(owners.get(column, "?") for column in expected_owners)
148-
configured = "/".join(expected_owners.values())
150+
found = "/".join(owners[column] for column in columns)
151+
configured = "/".join(owned_by[column] for column in columns)
149152
raise ValueError(
150153
f"The `{table}` table in {self.db_path} already belongs to '{found}', not to "
151154
f"the configured '{configured}'. Structure tables are shared per database "
152155
"file, so give each sessions_table/messages_table pair its own db_path."
153156
)
154157

158+
@staticmethod
159+
def _identifiers_equal(conn: sqlite3.Connection, left: str, right: str) -> bool:
160+
"""Compare two table names the way SQLite compares identifiers.
161+
162+
SQLite folds identifiers with ASCII rules only, so `NOCASE` is asked directly rather than
163+
reimplemented. Python's ``casefold()`` would equate names SQLite keeps distinct, for
164+
example ``ßsessions`` and ``sssessions``.
165+
"""
166+
row = conn.execute("SELECT ? = ? COLLATE NOCASE", (left, right)).fetchone()
167+
return bool(row[0])
168+
155169
def _init_structure_tables(self):
156170
"""Add structure and usage tracking tables.
157171
@@ -230,7 +244,7 @@ def _init_structure_tables(self):
230244
ON turn_usage(session_id, branch_id, user_turn_number)
231245
""")
232246

233-
self._verify_structure_tables_scope(conn)
247+
self._claim_structure_tables(conn)
234248

235249
conn.commit()
236250

‎tests/extensions/memory/test_advanced_sqlite_session.py‎

Lines changed: 142 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,34 @@
3333
pytestmark = pytest.mark.asyncio
3434

3535

36+
def _claim_structure_tables_in_process(
37+
db_path: str,
38+
sessions_table: str,
39+
messages_table: str,
40+
ready: Any,
41+
start: Any,
42+
results: Any,
43+
) -> None:
44+
"""Construct a create_tables session in a child process and report the outcome."""
45+
ready.set()
46+
start.wait(timeout=30)
47+
try:
48+
session = AdvancedSQLiteSession(
49+
session_id="concurrent",
50+
db_path=db_path,
51+
create_tables=True,
52+
sessions_table=sessions_table,
53+
messages_table=messages_table,
54+
)
55+
session.close()
56+
except ValueError:
57+
results.put(("rejected", sessions_table))
58+
except BaseException as exc: # pragma: no cover - surfaced in the assertion below
59+
results.put((f"error:{type(exc).__name__}", sessions_table))
60+
else:
61+
results.put(("claimed", sessions_table))
62+
63+
3664
def _multiprocessing_context() -> Any:
3765
method = "spawn" if sys.platform == "win32" else "forkserver"
3866
return multiprocessing.get_context(method)
@@ -3844,3 +3872,117 @@ async def test_structure_tables_accept_equivalent_identifier_casing(tmp_path: Pa
38443872
assert await recased.get_items() == [{"role": "user", "content": "first"}]
38453873
finally:
38463874
recased.close()
3875+
3876+
3877+
async def test_structure_tables_reject_distinct_non_ascii_identifiers(tmp_path: Path) -> None:
3878+
"""SQLite folds identifiers with ASCII rules, so these are two different pairs.
3879+
3880+
Python's `casefold()` equates `ßsessions` and `sssessions`, which would let the second pair
3881+
through and restore the cross-table mixing this change prevents.
3882+
"""
3883+
assert "ßsessions".casefold() == "sssessions".casefold()
3884+
3885+
db_path = tmp_path / "advanced_non_ascii_structure.db"
3886+
first = AdvancedSQLiteSession(
3887+
session_id="shared",
3888+
db_path=db_path,
3889+
create_tables=True,
3890+
sessions_table="ßsessions",
3891+
messages_table="ßmessages",
3892+
)
3893+
try:
3894+
await first.add_items([{"role": "user", "content": "first"}])
3895+
3896+
with pytest.raises(ValueError, match="ßsessions"):
3897+
AdvancedSQLiteSession(
3898+
session_id="shared",
3899+
db_path=db_path,
3900+
create_tables=True,
3901+
sessions_table="sssessions",
3902+
messages_table="ssmessages",
3903+
)
3904+
3905+
assert await first.get_items() == [{"role": "user", "content": "first"}]
3906+
finally:
3907+
first.close()
3908+
3909+
3910+
async def test_no_create_session_rejects_a_database_without_an_owner(tmp_path: Path) -> None:
3911+
"""A no-create session must not open a file before a pair has claimed the structure tables.
3912+
3913+
Accepting it would let another pair claim the tables afterwards, leaving this session writing
3914+
and reading structure rows owned by that other pair.
3915+
"""
3916+
db_path = tmp_path / "advanced_unclaimed_structure.db"
3917+
3918+
with pytest.raises(ValueError, match="create_tables=True"):
3919+
AdvancedSQLiteSession(session_id="shared", db_path=db_path, create_tables=False)
3920+
3921+
owner = AdvancedSQLiteSession(session_id="shared", db_path=db_path, create_tables=True)
3922+
try:
3923+
await owner.add_items([{"role": "user", "content": "first"}])
3924+
finally:
3925+
owner.close()
3926+
3927+
# Once a pair owns the layout, the same pair may open it without creating anything.
3928+
reader = AdvancedSQLiteSession(session_id="shared", db_path=db_path, create_tables=False)
3929+
try:
3930+
assert await reader.get_items() == [{"role": "user", "content": "first"}]
3931+
finally:
3932+
reader.close()
3933+
3934+
3935+
@pytest.mark.review_optional
3936+
async def test_concurrent_structure_table_claims_leave_one_coherent_owner(tmp_path: Path) -> None:
3937+
"""Two processes claiming a fresh file with different pairs must not split the layout."""
3938+
db_path = tmp_path / "advanced_concurrent_claim.db"
3939+
pairs = [("a_sessions", "a_messages"), ("b_sessions", "b_messages")]
3940+
3941+
# Create the file with WAL already enabled. Switching journal mode needs an exclusive lock,
3942+
# so two processes opening a brand-new file race there first, in SQLiteSession.__init__,
3943+
# before either reaches the structure tables. Settling it up front keeps this test about the
3944+
# claim.
3945+
with contextlib.closing(sqlite3.connect(db_path)) as setup_conn:
3946+
setup_conn.execute("PRAGMA journal_mode=WAL")
3947+
3948+
context = _multiprocessing_context()
3949+
start = context.Event()
3950+
results = context.Queue()
3951+
ready_events = [context.Event(), context.Event()]
3952+
processes = [
3953+
context.Process(
3954+
target=_claim_structure_tables_in_process,
3955+
args=(str(db_path), sessions_table, messages_table, ready, start, results),
3956+
)
3957+
for (sessions_table, messages_table), ready in zip(pairs, ready_events, strict=False)
3958+
]
3959+
3960+
try:
3961+
for process in processes:
3962+
process.start()
3963+
for ready in ready_events:
3964+
assert ready.wait(timeout=30)
3965+
start.set()
3966+
for process in processes:
3967+
process.join(timeout=30)
3968+
assert process.exitcode == 0
3969+
3970+
outcomes = [results.get(timeout=5), results.get(timeout=5)]
3971+
claimed = [pair for status, pair in outcomes if status == "claimed"]
3972+
assert len(claimed) == 1, outcomes
3973+
assert all(status in {"claimed", "rejected"} for status, _ in outcomes), outcomes
3974+
3975+
# Every owner-bearing structure table must name the one pair that won.
3976+
winner = claimed[0]
3977+
with contextlib.closing(sqlite3.connect(db_path)) as conn:
3978+
for table in ("message_structure", "turn_usage"):
3979+
owners = {
3980+
row[3]: row[2] for row in conn.execute(f"PRAGMA foreign_key_list({table})")
3981+
}
3982+
assert owners["session_id"] == winner, (table, owners)
3983+
finally:
3984+
start.set()
3985+
for process in processes:
3986+
if process.is_alive():
3987+
process.terminate()
3988+
process.join(timeout=5)

0 commit comments

Comments
 (0)