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

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

import (
	"context"
	"encoding/json"
	"io"
	"log"
	"net/http"
	"os"
	"strconv"
	"time"

	"github.com/gin-contrib/sse"
	"github.com/gin-gonic/gin"

	"github.com/sundynix/sundynix-gateway/internal/blob"
	"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
	blob  *blob.Store
}

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

// 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
	}
	// 成本护栏：单用户当日 token 日预算门控（USER_DAILY_TOKEN_BUDGET，0=不限）。超额则拒绝新任务。
	if budget := userDailyTokenBudget(); budget > 0 {
		uid := userID(c)
		used := h.cache.GetUsage(c.Request.Context(), uid, time.Now().Format("20060102"))
		if used >= int64(budget) {
			c.JSON(http.StatusPaymentRequired, gin.H{
				"error": "已达当日 token 预算上限", "used": used, "budget": budget,
			})
			return
		}
	}
	// 附上用户标识（召回偏好记忆）与会话标识（召回短期多轮历史）。
	// 真实场景由鉴权/会话中间件注入；此处用请求头，缺省匿名/默认会话。
	task.Meta[contract.MetaUserID] = userID(c)
	task.Meta[contract.MetaSessionID] = sessionID(c)
	// 输入护栏灰区升级：Tier1（中间件）判为疑似的输入打标，Dispatcher 执行前调 LLM 分类器裁决。
	if c.GetBool("guardrail_suspect") {
		task.Meta[contract.MetaSafetyCheck] = true
	}
	// 持久化任务提交（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
	}
	// 从提交即开始把 token 流 + 执行轨迹录进 Redis Stream（订阅早于 dispatcher 产出）→
	// SSE 可从中回放/断点续传，根治"连晚/重连丢 token / 丢轨迹事件"。
	h.startTokenRecorder(task.ID)
	h.startExecRecorder(task.ID)
	c.JSON(http.StatusAccepted, gin.H{"task_id": task.ID})
}

// startTokenRecorder 后台订阅 token 流并落 Redis Stream，与 SSE 客户端是否在线无关。
// Redis 降级时为空操作（SSE 自动回退到 live NATS 路径）。
func (h *Handler) startTokenRecorder(taskID string) {
	if !h.cache.Enabled() {
		return
	}
	ctx, cancel := context.WithTimeout(context.Background(), 6*time.Minute) // 兜底防泄漏
	unsub, err := h.bus.SubscribeTokens(taskID,
		func(tok []byte) { _ = h.cache.StreamAppend(ctx, store.ChannelToken, taskID, "token", string(tok)) },
		func() { _ = h.cache.StreamAppend(ctx, store.ChannelToken, taskID, "done", ""); cancel() },
	)
	if err != nil {
		cancel()
		return
	}
	go func() { <-ctx.Done(); _ = unsub() }()
}

// startExecRecorder 后台订阅执行轨迹事件并落 Redis Stream（与 SSE 是否在线无关），
// 使"运行·观测"轨迹可回放/断点续传——连晚或刷新重连不再丢节点点亮/工具调用事件。
// Redis 降级时为空操作（StreamExec 自动回退到 live NATS）。
func (h *Handler) startExecRecorder(taskID string) {
	if !h.cache.Enabled() {
		return
	}
	ctx, cancel := context.WithTimeout(context.Background(), 12*time.Minute) // 含 HITL 审批等待，给足
	unsub, err := h.bus.SubscribeExec(taskID,
		func(data []byte) { _ = h.cache.StreamAppend(ctx, store.ChannelExec, taskID, "exec", string(data)) },
		func() { _ = h.cache.StreamAppend(ctx, store.ChannelExec, taskID, "done", ""); cancel() },
	)
	if err != nil {
		cancel()
		return
	}
	go func() { <-ctx.Done(); _ = unsub() }()
}

