clean up symmetry bugs

src/common/constants.ts

@@ -236,6 +236,7 @@ export const STOP_QWEN = [
   '<|im_start|>',
   '<|repo_name|>',
   '<|fim_pad|>',
+  '<|cursor|>',
 ]
 
 export const STOP_CODEGEMMA = ['<|file_separator|>', '<|end_of_turn|>', '<eos>']
@@ -482,24 +483,6 @@ export const MULTILINE_TYPES = [...MULTILINE_OUTSIDE, ...MULTILINE_INSIDE]
 
 export const MULTI_LINE_DELIMITERS = ['\n\n', '\r\n\r\n']
 
-export const SYMMETRY_DATA_MESSAGE = {
-  disconnect: 'disconnect',
-  heartbeat: 'heartbeat',
-  inference: 'inference',
-  inferenceEnd: 'inferenceEnd',
-  join: 'join',
-  leave: 'leave',
-  newConversation: 'newConversation',
-  ping: 'ping',
-  pong: 'pong',
-  providerDetails: 'providerDetails',
-  reportCompletion: 'reportCompletion',
-  requestProvider: 'requestProvider',
-  sessionValid: 'sessionValid',
-  verifySession: 'verifySession'
-} as const
-
 export const SYMMETRY_EMITTER_KEY = {
-  conversationTitle: 'conversationTitle',
   inference: 'inference'
 }

src/common/types.ts

@@ -1,6 +1,7 @@
 import { InlineCompletionItem, InlineCompletionList } from 'vscode'
 import { CodeLanguageDetails } from './languages'
