MonocoEditor-With-Lsp-Backend/backend/internal/completion/manager.go

package completion

import (
	"context"
	"errors"
	"fmt"
	"slices"
	"strings"
	"sync"
	"time"
)

var ErrUnsupportedLanguage = errors.New("unsupported language")
var ErrTooManySessions = errors.New("too many active lsp sessions")

type RuntimeClient interface {
	Client
	Close() error
}

type SessionRegistry interface {
	ClaimSession(ctx context.Context, language, sessionID string) (ownerID string, ownerEndpoint string, err error)
	ReleaseSession(ctx context.Context, language, sessionID string) error
	Close() error
}

type LanguageServerSpec struct {
	Language   string
	LanguageID string
	Command    string
	Args       []string
}

type ClientFactory func(ctx context.Context, spec LanguageServerSpec, workspaceDir string) (RuntimeClient, error)

type ManagerConfig struct {
	WorkspaceDir    string
	MaxSessions     int
	SessionTTL      time.Duration
	CleanupInterval time.Duration
	InstanceID      string
	Registry        SessionRegistry
}

type Manager struct {
	mu sync.Mutex

	config      ManagerConfig
	specByLang  map[string]LanguageServerSpec
	sessions    map[string]*managedSession
	newClient   ClientFactory
	stopCh      chan struct{}
	stoppedOnce sync.Once
}

type managedSession struct {
	key       string
	sessionID string
	language  string
	service   *Service
	client    RuntimeClient
	lastUsed  time.Time
	createdAt time.Time
}

type ErrSessionOwnedByOtherInstance struct {
	OwnerID       string
	OwnerEndpoint string
}

func (e *ErrSessionOwnedByOtherInstance) Error() string {
	if e.OwnerEndpoint != "" {
		return fmt.Sprintf("session owned by another instance: %s (%s)", e.OwnerID, e.OwnerEndpoint)
	}
	return fmt.Sprintf("session owned by another instance: %s", e.OwnerID)
}

func NewManager(config ManagerConfig, specs []LanguageServerSpec, factory ClientFactory) *Manager {
	if config.MaxSessions <= 0 {
		config.MaxSessions = 256
	}
	if config.SessionTTL <= 0 {
		config.SessionTTL = 20 * time.Minute
	}
	if config.CleanupInterval <= 0 {
		config.CleanupInterval = 2 * time.Minute
	}
	if strings.TrimSpace(config.InstanceID) == "" {
		config.InstanceID = "instance-local"
	}

	specByLang := make(map[string]LanguageServerSpec)
	for _, spec := range specs {
		key := normalizeLanguage(spec.Language)
		if key == "" {
			continue
		}
		specByLang[key] = spec
	}

	m := &Manager{
		config:     config,
		specByLang: specByLang,
		sessions:   make(map[string]*managedSession),
		newClient:  factory,
		stopCh:     make(chan struct{}),
	}
	go m.cleanupLoop()
	return m
}

func (m *Manager) Complete(ctx context.Context, req Request) (Response, error) {
	language := normalizeLanguage(req.Language)
	if language == "" {
		return Response{}, ErrInvalidRequest
	}

	spec, ok := m.specByLang[language]
	if !ok {
		return Response{}, fmt.Errorf("%w: %s", ErrUnsupportedLanguage, req.Language)
	}

	sessionKey := buildSessionKey(language, req.SessionID)
	sessionID := normalizeSessionID(req.SessionID)

	if m.config.Registry != nil {
		ownerID, ownerEndpoint, err := m.config.Registry.ClaimSession(ctx, language, sessionID)
		if err != nil {
			return Response{}, err
		}
		if ownerID != m.config.InstanceID {
			return Response{}, &ErrSessionOwnedByOtherInstance{
				OwnerID:       ownerID,
				OwnerEndpoint: ownerEndpoint,
			}
		}
	}

	session, err := m.getOrCreateSession(ctx, sessionKey, sessionID, spec)
	if err != nil {
		return Response{}, err
	}

	resp, err := session.service.Complete(ctx, req)
	if err != nil {
		return Response{}, err
	}

	m.mu.Lock()
	if current, ok := m.sessions[sessionKey]; ok {
		current.lastUsed = time.Now()
	}
	m.mu.Unlock()
	return resp, nil
}

