fix(session): preserve queued commit on shutdown cancellation

openviking/session/session.py

@@ -2403,11 +2403,12 @@ async def _run_long_term_memory_extraction(
             telemetry.set_error("session.commit.phase2", "CANCELLED", "session commit cancelled")
             snapshot = telemetry.finish("cancelled")
             _publish_telemetry_summary_best_effort(snapshot)
-            await self._write_failed_marker(
-                archive_uri,
-                stage="cancelled",
-                error="session commit cancelled",
-            )
+            if tracker.is_cancellation_requested(task_id):
+                await self._write_failed_marker(
+                    archive_uri,
+                    stage="cancelled",
+                    error="session commit cancelled",
+                )
             raise
         except Exception as e:
             telemetry.set_error("session.commit.phase2", type(e).__name__, str(e))

tests/unit/session/test_session_commit_resume.py

@@ -1,6 +1,7 @@
 # Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd.
 # SPDX-License-Identifier: AGPL-3.0
 
+import asyncio
 import inspect
 import json
 from dataclasses import fields
@@ -137,6 +138,51 @@ def test_phase2_auto_commit_policy_parameters_are_appended():
     assert fields(SessionCommitMsg)[-1].name == "auto_commit_policy"
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("cancellation_requested", [False, True])
+async def test_phase2_cancellation_marks_archive_only_when_requested(
+    monkeypatch, cancellation_requested
+):
+    session_uri = "viking://user/default/sessions/session-1"
+    archive_uri = f"{session_uri}/history/archive_001"
+    tracker = TaskTracker(_TaskStore())
+    await tracker.create(
+        "session_commit",
+        account_id="default",
+        user_id="default",
+        task_id="task-1",
+    )
+    if cancellation_requested:
+        await tracker.record_cancelled("task-1", account_id="default", user_id="default")
+    else:
+        await tracker.start("task-1", account_id="default", user_id="default")
+
+    monkeypatch.setattr("openviking.service.task_tracker.get_task_tracker", lambda: tracker)
+    files = {}
+    session = Session(
+        viking_fs=_MemoryVikingFS(files),
+        session_id="session-1",
+        session_uri=session_uri,
+    )
+
+    async def cancel_prepare(*_args, **_kwargs):
+        raise asyncio.CancelledError
+
+    monkeypatch.setattr(session, "_prepare_phase2_archive_messages", cancel_prepare)
+
+    with pytest.raises(asyncio.CancelledError):
+        await session._run_memory_extraction(
+            task_id="task-1",
+            archive_uri=archive_uri,
+            messages=[],
+            first_message_id="",
+            last_message_id="",
+            memory_policy=None,
+        )
+
+    assert (f"{archive_uri}/.failed.json" in files) is cancellation_requested
+
+
 @pytest.mark.asyncio
 async def test_resume_queued_commit_continues_phase2(monkeypatch):
     session_uri = "viking://user/default/sessions/session-1"