From 9687f3d297a3efc19ccb1c150d7caf18f61739da Mon Sep 17 00:00:00 2001 From: yetone Date: Mon, 28 Sep 2026 16:31:06 +0800 Subject: [PATCH] =?UTF-8?q?provider:=20WorkBuddy=20is=20a=20subscription?= =?UTF-8?q?=20=E2=80=94=20Tencent's=20CodeBuddy=20plan,=20signed=20in=20as?= =?UTF-8?q?=20WorkBuddy=20does,=20its=20own=20account=20read=20too?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit WorkBuddy (CodeBuddy Code's desktop app) keeps its signed-in session in /auth/workbuddy-desktop.info, plain JSON while credential protection is off (its default); magpie reads that account, never writes it, and skips it when the tokens are encrypted at rest. Further accounts are signed in with WorkBuddy's own external-link flow (auth/state, then auth/token polled until ready, then login/account) and kept in logins.json. The plan is served OpenAI-compatibly at copilot.tencent.com /v2 under the account's access token, with its id and domain in headers (X-User-Id, X-Domain), exactly as WorkBuddy sends them; a near-expired JWT is refreshed as WorkBuddy does, and a magpie-signed-in one's refresh is written back beside it while WorkBuddy's own file is left untouched. The models are its CLI agent's coding picks, and the credits come from its resource summary as one window. A card in Providers' subscription list signs in and shows the accounts; the Provider is one of the account ids and joins the usage rollup. --- internal/gui/assets/app.js | 3 + internal/gui/assets/icons/workbuddy-color.svg | 1 + internal/provider/account.go | 3 + internal/provider/login_usage.go | 5 + internal/provider/logins.go | 7 + internal/provider/logins_on.go | 5 + internal/provider/provider.go | 2 +- internal/provider/side_logins.go | 2 +- internal/provider/signin.go | 5 + internal/provider/subscription_usage.go | 3 + internal/provider/workbuddy.go | 669 ++++++++++++++++++ internal/provider/workbuddy_test.go | 248 +++++++ 12 files changed, 951 insertions(+), 2 deletions(-) create mode 100644 internal/gui/assets/icons/workbuddy-color.svg create mode 100644 internal/provider/workbuddy.go create mode 100644 internal/provider/workbuddy_test.go 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") + } +}