diff --git a/e2e/specs/mock/model-replay-tools.spec.ts b/e2e/specs/mock/model-replay-tools.spec.ts index df7d1f1c50..d290073cd1 100644 --- a/e2e/specs/mock/model-replay-tools.spec.ts +++ b/e2e/specs/mock/model-replay-tools.spec.ts @@ -86,7 +86,7 @@ test.describe('recorded tool-call fixture replay', () => { await sendMessageAndWaitForCompletion(page, TOOL_PROMPT, { timeout: 120_000 }); - const conversationId = /\/c\/([^/]+)/.exec(page.url())?.[1]; + const conversationId = /\/c\/([^/]+)/.exec(new URL(page.url()).pathname)?.[1]; expect(conversationId, 'conversation should have a persisted id').toBeTruthy(); const token = await getAccessToken(page); const messages = await fetchJson( diff --git a/packages/api/src/stream/GenerationJobManager.ts b/packages/api/src/stream/GenerationJobManager.ts index 81b98f6764..9b7f979062 100644 --- a/packages/api/src/stream/GenerationJobManager.ts +++ b/packages/api/src/stream/GenerationJobManager.ts @@ -1125,10 +1125,10 @@ class GenerationJobManagerClass { this.reconcileInactiveGeneration(streamId, createdAt, currentJob, observedRuntime); } - private async reconcileDiscoveredSuccessor( + private async reconcileFencedRuntimeHandoff( streamId: string, runtime: RuntimeJobState, - successor: SerializableJobData, + currentJob: SerializableJobData | null, ): Promise { const ownsExactProvider = this.ownedJobs.get(streamId) === runtime.createdAt; if (!runtime.abortController.signal.aborted) { @@ -1140,7 +1140,7 @@ class GenerationJobManagerClass { ownsExactProvider, ); if (abortProofPersisted) { - this.reconcileInactiveGeneration(streamId, runtime.createdAt, successor, runtime); + this.reconcileInactiveGeneration(streamId, runtime.createdAt, currentJob, runtime); } this.preserveFencedRuntimeUntilHandoff(streamId, runtime); } @@ -1162,6 +1162,21 @@ class GenerationJobManagerClass { if (Date.now() - startedAt < TERMINAL_PERSISTENCE_TIMEOUT_MS) { return jobData; } + const observedRuntime = this.runtimeState.get(jobData.streamId); + let runtime = observedRuntime?.createdAt === jobData.createdAt ? observedRuntime : undefined; + /** Preserve the pre-await object when one exists, but also admit a predecessor + * attachment created while terminal finalization is in flight. Every late + * read remains generation-fenced, so a successor runtime is never captured. */ + const captureMatchingRuntime = (): RuntimeJobState | undefined => { + if (runtime) { + return runtime; + } + const candidate = this.runtimeState.get(jobData.streamId); + if (candidate?.createdAt === jobData.createdAt) { + runtime = candidate; + } + return runtime; + }; const reconcileEvent = buildTerminalPersistenceReconcile(jobData); const serialized = JSON.stringify(reconcileEvent); @@ -1171,19 +1186,33 @@ class GenerationJobManagerClass { serialized, ); if (!recovered) { - return this.jobStore.getJob(jobData.streamId); + const currentJob = await this.jobStore.getJob(jobData.streamId); + const predecessorRuntime = captureMatchingRuntime(); + if (predecessorRuntime && currentJob?.createdAt !== jobData.createdAt) { + await this.reconcileFencedRuntimeHandoff(jobData.streamId, predecessorRuntime, currentJob); + } + return currentJob; } - const runtime = this.runtimeState.get(jobData.streamId); - if (runtime?.createdAt === jobData.createdAt) { - runtime.finalEvent = reconcileEvent; + const recoveredRuntime = captureMatchingRuntime(); + if (recoveredRuntime) { + recoveredRuntime.finalEvent = reconcileEvent; } try { await this.eventTransport.emitDone(jobData.streamId, reconcileEvent, jobData.createdAt); } catch (err) { if (err instanceof GenerationPublicationFencedError) { const currentJob = await this.jobStore.getJob(jobData.streamId); - this.reconcileInactiveGeneration(jobData.streamId, jobData.createdAt, currentJob, runtime); + const predecessorRuntime = captureMatchingRuntime(); + if (predecessorRuntime) { + await this.reconcileFencedRuntimeHandoff( + jobData.streamId, + predecessorRuntime, + currentJob, + ); + } else { + this.reconcileInactiveGeneration(jobData.streamId, jobData.createdAt, currentJob); + } return currentJob; } logger.error( @@ -4935,7 +4964,7 @@ class GenerationJobManagerClass { let terminalJob = await this.jobStore.getJob(streamId); if (terminalJob?.createdAt !== runtime.createdAt) { if (terminalJob) { - await this.reconcileDiscoveredSuccessor(streamId, runtime, terminalJob); + await this.reconcileFencedRuntimeHandoff(streamId, runtime, terminalJob); } else { queueError(TERMINAL_PUBLICATION_RECONNECT_ERROR); } @@ -4948,7 +4977,7 @@ class GenerationJobManagerClass { terminalJob = await this.recoverStaleTerminalPersistence(terminalJob); if (terminalJob?.createdAt !== runtime.createdAt) { if (terminalJob) { - await this.reconcileDiscoveredSuccessor(streamId, runtime, terminalJob); + await this.reconcileFencedRuntimeHandoff(streamId, runtime, terminalJob); } else { queueError(TERMINAL_PUBLICATION_RECONNECT_ERROR); } @@ -8028,17 +8057,7 @@ class GenerationJobManagerClass { } catch (err) { if (err instanceof GenerationPublicationFencedError) { const currentJob = await this.jobStore.getJob(streamId); - if (currentJob) { - await this.reconcileDiscoveredSuccessor(streamId, observedRuntime, currentJob); - } else { - this.reconcileInactiveGeneration( - streamId, - observedRuntime.createdAt, - currentJob, - observedRuntime, - ); - this.preserveFencedRuntimeUntilHandoff(streamId, observedRuntime); - } + await this.reconcileFencedRuntimeHandoff(streamId, observedRuntime, currentJob); continue; } logger.error(`[GenerationJobManager] Failed to notify reaped stream ${streamId}:`, err); diff --git a/packages/api/src/stream/__tests__/startup.spec.ts b/packages/api/src/stream/__tests__/startup.spec.ts index c435fb4866..63d7980df1 100644 --- a/packages/api/src/stream/__tests__/startup.spec.ts +++ b/packages/api/src/stream/__tests__/startup.spec.ts @@ -3217,10 +3217,12 @@ describe('GenerationJobManager startup telemetry', () => { const streamId = 'stream-terminal-recovery-replaced'; const predecessor = await manager.createJob(streamId, 'user-1', streamId); const predecessorDone = jest.fn(); + const predecessorError = jest.fn(); const predecessorSubscription = await manager.subscribe( streamId, () => undefined, predecessorDone, + predecessorError, ); await jobStore.transitionStatusAndDrainSteers(streamId, { from: 'running', @@ -3238,20 +3240,178 @@ describe('GenerationJobManager startup telemetry', () => { const finalized = await finalize(...args); const successor = await jobStore.createJob(streamId, 'user-1', streamId); successorCreatedAt = successor.createdAt; + await manager.getJob(streamId); return finalized; }); jest.spyOn(eventTransport, 'emitDone').mockImplementation((_streamId, _event, generationId) => { throw new GenerationPublicationFencedError('done', streamId, generationId); }); + jest.useFakeTimers(); try { - await expect(manager.getJob(streamId)).resolves.toMatchObject({ + const successor = await manager.getJob(streamId); + expect(successor).toMatchObject({ createdAt: expect.any(Number), status: 'running', }); - expect((await manager.getJob(streamId))?.createdAt).toBe(successorCreatedAt); + expect(successor?.createdAt).toBe(successorCreatedAt); + expect(successor?.abortController.signal.aborted).toBe(false); + expect(predecessor.abortController.signal.aborted).toBe(true); expect(predecessorDone).not.toHaveBeenCalled(); + expect(predecessorError).not.toHaveBeenCalled(); + expect(manager.getRuntimeStats().fencedRuntimeRetirements).toBe(1); + + await jest.advanceTimersByTimeAsync( + REDIS_REPLACEMENT_HANDOFF_MAX_WAIT_MS + REDIS_EVENT_REORDER_TIMEOUT_MS * 2, + ); + + expect(predecessorDone).not.toHaveBeenCalled(); + expect(predecessorError).toHaveBeenCalledWith(TERMINAL_PUBLICATION_RECONNECT_ERROR); + expect(successor?.abortController.signal.aborted).toBe(false); + expect(manager.getRuntimeStats().runtimeStateSize).toBe(1); } finally { + jest.useRealTimers(); + predecessorSubscription?.unsubscribe(); + await manager.destroy(); + } + }); + + it('reconnect-closes a predecessor that attaches during stale finalization', async () => { + const jobStore = new InMemoryJobStore({ ttlAfterComplete: 60_000 }); + const eventTransport = new InMemoryEventTransport(); + const manager = new GenerationJobManagerClass(); + manager.configure({ jobStore, eventTransport, isRedis: false, cleanupOnComplete: false }); + manager.initialize(); + const streamId = 'stream-terminal-recovery-late-attachment'; + const predecessor = await jobStore.createJob(streamId, 'user-1', streamId); + await jobStore.transitionStatusAndDrainSteers(streamId, { + from: 'running', + to: 'aborted', + expectCreatedAt: predecessor.createdAt, + patch: { + completedAt: Date.now() - 60_000, + terminalPersistencePending: true, + terminalPersistenceStartedAt: Date.now() - 60_000, + }, + }); + const finalize = jobStore.finalizeTerminalPersistence.bind(jobStore); + let finalizeEntered!: () => void; + const finalizationStarted = new Promise((resolve) => { + finalizeEntered = resolve; + }); + let releaseFinalize!: () => void; + const finalizeReleased = new Promise((resolve) => { + releaseFinalize = resolve; + }); + let successorCreatedAt = 0; + jest.spyOn(jobStore, 'finalizeTerminalPersistence').mockImplementation(async (...args) => { + finalizeEntered(); + await finalizeReleased; + const finalized = await finalize(...args); + const successor = await jobStore.createJob(streamId, 'user-1', streamId); + successorCreatedAt = successor.createdAt; + return finalized; + }); + jest.spyOn(eventTransport, 'emitDone').mockImplementation((_streamId, _event, generationId) => { + throw new GenerationPublicationFencedError('done', streamId, generationId); + }); + const predecessorError = jest.fn(); + const recovery = manager.getJob(streamId); + await finalizationStarted; + const pendingPredecessor = await jobStore.getJob(streamId); + const lateRuntime = await ( + manager as unknown as { + getOrCreateRuntimeState: ( + id: string, + job: typeof pendingPredecessor, + ) => Promise<{ + createdAt: number; + abortController: AbortController; + localErrorHandlers: Set<(error: string) => void>; + } | null>; + } + ).getOrCreateRuntimeState(streamId, pendingPredecessor); + lateRuntime?.localErrorHandlers.add(predecessorError); + jest.useFakeTimers(); + + try { + expect(lateRuntime?.createdAt).toBe(predecessor.createdAt); + releaseFinalize(); + const successor = await recovery; + + expect(successor?.createdAt).toBe(successorCreatedAt); + expect(successor?.abortController.signal.aborted).toBe(false); + expect(predecessorError).not.toHaveBeenCalled(); + expect(manager.getRuntimeStats().fencedRuntimeRetirements).toBe(1); + + await jest.advanceTimersByTimeAsync( + REDIS_REPLACEMENT_HANDOFF_MAX_WAIT_MS + REDIS_EVENT_REORDER_TIMEOUT_MS * 2, + ); + + expect(predecessorError).toHaveBeenCalledWith(TERMINAL_PUBLICATION_RECONNECT_ERROR); + expect(successor?.abortController.signal.aborted).toBe(false); + expect(manager.getRuntimeStats().runtimeStateSize).toBe(1); + } finally { + releaseFinalize(); + jest.useRealTimers(); + await manager.destroy(); + } + }); + + it('reconnect-closes the predecessor when a successor wins stale finalization', async () => { + const jobStore = new InMemoryJobStore({ ttlAfterComplete: 60_000 }); + const eventTransport = new InMemoryEventTransport(); + const manager = new GenerationJobManagerClass(); + manager.configure({ jobStore, eventTransport, isRedis: false, cleanupOnComplete: false }); + manager.initialize(); + const streamId = 'stream-terminal-recovery-finalize-replaced'; + const predecessor = await manager.createJob(streamId, 'user-1', streamId); + const predecessorDone = jest.fn(); + const predecessorError = jest.fn(); + const predecessorSubscription = await manager.subscribe( + streamId, + () => undefined, + predecessorDone, + predecessorError, + ); + await jobStore.transitionStatusAndDrainSteers(streamId, { + from: 'running', + to: 'aborted', + expectCreatedAt: predecessor.createdAt, + patch: { + completedAt: Date.now() - 60_000, + terminalPersistencePending: true, + terminalPersistenceStartedAt: Date.now() - 60_000, + }, + }); + const finalize = jobStore.finalizeTerminalPersistence.bind(jobStore); + let successorCreatedAt = 0; + jest.spyOn(jobStore, 'finalizeTerminalPersistence').mockImplementation(async (...args) => { + const successor = await jobStore.createJob(streamId, 'user-1', streamId); + successorCreatedAt = successor.createdAt; + return finalize(...args); + }); + jest.useFakeTimers(); + + try { + const successor = await manager.getJob(streamId); + expect(successor?.createdAt).toBe(successorCreatedAt); + expect(successor?.abortController.signal.aborted).toBe(false); + expect(predecessor.abortController.signal.aborted).toBe(true); + expect(predecessorDone).not.toHaveBeenCalled(); + expect(predecessorError).not.toHaveBeenCalled(); + expect(manager.getRuntimeStats().fencedRuntimeRetirements).toBe(1); + + await jest.advanceTimersByTimeAsync( + REDIS_REPLACEMENT_HANDOFF_MAX_WAIT_MS + REDIS_EVENT_REORDER_TIMEOUT_MS * 2, + ); + + expect(predecessorDone).not.toHaveBeenCalled(); + expect(predecessorError).toHaveBeenCalledWith(TERMINAL_PUBLICATION_RECONNECT_ERROR); + expect(successor?.abortController.signal.aborted).toBe(false); + expect(manager.getRuntimeStats().runtimeStateSize).toBe(1); + } finally { + jest.useRealTimers(); predecessorSubscription?.unsubscribe(); await manager.destroy(); }