feat(govai): 0617 优化首批 — 安全/私有化/深度研究/服务层/可观测性

借鉴 odysseus 的能力设计，全程净室实现、零 AGPL 代码、不引入 AGPL 依赖。

T1 提示注入防护: pkg/promptguard 包裹外部/知识库内容为不可信数据，buildMessages 移出 system 指令区。 T2 安全 CI: .github/workflows(ci+security: govulncheck/gitleaks/actionlint/hadolint/trivy)+dependabot+.hadolint.yaml；go.mod 加 toolchain go1.25.11 修复 20 个 stdlib CVE。 T3 管理员 2FA: 迁移 000016 + RFC6238 TOTP/备份码(pkg/auth, 零依赖) + 登录流程集成(后端)。 T4 本地模型: LLM/embedding 支持本地 vLLM/Ollama(OpenAI 兼容, 鉴权头条件发送, NoAuth) + docs/local-deploy.md。 T6 深度研究: 迁移 000017 + Python research-worker(净室多步流水线, 检索避开 SearXNG) + Go research 服务/handler/路由。 T7 service 层: 新增 internal/service/{research,twofa}, 2FA 业务逻辑从胖 handler 下沉, 接口注入可单测。 T10 缓存/可观测性: internal/cache(Redis+内存, 优雅降级) 接入 store 热点列表; Prometheus 指标+/metrics; docs/openapi.yaml。 验证: go build/vet/test ./... 全绿(8 包); research-worker 12 单测过; 真实 PG 应用迁移并烟测。

server/internal/service/research/redis.go

@@ -0,0 +1,42 @@
+package research
+
+import (
+	"context"
+	"encoding/json"
+	"strconv"
+
+	"github.com/redis/go-redis/v9"
+)
+
+const (
+	taskQueueKey   = "research:tasks"
+	statusKeyPrefix = "research:status:"
+)
+
+// redisBackend 同时实现 Queue 与 Cache：任务下发到列表队列、快速状态读自 hash。
+// 与 research-worker 的 TASK_QUEUE / TASK_STATUS_PREFIX 约定保持一致。
+type redisBackend struct {
+	rdb *redis.Client
+}
+
+func NewRedisBackend(rdb *redis.Client) *redisBackend {
+	return &redisBackend{rdb: rdb}
+}
+
+func (b *redisBackend) Enqueue(ctx context.Context, taskID string) error {
+	msg, _ := json.Marshal(map[string]string{"task_id": taskID})
+	return b.rdb.LPush(ctx, taskQueueKey, msg).Err()
+}
+
+func (b *redisBackend) GetStatus(ctx context.Context, taskID string) (*CachedStatus, bool) {
+	m, err := b.rdb.HGetAll(ctx, statusKeyPrefix+taskID).Result()
+	if err != nil || len(m) == 0 {
+		return nil, false
+	}
+	progress, _ := strconv.Atoi(m["progress"])
+	return &CachedStatus{
+		Status:   m["status"],
+		Progress: progress,
+		Message:  m["message"],
+	}, true
+}

server/internal/service/research/repository.go

