chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,410 @@
|
||||
"""Tests for AsyncSQLiteSession functionality."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import tempfile
|
||||
from collections.abc import Sequence
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Any, cast
|
||||
|
||||
import pytest
|
||||
|
||||
pytest.importorskip("aiosqlite") # Skip tests if aiosqlite is not installed
|
||||
|
||||
from agents import Agent, Runner, TResponseInputItem
|
||||
from agents.extensions.memory import AsyncSQLiteSession
|
||||
from agents.memory import SessionSettings
|
||||
from tests.fake_model import FakeModel
|
||||
from tests.test_responses import get_text_message
|
||||
|
||||
pytestmark = pytest.mark.asyncio
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def agent() -> Agent:
|
||||
"""Fixture for a basic agent with a fake model."""
|
||||
return Agent(name="test", model=FakeModel())
|
||||
|
||||
|
||||
def _item_ids(items: Sequence[TResponseInputItem]) -> list[str]:
|
||||
result: list[str] = []
|
||||
for item in items:
|
||||
item_dict = cast(dict[str, Any], item)
|
||||
result.append(cast(str, item_dict["id"]))
|
||||
return result
|
||||
|
||||
|
||||
async def test_async_sqlite_session_basic_flow():
|
||||
"""Test AsyncSQLiteSession add/get/clear behavior."""
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
db_path = Path(temp_dir) / "async_basic.db"
|
||||
session = AsyncSQLiteSession("async_basic", db_path)
|
||||
|
||||
items: list[TResponseInputItem] = [
|
||||
{"role": "user", "content": "Hello"},
|
||||
{"role": "assistant", "content": "Hi there!"},
|
||||
]
|
||||
|
||||
await session.add_items(items)
|
||||
retrieved = await session.get_items()
|
||||
assert retrieved == items
|
||||
|
||||
await session.clear_session()
|
||||
assert await session.get_items() == []
|
||||
|
||||
await session.close()
|
||||
|
||||
|
||||
async def test_async_sqlite_session_pop_item():
|
||||
"""Test AsyncSQLiteSession pop_item behavior."""
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
db_path = Path(temp_dir) / "async_pop.db"
|
||||
session = AsyncSQLiteSession("async_pop", db_path)
|
||||
|
||||
assert await session.pop_item() is None
|
||||
|
||||
items: list[TResponseInputItem] = [
|
||||
{"role": "user", "content": "One"},
|
||||
{"role": "assistant", "content": "Two"},
|
||||
]
|
||||
await session.add_items(items)
|
||||
|
||||
popped = await session.pop_item()
|
||||
assert popped == items[-1]
|
||||
assert await session.get_items() == items[:-1]
|
||||
|
||||
await session.close()
|
||||
|
||||
|
||||
async def test_async_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) / "async_pop_corrupt.db"
|
||||
session = AsyncSQLiteSession("async_pop_corrupt", db_path)
|
||||
|
||||
valid_item: TResponseInputItem = {"role": "user", "content": "valid"}
|
||||
await session.add_items([valid_item])
|
||||
|
||||
conn = await session._get_connection()
|
||||
await conn.execute(
|
||||
f"INSERT INTO {session.messages_table} (session_id, message_data) VALUES (?, ?)",
|
||||
(session.session_id, "not valid json {{{"),
|
||||
)
|
||||
await conn.commit()
|
||||
|
||||
assert await session.pop_item() == valid_item
|
||||
assert await session.get_items() == []
|
||||
|
||||
await session.close()
|
||||
|
||||
|
||||
async def test_async_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) / "async_pop_only_corrupt.db"
|
||||
session = AsyncSQLiteSession("async_pop_only_corrupt", db_path)
|
||||
|
||||
conn = await session._get_connection()
|
||||
await conn.execute(
|
||||
f"INSERT INTO {session.messages_table} (session_id, message_data) VALUES (?, ?)",
|
||||
(session.session_id, "not valid json {{{"),
|
||||
)
|
||||
await conn.commit()
|
||||
|
||||
assert await session.pop_item() is None
|
||||
assert await session.get_items() == []
|
||||
|
||||
await session.close()
|
||||
|
||||
|
||||
async def test_async_sqlite_session_get_items_limit():
|
||||
"""Test AsyncSQLiteSession get_items limit handling."""
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
db_path = Path(temp_dir) / "async_limit.db"
|
||||
session = AsyncSQLiteSession("async_limit", db_path)
|
||||
|
||||
items: list[TResponseInputItem] = [
|
||||
{"role": "user", "content": "Message 1"},
|
||||
{"role": "assistant", "content": "Response 1"},
|
||||
{"role": "user", "content": "Message 2"},
|
||||
]
|
||||
await session.add_items(items)
|
||||
|
||||
latest = await session.get_items(limit=2)
|
||||
assert latest == items[-2:]
|
||||
|
||||
none = await session.get_items(limit=0)
|
||||
assert none == []
|
||||
|
||||
await session.close()
|
||||
|
||||
|
||||
async def test_async_sqlite_session_session_settings_default():
|
||||
"""Test that session_settings defaults to empty SessionSettings."""
|
||||
session = AsyncSQLiteSession("async_default_settings")
|
||||
|
||||
assert isinstance(session.session_settings, SessionSettings)
|
||||
assert session.session_settings.limit is None
|
||||
|
||||
await session.close()
|
||||
|
||||
|
||||
async def test_async_sqlite_session_session_settings_constructor():
|
||||
"""Test passing session_settings via constructor."""
|
||||
session = AsyncSQLiteSession(
|
||||
"async_constructor_settings",
|
||||
session_settings=SessionSettings(limit=5),
|
||||
)
|
||||
|
||||
assert session.session_settings is not None
|
||||
assert session.session_settings.limit == 5
|
||||
|
||||
await session.close()
|
||||
|
||||
|
||||
async def test_async_sqlite_session_get_items_uses_session_settings_limit():
|
||||
"""Test that get_items uses session_settings.limit as default."""
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
db_path = Path(temp_dir) / "async_settings_limit.db"
|
||||
session = AsyncSQLiteSession(
|
||||
"async_settings_limit",
|
||||
db_path,
|
||||
session_settings=SessionSettings(limit=3),
|
||||
)
|
||||
|
||||
items: list[TResponseInputItem] = [
|
||||
{"role": "user", "content": f"Message {i}"} for i in range(5)
|
||||
]
|
||||
await session.add_items(items)
|
||||
|
||||
retrieved = await session.get_items()
|
||||
assert retrieved == items[-3:]
|
||||
|
||||
await session.close()
|
||||
|
||||
|
||||
async def test_async_sqlite_session_explicit_limit_overrides_session_settings():
|
||||
"""Test that explicit limit parameter overrides session_settings."""
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
db_path = Path(temp_dir) / "async_settings_override.db"
|
||||
session = AsyncSQLiteSession(
|
||||
"async_settings_override",
|
||||
db_path,
|
||||
session_settings=SessionSettings(limit=5),
|
||||
)
|
||||
|
||||
items: list[TResponseInputItem] = [
|
||||
{"role": "user", "content": f"Message {i}"} for i in range(10)
|
||||
]
|
||||
await session.add_items(items)
|
||||
|
||||
retrieved = await session.get_items(limit=2)
|
||||
assert retrieved == items[-2:]
|
||||
|
||||
no_items = await session.get_items(limit=0)
|
||||
assert no_items == []
|
||||
|
||||
await session.close()
|
||||
|
||||
|
||||
async def test_async_sqlite_session_unicode_content():
|
||||
"""Test AsyncSQLiteSession stores unicode content."""
|
||||
session = AsyncSQLiteSession("async_unicode")
|
||||
items: list[TResponseInputItem] = [
|
||||
{"role": "user", "content": "こんにちは"},
|
||||
{"role": "assistant", "content": "Привет"},
|
||||
]
|
||||
await session.add_items(items)
|
||||
|
||||
retrieved = await session.get_items()
|
||||
assert retrieved == items
|
||||
|
||||
await session.close()
|
||||
|
||||
|
||||
async def test_async_sqlite_session_runner_integration(agent: Agent):
|
||||
"""Test that AsyncSQLiteSession works correctly with the agent Runner."""
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
db_path = Path(temp_dir) / "async_runner_integration.db"
|
||||
session = AsyncSQLiteSession("runner_integration_test", db_path)
|
||||
|
||||
assert isinstance(agent.model, FakeModel)
|
||||
|
||||
agent.model.set_next_output([get_text_message("San Francisco")])
|
||||
result1 = await Runner.run(
|
||||
agent,
|
||||
"What city is the Golden Gate Bridge in?",
|
||||
session=session,
|
||||
)
|
||||
assert result1.final_output == "San Francisco"
|
||||
|
||||
agent.model.set_next_output([get_text_message("California")])
|
||||
result2 = await Runner.run(agent, "What state is it in?", session=session)
|
||||
assert result2.final_output == "California"
|
||||
|
||||
last_input = agent.model.last_turn_args["input"]
|
||||
assert isinstance(last_input, list)
|
||||
assert len(last_input) > 1
|
||||
assert any("Golden Gate Bridge" in str(item.get("content", "")) for item in last_input)
|
||||
|
||||
await session.close()
|
||||
|
||||
|
||||
async def test_async_sqlite_session_session_isolation(agent: Agent):
|
||||
"""Test that different session IDs result in isolated conversation histories."""
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
db_path = Path(temp_dir) / "async_isolation.db"
|
||||
session1 = AsyncSQLiteSession("session_1", db_path)
|
||||
session2 = AsyncSQLiteSession("session_2", db_path)
|
||||
|
||||
assert isinstance(agent.model, FakeModel)
|
||||
agent.model.set_next_output([get_text_message("I like cats.")])
|
||||
await Runner.run(agent, "I like cats.", session=session1)
|
||||
|
||||
agent.model.set_next_output([get_text_message("I like dogs.")])
|
||||
await Runner.run(agent, "I like dogs.", session=session2)
|
||||
|
||||
agent.model.set_next_output([get_text_message("You said you like cats.")])
|
||||
result = await Runner.run(agent, "What animal did I say I like?", session=session1)
|
||||
assert "cats" in result.final_output.lower()
|
||||
assert "dogs" not in result.final_output.lower()
|
||||
|
||||
await session1.close()
|
||||
await session2.close()
|
||||
|
||||
|
||||
async def test_async_sqlite_session_add_empty_items_list():
|
||||
"""Test that adding an empty list of items is a no-op."""
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
db_path = Path(temp_dir) / "async_add_empty.db"
|
||||
session = AsyncSQLiteSession("add_empty_test", db_path)
|
||||
|
||||
assert await session.get_items() == []
|
||||
await session.add_items([])
|
||||
assert await session.get_items() == []
|
||||
|
||||
await session.close()
|
||||
|
||||
|
||||
async def test_async_sqlite_session_pop_from_empty_session():
|
||||
"""Test that pop_item returns None on an empty session."""
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
db_path = Path(temp_dir) / "async_pop_empty.db"
|
||||
session = AsyncSQLiteSession("empty_session", db_path)
|
||||
|
||||
popped = await session.pop_item()
|
||||
assert popped is None
|
||||
|
||||
await session.close()
|
||||
|
||||
|
||||
async def test_async_sqlite_session_get_items_with_limit_more_than_available():
|
||||
"""Test limit behavior when requesting more items than exist."""
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
db_path = Path(temp_dir) / "async_limit_more.db"
|
||||
session = AsyncSQLiteSession("limit_more_test", db_path)
|
||||
|
||||
items: list[TResponseInputItem] = [
|
||||
{"role": "user", "content": "1"},
|
||||
{"role": "assistant", "content": "2"},
|
||||
{"role": "user", "content": "3"},
|
||||
{"role": "assistant", "content": "4"},
|
||||
]
|
||||
await session.add_items(items)
|
||||
|
||||
retrieved = await session.get_items(limit=10)
|
||||
assert retrieved == items
|
||||
|
||||
await session.close()
|
||||
|
||||
|
||||
async def test_async_sqlite_session_get_items_same_timestamp_consistent_order():
|
||||
"""Test that items with identical timestamps keep insertion order."""
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
db_path = Path(temp_dir) / "async_same_timestamp.db"
|
||||
session = AsyncSQLiteSession("same_timestamp_test", db_path)
|
||||
|
||||
older_item = cast(
|
||||
TResponseInputItem, {"id": "older_same_ts", "role": "user", "content": "old"}
|
||||
)
|
||||
reasoning_item = cast(TResponseInputItem, {"id": "rs_same_ts", "type": "reasoning"})
|
||||
message_item = cast(
|
||||
TResponseInputItem,
|
||||
{"id": "msg_same_ts", "type": "message", "role": "assistant", "content": []},
|
||||
)
|
||||
|
||||
await session.add_items([older_item])
|
||||
await session.add_items([reasoning_item, message_item])
|
||||
|
||||
conn = await session._get_connection()
|
||||
cursor = await conn.execute(
|
||||
f"SELECT id, message_data FROM {session.messages_table} WHERE session_id = ?",
|
||||
(session.session_id,),
|
||||
)
|
||||
rows = await cursor.fetchall()
|
||||
await cursor.close()
|
||||
|
||||
id_map: dict[str, int] = {
|
||||
cast(str, json.loads(message_json)["id"]): cast(int, row_id)
|
||||
for row_id, message_json in rows
|
||||
}
|
||||
|
||||
shared = datetime(2025, 10, 15, 17, 26, 39, 132483)
|
||||
shared_str = shared.strftime("%Y-%m-%d %H:%M:%S.%f")
|
||||
await conn.execute(
|
||||
f"""
|
||||
UPDATE {session.messages_table}
|
||||
SET created_at = ?
|
||||
WHERE id IN (?, ?, ?)
|
||||
""",
|
||||
(
|
||||
shared_str,
|
||||
id_map["older_same_ts"],
|
||||
id_map["rs_same_ts"],
|
||||
id_map["msg_same_ts"],
|
||||
),
|
||||
)
|
||||
await conn.commit()
|
||||
|
||||
retrieved = await session.get_items()
|
||||
assert _item_ids(retrieved) == ["older_same_ts", "rs_same_ts", "msg_same_ts"]
|
||||
|
||||
latest_two = await session.get_items(limit=2)
|
||||
assert _item_ids(latest_two) == ["rs_same_ts", "msg_same_ts"]
|
||||
|
||||
await session.close()
|
||||
|
||||
|
||||
async def test_async_sqlite_session_pop_item_same_timestamp_returns_latest():
|
||||
"""Test that pop_item returns the newest item when timestamps tie."""
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
db_path = Path(temp_dir) / "async_same_timestamp_pop.db"
|
||||
session = AsyncSQLiteSession("same_timestamp_pop_test", db_path)
|
||||
|
||||
reasoning_item = cast(TResponseInputItem, {"id": "rs_pop_same_ts", "type": "reasoning"})
|
||||
message_item = cast(
|
||||
TResponseInputItem,
|
||||
{"id": "msg_pop_same_ts", "type": "message", "role": "assistant", "content": []},
|
||||
)
|
||||
|
||||
await session.add_items([reasoning_item, message_item])
|
||||
|
||||
conn = await session._get_connection()
|
||||
shared = datetime(2025, 10, 15, 17, 26, 39, 132483)
|
||||
shared_str = shared.strftime("%Y-%m-%d %H:%M:%S.%f")
|
||||
await conn.execute(
|
||||
f"UPDATE {session.messages_table} SET created_at = ? WHERE session_id = ?",
|
||||
(shared_str, session.session_id),
|
||||
)
|
||||
await conn.commit()
|
||||
|
||||
popped = await session.pop_item()
|
||||
assert popped is not None
|
||||
assert cast(dict[str, Any], popped)["id"] == "msg_pop_same_ts"
|
||||
|
||||
remaining = await session.get_items()
|
||||
assert _item_ids(remaining) == ["rs_pop_same_ts"]
|
||||
|
||||
await session.close()
|
||||
@@ -0,0 +1,594 @@
|
||||
"""
|
||||
Integration tests for DaprSession with real Dapr sidecar and Redis using testcontainers.
|
||||
|
||||
These tests use Docker containers for both Redis and Dapr, with proper networking.
|
||||
Tests are automatically skipped if dependencies (dapr, testcontainers, docker) are not available.
|
||||
|
||||
Run with: pytest tests/extensions/memory/test_dapr_redis_integration.py -v
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
import urllib.request
|
||||
|
||||
import docker # type: ignore[import-untyped]
|
||||
import pytest
|
||||
from docker.errors import DockerException # type: ignore[import-untyped]
|
||||
|
||||
# Skip tests if dependencies are not available
|
||||
pytest.importorskip("dapr") # Skip tests if Dapr is not installed
|
||||
pytest.importorskip("testcontainers") # Skip if testcontainers is not installed
|
||||
if sys.platform == "win32":
|
||||
pytest.skip(
|
||||
"Dapr Docker integration tests are not supported on Windows",
|
||||
allow_module_level=True,
|
||||
)
|
||||
if shutil.which("docker") is None:
|
||||
pytest.skip(
|
||||
"Docker executable is not available; skipping Dapr integration tests",
|
||||
allow_module_level=True,
|
||||
)
|
||||
try:
|
||||
client = docker.from_env()
|
||||
client.ping()
|
||||
except DockerException:
|
||||
pytest.skip(
|
||||
"Docker daemon is not available; skipping Dapr integration tests", allow_module_level=True
|
||||
)
|
||||
else:
|
||||
client.close()
|
||||
|
||||
from testcontainers.core.container import DockerContainer # type: ignore[import-untyped]
|
||||
from testcontainers.core.network import Network # type: ignore[import-untyped]
|
||||
from testcontainers.core.waiting_utils import wait_for_logs # type: ignore[import-untyped]
|
||||
|
||||
from agents import Agent, Runner, TResponseInputItem
|
||||
from agents.extensions.memory import (
|
||||
DAPR_CONSISTENCY_EVENTUAL,
|
||||
DAPR_CONSISTENCY_STRONG,
|
||||
DaprSession,
|
||||
)
|
||||
from tests.fake_model import FakeModel
|
||||
from tests.test_responses import get_text_message
|
||||
|
||||
# Docker-backed integration tests should stay on the serial test path.
|
||||
pytestmark = [pytest.mark.asyncio, pytest.mark.serial]
|
||||
|
||||
|
||||
def wait_for_dapr_health(host: str, port: int, timeout: int = 60) -> bool:
|
||||
"""
|
||||
Wait for Dapr sidecar to become healthy by checking the HTTP health endpoint.
|
||||
|
||||
Args:
|
||||
host: The host where Dapr is running
|
||||
port: The HTTP port (typically 3500)
|
||||
timeout: Maximum time to wait in seconds
|
||||
|
||||
Returns:
|
||||
True if Dapr becomes healthy, False otherwise
|
||||
"""
|
||||
health_url = f"http://{host}:{port}/v1.0/healthz/outbound"
|
||||
start_time = time.time()
|
||||
|
||||
while time.time() - start_time < timeout:
|
||||
try:
|
||||
with urllib.request.urlopen(health_url, timeout=5) as response:
|
||||
if 200 <= response.status < 300:
|
||||
print(f"✓ Dapr health check passed on {health_url}")
|
||||
return True
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
time.sleep(1)
|
||||
|
||||
print(f"✗ Dapr health check timed out after {timeout}s on {health_url}")
|
||||
return False
|
||||
|
||||
|
||||
def wait_for_dapr_component(host: str, port: int, component_name: str, timeout: int = 60) -> bool:
|
||||
"""Wait for a named component to appear in the Dapr metadata endpoint."""
|
||||
metadata_url = f"http://{host}:{port}/v1.0/metadata"
|
||||
start_time = time.time()
|
||||
|
||||
while time.time() - start_time < timeout:
|
||||
try:
|
||||
with urllib.request.urlopen(metadata_url, timeout=5) as response:
|
||||
if 200 <= response.status < 300:
|
||||
payload = json.load(response)
|
||||
components = payload.get("components", [])
|
||||
if any(component.get("name") == component_name for component in components):
|
||||
print(f"✓ Dapr component {component_name} loaded via {metadata_url}")
|
||||
return True
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
time.sleep(1)
|
||||
|
||||
print(f"✗ Dapr component {component_name} did not load after {timeout}s")
|
||||
return False
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def docker_network():
|
||||
"""Create a Docker network for container-to-container communication."""
|
||||
with Network() as network:
|
||||
yield network
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def redis_container(docker_network):
|
||||
"""Start Redis container on the shared network."""
|
||||
container = (
|
||||
DockerContainer("redis:7-alpine")
|
||||
.with_network(docker_network)
|
||||
.with_network_aliases("redis")
|
||||
.with_exposed_ports(6379)
|
||||
)
|
||||
container.start()
|
||||
wait_for_logs(container, "Ready to accept connections", timeout=30)
|
||||
try:
|
||||
yield container
|
||||
finally:
|
||||
container.stop()
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def dapr_container(redis_container, docker_network):
|
||||
"""Start Dapr sidecar container with Redis state store configuration."""
|
||||
# Create temporary components directory
|
||||
temp_dir = tempfile.mkdtemp()
|
||||
os.chmod(temp_dir, 0o755)
|
||||
components_path = os.path.join(temp_dir, "components")
|
||||
os.makedirs(components_path, exist_ok=True)
|
||||
os.chmod(components_path, 0o755)
|
||||
|
||||
# Write Redis state store component configuration
|
||||
# KEY: Use 'redis:6379' (network alias), NOT localhost!
|
||||
state_store_config = """
|
||||
apiVersion: dapr.io/v1alpha1
|
||||
kind: Component
|
||||
metadata:
|
||||
name: statestore
|
||||
spec:
|
||||
type: state.redis
|
||||
version: v1
|
||||
metadata:
|
||||
- name: redisHost
|
||||
value: redis:6379
|
||||
- name: redisPassword
|
||||
value: ""
|
||||
- name: actorStateStore
|
||||
value: "false"
|
||||
"""
|
||||
state_store_path = os.path.join(components_path, "statestore.yaml")
|
||||
with open(state_store_path, "w") as f:
|
||||
f.write(state_store_config)
|
||||
os.chmod(state_store_path, 0o644)
|
||||
|
||||
# Create Dapr container
|
||||
container = DockerContainer("daprio/daprd:latest")
|
||||
container = container.with_network(docker_network) # Join the same network
|
||||
container = container.with_volume_mapping(components_path, "/components", mode="ro")
|
||||
container = container.with_command(
|
||||
[
|
||||
"./daprd",
|
||||
"-app-id",
|
||||
"test-app",
|
||||
"-dapr-http-port",
|
||||
"3500", # HTTP API port for health checks
|
||||
"-dapr-grpc-port",
|
||||
"50001",
|
||||
"-resources-path",
|
||||
"/components",
|
||||
"-log-level",
|
||||
"info",
|
||||
]
|
||||
)
|
||||
container = container.with_exposed_ports(3500, 50001) # Expose both ports
|
||||
|
||||
container.start()
|
||||
|
||||
# Get the exposed HTTP port and host
|
||||
http_host = container.get_container_host_ip()
|
||||
http_port = container.get_exposed_port(3500)
|
||||
|
||||
# Wait for Dapr to become healthy
|
||||
if not wait_for_dapr_health(http_host, http_port, timeout=60):
|
||||
container.stop()
|
||||
pytest.fail("Dapr container failed to become healthy")
|
||||
|
||||
if not wait_for_dapr_component(http_host, http_port, "statestore", timeout=60):
|
||||
logs = container.get_wrapped_container().logs().decode("utf-8", errors="replace")
|
||||
container.stop()
|
||||
pytest.fail(f"Dapr state store component failed to load.\nContainer logs:\n{logs}")
|
||||
|
||||
# Set environment variables for Dapr SDK health checks
|
||||
# The Dapr SDK checks these when creating a client
|
||||
os.environ["DAPR_HTTP_PORT"] = str(http_port)
|
||||
os.environ["DAPR_RUNTIME_HOST"] = http_host
|
||||
|
||||
yield container
|
||||
|
||||
# Cleanup environment variables
|
||||
os.environ.pop("DAPR_HTTP_PORT", None)
|
||||
os.environ.pop("DAPR_RUNTIME_HOST", None)
|
||||
|
||||
container.stop()
|
||||
|
||||
# Cleanup
|
||||
import shutil
|
||||
|
||||
shutil.rmtree(temp_dir, ignore_errors=True)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def agent() -> Agent:
|
||||
"""Fixture for a basic agent with a fake model."""
|
||||
return Agent(name="test", model=FakeModel())
|
||||
|
||||
|
||||
async def test_dapr_redis_integration(dapr_container, monkeypatch):
|
||||
"""Test DaprSession with real Dapr sidecar and Redis backend."""
|
||||
# Get Dapr gRPC address (exposed to host)
|
||||
dapr_host = dapr_container.get_container_host_ip()
|
||||
dapr_port = dapr_container.get_exposed_port(50001)
|
||||
dapr_address = f"{dapr_host}:{dapr_port}"
|
||||
|
||||
# Monkeypatch the Dapr health check since we already verified it in the fixture
|
||||
from dapr.clients.health import DaprHealth
|
||||
|
||||
monkeypatch.setattr(DaprHealth, "wait_until_ready", lambda: None)
|
||||
|
||||
# Create session using from_address
|
||||
session = DaprSession.from_address(
|
||||
session_id="integration_test_session",
|
||||
state_store_name="statestore",
|
||||
dapr_address=dapr_address,
|
||||
)
|
||||
|
||||
try:
|
||||
# Test connectivity
|
||||
is_connected = await session.ping()
|
||||
assert is_connected is True
|
||||
|
||||
# Clear any existing data
|
||||
await session.clear_session()
|
||||
|
||||
# Test add_items
|
||||
items: list[TResponseInputItem] = [
|
||||
{"role": "user", "content": "Hello from integration test"},
|
||||
{"role": "assistant", "content": "Hi there!"},
|
||||
]
|
||||
await session.add_items(items)
|
||||
|
||||
# Test get_items
|
||||
retrieved = await session.get_items()
|
||||
assert len(retrieved) == 2
|
||||
assert retrieved[0].get("content") == "Hello from integration test"
|
||||
assert retrieved[1].get("content") == "Hi there!"
|
||||
|
||||
# Test get_items with limit
|
||||
latest_1 = await session.get_items(limit=1)
|
||||
assert len(latest_1) == 1
|
||||
assert latest_1[0].get("content") == "Hi there!"
|
||||
|
||||
# Test pop_item
|
||||
popped = await session.pop_item()
|
||||
assert popped is not None
|
||||
assert popped.get("content") == "Hi there!"
|
||||
|
||||
remaining = await session.get_items()
|
||||
assert len(remaining) == 1
|
||||
assert remaining[0].get("content") == "Hello from integration test"
|
||||
|
||||
# Test clear_session
|
||||
await session.clear_session()
|
||||
cleared = await session.get_items()
|
||||
assert len(cleared) == 0
|
||||
|
||||
finally:
|
||||
await session.close()
|
||||
|
||||
|
||||
async def test_dapr_runner_integration(agent: Agent, dapr_container, monkeypatch):
|
||||
"""Test DaprSession with agent Runner using real Dapr sidecar."""
|
||||
from dapr.clients.health import DaprHealth
|
||||
|
||||
monkeypatch.setattr(DaprHealth, "wait_until_ready", lambda: None)
|
||||
|
||||
dapr_host = dapr_container.get_container_host_ip()
|
||||
dapr_port = dapr_container.get_exposed_port(50001)
|
||||
dapr_address = f"{dapr_host}:{dapr_port}"
|
||||
|
||||
session = DaprSession.from_address(
|
||||
session_id="runner_integration_test",
|
||||
state_store_name="statestore",
|
||||
dapr_address=dapr_address,
|
||||
)
|
||||
|
||||
try:
|
||||
await session.clear_session()
|
||||
|
||||
# First turn
|
||||
assert isinstance(agent.model, FakeModel)
|
||||
agent.model.set_next_output([get_text_message("San Francisco")])
|
||||
result1 = await Runner.run(
|
||||
agent,
|
||||
"What city is the Golden Gate Bridge in?",
|
||||
session=session,
|
||||
)
|
||||
assert result1.final_output == "San Francisco"
|
||||
|
||||
# Second turn - should remember context
|
||||
agent.model.set_next_output([get_text_message("California")])
|
||||
result2 = await Runner.run(agent, "What state is it in?", session=session)
|
||||
assert result2.final_output == "California"
|
||||
|
||||
# Verify history
|
||||
last_input = agent.model.last_turn_args["input"]
|
||||
assert len(last_input) > 1
|
||||
assert any("Golden Gate Bridge" in str(item.get("content", "")) for item in last_input)
|
||||
|
||||
finally:
|
||||
await session.close()
|
||||
|
||||
|
||||
async def test_dapr_session_isolation(dapr_container, monkeypatch):
|
||||
"""Test that different session IDs are isolated with real Dapr."""
|
||||
from dapr.clients.health import DaprHealth
|
||||
|
||||
monkeypatch.setattr(DaprHealth, "wait_until_ready", lambda: None)
|
||||
|
||||
dapr_host = dapr_container.get_container_host_ip()
|
||||
dapr_port = dapr_container.get_exposed_port(50001)
|
||||
dapr_address = f"{dapr_host}:{dapr_port}"
|
||||
|
||||
session1 = DaprSession.from_address(
|
||||
session_id="isolated_session_1",
|
||||
state_store_name="statestore",
|
||||
dapr_address=dapr_address,
|
||||
)
|
||||
session2 = DaprSession.from_address(
|
||||
session_id="isolated_session_2",
|
||||
state_store_name="statestore",
|
||||
dapr_address=dapr_address,
|
||||
)
|
||||
|
||||
try:
|
||||
# Clear both sessions
|
||||
await session1.clear_session()
|
||||
await session2.clear_session()
|
||||
|
||||
# Add different data to each session
|
||||
await session1.add_items([{"role": "user", "content": "session 1 data"}])
|
||||
await session2.add_items([{"role": "user", "content": "session 2 data"}])
|
||||
|
||||
# Verify isolation
|
||||
items1 = await session1.get_items()
|
||||
items2 = await session2.get_items()
|
||||
|
||||
assert len(items1) == 1
|
||||
assert len(items2) == 1
|
||||
assert items1[0].get("content") == "session 1 data"
|
||||
assert items2[0].get("content") == "session 2 data"
|
||||
|
||||
finally:
|
||||
await session1.clear_session()
|
||||
await session2.clear_session()
|
||||
await session1.close()
|
||||
await session2.close()
|
||||
|
||||
|
||||
async def test_dapr_ttl_functionality(dapr_container, monkeypatch):
|
||||
"""Test TTL functionality with real Dapr and Redis (if supported by state store)."""
|
||||
from dapr.clients.health import DaprHealth
|
||||
|
||||
monkeypatch.setattr(DaprHealth, "wait_until_ready", lambda: None)
|
||||
|
||||
dapr_host = dapr_container.get_container_host_ip()
|
||||
dapr_port = dapr_container.get_exposed_port(50001)
|
||||
dapr_address = f"{dapr_host}:{dapr_port}"
|
||||
|
||||
# Create session with short TTL
|
||||
session = DaprSession.from_address(
|
||||
session_id="ttl_test_session",
|
||||
state_store_name="statestore",
|
||||
dapr_address=dapr_address,
|
||||
ttl=2, # 2 seconds TTL
|
||||
)
|
||||
|
||||
try:
|
||||
await session.clear_session()
|
||||
|
||||
# Add items with TTL
|
||||
items: list[TResponseInputItem] = [
|
||||
{"role": "user", "content": "This should expire soon"},
|
||||
]
|
||||
await session.add_items(items)
|
||||
|
||||
# Verify items exist immediately
|
||||
retrieved = await session.get_items()
|
||||
assert len(retrieved) == 1
|
||||
|
||||
# Note: Actual expiration testing depends on state store TTL support
|
||||
# Redis state store supports TTL via ttlInSeconds metadata
|
||||
|
||||
finally:
|
||||
await session.clear_session()
|
||||
await session.close()
|
||||
|
||||
|
||||
async def test_dapr_consistency_levels(dapr_container, monkeypatch):
|
||||
"""Test different consistency levels with real Dapr."""
|
||||
from dapr.clients.health import DaprHealth
|
||||
|
||||
monkeypatch.setattr(DaprHealth, "wait_until_ready", lambda: None)
|
||||
|
||||
dapr_host = dapr_container.get_container_host_ip()
|
||||
dapr_port = dapr_container.get_exposed_port(50001)
|
||||
dapr_address = f"{dapr_host}:{dapr_port}"
|
||||
|
||||
# Test eventual consistency
|
||||
session_eventual = DaprSession.from_address(
|
||||
session_id="eventual_consistency_test",
|
||||
state_store_name="statestore",
|
||||
dapr_address=dapr_address,
|
||||
consistency=DAPR_CONSISTENCY_EVENTUAL,
|
||||
)
|
||||
|
||||
# Test strong consistency
|
||||
session_strong = DaprSession.from_address(
|
||||
session_id="strong_consistency_test",
|
||||
state_store_name="statestore",
|
||||
dapr_address=dapr_address,
|
||||
consistency=DAPR_CONSISTENCY_STRONG,
|
||||
)
|
||||
|
||||
try:
|
||||
await session_eventual.clear_session()
|
||||
await session_strong.clear_session()
|
||||
|
||||
# Both should work correctly
|
||||
items: list[TResponseInputItem] = [{"role": "user", "content": "Consistency test"}]
|
||||
|
||||
await session_eventual.add_items(items)
|
||||
retrieved_eventual = await session_eventual.get_items()
|
||||
assert len(retrieved_eventual) == 1
|
||||
|
||||
await session_strong.add_items(items)
|
||||
retrieved_strong = await session_strong.get_items()
|
||||
assert len(retrieved_strong) == 1
|
||||
|
||||
finally:
|
||||
await session_eventual.clear_session()
|
||||
await session_strong.clear_session()
|
||||
await session_eventual.close()
|
||||
await session_strong.close()
|
||||
|
||||
|
||||
async def test_dapr_unicode_and_special_chars(dapr_container, monkeypatch):
|
||||
"""Test unicode and special characters with real Dapr and Redis."""
|
||||
from dapr.clients.health import DaprHealth
|
||||
|
||||
monkeypatch.setattr(DaprHealth, "wait_until_ready", lambda: None)
|
||||
|
||||
dapr_host = dapr_container.get_container_host_ip()
|
||||
dapr_port = dapr_container.get_exposed_port(50001)
|
||||
dapr_address = f"{dapr_host}:{dapr_port}"
|
||||
|
||||
session = DaprSession.from_address(
|
||||
session_id="unicode_test_session",
|
||||
state_store_name="statestore",
|
||||
dapr_address=dapr_address,
|
||||
)
|
||||
|
||||
try:
|
||||
await session.clear_session()
|
||||
|
||||
# Test unicode content
|
||||
items: list[TResponseInputItem] = [
|
||||
{"role": "user", "content": "こんにちは"},
|
||||
{"role": "assistant", "content": "😊👍"},
|
||||
{"role": "user", "content": "Привет"},
|
||||
{"role": "assistant", "content": '{"nested": "json"}'},
|
||||
{"role": "user", "content": "Line1\nLine2\tTabbed"},
|
||||
]
|
||||
await session.add_items(items)
|
||||
|
||||
# Retrieve and verify
|
||||
retrieved = await session.get_items()
|
||||
assert len(retrieved) == 5
|
||||
assert retrieved[0].get("content") == "こんにちは"
|
||||
assert retrieved[1].get("content") == "😊👍"
|
||||
assert retrieved[2].get("content") == "Привет"
|
||||
assert retrieved[3].get("content") == '{"nested": "json"}'
|
||||
assert retrieved[4].get("content") == "Line1\nLine2\tTabbed"
|
||||
|
||||
finally:
|
||||
await session.clear_session()
|
||||
await session.close()
|
||||
|
||||
|
||||
async def test_dapr_concurrent_writes_resolution(dapr_container, monkeypatch):
|
||||
"""
|
||||
Concurrent writes from multiple session instances should resolve via
|
||||
optimistic concurrency.
|
||||
"""
|
||||
from dapr.clients.health import DaprHealth
|
||||
|
||||
monkeypatch.setattr(DaprHealth, "wait_until_ready", lambda: None)
|
||||
|
||||
dapr_host = dapr_container.get_container_host_ip()
|
||||
dapr_port = dapr_container.get_exposed_port(50001)
|
||||
dapr_address = f"{dapr_host}:{dapr_port}"
|
||||
|
||||
# Use two different session objects pointing to the same logical session_id
|
||||
# to create real contention.
|
||||
session_id = "concurrent_integration_session"
|
||||
s1 = DaprSession.from_address(
|
||||
session_id=session_id,
|
||||
state_store_name="statestore",
|
||||
dapr_address=dapr_address,
|
||||
)
|
||||
s2 = DaprSession.from_address(
|
||||
session_id=session_id,
|
||||
state_store_name="statestore",
|
||||
dapr_address=dapr_address,
|
||||
)
|
||||
|
||||
try:
|
||||
# Clean slate.
|
||||
await s1.clear_session()
|
||||
|
||||
# Fire multiple parallel add_items calls from two different session instances.
|
||||
tasks: list[asyncio.Task[None]] = []
|
||||
for i in range(10):
|
||||
tasks.append(
|
||||
asyncio.create_task(
|
||||
s1.add_items(
|
||||
[
|
||||
{"role": "user", "content": f"A-{i}"},
|
||||
]
|
||||
)
|
||||
)
|
||||
)
|
||||
tasks.append(
|
||||
asyncio.create_task(
|
||||
s2.add_items(
|
||||
[
|
||||
{"role": "assistant", "content": f"B-{i}"},
|
||||
]
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
await asyncio.gather(*tasks)
|
||||
|
||||
# Validate all messages were persisted.
|
||||
# Use a fresh session object for readback to avoid any local caching
|
||||
# (none expected, but explicit).
|
||||
s_read = DaprSession.from_address(
|
||||
session_id=session_id,
|
||||
state_store_name="statestore",
|
||||
dapr_address=dapr_address,
|
||||
)
|
||||
try:
|
||||
items = await s_read.get_items()
|
||||
contents = [item.get("content") for item in items]
|
||||
# We expect 20 total messages: A-0..9 and B-0..9 (order unspecified).
|
||||
assert len(contents) == 20
|
||||
for i in range(10):
|
||||
assert f"A-{i}" in contents
|
||||
assert f"B-{i}" in contents
|
||||
finally:
|
||||
await s_read.close()
|
||||
finally:
|
||||
await s1.close()
|
||||
await s2.close()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,576 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from typing import cast
|
||||
|
||||
import pytest
|
||||
|
||||
pytest.importorskip("cryptography") # Skip tests if cryptography is not installed
|
||||
|
||||
from cryptography.fernet import Fernet
|
||||
|
||||
from agents import Agent, Runner, SessionSettings, SQLiteSession, TResponseInputItem
|
||||
from agents.extensions.memory.encrypt_session import EncryptedSession
|
||||
from tests.fake_model import FakeModel
|
||||
from tests.test_responses import get_text_message
|
||||
|
||||
# Mark all tests in this file as asyncio
|
||||
pytestmark = pytest.mark.asyncio
|
||||
|
||||
|
||||
def _invalid_encrypted_envelope() -> TResponseInputItem:
|
||||
return cast(
|
||||
TResponseInputItem,
|
||||
{"__enc__": 1, "v": 1, "kid": "hkdf-v1", "payload": "not-a-valid-token"},
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def agent() -> Agent:
|
||||
"""Fixture for a basic agent with a fake model."""
|
||||
return Agent(name="test", model=FakeModel())
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def encryption_key() -> str:
|
||||
"""Fixture for a valid Fernet encryption key."""
|
||||
return str(Fernet.generate_key().decode("utf-8"))
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def set_fernet_time(monkeypatch):
|
||||
"""Freeze Fernet TTL checks so expiration tests avoid real waiting."""
|
||||
current_time = 1_000
|
||||
|
||||
def _set_time(value: int) -> None:
|
||||
nonlocal current_time
|
||||
current_time = value
|
||||
|
||||
monkeypatch.setattr("cryptography.fernet.time.time", lambda: current_time)
|
||||
return _set_time
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def underlying_session():
|
||||
"""Fixture for an underlying SQLite session."""
|
||||
temp_dir = tempfile.mkdtemp()
|
||||
db_path = Path(temp_dir) / "test_encrypt.db"
|
||||
return SQLiteSession("test_session", db_path)
|
||||
|
||||
|
||||
async def test_encrypted_session_basic_functionality(
|
||||
agent: Agent, encryption_key: str, underlying_session: SQLiteSession
|
||||
):
|
||||
"""Test basic encryption/decryption functionality."""
|
||||
session = EncryptedSession(
|
||||
session_id="test_session",
|
||||
underlying_session=underlying_session,
|
||||
encryption_key=encryption_key,
|
||||
ttl=600,
|
||||
)
|
||||
|
||||
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("content") == "Hello"
|
||||
assert retrieved[1].get("content") == "Hi there!"
|
||||
|
||||
encrypted_items = await underlying_session.get_items()
|
||||
assert encrypted_items[0].get("__enc__") == 1
|
||||
assert "payload" in encrypted_items[0]
|
||||
assert encrypted_items[0].get("content") != "Hello"
|
||||
|
||||
underlying_session.close()
|
||||
|
||||
|
||||
async def test_encrypted_session_with_runner(
|
||||
agent: Agent, encryption_key: str, underlying_session: SQLiteSession
|
||||
):
|
||||
"""Test that EncryptedSession works with Runner."""
|
||||
session = EncryptedSession(
|
||||
session_id="test_session",
|
||||
underlying_session=underlying_session,
|
||||
encryption_key=encryption_key,
|
||||
)
|
||||
|
||||
assert isinstance(agent.model, FakeModel)
|
||||
agent.model.set_next_output([get_text_message("San Francisco")])
|
||||
result1 = await Runner.run(
|
||||
agent,
|
||||
"What city is the Golden Gate Bridge in?",
|
||||
session=session,
|
||||
)
|
||||
assert result1.final_output == "San Francisco"
|
||||
|
||||
agent.model.set_next_output([get_text_message("California")])
|
||||
result2 = await Runner.run(agent, "What state is it in?", session=session)
|
||||
assert result2.final_output == "California"
|
||||
|
||||
last_input = agent.model.last_turn_args["input"]
|
||||
assert len(last_input) > 1
|
||||
assert any("Golden Gate Bridge" in str(item.get("content", "")) for item in last_input)
|
||||
|
||||
underlying_session.close()
|
||||
|
||||
|
||||
async def test_encrypted_session_pop_item(encryption_key: str, underlying_session: SQLiteSession):
|
||||
"""Test pop_item functionality."""
|
||||
session = EncryptedSession(
|
||||
session_id="test_session",
|
||||
underlying_session=underlying_session,
|
||||
encryption_key=encryption_key,
|
||||
)
|
||||
|
||||
items: list[TResponseInputItem] = [
|
||||
{"role": "user", "content": "First"},
|
||||
{"role": "assistant", "content": "Second"},
|
||||
]
|
||||
await session.add_items(items)
|
||||
|
||||
popped = await session.pop_item()
|
||||
assert popped is not None
|
||||
assert popped.get("content") == "Second"
|
||||
|
||||
remaining = await session.get_items()
|
||||
assert len(remaining) == 1
|
||||
assert remaining[0].get("content") == "First"
|
||||
|
||||
underlying_session.close()
|
||||
|
||||
|
||||
async def test_encrypted_session_clear(encryption_key: str, underlying_session: SQLiteSession):
|
||||
"""Test clear_session functionality."""
|
||||
session = EncryptedSession(
|
||||
session_id="test_session",
|
||||
underlying_session=underlying_session,
|
||||
encryption_key=encryption_key,
|
||||
)
|
||||
|
||||
await session.add_items([{"role": "user", "content": "Test"}])
|
||||
await session.clear_session()
|
||||
|
||||
items = await session.get_items()
|
||||
assert len(items) == 0
|
||||
|
||||
underlying_session.close()
|
||||
|
||||
|
||||
async def test_encrypted_session_ttl_expiration(
|
||||
encryption_key: str, underlying_session: SQLiteSession, set_fernet_time
|
||||
):
|
||||
"""Test TTL expiration - expired items are silently skipped."""
|
||||
set_fernet_time(1_000)
|
||||
session = EncryptedSession(
|
||||
session_id="test_session",
|
||||
underlying_session=underlying_session,
|
||||
encryption_key=encryption_key,
|
||||
ttl=1, # 1 second TTL
|
||||
)
|
||||
|
||||
items: list[TResponseInputItem] = [
|
||||
{"role": "user", "content": "Hello"},
|
||||
{"role": "assistant", "content": "Hi"},
|
||||
]
|
||||
await session.add_items(items)
|
||||
|
||||
set_fernet_time(1_002)
|
||||
|
||||
retrieved = await session.get_items()
|
||||
assert len(retrieved) == 0
|
||||
|
||||
underlying_items = await underlying_session.get_items()
|
||||
assert len(underlying_items) == 2
|
||||
|
||||
underlying_session.close()
|
||||
|
||||
|
||||
async def test_encrypted_session_pop_expired(
|
||||
encryption_key: str, underlying_session: SQLiteSession, set_fernet_time
|
||||
):
|
||||
"""Test pop_item with expired data."""
|
||||
set_fernet_time(1_000)
|
||||
session = EncryptedSession(
|
||||
session_id="test_session",
|
||||
underlying_session=underlying_session,
|
||||
encryption_key=encryption_key,
|
||||
ttl=1,
|
||||
)
|
||||
|
||||
await session.add_items([{"role": "user", "content": "Test"}])
|
||||
set_fernet_time(1_002)
|
||||
|
||||
popped = await session.pop_item()
|
||||
assert popped is None
|
||||
|
||||
underlying_session.close()
|
||||
|
||||
|
||||
async def test_encrypted_session_pop_mixed_expired_valid(
|
||||
encryption_key: str, underlying_session: SQLiteSession, set_fernet_time
|
||||
):
|
||||
"""Test pop_item auto-retry with mixed expired and valid items."""
|
||||
set_fernet_time(1_000)
|
||||
session = EncryptedSession(
|
||||
session_id="test_session",
|
||||
underlying_session=underlying_session,
|
||||
encryption_key=encryption_key,
|
||||
ttl=2, # 2 second TTL
|
||||
)
|
||||
|
||||
await session.add_items(
|
||||
[
|
||||
{"role": "user", "content": "Old message 1"},
|
||||
{"role": "assistant", "content": "Old response 1"},
|
||||
]
|
||||
)
|
||||
|
||||
set_fernet_time(1_003)
|
||||
|
||||
await session.add_items(
|
||||
[
|
||||
{"role": "user", "content": "New message"},
|
||||
{"role": "assistant", "content": "New response"},
|
||||
]
|
||||
)
|
||||
|
||||
popped = await session.pop_item()
|
||||
assert popped is not None
|
||||
assert popped.get("content") == "New response"
|
||||
|
||||
popped2 = await session.pop_item()
|
||||
assert popped2 is not None
|
||||
assert popped2.get("content") == "New message"
|
||||
|
||||
popped3 = await session.pop_item()
|
||||
assert popped3 is None
|
||||
|
||||
underlying_session.close()
|
||||
|
||||
|
||||
async def test_encrypted_session_raw_string_key(underlying_session: SQLiteSession):
|
||||
"""Test using raw string as encryption key (not base64)."""
|
||||
session = EncryptedSession(
|
||||
session_id="test_session",
|
||||
underlying_session=underlying_session,
|
||||
encryption_key="my-secret-password", # Raw string, not Fernet key
|
||||
)
|
||||
|
||||
await session.add_items([{"role": "user", "content": "Test"}])
|
||||
items = await session.get_items()
|
||||
assert len(items) == 1
|
||||
assert items[0].get("content") == "Test"
|
||||
|
||||
underlying_session.close()
|
||||
|
||||
|
||||
async def test_encrypted_session_get_items_limit(
|
||||
encryption_key: str, underlying_session: SQLiteSession
|
||||
):
|
||||
"""Test get_items with limit parameter."""
|
||||
session = EncryptedSession(
|
||||
session_id="test_session",
|
||||
underlying_session=underlying_session,
|
||||
encryption_key=encryption_key,
|
||||
)
|
||||
|
||||
items: list[TResponseInputItem] = [
|
||||
{"role": "user", "content": f"Message {i}"} for i in range(5)
|
||||
]
|
||||
await session.add_items(items)
|
||||
|
||||
limited = await session.get_items(limit=2)
|
||||
assert len(limited) == 2
|
||||
assert limited[0].get("content") == "Message 3" # Latest 2
|
||||
assert limited[1].get("content") == "Message 4"
|
||||
|
||||
underlying_session.close()
|
||||
|
||||
|
||||
async def test_encrypted_session_get_items_limit_skips_invalid_latest_envelope(
|
||||
encryption_key: str, underlying_session: SQLiteSession
|
||||
):
|
||||
"""Test that limit counts valid decrypted items, not encrypted envelopes."""
|
||||
session = EncryptedSession(
|
||||
session_id="test_session",
|
||||
underlying_session=underlying_session,
|
||||
encryption_key=encryption_key,
|
||||
)
|
||||
|
||||
await session.add_items([{"role": "user", "content": "older valid"}])
|
||||
await underlying_session.add_items([_invalid_encrypted_envelope()])
|
||||
|
||||
all_items = await session.get_items()
|
||||
assert [item.get("content") for item in all_items] == ["older valid"]
|
||||
|
||||
limited = await session.get_items(limit=1)
|
||||
assert [item.get("content") for item in limited] == ["older valid"]
|
||||
|
||||
underlying_session.close()
|
||||
|
||||
|
||||
async def test_encrypted_session_get_items_limit_returns_latest_valid_items_after_invalids(
|
||||
encryption_key: str, underlying_session: SQLiteSession
|
||||
):
|
||||
"""Test that invalid envelopes do not hide earlier valid items from limit."""
|
||||
session = EncryptedSession(
|
||||
session_id="test_session",
|
||||
underlying_session=underlying_session,
|
||||
encryption_key=encryption_key,
|
||||
)
|
||||
|
||||
await session.add_items(
|
||||
[
|
||||
{"role": "user", "content": "valid 0"},
|
||||
{"role": "assistant", "content": "valid 1"},
|
||||
]
|
||||
)
|
||||
await underlying_session.add_items([_invalid_encrypted_envelope()])
|
||||
await session.add_items([{"role": "user", "content": "valid 2"}])
|
||||
|
||||
limited = await session.get_items(limit=2)
|
||||
assert [item.get("content") for item in limited] == ["valid 1", "valid 2"]
|
||||
|
||||
underlying_session.close()
|
||||
|
||||
|
||||
async def test_encrypted_session_get_items_session_settings_limit_skips_invalid_envelopes(
|
||||
encryption_key: str, underlying_session: SQLiteSession
|
||||
):
|
||||
"""Test that session settings limit counts valid decrypted items."""
|
||||
underlying_session.session_settings = SessionSettings(limit=3)
|
||||
session = EncryptedSession(
|
||||
session_id="test_session",
|
||||
underlying_session=underlying_session,
|
||||
encryption_key=encryption_key,
|
||||
)
|
||||
|
||||
await session.add_items(
|
||||
[
|
||||
{"role": "user", "content": "valid 0"},
|
||||
{"role": "assistant", "content": "valid 1"},
|
||||
{"role": "user", "content": "valid 2"},
|
||||
]
|
||||
)
|
||||
await underlying_session.add_items([_invalid_encrypted_envelope()])
|
||||
|
||||
items = await session.get_items()
|
||||
assert [item.get("content") for item in items] == ["valid 0", "valid 1", "valid 2"]
|
||||
|
||||
underlying_session.close()
|
||||
|
||||
|
||||
async def test_encrypted_session_unicode_content(
|
||||
encryption_key: str, underlying_session: SQLiteSession
|
||||
):
|
||||
"""Test encryption of international text content."""
|
||||
session = EncryptedSession(
|
||||
session_id="test_session",
|
||||
underlying_session=underlying_session,
|
||||
encryption_key=encryption_key,
|
||||
)
|
||||
|
||||
items: list[TResponseInputItem] = [
|
||||
{"role": "user", "content": "Hello world"},
|
||||
{"role": "assistant", "content": "Special chars: áéíóú"},
|
||||
{"role": "user", "content": "Numbers and symbols: 123!@#"},
|
||||
]
|
||||
await session.add_items(items)
|
||||
|
||||
retrieved = await session.get_items()
|
||||
assert retrieved[0].get("content") == "Hello world"
|
||||
assert retrieved[1].get("content") == "Special chars: áéíóú"
|
||||
assert retrieved[2].get("content") == "Numbers and symbols: 123!@#"
|
||||
|
||||
underlying_session.close()
|
||||
|
||||
|
||||
class CustomSession(SQLiteSession):
|
||||
"""Mock custom session with additional methods for testing delegation."""
|
||||
|
||||
def get_stats(self) -> dict[str, int]:
|
||||
"""Custom method that should be accessible through delegation."""
|
||||
return {"custom_method_calls": 42, "test_value": 123}
|
||||
|
||||
async def custom_async_method(self) -> str:
|
||||
"""Custom async method for testing delegation."""
|
||||
return "custom_async_result"
|
||||
|
||||
|
||||
async def test_encrypted_session_delegation():
|
||||
"""Test that custom methods on underlying session are accessible through delegation."""
|
||||
temp_dir = tempfile.mkdtemp()
|
||||
db_path = Path(temp_dir) / "test_delegation.db"
|
||||
underlying_session = CustomSession("test_session", db_path)
|
||||
|
||||
encryption_key = str(Fernet.generate_key().decode("utf-8"))
|
||||
session = EncryptedSession(
|
||||
session_id="test_session",
|
||||
underlying_session=underlying_session,
|
||||
encryption_key=encryption_key,
|
||||
)
|
||||
|
||||
stats = session.get_stats()
|
||||
assert stats == {"custom_method_calls": 42, "test_value": 123}
|
||||
|
||||
result = await session.custom_async_method()
|
||||
assert result == "custom_async_result"
|
||||
|
||||
await session.add_items([{"role": "user", "content": "Test delegation"}])
|
||||
items = await session.get_items()
|
||||
assert len(items) == 1
|
||||
assert items[0].get("content") == "Test delegation"
|
||||
|
||||
underlying_session.close()
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# SessionSettings Tests
|
||||
# ============================================================================
|
||||
|
||||
|
||||
async def test_session_settings_delegated_to_underlying(encryption_key: str):
|
||||
"""Test that session_settings is correctly delegated to underlying session."""
|
||||
from agents.memory import SessionSettings
|
||||
|
||||
temp_dir = tempfile.mkdtemp()
|
||||
db_path = Path(temp_dir) / "test_settings.db"
|
||||
underlying = SQLiteSession("test_session", db_path, session_settings=SessionSettings(limit=5))
|
||||
|
||||
session = EncryptedSession(
|
||||
session_id="test_session",
|
||||
underlying_session=underlying,
|
||||
encryption_key=encryption_key,
|
||||
)
|
||||
|
||||
# session_settings should be accessible through EncryptedSession
|
||||
assert session.session_settings is not None
|
||||
assert session.session_settings.limit == 5
|
||||
|
||||
underlying.close()
|
||||
|
||||
|
||||
async def test_session_settings_get_items_uses_underlying_limit(encryption_key: str):
|
||||
"""Test that get_items uses underlying session's session_settings.limit."""
|
||||
from agents.memory import SessionSettings
|
||||
|
||||
temp_dir = tempfile.mkdtemp()
|
||||
db_path = Path(temp_dir) / "test_settings_limit.db"
|
||||
underlying = SQLiteSession("test_session", db_path, session_settings=SessionSettings(limit=3))
|
||||
|
||||
session = EncryptedSession(
|
||||
session_id="test_session",
|
||||
underlying_session=underlying,
|
||||
encryption_key=encryption_key,
|
||||
)
|
||||
|
||||
# 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 underlying 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"
|
||||
|
||||
underlying.close()
|
||||
|
||||
|
||||
async def test_session_settings_explicit_limit_overrides_settings(encryption_key: str):
|
||||
"""Test that explicit limit parameter overrides session_settings."""
|
||||
from agents.memory import SessionSettings
|
||||
|
||||
temp_dir = tempfile.mkdtemp()
|
||||
db_path = Path(temp_dir) / "test_override.db"
|
||||
underlying = SQLiteSession("test_session", db_path, session_settings=SessionSettings(limit=5))
|
||||
|
||||
session = EncryptedSession(
|
||||
session_id="test_session",
|
||||
underlying_session=underlying,
|
||||
encryption_key=encryption_key,
|
||||
)
|
||||
|
||||
# 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"
|
||||
|
||||
underlying.close()
|
||||
|
||||
|
||||
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
|
||||
|
||||
|
||||
async def test_runner_with_session_settings_override(encryption_key: str):
|
||||
"""Test that RunConfig can override session's default settings."""
|
||||
from agents import Agent, RunConfig, Runner
|
||||
from agents.memory import SessionSettings
|
||||
from tests.fake_model import FakeModel
|
||||
from tests.test_responses import get_text_message
|
||||
|
||||
temp_dir = tempfile.mkdtemp()
|
||||
db_path = Path(temp_dir) / "test_runner_override.db"
|
||||
underlying = SQLiteSession("test_session", db_path, session_settings=SessionSettings(limit=100))
|
||||
|
||||
session = EncryptedSession(
|
||||
session_id="test_session",
|
||||
underlying_session=underlying,
|
||||
encryption_key=encryption_key,
|
||||
)
|
||||
|
||||
# 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
|
||||
|
||||
underlying.close()
|
||||
@@ -0,0 +1,154 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib.abc
|
||||
import sys
|
||||
from types import ModuleType
|
||||
|
||||
import pytest
|
||||
|
||||
_PACKAGE_EXPORTS: tuple[tuple[str, str, str, str, str], ...] = (
|
||||
(
|
||||
"EncryptedSession",
|
||||
"agents.extensions.memory.encrypt_session",
|
||||
"agents.extensions.memory.encrypt_session",
|
||||
"cryptography",
|
||||
"encrypt",
|
||||
),
|
||||
("RedisSession", "agents.extensions.memory.redis_session", "redis.asyncio", "redis", "redis"),
|
||||
(
|
||||
"SQLAlchemySession",
|
||||
"agents.extensions.memory.sqlalchemy_session",
|
||||
"agents.extensions.memory.sqlalchemy_session",
|
||||
"sqlalchemy",
|
||||
"sqlalchemy",
|
||||
),
|
||||
("DaprSession", "agents.extensions.memory.dapr_session", "dapr.aio.clients", "dapr", "dapr"),
|
||||
(
|
||||
"DAPR_CONSISTENCY_EVENTUAL",
|
||||
"agents.extensions.memory.dapr_session",
|
||||
"dapr.aio.clients",
|
||||
"dapr",
|
||||
"dapr",
|
||||
),
|
||||
(
|
||||
"DAPR_CONSISTENCY_STRONG",
|
||||
"agents.extensions.memory.dapr_session",
|
||||
"dapr.aio.clients",
|
||||
"dapr",
|
||||
"dapr",
|
||||
),
|
||||
(
|
||||
"MongoDBSession",
|
||||
"agents.extensions.memory.mongodb_session",
|
||||
"pymongo.asynchronous.collection",
|
||||
"mongodb",
|
||||
"mongodb",
|
||||
),
|
||||
)
|
||||
|
||||
_DIRECT_MODULE_IMPORTS: tuple[tuple[str, str, str, str], ...] = (
|
||||
("agents.extensions.memory.redis_session", "redis.asyncio", "redis", "redis"),
|
||||
("agents.extensions.memory.dapr_session", "dapr.aio.clients", "dapr", "dapr"),
|
||||
(
|
||||
"agents.extensions.memory.mongodb_session",
|
||||
"pymongo.asynchronous.collection",
|
||||
"mongodb",
|
||||
"mongodb",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class _BrokenImportFinder(importlib.abc.MetaPathFinder):
|
||||
def __init__(self, broken_module: str, error_cls: type[ImportError]) -> None:
|
||||
self._broken_module = broken_module
|
||||
self._error_cls = error_cls
|
||||
|
||||
def find_spec(
|
||||
self,
|
||||
fullname: str,
|
||||
path: object | None,
|
||||
target: ModuleType | None = None,
|
||||
) -> None:
|
||||
if fullname == self._broken_module:
|
||||
raise self._error_cls("simulated dependency import failure")
|
||||
return None
|
||||
|
||||
|
||||
def _reset_package_imports(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
memory_module: ModuleType,
|
||||
symbol: str,
|
||||
module_name: str,
|
||||
broken_module: str,
|
||||
) -> None:
|
||||
monkeypatch.delitem(memory_module.__dict__, symbol, raising=False)
|
||||
_reset_loaded_module(monkeypatch, module_name)
|
||||
_reset_loaded_module(monkeypatch, broken_module)
|
||||
|
||||
|
||||
def _reset_loaded_module(monkeypatch: pytest.MonkeyPatch, module_name: str) -> None:
|
||||
monkeypatch.delitem(sys.modules, module_name, raising=False)
|
||||
parent_name, short_name = module_name.rsplit(".", 1)
|
||||
parent_module = sys.modules.get(parent_name)
|
||||
if parent_module is not None:
|
||||
monkeypatch.delitem(parent_module.__dict__, short_name, raising=False)
|
||||
|
||||
|
||||
def _reset_module_imports(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
module_name: str,
|
||||
broken_module: str,
|
||||
) -> None:
|
||||
_reset_loaded_module(monkeypatch, module_name)
|
||||
_reset_loaded_module(monkeypatch, broken_module)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("symbol", "module_name", "broken_module", "dependency_name", "extra_name"),
|
||||
_PACKAGE_EXPORTS,
|
||||
)
|
||||
def test_memory_package_imports_point_to_optional_extra(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
symbol: str,
|
||||
module_name: str,
|
||||
broken_module: str,
|
||||
dependency_name: str,
|
||||
extra_name: str,
|
||||
) -> None:
|
||||
import agents.extensions.memory as memory_module
|
||||
|
||||
_reset_package_imports(monkeypatch, memory_module, symbol, module_name, broken_module)
|
||||
finder = _BrokenImportFinder(broken_module, ModuleNotFoundError)
|
||||
monkeypatch.setattr(sys, "meta_path", [finder, *sys.meta_path])
|
||||
|
||||
with pytest.raises(ImportError) as exc_info:
|
||||
getattr(memory_module, symbol)
|
||||
|
||||
assert f"requires the '{dependency_name}' extra" in str(exc_info.value)
|
||||
assert f"openai-agents[{extra_name}]" in str(exc_info.value)
|
||||
assert isinstance(exc_info.value.__cause__, ImportError)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("module_name", "broken_module", "dependency_name", "extra_name"),
|
||||
_DIRECT_MODULE_IMPORTS,
|
||||
)
|
||||
@pytest.mark.parametrize("error_cls", [ImportError, ModuleNotFoundError])
|
||||
def test_memory_direct_module_imports_point_to_optional_extra(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
module_name: str,
|
||||
broken_module: str,
|
||||
dependency_name: str,
|
||||
extra_name: str,
|
||||
error_cls: type[ImportError],
|
||||
) -> None:
|
||||
_reset_module_imports(monkeypatch, module_name, broken_module)
|
||||
finder = _BrokenImportFinder(broken_module, error_cls)
|
||||
monkeypatch.setattr(sys, "meta_path", [finder, *sys.meta_path])
|
||||
|
||||
with pytest.raises(ImportError) as exc_info:
|
||||
__import__(module_name)
|
||||
|
||||
assert f"requires the '{dependency_name}' extra" in str(exc_info.value)
|
||||
assert f"openai-agents[{extra_name}]" in str(exc_info.value)
|
||||
assert isinstance(exc_info.value.__cause__, ImportError)
|
||||
@@ -0,0 +1,831 @@
|
||||
"""Tests for MongoDBSession using in-process mock objects.
|
||||
|
||||
All tests run without a real MongoDB server — or even the ``pymongo``
|
||||
package — by injecting lightweight fake classes into ``sys.modules``
|
||||
before the module under test is imported. This keeps the suite fast and
|
||||
dependency-free while exercising the full session logic.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
import types
|
||||
from collections import defaultdict
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from agents import Agent, Runner, TResponseInputItem
|
||||
from agents.memory.session_settings import SessionSettings
|
||||
from tests.fake_model import FakeModel
|
||||
from tests.test_responses import get_text_message
|
||||
|
||||
pytestmark = pytest.mark.asyncio
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# In-memory fake pymongo async types
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class FakeObjectId:
|
||||
"""Minimal ObjectId stand-in with a monotonic counter for sort order."""
|
||||
|
||||
_counter = 0
|
||||
|
||||
def __init__(self) -> None:
|
||||
FakeObjectId._counter += 1
|
||||
self._value = FakeObjectId._counter
|
||||
|
||||
def __lt__(self, other: FakeObjectId) -> bool:
|
||||
return self._value < other._value
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"FakeObjectId({self._value})"
|
||||
|
||||
|
||||
class FakeCursor:
|
||||
"""Minimal async cursor returned by ``find()``."""
|
||||
|
||||
def __init__(self, docs: list[dict[str, Any]]) -> None:
|
||||
self._docs = docs
|
||||
|
||||
def sort(
|
||||
self,
|
||||
key: str | list[tuple[str, int]],
|
||||
direction: int | None = None,
|
||||
) -> FakeCursor:
|
||||
if isinstance(key, list):
|
||||
pairs = key
|
||||
else:
|
||||
direction = direction if direction is not None else 1
|
||||
pairs = [(key, direction)]
|
||||
|
||||
docs = list(self._docs)
|
||||
for field, dir_ in reversed(pairs):
|
||||
docs.sort(key=lambda d: d.get(field, 0), reverse=(dir_ == -1))
|
||||
self._docs = docs
|
||||
return self
|
||||
|
||||
def limit(self, n: int) -> FakeCursor:
|
||||
self._docs = self._docs[:n]
|
||||
return self
|
||||
|
||||
async def to_list(self) -> list[dict[str, Any]]:
|
||||
return list(self._docs)
|
||||
|
||||
|
||||
class FakeAsyncCollection:
|
||||
"""In-memory substitute for pymongo AsyncCollection."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._docs: dict[Any, dict[str, Any]] = {}
|
||||
|
||||
async def create_index(self, keys: Any, **kwargs: Any) -> str:
|
||||
return "fake_index"
|
||||
|
||||
def find(self, query: dict[str, Any] | None = None) -> FakeCursor:
|
||||
query = query or {}
|
||||
results = [doc for doc in self._docs.values() if self._matches(doc, query)]
|
||||
return FakeCursor(results)
|
||||
|
||||
async def find_one_and_delete(
|
||||
self,
|
||||
query: dict[str, Any],
|
||||
sort: list[tuple[str, int]] | None = None,
|
||||
) -> dict[str, Any] | None:
|
||||
matches = [doc for doc in self._docs.values() if self._matches(doc, query)]
|
||||
if not matches:
|
||||
return None
|
||||
if sort:
|
||||
field, dir_ = sort[0]
|
||||
matches.sort(key=lambda d: d.get(field, 0), reverse=(dir_ == -1))
|
||||
doc = matches[0]
|
||||
self._docs.pop(id(doc["_id"]))
|
||||
return doc
|
||||
|
||||
async def insert_many(
|
||||
self,
|
||||
documents: list[dict[str, Any]],
|
||||
ordered: bool = True,
|
||||
) -> Any:
|
||||
for doc in documents:
|
||||
if "_id" not in doc:
|
||||
doc["_id"] = FakeObjectId()
|
||||
self._docs[id(doc["_id"])] = dict(doc)
|
||||
|
||||
async def find_one_and_update(
|
||||
self,
|
||||
query: dict[str, Any],
|
||||
update: dict[str, Any],
|
||||
upsert: bool = False,
|
||||
return_document: bool = False,
|
||||
) -> dict[str, Any] | None:
|
||||
for doc in self._docs.values():
|
||||
if self._matches(doc, query):
|
||||
# Apply $inc fields.
|
||||
for field, delta in update.get("$inc", {}).items():
|
||||
doc[field] = doc.get(field, 0) + delta
|
||||
for field, value in update.get("$set", {}).items():
|
||||
doc[field] = value
|
||||
return dict(doc) if return_document else None
|
||||
if upsert:
|
||||
new_doc: dict[str, Any] = {"_id": FakeObjectId()}
|
||||
new_doc.update(update.get("$setOnInsert", {}))
|
||||
new_doc.update(update.get("$set", {}))
|
||||
for field, delta in update.get("$inc", {}).items():
|
||||
new_doc[field] = new_doc.get(field, 0) + delta
|
||||
self._docs[id(new_doc["_id"])] = new_doc
|
||||
return dict(new_doc) if return_document else None
|
||||
return None
|
||||
|
||||
async def update_one(
|
||||
self,
|
||||
query: dict[str, Any],
|
||||
update: dict[str, Any],
|
||||
upsert: bool = False,
|
||||
) -> None:
|
||||
for doc in self._docs.values():
|
||||
if self._matches(doc, query):
|
||||
return # Exists — $setOnInsert is a no-op on existing docs.
|
||||
if upsert:
|
||||
new_doc2: dict[str, Any] = {"_id": FakeObjectId()}
|
||||
new_doc2.update(update.get("$setOnInsert", {}))
|
||||
self._docs[id(new_doc2["_id"])] = new_doc2
|
||||
|
||||
async def delete_many(self, query: dict[str, Any]) -> None:
|
||||
to_remove = [k for k, d in self._docs.items() if self._matches(d, query)]
|
||||
for key in to_remove:
|
||||
del self._docs[key]
|
||||
|
||||
async def delete_one(self, query: dict[str, Any]) -> None:
|
||||
for key, doc in list(self._docs.items()):
|
||||
if self._matches(doc, query):
|
||||
del self._docs[key]
|
||||
return
|
||||
|
||||
@staticmethod
|
||||
def _matches(doc: dict[str, Any], query: dict[str, Any]) -> bool:
|
||||
return all(doc.get(k) == v for k, v in query.items())
|
||||
|
||||
|
||||
class FakeAsyncDatabase:
|
||||
"""In-memory substitute for a pymongo async Database."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._collections: dict[str, FakeAsyncCollection] = defaultdict(FakeAsyncCollection)
|
||||
|
||||
def __getitem__(self, name: str) -> FakeAsyncCollection:
|
||||
return self._collections[name]
|
||||
|
||||
|
||||
class FakeAdminDatabase:
|
||||
"""Minimal admin database used by ping()."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._closed = False
|
||||
|
||||
async def command(self, cmd: str) -> dict[str, Any]:
|
||||
if self._closed:
|
||||
raise ConnectionError("Client is closed.")
|
||||
return {"ok": 1}
|
||||
|
||||
|
||||
class FakeDriverInfo:
|
||||
"""Minimal stand-in for pymongo.driver_info.DriverInfo."""
|
||||
|
||||
def __init__(self, name: str, version: str | None = None) -> None:
|
||||
self.name = name
|
||||
self.version = version
|
||||
|
||||
|
||||
class FakeAsyncMongoClient:
|
||||
"""In-memory substitute for pymongo AsyncMongoClient."""
|
||||
|
||||
def __init__(self, *args: Any, **kwargs: Any) -> None:
|
||||
self._databases: dict[str, FakeAsyncDatabase] = defaultdict(FakeAsyncDatabase)
|
||||
self._closed = False
|
||||
self.admin = FakeAdminDatabase()
|
||||
self._metadata_calls: list[FakeDriverInfo] = []
|
||||
|
||||
def __getitem__(self, name: str) -> FakeAsyncDatabase:
|
||||
return self._databases[name]
|
||||
|
||||
def append_metadata(self, driver_info: FakeDriverInfo) -> None:
|
||||
"""Record append_metadata calls for test assertions."""
|
||||
self._metadata_calls.append(driver_info)
|
||||
|
||||
async def close(self) -> None:
|
||||
"""Async close — matches PyMongo's AsyncMongoClient.close() signature."""
|
||||
self._closed = True
|
||||
self.admin._closed = True
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Inject fake pymongo into sys.modules before importing the module under test
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_fake_pymongo_modules() -> None:
|
||||
"""Populate sys.modules with stub pymongo async modules."""
|
||||
pymongo_mod = sys.modules.get("pymongo") or types.ModuleType("pymongo")
|
||||
|
||||
async_pkg = types.ModuleType("pymongo.asynchronous")
|
||||
collection_mod = types.ModuleType("pymongo.asynchronous.collection")
|
||||
client_mod = types.ModuleType("pymongo.asynchronous.mongo_client")
|
||||
driver_info_mod = types.ModuleType("pymongo.driver_info")
|
||||
|
||||
collection_mod.AsyncCollection = FakeAsyncCollection # type: ignore[attr-defined]
|
||||
client_mod.AsyncMongoClient = FakeAsyncMongoClient # type: ignore[attr-defined]
|
||||
driver_info_mod.DriverInfo = FakeDriverInfo # type: ignore[attr-defined]
|
||||
|
||||
sys.modules["pymongo"] = pymongo_mod
|
||||
sys.modules["pymongo.asynchronous"] = async_pkg
|
||||
sys.modules["pymongo.asynchronous.collection"] = collection_mod
|
||||
sys.modules["pymongo.asynchronous.mongo_client"] = client_mod
|
||||
sys.modules["pymongo.driver_info"] = driver_info_mod
|
||||
|
||||
|
||||
_make_fake_pymongo_modules()
|
||||
|
||||
# Now it's safe to import the module under test.
|
||||
from agents.extensions.memory.mongodb_session import MongoDBSession # noqa: E402
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers / fixtures
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_session(session_id: str = "test-session", **kwargs: Any) -> MongoDBSession:
|
||||
"""Create a MongoDBSession backed by a FakeAsyncMongoClient."""
|
||||
client = FakeAsyncMongoClient()
|
||||
MongoDBSession._init_state.clear()
|
||||
return MongoDBSession(
|
||||
session_id,
|
||||
client=client, # type: ignore[arg-type]
|
||||
database="agents_test",
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def session() -> MongoDBSession:
|
||||
return _make_session()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def agent() -> Agent:
|
||||
return Agent(name="test", model=FakeModel())
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Core CRUD tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def test_add_and_get_items(session: MongoDBSession) -> None:
|
||||
"""Items added to the session are retrievable in insertion order."""
|
||||
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("content") == "Hello"
|
||||
assert retrieved[1].get("content") == "Hi there!"
|
||||
|
||||
|
||||
async def test_add_empty_list_is_noop(session: MongoDBSession) -> None:
|
||||
"""Adding an empty list must not create any documents."""
|
||||
await session.add_items([])
|
||||
assert await session.get_items() == []
|
||||
|
||||
|
||||
async def test_get_items_empty_session(session: MongoDBSession) -> None:
|
||||
"""Retrieving items from a brand-new session returns an empty list."""
|
||||
assert await session.get_items() == []
|
||||
|
||||
|
||||
async def test_pop_item_returns_last(session: MongoDBSession) -> None:
|
||||
"""pop_item must return and remove the most recently added item."""
|
||||
items: list[TResponseInputItem] = [
|
||||
{"role": "user", "content": "first"},
|
||||
{"role": "assistant", "content": "second"},
|
||||
]
|
||||
await session.add_items(items)
|
||||
|
||||
popped = await session.pop_item()
|
||||
assert popped is not None
|
||||
assert popped.get("content") == "second"
|
||||
|
||||
remaining = await session.get_items()
|
||||
assert len(remaining) == 1
|
||||
assert remaining[0].get("content") == "first"
|
||||
|
||||
|
||||
async def test_pop_item_empty_session(session: MongoDBSession) -> None:
|
||||
"""pop_item on an empty session must return None."""
|
||||
assert await session.pop_item() is None
|
||||
|
||||
|
||||
async def test_clear_session(session: MongoDBSession) -> None:
|
||||
"""clear_session must remove all items and session metadata."""
|
||||
await session.add_items([{"role": "user", "content": "x"}])
|
||||
await session.clear_session()
|
||||
assert await session.get_items() == []
|
||||
|
||||
|
||||
async def test_multiple_add_calls_accumulate(session: MongoDBSession) -> None:
|
||||
"""Items from separate add_items calls all appear in get_items."""
|
||||
await session.add_items([{"role": "user", "content": "a"}])
|
||||
await session.add_items([{"role": "assistant", "content": "b"}])
|
||||
await session.add_items([{"role": "user", "content": "c"}])
|
||||
|
||||
items = await session.get_items()
|
||||
assert [i.get("content") for i in items] == ["a", "b", "c"]
|
||||
|
||||
|
||||
async def test_session_metadata_timestamps_are_written(session: MongoDBSession) -> None:
|
||||
"""Session metadata records creation time and last update time."""
|
||||
created_at = datetime(2026, 1, 1, tzinfo=timezone.utc)
|
||||
updated_at = datetime(2026, 1, 2, tzinfo=timezone.utc)
|
||||
|
||||
with patch("agents.extensions.memory.mongodb_session.datetime") as mocked_datetime:
|
||||
mocked_datetime.now.side_effect = [created_at, updated_at]
|
||||
|
||||
await session.add_items([{"role": "user", "content": "first"}])
|
||||
session_doc: dict[str, Any] = next(iter(session._sessions._docs.values()))
|
||||
assert session_doc["session_id"] == session.session_id
|
||||
assert session_doc["created_at"] == created_at
|
||||
assert session_doc["updated_at"] == created_at
|
||||
|
||||
await session.add_items([{"role": "assistant", "content": "second"}])
|
||||
assert session_doc["created_at"] == created_at
|
||||
assert session_doc["updated_at"] == updated_at
|
||||
assert session_doc["_seq"] == 2
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Limit / SessionSettings tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def test_get_items_with_explicit_limit(session: MongoDBSession) -> None:
|
||||
"""Explicit limit returns the N most recent items in chronological order."""
|
||||
await session.add_items([{"role": "user", "content": str(i)} for i in range(6)])
|
||||
|
||||
result = await session.get_items(limit=3)
|
||||
assert len(result) == 3
|
||||
assert [r.get("content") for r in result] == ["3", "4", "5"]
|
||||
|
||||
|
||||
async def test_get_items_limit_zero(session: MongoDBSession) -> None:
|
||||
"""A limit of 0 must return an empty list immediately."""
|
||||
await session.add_items([{"role": "user", "content": "x"}])
|
||||
assert await session.get_items(limit=0) == []
|
||||
|
||||
|
||||
async def test_get_items_limit_exceeds_count(session: MongoDBSession) -> None:
|
||||
"""Requesting more items than exist returns all items without error."""
|
||||
await session.add_items([{"role": "user", "content": "only"}])
|
||||
result = await session.get_items(limit=100)
|
||||
assert len(result) == 1
|
||||
|
||||
|
||||
async def test_session_settings_limit_used_as_default() -> None:
|
||||
"""session_settings.limit is applied when no explicit limit is given."""
|
||||
MongoDBSession._init_state.clear()
|
||||
s = MongoDBSession(
|
||||
"ls-test",
|
||||
client=FakeAsyncMongoClient(), # type: ignore[arg-type]
|
||||
database="agents_test",
|
||||
session_settings=SessionSettings(limit=2),
|
||||
)
|
||||
await s.add_items([{"role": "user", "content": str(i)} for i in range(5)])
|
||||
|
||||
result = await s.get_items()
|
||||
assert len(result) == 2
|
||||
assert result[0].get("content") == "3"
|
||||
assert result[1].get("content") == "4"
|
||||
|
||||
|
||||
async def test_explicit_limit_overrides_session_settings() -> None:
|
||||
"""An explicit limit passed to get_items must override session_settings.limit."""
|
||||
MongoDBSession._init_state.clear()
|
||||
s = MongoDBSession(
|
||||
"override-test",
|
||||
client=FakeAsyncMongoClient(), # type: ignore[arg-type]
|
||||
database="agents_test",
|
||||
session_settings=SessionSettings(limit=10),
|
||||
)
|
||||
await s.add_items([{"role": "user", "content": str(i)} for i in range(8)])
|
||||
|
||||
result = await s.get_items(limit=2)
|
||||
assert len(result) == 2
|
||||
assert result[0].get("content") == "6"
|
||||
assert result[1].get("content") == "7"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Session isolation
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def test_sessions_are_isolated() -> None:
|
||||
"""Two sessions with different IDs must not share data."""
|
||||
MongoDBSession._init_state.clear()
|
||||
client = FakeAsyncMongoClient()
|
||||
s1 = MongoDBSession("alice", client=client, database="agents_test") # type: ignore[arg-type]
|
||||
s2 = MongoDBSession("bob", client=client, database="agents_test") # type: ignore[arg-type]
|
||||
|
||||
await s1.add_items([{"role": "user", "content": "alice msg"}])
|
||||
await s2.add_items([{"role": "user", "content": "bob msg"}])
|
||||
|
||||
assert [i.get("content") for i in await s1.get_items()] == ["alice msg"]
|
||||
assert [i.get("content") for i in await s2.get_items()] == ["bob msg"]
|
||||
|
||||
|
||||
async def test_clear_does_not_affect_other_sessions() -> None:
|
||||
"""Clearing one session must leave sibling sessions untouched."""
|
||||
MongoDBSession._init_state.clear()
|
||||
client = FakeAsyncMongoClient()
|
||||
s1 = MongoDBSession("s1", client=client, database="agents_test") # type: ignore[arg-type]
|
||||
s2 = MongoDBSession("s2", client=client, database="agents_test") # type: ignore[arg-type]
|
||||
|
||||
await s1.add_items([{"role": "user", "content": "keep"}])
|
||||
await s2.add_items([{"role": "user", "content": "delete"}])
|
||||
|
||||
await s2.clear_session()
|
||||
|
||||
assert len(await s1.get_items()) == 1
|
||||
assert await s2.get_items() == []
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Serialisation / unicode safety
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def test_unicode_content_roundtrip(session: MongoDBSession) -> None:
|
||||
"""Unicode and emoji content must survive the serialisation round-trip."""
|
||||
items: list[TResponseInputItem] = [
|
||||
{"role": "user", "content": "こんにちは"},
|
||||
{"role": "assistant", "content": "😊👍"},
|
||||
{"role": "user", "content": "Привет"},
|
||||
]
|
||||
await session.add_items(items)
|
||||
result = await session.get_items()
|
||||
assert result[0].get("content") == "こんにちは"
|
||||
assert result[1].get("content") == "😊👍"
|
||||
assert result[2].get("content") == "Привет"
|
||||
|
||||
|
||||
async def test_json_special_characters(session: MongoDBSession) -> None:
|
||||
"""Items containing JSON-special strings must be stored without corruption."""
|
||||
items: list[TResponseInputItem] = [
|
||||
{"role": "user", "content": '{"nested": "value"}'},
|
||||
{"role": "assistant", "content": 'Quote: "Hello"'},
|
||||
{"role": "user", "content": "Line1\nLine2\tTabbed"},
|
||||
]
|
||||
await session.add_items(items)
|
||||
result = await session.get_items()
|
||||
assert result[0].get("content") == '{"nested": "value"}'
|
||||
assert result[1].get("content") == 'Quote: "Hello"'
|
||||
assert result[2].get("content") == "Line1\nLine2\tTabbed"
|
||||
|
||||
|
||||
async def test_corrupted_document_is_skipped(session: MongoDBSession) -> None:
|
||||
"""Documents with invalid JSON in message_data are silently skipped."""
|
||||
await session.add_items([{"role": "user", "content": "valid"}])
|
||||
|
||||
# Inject a corrupted document directly into the fake collection.
|
||||
bad_doc = {
|
||||
"_id": FakeObjectId(),
|
||||
"session_id": session.session_id,
|
||||
"message_data": "not valid json {{{",
|
||||
}
|
||||
session._messages._docs[id(bad_doc["_id"])] = bad_doc
|
||||
|
||||
items = await session.get_items()
|
||||
assert len(items) == 1
|
||||
assert items[0].get("content") == "valid"
|
||||
|
||||
|
||||
async def test_missing_message_data_field_is_skipped(session: MongoDBSession) -> None:
|
||||
"""Documents without a message_data field are silently skipped."""
|
||||
await session.add_items([{"role": "user", "content": "valid"}])
|
||||
|
||||
bad_doc = {"_id": FakeObjectId(), "session_id": session.session_id}
|
||||
session._messages._docs[id(bad_doc["_id"])] = bad_doc
|
||||
|
||||
items = await session.get_items()
|
||||
assert len(items) == 1
|
||||
|
||||
|
||||
async def test_non_string_message_data_is_skipped(session: MongoDBSession) -> None:
|
||||
"""Documents whose message_data is a non-string BSON type are silently skipped."""
|
||||
await session.add_items([{"role": "user", "content": "valid"}])
|
||||
|
||||
# Inject a document where message_data is an integer — json.loads raises TypeError.
|
||||
bad_doc = {"_id": FakeObjectId(), "session_id": session.session_id, "message_data": 42}
|
||||
session._messages._docs[id(bad_doc["_id"])] = bad_doc
|
||||
|
||||
items = await session.get_items()
|
||||
assert len(items) == 1
|
||||
assert items[0].get("content") == "valid"
|
||||
|
||||
|
||||
async def test_pop_item_skips_corrupt_most_recent(session: MongoDBSession) -> None:
|
||||
"""pop_item must skip a corrupt most-recent document and return the next valid one."""
|
||||
await session.add_items([{"role": "user", "content": "valid"}])
|
||||
|
||||
# Inject a corrupt document with a higher seq so it sorts as "most recent".
|
||||
bad_doc = {
|
||||
"_id": FakeObjectId(),
|
||||
"session_id": session.session_id,
|
||||
"seq": 999,
|
||||
"message_data": "not valid json {{{",
|
||||
}
|
||||
session._messages._docs[id(bad_doc["_id"])] = bad_doc
|
||||
|
||||
popped = await session.pop_item()
|
||||
assert popped is not None
|
||||
assert popped.get("content") == "valid"
|
||||
|
||||
# Both the corrupt doc and the valid one are now gone.
|
||||
assert await session.get_items() == []
|
||||
|
||||
|
||||
async def test_pop_item_returns_none_when_only_corrupt_docs_remain(
|
||||
session: MongoDBSession,
|
||||
) -> None:
|
||||
"""pop_item must drop every corrupt doc and return None when nothing valid remains."""
|
||||
bad1 = {
|
||||
"_id": FakeObjectId(),
|
||||
"session_id": session.session_id,
|
||||
"seq": 1,
|
||||
"message_data": "garbage",
|
||||
}
|
||||
bad2 = {
|
||||
"_id": FakeObjectId(),
|
||||
"session_id": session.session_id,
|
||||
"seq": 2,
|
||||
"message_data": 42, # non-string — TypeError
|
||||
}
|
||||
session._messages._docs[id(bad1["_id"])] = bad1
|
||||
session._messages._docs[id(bad2["_id"])] = bad2
|
||||
|
||||
assert await session.pop_item() is None
|
||||
# Both corrupt docs must have been removed in the process.
|
||||
assert session._messages._docs == {}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Index initialisation (idempotency)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def test_index_creation_runs_only_once(session: MongoDBSession) -> None:
|
||||
"""_ensure_indexes must call create_index only on the very first call."""
|
||||
call_count = 0
|
||||
original_messages = session._messages.create_index
|
||||
original_sessions = session._sessions.create_index
|
||||
|
||||
async def counting(*args: Any, **kwargs: Any) -> str:
|
||||
nonlocal call_count
|
||||
call_count += 1
|
||||
return "fake_index"
|
||||
|
||||
session._messages.create_index = counting # type: ignore[method-assign]
|
||||
session._sessions.create_index = counting # type: ignore[method-assign]
|
||||
|
||||
await session._ensure_indexes()
|
||||
await session._ensure_indexes() # Second call must be a no-op.
|
||||
|
||||
# Exactly one call per collection (sessions + messages).
|
||||
assert call_count == 2
|
||||
|
||||
session._messages.create_index = original_messages # type: ignore[method-assign]
|
||||
session._sessions.create_index = original_sessions # type: ignore[method-assign]
|
||||
|
||||
|
||||
async def test_different_clients_each_run_index_init() -> None:
|
||||
"""Each distinct AsyncMongoClient gets its own index-creation pass."""
|
||||
MongoDBSession._init_state.clear()
|
||||
|
||||
client_a = FakeAsyncMongoClient()
|
||||
client_b = FakeAsyncMongoClient()
|
||||
|
||||
call_counts: dict[str, int] = {"a": 0, "b": 0}
|
||||
|
||||
async def counting_a(*args: Any, **kwargs: Any) -> str:
|
||||
call_counts["a"] += 1
|
||||
return "fake_index"
|
||||
|
||||
async def counting_b(*args: Any, **kwargs: Any) -> str:
|
||||
call_counts["b"] += 1
|
||||
return "fake_index"
|
||||
|
||||
s_a = MongoDBSession("x", client=client_a, database="agents_test") # type: ignore[arg-type]
|
||||
s_b = MongoDBSession("x", client=client_b, database="agents_test") # type: ignore[arg-type]
|
||||
|
||||
s_a._messages.create_index = counting_a # type: ignore[method-assign]
|
||||
s_a._sessions.create_index = counting_a # type: ignore[method-assign]
|
||||
s_b._messages.create_index = counting_b # type: ignore[method-assign]
|
||||
s_b._sessions.create_index = counting_b # type: ignore[method-assign]
|
||||
|
||||
await s_a._ensure_indexes()
|
||||
await s_b._ensure_indexes()
|
||||
|
||||
# Each client must trigger its own index creation (2 calls = sessions + messages).
|
||||
assert call_counts["a"] == 2
|
||||
assert call_counts["b"] == 2
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Connectivity and lifecycle
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def test_ping_success(session: MongoDBSession) -> None:
|
||||
"""ping() must return True when the client responds normally."""
|
||||
assert await session.ping() is True
|
||||
|
||||
|
||||
async def test_ping_failure(session: MongoDBSession) -> None:
|
||||
"""ping() must return False when the server raises an exception."""
|
||||
original = session._client.admin.command
|
||||
|
||||
async def _fail(*args: Any, **kwargs: Any) -> dict[str, Any]:
|
||||
raise ConnectionError("unreachable")
|
||||
|
||||
session._client.admin.command = _fail # type: ignore[method-assign, assignment]
|
||||
assert await session.ping() is False
|
||||
session._client.admin.command = original # type: ignore[method-assign]
|
||||
|
||||
|
||||
async def test_close_external_client_not_closed() -> None:
|
||||
"""close() must NOT close a client that was injected externally."""
|
||||
MongoDBSession._init_state.clear()
|
||||
client = FakeAsyncMongoClient()
|
||||
s = MongoDBSession("x", client=client, database="agents_test") # type: ignore[arg-type]
|
||||
assert s._owns_client is False
|
||||
|
||||
await s.close()
|
||||
assert not client._closed
|
||||
|
||||
|
||||
async def test_close_owned_client_is_closed() -> None:
|
||||
"""close() must close a client created by from_uri."""
|
||||
MongoDBSession._init_state.clear()
|
||||
fake_client = FakeAsyncMongoClient()
|
||||
with patch(
|
||||
"agents.extensions.memory.mongodb_session.AsyncMongoClient",
|
||||
return_value=fake_client,
|
||||
):
|
||||
s = MongoDBSession.from_uri("owned", uri="mongodb://localhost:27017", database="t")
|
||||
assert s._owns_client is True
|
||||
|
||||
await s.close()
|
||||
assert fake_client._closed
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Runner integration
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def test_runner_integration(agent: Agent) -> None:
|
||||
"""MongoDBSession must supply conversation history to the Runner."""
|
||||
session = _make_session("runner-test")
|
||||
|
||||
assert isinstance(agent.model, FakeModel)
|
||||
agent.model.set_next_output([get_text_message("San Francisco")])
|
||||
result1 = await Runner.run(agent, "Where is the Golden Gate Bridge?", session=session)
|
||||
assert result1.final_output == "San Francisco"
|
||||
|
||||
agent.model.set_next_output([get_text_message("California")])
|
||||
result2 = await Runner.run(agent, "What state is it in?", session=session)
|
||||
assert result2.final_output == "California"
|
||||
|
||||
last_input = agent.model.last_turn_args["input"]
|
||||
assert len(last_input) > 1
|
||||
assert any("Golden Gate Bridge" in str(item.get("content", "")) for item in last_input)
|
||||
|
||||
|
||||
async def test_runner_session_isolation(agent: Agent) -> None:
|
||||
"""Two independent sessions must not bleed history into each other."""
|
||||
MongoDBSession._init_state.clear()
|
||||
client = FakeAsyncMongoClient()
|
||||
s1 = MongoDBSession("user-a", client=client, database="agents_test") # type: ignore[arg-type]
|
||||
s2 = MongoDBSession("user-b", client=client, database="agents_test") # type: ignore[arg-type]
|
||||
|
||||
assert isinstance(agent.model, FakeModel)
|
||||
agent.model.set_next_output([get_text_message("I like cats.")])
|
||||
await Runner.run(agent, "I like cats.", session=s1)
|
||||
|
||||
agent.model.set_next_output([get_text_message("I like dogs.")])
|
||||
await Runner.run(agent, "I like dogs.", session=s2)
|
||||
|
||||
agent.model.set_next_output([get_text_message("You said you like cats.")])
|
||||
result = await Runner.run(agent, "What animal did I mention?", session=s1)
|
||||
assert "cats" in result.final_output.lower()
|
||||
assert "dogs" not in result.final_output.lower()
|
||||
|
||||
|
||||
async def test_runner_with_session_settings_limit(agent: Agent) -> None:
|
||||
"""RunConfig.session_settings.limit must cap the history sent to the model."""
|
||||
from agents import RunConfig
|
||||
|
||||
MongoDBSession._init_state.clear()
|
||||
session = MongoDBSession(
|
||||
"limit-test",
|
||||
client=FakeAsyncMongoClient(), # type: ignore[arg-type]
|
||||
database="agents_test",
|
||||
session_settings=SessionSettings(limit=100),
|
||||
)
|
||||
|
||||
history: list[TResponseInputItem] = [
|
||||
{"role": "user", "content": f"Turn {i}"} for i in range(10)
|
||||
]
|
||||
await session.add_items(history)
|
||||
|
||||
assert isinstance(agent.model, FakeModel)
|
||||
agent.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)),
|
||||
)
|
||||
|
||||
last_input = agent.model.last_turn_args["input"]
|
||||
history_items = [i for i in last_input if i.get("content") != "New question"]
|
||||
assert len(history_items) == 2
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Client metadata (driver handshake)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def test_injected_client_receives_append_metadata() -> None:
|
||||
"""Append_metadata is called on a caller-supplied client."""
|
||||
MongoDBSession._init_state.clear()
|
||||
client = FakeAsyncMongoClient()
|
||||
|
||||
MongoDBSession("meta-test", client=client, database="agents_test") # type: ignore[arg-type]
|
||||
|
||||
assert len(client._metadata_calls) == 1
|
||||
info = client._metadata_calls[0]
|
||||
assert info.name == "openai-agents"
|
||||
|
||||
|
||||
async def test_from_uri_passes_driver_info_to_constructor() -> None:
|
||||
"""driver=_DRIVER_INFO is forwarded to AsyncMongoClient via from_uri."""
|
||||
MongoDBSession._init_state.clear()
|
||||
|
||||
captured_kwargs: dict[str, Any] = {}
|
||||
|
||||
def _fake_client(uri: str, **kwargs: Any) -> FakeAsyncMongoClient:
|
||||
captured_kwargs.update(kwargs)
|
||||
return FakeAsyncMongoClient()
|
||||
|
||||
with patch(
|
||||
"agents.extensions.memory.mongodb_session.AsyncMongoClient",
|
||||
side_effect=_fake_client,
|
||||
):
|
||||
MongoDBSession.from_uri("uri-test", uri="mongodb://localhost:27017", database="t")
|
||||
|
||||
assert "driver" in captured_kwargs
|
||||
assert captured_kwargs["driver"].name == "openai-agents"
|
||||
|
||||
|
||||
async def test_caller_supplied_driver_info_is_not_overwritten() -> None:
|
||||
"""A caller-supplied driver kwarg must not be silently replaced."""
|
||||
MongoDBSession._init_state.clear()
|
||||
|
||||
captured_kwargs: dict[str, Any] = {}
|
||||
custom_info = FakeDriverInfo(name="MyApp")
|
||||
|
||||
def _fake_client(uri: str, **kwargs: Any) -> FakeAsyncMongoClient:
|
||||
captured_kwargs.update(kwargs)
|
||||
return FakeAsyncMongoClient()
|
||||
|
||||
with patch(
|
||||
"agents.extensions.memory.mongodb_session.AsyncMongoClient",
|
||||
side_effect=_fake_client,
|
||||
):
|
||||
MongoDBSession.from_uri(
|
||||
"uri-test",
|
||||
uri="mongodb://localhost:27017",
|
||||
database="t",
|
||||
client_kwargs={"driver": custom_info},
|
||||
)
|
||||
|
||||
# The caller's value must be preserved — setdefault must not overwrite it.
|
||||
assert captured_kwargs["driver"] is custom_info
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,956 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import threading
|
||||
from collections.abc import Iterable, Sequence
|
||||
from contextlib import asynccontextmanager
|
||||
from datetime import datetime, timedelta
|
||||
from typing import Any, cast
|
||||
|
||||
import pytest
|
||||
from openai.types.responses.response_output_message_param import ResponseOutputMessageParam
|
||||
from openai.types.responses.response_output_text_param import ResponseOutputTextParam
|
||||
from openai.types.responses.response_reasoning_item_param import (
|
||||
ResponseReasoningItemParam,
|
||||
Summary,
|
||||
)
|
||||
from sqlalchemy import insert, select, text, update
|
||||
from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
|
||||
from sqlalchemy.sql import Select
|
||||
|
||||
pytest.importorskip("sqlalchemy") # Skip tests if SQLAlchemy is not installed
|
||||
|
||||
from agents import Agent, Runner, TResponseInputItem
|
||||
from agents.extensions.memory.sqlalchemy_session import SQLAlchemySession
|
||||
from tests.fake_model import FakeModel
|
||||
from tests.test_responses import get_text_message
|
||||
|
||||
# Mark all tests in this file as asyncio
|
||||
pytestmark = pytest.mark.asyncio
|
||||
|
||||
# Use in-memory SQLite for tests
|
||||
DB_URL = "sqlite+aiosqlite:///:memory:"
|
||||
|
||||
|
||||
def _make_message_item(item_id: str, text_value: str) -> TResponseInputItem:
|
||||
content: ResponseOutputTextParam = {
|
||||
"type": "output_text",
|
||||
"text": text_value,
|
||||
"annotations": [],
|
||||
"logprobs": [],
|
||||
}
|
||||
message: ResponseOutputMessageParam = {
|
||||
"id": item_id,
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"status": "completed",
|
||||
"content": [content],
|
||||
}
|
||||
return cast(TResponseInputItem, message)
|
||||
|
||||
|
||||
def _make_reasoning_item(item_id: str, summary_text: str) -> TResponseInputItem:
|
||||
summary: Summary = {"type": "summary_text", "text": summary_text}
|
||||
reasoning: ResponseReasoningItemParam = {
|
||||
"id": item_id,
|
||||
"type": "reasoning",
|
||||
"summary": [summary],
|
||||
}
|
||||
return cast(TResponseInputItem, reasoning)
|
||||
|
||||
|
||||
def _item_ids(items: Sequence[TResponseInputItem]) -> list[str]:
|
||||
result: list[str] = []
|
||||
for item in items:
|
||||
item_dict = cast(dict[str, Any], item)
|
||||
result.append(cast(str, item_dict["id"]))
|
||||
return result
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def agent() -> Agent:
|
||||
"""Fixture for a basic agent with a fake model."""
|
||||
return Agent(name="test", model=FakeModel())
|
||||
|
||||
|
||||
async def test_sqlalchemy_session_direct_ops(agent: Agent):
|
||||
"""Test direct database operations of SQLAlchemySession."""
|
||||
session_id = "direct_ops_test"
|
||||
session = SQLAlchemySession.from_url(session_id, url=DB_URL, create_tables=True)
|
||||
|
||||
# 1. Add items
|
||||
items: list[TResponseInputItem] = [
|
||||
{"role": "user", "content": "Hello"},
|
||||
{"role": "assistant", "content": "Hi there!"},
|
||||
]
|
||||
await session.add_items(items)
|
||||
|
||||
# 2. Get items and verify
|
||||
retrieved = await session.get_items()
|
||||
assert len(retrieved) == 2
|
||||
assert retrieved[0].get("content") == "Hello"
|
||||
assert retrieved[1].get("content") == "Hi there!"
|
||||
|
||||
# 3. Pop item
|
||||
popped = await session.pop_item()
|
||||
assert popped is not None
|
||||
assert popped.get("content") == "Hi there!"
|
||||
retrieved_after_pop = await session.get_items()
|
||||
assert len(retrieved_after_pop) == 1
|
||||
assert retrieved_after_pop[0].get("content") == "Hello"
|
||||
|
||||
# 4. Clear session
|
||||
await session.clear_session()
|
||||
retrieved_after_clear = await session.get_items()
|
||||
assert len(retrieved_after_clear) == 0
|
||||
|
||||
|
||||
async def test_sqlalchemy_session_defaults_to_escaped_non_ascii_storage():
|
||||
"""Default storage keeps the historical escaped non-ASCII JSON representation."""
|
||||
session = SQLAlchemySession.from_url("default_ascii_storage", url=DB_URL, create_tables=True)
|
||||
item: TResponseInputItem = {"role": "user", "content": "café"}
|
||||
|
||||
await session.add_items([item])
|
||||
|
||||
async with session._session_factory() as sess:
|
||||
rows = await sess.execute(
|
||||
select(session._messages.c.message_data).where(
|
||||
session._messages.c.session_id == session.session_id
|
||||
)
|
||||
)
|
||||
stored = rows.scalar_one()
|
||||
|
||||
assert "\\u00e9" in stored
|
||||
assert "café" not in stored
|
||||
assert await session.get_items() == [item]
|
||||
|
||||
|
||||
async def test_sqlalchemy_session_can_store_non_ascii_without_escaping():
|
||||
"""ensure_ascii=False stores multilingual content readably while preserving round-trip data."""
|
||||
session = SQLAlchemySession.from_url(
|
||||
"non_ascii_storage",
|
||||
url=DB_URL,
|
||||
create_tables=True,
|
||||
ensure_ascii=False,
|
||||
)
|
||||
item: TResponseInputItem = {"role": "user", "content": "café"}
|
||||
|
||||
await session.add_items([item])
|
||||
|
||||
async with session._session_factory() as sess:
|
||||
rows = await sess.execute(
|
||||
select(session._messages.c.message_data).where(
|
||||
session._messages.c.session_id == session.session_id
|
||||
)
|
||||
)
|
||||
stored = rows.scalar_one()
|
||||
|
||||
assert "café" in stored
|
||||
assert "\\u00e9" not in stored
|
||||
assert await session.get_items() == [item]
|
||||
|
||||
|
||||
async def test_runner_integration(agent: Agent):
|
||||
"""Test that SQLAlchemySession works correctly with the agent Runner."""
|
||||
session_id = "runner_integration_test"
|
||||
session = SQLAlchemySession.from_url(session_id, url=DB_URL, create_tables=True)
|
||||
|
||||
# First turn
|
||||
assert isinstance(agent.model, FakeModel)
|
||||
agent.model.set_next_output([get_text_message("San Francisco")])
|
||||
result1 = await Runner.run(
|
||||
agent,
|
||||
"What city is the Golden Gate Bridge in?",
|
||||
session=session,
|
||||
)
|
||||
assert result1.final_output == "San Francisco"
|
||||
|
||||
# Second turn
|
||||
agent.model.set_next_output([get_text_message("California")])
|
||||
result2 = await Runner.run(agent, "What state is it in?", session=session)
|
||||
assert result2.final_output == "California"
|
||||
|
||||
# Verify history was passed to the model on the second turn
|
||||
last_input = agent.model.last_turn_args["input"]
|
||||
assert len(last_input) > 1
|
||||
assert any("Golden Gate Bridge" in str(item.get("content", "")) for item in last_input)
|
||||
|
||||
|
||||
async def test_session_isolation(agent: Agent):
|
||||
"""Test that different session IDs result in isolated conversation histories."""
|
||||
session_id_1 = "session_1"
|
||||
session1 = SQLAlchemySession.from_url(session_id_1, url=DB_URL, create_tables=True)
|
||||
|
||||
session_id_2 = "session_2"
|
||||
session2 = SQLAlchemySession.from_url(session_id_2, url=DB_URL, create_tables=True)
|
||||
|
||||
# Interact with session 1
|
||||
assert isinstance(agent.model, FakeModel)
|
||||
agent.model.set_next_output([get_text_message("I like cats.")])
|
||||
await Runner.run(agent, "I like cats.", session=session1)
|
||||
|
||||
# Interact with session 2
|
||||
agent.model.set_next_output([get_text_message("I like dogs.")])
|
||||
await Runner.run(agent, "I like dogs.", session=session2)
|
||||
|
||||
# Go back to session 1 and check its memory
|
||||
agent.model.set_next_output([get_text_message("You said you like cats.")])
|
||||
result = await Runner.run(agent, "What animal did I say I like?", session=session1)
|
||||
assert "cats" in result.final_output.lower()
|
||||
assert "dogs" not in result.final_output.lower()
|
||||
|
||||
|
||||
async def test_get_items_with_limit(agent: Agent):
|
||||
"""Test the limit parameter in get_items."""
|
||||
session_id = "limit_test"
|
||||
session = SQLAlchemySession.from_url(session_id, url=DB_URL, create_tables=True)
|
||||
|
||||
items: list[TResponseInputItem] = [
|
||||
{"role": "user", "content": "1"},
|
||||
{"role": "assistant", "content": "2"},
|
||||
{"role": "user", "content": "3"},
|
||||
{"role": "assistant", "content": "4"},
|
||||
]
|
||||
await session.add_items(items)
|
||||
|
||||
# Get last 2 items
|
||||
latest_2 = await session.get_items(limit=2)
|
||||
assert len(latest_2) == 2
|
||||
assert latest_2[0].get("content") == "3"
|
||||
assert latest_2[1].get("content") == "4"
|
||||
|
||||
# Get all items
|
||||
all_items = await session.get_items()
|
||||
assert len(all_items) == 4
|
||||
|
||||
# Get more than available
|
||||
more_than_all = await session.get_items(limit=10)
|
||||
assert len(more_than_all) == 4
|
||||
|
||||
|
||||
async def test_pop_from_empty_session():
|
||||
"""Test that pop_item returns None on an empty session."""
|
||||
session = SQLAlchemySession.from_url("empty_session", url=DB_URL, create_tables=True)
|
||||
popped = await session.pop_item()
|
||||
assert popped is None
|
||||
|
||||
|
||||
async def test_pop_item_skips_corrupt_most_recent():
|
||||
"""pop_item skips corrupt newest rows and returns the next valid item."""
|
||||
session = SQLAlchemySession.from_url("pop_corrupt", url=DB_URL, create_tables=True)
|
||||
|
||||
valid_item: TResponseInputItem = {"role": "user", "content": "valid"}
|
||||
await session.add_items([valid_item])
|
||||
|
||||
await session._ensure_tables()
|
||||
async with session._session_factory() as sess:
|
||||
async with sess.begin():
|
||||
await sess.execute(
|
||||
insert(session._messages).values(
|
||||
{"session_id": session.session_id, "message_data": "not valid json {{{"}
|
||||
)
|
||||
)
|
||||
|
||||
assert await session.pop_item() == valid_item
|
||||
assert await session.get_items() == []
|
||||
|
||||
|
||||
async def test_pop_item_returns_none_after_dropping_only_corrupt_rows():
|
||||
"""pop_item removes corrupt rows and returns None when no valid items remain."""
|
||||
session = SQLAlchemySession.from_url("pop_only_corrupt", url=DB_URL, create_tables=True)
|
||||
|
||||
await session._ensure_tables()
|
||||
async with session._session_factory() as sess:
|
||||
async with sess.begin():
|
||||
await sess.execute(
|
||||
insert(session._messages).values(
|
||||
{"session_id": session.session_id, "message_data": "not valid json {{{"}
|
||||
)
|
||||
)
|
||||
|
||||
assert await session.pop_item() is None
|
||||
assert await session.get_items() == []
|
||||
|
||||
|
||||
async def test_add_empty_items_list():
|
||||
"""Test that adding an empty list of items is a no-op."""
|
||||
session_id = "add_empty_test"
|
||||
session = SQLAlchemySession.from_url(session_id, url=DB_URL, create_tables=True)
|
||||
|
||||
initial_items = await session.get_items()
|
||||
assert len(initial_items) == 0
|
||||
|
||||
await session.add_items([])
|
||||
|
||||
items_after_add = await session.get_items()
|
||||
assert len(items_after_add) == 0
|
||||
|
||||
|
||||
async def test_add_items_concurrent_first_access_with_create_tables(tmp_path):
|
||||
"""Concurrent first writes should not race table creation or drop items."""
|
||||
db_url = f"sqlite+aiosqlite:///{tmp_path / 'concurrent_first_access.db'}"
|
||||
session = SQLAlchemySession.from_url(
|
||||
"concurrent_first_access",
|
||||
url=db_url,
|
||||
create_tables=True,
|
||||
)
|
||||
submitted = [f"msg-{i}" for i in range(25)]
|
||||
|
||||
async def worker(content: str) -> None:
|
||||
await session.add_items([{"role": "user", "content": content}])
|
||||
|
||||
results = await asyncio.gather(
|
||||
*(worker(content) for content in submitted),
|
||||
return_exceptions=True,
|
||||
)
|
||||
|
||||
assert [result for result in results if isinstance(result, Exception)] == []
|
||||
|
||||
stored = await session.get_items()
|
||||
assert len(stored) == len(submitted)
|
||||
stored_contents: list[str] = []
|
||||
for item in stored:
|
||||
content = item.get("content")
|
||||
assert isinstance(content, str)
|
||||
stored_contents.append(content)
|
||||
assert sorted(stored_contents) == sorted(submitted)
|
||||
|
||||
|
||||
async def test_add_items_concurrent_first_write_after_tables_exist(tmp_path):
|
||||
"""Concurrent first writes should not race parent session creation."""
|
||||
db_url = f"sqlite+aiosqlite:///{tmp_path / 'concurrent_first_write.db'}"
|
||||
setup_session = SQLAlchemySession.from_url(
|
||||
"concurrent_first_write",
|
||||
url=db_url,
|
||||
create_tables=True,
|
||||
)
|
||||
await setup_session.get_items()
|
||||
|
||||
session = SQLAlchemySession.from_url(
|
||||
"concurrent_first_write",
|
||||
url=db_url,
|
||||
create_tables=False,
|
||||
)
|
||||
submitted = [f"msg-{i}" for i in range(25)]
|
||||
|
||||
async def worker(content: str) -> None:
|
||||
await session.add_items([{"role": "user", "content": content}])
|
||||
|
||||
results = await asyncio.gather(
|
||||
*(worker(content) for content in submitted),
|
||||
return_exceptions=True,
|
||||
)
|
||||
|
||||
assert [result for result in results if isinstance(result, Exception)] == []
|
||||
|
||||
stored = await session.get_items()
|
||||
assert len(stored) == len(submitted)
|
||||
stored_contents: list[str] = []
|
||||
for item in stored:
|
||||
content = item.get("content")
|
||||
assert isinstance(content, str)
|
||||
stored_contents.append(content)
|
||||
assert sorted(stored_contents) == sorted(submitted)
|
||||
|
||||
|
||||
async def test_add_items_waits_for_transient_sqlite_write_lock(tmp_path):
|
||||
"""SQLite writes should wait briefly for a transient lock instead of failing."""
|
||||
db_url = f"sqlite+aiosqlite:///{tmp_path / 'sqlite_write_lock_retry.db'}"
|
||||
session = SQLAlchemySession.from_url(
|
||||
"sqlite_write_lock_retry",
|
||||
url=db_url,
|
||||
create_tables=True,
|
||||
)
|
||||
await session.get_items()
|
||||
|
||||
async with session.engine.connect() as conn:
|
||||
await conn.execute(text("BEGIN IMMEDIATE"))
|
||||
blocked_write = asyncio.create_task(
|
||||
session.add_items([{"role": "user", "content": "after-lock"}])
|
||||
)
|
||||
await asyncio.sleep(0.1)
|
||||
await conn.rollback()
|
||||
|
||||
await asyncio.wait_for(blocked_write, timeout=5)
|
||||
|
||||
stored = await session.get_items()
|
||||
assert len(stored) == 1
|
||||
assert stored[0].get("content") == "after-lock"
|
||||
|
||||
|
||||
async def test_add_items_concurrent_first_access_across_sessions_with_shared_engine(tmp_path):
|
||||
"""Concurrent first writes should not race table creation across session instances."""
|
||||
db_url = f"sqlite+aiosqlite:///{tmp_path / 'concurrent_shared_engine.db'}"
|
||||
engine = create_async_engine(db_url)
|
||||
try:
|
||||
session_a = SQLAlchemySession("shared_engine_a", engine=engine, create_tables=True)
|
||||
session_b = SQLAlchemySession("shared_engine_b", engine=engine, create_tables=True)
|
||||
|
||||
results = await asyncio.gather(
|
||||
session_a.add_items([{"role": "user", "content": "one"}]),
|
||||
session_b.add_items([{"role": "user", "content": "two"}]),
|
||||
return_exceptions=True,
|
||||
)
|
||||
|
||||
assert [result for result in results if isinstance(result, Exception)] == []
|
||||
|
||||
stored_a = await session_a.get_items()
|
||||
assert len(stored_a) == 1
|
||||
assert stored_a[0].get("content") == "one"
|
||||
|
||||
stored_b = await session_b.get_items()
|
||||
assert len(stored_b) == 1
|
||||
assert stored_b[0].get("content") == "two"
|
||||
finally:
|
||||
await engine.dispose()
|
||||
|
||||
|
||||
async def test_add_items_concurrent_first_access_across_from_url_sessions(tmp_path):
|
||||
"""Concurrent first writes should not race table creation across from_url sessions."""
|
||||
db_url = f"sqlite+aiosqlite:///{tmp_path / 'concurrent_from_url.db'}"
|
||||
session_a = SQLAlchemySession.from_url("from_url_a", url=db_url, create_tables=True)
|
||||
session_b = SQLAlchemySession.from_url("from_url_b", url=db_url, create_tables=True)
|
||||
try:
|
||||
results = await asyncio.gather(
|
||||
session_a.add_items([{"role": "user", "content": "one"}]),
|
||||
session_b.add_items([{"role": "user", "content": "two"}]),
|
||||
return_exceptions=True,
|
||||
)
|
||||
|
||||
assert [result for result in results if isinstance(result, Exception)] == []
|
||||
|
||||
stored_a = await session_a.get_items()
|
||||
assert len(stored_a) == 1
|
||||
assert stored_a[0].get("content") == "one"
|
||||
|
||||
stored_b = await session_b.get_items()
|
||||
assert len(stored_b) == 1
|
||||
assert stored_b[0].get("content") == "two"
|
||||
finally:
|
||||
await session_a.engine.dispose()
|
||||
await session_b.engine.dispose()
|
||||
|
||||
|
||||
async def test_add_items_concurrent_first_access_across_from_url_sessions_cross_loop(tmp_path):
|
||||
"""Concurrent first writes should not race or hang across event loops."""
|
||||
db_url = f"sqlite+aiosqlite:///{tmp_path / 'concurrent_from_url_cross_loop.db'}"
|
||||
barrier = threading.Barrier(2)
|
||||
results: list[tuple[str, str, Any]] = []
|
||||
results_lock = threading.Lock()
|
||||
|
||||
def worker(session_id: str, content: str) -> None:
|
||||
async def run() -> tuple[str, Any]:
|
||||
session = SQLAlchemySession.from_url(session_id, url=db_url, create_tables=True)
|
||||
barrier.wait()
|
||||
try:
|
||||
await asyncio.wait_for(
|
||||
session.add_items([{"role": "user", "content": content}]),
|
||||
timeout=5,
|
||||
)
|
||||
stored = await session.get_items()
|
||||
return ("ok", stored)
|
||||
finally:
|
||||
await session.engine.dispose()
|
||||
|
||||
try:
|
||||
status, payload = asyncio.run(run())
|
||||
except Exception as exc:
|
||||
status, payload = type(exc).__name__, str(exc)
|
||||
|
||||
with results_lock:
|
||||
results.append((session_id, status, payload))
|
||||
|
||||
threads = [
|
||||
threading.Thread(target=worker, args=("from_url_cross_loop_a", "one")),
|
||||
threading.Thread(target=worker, args=("from_url_cross_loop_b", "two")),
|
||||
]
|
||||
for thread in threads:
|
||||
thread.start()
|
||||
for thread in threads:
|
||||
await asyncio.to_thread(thread.join)
|
||||
|
||||
assert len(results) == 2
|
||||
assert [status for _, status, _ in results] == ["ok", "ok"]
|
||||
|
||||
stored_by_session = {
|
||||
session_id: cast(list[TResponseInputItem], payload) for session_id, _, payload in results
|
||||
}
|
||||
assert stored_by_session["from_url_cross_loop_a"][0].get("content") == "one"
|
||||
assert stored_by_session["from_url_cross_loop_b"][0].get("content") == "two"
|
||||
|
||||
|
||||
async def test_add_items_concurrent_first_access_with_shared_session_cross_loop(tmp_path):
|
||||
"""A shared session instance should not hang when used from two event loops."""
|
||||
db_url = f"sqlite+aiosqlite:///{tmp_path / 'shared_session_cross_loop.db'}"
|
||||
session = SQLAlchemySession.from_url(
|
||||
"shared_session_cross_loop",
|
||||
url=db_url,
|
||||
create_tables=True,
|
||||
)
|
||||
barrier = threading.Barrier(2)
|
||||
results: list[tuple[str, str]] = []
|
||||
results_lock = threading.Lock()
|
||||
|
||||
def worker(content: str) -> None:
|
||||
async def run() -> None:
|
||||
barrier.wait()
|
||||
await asyncio.wait_for(
|
||||
session.add_items([{"role": "user", "content": content}]),
|
||||
timeout=5,
|
||||
)
|
||||
|
||||
try:
|
||||
asyncio.run(run())
|
||||
status = "ok"
|
||||
except Exception as exc:
|
||||
status = type(exc).__name__
|
||||
|
||||
with results_lock:
|
||||
results.append((content, status))
|
||||
|
||||
threads = [
|
||||
threading.Thread(target=worker, args=("one",)),
|
||||
threading.Thread(target=worker, args=("two",)),
|
||||
]
|
||||
try:
|
||||
for thread in threads:
|
||||
thread.start()
|
||||
for thread in threads:
|
||||
await asyncio.to_thread(thread.join)
|
||||
|
||||
assert sorted(results) == [("one", "ok"), ("two", "ok")]
|
||||
|
||||
stored = await session.get_items()
|
||||
stored_contents: list[str] = []
|
||||
for item in stored:
|
||||
content = item.get("content")
|
||||
assert isinstance(content, str)
|
||||
stored_contents.append(content)
|
||||
assert sorted(stored_contents) == ["one", "two"]
|
||||
finally:
|
||||
await session.engine.dispose()
|
||||
|
||||
|
||||
async def test_add_items_cancelled_waiter_does_not_strand_table_init_lock(tmp_path):
|
||||
"""Cancelling a waiting initializer must not leave the shared init lock acquired."""
|
||||
db_url = f"sqlite+aiosqlite:///{tmp_path / 'cancelled_table_init_waiter.db'}"
|
||||
holder = SQLAlchemySession.from_url("holder", url=db_url, create_tables=True)
|
||||
waiter = SQLAlchemySession.from_url("waiter", url=db_url, create_tables=True)
|
||||
follower = SQLAlchemySession.from_url("follower", url=db_url, create_tables=True)
|
||||
|
||||
assert holder._init_lock is waiter._init_lock
|
||||
assert waiter._init_lock is follower._init_lock
|
||||
assert holder._init_lock is not None
|
||||
|
||||
acquired = holder._init_lock.acquire(blocking=False)
|
||||
assert acquired
|
||||
|
||||
try:
|
||||
blocked = asyncio.create_task(waiter.add_items([{"role": "user", "content": "waiter"}]))
|
||||
await asyncio.sleep(0.05)
|
||||
blocked.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await blocked
|
||||
finally:
|
||||
holder._init_lock.release()
|
||||
|
||||
try:
|
||||
await asyncio.wait_for(
|
||||
follower.add_items([{"role": "user", "content": "follower"}]),
|
||||
timeout=2,
|
||||
)
|
||||
stored = await follower.get_items()
|
||||
assert len(stored) == 1
|
||||
assert stored[0].get("content") == "follower"
|
||||
finally:
|
||||
await holder.engine.dispose()
|
||||
await waiter.engine.dispose()
|
||||
await follower.engine.dispose()
|
||||
|
||||
|
||||
async def test_create_tables_false_does_not_allocate_shared_init_lock(tmp_path):
|
||||
"""Sessions that skip auto-create should not populate the shared lock map."""
|
||||
db_url = f"sqlite+aiosqlite:///{tmp_path / 'no_create_tables_lock.db'}"
|
||||
before = len(SQLAlchemySession._table_init_locks)
|
||||
session = SQLAlchemySession.from_url("no_create_tables_lock", url=db_url, create_tables=False)
|
||||
try:
|
||||
assert session._init_lock is None
|
||||
assert len(SQLAlchemySession._table_init_locks) == before
|
||||
finally:
|
||||
await session.engine.dispose()
|
||||
|
||||
|
||||
async def test_get_items_same_timestamp_consistent_order():
|
||||
"""Test that items with identical timestamps keep insertion order."""
|
||||
session_id = "same_timestamp_test"
|
||||
session = SQLAlchemySession.from_url(session_id, url=DB_URL, create_tables=True)
|
||||
|
||||
older_item = _make_message_item("older_same_ts", "old")
|
||||
reasoning_item = _make_reasoning_item("rs_same_ts", "...")
|
||||
message_item = _make_message_item("msg_same_ts", "...")
|
||||
await session.add_items([older_item])
|
||||
await session.add_items([reasoning_item, message_item])
|
||||
|
||||
async with session._session_factory() as sess:
|
||||
rows = await sess.execute(
|
||||
select(session._messages.c.id, session._messages.c.message_data).where(
|
||||
session._messages.c.session_id == session.session_id
|
||||
)
|
||||
)
|
||||
id_map = {
|
||||
json.loads(message_json)["id"]: row_id for row_id, message_json in rows.fetchall()
|
||||
}
|
||||
shared = datetime(2025, 10, 15, 17, 26, 39, 132483)
|
||||
older = shared - timedelta(milliseconds=1)
|
||||
await sess.execute(
|
||||
update(session._messages)
|
||||
.where(
|
||||
session._messages.c.id.in_(
|
||||
[
|
||||
id_map["rs_same_ts"],
|
||||
id_map["msg_same_ts"],
|
||||
]
|
||||
)
|
||||
)
|
||||
.values(created_at=shared)
|
||||
)
|
||||
await sess.execute(
|
||||
update(session._messages)
|
||||
.where(session._messages.c.id == id_map["older_same_ts"])
|
||||
.values(created_at=older)
|
||||
)
|
||||
await sess.commit()
|
||||
|
||||
real_factory = session._session_factory
|
||||
|
||||
class FakeResult:
|
||||
def __init__(self, rows: Iterable[Any]):
|
||||
self._rows = list(rows)
|
||||
|
||||
def all(self) -> list[Any]:
|
||||
return list(self._rows)
|
||||
|
||||
def needs_shuffle(statement: Any) -> bool:
|
||||
if not isinstance(statement, Select):
|
||||
return False
|
||||
orderings = list(statement._order_by_clause)
|
||||
if not orderings:
|
||||
return False
|
||||
id_asc = session._messages.c.id.asc()
|
||||
id_desc = session._messages.c.id.desc()
|
||||
|
||||
def references_id(clause) -> bool:
|
||||
try:
|
||||
return bool(clause.compare(id_asc) or clause.compare(id_desc))
|
||||
except AttributeError:
|
||||
return False
|
||||
|
||||
if any(references_id(clause) for clause in orderings):
|
||||
return False
|
||||
# Only shuffle queries that target the messages table.
|
||||
target_tables: set[str] = set()
|
||||
for from_clause in statement.get_final_froms():
|
||||
name_attr = getattr(from_clause, "name", None)
|
||||
if isinstance(name_attr, str):
|
||||
target_tables.add(name_attr)
|
||||
table_name_obj = getattr(session._messages, "name", "")
|
||||
table_name = table_name_obj if isinstance(table_name_obj, str) else ""
|
||||
return bool(table_name in target_tables)
|
||||
|
||||
@asynccontextmanager
|
||||
async def shuffled_session():
|
||||
async with real_factory() as inner:
|
||||
original_execute = inner.execute
|
||||
|
||||
async def execute_with_shuffle(statement: Any, *args: Any, **kwargs: Any) -> Any:
|
||||
result = await original_execute(statement, *args, **kwargs)
|
||||
if needs_shuffle(statement):
|
||||
rows = result.all()
|
||||
shuffled = list(rows)
|
||||
shuffled.reverse()
|
||||
return FakeResult(shuffled)
|
||||
return result
|
||||
|
||||
cast(Any, inner).execute = execute_with_shuffle
|
||||
try:
|
||||
yield inner
|
||||
finally:
|
||||
cast(Any, inner).execute = original_execute
|
||||
|
||||
session._session_factory = cast(Any, shuffled_session)
|
||||
try:
|
||||
retrieved = await session.get_items()
|
||||
assert _item_ids(retrieved) == ["older_same_ts", "rs_same_ts", "msg_same_ts"]
|
||||
|
||||
latest_two = await session.get_items(limit=2)
|
||||
assert _item_ids(latest_two) == ["rs_same_ts", "msg_same_ts"]
|
||||
finally:
|
||||
session._session_factory = real_factory
|
||||
|
||||
|
||||
async def test_pop_item_same_timestamp_returns_latest():
|
||||
"""Test that pop_item returns the newest item when timestamps tie."""
|
||||
session_id = "same_timestamp_pop_test"
|
||||
session = SQLAlchemySession.from_url(session_id, url=DB_URL, create_tables=True)
|
||||
|
||||
reasoning_item = _make_reasoning_item("rs_pop_same_ts", "...")
|
||||
message_item = _make_message_item("msg_pop_same_ts", "...")
|
||||
await session.add_items([reasoning_item, message_item])
|
||||
|
||||
async with session._session_factory() as sess:
|
||||
await sess.execute(
|
||||
text(
|
||||
"UPDATE agent_messages SET created_at = :created_at WHERE session_id = :session_id"
|
||||
),
|
||||
{
|
||||
"created_at": "2025-10-15 17:26:39.132483",
|
||||
"session_id": session.session_id,
|
||||
},
|
||||
)
|
||||
await sess.commit()
|
||||
|
||||
popped = await session.pop_item()
|
||||
assert popped is not None
|
||||
assert cast(dict[str, Any], popped)["id"] == "msg_pop_same_ts"
|
||||
|
||||
remaining = await session.get_items()
|
||||
assert _item_ids(remaining) == ["rs_pop_same_ts"]
|
||||
|
||||
|
||||
async def test_get_items_orders_by_id_for_ties():
|
||||
"""Test that get_items adds id ordering to break timestamp ties."""
|
||||
session_id = "order_by_id_test"
|
||||
session = SQLAlchemySession.from_url(session_id, url=DB_URL, create_tables=True)
|
||||
|
||||
await session.add_items(
|
||||
[
|
||||
_make_reasoning_item("rs_first", "..."),
|
||||
_make_message_item("msg_second", "..."),
|
||||
]
|
||||
)
|
||||
|
||||
real_factory = session._session_factory
|
||||
recorded: list[Any] = []
|
||||
|
||||
@asynccontextmanager
|
||||
async def wrapped_session():
|
||||
async with real_factory() as inner:
|
||||
original_execute = inner.execute
|
||||
|
||||
async def recording_execute(statement: Any, *args: Any, **kwargs: Any) -> Any:
|
||||
recorded.append(statement)
|
||||
return await original_execute(statement, *args, **kwargs)
|
||||
|
||||
cast(Any, inner).execute = recording_execute
|
||||
try:
|
||||
yield inner
|
||||
finally:
|
||||
cast(Any, inner).execute = original_execute
|
||||
|
||||
session._session_factory = cast(Any, wrapped_session)
|
||||
try:
|
||||
retrieved_full = await session.get_items()
|
||||
retrieved_limited = await session.get_items(limit=2)
|
||||
finally:
|
||||
session._session_factory = real_factory
|
||||
|
||||
assert len(recorded) >= 2
|
||||
orderings_full = [str(clause) for clause in recorded[0]._order_by_clause]
|
||||
assert orderings_full == [
|
||||
"agent_messages.created_at ASC",
|
||||
"agent_messages.id ASC",
|
||||
]
|
||||
|
||||
orderings_limited = [str(clause) for clause in recorded[1]._order_by_clause]
|
||||
assert orderings_limited == [
|
||||
"agent_messages.created_at DESC",
|
||||
"agent_messages.id DESC",
|
||||
]
|
||||
|
||||
assert _item_ids(retrieved_full) == ["rs_first", "msg_second"]
|
||||
assert _item_ids(retrieved_limited) == ["rs_first", "msg_second"]
|
||||
|
||||
|
||||
async def test_engine_property_from_url():
|
||||
"""Test that the engine property returns the AsyncEngine from from_url."""
|
||||
session_id = "engine_property_test"
|
||||
session = SQLAlchemySession.from_url(session_id, url=DB_URL, create_tables=True)
|
||||
|
||||
# Verify engine property returns an AsyncEngine instance
|
||||
assert isinstance(session.engine, AsyncEngine)
|
||||
|
||||
# Verify we can use the engine for advanced operations
|
||||
# For example, check pool status
|
||||
assert session.engine.pool is not None
|
||||
|
||||
# Verify we can manually dispose the engine
|
||||
await session.engine.dispose()
|
||||
|
||||
|
||||
async def test_engine_property_from_external_engine():
|
||||
"""Test that the engine property returns the external engine."""
|
||||
session_id = "external_engine_test"
|
||||
|
||||
# Create engine externally
|
||||
external_engine = create_async_engine(DB_URL)
|
||||
|
||||
# Create session with external engine
|
||||
session = SQLAlchemySession(session_id, engine=external_engine, create_tables=True)
|
||||
|
||||
# Verify engine property returns the same engine instance
|
||||
assert session.engine is external_engine
|
||||
|
||||
# Verify we can use the engine
|
||||
assert isinstance(session.engine, AsyncEngine)
|
||||
|
||||
# Clean up - user is responsible for disposing external engine
|
||||
await external_engine.dispose()
|
||||
|
||||
|
||||
async def test_engine_property_is_read_only():
|
||||
"""Test that the engine property cannot be modified."""
|
||||
session_id = "readonly_engine_test"
|
||||
session = SQLAlchemySession.from_url(session_id, url=DB_URL, create_tables=True)
|
||||
|
||||
# Verify engine property exists
|
||||
assert hasattr(session, "engine")
|
||||
|
||||
# Verify it's a property (read-only, cannot be set)
|
||||
# Type ignore needed because mypy correctly detects this is read-only
|
||||
with pytest.raises(AttributeError):
|
||||
session.engine = create_async_engine(DB_URL) # type: ignore[misc]
|
||||
|
||||
# Clean up
|
||||
await session.engine.dispose()
|
||||
|
||||
|
||||
async def test_session_settings_default():
|
||||
"""Test that session_settings defaults to empty SessionSettings."""
|
||||
from agents.memory import SessionSettings
|
||||
|
||||
session = SQLAlchemySession.from_url("default_settings_test", url=DB_URL, create_tables=True)
|
||||
|
||||
# Should have default SessionSettings
|
||||
assert isinstance(session.session_settings, SessionSettings)
|
||||
assert session.session_settings.limit is None
|
||||
|
||||
|
||||
async def test_session_settings_from_url():
|
||||
"""Test passing session_settings via from_url."""
|
||||
from agents.memory import SessionSettings
|
||||
|
||||
session = SQLAlchemySession.from_url(
|
||||
"from_url_settings_test",
|
||||
url=DB_URL,
|
||||
create_tables=True,
|
||||
session_settings=SessionSettings(limit=5),
|
||||
)
|
||||
|
||||
assert session.session_settings is not None
|
||||
assert session.session_settings.limit == 5
|
||||
|
||||
|
||||
async def test_get_items_uses_session_settings_limit():
|
||||
"""Test that get_items uses session_settings.limit as default."""
|
||||
from agents.memory import SessionSettings
|
||||
|
||||
session = SQLAlchemySession.from_url(
|
||||
"uses_settings_limit_test",
|
||||
url=DB_URL,
|
||||
create_tables=True,
|
||||
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"
|
||||
|
||||
|
||||
async def test_get_items_explicit_limit_overrides_session_settings():
|
||||
"""Test that explicit limit parameter overrides session_settings."""
|
||||
from agents.memory import SessionSettings
|
||||
|
||||
session = SQLAlchemySession.from_url(
|
||||
"explicit_override_test",
|
||||
url=DB_URL,
|
||||
create_tables=True,
|
||||
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"
|
||||
|
||||
|
||||
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
|
||||
|
||||
|
||||
async def test_runner_with_session_settings_override(agent: Agent):
|
||||
"""Test that RunConfig can override session's default settings."""
|
||||
from agents import RunConfig
|
||||
from agents.memory import SessionSettings
|
||||
|
||||
# Session with default limit=100
|
||||
session = SQLAlchemySession.from_url(
|
||||
"runner_override_test",
|
||||
url=DB_URL,
|
||||
create_tables=True,
|
||||
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)
|
||||
|
||||
# Use RunConfig to override limit to 2
|
||||
assert isinstance(agent.model, FakeModel)
|
||||
agent.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 = agent.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
|
||||
Reference in New Issue
Block a user