MailUI4Agents/gateway/internal/handler/mail.go

package handler

import (
	"errors"
	"fmt"
	"net/http"
	"strings"

	"github.com/agentmail/gateway/internal/middleware"
	"github.com/agentmail/gateway/internal/models"
	"github.com/agentmail/gateway/internal/repo"
	"github.com/agentmail/gateway/internal/sse"
	"github.com/google/uuid"
)

// ---------- Mail ----------

type sendMailRequest struct {
	To      string `json:"to"`      // name@path.session（省略 session=默认会话，new=新建，别名=必须已存在）
	CC      string `json:"cc"`      // 逗号/分号/空格分隔的多个 name@path.session
	Subject string `json:"subject"`
	Body    string `json:"body"`
	ReplyTo string `json:"reply_to"`
	// SessionAlias 仅在本次投递【新建】会话时生效，为新会话命名，
	// 之后即可用 name@path.<alias> 续谈。命中已有会话时该字段被忽略。
	SessionAlias string `json:"session_alias"`
	// AttachmentIDs 先用 POST /attachments 上传拿到的 id；只能附加自己上传且未挂载的
	AttachmentIDs []string `json:"attachment_ids"`

	// Relay 标识本次发信是【插件代劳转发】而不是模型自主发信。
	//
	// 基本原则：**配额约束的是模型的自主发信，不是 harness 的转发**。
	// 平台原生的权限询问与本轮的最终总结都是插件搬运的，不计配额。
	//
	// RelayKey 必須是上游那条消息的稳定标识（permission id / assistant message id）：
	// 它由平台生成，模型伪造不出，而唯一约束保证同一条上游消息只能免费转一次。
	Relay    string `json:"relay"`     // "" | "permission" | "summary"
	RelayKey string `json:"relay_key"` // 上游消息 id；relay 非空时必填
}

// resolveTarget 根据三维地址 name@path.session 决定投递的会话。
//
// session 位三态语义（设计文档）：
//   - 省略（pi@root）      → 投递到 name@path 的默认会话；从未通信则建立
//   - new（pi@root.new）   → 强制新建一个会话
//   - 具体别名（pi@root.fix-leak）→ 必须已存在且该收件人参与过，否则 404 无法送达
//
// alias 为新建会话命名（仅新建时生效），使其之后可被 name@path.<alias> 寻址。
// reply_to 优先于地址：显式回复某封邮件时沿用该邮件的会话。
//
// byAgent 非空时表示这是 Agent 发起的投递，新建会话要过速率限制：
// 往返预算按会话计，Agent 用 .new 开一串会话就等于绕过预算。
// 人类不受此限（手工点「新建邮件」的频率天然受限，加限制只会在批量派活时误伤）。
func resolveTarget(r *http.Request, addr models.Address, replyTo, fromAgent, subject, alias string, byAgent string) (uuid.UUID, *uuid.UUID, error) {
	if replyTo != "" {
		replyID, err := uuid.Parse(replyTo)
		if err != nil {
			return uuid.Nil, nil, errBadRequest("Invalid reply_to UUID")
		}
		mail, err := repo.GetMailByID(r.Context(), replyID)
		if err != nil {
			return uuid.Nil, nil, errNotFound("Parent mail not found")
		}
		repo.TouchSession(r.Context(), mail.SessionID)
		return mail.SessionID, &replyID, nil
	}

	switch addr.Mode() {
	case models.SessionNew:
		// 新建会话：若调用方给了别名，当场命名，之后即可用 name@path.<alias> 续谈。
		// 别名全局唯一（负责寻址），已被占用时报 409 而不是静默吐出重名会话。
		var aliasPtr *string
		if a := strings.TrimSpace(alias); a != "" {
			if err := validateSessionAlias(a); err != nil {
				return uuid.Nil, nil, err
			}
			if _, err := repo.FindSessionByAlias(r.Context(), a); err == nil {
				return uuid.Nil, nil, errConflict(fmt.Sprintf(
					"会话别名 %q 已被占用；若要接着该会话谈请用 %s@%s.%s", a, addr.Name, addr.Path, a))
			}
			aliasPtr = &a
		}
		// Agent 主动开新线索要过速率限制
		if ok, retry := repo.AllowNewSession(byAgent); !ok {
			return uuid.Nil, nil, errRateLimited(fmt.Sprintf(
				"新建会话过于频繁（1 小时内已开 %d 条）。请在已有会话里继续，或 %d 秒后再试。",
				repo.SessionRateLimit(), retry))
		}
		id, err := repo.CreateSession(r.Context(), aliasPtr, fromAgent, subject)
		if err != nil {
			// 建失败要把名额还回去：那次新建实际上没有发生
			repo.ReleaseNewSession(byAgent)
		}
		return id, nil, err

	case models.SessionDefault:
		// 默认会话「从未通信则建立」也会产生新会话，但一个 name@path 只有一条，
		// 不构成暴开的手段，因此不计入速率限制。
		id, err := repo.FindOrCreateDefaultSession(r.Context(), addr.Name, addr.Path, fromAgent, subject)
		return id, nil, err

	default: // models.SessionNamed
		id, err := repo.FindNamedSessionFor(r.Context(), addr.Name, addr.Path, addr.Session)
		if errors.Is(err, repo.ErrSessionNotFound) {
			return uuid.Nil, nil, errNotFound(fmt.Sprintf(
				"无法送达：会话 %q 不存在于 %s@%s。若要新建会话请用 %s@%s.new，投递默认会话请省略 session 位",
				addr.Session, addr.Name, addr.Path, addr.Name, addr.Path))
		}
		if err != nil {
			return uuid.Nil, nil, err
		}
		repo.TouchSession(r.Context(), id)
		return id, nil, nil
	}
}

