From 6c23a8a604f2e1695d3dee7ff86767259d1043e4 Mon Sep 17 00:00:00 2001
From: Jeremy <git@jrmyphlmn.com>
Date: Mon, 13 Jan 2025 16:56:40 +0530
Subject: [PATCH] fix: handle suggestion stream parts (#700)

---
 components/data-stream-handler.tsx | 26 +++++++++++++++++++++++++-
 1 file changed, 25 insertions(+), 1 deletion(-)

diff --git a/components/data-stream-handler.tsx b/components/data-stream-handler.tsx
index 19721a5bf..6c50b77eb 100644
--- a/components/data-stream-handler.tsx
+++ b/components/data-stream-handler.tsx
@@ -1,12 +1,13 @@
 'use client';
 
 import { useChat } from 'ai/react';
-import { useEffect, useRef } from 'react';
+import { useEffect, useRef, useState } from 'react';
 import { BlockKind } from './block';
 import { Suggestion } from '@/lib/db/schema';
 import { initialBlockData, useBlock } from '@/hooks/use-block';
 import { useUserMessageId } from '@/hooks/use-user-message-id';
 import { cx } from 'class-variance-authority';
+import { useSWRConfig } from 'swr';
 
 type DataStreamDelta = {
   type:
@@ -28,6 +29,19 @@ export function DataStreamHandler({ id }: { id: string }) {
   const { setBlock } = useBlock();
   const lastProcessedIndex = useRef(-1);
 
+  const { mutate } = useSWRConfig();
+  const [optimisticSuggestions, setOptimisticSuggestions] = useState<
+    Array<Suggestion>
+  >([]);
+
+  useEffect(() => {
+    if (optimisticSuggestions && optimisticSuggestions.length > 0) {
+      const [optimisticSuggestion] = optimisticSuggestions;
+      const url = `/api/suggestions?documentId=${optimisticSuggestion.documentId}`;
+      mutate(url, optimisticSuggestions, false);
+    }
+  }, [optimisticSuggestions, mutate]);
+
   useEffect(() => {
     if (!dataStream?.length) return;
 
@@ -93,6 +107,16 @@ export function DataStreamHandler({ id }: { id: string }) {
               status: 'streaming',
             };
 
+          case 'suggestion':
+            setTimeout(() => {
+              setOptimisticSuggestions((currentSuggestions) => [
+                ...currentSuggestions,
+                delta.content as Suggestion,
+              ]);
+            }, 0);
+
+            return draftBlock;
+
           case 'clear':
             return {
               ...draftBlock,