sproutclaw-web/backend/internal/handlers/chat.go

package handlers

import (
	"encoding/json"
	"io"
	"net/http"

	"github.com/gin-gonic/gin"
	"sproutclaw-web/internal/models"
	"sproutclaw-web/internal/rpc"
	"sproutclaw-web/internal/services"
)

// RegisterChat registers all chat-related routes.
func RegisterChat(r *gin.RouterGroup, piClient *rpc.Client) {
	r.POST("/chat", handleChat(piClient))
	r.POST("/steer", handleSteer(piClient))
	r.POST("/follow-up", handleFollowUp(piClient))
	r.POST("/bash", handleBash(piClient))
	r.POST("/abort", handleAbort(piClient))
	r.POST("/abort-retry", handleAbortRetry(piClient))
	r.POST("/messages", handleMessages(piClient))
	r.GET("/events", handleSSE(piClient))
}

func handleChat(piClient *rpc.Client) gin.HandlerFunc {
	return func(c *gin.Context) {
		var req models.ChatRequest
		if err := c.ShouldBindJSON(&req); err != nil {
			c.JSON(http.StatusBadRequest, models.ErrorResponse{Error: err.Error()})
			return
		}

		// Validate images
		if len(req.Images) > 0 {
			if err := services.ValidateImages(req.Images); err != nil {
				c.JSON(http.StatusBadRequest, models.ErrorResponse{Error: err.Error()})
				return
			}
		}

		// Try slash command dispatch
		result := services.TrySlashDispatch(req.Message)
		if result.Handled {
			if result.Message != "" {
				// unsupported command – return friendly message
				c.JSON(http.StatusOK, gin.H{"handled": true, "message": result.Message})
				return
			}
			// supported built-in command – forward to pi CLI as a command type
			resp, err := piClient.SendCmd(models.RPCCommand{
				Type: result.Command,
				Args: result.Args,
			})
			if err != nil {
				c.JSON(http.StatusInternalServerError, models.ErrorResponse{Error: err.Error()})
				return
			}
			c.JSON(http.StatusOK, gin.H{"handled": true, "result": resp.Result})
			return
		}

		// Regular prompt
		if err := piClient.SubmitPrompt(req); err != nil {
			c.JSON(http.StatusInternalServerError, models.ErrorResponse{Error: err.Error()})
			return
		}
		c.JSON(http.StatusOK, models.OKResponse{OK: true})
	}
}

func handleSteer(piClient *rpc.Client) gin.HandlerFunc {
	return func(c *gin.Context) {
		var req models.SteerRequest
		if err := c.ShouldBindJSON(&req); err != nil {
			c.JSON(http.StatusBadRequest, models.ErrorResponse{Error: err.Error()})
			return
		}
		err := piClient.SubmitPrompt(models.ChatRequest{
			Message:           req.Message,
			StreamingBehavior: "steer",
		})
		if err != nil {
			c.JSON(http.StatusInternalServerError, models.ErrorResponse{Error: err.Error()})
			return
		}
		c.JSON(http.StatusOK, models.OKResponse{OK: true})
	}
}

func handleFollowUp(piClient *rpc.Client) gin.HandlerFunc {
	return func(c *gin.Context) {
		var req models.ChatRequest
		if err := c.ShouldBindJSON(&req); err != nil {
			c.JSON(http.StatusBadRequest, models.ErrorResponse{Error: err.Error()})
			return
		}
		req.StreamingBehavior = "follow_up"
		if err := piClient.SubmitPrompt(req); err != nil {
			c.JSON(http.StatusInternalServerError, models.ErrorResponse{Error: err.Error()})
			return
		}
		c.JSON(http.StatusOK, models.OKResponse{OK: true})
	}
}

func handleBash(piClient *rpc.Client) gin.HandlerFunc {
	return func(c *gin.Context) {
		var req models.BashRequest
		if err := c.ShouldBindJSON(&req); err != nil {
			c.JSON(http.StatusBadRequest, models.ErrorResponse{Error: err.Error()})
			return
		}
		resp, err := piClient.SendCmd(models.RPCCommand{
			Type: "bash",
			Args: req.Command,
		})
		if err != nil {
			c.JSON(http.StatusInternalServerError, models.ErrorResponse{Error: err.Error()})
			return
		}
		c.JSON(http.StatusOK, resp.Result)
	}
}

func handleAbort(piClient *rpc.Client) gin.HandlerFunc {
	return func(c *gin.Context) {
		_, err := piClient.SendCmd(models.RPCCommand{Type: "abort"})
		if err != nil {
			c.JSON(http.StatusInternalServerError, models.ErrorResponse{Error: err.Error()})
			return
		}
		c.JSON(http.StatusOK, models.OKResponse{OK: true})
	}
}

func handleAbortRetry(piClient *rpc.Client) gin.HandlerFunc {
	return func(c *gin.Context) {
		_, err := piClient.SendCmd(models.RPCCommand{Type: "abort_retry"})
		if err != nil {
			c.JSON(http.StatusInternalServerError, models.ErrorResponse{Error: err.Error()})
			return
		}
		c.JSON(http.StatusOK, models.OKResponse{OK: true})
	}
}

func handleMessages(piClient *rpc.Client) gin.HandlerFunc {
	return func(c *gin.Context) {
		var req models.MessagesRequest
		if err := c.ShouldBindJSON(&req); err != nil {
			c.JSON(http.StatusBadRequest, models.ErrorResponse{Error: err.Error()})
			return
		}
		snap := piClient.GetSnapshot()
		c.JSON(http.StatusOK, gin.H{
			"sessionFile": snap.SessionFile,
			"events":      snap.ReplayEvents,
		})
	}
}

func handleSSE(piClient *rpc.Client) gin.HandlerFunc {
	return func(c *gin.Context) {
		c.Header("Content-Type", "text/event-stream")
		c.Header("Cache-Control", "no-cache")
		c.Header("Connection", "keep-alive")
		c.Header("X-Accel-Buffering", "no")

		flusher, ok := c.Writer.(http.Flusher)
		if !ok {
			c.String(http.StatusInternalServerError, "streaming not supported")
			return
		}

		ch := make(chan []byte, 64)
		client := &rpc.SSEClient{
			Ch:   ch,
			Done: c.Request.Context().Done(),
		}
		piClient.RegisterSSEClient(client)
		defer piClient.UnregisterSSEClient(client)

		// send initial keepalive
		_, _ = io.WriteString(c.Writer, ": keepalive\n\n")
		flusher.Flush()

		for {
			select {
			case <-c.Request.Context().Done():
				return
			case raw, ok := <-ch:
				if !ok {
					return
				}
				frame := rpc.FormatSSEEvent(raw)
				if _, err := c.Writer.Write(frame); err != nil {
					return
				}
				flusher.Flush()
			}
		}
	}
}

// jsonRaw is a helper to return pre-serialized JSON.
func jsonRaw(raw json.RawMessage) any {
	if raw == nil {
		return nil
	}
	var v any
	_ = json.Unmarshal(raw, &v)
	return v
}