ModelRouter/internal/core/channels_trae_test.go

package core

import (
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync/atomic"
	"testing"

	"llmsproxy/internal/lua"
)

// traeHarness：TRAE callback 流的测试台。假上游模拟 ExchangeToken；
// 回调用 InvokeHandler 直打 oauth_callback（Query 表 = 平台重定向回来的参数）。
type traeHarness struct {
	t       *testing.T
	ps      *lua.Plugins
	ts      *httptest.Server
	exchHit int32
	exchOK  int32 // 1=成功 0=失败（atomic）
}

func newTraeHarness(t *testing.T) *traeHarness {
	t.Helper()
	h := &traeHarness{t: t, exchOK: 1}
	h.ts = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/cloudide/api/v3/trae/oauth/ExchangeToken" {
			w.WriteHeader(404)
			return
		}
		n := atomic.AddInt32(&h.exchHit, 1)
		w.Header().Set("Content-Type", "application/json")
		if atomic.LoadInt32(&h.exchOK) == 0 {
			w.WriteHeader(500)
			fmt.Fprint(w, `{"error":"exchange down"}`)
			return
		}
		// 校验请求体形状（对齐参考 exchangeTraeCallback 的请求构造）
		body, _ := io.ReadAll(r.Body)
		for _, want := range []string{`"ClientID"`, `"RefreshToken":"RT-FROM-CALLBACK"`, `"ClientSecret":"-"`} {
			if !strings.Contains(string(body), want) {
				t.Errorf("ExchangeToken 请求体缺 %s: %s", want, body)
			}
		}
		fmt.Fprintf(w, `{"Token":"ACCESS-%d","RefreshToken":"RT-NEW-%d"}`, n, n)
	}))
	t.Cleanup(h.ts.Close)

	c := kernelTestCore(t)
	h.ps = c.Plugins()
	src, err := readBundledChannels(t)
	if err != nil {
		t.Fatal(err)
	}
	if err := h.ps.LoadSource("channels", src); err != nil {
		t.Fatal(err)
	}
	_ = h.ps.SetPluginKernelPermission("channels", []string{"add_source", "status", "list_sources"}, false)
	// ExchangeToken 指到假上游（set_endpoints 既是运维能力也是测试缝）
	if _, err := h.ps.Call("channels", "set_endpoints", map[string]interface{}{
		"platform": "trae",
		"callback": map[string]interface{}{
			"exchange_url": h.ts.URL + "/cloudide/api/v3/trae/oauth/ExchangeToken",
		},
	}); err != nil {
		t.Fatalf("set_endpoints: %v", err)
	}
	return h
}

// beginAndCallback 发起登录（拿 nonce）后直接打 oauth_callback。
// cbQuery 是回调的 query（不含 state，内部补）；值即平台重定向回来的原始参数。
func (h *traeHarness) beginAndCallback(cbQuery map[string]string) (int, string) {
	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{})
	nonce := sess["nonce"].(string)

	q := map[string]string{"state": nonce}
	for k, v := range cbQuery {
		q[k] = v
	}
	res, err := h.ps.InvokeHandler("channels", "oauth_callback", &lua.HandlerRequest{
		Method: "GET", Query: q,
	})
	if err != nil {
		h.t.Fatal(err)
	}
	return res.Status, fmt.Sprintf("%v|%v", res.Body, res.JSON)
}

