business-assistant/src/step3_question_answer/tests/test_context_reuse.py

"""Synthetic Redis-only tests: no configuration or external connections."""

import json
import threading
import unittest
from concurrent.futures import ThreadPoolExecutor
from unittest.mock import patch

from langgraph.checkpoint.base import empty_checkpoint
from langgraph.types import Command
from redis.exceptions import ConnectionError

from src.step3_question_answer.context import (
    ContextConfig, ContextError, MIN_REUSE_BYTES, READ_CHECKPOINT, RedisContextSaver,
)
from src.step3_question_answer.graph import build_graph
from src.step3_question_answer.tests.test_context import FakeRedis, fake_saver
from src.step3_question_answer.timing import Timings, bind
from src.step4_api.runtime import GraphRuntime


def config(thread='synthetic-reuse', ns='', cid=None):
    values = {'thread_id': thread, 'checkpoint_ns': ns}
    if cid is not None:
        values['checkpoint_id'] = cid
    return {'configurable': values}


def checkpoint(cid='001', text=None):
    result = empty_checkpoint()
    result['id'] = cid
    result['channel_values'] = {
        'response': '仅限合成地区；有效期至合成日期；不含例外主体。',
        'recommendation': {
            'sources': [{'source_id': 'synthetic-source', 'path': '条件.例外',
                         'id': 'synthetic-source',
                         'content': ('合成原文条件、否定和尾部例外。' * 14000) if text is None else text}],
            'summary': {'text': '合成答复[依据](https://example.invalid/policy)'},
        },
    }
    return result


def measured_read(saver, cfg):
    timing = Timings()
    with bind(timing, 'session_result_load'):
        result = saver.get_tuple(cfg)
    return result, timing.snapshot()['measurements']


def total(counts, name):
    return counts[name + '.session_result_load']['total']


