1  import { LRUCache } from 'lru-cache'
  2  import { logError } from './log.js'
  3  import { jsonStringify } from './slowOperations.js'
  4  
  5  type CacheEntry<T> = {
  6    value: T
  7    timestamp: number
  8    refreshing: boolean
  9  }
 10  
 11  type MemoizedFunction<Args extends unknown[], Result> = {
 12    (...args: Args): Result
 13    cache: {
 14      clear: () => void
 15    }
 16  }
 17  
 18  type LRUMemoizedFunction<Args extends unknown[], Result> = {
 19    (...args: Args): Result
 20    cache: {
 21      clear: () => void
 22      size: () => number
 23      delete: (key: string) => boolean
 24      get: (key: string) => Result | undefined
 25      has: (key: string) => boolean
 26    }
 27  }
 28  
 29  /**
 30   * Creates a memoized function that returns cached values while refreshing in parallel.
 31   * This implements a write-through cache pattern:
 32   * - If cache is fresh, return immediately
 33   * - If cache is stale, return the stale value but refresh it in the background
 34   * - If no cache exists, block and compute the value
 35   *
 36   * @param f The function to memoize
 37   * @param cacheLifetimeMs The lifetime of cached values in milliseconds
 38   * @returns A memoized version of the function
 39   */
 40  export function memoizeWithTTL<Args extends unknown[], Result>(
 41    f: (...args: Args) => Result,
 42    cacheLifetimeMs: number = 5 * 60 * 1000, // Default 5 minutes
 43  ): MemoizedFunction<Args, Result> {
 44    const cache = new Map<string, CacheEntry<Result>>()
 45  
 46    const memoized = (...args: Args): Result => {
 47      const key = jsonStringify(args)
 48      const cached = cache.get(key)
 49      const now = Date.now()
 50  
 51      // Populate cache
 52      if (!cached) {
 53        const value = f(...args)
 54        cache.set(key, {
 55          value,
 56          timestamp: now,
 57          refreshing: false,
 58        })
 59        return value
 60      }
 61  
 62      // If we have a stale cache entry and it's not already refreshing
 63      if (
 64        cached &&
 65        now - cached.timestamp > cacheLifetimeMs &&
 66        !cached.refreshing
 67      ) {
 68        // Mark as refreshing to prevent multiple parallel refreshes
 69        cached.refreshing = true
 70  
 71        // Schedule async refresh (non-blocking). Both .then and .catch are
 72        // identity-guarded: a concurrent cache.clear() + cold-miss stores a
 73        // newer entry while this microtask is queued. .then overwriting with
 74        // the stale refresh's result is worse than .catch deleting (persists
 75        // wrong data for full TTL vs. self-correcting on next call).
 76        Promise.resolve()
 77          .then(() => {
 78            const newValue = f(...args)
 79            if (cache.get(key) === cached) {
 80              cache.set(key, {
 81                value: newValue,
 82                timestamp: Date.now(),
 83                refreshing: false,
 84              })
 85            }
 86          })
 87          .catch(e => {
 88            logError(e)
 89            if (cache.get(key) === cached) {
 90              cache.delete(key)
 91            }
 92          })
 93  
 94        // Return the stale value immediately
 95        return cached.value
 96      }
 97  
 98      return cache.get(key)!.value
 99    }
100  
101    // Add cache clear method
102    memoized.cache = {
103      clear: () => cache.clear(),
104    }
105  
106    return memoized
107  }
108  
109  /**
110   * Creates a memoized async function that returns cached values while refreshing in parallel.
111   * This implements a write-through cache pattern for async functions:
112   * - If cache is fresh, return immediately
113   * - If cache is stale, return the stale value but refresh it in the background
114   * - If no cache exists, block and compute the value
115   *
116   * @param f The async function to memoize
117   * @param cacheLifetimeMs The lifetime of cached values in milliseconds
118   * @returns A memoized version of the async function
119   */
120  export function memoizeWithTTLAsync<Args extends unknown[], Result>(
121    f: (...args: Args) => Promise<Result>,
122    cacheLifetimeMs: number = 5 * 60 * 1000, // Default 5 minutes
123  ): ((...args: Args) => Promise<Result>) & { cache: { clear: () => void } } {
124    const cache = new Map<string, CacheEntry<Result>>()
125    // In-flight cold-miss dedup. The old memoizeWithTTL (sync) accidentally
126    // provided this: it stored the Promise synchronously before the first
127    // await, so concurrent callers shared one f() invocation. This async
128    // variant awaits before cache.set, so concurrent cold-miss callers would
129    // each invoke f() independently without this map. For
130    // refreshAndGetAwsCredentials that means N concurrent `aws sso login`
131    // spawns. Same pattern as pending401Handlers in auth.ts:1171.
132    const inFlight = new Map<string, Promise<Result>>()
133  
134    const memoized = async (...args: Args): Promise<Result> => {
135      const key = jsonStringify(args)
136      const cached = cache.get(key)
137      const now = Date.now()
138  
139      // Populate cache - if this throws, nothing gets cached
140      if (!cached) {
141        const pending = inFlight.get(key)
142        if (pending) return pending
143        const promise = f(...args)
144        inFlight.set(key, promise)
145        try {
146          const result = await promise
147          // Identity-guard: cache.clear() during the await should discard this
148          // result (clear intent is to invalidate). If we're still in-flight,
149          // store it. clear() wipes inFlight too, so this check catches that.
150          if (inFlight.get(key) === promise) {
151            cache.set(key, {
152              value: result,
153              timestamp: now,
154              refreshing: false,
155            })
156          }
157          return result
158        } finally {
159          if (inFlight.get(key) === promise) {
160            inFlight.delete(key)
161          }
162        }
163      }
164  
165      // If we have a stale cache entry and it's not already refreshing
166      if (
167        cached &&
168        now - cached.timestamp > cacheLifetimeMs &&
169        !cached.refreshing
170      ) {
171        // Mark as refreshing to prevent multiple parallel refreshes
172        cached.refreshing = true
173  
174        // Schedule async refresh (non-blocking). Both .then and .catch are
175        // identity-guarded against a concurrent cache.clear() + cold-miss
176        // storing a newer entry while this refresh is in flight. .then
177        // overwriting with the stale refresh's result is worse than .catch
178        // deleting - wrong data persists for full TTL (e.g. credentials from
179        // the old awsAuthRefresh command after a settings change).
180        const staleEntry = cached
181        f(...args)
182          .then(newValue => {
183            if (cache.get(key) === staleEntry) {
184              cache.set(key, {
185                value: newValue,
186                timestamp: Date.now(),
187                refreshing: false,
188              })
189            }
190          })
191          .catch(e => {
192            logError(e)
193            if (cache.get(key) === staleEntry) {
194              cache.delete(key)
195            }
196          })
197  
198        // Return the stale value immediately
199        return cached.value
200      }
201  
202      return cache.get(key)!.value
203    }
204  
205    // Add cache clear method. Also clear inFlight: clear() during a cold-miss
206    // await should not let the stale in-flight promise be returned to the next
207    // caller (defeats the purpose of clear). The try/finally above
208    // identity-guards inFlight.delete so the stale promise doesn't delete a
209    // fresh one if clear+cold-miss happens before the finally fires.
210    memoized.cache = {
211      clear: () => {
212        cache.clear()
213        inFlight.clear()
214      },
215    }
216  
217    return memoized as ((...args: Args) => Promise<Result>) & {
218      cache: { clear: () => void }
219    }
220  }
221  
222  /**
223   * Creates a memoized function with LRU (Least Recently Used) eviction policy.
224   * This prevents unbounded memory growth by evicting the least recently used entries
225   * when the cache reaches its maximum size.
226   *
227   * Note: Cache size for memoized message processing functions
228   * Chosen to prevent unbounded memory growth (was 300MB+ with lodash memoize)
229   * while maintaining good cache hit rates for typical conversations.
230   *
231   * @param f The function to memoize
232   * @returns A memoized version of the function with cache management methods
233   */
234  export function memoizeWithLRU<
235    Args extends unknown[],
236    Result extends NonNullable<unknown>,
237  >(
238    f: (...args: Args) => Result,
239    cacheFn: (...args: Args) => string,
240    maxCacheSize: number = 100,
241  ): LRUMemoizedFunction<Args, Result> {
242    const cache = new LRUCache<string, Result>({
243      max: maxCacheSize,
244    })
245  
246    const memoized = (...args: Args): Result => {
247      const key = cacheFn(...args)
248      const cached = cache.get(key)
249      if (cached !== undefined) {
250        return cached
251      }
252  
253      const result = f(...args)
254      cache.set(key, result)
255      return result
256    }
257  
258    // Add cache management methods
259    memoized.cache = {
260      clear: () => cache.clear(),
261      size: () => cache.size,
262      delete: (key: string) => cache.delete(key),
263      // peek() avoids updating recency — we only want to observe, not promote
264      get: (key: string) => cache.peek(key),
265      has: (key: string) => cache.has(key),
266    }
267  
268    return memoized
269  }