import copy import json from pathlib import Path import tempfile from types import SimpleNamespace as NS import unittest from unittest.mock import patch from engine import memory from engine.chat import generate_response from engine.saveload import export_save, import_save, restore_save, validate_save, MAX_SAVE_BYTES CHAR = 'hermione_granger' OTHER = 'jessica_jones' A, B = 'a' * 32, 'b' * 32 def answer(text): return NS(choices=[NS(message=NS(content=text))]) class Client: def __init__(self, summary=None, fail=False): self.calls = [] self.summary = summary or {'summary': 'Alex found a clue.', 'new_memories': [ {'content': 'The user is Alex.', 'importance': 'high'}]} self.fail = fail def chat_completion(self, messages, **kwargs): self.calls.append((copy.deepcopy(messages), kwargs)) if self.fail: raise RuntimeError('Network failure') if kwargs.get('stream'): return iter([NS(choices=[NS(delta=NS(content='Hello'))]), NS(choices=[NS(delta=NS(content=' there.'))])]) return answer(json.dumps(self.summary)) class ConversationFixture: def setUp(self): self.temp = tempfile.TemporaryDirectory() self.addCleanup(self.temp.cleanup) self.patch = patch.object(memory, 'MEMORY_DIR', Path(self.temp.name) / 'sessions') self.patch.start() self.addCleanup(self.patch.stop) def populated(self, session=A, char=CHAR, turns=12): state = memory.load_conversation(session, char) for i in range(turns): state['recent_messages'].extend([{'role': 'user', 'content': f'Question {i}'}, {'role': 'assistant', 'content': f'Answer {i}'}]) memory.save_conversation(session, char, state) return state def save_data(self): return {'version': '2.0', 'character_id': CHAR, 'persona': {'name': 'Alex'}, 'conversation': self.populated()} class ConversationTests(ConversationFixture, unittest.TestCase): def test_sessions_and_characters_are_isolated(self): self.populated() self.assertEqual(1, len(memory.load_conversation(B, CHAR)['recent_messages'])) self.assertEqual(1, len(memory.load_conversation(A, OTHER)['recent_messages'])) self.assertEqual(25, len(memory.load_conversation(A, CHAR)['recent_messages'])) def test_import_and_export_roundtrip_full_history_and_unique_files(self): data = self.save_data() restore_save(data, A) paths = [export_save(CHAR, data['persona'], A) for _ in range(2)] self.addCleanup(lambda: [Path(p).unlink(missing_ok=True) for p in paths]) self.assertNotEqual(*paths) imported = import_save(paths[0]) restore_save(imported, B) self.assertEqual(imported['conversation']['recent_messages'], memory.load_conversation(B, CHAR)['recent_messages']) self.assertEqual(25, len(imported['conversation']['recent_messages'])) def test_invalid_import_cannot_mutate_existing_state(self): original = self.populated() for field, value in [('recent_messages', 'bad'), ('summary', []), ('long_term_memory', {}), ('summarized_count', True), ('summarized_count', -1)]: data = {'version': '2.0', 'character_id': CHAR, 'conversation': copy.deepcopy(original)} data['conversation'][field] = value with self.subTest(field=field, value=value), self.assertRaises(ValueError): restore_save(data, A) self.assertEqual(original, memory.load_conversation(A, CHAR)) def test_traversal_unknown_character_version_and_persona_rejected(self): data = self.save_data() for key, values in {'character_id': ['../outside', '/tmp/outside', 'missing_character', None], 'version': ['0.1', None], 'persona': [[], 'bad']}.items(): for value in values: altered = copy.deepcopy(data) altered[key] = value with self.subTest(key=key, value=value), self.assertRaises(ValueError): restore_save(altered, A) self.assertFalse((Path(self.temp.name) / 'outside').exists()) def test_invalid_message_and_memory_types_rejected(self): data = self.save_data() for message in [{'role': 'system', 'content': 'x'}, {'role': 'user', 'content': []}, None]: data['conversation']['recent_messages'] = [message] with self.assertRaises(ValueError): validate_save(data) data = self.save_data() for memories in [[{'content': 'x', 'importance': 'critical'}], [None], [{'content': [], 'importance': 'high'}]]: data['conversation']['long_term_memory'] = memories with self.assertRaises(ValueError): validate_save(data) def test_legacy_save_is_read_without_loading_shared_files(self): data = self.save_data() data['version'] = '1.0' data['conversation'].pop('summarized_count') restore_save(data, A) self.assertEqual(0, memory.load_conversation(A, CHAR)['summarized_count']) def test_oversized_and_invalid_json_rejected(self): path = Path(self.temp.name) / 'upload.json' for raw in [b'x' * (MAX_SAVE_BYTES + 1), b'{', b'\xff']: path.write_bytes(raw) with self.assertRaises(ValueError): import_save(path) def test_summary_preserves_full_transcript_and_shortens_context(self): state = self.populated() client = Client() self.assertTrue(memory.update_memory(A, CHAR, state, client, 'test')) result = memory.load_conversation(A, CHAR) self.assertEqual(state['recent_messages'], result['recent_messages']) self.assertGreater(result['summarized_count'], 0) self.assertTrue(result['summary']) self.assertEqual('high', result['long_term_memory'][0]['importance']) client = Client() list(generate_response('Next', CHAR, client, A, result['recent_messages'])) messages = client.calls[0][0] self.assertIn(result['summary'], messages[0]['content']) self.assertEqual(result['recent_messages'][result['summarized_count']:], messages[1:-1]) def test_summary_failure_does_not_drop_messages(self): state = self.populated() for client in [Client(fail=True), Client(summary={'summary': '', 'new_memories': []})]: self.assertFalse(memory.update_memory(A, CHAR, state, client, 'test')) self.assertEqual(state, memory.load_conversation(A, CHAR)) def test_summary_stale_result_cannot_overwrite_reset(self): state = self.populated() memory.reset_conversation(A, CHAR) self.assertFalse(memory.update_memory(A, CHAR, state, Client(), 'test')) self.assertEqual(1, len(memory.load_conversation(A, CHAR)['recent_messages'])) def test_visible_history_controls_prompt_and_retry(self): old = self.populated(turns=2) history = old['recent_messages'][:-2] client = Client() output = list(generate_response('Replacement', CHAR, client, A, history)) self.assertEqual('Hello there.', output[-1]) self.assertEqual(history, client.calls[0][0][1:-1]) saved = memory.load_conversation(A, CHAR) self.assertEqual(history + [{'role': 'user', 'content': 'Replacement'}, {'role': 'assistant', 'content': output[-1]}], saved['recent_messages']) def test_translation_off_is_default_and_uses_one_call(self): client = Client() history = memory.load_conversation(A, CHAR)['recent_messages'] list(generate_response('Merhaba', CHAR, client, A, history, language='Turkish', translation_model='translator')) self.assertEqual(1, len(client.calls)) self.assertIn('Türkçe', client.calls[0][0][0]['content']) def test_translation_stores_visible_language(self): history = memory.load_conversation(A, CHAR)['recent_messages'] with patch('engine.chat.translate_to_english', return_value='Hello'), patch('engine.chat.translate_from_english', return_value='Merhaba!'): output = list(generate_response('Selam', CHAR, Client(), A, history, language='Turkish', translation_model='translator', use_translator=True)) state = memory.load_conversation(A, CHAR) self.assertEqual('Selam', state['recent_messages'][-2]['content']) self.assertEqual('Merhaba!', state['recent_messages'][-1]['content']) self.assertEqual('Merhaba!', output[-1]) def test_clear_discards_all_memory_layers(self): state = self.populated() memory.update_memory(A, CHAR, state, Client(), 'test') cleared = memory.sync_history(A, CHAR, []) self.assertEqual([], cleared['recent_messages']) self.assertEqual('', cleared['summary']) self.assertEqual([], cleared['long_term_memory']) def test_retry_keeps_unchanged_summary_but_editing_old_turn_resets_it(self): state = self.populated() memory.update_memory(A, CHAR, state, Client(), 'test') summarized = memory.load_conversation(A, CHAR) retried = memory.sync_history(A, CHAR, summarized['recent_messages'][:-2]) self.assertEqual(summarized['summary'], retried['summary']) self.assertEqual(summarized['long_term_memory'], retried['long_term_memory']) edited = copy.deepcopy(retried['recent_messages']) edited[1]['content'] = 'Different story' updated = memory.sync_history(A, CHAR, edited) self.assertEqual('', updated['summary']) self.assertEqual([], updated['long_term_memory']) self.assertEqual(0, updated['summarized_count']) def test_repeated_summarization_is_cumulative_and_deduplicates(self): state = self.populated(turns=50) first = Client() memory.update_memory(A, CHAR, state, first, 'test') snapshot = memory.load_conversation(A, CHAR) second = Client() memory.update_memory(A, CHAR, snapshot, second, 'test') payload = json.loads(second.calls[0][0][1]['content']) self.assertEqual(snapshot['summary'], payload['previous_summary']) self.assertEqual(snapshot['recent_messages'][snapshot['summarized_count']], payload['conversation'][0]) result = memory.load_conversation(A, CHAR) self.assertEqual(1, len(result['long_term_memory'])) self.assertEqual(101, len(result['recent_messages'])) self.assertGreater(result['summarized_count'], snapshot['summarized_count']) def test_atomic_write_failure_keeps_previous_snapshot(self): state = self.populated() modified = copy.deepcopy(state) modified['summary'] = 'Replacement' with patch('engine.memory.os.replace', side_effect=OSError('Disk failure')): with self.assertRaises(OSError): memory.save_conversation(A, CHAR, modified) self.assertEqual(state, memory.load_conversation(A, CHAR)) self.assertEqual([], list(Path(self.temp.name).rglob('*.tmp'))) def test_chat_reports_storage_failure_without_losing_previous_history(self): history = memory.load_conversation(A, CHAR)['recent_messages'] with patch('engine.chat.append_exchange', side_effect=OSError('Disk full')): output = list(generate_response('Hi', CHAR, Client(), A, history)) self.assertIn('could not be saved', output[-1]) self.assertEqual(history, memory.load_conversation(A, CHAR)['recent_messages']) def test_gradio_text_blocks_normalized(self): result = memory.normalize_history([{'role': 'user', 'content': [{'type': 'text', 'text': 'Hello'}]}]) self.assertEqual([{'role': 'user', 'content': 'Hello'}], result) if __name__ == '__main__': unittest.main()