ollie/tools/lsp/bridge.go

package lsp

import (
	"encoding/json"
	"fmt"
	"net"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"
)

const idleTimeout = 5 * time.Minute

// BridgeRequest is the JSON structure sent to the bridge by tool clients.
type BridgeRequest struct {
	Method    string `json:"method"`
	File      string `json:"file"`
	Line      int    `json:"line"`
	Col       int    `json:"col"`
	NewName   string `json:"new_name,omitempty"`
	Query     string `json:"query,omitempty"`
	Workspace bool   `json:"workspace,omitempty"`
	Cwd       string `json:"cwd,omitempty"`
}

// BridgeResponse is the JSON structure returned by the bridge.
type BridgeResponse struct {
	Result json.RawMessage `json:"result,omitempty"`
	Error  string          `json:"error,omitempty"`
}

// Bridge manages LSP server processes and handles tool requests.
type Bridge struct {
	mu           sync.Mutex
	servers      map[string]*Server // key: "adapter:root"
	lastActivity time.Time
}

// NewBridge creates a new bridge instance.
func NewBridge() *Bridge {
	return &Bridge{
		servers:      make(map[string]*Server),
		lastActivity: time.Now(),
	}
}

func (b *Bridge) getServer(adapter *Adapter, root string) (*Server, error) {
	key := adapter.Name + ":" + root
	b.mu.Lock()
	srv := b.servers[key]
	b.mu.Unlock()
	if srv != nil {
		return srv, nil
	}

	srv, err := StartServer(adapter, root)
	if err != nil {
		return nil, err
	}
	b.mu.Lock()
	b.servers[key] = srv
	b.mu.Unlock()
	return srv, nil
}

// Handle processes a single bridge request.
func (b *Bridge) Handle(req *BridgeRequest) *BridgeResponse {
	b.mu.Lock()
	b.lastActivity = time.Now()
	b.mu.Unlock()

	if req.File == "" {
		return &BridgeResponse{Error: "missing 'file' parameter"}
	}

	adapter := Detect(req.File)
	if adapter == nil {
		return &BridgeResponse{Error: fmt.Sprintf("no adapter for file: %s", req.File)}
	}

	root := adapter.WorkspaceRoot(req.File, req.Cwd)
	srv, err := b.getServer(adapter, root)
	if err != nil {
		return &BridgeResponse{Error: fmt.Sprintf("start server: %v", err)}
	}

	var result json.RawMessage

	switch req.Method {
	case "definition":
		srv.EnsureOpen(req.File)
		result, err = srv.Request("textDocument/definition", map[string]any{
			"textDocument": map[string]any{"uri": FileURI(req.File)},
			"position":     map[string]any{"line": req.Line, "character": req.Col},
		}, 30*time.Second)

	case "references":
		srv.EnsureOpen(req.File)
		result, err = srv.Request("textDocument/references", map[string]any{
			"textDocument": map[string]any{"uri": FileURI(req.File)},
			"position":     map[string]any{"line": req.Line, "character": req.Col},
			"context":      map[string]any{"includeDeclaration": true},
		}, 30*time.Second)

	case "hover":
		srv.EnsureOpen(req.File)
		result, err = srv.Request("textDocument/hover", map[string]any{
			"textDocument": map[string]any{"uri": FileURI(req.File)},
			"position":     map[string]any{"line": req.Line, "character": req.Col},
		}, 30*time.Second)

	case "completion":
		srv.EnsureOpen(req.File)
		result, err = srv.Request("textDocument/completion", map[string]any{
			"textDocument": map[string]any{"uri": FileURI(req.File)},
			"position":     map[string]any{"line": req.Line, "character": req.Col},
		}, 30*time.Second)

	case "rename":
		srv.EnsureOpen(req.File)
		result, err = srv.Request("textDocument/rename", map[string]any{
			"textDocument": map[string]any{"uri": FileURI(req.File)},
			"position":     map[string]any{"line": req.Line, "character": req.Col},
			"newName":      req.NewName,
		}, 30*time.Second)

	case "symbols":
		if req.Workspace {
			result, err = srv.Request("workspace/symbol", map[string]any{
				"query": req.Query,
			}, 30*time.Second)
		} else {
			srv.EnsureOpen(req.File)
			result, err = srv.Request("textDocument/documentSymbol", map[string]any{
				"textDocument": map[string]any{"uri": FileURI(req.File)},
			}, 30*time.Second)
		}

	case "diagnostics":
		srv.EnsureOpen(req.File)
		// Wait for push diagnostics
		time.Sleep(1 * time.Second)
		diags := srv.Diagnostics(FileURI(req.File))
		if diags == nil {
			diags = json.RawMessage("[]")
		}
		return &BridgeResponse{Result: diags}

	default:
		return &BridgeResponse{Error: fmt.Sprintf("unknown method: %s", req.Method)}
	}

	if err != nil {
		return &BridgeResponse{Error: err.Error()}
	}
	return &BridgeResponse{Result: result}
}

// Shutdown stops all managed LSP servers.
func (b *Bridge) Shutdown() {
	b.mu.Lock()
	servers := make([]*Server, 0, len(b.servers))
	for _, s := range b.servers {
		servers = append(servers, s)
	}
	b.servers = make(map[string]*Server)
	b.mu.Unlock()

	for _, s := range servers {
		s.Stop()
	}
}

// RunBridge starts the bridge daemon listening on the Unix socket.
// Blocks until idle timeout or signal.
func RunBridge() error {
	sock := SockPath()
	os.Remove(sock)

	ln, err := net.Listen("unix", sock)
	if err != nil {
		return err
	}
	defer ln.Close()
	os.Chmod(sock, 0600)

	// Write PID file
	pidFile := PidPath()
	os.WriteFile(pidFile, []byte(fmt.Sprintf("%d", os.Getpid())), 0600)
	defer os.Remove(pidFile)
	defer os.Remove(sock)

	bridge := NewBridge()
	defer bridge.Shutdown()

	// Idle watchdog
	done := make(chan struct{})
	go func() {
		for {
			time.Sleep(30 * time.Second)
			bridge.mu.Lock()
			idle := time.Since(bridge.lastActivity) > idleTimeout
			bridge.mu.Unlock()
			if idle {
				ln.Close()
				return
			}
		}
	}()

	// Signal handler
	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
	go func() {
		<-sigCh
		ln.Close()
	}()

	// Accept loop
	go func() {
		for {
			conn, err := ln.Accept()
			if err != nil {
				close(done)
				return
			}
			go handleConn(bridge, conn)
		}
	}()

	<-done
	return nil
}

func handleConn(b *Bridge, conn net.Conn) {
	defer conn.Close()
	conn.SetDeadline(time.Now().Add(5 * time.Second))

	var req BridgeRequest
	dec := json.NewDecoder(conn)
	if err := dec.Decode(&req); err != nil {
		resp := BridgeResponse{Error: err.Error()}
		json.NewEncoder(conn).Encode(resp)
		return
	}

	resp := b.Handle(&req)
	json.NewEncoder(conn).Encode(resp)
}