Files
2026-07-13 12:39:17 +08:00

809 lines
29 KiB
Python

"""Tests for session memory functionality."""
import asyncio
import sqlite3
import tempfile
from pathlib import Path
import pytest
from agents import Agent, RunConfig, Runner, SQLiteSession, TResponseInputItem
from tests.fake_model import FakeModel
from tests.test_responses import get_text_message
# 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 = FakeModel()
agent = Agent(name="test", model=model)
# First turn
model.set_next_output([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.set_next_output([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.last_turn_args["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 = FakeModel()
agent = Agent(name="test", model=model)
# First turn
model.set_next_output([get_text_message("Hello")])
result1 = await run_agent_async(runner_method, agent, "Hi there", session=session)
assert result1.final_output == "Hello"
# Second turn
model.set_next_output([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 = FakeModel()
agent = Agent(name="test", model=model)
# First turn (no session parameters = disabled)
model.set_next_output([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.set_next_output([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.last_turn_args["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 = FakeModel()
agent = Agent(name="test", model=model)
# Session 1
session_id_1 = "session_1"
session_1 = SQLiteSession(session_id_1, db_path)
model.set_next_output([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.set_next_output([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.set_next_output([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_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.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 = FakeModel()
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.set_next_output([get_text_message("This should run")])
await run_agent_async(runner_method, agent, list_input, session=session)
assert model.last_turn_args["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 = FakeModel()
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.set_next_output([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.last_turn_args["input"]) == 2
assert model.last_turn_args["input"] == expected_model_input
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_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 = FakeModel()
agent = Agent(name="test", model=model)
model.set_next_output([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_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 = FakeModel()
agent = Agent(name="test", model=model)
model.set_next_output([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.last_turn_args["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()