// TestTraeCallbackBehavior 是 TRAE callback 流的行为对齐测试。
//
// 决策表逐行搬自参考实现 parseTraeCallbackDetailed + exchangeTraeCallback
// （trae-oauth.ts）。断言的是行为（凭据来源 / 错误文案 / 上游是否被调），不是实现细节。
func TestTraeCallbackBehavior(t *testing.T) {
	h := newTraeHarness(t)

	t.Run("行1 refreshToken直传_ExchangeToken换新", func(t *testing.T) {
		atomic.StoreInt32(&h.exchOK, 1)
		st, body := h.beginAndCallback(map[string]string{
			"refreshToken": "RT-FROM-CALLBACK",
			"userInfo":     `{"UserID":"u1","ScreenName":"张三","TenantID":"t1"}`,
		})
		if st != 200 {
			t.Fatalf("应 200: %d %s", st, body)
		}
		// 凭据必须是 ExchangeToken 的新 Token，不是回调原值。
		// 凭据是 JSON（Trae 需 device_id/uid 供签到派生设备身份），故解出 access_token。
		v, _, _ := h.ps.GetSecret("channels", "channels:trae")
		var cred map[string]interface{}
		if err := json.Unmarshal([]byte(v), &cred); err != nil {
			t.Fatalf("凭据应为 JSON（含 device_id/uid）: %v (%q)", err, v)
		}
		if at, _ := cred["access_token"].(string); !strings.HasPrefix(at, "ACCESS-") {
			t.Fatalf("凭据应是 ExchangeToken 的新 access, 实得 %q", at)
		}
		// uid 来自 userInfo.UserID —— 签到设备身份的派生种子，丢了就无法签到
		if cred["uid"] != "u1" {
			t.Fatalf("凭据应保留 uid=u1（签到设备身份的种子）, 实得 %v", cred["uid"])
		}
		if cred["enterprise_id"] != "t1" {
			t.Fatalf("凭据应保留 TenantID, 实得 %v", cred["enterprise_id"])
		}
		// 设备身份来自登录会话（授权 URL 里下发的那对 hex32）
		if d, _ := cred["device_id"].(string); len(d) != 32 {
			t.Fatalf("凭据应带 32 位 device_id（签到轮换用）, 实得 %q", d)
		}
	})

	t.Run("行2 无refreshToken_userJwtToken兜底_不打ExchangeToken", func(t *testing.T) {
		before := atomic.LoadInt32(&h.exchHit)
		st, body := h.beginAndCallback(map[string]string{
			"userJwt": `{"Token":"JWT-TOKEN-1","RefreshToken":""}`,
		})
		if st != 200 {
			t.Fatalf("应 200: %d %s", st, body)
		}
		v, _, _ := h.ps.GetSecret("channels", "channels:trae")
		var cred2 map[string]interface{}
		if err := json.Unmarshal([]byte(v), &cred2); err != nil {
			t.Fatalf("凭据应为 JSON: %v (%q)", err, v)
		}
		if cred2["access_token"] != "JWT-TOKEN-1" {
			t.Fatalf("兜底凭据应是 userJwt.Token, 实得 %v", cred2["access_token"])
		}
		if atomic.LoadInt32(&h.exchHit) != before {
			t.Fatal("无 refreshToken 不应调 ExchangeToken（参考实现分支2）")
		}
	})

	t.Run("行3 只有code_PKCE流_明确报暂不支持", func(t *testing.T) {
		st, body := h.beginAndCallback(map[string]string{"code": "PKCE-CODE-1"})
		if st != 400 {
			t.Fatalf("PKCE 流应 400: %d %s", st, body)
		}
		if !strings.Contains(body, "PKCE") {
			t.Fatalf("错误应明确指向 PKCE 流, 实得 %s", body)
		}
	})

	t.Run("行4 空回调_报缺失", func(t *testing.T) {
		st, body := h.beginAndCallback(map[string]string{})
		if st != 400 {
			t.Fatalf("空回调应 400: %d %s", st, body)
		}
		if !strings.Contains(body, "refreshToken") {
			t.Fatalf("错误应指出缺什么, 实得 %s", body)
		}
	})

	t.Run("行5 ExchangeToken失败_透传阶段", func(t *testing.T) {
		atomic.StoreInt32(&h.exchOK, 0)
		defer atomic.StoreInt32(&h.exchOK, 1)
		st, body := h.beginAndCallback(map[string]string{"refreshToken": "RT-FROM-CALLBACK"})
		if st != 400 {
			t.Fatalf("exchange 失败应 400: %d %s", st, body)
		}
		if !strings.Contains(body, "ExchangeToken") {
			t.Fatalf("错误应指明 ExchangeToken 阶段, 实得 %s", body)
		}
	})
}

// TestTraeCallbackPoolPath 回归：callback 流的**池路径**（UI 真实流程）。
//
// 与 buddy poll_login 的池路径同一缺陷：complete_login 不带 account_id 时，
// 凭据落到与占位账号无关的 `channels:<plat>` 键，池账号永远 has_credential=false。
// 此前只有 device-poll / state-poll 分支修了，callback 分支是漏网。
func TestTraeCallbackPoolPath(t *testing.T) {
	h := newTraeHarness(t)
	atomic.StoreInt32(&h.exchOK, 1)

	// UI 流程：先建占位账号，再带 account_id 发起登录
	add, err := h.ps.InvokeHandler("channels", "account_add", &lua.HandlerRequest{
		Method: "POST", Body: `{"platform":"trae"}`,
	})
	if err != nil {
		t.Fatal(err)
	}
	aid := add.JSON.(map[string]interface{})["account_id"].(string)

	out, err := h.ps.InvokeHandler("channels", "start_login", &lua.HandlerRequest{
		Method: "POST", Body: fmt.Sprintf(`{"platform":"trae","account_id":%q}`, aid),
	})
	if err != nil {
		t.Fatal(err)
	}
	nonce := out.JSON.(map[string]interface{})["session"].(map[string]interface{})["nonce"].(string)

	st, body := h.invokeCallback(map[string]string{
		"state": nonce, "refreshToken": "RT-FROM-CALLBACK", "userInfo": `{"UserID":"u9","ScreenName":"池用户"}`,
	})
	if st != 200 {
		t.Fatalf("回调应 200: %d %s", st, body)
	}

	// ① 凭据必须落在池账号的 ref
	v, ok, _ := h.ps.GetSecret("channels", "channels:trae:"+aid)
	if !ok || v == "" {
		t.Fatalf("callback 池路径凭据未落到 channels:trae:%s（用户报障的同一缺陷）", aid)
	}
	// ② account_list 必须报 has_credential=true
	lst, err := h.ps.InvokeHandler("channels", "account_list", &lua.HandlerRequest{
		Method: "POST", Body: `{"platform":"trae"}`,
	})
	if err != nil {
		t.Fatal(err)
	}
	for _, raw := range lst.JSON.(map[string]interface{})["accounts"].([]interface{}) {
		a := raw.(map[string]interface{})
		if a["id"] == aid && a["has_credential"] != true {
			t.Fatalf("has_credential 应为 true, 实得 %v", a["has_credential"])
		}
	}
}

// invokeCallback 直打 oauth_callback（Query 已含 state）。
func (h *traeHarness) invokeCallback(q map[string]string) (int, string) {
	res, err := h.ps.InvokeHandler("channels", "oauth_callback", &lua.HandlerRequest{
		Method: "GET", Query: q,
	})
	if err != nil {
		h.t.Fatal(err)
	}
	return res.Status, fmt.Sprintf("%v|%v", res.Body, res.JSON)
}