MailUI4Agents/gateway/cmd/server/main.go

package main

import (
	"context"
	"fmt"
	"log"
	"net/http"
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/agentmail/gateway/internal/blob"
	"github.com/agentmail/gateway/internal/config"
	"github.com/agentmail/gateway/internal/db"
	"github.com/agentmail/gateway/internal/handler"
	"github.com/agentmail/gateway/internal/middleware"
	"github.com/agentmail/gateway/internal/repo"
	"github.com/agentmail/gateway/internal/static"
	"github.com/go-chi/chi/v5"
	chimw "github.com/go-chi/chi/v5/middleware"
	"github.com/go-chi/cors"
)

func main() {
	cfg := config.Load()

	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
	defer cancel()

	if err := db.Connect(ctx, cfg.DatabaseURL); err != nil {
		log.Fatalf("Database connection failed: %v", err)
	}
	defer db.Close()

	// 第一轮迁移：建表（users 必须先存在，'human' 数据迁移才能找到管理员）
	if err := db.Migrate(ctx); err != nil {
		log.Fatalf("Migration failed: %v", err)
	}

	// 确保存在默认管理员
	bootstrapAdmin(ctx, cfg)

	// 第二轮迁移：此时管理员已存在，历史 'human' 字面量得以重写
	// （仅 PostgreSQL 有该历史包袹；SQLite 是新后端，这一轮是幂等的建表重跑）
	if err := db.Migrate(ctx); err != nil {
		log.Fatalf("Post-admin migration failed: %v", err)
	}

	// 附件存储：内容存盘，数据库只存元数据
	blobs, err := blob.New(cfg.AttachmentDir)
	if err != nil {
		log.Fatalf("Attachment store failed: %v", err)
	}
	handler.Blobs = blobs
	fmt.Printf("附件存储：%s（单个上限 %.0f MB）\n",
		blobs.Root(), float64(cfg.MaxAttachmentBytes)/(1<<20))

	// 后台 GC：清掉上传后未随邮件发出的孤立附件，
	// 否则取消发信与 Agent 崩溃留下的文件会让磁盘单调增长。
	go sweepOrphanAttachments(blobs)

	r := chi.NewRouter()

	r.Use(chimw.Logger)
	r.Use(chimw.Recoverer)
	r.Use(chimw.RequestID)
	r.Use(chimw.RealIP)
	r.Use(cors.Handler(cors.Options{
		AllowedOrigins: cfg.CORSOrigins,
		AllowedMethods: []string{"GET", "POST", "PUT", "DELETE", "OPTIONS"},
		AllowedHeaders: []string{"Accept", "Authorization", "Content-Type", "X-Agent-Name", "X-Agent-Secret"},
		// 附件下载靠 Content-Disposition 拿文件名；不暂存就拿不到。
		// Content-Length 给进度条用。
		ExposedHeaders:   []string{"Link", "Content-Disposition", "Content-Length"},
		AllowCredentials: true,
		MaxAge:           300,
	}))

	r.Get("/health", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		w.Write([]byte(`{"status":"ok"}`))
	})

	r.Route("/api/v1", func(r chi.Router) {
		// ---- 首次初始化（公开；仅系统无用户时可用） ----
		r.Get("/setup/status", handler.SetupStatus)
		r.Post("/setup/admin", handler.SetupAdmin)

		// ---- 认证（公开） ----
		r.Post("/auth/login", handler.Login)
		r.Post("/auth/logout", handler.Logout)

		// ---- Agent 注册（凭 secret，非人类登录态） ----
		r.Post("/agent/register", handler.RegisterAgent)

		// ---- Agent 侧（X-Agent-Name + X-Agent-Secret） ----
		r.Group(func(r chi.Router) {
			r.Use(middleware.AgentAuth)
			r.Post("/agent/heartbeat", handler.HeartbeatAgent)
			r.Post("/mail/send", handler.SendMail)
			r.Get("/mail/inbox", handler.GetInbox)
			// 批量标记已读：不给 mail_ids 就把收件箱全部未读标掉。
			// 没有它的话 Agent 每次拉收件箱都会重复捞同一批旧邮件。
			r.Post("/mail/read", handler.MarkInboxRead)
			r.Post("/mail/{id}/forward", handler.ForwardMail)
			r.Post("/permission/request", handler.RequestPermission)
			// 附件：先上传拿 id，再在发信时放进 attachment_ids
			r.Post("/attachments", handler.UploadAttachment)
			r.Get("/attachments/{id}", handler.DownloadAttachment)
			// 平台侧会话标题/slug 回写本侧（平台叫什么，本侧就叫什么）
			r.Post("/sessions/{id}/sync", handler.SyncSession)
			// 邮件场景下的可用模型范围。上报走心跳（agent/heartbeat 的 models 字段），
			// 这里只读 —— 给非插件的第三方客户端与排查用。
			r.Get("/agent/models/allowed", handler.GetAllowedModels)
		})

		// ---- 人类登录态 ----
		r.Group(func(r chi.Router) {
			r.Use(middleware.UserAuth)

			r.Get("/auth/me", handler.Me)
			r.Post("/auth/password", handler.ChangePassword)

			// 自己的邮箱
			r.Post("/me/mail/send", handler.MeSendMail)
			r.Get("/me/mail/inbox", handler.MeGetInbox)
			r.Get("/me/mail/sent", handler.MeGetSent)
			r.Get("/me/sessions", handler.MeGetSessions)
			r.Post("/me/mail/{id}/forward", handler.MeForwardMail)

			// 附件
			r.Post("/me/attachments", handler.MeUploadAttachment)
			r.Delete("/me/attachments/{id}", handler.MeDeleteAttachment)

			// 自己的客户端连接密钥（仅能用于 /me/* 与会话级接口，不可注册 Agent）
			r.Post("/me/keys", handler.CreateMyKey)
			r.Get("/me/keys", handler.ListMyKeys)
			r.Delete("/me/keys/{id}", handler.DeleteMyKey)

			// 邮件/会话（带会话级鉴权）
			r.Get("/mail/{id}", handler.GetMail)
			r.Get("/mail/{id}/thread", handler.GetMailThread)
			r.Post("/mail/{id}/read", handler.MarkMailRead)
			r.Get("/sessions/{id}", handler.GetSession)
			r.Get("/sessions/{id}/mails", handler.GetSessionMails)
			r.Put("/sessions/{id}/alias", handler.UpdateSessionAlias)
			// 本任务的往返预算：在对话页里随时可改
			r.Get("/sessions/{id}/budget", handler.GetSessionBudgetHandler)
			r.Put("/sessions/{id}/budget", handler.UpdateSessionBudget)
			// Agent 在正文里提的改名建议：读取与驳回（接受走上面的 PUT alias）
			r.Get("/sessions/{id}/rename-proposal", handler.GetRenameProposal)
			r.Post("/sessions/{id}/rename-proposal/dismiss", handler.DismissRenameProposal)

			// 联系人 name@path.session
			r.Get("/contacts", handler.ListContacts)
			r.Get("/contacts/suggest", handler.SuggestAddress)
			r.Post("/contacts/archive", handler.ArchiveContact)

			// 权限决策
			r.Post("/permission/decide", handler.DecidePermission)
			r.Get("/permission/pending", handler.ListPendingPermissions)

			// 在线 Agent 列表（补全用）
			r.Get("/agents", handler.ListAgents)

			// 管理员
			r.Group(func(r chi.Router) {
				r.Use(middleware.AdminOnly)
				r.Get("/admin/users", handler.AdminListUsers)
				r.Post("/admin/users", handler.AdminCreateUser)
				r.Put("/admin/users/{id}", handler.AdminUpdateUser)
				r.Delete("/admin/users/{id}", handler.AdminDisableUser)
				r.Post("/admin/users/{id}/reset", handler.AdminResetPassword)
				r.Get("/admin/scopes", handler.AdminListScopes)

				// Agent 接入密钥
				r.Post("/admin/agent-keys", handler.CreateAgentKey)
				r.Get("/admin/agent-keys", handler.ListAgentKeys)
				r.Delete("/admin/agent-keys/{id}", handler.DeleteAgentKey)
				r.Post("/admin/agent-keys/{id}/bind", handler.BindAgentKey)

				// Agent 发信配额
				r.Get("/admin/quotas", handler.AdminListQuotas)
				r.Put("/admin/quotas/{name}", handler.AdminSetQuota)

				// 邮件场景下每个 Agent 可用的模型范围（勾选平台上报的目录）
				r.Get("/admin/agents/{name}/models", handler.AdminListAgentModels)
				r.Put("/admin/agents/{name}/models", handler.AdminSetAgentModels)
			})
		})

		// ---- SSE：Agent 走 header，人类走 Cookie，内部自行分流 ----
		r.Get("/events/stream", handler.SSEStream)
		r.Get("/events/status", handler.SSEStatus)

		// ---- 附件下载：由浏览器直接发起（<a download>），无法带 Authorization 头，
		// 因此单独挂在允许 ?access_token= 的中间件下 ----
		r.Group(func(r chi.Router) {
			r.Use(middleware.UserAuthAllowQueryToken)
			r.Get("/me/attachments/{id}", handler.MeDownloadAttachment)
		})
	})

	// 静态前端
	staticRoot := os.Getenv("STATIC_DIR")
	var staticMux http.Handler
	if staticRoot != "" {
		staticMux = http.FileServer(http.Dir(staticRoot))
	} else {
		staticMux = static.Handler()
	}
	r.Handle("/assets/*", staticMux)
	r.Get("/", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/html; charset=utf-8")
		w.Write(static.GetIndex())
	})

	addr := fmt.Sprintf(":%s", cfg.Port)
	srv := &http.Server{
		Addr:         addr,
		Handler:      r,
		ReadTimeout:  15 * time.Second,
		WriteTimeout: 0, // SSE 长连接
		IdleTimeout:  120 * time.Second,
	}

	go func() {
		sigCh := make(chan os.Signal, 1)
		signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
		<-sigCh
		fmt.Println("\nShutting down...")
		c, cancel := context.WithTimeout(context.Background(), 10*time.Second)
		defer cancel()
		srv.Shutdown(c)
	}()

	fmt.Printf("AgentMail Gateway on %s\n", addr)
	fmt.Printf("  Health: http://localhost%s/health\n", addr)
	fmt.Printf("  API:    http://localhost%s/api/v1\n", addr)

	if err := srv.ListenAndServe(); err != http.ErrServerClosed {
		log.Fatalf("Server error: %v", err)
	}
	fmt.Println("Server stopped")
}

