feat: AgentMail —— 以邮件为统一范式的多智能体协作平台
Go 单二进制网关 + React 前端 + opencode 桥接插件。部署产物是 「一个二进制加一个 .db 文件」：前端经 go:embed 打进二进制， 数据库默认内置 SQLite，systemd 托管。 核心设计 - 三维寻址 name@path.session，按最后一个 . 切分；session 位三态： 省略=默认会话 / new=强制新建 / 具体别名=必须已存在（否则 404 无法送达） - 会话别名默认复用 Agent 平台自己的命名机制（opencode 的 slug 与模型生成的 标题），不在本侧另造一套；人显式定过的别名不被平台同步覆盖 - 对话树不建 tree_nodes 表：parent_mail_id 已完整编码树结构， 再维护一张表就是第二份真相。用递归 CTE 查，按方向分块加载 - 附件内容存磁盘、按 sha256 内容寻址，数据库只存元数据；天然去重， 且路径与用户 filename 无关，杜绝 ../ 穿越 - 配额约束的是模型的自主发信，不是 harness 的转发：插件代劳的权限询问与 最终总结走免配额通道，靠上游消息 id 做幂等键而非计数 - 往返预算下沉到会话（写信时给、对话页里改）+ Agent 全局配额，两层都要过 后端 gateway/ - models/repo/handler/middleware/sse/blob 分层；两方言（SQLite/PostgreSQL） 共用一份 repo 层 SQL，差异集中在 internal/db - 多用户认证（bcrypt cost12、登录限速、会话隔离、权限边界） - 密钥体系：Agent 密钥与用户密钥分表，三种生命周期；登记式密钥让全文 只从客户端流向服务器一次 - 所有「判断 + 自增」都在同一条 UPDATE 里（配额、预算、one_time 密钥、 附件挂载），并发下不会刷穿 前端 web/ - 三栏布局、三段式地址补全、权限卡片、密钥面板、配额面板、对话树、附件 - 全站纯 SVG 图标，不使用 emoji - api/ 即可复用的客户端 SDK：基地址与凭证集中在 api/config.ts 插件 plugins/opencode-mail-bridge/ - 六个工具 + 两类自动转发（permission.ask 钩子接管平台原生权限询问、 session.idle 时转发本轮总结）
This commit is contained in:
226
gateway/internal/repo/attachments.go
Normal file
226
gateway/internal/repo/attachments.go
Normal file
@ -0,0 +1,226 @@
|
||||
package repo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"errors"
|
||||
"time"
|
||||
|
||||
"github.com/agentmail/gateway/internal/db"
|
||||
"github.com/agentmail/gateway/internal/models"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
// ---------- 附件 ----------
|
||||
//
|
||||
// 元数据在库、内容在磁盘（internal/blob）。两者的一致性由调用顺序保证：
|
||||
// 先落盘再入库 —— 反过来会出现「库里有记录但文件不存在」的下载 500。
|
||||
// 落盘成功但入库失败时最多留下一个无引用的文件，由 GC 回收，不影响正确性。
|
||||
|
||||
var (
|
||||
// ErrAttachmentNotFound 附件不存在
|
||||
ErrAttachmentNotFound = errors.New("attachment not found")
|
||||
// ErrAttachmentNotOwned 附件不属于该上传者
|
||||
ErrAttachmentNotOwned = errors.New("attachment not owned by uploader")
|
||||
// ErrAttachmentAlreadyAttached 附件已挂到别的邮件上
|
||||
ErrAttachmentAlreadyAttached = errors.New("attachment already attached")
|
||||
)
|
||||
|
||||
const attachmentCols = `attachment_id, mail_id, uploader, filename, content_type, size_bytes, sha256, created_at`
|
||||
|
||||
func scanAttachment(sc interface{ Scan(...any) error }) (*models.Attachment, error) {
|
||||
var a models.Attachment
|
||||
if err := sc.Scan(&a.ID, &a.MailID, &a.Uploader, &a.Filename,
|
||||
&a.ContentType, &a.SizeBytes, &a.SHA256, &a.CreatedAt); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &a, nil
|
||||
}
|
||||
|
||||
// CreateAttachment 登记一条待挂载的附件（mail_id 为空）。
|
||||
func CreateAttachment(ctx context.Context, uploader, filename, contentType string, size int64, sum string) (*models.Attachment, error) {
|
||||
a := &models.Attachment{
|
||||
Uploader: uploader,
|
||||
Filename: filename,
|
||||
ContentType: contentType,
|
||||
SizeBytes: size,
|
||||
SHA256: sum,
|
||||
}
|
||||
err := db.DB.QueryRowContext(ctx, `
|
||||
INSERT INTO attachments (uploader, filename, content_type, size_bytes, sha256)
|
||||
VALUES ($1, $2, $3, $4, $5)
|
||||
RETURNING attachment_id, created_at
|
||||
`, uploader, filename, contentType, size, sum).Scan(&a.ID, &a.CreatedAt)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return a, nil
|
||||
}
|
||||
|
||||
// GetAttachment 读取一条附件元数据。
|
||||
func GetAttachment(ctx context.Context, id uuid.UUID) (*models.Attachment, error) {
|
||||
a, err := scanAttachment(db.DB.QueryRowContext(ctx,
|
||||
`SELECT `+attachmentCols+` FROM attachments WHERE attachment_id = $1`, id))
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
return nil, ErrAttachmentNotFound
|
||||
}
|
||||
return a, err
|
||||
}
|
||||
|
||||
// ListAttachmentsFor 列出某封邮件的附件。
|
||||
func ListAttachmentsFor(ctx context.Context, mailID uuid.UUID) ([]models.Attachment, error) {
|
||||
rows, err := db.DB.QueryContext(ctx,
|
||||
`SELECT `+attachmentCols+` FROM attachments WHERE mail_id = $1 ORDER BY created_at`, mailID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
out := []models.Attachment{}
|
||||
for rows.Next() {
|
||||
a, err := scanAttachment(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, *a)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// AttachToMail 把一批待挂载附件绑到某封邮件上。
|
||||
//
|
||||
// 每条都要求：存在、属于该上传者、且尚未挂载。
|
||||
// 用 WHERE mail_id IS NULL AND uploader = ? 一条 UPDATE 完成判断与写入，
|
||||
// 避免「先查后改」在并发下把同一个附件挂到两封邮件上。
|
||||
func AttachToMail(ctx context.Context, mailID uuid.UUID, ids []uuid.UUID, uploader string) error {
|
||||
for _, id := range ids {
|
||||
tag, err := db.DB.ExecContext(ctx, `
|
||||
UPDATE attachments SET mail_id = $1
|
||||
WHERE attachment_id = $2 AND uploader = $3 AND mail_id IS NULL
|
||||
`, mailID, id, uploader)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if n, _ := tag.RowsAffected(); n > 0 {
|
||||
continue
|
||||
}
|
||||
|
||||
// 没改到：查明原因，给调用方一个能照着修的错误
|
||||
a, gErr := GetAttachment(ctx, id)
|
||||
if gErr != nil {
|
||||
return gErr
|
||||
}
|
||||
if a.Uploader != uploader {
|
||||
return ErrAttachmentNotOwned
|
||||
}
|
||||
return ErrAttachmentAlreadyAttached
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// CopyAttachmentsTo 把源邮件的附件复制到目标邮件（转发时用）。
|
||||
//
|
||||
// 内容寻址下「复制」只是新增一条指向同一 sha256 的元数据，不拷磁盘文件。
|
||||
// uploader 记为转发人：附件随新邮件重新分发，其可见范围由新邮件的参与方决定，
|
||||
// 而不是沿用原上传者。返回复制的数量。
|
||||
func CopyAttachmentsTo(ctx context.Context, srcMailID, dstMailID uuid.UUID, forwarder string) (int, error) {
|
||||
src, err := ListAttachmentsFor(ctx, srcMailID)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
for _, a := range src {
|
||||
_, err := db.DB.ExecContext(ctx, `
|
||||
INSERT INTO attachments (mail_id, uploader, filename, content_type, size_bytes, sha256)
|
||||
VALUES ($1, $2, $3, $4, $5, $6)
|
||||
`, dstMailID, forwarder, a.Filename, a.ContentType, a.SizeBytes, a.SHA256)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
}
|
||||
return len(src), nil
|
||||
}
|
||||
|
||||
// DeleteAttachment 删除一条附件元数据，返回它的 sha256 以及该内容是否已无人引用。
|
||||
// 内容寻址下多条记录可能共享同一个文件，只有最后一条引用消失才能删磁盘文件。
|
||||
func DeleteAttachment(ctx context.Context, id uuid.UUID) (sum string, orphaned bool, err error) {
|
||||
a, err := GetAttachment(ctx, id)
|
||||
if err != nil {
|
||||
return "", false, err
|
||||
}
|
||||
if _, err = db.DB.ExecContext(ctx,
|
||||
`DELETE FROM attachments WHERE attachment_id = $1`, id); err != nil {
|
||||
return "", false, err
|
||||
}
|
||||
|
||||
var refs int
|
||||
if err = db.DB.QueryRowContext(ctx,
|
||||
`SELECT COUNT(*) FROM attachments WHERE sha256 = $1`, a.SHA256).Scan(&refs); err != nil {
|
||||
return "", false, err
|
||||
}
|
||||
return a.SHA256, refs == 0, nil
|
||||
}
|
||||
|
||||
// SweepOrphanAttachments 清理超过 age 仍未挂载到邮件的附件记录，
|
||||
// 返回可以从磁盘删除的 sha256 列表（已确认无任何记录引用）。
|
||||
//
|
||||
// 上传后没走完发信流程（用户取消、Agent 崩溃）会留下这类记录，
|
||||
// 不清理的话磁盘只会单调增长。
|
||||
func SweepOrphanAttachments(ctx context.Context, age time.Duration) ([]string, error) {
|
||||
cutoff := time.Now().Add(-age)
|
||||
|
||||
rows, err := db.DB.QueryContext(ctx,
|
||||
`SELECT attachment_id, sha256 FROM attachments
|
||||
WHERE mail_id IS NULL AND created_at < $1`, cutoff)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
type orphan struct {
|
||||
id uuid.UUID
|
||||
sum string
|
||||
}
|
||||
var found []orphan
|
||||
for rows.Next() {
|
||||
var o orphan
|
||||
if err := rows.Scan(&o.id, &o.sum); err != nil {
|
||||
rows.Close()
|
||||
return nil, err
|
||||
}
|
||||
found = append(found, o)
|
||||
}
|
||||
rows.Close()
|
||||
if err := rows.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var removable []string
|
||||
for _, o := range found {
|
||||
if _, err := db.DB.ExecContext(ctx,
|
||||
`DELETE FROM attachments WHERE attachment_id = $1`, o.id); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var refs int
|
||||
if err := db.DB.QueryRowContext(ctx,
|
||||
`SELECT COUNT(*) FROM attachments WHERE sha256 = $1`, o.sum).Scan(&refs); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if refs == 0 {
|
||||
removable = append(removable, o.sum)
|
||||
}
|
||||
}
|
||||
return removable, nil
|
||||
}
|
||||
|
||||
// AttachmentAccessible 判断某人是否有权读取某附件：
|
||||
// 已挂载的看邮件所属会话的参与关系，未挂载的只有上传者本人能看。
|
||||
func AttachmentAccessible(ctx context.Context, a *models.Attachment, name string) (bool, error) {
|
||||
if a.MailID == nil {
|
||||
return a.Uploader == name, nil
|
||||
}
|
||||
var n int
|
||||
err := db.DB.QueryRowContext(ctx, `
|
||||
SELECT COUNT(*) FROM mails m
|
||||
WHERE m.mail_id = $1
|
||||
AND (m.from_name = $2 OR m.to_name = $2 OR `+db.CCHas("m.cc_list", 2)+`)
|
||||
`, *a.MailID, name).Scan(&n)
|
||||
return n > 0, err
|
||||
}
|
||||
165
gateway/internal/repo/budget_test.go
Normal file
165
gateway/internal/repo/budget_test.go
Normal file
@ -0,0 +1,165 @@
|
||||
package repo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"errors"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
// setupBudgetDB 复用 quota_test.go 的临时库，再建一个会话。
|
||||
// 用真实 SQLite 而非 mock：预算的正确性核心是「判断与自增在同一条 UPDATE 里」，
|
||||
// 那正是只有真实数据库才能验证的部分。
|
||||
func setupBudgetDB(t *testing.T) uuid.UUID {
|
||||
t.Helper()
|
||||
setupTestDB(t)
|
||||
id, err := CreateSession(context.Background(), nil, "bot", "预算测试")
|
||||
if err != nil {
|
||||
t.Fatalf("create session: %v", err)
|
||||
}
|
||||
return id
|
||||
}
|
||||
|
||||
func TestSessionBudgetZeroMeansUnlimited(t *testing.T) {
|
||||
id := setupBudgetDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
// 默认 0 = 不限：引入预算不该把已在进行的会话卡死
|
||||
b, err := GetSessionBudget(ctx, id)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !b.Unlimited || b.Remaining != -1 {
|
||||
t.Fatalf("默认应为不限：%+v", b)
|
||||
}
|
||||
// 不限时反复占用都成功
|
||||
for i := 0; i < 5; i++ {
|
||||
if _, err := ConsumeSessionBudget(ctx, id); err != nil {
|
||||
t.Fatalf("不限额下第 %d 次占用失败: %v", i+1, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSessionBudgetExhausts(t *testing.T) {
|
||||
id := setupBudgetDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
if _, err := SetSessionBudget(ctx, id, 2); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for i := 1; i <= 2; i++ {
|
||||
b, err := ConsumeSessionBudget(ctx, id)
|
||||
if err != nil {
|
||||
t.Fatalf("第 %d 次应成功: %v", i, err)
|
||||
}
|
||||
if b.Used != i {
|
||||
t.Fatalf("第 %d 次后 used = %d", i, b.Used)
|
||||
}
|
||||
}
|
||||
b, err := ConsumeSessionBudget(ctx, id)
|
||||
if !errors.Is(err, ErrSessionBudgetExhausted) {
|
||||
t.Fatalf("第 3 次应耗尽，得到 err=%v b=%+v", err, b)
|
||||
}
|
||||
if b.Remaining != 0 {
|
||||
t.Fatalf("耗尽后剩余应为 0：%+v", b)
|
||||
}
|
||||
}
|
||||
|
||||
// 判断与自增必须在同一条 UPDATE 里，否则并发下会把预算刷穿
|
||||
func TestSessionBudgetConcurrentDoesNotOverdraw(t *testing.T) {
|
||||
id := setupBudgetDB(t)
|
||||
ctx := context.Background()
|
||||
const limit = 10
|
||||
if _, err := SetSessionBudget(ctx, id, limit); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
var wg sync.WaitGroup
|
||||
var mu sync.Mutex
|
||||
ok := 0
|
||||
for i := 0; i < 40; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
if _, err := ConsumeSessionBudget(ctx, id); err == nil {
|
||||
mu.Lock()
|
||||
ok++
|
||||
mu.Unlock()
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
if ok != limit {
|
||||
t.Fatalf("40 并发下成功 %d 次，期望恰好 %d 次（预算被刷穿或误拒）", ok, limit)
|
||||
}
|
||||
b, _ := GetSessionBudget(ctx, id)
|
||||
if b.Used != limit {
|
||||
t.Fatalf("used = %d，期望 %d", b.Used, limit)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSessionBudgetResetAndLowerBelowUsed(t *testing.T) {
|
||||
id := setupBudgetDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
SetSessionBudget(ctx, id, 5)
|
||||
for i := 0; i < 3; i++ {
|
||||
ConsumeSessionBudget(ctx, id)
|
||||
}
|
||||
|
||||
// 调到低于已用次数 = 「就到这里为止」，是人的合法意图，不该报错
|
||||
b, err := SetSessionBudget(ctx, id, 1)
|
||||
if err != nil {
|
||||
t.Fatalf("下调预算不该失败: %v", err)
|
||||
}
|
||||
if b.Remaining != 0 {
|
||||
t.Fatalf("已用 3 上限 1 时剩余应为 0：%+v", b)
|
||||
}
|
||||
if _, err := ConsumeSessionBudget(ctx, id); !errors.Is(err, ErrSessionBudgetExhausted) {
|
||||
t.Fatal("下调后应立即拦住")
|
||||
}
|
||||
|
||||
// 重置只清已用次数，不动上限
|
||||
b, err = ResetSessionBudget(ctx, id)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if b.Used != 0 || b.Max != 1 {
|
||||
t.Fatalf("重置后应为 0/1：%+v", b)
|
||||
}
|
||||
}
|
||||
|
||||
// 会话预算先扣、全局配额后扣；全局拦下时必须把会话那次退回去，
|
||||
// 否则那格白掉了 —— 那次往返实际上没有发生
|
||||
func TestRefundSessionBudget(t *testing.T) {
|
||||
id := setupBudgetDB(t)
|
||||
ctx := context.Background()
|
||||
SetSessionBudget(ctx, id, 3)
|
||||
ConsumeSessionBudget(ctx, id)
|
||||
|
||||
RefundSessionBudget(ctx, id)
|
||||
b, _ := GetSessionBudget(ctx, id)
|
||||
if b.Used != 0 {
|
||||
t.Fatalf("退还后 used 应为 0：%+v", b)
|
||||
}
|
||||
|
||||
// 已经是 0 时再退不该变成负数
|
||||
RefundSessionBudget(ctx, id)
|
||||
b, _ = GetSessionBudget(ctx, id)
|
||||
if b.Used != 0 {
|
||||
t.Fatalf("重复退还把 used 变成了 %d", b.Used)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetSessionBudgetMissingSession(t *testing.T) {
|
||||
setupBudgetDB(t)
|
||||
if _, err := GetSessionBudget(context.Background(), uuid.New()); err == nil {
|
||||
t.Fatal("不存在的会话应报错")
|
||||
} else if errors.Is(err, sql.ErrNoRows) {
|
||||
t.Fatal("应包装成可读错误而不是裸 sql.ErrNoRows")
|
||||
}
|
||||
}
|
||||
341
gateway/internal/repo/keys.go
Normal file
341
gateway/internal/repo/keys.go
Normal file
@ -0,0 +1,341 @@
|
||||
package repo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"errors"
|
||||
"time"
|
||||
|
||||
"github.com/agentmail/gateway/internal/db"
|
||||
"github.com/agentmail/gateway/internal/models"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
// ---------- 密钥认证 ----------
|
||||
//
|
||||
// 两类密钥共享一个全局唯一的 token 命名空间：验证时先查 agent_keys 再查 user_keys。
|
||||
// 这样一个 token 永远只有一种身份，不会出现「同一串字符既能注册 Agent 又能读人类邮箱」。
|
||||
|
||||
var (
|
||||
// ErrKeyNotFound 密钥不存在
|
||||
ErrKeyNotFound = errors.New("key not found")
|
||||
// ErrKeyUsed 一次性密钥已被使用
|
||||
ErrKeyUsed = errors.New("key already used")
|
||||
// ErrKeyExpired 定时密钥已过期
|
||||
ErrKeyExpired = errors.New("key expired")
|
||||
// ErrKeyTypeInvalid 密钥类型不受支持
|
||||
ErrKeyTypeInvalid = errors.New("invalid key type")
|
||||
// ErrKeyNeedsExpiry timed 密钥缺少有效的 expires_hours
|
||||
ErrKeyNeedsExpiry = errors.New("timed key requires positive expires_hours")
|
||||
// ErrKeyTooShort 登记的客户端密钥长度不足
|
||||
ErrKeyTooShort = errors.New("key token too short")
|
||||
)
|
||||
|
||||
// expiryFor 依据密钥类型算出过期时间。
|
||||
// 只有 timed 需要 expires_at；permanent 与 one_time 都是 NULL，
|
||||
// 各自的失效条件由 key_type 本身表达，不混用 expires_at。
|
||||
func expiryFor(keyType string, hours int) (*time.Time, error) {
|
||||
if !models.ValidKeyType(keyType) {
|
||||
return nil, ErrKeyTypeInvalid
|
||||
}
|
||||
if keyType != models.KeyTimed {
|
||||
return nil, nil
|
||||
}
|
||||
if hours <= 0 {
|
||||
return nil, ErrKeyNeedsExpiry
|
||||
}
|
||||
t := time.Now().Add(time.Duration(hours) * time.Hour)
|
||||
return &t, nil
|
||||
}
|
||||
|
||||
// checkKeyUsable 判断一条密钥记录当前是否可用。
|
||||
func checkKeyUsable(keyType string, expiresAt, usedAt *time.Time) error {
|
||||
switch keyType {
|
||||
case models.KeyOneTime:
|
||||
if usedAt != nil {
|
||||
return ErrKeyUsed
|
||||
}
|
||||
case models.KeyTimed:
|
||||
if expiresAt == nil || time.Now().After(*expiresAt) {
|
||||
return ErrKeyExpired
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// ---------- Agent 密钥（管理员签发） ----------
|
||||
|
||||
// ErrKeyTokenTaken 登记的密钥已被占用
|
||||
var ErrKeyTokenTaken = errors.New("key token already registered")
|
||||
|
||||
// CreateAgentKey 签发一条 Agent 接入密钥。agentName 为空表示待绑定。
|
||||
//
|
||||
// presetToken 非空时登记客户端已在本地生成的密钥（插件首装场景），
|
||||
// 这样密钥全文只从客户端往服务器走一次，不需要反方向传递；留空则由服务器生成。
|
||||
func CreateAgentKey(ctx context.Context, agentName, keyType, label string, expiresHours int, createdBy uuid.UUID, presetToken string) (*models.AgentKey, error) {
|
||||
expires, err := expiryFor(keyType, expiresHours)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
token := presetToken
|
||||
if token == "" {
|
||||
if token, err = newToken(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
} else if len(token) < 32 {
|
||||
// 太短的客户端密钥不接受，否则等于把弱口令当凭证
|
||||
return nil, ErrKeyTooShort
|
||||
}
|
||||
|
||||
var namePtr *string
|
||||
if agentName != "" {
|
||||
namePtr = &agentName
|
||||
}
|
||||
|
||||
k := &models.AgentKey{
|
||||
Token: token,
|
||||
TokenHint: models.TokenHint(token),
|
||||
AgentName: namePtr,
|
||||
KeyType: keyType,
|
||||
Label: label,
|
||||
ExpiresAt: expires,
|
||||
CreatedBy: &createdBy,
|
||||
}
|
||||
err = db.DB.QueryRowContext(ctx, `
|
||||
INSERT INTO agent_keys (key_token, agent_name, key_type, label, expires_at, created_by)
|
||||
VALUES ($1, $2, $3, $4, $5, $6)
|
||||
RETURNING key_id, created_at
|
||||
`, token, namePtr, keyType, label, expires, createdBy).Scan(&k.ID, &k.CreatedAt)
|
||||
if db.IsUniqueViolation(err) {
|
||||
return nil, ErrKeyTokenTaken
|
||||
}
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return k, nil
|
||||
}
|
||||
|
||||
// ListAgentKeys 列出 Agent 密钥；agentName 非空时按 Agent 过滤。
|
||||
// 返回值不含 token 全文，只有 hint。
|
||||
func ListAgentKeys(ctx context.Context, agentName string) ([]models.AgentKey, error) {
|
||||
q := `SELECT key_id, key_token, agent_name, key_type, label, expires_at, used_at, created_by, created_at
|
||||
FROM agent_keys`
|
||||
args := []any{}
|
||||
if agentName != "" {
|
||||
q += ` WHERE agent_name = $1`
|
||||
args = append(args, agentName)
|
||||
}
|
||||
q += ` ORDER BY created_at DESC`
|
||||
|
||||
rows, err := db.DB.QueryContext(ctx, q, args...)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
keys := []models.AgentKey{}
|
||||
for rows.Next() {
|
||||
var k models.AgentKey
|
||||
var token string
|
||||
if err := rows.Scan(&k.ID, &token, &k.AgentName, &k.KeyType, &k.Label,
|
||||
&k.ExpiresAt, &k.UsedAt, &k.CreatedBy, &k.CreatedAt); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
k.TokenHint = models.TokenHint(token) // 不回传全文
|
||||
keys = append(keys, k)
|
||||
}
|
||||
return keys, rows.Err()
|
||||
}
|
||||
|
||||
// DeleteAgentKey 吊销一条 Agent 密钥。
|
||||
func DeleteAgentKey(ctx context.Context, id uuid.UUID) error {
|
||||
tag, err := db.DB.ExecContext(ctx, `DELETE FROM agent_keys WHERE key_id = $1`, id)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if n, _ := tag.RowsAffected(); n == 0 {
|
||||
return ErrKeyNotFound
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// BindAgentKey 把一条密钥绑定到指定 Agent 名。
|
||||
func BindAgentKey(ctx context.Context, id uuid.UUID, agentName string) error {
|
||||
tag, err := db.DB.ExecContext(ctx,
|
||||
`UPDATE agent_keys SET agent_name = $2 WHERE key_id = $1`, id, agentName)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if n, _ := tag.RowsAffected(); n == 0 {
|
||||
return ErrKeyNotFound
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// VerifyAgentKey 校验 Agent 密钥并返回它绑定的 Agent 名（未绑定时返回空串）。
|
||||
//
|
||||
// 一次性密钥在校验通过时立刻写 used_at —— 用 WHERE used_at IS NULL 保证并发下
|
||||
// 只有一个请求能把它标记掉，避免两个 Agent 拿同一把一次性密钥同时注册成功。
|
||||
func VerifyAgentKey(ctx context.Context, token string) (string, error) {
|
||||
var (
|
||||
id uuid.UUID
|
||||
agentName *string
|
||||
keyType string
|
||||
expiresAt *time.Time
|
||||
usedAt *time.Time
|
||||
)
|
||||
err := db.DB.QueryRowContext(ctx, `
|
||||
SELECT key_id, agent_name, key_type, expires_at, used_at
|
||||
FROM agent_keys WHERE key_token = $1
|
||||
`, token).Scan(&id, &agentName, &keyType, &expiresAt, &usedAt)
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
return "", ErrKeyNotFound
|
||||
}
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if err := checkKeyUsable(keyType, expiresAt, usedAt); err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
if keyType == models.KeyOneTime {
|
||||
tag, err := db.DB.ExecContext(ctx,
|
||||
`UPDATE agent_keys SET used_at = NOW() WHERE key_id = $1 AND used_at IS NULL`, id)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if n, _ := tag.RowsAffected(); n == 0 {
|
||||
return "", ErrKeyUsed // 并发下被别人抢先用掉了
|
||||
}
|
||||
}
|
||||
|
||||
if agentName == nil {
|
||||
return "", nil
|
||||
}
|
||||
return *agentName, nil
|
||||
}
|
||||
|
||||
// ClaimAgentKey 在待绑定密钥首次注册时把它落定到该 Agent 名。
|
||||
// 已绑定的密钥不受影响（WHERE agent_name IS NULL），因此不能借一把已绑定的密钥改注册别的 Agent。
|
||||
func ClaimAgentKey(ctx context.Context, token, agentName string) error {
|
||||
_, err := db.DB.ExecContext(ctx,
|
||||
`UPDATE agent_keys SET agent_name = $2 WHERE key_token = $1 AND agent_name IS NULL`,
|
||||
token, agentName)
|
||||
return err
|
||||
}
|
||||
|
||||
// ---------- 用户密钥（用户自助签发） ----------
|
||||
|
||||
// CreateUserKey 为用户签发一条客户端连接密钥。
|
||||
func CreateUserKey(ctx context.Context, userID uuid.UUID, label, keyType string, expiresHours int) (*models.UserKey, error) {
|
||||
expires, err := expiryFor(keyType, expiresHours)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
token, err := newToken()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
k := &models.UserKey{
|
||||
Token: token,
|
||||
TokenHint: models.TokenHint(token),
|
||||
UserID: userID,
|
||||
Label: label,
|
||||
KeyType: keyType,
|
||||
ExpiresAt: expires,
|
||||
}
|
||||
err = db.DB.QueryRowContext(ctx, `
|
||||
INSERT INTO user_keys (key_token, user_id, label, key_type, expires_at)
|
||||
VALUES ($1, $2, $3, $4, $5)
|
||||
RETURNING key_id, created_at
|
||||
`, token, userID, label, keyType, expires).Scan(&k.ID, &k.CreatedAt)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return k, nil
|
||||
}
|
||||
|
||||
// ListUserKeys 列出某用户的连接密钥（不含 token 全文）。
|
||||
func ListUserKeys(ctx context.Context, userID uuid.UUID) ([]models.UserKey, error) {
|
||||
rows, err := db.DB.QueryContext(ctx, `
|
||||
SELECT key_id, key_token, user_id, label, key_type, expires_at, used_at, created_at
|
||||
FROM user_keys WHERE user_id = $1 ORDER BY created_at DESC
|
||||
`, userID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
keys := []models.UserKey{}
|
||||
for rows.Next() {
|
||||
var k models.UserKey
|
||||
var token string
|
||||
if err := rows.Scan(&k.ID, &token, &k.UserID, &k.Label, &k.KeyType,
|
||||
&k.ExpiresAt, &k.UsedAt, &k.CreatedAt); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
k.TokenHint = models.TokenHint(token)
|
||||
keys = append(keys, k)
|
||||
}
|
||||
return keys, rows.Err()
|
||||
}
|
||||
|
||||
// DeleteUserKey 删除自己的一条密钥。带 user_id 条件，避免删掉别人的。
|
||||
func DeleteUserKey(ctx context.Context, userID, keyID uuid.UUID) error {
|
||||
tag, err := db.DB.ExecContext(ctx,
|
||||
`DELETE FROM user_keys WHERE key_id = $1 AND user_id = $2`, keyID, userID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if n, _ := tag.RowsAffected(); n == 0 {
|
||||
return ErrKeyNotFound
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// VerifyUserKey 校验用户密钥并返回对应用户。
|
||||
// 用户必须仍处于 active 状态——禁用账号后其密钥应当立即失效。
|
||||
func VerifyUserKey(ctx context.Context, token string) (*models.User, error) {
|
||||
var (
|
||||
id uuid.UUID
|
||||
userID uuid.UUID
|
||||
keyType string
|
||||
expiresAt *time.Time
|
||||
usedAt *time.Time
|
||||
)
|
||||
err := db.DB.QueryRowContext(ctx, `
|
||||
SELECT key_id, user_id, key_type, expires_at, used_at
|
||||
FROM user_keys WHERE key_token = $1
|
||||
`, token).Scan(&id, &userID, &keyType, &expiresAt, &usedAt)
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
return nil, ErrKeyNotFound
|
||||
}
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := checkKeyUsable(keyType, expiresAt, usedAt); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if keyType == models.KeyOneTime {
|
||||
tag, err := db.DB.ExecContext(ctx,
|
||||
`UPDATE user_keys SET used_at = NOW() WHERE key_id = $1 AND used_at IS NULL`, id)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if n, _ := tag.RowsAffected(); n == 0 {
|
||||
return nil, ErrKeyUsed
|
||||
}
|
||||
}
|
||||
|
||||
u, err := GetUserByID(ctx, userID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if u.Status != "active" {
|
||||
return nil, ErrKeyNotFound // 账号已禁用，密钥一并失效
|
||||
}
|
||||
return u, nil
|
||||
}
|
||||
282
gateway/internal/repo/quota.go
Normal file
282
gateway/internal/repo/quota.go
Normal file
@ -0,0 +1,282 @@
|
||||
package repo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"errors"
|
||||
"fmt"
|
||||
|
||||
"github.com/agentmail/gateway/internal/db"
|
||||
"github.com/agentmail/gateway/internal/models"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
// ---------- 配额 ----------
|
||||
//
|
||||
// 配额限制的是 Agent「主动发信」的次数，不限制收信：
|
||||
// 收信是被动的，卡住收信只会让邮件凭空消失；卡住发信才能真正阻止 Agent 无限自我循环。
|
||||
//
|
||||
// agents.max_rounds = 0 表示不限额。used_rounds 单调递增，由管理员显式重置。
|
||||
|
||||
// ErrQuotaExhausted 表示 Agent 的发信配额已用尽。
|
||||
var ErrQuotaExhausted = errors.New("quota exhausted")
|
||||
|
||||
// Quota 是一个 Agent 的配额快照。
|
||||
type Quota struct {
|
||||
AgentName string `json:"agent_name"`
|
||||
Max int `json:"max_rounds"` // 0 = 不限
|
||||
Used int `json:"used_rounds"`
|
||||
Remaining int `json:"remaining"` // 不限时为 -1
|
||||
Unlimited bool `json:"unlimited"`
|
||||
}
|
||||
|
||||
func makeQuota(name string, max, used int) Quota {
|
||||
q := Quota{AgentName: name, Max: max, Used: used, Unlimited: max <= 0}
|
||||
if q.Unlimited {
|
||||
q.Remaining = -1
|
||||
return q
|
||||
}
|
||||
if r := max - used; r > 0 {
|
||||
q.Remaining = r
|
||||
}
|
||||
return q
|
||||
}
|
||||
|
||||
// GetQuota 读取某 Agent 的配额状态。
|
||||
func GetQuota(ctx context.Context, agentName string) (Quota, error) {
|
||||
var max, used int
|
||||
err := db.DB.QueryRowContext(ctx,
|
||||
`SELECT max_rounds, used_rounds FROM agents WHERE agent_name = $1`, agentName).Scan(&max, &used)
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
return Quota{}, fmt.Errorf("agent %q 不存在", agentName)
|
||||
}
|
||||
if err != nil {
|
||||
return Quota{}, err
|
||||
}
|
||||
return makeQuota(agentName, max, used), nil
|
||||
}
|
||||
|
||||
// ConsumeQuota 原子地占用一次发信配额，返回占用后的快照。
|
||||
//
|
||||
// 判断与自增必须在同一条 UPDATE 里完成（WHERE used_rounds < max_rounds），
|
||||
// 否则并发发信会双双通过检查再各自 +1，把配额刷穿。
|
||||
// 配额耗尽时返回 ErrQuotaExhausted，同时给出快照供调用方生成提示。
|
||||
func ConsumeQuota(ctx context.Context, agentName string) (Quota, error) {
|
||||
tag, err := db.DB.ExecContext(ctx, `
|
||||
UPDATE agents SET used_rounds = used_rounds + 1
|
||||
WHERE agent_name = $1
|
||||
AND (max_rounds <= 0 OR used_rounds < max_rounds)
|
||||
`, agentName)
|
||||
if err != nil {
|
||||
return Quota{}, err
|
||||
}
|
||||
n, _ := tag.RowsAffected()
|
||||
if n == 0 {
|
||||
// 要么 Agent 不存在，要么配额用尽——用快照区分
|
||||
q, qErr := GetQuota(ctx, agentName)
|
||||
if qErr != nil {
|
||||
return Quota{}, qErr
|
||||
}
|
||||
return q, ErrQuotaExhausted
|
||||
}
|
||||
return GetQuota(ctx, agentName)
|
||||
}
|
||||
|
||||
// SetQuota 设置某 Agent 的配额上限（0 = 不限）。
|
||||
func SetQuota(ctx context.Context, agentName string, max int) (Quota, error) {
|
||||
if max < 0 {
|
||||
max = 0
|
||||
}
|
||||
tag, err := db.DB.ExecContext(ctx,
|
||||
`UPDATE agents SET max_rounds = $2 WHERE agent_name = $1`, agentName, max)
|
||||
if err != nil {
|
||||
return Quota{}, err
|
||||
}
|
||||
if n, _ := tag.RowsAffected(); n == 0 {
|
||||
return Quota{}, fmt.Errorf("agent %q 不存在", agentName)
|
||||
}
|
||||
return GetQuota(ctx, agentName)
|
||||
}
|
||||
|
||||
// ResetQuota 把已用次数归零（配额上限不变）。
|
||||
func ResetQuota(ctx context.Context, agentName string) (Quota, error) {
|
||||
tag, err := db.DB.ExecContext(ctx,
|
||||
`UPDATE agents SET used_rounds = 0 WHERE agent_name = $1`, agentName)
|
||||
if err != nil {
|
||||
return Quota{}, err
|
||||
}
|
||||
if n, _ := tag.RowsAffected(); n == 0 {
|
||||
return Quota{}, fmt.Errorf("agent %q 不存在", agentName)
|
||||
}
|
||||
return GetQuota(ctx, agentName)
|
||||
}
|
||||
|
||||
// ListQuotas 列出所有 Agent 的配额（管理员视图）。
|
||||
func ListQuotas(ctx context.Context) ([]Quota, error) {
|
||||
rows, err := db.DB.QueryContext(ctx,
|
||||
`SELECT agent_name, max_rounds, used_rounds FROM agents ORDER BY agent_name`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
out := []Quota{}
|
||||
for rows.Next() {
|
||||
var name string
|
||||
var max, used int
|
||||
if err := rows.Scan(&name, &max, &used); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, makeQuota(name, max, used))
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// ---------- 转发 ----------
|
||||
|
||||
// ForwardSource 是被转发邮件的必要信息。
|
||||
type ForwardSource struct {
|
||||
Mail *models.Mail
|
||||
Session uuid.UUID
|
||||
}
|
||||
|
||||
// LoadForwardSource 读取待转发的邮件，并校验转发者确实参与过该邮件
|
||||
//（收件人、发件人或被抄送方之一）。防止凭 mail_id 转发别人的邮件。
|
||||
func LoadForwardSource(ctx context.Context, mailID uuid.UUID, actor string) (*models.Mail, error) {
|
||||
m, err := GetMailByID(ctx, mailID)
|
||||
if err != nil {
|
||||
return nil, ErrMailNotFound
|
||||
}
|
||||
|
||||
if m.FromName == actor || m.ToName == actor {
|
||||
return m, nil
|
||||
}
|
||||
for _, cc := range m.CCList {
|
||||
if cc.Name == actor {
|
||||
return m, nil
|
||||
}
|
||||
}
|
||||
return nil, ErrForwardNotAllowed
|
||||
}
|
||||
|
||||
var (
|
||||
// ErrMailNotFound 待转发的邮件不存在
|
||||
ErrMailNotFound = errors.New("mail not found")
|
||||
// ErrForwardNotAllowed 转发者未参与该邮件
|
||||
ErrForwardNotAllowed = errors.New("not a participant of that mail")
|
||||
)
|
||||
|
||||
// ---------- 会话级往返预算 ----------
|
||||
//
|
||||
// 配额的真实语义是「这件事值得多少个来回」——那是**任务**的属性，不是 Agent 的属性。
|
||||
// 只有 agents.max_rounds 一个全局计数器时有两个问题：
|
||||
// 1. 两个并行任务互相抢额度：给紧急任务留的份被另一条线索吃掉
|
||||
// 2. used_rounds 单调递增，跑满就得管理员手工重置才能再干活
|
||||
// 因此预算下沉到会话，由人在写信时给、在对话页里随时调。
|
||||
//
|
||||
// **两层都要过**：会话预算 + Agent 全局配额。少了后者，Agent 自己 `.new` 开一串会话
|
||||
// 每条都是全新预算，全局上限形同虚设；少了前者，就回到抢额度的老问题。
|
||||
|
||||
// SessionBudget 是一个会话的往返预算快照。
|
||||
type SessionBudget struct {
|
||||
SessionID uuid.UUID `json:"session_id"`
|
||||
Max int `json:"max_rounds"` // 0 = 本会话不限
|
||||
Used int `json:"used_rounds"`
|
||||
Remaining int `json:"remaining"` // 不限时为 -1
|
||||
Unlimited bool `json:"unlimited"`
|
||||
}
|
||||
|
||||
func makeSessionBudget(id uuid.UUID, max, used int) SessionBudget {
|
||||
b := SessionBudget{SessionID: id, Max: max, Used: used, Unlimited: max <= 0}
|
||||
if b.Unlimited {
|
||||
b.Remaining = -1
|
||||
return b
|
||||
}
|
||||
if r := max - used; r > 0 {
|
||||
b.Remaining = r
|
||||
}
|
||||
return b
|
||||
}
|
||||
|
||||
// ErrSessionBudgetExhausted 表示该会话的往返预算已用尽。
|
||||
var ErrSessionBudgetExhausted = errors.New("session budget exhausted")
|
||||
|
||||
// GetSessionBudget 读取会话预算。
|
||||
func GetSessionBudget(ctx context.Context, id uuid.UUID) (SessionBudget, error) {
|
||||
var max, used int
|
||||
err := db.DB.QueryRowContext(ctx,
|
||||
`SELECT COALESCE(max_rounds,0), COALESCE(used_rounds,0) FROM sessions WHERE session_id = $1`,
|
||||
id).Scan(&max, &used)
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
return SessionBudget{}, fmt.Errorf("会话 %s 不存在", id)
|
||||
}
|
||||
if err != nil {
|
||||
return SessionBudget{}, err
|
||||
}
|
||||
return makeSessionBudget(id, max, used), nil
|
||||
}
|
||||
|
||||
// ConsumeSessionBudget 原子地占用会话的一次往返。
|
||||
//
|
||||
// 与 ConsumeQuota 同理：判断与自增必须在同一条 UPDATE 里（WHERE used_rounds < max_rounds），
|
||||
// 否则并发发信会双双通过检查再各自 +1，把预算刷穿。
|
||||
func ConsumeSessionBudget(ctx context.Context, id uuid.UUID) (SessionBudget, error) {
|
||||
tag, err := db.DB.ExecContext(ctx, `
|
||||
UPDATE sessions SET used_rounds = COALESCE(used_rounds,0) + 1
|
||||
WHERE session_id = $1
|
||||
AND (COALESCE(max_rounds,0) <= 0 OR COALESCE(used_rounds,0) < max_rounds)
|
||||
`, id)
|
||||
if err != nil {
|
||||
return SessionBudget{}, err
|
||||
}
|
||||
if n, _ := tag.RowsAffected(); n == 0 {
|
||||
b, bErr := GetSessionBudget(ctx, id)
|
||||
if bErr != nil {
|
||||
return SessionBudget{}, bErr
|
||||
}
|
||||
return b, ErrSessionBudgetExhausted
|
||||
}
|
||||
return GetSessionBudget(ctx, id)
|
||||
}
|
||||
|
||||
// SetSessionBudget 设置会话预算上限（0 = 不限）。
|
||||
//
|
||||
// 允许把上限调到低于已用次数：那表示「就到这里为止」，是人的合法意图，
|
||||
// 不该因为算不出正的剩余量就拒绝。此时 Remaining 为 0，下次发信即被拦。
|
||||
func SetSessionBudget(ctx context.Context, id uuid.UUID, max int) (SessionBudget, error) {
|
||||
if max < 0 {
|
||||
max = 0
|
||||
}
|
||||
tag, err := db.DB.ExecContext(ctx,
|
||||
`UPDATE sessions SET max_rounds = $2, updated_at = NOW() WHERE session_id = $1`, id, max)
|
||||
if err != nil {
|
||||
return SessionBudget{}, err
|
||||
}
|
||||
if n, _ := tag.RowsAffected(); n == 0 {
|
||||
return SessionBudget{}, fmt.Errorf("会话 %s 不存在", id)
|
||||
}
|
||||
return GetSessionBudget(ctx, id)
|
||||
}
|
||||
|
||||
// ResetSessionBudget 把该会话的已用次数归零（上限不变）。
|
||||
func ResetSessionBudget(ctx context.Context, id uuid.UUID) (SessionBudget, error) {
|
||||
tag, err := db.DB.ExecContext(ctx,
|
||||
`UPDATE sessions SET used_rounds = 0, updated_at = NOW() WHERE session_id = $1`, id)
|
||||
if err != nil {
|
||||
return SessionBudget{}, err
|
||||
}
|
||||
if n, _ := tag.RowsAffected(); n == 0 {
|
||||
return SessionBudget{}, fmt.Errorf("会话 %s 不存在", id)
|
||||
}
|
||||
return GetSessionBudget(ctx, id)
|
||||
}
|
||||
|
||||
// RefundSessionBudget 退还一次往返。
|
||||
//
|
||||
// 会话预算先扣、Agent 全局配额后扣，全局那层拦下时必须把会话这次还回去，
|
||||
// 否则会话预算白掉一格 —— 那次往返实际上没有发生。
|
||||
func RefundSessionBudget(ctx context.Context, id uuid.UUID) {
|
||||
_, _ = db.DB.ExecContext(ctx,
|
||||
`UPDATE sessions SET used_rounds = COALESCE(used_rounds,0) - 1
|
||||
WHERE session_id = $1 AND COALESCE(used_rounds,0) > 0`, id)
|
||||
}
|
||||
173
gateway/internal/repo/quota_test.go
Normal file
173
gateway/internal/repo/quota_test.go
Normal file
@ -0,0 +1,173 @@
|
||||
package repo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"github.com/agentmail/gateway/internal/db"
|
||||
)
|
||||
|
||||
// setupTestDB 起一个临时 SQLite 库并建表，供配额测试使用。
|
||||
// 直接用真实的 SQLite 而非 mock：配额的正确性核心在于「判断与自增在同一条 UPDATE 里」，
|
||||
// 这正是只有真实数据库才能验证的部分。
|
||||
func setupTestDB(t *testing.T) {
|
||||
t.Helper()
|
||||
|
||||
dir := t.TempDir()
|
||||
if err := db.Connect(context.Background(), filepath.Join(dir, "test.db")); err != nil {
|
||||
t.Fatalf("connect: %v", err)
|
||||
}
|
||||
if err := db.Migrate(context.Background()); err != nil {
|
||||
t.Fatalf("migrate: %v", err)
|
||||
}
|
||||
t.Cleanup(db.Close)
|
||||
}
|
||||
|
||||
func seedAgent(t *testing.T, name string, max int) {
|
||||
t.Helper()
|
||||
_, err := db.DB.ExecContext(context.Background(),
|
||||
`INSERT INTO agents (agent_name, secret, platform, max_rounds) VALUES ($1, 'x', 'test', $2)`,
|
||||
name, max)
|
||||
if err != nil {
|
||||
t.Fatalf("seed agent: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestConsumeQuotaCountsDown(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
seedAgent(t, "bot", 3)
|
||||
ctx := context.Background()
|
||||
|
||||
for i := 1; i <= 3; i++ {
|
||||
q, err := ConsumeQuota(ctx, "bot")
|
||||
if err != nil {
|
||||
t.Fatalf("第 %d 次占用失败: %v", i, err)
|
||||
}
|
||||
if q.Used != i || q.Remaining != 3-i {
|
||||
t.Errorf("第 %d 次: used=%d remaining=%d, want used=%d remaining=%d",
|
||||
i, q.Used, q.Remaining, i, 3-i)
|
||||
}
|
||||
}
|
||||
|
||||
q, err := ConsumeQuota(ctx, "bot")
|
||||
if !errors.Is(err, ErrQuotaExhausted) {
|
||||
t.Fatalf("第 4 次应耗尽，得到 err=%v", err)
|
||||
}
|
||||
// 耗尽时仍要给出快照，调用方才能在错误文案里写清 used/max
|
||||
if q.Used != 3 || q.Max != 3 {
|
||||
t.Errorf("耗尽时快照不对: %+v", q)
|
||||
}
|
||||
}
|
||||
|
||||
func TestConsumeQuotaUnlimited(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
seedAgent(t, "free", 0) // 0 = 不限
|
||||
ctx := context.Background()
|
||||
|
||||
for i := 0; i < 5; i++ {
|
||||
q, err := ConsumeQuota(ctx, "free")
|
||||
if err != nil {
|
||||
t.Fatalf("不限额时不应失败: %v", err)
|
||||
}
|
||||
if !q.Unlimited || q.Remaining != -1 {
|
||||
t.Errorf("不限额快照不对: %+v", q)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 配额的核心保证：并发发信不能把额度刷穿。
|
||||
// 判断与自增若分成两步（先 SELECT 再 UPDATE），并发下两个请求会双双通过检查。
|
||||
func TestConsumeQuotaConcurrentDoesNotOverdraw(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
const limit = 10
|
||||
seedAgent(t, "racer", limit)
|
||||
ctx := context.Background()
|
||||
|
||||
const attempts = 40
|
||||
var (
|
||||
wg sync.WaitGroup
|
||||
mu sync.Mutex
|
||||
granted int
|
||||
)
|
||||
for i := 0; i < attempts; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
if _, err := ConsumeQuota(ctx, "racer"); err == nil {
|
||||
mu.Lock()
|
||||
granted++
|
||||
mu.Unlock()
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
if granted != limit {
|
||||
t.Errorf("并发 %d 次请求在上限 %d 下放行了 %d 次", attempts, limit, granted)
|
||||
}
|
||||
|
||||
q, err := GetQuota(ctx, "racer")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if q.Used != limit {
|
||||
t.Errorf("used_rounds = %d，应恰好等于上限 %d（未刷穿也未少记）", q.Used, limit)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSetAndResetQuota(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
seedAgent(t, "bot", 2)
|
||||
ctx := context.Background()
|
||||
|
||||
if _, err := ConsumeQuota(ctx, "bot"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
q, err := SetQuota(ctx, "bot", 5)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// 调上限不应清掉已用次数
|
||||
if q.Max != 5 || q.Used != 1 || q.Remaining != 4 {
|
||||
t.Errorf("SetQuota 后 %+v，want max=5 used=1 remaining=4", q)
|
||||
}
|
||||
|
||||
q, err = ResetQuota(ctx, "bot")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if q.Used != 0 || q.Remaining != 5 {
|
||||
t.Errorf("ResetQuota 后 %+v，want used=0 remaining=5", q)
|
||||
}
|
||||
|
||||
// 负数上限归一为 0（不限），而不是造出一个永远发不出信的 Agent
|
||||
if q, err = SetQuota(ctx, "bot", -3); err != nil {
|
||||
t.Fatal(err)
|
||||
} else if !q.Unlimited {
|
||||
t.Errorf("负数上限应视为不限，得到 %+v", q)
|
||||
}
|
||||
}
|
||||
|
||||
func TestQuotaUnknownAgent(t *testing.T) {
|
||||
setupTestDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
if _, err := GetQuota(ctx, "ghost"); err == nil {
|
||||
t.Error("不存在的 Agent 应报错")
|
||||
}
|
||||
if _, err := ConsumeQuota(ctx, "ghost"); err == nil {
|
||||
t.Error("不存在的 Agent 占用配额应报错")
|
||||
}
|
||||
if errors.Is(func() error { _, e := ConsumeQuota(ctx, "ghost"); return e }(), ErrQuotaExhausted) {
|
||||
t.Error("不存在的 Agent 不该被报成「配额耗尽」")
|
||||
}
|
||||
}
|
||||
|
||||
func TestMain(m *testing.M) {
|
||||
os.Exit(m.Run())
|
||||
}
|
||||
87
gateway/internal/repo/relay.go
Normal file
87
gateway/internal/repo/relay.go
Normal file
@ -0,0 +1,87 @@
|
||||
package repo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
|
||||
"github.com/agentmail/gateway/internal/db"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
// ---------- 插件自动转发（免配额通道） ----------
|
||||
//
|
||||
// **基本原则：配额约束的是模型的自主发信，不是 harness 的转发。**
|
||||
//
|
||||
// 配额存在的意义是防止 Agent 无限自我循环。而插件代劳搬运的两类消息不属于此列：
|
||||
// 1. 平台原生的权限询问（opencode 的 permission.updated）—— 不转给人，人就看不到，
|
||||
// Agent 卡在那里等一个永远不会来的回答
|
||||
// 2. 本轮的最终总结（session.idle 时最后一条 assistant 消息）—— 模型已经把话说完了，
|
||||
// 插件只是把它搬到邮件里；对它收费会导致配额用尽时 Agent 连交代都做不了
|
||||
//
|
||||
// 防滥用不靠计数，靠**幂等键**：relay_key 是上游那条消息的稳定标识
|
||||
// （permission id / assistant message id）。唯一约束让同一条上游消息只能转一次，
|
||||
// 于是插件重试与 SSE 重放不会产生第二封，想多转就得拿出不同的上游消息 id ——
|
||||
// 而那些 id 由平台生成，模型伪造不出来。
|
||||
|
||||
// ErrRelayDuplicate 表示这条上游消息已经转发过了。
|
||||
var ErrRelayDuplicate = errors.New("relay already recorded")
|
||||
|
||||
// ClaimRelay 占用一次免配额转发名额。
|
||||
//
|
||||
// 判断与占用在同一条 INSERT 里（靠主键唯一约束），并发重试下只有一个能成功 ——
|
||||
// 分成「先查有没有、再插入」两步的话，插件的两次重试会双双通过检查各插一条。
|
||||
//
|
||||
// 返回 ErrRelayDuplicate 表示重复，调用方应当据此跳过发信而不是报错：
|
||||
// 重复转发是插件重试的正常结果，不是故障。
|
||||
func ClaimRelay(ctx context.Context, agentName, relayKey, kind string) error {
|
||||
_, err := db.DB.ExecContext(ctx,
|
||||
`INSERT INTO relayed_mails (agent_name, relay_key, kind) VALUES ($1, $2, $3)`,
|
||||
agentName, relayKey, kind)
|
||||
if err != nil {
|
||||
if db.IsUniqueViolation(err) {
|
||||
return ErrRelayDuplicate
|
||||
}
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// BindRelayMail 把已占用的名额关联到真正发出的邮件，便于事后审计
|
||||
// 「这封免配额的信是从哪条上游消息来的」。
|
||||
//
|
||||
// 关联失败不该让发信失败：邮件已经入库，缺一条审计关联不影响功能。
|
||||
func BindRelayMail(ctx context.Context, agentName, relayKey string, mailID uuid.UUID) error {
|
||||
_, err := db.DB.ExecContext(ctx,
|
||||
`UPDATE relayed_mails SET mail_id = $1 WHERE agent_name = $2 AND relay_key = $3`,
|
||||
mailID, agentName, relayKey)
|
||||
return err
|
||||
}
|
||||
|
||||
// ReleaseRelay 撤销名额占用。
|
||||
//
|
||||
// 占用成功但发信失败时必须还回去，否则那条上游消息永远转不出来了 ——
|
||||
// 幂等键会一直认为它已经转过。
|
||||
func ReleaseRelay(ctx context.Context, agentName, relayKey string) error {
|
||||
_, err := db.DB.ExecContext(ctx,
|
||||
`DELETE FROM relayed_mails WHERE agent_name = $1 AND relay_key = $2 AND mail_id IS NULL`,
|
||||
agentName, relayKey)
|
||||
return err
|
||||
}
|
||||
|
||||
// RelayKeyForMail 反查某封邮件对应的上游消息 id。
|
||||
//
|
||||
// 人类决策一条权限请求后，插件需要知道该回复 opencode 的哪个 permission ——
|
||||
// 光有 AgentMail 的 mail_id 是不够的，两边的 id 空间不同。
|
||||
// 插件重启后内存映射会丢，所以这个映射必须在服务端持久化。
|
||||
//
|
||||
// 无记录时返回空串（例如旧数据，或压根没走 relay 通道的请求）。
|
||||
func RelayKeyForMail(ctx context.Context, mailID uuid.UUID) (string, string) {
|
||||
var key, kind string
|
||||
err := db.DB.QueryRowContext(ctx,
|
||||
`SELECT relay_key, kind FROM relayed_mails WHERE mail_id = $1 LIMIT 1`,
|
||||
mailID).Scan(&key, &kind)
|
||||
if err != nil {
|
||||
return "", ""
|
||||
}
|
||||
return key, kind
|
||||
}
|
||||
1046
gateway/internal/repo/repo.go
Normal file
1046
gateway/internal/repo/repo.go
Normal file
File diff suppressed because it is too large
Load Diff
188
gateway/internal/repo/thread.go
Normal file
188
gateway/internal/repo/thread.go
Normal file
@ -0,0 +1,188 @@
|
||||
package repo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
|
||||
"github.com/agentmail/gateway/internal/db"
|
||||
"github.com/agentmail/gateway/internal/models"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
// 对话树。
|
||||
//
|
||||
// **不另建 tree_nodes 表**：`mails.parent_mail_id` 已经完整编码了树结构 ——
|
||||
// 回复指向来信，转发指向被转发的原件。再维护一张 tree_nodes 就是第二份真相，
|
||||
// 两处不一致时无法判断谁对。这里直接用递归 CTE 在 mails 上查。
|
||||
//
|
||||
// 树可以跨会话：转发把线索引到新会话，但 parent 仍指向原件。这正是「对话树」比
|
||||
// 「会话内平铺」更有价值的地方 —— 能看出一条线索分叉去了哪里。
|
||||
// 也正因如此，读取时必须按会话逐个鉴权（见 handler）：
|
||||
// A 转发给 B 之后，B 与 C 在新会话里的往来不能回流给 A。
|
||||
//
|
||||
// **分块加载而非截断**：线索可以有几百封，一次全取要把几 MB 预览塞给前端。
|
||||
// 按方向分页 —— 祖先向上、子孙向下，各自带游标。
|
||||
//
|
||||
// 游标用「相对锚点的原始层号偏移」而不是 mail_id：
|
||||
// - 偏移量每次从锚点重走一遍，无状态、不可伪造，也不需要额外证明
|
||||
// 「这个 cursor 真的在这条线索上」
|
||||
// - 用 mail_id 做游标就必须允许传入**不可见**的邮件（不可见的中间段要穿过去），
|
||||
// 那就得单独校验它确实是锚点的祖先，反而更绕
|
||||
// - 祖先方向的层号天然稳定：新邮件只会追加成叶子，不会插进已有链条中间
|
||||
|
||||
// TreeMail 是树里的一个节点。正文只带预览：整棵线索带全文可能几百 KB，
|
||||
// 前端点开某封时再单取全文与附件清单。
|
||||
type TreeMail struct {
|
||||
models.Mail
|
||||
// Depth 是**相对锚点**的层级：0 = 锚点，-1 = 父，1 = 子。
|
||||
// 不用「距根深度」—— 分块加载时根可能还没取到，绝对深度无从得知。
|
||||
Depth int `json:"depth"`
|
||||
AttachmentCount int `json:"attachment_count"`
|
||||
}
|
||||
|
||||
// descendantDepthCap 只是数据损坏时的兜底。
|
||||
//
|
||||
// parent_mail_id 正常不成环（新邮件只能指向已存在的旧邮件），但一旦被外部工具改坏，
|
||||
// 无上限的递归 CTE 会把进程拖死。取得足够大，正常数据碰不到。
|
||||
const descendantDepthCap = 10000
|
||||
|
||||
const threadCols = `m.mail_id, m.session_id, m.parent_mail_id,
|
||||
m.from_name, m.from_workspace, m.to_name, m.to_workspace,
|
||||
m.cc_list, m.subject, m.body, m.mail_type,
|
||||
COALESCE(m.permission_result,'') AS permission_result,
|
||||
m.status, m.created_at, s.session_alias,
|
||||
(SELECT COUNT(*) FROM attachments a WHERE a.mail_id = m.mail_id) AS attach_count`
|
||||
|
||||
// AncestorsRaw 沿 parent_mail_id 上溯，取第 offset+1 .. offset+limit 层的祖先。
|
||||
// 层号 1 = 父，2 = 祖父；返回的 Depth 为负数。
|
||||
//
|
||||
// **不做可见性过滤** —— 不可见的中间段必须能穿过：转发把线索引进别人的会话，
|
||||
// 再往上却可能仍是自己参与的往来。过滤放在 handler 层（那里知道调用者是谁）。
|
||||
//
|
||||
// 第二个返回值表示 offset+limit 层之上还有节点。
|
||||
func AncestorsRaw(ctx context.Context, anchorID uuid.UUID, offset, limit int) ([]TreeMail, bool, error) {
|
||||
rows, err := db.DB.QueryContext(ctx, `
|
||||
WITH RECURSIVE up(mail_id, parent_mail_id, lvl) AS (
|
||||
SELECT mail_id, parent_mail_id, 0 FROM mails WHERE mail_id = $1
|
||||
UNION ALL
|
||||
SELECT m.mail_id, m.parent_mail_id, up.lvl + 1
|
||||
FROM mails m JOIN up ON m.mail_id = up.parent_mail_id
|
||||
WHERE up.lvl < $2
|
||||
)
|
||||
SELECT `+threadCols+`, u.lvl
|
||||
FROM up u
|
||||
JOIN mails m ON m.mail_id = u.mail_id
|
||||
JOIN sessions s ON m.session_id = s.session_id
|
||||
WHERE u.lvl > $3
|
||||
ORDER BY u.lvl ASC
|
||||
`, anchorID, offset+limit+1, offset)
|
||||
if err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
// 多取一层用来判断「上面还有没有」，不返回给调用方
|
||||
out, err := scanTreeRows(rows, true)
|
||||
if err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
hasMore := len(out) > limit
|
||||
if hasMore {
|
||||
out = out[:limit]
|
||||
}
|
||||
return out, hasMore, nil
|
||||
}
|
||||
|
||||
// DescendantsRaw 取锚点及其子孙，BFS 顺序（同层按时间），按节点数分页。
|
||||
//
|
||||
// offset = 0 时结果的第一个是锚点自己（Depth 0）。
|
||||
// 同样不做可见性过滤，理由同 AncestorsRaw：不可见的子节点下面可能挂着可见的孙节点
|
||||
// （别人把线索转走又转回来给我）。
|
||||
//
|
||||
// 注意 CTE 每次都会走完整棵子树，LIMIT 只截断输出。一封邮件的子孙通常很少
|
||||
// （分支来自转发，不是回复），这个代价可以接受；真出现巨型子树时再加物化。
|
||||
func DescendantsRaw(ctx context.Context, anchorID uuid.UUID, offset, limit int) ([]TreeMail, bool, error) {
|
||||
rows, err := db.DB.QueryContext(ctx, `
|
||||
WITH RECURSIVE down(mail_id, lvl) AS (
|
||||
SELECT mail_id, 0 FROM mails WHERE mail_id = $1
|
||||
UNION ALL
|
||||
SELECT m.mail_id, down.lvl + 1
|
||||
FROM mails m JOIN down ON m.parent_mail_id = down.mail_id
|
||||
WHERE down.lvl < $2
|
||||
)
|
||||
SELECT `+threadCols+`, d.lvl
|
||||
FROM down d
|
||||
JOIN mails m ON m.mail_id = d.mail_id
|
||||
JOIN sessions s ON m.session_id = s.session_id
|
||||
ORDER BY d.lvl ASC, m.created_at ASC, m.mail_id ASC
|
||||
LIMIT $3 OFFSET $4
|
||||
`, anchorID, descendantDepthCap, limit+1, offset)
|
||||
if err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
out, err := scanTreeRows(rows, false)
|
||||
if err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
hasMore := len(out) > limit
|
||||
if hasMore {
|
||||
out = out[:limit]
|
||||
}
|
||||
return out, hasMore, nil
|
||||
}
|
||||
|
||||
// scanTreeRows 读出节点。negate 为真时把层号取负（祖先方向）。
|
||||
func scanTreeRows(rows interface {
|
||||
Next() bool
|
||||
Scan(...interface{}) error
|
||||
Err() error
|
||||
Close() error
|
||||
}, negate bool) ([]TreeMail, error) {
|
||||
defer rows.Close()
|
||||
|
||||
out := []TreeMail{}
|
||||
for rows.Next() {
|
||||
var t TreeMail
|
||||
var alias *string
|
||||
var ccJSON []byte
|
||||
var lvl int
|
||||
if err := rows.Scan(&t.ID, &t.SessionID, &t.ParentMailID,
|
||||
&t.FromName, &t.FromWorkspace, &t.ToName, &t.ToWorkspace,
|
||||
&ccJSON, &t.Subject, &t.Body, &t.MailType, &t.PermResult,
|
||||
&t.Status, &t.CreatedAt, &alias, &t.AttachmentCount, &lvl); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(ccJSON) > 0 {
|
||||
json.Unmarshal(ccJSON, &t.CCList)
|
||||
}
|
||||
if t.CCList == nil {
|
||||
t.CCList = []models.Address{}
|
||||
}
|
||||
if alias != nil {
|
||||
t.SessionAlias = *alias
|
||||
}
|
||||
if negate {
|
||||
t.Depth = -lvl
|
||||
} else {
|
||||
t.Depth = lvl
|
||||
}
|
||||
t.BodyPreview = preview(t.Body, 240)
|
||||
t.Body = "" // 树视图只要预览，全文按需单取
|
||||
out = append(out, t)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
// preview 按 UTF-8 边界截断正文。
|
||||
// 直接切字节会把多字节字符切成半个，前端渲染出 U+FFFD 替换符。
|
||||
func preview(s string, max int) string {
|
||||
if len(s) <= max {
|
||||
return s
|
||||
}
|
||||
cut := max
|
||||
for cut > 0 && !utf8Start(s[cut]) {
|
||||
cut--
|
||||
}
|
||||
return s[:cut] + "..."
|
||||
}
|
||||
|
||||
// utf8Start 判断某字节是否为一个 UTF-8 序列的首字节
|
||||
func utf8Start(b byte) bool { return b&0xC0 != 0x80 }
|
||||
23
gateway/internal/repo/thread_test.go
Normal file
23
gateway/internal/repo/thread_test.go
Normal file
@ -0,0 +1,23 @@
|
||||
package repo
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestPreviewTruncatesOnUTF8Boundary(t *testing.T) {
|
||||
// 「巡」是 3 字节；在 max=4 处切会切进第 2 个字符中间
|
||||
s := "巡检报告"
|
||||
got := preview(s, 4)
|
||||
if got != "巡..." {
|
||||
t.Fatalf("按 UTF-8 边界截断失败：%q", got)
|
||||
}
|
||||
for i, r := range got {
|
||||
if r == 0xFFFD {
|
||||
t.Fatalf("位置 %d 出现替换符，说明切在了字符中间", i)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestPreviewKeepsShortBodyIntact(t *testing.T) {
|
||||
if got := preview("短正文", 240); got != "短正文" {
|
||||
t.Fatalf("未超长却被改动：%q", got)
|
||||
}
|
||||
}
|
||||
512
gateway/internal/repo/users.go
Normal file
512
gateway/internal/repo/users.go
Normal file
@ -0,0 +1,512 @@
|
||||
package repo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"database/sql"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/agentmail/gateway/internal/db"
|
||||
"github.com/agentmail/gateway/internal/models"
|
||||
"github.com/google/uuid"
|
||||
"golang.org/x/crypto/bcrypt"
|
||||
)
|
||||
|
||||
const (
|
||||
bcryptCost = 12
|
||||
sessionTTL = 7 * 24 * time.Hour
|
||||
userSelectCols = `user_id, username, display_name, password_hash, role, status, created_at, last_login, allowed_agents, allowed_paths`
|
||||
)
|
||||
|
||||
var (
|
||||
ErrUserNotFound = errors.New("user not found")
|
||||
ErrBadCredentials = errors.New("invalid username or password")
|
||||
ErrUserDisabled = errors.New("user disabled")
|
||||
ErrNameTaken = errors.New("name already taken by an agent or user")
|
||||
ErrSessionInvalid = errors.New("session invalid or expired")
|
||||
ErrInvalidUsername = errors.New("username must be 2-64 chars of [a-z0-9._-]")
|
||||
ErrAlreadySetup = errors.New("system already initialized")
|
||||
)
|
||||
|
||||
// ---------- 命名空间校验 ----------
|
||||
|
||||
// 三维地址的 name 位由人类用户与 Agent 共用，因此必须全局唯一
|
||||
func nameTaken(ctx context.Context, name string) (bool, error) {
|
||||
var n int
|
||||
err := db.DB.QueryRowContext(ctx, `
|
||||
SELECT (SELECT COUNT(*) FROM users WHERE username = $1)
|
||||
+ (SELECT COUNT(*) FROM agents WHERE agent_name = $1)
|
||||
`, name).Scan(&n)
|
||||
return n > 0, err
|
||||
}
|
||||
|
||||
// AgentNameAvailable 供 Agent 注册前校验（不与人类用户重名）
|
||||
func AgentNameAvailable(ctx context.Context, agentName string) (bool, error) {
|
||||
var n int
|
||||
err := db.DB.QueryRowContext(ctx,
|
||||
`SELECT COUNT(*) FROM users WHERE username = $1`, agentName).Scan(&n)
|
||||
return n == 0, err
|
||||
}
|
||||
|
||||
func validUsername(name string) bool {
|
||||
if len(name) < 2 || len(name) > 64 {
|
||||
return false
|
||||
}
|
||||
for _, r := range name {
|
||||
ok := (r >= 'a' && r <= 'z') || (r >= '0' && r <= '9') || r == '.' || r == '_' || r == '-'
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
}
|
||||
// 保留字：human 是兼容别名，不能被真实用户占用
|
||||
return name != "human"
|
||||
}
|
||||
|
||||
// ---------- User CRUD ----------
|
||||
|
||||
func scanUser(row *sql.Row) (*models.User, error) {
|
||||
var u models.User
|
||||
var agentsJSON, pathsJSON []byte
|
||||
err := row.Scan(&u.ID, &u.Username, &u.DisplayName, &u.PasswordHash,
|
||||
&u.Role, &u.Status, &u.CreatedAt, &u.LastLogin, &agentsJSON, &pathsJSON)
|
||||
if err != nil {
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
return nil, ErrUserNotFound
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
u.AllowedAgents = decodeStrList(agentsJSON)
|
||||
u.AllowedPaths = decodeStrList(pathsJSON)
|
||||
return &u, nil
|
||||
}
|
||||
|
||||
func decodeStrList(raw []byte) []string {
|
||||
out := []string{}
|
||||
if len(raw) > 0 {
|
||||
_ = json.Unmarshal(raw, &out)
|
||||
}
|
||||
if out == nil {
|
||||
out = []string{}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func CreateUser(ctx context.Context, username, password, displayName, role string,
|
||||
allowedAgents, allowedPaths []string) (*models.User, error) {
|
||||
username = strings.ToLower(strings.TrimSpace(username))
|
||||
if !validUsername(username) {
|
||||
return nil, ErrInvalidUsername
|
||||
}
|
||||
if role != "admin" {
|
||||
role = "user"
|
||||
}
|
||||
if displayName == "" {
|
||||
displayName = username
|
||||
}
|
||||
|
||||
taken, err := nameTaken(ctx, username)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if taken {
|
||||
return nil, ErrNameTaken
|
||||
}
|
||||
|
||||
hash, err := bcrypt.GenerateFromPassword([]byte(password), bcryptCost)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
agentsJSON, _ := json.Marshal(normalizeList(allowedAgents))
|
||||
pathsJSON, _ := json.Marshal(normalizeList(allowedPaths))
|
||||
|
||||
row := db.DB.QueryRowContext(ctx, `
|
||||
INSERT INTO users (username, display_name, password_hash, role, allowed_agents, allowed_paths)
|
||||
VALUES ($1, $2, $3, $4, $5, $6)
|
||||
RETURNING `+userSelectCols,
|
||||
username, displayName, string(hash), role, agentsJSON, pathsJSON)
|
||||
return scanUser(row)
|
||||
}
|
||||
|
||||
func GetUserByName(ctx context.Context, username string) (*models.User, error) {
|
||||
return scanUser(db.DB.QueryRowContext(ctx,
|
||||
`SELECT `+userSelectCols+` FROM users WHERE username = $1`,
|
||||
strings.ToLower(strings.TrimSpace(username))))
|
||||
}
|
||||
|
||||
func GetUserByID(ctx context.Context, id uuid.UUID) (*models.User, error) {
|
||||
return scanUser(db.DB.QueryRowContext(ctx,
|
||||
`SELECT `+userSelectCols+` FROM users WHERE user_id = $1`, id))
|
||||
}
|
||||
|
||||
func ListUsers(ctx context.Context) ([]models.User, error) {
|
||||
rows, err := db.DB.QueryContext(ctx,
|
||||
`SELECT `+userSelectCols+` FROM users ORDER BY created_at ASC`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
users := []models.User{}
|
||||
for rows.Next() {
|
||||
var u models.User
|
||||
var agentsJSON, pathsJSON []byte
|
||||
if err := rows.Scan(&u.ID, &u.Username, &u.DisplayName, &u.PasswordHash,
|
||||
&u.Role, &u.Status, &u.CreatedAt, &u.LastLogin, &agentsJSON, &pathsJSON); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
u.AllowedAgents = decodeStrList(agentsJSON)
|
||||
u.AllowedPaths = decodeStrList(pathsJSON)
|
||||
users = append(users, u)
|
||||
}
|
||||
return users, nil
|
||||
}
|
||||
|
||||
// UserUpdate 描述一次用户更新；nil 字段表示不改
|
||||
type UserUpdate struct {
|
||||
DisplayName *string
|
||||
Role *string
|
||||
Status *string
|
||||
AllowedAgents *[]string
|
||||
AllowedPaths *[]string
|
||||
}
|
||||
|
||||
func UpdateUser(ctx context.Context, id uuid.UUID, up UserUpdate) (*models.User, error) {
|
||||
var agentsJSON, pathsJSON *string
|
||||
if up.AllowedAgents != nil {
|
||||
b, _ := json.Marshal(normalizeList(*up.AllowedAgents))
|
||||
s := string(b)
|
||||
agentsJSON = &s
|
||||
}
|
||||
if up.AllowedPaths != nil {
|
||||
b, _ := json.Marshal(normalizeList(*up.AllowedPaths))
|
||||
s := string(b)
|
||||
pathsJSON = &s
|
||||
}
|
||||
|
||||
row := db.DB.QueryRowContext(ctx, `
|
||||
UPDATE users SET
|
||||
display_name = COALESCE($2, display_name),
|
||||
role = COALESCE($3, role),
|
||||
status = COALESCE($4, status),
|
||||
allowed_agents = COALESCE($5`+db.JSONCast()+`, allowed_agents),
|
||||
allowed_paths = COALESCE($6`+db.JSONCast()+`, allowed_paths)
|
||||
WHERE user_id = $1
|
||||
RETURNING `+userSelectCols,
|
||||
id, up.DisplayName, up.Role, up.Status, agentsJSON, pathsJSON)
|
||||
return scanUser(row)
|
||||
}
|
||||
|
||||
// normalizeList 去空白、去空项、去重，保持顺序
|
||||
func normalizeList(in []string) []string {
|
||||
out := make([]string, 0, len(in))
|
||||
seen := map[string]bool{}
|
||||
for _, s := range in {
|
||||
s = strings.TrimSpace(s)
|
||||
if s == "" || seen[s] {
|
||||
continue
|
||||
}
|
||||
seen[s] = true
|
||||
out = append(out, s)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func SetPassword(ctx context.Context, id uuid.UUID, newPassword string) error {
|
||||
hash, err := bcrypt.GenerateFromPassword([]byte(newPassword), bcryptCost)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
tag, err := db.DB.ExecContext(ctx,
|
||||
`UPDATE users SET password_hash = $2 WHERE user_id = $1`, id, string(hash))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if n, _ := tag.RowsAffected(); n == 0 {
|
||||
return ErrUserNotFound
|
||||
}
|
||||
// 改密后踢掉该用户所有会话
|
||||
_, _ = db.DB.ExecContext(ctx, `DELETE FROM user_sessions WHERE user_id = $1`, id)
|
||||
return nil
|
||||
}
|
||||
|
||||
func DisableUser(ctx context.Context, id uuid.UUID) error {
|
||||
tag, err := db.DB.ExecContext(ctx,
|
||||
`UPDATE users SET status = 'disabled' WHERE user_id = $1`, id)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if n, _ := tag.RowsAffected(); n == 0 {
|
||||
return ErrUserNotFound
|
||||
}
|
||||
_, _ = db.DB.ExecContext(ctx, `DELETE FROM user_sessions WHERE user_id = $1`, id)
|
||||
return nil
|
||||
}
|
||||
|
||||
func CountAdmins(ctx context.Context) (int, error) {
|
||||
var n int
|
||||
err := db.DB.QueryRowContext(ctx,
|
||||
`SELECT COUNT(*) FROM users WHERE role = 'admin' AND status = 'active'`).Scan(&n)
|
||||
return n, err
|
||||
}
|
||||
|
||||
// EnsureAdminUser 首次启动时创建默认管理员（幂等）
|
||||
func EnsureAdminUser(ctx context.Context, username, password string) (*models.User, bool, error) {
|
||||
if n, err := CountAdmins(ctx); err != nil {
|
||||
return nil, false, err
|
||||
} else if n > 0 {
|
||||
u, err := GetUserByName(ctx, username)
|
||||
if err != nil && !errors.Is(err, ErrUserNotFound) {
|
||||
return nil, false, err
|
||||
}
|
||||
return u, false, nil
|
||||
}
|
||||
u, err := CreateUser(ctx, username, password, "管理员", "admin", nil, nil)
|
||||
if err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
return u, true, nil
|
||||
}
|
||||
|
||||
// ---------- 登录 / 会话令牌 ----------
|
||||
|
||||
func Authenticate(ctx context.Context, username, password string) (*models.User, error) {
|
||||
u, err := GetUserByName(ctx, username)
|
||||
if err != nil {
|
||||
if errors.Is(err, ErrUserNotFound) {
|
||||
// 统一错误，避免暴露用户是否存在
|
||||
return nil, ErrBadCredentials
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
if u.Status != "active" {
|
||||
return nil, ErrUserDisabled
|
||||
}
|
||||
if bcrypt.CompareHashAndPassword([]byte(u.PasswordHash), []byte(password)) != nil {
|
||||
return nil, ErrBadCredentials
|
||||
}
|
||||
_, _ = db.DB.ExecContext(ctx, `UPDATE users SET last_login = NOW() WHERE user_id = $1`, u.ID)
|
||||
return u, nil
|
||||
}
|
||||
|
||||
func newToken() (string, error) {
|
||||
b := make([]byte, 32)
|
||||
if _, err := rand.Read(b); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return hex.EncodeToString(b), nil
|
||||
}
|
||||
|
||||
func CreateUserSession(ctx context.Context, userID uuid.UUID, userAgent string) (string, time.Time, error) {
|
||||
token, err := newToken()
|
||||
if err != nil {
|
||||
return "", time.Time{}, err
|
||||
}
|
||||
expires := time.Now().Add(sessionTTL)
|
||||
if len(userAgent) > 256 {
|
||||
userAgent = userAgent[:256]
|
||||
}
|
||||
_, err = db.DB.ExecContext(ctx, `
|
||||
INSERT INTO user_sessions (token, user_id, expires_at, user_agent)
|
||||
VALUES ($1, $2, $3, $4)`, token, userID, expires, userAgent)
|
||||
if err != nil {
|
||||
return "", time.Time{}, err
|
||||
}
|
||||
// 顺手清理过期令牌
|
||||
_, _ = db.DB.ExecContext(ctx, `DELETE FROM user_sessions WHERE expires_at < NOW()`)
|
||||
return token, expires, nil
|
||||
}
|
||||
|
||||
// ResolveUserSession 校验令牌并滑动续期
|
||||
func ResolveUserSession(ctx context.Context, token string) (*models.User, error) {
|
||||
if token == "" {
|
||||
return nil, ErrSessionInvalid
|
||||
}
|
||||
var userID uuid.UUID
|
||||
err := db.DB.QueryRowContext(ctx, `
|
||||
SELECT user_id FROM user_sessions
|
||||
WHERE token = $1 AND expires_at > NOW()`, token).Scan(&userID)
|
||||
if err != nil {
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
return nil, ErrSessionInvalid
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
|
||||
u, err := GetUserByID(ctx, userID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if u.Status != "active" {
|
||||
return nil, ErrUserDisabled
|
||||
}
|
||||
|
||||
_, _ = db.DB.ExecContext(ctx,
|
||||
`UPDATE user_sessions SET expires_at = $2 WHERE token = $1`,
|
||||
token, time.Now().Add(sessionTTL))
|
||||
return u, nil
|
||||
}
|
||||
|
||||
func DeleteUserSession(ctx context.Context, token string) error {
|
||||
_, err := db.DB.ExecContext(ctx, `DELETE FROM user_sessions WHERE token = $1`, token)
|
||||
return err
|
||||
}
|
||||
|
||||
// ---------- 人类用户候选（供地址补全） ----------
|
||||
|
||||
func ListActiveUsernames(ctx context.Context) ([]string, error) {
|
||||
rows, err := db.DB.QueryContext(ctx,
|
||||
`SELECT username FROM users WHERE status = 'active' ORDER BY username`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
out := []string{}
|
||||
for rows.Next() {
|
||||
var s string
|
||||
if err := rows.Scan(&s); err == nil {
|
||||
out = append(out, s)
|
||||
}
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// ---------- 会话归属 ----------
|
||||
|
||||
func SetSessionOwner(ctx context.Context, sessionID, userID uuid.UUID) error {
|
||||
_, err := db.DB.ExecContext(ctx,
|
||||
`UPDATE sessions SET owner_user_id = $2 WHERE session_id = $1 AND owner_user_id IS NULL`,
|
||||
sessionID, userID)
|
||||
return err
|
||||
}
|
||||
|
||||
// SessionOwnerUsername 返回会话归属人类用户名；无归属时返回空串
|
||||
func SessionOwnerUsername(ctx context.Context, sessionID uuid.UUID) (string, error) {
|
||||
var name *string
|
||||
err := db.DB.QueryRowContext(ctx, `
|
||||
SELECT u.username
|
||||
FROM sessions s LEFT JOIN users u ON u.user_id = s.owner_user_id
|
||||
WHERE s.session_id = $1`, sessionID).Scan(&name)
|
||||
if err != nil {
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
return "", fmt.Errorf("session %s not found", sessionID)
|
||||
}
|
||||
return "", err
|
||||
}
|
||||
if name == nil {
|
||||
return "", nil
|
||||
}
|
||||
return *name, nil
|
||||
}
|
||||
|
||||
// UserCanAccessSession 判断用户能否访问该会话：owner、或在邮件收发/抄送中出现，或 admin
|
||||
func UserCanAccessSession(ctx context.Context, u *models.User, sessionID uuid.UUID) (bool, error) {
|
||||
if u.IsAdmin() {
|
||||
return true, nil
|
||||
}
|
||||
var n int
|
||||
err := db.DB.QueryRowContext(ctx, `
|
||||
SELECT COUNT(*) FROM sessions s
|
||||
WHERE s.session_id = $1
|
||||
AND (s.owner_user_id = $2
|
||||
OR EXISTS (
|
||||
SELECT 1 FROM mails m
|
||||
WHERE m.session_id = s.session_id
|
||||
AND (m.from_name = $3 OR m.to_name = $3
|
||||
OR `+db.CCHas("m.cc_list", 3)+`)
|
||||
))
|
||||
`, sessionID, u.ID, u.Username).Scan(&n)
|
||||
return n > 0, err
|
||||
}
|
||||
|
||||
// FirstAdminUsername 返回最早创建的可用管理员用户名（用于无归属会话的兜底决策人）
|
||||
func FirstAdminUsername(ctx context.Context) (string, error) {
|
||||
var name string
|
||||
err := db.DB.QueryRowContext(ctx, `
|
||||
SELECT username FROM users
|
||||
WHERE role = 'admin' AND status = 'active'
|
||||
ORDER BY created_at ASC LIMIT 1`).Scan(&name)
|
||||
if err != nil {
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
return "", nil
|
||||
}
|
||||
return "", err
|
||||
}
|
||||
return name, nil
|
||||
}
|
||||
|
||||
// RandomPassword 生成一个随机初始密码（首次启动无 ADMIN_PASSWORD 时使用）
|
||||
func RandomPassword(n int) string {
|
||||
const charset = "abcdefghijkmnopqrstuvwxyzABCDEFGHJKLMNPQRSTUVWXYZ23456789"
|
||||
b := make([]byte, n)
|
||||
if _, err := rand.Read(b); err != nil {
|
||||
return "ChangeMe" + fmt.Sprint(time.Now().Unix())
|
||||
}
|
||||
for i := range b {
|
||||
b[i] = charset[int(b[i])%len(charset)]
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
|
||||
// ---------- Setup（首次初始化管理员） ----------
|
||||
|
||||
// NeedsSetup 返回系统是否尚未初始化（没有任何用户）
|
||||
func NeedsSetup(ctx context.Context) (bool, error) {
|
||||
var n int
|
||||
err := db.DB.QueryRowContext(ctx, `SELECT COUNT(*) FROM users`).Scan(&n)
|
||||
return n == 0, err
|
||||
}
|
||||
|
||||
// SetupFirstAdmin 在系统尚无任何用户时创建首个管理员。
|
||||
// 已初始化时返回 ErrAlreadySetup，避免被用作后门。
|
||||
func SetupFirstAdmin(ctx context.Context, username, password, displayName string) (*models.User, error) {
|
||||
empty, err := NeedsSetup(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if !empty {
|
||||
return nil, ErrAlreadySetup
|
||||
}
|
||||
if displayName == "" {
|
||||
displayName = username
|
||||
}
|
||||
return CreateUser(ctx, username, password, displayName, "admin", nil, nil)
|
||||
}
|
||||
|
||||
// ---------- 可选目录候选（供权限设置界面） ----------
|
||||
|
||||
// AllWorkspaceNames 汇总所有 Agent 注册过的工作区名，供管理员挑选可访问目录
|
||||
func AllWorkspaceNames(ctx context.Context) ([]string, error) {
|
||||
rows, err := db.DB.QueryContext(ctx, `
|
||||
SELECT DISTINCT ws->>'name' AS name
|
||||
FROM agents, jsonb_array_elements(workspaces) AS ws
|
||||
WHERE COALESCE(ws->>'name', '') <> ''
|
||||
ORDER BY name`)
|
||||
if err != nil {
|
||||
return []string{}, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
out := []string{}
|
||||
for rows.Next() {
|
||||
var s string
|
||||
if err := rows.Scan(&s); err == nil {
|
||||
out = append(out, s)
|
||||
}
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// IsHumanUser 判断某个三维地址 name 位是否为人类用户
|
||||
func IsHumanUser(ctx context.Context, name string) (bool, error) {
|
||||
var n int
|
||||
err := db.DB.QueryRowContext(ctx,
|
||||
`SELECT COUNT(*) FROM users WHERE username = $1`, name).Scan(&n)
|
||||
return n > 0, err
|
||||
}
|
||||
Reference in New Issue
Block a user