From ae322c53cc7f5a209206692a2d7e6e94a0911076 Mon Sep 17 00:00:00 2001
From: Shantur Rathore <i@shantur.com>
Date: Thu, 15 Jan 2026 08:26:49 +0000
Subject: [PATCH] fix(ui): correct Go to Session navigation across instances

---
 packages/ui/src/components/message-block.tsx     | 16 ++++++++++++++--
 .../src/components/permission-approval-modal.tsx | 11 +++++++++--
 2 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/packages/ui/src/components/message-block.tsx b/packages/ui/src/components/message-block.tsx
index eace1a0e..650df4ac 100644
--- a/packages/ui/src/components/message-block.tsx
+++ b/packages/ui/src/components/message-block.tsx
@@ -82,8 +82,20 @@ interface TaskSessionLocation {
   parentId: string | null
 }
 
-function findTaskSessionLocation(sessionId: string): TaskSessionLocation | null {
+function findTaskSessionLocation(sessionId: string, preferredInstanceId?: string): TaskSessionLocation | null {
   if (!sessionId) return null
+
+  if (preferredInstanceId) {
+    const session = sessions().get(preferredInstanceId)?.get(sessionId)
+    if (session) {
+      return {
+        sessionId: session.id,
+        instanceId: preferredInstanceId,
+        parentId: session.parentId ?? null,
+      }
+    }
+  }
+
   const allSessions = sessions()
   for (const [instanceId, sessionMap] of allSessions) {
     const session = sessionMap?.get(sessionId)
@@ -440,7 +452,7 @@ export default function MessageBlock(props: MessageBlockProps) {
                     const hasToolState =
                       Boolean(toolState) && (isToolStateRunning(toolState) || isToolStateCompleted(toolState) || isToolStateError(toolState))
                     const taskSessionId = hasToolState ? extractTaskSessionId(toolState) : ""
-                    const taskLocation = taskSessionId ? findTaskSessionLocation(taskSessionId) : null
+                    const taskLocation = taskSessionId ? findTaskSessionLocation(taskSessionId, props.instanceId) : null
                     const handleGoToTaskSession = (event: MouseEvent) => {
                       event.preventDefault()
                       event.stopPropagation()
diff --git a/packages/ui/src/components/permission-approval-modal.tsx b/packages/ui/src/components/permission-approval-modal.tsx
index 66e79be8..3a010503 100644
--- a/packages/ui/src/components/permission-approval-modal.tsx
+++ b/packages/ui/src/components/permission-approval-modal.tsx
@@ -10,7 +10,7 @@ import {
   setActivePermissionIdForInstance,
   setActiveQuestionIdForInstance,
 } from "../stores/instances"
-import { loadMessages, setActiveSession } from "../stores/sessions"
+import { ensureSessionParentExpanded, loadMessages, sessions as sessionStateSessions, setActiveSessionFromList } from "../stores/sessions"
 import { messageStoreBus } from "../stores/message-v2/bus"
 import ToolCall from "./tool-call"
 
@@ -201,7 +201,14 @@ const PermissionApprovalModal: Component<PermissionApprovalModalProps> = (props)
 
   function handleGoToSession(sessionId: string) {
     if (!sessionId) return
-    setActiveSession(props.instanceId, sessionId)
+
+    const session = sessionStateSessions().get(props.instanceId)?.get(sessionId)
+    const parentId = session?.parentId ?? session?.id
+    if (parentId) {
+      ensureSessionParentExpanded(props.instanceId, parentId)
+    }
+
+    setActiveSessionFromList(props.instanceId, sessionId)
     props.onClose()
   }