diff --git a/src/web/client/components/app.tsx b/src/web/client/components/app.tsx index a672d385..4698c988 100644 --- a/src/web/client/components/app.tsx +++ b/src/web/client/components/app.tsx @@ -1,4 +1,4 @@ -import { useState, useEffect, useCallback } from 'react' +import { useState, useEffect, useCallback, useRef } from 'react' import type { PTYSessionInfo } from 'opencode-pty/web/shared/types' import { useWebSocket } from '../hooks/use-web-socket.ts' @@ -11,40 +11,24 @@ import { api } from '../../shared/api-client.ts' export function App() { const [sessions, setSessions] = useState([]) const [activeSession, setActiveSession] = useState(null) - const [rawOutput, setRawOutput] = useState('') + const terminalRef = useRef(null) - const [connected, setConnected] = useState(false) const [wsMessageCount, setWsMessageCount] = useState(0) const [sessionUpdateCount, setSessionUpdateCount] = useState(0) - const { - connected: wsConnected, - subscribeWithRetry, - sendInput, - } = useWebSocket({ + const { connected, sendInput } = useWebSocket({ activeSession, + onSnapshot: useCallback((snapshot: string) => { + terminalRef.current?.reset(snapshot) + }, []), onRawData: useCallback((rawData: string) => { - setRawOutput((prev) => { - const newOutput = prev + rawData - return newOutput - }) + terminalRef.current?.write(rawData) setWsMessageCount((prev) => prev + 1) }, []), onSessionList: useCallback( (newSessions: PTYSessionInfo[], autoSelected: PTYSessionInfo | null) => { setSessions(newSessions) - if (!autoSelected) { - return - } - setActiveSession(autoSelected) - api.session.buffer - .raw({ id: autoSelected.id }) - .then((data) => { - setRawOutput(data.raw) - }) - .catch((error) => { - console.error('Failed to fetch initial raw buffer for auto-selected session', error) - }) + if (autoSelected) setActiveSession(autoSelected) }, [] ), @@ -65,11 +49,6 @@ export function App() { }, []), }) - // Update connected from wsConnected - useEffect(() => { - setConnected(wsConnected) - }, [wsConnected]) - // Periodic session list sync every 10 seconds useEffect(() => { const syncInterval = setInterval(async () => { @@ -83,15 +62,9 @@ export function App() { return () => clearInterval(syncInterval) }, []) - const { handleSessionClick, handleSendInput, handleKillSession } = useSessionManager({ + const { handleSendInput, handleKillSession } = useSessionManager({ activeSession, - setActiveSession, - subscribeWithRetry, sendInput, - wsConnected, - onRawOutputUpdate: useCallback((rawOutput: string) => { - setRawOutput(rawOutput) - }, []), }) return ( @@ -99,7 +72,7 @@ export function App() {
@@ -114,15 +87,15 @@ export function App() {
- Debug: {rawOutput.length} chars, active: {activeSession?.id || 'none'}, WS raw_data:{' '} - {wsMessageCount}, session_updates: {sessionUpdateCount} + Debug: active: {activeSession?.id || 'none'}, WS raw_data: {wsMessageCount}, + session_updates: {sessionUpdateCount}
) : ( diff --git a/src/web/client/components/terminal-renderer.tsx b/src/web/client/components/terminal-renderer.tsx index e79cc085..6eaa7035 100644 --- a/src/web/client/components/terminal-renderer.tsx +++ b/src/web/client/components/terminal-renderer.tsx @@ -13,7 +13,6 @@ declare global { } interface RawTerminalProps { - rawOutput: string onSendInput?: (data: string) => void onInterrupt?: () => void disabled?: boolean @@ -27,28 +26,15 @@ export class RawTerminal extends React.Component { override componentDidMount() { this.initializeTerminal() - if (this.xtermInstance && this.props.rawOutput) { - this.xtermInstance.write(this.props.rawOutput) - } } - override componentDidUpdate(prevProps: RawTerminalProps) { - if (!this.xtermInstance) return - - const currentData = this.props.rawOutput - const prevData = prevProps.rawOutput + // In band (RIS): Terminal.reset() keeps queued writes and any half-parsed escape sequence. + reset(snapshot: string): void { + this.write(`\x1bc${snapshot}`) + } - // Optimized diff-based writing - only write new content - if (currentData.startsWith(prevData)) { - const newData = currentData.slice(prevData.length) - if (newData) { - this.xtermInstance.write(newData) - } - } else { - // Session switch/truncate/etc - clear and rewrite - this.xtermInstance.clear() - this.xtermInstance.write(currentData) - } + write(chunk: string): void { + this.xtermInstance?.write(chunk) } override componentWillUnmount() { diff --git a/src/web/client/hooks/session-socket.ts b/src/web/client/hooks/session-socket.ts new file mode 100644 index 00000000..173ea43f --- /dev/null +++ b/src/web/client/hooks/session-socket.ts @@ -0,0 +1,73 @@ +import { RETRY_DELAY } from '../../shared/constants.ts' +import type { + PTYSessionInfo, + WSMessageServerRawData, + WSMessageServerSessionList, + WSMessageServerSessionUpdate, + WSMessageServerSubscribedSession, +} from '../../shared/types.ts' + +export function connectSessionSocket( + url: string, + sessionId: string | undefined, + callbacks: { + onConnected: (connected: boolean) => void + onSnapshot?: (snapshot: string) => void + onRawData?: (rawData: string) => void + onSessionList: (sessions: PTYSessionInfo[]) => void + onSessionUpdate?: (session: PTYSessionInfo) => void + } +) { + let ws: WebSocket + let retry: ReturnType | undefined + let closed = false + + const connect = () => { + ws = new WebSocket(url) + ws.onopen = () => { + if (closed) return + callbacks.onConnected(true) + ws.send(JSON.stringify({ type: 'session_list' })) + if (sessionId) ws.send(JSON.stringify({ type: 'subscribe', sessionId })) + } + ws.onmessage = (event) => { + if (closed) return + try { + const data = JSON.parse(event.data) as + | WSMessageServerSubscribedSession + | WSMessageServerRawData + | WSMessageServerSessionList + | WSMessageServerSessionUpdate + if (data.type === 'subscribed' && data.sessionId === sessionId) { + callbacks.onSnapshot?.(data.rawData) + } else if (data.type === 'raw_data' && data.session.id === sessionId) { + callbacks.onRawData?.(data.rawData) + } else if (data.type === 'session_list') { + callbacks.onSessionList(data.sessions || []) + } else if (data.type === 'session_update') { + callbacks.onSessionUpdate?.(data.session) + } + // Ignore malformed messages without interrupting the stream. + } catch {} + } + ws.onclose = () => { + if (closed) return + callbacks.onConnected(false) + retry = setTimeout(connect, RETRY_DELAY) + } + } + connect() + + return { + sendInput(sessionId: string, data: string): boolean { + if (ws.readyState !== WebSocket.OPEN) return false + ws.send(JSON.stringify({ type: 'input', sessionId, data })) + return true + }, + close() { + closed = true + clearTimeout(retry) + ws.close() + }, + } +} diff --git a/src/web/client/hooks/use-session-manager.ts b/src/web/client/hooks/use-session-manager.ts index 9411f9e7..b4d64163 100644 --- a/src/web/client/hooks/use-session-manager.ts +++ b/src/web/client/hooks/use-session-manager.ts @@ -5,75 +5,22 @@ import { api } from '../../shared/api-client' interface UseSessionManagerOptions { activeSession: PTYSessionInfo | null - setActiveSession: (session: PTYSessionInfo | null) => void - subscribeWithRetry: (sessionId: string) => void - sendInput?: (sessionId: string, data: string) => void - wsConnected?: boolean - onRawOutputUpdate?: (rawOutput: string) => void + sendInput: (sessionId: string, data: string) => boolean } -export function useSessionManager({ - activeSession, - setActiveSession, - subscribeWithRetry, - sendInput, - wsConnected, - onRawOutputUpdate, -}: UseSessionManagerOptions) { - const handleSessionClick = useCallback( - async (session: PTYSessionInfo) => { - try { - // Validate session object first - if (!session?.id) { - return - } - setActiveSession(session) - onRawOutputUpdate?.('') - // Subscribe to this session for live updates - subscribeWithRetry(session.id) - - try { - // Fetch raw buffer data only (processed output endpoint removed) - const rawData = await api.session.buffer - .raw({ id: session.id }) - .catch(() => ({ raw: '' })) - - // Call callback with raw data - onRawOutputUpdate?.(rawData.raw || '') - } catch { - onRawOutputUpdate?.('') - } - } catch { - // Ensure UI remains stable - onRawOutputUpdate?.('') - } - }, - [setActiveSession, subscribeWithRetry, onRawOutputUpdate] - ) - +export function useSessionManager({ activeSession, sendInput }: UseSessionManagerOptions) { const handleSendInput = useCallback( async (data: string) => { - if (!data || !activeSession) { + // WebSocket first; HTTP while it is not open (e.g. reconnecting). + if (!data || !activeSession || sendInput(activeSession.id, data)) { return } - - // Try WebSocket first if connected and available - if (wsConnected && sendInput) { - try { - sendInput(activeSession.id, data) - return - } catch (error) { - console.warn('WebSocket input failed, falling back to HTTP:', error) - } - } - - // HTTP fallback try { await api.session.input({ id: activeSession.id }, { data }) // eslint-disable-next-line no-empty } catch {} }, - [activeSession, wsConnected, sendInput] + [activeSession, sendInput] ) const handleKillSession = useCallback(async () => { @@ -97,7 +44,6 @@ export function useSessionManager({ }, [activeSession]) return { - handleSessionClick, handleSendInput, handleKillSession, } diff --git a/src/web/client/hooks/use-web-socket.ts b/src/web/client/hooks/use-web-socket.ts index a8201eed..b3957d94 100644 --- a/src/web/client/hooks/use-web-socket.ts +++ b/src/web/client/hooks/use-web-socket.ts @@ -1,17 +1,12 @@ import { useState, useEffect, useRef } from 'react' import type { PTYSessionInfo } from 'opencode-pty/web/shared/types' -import type { - WSMessageServer, - WSMessageServerRawData, - WSMessageServerSessionList, - WSMessageServerSessionUpdate, -} from 'opencode-pty/web/shared/types' -import { RETRY_DELAY, SKIP_AUTOSELECT_KEY } from 'opencode-pty/web/shared/constants' - +import { SKIP_AUTOSELECT_KEY } from 'opencode-pty/web/shared/constants' import { RouteBuilder } from 'opencode-pty/web/shared/route-builder' +import { connectSessionSocket } from './session-socket.ts' interface UseWebSocketOptions { activeSession: PTYSessionInfo | null + onSnapshot?: (snapshot: string) => void onRawData?: (rawData: string) => void onSessionList: (sessions: PTYSessionInfo[], autoSelected: PTYSessionInfo | null) => void onSessionUpdate?: (updatedSession: PTYSessionInfo) => void @@ -19,117 +14,43 @@ interface UseWebSocketOptions { export function useWebSocket({ activeSession, + onSnapshot, onRawData, onSessionList, onSessionUpdate, }: UseWebSocketOptions) { const [connected, setConnected] = useState(false) + const socketRef = useRef | null>(null) + const sessionId = activeSession?.id - const wsRef = useRef(null) - const activeSessionRef = useRef(null) - - // Keep ref in sync with activeSession useEffect(() => { - activeSessionRef.current = activeSession - }, [activeSession]) - - // Connect to WebSocket on mount - useEffect(() => { - const ws = new WebSocket( - `${RouteBuilder.websocket()}`.replace(/^\/ws/, `ws://${location.host}/ws`) - ) - ws.onopen = () => { - setConnected(true) - // Request initial session list - ws.send(JSON.stringify({ type: 'session_list' })) - // Resubscribe to active session if exists - if (activeSessionRef.current) { - ws.send(JSON.stringify({ type: 'subscribe', sessionId: activeSessionRef.current.id })) - } - } - ws.onmessage = (event) => { - try { - const data = JSON.parse(event.data) as WSMessageServer - if (data.type === 'session_list') { - const sessionListMsg = data as WSMessageServerSessionList - const sessions = sessionListMsg.sessions || [] - // Auto-select first running session if none selected (skip in tests that need empty state) - const shouldSkipAutoselect = localStorage.getItem(SKIP_AUTOSELECT_KEY) === 'true' - let autoSelected: PTYSessionInfo | null = null - if (sessions.length > 0 && !activeSession && !shouldSkipAutoselect) { - const runningSession = - sessions.find((s: PTYSessionInfo) => s.status === 'running') || null - autoSelected = runningSession || sessions[0] || null - if (autoSelected) { - activeSessionRef.current = autoSelected - // Subscribe to the auto-selected session for live updates - const readyState = wsRef.current?.readyState - - if (readyState === WebSocket.OPEN && wsRef.current) { - wsRef.current.send( - JSON.stringify({ type: 'subscribe', sessionId: autoSelected.id }) - ) - } else { - setTimeout( - (autoSelected) => { - const retryReadyState = wsRef.current?.readyState - if (retryReadyState === WebSocket.OPEN && wsRef.current) { - wsRef.current.send( - JSON.stringify({ type: 'subscribe', sessionId: autoSelected.id }) - ) - } - }, - RETRY_DELAY, - autoSelected - ) - } - } - } + const socket = connectSessionSocket( + RouteBuilder.websocket().replace(/^\/ws/, `ws://${location.host}/ws`), + sessionId, + { + onConnected: setConnected, + onSnapshot, + onRawData, + onSessionList: (sessions) => { + const skipAutoselect = localStorage.getItem(SKIP_AUTOSELECT_KEY) === 'true' + const autoSelected = + !sessionId && !skipAutoselect + ? sessions.find((s) => s.status === 'running') || sessions[0] || null + : null onSessionList(sessions, autoSelected) - } else if (data.type === 'session_update') { - const sessionUpdateMsg = data as WSMessageServerSessionUpdate - onSessionUpdate?.(sessionUpdateMsg.session) - } else if (data.type === 'raw_data') { - const rawDataMsg = data as WSMessageServerRawData - const isForActiveSession = rawDataMsg.session.id === activeSessionRef.current?.id - if (isForActiveSession) { - onRawData?.(rawDataMsg.rawData) - } - } - // eslint-disable-next-line no-empty - } catch {} - } - ws.onclose = () => { - setConnected(false) - } - ws.onerror = () => {} - wsRef.current = ws + }, + onSessionUpdate, + } + ) + socketRef.current = socket return () => { - ws.close() - } - }, [activeSession, onRawData, onSessionList, onSessionUpdate]) - - const subscribe = (sessionId: string) => { - if (wsRef.current?.readyState === WebSocket.OPEN) { - wsRef.current.send(JSON.stringify({ type: 'subscribe', sessionId })) - } - } - - const subscribeWithRetry = (sessionId: string) => { - if (wsRef.current?.readyState === WebSocket.OPEN) { - subscribe(sessionId) - } else { - setTimeout(() => { - subscribe(sessionId) - }, RETRY_DELAY) + socketRef.current = null + socket.close() } - } + }, [sessionId, onSnapshot, onRawData, onSessionList, onSessionUpdate]) - const sendInput = (sessionId: string, data: string) => { - if (wsRef.current?.readyState === WebSocket.OPEN) { - wsRef.current.send(JSON.stringify({ type: 'input', sessionId, data })) - } - } + const sendInput = (sessionId: string, data: string) => + socketRef.current?.sendInput(sessionId, data) ?? false - return { connected, subscribe, subscribeWithRetry, sendInput } + return { connected, sendInput } } diff --git a/src/web/server/handlers/websocket.ts b/src/web/server/handlers/websocket.ts index 1ad79567..3dfc1919 100644 --- a/src/web/server/handlers/websocket.ts +++ b/src/web/server/handlers/websocket.ts @@ -40,6 +40,7 @@ class WebSocketHandler { const response: WSMessageServerSubscribedSession = { type: 'subscribed', sessionId: message.sessionId, + rawData: manager.getRawBuffer(message.sessionId)?.raw ?? '', } ws.send(JSON.stringify(response)) } diff --git a/src/web/shared/types.ts b/src/web/shared/types.ts index aaa795b5..76c27c1e 100644 --- a/src/web/shared/types.ts +++ b/src/web/shared/types.ts @@ -66,6 +66,7 @@ export interface WSMessageServer { export interface WSMessageServerSubscribedSession extends WSMessageServer { type: 'subscribed' sessionId: string + rawData: string } export interface WSMessageServerUnsubscribedSession extends WSMessageServer { diff --git a/test/session-socket.test.ts b/test/session-socket.test.ts new file mode 100644 index 00000000..70ae3cd8 --- /dev/null +++ b/test/session-socket.test.ts @@ -0,0 +1,78 @@ +import { afterEach, beforeEach, expect, it, jest, mock } from 'bun:test' +import { connectSessionSocket } from '../src/web/client/hooks/session-socket.ts' +import { RETRY_DELAY } from '../src/web/shared/constants.ts' + +class FakeWebSocket { + static instances: FakeWebSocket[] = [] + onopen: (() => void) | null = null + onclose: (() => void) | null = null + onmessage: ((event: { data: string }) => void) | null = null + send = mock() + constructor() { + FakeWebSocket.instances.push(this) + } + close() { + this.onclose?.() + } + open() { + this.onopen?.() + } + message(data: unknown) { + this.onmessage?.({ data: JSON.stringify(data) }) + } +} + +const originalWebSocket = globalThis.WebSocket +let close: (() => void) | undefined +beforeEach(() => { + jest.useFakeTimers() + FakeWebSocket.instances = [] + globalThis.WebSocket = FakeWebSocket as unknown as typeof WebSocket +}) +afterEach(() => { + close?.() + globalThis.WebSocket = originalWebSocket + jest.useRealTimers() +}) + +it('reconnects and resubscribes after an unexpected close', () => { + const onConnected = mock() + close = connectSessionSocket('ws://localhost/ws', 'active', { + onConnected, + onSessionList: mock(), + }).close + FakeWebSocket.instances[0]?.open() + FakeWebSocket.instances[0]?.close() + expect(onConnected).toHaveBeenLastCalledWith(false) + jest.advanceTimersByTime(RETRY_DELAY - 1) + expect(FakeWebSocket.instances).toHaveLength(1) + jest.advanceTimersByTime(1) + expect(FakeWebSocket.instances).toHaveLength(2) + FakeWebSocket.instances[1]?.open() + expect(FakeWebSocket.instances[1]?.send).toHaveBeenCalledWith( + JSON.stringify({ type: 'subscribe', sessionId: 'active' }) + ) +}) + +it('ignores output from sessions and sockets that are no longer active', () => { + const callbacks = { + onConnected: mock(), + onSessionList: mock(), + onSnapshot: mock(), + onRawData: mock(), + } + const previous = connectSessionSocket('ws://localhost/ws', 'previous', callbacks) + const oldSocket = FakeWebSocket.instances[0] + previous.close() + close = connectSessionSocket('ws://localhost/ws', 'active', callbacks).close + const socket = FakeWebSocket.instances[1] + oldSocket?.message({ type: 'raw_data', session: { id: 'previous' }, rawData: 'late' }) + socket?.message({ type: 'subscribed', sessionId: 'previous', rawData: 'wrong snapshot' }) + socket?.message({ type: 'raw_data', session: { id: 'previous' }, rawData: 'wrong chunk' }) + socket?.message({ type: 'subscribed', sessionId: 'active', rawData: 'snapshot' }) + socket?.message({ type: 'raw_data', session: { id: 'active' }, rawData: 'chunk' }) + expect(callbacks.onRawData).toHaveBeenCalledTimes(1) + expect(callbacks.onRawData).toHaveBeenCalledWith('chunk') + expect(callbacks.onSnapshot).toHaveBeenCalledTimes(1) + expect(callbacks.onSnapshot).toHaveBeenCalledWith('snapshot') +}) diff --git a/test/websocket.test.ts b/test/websocket.test.ts index 19fd713e..561c8282 100644 --- a/test/websocket.test.ts +++ b/test/websocket.test.ts @@ -1,5 +1,9 @@ import { describe, it, expect, beforeAll, afterAll } from 'bun:test' -import { manager } from '../src/plugin/pty/manager.ts' +import { + manager, + registerRawOutputCallback, + removeRawOutputCallback, +} from '../src/plugin/pty/manager.ts' import type { CustomError, WSMessageServerError, @@ -67,6 +71,43 @@ describe('WebSocket Functionality', () => { }) describe('WebSocket Message Handling', () => { + it('should send the snapshot before subsequent live output without replaying it', async () => { + await using client = await ManagedTestClient.create(managedTestServer.server.getWsUrl()) + const snapshot = 'one\r\ntwo\r\n' + const session = manager.spawn({ + command: 'sh', + args: ['-c', 'stty -echo; printf "one\\ntwo\\n"; read line; printf "three\\n"'], + parentSessionId: managedTestServer.sessionId, + }) + // Output is delivered asynchronously, so registering after spawn misses nothing. + await new Promise((resolve) => { + const onOutput = () => { + if (manager.getRawBuffer(session.id)?.raw === snapshot) { + removeRawOutputCallback(onOutput) + resolve() + } + } + registerRawOutputCallback(onOutput) + }) + const subscribed = new Promise((resolve) => { + client.subscribedCallbacks.push(resolve) + }) + let liveOutput = '' + const live = new Promise((resolve) => { + client.rawDataCallbacks.push((message) => { + if (message.session.id === session.id) { + liveOutput += message.rawData + if (liveOutput.includes('three\r\n')) resolve() + } + }) + }) + client.send({ type: 'subscribe', sessionId: session.id }) + expect((await subscribed).rawData).toBe(snapshot) + manager.write(session.id, 'go\n') + await live + expect(liveOutput).toBe('three\r\n') + }, 1000) + it('should handle subscribe message', async () => { await using managedTestClient = await ManagedTestClient.create( managedTestServer.server.getWsUrl()