From 5aedb17cb5ce9aa3216b353e221fcf89408af13c Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Fri, 27 Feb 2026 03:55:56 +0000 Subject: [PATCH] feat(shellv2): implement websocket reconnection logic and status indicator - Modified `HeadlessWasmAdapter` to support reconnection logic and expose connection status. - Updated `useShellTerminal` hook to manage connection status state and block input when disconnected. - Updated `ShellStatusBar` to display connection status (Connected, Disconnected, Reconnecting) with tooltips. - Added unit tests for `ShellStatusBar` to verify rendering of different states. Co-authored-by: KCarretto <16250309+KCarretto@users.noreply.github.com> --- package-lock.json | 6 + .../internal/www/src/lib/headless-adapter.ts | 71 ++++++++- .../shellv2/components/ShellStatusBar.tsx | 63 ++++++-- .../__tests__/ShellStatusBar.test.tsx | 55 +++++++ .../pages/shellv2/hooks/useShellTerminal.ts | 141 ++++++++++-------- .../internal/www/src/pages/shellv2/index.tsx | 4 +- 6 files changed, 257 insertions(+), 83 deletions(-) create mode 100644 package-lock.json create mode 100644 tavern/internal/www/src/pages/shellv2/components/__tests__/ShellStatusBar.test.tsx diff --git a/package-lock.json b/package-lock.json new file mode 100644 index 000000000..4ca926f5d --- /dev/null +++ b/package-lock.json @@ -0,0 +1,6 @@ +{ + "name": "app", + "lockfileVersion": 3, + "requires": true, + "packages": {} +} diff --git a/tavern/internal/www/src/lib/headless-adapter.ts b/tavern/internal/www/src/lib/headless-adapter.ts index 526360536..379a0e0c5 100644 --- a/tavern/internal/www/src/lib/headless-adapter.ts +++ b/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; } } } diff --git a/tavern/internal/www/src/pages/shellv2/components/ShellStatusBar.tsx b/tavern/internal/www/src/pages/shellv2/components/ShellStatusBar.tsx index 35056d298..0335083b7 100644 --- a/tavern/internal/www/src/pages/shellv2/components/ShellStatusBar.tsx +++ b/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 = ({ portalId, timeUntilCallback, isMissedCallback }) => { +const ShellStatusBar: React.FC = ({ portalId, timeUntilCallback, isMissedCallback, connectionStatus }) => { + const getConnectionIcon = () => { + switch (connectionStatus) { + case "connected": + return ( + + + + + + ); + case "disconnected": + return ( + + + + + + ); + case "reconnecting": + return ( + + + + + + ); + } + }; + return (
-
- {portalId ? ( - Portal Active (ID: {portalId}) - ) : ( -
- non-interactive - - - -
- )} +
+
+ {getConnectionIcon()} +
+ +
+ {portalId ? ( + Portal Active (ID: {portalId}) + ) : ( +
+ non-interactive + + + +
+ )} +
{timeUntilCallback && ( diff --git a/tavern/internal/www/src/pages/shellv2/components/__tests__/ShellStatusBar.test.tsx b/tavern/internal/www/src/pages/shellv2/components/__tests__/ShellStatusBar.test.tsx new file mode 100644 index 000000000..620936c04 --- /dev/null +++ b/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(); + // Wifi icon should be present (green) + const connectedIcon = container.querySelector(".text-green-500"); + expect(connectedIcon).toBeInTheDocument(); + }); + + it("renders disconnected status correctly", () => { + const { container } = render(); + const disconnectedIcon = container.querySelector(".text-red-500"); + expect(disconnectedIcon).toBeInTheDocument(); + }); + + it("renders reconnecting status correctly", () => { + const { container } = render(); + const reconnectingIcon = container.querySelector(".text-yellow-500.animate-spin"); + expect(reconnectingIcon).toBeInTheDocument(); + }); + + it("displays portal active message when portalId is present", () => { + render(); + expect(screen.getByText("Portal Active (ID: 123)")).toBeInTheDocument(); + }); + + it("displays non-interactive message when portalId is null", () => { + render(); + expect(screen.getByText("non-interactive")).toBeInTheDocument(); + }); + + it("displays callback timer", () => { + render(); + expect(screen.getByText("in 5 minutes")).toBeInTheDocument(); + }); + + it("highlights missed callback", () => { + render(); + const timer = screen.getByText("expected 5 minutes ago"); + expect(timer).toHaveClass("text-red-500 font-bold"); + }); +}); diff --git a/tavern/internal/www/src/pages/shellv2/hooks/useShellTerminal.ts b/tavern/internal/www/src/pages/shellv2/hooks/useShellTerminal.ts index 063445a41..1f53619ca 100644 --- a/tavern/internal/www/src/pages/shellv2/hooks/useShellTerminal.ts +++ b/tavern/internal/www/src/pages/shellv2/hooks/useShellTerminal.ts @@ -2,7 +2,7 @@ import { useEffect, useRef, useState, useCallback } from "react"; import { Terminal } from "@xterm/xterm"; import { FitAddon } from "xterm-addon-fit"; import "@xterm/xterm/css/xterm.css"; -import { HeadlessWasmAdapter } from "../../../lib/headless-adapter"; +import { HeadlessWasmAdapter, ConnectionStatus } from "../../../lib/headless-adapter"; import { WebsocketControlFlowSignal, WebsocketMessage, WebsocketMessageKind } from "../websocket"; import docsData from "../../../assets/eldritch-docs.json"; import { moveWordLeft, moveWordRight, highlightPythonSyntax, loadHistory, saveHistory } from "./shellUtils"; @@ -32,6 +32,7 @@ export const useShellTerminal = ( const termInstance = useRef(null); const adapter = useRef(null); const [connectionError, setConnectionError] = useState(null); + const [connectionStatus, setConnectionStatus] = useState("disconnected"); // Shell state const shellState = useRef({ @@ -74,16 +75,19 @@ export const useShellTerminal = ( // Ref for late checkin to access in event handlers const isLateCheckinRef = useRef(isLateCheckin); + const connectionStatusRef = useRef(connectionStatus); useEffect(() => { isLateCheckinRef.current = isLateCheckin; + connectionStatusRef.current = connectionStatus; if (termInstance.current) { + const isDimmed = isLateCheckin || connectionStatus !== "connected"; termInstance.current.options.theme = { - foreground: isLateCheckin ? "#777777" : "#d4d4d4", + foreground: isDimmed ? "#777777" : "#d4d4d4", background: "#1e1e1e", }; } - }, [isLateCheckin]); + }, [isLateCheckin, connectionStatus]); const redrawLine = useCallback(() => { const term = termInstance.current; @@ -279,6 +283,9 @@ export const useShellTerminal = ( } }, [tooltipState.visible]); + const shellNodeId = shellData?.node?.id; + const shellClosedAt = shellData?.node?.closedAt; + useEffect(() => { if (!termRef.current || loading) return; @@ -292,12 +299,12 @@ export const useShellTerminal = ( return; } - if (!shellData?.node) { + if (!shellNodeId) { setConnectionError("Shell not found."); return; } - if (shellData.node.closedAt) { + if (shellClosedAt) { setConnectionError("This shell session is closed."); return; } @@ -308,7 +315,7 @@ export const useShellTerminal = ( macOptionIsMeta: true, theme: { background: "#1e1e1e", - foreground: isLateCheckin ? "#777777" : "#d4d4d4", + foreground: (isLateCheckinRef.current || connectionStatusRef.current !== "connected") ? "#777777" : "#d4d4d4", }, fontFamily: 'Menlo, Monaco, "Courier New", monospace', fontSize: 18, @@ -442,73 +449,82 @@ export const useShellTerminal = ( const scheme = window.location.protocol === "https:" ? "wss" : "ws"; const url = `${scheme}://${window.location.host}/shellv2/ws?shell_id=${shellId}`; - adapter.current = new HeadlessWasmAdapter(url, (msg: WebsocketMessage) => { - const term = termInstance.current; - if (!term) return; + adapter.current = new HeadlessWasmAdapter( + url, + (msg: WebsocketMessage) => { + const term = termInstance.current; + if (!term) return; - // Clear current input line(s) before printing output - const prevRows = lastBufferHeight.current; - if (prevRows > 0) { - term.write(`\x1b[${prevRows}A`); - } - term.write("\r\x1b[J"); + // Clear current input line(s) before printing output + const prevRows = lastBufferHeight.current; + if (prevRows > 0) { + term.write(`\x1b[${prevRows}A`); + } + term.write("\r\x1b[J"); - // Process message content - let content = ""; - let color = ""; - - switch (msg.kind) { - case WebsocketMessageKind.Output: - content = msg.output; - break; - case WebsocketMessageKind.TaskError: - content = msg.error; - color = "\x1b[31m"; // Red - break; - case WebsocketMessageKind.Error: - content = msg.error; - color = "\x1b[31m"; // Red - break; - case WebsocketMessageKind.ControlFlow: - if (msg.signal === WebsocketControlFlowSignal.TaskQueued && msg.message) { - content = msg.message + "\n"; - color = "\x1b[33m"; // Yellow - } else if (msg.signal === WebsocketControlFlowSignal.PortalUpgrade && msg.portal_id) { - setPortalId(msg.portal_id); - } - // Handle other control signals if needed - break; - case WebsocketMessageKind.OutputFromOtherStream: - content = msg.output; - break; - } + // Process message content + let content = ""; + let color = ""; - if (content) { - const formatted = content.replace(/\n/g, "\r\n"); - if (color) { - term.write(color + formatted + "\x1b[0m"); - } else { - term.write(formatted); + switch (msg.kind) { + case WebsocketMessageKind.Output: + content = msg.output; + break; + case WebsocketMessageKind.TaskError: + content = msg.error; + color = "\x1b[31m"; // Red + break; + case WebsocketMessageKind.Error: + content = msg.error; + color = "\x1b[31m"; // Red + break; + case WebsocketMessageKind.ControlFlow: + if (msg.signal === WebsocketControlFlowSignal.TaskQueued && msg.message) { + content = msg.message + "\n"; + color = "\x1b[33m"; // Yellow + } else if (msg.signal === WebsocketControlFlowSignal.PortalUpgrade && msg.portal_id) { + setPortalId(msg.portal_id); + } + // Handle other control signals if needed + break; + case WebsocketMessageKind.OutputFromOtherStream: + content = msg.output; + break; } - // Ensure there is a newline after output if not present, so prompt is on new line - if (!content.endsWith('\n')) { - term.write("\r\n"); + if (content) { + const formatted = content.replace(/\n/g, "\r\n"); + if (color) { + term.write(color + formatted + "\x1b[0m"); + } else { + term.write(formatted); + } + + // Ensure there is a newline after output if not present, so prompt is on new line + if (!content.endsWith('\n')) { + term.write("\r\n"); + } } - } - // Reset input line state and redraw it at the bottom - lastBufferHeight.current = 0; - redrawLine(); - }, () => { - termInstance.current?.write("Connected to Tavern.\r\n>>> "); - }); + // Reset input line state and redraw it at the bottom + lastBufferHeight.current = 0; + redrawLine(); + }, + () => { + termInstance.current?.write("Connected to Tavern.\r\n>>> "); + }, + (status: ConnectionStatus) => { + setConnectionStatus(status); + } + ); adapter.current.init(); termInstance.current.onData((data) => { // Check for late checkin and block input if (isLateCheckinRef.current) return; + // Check for connection status and block input + if (connectionStatusRef.current !== "connected") return; const code = data.charCodeAt(0); const state = shellState.current; @@ -834,7 +850,7 @@ export const useShellTerminal = ( termInstance.current?.dispose(); if (redrawTimeoutRef.current) clearTimeout(redrawTimeoutRef.current); }; - }, [shellId, loading, error, shellData, setPortalId, redrawLine, updateCompletionsUI, applyCompletion]); + }, [shellId, loading, error, shellNodeId, shellClosedAt, setPortalId, redrawLine, updateCompletionsUI, applyCompletion]); return { termRef, @@ -845,6 +861,7 @@ export const useShellTerminal = ( completionIndex, handleMouseMove, tooltipState, - handleCompletionSelect: applyCompletion + handleCompletionSelect: applyCompletion, + connectionStatus }; }; diff --git a/tavern/internal/www/src/pages/shellv2/index.tsx b/tavern/internal/www/src/pages/shellv2/index.tsx index 3222f59c1..b812bf0bc 100644 --- a/tavern/internal/www/src/pages/shellv2/index.tsx +++ b/tavern/internal/www/src/pages/shellv2/index.tsx @@ -32,7 +32,8 @@ const ShellV2 = () => { completionIndex, handleMouseMove, tooltipState, - handleCompletionSelect + handleCompletionSelect, + connectionStatus } = useShellTerminal(shellId, loading, error, shellData, setPortalId, isLateCheckin); if (connectionError) { @@ -67,6 +68,7 @@ const ShellV2 = () => { portalId={portalId} timeUntilCallback={timeUntilCallback} isMissedCallback={isMissedCallback} + connectionStatus={connectionStatus} />