feat: enhance completion service with session management and language support
- Introduced session management using Redis for tracking active sessions. - Added session claiming and releasing functionality in the completion manager. - Enhanced HTTP and WebSocket completion endpoints to support multiple languages. - Implemented request timeout and maximum body size configurations for API routes. - Updated client-side code to handle session IDs and language parameters in completion requests. - Improved error handling for unsupported languages and session conflicts. - Added tests for the completion manager to ensure proper session handling and cleanup.
This commit is contained in:
314
backend/internal/completion/manager.go
Normal file
314
backend/internal/completion/manager.go
Normal file
@@ -0,0 +1,314 @@
|
||||
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))
|
||||
}
|
||||
125
backend/internal/completion/manager_test.go
Normal file
125
backend/internal/completion/manager_test.go
Normal file
@@ -0,0 +1,125 @@
|
||||
package completion
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
type fakeRuntimeClient struct {
|
||||
closed bool
|
||||
}
|
||||
|
||||
func (f *fakeRuntimeClient) DidOpen(_ context.Context, _ string, _ string, _ int) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (f *fakeRuntimeClient) DidChange(_ context.Context, _ string, _ string, _ int) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (f *fakeRuntimeClient) Completion(_ context.Context, _ string, _ int, _ int) (Response, error) {
|
||||
return Response{
|
||||
Items: []Item{{Label: "ok"}},
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (f *fakeRuntimeClient) Close() error {
|
||||
f.closed = true
|
||||
return nil
|
||||
}
|
||||
|
||||
func TestManagerCompleteSelectsLanguageAndSession(t *testing.T) {
|
||||
createCount := 0
|
||||
factory := func(_ context.Context, spec LanguageServerSpec, _ string) (RuntimeClient, error) {
|
||||
createCount++
|
||||
if spec.Language != "go" {
|
||||
return nil, fmt.Errorf("unexpected language: %s", spec.Language)
|
||||
}
|
||||
return &fakeRuntimeClient{}, nil
|
||||
}
|
||||
|
||||
m := NewManager(ManagerConfig{WorkspaceDir: "."}, []LanguageServerSpec{
|
||||
{Language: "go", LanguageID: "go", Command: "gopls"},
|
||||
}, factory)
|
||||
defer m.Close()
|
||||
|
||||
for i := 0; i < 2; i++ {
|
||||
resp, err := m.Complete(context.Background(), Request{
|
||||
Language: "go",
|
||||
SessionID: "s1",
|
||||
URI: "file:///main.go",
|
||||
Text: "package main",
|
||||
Line: 0,
|
||||
Character: 0,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Complete() error = %v", err)
|
||||
}
|
||||
if len(resp.Items) != 1 || resp.Items[0].Label != "ok" {
|
||||
t.Fatalf("unexpected response: %+v", resp)
|
||||
}
|
||||
}
|
||||
|
||||
if createCount != 1 {
|
||||
t.Fatalf("expected one client creation, got %d", createCount)
|
||||
}
|
||||
}
|
||||
|
||||
func TestManagerCompleteUnsupportedLanguage(t *testing.T) {
|
||||
m := NewManager(ManagerConfig{}, []LanguageServerSpec{
|
||||
{Language: "go", LanguageID: "go", Command: "gopls"},
|
||||
}, func(_ context.Context, _ LanguageServerSpec, _ string) (RuntimeClient, error) {
|
||||
return &fakeRuntimeClient{}, nil
|
||||
})
|
||||
defer m.Close()
|
||||
|
||||
_, err := m.Complete(context.Background(), Request{
|
||||
Language: "python",
|
||||
URI: "file:///main.py",
|
||||
Text: "print('hi')",
|
||||
Line: 0,
|
||||
Character: 0,
|
||||
})
|
||||
if !errors.Is(err, ErrUnsupportedLanguage) {
|
||||
t.Fatalf("expected ErrUnsupportedLanguage, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestManagerCleanupIdleSession(t *testing.T) {
|
||||
client := &fakeRuntimeClient{}
|
||||
m := NewManager(ManagerConfig{
|
||||
WorkspaceDir: ".",
|
||||
SessionTTL: 30 * time.Millisecond,
|
||||
CleanupInterval: 10 * time.Millisecond,
|
||||
}, []LanguageServerSpec{
|
||||
{Language: "go", LanguageID: "go", Command: "gopls"},
|
||||
}, func(_ context.Context, _ LanguageServerSpec, _ string) (RuntimeClient, error) {
|
||||
return client, nil
|
||||
})
|
||||
defer m.Close()
|
||||
|
||||
_, err := m.Complete(context.Background(), Request{
|
||||
Language: "go",
|
||||
SessionID: "s2",
|
||||
URI: "file:///main.go",
|
||||
Text: "package main",
|
||||
Line: 0,
|
||||
Character: 0,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Complete() error = %v", err)
|
||||
}
|
||||
|
||||
time.Sleep(90 * time.Millisecond)
|
||||
|
||||
sessions := m.ActiveSessions()
|
||||
if sessions["go"] != 0 {
|
||||
t.Fatalf("expected session cleanup, got %+v", sessions)
|
||||
}
|
||||
if !client.closed {
|
||||
t.Fatal("expected client to be closed by cleanup")
|
||||
}
|
||||
}
|
||||
@@ -10,6 +10,8 @@ import (
|
||||
var ErrInvalidRequest = errors.New("invalid completion request")
|
||||
|
||||
type Request struct {
|
||||
Language string `json:"language,omitempty"`
|
||||
SessionID string `json:"sessionId,omitempty"`
|
||||
URI string `json:"uri"`
|
||||
Text string `json:"text"`
|
||||
Line int `json:"line"`
|
||||
|
||||
Reference in New Issue
Block a user