From 891e42beec49c7e1eddd7148644e4e97daf3c947 Mon Sep 17 00:00:00 2001
From: Byron <loong0306@163.com>
Date: Tue, 14 Apr 2026 04:24:56 +0800
Subject: [PATCH] fix(ui): preserve user-selected session on reconnect and tab
 switch (#59611) thanks @loong0306
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Fixes #57072 — chat UI state desync after route navigation.

- applySessionDefaults() now detects user-selected sessions and preserves them on reconnect
- Chat tab session switching consolidated to use switchChatSession() helper
- Overview session-key handler uses shared resetChatStateForSessionSwitch to prevent stale state leaks
- Session select dropdowns now set ?selected to reflect actual state

Co-authored-by: loong0306 <loong0306@gmail.com>
Co-authored-by: Nova <nova@openknot.ai>
---
 ui/src/ui/app-gateway.ts        | 21 +++++++++++++++++++++
 ui/src/ui/app-render.helpers.ts | 22 ++++++++++++++++++----
 ui/src/ui/app-render.ts         |  6 +++++-
 3 files changed, 44 insertions(+), 5 deletions(-)

diff --git a/ui/src/ui/app-gateway.ts b/ui/src/ui/app-gateway.ts
index f95a68014d8..ea1a30a0c8d 100644
--- a/ui/src/ui/app-gateway.ts
+++ b/ui/src/ui/app-gateway.ts
@@ -177,6 +177,27 @@ function applySessionDefaults(host: GatewayHost, defaults?: SessionDefaultsSnaps
   if (!defaults?.mainSessionKey) {
     return;
   }
+
+  // Detect if user has already selected a specific session (not an alias like "main").
+  // If normalization doesn't change the value, it's a user-selected session.
+  const normalizedSessionKey = normalizeSessionKeyForDefaults(host.sessionKey, defaults);
+  const isUserSelectedSession = normalizedSessionKey === host.sessionKey;
+
+  if (isUserSelectedSession) {
+    // User has selected a specific session; preserve their choice
+    // Only normalize lastActiveSessionKey, don't override current sessionKey
+    const resolvedLastActiveSessionKey = normalizeSessionKeyForDefaults(
+      host.settings.lastActiveSessionKey,
+      defaults,
+    );
+    if (resolvedLastActiveSessionKey !== host.settings.lastActiveSessionKey) {
+      applySettings(host as unknown as Parameters<typeof applySettings>[0], {
+        ...host.settings,
+        lastActiveSessionKey: resolvedLastActiveSessionKey,
+      });
+    }
+    return; // Keep user's session selection
+  }
   const resolvedSessionKey = normalizeSessionKeyForDefaults(host.sessionKey, defaults);
   const resolvedSettingsSessionKey = normalizeSessionKeyForDefaults(
     host.settings.sessionKey,
diff --git a/ui/src/ui/app-render.helpers.ts b/ui/src/ui/app-render.helpers.ts
index 35f3f21bdac..3fc87913a0e 100644
--- a/ui/src/ui/app-render.helpers.ts
+++ b/ui/src/ui/app-render.helpers.ts
@@ -117,9 +117,11 @@ export function renderTab(state: AppViewState, tab: Tab, opts?: { collapsed?: bo
         }
         event.preventDefault();
         if (tab === "chat") {
-          const mainSessionKey = resolveSidebarChatSessionKey(state);
-          if (state.sessionKey !== mainSessionKey) {
+          if (!state.sessionKey) {
+            const mainSessionKey = resolveSidebarChatSessionKey(state);
             resetChatStateForSessionSwitch(state, mainSessionKey);
+          }
+          if (state.tab !== "chat") {
             void state.loadAssistantIdentity();
           }
         }
@@ -202,7 +204,13 @@ export function renderChatSessionSelect(state: AppViewState) {
                   group.options,
                   (entry) => entry.key,
                   (entry) =>
-                    html`<option value=${entry.key} title=${entry.title}>${entry.label}</option>`,
+                    html`<option
+                      value=${entry.key}
+                      title=${entry.title}
+                      ?selected=${entry.key === state.sessionKey}
+                    >
+                      ${entry.label}
+                    </option>`,
                 )}
               </optgroup>`,
           )}
@@ -474,7 +482,13 @@ export function renderChatMobileToggle(state: AppViewState) {
                   <optgroup label=${group.label}>
                     ${group.options.map(
                       (opt) => html`
-                        <option value=${opt.key} title=${opt.title}>${opt.label}</option>
+                        <option
+                          value=${opt.key}
+                          title=${opt.title}
+                          ?selected=${opt.key === state.sessionKey}
+                        >
+                          ${opt.label}
+                        </option>
                       `,
                     )}
                   </optgroup>
diff --git a/ui/src/ui/app-render.ts b/ui/src/ui/app-render.ts
index 5cc294058d8..e5ecb0ee229 100644
--- a/ui/src/ui/app-render.ts
+++ b/ui/src/ui/app-render.ts
@@ -1113,13 +1113,17 @@ export function renderApp(state: AppViewState) {
               onSessionKeyChange: (next) => {
                 state.sessionKey = next;
                 state.chatMessage = "";
+                state.chatMessages = [];
+                state.chatToolMessages = [];
+                state.chatStream = null;
+                state.chatRunId = null;
+                state.chatQueue = [];
                 state.resetToolStream();
                 state.applySettings({
                   ...state.settings,
                   sessionKey: next,
                   lastActiveSessionKey: next,
                 });
-                void state.loadAssistantIdentity();
               },
               onToggleGatewayTokenVisibility: () => {
                 state.overviewShowGatewayToken = !state.overviewShowGatewayToken;