magpie/internal/provider/zcode_start.go

package provider

// ZCode's Start Plan (体验套餐) is the free allowance ZCode gives a Z.ai or
// BigModel account that has no GLM Coding Plan. It is not served where the
// Coding Plan is: ZCode sends its requests to zcode.z.ai itself,
// /api/v1/zcode-plan/anthropic, with its own session token (the
// credential store's zcodejwttoken, the `token` its sign-in's poll hands
// back) as the key, and reads what is left from
// /api/v1/zcode-plan/billing/balance with the same token. The account's
// zcode-api-key on api.z.ai/api/anthropic is billed to the Coding Plan, or
// else to the account's API balance, and refused with 1113 "Insufficient
// balance or no resource package" when there is neither.
//
// So an account goes to the Start Plan when it has ZCode's token and no
// Coding Plan (as /api/biz/subscription/list says), or has no Coding Plan
// key at all; to the Coding Plan otherwise, as before.

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"net/url"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/yetone/magpie/internal/catalog"
)

// zcodeStartBase is where the Start Plan is served.
func zcodeStartBase() string { return zcodeAPI + "/api/v1/zcode-plan/anthropic" }

// zcodeStartModels are the Start Plan's models before ZCode's config is
// read: its builtinModelIds for account:zai-start-plan.
var zcodeStartModels = func() []catalog.Model {
	var out []catalog.Model
	for _, m := range zcodeModels {
		if m.ID != "GLM-5.3" {
			out = append(out, m)
		}
	}
	return out
}()

// ---- which plan -----------------------------------------------------------------

var zcodeRoutes = struct {
	sync.Mutex
	m map[string]zcodeRoute
}{m: map[string]zcodeRoute{}}

type zcodeRoute struct {
	start bool
	at    time.Time
	ttl   time.Duration
}

// zcodeOnStart says whether k's requests go to the Start Plan, asking
// when the last answer is old; with ctx nil it only says what was found
// last (the Coding Plan when nothing was).
func zcodeOnStart(ctx context.Context, k zcodeKey) bool {
	if k.JWT == "" || k.team() { // a team's seat is on the team's plan (zcode_team.go)
		return false
	}
	if k.Key == "" {
		return true
	}
	id := k.Key + "\x00" + k.JWT
	zcodeRoutes.Lock()
	r, ok := zcodeRoutes.m[id]
	zcodeRoutes.Unlock()
	if ctx == nil || ok && time.Since(r.at) < r.ttl {
		return r.start
	}
	start, sure := zcodeDecide(ctx, k)
	ttl := 10 * time.Minute
	if !sure {
		ttl = time.Minute
	}
	zcodeRoutes.Lock()
	zcodeRoutes.m[id] = zcodeRoute{start: start, at: time.Now(), ttl: ttl}
	zcodeRoutes.Unlock()
	return start
}

// zcodeDecide asks whether the account has a Coding Plan: with one it is
// used; with none the Start Plan is all the account may have. When that
// can't be told, the Start Plan is used if it is there.
func zcodeDecide(ctx context.Context, k zcodeKey) (start, sure bool) {
	plan, err := zcodePlan(ctx, k)
	switch {
	case err == nil && plan != "":
		return false, true
	case err == nil:
		return true, true
	}
	if b, err := zcodeStartBalance(ctx, k.JWT); err == nil {
		if _, _, ok := b.active(); ok {
			return true, true
		}
	}
	return false, false
}

// zcodeRebase moves a request made for one plan's endpoint to another's.
func zcodeRebase(req *http.Request, from, to string) {
	if from == to || from == "" {
		return
	}
	rest, ok := strings.CutPrefix(req.URL.String(), from)
	if !ok {
		return
	}
	if u, err := url.Parse(to + rest); err == nil {
		req.URL, req.Host = u, u.Host
	}
}

// zcodeSourceHeaders are the headers ZCode names itself with on a model
// request.
func zcodeSourceHeaders(req *http.Request) {
	platform := runtime.GOOS
	if platform == "windows" {
		platform = "win32"
	}
	arch := runtime.GOARCH
	if arch == "amd64" {
		arch = "x64"
	}
	req.Header.Set("User-Agent", "ZCode/"+zcodeAppVersion)
	req.Header.Set("X-ZCode-App-Version", zcodeAppVersion)
	req.Header.Set("X-Title", "Z Code@electron")
	req.Header.Set("HTTP-Referer", zcodeAPI)
	req.Header.Set("X-Platform", platform+"-"+arch)
}