// POST /api/v1/mail/send
func SendMail(w http.ResponseWriter, r *http.Request) {
	agentName := middleware.GetAgentName(r)
	if agentName == "" {
		Error(w, http.StatusUnauthorized, "Unauthorized")
		return
	}

	var req sendMailRequest
	if err := Decode(r, &req); err != nil {
		Error(w, http.StatusBadRequest, "Invalid JSON")
		return
	}
	if req.To == "" || req.Subject == "" || req.Body == "" {
		Error(w, http.StatusBadRequest, "Missing to, subject, or body")
		return
	}

	to, err := models.ParseAddress(req.To)
	if err != nil {
		Error(w, http.StatusBadRequest, "Invalid to address: "+err.Error())
		return
	}
	ccList, err := models.ParseAddressList(req.CC)
	if err != nil {
		Error(w, http.StatusBadRequest, "Invalid cc address: "+err.Error())
		return
	}
	attachIDs, err := parseAttachmentIDs(req.AttachmentIDs)
	if err != nil {
		Error(w, http.StatusBadRequest, err.Error())
		return
	}

	sessionID, parentMailID, err := resolveTarget(r, to, req.ReplyTo, agentName, req.Subject, req.SessionAlias, agentName)
	if err != nil {
		writeErr(w, err, "Failed to resolve session")
		return
	}

	// 配额在建邮件之前扣：否则邮件已入库再报 403，收件方会看到一封发件方以为发失败的邮件。
	// 只限制主动发信，不限制收信（卡住收信只会让邮件凭空消失）。
	//
	// 插件代劳转发（relay）走免配额通道：配额约束的是模型的自主发信，
	// 不是 harness 把平台原生的权限询问与最终总结搬到邮件里。
	relay, relayKey, err := parseRelay(req.Relay, req.RelayKey)
	if err != nil {
		writeErr(w, err, "Invalid relay")
		return
	}

	var budget repo.SessionBudget
	if relay != "" {
		// 先占幂等键。重复则说明这条上游消息已经转过，
		// 这是插件重试 / SSE 重放的正常结果，不是故障 —— 幂等地返回成功。
		if cErr := repo.ClaimRelay(r.Context(), agentName, relayKey, relay); cErr != nil {
			if errors.Is(cErr, repo.ErrRelayDuplicate) {
				JSON(w, http.StatusOK, map[string]any{
					"status":    "duplicate_relay",
					"relay":     relay,
					"relay_key": relayKey,
					"detail":    "该上游消息已转发过，本次调用未产生新邮件",
				})
				return
			}
			Error(w, http.StatusInternalServerError, "Failed to claim relay")
			return
		}
		// 仅读快照用于回传，不扣预算
		budget, _ = repo.GetSessionBudget(r.Context(), sessionID)
	} else {
		// 额度只看【本任务】的往返预算。
		//
		// 不再叠一层 Agent 终身额度：那种额度跑满后要管理员手工重置才能再干活，
		// 而 Agent 是长期在线的。防止 Agent 用 .new 开一串新会话绕过预算，
		// 靠的是新建会话速率限制（resolveTarget 里）。
		budget, err = repo.ConsumeSessionBudget(r.Context(), sessionID)
		if errors.Is(err, repo.ErrSessionBudgetExhausted) {
			Error(w, http.StatusForbidden, fmt.Sprintf(
				"本任务的往返预算已用尽（%d/%d）。自动转发的总结与权限询问不占预算；"+
					"若需继续主动发信，请让人在对话页调高本任务的预算。",
				budget.Used, budget.Max))
			return
		}
		if err != nil {
			Error(w, http.StatusInternalServerError, "Failed to check session budget")
			return
		}
		// 纯统计，不拦请求；写失败也不该让邮件发不出去
		repo.BumpSentCount(r.Context(), agentName)
	}

	// Agent 可以在正文里提议改会话别名（<!-- agentmail:rename-session … -->）。
	// 标记从入库正文里剥掉：它是给系统看的元数据，不该出现在人读的正文里
	// （react-markdown 会把 HTML 注释转义成可见文本，不会自动吞掉）。
	//
	// 提议只是提议 —— 别名是人的寻址入口，Agent 干到一半自己改掉会让人
	// 上一秒记住的地址下一秒失效。真正改名要等用户在前端点「接受」。
	proposal, body := extractRenameProposal(req.Body)

	mailID, err := repo.CreateMail(r.Context(), sessionID, parentMailID,
		agentName, agentName, to.Name, to.Path, req.Subject, body, ccList)
	if err != nil {
		// 建邮件失败时必须把幂等键还回去，否则这条上游消息永远转不出来了
		if relay != "" {
			_ = repo.ReleaseRelay(r.Context(), agentName, relayKey)
		}
		Error(w, http.StatusInternalServerError, "Failed to create mail")
		return
	}
	if relay != "" {
		// 关联失败不影响功能，只是少一条审计记录
		_ = repo.BindRelayMail(r.Context(), agentName, relayKey, mailID)
	}
	if proposal != nil {
		// 记不上提议不该让发信失败：邮件本身已经入库，提议是旁支信息
		_ = repo.SetMailRenameProposal(r.Context(), mailID, proposal.Alias, proposal.Reason)
	}

	if !attachAll(w, r, mailID, attachIDs, agentName) {
		return
	}

	notifyRecipients(to, ccList, sessionID, mailID, agentName, req.Subject)

	// 回传会话别名与本任务剩余往返，让发件方知道后续用什么地址续谈、还能发几封
	resp := map[string]any{
		"mail_id":       mailID.String(),
		"session_id":    sessionID.String(),
		"session_alias": repo.SessionAliasOf(r.Context(), sessionID),
	}
	// 预算属于【本任务】，不限时不回传 —— 多给一个 -1 只会让插件去判断哪个值是哨兵
	if !budget.Unlimited {
		resp["budget_remaining"] = budget.Remaining
		resp["budget_used"] = budget.Used
		resp["budget_max"] = budget.Max
	}
	if relay != "" {
		// 告知本次未扣预算，否则插件看到 budget_remaining 没变会以为数据错了
		resp["relay"] = relay
		resp["budget_charged"] = false
	}
	if proposal != nil {
		// 回传规范化后的别名：Agent 提的名字可能含非法字符被改写过，
		// 让它知道最终会拿什么去问用户
		resp["rename_proposed"] = proposal.Alias
	}
	JSON(w, http.StatusOK, resp)
}

