test_trajectory_compressor_async.py
1 """Tests for trajectory_compressor AsyncOpenAI event loop binding. 2 3 The AsyncOpenAI client was created once at __init__ time and stored as an 4 instance attribute. When process_directory() calls asyncio.run() — which 5 creates and closes a fresh event loop — the client's internal httpx 6 transport remains bound to the now-closed loop. A second call to 7 process_directory() would fail with "Event loop is closed". 8 9 The fix creates the AsyncOpenAI client lazily via _get_async_client() so 10 each asyncio.run() gets a client bound to the current loop. 11 """ 12 13 import types 14 from types import SimpleNamespace 15 from unittest.mock import MagicMock, patch 16 17 import pytest 18 19 20 class TestAsyncClientLazyCreation: 21 """trajectory_compressor.py — _get_async_client()""" 22 23 def test_async_client_none_after_init(self): 24 """async_client should be None after __init__ (not eagerly created).""" 25 from trajectory_compressor import TrajectoryCompressor 26 27 comp = TrajectoryCompressor.__new__(TrajectoryCompressor) 28 comp.config = MagicMock() 29 comp.config.base_url = "https://api.example.com/v1" 30 comp.config.api_key_env = "TEST_API_KEY" 31 comp._use_call_llm = False 32 comp.async_client = None 33 comp._async_client_api_key = "test-key" 34 35 assert comp.async_client is None 36 37 def test_get_async_client_creates_new_client(self): 38 """_get_async_client() should create a fresh AsyncOpenAI instance.""" 39 from trajectory_compressor import TrajectoryCompressor 40 41 comp = TrajectoryCompressor.__new__(TrajectoryCompressor) 42 comp.config = MagicMock() 43 comp.config.base_url = "https://api.example.com/v1" 44 comp._async_client_api_key = "test-key" 45 comp.async_client = None 46 47 mock_async_openai = MagicMock() 48 with patch("openai.AsyncOpenAI", mock_async_openai): 49 client = comp._get_async_client() 50 51 mock_async_openai.assert_called_once_with( 52 api_key="test-key", 53 base_url="https://api.example.com/v1", 54 ) 55 assert comp.async_client is not None 56 57 def test_get_async_client_creates_fresh_each_call(self): 58 """Each call to _get_async_client() creates a NEW client instance, 59 so it binds to the current event loop.""" 60 from trajectory_compressor import TrajectoryCompressor 61 62 comp = TrajectoryCompressor.__new__(TrajectoryCompressor) 63 comp.config = MagicMock() 64 comp.config.base_url = "https://api.example.com/v1" 65 comp._async_client_api_key = "test-key" 66 comp.async_client = None 67 68 call_count = 0 69 instances = [] 70 71 def mock_constructor(**kwargs): 72 nonlocal call_count 73 call_count += 1 74 instance = MagicMock() 75 instances.append(instance) 76 return instance 77 78 with patch("openai.AsyncOpenAI", side_effect=mock_constructor): 79 client1 = comp._get_async_client() 80 client2 = comp._get_async_client() 81 82 # Should have created two separate instances 83 assert call_count == 2 84 assert instances[0] is not instances[1] 85 86 87 class TestSourceLineVerification: 88 """Verify the actual source has the lazy pattern applied.""" 89 90 @staticmethod 91 def _read_file() -> str: 92 import os 93 base = os.path.dirname(os.path.dirname(__file__)) 94 with open(os.path.join(base, "trajectory_compressor.py")) as f: 95 return f.read() 96 97 def test_no_eager_async_openai_in_init(self): 98 """__init__ should NOT create AsyncOpenAI eagerly.""" 99 src = self._read_file() 100 # The old pattern: self.async_client = AsyncOpenAI(...) in _init_summarizer 101 # should not exist — only self.async_client = None 102 lines = src.split("\n") 103 for i, line in enumerate(lines, 1): 104 if "self.async_client = AsyncOpenAI(" in line and "_get_async_client" not in lines[max(0,i-3):i+1]: 105 # Allow it inside _get_async_client method 106 # Check if we're inside _get_async_client by looking at context 107 context = "\n".join(lines[max(0,i-20):i+1]) 108 if "_get_async_client" not in context: 109 pytest.fail( 110 f"Line {i}: AsyncOpenAI created eagerly outside _get_async_client()" 111 ) 112 113 def test_get_async_client_method_exists(self): 114 """_get_async_client method should exist.""" 115 src = self._read_file() 116 assert "def _get_async_client(self)" in src 117 118 119 @pytest.mark.asyncio 120 async def test_generate_summary_async_kimi_omits_temperature(): 121 """Kimi models should have temperature omitted — server manages it.""" 122 from trajectory_compressor import CompressionConfig, TrajectoryCompressor, TrajectoryMetrics 123 124 config = CompressionConfig( 125 summarization_model="kimi-for-coding", 126 temperature=0.3, 127 summary_target_tokens=100, 128 max_retries=1, 129 ) 130 compressor = TrajectoryCompressor.__new__(TrajectoryCompressor) 131 compressor.config = config 132 compressor.logger = MagicMock() 133 compressor._use_call_llm = False 134 async_client = MagicMock() 135 async_client.chat.completions.create = MagicMock(return_value=SimpleNamespace( 136 choices=[SimpleNamespace(message=SimpleNamespace(content="[CONTEXT SUMMARY]: summary"))] 137 )) 138 compressor._get_async_client = MagicMock(return_value=async_client) 139 140 metrics = TrajectoryMetrics() 141 result = await compressor._generate_summary_async("tool output", metrics) 142 143 assert result.startswith("[CONTEXT SUMMARY]:") 144 assert "temperature" not in async_client.chat.completions.create.call_args.kwargs 145 146 147 @pytest.mark.asyncio 148 async def test_generate_summary_async_public_moonshot_kimi_k2_5_omits_temperature(): 149 """kimi-k2.5 on the public Moonshot API should not get a forced temperature.""" 150 from trajectory_compressor import CompressionConfig, TrajectoryCompressor, TrajectoryMetrics 151 152 config = CompressionConfig( 153 summarization_model="kimi-k2.5", 154 base_url="https://api.moonshot.ai/v1", 155 temperature=0.3, 156 summary_target_tokens=100, 157 max_retries=1, 158 ) 159 compressor = TrajectoryCompressor.__new__(TrajectoryCompressor) 160 compressor.config = config 161 compressor.logger = MagicMock() 162 compressor._use_call_llm = False 163 async_client = MagicMock() 164 async_client.chat.completions.create = MagicMock(return_value=SimpleNamespace( 165 choices=[SimpleNamespace(message=SimpleNamespace(content="[CONTEXT SUMMARY]: summary"))] 166 )) 167 compressor._get_async_client = MagicMock(return_value=async_client) 168 169 metrics = TrajectoryMetrics() 170 result = await compressor._generate_summary_async("tool output", metrics) 171 172 assert result.startswith("[CONTEXT SUMMARY]:") 173 assert "temperature" not in async_client.chat.completions.create.call_args.kwargs 174 175 176 @pytest.mark.asyncio 177 async def test_generate_summary_async_public_moonshot_cn_kimi_k2_5_omits_temperature(): 178 """kimi-k2.5 on api.moonshot.cn should not get a forced temperature.""" 179 from trajectory_compressor import CompressionConfig, TrajectoryCompressor, TrajectoryMetrics 180 181 config = CompressionConfig( 182 summarization_model="kimi-k2.5", 183 base_url="https://api.moonshot.cn/v1", 184 temperature=0.3, 185 summary_target_tokens=100, 186 max_retries=1, 187 ) 188 compressor = TrajectoryCompressor.__new__(TrajectoryCompressor) 189 compressor.config = config 190 compressor.logger = MagicMock() 191 compressor._use_call_llm = False 192 async_client = MagicMock() 193 async_client.chat.completions.create = MagicMock(return_value=SimpleNamespace( 194 choices=[SimpleNamespace(message=SimpleNamespace(content="[CONTEXT SUMMARY]: summary"))] 195 )) 196 compressor._get_async_client = MagicMock(return_value=async_client) 197 198 metrics = TrajectoryMetrics() 199 result = await compressor._generate_summary_async("tool output", metrics) 200 201 assert result.startswith("[CONTEXT SUMMARY]:") 202 assert "temperature" not in async_client.chat.completions.create.call_args.kwargs