@@ -0,0 +1,83 @@
+package research
+
+import (
+	"context"
+	"encoding/json"
+
+	"github.com/jackc/pgx/v5/pgxpool"
+)
+
+// pgxRepository 是基于 pgx 连接池的 Repository 实现。
+type pgxRepository struct {
+	pool *pgxpool.Pool
+}
+
+func NewPgxRepository(pool *pgxpool.Pool) Repository {
+	return &pgxRepository{pool: pool}
+}
+
+func (r *pgxRepository) Insert(ctx context.Context, id, userID string, appID *string, topic string, config map[string]any) error {
+	if config == nil {
+		config = map[string]any{}
+	}
+	configJSON, err := json.Marshal(config)
+	if err != nil {
+		return err
+	}
+	_, err = r.pool.Exec(ctx,
+		`INSERT INTO research_tasks (id, user_id, app_id, topic, config)
+		 VALUES ($1, $2, $3, $4, $5)`,
+		id, userID, appID, topic, configJSON,
+	)
+	return err
+}
+
+func (r *pgxRepository) Get(ctx context.Context, userID, taskID string) (*Task, error) {
+	var t Task
+	var sources []byte
+	err := r.pool.QueryRow(ctx,
+		`SELECT id, topic, status, progress, status_message, error_message, report, sources, tokens_used, created_at
+		 FROM research_tasks WHERE id = $1 AND user_id = $2`, taskID, userID,
+	).Scan(&t.ID, &t.Topic, &t.Status, &t.Progress, &t.StatusMessage, &t.ErrorMessage,
+		&t.Report, &sources, &t.TokensUsed, &t.CreatedAt)
+	if err != nil {
+		return nil, err
+	}
+	t.Sources = sources
+	return &t, nil
+}
+
+func (r *pgxRepository) List(ctx context.Context, userID string, limit int) ([]Task, error) {
+	rows, err := r.pool.Query(ctx,
+		`SELECT id, topic, status, progress, status_message, error_message, tokens_used, created_at
+		 FROM research_tasks WHERE user_id = $1 ORDER BY created_at DESC LIMIT $2`, userID, limit,
+	)
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+
+	var tasks []Task
+	for rows.Next() {
+		var t Task
+		if err := rows.Scan(&t.ID, &t.Topic, &t.Status, &t.Progress, &t.StatusMessage,
+			&t.ErrorMessage, &t.TokensUsed, &t.CreatedAt); err != nil {
+			continue
+		}
+		tasks = append(tasks, t)
+	}
+	return tasks, nil
+}
+
+func (r *pgxRepository) Cancel(ctx context.Context, userID, taskID string) (bool, error) {
+	tag, err := r.pool.Exec(ctx,
+		`UPDATE research_tasks SET status = 'canceled', updated_at = NOW()
+		 WHERE id = $1 AND user_id = $2
+		   AND status IN ('pending','planning','searching','reading','synthesizing')`,
+		taskID, userID,
+	)
+	if err != nil {
+		return false, err
+	}
+	return tag.RowsAffected() > 0, nil
+}

server/internal/service/research/service.go

