mirror of
https://github.com/NanmiCoder/cc-haha
synced 2026-07-26 15:03:34 +08:00
fix(api): add media size budget to prevent 20MB request overflow (#68)
* fix(api): add media size budget to prevent 20MB request limit overflow Conversations accumulate base64 images (from Read tool and image-gen plugin) that are never evicted by toolResultStorage. This adds shrinkMediaToSizeBudget() which enforces a 12MB total base64 media budget per request, removing oldest images first. Also upgrades the ToolCallBlock image grid to support click-to-fullscreen with left/right navigation via ImageGalleryModal. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> * test(desktop): add ImageBlockGallery tests for fullscreen and navigation Satisfies CI change-policy requiring desktop tests for product code changes. Tests cover grid rendering, click-to-fullscreen, left/right keyboard nav, and wrap-around behavior. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> --------- Co-authored-by: 你的姓名 <you@example.com> Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
小橙子
GitHub
parent
2c897d2c82
commit
f6872df48b
105
desktop/src/components/chat/ImageBlockGallery.test.tsx
Normal file
105
desktop/src/components/chat/ImageBlockGallery.test.tsx
Normal file
@ -0,0 +1,105 @@
|
||||
import '@testing-library/jest-dom'
|
||||
import { render, fireEvent } from '@testing-library/react'
|
||||
import { describe, expect, it } from 'vitest'
|
||||
|
||||
import { ImageBlockGallery, type ImageBlock } from './ToolCallBlock'
|
||||
|
||||
const singleImage: ImageBlock[] = [
|
||||
{ src: 'https://example.com/img1.png', mimeType: 'image/png' },
|
||||
]
|
||||
|
||||
const multipleImages: ImageBlock[] = [
|
||||
{ src: 'https://example.com/img1.png', mimeType: 'image/png' },
|
||||
{ src: 'https://example.com/img2.png', mimeType: 'image/png' },
|
||||
{ src: 'https://example.com/img3.png', mimeType: 'image/png' },
|
||||
]
|
||||
|
||||
/** Get the main (large) image inside the modal overlay. */
|
||||
function getModalMainImage(container: HTMLElement): HTMLImageElement | null {
|
||||
// The modal renders inside a [role=dialog] or a fixed overlay with max-h-[70vh]
|
||||
const modalImg = container.querySelector('.max-h-\\[70vh\\] img') as HTMLImageElement | null
|
||||
return modalImg
|
||||
}
|
||||
|
||||
describe('ImageBlockGallery', () => {
|
||||
it('renders images in a grid with correct count label', () => {
|
||||
const { container } = render(<ImageBlockGallery imageBlocks={multipleImages} />)
|
||||
const images = container.querySelectorAll('img')
|
||||
expect(images).toHaveLength(3)
|
||||
expect(container.textContent).toContain('3 images')
|
||||
})
|
||||
|
||||
it('renders single image with singular label', () => {
|
||||
const { container } = render(<ImageBlockGallery imageBlocks={singleImage} />)
|
||||
const images = container.querySelectorAll('img')
|
||||
expect(images).toHaveLength(1)
|
||||
expect(container.textContent).toContain('1 image')
|
||||
expect(container.textContent).not.toContain('1 images')
|
||||
})
|
||||
|
||||
it('uses 2-column grid for multiple images', () => {
|
||||
const { container } = render(<ImageBlockGallery imageBlocks={multipleImages} />)
|
||||
const grid = container.querySelector('.grid-cols-2')
|
||||
expect(grid).toBeInTheDocument()
|
||||
})
|
||||
|
||||
it('uses 1-column grid for single image', () => {
|
||||
const { container } = render(<ImageBlockGallery imageBlocks={singleImage} />)
|
||||
const grid = container.querySelector('.grid-cols-1')
|
||||
expect(grid).toBeInTheDocument()
|
||||
expect(container.querySelector('.grid-cols-2')).not.toBeInTheDocument()
|
||||
})
|
||||
|
||||
it('opens fullscreen modal on image click', () => {
|
||||
const { container, baseElement } = render(<ImageBlockGallery imageBlocks={multipleImages} />)
|
||||
|
||||
// Click the first image button
|
||||
const buttons = container.querySelectorAll('button[type="button"]')
|
||||
expect(buttons.length).toBe(3)
|
||||
fireEvent.click(buttons[0]!)
|
||||
|
||||
// Modal should now be open — find the large main image
|
||||
const modalImg = getModalMainImage(baseElement)
|
||||
expect(modalImg).toBeInTheDocument()
|
||||
expect(modalImg!.src).toBe('https://example.com/img1.png')
|
||||
})
|
||||
|
||||
it('supports left/right keyboard navigation in modal', () => {
|
||||
const { container, baseElement } = render(<ImageBlockGallery imageBlocks={multipleImages} />)
|
||||
|
||||
// Open modal on first image
|
||||
const buttons = container.querySelectorAll('button[type="button"]')
|
||||
fireEvent.click(buttons[0]!)
|
||||
|
||||
// Press right arrow to go to second image
|
||||
fireEvent.keyDown(document, { key: 'ArrowRight' })
|
||||
let modalImg = getModalMainImage(baseElement)
|
||||
expect(modalImg!.src).toBe('https://example.com/img2.png')
|
||||
|
||||
// Press right arrow again to go to third image
|
||||
fireEvent.keyDown(document, { key: 'ArrowRight' })
|
||||
modalImg = getModalMainImage(baseElement)
|
||||
expect(modalImg!.src).toBe('https://example.com/img3.png')
|
||||
|
||||
// Press left arrow to go back to second image
|
||||
fireEvent.keyDown(document, { key: 'ArrowLeft' })
|
||||
modalImg = getModalMainImage(baseElement)
|
||||
expect(modalImg!.src).toBe('https://example.com/img2.png')
|
||||
})
|
||||
|
||||
it('wraps around when navigating past last/first image', () => {
|
||||
const { container, baseElement } = render(<ImageBlockGallery imageBlocks={multipleImages} />)
|
||||
|
||||
// Open modal on last image (index 2)
|
||||
const buttons = container.querySelectorAll('button[type="button"]')
|
||||
fireEvent.click(buttons[2]!)
|
||||
|
||||
let modalImg = getModalMainImage(baseElement)
|
||||
expect(modalImg!.src).toBe('https://example.com/img3.png')
|
||||
|
||||
// Press right arrow — should wrap to first image
|
||||
fireEvent.keyDown(document, { key: 'ArrowRight' })
|
||||
modalImg = getModalMainImage(baseElement)
|
||||
expect(modalImg!.src).toBe('https://example.com/img1.png')
|
||||
})
|
||||
})
|
||||
@ -7,6 +7,7 @@ import { CopyButton } from '../shared/CopyButton'
|
||||
import { useTranslation } from '../../i18n'
|
||||
import type { TranslationKey } from '../../i18n'
|
||||
import { InlineImageGallery } from './InlineImageGallery'
|
||||
import { ImageGalleryModal } from './ImageGalleryModal'
|
||||
import type { AgentTaskNotification } from '../../types/chat'
|
||||
import { PlanPreviewCard, extractPlanPreview, isExitPlanModeTool } from './PlanModePreview'
|
||||
|
||||
@ -274,6 +275,57 @@ function getVisibleResultText(
|
||||
return text
|
||||
}
|
||||
|
||||
/** Renders extracted image blocks in a grid with click-to-fullscreen gallery. */
|
||||
export function ImageBlockGallery({ imageBlocks }: { imageBlocks: ImageBlock[] }) {
|
||||
const [activeIndex, setActiveIndex] = useState<number | null>(null)
|
||||
const galleryImages = useMemo(
|
||||
() => imageBlocks.map((img, i) => ({ src: img.src, name: `Image ${i + 1}` })),
|
||||
[imageBlocks],
|
||||
)
|
||||
|
||||
return (
|
||||
<div className="mt-2 space-y-2">
|
||||
<div className="flex items-center gap-1.5 text-[10px] font-semibold uppercase tracking-wider text-[var(--color-outline)]">
|
||||
<span className="material-symbols-outlined text-[12px]">image</span>
|
||||
{imageBlocks.length === 1 ? '1 image' : `${imageBlocks.length} images`}
|
||||
</div>
|
||||
<div className={`grid gap-2 ${imageBlocks.length === 1 ? 'grid-cols-1' : 'grid-cols-2'}`}>
|
||||
{imageBlocks.map((img, i) => (
|
||||
<button
|
||||
key={img.src}
|
||||
type="button"
|
||||
onClick={() => setActiveIndex(i)}
|
||||
className="group/image relative overflow-hidden rounded-xl border border-[var(--color-border)] bg-[var(--color-surface-container-low)] text-left shadow-sm transition-all hover:shadow-md hover:border-[var(--color-brand)]/40"
|
||||
>
|
||||
<img
|
||||
src={img.src}
|
||||
alt={`Generated image ${i + 1}`}
|
||||
loading="lazy"
|
||||
className="w-full object-contain"
|
||||
style={{ maxHeight: imageBlocks.length === 1 ? 400 : 240 }}
|
||||
/>
|
||||
<div className="absolute inset-0 flex items-center justify-center bg-black/0 opacity-0 transition-all group-hover/image:bg-black/20 group-hover/image:opacity-100">
|
||||
<span className="material-symbols-outlined rounded-full bg-white/90 p-2 text-[20px] text-[var(--color-text-primary)] shadow-lg">
|
||||
fullscreen
|
||||
</span>
|
||||
</div>
|
||||
</button>
|
||||
))}
|
||||
</div>
|
||||
|
||||
{activeIndex !== null && activeIndex >= 0 && (
|
||||
<ImageGalleryModal
|
||||
open={activeIndex !== null}
|
||||
images={galleryImages}
|
||||
activeIndex={activeIndex}
|
||||
onClose={() => setActiveIndex(null)}
|
||||
onSelect={setActiveIndex}
|
||||
/>
|
||||
)}
|
||||
</div>
|
||||
)
|
||||
}
|
||||
|
||||
function renderResultOutput(
|
||||
result: { content: unknown; isError: boolean },
|
||||
text: string,
|
||||
@ -283,28 +335,7 @@ function renderResultOutput(
|
||||
return (
|
||||
<>
|
||||
{imageBlocks.length > 0 && (
|
||||
<div className="mt-2 space-y-2">
|
||||
<div className="flex items-center gap-1.5 text-[10px] font-semibold uppercase tracking-wider text-[var(--color-outline)]">
|
||||
<span className="material-symbols-outlined text-[12px]">image</span>
|
||||
{imageBlocks.length === 1 ? '1 image' : `${imageBlocks.length} images`}
|
||||
</div>
|
||||
<div className={`grid gap-2 ${imageBlocks.length === 1 ? 'grid-cols-1' : 'grid-cols-2'}`}>
|
||||
{imageBlocks.map((img, i) => (
|
||||
<div
|
||||
key={i}
|
||||
className="overflow-hidden rounded-xl border border-[var(--color-border)] bg-[var(--color-surface-container-low)]"
|
||||
>
|
||||
<img
|
||||
src={img.src}
|
||||
alt={`Generated image ${i + 1}`}
|
||||
loading="lazy"
|
||||
className="w-full object-contain"
|
||||
style={{ maxHeight: 400 }}
|
||||
/>
|
||||
</div>
|
||||
))}
|
||||
</div>
|
||||
</div>
|
||||
<ImageBlockGallery imageBlocks={imageBlocks} />
|
||||
)}
|
||||
<InlineImageGallery text={text} />
|
||||
<div className={`overflow-hidden rounded-lg border ${
|
||||
@ -571,7 +602,7 @@ function extractTextContent(content: unknown): string | null {
|
||||
return null
|
||||
}
|
||||
|
||||
type ImageBlock = { src: string; mimeType: string }
|
||||
export type ImageBlock = { src: string; mimeType: string }
|
||||
|
||||
function extractImageBlocks(content: unknown): ImageBlock[] {
|
||||
if (!Array.isArray(content)) return []
|
||||
|
||||
@ -92,3 +92,17 @@ export const PDF_AT_MENTION_INLINE_THRESHOLD = 10
|
||||
* We validate client-side to provide a clear error message.
|
||||
*/
|
||||
export const API_MAX_MEDIA_PER_REQUEST = 100
|
||||
|
||||
/**
|
||||
* Maximum total base64-encoded media size budget per API request (bytes).
|
||||
*
|
||||
* The API has a ~20 MB total request size limit. Conversations accumulate
|
||||
* base64 images (from the Read tool viewing local files and from image-gen
|
||||
* plugin results) that never get evicted by toolResultStorage because
|
||||
* hasImageBlock() skips both persistence and budget eviction for image blocks.
|
||||
*
|
||||
* We enforce a 12 MB media budget to leave headroom for the rest of the
|
||||
* conversation context (system prompt, text messages, tool definitions).
|
||||
* When over budget, the oldest media blocks are replaced with text placeholders.
|
||||
*/
|
||||
export const API_MAX_MEDIA_BYTES_BUDGET = 12 * 1024 * 1024 // 12 MB
|
||||
|
||||
@ -191,7 +191,10 @@ import {
|
||||
isDeferredToolsDeltaEnabled,
|
||||
isToolSearchEnabled,
|
||||
} from "src/utils/toolSearch.js";
|
||||
import { API_MAX_MEDIA_PER_REQUEST } from "../../constants/apiLimits.js";
|
||||
import {
|
||||
API_MAX_MEDIA_BYTES_BUDGET,
|
||||
API_MAX_MEDIA_PER_REQUEST,
|
||||
} from "../../constants/apiLimits.js";
|
||||
import { ADVISOR_BETA_HEADER } from "../../constants/betas.js";
|
||||
import {
|
||||
formatDeferredToolLine,
|
||||
@ -954,6 +957,154 @@ function isToolResult(
|
||||
return block.type === "tool_result";
|
||||
}
|
||||
|
||||
/**
|
||||
* Returns the base64 payload size (in bytes) for a media block, or 0 if the
|
||||
* block uses a URL/file source (which doesn't contribute to request body size).
|
||||
*/
|
||||
function getMediaBase64Size(
|
||||
block: BetaImageBlockParam | BetaRequestDocumentBlock,
|
||||
): number {
|
||||
const source = block.source;
|
||||
if ("type" in source && source.type === "base64" && "data" in source) {
|
||||
return (source as { data: string }).data.length;
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
|
||||
/**
|
||||
* Tracked position of a media block in the message array, used to remove the
|
||||
* oldest items first when over budget.
|
||||
*/
|
||||
interface MediaPosition {
|
||||
msgIdx: number;
|
||||
/** Index in the message's top-level content array */
|
||||
blockIdx: number;
|
||||
/** If inside a tool_result, index within tool_result.content; otherwise -1 */
|
||||
nestedIdx: number;
|
||||
size: number;
|
||||
}
|
||||
|
||||
/**
|
||||
* Shrinks total base64 media payload to fit within `budgetBytes`.
|
||||
*
|
||||
* When the conversation accumulates large base64 images (from Read tool or
|
||||
* image-gen plugin), the total request can exceed the API's ~20 MB limit.
|
||||
* This function removes the OLDEST base64 media blocks first (replacing them
|
||||
* with a text placeholder) until total media size is within budget.
|
||||
*
|
||||
* URL-based images (size 0) are never removed — they don't affect request size.
|
||||
*/
|
||||
export function shrinkMediaToSizeBudget(
|
||||
messages: (UserMessage | AssistantMessage)[],
|
||||
budgetBytes: number,
|
||||
): (UserMessage | AssistantMessage)[] {
|
||||
// Pass 1: collect all media positions and their sizes
|
||||
const positions: MediaPosition[] = [];
|
||||
let totalSize = 0;
|
||||
|
||||
for (let msgIdx = 0; msgIdx < messages.length; msgIdx++) {
|
||||
const content = messages[msgIdx].message.content;
|
||||
if (!Array.isArray(content)) continue;
|
||||
|
||||
for (let blockIdx = 0; blockIdx < content.length; blockIdx++) {
|
||||
const block = content[blockIdx];
|
||||
if (isMedia(block)) {
|
||||
const size = getMediaBase64Size(block);
|
||||
if (size > 0) {
|
||||
positions.push({ msgIdx, blockIdx, nestedIdx: -1, size });
|
||||
totalSize += size;
|
||||
}
|
||||
}
|
||||
if (isToolResult(block) && Array.isArray(block.content)) {
|
||||
for (
|
||||
let nestedIdx = 0;
|
||||
nestedIdx < block.content.length;
|
||||
nestedIdx++
|
||||
) {
|
||||
const nested = block.content[nestedIdx];
|
||||
if (isMedia(nested)) {
|
||||
const size = getMediaBase64Size(nested);
|
||||
if (size > 0) {
|
||||
positions.push({ msgIdx, blockIdx, nestedIdx, size });
|
||||
totalSize += size;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (totalSize <= budgetBytes) return messages;
|
||||
|
||||
// Pass 2: mark oldest positions for removal until within budget
|
||||
const toRemove = new Set<MediaPosition>();
|
||||
for (const pos of positions) {
|
||||
if (totalSize <= budgetBytes) break;
|
||||
toRemove.add(pos);
|
||||
totalSize -= pos.size;
|
||||
}
|
||||
|
||||
// Pass 3: rebuild messages, replacing removed media with text placeholders
|
||||
return messages.map((msg, msgIdx) => {
|
||||
const content = msg.message.content;
|
||||
if (!Array.isArray(content)) return msg;
|
||||
|
||||
let changed = false;
|
||||
const newContent = content.map((block, blockIdx) => {
|
||||
// Check top-level media
|
||||
if (isMedia(block)) {
|
||||
for (const pos of toRemove) {
|
||||
if (
|
||||
pos.msgIdx === msgIdx &&
|
||||
pos.blockIdx === blockIdx &&
|
||||
pos.nestedIdx === -1
|
||||
) {
|
||||
changed = true;
|
||||
return {
|
||||
type: "text" as const,
|
||||
text: "[media removed: conversation media budget exceeded]",
|
||||
};
|
||||
}
|
||||
}
|
||||
return block;
|
||||
}
|
||||
|
||||
// Check nested media inside tool_results
|
||||
if (isToolResult(block) && Array.isArray(block.content)) {
|
||||
let nestedChanged = false;
|
||||
const newNested = block.content.map((nested, nestedIdx) => {
|
||||
if (isMedia(nested)) {
|
||||
for (const pos of toRemove) {
|
||||
if (
|
||||
pos.msgIdx === msgIdx &&
|
||||
pos.blockIdx === blockIdx &&
|
||||
pos.nestedIdx === nestedIdx
|
||||
) {
|
||||
nestedChanged = true;
|
||||
return {
|
||||
type: "text" as const,
|
||||
text: "[media removed: conversation media budget exceeded]",
|
||||
};
|
||||
}
|
||||
}
|
||||
}
|
||||
return nested;
|
||||
});
|
||||
if (nestedChanged) {
|
||||
changed = true;
|
||||
return { ...block, content: newNested };
|
||||
}
|
||||
}
|
||||
|
||||
return block;
|
||||
});
|
||||
|
||||
return changed
|
||||
? { ...msg, message: { ...msg.message, content: newContent } }
|
||||
: msg;
|
||||
}) as (UserMessage | AssistantMessage)[];
|
||||
}
|
||||
|
||||
/**
|
||||
* Ensures messages contain at most `limit` media items (images + documents).
|
||||
* Strips oldest media first to preserve the most recent.
|
||||
@ -1366,6 +1517,14 @@ async function* queryModel(
|
||||
API_MAX_MEDIA_PER_REQUEST,
|
||||
);
|
||||
|
||||
// Shrink total base64 media payload to stay under the API's ~20 MB request
|
||||
// size limit. Images accumulate in conversations (Read tool + image-gen)
|
||||
// and are never evicted by toolResultStorage; this is the safety net.
|
||||
messagesForAPI = shrinkMediaToSizeBudget(
|
||||
messagesForAPI,
|
||||
API_MAX_MEDIA_BYTES_BUDGET,
|
||||
);
|
||||
|
||||
// Instrumentation: Track message count after normalization
|
||||
logEvent("tengu_api_after_normalize", {
|
||||
postNormalizedMessageCount: messagesForAPI.length,
|
||||
|
||||
255
src/services/api/shrinkMediaBudget.test.ts
Normal file
255
src/services/api/shrinkMediaBudget.test.ts
Normal file
@ -0,0 +1,255 @@
|
||||
import { describe, expect, test } from 'bun:test'
|
||||
import { shrinkMediaToSizeBudget } from './claude.js'
|
||||
import type { UserMessage, AssistantMessage } from '../../types/message.js'
|
||||
|
||||
/**
|
||||
* Helper: create a base64 image block with a payload of the given byte length.
|
||||
* Uses 'A' repeated to simulate base64 data.
|
||||
*/
|
||||
function makeBase64ImageBlock(sizeBytes: number) {
|
||||
return {
|
||||
type: 'image' as const,
|
||||
source: {
|
||||
type: 'base64' as const,
|
||||
media_type: 'image/png' as const,
|
||||
data: 'A'.repeat(sizeBytes),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
/** Helper: create a URL image block (zero payload, should never be removed). */
|
||||
function makeURLImageBlock() {
|
||||
return {
|
||||
type: 'image' as const,
|
||||
source: {
|
||||
type: 'url' as const,
|
||||
url: 'https://example.com/image.png',
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
/** Helper: create a base64 PDF document block. */
|
||||
function makeBase64DocBlock(sizeBytes: number) {
|
||||
return {
|
||||
type: 'document' as const,
|
||||
source: {
|
||||
type: 'base64' as const,
|
||||
media_type: 'application/pdf' as const,
|
||||
data: 'B'.repeat(sizeBytes),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
/** Helper: wrap content blocks into a minimal UserMessage. */
|
||||
function userMsg(content: any[]): UserMessage {
|
||||
return {
|
||||
type: 'user',
|
||||
message: { role: 'user', content },
|
||||
uuid: crypto.randomUUID() as any,
|
||||
timestamp: new Date().toISOString(),
|
||||
} as UserMessage
|
||||
}
|
||||
|
||||
/** Helper: wrap tool_result content into a UserMessage. */
|
||||
function toolResultMsg(toolUseId: string, content: any[]): UserMessage {
|
||||
return {
|
||||
type: 'user',
|
||||
message: {
|
||||
role: 'user',
|
||||
content: [
|
||||
{
|
||||
type: 'tool_result' as const,
|
||||
tool_use_id: toolUseId,
|
||||
content,
|
||||
},
|
||||
],
|
||||
},
|
||||
uuid: crypto.randomUUID() as any,
|
||||
timestamp: new Date().toISOString(),
|
||||
} as UserMessage
|
||||
}
|
||||
|
||||
describe('shrinkMediaToSizeBudget', () => {
|
||||
test('returns messages unchanged when total media is under budget', () => {
|
||||
const messages = [
|
||||
userMsg([makeBase64ImageBlock(1000), { type: 'text', text: 'hello' }]),
|
||||
userMsg([makeBase64ImageBlock(2000)]),
|
||||
]
|
||||
|
||||
const result = shrinkMediaToSizeBudget(messages, 5000)
|
||||
expect(result).toBe(messages) // same reference, no copy
|
||||
})
|
||||
|
||||
test('removes oldest base64 images first when over budget', () => {
|
||||
const messages = [
|
||||
userMsg([makeBase64ImageBlock(3000)]), // oldest — should be removed
|
||||
userMsg([makeBase64ImageBlock(3000)]), // second oldest — should be removed
|
||||
userMsg([makeBase64ImageBlock(3000)]), // newest — kept
|
||||
]
|
||||
|
||||
const result = shrinkMediaToSizeBudget(messages, 4000)
|
||||
|
||||
// First two messages should have placeholder text
|
||||
const content0 = result[0].message.content as any[]
|
||||
expect(content0[0].type).toBe('text')
|
||||
expect(content0[0].text).toContain('media removed')
|
||||
|
||||
const content1 = result[1].message.content as any[]
|
||||
expect(content1[0].type).toBe('text')
|
||||
expect(content1[0].text).toContain('media removed')
|
||||
|
||||
// Third message (newest) should be preserved
|
||||
const content2 = result[2].message.content as any[]
|
||||
expect(content2[0].type).toBe('image')
|
||||
expect(content2[0].source.type).toBe('base64')
|
||||
})
|
||||
|
||||
test('never removes URL-based images (zero payload)', () => {
|
||||
const messages = [
|
||||
userMsg([makeURLImageBlock()]),
|
||||
userMsg([makeBase64ImageBlock(5000)]),
|
||||
userMsg([makeURLImageBlock()]),
|
||||
]
|
||||
|
||||
// Budget is 3000 — only the base64 image exceeds it
|
||||
const result = shrinkMediaToSizeBudget(messages, 3000)
|
||||
|
||||
// URL images untouched
|
||||
const content0 = result[0].message.content as any[]
|
||||
expect(content0[0].type).toBe('image')
|
||||
expect(content0[0].source.type).toBe('url')
|
||||
|
||||
const content2 = result[2].message.content as any[]
|
||||
expect(content2[0].type).toBe('image')
|
||||
expect(content2[0].source.type).toBe('url')
|
||||
|
||||
// Base64 image replaced
|
||||
const content1 = result[1].message.content as any[]
|
||||
expect(content1[0].type).toBe('text')
|
||||
expect(content1[0].text).toContain('media removed')
|
||||
})
|
||||
|
||||
test('handles media nested inside tool_result blocks', () => {
|
||||
const messages = [
|
||||
toolResultMsg('tool-1', [
|
||||
makeBase64ImageBlock(4000), // oldest nested — removed
|
||||
{ type: 'text', text: 'screenshot result' },
|
||||
]),
|
||||
toolResultMsg('tool-2', [
|
||||
makeBase64ImageBlock(4000), // newer nested — kept
|
||||
]),
|
||||
]
|
||||
|
||||
const result = shrinkMediaToSizeBudget(messages, 5000)
|
||||
|
||||
// First tool_result: image replaced, text preserved
|
||||
const tr0 = (result[0].message.content as any[])[0]
|
||||
expect(tr0.type).toBe('tool_result')
|
||||
expect(tr0.content[0].type).toBe('text')
|
||||
expect(tr0.content[0].text).toContain('media removed')
|
||||
expect(tr0.content[1].type).toBe('text')
|
||||
expect(tr0.content[1].text).toBe('screenshot result')
|
||||
|
||||
// Second tool_result: image preserved
|
||||
const tr1 = (result[1].message.content as any[])[0]
|
||||
expect(tr1.content[0].type).toBe('image')
|
||||
})
|
||||
|
||||
test('handles mixed media types (images + documents)', () => {
|
||||
const messages = [
|
||||
userMsg([makeBase64DocBlock(5000)]), // oldest PDF — removed
|
||||
userMsg([makeBase64ImageBlock(5000)]), // image — removed
|
||||
userMsg([makeBase64ImageBlock(3000)]), // newest — kept
|
||||
]
|
||||
|
||||
const result = shrinkMediaToSizeBudget(messages, 4000)
|
||||
|
||||
const content0 = result[0].message.content as any[]
|
||||
expect(content0[0].type).toBe('text')
|
||||
expect(content0[0].text).toContain('media removed')
|
||||
|
||||
const content1 = result[1].message.content as any[]
|
||||
expect(content1[0].type).toBe('text')
|
||||
expect(content1[0].text).toContain('media removed')
|
||||
|
||||
const content2 = result[2].message.content as any[]
|
||||
expect(content2[0].type).toBe('image')
|
||||
})
|
||||
|
||||
test('preserves non-media content blocks alongside removed media', () => {
|
||||
const messages = [
|
||||
userMsg([
|
||||
{ type: 'text', text: 'before' },
|
||||
makeBase64ImageBlock(6000),
|
||||
{ type: 'text', text: 'after' },
|
||||
]),
|
||||
]
|
||||
|
||||
const result = shrinkMediaToSizeBudget(messages, 3000)
|
||||
|
||||
const content = result[0].message.content as any[]
|
||||
expect(content[0]).toEqual({ type: 'text', text: 'before' })
|
||||
expect(content[1].type).toBe('text')
|
||||
expect(content[1].text).toContain('media removed')
|
||||
expect(content[2]).toEqual({ type: 'text', text: 'after' })
|
||||
})
|
||||
|
||||
test('removes exactly enough to stay within budget', () => {
|
||||
const messages = [
|
||||
userMsg([makeBase64ImageBlock(2000)]), // removed (cumulative 2000 > remaining budget)
|
||||
userMsg([makeBase64ImageBlock(2000)]), // kept
|
||||
userMsg([makeBase64ImageBlock(2000)]), // kept
|
||||
]
|
||||
|
||||
// Total = 6000, budget = 4000, need to remove 2000 → remove first only
|
||||
const result = shrinkMediaToSizeBudget(messages, 4000)
|
||||
|
||||
const content0 = result[0].message.content as any[]
|
||||
expect(content0[0].type).toBe('text')
|
||||
expect(content0[0].text).toContain('media removed')
|
||||
|
||||
const content1 = result[1].message.content as any[]
|
||||
expect(content1[0].type).toBe('image')
|
||||
|
||||
const content2 = result[2].message.content as any[]
|
||||
expect(content2[0].type).toBe('image')
|
||||
})
|
||||
|
||||
test('handles string content messages gracefully (no array)', () => {
|
||||
const messages = [
|
||||
{ type: 'user', message: { role: 'user', content: 'plain string' }, uuid: crypto.randomUUID(), timestamp: new Date().toISOString() } as unknown as UserMessage,
|
||||
userMsg([makeBase64ImageBlock(3000)]),
|
||||
]
|
||||
|
||||
const result = shrinkMediaToSizeBudget(messages, 5000)
|
||||
expect(result).toBe(messages) // under budget, no change
|
||||
})
|
||||
|
||||
test('handles empty messages array', () => {
|
||||
const result = shrinkMediaToSizeBudget([], 5000)
|
||||
expect(result).toEqual([])
|
||||
})
|
||||
|
||||
test('budget of 0 removes all base64 media', () => {
|
||||
const messages = [
|
||||
userMsg([makeBase64ImageBlock(100)]),
|
||||
userMsg([makeURLImageBlock()]), // URL — never removed
|
||||
userMsg([makeBase64ImageBlock(200)]),
|
||||
]
|
||||
|
||||
const result = shrinkMediaToSizeBudget(messages, 0)
|
||||
|
||||
const content0 = result[0].message.content as any[]
|
||||
expect(content0[0].type).toBe('text')
|
||||
expect(content0[0].text).toContain('media removed')
|
||||
|
||||
// URL image is always preserved
|
||||
const content1 = result[1].message.content as any[]
|
||||
expect(content1[0].type).toBe('image')
|
||||
expect(content1[0].source.type).toBe('url')
|
||||
|
||||
const content2 = result[2].message.content as any[]
|
||||
expect(content2[0].type).toBe('text')
|
||||
expect(content2[0].text).toContain('media removed')
|
||||
})
|
||||
})
|
||||
Loading…
x
Reference in New Issue
Block a user