diff --git a/tests/test_session_store.py b/tests/test_session_store.py new file mode 100644 index 0000000..de2b6b5 --- /dev/null +++ b/tests/test_session_store.py @@ -0,0 +1,553 @@ +from __future__ import annotations + +import json +import tempfile +import unittest +from pathlib import Path + +from src.agent_types import ( + AgentPermissions, + AgentRuntimeConfig, + BudgetConfig, + ModelConfig, + ModelPricing, + OutputSchemaConfig, + UsageStats, +) +from src.session_store import ( + StoredAgentSession, + StoredSession, + _deserialize_output_schema, + _optional_float, + _optional_int, + deserialize_model_config, + deserialize_runtime_config, + load_agent_session, + load_session, + save_agent_session, + save_session, + serialize_model_config, + serialize_runtime_config, + usage_from_payload, +) + + +class TestStoredSessionRoundTrip(unittest.TestCase): + """save_session then load_session preserves all fields.""" + + def test_round_trip(self) -> None: + session = StoredSession( + session_id='abc-123', + messages=('hello', 'world', 'foo'), + input_tokens=100, + output_tokens=200, + ) + with tempfile.TemporaryDirectory() as td: + directory = Path(td) + save_session(session, directory=directory) + loaded = load_session('abc-123', directory=directory) + + self.assertEqual(loaded.session_id, session.session_id) + self.assertEqual(loaded.messages, session.messages) + self.assertEqual(loaded.input_tokens, session.input_tokens) + self.assertEqual(loaded.output_tokens, session.output_tokens) + + def test_round_trip_empty_messages(self) -> None: + session = StoredSession( + session_id='empty', + messages=(), + input_tokens=0, + output_tokens=0, + ) + with tempfile.TemporaryDirectory() as td: + directory = Path(td) + save_session(session, directory=directory) + loaded = load_session('empty', directory=directory) + + self.assertEqual(loaded.messages, ()) + self.assertEqual(loaded.input_tokens, 0) + + +class TestStoredAgentSessionRoundTrip(unittest.TestCase): + """save_agent_session then load_agent_session preserves all fields.""" + + def _make_session(self, **overrides: object) -> StoredAgentSession: + defaults: dict = { + 'session_id': 'agent-001', + 'model_config': {'model': 'gpt-4', 'temperature': 0.5}, + 'runtime_config': {'cwd': '/home/user', 'max_turns': 20}, + 'system_prompt_parts': ('You are helpful.',), + 'user_context': {'lang': 'en'}, + 'system_context': {'os': 'linux'}, + 'messages': ({'role': 'user', 'content': 'hi'},), + 'turns': 3, + 'tool_calls': 7, + 'usage': {'input_tokens': 500, 'output_tokens': 300}, + 'total_cost_usd': 0.05, + 'file_history': ({'file': 'a.py', 'action': 'edit'},), + 'budget_state': {'remaining': 100}, + 'plugin_state': {'key': 'value'}, + 'scratchpad_directory': '/scratch/pad', + } + defaults.update(overrides) + return StoredAgentSession(**defaults) + + def test_round_trip_all_fields(self) -> None: + session = self._make_session() + with tempfile.TemporaryDirectory() as td: + directory = Path(td) + save_agent_session(session, directory=directory) + loaded = load_agent_session('agent-001', directory=directory) + + self.assertEqual(loaded.session_id, session.session_id) + self.assertEqual(loaded.model_config, session.model_config) + self.assertEqual(loaded.runtime_config, session.runtime_config) + self.assertEqual(loaded.system_prompt_parts, session.system_prompt_parts) + self.assertEqual(loaded.user_context, session.user_context) + self.assertEqual(loaded.system_context, session.system_context) + self.assertEqual(loaded.messages, session.messages) + self.assertEqual(loaded.turns, session.turns) + self.assertEqual(loaded.tool_calls, session.tool_calls) + self.assertEqual(loaded.usage, session.usage) + self.assertAlmostEqual(loaded.total_cost_usd, session.total_cost_usd) + self.assertEqual(loaded.file_history, session.file_history) + self.assertEqual(loaded.budget_state, session.budget_state) + self.assertEqual(loaded.plugin_state, session.plugin_state) + self.assertEqual(loaded.scratchpad_directory, session.scratchpad_directory) + + def test_round_trip_no_scratchpad(self) -> None: + session = self._make_session(scratchpad_directory=None) + with tempfile.TemporaryDirectory() as td: + directory = Path(td) + save_agent_session(session, directory=directory) + loaded = load_agent_session('agent-001', directory=directory) + + self.assertIsNone(loaded.scratchpad_directory) + + def test_load_filters_non_dict_messages(self) -> None: + """Non-dict entries in messages list are filtered out on load.""" + with tempfile.TemporaryDirectory() as td: + directory = Path(td) + path = directory / 'mixed.json' + data = { + 'session_id': 'mixed', + 'model_config': {}, + 'runtime_config': {'cwd': '/'}, + 'system_prompt_parts': [], + 'user_context': {}, + 'system_context': {}, + 'messages': [ + {'role': 'user', 'content': 'hi'}, + 'not a dict', + 42, + None, + {'role': 'assistant', 'content': 'hey'}, + ], + 'turns': 0, + 'tool_calls': 0, + 'usage': {}, + 'total_cost_usd': 0.0, + 'file_history': [], + 'budget_state': {}, + 'plugin_state': {}, + } + path.write_text(json.dumps(data)) + loaded = load_agent_session('mixed', directory=directory) + + self.assertEqual(len(loaded.messages), 2) + self.assertEqual(loaded.messages[0]['role'], 'user') + self.assertEqual(loaded.messages[1]['role'], 'assistant') + + def test_load_defaults_for_missing_optional_fields(self) -> None: + """Missing optional fields get sensible defaults.""" + with tempfile.TemporaryDirectory() as td: + directory = Path(td) + path = directory / 'minimal.json' + data = { + 'session_id': 'minimal', + 'model_config': {}, + 'runtime_config': {'cwd': '/'}, + 'system_prompt_parts': [], + 'user_context': {}, + 'system_context': {}, + 'messages': [], + 'turns': 1, + 'tool_calls': 2, + } + path.write_text(json.dumps(data)) + loaded = load_agent_session('minimal', directory=directory) + + self.assertEqual(loaded.usage, {}) + self.assertAlmostEqual(loaded.total_cost_usd, 0.0) + self.assertEqual(loaded.file_history, ()) + self.assertEqual(loaded.budget_state, {}) + self.assertEqual(loaded.plugin_state, {}) + self.assertIsNone(loaded.scratchpad_directory) + + def test_load_non_dict_budget_state_defaults_to_empty(self) -> None: + with tempfile.TemporaryDirectory() as td: + directory = Path(td) + path = directory / 'bad-budget.json' + data = { + 'session_id': 'bad-budget', + 'model_config': {}, + 'runtime_config': {'cwd': '/'}, + 'system_prompt_parts': [], + 'user_context': {}, + 'system_context': {}, + 'messages': [], + 'turns': 0, + 'tool_calls': 0, + 'budget_state': 'not-a-dict', + 'plugin_state': 123, + } + path.write_text(json.dumps(data)) + loaded = load_agent_session('bad-budget', directory=directory) + + self.assertEqual(loaded.budget_state, {}) + self.assertEqual(loaded.plugin_state, {}) + + +class TestModelConfigSerialization(unittest.TestCase): + """serialize_model_config + deserialize_model_config round-trip.""" + + def test_round_trip_preserves_pricing(self) -> None: + pricing = ModelPricing( + input_cost_per_million_tokens_usd=3.0, + output_cost_per_million_tokens_usd=15.0, + cache_creation_input_cost_per_million_tokens_usd=1.5, + cache_read_input_cost_per_million_tokens_usd=0.5, + ) + config = ModelConfig( + model='claude-3-sonnet', + base_url='https://api.example.com/v1', + api_key='sk-test-key', + temperature=0.7, + timeout_seconds=60.0, + pricing=pricing, + ) + payload = serialize_model_config(config) + restored = deserialize_model_config(payload) + + self.assertEqual(restored.model, config.model) + self.assertEqual(restored.base_url, config.base_url) + self.assertEqual(restored.api_key, config.api_key) + self.assertAlmostEqual(restored.temperature, config.temperature) + self.assertAlmostEqual(restored.timeout_seconds, config.timeout_seconds) + self.assertAlmostEqual( + restored.pricing.input_cost_per_million_tokens_usd, + pricing.input_cost_per_million_tokens_usd, + ) + self.assertAlmostEqual( + restored.pricing.output_cost_per_million_tokens_usd, + pricing.output_cost_per_million_tokens_usd, + ) + self.assertAlmostEqual( + restored.pricing.cache_creation_input_cost_per_million_tokens_usd, + pricing.cache_creation_input_cost_per_million_tokens_usd, + ) + self.assertAlmostEqual( + restored.pricing.cache_read_input_cost_per_million_tokens_usd, + pricing.cache_read_input_cost_per_million_tokens_usd, + ) + + def test_deserialize_defaults_for_missing_fields(self) -> None: + payload = {'model': 'gpt-4'} + config = deserialize_model_config(payload) + + self.assertEqual(config.model, 'gpt-4') + self.assertEqual(config.base_url, 'http://127.0.0.1:8000/v1') + self.assertEqual(config.api_key, 'local-token') + self.assertAlmostEqual(config.temperature, 0.0) + self.assertAlmostEqual(config.timeout_seconds, 120.0) + self.assertAlmostEqual(config.pricing.input_cost_per_million_tokens_usd, 0.0) + self.assertAlmostEqual(config.pricing.output_cost_per_million_tokens_usd, 0.0) + + def test_deserialize_with_non_dict_pricing(self) -> None: + payload = {'model': 'test', 'pricing': 'invalid'} + config = deserialize_model_config(payload) + self.assertAlmostEqual(config.pricing.input_cost_per_million_tokens_usd, 0.0) + + def test_deserialize_with_none_pricing(self) -> None: + payload = {'model': 'test', 'pricing': None} + config = deserialize_model_config(payload) + self.assertEqual(config.pricing, ModelPricing()) + + +class TestRuntimeConfigSerialization(unittest.TestCase): + """serialize_runtime_config + deserialize_runtime_config round-trip.""" + + def test_round_trip_preserves_all(self) -> None: + config = AgentRuntimeConfig( + cwd=Path('/home/user/project'), + max_turns=25, + command_timeout_seconds=45.0, + max_output_chars=8000, + stream_model_responses=True, + auto_snip_threshold_tokens=5000, + auto_compact_threshold_tokens=10000, + compact_preserve_messages=6, + permissions=AgentPermissions( + allow_file_write=True, + allow_shell_commands=True, + allow_destructive_shell_commands=False, + ), + additional_working_directories=(Path('/extra/dir'),), + disable_claude_md_discovery=True, + budget_config=BudgetConfig( + max_total_tokens=100000, + max_input_tokens=50000, + max_output_tokens=30000, + max_reasoning_tokens=20000, + max_total_cost_usd=5.0, + max_tool_calls=100, + max_delegated_tasks=10, + max_model_calls=200, + max_session_turns=50, + ), + output_schema=OutputSchemaConfig( + name='test_schema', + schema={'type': 'object', 'properties': {'answer': {'type': 'string'}}}, + strict=True, + ), + session_directory=Path('/sessions'), + scratchpad_root=Path('/scratch'), + ) + payload = serialize_runtime_config(config) + restored = deserialize_runtime_config(payload) + + self.assertEqual(restored.cwd, config.cwd.resolve()) + self.assertEqual(restored.max_turns, 25) + self.assertAlmostEqual(restored.command_timeout_seconds, 45.0) + self.assertEqual(restored.max_output_chars, 8000) + self.assertTrue(restored.stream_model_responses) + self.assertEqual(restored.auto_snip_threshold_tokens, 5000) + self.assertEqual(restored.auto_compact_threshold_tokens, 10000) + self.assertEqual(restored.compact_preserve_messages, 6) + self.assertTrue(restored.permissions.allow_file_write) + self.assertTrue(restored.permissions.allow_shell_commands) + self.assertFalse(restored.permissions.allow_destructive_shell_commands) + self.assertTrue(restored.disable_claude_md_discovery) + + self.assertEqual(restored.budget_config.max_total_tokens, 100000) + self.assertEqual(restored.budget_config.max_input_tokens, 50000) + self.assertEqual(restored.budget_config.max_output_tokens, 30000) + self.assertEqual(restored.budget_config.max_reasoning_tokens, 20000) + self.assertAlmostEqual(restored.budget_config.max_total_cost_usd, 5.0) + self.assertEqual(restored.budget_config.max_tool_calls, 100) + self.assertEqual(restored.budget_config.max_delegated_tasks, 10) + self.assertEqual(restored.budget_config.max_model_calls, 200) + self.assertEqual(restored.budget_config.max_session_turns, 50) + + self.assertIsNotNone(restored.output_schema) + assert restored.output_schema is not None + self.assertEqual(restored.output_schema.name, 'test_schema') + self.assertEqual(restored.output_schema.schema, config.output_schema.schema) + self.assertTrue(restored.output_schema.strict) + + def test_round_trip_none_output_schema(self) -> None: + config = AgentRuntimeConfig( + cwd=Path('/home/user'), + output_schema=None, + ) + payload = serialize_runtime_config(config) + restored = deserialize_runtime_config(payload) + self.assertIsNone(restored.output_schema) + + def test_deserialize_defaults_for_missing_fields(self) -> None: + payload = {'cwd': '/home/user'} + config = deserialize_runtime_config(payload) + + self.assertEqual(config.max_turns, 12) + self.assertAlmostEqual(config.command_timeout_seconds, 30.0) + self.assertEqual(config.max_output_chars, 12000) + self.assertFalse(config.stream_model_responses) + self.assertIsNone(config.auto_snip_threshold_tokens) + self.assertIsNone(config.auto_compact_threshold_tokens) + self.assertEqual(config.compact_preserve_messages, 4) + self.assertFalse(config.permissions.allow_file_write) + self.assertFalse(config.permissions.allow_shell_commands) + self.assertFalse(config.permissions.allow_destructive_shell_commands) + self.assertEqual(config.additional_working_directories, ()) + self.assertFalse(config.disable_claude_md_discovery) + self.assertIsNone(config.budget_config.max_total_tokens) + self.assertIsNone(config.output_schema) + + def test_deserialize_non_dict_permissions(self) -> None: + payload = {'cwd': '/home', 'permissions': 'invalid'} + config = deserialize_runtime_config(payload) + self.assertFalse(config.permissions.allow_file_write) + + def test_deserialize_non_dict_budget_config(self) -> None: + payload = {'cwd': '/home', 'budget_config': 42} + config = deserialize_runtime_config(payload) + self.assertIsNone(config.budget_config.max_total_tokens) + + +class TestUsageFromPayload(unittest.TestCase): + """usage_from_payload correctly maps fields including defaults.""" + + def test_full_payload(self) -> None: + payload = { + 'input_tokens': 1000, + 'output_tokens': 500, + 'cache_creation_input_tokens': 200, + 'cache_read_input_tokens': 100, + 'reasoning_tokens': 50, + } + usage = usage_from_payload(payload) + self.assertEqual(usage.input_tokens, 1000) + self.assertEqual(usage.output_tokens, 500) + self.assertEqual(usage.cache_creation_input_tokens, 200) + self.assertEqual(usage.cache_read_input_tokens, 100) + self.assertEqual(usage.reasoning_tokens, 50) + + def test_partial_payload_uses_defaults(self) -> None: + payload = {'input_tokens': 42} + usage = usage_from_payload(payload) + self.assertEqual(usage.input_tokens, 42) + self.assertEqual(usage.output_tokens, 0) + self.assertEqual(usage.cache_creation_input_tokens, 0) + self.assertEqual(usage.cache_read_input_tokens, 0) + self.assertEqual(usage.reasoning_tokens, 0) + + def test_none_returns_empty(self) -> None: + usage = usage_from_payload(None) + self.assertEqual(usage, UsageStats()) + + def test_empty_dict_returns_defaults(self) -> None: + usage = usage_from_payload({}) + self.assertEqual(usage, UsageStats()) + + def test_non_dict_returns_empty(self) -> None: + usage = usage_from_payload('not a dict') # type: ignore[arg-type] + self.assertEqual(usage, UsageStats()) + + def test_string_token_values_parsed(self) -> None: + payload = {'input_tokens': '99', 'output_tokens': '77'} + usage = usage_from_payload(payload) + self.assertEqual(usage.input_tokens, 99) + self.assertEqual(usage.output_tokens, 77) + + +class TestOptionalInt(unittest.TestCase): + """_optional_int handles int, str, float, None, bool correctly.""" + + def test_int_value(self) -> None: + self.assertEqual(_optional_int(42), 42) + + def test_zero(self) -> None: + self.assertEqual(_optional_int(0), 0) + + def test_negative(self) -> None: + self.assertEqual(_optional_int(-5), -5) + + def test_str_numeric(self) -> None: + self.assertEqual(_optional_int('123'), 123) + + def test_float_value(self) -> None: + self.assertEqual(_optional_int(3.9), 3) + + def test_none_returns_none(self) -> None: + self.assertIsNone(_optional_int(None)) + + def test_bool_true_returns_none(self) -> None: + self.assertIsNone(_optional_int(True)) + + def test_bool_false_returns_none(self) -> None: + self.assertIsNone(_optional_int(False)) + + def test_non_numeric_string_returns_none(self) -> None: + self.assertIsNone(_optional_int('hello')) + + def test_empty_string_returns_none(self) -> None: + self.assertIsNone(_optional_int('')) + + +class TestOptionalFloat(unittest.TestCase): + """_optional_float handles int, str, float, None, bool correctly.""" + + def test_float_value(self) -> None: + self.assertAlmostEqual(_optional_float(3.14), 3.14) + + def test_int_value(self) -> None: + self.assertAlmostEqual(_optional_float(42), 42.0) + + def test_zero(self) -> None: + self.assertAlmostEqual(_optional_float(0), 0.0) + + def test_str_numeric(self) -> None: + self.assertAlmostEqual(_optional_float('2.5'), 2.5) + + def test_none_returns_none(self) -> None: + self.assertIsNone(_optional_float(None)) + + def test_bool_true_returns_none(self) -> None: + self.assertIsNone(_optional_float(True)) + + def test_bool_false_returns_none(self) -> None: + self.assertIsNone(_optional_float(False)) + + def test_non_numeric_string_returns_none(self) -> None: + self.assertIsNone(_optional_float('abc')) + + def test_empty_string_returns_none(self) -> None: + self.assertIsNone(_optional_float('')) + + +class TestDeserializeOutputSchema(unittest.TestCase): + """_deserialize_output_schema with valid, None, invalid data.""" + + def test_valid_payload(self) -> None: + payload = { + 'name': 'my_schema', + 'schema': {'type': 'object'}, + 'strict': True, + } + result = _deserialize_output_schema(payload) + self.assertIsNotNone(result) + assert result is not None + self.assertEqual(result.name, 'my_schema') + self.assertEqual(result.schema, {'type': 'object'}) + self.assertTrue(result.strict) + + def test_strict_defaults_false(self) -> None: + payload = { + 'name': 'basic', + 'schema': {'type': 'string'}, + } + result = _deserialize_output_schema(payload) + self.assertIsNotNone(result) + assert result is not None + self.assertFalse(result.strict) + + def test_none_payload(self) -> None: + self.assertIsNone(_deserialize_output_schema(None)) + + def test_non_dict_payload(self) -> None: + self.assertIsNone(_deserialize_output_schema('not a dict')) + self.assertIsNone(_deserialize_output_schema(42)) + self.assertIsNone(_deserialize_output_schema([])) + + def test_missing_schema_key(self) -> None: + self.assertIsNone(_deserialize_output_schema({'name': 'test'})) + + def test_non_dict_schema(self) -> None: + self.assertIsNone(_deserialize_output_schema({'name': 'test', 'schema': 'bad'})) + + def test_missing_name(self) -> None: + self.assertIsNone(_deserialize_output_schema({'schema': {'type': 'object'}})) + + def test_empty_name(self) -> None: + self.assertIsNone( + _deserialize_output_schema({'name': '', 'schema': {'type': 'object'}}) + ) + + def test_non_string_name(self) -> None: + self.assertIsNone( + _deserialize_output_schema({'name': 123, 'schema': {'type': 'object'}}) + ) + + +if __name__ == '__main__': + unittest.main()