Files

317 lines
8.9 KiB
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
}
}
// Intercept slash commands before forwarding to agent.
if dispatch := services.TrySlashDispatch(req.Message); dispatch.Handled {
dispatchSlashCommand(c, dispatch, piClient)
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})
}
}
// dispatchSlashCommand handles slash commands intercepted from /api/chat.
func dispatchSlashCommand(c *gin.Context, d services.SlashDispatchResult, piClient *rpc.Client) {
// Unsupported command (not in the whitelist)
if d.Message != "" {
c.JSON(http.StatusOK, gin.H{"ok": true, "slash": true, "message": d.Message})
return
}
switch d.Command {
case "new":
slashNew(c, piClient)
case "reload":
slashReload(c, piClient)
case "copy":
slashCopy(c, piClient)
default:
// Forward to agent as a prompt; SSE events carry the result back.
msg := "/" + d.Command
if d.Args != "" {
msg += " " + d.Args
}
if err := piClient.SubmitPrompt(models.ChatRequest{Message: msg}); err != nil {
c.JSON(http.StatusInternalServerError, models.ErrorResponse{Error: err.Error()})
return
}
c.JSON(http.StatusAccepted, gin.H{"ok": true, "slash": true, "accepted": true})
}
}
func slashNew(c *gin.Context, piClient *rpc.Client) {
resp, err := piClient.SendCmd(models.RPCCommand{Type: "new_session"})
if err != nil || !resp.Success {
msg := "新建会话失败"
if err != nil {
msg = err.Error()
} else if resp.Error != "" {
msg = resp.Error
}
c.JSON(http.StatusInternalServerError, models.ErrorResponse{Error: msg})
return
}
out := gin.H{"ok": true, "slash": true, "action": "new_session"}
if state, serr := piClient.SendCmd(models.RPCCommand{Type: "get_state"}); serr == nil && state.Success {
var sd struct {
SessionFile string `json:"sessionFile"`
}
if json.Unmarshal(state.Data, &sd) == nil && sd.SessionFile != "" {
out["sessionFile"] = sd.SessionFile
}
}
c.JSON(http.StatusOK, out)
}
func slashReload(c *gin.Context, piClient *rpc.Client) {
reloadAgent(piClient)
c.JSON(http.StatusOK, gin.H{"ok": true, "slash": true, "action": "reload_sessions", "message": "已重新加载"})
}
func slashCopy(c *gin.Context, piClient *rpc.Client) {
resp, err := piClient.SendCmd(models.RPCCommand{Type: "get_messages"})
if err != nil || !resp.Success {
c.JSON(http.StatusOK, gin.H{"ok": true, "slash": true, "message": "获取消息失败"})
return
}
var msgs []struct {
Role string `json:"role"`
Content any `json:"content"`
}
if json.Unmarshal(resp.Data, &msgs) != nil {
c.JSON(http.StatusOK, gin.H{"ok": true, "slash": true, "message": "没有可复制的消息"})
return
}
for i := len(msgs) - 1; i >= 0; i-- {
if msgs[i].Role == "assistant" {
text := ""
switch v := msgs[i].Content.(type) {
case string:
text = v
default:
b, _ := json.Marshal(v)
text = string(b)
}
c.JSON(http.StatusOK, gin.H{"ok": true, "slash": true, "action": "copy", "message": text})
return
}
}
c.JSON(http.StatusOK, gin.H{"ok": true, "slash": true, "message": "没有助手消息可复制"})
}
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
}