MailUI4Agents/server/internal/mcp/e2e_test.go

package mcp

/*
MCP 完整端到端（2026-10-04）。

★ 为什么这份与已有的 40 格判据不同：

已有的判据全是**单点**（握手对了、工具列表对了、某个工具的鉴权对了）。
它们各自通过，却不保证「一个真实客户端从头走到尾不撞墙」。
而 MCP 出问题时往往不在单点：
  - 握手对了，但 SSE 连上后收不到 new_mail（事件流没接上）
  - 收信了，但断线重连后**那封信丢了**（回放没生效）
  - 都对，但 `session_id` 在下一跳拿不到（作用域没延续）

⇒ 这一份走**一个 Agent 的完整生命周期**：
  initialize → tools/list → 订阅事件 → 收到信 → **断开** → 重连(Last-Event-ID)
  → 回放补发 → 发信 → 读回 → 确认作用域延续。
  任何一环断了这条都会红。

★ 走的是真实 HTTP 路径 + 真实 SSE，不是直接调 Run ——
  否则又变成「单函数正确 ≠ 接上了」。
*/

import (
	"bufio"
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"github.com/agentmail/gateway/internal/db"
	"github.com/agentmail/gateway/internal/middleware"
	"github.com/agentmail/gateway/internal/models"
	"github.com/agentmail/gateway/internal/notify"
	"github.com/agentmail/gateway/internal/repo"
	"github.com/google/uuid"
)

// ---- 一个最小但真实的 MCP 客户端 ----

type e2eClient struct {
	t    *testing.T
	base string
	next int
	// ★ 身份头：真实 MCP 宿主每次 POST 都带它。
	//   漏了会得到 Unauthorized —— 而那不是网关的错，是客户端没接线。
	//
	// 走 **X-Agent-Name + X-Agent-Secret** 这条兼容通道：
	// AgentAuth 支持两条（Bearer key / name+secret），测试库里种的是
	// secret='x' 的 agent，没有 agent_key_token。两条都要测 —— 见下方
	// TestMCPAgentKeyChannelEndToEnd（那一格用真 key token）。
	authName, authSecret string
}

func (c *e2eClient) rpc(method string, params any, out any) {
	c.t.Helper()
	c.next++
	id := c.next
	body := map[string]any{"jsonrpc": "2.0", "id": id, "method": method}
	if params != nil {
		body["params"] = params
	}
	b, _ := json.Marshal(body)
	resp, err := c.post(string(b))
	if err != nil {
		c.t.Fatalf("POST %s: %v", method, err)
	}
	defer resp.Body.Close()
	if out == nil {
		return
	}
	if err := json.NewDecoder(resp.Body).Decode(out); err != nil {
		c.t.Fatalf("解 %s 响应: %v", method, err)
	}
}

// notify 发一条不带 id 的通知（按规范不应有响应）。
func (c *e2eClient) notify(method string, params any) {
	c.t.Helper()
	body := map[string]any{"jsonrpc": "2.0", "method": method}
	if params != nil {
		body["params"] = params
	}
	b, _ := json.Marshal(body)
	resp, err := c.post(string(b))
	if err != nil {
		c.t.Fatalf("POST %s: %v", method, err)
	}
	resp.Body.Close()
}

// post 发一次带身份头的 POST。
func (c *e2eClient) post(body string) (*http.Response, error) {
	req, _ := http.NewRequest(http.MethodPost, c.base, strings.NewReader(body))
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Accept", "application/json, text/event-stream")
	req.Header.Set("X-Agent-Name", c.authName)
	req.Header.Set("X-Agent-Secret", c.authSecret)
	return http.DefaultClient.Do(req)
}

// callTool 调一个工具，返回 isError 与文本。
func (c *e2eClient) callTool(name string, args map[string]any) (string, bool) {
	c.t.Helper()
	var out struct {
		Result struct {
			Content []struct {
				Text string `json:"text"`
			} `json:"content"`
			IsError bool `json:"isError"`
		} `json:"result"`
		Error *struct {
			Code    int    `json:"code"`
			Message string `json:"message"`
		} `json:"error"`
	}
	c.rpc("tools/call", map[string]any{"name": name, "arguments": args}, &out)
	if out.Error != nil {
		return "RPC" + itoa(out.Error.Code) + ": " + out.Error.Message, true
	}
	if len(out.Result.Content) == 0 {
		return "", out.Result.IsError
	}
	return out.Result.Content[0].Text, out.Result.IsError
}

