Merge pull request #22 from rjmacarthy/feature/persist-chats

Feature: persist chat and new session

src/context.ts

@@ -0,0 +1,13 @@
+import { ExtensionContext } from 'vscode'
+
+let context: ExtensionContext | null = null
+
+export function setContext(newContext: ExtensionContext) {
+  context = newContext
+}
+
+export function getContext() {
+  return context
+}
+
+module.exports = { setContext, getContext }

src/extension.ts

@@ -11,12 +11,14 @@ import { CompletionProvider } from './providers/completion'
 import { init } from './init'
 import { SidebarProvider } from './providers/sidebar'
 import { chatCompletion, delayExecution, deleteTempFiles } from './utils'
+import { setContext } from './context'
 
 export async function activate(context: ExtensionContext) {
   const config = workspace.getConfiguration('twinny')
   const fimModel = config.get('fimModelName') as string
   const chatModel = config.get('chatModelName') as string
   const statusBar = window.createStatusBarItem(StatusBarAlignment.Right)
+  setContext(context)
 
   try {
     await init()

src/providers/sidebar.ts

@@ -1,6 +1,7 @@
 import * as vscode from 'vscode'
 import { chatCompletion, getTextSelection, openDiffView } from '../utils'
 import { chatMessage } from '../prompts'
+import { getContext } from '../context'
 
 export class SidebarProvider implements vscode.WebviewViewProvider {
   view?: vscode.WebviewView
@@ -33,7 +34,10 @@ export class SidebarProvider implements vscode.WebviewViewProvider {
     )
 
     webviewView.webview.onDidReceiveMessage(
-      (data: { type: string; data: Message[] | string }) => {
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      (data: any) => {
+        const context = getContext()
+
         if (data.type === 'chatMessage') {
           chatCompletion('chat', this.view, (selection: string) =>
             chatMessage(data.data as Message[], selection)
@@ -63,15 +67,23 @@ export class SidebarProvider implements vscode.WebviewViewProvider {
         if (data.type === 'acceptSolution') {
           const editor = vscode.window.activeTextEditor
           const selection = editor?.selection
-
-          if (!selection) {
-            return
-          }
-
+          if (!selection) return
           vscode.window.activeTextEditor?.edit((editBuilder) => {
             editBuilder.replace(selection, data.data as string)
           })
         }
+        if (data.type === 'getTwinnyWorkspaceContext') {
+          this.view?.webview.postMessage({
+            type: `twinnyWorkSpaceContext-${data.key}`,
+            value: context?.workspaceState.get(`twinnyWorkSpaceContext-${data.key}`) || ''
+          })
+        }
+        if (data.type === 'setTwinnyWorkSpaceContext') {
+          context?.workspaceState.update(
+            `twinnyWorkSpaceContext-${data.key}`,
+            data.data
+          )
+        }
       }
     )
   }

src/webview/chat.tsx

@@ -12,14 +12,17 @@ import { Selection } from './selection'
 import { BOT_NAME, USER_NAME } from './constants'
 
 import styles from './index.module.css'
+import { useWorkSpaceContext } from './hooks'
+//import { useWorkSpaceContext } from './hooks'
 
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
 const global = globalThis as any
 export const Chat = () => {
   const [inputText, setInputText] = useState('')
   const [isGenerating, setIsGenerating] = useState(false)
   const [loading, setLoading] = useState(false)
-  const [messages, setMessages] = useState<Message[]>([])
+  const lastConversation = useWorkSpaceContext<Message[]>('lastConversation')
+  const [messages, setMessages] = useState<Message[] | undefined>()
   const [completion, setCompletion] = useState<Message | null>()
   const divRef = useRef<HTMLDivElement>(null)
 
@@ -36,10 +39,9 @@ export const Chat = () => {
 
     if (inputText.trim()) {
       setInputText('')
-
       global.vscode.postMessage({
         type: 'chatMessage',
-        data: messages.length
+        data: messages?.length
           ? [
               ...messages,
               {
@@ -58,7 +60,7 @@ export const Chat = () => {
       })
 
       setMessages((prev) => [
-        ...prev,
+        ...(prev || []),
         { role: USER_NAME, content: inputText, type: '' }
       ])
 
@@ -87,14 +89,20 @@ export const Chat = () => {
         }
         case 'onEnd': {
           setMessages((prev) => {
-            return [
-              ...prev,
+            const update = [
+              ...(prev || []),
               {
                 role: BOT_NAME,
                 content: message.value.completion,
                 type: message.value.type
               }
             ]
+            global.vscode.postMessage({
+              type: 'setTwinnyWorkSpaceContext',
+              key: 'lastConversation',
+              data: update
+            })
+            return update
           })
           setCompletion(null)
           setIsGenerating(false)
@@ -103,16 +111,25 @@ export const Chat = () => {
     })
   }, [])
 
+  useEffect(() => {
+    if (lastConversation?.length) {
+      return setMessages(lastConversation)
+    }
+    setMessages([])
+  }, [lastConversation])
+
   return (
     <VSCodePanelView>
       <div className={styles.container}>
         <div className={styles.markdown} ref={divRef}>
-          {messages.map((message) => (
-            <Message
-              completionType={message.type}
-              sender={message.role}
-              message={message.content}
-            />
+          {messages?.map((message, index) => (
+            <div key={`message-${index}`}>
+              <Message
+                completionType={message.type}
+                sender={message.role}
+                message={message.content}
+              />
+            </div>
           ))}
           {loading && (
             <div className={styles.loading}>

src/webview/hooks.ts

@@ -5,18 +5,44 @@ const global = globalThis as any
 
 export const useSelection = (onSelect: () => void) => {
   const [selection, setSelection] = useState('')
+  const handler = (event: MessageEvent) => {
+    const message: PostMessage = event.data
+    if (message?.type === 'textSelection') {
+      setSelection(message?.value.completion.trim())
+      onSelect?.()
+    }
+  }
+
   useEffect(() => {
-    window.addEventListener('message', (event) => {
-      const message: PostMessage = event.data
-      if (message?.type === 'textSelection') {
-        setSelection(message?.value.completion.trim())
-        onSelect?.()
-      }
-    })
+    window.addEventListener('message', handler)
     global.vscode.postMessage({
       type: 'getTextSelection'
     })
+    return () => window.removeEventListener('message', handler)
   }, [])
 
   return selection
 }
+
+export const useWorkSpaceContext = <T>(key: string) => {
+  const [context, setContext] = useState<T>()
+
+  const handler = (event: MessageEvent) => {
+    const message: PostMessage = event.data
+    if (message?.type === `twinnyWorkSpaceContext-${key}`) {
+      setContext(event.data.value)
+    }
+  }
+
+  useEffect(() => {
+    window.addEventListener('message', handler)
+    global.vscode.postMessage({
+      type: 'getTwinnyWorkspaceContext',
+      key
+    })
+
+    return () => window.removeEventListener('message', handler)
+  }, [])
+
+  return context
+}

src/webview/index.module.css

@@ -16,6 +16,13 @@ code {
   height: 95vh;
 }
 
+.controlBar {
+  display: flex;
+  align-items: center;
+  justify-content: flex-end;
+  gap: 4px;
+}
+
 .markdown {
   flex-grow: 1;
   overflow-y: auto;

src/webview/index.tsx

@@ -2,17 +2,23 @@ import { createRoot } from 'react-dom/client'
 
 import { Chat } from './chat'
 import { Settings } from './settings'
+import { NewSession } from './new-session'
+
+import styles from './index.module.css'
 
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
-(globalThis as any).vscode = window.acquireVsCodeApi()
+;(globalThis as any).vscode = window.acquireVsCodeApi()
 
 const container = document.querySelector('#root')
 
 if (container) {
   const root = createRoot(container)
   root.render(
     <>
-      <Settings />
+      <div className={styles.controlBar}>
+        <NewSession />
+        <Settings />
+      </div>
       <Chat />
     </>
   )

src/webview/new-session.tsx

@@ -0,0 +1,59 @@
+import { VSCodeButton } from '@vscode/webview-ui-toolkit/react'
+
+import styles from './index.module.css'
+
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+const global = globalThis as any
+
+const key = 'lastConversation'
+
+export const NewSession = () => {
+  const handleClickClearSession = (): void => {
+    global.vscode.postMessage({
+      type: 'setTwinnyWorkSpaceContext',
+      key,
+      data: []
+    })
+    global.vscode.postMessage({
+      type: 'getTwinnyWorkspaceContext',
+      key
+    })
+  }
+
+  return (
+    <div className={styles.settings}>
+      <VSCodeButton
+        title="New chat session"
+        appearance="icon"
+        onClick={handleClickClearSession}
+      >
+        <ClearChatIcon />
+      </VSCodeButton>
+    </div>
+  )
+}
+
+function ClearChatIcon() {
+  return (
+    <svg
+      width="24"
+      height="24"
+      viewBox="0 0 24 24"
+      fill="none"
+      xmlns="http://www.w3.org/2000/svg"
+    >
+      <path
+        d="M2 21l4-4h12a2 2 0 002-2V3a2 2 0 00-2-2H4a2 2 0 00-2 2v16z"
+        stroke="currentColor"
+        strokeWidth="2"
+      />
+      <path
+        d="M9 6l6 6M15 6l-6 6"
+        stroke="currentColor"
+        strokeWidth="2"
+        strokeLinecap="round"
+        strokeLinejoin="round"
+      />
+    </svg>
+  )
+}