Vectrum
Sign In
232 lines · 7.7 KB · Go
package claudecode

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"sync"
	"time"
)

// ClientID identifies Claude Code to the token endpoint.
const ClientID = "9d1c250a-e61b-44d9-88ed-5944d1962f5e"

// tokenURLEnv overrides the token endpoint. It exists so tests can point at an
// httptest server; nothing in normal operation sets it.
const tokenURLEnv = "CLAUDECODE_TOKEN_URL"

const defaultTokenURL = "https://console.anthropic.com/v1/oauth/token"

// tokenURL is the endpoint a refresh is sent to.
func tokenURL() string {
	if u := os.Getenv(tokenURLEnv); u != "" {
		return u
	}
	return defaultTokenURL
}

// ErrExpired means the credential has lapsed and cannot be renewed here —
// either the refresh token is gone or itself expired. Only logging in again
// fixes it, and the tool that owns the credential is the one to do it.
var ErrExpired = errors.New("claudecode: your Claude Code login has expired, run `claude` to refresh it")

// Refresh exchanges the refresh token for a new access token. It returns new
// credentials and writes nothing — persisting is a separate decision, made by
// Source, because whether it is safe depends on whether the server rotated.
//
// Every failure says to run `claude`, including a 429. On these OAuth surfaces
// a 429 is not reliable evidence of a rate limit: the Messages API answers a
// request shape it declines the same way, and every probe of this endpoint has
// returned one without ever completing an exchange. Advising someone to wait
// would be advice that may never come good.
func Refresh(ctx context.Context, h *http.Client, c *Credentials) (*Credentials, error) {
	body, err := json.Marshal(map[string]string{
		"grant_type":    "refresh_token",
		"refresh_token": c.RefreshToken,
		"client_id":     ClientID,
	})
	if err != nil {
		return nil, err
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, tokenURL(), bytes.NewReader(body))
	if err != nil {
		return nil, err
	}
	req.Header.Set("content-type", "application/json")

	resp, err := h.Do(req)
	if err != nil {
		return nil, fmt.Errorf("claudecode: refreshing: %w", err)
	}
	defer resp.Body.Close()
	raw, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("claudecode: reading refresh response: %w", err)
	}
	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("claudecode: refresh refused (%d), run `claude` to log in again: %s",
			resp.StatusCode, bytes.TrimSpace(raw))
	}

	var out struct {
		AccessToken  string `json:"access_token"`
		RefreshToken string `json:"refresh_token"`
		ExpiresIn    int64  `json:"expires_in"`
	}
	if err := json.Unmarshal(raw, &out); err != nil {
		return nil, fmt.Errorf("claudecode: decoding refresh response: %w", err)
	}
	if out.AccessToken == "" {
		return nil, errors.New("claudecode: refresh returned no access token")
	}

	next := *c
	next.AccessToken = out.AccessToken
	// An omitted refresh_token means the old one is still good. Blanking it
	// would throw away the only thing that can refresh again.
	if out.RefreshToken != "" {
		next.RefreshToken = out.RefreshToken
	}
	next.ExpiresAt = time.Now().Add(time.Duration(out.ExpiresIn) * time.Second)
	return &next, nil
}