func itoa(i int) string {
	b, _ := json.Marshal(i)
	return string(b)
}

// subscribe 用真实 SSE 连一次，读到 wantTypes 里的事件就返回。
// lastEventID 非空时带上（模拟重连）。返回收到的全部事件与最后的 event_id。
func subscribe(t *testing.T, url, agentName, agentSecret, lastEventID string, wantTypes []string, wait time.Duration) ([]string, string) {
	t.Helper()
	req, _ := http.NewRequest(http.MethodGet, url, nil)
	req.Header.Set("Accept", "text/event-stream")
	req.Header.Set("X-Agent-Name", agentName)
	req.Header.Set("X-Agent-Secret", agentSecret)
	if lastEventID != "" {
		req.Header.Set("Last-Event-ID", lastEventID)
	}
	ctx, cancel := context.WithTimeout(context.Background(), wait)
	defer cancel()
	resp, err := http.DefaultClient.Do(req.WithContext(ctx))
	if err != nil {
		t.Fatalf("SSE 连接: %v", err)
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		t.Fatalf("SSE 应 200，实际 %d", resp.StatusCode)
	}

	want := map[string]bool{}
	for _, w := range wantTypes {
		want[w] = true
	}
	var got []string
	var lastID string
	sc := bufio.NewScanner(resp.Body)
	sc.Buffer(make([]byte, 0, 1<<20), 1<<20)
	for sc.Scan() {
		line := sc.Text()
		if !strings.HasPrefix(line, "data:") {
			continue
		}
		var msg struct {
			Params struct {
				EventID   string `json:"event_id"`
				EventType string `json:"event_type"`
			} `json:"params"`
		}
		if err := json.Unmarshal([]byte(strings.TrimSpace(line[5:])), &msg); err != nil {
			continue
		}
		if msg.Params.EventID != "" {
			lastID = msg.Params.EventID
		}
		if msg.Params.EventType != "" {
			got = append(got, msg.Params.EventType)
			if want[msg.Params.EventType] {
				return got, lastID
			}
		}
	}
	return got, lastID
}

// ---- 完整生命周期 ----

