orch-code 0.1.3 → 0.1.5

package/internal/auth/account_session.go

@@ -0,0 +1,154 @@
+package auth
+
+import (
+	"context"
+	"fmt"
+	"strings"
+	"time"
+)
+
+type AccountSession struct {
+	repoRoot  string
+	provider  string
+	currentID string
+	excluded  map[string]struct{}
+}
+
+func NewAccountSession(repoRoot, provider string) *AccountSession {
+	return &AccountSession{
+		repoRoot: strings.TrimSpace(repoRoot),
+		provider: strings.ToLower(strings.TrimSpace(provider)),
+		excluded: map[string]struct{}{},
+	}
+}
+
+func (s *AccountSession) ResolveToken(ctx context.Context) (string, error) {
+	_ = ctx
+	cred, err := s.pickCredential()
+	if err != nil {
+		return "", err
+	}
+	if cred == nil {
+		return "", fmt.Errorf("no active oauth credential available for provider %s", s.provider)
+	}
+	s.currentID = cred.ID
+	return strings.TrimSpace(cred.AccessToken), nil
+}
+
+func (s *AccountSession) Failover(ctx context.Context, cooldown time.Duration, reason string) (string, bool, error) {
+	_ = ctx
+	if strings.TrimSpace(s.currentID) == "" {
+		return "", false, nil
+	}
+	if err := mutateCredential(s.repoRoot, s.provider, s.currentID, func(cred *Credential) error {
+		if cooldown > 0 {
+			cred.CooldownUntil = time.Now().UTC().Add(cooldown)
+		}
+		cred.LastError = strings.TrimSpace(reason)
+		cred.UpdatedAt = time.Now().UTC()
+		return nil
+	}); err != nil {
+		return "", false, err
+	}
+	s.excluded[s.currentID] = struct{}{}
+
+	cred, err := s.pickCredential()
+	if err != nil {
+		return "", false, err
+	}
+	if cred == nil {
+		return "", false, nil
+	}
+	s.currentID = cred.ID
+	return strings.TrimSpace(cred.AccessToken), true, nil
+}
+
+func (s *AccountSession) MarkSuccess(ctx context.Context) {
+	_ = ctx
+	if strings.TrimSpace(s.currentID) == "" {
+		return
+	}
+	_ = mutateCredential(s.repoRoot, s.provider, s.currentID, func(cred *Credential) error {
+		cred.LastError = ""
+		cred.CooldownUntil = time.Time{}
+		cred.LastUsedAt = time.Now().UTC()
+		cred.UpdatedAt = time.Now().UTC()
+		return nil
+	})
+	s.excluded = map[string]struct{}{}
+}
+
+func (s *AccountSession) pickCredential() (*Credential, error) {
+	credentials, activeID, err := List(s.repoRoot, s.provider)
+	if err != nil {
+		return nil, err
+	}
+	if len(credentials) == 0 {
+		return nil, nil
+	}
+
+	now := time.Now().UTC()
+	ordered := orderCredentials(credentials, activeID, s.currentID)
+	for _, candidate := range ordered {
+		if candidate.Type != "oauth" {
+			continue
+		}
+		if _, skip := s.excluded[candidate.ID]; skip {
+			continue
+		}
+		if !candidate.CooldownUntil.IsZero() && candidate.CooldownUntil.After(now) {
+			continue
+		}
+		if candidate.ID != activeID {
+			if err := SetActive(s.repoRoot, s.provider, candidate.ID); err != nil {
+				return nil, err
+			}
+		}
+		resolved, err := resolveAccountCredentialByID(s.repoRoot, s.provider, candidate.ID)
+		if err == nil {
+			return resolved, nil
+		}
+		_ = mutateCredential(s.repoRoot, s.provider, candidate.ID, func(cred *Credential) error {
+			cred.LastError = strings.TrimSpace(err.Error())
+			cred.CooldownUntil = now.Add(5 * time.Minute)
+			cred.UpdatedAt = now
+			return nil
+		})
+		s.excluded[candidate.ID] = struct{}{}
+	}
+
+	return nil, nil
+}
+
+func orderCredentials(credentials []Credential, activeID, currentID string) []Credential {
+	ordered := make([]Credential, 0, len(credentials))
+	appendByID := func(id string) {
+		if strings.TrimSpace(id) == "" {
+			return
+		}
+		for _, cred := range credentials {
+			if cred.ID == id && !containsCredential(ordered, id) {
+				ordered = append(ordered, cred)
+				return
+			}
+		}
+	}
+	appendByID(currentID)
+	appendByID(activeID)
+	for _, cred := range credentials {
+		if containsCredential(ordered, cred.ID) {
+			continue
+		}
+		ordered = append(ordered, cred)
+	}
+	return ordered
+}
+
+func containsCredential(credentials []Credential, credentialID string) bool {
+	for _, cred := range credentials {
+		if cred.ID == credentialID {
+			return true
+		}
+	}
+	return false
+}