// zcodeJWTExpired says whether ZCode's token has run out; ZCode then asks
// to sign in again, and so does magpie.
func zcodeJWTExpired(jwt string) bool {
	parts := strings.Split(jwt, ".")
	if len(parts) != 3 {
		return false
	}
	b, err := base64.RawURLEncoding.DecodeString(strings.TrimRight(parts[1], "="))
	if err != nil {
		return false
	}
	var claims struct {
		Exp float64 `json:"exp"`
	}
	if json.Unmarshal(b, &claims) != nil || claims.Exp == 0 {
		return false
	}
	return time.Now().After(time.Unix(int64(claims.Exp), 0))
}

var errZCodeExpired = errors.New("ZCode's sign-in has expired; sign in to ZCode again (or add the account again in magpie)")

// ---- the allowance ----------------------------------------------------------------

type zcodeBalance struct {
	ServerTime any `json:"server_time"`
	Plans      []struct {
		PlanID       string `json:"plan_id"`
		UserPlanID   string `json:"user_plan_id"`
		Name         string `json:"name"`
		Status       string `json:"status"`
		EndsAt       any    `json:"ends_at"`
		Entitlements []struct {
			ID     string `json:"entitlement_id"`
			Period string `json:"period"`
		} `json:"entitlements"`
	} `json:"plans"`
	Balances []struct {
		PlanID       string   `json:"plan_id"`
		UserPlanID   string   `json:"user_plan_id"`
		Entitlement  string   `json:"entitlement_id"`
		ShowName     string   `json:"show_name"`
		Capabilities []string `json:"capabilities"`
		Total        any      `json:"total_units"`
		Used         any      `json:"used_units"`
		Remaining    any      `json:"remaining_units"`
		ExpiresAt    any      `json:"expires_at"`
		PeriodStart  any      `json:"period_start"`
		PeriodEnd    any      `json:"period_end"`
	} `json:"balances"`
}

// zcodeNum reads a number the balance gives as a number or a string.
func zcodeNum(v any) (float64, bool) {
	switch x := v.(type) {
	case float64:
		return x, true
	case string:
		f, err := strconv.ParseFloat(strings.TrimSpace(x), 64)
		return f, err == nil
	}
	return 0, false
}

// zcodeStartBalance is the Start Plan's balance, as ZCode reads it: a plan
// still "active" past its end is over, and the buckets of plans that are
// over are left out.
func zcodeStartBalance(ctx context.Context, jwt string) (zcodeBalance, error) {
	var b zcodeBalance
	if jwt == "" {
		return b, errors.New("not signed in to ZCode")
	}
	if zcodeJWTExpired(jwt) {
		return b, errZCodeExpired
	}
	u := zcodeAPI + "/api/v1/zcode-plan/billing/balance?" + url.Values{"app_version": {zcodeAppVersion}}.Encode()
	if err := zcodeCall(ctx, http.MethodGet, u, "Bearer "+jwt, nil, &b); err != nil {
		return b, err
	}
	now := float64(time.Now().Unix())
	if t, ok := zcodeNum(b.ServerTime); ok && t > 0 {
		now = t
	}
	over := map[string]bool{}
	for i, p := range b.Plans {
		if end, ok := zcodeNum(p.EndsAt); ok && end > 0 && end <= now && strings.EqualFold(strings.TrimSpace(p.Status), "active") {
			b.Plans[i].Status = "expired"
		}
		if strings.EqualFold(strings.TrimSpace(b.Plans[i].Status), "expired") {
			over[p.UserPlanID+"\x00"+p.PlanID] = true
		}
	}
	bs := b.Balances[:0]
	for _, x := range b.Balances {
		keep := true
		for _, p := range b.Plans {
			if x.UserPlanID != "" && p.UserPlanID != "" && x.UserPlanID == p.UserPlanID || (x.UserPlanID == "" || p.UserPlanID == "") && x.PlanID == p.PlanID {
				keep = !over[p.UserPlanID+"\x00"+p.PlanID]
				if keep {
					break
				}
			}
		}
		if keep {
			bs = append(bs, x)
		}
	}
	b.Balances = bs
	return b, nil
}

