diff --git a/integrations/mason/src/databricks_mason/client.py b/integrations/mason/src/databricks_mason/client.py index 1b584f91..77330fa8 100644 --- a/integrations/mason/src/databricks_mason/client.py +++ b/integrations/mason/src/databricks_mason/client.py @@ -19,7 +19,7 @@ from databricks.sdk import WorkspaceClient from databricks_mason import models -from databricks_mason.errors import AgentCliError, wrap_api_error +from databricks_mason.errors import AgentCliError, auth_hint, wrap_api_error _BASE = "/api/agents/v1" _MCP_SERVICES_PATH = "/api/2.1/unity-catalog/mcp-services" @@ -118,8 +118,7 @@ def __init__(self, profile: Optional[str] = None): except Exception as exc: # noqa: BLE001 - surfaced as a clean CLI error raise AgentCliError( f"Could not initialize Databricks auth: {exc}", - hint="Select an existing profile with `mason --profile ` " - "or authenticate and save it with `mason login --profile `.", + hint=auth_hint(profile), ) from exc @property diff --git a/integrations/mason/src/databricks_mason/errors.py b/integrations/mason/src/databricks_mason/errors.py index c9ec14a5..ffd12012 100644 --- a/integrations/mason/src/databricks_mason/errors.py +++ b/integrations/mason/src/databricks_mason/errors.py @@ -7,6 +7,7 @@ from __future__ import annotations import json +import os from typing import Optional import click @@ -72,3 +73,28 @@ def wrap_api_error(exc: Exception) -> AgentCliError: message = str(exc).strip() or exc.__class__.__name__ hint = _PREVIEW_HINT if error_code in _PREVIEW_ERROR_CODES else None return AgentCliError(message, error_code=error_code, hint=hint) + + +def auth_hint(profile: Optional[str] = None) -> str: + """The remediation to append to any Databricks auth failure. + + We deliberately do not parse the SDK's error text — its wording varies by version and + failure mode, and the message is already surfaced verbatim as the error itself (see + ``MasonClient.__init__``). This mirrors ``wrap_api_error``, which passes the upstream + message through and only adds a supplemental hint. Here that hint is a single, + always-valid next step — ``mason login``, which both authenticates and remembers the + profile — plus a note for the one case the SDK message rarely makes obvious: a stray + ``DATABRICKS_TOKEN`` in the environment silently overrides profile auth. That token + check is a structural environment probe, not string matching on the error. + """ + sel = f"--profile {profile}" if profile else "--profile " + lead = "" + if os.environ.get("DATABRICKS_TOKEN"): + lead = ( + "A DATABRICKS_TOKEN is set and overrides profile auth — `unset DATABRICKS_TOKEN` " + "if you meant to use a profile. " + ) + return ( + f"{lead}Run `mason login {sel}` to authenticate " + "(list profiles with `databricks auth profiles`)." + ) diff --git a/integrations/mason/templates/agent-langgraph/agent/agent.py b/integrations/mason/templates/agent-langgraph/agent/agent.py index d0d02499..866f589b 100644 --- a/integrations/mason/templates/agent-langgraph/agent/agent.py +++ b/integrations/mason/templates/agent-langgraph/agent/agent.py @@ -1,5 +1,4 @@ import logging -import os from collections.abc import AsyncGenerator, AsyncIterator from typing import Any @@ -59,16 +58,10 @@ def _check_databricks_auth() -> None: try: workspace_client() except Exception as e: - profile = os.getenv("DATABRICKS_CONFIG_PROFILE") - target = ( - f"profile {profile!r}" if profile else "the DEFAULT profile / DATABRICKS_HOST+TOKEN" - ) raise RuntimeError( - f"Databricks auth is not configured — the agent can't call the model. Tried {target}.\n" - "Fix one of:\n" - " • set DATABRICKS_CONFIG_PROFILE in .env to a profile from `databricks auth profiles`, or\n" - " • run `databricks auth login --profile ` to create one, or\n" - " • set DATABRICKS_HOST and DATABRICKS_TOKEN in .env.\n" + "Databricks auth is not configured — the agent can't call the model.\n" + "Log in with `mason login --profile `, then set " + "DATABRICKS_CONFIG_PROFILE= in .env (or set DATABRICKS_HOST and DATABRICKS_TOKEN).\n" f"(underlying error: {e})" ) from e diff --git a/integrations/mason/templates/agent-openai/agent/agent.py b/integrations/mason/templates/agent-openai/agent/agent.py index ebc97dd7..24d8cbec 100644 --- a/integrations/mason/templates/agent-openai/agent/agent.py +++ b/integrations/mason/templates/agent-openai/agent/agent.py @@ -1,5 +1,4 @@ import logging -import os from collections.abc import AsyncGenerator from contextlib import AsyncExitStack from typing import Any @@ -56,16 +55,10 @@ def _check_databricks_auth() -> None: try: workspace_client() except Exception as e: - profile = os.getenv("DATABRICKS_CONFIG_PROFILE") - target = ( - f"profile {profile!r}" if profile else "the DEFAULT profile / DATABRICKS_HOST+TOKEN" - ) raise RuntimeError( - f"Databricks auth is not configured — the agent can't call the model. Tried {target}.\n" - "Fix one of:\n" - " • set DATABRICKS_CONFIG_PROFILE in .env to a profile from `databricks auth profiles`, or\n" - " • run `databricks auth login --profile ` to create one, or\n" - " • set DATABRICKS_HOST and DATABRICKS_TOKEN in .env.\n" + "Databricks auth is not configured — the agent can't call the model.\n" + "Log in with `mason login --profile `, then set " + "DATABRICKS_CONFIG_PROFILE= in .env (or set DATABRICKS_HOST and DATABRICKS_TOKEN).\n" f"(underlying error: {e})" ) from e diff --git a/integrations/mason/tests/unit_tests/errors_test.py b/integrations/mason/tests/unit_tests/errors_test.py index 7ddeb102..9cec6438 100644 --- a/integrations/mason/tests/unit_tests/errors_test.py +++ b/integrations/mason/tests/unit_tests/errors_test.py @@ -7,7 +7,7 @@ import pytest from databricks_mason import errors -from databricks_mason.errors import AgentCliError +from databricks_mason.errors import AgentCliError, auth_hint @pytest.fixture(autouse=True) @@ -43,3 +43,26 @@ def test_error_renders_text_in_text_mode(capsys): # Not JSON in text mode. with pytest.raises(json.JSONDecodeError): json.loads(err) + + +def test_auth_hint_steers_to_a_single_mason_login(monkeypatch): + monkeypatch.delenv("DATABRICKS_TOKEN", raising=False) + hint = auth_hint("my-workspace") + assert "mason login --profile my-workspace" in hint + assert "databricks auth profiles" in hint + # No stray-token note when the env var is absent. + assert "DATABRICKS_TOKEN" not in hint + + +def test_auth_hint_uses_placeholder_without_profile(monkeypatch): + monkeypatch.delenv("DATABRICKS_TOKEN", raising=False) + assert "mason login --profile " in auth_hint() + + +def test_auth_hint_flags_a_stray_token(monkeypatch): + monkeypatch.setenv("DATABRICKS_TOKEN", "dapiXXXX") + hint = auth_hint("my-workspace") + # Calls out the override that the SDK message alone rarely makes obvious... + assert "unset DATABRICKS_TOKEN" in hint + # ...while still funneling to the one command. + assert "mason login --profile my-workspace" in hint