all

backend/internal/api/ws_handler.go

@@ -0,0 +1,101 @@
+package api
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"net/http"
+	"sync"
+	"time"
+
+	"github.com/gin-gonic/gin"
+	"github.com/gorilla/websocket"
+
+	"monica-go-completion-backend/internal/completion"
+)
+
+var wsUpgrader = websocket.Upgrader{
+	ReadBufferSize:  4096,
+	WriteBufferSize: 4096,
+	CheckOrigin: func(_ *http.Request) bool {
+		return true
+	},
+}
+
+type wsCompletionRequest struct {
+	ID        string `json:"id"`
+	URI       string `json:"uri"`
+	Text      string `json:"text"`
+	Line      int    `json:"line"`
+	Character int    `json:"character"`
+}
+
+type wsCompletionResponse struct {
+	ID           string            `json:"id"`
+	Items        []completion.Item `json:"items,omitempty"`
+	IsIncomplete bool              `json:"isIncomplete,omitempty"`
+	Error        string            `json:"error,omitempty"`
+}
+
+func registerWSRoutes(router *gin.Engine, service CompletionService) {
+	router.GET("/ws/completions/go", func(c *gin.Context) {
+		conn, err := wsUpgrader.Upgrade(c.Writer, c.Request, nil)
+		if err != nil {
+			return
+		}
+		defer conn.Close()
+
+		var writeMu sync.Mutex
+
+		for {
+			_, payload, err := conn.ReadMessage()
+			if err != nil {
+				break
+			}
+
+			var req wsCompletionRequest
+			if err := json.Unmarshal(payload, &req); err != nil {
+				sendWSResponse(conn, &writeMu, wsCompletionResponse{
+					ID:    "",
+					Error: "invalid JSON payload",
+				})
+				continue
+			}
+
+			go func(r wsCompletionRequest) {
+				ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+				defer cancel()
+
+				resp, err := service.Complete(ctx, completion.Request{
+					URI:       r.URI,
+					Text:      r.Text,
+					Line:      r.Line,
+					Character: r.Character,
+				})
+				if err != nil {
+					msg := "completion failed"
+					if errors.Is(err, completion.ErrInvalidRequest) {
+						msg = err.Error()
+					}
+					sendWSResponse(conn, &writeMu, wsCompletionResponse{
+						ID:    r.ID,
+						Error: msg,
+					})
+					return
+				}
+
+				sendWSResponse(conn, &writeMu, wsCompletionResponse{
+					ID:           r.ID,
+					Items:        resp.Items,
+					IsIncomplete: resp.IsIncomplete,
+				})
+			}(req)
+		}
+	})
+}
+
+func sendWSResponse(conn *websocket.Conn, writeMu *sync.Mutex, resp wsCompletionResponse) {
+	writeMu.Lock()
+	defer writeMu.Unlock()
+	_ = conn.WriteJSON(resp)
+}