kandev/apps/backend/internal/gateway/websocket/lsp_handler.go

package websocket

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"sync"
	"time"

	"github.com/gin-gonic/gin"
	gorillaws "github.com/gorilla/websocket"
	"github.com/kandev/kandev/internal/auth/authn"
	"go.uber.org/zap"

	"github.com/kandev/kandev/internal/agent/runtime/lifecycle"
	"github.com/kandev/kandev/internal/agentruntime"
	"github.com/kandev/kandev/internal/common/logger"
	"github.com/kandev/kandev/internal/events"
	"github.com/kandev/kandev/internal/events/bus"
	"github.com/kandev/kandev/internal/lsp/installer"
	"github.com/kandev/kandev/internal/lsp/protocol"
	"github.com/kandev/kandev/internal/user/models"
)

// Custom WebSocket close codes for LSP connections.
const (
	lspCloseBinaryNotFound       = 4001
	lspCloseSessionNotFound      = 4002
	lspCloseInstallFailed        = 4003
	lspCloseUnsupportedExecutor  = 4004
	lspCloseCapacityExceeded     = 4005
	lspCloseStreamError          = 4006
	lspCloseRuntimeStopped       = 4010
	lspCloseUnsupportedCloseText = "LSP is only supported for local_pc and local_docker tasks in this release"
	lspCloseSessionNotFoundText  = "session not found"
	lspProxyWriteTimeout         = 10 * time.Second
)

type lspMessageWriter interface {
	SetWriteDeadline(time.Time) error
	WriteMessage(int, []byte) error
}

// LSPUserService provides user settings for the LSP handler.
type LSPUserService interface {
	GetUserSettings(ctx context.Context) (*models.UserSettings, error)
}

type lspLifecycleManager interface {
	CheckSessionAccess(ctx context.Context, sessionID string) error
	ResolveSessionRuntime(ctx context.Context, sessionID string) (agentruntime.Runtime, error)
	GetOrEnsureExecution(ctx context.Context, sessionID string) (*lifecycle.AgentExecution, error)
}

var (
	errLSPUnsupportedExecutor = errors.New("unsupported LSP executor")
	errLSPCapacityExceeded    = errors.New("active LSP connection cap exceeded")
)

// LSPHandler handles browser-facing LSP WebSocket connections.
// The backend owns session/runtime policy and proxies raw LSP traffic to the
// task host's agentctl instance, where the language server process runs.
type LSPHandler struct {
	lifecycleMgr      lspLifecycleManager
	userService       LSPUserService
	capacity          *lspCapacityLimiter
	logger            *logger.Logger
	continuityEnabled bool
	leases            *lspLeaseManager
	sessionFence      func(string) func()
	settingsSub       bus.Subscription
}

// EnableContinuity installs the runtime-owned lease manager. The session
// fence serializes lease admission with idle execution reclaim.
func (h *LSPHandler) EnableContinuity(acquireFence func(string) func(), eventBus bus.EventBus) {
	h.continuityEnabled = true
	h.sessionFence = acquireFence
	h.leases = newLSPLeaseManager(h.capacity.max, h.logger)
	if eventBus == nil {
		return
	}
	subscription, err := eventBus.Subscribe(events.UserSettingsUpdated, h.handleUserSettingsUpdated)
	if err != nil {
		h.logger.Warn("failed to subscribe LSP leases to user settings updates", zap.Error(err))
		return
	}
	h.settingsSub = subscription
}

// Close releases retained language servers and the settings subscription.
func (h *LSPHandler) Close() error {
	if h.settingsSub != nil && h.settingsSub.IsValid() {
		_ = h.settingsSub.Unsubscribe()
	}
	if h.leases != nil {
		h.leases.Close()
	}
	return nil
}

var lspUpgrader = gorillaws.Upgrader{
	ReadBufferSize:  8192,
	WriteBufferSize: 8192,
	CheckOrigin:     checkWebSocketOrigin,
}

// NewLSPHandler creates a new LSPHandler.
func NewLSPHandler(lifecycleMgr *lifecycle.Manager, userService LSPUserService, log *logger.Logger, configuredMax ...int) *LSPHandler {
	capacity := newLSPCapacityLimiterFromEnv()
	if len(configuredMax) > 0 {
		capacity = newLSPCapacityLimiter(configuredMax[0])
	}
	return &LSPHandler{
		lifecycleMgr: lifecycleMgr,
		userService:  userService,
		capacity:     capacity,
		logger:       log.WithFields(zap.String("component", "lsp_handler")),
	}
}

