feat(channels): TRAE callback 真实移植 —— token 直传 + ExchangeToken（照参考逐行）

此前 oauth_callback 是通用 OAuth 形态（直接拿 code 当凭据），与 TRAE 真实行为不符。参考实现 trae-oauth.ts 的回调是 **token 直传**，本次照其行为移植： parseTraeCallbackDetailed（决策表 5 行）：行1 refreshToken 直传（或 userJwt.RefreshToken 回退）→ ExchangeToken 换新 access（请求体 ClientID/RefreshToken/ClientSecret:'-'/UserID 照抄，响应取 Token）→ 凭据 = 新 access 行2 无 refreshToken 但有 userJwt.Token → 直接兜底，**不打** ExchangeToken 行3 只有 code/authCodeInfo → PKCE 流，明确报"暂不支持"而非含糊失败（参考注释强调：判为无效会让用户往错误方向排查）行4 全空 → 报"回调未携带 refreshToken / userJwt.Token / code" 行5 ExchangeToken HTTP 失败 → 错误指明 ExchangeToken 阶段配套： - PLATFORMS.trae 增加 callback = {mode="trae", exchange_url, exchange_body}； - oauth_callback 按平台 callback.mode 分派（nil→code 旧路径不变，googleish 不受影响）； - traeCallback/jsonField 新函数；set_endpoints 支持 callback.exchange_url 覆盖（既是运维能力也是测试缝，与 device 端点同语义）； - TestTraeCallbackBehavior：5 行决策表 + 行1 断言凭据来自假上游新 Token （防"把回调原值当凭据"）、行2 断言 ExchangeToken 零调用（防多余上游请求）。回归：channels/trae/device 变异/压力/legacy/保险库/e2e 全绿。
2026-11-07 09:47:23 +00:00 · 2026-10-11 22:35:18 +08:00
parent 2f5cde7fea
commit 893857464d
2 changed files with 282 additions and 8 deletions
--- a/internal/core/channels_trae_test.go
+++ b/internal/core/channels_trae_test.go
@ -0,0 +1,173 @@
+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)
+		}
+	})
+}
--- a/internal/lua/plugins/channels.lua
+++ b/internal/lua/plugins/channels.lua
@ -121,6 +121,17 @@ local PLATFORMS = {
      url = "https://api.trae.cn/trae/api/v2/ug/checkin_credits/claim",
      method = "POST", body = "{}",
    },
+    -- callback 解析模式（移植自参考 trae-oauth.ts parseTraeCallbackDetailed）：
+    --   trae    token 直传：query 带 refreshToken / userJwt(JSON) / userInfo(JSON)。
+    --           refreshToken 存在时调 ExchangeToken 换新（轮换旧 refresh）；
+    --           无 refreshToken 用 userJwt.Token 兑底；
+    --           只有 code/authCodeInfo 时是 PKCE 流 → 明确报"暂不支持"而不是含糊失败。
+    callback = {
+      mode = "trae",
+      exchange_url = "https://api.trae.com.cn/cloudide/api/v3/trae/oauth/ExchangeToken",
+      exchange_body = '{"ClientID":"{client_id}","RefreshToken":"{refresh_token}","ClientSecret":"-","UserID":""}',
+      client_id = "",  -- 参考用 product.clientId；TRAE 回调直传形态不强制，留空
+    },
  },

  -- MiniMax —— PKCE 设备码流（RFC 8628）。三步：申请码 → 用户授权 → 轮询换
@ -731,6 +742,11 @@ function plugin.set_endpoints(req)
      if args.checkin[k] then plat.checkin[k] = tostring(args.checkin[k]) changed[#changed+1] = "checkin." .. k end
    end
  end
+  if type(args.callback) == "table" and plat.callback then
+    for _, k in ipairs({ "exchange_url" }) do
+      if args.callback[k] then plat.callback[k] = tostring(args.callback[k]) changed[#changed+1] = "callback." .. k end
+    end
+  end
  if #changed == 0 then return { status = 400, json = { error = "无可识别的端点字段" } } end
  logEvent("端点已改写 " .. plat.id .. ": " .. table.concat(changed, ","))
  return { status = 200, json = { ok = true, changed = changed } }
@ -761,10 +777,15 @@ function plugin.status_page(req)
  } }
 end

-- oauth_callback 处理平台重定向：校验 nonce，然后交给 complete_login。
+-- oauth_callback 处理平台重定向：校验 nonce，然后按平台的 callback 模式解析。
 --
 -- nonce 校验是必需的：没有它，别人构造一个回调就能把**自己的**凭据塞进
 -- 这个网关的登录流程。
+--
+-- callback 模式（platform.callback.mode）决定如何从回调 query 解出凭据：
+--   nil/"code"  标准 OAuth：直接把 code 当凭据交给 complete_login。
+--   "trae"      字节 TRAE 的 token 直传形态，解析逐字对齐参考实现
+--               parseTraeCallbackDetailed + exchangeTraeCallback（见 traeCallback）。
 function plugin.oauth_callback(req)
  local nonce = req.query and req.query.state
  if not nonce or nonce == "" then