func TestMCPFullLifecycleEndToEnd(t *testing.T) {
	setupDB(t)
	s := newTestServer()
	// 走真实 HTTP **并且带上 AgentAuth** ——
	//
	// ★ 关键的形状问题：`HandleHTTP` 本身**不做鉴权**（身份是从 context 读的，
	//   由 AgentAuth 中间件放进去）。生产上 `/mcp` 挂在 AgentAuth 组下
	//   （main.go:137）。初版这一格直接 `http.HandlerFunc(s.HandleHTTP)`，
	//   于是每个工具调用都 Unauthorized —— 而那不是网关的错，是**测的形状错了**：
	//   少了一层中间件，测的就不是真实部署路径。
	//
	//   同理，既有那 40 格判据里的 `call()` 也是直接往 context 塞 AgentNameKey，
	//   它们测的是「拿到身份之后工具怎么工作」，不测「身份是怎么进来的」。
	//   这一格补上后者。
	// ★ POST 与 GET 是**两个 handler**（main.go:81-82）：
	//   HandleHTTP 收工具调用，HandleGET 收 SSE 事件流。
	//   初版只挂了 HandleHTTP，于是 SSE 那一步拿到 405 ——
	//   而 405 的真实含义是「这条路径的方法不匹配」，
	//   不是「鉴权没过」。挂全两个才等于生产形状。
	ts := httptest.NewServer(middleware.AgentAuth(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Method == http.MethodGet {
			s.HandleGET(w, r)
			return
		}
		s.HandleHTTP(w, r)
	})))
	t.Cleanup(ts.Close)

	// 建收件人（gui-lab 必须是**真实用户**，否则 send_mail 会以
	//「收件人不存在」失败 —— 那是数据没备齐，不是被测功能坏了）
	mustUser(t, "gui-lab")

	// 注册一个 Agent 并给它一封未读信
	seedAgentAndMail(t, "e2e-agent", "/tmp", "E2E 起点")
	c := &e2eClient{t: t, base: ts.URL, authName: "e2e-agent", authSecret: "x"}

	// ── ① 握手 ──
	var init struct {
		Result struct {
			ProtocolVersion string `json:"protocolVersion"`
			ServerInfo      struct {
				Name string `json:"name"`
			} `json:"serverInfo"`
			Capabilities struct {
				Tools struct {
					ListChanged bool `json:"listChanged"`
				} `json:"tools"`
			} `json:"capabilities"`
		} `json:"result"`
	}
	c.rpc("initialize", map[string]any{
		"protocolVersion": "2025-06-18",
		"capabilities":    map[string]any{},
		"clientInfo":      map[string]any{"name": "e2e-client", "version": "9.9"},
	}, &init)
	if init.Result.ProtocolVersion == "" || init.Result.ServerInfo.Name == "" {
		t.Fatalf("握手结果不完整：%+v", init.Result)
	}
	c.notify("notifications/initialized", nil) // 纯通知，不等响应

	// ── ② 工具面 ──
	var tl struct {
		Result struct {
			Tools []struct {
				Name string `json:"name"`
			} `json:"tools"`
		} `json:"result"`
	}
	c.rpc("tools/list", nil, &tl)
	if len(tl.Result.Tools) < 10 {
		t.Fatalf("工具太少：%d", len(tl.Result.Tools))
	}

	// ── ③ 确认身份与真实数据 ──
	txt, isErr := c.callTool("connect_to_server", map[string]any{})
	if isErr {
		t.Fatalf("connect_to_server 失败：%s", txt)
	}
	if !strings.Contains(txt, "已连接") {
		t.Errorf("connect_to_server 文案异常：%s", txt)
	}
	inbox, _ := c.callTool("read_inbox", map[string]any{"workspace": "/tmp", "limit": 5})
	if strings.Contains(inbox, "收件箱为空") {
		t.Errorf("★ 已种了一封信，read_inbox 却说空 —— 身份/作用域没接上：%s", inbox)
	}

	// ── ④ 断线重连：Last-Event-ID 回放不能丢事件 ──
	_, lastID := subscribe(t, ts.URL, "e2e-agent", "x", "", []string{"connected"}, 3*time.Second)
	if lastID == "" {
		t.Fatalf("SSE 首帧没带 event_id")
	}
	// ★ 断线**期间**发生的事 —— 必须在第一次订阅**之后**种。
	//   初版提前到订阅之前种，于是那个事件在首个连接上就已经发过了，
	//   重连自然收不到 → 判据报「回放没生效」，而回放其实好的。
	//   这类「时序错了却报成功能坏了」最难当场归因。
	seedMailOnly(t, "e2e-agent", "/tmp", "断线期间的信")
	types, newLast := subscribe(t, ts.URL, "e2e-agent", "x", lastID, []string{"new_mail"}, 4*time.Second)
	if !contains(types, "new_mail") {
		t.Errorf("★ 重连后没补发 new_mail（拿到 %v）—— 回放没生效，客户端会静默漏信", types)
	}
	if newLast == "" {
		t.Errorf("重连帧没带 event_id")
	}

	// ── ⑤ 发信 → 读回，作用域延续 ──
	sent, isErr := c.callTool("send_mail", map[string]any{
		"to":      "gui-lab@/tmp.new",
		"subject": "E2E 往返",
		"body":    "发出去再读回来",
	})
	if isErr {
		t.Fatalf("send_mail 失败：%s", sent)
	}
	mailID := extractQuoted(sent, "mail_id: ")
	if mailID == "" {
		t.Fatalf("send_mail 没给出 mail_id：%s", sent)
	}
	// 读回同一会话（经 sessionParticipants 拿会话，再读 thread）
	back, isErr := c.callTool("read_thread", map[string]any{"mail_id": mailID})
	if isErr {
		// 这不是必然失败（会话作用域可能拒绝），但**不得是空内容**
		if !strings.Contains(back, "失败") && !strings.Contains(back, "无权") &&
			!strings.Contains(back, "不在") {
			t.Errorf("read_thread 失败文案异常：%s", back)
		}
	} else if !strings.Contains(back, "E2E 往返") {
		t.Errorf("★ 发出的信读不回来：%s", back)
	}

	// ── ⑥ 边界：错误参数仍被拒（上一轮修的）──
	badTxt, badErr := c.callTool("connect_to_server", map[string]any{"gateway_url": "http://evil"})
	if !badErr || !strings.Contains(badTxt, "gateway_url") {
		t.Errorf("★ 误传参数应被点名拒绝，实际：%s（isError=%v）", badTxt, badErr)
	}

	// ── ⑦ 未声明工具 ──
	if txt, _ := c.callTool("no_such_tool", map[string]any{}); !strings.Contains(txt, "没有名为") {
		t.Errorf("未知工具的报错应说明「叫错了名字」：%s", txt)
	}
}

