MailUI4Agents/gateway/internal/repo/sessionrate.go

package repo

import (
	"errors"
	"sync"
	"time"
)

// ---------- 新建会话速率限制 ----------
//
// 会话往返预算（sessions.max_rounds）管住了「一条线索能来回多少次」，
// 但 Agent 仍可以用 `name@path.new` 开一串新会话，每条都是全新预算 ——
// 预算就被绕过了。
//
// 为什么用速率限制而不是「终身额度」：
// 终身额度跑满后要管理员手工重置才能再干活，而 Agent 是长期在线的 ——
// 那是把一次性资源的模型套在长期服务上。速率限制只压住「短时间内暴开」这个
// 真正的滥用形态，过一个窗口自动恢复，无需人工介入。
//
// 为什么不禁止 Agent 主动开新会话：那会堵死 Agent 之间的主动协作
//（A 发现问题主动找 B），而这正是这个平台存在的意义。

// ErrSessionRateLimited 表示该 Agent 短时间内新建会话过多。
var ErrSessionRateLimited = errors.New("session creation rate limited")

const (
	// sessionRateWindow 是滑动窗口长度
	sessionRateWindow = time.Hour
	// sessionRateLimit 是窗口内允许新建的会话数。
	//
	// 取 20：正常协作里 Agent 一小时开二十条新线索已经很多了；
	// 真到了这个量级，更可能是循环而不是在干活。
	sessionRateLimit = 20
)

// sessionRateLimiter 记录每个 Agent 新建会话的时间戳。
//
// 进程内内存计数，与登录限速（handler/ratelimit.go）同一取舍：
// 单实例部署下够用；多实例时各自计数，等效上限变成 N 倍 ——
// 那时应当换成数据库计数或 Redis。这个限制记在 README 的已知取舍里。
type sessionRateLimiter struct {
	mu    sync.Mutex
	marks map[string][]time.Time
}

var sessionLimiter = &sessionRateLimiter{marks: make(map[string][]time.Time)}

// Allow 判断是否允许新建，允许则记账。
//
// 判断与记账在同一把锁里：分开的话并发请求会双双通过检查，把上限刷穿 ——
// 与配额那条 UPDATE 同样的道理。
func (l *sessionRateLimiter) Allow(name string) (bool, int) {
	l.mu.Lock()
	defer l.mu.Unlock()

	now := time.Now()
	cutoff := now.Add(-sessionRateWindow)

	kept := l.marks[name][:0]
	for _, t := range l.marks[name] {
		if t.After(cutoff) {
			kept = append(kept, t)
		}
	}
	l.marks[name] = kept

	if len(kept) >= sessionRateLimit {
		// 最早那条何时过期 = 何时能再开一条
		retry := int(kept[0].Add(sessionRateWindow).Sub(now).Seconds()) + 1
		if retry < 1 {
			retry = 1
		}
		return false, retry
	}
	l.marks[name] = append(l.marks[name], now)
	return true, 0
}

// Release 撤销一次记账。
//
// 允许之后建会话失败时必须还回去，否则那次没发生的新建也占了名额。
func (l *sessionRateLimiter) Release(name string) {
	l.mu.Lock()
	defer l.mu.Unlock()
	m := l.marks[name]
	if len(m) > 0 {
		l.marks[name] = m[:len(m)-1]
	}
}

// AllowNewSession 供 handler 调用：Agent 新建会话前先过速率限制。
//
// 第二个返回值是建议的重试等待秒数（供 Retry-After 使用）。
// 人类用户不走这条路径 —— 人手工点「新建邮件」的频率天然受限，
// 给它加限制只会在批量派活时误伤。
func AllowNewSession(agentName string) (bool, int) {
	if agentName == "" {
		return true, 0
	}
	return sessionLimiter.Allow(agentName)
}

// ReleaseNewSession 建会话失败后归还名额。
func ReleaseNewSession(agentName string) {
	if agentName == "" {
		return
	}
	sessionLimiter.Release(agentName)
}

// 定期清理空闲 Agent 的记录，避免 map 随 Agent 名无限增长。
func init() {
	go func() {
		t := time.NewTicker(sessionRateWindow)
		defer t.Stop()
		for range t.C {
			sessionLimiter.mu.Lock()
			cutoff := time.Now().Add(-sessionRateWindow)
			for name, marks := range sessionLimiter.marks {
				fresh := marks[:0]
				for _, ts := range marks {
					if ts.After(cutoff) {
						fresh = append(fresh, ts)
					}
				}
				if len(fresh) == 0 {
					delete(sessionLimiter.marks, name)
				} else {
					sessionLimiter.marks[name] = fresh
				}
			}
			sessionLimiter.mu.Unlock()
		}
	}()
}

// SessionRateLimit 暴露窗口内的新建上限，供错误文案使用。
// 不导出常量本身：外部只该读它，不该改它。
func SessionRateLimit() int { return sessionRateLimit }