@ -781,19 +802,99 @@ function plugin.oauth_callback(req)
  if s.flow ~= "callback" then
    return { status = 400, body = "session is not a callback flow" }
  end
-  local code = req.query.code
-  if not code or code == "" then
-    return { status = 400, body = "missing code" }
+  local plat = platformById(s.platform)
+  local mode = plat and plat.callback and plat.callback.mode or "code"
+
+  local credential, account, cerr
+  if mode == "trae" then
+    credential, account, cerr = traeCallback(plat, req.query)
+  else
+    credential = req.query.code
+    account = "oauth"
+    if not credential or credential == "" then cerr = "missing code" end
  end
-  local res = plugin.complete_login({ platform = s.platform, credential = code, account = "oauth" })
-  if res.error then
-    return { status = 400, body = res.error }
+  if cerr then
+    s.state = "failed"
+    s.message = cerr
+    return { status = 400, body = cerr }
+  end
+
+  local res = plugin.complete_login({ platform = s.platform, credential = credential, account = account })
+  if res.json and res.json.error then
+    return { status = 400, body = res.json.error }
  end
  s.state = "completed"
-  s.account = "oauth"
+  s.account = account
  return { status = 200, body = "ok" }
 end

+-- traeCallback 解析 TRAE 的 token 直传回调并（必要时）走 ExchangeToken。
+-- 返回 (credential, account, nil) 或 (nil, nil, err)。
+--
+-- 行为逐行对齐参考实现 trae-oauth.ts：
+--   refreshToken（query 或 userJwt.RefreshToken 回退）存在 → ExchangeToken 换新
+--     access（并轮换 refreshToken）；凭据 = 新 access。
+--   无 refreshToken 但有 userJwt.Token → 直接用它兜底（不走 ExchangeToken）。
+--   二者皆无、但有 code/authCodeInfo → **PKCE 流**，明确报"暂不支持"而非含糊失败。
+--   都没有 → 回调无效。
+local function jsonField(raw, key)
+  if not raw or raw == "" then return "" end
+  local ok, v = pcall(json.decode, raw)
+  if not ok or type(v) ~= "table" then return "" end
+  return v[key] or ""
+end
+
+function traeCallback(plat, q)
+  q = q or {}
+  local userJwt = q.userJwt or ""
+  local refreshToken = q.refreshToken
+  if not refreshToken or refreshToken == "" then
+    refreshToken = jsonField(userJwt, "RefreshToken")
+  end
+  local jwtToken = jsonField(userJwt, "Token")
+  local account = jsonField(q.userInfo or "", "ScreenName")
+  if account == "" then account = "trae" end
+
+  -- PKCE 形态探测：仅在 refreshToken / jwtToken 都空时才判，与参考一致。
+  local authCode = q.code or q.authCode or jsonField(q.authCodeInfo or "", "code")
+
+  if (not refreshToken or refreshToken == "") and (not jwtToken or jwtToken == "") then
+    if authCode and authCode ~= "" then
+      return nil, nil, "上游返回了 PKCE 授权码（code/authCodeInfo），本实现暂不支持该流程"
+    end
+    return nil, nil, "回调未携带 refreshToken / userJwt.Token / code"
+  end
+
+  -- 分支 1：有 refreshToken → ExchangeToken 换新 access（轮换 refresh）
+  if refreshToken and refreshToken ~= "" and plat.callback.exchange_url then
+    local body = (plat.callback.exchange_body
+      :gsub("{client_id}", plat.callback.client_id or "")
+      :gsub("{refresh_token}", refreshToken))
+    local resp, err = __llmsproxy_http({
+      url = plat.callback.exchange_url,
+      method = "POST",
+      headers = { ["Content-Type"] = "application/json", ["Accept"] = "application/json" },
+      body = body,
+    })
+    if err then return nil, nil, "TRAE ExchangeToken 失败: " .. err end
+    if resp.status ~= 200 then
+      return nil, nil, "TRAE ExchangeToken HTTP " .. tostring(resp.status)
+    end
+    local ok, rec = pcall(json.decode, resp.body or "")
+    if not ok or type(rec) ~= "table" then
+      return nil, nil, "TRAE ExchangeToken 响应无法解析"
+    end
+    local access = rec.Token or rec.token or rec.AccessToken or rec.access_token
+    if not access or access == "" then
+      return nil, nil, "TRAE ExchangeToken 响应缺少 Token"
+    end
+    return access, account, nil
+  end
+
+  -- 分支 2：无 refreshToken，用 userJwt.Token 兜底
+  return jwtToken, account, nil
+end
+
 -- plugin.on_unload：插件被卸载/重载时回收自己拉起的浏览器栈。进程是插件起的，
 -- 收尾也是插件的责任 —— 内核不应知道这些进程的存在。
 plugin.on_unload = function()