MailUI4Agents/server/internal/handler/relay_budget_test.go

package handler

/*
relay 的配额回传必须与实际扣费一致（2026-10-04 实测驱动）。

★ 缺陷：`resp["budget_charged"] = false` 只看 `relay != ""`，
  于是**任何** relay 都回「没扣预算」。而实际：

    relay → 人类      relayFree=true   确实没扣   回传 false  ✓
    relay → 另一个 Agent relayFree=false **扣了**   回传 false  ✗ 撒谎

★ 为什么这个字段的错很要紧（不是文案瑕疵）：
  插件据 budget_charged 判断「还有额度」，Agent↔Agent 的 relay
  明明扣了却被告知没扣 —— 而「Agent↔Agent 双方都不计数」正是
  生产上那个 41 封、间隔 15 分钟→5 秒的无穷循环（会话 f3d824ce）。
  这个字段就是那条链上的判定依据之一。

★ 实测证据：relay→Agent 的会话 used_rounds=1（确实扣了），
  同一响应里 budget_charged=false。
*/

import (
	"context"
	"net/http"
	"net/http/httptest"
	"path/filepath"
	"strings"
	"testing"

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

func setupRelayBudgetDB(t *testing.T) {
	t.Helper()
	db.Close()
	p := filepath.Join(t.TempDir(), "relay-budget.db")
	if err := db.Connect(context.Background(), "sqlite://"+p); err != nil {
		t.Fatalf("连接测试库: %v", err)
	}
	if err := db.Migrate(context.Background()); err != nil {
		t.Fatalf("迁移: %v", err)
	}
	t.Cleanup(db.Close)
}

func mustAgentForBudget(t *testing.T, name string) {
	t.Helper()
	if _, err := db.DB.ExecContext(context.Background(),
		`INSERT INTO agents (agent_name, secret, platform) VALUES ($1,'x','test')
		 ON CONFLICT (agent_name) DO NOTHING`, name); err != nil {
		t.Fatalf("建 agent %s: %v", name, err)
	}
}

func mustHumanForBudget(t *testing.T, name string) {
	t.Helper()
	if _, err := db.DB.ExecContext(context.Background(),
		`INSERT INTO users (username, password_hash, role) VALUES ($1,'x','user')
		 ON CONFLICT (username) DO NOTHING`, name); err != nil {
		t.Fatalf("建用户 %s: %v", name, err)
	}
}

func postMailAs(t *testing.T, agent, body string) *httptest.ResponseRecorder {
	t.Helper()
	req := httptest.NewRequest(http.MethodPost, "/api/v1/mail/send", strings.NewReader(body))
	req = req.WithContext(context.WithValue(req.Context(), middleware.AgentNameKey, agent))
	rec := httptest.NewRecorder()
	SendMail(rec, req)
	return rec
}

// relay 发给人类：确实免配额，且回传 budget_charged=false
func TestRelayToHumanIsFreeAndReportsSo(t *testing.T) {
	setupRelayBudgetDB(t)
	mustAgentForBudget(t, "sender-h")
	mustHumanForBudget(t, "human-h")

	rec := postMailAs(t, "sender-h",
		`{"to":"human-h@/tmp.new","subject":"总结","body":"x","relay":"summary","relay_key":"h1"}`)
	if rec.Code != http.StatusOK {
		t.Fatalf("应 200，实际 %d：%s", rec.Code, rec.Body.String())
	}
	body := rec.Body.String()
	if !strings.Contains(body, `"budget_charged":false`) {
		t.Errorf("relay 发给人类应回 budget_charged=false，实际：%s", body)
	}
}

// ★ relay 发给另一个 Agent：**要扣**预算，且回传必须说 true
func TestRelayToAgentIsChargedAndReportsTruthfully(t *testing.T) {
	setupRelayBudgetDB(t)
	mustAgentForBudget(t, "sender-a")
	mustAgentForBudget(t, "receiver-a")

	rec := postMailAs(t, "sender-a",
		`{"to":"receiver-a@/tmp.new","subject":"relay给Agent","body":"x","relay":"summary","relay_key":"a1"}`)
	if rec.Code != http.StatusOK {
		t.Fatalf("应 200，实际 %d：%s", rec.Code, rec.Body.String())
	}
	body := rec.Body.String()

	// ① 回传不能说谎
	if !strings.Contains(body, `"budget_charged":true`) {
		t.Errorf("★ relay 发给 Agent **扣了**预算，回传却说 budget_charged=false"+
			"（插件据此以为还有额度而继续发）：%s", body)
	}

	// ② 而且库里的计数必须真的动了
	var sid string
	if err := db.DB.QueryRowContext(context.Background(),
		`SELECT session_id FROM mails ORDER BY created_at DESC LIMIT 1`).Scan(&sid); err != nil {
		t.Fatalf("取会话: %v", err)
	}
	var used int
	if err := db.DB.QueryRowContext(context.Background(),
		`SELECT used_rounds FROM sessions WHERE session_id = $1`, sid).Scan(&used); err != nil {
		t.Fatalf("读 used_rounds: %v", err)
	}
	if used != 1 {
		t.Errorf("★ relay 发给 Agent 应扣 1 次预算，实际 used_rounds=%d", used)
	}
}

var _ = models.Mail{}