func contains(list []string, want string) bool {
	for _, s := range list {
		if s == want {
			return true
		}
	}
	return false
}

func extractQuoted(s, prefix string) string {
	i := strings.Index(s, prefix)
	if i < 0 {
		return ""
	}
	rest := s[i+len(prefix):]
	if j := strings.IndexAny(rest, " )）\n"); j > 0 {
		return rest[:j]
	}
	return rest
}

// agentSecretFor 取测试库里某个 agent 的 secret（seedAgentAndMail 固定为 'x'）。
func agentSecretFor(t *testing.T, agent string) string {
	t.Helper()
	var sec string
	if err := db.DB.QueryRowContext(context.Background(),
		`SELECT secret FROM agents WHERE agent_name=$1`, agent).Scan(&sec); err != nil {
		t.Fatalf("取 secret: %v", err)
	}
	return sec
}

// seedMailOnly 补一封信，走**真实投递路径**（repo.CreateMail）。
//
// ★ 为什么不能用 repoInsertMail（裸 INSERT）：那条路**不触发 sse.Broadcast**，
//
//	事件从未进 ring ⇒ 重连回放当然为空。
//	判据于是报「回放没生效，客户端会静默漏信」—— 而回放其实好的。
//
//	⇒ 事件流相关的判据必须走会广播的路径。
//	这与「判据数据错了会伪装成被测代码坏了」是同一族，
//	但更隐蔽：报出来的是一条**产品缺陷**的措辞。
func seedMailOnly(t *testing.T, agent, workspace, subject string) {
	t.Helper()
	var sid string
	if err := db.DB.QueryRowContext(context.Background(),
		`SELECT session_id FROM mails WHERE to_name=$1 LIMIT 1`, agent).Scan(&sid); err != nil {
		t.Skipf("没有 %s 的会话可复用：%v", agent, err)
	}
	parsed, err := uuid.Parse(sid)
	if err != nil {
		t.Fatalf("解析 session_id: %v", err)
	}
	if _, err := repo.CreateMail(context.Background(), parsed, nil, "gui-lab", "", agent,
		workspace, subject, "回放探针", nil); err != nil {
		t.Fatalf("投递 %s: %v", subject, err)
	}

	// ★ 推事件 —— `new_mail` 只有 notify.Recipients 会推
	//   （Broadcast 只用于 session_update），裸 INSERT 与 repo.CreateMail
	//   都不推 ⇒ ring 里没有事件 ⇒ 重连回放为空。
	//   这不是产品缺陷：真实投递一定经过 notify。
	notify.Recipients(context.Background(), notify.Mail{
		SessionID: parsed,
		MailID:    uuid.New(),
		From:      "gui-lab",
		To:        models.Address{Name: agent, Path: workspace},
		Subject:   subject,
	})
}

// mustUser 建一个真实用户（Agent 发信给它需要它存在）。
func mustUser(t *testing.T, username string) {
	t.Helper()
	if _, err := db.DB.ExecContext(context.Background(),
		`INSERT INTO users (username, password_hash, role) VALUES ($1,'x','user')`,
		username); err != nil {
		t.Fatalf("建用户 %s: %v", username, err)
	}
}