diff --git a/plugins/scaffolder-backend/api-report.md b/plugins/scaffolder-backend/api-report.md
index b226372d86..0bcd8ce33c 100644
--- a/plugins/scaffolder-backend/api-report.md
+++ b/plugins/scaffolder-backend/api-report.md
@@ -411,17 +411,15 @@ export class DatabaseTaskStore implements TaskStore {
   // (undocumented)
   getTaskState({ taskId }: { taskId: string }): Promise<
     | {
-        state: {
-          [key: string]:
-            | {
-                status: 'failed';
-                reason: string;
-              }
-            | {
-                status: 'success';
-                value: JsonValue;
-              };
-        };
+        [key: string]:
+          | {
+              status: 'failed';
+              reason: string;
+            }
+          | {
+              status: 'success';
+              value: JsonValue;
+            };
       }
     | undefined
   >;
@@ -567,17 +565,15 @@ export class TaskManager implements TaskContext {
   // (undocumented)
   getTaskState?(): Promise<
     | {
-        state: {
-          [key: string]:
-            | {
-                status: 'failed';
-                reason: string;
-              }
-            | {
-                status: 'success';
-                value: JsonValue;
-              };
-        };
+        [key: string]:
+          | {
+              status: 'failed';
+              reason: string;
+            }
+          | {
+              status: 'success';
+              value: JsonValue;
+            };
       }
     | undefined
   >;
@@ -632,17 +628,15 @@ export interface TaskStore {
   // (undocumented)
   getTaskState?({ taskId }: { taskId: string }): Promise<
     | {
-        state: {
-          [key: string]:
-            | {
-                status: 'failed';
-                reason: string;
-              }
-            | {
-                status: 'success';
-                value: JsonValue;
-              };
-        };
+        [key: string]:
+          | {
+              status: 'failed';
+              reason: string;
+            }
+          | {
+              status: 'success';
+              value: JsonValue;
+            };
       }
     | undefined
   >;
diff --git a/plugins/scaffolder-backend/src/scaffolder/tasks/DatabaseTaskStore.ts b/plugins/scaffolder-backend/src/scaffolder/tasks/DatabaseTaskStore.ts
index 0d6521dfd2..db890c330e 100644
--- a/plugins/scaffolder-backend/src/scaffolder/tasks/DatabaseTaskStore.ts
+++ b/plugins/scaffolder-backend/src/scaffolder/tasks/DatabaseTaskStore.ts
@@ -397,14 +397,12 @@ export class DatabaseTaskStore implements TaskStore {
 
   async getTaskState({ taskId }: { taskId: string }): Promise<
     | {
-        state: {
-          [key: string]:
-            | { status: 'failed'; reason: string }
-            | {
-                status: 'success';
-                value: JsonValue;
-              };
-        };
+        [key: string]:
+          | { status: 'failed'; reason: string }
+          | {
+              status: 'success';
+              value: JsonValue;
+            };
       }
     | undefined
   > {
@@ -412,16 +410,14 @@ export class DatabaseTaskStore implements TaskStore {
       .where({ id: taskId })
       .select('state');
     return result.state
-      ? {
-          state: JSON.parse(result.state) as unknown as {
-            [key: string]:
-              | { status: 'failed'; reason: string }
-              | {
-                  status: 'success';
-                  value: JsonValue;
-                };
-          },
-        }
+      ? (JSON.parse(result.state) as unknown as {
+          [key: string]:
+            | { status: 'failed'; reason: string }
+            | {
+                status: 'success';
+                value: JsonValue;
+              };
+        })
       : undefined;
   }
 
diff --git a/plugins/scaffolder-backend/src/scaffolder/tasks/NunjucksWorkflowRunner.test.ts b/plugins/scaffolder-backend/src/scaffolder/tasks/NunjucksWorkflowRunner.test.ts
index e6abfcb8de..2b8db0d0ba 100644
--- a/plugins/scaffolder-backend/src/scaffolder/tasks/NunjucksWorkflowRunner.test.ts
+++ b/plugins/scaffolder-backend/src/scaffolder/tasks/NunjucksWorkflowRunner.test.ts
@@ -18,6 +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 { ConfigReader } from '@backstage/config';
 import { TaskContext } from './types';
 import { TaskSpec } from '@backstage/plugin-scaffolder-common';
@@ -137,29 +138,19 @@ describe('NunjucksWorkflowRunner', () => {
       id: 'checkpoints-action',
       description: 'Mock action with checkpoints',
       handler: async ctx => {
-        let key1 = 0;
-        let key2 = '';
-        let i = 0;
-
-        const incrementKey1 = async () => {
-          key1 += 1;
-          return key1;
-        };
-
-        const appendKey2 = async () => {
-          key2 += 'k';
-          return key2;
-        };
-
-        while (i < 3) {
-          i += 1;
-          ctx.checkpoint?.('key1', incrementKey1);
-        }
-
-        ctx.checkpoint?.('key2', appendKey2);
+        const key1 = await ctx.checkpoint?.('key1', async () => {
+          return 'updated';
+        });
+        const key2 = await ctx.checkpoint?.('key2', async () => {
+          return 'updated';
+        });
+        const key3 = await ctx.checkpoint?.('key3', async () => {
+          return 'updated';
+        });
 
         ctx.output('key1', key1);
         ctx.output('key2', key2);
+        ctx.output('key3', key3);
       },
     });
 
@@ -569,26 +560,52 @@ describe('NunjucksWorkflowRunner', () => {
     });
 
     it('should deal with checkpoints', async () => {
-      const task = createMockTaskWithSpec({
-        apiVersion: 'scaffolder.backstage.io/v1beta3',
-        parameters: {},
-        steps: [
-          {
-            id: 'test',
-            name: 'name',
-            action: 'checkpoints-action',
-            input: { foo: 1 },
+      const task = {
+        ...createMockTaskWithSpec({
+          apiVersion: 'scaffolder.backstage.io/v1beta3',
+          parameters: {},
+          steps: [
+            {
+              id: 'test',
+              name: 'name',
+              action: 'checkpoints-action',
+              input: { foo: 1 },
+            },
+          ],
+          output: {
+            key1: '${{steps.test.output.key1}}',
+            key2: '${{steps.test.output.key2}}',
+            key3: '${{steps.test.output.key3}}',
           },
-        ],
-        output: {
-          key1: '${{steps.test.output.key1}}',
-          key2: '${{steps.test.output.key2}}',
+        }),
+        getTaskState: (): Promise<
+          | {
+              [key: string]:
+                | { status: 'failed'; reason: string }
+                | {
+                    status: 'success';
+                    value: JsonValue;
+                  };
+            }
+          | undefined
+        > => {
+          return Promise.resolve({
+            ['v1.task.checkpoint.key1']: {
+              status: 'success',
+              value: 'initial',
+            },
+            ['v1.task.checkpoint.key2']: {
+              status: 'failed',
+              reason: 'fatal error',
+            },
+          });
         },
-      });
+      };
       const result = await runner.execute(task);
 
-      expect(result.output.key1).toEqual(3);
-      expect(result.output.key2).toEqual('k');
+      expect(result.output.key1).toEqual('initial');
+      expect(result.output.key2).toEqual('updated');
+      expect(result.output.key3).toEqual('updated');
     });
 
     it('should template the output from simple actions', async () => {
diff --git a/plugins/scaffolder-backend/src/scaffolder/tasks/NunjucksWorkflowRunner.ts b/plugins/scaffolder-backend/src/scaffolder/tasks/NunjucksWorkflowRunner.ts
index b5ffc73189..1dd14bc905 100644
--- a/plugins/scaffolder-backend/src/scaffolder/tasks/NunjucksWorkflowRunner.ts
+++ b/plugins/scaffolder-backend/src/scaffolder/tasks/NunjucksWorkflowRunner.ts
@@ -358,18 +358,21 @@ export class NunjucksWorkflowRunner implements WorkflowRunner {
             try {
               let prevValue: U | undefined;
               if (prevTaskState) {
-                const prevState = prevTaskState.state[key];
-                if (prevState.status === 'success') {
+                const prevState = prevTaskState[key];
+                if (prevState && prevState.status === 'success') {
                   prevValue = prevState.value as U;
                 }
               }
 
               const value = prevValue ? prevValue : await fn();
-              task.updateCheckpoint?.({
-                key,
-                status: 'success',
-                value,
-              });
+
+              if (!prevValue) {
+                task.updateCheckpoint?.({
+                  key,
+                  status: 'success',
+                  value,
+                });
+              }
               return value;
             } catch (err) {
               task.updateCheckpoint?.({
diff --git a/plugins/scaffolder-backend/src/scaffolder/tasks/StorageTaskBroker.ts b/plugins/scaffolder-backend/src/scaffolder/tasks/StorageTaskBroker.ts
index 453df09c77..0404f143de 100644
--- a/plugins/scaffolder-backend/src/scaffolder/tasks/StorageTaskBroker.ts
+++ b/plugins/scaffolder-backend/src/scaffolder/tasks/StorageTaskBroker.ts
@@ -93,14 +93,12 @@ export class TaskManager implements TaskContext {
 
   async getTaskState?(): Promise<
     | {
-        state: {
-          [key: string]:
-            | { status: 'failed'; reason: string }
-            | {
-                status: 'success';
-                value: JsonValue;
-              };
-        };
+        [key: string]:
+          | { status: 'failed'; reason: string }
+          | {
+              status: 'success';
+              value: JsonValue;
+            };
       }
     | undefined
   > {
diff --git a/plugins/scaffolder-backend/src/scaffolder/tasks/types.ts b/plugins/scaffolder-backend/src/scaffolder/tasks/types.ts
index 215defd280..45c4774299 100644
--- a/plugins/scaffolder-backend/src/scaffolder/tasks/types.ts
+++ b/plugins/scaffolder-backend/src/scaffolder/tasks/types.ts
@@ -196,14 +196,12 @@ export interface TaskStore {
 
   getTaskState?({ taskId }: { taskId: string }): Promise<
     | {
-        state: {
-          [key: string]:
-            | { status: 'failed'; reason: string }
-            | {
-                status: 'success';
-                value: JsonValue;
-              };
-        };
+        [key: string]:
+          | { status: 'failed'; reason: string }
+          | {
+              status: 'success';
+              value: JsonValue;
+            };
       }
     | undefined
   >;
diff --git a/plugins/scaffolder-node/api-report.md b/plugins/scaffolder-node/api-report.md
index 7e29e16aa0..7c024a0043 100644
--- a/plugins/scaffolder-node/api-report.md
+++ b/plugins/scaffolder-node/api-report.md
@@ -351,17 +351,15 @@ export interface TaskContext {
   // (undocumented)
   getTaskState?(): Promise<
     | {
-        state: {
-          [key: string]:
-            | {
-                status: 'failed';
-                reason: string;
-              }
-            | {
-                status: 'success';
-                value: JsonValue;
-              };
-        };
+        [key: string]:
+          | {
+              status: 'failed';
+              reason: string;
+            }
+          | {
+              status: 'success';
+              value: JsonValue;
+            };
       }
     | undefined
   >;
diff --git a/plugins/scaffolder-node/src/tasks/types.ts b/plugins/scaffolder-node/src/tasks/types.ts
index 5a4838737c..18383dbac3 100644
--- a/plugins/scaffolder-node/src/tasks/types.ts
+++ b/plugins/scaffolder-node/src/tasks/types.ts
@@ -128,14 +128,12 @@ export interface TaskContext {
 
   getTaskState?(): Promise<
     | {
-        state: {
-          [key: string]:
-            | { status: 'failed'; reason: string }
-            | {
-                status: 'success';
-                value: JsonValue;
-              };
-        };
+        [key: string]:
+          | { status: 'failed'; reason: string }
+          | {
+              status: 'success';
+              value: JsonValue;
+            };
       }
     | undefined
   >;