Vectrum
Sign In
CodeyPublic
Like0 likesFork0 forks

llm: resolve the credential per request

A frozen authToken cannot outlive its own expiry. WithTokenSource makes
the bearer a function; WithAuthToken is now its constant case.

Resolution happens before the retry loop because a credential error is
not an *APIError, and the loop only returns early for those — resolving
inside attempt would retry a dead credential three times with backoff.
Both request paths now share one apply instead of repeating the headers.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Browse Files
Jjonicommitted 1 parent d33f0becommit28d36f5

Showing 3 changed files with 186 additions and 21 deletions.

+60−13internal/llm/client.goView File
@@ -35,8 +35,9 @@ type Client struct {
3535 RetryDelay time.Duration
3636 HTTP *http.Client
3737
38- apiKey string
39- authToken string
38+ apiKey string
39+ authToken string
40+ tokenSource func(context.Context) (string, error)
4041 }
4142
4243 // Option configures a Client.
@@ -80,6 +81,17 @@ func WithAuthToken(t string) Option {
8081 }
8182 }
8283 }
84+
85+ // WithTokenSource sets a bearer credential resolved per request, so a token
86+ // that expires mid-session can be replaced without rebuilding the client. A
87+ // nil function is ignored.
88+ func WithTokenSource(f func(context.Context) (string, error)) Option {
89+ return func(c *Client) {
90+ if f != nil {
91+ c.tokenSource, c.apiKey, c.authToken = f, "", ""
92+ }
93+ }
94+ }
8395 func WithHTTPClient(h *http.Client) Option { return func(c *Client) { c.HTTP = h } }
8496
8597 // New builds a client, reading credentials from the environment unless an
@@ -120,7 +132,9 @@ func normalizeBaseURL(u string) string {
120132 func (c *Client) IsDefaultEndpoint() bool { return c.BaseURL == DefaultBaseURL }
121133
122134 // HasCredentials reports whether the client can authenticate.
123- func (c *Client) HasCredentials() bool { return c.apiKey != "" || c.authToken != "" }
135+ func (c *Client) HasCredentials() bool {
136+ return c.apiKey != "" || c.authToken != "" || c.tokenSource != nil
137+ }
124138
125139 // APIError is a non-2xx response.
126140 type APIError struct {
@@ -145,6 +159,42 @@ func (e *APIError) Retryable() bool {
145159 e.StatusCode >= 500
146160 }
147161
162+ // auth is one request's resolved credential. Resolving returns a value rather
163+ // than mutating the client so concurrent requests cannot interleave.
164+ type auth struct {
165+ apiKey string
166+ bearer string
167+ }
168+
169+ // apply writes the credential headers. An OAuth bearer also needs the beta
170+ // flag; an API key does not.
171+ func (a auth) apply(h http.Header) {
172+ if a.apiKey != "" {
173+ h.Set("x-api-key", a.apiKey)
174+ return
175+ }
176+ h.Set("authorization", "Bearer "+a.bearer)
177+ h.Set("anthropic-beta", oauthBeta)
178+ }
179+
180+ // resolveAuth produces the credential for one request. It runs before the
181+ // retry loop: a credential that will not resolve is permanent, and the loop
182+ // only short-circuits on a non-retryable *APIError, so resolving inside
183+ // attempt would retry a dead credential three times with backoff.
184+ func (c *Client) resolveAuth(ctx context.Context) (auth, error) {
185+ if c.apiKey != "" {
186+ return auth{apiKey: c.apiKey}, nil
187+ }
188+ if c.tokenSource != nil {
189+ tok, err := c.tokenSource(ctx)
190+ if err != nil {
191+ return auth{}, err
192+ }
193+ return auth{bearer: tok}, nil
194+ }
195+ return auth{bearer: c.authToken}, nil
196+ }
197+
148198 // CreateMessage sends one request, filling in the client's defaults and
149199 // retrying the errors worth retrying.
150200 func (c *Client) CreateMessage(ctx context.Context, req Request) (*Response, error) {
@@ -153,6 +203,10 @@ func (c *Client) CreateMessage(ctx context.Context, req Request) (*Response, err
153203 if !c.HasCredentials() {
154204 return nil, fmt.Errorf("llm: no credentials: set ANTHROPIC_API_KEY, or run `ant auth login`")
155205 }
206+ a, err := c.resolveAuth(ctx)
207+ if err != nil {
208+ return nil, err
209+ }
156210 if req.Model == "" {
157211 req.Model = c.Model
158212 }
@@ -171,7 +225,7 @@ func (c *Client) CreateMessage(ctx context.Context, req Request) (*Response, err
171225 return nil, err
172226 }
173227 }
174- resp, err := c.attempt(ctx, body)
228+ resp, err := c.attempt(ctx, body, a)
175229 if err == nil {
176230 return resp, nil
177231 }
@@ -187,21 +241,14 @@ func (c *Client) CreateMessage(ctx context.Context, req Request) (*Response, err
187241 return nil, lastErr
188242 }
189243
190- func (c *Client) attempt(ctx context.Context, body []byte) (*Response, error) {
244+ func (c *Client) attempt(ctx context.Context, body []byte, a auth) (*Response, error) {
191245 httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, c.BaseURL+"/v1/messages", bytes.NewReader(body))
192246 if err != nil {
193247 return nil, err
194248 }
195249 httpReq.Header.Set("content-type", "application/json")
196250 httpReq.Header.Set("anthropic-version", apiVersion)
197- switch {
198- case c.apiKey != "":
199- httpReq.Header.Set("x-api-key", c.apiKey)
200- default:
201- // OAuth tokens go on Authorization, not x-api-key, and need the beta flag.
202- httpReq.Header.Set("authorization", "Bearer "+c.authToken)
203- httpReq.Header.Set("anthropic-beta", oauthBeta)
204- }
251+ a.apply(httpReq.Header)
205252
206253 httpResp, err := c.HTTP.Do(httpReq)
207254 if err != nil {
+119−0internal/llm/client_test.goView File
@@ -3,6 +3,8 @@ package llm
33 import (
44 "context"
55 "encoding/json"
6+ "errors"
7+ "fmt"
68 "io"
79 "net/http"
810 "net/http/httptest"
@@ -284,3 +286,120 @@ func TestCustomBaseURLIsActuallyUsed(t *testing.T) {
284286 t.Errorf("path = %q, want /v1/messages", gotPath)
285287 }
286288 }
289+
290+ // A token source is consulted per request, so a credential that rotates
291+ // mid-session reaches the wire without rebuilding the client.
292+ func TestTokenSourceIsCalledPerRequest(t *testing.T) {
293+ var got []string
294+ srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
295+ got = append(got, r.Header.Get("authorization"))
296+ io.WriteString(w, okResponse)
297+ }))
298+ defer srv.Close()
299+
300+ n := 0
301+ c := New(WithBaseURL(srv.URL), WithTokenSource(func(context.Context) (string, error) {
302+ n++
303+ return fmt.Sprintf("tok-%d", n), nil
304+ }))
305+ for i := 0; i < 2; i++ {
306+ if _, err := c.CreateMessage(context.Background(), Request{}); err != nil {
307+ t.Fatalf("CreateMessage: %v", err)
308+ }
309+ }
310+ if len(got) != 2 || got[0] != "Bearer tok-1" || got[1] != "Bearer tok-2" {
311+ t.Errorf("authorization headers = %v, want Bearer tok-1 then tok-2", got)
312+ }
313+ }
314+
315+ func TestTokenSourceGetsTheOAuthBetaHeader(t *testing.T) {
316+ var got http.Header
317+ srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
318+ got = r.Header.Clone()
319+ io.WriteString(w, okResponse)
320+ }))
321+ defer srv.Close()
322+
323+ c := New(WithBaseURL(srv.URL), WithTokenSource(func(context.Context) (string, error) {
324+ return "tok", nil
325+ }))
326+ if _, err := c.CreateMessage(context.Background(), Request{}); err != nil {
327+ t.Fatalf("CreateMessage: %v", err)
328+ }
329+ if got.Get("anthropic-beta") != oauthBeta {
330+ t.Errorf("anthropic-beta = %q, want %s", got.Get("anthropic-beta"), oauthBeta)
331+ }
332+ if got.Get("x-api-key") != "" {
333+ t.Errorf("x-api-key should be unset, got %q", got.Get("x-api-key"))
334+ }
335+ }
336+
337+ // A credential that cannot be resolved is permanent. Retrying it three times
338+ // with backoff just delays the error, so resolution happens before the loop.
339+ func TestTokenSourceErrorIsNotRetried(t *testing.T) {
340+ var hits int
341+ srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
342+ hits++
343+ io.WriteString(w, okResponse)
344+ }))
345+ defer srv.Close()
346+
347+ var calls int
348+ c := New(WithBaseURL(srv.URL), WithTokenSource(func(context.Context) (string, error) {
349+ calls++
350+ return "", errors.New("login expired")
351+ }))
352+ c.RetryDelay = time.Millisecond
353+ _, err := c.CreateMessage(context.Background(), Request{})
354+ if err == nil || !strings.Contains(err.Error(), "login expired") {
355+ t.Fatalf("err = %v, want the source's error", err)
356+ }
357+ if calls != 1 {
358+ t.Errorf("token source called %d times, want 1", calls)
359+ }
360+ if hits != 0 {
361+ t.Errorf("server hit %d times, want 0 — no request should have been sent", hits)
362+ }
363+ }
364+
365+ func TestHasCredentialsCountsATokenSource(t *testing.T) {
366+ c := New(WithBaseURL("http://127.0.0.1:1"))
367+ c.apiKey, c.authToken = "", ""
368+ if c.HasCredentials() {
369+ t.Fatal("no credentials expected before the source is set")
370+ }
371+ c.tokenSource = func(context.Context) (string, error) { return "tok", nil }
372+ if !c.HasCredentials() {
373+ t.Error("a token source is a credential")
374+ }
375+ }
376+
377+ // The streaming path must authenticate identically to the buffered one.
378+ func TestStreamUsesTheTokenSource(t *testing.T) {
379+ streamOK := sse(
380+ `{"type":"message_start","message":{"id":"msg_1","model":"claude-opus-5","usage":{"input_tokens":100,"cache_read_input_tokens":40}}}`,
381+ `{"type":"message_delta","delta":{"stop_reason":"tool_use"},"usage":{"output_tokens":55}}`,
382+ `{"type":"message_stop"}`,
383+ )
384+
385+ var got http.Header
386+ srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
387+ got = r.Header.Clone()
388+ w.Header().Set("content-type", "text/event-stream")
389+ io.WriteString(w, streamOK)
390+ }))
391+ defer srv.Close()
392+
393+ c := New(WithBaseURL(srv.URL), WithTokenSource(func(context.Context) (string, error) {
394+ return "stream-tok", nil
395+ }))
396+ if _, err := c.StreamMessage(context.Background(), Request{}, func(StreamEvent) {}); err != nil {
397+ t.Fatalf("StreamMessage: %v", err)
398+ }
399+ if got.Get("authorization") != "Bearer stream-tok" {
400+ t.Errorf("authorization = %q", got.Get("authorization"))
401+ }
402+ if got.Get("anthropic-beta") != oauthBeta {
403+ t.Errorf("anthropic-beta = %q, want %s", got.Get("anthropic-beta"), oauthBeta)
404+ }
405+ }
+7−8internal/llm/stream.goView File
@@ -37,6 +37,10 @@ func (c *Client) StreamMessage(ctx context.Context, req Request, onEvent func(St
3737 if !c.HasCredentials() {
3838 return nil, fmt.Errorf("llm: no credentials: set ANTHROPIC_API_KEY, or run `ant auth login`")
3939 }
40+ a, err := c.resolveAuth(ctx)
41+ if err != nil {
42+ return nil, err
43+ }
4044 if req.Model == "" {
4145 req.Model = c.Model
4246 }
@@ -57,7 +61,7 @@ func (c *Client) StreamMessage(ctx context.Context, req Request, onEvent func(St
5761 return nil, err
5862 }
5963 }
60- resp, err := c.streamAttempt(ctx, body, onEvent)
64+ resp, err := c.streamAttempt(ctx, body, a, onEvent)
6165 if err == nil {
6266 return resp, nil
6367 }
@@ -78,7 +82,7 @@ func (c *Client) StreamMessage(ctx context.Context, req Request, onEvent func(St
7882 return nil, lastErr
7983 }
8084
81- func (c *Client) streamAttempt(ctx context.Context, body []byte, onEvent func(StreamEvent)) (*Response, error) {
85+ func (c *Client) streamAttempt(ctx context.Context, body []byte, a auth, onEvent func(StreamEvent)) (*Response, error) {
8286 httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, c.BaseURL+"/v1/messages", bytes.NewReader(body))
8387 if err != nil {
8488 return nil, err
@@ -86,12 +90,7 @@ func (c *Client) streamAttempt(ctx context.Context, body []byte, onEvent func(St
8690 httpReq.Header.Set("content-type", "application/json")
8791 httpReq.Header.Set("accept", "text/event-stream")
8892 httpReq.Header.Set("anthropic-version", apiVersion)
89- if c.apiKey != "" {
90- httpReq.Header.Set("x-api-key", c.apiKey)
91- } else {
92- httpReq.Header.Set("authorization", "Bearer "+c.authToken)
93- httpReq.Header.Set("anthropic-beta", oauthBeta)
94- }
93+ a.apply(httpReq.Header)
9594
9695 httpResp, err := c.HTTP.Do(httpReq)
9796 if err != nil {