when connecting to a rogue server, send close message

microsoft · Jul 29, 2024 · 546e550 · 546e550
1 parent c0b2d9f
commit 546e550
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 5 deletions.
diff --git a/packages/cli/src/server.ts b/packages/cli/src/server.ts
@@ -202,7 +202,7 @@ export async function startServer(options: { port: string }) {
                     }
                     case "script.abort": {
                         const { runId, reason } = data
-                        console.log(`abort run ${runId}`)
+                        console.log(`run ${runId}: abort`)
                         const run = runs[runId]
                         if (run) {
                             delete runs[runId]

diff --git a/packages/core/src/constants.ts b/packages/core/src/constants.ts
@@ -1,6 +1,7 @@
 export const CHANGE = "change"
 export const TRACE_CHUNK = "traceChunk"
 export const RECONNECT = "reconnect"
+export const OPEN = "open"
 export const MAX_CACHED_TEMPERATURE = 0.5
 export const MAX_CACHED_TOP_P = 0.5
 export const MAX_TOOL_CALLS = 100

diff --git a/packages/core/src/server/client.ts b/packages/core/src/server/client.ts
@@ -1,5 +1,5 @@
 import { ChatCompletionsProgressReport } from "../chattypes"
-import { CLIENT_RECONNECT_DELAY, RECONNECT } from "../constants"
+import { CLIENT_RECONNECT_DELAY, OPEN, RECONNECT } from "../constants"
 import { randomHex } from "../crypto"
 import { errorMessage } from "../error"
 import { GenerationResult } from "../generation"
@@ -97,6 +97,7 @@ export class WebSocketClient extends EventTarget {
                 (m = this._pendingMessages.pop())
             )
                 this._ws.send(m)
+            this.dispatchEvent(new Event(OPEN))
         })
         this._ws.addEventListener("error", (ev) => {
             this.reconnect()

diff --git a/packages/vscode/src/servermanager.ts b/packages/vscode/src/servermanager.ts
@@ -3,6 +3,7 @@ import { ExtensionState } from "./state"
 import {
     SERVER_PORT,
     RECONNECT,
+    OPEN,
     TOOL_NAME,
     ICON_LOGO_NAME,
     CLIENT_RECONNECT_MAX_ATTEMPTS,
@@ -11,7 +12,7 @@ import {
     VSCODE_CONFIG_CLI_PATH,
 } from "../../core/src/constants"
 import { ServerManager, host } from "../../core/src/host"
-import { logError } from "../../core/src/util"
+import { logError, logVerbose } from "../../core/src/util"
 import { WebSocketClient } from "../../core/src/server/client"
 import { CORE_VERSION } from "../../core/src/version"
 
@@ -42,9 +43,19 @@ export class TerminalServerManager implements ServerManager {
         )
 
         this.client = new WebSocketClient(`http://localhost:${SERVER_PORT}`)
+        this.client.addEventListener(OPEN, () => {
+            // client connected to a rogue server
+            if (!this._terminal) {
+                logVerbose("found rogue server, closing...")
+                this.client?.kill()
+            }
+        })
         this.client.addEventListener(RECONNECT, () => {
             // server process died somehow
-            if (this.client.connectedOnce && this.client.reconnectAttempts > CLIENT_RECONNECT_MAX_ATTEMPTS) {
+            if (
+                this.client.connectedOnce &&
+                this.client.reconnectAttempts > CLIENT_RECONNECT_MAX_ATTEMPTS
+            ) {
                 this.closeTerminal()
                 this.start()
             }
@@ -65,7 +76,9 @@ export class TerminalServerManager implements ServerManager {
         const cliPath = config.get(VSCODE_CONFIG_CLI_PATH) as string
         if (cliPath) this._terminal.sendText(`node "${cliPath}" serve`)
         else {
-            const cliVersion = (config.get(VSCODE_CONFIG_CLI_VERSION) as string) || CORE_VERSION
+            const cliVersion =
+                (config.get(VSCODE_CONFIG_CLI_VERSION) as string) ||
+                CORE_VERSION
             this._terminal.sendText(`npx --yes ${TOOL_ID}@${cliVersion} serve`)
         }
         this._terminal.show()