[ML-69282] [s] Explain store creation permissions across CLI and SDK

integrations/agentbricks/src/databricks_agentbricks/clients/conversation_store_client.py

@@ -10,14 +10,10 @@
 from typing import Any, Optional
 
 from databricks_agentbricks.clients.api_client_provider import ApiClientProvider
-from databricks_agentbricks.errors import AgentCliError
+from databricks_agentbricks.errors import AgentCliError, store_create_permission_error
 
 _MEMORY_STORE_PAGE_SIZE = 100  # the memory-stores list API caps page_size at 100
 
-_LAKEBASE_PERMISSION_DOCS = (
-    "https://docs.databricks.com/aws/en/oltp/projects/manage-project-permissions"
-)
-
 
 def _field(payload: Any, name: str) -> Any:
     """Read snake_case or camelCase fields without importing the CLI rendering layer."""
@@ -46,27 +42,14 @@ class SessionStoreReconcileResult:
     created: bool
 
 
-def _store_create_permission_error(name: str, kind: str, cause: AgentCliError) -> AgentCliError:
-    """PERMISSION_DENIED on create: the workspace admin has restricted Lakebase project creation."""
-    return AgentCliError(
-        f"You don't have permission to create {kind} store '{name}'.",
-        error_code=cause.error_code,
-        hint=(
-            "Creating a managed store provisions a Lakebase project, which your workspace admin "
-            "has restricted. Ask your workspace admin to grant you permission to create Lakebase "
-            f"projects ({_LAKEBASE_PERMISSION_DOCS}), or bind an existing store you can access "
-            "with --no-create-stores."
-        ),
-    )
-
-
 def _store_access_error(name: str, kind: str) -> AgentCliError:
     """The store already exists but isn't accessible to the caller."""
+    command = "memory" if kind == "memory" else "sessions"
     return AgentCliError(
         f"{kind.capitalize()} store '{name}' already exists but you don't have access to it.",
         hint=(
             "Ask the store's owner or your workspace admin to grant you access, or bind a "
-            "different store you can access with --no-create-stores."
+            f"different store with `agentbricks {command} bind <existing-store>` and deploy again."
         ),
     )
 
@@ -108,7 +91,7 @@ def ensure(self, display_name: str) -> tuple[dict, bool]:
             ), True
         except AgentCliError as exc:
             if exc.error_code == "PERMISSION_DENIED":
-                raise _store_create_permission_error(display_name, "memory", exc) from exc
+                raise store_create_permission_error(display_name, "memory", exc) from exc
             if exc.error_code != "ALREADY_EXISTS":
                 raise
         store = self.resolve(display_name)
@@ -167,7 +150,7 @@ def ensure(self, name: str) -> tuple[dict, bool]:
             ), True
         except AgentCliError as exc:
             if exc.error_code == "PERMISSION_DENIED":
-                raise _store_create_permission_error(name, "session", exc) from exc
+                raise store_create_permission_error(name, "session", exc) from exc
             if exc.error_code != "ALREADY_EXISTS":
                 raise
         try:

integrations/agentbricks/src/databricks_agentbricks/errors.py

@@ -43,6 +43,10 @@ def set_output_mode(mode: str) -> None:
     "This is usually a transient backend issue — re-running the command often succeeds."
 )
 
+_LAKEBASE_PERMISSION_DOCS = (
+    "https://docs.databricks.com/aws/en/oltp/projects/manage-project-permissions"
+)
+
 
 class AgentCliError(click.ClickException):
     """A user-facing CLI error rendered without a Python traceback."""
@@ -76,6 +80,21 @@ def show(self, file=None) -> None:
         )
 
 