// TaskStatus: GET /api/v1/tasks/:id —— 返回任务生命周期状态（供 UI 轮询，根治"卡运行中看不出来"）。
func (h *Handler) TaskStatus(c *gin.Context) {
	id := c.Param("id")
	status, detail := h.db.GetTaskStatus(c.Request.Context(), id)
	if status == "" {
		c.JSON(http.StatusNotFound, gin.H{"error": "task not found"})
		return
	}
	c.JSON(http.StatusOK, gin.H{"task_id": id, "status": status, "detail": detail})
}

// TaskEval: GET /api/v1/tasks/:id/eval —— 取一次任务的自动化评测结果（综合/质量/忠实度/分级/flags）。
func (h *Handler) TaskEval(c *gin.Context) {
	e := h.db.GetEval(c.Request.Context(), c.Param("id"))
	if e == nil {
		c.JSON(http.StatusNotFound, gin.H{"error": "尚无评测结果（任务未完成或评测进行中）"})
		return
	}
	var flags []string
	_ = json.Unmarshal([]byte(e.Flags), &flags)
	c.JSON(http.StatusOK, gin.H{
		"task_id": e.TaskID, "overall": e.Overall, "rule": e.Rule, "llm": e.LLM,
		"faithful": e.Faithful, "level": e.Level, "flags": flags, "reason": e.Reason,
		"sources": e.Sources, "corrected": e.Corrected,
	})
}

// ApproveTask: POST /api/v1/tasks/:id/approve {approved, node?, note?} —— 人工审批决定（HITL）。
// 把决定经 NATS 发给 dispatcher，解除审批节点的阻塞（批准放行 / 拒绝中止）。
func (h *Handler) ApproveTask(c *gin.Context) {
	id := c.Param("id")
	var body struct {
		Approved bool   `json:"approved"`
		Node     string `json:"node"`
		Note     string `json:"note"`
	}
	if err := c.ShouldBindJSON(&body); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}
	// 仅在任务确为等待审批时受理（幂等：重复/迟到的决定不报错，dispatcher 侧已只取首条）。
	if status, _ := h.db.GetTaskStatus(c.Request.Context(), id); status != contract.TaskWaiting {
		c.JSON(http.StatusConflict, gin.H{"error": "任务当前非待审批状态", "status": status})
		return
	}
	if err := h.bus.PublishApproval(&contract.ApprovalDecision{
		TaskID: id, Node: body.Node, Approved: body.Approved, Note: body.Note,
		By: userID(c), TS: time.Now().UnixMilli(),
	}); err != nil {
		c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()})
		return
	}
	c.JSON(http.StatusOK, gin.H{"task_id": id, "approved": body.Approved})
}

// StreamTask: 以 SSE 把 Token Stream 推给客户端。
// 优先从 Redis Stream 读（可回放 + 断点续传，根治连晚/重连丢 token）；Redis 降级时回退 live NATS。
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")

	if !h.cache.Enabled() {
		h.streamTaskLive(c, taskID) // Redis 不可用 → 旧的 live-NATS 路径兜底
		return
	}

	// 断点续传：浏览器 EventSource 重连会带 Last-Event-ID；缺省 "0" 从头回放。
	lastID := c.GetHeader("Last-Event-ID")
	if lastID == "" {
		lastID = "0"
	}
	ctx := c.Request.Context()
	c.Stream(func(w io.Writer) bool {
		entries, nl, err := h.cache.StreamRead(ctx, store.ChannelToken, taskID, lastID, 20*time.Second)
		if err != nil {
			return false // ctx 取消（客户端断开）或后端故障
		}
		lastID = nl
		for _, e := range entries {
			if e.Kind == "done" {
				_ = sse.Encode(w, sse.Event{Id: e.ID, Event: "done", Data: taskID})
				return false
			}
			_ = sse.Encode(w, sse.Event{Id: e.ID, Event: "token", Data: e.Data})
		}
		return true // 继续阻塞读取后续 token
	})
}

