wushenghua/hapi
1
0
Fork 0
mirror of https://github.com/wu736139669/hapi.git synced 2026-10-06 18:39:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(web): bound loaded generated-media memory with a small LRU

Browse Source 
Loaded video/audio/file blobs stayed in memory for the whole session
(object URLs were only revoked on unmount), so opening several tens-of-MB
videos in one chat grew memory without bound. Explicitly loaded media now
goes through a shared cache that keeps at most three entries; the least
recently used object URL is revoked on overflow and its card falls back to
the Load button (a sticky eviction flag prevents a slow load from
re-publishing an already-revoked URL). Auto-loaded images keep their
existing per-card lifetime.
This commit is contained in:
wushenghua
2026-10-01 00:13:38 +08:00
parent cc1fa87bd4
commit 587705dedb
4 changed files with 266 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
web/src/components/AssistantChat/messages/ToolMessage.generatedMedia.test.tsx
+62 -1
View File
@@ -1,10 +1,11 @@
import { describe, expect, it, vi } from 'vitest'
import { beforeEach, describe, expect, it, vi } from 'vitest'
import { fireEvent, render, screen, waitFor } from '@testing-library/react'
import { HappyChatProvider } from '@/components/AssistantChat/context'
import { GeneratedImageCard } from '@/components/AssistantChat/messages/ToolMessage'
import { I18nProvider } from '@/lib/i18n-context'
import type { ApiClient } from '@/api/client'
import type { HappyChatContextValue } from '@/components/AssistantChat/context'
import { _resetGeneratedMediaCacheForTests } from '@/lib/generatedMediaCache'
function renderCard(options: {
mimeType: string | null
@@ -54,7 +55,52 @@ function renderCard(options: {
return { getGeneratedImageBlob }
}
function renderVideoCards(imageIds: string[]) {
const getGeneratedImageBlob = vi.fn(async () => new Blob(['x'], { type: 'video/mp4' }))
const api = { getGeneratedImageBlob } as unknown as ApiClient
const value: HappyChatContextValue = {
api,
sessionId: 'session-1',
metadata: null,
terminalToolDisplayMode: 'compact',
showSessionSummaryInChat: false,
disabled: false,
onRefresh: () => {},
hasMoreMessages: false,
isSyncingTail: false,
isLoadingMoreMessages: false,
loadOlderMessagesPreservingScroll: async () => 'loaded',
}
render(
<I18nProvider>
<HappyChatProvider value={value}>
{imageIds.map((imageId) => (
<GeneratedImageCard
key={imageId}
block={{
kind: 'generated-image',
id: `block-${imageId}`,
localId: null,
createdAt: 1,
imageId,
fileName: `${imageId}.mp4`,
mimeType: 'video/mp4',
}}
/>
))}
</HappyChatProvider>
</I18nProvider>
)
return { getGeneratedImageBlob }
}
describe('GeneratedImageCard video fetch', () => {
beforeEach(() => {
_resetGeneratedMediaCacheForTests()
})
it('labels displayed images in English without implying AI generation', () => {
renderCard({ mimeType: 'image/png', locale: 'en' })
@@ -105,6 +151,21 @@ describe('GeneratedImageCard video fetch', () => {
})
})
it('returns an evicted video card to its Load button to bound memory', async () => {
renderVideoCards(['v1', 'v2', 'v3', 'v4'])
const buttons = screen.getAllByRole('button', { name: 'Load video' })
expect(buttons).toHaveLength(4)
buttons.forEach((button) => fireEvent.click(button))
// Cap is 3: the oldest load is evicted, so its card falls back to the
// Load button while three videos stay rendered.
await waitFor(() => {
expect(document.querySelectorAll('video')).toHaveLength(3)
expect(screen.getAllByRole('button', { name: 'Load video' })).toHaveLength(1)
})
})
it('loads unknown files on demand and renders a download link', async () => {
renderCard({ mimeType: 'application/octet-stream' })
web/src/components/AssistantChat/messages/ToolMessage.tsx
+33 -1
View File
@@ -18,6 +18,7 @@ import { ImagePreview } from '@/components/ImagePreview'
import { FileIcon } from '@/components/FileIcon'
import { useTranslation } from '@/lib/use-translation'
import { inlineMediaLabelKey, isInlineAudioMimeType, isInlineImageMimeType, isInlineVideoMimeType } from '@/lib/generatedInlineMedia'
import { loadGeneratedMediaObjectUrl, subscribeGeneratedMediaEviction } from '@/lib/generatedMediaCache'
function isToolCallBlock(value: unknown): value is ToolCallBlock {
if (!isObject(value)) return false
@@ -72,6 +73,7 @@ export function GeneratedImageCard(props: { block: GeneratedImageBlock }) {
const [imageStyle, setImageStyle] = useState<CSSProperties | undefined>(undefined)
const [loadMedia, setLoadMedia] = useState(false)
const objectUrlRef = useRef<string | null>(null)
const evictedRef = useRef(false)
const isVideo = isInlineVideoMimeType(props.block.mimeType)
const isAudio = isInlineAudioMimeType(props.block.mimeType)
const isImage = isInlineImageMimeType(props.block.mimeType)
@@ -90,6 +92,8 @@ export function GeneratedImageCard(props: { block: GeneratedImageBlock }) {
}
}, [])
const mediaKey = `${ctx.sessionId}:${props.block.imageId}`
useEffect(() => {
if (!shouldFetch) {
return
@@ -97,6 +101,34 @@ export function GeneratedImageCard(props: { block: GeneratedImageBlock }) {
let disposed = false
if (!isImage) {
// Explicitly loaded media lives in a bounded shared cache: the cap
// keeps at most a few tens-of-MB blobs alive per session, and the
// listener brings this card back to its Load button when evicted.
setError(null)
evictedRef.current = false
const unsubscribe = subscribeGeneratedMediaEviction(mediaKey, () => {
evictedRef.current = true
setObjectUrl(null)
setLoadMedia(false)
})
void loadGeneratedMediaObjectUrl(mediaKey, () => ctx.api.getGeneratedImageBlob(ctx.sessionId, props.block.imageId))
.then((url) => {
if (!disposed && !evictedRef.current) {
setObjectUrl(url)
}
})
.catch((err: unknown) => {
if (!disposed) {
setError(err instanceof Error ? err.message : 'Failed to load inline media')
}
})
return () => {
disposed = true
unsubscribe()
}
}
if (objectUrlRef.current) {
URL.revokeObjectURL(objectUrlRef.current)
objectUrlRef.current = null
@@ -133,7 +165,7 @@ export function GeneratedImageCard(props: { block: GeneratedImageBlock }) {
return () => {
disposed = true
}
}, [ctx.api, ctx.sessionId, props.block.imageId, isImage, shouldFetch])
}, [ctx.api, ctx.sessionId, props.block.imageId, isImage, shouldFetch, mediaKey])
return (
<div className="max-w-[92%] rounded-2xl border border-[var(--app-border)] bg-[var(--app-tool-card-bg)] p-3">
web/src/lib/generatedMediaCache.test.ts
+81
View File
@@ -0,0 +1,81 @@
import { afterEach, describe, expect, it, vi } from 'vitest'
import {
_resetGeneratedMediaCacheForTests,
loadGeneratedMediaObjectUrl,
peekGeneratedMedia,
subscribeGeneratedMediaEviction
} from './generatedMediaCache'
afterEach(() => {
_resetGeneratedMediaCacheForTests()
vi.restoreAllMocks()
})
describe('generatedMediaCache', () => {
it('reuses a cached object URL without loading the blob again', async () => {
const load = vi.fn(async () => new Blob(['a']))
const first = await loadGeneratedMediaObjectUrl('s:1', load)
const second = await loadGeneratedMediaObjectUrl('s:1', load)
expect(second).toBe(first)
expect(load).toHaveBeenCalledTimes(1)
})
it('revokes the least recently used blob once the cap is exceeded', async () => {
const revoke = vi.spyOn(URL, 'revokeObjectURL')
const urls: string[] = []
for (const key of ['a', 'b', 'c', 'd']) {
urls.push(await loadGeneratedMediaObjectUrl(`s:${key}`, async () => new Blob([key])))
}
expect(peekGeneratedMedia('s:a')).toBeNull()
expect(peekGeneratedMedia('s:b')).toBe(urls[1])
expect(peekGeneratedMedia('s:c')).toBe(urls[2])
expect(peekGeneratedMedia('s:d')).toBe(urls[3])
expect(revoke).toHaveBeenCalledWith(urls[0])
})
it('protects a cache hit from eviction by refreshing its position', async () => {
const revoke = vi.spyOn(URL, 'revokeObjectURL')
const a = await loadGeneratedMediaObjectUrl('s:a', async () => new Blob(['a']))
await loadGeneratedMediaObjectUrl('s:b', async () => new Blob(['b']))
const c = await loadGeneratedMediaObjectUrl('s:c', async () => new Blob(['c']))
// Touch a without reloading; b is now the oldest.
await loadGeneratedMediaObjectUrl('s:a', async () => {
throw new Error('cache hit must not load')
})
const d = await loadGeneratedMediaObjectUrl('s:d', async () => new Blob(['d']))
expect(peekGeneratedMedia('s:b')).toBeNull()
expect(peekGeneratedMedia('s:a')).toBe(a)
expect(peekGeneratedMedia('s:c')).toBe(c)
expect(peekGeneratedMedia('s:d')).toBe(d)
expect(revoke).toHaveBeenCalledTimes(1)
})
it('notifies holders when their entry is evicted', async () => {
const evicted = vi.fn()
await loadGeneratedMediaObjectUrl('s:a', async () => new Blob(['a']))
subscribeGeneratedMediaEviction('s:a', evicted)
await loadGeneratedMediaObjectUrl('s:b', async () => new Blob(['b']))
await loadGeneratedMediaObjectUrl('s:c', async () => new Blob(['c']))
expect(evicted).not.toHaveBeenCalled()
await loadGeneratedMediaObjectUrl('s:d', async () => new Blob(['d']))
expect(evicted).toHaveBeenCalledTimes(1)
})
it('stops notifying after unsubscribe', async () => {
const evicted = vi.fn()
await loadGeneratedMediaObjectUrl('s:a', async () => new Blob(['a']))
const unsubscribe = subscribeGeneratedMediaEviction('s:a', evicted)
unsubscribe()
await loadGeneratedMediaObjectUrl('s:b', async () => new Blob(['b']))
await loadGeneratedMediaObjectUrl('s:c', async () => new Blob(['c']))
await loadGeneratedMediaObjectUrl('s:d', async () => new Blob(['d']))
expect(evicted).not.toHaveBeenCalled()
})
})
web/src/lib/generatedMediaCache.ts
+90
View File
@@ -0,0 +1,90 @@
/**
* Bounded cache for generated-media object URLs (videos/audio/files the user
* explicitly loads).
*
* Blobs are tens of MB and the chat keeps every message mounted, so without a
* cap each loaded video would stay in memory for the whole session. The least
* recently used entry is revoked once a new load exceeds the cap; evicted
* holders are notified so their card can fall back to its Load button.
*
* Auto-loaded images are intentionally NOT routed through here: they are small
* and always rendered, so caching/evicting them would only cause reload churn.
*/
const MAX_CACHED_MEDIA = 3
type Entry = {
objectUrl: string
}
/** Insertion order doubles as LRU order (Map preserves it; hits re-insert). */
const entries = new Map<string, Entry>()
const listeners = new Map<string, Set<() => void>>()
function notify(key: string): void {
const set = listeners.get(key)
if (!set) return
for (const listener of [...set]) {
listener()
}
}
function evictOverflow(): void {
while (entries.size > MAX_CACHED_MEDIA) {
const oldest = entries.entries().next().value as [string, Entry] | undefined
if (!oldest) return
entries.delete(oldest[0])
notify(oldest[0])
URL.revokeObjectURL(oldest[1].objectUrl)
}
}
/** Cache hit without loading; refreshes LRU position. */
export function peekGeneratedMedia(key: string): string | null {
const entry = entries.get(key)
if (!entry) return null
entries.delete(key)
entries.set(key, entry)
return entry.objectUrl
}
/** Load (or reuse) the object URL for one media key, evicting the oldest over the cap. */
export async function loadGeneratedMediaObjectUrl(key: string, load: () => Promise<Blob>): Promise<string> {
const cached = peekGeneratedMedia(key)
if (cached) return cached
const blob = await load()
// A concurrent load for the same key may have finished first.
const raced = peekGeneratedMedia(key)
if (raced) return raced
const objectUrl = URL.createObjectURL(blob)
entries.set(key, { objectUrl })
evictOverflow()
return objectUrl
}
/** Called when this key is evicted; the holder should drop its rendered URL. */
export function subscribeGeneratedMediaEviction(key: string, listener: () => void): () => void {
let set = listeners.get(key)
if (!set) {
set = new Set()
listeners.set(key, set)
}
set.add(listener)
return () => {
set.delete(listener)
if (set.size === 0) {
listeners.delete(key)
}
}
}
/** Test helper: drop everything and revoke outstanding URLs. */
export function _resetGeneratedMediaCacheForTests(): void {
for (const entry of entries.values()) {
URL.revokeObjectURL(entry.objectUrl)
}
entries.clear()
listeners.clear()
}