Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
130 changes: 83 additions & 47 deletions packages/opencode/src/provider/image-files.ts
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ type RunCache = {
ready: Map<string, Uploaded>
busy: boolean
cleanupQueued: boolean
cleaning?: Promise<boolean>
refreshed: Set<string>
retired: Uploaded[]
waiting: Array<() => void>
Expand Down Expand Up @@ -102,56 +103,91 @@ export function withImageFiles(model: LanguageModelV3, options: Options) {
return { bytes, key: `${mediaType}:${new Bun.CryptoHasher("sha256").update(bytes).digest("hex")}` }
}

async function cleanup(signal?: AbortSignal, preserveRegistry = false) {
const pinned = new Set([...state.leases].flatMap((lease) => [...lease.files]))
if (state.closed) {
for (const [key, pending] of state.cache) {
if ([...state.leases].some((lease) => lease.keys.has(key))) continue
if (state.retired.length >= 300) break
const file = await pending.catch(() => undefined)
state.cache.delete(key)
state.ready.delete(key)
state.refreshed.delete(key)
if (file) state.retired.push(file)
}
// One shared drain prevents concurrent wrappers from multiplying DELETE work.
// Its per-file timeout is independent of a caller's bounded foreground wait.
async function cleanup(signal?: AbortSignal, preserveRegistry = false, budget = { remaining: 1000 }) {
const forget = () => {
if (
state.closed &&
!state.cache.size &&
!state.retired.length &&
!state.leases.size &&
cacheID &&
runCaches.get(cacheID) === state
)
runCaches.delete(cacheID)
}
// Cleanup is best effort with one total deadline, including retries. A
// deletion outage must not hold the request lock for one timeout per file.
const deadline = AbortSignal.any([AbortSignal.timeout(1000), ...(signal ? [signal] : [])])
const removed = new Set<Uploaded>()
let attempted = 0
for (const file of state.retired) {
if (pinned.has(file.file_id) || (state.chain && !state.closed)) continue
if (file.expires_at <= now() / 1000) {
removed.add(file)
continue
}
if (deadline.aborted || attempted++ >= (state.closed ? 300 : 8)) break
const response = await request(
`${endpoint}/${encodeURIComponent(file.file_id)}?model=${encodeURIComponent(model.modelId)}`,
{
method: "DELETE",
headers: { ...routeHeaders, "x-bu-image-scope": scope! },
signal: deadline,
},
).catch(() => undefined)
if (!response?.ok && response?.status !== 404) break
removed.add(file)
const joined = Boolean(state.cleaning)
if (!state.cleaning) {
state.cleaning = (async () => {
for (;;) {
if (state.closed) {
for (const [key, file] of state.ready) {
if ([...state.leases].some((lease) => lease.keys.has(key))) continue
if (state.retired.length >= 300) break
state.cache.delete(key)
state.ready.delete(key)
state.refreshed.delete(key)
state.retired.push(file)
}
}
const pinned = new Set([...state.leases].flatMap((lease) => [...lease.files]))
const batch = state.retired
.filter((file) => !pinned.has(file.file_id) && (!state.chain || state.closed))
.slice(0, 16)
if (!batch.length) return true
const removed = await Promise.all(
batch.map(async (file) => {
if (file.expires_at <= now() / 1000) return file
const response = await request(
`${endpoint}/${encodeURIComponent(file.file_id)}?model=${encodeURIComponent(model.modelId)}`,
{
method: "DELETE",
headers: { ...routeHeaders, "x-bu-image-scope": scope! },
signal: AbortSignal.timeout(1000),
},
).catch(() => undefined)
await response?.body?.cancel().catch(() => {})
return response?.ok || response?.status === 404 ? file : undefined
}),
)
// Filter the current queue: other requests can retire files while a
// batch is in flight. Failed tombstones survive for an explicit retry.
const done = new Set(removed.filter((file) => file !== undefined))
state.retired = state.retired.filter((file) => !done.has(file))
if (done.size !== batch.length) return false
}
})().finally(() => {
state.cleaning = undefined
// A caller reopening this cache holds the lock and keeps its registry
// entry stable until it has either reopened it or reported backpressure.
if (!state.busy) forget()
})
}
state.retired = state.retired.filter((file) => !removed.has(file))
if (
!preserveRegistry &&
state.closed &&
!state.cache.size &&
!state.retired.length &&
cacheID &&
runCaches.get(cacheID) === state
)
runCaches.delete(cacheID)
const started = performance.now()
const result = await new Promise<boolean | undefined>((resolve) => {
const finish = (result?: boolean) => {
clearTimeout(timer)
signal?.removeEventListener("abort", abort)
resolve(result)
}
const abort = () => finish()
const timer = setTimeout(abort, Math.max(0, budget.remaining))
signal?.addEventListener("abort", abort, { once: true })
void state.cleaning!.then(finish, abort)
if (signal?.aborted) abort()
})
budget.remaining = Math.max(0, budget.remaining - (performance.now() - started))
// An explicit retry may have joined the tail of a failed drain. Use only
// its remaining budget to start one fresh attempt after that drain ends.
if (joined && result === false && budget.remaining > 0 && !signal?.aborted)
await cleanup(signal, preserveRegistry, budget)
if (!preserveRegistry) forget()
}

async function prepare(params: LanguageModelV3CallOptions, lease: Lease): Promise<LanguageModelV3CallOptions> {
if (!enabled) return params
const cleanupBudget = { remaining: 1000 }
const stateful =
!anthropic &&
(typeof params.providerOptions?.openai?.previousResponseId === "string" ||
Expand Down Expand Up @@ -223,7 +259,7 @@ export function withImageFiles(model: LanguageModelV3, options: Options) {
if (state.closed)
await locked(async () => {
if (!state.closed) return
await cleanup(signal, true)
await cleanup(signal, true, cleanupBudget)
if (state.cache.size || state.retired.length || state.leases.size)
throw new Error("Screenshot file cache is closing; retry cleanup first")
state.closed = false
Expand Down Expand Up @@ -273,7 +309,7 @@ export function withImageFiles(model: LanguageModelV3, options: Options) {
)
lease.keys = keys
state.leases.add(lease)
if (!cached) await cleanup(signal)
if (!cached) await cleanup(signal, false, cleanupBudget)
const missing = [...keys].filter((key) => !state.cache.has(key)).length
for (const key of state.cache.keys()) {
if (state.cache.size + missing <= 300) break
Expand All @@ -291,7 +327,7 @@ export function withImageFiles(model: LanguageModelV3, options: Options) {
// files. Keep explicit stateful chains until close; never assume a copy.
if (state.chain) for (const key of keys) state.chainKeys.add(key)
// Reclaim files newly retired by eviction, after the pre-eviction backlog pass.
if (!cached) await cleanup(signal)
if (!cached) await cleanup(signal, false, cleanupBudget)
const headers = { ...currentHeaders, "x-bu-image-scope": scope! }
const proofs: Record<string, { expires_at: number; signature: string }> = {}
async function upload(data: string | Uint8Array | URL, mediaType: string) {
Expand Down
97 changes: 89 additions & 8 deletions packages/opencode/test/provider/image-files-lifecycle.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,10 @@ function fixture(options: { legacy?: boolean; anthropic?: boolean } = {}) {
deleteStatus: 0,
deleteGate: undefined as Promise<void> | undefined,
deleteStarted: 0,
deleteDelay: 0,
deleteActive: 0,
deletePeak: 0,
uploadDelay: 0,
hangDeletes: false,
uploadGate: undefined as Promise<void> | undefined,
uploads: 0,
Expand Down Expand Up @@ -50,17 +54,25 @@ function fixture(options: { legacy?: boolean; anthropic?: boolean } = {}) {
}
if (req.method === "DELETE") {
state.deleteStarted++
await state.deleteGate
if (state.hangDeletes) return new Promise<Response>(() => {})
if (state.deleteStatus) return new Response("failed", { status: state.deleteStatus })
if (state.deleteFailures-- > 0) return new Response("unavailable", { status: 503 })
const id = new URL(req.url).pathname.split("/").at(-1)!
state.deleted.push(id)
state.live.delete(id)
return Response.json({ deleted: true })
state.deleteActive++
state.deletePeak = Math.max(state.deletePeak, state.deleteActive)
try {
await Bun.sleep(state.deleteDelay)
await state.deleteGate
if (state.hangDeletes) return new Promise<Response>(() => {})
if (state.deleteStatus) return new Response("failed", { status: state.deleteStatus })
if (state.deleteFailures-- > 0) return new Response("unavailable", { status: 503 })
const id = new URL(req.url).pathname.split("/").at(-1)!
state.deleted.push(id)
state.live.delete(id)
return Response.json({ deleted: true })
} finally {
state.deleteActive--
}
}
state.uploads++
const upload = state.uploads
await Bun.sleep(state.uploadDelay)
await state.uploadGate
if (state.uploadFailures-- > 0) return new Response("unavailable", { status: 503 })
const id = `${options.anthropic ? "file_" : "file-"}${upload}`
Expand Down Expand Up @@ -464,3 +476,72 @@ test("last stream ending during close cleanup schedules a final cleanup pass", a
await f.close()
}
})

for (const latency of [false, true]) {
test(`healthy cleanup keeps up with 1000 lifetime images and 100 active (latency=${latency})`, async () => {
const f = fixture()
f.state.deleteDelay = latency ? 200 : 0
f.state.uploadDelay = latency ? 20 : 0
try {
for (let turn = 0; turn < 10; turn++) {
await f.model.doGenerate(images(turn * 100, 100))
expect(f.state.live.size).toBeLessThanOrEqual(600)
if (latency) await Bun.sleep(50)
}
expect(f.state.uploads).toBe(1000)
expect(f.state.calls).toHaveLength(10)
await f.model.closeImageFiles()
const deadline = Date.now() + 6000
while (f.state.live.size && Date.now() < deadline) await Bun.sleep(10)
expect(f.state.live.size).toBe(0)
expect(f.state.deleted).toHaveLength(1000)
expect(new Set(f.state.deleted).size).toBe(1000)
expect(f.state.deletePeak).toBeLessThanOrEqual(16)
} finally {
await f.close()
}
}, 30000)
}

test("preparation shares one second of cleanup wait across eviction passes", async () => {
const f = fixture()
try {
await f.model.doGenerate(images(0, 300))
f.state.deleteStatus = 400
await f.model.doGenerate(images(300, 100))
f.state.deleteStatus = 0
f.state.deleteDelay = 200
const started = performance.now()
const result = await f.model.doStream(images(400, 100))
expect(performance.now() - started).toBeLessThan(1900)
expect(f.state.uploads).toBe(500)
expect(f.state.deletePeak).toBeLessThanOrEqual(16)
await result.stream.cancel()
await f.model.closeImageFiles()
const deadline = Date.now() + 7000
while (f.state.live.size && Date.now() < deadline) await Bun.sleep(10)
expect(f.state.live.size).toBe(0)
} finally {
await f.close()
}
}, 15000)

test("background cleanup stops on failure and keeps tombstones for explicit recovery", async () => {
const f = fixture()
try {
await f.model.doGenerate(images(0, 100))
f.state.deleteStatus = 400
await f.model.closeImageFiles()
const attempts = f.state.deleteStarted
expect(attempts).toBeGreaterThan(0)
expect(attempts).toBeLessThanOrEqual(16)
await Bun.sleep(100)
expect(f.state.deleteStarted).toBe(attempts)
expect(f.state.live.size).toBe(100)
f.state.deleteStatus = 0
await f.model.closeImageFiles()
expect(f.state.live.size).toBe(0)
} finally {
await f.close()
}
})
Loading