MailUI4Agents/server/internal/repo/agentlock_test.go

package repo

import (
	"context"
	"path/filepath"
	"testing"
	"time"

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

/*
Agent↔Agent 回路的**冷静期**（2026-10-01）

# 缺口

`maxAgentPingPong = 8` 撞闸之后**计数永不回落** —— 只有人类插话才归零。
旧文案把出路指向"请由人类插一句话"，而那条线索上**常常根本没有人类**，
于是那一侧永久发不出信（2026-10-01 报告：agent 一封都发不出）。

# 这几条钉的是"有出路"，不是"有锁"

`TestSessionLockBlocksThenRecovers`      —— 锁住 / 到期放行
`TestSessionLockNotExtendedByRetry`     —— 重试不延长（否则永不恢复）
`TestHumanTouchClearsLock`               —— 人插话立刻解锁
`TestExpiredLockRowIsCleaned`           —— 过期行清理（不留两个答案）
`TestSessionLockSurvivesReconnect`       —— 落库而非放内存
*/

func TestSessionLockBlocksThenRecovers(t *testing.T) {
	setupTestDB(t)
	ctx := context.Background()
	sid, _ := CreateSession(ctx, nil, "human", "冷静期", "")

	// 初始：未锁定
	if _, locked, err := SessionLockOf(ctx, sid); err != nil || locked {
		t.Fatalf("初始应为未锁定：locked=%v err=%v", locked, err)
	}

	lock, err := SetSessionLock(ctx, sid, 8, "test")
	if err != nil {
		t.Fatal(err)
	}
	if !lock.Active() {
		t.Fatal("刚设的锁应当生效")
	}

	// 在锁：Agent 的自主回信被拒
	if _, locked, _ := SessionLockOf(ctx, sid); !locked {
		t.Fatal("★ 锁定期内必须拦住（这正是「一封都发不出」那个缺陷）")
	}

	// 到期：自动恢复 —— **不需要任何人介入**
	if _, err := db.DB.ExecContext(ctx,
		`UPDATE session_agent_locks SET until_at = ? WHERE session_id = ?`,
		time.Now().Add(-time.Minute), sid); err != nil {
		t.Fatal(err)
	}
	if _, locked, _ := SessionLockOf(ctx, sid); locked {
		t.Fatal("★ 到期后必须自动放行 —— 恢复不能依赖人类在场")
	}
}

func TestSessionLockNotExtendedByRetry(t *testing.T) {
	setupTestDB(t)
	ctx := context.Background()
	sid, _ := CreateSession(ctx, nil, "human", "重试", "")

	if _, err := SetSessionLock(ctx, sid, 8, "first"); err != nil {
		t.Fatal(err)
	}
	first, _, _ := SessionLockOf(ctx, sid)

	// 一个每 30 分钟试一次的 Agent：若每次都 +2h ⇒ 永不解锁
	time.Sleep(1100 * time.Millisecond)
	if _, err := SetSessionLock(ctx, sid, 9, "retry"); err != nil {
		t.Fatal(err)
	}
	second, _, _ := SessionLockOf(ctx, sid)

	if second.Until.After(first.Until) {
		t.Fatalf("★ 重试不得延长锁定期（否则反复触发=永不解锁）：%s → %s",
			first.Until, second.Until)
	}
	if second.Hops != 9 {
		t.Fatalf("hops 应更新为最新值，实际 %d", second.Hops)
	}
}

func TestHumanTouchClearsLock(t *testing.T) {
	setupTestDB(t)
	ctx := context.Background()
	sid, _ := CreateSession(ctx, nil, "human", "人类插话", "")

	if _, err := SetSessionLock(ctx, sid, 8, "test"); err != nil {
		t.Fatal(err)
	}
	cleared, err := SessionLockTouch(ctx, sid)
	if err != nil || !cleared {
		t.Fatalf("人类插话应立刻清锁：cleared=%v err=%v", cleared, err)
	}
	if _, locked, _ := SessionLockOf(ctx, sid); locked {
		t.Fatal("★ 人参与了就不该再等冷静期")
	}

	// 再插一次：本来就没锁 ⇒ 不报错，只是"没清掉什么"
	if cleared, err := SessionLockTouch(ctx, sid); err != nil || cleared {
		t.Fatalf("未锁定时重复 touch 应为 no-op：cleared=%v err=%v", cleared, err)
	}
}

func TestExpiredLockRowIsCleaned(t *testing.T) {
	setupTestDB(t)
	ctx := context.Background()
	sid, _ := CreateSession(ctx, nil, "human", "过期", "")

	if _, err := SetSessionLock(ctx, sid, 8, "test"); err != nil {
		t.Fatal(err)
	}
	if _, err := db.DB.ExecContext(ctx,
		`UPDATE session_agent_locks SET until_at = ? WHERE session_id = ?`,
		time.Now().Add(-time.Hour), sid); err != nil {
		t.Fatal(err)
	}
	if _, locked, _ := SessionLockOf(ctx, sid); locked {
		t.Fatal("过期行应视为未锁定")
	}

	// 读数只有一个答案：行也必须被回收
	var n int
	if err := db.DB.QueryRowContext(ctx,
		`SELECT COUNT(*) FROM session_agent_locks WHERE session_id = ?`, sid).Scan(&n); err != nil {
		t.Fatal(err)
	}
	if n != 0 {
		t.Fatalf("★ 过期行应顺手清理（留两个答案会让「是否锁定」不可判）：剩 %d 行", n)
	}
}

func TestSessionLockSurvivesReconnect(t *testing.T) {
	// 落库而非放内存：内存态一重启就"恢复"，而锁定期的意义正是
	// 别在短时间内又炸一轮；且多副本部署时各算各的。
	dir := t.TempDir()
	path := filepath.Join(dir, "lock.db")
	ctx := context.Background()

	if err := db.Connect(ctx, path); err != nil {
		t.Fatal(err)
	}
	if err := db.Migrate(ctx); err != nil {
		t.Fatal(err)
	}
	sid, _ := CreateSession(ctx, nil, "human", "重启", "")
	if _, err := SetSessionLock(ctx, sid, 8, "test"); err != nil {
		t.Fatal(err)
	}
	db.Close() // 模拟重启

	if err := db.Connect(ctx, path); err != nil {
		t.Fatal(err)
	}
	t.Cleanup(db.Close)
	if _, locked, err := SessionLockOf(ctx, sid); err != nil || !locked {
		t.Fatalf("★ 重启后冷却期必须仍在（否则它只是个瞬时现象）：locked=%v err=%v", locked, err)
	}
}

// 确认表真的建在两种方言的 schema 里（老库靠 migrate，新库靠 init.sql）。
func TestSessionLockTableExistsOnSQLite(t *testing.T) {
	setupTestDB(t)
	var name string
	if err := db.DB.QueryRowContext(context.Background(),
		`SELECT name FROM sqlite_master WHERE type='table' AND name='session_agent_locks'`).Scan(&name); err != nil {
		t.Fatalf("session_agent_locks 未建表：%v", err)
	}
}

/*
★ 读失败**不得**被读成"未锁定"（变异②补：原本 0 格能抓，漏网了）

# 为什么单独一格

我把 `err != nil` 合并进 `return {}, false, nil`（即「读不出来 = 没锁」）。
上面五条判据**全都抓不住它** —— 它们都不制造读失败，只是各走各的正常路径。

而这个形状是本仓记录过的同一个错：「`getWorkspace` 拿不到时不带，让服务端报 400」
背后的原则是**宁可报错、不可静默降级**；这里反过来：**读不出来却静默放行**。
把"不知道"读成"安全"，等于静默拆掉防护，且**没有任何症状**。

# 怎么钉

让读真的失败（表不存在），然后断言「**仍报在锁**」——
若实现退化成 `return {}, false, nil`，locked 会变 false ⇒ 这格变红。
*/
func TestSessionLockReadFailureIsNotTreatedAsUnlocked(t *testing.T) {
	setupTestDB(t)
	ctx := context.Background()
	sid, _ := CreateSession(ctx, nil, "human", "读失败", "")

	// 制造一次真实的读失败：把表改名，Scan 就会报 no such table
	if _, err := db.DB.ExecContext(ctx,
		`ALTER TABLE session_agent_locks RENAME TO session_agent_locks_gone`); err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() {
		_, _ = db.DB.ExecContext(ctx,
			`ALTER TABLE session_agent_locks_gone RENAME TO session_agent_locks`)
	})

	lock, locked, err := SessionLockOf(ctx, sid)
	if !locked {
		t.Fatalf("★ 读失败必须当「仍在锁」而非「未锁定」"+
			"（否则一次 DB 抖动就静默拆掉整个防护，且无任何症状）：locked=%v lock=%+v err=%v",
			locked, lock, err)
	}
	if err == nil {
		t.Fatal("读失败必须把错误往上带（调用方需要知道这是异常而不是「没锁」）")
	}
}

var _ = uuid.New