fix: restore code block loading states (#728)

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: () => ({
+  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 (
       <>
+        <div className="px-1">
           <CodeEditor {...props} />
+        </div>
 
         {metadata?.outputs && (
           <Console
@@ -64,8 +118,21 @@ export const codeBlock = new Block<'code', Metadata>({
       description: 'Execute code',
       onClick: async ({ content, setMetadata }) => {
         const runId = generateUUID();
-        const outputs: any[] = [];
+        const outputContent: Array<ConsoleOutputContent> = [];
 
+        setMetadata((metadata) => ({
+          ...metadata,
+          outputs: [
+            ...metadata.outputs,
+            {
+              id: runId,
+              contents: [],
+              status: 'in_progress',
+            },
+          ],
+        }));
+
+        try {
           // @ts-expect-error - loadPyodide is not defined
           const currentPyodideInstance = await globalThis.loadPyodide({
             indexURL: 'https://cdn.jsdelivr.net/pyodide/v0.23.4/full/',
@@ -73,37 +140,72 @@ export const codeBlock = new Block<'code', Metadata>({
 
           currentPyodideInstance.setStdout({
             batched: (output: string) => {
-            outputs.push({
+              outputContent.push({
-              id: runId,
-              contents: [
-                {
                 type: output.startsWith('data:image/png;base64')
                   ? 'image'
                   : 'text',
                 value: output,
-                },
-              ],
-              status: 'completed',
               });
             },
           });
 
           await currentPyodideInstance.loadPackagesFromImports(content, {
             messageCallback: (message: string) => {
-            outputs.push({
+              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: any) => ({
+          setMetadata((metadata) => ({
             ...metadata,
-          outputs,
+            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,6 +214,13 @@ export const codeBlock = new Block<'code', Metadata>({
       onClick: ({ handleVersionChange }) => {
         handleVersionChange('prev');
       },
+      isDisabled: ({ currentVersionIndex }) => {
+        if (currentVersionIndex === 0) {
+          return true;
+        }
+
+        return false;
+      },
     },
     {
       icon: <RedoIcon size={18} />,
@@ -119,6 +228,13 @@ export const codeBlock = new Block<'code', Metadata>({
       onClick: ({ handleVersionChange }) => {
         handleVersionChange('next');
       },
+      isDisabled: ({ isCurrentVersion }) => {
+        if (isCurrentVersion) {
+          return true;
+        }
+
+        return false;
+      },
     },
     {
       icon: <CopyIcon size={18} />,

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,6 +24,13 @@ export const imageBlock = new Block({
       onClick: ({ handleVersionChange }) => {
         handleVersionChange('prev');
       },
+      isDisabled: ({ currentVersionIndex }) => {
+        if (currentVersionIndex === 0) {
+          return true;
+        }
+
+        return false;
+      },
     },
     {
       icon: <RedoIcon size={18} />,
@@ -30,6 +38,13 @@ export const imageBlock = new Block({
       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!');
       },
     },
   ],

blocks/text.tsx

@@ -80,6 +80,7 @@ export const textBlock = new Block<'text', TextBlockMetadata>({
 
     return (
       <>
+        <div className="flex flex-row py-8 md:p-20 px-4">
           <Editor
             content={content}
             suggestions={metadata ? metadata.suggestions : []}
@@ -89,9 +90,12 @@ export const textBlock = new Block<'text', TextBlockMetadata>({
             onSaveContent={onSaveContent}
           />
 
-        {metadata && metadata.suggestions && metadata.suggestions.length > 0 ? (
+          {metadata &&
+          metadata.suggestions &&
+          metadata.suggestions.length > 0 ? (
             <div className="md:hidden h-dvh w-12 shrink-0" />
           ) : null}
+        </div>
       </>
     );
   },

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}

components/block.tsx

@@ -17,7 +17,7 @@ import {
 import useSWR, { useSWRConfig } from 'swr';
 import { useDebounceCallback, useWindowSize } from 'usehooks-ts';
 import type { Document, Vote } from '@/lib/db/schema';
-import { cn, fetcher } from '@/lib/utils';
+import { fetcher } from '@/lib/utils';
 import { MultimodalInput } from './multimodal-input';
 import { Toolbar } from './toolbar';
 import { VersionFooter } from './version-footer';
@@ -26,10 +26,10 @@ import { BlockCloseButton } from './block-close-button';
 import { BlockMessages } from './block-messages';
 import { useSidebar } from './ui/sidebar';
 import { useBlock } from '@/hooks/use-block';
-import equal from 'fast-deep-equal';
 import { textBlock } from '@/blocks/text';
 import { imageBlock } from '@/blocks/image';
 import { codeBlock } from '@/blocks/code';
+import equal from 'fast-deep-equal';
 
 export const blockDefinitions = [textBlock, codeBlock, imageBlock] as const;
 export type BlockKind = (typeof blockDefinitions)[number]['kind'];
@@ -250,11 +250,13 @@ function PureBlock({
 
   useEffect(() => {
     if (block && block.documentId !== 'init') {
+      if (blockDefinition.initialize) {
         blockDefinition.initialize({
           documentId: block.documentId,
           setMetadata,
         });
       }
+    }
   }, [block, blockDefinition, setMetadata]);
 
   return (
@@ -451,21 +453,7 @@ function PureBlock({
               />
             </div>
 
-            <div
+            <div className="dark:bg-muted bg-background h-full overflow-y-scroll !max-w-full items-center">
-              className={cn(
-                'dark:bg-muted bg-background h-full overflow-y-scroll !max-w-full pb-40 items-center',
-                {
-                  '': block.kind === 'code',
-                  'py-8 md:p-20 px-4': block.kind === 'text',
-                },
-              )}
-            >
-              <div
-                className={cn('flex flex-row', {
-                  '': block.kind === 'code',
-                  'mx-auto max-w-[600px]': block.kind === 'text',
-                })}
-              >
               <blockDefinition.content
                 title={block.title}
                 content={
@@ -500,7 +488,6 @@ function PureBlock({
                 )}
               </AnimatePresence>
             </div>
-            </div>
 
             <AnimatePresence>
               {!isCurrentVersion && (

components/create-block.tsx

@@ -4,22 +4,22 @@ import { ComponentType, Dispatch, ReactNode, SetStateAction } from 'react';
 import { DataStreamDelta } from './data-stream-handler';
 import { UIBlock } from './block';
 
-export type BlockActionContext = {
+export type BlockActionContext<M = any> = {
   content: string;
   handleVersionChange: (type: 'next' | 'prev' | 'toggle' | 'latest') => void;
   currentVersionIndex: number;
   isCurrentVersion: boolean;
   mode: 'edit' | 'diff';
-  metadata: any;
+  metadata: M;
-  setMetadata: Dispatch<SetStateAction<any>>;
+  setMetadata: Dispatch<SetStateAction<M>>;
 };
 
-type BlockAction = {
+type BlockAction<M = any> = {
   icon: ReactNode;
   label?: string;
   description: string;
-  onClick: (context: BlockActionContext) => void;
+  onClick: (context: BlockActionContext<M>) => Promise<void> | void;
-  isDisabled?: (context: BlockActionContext) => boolean;
+  isDisabled?: (context: BlockActionContext<M>) => boolean;
 };
 
 export type BlockToolbarContext = {
@@ -32,7 +32,7 @@ export type BlockToolbarItem = {
   onClick: (context: BlockToolbarContext) => void;
 };
 
-type BlockContent = {
+interface BlockContent<M = any> {
   title: string;
   content: string;
   mode: 'edit' | 'diff';
@@ -44,9 +44,9 @@ type BlockContent = {
   isInline: boolean;
   getDocumentContentById: (index: number) => string;
   isLoading: boolean;
-  metadata: any;
+  metadata: M;
-  setMetadata: Dispatch<SetStateAction<any>>;
+  setMetadata: Dispatch<SetStateAction<M>>;
-};
+}
 
 interface InitializeParameters<M = any> {
   documentId: string;
@@ -56,17 +56,11 @@ interface InitializeParameters<M = any> {
 type BlockConfig<T extends string, M = any> = {
   kind: T;
   description: string;
-  content: ComponentType<
+  content: ComponentType<BlockContent<M>>;
-    Omit<BlockContent, 'metadata' | 'setMetadata'> & {
+  actions: Array<BlockAction<M>>;
-      metadata: M;
+  toolbar: BlockToolbarItem[];
-      setMetadata: Dispatch<SetStateAction<M>>;
-    }
-  >;
-  actions?: BlockAction[];
-  toolbar?: BlockToolbarItem[];
-  metadata?: M;
   initialize?: (parameters: InitializeParameters<M>) => void;
-  onStreamPart?: (args: {
+  onStreamPart: (args: {
     setMetadata: Dispatch<SetStateAction<M>>;
     setBlock: Dispatch<SetStateAction<UIBlock>>;
     streamPart: DataStreamDelta;
@@ -76,12 +70,11 @@ type BlockConfig<T extends string, M = any> = {
 export class Block<T extends string, M = any> {
   readonly kind: T;
   readonly description: string;
-  readonly content: ComponentType<BlockContent>;
+  readonly content: ComponentType<BlockContent<M>>;
-  readonly actions: BlockAction[];
+  readonly actions: Array<BlockAction<M>>;
   readonly toolbar: BlockToolbarItem[];
-  readonly metadata: M;
+  readonly initialize?: (parameters: InitializeParameters) => void;
-  readonly initialize: (parameters: InitializeParameters) => void;
+  readonly onStreamPart: (args: {
-  readonly onStreamPart?: (args: {
     setMetadata: Dispatch<SetStateAction<M>>;
     setBlock: Dispatch<SetStateAction<UIBlock>>;
     streamPart: DataStreamDelta;
@@ -93,7 +86,6 @@ export class Block<T extends string, M = any> {
     this.content = config.content;
     this.actions = config.actions || [];
     this.toolbar = config.toolbar || [];
-    this.metadata = config.metadata as M;
     this.initialize = config.initialize || (async () => ({}));
     this.onStreamPart = config.onStreamPart;
   }