Files
2026-08-13 11:20:33 +09:00

1215 lines
44 KiB
Python

"""Tests for session memory functionality."""
import asyncio
import sqlite3
import tempfile
import threading
from pathlib import Path
from typing import Any, cast
import pytest
from agents import Agent, RunConfig, Runner, SessionSettings, SQLiteSession, TResponseInputItem
from agents.memory.sqlite_session import _await_mutation
from agents.testing import ScriptedModel
from tests.test_responses import get_text_message
@pytest.mark.asyncio
async def test_await_mutation_cancellation_hides_later_failure_without_loop_error() -> None:
"""A failed mutation must not leak a false loop error after caller cancellation."""
mutation_started = asyncio.Event()
allow_failure = asyncio.Event()
loop = asyncio.get_running_loop()
previous_exception_handler = loop.get_exception_handler()
loop_errors: list[dict[str, Any]] = []
async def mutation() -> None:
mutation_started.set()
await allow_failure.wait()
raise RuntimeError("mutation failed")
loop.set_exception_handler(lambda _loop, context: loop_errors.append(context))
task = asyncio.create_task(_await_mutation(mutation()))
try:
await mutation_started.wait()
task.cancel("caller-cancelled")
allow_failure.set()
with pytest.raises(asyncio.CancelledError):
await task
await asyncio.sleep(0)
assert loop_errors == []
finally:
loop.set_exception_handler(previous_exception_handler)
allow_failure.set()
if not task.done():
task.cancel()
await asyncio.gather(task, return_exceptions=True)
# Helper functions for parametrized testing of different Runner methods
def _run_sync_wrapper(agent, input_data, **kwargs):
"""Wrapper for run_sync that properly sets up an event loop."""
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
try:
return Runner.run_sync(agent, input_data, **kwargs)
finally:
loop.close()
async def run_agent_async(runner_method: str, agent, input_data, **kwargs):
"""Helper function to run agent with different methods."""
if runner_method == "run":
return await Runner.run(agent, input_data, **kwargs)
elif runner_method == "run_sync":
# For run_sync, we need to run it in a thread with its own event loop
return await asyncio.to_thread(_run_sync_wrapper, agent, input_data, **kwargs)
elif runner_method == "run_streamed":
result = Runner.run_streamed(agent, input_data, **kwargs)
# For streaming, we first try to get at least one event to trigger any early exceptions
# If there's an exception in setup (like memory validation), it will be raised here
try:
first_event = None
async for event in result.stream_events():
if first_event is None:
first_event = event
# Continue consuming all events
pass
except Exception:
# If an exception occurs during streaming, we let it propagate up
raise
return result
else:
raise ValueError(f"Unknown runner method: {runner_method}")
# Parametrized tests for different runner methods
@pytest.mark.parametrize("runner_method", ["run", "run_sync", "run_streamed"])
@pytest.mark.asyncio
async def test_session_memory_basic_functionality_parametrized(runner_method):
"""Test basic session memory functionality with SQLite backend across all runner methods."""
with tempfile.TemporaryDirectory() as temp_dir:
db_path = Path(temp_dir) / "test_memory.db"
session_id = "test_session_123"
session = SQLiteSession(session_id, db_path)
model = ScriptedModel()
agent = Agent(name="test", model=model)
# First turn
model.enqueue([get_text_message("San Francisco")])
result1 = await run_agent_async(
runner_method,
agent,
"What city is the Golden Gate Bridge in?",
session=session,
)
assert result1.final_output == "San Francisco"
# Second turn - should have conversation history
model.enqueue([get_text_message("California")])
result2 = await run_agent_async(
runner_method,
agent,
"What state is it in?",
session=session,
)
assert result2.final_output == "California"
# Verify that the input to the second turn includes the previous conversation
# The model should have received the full conversation history
last_input = model.calls[-1].input
assert len(last_input) > 1 # Should have more than just the current message
session.close()
@pytest.mark.parametrize("runner_method", ["run", "run_sync", "run_streamed"])
@pytest.mark.asyncio
async def test_session_memory_with_explicit_instance_parametrized(runner_method):
"""Test session memory with an explicit SQLiteSession instance across all runner methods."""
with tempfile.TemporaryDirectory() as temp_dir:
db_path = Path(temp_dir) / "test_memory.db"
session_id = "test_session_456"
session = SQLiteSession(session_id, db_path)
model = ScriptedModel()
agent = Agent(name="test", model=model)
# First turn
model.enqueue([get_text_message("Hello")])
result1 = await run_agent_async(runner_method, agent, "Hi there", session=session)
assert result1.final_output == "Hello"
# Second turn
model.enqueue([get_text_message("I remember you said hi")])
result2 = await run_agent_async(
runner_method,
agent,
"Do you remember what I said?",
session=session,
)
assert result2.final_output == "I remember you said hi"
session.close()
@pytest.mark.parametrize("runner_method", ["run", "run_sync", "run_streamed"])
@pytest.mark.asyncio
async def test_session_memory_disabled_parametrized(runner_method):
"""Test that session memory is disabled when session=None across all runner methods."""
model = ScriptedModel()
agent = Agent(name="test", model=model)
# First turn (no session parameters = disabled)
model.enqueue([get_text_message("Hello")])
result1 = await run_agent_async(runner_method, agent, "Hi there")
assert result1.final_output == "Hello"
# Second turn - should NOT have conversation history
model.enqueue([get_text_message("I don't remember")])
result2 = await run_agent_async(runner_method, agent, "Do you remember what I said?")
assert result2.final_output == "I don't remember"
# Verify that the input to the second turn is just the current message
last_input = model.calls[-1].input
assert len(last_input) == 1 # Should only have the current message
@pytest.mark.parametrize("runner_method", ["run", "run_sync", "run_streamed"])
@pytest.mark.asyncio
async def test_session_memory_different_sessions_parametrized(runner_method):
"""Test that different session IDs maintain separate conversation histories across all runner
methods."""
with tempfile.TemporaryDirectory() as temp_dir:
db_path = Path(temp_dir) / "test_memory.db"
model = ScriptedModel()
agent = Agent(name="test", model=model)
# Session 1
session_id_1 = "session_1"
session_1 = SQLiteSession(session_id_1, db_path)
model.enqueue([get_text_message("I like cats")])
result1 = await run_agent_async(runner_method, agent, "I like cats", session=session_1)
assert result1.final_output == "I like cats"
# Session 2 - different session
session_id_2 = "session_2"
session_2 = SQLiteSession(session_id_2, db_path)
model.enqueue([get_text_message("I like dogs")])
result2 = await run_agent_async(runner_method, agent, "I like dogs", session=session_2)
assert result2.final_output == "I like dogs"
# Back to Session 1 - should remember cats, not dogs
model.enqueue([get_text_message("Yes, you mentioned cats")])
result3 = await run_agent_async(
runner_method,
agent,
"What did I say I like?",
session=session_1,
)
assert result3.final_output == "Yes, you mentioned cats"
session_1.close()
session_2.close()
@pytest.mark.asyncio
async def test_sqlite_session_memory_direct():
"""Test SQLiteSession class directly."""
with tempfile.TemporaryDirectory() as temp_dir:
db_path = Path(temp_dir) / "test_direct.db"
session_id = "direct_test"
session = SQLiteSession(session_id, db_path)
# Test adding and retrieving items
items: list[TResponseInputItem] = [
{"role": "user", "content": "Hello"},
{"role": "assistant", "content": "Hi there!"},
]
await session.add_items(items)
retrieved = await session.get_items()
assert len(retrieved) == 2
assert retrieved[0].get("role") == "user"
assert retrieved[0].get("content") == "Hello"
assert retrieved[1].get("role") == "assistant"
assert retrieved[1].get("content") == "Hi there!"
# Test clearing session
await session.clear_session()
retrieved_after_clear = await session.get_items()
assert len(retrieved_after_clear) == 0
session.close()
@pytest.mark.asyncio
async def test_sqlite_session_close_closes_worker_thread_connections():
"""Test that close cleans up connections opened by async worker threads."""
with tempfile.TemporaryDirectory() as temp_dir:
db_path = Path(temp_dir) / "test_worker_thread_close.db"
session = SQLiteSession("worker_thread_close", db_path)
await session.add_items([{"role": "user", "content": "Hello"}])
connections = list(session._connections)
assert connections
session.close()
assert session._connections == set()
with pytest.raises(sqlite3.ProgrammingError):
connections[0].execute("SELECT 1")
@pytest.mark.asyncio
async def test_sqlite_session_closed_rejects_empty_add_items():
"""add_items([]) must not bypass the closed check through the empty-list fast path."""
with tempfile.TemporaryDirectory() as temp_dir:
db_path = Path(temp_dir) / "closed_empty_add.db"
session = SQLiteSession("closed_empty_add_test", db_path)
session.close()
with pytest.raises(RuntimeError, match="SQLiteSession is closed"):
await session.add_items([])
@pytest.mark.asyncio
async def test_sqlite_session_memory_pop_item():
"""Test SQLiteSession pop_item functionality."""
with tempfile.TemporaryDirectory() as temp_dir:
db_path = Path(temp_dir) / "test_pop.db"
session_id = "pop_test"
session = SQLiteSession(session_id, db_path)
# Test popping from empty session
popped = await session.pop_item()
assert popped is None
# Add items
items: list[TResponseInputItem] = [
{"role": "user", "content": "Hello"},
{"role": "assistant", "content": "Hi there!"},
{"role": "user", "content": "How are you?"},
]
await session.add_items(items)
# Verify all items are there
retrieved = await session.get_items()
assert len(retrieved) == 3
# Pop the most recent item
popped = await session.pop_item()
assert popped is not None
assert popped.get("role") == "user"
assert popped.get("content") == "How are you?"
# Verify item was removed
retrieved_after_pop = await session.get_items()
assert len(retrieved_after_pop) == 2
assert retrieved_after_pop[-1].get("content") == "Hi there!"
# Pop another item
popped2 = await session.pop_item()
assert popped2 is not None
assert popped2.get("role") == "assistant"
assert popped2.get("content") == "Hi there!"
# Pop the last item
popped3 = await session.pop_item()
assert popped3 is not None
assert popped3.get("role") == "user"
assert popped3.get("content") == "Hello"
# Try to pop from empty session again
popped4 = await session.pop_item()
assert popped4 is None
# Verify session is empty
final_items = await session.get_items()
assert len(final_items) == 0
session.close()
@pytest.mark.asyncio
async def test_session_memory_pop_different_sessions():
"""Test that pop_item only affects the specified session."""
with tempfile.TemporaryDirectory() as temp_dir:
db_path = Path(temp_dir) / "test_pop_sessions.db"
session_1_id = "session_1"
session_2_id = "session_2"
session_1 = SQLiteSession(session_1_id, db_path)
session_2 = SQLiteSession(session_2_id, db_path)
# Add items to both sessions
items_1: list[TResponseInputItem] = [
{"role": "user", "content": "Session 1 message"},
]
items_2: list[TResponseInputItem] = [
{"role": "user", "content": "Session 2 message 1"},
{"role": "user", "content": "Session 2 message 2"},
]
await session_1.add_items(items_1)
await session_2.add_items(items_2)
# Pop from session 2
popped = await session_2.pop_item()
assert popped is not None
assert popped.get("content") == "Session 2 message 2"
# Verify session 1 is unaffected
session_1_items = await session_1.get_items()
assert len(session_1_items) == 1
assert session_1_items[0].get("content") == "Session 1 message"
# Verify session 2 has one item left
session_2_items = await session_2.get_items()
assert len(session_2_items) == 1
assert session_2_items[0].get("content") == "Session 2 message 1"
session_1.close()
session_2.close()
@pytest.mark.asyncio
async def test_sqlite_session_pop_item_skips_corrupt_most_recent():
"""pop_item skips corrupt newest rows and returns the next valid item."""
with tempfile.TemporaryDirectory() as temp_dir:
db_path = Path(temp_dir) / "test_pop_corrupt.db"
session = SQLiteSession("pop_corrupt", db_path)
valid_item: TResponseInputItem = {"role": "user", "content": "valid"}
await session.add_items([valid_item])
with session._locked_connection() as conn:
conn.execute(
f"INSERT INTO {session.messages_table} (session_id, message_data) VALUES (?, ?)",
(session.session_id, "not valid json {{{"),
)
conn.commit()
assert await session.pop_item() == valid_item
assert await session.get_items() == []
session.close()
@pytest.mark.asyncio
async def test_sqlite_session_pop_item_returns_none_after_dropping_only_corrupt_rows():
"""pop_item removes corrupt rows and returns None when no valid items remain."""
with tempfile.TemporaryDirectory() as temp_dir:
db_path = Path(temp_dir) / "test_pop_only_corrupt.db"
session = SQLiteSession("pop_only_corrupt", db_path)
with session._locked_connection() as conn:
conn.execute(
f"INSERT INTO {session.messages_table} (session_id, message_data) VALUES (?, ?)",
(session.session_id, "not valid json {{{"),
)
conn.commit()
assert await session.pop_item() is None
assert await session.get_items() == []
session.close()
@pytest.mark.asyncio
async def test_sqlite_session_get_items_with_limit():
"""Test SQLiteSession get_items with limit parameter."""
with tempfile.TemporaryDirectory() as temp_dir:
db_path = Path(temp_dir) / "test_count.db"
session_id = "count_test"
session = SQLiteSession(session_id, db_path)
# Add multiple items
items: list[TResponseInputItem] = [
{"role": "user", "content": "Message 1"},
{"role": "assistant", "content": "Response 1"},
{"role": "user", "content": "Message 2"},
{"role": "assistant", "content": "Response 2"},
{"role": "user", "content": "Message 3"},
{"role": "assistant", "content": "Response 3"},
]
await session.add_items(items)
# Test getting all items (default behavior)
all_items = await session.get_items()
assert len(all_items) == 6
assert all_items[0].get("content") == "Message 1"
assert all_items[-1].get("content") == "Response 3"
# Test getting latest 2 items
latest_2 = await session.get_items(limit=2)
assert len(latest_2) == 2
assert latest_2[0].get("content") == "Message 3"
assert latest_2[1].get("content") == "Response 3"
# Test getting latest 4 items
latest_4 = await session.get_items(limit=4)
assert len(latest_4) == 4
assert latest_4[0].get("content") == "Message 2"
assert latest_4[1].get("content") == "Response 2"
assert latest_4[2].get("content") == "Message 3"
assert latest_4[3].get("content") == "Response 3"
# Test getting more items than available
latest_10 = await session.get_items(limit=10)
assert len(latest_10) == 6 # Should return all available items
assert latest_10[0].get("content") == "Message 1"
assert latest_10[-1].get("content") == "Response 3"
# Test getting 0 items
latest_0 = await session.get_items(limit=0)
assert len(latest_0) == 0
session.close()
@pytest.mark.asyncio
async def test_sqlite_session_get_items_limit_skips_corrupt_newest_rows():
"""limit counts valid items, expanding past corrupt newest rows."""
with tempfile.TemporaryDirectory() as temp_dir:
db_path = Path(temp_dir) / "test_limit_corrupt.db"
session = SQLiteSession("limit_corrupt", db_path)
await session.add_items(
[
{"role": "user", "content": "valid 0"},
{"role": "assistant", "content": "valid 1"},
{"role": "user", "content": "valid 2"},
]
)
with session._locked_connection() as conn:
conn.execute(
f"INSERT INTO {session.messages_table} (session_id, message_data) VALUES (?, ?)",
(session.session_id, "not valid json {{{"),
)
conn.commit()
# Newest row is corrupt; limit=2 should still return the two latest valid items.
limited = await session.get_items(limit=2)
assert [item.get("content") for item in limited] == ["valid 1", "valid 2"]
session.close()
@pytest.mark.asyncio
async def test_sqlite_session_get_items_session_settings_limit_skips_corrupt_rows():
"""session_settings.limit also counts valid items when newest rows are corrupt."""
with tempfile.TemporaryDirectory() as temp_dir:
db_path = Path(temp_dir) / "test_settings_limit_corrupt.db"
session = SQLiteSession(
"settings_limit_corrupt",
db_path,
session_settings=SessionSettings(limit=2),
)
await session.add_items(
[
{"role": "user", "content": "valid 0"},
{"role": "assistant", "content": "valid 1"},
{"role": "user", "content": "valid 2"},
]
)
with session._locked_connection() as conn:
conn.execute(
f"INSERT INTO {session.messages_table} (session_id, message_data) VALUES (?, ?)",
(session.session_id, "not valid json {{{"),
)
conn.commit()
limited = await session.get_items()
assert [item.get("content") for item in limited] == ["valid 1", "valid 2"]
session.close()
@pytest.mark.parametrize("runner_method", ["run", "run_sync", "run_streamed"])
@pytest.mark.asyncio
async def test_session_memory_appends_list_input_by_default(runner_method):
"""Test that list inputs are appended to session history when no callback is provided."""
with tempfile.TemporaryDirectory() as temp_dir:
db_path = Path(temp_dir) / "test_validation.db"
session_id = "test_validation_parametrized"
session = SQLiteSession(session_id, db_path)
model = ScriptedModel()
agent = Agent(name="test", model=model)
initial_history: list[TResponseInputItem] = [
{"role": "user", "content": "Earlier message"},
{"role": "assistant", "content": "Saved reply"},
]
await session.add_items(initial_history)
list_input = [{"role": "user", "content": "Test message"}]
model.enqueue([get_text_message("This should run")])
await run_agent_async(runner_method, agent, list_input, session=session)
assert model.calls[-1].input == initial_history + list_input
session.close()
@pytest.mark.parametrize("runner_method", ["run", "run_sync", "run_streamed"])
@pytest.mark.asyncio
async def test_session_callback_prepared_input(runner_method):
"""Test if the user passes a list of items and want to append them."""
with tempfile.TemporaryDirectory() as temp_dir:
db_path = Path(temp_dir) / "test_memory.db"
model = ScriptedModel()
agent = Agent(name="test", model=model)
# Session
session_id = "session_1"
session = SQLiteSession(session_id, db_path)
# Add first messages manually
initial_history: list[TResponseInputItem] = [
{"role": "user", "content": "Hello there."},
{"role": "assistant", "content": "Hi, I'm here to assist you."},
]
try:
await session.add_items(initial_history)
def filter_assistant_messages(history, new_input):
# Only include user messages from history
return [item for item in history if item["role"] == "user"] + new_input
new_turn_input = [{"role": "user", "content": "What your name?"}]
model.enqueue([get_text_message("I'm gpt-4o")])
# Run the agent with the callable
await run_agent_async(
runner_method,
agent,
new_turn_input,
session=session,
run_config=RunConfig(session_input_callback=filter_assistant_messages),
)
expected_model_input = [
initial_history[0], # From history
new_turn_input[0], # New input
]
assert len(model.calls[-1].input) == 2
assert model.calls[-1].input == expected_model_input
finally:
session.close()
@pytest.mark.parametrize("runner_method", ["run", "run_sync", "run_streamed"])
@pytest.mark.asyncio
async def test_session_callback_repeating_history_does_not_grow_session(runner_method):
with tempfile.TemporaryDirectory() as temp_dir:
db_path = Path(temp_dir) / "test_memory.db"
model = ScriptedModel()
agent = Agent(name="test", model=model)
session = SQLiteSession("session_repeat", db_path)
def repeat_first(history, new_input):
if not history:
return new_input
return history + [history[0]] + new_input
try:
for turn in range(3):
model.enqueue([get_text_message(f"assistant {turn}")])
await run_agent_async(
runner_method,
agent,
f"user {turn}",
session=session,
run_config=RunConfig(session_input_callback=repeat_first),
)
stored = await session.get_items()
user_messages = [item for item in stored if item.get("role") == "user"]
assert [item.get("content") for item in user_messages] == [
"user 0",
"user 1",
"user 2",
]
assert len(stored) == 6
finally:
session.close()
@pytest.mark.asyncio
async def test_sqlite_session_unicode_content():
"""Test that session correctly stores and retrieves unicode/non-ASCII content."""
with tempfile.TemporaryDirectory() as temp_dir:
db_path = Path(temp_dir) / "test_unicode.db"
session_id = "unicode_test"
session = SQLiteSession(session_id, db_path)
# Add unicode content to the session
items: list[TResponseInputItem] = [
{"role": "user", "content": "こんにちは"},
{"role": "assistant", "content": "😊👍"},
{"role": "user", "content": "Привет"},
]
await session.add_items(items)
# Retrieve items and verify unicode content
retrieved = await session.get_items()
assert retrieved[0].get("content") == "こんにちは"
assert retrieved[1].get("content") == "😊👍"
assert retrieved[2].get("content") == "Привет"
session.close()
@pytest.mark.asyncio
async def test_sqlite_session_special_characters_and_sql_injection():
"""
Test that session safely stores and retrieves items with special characters and SQL keywords.
"""
with tempfile.TemporaryDirectory() as temp_dir:
db_path = Path(temp_dir) / "test_special_chars.db"
session_id = "special_chars_test"
session = SQLiteSession(session_id, db_path)
# Add items with special characters and SQL keywords
items: list[TResponseInputItem] = [
{"role": "user", "content": "O'Reilly"},
{"role": "assistant", "content": "DROP TABLE sessions;"},
{"role": "user", "content": ('"SELECT * FROM users WHERE name = "admin";"')},
{"role": "assistant", "content": "Robert'); DROP TABLE students;--"},
{"role": "user", "content": "Normal message"},
]
await session.add_items(items)
# Retrieve all items and verify they are stored correctly
retrieved = await session.get_items()
assert len(retrieved) == len(items)
assert retrieved[0].get("content") == "O'Reilly"
assert retrieved[1].get("content") == "DROP TABLE sessions;"
assert retrieved[2].get("content") == '"SELECT * FROM users WHERE name = "admin";"'
assert retrieved[3].get("content") == "Robert'); DROP TABLE students;--"
assert retrieved[4].get("content") == "Normal message"
session.close()
@pytest.mark.asyncio
async def test_sqlite_session_concurrent_access():
"""
Test concurrent access to the same session to verify data integrity.
"""
import concurrent.futures
with tempfile.TemporaryDirectory() as temp_dir:
db_path = Path(temp_dir) / "test_concurrent.db"
session_id = "concurrent_test"
session = SQLiteSession(session_id, db_path)
# Add initial item
items: list[TResponseInputItem] = [
{"role": "user", "content": f"Message {i}"} for i in range(10)
]
# Use ThreadPoolExecutor to simulate concurrent writes
def add_item(item):
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
loop.run_until_complete(session.add_items([item]))
loop.close()
with concurrent.futures.ThreadPoolExecutor(max_workers=5) as executor:
executor.map(add_item, items)
# Retrieve all items and verify all are present
retrieved = await session.get_items()
contents = {
content
for item in retrieved
for content in [item.get("content")]
if isinstance(content, str)
}
expected = {f"Message {i}" for i in range(10)}
assert contents == expected
session.close()
@pytest.mark.asyncio
async def test_sqlite_session_file_lock_is_shared_across_instances():
"""File-backed sessions pointing at the same DB path should reuse one process-local lock."""
with tempfile.TemporaryDirectory() as temp_dir:
db_path = Path(temp_dir) / "test_shared_lock.db"
lock_path = db_path.resolve()
session_1 = SQLiteSession("session_1", db_path)
session_2 = SQLiteSession("session_2", db_path)
assert session_1._lock is session_2._lock
assert SQLiteSession._file_lock_counts[lock_path] == 2
await asyncio.gather(
session_1.add_items([{"role": "user", "content": "session_1"}]),
session_2.add_items([{"role": "user", "content": "session_2"}]),
)
assert [item.get("content") for item in await session_1.get_items()] == ["session_1"]
assert [item.get("content") for item in await session_2.get_items()] == ["session_2"]
session_1.close()
assert SQLiteSession._file_lock_counts[lock_path] == 1
assert lock_path in SQLiteSession._file_locks
session_2.close()
assert lock_path not in SQLiteSession._file_lock_counts
assert lock_path not in SQLiteSession._file_locks
@pytest.mark.asyncio
async def test_sqlite_session_failed_add_items_releases_write_lock():
"""A failed add_items must not leave an open write transaction on the cached connection."""
with tempfile.TemporaryDirectory() as temp_dir:
db_path = Path(temp_dir) / "test_rollback.db"
session = SQLiteSession("rollback_test", db_path)
# json.dumps() fails only after _insert_items() has already opened a write
# transaction with the sessions-table upsert.
unserializable = cast(TResponseInputItem, {"role": "user", "content": object()})
with pytest.raises(TypeError):
await session.add_items([unserializable])
# timeout=0 disables the busy handler, so this raises immediately if the failed
# write is still holding the SQLite write lock.
probe = sqlite3.connect(str(db_path), timeout=0)
try:
probe.execute("INSERT INTO agent_sessions (session_id) VALUES ('probe')")
probe.commit()
rolled_back = probe.execute(
"SELECT COUNT(*) FROM agent_sessions WHERE session_id = 'rollback_test'"
).fetchone()[0]
finally:
probe.close()
assert rolled_back == 0
# The session must remain usable after the failure.
await session.add_items([{"role": "user", "content": "after failure"}])
assert [item.get("content") for item in await session.get_items()] == ["after failure"]
session.close()
@pytest.mark.asyncio
async def test_session_add_items_exception_propagates_in_streamed():
"""Test that exceptions from session.add_items are properly propagated
in run_streamed instead of causing the stream to hang forever.
Regression test for https://github.com/openai/openai-agents-python/issues/2130
"""
session = SQLiteSession("test_exception_session")
async def _failing_add_items(_items):
raise RuntimeError("Simulated session.add_items failure")
session.add_items = _failing_add_items # type: ignore[method-assign]
model = ScriptedModel()
agent = Agent(name="test", model=model)
model.enqueue([get_text_message("This should not be reached")])
result = Runner.run_streamed(agent, "Hello", session=session)
async def consume_stream():
async for _event in result.stream_events():
pass
with pytest.raises(RuntimeError, match="Simulated session.add_items failure"):
# Timeout ensures test fails fast instead of hanging forever if bug regresses
await asyncio.wait_for(consume_stream(), timeout=5.0)
session.close()
# ============================================================================
# SessionSettings Tests
# ============================================================================
@pytest.mark.asyncio
async def test_session_settings_default():
"""Test that session_settings defaults to empty SessionSettings."""
from agents.memory import SessionSettings
session = SQLiteSession("default_settings_test")
# Should have default SessionSettings
assert isinstance(session.session_settings, SessionSettings)
assert session.session_settings.limit is None
session.close()
@pytest.mark.asyncio
async def test_session_settings_constructor():
"""Test passing session_settings via constructor."""
from agents.memory import SessionSettings
session = SQLiteSession("constructor_settings_test", session_settings=SessionSettings(limit=5))
assert session.session_settings is not None
assert session.session_settings.limit == 5
session.close()
@pytest.mark.asyncio
async def test_session_settings_constructor_normalizes_dictionary() -> None:
session = SQLiteSession("dictionary_settings_test", session_settings={"limit": 0})
assert isinstance(session.session_settings, SessionSettings)
assert session.session_settings.limit == 0
assert session.session_settings.resolve({"limit": 4}).limit == 4
session.close()
def test_session_settings_rejects_unknown_dictionary_fields() -> None:
with pytest.raises(TypeError, match="Unknown session settings: limitt"):
SQLiteSession("invalid_settings_test", session_settings={"limitt": 1})
@pytest.mark.asyncio
async def test_get_items_uses_session_settings_limit():
"""Test that get_items uses session_settings.limit as default."""
from agents.memory import SessionSettings
with tempfile.TemporaryDirectory() as temp_dir:
db_path = Path(temp_dir) / "test_settings_limit.db"
session = SQLiteSession(
"uses_settings_limit_test", db_path, session_settings=SessionSettings(limit=3)
)
# Add 5 items
items: list[TResponseInputItem] = [
{"role": "user", "content": f"Message {i}"} for i in range(5)
]
await session.add_items(items)
# get_items() with no limit should use session_settings.limit=3
retrieved = await session.get_items()
assert len(retrieved) == 3
# Should get the last 3 items
assert retrieved[0].get("content") == "Message 2"
assert retrieved[1].get("content") == "Message 3"
assert retrieved[2].get("content") == "Message 4"
session.close()
@pytest.mark.asyncio
async def test_get_items_explicit_limit_overrides_session_settings():
"""Test that explicit limit parameter overrides session_settings."""
from agents.memory import SessionSettings
with tempfile.TemporaryDirectory() as temp_dir:
db_path = Path(temp_dir) / "test_override.db"
session = SQLiteSession(
"explicit_override_test", db_path, session_settings=SessionSettings(limit=5)
)
# Add 10 items
items: list[TResponseInputItem] = [
{"role": "user", "content": f"Message {i}"} for i in range(10)
]
await session.add_items(items)
# Explicit limit=2 should override session_settings.limit=5
retrieved = await session.get_items(limit=2)
assert len(retrieved) == 2
assert retrieved[0].get("content") == "Message 8"
assert retrieved[1].get("content") == "Message 9"
session.close()
@pytest.mark.asyncio
async def test_session_settings_resolve():
"""Test SessionSettings.resolve() method."""
from agents.memory import SessionSettings
base = SessionSettings(limit=100)
override = SessionSettings(limit=50)
final = base.resolve(override)
assert final.limit == 50 # Override wins
assert base.limit == 100 # Original unchanged
# Resolving with None returns self
final_none = base.resolve(None)
assert final_none.limit == 100
@pytest.mark.asyncio
async def test_runner_with_session_settings_override():
"""Test that RunConfig can override session's default settings."""
from agents.memory import SessionSettings
with tempfile.TemporaryDirectory() as temp_dir:
db_path = Path(temp_dir) / "test_runner_override.db"
# Session with default limit=100
session = SQLiteSession(
"runner_override_test", db_path, session_settings=SessionSettings(limit=100)
)
# Add some history
items: list[TResponseInputItem] = [
{"role": "user", "content": f"Turn {i}"} for i in range(10)
]
await session.add_items(items)
model = ScriptedModel()
agent = Agent(name="test", model=model)
model.enqueue([get_text_message("Got it")])
await Runner.run(
agent,
"New question",
session=session,
run_config=RunConfig(
session_settings=SessionSettings(limit=2) # Override to 2
),
)
# Verify the agent received only the last 2 history items + new question
last_input = model.calls[-1].input
# Filter out the new "New question" input
history_items = [item for item in last_input if item.get("content") != "New question"]
# Should have 2 history items (last two from the 10 we added)
assert len(history_items) == 2
session.close()
def _drop_sqlite_table(db_path: Path, table: str) -> None:
"""Drop a table from an independent connection to make a later statement fail."""
helper = sqlite3.connect(str(db_path))
try:
helper.execute(f"DROP TABLE {table}")
helper.commit()
finally:
helper.close()
def _sqlite_write_lock_is_free(db_path: Path) -> bool:
"""Return whether an independent writer can take the SQLite write lock."""
probe = sqlite3.connect(str(db_path), timeout=0)
try:
probe.execute("CREATE TABLE IF NOT EXISTS probe_lock (x INTEGER)")
probe.commit()
return True
except sqlite3.OperationalError:
return False
finally:
probe.close()
@pytest.mark.asyncio
async def test_sqlite_session_failed_clear_session_rolls_back():
"""A failed clear must restore earlier statements and release the cached write lock."""
with tempfile.TemporaryDirectory() as temp_dir:
db_path = Path(temp_dir) / "clear_rollback.db"
session = SQLiteSession("clear_rollback", db_path)
await session.add_items([{"role": "user", "content": "kept"}])
_drop_sqlite_table(db_path, "agent_sessions")
with pytest.raises(sqlite3.OperationalError):
await session.clear_session()
assert all(not conn.in_transaction for conn in session._connections)
assert _sqlite_write_lock_is_free(db_path)
session.close()
@pytest.mark.asyncio
async def test_sqlite_session_failed_pop_item_releases_write_lock():
"""A failed pop must not leave a write transaction on the cached connection."""
with tempfile.TemporaryDirectory() as temp_dir:
db_path = Path(temp_dir) / "pop_rollback.db"
session = SQLiteSession("pop_rollback", db_path)
await session.add_items([{"role": "user", "content": "kept"}])
_drop_sqlite_table(db_path, "agent_messages")
with pytest.raises(sqlite3.OperationalError):
await session.pop_item()
assert all(not conn.in_transaction for conn in session._connections)
assert _sqlite_write_lock_is_free(db_path)
session.close()
@pytest.mark.asyncio
async def test_sqlite_session_rollback_failure_evicts_connection(
monkeypatch: pytest.MonkeyPatch,
):
"""A file connection that cannot roll back must be closed and replaced."""
class FailingRollbackConnection(sqlite3.Connection):
def rollback(self) -> None:
raise RuntimeError("rollback failed")
with tempfile.TemporaryDirectory() as temp_dir:
db_path = Path(temp_dir) / "rollback_failure.db"
session = SQLiteSession("rollback_failure", db_path)
conn = sqlite3.connect(
str(db_path),
check_same_thread=False,
factory=FailingRollbackConnection,
)
with session._connections_lock:
session._connections.add(conn)
real_get_connection = session._get_connection
monkeypatch.setattr(session, "_get_connection", lambda: conn)
unserializable = cast(TResponseInputItem, {"role": "user", "content": object()})
with pytest.raises(TypeError):
await session.add_items([unserializable])
assert conn not in session._connections
assert _sqlite_write_lock_is_free(db_path)
monkeypatch.setattr(session, "_get_connection", real_get_connection)
await session.add_items([{"role": "user", "content": "after failure"}])
assert [item.get("content") for item in await session.get_items()] == ["after failure"]
session.close()
@pytest.mark.asyncio
async def test_sqlite_session_close_retries_quarantined_connection(
monkeypatch: pytest.MonkeyPatch,
):
"""A failed invalidation close must remain owned until a later close succeeds."""
class FailingRollbackAndCloseConnection(sqlite3.Connection):
fail_close = True
def rollback(self) -> None:
raise RuntimeError("rollback failed")
def close(self) -> None:
if self.fail_close:
raise RuntimeError("close failed")
super().close()
with tempfile.TemporaryDirectory() as temp_dir:
db_path = Path(temp_dir) / "close_retry.db"
session = SQLiteSession("close_retry", db_path)
conn = sqlite3.connect(
str(db_path),
check_same_thread=False,
factory=FailingRollbackAndCloseConnection,
)
with session._connections_lock:
session._connections.add(conn)
monkeypatch.setattr(session, "_get_connection", lambda: conn)
unserializable = cast(TResponseInputItem, {"role": "user", "content": object()})
with pytest.raises(TypeError):
await session.add_items([unserializable])
assert session._closed is True
assert conn in session._quarantined_connections
assert _sqlite_write_lock_is_free(db_path) is False
conn.fail_close = False
session.close()
assert session._quarantined_connections == set()
assert _sqlite_write_lock_is_free(db_path)
with pytest.raises(sqlite3.ProgrammingError):
conn.execute("SELECT 1")
@pytest.mark.asyncio
@pytest.mark.parametrize("operation", ["add", "pop", "clear"])
async def test_sqlite_session_post_commit_cancellation_propagates_after_known_outcome(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
operation: str,
):
"""Cancellation after a worker commit must propagate without inviting a retry."""
class PausingCommitConnection(sqlite3.Connection):
pause_commit = False
commit_finished = threading.Event()
allow_return = threading.Event()
def commit(self) -> None:
super().commit()
if self.pause_commit:
self.pause_commit = False
self.commit_finished.set()
assert self.allow_return.wait(timeout=10)
db_path = tmp_path / f"post_commit_{operation}.db"
session = SQLiteSession(f"post_commit_{operation}", db_path)
item: TResponseInputItem = {"role": "user", "content": "once"}
if operation != "add":
await session.add_items([item])
conn = sqlite3.connect(
str(db_path),
check_same_thread=False,
factory=PausingCommitConnection,
)
with session._connections_lock:
session._connections.add(conn)
monkeypatch.setattr(session, "_get_connection", lambda: conn)
conn.pause_commit = True
if operation == "add":
mutation: asyncio.Task[Any] = asyncio.create_task(session.add_items([item]))
elif operation == "pop":
mutation = asyncio.create_task(session.pop_item())
else:
mutation = asyncio.create_task(session.clear_session())
try:
assert await asyncio.to_thread(conn.commit_finished.wait, 10)
mutation.cancel()
await asyncio.sleep(0)
mutation.cancel()
await asyncio.sleep(0)
conn.allow_return.set()
with pytest.raises(asyncio.CancelledError):
await mutation
finally:
conn.allow_return.set()
if not mutation.done():
mutation.cancel()
await asyncio.gather(mutation, return_exceptions=True)
if operation == "add":
assert await session.get_items() == [item]
elif operation == "pop":
assert await session.get_items() == []
else:
assert await session.get_items() == []
assert mutation.cancelled()
session.close()