package ws import ( "encoding/json" "log/slog" "net/http" "sync" "time" "github.com/gorilla/websocket" "github.com/sandrodach/mindsift/internal/domain" ) // CursorPresence represents an active user's ephemeral cursor position in a document room. type CursorPresence struct { ClientID string `json:"clientId"` UserID string `json:"userId"` UserName string `json:"userName"` UserColor string `json:"userColor"` BlockID string `json:"blockId,omitempty"` SelectionStart *int `json:"selectionStart,omitempty"` SelectionEnd *int `json:"selectionEnd,omitempty"` UpdatedAt int64 `json:"updatedAt"` } // DocWSClient represents an active WebSocket connection to a specific document room. type DocWSClient struct { hub *Hub room *DocRoom conn *websocket.Conn userID string clientID string canWrite bool send chan []byte } // DocSyncEvent is the lightweight JSON payload exchanged over /api/v1/ws/docs/{id}. type DocSyncEvent struct { Type string `json:"type"` DocID string `json:"docId,omitempty"` ClientID string `json:"clientId,omitempty"` UserID string `json:"userId,omitempty"` UserName string `json:"userName,omitempty"` UserColor string `json:"userColor,omitempty"` SelectionStart *int `json:"selectionStart,omitempty"` SelectionEnd *int `json:"selectionEnd,omitempty"` ActiveCursors []CursorPresence `json:"activeCursors,omitempty"` Title string `json:"title,omitempty"` Block *domain.DocumentBlock `json:"block,omitempty"` BlockID string `json:"blockId,omitempty"` AfterBlockID string `json:"afterBlockId,omitempty"` NewIndex *int `json:"newIndex,omitempty"` Content *string `json:"content,omitempty"` Level *int `json:"level,omitempty"` Checked *bool `json:"checked,omitempty"` Collapsed *bool `json:"collapsed,omitempty"` BlockType *string `json:"blockType,omitempty"` Blocks []domain.DocumentBlock `json:"blocks,omitempty"` } // DocRoom coordinates in-memory real-time block state and 1.5s debounced SQLite persistence for one document. type DocRoom struct { mu sync.Mutex docID string hub *Hub title string blocks []domain.DocumentBlock loaded bool wsClients map[*DocWSClient]bool activeCursors map[string]*CursorPresence debounceTimer *time.Timer dirty bool } // GetOrCreateDocRoom returns the active in-memory room for docID or initializes one from SQLite. func (h *Hub) GetOrCreateDocRoom(docID string) *DocRoom { h.mu.Lock() if h.docRooms == nil { h.docRooms = make(map[string]*DocRoom) } room, exists := h.docRooms[docID] if !exists { room = &DocRoom{ docID: docID, hub: h, wsClients: make(map[*DocWSClient]bool), activeCursors: make(map[string]*CursorPresence), } h.docRooms[docID] = room } store := h.docStore h.mu.Unlock() room.mu.Lock() if !room.loaded && store != nil { title, blocks, err := store.LoadDocumentState(docID) if err == nil { room.title = title room.blocks = normalizeRoomBlocks(docID, blocks) room.loaded = true } } room.mu.Unlock() return room } // GetActiveDocState returns the live in-memory title and blocks if an active DocRoom exists in RAM. func (h *Hub) GetActiveDocState(docID string) (string, []domain.DocumentBlock, bool) { h.mu.RLock() room, exists := h.docRooms[docID] h.mu.RUnlock() if !exists || room == nil { return "", nil, false } room.mu.Lock() defer room.mu.Unlock() if !room.loaded { return "", nil, false } blocksCopy := make([]domain.DocumentBlock, len(room.blocks)) copy(blocksCopy, room.blocks) return room.title, blocksCopy, true } // SyncActiveDocRoom updates an in-memory room if one is currently loaded (e.g. after REST PUT/PATCH). func (h *Hub) SyncActiveDocRoom(docID string, title string, blocks []domain.DocumentBlock) { h.mu.RLock() room, exists := h.docRooms[docID] h.mu.RUnlock() if !exists || room == nil { return } room.mu.Lock() if title != "" { room.title = title } if blocks != nil { room.blocks = normalizeRoomBlocks(docID, blocks) } room.loaded = true room.mu.Unlock() } // CloseDocRoom removes a document room from memory when a document is deleted. func (h *Hub) CloseDocRoom(docID string) { h.mu.Lock() room, exists := h.docRooms[docID] if exists { delete(h.docRooms, docID) } h.mu.Unlock() if exists && room != nil { room.mu.Lock() if room.debounceTimer != nil { room.debounceTimer.Stop() } room.dirty = false for c := range room.wsClients { _ = c.conn.Close() } room.wsClients = make(map[*DocWSClient]bool) room.mu.Unlock() } } func normalizeRoomBlocks(docID string, blocks []domain.DocumentBlock) []domain.DocumentBlock { if len(blocks) == 0 { return []domain.DocumentBlock{ { ID: docID + "-b0", DocumentID: docID, Type: "text", Content: "", Level: 0, Position: 1000, }, } } out := make([]domain.DocumentBlock, len(blocks)) for i, b := range blocks { if b.ID == "" { b.ID = domain.NewID("blk") } b.DocumentID = docID if b.Type == "" { b.Type = "text" } if b.Level < 0 { b.Level = 0 } b.Position = float64((i + 1) * 1000) out[i] = b } return out } // GetStateSnapshot returns a copy of the room's current live in-memory title and blocks. func (r *DocRoom) GetStateSnapshot() (string, []domain.DocumentBlock) { r.mu.Lock() defer r.mu.Unlock() blocksCopy := make([]domain.DocumentBlock, len(r.blocks)) copy(blocksCopy, r.blocks) return r.title, blocksCopy } // HandleEvent applies a granular block or title mutation to the in-memory room, // broadcasts it to all OTHER clients in the room (excluding senderClientID for echo protection), // and resets the 1.5-second inactivity debounce timer for SQLite persistence. func (r *DocRoom) HandleEvent(evt DocSyncEvent, sender *DocWSClient) { r.mu.Lock() evt.DocID = r.docID if sender != nil { if evt.ClientID == "" { evt.ClientID = sender.clientID } evt.UserID = sender.userID } mutated := false switch evt.Type { case "cursor_moved": if evt.ClientID != "" { cp := &CursorPresence{ ClientID: evt.ClientID, UserID: evt.UserID, UserName: evt.UserName, UserColor: evt.UserColor, BlockID: evt.BlockID, SelectionStart: evt.SelectionStart, SelectionEnd: evt.SelectionEnd, UpdatedAt: time.Now().UnixMilli(), } if r.activeCursors == nil { r.activeCursors = make(map[string]*CursorPresence) } r.activeCursors[evt.ClientID] = cp msgBytes, err := json.Marshal(evt) if err == nil { for client := range r.wsClients { if sender != nil && (client == sender || (sender.clientID != "" && client.clientID == sender.clientID)) { continue } if evt.ClientID != "" && client.clientID == evt.ClientID { continue } select { case client.send <- msgBytes: default: } } } } r.mu.Unlock() return case "title_updated": if evt.Title != "" || r.title != evt.Title { r.title = evt.Title mutated = true } case "block_created": if evt.Block != nil { newBlock := *evt.Block if newBlock.ID == "" { newBlock.ID = domain.NewID("blk") } newBlock.DocumentID = r.docID if newBlock.Type == "" { newBlock.Type = "text" } if newBlock.Level < 0 { newBlock.Level = 0 } // Avoid duplicate ID if already created existsIdx := -1 for i, b := range r.blocks { if b.ID == newBlock.ID { existsIdx = i break } } if existsIdx >= 0 { r.blocks[existsIdx] = newBlock } else { insertIdx := len(r.blocks) if evt.AfterBlockID != "" { for i, b := range r.blocks { if b.ID == evt.AfterBlockID { insertIdx = i + 1 break } } } else if evt.NewIndex != nil && *evt.NewIndex >= 0 && *evt.NewIndex <= len(r.blocks) { insertIdx = *evt.NewIndex } updated := make([]domain.DocumentBlock, 0, len(r.blocks)+1) updated = append(updated, r.blocks[:insertIdx]...) updated = append(updated, newBlock) updated = append(updated, r.blocks[insertIdx:]...) r.blocks = normalizeRoomBlocks(r.docID, updated) } evt.Block = &newBlock mutated = true } case "block_updated": targetID := evt.BlockID if targetID == "" && evt.Block != nil { targetID = evt.Block.ID } if targetID != "" { for i := range r.blocks { if r.blocks[i].ID == targetID { if evt.Block != nil { r.blocks[i].Content = evt.Block.Content r.blocks[i].Level = evt.Block.Level if evt.Block.Type != "" { r.blocks[i].Type = evt.Block.Type } r.blocks[i].Checked = evt.Block.Checked r.blocks[i].Collapsed = evt.Block.Collapsed r.blocks[i].Lang = evt.Block.Lang } else { if evt.Content != nil { r.blocks[i].Content = *evt.Content } if evt.Level != nil { lvl := *evt.Level if lvl < 0 { lvl = 0 } r.blocks[i].Level = lvl } if evt.Checked != nil { r.blocks[i].Checked = *evt.Checked } if evt.Collapsed != nil { r.blocks[i].Collapsed = *evt.Collapsed } if evt.BlockType != nil && *evt.BlockType != "" { r.blocks[i].Type = *evt.BlockType } } bCopy := r.blocks[i] evt.Block = &bCopy evt.BlockID = targetID mutated = true break } } } case "block_deleted": if evt.BlockID != "" { filtered := make([]domain.DocumentBlock, 0, len(r.blocks)) for _, b := range r.blocks { if b.ID != evt.BlockID { filtered = append(filtered, b) } } r.blocks = normalizeRoomBlocks(r.docID, filtered) mutated = true } case "block_moved": if evt.BlockID != "" && evt.NewIndex != nil { fromIdx := -1 for i, b := range r.blocks { if b.ID == evt.BlockID { fromIdx = i break } } if fromIdx >= 0 { moving := r.blocks[fromIdx] if evt.Level != nil { lvl := *evt.Level if lvl < 0 { lvl = 0 } moving.Level = lvl } without := make([]domain.DocumentBlock, 0, len(r.blocks)-1) without = append(without, r.blocks[:fromIdx]...) without = append(without, r.blocks[fromIdx+1:]...) toIdx := *evt.NewIndex if toIdx < 0 { toIdx = 0 } if toIdx > len(without) { toIdx = len(without) } reordered := make([]domain.DocumentBlock, 0, len(r.blocks)) reordered = append(reordered, without[:toIdx]...) reordered = append(reordered, moving) reordered = append(reordered, without[toIdx:]...) r.blocks = normalizeRoomBlocks(r.docID, reordered) mutated = true } } case "blocks_replaced": if evt.Blocks != nil { r.blocks = normalizeRoomBlocks(r.docID, evt.Blocks) evt.Blocks = r.blocks mutated = true } } if !mutated { r.mu.Unlock() return } r.dirty = true r.loaded = true msgBytes, err := json.Marshal(evt) if err == nil { for client := range r.wsClients { // Strict Echo Protection: never send back to the originating connection or clientId if sender != nil && (client == sender || (sender.clientID != "" && client.clientID == sender.clientID)) { continue } if evt.ClientID != "" && client.clientID == evt.ClientID { continue } select { case client.send <- msgBytes: default: } } } r.scheduleDebouncedFlushLocked() r.mu.Unlock() } func (r *DocRoom) scheduleDebouncedFlushLocked() { if r.debounceTimer != nil { r.debounceTimer.Stop() } r.debounceTimer = time.AfterFunc(1500*time.Millisecond, func() { r.Flush() }) } // Flush persists the current in-memory document state (blocks, #tags in block_tags, and markdown_cache) to SQLite. func (r *DocRoom) Flush() { r.mu.Lock() if !r.dirty { r.mu.Unlock() return } if r.debounceTimer != nil { r.debounceTimer.Stop() r.debounceTimer = nil } docID := r.docID title := r.title blocksCopy := make([]domain.DocumentBlock, len(r.blocks)) copy(blocksCopy, r.blocks) r.dirty = false r.mu.Unlock() tagsCount := 0 for _, b := range blocksCopy { tagsCount += len(domain.ExtractTags(b.Content)) } start := time.Now() r.hub.mu.RLock() store := r.hub.docStore r.hub.mu.RUnlock() if store != nil { if err := store.FlushDocumentState(docID, title, blocksCopy); err != nil { slog.Error("SQLite document flush failed", "doc_id", docID, "error", err, ) } else { durationMs := float64(time.Since(start).Microseconds()) / 1000.0 slog.Debug("SQLite document flushed", "doc_id", docID, "blocks_count", len(blocksCopy), "tags_count", tagsCount, "duration_ms", durationMs, ) } } } // ServeDocWS handles WebSocket connections for GET /api/v1/ws/docs/{id} func (r *DocRoom) ServeDocWS(w http.ResponseWriter, req *http.Request, userID string, clientID string, canWrite bool) { conn, err := upgrader.Upgrade(w, req, nil) if err != nil { slog.Error("WebSocket upgrade error", "doc_id", r.docID, "error", err) return } if clientID == "" { clientID = req.URL.Query().Get("clientId") } if clientID == "" { clientID = domain.NewID("cli") } userName := req.URL.Query().Get("userName") userColor := req.URL.Query().Get("userColor") client := &DocWSClient{ hub: r.hub, room: r, conn: conn, userID: userID, clientID: clientID, canWrite: canWrite, send: make(chan []byte, 256), } r.mu.Lock() r.wsClients[client] = true activeClients := len(r.wsClients) if r.activeCursors == nil { r.activeCursors = make(map[string]*CursorPresence) } if userName != "" { r.activeCursors[clientID] = &CursorPresence{ ClientID: clientID, UserID: userID, UserName: userName, UserColor: userColor, UpdatedAt: time.Now().UnixMilli(), } } title := r.title blocksCopy := make([]domain.DocumentBlock, len(r.blocks)) copy(blocksCopy, r.blocks) activeList := make([]CursorPresence, 0, len(r.activeCursors)) for _, cp := range r.activeCursors { if cp != nil { activeList = append(activeList, *cp) } } r.mu.Unlock() slog.Debug("WS client joined doc room", "doc_id", r.docID, "client_id", clientID, "user_id", userID, "active_clients", activeClients, ) // Immediately send live in-memory state and current cursors to the newly connected client initEvt := DocSyncEvent{ Type: "doc_state", DocID: r.docID, ClientID: clientID, Title: title, Blocks: blocksCopy, ActiveCursors: activeList, } if initBytes, err := json.Marshal(initEvt); err == nil { client.send <- initBytes } // If userName was provided, broadcast initial presence to existing clients in room if userName != "" { joinEvt := DocSyncEvent{ Type: "cursor_moved", DocID: r.docID, ClientID: clientID, UserID: userID, UserName: userName, UserColor: userColor, } if joinBytes, err := json.Marshal(joinEvt); err == nil { r.mu.Lock() for c := range r.wsClients { if c != client && c.clientID != clientID { select { case c.send <- joinBytes: default: } } } r.mu.Unlock() } } go client.writePump() go client.readPump() } func (c *DocWSClient) readPump() { defer func() { c.room.mu.Lock() delete(c.room.wsClients, c) delete(c.room.activeCursors, c.clientID) remaining := len(c.room.wsClients) // Broadcast client_left to all remaining clients in room leftEvt := DocSyncEvent{ Type: "client_left", DocID: c.room.docID, ClientID: c.clientID, UserID: c.userID, } if leftBytes, err := json.Marshal(leftEvt); err == nil { for client := range c.room.wsClients { if client != c && client.clientID != c.clientID { select { case client.send <- leftBytes: default: } } } } c.room.mu.Unlock() slog.Debug("WS client left doc room", "doc_id", c.room.docID, "client_id", c.clientID, "user_id", c.userID, "active_clients", remaining, ) _ = c.conn.Close() // Flush immediately when the last client leaves the document room if remaining == 0 { c.room.Flush() } }() c.conn.SetReadLimit(4 * 1024 * 1024) _ = c.conn.SetReadDeadline(time.Now().Add(60 * time.Second)) c.conn.SetPongHandler(func(string) error { _ = c.conn.SetReadDeadline(time.Now().Add(60 * time.Second)) return nil }) for { _, message, err := c.conn.ReadMessage() if err != nil { break } var evt DocSyncEvent if err := json.Unmarshal(message, &evt); err == nil && evt.Type != "" { targetBlockID := evt.BlockID if targetBlockID == "" && evt.Block != nil { targetBlockID = evt.Block.ID } slog.Debug("WS doc event received", "type", evt.Type, "doc_id", c.room.docID, "block_id", targetBlockID, "sender_client_id", c.clientID, ) // Read-only clients can still move their cursor and send presence updates if !c.canWrite && evt.Type != "cursor_moved" { continue } c.room.HandleEvent(evt, c) } } } func (c *DocWSClient) writePump() { ticker := time.NewTicker(25 * time.Second) defer func() { ticker.Stop() _ = c.conn.Close() }() for { select { case message, ok := <-c.send: _ = c.conn.SetWriteDeadline(time.Now().Add(10 * time.Second)) if !ok { _ = c.conn.WriteMessage(websocket.CloseMessage, []byte{}) return } if err := c.conn.WriteMessage(websocket.TextMessage, message); err != nil { return } case <-ticker.C: _ = c.conn.SetWriteDeadline(time.Now().Add(10 * time.Second)) if err := c.conn.WriteMessage(websocket.PingMessage, nil); err != nil { return } } } }