fix: stream auto-reconnect and preserve scene entities during disconnect

.changeset/fix-scene-drops.md

@@ -0,0 +1,5 @@
+---
+'@viamrobotics/motion-tools': patch
+---
+
+fix: stream auto-reconnect for draw service and preserve scene entities during disconnect

package.json

@@ -485,7 +485,7 @@
 		"@typescript-eslint/parser": "8.56.1",
 		"@viamrobotics/prime-core": "0.1.5",
 		"@viamrobotics/sdk": "0.69.0",
-		"@viamrobotics/svelte-sdk": "1.2.2",
+		"@viamrobotics/svelte-sdk": "1.2.3",
 		"@viamrobotics/tweakpane-config": "0.1.1",
 		"@vitest/browser-playwright": "4.1.9",
 		"@vitest/coverage-v8": "4.1.9",

src/lib/__tests__/retry-stream.spec.ts

@@ -0,0 +1,196 @@
+import { describe, expect, it, vi } from 'vitest'
+
+import { retryStream } from '../retry-stream'
+
+describe('retryStream', () => {
+	it('calls run and resolves when run succeeds', async () => {
+		const run = vi.fn().mockResolvedValue(undefined)
+		const controller = new AbortController()
+
+		// run resolves once, retryStream will call it again — abort after first call
+		run.mockImplementation(async () => {
+			controller.abort()
+		})
+
+		await retryStream(run, controller.signal)
+
+		expect(run).toHaveBeenCalledTimes(1)
+	})
+
+	it('retries when run throws', async () => {
+		vi.useFakeTimers()
+
+		const controller = new AbortController()
+		let callCount = 0
+
+		const run = vi.fn().mockImplementation(async () => {
+			callCount++
+			if (callCount < 3) {
+				throw new Error('stream error')
+			}
+			controller.abort()
+		})
+
+		const promise = retryStream(run, controller.signal)
+		// Advance through the backoff delays
+		await vi.advanceTimersByTimeAsync(1_000)
+		await vi.advanceTimersByTimeAsync(2_000)
+
+		await promise
+
+		expect(run).toHaveBeenCalledTimes(3)
+
+		vi.useRealTimers()
+	})
+
+	it('calls onRetry but does not retry run when signal is aborted during backoff', async () => {
+		vi.useFakeTimers()
+
+		const controller = new AbortController()
+		const run = vi.fn().mockRejectedValue(new Error('stream error'))
+		const onRetry = vi.fn()
+
+		const promise = retryStream(run, controller.signal, onRetry)
+
+		// First call fails immediately, then waits for backoff
+		await vi.advanceTimersByTimeAsync(0)
+		expect(run).toHaveBeenCalledTimes(1)
+
+		// Abort during backoff wait
+		controller.abort()
+		await vi.advanceTimersByTimeAsync(1_000)
+
+		await promise
+
+		// Should have called onRetry once, but not retried run
+		expect(onRetry).toHaveBeenCalledTimes(1)
+		expect(run).toHaveBeenCalledTimes(1)
+
+		vi.useRealTimers()
+	})
+
+	it('calls onRetry with the current delay', async () => {
+		vi.useFakeTimers()
+
+		const controller = new AbortController()
+		let callCount = 0
+
+		const run = vi.fn().mockImplementation(async () => {
+			callCount++
+			if (callCount < 3) {
+				throw new Error('stream error')
+			}
+			controller.abort()
+		})
+
+		const onRetry = vi.fn()
+		const promise = retryStream(run, controller.signal, onRetry)
+
+		await vi.advanceTimersByTimeAsync(1_000)
+		await vi.advanceTimersByTimeAsync(2_000)
+
+		await promise
+
+		expect(onRetry).toHaveBeenCalledTimes(2)
+		expect(onRetry).toHaveBeenNthCalledWith(1, 1_000)
+		expect(onRetry).toHaveBeenNthCalledWith(2, 2_000)
+
+		vi.useRealTimers()
+	})
+
+	it('does not call onRetry and restarts immediately on clean stream end', async () => {
+		const controller = new AbortController()
+		let callCount = 0
+
+		const run = vi.fn().mockImplementation(async () => {
+			callCount++
+			if (callCount === 1) return // clean end — server closed the stream
+			controller.abort()
+		})
+
+		const onRetry = vi.fn()
+		await retryStream(run, controller.signal, onRetry)
+
+		expect(run).toHaveBeenCalledTimes(2)
+		expect(onRetry).not.toHaveBeenCalled()
+	})
+
+	it('caps delay at MAX_DELAY_MS (30s) and does not double beyond it', async () => {
+		vi.useFakeTimers()
+
+		const controller = new AbortController()
+		let callCount = 0
+
+		const run = vi.fn().mockImplementation(async () => {
+			callCount++
+			if (callCount <= 6) throw new Error('stream error')
+			controller.abort()
+		})
+
+		const onRetry = vi.fn()
+		const promise = retryStream(run, controller.signal, onRetry)
+
+		// Advance through each exponential backoff step
+		await vi.advanceTimersByTimeAsync(1_000) // delay: 1000
+		await vi.advanceTimersByTimeAsync(2_000) // delay: 2000
+		await vi.advanceTimersByTimeAsync(4_000) // delay: 4000
+		await vi.advanceTimersByTimeAsync(8_000) // delay: 8000
+		await vi.advanceTimersByTimeAsync(16_000) // delay: 16000
+		await vi.advanceTimersByTimeAsync(30_000) // delay: capped at 30000 (not 32000)
+
+		await promise
+
+		expect(onRetry).toHaveBeenCalledTimes(6)
+		expect(onRetry).toHaveBeenNthCalledWith(5, 16_000)
+		expect(onRetry).toHaveBeenNthCalledWith(6, 30_000)
+
+		vi.useRealTimers()
+	})
+
+	it('does not call run when signal is already aborted before retryStream is called', async () => {
+		const controller = new AbortController()
+		controller.abort()
+
+		const run = vi.fn()
+		await retryStream(run, controller.signal)
+
+		expect(run).not.toHaveBeenCalled()
+	})
+
+	it('resets delay after a successful run', async () => {
+		vi.useFakeTimers()
+
+		const controller = new AbortController()
+		let callCount = 0
+
+		const run = vi.fn().mockImplementation(async () => {
+			callCount++
+			// First call: fail
+			if (callCount === 1) throw new Error('fail')
+			// Second call: succeed (stream ended cleanly)
+			if (callCount === 2) return
+			// Third call: fail
+			if (callCount === 3) throw new Error('fail')
+			// Fourth call: abort
+			controller.abort()
+		})
+
+		const onRetry = vi.fn()
+		const promise = retryStream(run, controller.signal, onRetry)
+
+		// First failure + 1s backoff
+		await vi.advanceTimersByTimeAsync(1_000)
+		// Second call succeeds, delay resets. Third call fails, should use 1s again
+		await vi.advanceTimersByTimeAsync(1_000)
+		// Fourth call - abort
+		await vi.advanceTimersByTimeAsync(2_000)
+
+		await promise
+
+		// Both retries should have used 1000ms (reset after success)
+		expect(onRetry).toHaveBeenNthCalledWith(1, 1_000)
+		expect(onRetry).toHaveBeenNthCalledWith(2, 1_000)
+
+		vi.useRealTimers()
+	})
+})