func (m *Manager) ActiveSessions() map[string]int {
	m.mu.Lock()
	defer m.mu.Unlock()

	out := make(map[string]int)
	for _, session := range m.sessions {
		out[session.language]++
	}
	return out
}

func (m *Manager) Close() error {
	m.stoppedOnce.Do(func() {
		close(m.stopCh)

		m.mu.Lock()
		defer m.mu.Unlock()

		for key, session := range m.sessions {
			if m.config.Registry != nil {
				_ = m.config.Registry.ReleaseSession(context.Background(), session.language, session.sessionID)
			}
			_ = session.client.Close()
			delete(m.sessions, key)
		}
		if m.config.Registry != nil {
			_ = m.config.Registry.Close()
		}
	})
	return nil
}

func (m *Manager) cleanupLoop() {
	ticker := time.NewTicker(m.config.CleanupInterval)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			m.cleanupIdleSessions()
		case <-m.stopCh:
			return
		}
	}
}

func (m *Manager) cleanupIdleSessions() {
	cutoff := time.Now().Add(-m.config.SessionTTL)

	m.mu.Lock()
	defer m.mu.Unlock()

	for key, session := range m.sessions {
		if session.lastUsed.After(cutoff) {
			continue
		}
		if m.config.Registry != nil {
			_ = m.config.Registry.ReleaseSession(context.Background(), session.language, session.sessionID)
		}
		_ = session.client.Close()
		delete(m.sessions, key)
	}
}

func (m *Manager) getOrCreateSession(
	ctx context.Context,
	sessionKey string,
	sessionID string,
	spec LanguageServerSpec,
) (*managedSession, error) {
	m.mu.Lock()
	if existing, ok := m.sessions[sessionKey]; ok {
		existing.lastUsed = time.Now()
		m.mu.Unlock()
		return existing, nil
	}

	if len(m.sessions) >= m.config.MaxSessions {
		if !m.evictLeastRecentlyUsedLocked() {
			m.mu.Unlock()
			return nil, ErrTooManySessions
		}
	}
	m.mu.Unlock()

	client, err := m.newClient(ctx, spec, m.config.WorkspaceDir)
	if err != nil {
		return nil, err
	}

	now := time.Now()
	newSession := &managedSession{
		key:       sessionKey,
		sessionID: sessionID,
		language:  normalizeLanguage(spec.Language),
		service:   NewService(client),
		client:    client,
		lastUsed:  now,
		createdAt: now,
	}

	m.mu.Lock()
	defer m.mu.Unlock()
	if existing, ok := m.sessions[sessionKey]; ok {
		_ = client.Close()
		existing.lastUsed = now
		return existing, nil
	}
	m.sessions[sessionKey] = newSession
	return newSession, nil
}

func (m *Manager) evictLeastRecentlyUsedLocked() bool {
	if len(m.sessions) == 0 {
		return false
	}

	keys := make([]string, 0, len(m.sessions))
	for key := range m.sessions {
		keys = append(keys, key)
	}
	slices.SortFunc(keys, func(a, b string) int {
		as := m.sessions[a]
		bs := m.sessions[b]
		if as.lastUsed.Before(bs.lastUsed) {
			return -1
		}
		if as.lastUsed.After(bs.lastUsed) {
			return 1
		}
		return strings.Compare(as.key, bs.key)
	})

	victimKey := keys[0]
	victim := m.sessions[victimKey]
	if m.config.Registry != nil {
		_ = m.config.Registry.ReleaseSession(context.Background(), victim.language, victim.sessionID)
	}
	_ = victim.client.Close()
	delete(m.sessions, victimKey)
	return true
}

func buildSessionKey(language, sessionID string) string {
	return language + ":" + normalizeSessionID(sessionID)
}

func normalizeSessionID(sessionID string) string {
	sid := strings.TrimSpace(sessionID)
	if sid == "" {
		return "default"
	}
	return sid
}

func normalizeLanguage(language string) string {
	return strings.ToLower(strings.TrimSpace(language))
}