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
		}

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

		// normalize streamingBehavior: only "steer" / "followUp" are forwarded
		if req.StreamingBehavior != "steer" && req.StreamingBehavior != "followUp" {
			req.StreamingBehavior = ""
		}

		// Submit prompt; agent output streams back over SSE.
		if err := piClient.SubmitPrompt(req); err != nil {
			c.JSON(http.StatusInternalServerError, models.ErrorResponse{Error: err.Error()})
			return
		}
		c.JSON(http.StatusAccepted, gin.H{"ok": true, "accepted": true})
	}
}

func handleSteer(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
		}
		resp, err := piClient.SendCmd(models.RPCCommand{
			Type:    "steer",
			Message: req.Message,
			Images:  req.Images,
		})
		if err != nil {
			c.JSON(http.StatusInternalServerError, models.ErrorResponse{Error: err.Error()})
			return
		}
		if !resp.Success {
			c.JSON(http.StatusInternalServerError, models.ErrorResponse{Error: orDefault(resp.Error, "steer 失败")})
			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
		}
		resp, err := piClient.SendCmd(models.RPCCommand{
			Type:    "follow_up",
			Message: req.Message,
			Images:  req.Images,
		})
		if err != nil {
			c.JSON(http.StatusInternalServerError, models.ErrorResponse{Error: err.Error()})
			return
		}
		if !resp.Success {
			c.JSON(http.StatusInternalServerError, models.ErrorResponse{Error: orDefault(resp.Error, "follow_up 失败")})
			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",
			Command: req.Command,
		})
		if err != nil {
			c.JSON(http.StatusInternalServerError, models.ErrorResponse{Error: err.Error()})
			return
		}
		if !resp.Success {
			c.JSON(http.StatusInternalServerError, models.ErrorResponse{Error: orDefault(resp.Error, "命令执行失败")})
			return
		}
		c.JSON(http.StatusOK, jsonRaw(resp.Data))
	}
}

func handleAbort(piClient *rpc.Client) gin.HandlerFunc {
	return func(c *gin.Context) {
		if _, err := piClient.SendCmd(models.RPCCommand{Type: "abort"}); 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) {
		if _, err := piClient.SendCmd(models.RPCCommand{Type: "abort_retry"}); 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) {
		resp, err := piClient.SendCmd(models.RPCCommand{Type: "get_messages"})
		if err != nil {
			c.JSON(http.StatusInternalServerError, models.ErrorResponse{Error: err.Error()})
			return
		}
		if !resp.Success {
			c.JSON(http.StatusInternalServerError, models.ErrorResponse{Error: resp.Error})
			return
		}
		c.JSON(http.StatusOK, jsonRaw(resp.Data))
	}
}

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, 256)
		client := &rpc.SSEClient{
			Ch:   ch,
			Done: c.Request.Context().Done(),
		}
		piClient.RegisterSSEClient(client)
		defer piClient.UnregisterSSEClient(client)

		// initial comment to open the stream
		_, _ = io.WriteString(c.Writer, ": connected\n\n")
		flusher.Flush()

		for {
			select {
			case <-c.Request.Context().Done():
				return
			case frame, ok := <-ch:
				if !ok {
					return
				}
				// frame is already a fully-formatted "data: ...\n\n" SSE frame
				if _, err := c.Writer.Write(frame); err != nil {
					return
				}
				flusher.Flush()
			}
		}
	}
}

// jsonRaw decodes pre-serialized JSON for re-emission via gin.
func jsonRaw(raw json.RawMessage) any {
	if len(raw) == 0 {
		return gin.H{}
	}
	var v any
	if err := json.Unmarshal(raw, &v); err != nil {
		return gin.H{}
	}
	return v
}

func orDefault(s, fallback string) string {
	if s == "" {
		return fallback
	}
	return s
}