// notifyRecipients 向主收件人与抄送方推送 new_mail，并刷新相关方的会话列表。
// 收件人可能是 Agent 也可能是人类用户（三维地址 name 位共享命名空间），
// 因此统一用 SendToRecipient 同时试 Agent 通道与用户通道。
func notifyRecipients(to models.Address, cc []models.Address, sessionID, mailID uuid.UUID, from, subject string) {
	payload := func(role string) map[string]interface{} {
		return map[string]interface{}{
			"mail_id":    mailID.String(),
			"session_id": sessionID.String(),
			"from_name":  from,
			"subject":    subject,
			"mail_type":  "normal",
			"role":       role, // to / cc
		}
	}

	update := map[string]interface{}{
		"session_id": sessionID.String(),
		"status":     "active",
	}

	// 参与方去重：收件人 + 所有抄送 + 发件人自己（刷新他的发件箱）
	seen := map[string]bool{}

	sse.Default.SendToRecipient(to.Name, "new_mail", payload("to"))
	sse.Default.SendToRecipient(to.Name, "session_update", update)
	seen[to.Name] = true

	for _, c := range cc {
		if seen[c.Name] {
			continue
		}
		seen[c.Name] = true
		sse.Default.SendToRecipient(c.Name, "new_mail", payload("cc"))
		sse.Default.SendToRecipient(c.Name, "session_update", update)
	}

	if !seen[from] {
		sse.Default.SendToRecipient(from, "session_update", update)
	}
}