@@ -0,0 +1,130 @@
+// Package research 提供深度研究任务的业务编排（service 层）。
+//
+// 该层与具体存储/队列解耦：依赖 Repository（任务持久化）、Queue（任务下发）、
+// Cache（快速状态）三个接口，便于单测与替换实现。HTTP handler 仅做参数解析与
+// 响应，业务规则集中在这里。
+package research
+
+import (
+	"context"
+	"errors"
+	"strings"
+	"time"
+
+	"github.com/google/uuid"
+)
+
+// Task 是研究任务的领域模型（用于查询/列表返回）。
+type Task struct {
+	ID            string
+	Topic         string
+	Status        string
+	Progress      int
+	StatusMessage *string
+	ErrorMessage  *string
+	Report        *string
+	Sources       []byte // 原始 JSON（[{title,url,snippet}]）
+	TokensUsed    int
+	CreatedAt     time.Time
+}
+
+// CreateInput 创建研究任务的入参。
+type CreateInput struct {
+	UserID string
+	AppID  *string
+	Topic  string
+	Config map[string]any
+}
+
+// CachedStatus 来自 Redis 的快速状态（worker 实时写入）。
+type CachedStatus struct {
+	Status   string
+	Progress int
+	Message  string
+}
+
+// Repository 任务持久化接口。
+type Repository interface {
+	Insert(ctx context.Context, id, userID string, appID *string, topic string, config map[string]any) error
+	Get(ctx context.Context, userID, taskID string) (*Task, error)
+	List(ctx context.Context, userID string, limit int) ([]Task, error)
+	// Cancel 仅取消进行中的任务；found 表示是否有可取消的任务被更新。
+	Cancel(ctx context.Context, userID, taskID string) (found bool, err error)
+}
+
+// Queue 任务下发接口（worker 消费）。
+type Queue interface {
+	Enqueue(ctx context.Context, taskID string) error
+}
+
+// Cache 快速状态读取接口（可选）。
+type Cache interface {
+	GetStatus(ctx context.Context, taskID string) (*CachedStatus, bool)
+}
+
+var (
+	ErrEmptyTopic = errors.New("研究题目不能为空")
+	ErrNotFound   = errors.New("任务不存在")
+)
+
+type Service struct {
+	repo  Repository
+	queue Queue
+	cache Cache // 可为 nil
+}
+
+func NewService(repo Repository, queue Queue, cache Cache) *Service {
+	return &Service{repo: repo, queue: queue, cache: cache}
+}
+
+// Create 校验入参、落库并下发到队列，返回任务 ID。
+func (s *Service) Create(ctx context.Context, in CreateInput) (string, error) {
+	topic := strings.TrimSpace(in.Topic)
+	if topic == "" {
+		return "", ErrEmptyTopic
+	}
+	id := uuid.New().String()
+	if err := s.repo.Insert(ctx, id, in.UserID, in.AppID, topic, in.Config); err != nil {
+		return "", err
+	}
+	if err := s.queue.Enqueue(ctx, id); err != nil {
+		return "", err
+	}
+	return id, nil
+}
+
+// Status 返回任务（先按所有权从库中取，再用 Redis 快速状态覆盖以保证新鲜度）。
+func (s *Service) Status(ctx context.Context, userID, taskID string) (*Task, error) {
+	t, err := s.repo.Get(ctx, userID, taskID)
+	if err != nil {
+		return nil, err
+	}
+	if s.cache != nil {
+		if cs, ok := s.cache.GetStatus(ctx, taskID); ok {
+			t.Status = cs.Status
+			t.Progress = cs.Progress
+			if cs.Message != "" {
+				msg := cs.Message
+				t.StatusMessage = &msg
+			}
+		}
+	}
+	return t, nil
+}
+
+// List 返回用户最近的研究任务。
+func (s *Service) List(ctx context.Context, userID string) ([]Task, error) {
+	return s.repo.List(ctx, userID, 50)
+}
+
+// Cancel 取消进行中的任务；任务不存在/不可取消时返回 ErrNotFound。
+func (s *Service) Cancel(ctx context.Context, userID, taskID string) error {
+	found, err := s.repo.Cancel(ctx, userID, taskID)
+	if err != nil {
+		return err
+	}
+	if !found {
+		return ErrNotFound
+	}
+	return nil
+}

server/internal/service/research/service_test.go

