diff --git a/src/Versioned.php b/src/Versioned.php
index 87858e5a..f4e56f9e 100644
--- a/src/Versioned.php
+++ b/src/Versioned.php
@@ -2981,4 +2981,42 @@ public function Publisher()
         $member = DataObject::get_by_id(Member::class, $this->owner->PublisherID);
         return $member;
     }
+
+    protected function updateStatusFlags(array &$flags): void
+    {
+        if ($this->isOnLiveOnly()) {
+            $flags['removedfromdraft'] = [
+                'text' => _t(__CLASS__ . '.FLAG_ONLIVEONLY_SHORT', 'On live only'),
+                'title' => _t(
+                    __CLASS__ . '.FLAG_ONLIVEONLYSHORT_HELP',
+                    'Item is published, but has been deleted from draft'
+                ),
+            ];
+            return;
+        }
+
+        if ($this->isArchived()) {
+            $flags['archived'] = [
+                'text' => _t(__CLASS__ . '.FLAG_ARCHIVED_SHORT', 'Archived'),
+                'title' => _t(__CLASS__ . '.FLAG_ARCHIVED_HELP', 'Item is removed from draft and live'),
+            ];
+            return;
+        }
+
+        if ($this->isOnDraftOnly()) {
+            $flags['addedtodraft'] = [
+                'text' => _t(__CLASS__ . '.FLAG_ADDEDTODRAFT_SHORT', 'Draft'),
+                'title' => _t(__CLASS__ . '.FLAG_ADDEDTODRAFT_HELP', 'Item has not been published yet')
+            ];
+            return;
+        }
+
+        if ($this->isModifiedOnDraft()) {
+            $flags['modified'] = [
+                'text' => _t(__CLASS__ . '.FLAG_MODIFIEDONDRAFT_SHORT', 'Modified'),
+                'title' => _t(__CLASS__ . '.FLAG_MODIFIEDONDRAFT_HELP', 'Item has unpublished changes'),
+            ];
+            return;
+        }
+    }
 }
