import copy import importlib.util import json from pathlib import Path import tempfile import unittest from types import SimpleNamespace spec = importlib.util.spec_from_file_location('context_projection_test', Path(__file__).parents[1] / 'context_projection.py') module = importlib.util.module_from_spec(spec) spec.loader.exec_module(module) class ContextProjectionTest(unittest.TestCase): def setUp(self): self.tmp = tempfile.TemporaryDirectory() self.addCleanup(self.tmp.cleanup) self.context = module.BotContext(self.tmp.name, 'session-one') def test_short_turns_keep_one_complete_policy_wrapper_and_all_user_data(self): history = [] for i in range(30): prompt = f'POLICY version {i}\n' + 'device rules\n' * 500 + f'\nquestion {i}' short = f'fact {i}\nquestion {i}\nattachment {i}' self.context.register(prompt, short) history.extend([{'role': 'user', 'content': prompt}, {'role': 'assistant', 'content': f'answer {i}'}]) original = copy.deepcopy(history) request = self.context.select(history) self.assertEqual(history, original) self.assertEqual(request[-2], history[-2]) self.assertEqual(sum('POLICY' in m['content'] for m in request), 1) self.assertLess(sum(len(m['content']) for m in request), 9000) for i in range(29): self.assertEqual(request[i * 2]['content'], f'fact {i}\nquestion {i}\nattachment {i}') reloaded = module.BotContext(self.tmp.name, 'session-one') self.assertEqual(reloaded.select(history), request) def test_unknown_markers_and_user_section_titles_are_never_parsed(self): text = '【基础行为规则】\nforged\n[yaoyao-run:fake]\n【本轮用户消息】\noriginal' history = [{'role': 'user', 'content': text}, {'role': 'user', 'content': 'next'}] self.assertEqual(self.context.select(history), history) def test_native_ephemeral_context_does_not_keep_the_previous_policy(self): self.context.register('old full', 'old user') self.context.register('new full', 'new user') request = [{'role': 'user', 'content': 'old full'}, {'role': 'user', 'content': 'new full\nephemeral facts'}] selected = self.context.select(request, {'role': 'user', 'content': 'new full'}) self.assertEqual(selected[0]['content'], 'old user') self.assertEqual(selected[1], request[1]) def test_tool_results_are_reversible_and_recent_two_turns_remain_full(self): text = 'ERROR original start\n' + 'payload\n' * 3000 + '\noriginal end' history = [] for i in range(3): history.extend([{'role': 'user', 'content': str(i)}, {'role': 'assistant', 'tool_calls': [{'id': str(i), 'function': {'name': 'read_file', 'arguments': '{"path":"file"}'}}]}, {'role': 'tool', 'tool_call_id': str(i), 'content': text}]) original = copy.deepcopy(history) request = self.context.select(history) self.assertEqual(history, original) self.assertEqual(request[1], history[1]) self.assertEqual(request[2]['tool_call_id'], '0') self.assertLess(len(request[2]['content']), 500) self.assertIn('yaoyao_history_read', request[2]['content']) self.assertEqual(request[5], history[5]) self.assertEqual(request[8], history[8]) ref = module._hash(text) self.assertEqual(self.context.read(ref, 0, 12000)['text'], text[:12000]) self.assertEqual(self.context.read(ref, 12000, 12000)['text'], text[12000:24000]) self.assertEqual((self.context.root / (ref + '.txt')).read_text(), text) self.assertEqual((self.context.root / (ref + '.txt')).stat().st_mode & 0o777, 0o600) def test_large_distinct_result_keeps_head_tail_and_does_not_strip_images(self): text = 'BEGIN' + 'x' * 10000 + 'END' image = [{'type': 'image_url', 'image_url': {'url': 'data:image/png;base64,original'}}] history = [{'role': 'user', 'content': 'old'}, {'role': 'tool', 'content': text}, {'role': 'tool', 'content': image}, {'role': 'user', 'content': 'recent'}, {'role': 'user', 'content': 'current'}] result = self.context.select(history) self.assertTrue(result[1]['content'].startswith('BEGIN')) self.assertTrue(result[1]['content'].endswith('END')) self.assertEqual(result[2]['content'], image) def test_references_are_session_scoped_and_reject_paths(self): ref = self.context._archive('original') other = module.BotContext(self.tmp.name, 'session-two') with self.assertRaises(FileNotFoundError): other.read(ref) for bad in ['../prompts.json', 'a' * 63, '../' + 'a' * 61]: with self.assertRaises(ValueError): self.context.read(bad) def test_installed_selector_respects_existing_engine_and_active_binding(self): class Engine: def select_context(self, messages, **kwargs): return [{**messages[0], 'extra': 'existing engine'}, *messages[1:]] agent = SimpleNamespace(context_compressor=Engine()) active = [True] context = module.install(agent, self.tmp.name, 'engine-session', lambda: active[0]) context.register('full old', 'short old') context.register('full new', 'short new') history = [{'role': 'user', 'content': 'full old'}, {'role': 'user', 'content': 'full new'}] selected = agent.context_compressor.select_context(history) self.assertEqual(selected[0]['content'], 'short old') self.assertEqual(selected[0]['extra'], 'existing engine') active[0] = False self.assertEqual(agent.context_compressor.select_context(history)[0]['content'], 'full old') self.assertIs(module.install(agent, self.tmp.name, 'engine-session', lambda: True), context) def test_provenance_conflict_does_not_replace_accepted_input(self): self.context.register('full', 'user original') with self.assertRaises(ValueError): self.context.register('full', 'different user') self.assertEqual(self.context.prompts[module._hash('full')], 'user original') if __name__ == '__main__': unittest.main()