Initial sanitized code sync

internal/task/membership.go

@@ -0,0 +1,229 @@
+package task
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"gpt-plus/pkg/auth"
+	"gpt-plus/pkg/chatgpt"
+)
+
+const (
+	finalMembershipPolls     = 5
+	finalMembershipPollDelay = 2 * time.Second
+)
+
+type finalMembershipState struct {
+	Personal  *chatgpt.AccountInfo
+	Workspace *chatgpt.AccountInfo
+}
+
+func (s *finalMembershipState) plusActive() bool {
+	return plusMembershipActive(s.Personal)
+}
+
+func (s *finalMembershipState) teamActive() bool {
+	return teamMembershipActive(s.Workspace)
+}
+
+func (s *finalMembershipState) satisfied(taskType string) bool {
+	switch taskType {
+	case TaskTypePlus:
+		return s.plusActive()
+	case TaskTypeTeam:
+		return s.teamActive()
+	case TaskTypeBoth:
+		return s.plusActive() && s.teamActive()
+	default:
+		return false
+	}
+}
+
+func (s *finalMembershipState) resultPlanForTask(taskType string) string {
+	if s != nil && s.satisfied(taskType) {
+		return taskType
+	}
+	return s.actualPlan()
+}
+
+func (s *finalMembershipState) actualPlan() string {
+	switch {
+	case s == nil:
+		return "unknown"
+	case s.plusActive() && s.teamActive():
+		return TaskTypeBoth
+	case s.teamActive():
+		return "team"
+	case s.plusActive():
+		return "plus"
+	case s.Personal != nil && s.Personal.PlanType != "":
+		return s.Personal.PlanType
+	case s.Workspace != nil && s.Workspace.PlanType != "":
+		return s.Workspace.PlanType
+	default:
+		return "unknown"
+	}
+}
+
+func (s *finalMembershipState) describe() string {
+	return fmt.Sprintf("personal=%s workspace=%s", describeMembership(s.Personal), describeMembership(s.Workspace))
+}
+
+func plusMembershipActive(info *chatgpt.AccountInfo) bool {
+	return info != nil &&
+		info.Structure == "personal" &&
+		info.PlanType == "plus" &&
+		info.HasActiveSubscription &&
+		info.SubscriptionID != ""
+}
+
+func teamMembershipActive(info *chatgpt.AccountInfo) bool {
+	return info != nil &&
+		info.Structure == "workspace" &&
+		info.PlanType == "team"
+}
+
+func describeMembership(info *chatgpt.AccountInfo) string {
+	if info == nil {
+		return "none"
+	}
+	return fmt.Sprintf("%s/%s(active=%v sub=%t id=%s)",
+		info.Structure, info.PlanType, info.HasActiveSubscription, info.SubscriptionID != "", info.AccountID)
+}
+
+func selectPersonalMembership(accounts []*chatgpt.AccountInfo) *chatgpt.AccountInfo {
+	for _, acct := range accounts {
+		if acct.Structure == "personal" {
+			return acct
+		}
+	}
+	return nil
+}
+
+func selectWorkspaceMembership(accounts []*chatgpt.AccountInfo, preferredID string) *chatgpt.AccountInfo {
+	if preferredID != "" {
+		for _, acct := range accounts {
+			if acct.Structure == "workspace" && acct.AccountID == preferredID {
+				return acct
+			}
+		}
+	}
+	for _, acct := range accounts {
+		if acct.Structure == "workspace" && acct.PlanType == "team" {
+			return acct
+		}
+	}
+	for _, acct := range accounts {
+		if acct.Structure == "workspace" {
+			return acct
+		}
+	}
+	return nil
+}
+
+func (r *TaskRunner) verifyTaskMembership(
+	ctx context.Context,
+	taskType string,
+	session *chatgpt.Session,
+	teamAccountID string,
+	statusFn chatgpt.StatusFunc,
+) (*finalMembershipState, error) {
+	var lastState *finalMembershipState
+	var lastErr error
+
+	for attempt := 1; attempt <= finalMembershipPolls; attempt++ {
+		if err := refreshVerificationTokens(ctx, session); err != nil && statusFn != nil {
+			statusFn("  -> Final token refresh %d/%d failed: %v", attempt, finalMembershipPolls, err)
+		}
+
+		accounts, err := chatgpt.CheckAccountFull(session.Client, session.AccessToken, session.DeviceID)
+		if err != nil {
+			lastErr = fmt.Errorf("accounts/check failed: %w", err)
+		} else {
+			lastState = &finalMembershipState{
+				Personal:  selectPersonalMembership(accounts),
+				Workspace: selectWorkspaceMembership(accounts, teamAccountID),
+			}
+		}
+
+		if (taskType == TaskTypeTeam || taskType == TaskTypeBoth) &&
+			(lastState == nil || !lastState.teamActive()) &&
+			teamAccountID != "" {
+			workspaceToken, wsErr := chatgpt.GetWorkspaceAccessToken(session.Client, teamAccountID)
+			if wsErr == nil {
+				if wsAccounts, wsCheckErr := chatgpt.CheckAccountFull(session.Client, workspaceToken, session.DeviceID); wsCheckErr == nil {
+					if lastState == nil {
+						lastState = &finalMembershipState{}
+					}
+					lastState.Workspace = selectWorkspaceMembership(wsAccounts, teamAccountID)
+				} else if lastErr == nil {
+					lastErr = fmt.Errorf("workspace accounts/check failed: %w", wsCheckErr)
+				}
+			} else if lastErr == nil {
+				lastErr = fmt.Errorf("workspace token refresh failed: %w", wsErr)
+			}
+		}
+
+		if lastState != nil && statusFn != nil {
+			statusFn("  -> Final membership %d/%d: %s", attempt, finalMembershipPolls, lastState.describe())
+		}
+		if lastState != nil && lastState.satisfied(taskType) {
+			return lastState, nil
+		}
+		if lastState != nil {
+			lastErr = fmt.Errorf("membership mismatch for %s: %s", taskType, lastState.describe())
+		}
+
+		if attempt < finalMembershipPolls {
+			if err := sleepWithContext(ctx, finalMembershipPollDelay); err != nil {
+				return lastState, err
+			}
+		}
+	}
+
+	if lastState != nil {
+		return lastState, fmt.Errorf("final membership mismatch for %s: %s", taskType, lastState.describe())
+	}
+	if lastErr != nil {
+		return nil, fmt.Errorf("final membership check failed for %s: %w", taskType, lastErr)
+	}
+	return nil, fmt.Errorf("final membership check failed for %s", taskType)
+}
+
+func refreshVerificationTokens(ctx context.Context, session *chatgpt.Session) error {
+	if err := session.RefreshSession(); err == nil {
+		return nil
+	} else {
+		tokens, tokenErr := auth.ObtainCodexTokens(ctx, session.Client, session.DeviceID, "")
+		if tokenErr != nil {
+			return fmt.Errorf("session refresh failed: %v; codex refresh failed: %w", err, tokenErr)
+		}
+		if tokens.AccessToken == "" {
+			return fmt.Errorf("codex refresh returned empty access token")
+		}
+		session.AccessToken = tokens.AccessToken
+		if tokens.RefreshToken != "" {
+			session.RefreshToken = tokens.RefreshToken
+		}
+		if tokens.IDToken != "" {
+			session.IDToken = tokens.IDToken
+		}
+		if tokens.ChatGPTAccountID != "" {
+			session.AccountID = tokens.ChatGPTAccountID
+		}
+		return nil
+	}
+}
+
+func sleepWithContext(ctx context.Context, delay time.Duration) error {
+	timer := time.NewTimer(delay)
+	defer timer.Stop()
+
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-timer.C:
+		return nil
+	}
+}