MailUI4Agents/server/internal/repo/push_tokens.go

package repo

import (
	"context"
	"fmt"
	"strings"
	"time"

	"github.com/agentmail/gateway/internal/db"
	"github.com/google/uuid"
)

/*
设备推送 token 的读写（可选通道，见 internal/push）。

# 语义是「设备 + 注册者」

同一个 push token 只属于一个注册者：换人登录是**转移**，不是并存。这不是洁癖 ——
同一台设备上如果两个账号各存一份同一个 token，前一个人没退干净时，他的新邮件
会推到后一个人手里的那台设备上。

# 为什么没配推送也要能写

推送是自部署后端的可选项（用户 2026-09-15 明确要求）。所以登记 token 不依赖
「当前是否配了推送渠道」：端点照存，管理员之后把凭证配上就立刻生效，
不需要客户端重新登记一遍（客户端那时可能已经不在前台了）。
*/

// PushToken 是一台设备为某个注册者登记的推送地址。
type PushToken struct {
	TokenID    string
	Provider   string
	Token      string
	OwnerName  string
	SessionID  string
	DeviceName string
}

// UpsertPushToken 登记/刷新一台设备的推送地址。
//
// 同一个 (provider, token) 重复登记时**改归属**并刷新 session/device_name：
// 客户端每次启动都会登记一次，若这里报错或插重复行，表会随启动次数膨胀。
func UpsertPushToken(ctx context.Context, provider, token, owner, sessionID, deviceName string) error {
	_, err := db.DB.ExecContext(ctx, `
		INSERT INTO push_tokens (token_id, provider, token, owner_name, session_id, device_name, created_at, updated_at)
		VALUES ($1, $2, $3, $4, $5, $6, NOW(), NOW())
		ON CONFLICT (provider, token) DO UPDATE
		   SET owner_name = $4, session_id = $5, device_name = $6, updated_at = NOW()`,
		uuid.New().String(), provider, token, owner, sessionID, deviceName)
	return err
}

// DeletePushToken 注销一个 token。只允许**注册者本人**注销：不能凭一个 token
// 字符串把别人的设备推送掐掉。
//
// 返回是否真的删到了行 —— 客户端拿它区分「已注销」与「本来就没登记」，
// 但两者对客户端都不是错误（推送是可选通道，注销失败不该弹提示）。
func DeletePushToken(ctx context.Context, provider, token, owner string) (bool, error) {
	res, err := db.DB.ExecContext(ctx, `
		DELETE FROM push_tokens WHERE provider = $1 AND token = $2 AND owner_name = $3`,
		provider, token, owner)
	if err != nil {
		return false, err
	}
	n, err := res.RowsAffected()
	if err != nil {
		return false, err
	}
	return n > 0, nil
}

// ListPushTokensOf 列出某个注册者的全部推送地址（各 provider 都有）。
func ListPushTokensOf(ctx context.Context, owner string) ([]PushToken, error) {
	rows, err := db.DB.QueryContext(ctx, `
		SELECT token_id, provider, token, owner_name, session_id, device_name
		  FROM push_tokens WHERE owner_name = $1 ORDER BY provider, created_at`, owner)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	return scanPushTokens(rows)
}

// ListPushTokensOfOwners 一次取回多个注册者的推送地址。
//
// 为什么批量：一封邮件可能同时推给收件人 + 若干抄送方，逐个查库就是每个参与方
// 一次往返（`new_mail` 的分发已经因为同类原因重排过一次：会话级字段不许每人查一次）。
// 参与方为空时直接返回，不拼 `IN ()` 这种非法 SQL。
func ListPushTokensOfOwners(ctx context.Context, owners []string) ([]PushToken, error) {
	if len(owners) == 0 {
		return nil, nil
	}
	ph := make([]string, len(owners))
	args := make([]any, 0, len(owners))
	for i, o := range owners {
		ph[i] = fmt.Sprintf("$%d", i+1)
		args = append(args, o)
	}
	rows, err := db.DB.QueryContext(ctx, `
		SELECT token_id, provider, token, owner_name, session_id, device_name
		  FROM push_tokens WHERE owner_name IN (`+strings.Join(ph, ",")+`)
		 ORDER BY provider, created_at`, args...)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	return scanPushTokens(rows)
}

// PruneStalePushTokens 删掉 90 天没刷新过的登记。
//
// 客户端只在启动/换 token 时登记，所以「长期不刷新」就等于「这台设备不再用了」
// （App 卸载、token 轮换、换机）。留着它们的代价是每次发信都向一批死 token 发推送，
// 而华为对测试消息的额度是**项目级**的（1000 条/天，未上架应用），死 token 会白吃额度。
func PruneStalePushTokens(ctx context.Context) (int64, error) {
	// 截止时间在 Go 里算，不用 SQL 的日期运算：两种方言的写法不同
	// （PG 是 INTERVAL，SQLite 没有），而库里存的就是 NOW() 写的文本
	// "YYYY-MM-DD HH:MM:SS.ffffff"（见 db.go 的 now 注册），UTC 字符串比较即正确。
	cutoff := time.Now().UTC().Add(-90 * 24 * time.Hour).Format("2006-01-02 15:04:05.000000")
	res, err := db.DB.ExecContext(ctx,
		`DELETE FROM push_tokens WHERE updated_at < $1`, cutoff)
	if err != nil {
		return 0, err
	}
	return res.RowsAffected()
}

// DeletePushTokensByValue 按 token 值删除（不看注册者）。
//
// 只给**系统自愈**用：厂商回「这些 token 无效」（设备卸了 App、token 轮换）时，
// 留着它们的结果是每次发信都白吃额度（华为测试消息是**项目级** 1000 条/天），
// 而且用户那边永远收不到。用户主动注销走 DeletePushToken（要带 owner）。
func DeletePushTokensByValue(ctx context.Context, provider string, tokens []string) (int64, error) {
	if len(tokens) == 0 {
		return 0, nil
	}
	ph := make([]string, len(tokens))
	args := make([]any, 0, len(tokens)+1)
	args = append(args, provider)
	for i, t := range tokens {
		ph[i] = fmt.Sprintf("$%d", i+2)
		args = append(args, t)
	}
	res, err := db.DB.ExecContext(ctx,
		`DELETE FROM push_tokens WHERE provider = $1 AND token IN (`+strings.Join(ph, ",")+`)`, args...)
	if err != nil {
		return 0, err
	}
	return res.RowsAffected()
}

func scanPushTokens(rows interface {
	Next() bool
	Scan(...any) error
	Err() error
}) ([]PushToken, error) {
	var out []PushToken
	for rows.Next() {
		var t PushToken
		if err := rows.Scan(&t.TokenID, &t.Provider, &t.Token, &t.OwnerName, &t.SessionID, &t.DeviceName); err != nil {
			return nil, err
		}
		out = append(out, t)
	}
	return out, rows.Err()
}