-import { ALL_BRACKETS, SYMMETRY_DATA_MESSAGE } from './constants'
+import { ALL_BRACKETS } from './constants'
+import { serverMessageKeys } from 'symmetry-core'
 
 export interface RequestBodyBase {
   stream: boolean
@@ -33,7 +34,7 @@ export interface StreamResponse {
   response: string
   content: string
   message: {
-    content: string
+    role: 'assistant'
   }
   done: boolean
   context: number[]
@@ -240,7 +241,7 @@ export interface SymmetryMessage<T> {
   data: T
 }
 
-export type ServerMessageKey = keyof typeof SYMMETRY_DATA_MESSAGE
+export type ServerMessageKey = keyof typeof serverMessageKeys
 
 export interface SymmetryConnection {
   sessionToken?: string

src/extension/conversation-history.ts

@@ -12,15 +12,13 @@ import { v4 as uuidv4 } from 'uuid'
 import { createStreamRequestBody } from './provider-options'
 import { TwinnyProvider } from './provider-manager'
 import { streamResponse } from './stream'
-import { createSymmetryMessage, getChatDataFromProvider } from './utils'
+import { getChatDataFromProvider } from './utils'
 import {
   ACTIVE_CHAT_PROVIDER_STORAGE_KEY,
   ACTIVE_CONVERSATION_STORAGE_KEY,
   CONVERSATION_EVENT_NAME,
   CONVERSATION_STORAGE_KEY,
   EXTENSION_SESSION_NAME,
-  SYMMETRY_EMITTER_KEY,
-  SYMMETRY_DATA_MESSAGE,
   TITLE_GENERATION_PROMPT_MESAGE,
   USER,
 } from '../common/constants'
@@ -58,22 +56,6 @@ export class ConversationHistory {
         this.handleMessage(message)
       }
     )
-
-    this._symmetryService.on(
-      SYMMETRY_EMITTER_KEY.conversationTitle,
-      (completion: string) => {
-        const activeConversation = this.getActiveConversation()
-        this.webView?.postMessage({
-          type: CONVERSATION_EVENT_NAME.getActiveConversation,
-          value: {
-            data: {
-              ...activeConversation,
-              title: completion
-            }
-          }
-        } as ServerMessage<Conversation>)
-      }
-    )
   }
 
   handleMessage(message: ClientMessage<Conversation>) {
@@ -194,6 +176,12 @@ export class ConversationHistory {
   }
 
   async getConversationTitle(messages: Message[]): Promise<string> {
+    const symmetryConnected = this._sessionManager?.get(
+      EXTENSION_SESSION_NAME.twinnySymmetryConnection
+    )
+    if (symmetryConnected) {
+      return Promise.resolve(`${messages[0].content?.substring(0, 50)}...`)
+    }
     const request = this.buildStreamRequestTitle(messages)
     if (!request) return ''
     return await this.streamConversationTitle(request)

src/extension/providers/base.ts

@@ -1,4 +1,5 @@
 import * as vscode from 'vscode'
+import { serverMessageKeys } from 'symmetry-core'
 
 import {
   createSymmetryMessage,
@@ -13,7 +14,6 @@ import {
   EXTENSION_SESSION_NAME,
   EVENT_NAME,
   TWINNY_COMMAND_NAME,
-  SYMMETRY_DATA_MESSAGE,
   SYMMETRY_EMITTER_KEY,
   SYSTEM
 } from '../../common/constants'
@@ -24,7 +24,6 @@ import {
   ApiModel,
   ServerMessage,
   InferenceRequest,
-  SymmetryModelProvider,
   FileItem
 } from '../../common/types'
 import { TemplateProvider } from '../template-provider'
@@ -147,8 +146,6 @@ export class BaseProvider {
         [EVENT_NAME.twinnyAcceptSolution]: this.acceptSolution,
         [EVENT_NAME.twinnyChatMessage]: this.streamChatCompletion,
         [EVENT_NAME.twinnyClickSuggestion]: this.clickSuggestion,
-        [EVENT_NAME.twinnyConnectSymmetry]: this.connectToSymmetry,
-        [EVENT_NAME.twinnyDisconnectSymmetry]: this.disconnectSymmetry,
         [EVENT_NAME.twinnyEmbedDocuments]: this.embedDocuments,
         [EVENT_NAME.twinnyFetchOllamaModels]: this.fetchOllamaModels,
         [EVENT_NAME.twinnyGetConfigValue]: this.getConfigurationValue,
@@ -167,8 +164,6 @@ export class BaseProvider {
         [EVENT_NAME.twinnySetGlobalContext]: this.setGlobalContext,
         [EVENT_NAME.twinnySetTab]: this.setTab,
         [EVENT_NAME.twinnySetWorkspaceContext]: this.setWorkspaceContext,
-        [EVENT_NAME.twinnyStartSymmetryProvider]: this.createSymmetryProvider,
-        [EVENT_NAME.twinnyStopSymmetryProvider]: this.stopSymmetryProvider,
         [EVENT_NAME.twinnyTextSelection]: this.getSelectedText,
         [EVENT_NAME.twinnyFileListRequest]: this.fileListRequest,
         [EVENT_NAME.twinnyNewConversation]: this.twinnyNewConversation,
@@ -180,7 +175,7 @@ export class BaseProvider {
 
   public newConversation() {
     this._symmetryService?.write(
-      createSymmetryMessage(SYMMETRY_DATA_MESSAGE.newConversation)
+      createSymmetryMessage(serverMessageKeys.newConversation)
     )
   }
 
@@ -315,7 +310,7 @@ export class BaseProvider {
 
       return this._symmetryService?.write(
         createSymmetryMessage<InferenceRequest>(
-          SYMMETRY_DATA_MESSAGE.inference,
+          serverMessageKeys.inference,
           {
             messages,
             key: SYMMETRY_EMITTER_KEY.inference
@@ -343,7 +338,7 @@ export class BaseProvider {
       `)
       return this._symmetryService?.write(
         createSymmetryMessage<InferenceRequest>(
-          SYMMETRY_DATA_MESSAGE.inference,
+          serverMessageKeys.inference,
           {
             messages,
             key: SYMMETRY_EMITTER_KEY.inference
@@ -465,30 +460,6 @@ export class BaseProvider {
     })
   }
 
-  private connectToSymmetry = (data: ClientMessage<SymmetryModelProvider>) => {
-    if (this._config.symmetryServerKey) {
-      this._symmetryService?.connect(
-        this._config.symmetryServerKey,
-        data.data?.model_name,
-        data.data?.provider
-      )
-    }
-  }
-
-  private disconnectSymmetry = async () => {
-    if (this._config.symmetryServerKey) {
-      await this._symmetryService?.disconnect()
-    }
-  }
-
-  private createSymmetryProvider = () => {
-    this._symmetryService?.startSymmetryProvider()
-  }
-
-  private stopSymmetryProvider = () => {
-    this._symmetryService?.stopSymmetryProvider()
-  }
-
   private twinnyHideBackButton() {
     vscode.commands.executeCommand(TWINNY_COMMAND_NAME.hideBackButton)
   }

src/extension/symmetry-service.ts

@@ -2,7 +2,7 @@
 import { workspace, commands, ExtensionContext, Webview } from 'vscode'
 import Hyperswarm from 'hyperswarm'
 import crypto from 'hypercore-crypto'
-import { SymmetryProvider } from 'symmetry-core'
+import { serverMessageKeys, SymmetryProvider } from 'symmetry-core'
 import path from 'path'
 import os from 'os'
 import fs from 'fs'
@@ -21,14 +21,15 @@ import {
   ServerMessage,
   Peer,
   SymmetryMessage,
-  SymmetryConnection
+  SymmetryConnection,
+  ClientMessage,
+  SymmetryModelProvider
 } from '../common/types'
 import {
   EVENT_NAME,
   EXTENSION_CONTEXT_NAME,
   EXTENSION_SESSION_NAME,
   SYMMETRY_EMITTER_KEY,
-  SYMMETRY_DATA_MESSAGE,
   WEBUI_TABS,
   ACTIVE_CHAT_PROVIDER_STORAGE_KEY,
   GLOBAL_STORAGE_KEY
@@ -79,15 +80,26 @@ export class SymmetryService extends EventEmitter {
 
     this._ws = new SymmetryWs(this._webView)
     this._ws.connectSymmetryWs()
+    this.setupEventListeners()
   }
 
-  public connect = async (
-    key: string,
-    model: string | undefined,
-    provider: string | undefined
-  ) => {
-    if (!model || !key) return
-    this._symmetryProvider = provider
+  private setupEventListeners() {
+    this._webView?.onDidReceiveMessage((message) => {
+      const eventHandlers = {
+        [EVENT_NAME.twinnyConnectSymmetry]: this.connect,
+        [EVENT_NAME.twinnyDisconnectSymmetry]: this.disconnect,
+        [EVENT_NAME.twinnyStartSymmetryProvider]: this.startSymmetryProvider,
+        [EVENT_NAME.twinnyStopSymmetryProvider]: this.stopSymmetryProvider,
+      }
+      eventHandlers[message.type as string]?.(message)
+    })
+  }
+
+  public connect = async (data: ClientMessage<SymmetryModelProvider>) => {
+    const key = this._config.symmetryServerKey
+    const model = data.data?.model_name
+    if (!data.data?.model_name || !key) return
+    this._symmetryProvider = data.data.provider
     this._serverSwarm = new Hyperswarm()
     const serverKey = Buffer.from(key)
     const discoveryKey = crypto.discoveryKey(serverKey)
@@ -97,7 +109,7 @@ export class SymmetryService extends EventEmitter {
     this._serverSwarm.on('connection', (peer: Peer) => {
       this._serverPeer = peer
       peer.write(
-        createSymmetryMessage(SYMMETRY_DATA_MESSAGE.requestProvider, {
+        createSymmetryMessage(serverMessageKeys.requestProvider, {
           modelName: model
         })
       )
@@ -107,20 +119,20 @@ export class SymmetryService extends EventEmitter {
         )
         if (data && data.key) {
           switch (data?.key) {
-            case SYMMETRY_DATA_MESSAGE.ping:
+            case serverMessageKeys.ping:
               this._providerPeer?.write(
-                createSymmetryMessage(SYMMETRY_DATA_MESSAGE.pong)
+                createSymmetryMessage(serverMessageKeys.pong)
               )
               break
-            case SYMMETRY_DATA_MESSAGE.providerDetails:
+            case serverMessageKeys.providerDetails:
               peer.write(
                 createSymmetryMessage(
-                  SYMMETRY_DATA_MESSAGE.verifySession,
+                  serverMessageKeys.verifySession,
                   data.data?.sessionToken
                 )
               )
               break
-            case SYMMETRY_DATA_MESSAGE.sessionValid:
+            case serverMessageKeys.sessionValid:
               this.connectToProvider(data.data)
           }
         }
@@ -187,7 +199,7 @@ export class SymmetryService extends EventEmitter {
 
       if (!this._emitterKey) return
 
-      if (str.includes(SYMMETRY_DATA_MESSAGE.inferenceEnd))
+      if (str.includes(serverMessageKeys.inferenceEnded))
         this.handleInferenceEnd()
 
       this.handleIncomingData(chunk, (response: StreamResponse) => {

src/webview/chat.tsx

@@ -16,7 +16,6 @@ import {
   WORKSPACE_STORAGE_KEY,
   EVENT_NAME,
   USER,
-  SYMMETRY_EMITTER_KEY,
   EXTENSION_CONTEXT_NAME
 } from '../common/constants'
 
@@ -105,10 +104,6 @@ export const Chat = (props: ChatProps): JSX.Element => {
           }
         ]
 
-        if (message.value.type === SYMMETRY_EMITTER_KEY.conversationTitle) {
-          return messages
-        }
-
         saveLastConversation({
           ...conversation,
           messages: messages
@@ -553,7 +548,7 @@ export const Chat = (props: ChatProps): JSX.Element => {
             <VSCodeBadge>{selection?.length}</VSCodeBadge>
           </div>
           <div>
-            {generatingRef.current && (
+            {generatingRef.current && !symmetryConnection && (
               <VSCodeButton
                 type="button"
                 appearance="icon"