class CheckpointReuseTests(unittest.TestCase):
    def setUp(self):
        self.saver = fake_saver()
        self.cfg = config()
        self.key = self.saver.key('synthetic-reuse')

    def put(self, cid='001', cfg=None, text=None):
        return self.saver.put(cfg or self.cfg, checkpoint(cid, text), {'step': 1}, {})

    def test_large_body_wire_reduction_full_answer_and_fresh_pending(self):
        with self.saver.reuse_written_checkpoint('synthetic-reuse'):
            saved = self.put()
            # New writes after checkpoint storage must never come from cache.
            self.saver.put_writes(saved, [('answer', '完整合成回答'), ('__error__', '旧错误')], 'task')
            self.saver.put_writes(saved, [('__error__', '新错误')], 'task')
            cached, counts = measured_read(self.saver, self.cfg)
            self.assertEqual(total(counts, 'redis_checkpoint_cache_hit'), 1)
            self.assertIn(('task', '__error__', '新错误'), cached.pending_writes)
            self.assertIn(('task', 'answer', '完整合成回答'), cached.pending_writes)
            self.assertEqual(len(self.saver._reuse_scopes[self.key]), 1)
        normal, full = measured_read(self.saver, self.cfg)
        self.assertEqual(cached, normal)
        self.assertEqual(self.saver._reuse_scopes, {})
        self.assertLess(total(counts, 'redis_returned_bytes'), total(full, 'redis_returned_bytes') // 100)
        self.assertEqual(total(full, 'redis_returned_bytes') - total(counts, 'redis_returned_bytes'),
                         total(counts, 'redis_checkpoint_cache_saved_bytes'))
        self.assertEqual(total(counts, 'checkpoint_decoded_input_bytes'), total(full, 'checkpoint_decoded_input_bytes'))
        self.assertNotIn('合成原文', json.dumps(counts, ensure_ascii=False))

    def test_new_request_restart_and_old_storage_use_full_read(self):
        with self.saver.reuse_written_checkpoint('synthetic-reuse'):
            self.put()
            original, counts = measured_read(self.saver, self.cfg)
            self.assertEqual(total(counts, 'redis_checkpoint_cache_hit'), 1)
        second = RedisContextSaver(ContextConfig('redis://localhost/0', 30), self.saver.client)
        for saver in (self.saver, second):
            with saver.reuse_written_checkpoint('synthetic-reuse'):
                restored, counts = measured_read(saver, self.cfg)
            self.assertEqual(restored, original)
            self.assertEqual(total(counts, 'redis_checkpoint_cache_hit'), 0)
            self.assertGreater(total(counts, 'redis_returned_bytes'), MIN_REUSE_BYTES)
        self.assertTrue(all(json.loads(name)[0] == 'c' for name in self.saver.client.rows[self.key]))

    def test_external_newer_checkpoint_and_same_id_rewrite_are_authoritative(self):
        other = RedisContextSaver(ContextConfig('redis://localhost/0', 30), self.saver.client)
        with self.saver.reuse_written_checkpoint('synthetic-reuse'):
            self.put()
            # Same field and same byte length, different content: length/ID alone
            # would miss this update, the Redis fingerprint must reject reuse.
            other.put(self.cfg, checkpoint('001', '合成新原文。' * 30000), {}, {})
            state, counts = measured_read(self.saver, self.cfg)
            self.assertEqual(total(counts, 'redis_checkpoint_cache_hit'), 0)
            self.assertEqual(state.checkpoint['channel_values']['recommendation']['sources'][0]['content'], '合成新原文。' * 30000)
            other.put(self.cfg, checkpoint('002'), {}, {})
            state, counts = measured_read(self.saver, self.cfg)
            self.assertEqual(state.checkpoint['id'], '002')
            self.assertEqual(total(counts, 'redis_checkpoint_cache_hit'), 0)

    def test_same_length_changed_body_corrupt_body_and_marker_string_fail_closed(self):
        with self.saver.reuse_written_checkpoint('synthetic-reuse'):
            self.put(text='甲' * 60000)
            name = self.saver._field('c', '', '001')
            before = self.saver.client.rows[self.key][name]
            # Reuse exact metadata from the original to hold encoded length fixed.
            value, metadata, parent = self.saver._decode(before)
            value['channel_values']['recommendation']['sources'][0]['content'] = '乙' * 60000
            replaced = self.saver._encode((value, metadata, parent))
            self.assertEqual(len(before), len(replaced))
            self.saver.client.rows[self.key][name] = replaced.encode()
            changed, counts = measured_read(self.saver, self.cfg)
            self.assertEqual(total(counts, 'redis_checkpoint_cache_hit'), 0)
            self.assertEqual(changed.checkpoint['channel_values']['recommendation']['sources'][0]['content'], '乙' * 60000)
            for invalid in ('1', '__reuse_written_checkpoint_v1__', 'invalid-json'):
                self.saver.client.rows[self.key][name] = invalid
                with self.assertRaisesRegex(ContextError, '^context_data_invalid$'):
                    self.saver.get_tuple(self.cfg)

    def test_namespace_explicit_history_parent_and_latest_selection(self):
        with self.saver.reuse_written_checkpoint('synthetic-reuse'):
            first = self.put('001')
            second = self.put('002', first)
            current, counts = measured_read(self.saver, second)
            self.assertEqual(current.parent_config, first)
            self.assertEqual(total(counts, 'redis_checkpoint_cache_hit'), 1)
            old, counts = measured_read(self.saver, first)
            self.assertEqual(old.checkpoint['id'], '001')
            self.assertEqual(total(counts, 'redis_checkpoint_cache_hit'), 0)
            self.put('999', config(ns='other'))
            main, counts = measured_read(self.saver, self.cfg)
            self.assertEqual(main.checkpoint['id'], '002')
            self.assertEqual(total(counts, 'redis_checkpoint_cache_hit'), 0)
            branch, counts = measured_read(self.saver, config(ns='other'))
            self.assertEqual(branch.checkpoint['id'], '999')
            self.assertEqual(total(counts, 'redis_checkpoint_cache_hit'), 1)
            self.assertIsNone(self.saver.get_tuple(config(ns='missing')))
            self.assertIsNone(self.saver.get_tuple(config(cid='missing')))
            self.assertEqual(len(list(self.saver.list(self.cfg))), 2)

    def test_failure_no_stale_fallback_and_scope_cleanup(self):
        with self.assertRaisesRegex(RuntimeError, 'synthetic-failure'):
            with self.saver.reuse_written_checkpoint('synthetic-reuse'):
                self.put()
                with patch.object(self.saver.client, 'eval', side_effect=ConnectionError('private')):
                    with self.assertRaisesRegex(ContextError, '^redis_connection_failed$'):
                        self.put('002')
                    with self.assertRaisesRegex(ContextError, '^redis_connection_failed$'):
                        self.saver.get_tuple(self.cfg)
                state, counts = measured_read(self.saver, self.cfg)
                self.assertEqual(state.checkpoint['id'], '001')
                self.assertEqual(total(counts, 'redis_checkpoint_cache_hit'), 1)
                raise RuntimeError('synthetic-failure')
        self.assertEqual(self.saver._reuse_scopes, {})

    def test_expiry_deletion_and_reads_do_not_renew_ttl(self):
        with self.saver.reuse_written_checkpoint('synthetic-reuse'):
            self.put()
            self.saver.client.now = 29
            self.assertIsNotNone(self.saver.get_tuple(self.cfg))
            self.assertEqual(self.saver.client.expiry[self.key], 30)
            self.saver.client.now = 30
            self.assertIsNone(self.saver.get_tuple(self.cfg))
            self.put('002')
            self.saver.delete_thread('synthetic-reuse')
            self.assertIsNone(self.saver.get_tuple(self.cfg))

    def test_small_body_clears_large_cached_body_and_close_cleans_scope(self):
        with self.saver.reuse_written_checkpoint('synthetic-reuse'):
            self.put()
            self.put('002', text='很短')
            _, counts = measured_read(self.saver, self.cfg)
            self.assertEqual(total(counts, 'redis_checkpoint_cache_hit'), 0)
            self.assertIsNone(self.saver._reuse_scopes[self.key][0].checkpoint)
            self.saver.close()
            self.assertEqual(self.saver._reuse_scopes, {})

    def test_worker_thread_writes_visible_and_returned_state_is_not_shared(self):
        with self.saver.reuse_written_checkpoint('synthetic-reuse'):
            with ThreadPoolExecutor(1) as pool:
                pool.submit(self.put).result()
            first, counts = measured_read(self.saver, self.cfg)
            self.assertEqual(total(counts, 'redis_checkpoint_cache_hit'), 1)
            first.checkpoint['channel_values']['response'] = '客户端修改'
            second = self.saver.get_tuple(self.cfg)
            self.assertNotEqual(second.checkpoint['channel_values']['response'], '客户端修改')

    def test_intermediate_writes_do_not_hash_bodies_until_read(self):
        with self.saver.reuse_written_checkpoint('synthetic-reuse'):
            with patch('src.step3_question_answer.context.hashlib.sha1', side_effect=AssertionError('unnecessary-hash')):
                for cid in ('001', '002', '003'):
                    self.put(cid)
            self.assertEqual(self.saver.get_tuple(self.cfg).checkpoint['id'], '003')

    def test_parallel_sessions_overlap_fallback_and_two_agents_are_isolated(self):
        barrier = threading.Barrier(2)
        def request(thread):
            with self.saver.reuse_written_checkpoint(thread):
                self.put(cfg=config(thread), text=thread * 30000)
                barrier.wait(5)
                return measured_read(self.saver, config(thread))
        with ThreadPoolExecutor(2) as pool:
            result = list(pool.map(request, ('AAA', 'BBB')))
        for thread, (state, counts) in zip(('AAA', 'BBB'), result):
            self.assertEqual(state.checkpoint['channel_values']['recommendation']['sources'][0]['content'], thread * 30000)
            self.assertEqual(total(counts, 'redis_checkpoint_cache_hit'), 1)
        with self.saver.reuse_written_checkpoint('synthetic-reuse'):
            self.put()
            with self.saver.reuse_written_checkpoint('synthetic-reuse'):
                self.put('002')
                self.assertEqual(total(measured_read(self.saver, self.cfg)[1], 'redis_checkpoint_cache_hit'), 0)
            self.put('003')
            self.assertEqual(total(measured_read(self.saver, self.cfg)[1], 'redis_checkpoint_cache_hit'), 0)
        from src.step4_api.talent import TalentContextSaver
        talent = TalentContextSaver(ContextConfig('redis://localhost/0', 30), self.saver.client)
        with talent.reuse_written_checkpoint('synthetic-reuse'):
            talent.put(self.cfg, checkpoint('900'), {}, {})
            self.assertEqual(talent.get_tuple(self.cfg).checkpoint['id'], '900')
            self.assertEqual(self.saver.get_tuple(self.cfg).checkpoint['id'], '003')
        self.assertEqual(self.saver._reuse_scopes, {})
        self.assertEqual(talent._reuse_scopes, {})

    def test_invalid_conditional_responses_do_not_hide_bad_state(self):
        with self.saver.reuse_written_checkpoint('synthetic-reuse'):
            self.put()
            for response in ([1, 'bad', 1], [1, self.saver._field('c', '', '001'), None],
                             [1, self.saver._field('c', '', '001'), True]):
                with patch.object(self.saver.client, 'eval', return_value=response):
                    with self.assertRaisesRegex(ContextError, '^context_data_invalid$'):
                        self.saver.get_tuple(self.cfg)
        with patch.object(self.saver.client, 'eval', return_value=[1, self.saver._field('c', '', '001'), 1]):
            with self.assertRaisesRegex(ContextError, '^context_data_invalid$'):
                self.saver.get_tuple(self.cfg)

    def test_bytes_transport_and_pending_updates_during_read(self):
        with self.saver.reuse_written_checkpoint('synthetic-reuse'):
            saved = self.put()
            read = self.saver.client.eval
            def server(script, *args):
                if script == READ_CHECKPOINT:
                    self.saver.put_writes(saved, [('__interrupt__', {'value': '新补问'})], 'task')
                result = read(script, *args)
                return ([v.encode() if isinstance(v, str) else v for v in result]
                        if isinstance(result, list) else result)
            with patch.object(self.saver.client, 'eval', side_effect=server):
                state, counts = measured_read(self.saver, self.cfg)
            self.assertEqual(total(counts, 'redis_checkpoint_cache_hit'), 1)
            self.assertIn(('task', '__interrupt__', {'value': '新补问'}), state.pending_writes)


class ReuseGraphTests(unittest.TestCase):
    def test_runtime_failure_retry_and_stopped_requests_discard_cached_body(self):
        from types import SimpleNamespace
        saver = fake_saver()
        fail = True
        def plan(*_):
            if fail:
                raise RuntimeError('synthetic-plan-failure')
            return {'queries': [], 'error': None}
        def runtime():
            return GraphRuntime(build_graph(None,
                lambda _: {'conversation_kind': 'policy_overview', 'needs_company': False},
                checkpointer=saver, retriever=SimpleNamespace(search=lambda *_: {'candidates': []}),
                recommender=SimpleNamespace(plan=plan, answer=lambda *_: {'overview': '合成答复', 'items': [], 'sources': []})))
        question = '合成问题' * 12000
        with self.assertRaisesRegex(RuntimeError, 'synthetic-plan-failure'):
            runtime().run('synthetic-reuse', question, lambda *_: None, threading.Event())
        self.assertEqual(saver._reuse_scopes, {})
        # Restore via a fresh graph as the existing /retry contract requires.
        fail = False
        rows = []
        runtime().run('synthetic-reuse', '/retry', lambda e, d: rows.append((e, d)), threading.Event())
        self.assertEqual(rows[-1], ('done', {'status': 'completed'}))
        self.assertEqual(saver._reuse_scopes, {})
        stopped = threading.Event()
        def stop(event, data):
            if event == 'progress' and data.get('stage') == 'analyze' and data.get('status') == 'completed':
                stopped.set()
        runtime().run('synthetic-stop', question, stop, stopped)
        self.assertTrue(stopped.is_set())
        self.assertEqual(saver._reuse_scopes, {})

    def test_interrupt_resume_new_process_and_complete_state(self):
        saver = fake_saver()
        graph = build_graph(None, checkpointer=saver)
        question = '合成人才培训需求' * 9000
        with saver.reuse_written_checkpoint('synthetic-reuse'):
            graph.invoke({'question': question}, config())
            state = graph.get_state(config())
            self.assertTrue(state.interrupts)
            before = saver.get_tuple(config())
        restored_saver = RedisContextSaver(ContextConfig('redis://localhost/0', 30), saver.client)
        restored = build_graph(None, checkpointer=restored_saver)
        self.assertEqual(before, restored_saver.get_tuple(config()))
        with restored_saver.reuse_written_checkpoint('synthetic-reuse'):
            restored.invoke(Command(resume={'needs_company': False}), config())
            state = restored.get_state(config())
            self.assertFalse(state.next)
            self.assertFalse(state.interrupts)
            self.assertEqual(state.values['original_question'], question)
            cached = restored_saver.get_tuple(config())
        self.assertEqual(cached, restored_saver.get_tuple(config()))

    def test_api_full_recommendation_matches_unoptimized_final_result(self):
        from types import SimpleNamespace
        answer = checkpoint()['channel_values']['recommendation']
        answer.update(items=[{'source_id': f'synthetic-card-{i}', 'title': f'合成政策{i}',
                              'card': {'kind': 'policy', 'conditions': ['合成地区', '不含例外主体']}}
                             for i in (1, 2)], overview='完整合成结果', complete=True)
        def runtime(saver):
            return GraphRuntime(build_graph(None,
                lambda _: {'conversation_kind': 'policy_overview', 'needs_company': False},
                checkpointer=saver,
                retriever=SimpleNamespace(search=lambda *_: {'candidates': []}),
                recommender=SimpleNamespace(plan=lambda *_: {'queries': [], 'error': None}, answer=lambda *_: answer)))
        results = []
        for optimized in (False, True):
            saver = fake_saver()
            rows, timing = [], Timings()
            # Disable only the optimization threshold for the control, keeping
            # the same real graph, serialization and synthetic Redis adapter.
            with patch('src.step3_question_answer.context.MIN_REUSE_BYTES', 65536 if optimized else 10**12), bind(timing):
                runtime(saver).run('synthetic-reuse', '合成政策需求', lambda e, d: rows.append((e, d)), threading.Event())
            results.append(next(data for event, data in rows if event == 'result'))
            self.assertEqual(rows[-1], ('done', {'status': 'completed'}))
            self.assertEqual(saver._reuse_scopes, {})
            counts = timing.snapshot()['measurements']
            self.assertEqual(total(counts, 'redis_checkpoint_cache_hit'), int(optimized))
        self.assertEqual(results[0], results[1])
        self.assertEqual(results[1]['recommendation']['sources'], answer['sources'])


if __name__ == '__main__':
    unittest.main()