@@ -0,0 +1,130 @@
+package research
+
+import (
+	"context"
+	"errors"
+	"testing"
+)
+
+// ---- 测试替身 ----
+
+type fakeRepo struct {
+	inserted   map[string]bool
+	getResult  *Task
+	getErr     error
+	cancelOK   bool
+	cancelErr  error
+	lastInsert struct {
+		id, userID, topic string
+	}
+}
+
+func newFakeRepo() *fakeRepo { return &fakeRepo{inserted: map[string]bool{}} }
+
+func (f *fakeRepo) Insert(ctx context.Context, id, userID string, appID *string, topic string, config map[string]any) error {
+	f.inserted[id] = true
+	f.lastInsert.id = id
+	f.lastInsert.userID = userID
+	f.lastInsert.topic = topic
+	return nil
+}
+func (f *fakeRepo) Get(ctx context.Context, userID, taskID string) (*Task, error) {
+	return f.getResult, f.getErr
+}
+func (f *fakeRepo) List(ctx context.Context, userID string, limit int) ([]Task, error) {
+	return nil, nil
+}
+func (f *fakeRepo) Cancel(ctx context.Context, userID, taskID string) (bool, error) {
+	return f.cancelOK, f.cancelErr
+}
+
+type fakeQueue struct{ enqueued []string }
+
+func (q *fakeQueue) Enqueue(ctx context.Context, taskID string) error {
+	q.enqueued = append(q.enqueued, taskID)
+	return nil
+}
+
+type fakeCache struct{ cs *CachedStatus }
+
+func (c *fakeCache) GetStatus(ctx context.Context, taskID string) (*CachedStatus, bool) {
+	if c.cs == nil {
+		return nil, false
+	}
+	return c.cs, true
+}
+
+// ---- 测试 ----
+
+func TestCreate_EmptyTopicRejected(t *testing.T) {
+	svc := NewService(newFakeRepo(), &fakeQueue{}, nil)
+	_, err := svc.Create(context.Background(), CreateInput{UserID: "u1", Topic: "   "})
+	if !errors.Is(err, ErrEmptyTopic) {
+		t.Fatalf("空题目应返回 ErrEmptyTopic，实际: %v", err)
+	}
+}
+
+func TestCreate_InsertsAndEnqueues(t *testing.T) {
+	repo := newFakeRepo()
+	q := &fakeQueue{}
+	svc := NewService(repo, q, nil)
+
+	id, err := svc.Create(context.Background(), CreateInput{UserID: "u1", Topic: "数字政府研究"})
+	if err != nil {
+		t.Fatalf("Create 出错: %v", err)
+	}
+	if id == "" || !repo.inserted[id] {
+		t.Fatal("应已插入任务记录")
+	}
+	if len(q.enqueued) != 1 || q.enqueued[0] != id {
+		t.Fatalf("应已用相同 id 下发到队列，实际: %v", q.enqueued)
+	}
+	if repo.lastInsert.topic != "数字政府研究" {
+		t.Fatalf("题目透传错误: %q", repo.lastInsert.topic)
+	}
+}
+
+func TestStatus_CacheOverlaysDB(t *testing.T) {
+	repo := newFakeRepo()
+	repo.getResult = &Task{ID: "t1", Status: "pending", Progress: 0}
+	cache := &fakeCache{cs: &CachedStatus{Status: "searching", Progress: 30, Message: "检索中"}}
+	svc := NewService(repo, &fakeQueue{}, cache)
+
+	got, err := svc.Status(context.Background(), "u1", "t1")
+	if err != nil {
+		t.Fatalf("Status 出错: %v", err)
+	}
+	if got.Status != "searching" || got.Progress != 30 {
+		t.Fatalf("缓存状态应覆盖 DB，实际 status=%s progress=%d", got.Status, got.Progress)
+	}
+	if got.StatusMessage == nil || *got.StatusMessage != "检索中" {
+		t.Fatal("应带上缓存的状态消息")
+	}
+}
+
+func TestStatus_DBErrorPropagates(t *testing.T) {
+	repo := newFakeRepo()
+	repo.getErr = errors.New("not found")
+	svc := NewService(repo, &fakeQueue{}, &fakeCache{})
+	if _, err := svc.Status(context.Background(), "u1", "missing"); err == nil {
+		t.Fatal("DB 错误应向上传播")
+	}
+}
+
+func TestCancel_NotFound(t *testing.T) {
+	repo := newFakeRepo()
+	repo.cancelOK = false
+	svc := NewService(repo, &fakeQueue{}, nil)
+	if err := svc.Cancel(context.Background(), "u1", "t1"); !errors.Is(err, ErrNotFound) {
+		t.Fatalf("不可取消时应返回 ErrNotFound，实际: %v", err)
+	}
+}
+
+func TestCancel_Success(t *testing.T) {
+	repo := newFakeRepo()
+	repo.cancelOK = true
+	svc := NewService(repo, &fakeQueue{}, nil)
+	if err := svc.Cancel(context.Background(), "u1", "t1"); err != nil {
+		t.Fatalf("可取消时应返回 nil，实际: %v", err)
+	}
+}

server/internal/service/twofa/service.go

