claude-web-chat/ws.go

package main

import (
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"time"

	"github.com/gorilla/websocket"
)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  4096,
	WriteBufferSize: 4096,
	CheckOrigin:     func(r *http.Request) bool { return true },
}

// resizeMsg is sent from the browser when the terminal size changes.
type resizeMsg struct {
	Type string `json:"type"`
	Cols uint16 `json:"cols"`
	Rows uint16 `json:"rows"`
}

// TerminalHandler handles WebSocket connections for terminal sessions.
type TerminalHandler struct {
	ptyMgr *PTYSessionManager
}

func NewTerminalHandler(ptyMgr *PTYSessionManager) *TerminalHandler {
	return &TerminalHandler{ptyMgr: ptyMgr}
}

func (h *TerminalHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	project := r.URL.Query().Get("project")
	projectDir := r.URL.Query().Get("project_dir")
	if project == "" || projectDir == "" {
		http.Error(w, "missing project params", http.StatusBadRequest)
		return
	}
	sessionKey := project

	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Printf("WebSocket upgrade error: %v", err)
		return
	}
	defer conn.Close()

	sess, isNew, err := h.ptyMgr.GetOrCreate(sessionKey, projectDir)
	if err != nil {
		log.Printf("PTY session error: %v", err)
		conn.WriteMessage(websocket.TextMessage, []byte(fmt.Sprintf("\r\nGreška: %v\r\n", err)))
		return
	}

	subID := fmt.Sprintf("ws-%d", time.Now().UnixNano())
	outputCh := sess.Subscribe(subID)
	defer sess.Unsubscribe(subID)

	// Send buffered output for reconnect (replay)
	if !isNew {
		buffered := sess.GetBuffer()
		if len(buffered) > 0 {
			conn.WriteMessage(websocket.BinaryMessage, buffered)
		}
	}

	// Serialized write channel to prevent concurrent WebSocket writes
	writeCh := make(chan []byte, 256)
	writeDone := make(chan struct{})
	go func() {
		defer close(writeDone)
		for data := range writeCh {
			if err := conn.WriteMessage(websocket.BinaryMessage, data); err != nil {
				return
			}
		}
	}()

	// PTY output → WebSocket (via write channel)
	go func() {
		for data := range outputCh {
			select {
			case writeCh <- data:
			default:
				// Drop if write channel is full
			}
		}
	}()

	// Watch for process exit
	go func() {
		<-sess.Done()
		// Send exit message and close
		select {
		case writeCh <- []byte("\r\n\033[33m[Sesija završena]\033[0m\r\n"):
		default:
		}
	}()

	// WebSocket → PTY (read pump)
	for {
		_, msg, err := conn.ReadMessage()
		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseNormalClosure) {
				log.Printf("WebSocket read error: %v", err)
			}
			break
		}

		// Check for resize message
		var resize resizeMsg
		if json.Unmarshal(msg, &resize) == nil && resize.Type == "resize" && resize.Cols > 0 && resize.Rows > 0 {
			if err := sess.Resize(resize.Rows, resize.Cols); err != nil {
				log.Printf("PTY resize error: %v", err)
			}
			continue
		}

		// Regular keyboard input → PTY
		if _, err := sess.WriteInput(msg); err != nil {
			log.Printf("PTY write error: %v", err)
			break
		}
	}

	close(writeCh)
	<-writeDone
}