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 优先于地址：显式回复某封邮件时沿用该邮件的会话。
func resolveTarget(r *http.Request, addr models.Address, replyTo, fromAgent, subject, alias 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
		}
		id, err := repo.CreateSession(r.Context(), aliasPtr, fromAgent, subject)
		return id, nil, err

	case models.SessionDefault:
		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)
	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 quota repo.Quota
	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
		}
		// 仅读快照用于回传，不扣任何一层
		quota, _ = repo.GetQuota(r.Context(), agentName)
		budget, _ = repo.GetSessionBudget(r.Context(), sessionID)
	} else {
		// 两层都要过：会话预算管「这件事值得多少个来回」，
		// Agent 全局配额管「这个 Agent 总共能发多少」。
		// 先扣会话、后扣全局；全局拦下时把会话那次退回去 ——
		// 那次往返实际上没有发生，不能白掉一格。
		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
		}

		quota, err = repo.ConsumeQuota(r.Context(), agentName)
		if errors.Is(err, repo.ErrQuotaExhausted) {
			repo.RefundSessionBudget(r.Context(), sessionID)
			Error(w, http.StatusForbidden, fmt.Sprintf(
				"Agent 全局发信配额已用尽（%d/%d）。插件代劳转发的权限询问与最终总结不占配额；"+
					"若需继续主动发信请联系管理员重置配额。",
				quota.Used, quota.Max))
			return
		}
		if err != nil {
			repo.RefundSessionBudget(r.Context(), sessionID)
			Error(w, http.StatusInternalServerError, "Failed to check quota")
			return
		}
	}

	// 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),
	}
	if !quota.Unlimited {
		resp["quota_remaining"] = quota.Remaining
		resp["quota_used"] = quota.Used
		resp["quota_max"] = quota.Max
	}
	// 会话预算是【本任务】的剩余往返，Agent 更应该看这个而不是全局配额
	if !budget.Unlimited {
		resp["budget_remaining"] = budget.Remaining
		resp["budget_used"] = budget.Used
		resp["budget_max"] = budget.Max
	}
	if relay != "" {
		// 告知本次未扣配额，否则插件看到 quota_remaining 没变会以为数据错了
		resp["relay"] = relay
		resp["quota_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
	for i := range mails {
		fillAttachments(r, &mails[i])
	}
	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
}