databricks-mason 0.1.2.dev0__tar.gz → 0.1.3.dev0__tar.gz

{databricks_mason-0.1.2.dev0 → databricks_mason-0.1.3.dev0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: databricks-mason
-Version: 0.1.2.dev0
+Version: 0.1.3.dev0
 Summary: Databricks integration for Mason
 Author-email: Databricks <agent-feedback@databricks.com>
 License-File: NOTICE
@@ -27,7 +27,7 @@ Requires-Dist: databricks-agents>=1.9.3; extra == 'runtime-openai'
 Requires-Dist: databricks-openai>=0.13.0; extra == 'runtime-openai'
 Requires-Dist: fastapi>=0.129.0; extra == 'runtime-openai'
 Requires-Dist: mlflow>=3.10.1; extra == 'runtime-openai'
-Requires-Dist: openai-agents>=0.4.1; extra == 'runtime-openai'
+Requires-Dist: openai-agents>=0.7.0; extra == 'runtime-openai'
 Requires-Dist: opentelemetry-exporter-otlp-proto-grpc>=1.25.0; extra == 'runtime-openai'
 Requires-Dist: uuid-utils>=0.10.0; extra == 'runtime-openai'
 Provides-Extra: tracing
@@ -62,23 +62,31 @@ For tracing commands, install Mason with tracing extras:
 pip install 'databricks-mason[tracing]'
 ```
 
+## Shell completion
+Add this to `~/.zshrc`:
+```sh
+eval "$(_MASON_COMPLETE=zsh_source mason)"
+```
+
 ## Authentication
 
 Mason uses [Databricks authentication](https://docs.databricks.com/aws/en/dev-tools/cli/authentication).
-If you do not already have credentials, authenticate a named profile first. You can
-then ask Mason to validate and remember that profile:
+Ask Mason to authenticate and remember a named profile:
 
 ```sh
-databricks auth login --profile <profile>
 mason login --profile <profile>
 mason sessions stores list
 ```
 
-`mason login` does not create credentials; it stores the selected profile in
-`~/.mason/config.json`. `mason logout` forgets that selection without revoking the
-underlying credentials. If Databricks SDK default authentication is already configured,
-you can skip `mason login`. You can also pass `--profile/-p` for an individual command.
-Use `--output json` for scripting.
+`mason login` validates existing credentials first. If credentials are missing or rejected in
+an interactive terminal, Mason runs `databricks auth login --profile <profile>`, revalidates the
+profile, and stores the selection in `~/.mason/config.json`. This browser-based setup requires
+the Databricks CLI. In non-interactive environments, authenticate the profile before running
+Mason. `mason logout` forgets the saved selection without revoking the underlying credentials.
+
+If Databricks SDK default authentication is already configured, you can skip `mason login`.
+You can also pass the global `--profile/-p` option before an individual command, for example
+`mason --profile <profile> mcp list`. Use `--output json` for scripting.
 
 ## Python SDK
 
@@ -153,7 +161,7 @@ mason sessions items append --help
 For the shortest path from a blank directory to a running and deployed agent:
 
 ```sh
-mason login --profile my-workspace
+mason login --profile <profile>
 mason init my-agent
 cd my-agent
 mason dev
@@ -176,9 +184,14 @@ mason tools add sandbox --scope table:samples.nyctaxi.trips
 mason tools add mcp system.ai.web_search
 mason tools add uc-function catalog.schema.lookup_ticket
 mason tools add python lookup-ticket
+mason tools remove mcp system.ai.web_search
 mason tools list
 ```
 
+For MCP services, the remove command accepts the same service name as the add command. You can also
+remove any binding by the ID shown in `mason tools list`, for example `mason tools remove
+web_search`. Removal updates only `agent.toml`; Python source and test files remain user-owned.
+
 Discover the MCP Services available to your user before adding one. By default Mason lists the
 Databricks-managed services in `system.ai`; pass `--schema catalog.schema` for another Unity Catalog
 schema. Text output includes a copyable add command, while `--output json` returns normalized service

{databricks_mason-0.1.2.dev0 → databricks_mason-0.1.3.dev0}/README.md

@@ -26,23 +26,31 @@ For tracing commands, install Mason with tracing extras:
 pip install 'databricks-mason[tracing]'
 ```
 
+## Shell completion
+Add this to `~/.zshrc`:
+```sh
+eval "$(_MASON_COMPLETE=zsh_source mason)"
+```
+
 ## Authentication
 
 Mason uses [Databricks authentication](https://docs.databricks.com/aws/en/dev-tools/cli/authentication).
-If you do not already have credentials, authenticate a named profile first. You can
-then ask Mason to validate and remember that profile:
+Ask Mason to authenticate and remember a named profile:
 
 ```sh
-databricks auth login --profile <profile>
 mason login --profile <profile>
 mason sessions stores list
 ```
 
-`mason login` does not create credentials; it stores the selected profile in
-`~/.mason/config.json`. `mason logout` forgets that selection without revoking the
-underlying credentials. If Databricks SDK default authentication is already configured,
-you can skip `mason login`. You can also pass `--profile/-p` for an individual command.
-Use `--output json` for scripting.
+`mason login` validates existing credentials first. If credentials are missing or rejected in
+an interactive terminal, Mason runs `databricks auth login --profile <profile>`, revalidates the
+profile, and stores the selection in `~/.mason/config.json`. This browser-based setup requires
+the Databricks CLI. In non-interactive environments, authenticate the profile before running
+Mason. `mason logout` forgets the saved selection without revoking the underlying credentials.
+
+If Databricks SDK default authentication is already configured, you can skip `mason login`.
+You can also pass the global `--profile/-p` option before an individual command, for example
+`mason --profile <profile> mcp list`. Use `--output json` for scripting.
 
 ## Python SDK
 
@@ -117,7 +125,7 @@ mason sessions items append --help
 For the shortest path from a blank directory to a running and deployed agent:
 
 ```sh
-mason login --profile my-workspace
+mason login --profile <profile>
 mason init my-agent
 cd my-agent
 mason dev
@@ -140,9 +148,14 @@ mason tools add sandbox --scope table:samples.nyctaxi.trips
 mason tools add mcp system.ai.web_search
 mason tools add uc-function catalog.schema.lookup_ticket
 mason tools add python lookup-ticket
+mason tools remove mcp system.ai.web_search
 mason tools list
 ```
 
+For MCP services, the remove command accepts the same service name as the add command. You can also
+remove any binding by the ID shown in `mason tools list`, for example `mason tools remove
+web_search`. Removal updates only `agent.toml`; Python source and test files remain user-owned.
+
 Discover the MCP Services available to your user before adding one. By default Mason lists the
 Databricks-managed services in `system.ai`; pass `--schema catalog.schema` for another Unity Catalog
 schema. Text output includes a copyable add command, while `--output json` returns normalized service

{databricks_mason-0.1.2.dev0 → databricks_mason-0.1.3.dev0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "databricks-mason"
-version = "0.1.2.dev0"
+version = "0.1.3.dev0"
 description = "Databricks integration for Mason"
 authors = [
     { name="Databricks", email="agent-feedback@databricks.com" },
@@ -38,7 +38,7 @@ runtime = [
     "databricks-agents>=1.9.3",
 ]
 runtime-openai = [
-    "openai-agents>=0.4.1",
+    "openai-agents>=0.7.0",
     "databricks-openai>=0.13.0",
     "fastapi>=0.129.0",
     "mlflow>=3.10.1",

{databricks_mason-0.1.2.dev0 → databricks_mason-0.1.3.dev0}/src/databricks_mason/agent_project.py

@@ -15,6 +15,7 @@ from tomlkit import TOMLDocument
 from tomlkit.exceptions import ParseError
 
 from databricks_mason.errors import AgentCliError
+from databricks_mason.runtime.tool_manifest import MEMORY_STORE_TABLE, SESSION_STORE_TABLE
 
 _SCHEMA_VERSION = 1
 _SUPPORTED_FRAMEWORKS = {"langgraph", "openai"}
@@ -191,6 +192,15 @@ def _required_string(value: object, description: str) -> str:
     return value
 
 
+def _store_name_from_manifest(value: object, table: str) -> str | None:
+    """Read the ``name`` from a ``[memory_store]`` / ``[session_store]`` table, or None if absent."""
+    if value is None:
+        return None
+    if not isinstance(value, Mapping):
+        raise AgentCliError(f"agent.toml [{table}] must be a table.")
+    return _required_string(cast(Mapping[str, Any], value).get("name"), f"[{table}] name")
+
+
 def _scope_from_manifest(value: object) -> Scope:
     if not isinstance(value, Mapping):
         raise AgentCliError("Sandbox downscope entries must be TOML tables.")
@@ -275,12 +285,17 @@ class AgentProject:
         document: TOMLDocument,
         framework: str,
         tools: list[ToolSpec],
+        memory_store: str | None = None,
+        session_store: str | None = None,
     ) -> None:
         self.root = root
         self.path = root / "agent.toml"
         self._document = document
         self.framework = framework
         self.tools = tools
+        # Managed store bindings declared in agent.toml; None = unbound.
+        self.memory_store = memory_store
+        self.session_store = session_store
 
     @classmethod
     def load(cls, root: pathlib.Path | str) -> "AgentProject":
@@ -314,7 +329,13 @@ class AgentProject:
         ids = [tool.id for tool in tools]
         if len(ids) != len(set(ids)):
             raise AgentCliError("agent.toml tool ids must be unique.")
-        return cls(project_root, document, framework, tools)
+        memory_store = _store_name_from_manifest(
+            document.get(MEMORY_STORE_TABLE), MEMORY_STORE_TABLE
+        )
+        session_store = _store_name_from_manifest(
+            document.get(SESSION_STORE_TABLE), SESSION_STORE_TABLE
+        )
+        return cls(project_root, document, framework, tools, memory_store, session_store)
 
     @classmethod
     def create(cls, root: pathlib.Path | str, *, framework: str) -> "AgentProject":
@@ -367,6 +388,44 @@ class AgentProject:
         del self.tools[index]
         return True
 
+    def bind_memory_store(self, name: str) -> bool:
+        """Declare the memory store binding in agent.toml. Returns True if it changed."""
+        return self._set_store(MEMORY_STORE_TABLE, name)
+
+    def bind_session_store(self, name: str) -> bool:
+        """Declare the session store binding in agent.toml. Returns True if it changed."""
+        return self._set_store(SESSION_STORE_TABLE, name)
+
+    def unbind_memory_store(self) -> bool:
+        """Remove the memory store binding from agent.toml. Returns True if it was present."""
+        return self._clear_store(MEMORY_STORE_TABLE)
+
+    def unbind_session_store(self) -> bool:
+        """Remove the session store binding from agent.toml. Returns True if it was present."""
+        return self._clear_store(SESSION_STORE_TABLE)
+
+    def _set_store(self, table: str, name: str) -> bool:
+        name = _required_string(name, f"[{table}] name")
+        if getattr(self, table) == name:
+            return False
+        existing = self._document.get(table)
+        if isinstance(existing, Mapping):
+            existing["name"] = name
+        else:
+            store_table = tomlkit.table()
+            store_table.add("name", name)
+            self._document.append(table, store_table)
+        setattr(self, table, name)
+        return True
+
+    def _clear_store(self, table: str) -> bool:
+        if getattr(self, table) is None:
+            return False
+        if table in self._document:
+            del self._document[table]
+        setattr(self, table, None)
+        return True
+
     def write(self) -> pathlib.Path:
         self.path.parent.mkdir(parents=True, exist_ok=True)
         temporary: pathlib.Path | None = None

databricks_mason-0.1.3.dev0/src/databricks_mason/auth.py

@@ -0,0 +1,139 @@
+"""`mason login` / `logout` — remember an optional Databricks profile.
+
+`login` validates a named profile and persists the selection; when credentials are missing or
+rejected in an interactive terminal, it delegates setup to `databricks auth login` and retries.
+The root group falls back to the saved profile whenever `-p` is omitted. Without a saved profile,
+the Databricks SDK performs its normal default authentication resolution. `logout` removes only
+Mason's saved selection, not the underlying credentials. State lives in a small JSON file under
+`~/.mason` (override the directory with `MASON_CONFIG_HOME`, mainly for tests).
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import pathlib
+import subprocess
+import sys
+from typing import Optional
+
+import click
+from databricks.sdk.errors import Unauthenticated
+
+from databricks_mason import render
+from databricks_mason.client import MasonClient
+from databricks_mason.errors import AgentCliError
+
+
+def _config_file() -> pathlib.Path:
+    base = os.environ.get("MASON_CONFIG_HOME")
+    root = pathlib.Path(base) if base else pathlib.Path.home() / ".mason"
+    return root / "config.json"
+
+
+def load_default_profile() -> Optional[str]:
+    """The profile saved by `mason login`, or None if the user never logged in."""
+    try:
+        return json.loads(_config_file().read_text()).get("profile")
+    except (OSError, json.JSONDecodeError):
+        return None
+
+
+def _save_default_profile(profile: str) -> None:
+    path = _config_file()
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(json.dumps({"profile": profile}, indent=2) + "\n")
+
+
+def _validate_profile(profile: str) -> tuple[MasonClient, str]:
+    client = MasonClient(profile)
+    return client, client.current_user
+
+
+def _is_interactive() -> bool:
+    return sys.stdin.isatty()
+
+
+def _run_databricks_login(profile: str) -> None:
+    command = ["databricks", "auth", "login", "--profile", profile]
+    try:
+        # Keep the child process interactive while preserving stdout for Mason's JSON output.
+        result = subprocess.run(command, text=True, check=False, stdout=sys.stderr)
+    except FileNotFoundError as exc:
+        raise AgentCliError(
+            "Could not configure Databricks authentication: the `databricks` CLI was not found.",
+            hint=f"Install the Databricks CLI, then retry `mason login --profile {profile}`.",
+        ) from exc
+    if result.returncode != 0:
+        raise AgentCliError(
+            f"`databricks auth login --profile {profile}` failed (exit {result.returncode})."
+        )
+
+
+def _authenticate_profile(profile: str) -> tuple[MasonClient, str]:
+    try:
+        return _validate_profile(profile)
+    except (AgentCliError, Unauthenticated) as initial_error:
+        if not _is_interactive():
+            raise AgentCliError(
+                f"Could not validate Databricks profile {profile!r}: {initial_error}",
+                hint="Run this command in an interactive terminal so Mason can open "
+                "Databricks login, or authenticate first with "
+                f"`databricks auth login --profile {profile}`.",
+            ) from initial_error
+    except Exception as validation_error:  # noqa: BLE001 - normalize unexpected API failures
+        raise AgentCliError(
+            f"Could not validate Databricks profile {profile!r}: {validation_error}"
+        ) from validation_error
+
+    _run_databricks_login(profile)
+    try:
+        return _validate_profile(profile)
+    except Exception as retry_error:  # noqa: BLE001 - normalize the post-login failure
+        raise AgentCliError(
+            f"Databricks login completed, but profile {profile!r} could not be validated: "
+            f"{retry_error}"
+        ) from retry_error
+
+
+@click.command()
+@click.option(
+    "--profile",
+    "-p",
+    default=None,
+    help="Profile to authenticate with and remember as the default.",
+)
+@click.pass_obj
+def login(obj, profile) -> None:
+    """Authenticate a profile and save it as the default, so later commands can omit -p."""
+    profile = profile or obj.profile
+    if not profile:
+        raise AgentCliError(
+            "No profile to save.",
+            hint="Pass one to remember, e.g. `mason login --profile <profile>`.",
+        )
+    client, user = _authenticate_profile(profile)
+    _save_default_profile(profile)
+    if obj.output == "json":
+        render.emit_json({"profile": profile, "user": user, "host": client.host})
+        return
+    render.success(
+        f"Logged in as {user}",
+        fields={"Profile": profile, "Host": client.host},
+        next_steps=[
+            ("mason init my-agent", "Scaffold a new agent project"),
+        ],
+    )
+
+
+@click.command()
+@click.pass_obj
+def logout(obj) -> None:
+    """Forget the saved profile selection without deleting its credentials."""
+    path = _config_file()
+    existed = path.exists()
+    path.unlink(missing_ok=True)
+    if obj.output == "json":
+        render.emit_json({"logged_out": existed})
+        return
+    render.success("Logged out" if existed else "No saved login to clear")

{databricks_mason-0.1.2.dev0 → databricks_mason-0.1.3.dev0}/src/databricks_mason/client.py

@@ -14,16 +14,23 @@ from __future__ import annotations
 import configparser
 import os
 import pathlib
+import time
 from typing import Any, Optional
 
 from databricks.sdk import WorkspaceClient
 
 from databricks_mason import models
-from databricks_mason.errors import AgentCliError, wrap_api_error
+from databricks_mason.errors import TRANSIENT_ERROR_CODES, AgentCliError, wrap_api_error
 
 _BASE = "/api/agents/v1"
 _MCP_SERVICES_PATH = "/api/2.1/unity-catalog/mcp-services"
 
+# Transient backend failures (e.g. a CANCELLED RPC) usually clear on a retry, so retry safe
+# requests once before surfacing them. Mutating requests must opt in explicitly: their first
+# attempt may have committed even when its response was lost.
+_MAX_ATTEMPTS = 2
+_RETRY_BASE_DELAY_S = 0.2
+
 
 def _query(**kwargs: Any) -> dict[str, Any]:
     """Build a query dict, dropping None and empty values."""
@@ -118,8 +125,8 @@ class MasonClient:
         except Exception as exc:  # noqa: BLE001 - surfaced as a clean CLI error
             raise AgentCliError(
                 f"Could not initialize Databricks auth: {exc}",
-                hint="Check your profile (`databricks auth login --profile <name>`) "
-                "or pass --profile.",
+                hint="Select an existing profile with `mason --profile <name> <command>` "
+                "or authenticate and save it with `mason login --profile <name>`.",
             ) from exc
 
     @property
@@ -140,12 +147,27 @@ class MasonClient:
         self._w.workspace.mkdirs(path)
 
     def _do(
-        self, method: str, path: str, *, query: Optional[dict] = None, body: Optional[dict] = None
+        self,
+        method: str,
+        path: str,
+        *,
+        query: Optional[dict] = None,
+        body: Optional[dict] = None,
+        safe_to_retry: bool = False,
     ) -> Any:
-        try:
-            return self._w.api_client.do(method, path, query=query, body=body)
-        except Exception as exc:  # noqa: BLE001 - normalized to AgentCliError
-            raise wrap_api_error(exc) from exc
+        retry_allowed = method == "GET" or safe_to_retry
+        delay = _RETRY_BASE_DELAY_S
+        for attempt in range(1, _MAX_ATTEMPTS + 1):
+            try:
+                return self._w.api_client.do(method, path, query=query, body=body)
+            except Exception as exc:  # noqa: BLE001 - normalized to AgentCliError
+                retryable = (
+                    retry_allowed and getattr(exc, "error_code", None) in TRANSIENT_ERROR_CODES
+                )
+                if not retryable or attempt == _MAX_ATTEMPTS:
+                    raise wrap_api_error(exc) from exc
+                time.sleep(delay)
+                delay *= 2
 
     # --- Unity Catalog MCP Services -----------------------------------------
 
@@ -162,10 +184,22 @@ class MasonClient:
     # --- memory stores -------------------------------------------------------
 
     def create_memory_store(
-        self, display_name: str, description: Optional[str] = None
+        self,
+        display_name: str,
+        description: Optional[str] = None,
+        *,
+        retry_transient: bool = False,
     ) -> models.MemoryStore:
         body = _query(display_name=display_name, description=description)
-        return _as(models.MemoryStore, self._do("POST", f"{_BASE}/memory-stores", body=body))
+        return _as(
+            models.MemoryStore,
+            self._do(
+                "POST",
+                f"{_BASE}/memory-stores",
+                body=body,
+                safe_to_retry=retry_transient,
+            ),
+        )
 
     def get_memory_store(self, name: str) -> models.MemoryStore:
         return _as(models.MemoryStore, self._do("GET", f"{_BASE}/{memory_store_path(name)}"))
@@ -262,7 +296,12 @@ class MasonClient:
         body = _query(actor_id=actor_id, query=query, limit=limit)
         return _as(
             models.MemorySearchResult,
-            self._do("POST", f"{_BASE}/{memory_store_path(store)}/entries:search", body=body),
+            self._do(
+                "POST",
+                f"{_BASE}/{memory_store_path(store)}/entries:search",
+                body=body,
+                safe_to_retry=True,
+            ),
         )
 
     def update_memory_entry(
@@ -286,13 +325,22 @@ class MasonClient:
     # --- session stores ------------------------------------------------------
 
     def create_session_store(
-        self, name: str, description: Optional[str] = None, metadata: Optional[dict] = None
+        self,
+        name: str,
+        description: Optional[str] = None,
+        metadata: Optional[dict] = None,
+        *,
+        retry_transient: bool = False,
     ) -> models.SessionStore:
         body = _query(description=description, metadata=metadata)
         return _as(
             models.SessionStore,
             self._do(
-                "POST", f"{_BASE}/session-stores", query={"session_store_name": name}, body=body
+                "POST",
+                f"{_BASE}/session-stores",
+                query={"session_store_name": name},
+                body=body,
+                safe_to_retry=retry_transient,
             ),
         )