business-assistant/src/step4_api/tests/test_session_archive.py

import json
import os
import threading
import unittest
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from tempfile import TemporaryDirectory
from types import SimpleNamespace
from unittest.mock import patch

from src.step3_question_answer.graph import build_graph
from src.step3_question_answer.session_archive import (
    SessionArchive, SessionArchiveError, invoke_recorded, load_session_dir, resolve_session_dir,
)
from src.step3_question_answer.tests.test_context import fake_saver
from src.step4_api.runtime import GraphRuntime


class SessionArchiveTests(unittest.TestCase):
    def setUp(self):
        self.temp = TemporaryDirectory()
        self.addCleanup(self.temp.cleanup)
        self.root = Path(self.temp.name)
        self.archive = SessionArchive(self.root / '中文 会话')

    def read(self, thread='thread'):
        return json.loads(self.archive.path_for(thread).read_text(encoding='utf-8'))

    def run_turn(self, runtime, question, request_id, thread='thread', stopped=None):
        events = []
        def emit(event, data):
            if event == 'done':
                self.assertEqual(self.read(thread)['turns'][-1]['status'], data['status'])
            events.append((event, data))
        runtime.run(thread, question, emit, stopped or threading.Event(), request_id=request_id)
        return events

    def test_paths_config_quotes_env_override_and_foreign_absolute(self):
        config = self.root / 'synthetic.env'
        config.write_text('SESSION_DIR="资料\\会话 记录" # comment\nUNRELATED=unused\n', encoding='utf-8')
        self.assertEqual(load_session_dir(config, {}, self.root), self.root / '资料/会话 记录')
        self.assertEqual(load_session_dir(config, {'SESSION_DIR': 'other/session'}, self.root), self.root / 'other/session')
        self.assertEqual(resolve_session_dir(str(self.root / 'absolute'), self.root), self.root / 'absolute')
        self.assertEqual(load_session_dir(self.root / 'missing', {}, self.root), self.root / 'session')
        for value in ('', 'C:relative', '/var/lib/session' if os.name == 'nt' else 'C:\\session'):
            with self.subTest(value=value), self.assertRaises(SessionArchiveError):
                resolve_session_dir(value, self.root)
        if os.name == 'nt':
            self.assertEqual(resolve_session_dir('C:/archive/session'), Path('C:/archive/session'))
            self.assertEqual(resolve_session_dir(r'\\server\share\session'), Path(r'\\server\share\session'))
        else:
            self.assertEqual(resolve_session_dir('/var/lib/session'), Path('/var/lib/session'))

    def test_restart_clarifications_and_isolation_no_truncation(self):
        saver = fake_saver()
        runtime = GraphRuntime(build_graph(None, checkpointer=saver), self.archive)
        self.run_turn(runtime, '合成需求', 'r1')
        runtime = GraphRuntime(build_graph(None, checkpointer=saver), SessionArchive(self.archive.directory))
        self.run_turn(runtime, '不需要', 'r2')
        answer = '中文回答'
        runtime = GraphRuntime(build_graph(None, lambda _: {'conversation_kind': 'smalltalk', 'direct_response': answer},
                                          checkpointer=saver), self.archive)
        self.run_turn(runtime, '新问题', 'r3')
        self.run_turn(runtime, '其他窗口', 'r4', thread='../CON/另一个')
        doc = self.read()
        self.assertEqual(doc['thread_id'], 'thread')
        self.assertEqual([r['question'] for r in doc['turns']], ['合成需求', '不需要', '新问题'])
        self.assertEqual(doc['turns'][0]['status'], 'needs_input')
        self.assertTrue(doc['turns'][0]['ai_answer'])
        self.assertEqual(doc['turns'][0]['events'][-1]['event'], 'interrupt')
        self.assertEqual(doc['turns'][2]['ai_answer'], answer)
        self.assertEqual(doc['turns'][2]['result']['response'], answer)
        self.assertEqual(len(self.read('../CON/另一个')['turns']), 1)
        self.assertEqual(self.archive.path_for('../CON/另一个').parent, self.archive.directory)
        saver.delete_thread('thread')
        self.assertEqual(len(self.read()['turns']), 3)

    def test_full_structured_answer_snapshot_and_error(self):
        class Runtime(GraphRuntime):
            def _run(self, thread_id, question, emit, stopped, record):
                emit('answer', {'text': '开场'})
                emit('summary', {'text': '**结论**\n下一步' * 5000, 'source_ids': ['s1']})
                emit('result', {'response': '开场', 'recommendation': {
                    'items': [{'id': 'i1', 'title': '合成卡片'}],
                    'sources': [{'id': 's1', 'text': '合成原文'}]}})
                emit('done', {'status': 'completed'})
        runtime = Runtime(SimpleNamespace(checkpointer=None), self.archive)
        events = self.run_turn(runtime, '问题', 'r1')
        row = self.read()['turns'][0]
        self.assertEqual(row['ai_answer'], '开场\n\n' + '**结论**\n下一步' * 5000)
        self.assertEqual(row['result'], dict(events)['result'])
        def fail(*args):
            raise RuntimeError('PRIVATE_EXCEPTION')
        with patch.object(runtime, '_run', side_effect=fail), self.assertRaises(RuntimeError):
            self.run_turn(runtime, '失败问题', 'r2')
        row = self.read()['turns'][1]
        self.assertEqual(row['status'], 'failed')
        self.assertTrue(row['ai_answer'])
        self.assertNotIn('PRIVATE_EXCEPTION', json.dumps(row))

    def test_stop_and_write_failure_never_announce_done(self):
        runtime = GraphRuntime(build_graph(None), self.archive)
        stop = threading.Event()
        stop.set()
        self.run_turn(runtime, '停止问题', 'r1', stopped=stop)
        self.assertEqual(self.read()['turns'][0]['status'], 'stopped')
        runtime = GraphRuntime(build_graph(None, lambda _: {'conversation_kind': 'smalltalk', 'direct_response': '你好'}), self.archive)
        put = self.archive.put
        def fail_final(record):
            if record['finished_at']:
                raise SessionArchiveError('session_archive_write_failed')
            put(record)
        events = []
        with patch.object(self.archive, 'put', side_effect=fail_final), self.assertRaises(SessionArchiveError):
            runtime.run('thread', '你好', lambda e, d: events.append(e), threading.Event(), request_id='r2')
        self.assertNotIn('done', events)
        self.assertEqual(self.read()['turns'][1]['status'], 'running')

    def test_atomic_failure_corruption_and_idempotent_concurrency(self):
        def write(index):
            archive = SessionArchive(self.archive.directory)
            archive.put({'thread_id': 'thread', 'request_id': str(index), 'question': '合成'})
        with ThreadPoolExecutor(max_workers=8) as pool:
            list(pool.map(write, range(20)))
        write(0)
        self.assertEqual(len(self.read()['turns']), 20)
        before = self.archive.path_for('thread').read_bytes()
        with patch('src.step3_question_answer.session_archive.os.replace', side_effect=OSError), self.assertRaises(SessionArchiveError):
            write(21)
        self.assertEqual(self.archive.path_for('thread').read_bytes(), before)
        self.assertEqual(list(self.archive.directory.glob('*.tmp')), [])
        path = self.archive.path_for('thread')
        path.write_text('{broken', encoding='utf-8')
        with self.assertRaises(SessionArchiveError):
            write(22)
        self.assertEqual(path.read_text(), '{broken')

    def test_readable_filename_and_legacy_migration(self):
        import hashlib
        legacy = self.archive.directory / (hashlib.sha256(b'thread').hexdigest() + '.json')
        legacy.write_text(json.dumps({'schema_version': 1, 'thread_id': 'thread', 'turns': [
            {'request_id': 'old', 'thread_id': 'thread', 'question': 'old'}]}), encoding='utf-8')
        self.archive.put({'request_id': 'new', 'thread_id': 'thread', 'question': 'new'})
        self.assertEqual(self.archive.path_for('thread').name, 'thread.json')
        self.assertFalse(legacy.exists())
        self.assertEqual([r['request_id'] for r in self.read()['turns']], ['old', 'new'])
        for identity in ('CON', '../bad', 'thread/unsafe'):
            self.assertTrue(self.archive.path_for(identity).name.startswith('~'))
        self.archive.put({'request_id': 'upper', 'thread_id': 'THREAD', 'question': 'upper'})
        self.assertEqual(self.read()['thread_id'], 'thread')
        self.assertEqual(self.read('THREAD')['thread_id'], 'THREAD')

    def test_cli_records_question_and_clarification(self):
        from langgraph.types import Command
        graph = build_graph(None)
        config = {'configurable': {'thread_id': 'thread'}}
        invoke_recorded(graph, {'question': '合成问题'}, config, self.archive, '合成问题')
        invoke_recorded(graph, Command(resume={'needs_company': False, 'company_hint': ''}),
                        config, self.archive, '', 'clarification_reply')
        rows = self.read()['turns']
        self.assertEqual([r['status'] for r in rows], ['needs_input', 'completed'])
        self.assertEqual(rows[1]['question'], '')
        self.assertTrue(all(r['ai_answer'] for r in rows))

    def test_http_sse_runtime_uses_dms_observer_without_local_archive(self):
        from starlette.testclient import TestClient
        from src.step4_api.app import create_app
        from src.step4_api.config import Settings
        from src.step4_api.runtime import Runtime
        from src.step3_question_answer.context import RedisContextSaver, ContextConfig
        from src.step3_question_answer.tests.test_context import FakeRedis
        class Saver(RedisContextSaver):
            def __init__(self, config):
                client = FakeRedis()
                client.ping = lambda: True
                super().__init__(ContextConfig('redis://localhost/0', 30), client)
        with patch('src.step4_api.runtime.RedisContextSaver', Saver), \
                patch('src.step4_api.runtime.TalentRuntime'), \
                patch('src.step4_api.runtime.load_context_config'), \
                patch('src.step4_api.runtime.Retriever'), \
                patch('src.step4_api.runtime.Recommender'), \
                patch('src.step4_api.runtime.JsonModel'), \
                patch('src.step4_api.runtime.deepseek_analyzer', return_value=lambda _: {
                    'conversation_kind': 'smalltalk', 'direct_response': '合成完整回答'}), \
                patch('src.step4_api.runtime.create_dms_observer') as observer_factory, \
                patch('src.step3_question_answer.session_archive.SessionArchive') as archive_factory, \
                TestClient(create_app(Settings(), Runtime)) as client:
            response = client.post('/api/chat', json={'thread_id': 'thread', 'question': '你好'})
        self.assertEqual(response.status_code, 200)
        self.assertIn('合成完整回答', response.text)
        self.assertFalse(self.archive.path_for('thread').exists())
        observer_factory.return_value.start.assert_called_once()
        observer_factory.return_value.close.assert_called_once()
        archive_factory.assert_not_called()


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