diff --git a/apps/src/base/appController.ts b/apps/src/base/appController.ts
index 4aae122..2790104 100644
--- a/apps/src/base/appController.ts
+++ b/apps/src/base/appController.ts
@@ -26,6 +26,10 @@ export function Action(metadata: ActionMetadata) {
     Reflect.defineMetadata('actionMetadata', metadata, target, propertyKey);
   };
 }
+
+// export type ActionReturn = string | void | Promise<string | void>
+// export type CheckMethods<T> = { [K in keyof T]: T[K] extends Function ?
+//   ((...args: any) => ActionReturn) : T[K] }
 export abstract class AppController<T> {
   protected app: App<T>;
 
@@ -109,8 +113,26 @@ export abstract class AppController<T> {
   }
 
   async runAction(fn: string, args: any) {
+    let renderInfo = {
+      text: null,
+      code: null,
+      oldCode: null
+    }
+    // get render info if it exists
+    const metadata = Reflect.getMetadata('actionMetadata', this, fn);
+    if (metadata) {
+      let renderBody = metadata['renderBody']
+      if (typeof renderBody === 'function') {
+        renderInfo = await renderBody(args, await this.app.getState())
+      }
+    }
     // @ts-ignore: Check if controller has function and execute!
-    return await this[fn](args);
+    let result = await this[fn](args);
+    return {
+      type: 'BLANK',
+      ...result,
+      renderInfo
+    }
   }
 }
 
