fix: restore code block loading states (#728)

vercel · Jan 27, 2025 · 085f4a8 · 085f4a8
1 parent 38527ff
commit 085f4a8
Show file tree

Hide file tree

Showing 6 changed files with 267 additions and 134 deletions.
diff --git a/blocks/code.tsx b/blocks/code.tsx
@@ -10,7 +10,57 @@ import {
 } from '@/components/icons';
 import { toast } from 'sonner';
 import { generateUUID } from '@/lib/utils';
-import { Console, ConsoleOutput } from '@/components/console';
+import {
+  Console,
+  ConsoleOutput,
+  ConsoleOutputContent,
+} from '@/components/console';
+
+const OUTPUT_HANDLERS = {
+  matplotlib: `
+    import io
+    import base64
+    from matplotlib import pyplot as plt
+
+    # Clear any existing plots
+    plt.clf()
+    plt.close('all')
+
+    # Switch to agg backend
+    plt.switch_backend('agg')
+
+    def setup_matplotlib_output():
+        def custom_show():
+            if plt.gcf().get_size_inches().prod() * plt.gcf().dpi ** 2 > 25_000_000:
+                print("Warning: Plot size too large, reducing quality")
+                plt.gcf().set_dpi(100)
+
+            png_buf = io.BytesIO()
+            plt.savefig(png_buf, format='png')
+            png_buf.seek(0)
+            png_base64 = base64.b64encode(png_buf.read()).decode('utf-8')
+            print(f'data:image/png;base64,{png_base64}')
+            png_buf.close()
+
+            plt.clf()
+            plt.close('all')
+
+        plt.show = custom_show
+  `,
+  basic: `
+    # Basic output capture setup
+  `,
+};
+
+function detectRequiredHandlers(code: string): string[] {
+  const handlers: string[] = ['basic'];
+
+  if (code.includes('matplotlib') || code.includes('plt.')) {
+    handlers.push('matplotlib');
+  }
+
+  return handlers;
+}
 
 interface Metadata {
   outputs: Array<ConsoleOutput>;
@@ -20,9 +70,11 @@ export const codeBlock = new Block<'code', Metadata>({
   kind: 'code',
   description:
     'Useful for code generation; Code execution is only available for python code.',
-  initialize: () => ({
-    outputs: [],
-  }),
+  initialize: async ({ setMetadata }) => {
+    setMetadata({
+      outputs: [],
+    });
+  },
   onStreamPart: ({ streamPart, setBlock }) => {
     if (streamPart.type === 'code-delta') {
       setBlock((draftBlock) => ({
@@ -41,7 +93,9 @@ export const codeBlock = new Block<'code', Metadata>({
   content: ({ metadata, setMetadata, ...props }) => {
     return (
       <>
-        <CodeEditor {...props} />
+        <div className="px-1">
+          <CodeEditor {...props} />
+        </div>
 
         {metadata?.outputs && (
           <Console
@@ -64,46 +118,94 @@ export const codeBlock = new Block<'code', Metadata>({
       description: 'Execute code',
       onClick: async ({ content, setMetadata }) => {
         const runId = generateUUID();
-        const outputs: any[] = [];
-
-        // @ts-expect-error - loadPyodide is not defined
-        const currentPyodideInstance = await globalThis.loadPyodide({
-          indexURL: 'https://cdn.jsdelivr.net/pyodide/v0.23.4/full/',
-        });
+        const outputContent: Array<ConsoleOutputContent> = [];
 
-        currentPyodideInstance.setStdout({
-          batched: (output: string) => {
-            outputs.push({
+        setMetadata((metadata) => ({
+          ...metadata,
+          outputs: [
+            ...metadata.outputs,
+            {
               id: runId,
-              contents: [
-                {
-                  type: output.startsWith('data:image/png;base64')
-                    ? 'image'
-                    : 'text',
-                  value: output,
-                },
-              ],
-              status: 'completed',
-            });
-          },
-        });
+              contents: [],
+              status: 'in_progress',
+            },
+          ],
+        }));
 
-        await currentPyodideInstance.loadPackagesFromImports(content, {
-          messageCallback: (message: string) => {
-            outputs.push({
-              id: runId,
-              contents: [{ type: 'text', value: message }],
-              status: 'loading_packages',
-            });
-          },
-        });
+        try {
+          // @ts-expect-error - loadPyodide is not defined
+          const currentPyodideInstance = await globalThis.loadPyodide({
+            indexURL: 'https://cdn.jsdelivr.net/pyodide/v0.23.4/full/',
+          });
 
-        await currentPyodideInstance.runPythonAsync(content);
+          currentPyodideInstance.setStdout({
+            batched: (output: string) => {
+              outputContent.push({
+                type: output.startsWith('data:image/png;base64')
+                  ? 'image'
+                  : 'text',
+                value: output,
+              });
+            },
+          });
 
-        setMetadata((metadata: any) => ({
-          ...metadata,
-          outputs,
-        }));
+          await currentPyodideInstance.loadPackagesFromImports(content, {
+            messageCallback: (message: string) => {
+              setMetadata((metadata) => ({
+                ...metadata,
+                outputs: [
+                  ...metadata.outputs.filter((output) => output.id !== runId),
+                  {
+                    id: runId,
+                    contents: [{ type: 'text', value: message }],
+                    status: 'loading_packages',
+                  },
+                ],
+              }));
+            },
+          });
+
+          const requiredHandlers = detectRequiredHandlers(content);
+          for (const handler of requiredHandlers) {
+            if (OUTPUT_HANDLERS[handler as keyof typeof OUTPUT_HANDLERS]) {
+              await currentPyodideInstance.runPythonAsync(
+                OUTPUT_HANDLERS[handler as keyof typeof OUTPUT_HANDLERS],
+              );
+
+              if (handler === 'matplotlib') {
+                await currentPyodideInstance.runPythonAsync(
+                  'setup_matplotlib_output()',
+                );
+              }
+            }
+          }
+
+          await currentPyodideInstance.runPythonAsync(content);
+
+          setMetadata((metadata) => ({
+            ...metadata,
+            outputs: [
+              ...metadata.outputs.filter((output) => output.id !== runId),
+              {
+                id: runId,
+                contents: outputContent,
+                status: 'completed',
+              },
+            ],
+          }));
+        } catch (error: any) {
+          setMetadata((metadata) => ({
+            ...metadata,
+            outputs: [
+              ...metadata.outputs.filter((output) => output.id !== runId),
+              {
+                id: runId,
+                contents: [{ type: 'text', value: error.message }],
+                status: 'failed',
+              },
+            ],
+          }));
+        }
       },
     },
     {
@@ -112,13 +214,27 @@ export const codeBlock = new Block<'code', Metadata>({
       onClick: ({ handleVersionChange }) => {
         handleVersionChange('prev');
       },
+      isDisabled: ({ currentVersionIndex }) => {
+        if (currentVersionIndex === 0) {
+          return true;
+        }
+
+        return false;
+      },
     },
     {
       icon: <RedoIcon size={18} />,
       description: 'View Next version',
       onClick: ({ handleVersionChange }) => {
         handleVersionChange('next');
       },
+      isDisabled: ({ isCurrentVersion }) => {
+        if (isCurrentVersion) {
+          return true;
+        }
+
+        return false;
+      },
     },
     {
       icon: <CopyIcon size={18} />,

diff --git a/blocks/image.tsx b/blocks/image.tsx
@@ -1,6 +1,7 @@
 import { Block } from '@/components/create-block';
 import { CopyIcon, RedoIcon, UndoIcon } from '@/components/icons';
 import { ImageEditor } from '@/components/image-editor';
+import { toast } from 'sonner';
 
 export const imageBlock = new Block({
   kind: 'image',
@@ -23,13 +24,27 @@ export const imageBlock = new Block({
       onClick: ({ handleVersionChange }) => {
         handleVersionChange('prev');
       },
+      isDisabled: ({ currentVersionIndex }) => {
+        if (currentVersionIndex === 0) {
+          return true;
+        }
+
+        return false;
+      },
     },
     {
       icon: <RedoIcon size={18} />,
       description: 'View Next version',
       onClick: ({ handleVersionChange }) => {
         handleVersionChange('next');
       },
+      isDisabled: ({ isCurrentVersion }) => {
+        if (isCurrentVersion) {
+          return true;
+        }
+
+        return false;
+      },
     },
     {
       icon: <CopyIcon size={18} />,
@@ -52,6 +67,8 @@ export const imageBlock = new Block({
             }
           }, 'image/png');
         };
+
+        toast.success('Copied image to clipboard!');
       },
     },
   ],

diff --git a/blocks/text.tsx b/blocks/text.tsx
@@ -80,18 +80,22 @@ export const textBlock = new Block<'text', TextBlockMetadata>({
 
     return (
       <>
-        <Editor
-          content={content}
-          suggestions={metadata ? metadata.suggestions : []}
-          isCurrentVersion={isCurrentVersion}
-          currentVersionIndex={currentVersionIndex}
-          status={status}
-          onSaveContent={onSaveContent}
-        />
+        <div className="flex flex-row py-8 md:p-20 px-4">
+          <Editor
+            content={content}
+            suggestions={metadata ? metadata.suggestions : []}
+            isCurrentVersion={isCurrentVersion}
+            currentVersionIndex={currentVersionIndex}
+            status={status}
+            onSaveContent={onSaveContent}
+          />
 
-        {metadata && metadata.suggestions && metadata.suggestions.length > 0 ? (
-          <div className="md:hidden h-dvh w-12 shrink-0" />
-        ) : null}
+          {metadata &&
+          metadata.suggestions &&
+          metadata.suggestions.length > 0 ? (
+            <div className="md:hidden h-dvh w-12 shrink-0" />
+          ) : null}
+        </div>
       </>
     );
   },

diff --git a/components/block-actions.tsx b/components/block-actions.tsx
@@ -1,9 +1,10 @@
 import { Button } from './ui/button';
 import { Tooltip, TooltipContent, TooltipTrigger } from './ui/tooltip';
 import { blockDefinitions, UIBlock } from './block';
-import { Dispatch, memo, SetStateAction } from 'react';
+import { Dispatch, memo, SetStateAction, useState } from 'react';
 import { BlockActionContext } from './create-block';
 import { cn } from '@/lib/utils';
+import { toast } from 'sonner';
 
 interface BlockActionsProps {
   block: UIBlock;
@@ -24,6 +25,8 @@ function PureBlockActions({
   metadata,
   setMetadata,
 }: BlockActionsProps) {
+  const [isLoading, setIsLoading] = useState(false);
+
   const blockDefinition = blockDefinitions.find(
     (definition) => definition.kind === block.kind,
   );
@@ -53,9 +56,23 @@ function PureBlockActions({
                 'p-2': !action.label,
                 'py-1.5 px-2': action.label,
               })}
-              onClick={() => action.onClick(actionContext)}
+              onClick={async () => {
+                setIsLoading(true);
+
+                try {
+                  await Promise.resolve(action.onClick(actionContext));
+                } catch (error) {
+                  toast.error('Failed to execute action');
+                } finally {
+                  setIsLoading(false);
+                }
+              }}
               disabled={
-                action.isDisabled ? action.isDisabled(actionContext) : false
+                isLoading || block.status === 'streaming'
+                  ? true
+                  : action.isDisabled
+                    ? action.isDisabled(actionContext)
+                    : false
               }
             >
               {action.icon}