ModelRouter/internal/core/channels_trae_test.go

package core

import (
	"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，不是回调原值
		v, _, _ := h.ps.GetSecret("channels", "channels:trae")
		if !strings.HasPrefix(v, "ACCESS-") {
			t.Fatalf("凭据应是 ExchangeToken 的新 access, 实得 %q", v)
		}
	})

	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")
		if v != "JWT-TOKEN-1" {
			t.Fatalf("兜底凭据应是 userJwt.Token, 实得 %q", v)
		}
		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)
		}
	})
}