From 3c0aa7e2372d0655861bd4b545b6c592c97603a0 Mon Sep 17 00:00:00 2001
From: TT <12365@566>
Date: Sat, 3 Oct 2026 18:26:09 +0800
Subject: [PATCH] fix(session): preserve queued commit on shutdown cancellation

---
 openviking/session/session.py                 | 11 +++--
 .../session/test_session_commit_resume.py     | 46 +++++++++++++++++++
 2 files changed, 52 insertions(+), 5 deletions(-)

diff --git a/openviking/session/session.py b/openviking/session/session.py
index 2c97e777ef..5244d2a14c 100644
--- a/openviking/session/session.py
+++ b/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))
diff --git a/tests/unit/session/test_session_commit_resume.py b/tests/unit/session/test_session_commit_resume.py
index 285a14e0a1..dbf2b54391 100644
--- a/tests/unit/session/test_session_commit_resume.py
+++ b/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"