| 1 | import sys |
| 2 | from pathlib import Path |
| 3 | |
| 4 | import pytest |
| 5 | |
| 6 | PROJECT_ROOT = Path(__file__).resolve().parents[1] |
| 7 | if str(PROJECT_ROOT) not in sys.path: |
| 8 | sys.path.insert(0, str(PROJECT_ROOT)) |
| 9 | |
| 10 | from extensions.python.message_loop_prompts_before._90_organize_history_wait import ( |
| 11 | MAX_SYNC_COMPRESSION_PASSES, |
| 12 | OrganizeHistoryWait, |
| 13 | ) |
| 14 | |
| 15 | |
| 16 | class _StalledHistory: |
| 17 | def __init__(self): |
| 18 | self.compress_calls = 0 |
| 19 | |
| 20 | def is_over_limit(self): |
| 21 | return self.compress_calls < 2 |
| 22 | |
| 23 | def get_tokens(self): |
| 24 | return 1234 |
| 25 | |
| 26 | async def compress(self): |
| 27 | self.compress_calls += 1 |
| 28 | return False |
| 29 | |
| 30 | |
| 31 | class _MaxPassHistory: |
| 32 | def __init__(self): |
| 33 | self.compress_calls = 0 |
| 34 | self.tokens = 2000 |
| 35 | |
| 36 | def is_over_limit(self): |
| 37 | return True |
| 38 | |
| 39 | def get_tokens(self): |
| 40 | return self.tokens |
| 41 | |
| 42 | async def compress(self): |
| 43 | self.compress_calls += 1 |
| 44 | self.tokens -= 1 |
| 45 | return True |
| 46 | |
| 47 | |
| 48 | class _CompressOnceHistory: |
| 49 | def __init__(self): |
| 50 | self.compress_calls = 0 |
| 51 | self.tokens = 2000 |
| 52 | |
| 53 | def is_over_limit(self): |
| 54 | return self.compress_calls == 0 |
| 55 | |
| 56 | def get_tokens(self): |
| 57 | return self.tokens |
| 58 | |
| 59 | async def compress(self): |
| 60 | self.compress_calls += 1 |
| 61 | self.tokens -= 1000 |
| 62 | return True |
| 63 | |
| 64 | |
| 65 | class _FakeLog: |
| 66 | def __init__(self): |
| 67 | self.entries = [] |
| 68 | |
| 69 | def set_progress(self, *args, **kwargs): |
| 70 | pass |
| 71 | |
| 72 | def log(self, **kwargs): |
| 73 | self.entries.append(kwargs) |
| 74 | |
| 75 | |
| 76 | class _FakeAgent: |
| 77 | def __init__(self, history=None): |
| 78 | self.data = {} |
| 79 | self.history = history or _StalledHistory() |
| 80 | self.context = type("Context", (), {"log": _FakeLog()})() |
| 81 | |
| 82 | def get_data(self, key): |
| 83 | return self.data.get(key) |
| 84 | |
| 85 | def set_data(self, key, value): |
| 86 | self.data[key] = value |
| 87 | |
| 88 | |
| 89 | @pytest.mark.asyncio |
| 90 | async def test_history_wait_stops_when_compression_makes_no_progress(): |
| 91 | agent = _FakeAgent() |
| 92 | |
| 93 | await OrganizeHistoryWait(agent).execute() |
| 94 | |
| 95 | assert agent.history.compress_calls == 1 |
| 96 | assert agent.context.log.entries |
| 97 | assert agent.context.log.entries[-1]["heading"] == "History compression stalled" |
| 98 | |
| 99 | |
| 100 | @pytest.mark.asyncio |
| 101 | async def test_history_wait_stops_after_max_sync_compression_passes(): |
| 102 | history = _MaxPassHistory() |
| 103 | agent = _FakeAgent(history) |
| 104 | |
| 105 | await OrganizeHistoryWait(agent).execute() |
| 106 | |
| 107 | assert history.compress_calls == MAX_SYNC_COMPRESSION_PASSES |
| 108 | assert agent.context.log.entries |
| 109 | assert agent.context.log.entries[-1]["heading"] == "History compression stalled" |
| 110 | assert ( |
| 111 | f"stopped after {MAX_SYNC_COMPRESSION_PASSES} passes" |
| 112 | in agent.context.log.entries[-1]["content"] |
| 113 | ) |
| 114 | |
| 115 | |
| 116 | @pytest.mark.asyncio |
| 117 | async def test_history_compression_clears_active_responses_state(): |
| 118 | agent = _FakeAgent(_CompressOnceHistory()) |
| 119 | agent.data["responses_state"] = { |
| 120 | "response_id": "resp_current", |
| 121 | "previous_response_id": "resp_previous", |
| 122 | "response_ids": ["resp_previous", "resp_current"], |
| 123 | } |
| 124 | |
| 125 | await OrganizeHistoryWait(agent).execute() |
| 126 | |
| 127 | state = agent.data["responses_state"] |
| 128 | assert "response_id" not in state |
| 129 | assert "previous_response_id" not in state |
| 130 | assert state["response_ids"] == ["resp_previous", "resp_current"] |