sundynix-agentix/sundynix-gateway/internal/handler/task_handler.go

// Package handler 实现网关的 HTTP 处理器。
package handler

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

	"github.com/gin-gonic/gin"

	"github.com/sundynix/sundynix-gateway/internal/dsl"
	"github.com/sundynix/sundynix-gateway/internal/nats"
	"github.com/sundynix/sundynix-gateway/internal/store"
	"github.com/sundynix/sundynix-shared/contract"
)

type Handler struct {
	db    *store.Postgres
	cache *store.Redis
	bus   *nats.Bus
}

func New(db *store.Postgres, cache *store.Redis, bus *nats.Bus) *Handler {
	return &Handler{db: db, cache: cache, bus: bus}
}

// SubmitTask: 解析客户端导出的 JSON DSL，组装为 Task，Publish 到 sundynix.tasks.*。
func (h *Handler) SubmitTask(c *gin.Context) {
	var raw json.RawMessage
	if err := c.ShouldBindJSON(&raw); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}
	task, err := dsl.ParseAndAssemble(raw) // Task DSL Parser & Assembly
	if err != nil {
		c.JSON(http.StatusUnprocessableEntity, gin.H{"error": err.Error()})
		return
	}
	// 附上已登录用户标识，供 Dispatcher 召回其偏好记忆。
	// 真实场景由鉴权中间件注入；此处用 X-User-ID 头，缺省匿名。
	task.Meta[contract.MetaUserID] = userID(c)
	// 持久化任务提交（best-effort：降级模式下静默跳过，不阻断发布）。
	if err := h.db.SaveTask(c.Request.Context(), task.ID, string(task.Graph)); err != nil {
		log.Printf("[gateway] save task %s failed: %v", task.ID, err)
	}
	if err := h.bus.PublishTask(c.Request.Context(), task); err != nil {
		c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()})
		return
	}
	c.JSON(http.StatusAccepted, gin.H{"task_id": task.ID})
}

// StreamTask: 订阅 sundynix.streams.<task_id>，以 SSE 把零拷贝 Token Stream 推给客户端。
func (h *Handler) StreamTask(c *gin.Context) {
	taskID := c.Param("id")
	c.Writer.Header().Set("Content-Type", "text/event-stream")
	c.Writer.Header().Set("Cache-Control", "no-cache")
	c.Writer.Header().Set("Connection", "keep-alive")

	tokens := make(chan []byte, 256)
	done := make(chan struct{})
	unsub, err := h.bus.SubscribeTokens(taskID,
		func(tok []byte) {
			select {
			case tokens <- tok:
			default: // 背压保护：客户端过慢则丢弃，避免阻塞 NATS 回调
			}
		},
		func() { close(done) },
	)
	if err != nil {
		c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()})
		return
	}
	defer func() { _ = unsub() }()

	// gin 的流式写：返回 false 即结束响应。
	c.Stream(func(w io.Writer) bool {
		select {
		case tok := <-tokens:
			c.SSEvent("token", string(tok))
			return true
		case <-done:
			c.SSEvent("done", taskID)
			return false
		case <-c.Request.Context().Done():
			return false
		}
	})
}

// SetMemory: 写入/更新一条用户偏好记忆，经 NATS 调 mcp-go 的 memory_upsert 工具。
// 桌面端"偏好记忆面板"可用它让用户显式登记/纠正模型对自己的记忆。
func (h *Handler) SetMemory(c *gin.Context) {
	var body struct {
		Key   string `json:"key"`
		Value string `json:"value"`
	}
	if err := c.ShouldBindJSON(&body); err != nil || body.Key == "" {
		c.JSON(http.StatusBadRequest, gin.H{"error": "key/value required"})
		return
	}
	res, err := h.bus.CallTool(c.Request.Context(), contract.ToolSubjectGo("memory_upsert"),
		&contract.ToolCall{Tool: "memory_upsert", Args: map[string]any{
			"user_id": userID(c), "key": body.Key, "value": body.Value,
		}})
	if err != nil {
		c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()})
		return
	}
	if !res.OK {
		c.JSON(http.StatusUnprocessableEntity, gin.H{"error": res.Error})
		return
	}
	c.JSON(http.StatusOK, gin.H{"status": "ok", "message": res.Content})
}

// userID 从请求取已登录用户标识（真实场景应由鉴权中间件注入）。
func userID(c *gin.Context) string {
	if u := c.GetHeader("X-User-ID"); u != "" {
		return u
	}
	return "anonymous"
}

func (h *Handler) Billing(c *gin.Context) {
	// TODO: 商业化与计费模块；暂以已提交任务计数演示真实读库。
	n, err := h.db.CountTasks(c.Request.Context())
	if err != nil {
		c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()})
		return
	}
	c.JSON(http.StatusOK, gin.H{"status": "ok", "tasks_submitted": n, "persisted": h.db.Enabled()})
}