// streamTaskLive 是 Redis 降级时的兜底：直接订阅 NATS token 流转 SSE（无回放/续传能力）。
func (h *Handler) streamTaskLive(c *gin.Context, taskID string) {
	tokens := make(chan []byte, 256)
	done := make(chan struct{})
	unsub, err := h.bus.SubscribeTokens(taskID,
		func(tok []byte) {
			select {
			case tokens <- tok:
			default:
			}
		},
		func() { close(done) },
	)
	if err != nil {
		c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()})
		return
	}
	defer func() { _ = unsub() }()
	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
		}
	})
}

// Healthz: GET /healthz —— 存活探针（liveness）：进程能应答即 200，不查依赖。
func (h *Handler) Healthz(c *gin.Context) {
	c.JSON(http.StatusOK, gin.H{"status": "ok"})
}

// Readyz: GET /readyz —— 就绪探针（readiness）：核心依赖(DB/Redis)可用才 200，否则 503。
// 供 k8s 等编排器在依赖未就绪时暂不导流。NATS 在启动时即连（连不上会 fatal），故不单列。
func (h *Handler) Readyz(c *gin.Context) {
	deps := gin.H{"db": h.db.Enabled(), "redis": h.cache.Enabled()}
	if h.db.Enabled() && h.cache.Enabled() {
		c.JSON(http.StatusOK, gin.H{"status": "ready", "deps": deps})
		return
	}
	c.JSON(http.StatusServiceUnavailable, gin.H{"status": "not_ready", "deps": deps})
}

// Health: GET /api/v1/health —— 聚合各依赖子系统健康，供桌面端顶栏五盏灯实时点亮。
// gateway/db/redis/nats 网关本地可判；milvus/neo4j 经 mcp-go health 工具取（不可用则置否）。
func (h *Handler) Health(c *gin.Context) {
	status := gin.H{
		"gateway": true,              // 能应答即在线
		"nats":    true,              // 网关启动即连上 NATS（连不上会 fatal）
		"db":      h.db.Enabled(),    // Postgres
		"redis":   h.cache.Enabled(), // Redis
		"milvus":  false,
		"neo4j":   false,
	}
	cctx, cancel := context.WithTimeout(c.Request.Context(), 2*time.Second)
	defer cancel()
	if res, err := h.bus.CallTool(cctx, contract.ToolSubjectGo("health"),
		&contract.ToolCall{Tool: "health"}); err == nil && res != nil && res.OK {
		var sub map[string]bool
		if json.Unmarshal([]byte(res.Content), &sub) == nil {
			status["milvus"] = sub["milvus"]
			status["neo4j"] = sub["neo4j"]
		}
	}
	c.JSON(http.StatusOK, status)
}

// StreamExec: 以 SSE 把执行轨迹事件推给客户端（运行·观测）。
// 与 StreamTask（token 流）并行：前端同时连两路，token 走输出、exec 走轨迹/工具面板。
// 优先从 Redis Stream 读（可回放 + 断点续传，根治连晚/刷新重连丢轨迹事件）；Redis 降级时回退 live NATS。
func (h *Handler) StreamExec(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")

	if !h.cache.Enabled() {
		h.streamExecLive(c, taskID) // Redis 不可用 → live-NATS 兜底
		return
	}

	// 断点续传：EventSource 重连带 Last-Event-ID；缺省 "0" 从头回放（连晚也能补齐已发生的轨迹）。
	lastID := c.GetHeader("Last-Event-ID")
	if lastID == "" {
		lastID = "0"
	}
	ctx := c.Request.Context()
	c.Stream(func(w io.Writer) bool {
		entries, nl, err := h.cache.StreamRead(ctx, store.ChannelExec, taskID, lastID, 20*time.Second)
		if err != nil {
			return false
		}
		lastID = nl
		for _, e := range entries {
			if e.Kind == "done" {
				_ = sse.Encode(w, sse.Event{Id: e.ID, Event: "done", Data: taskID})
				return false
			}
			_ = sse.Encode(w, sse.Event{Id: e.ID, Event: "exec", Data: e.Data})
		}
		return true
	})
}