// bootstrapAdmin 首次启动时创建默认管理员
func bootstrapAdmin(ctx context.Context, cfg *config.Config) {
	n, err := repo.CountAdmins(ctx)
	if err != nil {
		log.Fatalf("Failed to count admins: %v", err)
	}
	if n > 0 {
		return
	}

	pw := cfg.AdminPassword
	generated := false
	if pw == "" {
		pw = repo.RandomPassword(16)
		generated = true
	}

	u, created, err := repo.EnsureAdminUser(ctx, cfg.AdminUser, pw)
	if err != nil {
		log.Fatalf("Failed to create admin user: %v", err)
	}
	if created && u != nil {
		fmt.Println("========================================")
		fmt.Printf("  已创建默认管理员: %s\n", u.Username)
		if generated {
			fmt.Printf("  初始密码（仅本次显示）: %s\n", pw)
			fmt.Println("  请登录后立即通过 /auth/password 修改")
		} else {
			fmt.Println("  密码来自环境变量 ADMIN_PASSWORD")
		}
		fmt.Println("========================================")
	}
}

// sweepOrphanAttachments 周期清理「已上传但从未随邮件发出」的附件。
//
// 上传与发信是两步，中间放弃（用户取消写信、Agent 崩溃）就会留下孤立记录与文件。
// 保留 24 小时再清：足以覆盖一次正常的写信过程，也不至于让废弃文件长期占盘。
func sweepOrphanAttachments(blobs *blob.Store) {
	const (
		interval = 1 * time.Hour
		keepFor  = 24 * time.Hour
	)

	sweep := func() {
		ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
		defer cancel()

		sums, err := repo.SweepOrphanAttachments(ctx, keepFor)
		if err != nil {
			log.Printf("附件 GC 失败: %v", err)
			return
		}
		for _, sum := range sums {
			if err := blobs.Remove(sum); err != nil {
				log.Printf("附件 GC 删除 %s 失败: %v", sum[:8], err)
			}
		}
		if len(sums) > 0 {
			log.Printf("附件 GC 清理了 %d 个孤立文件", len(sums))
		}
	}

	// 启动时先扫一遍：上次进程可能是被 kill 掉的，留下的孤立文件不该等一小时
	sweep()
	for range time.Tick(interval) {
		sweep()
	}
}