[fix] allow for host cloud uploads of state (#25773)

## Summary & Motivation

In host cloud, calls to store state would fail because the state storage commands weren't implemented (as previously all state storage would have happened from user code side of the world)

This adds in that functionality and does a small refactor to reduce duplicated code

## Test Plan

## Changelog

> The changelog is generated by an agent that examines merged PRs and
> summarizes/categorizes user-facing changes. You can optionally replace
> this text with a terse description of any user-facing changes in your PR,
> which the agent will prioritize. Otherwise, delete this section.

Internal-RevId: 4c23bd2d5ff818bbd9996f3b9c00f9f32bc287f8

Author: OwenKephart

python_modules/dagster-cloud/dagster_cloud/storage/defs_state/storage.py

@@ -77,9 +77,6 @@ def _execute_query(self, query, variables=None, idempotent_mutation=False):
             query, variable_values=variables, idempotent_mutation=idempotent_mutation
         )
 
-    def _get_artifact_key(self, key: str, version: str) -> str:
-        return f"__state__/{self._sanitize_key(key)}/{version}"
-
     def download_state_to_path(self, key: str, version: str, path: Path) -> None:
         download_artifact(
             url=self.url,

python_modules/dagster/dagster/_core/storage/defs_state/base.py

@@ -117,13 +117,18 @@ def set_latest_version(self, key: str, version: str | None) -> None:
     def _sanitize_key(self, key: str) -> str:
         return re.sub(r"[^A-Za-z0-9._-]", "__", key)
 
-    def _get_version_key(self, key: str) -> str:
-        """Returns a storage key under which the latest version of a given key's state is stored."""
-        return f"__version__/{key}"
-
-    def _get_state_key(self, key: str, version: str) -> str:
-        """Returns a storage key under which a given key's state at a given version is stored."""
-        return f"__state__/{key}/{version}"
+    def _get_artifact_key(self, key: str, version: str) -> str:
+        """Returns the storage key under which a given defs key's state is stored at a given version.
+
+        The ``key`` component is sanitized so the result is safe to use as an S3 key or path
+        segment. This is the single canonical key layout shared by every remote
+        ``DefsStateStorage``: producers (the agent and ``dg`` CLI) and the consumer (host cloud)
+        must derive identical keys, so the layout lives here on the base class rather than being
+        reimplemented per subclass. It is effectively frozen -- changing the prefix or the
+        sanitization orphans every previously-uploaded state blob. See
+        ``test_artifact_key_layout_is_frozen``.
+        """
+        return f"__state__/{self._sanitize_key(key)}/{version}"
 
     @classmethod
     def get(cls) -> Optional["DefsStateStorage"]:

python_modules/dagster/dagster_tests/storage_tests/test_defs_state_storage.py

@@ -1,14 +1,44 @@
 import tempfile
+from pathlib import Path
 
 import pytest
 from dagster._core.instance_for_test import instance_for_test
+from dagster._core.storage.defs_state.base import DefsStateStorage
 from dagster._core.storage.defs_state.blob_storage_state_storage import UPathDefsStateStorage
 from dagster_shared import check
+from dagster_shared.serdes.objects.models.defs_state_info import DefsStateInfo
 from upath import UPath
 
 from dagster_tests.storage_tests.utils.defs_state_storage import TestDefsStateStorage
 
 
+def test_artifact_key_layout_is_frozen():
+    """The remote DefsStateStorage key layout is a cross-version contract: the agent and the
+    ``dg`` CLI (producers) and host cloud (consumer) all derive blob keys from the shared
+    base-class ``_get_artifact_key``. Changing the prefix or the sanitization orphans every
+    previously-uploaded state blob, so this layout is effectively frozen -- if this test fails
+    because you changed the layout, you are also signing up to migrate already-stored state.
+    """
+
+    class _StubDefsStateStorage(DefsStateStorage):
+        def get_latest_defs_state_info(self) -> DefsStateInfo | None:
+            raise NotImplementedError
+
+        def download_state_to_path(self, key: str, version: str, path: Path) -> None:
+            raise NotImplementedError
+
+        def upload_state_from_path(self, key: str, version: str, path: Path) -> None:
+            raise NotImplementedError
+
+        def set_latest_version(self, key: str, version: str | None) -> None:
+            raise NotImplementedError
+
+    storage = _StubDefsStateStorage()
+    assert storage._get_artifact_key("MyComponent", "v1") == "__state__/MyComponent/v1"  # noqa: SLF001
+    # The key component is sanitized: anything outside [A-Za-z0-9._-] collapses to "__".
+    assert storage._get_artifact_key("a/b c", "v1") == "__state__/a__b__c/v1"  # noqa: SLF001
+
+
 class TestDefaultDefsStateStorage(TestDefsStateStorage):
     """Tests the default state storage implementation."""
 

python_modules/libraries/dagster-dg-cli/dagster_dg_cli/utils/plus/defs_state_storage.py

@@ -136,9 +136,6 @@ def graphql_client(self) -> IGraphQLClient:
     def _execute_query(self, query, variables=None):
         return self.graphql_client.execute_arbitrary(query, variables=variables)
 
-    def _get_artifact_key(self, key: str, version: str) -> str:
-        return f"__state__/{self._sanitize_key(key)}/{version}"
-
     def download_state_to_path(self, key: str, version: str, path: Path) -> None:
         download_artifact(
             url=self.url,