magpie/internal/provider/commandcode_plan_test.go

package provider

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"net/url"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"
)

func TestCommandCodePlan(t *testing.T) {
	home := signIn(t)
	writeFile(t, filepath.Join(home, ".commandcode", "auth.json"), map[string]any{
		"apiKey": "own-key", "userId": "u1", "userName": "ownuser", "keyName": "cli", "authenticatedAt": "2026-09-01T00:00:00Z",
	})
	reset := time.Now().Add(2 * time.Hour).Unix()
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		key := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ")
		if key != "own-key" && key != "two-key" {
			w.WriteHeader(401)
			return
		}
		ok := func(v any) { _ = json.NewEncoder(w).Encode(v) }
		switch r.URL.Path {
		case "/alpha/whoami":
			ok(map[string]any{"user": map[string]any{"id": "u2", "userName": "two", "email": "two@example.com"}})
		case "/alpha/billing/credits":
			ok(map[string]any{
				"credits": map[string]any{"planId": "individual-max-monthly", "monthlyCredits": 150.5, "purchasedCredits": 10, "freeCredits": 0},
				"windowLimits": map[string]any{"limited": false,
					"fiveHour": map[string]any{"used": 25, "cap": 100, "resetAt": reset},
					"weekly":   map[string]any{"used": 30, "cap": 300, "resetAt": time.Now().Add(72 * time.Hour).Format(time.RFC3339)},
				},
			})
		default:
			w.WriteHeader(404)
		}
	}))
	defer srv.Close()
	oldAPI := cmdAPI
	cmdAPI = srv.URL
	defer func() { cmdAPI = oldAPI }()

	if who, a, ok := cmdOwn(); !ok || who != "ownuser" || a.APIKey != "own-key" {
		t.Fatalf("own: %v %q", ok, who)
	}

	st, err := StartSignIn(CommandCodePlanID)
	if err != nil {
		t.Fatal(err)
	}
	u, _ := url.Parse(st.URL)
	cb := u.Query().Get("callback")
	if u.Host != "commandcode.ai" || u.Path != "/studio/auth/cli" || u.Query().Get("mode") != "redirect" ||
		!strings.HasPrefix(cb, "http://127.0.0.1:") || !strings.HasSuffix(cb, "/callback") {
		t.Fatalf("sign-in page: %s", st.URL)
	}
	state := u.Query().Get("state")
	noFollow := &http.Client{CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }}

	// the browser's preflight, from Studio
	pre, _ := http.NewRequest(http.MethodOptions, cb, nil)
	pre.Header.Set("Origin", "https://commandcode.ai")
	pre.Header.Set("Access-Control-Request-Private-Network", "true")
	if res, err := noFollow.Do(pre); err != nil || res.StatusCode != 204 ||
		res.Header.Get("Access-Control-Allow-Origin") != "https://commandcode.ai" || res.Header.Get("Access-Control-Allow-Private-Network") != "true" {
		t.Fatalf("preflight: %v %+v", err, res)
	}
	// someone else's state is turned away, and the sign-in goes on
	if res, err := noFollow.PostForm(cb, url.Values{"apiKey": {"evil"}, "state": {"nope"}}); err != nil || res.StatusCode != 403 {
		t.Fatalf("wrong state: %v %+v", err, res)
	}
	if res, _ := noFollow.Get(cb); res.StatusCode != 405 {
		t.Fatalf("GET callback: %d", res.StatusCode)
	}

	res, err := noFollow.PostForm(cb, url.Values{"apiKey": {"two-key"}, "state": {state}, "userId": {"u2"}, "userName": {"two"}, "keyName": {"magpie"}})
	if err != nil || res.StatusCode != 303 || !strings.HasPrefix(res.Header.Get("Location"), "/callback/complete?state=") {
		t.Fatalf("callback: %v %+v", err, res)
	}
	page, err := http.Get(strings.TrimSuffix(cb, "/callback") + res.Header.Get("Location"))
	if err != nil || page.StatusCode != 200 {
		t.Fatalf("complete page: %v %+v", err, page)
	}
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	st, _ = WaitSignIn(ctx, st.ID)
	if st.State != "done" || st.User != "two" || st.Plan != "Max" || st.Using {
		t.Fatalf("done: %+v", st)
	}

	var users []string
	for _, l := range Logins(CommandCodePlanID) {
		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 != "ownuser*+ two+" {
		t.Fatalf("logins: %s", got)
	}
	p, ok := find(All(), CommandCodePlanID)
	if !ok || p.Account.User != "ownuser" || p.Chat != srv.URL+"/provider/v1" || p.Anthropic != srv.URL+"/provider" {
		t.Fatalf("first: %+v", p)
	}
	also := p.AlsoOn()
	if len(also) != 1 || also[0].Account.User != "two" {
		t.Fatalf("also on: %+v", also)
	}
	req, _ := http.NewRequest("POST", also[0].Chat+"/chat/completions", nil)
	req.Header.Set("Authorization", "Bearer magpie")
	if err := also[0].Sign(context.Background(), req, Chat, []byte(`{}`)); err != nil ||
		req.Header.Get("Authorization") != "Bearer two-key" || req.Header.Get("x-api-key") != "two-key" {
		t.Fatalf("signs with its own key: %v %v", err, req.Header)
	}

	q := LoginUsage(context.Background(), CommandCodePlanID)["two"]
	if q.Error != "" || q.Plan != "Max" || len(q.Windows) != 2 || q.Windows[0].Name != "5 hours" || q.Windows[0].Used != 25 ||
		q.Windows[1].Name != "Weekly" || q.Windows[1].Used != 10 || q.Windows[0].ResetsAt == nil || q.Windows[1].ResetsAt == nil ||
		q.Windows[0].ResetsAt.Unix() != reset {
		t.Fatalf("usage: %+v", q)
	}
	if !strings.Contains(q.Balance, "160.5") {
		t.Fatalf("balance: %q", q.Balance)
	}

	if err := ForgetLogin(CommandCodePlanID, "ownuser"); err == nil {
		t.Fatal("forgot the CLI's own")
	}
	if err := SwitchLogin(CommandCodePlanID, "two"); err != nil {
		t.Fatal(err)
	}
	if p, _ := find(All(), CommandCodePlanID); p.Account.User != "two" {
		t.Fatalf("after switch: %+v", p.Account)
	}
	if err := SwitchLogin(CommandCodePlanID, "ownuser"); err != nil {
		t.Fatal(err)
	}
	if err := ForgetLogin(CommandCodePlanID, "two"); err != nil {
		t.Fatal(err)
	}
	if ls := Logins(CommandCodePlanID); len(ls) != 1 {
		t.Fatalf("after forget: %+v", ls)
	}

	// a sign-in the user turns down fails
	st, _ = StartSignIn(CommandCodePlanID)
	u, _ = url.Parse(st.URL)
	body, _ := json.Marshal(map[string]any{"state": u.Query().Get("state"), "error": "access_denied"})
	if res, err := http.Post(u.Query().Get("callback"), "application/json", strings.NewReader(string(body))); err != nil || res.StatusCode != 400 {
		t.Fatalf("denied: %v %+v", err, res)
	}
	st, _ = WaitSignIn(ctx, st.ID)
	if st.State != "failed" || st.Error != "the sign-in was denied" {
		t.Fatalf("denied: %+v", st)
	}

	// signed out of the CLI, with none of magpie's: no Command Code subscription
	os.Remove(filepath.Join(home, ".commandcode", "auth.json"))
	if _, ok := find(All(), CommandCodePlanID); ok {
		t.Fatal("command code without an account")
	}
}

func TestCommandCodePlanName(t *testing.T) {
	for id, want := range map[string]string{
		"individual-go": "Go", "individual-goat": "GOAT", "individual-pro-v1": "Pro", "individual_ultra": "Ultra",
		"teams-pro": "Teams Pro", "individual-provider": "Provider", "enterprise": "",
	} {
		if got := cmdPlanName(id); got != want {
			t.Errorf("%s: %q, want %q", id, got, want)
		}
	}
	if tm := cmdTime(float64(1_900_000_000_000)); tm == nil || tm.Unix() != 1_900_000_000 {
		t.Fatalf("ms: %v", tm)
	}
}