diff --git a/web/bin/vt b/web/bin/vt index 7ba4e9a7..d4a5c65f 100755 --- a/web/bin/vt +++ b/web/bin/vt @@ -141,6 +141,13 @@ if [ -n "$VIBETUNNEL_PREFER_DERIVED_DATA" ] && [ -n "$VIBETUNNEL_BIN" ]; then fi fi +# Handle safe commands first that work both inside and outside sessions +# This must come BEFORE the session check to avoid the recursive session error +if [[ "$1" == "status" || "$1" == "version" || "$1" == "--version" ]]; then + # These commands can run safely inside or outside a session + exec "$VIBETUNNEL_BIN" "$@" +fi + # Check if we're already inside a VibeTunnel session if [ -n "$VIBETUNNEL_SESSION_ID" ]; then # Special case: handle 'vt title' command inside a session @@ -159,6 +166,7 @@ if [ -n "$VIBETUNNEL_SESSION_ID" ]; then exit 1 fi + # For all other commands, block recursive sessions echo "Error: Already inside a VibeTunnel session (ID: $VIBETUNNEL_SESSION_ID). Recursive VibeTunnel sessions are not supported." >&2 echo "If you need to run commands, use them directly without the 'vt' prefix." >&2 exit 1 @@ -365,6 +373,7 @@ if [[ $# -eq 0 || "$1" == "--help" || "$1" == "-h" ]]; then exit 0 fi + # Handle 'vt title' command when not inside a session if [[ "$1" == "title" ]]; then echo "Error: 'vt title' can only be used inside a VibeTunnel session." >&2 @@ -372,12 +381,6 @@ if [[ "$1" == "title" ]]; then exit 1 fi -# Handle 'vt status' command -if [[ "$1" == "status" ]]; then - # Use vibetunnel CLI to show status via socket - exec "$VIBETUNNEL_BIN" status -fi - # Handle 'vt follow' command if [[ "$1" == "follow" ]]; then # Detect if we're in a worktree diff --git a/web/src/server/pty/socket-client.ts b/web/src/server/pty/socket-client.ts index dad37209..90ac8fa6 100644 --- a/web/src/server/pty/socket-client.ts +++ b/web/src/server/pty/socket-client.ts @@ -6,12 +6,20 @@ import { EventEmitter } from 'events'; import * as net from 'net'; import { createLogger } from '../utils/logger.js'; import { + type ControlCommand, type ErrorMessage, + frameMessage, + type GitEventNotify, + type GitFollowRequest, + type KillCommand, MessageBuilder, MessageParser, + type MessagePayload, MessageType, parsePayload, + type ResizeCommand, type StatusUpdate, + type UpdateTitleCommand, } from './socket-protocol.js'; const logger = createLogger('socket-client'); @@ -20,8 +28,8 @@ export interface SocketClientEvents { connect: () => void; disconnect: (error?: Error) => void; error: (error: Error) => void; - status: (status: StatusUpdate) => void; - serverError: (error: ErrorMessage) => void; + // Message-specific events are emitted using MessageType enum names + // e.g., 'STATUS_UPDATE', 'ERROR', 'HEARTBEAT', etc. } /** @@ -156,23 +164,14 @@ export class VibeTunnelSocketClient extends EventEmitter { try { const data = parsePayload(type, payload); - switch (type) { - case MessageType.STATUS_UPDATE: - this.emit('status', data as StatusUpdate); - break; + // Emit event with message type enum name + this.emit(MessageType[type], data); - case MessageType.ERROR: - this.emit('serverError', data as ErrorMessage); - break; - - case MessageType.HEARTBEAT: - this.lastHeartbeat = Date.now(); - // Echo heartbeat back - this.sendHeartbeat(); - break; - - default: - logger.debug(`Received unexpected message type: ${type}`); + // Handle heartbeat + if (type === MessageType.HEARTBEAT) { + this.lastHeartbeat = Date.now(); + // Echo heartbeat back + this.sendHeartbeat(); } } catch (error) { logger.error('Failed to parse message:', error); @@ -261,6 +260,104 @@ export class VibeTunnelSocketClient extends EventEmitter { return this.send(MessageBuilder.status(app, status, extra)); } + /** + * Send a message with type-safe payload + */ + public sendMessage(type: T, payload: MessagePayload): boolean { + const message = this.buildMessage(type, payload); + return this.send(message); + } + + /** + * Send a message and wait for a response + */ + public async sendMessageWithResponse( + requestType: TRequest, + payload: MessagePayload, + responseType: TResponse, + timeout = 5000 + ): Promise> { + return new Promise((resolve, reject) => { + const timer = setTimeout(() => { + this.off(MessageType[responseType], handleResponse); + this.off('error', handleError); + reject(new Error(`Request timeout waiting for ${MessageType[responseType]}`)); + }, timeout); + + const handleResponse = (data: MessagePayload) => { + clearTimeout(timer); + this.off('error', handleError); + resolve(data); + }; + + const handleError = (error: Error | ErrorMessage) => { + clearTimeout(timer); + this.off(MessageType[responseType], handleResponse); + if ('message' in error) { + reject(new Error(error.message)); + } else { + reject(error); + } + }; + + // Listen for response + this.once(MessageType[responseType], handleResponse); + this.once('error', handleError); + + const sent = this.sendMessage(requestType, payload); + if (!sent) { + clearTimeout(timer); + this.off(MessageType[responseType], handleResponse); + this.off('error', handleError); + reject(new Error('Failed to send message')); + } + }); + } + + /** + * Build a message buffer from type and payload + */ + private buildMessage(type: T, payload: MessagePayload): Buffer { + switch (type) { + case MessageType.STDIN_DATA: + return MessageBuilder.stdin(payload as string); + case MessageType.CONTROL_CMD: { + const cmd = payload as ControlCommand; + switch (cmd.cmd) { + case 'resize': + return MessageBuilder.resize((cmd as ResizeCommand).cols, (cmd as ResizeCommand).rows); + case 'kill': + return MessageBuilder.kill((cmd as KillCommand).signal); + case 'reset-size': + return MessageBuilder.resetSize(); + case 'update-title': + return MessageBuilder.updateTitle((cmd as UpdateTitleCommand).title); + default: + // For generic control commands, use frameMessage directly + return frameMessage(MessageType.CONTROL_CMD, cmd); + } + } + case MessageType.STATUS_UPDATE: { + const statusPayload = payload as StatusUpdate; + return MessageBuilder.status( + statusPayload.app, + statusPayload.status, + statusPayload.extra as Record | undefined + ); + } + case MessageType.HEARTBEAT: + return MessageBuilder.heartbeat(); + case MessageType.STATUS_REQUEST: + return MessageBuilder.statusRequest(); + case MessageType.GIT_FOLLOW_REQUEST: + return MessageBuilder.gitFollowRequest(payload as GitFollowRequest); + case MessageType.GIT_EVENT_NOTIFY: + return MessageBuilder.gitEventNotify(payload as GitEventNotify); + default: + throw new Error(`Unsupported message type: ${type}`); + } + } + /** * Send heartbeat */ diff --git a/web/src/server/pty/socket-protocol.ts b/web/src/server/pty/socket-protocol.ts index 85d20e78..c31e24b8 100644 --- a/web/src/server/pty/socket-protocol.ts +++ b/web/src/server/pty/socket-protocol.ts @@ -136,6 +136,30 @@ export interface GitEventAck { handled: boolean; } +/** + * Type-safe mapping of message types to their payload types + */ +export type MessagePayloadMap = { + [MessageType.STDIN_DATA]: string; + [MessageType.CONTROL_CMD]: ControlCommand; + [MessageType.STATUS_UPDATE]: StatusUpdate; + [MessageType.HEARTBEAT]: Record; + [MessageType.ERROR]: ErrorMessage; + [MessageType.STATUS_REQUEST]: StatusRequest; + [MessageType.STATUS_RESPONSE]: StatusResponse; + [MessageType.GIT_FOLLOW_REQUEST]: GitFollowRequest; + [MessageType.GIT_FOLLOW_RESPONSE]: GitFollowResponse; + [MessageType.GIT_EVENT_NOTIFY]: GitEventNotify; + [MessageType.GIT_EVENT_ACK]: GitEventAck; +}; + +/** + * Get the payload type for a given message type + */ +export type MessagePayload = T extends keyof MessagePayloadMap + ? MessagePayloadMap[T] + : never; + /** * Frame a message for transmission */ diff --git a/web/src/server/socket-api-client.ts b/web/src/server/socket-api-client.ts index b70e47cc..364b0926 100644 --- a/web/src/server/socket-api-client.ts +++ b/web/src/server/socket-api-client.ts @@ -12,6 +12,7 @@ import { type GitEventNotify, type GitFollowRequest, type GitFollowResponse, + type MessagePayload, MessageType, } from './pty/socket-protocol.js'; import { createLogger } from './utils/logger.js'; @@ -52,9 +53,9 @@ export class SocketApiClient { /** * Send a request and wait for response */ - private async sendRequest( - type: MessageType, - payload: TRequest, + private async sendRequest( + type: TRequest, + payload: MessagePayload, responseType: MessageType, timeout = 5000 ): Promise { @@ -64,106 +65,18 @@ export class SocketApiClient { const client = new VibeTunnelSocketClient(this.controlSocketPath); - return new Promise((resolve, reject) => { - const timer = setTimeout(() => { - client.disconnect(); - reject(new Error('Request timeout')); - }, timeout); - - let responseReceived = false; - - client.on('error', (error) => { - clearTimeout(timer); - if (!responseReceived) { - reject(error); - } - }); - - // Handle the specific response type we're expecting - const handleMessage = (msgType: MessageType, data: unknown) => { - if (msgType === responseType) { - responseReceived = true; - clearTimeout(timer); - client.disconnect(); - resolve(data as TResponse); - } else if (msgType === MessageType.ERROR) { - responseReceived = true; - clearTimeout(timer); - client.disconnect(); - reject(new Error((data as { message?: string }).message || 'Server error')); - } - }; - - // Override the handleMessage method to intercept messages - (client as unknown as { handleMessage: typeof handleMessage }).handleMessage = handleMessage; - - client - .connect() - .then(() => { - // Send the request - let message: unknown; - switch (type) { - case MessageType.STATUS_REQUEST: - message = (client as unknown as { send: (msg: unknown) => unknown }).send( - ( - client as unknown as { - constructor: { - prototype: { - constructor: { - MessageBuilder: Record unknown>; - }; - }; - }; - } - ).constructor.prototype.constructor.MessageBuilder.statusRequest() - ); - break; - case MessageType.GIT_FOLLOW_REQUEST: - message = (client as unknown as { send: (msg: unknown) => unknown }).send( - ( - client as unknown as { - constructor: { - prototype: { - constructor: { - MessageBuilder: Record unknown>; - }; - }; - }; - } - ).constructor.prototype.constructor.MessageBuilder.gitFollowRequest(payload) - ); - break; - case MessageType.GIT_EVENT_NOTIFY: - message = (client as unknown as { send: (msg: unknown) => unknown }).send( - ( - client as unknown as { - constructor: { - prototype: { - constructor: { - MessageBuilder: Record unknown>; - }; - }; - }; - } - ).constructor.prototype.constructor.MessageBuilder.gitEventNotify(payload) - ); - break; - default: - clearTimeout(timer); - reject(new Error(`Unsupported message type: ${type}`)); - return; - } - - if (!message) { - clearTimeout(timer); - reject(new Error('Failed to send request')); - } - }) - .catch((error) => { - clearTimeout(timer); - reject(error); - }); - }); + try { + await client.connect(); + const response = await client.sendMessageWithResponse(type, payload, responseType, timeout); + return response as TResponse; + } catch (error) { + if (error instanceof Error && error.message.includes('ENOENT')) { + throw new Error('VibeTunnel server is not running'); + } + throw error; + } finally { + client.disconnect(); + } } /** @@ -176,7 +89,7 @@ export class SocketApiClient { try { // Send STATUS_REQUEST and wait for STATUS_RESPONSE - const response = await this.sendRequest, ServerStatus>( + const response = await this.sendRequest( MessageType.STATUS_REQUEST, {}, MessageType.STATUS_RESPONSE @@ -192,7 +105,7 @@ export class SocketApiClient { * Enable or disable Git follow mode */ async setFollowMode(request: GitFollowRequest): Promise { - return this.sendRequest( + return this.sendRequest( MessageType.GIT_FOLLOW_REQUEST, request, MessageType.GIT_FOLLOW_RESPONSE @@ -203,7 +116,7 @@ export class SocketApiClient { * Send Git event notification */ async sendGitEvent(event: GitEventNotify): Promise { - return this.sendRequest( + return this.sendRequest( MessageType.GIT_EVENT_NOTIFY, event, MessageType.GIT_EVENT_ACK