// active is the Start Plan the account has now, as ZCode tells it: an
// active plan whose id or name says start plan (or that has neither).
func (b zcodeBalance) active() (name string, until *time.Time, ok bool) {
	for _, p := range b.Plans {
		if !strings.EqualFold(strings.TrimSpace(p.Status), "active") {
			continue
		}
		id, n := strings.ToLower(strings.TrimSpace(p.PlanID)), strings.ToLower(strings.TrimSpace(p.Name))
		isStart := func(s string) bool { return strings.Contains(s, "start-plan") || strings.Contains(s, "start plan") }
		if id != "" || n != "" {
			if !isStart(id) && !isStart(n) {
				continue
			}
		}
		if end, ok := zcodeNum(p.EndsAt); ok && end > 0 {
			t := time.Unix(int64(end), 0)
			until = &t
		}
		return firstNonEmpty(p.Name, "Start Plan"), until, true
	}
	return "", nil, false
}

// zcodePeriod reads an entitlement's period: "daily", "weekly", "monthly".
func zcodePeriod(p string) time.Duration {
	p = strings.ToLower(p)
	switch {
	case strings.Contains(p, "day") || strings.Contains(p, "daily"):
		return 24 * time.Hour
	case strings.Contains(p, "week"):
		return 7 * 24 * time.Hour
	case strings.Contains(p, "month"):
		return 30 * 24 * time.Hour
	}
	return 0
}

// zcodeStartQuota is the Start Plan's allowance: a window for each of its
// buckets, a model's tokens for the day or for the plan's time, as ZCode
// shows them.
func zcodeStartQuota(ctx context.Context, l Login, jwt string) SubscriptionQuota {
	q := SubscriptionQuota{Provider: "zcode", Name: "ZCode", Icon: "zcode", Plan: l.Plan, User: l.User, Windows: []QuotaWindow{}}
	b, err := zcodeStartBalance(ctx, jwt)
	if err != nil {
		q.Error = err.Error()
		return q
	}
	name, until, ok := b.active()
	if !ok {
		q.Error = "this account has no GLM Coding Plan, and ZCode's Start Plan has ended or was never started"
		return q
	}
	q.Plan = name
	if until != nil {
		q.Until, q.Renew = until, "off"
	}
	for _, x := range b.Balances {
		total, hasTotal := zcodeNum(x.Total)
		used, hasUsed := zcodeNum(x.Used)
		left, hasLeft := zcodeNum(x.Remaining)
		if !hasTotal && !hasUsed && !hasLeft {
			continue
		}
		var models []string
		for _, c := range x.Capabilities {
			if m := strings.TrimSpace(strings.TrimPrefix(strings.TrimSpace(c), "model:")); m != "" {
				models = append(models, m)
			}
		}
		w := QuotaWindow{Name: firstNonEmpty(x.ShowName, strings.Join(models, ", "), "Credits")}
		if !hasUsed && hasTotal && hasLeft {
			used = total - left
		}
		if hasTotal && total > 0 {
			w.Used = 100 * used / total
			w.Display = fmt.Sprintf("%s / %s", compactNumber(used), compactNumber(total))
		}
		if t, ok := zcodeNum(x.ExpiresAt); ok && t > 0 {
			r := time.Unix(int64(t), 0)
			w.ResetsAt = &r
		}
		for _, p := range b.Plans {
			if x.UserPlanID != "" && p.UserPlanID != "" && x.UserPlanID != p.UserPlanID || (x.UserPlanID == "" || p.UserPlanID == "") && x.PlanID != p.PlanID {
				continue
			}
			for _, e := range p.Entitlements {
				if e.ID == x.Entitlement {
					w.Span = zcodePeriod(e.Period)
				}
			}
		}
		if w.Span == 0 {
			s, ok1 := zcodeNum(x.PeriodStart)
			e, ok2 := zcodeNum(x.PeriodEnd)
			if ok1 && ok2 && e > s {
				w.Span = time.Duration(e-s) * time.Second
			}
		}
		if len(models) > 0 {
			w.matches = func(model string) bool {
				for _, m := range models {
					if strings.EqualFold(m, model) {
						return true
					}
				}
				return false
			}
		}
		q.Windows = append(q.Windows, w)
	}
	return q
}