diff --git a/internal/gui/assets/app.js b/internal/gui/assets/app.js index 30b39564..423294b4 100644 --- a/internal/gui/assets/app.js +++ b/internal/gui/assets/app.js @@ -1392,6 +1392,7 @@ function accountPlan(a) { if (a.agent === "grok") return a.plan || "SuperGrok"; if (a.agent === "gemini" || a.agent === "antigravity") return a.plan || "Google"; if (a.agent === "zcode") return a.plan || "GLM Coding Plan"; + if (a.agent === "workbuddy") return a.plan || "WorkBuddy"; return t("signed in"); } @@ -3322,6 +3323,8 @@ const SUBS = [ { agent: "copilot", name: "Copilot", icon: "githubcopilot", plans: "Pro · Pro+ · Business", own: true }, // Z.ai's GLM Coding Plan, signed in as ZCode does; ZCode's own account is read too { agent: "zcode", name: "ZCode (GLM Coding Plan)", icon: "zcode", plans: "Lite · Pro · Max", own: true }, + // Tencent's CodeBuddy plan, signed in as WorkBuddy does; WorkBuddy's own account is read too + { agent: "workbuddy", name: "WorkBuddy (CodeBuddy)", icon: "workbuddy-color", plans: "Free · Pro", own: true }, // devin's credentials.toml keeps one account too { agent: "devin", name: "Devin", icon: "devin", plans: "Pro · Enterprise", single: true }, // Google's sign-ins; Gemini CLI's own account is read too diff --git a/internal/gui/assets/icons/workbuddy-color.svg b/internal/gui/assets/icons/workbuddy-color.svg new file mode 100644 index 00000000..faf7e5df --- /dev/null +++ b/internal/gui/assets/icons/workbuddy-color.svg @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/internal/provider/account.go b/internal/provider/account.go index 8d39046b..f0167db6 100644 --- a/internal/provider/account.go +++ b/internal/provider/account.go @@ -703,6 +703,9 @@ func Accounts() []Provider { if p, ok := zcodeAccount(); ok { out = append(out, p) } + if p, ok := workBuddyAccount(); ok { + out = append(out, p) + } for _, agent := range []string{"gemini", "antigravity"} { if p, ok := googleAccountOf(agent); ok { out = append(out, p) diff --git a/internal/provider/login_usage.go b/internal/provider/login_usage.go index abdd3620..49c694f1 100644 --- a/internal/provider/login_usage.go +++ b/internal/provider/login_usage.go @@ -38,6 +38,8 @@ func LoginUsage(ctx context.Context, agent string) map[string]SubscriptionQuota logins = copilotLoginList() case "zcode": logins = zcodeLoginList() + case "workbuddy": + logins = wbLoginList() case "gemini", "antigravity": logins = googleLoginList(agent) case "cursor": // one account, the one cursor-agent is signed in to @@ -93,6 +95,9 @@ func loginQuota(ctx context.Context, l Login) SubscriptionQuota { if l.Agent == "zcode" { return zcodeLoginQuota(ctx, l) } + if l.Agent == "workbuddy" { + return wbLoginQuota(ctx, l) + } if l.Agent == "copilot" { for _, c := range copilotLogins(copilotConfigDir()) { if strings.EqualFold(c.User, l.User) { diff --git a/internal/provider/logins.go b/internal/provider/logins.go index 69643fa3..aab7fcde 100644 --- a/internal/provider/logins.go +++ b/internal/provider/logins.go @@ -423,11 +423,14 @@ func Logins(agent string) []Login { return copilotLoginList() case "zcode": return zcodeLoginList() + case "workbuddy": + return wbLoginList() case "gemini", "antigravity": return googleLoginList(agent) case "": side = append(grokLoginList(), copilotLoginList()...) side = append(side, zcodeLoginList()...) + side = append(side, wbLoginList()...) side = append(side, googleLoginList("gemini")...) side = append(side, googleLoginList("antigravity")...) } @@ -466,6 +469,8 @@ func SwitchLogin(agent, user string) error { return switchCopilotLogin(user) case "zcode": return switchZCodeLogin(user) + case "workbuddy": + return switchWorkBuddyLogin(user) case "gemini", "antigravity": return switchGoogleLogin(agent, user) } @@ -575,6 +580,8 @@ func ForgetLogin(agent, user string) error { return forgetCopilotLogin(user) case "zcode": return forgetZCodeLogin(user) + case "workbuddy": + return forgetWorkBuddyLogin(user) case "gemini", "antigravity": return forgetGoogleLogin(agent, user) } diff --git a/internal/provider/logins_on.go b/internal/provider/logins_on.go index 3cbd795e..405062a0 100644 --- a/internal/provider/logins_on.go +++ b/internal/provider/logins_on.go @@ -31,6 +31,8 @@ func SetLoginOn(agent, user string, on bool) error { return setCopilotLoginOn(user, on) case "zcode": return setZCodeLoginOn(user, on) + case "workbuddy": + return setWorkBuddyLoginOn(user, on) case "gemini", "antigravity": return setGoogleLoginOn(agent, user, on) } @@ -64,6 +66,9 @@ func (p Provider) AlsoOn() []Provider { if p.Account != nil && p.Account.Agent == "zcode" { return zcodeAlsoOn() } + if p.Account != nil && p.Account.Agent == "workbuddy" { + return workBuddyAlsoOn() + } if p.Account != nil && (p.Account.Agent == "gemini" || p.Account.Agent == "antigravity") { return googleAlsoOn(p.Account.Agent) } diff --git a/internal/provider/provider.go b/internal/provider/provider.go index 88ae1fe1..2698a59d 100644 --- a/internal/provider/provider.go +++ b/internal/provider/provider.go @@ -359,7 +359,7 @@ func freeName(name string) string { } // accountIDs are the ids of the subscriptions magpie can list (account.go). -var accountIDs = []string{"antigravity", "claude", "codex", "copilot", "cursor", "devin", "gemini", "grok", "kiro", "zcode"} +var accountIDs = []string{"antigravity", "claude", "codex", "copilot", "cursor", "devin", "gemini", "grok", "kiro", "workbuddy", "zcode"} func stored(id string) bool { for _, p := range load().Providers { diff --git a/internal/provider/side_logins.go b/internal/provider/side_logins.go index aeb86e46..7fb181c1 100644 --- a/internal/provider/side_logins.go +++ b/internal/provider/side_logins.go @@ -184,7 +184,7 @@ func addSideLogin(l savedLogin, dup func(savedLogin)) error { // agent's own store as Claude Code's and Codex's are. func sideAgent(agent string) bool { switch agent { - case "grok", "copilot", "zcode", "gemini", "antigravity": + case "grok", "copilot", "zcode", "workbuddy", "gemini", "antigravity": return true } return false diff --git a/internal/provider/signin.go b/internal/provider/signin.go index 740925ed..d3a87457 100644 --- a/internal/provider/signin.go +++ b/internal/provider/signin.go @@ -233,6 +233,11 @@ func (s *signInFlow) begin() error { if err := startCopilotSignIn(s); err != nil { return err } + case "workbuddy": + // Tencent's sign-in, as WorkBuddy (CodeBuddy) makes it + if err := startWorkBuddySignIn(s); err != nil { + return err + } case "zcode": // Z.ai's sign-in, as ZCode makes it if err := startZCodeSignIn(s); err != nil { diff --git a/internal/provider/subscription_usage.go b/internal/provider/subscription_usage.go index 6880b1d7..876c5433 100644 --- a/internal/provider/subscription_usage.go +++ b/internal/provider/subscription_usage.go @@ -208,6 +208,9 @@ func fetchSubscriptionUsage() []SubscriptionQuota { if !hidden["zcode"] { fetches = append(fetches, perLogin(ctx, zcodeLoginList(), "ZCode", "zcode")...) } + if !hidden["workbuddy"] { + fetches = append(fetches, perLogin(ctx, wbLoginList(), "WorkBuddy", "workbuddy-color")...) + } for _, agent := range []string{"gemini", "antigravity"} { if hidden[agent] { continue diff --git a/internal/provider/workbuddy.go b/internal/provider/workbuddy.go new file mode 100644 index 00000000..34fcceed --- /dev/null +++ b/internal/provider/workbuddy.go @@ -0,0 +1,669 @@ +package provider + +// A WorkBuddy subscription is Tencent's CodeBuddy plan, which WorkBuddy +// (its desktop app, packaged from CodeBuddy Code) signs in to. The plan is +// served on an OpenAI-compatible endpoint under the account's own access +// token — a JWT the account refreshes — with the account's id and domain in +// headers, exactly as WorkBuddy sends them. +// +// WorkBuddy's own account is read, never changed, from its auth store, +// /auth/workbuddy-desktop.info (plain JSON while credential +// protection is off, which it is by default); its access and refresh tokens +// and the account are taken as they sit there. Further accounts are signed +// in by magpie with WorkBuddy's own external-link flow (auth/state, then +// auth/token polled until ready, then login/account), and their tokens are +// kept in logins.json. A near-expired token is refreshed as WorkBuddy does, +// and a magpie-signed-in one's refresh is written back beside it. +// +// The plan has no small model list to ask; the models are WorkBuddy's, its +// CLI agent's coding picks. + +import ( + "bytes" + "context" + "crypto/rand" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "net/url" + "os" + "path/filepath" + "runtime" + "strings" + "sync" + "time" + + "github.com/yetone/magpie/internal/catalog" +) + +// Where WorkBuddy's plan and its sign-in are; vars so tests can point them +// elsewhere. +var ( + wbEndpoint = "https://copilot.tencent.com" + // wbAppVersion is the plugin version the sign-in page is told. + wbAppVersion = "2.0.0" + // wbPollInterval is how often the sign-in asks whether it is ready. + wbPollInterval = time.Second +) + +// wbModels are the plan's models, WorkBuddy's CLI agent's coding picks, with +// the names and context windows its config gives them. +var wbModels = []catalog.Model{ + {ID: "auto", Name: "Auto", Context: 168_000}, + {ID: "hy4-preview-f", Name: "Hy4 preview", Context: 1_000_000}, + {ID: "hy3", Name: "Hy3", Context: 192_000}, + {ID: "hy3-x", Name: "Hy3-X", Context: 192_000}, + {ID: "deepseek-v4.1-flash", Name: "Deepseek-V4.1-Flash", Context: 1_000_000}, + {ID: "glm-5.3", Name: "GLM-5.3", Context: 1_000_000}, + {ID: "glm-5.3-flash", Name: "GLM-5.3-Flash", Context: 1_000_000}, + {ID: "glm-5.2", Name: "GLM-5.2", Context: 1_000_000}, + {ID: "glm-5.1", Name: "GLM-5.1", Context: 200_000}, + {ID: "glm-5v-turbo", Name: "GLM-5v-Turbo", Context: 200_000}, + {ID: "minimax-m3", Name: "MiniMax-M3", Context: 512_000}, + {ID: "kimi-k3-1", Name: "Kimi-K3", Context: 1_000_000}, + {ID: "kimi-k2.7", Name: "Kimi-K2.7-Code", Context: 256_000}, + {ID: "kimi-k2.6", Name: "Kimi-K2.6", Context: 256_000}, + {ID: "deepseek-v4-pro", Name: "Deepseek-V4-Pro", Context: 1_000_000}, +} + +// wbCreds is a WorkBuddy account's tokens and where they are served, as the +// auth store and the sign-in name them. +type wbCreds struct { + UID string `json:"uid"` + Access string `json:"accessToken"` + Refresh string `json:"refreshToken"` + ExpiresAt int64 `json:"expiresAt"` // unix ms, when the access token lapses + RefreshExpiresAt int64 `json:"refreshExpiresAt"` // unix ms, when the refresh token lapses + Domain string `json:"domain"` + TokenType string `json:"tokenType,omitempty"` +} + +// wbAccount is a signed-in WorkBuddy account: which login it is (who, its +// plan, and whether it is the one in use), its tokens, and whether it is +// WorkBuddy's own sign-in (read-only) or one magpie added (in logins.json). +type wbAccount struct { + Login + creds wbCreds + own bool +} + +// wbAPI is the API root for a domain: auth and billing sit under it. The +// account's domain (www.codebuddy.cn) names the site; requests otherwise go +// to the product endpoint. +func wbAPI() string { return strings.TrimRight(wbEndpoint, "/") } + +// ---- WorkBuddy's own account -------------------------------------------------- + +// wbAuthPath is where WorkBuddy keeps the signed-in session, per platform, +// as its file-authentication-storage does: /auth/.info. +func wbAuthPath() string { + home, _ := os.UserHomeDir() + var base string + switch runtime.GOOS { + case "darwin": + base = filepath.Join(home, "Library", "Application Support", "CodeBuddyExtension") + case "windows": + base = filepath.Join(home, "AppData", "Local", "CodeBuddyExtension") + default: + base = filepath.Join(home, ".local", "share", "CodeBuddyExtension") + } + return filepath.Join(base, "Data", "Public", "auth", "workbuddy-desktop.info") +} + +// wbStoredSession is the shape of the auth store's session: the account and +// its tokens. The tokens are plain strings while credential protection is +// off (the default); when it is on they are objects magpie can't read, and +// the account is then taken as not readable. +type wbStoredSession struct { + Account struct { + UID string `json:"uid"` + Nickname json.RawMessage `json:"nickname"` + PhoneNumber json.RawMessage `json:"phoneNumber"` + } `json:"account"` + Auth struct { + AccessToken json.RawMessage `json:"accessToken"` + RefreshToken json.RawMessage `json:"refreshToken"` + ExpiresAt int64 `json:"expiresAt"` + ExpiresIn int64 `json:"expiresIn"` + RefreshExpiresAt int64 `json:"refreshExpiresAt"` + RefreshExpiresIn int64 `json:"refreshExpiresIn"` + Domain string `json:"domain"` + TokenType string `json:"tokenType"` + } `json:"auth"` +} + +// wbString reads a JSON value that is a plain string, "" for anything else +// (an encrypted-field object magpie doesn't decrypt). +func wbString(raw json.RawMessage) string { + var s string + if json.Unmarshal(raw, &s) == nil { + return s + } + return "" +} + +// wbOwn is the account WorkBuddy is signed in to and its tokens; ok is +// false when it has none, or they are encrypted at rest. +func wbOwn() (who string, c wbCreds, ok bool) { + var s wbStoredSession + if !readJSON(wbAuthPath(), &s) { + return "", wbCreds{}, false + } + access, refresh := wbString(s.Auth.AccessToken), wbString(s.Auth.RefreshToken) + if s.Account.UID == "" || access == "" { + return "", wbCreds{}, false + } + c = wbCreds{ + UID: s.Account.UID, + Access: access, + Refresh: refresh, + ExpiresAt: s.Auth.ExpiresAt, + RefreshExpiresAt: s.Auth.RefreshExpiresAt, + Domain: s.Auth.Domain, + TokenType: s.Auth.TokenType, + } + now := time.Now().UnixMilli() + if c.ExpiresAt == 0 && s.Auth.ExpiresIn > 0 { + c.ExpiresAt = now + s.Auth.ExpiresIn*1000 + } + if c.RefreshExpiresAt == 0 && s.Auth.RefreshExpiresIn > 0 { + c.RefreshExpiresAt = now + s.Auth.RefreshExpiresIn*1000 + } + return wbWho(wbString(s.Account.Nickname), wbString(s.Account.PhoneNumber), s.Account.UID), c, true +} + +// wbWho names a WorkBuddy account: its nickname, its phone number, or its id. +func wbWho(nickname, phone, id string) string { + return firstNonEmpty(strings.TrimSpace(nickname), strings.TrimSpace(phone), id, "WorkBuddy") +} + +// ---- the accounts ------------------------------------------------------------- + +func wbSavedCreds(l savedLogin) (wbCreds, bool) { + var c wbCreds + if json.Unmarshal(l.Auth, &c) != nil || c.Access == "" || c.UID == "" { + return wbCreds{}, false + } + return c, true +} + +// wbLogins is every WorkBuddy account signed in, the first in use first. +func wbLogins() []wbAccount { + ownUser, own, hasOwn := wbOwn() + if !hasOwn { + ownUser = "" + } + var out []wbAccount + for _, l := range sideLogins("workbuddy", ownUser, func(l savedLogin) bool { + _, ok := wbSavedCreds(l) + return ok + }) { + a := wbAccount{Login: l.Login, own: l.saved.own()} + if a.own { + a.creds = own + } else { + a.creds, _ = wbSavedCreds(l.saved) + } + out = append(out, a) + } + return out +} + +func wbSide() []sideLogin { + var out []sideLogin + for _, a := range wbLogins() { + out = append(out, sideLogin{Login: a.Login}) + } + return out +} + +func wbLoginList() []Login { return loginsOf(wbSide()) } + +func switchWorkBuddyLogin(user string) error { + return switchSideLogin("workbuddy", user, wbSide()) +} + +func setWorkBuddyLoginOn(user string, on bool) error { + return setSideLoginOn("workbuddy", user, on, wbSide()) +} + +func forgetWorkBuddyLogin(user string) error { + return forgetSideLogin("workbuddy", user, "WorkBuddy's own sign-in; sign out in WorkBuddy", wbSide(), nil) +} + +func workBuddyAccount() (Provider, bool) { + ls := wbLogins() + if len(ls) == 0 { + return Provider{}, false + } + return wbProvider(ls[0]), true +} + +// workBuddyAlsoOn is the WorkBuddy accounts in use behind the first. +func workBuddyAlsoOn() []Provider { + var out []Provider + for _, a := range wbLogins() { + if !a.Active && a.On { + out = append(out, wbProvider(a)) + } + } + return out +} + +func wbProvider(a wbAccount) Provider { + acct := &Account{Agent: "workbuddy", User: a.User, Plan: a.Plan} + acct.sign = func(ctx context.Context, req *http.Request, body []byte) error { + c, err := wbFresh(ctx, a) + if err != nil { + return err + } + req.Header.Del("Authorization") + req.Header.Set("Authorization", "Bearer "+c.Access) + req.Header.Set("X-User-Id", c.UID) + req.Header.Set("X-Domain", wbDomain(c)) + req.Header.Set("X-Product", "SaaS") + req.Header.Set("X-IDE-Type", "WorkBuddy") + return nil + } + acct.models = func() []catalog.Model { return wbModels } + return Provider{ID: "workbuddy", Name: "WorkBuddy", Icon: "workbuddy-color", Chat: wbAPI() + "/v2", Website: "https://www.codebuddy.cn", Account: acct} +} + +// wbDomain is the X-Domain a request carries: the account's own domain, or +// the endpoint's authority. +func wbDomain(c wbCreds) string { + if c.Domain != "" { + return c.Domain + } + if u, err := url.Parse(wbAPI()); err == nil && u.Host != "" { + return u.Host + } + return "" +} + +// ---- tokens ------------------------------------------------------------------- + +// wbTokens caches each account's freshest tokens (by uid), so a token +// refreshed for one request is used by the next; WorkBuddy's own file is +// never written, and a magpie-added account's refresh goes to logins.json. +var wbTokens = struct { + sync.Mutex + m map[string]wbCreds +}{m: map[string]wbCreds{}} + +// wbFresh is a's tokens, refreshed when the access token is about to lapse. +func wbFresh(ctx context.Context, a wbAccount) (wbCreds, error) { + wbTokens.Lock() + c := a.creds + if cached, ok := wbTokens.m[a.creds.UID]; ok && cached.ExpiresAt >= c.ExpiresAt { + c = cached + } + wbTokens.Unlock() + + if c.Access != "" && !wbNearExpiry(c.ExpiresAt) { + return c, nil + } + if c.Refresh == "" || (c.RefreshExpiresAt > 0 && time.Now().UnixMilli() >= c.RefreshExpiresAt) { + if c.Access != "" { + return c, nil // no way to refresh; let the request try what there is + } + return wbCreds{}, errors.New("this WorkBuddy account is signed out; sign in again") + } + refreshed, err := wbRefresh(ctx, c) + if err != nil { + if c.Access != "" { + return c, nil // a refresh hiccup: the current token may still work + } + return wbCreds{}, err + } + wbTokens.Lock() + wbTokens.m[refreshed.UID] = refreshed + wbTokens.Unlock() + if !a.own { + wbSaveCreds(a.User, refreshed) + } + return refreshed, nil +} + +// wbNearExpiry is true within a minute of a token's end (or when unknown). +func wbNearExpiry(expiresAt int64) bool { + if expiresAt == 0 { + return false // the store didn't say; trust it until refused + } + return time.Now().UnixMilli() >= expiresAt-60_000 +} + +// wbRefresh trades a refresh token for a fresh access token, as WorkBuddy's +// auth provider does: POST /v2/plugin/auth/token/refresh with the refresh +// token in a header. +func wbRefresh(ctx context.Context, c wbCreds) (wbCreds, error) { + var got wbRefreshed + err := wbCall(ctx, http.MethodPost, wbAPI()+"/v2/plugin/auth/token/refresh", map[string]string{ + "X-Refresh-Token": c.Refresh, + "X-Auth-Refresh-Source": "plugin", + "X-Domain": wbDomain(c), + }, map[string]any{}, &got) + if err != nil { + return wbCreds{}, fmt.Errorf("WorkBuddy token refresh: %w", err) + } + if got.AccessToken == "" { + return wbCreds{}, errors.New("WorkBuddy gave no refreshed token") + } + return wbMergeRefreshed(c, got), nil +} + +type wbRefreshed struct { + AccessToken string `json:"accessToken"` + RefreshToken string `json:"refreshToken"` + ExpiresAt int64 `json:"expiresAt"` + ExpiresIn int64 `json:"expiresIn"` + RefreshExpiresAt int64 `json:"refreshExpiresAt"` + RefreshExpiresIn int64 `json:"refreshExpiresIn"` + Domain string `json:"domain"` + TokenType string `json:"tokenType"` +} + +func wbMergeRefreshed(c wbCreds, got wbRefreshed) wbCreds { + now := time.Now().UnixMilli() + out := c + out.Access = got.AccessToken + if got.RefreshToken != "" { + out.Refresh = got.RefreshToken + } + if got.Domain != "" { + out.Domain = got.Domain + } + if got.TokenType != "" { + out.TokenType = got.TokenType + } + switch { + case got.ExpiresAt > 0: + out.ExpiresAt = got.ExpiresAt + case got.ExpiresIn > 0: + out.ExpiresAt = now + got.ExpiresIn*1000 + } + switch { + case got.RefreshExpiresAt > 0: + out.RefreshExpiresAt = got.RefreshExpiresAt + case got.RefreshExpiresIn > 0: + out.RefreshExpiresAt = now + got.RefreshExpiresIn*1000 + } + return out +} + +// wbSaveCreds writes a magpie-added account's refreshed tokens back to +// logins.json, so the next run starts from them. +func wbSaveCreds(user string, c wbCreds) { + _ = editSideLogin("workbuddy", user, func(ls []savedLogin, i int) ([]savedLogin, error) { + if ls[i].own() { + return ls, nil // never write WorkBuddy's own file + } + auth, err := json.Marshal(c) + if err != nil { + return ls, err + } + ls[i].Auth = auth + ls[i].Renewed = time.Now().UTC().Truncate(time.Second) + return ls, nil + }) +} + +// ---- allowance ---------------------------------------------------------------- + +// wbQuota is a WorkBuddy account's credit allowance, from its resource +// summary: the plan's credits used against what the cycle grants. +func wbQuota(ctx context.Context, a wbAccount) SubscriptionQuota { + q := SubscriptionQuota{Provider: "workbuddy", Name: "WorkBuddy", Icon: "workbuddy-color", Plan: a.Plan, User: a.User, Windows: []QuotaWindow{}} + c, err := wbFresh(ctx, a) + if err != nil { + q.Error = err.Error() + return q + } + var sum wbResourceSummary + if err := wbCall(ctx, http.MethodPost, wbAPI()+"/v2/billing/meter/get-user-resource-summary", wbAuthHeaders(c), map[string]any{}, &sum); err != nil { + q.Error = err.Error() + return q + } + if sum.IsPaidUser { + q.Plan = firstNonEmpty(a.Plan, "Pro") + } else { + q.Plan = firstNonEmpty(a.Plan, "Free") + } + var total, used float64 + for _, p := range sum.Packages { + total += p.CycleTotalCapacity + used += p.CycleUsedCapacity + } + if total > 0 { + w := QuotaWindow{Name: "Credits", Used: 100 * used / total, Display: fmt.Sprintf("%s / %s", compactNumber(used), compactNumber(total))} + q.Windows = append(q.Windows, w) + } + return q +} + +type wbResourceSummary struct { + Packages []struct { + PackageCode string `json:"PackageCode"` + CycleTotalCapacity float64 `json:"CycleTotalCapacity"` + CycleRemainCapacity float64 `json:"CycleRemainCapacity"` + CycleUsedCapacity float64 `json:"CycleUsedCapacity"` + } `json:"Packages"` + SubscriptionPackageCode string `json:"SubscriptionPackageCode"` + IsPaidUser bool `json:"IsPaidUser"` +} + +func wbAuthHeaders(c wbCreds) map[string]string { + return map[string]string{ + "Authorization": "Bearer " + c.Access, + "X-User-Id": c.UID, + "X-Domain": wbDomain(c), + "X-Product": "SaaS", + "X-IDE-Type": "WorkBuddy", + } +} + +func wbLoginQuota(ctx context.Context, l Login) SubscriptionQuota { + for _, a := range wbLogins() { + if strings.EqualFold(a.User, l.User) { + return wbQuota(ctx, a) + } + } + return SubscriptionQuota{Provider: "workbuddy", Plan: l.Plan, Windows: []QuotaWindow{}, Error: "not signed in"} +} + +// ---- WorkBuddy's API ---------------------------------------------------------- + +// wbError is a WorkBuddy business error: its code lets a poll loop tell a +// "come back" answer (retry) from a real failure. +type wbError struct { + code int + msg string +} + +func (e *wbError) Error() string { + if e.msg != "" { + return e.msg + } + return fmt.Sprintf("error %d", e.code) +} + +// WorkBuddy's retry codes: the token or the account isn't ready yet. +const ( + wbRetryToken = 11217 + wbRetryAccount = 12151 +) + +// wbCall asks one of WorkBuddy's JSON endpoints, which wrap what they say in +// {code, msg, data}: code 0 is a success. On a non-zero code it returns a +// *wbError carrying it. +func wbCall(ctx context.Context, method, u string, headers map[string]string, body, dst any) error { + var rd io.Reader + if body != nil { + b, _ := json.Marshal(body) + rd = bytes.NewReader(b) + } + req, err := http.NewRequestWithContext(ctx, method, u, rd) + if err != nil { + return err + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json") + req.Header.Set("User-Agent", "WorkBuddy/"+wbAppVersion) + for k, v := range headers { + req.Header.Set(k, v) + } + res, err := http.DefaultClient.Do(req) + if err != nil { + return err + } + defer res.Body.Close() + b, _ := io.ReadAll(io.LimitReader(res.Body, 1<<20)) + var env struct { + Code json.Number `json:"code"` + Msg string `json:"msg"` + Data json.RawMessage `json:"data"` + } + _ = json.Unmarshal(b, &env) + if res.StatusCode < 200 || res.StatusCode >= 300 { + if code, _ := env.Code.Int64(); code != 0 { + return &wbError{code: int(code), msg: env.Msg} + } + if env.Msg != "" { + return &wbError{code: res.StatusCode, msg: env.Msg} + } + return &accountStatusError{status: res.StatusCode} + } + if code, _ := env.Code.Int64(); code != 0 { + return &wbError{code: int(code), msg: env.Msg} + } + if dst == nil || len(env.Data) == 0 || string(env.Data) == "null" { + return nil + } + return json.Unmarshal(env.Data, dst) +} + +// ---- signing in --------------------------------------------------------------- + +// startWorkBuddySignIn is WorkBuddy's own external-link sign-in: the app +// asks for a state and a page, opens the page for the user to sign in, then +// polls for the token and the account. +func startWorkBuddySignIn(s *signInFlow) error { + ctx, cancel := context.WithCancel(context.Background()) + var state struct { + State string `json:"state"` + AuthURL string `json:"authUrl"` + } + stateHeaders := map[string]string{ + "X-No-Authorization": "true", + "X-No-User-Id": "true", + "X-No-Enterprise-Id": "true", + "X-No-Department-Info": "true", + } + if err := wbCall(ctx, http.MethodPost, wbAPI()+"/v2/plugin/auth/state?platform=workbuddy", stateHeaders, map[string]any{}, &state); err != nil { + cancel() + return fmt.Errorf("WorkBuddy sign-in: %w", err) + } + u, err := url.Parse(state.AuthURL) + if state.State == "" || err != nil || u.Scheme != "https" { + cancel() + return errors.New("WorkBuddy gave no sign-in page") + } + sid := make([]byte, 16) + _, _ = rand.Read(sid) + q := u.Query() + q.Set("version", wbAppVersion) + q.Set("loginSessionId", hex.EncodeToString(sid)) + u.RawQuery = q.Encode() + s.mu.Lock() + s.st.URL = u.String() + s.stop = cancel + s.mu.Unlock() + + go func() { + defer cancel() + fail := func(msg string) { s.finish(SignInState{State: "failed", Error: msg}) } + deadline := time.Now().Add(5 * time.Minute) + token, ok := wbPoll(ctx, deadline, "/v2/plugin/auth/token?state="+url.QueryEscape(state.State), nil, wbRetryToken, fail) + if !ok { + return + } + var tok wbRefreshed + if json.Unmarshal(token, &tok) != nil || tok.AccessToken == "" { + fail("WorkBuddy gave no token") + return + } + c := wbMergeRefreshed(wbCreds{}, tok) + acctHeaders := map[string]string{ + "Authorization": "Bearer " + c.Access, + "X-Domain": wbDomain(c), + "X-No-User-Id": "true", + "X-No-Enterprise-Id": "true", + } + account, ok := wbPoll(ctx, deadline, "/v2/plugin/login/account?state="+url.QueryEscape(state.State), acctHeaders, wbRetryAccount, fail) + if !ok { + return + } + var acc struct { + UID string `json:"uid"` + Nickname string `json:"nickname"` + PhoneNumber string `json:"phoneNumber"` + } + if json.Unmarshal(account, &acc) != nil || acc.UID == "" { + fail("WorkBuddy gave no account") + return + } + c.UID = acc.UID + who := wbWho(acc.Nickname, acc.PhoneNumber, acc.UID) + auth, _ := json.Marshal(c) + if err := addSideLogin(savedLogin{Agent: "workbuddy", User: who, Auth: auth}, func(savedLogin) {}); err != nil { + fail(err.Error()) + return + } + wbTokens.Lock() + wbTokens.m[c.UID] = c + wbTokens.Unlock() + ownUser, _, hasOwn := wbOwn() + s.finish(SignInState{State: "done", User: who, Using: hasOwn && strings.EqualFold(ownUser, who)}) + }() + return nil +} + +// wbPoll asks path every second until it answers with data, giving up at +// the deadline. A retry code (the answer isn't ready) waits and asks again; +// any other error ends the sign-in through fail. +func wbPoll(ctx context.Context, deadline time.Time, path string, headers map[string]string, retryCode int, fail func(string)) (json.RawMessage, bool) { + for { + select { + case <-ctx.Done(): + return nil, false + case <-time.After(wbPollInterval): + } + if time.Now().After(deadline) { + fail("the sign-in expired; start again") + return nil, false + } + var raw json.RawMessage + err := wbCall(ctx, http.MethodGet, wbAPI()+path, headers, nil, &raw) + switch { + case ctx.Err() != nil: + return nil, false + case err == nil && len(raw) > 0 && string(raw) != "null": + return raw, true + case err == nil: + continue // ready, but empty: ask again + } + var we *wbError + if errors.As(err, &we) && we.code == retryCode { + continue // not ready yet + } + var st *accountStatusError + if errors.As(err, &st) && (st.status == 408 || st.status == 429) { + continue // a hiccup: ask again + } + fail("WorkBuddy sign-in: " + err.Error()) + return nil, false + } +} diff --git a/internal/provider/workbuddy_test.go b/internal/provider/workbuddy_test.go new file mode 100644 index 00000000..80f42118 --- /dev/null +++ b/internal/provider/workbuddy_test.go @@ -0,0 +1,248 @@ +package provider + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "net/url" + "os" + "path/filepath" + "strings" + "sync/atomic" + "testing" + "time" +) + +// wbAuthFile is where WorkBuddy keeps its signed-in session under a temp home. +func wbAuthFile(home string) string { + // mirrors wbAuthPath's darwin layout; the test runs on darwin + return filepath.Join(home, "Library", "Application Support", "CodeBuddyExtension", "Data", "Public", "auth", "workbuddy-desktop.info") +} + +// WorkBuddy's own account is read from its auth store; a second, signed in +// from magpie with WorkBuddy's own flow, is kept beside it with its tokens. +func TestWorkBuddyAccounts(t *testing.T) { + home := signIn(t) + wbTokens.Lock() + wbTokens.m = map[string]wbCreds{} + wbTokens.Unlock() + + // far-future so the own token isn't seen as near expiry + future := time.Now().Add(24 * time.Hour).UnixMilli() + writeFile(t, wbAuthFile(home), map[string]any{ + "account": map[string]any{"uid": "u1", "nickname": "旅行者", "phoneNumber": "13800000000"}, + "auth": map[string]any{ + "accessToken": "own-access", "refreshToken": "own-refresh", + "expiresAt": future, "refreshExpiresAt": future, + "domain": "www.codebuddy.cn", "tokenType": "Bearer", + }, + }) + + var polls atomic.Int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + ok := func(data any) { json.NewEncoder(w).Encode(map[string]any{"code": 0, "data": data}) } + retry := func(code int) { json.NewEncoder(w).Encode(map[string]any{"code": code, "msg": "pending"}) } + switch r.URL.Path { + case "/v2/plugin/auth/state": + ok(map[string]any{"state": "s1", "authUrl": "https://login.codebuddy.cn/oauth?client_id=c"}) + case "/v2/plugin/auth/token": + if polls.Add(1) < 2 { + retry(wbRetryToken) + return + } + ok(map[string]any{"accessToken": "two-access", "refreshToken": "two-refresh", + "expiresIn": 3600, "refreshExpiresIn": 7200, "domain": "www.codebuddy.cn", "tokenType": "Bearer"}) + case "/v2/plugin/login/account": + if r.Header.Get("Authorization") != "Bearer two-access" { + w.WriteHeader(401) + return + } + ok(map[string]any{"uid": "u2", "nickname": "Two", "phoneNumber": ""}) + case "/v2/billing/meter/get-user-resource-summary": + auth := r.Header.Get("Authorization") + if auth != "Bearer own-access" && auth != "Bearer two-access" { + w.WriteHeader(401) + return + } + ok(map[string]any{"IsPaidUser": true, "Packages": []any{ + map[string]any{"PackageCode": "coding", "CycleTotalCapacity": 10000, "CycleRemainCapacity": 7500, "CycleUsedCapacity": 2500}, + }}) + default: + w.WriteHeader(404) + } + })) + defer srv.Close() + oldEnd, oldInt := wbEndpoint, wbPollInterval + wbEndpoint, wbPollInterval = srv.URL, 5*time.Millisecond + defer func() { wbEndpoint, wbPollInterval = oldEnd, oldInt }() + + // WorkBuddy's own, named by its nickname + who, c, ok := wbOwn() + if !ok || who != "旅行者" || c.Access != "own-access" || c.UID != "u1" { + t.Fatalf("own: %v %q %+v", ok, who, c) + } + + st, err := StartSignIn("workbuddy") + if err != nil { + t.Fatal(err) + } + u, _ := url.Parse(st.URL) + if u.Host != "login.codebuddy.cn" || u.Query().Get("version") == "" || u.Query().Get("loginSessionId") == "" { + t.Fatalf("sign-in page: %s", st.URL) + } + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + st, _ = WaitSignIn(ctx, st.ID) + if st.State != "done" || st.User != "Two" || st.Using { + t.Fatalf("done: %+v", st) + } + + var users []string + for _, l := range Logins("workbuddy") { + users = append(users, l.User+map[bool]string{true: "*", false: ""}[l.Active]+map[bool]string{true: "+", false: ""}[l.On]) + } + if got := strings.Join(users, " "); got != "旅行者*+ Two+" { + t.Fatalf("logins: %s", got) + } + + p, ok := find(All(), "workbuddy") + if !ok || p.Account.User != "旅行者" || p.Chat != srv.URL+"/v2" || p.Icon != "workbuddy-color" { + t.Fatalf("first: %+v", p) + } + + // the first signs with the own account's token, id and domain + req, _ := http.NewRequest("POST", p.Chat+"/chat/completions", nil) + if err := p.Sign(context.Background(), req, Chat, []byte(`{}`)); err != nil || + req.Header.Get("Authorization") != "Bearer own-access" || + req.Header.Get("X-User-Id") != "u1" || req.Header.Get("X-Domain") != "www.codebuddy.cn" { + t.Fatalf("own sign: %v %v", err, req.Header) + } + + also := p.AlsoOn() + if len(also) != 1 || also[0].Account.User != "Two" { + t.Fatalf("also on: %+v", also) + } + req2, _ := http.NewRequest("POST", also[0].Chat+"/chat/completions", nil) + if err := also[0].Sign(context.Background(), req2, Chat, []byte(`{}`)); err != nil || + req2.Header.Get("Authorization") != "Bearer two-access" || req2.Header.Get("X-User-Id") != "u2" { + t.Fatalf("second signs with its own token: %v %v", err, req2.Header) + } + if ms := p.Account.models(); len(ms) != len(wbModels) || ms[0].ID != "auto" { + t.Fatalf("models: %+v", ms) + } + + // credits, aggregated from the resource summary + use := LoginUsage(context.Background(), "workbuddy") + q := use["旅行者"] + if q.Error != "" || len(q.Windows) != 1 || q.Windows[0].Name != "Credits" || q.Windows[0].Used != 25 { + t.Fatalf("usage: %+v", q) + } + + if err := ForgetLogin("workbuddy", "旅行者"); err == nil { + t.Fatal("forgot the one in use") + } + if err := SwitchLogin("workbuddy", "Two"); err != nil { + t.Fatal(err) + } + if p, _ := find(All(), "workbuddy"); p.Account.User != "Two" { + t.Fatalf("after switch: %+v", p.Account) + } + if err := SwitchLogin("workbuddy", "旅行者"); err != nil { + t.Fatal(err) + } + if err := ForgetLogin("workbuddy", "Two"); err != nil { + t.Fatal(err) + } + if ls := Logins("workbuddy"); len(ls) != 1 { + t.Fatalf("after forget: %+v", ls) + } + + // signed out of WorkBuddy, with none of magpie's: no WorkBuddy subscription + os.Remove(wbAuthFile(home)) + if _, ok := find(All(), "workbuddy"); ok { + t.Fatal("workbuddy without an account") + } +} + +// A near-expired own token is refreshed before use, and the refreshed one +// is kept in memory (WorkBuddy's file is never written). +func TestWorkBuddyRefresh(t *testing.T) { + home := signIn(t) + wbTokens.Lock() + wbTokens.m = map[string]wbCreds{} + wbTokens.Unlock() + + past := time.Now().Add(-time.Minute).UnixMilli() // access already lapsed + far := time.Now().Add(24 * time.Hour).UnixMilli() // refresh still good + writeFile(t, wbAuthFile(home), map[string]any{ + "account": map[string]any{"uid": "u1", "nickname": "R", "phoneNumber": ""}, + "auth": map[string]any{ + "accessToken": "stale", "refreshToken": "good-refresh", + "expiresAt": past, "refreshExpiresAt": far, "domain": "www.codebuddy.cn", "tokenType": "Bearer", + }, + }) + + var refreshed atomic.Int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/v2/plugin/auth/token/refresh" { + if r.Header.Get("X-Refresh-Token") != "good-refresh" { + w.WriteHeader(401) + return + } + refreshed.Add(1) + json.NewEncoder(w).Encode(map[string]any{"code": 0, "data": map[string]any{ + "accessToken": "fresh", "expiresIn": 3600, + }}) + return + } + w.WriteHeader(404) + })) + defer srv.Close() + old := wbEndpoint + wbEndpoint = srv.URL + defer func() { wbEndpoint = old }() + + p, ok := find(All(), "workbuddy") + if !ok { + t.Fatal("no workbuddy account") + } + req, _ := http.NewRequest("POST", p.Chat+"/chat/completions", nil) + if err := p.Sign(context.Background(), req, Chat, []byte(`{}`)); err != nil { + t.Fatal(err) + } + if got := req.Header.Get("Authorization"); got != "Bearer fresh" { + t.Fatalf("Authorization = %q, want the refreshed token", got) + } + if refreshed.Load() != 1 { + t.Fatalf("refreshed %d times, want 1", refreshed.Load()) + } + // a second request reuses the cached fresh token, no new refresh + req2, _ := http.NewRequest("POST", p.Chat+"/chat/completions", nil) + if err := p.Sign(context.Background(), req2, Chat, []byte(`{}`)); err != nil { + t.Fatal(err) + } + if refreshed.Load() != 1 { + t.Fatalf("refreshed again: %d", refreshed.Load()) + } + // WorkBuddy's own file is left as it was + var s wbStoredSession + if !readJSON(wbAuthFile(home), &s) || wbString(s.Auth.AccessToken) != "stale" { + t.Fatalf("own file was written: %+v", s.Auth) + } +} + +// An encrypted access token at rest means magpie can't read the account. +func TestWorkBuddyEncryptedOwn(t *testing.T) { + home := signIn(t) + writeFile(t, wbAuthFile(home), map[string]any{ + "account": map[string]any{"uid": "u1", "nickname": "E"}, + "auth": map[string]any{ + "accessToken": map[string]any{"$wbEncrypted": 1, "envelope": "x", "scheme": "asym-v1"}, + "domain": "www.codebuddy.cn", + }, + }) + if _, _, ok := wbOwn(); ok { + t.Fatal("read an encrypted own account") + } +}