diff --git a/packages/@n8n/ai-workflow-builder.ee/src/test/checkpoint-persistence.test.ts b/packages/@n8n/ai-workflow-builder.ee/src/test/checkpoint-persistence.test.ts
index d4f1b10b2c4..e193be63d9b 100644
--- a/packages/@n8n/ai-workflow-builder.ee/src/test/checkpoint-persistence.test.ts
+++ b/packages/@n8n/ai-workflow-builder.ee/src/test/checkpoint-persistence.test.ts
@@ -11,7 +11,6 @@
  */
 import { AIMessage, HumanMessage } from '@langchain/core/messages';
 import type { BaseMessage } from '@langchain/core/messages';
-import type { CheckpointTuple } from '@langchain/langgraph';
 import {
 	Annotation,
 	Command,
@@ -42,30 +41,11 @@ async function getCheckpointMessages(
 	checkpointer: MemorySaver,
 	threadId: string,
 ): Promise<BaseMessage[]> {
-	let latestTuple: CheckpointTuple | undefined;
-	for await (const tuple of checkpointer.list({ configurable: { thread_id: threadId } })) {
-		if (!latestTuple || isLaterCheckpoint(tuple, latestTuple)) latestTuple = tuple;
-	}
-
-	const messages = latestTuple?.checkpoint?.channel_values?.messages;
+	const tuple = await checkpointer.getTuple({ configurable: { thread_id: threadId } });
+	const messages = tuple?.checkpoint?.channel_values?.messages;
 	return Array.isArray(messages) ? (messages as BaseMessage[]) : [];
 }
 
-function isLaterCheckpoint(candidate: CheckpointTuple, current: CheckpointTuple): boolean {
-	const candidateStep = getCheckpointStep(candidate);
-	const currentStep = getCheckpointStep(current);
-
-	if (candidateStep !== currentStep) return candidateStep > currentStep;
-
-	return candidate.checkpoint.ts > current.checkpoint.ts;
-}
-
-function getCheckpointStep(tuple: CheckpointTuple): number {
-	const step = tuple.metadata?.step;
-
-	return typeof step === 'number' ? step : Number.NEGATIVE_INFINITY;
-}
-
 // ============================================================================
 // Tests
 // ============================================================================