// Save writes credentials back to the store.
//
// This rewrites a file Claude Code owns, so it is a read-modify-write: the
// document is parsed, only the three fields we came to change are replaced,
// and everything else — keys this package does not model, keys it has never
// heard of — goes back exactly as it was found. Numbers are decoded as literals
// rather than float64s, or the millisecond timestamps would come back as
// 1.78966826946e+12 and Claude Code would be reading a file we corrupted.
//
// The write is atomic: a temp file in the same directory, then a rename, so a
// crash midway leaves the original rather than a truncated credential.
func Save(home string, c *Credentials) error {
	path := filepath.Join(home, Filename)
	raw, err := os.ReadFile(path)
	if err != nil {
		return fmt.Errorf("claudecode: reading %s: %w", path, err)
	}

	var doc map[string]json.RawMessage
	dec := json.NewDecoder(bytes.NewReader(raw))
	dec.UseNumber()
	if err := dec.Decode(&doc); err != nil {
		return fmt.Errorf("claudecode: parsing %s: %w", path, err)
	}
	blob, ok := doc["claudeAiOauth"]
	if !ok {
		return fmt.Errorf("claudecode: %s has no claudeAiOauth object", path)
	}
	var oauth map[string]any
	dec = json.NewDecoder(bytes.NewReader(blob))
	dec.UseNumber()
	if err := dec.Decode(&oauth); err != nil {
		return fmt.Errorf("claudecode: parsing claudeAiOauth in %s: %w", path, err)
	}

	oauth["accessToken"] = c.AccessToken
	oauth["refreshToken"] = c.RefreshToken
	oauth["expiresAt"] = json.Number(strconv.FormatInt(c.ExpiresAt.UnixMilli(), 10))

	if blob, err = json.Marshal(oauth); err != nil {
		return fmt.Errorf("claudecode: encoding claudeAiOauth: %w", err)
	}
	doc["claudeAiOauth"] = blob
	out, err := json.Marshal(doc)
	if err != nil {
		return fmt.Errorf("claudecode: encoding %s: %w", path, err)
	}

	tmp, err := os.CreateTemp(filepath.Dir(path), ".credentials-*.tmp")
	if err != nil {
		return fmt.Errorf("claudecode: writing %s: %w", path, err)
	}
	defer os.Remove(tmp.Name())
	if err := tmp.Chmod(0o600); err != nil {
		tmp.Close()
		return fmt.Errorf("claudecode: securing %s: %w", tmp.Name(), err)
	}
	if _, err := tmp.Write(out); err != nil {
		tmp.Close()
		return fmt.Errorf("claudecode: writing %s: %w", tmp.Name(), err)
	}
	if err := tmp.Close(); err != nil {
		return fmt.Errorf("claudecode: closing %s: %w", tmp.Name(), err)
	}
	if err := os.Rename(tmp.Name(), path); err != nil {
		return fmt.Errorf("claudecode: replacing %s: %w", path, err)
	}
	return nil
}

// Source hands out an access token, refreshing it when it lapses. It satisfies
// the shape llm.WithTokenSource wants.
type Source struct {
	mu    sync.Mutex
	creds *Credentials
	http  *http.Client
	// home is where a rotated credential is written back. Empty disables
	// write-back entirely.
	home string
}

// NewSource wraps credentials in a token source. h is the client refreshes are
// sent with; home is the directory holding the credential store, or "" to
// refresh without ever writing.
func NewSource(c *Credentials, h *http.Client, home string) *Source {
	if h == nil {
		h = http.DefaultClient
	}
	return &Source{creds: c, http: h, home: home}
}

// Token returns a valid access token, refreshing first if the current one has
// lapsed.
//
// The refreshed credential is kept in memory. It is written back only when the
// server rotated the refresh token, and that is the point: if rotation
// happened, the token still on disk is already dead, so leaving it there is
// what breaks the user's Claude Code login. Writing is the safe act in that
// case and the risky one in every other, which is why it is conditional rather
// than routine.
//
// The lock means concurrent requests — streaming and buffered can be in flight
// together — produce one refresh rather than a thundering herd.
func (s *Source) Token(ctx context.Context) (string, error) {
	s.mu.Lock()
	defer s.mu.Unlock()

	now := time.Now()
	if !s.creds.Expired(now) {
		return s.creds.AccessToken, nil
	}
	if s.creds.RefreshToken == "" || s.creds.RefreshExpired(now) {
		return "", ErrExpired
	}

	next, err := Refresh(ctx, s.http, s.creds)
	if err != nil {
		return "", err
	}
	rotated := next.RefreshToken != s.creds.RefreshToken
	s.creds = next

	if rotated && s.home != "" {
		// Best effort. A store we cannot write is a worse next session, not a
		// failed this one — we are holding a token that works, and refusing to
		// use it would be a strictly worse outcome than the stale file.
		if err := Save(s.home, s.creds); err != nil {
			return s.creds.AccessToken, nil
		}
	}
	return s.creds.AccessToken, nil
}