// HandleLSPConnection handles WebSocket connections at /lsp/:sessionId?language=...
func (h *LSPHandler) HandleLSPConnection(c *gin.Context) {
	sessionID := c.Param("sessionId")
	if sessionID == "" {
		c.JSON(http.StatusBadRequest, gin.H{"error": "sessionId is required"})
		return
	}

	language := c.Query("language")
	if language == "" {
		c.JSON(http.StatusBadRequest, gin.H{"error": "language query parameter is required"})
		return
	}
	if !installer.IsSupported(language) {
		c.JSON(http.StatusBadRequest, gin.H{"error": fmt.Sprintf("unsupported language: %s", language)})
		return
	}
	if h.continuityEnabled {
		h.handleContinuityConnection(c, sessionID, language)
		return
	}

	h.logger.Info("LSP WebSocket connection request",
		zap.String("session_id", sessionID),
		zap.String("language", language))

	execution, runtimeName, err := h.resolveLSPExecution(c.Request.Context(), sessionID)
	if err != nil {
		if errors.Is(err, errLSPCapacityExceeded) {
			h.logger.Info("LSP: capacity exceeded",
				zap.String("session_id", sessionID),
				zap.String("language", language))
			h.closeWithCode(c, lspCloseCapacityExceeded, errLSPCapacityExceeded.Error())
			return
		}
		if errors.Is(err, errLSPUnsupportedExecutor) {
			h.logger.Info("LSP: unsupported runtime",
				zap.String("session_id", sessionID),
				zap.Stringer("runtime", runtimeName))
			h.closeWithCode(c, lspCloseUnsupportedExecutor, lspCloseUnsupportedCloseText)
			return
		}
		h.logger.Warn("LSP: session not found in lifecycle manager",
			zap.String("session_id", sessionID),
			zap.Error(err))
		h.closeWithCode(c, lspCloseSessionNotFound, lspCloseSessionNotFoundText)
		return
	}
	defer h.capacity.Release()
	agentctlClient, releaseClient := execution.AcquireAgentCtlClient()
	if agentctlClient == nil {
		releaseClient()
		h.logger.Warn("LSP: execution has no agentctl client", zap.String("session_id", sessionID))
		h.closeWithCode(c, lspCloseSessionNotFound, "agentctl unavailable")
		return
	}
	browserConn, upgradeErr := lspUpgrader.Upgrade(c.Writer, c.Request, nil)
	if upgradeErr != nil {
		releaseClient()
		h.logger.Error("LSP: failed to upgrade to WebSocket",
			zap.String("session_id", sessionID),
			zap.Error(upgradeErr))
		return
	}
	browserConn.SetReadLimit(protocol.MaxMessageBytes)

	autoInstall := h.shouldAutoInstall(c.Request.Context(), language)
	upstreamConn, resp, dialErr := agentctlClient.DialLSP(c.Request.Context(), language, autoInstall)
	releaseClient()
	if dialErr != nil {
		status := 0
		if resp != nil {
			status = resp.StatusCode
		}
		h.logger.Warn("LSP: failed to connect to agentctl LSP stream",
			zap.String("session_id", sessionID),
			zap.String("language", language),
			zap.Int("status", status),
			zap.Error(dialErr))
		closeLSPConnWithCode(browserConn, lspCloseSessionNotFound, "failed to connect to task host LSP stream")
		return
	}
	upstreamConn.SetReadLimit(protocol.MaxMessageBytes)
	defer func() { _ = upstreamConn.Close() }()

	h.proxyLSPConnections(browserConn, upstreamConn, sessionID, language)
}

func (h *LSPHandler) handleContinuityConnection(c *gin.Context, sessionID, language string) {
	releaseFence := func() {}
	fenceHeld := true
	if h.sessionFence != nil {
		releaseFence = h.sessionFence(sessionID)
	}
	defer func() {
		if fenceHeld {
			releaseFence()
		}
	}()
	execution, closeCode, closeMessage := h.resolveContinuityExecution(c.Request.Context(), sessionID)
	if closeCode != 0 {
		h.closeWithCode(c, closeCode, closeMessage)
		return
	}
	userID, configuration, autoInstall := h.continuityUserSettings(c.Request.Context(), language)

	browserConn, err := lspUpgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		h.logger.Error("LSP: failed to upgrade to WebSocket", zap.String("session_id", sessionID), zap.Error(err))
		return
	}
	browserConn.SetReadLimit(protocol.MaxMessageBytes)
	lease, generation, _, err := h.leases.attachOrCreate(
		c.Request.Context(),
		lspLeaseExecutionFromAgentExecution(execution),
		language,
		userID,
		c.Query("leaseId"),
		configuration,
		autoInstall,
		browserConn,
	)
	if err != nil {
		code, message := lspLeaseAdmissionClose(err)
		closeLSPConnWithCode(browserConn, code, message)
		return
	}
	// Admission is now represented by the live lease pin. Do not hold the
	// session lifecycle fence for the lifetime of the browser socket.
	releaseFence()
	fenceHeld = false

	h.serveContinuityConnection(browserConn, lease, generation, sessionID, language)
}

