feat(ai): add version-aware prompt cache clearing (#3210)

* feat(ai): add version-aware prompt cache clearing

* chore(ai): add changeset for prompt cache clearing

* test(ai): cover prompt cache cleanup path

Author: Radu-Raicea

.changeset/poor-lobsters-tease.md

@@ -0,0 +1,5 @@
+---
+'@posthog/ai': minor
+---
+
+Add version-specific `Prompts.clearCache(name, version)` support and switch prompt caching to structured name/version entries instead of string-encoded cache keys.

packages/ai/src/prompts.ts

@@ -4,6 +4,7 @@ import type { PostHog } from 'posthog-node'
 import type { CachedPrompt, GetPromptOptions, PromptApiResponse, PromptVariables, PromptsDirectOptions } from './types'
 
 const DEFAULT_CACHE_TTL_SECONDS = 300 // 5 minutes
+type PromptVersionCache = Map<number | undefined, CachedPrompt>
 
 function isPromptApiResponse(data: unknown): data is PromptApiResponse {
   return (
@@ -63,7 +64,7 @@ export class Prompts {
   private projectApiKey: string
   private host: string
   private defaultCacheTtlSeconds: number
-  private cache: Map<string, CachedPrompt> = new Map()
+  private cache: Map<string, PromptVersionCache> = new Map()
 
   constructor(options: PromptsOptions) {
     this.defaultCacheTtlSeconds = options.defaultCacheTtlSeconds ?? DEFAULT_CACHE_TTL_SECONDS
@@ -80,8 +81,19 @@ export class Prompts {
     }
   }
 
-  private getCacheKey(name: string, version?: number): string {
-    return version === undefined ? `${name}::latest` : `${name}::version:${version}`
+  private getPromptCache(name: string): PromptVersionCache | undefined {
+    return this.cache.get(name)
+  }
+
+  private getOrCreatePromptCache(name: string): PromptVersionCache {
+    const cachedPromptVersions = this.cache.get(name)
+    if (cachedPromptVersions) {
+      return cachedPromptVersions
+    }
+
+    const promptVersions: PromptVersionCache = new Map()
+    this.cache.set(name, promptVersions)
+    return promptVersions
   }
 
   private getPromptLabel(name: string, version?: number): string {
@@ -100,11 +112,10 @@ export class Prompts {
     const cacheTtlSeconds = options?.cacheTtlSeconds ?? this.defaultCacheTtlSeconds
     const fallback = options?.fallback
     const version = options?.version
-    const cacheKey = this.getCacheKey(name, version)
     const promptLabel = this.getPromptLabel(name, version)
 
     // Check cache first
-    const cached = this.cache.get(cacheKey)
+    const cached = this.getPromptCache(name)?.get(version)
     const now = Date.now()
 
     if (cached) {
@@ -121,7 +132,7 @@ export class Prompts {
       const fetchedAt = Date.now()
 
       // Update cache
-      this.cache.set(cacheKey, {
+      this.getOrCreatePromptCache(name).set(version, {
         prompt,
         fetchedAt,
       })
@@ -169,24 +180,29 @@ export class Prompts {
   /**
    * Clear the cache for a specific prompt or all prompts
    *
-   * @param name - Optional prompt name to clear. If provided, clears all cached versions for that prompt.
+   * @param name - Optional prompt name to clear. If provided, clears all cached versions for that prompt unless a version is also provided.
+   * @param version - Optional prompt version to clear. Requires a prompt name.
    */
-  clearCache(name?: string): void {
-    if (name !== undefined) {
-      const latestKey = this.getCacheKey(name)
-      const versionPrefix = `${name}::version:`
-      for (const key of this.cache.keys()) {
-        if (key === latestKey) {
-          this.cache.delete(key)
-          continue
-        }
-
-        if (key.startsWith(versionPrefix) && /^\d+$/.test(key.slice(versionPrefix.length))) {
-          this.cache.delete(key)
-        }
-      }
-    } else {
+  clearCache(name?: string, version?: number): void {
+    if (version !== undefined && name === undefined) {
+      throw new Error("'version' requires 'name' to be provided")
+    }
+
+    if (name === undefined) {
       this.cache.clear()
+      return
+    }
+
+    if (version === undefined) {
+      this.cache.delete(name)
+      return
+    }
+
+    const promptVersions = this.getPromptCache(name)
+    promptVersions?.delete(version)
+
+    if (promptVersions?.size === 0) {
+      this.cache.delete(name)
     }
   }
 

packages/ai/tests/prompts.test.ts

@@ -588,6 +588,13 @@ describe('Prompts', () => {
   })
 
   describe('clearCache()', () => {
+    it('should throw when clearing a specific version without a prompt name', () => {
+      const posthog = createMockPostHog()
+      const prompts = new Prompts({ posthog })
+
+      expect(() => prompts.clearCache(undefined, 1)).toThrow("'version' requires 'name' to be provided")
+    })
+
     it('should clear a specific prompt from cache', async () => {
       mockFetch
         .mockResolvedValueOnce({
@@ -703,6 +710,63 @@ describe('Prompts', () => {
       expect(mockFetch).toHaveBeenCalledTimes(4)
     })
 
+    it('should clear only the requested cached version when name and version are provided', async () => {
+      mockFetch
+        .mockResolvedValueOnce({
+          ok: true,
+          status: 200,
+          json: () => Promise.resolve({ ...mockPromptResponse, version: 2, prompt: 'Latest prompt' }),
+        })
+        .mockResolvedValueOnce({
+          ok: true,
+          status: 200,
+          json: () => Promise.resolve({ ...mockPromptResponse, version: 1, prompt: 'Version 1 prompt' }),
+        })
+        .mockResolvedValueOnce({
+          ok: true,
+          status: 200,
+          json: () => Promise.resolve({ ...mockPromptResponse, version: 1, prompt: 'Version 1 prompt refreshed' }),
+        })
+
+      const posthog = createMockPostHog()
+      const prompts = new Prompts({ posthog })
+
+      await prompts.get('test-prompt')
+      await prompts.get('test-prompt', { version: 1 })
+      expect(mockFetch).toHaveBeenCalledTimes(2)
+
+      prompts.clearCache('test-prompt', 1)
+
+      await expect(prompts.get('test-prompt')).resolves.toBe('Latest prompt')
+      await expect(prompts.get('test-prompt', { version: 1 })).resolves.toBe('Version 1 prompt refreshed')
+      expect(mockFetch).toHaveBeenCalledTimes(3)
+    })
+
+    it('should remove the outer cache entry when the last version is cleared', async () => {
+      mockFetch
+        .mockResolvedValueOnce({
+          ok: true,
+          status: 200,
+          json: () => Promise.resolve({ ...mockPromptResponse, version: 1, prompt: 'Version 1 prompt' }),
+        })
+        .mockResolvedValueOnce({
+          ok: true,
+          status: 200,
+          json: () => Promise.resolve({ ...mockPromptResponse, version: 1, prompt: 'Version 1 prompt refreshed' }),
+        })
+
+      const posthog = createMockPostHog()
+      const prompts = new Prompts({ posthog })
+
+      await prompts.get('test-prompt', { version: 1 })
+      expect(mockFetch).toHaveBeenCalledTimes(1)
+
+      prompts.clearCache('test-prompt', 1)
+
+      await expect(prompts.get('test-prompt', { version: 1 })).resolves.toBe('Version 1 prompt refreshed')
+      expect(mockFetch).toHaveBeenCalledTimes(2)
+    })
+
     it('should not clear cache entries for other prompt names that share the same prefix', async () => {
       mockFetch
         .mockResolvedValueOnce({