ModelRouter/internal/core/channels_trae_pkce_test.go

package core

import (
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"sync/atomic"
	"testing"

	"llmsproxy/internal/lua"
)

// ★ TRAE 回调只带 code（PKCE 形态）时必须**明确报出**「暂不支持 PKCE」，
// 而不是去调 ExchangeToken 换一个不存在的码。
//
// 真实缺陷（用户报「登录后显示失败」）：
// poll_login 的 trae 分支曾自己发明了一段「授权码交换」——把 {code} 塞进
// exchange_body 再打 ExchangeToken。但：
//
//	· ExchangeToken 的签名是 RefreshToken，不是 code；
//	· exchange_body 里根本没有 {code} 占位符，那次替换是空操作；
//	· client_id 当时是空串。
//
// 于是稳定得到 400，日志回一句误导性的「Trae 授权码交换失败: 400」，
// 把真因「上游走了 PKCE 流程」掩盖成「交换出错了」——排查方向完全错。
//
// 参考实现（trae-oauth.ts parseTraeCallbackDetailed）的做法是：带 code 的回调
// 是**合法回调**，只是不支持的 PKCE 分支，必须明确报出来。
func TestTraePkceCallbackIsReportedNotExchanged(t *testing.T) {
	h := newTraeHarness(t)
	// ★ 必须打**真实登录用的那条路径**：poll_login（loopback 回调文件）。
	// oauth_callback 那条一直是对的，判据打在那里会漏掉缺陷 —— 我第一版就
	// 犯了这个错（变异测试揭穿：注入缺陷后那条判据仍绿）。
	//
	// poll_login 从 plugin.state.loopbacks[sessionId].file 读 {code,state}，
	// 故判据直接摆好那个文件再调 poll_login。
	status, body := h.beginAndPollWithCallback(map[string]string{
		"code": "PKCE_AUTH_CODE_XYZ",
	})

	// poll_login 的契约：终态用 HTTP 200 + json.state 表达（不是非 200 状态码）。
	// 所以「失败」的判据是 state=failed，而不是 status != 200 —— 我第一版把
	// 前者写成了后者，于是把正确的实现判成错。
	if !strings.Contains(body, "state:failed") {
		t.Fatalf("PKCE 回调应落定为 failed（poll_login 用 200+state 表达终态）：status=%d body=%s",
			status, body)
	}
	if n := atomic.LoadInt32(&h.exchHit); n != 0 {
		t.Fatalf("不得去调 ExchangeToken —— 那是拿 code 换一个不存在的码（真实缺陷）。"+
			"实际打了 %d 次", n)
	}
	// 错误信息必须指向真因（PKCE），而不是含糊的交换失败
	blob := body
	if !strings.Contains(blob, "PKCE") {
		t.Fatalf("错误信息应点名 PKCE（让排查方向正确），实得 %q", blob)
	}
	if strings.Contains(blob, "授权码交换失败") {
		t.Fatalf("不应再出现误导性的「授权码交换失败」：%q", blob)
	}
	t.Logf("✓ PKCE 回调被正确识别（poll_login 路径）: %s", strings.TrimSpace(blob))
}

// ★ client_id 必须是实测真值 —— ExchangeToken 的请求体里 ClientID 是必填。
// 留空会让真实 token 直传场景的交换稳定 400（夹具自身也校验 ClientID 存在）。
func TestTraeExchangeSendsRealClientID(t *testing.T) {
	h := newTraeHarness(t)
	// 真实回调形态：带 refreshToken（token 直传）⇒ 会走 ExchangeToken。
	// 夹具的 handler 已校验请求体含 ClientID / RefreshToken:<回调值> / ClientSecret。
	status, body := h.beginAndCallback(map[string]string{
		"refreshToken": "RT-FROM-CALLBACK",
		"userInfo":     `{"UserID":"u1","ScreenName":"tester","TenantID":"t1"}`,
	})
	if status != 200 {
		t.Fatalf("token 直传回调应成功: status=%d body=%s", status, body)
	}
	if n := atomic.LoadInt32(&h.exchHit); n == 0 {
		t.Fatal("token 直传场景必须真的调用 ExchangeToken")
	}
	t.Logf("✓ ExchangeToken 被调用 %d 次，请求体形状正确", atomic.LoadInt32(&h.exchHit))
}

// client_id 的真值必须是实测值 —— 这是「交换必然 400」的两个成因之一。
func TestTraeClientIDIsNotBlank(t *testing.T) {
	b, err := readBundledChannels(t)
	if err != nil {
		t.Fatal(err)
	}
	src := string(b)
	if !strings.Contains(src, `client_id = "en1oxy7wnw8j9n"`) {
		t.Fatal("trae 的 client_id 必须是实测真值 en1oxy7wnw8j9n —— " +
			"留空/占位会让 ExchangeToken 的 ClientID 为空而稳定 400")
	}
}

