diff --git a/openviking/session/session.py b/openviking/session/session.py index 084115951a..bd5e0b7c9a 100644 --- a/openviking/session/session.py +++ b/openviking/session/session.py @@ -62,6 +62,7 @@ WM_UPDATE_TOOL, ) from openviking.storage.abstract_overview import render_abstract_overview +from openviking.storage.errors import LockAcquisitionError from openviking.telemetry import get_current_telemetry, tracer from openviking.telemetry.request_wait_tracker import get_request_wait_tracker from openviking.utils.model_retry import is_retryable_api_error, retry_async @@ -1824,20 +1825,23 @@ async def resume_queued_commit(self, msg: "SessionCommitMsg") -> bool: if user_config_error is not None: user_config_error = str(user_config_error) - await self._run_memory_extraction( - task_id=msg.task_id, - archive_uri=msg.archive_uri, - messages=archive_messages, - first_message_id=archive_messages[0].id, - last_message_id=archive_messages[-1].id, - memory_policy=msg.memory_policy, - agent_evolution_enabled=agent_evolution_enabled, - agent_memory_skip_reason=agent_memory_skip_reason, - user_config_error=user_config_error, - record_auto_commit_success=msg.record_auto_commit_success, - event_search_tags=list(msg.event_search_tags or []), - auto_commit_policy=msg.auto_commit_policy, - ) + try: + await self._run_memory_extraction( + task_id=msg.task_id, + archive_uri=msg.archive_uri, + messages=archive_messages, + first_message_id=archive_messages[0].id, + last_message_id=archive_messages[-1].id, + memory_policy=msg.memory_policy, + agent_evolution_enabled=agent_evolution_enabled, + agent_memory_skip_reason=agent_memory_skip_reason, + user_config_error=user_config_error, + record_auto_commit_success=msg.record_auto_commit_success, + event_search_tags=list(msg.event_search_tags or []), + auto_commit_policy=msg.auto_commit_policy, + ) + except LockAcquisitionError: + return False # requeued by SessionCommitProcessor return True async def _run_usage_reporting( @@ -2445,6 +2449,13 @@ async def _run_long_term_memory_extraction( error="session commit cancelled", ) raise + except LockAcquisitionError as e: + # Lock contention is transient and completed steps are checkpointed: + # hand the job back to the queue instead of failing the archive. + telemetry.set_error("session.commit.phase2", type(e).__name__, str(e)) + _publish_telemetry_summary_best_effort(telemetry.finish("error")) + logger.warning("Phase 2 for %s hit lock contention, requeueing: %s", archive_uri, e) + raise except Exception as e: telemetry.set_error("session.commit.phase2", type(e).__name__, str(e)) snapshot = telemetry.finish("error") diff --git a/openviking/storage/queuefs/session_commit_processor.py b/openviking/storage/queuefs/session_commit_processor.py index 9eb14088b6..178fb06a84 100644 --- a/openviking/storage/queuefs/session_commit_processor.py +++ b/openviking/storage/queuefs/session_commit_processor.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: AGPL-3.0 """Queue consumer for restart-safe Session Phase 2 work.""" +import asyncio import json from typing import TYPE_CHECKING, Any, Dict, Optional @@ -21,6 +22,10 @@ if TYPE_CHECKING: from openviking.service.session_service import SessionService +# A commit that waits for its predecessor (or a lock) is requeued; without a pause +# each waiting message cycles through the queue continuously. +_REQUEUE_DELAY_S = 2.0 + class SessionCommitProcessor(DequeueHandlerBase): def __init__( @@ -84,6 +89,7 @@ async def _process(self, msg: SessionCommitMsg, ctx: RequestContext) -> bool: if not processed: from openviking.storage.queuefs import QueueManager, get_queue_manager + await asyncio.sleep(_REQUEUE_DELAY_S) await get_queue_manager().enqueue( QueueManager.SESSION_COMMIT, msg.to_dict(), diff --git a/tests/storage/test_session_commit_processor_identity.py b/tests/storage/test_session_commit_processor_identity.py index 5c184413b7..64a2a57182 100644 --- a/tests/storage/test_session_commit_processor_identity.py +++ b/tests/storage/test_session_commit_processor_identity.py @@ -91,11 +91,16 @@ async def test_process_binds_committing_identity_to_root_context(): async def test_process_requeues_deferred_commit_and_resets_root_context(monkeypatch): queued = [] + events = [] class _QueueManager: async def enqueue(self, queue_name, data): + events.append("enqueue") queued.append((queue_name, data)) + async def fake_sleep(seconds): + events.append(("sleep", seconds)) + processor = SessionCommitProcessor( _FakeSessionService({}, processed=False), ) @@ -103,12 +108,17 @@ async def enqueue(self, queue_name, data): "openviking.storage.queuefs.get_queue_manager", lambda: _QueueManager(), ) + monkeypatch.setattr( + "openviking.storage.queuefs.session_commit_processor.asyncio.sleep", fake_sleep + ) ctx = RequestContext(user=UserIdentifier("acme", "alice"), role=Role.USER) result = await processor._process(_make_msg(), ctx) assert result is False assert queued == [("SessionCommit", _make_msg().to_dict())] + # A waiting commit pauses before it goes back to the queue instead of cycling. + assert events == [("sleep", 2.0), "enqueue"] assert get_root_observability_context() is None diff --git a/tests/unit/session/test_session_commit_resume.py b/tests/unit/session/test_session_commit_resume.py index 90adac7c7a..1ee0e496c1 100644 --- a/tests/unit/session/test_session_commit_resume.py +++ b/tests/unit/session/test_session_commit_resume.py @@ -14,6 +14,7 @@ from openviking.service.task_tracker import TaskStatus, TaskTracker, set_task_tracker from openviking.session.memory_policy import MemoryPolicy from openviking.session.session import Session +from openviking.storage.errors import LockAcquisitionError from openviking.storage.queuefs.session_commit_msg import SessionCommitMsg @@ -218,6 +219,47 @@ async def cancel_prepare(*_args, **_kwargs): assert (f"{archive_uri}/.failed.json" in files) is cancellation_requested +@pytest.mark.asyncio +async def test_phase2_lock_contention_requeues_without_failing_archive(monkeypatch): + session_uri = "viking://user/default/sessions/session-1" + archive_uri = f"{session_uri}/history/archive_001" + archived = Message(id="archived", role="user", parts=[TextPart("old")]) + files = { + f"{session_uri}/.meta.json": json.dumps( + {"session_id": "session-1", "message_count": 0, "commit_count": 1} + ), + f"{archive_uri}/messages.jsonl": f"{archived.to_jsonl()}\n", + } + session = Session( + viking_fs=_MemoryVikingFS(files), session_id="session-1", session_uri=session_uri + ) + tracker = TaskTracker(_TaskStore()) + set_task_tracker(tracker) + + async def locked(*_args, **_kwargs): + raise LockAcquisitionError("lock acquire timed out after 30000ms: /local/session-1") + + monkeypatch.setattr(session, "_prepare_phase2_archive_messages", locked) + message = SessionCommitMsg( + task_id="task-1", + session_id="session-1", + session_uri=session_uri, + archive_uri=archive_uri, + user={"account_id": "default", "user_id": "default"}, + memory_policy={"memory_types": []}, + ) + + try: + processed = await session.resume_queued_commit(message) + task = await tracker.get("task-1", account_id="default", user_id="default") + finally: + set_task_tracker(None) + + assert processed is False # SessionCommitProcessor requeues it + assert f"{archive_uri}/.failed.json" not in files + assert task.status != TaskStatus.FAILED + + @pytest.mark.asyncio async def test_resume_queued_commit_continues_phase2(monkeypatch): session_uri = "viking://user/default/sessions/session-1"