From c7dffb828ac4611aa5dd8bfa41c8af3b2cb79d73 Mon Sep 17 00:00:00 2001
From: Matt Rossman <22670878+mattrossman@users.noreply.github.com>
Date: Mon, 4 May 2026 16:57:40 -0400
Subject: [PATCH] refactor(studio): extract shared AI message utilities and
 simplify (AI-658)

- Add isApprovalContinuation/getLastUserText helpers to message-utils.ts
- Merge duplicate JSX branches in MessagePartExecuteSql
- Fix existingFunction race in EdgeFunctionRenderer (restore eager fetch)
---
 .../AIAssistantPanel/EdgeFunctionRenderer.tsx |  2 +-
 .../ui/AIAssistantPanel/Message.Parts.tsx     | 43 +++++++------------
 apps/studio/lib/ai/message-utils.ts           | 20 ++++++++-
 3 files changed, 36 insertions(+), 29 deletions(-)

diff --git a/apps/studio/components/ui/AIAssistantPanel/EdgeFunctionRenderer.tsx b/apps/studio/components/ui/AIAssistantPanel/EdgeFunctionRenderer.tsx
index bfb76785b72..89958761d2d 100644
--- a/apps/studio/components/ui/AIAssistantPanel/EdgeFunctionRenderer.tsx
+++ b/apps/studio/components/ui/AIAssistantPanel/EdgeFunctionRenderer.tsx
@@ -35,7 +35,7 @@ export const EdgeFunctionRenderer = ({
   const { data: settings } = useProjectSettingsV2Query({ projectRef: ref }, { enabled: !!ref })
   const { data: existingFunction } = useEdgeFunctionQuery(
     { projectRef: ref, slug: functionName },
-    { enabled: !!ref && !!functionName && showApprovalFooter }
+    { enabled: !!ref && !!functionName }
   )
 
   const functionUrl = useMemo(() => {
diff --git a/apps/studio/components/ui/AIAssistantPanel/Message.Parts.tsx b/apps/studio/components/ui/AIAssistantPanel/Message.Parts.tsx
index e5aac34898d..6a9024df564 100644
--- a/apps/studio/components/ui/AIAssistantPanel/Message.Parts.tsx
+++ b/apps/studio/components/ui/AIAssistantPanel/Message.Parts.tsx
@@ -142,8 +142,12 @@ function MessagePartExecuteSql({
   const { data: chart, success } = parseExecuteSqlChartResult(input)
   if (!success) return null
 
-  if (toolPart.state === 'approval-requested') {
-    const approvalId = toolPart.approval.id
+  if (
+    toolPart.state === 'approval-requested' ||
+    state === 'input-available' ||
+    state === 'output-available'
+  ) {
+    const approvalId = toolPart.state === 'approval-requested' ? toolPart.approval.id : undefined
     return (
       <div className="w-auto overflow-x-hidden my-4 space-y-2">
         <DisplayBlockRenderer
@@ -161,31 +165,16 @@ function MessagePartExecuteSql({
           toolState={toolPart.state}
           isLastPart={isLastPart}
           isLastMessage={isLastMessage}
-          onApprove={() => addToolApprovalResponse?.({ id: approvalId, approved: true })}
-          onDeny={() => addToolApprovalResponse?.({ id: approvalId, approved: false })}
-        />
-      </div>
-    )
-  }
-
-  if (state === 'input-available' || state === 'output-available') {
-    return (
-      <div className="w-auto overflow-x-hidden my-4 space-y-2">
-        <DisplayBlockRenderer
-          messageId={id}
-          toolCallId={toolCallId}
-          initialArgs={{
-            sql: chart.sql,
-            label: chart.label,
-            isWriteQuery: chart.isWriteQuery,
-            view: chart.view,
-            xAxis: chart.xAxis,
-            yAxis: chart.yAxis,
-          }}
-          initialResults={output}
-          toolState={state}
-          isLastPart={isLastPart}
-          isLastMessage={isLastMessage}
+          onApprove={
+            approvalId
+              ? () => addToolApprovalResponse?.({ id: approvalId, approved: true })
+              : undefined
+          }
+          onDeny={
+            approvalId
+              ? () => addToolApprovalResponse?.({ id: approvalId, approved: false })
+              : undefined
+          }
         />
       </div>
     )
diff --git a/apps/studio/lib/ai/message-utils.ts b/apps/studio/lib/ai/message-utils.ts
index 2b0cbd62924..22287b495c6 100644
--- a/apps/studio/lib/ai/message-utils.ts
+++ b/apps/studio/lib/ai/message-utils.ts
@@ -1,4 +1,22 @@
-import type { UIMessage } from 'ai'
+import { isToolUIPart, type UIMessage } from 'ai'
+
+export function isApprovalContinuation(messages: UIMessage[]): boolean {
+  return messages.some(
+    (msg) =>
+      msg.role === 'assistant' &&
+      msg.parts?.some((part) => isToolUIPart(part) && part.state === 'approval-responded')
+  )
+}
+
+export function getLastUserText(messages: UIMessage[]): string {
+  const lastUserMessage = messages.findLast((m) => m.role === 'user')
+  return (
+    lastUserMessage?.parts
+      ?.filter((p): p is { type: 'text'; text: string } => p.type === 'text')
+      .map((p) => p.text)
+      .join('\n') ?? ''
+  )
+}
 
 /**
  * Prepares messages for API transmission by cleaning and limiting history