317 lines
8.9 KiB
Go
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
|
|
}
|