@juspay/neurolink
Version:
Universal AI Development Platform with working MCP integration, multi-provider support, voice (TTS/STT/realtime), and professional CLI. 58+ external MCP servers discoverable, multimodal file processing, RAG pipelines. Build, test, and deploy AI applicatio
450 lines (449 loc) • 16.3 kB
JavaScript
/**
* WebSocket Client for NeuroLink SDK
*
* Provides a dedicated WebSocket client for real-time streaming connections
* to NeuroLink servers. Supports bidirectional communication, automatic
* reconnection, and message queuing.
*
* @module @neurolink/client/wsClient
*/
import { SpanKind, SpanStatusCode } from "@opentelemetry/api";
import { tracers } from "../telemetry/tracers.js";
// =============================================================================
// WebSocket Client
// =============================================================================
/**
* WebSocket streaming client for NeuroLink
*
* Provides real-time bidirectional communication with NeuroLink servers.
*
* @example Basic usage
* ```typescript
* const wsClient = new NeuroLinkWebSocket({
* baseUrl: 'wss://api.neurolink.example.com/ws',
* apiKey: 'your-api-key',
* });
*
* wsClient.connect({
* onMessage: (event) => console.log('Received:', event),
* onError: (error) => console.error('Error:', error),
* });
*
* // Send a message
* wsClient.send({
* type: 'message',
* channel: 'chat',
* payload: { prompt: 'Hello!' },
* });
* ```
*/
export class NeuroLinkWebSocket {
ws = null;
config;
state = "disconnected";
reconnectAttempts = 0;
heartbeatTimer = null;
messageQueue = [];
eventHandlers = {};
subscriptions = new Map();
pendingAuth = false;
/**
* Active OTel span for the current WebSocket connection lifecycle.
* Created at connect() and ended on close/error so we capture connection
* lifetime, reconnect counts, and error attribution in Langfuse.
*/
connectionSpan = null;
/**
* Local flag to suppress reconnection during an explicit disconnect().
* Unlike mutating config.autoReconnect, this preserves the user's
* original configuration so that a subsequent connect() still honours it.
*/
disconnectRequested = false;
constructor(config) {
this.config = {
baseUrl: config.baseUrl,
apiKey: config.apiKey ?? "",
token: config.token ?? "",
timeout: config.timeout ?? 30000,
headers: config.headers ?? {},
autoReconnect: config.autoReconnect ?? true,
maxReconnectAttempts: config.maxReconnectAttempts ?? 10,
reconnectDelay: config.reconnectDelay ?? 1000,
maxReconnectDelay: config.maxReconnectDelay ?? 30000,
heartbeatInterval: config.heartbeatInterval ?? 30000,
queueSize: config.queueSize ?? 100,
};
}
/**
* Get current connection state
*/
getState() {
return this.state;
}
/**
* Check if connected
*/
isConnected() {
return this.state === "connected" && this.ws?.readyState === WebSocket.OPEN;
}
/**
* Connect to WebSocket server
*/
connect(handlers) {
if (this.ws && this.ws.readyState === WebSocket.OPEN) {
return;
}
// Reset the disconnect flag so reconnection logic works again
this.disconnectRequested = false;
this.eventHandlers = handlers ?? {};
this.setState("connecting");
// End any orphaned span from a prior connect() attempt (e.g., re-entrant call
// while a previous attempt was still connecting).
if (this.connectionSpan) {
this.connectionSpan.setAttribute("ws.superseded", true);
this.connectionSpan.setStatus({
code: SpanStatusCode.ERROR,
message: "Connection attempt superseded by new connect() call",
});
this.connectionSpan.end();
this.connectionSpan = null;
}
// Start an OTel span that tracks the lifetime of this connection attempt.
// Ended in onclose/onerror/disconnect so metrics capture connection
// duration and error attribution.
this.connectionSpan = tracers.http.startSpan("neurolink.client.ws.connect", {
kind: SpanKind.CLIENT,
attributes: {
"http.url": this.config.baseUrl,
"ws.auto_reconnect": this.config.autoReconnect,
"ws.reconnect_attempt": this.reconnectAttempts,
},
});
// Build WebSocket URL (credentials are sent via headers, not query params,
// to avoid leaking secrets in server logs, browser history, and HTTP referers)
const url = new URL(this.config.baseUrl);
// Build auth headers matching httpClient.ts conventions
const authHeaders = {
...this.config.headers,
};
if (this.config.apiKey) {
authHeaders["X-API-Key"] = this.config.apiKey;
}
if (this.config.token) {
authHeaders["Authorization"] = `Bearer ${this.config.token}`;
}
// In Node.js (ws package), pass headers via the options parameter.
// In browsers, the WebSocket API does not support custom headers, so
// credentials are sent as the first message after the connection opens.
const isNode = typeof globalThis.process !== "undefined" &&
typeof globalThis.process.versions?.node === "string";
try {
if (isNode && Object.keys(authHeaders).length > 0) {
this.ws = new WebSocket(url.toString(), { headers: authHeaders });
}
else {
this.ws = new WebSocket(url.toString());
}
}
catch (error) {
if (this.connectionSpan) {
this.connectionSpan.recordException(error instanceof Error ? error : new Error(String(error)));
this.connectionSpan.setStatus({
code: SpanStatusCode.ERROR,
message: error instanceof Error ? error.message : String(error),
});
this.connectionSpan.end();
this.connectionSpan = null;
}
throw error;
}
this.pendingAuth = !isNode && (!!this.config.apiKey || !!this.config.token);
this.setupEventListeners();
}
/**
* Disconnect from WebSocket server
*
* Sets a local flag to prevent the onclose handler from triggering
* reconnection, without mutating the shared config.
*/
disconnect() {
this.stopHeartbeat();
this.disconnectRequested = true;
if (this.ws) {
this.ws.close(1000, "Client disconnect");
this.ws = null;
}
this.setState("disconnected");
// Let onclose finalize the span — it fires from ws.close(1000,...) and
// has the close code context. We only end here if ws is already null
// (e.g. connect was never called) to avoid leaking.
if (this.connectionSpan && !this.ws) {
this.connectionSpan.setAttribute("ws.close_reason", "client_disconnect");
this.connectionSpan.setStatus({ code: SpanStatusCode.OK });
this.connectionSpan.end();
this.connectionSpan = null;
}
}
/**
* Send a message through WebSocket
*/
send(message) {
if (this.isConnected() && this.ws) {
this.ws.send(JSON.stringify(message));
}
else {
// Queue message for when connected
if (this.messageQueue.length < (this.config.queueSize ?? 100)) {
this.messageQueue.push(message);
}
}
}
/**
* Subscribe to a channel with streaming callbacks
*/
subscribe(channel, callbacks) {
this.subscriptions.set(channel, callbacks);
this.send({
type: "subscribe",
channel,
});
}
/**
* Unsubscribe from a channel
*/
unsubscribe(channel) {
this.subscriptions.delete(channel);
this.send({
type: "unsubscribe",
channel,
});
}
/**
* Stream a prompt with callbacks
*/
stream(prompt, options) {
const channel = options?.channel ?? `stream_${Date.now()}`;
if (options) {
this.subscribe(channel, {
onText: options.onText,
onToolCall: options.onToolCall,
onToolResult: options.onToolResult,
onDone: options.onDone,
onError: options.onError,
});
}
this.send({
type: "message",
channel,
payload: { prompt },
});
}
// =============================================================================
// Private Methods
// =============================================================================
setupEventListeners() {
if (!this.ws) {
return;
}
this.ws.onopen = () => {
this.setState("connected");
this.reconnectAttempts = 0;
if (this.connectionSpan) {
this.connectionSpan.setAttribute("ws.connected", true);
}
// In browser environments, send credentials as the first message
// since the browser WebSocket API does not support custom headers.
if (this.pendingAuth && this.ws) {
const authPayload = { type: "auth" };
if (this.config.apiKey) {
authPayload["apiKey"] = this.config.apiKey;
}
if (this.config.token) {
authPayload["token"] = this.config.token;
}
this.ws.send(JSON.stringify(authPayload));
this.pendingAuth = false;
}
this.eventHandlers.onOpen?.();
this.startHeartbeat();
this.flushMessageQueue();
// Re-subscribe to all active channels after a reconnect so that
// channel listeners are restored transparently.
this.replaySubscriptions();
};
this.ws.onclose = (event) => {
this.setState("disconnected");
this.stopHeartbeat();
this.eventHandlers.onClose?.(event.code, event.reason);
if (this.connectionSpan) {
this.connectionSpan.setAttribute("ws.close_code", event.code);
if (event.reason) {
this.connectionSpan.setAttribute("ws.close_reason", event.reason);
}
// 1000 = normal closure; other codes are abnormal.
if (event.code === 1000) {
this.connectionSpan.setStatus({ code: SpanStatusCode.OK });
}
else {
this.connectionSpan.setStatus({
code: SpanStatusCode.ERROR,
message: `WebSocket closed with code ${event.code}${event.reason ? `: ${event.reason}` : ""}`,
});
}
this.connectionSpan.end();
this.connectionSpan = null;
}
// Only attempt reconnection when auto-reconnect is enabled AND this
// was not an intentional disconnect (code 1000 or explicit call).
if (this.config.autoReconnect &&
!this.disconnectRequested &&
event.code !== 1000) {
this.attemptReconnect();
}
};
this.ws.onerror = () => {
this.setState("error");
const error = new Error("WebSocket connection error");
this.eventHandlers.onError?.(error);
if (this.connectionSpan) {
this.connectionSpan.recordException(error);
this.connectionSpan.setStatus({
code: SpanStatusCode.ERROR,
message: error.message,
});
// Do not end here — onclose will fire next and end the span with
// the precise close code. Keeping the span open until close gives
// us the full connection lifetime on Langfuse.
}
};
this.ws.onmessage = (event) => {
try {
const data = JSON.parse(event.data);
this.handleMessage(data);
}
catch {
// Handle non-JSON messages
this.eventHandlers.onMessage?.({
type: "text",
content: event.data,
timestamp: Date.now(),
});
}
};
}
handleMessage(data) {
// Notify global handler
this.eventHandlers.onMessage?.(data);
// Notify channel-specific subscribers
if (data.channel) {
const callbacks = this.subscriptions.get(data.channel);
if (callbacks) {
this.dispatchToCallbacks(data, callbacks);
}
}
}
dispatchToCallbacks(event, callbacks) {
switch (event.type) {
case "text":
callbacks.onText?.(event.content ?? "");
break;
case "tool-call":
if (event.toolCall) {
callbacks.onToolCall?.(event.toolCall);
}
break;
case "tool-result":
if (event.toolResult) {
callbacks.onToolResult?.(event.toolResult);
}
break;
case "done": {
// Create a StreamResult from the event
const result = {
content: event.content ?? "",
finishReason: "stop",
};
callbacks.onDone?.(result);
break;
}
case "error":
if (event.error) {
callbacks.onError?.(event.error);
}
break;
}
}
setState(state) {
this.state = state;
this.eventHandlers.onStateChange?.(state);
}
startHeartbeat() {
this.stopHeartbeat();
this.heartbeatTimer = setInterval(() => {
if (this.isConnected()) {
this.send({ type: "ping" });
}
}, this.config.heartbeatInterval ?? 30000);
}
stopHeartbeat() {
if (this.heartbeatTimer) {
clearInterval(this.heartbeatTimer);
this.heartbeatTimer = null;
}
}
attemptReconnect() {
if (this.reconnectAttempts >= this.config.maxReconnectAttempts) {
this.eventHandlers.onError?.(new Error(`Max reconnection attempts (${this.config.maxReconnectAttempts}) exceeded`));
return;
}
this.reconnectAttempts++;
const delay = Math.min(this.config.reconnectDelay * Math.pow(2, this.reconnectAttempts - 1), this.config.maxReconnectDelay);
this.eventHandlers.onReconnect?.(this.reconnectAttempts);
setTimeout(() => {
this.connect(this.eventHandlers);
}, delay);
}
flushMessageQueue() {
while (this.messageQueue.length > 0 && this.isConnected()) {
const message = this.messageQueue.shift();
if (message) {
this.send(message);
}
}
}
/**
* Re-send subscribe messages for every active subscription.
* Called after a successful reconnect so that channel listeners
* resume working without the caller needing to re-subscribe manually.
*/
replaySubscriptions() {
for (const channel of this.subscriptions.keys()) {
this.send({
type: "subscribe",
channel,
});
}
}
}
// =============================================================================
// Factory Function
// =============================================================================
/**
* Create a WebSocket client instance
*
* @example
* ```typescript
* const client = createWebSocketClient({
* baseUrl: 'wss://api.neurolink.example.com/ws',
* apiKey: 'your-api-key',
* autoReconnect: true,
* });
*
* client.connect({
* onMessage: (event) => console.log('Received:', event),
* });
* ```
*/
export function createWebSocketClient(config) {
return new NeuroLinkWebSocket(config);
}