diff --git a/apps/src/metabase/appController.ts b/apps/src/metabase/appController.ts
index 82d5ccb..8631206 100644
--- a/apps/src/metabase/appController.ts
+++ b/apps/src/metabase/appController.ts
@@ -41,8 +41,9 @@ export class MetabaseController extends AppController<MetabaseAppState> {
     labelRunning: "Updating SQL query",
     labelDone: "Updated query",
     description: "Updates the SQL query in the Metabase SQL editor and executes it.",
-    renderBody: ({ sql }: { sql: string }) => {
-      return {text: null, code: sql}
+    renderBody: ({ sql }: { sql: string }, appState: MetabaseAppStateSQLEditor) => {
+      const sqlQuery = appState?.sqlQuery
+      return {text: null, code: sql, oldCode: sqlQuery}
     }
   })
   async updateSQLQuery({ sql, executeImmediately = true }: { sql: string, executeImmediately?: boolean }) {
diff --git a/web/src/components/common/ActionStack.tsx b/web/src/components/common/ActionStack.tsx
index da67b3e..22733d2 100644
--- a/web/src/components/common/ActionStack.tsx
+++ b/web/src/components/common/ActionStack.tsx
@@ -15,6 +15,7 @@ import { getApp } from "../../helpers/app";
 import 'reflect-metadata';
 import { parseArguments } from '../../planner/plannerActions';
 import { CodeBlock } from './CodeBlock';
+import { ActionRenderInfo } from '../../state/chat/reducer';
 
 function removeThinkingTags(input: string): string {
   return input ? input.replace(/<thinking>[\s\S]*?<\/thinking>/g, '') : input;
@@ -25,7 +26,10 @@ function extractMessageContent(input: string): string {
   return match ? match[1] : "";
 }
 
-export type ActionStatusView = Pick<Action, 'finished' | 'function' | 'status'>
+export type ActionStatusView = Pick<Action, 'finished' | 'function' | 'status'> & {
+  renderInfo: ActionRenderInfo
+}
+
 export const ActionStack: React.FC<{status: string, actions: Array<ActionStatusView>, index:number, content: string, latency: number}> = ({
   actions,
   status,
@@ -46,16 +50,6 @@ export const ActionStack: React.FC<{status: string, actions: Array<ActionStatusV
     return action;
   }
 
-  const renderActionArgs = (action: string, args: string) => {
-    if (controller) {
-      const metadata = Reflect.getMetadata('actionMetadata', controller, action);
-      if (metadata) {
-        return metadata['renderBody'](parseArguments(args, action));
-      }
-    }
-    return {text: null, code: null};
-  }
-  
   let title: string = "";
   if (status == 'FINISHED') {
     let titles = actions.map(action => getActionLabels(action.function.name, 'labelDone')) 
@@ -116,7 +110,7 @@ export const ActionStack: React.FC<{status: string, actions: Array<ActionStatusV
           
         </HStack>
         {isExpanded && actions.map((action, index) => {
-          const { text, code } = renderActionArgs(action.function.name, action.function.arguments)
+          const { text, code, oldCode } = action.renderInfo || {text: null, code: null, oldCode: undefined}
           return (
           <VStack className={'action'} padding={'2px'} key={index} alignItems={"start"}>
             <HStack>
@@ -132,7 +126,7 @@ export const ActionStack: React.FC<{status: string, actions: Array<ActionStatusV
             </HStack>
             
             { code && <Box width={"100%"} p={2} bg={"#1e1e1e"} borderRadius={5}>
-              <CodeBlock code={code || ""} tool={currentTool} oldCode={undefined}/>
+              <CodeBlock code={code || ""} tool={currentTool} oldCode={oldCode || undefined}/>
              </Box>
             }
             
diff --git a/web/src/components/common/Chat.tsx b/web/src/components/common/Chat.tsx
index 2b87ba1..e56b627 100644
--- a/web/src/components/common/Chat.tsx
+++ b/web/src/components/common/Chat.tsx
@@ -10,7 +10,10 @@ import { ActionStack, ActionStatusView, OngoingActionStack } from './ActionStack
 import { ChatContent } from './ChatContent';
 import { getApp } from '../../helpers/app';
 
-function addStatusInfoToActionPlanMessages(messages: Array<ChatMessage>) {
+// adds tool information like execution status and rendering info
+// this stuff is in the 'tool' messages, but we're ony rendering 'assistant' messages
+// so this copy needs to be done while rendering.
+function addToolInfoToActionPlanMessages(messages: Array<ChatMessage>) {
   const toolMessages = messages.filter(message => message.role == 'tool') as Array<ActionChatMessage>
   const toolMessageMap = new Map(toolMessages.map((message: ActionChatMessage) => [message.action.id, message]))
   return messages.map(message => {
@@ -20,7 +23,8 @@ function addStatusInfoToActionPlanMessages(messages: Array<ChatMessage>) {
         if (toolMessage) {
           return {
             ...toolCall,
-            status: toolMessage.action.status
+            status: toolMessage.action.status,
+            renderInfo: toolMessage.content.renderInfo
           }
         } else {
           return toolCall
@@ -39,7 +43,7 @@ function addStatusInfoToActionPlanMessages(messages: Array<ChatMessage>) {
   })
 }
 
-const Chat: React.FC<ReturnType<typeof addStatusInfoToActionPlanMessages>[number]> = ({
+const Chat: React.FC<ReturnType<typeof addToolInfoToActionPlanMessages>[number]> = ({
   index,
   role,
   content,
@@ -63,7 +67,8 @@ const Chat: React.FC<ReturnType<typeof addStatusInfoToActionPlanMessages>[number
       actions.push({
         finished: true,
         function: toolCall.function,
-        status: toolCall.status
+        status: toolCall.status,
+        renderInfo: toolCall.renderInfo
       })
     })
     const latency = ('latency' in debug)? Math.round(debug.latency as number /100)/10 : 0
@@ -175,7 +180,7 @@ export const ChatSection = () => {
   // just create a map of all role='tool' messages by their id, and for each
   // tool call in each assistant message, add the status from the corresponding
   // tool message
-  const messagesWithStatus = addStatusInfoToActionPlanMessages(messages)
+  const messagesWithStatus = addToolInfoToActionPlanMessages(messages)
   const Chats = isEmpty(messagesWithStatus) ?
     <HelperMessage /> :
     messagesWithStatus.map((message, key) => (<Chat key={key} {...message} />))
diff --git a/web/src/state/chat/reducer.ts b/web/src/state/chat/reducer.ts
index fd28bd9..69882c3 100644
--- a/web/src/state/chat/reducer.ts
+++ b/web/src/state/chat/reducer.ts
@@ -83,7 +83,16 @@ export interface ActionPlanChatMessage extends BaseChatMessage {
   content: ActionPlanMessageContent
 }
 
-export type ActionChatMessageContent = Subset<ChatMessageContent, DefaultMessageContent | BlankMessageContent>
+
+export type ActionRenderInfo = {
+  text: string | null,
+  code: string | null,
+  oldCode?: string | null
+}
+
+export type ActionChatMessageContent = Subset<ChatMessageContent, DefaultMessageContent | BlankMessageContent> & {
+  renderInfo: ActionRenderInfo
+}
 export interface ActionChatMessage extends BaseChatMessage {
   role: 'tool'
   action: Action
@@ -220,6 +229,11 @@ export const chatSlice = createSlice({
           index: actionMessageID,
           content: {
             type: 'BLANK',
+            renderInfo: {
+              text: null,
+              code: null,
+              oldCode: null
+            }
           },
           createdAt: timestamp,
           updatedAt: timestamp,