// 两条路径（oauth_callback 与 poll_login 的 loopback）必须给**同一个**结论，
// 否则修复只覆盖一条、另一条继续漂移（这正是缺陷能藏住的原因）。
func TestTraeCallbackPathsAgreeOnPkce(t *testing.T) {
	b, err := readBundledChannels(t)
	if err != nil {
		t.Fatal(err)
	}
	src := string(b)
	// ① 不得再出现自己发明的交换（注释里说明历史的那句除外）
	if strings.Contains(src, `cerr = "Trae 授权码交换失败`) {
		t.Fatal("已删除的「授权码交换」又回来了 —— 那是拿 code 去 ExchangeToken，" +
			"必然 400 且把真因掩盖成交换失败")
	}
	// ② 两条路径都必须交给 traeCallback（它内建 PKCE 判定）
	if n := strings.Count(src, "traeCallback(plat,"); n < 2 {
		t.Fatalf("traeCallback 应被两条路径共用（当前 %d 处）—— 只有一条会漂移", n)
	}
	// ③ PKCE 判定必须仍在 traeCallback 内
	if !strings.Contains(src, "本实现暂不支持该流程") {
		t.Fatal("traeCallback 必须保留 PKCE 明确报错（参考实现的语义）")
	}
}

// beginAndPollWithCallback 走**真实登录用的那条路径**：
// start_login 拿会话 → 摆好 loopback 回调文件（平台重定向的落地）→ poll_login。
//
// 与 beginAndCallback（打 oauth_callback）的区别很关键：poll_login 才是
// callback 流在服务端浏览器（VNC）场景下真正走的那条，真实缺陷就藏在它的
// trae 分支里。
func (h *traeHarness) beginAndPollWithCallback(cbQuery map[string]string) (int, string) {
	h.t.Helper()
	out, err := h.ps.InvokeHandler("channels", "start_login", &lua.HandlerRequest{
		Method: "POST", Body: `{"platform":"trae"}`,
	})
	if err != nil {
		h.t.Fatal(err)
	}
	j, ok := out.JSON.(map[string]interface{})
	if !ok {
		h.t.Fatalf("start_login 返回非对象: %#v", out.JSON)
	}
	sess := j["session"].(map[string]interface{})
	sid, _ := sess["id"].(string)
	nonce, _ := sess["nonce"].(string)
	if sid == "" {
		h.t.Fatalf("会话缺 id: %#v", sess)
	}

	// 摆好 loopback 回调文件（内容 = 平台重定向带回的 query）
	rec := map[string]string{"state": nonce}
	for k, v := range cbQuery {
		rec[k] = v
	}
	b, _ := json.Marshal(rec)
	cbPath := filepath.Join(h.t.TempDir(), "cb-"+sid+".json")
	if err := os.WriteFile(cbPath, b, 0600); err != nil {
		h.t.Fatal(err)
	}
	// 告诉插件「这个会话的 loopback 接收器把结果写在这里了」
	if _, err := h.ps.Call("channels", "__set_loopback_probe", map[string]interface{}{
		"session_id": sid, "file": cbPath,
	}); err != nil {
		h.t.Fatalf("设置 loopback 文件: %v", err)
	}

	res, err := h.ps.InvokeHandler("channels", "poll_login", &lua.HandlerRequest{
		Method: "GET", Query: map[string]string{"id": sid},
	})
	if err != nil {
		h.t.Fatal(err)
	}
	return res.Status, fmt.Sprintf("%v|%v", res.Body, res.JSON)
}

// ★ PKCE 探测的候选来源必须覆盖参考的 5 个（含 authCodeInfo 为**纯 code 字符串**
// 这一形态 —— 参考注释点名，解析不出 JSON 时必须回退到原始值）。
//
// 反例（我第一版就是错的）：用 `q.code or q.authCode or jsonField(...)` 链，
// 而 jsonField 取不到时返回**空串**；Lua 里空串是真值 ⇒ 链在第一个空串处短路，
// 后面的来源永远取不到。
func TestTraePkceDetectsPlainStringAuthCodeInfo(t *testing.T) {
	for _, tc := range []struct {
		name string
		q    map[string]string
	}{
		{"code 参数", map[string]string{"code": "C1"}},
		{"authCode 参数", map[string]string{"authCode": "C2"}},
		{"authCodeInfo 是 JSON", map[string]string{"authCodeInfo": `{"code":"C3"}`}},
		{"authCodeInfo 是 JSON(authCode)", map[string]string{"authCodeInfo": `{"authCode":"C4"}`}},
		{"authCodeInfo 是纯字符串", map[string]string{"authCodeInfo": "C5"}},
	} {
		t.Run(tc.name, func(t *testing.T) {
			h := newTraeHarness(t)
			status, body := h.beginAndPollWithCallback(tc.q)
			_ = status
			if n := atomic.LoadInt32(&h.exchHit); n != 0 {
				t.Fatalf("PKCE 形态不得去调 ExchangeToken（打了 %d 次）", n)
			}
			if !strings.Contains(body, "PKCE") {
				t.Fatalf("应识别为 PKCE 并明确报出，实得 %q", body)
			}
		})
	}
}