102 lines
2.2 KiB
Go
102 lines
2.2 KiB
Go
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)
|
|
}
|