diff --git a/plugins/scaffolder-backend/src/scaffolder/tasks/NunjucksWorkflowRunner.test.ts b/plugins/scaffolder-backend/src/scaffolder/tasks/NunjucksWorkflowRunner.test.ts
index eaf477f35f..420ec69427 100644
--- a/plugins/scaffolder-backend/src/scaffolder/tasks/NunjucksWorkflowRunner.test.ts
+++ b/plugins/scaffolder-backend/src/scaffolder/tasks/NunjucksWorkflowRunner.test.ts
@@ -18,7 +18,7 @@ import { getVoidLogger } from '@backstage/backend-common';
 import { NunjucksWorkflowRunner } from './NunjucksWorkflowRunner';
 import { TemplateActionRegistry } from '../actions';
 import { ScmIntegrations } from '@backstage/integration';
-import { JsonValue } from '@backstage/types';
+import { JsonObject } from '@backstage/types';
 import { ConfigReader } from '@backstage/config';
 import { TaskSpec } from '@backstage/plugin-scaffolder-common';
 import {
@@ -580,26 +580,21 @@ describe('NunjucksWorkflowRunner', () => {
         }),
         getTaskState: (): Promise<
           | {
-              state: {
-                [key: string]:
-                  | { status: 'failed'; reason: string }
-                  | {
-                      status: 'success';
-                      value: JsonValue;
-                    };
-              };
+              state: JsonObject;
             }
           | undefined
         > => {
           return Promise.resolve({
             state: {
-              ['v1.task.checkpoint.key1']: {
-                status: 'success',
-                value: 'initial',
-              },
-              ['v1.task.checkpoint.key2']: {
-                status: 'failed',
-                reason: 'fatal error',
+              checkpoints: {
+                ['v1.task.checkpoint.key1']: {
+                  status: 'success',
+                  value: 'initial',
+                },
+                ['v1.task.checkpoint.key2']: {
+                  status: 'failed',
+                  reason: 'fatal error',
+                },
               },
             },
           });
diff --git a/plugins/scaffolder-backend/src/scaffolder/tasks/NunjucksWorkflowRunner.ts b/plugins/scaffolder-backend/src/scaffolder/tasks/NunjucksWorkflowRunner.ts
index e3bf094009..634ac420aa 100644
--- a/plugins/scaffolder-backend/src/scaffolder/tasks/NunjucksWorkflowRunner.ts
+++ b/plugins/scaffolder-backend/src/scaffolder/tasks/NunjucksWorkflowRunner.ts
@@ -368,7 +368,9 @@ export class NunjucksWorkflowRunner implements WorkflowRunner {
             try {
               let prevValue: U | undefined;
               if (prevTaskState) {
-                const prevState = (prevTaskState.state as TaskState)?.[key];
+                const prevState = (
+                  prevTaskState.state?.checkpoints as TaskState
+                )?.[key];
                 if (prevState && prevState.status === 'success') {
                   prevValue = prevState.value as U;
                 }