// GET /api/v1/mail/inbox
func GetInbox(w http.ResponseWriter, r *http.Request) {
	agentName := middleware.GetAgentName(r)
	if agentName == "" {
		Error(w, http.StatusUnauthorized, "Unauthorized")
		return
	}

	status := r.URL.Query().Get("status")
	if status == "" {
		status = "unread"
	}
	limit := 10
	if l := r.URL.Query().Get("limit"); l != "" {
		if n, err := parseInt(l); err == nil && n > 0 {
			limit = n
		}
	}

	mails, err := repo.ListInbox(r.Context(), agentName, status, limit)
	if err != nil {
		Error(w, http.StatusInternalServerError, "Failed to list inbox")
		return
	}
	// Agent 靠收件箱列表得知有哪些附件可下载，否则它不知道该调 attachment_id
	ptrs := make([]*models.Mail, len(mails))
	for i := range mails {
		ptrs[i] = &mails[i]
	}
	fillAttachments(r, ptrs...)
	total, _ := repo.CountUnread(r.Context(), agentName)

	JSON(w, http.StatusOK, map[string]interface{}{
		"mails": emptySlice(mails),
		"total": total,
	})
}

// GET /api/v1/mail/{id} —— 需登录，且需对所属会话有权限
func GetMail(w http.ResponseWriter, r *http.Request) {
	user := middleware.GetUser(r)
	if user == nil {
		Error(w, http.StatusUnauthorized, "not authenticated")
		return
	}
	mailID, ok := pathUUID(w, r, "id")
	if !ok {
		return
	}

	mail, err := repo.GetMailByID(r.Context(), mailID)
	if err != nil {
		Error(w, http.StatusNotFound, "Mail not found")
		return
	}

	allowed, err := repo.UserCanAccessSession(r.Context(), user, mail.SessionID)
	if err != nil {
		Error(w, http.StatusInternalServerError, "Failed to check permission")
		return
	}
	if !allowed {
		Error(w, http.StatusForbidden, "无权访问该邮件")
		return
	}
	fillAttachments(r, mail)
	JSON(w, http.StatusOK, mail)
}

