diff --git a/raven/importer/orchestrator.py b/raven/importer/orchestrator.py index a38710e80..b1043da91 100644 --- a/raven/importer/orchestrator.py +++ b/raven/importer/orchestrator.py @@ -2,6 +2,7 @@ from __future__ import annotations +import asyncio from collections.abc import Callable, Sequence from dataclasses import dataclass from pathlib import Path @@ -17,13 +18,27 @@ # message_id from (session_id, timestamp_ms, index-within-batch), so those # boundaries are part of the id: two messages sharing a millisecond collide, # and one is dropped, if they land at the same index in different batches. -# Fifty, not a hundred: EverOS extracts on every add, and that cost is -# superlinear in the message count -- against a real service a 15-message -# batch took 12s and a 52-message batch 24s, while a batch of 100 ran past the -# six-minute extraction budget and failed every memory-file source. -_BATCH_MSG_LIMIT = 50 +# Ten: EverOS extracts on every add, and that cost is superlinear in the +# message count -- against a real service a 15-message batch took 12s and a +# 52-message batch 24s, while a batch of 100 ran past the six-minute +# extraction budget and failed every memory-file source. With a slower +# extraction model, batches of 50 took 2.4-7.4 minutes and six of seven +# memory-file sources died on that same budget. Ten is the maintainer's +# call: a batch that finishes well inside the budget on any model matters +# more than the fixed cost of about 7s that every add carries -- which is +# also why it is not one message per add. +_BATCH_MSG_LIMIT = 10 _BATCH_CHAR_LIMIT = 30_000 +# A batch the memory service refuses is sent again before its source is given +# up on. The refusals seen against a real service were transient -- a rate +# limit at the extraction provider, an answer the extractor could not parse, a +# slow answer past the budget -- and a source that fails on one of them takes +# every message behind it down with it, then the next source runs straight +# into the same wall. Three retries with these waits cover a rate-limit window +# of a couple of minutes; the wait polls the stop file so a stop lands in it. +_STORE_RETRY_BACKOFF_S: tuple[float, ...] = (30.0, 60.0, 120.0) + @dataclass(frozen=True) class ImportFailure: @@ -206,8 +221,10 @@ async def _feed_session( cancel_path: Path | None = None, ) -> bool: """Store the session in batches. Returns False when a stop request arrived - between two batches, leaving the rest unsent; a long conversation is many - batches, and a stop that waited for the whole source was not a stop.""" + between two batches or during a retry wait, leaving the rest unsent; a long + conversation is many batches, and a stop that waited for the whole source + was not a stop. A batch the backend refuses is retried on + ``_STORE_RETRY_BACKOFF_S`` before the source counts as failed.""" if not session.messages: return True all_dicts = [_to_store_dict(m) for m in session.messages] @@ -217,36 +234,71 @@ async def _feed_session( def _cancelled() -> bool: return cancel_path is not None and cancel_path.exists() - async def _flush(*, is_final: bool) -> None: + async def _flush(*, is_final: bool) -> bool: nonlocal batch, batch_chars # bulk: nothing waits on an import write; the backend budgets it as extraction. metadata: dict[str, Any] = {"is_final": is_final, "bulk": True} _log_store_request(session.session_id, batch, metadata, batch_chars) - landed = await backend.store(session.session_id, batch, metadata=metadata) - if landed is False: - raise MemoryWriteDroppedError( - f"memory service did not accept a batch for {session.session_id}; source left unsubmitted" + for attempt, wait in enumerate((*_STORE_RETRY_BACKOFF_S, None), start=1): + try: + landed = await backend.store(session.session_id, batch, metadata=metadata) + reason = "memory service did not accept a batch" + except Exception as exc: + landed, reason = False, (str(exc) or repr(exc)) + if landed is not False: + break + if wait is None: + raise MemoryWriteDroppedError( + f"{reason} for {session.session_id} after {attempt} attempts; source left unsubmitted" + ) + logger.warning( + "batch for {} not accepted ({}); retrying in {}s ({}/{})", + session.session_id, + reason, + int(wait), + attempt, + len(_STORE_RETRY_BACKOFF_S), ) + if not await _pause(wait, _cancelled): + return False logger.debug("store completed: session_id={}", session.session_id) batch = [] batch_chars = 0 + return True for msg_dict in all_dicts: msg_chars = len(msg_dict["content"]) if batch and (len(batch) >= _BATCH_MSG_LIMIT or batch_chars + msg_chars > _BATCH_CHAR_LIMIT): if _cancelled(): return False - await _flush(is_final=False) + if not await _flush(is_final=False): + return False batch.append(msg_dict) batch_chars += msg_chars if batch: if _cancelled(): return False - await _flush(is_final=True) + if not await _flush(is_final=True): + return False return True +async def _pause(seconds: float, cancelled: Callable[[], bool]) -> bool: + """Wait out a retry backoff a second at a time, so a stop lands inside it. + + Returns False when the stop arrived.""" + loop = asyncio.get_running_loop() + deadline = loop.time() + seconds + while True: + if cancelled(): + return False + remaining = deadline - loop.time() + if remaining <= 0: + return True + await asyncio.sleep(min(1.0, remaining)) + + def _log_store_request( session_id: str, batch: list[dict[str, Any]], diff --git a/tests/integration/test_import_e2e.py b/tests/integration/test_import_e2e.py index 601b6ccbc..a5af3595a 100644 --- a/tests/integration/test_import_e2e.py +++ b/tests/integration/test_import_e2e.py @@ -237,7 +237,7 @@ async def test_full_pipeline_memory_files(scanner: ClaudeCodeScanner, tmp_path: @pytest.mark.asyncio async def test_batching_large_conversation(scanner: ClaudeCodeScanner, tmp_path: Path) -> None: - """160 messages -> four store calls of 50, 50, 50 and 10, is_final only on the last, every one bulk.""" + """160 messages -> sixteen store calls of 10, is_final only on the last, every one bulk.""" results = await scanner.scan() items = _items_of_kind(scanner, results, SourceKind.CONVERSATION, source_key="sess-large") assert len(items) == 1 @@ -249,8 +249,8 @@ async def test_batching_large_conversation(scanner: ClaudeCodeScanner, tmp_path: assert summary.submitted == 1 sizes = [len(call["messages"]) for call in backend.store_calls] - assert sizes == [50, 50, 50, 10] - assert [call["metadata"]["is_final"] for call in backend.store_calls] == [False, False, False, True] + assert sizes == [10] * 16 + assert [call["metadata"]["is_final"] for call in backend.store_calls] == [False] * 15 + [True] assert all(call["metadata"]["bulk"] is True for call in backend.store_calls) assert sum(sizes) == 160 diff --git a/tests/test_importer_orchestrator.py b/tests/test_importer_orchestrator.py index c803bd390..cdda04b70 100644 --- a/tests/test_importer_orchestrator.py +++ b/tests/test_importer_orchestrator.py @@ -7,6 +7,7 @@ import pytest +from raven.importer import orchestrator from raven.importer.orchestrator import ImportSummary, ProgressEvent, run_import from raven.importer.state import ImportState from raven.importer.types import ( @@ -22,15 +23,30 @@ # --------------------------------------------------------------------------- +@pytest.fixture(autouse=True) +def _no_backoff(monkeypatch: pytest.MonkeyPatch) -> None: + """The retry waits are real seconds in production; tests that care set them.""" + monkeypatch.setattr(orchestrator, "_STORE_RETRY_BACKOFF_S", (0.0, 0.0, 0.0)) + + class FakeBackend: """Records store() calls for assertion.""" - def __init__(self, *, fail_on: set[str] | None = None, drop_on: set[str] | None = None) -> None: + def __init__( + self, + *, + fail_on: set[str] | None = None, + drop_on: set[str] | None = None, + drop_first: dict[str, int] | None = None, + ) -> None: self.calls: list[dict[str, Any]] = [] + self.attempts = 0 self._fail_on = fail_on or set() # A backend that reports a dropped write instead of raising: the shape # a real EverosBackend takes when the memory service is unavailable. self._drop_on = drop_on or set() + # Refuse the first n writes of a session, then accept: a transient fault. + self._drop_first = dict(drop_first or {}) async def recall(self, query: str, *, user_id: str | None = None, agent_id: str | None = None, top_k: int) -> list: return [] @@ -38,10 +54,14 @@ async def recall(self, query: str, *, user_id: str | None = None, agent_id: str async def store( self, session_id: str, messages: list[dict[str, Any]], *, metadata: dict[str, Any] | None = None ) -> bool: + self.attempts += 1 if session_id in self._fail_on: raise RuntimeError(f"store failed for {session_id}") if session_id in self._drop_on: return False + if self._drop_first.get(session_id, 0) > 0: + self._drop_first[session_id] -= 1 + return False self.calls.append({"session_id": session_id, "messages": messages, "metadata": metadata}) return True @@ -247,18 +267,116 @@ async def test_a_dropped_write_is_treated_exactly_like_a_raised_one(self, tmp_pa assert state.is_submitted("claude_code", "b") +class TestStoreRetry: + """A refused batch is sent again, with a wait, before its source is given up on.""" + + @staticmethod + def _recording_pause(waits: list[float]): + async def _pause(seconds: float, cancelled: Any) -> bool: + waits.append(seconds) + return True + + return _pause + + @pytest.mark.asyncio + async def test_a_refused_batch_is_retried_after_each_wait_and_lands( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + waits: list[float] = [] + monkeypatch.setattr(orchestrator, "_STORE_RETRY_BACKOFF_S", (5.0, 7.0, 9.0)) + monkeypatch.setattr(orchestrator, "_pause", self._recording_pause(waits)) + state = ImportState(path=tmp_path / "state.json") + backend = FakeBackend(drop_first={"import-a": 2}) + scanner = FakeScanner({"a": _session(n_msgs=1, session_id="import-a")}) + + summary = await run_import([(scanner, _scan_result("a"))], backend, state) + + assert (summary.submitted, summary.failed) == (1, 0) + assert state.is_submitted("claude_code", "a") + assert backend.attempts == 3 + assert len(backend.calls) == 1 + assert waits == [5.0, 7.0] + + @pytest.mark.asyncio + async def test_a_batch_refused_every_time_fails_its_source_after_the_last_wait( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + waits: list[float] = [] + monkeypatch.setattr(orchestrator, "_STORE_RETRY_BACKOFF_S", (1.0, 2.0, 3.0)) + monkeypatch.setattr(orchestrator, "_pause", self._recording_pause(waits)) + state = ImportState(path=tmp_path / "state.json") + backend = FakeBackend(drop_on={"import-a"}) + scanner = FakeScanner({"a": _session(n_msgs=1, session_id="import-a")}) + + summary = await run_import([(scanner, _scan_result("a"))], backend, state) + + assert (summary.submitted, summary.failed) == (0, 1) + assert backend.attempts == 4 + assert waits == [1.0, 2.0, 3.0] + assert "after 4 attempts" in summary.errors[0].error + + @pytest.mark.asyncio + async def test_a_raised_store_error_is_retried_the_same_way(self, tmp_path: Path) -> None: + state = ImportState(path=tmp_path / "state.json") + backend = FakeBackend(fail_on={"import-a"}) + scanner = FakeScanner({"a": _session(n_msgs=1, session_id="import-a")}) + + summary = await run_import([(scanner, _scan_result("a"))], backend, state) + + assert backend.attempts == 4 + assert summary.errors[0].error.startswith("store failed for import-a for import-a after 4 attempts") + + @pytest.mark.asyncio + async def test_a_stop_during_the_wait_ends_the_run_and_leaves_the_source_unmarked( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + cancel = tmp_path / "cancel" + + async def _stopped_pause(seconds: float, cancelled: Any) -> bool: + cancel.touch() + return False + + monkeypatch.setattr(orchestrator, "_pause", _stopped_pause) + state = ImportState(path=tmp_path / "state.json") + backend = FakeBackend(drop_first={"import-a": 1}) + scanner = FakeScanner( + {"a": _session(n_msgs=1, session_id="import-a"), "b": _session(n_msgs=1, session_id="import-b")} + ) + items = [(scanner, _scan_result("a")), (scanner, _scan_result("b"))] + + summary = await run_import(items, backend, state, cancel_path=cancel) + + assert (summary.submitted, summary.failed) == (0, 0) + assert not state.is_submitted("claude_code", "a") + assert state.get_progress()["entries"] == {} + assert [c["session_id"] for c in backend.calls] == [] + + @pytest.mark.asyncio + async def test_pause_returns_early_when_the_stop_arrives(self) -> None: + seen = 0 + + def _cancelled() -> bool: + nonlocal seen + seen += 1 + return seen > 1 + + assert await orchestrator._pause(0.0, lambda: False) is True + assert await orchestrator._pause(30.0, _cancelled) is False + assert seen == 2 + + class TestBatching: @pytest.mark.asyncio async def test_msg_count_limit(self, tmp_path: Path) -> None: - """120 messages -> 3 batches (50 + 50 + 20), only the last one final, every one bulk.""" + """120 messages -> 12 batches of 10, only the last one final, every one bulk.""" state = ImportState(path=tmp_path / "state.json") backend = FakeBackend() scanner = FakeScanner({"k1": _session(n_msgs=120, session_id="s1", content="x")}) await run_import([(scanner, _scan_result("k1"))], backend, state) - assert [len(c["messages"]) for c in backend.calls] == [50, 50, 20] - assert [c["metadata"]["is_final"] for c in backend.calls] == [False, False, True] + assert [len(c["messages"]) for c in backend.calls] == [10] * 12 + assert [c["metadata"]["is_final"] for c in backend.calls] == [False] * 11 + [True] assert all(c["metadata"]["bulk"] is True for c in backend.calls) def test_a_batch_stays_inside_the_zone_everos_extracts_linearly(self) -> None: @@ -269,7 +387,7 @@ def test_a_batch_stays_inside_the_zone_everos_extracts_linearly(self) -> None: memory-file source it belonged to.""" from raven.importer.orchestrator import _BATCH_MSG_LIMIT - assert _BATCH_MSG_LIMIT <= 50 + assert _BATCH_MSG_LIMIT <= 10 @pytest.mark.asyncio async def test_char_limit_fallback(self, tmp_path: Path) -> None: @@ -288,10 +406,10 @@ async def test_char_limit_fallback(self, tmp_path: Path) -> None: @pytest.mark.asyncio async def test_is_final_only_on_last_batch(self, tmp_path: Path) -> None: - """Exactly 50 messages -> 1 batch with is_final=True.""" + """Exactly 10 messages -> 1 batch with is_final=True.""" state = ImportState(path=tmp_path / "state.json") backend = FakeBackend() - scanner = FakeScanner({"k1": _session(n_msgs=50, session_id="s1", content="x")}) + scanner = FakeScanner({"k1": _session(n_msgs=10, session_id="s1", content="x")}) await run_import([(scanner, _scan_result("k1"))], backend, state) diff --git a/ui-web/src/features/importSync/ImportSyncPage.test.tsx b/ui-web/src/features/importSync/ImportSyncPage.test.tsx index 437e8f1ec..19008d26d 100644 --- a/ui-web/src/features/importSync/ImportSyncPage.test.tsx +++ b/ui-web/src/features/importSync/ImportSyncPage.test.tsx @@ -84,14 +84,18 @@ describe('the import row', () => { expect(runs).toEqual([[['claude_code'], 'memory_files']]) }) - it('says how many did not make it and offers a retry, with nothing to dismiss', () => { + it('says how many did not make it, offers a retry, and a dismiss that takes the row down', async () => { draw(status({ total: 18, submitted: 15, failed: 3, tier: 'full', platforms: ['claude_code'], phases: { status: 'done', errors: [] } })) expect(row()?.className).toContain('importSync-done') expect(row()?.className).toContain('importSync-warn') expect(row()?.textContent).toContain('gui.importSync.failed_n:{"n":3}') expect(main().disabled).toBe(false) expect(main().getAttribute('aria-label')).toBe('gui.importSync.retry') - expect(x()).toBeNull() + expect(x()?.getAttribute('aria-label')).toBe('gui.importSync.dismiss') + + await act(async () => { fireEvent.click(x()!) }) + + expect(row()).toBeNull() }) it('shows a stopped run it cannot ask for again as paused with nothing to click', () => { diff --git a/ui-web/src/features/importSync/ImportSyncPage.tsx b/ui-web/src/features/importSync/ImportSyncPage.tsx index b59b462bb..a1b37285d 100644 --- a/ui-web/src/features/importSync/ImportSyncPage.tsx +++ b/ui-web/src/features/importSync/ImportSyncPage.tsx @@ -61,7 +61,7 @@ export function ImportSyncApp(): JSX.Element | null { {running ? ( - ) : v.kind === 'done' && !v.failed ? ( + ) : v.kind === 'done' ? ( ) : null}