// streamExecLive 是 Redis 降级时的兜底：直接订阅 NATS 轨迹流转 SSE（无回放/续传能力）。
func (h *Handler) streamExecLive(c *gin.Context, taskID string) {
	events := make(chan []byte, 256)
	done := make(chan struct{})
	unsub, err := h.bus.SubscribeExec(taskID,
		func(ev []byte) {
			select {
			case events <- ev:
			default: // 背压保护：客户端过慢则丢弃，避免阻塞 NATS 回调
			}
		},
		func() { close(done) },
	)
	if err != nil {
		c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()})
		return
	}
	defer func() { _ = unsub() }()

	c.Stream(func(w io.Writer) bool {
		select {
		case ev := <-events:
			c.SSEvent("exec", string(ev))
			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})
}

// ListMemory: GET /api/v1/memory —— 列出当前用户的全部偏好（结构化，供记忆面板）。
func (h *Handler) ListMemory(c *gin.Context) {
	res, err := h.bus.CallTool(c.Request.Context(), contract.ToolSubjectGo("memory_list"),
		&contract.ToolCall{Tool: "memory_list", Args: map[string]any{"user_id": userID(c)}})
	if err != nil || res == nil || !res.OK {
		msg := "记忆列举失败"
		if res != nil && res.Error != "" {
			msg = res.Error
		}
		c.JSON(http.StatusBadGateway, gin.H{"error": msg})
		return
	}
	var items []map[string]any
	_ = json.Unmarshal([]byte(res.Content), &items)
	c.JSON(http.StatusOK, gin.H{"memories": items})
}

// DeleteMemory: DELETE /api/v1/memory?key= —— 软删当前用户的一条偏好。
func (h *Handler) DeleteMemory(c *gin.Context) {
	key := c.Query("key")
	if key == "" {
		c.JSON(http.StatusBadRequest, gin.H{"error": "key required"})
		return
	}
	res, err := h.bus.CallTool(c.Request.Context(), contract.ToolSubjectGo("memory_delete"),
		&contract.ToolCall{Tool: "memory_delete", Args: map[string]any{"user_id": userID(c), "key": key}})
	if err != nil || res == nil || !res.OK {
		c.JSON(http.StatusBadGateway, gin.H{"error": "删除失败"})
		return
	}
	c.JSON(http.StatusOK, gin.H{"status": "ok"})
}

// userID 取当前用户标识 —— 仅信任 JWT 鉴权中间件注入的已验证 uid（不再认 header）。
// 受保护路由有 RequireAuth 兜底，此处理论上不会返回 anonymous。
func userID(c *gin.Context) string {
	if v, ok := c.Get("uid"); ok {
		if s, _ := v.(string); s != "" {
			return s
		}
	}
	return "anonymous"
}

// sessionID 从请求取会话标识（真实场景应由会话中间件注入）。
func sessionID(c *gin.Context) string {
	if s := c.GetHeader("X-Session-ID"); s != "" {
		return s
	}
	return "default"
}

// userDailyTokenBudget 读单用户当日 token 预算（env USER_DAILY_TOKEN_BUDGET，缺省/非法=0 即不限）。
func userDailyTokenBudget() int {
	if v := os.Getenv("USER_DAILY_TOKEN_BUDGET"); v != "" {
		if n, err := strconv.Atoi(v); err == nil && n > 0 {
			return n
		}
	}
	return 0
}

func (h *Handler) Billing(c *gin.Context) {
	n, err := h.db.CountTasks(c.Request.Context())
	if err != nil {
		c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()})
		return
	}
	// 成本护栏：当日 token 用量与预算（按用户隔离）。
	uid := userID(c)
	used := h.cache.GetUsage(c.Request.Context(), uid, time.Now().Format("20060102"))
	budget := userDailyTokenBudget()
	remaining := -1 // -1 表示不限额
	if budget > 0 {
		if r := budget - int(used); r > 0 {
			remaining = r
		} else {
			remaining = 0
		}
	}
	c.JSON(http.StatusOK, gin.H{
		"status": "ok", "tasks_submitted": n, "persisted": h.db.Enabled(),
		"token_used_today": used, "daily_budget": budget, "remaining": remaining,
	})
}