3 files changed, 25 insertions, 22 deletions

diff --git a/packages/ui/src/components/message-part.tsx b/packages/ui/src/components/message-part.tsx
index ac80dada7..d0e8afefd 100644
--- a/packages/ui/src/components/message-part.tsx
+++ b/packages/ui/src/components/message-part.tsx
@@ -455,8 +455,8 @@ PART_MAPPING["tool"] = function ToolPartDisplay(props) {
 
   const permission = createMemo(() => {
     const next = data.store.permission?.[props.message.sessionID]?.[0]
-    if (!next) return undefined
-    if (next.callID !== part.callID) return undefined
+    if (!next || !next.tool) return undefined
+    if (next.tool!.callID !== part.callID) return undefined
     return next
   })
 
@@ -732,19 +732,20 @@ ToolRegistry.register({
 
     const childToolPart = createMemo(() => {
       const perm = childPermission()
-      if (!perm) return undefined
+      if (!perm || !perm.tool) return undefined
       const sessionId = childSessionId()
       if (!sessionId) return undefined
       // Find the tool part that matches the permission's callID
       const messages = data.store.message[sessionId] ?? []
-      for (const msg of messages) {
-        const parts = data.store.part[msg.id] ?? []
-        for (const part of parts) {
-          if (part.type === "tool" && (part as ToolPart).callID === perm.callID) {
-            return { part: part as ToolPart, message: msg }
-          }
+      const message = messages.findLast((m) => m.id === perm.tool!.messageID)
+      if (!message) return undefined
+      const parts = data.store.part[message.id] ?? []
+      for (const part of parts) {
+        if (part.type === "tool" && (part as ToolPart).callID === perm.tool!.callID) {
+          return { part: part as ToolPart, message }
         }
       }
+
       return undefined
     })
 
diff --git a/packages/ui/src/components/session-turn.tsx b/packages/ui/src/components/session-turn.tsx
index 0ef1f135c..8285b9822 100644
--- a/packages/ui/src/components/session-turn.tsx
+++ b/packages/ui/src/components/session-turn.tsx
@@ -2,7 +2,7 @@ import {
   AssistantMessage,
   Message as MessageType,
   Part as PartType,
-  type Permission,
+  type PermissionRequest,
   TextPart,
   ToolPart,
 } from "@opencode-ai/sdk/v2/client"
@@ -132,7 +132,7 @@ export function SessionTurn(
   const emptyMessages: MessageType[] = []
   const emptyParts: PartType[] = []
   const emptyAssistant: AssistantMessage[] = []
-  const emptyPermissions: Permission[] = []
+  const emptyPermissions: PermissionRequest[] = []
   const emptyPermissionParts: { part: ToolPart; message: AssistantMessage }[] = []
   const idle = { type: "idle" as const }
 
@@ -235,16 +235,18 @@ export function SessionTurn(
     if (props.stepsExpanded) return emptyPermissionParts
 
     const next = nextPermission()
-    if (!next) return emptyPermissionParts
-
-    for (const message of assistantMessages()) {
-      const parts = data.store.part[message.id] ?? emptyParts
-      for (const part of parts) {
-        if (part?.type !== "tool") continue
-        const tool = part as ToolPart
-        if (tool.callID === next.callID) return [{ part: tool, message }]
-      }
+    if (!next || !next.tool) return emptyPermissionParts
+
+    const message = assistantMessages().findLast((m) => m.id === next.tool!.messageID)
+    if (!message) return emptyPermissionParts
+
+    const parts = data.store.part[message.id] ?? emptyParts
+    for (const part of parts) {
+      if (part?.type !== "tool") continue
+      const tool = part as ToolPart
+      if (tool.callID === next.tool?.callID) return [{ part: tool, message }]
     }
+
     return emptyPermissionParts
   })
 
diff --git a/packages/ui/src/context/data.tsx b/packages/ui/src/context/data.tsx
index 3292ba579..9f7ec813f 100644
--- a/packages/ui/src/context/data.tsx
+++ b/packages/ui/src/context/data.tsx
@@ -1,4 +1,4 @@
-import type { Message, Session, Part, FileDiff, SessionStatus, Permission } from "@opencode-ai/sdk/v2"
+import type { Message, Session, Part, FileDiff, SessionStatus, PermissionRequest } from "@opencode-ai/sdk/v2"
 import { createSimpleContext } from "./helper"
 import { PreloadMultiFileDiffResult } from "@pierre/diffs/ssr"
 
@@ -14,7 +14,7 @@ type Data = {
     [sessionID: string]: PreloadMultiFileDiffResult<any>[]
   }
   permission?: {
-    [sessionID: string]: Permission[]
+    [sessionID: string]: PermissionRequest[]
   }
   message: {
     [sessionID: string]: Message[]