From 17a5eb5e081777771bc3f86165aa704c4370bf4f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Wed, 16 Sep 2026 18:25:35 +0200 Subject: [PATCH 1/3] fix(kiloclaw): size snapshot restore volume from snapshot source Fly requires an explicit size_gb for a snapshot restore volume and rejects a size smaller than the snapshot's source volume. The queue worker sized the new volume from the current volume, so restoring a snapshot whose source volume is larger failed with a 400 and reverted the instance to stopped. Resolve the snapshot's source volume size by scanning the app volumes and create the restore volume at the larger of the current and source sizes. Add queue worker tests for both sizing directions and the not-found fallback. --- .../src/queue/snapshot-restore.test.ts | 159 ++++++++++++++++++ .../kiloclaw/src/queue/snapshot-restore.ts | 36 +++- 2 files changed, 193 insertions(+), 2 deletions(-) create mode 100644 services/kiloclaw/src/queue/snapshot-restore.test.ts diff --git a/services/kiloclaw/src/queue/snapshot-restore.test.ts b/services/kiloclaw/src/queue/snapshot-restore.test.ts new file mode 100644 index 0000000000..d854b2530a --- /dev/null +++ b/services/kiloclaw/src/queue/snapshot-restore.test.ts @@ -0,0 +1,159 @@ +import { describe, it, expect, vi, beforeEach } from 'vitest'; +import { handleSnapshotRestoreQueue } from './snapshot-restore'; +import * as fly from '../fly/client'; +import type * as flyClientModule from '../fly/client'; +import type { FlyVolume, FlyVolumeSnapshot } from '../fly/types'; +import type { KiloClawEnv } from '../types'; +import type { SnapshotRestoreMessage } from '../schemas/snapshot-restore'; + +vi.mock('../fly/client', async importOriginal => { + const actual = await importOriginal(); + return { + ...actual, + getVolume: vi.fn(), + listVolumes: vi.fn(), + listVolumeSnapshots: vi.fn(), + createVolume: vi.fn(), + }; +}); + +const flyMock = vi.mocked(fly); +const GIB = 1024 ** 3; + +const MESSAGE: SnapshotRestoreMessage = { + userId: 'user-1', + snapshotId: 'vs-1', + previousVolumeId: 'vol-current', + region: 'ord', + instanceId: 'inst-1', +}; + +function volume(overrides: Partial = {}): FlyVolume { + return { + id: 'vol-current', + name: 'kiloclaw_volume', + state: 'created', + size_gb: 10, + region: 'ord', + attached_machine_id: null, + created_at: '2026-01-01T00:00:00Z', + ...overrides, + }; +} + +function snapshot(overrides: Partial = {}): FlyVolumeSnapshot { + return { + id: 'vs-1', + created_at: '2026-01-01T00:00:00Z', + digest: 'digest', + retention_days: 5, + size: 1, + status: 'created', + volume_size: 10 * GIB, + ...overrides, + }; +} + +function makeStub() { + return { + getStatus: vi.fn().mockResolvedValue({ flyVolumeId: 'vol-current', flyAppName: 'acct-test' }), + stop: vi.fn().mockResolvedValue(undefined), + destroyMachineForRestore: vi.fn().mockResolvedValue(undefined), + getDebugState: vi.fn().mockResolvedValue({ pendingRestoreVolumeId: null }), + setPendingRestoreVolumeId: vi.fn().mockResolvedValue(undefined), + completeSnapshotRestore: vi.fn().mockResolvedValue(undefined), + start: vi.fn().mockResolvedValue(undefined), + }; +} + +function makeEnv(stub: ReturnType): KiloClawEnv { + return { + KILOCLAW_INSTANCE: { + idFromName: (name: string) => name, + get: () => stub, + }, + FLY_API_TOKEN: 'token', + FLY_APP_NAME: 'fallback-app', + } as unknown as KiloClawEnv; +} + +function makeBatch(body: SnapshotRestoreMessage = MESSAGE) { + const message = { body, ack: vi.fn(), retry: vi.fn(), attempts: 1 }; + const batch = { + messages: [message], + } as unknown as Parameters[0]; + return { batch, message }; +} + +beforeEach(() => { + vi.clearAllMocks(); + flyMock.getVolume.mockResolvedValue(volume()); + flyMock.listVolumes.mockResolvedValue([volume(), volume({ id: 'vol-previous', size_gb: 11 })]); + flyMock.listVolumeSnapshots.mockResolvedValue([]); + flyMock.createVolume.mockResolvedValue(volume({ id: 'vol-restored' })); +}); + +describe('handleSnapshotRestoreQueue', () => { + it('sizes the restore volume from the snapshot source when it is larger than the current volume', async () => { + flyMock.listVolumeSnapshots.mockImplementation(async (_config, volumeId) => + volumeId === 'vol-previous' ? [snapshot({ volume_size: 11 * GIB })] : [] + ); + + const { batch, message } = makeBatch(); + await handleSnapshotRestoreQueue(batch, makeEnv(makeStub())); + + expect(flyMock.createVolume).toHaveBeenCalledWith( + expect.anything(), + expect.objectContaining({ snapshot_id: 'vs-1', size_gb: 11 }) + ); + expect(message.ack).toHaveBeenCalled(); + }); + + it('keeps the current volume size when the snapshot source is smaller', async () => { + flyMock.getVolume.mockResolvedValue(volume({ size_gb: 20 })); + flyMock.listVolumeSnapshots.mockResolvedValue([snapshot({ volume_size: 11 * GIB })]); + + const { batch } = makeBatch(); + await handleSnapshotRestoreQueue(batch, makeEnv(makeStub())); + + expect(flyMock.createVolume).toHaveBeenCalledWith( + expect.anything(), + expect.objectContaining({ size_gb: 20 }) + ); + }); + + it('falls back to the current volume size when the snapshot is not found', async () => { + const { batch } = makeBatch(); + await handleSnapshotRestoreQueue(batch, makeEnv(makeStub())); + + expect(flyMock.createVolume).toHaveBeenCalledWith( + expect.anything(), + expect.objectContaining({ size_gb: 10 }) + ); + }); + + it('swaps the volume and starts the machine on success', async () => { + flyMock.listVolumeSnapshots.mockResolvedValue([snapshot({ volume_size: 10 * GIB })]); + flyMock.createVolume.mockResolvedValue(volume({ id: 'vol-restored', region: 'ord' })); + const stub = makeStub(); + + const { batch, message } = makeBatch(); + await handleSnapshotRestoreQueue(batch, makeEnv(stub)); + + expect(stub.setPendingRestoreVolumeId).toHaveBeenCalledWith('vol-restored'); + expect(stub.completeSnapshotRestore).toHaveBeenCalledWith('vol-restored', 'ord'); + expect(stub.start).toHaveBeenCalledWith('user-1', { reason: 'snapshot_restore' }); + expect(message.ack).toHaveBeenCalled(); + }); + + it('acks without restoring when the volume was already swapped', async () => { + const stub = makeStub(); + stub.getStatus.mockResolvedValue({ flyVolumeId: 'vol-other', flyAppName: 'acct-test' }); + + const { batch, message } = makeBatch(); + await handleSnapshotRestoreQueue(batch, makeEnv(stub)); + + expect(flyMock.createVolume).not.toHaveBeenCalled(); + expect(message.ack).toHaveBeenCalled(); + }); +}); diff --git a/services/kiloclaw/src/queue/snapshot-restore.ts b/services/kiloclaw/src/queue/snapshot-restore.ts index aa4827f2bf..d43142d68b 100644 --- a/services/kiloclaw/src/queue/snapshot-restore.ts +++ b/services/kiloclaw/src/queue/snapshot-restore.ts @@ -22,6 +22,37 @@ import { SnapshotRestoreMessageSchema } from '../schemas/snapshot-restore'; import { writeEvent } from '../utils/analytics'; import * as fly from '../fly/client'; +const BYTES_PER_GIB = 1024 ** 3; + +/** + * Resolve the source volume size of a snapshot, in GiB, by scanning the app's volumes. + * + * Fly requires an explicit `size_gb` when creating a volume from a snapshot, and it + * rejects a size smaller than the snapshot's source volume. The snapshot can come from + * a volume other than the current one (for example a retained previous or recovery + * volume), so the restore volume must not be sized from the current volume alone. + */ +async function resolveSnapshotVolumeSizeGb( + flyConfig: { apiToken: string; appName: string }, + snapshotId: string +): Promise { + const volumes = await fly.listVolumes(flyConfig); + const sizes = await Promise.all( + volumes.map(async volume => { + try { + const snapshots = await fly.listVolumeSnapshots(flyConfig, volume.id); + const snapshot = snapshots.find(candidate => candidate.id === snapshotId); + return snapshot ? Math.ceil(snapshot.volume_size / BYTES_PER_GIB) : null; + } catch (err) { + // A volume can disappear between listing and snapshot lookup. + if (fly.isFlyNotFound(err)) return null; + throw err; + } + }) + ); + return sizes.find((size): size is number => size !== null) ?? null; +} + async function createRestoreVolume( flyConfig: { apiToken: string; appName: string }, previousVolumeId: string, @@ -29,15 +60,16 @@ async function createRestoreVolume( region: string ) { const existingVolume = await fly.getVolume(flyConfig, previousVolumeId); + const snapshotVolumeSizeGb = await resolveSnapshotVolumeSizeGb(flyConfig, snapshotId); const newVolume = await fly.createVolume(flyConfig, { name: existingVolume.name, region, snapshot_id: snapshotId, - size_gb: existingVolume.size_gb, + size_gb: Math.max(existingVolume.size_gb, snapshotVolumeSizeGb ?? 0), snapshot_retention: 5, }); console.log( - `[queue] New volume created: id=${newVolume.id} region=${newVolume.region} from snapshot=${snapshotId}` + `[queue] New volume created: id=${newVolume.id} region=${newVolume.region} size=${newVolume.size_gb}GB from snapshot=${snapshotId}` ); return newVolume; } From c7566e44f8f2db727bac63602e9b45bd7bcadd15 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Wed, 16 Sep 2026 18:29:39 +0200 Subject: [PATCH 2/3] fix(kiloclaw): tolerate vanished volumes and bound snapshot scan Address review remarks on the snapshot restore sizing fix: - Match vanished volumes with isFlyMissingVolume so the scan also tolerates the 400 and 422 responses Fly uses, not only 404. - Skip volumes smaller than the current volume, which cannot raise the restore size, to bound the scan as retained volumes accumulate. Add tests for a volume that vanishes mid-scan and for skipped small volumes. --- .../src/queue/snapshot-restore.test.ts | 43 ++++++++++++++++++- .../kiloclaw/src/queue/snapshot-restore.ts | 22 +++++++--- 2 files changed, 58 insertions(+), 7 deletions(-) diff --git a/services/kiloclaw/src/queue/snapshot-restore.test.ts b/services/kiloclaw/src/queue/snapshot-restore.test.ts index d854b2530a..06edbcf75c 100644 --- a/services/kiloclaw/src/queue/snapshot-restore.test.ts +++ b/services/kiloclaw/src/queue/snapshot-restore.test.ts @@ -109,8 +109,9 @@ describe('handleSnapshotRestoreQueue', () => { expect(message.ack).toHaveBeenCalled(); }); - it('keeps the current volume size when the snapshot source is smaller', async () => { + it('keeps the current volume size when the snapshot is smaller than the current volume', async () => { flyMock.getVolume.mockResolvedValue(volume({ size_gb: 20 })); + flyMock.listVolumes.mockResolvedValue([volume({ id: 'vol-current', size_gb: 20 })]); flyMock.listVolumeSnapshots.mockResolvedValue([snapshot({ volume_size: 11 * GIB })]); const { batch } = makeBatch(); @@ -122,6 +123,46 @@ describe('handleSnapshotRestoreQueue', () => { ); }); + it('tolerates a volume that vanishes during the snapshot scan', async () => { + flyMock.listVolumes.mockResolvedValue([ + volume({ id: 'vol-current', size_gb: 10 }), + volume({ id: 'vol-gone', size_gb: 11 }), + volume({ id: 'vol-previous', size_gb: 11 }), + ]); + flyMock.listVolumeSnapshots.mockImplementation(async (_config, volumeId) => { + if (volumeId === 'vol-gone') { + throw new fly.FlyApiError('volume not found', 400, '{"error":"volume not found"}'); + } + return volumeId === 'vol-previous' ? [snapshot({ volume_size: 11 * GIB })] : []; + }); + + const { batch, message } = makeBatch(); + await handleSnapshotRestoreQueue(batch, makeEnv(makeStub())); + + expect(flyMock.createVolume).toHaveBeenCalledWith( + expect.anything(), + expect.objectContaining({ size_gb: 11 }) + ); + expect(message.ack).toHaveBeenCalled(); + }); + + it('does not scan volumes smaller than the current volume', async () => { + flyMock.listVolumes.mockResolvedValue([ + volume({ id: 'vol-current', size_gb: 10 }), + volume({ id: 'vol-small', size_gb: 5 }), + ]); + + const { batch } = makeBatch(); + await handleSnapshotRestoreQueue(batch, makeEnv(makeStub())); + + expect(flyMock.listVolumeSnapshots).toHaveBeenCalledWith(expect.anything(), 'vol-current'); + expect(flyMock.listVolumeSnapshots).not.toHaveBeenCalledWith(expect.anything(), 'vol-small'); + expect(flyMock.createVolume).toHaveBeenCalledWith( + expect.anything(), + expect.objectContaining({ size_gb: 10 }) + ); + }); + it('falls back to the current volume size when the snapshot is not found', async () => { const { batch } = makeBatch(); await handleSnapshotRestoreQueue(batch, makeEnv(makeStub())); diff --git a/services/kiloclaw/src/queue/snapshot-restore.ts b/services/kiloclaw/src/queue/snapshot-restore.ts index d43142d68b..f0025b98ca 100644 --- a/services/kiloclaw/src/queue/snapshot-restore.ts +++ b/services/kiloclaw/src/queue/snapshot-restore.ts @@ -25,27 +25,33 @@ import * as fly from '../fly/client'; const BYTES_PER_GIB = 1024 ** 3; /** - * Resolve the source volume size of a snapshot, in GiB, by scanning the app's volumes. + * Resolve the source volume size of a snapshot, in GiB. * * Fly requires an explicit `size_gb` when creating a volume from a snapshot, and it * rejects a size smaller than the snapshot's source volume. The snapshot can come from * a volume other than the current one (for example a retained previous or recovery * volume), so the restore volume must not be sized from the current volume alone. + * + * Only volumes at least as large as the current volume can raise the restore size, so + * smaller volumes are skipped to bound the scan as retained volumes accumulate. */ async function resolveSnapshotVolumeSizeGb( flyConfig: { apiToken: string; appName: string }, - snapshotId: string + snapshotId: string, + currentSizeGb: number ): Promise { const volumes = await fly.listVolumes(flyConfig); + const candidates = volumes.filter(volume => volume.size_gb >= currentSizeGb); const sizes = await Promise.all( - volumes.map(async volume => { + candidates.map(async volume => { try { const snapshots = await fly.listVolumeSnapshots(flyConfig, volume.id); const snapshot = snapshots.find(candidate => candidate.id === snapshotId); return snapshot ? Math.ceil(snapshot.volume_size / BYTES_PER_GIB) : null; } catch (err) { - // A volume can disappear between listing and snapshot lookup. - if (fly.isFlyNotFound(err)) return null; + // A volume can disappear between listing and snapshot lookup. Fly reports a + // vanished volume as 400, 404, or 422 depending on the endpoint. + if (fly.isFlyMissingVolume(err)) return null; throw err; } }) @@ -60,7 +66,11 @@ async function createRestoreVolume( region: string ) { const existingVolume = await fly.getVolume(flyConfig, previousVolumeId); - const snapshotVolumeSizeGb = await resolveSnapshotVolumeSizeGb(flyConfig, snapshotId); + const snapshotVolumeSizeGb = await resolveSnapshotVolumeSizeGb( + flyConfig, + snapshotId, + existingVolume.size_gb + ); const newVolume = await fly.createVolume(flyConfig, { name: existingVolume.name, region, From b2285bf5f7b7c20009fa57ca11d52747610ff983 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Wed, 16 Sep 2026 18:30:24 +0200 Subject: [PATCH 3/3] fix(kiloclaw): keep the 404 check alongside isFlyMissingVolume isFlyMissingVolume requires both a missing-volume status code and a known marker, so a bare 404 without a marker would have thrown. Match the platform route pattern and accept either check. --- services/kiloclaw/src/queue/snapshot-restore.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/services/kiloclaw/src/queue/snapshot-restore.ts b/services/kiloclaw/src/queue/snapshot-restore.ts index f0025b98ca..831cfe22e4 100644 --- a/services/kiloclaw/src/queue/snapshot-restore.ts +++ b/services/kiloclaw/src/queue/snapshot-restore.ts @@ -51,7 +51,7 @@ async function resolveSnapshotVolumeSizeGb( } catch (err) { // A volume can disappear between listing and snapshot lookup. Fly reports a // vanished volume as 400, 404, or 422 depending on the endpoint. - if (fly.isFlyMissingVolume(err)) return null; + if (fly.isFlyNotFound(err) || fly.isFlyMissingVolume(err)) return null; throw err; } })