+def store_create_permission_error(name: str, kind: str, cause: AgentCliError) -> AgentCliError:
+    """Explain managed-store creation permissions for CLI and SDK callers."""
+    command = "memory" if kind == "memory" else "sessions"
+    return AgentCliError(
+        f"You don't have permission to create {kind} store '{name}'.",
+        error_code=cause.error_code,
+        hint=(
+            "Managed store creation may require permission to create Lakebase projects. "
+            "Ask your workspace admin to check your access "
+            f"({_LAKEBASE_PERMISSION_DOCS}). To reuse an existing store you can access, run "
+            f"`agentbricks {command} bind <existing-store>` and deploy again."
+        ),
+    )
+
+
 def wrap_api_error(exc: Exception) -> AgentCliError:
     """Convert a databricks-sdk error (or any exception) into an `AgentCliError`.
 

integrations/agentbricks/src/databricks_agentkit/_api_client.py

@@ -15,7 +15,12 @@
 from typing import TYPE_CHECKING, Any, Optional
 from urllib.parse import quote
 
-from databricks_agentbricks.errors import TRANSIENT_ERROR_CODES, AgentCliError, wrap_api_error
+from databricks_agentbricks.errors import (
+    TRANSIENT_ERROR_CODES,
+    AgentCliError,
+    store_create_permission_error,
+    wrap_api_error,
+)
 from databricks_agentkit import models
 
 if TYPE_CHECKING:
@@ -282,16 +287,19 @@ def create_memory_store(
         retry_transient: bool = False,
     ) -> models.MemoryStore:
         body = _body(display_name=display_name, description=description)
-        return _as(
-            models.MemoryStore,
-            self._do(
+        try:
+            response = self._do(
                 "POST",
                 f"{_BASE}/memory-stores",
                 query={"managed_memory_store_id": display_name},
                 body=body,
                 safe_to_retry=retry_transient,
-            ),
-        )
+            )
+        except AgentCliError as exc:
+            if exc.error_code == "PERMISSION_DENIED":
+                raise store_create_permission_error(display_name, "memory", exc) from exc
+            raise
+        return _as(models.MemoryStore, response)
 
     def get_memory_store(self, name: str) -> models.MemoryStore:
         return _as(models.MemoryStore, self._do("GET", f"{_BASE}/{memory_store_path(name)}"))
@@ -458,16 +466,19 @@ def create_session_store(
         retry_transient: bool = False,
     ) -> models.SessionStore:
         body = _body(description=description, metadata=metadata)
-        return _as(
-            models.SessionStore,
-            self._do(
+        try:
+            response = self._do(
                 "POST",
                 f"{_BASE}/session-stores",
                 query={"session_store_name": name, "session_store_id": name},
                 body=body,
                 safe_to_retry=retry_transient,
-            ),
-        )
+            )
+        except AgentCliError as exc:
+            if exc.error_code == "PERMISSION_DENIED":
+                raise store_create_permission_error(name, "session", exc) from exc
+            raise
+        return _as(models.SessionStore, response)
 
     def get_session_store(self, name: str) -> models.SessionStore:
         return _as(models.SessionStore, self._do("GET", f"{_BASE}/{session_store_path(name)}"))

integrations/agentbricks/tests/unit_tests/client_test.py

@@ -168,6 +168,42 @@ def test_create_memory_store_uses_transitional_request(workspace_client):
     )
 
 
+@pytest.mark.parametrize("kind", ["memory", "session"])
+def test_store_create_permission_error_has_current_recovery_commands(kind):
+    workspace_client = mock.Mock()
+    workspace_client.api_client.do.side_effect = AgentCliError(
+        "Lakebase project creation is restricted", error_code="PERMISSION_DENIED"
+    )
+    client = _AgentBricksApiClient(workspace_client=workspace_client)
+
+    with pytest.raises(AgentCliError) as raised:
+        getattr(client, f"create_{kind}_store")("restricted-store")
+
+    error = raised.value
+    assert error.error_code == "PERMISSION_DENIED"
+    assert error.hint is not None
+    assert "Lakebase" in error.hint and "workspace admin" in error.hint
+    command = "memory" if kind == "memory" else "sessions"
+    assert f"agentbricks {command} bind <existing-store>" in error.hint
+    assert "--no-create-stores" not in error.hint
+    workspace_client.api_client.do.assert_called_once()
+
+
+@pytest.mark.parametrize("kind", ["memory", "session"])
+@pytest.mark.parametrize("code", ["ALREADY_EXISTS", "INVALID_PARAMETER_VALUE"])
+def test_store_create_does_not_rewrite_other_errors(kind, code):
+    workspace_client = mock.Mock()
+    workspace_client.api_client.do.side_effect = AgentCliError("original detail", error_code=code)
+    client = _AgentBricksApiClient(workspace_client=workspace_client)
+
+    with pytest.raises(AgentCliError, match="original detail") as raised:
+        getattr(client, f"create_{kind}_store")("example-store")
+
+    assert raised.value.error_code == code
+    assert raised.value.hint is None
+    workspace_client.api_client.do.assert_called_once()
+
+
 @mock.patch("databricks.sdk.WorkspaceClient")
 def test_create_runtime_store_uses_v2_api_and_app_owner(workspace_client):
     client, do = _client(workspace_client)

integrations/agentbricks/tests/unit_tests/deploy_test.py

@@ -290,9 +290,13 @@ def test_memory_store_client_distinguishes_inaccessible_existing_store():
     provider = mock.Mock(spec=ApiClientProvider)
     provider.get.return_value = api
 
-    with pytest.raises(AgentCliError, match="already exists but you don't have access"):
+    with pytest.raises(AgentCliError, match="already exists but you don't have access") as exc:
         MemoryStoreClient(provider).ensure("mem")
 
+    assert exc.value.hint is not None
+    assert "agentbricks memory bind <existing-store>" in exc.value.hint
+    assert "--no-create-stores" not in exc.value.hint
+
 
 def test_session_store_client_reuses_existing_store():
     api = mock.Mock()
@@ -318,6 +322,21 @@ def test_session_store_client_reports_permission_hint():
         SessionStoreClient(provider).ensure("sessions")
 
 
+def test_session_store_client_distinguishes_inaccessible_existing_store():
+    api = mock.Mock()
+    api.create_session_store.side_effect = AgentCliError("exists", error_code="ALREADY_EXISTS")
+    api.get_session_store.side_effect = AgentCliError("denied", error_code="PERMISSION_DENIED")
+    provider = mock.Mock(spec=ApiClientProvider)
+    provider.get.return_value = api
+
+    with pytest.raises(AgentCliError, match="already exists but you don't have access") as exc:
+        SessionStoreClient(provider).ensure("sessions")
+
+    assert exc.value.hint is not None
+    assert "agentbricks sessions bind <existing-store>" in exc.value.hint
+    assert "--no-create-stores" not in exc.value.hint
+
+
 def test_deploy_help_exposes_instance_count_and_routing_key():
     result = CliRunner().invoke(deploy_mod.deploy, ["--help"])