kiro backend: cache Models() for 24h, cache sqlite tokens for 60s

Models() now caches at the backend level (shared across sessions) with
a 24h TTL, eliminating repeated expensive HTTP calls to AWS.

The sqliteKiroAuth.load() now caches token reads for 60s, avoiding
shelling out to sqlite3 on every API call (previously 3x per Models
fetch).
This commit is contained in:
ollie 2026-07-18 01:40:19 +02:00
parent 42361e6505
commit 93f672df20
1 changed files with 41 additions and 10 deletions

View File

@ -41,6 +41,11 @@ type CodeWhispererBackend struct {
authInitErr error
httpClient *http.Client
ctxLength int // cached
// Models cache at the backend level (shared across all sessions).
modelsMu sync.Mutex
modelsCache []string
modelsCacheAt time.Time
}
// NewCodeWhisperer returns a CodeWhisperer backend. See package doc for apiKey semantics.
@ -63,6 +68,14 @@ func (b *CodeWhispererBackend) DefaultModel() string { return "auto" }
func (b *CodeWhispererBackend) Model() string { return b.model }
func (b *CodeWhispererBackend) SetModel(m string) { b.model = m }
func (b *CodeWhispererBackend) Models(ctx context.Context) []string {
b.modelsMu.Lock()
if b.modelsCache != nil && time.Since(b.modelsCacheAt) < 24*time.Hour {
result := b.modelsCache
b.modelsMu.Unlock()
return result
}
b.modelsMu.Unlock()
resp := b.fetchModels(ctx)
if resp == nil {
return nil
@ -71,6 +84,11 @@ func (b *CodeWhispererBackend) Models(ctx context.Context) []string {
for i, m := range resp.Models {
ids[i] = m.ModelID
}
b.modelsMu.Lock()
b.modelsCache = ids
b.modelsCacheAt = time.Now()
b.modelsMu.Unlock()
return ids
}
@ -677,11 +695,12 @@ func (s *staticKiroAuth) CanRefresh() bool { return false }
// sqliteKiroAuth reads tokens from the Kiro CLI SQLite database and refreshes via OIDC.
type sqliteKiroAuth struct {
path string
state kiroSQLiteState
loaded bool
initMu sync.Mutex
mu sync.Mutex
path string
state kiroSQLiteState
loaded bool
initMu sync.Mutex
mu sync.Mutex
tokenReadAt time.Time // last time tokens were read from sqlite
// jitter is computed once at creation time so all checks within this
// process use the same refresh threshold. Different processes get
// different jitter values, spreading refresh attempts over time.
@ -871,17 +890,29 @@ func (s *sqliteKiroAuth) load(ctx context.Context) (kiroSQLiteState, error) {
}
s.initMu.Unlock()
// Always re-read tokens — they may have been refreshed by another process.
// Re-read tokens if stale (>60s). Within that window, use cached values
// to avoid shelling out to sqlite3 on every API call.
s.mu.Lock()
if s.state.AccessToken != "" && time.Since(s.tokenReadAt) < 60*time.Second {
state := s.state
s.mu.Unlock()
return state, nil
}
s.mu.Unlock()
access, refresh, expiresAt, err := kiroReadSQLiteTokens(ctx, s.path)
if err != nil {
return s.state, nil // fall back to cached state
}
state := s.state
state.AccessToken = access
state.ExpiresAt = expiresAt
s.mu.Lock()
s.state.AccessToken = access
s.state.ExpiresAt = expiresAt
if refresh != "" {
state.RefreshToken = refresh
s.state.RefreshToken = refresh
}
s.tokenReadAt = time.Now()
state := s.state
s.mu.Unlock()
return state, nil
}