fix: restore code block loading states (#728)

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,10 +250,12 @@ function PureBlock({
 
   useEffect(() => {
     if (block && block.documentId !== 'init') {
-      blockDefinition.initialize({
-        documentId: block.documentId,
-        setMetadata,
-      });
+      if (blockDefinition.initialize) {
+        blockDefinition.initialize({
+          documentId: block.documentId,
+          setMetadata,
+        });
+      }
     }
   }, [block, blockDefinition, setMetadata]);
 
@@ -451,55 +453,40 @@ function PureBlock({
               />
             </div>
 
-            <div
-              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={
-                    isCurrentVersion
-                      ? block.content
-                      : getDocumentContentById(currentVersionIndex)
-                  }
-                  mode={mode}
-                  status={block.status}
-                  currentVersionIndex={currentVersionIndex}
-                  suggestions={[]}
-                  onSaveContent={saveContent}
-                  isInline={false}
-                  isCurrentVersion={isCurrentVersion}
-                  getDocumentContentById={getDocumentContentById}
-                  isLoading={isDocumentsFetching && !block.content}
-                  metadata={metadata}
-                  setMetadata={setMetadata}
-                />
+            <div className="dark:bg-muted bg-background h-full overflow-y-scroll !max-w-full items-center">
+              <blockDefinition.content
+                title={block.title}
+                content={
+                  isCurrentVersion
+                    ? block.content
+                    : getDocumentContentById(currentVersionIndex)
+                }
+                mode={mode}
+                status={block.status}
+                currentVersionIndex={currentVersionIndex}
+                suggestions={[]}
+                onSaveContent={saveContent}
+                isInline={false}
+                isCurrentVersion={isCurrentVersion}
+                getDocumentContentById={getDocumentContentById}
+                isLoading={isDocumentsFetching && !block.content}
+                metadata={metadata}
+                setMetadata={setMetadata}
+              />
 
-                <AnimatePresence>
-                  {isCurrentVersion && (
-                    <Toolbar
-                      isToolbarVisible={isToolbarVisible}
-                      setIsToolbarVisible={setIsToolbarVisible}
-                      append={append}
-                      isLoading={isLoading}
-                      stop={stop}
-                      setMessages={setMessages}
-                      blockKind={block.kind}
-                    />
-                  )}
-                </AnimatePresence>
-              </div>
+              <AnimatePresence>
+                {isCurrentVersion && (
+                  <Toolbar
+                    isToolbarVisible={isToolbarVisible}
+                    setIsToolbarVisible={setIsToolbarVisible}
+                    append={append}
+                    isLoading={isLoading}
+                    stop={stop}
+                    setMessages={setMessages}
+                    blockKind={block.kind}
+                  />
+                )}
+              </AnimatePresence>
             </div>
 
             <AnimatePresence>

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;
-  setMetadata: Dispatch<SetStateAction<any>>;
+  metadata: M;
+  setMetadata: Dispatch<SetStateAction<M>>;
 };
 
-type BlockAction = {
+type BlockAction<M = any> = {
   icon: ReactNode;
   label?: string;
   description: string;
-  onClick: (context: BlockActionContext) => void;
-  isDisabled?: (context: BlockActionContext) => boolean;
+  onClick: (context: BlockActionContext<M>) => Promise<void> | void;
+  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;
-  setMetadata: Dispatch<SetStateAction<any>>;
-};
+  metadata: M;
+  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<
-    Omit<BlockContent, 'metadata' | 'setMetadata'> & {
-      metadata: M;
-      setMetadata: Dispatch<SetStateAction<M>>;
-    }
-  >;
-  actions?: BlockAction[];
-  toolbar?: BlockToolbarItem[];
-  metadata?: M;
+  content: ComponentType<BlockContent<M>>;
+  actions: Array<BlockAction<M>>;
+  toolbar: BlockToolbarItem[];
   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 actions: BlockAction[];
+  readonly content: ComponentType<BlockContent<M>>;
+  readonly actions: Array<BlockAction<M>>;
   readonly toolbar: BlockToolbarItem[];
-  readonly metadata: M;
-  readonly initialize: (parameters: InitializeParameters) => void;
-  readonly onStreamPart?: (args: {
+  readonly initialize?: (parameters: InitializeParameters) => void;
+  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;
   }