diff --git a/changelog/unreleased/trae-login-checkin-fixes.md b/changelog/unreleased/trae-login-checkin-fixes.md new file mode 100644 index 0000000..e4faac8 --- /dev/null +++ b/changelog/unreleased/trae-login-checkin-fixes.md @@ -0,0 +1,7 @@ +### English + +- Trae browser login now completes on its own. The loopback listener only parked the authorization code and waited for a console poll to redeem it, but the console's poll window is far shorter than a real browser login (an SMS login can take minutes), so the poll expired first and the code was never exchanged. The callback now exchanges the code and stores the credential as soon as the browser reaches the loopback listener, so closing the browser or missing the poll window no longer costs the login. A pasted callback URL still works as a fallback and now actually redeems the parked code; a callback URL carrying only `loginTraceID` is accepted for that path; and the single-use code is exchanged exactly once even when the callback, a poll, and a paste overlap. + +### 中文 + +- Trae 浏览器登录现在会自行完成。此前 loopback 监听只把授权码暂存下来，等待控制台轮询去换取令牌，但控制台的轮询窗口远短于真实浏览器登录（短信登录可能耗时数分钟），轮询先超时退出，授权码因此从未被换取。现在浏览器一到达 loopback 监听，回调就会立即换取授权码并写入凭据，关闭浏览器或错过轮询窗口都不再导致登录失败。粘贴回调 URL 仍可作为兜底，并且现在会真正换取已暂存的授权码；仅携带 `loginTraceID` 的回调 URL 也被接受用于该路径；即便回调、轮询与粘贴同时发生，单次使用的授权码也只会被换取一次。 diff --git a/internal/providers/trae/client.go b/internal/providers/trae/client.go index 9f6625d..f2de70d 100644 --- a/internal/providers/trae/client.go +++ b/internal/providers/trae/client.go @@ -4,6 +4,7 @@ import ( "bytes" "context" "encoding/json" + "errors" "fmt" "io" "log" @@ -47,6 +48,7 @@ type ModelMaxModeReader interface { } type loginPending struct { + accountID string machineID string deviceID string deviceKey string @@ -59,6 +61,12 @@ type loginPending struct { failed bool message string credential Credential + // redeemDone closes once a redemption attempt finishes, so a console poll + // or a pasted callback can wait for an exchange already in flight instead + // of starting a second one against the same single-use code. redeemErr + // carries that attempt's outcome. + redeemDone chan struct{} + redeemErr error } type Client struct { @@ -200,6 +208,7 @@ func (c *Client) StartLogin(ctx context.Context, accountID string) (providers.Lo c.pendingTraces = map[string]*loginPending{} } entry := &loginPending{ + accountID: accountID, machineID: credential.MachineID, deviceID: credential.DeviceID, deviceKey: credential.DevicePrivateKey, @@ -220,12 +229,18 @@ func (c *Client) PollLogin(ctx context.Context, accountID string) (bool, string, pending := c.pending[accountID] c.mu.Unlock() if pending == nil { + // No session in flight. A pasted callback may have redeemed this login + // already, in which case the account is simply done; otherwise there is + // nothing to poll. + if _, payload, err := c.store.LoadCredentialPayload(ctx, accountID); err == nil { + if decoded, err := DecodeCredential(payload); err == nil && decoded.Ready() { + return true, "login complete", nil + } + } return false, "", fmt.Errorf("login not started for account %s", accountID) } if time.Since(pending.createdAt) > loginPendingTTL { - c.mu.Lock() - delete(c.pending, accountID) - c.mu.Unlock() + c.forgetPending(accountID, pending) return false, "", fmt.Errorf("login expired; start again") } c.mu.Lock() @@ -242,13 +257,61 @@ func (c *Client) PollLogin(ctx context.Context, accountID string) (bool, string, if !done { return false, firstNonEmpty(message, "waiting for authorization"), nil } - if err := c.finishCredential(ctx, accountID, credential, authCode, codeVerifier); err != nil { + if err := c.redeem(ctx, accountID, pending, credential, authCode, codeVerifier); err != nil { return false, "", err } + c.forgetPending(accountID, pending) + return true, "login complete", nil +} + +// redeem exchanges the authorization code parked on a pending entry exactly +// once. The loopback callback, a console poll, and a pasted callback can all +// reach the same code within the same second, and the code is single-use, so +// the first caller performs the exchange and the others wait for its outcome. +// A failed attempt clears the guard: a transient upstream error must not burn +// the operator's only chance to retry the code. +func (c *Client) redeem(ctx context.Context, accountID string, pending *loginPending, credential Credential, authCode, codeVerifier string) error { c.mu.Lock() - delete(c.pending, accountID) + if inFlight := pending.redeemDone; inFlight != nil { + c.mu.Unlock() + select { + case <-inFlight: + c.mu.Lock() + err := pending.redeemErr + c.mu.Unlock() + return err + case <-ctx.Done(): + return ctx.Err() + } + } + done := make(chan struct{}) + pending.redeemDone = done c.mu.Unlock() - return true, "login complete", nil + + err := c.finishCredential(ctx, accountID, credential, authCode, codeVerifier) + + c.mu.Lock() + pending.redeemErr = err + close(done) + if err != nil { + pending.redeemDone = nil + } + c.mu.Unlock() + return err +} + +// forgetPending drops a pending entry together with the trace index pointing at +// it, so a redeemed or expired login leaves nothing behind for a later pasted +// callback to match. +func (c *Client) forgetPending(accountID string, pending *loginPending) { + c.mu.Lock() + defer c.mu.Unlock() + if current := c.pending[accountID]; current == nil || current == pending { + delete(c.pending, accountID) + } + if pending != nil && pending.trace != "" && c.pendingTraces[pending.trace] == pending { + delete(c.pendingTraces, pending.trace) + } } func (c *Client) CompleteLogin(ctx context.Context, accountID, callbackURL string) error { @@ -279,16 +342,39 @@ func (c *Client) CompleteLogin(ctx context.Context, accountID, callbackURL strin c.mu.Unlock() codeVerifier := "" if pending != nil { - if pending.done { - c.mu.Lock() - delete(c.pending, accountID) - c.mu.Unlock() - return nil - } credential.MachineID = pending.machineID credential.DeviceID = pending.deviceID credential.DevicePrivateKey = pending.deviceKey codeVerifier = pending.codeVerifier + // The loopback listener may have captured this callback already and + // parked the code on the pending entry, while the console's poll window + // expired before PollLogin could redeem it. That code is still + // unexchanged, so fill the gaps from the captured callback and let + // finishCredential redeem it. Returning early here reported success + // while storing no credential at all. + if pending.done { + if strings.TrimSpace(info.AuthCode) == "" { + info.AuthCode = pending.authCode + } + if credential.AccessToken == "" { + credential.AccessToken = pending.credential.AccessToken + } + if credential.RefreshToken == "" { + credential.RefreshToken = pending.credential.RefreshToken + } + if credential.ExpiresAt == 0 { + credential.ExpiresAt = pending.credential.ExpiresAt + } + if credential.UID == "" { + credential.UID = pending.credential.UID + } + if credential.Nickname == "" { + credential.Nickname = pending.credential.Nickname + } + if credential.EnterpriseID == "" { + credential.EnterpriseID = pending.credential.EnterpriseID + } + } } else if _, payload, err := c.store.LoadCredentialPayload(ctx, accountID); err == nil { if decoded, err := DecodeCredential(payload); err == nil { if decoded.Ready() { @@ -298,15 +384,16 @@ func (c *Client) CompleteLogin(ctx context.Context, accountID, callbackURL strin credential.DeviceID = decoded.DeviceID } } + if pending != nil { + if err := c.redeem(ctx, accountID, pending, credential, info.AuthCode, codeVerifier); err != nil { + return err + } + c.forgetPending(accountID, pending) + return nil + } if err := c.finishCredential(ctx, accountID, credential, info.AuthCode, codeVerifier); err != nil { return err } - c.mu.Lock() - delete(c.pending, accountID) - if pending != nil && pending.trace != "" { - delete(c.pendingTraces, pending.trace) - } - c.mu.Unlock() return nil } @@ -325,6 +412,10 @@ func (c *Client) finishCredential(ctx context.Context, accountID string, credent return err } credential = refreshed + case strings.TrimSpace(credential.AccessToken) == "": + // Neither a code to redeem nor a token to verify. Storing this would + // persist an empty credential and report the login as complete. + return fmt.Errorf("callback carried no authorization code or token") } if strings.TrimSpace(credential.UID) == "" && strings.TrimSpace(credential.AccessToken) != "" { info, err := c.GetUserInfo(ctx, accountID, credential) @@ -345,6 +436,7 @@ func (c *Client) finishCredential(ctx context.Context, accountID string, credent return err } if err := c.store.SaveCredentialPayload(ctx, accountID, CredentialFormat, payload); err != nil { + log.Printf("trae: SaveCredentialPayload failed account=%s: %v", accountID, err) return err } _ = c.store.Observe(ctx, accountID, credential.UID, "ready", "", "") @@ -373,6 +465,9 @@ func (c *Client) ensureCallback() (string, error) { func (c *Client) acceptCallback(ctx context.Context, rawURL string) error { info, err := ParseCallback(rawURL) if err != nil { + // Do not log the parse error verbatim: it embeds the raw callback URL, + // which carries the one-shot auth code and OAuth state. + log.Printf("trae: acceptCallback parse failed") c.markPendingFailed(err.Error()) return err } @@ -387,22 +482,74 @@ func (c *Client) acceptCallback(ctx context.Context, rawURL string) error { APIHost: OAuthHost, } c.mu.Lock() - for _, pending := range c.pending { - if pending.done || pending.failed { - continue + // The listener is shared by every account, so pair the callback with the + // pending entry that minted its challenge. Falling back to map order would + // hand one account's code to another account's login when two are in flight. + var claimed *loginPending + if info.Trace != "" { + if traced := c.pendingTraces[info.Trace]; traced != nil && !traced.done && !traced.failed { + claimed = traced } - credential.MachineID = pending.machineID - credential.DeviceID = pending.deviceID - pending.credential = credential - pending.authCode = info.AuthCode - pending.done = true - pending.message = "authorization received" - break + } + if claimed == nil { + for _, pending := range c.pending { + if pending.done || pending.failed { + continue + } + claimed = pending + break + } + } + if claimed != nil { + credential.MachineID = claimed.machineID + credential.DeviceID = claimed.deviceID + // The v3 exchange signs with the device key minted for this login; the + // public half goes out as DeviceInfo.DevicePublicKey. Dropping it here + // sent the exchange without a device proof, matching CompleteLogin's + // handling of the same field. + credential.DevicePrivateKey = claimed.deviceKey + claimed.credential = credential + claimed.authCode = info.AuthCode + claimed.done = true + claimed.message = "authorization received" } c.mu.Unlock() + if claimed == nil { + return nil + } + // Redeem here rather than waiting for a poller. The console's poll window is + // far shorter than a real browser login, so the code would otherwise sit + // unexchanged until the operator pasted the callback URL by hand. The + // request context is the loopback request's, which ends when the browser + // disconnects, so the exchange gets its own budget. + go c.redeemAsync(claimed) return nil } +// redeemAsync redeems a captured callback on the callback path. The account id +// comes from the pending entry, which is registered under it in StartLogin. +func (c *Client) redeemAsync(pending *loginPending) { + ctx, cancel := context.WithTimeout(context.WithoutCancel(context.Background()), redeemTimeout) + defer cancel() + c.mu.Lock() + credential, authCode, codeVerifier := pending.credential, pending.authCode, pending.codeVerifier + accountID := pending.accountID + c.mu.Unlock() + if err := c.redeem(ctx, accountID, pending, credential, authCode, codeVerifier); err != nil { + // Log only the classified status/code: the full error message can embed + // the upstream response payload. + var perr *providers.Error + if errors.As(err, &perr) { + log.Printf("trae: redeemAsync failed account=%s status=%d code=%s", accountID, perr.Status, perr.Code) + } else { + log.Printf("trae: redeemAsync failed account=%s (transport error)", accountID) + } + c.markPendingFailedFor(pending, err.Error()) + return + } + c.forgetPending(accountID, pending) +} + func (c *Client) markPendingFailed(message string) { c.mu.Lock() defer c.mu.Unlock() @@ -416,6 +563,15 @@ func (c *Client) markPendingFailed(message string) { } } +// markPendingFailedFor records a failure against one specific entry, leaving +// other accounts' in-flight logins untouched. +func (c *Client) markPendingFailedFor(pending *loginPending, message string) { + c.mu.Lock() + defer c.mu.Unlock() + pending.failed = true + pending.message = message +} + type callbackInfo struct { AuthCode string RefreshToken string @@ -497,7 +653,10 @@ func ParseCallback(rawURL string) (callbackInfo, error) { } } } - if info.RefreshToken == "" && info.AccessToken == "" && info.AuthCode == "" { + // A callback carrying only loginTraceID is still usable: CompleteLogin + // matches the trace to the pending entry and redeems the code the loopback + // listener parked there. Anything without a trace needs its own material. + if info.RefreshToken == "" && info.AccessToken == "" && info.AuthCode == "" && info.Trace == "" { return callbackInfo{}, fmt.Errorf("callback missing code, refreshToken and userJwt.Token") } return info, nil @@ -613,6 +772,7 @@ func (c *Client) ExchangeAuthCode(ctx context.Context, accountID string, credent "DeviceModel": "", "ClientVersion": IdeVersion, "DeviceBrand": "", + "DeviceCPU": "", "OSInfo": "linux", "OSVersion": "Ubuntu 24.04.4 LTS", } @@ -629,9 +789,32 @@ func (c *Client) ExchangeAuthCode(ctx context.Context, accountID string, credent if err != nil { return credential, err } - payload, status, err := c.do(ctx, accountID, http.MethodPost, credential.AuthBase()+pathExchangeCode, body, SetOAuthHeaders) - if err != nil { - return credential, err + // The official CN client exchanges at the host product.json configures + // (bootConfig.account.trae.normal = https://api.trae.cn), while the callback + // advertises api.trae.com.cn. Try the configured host first and fall back to + // the callback's only when that host rejects the request as invalid — a + // transient upstream failure must surface instead of being masked by a + // second attempt. + hosts := []string{UgHost} + if base := credential.AuthBase(); base != "" && base != UgHost { + hosts = append(hosts, base) + } + var payload []byte + var status int + for i, host := range hosts { + payload, status, err = c.do(ctx, accountID, http.MethodPost, host+pathExchangeCode, body, SetOAuthHeaders) + if err != nil { + log.Printf("trae: ExchangeAuthCode request failed account=%s host=%s: %v", accountID, host, err) + return credential, err + } + if status < 300 { + break + } + // Log status and the extracted error code only — never the raw payload. + log.Printf("trae: ExchangeAuthCode upstream error account=%s host=%s status=%d code=%s", accountID, host, status, firstNonEmpty(extractMetadataCode(payload), extractCode(string(payload)))) + if i == len(hosts)-1 || !isInvalidParameter(status, payload) { + return credential, classifiedError(status, payload) + } } if status >= 300 { return credential, classifiedError(status, payload) @@ -1031,6 +1214,41 @@ func classifiedError(status int, body []byte) error { return wrapClassified(Classify(status, string(body)), extractCode(string(body))) } +// isInvalidParameter reports the upstream's "无效参数" rejection (code 10101). +// The CN front doors disagree about which host mints a code, and that +// rejection is what the wrong host returns, so it is the one case worth +// retrying elsewhere. Everything else (rate limits, 5xx, auth) is real and +// must surface. +func isInvalidParameter(status int, body []byte) bool { + if status != http.StatusBadRequest { + return false + } + return extractMetadataCode(body) == "10101" +} + +// extractMetadataCode reads the code Trae puts under +// ResponseMetadata.Error.Code. extractCode handles the flat {"code": ...} +// shape other endpoints use; this is the envelope the v3 exchange returns. +func extractMetadataCode(body []byte) string { + var env struct { + ResponseMetadata struct { + Error struct { + Code any `json:"Code"` + } `json:"Error"` + } `json:"ResponseMetadata"` + } + if json.Unmarshal(body, &env) != nil { + return "" + } + switch v := env.ResponseMetadata.Error.Code.(type) { + case string: + return v + case float64: + return fmt.Sprintf("%.0f", v) + } + return "" +} + func wrapClassified(classified providers.ClassifiedError, code string) error { return &providers.Error{ Kind: classified.Kind, diff --git a/internal/providers/trae/client_test.go b/internal/providers/trae/client_test.go index 3e7638b..4090daf 100644 --- a/internal/providers/trae/client_test.go +++ b/internal/providers/trae/client_test.go @@ -8,6 +8,8 @@ import ( "net/http" "net/http/httptest" "strings" + "sync" + "sync/atomic" "testing" "time" @@ -17,6 +19,9 @@ import ( ) type memStore struct { + // mu guards every field below. Redemption runs on a background goroutine + // after the loopback callback, so the store is reached concurrently. + mu sync.Mutex items map[string][]byte region string settings map[string]accounts.ProviderModelSetting @@ -30,6 +35,8 @@ type memStore struct { } func (s *memStore) Get(ctx context.Context, id string) (accounts.Account, error) { + s.mu.Lock() + defer s.mu.Unlock() region := s.region if region == "" { region = "cn" @@ -37,6 +44,8 @@ func (s *memStore) Get(ctx context.Context, id string) (accounts.Account, error) return accounts.Account{ID: id, Provider: "trae", ProviderRegion: region, ProxyURL: s.accountProxyURL}, nil } func (s *memStore) GetSecret(ctx context.Context, key string) (string, bool, error) { + s.mu.Lock() + defer s.mu.Unlock() s.secretCalls++ if s.secretErr != nil { return "", false, s.secretErr @@ -44,6 +53,8 @@ func (s *memStore) GetSecret(ctx context.Context, key string) (string, bool, err return s.secretValue, s.secretFound, nil } func (s *memStore) LoadCredentialPayload(ctx context.Context, accountID string) (string, []byte, error) { + s.mu.Lock() + defer s.mu.Unlock() payload, ok := s.items[accountID] if !ok { return "", nil, accounts.ErrAccountNotFound @@ -51,6 +62,8 @@ func (s *memStore) LoadCredentialPayload(ctx context.Context, accountID string) return CredentialFormat, payload, nil } func (s *memStore) SaveCredentialPayload(ctx context.Context, accountID, format string, payload []byte) error { + s.mu.Lock() + defer s.mu.Unlock() if s.items == nil { s.items = map[string][]byte{} } @@ -61,6 +74,8 @@ func (s *memStore) Observe(ctx context.Context, id, remoteUID, status, lastError return nil } func (s *memStore) GetProviderModelSetting(ctx context.Context, provider, modelID string) (accounts.ProviderModelSetting, error) { + s.mu.Lock() + defer s.mu.Unlock() s.lookups = append(s.lookups, provider+"/"+modelID) if s.settings == nil { return accounts.ProviderModelSetting{}, nil @@ -209,6 +224,485 @@ type roundTripFunc func(*http.Request) (*http.Response, error) func (f roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) { return f(req) } +// mustKey returns a fresh EC device private key, the shape a real login stores +// and the code exchange signs with. +func mustKey(t *testing.T) string { + t.Helper() + priv, _, err := generateDeviceKey() + if err != nil { + t.Fatal(err) + } + return priv +} + +// A callback captured by the loopback listener but never redeemed by the poll +// loop (the console's poll window can expire first) must still be exchangeable +// through a pasted callback URL. Reporting success without storing a credential +// left the account permanently at "account not found". +func TestCompleteLoginRedeemsCodeCapturedByLoopback(t *testing.T) { + client, store := newTestClient(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case pathExchangeCode: + _ = json.NewEncoder(w).Encode(map[string]any{"Result": map[string]any{ + "Token": "at", "RefreshToken": "rt2", "TokenExpireAt": time.Now().Add(time.Hour).Unix(), + }}) + case pathUserInfo: + _ = json.NewEncoder(w).Encode(map[string]any{"Result": map[string]any{ + "UserID": "u1", "ScreenName": "Tester", "EnterpriseID": "e1", + }}) + default: + t.Fatalf("unexpected path %s", r.URL.Path) + } + })) + session, err := client.StartLogin(context.Background(), "acc1") + if err != nil { + t.Fatal(err) + } + // The browser hits the loopback listener: the callback is parked as done + // with the code attached, exactly as acceptCallback leaves it. + client.mu.Lock() + pending := client.pending["acc1"] + pending.done = true + pending.authCode = "ABC-123" + pending.credential = Credential{ + UID: "u1", Nickname: "Tester", EnterpriseID: "e1", + MachineID: pending.machineID, DeviceID: pending.deviceID, + } + machine, device := pending.machineID, pending.deviceID + client.mu.Unlock() + + // The operator pastes the callback URL afterwards; the poll loop never ran. + callback := "http://127.0.0.1:49519/authorize?loginTraceID=" + session.State + + "&authCodeInfo=%7B%22AuthCode%22%3A%22ABC-123%22%7D" + if err := client.CompleteLogin(context.Background(), "acc1", callback); err != nil { + t.Fatal(err) + } + _, payload, err := store.LoadCredentialPayload(context.Background(), "acc1") + if err != nil { + t.Fatalf("credential must be stored: %v", err) + } + credential, err := DecodeCredential(payload) + if err != nil { + t.Fatal(err) + } + if credential.UID != "u1" || credential.AccessToken != "at" || credential.RefreshToken != "rt2" { + t.Fatalf("credential=%+v", credential) + } + if credential.MachineID != machine || credential.DeviceID != device { + t.Fatalf("device identity must survive: got machine=%s device=%s want machine=%s device=%s", + credential.MachineID, credential.DeviceID, machine, device) + } + if credential.DevicePrivateKey == "" { + t.Fatal("device key must be stored for later code exchanges") + } +} + +// The pasted URL may carry nothing but the trace id; the code parked by the +// loopback listener is the only source of the authorization code. +func TestCompleteLoginFallsBackToCapturedAuthCode(t *testing.T) { + client, store := newTestClient(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case pathExchangeCode: + var body map[string]any + _ = json.NewDecoder(r.Body).Decode(&body) + if body["AuthCode"] != "PARKED-CODE" { + t.Fatalf("exchange must redeem the parked code, body=%v", body) + } + _ = json.NewEncoder(w).Encode(map[string]any{"Result": map[string]any{ + "Token": "at", "RefreshToken": "rt2", "TokenExpireAt": time.Now().Add(time.Hour).Unix(), + }}) + case pathUserInfo: + _ = json.NewEncoder(w).Encode(map[string]any{"Result": map[string]any{"UserID": "u1"}}) + default: + t.Fatalf("unexpected path %s", r.URL.Path) + } + })) + session, err := client.StartLogin(context.Background(), "acc1") + if err != nil { + t.Fatal(err) + } + client.mu.Lock() + pending := client.pending["acc1"] + pending.done = true + pending.authCode = "PARKED-CODE" + pending.credential = Credential{MachineID: pending.machineID, DeviceID: pending.deviceID} + client.mu.Unlock() + + callback := "http://127.0.0.1:49519/authorize?loginTraceID=" + session.State + if err := client.CompleteLogin(context.Background(), "acc1", callback); err != nil { + t.Fatal(err) + } + if _, _, err := store.LoadCredentialPayload(context.Background(), "acc1"); err != nil { + t.Fatalf("credential must be stored: %v", err) + } +} + +// A callback with no code and no token must fail loudly rather than persist an +// empty credential the probe would then report as ready. +func TestCompleteLoginRejectsCallbackWithoutCodeOrToken(t *testing.T) { + client, store := newTestClient(t, http.HandlerFunc(func(http.ResponseWriter, *http.Request) { + t.Fatal("no upstream call without a code or token") + })) + if _, err := client.StartLogin(context.Background(), "acc1"); err != nil { + t.Fatal(err) + } + err := client.CompleteLogin(context.Background(), "acc1", "http://127.0.0.1:49519/authorize?isRedirect=true") + if err == nil { + t.Fatal("expected an error for a callback with neither code nor token") + } + if _, _, loadErr := store.LoadCredentialPayload(context.Background(), "acc1"); loadErr == nil { + t.Fatal("no credential may be stored for an empty callback") + } +} + +// A redeemed login must not leave a trace index behind for a later callback to +// match against a credential that was already exchanged. +func TestPollLoginClearsTraceIndexAfterRedeem(t *testing.T) { + client, _ := newTestClient(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case pathExchangeCode: + _ = json.NewEncoder(w).Encode(map[string]any{"Result": map[string]any{ + "Token": "at", "RefreshToken": "rt2", "TokenExpireAt": time.Now().Add(time.Hour).Unix(), + }}) + case pathUserInfo: + _ = json.NewEncoder(w).Encode(map[string]any{"Result": map[string]any{"UserID": "u1"}}) + default: + t.Fatalf("unexpected path %s", r.URL.Path) + } + })) + session, err := client.StartLogin(context.Background(), "acc1") + if err != nil { + t.Fatal(err) + } + client.mu.Lock() + pending := client.pending["acc1"] + pending.done = true + pending.authCode = "ABC-123" + pending.credential = Credential{MachineID: pending.machineID, DeviceID: pending.deviceID} + client.mu.Unlock() + + if done, _, err := client.PollLogin(context.Background(), "acc1"); err != nil || !done { + t.Fatalf("done=%v err=%v", done, err) + } + client.mu.Lock() + _, stillPending := client.pending["acc1"] + _, stillTraced := client.pendingTraces[session.State] + client.mu.Unlock() + if stillPending || stillTraced { + t.Fatalf("redeemed login must clear pending (pending=%v trace=%v)", stillPending, stillTraced) + } +} + +// After a pasted callback completes the login, a poll still in flight from the +// browser flow must report success rather than "login not started". +func TestPollLoginReportsReadyAfterPastedCallback(t *testing.T) { + client, store := newTestClient(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case pathExchangeCode: + _ = json.NewEncoder(w).Encode(map[string]any{"Result": map[string]any{ + "Token": "at", "RefreshToken": "rt2", "TokenExpireAt": time.Now().Add(time.Hour).Unix(), + }}) + case pathUserInfo: + _ = json.NewEncoder(w).Encode(map[string]any{"Result": map[string]any{"UserID": "u1"}}) + default: + t.Fatalf("unexpected path %s", r.URL.Path) + } + })) + session, err := client.StartLogin(context.Background(), "acc1") + if err != nil { + t.Fatal(err) + } + client.mu.Lock() + pending := client.pending["acc1"] + pending.done = true + pending.authCode = "ABC-123" + pending.credential = Credential{MachineID: pending.machineID, DeviceID: pending.deviceID} + client.mu.Unlock() + + callback := "http://127.0.0.1:49519/authorize?loginTraceID=" + session.State + if err := client.CompleteLogin(context.Background(), "acc1", callback); err != nil { + t.Fatal(err) + } + if _, _, err := store.LoadCredentialPayload(context.Background(), "acc1"); err != nil { + t.Fatalf("credential must be stored: %v", err) + } + done, message, err := client.PollLogin(context.Background(), "acc1") + if err != nil { + t.Fatalf("poll after paste must not fail: %v", err) + } + if !done { + t.Fatalf("poll after paste must report done, message=%q", message) + } +} + +// The loopback callback must redeem the code on its own. The console's poll +// window is much shorter than a real browser login, so waiting for a poller +// left the code unexchanged and the account at "account not found". +func TestAcceptCallbackRedeemsWithoutPolling(t *testing.T) { + var exchanges int32 + client, store := newTestClient(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case pathExchangeCode: + atomic.AddInt32(&exchanges, 1) + _ = json.NewEncoder(w).Encode(map[string]any{"Result": map[string]any{ + "Token": "at", "RefreshToken": "rt2", "TokenExpireAt": time.Now().Add(time.Hour).Unix(), + }}) + case pathUserInfo: + _ = json.NewEncoder(w).Encode(map[string]any{"Result": map[string]any{ + "UserID": "u1", "ScreenName": "Tester", "EnterpriseID": "e1", + }}) + default: + t.Fatalf("unexpected path %s", r.URL.Path) + } + })) + session, err := client.StartLogin(context.Background(), "acc1") + if err != nil { + t.Fatal(err) + } + // The browser lands on the loopback listener; no poll ever runs. + callback := "http://127.0.0.1:9/authorize?loginTraceID=" + session.State + + "&authCodeInfo=%7B%22AuthCode%22%3A%22ABC-123%22%7D" + if err := client.acceptCallback(context.Background(), callback); err != nil { + t.Fatal(err) + } + // The exchange runs in the background, so wait for the credential. + var credential Credential + deadline := time.Now().Add(5 * time.Second) + for time.Now().Before(deadline) { + if _, payload, err := store.LoadCredentialPayload(context.Background(), "acc1"); err == nil { + if decoded, err := DecodeCredential(payload); err == nil && decoded.Ready() { + credential = decoded + break + } + } + time.Sleep(10 * time.Millisecond) + } + if !credential.Ready() { + t.Fatal("callback must redeem the code and store a credential without any poll") + } + if credential.UID != "u1" || credential.AccessToken != "at" || credential.RefreshToken != "rt2" { + t.Fatalf("credential=%+v", credential) + } + if credential.DevicePrivateKey == "" { + t.Fatal("device key must be stored for later code exchanges") + } + if got := atomic.LoadInt32(&exchanges); got != 1 { + t.Fatalf("code must be exchanged exactly once, got %d", got) + } + // A poll racing the background exchange must report the same success, not + // "login not started". + done, _, err := client.PollLogin(context.Background(), "acc1") + if err != nil || !done { + t.Fatalf("poll after callback done=%v err=%v", done, err) + } +} + +// The authorization code is single-use, so overlapping redeemers must produce +// exactly one exchange and share its result. +func TestRedeemExchangesCodeOnceUnderConcurrency(t *testing.T) { + var exchanges int32 + release := make(chan struct{}) + client, _ := newTestClient(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case pathExchangeCode: + atomic.AddInt32(&exchanges, 1) + <-release // hold the exchange open so both callers overlap + _ = json.NewEncoder(w).Encode(map[string]any{"Result": map[string]any{ + "Token": "at", "RefreshToken": "rt2", "TokenExpireAt": time.Now().Add(time.Hour).Unix(), + }}) + case pathUserInfo: + _ = json.NewEncoder(w).Encode(map[string]any{"Result": map[string]any{"UserID": "u1"}}) + default: + t.Fatalf("unexpected path %s", r.URL.Path) + } + })) + if _, err := client.StartLogin(context.Background(), "acc1"); err != nil { + t.Fatal(err) + } + client.mu.Lock() + pending := client.pending["acc1"] + pending.done = true + pending.authCode = "ABC-123" + pending.credential = Credential{MachineID: pending.machineID, DeviceID: pending.deviceID} + credential, authCode, verifier := pending.credential, pending.authCode, pending.codeVerifier + client.mu.Unlock() + + errs := make(chan error, 2) + go func() { + errs <- client.redeem(context.Background(), "acc1", pending, credential, authCode, verifier) + }() + // Wait for the first exchange to be in flight before starting the second. + deadline := time.Now().Add(5 * time.Second) + for atomic.LoadInt32(&exchanges) == 0 && time.Now().Before(deadline) { + time.Sleep(5 * time.Millisecond) + } + go func() { + errs <- client.redeem(context.Background(), "acc1", pending, credential, authCode, verifier) + }() + close(release) + for i := 0; i < 2; i++ { + if err := <-errs; err != nil { + t.Fatalf("concurrent redeem failed: %v", err) + } + } + if got := atomic.LoadInt32(&exchanges); got != 1 { + t.Fatalf("single-use code must be exchanged once, got %d exchanges", got) + } +} + +// A failed exchange must not burn the guard, so the operator can retry the +// same callback instead of restarting the whole login. +func TestRedeemAllowsRetryAfterFailure(t *testing.T) { + var exchanges int32 + client, store := newTestClient(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case pathExchangeCode: + if atomic.AddInt32(&exchanges, 1) == 1 { + w.WriteHeader(http.StatusBadGateway) + _, _ = w.Write([]byte(`{"error":"upstream blip"}`)) + return + } + _ = json.NewEncoder(w).Encode(map[string]any{"Result": map[string]any{ + "Token": "at", "RefreshToken": "rt2", "TokenExpireAt": time.Now().Add(time.Hour).Unix(), + }}) + case pathUserInfo: + _ = json.NewEncoder(w).Encode(map[string]any{"Result": map[string]any{"UserID": "u1"}}) + default: + t.Fatalf("unexpected path %s", r.URL.Path) + } + })) + if _, err := client.StartLogin(context.Background(), "acc1"); err != nil { + t.Fatal(err) + } + client.mu.Lock() + pending := client.pending["acc1"] + pending.done = true + pending.authCode = "ABC-123" + pending.credential = Credential{MachineID: pending.machineID, DeviceID: pending.deviceID} + credential, authCode, verifier := pending.credential, pending.authCode, pending.codeVerifier + client.mu.Unlock() + + if err := client.redeem(context.Background(), "acc1", pending, credential, authCode, verifier); err == nil { + t.Fatal("first exchange must surface the upstream failure") + } + if err := client.redeem(context.Background(), "acc1", pending, credential, authCode, verifier); err != nil { + t.Fatalf("retry after a transient failure must succeed: %v", err) + } + if _, _, err := store.LoadCredentialPayload(context.Background(), "acc1"); err != nil { + t.Fatalf("credential must be stored after the retry: %v", err) + } +} + +// The loopback listener is shared across accounts, so a callback must be paired +// with the login whose challenge it echoes, not with whichever entry the map +// happens to yield first. +func TestAcceptCallbackPairsByTraceAcrossAccounts(t *testing.T) { + client, store := newTestClient(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case pathExchangeCode: + _ = json.NewEncoder(w).Encode(map[string]any{"Result": map[string]any{ + "Token": "at", "RefreshToken": "rt2", "TokenExpireAt": time.Now().Add(time.Hour).Unix(), + }}) + case pathUserInfo: + _ = json.NewEncoder(w).Encode(map[string]any{"Result": map[string]any{"UserID": "u1"}}) + default: + t.Fatalf("unexpected path %s", r.URL.Path) + } + })) + first, err := client.StartLogin(context.Background(), "acc-first") + if err != nil { + t.Fatal(err) + } + if _, err := client.StartLogin(context.Background(), "acc-second"); err != nil { + t.Fatal(err) + } + client.mu.Lock() + wantMachine := client.pending["acc-first"].machineID + client.mu.Unlock() + + // The callback echoes the first login's trace id. + callback := "http://127.0.0.1:9/authorize?loginTraceID=" + first.State + + "&authCodeInfo=%7B%22AuthCode%22%3A%22ABC-123%22%7D" + if err := client.acceptCallback(context.Background(), callback); err != nil { + t.Fatal(err) + } + var credential Credential + deadline := time.Now().Add(5 * time.Second) + for time.Now().Before(deadline) { + if _, payload, err := store.LoadCredentialPayload(context.Background(), "acc-first"); err == nil { + if decoded, err := DecodeCredential(payload); err == nil && decoded.Ready() { + credential = decoded + break + } + } + time.Sleep(10 * time.Millisecond) + } + if !credential.Ready() { + t.Fatal("the traced account must receive the credential") + } + if credential.MachineID != wantMachine { + t.Fatalf("credential machine=%s want %s", credential.MachineID, wantMachine) + } + if _, _, err := store.LoadCredentialPayload(context.Background(), "acc-second"); err == nil { + t.Fatal("the other account's login must stay untouched") + } +} + +// The CN front doors disagree about which host mints a code: the official +// client exchanges at the host product.json configures (api.trae.cn) while the +// callback advertises api.trae.com.cn. A code minted by one must be redeemable +// even when the other host rejects it as an invalid parameter. +func TestExchangeAuthCodeFallsBackAcrossCNHosts(t *testing.T) { + var hosts []string + client, _ := newTestClient(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + hosts = append(hosts, r.Host) + // The first host (api.trae.cn) rejects the code; the callback's host accepts it. + if len(hosts) == 1 { + w.WriteHeader(http.StatusBadRequest) + _, _ = w.Write([]byte(`{"ResponseMetadata":{"Error":{"Code":"10101","Message":"无效参数"}}}`)) + return + } + _ = json.NewEncoder(w).Encode(map[string]any{"Result": map[string]any{ + "Token": "at", "RefreshToken": "rt2", "TokenExpireAt": time.Now().Add(time.Hour).Unix(), + }}) + })) + credential := Credential{ + MachineID: "m1", DeviceID: "1234567890123456", + DevicePrivateKey: mustKey(t), APIHost: "https://api.trae.com.cn", Domain: DomainCN, + } + refreshed, err := client.ExchangeAuthCode(context.Background(), "acc1", credential, "CODE-XYZ", "verifier") + if err != nil { + t.Fatalf("a code rejected by the configured host must still be tried on the callback host: %v", err) + } + if refreshed.AccessToken != "at" { + t.Fatalf("credential=%+v", refreshed) + } + if len(hosts) != 2 { + t.Fatalf("expected two host attempts, got %v", hosts) + } +} + +// A transient upstream failure must surface, not be masked by a second host +// attempt that happens to succeed. +func TestExchangeAuthCodeSurfacesTransientFailure(t *testing.T) { + var attempts int32 + client, _ := newTestClient(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + atomic.AddInt32(&attempts, 1) + w.WriteHeader(http.StatusBadGateway) + _, _ = w.Write([]byte(`{"error":"upstream blip"}`)) + })) + credential := Credential{ + MachineID: "m1", DeviceID: "1234567890123456", + DevicePrivateKey: mustKey(t), APIHost: "https://api.trae.com.cn", Domain: DomainCN, + } + if _, err := client.ExchangeAuthCode(context.Background(), "acc1", credential, "CODE-XYZ", "verifier"); err == nil { + t.Fatal("a 502 must surface as an error") + } + if got := atomic.LoadInt32(&attempts); got != 1 { + t.Fatalf("a transient failure must not be retried on another host, got %d attempts", got) + } +} + func TestChatNonStreamAggregatesToolsAndReasoning(t *testing.T) { payload, _ := Credential{AccessToken: "at", RefreshToken: "rt", UID: "u1", Domain: DomainCN, ExpiresAt: 4102444800, MachineID: "m1", DeviceID: "d1"}.Encode() store := &memStore{items: map[string][]byte{"acc1": payload}} diff --git a/internal/providers/trae/credential.go b/internal/providers/trae/credential.go index 43a6eca..d48e625 100644 --- a/internal/providers/trae/credential.go +++ b/internal/providers/trae/credential.go @@ -65,6 +65,9 @@ const ( const ( refreshLead = 24 * time.Hour loginPendingTTL = 10 * time.Minute + // redeemTimeout bounds the background code exchange started by the loopback + // callback, which outlives the browser request that delivered it. + redeemTimeout = 90 * time.Second ) // Credential is the canonical storage payload shape.