func (h *LSPHandler) resolveContinuityExecution(
	ctx context.Context,
	sessionID string,
) (*lifecycle.AgentExecution, int, string) {
	if h.lifecycleMgr == nil || h.leases == nil {
		return nil, lspCloseSessionNotFound, "LSP lease manager unavailable"
	}
	if h.lifecycleMgr.CheckSessionAccess(ctx, sessionID) != nil {
		return nil, lspCloseSessionNotFound, lspCloseSessionNotFoundText
	}
	runtimeName, err := h.lifecycleMgr.ResolveSessionRuntime(ctx, sessionID)
	if err != nil {
		return nil, lspCloseSessionNotFound, lspCloseSessionNotFoundText
	}
	if !lspRuntimeSupported(runtimeName) {
		return nil, lspCloseUnsupportedExecutor, lspCloseUnsupportedCloseText
	}
	execution, err := h.lifecycleMgr.GetOrEnsureExecution(ctx, sessionID)
	if err != nil || execution == nil || execution.SessionID != sessionID {
		return nil, lspCloseSessionNotFound, "agentctl unavailable"
	}
	if !lspRuntimeSupported(execution.RuntimeName) {
		return nil, lspCloseUnsupportedExecutor, lspCloseUnsupportedCloseText
	}
	return execution, 0, ""
}

func (h *LSPHandler) continuityUserSettings(
	ctx context.Context,
	language string,
) (string, map[string]any, bool) {
	userID := ""
	if identity, ok := authn.IdentityFromContext(ctx); ok {
		userID = identity.UserID
	}
	if h.userService == nil {
		return userID, make(map[string]any), false
	}
	settings, err := h.userService.GetUserSettings(ctx)
	if err != nil || settings == nil {
		if err != nil {
			h.logger.Debug("LSP: failed to load settings for a retained lease", zap.Error(err))
		}
		return userID, make(map[string]any), false
	}
	if userID == "" {
		userID = settings.UserID
	}
	configuration := make(map[string]any)
	for key, value := range settings.LspServerConfigs[language] {
		configuration[key] = value
	}
	return userID, configuration, h.settingsAllowAutoInstall(settings, language)
}

func lspLeaseAdmissionClose(err error) (int, string) {
	if errors.Is(err, errLSPCapacityExceeded) {
		return lspCloseCapacityExceeded, errLSPCapacityExceeded.Error()
	}
	return lspCloseSessionNotFound, "failed to connect to task host LSP stream"
}

func (h *LSPHandler) serveContinuityConnection(
	browserConn *gorillaws.Conn,
	lease *lspLease,
	generation uint64,
	sessionID, language string,
) {
	defer lease.detach(generation)
	for {
		messageType, message, readErr := browserConn.ReadMessage()
		if readErr != nil {
			return
		}
		if messageType != gorillaws.TextMessage {
			continue
		}
		if err := lease.handleBrowserMessage(generation, message); err != nil {
			h.logger.Debug("LSP browser frame failed", zap.String("session_id", sessionID), zap.String("language", language), zap.Error(err))
			return
		}
	}
}

func lspLeaseExecutionFromAgentExecution(execution *lifecycle.AgentExecution) lspLeaseExecution {
	if execution == nil {
		return lspLeaseExecution{}
	}
	return lspLeaseExecution{
		ID:        execution.ID,
		SessionID: execution.SessionID,
		TaskID:    execution.TaskID,
		acquireAgentCtlClient: func() (lspLeaseAgentCtlClient, func()) {
			client, release := execution.AcquireAgentCtlClient()
			return client, release
		},
	}
}

func (h *LSPHandler) resolveLSPExecution(
	ctx context.Context,
	sessionID string,
) (*lifecycle.AgentExecution, agentruntime.Runtime, error) {
	runtimeName, err := h.lifecycleMgr.ResolveSessionRuntime(ctx, sessionID)
	if err != nil {
		return nil, runtimeName, err
	}
	if !lspRuntimeSupported(runtimeName) {
		return nil, runtimeName, errLSPUnsupportedExecutor
	}
	if !h.capacity.TryAcquire() {
		return nil, runtimeName, errLSPCapacityExceeded
	}
	releaseCapacity := true
	defer func() {
		if releaseCapacity {
			h.capacity.Release()
		}
	}()
	execution, err := h.lifecycleMgr.GetOrEnsureExecution(ctx, sessionID)
	if err != nil {
		return nil, runtimeName, err
	}
	if execution == nil {
		return nil, runtimeName, errors.New("lifecycle manager returned no execution")
	}
	if !lspRuntimeSupported(execution.RuntimeName) {
		return nil, execution.RuntimeName, errLSPUnsupportedExecutor
	}
	releaseCapacity = false
	return execution, execution.RuntimeName, nil
}

