solpbc.org
personal memory agent
1# SPDX-License-Identifier: AGPL-3.0-only
2# Copyright (c) 2026 sol pbc
3
4"""Tests for the Batch async batch processor."""
5
6import asyncio
7from unittest.mock import AsyncMock, patch
8
9import pytest
10
11from solstone.think.batch import Batch, BatchRequest
12from solstone.think.models import GEMINI_FLASH, SchemaValidationError
13
14
15def _result(text: str = "Response", finish_reason: str = "stop", **extra):
16 return {"text": text, "finish_reason": finish_reason, **extra}
17
18
19def test_batch_request_creation():
20 """Test BatchRequest can be created with required and custom params."""
21 # Required params only
22 req = BatchRequest(contents="Test prompt", context="test.context")
23 assert req.contents == "Test prompt"
24 assert req.context == "test.context"
25 assert req.temperature == 0.3
26 assert req.response is None
27 assert req.error is None
28
29 # With generation options
30 req2 = BatchRequest(
31 contents=["Part 1", "Part 2"],
32 context="test.context",
33 temperature=0.7,
34 json_output=True,
35 )
36 assert req2.contents == ["Part 1", "Part 2"]
37 assert req2.temperature == 0.7
38 assert req2.json_output is True
39
40
41def test_batch_request_custom_attributes():
42 """Test that arbitrary attributes can be added to BatchRequest."""
43 req = BatchRequest(contents="Test", context="test.context")
44 req.frame_id = 123
45 req.stage = "initial"
46 req.custom_data = {"foo": "bar"}
47
48 assert req.frame_id == 123
49 assert req.stage == "initial"
50 assert req.custom_data == {"foo": "bar"}
51
52
53@pytest.mark.asyncio
54@patch("solstone.think.batch.agenerate_with_result", new_callable=AsyncMock)
55async def test_batch_basic(mock_agenerate):
56 """Test basic batch execution with single request."""
57 mock_agenerate.return_value = _result("Response 1")
58
59 # Create batch and add request
60 batch = Batch(max_concurrent=5)
61 req = batch.create(contents="What is 2+2?", context="test.calc")
62 req.task_id = "calc1"
63 batch.add(req)
64
65 # Iterate and verify
66 results = []
67 async for completed_req in batch.drain_batch():
68 results.append(completed_req)
69
70 assert len(results) == 1
71 assert results[0].task_id == "calc1"
72 assert results[0].response == "Response 1"
73 assert results[0].error is None
74 assert results[0].duration > 0
75
76 # Verify API was called correctly
77 mock_agenerate.assert_called_once()
78 call_kwargs = mock_agenerate.call_args[1]
79 assert call_kwargs["contents"] == "What is 2+2?"
80 assert call_kwargs["context"] == "test.calc"
81
82
83@pytest.mark.asyncio
84@patch("solstone.think.batch.agenerate_with_result", new_callable=AsyncMock)
85async def test_batch_records_resolved_model(mock_agenerate):
86 """The transport-reported model is recorded without a request override."""
87 mock_agenerate.return_value = _result("Response", model=GEMINI_FLASH)
88
89 batch = Batch(max_concurrent=5)
90 req = batch.create(contents="Test", context="test.context")
91 batch.add(req)
92
93 results = []
94 async for completed_req in batch.drain_batch():
95 results.append(completed_req)
96
97 assert len(results) == 1
98 assert results[0].model_used == GEMINI_FLASH
99
100 call_kwargs = mock_agenerate.call_args[1]
101 assert "model" not in call_kwargs
102
103
104@pytest.mark.asyncio
105@patch("solstone.think.batch.agenerate_with_result", new_callable=AsyncMock)
106async def test_batch_multiple_requests(mock_agenerate):
107 """Test batch with multiple requests."""
108 mock_agenerate.side_effect = [
109 _result("Response 1"),
110 _result("Response 2"),
111 _result("Response 3"),
112 ]
113
114 batch = Batch(max_concurrent=2)
115
116 # Add multiple requests
117 req1 = batch.create(contents="Prompt 1", context="test.context")
118 req1.id = 1
119 batch.add(req1)
120
121 req2 = batch.create(contents="Prompt 2", context="test.context")
122 req2.id = 2
123 batch.add(req2)
124
125 req3 = batch.create(contents="Prompt 3", context="test.context")
126 req3.id = 3
127 batch.add(req3)
128
129 # Collect results
130 results = []
131 async for req in batch.drain_batch():
132 results.append(req)
133
134 # Should have all 3 results
135 assert len(results) == 3
136
137 # Results may come in any order (concurrent execution)
138 result_ids = {r.id for r in results}
139 assert result_ids == {1, 2, 3}
140
141 # All should have responses
142 for r in results:
143 assert r.response is not None
144 assert r.error is None
145
146
147@pytest.mark.asyncio
148@patch("solstone.think.batch.agenerate_with_result", new_callable=AsyncMock)
149async def test_batch_error_handling(mock_agenerate):
150 """Test that errors are captured in request.error."""
151 mock_agenerate.side_effect = ValueError("API error")
152
153 batch = Batch(max_concurrent=5)
154 req = batch.create(contents="Bad prompt", context="test.context")
155 req.id = "error_test"
156 batch.add(req)
157
158 # Get result
159 results = []
160 async for r in batch.drain_batch():
161 results.append(r)
162
163 assert len(results) == 1
164 assert results[0].id == "error_test"
165 assert results[0].response is None
166 assert results[0].error == "API error"
167 assert results[0].duration > 0
168
169
170@pytest.mark.asyncio
171@patch("solstone.think.batch.resolve_provider")
172@patch("solstone.think.batch.agenerate_with_result", new_callable=AsyncMock)
173async def test_batch_preserves_exception_metadata(
174 mock_agenerate, mock_resolve_provider
175):
176 """Structured provider metadata should survive the batch error path."""
177
178 class ProviderError(RuntimeError):
179 reason_code = "provider_key_invalid"
180 reset_at_ms = 12345
181 provider = "anthropic"
182
183 mock_agenerate.side_effect = ProviderError("bad key")
184 mock_resolve_provider.return_value = ("google", "unused")
185
186 batch = Batch(max_concurrent=5)
187 req = batch.create(contents="Bad prompt", context="test.context")
188 batch.add(req)
189
190 results = []
191 async for r in batch.drain_batch():
192 results.append(r)
193
194 assert len(results) == 1
195 assert results[0].error == "bad key"
196 assert results[0].reason_code == "provider_key_invalid"
197 assert results[0].reset_at_ms == 12345
198 assert results[0].provider == "anthropic"
199 mock_resolve_provider.assert_not_called()
200
201
202@pytest.mark.asyncio
203@patch("solstone.think.batch.resolve_provider")
204@patch("solstone.think.batch.agenerate_with_result", new_callable=AsyncMock)
205async def test_batch_classifies_exception_metadata_when_attrs_missing(
206 mock_agenerate, mock_resolve_provider
207):
208 """Batch should classify generic exceptions and best-effort provider."""
209 mock_agenerate.side_effect = ConnectionError("network down")
210 mock_resolve_provider.return_value = ("google", "gemini-test")
211
212 batch = Batch(max_concurrent=5)
213 req = batch.create(contents="Bad prompt", context="test.context")
214 batch.add(req)
215
216 results = []
217 async for r in batch.drain_batch():
218 results.append(r)
219
220 assert len(results) == 1
221 assert results[0].reason_code == "network_unreachable"
222 assert results[0].reset_at_ms is None
223 assert results[0].provider == "google"
224 mock_resolve_provider.assert_called_once_with("generate")
225
226
227def test_batch_update_clears_error_metadata():
228 batch = Batch(max_concurrent=5)
229 req = batch.create(contents="Retry prompt", context="test.context")
230 req.reason_code = "provider_key_invalid"
231 req.provider = "anthropic"
232 req.reset_at_ms = 12345
233 req.error = "bad key"
234
235 with patch.object(batch, "add") as mock_add:
236 batch.update(req)
237
238 assert req.reason_code is None
239 assert req.provider is None
240 assert req.reset_at_ms is None
241 assert req.error is None
242 mock_add.assert_called_once_with(req)
243
244
245@pytest.mark.asyncio
246@patch("solstone.think.batch.agenerate_with_result", new_callable=AsyncMock)
247async def test_batch_dynamic_adding(mock_agenerate):
248 """Test adding requests dynamically during iteration."""
249 mock_agenerate.return_value = _result("Response")
250
251 batch = Batch(max_concurrent=5)
252
253 # Add initial request
254 req1 = batch.create(contents="Initial", context="test.context")
255 req1.stage = "initial"
256 batch.add(req1)
257
258 # Process and add more during iteration
259 results = []
260 added_followup = False
261
262 async for req in batch.drain_batch():
263 results.append(req)
264
265 # After first result, add a follow-up
266 if req.stage == "initial" and not added_followup:
267 req2 = batch.create(contents="Follow-up", context="test.context")
268 req2.stage = "followup"
269 batch.add(req2)
270 added_followup = True
271
272 # Should have both results
273 assert len(results) == 2
274 stages = {r.stage for r in results}
275 assert stages == {"initial", "followup"}
276
277
278@pytest.mark.asyncio
279@patch("solstone.think.batch.agenerate_with_result", new_callable=AsyncMock)
280async def test_batch_retry_pattern(mock_agenerate):
281 """Test retry pattern without bypassing the active model profile."""
282 # First call fails, second succeeds
283 call_count = 0
284
285 async def mock_response(*args, **kwargs):
286 nonlocal call_count
287 call_count += 1
288 if call_count == 1:
289 raise ValueError("Transient error")
290 return _result("Success on retry")
291
292 mock_agenerate.side_effect = mock_response
293
294 batch = Batch(max_concurrent=5)
295
296 # Add initial request
297 req1 = batch.create(contents="Test", context="test.context")
298 req1.attempt = 1
299 batch.add(req1)
300
301 results = []
302 async for req in batch.drain_batch():
303 results.append(req)
304
305 # If error, retry through the same active profile.
306 if req.error and req.attempt == 1:
307 retry = batch.create(
308 contents=req.contents,
309 context="test.context",
310 )
311 retry.attempt = 2
312 batch.add(retry)
313
314 # Should have both attempts
315 assert len(results) == 2
316 assert results[0].error is not None
317 assert results[0].attempt == 1
318 assert results[1].response == "Success on retry"
319 assert results[1].attempt == 2
320
321
322@pytest.mark.asyncio
323@patch("solstone.think.batch.agenerate_with_result", new_callable=AsyncMock)
324async def test_batch_factory_method(mock_agenerate):
325 """Test that batch.create() factory method works correctly."""
326 mock_agenerate.return_value = _result("Response")
327
328 batch = Batch()
329
330 # Use factory method
331 req = batch.create(
332 contents="Test",
333 context="test.context",
334 temperature=0.8,
335 json_output=True,
336 )
337
338 assert isinstance(req, BatchRequest)
339 assert req.contents == "Test"
340 assert req.context == "test.context"
341 assert req.temperature == 0.8
342 assert req.json_output is True
343
344
345@pytest.mark.asyncio
346@patch("solstone.think.batch.agenerate_with_result", new_callable=AsyncMock)
347async def test_batch_can_add_after_draining(mock_agenerate):
348 """Test that adding after draining works (reusable batch)."""
349 mock_agenerate.side_effect = [_result("Response 1"), _result("Response 2")]
350
351 batch = Batch()
352
353 # First batch
354 req1 = batch.create(contents="First", context="test.context")
355 req1.id = 1
356 batch.add(req1)
357
358 results = []
359 async for req in batch.drain_batch():
360 results.append(req)
361
362 assert len(results) == 1
363 assert results[0].id == 1
364
365 # Add more work after draining
366 req2 = batch.create(contents="Second", context="test.context")
367 req2.id = 2
368 batch.add(req2)
369
370 async for req in batch.drain_batch():
371 results.append(req)
372
373 # Should have both results
374 assert len(results) == 2
375 assert {r.id for r in results} == {1, 2}
376
377
378@pytest.mark.asyncio
379@patch("solstone.think.batch.agenerate_with_result", new_callable=AsyncMock)
380async def test_batch_empty_batch(mock_agenerate):
381 """Test that empty batch (no requests) completes immediately."""
382 batch = Batch()
383
384 results = []
385 async for req in batch.drain_batch():
386 results.append(req)
387
388 assert len(results) == 0
389
390
391@pytest.mark.asyncio
392@patch("solstone.think.batch.agenerate_with_result", new_callable=AsyncMock)
393async def test_batch_concurrency_limit(mock_agenerate):
394 """Test that semaphore limits concurrent requests."""
395 # Track concurrent calls
396 concurrent_calls = 0
397 max_concurrent_seen = 0
398 lock = asyncio.Lock()
399
400 async def mock_with_tracking(*args, **kwargs):
401 nonlocal concurrent_calls, max_concurrent_seen
402 async with lock:
403 concurrent_calls += 1
404 max_concurrent_seen = max(max_concurrent_seen, concurrent_calls)
405
406 await asyncio.sleep(0.1) # Simulate API call
407
408 async with lock:
409 concurrent_calls -= 1
410
411 return _result("Response")
412
413 mock_agenerate.side_effect = mock_with_tracking
414
415 # Create batch with max_concurrent=2
416 batch = Batch(max_concurrent=2)
417
418 # Add 5 requests
419 for i in range(5):
420 req = batch.create(contents=f"Request {i}", context="test.context")
421 batch.add(req)
422
423 results = []
424 async for req in batch.drain_batch():
425 results.append(req)
426
427 assert len(results) == 5
428 # Should never exceed max_concurrent=2
429 assert max_concurrent_seen <= 2
430
431
432@pytest.mark.asyncio
433@patch("solstone.think.batch.agenerate_with_result", new_callable=AsyncMock)
434async def test_batch_update_method(mock_agenerate):
435 """Test batch.update() method for modifying and re-adding requests."""
436 call_temperatures = []
437
438 async def mock_track_options(*args, **kwargs):
439 call_temperatures.append(kwargs["temperature"])
440 return _result(f"Response at {kwargs['temperature']}")
441
442 mock_agenerate.side_effect = mock_track_options
443
444 batch = Batch(max_concurrent=5)
445
446 # Create initial request
447 req = batch.create(contents="Initial prompt", context="test.context")
448 req.stage = "initial"
449 batch.add(req)
450
451 results = []
452 result_count = 0
453 async for completed_req in batch.drain_batch():
454 result_count += 1
455 # Capture the response at this point
456 results.append((result_count, completed_req.response, completed_req.stage))
457
458 # After first result, update generation options and re-add.
459 if result_count == 1:
460 batch.update(
461 completed_req,
462 contents="Updated prompt",
463 temperature=0.8,
464 stage="updated", # Update custom attribute too
465 custom_field="test_value", # Add new custom attribute
466 )
467
468 # Should have both results
469 assert len(results) == 2
470 assert results[0][2] == "initial" # First result was initial stage
471 assert results[1][2] == "updated" # Second result was updated stage
472
473 assert call_temperatures == [0.3, 0.8]
474
475 # Verify correct responses at each stage
476 assert results[0][1] == "Response at 0.3"
477 assert results[1][1] == "Response at 0.8"
478
479 # Verify custom attribute was set
480 assert req.custom_field == "test_value"
481
482
483def test_batch_request_with_timeout():
484 """Test BatchRequest can be created with timeout_s parameter."""
485 req = BatchRequest(contents="Test prompt", context="test.context", timeout_s=30)
486 assert req.timeout_s == 30
487
488 req2 = BatchRequest(contents="Test prompt", context="test.context", timeout_s=60.5)
489 assert req2.timeout_s == 60.5
490
491 # Default is None
492 req3 = BatchRequest(contents="Test prompt", context="test.context")
493 assert req3.timeout_s is None
494
495
496@pytest.mark.asyncio
497@patch("solstone.think.batch.agenerate_with_result", new_callable=AsyncMock)
498async def test_batch_timeout_passthrough(mock_agenerate):
499 """Test that timeout_s is passed through to agenerate_with_result."""
500 mock_agenerate.return_value = _result("Response")
501
502 batch = Batch(max_concurrent=5)
503
504 # Create request with timeout_s
505 req = batch.create(contents="Test prompt", context="test.context", timeout_s=45)
506 batch.add(req)
507
508 results = []
509 async for completed_req in batch.drain_batch():
510 results.append(completed_req)
511
512 assert len(results) == 1
513
514 # Verify timeout_s was passed to agenerate_with_result
515 mock_agenerate.assert_called_once()
516 call_kwargs = mock_agenerate.call_args[1]
517 assert call_kwargs["timeout_s"] == 45
518
519
520def test_batch_rejects_provider_specific_client_override():
521 with pytest.raises(TypeError, match="unexpected keyword argument 'client'"):
522 Batch(max_concurrent=5, client=object())
523
524
525@pytest.mark.asyncio
526@patch("solstone.think.batch.agenerate_with_result", new_callable=AsyncMock)
527async def test_batch_passes_json_schema_to_agenerate(mock_agenerate):
528 """Test that json_schema is passed through to agenerate_with_result."""
529 mock_agenerate.return_value = _result("Response")
530
531 batch = Batch(max_concurrent=5)
532 req = batch.create(
533 contents="Test prompt",
534 context="test.context",
535 json_schema={"type": "object"},
536 )
537 batch.add(req)
538
539 results = []
540 async for completed_req in batch.drain_batch():
541 results.append(completed_req)
542
543 assert len(results) == 1
544
545 call_kwargs = mock_agenerate.call_args[1]
546 assert call_kwargs["json_schema"] == {"type": "object"}
547
548
549@pytest.mark.asyncio
550@patch("solstone.think.batch.agenerate_with_result", new_callable=AsyncMock)
551async def test_batch_schema_validation_error_populates_request_error(mock_agenerate):
552 mock_agenerate.side_effect = SchemaValidationError(
553 [{"path": "", "constraint": "json_parse", "message": "empty"}],
554 "",
555 )
556
557 batch = Batch(max_concurrent=5)
558 req = batch.create(
559 contents="Test prompt",
560 context="test.context",
561 json_schema={"type": "object"},
562 )
563 batch.add(req)
564
565 results = []
566 async for completed_req in batch.drain_batch():
567 results.append(completed_req)
568
569 assert len(results) == 1
570 assert results[0].response is None
571 assert "schema validation" in results[0].error
572
573
574@pytest.mark.asyncio
575@patch("solstone.think.batch.agenerate_with_result", new_callable=AsyncMock)
576async def test_batch_non_json_length_finish_populates_text_length_error(mock_agenerate):
577 mock_agenerate.return_value = _result("partial text", finish_reason="max_tokens")
578
579 batch = Batch(max_concurrent=5)
580 req = batch.create(contents="Test prompt", context="test.context")
581 batch.add(req)
582
583 results = []
584 async for completed_req in batch.drain_batch():
585 results.append(completed_req)
586
587 assert len(results) == 1
588 assert results[0].response is None
589 assert results[0].reason_code == "incomplete_text_length"
590 assert "Text response incomplete" in results[0].error
591
592
593@pytest.mark.asyncio
594@patch("solstone.think.batch.agenerate_with_result", new_callable=AsyncMock)
595async def test_batch_non_json_stop_finish_succeeds(mock_agenerate):
596 mock_agenerate.return_value = _result("complete text", finish_reason="stop")
597
598 batch = Batch(max_concurrent=1)
599 req = batch.create(
600 contents="Test",
601 context="test.context",
602 json_output=False,
603 )
604 batch.add(req)
605
606 results = []
607 async for completed_req in batch.drain_batch():
608 results.append(completed_req)
609
610 assert len(results) == 1
611 assert req.response == "complete text"
612 assert req.error is None
613 assert req.reason_code is None
614
615
616@pytest.mark.asyncio
617@patch("solstone.think.batch.agenerate_with_result", new_callable=AsyncMock)
618async def test_batch_non_json_non_length_finish_is_provider_response_invalid(
619 mock_agenerate,
620):
621 mock_agenerate.return_value = _result("", finish_reason="content_filter")
622
623 batch = Batch(max_concurrent=5)
624 req = batch.create(contents="Test prompt", context="test.context")
625 batch.add(req)
626
627 results = []
628 async for completed_req in batch.drain_batch():
629 results.append(completed_req)
630
631 assert len(results) == 1
632 assert results[0].response is None
633 assert results[0].reason_code == "provider_response_invalid"
634
635
636@pytest.mark.asyncio
637@patch("solstone.think.batch.agenerate_with_result", new_callable=AsyncMock)
638async def test_batch_full_result_schema_invalid_stays_hard_failure(mock_agenerate):
639 mock_agenerate.return_value = _result(
640 '{"field": "bad"}',
641 schema_validation={
642 "valid": False,
643 "errors": [{"path": "/field", "constraint": "type", "message": "bad"}],
644 },
645 )
646
647 batch = Batch(max_concurrent=5)
648 req = batch.create(
649 contents="Test prompt",
650 context="test.context",
651 json_schema={"type": "object"},
652 )
653 batch.add(req)
654
655 results = []
656 async for completed_req in batch.drain_batch():
657 results.append(completed_req)
658
659 assert len(results) == 1
660 assert results[0].response is None
661 assert "schema validation" in results[0].error