diff --git a/internal/provider/codex_workspaces_test.go b/internal/provider/codex_workspaces_test.go new file mode 100644 index 00000000..870528fd --- /dev/null +++ b/internal/provider/codex_workspaces_test.go @@ -0,0 +1,168 @@ +package provider + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "testing" + "time" +) + +// codexWorkspaceSignIn signs Codex in to a seat of email's in a Team +// workspace: two of them read alike, "email · Team". +func codexWorkspaceSignIn(t *testing.T, home, email, workspace, refresh string) { + t.Helper() + exp := float64(time.Now().Add(time.Hour).Unix()) + writeFile(t, filepath.Join(home, ".codex", "auth.json"), map[string]any{ + "auth_mode": "chatgpt", + "tokens": map[string]any{ + "id_token": fakeJWT(map[string]any{"email": email, "https://api.openai.com/auth": map[string]any{ + "chatgpt_plan_type": "team", "chatgpt_account_id": workspace}}), + "access_token": fakeJWT(map[string]any{"exp": exp}), + "refresh_token": refresh, "account_id": workspace, + }, + }) +} + +// workspaceOf is the workspace a saved Codex account's credentials are +// for, and their refresh token. +func workspaceOf(t *testing.T, user string) (workspace, refresh string) { + t.Helper() + for _, l := range readLogins() { + if l.Agent == "codex" && l.User == user { + var a codexAuth + json.Unmarshal(l.Auth, &a) + return a.Tokens.AccountID, a.Tokens.RefreshToken + } + } + t.Fatalf("%q isn't saved: %+v", user, readLogins()) + return "", "" +} + +// Two Codex accounts of one email in two Team workspaces are two accounts +// in every way (vincentzhang on Discord: removing one, or codex logout on +// one, had the other ask to be signed in again). They read alike, and +// everything that went by the name took both: removing the one not in use +// signed Codex out of the other (it looked like the last), and a refresh +// of one was written over the other's credentials too. +func TestCodexWorkspacesOfOneEmail(t *testing.T) { + home := signIn(t) + auth := filepath.Join(home, ".codex", "auth.json") + codexWorkspaceSignIn(t, home, "me@example.com", "ws-one", "r-one") + rememberLogins(true) + codexWorkspaceSignIn(t, home, "me@example.com", "ws-two", "r-two") + rememberLogins(true) + + users, active := loginUsers(Logins("codex")) + if len(users) != 2 { + t.Fatalf("accounts: %v", users) + } + seats := []string{} + for _, u := range users { + if strings.HasPrefix(u, "me@example.com · Team") { + seats = append(seats, u) + } + } + if len(seats) != 2 || strings.EqualFold(seats[0], seats[1]) { + t.Fatalf("the two seats read alike: %v", users) + } + actives := 0 + for _, l := range Logins("codex") { + if l.Active { + actives++ + } + } + if actives != 1 { + t.Fatalf("%d accounts in use: %+v", actives, Logins("codex")) + } + if ws, _ := workspaceOf(t, active); ws != "ws-two" { + t.Fatalf("in use: %q is %s", active, ws) + } + other := seats[0] + if other == active { + other = seats[1] + } + if ws, rt := workspaceOf(t, other); ws != "ws-one" || rt != "r-one" { + t.Fatalf("the other: %s %s", ws, rt) + } + + // a refresh of the one not in use is its own, not the other's too + fake := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var body map[string]string + json.NewDecoder(r.Body).Decode(&body) + json.NewEncoder(w).Encode(map[string]any{"access_token": fakeJWT(map[string]any{"exp": float64(time.Now().Add(time.Hour).Unix())}), + "refresh_token": body["refresh_token"] + "-2"}) + })) + defer fake.Close() + old := codexTokenURL + codexTokenURL = fake.URL + t.Cleanup(func() { codexTokenURL = old }) + if _, _, err := renewSavedLogin(context.Background(), "codex", other, true); err != nil { + t.Fatal(err) + } + if ws, rt := workspaceOf(t, other); ws != "ws-one" || rt != "r-one-2" { + t.Fatalf("the one refreshed: %s %s", ws, rt) + } + if ws, rt := workspaceOf(t, active); ws != "ws-two" || rt != "r-two" { + t.Fatalf("the one in use was written over: %s %s", ws, rt) + } + + // removing the one not in use leaves Codex signed in to the other + if err := ForgetLogin("codex", other); err != nil { + t.Fatal(err) + } + if _, err := os.Stat(auth); err != nil { + t.Fatalf("Codex was signed out of %s: %v", active, err) + } + if ws, rt := workspaceOf(t, active); ws != "ws-two" || rt != "r-two" { + t.Fatalf("the one in use: %s %s", ws, rt) + } + users, now := loginUsers(Logins("codex")) + if len(users) != 1 || now != active { + t.Fatalf("after removing %s: %v, in use %q", other, users, now) + } +} + +// Two saved by one name before this are told apart when read, and the one +// Codex isn't on can be removed without signing it out of the other. +func TestCodexWorkspacesSavedAlike(t *testing.T) { + home := signIn(t) + auth := filepath.Join(home, ".codex", "auth.json") + var saved []savedLogin + for _, ws := range []string{"ws-one", "ws-two"} { + codexWorkspaceSignIn(t, home, "me@example.com", ws, "r-"+ws) + b, err := os.ReadFile(auth) + if err != nil { + t.Fatal(err) + } + saved = append(saved, savedLogin{Agent: "codex", User: "me@example.com · Team", Plan: "team", Auth: b}) + } + if err := writeLogins(saved); err != nil { + t.Fatal(err) + } + ls := Logins("codex") + users, active := loginUsers(ls) + if len(users) != 2 || strings.EqualFold(users[0], users[1]) { + t.Fatalf("accounts: %v", users) + } + if ws, _ := workspaceOf(t, active); ws != "ws-two" { + t.Fatalf("in use: %q is %s", active, ws) + } + other := users[0] + if other == active { + other = users[1] + } + if err := ForgetLogin("codex", other); err != nil { + t.Fatal(err) + } + if _, err := os.Stat(auth); err != nil { + t.Fatalf("Codex was signed out of %s: %v", active, err) + } + if users, now := loginUsers(Logins("codex")); len(users) != 1 || now != active { + t.Fatalf("after removing %s: %v, in use %q", other, users, now) + } +} diff --git a/internal/provider/logins.go b/internal/provider/logins.go index 17e60cca..12655e82 100644 --- a/internal/provider/logins.go +++ b/internal/provider/logins.go @@ -124,7 +124,7 @@ func readLogins() []savedLogin { // doesn't allow its subscription used outside its client), so one // signed in before is left out, and gone from the file at its next write out = slices.DeleteFunc(out, func(l savedLogin) bool { return l.Agent == "dimagent" }) - return dedupeLogins(out), nil + return nameAlike(dedupeLogins(out)), nil }) return slices.Clone(ls) // callers change theirs } @@ -184,6 +184,7 @@ func writePrivate(path string, b []byte) error { } func upsertLogin(ls []savedLogin, l savedLogin) []savedLogin { + l.User = codexName(ls, l) for i := range ls { if sameLogin(ls[i], l) { // a refused Claude credential stays refused while it is the @@ -279,6 +280,58 @@ func codexWho(auth json.RawMessage) (email, workspace string) { return claimString(id, "email"), workspace } +// codexName is the name a Codex account goes by among the saved ones ls: +// codexUser's, unless another account goes by that already — two seats of +// one email in two Team workspaces read alike — and then the one it was +// saved under, or for a new one the name with its workspace after it. +// Every account is told by its name (switched to, refreshed, removed), and +// two by one name were taken as one: removing the one not in use signed +// Codex out of the other, as the last account, and a refresh of one was +// written over the other's credentials (vincentzhang on Discord). +func codexName(ls []savedLogin, l savedLogin) string { + if l.Agent != "codex" { + return l.User + } + taken := func(user string) bool { + return slices.ContainsFunc(ls, func(x savedLogin) bool { + return x.Agent == "codex" && strings.EqualFold(x.User, user) && !sameLogin(x, l) + }) + } + for _, x := range ls { + // told apart once, it keeps that name; else it takes a new plan's + if x.Agent == "codex" && sameLogin(x, l) && (taken(l.User) || strings.HasPrefix(strings.ToLower(x.User), strings.ToLower(l.User)+" · ")) { + return x.User + } + } + if !taken(l.User) { + return l.User + } + _, ws := codexWho(l.Auth) + if len(ws) > 8 { + ws = ws[:8] + } + name := l.User + if ws != "" { + name += " · " + ws + } + for n := 2; taken(name); n++ { + name = fmt.Sprintf("%s · %s (%d)", l.User, ws, n) + } + return name +} + +// nameAlike gives each Codex account in ls a name of its own (codexName), +// the first by a name keeping it: two saved by one name before are told +// apart from the next write on. +func nameAlike(ls []savedLogin) []savedLogin { + for i := range ls { + if ls[i].Agent == "codex" { + ls[i].User = codexName(ls[:i], ls[i]) + } + } + return ls +} + // codexUser names a ChatGPT account from its ID token's claims: its email, // and for a seat in a workspace the plan too, so it reads apart from a // personal plan of the same email. @@ -433,8 +486,11 @@ func liveLogin(agent string) (savedLogin, bool) { if user == "" { return savedLogin{}, false } - return savedLogin{Agent: agent, User: user, Plan: claimString(id, "https://api.openai.com/auth", "chatgpt_plan_type"), - Auth: json.RawMessage(bytes.TrimSpace(b))}, true + l := savedLogin{Agent: agent, User: user, Plan: claimString(id, "https://api.openai.com/auth", "chatgpt_plan_type"), + Auth: json.RawMessage(bytes.TrimSpace(b))} + // by the name it is saved under, which may not be codexUser's + l.User = codexName(readLogins(), l) + return l, true case "claude": c, _, ok := claudeCredential() if !ok { diff --git a/internal/provider/signin.go b/internal/provider/signin.go index c2549d06..a33b3aee 100644 --- a/internal/provider/signin.go +++ b/internal/provider/signin.go @@ -765,6 +765,7 @@ func addLogin(l savedLogin) (using bool, err error) { l.Seen = time.Now().UTC().Truncate(time.Second) live, signedIn := liveLogin(l.Agent) ls := readLogins() + l.User = codexName(ls, l) using = !signedIn || sameLogin(live, l) if first := claudeStandIn(ls); !signedIn && l.Agent == "claude" && first != "" { // logged out of Claude Code with accounts in magpie: it stays so,