databricks-agentbricks 0.2.0__py3-none-any.whl

databricks_agentbricks/app_resources.py

@@ -0,0 +1,139 @@
+"""Lakebase/Apps resource plumbing for a deployed app.
+
+Binds Databricks Apps resources onto an app so its service principal gets platform-managed grants:
+a `postgres` resource for the legacy per-app Runtime Store and the tracing `experiment` resource.
+The service-managed Runtime Store path grants database access through Conversation Store instead.
+
+Managed-store (session/memory) table access is NOT granted here. The deployed app reaches those
+stores over the conversation-store REST API, which grants the app's service principal read/write
+server-side (see `deploy._grant_store_access`), so no direct Lakebase grant is needed.
+"""
+
+from __future__ import annotations
+
+import json
+import subprocess
+from dataclasses import dataclass
+from typing import Optional
+
+from databricks_agentbricks.databricks_cli import _databricks
+
+
+@dataclass(frozen=True)
+class LakebaseBackend:
+    """A Lakebase database and endpoint bound as a `postgres` app resource."""
+
+    project: str
+    branch: str
+    endpoint_id: str
+    database: str
+    schema: str
+    tables: tuple[str, ...]
+    resource_name: str  # the app-resource name (must be unique across an app's resources)
+
+    @property
+    def branch_path(self) -> str:
+        return f"projects/{self.project}/branches/{self.branch}"
+
+    @property
+    def database_path(self) -> str:
+        return f"{self.branch_path}/databases/{self.database}"
+
+    @property
+    def endpoint_path(self) -> str:
+        return f"{self.branch_path}/endpoints/{self.endpoint_id}"
+
+    def postgres_resource(self) -> dict:
+        """The `postgres` app-resource entry that grants the SP a Lakebase role + CONNECT."""
+        return {
+            "name": self.resource_name,
+            "postgres": {
+                "branch": self.branch_path,
+                "database": self.database_path,
+                "permission": "CAN_CONNECT_AND_CREATE",
+            },
+        }
+
+
+def _current_app_resources(app: str, profile: Optional[str]) -> list[dict]:
+    """Read the app's existing resources array (empty list if it can't be read)."""
+    result = _databricks(["apps", "get", app, "-o", "json"], profile, capture=True, check=False)
+    if result.returncode != 0:
+        return []
+    try:
+        resources = json.loads(result.stdout or "{}").get("resources", [])
+    except (json.JSONDecodeError, AttributeError):
+        return []
+    return resources if isinstance(resources, list) else []
+
+
+# The app-resource name for the trace experiment (unique across an app's resources, like a store's).
+_TRACE_EXPERIMENT_RESOURCE = "agentbricks-trace-experiment"
+
+
+def apply_experiment_resource(
+    app: str, experiment_id: str, profile: Optional[str]
+) -> Optional[str]:
+    """Bind the trace experiment as an `experiment` app resource so the SP can write traces.
+
+    This is the platform-managed grant: declaring the experiment as a `CAN_EDIT` resource lets the
+    app's service principal log traces to it (no manual SQL grant). Uses the same masked
+    read-modify-write as `apply_postgres_resources`: replace the whole resource array (preserving
+    every resource we don't own) while `update_mask` keeps the write from touching any other app
+    field. Returns None on success or a human-readable reason on failure.
+    """
+    ours = {
+        "name": _TRACE_EXPERIMENT_RESOURCE,
+        "experiment": {"experiment_id": experiment_id, "permission": "CAN_EDIT"},
+    }
+    preserved = [
+        r
+        for r in _current_app_resources(app, profile)
+        if isinstance(r, dict) and r.get("name") != _TRACE_EXPERIMENT_RESOURCE
+    ]
+    result = _update_app_resources(app, preserved + [ours], profile)
+    if result.returncode == 0:
+        return None
+    return (result.stderr or result.stdout or "").strip() or "unknown error"
+
+
+def apply_postgres_resources(
+    app: str, backends: list[LakebaseBackend], profile: Optional[str]
+) -> Optional[str]:
+    """Bind each backend's database as a `postgres` app resource in one update.
+
+    Writing the resources field REPLACES the whole array, so we must send the complete set:
+    read the app's current resources, drop the ones we manage (matched by name) so re-deploys
+    update rather than duplicate them, keep every other (user-owned) resource, and append ours.
+    Returns None on success or a human-readable reason on failure.
+    """
+    ours = [b.postgres_resource() for b in backends]
+    our_names = {r["name"] for r in ours}
+    preserved = [
+        r
+        for r in _current_app_resources(app, profile)
+        if isinstance(r, dict) and r.get("name") not in our_names
+    ]
+    result = _update_app_resources(app, preserved + ours, profile)
+    if result.returncode == 0:
+        return None
+    return (result.stderr or result.stdout or "").strip() or "unknown error"
+
+
+def _update_app_resources(
+    app: str, resources: list[dict], profile: Optional[str]
+) -> subprocess.CompletedProcess:
+    """Set the app's resources array without disturbing any other app field.
+
+    `update_mask` scopes the write to `resources` only. A bare `apps update` replaces the whole app
+    spec, so fields we don't send — notably `user_api_scopes` — would reset to their defaults every
+    deploy, silently breaking OBO. The masked `create-update` still replaces the resources array
+    (callers pass the complete desired set), but leaves everything else untouched.
+    """
+    payload = {"app": {"resources": resources}, "update_mask": "resources"}
+    return _databricks(
+        ["apps", "create-update", app, "--json", json.dumps(payload)],
+        profile,
+        capture=True,
+        check=False,
+    )