@@ -0,0 +1,139 @@
+// Package twofa 提供两步验证（2FA / TOTP）的业务编排（service 层）。
+//
+// 业务规则（生成密钥/备份码、校验、启用/关闭判定）集中在此，DB 操作通过 Store 接口
+// 注入，便于单测。TOTP/备份码算法复用 pkg/auth。HTTP handler 仅做参数解析与响应。
+package twofa
+
+import (
+	"context"
+	"errors"
+	"time"
+
+	"github.com/enterprise-ai-platform/server/pkg/auth"
+)
+
+// Issuer 显示在认证器 App 中的发行方名称。
+const Issuer = "政智通 GovAI"
+
+// Store 2FA 持久化接口。
+type Store interface {
+	// GetStatus 返回是否启用与剩余可用备份码数量。
+	GetStatus(ctx context.Context, userID string) (enabled bool, remaining int, err error)
+	// GetSecret 返回用户的 TOTP 密钥（可能为空）。
+	GetSecret(ctx context.Context, userID string) (secret string, err error)
+	// GetSecretAndEnabled 返回密钥与启用状态。
+	GetSecretAndEnabled(ctx context.Context, userID string) (secret string, enabled bool, err error)
+	// SaveEnrollment 原子写入新密钥并重置备份码（未启用）。
+	SaveEnrollment(ctx context.Context, userID, secret string, codeHashes []string) error
+	// Enable 置 totp_enabled=true。
+	Enable(ctx context.Context, userID string) error
+	// Disable 关闭 2FA：清除密钥并删除所有备份码。
+	Disable(ctx context.Context, userID string) error
+	// ConsumeBackupCode 校验并一次性消费备份码，命中返回 true。
+	ConsumeBackupCode(ctx context.Context, userID, code string) (bool, error)
+}
+
+var (
+	ErrAlreadyEnabled = errors.New("两步验证已启用")
+	ErrNotSetup       = errors.New("尚未开始两步验证设置")
+	ErrBadCode        = errors.New("验证码或备份码错误")
+)
+
+// EnrollResult 是开始设置 2FA 的返回。
+type EnrollResult struct {
+	Secret      string
+	OtpauthURI  string
+	BackupCodes []string // 明文，仅返回一次
+}
+
+type Service struct {
+	store       Store
+	backupCount int
+	nowUnix     func() int64
+}
+
+func NewService(store Store) *Service {
+	return &Service{store: store, backupCount: 8, nowUnix: func() int64 { return time.Now().Unix() }}
+}
+
+// Status 返回当前 2FA 状态。
+func (s *Service) Status(ctx context.Context, userID string) (enabled bool, remaining int, err error) {
+	return s.store.GetStatus(ctx, userID)
+}
+
+// Enroll 生成新密钥与备份码并落库（未启用，需 Verify 确认）。
+func (s *Service) Enroll(ctx context.Context, userID, email string) (*EnrollResult, error) {
+	_, enabled, err := s.store.GetSecretAndEnabled(ctx, userID)
+	if err != nil {
+		return nil, err
+	}
+	if enabled {
+		return nil, ErrAlreadyEnabled
+	}
+
+	secret, err := auth.GenerateTOTPSecret()
+	if err != nil {
+		return nil, err
+	}
+	plain, hashes, err := auth.GenerateBackupCodes(s.backupCount)
+	if err != nil {
+		return nil, err
+	}
+	if err := s.store.SaveEnrollment(ctx, userID, secret, hashes); err != nil {
+		return nil, err
+	}
+	return &EnrollResult{
+		Secret:      secret,
+		OtpauthURI:  auth.TOTPProvisioningURI(secret, email, Issuer),
+		BackupCodes: plain,
+	}, nil
+}
+
+// EnableAfterVerify 校验首个验证码并启用 2FA。
+func (s *Service) EnableAfterVerify(ctx context.Context, userID, code string) error {
+	secret, err := s.store.GetSecret(ctx, userID)
+	if err != nil {
+		return err
+	}
+	if secret == "" {
+		return ErrNotSetup
+	}
+	if !auth.ValidateTOTP(secret, code, s.nowUnix()) {
+		return ErrBadCode
+	}
+	return s.store.Enable(ctx, userID)
+}
+
+// Disable 校验 TOTP 或备份码后关闭 2FA。未启用时视为成功（幂等）。
+func (s *Service) Disable(ctx context.Context, userID, code, backupCode string) error {
+	secret, enabled, err := s.store.GetSecretAndEnabled(ctx, userID)
+	if err != nil {
+		return err
+	}
+	if !enabled {
+		return nil
+	}
+	if !s.verify(ctx, userID, secret, code, backupCode) {
+		return ErrBadCode
+	}
+	return s.store.Disable(ctx, userID)
+}
+
+// VerifyLogin 在登录流程中校验 2FA：先试 TOTP，再试备份码（一次性消费）。
+// secret 由调用方在登录查询时一并取出，避免重复查库。
+func (s *Service) VerifyLogin(ctx context.Context, userID, secret, totpCode, backupCode string) bool {
+	return s.verify(ctx, userID, secret, totpCode, backupCode)
+}
+
+func (s *Service) verify(ctx context.Context, userID, secret, totpCode, backupCode string) bool {
+	if totpCode != "" && secret != "" && auth.ValidateTOTP(secret, totpCode, s.nowUnix()) {
+		return true
+	}
+	if backupCode != "" {
+		ok, err := s.store.ConsumeBackupCode(ctx, userID, backupCode)
+		if err == nil && ok {
+			return true
+		}
+	}
+	return false
+}

