"""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()