feat(shellv2): implement websocket reconnection logic and status indicator

tavern/internal/www/src/lib/headless-adapter.ts

@@ -6,27 +6,51 @@ export interface ExecutionResult {
     message?: string;
 }
 
+export type ConnectionStatus = "connected" | "disconnected" | "reconnecting";
+
 export class HeadlessWasmAdapter {
     private repl: any; // HeadlessRepl instance
-    private ws: WebSocket;
+    private ws: WebSocket | null = null;
+    private url: string;
     private onMessageCallback: (msg: WebsocketMessage) => void;
     private onReadyCallback?: () => void;
+    private onStatusChangeCallback?: (status: ConnectionStatus) => void;
     private isWsOpen: boolean = false;
+    private reconnectTimer: NodeJS.Timeout | null = null;
+    private isClosedExplicitly: boolean = false;
 
-    constructor(url: string, onMessage: (msg: WebsocketMessage) => void, onReady?: () => void) {
+    constructor(
+        url: string,
+        onMessage: (msg: WebsocketMessage) => void,
+        onReady?: () => void,
+        onStatusChange?: (status: ConnectionStatus) => void
+    ) {
+        this.url = url;
         this.onMessageCallback = onMessage;
         this.onReadyCallback = onReady;
-        this.ws = new WebSocket(url);
+        this.onStatusChangeCallback = onStatusChange;
+
+        this.connect();
+    }
+
+    private connect() {
+        if (this.isClosedExplicitly) return;
+
+        this.ws = new WebSocket(this.url);
 
         this.ws.onopen = () => {
             this.isWsOpen = true;
+            if (this.reconnectTimer) {
+                clearTimeout(this.reconnectTimer);
+                this.reconnectTimer = null;
+            }
+            this.onStatusChangeCallback?.("connected");
             this.checkReady();
         };
 
         this.ws.onmessage = (event) => {
             try {
                 const msg = JSON.parse(event.data) as WebsocketMessage;
-                // Basic validation or filtering could happen here, but we pass it all
                 this.onMessageCallback(msg);
             } catch (e) {
                 console.error("Failed to parse WebSocket message", e);
@@ -35,9 +59,31 @@ export class HeadlessWasmAdapter {
 
         this.ws.onclose = () => {
             this.isWsOpen = false;
+            if (this.isClosedExplicitly) {
+                 this.onStatusChangeCallback?.("disconnected");
+                 return;
+            }
+
+            this.onStatusChangeCallback?.("disconnected");
+            this.scheduleReconnect();
+        };
+
+        this.ws.onerror = (e) => {
+            console.error("WebSocket error:", e);
+            // onError usually precedes onClose, so we handle reconnection in onClose
         };
     }
 
+    private scheduleReconnect() {
+        if (this.reconnectTimer || this.isClosedExplicitly) return;
+
+        this.onStatusChangeCallback?.("reconnecting");
+        this.reconnectTimer = setTimeout(() => {
+            this.reconnectTimer = null;
+            this.connect();
+        }, 3000); // Retry every 3 seconds
+    }
+
     async init() {
         try {
             // Load the WASM module dynamically from public folder
@@ -54,7 +100,13 @@ export class HeadlessWasmAdapter {
     private checkReady() {
         if (this.repl && this.isWsOpen && this.onReadyCallback) {
             this.onReadyCallback();
-            this.onReadyCallback = undefined; // Only call once
+            // We don't clear onReadyCallback because we might need it again on reconnect?
+            // The original code cleared it. If the WASM REPL is already initialized,
+            // we probably want to signal "ready" again if we reconnect so the UI can perhaps re-print the prompt or status.
+            // But usually "ready" means "initial load done".
+            // Let's keep the original behavior of calling it once for now,
+            // assuming the UI handles reconnection status separately.
+            this.onReadyCallback = undefined;
         }
     }
 
@@ -68,7 +120,7 @@ export class HeadlessWasmAdapter {
             const result = JSON.parse(resultJson);
 
             if (result.status === "complete") {
-                if (this.isWsOpen) {
+                if (this.isWsOpen && this.ws) {
                     this.ws.send(JSON.stringify({
                         kind: WebsocketMessageKind.Input,
                         input: result.payload
@@ -90,7 +142,6 @@ export class HeadlessWasmAdapter {
 
     complete(line: string, cursor: number): { suggestions: string[], start: number } {
         if (!this.repl) return { suggestions: [], start: cursor };
-        // complete returns a JSON object { suggestions: [...], start: number }
         const resultJson = this.repl.complete(line, cursor);
         try {
             return JSON.parse(resultJson);
@@ -107,8 +158,14 @@ export class HeadlessWasmAdapter {
     }
 
     close() {
+        this.isClosedExplicitly = true;
+        if (this.reconnectTimer) {
+            clearTimeout(this.reconnectTimer);
+            this.reconnectTimer = null;
+        }
         if (this.ws) {
             this.ws.close();
+            this.ws = null;
         }
     }
 }

tavern/internal/www/src/pages/shellv2/components/ShellStatusBar.tsx

@@ -1,27 +1,64 @@
 import React from "react";
-import { Info } from "lucide-react";
+import { Info, Wifi, WifiOff, RefreshCw } from "lucide-react";
 import { Tooltip } from "@chakra-ui/react";
+import { ConnectionStatus } from "../../../lib/headless-adapter";
 
 interface ShellStatusBarProps {
   portalId: number | null;
   timeUntilCallback: string;
   isMissedCallback: boolean;
+  connectionStatus: ConnectionStatus;
 }
 
-const ShellStatusBar: React.FC<ShellStatusBarProps> = ({ portalId, timeUntilCallback, isMissedCallback }) => {
+const ShellStatusBar: React.FC<ShellStatusBarProps> = ({ portalId, timeUntilCallback, isMissedCallback, connectionStatus }) => {
+  const getConnectionIcon = () => {
+    switch (connectionStatus) {
+      case "connected":
+        return (
+          <Tooltip label="Connected to Shell">
+            <span className="text-green-500">
+              <Wifi size={18} />
+            </span>
+          </Tooltip>
+        );
+      case "disconnected":
+        return (
+          <Tooltip label="Disconnected from Shell">
+            <span className="text-red-500">
+              <WifiOff size={18} />
+            </span>
+          </Tooltip>
+        );
+      case "reconnecting":
+        return (
+          <Tooltip label="Reconnecting...">
+            <span className="text-yellow-500 animate-spin">
+              <RefreshCw size={18} />
+            </span>
+          </Tooltip>
+        );
+    }
+  };
+
   return (
     <div className="flex justify-between items-center mt-2 text-sm text-gray-400 h-6">
-      <div className="flex items-center gap-2">
-        {portalId ? (
-          <span className="text-green-500 font-semibold">Portal Active (ID: {portalId})</span>
-        ) : (
-          <div className="flex items-center gap-1 group relative cursor-help">
-            <span>non-interactive</span>
-            <Tooltip label="This shell is currently in non-interactive mode. Input will be asynchronously queued for the beacon and output will be submitted through beacon callbacks. To upgrade to an interactive low-latency shell, you may open a 'Portal' on the beacon, which leverages an established connection to provide low-latency interactivity.">
-              <span><Info size={14} /></span>
-            </Tooltip>
-          </div>
-        )}
+      <div className="flex items-center gap-4">
+        <div className="flex items-center">
+            {getConnectionIcon()}
+        </div>
+
+        <div className="flex items-center gap-2">
+          {portalId ? (
+            <span className="text-green-500 font-semibold">Portal Active (ID: {portalId})</span>
+          ) : (
+            <div className="flex items-center gap-1 group relative cursor-help">
+              <span>non-interactive</span>
+              <Tooltip label="This shell is currently in non-interactive mode. Input will be asynchronously queued for the beacon and output will be submitted through beacon callbacks. To upgrade to an interactive low-latency shell, you may open a 'Portal' on the beacon, which leverages an established connection to provide low-latency interactivity.">
+                <span><Info size={14} /></span>
+              </Tooltip>
+            </div>
+          )}
+        </div>
       </div>
 
       {timeUntilCallback && (

tavern/internal/www/src/pages/shellv2/components/__tests__/ShellStatusBar.test.tsx

@@ -0,0 +1,55 @@
+import { render, screen } from "@testing-library/react";
+import ShellStatusBar from "../ShellStatusBar";
+import { ConnectionStatus } from "../../../../lib/headless-adapter";
+import React from "react";
+import { expect, describe, it } from "vitest";
+import "@testing-library/jest-dom";
+
+describe("ShellStatusBar", () => {
+    const defaultProps = {
+        portalId: null,
+        timeUntilCallback: "",
+        isMissedCallback: false,
+        connectionStatus: "connected" as ConnectionStatus,
+    };
+
+    it("renders connected status correctly", () => {
+        const { container } = render(<ShellStatusBar {...defaultProps} connectionStatus="connected" />);
+        // Wifi icon should be present (green)
+        const connectedIcon = container.querySelector(".text-green-500");
+        expect(connectedIcon).toBeInTheDocument();
+    });
+
+    it("renders disconnected status correctly", () => {
+        const { container } = render(<ShellStatusBar {...defaultProps} connectionStatus="disconnected" />);
+        const disconnectedIcon = container.querySelector(".text-red-500");
+        expect(disconnectedIcon).toBeInTheDocument();
+    });
+
+    it("renders reconnecting status correctly", () => {
+        const { container } = render(<ShellStatusBar {...defaultProps} connectionStatus="reconnecting" />);
+        const reconnectingIcon = container.querySelector(".text-yellow-500.animate-spin");
+        expect(reconnectingIcon).toBeInTheDocument();
+    });
+
+    it("displays portal active message when portalId is present", () => {
+        render(<ShellStatusBar {...defaultProps} portalId={123} />);
+        expect(screen.getByText("Portal Active (ID: 123)")).toBeInTheDocument();
+    });
+
+    it("displays non-interactive message when portalId is null", () => {
+        render(<ShellStatusBar {...defaultProps} portalId={null} />);
+        expect(screen.getByText("non-interactive")).toBeInTheDocument();
+    });
+
+    it("displays callback timer", () => {
+        render(<ShellStatusBar {...defaultProps} timeUntilCallback="in 5 minutes" />);
+        expect(screen.getByText("in 5 minutes")).toBeInTheDocument();
+    });
+
+    it("highlights missed callback", () => {
+        render(<ShellStatusBar {...defaultProps} timeUntilCallback="expected 5 minutes ago" isMissedCallback={true} />);
+        const timer = screen.getByText("expected 5 minutes ago");
+        expect(timer).toHaveClass("text-red-500 font-bold");
+    });
+});