// POST /api/v1/mail/{id}/read —— 需登录，只能标记自己可见的邮件
func MarkMailRead(w http.ResponseWriter, r *http.Request) {
	user := middleware.GetUser(r)
	if user == nil {
		Error(w, http.StatusUnauthorized, "not authenticated")
		return
	}
	mailID, ok := pathUUID(w, r, "id")
	if !ok {
		return
	}

	mail, err := repo.GetMailByID(r.Context(), mailID)
	if err != nil {
		Error(w, http.StatusNotFound, "Mail not found")
		return
	}
	allowed, err := repo.UserCanAccessSession(r.Context(), user, mail.SessionID)
	if err != nil {
		Error(w, http.StatusInternalServerError, "Failed to check permission")
		return
	}
	if !allowed {
		Error(w, http.StatusForbidden, "无权操作该邮件")
		return
	}

	if err := repo.MarkMailRead(r.Context(), mailID); err != nil {
		Error(w, http.StatusInternalServerError, "Failed to mark read")
		return
	}
	JSON(w, http.StatusOK, map[string]string{"status": "read"})
}

func parseInt(s string) (int, error) {
	n := 0
	for _, c := range s {
		if c < '0' || c > '9' {
			return 0, nil
		}
		n = n*10 + int(c-'0')
	}
	return n, nil
}

type markReadRequest struct {
	// MailIDs 要标记为已读的邮件；省略/为空 = 把收件箱里全部未读标掉。
	MailIDs []string `json:"mail_ids"`
}

// POST /api/v1/mail/read —— Agent 侧批量标记已读
//
// 为什么需要它：Agent 读完 read_inbox 后没有任何办法把邮件标掉，
// 于是每次拉收件箱都把同一批旧邮件重新捞出来 —— 处理过的信和新来的信混在一起，
// 模型分不清哪封该回。心跳里的未读数也永远只增不减。
//
// 鉴权写进 UPDATE 的 WHERE 而不是先查后改：不是发给自己的邮件根本改不动，
// 既省一次查询，也没有「查完到改之间邮件被转走」的时间窗。
func MarkInboxRead(w http.ResponseWriter, r *http.Request) {
	agentName := middleware.GetAgentName(r)
	if agentName == "" {
		Error(w, http.StatusUnauthorized, "Unauthorized")
		return
	}

	var req markReadRequest
	// 允许空 body：`POST /mail/read` 不带任何内容 = 全部标掉
	if r.ContentLength > 0 {
		if err := Decode(r, &req); err != nil {
			Error(w, http.StatusBadRequest, "Invalid JSON")
			return
		}
	}

	// 不给 id 就把收件箱里全部未读标掉。
	// 这是 Agent 最常见的用法：一轮处理完，剩下的都不必再看。
	if len(req.MailIDs) == 0 {
		n, err := repo.MarkAllInboxReadFor(r.Context(), agentName)
		if err != nil {
			Error(w, http.StatusInternalServerError, "Failed to mark read")
			return
		}
		JSON(w, http.StatusOK, map[string]any{"status": "read", "marked": n, "scope": "all"})
		return
	}

	const maxBatch = 200
	if len(req.MailIDs) > maxBatch {
		Error(w, http.StatusBadRequest, fmt.Sprintf("一次最多标记 %d 封", maxBatch))
		return
	}
	ids := make([]uuid.UUID, 0, len(req.MailIDs))
	for _, s := range req.MailIDs {
		id, err := uuid.Parse(strings.TrimSpace(s))
		if err != nil {
			Error(w, http.StatusBadRequest, "非法的 mail_id: "+s)
			return
		}
		ids = append(ids, id)
	}

	n, err := repo.MarkMailsReadFor(r.Context(), agentName, ids)
	if err != nil {
		Error(w, http.StatusInternalServerError, "Failed to mark read")
		return
	}
	// 不因为「有些 id 不是发给你的」而报错：那些 id 只是没被标掉。
	// 报错会让整批失败，而 Agent 通常是把上一轮列出的 id 原样传回来，
	// 其中可能混着已读的（幂等）——那不该是错误。
	JSON(w, http.StatusOK, map[string]any{
		"status":    "read",
		"marked":    n,
		"requested": len(ids),
	})
}