jarvis/lsp.go

package main

import (
	"bufio"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"sync"

	"github.com/gorilla/websocket"
)

var upgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

type LSPSession struct {
	cmd    *exec.Cmd
	stdin  io.WriteCloser
	stdout io.ReadCloser
	conn   *websocket.Conn
	mu     sync.Mutex
}

func lspHandler(w http.ResponseWriter, r *http.Request) {
	language := r.URL.Query().Get("language")
	if language == "" {
		language = "typescript"
	}

	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Printf("WebSocket upgrade failed: %v", err)
		return
	}
	defer conn.Close()

	session, err := startLanguageServer(language)
	if err != nil {
		log.Printf("Failed to start language server for %s: %v", language, err)
		conn.WriteMessage(websocket.TextMessage, []byte(fmt.Sprintf(`{"error": "Failed to start %s language server"}`, language)))
		return
	}
	session.conn = conn

	defer func() {
		session.stdin.Close()
		session.cmd.Process.Kill()
		session.cmd.Wait()
	}()

	var wg sync.WaitGroup
	wg.Add(2)

	go func() {
		defer wg.Done()
		session.readFromServer()
	}()

	go func() {
		defer wg.Done()
		session.writeToServer()
	}()

	wg.Wait()
}

func startLanguageServer(language string) (*LSPSession, error) {
	var cmd *exec.Cmd

	switch language {
	case "typescript":
		workDir := createTypeScriptWorkspace()
		cmd = exec.Command("typescript-language-server", "--stdio")
		cmd.Dir = workDir
	case "go":
		workDir := createGoWorkspace()
		cmd = exec.Command("gopls", "serve")
		cmd.Dir = workDir
	default:
		return nil, fmt.Errorf("unsupported language: %s", language)
	}

	stdin, err := cmd.StdinPipe()
	if err != nil {
		return nil, err
	}

	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return nil, err
	}

	cmd.Stderr = os.Stderr

	if err := cmd.Start(); err != nil {
		return nil, err
	}

	return &LSPSession{
		cmd:    cmd,
		stdin:  stdin,
		stdout: stdout,
	}, nil
}

func createTypeScriptWorkspace() string {
	dir, _ := os.MkdirTemp("", "lsp-ts-*")

	sdkPath := "/app/sdk/writekit.d.ts"
	if _, err := os.Stat(sdkPath); err != nil {
		sdkPath = "sdk/writekit.d.ts"
	}
	sdkContent, _ := os.ReadFile(sdkPath)

	os.WriteFile(filepath.Join(dir, "writekit.d.ts"), sdkContent, 0644)
	os.WriteFile(filepath.Join(dir, "tsconfig.json"), []byte(`{
  "compilerOptions": {
    "target": "ESNext",
    "module": "ESNext",
    "moduleResolution": "node",
    "strict": true,
    "noEmit": true
  },
  "include": ["*.ts"]
}`), 0644)

	return dir
}

func createGoWorkspace() string {
	dir, _ := os.MkdirTemp("", "lsp-go-*")

	os.WriteFile(filepath.Join(dir, "go.mod"), []byte(`module plugin

go 1.22

require github.com/extism/go-pdk v1.0.6
`), 0644)

	os.WriteFile(filepath.Join(dir, "main.go"), []byte(`package main

import "github.com/extism/go-pdk"

func main() {}

var _ = pdk.Log
`), 0644)

	return dir
}

func (s *LSPSession) readFromServer() {
	reader := bufio.NewReader(s.stdout)

	for {
		header, err := reader.ReadString('\n')
		if err != nil {
			return
		}

		if !strings.HasPrefix(header, "Content-Length:") {
			continue
		}

		lengthStr := strings.TrimSpace(strings.TrimPrefix(header, "Content-Length:"))
		length, err := strconv.Atoi(lengthStr)
		if err != nil {
			continue
		}

		reader.ReadString('\n')

		body := make([]byte, length)
		_, err = io.ReadFull(reader, body)
		if err != nil {
			return
		}

		s.mu.Lock()
		err = s.conn.WriteMessage(websocket.TextMessage, body)
		s.mu.Unlock()

		if err != nil {
			return
		}
	}
}

func (s *LSPSession) writeToServer() {
	for {
		_, message, err := s.conn.ReadMessage()
		if err != nil {
			return
		}

		var msg json.RawMessage
		if err := json.Unmarshal(message, &msg); err != nil {
			continue
		}

		lspMessage := fmt.Sprintf("Content-Length: %d\r\n\r\n%s", len(message), message)
		_, err = s.stdin.Write([]byte(lspMessage))
		if err != nil {
			return
		}
	}
}