import type {
BroadcastMessage,
PendingRequest,
SystemMessage,
} from "@caelinsutch/types";
import { v4 as uuidv4 } from "uuid";
import WebSocket from "ws";
import { logger } from "../logger";
// Connection state
let ws: WebSocket | null = null;
let channel: string | null = null;
let isConnected = false;
let reconnectAttempts = 0;
const MAX_RECONNECT_ATTEMPTS = 5;
const RECONNECT_DELAY_MS = 2000;
// Pending requests waiting for responses
const pendingRequests = new Map<string, PendingRequest>();
// Default timeout for requests (30 seconds)
const DEFAULT_TIMEOUT_MS = 30000;
/**
* Connect to the relay server
*/
export async function connectToRelay(
urlOrPort: string | number = 3055,
channelName?: string,
): Promise<void> {
return new Promise((resolve, reject) => {
if (isConnected && ws) {
logger.info("Already connected to relay server");
resolve();
return;
}
const url =
typeof urlOrPort === "number" ? `ws://localhost:${urlOrPort}` : urlOrPort;
logger.info(`Connecting to relay server at ${url}...`);
ws = new WebSocket(url);
ws.on("open", () => {
logger.info("Connected to relay server");
isConnected = true;
reconnectAttempts = 0;
// Generate channel name if not provided
channel = channelName || generateChannelName();
// Join the channel
joinChannel(channel)
.then(() => {
logger.info(`Joined channel: ${channel}`);
resolve();
})
.catch(reject);
});
ws.on("message", (data) => {
handleMessage(data.toString());
});
ws.on("close", () => {
logger.info("Disconnected from relay server");
isConnected = false;
ws = null;
handleDisconnect();
});
ws.on("error", (error) => {
logger.error("WebSocket error", error.message);
if (!isConnected) {
reject(error);
}
});
});
}
/**
* Generate a random channel name
*/
function generateChannelName(): string {
const chars = "abcdefghijklmnopqrstuvwxyz0123456789";
let result = "";
for (let i = 0; i < 8; i++) {
result += chars.charAt(Math.floor(Math.random() * chars.length));
}
return result;
}
/**
* Join a channel on the relay server
*/
async function joinChannel(channelName: string): Promise<void> {
return new Promise((resolve, reject) => {
if (!ws || !isConnected) {
reject(new Error("Not connected to relay server"));
return;
}
const id = uuidv4();
// Set up handler for join response
const timeout = setTimeout(() => {
pendingRequests.delete(id);
reject(new Error("Join channel timeout"));
}, 10000);
pendingRequests.set(id, {
resolve: () => {
clearTimeout(timeout);
resolve();
},
reject: (err: unknown) => {
clearTimeout(timeout);
reject(err);
},
timeout,
lastActivity: Date.now(),
});
ws.send(
JSON.stringify({
id,
type: "join",
channel: channelName,
}),
);
});
}
/**
* Handle incoming WebSocket messages
*/
function handleMessage(data: string): void {
try {
const message = JSON.parse(data);
logger.debug("Received message", { type: message.type });
// Handle system messages (join confirmations, etc.)
if (message.type === "system") {
handleSystemMessage(message as SystemMessage);
return;
}
// Handle broadcast messages (responses from Figma plugin)
if (message.type === "broadcast") {
handleBroadcastMessage(message as BroadcastMessage);
return;
}
// Handle error messages
if (message.type === "error") {
logger.error("Relay error", message.message);
return;
}
} catch (error) {
logger.error("Error parsing message", error);
}
}
/**
* Handle system messages from relay
*/
function handleSystemMessage(message: SystemMessage): void {
// Check if it's a join confirmation
if (
typeof message.message === "object" &&
message.message !== null &&
"id" in message.message &&
"result" in message.message
) {
const { id, result } = message.message as { id: string; result: string };
const pending = pendingRequests.get(id);
if (pending) {
pendingRequests.delete(id);
pending.resolve(result);
}
}
}
/**
* Handle broadcast messages (responses from Figma plugin)
*/
function handleBroadcastMessage(message: BroadcastMessage): void {
const { id, result, error } = message.message || {};
if (!id) {
logger.debug("Broadcast message without ID (likely a command)");
return;
}
const pending = pendingRequests.get(id);
if (!pending) {
logger.debug("No pending request for message ID", { id });
return;
}
pendingRequests.delete(id);
if (error) {
pending.reject(new Error(error));
} else {
pending.resolve(result);
}
}
/**
* Handle disconnect and attempt reconnection
*/
function handleDisconnect(): void {
// Clear all pending requests
for (const [id, pending] of pendingRequests) {
pending.reject(new Error("Disconnected from relay server"));
pendingRequests.delete(id);
}
// Attempt reconnection
if (reconnectAttempts < MAX_RECONNECT_ATTEMPTS) {
reconnectAttempts++;
logger.info(
`Attempting reconnection (${reconnectAttempts}/${MAX_RECONNECT_ATTEMPTS})...`,
);
setTimeout(() => {
connectToRelay().catch((err) => {
logger.error("Reconnection failed", err.message);
});
}, RECONNECT_DELAY_MS);
} else {
logger.error("Max reconnection attempts reached");
}
}
/**
* Send a command to Figma and wait for response
*/
export async function sendCommandToFigma<T = unknown>(
command: string,
params: unknown,
timeoutMs: number = DEFAULT_TIMEOUT_MS,
): Promise<T> {
return new Promise((resolve, reject) => {
if (!ws || !isConnected || !channel) {
reject(new Error("Not connected to relay server"));
return;
}
const id = uuidv4();
const timeout = setTimeout(() => {
pendingRequests.delete(id);
reject(new Error(`Command timeout: ${command}`));
}, timeoutMs);
pendingRequests.set(id, {
resolve: (result: unknown) => {
clearTimeout(timeout);
resolve(result as T);
},
reject: (err: unknown) => {
clearTimeout(timeout);
reject(err);
},
timeout,
lastActivity: Date.now(),
});
const message = {
id,
type: "message",
channel,
message: {
id,
command,
params,
},
};
logger.debug("Sending command to Figma", { command, id });
ws.send(JSON.stringify(message));
});
}
/**
* Check if connected to relay
*/
export function isRelayConnected(): boolean {
return isConnected && ws !== null;
}
/**
* Get the current channel name
*/
export function getChannel(): string | null {
return channel;
}
/**
* Disconnect from relay server
*/
export function disconnectFromRelay(): void {
if (ws) {
ws.close();
ws = null;
isConnected = false;
channel = null;
}
}