databricks_agentbricks/cli/__init__.py

@@ -0,0 +1,12 @@
+"""Agent Bricks CLI: the ``ab`` command and its subcommands.
+
+The console entrypoint (``project.scripts``) is ``databricks_agentbricks.cli:main``, re-exported here from
+``databricks_agentbricks.cli.app`` (the command tree) so that path stays stable. The entry function is named
+``app`` there rather than ``main`` so it does not shadow this package's re-exported ``main``. These
+modules are the command surface; the SDK and deployed-agent runtime live under
+``databricks_agentkit``.
+"""
+
+from databricks_agentbricks.cli.app import main
+
+__all__ = ["main"]

databricks_agentbricks/cli/app.py

@@ -0,0 +1,108 @@
+"""`ab` — the Databricks CLI for agent deployment, memory, and sessions.
+
+Root Click group. Global `--profile` and `--output` flow to every subcommand via
+`CliContext` on `ctx.obj`; subcommands build an authenticated API client on demand.
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import click
+
+from databricks_agentbricks import errors
+from databricks_agentbricks.cli.auth import load_default_profile, login, logout
+from databricks_agentbricks.cli.deploy import deploy, deployments
+from databricks_agentbricks.cli.dev import dev
+from databricks_agentbricks.cli.endpoint import endpoint
+from databricks_agentbricks.cli.help import configure_help
+from databricks_agentbricks.cli.init import init
+from databricks_agentbricks.cli.memory import memory
+from databricks_agentbricks.cli.sessions import sessions
+from databricks_agentbricks.cli.tools import tools
+from databricks_agentbricks.cli.tracing import tracing
+from databricks_agentkit._api_client import _AgentBricksApiClient
+
+
+class CliContext:
+    """Shared per-invocation state: selected profile, output mode, lazily-built client."""
+
+    def __init__(self, profile: Optional[str], output: str):
+        self.profile = profile
+        self.output = output
+        self._client: Optional[_AgentBricksApiClient] = None
+
+    def client(self) -> _AgentBricksApiClient:
+        if self._client is None:
+            self._client = _AgentBricksApiClient(self.profile)
+        return self._client
+
+
+@click.group(name="ab", context_settings={"help_option_names": ["-h", "--help"]})
+@click.option(
+    "--profile", "-p", default=None, help="~/.databrickscfg profile to authenticate with."
+)
+@click.option(
+    "--output",
+    "-o",
+    type=click.Choice(["text", "json"]),
+    default="text",
+    help="Output format (default: text).",
+)
+@click.version_option(package_name="databricks-agentbricks")
+@click.pass_context
+def agentbricks(ctx: click.Context, profile: Optional[str], output: str) -> None:
+    """Agent Bricks is a CLI for building and deploying custom AI agents on Databricks.
+
+    The Agent Bricks CLI is experimental: its commands and the underlying agent APIs are all in
+    preview, may need to be enabled for your workspace, and are likely to change in
+    backward-incompatible ways.
+
+    Scaffold an agent project from a template, run it locally with a chat UI, and deploy it to
+    Databricks Apps — then manage the tools, memory, sessions, and tracing behind it, all from one
+    authenticated command.
+
+    New here? The examples below take you from an empty directory to a deployed agent. Agent Bricks
+    authenticates with a Databricks profile: run `ab login` once to save a default, or pass
+    --profile / -p (without one, the Databricks SDK's default authentication is used).
+
+    Agents built with Agent Bricks combine the platform's capabilities:
+
+    \b
+      Models       Call Databricks model serving out of the box, routed through
+                   the AI Gateway for capacity on your existing Databricks auth.
+      Tools        Data sandboxes, managed MCP services, Unity Catalog
+                   functions, and local Python tools the agent can call.
+      Memory       Long-term memory the agent recalls across conversations.
+      Sessions     The transcript, history, and state of a single conversation.
+      Tracing      MLflow traces in Unity Catalog to debug and evaluate runs.
+      Deployment   Hosting on Databricks Apps, with scaling and sticky routing.
+
+    `ab deploy` provisions and wires these into a single agent hosted on Databricks Apps.
+    """
+    # Let errors render to match the selected output mode (JSON errors for -o json).
+    errors.set_output_mode(output)
+    ctx.obj = CliContext(profile=profile or load_default_profile(), output=output)
+
+
+agentbricks.add_command(login)
+agentbricks.add_command(logout)
+agentbricks.add_command(init)
+agentbricks.add_command(dev)
+agentbricks.add_command(memory)
+agentbricks.add_command(sessions)
+agentbricks.add_command(tracing)
+agentbricks.add_command(deploy)
+agentbricks.add_command(deployments)
+agentbricks.add_command(endpoint)
+agentbricks.add_command(tools)
+configure_help(agentbricks)
+
+
+def main() -> None:
+    # Click derives the display name from argv[0] so the `ab` script is shown in help output.
+    agentbricks()
+
+
+if __name__ == "__main__":
+    main()

databricks_agentbricks/cli/app_auth.py

@@ -0,0 +1,198 @@
+"""Preflight and conservative Apps scope reconciliation for request-user tools."""
+
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass
+from datetime import timedelta
+
+from databricks.sdk import WorkspaceClient
+from databricks.sdk.errors import DatabricksError, NotFound
+from databricks.sdk.service.apps import App, AppsAPI
+
+from databricks_agentbricks.agent_project import AgentProject
+from databricks_agentbricks.errors import AgentCliError
+from databricks_agentbricks.project_types import AgentServer
+
+_IDENTITY_DEFAULT_SCOPES = frozenset({"iam.access-control:read", "iam.current-user:read"})
+
+
+def _validate_forwarding(app: App) -> None:
+    if app.forward_user_access_token is False:
+        raise AgentCliError(
+            f"App '{app.name}' has forward_user_access_token=False; user auth requires forwarding.",
+            hint="Enable user access token forwarding in Databricks Apps, then stop and start "
+            "the App compute for the setting to take effect before retrying deployment.",
+        )
+
+
+def _validate_implicit_identity_scopes(app: App) -> None:
+    if app.user_api_scopes is not None:
+        return
+    effective = set(app.effective_user_api_scopes or [])
+    if effective - _IDENTITY_DEFAULT_SCOPES:
+        raise AgentCliError(
+            "Apps did not return configured scopes for an App with unexplained effective grants.",
+            hint="Inspect the App's scope configuration with its owner before retrying the scope "
+            "update.",
+        )
+
+
+def requires_user_auth(project: AgentProject | None) -> bool:
+    """Infer request-user auth from tool bindings before any deployment mutation."""
+    if project is None or not project.tools:
+        return False
+    managed = [
+        tool
+        for tool in project.tools
+        if tool.source.kind in ("mcp", "sandbox", "genie_one", "genie_agent")
+    ]
+    user_auth = any(tool.auth == "user" for tool in managed)
+    if user_auth and project.server != AgentServer.AGENTBRICKS:
+        raise AgentCliError(
+            "Managed tools with auth = 'user' require [agent].server = 'agentbricks'.",
+            hint="Migrate to the request-auth-aware Agent Bricks DurableAgentServer template before enabling user "
+            "auth. Failure recovery is unsupported for request-user attempts because the credential "
+            "is transient.",
+        )
+    if user_auth:
+        unspecified = [tool.id for tool in managed if tool.auth is None]
+        if unspecified:
+            raise AgentCliError(
+                "Request-user tools require explicit auth on every managed "
+                f"tool binding: {', '.join(unspecified)}.",
+                hint="Choose auth = 'app' to preserve legacy identity, or explicitly choose 'user'.",
+            )
+    return user_auth
+
+
+@dataclass(frozen=True)
+class AppUserScopeUpdatePlan:
+    apps: AppsAPI
+    name: str
+    existing_scopes: tuple[str, ...] | None
+    scopes: tuple[str, ...]
+
+
+def required_user_api_scopes(project: AgentProject | None) -> set[str]:
+    """Return Databricks Apps user API scopes required by request-user tools."""
+    scopes: set[str] = set()
+    # TODO: Extend this least-privilege mapping for each supported request-user tool kind/service.
+    for tool in project.tools if project else ():
+        if tool.auth != "user":
+            continue
+        if tool.source.kind in ("genie_one", "genie_agent"):
+            scopes.add("genie")
+            continue
+        if tool.source.kind not in ("mcp", "sandbox"):
+            continue
+        scopes.add("ai-gateway")
+        if tool.source.service == "system.ai.dbsql":
+            scopes.add("sql")
+        if tool.source.service == "system.ai.genie_one_mcp":
+            scopes.add("genie")
+        if tool.source.kind == "sandbox" and any(
+            scope.kind == "volume" for scope in tool.policy.downscope
+        ):
+            scopes.add("files")
+    return scopes
+
+
+def plan_app_user_scope_update(
+    name: str,
+    profile: str | None,
+    *,
+    allow_existing_app_update: bool,
+    required_scopes: set[str] | None = None,
+) -> AppUserScopeUpdatePlan:
+    """Plan user-scope creation or addition without changing the Databricks App.
+
+    A new App can be created with the requested user API scopes automatically. An existing App is
+    left unchanged unless all requested scopes are already present or the caller explicitly allows
+    Agent Bricks to add the missing scopes.
+    """
+    try:
+        apps = WorkspaceClient(profile=profile).apps
+        try:
+            existing = apps.get(name)
+        except NotFound:
+            existing = None
+    except (DatabricksError, ValueError) as exc:
+        raise AgentCliError(f"Could not read Apps user scopes for '{name}'.") from exc
+    if existing is not None:
+        _validate_forwarding(existing)
+    if existing is not None:
+        _validate_implicit_identity_scopes(existing)
+    configured = tuple(sorted(set(existing.user_api_scopes or []))) if existing else None
+    requested = {"ai-gateway"} if required_scopes is None else required_scopes
+    scopes = tuple(sorted({*(configured or ()), *requested}))
+    if existing is not None and scopes != configured and not allow_existing_app_update:
+        raise AgentCliError(
+            f"App '{name}' is missing required user API scopes; re-run with "
+            "--allow-user-scope-update to add them.",
+            hint="Review the existing App scopes and coordinate with other owners first. Agent Bricks "
+            "preserves unrelated scopes; later deploys do not need the flag once all required "
+            "scopes are present.",
+        )
+    return AppUserScopeUpdatePlan(apps=apps, name=name, existing_scopes=configured, scopes=scopes)
+
+
+def apply_app_user_scope_update(
+    plan: AppUserScopeUpdatePlan, *, instances: int | None = None, attempts: int = 12
+) -> None:
+    """Create a scoped App or add missing user API scopes to an existing App.
+
+    New Apps enable user-token forwarding at creation. Existing Apps retain unrelated scopes and
+    settings; the update mask includes only user API scopes and explicitly requested instance
+    fields. The read-before-write check detects known drift, but Apps has no compare-and-swap
+    contract, so owners must still coordinate concurrent updates. Scope removal is never automatic.
+    """
+    if not plan.scopes or not set(plan.existing_scopes or ()).issubset(plan.scopes):
+        raise AgentCliError(
+            "Automatic removal of Apps user scopes is unsupported.",
+            hint="Remove scopes explicitly in Databricks Apps and verify the effective scopes. "
+            "The SDK omits empty lists when serializing App; no removal was sent.",
+        )
+    desired = App(
+        name=plan.name,
+        user_api_scopes=list(plan.scopes),
+        forward_user_access_token=True if plan.existing_scopes is None else None,
+    )
+    mask = ["user_api_scopes"]
+    if instances is not None:
+        desired.compute_min_instances = instances
+        desired.compute_max_instances = instances
+        mask.extend(["compute_min_instances", "compute_max_instances"])
+    try:
+        if plan.existing_scopes is None:
+            plan.apps.create(desired)
+        else:
+            current = plan.apps.get(plan.name)
+            _validate_forwarding(current)
+            _validate_implicit_identity_scopes(current)
+            if tuple(sorted(set(current.user_api_scopes or []))) != plan.existing_scopes:
+                raise AgentCliError("Apps user scopes changed since preflight; review and retry.")
+            if plan.scopes != plan.existing_scopes or instances is not None:
+                plan.apps.create_update(plan.name, update_mask=",".join(mask), app=desired).result(
+                    timeout=timedelta(minutes=5)
+                )
+        for attempt in range(attempts):
+            current = plan.apps.get(plan.name)
+            _validate_forwarding(current)
+            desired_scopes = set(plan.scopes)
+            effective = set(current.effective_user_api_scopes or [])
+            if (
+                set(current.user_api_scopes or []) == desired_scopes
+                and desired_scopes.issubset(effective)
+                and effective.issubset(desired_scopes | _IDENTITY_DEFAULT_SCOPES)
+            ):
+                return
+            if attempt + 1 < attempts:
+                time.sleep(5)
+    except (DatabricksError, TimeoutError) as exc:
+        raise AgentCliError(f"Could not reconcile Apps user scopes for '{plan.name}'.") from exc
+    raise AgentCliError(
+        f"App '{plan.name}' effective user scopes did not converge after {attempts} checks.",
+        hint="Source deployment was stopped. Inspect requested/effective scopes in Databricks "
+        "Apps; contact your platform administrator if propagation remains blocked. Then retry.",
+    )

databricks_agentbricks/cli/auth.py

@@ -0,0 +1,142 @@
+"""`ab 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
+Agent Bricks' saved selection, not the underlying credentials. State lives in a small JSON file under
+`~/.agentbricks` (override the directory with `AGENTBRICKS_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_agentbricks import render
+from databricks_agentbricks.errors import AgentCliError
+from databricks_agentkit._api_client import _AgentBricksApiClient
+
+
+def _config_file() -> pathlib.Path:
+    base = os.environ.get("AGENTBRICKS_CONFIG_HOME")
+    root = pathlib.Path(base) if base else pathlib.Path.home() / ".agentbricks"
+    return root / "config.json"
+
+
+def load_default_profile() -> Optional[str]:
+    """The profile saved by `ab 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[_AgentBricksApiClient, str]:
+    client = _AgentBricksApiClient(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 Agent Bricks 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 `ab 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[_AgentBricksApiClient, str]:
+    # Local import: pulling databricks.sdk.errors loads the full SDK (~0.7s), which we defer off
+    # the CLI startup path. This function already builds a client, so the cost lands here anyway.
+    from databricks.sdk.errors import Unauthenticated
+
+    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 Agent Bricks 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. `ab 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=[
+            ("ab 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")