Add comprehensive tests for session_store module

Cover all 12 areas: StoredSession/StoredAgentSession round-trips,
model/runtime config serialization, usage_from_payload, _optional_int,
_optional_float, _deserialize_output_schema, default handling, and
non-dict message filtering. 50 test cases total.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Co-authored-by: abdoelsayed2016 <27821589+abdoelsayed2016@users.noreply.github.com>

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