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) } }) } // TestTraeCallbackPoolPath 回归:callback 流的**池路径**(UI 真实流程)。 // // 与 buddy poll_login 的池路径同一缺陷:complete_login 不带 account_id 时, // 凭据落到与占位账号无关的 `channels:` 键,池账号永远 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) }