From 34cef2b370ac13844b56cf94b271ba4cdfdee644 Mon Sep 17 00:00:00 2001
From: liangfung <1098486429@qq.com>
Date: Mon, 1 Apr 2024 18:52:49 +0800
Subject: [PATCH] iframe

---
 .../action-bar-widget-extension.tsx           |   7 +-
 .../app/files/components/completion-panel.tsx | 151 +++++++-----------
 .../components/file-directory-breadcrumb.tsx  |  64 ++++----
 .../files/components/source-code-browser.tsx  |  53 ++++--
 ee/tabby-ui/app/files/lib/event-emitter.ts    |   2 +-
 .../app/playground/components/chats.tsx       |  31 +++-
 ee/tabby-ui/components/chat-panel.tsx         |  23 ++-
 ee/tabby-ui/components/chat.tsx               |   1 +
 8 files changed, 182 insertions(+), 150 deletions(-)

diff --git a/ee/tabby-ui/app/files/components/action-bar-widget/action-bar-widget-extension.tsx b/ee/tabby-ui/app/files/components/action-bar-widget/action-bar-widget-extension.tsx
index 747ed4f2e7fe..0b83af12bebf 100644
--- a/ee/tabby-ui/app/files/components/action-bar-widget/action-bar-widget-extension.tsx
+++ b/ee/tabby-ui/app/files/components/action-bar-widget/action-bar-widget-extension.tsx
@@ -41,10 +41,9 @@ function createActionBarWidget(state: EditorState): Tooltip {
   const lineTo = state.doc.lineAt(selection.main.to)
   const isMultiline = lineFrom.number !== lineTo.number
   const pos = isMultiline ? lineTo.from : selection.main.from
-  const text = state.doc.sliceString(
-    state.selection.main.from,
-    state.selection.main.to
-  ) || ''
+  const text =
+    state.doc.sliceString(state.selection.main.from, state.selection.main.to) ||
+    ''
 
   return {
     pos,
diff --git a/ee/tabby-ui/app/files/components/completion-panel.tsx b/ee/tabby-ui/app/files/components/completion-panel.tsx
index d1550539a749..7b697a0850cd 100644
--- a/ee/tabby-ui/app/files/components/completion-panel.tsx
+++ b/ee/tabby-ui/app/files/components/completion-panel.tsx
@@ -28,19 +28,14 @@ import {
 import { Chat, ChatRef } from '@/components/chat'
 import { ClearChatsButton } from '@/components/clear-chats-button'
 import { EditChatTitleDialog } from '@/components/edit-chat-title-dialog'
+import LoadingWrapper from '@/components/loading-wrapper'
 import { ListSkeleton } from '@/components/skeleton'
 
-import { SourceCodeBrowserContext } from './source-code-browser'
-import LoadingWrapper from '@/components/loading-wrapper'
 import { CodeBrowserQuickAction } from '../lib/event-emitter'
+import { SourceCodeBrowserContext } from './source-code-browser'
 
 interface CompletionPanelProps
-  extends Omit<React.HTMLAttributes<HTMLDivElement>, 'children'> {
-  // open: boolean
-  // onOpenChange: (v: boolean) => void
-}
-
-const emptyMessages: Message[] = []
+  extends Omit<React.HTMLAttributes<HTMLDivElement>, 'children'> {}
 
 enum CompletionPanelView {
   CHAT,
@@ -57,47 +52,20 @@ export const CompletionPanel: React.FC<CompletionPanelProps> = ({
     setPendingEvent,
     setCompletionPanelViewType
   } = React.useContext(SourceCodeBrowserContext)
-  const _hasHydrated = useStore(useChatStore, state => state._hasHydrated)
   const chats = useStore(useChatStore, state => state.chats)
   const activeChatId = useStore(useChatStore, state => state.activeChatId)
   const chatId = activeChatId
   const chat = getChatById(chats, chatId)
-  const chatRef = React.useRef<ChatRef>(null)
   const appending = React.useRef(false)
-
-  // const quickActionBarCallback = (action: CodeBrowserQuickAction) => {
-  //   let builtInPrompt = ''
-  //   switch (action) {
-  //     case 'explain':
-  //       builtInPrompt = 'Explain the following code:'
-  //       break
-  //     case 'generate_unittest':
-  //       builtInPrompt = 'Generate a unit test for the following code:'
-  //       break
-  //     case 'generate_doc':
-  //       builtInPrompt = 'Generate documentation for the following code:'
-  //       break
-  //     default:
-  //       break
-  //   }
-  //   const view = editorRef.current?.editorView
-  //   const text =
-  //     view?.state.doc.sliceString(
-  //       view?.state.selection.main.from,
-  //       view?.state.selection.main.to
-  //     ) || ''
-
-  //   const initialMessage = `${builtInPrompt}\n${'```'}${
-  //     language ?? ''
-  //   }\n${text}\n${'```'}\n`
-  //   if (initialMessage) {
-  //     window.open(
-  //       `/playground?initialMessage=${encodeURIComponent(initialMessage)}`
-  //     )
-  //   }
-  // }
-
-  const getPrompt = ({ action, payload }: { action: CodeBrowserQuickAction, payload: string }) => {
+  const iframeRef = React.useRef<HTMLIFrameElement>(null)
+
+  const getPrompt = ({
+    action,
+    payload
+  }: {
+    action: CodeBrowserQuickAction
+    payload: string
+  }) => {
     let builtInPrompt = ''
     switch (action) {
       case 'explain':
@@ -116,43 +84,29 @@ export const CompletionPanel: React.FC<CompletionPanelProps> = ({
   }
 
   React.useEffect(() => {
-    if (chatRef.current?.append && pendingEvent) {
-      setCompletionPanelViewType(CompletionPanelView.CHAT)
-      if (!appending.current) {
-        appending.current = true
+    const contentWindow = iframeRef.current?.contentWindow
 
-        const prompt = getPrompt(pendingEvent as any)
+    if (pendingEvent) {
+      setCompletionPanelViewType(CompletionPanelView.CHAT)
 
-        chatRef.current
-          ?.append({
-            role: 'user',
-            content: prompt
-          })
-          .then(() => {
-            setPendingEvent(undefined)
-            appending.current = false
-          })
-        if (!chat) {
-          addChat(activeChatId, truncateText(prompt))
-        }
-      }
+      contentWindow?.postMessage({
+        action: 'append',
+        payload: getPrompt(pendingEvent)
+      })
+      setPendingEvent(undefined)
     }
-  }, [pendingEvent, chatRef.current?.append])
+  }, [pendingEvent, iframeRef.current?.contentWindow])
 
   return (
-    <div className={cn('h-full overflow relative', className)} {...props}>
+    <div className={cn('h-full flex flex-col', className)} {...props}>
       <Header />
       {completionPanelViewType === CompletionPanelView.CHAT && (
-        <LoadingWrapper loading={!activeChatId}>
-          <Chat
-            id={activeChatId}
-            // loading={_hasHydrated}
-            chatPanelClassName="w-full bottom-0 absolute lg:ml-0"
-            initialMessages={chat?.messages ?? emptyMessages}
-            key={activeChatId}
-            ref={chatRef}
-          />
-        </LoadingWrapper>
+        <iframe
+          src={`/playground`}
+          className="border-0 w-full flex-1"
+          key={activeChatId}
+          ref={iframeRef}
+        />
       )}
       {completionPanelViewType === CompletionPanelView.SESSIONS && (
         <ChatSessions />
@@ -247,30 +201,28 @@ function ChatSessions({ className }: { className?: string }) {
 }
 
 function Header() {
-  const {
-    completionPanelViewType,
-    setCompletionPanelViewType,
-    setCompletionPanelVisible
-  } = React.useContext(SourceCodeBrowserContext)
+  const { setCompletionPanelVisible } = React.useContext(
+    SourceCodeBrowserContext
+  )
 
-  const onToggleChatHistory = () => {
-    setCompletionPanelViewType(
-      completionPanelViewType === CompletionPanelView.CHAT
-        ? CompletionPanelView.SESSIONS
-        : CompletionPanelView.CHAT
-    )
-  }
+  // const onToggleChatHistory = () => {
+  //   setCompletionPanelViewType(
+  //     completionPanelViewType === CompletionPanelView.CHAT
+  //       ? CompletionPanelView.SESSIONS
+  //       : CompletionPanelView.CHAT
+  //   )
+  // }
 
-  const onNewChatClick = (e: React.MouseEvent<HTMLButtonElement>) => {
-    setActiveChatId(nanoid())
-    if (completionPanelViewType === CompletionPanelView.SESSIONS) {
-      setCompletionPanelViewType(CompletionPanelView.CHAT)
-    }
-  }
+  // const onNewChatClick = (e: React.MouseEvent<HTMLButtonElement>) => {
+  //   setActiveChatId(nanoid())
+  //   if (completionPanelViewType === CompletionPanelView.SESSIONS) {
+  //     setCompletionPanelViewType(CompletionPanelView.CHAT)
+  //   }
+  // }
 
   return (
-    <div className="flex items-center justify-between bg-secondary px-2 py-1">
-      <div className="flex items-center gap-2">
+    <div className="flex items-center justify-between bg-secondary px-2 py-1 sticky top-0">
+      {/* <div className="flex items-center gap-2">
         <Tooltip>
           <TooltipTrigger asChild>
             <Button size="icon" variant="ghost" onClick={onToggleChatHistory}>
@@ -287,9 +239,16 @@ function Header() {
           </TooltipTrigger>
           <TooltipContent side="bottom">Start a new chat</TooltipContent>
         </Tooltip>
-      </div>
+      </div> */}
+      <div className="w-8"></div>
       <Image src={tabbyLogo} alt="logo" width={32} />
-      <IconClose onClick={e => setCompletionPanelVisible(false)} />
+      <Button
+        size="icon"
+        variant="ghost"
+        onClick={e => setCompletionPanelVisible(false)}
+      >
+        <IconClose />
+      </Button>
     </div>
   )
 }
diff --git a/ee/tabby-ui/app/files/components/file-directory-breadcrumb.tsx b/ee/tabby-ui/app/files/components/file-directory-breadcrumb.tsx
index 1e84487f33d3..437a07f49e0c 100644
--- a/ee/tabby-ui/app/files/components/file-directory-breadcrumb.tsx
+++ b/ee/tabby-ui/app/files/components/file-directory-breadcrumb.tsx
@@ -16,40 +16,40 @@ const FileDirectoryBreadcrumb: React.FC<FileDirectoryBreadcrumbProps> = ({
   )
 
   return (
-    <div
-      className={cn('flex flex-nowrap items-center gap-1 leading-8', className)}
-    >
-      <div
-        className="cursor-pointer font-medium text-primary hover:underline"
-        onClick={e => setActivePath(undefined)}
-      >
-        Repositories
-      </div>
-      <div>/</div>
-      {currentFileRoutes?.map((route, idx) => {
-        const isRepo = idx === 0 && currentFileRoutes?.length > 1
-        const isActiveFile = idx === currentFileRoutes.length - 1
+    <div className={cn('flex flex-nowrap items-center gap-1', className)}>
+      <div className="flex items-center gap-1 flex-1 truncate">
+        <div
+          className="cursor-pointer font-medium text-primary hover:underline"
+          onClick={e => setActivePath(undefined)}
+        >
+          Repositories
+        </div>
+        <div>/</div>
+        {currentFileRoutes?.map((route, idx) => {
+          const isRepo = idx === 0 && currentFileRoutes?.length > 1
+          const isActiveFile = idx === currentFileRoutes.length - 1
 
-        return (
-          <React.Fragment key={route.fullPath}>
-            <div
-              className={cn(
-                isRepo || isActiveFile ? 'font-bold' : 'font-medium',
-                isActiveFile
-                  ? ''
-                  : 'cursor-pointer text-primary hover:underline',
-                isRepo ? 'hover:underline' : undefined
-              )}
-              onClick={e => setActivePath(route.fullPath)}
-            >
-              {route.name}
-            </div>
-            {route.file.kind !== 'file' && <div>/</div>}
-          </React.Fragment>
-        )
-      })}
+          return (
+            <React.Fragment key={route.fullPath}>
+              <div
+                className={cn(
+                  isRepo || isActiveFile ? 'font-bold' : 'font-medium',
+                  isActiveFile
+                    ? ''
+                    : 'cursor-pointer text-primary hover:underline',
+                  isRepo ? 'hover:underline' : undefined
+                )}
+                onClick={e => setActivePath(route.fullPath)}
+              >
+                {route.name}
+              </div>
+              {route.file.kind !== 'file' && <div>/</div>}
+            </React.Fragment>
+          )
+        })}
+      </div>
       {!!currentFileRoutes?.length && !!activePath && (
-        <CopyButton value={activePath} />
+        <CopyButton className="shrink-0" value={activePath} />
       )}
     </div>
   )
diff --git a/ee/tabby-ui/app/files/components/source-code-browser.tsx b/ee/tabby-ui/app/files/components/source-code-browser.tsx
index b4f30940b606..2a1543e9ef2f 100644
--- a/ee/tabby-ui/app/files/components/source-code-browser.tsx
+++ b/ee/tabby-ui/app/files/components/source-code-browser.tsx
@@ -3,6 +3,7 @@
 import React, { PropsWithChildren } from 'react'
 import filename2prism from 'filename2prism'
 import { compact, findIndex, toNumber } from 'lodash-es'
+import { ImperativePanelHandle } from 'react-resizable-panels'
 import { toast } from 'sonner'
 import { SWRResponse } from 'swr'
 import useSWRImmutable from 'swr/immutable'
@@ -18,6 +19,7 @@ import {
 } from '@/components/ui/resizable'
 import { useTopbarProgress } from '@/components/topbar-progress-indicator'
 
+import { CodeBrowserQuickAction, emitter } from '../lib/event-emitter'
 import { CompletionPanel } from './completion-panel'
 import { FileDirectoryBreadcrumb } from './file-directory-breadcrumb'
 import { DirectoryView } from './file-directory-view'
@@ -31,7 +33,6 @@ import {
   resolveFileNameFromPath,
   resolveRepoNameFromPath
 } from './utils'
-import { emitter } from '../lib/event-emitter'
 
 /**
  * FileMap example
@@ -80,8 +81,10 @@ type SourceCodeBrowserContextValue = {
   setCompletionPanelViewType: React.Dispatch<
     React.SetStateAction<CompletionPanelView>
   >
-  pendingEvent: { action: string; payload: string } | undefined
-  setPendingEvent: (d: { action: string; payload: string } | undefined) => void
+  pendingEvent: { action: CodeBrowserQuickAction; payload: string } | undefined
+  setPendingEvent: (
+    d: { action: CodeBrowserQuickAction; payload: string } | undefined
+  ) => void
 }
 
 const SourceCodeBrowserContext =
@@ -115,7 +118,7 @@ const SourceCodeBrowserContextProvider: React.FC<PropsWithChildren> = ({
     CompletionPanelView.CHAT
   )
   const [pendingEvent, setPendingEvent] = React.useState<
-    { action: string; payload: string } | undefined
+    { action: CodeBrowserQuickAction; payload: string } | undefined
   >()
 
   const updateFileMap = (map: TFileMap) => {
@@ -201,6 +204,8 @@ const SourceCodeBrowserRenderer: React.FC<SourceCodeBrowserProps> = ({
     setPendingEvent
   } = React.useContext(SourceCodeBrowserContext)
   const { setProgress } = useTopbarProgress()
+  const completionPanelRef = React.useRef<ImperativePanelHandle>(null)
+  const [completionPanelSize, setCompletionPanelSize] = React.useState(30)
 
   const activeRepoName = React.useMemo(() => {
     return resolveRepoNameFromPath(activePath)
@@ -291,6 +296,12 @@ const SourceCodeBrowserRenderer: React.FC<SourceCodeBrowserProps> = ({
   const showRawFileView =
     isFileSelected && (fileViewType === 'image' || fileViewType === 'raw')
 
+  const onPanelLayout = (sizes: number[]) => {
+    if (sizes?.[2]) {
+      setCompletionPanelSize(sizes[2])
+    }
+  }
+
   React.useEffect(() => {
     const init = async () => {
       const { patchMap, expandedKeys } = await initFileMap(activePath)
@@ -362,8 +373,21 @@ const SourceCodeBrowserRenderer: React.FC<SourceCodeBrowserProps> = ({
     }
   }, [])
 
+  React.useEffect(() => {
+    if (completionPanelVisible) {
+      completionPanelRef.current?.expand()
+      completionPanelRef.current?.resize(completionPanelSize)
+    } else {
+      completionPanelRef.current?.collapse()
+    }
+  }, [completionPanelVisible])
+
   return (
-    <ResizablePanelGroup direction="horizontal" className={cn(className)}>
+    <ResizablePanelGroup
+      direction="horizontal"
+      className={cn(className)}
+      onLayout={onPanelLayout}
+    >
       <ResizablePanel defaultSize={20} minSize={20}>
         <FileTreePanel />
       </ResizablePanel>
@@ -396,14 +420,17 @@ const SourceCodeBrowserRenderer: React.FC<SourceCodeBrowserProps> = ({
           </div>
         </div>
       </ResizablePanel>
-      {completionPanelVisible && (
-        <>
-          <ResizableHandle className="w-1 bg-border/40 hover:bg-border active:bg-border" />
-          <ResizablePanel defaultSize={30}>
-            <CompletionPanel />
-          </ResizablePanel>
-        </>
-      )}
+      <>
+        <ResizableHandle
+          className={cn(
+            'w-1 bg-border/40 hover:bg-border active:bg-border hidden',
+            completionPanelVisible && 'block'
+          )}
+        />
+        <ResizablePanel defaultSize={0} collapsible ref={completionPanelRef}>
+          <CompletionPanel />
+        </ResizablePanel>
+      </>
     </ResizablePanelGroup>
   )
 }
diff --git a/ee/tabby-ui/app/files/lib/event-emitter.ts b/ee/tabby-ui/app/files/lib/event-emitter.ts
index a40241c57157..ecea01786a5b 100644
--- a/ee/tabby-ui/app/files/lib/event-emitter.ts
+++ b/ee/tabby-ui/app/files/lib/event-emitter.ts
@@ -3,7 +3,7 @@ import mitt from 'mitt'
 type CodeBrowserQuickAction = 'explain' | 'generate_unittest' | 'generate_doc'
 
 type CodeBrowserQuickActionEvents = {
-  code_browser_quick_action: { action: CodeBrowserQuickAction, payload: string }
+  code_browser_quick_action: { action: CodeBrowserQuickAction; payload: string }
 }
 
 const emitter = mitt<CodeBrowserQuickActionEvents>()
diff --git a/ee/tabby-ui/app/playground/components/chats.tsx b/ee/tabby-ui/app/playground/components/chats.tsx
index 3dc5d16eab0a..62a635d4e339 100644
--- a/ee/tabby-ui/app/playground/components/chats.tsx
+++ b/ee/tabby-ui/app/playground/components/chats.tsx
@@ -1,7 +1,7 @@
 'use client'
 
 import React from 'react'
-import type { Message } from 'ai'
+import { nanoid, type Message } from 'ai'
 
 import useRouterStuff from '@/lib/hooks/use-router-stuff'
 import { useStore } from '@/lib/hooks/use-store'
@@ -52,6 +52,35 @@ export default function Chats() {
     }
   }, [chatRef.current?.append])
 
+  React.useEffect(() => {
+    const onMessage = async (event: MessageEvent) => {
+      if (event.origin !== window.origin || !event.data) {
+        return
+      }
+      if (!chatRef.current || chatRef.current.isLoading) return
+
+      const { data } = event
+      if (data.action === 'append') {
+        chatRef.current.append({
+          id: nanoid(),
+          role: 'user',
+          content: data.payload
+        })
+        return
+      }
+
+      if (data.action === 'stop') {
+        chatRef.current.stop()
+      }
+    }
+
+    window.addEventListener('message', onMessage)
+
+    return () => {
+      window.removeEventListener('message', onMessage)
+    }
+  }, [chatRef.current])
+
   return (
     <div className="grid flex-1 overflow-hidden lg:grid-cols-[280px_1fr]">
       <ChatSessions className="hidden w-[280px] border-r lg:block" />
diff --git a/ee/tabby-ui/components/chat-panel.tsx b/ee/tabby-ui/components/chat-panel.tsx
index a49098a48182..831a7a2b9386 100644
--- a/ee/tabby-ui/components/chat-panel.tsx
+++ b/ee/tabby-ui/components/chat-panel.tsx
@@ -3,7 +3,7 @@ import type { UseChatHelpers } from 'ai/react'
 
 import { cn } from '@/lib/utils'
 import { Button } from '@/components/ui/button'
-import { IconRefresh, IconStop } from '@/components/ui/icons'
+import { IconRefresh, IconStop, IconTrash } from '@/components/ui/icons'
 import { ButtonScrollToBottom } from '@/components/button-scroll-to-bottom'
 import { FooterText } from '@/components/footer'
 import { PromptForm, PromptFormRef } from '@/components/prompt-form'
@@ -18,6 +18,7 @@ export interface ChatPanelProps
     | 'stop'
     | 'input'
     | 'setInput'
+    | 'setMessages'
   > {
   id?: string
   className?: string
@@ -34,13 +35,19 @@ export function ChatPanel({
   setInput,
   messages,
   className,
-  onSubmit
+  onSubmit,
+  setMessages
 }: ChatPanelProps) {
   const promptFormRef = React.useRef<PromptFormRef>(null)
   React.useEffect(() => {
     promptFormRef?.current?.focus()
   }, [id])
 
+  const onClearContext = () => {
+    stop()
+    setMessages([])
+  }
+
   return (
     <div
       className={cn(
@@ -50,7 +57,7 @@ export function ChatPanel({
     >
       <ButtonScrollToBottom />
       <div className="mx-auto sm:max-w-2xl sm:px-4">
-        <div className="flex h-10 items-center justify-center">
+        <div className="flex h-10 items-center justify-center gap-2">
           {isLoading ? (
             <Button
               variant="outline"
@@ -72,6 +79,16 @@ export function ChatPanel({
               </Button>
             )
           )}
+          {messages?.length > 0 && (
+            <Button
+              variant="outline"
+              onClick={() => onClearContext()}
+              className="bg-background lg:hidden"
+            >
+              <IconTrash className="mr-2" />
+              Clear
+            </Button>
+          )}
         </div>
         <div className="space-y-4 border-t bg-background px-4 py-2 shadow-lg sm:rounded-t-xl sm:border md:py-4">
           <PromptForm
diff --git a/ee/tabby-ui/components/chat.tsx b/ee/tabby-ui/components/chat.tsx
index dfabdad2f405..4eeded376e5b 100644
--- a/ee/tabby-ui/components/chat.tsx
+++ b/ee/tabby-ui/components/chat.tsx
@@ -167,6 +167,7 @@ function ChatRenderer(
           messages={messages}
           input={input}
           setInput={setInput}
+          setMessages={setMessages}
         />
       </div>
     </div>