|
24 | 24 | query, |
25 | 25 | ) |
26 | 26 | from claude_agent_sdk._internal.query import Query |
| 27 | +from claude_agent_sdk._internal.session_resume import build_mirror_batcher |
27 | 28 | from claude_agent_sdk._internal.session_store import file_path_to_session_key |
28 | 29 | from claude_agent_sdk._internal.sessions import _get_projects_dir |
29 | 30 | from claude_agent_sdk._internal.transcript_mirror_batcher import ( |
@@ -503,6 +504,62 @@ async def append(self, key, entries): |
503 | 504 | assert order == [1, 2, 3] |
504 | 505 |
|
505 | 506 |
|
| 507 | +# --------------------------------------------------------------------------- |
| 508 | +# build_mirror_batcher / session_store_flush |
| 509 | +# --------------------------------------------------------------------------- |
| 510 | + |
| 511 | + |
| 512 | +class TestBuildMirrorBatcherFlushMode: |
| 513 | + """``session_store_flush`` threads through ``build_mirror_batcher`` to the |
| 514 | + batcher's pending thresholds: ``"batched"`` keeps the defaults, |
| 515 | + ``"eager"`` zeroes them so every enqueue schedules a background flush.""" |
| 516 | + |
| 517 | + @pytest.mark.parametrize( |
| 518 | + ("kwargs", "want_entries", "want_bytes"), |
| 519 | + [ |
| 520 | + ({}, MAX_PENDING_ENTRIES, MAX_PENDING_BYTES), |
| 521 | + ({"flush_mode": "batched"}, MAX_PENDING_ENTRIES, MAX_PENDING_BYTES), |
| 522 | + ({"flush_mode": "eager"}, 0, 0), |
| 523 | + ], |
| 524 | + ids=["default", "batched", "eager"], |
| 525 | + ) |
| 526 | + def test_flush_mode_sets_thresholds( |
| 527 | + self, kwargs: dict[str, Any], want_entries: int, want_bytes: int |
| 528 | + ) -> None: |
| 529 | + batcher = build_mirror_batcher( |
| 530 | + store=InMemorySessionStore(), |
| 531 | + materialized=None, |
| 532 | + env={"CLAUDE_CONFIG_DIR": str(Path(PROJECTS_DIR).parent)}, |
| 533 | + on_error=_noop_error, |
| 534 | + **kwargs, |
| 535 | + ) |
| 536 | + assert batcher.max_pending_entries == want_entries |
| 537 | + assert batcher.max_pending_bytes == want_bytes |
| 538 | + |
| 539 | + @pytest.mark.asyncio |
| 540 | + async def test_eager_mode_flushes_per_frame(self) -> None: |
| 541 | + store = _RecordingStore() |
| 542 | + batcher = build_mirror_batcher( |
| 543 | + store=store, |
| 544 | + materialized=None, |
| 545 | + env={"CLAUDE_CONFIG_DIR": str(Path(PROJECTS_DIR).parent)}, |
| 546 | + on_error=_noop_error, |
| 547 | + flush_mode="eager", |
| 548 | + ) |
| 549 | + batcher.enqueue(_main_path(), [{"type": "user", "n": 1}]) |
| 550 | + await asyncio.sleep(0) |
| 551 | + await asyncio.sleep(0) |
| 552 | + assert len(store.append_calls) == 1 |
| 553 | + batcher.enqueue(_main_path(), [{"type": "assistant", "n": 2}]) |
| 554 | + await asyncio.sleep(0) |
| 555 | + await asyncio.sleep(0) |
| 556 | + assert len(store.append_calls) == 2 |
| 557 | + assert [e["n"] for c in store.append_calls for e in c[1]] == [1, 2] |
| 558 | + |
| 559 | + def test_options_default_is_batched(self) -> None: |
| 560 | + assert ClaudeAgentOptions().session_store_flush == "batched" |
| 561 | + |
| 562 | + |
506 | 563 | # --------------------------------------------------------------------------- |
507 | 564 | # --session-mirror CLI flag |
508 | 565 | # --------------------------------------------------------------------------- |
@@ -533,11 +590,15 @@ def test_flag_absent_when_session_store_unset(self) -> None: |
533 | 590 | # --------------------------------------------------------------------------- |
534 | 591 |
|
535 | 592 |
|
536 | | -def _make_mock_transport(messages: list[dict[str, Any]]) -> Any: |
| 593 | +def _make_mock_transport( |
| 594 | + messages: list[dict[str, Any]], *, yield_between: bool = False |
| 595 | +) -> Any: |
537 | 596 | mock_transport = AsyncMock() |
538 | 597 |
|
539 | 598 | async def mock_receive(): |
540 | 599 | for msg in messages: |
| 600 | + if yield_between: |
| 601 | + await anyio.sleep(0) |
541 | 602 | yield msg |
542 | 603 |
|
543 | 604 | mock_transport.read_messages = mock_receive |
@@ -710,6 +771,67 @@ async def _test() -> None: |
710 | 771 |
|
711 | 772 | anyio.run(_test) |
712 | 773 |
|
| 774 | + def test_eager_flush_mode_appends_per_frame_before_result(self) -> None: |
| 775 | + """With ``session_store_flush="eager"`` each ``transcript_mirror`` frame |
| 776 | + is flushed as it arrives, so the store sees one ``append()`` per frame |
| 777 | + rather than a single coalesced batch at ``result`` time.""" |
| 778 | + |
| 779 | + async def _test() -> None: |
| 780 | + store = _RecordingStore() |
| 781 | + frame1 = { |
| 782 | + "type": "transcript_mirror", |
| 783 | + "filePath": _main_path("p", "s"), |
| 784 | + "entries": [{"type": "user", "uuid": "u1"}], |
| 785 | + } |
| 786 | + frame2 = { |
| 787 | + "type": "transcript_mirror", |
| 788 | + "filePath": _main_path("p", "s"), |
| 789 | + "entries": [{"type": "assistant", "uuid": "a1"}], |
| 790 | + } |
| 791 | + # Yield to the event loop between frames so the eager background |
| 792 | + # drain scheduled by enqueue() can run before the next frame |
| 793 | + # arrives — models the await on real stdout I/O. Without this the |
| 794 | + # mock delivers both frames synchronously and they coalesce, which |
| 795 | + # is correct back-pressure behaviour but not what we're asserting. |
| 796 | + mock_transport = _make_mock_transport( |
| 797 | + [frame1, frame2, _ASSISTANT_MSG, _RESULT_MSG], |
| 798 | + yield_between=True, |
| 799 | + ) |
| 800 | + |
| 801 | + with ( |
| 802 | + patch( |
| 803 | + "claude_agent_sdk._internal.client.SubprocessCLITransport" |
| 804 | + ) as mock_cls, |
| 805 | + patch( |
| 806 | + "claude_agent_sdk._internal.query.Query.initialize", |
| 807 | + new_callable=AsyncMock, |
| 808 | + ), |
| 809 | + patch( |
| 810 | + "claude_agent_sdk._internal.session_resume._get_projects_dir", |
| 811 | + return_value=PROJECTS_DIR, |
| 812 | + ), |
| 813 | + ): |
| 814 | + mock_cls.return_value = mock_transport |
| 815 | + appends_at_assistant = None |
| 816 | + async for msg in query( |
| 817 | + prompt="Hello", |
| 818 | + options=ClaudeAgentOptions( |
| 819 | + session_store=store, session_store_flush="eager" |
| 820 | + ), |
| 821 | + ): |
| 822 | + if isinstance(msg, AssistantMessage): |
| 823 | + appends_at_assistant = len(store.append_calls) |
| 824 | + |
| 825 | + # Both frames flushed individually before the assistant message |
| 826 | + # was yielded (eager background flush ran while the read loop |
| 827 | + # awaited the next stdout line). |
| 828 | + assert appends_at_assistant == 2 |
| 829 | + assert len(store.append_calls) == 2 |
| 830 | + assert store.append_calls[0][1] == [{"type": "user", "uuid": "u1"}] |
| 831 | + assert store.append_calls[1][1] == [{"type": "assistant", "uuid": "a1"}] |
| 832 | + |
| 833 | + anyio.run(_test) |
| 834 | + |
713 | 835 | def test_mirror_frames_dropped_when_no_session_store(self) -> None: |
714 | 836 | """Without a session_store the batcher isn't attached; frames are |
715 | 837 | peeled and dropped (still not yielded), normal messages flow.""" |
|
0 commit comments