From cf98059b8dd6605757762ffb50ea5f3fa294cbf3 Mon Sep 17 00:00:00 2001
From: liruifengv <liruifeng1024@gmail.com>
Date: Mon, 25 May 2026 18:29:39 +0800
Subject: [PATCH] fix: wait for background shutdown notifications

---
 .changeset/persist-runtime-model-defaults.md      |  2 +-
 packages/agent-core/src/agent/background/index.ts | 15 ++++++++++++++-
 packages/agent-core/src/session/index.ts          |  7 ++++---
 3 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/.changeset/persist-runtime-model-defaults.md b/.changeset/persist-runtime-model-defaults.md
index 997cbcff6..b7b6584a7 100644
--- a/.changeset/persist-runtime-model-defaults.md
+++ b/.changeset/persist-runtime-model-defaults.md
@@ -3,4 +3,4 @@
 "@moonshot-ai/kimi-code": patch
 ---
 
-Persist model selections from the terminal UI to the default configuration, and honor the configured default thinking state for new sessions.
+Persist model selections from the terminal UI to the default configuration, honor the configured default thinking state for new sessions, and wait for background task terminal notifications during session shutdown before flushing records.
diff --git a/packages/agent-core/src/agent/background/index.ts b/packages/agent-core/src/agent/background/index.ts
index 14af9fa52..b9f0f9e0a 100644
--- a/packages/agent-core/src/agent/background/index.ts
+++ b/packages/agent-core/src/agent/background/index.ts
@@ -35,6 +35,7 @@ const NOTIFICATION_TAIL_BYTES = 3_000;
 export class BackgroundManager extends BackgroundProcessManager {
   private readonly scheduledNotificationKeys = new Set<string>();
   private readonly deliveredNotificationKeys = new Set<string>();
+  private readonly terminalNotificationPromises = new Set<Promise<void>>();
 
   constructor(
     public readonly agent: Agent,
@@ -85,7 +86,18 @@ export class BackgroundManager extends BackgroundProcessManager {
   }
 
   protected override onLiveTaskTerminal(info: BackgroundTaskInfo): void | Promise<void> {
-    return this.notifyBackgroundTask(info);
+    const promise = this.notifyBackgroundTask(info).catch(() => {});
+    this.terminalNotificationPromises.add(promise);
+    void promise.finally(() => {
+      this.terminalNotificationPromises.delete(promise);
+    });
+    return promise;
+  }
+
+  async settleTerminalNotifications(): Promise<void> {
+    while (this.terminalNotificationPromises.size > 0) {
+      await Promise.all(Array.from(this.terminalNotificationPromises));
+    }
   }
 
   private async restoreBackgroundTaskNotifications(): Promise<void> {
@@ -183,6 +195,7 @@ export class BackgroundManager extends BackgroundProcessManager {
     super._reset();
     this.scheduledNotificationKeys.clear();
     this.deliveredNotificationKeys.clear();
+    this.terminalNotificationPromises.clear();
   }
 }
 
diff --git a/packages/agent-core/src/session/index.ts b/packages/agent-core/src/session/index.ts
index b4c0adba5..51184027e 100644
--- a/packages/agent-core/src/session/index.ts
+++ b/packages/agent-core/src/session/index.ts
@@ -194,9 +194,10 @@ export class Session {
     });
     if (keepAliveOnExit) return;
     await Promise.all(
-      Array.from(this.agents.values(), (agent) =>
-        agent.background.stopAll('Session closed'),
-      ),
+      Array.from(this.agents.values(), async (agent) => {
+        await agent.background.stopAll('Session closed');
+        await agent.background.settleTerminalNotifications();
+      }),
     );
   }