diff --git a/src/VersionedGridFieldState/VersionedGridFieldState.php b/src/VersionedGridFieldState/VersionedGridFieldState.php
index 65b5dda6..b6b1ac0c 100644
--- a/src/VersionedGridFieldState/VersionedGridFieldState.php
+++ b/src/VersionedGridFieldState/VersionedGridFieldState.php
@@ -124,7 +124,7 @@ public function getColumnsHandled($gridField)
     public function getColumnContent($gridField, $record, $columnName)
     {
         $flagContent = '';
-        $flags = $this->getStatusFlags($record);
+        $flags = $record->getStatusFlags();
         foreach ($flags as $class => $data) {
             $flagAttributes = [
                 'class' => "ss-gridfield-badge badge status-{$class}",
@@ -162,78 +162,4 @@ public function getColumnMetadata($gridField, $columnName)
     {
         return [];
     }
-
-
-    /**
-     * A flag provides the user with additional data about the current item
-     * status, for example a "removed from draft" status. Each item can have
-     * more than one status flag. Returns a map of a unique key to a
-     * (localized) title for the flag. The unique key can be reused as a CSS
-     * class.
-     *
-     * Example (simple):
-     *
-     * ```php
-     *   "deletedonlive" => "Deleted"
-     * ```
-     *
-     * Example (with optional title attribute):
-     *
-     * ```php
-     *   "deletedonlive" => array(
-     *      'text' => "Deleted",
-     *      'title' => 'This page has been deleted'
-     *   )
-     * ```
-     *
-     * @param Versioned|DataObject $record - the record to check status for
-     * @return array
-     */
-    protected function getStatusFlags($record)
-    {
-        if (!$record->hasExtension(Versioned::class)) {
-            return [];
-        }
-
-        if ($record->isOnLiveOnly()) {
-            return [
-                'removedfromdraft' => [
-                    'text' => _t(__CLASS__ . '.ONLIVEONLYSHORT', 'On live only'),
-                    'title' => _t(
-                        __CLASS__ . '.ONLIVEONLYSHORTHELP',
-                        'Item is published, but has been deleted from draft'
-                    ),
-                ]
-            ];
-        }
-
-        if ($record->isArchived()) {
-            return [
-                'archived' => [
-                    'text' => _t(__CLASS__ . '.ARCHIVEDPAGESHORT', 'Archived'),
-                    'title' => _t(__CLASS__ . '.ARCHIVEDPAGEHELP', 'Item is removed from draft and live'),
-                ]
-            ];
-        }
-
-        if ($record->isOnDraftOnly()) {
-            return [
-                'addedtodraft' => [
-                    'text' => _t(__CLASS__ . '.ADDEDTODRAFTSHORT', 'Draft'),
-                    'title' => _t(__CLASS__ . '.ADDEDTODRAFTHELP', "Item has not been published yet")
-                ]
-            ];
-        }
-
-        if ($record->isModifiedOnDraft()) {
-            return [
-                'modified' => [
-                    'text' => _t(__CLASS__ . '.MODIFIEDONDRAFTSHORT', 'Modified'),
-                    'title' => _t(__CLASS__ . '.MODIFIEDONDRAFTHELP', 'Item has unpublished changes'),
-                ]
-            ];
-        }
-
-        return [];
-    }
 }
diff --git a/tests/php/VersionedTest.php b/tests/php/VersionedTest.php
index 3539fd78..9f6e47c2 100644
--- a/tests/php/VersionedTest.php
+++ b/tests/php/VersionedTest.php
@@ -6,6 +6,8 @@
 use InvalidArgumentException;
 use ReflectionMethod;
 use DMS\PHPUnitExtensions\ArraySubset\ArraySubsetAsserts;
+use LogicException;
+use PHPUnit\Framework\Attributes\DataProvider;
 use SilverStripe\Control\Controller;
 use SilverStripe\Control\Director;
 use SilverStripe\Control\HTTPRequest;
@@ -1700,4 +1702,94 @@ public function testLiveObjectDeletedOnUnpublish()
         $obj->doUnpublish();
         $this->assertEquals('First name', VersionedTest\TestObject::$nameValueOfObjectJustDeleted);
     }
+
+    public static function provideUpdateStatusFlags(): array
+    {
+        return [
+            'new record no status' => [
+                'status' => 'new',
+                'expected' => [],
+            ],
+            'saved to draft' => [
+                'status' => 'draft',
+                'expected' => [
+                    'addedtodraft' => [
+                        'text' => 'Draft',
+                        'title' => 'Item has not been published yet',
+                    ],
+                ],
+            ],
+            'saved and published' => [
+                'status' => 'published',
+                'expected' => [],
+            ],
+            'published but deleted on draft' => [
+                'status' => 'published only',
+                'expected' => [
+                    'removedfromdraft' => [
+                        'text' => 'On live only',
+                        'title' => 'Item is published, but has been deleted from draft',
+                    ],
+                ],
+            ],
+            'published but modified on draft' => [
+                'status' => 'modified',
+                'expected' => [
+                    'modified' => [
+                        'text' => 'Modified',
+                        'title' => 'Item has unpublished changes',
+                    ],
+                ],
+            ],
+            'archived' => [
+                'status' => 'archived',
+                'expected' => [
+                    'archived' => [
+                        'text' => 'Archived',
+                        'title' => 'Item is removed from draft and live',
+                    ],
+                ],
+            ],
+        ];
+    }
+
+    #[DataProvider('provideUpdateStatusFlags')]
+    public function testUpdateStatusFlags(string $status, array $expected): void
+    {
+        /** @var VersionedTest\TestObject&Versioned */
+        $record = new VersionedTest\TestObject();
+
+        switch ($status) {
+            case 'new':
+                // no-op
+                break;
+            case 'draft':
+                $record->write();
+                break;
+            case 'published':
+                $record->write();
+                $record->publishSingle();
+                break;
+            case 'published only':
+                $record->write();
+                $record->publishSingle();
+                $record->deleteFromStage(Versioned::DRAFT);
+                break;
+            case 'modified':
+                $record->write();
+                $record->publishSingle();
+                $record->Name = 'modified record';
+                $record->write();
+                break;
+            case 'archived':
+                $record->write();
+                $record->publishSingle();
+                $record->doArchive();
+                break;
+            default:
+                throw new LogicException("Status $status not accounted for");
+        }
+
+        $this->assertSame($expected, $record->getStatusFlags());
+    }
 }