func lspRuntimeSupported(runtimeName agentruntime.Runtime) bool {
	return runtimeName == agentruntime.RuntimeStandalone || runtimeName == agentruntime.RuntimeDocker
}

func (h *LSPHandler) shouldAutoInstall(ctx context.Context, language string) bool {
	if h.userService == nil || !installer.SupportsAutoInstall(language) {
		return false
	}
	settings, err := h.userService.GetUserSettings(ctx)
	if err != nil || settings == nil {
		if err != nil {
			h.logger.Debug("LSP: failed to load user settings for auto-install", zap.Error(err))
		}
		return false
	}
	return h.settingsAllowAutoInstall(settings, language)
}

func (h *LSPHandler) settingsAllowAutoInstall(settings *models.UserSettings, language string) bool {
	if settings == nil || !installer.SupportsAutoInstall(language) {
		return false
	}
	for _, lang := range settings.LspAutoInstallLanguages {
		if lang == language {
			return true
		}
	}
	return false
}

// closeWithCode upgrades the WebSocket and immediately closes it with the given code.
func (h *LSPHandler) closeWithCode(c *gin.Context, code int, text string) {
	conn, err := lspUpgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		h.logger.Error("LSP: failed to upgrade WebSocket for close", zap.Error(err))
		return
	}
	closeLSPConnWithCode(conn, code, text)
}

func (h *LSPHandler) proxyLSPConnections(browserConn, upstreamConn *gorillaws.Conn, sessionID, language string) {
	var closeOnce sync.Once
	closeBoth := func() {
		closeOnce.Do(func() {
			_ = browserConn.Close()
			_ = upstreamConn.Close()
		})
	}

	done := make(chan struct{}, 2)
	go func() {
		h.copyLSPMessages("agentctl->browser", upstreamConn, browserConn, sessionID, language)
		done <- struct{}{}
	}()
	go func() {
		h.copyLSPMessages("browser->agentctl", browserConn, upstreamConn, sessionID, language)
		done <- struct{}{}
	}()

	<-done
	closeBoth()
	<-done

	h.logger.Info("LSP: connection closed",
		zap.String("session_id", sessionID),
		zap.String("language", language))
}

func (h *LSPHandler) copyLSPMessages(
	direction string,
	src *gorillaws.Conn,
	dst lspMessageWriter,
	sessionID, language string,
) {
	for {
		messageType, msg, err := src.ReadMessage()
		if err != nil {
			h.forwardLSPClose(dst, err)
			if !gorillaws.IsCloseError(err, gorillaws.CloseNormalClosure, gorillaws.CloseGoingAway) {
				h.logger.Debug("LSP proxy read error",
					zap.String("direction", direction),
					zap.String("session_id", sessionID),
					zap.String("language", language),
					zap.Error(err))
			}
			return
		}
		if err := writeLSPProxyMessage(dst, messageType, msg); err != nil {
			h.logger.Debug("LSP proxy write error",
				zap.String("direction", direction),
				zap.String("session_id", sessionID),
				zap.String("language", language),
				zap.Error(err))
			return
		}
	}
}

func (h *LSPHandler) forwardLSPClose(dst lspMessageWriter, err error) {
	if closeErr, ok := err.(*gorillaws.CloseError); ok {
		if closeErr.Code == gorillaws.CloseNoStatusReceived || closeErr.Code == gorillaws.CloseAbnormalClosure {
			_ = writeLSPProxyMessage(
				dst,
				gorillaws.CloseMessage,
				gorillaws.FormatCloseMessage(lspCloseTransport, "LSP transport failed"),
			)
			return
		}
		_ = writeLSPProxyMessage(
			dst,
			gorillaws.CloseMessage,
			gorillaws.FormatCloseMessage(closeErr.Code, closeErr.Text),
		)
		return
	}
	_ = writeLSPProxyMessage(
		dst,
		gorillaws.CloseMessage,
		gorillaws.FormatCloseMessage(lspCloseStreamError, "LSP stream closed"),
	)
}

func writeLSPProxyMessage(dst lspMessageWriter, messageType int, payload []byte) error {
	if err := dst.SetWriteDeadline(time.Now().Add(lspProxyWriteTimeout)); err != nil {
		return err
	}
	return dst.WriteMessage(messageType, payload)
}

func closeLSPConnWithCode(conn *gorillaws.Conn, code int, text string) {
	_ = writeLSPProxyMessage(conn, gorillaws.CloseMessage, gorillaws.FormatCloseMessage(code, text))
	_ = conn.Close()
}