From 81fc424592d628b3bc0e4cb5624eba4dc2312c87 Mon Sep 17 00:00:00 2001 From: Chris Korhonen Date: Sat, 18 Jul 2026 22:19:41 +0800 Subject: [PATCH] fix(openviking): chunk structured session sync Preserve ordered structured turns across OpenViking's 100-message batch limit and resume retries from the first unconfirmed message. Based on the OpenViking batching work from commit 1a567f706703b8005e3fb915548f8a3cf137e581 in #58981. --- contributors/emails/ckorhonen@gmail.com | 2 + plugins/memory/openviking/__init__.py | 84 ++++++++-- .../memory/test_openviking_provider.py | 144 ++++++++++++++++++ 3 files changed, 215 insertions(+), 15 deletions(-) create mode 100644 contributors/emails/ckorhonen@gmail.com diff --git a/contributors/emails/ckorhonen@gmail.com b/contributors/emails/ckorhonen@gmail.com new file mode 100644 index 000000000000..4114b0a65209 --- /dev/null +++ b/contributors/emails/ckorhonen@gmail.com @@ -0,0 +1,2 @@ +ckorhonen +# PR #58981 OpenViking batching salvage diff --git a/plugins/memory/openviking/__init__.py b/plugins/memory/openviking/__init__.py index 9b1dcb58c2de..d957d889e74b 100644 --- a/plugins/memory/openviking/__init__.py +++ b/plugins/memory/openviking/__init__.py @@ -76,6 +76,7 @@ _OPENVIKING_ENV_KEYS = ( _TIMEOUT = 30.0 _SESSION_DRAIN_TIMEOUT = 10.0 _DEFERRED_COMMIT_TIMEOUT = (_TIMEOUT * 2) + 5.0 +_SESSION_MESSAGE_BATCH_LIMIT = 100 _REMOTE_RESOURCE_PREFIXES = ("http://", "https://", "git@", "ssh://", "git://") _SYNC_TRACE_ENV = "HERMES_OPENVIKING_SYNC_TRACE" _DEFAULT_RECALL_LIMIT = 6 @@ -3430,23 +3431,54 @@ class OpenVikingMemoryProvider(MemoryProvider): self._mark_session_pending(sid) def _sync(): - def _post_turn(client: _VikingClient) -> None: - if batch_messages: - payload = {"messages": batch_messages} + next_batch_index = 0 + + def _post_unsent_messages_individually(client: _VikingClient) -> None: + nonlocal next_batch_index + path = f"/api/v1/sessions/{sid}/messages" + while next_batch_index < len(batch_messages): if _sync_trace_enabled(): logger.info( - "OpenViking sync_turn trace: POST /api/v1/sessions/%s/messages/batch payload=%s", - sid, - json.dumps(payload, ensure_ascii=False), + "OpenViking sync_turn trace: POST %s message_index=%d payload=%s", + path, + next_batch_index, + json.dumps(batch_messages[next_batch_index], ensure_ascii=False), ) - try: - client.post(f"/api/v1/sessions/{sid}/messages/batch", payload) + client.post(path, batch_messages[next_batch_index]) + next_batch_index += 1 + + def _post_turn(client: _VikingClient) -> None: + nonlocal next_batch_index + if batch_messages: + while next_batch_index < len(batch_messages): + batch_end = min( + next_batch_index + _SESSION_MESSAGE_BATCH_LIMIT, + len(batch_messages), + ) + payload = {"messages": batch_messages[next_batch_index:batch_end]} + if _sync_trace_enabled(): + logger.info( + "OpenViking sync_turn trace: POST " + "/api/v1/sessions/%s/messages/batch range=%d:%d payload=%s", + sid, + next_batch_index, + batch_end, + json.dumps(payload, ensure_ascii=False), + ) + try: + client.post(f"/api/v1/sessions/{sid}/messages/batch", payload) + except Exception as batch_error: + if next_batch_index: + raise + logger.warning( + "OpenViking structured sync failed; falling back to text sync: %s", + batch_error, + ) + break + next_batch_index = batch_end + + if next_batch_index == len(batch_messages): return - except Exception as batch_error: - logger.warning( - "OpenViking structured sync failed; falling back to text sync: %s", - batch_error, - ) self._post_session_turn( client, @@ -3460,10 +3492,32 @@ class OpenVikingMemoryProvider(MemoryProvider): _post_turn(client) except Exception as e: logger.debug("OpenViking sync_turn failed, reconnecting: %s", e) + retry_client = None try: - client = self._new_client() - _post_turn(client) + retry_client = self._new_client() + _post_turn(retry_client) except Exception as retry_error: + if ( + retry_client is not None + and batch_messages + and next_batch_index < len(batch_messages) + ): + logger.warning( + "OpenViking structured sync retry failed; writing %d remaining " + "messages individually: %s", + len(batch_messages) - next_batch_index, + retry_error, + ) + try: + _post_unsent_messages_individually(retry_client) + return + except Exception as fallback_error: + logger.warning( + "OpenViking sync_turn failed during individual-message " + "fallback: %s", + fallback_error, + ) + return logger.warning("OpenViking sync_turn failed: %s", retry_error) self._spawn_writer(sid, _sync, name="openviking-sync") diff --git a/tests/plugins/memory/test_openviking_provider.py b/tests/plugins/memory/test_openviking_provider.py index 45e868f4a715..f4afea19aa56 100644 --- a/tests/plugins/memory/test_openviking_provider.py +++ b/tests/plugins/memory/test_openviking_provider.py @@ -2405,6 +2405,150 @@ def test_sync_turn_retries_batch_write_with_fresh_client(): )] +def _long_structured_turn(assistant_count=204): + return [ + {"role": "user", "content": "u"}, + *[ + {"role": "assistant", "content": f"assistant-{index}"} + for index in range(assistant_count) + ], + ] + + +@pytest.mark.parametrize( + ("assistant_count", "expected_chunk_sizes"), + [ + (99, [100]), + (100, [100, 1]), + (204, [100, 100, 5]), + ], +) +def test_sync_turn_chunks_structured_messages_to_openviking_limit( + monkeypatch, + assistant_count, + expected_chunk_sizes, +): + provider = OpenVikingMemoryProvider() + provider._client = MagicMock() + provider._endpoint = "http://test" + provider._api_key = "" + provider._account = "acct" + provider._user = "usr" + provider._agent = "hermes" + provider._session_id = "sid-chunked" + + captured = [] + + class StubClient: + def __init__(self, *args, **kwargs): + pass + + def post(self, path, payload=None, **kwargs): + captured.append((path, payload)) + return {} + + monkeypatch.setattr(openviking_module, "_VikingClient", StubClient) + messages = _long_structured_turn(assistant_count) + + provider.sync_turn("u", f"assistant-{assistant_count - 1}", messages=messages) + assert provider._drain_writers("sid-chunked", timeout=2.0) + + assert [ + len(payload["messages"]) + for _path, payload in captured + ] == expected_chunk_sizes + assert all(path == "/api/v1/sessions/sid-chunked/messages/batch" for path, _ in captured) + assert [ + message + for _path, payload in captured + for message in payload["messages"] + ] == provider._messages_to_openviking_batch(messages, assistant_peer_id="hermes") + + +def test_sync_turn_retries_only_unsent_chunks_with_fresh_client(monkeypatch): + provider = OpenVikingMemoryProvider() + provider._client = MagicMock() + provider._endpoint = "http://test" + provider._api_key = "" + provider._account = "acct" + provider._user = "usr" + provider._agent = "hermes" + provider._session_id = "sid-resume" + + clients = [] + attempts = [] + accepted = [] + + class StubClient: + def __init__(self, *args, **kwargs): + self.index = len(clients) + clients.append(self) + + def post(self, path, payload=None, **kwargs): + attempts.append((self.index, path, payload)) + if self.index == 0 and len([item for item in attempts if item[0] == 0]) == 2: + raise RuntimeError("transient second chunk failure") + accepted.extend(payload["messages"]) + return {} + + monkeypatch.setattr(openviking_module, "_VikingClient", StubClient) + messages = _long_structured_turn() + + provider.sync_turn("u", "assistant-203", messages=messages) + assert provider._drain_writers("sid-resume", timeout=2.0) + + assert len(clients) == 2 + assert [ + (client_index, len(payload["messages"])) + for client_index, _path, payload in attempts + ] == [(0, 100), (0, 100), (1, 100), (1, 5)] + assert accepted == provider._messages_to_openviking_batch( + messages, + assistant_peer_id="hermes", + ) + + +def test_sync_turn_falls_back_to_individual_writes_for_unsent_chunks(monkeypatch): + provider = OpenVikingMemoryProvider() + provider._client = MagicMock() + provider._endpoint = "http://test" + provider._api_key = "" + provider._account = "acct" + provider._user = "usr" + provider._agent = "hermes" + provider._session_id = "sid-individual-fallback" + + clients = [] + accepted = [] + + class StubClient: + def __init__(self, *args, **kwargs): + self.index = len(clients) + clients.append(self) + + def post(self, path, payload=None, **kwargs): + if path.endswith("/messages/batch"): + if self.index == 0 and not accepted: + accepted.extend(payload["messages"]) + return {} + raise RuntimeError("persistent batch failure") + assert path == "/api/v1/sessions/sid-individual-fallback/messages" + accepted.append(payload) + return {} + + monkeypatch.setattr(openviking_module, "_VikingClient", StubClient) + messages = _long_structured_turn() + + provider.sync_turn("u", "assistant-203", messages=messages) + assert provider._drain_writers("sid-individual-fallback", timeout=2.0) + + assert len(clients) == 2 + assert accepted == provider._messages_to_openviking_batch( + messages, + assistant_peer_id="hermes", + ) + + def test_sync_turn_structured_messages_include_assistant_peer_id(): provider = OpenVikingMemoryProvider() provider._client = MagicMock()