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)) }