MailUI4Agents/server/internal/handler/permission_relay_release_test.go

package handler

import (
	"bytes"
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/agentmail/gateway/internal/db"
	"github.com/agentmail/gateway/internal/middleware"
	"github.com/agentmail/gateway/internal/repo"
)

// 占住幂等键之后，**任何**一条早退路径都必须把键还回去；同时幂等性本身不能坏掉。
//
// 背景（真实数据）：现库里 `relayed_mails` 有一行孤儿键 ——
//
//	agent_name = zcode
//	relay_key  = no-such-session-0000:toolu-nohuman-1789193173578
//	mail_id    = NULL          （= 占位成功，但从未绑定到任何邮件）
//	created_at = 2026-09-12 06:06:13
//
// 该键的 session 位不是 UUID，于是 `RequestPermission` 走到
// `uuid.Parse` 失败那条 `Invalid session_id`（permission.go:112）时直接 return，
// **没有调用 ReleaseRelay**。键因此永久留在表里：之后任何用同一个
// (agent_name, relay_key) 的重试都会拿到 `ErrRelayDuplicate`（→ 200
// duplicate_relay），而那次询问其实**从未发出去**。
//
// 为什么必须是两个独立的测试：这两件事在**早退路径**上观测等价。
// 早退时正确实现与"把 ClaimRelay 换成空操作"的坏实现**都**在表里留下一个空键，
// 所以任何单条用例都无法同时钉住"该退的时候退"和"该占的时候占"。
// 前者只能在**失败路径**上观测（键本该被还回来），
// 后者只能在**成功路径**上观测（键本该被占住）。

// 失败路径：早退之后同一个 relay_key 必须还能再次被占用（键被还回来了）。
func TestRequestPermissionReleasesRelayKeyOnEarlyExit(t *testing.T) {
	setupPermissionHandlerDB(t)
	ctx := context.Background()

	if err := repo.CreateOrUpdateAgent(ctx, "pi", "secret", "test", nil); err != nil {
		t.Fatalf("注册 pi: %v", err)
	}

	// 非 UUID 的 session 位：这正是野外那行孤儿键的形态，也是走到早退的入口。
	const relayKey = "no-such-session-0000:toolu-nohuman-1789193173578"
	resp := requestPermissionRaw(t, map[string]any{
		"session_id": "no-such-session-0000",
		"to":         "human",
		"question":   "是否允许执行 bash？",
		"options":    []string{"同意", "拒绝"},
		"relay_key":  relayKey,
	})

	// 前提：请求确实失败（否则这个用例没验到早退那条路）。
	if resp.Code != http.StatusBadRequest {
		t.Fatalf("非 UUID 的 session_id 必须 400，实际 HTTP %d: %s", resp.Code, resp.Body.String())
	}

	// 断言：键被还回来了 —— 同一个键必须能再次占用。
	if err := repo.ClaimRelay(ctx, "pi", relayKey, "permission"); err != nil {
		t.Fatalf("早退路径没把幂等键还回去：再次占用同一个键被拒（%v）—— "+
			"该 (agent, relay_key) 已被永久占用，之后的重试只会拿到 duplicate_relay，"+
			"而那次询问从未发出", err)
	}
}

// 成功路径：幂等键必须真的被占住 —— 同一条询问重复投递只能生成一封邮件。
//
// 这条与上一条是**互补**的，不是重复：把 `ClaimRelay` 换成空操作会让上一条全绿
// （键从未被占 ⇒ "还能再占"天然成立），只有这条会红。
func TestRequestPermissionClaimsRelayKeyOnSuccess(t *testing.T) {
	setupPermissionHandlerDB(t)
	ctx := context.Background()

	if err := repo.CreateOrUpdateAgent(ctx, "pi", "secret", "test", nil); err != nil {
		t.Fatalf("注册 pi: %v", err)
	}
	mustPermissionUser(t, "owner")
	sid, err := repo.CreateSession(ctx, nil, "pi", "权限幂等", "/tmp")
	if err != nil {
		t.Fatalf("建会话: %v", err)
	}
	if err := repo.SetSessionOwner(ctx, sid, permissionUserID(t, "owner")); err != nil {
		t.Fatalf("设置 owner: %v", err)
	}
	if _, err := repo.CreateMail(ctx, sid, nil, "owner", "", "pi", "/tmp",
		"任务", "请执行", nil); err != nil {
		t.Fatalf("建线索邮件: %v", err)
	}

	const relayKey = "sess-1:toolu-repeat-1"
	body := map[string]any{
		"session_id": sid.String(),
		"to":         "owner",
		"question":   "是否允许执行 bash？",
		"options":    []string{"同意", "拒绝"},
		"relay_key":  relayKey,
	}

	first := requestPermissionRaw(t, body)
	if first.Code != http.StatusOK {
		t.Fatalf("首次询问必须受理，HTTP %d: %s", first.Code, first.Body.String())
	}

	// 第二次同一 relay_key：幂等要求它**不产生新邮件**。
	second := requestPermissionRaw(t, body)
	if second.Code != http.StatusOK {
		t.Fatalf("重复投递应幂等返回 200，HTTP %d: %s", second.Code, second.Body.String())
	}
	var got map[string]any
	if err := json.Unmarshal(second.Body.Bytes(), &got); err != nil {
		t.Fatalf("解析响应: %v", err)
	}
	if got["status"] != "duplicate_relay" {
		t.Fatalf("重复投递必须被幂等键挡住（status=duplicate_relay），实际 %v —— "+
			"幂等键没生效，SSE 重放与插件重试会给同一次询问发多封邮件", got["status"])
	}

	// 并且库里那条键确实绑到了第一封邮件上（不是留在空占位状态）。
	var bound string
	if err := db.DB.QueryRowContext(ctx,
		`SELECT mail_id FROM relayed_mails WHERE agent_name = $1 AND relay_key = $2`,
		"pi", relayKey).Scan(&bound); err != nil {
		t.Fatalf("查幂等键绑定: %v", err)
	}
	if bound == "" {
		t.Fatal("成功路径没有把幂等键绑定到邮件上（mail_id 为空）—— " +
			"该键会被 ReleaseRelay 当成占位残留删掉，幂等随之失效")
	}
}

// requestPermissionRaw 发一次 /permission/request，body 由调用方给定。
func requestPermissionRaw(t *testing.T, body map[string]any) *httptest.ResponseRecorder {
	t.Helper()
	payload, _ := json.Marshal(body)
	req := httptest.NewRequest(http.MethodPost, "/api/v1/permission/request", bytes.NewReader(payload))
	req.Header.Set("Content-Type", "application/json")
	ctx := context.WithValue(req.Context(), middleware.AgentNameKey, "pi")
	resp := httptest.NewRecorder()
	RequestPermission(resp, req.WithContext(ctx))
	return resp
}