// mcp/src/client/SocketClient.ts
import { io, Socket } from 'socket.io-client';
import {
EventType,
NoteTreeManager,
type BaseEvent,
type ImageInfo,
type PDFInfo,
type RoomSnapshot,
type TreeNode,
} from '@whkerdb/shared';
interface SocketClientOptions {
serverUrl: string;
userId: string;
username: string;
clientId?: string;
}
interface JoinRoomResult {
snapshot: RoomSnapshot;
userId: string;
}
/**
* WhkerDB WebSocket 客户端
*/
export class SocketClient {
private socket: Socket | null = null;
private serverUrl: string;
private userId: string;
private username: string;
private clientId: string;
private currentRoomId: string | null = null;
private currentSnapshot: RoomSnapshot | null = null;
private noteTree: NoteTreeManager | null = null;
private pdfs: Map<string, PDFInfo> = new Map();
private images: Map<string, ImageInfo> = new Map();
private eventHandlers: Map<string, (event: BaseEvent) => void> = new Map();
constructor(options: SocketClientOptions) {
this.serverUrl = options.serverUrl;
this.userId = options.userId;
this.username = options.username;
this.clientId = options.clientId || `mcp-${Date.now()}`;
}
private setSnapshot(snapshot: RoomSnapshot): void {
this.currentRoomId = snapshot.id;
this.currentSnapshot = snapshot;
this.noteTree = new NoteTreeManager();
this.noteTree.import(snapshot.noteTree);
this.pdfs = new Map(snapshot.pdfs);
this.images = new Map(snapshot.images);
this.rebuildSnapshot();
}
private rebuildSnapshot(): void {
if (!this.currentSnapshot) return;
if (this.noteTree) {
this.currentSnapshot.noteTree = this.noteTree.export();
}
this.currentSnapshot.pdfs = Array.from(this.pdfs.entries());
this.currentSnapshot.images = Array.from(this.images.entries());
}
private applyServerEvent(event: BaseEvent): void {
if (!this.currentSnapshot) return;
// Keep room version monotonic.
if (event.version > this.currentSnapshot.version) {
this.currentSnapshot.version = event.version;
}
if (event.timestamp > this.currentSnapshot.updatedAt) {
this.currentSnapshot.updatedAt = event.timestamp;
}
if (!this.noteTree) {
// Shouldn't happen after join-room, but stay safe.
this.noteTree = new NoteTreeManager();
this.noteTree.import(this.currentSnapshot.noteTree);
}
switch (event.type) {
case EventType.NODE_ADDED: {
const { node, parentId, position } = (event as any).payload as {
node: TreeNode;
parentId: string | null;
position: number;
};
this.noteTree.addNodeDirect(node, parentId ?? null, position ?? 0);
// Link PDF metadata to its file node when possible.
if (node.type === 'file' && node.pdfId) {
const current = this.pdfs.get(node.pdfId);
if (current) {
this.pdfs.set(node.pdfId, { ...current, fileNodeId: node.id });
}
}
break;
}
case EventType.NODE_DELETED: {
const { nodeId } = (event as any).payload as { nodeId: string };
const node = this.noteTree.getNode(nodeId);
this.noteTree.deleteNode(nodeId, true);
if (node?.type === 'file' && node.pdfId) {
this.pdfs.delete(node.pdfId);
}
if (
this.currentSnapshot.viewState?.currentNodeId &&
!this.noteTree.getNode(this.currentSnapshot.viewState.currentNodeId)
) {
this.currentSnapshot.viewState = { ...this.currentSnapshot.viewState, currentNodeId: '' };
}
break;
}
case EventType.NODE_MOVED: {
const { nodeId, newParentId, newPosition } = (event as any).payload as {
nodeId: string;
newParentId: string | null;
newPosition: number;
};
this.noteTree.moveNode(nodeId, newParentId ?? null, newPosition ?? 0);
break;
}
case EventType.NODE_UPDATED: {
const { nodeId, changes } = (event as any).payload as {
nodeId: string;
changes: Partial<TreeNode>;
};
const before = this.noteTree.getNode(nodeId);
this.noteTree.updateNode(nodeId, changes);
const after = this.noteTree.getNode(nodeId);
if (before?.type === 'file') {
const prevPdfId = before.pdfId;
const nextPdfId = after?.type === 'file' ? after.pdfId : undefined;
if (prevPdfId && prevPdfId !== nextPdfId) {
const prevInfo = this.pdfs.get(prevPdfId);
if (prevInfo && prevInfo.fileNodeId === nodeId) {
this.pdfs.set(prevPdfId, { ...prevInfo, fileNodeId: '' });
}
}
if (nextPdfId) {
const info = this.pdfs.get(nextPdfId);
if (info && (!info.fileNodeId || info.fileNodeId === nodeId)) {
this.pdfs.set(nextPdfId, { ...info, fileNodeId: nodeId });
}
}
}
break;
}
case EventType.CHILDREN_REORDERED: {
const { parentId, childrenIds } = (event as any).payload as {
parentId: string | null;
childrenIds: string[];
};
this.noteTree.reorderChildren(parentId ?? null, childrenIds);
break;
}
case EventType.OBJECT_ADDED: {
const { nodeId, object } = (event as any).payload as { nodeId: string; object: any };
this.noteTree.addObject(nodeId, object);
break;
}
case EventType.OBJECT_UPDATED: {
const { nodeId, objectId, changes } = (event as any).payload as {
nodeId: string;
objectId: string;
changes: any;
};
this.noteTree.updateObject(nodeId, objectId, changes);
break;
}
case EventType.OBJECT_DELETED: {
const { nodeId, objectId } = (event as any).payload as { nodeId: string; objectId: string };
this.noteTree.deleteObject(nodeId, objectId);
break;
}
case EventType.VIEW_STATE_CHANGED: {
const { viewState } = (event as any).payload as { viewState: RoomSnapshot['viewState'] };
this.currentSnapshot.viewState = viewState;
break;
}
case EventType.PDF_UPLOADED: {
const { pdfId, filename, totalPages, fileSize } = (event as any).payload as {
pdfId: string;
filename: string;
totalPages: number;
fileSize: number;
};
const existing = this.pdfs.get(pdfId);
this.pdfs.set(pdfId, {
id: pdfId,
filename,
totalPages,
fileSize,
uploadedAt: event.timestamp,
fileNodeId: existing?.fileNodeId || '',
});
break;
}
case EventType.PDF_DELETED: {
const { pdfId } = (event as any).payload as { pdfId: string };
this.pdfs.delete(pdfId);
break;
}
case EventType.IMAGE_UPLOADED: {
const { imageId, filename, mimeType, width, height, fileSize } = (event as any).payload as {
imageId: string;
filename: string;
mimeType: string;
width: number;
height: number;
fileSize: number;
};
this.images.set(imageId, {
id: imageId,
filename,
mimeType,
width,
height,
fileSize,
uploadedAt: event.timestamp,
});
break;
}
case EventType.IMAGE_DELETED: {
const { imageId } = (event as any).payload as { imageId: string };
this.images.delete(imageId);
break;
}
}
this.rebuildSnapshot();
}
/**
* 连接到服务器
*/
connect(): Promise<void> {
return new Promise((resolve, reject) => {
if (this.socket?.connected) {
resolve();
return;
}
this.socket = io(this.serverUrl, {
transports: ['websocket', 'polling'],
reconnection: true,
reconnectionAttempts: 5,
reconnectionDelay: 1000,
});
const timeout = setTimeout(() => {
reject(new Error('Connection timeout'));
}, 10000);
this.socket.on('connect', () => {
clearTimeout(timeout);
this.setupEventListeners();
resolve();
});
this.socket.on('connect_error', (error) => {
clearTimeout(timeout);
reject(new Error(`Connection failed: ${error.message}`));
});
});
}
/**
* 设置事件监听器
*/
private setupEventListeners(): void {
if (!this.socket) return;
this.socket.on('room-snapshot', (data: JoinRoomResult) => {
this.setSnapshot(data.snapshot);
});
this.socket.on('event', (event: BaseEvent) => {
this.applyServerEvent(event);
// 通知事件处理器
for (const handler of this.eventHandlers.values()) {
handler(event);
}
});
this.socket.on('error', (error: { message: string }) => {
console.error('Socket error:', error.message);
});
}
/**
* 加入房间
*/
joinRoom(roomId: string, roomName?: string): Promise<JoinRoomResult> {
return new Promise((resolve, reject) => {
if (!this.socket?.connected) {
reject(new Error('Not connected to server'));
return;
}
const timeout = setTimeout(() => {
reject(new Error('Join room timeout'));
}, 10000);
this.socket.once('room-snapshot', (data: JoinRoomResult) => {
clearTimeout(timeout);
this.currentRoomId = roomId;
this.setSnapshot(data.snapshot);
resolve(data);
});
this.socket.once('error', (error: { message: string }) => {
clearTimeout(timeout);
reject(new Error(error.message));
});
this.socket.emit('join-room', roomId, this.userId, this.username, roomName, this.clientId);
});
}
/**
* 离开房间
*/
leaveRoom(): void {
if (this.socket?.connected && this.currentRoomId) {
this.socket.emit('leave-room');
this.currentRoomId = null;
this.currentSnapshot = null;
this.noteTree = null;
this.pdfs = new Map();
this.images = new Map();
}
}
/**
* 发送事件
*/
applyEvent(event: Partial<BaseEvent>): Promise<void> {
return new Promise((resolve, reject) => {
if (!this.socket?.connected) {
reject(new Error('Not connected to server'));
return;
}
if (!event.id || typeof event.id !== 'string') {
reject(new Error('Event id is required'));
return;
}
let done = false;
const cleanup = () => {
if (!this.socket) return;
this.socket.off('event', onEvent);
this.socket.off('event-rejected', onRejected);
};
const onRejected = (data: { eventId: string; reason: string }) => {
if (data.eventId !== event.id) return;
if (done) return;
done = true;
clearTimeout(timeout);
cleanup();
if (data.reason === 'DUPLICATE_EVENT_ID') {
resolve();
return;
}
reject(new Error(`Event rejected: ${data.reason}`));
};
const onEvent = (broadcastedEvent: BaseEvent) => {
if (broadcastedEvent.id !== event.id) return;
if (done) return;
done = true;
clearTimeout(timeout);
cleanup();
resolve();
};
const timeout = setTimeout(() => {
if (done) return;
done = true;
cleanup();
reject(new Error('Event confirmation timeout'));
}, 15000);
this.socket.on('event-rejected', onRejected);
this.socket.on('event', onEvent);
this.socket.emit('apply-event', event);
});
}
/**
* 获取当前房间快照
*/
getSnapshot(): RoomSnapshot | null {
return this.currentSnapshot;
}
/**
* 获取当前房间 ID
*/
getCurrentRoomId(): string | null {
return this.currentRoomId;
}
/**
* 获取当前版本号
*/
getCurrentVersion(): number {
return this.currentSnapshot?.version ?? 0;
}
/**
* 获取下一个版本号
*/
getNextVersion(): number {
return this.getCurrentVersion() + 1;
}
/**
* 添加事件处理器
*/
onEvent(id: string, handler: (event: BaseEvent) => void): void {
this.eventHandlers.set(id, handler);
}
/**
* 移除事件处理器
*/
offEvent(id: string): void {
this.eventHandlers.delete(id);
}
/**
* 断开连接
*/
disconnect(): void {
if (this.socket) {
this.leaveRoom();
this.socket.disconnect();
this.socket = null;
}
}
/**
* 检查是否已连接
*/
isConnected(): boolean {
return this.socket?.connected ?? false;
}
/**
* 获取用户 ID
*/
getUserId(): string {
return this.userId;
}
/**
* 获取客户端 ID
*/
getClientId(): string {
return this.clientId;
}
}