package/internal/auth/account_session_test.go

@@ -0,0 +1,71 @@
+package auth
+
+import (
+	"context"
+	"encoding/base64"
+	"fmt"
+	"testing"
+	"time"
+)
+
+func TestAccountSessionFailsOverToNextCredential(t *testing.T) {
+	repoRoot := t.TempDir()
+	if err := Set(repoRoot, "openai", Credential{Type: "oauth", AccessToken: testSessionAccountToken("acc-1"), RefreshToken: "refresh-1", AccountID: "acc-1", Email: "one@example.com"}); err != nil {
+		t.Fatalf("set first account: %v", err)
+	}
+	if err := Set(repoRoot, "openai", Credential{Type: "oauth", AccessToken: testSessionAccountToken("acc-2"), RefreshToken: "refresh-2", AccountID: "acc-2", Email: "two@example.com"}); err != nil {
+		t.Fatalf("set second account: %v", err)
+	}
+	if err := SetActive(repoRoot, "openai", "acc-1"); err != nil {
+		t.Fatalf("set active: %v", err)
+	}
+
+	session := NewAccountSession(repoRoot, "openai")
+	token, err := session.ResolveToken(context.Background())
+	if err != nil {
+		t.Fatalf("resolve token: %v", err)
+	}
+	if token != testSessionAccountToken("acc-1") {
+		t.Fatalf("expected first token, got %q", token)
+	}
+
+	nextToken, switched, err := session.Failover(context.Background(), time.Minute, "rate limited")
+	if err != nil {
+		t.Fatalf("failover: %v", err)
+	}
+	if !switched {
+		t.Fatalf("expected failover switch")
+	}
+	if nextToken != testSessionAccountToken("acc-2") {
+		t.Fatalf("expected second token, got %q", nextToken)
+	}
+
+	active, err := Get(repoRoot, "openai")
+	if err != nil {
+		t.Fatalf("get active credential: %v", err)
+	}
+	if active == nil || active.ID != "acc-2" {
+		t.Fatalf("expected acc-2 active after failover, got %#v", active)
+	}
+	credentials, _, err := List(repoRoot, "openai")
+	if err != nil {
+		t.Fatalf("list credentials: %v", err)
+	}
+	var first *Credential
+	for i := range credentials {
+		if credentials[i].ID == "acc-1" {
+			first = &credentials[i]
+			break
+		}
+	}
+	if first == nil || first.CooldownUntil.IsZero() {
+		t.Fatalf("expected first credential to have cooldown set, got %#v", first)
+	}
+}
+
+func testSessionAccountToken(accountID string) string {
+	header := base64.RawURLEncoding.EncodeToString([]byte(`{"alg":"none"}`))
+	payload := fmt.Sprintf(`{"https://api.openai.com/auth":{"chatgpt_account_id":"%s"}}`, accountID)
+	body := base64.RawURLEncoding.EncodeToString([]byte(payload))
+	return header + "." + body + ".sig"
+}