server/internal/service/twofa/service_test.go

@@ -0,0 +1,139 @@
+package twofa
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	"github.com/enterprise-ai-platform/server/pkg/auth"
+)
+
+type fakeStore struct {
+	enabled    bool
+	remaining  int
+	secret     string
+	saved      bool
+	enabled2   bool // Enable 被调用
+	disabled   bool // Disable 被调用
+	backupOK   bool
+	getErr     error
+}
+
+func (f *fakeStore) GetStatus(ctx context.Context, userID string) (bool, int, error) {
+	return f.enabled, f.remaining, f.getErr
+}
+func (f *fakeStore) GetSecret(ctx context.Context, userID string) (string, error) {
+	return f.secret, f.getErr
+}
+func (f *fakeStore) GetSecretAndEnabled(ctx context.Context, userID string) (string, bool, error) {
+	return f.secret, f.enabled, f.getErr
+}
+func (f *fakeStore) SaveEnrollment(ctx context.Context, userID, secret string, codeHashes []string) error {
+	f.saved = true
+	f.secret = secret
+	return nil
+}
+func (f *fakeStore) Enable(ctx context.Context, userID string) error  { f.enabled2 = true; return nil }
+func (f *fakeStore) Disable(ctx context.Context, userID string) error { f.disabled = true; return nil }
+func (f *fakeStore) ConsumeBackupCode(ctx context.Context, userID, code string) (bool, error) {
+	return f.backupOK, nil
+}
+
+const fixedNow int64 = 1_700_000_000
+
+func newSvc(store Store) *Service {
+	s := NewService(store)
+	s.nowUnix = func() int64 { return fixedNow }
+	return s
+}
+
+func TestEnroll_RejectsWhenAlreadyEnabled(t *testing.T) {
+	svc := newSvc(&fakeStore{enabled: true})
+	if _, err := svc.Enroll(context.Background(), "u1", "a@b.c"); !errors.Is(err, ErrAlreadyEnabled) {
+		t.Fatalf("已启用应返回 ErrAlreadyEnabled，实际: %v", err)
+	}
+}
+
+func TestEnroll_GeneratesAndSaves(t *testing.T) {
+	store := &fakeStore{}
+	svc := newSvc(store)
+	res, err := svc.Enroll(context.Background(), "u1", "admin@govai.gov.cn")
+	if err != nil {
+		t.Fatalf("Enroll 出错: %v", err)
+	}
+	if res.Secret == "" || len(res.BackupCodes) != 8 {
+		t.Fatalf("应返回密钥与 8 个备份码，实际 codes=%d", len(res.BackupCodes))
+	}
+	if !store.saved {
+		t.Fatal("应调用 SaveEnrollment")
+	}
+	if res.OtpauthURI == "" {
+		t.Fatal("应返回 otpauth URI")
+	}
+}
+
+func TestEnableAfterVerify(t *testing.T) {
+	secret, _ := auth.GenerateTOTPSecret()
+	code, _ := auth.TOTPCodeAt(secret, fixedNow)
+
+	// 未设置密钥
+	if err := newSvc(&fakeStore{secret: ""}).EnableAfterVerify(context.Background(), "u1", code); !errors.Is(err, ErrNotSetup) {
+		t.Fatalf("无密钥应返回 ErrNotSetup，实际: %v", err)
+	}
+	// 错误验证码
+	if err := newSvc(&fakeStore{secret: secret}).EnableAfterVerify(context.Background(), "u1", "000000"); !errors.Is(err, ErrBadCode) {
+		t.Fatalf("错误码应返回 ErrBadCode，实际: %v", err)
+	}
+	// 正确验证码
+	store := &fakeStore{secret: secret}
+	if err := newSvc(store).EnableAfterVerify(context.Background(), "u1", code); err != nil {
+		t.Fatalf("正确码应成功，实际: %v", err)
+	}
+	if !store.enabled2 {
+		t.Fatal("应调用 Enable")
+	}
+}
+
+func TestDisable(t *testing.T) {
+	secret, _ := auth.GenerateTOTPSecret()
+	code, _ := auth.TOTPCodeAt(secret, fixedNow)
+
+	// 未启用 → 幂等成功，不调用 Disable
+	store0 := &fakeStore{enabled: false}
+	if err := newSvc(store0).Disable(context.Background(), "u1", "", ""); err != nil || store0.disabled {
+		t.Fatalf("未启用应幂等返回 nil 且不调用 Disable，err=%v disabled=%v", err, store0.disabled)
+	}
+	// 启用 + 正确 TOTP
+	store1 := &fakeStore{enabled: true, secret: secret}
+	if err := newSvc(store1).Disable(context.Background(), "u1", code, ""); err != nil {
+		t.Fatalf("正确 TOTP 应成功: %v", err)
+	}
+	if !store1.disabled {
+		t.Fatal("应调用 Disable")
+	}
+	// 启用 + 备份码
+	store2 := &fakeStore{enabled: true, secret: secret, backupOK: true}
+	if err := newSvc(store2).Disable(context.Background(), "u1", "", "backup-xxxx"); err != nil || !store2.disabled {
+		t.Fatalf("备份码应可关闭，err=%v disabled=%v", err, store2.disabled)
+	}
+	// 启用 + 错误码
+	store3 := &fakeStore{enabled: true, secret: secret, backupOK: false}
+	if err := newSvc(store3).Disable(context.Background(), "u1", "000000", "bad"); !errors.Is(err, ErrBadCode) {
+		t.Fatalf("错误码应返回 ErrBadCode，实际: %v", err)
+	}
+}
+
+func TestVerifyLogin(t *testing.T) {
+	secret, _ := auth.GenerateTOTPSecret()
+	code, _ := auth.TOTPCodeAt(secret, fixedNow)
+
+	if !newSvc(&fakeStore{}).VerifyLogin(context.Background(), "u1", secret, code, "") {
+		t.Fatal("正确 TOTP 应通过")
+	}
+	if !newSvc(&fakeStore{backupOK: true}).VerifyLogin(context.Background(), "u1", secret, "", "backup") {
+		t.Fatal("有效备份码应通过")
+	}
+	if newSvc(&fakeStore{backupOK: false}).VerifyLogin(context.Background(), "u1", secret, "000000", "bad") {
+		t.Fatal("错误码应不通过")
+	}
+}

