diff --git a/.changeset/quiet-reconnect-claim.md b/.changeset/quiet-reconnect-claim.md new file mode 100644 index 0000000000..dfaf4772c3 --- /dev/null +++ b/.changeset/quiet-reconnect-claim.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-sandbox': patch +--- + +A reconnect during a long tool call no longer writes the rest of a durable sandbox run twice. The host that started the run holds the driver claim until the run ends, so the reconnect tails the log. diff --git a/docs/config.json b/docs/config.json index f75826e68f..fb945e5268 100644 --- a/docs/config.json +++ b/docs/config.json @@ -1509,7 +1509,8 @@ { "label": "Takeover Internals", "to": "sandbox/takeover-advanced", - "addedAt": "2026-10-21" + "addedAt": "2026-10-21", + "updatedAt": "2026-11-09" }, { "label": "Reaping & Retention", diff --git a/docs/sandbox/takeover-advanced.md b/docs/sandbox/takeover-advanced.md index 3256503b80..11c08fbc8c 100644 --- a/docs/sandbox/takeover-advanced.md +++ b/docs/sandbox/takeover-advanced.md @@ -106,6 +106,8 @@ Only one host can write a run. The client has no safety net below its offset de- 2. **An epoch.** Each claim bumps `RunRecord.driverEpoch`. The log re-reads it from time to time and refuses to append once a higher epoch exists. This covers a lock that renews slower than the run appends, or one whose signal never fires. 3. **Quiet time.** Before its first append, the new host waits for the stored log to stop growing. So it sees an old host that still writes, and does not race it. The wait is `DEFAULT_FENCE_QUIET_MS` (5 seconds). Change it with `fenceQuietMs` on `sandboxRunDriver`. If the log never goes quiet, the drive fails and appends nothing. +The host that starts the run holds this same claim until that run ends. A reconnect while that host still writes only tails the log. The first host stays the only writer. A tool call can add no events for longer than the quiet window. The claim blocks a second writer for that whole call. When the first host is gone, the quiet window applies again. + ### Both channels are fenced A run's facts live in two places: its event log and its record. If only the log is fenced, a superseded host turns its refused append into a terminal `runs.update`. The record then says `'failed'` for a run the new host streams without problems. Then anything that checks status (`isTerminalRunStatus`, `findActiveRun`, a status poller, the reaper) believes that a live run died. diff --git a/packages/ai-sandbox/src/claim.ts b/packages/ai-sandbox/src/claim.ts index 54d1c1b297..4f424d65fc 100644 --- a/packages/ai-sandbox/src/claim.ts +++ b/packages/ai-sandbox/src/claim.ts @@ -97,6 +97,20 @@ export function runDriverLockKey(runId: string): string { return `run-driver:${runId}` } +/** + * Run ids whose driver claim is held in THIS process. + * + * `InMemoryLockStore.withLock` waits forever if the holder calls it again, so + * a chat that is already inside {@link withRunClaim} must not take the claim + * a second time. A different process has its own set and its own lock. + */ +const locallyHeldDriverClaims = new Set() + +/** True when this process already holds `runId`'s driver claim. */ +export function runDriverClaimHeldLocally(runId: string): boolean { + return locallyHeldDriverClaims.has(runId) +} + /** The claim was never acquired, so the caller must not drive the run. */ export class RunClaimNotAcquiredError extends Error { constructor( @@ -185,7 +199,45 @@ export async function withRunClaim( runId, epoch, }) - return fn({ runId, epoch, signal }) + const claim: RunClaim = { runId, epoch, signal } + locallyHeldDriverClaims.add(runId) + try { + return await fn(claim) + } finally { + locallyHeldDriverClaims.delete(runId) + } + }) +} + +/** + * Hold the driver claim until `release` runs. + * + * The fresh producer uses this. Its delivery sink keeps appending after a + * client disconnect, and that sink does not take the claim. A takeover that + * only waits for a quiet log will start a second writer during a tool call + * that appends nothing for longer than {@link DEFAULT_FENCE_QUIET_MS}. + */ +export function holdRunDriverClaim( + options: WithRunClaimOptions, +): Promise<() => void> { + let release = (): void => {} + const acquired = new Promise((resolve, reject) => { + void withRunClaim(options, () => { + return new Promise((done) => { + let released = false + release = () => { + if (released) return + released = true + done() + } + resolve() + }) + }).catch((error: unknown) => { + reject(error) + }) + }) + return acquired.then(() => () => { + release() }) } diff --git a/packages/ai-sandbox/src/middleware.ts b/packages/ai-sandbox/src/middleware.ts index caf0791bae..e1a0c49450 100644 --- a/packages/ai-sandbox/src/middleware.ts +++ b/packages/ai-sandbox/src/middleware.ts @@ -44,6 +44,7 @@ import { computeWorkspaceHash } from './key' import { buildFileHookEvent, resolveFileEvents } from './file-diff' import { ProjectionCapability, provideWorkspaceProjection } from './projection' import { resolveAllSecrets, resolveSecret } from './secrets' +import { holdRunDriverClaim, runDriverClaimHeldLocally } from './claim' import { createToolHistoryRecorder, stripObservedToolCalls, @@ -135,6 +136,11 @@ interface SandboxRunState { pendingDiffs: Array> /** Logger captured at setup, so terminal hooks can log watcher teardown. */ logger?: InternalLogger + /** + * Ends the fresh producer's driver claim. Absent on an attach, and absent + * when this chat is already inside a claim (the takeover's own `drive`). + */ + releaseDriverClaim?: () => void /** * Durability resolved once at setup (absent when the run is not durable), so * `onAbort` cannot reach a different verdict than the one `setup` published @@ -454,6 +460,13 @@ function buildEnsureCtx( } } +function releaseDriverClaim(state: SandboxRunState | undefined): void { + if (state === undefined || state.releaseDriverClaim === undefined) return + const release = state.releaseDriverClaim + state.releaseDriverClaim = undefined + release() +} + /** * Dispatch a sandbox file event to the per-type hooks declared on the * definition. Errors in individual hooks are swallowed so one bad hook @@ -461,6 +474,7 @@ function buildEnsureCtx( * a throwing hook is observable (matching the run-scoped path in the engine * and the behavior the observability docs promise). */ + async function dispatchDefinitionHooks( hooks: SandboxHooks | undefined, event: SandboxFileHookEvent, @@ -667,12 +681,14 @@ export function withSandbox( // `createOrResume` is idempotent and never resurrects a finished run, so // persistence's own later call stays correct and simply finds this record. if (durability !== undefined) { + let durableRecordReady = false try { await durability.runs.createOrResume({ runId: ctx.runId, threadId: ctx.threadId, startedAt: Date.now(), }) + durableRecordReady = true } catch (error) { // Best-effort: a store blip must not stop a run that is otherwise fine. // The run is simply invisible until persistence's own `onConfig` call. @@ -715,6 +731,25 @@ export function withSandbox( error, }) } + + // The delivery sink keeps appending after a disconnect, and it never + // takes this lock. Hold the claim for the whole fresh run so a takeover + // cannot start a second writer while a tool call leaves the log quiet + // for longer than `fenceQuietMs`. An attach, and a chat that is already + // inside a claim, must not take it again: `withLock` is not re-entrant. + if ( + durableRecordReady && + durability.attach !== true && + ensureCtx.locks !== undefined && + !runDriverClaimHeldLocally(ctx.runId) + ) { + state.releaseDriverClaim = await holdRunDriverClaim({ + runs: durability.runs, + locks: ensureCtx.locks, + runId: ctx.runId, + ...(logger === undefined ? {} : { logger }), + }) + } } // SUBSCRIBE BEFORE `ensure`, for the same reason the state is registered @@ -1045,248 +1080,262 @@ export function withSandbox( async onFinish(ctx) { const state = runState.get(ctx) - // `chat()` also calls `onFinish` when the run pauses at an interrupt. The - // run is not complete, so keep the sandbox: no after-run snapshot and no - // `destroyOnComplete`. `onChunk` already stopped the watcher and the lease. - if (!state || state.interrupted) return - const { handle, ensureCtx } = state + try { + // `chat()` also calls `onFinish` when the run pauses at an interrupt. The + // run is not complete, so keep the sandbox: no after-run snapshot and no + // `destroyOnComplete`. `onChunk` already stopped the watcher and the lease. + if (!state || state.interrupted) return + const { handle, ensureCtx } = state - // Last chance before persistence writes the transcript. Only matters if a - // config sync landed after the final tool chunk; the recorder is idempotent, so - // in the normal case this changes nothing. - state.toolHistory.reconcile(ctx) + // Last chance before persistence writes the transcript. Only matters if a + // config sync landed after the final tool chunk; the recorder is idempotent, so + // in the normal case this changes nothing. + state.toolHistory.reconcile(ctx) - await drainWatcher(state, 'finish') + await drainWatcher(state, 'finish') - let primaryError: unknown - try { - const snapshotCaptureTask = Promise.resolve().then( - async (): Promise => { - const config = state.snapshotConfig - const runtime = state.snapshotRuntime - const lease = state.snapshotLease - if (!config || !runtime || !handle || !lease) { - if (state.snapshotLost) throw state.snapshotLost - return - } - if (!canPublishPortableSnapshot(state, lease)) return + let primaryError: unknown + try { + const snapshotCaptureTask = Promise.resolve().then( + async (): Promise => { + const config = state.snapshotConfig + const runtime = state.snapshotRuntime + const lease = state.snapshotLease + if (!config || !runtime || !handle || !lease) { + if (state.snapshotLost) throw state.snapshotLost + return + } + if (!canPublishPortableSnapshot(state, lease)) return - await runtime.completion.waitForRunCompletion() - if (!canPublishPortableSnapshot(state, lease)) return + await runtime.completion.waitForRunCompletion() + if (!canPublishPortableSnapshot(state, lease)) return - const conversation = - await runtime.persistence.stores.messages.loadThread(ctx.threadId) - if (!canPublishPortableSnapshot(state, lease)) return + const conversation = + await runtime.persistence.stores.messages.loadThread( + ctx.threadId, + ) + if (!canPublishPortableSnapshot(state, lease)) return - const files = await captureSandboxFiles( - handle, - { - blobs: config.persistence.stores.blobs, - workspaceRoot: - definition.workspace?.root ?? DEFAULT_WORKSPACE_ROOT, - }, - state.snapshotPolicy, - definition.workspace?.secrets !== undefined - ? resolveAllSecrets(definition.workspace.secrets) - : {}, - ) - if (!canPublishPortableSnapshot(state, lease)) return + const files = await captureSandboxFiles( + handle, + { + blobs: config.persistence.stores.blobs, + workspaceRoot: + definition.workspace?.root ?? DEFAULT_WORKSPACE_ROOT, + }, + state.snapshotPolicy, + definition.workspace?.secrets !== undefined + ? resolveAllSecrets(definition.workspace.secrets) + : {}, + ) + if (!canPublishPortableSnapshot(state, lease)) return - const artifacts = await captureSandboxArtifacts( - { - blobs: config.persistence.stores.blobs, - artifacts: config.persistence.stores.artifacts, - }, - ctx.threadId, - definition.workspace?.secrets !== undefined - ? resolveAllSecrets(definition.workspace.secrets) - : {}, - ) - if (!canPublishPortableSnapshot(state, lease)) return + const artifacts = await captureSandboxArtifacts( + { + blobs: config.persistence.stores.blobs, + artifacts: config.persistence.stores.artifacts, + }, + ctx.threadId, + definition.workspace?.secrets !== undefined + ? resolveAllSecrets(definition.workspace.secrets) + : {}, + ) + if (!canPublishPortableSnapshot(state, lease)) return - const parentCheckpointId = await config.checkpoints.getHead( - ctx.threadId, - ) - if (!canPublishPortableSnapshot(state, lease)) return + const parentCheckpointId = await config.checkpoints.getHead( + ctx.threadId, + ) + if (!canPublishPortableSnapshot(state, lease)) return - try { - await config.checkpoints.append({ - checkpoint: { - id: `checkpoint-${ctx.runId}`, - threadId: ctx.threadId, - parentCheckpointId, - createdAt: Date.now(), - reason: 'automatic', - sourceRunId: ctx.runId, - files: files.files, - conversation, - artifacts, - }, - expectedHeadId: parentCheckpointId, - writer: lease, - }) - } catch (error) { + try { + await config.checkpoints.append({ + checkpoint: { + id: `checkpoint-${ctx.runId}`, + threadId: ctx.threadId, + parentCheckpointId, + createdAt: Date.now(), + reason: 'automatic', + sourceRunId: ctx.runId, + files: files.files, + conversation, + artifacts, + }, + expectedHeadId: parentCheckpointId, + writer: lease, + }) + } catch (error) { + if (state.snapshotLost) throw state.snapshotLost + throw error + } if (state.snapshotLost) throw state.snapshotLost - throw error - } - if (state.snapshotLost) throw state.snapshotLost - }, - ) - state.snapshotCaptureTask = snapshotCaptureTask - try { - await snapshotCaptureTask - } finally { - if (state.snapshotCaptureTask === snapshotCaptureTask) - state.snapshotCaptureTask = undefined - } + }, + ) + state.snapshotCaptureTask = snapshotCaptureTask + try { + await snapshotCaptureTask + } finally { + if (state.snapshotCaptureTask === snapshotCaptureTask) + state.snapshotCaptureTask = undefined + } - const lifecycle = definition.lifecycle + const lifecycle = definition.lifecycle - // `handle` is absent only if `setup` never got past `definition.ensure`, in - // which case there is no sandbox to snapshot. - if ( - lifecycle?.snapshot === 'after-run' && - handle?.capabilities.snapshots && - handle.snapshot - ) { - const snapshot = await handle.snapshot(`after-run-${ctx.runId}`) - const store = ensureCtx.store - if (store) { - const key = definition.key(ensureCtx) - const existing = await store.get(key) - if (existing) { - await store.upsert({ - ...existing, - latestSnapshotId: snapshot.id, - updatedAt: Date.now(), - }) + // `handle` is absent only if `setup` never got past `definition.ensure`, in + // which case there is no sandbox to snapshot. + if ( + lifecycle?.snapshot === 'after-run' && + handle?.capabilities.snapshots && + handle.snapshot + ) { + const snapshot = await handle.snapshot(`after-run-${ctx.runId}`) + const store = ensureCtx.store + if (store) { + const key = definition.key(ensureCtx) + const existing = await store.get(key) + if (existing) { + await store.upsert({ + ...existing, + latestSnapshotId: snapshot.id, + updatedAt: Date.now(), + }) + } } } - } - if (lifecycle?.destroyOnComplete) { - await definition.destroy(ensureCtx) - await definition.hooks?.onDestroy?.() - } - } catch (error) { - primaryError = error - if (definition.lifecycle?.destroyOnComplete) { - try { + if (lifecycle?.destroyOnComplete) { await definition.destroy(ensureCtx) await definition.hooks?.onDestroy?.() - } catch (cleanupError) { - state.logger?.warn( - 'sandbox destroy after terminal failure failed', - { - runId: ctx.runId, - phase: 'finish', - error: cleanupError, - }, - ) + } + } catch (error) { + primaryError = error + if (definition.lifecycle?.destroyOnComplete) { + try { + await definition.destroy(ensureCtx) + await definition.hooks?.onDestroy?.() + } catch (cleanupError) { + state.logger?.warn( + 'sandbox destroy after terminal failure failed', + { + runId: ctx.runId, + phase: 'finish', + error: cleanupError, + }, + ) + } } } - } - let snapshotCleanupError: unknown - try { - await stopSnapshotLease(state, { closePortable: true }) - } catch (error) { - snapshotCleanupError = error - } - if (primaryError !== undefined) { - if (snapshotCleanupError !== undefined) - state.logger?.warn('sandbox snapshot writer release failed', { - runId: ctx.runId, - phase: 'finish', - error: snapshotCleanupError, - }) - throw primaryError + let snapshotCleanupError: unknown + try { + await stopSnapshotLease(state, { closePortable: true }) + } catch (error) { + snapshotCleanupError = error + } + if (primaryError !== undefined) { + if (snapshotCleanupError !== undefined) + state.logger?.warn('sandbox snapshot writer release failed', { + runId: ctx.runId, + phase: 'finish', + error: snapshotCleanupError, + }) + throw primaryError + } + if (snapshotCleanupError !== undefined) throw snapshotCleanupError + } finally { + releaseDriverClaim(state) } - if (snapshotCleanupError !== undefined) throw snapshotCleanupError }, async onAbort(ctx, info: AbortInfo) { const state = runState.get(ctx) - if (!state) return - - // First on BOTH branches: a diff still in flight must be drained whether - // the sandbox is about to be destroyed or merely detached, or the final - // file's diff is dropped. - await drainWatcher(state, 'abort') - let releaseError: unknown try { - await stopSnapshotLease(state, { closePortable: true }) - } catch (error) { - releaseError = error - } + if (!state) return + + // First on BOTH branches: a diff still in flight must be drained whether + // the sandbox is about to be destroyed or merely detached, or the final + // file's diff is dropped. + await drainWatcher(state, 'abort') + let releaseError: unknown + try { + await stopSnapshotLease(state, { closePortable: true }) + } catch (error) { + releaseError = error + } - const durability = state.durability - const cancelled = await cancelIntent( - durability, - ctx.runId, - info.cancelRequested === true, - ) + const durability = state.durability + const cancelled = await cancelIntent( + durability, + ctx.runId, + info.cancelRequested === true, + ) - if ( - durability !== undefined && - !cancelled && - durability.detachOnDisconnect - ) { - // DETACH on the teardown path. Reached when the run is aborted for a - // reason that is NOT an out-of-band cancel while detachable — a genuine - // stop from elsewhere, or a host going down. The ordinary disconnect is - // handled by the disconnect subscriber in `setup`, which does not end the - // run at all. - // - // On a failed record write this branch is ABANDONED for the destroy one - // below, because a rejection here is the worst shape available: the - // verdict is unpublished, so core terminalizes the log and records a - // healthy detached run as failed; `detachedSince`/`sandboxKey` are - // unwritten, so `listReclaimable` can never surface the run and - // `reapDetachedRuns` can never reclaim it. A DESTROYED sandbox beats an - // unreachable one — the same reasoning `drainWatcher` applies to its own - // guarded `stop()`. - if (await recordDetach(definition, state, durability, ctx, 'abort')) { + if ( + durability !== undefined && + !cancelled && + durability.detachOnDisconnect + ) { + // DETACH on the teardown path. Reached when the run is aborted for a + // reason that is NOT an out-of-band cancel while detachable — a genuine + // stop from elsewhere, or a host going down. The ordinary disconnect is + // handled by the disconnect subscriber in `setup`, which does not end the + // run at all. + // + // On a failed record write this branch is ABANDONED for the destroy one + // below, because a rejection here is the worst shape available: the + // verdict is unpublished, so core terminalizes the log and records a + // healthy detached run as failed; `detachedSince`/`sandboxKey` are + // unwritten, so `listReclaimable` can never surface the run and + // `reapDetachedRuns` can never reclaim it. A DESTROYED sandbox beats an + // unreachable one — the same reasoning `drainWatcher` applies to its own + // guarded `stop()`. + if (await recordDetach(definition, state, durability, ctx, 'abort')) { + if (releaseError) throw releaseError + return + } + await definition.destroy(state.ensureCtx) + await definition.hooks?.onDestroy?.() if (releaseError) throw releaseError return } + + // ALWAYS tear down on an explicit abort, regardless of `destroyOnComplete`. + // The in-sandbox agent process is not killed by closing its IO stream + // (e.g. a Docker exec survives client disconnect), so the only reliable way + // to stop it — and the token/cost drain of its ongoing API calls — is to + // destroy the sandbox (stop the container/VM). `keepAlive` / + // `destroyOnComplete:false` governs *successful completion*, never cancel. await definition.destroy(state.ensureCtx) await definition.hooks?.onDestroy?.() if (releaseError) throw releaseError - return + } finally { + releaseDriverClaim(state) } - - // ALWAYS tear down on an explicit abort, regardless of `destroyOnComplete`. - // The in-sandbox agent process is not killed by closing its IO stream - // (e.g. a Docker exec survives client disconnect), so the only reliable way - // to stop it — and the token/cost drain of its ongoing API calls — is to - // destroy the sandbox (stop the container/VM). `keepAlive` / - // `destroyOnComplete:false` governs *successful completion*, never cancel. - await definition.destroy(state.ensureCtx) - await definition.hooks?.onDestroy?.() - if (releaseError) throw releaseError }, async onError(ctx, info) { const state = runState.get(ctx) - if (!state) return - - await drainWatcher(state, 'error') - let releaseError: unknown try { - await stopSnapshotLease(state) - } catch (error) { - releaseError = error - } - await definition.hooks?.onError?.(info.error) + if (!state) return - // On failure, only tear down when the lifecycle says so; otherwise leave - // the sandbox for a resumed retry. - if (definition.lifecycle?.destroyOnComplete) { - await definition.destroy(state.ensureCtx) - await definition.hooks?.onDestroy?.() + await drainWatcher(state, 'error') + let releaseError: unknown + try { + await stopSnapshotLease(state) + } catch (error) { + releaseError = error + } + await definition.hooks?.onError?.(info.error) + + // On failure, only tear down when the lifecycle says so; otherwise leave + // the sandbox for a resumed retry. + if (definition.lifecycle?.destroyOnComplete) { + await definition.destroy(state.ensureCtx) + await definition.hooks?.onDestroy?.() + } + if (releaseError) throw releaseError + } finally { + releaseDriverClaim(state) } - if (releaseError) throw releaseError }, }) } diff --git a/packages/ai-sandbox/tests/producer-claim.test.ts b/packages/ai-sandbox/tests/producer-claim.test.ts new file mode 100644 index 0000000000..d455d64965 --- /dev/null +++ b/packages/ai-sandbox/tests/producer-claim.test.ts @@ -0,0 +1,114 @@ +/** + * The fresh producer must hold the run-driver claim until its terminal hook. + * A takeover waits on that claim. A quiet tool call must not let it start. + */ +import { describe, expect, it } from 'vitest' +import { InMemoryRunStore, memoryStream } from '@tanstack/ai' +import { InMemoryLockStore } from '@tanstack/ai/locks' +import { provideSandboxRuntime } from '@tanstack/ai/adapter-internals' +import { defineSandbox } from '../src/sandbox' +import { withSandbox } from '../src/middleware' +import { withRunClaim } from '../src/claim' +import { captureLogger, makeFakeProvider, makeMiddlewareCtx } from './fakes' +import type { LockStore } from '@tanstack/ai/locks' +import type { RunStore } from '@tanstack/ai' + +function harness(input: { + runId: string + locks: LockStore + runs: RunStore + attach?: boolean +}) { + const ctx = makeMiddlewareCtx({ threadId: 't1', runId: input.runId }) + provideSandboxRuntime(ctx, { + logger: captureLogger().logger, + emit: () => undefined, + emitFileDiff: () => undefined, + }) + const middleware = withSandbox( + defineSandbox({ + id: 's', + provider: makeFakeProvider(), + lifecycle: { snapshot: 'none' }, + fileEvents: false, + }), + { + locks: input.locks, + runs: input.runs, + durability: { + adapter: memoryStream({ runId: input.runId }), + ...(input.attach === true ? { attach: true } : {}), + }, + }, + ) + return { ctx, middleware } +} + +function wait(ms: number): Promise { + return new Promise((resolve) => setTimeout(resolve, ms)) +} + +describe('fresh producer driver claim', () => { + it('keeps a takeover out until the run ends', async () => { + const locks = new InMemoryLockStore() + const runs = new InMemoryRunStore() + const { ctx, middleware } = harness({ runId: 'quiet', locks, runs }) + await middleware.setup?.(ctx) + expect((await runs.get('quiet'))?.driverEpoch).toBe(1) + + let entered = false + const takeover = withRunClaim({ runs, locks, runId: 'quiet' }, async () => { + entered = true + }) + await wait(30) + expect(entered).toBe(false) + + await middleware.onFinish?.(ctx, { + finishReason: 'stop', + duration: 1, + content: '', + }) + await takeover + expect(entered).toBe(true) + expect((await runs.get('quiet'))?.driverEpoch).toBe(2) + }) + + it('does not take the claim again from inside an existing claim', async () => { + const locks = new InMemoryLockStore() + const runs = new InMemoryRunStore() + await runs.createOrResume({ + runId: 'inner', + threadId: 't1', + startedAt: 1, + }) + await withRunClaim({ runs, locks, runId: 'inner' }, async () => { + const { ctx, middleware } = harness({ runId: 'inner', locks, runs }) + await middleware.setup?.(ctx) + expect((await runs.get('inner'))?.driverEpoch).toBe(1) + await middleware.onFinish?.(ctx, { + finishReason: 'stop', + duration: 1, + content: '', + }) + }) + }) + + it('an attach does not hold the claim', async () => { + const locks = new InMemoryLockStore() + const runs = new InMemoryRunStore() + const { ctx, middleware } = harness({ + runId: 'attach', + locks, + runs, + attach: true, + }) + await middleware.setup?.(ctx) + expect((await runs.get('attach'))?.driverEpoch).toBeUndefined() + + let entered = false + await withRunClaim({ runs, locks, runId: 'attach' }, async () => { + entered = true + }) + expect(entered).toBe(true) + }) +}) diff --git a/testing/e2e/tests/durable-takeover.spec.ts b/testing/e2e/tests/durable-takeover.spec.ts index 08c02feccc..0e788c1026 100644 --- a/testing/e2e/tests/durable-takeover.spec.ts +++ b/testing/e2e/tests/durable-takeover.spec.ts @@ -728,6 +728,46 @@ test.describe('durable runs — takeover', () => { // stored RUN_ERROR. expect(after.log.map((entry) => entry.type)).toEqual(FRESH_FULL_SEQUENCE) }) + + test('a reconnect during a quiet wait does not start a second driver', async () => { + const runId = uniqueRunId('quiet-reconnect') + const run = await SseStream.open(routeUrl({ runId, total: '6' })) + await tick(runId, 2) + expect(contentDeltas(await run.take(7))).toEqual(['1', '2']) + // Close the client only. `drop()` aborts the producer, which hides the + // bug: the producer must keep the claim while the journal is quiet. + run.close() + + const attach = await SseStream.open( + routeUrl({ runId, offset: '-1', fenceQuietMs: '50' }), + { method: 'GET' }, + ) + // The replay proves the GET handler is up. The driver starts beside it. + // One quiet window is 50ms, so 300ms is long enough for a second driver + // to claim the run when the first host does not hold the lock. + expect(contentDeltas(await attach.take(7))).toEqual(['1', '2']) + await new Promise((resolve) => setTimeout(resolve, 300)) + const mid = await state(runId) + expect(mid.record?.status).toBe('running') + expect(mid.record?.driverEpoch).toBe(1) + expect(mid.attachDrives).toBe(0) + + await tick(runId, 4) + await attach.drain() + const after = await stateUntil( + runId, + 'the original producer to finish the quiet run', + (body) => body.record?.status === 'completed', + ) + expect(after.attachDrives).toBe(0) + expect(after.record?.driverEpoch).toBe(1) + expect(after.log.map((entry) => entry.type)).toEqual(FRESH_FULL_SEQUENCE) + expect( + after.log + .filter((entry) => entry.type === 'TEXT_MESSAGE_CONTENT') + .map((entry) => entry.delta), + ).toEqual(['1', '2', '3', '4', '5', '6']) + }) }) // ---------------------------------------------------------------------------