renjue
d440fb1ae7
CI / docker (push) Has been cancelled
引入用户端对话/生图与管理端 Provider 配置，统一 invoke 编排、Gateway、请求日志与模型管理；前端按设计规范拆分布局，Go module 更名为 github.com/rose_cat707/Atlas；生图长任务 SSE 推送；管理端登录防爆破；协议与服务类型联动；Replicate 模型同步修复；用户端创作台 UI 重构；管理端创作历史替代请求日志。 Co-authored-by: Cursor <cursoragent@cursor.com>
85 lines
1.8 KiB
Go
85 lines
1.8 KiB
Go
package api
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"net/http"
|
|
"strconv"
|
|
"time"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
|
|
"github.com/rose_cat707/Atlas/internal/service/taskevents"
|
|
"github.com/rose_cat707/Atlas/internal/store"
|
|
)
|
|
|
|
func (s *Server) streamTaskEvents(c *gin.Context) {
|
|
id, err := strconv.ParseInt(c.Param("id"), 10, 64)
|
|
if err != nil {
|
|
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid task id"})
|
|
return
|
|
}
|
|
ctx := c.Request.Context()
|
|
sid := sessionID(c)
|
|
task, err := s.resolveTask(ctx, sid, id)
|
|
if err != nil {
|
|
c.JSON(http.StatusNotFound, gin.H{"error": "not found"})
|
|
return
|
|
}
|
|
|
|
c.Header("Content-Type", "text/event-stream")
|
|
c.Header("Cache-Control", "no-cache")
|
|
c.Header("Connection", "keep-alive")
|
|
c.Status(http.StatusOK)
|
|
|
|
writeTaskSSE(c, task)
|
|
if taskevents.Terminal(task.Status) {
|
|
writeSSEDone(c)
|
|
return
|
|
}
|
|
|
|
ch := s.app.TaskEvents.Subscribe(id)
|
|
defer s.app.TaskEvents.Unsubscribe(id, ch)
|
|
|
|
heartbeat := time.NewTicker(15 * time.Second)
|
|
defer heartbeat.Stop()
|
|
|
|
for {
|
|
select {
|
|
case <-ctx.Done():
|
|
return
|
|
case t := <-ch:
|
|
writeTaskSSE(c, &t)
|
|
if taskevents.Terminal(t.Status) {
|
|
writeSSEDone(c)
|
|
return
|
|
}
|
|
case <-heartbeat.C:
|
|
fmt.Fprintf(c.Writer, ": heartbeat\n\n")
|
|
c.Writer.Flush()
|
|
}
|
|
}
|
|
}
|
|
|
|
func (s *Server) resolveTask(ctx context.Context, sessionID string, id int64) (*store.Task, error) {
|
|
key := "task:" + strconv.FormatInt(id, 10)
|
|
if v, ok := s.app.Store.Cache.Get(key); ok {
|
|
if t, ok := v.(*store.Task); ok {
|
|
return t, nil
|
|
}
|
|
}
|
|
return s.app.Store.GetTask(ctx, sessionID, id)
|
|
}
|
|
|
|
func writeTaskSSE(c *gin.Context, task *store.Task) {
|
|
data, _ := json.Marshal(task)
|
|
fmt.Fprintf(c.Writer, "event: task\ndata: %s\n\n", data)
|
|
c.Writer.Flush()
|
|
}
|
|
|
|
func writeSSEDone(c *gin.Context) {
|
|
fmt.Fprintf(c.Writer, "event: done\ndata: [DONE]\n\n")
|
|
c.Writer.Flush()
|
|
}
|