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:
parent
42361e6505
commit
93f672df20
|
|
@ -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
|
||||
}
|
||||
|
||||
|
|
|
|||
Reference in New Issue