diff --git a/integrations/agentbricks/src/databricks_agentbricks/clients/conversation_store_client.py b/integrations/agentbricks/src/databricks_agentbricks/clients/conversation_store_client.py index 0444551d7..2483457db 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/clients/conversation_store_client.py +++ b/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 ` 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: diff --git a/integrations/agentbricks/src/databricks_agentbricks/errors.py b/integrations/agentbricks/src/databricks_agentbricks/errors.py index 5bb5d6616..4d3615031 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/errors.py +++ b/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 ` and deploy again." + ), + ) + + def wrap_api_error(exc: Exception) -> AgentCliError: """Convert a databricks-sdk error (or any exception) into an `AgentCliError`. diff --git a/integrations/agentbricks/src/databricks_agentkit/_api_client.py b/integrations/agentbricks/src/databricks_agentkit/_api_client.py index 38b22eb25..59618649d 100644 --- a/integrations/agentbricks/src/databricks_agentkit/_api_client.py +++ b/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)}")) diff --git a/integrations/agentbricks/tests/unit_tests/client_test.py b/integrations/agentbricks/tests/unit_tests/client_test.py index 4576da60b..fbdcaf205 100644 --- a/integrations/agentbricks/tests/unit_tests/client_test.py +++ b/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 " 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) diff --git a/integrations/agentbricks/tests/unit_tests/deploy_test.py b/integrations/agentbricks/tests/unit_tests/deploy_test.py index 060dded7d..ae6003de9 100644 --- a/integrations/agentbricks/tests/unit_tests/deploy_test.py +++ b/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 " 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 " 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"])