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

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

from redis.exceptions import ConnectionError
from src.step3_question_answer.context import ContextError, stored_bytes
from src.step3_question_answer.tests.test_context import fake_saver
from src.step3_question_answer.timing import Timings, bind, span
from src.step3_question_answer.graph import build_graph
from src.step4_api.runtime import GraphRuntime


class Clock:
    def __init__(self):
        self.now = 0
    def __call__(self):
        return self.now


class ContextTimingTests(unittest.TestCase):
    def test_directed_read_matches_legacy_history_and_does_not_mutate(self):
        saver, cfg, fields = self.seed()
        ns = '子图"\\namespace'
        fields[saver._field('c', ns, '999')] = saver._encode(({'id': '999'}, {}, None))
        fields[saver._field('w', ns, '999', 'task', -1)] = saver._encode(('task', '__error__', 'other'))
        fields[saver._field('w', '', '001', 'old', 0)] = saver._encode(('old', 'answer', 'old'))
        before = dict(fields)
        history = list(saver.list(cfg))
        with patch.object(saver.client, 'hgetall', side_effect=AssertionError('full read forbidden')):
            self.assertEqual(saver.get_tuple(cfg), next(item for item in history if item.checkpoint['id'] == '002'))
            for item in history:
                self.assertEqual(saver.get_tuple(item.config), item)
            other = {'configurable': dict(cfg['configurable'], checkpoint_ns=ns)}
            self.assertEqual(saver.get_tuple(other).pending_writes, [('task', '__error__', 'other')])
            for overrides in ({'checkpoint_id': 'absent'}, {'checkpoint_ns': 'absent'}):
                self.assertIsNone(saver.get_tuple({'configurable': dict(cfg['configurable'], **overrides)}))
        self.assertEqual(fields, before)
        self.assertEqual(saver.client.expiry[saver.key(cfg['configurable']['thread_id'])], 30)

    def seed(self, thread='PRIVATE_THREAD'):
        saver = fake_saver()
        cfg = {'configurable': {'thread_id': thread}}
        key = saver.key(thread)
        fields = {
            saver._field('c', '', '001'): saver._encode(({'id': '001', 'data': 'PRIVATE_HISTORY' * 1000}, {}, None)),
            saver._field('c', '', '002'): saver._encode(({'id': '002', 'data': 'PRIVATE_CURRENT'}, {}, '001')),
            saver._field('w', '', '002', 'task', 0): saver._encode(('task', 'channel', 'PRIVATE_PENDING')),
            saver._field('request', 'v1', 'r'): json.dumps({'question': 'PRIVATE_QUESTION'}),
        }
        saver.client.rows[key] = fields
        saver.client.expiry[key] = 30
        return saver, cfg, fields

    def test_io_decode_and_full_hash_vs_selected_bytes_are_distinct(self):
        saver, cfg, fields = self.seed()
        clock = Clock()
        timing = Timings(clock=clock)
        original_get = saver.client.eval
        original_load = saver.serde.loads_typed
        original_b64 = base64.b64decode
        def get(*args):
            clock.now += .2
            return original_get(*args)
        def load(value):
            clock.now += .03
            return original_load(value)
        def b64(*args, **kwargs):
            clock.now += .02
            return original_b64(*args, **kwargs)
        with patch.object(saver.client, 'eval', side_effect=get), \
             patch.object(saver.serde, 'loads_typed', side_effect=load), \
             patch('src.step3_question_answer.context.base64.b64decode', side_effect=b64), \
             bind(timing), span('session_result_load', 'session_result_load'):
            result = saver.get_tuple(cfg)
        report = timing.snapshot()
        spans, counts = report['spans'], report['measurements']
        suffix = '.session_result_load'
        self.assertEqual(spans['redis_checkpoint_read' + suffix]['total_ms'], 200)
        self.assertEqual(spans['checkpoint_deserialize' + suffix]['total_ms'], 60)
        self.assertEqual(spans['checkpoint_base64_decode' + suffix]['total_ms'], 40)
        self.assertEqual(spans['session_result_load']['total_ms'], 300)
        self.assertEqual(result.checkpoint['id'], '002')
        self.assertEqual(result.pending_writes, [('task', 'channel', 'PRIVATE_PENDING')])
        self.assertEqual(counts['redis_scanned_fields' + suffix]['total'], 4)
        self.assertEqual(counts['redis_returned_fields' + suffix]['total'], 2)
        self.assertEqual(counts['redis_checkpoint_fields' + suffix]['total'], 1)
        self.assertEqual(counts['redis_write_fields' + suffix]['total'], 1)
        self.assertEqual(counts['redis_request_fields' + suffix]['total'], 0)
        received = sum(stored_bytes(k) + stored_bytes(v) for k, v in fields.items() if json.loads(k)[:3] in (['c', '', '002'], ['w', '', '002']))
        self.assertEqual(counts['redis_returned_bytes' + suffix]['total'], received)
        self.assertEqual(counts['checkpoint_decoded_input_bytes' + suffix]['count'], 2)
        self.assertLess(received, sum(stored_bytes(k) + stored_bytes(v) for k, v in fields.items()) // 10)
        self.assertNotIn('PRIVATE', json.dumps(report))

    def test_connection_and_corrupt_value_failures_still_have_safe_measurements(self):
        saver, cfg, _ = self.seed()
        timing = Timings()
        with bind(timing, 'session_result_load'), patch.object(saver.client, 'eval', side_effect=ConnectionError('PRIVATE')):
            with self.assertRaisesRegex(ContextError, '^redis_connection_failed$'):
                saver.get_tuple(cfg)
        self.assertEqual(timing.snapshot()['spans']['redis_checkpoint_read.session_result_load']['failed_count'], 1)
        timing = Timings()
        saver.client.rows[saver.key(cfg['configurable']['thread_id'])][saver._field('c', '', '002')] = 'PRIVATE_INVALID_JSON'
        with bind(timing, 'session_result_load'), self.assertRaisesRegex(ContextError, '^context_data_invalid$'):
            saver.get_tuple(cfg)
        report = timing.snapshot()
        self.assertEqual(report['spans']['checkpoint_json_decode.session_result_load']['failed_count'], 1)
        self.assertNotIn('PRIVATE', json.dumps(report))

    def test_parallel_telemetry_is_isolated_and_empty_session_is_counted(self):
        barrier = threading.Barrier(2)
        def read(present):
            saver, cfg, _ = self.seed(str(present))
            if not present:
                saver.client.rows.clear()
            timing = Timings()
            with bind(timing, 'session_result_load'):
                barrier.wait(timeout=3)
                saver.get_tuple(cfg)
            return timing.snapshot()
        with ThreadPoolExecutor(2) as pool:
            reports = list(pool.map(read, [True, False]))
        self.assertEqual([r['measurements']['redis_returned_fields.session_result_load']['total'] for r in reports], [2, 0])
        self.assertNotIn('checkpoint_deserialize.session_result_load', reports[1]['spans'])

    def test_real_graph_final_read_is_attributed_separately_from_initial_read(self):
        saver = fake_saver()
        graph = build_graph(None, lambda _: {'conversation_kind': 'smalltalk', 'direct_response': '合成答复'}, checkpointer=saver)
        timing = Timings()
        events = []
        with bind(timing):
            GraphRuntime(graph).run('synthetic', '合成问题', lambda e, d: events.append((e, d)), threading.Event())
        result = timing.snapshot()
        self.assertEqual(events[-1], ('done', {'status': 'completed'}))
        self.assertIn('redis_checkpoint_read.session_load', result['spans'])
        self.assertIn('redis_checkpoint_read.session_result_load', result['spans'])
        self.assertEqual(result['measurements']['redis_checkpoint_fields.session_result_load']['total'], 1)
        self.assertEqual(result['measurements']['redis_returned_fields.session_load']['total'], 0)
        self.assertIn('checkpoint_deserialize.session_result_load', result['spans'])