server/internal/service/twofa/store.go

@@ -0,0 +1,121 @@
+package twofa
+
+import (
+	"context"
+
+	"github.com/enterprise-ai-platform/server/pkg/auth"
+	"github.com/jackc/pgx/v5/pgxpool"
+)
+
+// pgxStore 基于 pgx 连接池的 Store 实现。
+type pgxStore struct {
+	pool *pgxpool.Pool
+}
+
+func NewPgxStore(pool *pgxpool.Pool) Store {
+	return &pgxStore{pool: pool}
+}
+
+func (s *pgxStore) GetStatus(ctx context.Context, userID string) (bool, int, error) {
+	var enabled bool
+	var remaining int
+	err := s.pool.QueryRow(ctx,
+		`SELECT u.totp_enabled,
+		        (SELECT COUNT(*) FROM user_backup_codes b WHERE b.user_id = u.id AND b.used_at IS NULL)
+		 FROM users u WHERE u.id = $1`, userID).Scan(&enabled, &remaining)
+	return enabled, remaining, err
+}
+
+func (s *pgxStore) GetSecret(ctx context.Context, userID string) (string, error) {
+	var secret *string
+	if err := s.pool.QueryRow(ctx,
+		`SELECT totp_secret FROM users WHERE id = $1`, userID).Scan(&secret); err != nil {
+		return "", err
+	}
+	if secret == nil {
+		return "", nil
+	}
+	return *secret, nil
+}
+
+func (s *pgxStore) GetSecretAndEnabled(ctx context.Context, userID string) (string, bool, error) {
+	var secret *string
+	var enabled bool
+	if err := s.pool.QueryRow(ctx,
+		`SELECT totp_secret, totp_enabled FROM users WHERE id = $1`, userID).Scan(&secret, &enabled); err != nil {
+		return "", false, err
+	}
+	if secret == nil {
+		return "", enabled, nil
+	}
+	return *secret, enabled, nil
+}
+
+func (s *pgxStore) SaveEnrollment(ctx context.Context, userID, secret string, codeHashes []string) error {
+	tx, err := s.pool.Begin(ctx)
+	if err != nil {
+		return err
+	}
+	defer tx.Rollback(ctx)
+
+	if _, err = tx.Exec(ctx,
+		`UPDATE users SET totp_secret = $2, totp_enabled = false WHERE id = $1`, userID, secret); err != nil {
+		return err
+	}
+	if _, err = tx.Exec(ctx, `DELETE FROM user_backup_codes WHERE user_id = $1`, userID); err != nil {
+		return err
+	}
+	for _, h := range codeHashes {
+		if _, err = tx.Exec(ctx,
+			`INSERT INTO user_backup_codes (user_id, code_hash) VALUES ($1, $2)`, userID, h); err != nil {
+			return err
+		}
+	}
+	return tx.Commit(ctx)
+}
+
+func (s *pgxStore) Enable(ctx context.Context, userID string) error {
+	_, err := s.pool.Exec(ctx, `UPDATE users SET totp_enabled = true WHERE id = $1`, userID)
+	return err
+}
+
+func (s *pgxStore) Disable(ctx context.Context, userID string) error {
+	tx, err := s.pool.Begin(ctx)
+	if err != nil {
+		return err
+	}
+	defer tx.Rollback(ctx)
+	if _, err = tx.Exec(ctx,
+		`UPDATE users SET totp_enabled = false, totp_secret = NULL WHERE id = $1`, userID); err != nil {
+		return err
+	}
+	if _, err = tx.Exec(ctx, `DELETE FROM user_backup_codes WHERE user_id = $1`, userID); err != nil {
+		return err
+	}
+	return tx.Commit(ctx)
+}
+
+func (s *pgxStore) ConsumeBackupCode(ctx context.Context, userID, code string) (bool, error) {
+	rows, err := s.pool.Query(ctx,
+		`SELECT id, code_hash FROM user_backup_codes WHERE user_id = $1 AND used_at IS NULL`, userID)
+	if err != nil {
+		return false, err
+	}
+	type bc struct{ id, hash string }
+	var list []bc
+	for rows.Next() {
+		var x bc
+		if rows.Scan(&x.id, &x.hash) == nil {
+			list = append(list, x)
+		}
+	}
+	rows.Close() // 先释放连接再执行更新
+
+	for _, x := range list {
+		if auth.CheckBackupCode(code, x.hash) {
+			_, _ = s.pool.Exec(ctx, `UPDATE user_backup_codes SET used_at = NOW() WHERE id = $1`, x.id)
+			return true, nil
+		}
+	}
+	return false, nil
+}