diff --git a/internal/core/channels_trae_test.go b/internal/core/channels_trae_test.go new file mode 100644 index 0000000..9666391 --- /dev/null +++ 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) + } + }) +} diff --git a/internal/lua/plugins/channels.lua b/internal/lua/plugins/channels.lua index c18c4a6..33e2335 100644 --- 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()