wehub-resource-sync
577 lines
18 KiB
Python
577 lines
18 KiB
Python
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()
|