diff --git a/integrations/agentbricks/README.md b/integrations/agentbricks/README.md index ae062a2ac..f72db48a3 100644 --- a/integrations/agentbricks/README.md +++ b/integrations/agentbricks/README.md @@ -77,7 +77,7 @@ Custom servers own their authentication and approval protocol. runtime. The store/session/tools commands and `agentbricks tracing bind`/`unbind` don't need it; `agentbricks tracing list`/`get` do, to read `agentbricks dev`'s local trace store. - **[Databricks CLI](https://docs.databricks.com/dev-tools/cli/)** — needed for - browser-based `agentbricks login`. If a profile is already authenticated, the CLI uses it + browser-based `agentbricks auth login`. If a profile is already authenticated, the CLI uses it directly and the Databricks CLI is optional. ## Installation @@ -105,23 +105,41 @@ eval "$(_AGENTBRICKS_COMPLETE=zsh_source agentbricks)" ## Authentication Agent Bricks CLI uses [Databricks authentication](https://docs.databricks.com/aws/en/dev-tools/cli/authentication). -Ask the CLI to authenticate and remember a named profile: +Create a project with the profile you want it to use (replace `my-profile` with your profile name): ```sh -agentbricks login --profile -agentbricks sessions stores list +agentbricks init my-agent --profile my-profile ``` -`agentbricks login` validates existing credentials first. If credentials are missing or rejected in -an interactive terminal, the CLI runs `databricks auth login --profile `, revalidates the -profile, and stores the selection in the existing `~/.agentbricks/config.json` state file. This -browser-based setup requires the Databricks CLI. In non-interactive environments, authenticate the -profile before running `agentbricks`. `agentbricks logout` forgets the saved selection without revoking the underlying -credentials. - -If Databricks SDK default authentication is already configured, you can skip `agentbricks login`. -You can also pass the global `--profile/-p` option before an individual command, for example -`agentbricks --profile tools list`. Use `--output json` for scripting. +`init` records the profile in the new project's `.env`. In an interactive terminal, it checks +whether you are signed in and runs `databricks auth login --profile my-profile` if needed. In CI or +without a terminal, configure authentication beforehand; `init` does not open a login flow there. +Use `agentbricks auth describe` in the project to see which profile and workspace a command would +use. `agentbricks auth login ` signs in without changing the project's choice; to change +that choice, run `agentbricks auth switch ` in the project. An interactive `auth switch` +also signs in if needed. + +Inside an Agent Bricks project (a directory with `.agentbricks/project.toml` or `app.yaml`), every +command uses the profile in the project's `.env`. Outside a project, pass the global `--profile/-p` +option before a command, for example `agentbricks --profile tools list`, or set +`DATABRICKS_CONFIG_PROFILE`. `-p` can go before or after the command — +`agentbricks tools list -p ` works too (the trailing form goes after the full subcommand). +Use `--output json` for scripting. + +The profile is resolved in this order, most to least specific: the `--profile/-p` flag, the project +`.env`'s `DATABRICKS_CONFIG_PROFILE` (inside an Agent Bricks project; `dev` and `deploy` resolve +against their source directory so the CLI and the locally running agent use one profile), the +`DATABRICKS_CONFIG_PROFILE` environment variable, and finally the Databricks SDK's default +authentication resolution — which also covers `DATABRICKS_HOST`/`DATABRICKS_TOKEN` environment +authentication, not only `[DEFAULT]` in `~/.databrickscfg`. `dev` passes the resolved profile to +the agent (unless `.env` sets `DATABRICKS_TOKEN` or `DATABRICKS_CLIENT_ID`/ +`DATABRICKS_CLIENT_SECRET`). Every command that reaches the workspace prints one stderr line naming +the profile, where it came from, and its host, so the target workspace is never implicit. + +Selecting a profile preserves the SDK's precedence for exported credential fields such as +`DATABRICKS_HOST` and `DATABRICKS_TOKEN`; the workspace notice shows the environment host when it +overrides the profile's host. With no selected profile or environment credentials, the default +chosen with `databricks auth switch` takes precedence over `[DEFAULT]`. The CLI sends best effort command usage telemetry after leaf commands. Its log records the static command path, package version, operating system, duration, and process outcome (exit code and, for @@ -143,7 +161,7 @@ messages are omitted. Set Telemetry runs only when a leaf command is invoked, whether its callback succeeds or fails. Help and errors rejected during Click parsing, before a leaf starts, do not produce an event. Uploads reuse the command's workspace client or let the Databricks SDK resolve the selected or default -profile. Authenticate first with `agentbricks login --profile ` or configure SDK auth; +profile. Authenticate first with `agentbricks auth login ` or configure SDK auth; an SDK `external-browser` profile may prompt during telemetry even for local commands such as `init`. The CLI waits up to two seconds for client setup, authentication, and upload on a daemon thread, returning sooner if the worker finishes. This is not a delivery guarantee: unfinished @@ -152,16 +170,27 @@ failed upload never change the command result. ## Quickstart -The shortest path from a blank directory to a running and deployed agent: +A typical path from a blank directory to a running and deployed agent: ```sh -agentbricks login --profile -agentbricks init my-agent # Defaults to LangGraph; add --framework=openai for OpenAI Agents SDK +agentbricks init my-agent --profile my-profile # LangGraph by default; add --framework=openai for OpenAI Agents SDK cd my-agent +agentbricks auth describe # optional: check the selected profile and workspace agentbricks dev # optional: develop and test locally agentbricks deploy # deploy to Databricks as agent-bricks-my-agent ``` +Replace `my-profile` with the Databricks profile for your workspace. In an interactive terminal, +`init` signs in to it if needed. + +To move an existing project to a different workspace later, run `agentbricks auth switch ` in +its directory. Rerunning `agentbricks init` on a project it scaffolded changes nothing. (`--existing` +is different: it prepares a migration for a project Agent Bricks didn't generate.) + +`dev` and `deploy` validate the resolved profile's authentication up front and fail fast with the +`agentbricks auth login` command to run — an unauthenticated external-browser profile would +otherwise hang on a browser login that never completes. + `agentbricks dev` runs the agent locally on `http://localhost:8000`, wrapping the Databricks Apps local runtime so local behavior matches a deployment. diff --git a/integrations/agentbricks/cli.md b/integrations/agentbricks/cli.md index a109dabde..285278f5e 100644 --- a/integrations/agentbricks/cli.md +++ b/integrations/agentbricks/cli.md @@ -27,14 +27,47 @@ See [Installation](README.md#installation) for installation details and shell co ## Authentication Agent Bricks CLI authenticates with a [Databricks configuration profile](https://docs.databricks.com/aws/en/dev-tools/cli/authentication). -Run `agentbricks login` once to save a default profile, or pass `--profile` / `-p` on any command. Without -a profile, the Databricks SDK's default authentication resolution is used. See +Start with `agentbricks init my-agent --profile my-profile` (using your profile name), then run +`agentbricks dev` and `agentbricks deploy` from `my-agent`. `init` records the profile in the +project's `.env` and signs in if needed in an interactive terminal. In CI or without a terminal, +configure authentication beforehand. Run `agentbricks auth describe` to check the selected profile +and workspace. +`agentbricks auth login ` signs in without changing the project; +`agentbricks auth switch ` changes its `.env` profile and signs in if needed in an +interactive terminal. Use `--profile` / `-p` to override the choice for one command. See the +[Quickstart](README.md#quickstart) for the full journey. + +Inside an Agent Bricks project (a directory with `.agentbricks/project.toml` or `app.yaml`) every +command uses that project's `.env` profile, and so does a command run from any folder inside it (the +nearest parent with `.agentbricks/project.toml`). Agent Bricks saves no separate default profile +outside the project. The profile is resolved in this order, most to least specific: + +1. `--profile` / `-p` +2. the project `.env`'s `DATABRICKS_CONFIG_PROFILE` — only inside an Agent Bricks project; `dev` and + `deploy` resolve against their `--source` directory so the CLI and the locally running agent + (which reads the same `.env`) use one profile +3. the `DATABRICKS_CONFIG_PROFILE` environment variable +4. none — the Databricks SDK's default authentication resolution, including environment + credentials and the default set with `databricks auth switch`, then `[DEFAULT]` + +Selecting a profile preserves the SDK's precedence for exported credential fields such as +`DATABRICKS_HOST` and `DATABRICKS_TOKEN`; the workspace notice reports an overriding environment +host. An `auth_type = external-browser` profile must sign in with `agentbricks auth login `. + +`dev` also passes the resolved profile to the agent (unless `.env` sets +`DATABRICKS_TOKEN` or `DATABRICKS_CLIENT_ID`/`DATABRICKS_CLIENT_SECRET`, in which case the agent +uses those credentials). Every command that reaches the workspace prints one line to stderr naming +the profile, where it came from, and its host, e.g. +`Using profile 'ml' (from .env) · https://my-workspace.cloud.databricks.com` (text output only). See [Authentication](README.md#authentication) for details. ## Global options These options apply to every command. Pass them before the command name, for example -`agentbricks -p my-profile -o json sessions stores list`. +`agentbricks -p my-profile -o json sessions stores list`. `--profile` / `-p` can also follow the +command (`agentbricks sessions stores list -p my-profile`); given in both places, the later one wins. +A trailing `-p` goes after the full subcommand, not after a group name — `agentbricks auth +-p my-profile login` isn't recognized; write `agentbricks auth login -p my-profile`. | Option | Values | Default | Description | | --- | --- | --- | --- | @@ -45,6 +78,10 @@ These options apply to every command. Pass them before the command name, for exa ## How to read this reference +- Every leaf command accepts `--profile ` / `-p ` after its full command path, + as described under [Global options](#global-options). The option tables below list options + specific to each command; `init` lists its own `--profile` because it also writes the selected + profile to the new project's `.env`. - **Arguments** are positional; **options** are named flags. In a synopsis, `NAME` is required and `[NAME]` is optional; `` marks a group that requires a subcommand; `[options]` stands in for the option list documented below each command. @@ -57,9 +94,8 @@ These options apply to every command. Pass them before the command name, for exa | Command | Description | | --- | --- | -| [`login`](#agentbricks-login) | Authenticate and save a default profile | -| [`logout`](#agentbricks-logout) | Forget the saved default profile | | [`init`](#agentbricks-init) | Scaffold a new agent project | +| [`auth`](#agentbricks-auth) | Choose and sign in to a Databricks profile | | [`doctor`](#agentbricks-doctor) | Check an existing agent's Agent Bricks onboarding | | [`dev`](#agentbricks-dev) | Run the agent locally with a chat UI | | [`memory`](#agentbricks-memory) | Manage an agent's long-term memory | @@ -74,40 +110,15 @@ These options apply to every command. Pass them before the command name, for exa ## Commands -### `agentbricks login` - -Save a Databricks workspace sign-in as your default, so later commands can omit -p. - -Create the sign-in first with `databricks auth login --profile `. - -``` -agentbricks login [options] -``` - - -_Options_ - -| Option | Values | Default | Required | Description | -| --- | --- | --- | --- | --- | -| `--profile ` (`-p`) | string | - | no | Name of a saved Databricks workspace sign-in to use by default. Run `databricks auth profiles` to list yours. | - -### `agentbricks logout` - -Forget the saved profile selection without deleting its credentials. - -``` -agentbricks logout -``` - ### `agentbricks init` Scaffold a local agent project from an Agent Bricks CLI template. -DIRECTORY is the target path to create (defaults to the template's own name). The directory must not already exist. init records a deployment name that fits the Databricks Apps limit (--name, or the directory name shortened to 17 characters), so `agentbricks deploy` from the project directory needs no NAME. +DIRECTORY is the target path to create (defaults to the template's own name). The directory must not already exist — except a project `init` itself scaffolded, where a rerun changes nothing and points at `agentbricks dev` / `agentbricks deploy` (use `agentbricks auth switch ` to change its workspace). Use `--existing` to migrate a project Agent Bricks did not generate. init records a deployment name that fits the Databricks Apps limit (--name, or the directory name shortened to 17 characters), so `agentbricks deploy` from the project directory needs no NAME. Databricks App names cap at 30 characters and deploy adds the `agent-bricks-` prefix, so a deployment name can be up to 17 lowercase letters, digits, and hyphens. init records `[agent].deployment_name` in `agent.toml` and shows it as `Deploy name`. A directory name that is valid and fits is recorded unchanged. Otherwise init sanitizes it and, if it is still over 17 characters, keeps its first 10 characters plus `-`, the same 6-letter token as the default stores (`caj-langgraph-agent` becomes `caj-langgr-`). `--name` is validated before any file is written; an invalid name fails with a suggestion. JSON results include `deployment_name` and `deployment_name_truncated` (true when init changed the directory name to fit). -Pass --profile (or set a default via `agentbricks login` / -p) to seed a local `.env` so the scaffolded project runs with `agentbricks dev` right away. +The project's `.env` pins a Databricks profile: --profile (or the global -p, or DATABRICKS_CONFIG_PROFILE), else the default set with `databricks auth switch`, else `DEFAULT`, so the project keeps its workspace however the shell is set up later. With env-var auth set (`DATABRICKS_HOST`, `DATABRICKS_TOKEN` or `DATABRICKS_CLIENT_ID`, e.g. CI) and no profile, `init` writes no `.env` and every command keeps using those environment variables (`agentbricks dev` passes them to the local agent through `databricks apps run-local`). In an interactive terminal, `init` then signs in to it, opening `databricks auth login` only if it isn't authenticated yet (which also creates a missing profile, asking for the workspace URL); a failed sign-in keeps the scaffold and lists `agentbricks auth login ` as a next step. Non-interactive and CI runs skip sign-in (and warn if the profile isn't in your Databricks config yet). If nothing chose a profile and `[DEFAULT]` doesn't exist but other profiles do, `init` doesn't sign in, since that would add a machine-wide `[DEFAULT]`, and lists those profiles for `agentbricks auth switch` instead. Change the profile later with `agentbricks auth switch `. With `-o json`, the result includes `env_profile`, `workspace_host`, and `signed_in_user`. The scaffold is preconfigured to call Databricks model serving through the AI Gateway using that profile, so it can talk to a model with no separate endpoint or API key to set up. @@ -141,7 +152,7 @@ _Options_ | --- | --- | --- | --- | --- | | `--framework ` | `langgraph` \| `openai` | `langgraph` | no | Agent framework: langgraph (LangGraph, default) or openai (OpenAI Agents SDK). | | `--server ` | `agentbricks` \| `custom` | `agentbricks` | no | Use the managed invocation server (`agentbricks` is the existing `agent.toml` value) or a minimal custom FastAPI server. | -| `--profile ` | string | - | no | Seed a local .env with this DATABRICKS_CONFIG_PROFILE so `agentbricks dev` works immediately (defaults to the profile from -p / `agentbricks login`). | +| `--profile ` (`-p`) | string | - | no | DATABRICKS_CONFIG_PROFILE to seed the project's .env with, so `agentbricks dev` works right away (default: the global -p, then DATABRICKS_CONFIG_PROFILE, then DEFAULT; with no profile and env-var auth set (DATABRICKS_HOST, DATABRICKS_TOKEN or DATABRICKS_CLIENT_ID), no .env is written). | | `--disable-chat-app` | flag | - | no | Scaffold the API-only backend, without the browser chat app. | | `--enable-chat-app` | flag | - | no | Deprecated: the chat app is included by default; this flag is a no-op. | | `--memory-store ` | string | - | no | Name for the declared memory store (default: a unique name derived from --name or the directory, --memory). Only --server agentbricks declares stores by default. | @@ -149,6 +160,84 @@ _Options_ | `--name ` | string | - | no | Deployment name, recorded in agent.toml as [agent].deployment_name. Deploys as agent-bricks-, so it can be up to 17 lowercase letters, digits, and hyphens. Also names the default stores and tracing experiment (default: derived from the directory, shortened to fit). | | `--existing` | flag | - | no | Prepare a coding-agent migration bundle for an existing LangGraph or OpenAI Agents SDK project (defaults to `.`). Requires `--server agentbricks`. | +### `agentbricks auth` + +Choose, inspect, and sign in to the Databricks profile a project uses. + +A project records its profile in its `.env`, and every command run inside the project uses it (the global `-p` overrides it). Outside a project the DATABRICKS_CONFIG_PROFILE environment variable applies, then the Databricks SDK's default authentication. `login` does not save anything: `switch` is what changes the project's profile (unlike `databricks auth switch`, it writes the project's `.env`, not a global default). + +| Subcommand | Description | +| --- | --- | +| [`auth switch`](#agentbricks-auth-switch) | Set the Databricks profile an Agent Bricks project uses. | +| [`auth describe`](#agentbricks-auth-describe) | Show the profile the CLI would use in a directory, and where it came from. | +| [`auth login`](#agentbricks-auth-login) | Sign in to a Databricks profile, opening a browser login if it isn't authenticated yet. | + +#### `agentbricks auth switch` + +Set the Databricks profile an Agent Bricks project uses. + +Updates DATABRICKS_CONFIG_PROFILE in the project's `.env` in place, leaving every other line alone; a project without a `.env` gets one seeded from `.env.example`. Nothing else about the project changes. In an interactive terminal it then signs in to the profile, opening `databricks auth login` only if it isn't authenticated yet (which also creates a profile missing from your Databricks config); a failed sign-in still keeps the new `.env` profile. Non-interactive and CI runs make no workspace calls, warn if the profile isn't in your Databricks config yet, and list `agentbricks auth login ` as a next step. The directory must be an Agent Bricks project (it contains `.agentbricks/project.toml` or `app.yaml`). With `-o json`, the result includes `signed_in_user`. + +``` +agentbricks auth switch [options] NAME +``` + +_Arguments_ + +| Argument | Required | Description | +| --- | --- | --- | +| `NAME` | yes | - | + +_Options_ + +| Option | Values | Default | Required | Description | +| --- | --- | --- | --- | --- | +| `--source ` | path | `.` | no | Project directory whose .env to update. Defaults to the current directory. | + +With `-o json`: `{"directory", "env_profile", "changed", "workspace_host", "signed_in_user"}`. + +#### `agentbricks auth describe` + +Show the profile the CLI would use in a directory, and where it came from. + +Resolution order: the global `-p`, the project's `.env` (inside an Agent Bricks project), the DATABRICKS_CONFIG_PROFILE environment variable, then the Databricks SDK's default authentication. Read-only; no workspace calls are made. The output shows the profile (or `none (Databricks SDK default)`), its source, the workspace host (for the SDK default: `DATABRICKS_HOST`, else the default chosen with `databricks auth switch`, then `[DEFAULT]`), and the project directory. + +``` +agentbricks auth describe [options] +``` + +_Options_ + +| Option | Values | Default | Required | Description | +| --- | --- | --- | --- | --- | +| `--source ` | path | `.` | no | Project directory. Defaults to the current directory. | + +With `-o json`: `{"profile", "source", "workspace_host", "project"}`. + +#### `agentbricks auth login` + +Sign in to a Databricks profile, opening a browser login if it isn't authenticated yet. + +The profile is NAME, else the global `-p`, else the `.env` profile of the project in the current (or --source) directory; DATABRICKS_CONFIG_PROFILE and the SDK default are not used. With none of those, an interactive terminal hands off to `databricks auth login`, which lets you pick a profile (outside a terminal it fails with a hint). An authenticated profile is only validated. An unauthenticated one is signed in with `databricks auth login --profile NAME` in an interactive terminal (the Databricks CLI is required for that) and validated again; without a terminal, or in CI, it fails with a hint instead. Nothing is saved: use `agentbricks auth switch` to change the project's profile. The profile and where it came from are printed to stderr before anything else happens. + +``` +agentbricks auth login [options] [NAME] +``` + +_Arguments_ + +| Argument | Required | Description | +| --- | --- | --- | +| `NAME` | no | - | + +_Options_ + +| Option | Values | Default | Required | Description | +| --- | --- | --- | --- | --- | +| `--source ` | path | `.` | no | Project directory whose .env picks the profile when NAME and -p are omitted. Defaults to the current directory. | + +With `-o json`: `{"profile", "source", "user", "host"}`. + ### `agentbricks doctor` Check whether an existing agent repository is onboarded to Agent Bricks. DIRECTORY defaults to the @@ -185,7 +274,9 @@ Run your agent locally so you can try it before deploying. Starts the agent on a local server - by default http://localhost:8000 - and, once it answers, prints where to reach it: the chat UI if the project has one, otherwise a sample request against the agent's API. If the agent fails to start, you get an error instead. -Auth uses your Databricks profile (`-p` / `agentbricks login`), and the agent reaches Databricks model serving through the AI Gateway on that profile - so there are no model keys to set up. +Auth uses your Databricks profile, resolved in this order: `-p`, the project's `.env`, then `DATABRICKS_CONFIG_PROFILE`. The agent reaches Databricks model serving through the AI Gateway on that profile - so there are no model keys to set up. + +`dev` first checks that the resolved profile can authenticate and fails fast with the `agentbricks auth login` command to run - without the check, an unauthenticated external-browser profile hangs the agent on a browser login that never completes. (Skipped when `.env` sets `DATABRICKS_TOKEN` or `DATABRICKS_CLIENT_ID`/`DATABRICKS_CLIENT_SECRET`, which authenticate the agent directly.) Under the hood this wraps `databricks apps run-local`: it reads the command + env from `app.yaml` and runs the app the way the Apps runtime would, so local behavior matches a deployment. The environment is built on the first run and reused after; pass `--prepare-environment` to force a rebuild (e.g. after changing dependencies). @@ -1652,7 +1743,7 @@ _Options_ Deploy your agent to Databricks Apps and get back a hosted URL to try it. -Rolls the agent out to Databricks Apps and prints the URL where you (or anyone you share it with) can use it. The deployed agent reaches Databricks model serving through the AI Gateway using the app's own identity - no model keys to configure - and `deploy` also reconciles the stores declared in agent.toml and wires in any tracing. +Rolls the agent out to Databricks Apps and prints the URL where you (or anyone you share it with) can use it. The deployed agent reaches Databricks model serving through the AI Gateway using the app's own identity - no model keys to configure - and `deploy` also reconciles the stores declared in agent.toml and wires in any tracing. `deploy` first checks that the resolved profile can authenticate and fails fast with the `agentbricks auth login` command to run. NAME is recorded in agent.toml on the first deploy, so a later `agentbricks deploy` from the project directory can omit it (passing NAME again updates the recorded name). Deployed apps are named `agent-bricks-`. App names cap at 30 characters, so NAME can be up to 17 (`agentbricks init` records one that fits). A longer name fails before anything is provisioned, with a suggested name that fits. A project with a recorded base name reuses that app when NAME is omitted; you can also pass the full app name to update it. Use the full app name with the `agentbricks deployments` commands. diff --git a/integrations/agentbricks/pyproject.toml b/integrations/agentbricks/pyproject.toml index 618c3f3bf..e8ce332e9 100644 --- a/integrations/agentbricks/pyproject.toml +++ b/integrations/agentbricks/pyproject.toml @@ -24,6 +24,8 @@ dependencies = [ "mlflow-skinny>=3.10.1", "packaging>=23.2", "psycopg[binary]>=3.1", + # The CLI must parse a project `.env` exactly like the generated agent's `load_dotenv`. + "python-dotenv>=0.19.0", "PyYAML>=6.0", # app_provisioner retries the SDK's `requests` connection errors; same range databricks-sdk requires. "requests>=2.28.1,<3", diff --git a/integrations/agentbricks/src/databricks_agentbricks/cli/app.py b/integrations/agentbricks/src/databricks_agentbricks/cli/app.py index e859356a7..54a21d79d 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/cli/app.py +++ b/integrations/agentbricks/src/databricks_agentbricks/cli/app.py @@ -6,13 +6,15 @@ from __future__ import annotations +import functools +import pathlib import sys -from typing import Optional +from typing import Any, Optional import click from databricks_agentbricks import errors -from databricks_agentbricks.cli.auth import load_default_profile, login, logout +from databricks_agentbricks.cli.auth import auth_group from databricks_agentbricks.cli.deploy import deploy, deployments from databricks_agentbricks.cli.dev import dev from databricks_agentbricks.cli.doctor import doctor @@ -25,16 +27,97 @@ from databricks_agentbricks.cli.tracing import tracing from databricks_agentbricks.clients.api_client_provider import ApiClientProvider from databricks_agentbricks.presentation.help import configure_help +from databricks_agentbricks.presentation.profile import ( + announce_workspace, + warn_legacy_default_profile, +) +from databricks_agentbricks.projects.config import find_project_root, is_agentbricks_project +from databricks_agentbricks.projects.env_file import read_env +from databricks_agentbricks.projects.profiles import ( + SOURCE_SDK_DEFAULT, + ProfileInfo, + auth_login_command, + auth_switch_command, + legacy_default_profile, + resolve_profile, +) from databricks_agentkit._api_client import _AgentBricksApiClient class CliContext: - """Shared per-invocation state: selected profile, output mode, lazily-built client.""" + """Shared per-invocation state: selected profile, output mode, lazily-built client. + + The profile is decided once per command, by `select_profile`, before the command body runs. + """ def __init__(self, profile: Optional[str], output: str): - self.profile = profile self.output = output - self.api_client_provider = ApiClientProvider(profile) + self.global_profile = profile + # A `-p/--profile` given after the subcommand; recorded by its option callback. + self.subcommand_profile: Optional[str] = None + # The current directory, when it's an Agent Bricks project; recorded by the root group. + self.cwd_project: Optional[pathlib.Path] = None + # Unset until `select_profile` runs, so reading it earlier fails instead of going stale. + self.profile_info: ProfileInfo + self._provider: Optional[ApiClientProvider] = None + self._workspace_announced = False + self._project_dir: Optional[pathlib.Path] = None + self._project_env: dict[str, str] = {} + + @property + def profile(self) -> Optional[str]: + return self.profile_info.name + + def select_profile(self, source_project: Optional[pathlib.Path] = None) -> None: + """Resolve the command's profile once (see `resolve_profile`); the client provider follows it. + + The project whose `.env` applies is `source_project` (a `--source` project) if given, else + the current directory's project; `-p` (global or after the subcommand) outranks it. + """ + flag = self.subcommand_profile or self.global_profile + project_dir = source_project or self.cwd_project + env_values = None + if project_dir is not None: + directory = project_dir.resolve() + if directory != self._project_dir: + self._project_dir = directory + self._project_env = read_env(directory / ".env") + project_dir, env_values = directory, self._project_env + self.profile_info = resolve_profile(flag, project_dir, env_values) + + @property + def api_client_provider(self) -> ApiClientProvider: + """Built on first use, so it always reflects the profile `select_profile` settled on.""" + if self._provider is None: + self._provider = ApiClientProvider(self.profile, on_first_use=self.announce_workspace) + return self._provider + + @api_client_provider.setter + def api_client_provider(self, provider: Any) -> None: + """Lets tests inject a fake provider.""" + self._provider = provider + + def project_env(self, project_dir: pathlib.Path) -> dict[str, str]: + """`project_dir`'s parsed `.env`, reusing the one `select_profile` already read.""" + directory = project_dir.resolve() + if directory == self._project_dir: + return self._project_env + return read_env(directory / ".env") + + def announce_workspace(self, resolved_host: Optional[str] = None) -> None: + """Name the profile and workspace once, when the command first reaches the workspace.""" + if not self._workspace_announced: + self._workspace_announced = True + announce_workspace( + self.profile, + self.profile_info.source, + output=self.output, + resolved_host=resolved_host, + ) + # Before profiles moved into `.env`, this case used the profile `agentbricks login` saved. + legacy_profile = legacy_default_profile() + if legacy_profile and self.profile_info.source == SOURCE_SDK_DEFAULT: + warn_legacy_default_profile(legacy_profile, self.output) def client(self) -> _AgentBricksApiClient: """Compatibility accessor for commands not yet migrated to ``api_client_provider``.""" @@ -70,8 +153,11 @@ def agentbricks(ctx: click.Context, profile: Optional[str], output: str) -> None 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 `agentbricks login` once to save a default, or pass - --profile / -p (without one, the Databricks SDK's default authentication is used). + authenticates with a Databricks profile: sign in with `agentbricks auth login `, + then run `agentbricks init --profile ` to record it in the project's .env. + Inside a project, its .env profile is used (change it with `agentbricks auth switch`); + --profile / -p overrides it, and outside a project DATABRICKS_CONFIG_PROFILE applies (without + any of these, the Databricks SDK's default authentication is used). Agents built with Agent Bricks combine the platform's capabilities: @@ -91,12 +177,94 @@ def agentbricks(ctx: click.Context, profile: Optional[str], output: str) -> None """ # 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) + ctx.obj = CliContext(profile=profile, output=output) + project = find_project_root(pathlib.Path.cwd()) + if project is not None: + # Every command run inside a project (including its subfolders) uses its `.env` profile, so + # e.g. `endpoint invoke` targets the same workspace that `deploy` just used. + ctx.obj.cwd_project = project + + +_REMOVED_COMMANDS = { + "login": f"Sign in with `{auth_login_command(None)}`, and pick a project's profile with " + f"`{auth_switch_command(None)}`.", + "logout": f"Profiles are now chosen per project in its .env (`{auth_switch_command(None)}`); " + "there is no saved default to clear.", +} + + +def _removed_command(name: str, hint: str) -> click.Command: + @click.command( + name, + hidden=True, + context_settings={"ignore_unknown_options": True, "allow_extra_args": True}, + ) + def removed() -> None: + raise errors.AgentCliError(f"`agentbricks {name}` was removed.", hint=hint) + + return removed + + +def _apply_subcommand_profile( + ctx: click.Context, _param: click.Parameter, value: Optional[str] +) -> None: + obj = ctx.find_object(CliContext) + if value and obj is not None: + obj.subcommand_profile = value + + +def _accept_profile_anywhere(command: click.Command, *, source_profile: bool = True) -> None: + """Let `-p/--profile` follow the subcommand too (`agentbricks dev -p X`), not only precede it. + + Commands that declare their own `--profile` (`init`) keep it. + """ + if isinstance(command, click.Group): + for subcommand in command.commands.values(): + _accept_profile_anywhere(subcommand, source_profile=source_profile) + return + _select_profile_before_running( + command, + use_source=source_profile and any(param.name == "source" for param in command.params), + ) + if any("--profile" in param.opts for param in command.params): + return + command.params.append( + click.Option( + ["--profile", "-p"], + default=None, + expose_value=False, + callback=_apply_subcommand_profile, + help="Databricks config profile to authenticate with (same as the global -p).", + ) + ) + + +def _select_profile_before_running(command: click.Command, *, use_source: bool) -> None: + """Decide the profile once, just before `command` runs, from `-p`, `--source`'s project (so e.g. + `tracing bind --source proj` reads the workspace `deploy --source proj` writes to), the + current directory's project, then the environment.""" + callback = command.callback + if callback is None: + return + + @functools.wraps(callback) + def with_selected_profile(*args, **kwargs): + obj = click.get_current_context().find_object(CliContext) + if obj is not None: + source = kwargs.get("source") if use_source else None + source_project = ( + pathlib.Path(source) + if source and is_agentbricks_project(pathlib.Path(source)) + else None + ) + obj.select_profile(source_project) + return callback(*args, **kwargs) + + command.callback = with_selected_profile -agentbricks.add_command(login) -agentbricks.add_command(logout) agentbricks.add_command(init) +agentbricks.add_command(auth_group) agentbricks.add_command(doctor) agentbricks.add_command(dev) agentbricks.add_command(memory) @@ -107,6 +275,11 @@ def agentbricks(ctx: click.Context, profile: Optional[str], output: str) -> None agentbricks.add_command(endpoint) agentbricks.add_command(tools) agentbricks.add_command(experimental) +for _name, _hint in _REMOVED_COMMANDS.items(): + agentbricks.add_command(_removed_command(_name, _hint)) +# `auth` commands resolve their own `--source` directory and must not move the session to it. +for _command in agentbricks.commands.values(): + _accept_profile_anywhere(_command, source_profile=_command is not auth_group) configure_help(agentbricks) diff --git a/integrations/agentbricks/src/databricks_agentbricks/cli/auth.py b/integrations/agentbricks/src/databricks_agentbricks/cli/auth.py index 18bbfbe39..1a58a4b2d 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/cli/auth.py +++ b/integrations/agentbricks/src/databricks_agentbricks/cli/auth.py @@ -1,146 +1,191 @@ -"""`agentbricks 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). -""" +"""`agentbricks auth` — choose, inspect, and sign in to the Databricks profile a project uses.""" from __future__ import annotations -import json import os import pathlib -import subprocess import sys from typing import Optional import click +from databricks_agentbricks.clients.databricks_cli import run_databricks_login from databricks_agentbricks.errors import AgentCliError -from databricks_agentbricks.presentation import render -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" +from databricks_agentbricks.presentation.profile import ( + announce_login, + show_databricks_login_result, + show_login_result, + show_profile_description, + show_profile_switched, + warn_sign_in_failed, + warn_unknown_profile, +) +from databricks_agentbricks.projects.config import find_project_root +from databricks_agentbricks.projects.env_file import read_env, update_env_profile +from databricks_agentbricks.projects.profiles import ( + SOURCE_FLAG, + ProfileInfo, + auth_login_command, + auth_switch_command, + login_target, + profile_host, + resolve_profile, + workspace_for_profile, +) +from databricks_agentbricks.services.workspace_auth import authenticate_profile +_SOURCE_HELP = "Project directory. Defaults to the current directory." -def load_default_profile() -> Optional[str]: - """The profile saved by `agentbricks 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 _is_interactive() -> bool: + return sys.stdin.isatty() -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 can_prompt_login() -> bool: + """Whether someone is at a terminal to complete a browser login (and this isn't CI).""" + return not os.getenv("CI") and _is_interactive() -def _validate_profile(profile: str) -> tuple[_AgentBricksApiClient, str]: - client = _AgentBricksApiClient(profile) - return client, client.current_user +def sign_in_chosen_profile(obj, profile: str, source: str) -> Optional[str]: + """Sign in to `profile` right after it's chosen, so `agentbricks dev` works right away. -def _is_interactive() -> bool: - return sys.stdin.isatty() + Interactive terminals only: a scripted or CI run never opens a browser or waits on the + workspace, and only warns when the profile isn't configured yet (an interactive sign-in + creates it). A failed sign-in is reported, not raised, since the profile choice still stands. + Returns the signed-in user, or None when it was skipped or failed. + """ + if not can_prompt_login(): + warn_unknown_profile(profile, obj) + return None + announce_login(obj, profile, source) + try: + _, user = authenticate_profile(profile, source, interactive=True) + except AgentCliError as exc: + warn_sign_in_failed(profile, exc, obj) + return None + return user -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 `agentbricks login --profile {profile}`.", - ) from exc - if result.returncode != 0: +def _require_project(source: str) -> pathlib.Path: + """The project root containing `source`, which may be any folder inside the project.""" + directory = pathlib.Path(source).resolve() + project = find_project_root(directory) + if project is None: raise AgentCliError( - f"`databricks auth login --profile {profile}` failed (exit {result.returncode})." + f"'{directory}' isn't an Agent Bricks project.", + hint="Run this from the project directory, pass --source , or create a project " + "with `agentbricks init`.", ) + return project -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 +def _resolve_for(obj, directory: pathlib.Path, in_project: bool) -> ProfileInfo: + """The profile the CLI would use for `directory`: `-p`, its project `.env`, the env var, default.""" + if obj.profile_info.source == SOURCE_FLAG: + return obj.profile_info + return resolve_profile(None, directory if in_project else None) - 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.group("auth") +def auth_group() -> None: + """Choose, inspect, and sign in to the Databricks profile a project uses. + A project records its profile in its `.env`, and every command run inside the project uses it + (the global `-p` overrides it). Outside a project the DATABRICKS_CONFIG_PROFILE environment + variable applies, then the Databricks SDK's default authentication. `login` does not save + anything: `switch` is what changes the project's profile (unlike `databricks auth switch`, it + writes the project's `.env`, not a global default). + """ -@click.command() + +@auth_group.command("switch") +@click.argument("name") @click.option( - "--profile", - "-p", - default=None, - help="Name of a saved Databricks workspace sign-in to use by default. Run " - "`databricks auth profiles` to list yours.", + "--source", + default=".", + type=click.Path(file_okay=False), + help="Project directory whose .env to update. Defaults to the current directory.", ) @click.pass_obj -def login(obj, profile) -> None: - """Save a Databricks workspace sign-in as your default, so later commands can omit -p. +def auth_switch(obj, name: str, source: str) -> None: + """Set the Databricks profile an Agent Bricks project uses. - Create the sign-in first with `databricks auth login --profile `. + Updates DATABRICKS_CONFIG_PROFILE in the project's `.env` in place, leaving every other line + alone; a project without a `.env` gets one seeded from `.env.example`. Nothing else about the + project changes. In an interactive terminal it then signs in to the profile (see `auth + login`); non-interactive and CI runs make no workspace calls. """ - profile = profile or obj.profile - if not profile: - raise AgentCliError( - "No profile to save.", - hint="Pass one to remember, e.g. `agentbricks login --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=[ - ("agentbricks init my-agent", "Scaffold a new agent project"), - ], - ) + directory = _require_project(source) + changed = update_env_profile(directory, name) + user = sign_in_chosen_profile(obj, name, "argument") + show_profile_switched(directory, name, changed, profile_host(name), user, obj.output) + + +@auth_group.command("describe") +@click.option( + "--source", + default=".", + type=click.Path(file_okay=False), + help=_SOURCE_HELP, +) +@click.pass_obj +def auth_describe(obj, source: str) -> None: + """Show the profile the CLI would use in a directory, and where it came from. + + Resolution order: the global `-p`, the project's `.env` (inside an Agent Bricks project), the + DATABRICKS_CONFIG_PROFILE environment variable, then the Databricks SDK's default + authentication. Read-only; no workspace calls are made. + """ + directory = pathlib.Path(source).resolve() + project_root = find_project_root(directory) + info = _resolve_for(obj, project_root or directory, project_root is not None) + host, workspace = workspace_for_profile(info.name) + project = str(project_root) if project_root is not None else None + show_profile_description(info.name, info.source, host, workspace, project, obj.output) -@click.command() +@auth_group.command("login") +@click.argument("name", required=False) +@click.option( + "--source", + default=".", + type=click.Path(file_okay=False), + help="Project directory whose .env picks the profile when NAME and -p are omitted. Defaults " + "to the current directory.", +) @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}) +def auth_login(obj, name: Optional[str], source: str) -> None: + """Sign in to a Databricks profile, opening a browser login if it isn't authenticated yet. + + The profile is NAME, else the global `-p`, else the `.env` profile of the project in the + current (or --source) directory; DATABRICKS_CONFIG_PROFILE and the SDK default are not used. + With none of those, an interactive terminal hands off to `databricks auth login`, which lets + you pick a profile. + An authenticated profile is only validated. An unauthenticated one is signed in with + `databricks auth login --profile NAME` in an interactive terminal (the Databricks CLI is + required for that) and validated again; without a terminal, or in CI, it fails with a hint + instead. Nothing is saved: use `agentbricks auth switch` to change the project's profile. + """ + project_root = find_project_root(pathlib.Path(source).resolve()) + env_profile = ( + read_env(project_root / ".env").get("DATABRICKS_CONFIG_PROFILE") + if project_root is not None + else None + ) + target = login_target(name, obj.profile_info, env_profile) + if target is None: + if not can_prompt_login(): + raise AgentCliError( + "No profile to log in to.", + hint=f"Pass one (`{auth_login_command(None)}`), or set the project's profile " + f"with `{auth_switch_command(None)}`.", + ) + # `databricks auth login` picks the profile; it doesn't say which, so there's nothing to + # validate or report beyond its own output. + run_databricks_login(None) + show_databricks_login_result(obj.output) return - render.success("Logged out" if existed else "No saved login to clear") + chosen, origin = target + announce_login(obj, chosen, origin) + client, user = authenticate_profile(chosen, origin, interactive=can_prompt_login()) + show_login_result(chosen, origin, user, client.host, env_profile, obj.output) diff --git a/integrations/agentbricks/src/databricks_agentbricks/cli/deploy.py b/integrations/agentbricks/src/databricks_agentbricks/cli/deploy.py index 028dd69ad..5eb323602 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/cli/deploy.py +++ b/integrations/agentbricks/src/databricks_agentbricks/cli/deploy.py @@ -81,6 +81,7 @@ ) from databricks_agentbricks.services.deployment.tool_access_provisioner import ToolAccessProvisioner from databricks_agentbricks.services.invoke.auth import WorkspaceOAuthAuthenticator +from databricks_agentbricks.services.workspace_auth import preflight_auth # `deployments wait` returns before a coding agent's usual 120 s command timeout, and exits with # EX_TEMPFAIL from sysexits.h when the deploy is still running so the caller knows to wait again. @@ -296,6 +297,14 @@ def deploy( """ if allowed_origin and clear_allowed_origins: raise click.UsageError("Use --allowed-origin or --clear-allowed-origins, not both.") + # Validate auth before any deployment work: an unauthenticated profile would otherwise fail + # deep into the deploy (or hang on a browser flow), far from anything actionable. + try: + validated = preflight_auth(obj.profile, obj.profile_info.source) + except Exception: + obj.announce_workspace() + raise + obj.api_client_provider.use(validated) request = DeployRequest( name=name, source=source, @@ -337,6 +346,7 @@ def deployments() -> None: @click.pass_obj def deployments_list(obj) -> None: """List Agent Bricks deployments (apps named `agent-bricks-*`).""" + obj.announce_workspace() present_deployments_list(build_deploy_service(obj).list_deployments(), output=obj.output) @@ -345,6 +355,7 @@ def deployments_list(obj) -> None: @click.pass_obj def deployments_get(obj, name) -> None: """Get an agent deployment's details.""" + obj.announce_workspace() present_deployment_detail(build_deploy_service(obj).get(name), name, output=obj.output) @@ -390,6 +401,7 @@ def deployments_wait(ctx: click.Context, name, timeout_s, source) -> None: @click.pass_obj def deployments_logs(obj, name) -> None: """Stream a deployment's logs.""" + obj.announce_workspace() build_deploy_service(obj).logs(name) @@ -398,6 +410,7 @@ def deployments_logs(obj, name) -> None: @click.pass_obj def deployments_start(obj, name) -> None: """Start a deployment.""" + obj.announce_workspace() build_deploy_service(obj).start(name) present_started(name, output=obj.output) @@ -412,6 +425,7 @@ def deployments_stop(obj, name, yes) -> None: click.confirm( f"Stop deployment '{name}'? This cannot be undone.", default=False, abort=True ) + obj.announce_workspace() build_deploy_service(obj).stop(name) present_stopped(name, output=obj.output) @@ -429,7 +443,15 @@ def deployments_stop(obj, name, yes) -> None: def deployments_delete(obj, name, delete_source, yes) -> None: """Delete a deployment, and its Runtime Store when managed provisioning is enabled.""" service = build_deploy_service(obj) - target = service.delete_target(name) + # Managed cleanup reaches the shared API provider, whose notice uses the resolved host. + # A CLI-only delete has no SDK client to consult, so use the offline lookup there. + if not service.deletes_runtime_store_data(): + obj.announce_workspace() + try: + target = service.delete_target(name) + except Exception: + obj.announce_workspace() + raise parts = [f"deployment '{name}'"] if not yes and service.deletes_runtime_store_data(): parts.append("its Runtime Store data") @@ -437,7 +459,11 @@ def deployments_delete(obj, name, delete_source, yes) -> None: parts.append(f"its source folder {target.source_path}") described = parts[0] if len(parts) == 1 else ", ".join(parts[:-1]) + " and " + parts[-1] render.confirm_destroy(described, assume_yes=yes) - result = service.delete(target, delete_source=delete_source) + try: + result = service.delete(target, delete_source=delete_source) + except Exception: + obj.announce_workspace() + raise present_deleted(result, output=obj.output) if not result.confirmed or result.source_folder_status == "failed": raise click.exceptions.Exit(1) diff --git a/integrations/agentbricks/src/databricks_agentbricks/cli/dev.py b/integrations/agentbricks/src/databricks_agentbricks/cli/dev.py index e3fa08748..5993d4f33 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/cli/dev.py +++ b/integrations/agentbricks/src/databricks_agentbricks/cli/dev.py @@ -9,6 +9,7 @@ from __future__ import annotations +import pathlib from typing import Optional import click @@ -20,6 +21,7 @@ from databricks_agentbricks.errors import AgentCliError from databricks_agentbricks.presentation import render from databricks_agentbricks.presentation.dev import ( + announce_agent_profile, announce_bound_resources, announce_starting, local_trace_label, @@ -33,6 +35,21 @@ DevService, _dev_entry_point, # noqa: F401 - compatibility re-export for existing callers ) +from databricks_agentbricks.services.workspace_auth import preflight_auth + +# `.env` credential sets that authenticate the agent on their own, without a profile. +_ENV_CREDENTIALS = ( + ("DATABRICKS_TOKEN",), + ("DATABRICKS_CLIENT_ID", "DATABRICKS_CLIENT_SECRET"), +) + + +def _env_credential(env_file: dict[str, str]) -> Optional[str]: + """The explicit credential `.env` sets (e.g. ``DATABRICKS_TOKEN``), or None.""" + for names in _ENV_CREDENTIALS: + if all(env_file.get(name) for name in names): + return "/".join(names) + return None def build_dev_service(obj) -> DevService: @@ -73,8 +90,9 @@ def dev( prints where to reach it: the chat UI if the project has one, otherwise a sample request against the agent's API. If the agent fails to start, you get an error instead. - Auth uses your Databricks profile (`-p` / `agentbricks login`), and the agent reaches Databricks model - serving through the AI Gateway on that profile — so there are no model keys to set up. + Auth uses your Databricks profile, resolved in this order: `-p`, the project's `.env`, + then `DATABRICKS_CONFIG_PROFILE`. The agent reaches Databricks model serving through the AI + Gateway on that profile — so there are no model keys to set up. Under the hood this wraps `databricks apps run-local`: it reads the command + env from `app.yaml` and runs the app the way the Apps runtime would, so local behavior matches a @@ -90,15 +108,51 @@ def dev( `agentbricks deploy`, not here. So there's nothing to provision and no service-principal grant to make; that all happens at `agentbricks deploy` time. """ + # The project's `.env` profile is already selected (see `CliContext.select_profile`); the + # locally running agent reads the same file, which dev also consults for explicit credentials. + env_file = obj.project_env(pathlib.Path(source)) + env_credential = _env_credential(env_file) + # Explicit `.env` credentials authenticate the agent directly, so the profile needs neither + # validation nor a place in the dev manifest (a bare DATABRICKS_HOST is not a credential). + uses_env_credentials = bool(obj.profile and env_credential) + local_env: dict[str, str] = {} + if obj.profile and not uses_env_credentials: + # Landing the profile in the dev manifest puts it in the agent's process env, where it + # beats `.env` because the template loads dotenv with override=False. + local_env["DATABRICKS_CONFIG_PROFILE"] = obj.profile + service = build_dev_service(obj) + service.check_project(source) + # Auth is checked before anything is started: an unauthenticated external-browser profile + # would otherwise hang the local agent on a browser flow with no error. + if not env_credential: + try: + validated = preflight_auth(obj.profile, obj.profile_info.source) + except Exception: + obj.announce_workspace() + raise + obj.announce_workspace(validated.host) + elif not uses_env_credentials: + obj.announce_workspace() with service.prepare( - DevRequest(source=source, prepare_environment=prepare_environment, app_port=app_port) + DevRequest( + source=source, + prepare_environment=prepare_environment, + app_port=app_port, + local_env=local_env, + ) ) as plan: if plan.preview.tracing_warning is not None: render.diagnostic( "warning", plan.preview.tracing_warning, help=plan.preview.tracing_help ) announce_bound_resources(plan.preview) + announce_agent_profile( + obj.profile, + obj.profile_info.source, + env_file.get("DATABRICKS_CONFIG_PROFILE"), + env_credential=env_credential if uses_env_credentials else None, + ) announce_starting() # The green panel waits until the agent answers on its port, so a failed start (bad auth, a # build error, a crash on import) never looks like a successful one. diff --git a/integrations/agentbricks/src/databricks_agentbricks/cli/endpoint.py b/integrations/agentbricks/src/databricks_agentbricks/cli/endpoint.py index 81e3a6e9b..10dd2168e 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/cli/endpoint.py +++ b/integrations/agentbricks/src/databricks_agentbricks/cli/endpoint.py @@ -88,6 +88,8 @@ def invoke( auth, ) -> None: """Send one HTTP request to a Databricks App or arbitrary URL.""" + if app: + obj.announce_workspace() service = build_endpoint_service(obj) request = InvokeRequest( app=app, @@ -194,6 +196,8 @@ def loadtest( responses_file=responses_file, warmup=warmup, ) + if app: + obj.announce_workspace() result = build_endpoint_service(obj).loadtest(plan) render_loadtest_summary(result.summary, report_dir=report_dir, output=obj.output) if obj.output == "json": diff --git a/integrations/agentbricks/src/databricks_agentbricks/cli/init.py b/integrations/agentbricks/src/databricks_agentbricks/cli/init.py index 389095bab..0e9b5c085 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/cli/init.py +++ b/integrations/agentbricks/src/databricks_agentbricks/cli/init.py @@ -11,6 +11,7 @@ from __future__ import annotations import json +import os import pathlib import secrets import shlex @@ -25,14 +26,32 @@ import click +from databricks_agentbricks.cli.auth import sign_in_chosen_profile from databricks_agentbricks.errors import AgentCliError from databricks_agentbricks.presentation import render +from databricks_agentbricks.presentation.profile import note_choose_profile from databricks_agentbricks.projects.agent_project import ( AgentProject, default_experiment_name, default_store_name, ) -from databricks_agentbricks.projects.config import write_project_metadata +from databricks_agentbricks.projects.config import ( + _load_persisted_metadata, + write_project_metadata, +) +from databricks_agentbricks.projects.env_file import write_env +from databricks_agentbricks.projects.profiles import ( + SOURCE_CONFIG_DEFAULT, + SOURCE_ENV_VAR, + SOURCE_FLAG, + auth_login_command, + auth_switch_command, + configured_profiles, + env_auth_configured, + profile_exists, + profile_host, + sdk_default_profile, +) from databricks_agentbricks.projects.types import ( AgentFramework, AgentServer, @@ -129,29 +148,51 @@ def _bundled_template_ref() -> str: return "bundled" -def _write_env(dest: pathlib.Path, profile: str) -> bool: - """Seed a local `.env` from `.env.example` with DATABRICKS_CONFIG_PROFILE=. +def _is_generated_project(dest: pathlib.Path) -> bool: + """Whether `init` (this CLI family) scaffolded `dest`, by its persisted metadata marker.""" + try: + return _load_persisted_metadata(dest) is not None + except AgentCliError: + return False + + +def _report_already_initialized(obj, dest: pathlib.Path) -> None: + """Re-run of `init` on a project this CLI scaffolded: nothing is written. - Returns True if a `.env` was written. Skips if `.env` already exists (never clobbers). The - template reads DATABRICKS_CONFIG_PROFILE for local model auth, so this makes the scaffolded - project runnable with `agentbricks dev` without a manual `cp .env.example .env` step. + Re-scaffolding would clobber user code, so the project's identity and `.env` stay as they + are; `--existing` migration is for projects Agent Bricks did NOT generate. """ - env_path = dest / ".env" - if env_path.exists(): - return False - example = dest / ".env.example" - base = example.read_text() if example.exists() else "" - lines, replaced = [], False - for line in base.splitlines(): - if line.startswith("DATABRICKS_CONFIG_PROFILE="): - lines.append(f"DATABRICKS_CONFIG_PROFILE={profile}") - replaced = True - else: - lines.append(line) - if not replaced: - lines.insert(0, f"DATABRICKS_CONFIG_PROFILE={profile}") - env_path.write_text("\n".join(lines) + "\n") - return True + if obj.output == "json": + render.emit_json({"already_initialized": True, "directory": str(dest)}) + return + render.console().print( + f"{dest} is already an Agent Bricks project; nothing was changed. Run `agentbricks dev` / " + f"`agentbricks deploy`, or `{auth_switch_command(None)} --source {dest}` to change " + "its workspace." + ) + + +def _project_profile(obj, flag_profile: Optional[str]) -> Optional[tuple[str, str]]: + """(profile, source) to pin in a new project's `.env`: `--profile`, the global `-p`, the + DATABRICKS_CONFIG_PROFILE env var, else None when env-var auth is configured (DATABRICKS_HOST, + DATABRICKS_TOKEN, or DATABRICKS_CLIENT_ID), else the default chosen with `databricks auth + switch`, else `DEFAULT` (the order the SDK itself falls back in). + + Never the cwd project's `.env` — `init` may run inside another project, and the new project + must not inherit its workspace. A pinned profile keeps the project on its workspace however + the shell is set up later; None under env-var auth because it (CI, containers) already names + the workspace and a `.env` profile would override it. + """ + if flag_profile: + return flag_profile, SOURCE_FLAG + if obj.profile_info.source == SOURCE_FLAG and obj.profile: + return obj.profile, SOURCE_FLAG + env_profile = os.environ.get("DATABRICKS_CONFIG_PROFILE") + if env_profile: + return env_profile, SOURCE_ENV_VAR + if env_auth_configured(): + return None + return sdk_default_profile() def _scaffold_token() -> str: @@ -383,9 +424,11 @@ def _prepare_migration( ) @click.option( "--profile", + "-p", default=None, - help="Seed a local .env with this DATABRICKS_CONFIG_PROFILE so `agentbricks dev` works " - "immediately (defaults to the profile from -p / `agentbricks login`).", + help="DATABRICKS_CONFIG_PROFILE to seed the project's .env with, so `agentbricks dev` works " + "right away (default: the global -p, then DATABRICKS_CONFIG_PROFILE, then DEFAULT; with no " + "profile and env-var auth set (DATABRICKS_HOST, DATABRICKS_TOKEN or DATABRICKS_CLIENT_ID), no .env is written).", ) @click.option( "--disable-chat-app", @@ -438,12 +481,23 @@ def init( """Scaffold a local agent project from an Agent Bricks template. DIRECTORY is the target path to create (defaults to the template's own name). The - directory must not already exist unless --existing is supplied. init records a deployment name - that fits the Databricks Apps limit (--name, or the directory name shortened to 17 characters), - so `agentbricks deploy` from the project directory needs no NAME. - - Pass --profile (or set a default via `agentbricks login` / -p) to seed a local `.env` so the - scaffolded project runs with `agentbricks dev` right away. + directory must not already exist — except a project `init` itself scaffolded, where a rerun + changes nothing and points at `agentbricks dev` / `agentbricks deploy` (use + `agentbricks auth switch ` to change its workspace). Use --existing to migrate a + project Agent Bricks did not generate. Once scaffolded, deploy it with + `agentbricks deploy --source `. init records a deployment name that fits the + Databricks Apps limit (--name, or the directory name shortened to 17 characters), so + `agentbricks deploy` from the project directory needs no NAME. + + The project's `.env` pins a Databricks profile: --profile (or the global -p, or + DATABRICKS_CONFIG_PROFILE), else the default set with `databricks auth switch`, else + `DEFAULT`. With env-var auth set (DATABRICKS_HOST, DATABRICKS_TOKEN or DATABRICKS_CLIENT_ID, e.g. CI) and no + profile, no `.env` is written and commands keep using those environment variables. In an interactive terminal, + init then signs in to it, opening `databricks auth login` only if it isn't authenticated yet + (which also creates a missing profile); non-interactive and CI runs skip sign-in. If nothing + chose a profile and `[DEFAULT]` doesn't exist but other profiles do, init doesn't sign in + (that would add a machine-wide `[DEFAULT]`) and lists them instead. Change it later with + `agentbricks auth switch `. The scaffold is preconfigured to call Databricks model serving through the AI Gateway using that profile, so it can talk to a model with no separate endpoint or API key to set up. @@ -487,11 +541,14 @@ def init( dest = pathlib.Path(directory) if directory else pathlib.Path(template_name) if dest.exists(): - raise AgentCliError( - f"Destination '{dest}' already exists.", - hint=f"Use --existing to prepare a migration from an existing " - f"{_FRAMEWORK_LABEL_PHRASE} project, or choose a new directory to scaffold.", - ) + if not _is_generated_project(dest): + raise AgentCliError( + f"Destination '{dest}' already exists.", + hint=f"Use --existing to prepare a migration from an existing " + f"{_FRAMEWORK_LABEL_PHRASE} project, or choose a new directory to scaffold.", + ) + _report_already_initialized(obj, dest) + return overlay_names = (template.chat_app,) if chat_app_enabled else () try: @@ -528,13 +585,37 @@ def init( # Record a name that fits the Apps limit so a bare `agentbricks deploy` works. project.set_deployment_name(deployment_name) project.write() - env_profile = profile or obj.profile - wrote_env = _write_env(dest, env_profile) if env_profile else False + project_profile = _project_profile(obj, profile) + env_profile = project_profile[0] if project_profile else None + if env_profile: + write_env(dest, env_profile) workflow_skills = install_workflow_skill(dest) except Exception: shutil.rmtree(dest, ignore_errors=True) raise + # Nothing named a profile and `[DEFAULT]` doesn't exist, but others do: signing in would add a + # new machine-wide `[DEFAULT]` the user likely didn't want, so point at their profiles instead. + other_profiles = ( + configured_profiles() + if ( + project_profile + and project_profile[1] == SOURCE_CONFIG_DEFAULT + and not profile_exists(project_profile[0]) + ) + else [] + ) + if other_profiles: + note_choose_profile(other_profiles, obj) + + # A failed sign-in doesn't undo the scaffold; the summary says how to retry. + signed_in_user = ( + sign_in_chosen_profile(obj, *project_profile) + if project_profile and not other_profiles + else None + ) + workspace_host = profile_host(env_profile) if env_profile else os.environ.get("DATABRICKS_HOST") + if obj.output == "json": render.emit_json( { @@ -544,7 +625,9 @@ def init( "directory": str(dest), "server": selected_server.value, "chat_app_enabled": chat_app_enabled, - "env_profile": env_profile if wrote_env else None, + "env_profile": env_profile, + "workspace_host": workspace_host, + "signed_in_user": signed_in_user, "memory_store": memory_store, "session_store": session_store, "experiment_name": experiment_name, @@ -574,15 +657,20 @@ def init( fields["Traces experiment"] = experiment_name fields["Deploy name"] = _deployment_name_note(dest.name, deployment_name, truncated) steps: list[str | tuple[str, str]] = [(f"cd {dest}", "Enter the project directory")] - if wrote_env: + if env_profile: fields["Profile (.env)"] = env_profile + fields["Workspace"] = workspace_host or "not configured yet" else: - # No profile resolved, so no .env was seeded — call out the auth step explicitly rather - # than burying it, since running locally fails without a Databricks profile. - steps += [ - ("cp .env.example .env", "Create your local env file"), - "Set DATABRICKS_CONFIG_PROFILE in .env (or re-run `agentbricks init --profile `)", - ] + fields["Auth"] = "environment variables; no .env written" + fields["Workspace"] = workspace_host or "not configured yet" + if signed_in_user: + fields["Signed in as"] = signed_in_user + elif other_profiles: + steps.append((auth_switch_command(None), f"Pick one of: {', '.join(other_profiles)}")) + elif env_profile: + steps.append((auth_login_command(env_profile), "Sign in before running the agent")) + if not other_profiles: + steps.append((auth_switch_command(None), "Use a different profile (optional)")) steps.append(("agentbricks dev", "Optional: develop and test the agent locally")) if chat_app_enabled: steps.append("Open http://localhost:8000 to chat with it") diff --git a/integrations/agentbricks/src/databricks_agentbricks/cli/tracing.py b/integrations/agentbricks/src/databricks_agentbricks/cli/tracing.py index a5df37e6f..ba97d4ecd 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/cli/tracing.py +++ b/integrations/agentbricks/src/databricks_agentbricks/cli/tracing.py @@ -75,6 +75,7 @@ def _build_mlflow_trace_client(obj) -> MLflowTraceClient: return MLflowTraceClient( obj.profile, mlflow_loader=_mlflow, + on_workspace=obj.announce_workspace, ) diff --git a/integrations/agentbricks/src/databricks_agentbricks/clients/api_client_provider.py b/integrations/agentbricks/src/databricks_agentbricks/clients/api_client_provider.py index dd511062a..c510497aa 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/clients/api_client_provider.py +++ b/integrations/agentbricks/src/databricks_agentbricks/clients/api_client_provider.py @@ -11,19 +11,37 @@ from __future__ import annotations from threading import Lock -from typing import Optional +from typing import Callable, Optional from databricks_agentkit._api_client import _AgentBricksApiClient +def validate_profile(profile: Optional[str]) -> tuple[_AgentBricksApiClient, str]: + """Build a workspace client and confirm its credentials identify a user.""" + client = _AgentBricksApiClient(profile) + return client, client.current_user + + class ApiClientProvider: """Own exactly one lazily created API client for a CLI invocation.""" - def __init__(self, profile: Optional[str]) -> None: + def __init__( + self, + profile: Optional[str], + on_first_use: Optional[Callable[[Optional[str]], None]] = None, + ) -> None: self._profile = profile + self._on_first_use = on_first_use self._client: Optional[_AgentBricksApiClient] = None self._lock = Lock() + def use(self, client: _AgentBricksApiClient) -> None: + """Adopt a client the command already validated, so first use doesn't build another.""" + with self._lock: + self._client = client + if self._on_first_use is not None: + self._on_first_use(client.host) + def get(self) -> _AgentBricksApiClient: """Return the shared client, constructing it on first use.""" client = self._client @@ -32,5 +50,12 @@ def get(self) -> _AgentBricksApiClient: with self._lock: if self._client is None: - self._client = _AgentBricksApiClient(self._profile) + try: + self._client = _AgentBricksApiClient(self._profile) + except Exception: + if self._on_first_use is not None: + self._on_first_use(None) + raise + if self._on_first_use is not None: + self._on_first_use(self._client.host) return self._client diff --git a/integrations/agentbricks/src/databricks_agentbricks/clients/apps_client.py b/integrations/agentbricks/src/databricks_agentbricks/clients/apps_client.py index e754b5f0a..de09d1d96 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/clients/apps_client.py +++ b/integrations/agentbricks/src/databricks_agentbricks/clients/apps_client.py @@ -24,6 +24,7 @@ from databricks_agentbricks.clients.legacy_runtime_store import LakebaseBackend from databricks_agentbricks.clients.local_ports import accepts_connections from databricks_agentbricks.errors import AgentCliError +from databricks_agentbricks.projects.profiles import auth_login_command DatabricksRunner = Callable[..., subprocess.CompletedProcess] @@ -257,7 +258,7 @@ def _raise_unless_missing(self, name: str, result: subprocess.CompletedProcess) raise self._existence_check_failed(name, detail) def _existence_check_failed(self, name: str, detail: str) -> AgentCliError: - login = f"agentbricks login --profile {self._profile or ''}" + login = auth_login_command(self._profile) return AgentCliError( f"Could not check whether deployment '{name}' exists.", hint=f"{detail or 'The `databricks apps get` command failed without a message.'}\n" diff --git a/integrations/agentbricks/src/databricks_agentbricks/clients/databricks_cli.py b/integrations/agentbricks/src/databricks_agentbricks/clients/databricks_cli.py index e9a138526..92f9df97b 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/clients/databricks_cli.py +++ b/integrations/agentbricks/src/databricks_agentbricks/clients/databricks_cli.py @@ -1,7 +1,7 @@ """Thin wrapper around the ``databricks`` CLI. Agent Bricks shells out to the ``databricks`` CLI for the operations it doesn't call over the REST client -(``apps`` lifecycle, ``sync``, ``postgres`` endpoint/credential lookups). This is the single place +(``apps`` lifecycle, ``sync``, ``postgres`` endpoint/credential lookups, ``auth login``). This is the single place that builds the argv, threads the ``--profile`` through, and normalizes a non-zero exit into an ``AgentCliError``. Kept separate from any one caller so ``deploy``, ``dev``, and ``app_resources`` share it without importing each other. @@ -17,6 +17,7 @@ from typing import Optional from databricks_agentbricks.errors import AgentCliError +from databricks_agentbricks.projects.profiles import auth_login_command # Lines of relayed output a `tee` run keeps for callers that inspect a failure. _TEE_TAIL_LINES = 200 @@ -98,6 +99,22 @@ def _run_tee( return subprocess.CompletedProcess(cmd, proc.returncode, stdout="".join(tail), stderr="") +def run_databricks_login(profile: Optional[str]) -> None: + """Run `databricks auth login`; with no profile it shows its own interactive profile picker.""" + command = ["databricks", "auth", "login", *(["--profile", profile] if profile else [])] + retry = auth_login_command(profile) if profile else "agentbricks auth login" + 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 `{retry}`.", + ) from exc + if result.returncode != 0: + raise AgentCliError(f"`{' '.join(command)}` failed (exit {result.returncode}).") + + def _run_until_ready( cmd: list[str], *, diff --git a/integrations/agentbricks/src/databricks_agentbricks/clients/mlflow_trace_client.py b/integrations/agentbricks/src/databricks_agentbricks/clients/mlflow_trace_client.py index 3aa44e8ee..cc684aa2b 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/clients/mlflow_trace_client.py +++ b/integrations/agentbricks/src/databricks_agentbricks/clients/mlflow_trace_client.py @@ -32,9 +32,11 @@ def __init__( profile: Optional[str], *, mlflow_loader: Callable[[], Any] = _mlflow, + on_workspace: Callable[[], None] = lambda: None, ) -> None: self._profile = profile self._mlflow_loader = mlflow_loader + self._on_workspace = on_workspace self._module: Any | None = None @property @@ -48,6 +50,7 @@ def workspace_uri(self) -> str: return _workspace_uri(self._profile) def select_workspace(self) -> None: + self._on_workspace() self.module.set_tracking_uri(self.workspace_uri) def select_local(self, base_url: str) -> None: diff --git a/integrations/agentbricks/src/databricks_agentbricks/presentation/deploy.py b/integrations/agentbricks/src/databricks_agentbricks/presentation/deploy.py index 0f62142a7..2b8f535cf 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/presentation/deploy.py +++ b/integrations/agentbricks/src/databricks_agentbricks/presentation/deploy.py @@ -125,7 +125,7 @@ def present_deleted(result: DeleteResult, *, output: Optional[str]) -> None: elif result.wait_error: render.warning( f"Could not confirm deployment '{result.deployment}' was deleted: {result.wait_error} " - f"Check {check}; if the profile isn't signed in, run `agentbricks login` first." + f"Check {check}; if the profile isn't signed in, run `agentbricks auth login ` first." ) else: render.success(f"Deleted deployment '{result.deployment}'") diff --git a/integrations/agentbricks/src/databricks_agentbricks/presentation/dev.py b/integrations/agentbricks/src/databricks_agentbricks/presentation/dev.py index ddd438221..d718042b6 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/presentation/dev.py +++ b/integrations/agentbricks/src/databricks_agentbricks/presentation/dev.py @@ -38,6 +38,26 @@ def announce_bound_resources(preview: DevPreview) -> None: ) +def announce_agent_profile( + profile: str | None, + profile_source: str, + env_file_profile: str | None, + *, + env_credential: str | None, +) -> None: + """Note when the agent's credentials differ from what `.env` alone would suggest.""" + if env_credential: + render.console().print( + f"[dim]The agent authenticates with the {env_credential} credentials in .env, " + "not the resolved profile.[/]" + ) + elif profile and env_file_profile and env_file_profile != profile: + render.console().print( + f"[dim]The agent will run with profile '{profile}' (from {profile_source}), " + f"not .env's '{env_file_profile}'.[/]" + ) + + def announce_starting() -> None: """Say the agent is starting, in a neutral style: success is shown only once it answers. @@ -70,7 +90,7 @@ def announce_local_url( uses_runtime_api = has_chat_ui or server == AgentServer.AGENTBRICKS invoke_step = agent_invoke_step(f"--url {base}", uses_runtime_api=uses_runtime_api) if has_chat_ui: - fields = {"Chat UI": base} + fields: dict[str, str] = {"Chat UI": base} if trace_url: fields["Traces"] = trace_url render.success( diff --git a/integrations/agentbricks/src/databricks_agentbricks/presentation/help.py b/integrations/agentbricks/src/databricks_agentbricks/presentation/help.py index 174f881b7..17556a09c 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/presentation/help.py +++ b/integrations/agentbricks/src/databricks_agentbricks/presentation/help.py @@ -16,7 +16,7 @@ # getting-started path. Any command missing here still lists under "Other commands" (see # `_group.AgentBricksGroup`). _COMMAND_SECTIONS: tuple[tuple[str, tuple[str, ...]], ...] = ( - ("SETUP", ("login", "logout", "init", "doctor")), + ("SETUP", ("init", "auth", "doctor")), ("DEVELOP", ("dev", "tools", "memory", "sessions", "tracing")), ("SHIP", ("deploy", "deployments")), ("EXPERIMENTAL", ("experimental",)), @@ -29,14 +29,15 @@ _EXAMPLES: dict[CommandPath, tuple[Example, ...]] = { (): ( - ("agentbricks login --profile my-workspace", "authenticate and save a default profile"), - ("agentbricks init my-agent", "scaffold a new agent project"), + ("agentbricks auth login ", "sign in to a Databricks profile"), + ( + "agentbricks init my-agent --profile ", + "scaffold a new agent project", + ), ("cd my-agent", "enter the project directory"), ("agentbricks dev", "run the agent locally with a chat UI"), ("agentbricks deploy", "deploy the agent to Databricks Apps"), ), - ("login",): (("agentbricks login --profile my-workspace", "save a profile as your default"),), - ("logout",): (("agentbricks logout", "forget the saved default profile"),), ("init",): ( ("agentbricks init my-agent", "scaffold a new agent project"), ( @@ -48,6 +49,26 @@ "prepare a coding-agent migration bundle for an existing agent", ), ), + ("auth",): ( + ("agentbricks auth login ", "sign in to a Databricks profile"), + ("agentbricks auth switch ", "use that profile for the current project"), + ("agentbricks auth describe", "show the profile this directory would use"), + ), + ("auth", "switch"): ( + ("agentbricks auth switch ", "write the profile to the project's .env"), + ( + "agentbricks auth switch --source my-agent", + "set it for a project in another directory", + ), + ), + ("auth", "describe"): ( + ("agentbricks auth describe", "show the profile, its source, and its workspace"), + ("agentbricks -o json auth describe --source my-agent", "script against another project"), + ), + ("auth", "login"): ( + ("agentbricks auth login ", "validate it, signing in if needed"), + ("agentbricks auth login", "sign in to the project's .env profile"), + ), ("doctor",): ( ("agentbricks doctor .", "check an existing repository's Agent Bricks onboarding"), ), @@ -627,9 +648,8 @@ # crisp while the command's own `--help` page still shows its full docstring. Keep these under ~45 # chars so they never truncate. _SHORT_HELP: dict[CommandPath, str] = { - ("login",): "Authenticate and save a default profile", - ("logout",): "Forget the saved default profile", ("init",): "Scaffold a new agent project", + ("auth",): "Choose and sign in to a Databricks profile", ("doctor",): "Check an existing agent's onboarding", ("dev",): "Run the agent locally with a chat UI", ("deploy",): "Deploy an agent to Databricks Apps", @@ -686,7 +706,7 @@ def _example_epilog(examples: tuple[Example, ...]) -> str: def _getting_started_epilog() -> str: - """The root's numbered "Getting started" path: login → init → cd → dev → deploy. + """The root's numbered "Getting started" path: sign in → init → cd → dev → deploy. A numbered, ordered path — rather than an unlabeled grab-bag of examples — removes the "blank-page problem" for a first-time reader: it says *start here, in this order*. Reuses the @@ -714,9 +734,9 @@ def _root_epilog() -> str: auth = "\n".join( [ "\b", - "Not authenticated yet? Create a profile with the Databricks CLI first:", - " databricks auth login --profile my-workspace", - "Then `agentbricks login --profile my-workspace` saves it as your default.", + "Not authenticated yet? Sign in to a profile first (creates it if needed):", + " agentbricks auth login ", + "Then `agentbricks init --profile ` records it in the project's .env.", ] ) links = "\n".join(["\b", f"Docs: {_DOCS_URL}", f"Issues: {_ISSUES_URL}"]) diff --git a/integrations/agentbricks/src/databricks_agentbricks/presentation/profile.py b/integrations/agentbricks/src/databricks_agentbricks/presentation/profile.py new file mode 100644 index 000000000..7e8c82d04 --- /dev/null +++ b/integrations/agentbricks/src/databricks_agentbricks/presentation/profile.py @@ -0,0 +1,185 @@ +"""Human-facing output for the profile commands.""" + +from __future__ import annotations + +import pathlib + +import click + +from databricks_agentbricks.presentation import render +from databricks_agentbricks.projects.profiles import ( + auth_login_command, + auth_switch_command, + profile_exists, + workspace_for_profile, +) + + +def show_profile_switched( + directory: pathlib.Path, + name: str, + changed: bool, + host: str | None, + user: str | None, + output: str, +) -> None: + if output == "json": + render.emit_json( + { + "directory": str(directory), + "env_profile": name, + "changed": changed, + "workspace_host": host, + "signed_in_user": user, + } + ) + return + fields = {"Profile (.env)": name} + if host: + fields["Workspace"] = host + if user: + fields["Signed in as"] = user + render.success( + f"Project profile set → '{name}'" if changed else f"Project profile is already '{name}'", + fields=fields, + next_steps=None if user else [(auth_login_command(name), "Sign in")], + ) + + +def show_profile_description( + profile: str | None, + source: str, + host: str | None, + workspace: str | None, + project: str | None, + output: str, +) -> None: + if output == "json": + render.emit_json( + { + "profile": profile, + "source": source, + "workspace_host": host, + "project": project, + } + ) + return + fields = { + "Profile": profile or "none (Databricks SDK default)", + "Source": source, + } + if workspace: + fields["Workspace"] = workspace + fields["Project"] = project or "not in a project" + render.detail("Profile", profile or "SDK default", fields) + + +def show_login_result( + profile: str, + source: str, + user: str | None, + host: str | None, + env_profile: str | None, + output: str, +) -> None: + if output == "json": + render.emit_json({"profile": profile, "source": source, "user": user, "host": host}) + return + render.success( + f"Logged in to '{profile}'", + fields={"Profile": profile, "Source": source, "User": user, "Workspace": host}, + ) + if env_profile and env_profile != profile: + render.console().print( + f"[dim]This project still uses profile '{env_profile}' (from .env). " + f"Run `{auth_switch_command(profile)}` to switch it.[/]" + ) + + +def show_databricks_login_result(output: str) -> None: + """The result when `databricks auth login` chose the profile, so nothing about it is known here.""" + if output == "json": + render.emit_json({"profile": None, "source": "databricks auth login"}) + else: + render.success("Logged in with `databricks auth login`") + + +def announce_workspace( + profile: str | None, source: str, *, output: str, resolved_host: str | None = None +) -> None: + """One dim stderr line naming the profile and workspace a command is about to use. + + Every workspace command shares it so the target is never implicit. Text-only, like the + other auth notices. + """ + if output != "text": + return + host, workspace = workspace_for_profile(profile) + if resolved_host and resolved_host != host: + workspace = resolved_host + who = f"profile '{profile}' (from {source})" if profile else "Databricks SDK default auth" + click.echo( + click.style(f"Using {who} · {workspace or 'no host configured'}", dim=True), err=True + ) + + +def note_choose_profile(profiles: list[str], obj) -> None: + """Explain why a new project's `DEFAULT` profile wasn't signed in to: pick an existing one.""" + if obj.output == "text": + render.console().print( + "[dim]No profile was chosen, so .env uses 'DEFAULT', which isn't in your Databricks " + f"config. Run `{auth_switch_command(None)}` with one of: {', '.join(profiles)}.[/]" + ) + + +def warn_unknown_profile(profile: str, obj) -> None: + """Dim warning when a profile written to `.env` isn't in the Databricks config file yet. + + Warn, don't fail: a profile is often chosen before it has been created. Text-only so + `-o json` output stays machine-readable. + """ + if obj.output == "text" and not profile_exists(profile): + render.console().print( + f"[dim].env set to profile '{profile}', but it isn't in your Databricks " + f"config yet. Run `{auth_login_command(profile)}` before `agentbricks dev`.[/]" + ) + + +def announce_login(obj, profile: str, source: str) -> None: + """Say which profile is being signed in to, before any browser flow starts.""" + if obj.output != "text": + return + host, workspace = workspace_for_profile(profile) + where = ( + workspace + if host + else ( + "not in your Databricks config yet; `databricks auth login` will create it" + if not profile_exists(profile) + else "no host configured" + ) + ) + # stderr, not render's console (stdout-only): visible before any browser flow without + # polluting `-o json` stdout. + click.echo(f"Logging in to profile '{profile}' (from {source}) — {where}", err=True) + + +def warn_sign_in_failed(profile: str, exc: Exception, obj) -> None: + """Report a failed sign-in without raising: the profile choice still stands.""" + if obj.output == "text": + render.console().print(f"[dim]Sign-in to profile '{profile}' failed: {exc}[/]") + + +def warn_legacy_default_profile(profile: str, output: str) -> None: + """Point users of the removed `agentbricks login` at the profile it saved, which is now unused.""" + if output != "text": + return + click.echo( + click.style( + f"The default profile '{profile}' saved by the old `agentbricks login` is no longer " + f"used. Pass `-p {profile}`, or run `{auth_switch_command(profile)}` in the project " + "(delete ~/.agentbricks/config.json to hide this).", + dim=True, + ), + err=True, + ) diff --git a/integrations/agentbricks/src/databricks_agentbricks/projects/config.py b/integrations/agentbricks/src/databricks_agentbricks/projects/config.py index d45270b5e..a03c3c72f 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/projects/config.py +++ b/integrations/agentbricks/src/databricks_agentbricks/projects/config.py @@ -6,7 +6,7 @@ import pathlib import re from dataclasses import dataclass -from typing import Any +from typing import Any, Optional import tomli @@ -26,6 +26,25 @@ class ProjectMetadata: template: str | None +def is_agentbricks_project(path: pathlib.Path) -> bool: + """Whether `path` is an Agent Bricks project root (scaffold metadata or an Apps manifest).""" + return (path / _CONFIG_PATH).is_file() or (path / "app.yaml").is_file() + + +def find_project_root(start: pathlib.Path) -> Optional[pathlib.Path]: + """The Agent Bricks project `start` is in, so commands run from a subfolder use its `.env`. + + Only `start` itself may be an `app.yaml` project: an unrelated Databricks Apps folder above it + must not silently choose the profile, so ancestors count only with `init`'s metadata. + """ + if is_agentbricks_project(start): + return start + for ancestor in start.resolve().parents: + if (ancestor / _CONFIG_PATH).is_file(): + return ancestor + return None + + def write_project_metadata( project: pathlib.Path, *, diff --git a/integrations/agentbricks/src/databricks_agentbricks/projects/env_file.py b/integrations/agentbricks/src/databricks_agentbricks/projects/env_file.py new file mode 100644 index 000000000..81a3b746f --- /dev/null +++ b/integrations/agentbricks/src/databricks_agentbricks/projects/env_file.py @@ -0,0 +1,114 @@ +"""Reading and read-modify-writing of a project's local `.env` profile entry.""" + +from __future__ import annotations + +import io +import os +import pathlib +import re +import tempfile + +from dotenv import dotenv_values + +# Every line form dotenv reads as this key (optional `export`, optionally single-quoted key), so a +# rewrite can't leave an older value that dotenv would still pick up. +_PROFILE_LINE = re.compile( + r"\s*(export\s+)?(DATABRICKS_CONFIG_PROFILE|'DATABRICKS_CONFIG_PROFILE')\s*=" +) + + +_BOM = "\ufeff" + + +def _read_text(path: pathlib.Path) -> str: + # surrogateescape keeps bytes that aren't UTF-8 (e.g. in a secret) intact through a rewrite. + return path.read_text(encoding="utf-8", errors="surrogateescape") + + +def _lines(text: str) -> list[str]: + # Split on newlines only: str.splitlines() also breaks on control characters inside values. + lines = [line.removesuffix("\r") for line in text.removeprefix(_BOM).split("\n")] + return lines[:-1] if lines and lines[-1] == "" else lines + + +def _write_text_atomically(path: pathlib.Path, text: str) -> None: + """Replace `path` in one step so an interrupted write can't truncate a `.env` with secrets.""" + fd, tmp = tempfile.mkstemp(dir=path.parent, prefix=f".{path.name}.", suffix=".tmp") + try: + with os.fdopen(fd, "w", encoding="utf-8", errors="surrogateescape", newline="") as f: + f.write(text) + if path.exists(): + os.chmod(tmp, path.stat().st_mode & 0o7777) + os.replace(tmp, path) + except BaseException: + pathlib.Path(tmp).unlink(missing_ok=True) + raise + + +def read_env(path: pathlib.Path) -> dict[str, str]: + """Read a project `.env` the way the generated agent's `load_dotenv` does.""" + try: + text = _read_text(path) + except OSError: + return {} + return _parse_env(text) + + +def _parse_env(text: str, *, interpolate: bool = True) -> dict[str, str]: + # Same parser and default interpolation as the agent's `load_dotenv`, so both resolve one value. + values = dotenv_values(stream=io.StringIO(text.removeprefix(_BOM)), interpolate=interpolate) + return {key: value for key, value in values.items() if value is not None} + + +def _with_profile_line(lines: list[str], profile: str) -> list[str]: + """`lines` with the DATABRICKS_CONFIG_PROFILE entry set to `profile` (inserted when absent).""" + updated, replaced = [], False + for line in lines: + if _PROFILE_LINE.match(line): + if replaced: + # Drop later duplicates so the reader can't pick up an old value. + continue + updated.append(f"DATABRICKS_CONFIG_PROFILE={profile}") + replaced = True + else: + updated.append(line) + if not replaced: + updated.insert(0, f"DATABRICKS_CONFIG_PROFILE={profile}") + return updated + + +def write_env(dest: pathlib.Path, profile: str) -> bool: + """Seed a local `.env` from `.env.example` with DATABRICKS_CONFIG_PROFILE=. + + Returns True if a `.env` was written. Skips if `.env` already exists (never clobbers). The + template reads DATABRICKS_CONFIG_PROFILE for local model auth, so this makes the scaffolded + project runnable with `agentbricks dev` without a manual `cp .env.example .env` step. + """ + env_path = dest / ".env" + if env_path.exists(): + return False + example = dest / ".env.example" + base = _lines(_read_text(example)) if example.exists() else [] + _write_text_atomically(env_path, "\n".join(_with_profile_line(base, profile)) + "\n") + return True + + +def update_env_profile(dest: pathlib.Path, profile: str) -> bool: + """Point a project's `.env` at a different profile, preserving every other line. + + Returns True if the file changed (same profile -> False). A project without a `.env` gets one + seeded from `.env.example`, matching what a fresh scaffold produces. + """ + env_path = dest / ".env" + if not env_path.exists(): + return write_env(dest, profile) + before = _read_text(env_path) + if _parse_env(before, interpolate=False).get("DATABRICKS_CONFIG_PROFILE") == profile: + # Compare the parsed profile, not the file text: a lone CRLF or trailing-newline + # difference must not count as a change. The literal value is compared, not the + # interpolated one: `${VAR}` resolves to something else when the agent loads it. + return False + bom = _BOM if before.startswith(_BOM) else "" + lines = _with_profile_line(_lines(before), profile) + _write_text_atomically(env_path, bom + "\n".join(lines) + "\n") + return True diff --git a/integrations/agentbricks/src/databricks_agentbricks/projects/profiles.py b/integrations/agentbricks/src/databricks_agentbricks/projects/profiles.py new file mode 100644 index 000000000..4184af131 --- /dev/null +++ b/integrations/agentbricks/src/databricks_agentbricks/projects/profiles.py @@ -0,0 +1,224 @@ +"""Local Databricks profile configuration: which profile a command uses, and what the config says. + +Everything here reads local files and the environment only; it never contacts a workspace. +""" + +from __future__ import annotations + +import configparser +import dataclasses +import json +import os +import pathlib +from typing import Optional + +from databricks_agentbricks.projects.env_file import read_env + +# Where a resolved profile came from; the strings are user-facing too. +SOURCE_FLAG = "--profile" +SOURCE_ENV_FILE = ".env" +SOURCE_ENV_VAR = "DATABRICKS_CONFIG_PROFILE" +SOURCE_SDK_DEFAULT = "Databricks SDK default" +SOURCE_SETTINGS_DEFAULT = "databricks auth switch" +SOURCE_CONFIG_DEFAULT = "default" + +_SETTINGS_SECTION = "__settings__" +# Any of these without a profile means the SDK authenticates from env vars and never reads the +# config file. +_ENV_AUTH_VARS = ("DATABRICKS_HOST", "DATABRICKS_TOKEN", "DATABRICKS_CLIENT_ID") + + +def auth_login_command(profile: Optional[str]) -> str: + return f"agentbricks auth login {profile or ''}" + + +def auth_switch_command(profile: Optional[str]) -> str: + return f"agentbricks auth switch {profile or ''}" + + +@dataclasses.dataclass(frozen=True) +class ProfileInfo: + """The profile a command runs with, and where it came from (for display and re-resolution).""" + + name: Optional[str] + source: str + + +def resolve_profile( + flag: Optional[str], + project_dir: Optional[pathlib.Path] = None, + env_values: Optional[dict[str, str]] = None, +) -> ProfileInfo: + """Pick the Databricks profile for a command, most to least specific. + + 1. the `-p/--profile` flag, + 2. the project `.env`'s `DATABRICKS_CONFIG_PROFILE` (only when `project_dir` is given, so the CLI + and the locally running agent, which reads the same `.env`, agree on one profile), + 3. the `DATABRICKS_CONFIG_PROFILE` environment variable, + 4. none — the Databricks SDK's own default authentication. + + ``env_values`` lets a caller that already parsed ``/.env`` pass it in instead of + having it read again. + """ + if flag: + return ProfileInfo(flag, SOURCE_FLAG) + if project_dir is not None: + if env_values is None: + env_values = read_env(project_dir / ".env") + file_profile = env_values.get("DATABRICKS_CONFIG_PROFILE") + if file_profile: + return ProfileInfo(file_profile, SOURCE_ENV_FILE) + env_profile = os.environ.get("DATABRICKS_CONFIG_PROFILE") + if env_profile: + return ProfileInfo(env_profile, SOURCE_ENV_VAR) + return ProfileInfo(None, SOURCE_SDK_DEFAULT) + + +def login_target( + name: Optional[str], flag: ProfileInfo, env_profile: Optional[str] +) -> Optional[tuple[str, str]]: + """(profile, source) for `auth login`: NAME, the global `-p`, or the project's `.env` only. + + Signing in is an explicit act, so the ambient DATABRICKS_CONFIG_PROFILE and the SDK default + are deliberately not consulted. None means nothing chose one. + """ + if name: + return name, "argument" + if flag.source == SOURCE_FLAG and flag.name: + return flag.name, SOURCE_FLAG + if env_profile: + return env_profile, SOURCE_ENV_FILE + return None + + +def _config_path() -> pathlib.Path: + # expanduser: the SDK expands `~` in DATABRICKS_CONFIG_FILE; without it a tilde path reads nothing. + return pathlib.Path( + os.getenv("DATABRICKS_CONFIG_FILE", str(pathlib.Path.home() / ".databrickscfg")) + ).expanduser() + + +_config_cache: Optional[tuple[tuple[str, int, int], dict[str, dict[str, str]]]] = None + + +def _config_sections() -> dict[str, dict[str, str]]: + """Every section's own keys, `[DEFAULT]` included as an ordinary section; never raises. + + Sections don't inherit `[DEFAULT]`, matching how the SDK reads a named profile. No + interpolation, so a `%` in a value is taken verbatim. Cached on the file's identity, since one + command looks the config up several times and tests rewrite it between commands. + """ + global _config_cache + path = _config_path() + try: + stat = path.stat() + except OSError: + return {} + key = (str(path), stat.st_mtime_ns, stat.st_size) + if _config_cache is not None and _config_cache[0] == key: + return _config_cache[1] + parser = configparser.ConfigParser(interpolation=None, default_section="\0no-fallback") + try: + parser.read(path) + except (OSError, configparser.Error): + pass # keep whatever parsed before the problem, as the SDK-facing lookups always have + sections = {name: dict(parser[name]) for name in parser.sections()} + _config_cache = (key, sections) + return sections + + +def _profile_section(profile: str) -> dict[str, str]: + return _config_sections().get(profile, {}) + + +def profile_exists(profile: str) -> bool: + """Whether the profile is configured in the Databricks config file; never raises. + + `[DEFAULT]` counts only once it has a key, the way the SDK falls back to it. + """ + sections = _config_sections() + if profile == configparser.DEFAULTSECT: + return bool(sections.get(profile)) + return profile in sections + + +def profile_host(profile: Optional[str]) -> Optional[str]: + """Host configured for `profile` in the Databricks config file; None when unknown. Never raises.""" + if not profile: + return None + return _profile_section(profile).get("host") + + +def configured_profiles() -> list[str]: + """Profile names in the Databricks config file, `[DEFAULT]` first when it exists.""" + names = [ + name + for name in _config_sections() + if name not in (_SETTINGS_SECTION, configparser.DEFAULTSECT) + ] + return ([configparser.DEFAULTSECT] if profile_exists(configparser.DEFAULTSECT) else []) + names + + +def settings_default_profile() -> Optional[str]: + """The default profile chosen with `databricks auth switch` ([__settings__] default_profile).""" + return _profile_section(_SETTINGS_SECTION).get("default_profile", "").strip() or None + + +def env_auth_configured() -> bool: + """Whether env-var auth is set, so without a profile the SDK never reads the config file.""" + return any(os.environ.get(name) for name in _ENV_AUTH_VARS) + + +def sdk_default_profile() -> tuple[str, str]: + """(profile, source) the SDK falls back to without one: the switched default, else DEFAULT. + + Only meaningful when `env_auth_configured()` is False; otherwise the config file isn't read. + """ + switched = settings_default_profile() + if switched: + return switched, SOURCE_SETTINGS_DEFAULT + return configparser.DEFAULTSECT, SOURCE_CONFIG_DEFAULT + + +def workspace_for_profile(profile: Optional[str]) -> tuple[Optional[str], Optional[str]]: + """An offline workspace estimate as (host, display text), for notices and `auth describe`. + + An exported host wins even with an explicit profile. Without a profile, common environment + credentials suppress the file; otherwise prefer the switched default, then `[DEFAULT]`. + Callers with a constructed SDK client should use its resolved host instead. + """ + env_host = os.environ.get("DATABRICKS_HOST") + if env_host: + return env_host, f"{env_host} (from DATABRICKS_HOST)" + if profile: + host = profile_host(profile) + return host, host + if env_auth_configured(): + return None, "none configured" + default_profile, _ = sdk_default_profile() + default_host = profile_host(default_profile) + if default_host: + return default_host, f"{default_host} ([{default_profile}] profile)" + return None, "none configured" + + +def legacy_default_profile() -> Optional[str]: + """The default profile the removed `agentbricks login` saved, if its file is still around.""" + base = os.environ.get("AGENTBRICKS_CONFIG_HOME") + path = (pathlib.Path(base) if base else pathlib.Path.home() / ".agentbricks") / "config.json" + try: + return json.loads(path.read_text()).get("profile") or None + except (OSError, ValueError, AttributeError): + return None + + +def sdk_auth_type(profile: Optional[str]) -> Optional[str]: + """The auth_type the Databricks SDK would use for `profile`, resolved the way the SDK does.""" + env_auth_type = os.environ.get("DATABRICKS_AUTH_TYPE") + if env_auth_type: + return env_auth_type + if not profile: + if env_auth_configured(): + return None + profile, _ = sdk_default_profile() + return _profile_section(profile).get("auth_type") diff --git a/integrations/agentbricks/src/databricks_agentbricks/services/deploy_service.py b/integrations/agentbricks/src/databricks_agentbricks/services/deploy_service.py index a36e01a06..2d44dccfe 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/services/deploy_service.py +++ b/integrations/agentbricks/src/databricks_agentbricks/services/deploy_service.py @@ -616,7 +616,7 @@ def check_auth(self) -> None: except Exception as exc: # noqa: BLE001 - any auth/transport failure reads the same way raise AgentCliError( "Could not authenticate with the Databricks workspace.", - hint=f"{exc}\nRun `agentbricks login --profile ` (or pass --profile), " + hint=f"{exc}\nRun `agentbricks auth login ` (or pass --profile), " "then rerun the command.", ) from exc diff --git a/integrations/agentbricks/src/databricks_agentbricks/services/deployment/app_provisioner.py b/integrations/agentbricks/src/databricks_agentbricks/services/deployment/app_provisioner.py index 79512d8a9..24fb2e280 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/services/deployment/app_provisioner.py +++ b/integrations/agentbricks/src/databricks_agentbricks/services/deployment/app_provisioner.py @@ -80,7 +80,7 @@ def _app_lookup_error(name: DeploymentName, exc: Exception) -> AgentCliError: return AgentCliError( f"Could not check whether deployment '{name}' exists.", error_code=getattr(exc, "error_code", None), - hint=f"{exc}\nIf the profile isn't signed in, run `agentbricks login`; otherwise check the " + hint=f"{exc}\nIf the profile isn't signed in, run `agentbricks auth login `; otherwise check the " "network and the profile's host, then rerun the command.", ) diff --git a/integrations/agentbricks/src/databricks_agentbricks/services/dev_service.py b/integrations/agentbricks/src/databricks_agentbricks/services/dev_service.py index be3b791f8..d1befa60b 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/services/dev_service.py +++ b/integrations/agentbricks/src/databricks_agentbricks/services/dev_service.py @@ -10,9 +10,9 @@ import json import pathlib -from collections.abc import Callable, Iterator +from collections.abc import Callable, Iterator, Mapping from contextlib import contextmanager -from dataclasses import dataclass +from dataclasses import dataclass, field from databricks_agentbricks.clients.apps_client import AppsClient from databricks_agentbricks.clients.local_tracing_client import LocalTracingClient @@ -50,6 +50,8 @@ class DevRequest: source: str prepare_environment: bool | None app_port: int | None + # Extra env for the local-only manifest (e.g. the resolved Databricks profile). + local_env: Mapping[str, str] = field(default_factory=dict) @dataclass(frozen=True) @@ -102,6 +104,16 @@ def __init__( self._port_in_use = port_in_use self._free_port = free_port + def check_project(self, source: str) -> pathlib.Path: + """Fail unless ``source`` has an ``app.yaml`` to run; returns that path. Local only.""" + app_yaml = pathlib.Path(source) / "app.yaml" + if not app_yaml.exists(): + raise AgentCliError( + f"No app.yaml found at {app_yaml}.", + hint="Run from a scaffolded project, or pass --source (see `agentbricks init`).", + ) + return app_yaml + @contextmanager def prepare(self, request: DevRequest) -> Iterator[DevPlan]: """Prepare a local run and clean up its temporary resources on context exit. @@ -111,12 +123,7 @@ def prepare(self, request: DevRequest) -> Iterator[DevPlan]: is removed along with the local tracing process even when either of those operations fails. """ source_dir = pathlib.Path(request.source) - app_yaml = source_dir / "app.yaml" - if not app_yaml.exists(): - raise AgentCliError( - f"No app.yaml found at {app_yaml}.", - hint="Run from a scaffolded project, or pass --source (see `agentbricks init`).", - ) + app_yaml = self.check_project(request.source) project = self._project_resolver.load(source_dir) if project is not None and project.tools: @@ -147,9 +154,10 @@ def prepare(self, request: DevRequest) -> Iterator[DevPlan]: if prepare_environment is None: prepare_environment = not (source_dir / ".venv").exists() + manifest_env = {**request.local_env, **tracing_env} entry_point = _dev_entry_point( app_yaml, - tracing_env or None, + manifest_env or None, allowed_origins=project.allowed_origins if project is not None else None, ) preview = DevPreview( diff --git a/integrations/agentbricks/src/databricks_agentbricks/services/workspace_auth.py b/integrations/agentbricks/src/databricks_agentbricks/services/workspace_auth.py new file mode 100644 index 000000000..fc679c583 --- /dev/null +++ b/integrations/agentbricks/src/databricks_agentbricks/services/workspace_auth.py @@ -0,0 +1,218 @@ +"""Workspace authentication for Agent Bricks commands: validation, the preflight, and sign-in. + +Which profile a command uses, and what the local config says about it, live in +`projects.profiles`; this is where a profile meets the workspace. The service never reads the +terminal: callers say whether a person is there to complete a browser login. +""" + +from __future__ import annotations + +import enum +import threading +from typing import Callable, Optional + +from databricks_agentbricks.clients.api_client_provider import validate_profile as _validate_profile +from databricks_agentbricks.clients.databricks_cli import run_databricks_login +from databricks_agentbricks.errors import AgentCliError +from databricks_agentbricks.projects.profiles import ( + auth_login_command, + auth_switch_command, + sdk_auth_type, +) +from databricks_agentkit._api_client import _AgentBricksApiClient + + +class _ValidationFailure(enum.Enum): + NOT_SIGNED_IN = "not signed in" # a login could fix it + UNREACHABLE = "unreachable" # the network or the server; signing in wouldn't fix it + OTHER = "other" + + +def _classify_validation_error(error: Exception) -> _ValidationFailure: + from databricks.sdk.errors import ( + DeadlineExceeded, + InternalError, + TemporarilyUnavailable, + Unauthenticated, + ) + from requests.exceptions import ConnectionError as RequestsConnectionError + + # A bounded-validation timeout is a stall, not a credential problem: a login wouldn't unstick it. + if isinstance( + error, + ( + RequestsConnectionError, + InternalError, + TemporarilyUnavailable, + DeadlineExceeded, + TimeoutError, + ), + ): + return _ValidationFailure.UNREACHABLE + # Client construction wraps every failure (e.g. an unknown profile) in AgentCliError. + if isinstance(error, (AgentCliError, Unauthenticated)): + return _ValidationFailure.NOT_SIGNED_IN + return _ValidationFailure.OTHER + + +def _profile_label(profile: str, source: str) -> str: + return f"Databricks profile {profile!r} (from {source})" + + +def _sign_in_hint(profile: Optional[str], *, needs_terminal: bool = False) -> str: + where = " in an interactive terminal" if needs_terminal else "" + if profile: + return f"Run `{auth_login_command(profile)}`{where} to sign in, then rerun the command." + return ( + f"Run `{auth_login_command(None)}`{where} to sign in, then pass `-p ` or " + f"choose one for the project with `{auth_switch_command(None)}` (or " + "DATABRICKS_CONFIG_PROFILE), and rerun the command." + ) + + +def _not_signed_in_error(profile: Optional[str], source: str) -> AgentCliError: + if profile: + return AgentCliError( + f"{_profile_label(profile, source)} isn't authenticated.", hint=_sign_in_hint(profile) + ) + return AgentCliError( + "No Databricks profile is configured and the Databricks SDK's default authentication " + "didn't authenticate.", + hint=_sign_in_hint(None), + ) + + +def _unreachable_error(profile: Optional[str], source: str, error: Exception) -> AgentCliError: + target = f"profile {profile!r} (from {source})" if profile else "the default auth" + if isinstance(error, TimeoutError): + return AgentCliError( + f"The Databricks workspace for {target} didn't respond within " + f"{_PREFLIGHT_TIMEOUT_S:g}s.", + hint="Check the network (VPN/proxy) and the workspace, then rerun the command.", + ) + return AgentCliError( + f"Could not reach the Databricks workspace for {target}: {error}", + hint="Check the network (VPN/proxy) and the profile's host, then rerun the command.", + ) + + +# A preflight auth check must give up long before a user could plausibly complete a browser login: +# its job is to fail fast, not to host the login. `auth login` is the interactive path. +_PREFLIGHT_TIMEOUT_S = 20.0 + + +def _needs_browser_login(profile: Optional[str]) -> bool: + """An `external-browser` profile: only a person at a browser can sign it in. + + Validating it in-process would start that browser flow and leave it running past the bounded + wait, so every path that doesn't host a login fails fast instead. `databricks auth login` + switches the profile to `databricks-cli`, after which it validates normally. + """ + return sdk_auth_type(profile) == "external-browser" + + +def _validate_bounded( + profile: Optional[str], +) -> tuple[Optional[tuple[_AgentBricksApiClient, str]], Optional[Exception]]: + """`_validate_profile` on a daemon worker thread, giving up after `_PREFLIGHT_TIMEOUT_S`. + + The SDK call can stall on an unresponsive workspace, so it is bounded. Daemon because executor + threads are joined at interpreter exit and a stalled call must not block the CLI's exit too. + Returns (result, None) on success and (None, error) on failure or timeout. + """ + outcome: list[tuple[Optional[tuple[_AgentBricksApiClient, str]], Optional[Exception]]] = [] + + def _validate() -> None: + try: + outcome.append((_validate_profile(profile), None)) + except Exception as exc: # carried to the caller, never re-raised raw + outcome.append((None, exc)) + + worker = threading.Thread(target=_validate, daemon=True, name="agentbricks-auth-preflight") + worker.start() + worker.join(_PREFLIGHT_TIMEOUT_S) + if not outcome: + return None, TimeoutError(f"no response within {_PREFLIGHT_TIMEOUT_S:g}s") + return outcome[0] + + +ValidationResult = tuple[Optional[tuple[_AgentBricksApiClient, str]], Optional[Exception]] + + +class WorkspaceAuthService: + """Coordinate bounded validation and login with supplied workspace collaborators.""" + + def __init__( + self, + validate_bounded: Callable[[Optional[str]], ValidationResult], + needs_browser_login: Callable[[Optional[str]], bool], + login: Callable[[Optional[str]], None], + ) -> None: + self._validate_bounded = validate_bounded + self._needs_browser_login = needs_browser_login + self._login = login + + def preflight_auth(self, profile: Optional[str], source: str) -> _AgentBricksApiClient: + """Fail before starting work when the resolved profile cannot authenticate.""" + if self._needs_browser_login(profile): + # Validation would start a browser flow that this command cannot complete. + raise _not_signed_in_error(profile, source) + + client_and_user, error = self._validate_bounded(profile) + if client_and_user is not None: + return client_and_user[0] + assert error is not None + if _classify_validation_error(error) is _ValidationFailure.UNREACHABLE: + raise _unreachable_error(profile, source, error) from error + raise _not_signed_in_error(profile, source) from error + + def authenticate_profile( + self, profile: str, source: str, *, interactive: bool + ) -> tuple[_AgentBricksApiClient, str]: + """Validate credentials, signing in once when a person can complete login.""" + label = _profile_label(profile, source) + error: Optional[Exception] + if self._needs_browser_login(profile): + # Validation would start a browser flow that only the login command should host. + error = None + else: + client_and_user, error = self._validate_bounded(profile) + if client_and_user is not None: + return client_and_user + assert error is not None + failure = _classify_validation_error(error) + if failure is _ValidationFailure.UNREACHABLE: + raise _unreachable_error(profile, source, error) from error + if failure is _ValidationFailure.OTHER: + raise AgentCliError(f"Could not validate {label}: {error}") from error + if not interactive: + detail = ( + "external-browser auth needs an interactive login" if error is None else str(error) + ) + raise AgentCliError( + f"Could not validate {label}: {detail}", + hint=_sign_in_hint(profile, needs_terminal=True), + ) from error + self._login(profile) + client_and_user, error = self._validate_bounded(profile) + if client_and_user is None: + raise AgentCliError( + f"Databricks login completed, but {label} could not be validated: {error}" + ) from error + return client_and_user + + +def _service() -> WorkspaceAuthService: + return WorkspaceAuthService(_validate_bounded, _needs_browser_login, run_databricks_login) + + +def preflight_auth(profile: Optional[str], source: str) -> _AgentBricksApiClient: + """Validate auth before a command spins anything up, reusing the validated client.""" + return _service().preflight_auth(profile, source) + + +def authenticate_profile( + profile: str, source: str, *, interactive: bool +) -> tuple[_AgentBricksApiClient, str]: + """Validate a profile and offer interactive login when credentials are missing.""" + return _service().authenticate_profile(profile, source, interactive=interactive) diff --git a/integrations/agentbricks/src/databricks_agentkit/_api_client.py b/integrations/agentbricks/src/databricks_agentkit/_api_client.py index 26f981afa..ce33c4b04 100644 --- a/integrations/agentbricks/src/databricks_agentkit/_api_client.py +++ b/integrations/agentbricks/src/databricks_agentkit/_api_client.py @@ -8,9 +8,6 @@ from __future__ import annotations -import configparser -import os -import pathlib import time from typing import TYPE_CHECKING, Any, Optional from urllib.parse import quote @@ -113,15 +110,10 @@ def memory_entry_path(store: str, entry: str) -> str: def _profile_host(profile: str) -> Optional[str]: - config_path = pathlib.Path( - os.getenv("DATABRICKS_CONFIG_FILE", pathlib.Path.home() / ".databrickscfg") - ) - parser = configparser.ConfigParser() - try: - parser.read(config_path) - except (OSError, configparser.Error): - return None - return parser.get(profile, "host", fallback=None) + # Imported here so the runtime import of this module doesn't load the CLI's profile lookups. + from databricks_agentbricks.projects.profiles import profile_host + + return profile_host(profile) def _bound_retry_timeout(client: WorkspaceClient) -> WorkspaceClient: @@ -174,7 +166,7 @@ def __init__( raise AgentCliError( f"Could not initialize Databricks auth: {exc}", hint="Select an existing profile with `agentbricks --profile ` " - "or authenticate and save it with `agentbricks login --profile `.", + "or sign in with `agentbricks auth login `.", ) from exc @property diff --git a/integrations/agentbricks/tests/unit_tests/apps_client_test.py b/integrations/agentbricks/tests/unit_tests/apps_client_test.py index ce6d29110..8689ce064 100644 --- a/integrations/agentbricks/tests/unit_tests/apps_client_test.py +++ b/integrations/agentbricks/tests/unit_tests/apps_client_test.py @@ -57,7 +57,7 @@ def test_exists_raises_when_the_probe_fails_for_another_reason(stderr): client.exists("agent-bricks-myapp") hint = raised.value.hint or "" assert hint.startswith(stderr or "The `databricks apps get` command failed without a message.") - assert "`agentbricks login --profile prof`" in hint + assert "`agentbricks auth login prof`" in hint def test_service_principal_parses_from_json(): diff --git a/integrations/agentbricks/tests/unit_tests/auth_test.py b/integrations/agentbricks/tests/unit_tests/auth_test.py index 5de921638..6937389c0 100644 --- a/integrations/agentbricks/tests/unit_tests/auth_test.py +++ b/integrations/agentbricks/tests/unit_tests/auth_test.py @@ -1,193 +1,385 @@ -"""Unit tests for `agentbricks login` / `logout` and the saved-profile helpers. +"""Unit tests for `agentbricks auth` (switch / describe / login) and the profile-aware root context. -`AGENTBRICKS_CONFIG_HOME` redirects the config file into a tmp dir, and `_AgentBricksApiClient` -is stubbed so login never touches the network. +The workspace boundary (`workspace_auth._validate_bounded`), `subprocess.run`, and the terminal check are +stubbed; `DATABRICKS_CONFIG_FILE` points at a tmp file (see conftest.py). """ from __future__ import annotations import json -import sys +import pathlib from unittest import mock import pytest from click.testing import CliRunner from databricks.sdk.errors import Unauthenticated +from databricks_agentbricks.cli import app as cli from databricks_agentbricks.cli import auth +from databricks_agentbricks.cli.app import CliContext +from databricks_agentbricks.clients import api_client_provider +from databricks_agentbricks.errors import AgentCliError +from databricks_agentbricks.projects.config import write_project_metadata +from databricks_agentbricks.projects.env_file import read_env +from databricks_agentbricks.projects.profiles import ProfileInfo +from databricks_agentbricks.services import workspace_auth -@pytest.fixture(autouse=True) -def _disable_telemetry(monkeypatch): - monkeypatch.setenv("AGENTBRICKS_DISABLE_TELEMETRY", "1") +def _write_env(directory: pathlib.Path, profile: str) -> None: + (directory / ".env").write_text(f"DATABRICKS_CONFIG_PROFILE={profile}\n") -class _Ctx: - """Stand-in for CliContext: auth commands read only .profile and .output.""" +@pytest.fixture +def project(tmp_path: pathlib.Path, monkeypatch: pytest.MonkeyPatch) -> pathlib.Path: + """An Agent Bricks project directory (has app.yaml) that is also the cwd.""" + (tmp_path / "app.yaml").write_text("command: []\n") + monkeypatch.chdir(tmp_path) + return tmp_path - def __init__(self, profile=None, output="text"): - self.profile = profile - self.output = output +@pytest.fixture +def env_project(project: pathlib.Path) -> pathlib.Path: + """The `project`, whose .env selects profile 'proj'.""" + _write_env(project, "proj") + return project -def _stub_client(monkeypatch, user="me@example.com", host="https://ws"): - fake = mock.Mock() - fake.current_user = user - fake.host = host - monkeypatch.setattr(auth, "_AgentBricksApiClient", lambda profile: fake) +@pytest.fixture +def outside_project(tmp_path: pathlib.Path, monkeypatch: pytest.MonkeyPatch) -> pathlib.Path: + """A cwd that is not an Agent Bricks project (no app.yaml).""" + monkeypatch.chdir(tmp_path) + return tmp_path -def test_load_default_profile_missing_returns_none(tmp_path, monkeypatch): - monkeypatch.setenv("AGENTBRICKS_CONFIG_HOME", str(tmp_path)) - assert auth.load_default_profile() is None +@pytest.fixture +def fake_api_client(monkeypatch: pytest.MonkeyPatch) -> mock.Mock: + """Replace the workspace client the provider builds, so first use makes no SDK calls.""" + client_cls = mock.Mock() + client_cls.return_value.host = None + monkeypatch.setattr(api_client_provider, "_AgentBricksApiClient", client_cls) + return client_cls -def test_login_persists_and_load_round_trips(tmp_path, monkeypatch): - monkeypatch.setenv("AGENTBRICKS_CONFIG_HOME", str(tmp_path)) - _stub_client(monkeypatch) - databricks_login = mock.Mock() - monkeypatch.setattr(auth.subprocess, "run", databricks_login) - result = CliRunner().invoke(auth.login, ["--profile", "my-workspace"], obj=_Ctx()) - assert result.exit_code == 0, result.output - assert auth.load_default_profile() == "my-workspace" - databricks_login.assert_not_called() +def _ctx(profile=None, output="text") -> CliContext: + ctx = CliContext(profile, output) + ctx.select_profile() + return ctx -def test_login_json_output(tmp_path, monkeypatch): - monkeypatch.setenv("AGENTBRICKS_CONFIG_HOME", str(tmp_path)) - _stub_client(monkeypatch) - result = CliRunner().invoke(auth.login, ["-p", "prof"], obj=_Ctx(output="json")) - assert result.exit_code == 0, result.output - assert json.loads(result.output) == { - "profile": "prof", - "user": "me@example.com", - "host": "https://ws", - } + +def _invoke_profile(args, **ctx_kwargs): + return CliRunner().invoke(auth.auth_group, args, obj=_ctx(**ctx_kwargs)) -def test_login_falls_back_to_global_profile(tmp_path, monkeypatch): - # When the command's own --profile is omitted, the global -p (obj.profile) is saved. - monkeypatch.setenv("AGENTBRICKS_CONFIG_HOME", str(tmp_path)) - _stub_client(monkeypatch) - result = CliRunner().invoke(auth.login, [], obj=_Ctx(profile="from-global")) +def _get_json(args: list[str]) -> dict: + result = CliRunner().invoke(cli.agentbricks, ["-o", "json", *args]) assert result.exit_code == 0, result.output - assert auth.load_default_profile() == "from-global" + return json.loads(result.output) -def test_login_without_any_profile_errors(tmp_path, monkeypatch): - monkeypatch.setenv("AGENTBRICKS_CONFIG_HOME", str(tmp_path)) - result = CliRunner().invoke(auth.login, [], obj=_Ctx(profile=None)) - assert result.exit_code != 0 - assert auth.load_default_profile() is None +def _flat(output: str) -> str: + return " ".join(output.split()) -def test_login_configures_invalid_profile_then_revalidates(tmp_path, monkeypatch): - monkeypatch.setenv("AGENTBRICKS_CONFIG_HOME", str(tmp_path)) - validated = mock.Mock(current_user="me@example.com", host="https://ws") - agentbricks_client = mock.Mock(side_effect=[auth.AgentCliError("no credentials"), validated]) - monkeypatch.setattr(auth, "_AgentBricksApiClient", agentbricks_client) - monkeypatch.setattr(auth, "_is_interactive", lambda: True) - databricks_login = mock.Mock(return_value=mock.Mock(returncode=0)) - monkeypatch.setattr(auth.subprocess, "run", databricks_login) +def test_context_precedence_is_flag_then_env_file_then_env_var(tmp_path, monkeypatch): + monkeypatch.setenv("DATABRICKS_CONFIG_PROFILE", "from-env-var") + _write_env(tmp_path, "proj") + ambient, project, flagged = _ctx(), _ctx(), _ctx("from-flag") - result = CliRunner().invoke(auth.login, ["--profile", "prof"], obj=_Ctx()) + project.select_profile(tmp_path) + flagged.select_profile(tmp_path) - assert result.exit_code == 0, result.output - assert auth.load_default_profile() == "prof" - assert agentbricks_client.call_args_list == [mock.call("prof"), mock.call("prof")] - databricks_login.assert_called_once_with( - ["databricks", "auth", "login", "--profile", "prof"], - text=True, - check=False, - stdout=mock.ANY, + assert ambient.profile_info == ProfileInfo("from-env-var", "DATABRICKS_CONFIG_PROFILE") + assert project.profile_info == ProfileInfo("proj", ".env") + assert flagged.profile_info == ProfileInfo("from-flag", "--profile") + + +def test_profile_flag_after_subcommand_wins(env_project): + payload = _get_json(["-p", "before", "auth", "describe", "-p", "after"]) + + assert (payload["profile"], payload["source"]) == ("after", "--profile") + + +def test_outside_a_project_ignores_env_file(outside_project, monkeypatch): + _write_env(outside_project, "proj") + + payload = _get_json(["auth", "describe"]) + assert (payload["profile"], payload["source"]) == (None, "Databricks SDK default") + + monkeypatch.setenv("DATABRICKS_CONFIG_PROFILE", "ambient") + payload = _get_json(["auth", "describe"]) + assert (payload["profile"], payload["source"], payload["project"]) == ( + "ambient", + "DATABRICKS_CONFIG_PROFILE", + None, ) -def test_databricks_login_routes_child_stdout_to_stderr(monkeypatch): - databricks_login = mock.Mock(return_value=mock.Mock(returncode=0)) - monkeypatch.setattr(auth.subprocess, "run", databricks_login) +def test_first_workspace_call_names_profile_and_host_once( + capsys, write_databrickscfg, fake_api_client +): + write_databrickscfg("[ws]\nhost = https://ws.example\n") + ctx = _ctx("ws") + + ctx.client() + ctx.client() + + captured = capsys.readouterr() + assert captured.out == "" + assert captured.err.count("Using profile 'ws' (from --profile) · https://ws.example") == 1 + + +def test_first_workspace_call_reports_resolved_host_once( + capsys, write_databrickscfg, fake_api_client +): + write_databrickscfg("[ws]\nhost = https://profile.example\n") + fake_api_client.return_value.host = "https://resolved.example" + ctx = _ctx("ws") + + ctx.client() + ctx.client() + + error_output = capsys.readouterr().err + assert error_output.count("Using profile 'ws' (from --profile) · https://resolved.example") == 1 + assert "https://profile.example" not in error_output + - auth._run_databricks_login("prof") +def test_failed_client_construction_still_names_offline_workspace( + capsys, write_databrickscfg, fake_api_client +): + write_databrickscfg("[ws]\nhost = https://profile.example\n") + fake_api_client.side_effect = AgentCliError("synthetic authentication failure") - databricks_login.assert_called_once_with( - ["databricks", "auth", "login", "--profile", "prof"], - text=True, - check=False, - stdout=sys.stderr, + with pytest.raises(AgentCliError): + _ctx("ws").client() + + assert ( + "Using profile 'ws' (from --profile) · https://profile.example" in capsys.readouterr().err ) -def test_login_reauthenticates_unauthenticated_api_response(tmp_path, monkeypatch): - monkeypatch.setenv("AGENTBRICKS_CONFIG_HOME", str(tmp_path)) - validated = mock.Mock(current_user="me@example.com", host="https://ws") - validate_profile = mock.Mock( - side_effect=[Unauthenticated("expired credentials"), (validated, "me@example.com")] +def test_workspace_notice_names_sdk_default_auth(capsys, monkeypatch, fake_api_client): + monkeypatch.setenv("DATABRICKS_HOST", "https://env.example") + + _ctx().client() + + assert ( + "Using Databricks SDK default auth · https://env.example (from DATABRICKS_HOST)" + in capsys.readouterr().err ) - monkeypatch.setattr(auth, "_validate_profile", validate_profile) - monkeypatch.setattr(auth, "_is_interactive", lambda: True) - databricks_login = mock.Mock(return_value=mock.Mock(returncode=0)) - monkeypatch.setattr(auth.subprocess, "run", databricks_login) - result = CliRunner().invoke(auth.login, ["--profile", "prof"], obj=_Ctx()) - assert result.exit_code == 0, result.output - assert validate_profile.call_count == 2 - databricks_login.assert_called_once() +def test_workspace_notice_reports_environment_host_for_named_profile( + capsys, monkeypatch, write_databrickscfg, fake_api_client +): + write_databrickscfg("[ws]\nhost = https://profile.example\n") + monkeypatch.setenv("DATABRICKS_HOST", "https://env.example") + monkeypatch.setenv("DATABRICKS_TOKEN", "test-token") + _ctx("ws").client() -def test_login_does_not_reauthenticate_non_auth_validation_error(tmp_path, monkeypatch): - monkeypatch.setenv("AGENTBRICKS_CONFIG_HOME", str(tmp_path)) - monkeypatch.setattr( - auth, "_validate_profile", mock.Mock(side_effect=RuntimeError("service unavailable")) + assert "Using profile 'ws' (from --profile) · https://env.example (from DATABRICKS_HOST)" in ( + capsys.readouterr().err ) - monkeypatch.setattr(auth, "_is_interactive", lambda: True) - databricks_login = mock.Mock() - monkeypatch.setattr(auth.subprocess, "run", databricks_login) - result = CliRunner().invoke(auth.login, ["--profile", "prof"], obj=_Ctx()) - assert result.exit_code != 0 - assert "service unavailable" in result.output - assert auth.load_default_profile() is None - databricks_login.assert_not_called() +def test_describe_and_workspace_notice_report_switched_default( + outside_project, capsys, write_databrickscfg, fake_api_client, validated_profile +): + write_databrickscfg( + "[DEFAULT]\nhost = https://default.example\n" + "[__settings__]\ndefault_profile = switched\n" + "[switched]\nhost = https://switched.example\n" + ) + + assert _get_json(["auth", "describe"])["workspace_host"] == "https://switched.example" + validated_profile.assert_not_called() + fake_api_client.assert_not_called() + _ctx().client() + + assert "https://switched.example ([switched] profile)" in capsys.readouterr().err + + +def test_describe_reports_environment_host_even_with_a_flagged_profile( + outside_project, monkeypatch, write_databrickscfg, validated_profile +): + write_databrickscfg("[ws]\nhost = https://profile.example\n") + monkeypatch.setenv("DATABRICKS_HOST", "https://env.example") + monkeypatch.setenv("DATABRICKS_TOKEN", "test-token") + + assert _get_json(["auth", "describe", "-p", "ws"])["workspace_host"] == "https://env.example" + validated_profile.assert_not_called() + + +def test_login_notice_reports_environment_host( + outside_project, monkeypatch, write_databrickscfg, validated_profile +): + write_databrickscfg("[ws]\nhost = https://profile.example\n") + monkeypatch.setenv("DATABRICKS_HOST", "https://env.example") + + result = _invoke_profile(["login", "ws"]) + + assert result.exit_code == 0, result.output + assert "Logging in to profile 'ws' (from argument) — https://env.example" in result.output + + +def test_workspace_notice_is_silent_for_json(capsys, fake_api_client): + _ctx("ws", output="json").client() + + assert capsys.readouterr().err == "" -def test_login_does_not_launch_browser_when_noninteractive(tmp_path, monkeypatch): - monkeypatch.setenv("AGENTBRICKS_CONFIG_HOME", str(tmp_path)) - monkeypatch.setattr( - auth, "_AgentBricksApiClient", mock.Mock(side_effect=auth.AgentCliError("no credentials")) +def test_get_reports_profile_source_host_and_project_without_workspace_calls( + env_project, write_databrickscfg, validated_profile +): + write_databrickscfg("[proj]\nhost = https://proj.example\n") + + assert _get_json(["auth", "describe"]) == { + "profile": "proj", + "source": ".env", + "workspace_host": "https://proj.example", + "project": str(env_project), + } + validated_profile.assert_not_called() + + +def test_commands_in_a_project_subfolder_use_the_root_env( + env_project, monkeypatch, write_databrickscfg, validated_profile +): + write_databrickscfg("[proj]\nhost = https://proj.example\n") + write_project_metadata(env_project, framework="openai", template="agent-template") + subfolder = env_project / "agent" + subfolder.mkdir() + monkeypatch.chdir(subfolder) + + assert _get_json(["auth", "describe"]) == { + "profile": "proj", + "source": ".env", + "workspace_host": "https://proj.example", + "project": str(env_project), + } + + +def test_app_yaml_in_a_parent_does_not_make_a_subfolder_a_project(tmp_path, monkeypatch): + (tmp_path / "app.yaml").write_text("command: []\n") + _write_env(tmp_path, "proj") + subfolder = tmp_path / "agent" + subfolder.mkdir() + monkeypatch.chdir(subfolder) + + payload = _get_json(["auth", "describe"]) + + assert (payload["profile"], payload["source"], payload["project"]) == ( + None, + "Databricks SDK default", + None, ) - monkeypatch.setattr(auth, "_is_interactive", lambda: False) - databricks_login = mock.Mock() - monkeypatch.setattr(auth.subprocess, "run", databricks_login) - result = CliRunner().invoke(auth.login, ["--profile", "prof"], obj=_Ctx()) + +def test_switch_from_a_subfolder_rewrites_the_root_env(env_project, monkeypatch): + write_project_metadata(env_project, framework="openai", template="agent-template") + subfolder = env_project / "agent" + subfolder.mkdir() + monkeypatch.chdir(subfolder) + + result = _invoke_profile(["switch", "new"]) + + assert result.exit_code == 0, result.output + assert read_env(env_project / ".env")["DATABRICKS_CONFIG_PROFILE"] == "new" + assert not (subfolder / ".env").exists() + + +def test_set_rewrites_profile_line_and_keeps_other_lines(project): + (project / ".env").write_text("OTHER=1\nDATABRICKS_CONFIG_PROFILE=old\nLAST=2\n") + + result = _invoke_profile(["switch", "new"]) + + assert result.exit_code == 0, result.output + assert (project / ".env").read_text() == "OTHER=1\nDATABRICKS_CONFIG_PROFILE=new\nLAST=2\n" + + +def test_set_requires_a_project(outside_project): + result = _invoke_profile(["switch", "ml"]) assert result.exit_code != 0 - assert "interactive terminal" in result.output - assert auth.load_default_profile() is None - databricks_login.assert_not_called() + assert "isn't an Agent Bricks project" in _flat(result.output) + assert not (outside_project / ".env").exists() -def test_login_reports_when_databricks_cli_is_missing(tmp_path, monkeypatch): - monkeypatch.setenv("AGENTBRICKS_CONFIG_HOME", str(tmp_path)) - monkeypatch.setattr( - auth, "_AgentBricksApiClient", mock.Mock(side_effect=auth.AgentCliError("no credentials")) - ) - monkeypatch.setattr(auth, "_is_interactive", lambda: True) - monkeypatch.setattr(auth.subprocess, "run", mock.Mock(side_effect=FileNotFoundError)) +def test_set_signs_in_when_interactive(project, interactive_terminal, validated_profile): + result = _invoke_profile(["switch", "ml"], output="json") - result = CliRunner().invoke(auth.login, ["--profile", "prof"], obj=_Ctx()) + assert result.exit_code == 0, result.output + validated_profile.assert_called_once_with("ml") + assert json.loads(result.output)["signed_in_user"] == "me@example.com" + + +def test_set_noninteractive_suggests_login_without_signing_in(project, validated_profile): + result = _invoke_profile(["switch", "ml"]) + + assert result.exit_code == 0, result.output + assert "agentbricks auth login ml" in _flat(result.output) + validated_profile.assert_not_called() + + +@pytest.mark.parametrize( + ("args", "flag", "expected_profile", "expected_source"), + [ + (["login", "named"], "flagged", "named", "argument"), + (["login"], "flagged", "flagged", "--profile"), + (["login"], None, "proj", ".env"), + ], + ids=["argument-beats-flag", "global-flag-beats-env-file", "project-env-file"], +) +def test_login_target_precedence( + env_project, validated_profile, args, flag, expected_profile, expected_source +): + result = _invoke_profile(args, profile=flag, output="json") + + assert result.exit_code == 0, result.output + assert json.loads(result.output)["source"] == expected_source + validated_profile.assert_called_once_with(expected_profile) + + +def test_login_without_a_target_errors_with_hint_and_ignores_env_var( + outside_project, monkeypatch, validated_profile +): + monkeypatch.setenv("DATABRICKS_CONFIG_PROFILE", "ambient") + + out = _flat(_invoke_profile(["login"]).output) + + assert "No profile to log in to." in out + assert "agentbricks auth login " in out + validated_profile.assert_not_called() + + +def test_login_runs_databricks_cli_when_unauthenticated_and_interactive( + monkeypatch, interactive_terminal, databricks_login +): + client = mock.Mock(host="https://ws.example") + validate = mock.Mock(side_effect=[(None, Unauthenticated("expired")), ((client, "me"), None)]) + monkeypatch.setattr(workspace_auth, "_validate_bounded", validate) + + result = _invoke_profile(["login", "ml"]) + + assert result.exit_code == 0, result.output + assert databricks_login.call_args.args[0] == ["databricks", "auth", "login", "--profile", "ml"] + assert "Logged in to 'ml'" in _flat(result.output) + + +def test_login_noninteractive_unauthenticated_fails_without_browser( + fail_validation, databricks_login +): + fail_validation(Unauthenticated("expired")) + + result = _invoke_profile(["login", "ml"]) assert result.exit_code != 0 - assert "Could not configure Databricks authentication" in result.output - assert "not found" in result.output - assert auth.load_default_profile() is None + assert "interactive terminal" in _flat(result.output) + databricks_login.assert_not_called() + +def test_login_leaves_project_env_file_alone(env_project, validated_profile): + result = _invoke_profile(["login", "other"]) -def test_logout_clears_saved_profile(tmp_path, monkeypatch): - monkeypatch.setenv("AGENTBRICKS_CONFIG_HOME", str(tmp_path)) - (tmp_path / "config.json").write_text(json.dumps({"profile": "x"})) - result = CliRunner().invoke(auth.logout, [], obj=_Ctx()) assert result.exit_code == 0, result.output - assert auth.load_default_profile() is None + assert read_env(env_project / ".env").get("DATABRICKS_CONFIG_PROFILE") == "proj" diff --git a/integrations/agentbricks/tests/unit_tests/cli_test.py b/integrations/agentbricks/tests/unit_tests/cli_test.py index 973e400cb..515bb8505 100644 --- a/integrations/agentbricks/tests/unit_tests/cli_test.py +++ b/integrations/agentbricks/tests/unit_tests/cli_test.py @@ -12,6 +12,8 @@ def _command_paths(group: click.Group, prefix: tuple[str, ...] = ()): for name, command in group.commands.items(): + if command.hidden: + continue path = (*prefix, name) yield path if isinstance(command, click.Group): @@ -30,8 +32,7 @@ def test_sessions_verbs_are_flat_no_redundant_subgroup(): def test_root_registers_supported_commands(): names = set(cli.agentbricks.commands) assert { - "login", - "logout", + "auth", "init", "doctor", "dev", @@ -43,6 +44,8 @@ def test_root_registers_supported_commands(): "endpoint", "tools", } <= names + # Removed commands stay only as hidden pointers to their replacements. + assert all(cli.agentbricks.commands[name].hidden for name in ("login", "logout")) assert "eval" not in names assert "mcp" not in names assert "pipeline" not in names @@ -169,8 +172,8 @@ def test_help_examples_recommend_the_default_happy_path(): runner = CliRunner() expected_examples = { (): ( - "agentbricks login --profile my-workspace", - "agentbricks init my-agent", + "agentbricks auth login ", + "agentbricks init my-agent --profile ", "cd my-agent", "agentbricks dev", "agentbricks deploy", @@ -281,7 +284,7 @@ def test_unknown_command_suggests_close_match(): def test_doubled_invocation_is_named_directly(): - result = CliRunner().invoke(cli.agentbricks, ["agentbricks", "login"]) + result = CliRunner().invoke(cli.agentbricks, ["agentbricks", "auth"]) assert result.exit_code != 0 assert "you typed `agentbricks` twice" in result.output @@ -388,12 +391,12 @@ def test_root_help_shows_numbered_getting_started_path(): assert result.exit_code == 0, result.output out = result.output assert "GETTING STARTED" in out - # Read just the numbered block (login/init/deploy also appear in the prose description above). + # Read just the numbered block (auth login/init/deploy also appear in the prose description above). block = out[out.index("GETTING STARTED") :] block = block[: block.index("Not authenticated")] numbered = [ln.strip() for ln in block.splitlines() if ln.strip()[:1].isdigit()] - # The path is numbered and ordered: login (1) → init (2) → cd (3) → dev (4) → deploy (5). - assert numbered[0].startswith("1") and "agentbricks login --profile" in numbered[0] + # The path is numbered and ordered: sign in (1) → init (2) → cd (3) → dev (4) → deploy (5). + assert numbered[0].startswith("1") and "agentbricks auth login" in numbered[0] assert numbered[1].startswith("2") and "agentbricks init my-agent" in numbered[1] assert numbered[4].startswith("5") and "agentbricks deploy " in numbered[4] @@ -413,9 +416,9 @@ def test_help_dims_headings_and_descriptions_not_names(): setup = next(ln for ln in lines if click.unstyle(ln).strip() == "SETUP") assert setup.startswith(dim) # A command row dims the description but not the name. - login = next(ln for ln in lines if click.unstyle(ln).strip().startswith("login")) - assert dim in login # the description is dimmed - assert not login.startswith(dim) # the `login` name is at full intensity + auth_row = next(ln for ln in lines if click.unstyle(ln).strip().startswith("auth ")) + assert dim in auth_row # the description is dimmed + assert not auth_row.startswith(dim) # the `auth` name is at full intensity def test_epilog_headings_align_flush_left_with_sections(): @@ -434,7 +437,7 @@ def indent(line: str) -> int: line = next(ln for ln in lines if ln.strip() == name) assert indent(line) == 0, (name, line) # A numbered getting-started row aligns with an OPTIONS/command row at column 2. - row = next(ln for ln in lines if ln.strip().startswith("1 agentbricks login")) + row = next(ln for ln in lines if ln.strip().startswith("1 agentbricks auth login")) assert indent(row) == 2, row @@ -467,6 +470,6 @@ def test_root_help_groups_commands_by_intent(): assert out.index("SETUP") < out.index("DEVELOP") < out.index("SHIP") # Commands land under their section. setup_to_develop = out[out.index("SETUP") : out.index("DEVELOP")] - assert "login" in setup_to_develop and "init" in setup_to_develop + assert "auth" in setup_to_develop and "init" in setup_to_develop ship_onward = out[out.index("SHIP") :] assert "deploy" in ship_onward diff --git a/integrations/agentbricks/tests/unit_tests/client_test.py b/integrations/agentbricks/tests/unit_tests/client_test.py index 6a90afe4f..f720686e1 100644 --- a/integrations/agentbricks/tests/unit_tests/client_test.py +++ b/integrations/agentbricks/tests/unit_tests/client_test.py @@ -132,6 +132,16 @@ def _client(workspace_client): return _AgentBricksApiClient(profile="p"), inst.api_client.do +@pytest.fixture +def sleep(monkeypatch): + retry_sleep = mock.Mock(name="agentkit_retry_sleep") + # Patching time.sleep would also capture SDK retries on background telemetry threads. + monkeypatch.setattr( + "databricks_agentkit._api_client.time", types.SimpleNamespace(sleep=retry_sleep) + ) + return retry_sleep + + @pytest.mark.parametrize( "client_class", [AgentKitClient], @@ -785,7 +795,7 @@ def test_auth_error_hint_explains_profile_selection_and_login(_workspace_client) hint = exc_info.value.hint assert hint is not None assert "`agentbricks --profile `" in hint - assert "`agentbricks login --profile `" in hint + assert "`agentbricks auth login `" in hint assert "`databricks auth login --profile `" not in hint @@ -855,7 +865,6 @@ class _TransientError(RuntimeError): error_code = "CANCELLED" -@mock.patch("databricks_agentkit._api_client.time.sleep") @mock.patch("databricks.sdk.WorkspaceClient") def test_delete_runtime_store_retries_transient_errors(workspace_client, sleep): client, do = _client(workspace_client) @@ -866,7 +875,6 @@ def test_delete_runtime_store_retries_transient_errors(workspace_client, sleep): sleep.assert_called_once() -@mock.patch("databricks_agentkit._api_client.time.sleep") @mock.patch("databricks.sdk.WorkspaceClient") def test_do_retries_transient_error_then_succeeds(workspace_client, sleep): client, do = _client(workspace_client) @@ -879,7 +887,6 @@ def test_do_retries_transient_error_then_succeeds(workspace_client, sleep): sleep.assert_called_once() -@mock.patch("databricks_agentkit._api_client.time.sleep") @mock.patch("databricks.sdk.WorkspaceClient") @pytest.mark.parametrize( ("method_name", "args"), @@ -901,7 +908,6 @@ def test_create_store_does_not_retry_transient_error_by_default( sleep.assert_not_called() -@mock.patch("databricks_agentkit._api_client.time.sleep") @mock.patch("databricks.sdk.WorkspaceClient") def test_do_stops_after_max_attempts(workspace_client, sleep): client, do = _client(workspace_client) @@ -915,7 +921,6 @@ def test_do_stops_after_max_attempts(workspace_client, sleep): assert "None" not in exc_info.value.message -@mock.patch("databricks_agentkit._api_client.time.sleep") @mock.patch("databricks.sdk.WorkspaceClient") def test_do_does_not_retry_non_transient_error(workspace_client, sleep): client, do = _client(workspace_client) @@ -932,7 +937,6 @@ class NotFoundError(RuntimeError): sleep.assert_not_called() -@mock.patch("databricks_agentkit._api_client.time.sleep") @mock.patch("databricks.sdk.WorkspaceClient") def test_do_does_not_retry_transient_error_for_pop(workspace_client, sleep): client, do = _client(workspace_client) diff --git a/integrations/agentbricks/tests/unit_tests/conftest.py b/integrations/agentbricks/tests/unit_tests/conftest.py index b501cadc7..ca44bee3a 100644 --- a/integrations/agentbricks/tests/unit_tests/conftest.py +++ b/integrations/agentbricks/tests/unit_tests/conftest.py @@ -2,11 +2,86 @@ Service tests inject their collaborators directly; a directory-wide patch of the old deploy module would mask the new boundaries and fails because that function no longer exists. + +Every test runs against an empty Databricks config file and none of the ambient credential or CI +variables, so profile resolution never depends on the developer's machine or the CI runner. """ +from __future__ import annotations + +import pathlib +import textwrap +from collections.abc import Callable +from unittest import mock + import pytest from databricks_agentbricks.agent_hint import AGENT_ENV_MARKERS, AGENT_ENV_VALUES +from databricks_agentbricks.cli import auth +from databricks_agentbricks.clients import databricks_cli +from databricks_agentbricks.services import workspace_auth + +_AMBIENT_ENV = ("DATABRICKS_CONFIG_PROFILE", "DATABRICKS_HOST", "DATABRICKS_TOKEN", "CI") + + +@pytest.fixture(autouse=True) +def databrickscfg_path( + monkeypatch: pytest.MonkeyPatch, tmp_path_factory: pytest.TempPathFactory +) -> pathlib.Path: + """Isolated (initially absent) `DATABRICKS_CONFIG_FILE` and `HOME`; the terminal is non-interactive. + + HOME points at a tmp dir so nothing reads the developer's real config or token cache. + """ + for name in _AMBIENT_ENV: + monkeypatch.delenv(name, raising=False) + path = tmp_path_factory.mktemp("databricks-config") / ".databrickscfg" + monkeypatch.setenv("DATABRICKS_CONFIG_FILE", str(path)) + monkeypatch.setenv("HOME", str(tmp_path_factory.mktemp("agentbricks-home"))) + monkeypatch.setattr(auth, "_is_interactive", lambda: False) + return path + + +@pytest.fixture +def write_databrickscfg(databrickscfg_path: pathlib.Path) -> Callable[[str], None]: + """Write the isolated Databricks config file from an INI snippet.""" + + def _write(text: str) -> None: + databrickscfg_path.write_text(textwrap.dedent(text)) + + return _write + + +@pytest.fixture +def validated_profile(monkeypatch: pytest.MonkeyPatch) -> mock.Mock: + """Make the bounded SDK validation (the workspace boundary) succeed for any profile.""" + client = mock.Mock(host="https://ws.example", current_user="me@example.com") + validate = mock.Mock(return_value=((client, "me@example.com"), None)) + monkeypatch.setattr(workspace_auth, "_validate_bounded", validate) + return validate + + +@pytest.fixture +def interactive_terminal(monkeypatch: pytest.MonkeyPatch) -> None: + """Someone is at the terminal, so a browser login could be completed.""" + monkeypatch.setattr(auth, "_is_interactive", lambda: True) + + +@pytest.fixture +def fail_validation(monkeypatch: pytest.MonkeyPatch) -> Callable[[Exception], None]: + """Make the bounded SDK validation fail with the given error for any profile.""" + + def _fail(error: Exception) -> None: + monkeypatch.setattr(workspace_auth, "_validate_bounded", lambda profile: (None, error)) + + return _fail + + +@pytest.fixture +def databricks_login(monkeypatch: pytest.MonkeyPatch) -> mock.Mock: + """Stub the `databricks auth login` subprocess; it succeeds unless a test reconfigures it.""" + login = mock.Mock(return_value=mock.Mock(returncode=0)) + monkeypatch.setattr(databricks_cli.subprocess, "run", login) + return login @pytest.fixture(autouse=True) diff --git a/integrations/agentbricks/tests/unit_tests/deploy_bugfix_test.py b/integrations/agentbricks/tests/unit_tests/deploy_bugfix_test.py index 53dd23e71..564b042b8 100644 --- a/integrations/agentbricks/tests/unit_tests/deploy_bugfix_test.py +++ b/integrations/agentbricks/tests/unit_tests/deploy_bugfix_test.py @@ -20,6 +20,7 @@ class _Ctx: profile = "prof" output = "text" api_client_provider = types.SimpleNamespace(get=lambda: types.SimpleNamespace()) + announce_workspace = staticmethod(lambda: None) # --- ML-69259 / 69247: deployment-name validation --------------------------- @@ -75,6 +76,9 @@ class _JsonCtx: profile = "prof" output = "json" api_client_provider = types.SimpleNamespace(get=lambda: types.SimpleNamespace()) + announce_workspace = staticmethod(lambda: None) + + announce_workspace = staticmethod(lambda: None) result = CliRunner().invoke(deploy_mod.deployments_list, [], obj=_JsonCtx()) assert result.exit_code == 0, result.output diff --git a/integrations/agentbricks/tests/unit_tests/deploy_test.py b/integrations/agentbricks/tests/unit_tests/deploy_test.py index 6d43d1b70..1c854382a 100644 --- a/integrations/agentbricks/tests/unit_tests/deploy_test.py +++ b/integrations/agentbricks/tests/unit_tests/deploy_test.py @@ -19,6 +19,7 @@ from click.testing import CliRunner from databricks_agentbricks.cli import deploy as deploy_mod +from databricks_agentbricks.cli.app import CliContext from databricks_agentbricks.clients.api_client_provider import ApiClientProvider from databricks_agentbricks.clients.apps_client import AppsClient from databricks_agentbricks.clients.conversation_store_client import ( @@ -34,8 +35,11 @@ from databricks_agentbricks.presentation.endpoint import agent_invoke_command from databricks_agentbricks.projects.agent_project import AgentProject from databricks_agentbricks.projects.app_manifest import AppManifest +from databricks_agentbricks.projects.profiles import ProfileInfo from databricks_agentbricks.projects.resolver import ProjectResolver from databricks_agentbricks.services.deploy_service import ( + DeleteResult, + DeleteTarget, DeployRequest, DeployResult, ToolAccessSummary, @@ -55,13 +59,17 @@ class _FakeApiClient: current_user = "me@example.com" -class _Ctx: - profile = "prof" - output = "text" +class _Ctx(CliContext): + def __init__(self, *, output: str = "text", profile: str | None = "prof"): + super().__init__(profile, output) + self.api_client_provider = types.SimpleNamespace( + get=lambda: _FakeApiClient(), use=lambda client: self.announce_workspace(client.host) + ) + - def __init__(self, *, output: str = "text"): - self.output = output - self.api_client_provider = types.SimpleNamespace(get=lambda: _FakeApiClient()) +@pytest.fixture(autouse=True) +def _authenticated_profile(validated_profile, monkeypatch: pytest.MonkeyPatch, tmp_path): + monkeypatch.chdir(tmp_path) def _result(*, deployment: str = "agent-bricks-demo") -> DeployResult: @@ -232,11 +240,10 @@ def _no_real_sleep(monkeypatch: pytest.MonkeyPatch) -> None: def _real_deploy_context(api: _DeployApi, *, output: str = "text"): - return types.SimpleNamespace( - profile="prof", - output=output, - api_client_provider=types.SimpleNamespace(get=lambda: api), - ) + ctx = CliContext("prof", output) + # The preflight's validated client is discarded so every call reaches the fake `api`. + ctx.api_client_provider = types.SimpleNamespace(get=lambda: api, use=lambda client: None) + return ctx def test_manifest_env_scaffolds_when_missing(tmp_path: pathlib.Path): @@ -392,6 +399,40 @@ def test_deploy_help_exposes_instance_count_and_routing_key(): assert "X-Routing-Key" in result.output +@pytest.mark.parametrize("managed", [True, False], ids=["managed", "cli-only"]) +def test_delete_announces_the_workspace_used_by_its_client(write_databrickscfg, managed): + write_databrickscfg("[prof]\nhost = https://profile.example\n") + ctx = CliContext("prof", "text") + service = mock.Mock() + service.deletes_runtime_store_data.return_value = managed + target = DeleteTarget(DeploymentName("demo"), source_path=None, service_principal=None) + service.delete_target.return_value = target + deleted = DeleteResult( + "demo", None, "not_recorded", None, still_deleting=False, wait_error=None + ) + + def delete(_target, *, delete_source): + if managed: + ctx.client() + return deleted + + service.delete.side_effect = delete + with ( + mock.patch.object(deploy_mod, "build_deploy_service", return_value=service), + mock.patch( + "databricks_agentbricks.clients.api_client_provider._AgentBricksApiClient", + return_value=_FakeApiClient(), + ) as client, + ): + result = CliRunner().invoke(deploy_mod.deployments, ["delete", "demo", "--yes"], obj=ctx) + + assert result.exit_code == 0, result.output + host = "https://workspace.example" if managed else "https://profile.example" + assert result.output.count(f"Using profile 'prof' (from --profile) · {host}") == 1 + assert client.call_count == int(managed) + service.delete.assert_called_once_with(target, delete_source=False) + + def test_deploy_rejects_instance_count_above_platform_limit(): result = CliRunner().invoke(deploy_mod.deploy, ["demo", "--instances", "6"], obj=_Ctx()) @@ -509,6 +550,33 @@ def test_deploy_cli_rejects_origins_with_clear(): assert "not both" in result.output +@pytest.mark.parametrize( + ("flag_profile", "expected", "source"), [(None, "proj", ".env"), ("flag", "flag", "--profile")] +) +def test_deploy_resolves_profile_from_flag_then_source_env( + tmp_path: pathlib.Path, validated_profile, write_databrickscfg, flag_profile, expected, source +): + write_databrickscfg( + "[proj]\nhost = https://proj.example\n[flag]\nhost = https://flag.example\n" + ) + (tmp_path / "app.yaml").write_text("command: []\n") + (tmp_path / ".env").write_text("DATABRICKS_CONFIG_PROFILE=proj\n") + ctx = _Ctx(profile=flag_profile) + + service = mock.Mock() + service.deploy.return_value = _result() + with mock.patch.object(deploy_mod, "build_deploy_service", return_value=service) as build: + result = CliRunner().invoke(deploy_mod.deploy, ["demo", "--source", str(tmp_path)], obj=ctx) + + assert result.exit_code == 0, result.output + assert ctx.profile_info == ProfileInfo(expected, source) + assert build.call_args.args[0] is ctx + validated_profile.assert_called_once_with(expected) + out = " ".join(result.output.split()) + assert f"Using profile '{expected}' (from {source}) · https://ws.example" in out + assert f"https://{expected}.example" not in out + + def test_deploy_json_preserves_additive_tool_access_contract(tmp_path: pathlib.Path): service = mock.Mock() service.deploy.return_value = replace( @@ -537,6 +605,7 @@ def test_deploy_json_preserves_additive_tool_access_contract(tmp_path: pathlib.P def test_build_deploy_service_composes_current_clients_without_opening_api_client(): ctx = _Ctx() + ctx.select_profile() service = deploy_mod.build_deploy_service(ctx) diff --git a/integrations/agentbricks/tests/unit_tests/deploy_wait_test.py b/integrations/agentbricks/tests/unit_tests/deploy_wait_test.py index 53b313a8b..f67eb2021 100644 --- a/integrations/agentbricks/tests/unit_tests/deploy_wait_test.py +++ b/integrations/agentbricks/tests/unit_tests/deploy_wait_test.py @@ -26,6 +26,11 @@ _DEAD_PID = 2**22 + 1 +@pytest.fixture(autouse=True) +def _authenticated_profile(validated_profile): + """Deploy checks the profile's auth before any work; these tests exercise what follows.""" + + @pytest.fixture def project(tmp_path: pathlib.Path) -> pathlib.Path: AgentProject.create(tmp_path, framework="langgraph", server="custom").write() @@ -254,7 +259,7 @@ def test_no_wait_passes_through_non_default_options( events: list[tuple] = [] monkeypatch.setattr(deploy_mod, "_databricks", _DeployRunner(events)) ctx = _real_deploy_context(_DeployApi(events)) - ctx.profile = None + ctx.global_profile = None result = CliRunner().invoke( deploy_mod.deploy, diff --git a/integrations/agentbricks/tests/unit_tests/dev_service_test.py b/integrations/agentbricks/tests/unit_tests/dev_service_test.py index 86b721ad0..7715b55dc 100644 --- a/integrations/agentbricks/tests/unit_tests/dev_service_test.py +++ b/integrations/agentbricks/tests/unit_tests/dev_service_test.py @@ -162,6 +162,24 @@ def test_prepare_builds_local_plan_and_run_uses_injected_apps_client( ) +def test_prepare_puts_local_env_in_the_local_manifest( + service_fixture: SimpleNamespace, tmp_path: pathlib.Path +) -> None: + request = DevRequest( + source=str(tmp_path), + prepare_environment=None, + app_port=None, + local_env={"DATABRICKS_CONFIG_PROFILE": "ml"}, + ) + + with service_fixture.service.prepare(request) as plan: + env = _manifest_env(plan.entry_point) + + assert env["DATABRICKS_CONFIG_PROFILE"] == "ml" + assert env["MLFLOW_TRACKING_URI"] == "http://127.0.0.1:5599" + assert "DATABRICKS_CONFIG_PROFILE" not in (tmp_path / "app.yaml").read_text() + + def test_prepare_honors_explicit_environment_flag_and_default_port( service_fixture: SimpleNamespace, tmp_path: pathlib.Path ) -> None: diff --git a/integrations/agentbricks/tests/unit_tests/dev_test.py b/integrations/agentbricks/tests/unit_tests/dev_test.py index 5977d2b40..551f52153 100644 --- a/integrations/agentbricks/tests/unit_tests/dev_test.py +++ b/integrations/agentbricks/tests/unit_tests/dev_test.py @@ -12,6 +12,7 @@ from click.testing import CliRunner from databricks_agentbricks.cli import dev as dev_mod +from databricks_agentbricks.cli.app import CliContext from databricks_agentbricks.clients.local_tracing_client import LocalTracingStart from databricks_agentbricks.errors import AgentCliError from databricks_agentbricks.presentation.endpoint import agent_invoke_command @@ -37,10 +38,9 @@ def _write_agent_manifest( (source / "agent.toml").write_text(body) -class _Ctx: +class _Ctx(CliContext): def __init__(self, output: str = "text", profile=None): - self.output = output - self.profile = profile + super().__init__(profile, output) self.api_client_provider = types.SimpleNamespace( get=lambda: mock.Mock( current_user="me@example.com", host="https://my-workspace.databricks.com" @@ -63,6 +63,11 @@ def _panel_text(output: str) -> str: return " ".join(line.strip("│ ") for line in output.splitlines() if line.startswith("│")) +@pytest.fixture(autouse=True) +def _authenticated_profile(validated_profile): + """Preflight validation is the workspace boundary; these tests are not about workspace_auth.""" + + @pytest.fixture(autouse=True) def _stub_local_tracing(monkeypatch): """`agentbricks dev` starts a local MLflow tracking server for Agent Bricks projects; @@ -393,12 +398,12 @@ def test_dev_runs_offline_when_client_unavailable(tmp_path: pathlib.Path): (tmp_path / "app.yaml").write_text("command: []\n") (tmp_path / ".venv").mkdir() - class _OfflineCtx: - output = "text" - profile = None - api_client_provider = types.SimpleNamespace( - get=lambda: (_ for _ in ()).throw(AgentCliError("no databricks auth configured")) - ) + class _OfflineCtx(CliContext): + def __init__(self): + super().__init__(None, "text") + self.api_client_provider = types.SimpleNamespace( + get=lambda: (_ for _ in ()).throw(AgentCliError("no databricks auth configured")) + ) def client(self): raise AgentCliError("no databricks auth configured") @@ -693,6 +698,91 @@ def test_dev_hides_active_venv_and_port_from_run_local(tmp_path: pathlib.Path, m assert env["AGENTBRICKS_DEV_TEST_MARKER"] == "kept" +def _write_project(source: pathlib.Path, env_file: str) -> None: + (source / "app.yaml").write_text("command: []\n") + (source / ".venv").mkdir() + (source / ".env").write_text(env_file) + + +def _run_dev(source: pathlib.Path, monkeypatch, env_file: str, profile=None): + """Run `dev` on a fresh project; returns (result, run-local profiles, local-manifest env).""" + _write_project(source, env_file) + profiles: list = [] + manifest_env: dict[str, str] = {} + + def _fake_databricks(cmd, profile, **kwargs): + profiles.append(profile) + manifest = yaml.safe_load( + (pathlib.Path(kwargs["cwd"]) / "app.agentbricksdev.yaml").read_text() + ) + manifest_env.update({e["name"]: e["value"] for e in manifest["env"]}) + return types.SimpleNamespace(returncode=0) + + monkeypatch.setattr(dev_mod, "_databricks", _fake_databricks) + result = CliRunner().invoke(dev_mod.dev, ["--source", str(source)], obj=_Ctx(profile=profile)) + return result, profiles, manifest_env + + +@pytest.mark.parametrize( + ("flag_profile", "expected", "note"), + [ + (None, "proj", "Using profile 'proj' (from .env) · https://ws.example"), + ("flag", "flag", "will run with profile 'flag' (from --profile), not .env's 'proj'"), + ], +) +def test_dev_runs_with_resolved_profile( + tmp_path, monkeypatch, validated_profile, write_databrickscfg, flag_profile, expected, note +): + write_databrickscfg("[proj]\nhost = https://proj.example\n") + result, profiles, manifest_env = _run_dev( + tmp_path, monkeypatch, "DATABRICKS_CONFIG_PROFILE=proj\n", profile=flag_profile + ) + + assert result.exit_code == 0, result.output + assert profiles == [expected] + assert manifest_env["DATABRICKS_CONFIG_PROFILE"] == expected + validated_profile.assert_called_once_with(expected) + assert note in " ".join(result.output.split()) + + +# A client id without its secret can't authenticate on its own, so the profile still applies. +@pytest.mark.parametrize( + ("credentials", "uses_profile"), + [ + ("DATABRICKS_TOKEN=dapi123\n", False), + ("DATABRICKS_CLIENT_ID=sp\nDATABRICKS_CLIENT_SECRET=secret\n", False), + ("DATABRICKS_CLIENT_ID=sp\n", True), + ], +) +def test_dev_env_credentials_replace_the_profile_only_when_complete( + tmp_path, monkeypatch, validated_profile, credentials, uses_profile +): + env_file = f"DATABRICKS_CONFIG_PROFILE=proj\n{credentials}" + result, _, manifest_env = _run_dev(tmp_path, monkeypatch, env_file) + + assert result.exit_code == 0, result.output + assert manifest_env.get("DATABRICKS_CONFIG_PROFILE") == ("proj" if uses_profile else None) + assert validated_profile.call_args_list == ([mock.call("proj")] if uses_profile else []) + + +def test_dev_unauthenticated_profile_fails_before_run_local( + tmp_path: pathlib.Path, monkeypatch, fail_validation +): + _write_project(tmp_path, "DATABRICKS_CONFIG_PROFILE=proj\n") + fail_validation(RuntimeError("expired")) + run_local = mock.Mock() + monkeypatch.setattr(dev_mod, "_databricks", run_local) + + result = CliRunner().invoke(dev_mod.dev, ["--source", str(tmp_path)], obj=_Ctx()) + + out = " ".join(result.output.split()) + assert result.exit_code != 0 + assert "Using profile 'proj' (from .env)" in out + assert "'proj' (from .env) isn't authenticated" in out + assert "agentbricks auth login proj" in out + run_local.assert_not_called() + + def test_dev_busy_port_renders_error_and_hint(tmp_path: pathlib.Path, monkeypatch): (tmp_path / "app.yaml").write_text("command: []\n") monkeypatch.setattr(dev_mod, "port_in_use", lambda port: port == 9100) diff --git a/integrations/agentbricks/tests/unit_tests/endpoint_test.py b/integrations/agentbricks/tests/unit_tests/endpoint_test.py index 0ff598b73..c8eb9f819 100644 --- a/integrations/agentbricks/tests/unit_tests/endpoint_test.py +++ b/integrations/agentbricks/tests/unit_tests/endpoint_test.py @@ -21,6 +21,9 @@ class _Ctx: profile = "profile" output = "json" + def announce_workspace(self) -> None: + pass + def _response( body, diff --git a/integrations/agentbricks/tests/unit_tests/env_file_test.py b/integrations/agentbricks/tests/unit_tests/env_file_test.py new file mode 100644 index 000000000..340334541 --- /dev/null +++ b/integrations/agentbricks/tests/unit_tests/env_file_test.py @@ -0,0 +1,76 @@ +"""Unit tests for the project `.env` reader and DATABRICKS_CONFIG_PROFILE helpers.""" + +from __future__ import annotations + +import pathlib + +from databricks_agentbricks.projects.env_file import read_env, update_env_profile, write_env + + +def test_read_env_parses_pairs_and_export(tmp_path: pathlib.Path): + (tmp_path / ".env").write_text("# comment\n\nA=1\nexport B=2\nC='three'\n") + + assert read_env(tmp_path / ".env") == {"A": "1", "B": "2", "C": "three"} + + +def test_write_env_seeds_from_example(tmp_path: pathlib.Path): + (tmp_path / ".env.example").write_text("DATABRICKS_CONFIG_PROFILE=DEFAULT\nOTHER=1\n") + + assert write_env(tmp_path, "ml") is True + + assert (tmp_path / ".env").read_text() == "DATABRICKS_CONFIG_PROFILE=ml\nOTHER=1\n" + + +def test_write_env_never_clobbers_existing(tmp_path: pathlib.Path): + (tmp_path / ".env").write_text("DATABRICKS_CONFIG_PROFILE=keepme\n") + + assert write_env(tmp_path, "ml") is False + + assert (tmp_path / ".env").read_text() == "DATABRICKS_CONFIG_PROFILE=keepme\n" + + +def test_update_env_profile_replaces_in_place_keeping_other_lines(tmp_path: pathlib.Path): + (tmp_path / ".env").write_text("A=1\nDATABRICKS_CONFIG_PROFILE=old\n# note\nB=2\n") + + assert update_env_profile(tmp_path, "new") is True + + assert (tmp_path / ".env").read_text() == "A=1\nDATABRICKS_CONFIG_PROFILE=new\n# note\nB=2\n" + + +# The writer matches what the reader honors (`export`, spacing, indentation), and drops later +# duplicates so the reader can't pick up a stale value. +def test_update_env_profile_rewrites_export_and_drops_duplicates(tmp_path: pathlib.Path): + (tmp_path / ".env").write_text( + "export DATABRICKS_CONFIG_PROFILE=old\nOTHER=1\n DATABRICKS_CONFIG_PROFILE = stale\n" + ) + + assert update_env_profile(tmp_path, "new") is True + + assert (tmp_path / ".env").read_text() == "DATABRICKS_CONFIG_PROFILE=new\nOTHER=1\n" + + +def test_update_env_profile_adds_missing_line(tmp_path: pathlib.Path): + (tmp_path / ".env").write_text("A=1\n") + + assert update_env_profile(tmp_path, "new") is True + + assert (tmp_path / ".env").read_text() == "DATABRICKS_CONFIG_PROFILE=new\nA=1\n" + + +def test_update_env_profile_rewrites_interpolated_value_to_literal(tmp_path: pathlib.Path): + (tmp_path / ".env").write_text( + "DATABRICKS_CONFIG_PROFILE=${CHOSEN_PROFILE}\nCHOSEN_PROFILE=first\n" + ) + + assert update_env_profile(tmp_path, "first") is True + + assert "DATABRICKS_CONFIG_PROFILE=first\n" in (tmp_path / ".env").read_text() + + +def test_update_env_profile_same_value_leaves_file_untouched(tmp_path: pathlib.Path): + body = b"A=1\nDATABRICKS_CONFIG_PROFILE=ml\n" + (tmp_path / ".env").write_bytes(body) + + assert update_env_profile(tmp_path, "ml") is False + + assert (tmp_path / ".env").read_bytes() == body diff --git a/integrations/agentbricks/tests/unit_tests/init_test.py b/integrations/agentbricks/tests/unit_tests/init_test.py index 15e480394..97c39c7f2 100644 --- a/integrations/agentbricks/tests/unit_tests/init_test.py +++ b/integrations/agentbricks/tests/unit_tests/init_test.py @@ -15,7 +15,9 @@ import tomli from click.testing import CliRunner +from databricks_agentbricks.cli import app as cli_app from databricks_agentbricks.cli import init as init_mod +from databricks_agentbricks.cli.app import CliContext from databricks_agentbricks.errors import AgentCliError from databricks_agentbricks.projects.agent_project import AgentProject from databricks_agentbricks.projects.resolver import ProjectResolver @@ -23,12 +25,11 @@ from databricks_agentbricks.services.deployment.names import DeploymentName, _prefixed_name -class _Ctx: - """Stand-in for CliContext: init reads .output and .profile.""" +class _Ctx(CliContext): + """CliContext with the (output, profile) signature init tests pass.""" def __init__(self, output: str = "text", profile=None): - self.output = output - self.profile = profile + super().__init__(profile, output) def _default_store_token(manifest: dict, slug: str = "proj") -> str: @@ -281,56 +282,86 @@ def test_init_json_keeps_lowercase_framework_and_server( assert payload["server"] == server.value -def test_write_env_seeds_profile_from_example(tmp_path: pathlib.Path): - (tmp_path / ".env.example").write_text( - "DATABRICKS_CONFIG_PROFILE=DEFAULT\n# MLFLOW_EXPERIMENT_ID=\n" - ) - assert init_mod._write_env(tmp_path, "ml") is True - body = (tmp_path / ".env").read_text() - assert "DATABRICKS_CONFIG_PROFILE=ml" in body - assert "# MLFLOW_EXPERIMENT_ID=" in body # rest of the example preserved +_ENV_EXAMPLE = {".env.example": "DATABRICKS_CONFIG_PROFILE=DEFAULT\n# MLFLOW_EXPERIMENT_ID=\n"} -def test_write_env_never_clobbers_existing(tmp_path: pathlib.Path): - (tmp_path / ".env").write_text("DATABRICKS_CONFIG_PROFILE=keepme\n") - assert init_mod._write_env(tmp_path, "ml") is False - assert "keepme" in (tmp_path / ".env").read_text() +def _init(dest: pathlib.Path, *args: str, output: str = "text"): + copy = _copy_writing(_ENV_EXAMPLE) + with mock.patch.object(init_mod, "_copy_packaged_template", side_effect=copy): + return CliRunner().invoke(init_mod.init, [*args, str(dest)], obj=_Ctx(output=output)) -def test_init_profile_flag_writes_env(tmp_path: pathlib.Path): +@pytest.mark.parametrize( + ("args", "env", "expected"), + [ + (["-p", "ml"], {}, "ml"), + ([], {"DATABRICKS_CONFIG_PROFILE": "from-env-var"}, "from-env-var"), + ], +) +def test_init_writes_env_with_resolved_profile( + tmp_path: pathlib.Path, monkeypatch, args, env, expected +): + for name, value in env.items(): + monkeypatch.setenv(name, value) dest = tmp_path / "proj" - with mock.patch.object( - init_mod, - "_copy_packaged_template", - side_effect=_copy_writing({".env.example": "DATABRICKS_CONFIG_PROFILE=DEFAULT\n"}), - ): - result = CliRunner().invoke(init_mod.init, ["--profile", "ml", str(dest)], obj=_Ctx()) + result = _init(dest, *args) assert result.exit_code == 0, result.output - assert "DATABRICKS_CONFIG_PROFILE=ml" in (dest / ".env").read_text() + body = (dest / ".env").read_text() + assert f"DATABRICKS_CONFIG_PROFILE={expected}\n" in body + assert "# MLFLOW_EXPERIMENT_ID=" in body # rest of the example preserved -def test_init_uses_ctx_profile_when_flag_absent(tmp_path: pathlib.Path): - dest = tmp_path / "proj" - with mock.patch.object( - init_mod, - "_copy_packaged_template", - side_effect=_copy_writing({".env.example": "DATABRICKS_CONFIG_PROFILE=DEFAULT\n"}), - ): - result = CliRunner().invoke(init_mod.init, [str(dest)], obj=_Ctx(profile="from-login")) +def test_init_inside_another_project_does_not_inherit_its_env_profile( + tmp_path: pathlib.Path, monkeypatch +): + # The root group folds the cwd project's `.env` into the context, but the new project + # must not inherit that workspace. + monkeypatch.chdir(tmp_path) + (tmp_path / "app.yaml").write_text("command: []\n") + (tmp_path / ".env").write_text("DATABRICKS_CONFIG_PROFILE=cwd-project\n") + dest = tmp_path / "new-proj" + + result = CliRunner().invoke(cli_app.agentbricks, ["init", str(dest)]) + assert result.exit_code == 0, result.output - assert "DATABRICKS_CONFIG_PROFILE=from-login" in (dest / ".env").read_text() + body = (dest / ".env").read_text() + assert "cwd-project" not in body + assert "DATABRICKS_CONFIG_PROFILE=DEFAULT" in body -def test_init_no_profile_writes_no_env(tmp_path: pathlib.Path): +def test_init_with_databricks_host_and_no_profile_writes_no_env( + tmp_path: pathlib.Path, monkeypatch, validated_profile, interactive_terminal +): + monkeypatch.setenv("DATABRICKS_HOST", "https://env.example") dest = tmp_path / "proj" - with mock.patch.object( - init_mod, - "_copy_packaged_template", - side_effect=_copy_writing({".env.example": "DATABRICKS_CONFIG_PROFILE=DEFAULT\n"}), - ): - result = CliRunner().invoke(init_mod.init, [str(dest)], obj=_Ctx()) + result = _init(dest, output="json") + assert result.exit_code == 0, result.output + assert not (dest / ".env").exists() + payload = json.loads(result.output) + assert payload["env_profile"] is None + assert payload["workspace_host"] == "https://env.example" + validated_profile.assert_not_called() # interactive, but host env auth never signs in + + +def test_init_signs_in_when_interactive( + tmp_path: pathlib.Path, validated_profile, interactive_terminal +): + result = _init(tmp_path / "proj", "--profile", "ml", output="json") + assert result.exit_code == 0, result.output + validated_profile.assert_called_once_with("ml") + assert json.loads(result.output)["signed_in_user"] == "me@example.com" + + +def test_init_noninteractive_shows_login_next_step( + tmp_path: pathlib.Path, write_databrickscfg, validated_profile +): + write_databrickscfg("[ml]\nhost = https://ml.example\n") + result = _init(tmp_path / "proj", "--profile", "ml") assert result.exit_code == 0, result.output - assert not (dest / ".env").exists() # no profile -> scaffold-only, no .env + out = " ".join(result.output.split()) + assert "Signed in as" not in out + assert "agentbricks auth login ml" in out + validated_profile.assert_not_called() def test_init_store_name_overrides(tmp_path: pathlib.Path): diff --git a/integrations/agentbricks/tests/unit_tests/runtime_store_deploy_test.py b/integrations/agentbricks/tests/unit_tests/runtime_store_deploy_test.py index 0dc4b9d6e..beb9fcdaa 100644 --- a/integrations/agentbricks/tests/unit_tests/runtime_store_deploy_test.py +++ b/integrations/agentbricks/tests/unit_tests/runtime_store_deploy_test.py @@ -191,6 +191,7 @@ def _ctx(client, workspace_client, *, output="text"): profile="prof", output=output, api_client_provider=types.SimpleNamespace(get=lambda: client), + announce_workspace=lambda: None, ) @@ -481,7 +482,7 @@ def test_unconfirmed_delete_reports_the_result_and_fails(monkeypatch, wait, stil ("wait", "expected"), [ (AppDeleteWaitResult(timed_out=True), "it may finish on its own"), - (AppDeleteWaitResult(error="expired"), "agentbricks login"), + (AppDeleteWaitResult(error="expired"), "agentbricks auth login"), ], ) def test_unconfirmed_delete_text_says_what_to_do(monkeypatch, wait, expected): diff --git a/integrations/agentbricks/tests/unit_tests/tracing_test.py b/integrations/agentbricks/tests/unit_tests/tracing_test.py index 940184eb5..1588ee701 100644 --- a/integrations/agentbricks/tests/unit_tests/tracing_test.py +++ b/integrations/agentbricks/tests/unit_tests/tracing_test.py @@ -45,6 +45,9 @@ def __init__(self, output: str = "text", profile=None, user="me@example.com"): def client(self): return mock.Mock(current_user=self._user, host="https://ws") + def announce_workspace(self) -> None: + pass + def _project(tmp_path: pathlib.Path, *, experiment_name: str | None = None): body = _AGENT_TOML diff --git a/integrations/agentbricks/tests/unit_tests/workspace_auth_test.py b/integrations/agentbricks/tests/unit_tests/workspace_auth_test.py new file mode 100644 index 000000000..e755ec26b --- /dev/null +++ b/integrations/agentbricks/tests/unit_tests/workspace_auth_test.py @@ -0,0 +1,271 @@ +"""Unit tests for profile resolution, config-file lookups, and the auth preflight/login helpers. + +The workspace boundary (`_validate_profile` / `_validate_bounded`), the `databricks` CLI +(`subprocess.run`), and the terminal check (`_is_interactive`) are stubbed (see conftest.py) so +nothing touches the network or a browser. `DATABRICKS_CONFIG_FILE` points at a tmp file (see conftest.py). +""" + +from __future__ import annotations + +import sys +import threading +from unittest import mock + +import pytest +from databricks.sdk.errors import Unauthenticated + +from databricks_agentbricks.cli import auth +from databricks_agentbricks.errors import AgentCliError +from databricks_agentbricks.projects import profiles +from databricks_agentbricks.services import workspace_auth + +_PROF_CFG = "[prof]\nhost = https://prof.example\n" + + +@pytest.fixture(autouse=True) +def _disable_telemetry(monkeypatch): + monkeypatch.setenv("AGENTBRICKS_DISABLE_TELEMETRY", "1") + + +# --- resolve_profile ------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("flag", "env_file", "env_var", "expected"), + [ + ( + "from-flag", + "DATABRICKS_CONFIG_PROFILE=from-file\n", + "from-var", + ("from-flag", "--profile"), + ), + (None, "DATABRICKS_CONFIG_PROFILE=from-file\n", "from-var", ("from-file", ".env")), + (None, None, "from-var", ("from-var", "DATABRICKS_CONFIG_PROFILE")), + (None, None, None, (None, "Databricks SDK default")), + ], + ids=["flag", "env-file", "env-var", "sdk-default"], +) +def test_resolve_profile_order(tmp_path, monkeypatch, flag, env_file, env_var, expected): + if env_file is not None: + (tmp_path / ".env").write_text(env_file) + if env_var is not None: + monkeypatch.setenv("DATABRICKS_CONFIG_PROFILE", env_var) + + assert profiles.resolve_profile(flag, tmp_path) == profiles.ProfileInfo(*expected) + + +# --- config-file lookups --------------------------------------------------------------------- + + +def test_config_file_path_expands_tilde_and_tolerates_percent_in_values(tmp_path, monkeypatch): + monkeypatch.setenv("HOME", str(tmp_path)) + (tmp_path / "cfg").write_text(_PROF_CFG + "token = 100%secret\n") + monkeypatch.setenv("DATABRICKS_CONFIG_FILE", "~/cfg") + + assert profiles.profile_host("prof") == "https://prof.example" + + +@pytest.mark.parametrize( + ("cfg", "env_host", "profile", "expected"), + [ + (_PROF_CFG, None, "prof", ("https://prof.example",) * 2), + ("[DEFAULT]\nhost = https://d.example\n", "https://env.example", None, "env"), + (_PROF_CFG, "https://env.example", "prof", "env"), + ("[DEFAULT]\nhost = https://d.example\n", None, None, "default"), + ( + "[DEFAULT]\nhost = https://d.example\n" + "[__settings__]\ndefault_profile = prof\n" + _PROF_CFG, + None, + None, + ("https://prof.example", "https://prof.example ([prof] profile)"), + ), + ( + "[__settings__]\ndefault_profile = prof\n" + _PROF_CFG, + None, + None, + ("https://prof.example", "https://prof.example ([prof] profile)"), + ), + ], + ids=[ + "named-profile", + "databricks-host", + "databricks-host-overrides-named-profile", + "default-section", + "switched-default-overrides-default-section", + "switched-default-without-default-section", + ], +) +def test_workspace_for_profile(write_databrickscfg, monkeypatch, cfg, env_host, profile, expected): + write_databrickscfg(cfg) + if env_host is not None: + monkeypatch.setenv("DATABRICKS_HOST", env_host) + if expected == "env": + expected = ("https://env.example", "https://env.example (from DATABRICKS_HOST)") + elif expected == "default": + expected = ("https://d.example", "https://d.example ([DEFAULT] profile)") + + assert profiles.workspace_for_profile(profile) == expected + + +def test_environment_auth_without_a_host_does_not_fall_back_to_config( + write_databrickscfg, monkeypatch +): + write_databrickscfg("[DEFAULT]\nhost = https://d.example\n") + monkeypatch.setenv("DATABRICKS_TOKEN", "test-token") + + assert profiles.workspace_for_profile(None) == (None, "none configured") + + +# --- preflight_auth -------------------------------------------------------------------------- + + +def test_preflight_passes_when_profile_validates(validated_profile): + client = workspace_auth.preflight_auth("prof", "--profile") + + validated_profile.assert_called_once_with("prof") + assert client is validated_profile.return_value[0][0] + + +def test_preflight_failure_hints_at_login_and_retry(fail_validation): + fail_validation(Unauthenticated("expired")) + + with pytest.raises(AgentCliError) as raised: + workspace_auth.preflight_auth("prof", ".env") + + assert "'prof' (from .env) isn't authenticated" in str(raised.value) + assert "Run `agentbricks auth login prof` to sign in, then rerun the command." == ( + str(raised.value.hint) + ) + + +def test_preflight_timeout_reports_an_unresponsive_workspace_not_a_sign_in(fail_validation): + fail_validation(TimeoutError("no response within 20s")) + + with pytest.raises(AgentCliError) as raised: + workspace_auth.preflight_auth("prof", ".env") + + assert "didn't respond within 20s" in str(raised.value) + assert "auth login" not in str(raised.value.hint or "") + + +def test_validate_bounded_gives_up_after_the_timeout(monkeypatch): + release = threading.Event() + monkeypatch.setattr(workspace_auth, "_PREFLIGHT_TIMEOUT_S", 0.05) + monkeypatch.setattr(workspace_auth, "_validate_profile", lambda profile: release.wait(5)) + + result, error = workspace_auth._validate_bounded("prof") + release.set() + + assert result is None + assert isinstance(error, TimeoutError) + + +def test_preflight_external_browser_fails_fast(write_databrickscfg, validated_profile, monkeypatch): + write_databrickscfg("[prof]\nhost = https://ws.example\nauth_type = external-browser\n") + + with pytest.raises(AgentCliError): + workspace_auth.preflight_auth("prof", "--profile") + + validated_profile.assert_not_called() + + +# --- authenticate_profile -------------------------------------------------------------------- + + +def test_authenticate_profile_returns_validated_client_without_login( + validated_profile, databricks_login +): + assert ( + workspace_auth.authenticate_profile("prof", "argument", interactive=False) + == validated_profile.return_value[0] + ) + databricks_login.assert_not_called() + + +def test_authenticate_profile_logs_in_then_revalidates(monkeypatch, databricks_login): + client = mock.Mock() + validate = mock.Mock( + side_effect=[(None, Unauthenticated("expired")), ((client, "me@example.com"), None)] + ) + monkeypatch.setattr(workspace_auth, "_validate_bounded", validate) + + assert workspace_auth.authenticate_profile("prof", "argument", interactive=True) == ( + client, + "me@example.com", + ) + + assert validate.call_count == 2 + # The child's stdout goes to stderr so it can't corrupt a command's machine-readable output. + databricks_login.assert_called_once_with( + ["databricks", "auth", "login", "--profile", "prof"], + text=True, + check=False, + stdout=sys.stderr, + ) + + +def test_auth_service_logs_in_then_revalidates_through_supplied_collaborators(): + client = mock.Mock() + validate = mock.Mock( + side_effect=[(None, Unauthenticated("expired")), ((client, "me@example.com"), None)] + ) + login = mock.Mock() + needs_browser_login = mock.Mock(return_value=False) + service = workspace_auth.WorkspaceAuthService(validate, needs_browser_login, login) + + assert service.authenticate_profile("prof", "argument", interactive=True) == ( + client, + "me@example.com", + ) + + assert validate.call_args_list == [mock.call("prof"), mock.call("prof")] + needs_browser_login.assert_called_once_with("prof") + login.assert_called_once_with("prof") + + +def test_authenticate_profile_timeout_does_not_start_a_login(fail_validation, databricks_login): + fail_validation(TimeoutError("no response within 20s")) + + with pytest.raises(AgentCliError, match="didn't respond within 20s"): + workspace_auth.authenticate_profile("prof", ".env", interactive=True) + + databricks_login.assert_not_called() + + +def test_authenticate_profile_reports_a_missing_databricks_cli(fail_validation, databricks_login): + fail_validation(Unauthenticated("expired")) + databricks_login.side_effect = FileNotFoundError + + with pytest.raises(AgentCliError, match="`databricks` CLI was not found"): + workspace_auth.authenticate_profile("prof", "argument", interactive=True) + + +# A non-auth error isn't fixed by a login, so it never opens one. An auth error would be, but only +# with someone at a terminal outside CI; otherwise the hint says how to sign in. +@pytest.mark.parametrize( + ("error", "ci", "interactive", "hint"), + [ + (RuntimeError("no creds"), False, True, None), + (Unauthenticated("expired"), False, False, "agentbricks auth login prof"), + (Unauthenticated("expired"), True, True, "agentbricks auth login prof"), + ], + ids=["non-auth-error", "no-terminal", "ci"], +) +def test_authenticate_profile_raises_without_opening_a_login( + request, monkeypatch, fail_validation, databricks_login, error, ci, interactive, hint +): + if ci: + monkeypatch.setenv("CI", "true") + if interactive: + request.getfixturevalue("interactive_terminal") + fail_validation(error) + + with pytest.raises(AgentCliError) as raised: + workspace_auth.authenticate_profile("prof", ".env", interactive=auth.can_prompt_login()) + + assert f"Could not validate Databricks profile 'prof' (from .env): {error}" in str(raised.value) + if hint is None: + assert raised.value.hint is None + else: + assert hint in (raised.value.hint or "") + databricks_login.assert_not_called()