src/lib/hooks/useGeometries.svelte.ts

@@ -222,19 +222,27 @@ export const provideGeometries = (partID: () => string) => {
 			})
 		}
 
-		// Clean up owners whose queries disappeared entirely
+		// Clean up owners whose queries disappeared entirely.
+		// Guard: if ALL queries are gone (activeQueryKeys empty), the machine is likely
+		// temporarily disconnected — preserve entities so they reappear on reconnect.
+		// Only destroy when the partID changed (old-partID entities) or other queries
+		// are still active (connected machine, resource legitimately removed).
+		const anyQueriesActive = activeQueryKeys.size > 0
 		for (const [queryKey, keys] of queryEntityKeys) {
 			if (!activeQueryKeys.has(queryKey)) {
-				for (const key of keys) {
-					const entity = entities.get(key)
-					if (entity && world.has(entity)) {
-						entity.destroy()
+				const queryPartID = queryKey.split(':')[0]!
+				if (queryPartID !== currentPartID || anyQueriesActive) {
+					for (const key of keys) {
+						const entity = entities.get(key)
+						if (entity && world.has(entity)) {
+							entity.destroy()
+						}
+
+						entities.delete(key)
 					}
 
-					entities.delete(key)
+					queryEntityKeys.delete(queryKey)
 				}
-
-				queryEntityKeys.delete(queryKey)
 			}
 		}
 	})

src/lib/hooks/usePointclouds.svelte.ts

@@ -188,13 +188,21 @@ export const providePointclouds = (partID: () => string) => {
 			})
 		}
 
-		// clean up queries that disappeared entirely
+		// clean up queries that disappeared entirely.
+		// Guard: if ALL queries are gone (activeQueryKeys empty), the machine is likely
+		// temporarily disconnected — preserve entities so they reappear on reconnect.
+		// Only destroy when the partID changed (old-partID entities) or other queries
+		// are still active (connected machine, camera legitimately removed).
+		const anyQueriesActive = activeQueryKeys.size > 0
 		for (const [queryKey, entity] of entities) {
 			if (!activeQueryKeys.has(queryKey)) {
-				if (world.has(entity)) {
-					entity.destroy()
+				const queryPartID = queryKey.split(':')[0]!
+				if (queryPartID !== currentPartID || anyQueriesActive) {
+					if (world.has(entity)) {
+						entity.destroy()
+					}
+					entities.delete(queryKey)
 				}
-				entities.delete(queryKey)
 			}
 		}
 	})