"""§8.8 Phase 2 折叠逻辑单测:切点选择 / 内存-prepared 映射 / Session.apply_fold / maybe_fold 编排(DB 与 LLM 全打桩,不碰真实 PG)。""" import unittest from types import SimpleNamespace from unittest.mock import patch from uuid import uuid4 import core.context_fold as cf from core.context_fold import ( FOLD_TAIL_RATIO, MIN_FOLD_MESSAGES, SUMMARY_NOTE_HEADER, _kth_user_index, _user_ordinal, build_summary_note, find_cutoff, ) from core.session import Session def _mk_window(n_turns: int, content_chars: int = 200) -> list: """构造 n_turns 轮 user/assistant 对话(等长消息,方便按条数推算体量占比)。""" msgs = [] for i in range(n_turns): msgs.append({"role": "user", "content": f"u{i:03d}" + "x" * content_chars}) msgs.append({"role": "assistant", "content": f"a{i:03d}" + "y" * content_chars}) return msgs class FindCutoffTests(unittest.TestCase): def test_cutoff_is_user_message_and_tail_ratio_respected(self) -> None: messages = [{"role": "system", "content": "rules"}] + _mk_window(30) cutoff = find_cutoff(messages, head_len=1) self.assertIsNotNone(cutoff) self.assertEqual(messages[cutoff].get("role"), "user") # 尾部(切点起)体量 >= FOLD_TAIL_RATIO(等长消息按条数近似) tail = len(messages) - cutoff window = len(messages) - 1 self.assertGreaterEqual(tail / window, FOLD_TAIL_RATIO * 0.9) # 可折叠段不为空且够长 self.assertGreaterEqual(cutoff - 1, MIN_FOLD_MESSAGES) def test_too_few_messages_returns_none(self) -> None: messages = [{"role": "system", "content": "rules"}] + _mk_window(3) self.assertIsNone(find_cutoff(messages, head_len=1)) def test_empty_window_returns_none(self) -> None: self.assertIsNone(find_cutoff([{"role": "system", "content": "s"}], head_len=1)) def test_cutoff_never_splits_tool_pairing(self) -> None: # assistant.tool_calls 后面跟 tool 结果 —— 切点是 user 消息,天然不落在中间 messages = [{"role": "system", "content": "rules"}] for i in range(20): messages.append({"role": "user", "content": f"u{i}" + "x" * 200}) messages.append({ "role": "assistant", "content": None, "tool_calls": [{"id": f"tc{i}", "function": {"name": "shell"}}], }) messages.append({"role": "tool", "tool_call_id": f"tc{i}", "content": "r" * 200}) cutoff = find_cutoff(messages, head_len=1) self.assertIsNotNone(cutoff) self.assertEqual(messages[cutoff]["role"], "user") class UserOrdinalMappingTests(unittest.TestCase): def test_ordinal_roundtrip_across_prepared_list(self) -> None: """prepare 补占位 tool 消息后 index 平移,但「第 k 条 user」两边一致。""" messages = [ {"role": "system", "content": "s"}, {"role": "user", "content": "q1"}, # 悬空 tool_call:prepare 的 repair 会在其后补一条占位 tool 消息 {"role": "assistant", "content": None, "tool_calls": [{"id": "tc1", "function": {"name": "shell"}}]}, {"role": "user", "content": "q2"}, {"role": "assistant", "content": "a2"}, ] from core.context import prepare_messages_with_stats prepared, stats = prepare_messages_with_stats(messages) self.assertEqual(stats["repaired_tool_calls"], 1) k = _user_ordinal(messages, 3) # q2 是第 2 条 user self.assertEqual(k, 2) p_idx = _kth_user_index(prepared, k) self.assertEqual(prepared[p_idx]["content"], "q2") self.assertEqual(p_idx, 4) # 占位 tool 插入使 index +1 def test_kth_user_missing_returns_none(self) -> None: self.assertIsNone(_kth_user_index([{"role": "assistant", "content": "a"}], 1)) class ApplyFoldTests(unittest.TestCase): def test_apply_fold_rebuilds_head_and_advances_base(self) -> None: sess = Session(task_id=uuid4(), system_prompt="SYS") window = _mk_window(10) sess.messages.extend(window) # 模拟 load:base=0,head=1(system) cutoff = 11 # 内存 index,某条 user 消息 self.assertEqual(sess.messages[cutoff]["role"], "user") sess.apply_fold(cutoff, "SUMMARY-TEXT") # 窗口 = [system] + [摘要注入] + 尾段 self.assertEqual(sess.messages[0]["role"], "system") self.assertTrue(sess.messages[1]["content"].startswith(SUMMARY_NOTE_HEADER)) self.assertIn("SUMMARY-TEXT", sess.messages[1]["content"]) self.assertEqual(sess.messages[2]["content"], window[cutoff - 1]["content"]) # base 推进 = 折叠条数(cutoff - 旧 head 1);head 变 2(system + 摘要) self.assertEqual(sess.context_base, cutoff - 1) self.assertEqual(sess.context_head_len, 2) def test_second_fold_replaces_summary_note(self) -> None: sess = Session(task_id=uuid4(), system_prompt="SYS") sess.messages.extend(_mk_window(10)) sess.apply_fold(11, "FIRST") base_after_first = sess.context_base sess.messages.extend(_mk_window(10)) # head=2,再折 10 条(切到 index 12 的 user 消息) self.assertEqual(sess.messages[12]["role"], "user") sess.apply_fold(12, "SECOND") notes = [m for m in sess.messages if str(m.get("content", "")).startswith(SUMMARY_NOTE_HEADER)] self.assertEqual(len(notes), 1) # 旧摘要被折进新摘要,不叠加 self.assertIn("SECOND", notes[0]["content"]) self.assertEqual(sess.context_base, base_after_first + 10) self.assertEqual(sess.context_head_len, 2) def test_summary_note_not_counted_as_user_msg(self) -> None: sess = Session(task_id=uuid4(), system_prompt="SYS") sess.messages.extend(_mk_window(10)) before = sess.n_user_msgs() sess.apply_fold(11, "S") # 折掉 5 轮 user(index 1..10 里 5 条 user),注入的摘要不计入 self.assertEqual(sess.n_user_msgs(), before - 5) class _FakeLLM: def __init__(self, reply: str) -> None: self.reply = reply self.last_messages = None def chat(self, messages, **kwargs): self.last_messages = messages return SimpleNamespace( choices=[SimpleNamespace(message=SimpleNamespace(content=self.reply))], usage=SimpleNamespace(prompt_tokens=100, completion_tokens=20), ) _FAKE_CAPS = SimpleNamespace( reliable_context=100, # 极小上限 → 体量必然超阈值,不用造几十万字符 family="deepseek_v4", variant="flash", input_cny_per_mtoken=1.0, output_cny_per_mtoken=2.0, ) class MaybeFoldTests(unittest.TestCase): def _session(self) -> Session: sess = Session(task_id=uuid4(), system_prompt="SYS") sess.messages.extend(_mk_window(30)) return sess def test_fold_happy_path(self) -> None: sess = self._session() llm = _FakeLLM("## 目标\n...") events = [] with patch.object(cf, "persist_fold") as persist, \ patch.object(cf, "record_chat_usage") as usage: stats = cf.maybe_fold(sess, llm, _FAKE_CAPS, user_id=uuid4(), emit=events.append) self.assertIsNotNone(stats) self.assertEqual(stats["phase"], "done") # 摘要指令是折叠调用的最后一条消息(前缀复用构造) self.assertEqual(llm.last_messages[-1]["content"], cf.FOLD_INSTRUCTION) # DB 持久化先行,参数与内存变更一致 persist.assert_called_once_with(sess.task_id, stats["new_base"], "## 目标\n...") usage.assert_called_once() # 独立 kind:折叠次数 = count(usage_events kind='context_fold'),task 详情接口依赖 self.assertEqual(usage.call_args.kwargs.get("kind"), "context_fold") self.assertEqual(sess.context_base, stats["new_base"]) phases = [e.get("phase") for e in events if e.get("type") == "context_fold"] self.assertEqual(phases, ["start", "done"]) self.assertLess(stats["after_chars"], stats["before_chars"]) def test_below_threshold_no_fold(self) -> None: sess = self._session() caps = SimpleNamespace(**{**_FAKE_CAPS.__dict__, "reliable_context": 10_000_000}) llm = _FakeLLM("unused") with patch.object(cf, "persist_fold") as persist: self.assertIsNone(cf.maybe_fold(sess, llm, caps, user_id=uuid4(), emit=lambda e: None)) persist.assert_not_called() self.assertIsNone(llm.last_messages) def test_empty_summary_skips_fold(self) -> None: sess = self._session() before = list(sess.messages) events = [] with patch.object(cf, "persist_fold") as persist: result = cf.maybe_fold( sess, _FakeLLM(" "), _FAKE_CAPS, user_id=uuid4(), emit=events.append ) self.assertIsNone(result) persist.assert_not_called() self.assertEqual(sess.messages, before) # 内存零污染 self.assertTrue(any(e.get("type") == "warn" for e in events)) if __name__ == "__main__": unittest.main()