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'])