From 07a23b55d27b23547c882d7d20149568ddd6a2d3 Mon Sep 17 00:00:00 2001 From: Elaine Wang Date: Tue, 13 Oct 2026 21:22:20 +0000 Subject: [PATCH 01/12] [agentbricks] Resolve one profile across CLI and local agent Co-authored-by: Isaac --- integrations/agentbricks/README.md | 52 ++- integrations/agentbricks/cli.md | 144 ++++++-- .../src/databricks_agentbricks/cli/app.py | 155 ++++++++- .../src/databricks_agentbricks/cli/auth.py | 308 +++++++++++------ .../src/databricks_agentbricks/cli/deploy.py | 11 + .../src/databricks_agentbricks/cli/dev.py | 67 +++- .../databricks_agentbricks/cli/endpoint.py | 2 + .../src/databricks_agentbricks/cli/init.py | 172 ++++++--- .../src/databricks_agentbricks/cli/tracing.py | 5 + .../cli/workspace_auth.py | 302 ++++++++++++++++ .../clients/api_client_provider.py | 9 +- .../clients/apps_client.py | 2 +- .../presentation/deploy.py | 6 +- .../presentation/dev.py | 22 +- .../presentation/help.py | 42 ++- .../presentation/profile.py | 85 +++++ .../databricks_agentbricks/projects/config.py | 5 + .../projects/env_file.py | 123 +++++++ .../services/dev_service.py | 18 +- .../src/databricks_agentkit/_api_client.py | 37 +- .../tests/functional/cli_smoke_test.py | 37 +- .../tests/functional/tracing_e2e_test.py | 4 +- .../tests/unit_tests/apps_client_test.py | 2 +- .../agentbricks/tests/unit_tests/auth_test.py | 325 +++++++++++------- .../agentbricks/tests/unit_tests/cli_test.py | 29 +- .../tests/unit_tests/client_test.py | 2 +- .../agentbricks/tests/unit_tests/conftest.py | 49 +++ .../tests/unit_tests/deploy_bugfix_test.py | 4 + .../tests/unit_tests/deploy_test.py | 49 ++- .../tests/unit_tests/dev_service_test.py | 18 + .../agentbricks/tests/unit_tests/dev_test.py | 110 +++++- .../tests/unit_tests/endpoint_test.py | 3 + .../tests/unit_tests/env_file_test.py | 66 ++++ .../agentbricks/tests/unit_tests/init_test.py | 114 +++--- .../unit_tests/runtime_store_deploy_test.py | 1 + .../tests/unit_tests/tracing_test.py | 3 + .../tests/unit_tests/workspace_auth_test.py | 215 ++++++++++++ 37 files changed, 2132 insertions(+), 466 deletions(-) create mode 100644 integrations/agentbricks/src/databricks_agentbricks/cli/workspace_auth.py create mode 100644 integrations/agentbricks/src/databricks_agentbricks/presentation/profile.py create mode 100644 integrations/agentbricks/src/databricks_agentbricks/projects/env_file.py create mode 100644 integrations/agentbricks/tests/unit_tests/env_file_test.py create mode 100644 integrations/agentbricks/tests/unit_tests/workspace_auth_test.py diff --git a/integrations/agentbricks/README.md b/integrations/agentbricks/README.md index 16a06a345..eec329820 100644 --- a/integrations/agentbricks/README.md +++ b/integrations/agentbricks/README.md @@ -62,7 +62,7 @@ The two ways to run an agent: 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 @@ -90,23 +90,35 @@ 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: +Sign in to a profile, then scaffold a project that records it in its `.env`: ```sh -agentbricks login --profile -agentbricks sessions stores list +agentbricks auth login +agentbricks init my-agent --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. +`agentbricks auth login` validates the profile and, in an interactive terminal, runs +`databricks auth login --profile ` if it is not authenticated yet (this needs the Databricks +CLI; in CI or without a terminal, authenticate the profile beforehand). It saves nothing. To change a +project's profile later, run `agentbricks auth switch ` in the project directory, and +`agentbricks auth describe` shows which profile (and workspace) a directory would use. + +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. 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 @@ -135,13 +147,21 @@ failed upload never change the command result. The shortest 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 auth login +agentbricks init my-agent --profile # LangGraph by default; add --framework=openai for OpenAI Agents SDK cd my-agent agentbricks dev # optional: develop and test locally agentbricks deploy my-agent # deploy to Databricks ``` +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 b9579bd7f..5e8049d1f 100644 --- a/integrations/agentbricks/cli.md +++ b/integrations/agentbricks/cli.md @@ -27,18 +27,38 @@ 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 +Sign in with `agentbricks auth login ` (a wrapper over `databricks auth login`), then +select the profile with `--profile` / `-p` on any command or record it in the project's `.env` +(`agentbricks init --profile `, or `agentbricks auth switch ` later). +Inside an Agent Bricks project (a directory with `.agentbricks/project.toml` or `app.yaml`) every +command uses that project's `.env` profile. Nothing is saved outside the project: there is no +global default profile. 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 (`[DEFAULT]` in `~/.databrickscfg`) + +`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 | | --- | --- | --- | --- | -| `--profile ` (`-p`) | string | - | `~/.databrickscfg` profile to authenticate with. | +| `--profile ` (`-p`) | string | - | Databricks config profile to authenticate with. Highest-precedence input to the resolution described under [Authentication](#authentication). | | `--output ` (`-o`) | `text` \| `json` | `text` | Output format. Use `json` for scripting. | | `--version` | flag | - | Show the version and exit. | | `--help` (`-h`) | flag | - | Show help for the command and exit. Works at every level. | @@ -57,9 +77,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 | @@ -73,36 +92,13 @@ These options apply to every command. Pass them before the command name, for exa ## Commands -### `agentbricks login` - -Authenticate a profile and save it as the default, so later commands can omit -p. - -``` -agentbricks login [options] -``` - - -_Options_ - -| Option | Values | Default | Required | Description | -| --- | --- | --- | --- | --- | -| `--profile ` (`-p`) | string | - | no | Profile to authenticate with and remember as the default. | - -### `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. Once scaffolded, deploy it with `agentbricks deploy --source `. +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. -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 `DATABRICKS_HOST` set (env-var auth, 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. @@ -136,13 +132,91 @@ _Options_ | --- | --- | --- | --- | --- | | `--framework ` | `langgraph` \| `openai` | - | 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 DATABRICKS_HOST set, 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: derived from the directory, -memory). Only --server agentbricks declares stores by default. | | `--session-store ` | string | - | no | Name for the declared session store (default: derived from the directory, -session). | | `--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]` profile's host), 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 @@ -175,7 +249,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 prints where to reach it: the chat UI if the project has one, otherwise a sample request against the agent's API. -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). @@ -1113,7 +1189,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-`. 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/src/databricks_agentbricks/cli/app.py b/integrations/agentbricks/src/databricks_agentbricks/cli/app.py index 00332909d..4aae1ae7b 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/cli/app.py +++ b/integrations/agentbricks/src/databricks_agentbricks/cli/app.py @@ -6,12 +6,14 @@ from __future__ import annotations +import functools +import pathlib from typing import Optional import click from databricks_agentbricks import errors -from databricks_agentbricks.cli.auth import load_default_profile, login, logout +from databricks_agentbricks.cli.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 @@ -21,8 +23,20 @@ from databricks_agentbricks.cli.sessions import sessions from databricks_agentbricks.cli.tools import tools from databricks_agentbricks.cli.tracing import tracing +from databricks_agentbricks.cli.workspace_auth import ( + SOURCE_FLAG, + SOURCE_SDK_DEFAULT, + ProfileInfo, + legacy_default_profile, + resolve_profile, +) 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 is_agentbricks_project from databricks_agentkit._api_client import _AgentBricksApiClient @@ -30,9 +44,50 @@ class CliContext: """Shared per-invocation state: selected profile, output mode, lazily-built client.""" def __init__(self, profile: Optional[str], output: str): - self.profile = profile + self.profile_info = resolve_profile(profile) self.output = output - self.api_client_provider = ApiClientProvider(profile) + self._workspace_announced = False + self.api_client_provider = self._client_provider(self.profile) + + @property + def profile(self) -> Optional[str]: + return self.profile_info.name + + def use_project_profile( + self, project_dir: pathlib.Path, env_values: Optional[dict[str, str]] = None + ) -> None: + """Fold in the project's `.env` profile (see `resolve_profile`) unless `-p` was given. + + The root group calls this for the current directory when it's an Agent Bricks project; + `dev` and `deploy` call it again with their source dir, whose `.env` then replaces the + current directory's. Only an explicit `-p` outranks the `.env`. + ``env_values`` passes an already-parsed ``/.env`` so it isn't read twice. + """ + if self.profile_info.source == SOURCE_FLAG: + return + resolved = resolve_profile(None, project_dir, env_values) + if resolved.name != self.profile_info.name: + self.api_client_provider = self._client_provider(resolved.name) + self.profile_info = resolved + + def use_flag_profile(self, profile: str) -> None: + """Apply a `-p/--profile` given after the subcommand; like the global one, it beats `.env`.""" + if profile != self.profile: + self.api_client_provider = self._client_provider(profile) + self.profile_info = ProfileInfo(profile, SOURCE_FLAG) + + def announce_workspace(self) -> 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) + # 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_provider(self, profile: Optional[str]) -> ApiClientProvider: + return ApiClientProvider(profile, on_first_use=self.announce_workspace) def client(self) -> _AgentBricksApiClient: """Compatibility accessor for commands not yet migrated to ``api_client_provider``.""" @@ -41,7 +96,7 @@ def client(self) -> _AgentBricksApiClient: @click.group(name="agentbricks", context_settings={"help_option_names": ["-h", "--help"]}) @click.option( - "--profile", "-p", default=None, help="~/.databrickscfg profile to authenticate with." + "--profile", "-p", default=None, help="Databricks config profile to authenticate with." ) @click.option( "--output", @@ -64,8 +119,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: @@ -83,12 +141,86 @@ 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) + cwd = pathlib.Path.cwd() + if is_agentbricks_project(cwd): + # Every command run inside a project uses its `.env` profile, so e.g. `endpoint invoke` + # targets the same workspace that `deploy` just used. + ctx.obj.use_project_profile(cwd) + + +_REMOVED_COMMANDS = { + "login": "Sign in with `agentbricks auth login `, and pick a project's profile with " + "`agentbricks auth switch `.", + "logout": "Profiles are now chosen per project in its .env (`agentbricks auth switch `); " + "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.use_flag_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 + if source_profile and any(param.name == "source" for param in command.params): + _use_source_project_profile(command) + 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 _use_source_project_profile(command: click.Command) -> None: + """Resolve the profile from `--source`'s project, so e.g. `tracing bind --source proj` reads + the same workspace that `deploy --source proj` writes to (only `-p` outranks it).""" + callback = command.callback + if callback is None: + return + + @functools.wraps(callback) + def with_source_profile(*args, **kwargs): + obj = click.get_current_context().find_object(CliContext) + source = kwargs.get("source") + if obj is not None and source and is_agentbricks_project(pathlib.Path(source)): + obj.use_project_profile(pathlib.Path(source)) + return callback(*args, **kwargs) + + command.callback = with_source_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) @@ -98,6 +230,11 @@ def agentbricks(ctx: click.Context, profile: Optional[str], output: str) -> None agentbricks.add_command(deployments) agentbricks.add_command(endpoint) agentbricks.add_command(tools) +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 faa4c5e0d..72d18bfa9 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/cli/auth.py +++ b/integrations/agentbricks/src/databricks_agentbricks/cli/auth.py @@ -1,142 +1,240 @@ -"""`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.cli.workspace_auth import ( + SOURCE_ENV_FILE, + SOURCE_FLAG, + ProfileInfo, + authenticate_profile, + can_prompt_login, + profile_host, + resolve_profile, + run_databricks_login, + workspace_for_profile, +) from databricks_agentbricks.errors import AgentCliError from databricks_agentbricks.presentation import render -from databricks_agentkit._api_client import _AgentBricksApiClient +from databricks_agentbricks.presentation.profile import ( + announce_login, + warn_sign_in_failed, + warn_unknown_profile, +) +from databricks_agentbricks.projects.config import is_agentbricks_project +from databricks_agentbricks.projects.env_file import read_env, update_env_profile +_SOURCE_HELP = "Project directory. Defaults to the current directory." -def _config_file() -> pathlib.Path: - base = os.environ.get("AGENTBRICKS_CONFIG_HOME") - root = pathlib.Path(base) if base else pathlib.Path.home() / ".agentbricks" - return root / "config.json" +def sign_in_if_interactive(obj, profile: str, source: str) -> Optional[str]: + """Sign in to `profile` right after it's chosen, so `agentbricks dev` works right away. -def load_default_profile() -> Optional[str]: - """The profile saved by `agentbricks login`, or None if the user never logged in.""" + Interactive terminals only: a scripted or CI run never opens a browser or waits on the + workspace. 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(): + return None + announce_login(obj, profile, source) try: - return json.loads(_config_file().read_text()).get("profile") - except (OSError, json.JSONDecodeError): + _, user = authenticate_profile(profile, source) + except AgentCliError as exc: + warn_sign_in_failed(profile, exc, obj) return None + return user -def _save_default_profile(profile: str) -> None: - path = _config_file() - path.parent.mkdir(parents=True, exist_ok=True) - path.write_text(json.dumps({"profile": profile}, indent=2) + "\n") - - -def _validate_profile(profile: str) -> tuple[_AgentBricksApiClient, str]: - client = _AgentBricksApiClient(profile) - return client, client.current_user - - -def _is_interactive() -> bool: - return sys.stdin.isatty() - - -def _run_databricks_login(profile: str) -> None: - command = ["databricks", "auth", "login", "--profile", profile] - try: - # Keep the child process interactive while preserving stdout for Agent Bricks JSON output. - result = subprocess.run(command, text=True, check=False, stdout=sys.stderr) - except FileNotFoundError as exc: +def _require_project(source: str) -> pathlib.Path: + directory = pathlib.Path(source).resolve() + if not is_agentbricks_project(directory): 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: - 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 directory + + +def _resolve_for(obj, directory: pathlib.Path) -> 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 is_agentbricks_project(directory) else None) + + +def _login_target(obj, name: Optional[str], directory: pathlib.Path) -> 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 but a terminal can pick one. + """ + if name: + return name, "argument" + if obj.profile_info.source == SOURCE_FLAG and obj.profile_info.name: + return obj.profile_info.name, SOURCE_FLAG + if is_agentbricks_project(directory): + env_profile = read_env(directory / ".env").get("DATABRICKS_CONFIG_PROFILE") + if env_profile: + return env_profile, SOURCE_ENV_FILE + if can_prompt_login(): + return None + raise AgentCliError( + "No profile to log in to.", + hint="Pass one (`agentbricks auth login `), or set the project's profile " + "with `agentbricks auth switch `.", + ) -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 +@click.group("auth") +def auth_group() -> None: + """Choose, inspect, and sign in to the Databricks profile a project uses. - 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 + 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="Profile to authenticate with and remember as the default.", + "--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: - """Authenticate a profile and save it as the default, so later commands can omit -p.""" - profile = profile or obj.profile - if not profile: - raise AgentCliError( - "No profile to save.", - hint="Pass one to remember, e.g. `agentbricks login --profile `.", - ) - client, user = _authenticate_profile(profile) - _save_default_profile(profile) +def auth_switch(obj, name: str, source: str) -> None: + """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 (see `auth + login`); non-interactive and CI runs make no workspace calls. + """ + directory = _require_project(source) + changed = update_env_profile(directory, name) + if not can_prompt_login(): + # Interactive runs sign in next, which creates a missing profile. + warn_unknown_profile(name, obj) + user = sign_in_if_interactive(obj, name, "argument") + host = profile_host(name) if obj.output == "json": - render.emit_json({"profile": profile, "user": user, "host": client.host}) + 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"Logged in as {user}", - fields={"Profile": profile, "Host": client.host}, - next_steps=[ - ("agentbricks init my-agent", "Scaffold a new agent project"), - ], + f"Project profile set → '{name}'" if changed else f"Project profile is already '{name}'", + fields=fields, + next_steps=None if user else [(f"agentbricks auth login {name}", "Sign in")], ) -@click.command() +@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() + info = _resolve_for(obj, directory) + host, workspace = workspace_for_profile(info.name) + project = str(directory) if is_agentbricks_project(directory) else None + if obj.output == "json": + render.emit_json( + { + "profile": info.name, + "source": info.source, + "workspace_host": host, + "project": project, + } + ) + return + fields = { + "Profile": info.name or "none (Databricks SDK default)", + "Source": info.source, + } + if workspace: + fields["Workspace"] = workspace + fields["Project"] = project or "not in a project" + render.detail("Profile", info.name or "SDK default", fields) + + +@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) +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. + """ + directory = pathlib.Path(source).resolve() + target = _login_target(obj, name, directory) + if target is 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) + if obj.output == "json": + render.emit_json({"profile": None, "source": "databricks auth login"}) + else: + render.success("Logged in with `databricks auth login`") + return + chosen, origin = target + announce_login(obj, chosen, origin) + client, user = authenticate_profile(chosen, origin) if obj.output == "json": - render.emit_json({"logged_out": existed}) + render.emit_json({"profile": chosen, "source": origin, "user": user, "host": client.host}) return - render.success("Logged out" if existed else "No saved login to clear") + render.success( + f"Logged in to '{chosen}'", + fields={"Profile": chosen, "Source": origin, "User": user, "Workspace": client.host}, + ) + if is_agentbricks_project(directory): + env_profile = read_env(directory / ".env").get("DATABRICKS_CONFIG_PROFILE") + if env_profile and env_profile != chosen: + render.console().print( + f"[dim]This project still uses profile '{env_profile}' (from .env). " + f"Run `agentbricks auth switch {chosen}` to switch it.[/]" + ) diff --git a/integrations/agentbricks/src/databricks_agentbricks/cli/deploy.py b/integrations/agentbricks/src/databricks_agentbricks/cli/deploy.py index 718c93ac9..85fb1a25f 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/cli/deploy.py +++ b/integrations/agentbricks/src/databricks_agentbricks/cli/deploy.py @@ -18,8 +18,11 @@ from __future__ import annotations +import pathlib + import click +from databricks_agentbricks.cli.workspace_auth import preflight_auth from databricks_agentbricks.clients.apps_client import AppsClient from databricks_agentbricks.clients.apps_user_auth_client import AppsUserAuthClient from databricks_agentbricks.clients.conversation_store_client import ( @@ -64,6 +67,7 @@ def build_deploy_service(obj) -> DeployService: """Compose a deployment service for this CLI invocation without opening an API client.""" + obj.announce_workspace() api_client_provider = obj.api_client_provider apps_client = AppsClient(obj.profile, runner=_databricks) apps_user_auth_client = AppsUserAuthClient(obj.profile) @@ -162,6 +166,13 @@ def deploy( not authentication and not the session id itself (send the session id in the request body): X-Routing-Key: """ + # Fold the project's `.env` into the profile resolution before the profile is used, matching + # `agentbricks dev` so both commands deploy and run against the same workspace. + obj.use_project_profile(pathlib.Path(source)) + obj.announce_workspace() + # 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. + preflight_auth(obj.profile, obj.profile_info.source) request = DeployRequest( name=name, source=source, diff --git a/integrations/agentbricks/src/databricks_agentbricks/cli/dev.py b/integrations/agentbricks/src/databricks_agentbricks/cli/dev.py index 15333cebc..0625bf0be 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/cli/dev.py +++ b/integrations/agentbricks/src/databricks_agentbricks/cli/dev.py @@ -11,20 +11,23 @@ import pathlib import subprocess -from typing import Optional +from typing import Callable, Optional import click from databricks_agentbricks.cli.tracing import start_local_tracing_server, stop_local_tracing_server +from databricks_agentbricks.cli.workspace_auth import preflight_auth from databricks_agentbricks.clients.apps_client import AppsClient from databricks_agentbricks.clients.databricks_cli import _databricks from databricks_agentbricks.presentation.dev import ( + announce_agent_profile, announce_bound_resources, local_trace_label, ) from databricks_agentbricks.presentation.dev import ( announce_local_url as _announce_local_url, ) +from databricks_agentbricks.projects.env_file import read_env from databricks_agentbricks.projects.resolver import ProjectResolver from databricks_agentbricks.services.dev_service import ( DevRequest, @@ -43,12 +46,28 @@ def stop(self, server: subprocess.Popen) -> None: stop_local_tracing_server(server) -def build_dev_service(obj) -> DevService: +# `.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, preflight: Optional[Callable[[], None]] = None) -> DevService: """Compose the local workflow without opening a workspace API client.""" return DevService( project_resolver=ProjectResolver(), apps_client=AppsClient(obj.profile, runner=_databricks), local_tracing=_CliLocalTracing(), + preflight=preflight, ) @@ -79,8 +98,9 @@ def dev( reach it: the chat UI if the project has one, otherwise a sample request against the agent's API. - 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 @@ -96,11 +116,46 @@ 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. """ - service = build_dev_service(obj) + source_dir = pathlib.Path(source) + # The CLI and the locally running agent read the same `.env`, so resolve the profile from it + # once here (dev also consults it below for explicit credentials). + env_file = read_env(source_dir / ".env") + obj.use_project_profile(source_dir, env_file) + 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 + + def _preflight() -> None: + if not env_credential: + preflight_auth(obj.profile, obj.profile_info.source) + + # Announced before the preflight so an auth failure still names its target. When `.env` + # credentials win, the resolved profile would name a different workspace than the agent + # actually uses, so it isn't announced. + if not uses_env_credentials: + obj.announce_workspace() + service = build_dev_service(obj, _preflight) 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: 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_local_url( plan.preview.source_dir, plan.preview.port, diff --git a/integrations/agentbricks/src/databricks_agentbricks/cli/endpoint.py b/integrations/agentbricks/src/databricks_agentbricks/cli/endpoint.py index 67dfcf971..f07bf0c44 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/cli/endpoint.py +++ b/integrations/agentbricks/src/databricks_agentbricks/cli/endpoint.py @@ -60,6 +60,8 @@ def invoke( auth, ) -> None: """Send one HTTP request to a Databricks App or arbitrary URL.""" + if app: + obj.announce_workspace() service = build_invoke_service(obj) request = InvokeRequest( app=app, diff --git a/integrations/agentbricks/src/databricks_agentbricks/cli/init.py b/integrations/agentbricks/src/databricks_agentbricks/cli/init.py index 158e23dfb..140d39623 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,11 +26,27 @@ import click +from databricks_agentbricks.cli.auth import sign_in_if_interactive from databricks_agentbricks.cli.tracing import default_experiment_name +from databricks_agentbricks.cli.workspace_auth import ( + SOURCE_ENV_VAR, + SOURCE_FLAG, + SOURCE_SETTINGS_DEFAULT, + can_prompt_login, + configured_profiles, + profile_exists, + profile_host, + settings_default_profile, +) from databricks_agentbricks.errors import AgentCliError from databricks_agentbricks.presentation import render +from databricks_agentbricks.presentation.profile import note_choose_profile, warn_unknown_profile from databricks_agentbricks.projects.agent_project import AgentProject, 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.types import ( AgentFramework, AgentServer, @@ -118,29 +135,53 @@ 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 `agentbricks auth switch --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 DATABRICKS_HOST is set, 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 when `DATABRICKS_HOST` is set because env-var auth (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 os.environ.get("DATABRICKS_HOST"): + return None + switched = settings_default_profile() + if switched: + return switched, SOURCE_SETTINGS_DEFAULT + return "DEFAULT", "default" def _migration_paths(dest: pathlib.Path, target: pathlib.Path) -> tuple[pathlib.Path, ...]: @@ -320,9 +361,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 DATABRICKS_HOST set, no .env is written).", ) @click.option( "--disable-chat-app", @@ -365,11 +408,21 @@ 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. Once scaffolded, deploy it with + 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 `. - 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`. With DATABRICKS_HOST set (env-var auth, 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. @@ -411,11 +464,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: @@ -446,13 +502,36 @@ def init( experiment_name=experiment_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] == "default" and not profile_exists("DEFAULT") + else [] + ) + if other_profiles: + note_choose_profile(other_profiles, obj) + elif env_profile and not can_prompt_login(): + # Interactive runs sign in right after scaffolding, which creates a missing profile. + warn_unknown_profile(env_profile, obj) + + # A failed sign-in doesn't undo the scaffold; the summary says how to retry. + signed_in_user = ( + sign_in_if_interactive(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( { @@ -462,7 +541,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, @@ -489,15 +570,22 @@ def init( if experiment_name: fields["Traces experiment"] = experiment_name 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 (DATABRICKS_HOST); 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( + ("agentbricks auth switch ", f"Pick one of: {', '.join(other_profiles)}") + ) + elif env_profile: + steps.append((f"agentbricks auth login {env_profile}", "Sign in before running the agent")) + if not other_profiles: + steps.append(("agentbricks auth switch ", "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 e4ef758e5..e7f65dba0 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/cli/tracing.py +++ b/integrations/agentbricks/src/databricks_agentbricks/cli/tracing.py @@ -555,6 +555,7 @@ def tracing_bind(obj, experiment_name, experiment_id, source) -> None: # service principal MODIFY on its UC OTEL tables through TracingClient. name = experiment_name if experiment_id: + obj.announce_workspace() mlflow = _mlflow() _set_tracking_uri(mlflow, obj.profile) experiment = _get_experiment_by_id(mlflow, experiment_id) @@ -651,6 +652,8 @@ def tracing_list(obj, experiment_name, experiment_id, warehouse_id, limit, sourc with _with_trace_read_target( source, obj.profile, experiment_name, experiment_id, _resolve_warehouse_id(warehouse_id) ) as target: + if not target.local and target.tracking_uri: + obj.announce_workspace() traces = [] if target.experiment_id: traces = mlflow.search_traces( @@ -709,6 +712,8 @@ def tracing_get(obj, trace_id, experiment_name, experiment_id, warehouse_id, sou ) as target: # get_trace resolves by id and needs only the tracking URI; when nothing resolved, fall back to # the workspace so an id still looks there. + if not target.local: + obj.announce_workspace() if target.tracking_uri is None: mlflow.set_tracking_uri(_workspace_uri(obj.profile)) if warehouse: diff --git a/integrations/agentbricks/src/databricks_agentbricks/cli/workspace_auth.py b/integrations/agentbricks/src/databricks_agentbricks/cli/workspace_auth.py new file mode 100644 index 000000000..60dc2aacc --- /dev/null +++ b/integrations/agentbricks/src/databricks_agentbricks/cli/workspace_auth.py @@ -0,0 +1,302 @@ +"""Profile resolution, workspace lookup, and authentication checks for the Agent Bricks CLI.""" + +from __future__ import annotations + +import configparser +import dataclasses +import json +import os +import pathlib +import subprocess +import sys +import threading +from typing import Optional + +from databricks_agentbricks.errors import AgentCliError +from databricks_agentbricks.projects.env_file import read_env +from databricks_agentkit._api_client import ( + _AgentBricksApiClient, + _databricks_config_parser, + _profile_host, + _profile_section, +) + +# 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" + + +@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 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_host(profile) + + +def profile_exists(profile: str) -> bool: + """Whether the profile is configured in the Databricks config file; never raises.""" + parser = _databricks_config_parser() + if profile == parser.default_section: + # configparser never reports [DEFAULT] via has_section(); it exists once it has any keys. + return bool(parser.defaults()) + return parser.has_section(profile) + + +def workspace_for_profile(profile: Optional[str]) -> tuple[Optional[str], Optional[str]]: + """The workspace a profile points at, as (host, display text); (None, None) when unknown. + + For display only: without a profile the text names where the host came from + (DATABRICKS_HOST, else the config file's `[DEFAULT]`), because "no workspace" is the answer + the caller needs to surface. It does not claim to mirror the SDK's full default resolution. + """ + if profile: + host = profile_host(profile) + return host, host + env_host = os.environ.get("DATABRICKS_HOST") + if env_host: + return env_host, f"{env_host} (from DATABRICKS_HOST)" + default_host = profile_host(_databricks_config_parser().default_section) + if default_host: + return default_host, f"{default_host} ([DEFAULT] profile)" + return None, "none configured" + + +def _validate_profile(profile: Optional[str]) -> tuple[_AgentBricksApiClient, str]: + client = _AgentBricksApiClient(profile) + return client, client.current_user + + +def _is_interactive() -> bool: + return sys.stdin.isatty() + + +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 _is_auth_error(error: Optional[Exception]) -> bool: + """Whether a validation failure means the profile isn't signed in, so a login could fix it.""" + from databricks.sdk.errors import Unauthenticated + + return isinstance(error, (AgentCliError, Unauthenticated)) + + +def _is_unreachable(error: Exception) -> bool: + """Whether validation failed on the network or the server, so signing in wouldn't fix it.""" + from databricks.sdk.errors import DeadlineExceeded, InternalError, TemporarilyUnavailable + from requests.exceptions import ConnectionError as RequestsConnectionError + + return isinstance( + error, (RequestsConnectionError, InternalError, TemporarilyUnavailable, DeadlineExceeded) + ) + + +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 + + +# 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 _preflight_error(profile: Optional[str], source: str) -> AgentCliError: + if profile: + return AgentCliError( + f"Databricks profile {profile!r} (from {source}) isn't authenticated.", + hint=f"Run `agentbricks auth login {profile}` to sign in, then rerun the command.", + ) + return AgentCliError( + "No Databricks profile is configured and the Databricks SDK's default authentication " + "didn't authenticate.", + hint="Run `agentbricks auth login ` to sign in, then pass `-p ` or " + "choose one for the project with `agentbricks auth switch ` (or " + "DATABRICKS_CONFIG_PROFILE), and rerun the command.", + ) + + +SOURCE_SETTINGS_DEFAULT = "databricks auth switch" + + +def settings_default_profile() -> Optional[str]: + """The default profile chosen with `databricks auth switch` ([__settings__] default_profile).""" + return _profile_section("__settings__").get("default_profile", "").strip() or None + + +def configured_profiles() -> list[str]: + """Profile names in the Databricks config file, `[DEFAULT]` included when it has keys.""" + parser = _databricks_config_parser() + names = [name for name in parser.sections() if name != "__settings__"] + return ([parser.default_section] if parser.defaults() else []) + names + + +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 and any( + os.environ.get(name) + for name in ("DATABRICKS_HOST", "DATABRICKS_TOKEN", "DATABRICKS_CLIENT_ID") + ): + return None # env-var auth without a profile: the SDK never reads the config file + if not profile: + profile = settings_default_profile() or configparser.DEFAULTSECT + return _profile_section(profile).get("auth_type") + + +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`. + + A browser-OAuth flow can block forever, so the call is bounded. Daemon because executor + threads are joined at interpreter exit and a hung flow 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] + + +def preflight_auth(profile: Optional[str], source: str) -> None: + """Fail fast when the resolved profile can't authenticate, before a command spins anything up. + + Without this, an unauthenticated `external-browser` profile hangs on a browser OAuth flow that + never completes — the local agent sits stuck with no error. The validation call itself can + hang that way, so it runs bounded (see `_validate_bounded`). + """ + if _needs_browser_login(profile): + # The bounded validation would open a browser flow this command will abandon at its + # timeout, so fail before it starts — interactive or not. + raise _preflight_error(profile, source) + + _, error = _validate_bounded(profile) + if error is None: + return + if _is_unreachable(error): + target = f"profile {profile!r} (from {source})" if profile else "the default auth" + raise 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.", + ) from error + raise _preflight_error(profile, source) from error + + +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 = f"agentbricks auth login {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 authenticate_profile(profile: str, source: str) -> tuple[_AgentBricksApiClient, str]: + """Validate `profile`, signing in through `databricks auth login` once if it isn't yet. + + Returns the client and the signed-in user; `source` is only named in errors. Without a + terminal (or in CI) nobody can complete a browser login, so an unauthenticated profile fails + with a hint instead of opening one. + """ + label = f"Databricks profile {profile!r} (from {source})" + error: Optional[Exception] + if _needs_browser_login(profile): + # Validation would open a browser flow it then abandons; head straight for the login. + error = None + else: + client_and_user, error = _validate_bounded(profile) + if client_and_user is not None: + return client_and_user + if not _is_auth_error(error): + # A timeout or a network/5xx failure won't be fixed by a login; don't open one. + raise AgentCliError(f"Could not validate {label}: {error}") from error + if not can_prompt_login(): + detail = "external-browser auth needs an interactive login" if error is None else str(error) + raise AgentCliError( + f"Could not validate {label}: {detail}", + hint="Run this command in an interactive terminal so Agent Bricks can open " + f"Databricks login, or authenticate first with `databricks auth login --profile {profile}`.", + ) from error + run_databricks_login(profile) + client_and_user, error = _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 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..0008a05b4 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/clients/api_client_provider.py +++ b/integrations/agentbricks/src/databricks_agentbricks/clients/api_client_provider.py @@ -11,7 +11,7 @@ from __future__ import annotations from threading import Lock -from typing import Optional +from typing import Callable, Optional from databricks_agentkit._api_client import _AgentBricksApiClient @@ -19,8 +19,11 @@ 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[[], None]] = None + ) -> None: self._profile = profile + self._on_first_use = on_first_use self._client: Optional[_AgentBricksApiClient] = None self._lock = Lock() @@ -32,5 +35,7 @@ def get(self) -> _AgentBricksApiClient: with self._lock: if self._client is None: + if self._on_first_use is not None: + self._on_first_use() self._client = _AgentBricksApiClient(self._profile) 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 b39afb10c..9fe88c92a 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/clients/apps_client.py +++ b/integrations/agentbricks/src/databricks_agentbricks/clients/apps_client.py @@ -145,7 +145,7 @@ def exists(self, name: str) -> bool: marker in detail.lower() for marker in ("resource_does_not_exist", "does not exist") ): return False - login = f"agentbricks login --profile {self._profile or ''}" + login = f"agentbricks auth login {self._profile or ''}" raise 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/presentation/deploy.py b/integrations/agentbricks/src/databricks_agentbricks/presentation/deploy.py index 9a651b1b6..0e43154e9 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/presentation/deploy.py +++ b/integrations/agentbricks/src/databricks_agentbricks/presentation/deploy.py @@ -98,7 +98,11 @@ def present_deleted(name: str, *, output: Optional[str]) -> None: render.success(f"Deleted deployment '{name}'") -def present_deploy_result(result: DeployResult, *, output: Optional[str]) -> None: +def present_deploy_result( + result: DeployResult, + *, + output: Optional[str], +) -> None: """Report a finished deploy: the machine payload for `--output json`, else the success panel.""" if output == "json": store_grant_attempted = result.memory_grant_attempted or result.session_grant_attempted diff --git a/integrations/agentbricks/src/databricks_agentbricks/presentation/dev.py b/integrations/agentbricks/src/databricks_agentbricks/presentation/dev.py index af4782c92..4b3c92267 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_local_url( source_dir: pathlib.Path, port: int, @@ -57,7 +77,7 @@ def announce_local_url( else ("agentbricks tools add mcp ", "Give the agent a tool") ) 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 2e000284a..d775347b5 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")), ) @@ -28,14 +28,15 @@ _EXAMPLES: dict[CommandPath, tuple[Example, ...]] = { (): ( - ("agentbricks login --profile ", "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 my-agent", "deploy the agent to Databricks Apps"), ), - ("login",): (("agentbricks login --profile ", "save a profile as your default"),), - ("logout",): (("agentbricks logout", "forget the saved default profile"),), ("init",): ( ("agentbricks init my-agent", "scaffold a new agent project"), ( @@ -43,6 +44,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"), ), @@ -423,9 +444,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", @@ -482,7 +502,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 @@ -510,9 +530,9 @@ def _root_epilog() -> str: auth = "\n".join( [ "\b", - "Not authenticated yet? Create a profile with the Databricks CLI first:", - " databricks auth login --profile ", - "Then `agentbricks login --profile ` 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..3388a6f78 --- /dev/null +++ b/integrations/agentbricks/src/databricks_agentbricks/presentation/profile.py @@ -0,0 +1,85 @@ +"""Human-facing output for the profile commands.""" + +from __future__ import annotations + +import click + +from databricks_agentbricks.cli.workspace_auth import ( + profile_exists, + profile_host, + workspace_for_profile, +) +from databricks_agentbricks.presentation import render + + +def announce_workspace(profile: str | None, source: str, *, output: str) -> 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 + _, workspace = workspace_for_profile(profile) + 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 `agentbricks auth switch ` 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 `agentbricks auth login {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 = profile_host(profile) + where = host or ( + "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 `agentbricks auth switch {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..1035f55a0 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/projects/config.py +++ b/integrations/agentbricks/src/databricks_agentbricks/projects/config.py @@ -26,6 +26,11 @@ 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 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..82d9dee9a --- /dev/null +++ b/integrations/agentbricks/src/databricks_agentbricks/projects/env_file.py @@ -0,0 +1,123 @@ +"""Reading and read-modify-writing of a project's local `.env` profile entry.""" + +from __future__ import annotations + +import os +import pathlib +import re +import tempfile + +# Same shape the reader accepts, so a line the CLI honors is a line it can rewrite. +_PROFILE_LINE = re.compile(r"\s*(export\s+)?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]: + """Minimal KEY=VALUE reader for a project `.env` (python-dotenv is not a CLI dependency).""" + values: dict[str, str] = {} + try: + text = _read_text(path) + except OSError: + return values + for line in _lines(text): + line = line.strip() + if not line or line.startswith("#") or "=" not in line: + continue + if line.startswith("export "): + line = line[len("export ") :].strip() + key, _, value = line.partition("=") + key, value = key.strip(), value.strip() + if value[:1] in ("'", '"'): + end = value.find(value[0], 1) + if end != -1: + # A quoted value ends at its matching close quote: a `#` inside is literal, and + # anything after (an inline comment) is dropped. + value = value[1:end] + else: + value = re.split(r"\s+#", value, maxsplit=1)[0].rstrip() + else: + # Match dotenv for unquoted values: a `#` that follows whitespace starts a comment. + value = re.split(r"\s+#", value, maxsplit=1)[0].rstrip() + if key: + values[key] = value + return values + + +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) + if read_env(env_path).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. + return False + before = _read_text(env_path) + 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/services/dev_service.py b/integrations/agentbricks/src/databricks_agentbricks/services/dev_service.py index 1e5995e5f..e18d29f2a 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/services/dev_service.py +++ b/integrations/agentbricks/src/databricks_agentbricks/services/dev_service.py @@ -10,10 +10,10 @@ import pathlib import subprocess -from collections.abc import Iterator +from collections.abc import Iterator, Mapping from contextlib import contextmanager -from dataclasses import dataclass -from typing import Protocol +from dataclasses import dataclass, field +from typing import Callable, Optional, Protocol from databricks_agentbricks.clients.apps_client import AppsClient from databricks_agentbricks.errors import AgentCliError @@ -49,6 +49,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) @@ -99,10 +101,12 @@ def __init__( project_resolver: ProjectResolver, apps_client: AppsClient, local_tracing: LocalTracing, + preflight: Optional[Callable[[], None]] = None, ) -> None: self._project_resolver = project_resolver self._apps_client = apps_client self._local_tracing = local_tracing + self._preflight = preflight @contextmanager def prepare(self, request: DevRequest) -> Iterator[DevPlan]: @@ -120,6 +124,11 @@ def prepare(self, request: DevRequest) -> Iterator[DevPlan]: hint="Run from a scaffolded project, or pass --source (see `agentbricks init`).", ) + # 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 self._preflight is not None: + self._preflight() + project = self._project_resolver.load(source_dir) if project is not None and project.tools: require_managed_tool_support(source_dir) @@ -141,7 +150,8 @@ def prepare(self, request: DevRequest) -> Iterator[DevPlan]: if prepare_environment is None: prepare_environment = not (source_dir / ".venv").exists() - entry_point = _dev_entry_point(app_yaml, tracing_env or None) + manifest_env = {**request.local_env, **tracing_env} + entry_point = _dev_entry_point(app_yaml, manifest_env or None) preview = DevPreview( source_dir=source_dir, # Preserve the CLI's historical display/default behavior while retaining the diff --git a/integrations/agentbricks/src/databricks_agentkit/_api_client.py b/integrations/agentbricks/src/databricks_agentkit/_api_client.py index de14a3d03..a47bb9bd8 100644 --- a/integrations/agentbricks/src/databricks_agentkit/_api_client.py +++ b/integrations/agentbricks/src/databricks_agentkit/_api_client.py @@ -104,16 +104,35 @@ def memory_entry_path(store: str, entry: str) -> str: return f"{memory_store_path(store)}/entries/{entry}" -def _profile_host(profile: str) -> Optional[str]: - config_path = pathlib.Path( - os.getenv("DATABRICKS_CONFIG_FILE", pathlib.Path.home() / ".databrickscfg") - ) - parser = configparser.ConfigParser() +def _databricks_config_parser( + default_section: str = configparser.DEFAULTSECT, +) -> configparser.ConfigParser: + """The parsed Databricks config file, or an empty parser on any problem (never raises). + + expanduser: the SDK expands `~` in DATABRICKS_CONFIG_FILE; without it a tilde path reads + nothing. No interpolation: values are taken verbatim — a `%` must not raise at get-time. + """ + parser = configparser.ConfigParser(interpolation=None, default_section=default_section) try: - parser.read(config_path) + parser.read( + pathlib.Path( + os.getenv("DATABRICKS_CONFIG_FILE", str(pathlib.Path.home() / ".databrickscfg")) + ).expanduser() + ) except (OSError, configparser.Error): - return None - return parser.get(profile, "host", fallback=None) + pass + return parser + + +def _profile_section(profile: str) -> dict[str, str]: + """`profile`'s own keys, as the SDK reads them: a named profile doesn't inherit [DEFAULT].""" + # With an unused fallback section, [DEFAULT] parses as an ordinary section like any other. + parser = _databricks_config_parser(default_section="\0no-fallback") + return dict(parser[profile]) if parser.has_section(profile) else {} + + +def _profile_host(profile: str) -> Optional[str]: + return _profile_section(profile).get("host") def _bound_retry_timeout(client: WorkspaceClient) -> WorkspaceClient: @@ -166,7 +185,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/functional/cli_smoke_test.py b/integrations/agentbricks/tests/functional/cli_smoke_test.py index fe7fb2850..725f3f3ce 100644 --- a/integrations/agentbricks/tests/functional/cli_smoke_test.py +++ b/integrations/agentbricks/tests/functional/cli_smoke_test.py @@ -29,9 +29,8 @@ # Every top-level command; ` --help` proves each module imports and wires up as installed. _COMMANDS = ( - "login", - "logout", "init", + "auth", "doctor", "dev", "memory", @@ -57,8 +56,8 @@ def run_agentbricks(tmp_path: pathlib.Path): home.mkdir() empty_cfg = tmp_path / "empty.databrickscfg" empty_cfg.write_text("") - # Isolate HOME so login/logout can't touch the real ~/.agentbricks; empty config file so no real - # profile is reachable. env= replaces the environment wholesale (no ambient DATABRICKS_* leaks). + # Isolate HOME so no real token cache is visible; empty config file so no real profile is + # reachable. env= replaces the environment wholesale (no ambient DATABRICKS_* leaks). env = { "PATH": f"{agentbricks.parent}:/usr/bin:/bin", "HOME": str(home), @@ -86,7 +85,7 @@ def run( return run -@pytest.mark.parametrize("profile_source", ["saved", "sdk"]) +@pytest.mark.parametrize("profile_source", ["flag", "sdk"]) @pytest.mark.parametrize("command", ["memory-list", "slow-init"]) def test_cli_posts_telemetry_to_local_receiver( run_agentbricks, tmp_path, profile_source, command @@ -131,15 +130,9 @@ def do_POST(self) -> None: f"[{profile_name}]\nhost = http://127.0.0.1:{server.server_port}\n" "token = local-test-token\n" ) - agentbricks_home = tmp_path / "agentbricks-home" - agentbricks_home.mkdir() - env_override = { - "DATABRICKS_CONFIG_FILE": str(config_file), - "AGENTBRICKS_CONFIG_HOME": str(agentbricks_home), - } - if profile_source == "saved": - (agentbricks_home / "config.json").write_text(json.dumps({"profile": profile_name})) - else: + env_override = {"DATABRICKS_CONFIG_FILE": str(config_file)} + profile_args = ["-p", profile_name] if profile_source == "flag" else [] + if profile_source == "sdk": env_override["DATABRICKS_CONFIG_PROFILE"] = profile_name if command == "slow-init": @@ -165,7 +158,9 @@ def do_POST(self) -> None: ] else: arguments = ["memory", "stores", "list", "--page-size", "7"] - result = run_agentbricks("-o", "json", *arguments, env_override=env_override) + result = run_agentbricks( + *profile_args, "-o", "json", *arguments, env_override=env_override + ) assert result.stdout if command == "memory-list": assert reads.get(timeout=5).startswith("/api/2.0/agents/memory-stores") @@ -210,6 +205,9 @@ def do_POST(self) -> None: "name": "agentbricks memory stores list.page_size", "int_value": 7, } + if profile_source == "flag": + # Free-text values aren't logged, only that the option was passed. + expected_parameters["agentbricks.profile"] = {"name": "agentbricks.profile"} assert {parameter["name"]: parameter for parameter in log["parameters"]} == expected_parameters @@ -455,6 +453,9 @@ def test_tracing_unbind_then_bind_requires_an_experiment( assert "Pass --experiment-name or --experiment-id" in result.stderr -def test_logout_runs_cleanly(run_agentbricks) -> None: - # Isolated HOME: there's no saved selection to forget, but it must still exit cleanly. - run_agentbricks("logout") +def test_profile_get_reports_sdk_default_outside_a_project(run_agentbricks) -> None: + # No -p, no DATABRICKS_CONFIG_PROFILE, no project: resolution falls through to the SDK default. + result = run_agentbricks("--output", "json", "auth", "describe") + payload = json.loads(result.stdout) + assert payload["profile"] is None + assert payload["source"] == "Databricks SDK default" diff --git a/integrations/agentbricks/tests/functional/tracing_e2e_test.py b/integrations/agentbricks/tests/functional/tracing_e2e_test.py index a17abf503..2a81d3a5c 100644 --- a/integrations/agentbricks/tests/functional/tracing_e2e_test.py +++ b/integrations/agentbricks/tests/functional/tracing_e2e_test.py @@ -69,8 +69,8 @@ def local_trace_store(tmp_path: pathlib.Path) -> pathlib.Path: def _agentbricks_json(*args: str, project: pathlib.Path, home: pathlib.Path): """Run the real `agentbricks` CLI with `--output json`, isolated HOME, and return the parsed stdout. - Inherits PATH so the read command can find `uvx`, but points HOME at an empty dir (no agentbricks login / - real ~/.agentbricks) and silences MLflow's agent hint so stdout is clean JSON. + Inherits PATH so the read command can find `uvx`, but points HOME at an empty dir (no real token + cache) and silences MLflow's agent hint so stdout is clean JSON. """ agentbricks = pathlib.Path(sys.executable).with_name("agentbricks") if not agentbricks.is_file(): diff --git a/integrations/agentbricks/tests/unit_tests/apps_client_test.py b/integrations/agentbricks/tests/unit_tests/apps_client_test.py index 22927edfc..872a593b0 100644 --- a/integrations/agentbricks/tests/unit_tests/apps_client_test.py +++ b/integrations/agentbricks/tests/unit_tests/apps_client_test.py @@ -56,7 +56,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..02f4a0e88 100644 --- a/integrations/agentbricks/tests/unit_tests/auth_test.py +++ b/integrations/agentbricks/tests/unit_tests/auth_test.py @@ -1,193 +1,260 @@ -"""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 auth +from databricks_agentbricks.cli import app as cli +from databricks_agentbricks.cli import auth, workspace_auth +from databricks_agentbricks.cli.app import CliContext +from databricks_agentbricks.clients import api_client_provider +from databricks_agentbricks.projects.env_file import read_env -@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 interactive_terminal(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(workspace_auth, "_is_interactive", lambda: True) -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() +@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() + monkeypatch.setattr(api_client_provider, "_AgentBricksApiClient", client_cls) + return client_cls -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 _ctx(profile=None, output="text") -> CliContext: + return CliContext(profile, output) + + +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.use_project_profile(tmp_path) + flagged.use_project_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 == workspace_auth.ProfileInfo( + "from-env-var", "DATABRICKS_CONFIG_PROFILE" ) + assert project.profile_info == workspace_auth.ProfileInfo("proj", ".env") + assert flagged.profile_info == workspace_auth.ProfileInfo("from-flag", "--profile") -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_profile_flag_after_subcommand_wins(env_project): + payload = _get_json(["-p", "before", "auth", "describe", "-p", "after"]) - auth._run_databricks_login("prof") + assert (payload["profile"], payload["source"]) == ("after", "--profile") - databricks_login.assert_called_once_with( - ["databricks", "auth", "login", "--profile", "prof"], - text=True, - check=False, - stdout=sys.stderr, - ) +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") -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")] + monkeypatch.setenv("DATABRICKS_CONFIG_PROFILE", "ambient") + payload = _get_json(["auth", "describe"]) + assert (payload["profile"], payload["source"], payload["project"]) == ( + "ambient", + "DATABRICKS_CONFIG_PROFILE", + None, ) - 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_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() -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")) - ) - monkeypatch.setattr(auth, "_is_interactive", lambda: True) - databricks_login = mock.Mock() - monkeypatch.setattr(auth.subprocess, "run", databricks_login) + captured = capsys.readouterr() + assert captured.out == "" + assert captured.err.count("Using profile 'ws' (from --profile) · https://ws.example") == 1 - 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_workspace_notice_names_sdk_default_auth(capsys, monkeypatch, fake_api_client): + monkeypatch.setenv("DATABRICKS_HOST", "https://env.example") + _ctx().client() -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")) + assert ( + "Using Databricks SDK default auth · https://env.example (from DATABRICKS_HOST)" + in capsys.readouterr().err ) - 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_workspace_notice_is_silent_for_json(capsys, fake_api_client): + _ctx("ws", output="json").client() + + assert capsys.readouterr().err == "" + + +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_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_set_signs_in_when_interactive(project, interactive_terminal, validated_profile): + result = _invoke_profile(["switch", "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_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 +): + 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) + databricks_login = mock.Mock(return_value=mock.Mock(returncode=0)) + monkeypatch.setattr(workspace_auth.subprocess, "run", databricks_login) -def test_login_reports_when_databricks_cli_is_missing(tmp_path, monkeypatch): - monkeypatch.setenv("AGENTBRICKS_CONFIG_HOME", str(tmp_path)) + 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(monkeypatch): monkeypatch.setattr( - auth, "_AgentBricksApiClient", mock.Mock(side_effect=auth.AgentCliError("no credentials")) + workspace_auth, "_validate_bounded", lambda profile: (None, Unauthenticated("expired")) ) - monkeypatch.setattr(auth, "_is_interactive", lambda: True) - monkeypatch.setattr(auth.subprocess, "run", mock.Mock(side_effect=FileNotFoundError)) + databricks_login = mock.Mock() + monkeypatch.setattr(workspace_auth.subprocess, "run", databricks_login) - result = CliRunner().invoke(auth.login, ["--profile", "prof"], obj=_Ctx()) + 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 df4898d4d..51bd6e143 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 "mcp" not in names assert "pipeline" not in names assert "durability" not in names @@ -166,8 +169,8 @@ def test_help_examples_recommend_the_default_happy_path(): runner = CliRunner() expected_examples = { (): ( - "agentbricks login --profile ", - "agentbricks init my-agent", + "agentbricks auth login ", + "agentbricks init my-agent --profile ", "cd my-agent", "agentbricks dev", "agentbricks deploy my-agent", @@ -275,7 +278,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 @@ -320,12 +323,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 my-agent" in numbered[4] @@ -345,9 +348,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(): @@ -366,7 +369,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 @@ -399,6 +402,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 cb344f3e0..86e4252dd 100644 --- a/integrations/agentbricks/tests/unit_tests/client_test.py +++ b/integrations/agentbricks/tests/unit_tests/client_test.py @@ -756,7 +756,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 diff --git a/integrations/agentbricks/tests/unit_tests/conftest.py b/integrations/agentbricks/tests/unit_tests/conftest.py index c95115734..7a497516e 100644 --- a/integrations/agentbricks/tests/unit_tests/conftest.py +++ b/integrations/agentbricks/tests/unit_tests/conftest.py @@ -2,11 +2,60 @@ 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 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(workspace_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(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 3ef7d2c5c..a0e978aa3 100644 --- a/integrations/agentbricks/tests/unit_tests/deploy_bugfix_test.py +++ b/integrations/agentbricks/tests/unit_tests/deploy_bugfix_test.py @@ -19,6 +19,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 --------------------------- @@ -74,6 +75,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 060dded7d..11f4b1384 100644 --- a/integrations/agentbricks/tests/unit_tests/deploy_test.py +++ b/integrations/agentbricks/tests/unit_tests/deploy_test.py @@ -18,6 +18,8 @@ from click.testing import CliRunner from databricks_agentbricks.cli import deploy as deploy_mod +from databricks_agentbricks.cli.app import CliContext +from databricks_agentbricks.cli.workspace_auth import ProfileInfo 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 ( @@ -53,15 +55,17 @@ class _FakeApiClient: current_user = "me@example.com" -class _Ctx: - profile = "prof" - output = "text" - - def __init__(self, *, output: str = "text"): - self.output = output +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()) +@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: return DeployResult( deployment=deployment, @@ -169,11 +173,9 @@ def __call__(self, args, profile, **kwargs): 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) + ctx.api_client_provider = types.SimpleNamespace(get=lambda: api) + return ctx def test_manifest_env_scaffolds_when_missing(tmp_path: pathlib.Path): @@ -372,6 +374,31 @@ def test_deploy_cli_passes_request_to_service_and_presents_json(tmp_path: pathli assert json.loads(result.output)["deployment"] == "agent-bricks-demo" +@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 / ".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://{expected}.example" in out + + def test_deploy_json_preserves_additive_tool_access_contract(tmp_path: pathlib.Path): service = mock.Mock() service.deploy.return_value = replace( diff --git a/integrations/agentbricks/tests/unit_tests/dev_service_test.py b/integrations/agentbricks/tests/unit_tests/dev_service_test.py index c88e214e3..f2dbd8898 100644 --- a/integrations/agentbricks/tests/unit_tests/dev_service_test.py +++ b/integrations/agentbricks/tests/unit_tests/dev_service_test.py @@ -119,6 +119,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 43b58cc74..4027882d2 100644 --- a/integrations/agentbricks/tests/unit_tests/dev_test.py +++ b/integrations/agentbricks/tests/unit_tests/dev_test.py @@ -11,6 +11,8 @@ from click.testing import CliRunner from databricks_agentbricks.cli import dev as dev_mod +from databricks_agentbricks.cli import workspace_auth +from databricks_agentbricks.cli.app import CliContext from databricks_agentbricks.errors import AgentCliError from databricks_agentbricks.projects.agent_project import AgentProject, ToolSpec from databricks_agentbricks.projects.config import write_project_metadata @@ -34,10 +36,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" @@ -48,6 +49,11 @@ def client(self): return mock.Mock(current_user="me@example.com", host="https://my-workspace.databricks.com") +@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 (via cli.tracing) for Agent Bricks projects; @@ -346,12 +352,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") @@ -630,3 +636,89 @@ def test_dev_hides_active_venv_and_port_from_run_local(tmp_path: pathlib.Path, m assert "VIRTUAL_ENV" not in env assert "PORT" not in env 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://proj.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): + _write_project(tmp_path, "DATABRICKS_CONFIG_PROFILE=proj\n") + monkeypatch.setattr( + workspace_auth, "_validate_bounded", lambda profile: (None, 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() + assert not (tmp_path / "app.agentbricksdev.yaml").exists() diff --git a/integrations/agentbricks/tests/unit_tests/endpoint_test.py b/integrations/agentbricks/tests/unit_tests/endpoint_test.py index 18f7de45a..fba1ea051 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..bd4a2b53d --- /dev/null +++ b/integrations/agentbricks/tests/unit_tests/env_file_test.py @@ -0,0 +1,66 @@ +"""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_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 c47a25ace..a221bc38c 100644 --- a/integrations/agentbricks/tests/unit_tests/init_test.py +++ b/integrations/agentbricks/tests/unit_tests/init_test.py @@ -15,17 +15,19 @@ 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 import workspace_auth +from databricks_agentbricks.cli.app import CliContext from databricks_agentbricks.errors import AgentCliError from databricks_agentbricks.projects.types import AgentFramework, AgentServer -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: @@ -278,56 +280,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 +): + monkeypatch.setenv("DATABRICKS_HOST", "https://env.example") + monkeypatch.setattr(workspace_auth, "_is_interactive", lambda: True) 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, monkeypatch, validated_profile): + monkeypatch.setattr(workspace_auth, "_is_interactive", lambda: True) + 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 7ddd8f362..b6a28fda1 100644 --- a/integrations/agentbricks/tests/unit_tests/runtime_store_deploy_test.py +++ b/integrations/agentbricks/tests/unit_tests/runtime_store_deploy_test.py @@ -36,6 +36,7 @@ def _ctx(client, *, output="text"): profile="prof", output=output, api_client_provider=types.SimpleNamespace(get=lambda: client), + announce_workspace=lambda: None, ) diff --git a/integrations/agentbricks/tests/unit_tests/tracing_test.py b/integrations/agentbricks/tests/unit_tests/tracing_test.py index b6a4d703a..4dc97ea99 100644 --- a/integrations/agentbricks/tests/unit_tests/tracing_test.py +++ b/integrations/agentbricks/tests/unit_tests/tracing_test.py @@ -44,6 +44,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..db7dcf206 --- /dev/null +++ b/integrations/agentbricks/tests/unit_tests/workspace_auth_test.py @@ -0,0 +1,215 @@ +"""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 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 workspace_auth +from databricks_agentbricks.errors import AgentCliError + +_PROF_CFG = "[prof]\nhost = https://prof.example\n" + + +@pytest.fixture(autouse=True) +def _disable_telemetry(monkeypatch): + monkeypatch.setenv("AGENTBRICKS_DISABLE_TELEMETRY", "1") + + +def _set_interactive(monkeypatch, interactive: bool = True) -> None: + monkeypatch.setattr(workspace_auth, "_is_interactive", lambda: interactive) + + +def _fail_validation(monkeypatch, error: Exception) -> None: + monkeypatch.setattr(workspace_auth, "_validate_bounded", lambda profile: (None, error)) + + +def _stub_databricks_cli(monkeypatch, **mock_kwargs) -> mock.Mock: + databricks_login = mock.Mock(**mock_kwargs) + monkeypatch.setattr(workspace_auth.subprocess, "run", databricks_login) + return databricks_login + + +# --- 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 workspace_auth.resolve_profile(flag, tmp_path) == workspace_auth.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 workspace_auth.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"), + ("[DEFAULT]\nhost = https://d.example\n", None, None, "default"), + ], + ids=["named-profile", "databricks-host", "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 workspace_auth.workspace_for_profile(profile) == expected + + +# --- preflight_auth -------------------------------------------------------------------------- + + +def test_preflight_passes_when_profile_validates(validated_profile): + workspace_auth.preflight_auth("prof", "--profile") + + validated_profile.assert_called_once_with("prof") + + +def test_preflight_failure_hints_at_login_and_retry(monkeypatch): + _fail_validation(monkeypatch, 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_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, monkeypatch +): + databricks_login = _stub_databricks_cli(monkeypatch) + + assert ( + workspace_auth.authenticate_profile("prof", "argument") == validated_profile.return_value[0] + ) + databricks_login.assert_not_called() + + +def test_authenticate_profile_logs_in_then_revalidates(monkeypatch): + client = mock.Mock() + validate = mock.Mock( + side_effect=[(None, Unauthenticated("expired")), ((client, "me@example.com"), None)] + ) + monkeypatch.setattr(workspace_auth, "_validate_bounded", validate) + _set_interactive(monkeypatch) + databricks_login = _stub_databricks_cli(monkeypatch, return_value=mock.Mock(returncode=0)) + + assert workspace_auth.authenticate_profile("prof", "argument") == (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_authenticate_profile_reports_a_missing_databricks_cli(monkeypatch): + _fail_validation(monkeypatch, Unauthenticated("expired")) + _set_interactive(monkeypatch) + _stub_databricks_cli(monkeypatch, side_effect=FileNotFoundError) + + with pytest.raises(AgentCliError, match="`databricks` CLI was not found"): + workspace_auth.authenticate_profile("prof", "argument") + + +# 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, "databricks auth login --profile prof"), + (Unauthenticated("expired"), True, True, "databricks auth login --profile prof"), + ], + ids=["non-auth-error", "no-terminal", "ci"], +) +def test_authenticate_profile_raises_without_opening_a_login( + monkeypatch, error, ci, interactive, hint +): + if ci: + monkeypatch.setenv("CI", "true") + _set_interactive(monkeypatch, interactive) + _fail_validation(monkeypatch, error) + databricks_login = _stub_databricks_cli(monkeypatch) + + with pytest.raises(AgentCliError) as raised: + workspace_auth.authenticate_profile("prof", ".env") + + 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() From 7fde8f8a3a67661761481dfc2c95a4445c85f5c1 Mon Sep 17 00:00:00 2001 From: Elaine Wang Date: Wed, 14 Oct 2026 23:01:04 +0000 Subject: [PATCH 02/12] [agentbricks] Move workspace auth into a service and consolidate profile lookups Co-authored-by: Isaac --- integrations/agentbricks/cli.md | 4 +- .../src/databricks_agentbricks/cli/app.py | 47 ++- .../src/databricks_agentbricks/cli/auth.py | 93 +++--- .../src/databricks_agentbricks/cli/deploy.py | 16 +- .../src/databricks_agentbricks/cli/dev.py | 27 +- .../src/databricks_agentbricks/cli/init.py | 71 ++-- .../cli/workspace_auth.py | 302 ------------------ .../clients/api_client_provider.py | 5 + .../clients/apps_client.py | 3 +- .../clients/databricks_cli.py | 20 +- .../presentation/profile.py | 12 +- .../projects/env_file.py | 12 +- .../projects/profiles.py | 204 ++++++++++++ .../services/dev_service.py | 26 +- .../services/workspace_auth.py | 189 +++++++++++ .../agentbricks/tests/unit_tests/auth_test.py | 31 +- .../agentbricks/tests/unit_tests/conftest.py | 30 +- .../tests/unit_tests/deploy_test.py | 10 +- .../agentbricks/tests/unit_tests/dev_test.py | 9 +- .../agentbricks/tests/unit_tests/init_test.py | 9 +- .../tests/unit_tests/workspace_auth_test.py | 74 ++--- 21 files changed, 672 insertions(+), 522 deletions(-) delete mode 100644 integrations/agentbricks/src/databricks_agentbricks/cli/workspace_auth.py create mode 100644 integrations/agentbricks/src/databricks_agentbricks/projects/profiles.py create mode 100644 integrations/agentbricks/src/databricks_agentbricks/services/workspace_auth.py diff --git a/integrations/agentbricks/cli.md b/integrations/agentbricks/cli.md index 5e8049d1f..a6e8b108a 100644 --- a/integrations/agentbricks/cli.md +++ b/integrations/agentbricks/cli.md @@ -98,7 +98,7 @@ 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 — 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. -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 `DATABRICKS_HOST` set (env-var auth, 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 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. @@ -132,7 +132,7 @@ _Options_ | --- | --- | --- | --- | --- | | `--framework ` | `langgraph` \| `openai` | - | 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 ` (`-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 DATABRICKS_HOST set, no .env is written). | +| `--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: derived from the directory, -memory). Only --server agentbricks declares stores by default. | diff --git a/integrations/agentbricks/src/databricks_agentbricks/cli/app.py b/integrations/agentbricks/src/databricks_agentbricks/cli/app.py index 4aae1ae7b..16f9885a4 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/cli/app.py +++ b/integrations/agentbricks/src/databricks_agentbricks/cli/app.py @@ -23,13 +23,6 @@ from databricks_agentbricks.cli.sessions import sessions from databricks_agentbricks.cli.tools import tools from databricks_agentbricks.cli.tracing import tracing -from databricks_agentbricks.cli.workspace_auth import ( - SOURCE_FLAG, - SOURCE_SDK_DEFAULT, - ProfileInfo, - legacy_default_profile, - resolve_profile, -) from databricks_agentbricks.clients.api_client_provider import ApiClientProvider from databricks_agentbricks.presentation.help import configure_help from databricks_agentbricks.presentation.profile import ( @@ -37,6 +30,16 @@ warn_legacy_default_profile, ) from databricks_agentbricks.projects.config import is_agentbricks_project +from databricks_agentbricks.projects.env_file import read_env +from databricks_agentbricks.projects.profiles import ( + SOURCE_FLAG, + SOURCE_SDK_DEFAULT, + ProfileInfo, + auth_login_command, + auth_switch_command, + legacy_default_profile, + resolve_profile, +) from databricks_agentkit._api_client import _AgentBricksApiClient @@ -47,29 +50,39 @@ def __init__(self, profile: Optional[str], output: str): self.profile_info = resolve_profile(profile) self.output = output self._workspace_announced = False + self._project_dir: Optional[pathlib.Path] = None + self._project_env: dict[str, str] = {} self.api_client_provider = self._client_provider(self.profile) @property def profile(self) -> Optional[str]: return self.profile_info.name - def use_project_profile( - self, project_dir: pathlib.Path, env_values: Optional[dict[str, str]] = None - ) -> None: + def use_project_profile(self, project_dir: pathlib.Path) -> None: """Fold in the project's `.env` profile (see `resolve_profile`) unless `-p` was given. - The root group calls this for the current directory when it's an Agent Bricks project; - `dev` and `deploy` call it again with their source dir, whose `.env` then replaces the + The root group calls this for the current directory when it's an Agent Bricks project, and + commands with `--source` call it again for that project, whose `.env` then replaces the current directory's. Only an explicit `-p` outranks the `.env`. - ``env_values`` passes an already-parsed ``/.env`` so it isn't read twice. """ + directory = project_dir.resolve() + if directory != self._project_dir: + self._project_dir = directory + self._project_env = read_env(directory / ".env") if self.profile_info.source == SOURCE_FLAG: return - resolved = resolve_profile(None, project_dir, env_values) + resolved = resolve_profile(None, directory, self._project_env) if resolved.name != self.profile_info.name: self.api_client_provider = self._client_provider(resolved.name) self.profile_info = resolved + def project_env(self, project_dir: pathlib.Path) -> dict[str, str]: + """`project_dir`'s parsed `.env`, reusing the one `use_project_profile` already read.""" + directory = project_dir.resolve() + if directory == self._project_dir: + return self._project_env + return read_env(directory / ".env") + def use_flag_profile(self, profile: str) -> None: """Apply a `-p/--profile` given after the subcommand; like the global one, it beats `.env`.""" if profile != self.profile: @@ -150,9 +163,9 @@ def agentbricks(ctx: click.Context, profile: Optional[str], output: str) -> None _REMOVED_COMMANDS = { - "login": "Sign in with `agentbricks auth login `, and pick a project's profile with " - "`agentbricks auth switch `.", - "logout": "Profiles are now chosen per project in its .env (`agentbricks auth switch `); " + "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.", } diff --git a/integrations/agentbricks/src/databricks_agentbricks/cli/auth.py b/integrations/agentbricks/src/databricks_agentbricks/cli/auth.py index 72d18bfa9..47e56cdce 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/cli/auth.py +++ b/integrations/agentbricks/src/databricks_agentbricks/cli/auth.py @@ -2,22 +2,14 @@ from __future__ import annotations +import os import pathlib +import sys from typing import Optional import click -from databricks_agentbricks.cli.workspace_auth import ( - SOURCE_ENV_FILE, - SOURCE_FLAG, - ProfileInfo, - authenticate_profile, - can_prompt_login, - profile_host, - resolve_profile, - run_databricks_login, - workspace_for_profile, -) +from databricks_agentbricks.clients.databricks_cli import run_databricks_login from databricks_agentbricks.errors import AgentCliError from databricks_agentbricks.presentation import render from databricks_agentbricks.presentation.profile import ( @@ -27,22 +19,44 @@ ) from databricks_agentbricks.projects.config import is_agentbricks_project from databricks_agentbricks.projects.env_file import read_env, update_env_profile +from databricks_agentbricks.projects.profiles import ( + SOURCE_ENV_FILE, + SOURCE_FLAG, + ProfileInfo, + auth_login_command, + auth_switch_command, + 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 sign_in_if_interactive(obj, profile: str, source: str) -> Optional[str]: +def _is_interactive() -> bool: + return sys.stdin.isatty() + + +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 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. Interactive terminals only: a scripted or CI run never opens a browser or waits on the - workspace. A failed sign-in is reported, not raised, since the profile choice still stands. + 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) + _, user = authenticate_profile(profile, source, interactive=True) except AgentCliError as exc: warn_sign_in_failed(profile, exc, obj) return None @@ -60,14 +74,16 @@ def _require_project(source: str) -> pathlib.Path: return directory -def _resolve_for(obj, directory: pathlib.Path) -> ProfileInfo: +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 is_agentbricks_project(directory) else None) + return resolve_profile(None, directory if in_project else None) -def _login_target(obj, name: Optional[str], directory: pathlib.Path) -> Optional[tuple[str, str]]: +def _login_target( + obj, name: Optional[str], 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 @@ -77,16 +93,14 @@ def _login_target(obj, name: Optional[str], directory: pathlib.Path) -> Optional return name, "argument" if obj.profile_info.source == SOURCE_FLAG and obj.profile_info.name: return obj.profile_info.name, SOURCE_FLAG - if is_agentbricks_project(directory): - env_profile = read_env(directory / ".env").get("DATABRICKS_CONFIG_PROFILE") - if env_profile: - return env_profile, SOURCE_ENV_FILE + if env_profile: + return env_profile, SOURCE_ENV_FILE if can_prompt_login(): return None raise AgentCliError( "No profile to log in to.", - hint="Pass one (`agentbricks auth login `), or set the project's profile " - "with `agentbricks auth switch `.", + hint=f"Pass one (`{auth_login_command(None)}`), or set the project's profile " + f"with `{auth_switch_command(None)}`.", ) @@ -121,10 +135,7 @@ def auth_switch(obj, name: str, source: str) -> None: """ directory = _require_project(source) changed = update_env_profile(directory, name) - if not can_prompt_login(): - # Interactive runs sign in next, which creates a missing profile. - warn_unknown_profile(name, obj) - user = sign_in_if_interactive(obj, name, "argument") + user = sign_in_chosen_profile(obj, name, "argument") host = profile_host(name) if obj.output == "json": render.emit_json( @@ -145,7 +156,7 @@ def auth_switch(obj, name: str, source: str) -> None: render.success( f"Project profile set → '{name}'" if changed else f"Project profile is already '{name}'", fields=fields, - next_steps=None if user else [(f"agentbricks auth login {name}", "Sign in")], + next_steps=None if user else [(auth_login_command(name), "Sign in")], ) @@ -165,9 +176,10 @@ def auth_describe(obj, source: str) -> None: authentication. Read-only; no workspace calls are made. """ directory = pathlib.Path(source).resolve() - info = _resolve_for(obj, directory) + in_project = is_agentbricks_project(directory) + info = _resolve_for(obj, directory, in_project) host, workspace = workspace_for_profile(info.name) - project = str(directory) if is_agentbricks_project(directory) else None + project = str(directory) if in_project else None if obj.output == "json": render.emit_json( { @@ -211,7 +223,12 @@ def auth_login(obj, name: Optional[str], source: str) -> None: instead. Nothing is saved: use `agentbricks auth switch` to change the project's profile. """ directory = pathlib.Path(source).resolve() - target = _login_target(obj, name, directory) + env_profile = ( + read_env(directory / ".env").get("DATABRICKS_CONFIG_PROFILE") + if is_agentbricks_project(directory) + else None + ) + target = _login_target(obj, name, env_profile) if target is None: # `databricks auth login` picks the profile; it doesn't say which, so there's nothing to # validate or report beyond its own output. @@ -223,7 +240,7 @@ def auth_login(obj, name: Optional[str], source: str) -> None: return chosen, origin = target announce_login(obj, chosen, origin) - client, user = authenticate_profile(chosen, origin) + client, user = authenticate_profile(chosen, origin, interactive=can_prompt_login()) if obj.output == "json": render.emit_json({"profile": chosen, "source": origin, "user": user, "host": client.host}) return @@ -231,10 +248,8 @@ def auth_login(obj, name: Optional[str], source: str) -> None: f"Logged in to '{chosen}'", fields={"Profile": chosen, "Source": origin, "User": user, "Workspace": client.host}, ) - if is_agentbricks_project(directory): - env_profile = read_env(directory / ".env").get("DATABRICKS_CONFIG_PROFILE") - if env_profile and env_profile != chosen: - render.console().print( - f"[dim]This project still uses profile '{env_profile}' (from .env). " - f"Run `agentbricks auth switch {chosen}` to switch it.[/]" - ) + if env_profile and env_profile != chosen: + render.console().print( + f"[dim]This project still uses profile '{env_profile}' (from .env). " + f"Run `{auth_switch_command(chosen)}` to switch it.[/]" + ) diff --git a/integrations/agentbricks/src/databricks_agentbricks/cli/deploy.py b/integrations/agentbricks/src/databricks_agentbricks/cli/deploy.py index 85fb1a25f..ecbd5dde8 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/cli/deploy.py +++ b/integrations/agentbricks/src/databricks_agentbricks/cli/deploy.py @@ -18,11 +18,8 @@ from __future__ import annotations -import pathlib - import click -from databricks_agentbricks.cli.workspace_auth import preflight_auth from databricks_agentbricks.clients.apps_client import AppsClient from databricks_agentbricks.clients.apps_user_auth_client import AppsUserAuthClient from databricks_agentbricks.clients.conversation_store_client import ( @@ -63,11 +60,11 @@ TracingProvisioner, ) from databricks_agentbricks.services.deployment.tool_access_provisioner import ToolAccessProvisioner +from databricks_agentbricks.services.workspace_auth import preflight_auth def build_deploy_service(obj) -> DeployService: """Compose a deployment service for this CLI invocation without opening an API client.""" - obj.announce_workspace() api_client_provider = obj.api_client_provider apps_client = AppsClient(obj.profile, runner=_databricks) apps_user_auth_client = AppsUserAuthClient(obj.profile) @@ -166,13 +163,10 @@ def deploy( not authentication and not the session id itself (send the session id in the request body): X-Routing-Key: """ - # Fold the project's `.env` into the profile resolution before the profile is used, matching - # `agentbricks dev` so both commands deploy and run against the same workspace. - obj.use_project_profile(pathlib.Path(source)) obj.announce_workspace() # 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. - preflight_auth(obj.profile, obj.profile_info.source) + obj.api_client_provider.use(preflight_auth(obj.profile, obj.profile_info.source)) request = DeployRequest( name=name, source=source, @@ -197,6 +191,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) @@ -205,6 +200,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) @@ -213,6 +209,7 @@ def deployments_get(obj, name) -> None: @click.pass_obj def deployments_logs(obj, name) -> None: """Stream a deployment's logs.""" + obj.announce_workspace() build_deploy_service(obj).logs(name) @@ -221,6 +218,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) @@ -235,6 +233,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) @@ -245,6 +244,7 @@ def deployments_stop(obj, name, yes) -> None: @click.pass_obj def deployments_delete(obj, name, yes) -> None: """Delete a deployment and, when managed provisioning is enabled, its Runtime Store.""" + obj.announce_workspace() service = build_deploy_service(obj) target = f"deployment '{name}'" if not yes and service.deletes_runtime_store_data(): diff --git a/integrations/agentbricks/src/databricks_agentbricks/cli/dev.py b/integrations/agentbricks/src/databricks_agentbricks/cli/dev.py index 0625bf0be..827133d8f 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/cli/dev.py +++ b/integrations/agentbricks/src/databricks_agentbricks/cli/dev.py @@ -11,12 +11,11 @@ import pathlib import subprocess -from typing import Callable, Optional +from typing import Optional import click from databricks_agentbricks.cli.tracing import start_local_tracing_server, stop_local_tracing_server -from databricks_agentbricks.cli.workspace_auth import preflight_auth from databricks_agentbricks.clients.apps_client import AppsClient from databricks_agentbricks.clients.databricks_cli import _databricks from databricks_agentbricks.presentation.dev import ( @@ -27,13 +26,13 @@ from databricks_agentbricks.presentation.dev import ( announce_local_url as _announce_local_url, ) -from databricks_agentbricks.projects.env_file import read_env from databricks_agentbricks.projects.resolver import ProjectResolver from databricks_agentbricks.services.dev_service import ( DevRequest, DevService, _dev_entry_point, # noqa: F401 - compatibility re-export for existing callers ) +from databricks_agentbricks.services.workspace_auth import preflight_auth class _CliLocalTracing: @@ -61,13 +60,12 @@ def _env_credential(env_file: dict[str, str]) -> Optional[str]: return None -def build_dev_service(obj, preflight: Optional[Callable[[], None]] = None) -> DevService: +def build_dev_service(obj) -> DevService: """Compose the local workflow without opening a workspace API client.""" return DevService( project_resolver=ProjectResolver(), apps_client=AppsClient(obj.profile, runner=_databricks), local_tracing=_CliLocalTracing(), - preflight=preflight, ) @@ -116,11 +114,9 @@ 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. """ - source_dir = pathlib.Path(source) - # The CLI and the locally running agent read the same `.env`, so resolve the profile from it - # once here (dev also consults it below for explicit credentials). - env_file = read_env(source_dir / ".env") - obj.use_project_profile(source_dir, env_file) + # The project's `.env` profile is already applied (see `_use_source_project_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). @@ -131,16 +127,17 @@ def dev( # beats `.env` because the template loads dotenv with override=False. local_env["DATABRICKS_CONFIG_PROFILE"] = obj.profile - def _preflight() -> None: - if not env_credential: - preflight_auth(obj.profile, obj.profile_info.source) - # Announced before the preflight so an auth failure still names its target. When `.env` # credentials win, the resolved profile would name a different workspace than the agent # actually uses, so it isn't announced. if not uses_env_credentials: obj.announce_workspace() - service = build_dev_service(obj, _preflight) + 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: + preflight_auth(obj.profile, obj.profile_info.source) with service.prepare( DevRequest( source=source, diff --git a/integrations/agentbricks/src/databricks_agentbricks/cli/init.py b/integrations/agentbricks/src/databricks_agentbricks/cli/init.py index 140d39623..5169a4b60 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/cli/init.py +++ b/integrations/agentbricks/src/databricks_agentbricks/cli/init.py @@ -26,27 +26,29 @@ import click -from databricks_agentbricks.cli.auth import sign_in_if_interactive +from databricks_agentbricks.cli.auth import sign_in_chosen_profile from databricks_agentbricks.cli.tracing import default_experiment_name -from databricks_agentbricks.cli.workspace_auth import ( - SOURCE_ENV_VAR, - SOURCE_FLAG, - SOURCE_SETTINGS_DEFAULT, - can_prompt_login, - configured_profiles, - profile_exists, - profile_host, - settings_default_profile, -) from databricks_agentbricks.errors import AgentCliError from databricks_agentbricks.presentation import render -from databricks_agentbricks.presentation.profile import note_choose_profile, warn_unknown_profile +from databricks_agentbricks.presentation.profile import note_choose_profile from databricks_agentbricks.projects.agent_project import AgentProject, default_store_name 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, @@ -154,20 +156,21 @@ def _report_already_initialized(obj, dest: pathlib.Path) -> None: return render.console().print( f"{dest} is already an Agent Bricks project; nothing was changed. Run `agentbricks dev` / " - f"`agentbricks deploy`, or `agentbricks auth switch --source {dest}` to change " + 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 DATABRICKS_HOST is set, else the default - chosen with `databricks auth switch`, else `DEFAULT` (the order the SDK itself falls back in). + 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 when `DATABRICKS_HOST` is set because env-var auth (CI, - containers) already names the workspace and a `.env` profile would override it. + 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 @@ -176,12 +179,9 @@ def _project_profile(obj, flag_profile: Optional[str]) -> Optional[tuple[str, st env_profile = os.environ.get("DATABRICKS_CONFIG_PROFILE") if env_profile: return env_profile, SOURCE_ENV_VAR - if os.environ.get("DATABRICKS_HOST"): + if env_auth_configured(): return None - switched = settings_default_profile() - if switched: - return switched, SOURCE_SETTINGS_DEFAULT - return "DEFAULT", "default" + return sdk_default_profile() def _migration_paths(dest: pathlib.Path, target: pathlib.Path) -> tuple[pathlib.Path, ...]: @@ -365,7 +365,7 @@ def _prepare_migration( default=None, 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 DATABRICKS_HOST set, no .env is written).", + "profile and env-var auth set (DATABRICKS_HOST, DATABRICKS_TOKEN or DATABRICKS_CLIENT_ID), no .env is written).", ) @click.option( "--disable-chat-app", @@ -416,8 +416,8 @@ def init( 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 DATABRICKS_HOST set (env-var auth, e.g. CI) and no profile, no `.env` is - written and commands keep using those environment variables. In an interactive terminal, + `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 @@ -515,18 +515,19 @@ def init( # 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] == "default" and not profile_exists("DEFAULT") + 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) - elif env_profile and not can_prompt_login(): - # Interactive runs sign in right after scaffolding, which creates a missing profile. - warn_unknown_profile(env_profile, obj) # A failed sign-in doesn't undo the scaffold; the summary says how to retry. signed_in_user = ( - sign_in_if_interactive(obj, *project_profile) + sign_in_chosen_profile(obj, *project_profile) if project_profile and not other_profiles else None ) @@ -574,18 +575,16 @@ def init( fields["Profile (.env)"] = env_profile fields["Workspace"] = workspace_host or "not configured yet" else: - fields["Auth"] = "environment variables (DATABRICKS_HOST); no .env written" + 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( - ("agentbricks auth switch ", f"Pick one of: {', '.join(other_profiles)}") - ) + steps.append((auth_switch_command(None), f"Pick one of: {', '.join(other_profiles)}")) elif env_profile: - steps.append((f"agentbricks auth login {env_profile}", "Sign in before running the agent")) + steps.append((auth_login_command(env_profile), "Sign in before running the agent")) if not other_profiles: - steps.append(("agentbricks auth switch ", "Use a different profile (optional)")) + 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/workspace_auth.py b/integrations/agentbricks/src/databricks_agentbricks/cli/workspace_auth.py deleted file mode 100644 index 60dc2aacc..000000000 --- a/integrations/agentbricks/src/databricks_agentbricks/cli/workspace_auth.py +++ /dev/null @@ -1,302 +0,0 @@ -"""Profile resolution, workspace lookup, and authentication checks for the Agent Bricks CLI.""" - -from __future__ import annotations - -import configparser -import dataclasses -import json -import os -import pathlib -import subprocess -import sys -import threading -from typing import Optional - -from databricks_agentbricks.errors import AgentCliError -from databricks_agentbricks.projects.env_file import read_env -from databricks_agentkit._api_client import ( - _AgentBricksApiClient, - _databricks_config_parser, - _profile_host, - _profile_section, -) - -# 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" - - -@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 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_host(profile) - - -def profile_exists(profile: str) -> bool: - """Whether the profile is configured in the Databricks config file; never raises.""" - parser = _databricks_config_parser() - if profile == parser.default_section: - # configparser never reports [DEFAULT] via has_section(); it exists once it has any keys. - return bool(parser.defaults()) - return parser.has_section(profile) - - -def workspace_for_profile(profile: Optional[str]) -> tuple[Optional[str], Optional[str]]: - """The workspace a profile points at, as (host, display text); (None, None) when unknown. - - For display only: without a profile the text names where the host came from - (DATABRICKS_HOST, else the config file's `[DEFAULT]`), because "no workspace" is the answer - the caller needs to surface. It does not claim to mirror the SDK's full default resolution. - """ - if profile: - host = profile_host(profile) - return host, host - env_host = os.environ.get("DATABRICKS_HOST") - if env_host: - return env_host, f"{env_host} (from DATABRICKS_HOST)" - default_host = profile_host(_databricks_config_parser().default_section) - if default_host: - return default_host, f"{default_host} ([DEFAULT] profile)" - return None, "none configured" - - -def _validate_profile(profile: Optional[str]) -> tuple[_AgentBricksApiClient, str]: - client = _AgentBricksApiClient(profile) - return client, client.current_user - - -def _is_interactive() -> bool: - return sys.stdin.isatty() - - -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 _is_auth_error(error: Optional[Exception]) -> bool: - """Whether a validation failure means the profile isn't signed in, so a login could fix it.""" - from databricks.sdk.errors import Unauthenticated - - return isinstance(error, (AgentCliError, Unauthenticated)) - - -def _is_unreachable(error: Exception) -> bool: - """Whether validation failed on the network or the server, so signing in wouldn't fix it.""" - from databricks.sdk.errors import DeadlineExceeded, InternalError, TemporarilyUnavailable - from requests.exceptions import ConnectionError as RequestsConnectionError - - return isinstance( - error, (RequestsConnectionError, InternalError, TemporarilyUnavailable, DeadlineExceeded) - ) - - -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 - - -# 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 _preflight_error(profile: Optional[str], source: str) -> AgentCliError: - if profile: - return AgentCliError( - f"Databricks profile {profile!r} (from {source}) isn't authenticated.", - hint=f"Run `agentbricks auth login {profile}` to sign in, then rerun the command.", - ) - return AgentCliError( - "No Databricks profile is configured and the Databricks SDK's default authentication " - "didn't authenticate.", - hint="Run `agentbricks auth login ` to sign in, then pass `-p ` or " - "choose one for the project with `agentbricks auth switch ` (or " - "DATABRICKS_CONFIG_PROFILE), and rerun the command.", - ) - - -SOURCE_SETTINGS_DEFAULT = "databricks auth switch" - - -def settings_default_profile() -> Optional[str]: - """The default profile chosen with `databricks auth switch` ([__settings__] default_profile).""" - return _profile_section("__settings__").get("default_profile", "").strip() or None - - -def configured_profiles() -> list[str]: - """Profile names in the Databricks config file, `[DEFAULT]` included when it has keys.""" - parser = _databricks_config_parser() - names = [name for name in parser.sections() if name != "__settings__"] - return ([parser.default_section] if parser.defaults() else []) + names - - -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 and any( - os.environ.get(name) - for name in ("DATABRICKS_HOST", "DATABRICKS_TOKEN", "DATABRICKS_CLIENT_ID") - ): - return None # env-var auth without a profile: the SDK never reads the config file - if not profile: - profile = settings_default_profile() or configparser.DEFAULTSECT - return _profile_section(profile).get("auth_type") - - -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`. - - A browser-OAuth flow can block forever, so the call is bounded. Daemon because executor - threads are joined at interpreter exit and a hung flow 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] - - -def preflight_auth(profile: Optional[str], source: str) -> None: - """Fail fast when the resolved profile can't authenticate, before a command spins anything up. - - Without this, an unauthenticated `external-browser` profile hangs on a browser OAuth flow that - never completes — the local agent sits stuck with no error. The validation call itself can - hang that way, so it runs bounded (see `_validate_bounded`). - """ - if _needs_browser_login(profile): - # The bounded validation would open a browser flow this command will abandon at its - # timeout, so fail before it starts — interactive or not. - raise _preflight_error(profile, source) - - _, error = _validate_bounded(profile) - if error is None: - return - if _is_unreachable(error): - target = f"profile {profile!r} (from {source})" if profile else "the default auth" - raise 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.", - ) from error - raise _preflight_error(profile, source) from error - - -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 = f"agentbricks auth login {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 authenticate_profile(profile: str, source: str) -> tuple[_AgentBricksApiClient, str]: - """Validate `profile`, signing in through `databricks auth login` once if it isn't yet. - - Returns the client and the signed-in user; `source` is only named in errors. Without a - terminal (or in CI) nobody can complete a browser login, so an unauthenticated profile fails - with a hint instead of opening one. - """ - label = f"Databricks profile {profile!r} (from {source})" - error: Optional[Exception] - if _needs_browser_login(profile): - # Validation would open a browser flow it then abandons; head straight for the login. - error = None - else: - client_and_user, error = _validate_bounded(profile) - if client_and_user is not None: - return client_and_user - if not _is_auth_error(error): - # A timeout or a network/5xx failure won't be fixed by a login; don't open one. - raise AgentCliError(f"Could not validate {label}: {error}") from error - if not can_prompt_login(): - detail = "external-browser auth needs an interactive login" if error is None else str(error) - raise AgentCliError( - f"Could not validate {label}: {detail}", - hint="Run this command in an interactive terminal so Agent Bricks can open " - f"Databricks login, or authenticate first with `databricks auth login --profile {profile}`.", - ) from error - run_databricks_login(profile) - client_and_user, error = _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 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 0008a05b4..27bc04ba8 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/clients/api_client_provider.py +++ b/integrations/agentbricks/src/databricks_agentbricks/clients/api_client_provider.py @@ -27,6 +27,11 @@ def __init__( 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 + def get(self) -> _AgentBricksApiClient: """Return the shared client, constructing it on first use.""" client = 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 9fe88c92a..8b9972e73 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/clients/apps_client.py +++ b/integrations/agentbricks/src/databricks_agentbricks/clients/apps_client.py @@ -22,6 +22,7 @@ from databricks_agentbricks.clients.databricks_cli import _databricks from databricks_agentbricks.clients.legacy_runtime_store import LakebaseBackend from databricks_agentbricks.errors import AgentCliError +from databricks_agentbricks.projects.profiles import auth_login_command DatabricksRunner = Callable[..., subprocess.CompletedProcess] @@ -145,7 +146,7 @@ def exists(self, name: str) -> bool: marker in detail.lower() for marker in ("resource_does_not_exist", "does not exist") ): return False - login = f"agentbricks auth login {self._profile or ''}" + login = auth_login_command(self._profile) raise 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 5f05af366..ee6964c2b 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. @@ -10,10 +10,12 @@ from __future__ import annotations import subprocess +import sys from collections.abc import Mapping from typing import Optional from databricks_agentbricks.errors import AgentCliError +from databricks_agentbricks.projects.profiles import auth_login_command def _databricks( @@ -38,3 +40,19 @@ def _databricks( detail = (result.stderr or result.stdout or "").strip() if capture else None raise AgentCliError(action or "A Databricks CLI command failed.", hint=detail) return result + + +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}).") diff --git a/integrations/agentbricks/src/databricks_agentbricks/presentation/profile.py b/integrations/agentbricks/src/databricks_agentbricks/presentation/profile.py index 3388a6f78..9075df0cb 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/presentation/profile.py +++ b/integrations/agentbricks/src/databricks_agentbricks/presentation/profile.py @@ -4,12 +4,14 @@ import click -from databricks_agentbricks.cli.workspace_auth import ( +from databricks_agentbricks.presentation import render +from databricks_agentbricks.projects.profiles import ( + auth_login_command, + auth_switch_command, profile_exists, profile_host, workspace_for_profile, ) -from databricks_agentbricks.presentation import render def announce_workspace(profile: str | None, source: str, *, output: str) -> None: @@ -32,7 +34,7 @@ def note_choose_profile(profiles: list[str], obj) -> None: 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 `agentbricks auth switch ` with one of: {', '.join(profiles)}.[/]" + f"config. Run `{auth_switch_command(None)}` with one of: {', '.join(profiles)}.[/]" ) @@ -45,7 +47,7 @@ def warn_unknown_profile(profile: str, obj) -> None: 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 `agentbricks auth login {profile}` before `agentbricks dev`.[/]" + f"config yet. Run `{auth_login_command(profile)}` before `agentbricks dev`.[/]" ) @@ -77,7 +79,7 @@ def warn_legacy_default_profile(profile: str, output: str) -> None: 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 `agentbricks auth switch {profile}` in the project " + f"used. Pass `-p {profile}`, or run `{auth_switch_command(profile)}` in the project " "(delete ~/.agentbricks/config.json to hide this).", dim=True, ), diff --git a/integrations/agentbricks/src/databricks_agentbricks/projects/env_file.py b/integrations/agentbricks/src/databricks_agentbricks/projects/env_file.py index 82d9dee9a..71babc640 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/projects/env_file.py +++ b/integrations/agentbricks/src/databricks_agentbricks/projects/env_file.py @@ -41,11 +41,15 @@ def _write_text_atomically(path: pathlib.Path, text: str) -> None: def read_env(path: pathlib.Path) -> dict[str, str]: """Minimal KEY=VALUE reader for a project `.env` (python-dotenv is not a CLI dependency).""" - values: dict[str, str] = {} try: text = _read_text(path) except OSError: - return values + return {} + return _parse_env(text) + + +def _parse_env(text: str) -> dict[str, str]: + values: dict[str, str] = {} for line in _lines(text): line = line.strip() if not line or line.startswith("#") or "=" not in line: @@ -112,11 +116,11 @@ def update_env_profile(dest: pathlib.Path, profile: str) -> bool: env_path = dest / ".env" if not env_path.exists(): return write_env(dest, profile) - if read_env(env_path).get("DATABRICKS_CONFIG_PROFILE") == profile: + before = _read_text(env_path) + if _parse_env(before).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. return False - before = _read_text(env_path) bom = _BOM if before.startswith(_BOM) else "" lines = _with_profile_line(_lines(before), profile) _write_text_atomically(env_path, bom + "\n".join(lines) + "\n") 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..b72d5dc2f --- /dev/null +++ b/integrations/agentbricks/src/databricks_agentbricks/projects/profiles.py @@ -0,0 +1,204 @@ +"""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 _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]]: + """The workspace a profile points at, as (host, display text); (None, None) when unknown. + + For display only: without a profile the text names where the host came from + (DATABRICKS_HOST, else the config file's `[DEFAULT]`), because "no workspace" is the answer + the caller needs to surface. It does not claim to mirror the SDK's full default resolution. + """ + if profile: + host = profile_host(profile) + return host, host + env_host = os.environ.get("DATABRICKS_HOST") + if env_host: + return env_host, f"{env_host} (from DATABRICKS_HOST)" + default_host = profile_host(configparser.DEFAULTSECT) + if default_host: + return default_host, f"{default_host} ([DEFAULT] 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/dev_service.py b/integrations/agentbricks/src/databricks_agentbricks/services/dev_service.py index e18d29f2a..b4a0f829d 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/services/dev_service.py +++ b/integrations/agentbricks/src/databricks_agentbricks/services/dev_service.py @@ -13,7 +13,7 @@ from collections.abc import Iterator, Mapping from contextlib import contextmanager from dataclasses import dataclass, field -from typing import Callable, Optional, Protocol +from typing import Protocol from databricks_agentbricks.clients.apps_client import AppsClient from databricks_agentbricks.errors import AgentCliError @@ -101,12 +101,20 @@ def __init__( project_resolver: ProjectResolver, apps_client: AppsClient, local_tracing: LocalTracing, - preflight: Optional[Callable[[], None]] = None, ) -> None: self._project_resolver = project_resolver self._apps_client = apps_client self._local_tracing = local_tracing - self._preflight = preflight + + 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]: @@ -117,17 +125,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`).", - ) - - # 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 self._preflight is not None: - self._preflight() + app_yaml = self.check_project(request.source) project = self._project_resolver.load(source_dir) if project is not None and project.tools: 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..53b7dd55e --- /dev/null +++ b/integrations/agentbricks/src/databricks_agentbricks/services/workspace_auth.py @@ -0,0 +1,189 @@ +"""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 Optional + +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 + + +def _validate_profile(profile: Optional[str]) -> tuple[_AgentBricksApiClient, str]: + client = _AgentBricksApiClient(profile) + return client, client.current_user + + +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 + + if isinstance( + error, (RequestsConnectionError, InternalError, TemporarilyUnavailable, DeadlineExceeded) + ): + 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" + 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`. + + A browser-OAuth flow can block forever, so the call is bounded. Daemon because executor + threads are joined at interpreter exit and a hung flow 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] + + +def preflight_auth(profile: Optional[str], source: str) -> _AgentBricksApiClient: + """Fail fast when the resolved profile can't authenticate, before a command spins anything up. + + Returns the validated client so the command can reuse it. Without this, an unauthenticated + `external-browser` profile hangs on a browser OAuth flow that never completes — the local + agent sits stuck with no error. The validation call itself can hang that way, so it runs + bounded (see `_validate_bounded`). + """ + if _needs_browser_login(profile): + # The bounded validation would open a browser flow this command will abandon at its + # timeout, so fail before it starts — interactive or not. + raise _not_signed_in_error(profile, source) + + client_and_user, error = _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 + # A timeout here is most likely a browser flow nobody completed, so it reads as signed out too. + raise _not_signed_in_error(profile, source) from error + + +def authenticate_profile( + profile: str, source: str, *, interactive: bool +) -> tuple[_AgentBricksApiClient, str]: + """Validate `profile`, signing in through `databricks auth login` once if it isn't yet. + + Returns the client and the signed-in user; `source` is only named in errors. Without + `interactive` (no terminal, or CI) nobody can complete a browser login, so an unauthenticated + profile fails with a hint instead of opening one. + """ + label = _profile_label(profile, source) + error: Optional[Exception] + if _needs_browser_login(profile): + # Validation would open a browser flow it then abandons; head straight for the login. + error = None + else: + client_and_user, error = _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: + # A timeout or an unexpected failure won't be fixed by a login; don't open one. + 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 + run_databricks_login(profile) + client_and_user, error = _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 diff --git a/integrations/agentbricks/tests/unit_tests/auth_test.py b/integrations/agentbricks/tests/unit_tests/auth_test.py index 02f4a0e88..b7561f525 100644 --- a/integrations/agentbricks/tests/unit_tests/auth_test.py +++ b/integrations/agentbricks/tests/unit_tests/auth_test.py @@ -15,10 +15,12 @@ from databricks.sdk.errors import Unauthenticated from databricks_agentbricks.cli import app as cli -from databricks_agentbricks.cli import auth, workspace_auth +from databricks_agentbricks.cli import auth from databricks_agentbricks.cli.app import CliContext from databricks_agentbricks.clients import api_client_provider from databricks_agentbricks.projects.env_file import read_env +from databricks_agentbricks.projects.profiles import ProfileInfo +from databricks_agentbricks.services import workspace_auth def _write_env(directory: pathlib.Path, profile: str) -> None: @@ -47,11 +49,6 @@ def outside_project(tmp_path: pathlib.Path, monkeypatch: pytest.MonkeyPatch) -> return tmp_path -@pytest.fixture -def interactive_terminal(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(workspace_auth, "_is_interactive", lambda: True) - - @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.""" @@ -86,11 +83,9 @@ def test_context_precedence_is_flag_then_env_file_then_env_var(tmp_path, monkeyp project.use_project_profile(tmp_path) flagged.use_project_profile(tmp_path) - assert ambient.profile_info == workspace_auth.ProfileInfo( - "from-env-var", "DATABRICKS_CONFIG_PROFILE" - ) - assert project.profile_info == workspace_auth.ProfileInfo("proj", ".env") - assert flagged.profile_info == workspace_auth.ProfileInfo("from-flag", "--profile") + 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): @@ -224,13 +219,11 @@ def test_login_without_a_target_errors_with_hint_and_ignores_env_var( def test_login_runs_databricks_cli_when_unauthenticated_and_interactive( - monkeypatch, interactive_terminal + 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) - databricks_login = mock.Mock(return_value=mock.Mock(returncode=0)) - monkeypatch.setattr(workspace_auth.subprocess, "run", databricks_login) result = _invoke_profile(["login", "ml"]) @@ -239,12 +232,10 @@ def test_login_runs_databricks_cli_when_unauthenticated_and_interactive( assert "Logged in to 'ml'" in _flat(result.output) -def test_login_noninteractive_unauthenticated_fails_without_browser(monkeypatch): - monkeypatch.setattr( - workspace_auth, "_validate_bounded", lambda profile: (None, Unauthenticated("expired")) - ) - databricks_login = mock.Mock() - monkeypatch.setattr(workspace_auth.subprocess, "run", databricks_login) +def test_login_noninteractive_unauthenticated_fails_without_browser( + fail_validation, databricks_login +): + fail_validation(Unauthenticated("expired")) result = _invoke_profile(["login", "ml"]) diff --git a/integrations/agentbricks/tests/unit_tests/conftest.py b/integrations/agentbricks/tests/unit_tests/conftest.py index 7a497516e..650b50817 100644 --- a/integrations/agentbricks/tests/unit_tests/conftest.py +++ b/integrations/agentbricks/tests/unit_tests/conftest.py @@ -17,7 +17,9 @@ import pytest from databricks_agentbricks.agent_hint import AGENT_ENV_MARKERS, AGENT_ENV_VALUES -from databricks_agentbricks.cli import workspace_auth +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") @@ -35,7 +37,7 @@ def databrickscfg_path( 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(workspace_auth, "_is_interactive", lambda: False) + monkeypatch.setattr(auth, "_is_interactive", lambda: False) return path @@ -58,6 +60,30 @@ def validated_profile(monkeypatch: pytest.MonkeyPatch) -> mock.Mock: 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) def human_cli_environment(monkeypatch): """Keep ordinary CLI assertions independent of the agent running the tests.""" diff --git a/integrations/agentbricks/tests/unit_tests/deploy_test.py b/integrations/agentbricks/tests/unit_tests/deploy_test.py index 11f4b1384..1d84ed6b8 100644 --- a/integrations/agentbricks/tests/unit_tests/deploy_test.py +++ b/integrations/agentbricks/tests/unit_tests/deploy_test.py @@ -19,7 +19,6 @@ from databricks_agentbricks.cli import deploy as deploy_mod from databricks_agentbricks.cli.app import CliContext -from databricks_agentbricks.cli.workspace_auth import ProfileInfo 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,6 +33,7 @@ from databricks_agentbricks.errors import AgentCliError 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 ( DeployRequest, @@ -58,7 +58,9 @@ class _FakeApiClient: 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()) + self.api_client_provider = types.SimpleNamespace( + get=lambda: _FakeApiClient(), use=lambda client: None + ) @pytest.fixture(autouse=True) @@ -174,7 +176,8 @@ def __call__(self, args, profile, **kwargs): def _real_deploy_context(api: _DeployApi, *, output: str = "text"): ctx = CliContext("prof", output) - ctx.api_client_provider = types.SimpleNamespace(get=lambda: api) + # 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 @@ -383,6 +386,7 @@ def test_deploy_resolves_profile_from_flag_then_source_env( 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) diff --git a/integrations/agentbricks/tests/unit_tests/dev_test.py b/integrations/agentbricks/tests/unit_tests/dev_test.py index 4027882d2..9bcba51b3 100644 --- a/integrations/agentbricks/tests/unit_tests/dev_test.py +++ b/integrations/agentbricks/tests/unit_tests/dev_test.py @@ -11,7 +11,6 @@ from click.testing import CliRunner from databricks_agentbricks.cli import dev as dev_mod -from databricks_agentbricks.cli import workspace_auth from databricks_agentbricks.cli.app import CliContext from databricks_agentbricks.errors import AgentCliError from databricks_agentbricks.projects.agent_project import AgentProject, ToolSpec @@ -705,11 +704,11 @@ def test_dev_env_credentials_replace_the_profile_only_when_complete( 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): +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") - monkeypatch.setattr( - workspace_auth, "_validate_bounded", lambda profile: (None, RuntimeError("expired")) - ) + fail_validation(RuntimeError("expired")) run_local = mock.Mock() monkeypatch.setattr(dev_mod, "_databricks", run_local) diff --git a/integrations/agentbricks/tests/unit_tests/init_test.py b/integrations/agentbricks/tests/unit_tests/init_test.py index a221bc38c..bc2d60eae 100644 --- a/integrations/agentbricks/tests/unit_tests/init_test.py +++ b/integrations/agentbricks/tests/unit_tests/init_test.py @@ -17,7 +17,6 @@ from databricks_agentbricks.cli import app as cli_app from databricks_agentbricks.cli import init as init_mod -from databricks_agentbricks.cli import workspace_auth from databricks_agentbricks.cli.app import CliContext from databricks_agentbricks.errors import AgentCliError from databricks_agentbricks.projects.types import AgentFramework, AgentServer @@ -328,10 +327,9 @@ def test_init_inside_another_project_does_not_inherit_its_env_profile( def test_init_with_databricks_host_and_no_profile_writes_no_env( - tmp_path: pathlib.Path, monkeypatch, validated_profile + tmp_path: pathlib.Path, monkeypatch, validated_profile, interactive_terminal ): monkeypatch.setenv("DATABRICKS_HOST", "https://env.example") - monkeypatch.setattr(workspace_auth, "_is_interactive", lambda: True) dest = tmp_path / "proj" result = _init(dest, output="json") assert result.exit_code == 0, result.output @@ -342,8 +340,9 @@ def test_init_with_databricks_host_and_no_profile_writes_no_env( validated_profile.assert_not_called() # interactive, but host env auth never signs in -def test_init_signs_in_when_interactive(tmp_path: pathlib.Path, monkeypatch, validated_profile): - monkeypatch.setattr(workspace_auth, "_is_interactive", lambda: True) +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") diff --git a/integrations/agentbricks/tests/unit_tests/workspace_auth_test.py b/integrations/agentbricks/tests/unit_tests/workspace_auth_test.py index db7dcf206..e0c7ab0b6 100644 --- a/integrations/agentbricks/tests/unit_tests/workspace_auth_test.py +++ b/integrations/agentbricks/tests/unit_tests/workspace_auth_test.py @@ -1,8 +1,8 @@ """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 so nothing touches the -network or a browser. `DATABRICKS_CONFIG_FILE` points at a tmp file (see conftest.py). +(`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 @@ -14,8 +14,10 @@ import pytest from databricks.sdk.errors import Unauthenticated -from databricks_agentbricks.cli import workspace_auth +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" @@ -25,20 +27,6 @@ def _disable_telemetry(monkeypatch): monkeypatch.setenv("AGENTBRICKS_DISABLE_TELEMETRY", "1") -def _set_interactive(monkeypatch, interactive: bool = True) -> None: - monkeypatch.setattr(workspace_auth, "_is_interactive", lambda: interactive) - - -def _fail_validation(monkeypatch, error: Exception) -> None: - monkeypatch.setattr(workspace_auth, "_validate_bounded", lambda profile: (None, error)) - - -def _stub_databricks_cli(monkeypatch, **mock_kwargs) -> mock.Mock: - databricks_login = mock.Mock(**mock_kwargs) - monkeypatch.setattr(workspace_auth.subprocess, "run", databricks_login) - return databricks_login - - # --- resolve_profile ------------------------------------------------------------------------- @@ -63,7 +51,7 @@ def test_resolve_profile_order(tmp_path, monkeypatch, flag, env_file, env_var, e if env_var is not None: monkeypatch.setenv("DATABRICKS_CONFIG_PROFILE", env_var) - assert workspace_auth.resolve_profile(flag, tmp_path) == workspace_auth.ProfileInfo(*expected) + assert profiles.resolve_profile(flag, tmp_path) == profiles.ProfileInfo(*expected) # --- config-file lookups --------------------------------------------------------------------- @@ -74,7 +62,7 @@ def test_config_file_path_expands_tilde_and_tolerates_percent_in_values(tmp_path (tmp_path / "cfg").write_text(_PROF_CFG + "token = 100%secret\n") monkeypatch.setenv("DATABRICKS_CONFIG_FILE", "~/cfg") - assert workspace_auth.profile_host("prof") == "https://prof.example" + assert profiles.profile_host("prof") == "https://prof.example" @pytest.mark.parametrize( @@ -95,20 +83,21 @@ def test_workspace_for_profile(write_databrickscfg, monkeypatch, cfg, env_host, elif expected == "default": expected = ("https://d.example", "https://d.example ([DEFAULT] profile)") - assert workspace_auth.workspace_for_profile(profile) == expected + assert profiles.workspace_for_profile(profile) == expected # --- preflight_auth -------------------------------------------------------------------------- def test_preflight_passes_when_profile_validates(validated_profile): - workspace_auth.preflight_auth("prof", "--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(monkeypatch): - _fail_validation(monkeypatch, Unauthenticated("expired")) +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") @@ -144,26 +133,26 @@ def test_preflight_external_browser_fails_fast(write_databrickscfg, validated_pr def test_authenticate_profile_returns_validated_client_without_login( - validated_profile, monkeypatch + validated_profile, databricks_login ): - databricks_login = _stub_databricks_cli(monkeypatch) - assert ( - workspace_auth.authenticate_profile("prof", "argument") == validated_profile.return_value[0] + 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): +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) - _set_interactive(monkeypatch) - databricks_login = _stub_databricks_cli(monkeypatch, return_value=mock.Mock(returncode=0)) - assert workspace_auth.authenticate_profile("prof", "argument") == (client, "me@example.com") + 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. @@ -175,13 +164,12 @@ def test_authenticate_profile_logs_in_then_revalidates(monkeypatch): ) -def test_authenticate_profile_reports_a_missing_databricks_cli(monkeypatch): - _fail_validation(monkeypatch, Unauthenticated("expired")) - _set_interactive(monkeypatch) - _stub_databricks_cli(monkeypatch, side_effect=FileNotFoundError) +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") + 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 @@ -190,22 +178,22 @@ def test_authenticate_profile_reports_a_missing_databricks_cli(monkeypatch): ("error", "ci", "interactive", "hint"), [ (RuntimeError("no creds"), False, True, None), - (Unauthenticated("expired"), False, False, "databricks auth login --profile prof"), - (Unauthenticated("expired"), True, True, "databricks auth login --profile prof"), + (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( - monkeypatch, error, ci, interactive, hint + request, monkeypatch, fail_validation, databricks_login, error, ci, interactive, hint ): if ci: monkeypatch.setenv("CI", "true") - _set_interactive(monkeypatch, interactive) - _fail_validation(monkeypatch, error) - databricks_login = _stub_databricks_cli(monkeypatch) + if interactive: + request.getfixturevalue("interactive_terminal") + fail_validation(error) with pytest.raises(AgentCliError) as raised: - workspace_auth.authenticate_profile("prof", ".env") + 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: From 1dc8a72e82a03d52bd12f7b36eac31a29146301a Mon Sep 17 00:00:00 2001 From: jamesbxwu Date: Thu, 15 Oct 2026 15:14:39 +0000 Subject: [PATCH 03/12] Fix resolved workspace notices and cached browser authentication --- integrations/agentbricks/README.md | 8 + integrations/agentbricks/cli.md | 8 +- .../src/databricks_agentbricks/cli/app.py | 9 +- .../src/databricks_agentbricks/cli/deploy.py | 19 +- .../src/databricks_agentbricks/cli/dev.py | 14 +- .../clients/api_client_provider.py | 15 +- .../clients/cached_browser_auth.py | 60 ++++++ .../presentation/profile.py | 23 +- .../projects/profiles.py | 21 +- .../services/workspace_auth.py | 56 ++--- .../src/databricks_agentkit/_api_client.py | 20 +- .../agentbricks/tests/unit_tests/auth_test.py | 85 ++++++++ .../unit_tests/cached_browser_auth_test.py | 202 ++++++++++++++++++ .../tests/unit_tests/client_test.py | 13 +- .../tests/unit_tests/deploy_test.py | 29 ++- .../agentbricks/tests/unit_tests/dev_test.py | 2 +- .../tests/unit_tests/workspace_auth_test.py | 43 +++- 17 files changed, 544 insertions(+), 83 deletions(-) create mode 100644 integrations/agentbricks/src/databricks_agentbricks/clients/cached_browser_auth.py create mode 100644 integrations/agentbricks/tests/unit_tests/cached_browser_auth_test.py diff --git a/integrations/agentbricks/README.md b/integrations/agentbricks/README.md index eec329820..d656fe0d6 100644 --- a/integrations/agentbricks/README.md +++ b/integrations/agentbricks/README.md @@ -120,6 +120,11 @@ 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 nonzero exits, a bounded error category). Parameter entries are limited to options explicitly @@ -161,6 +166,9 @@ 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. +An external-browser profile's valid or refreshable Python SDK token cache is reused without +opening a browser. If that cache is missing or unusable, run the suggested `auth login` command +in an interactive terminal. `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 a6e8b108a..0816ede0f 100644 --- a/integrations/agentbricks/cli.md +++ b/integrations/agentbricks/cli.md @@ -39,7 +39,13 @@ global default profile. The profile is resolved in this order, most to least spe `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 (`[DEFAULT]` in `~/.databrickscfg`) +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. Authentication preflight reuses a valid or refreshable external-browser Python SDK token +cache without opening a browser; a missing or unusable cache requires `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 diff --git a/integrations/agentbricks/src/databricks_agentbricks/cli/app.py b/integrations/agentbricks/src/databricks_agentbricks/cli/app.py index 16f9885a4..47bd2fb0a 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/cli/app.py +++ b/integrations/agentbricks/src/databricks_agentbricks/cli/app.py @@ -89,11 +89,16 @@ def use_flag_profile(self, profile: str) -> None: self.api_client_provider = self._client_provider(profile) self.profile_info = ProfileInfo(profile, SOURCE_FLAG) - def announce_workspace(self) -> None: + 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) + 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: diff --git a/integrations/agentbricks/src/databricks_agentbricks/cli/deploy.py b/integrations/agentbricks/src/databricks_agentbricks/cli/deploy.py index ecbd5dde8..fd69ce6b4 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/cli/deploy.py +++ b/integrations/agentbricks/src/databricks_agentbricks/cli/deploy.py @@ -163,10 +163,14 @@ def deploy( not authentication and not the session id itself (send the session id in the request body): X-Routing-Key: """ - obj.announce_workspace() # 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. - obj.api_client_provider.use(preflight_auth(obj.profile, obj.profile_info.source)) + 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, @@ -244,11 +248,18 @@ def deployments_stop(obj, name, yes) -> None: @click.pass_obj def deployments_delete(obj, name, yes) -> None: """Delete a deployment and, when managed provisioning is enabled, its Runtime Store.""" - obj.announce_workspace() service = build_deploy_service(obj) + # 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() target = f"deployment '{name}'" if not yes and service.deletes_runtime_store_data(): target += " and its Runtime Store data" render.confirm_destroy(target, assume_yes=yes) - service.delete(name) + try: + service.delete(name) + except Exception: + obj.announce_workspace() + raise present_deleted(name, output=obj.output) diff --git a/integrations/agentbricks/src/databricks_agentbricks/cli/dev.py b/integrations/agentbricks/src/databricks_agentbricks/cli/dev.py index 827133d8f..c4b30a0a0 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/cli/dev.py +++ b/integrations/agentbricks/src/databricks_agentbricks/cli/dev.py @@ -127,17 +127,19 @@ def dev( # beats `.env` because the template loads dotenv with override=False. local_env["DATABRICKS_CONFIG_PROFILE"] = obj.profile - # Announced before the preflight so an auth failure still names its target. When `.env` - # credentials win, the resolved profile would name a different workspace than the agent - # actually uses, so it isn't announced. - if not uses_env_credentials: - obj.announce_workspace() 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: - preflight_auth(obj.profile, obj.profile_info.source) + 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, 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 27bc04ba8..7fef07d83 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/clients/api_client_provider.py +++ b/integrations/agentbricks/src/databricks_agentbricks/clients/api_client_provider.py @@ -20,7 +20,9 @@ class ApiClientProvider: """Own exactly one lazily created API client for a CLI invocation.""" def __init__( - self, profile: Optional[str], on_first_use: Optional[Callable[[], None]] = None + self, + profile: Optional[str], + on_first_use: Optional[Callable[[Optional[str]], None]] = None, ) -> None: self._profile = profile self._on_first_use = on_first_use @@ -31,6 +33,8 @@ 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.""" @@ -40,7 +44,12 @@ def get(self) -> _AgentBricksApiClient: with self._lock: if self._client is None: + 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 = _AgentBricksApiClient(self._profile) + self._on_first_use(self._client.host) return self._client diff --git a/integrations/agentbricks/src/databricks_agentbricks/clients/cached_browser_auth.py b/integrations/agentbricks/src/databricks_agentbricks/clients/cached_browser_auth.py new file mode 100644 index 000000000..b6724abe5 --- /dev/null +++ b/integrations/agentbricks/src/databricks_agentbricks/clients/cached_browser_auth.py @@ -0,0 +1,60 @@ +"""Use the Python SDK's external-browser cache without ever starting OAuth consent.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Optional + +from databricks_agentkit._api_client import _AgentBricksApiClient + +if TYPE_CHECKING: + from databricks.sdk.config import Config + from databricks.sdk.credentials_provider import CredentialsProvider + + +def _cached_credentials(config: Config) -> CredentialsProvider: + from databricks.sdk import oauth + + # Group assumption was added after the minimum supported SDK (0.124). + group_id = getattr(config, "group_id", None) + # Match the SDK's external_browser cache identity, including Azure Entra profiles, scopes, + # and group assumption. Using a different identity could reuse another profile's session. + if config.client_id: + client_id, client_secret = config.client_id, config.client_secret + endpoints = config.databricks_oidc_endpoints + elif config.azure_client_id: + if group_id: + raise ValueError("external-browser with Azure Entra ID cannot assume a group") + client_id, client_secret = config.azure_client_id, config.azure_client_secret + endpoints = oauth.get_azure_entra_id_workspace_endpoints(config.host) + else: + client_id, client_secret = "databricks-cli", None + endpoints = config.databricks_oidc_endpoints + if endpoints is None: + raise ValueError("The workspace does not support OAuth") + scopes = list(config.get_scopes()) + if not config.disable_oauth_refresh_token and "offline_access" not in scopes: + scopes.append("offline_access") + cache = oauth.TokenCache( + host=config.host, + oidc_endpoints=endpoints, + client_id=client_id, + client_secret=client_secret, + redirect_url="http://localhost:8020", + scopes=scopes, + profile=config.profile, + **({"group_id": group_id} if group_id else {}), + ) + credentials = cache.load() + if credentials is None: + raise ValueError("No cached external-browser session is available") + # A refresh failure propagates to the caller; unlike the SDK's interactive strategy there + # is deliberately no fallback to OAuthClient, consent, or a browser listener here. + credentials.token() + return credentials(config) + + +def cached_browser_client(profile: Optional[str]) -> _AgentBricksApiClient: + from databricks.sdk.credentials_provider import credentials_strategy + + strategy = credentials_strategy("external-browser", ["host"])(_cached_credentials) + return _AgentBricksApiClient(profile, credentials_strategy=strategy) diff --git a/integrations/agentbricks/src/databricks_agentbricks/presentation/profile.py b/integrations/agentbricks/src/databricks_agentbricks/presentation/profile.py index 9075df0cb..4bc2550ac 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/presentation/profile.py +++ b/integrations/agentbricks/src/databricks_agentbricks/presentation/profile.py @@ -9,12 +9,13 @@ auth_login_command, auth_switch_command, profile_exists, - profile_host, workspace_for_profile, ) -def announce_workspace(profile: str | None, source: str, *, output: str) -> None: +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 @@ -22,7 +23,9 @@ def announce_workspace(profile: str | None, source: str, *, output: str) -> None """ if output != "text": return - _, workspace = workspace_for_profile(profile) + 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 @@ -55,11 +58,15 @@ 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 = profile_host(profile) - where = host or ( - "not in your Databricks config yet; `databricks auth login` will create it" - if not profile_exists(profile) - else "no host configured" + 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. diff --git a/integrations/agentbricks/src/databricks_agentbricks/projects/profiles.py b/integrations/agentbricks/src/databricks_agentbricks/projects/profiles.py index b72d5dc2f..b0d8d797b 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/projects/profiles.py +++ b/integrations/agentbricks/src/databricks_agentbricks/projects/profiles.py @@ -164,21 +164,24 @@ def sdk_default_profile() -> tuple[str, str]: def workspace_for_profile(profile: Optional[str]) -> tuple[Optional[str], Optional[str]]: - """The workspace a profile points at, as (host, display text); (None, None) when unknown. + """An offline workspace estimate as (host, display text), for notices and `auth describe`. - For display only: without a profile the text names where the host came from - (DATABRICKS_HOST, else the config file's `[DEFAULT]`), because "no workspace" is the answer - the caller needs to surface. It does not claim to mirror the SDK's full default resolution. + 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. """ - if profile: - host = profile_host(profile) - return host, host env_host = os.environ.get("DATABRICKS_HOST") if env_host: return env_host, f"{env_host} (from DATABRICKS_HOST)" - default_host = profile_host(configparser.DEFAULTSECT) + 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)" + return default_host, f"{default_host} ([{default_profile}] profile)" return None, "none configured" diff --git a/integrations/agentbricks/src/databricks_agentbricks/services/workspace_auth.py b/integrations/agentbricks/src/databricks_agentbricks/services/workspace_auth.py index 53b7dd55e..22a2481ae 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/services/workspace_auth.py +++ b/integrations/agentbricks/src/databricks_agentbricks/services/workspace_auth.py @@ -11,6 +11,7 @@ import threading from typing import Optional +from databricks_agentbricks.clients.cached_browser_auth import cached_browser_client from databricks_agentbricks.clients.databricks_cli import run_databricks_login from databricks_agentbricks.errors import AgentCliError from databricks_agentbricks.projects.profiles import ( @@ -22,7 +23,11 @@ def _validate_profile(profile: Optional[str]) -> tuple[_AgentBricksApiClient, str]: - client = _AgentBricksApiClient(profile) + client = ( + cached_browser_client(profile) + if sdk_auth_type(profile) == "external-browser" + else _AgentBricksApiClient(profile) + ) return client, client.current_user @@ -91,23 +96,14 @@ def _unreachable_error(profile: Optional[str], source: str, error: Exception) -> _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`. - A browser-OAuth flow can block forever, so the call is bounded. Daemon because executor - threads are joined at interpreter exit and a hung flow must not block the CLI's exit too. + Auth/network resolution can block, so the call is bounded. External-browser profiles use + cached credentials only and never start consent. Daemon because executor threads are joined + at interpreter exit and a hung request 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]]] = [] @@ -134,18 +130,12 @@ def preflight_auth(profile: Optional[str], source: str) -> _AgentBricksApiClient agent sits stuck with no error. The validation call itself can hang that way, so it runs bounded (see `_validate_bounded`). """ - if _needs_browser_login(profile): - # The bounded validation would open a browser flow this command will abandon at its - # timeout, so fail before it starts — interactive or not. - raise _not_signed_in_error(profile, source) - client_and_user, error = _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 - # A timeout here is most likely a browser flow nobody completed, so it reads as signed out too. raise _not_signed_in_error(profile, source) from error @@ -159,25 +149,19 @@ def authenticate_profile( profile fails with a hint instead of opening one. """ label = _profile_label(profile, source) - error: Optional[Exception] - if _needs_browser_login(profile): - # Validation would open a browser flow it then abandons; head straight for the login. - error = None - else: - client_and_user, error = _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: - # A timeout or an unexpected failure won't be fixed by a login; don't open one. - raise AgentCliError(f"Could not validate {label}: {error}") from error + client_and_user, error = _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: + # A timeout or an unexpected failure won't be fixed by a login; don't open one. + 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}", + f"Could not validate {label}: {error}", hint=_sign_in_hint(profile, needs_terminal=True), ) from error run_databricks_login(profile) diff --git a/integrations/agentbricks/src/databricks_agentkit/_api_client.py b/integrations/agentbricks/src/databricks_agentkit/_api_client.py index a47bb9bd8..b86fb2deb 100644 --- a/integrations/agentbricks/src/databricks_agentkit/_api_client.py +++ b/integrations/agentbricks/src/databricks_agentkit/_api_client.py @@ -20,6 +20,7 @@ if TYPE_CHECKING: from databricks.sdk import WorkspaceClient + from databricks.sdk.credentials_provider import CredentialsStrategy _BASE = "/api/2.0/agents" _MCP_SERVICES_PATH = "/api/2.1/unity-catalog/mcp-services" @@ -144,12 +145,17 @@ def _bound_retry_timeout(client: WorkspaceClient) -> WorkspaceClient: return client -def _workspace_client(profile: Optional[str]) -> WorkspaceClient: +def _workspace_client( + profile: Optional[str], *, credentials_strategy: Optional[CredentialsStrategy] = None +) -> WorkspaceClient: # Imported here (not at module top) so the ~0.7s databricks.sdk import is paid only when a # command actually builds a client, not on every CLI invocation. from databricks.sdk import WorkspaceClient - client = _bound_retry_timeout(WorkspaceClient(profile=profile)) + auth_kwargs: dict[str, Any] = ( + {"credentials_strategy": credentials_strategy} if credentials_strategy is not None else {} + ) + client = _bound_retry_timeout(WorkspaceClient(profile=profile, **auth_kwargs)) if not profile or not client.config.workspace_id: return client @@ -164,6 +170,7 @@ def _workspace_client(profile: Optional[str]) -> WorkspaceClient: profile=profile, host=configured_host, custom_headers={"X-Databricks-Org-Id": workspace_id}, + **auth_kwargs, ) ) @@ -176,11 +183,18 @@ def __init__( profile: Optional[str] = None, *, workspace_client: Optional[WorkspaceClient] = None, + credentials_strategy: Optional[CredentialsStrategy] = None, ) -> None: if profile is not None and workspace_client is not None: raise ValueError("profile and workspace_client are mutually exclusive") + if credentials_strategy is not None and workspace_client is not None: + raise ValueError("credentials_strategy and workspace_client are mutually exclusive") try: - self._w = workspace_client or _workspace_client(profile) + self._w = workspace_client or ( + _workspace_client(profile, credentials_strategy=credentials_strategy) + if credentials_strategy is not None + else _workspace_client(profile) + ) except Exception as exc: # noqa: BLE001 - surfaced as a clean CLI error raise AgentCliError( f"Could not initialize Databricks auth: {exc}", diff --git a/integrations/agentbricks/tests/unit_tests/auth_test.py b/integrations/agentbricks/tests/unit_tests/auth_test.py index b7561f525..3a5c92984 100644 --- a/integrations/agentbricks/tests/unit_tests/auth_test.py +++ b/integrations/agentbricks/tests/unit_tests/auth_test.py @@ -18,6 +18,7 @@ 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.env_file import read_env from databricks_agentbricks.projects.profiles import ProfileInfo from databricks_agentbricks.services import workspace_auth @@ -53,6 +54,7 @@ def outside_project(tmp_path: pathlib.Path, monkeypatch: pytest.MonkeyPatch) -> 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 @@ -123,6 +125,35 @@ def test_first_workspace_call_names_profile_and_host_once( 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 + + +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") + + with pytest.raises(AgentCliError): + _ctx("ws").client() + + assert ( + "Using profile 'ws' (from --profile) · https://profile.example" in capsys.readouterr().err + ) + + def test_workspace_notice_names_sdk_default_auth(capsys, monkeypatch, fake_api_client): monkeypatch.setenv("DATABRICKS_HOST", "https://env.example") @@ -134,6 +165,60 @@ def test_workspace_notice_names_sdk_default_auth(capsys, monkeypatch, fake_api_c ) +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() + + assert "Using profile 'ws' (from --profile) · https://env.example (from DATABRICKS_HOST)" in ( + capsys.readouterr().err + ) + + +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() diff --git a/integrations/agentbricks/tests/unit_tests/cached_browser_auth_test.py b/integrations/agentbricks/tests/unit_tests/cached_browser_auth_test.py new file mode 100644 index 000000000..fe8d8f66a --- /dev/null +++ b/integrations/agentbricks/tests/unit_tests/cached_browser_auth_test.py @@ -0,0 +1,202 @@ +"""Cached OAuth validation with synthetic credentials and no network/consent fallback.""" + +from __future__ import annotations + +import inspect +import json +import pathlib +from datetime import datetime, timedelta +from types import SimpleNamespace +from typing import Any +from unittest import mock + +import pytest +from databricks.sdk import oauth +from databricks.sdk.config import Config +from databricks.sdk.service.iam import CurrentUserAPI, User + +from databricks_agentbricks.clients.cached_browser_auth import _cached_credentials +from databricks_agentbricks.errors import AgentCliError +from databricks_agentbricks.services import workspace_auth + + +@pytest.fixture(autouse=True) +def isolate_auth_and_forbid_network(monkeypatch): + # Config/HOME already point at temp files (conftest); also remove alternate SDK credential + # env vars so no real secret can influence these tests. The only tokens below are synthetic. + for attr in Config.attributes(): + for name in (attr.env, *getattr(attr, "env_aliases", ())): + if name and name != "DATABRICKS_CONFIG_FILE": + monkeypatch.delenv(name, raising=False) + with ( + mock.patch("requests.Session.request", side_effect=AssertionError("unexpected network")), + mock.patch.object(oauth, "OAuthClient") as consent_client, + mock.patch.object(oauth.Consent, "launch_external_browser") as browser, + ): + yield + consent_client.assert_not_called() + browser.assert_not_called() + + +@pytest.fixture +def config(): + return SimpleNamespace( + host="https://ws.example", + profile="prof", + client_id=None, + client_secret=None, + azure_client_id=None, + azure_client_secret=None, + databricks_oidc_endpoints=oauth.OidcEndpoints( + authorization_endpoint="https://ws.example/oidc/v1/authorize", + token_endpoint="https://ws.example/oidc/v1/token", + ), + get_scopes=lambda: ["all-apis"], + disable_oauth_refresh_token=False, + ) + + +def _seed_cache(config, *, expired=False, **identity): + cache_kwargs: dict[str, Any] = { + "host": config.host, + "oidc_endpoints": config.databricks_oidc_endpoints, + "client_id": "databricks-cli", + "redirect_url": "http://localhost:8020", + "scopes": ["all-apis", "offline_access"], + "profile": config.profile, + **identity, + } + cache = oauth.TokenCache(**cache_kwargs) + token = oauth.Token( + access_token="synthetic-cached-token", + refresh_token="synthetic-refresh-token", + token_type="Bearer", + expiry=datetime.now() + timedelta(hours=-1 if expired else 1), + ) + path = pathlib.Path(cache.filename) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps({"token": token.as_dict()})) + return path + + +@pytest.fixture +def sdk_workspace(config, write_databrickscfg): + write_databrickscfg("[prof]\nhost = https://ws.example\nauth_type = external-browser\n") + metadata = SimpleNamespace( + account_id=None, + workspace_id=None, + oidc_endpoint="", + cloud=None, + token_federation_default_oidc_audiences=[], + ) + with ( + mock.patch("databricks.sdk.config.get_host_metadata", return_value=metadata), + mock.patch.object( + Config, + "databricks_oidc_endpoints", + new_callable=mock.PropertyMock, + return_value=config.databricks_oidc_endpoints, + ), + mock.patch.object( + CurrentUserAPI, "me", return_value=User(user_name="cached@example.com") + ) as current_user, + ): + yield current_user + + +@pytest.mark.parametrize("expired", [False, True], ids=["valid", "refreshable"]) +def test_preflight_and_login_accept_sdk_cached_session( + config, sdk_workspace, databricks_login, expired +): + _seed_cache(config, expired=expired) + refreshed = oauth.Token( + access_token="synthetic-refreshed-token", + refresh_token="synthetic-refresh-token", + token_type="Bearer", + expiry=datetime.now() + timedelta(hours=1), + ) + with mock.patch.object(oauth, "retrieve_token", return_value=refreshed) as refresh: + client = workspace_auth.preflight_auth("prof", "--profile") + expected_token = "synthetic-refreshed-token" if expired else "synthetic-cached-token" + assert client.workspace_client.config.authenticate() == { + "Authorization": f"Bearer {expected_token}" + } + if expired: + refresh.assert_called_once_with( + client_id="databricks-cli", + client_secret=None, + token_url=config.databricks_oidc_endpoints.token_endpoint, + params={ + "grant_type": "refresh_token", + "refresh_token": "synthetic-refresh-token", + }, + use_params=True, + headers={}, + ) + else: + refresh.assert_not_called() + _, user = workspace_auth.authenticate_profile("prof", "argument", interactive=True) + + assert user == "cached@example.com" + assert sdk_workspace.call_count == 2 + databricks_login.assert_not_called() + + +@pytest.mark.parametrize("state", ["missing", "invalid", "expired"]) +def test_unusable_cache_fails_without_starting_consent( + config, sdk_workspace, databricks_login, state +): + if state != "missing": + path = _seed_cache(config, expired=state == "expired") + if state == "invalid": + path.write_text("invalid JSON") + with mock.patch.object( + oauth, "retrieve_token", side_effect=ValueError("synthetic revoked token") + ): + with pytest.raises(AgentCliError, match="isn't authenticated"): + workspace_auth.preflight_auth("prof", "--profile") + with pytest.raises(AgentCliError) as raised: + workspace_auth.authenticate_profile("prof", "argument", interactive=False) + + assert raised.value.hint is not None + assert "interactive terminal" in raised.value.hint + sdk_workspace.assert_not_called() + databricks_login.assert_not_called() + + +@pytest.mark.parametrize( + "identity", + [ + {"profile": "other"}, + {"host": "https://other.example"}, + {"scopes": ["other-scope", "offline_access"]}, + {"group_id": "other-group"}, + ], + ids=["profile", "host", "scopes", "group"], +) +def test_a_different_cache_identity_cannot_authenticate(config, identity): + if "group_id" in identity and "group_id" not in inspect.signature(oauth.TokenCache).parameters: + pytest.skip("Group assumption is unavailable in the minimum supported SDK") + _seed_cache(config, **identity) + + with pytest.raises(ValueError, match="No cached external-browser session"): + _cached_credentials(config) + + +def test_cache_identity_includes_group_when_supported(config): + if "group_id" not in inspect.signature(oauth.TokenCache).parameters: + pytest.skip("Group assumption is unavailable in the minimum supported SDK") + config.group_id = "expected-group" + _seed_cache(config, group_id=config.group_id) + + assert _cached_credentials(config)() == {"Authorization": "Bearer synthetic-cached-token"} + + +def test_minimum_sdk_cache_constructor_does_not_need_group_argument(config): + # SDK 0.124 has this exact constructor shape and no Config.group_id attribute. + def old_cache(*, host, oidc_endpoints, client_id, client_secret, redirect_url, scopes, profile): + return SimpleNamespace(load=lambda: credentials) + + credentials = mock.Mock(return_value=lambda: {"Authorization": "Bearer synthetic-token"}) + with mock.patch.object(oauth, "TokenCache", side_effect=old_cache): + assert _cached_credentials(config)() == {"Authorization": "Bearer synthetic-token"} diff --git a/integrations/agentbricks/tests/unit_tests/client_test.py b/integrations/agentbricks/tests/unit_tests/client_test.py index 86e4252dd..d7794a6f7 100644 --- a/integrations/agentbricks/tests/unit_tests/client_test.py +++ b/integrations/agentbricks/tests/unit_tests/client_test.py @@ -916,7 +916,8 @@ def test_do_does_not_retry_transient_error_for_pop(workspace_client, sleep): sleep.assert_not_called() -def test_account_routed_profile_uses_configured_host_and_workspace_header(): +@pytest.mark.parametrize("strategy", [None, mock.Mock()], ids=["default-auth", "injected-auth"]) +def test_account_routed_profile_uses_configured_host_and_workspace_header(strategy): resolved = mock.Mock() resolved.config.host = "https://workspace.example.com" resolved.config.workspace_id = "123" @@ -929,15 +930,21 @@ def test_account_routed_profile_uses_configured_host_and_workspace_header(): return_value="https://account.example.com", ), ): - client = _workspace_client("p") + client = ( + _workspace_client("p", credentials_strategy=strategy) + if strategy is not None + else _workspace_client("p") + ) assert client is routed + auth_kwargs = {"credentials_strategy": strategy} if strategy is not None else {} assert wc.call_args_list == [ - mock.call(profile="p"), + mock.call(profile="p", **auth_kwargs), mock.call( profile="p", host="https://account.example.com", custom_headers={"X-Databricks-Org-Id": "123"}, + **auth_kwargs, ), ] diff --git a/integrations/agentbricks/tests/unit_tests/deploy_test.py b/integrations/agentbricks/tests/unit_tests/deploy_test.py index 1d84ed6b8..b203cafe9 100644 --- a/integrations/agentbricks/tests/unit_tests/deploy_test.py +++ b/integrations/agentbricks/tests/unit_tests/deploy_test.py @@ -59,7 +59,7 @@ 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: None + get=lambda: _FakeApiClient(), use=lambda client: self.announce_workspace(client.host) ) @@ -334,6 +334,30 @@ 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 + if managed: + service.delete.side_effect = lambda name: ctx.client() + 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("demo") + + def test_deploy_rejects_instance_count_above_platform_limit(): result = CliRunner().invoke(deploy_mod.deploy, ["demo", "--instances", "6"], obj=_Ctx()) @@ -400,7 +424,8 @@ def test_deploy_resolves_profile_from_flag_then_source_env( 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://{expected}.example" in out + 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): diff --git a/integrations/agentbricks/tests/unit_tests/dev_test.py b/integrations/agentbricks/tests/unit_tests/dev_test.py index 9bcba51b3..769c2fa9c 100644 --- a/integrations/agentbricks/tests/unit_tests/dev_test.py +++ b/integrations/agentbricks/tests/unit_tests/dev_test.py @@ -665,7 +665,7 @@ def _fake_databricks(cmd, profile, **kwargs): @pytest.mark.parametrize( ("flag_profile", "expected", "note"), [ - (None, "proj", "Using profile 'proj' (from .env) · https://proj.example"), + (None, "proj", "Using profile 'proj' (from .env) · https://ws.example"), ("flag", "flag", "will run with profile 'flag' (from --profile), not .env's 'proj'"), ], ) diff --git a/integrations/agentbricks/tests/unit_tests/workspace_auth_test.py b/integrations/agentbricks/tests/unit_tests/workspace_auth_test.py index e0c7ab0b6..876c7b890 100644 --- a/integrations/agentbricks/tests/unit_tests/workspace_auth_test.py +++ b/integrations/agentbricks/tests/unit_tests/workspace_auth_test.py @@ -70,9 +70,30 @@ def test_config_file_path_expands_tilde_and_tolerates_percent_in_values(tmp_path [ (_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", ], - ids=["named-profile", "databricks-host", "default-section"], ) def test_workspace_for_profile(write_databrickscfg, monkeypatch, cfg, env_host, profile, expected): write_databrickscfg(cfg) @@ -86,6 +107,15 @@ def test_workspace_for_profile(write_databrickscfg, monkeypatch, cfg, env_host, 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 -------------------------------------------------------------------------- @@ -120,13 +150,16 @@ def test_validate_bounded_gives_up_after_the_timeout(monkeypatch): assert isinstance(error, TimeoutError) -def test_preflight_external_browser_fails_fast(write_databrickscfg, validated_profile, monkeypatch): +def test_preflight_external_browser_validates_cached_credentials( + write_databrickscfg, validated_profile +): write_databrickscfg("[prof]\nhost = https://ws.example\nauth_type = external-browser\n") - with pytest.raises(AgentCliError): - workspace_auth.preflight_auth("prof", "--profile") + assert ( + workspace_auth.preflight_auth("prof", "--profile") is validated_profile.return_value[0][0] + ) - validated_profile.assert_not_called() + validated_profile.assert_called_once_with("prof") # --- authenticate_profile -------------------------------------------------------------------- From 664bac76dc84df6a7ceb9469ef9a7dddbec24f3b Mon Sep 17 00:00:00 2001 From: jamesbxwu Date: Thu, 15 Oct 2026 15:21:03 +0000 Subject: [PATCH 04/12] Isolate AgentKit retry sleep mocks from background SDK calls --- .../agentbricks/tests/unit_tests/client_test.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/integrations/agentbricks/tests/unit_tests/client_test.py b/integrations/agentbricks/tests/unit_tests/client_test.py index d7794a6f7..c5230b481 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], @@ -826,7 +836,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) @@ -837,7 +846,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) @@ -850,7 +858,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"), @@ -872,7 +879,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) @@ -886,7 +892,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) @@ -903,7 +908,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) From b6b2cda30293bece4b9702a94652dfcfe3848fa1 Mon Sep 17 00:00:00 2001 From: Elaine Wang Date: Tue, 20 Oct 2026 19:05:21 +0000 Subject: [PATCH 05/12] [agentbricks] Read .env with python-dotenv and resolve the profile once per command Co-authored-by: Isaac --- integrations/agentbricks/pyproject.toml | 2 + .../src/databricks_agentbricks/cli/app.py | 102 ++++++++++-------- .../projects/env_file.py | 38 +++---- .../agentbricks/tests/unit_tests/auth_test.py | 8 +- .../tests/unit_tests/deploy_test.py | 1 + 5 files changed, 80 insertions(+), 71 deletions(-) diff --git a/integrations/agentbricks/pyproject.toml b/integrations/agentbricks/pyproject.toml index fad3b192d..390da15f5 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", "rich>=13.7", "tomli>=2.0", diff --git a/integrations/agentbricks/src/databricks_agentbricks/cli/app.py b/integrations/agentbricks/src/databricks_agentbricks/cli/app.py index 47bd2fb0a..a3c885d00 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/cli/app.py +++ b/integrations/agentbricks/src/databricks_agentbricks/cli/app.py @@ -8,7 +8,7 @@ import functools import pathlib -from typing import Optional +from typing import Any, Optional import click @@ -32,7 +32,6 @@ from databricks_agentbricks.projects.config import is_agentbricks_project from databricks_agentbricks.projects.env_file import read_env from databricks_agentbricks.projects.profiles import ( - SOURCE_FLAG, SOURCE_SDK_DEFAULT, ProfileInfo, auth_login_command, @@ -44,51 +43,65 @@ 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_info = resolve_profile(profile) self.output = output + 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] = {} - self.api_client_provider = self._client_provider(self.profile) @property def profile(self) -> Optional[str]: return self.profile_info.name - def use_project_profile(self, project_dir: pathlib.Path) -> None: - """Fold in the project's `.env` profile (see `resolve_profile`) unless `-p` was given. + 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 root group calls this for the current directory when it's an Agent Bricks project, and - commands with `--source` call it again for that project, whose `.env` then replaces the - current directory's. Only an explicit `-p` outranks the `.env`. + 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. """ - directory = project_dir.resolve() - if directory != self._project_dir: - self._project_dir = directory - self._project_env = read_env(directory / ".env") - if self.profile_info.source == SOURCE_FLAG: - return - resolved = resolve_profile(None, directory, self._project_env) - if resolved.name != self.profile_info.name: - self.api_client_provider = self._client_provider(resolved.name) - self.profile_info = resolved + 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 `use_project_profile` already read.""" + """`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 use_flag_profile(self, profile: str) -> None: - """Apply a `-p/--profile` given after the subcommand; like the global one, it beats `.env`.""" - if profile != self.profile: - self.api_client_provider = self._client_provider(profile) - self.profile_info = ProfileInfo(profile, SOURCE_FLAG) - 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: @@ -104,9 +117,6 @@ def announce_workspace(self, resolved_host: Optional[str] = None) -> None: if legacy_profile and self.profile_info.source == SOURCE_SDK_DEFAULT: warn_legacy_default_profile(legacy_profile, self.output) - def _client_provider(self, profile: Optional[str]) -> ApiClientProvider: - return ApiClientProvider(profile, on_first_use=self.announce_workspace) - def client(self) -> _AgentBricksApiClient: """Compatibility accessor for commands not yet migrated to ``api_client_provider``.""" return self.api_client_provider.get() @@ -164,7 +174,7 @@ def agentbricks(ctx: click.Context, profile: Optional[str], output: str) -> None if is_agentbricks_project(cwd): # Every command run inside a project uses its `.env` profile, so e.g. `endpoint invoke` # targets the same workspace that `deploy` just used. - ctx.obj.use_project_profile(cwd) + ctx.obj.cwd_project = cwd _REMOVED_COMMANDS = { @@ -192,7 +202,7 @@ def _apply_subcommand_profile( ) -> None: obj = ctx.find_object(CliContext) if value and obj is not None: - obj.use_flag_profile(value) + obj.subcommand_profile = value def _accept_profile_anywhere(command: click.Command, *, source_profile: bool = True) -> None: @@ -204,8 +214,10 @@ def _accept_profile_anywhere(command: click.Command, *, source_profile: bool = T for subcommand in command.commands.values(): _accept_profile_anywhere(subcommand, source_profile=source_profile) return - if source_profile and any(param.name == "source" for param in command.params): - _use_source_project_profile(command) + _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( @@ -219,22 +231,28 @@ def _accept_profile_anywhere(command: click.Command, *, source_profile: bool = T ) -def _use_source_project_profile(command: click.Command) -> None: - """Resolve the profile from `--source`'s project, so e.g. `tracing bind --source proj` reads - the same workspace that `deploy --source proj` writes to (only `-p` outranks it).""" +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_source_profile(*args, **kwargs): + def with_selected_profile(*args, **kwargs): obj = click.get_current_context().find_object(CliContext) - source = kwargs.get("source") - if obj is not None and source and is_agentbricks_project(pathlib.Path(source)): - obj.use_project_profile(pathlib.Path(source)) + 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_source_profile + command.callback = with_selected_profile agentbricks.add_command(init) diff --git a/integrations/agentbricks/src/databricks_agentbricks/projects/env_file.py b/integrations/agentbricks/src/databricks_agentbricks/projects/env_file.py index 71babc640..70436f804 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/projects/env_file.py +++ b/integrations/agentbricks/src/databricks_agentbricks/projects/env_file.py @@ -2,13 +2,19 @@ from __future__ import annotations +import io import os import pathlib import re import tempfile -# Same shape the reader accepts, so a line the CLI honors is a line it can rewrite. -_PROFILE_LINE = re.compile(r"\s*(export\s+)?DATABRICKS_CONFIG_PROFILE\s*=") +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" @@ -40,7 +46,7 @@ def _write_text_atomically(path: pathlib.Path, text: str) -> None: def read_env(path: pathlib.Path) -> dict[str, str]: - """Minimal KEY=VALUE reader for a project `.env` (python-dotenv is not a CLI dependency).""" + """Read a project `.env` the way the generated agent's `load_dotenv` does.""" try: text = _read_text(path) except OSError: @@ -49,29 +55,9 @@ def read_env(path: pathlib.Path) -> dict[str, str]: def _parse_env(text: str) -> dict[str, str]: - values: dict[str, str] = {} - for line in _lines(text): - line = line.strip() - if not line or line.startswith("#") or "=" not in line: - continue - if line.startswith("export "): - line = line[len("export ") :].strip() - key, _, value = line.partition("=") - key, value = key.strip(), value.strip() - if value[:1] in ("'", '"'): - end = value.find(value[0], 1) - if end != -1: - # A quoted value ends at its matching close quote: a `#` inside is literal, and - # anything after (an inline comment) is dropped. - value = value[1:end] - else: - value = re.split(r"\s+#", value, maxsplit=1)[0].rstrip() - else: - # Match dotenv for unquoted values: a `#` that follows whitespace starts a comment. - value = re.split(r"\s+#", value, maxsplit=1)[0].rstrip() - if key: - values[key] = value - return values + # 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))) + 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]: diff --git a/integrations/agentbricks/tests/unit_tests/auth_test.py b/integrations/agentbricks/tests/unit_tests/auth_test.py index 3a5c92984..0b7bb85ce 100644 --- a/integrations/agentbricks/tests/unit_tests/auth_test.py +++ b/integrations/agentbricks/tests/unit_tests/auth_test.py @@ -60,7 +60,9 @@ def fake_api_client(monkeypatch: pytest.MonkeyPatch) -> mock.Mock: def _ctx(profile=None, output="text") -> CliContext: - return CliContext(profile, output) + ctx = CliContext(profile, output) + ctx.select_profile() + return ctx def _invoke_profile(args, **ctx_kwargs): @@ -82,8 +84,8 @@ def test_context_precedence_is_flag_then_env_file_then_env_var(tmp_path, monkeyp _write_env(tmp_path, "proj") ambient, project, flagged = _ctx(), _ctx(), _ctx("from-flag") - project.use_project_profile(tmp_path) - flagged.use_project_profile(tmp_path) + project.select_profile(tmp_path) + flagged.select_profile(tmp_path) assert ambient.profile_info == ProfileInfo("from-env-var", "DATABRICKS_CONFIG_PROFILE") assert project.profile_info == ProfileInfo("proj", ".env") diff --git a/integrations/agentbricks/tests/unit_tests/deploy_test.py b/integrations/agentbricks/tests/unit_tests/deploy_test.py index b203cafe9..dea2d8fff 100644 --- a/integrations/agentbricks/tests/unit_tests/deploy_test.py +++ b/integrations/agentbricks/tests/unit_tests/deploy_test.py @@ -456,6 +456,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) From 7ac0f53b0f83ca19aadbc990b81b2aeee1de68e3 Mon Sep 17 00:00:00 2001 From: Elaine Wang Date: Tue, 20 Oct 2026 21:06:40 +0000 Subject: [PATCH 06/12] [agentbricks] Move cached browser auth out of the profile resolution change Co-authored-by: Isaac --- integrations/agentbricks/README.md | 3 - integrations/agentbricks/cli.md | 3 +- .../clients/cached_browser_auth.py | 60 ------ .../services/workspace_auth.py | 56 +++-- .../src/databricks_agentkit/_api_client.py | 20 +- .../unit_tests/cached_browser_auth_test.py | 202 ------------------ .../tests/unit_tests/client_test.py | 13 +- .../tests/unit_tests/workspace_auth_test.py | 11 +- 8 files changed, 47 insertions(+), 321 deletions(-) delete mode 100644 integrations/agentbricks/src/databricks_agentbricks/clients/cached_browser_auth.py delete mode 100644 integrations/agentbricks/tests/unit_tests/cached_browser_auth_test.py diff --git a/integrations/agentbricks/README.md b/integrations/agentbricks/README.md index 8df940e3f..2f369263a 100644 --- a/integrations/agentbricks/README.md +++ b/integrations/agentbricks/README.md @@ -186,9 +186,6 @@ 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. -An external-browser profile's valid or refreshable Python SDK token cache is reused without -opening a browser. If that cache is missing or unusable, run the suggested `auth login` command -in an interactive terminal. `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 576bcc3ac..d974d5a0f 100644 --- a/integrations/agentbricks/cli.md +++ b/integrations/agentbricks/cli.md @@ -44,8 +44,7 @@ global default profile. The profile is resolved in this order, most to least spe 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. Authentication preflight reuses a valid or refreshable external-browser Python SDK token -cache without opening a browser; a missing or unusable cache requires `agentbricks auth login`. +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 diff --git a/integrations/agentbricks/src/databricks_agentbricks/clients/cached_browser_auth.py b/integrations/agentbricks/src/databricks_agentbricks/clients/cached_browser_auth.py deleted file mode 100644 index b6724abe5..000000000 --- a/integrations/agentbricks/src/databricks_agentbricks/clients/cached_browser_auth.py +++ /dev/null @@ -1,60 +0,0 @@ -"""Use the Python SDK's external-browser cache without ever starting OAuth consent.""" - -from __future__ import annotations - -from typing import TYPE_CHECKING, Optional - -from databricks_agentkit._api_client import _AgentBricksApiClient - -if TYPE_CHECKING: - from databricks.sdk.config import Config - from databricks.sdk.credentials_provider import CredentialsProvider - - -def _cached_credentials(config: Config) -> CredentialsProvider: - from databricks.sdk import oauth - - # Group assumption was added after the minimum supported SDK (0.124). - group_id = getattr(config, "group_id", None) - # Match the SDK's external_browser cache identity, including Azure Entra profiles, scopes, - # and group assumption. Using a different identity could reuse another profile's session. - if config.client_id: - client_id, client_secret = config.client_id, config.client_secret - endpoints = config.databricks_oidc_endpoints - elif config.azure_client_id: - if group_id: - raise ValueError("external-browser with Azure Entra ID cannot assume a group") - client_id, client_secret = config.azure_client_id, config.azure_client_secret - endpoints = oauth.get_azure_entra_id_workspace_endpoints(config.host) - else: - client_id, client_secret = "databricks-cli", None - endpoints = config.databricks_oidc_endpoints - if endpoints is None: - raise ValueError("The workspace does not support OAuth") - scopes = list(config.get_scopes()) - if not config.disable_oauth_refresh_token and "offline_access" not in scopes: - scopes.append("offline_access") - cache = oauth.TokenCache( - host=config.host, - oidc_endpoints=endpoints, - client_id=client_id, - client_secret=client_secret, - redirect_url="http://localhost:8020", - scopes=scopes, - profile=config.profile, - **({"group_id": group_id} if group_id else {}), - ) - credentials = cache.load() - if credentials is None: - raise ValueError("No cached external-browser session is available") - # A refresh failure propagates to the caller; unlike the SDK's interactive strategy there - # is deliberately no fallback to OAuthClient, consent, or a browser listener here. - credentials.token() - return credentials(config) - - -def cached_browser_client(profile: Optional[str]) -> _AgentBricksApiClient: - from databricks.sdk.credentials_provider import credentials_strategy - - strategy = credentials_strategy("external-browser", ["host"])(_cached_credentials) - return _AgentBricksApiClient(profile, credentials_strategy=strategy) diff --git a/integrations/agentbricks/src/databricks_agentbricks/services/workspace_auth.py b/integrations/agentbricks/src/databricks_agentbricks/services/workspace_auth.py index 22a2481ae..53b7dd55e 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/services/workspace_auth.py +++ b/integrations/agentbricks/src/databricks_agentbricks/services/workspace_auth.py @@ -11,7 +11,6 @@ import threading from typing import Optional -from databricks_agentbricks.clients.cached_browser_auth import cached_browser_client from databricks_agentbricks.clients.databricks_cli import run_databricks_login from databricks_agentbricks.errors import AgentCliError from databricks_agentbricks.projects.profiles import ( @@ -23,11 +22,7 @@ def _validate_profile(profile: Optional[str]) -> tuple[_AgentBricksApiClient, str]: - client = ( - cached_browser_client(profile) - if sdk_auth_type(profile) == "external-browser" - else _AgentBricksApiClient(profile) - ) + client = _AgentBricksApiClient(profile) return client, client.current_user @@ -96,14 +91,23 @@ def _unreachable_error(profile: Optional[str], source: str, error: Exception) -> _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`. - Auth/network resolution can block, so the call is bounded. External-browser profiles use - cached credentials only and never start consent. Daemon because executor threads are joined - at interpreter exit and a hung request must not block the CLI's exit too. + A browser-OAuth flow can block forever, so the call is bounded. Daemon because executor + threads are joined at interpreter exit and a hung flow 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]]] = [] @@ -130,12 +134,18 @@ def preflight_auth(profile: Optional[str], source: str) -> _AgentBricksApiClient agent sits stuck with no error. The validation call itself can hang that way, so it runs bounded (see `_validate_bounded`). """ + if _needs_browser_login(profile): + # The bounded validation would open a browser flow this command will abandon at its + # timeout, so fail before it starts — interactive or not. + raise _not_signed_in_error(profile, source) + client_and_user, error = _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 + # A timeout here is most likely a browser flow nobody completed, so it reads as signed out too. raise _not_signed_in_error(profile, source) from error @@ -149,19 +159,25 @@ def authenticate_profile( profile fails with a hint instead of opening one. """ label = _profile_label(profile, source) - client_and_user, error = _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: - # A timeout or an unexpected failure won't be fixed by a login; don't open one. - raise AgentCliError(f"Could not validate {label}: {error}") from error + error: Optional[Exception] + if _needs_browser_login(profile): + # Validation would open a browser flow it then abandons; head straight for the login. + error = None + else: + client_and_user, error = _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: + # A timeout or an unexpected failure won't be fixed by a login; don't open one. + 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}: {error}", + f"Could not validate {label}: {detail}", hint=_sign_in_hint(profile, needs_terminal=True), ) from error run_databricks_login(profile) diff --git a/integrations/agentbricks/src/databricks_agentkit/_api_client.py b/integrations/agentbricks/src/databricks_agentkit/_api_client.py index 13fe3d2cb..eae2c9ebb 100644 --- a/integrations/agentbricks/src/databricks_agentkit/_api_client.py +++ b/integrations/agentbricks/src/databricks_agentkit/_api_client.py @@ -21,7 +21,6 @@ if TYPE_CHECKING: from databricks.sdk import WorkspaceClient - from databricks.sdk.credentials_provider import CredentialsStrategy _BASE = "/api/2.0/agents" _MCP_SERVICES_PATH = "/api/2.1/unity-catalog/mcp-services" @@ -147,17 +146,12 @@ def _bound_retry_timeout(client: WorkspaceClient) -> WorkspaceClient: return client -def _workspace_client( - profile: Optional[str], *, credentials_strategy: Optional[CredentialsStrategy] = None -) -> WorkspaceClient: +def _workspace_client(profile: Optional[str]) -> WorkspaceClient: # Imported here (not at module top) so the ~0.7s databricks.sdk import is paid only when a # command actually builds a client, not on every CLI invocation. from databricks.sdk import WorkspaceClient - auth_kwargs: dict[str, Any] = ( - {"credentials_strategy": credentials_strategy} if credentials_strategy is not None else {} - ) - client = _bound_retry_timeout(WorkspaceClient(profile=profile, **auth_kwargs)) + client = _bound_retry_timeout(WorkspaceClient(profile=profile)) if not profile or not client.config.workspace_id: return client @@ -172,7 +166,6 @@ def _workspace_client( profile=profile, host=configured_host, custom_headers={"X-Databricks-Org-Id": workspace_id}, - **auth_kwargs, ) ) @@ -185,18 +178,11 @@ def __init__( profile: Optional[str] = None, *, workspace_client: Optional[WorkspaceClient] = None, - credentials_strategy: Optional[CredentialsStrategy] = None, ) -> None: if profile is not None and workspace_client is not None: raise ValueError("profile and workspace_client are mutually exclusive") - if credentials_strategy is not None and workspace_client is not None: - raise ValueError("credentials_strategy and workspace_client are mutually exclusive") try: - self._w = workspace_client or ( - _workspace_client(profile, credentials_strategy=credentials_strategy) - if credentials_strategy is not None - else _workspace_client(profile) - ) + self._w = workspace_client or _workspace_client(profile) except Exception as exc: # noqa: BLE001 - surfaced as a clean CLI error raise AgentCliError( f"Could not initialize Databricks auth: {exc}", diff --git a/integrations/agentbricks/tests/unit_tests/cached_browser_auth_test.py b/integrations/agentbricks/tests/unit_tests/cached_browser_auth_test.py deleted file mode 100644 index fe8d8f66a..000000000 --- a/integrations/agentbricks/tests/unit_tests/cached_browser_auth_test.py +++ /dev/null @@ -1,202 +0,0 @@ -"""Cached OAuth validation with synthetic credentials and no network/consent fallback.""" - -from __future__ import annotations - -import inspect -import json -import pathlib -from datetime import datetime, timedelta -from types import SimpleNamespace -from typing import Any -from unittest import mock - -import pytest -from databricks.sdk import oauth -from databricks.sdk.config import Config -from databricks.sdk.service.iam import CurrentUserAPI, User - -from databricks_agentbricks.clients.cached_browser_auth import _cached_credentials -from databricks_agentbricks.errors import AgentCliError -from databricks_agentbricks.services import workspace_auth - - -@pytest.fixture(autouse=True) -def isolate_auth_and_forbid_network(monkeypatch): - # Config/HOME already point at temp files (conftest); also remove alternate SDK credential - # env vars so no real secret can influence these tests. The only tokens below are synthetic. - for attr in Config.attributes(): - for name in (attr.env, *getattr(attr, "env_aliases", ())): - if name and name != "DATABRICKS_CONFIG_FILE": - monkeypatch.delenv(name, raising=False) - with ( - mock.patch("requests.Session.request", side_effect=AssertionError("unexpected network")), - mock.patch.object(oauth, "OAuthClient") as consent_client, - mock.patch.object(oauth.Consent, "launch_external_browser") as browser, - ): - yield - consent_client.assert_not_called() - browser.assert_not_called() - - -@pytest.fixture -def config(): - return SimpleNamespace( - host="https://ws.example", - profile="prof", - client_id=None, - client_secret=None, - azure_client_id=None, - azure_client_secret=None, - databricks_oidc_endpoints=oauth.OidcEndpoints( - authorization_endpoint="https://ws.example/oidc/v1/authorize", - token_endpoint="https://ws.example/oidc/v1/token", - ), - get_scopes=lambda: ["all-apis"], - disable_oauth_refresh_token=False, - ) - - -def _seed_cache(config, *, expired=False, **identity): - cache_kwargs: dict[str, Any] = { - "host": config.host, - "oidc_endpoints": config.databricks_oidc_endpoints, - "client_id": "databricks-cli", - "redirect_url": "http://localhost:8020", - "scopes": ["all-apis", "offline_access"], - "profile": config.profile, - **identity, - } - cache = oauth.TokenCache(**cache_kwargs) - token = oauth.Token( - access_token="synthetic-cached-token", - refresh_token="synthetic-refresh-token", - token_type="Bearer", - expiry=datetime.now() + timedelta(hours=-1 if expired else 1), - ) - path = pathlib.Path(cache.filename) - path.parent.mkdir(parents=True, exist_ok=True) - path.write_text(json.dumps({"token": token.as_dict()})) - return path - - -@pytest.fixture -def sdk_workspace(config, write_databrickscfg): - write_databrickscfg("[prof]\nhost = https://ws.example\nauth_type = external-browser\n") - metadata = SimpleNamespace( - account_id=None, - workspace_id=None, - oidc_endpoint="", - cloud=None, - token_federation_default_oidc_audiences=[], - ) - with ( - mock.patch("databricks.sdk.config.get_host_metadata", return_value=metadata), - mock.patch.object( - Config, - "databricks_oidc_endpoints", - new_callable=mock.PropertyMock, - return_value=config.databricks_oidc_endpoints, - ), - mock.patch.object( - CurrentUserAPI, "me", return_value=User(user_name="cached@example.com") - ) as current_user, - ): - yield current_user - - -@pytest.mark.parametrize("expired", [False, True], ids=["valid", "refreshable"]) -def test_preflight_and_login_accept_sdk_cached_session( - config, sdk_workspace, databricks_login, expired -): - _seed_cache(config, expired=expired) - refreshed = oauth.Token( - access_token="synthetic-refreshed-token", - refresh_token="synthetic-refresh-token", - token_type="Bearer", - expiry=datetime.now() + timedelta(hours=1), - ) - with mock.patch.object(oauth, "retrieve_token", return_value=refreshed) as refresh: - client = workspace_auth.preflight_auth("prof", "--profile") - expected_token = "synthetic-refreshed-token" if expired else "synthetic-cached-token" - assert client.workspace_client.config.authenticate() == { - "Authorization": f"Bearer {expected_token}" - } - if expired: - refresh.assert_called_once_with( - client_id="databricks-cli", - client_secret=None, - token_url=config.databricks_oidc_endpoints.token_endpoint, - params={ - "grant_type": "refresh_token", - "refresh_token": "synthetic-refresh-token", - }, - use_params=True, - headers={}, - ) - else: - refresh.assert_not_called() - _, user = workspace_auth.authenticate_profile("prof", "argument", interactive=True) - - assert user == "cached@example.com" - assert sdk_workspace.call_count == 2 - databricks_login.assert_not_called() - - -@pytest.mark.parametrize("state", ["missing", "invalid", "expired"]) -def test_unusable_cache_fails_without_starting_consent( - config, sdk_workspace, databricks_login, state -): - if state != "missing": - path = _seed_cache(config, expired=state == "expired") - if state == "invalid": - path.write_text("invalid JSON") - with mock.patch.object( - oauth, "retrieve_token", side_effect=ValueError("synthetic revoked token") - ): - with pytest.raises(AgentCliError, match="isn't authenticated"): - workspace_auth.preflight_auth("prof", "--profile") - with pytest.raises(AgentCliError) as raised: - workspace_auth.authenticate_profile("prof", "argument", interactive=False) - - assert raised.value.hint is not None - assert "interactive terminal" in raised.value.hint - sdk_workspace.assert_not_called() - databricks_login.assert_not_called() - - -@pytest.mark.parametrize( - "identity", - [ - {"profile": "other"}, - {"host": "https://other.example"}, - {"scopes": ["other-scope", "offline_access"]}, - {"group_id": "other-group"}, - ], - ids=["profile", "host", "scopes", "group"], -) -def test_a_different_cache_identity_cannot_authenticate(config, identity): - if "group_id" in identity and "group_id" not in inspect.signature(oauth.TokenCache).parameters: - pytest.skip("Group assumption is unavailable in the minimum supported SDK") - _seed_cache(config, **identity) - - with pytest.raises(ValueError, match="No cached external-browser session"): - _cached_credentials(config) - - -def test_cache_identity_includes_group_when_supported(config): - if "group_id" not in inspect.signature(oauth.TokenCache).parameters: - pytest.skip("Group assumption is unavailable in the minimum supported SDK") - config.group_id = "expected-group" - _seed_cache(config, group_id=config.group_id) - - assert _cached_credentials(config)() == {"Authorization": "Bearer synthetic-cached-token"} - - -def test_minimum_sdk_cache_constructor_does_not_need_group_argument(config): - # SDK 0.124 has this exact constructor shape and no Config.group_id attribute. - def old_cache(*, host, oidc_endpoints, client_id, client_secret, redirect_url, scopes, profile): - return SimpleNamespace(load=lambda: credentials) - - credentials = mock.Mock(return_value=lambda: {"Authorization": "Bearer synthetic-token"}) - with mock.patch.object(oauth, "TokenCache", side_effect=old_cache): - assert _cached_credentials(config)() == {"Authorization": "Bearer synthetic-token"} diff --git a/integrations/agentbricks/tests/unit_tests/client_test.py b/integrations/agentbricks/tests/unit_tests/client_test.py index cfbcb70e0..d9b0b5f93 100644 --- a/integrations/agentbricks/tests/unit_tests/client_test.py +++ b/integrations/agentbricks/tests/unit_tests/client_test.py @@ -936,8 +936,7 @@ def test_do_does_not_retry_transient_error_for_pop(workspace_client, sleep): sleep.assert_not_called() -@pytest.mark.parametrize("strategy", [None, mock.Mock()], ids=["default-auth", "injected-auth"]) -def test_account_routed_profile_uses_configured_host_and_workspace_header(strategy): +def test_account_routed_profile_uses_configured_host_and_workspace_header(): resolved = mock.Mock() resolved.config.host = "https://workspace.example.com" resolved.config.workspace_id = "123" @@ -950,21 +949,15 @@ def test_account_routed_profile_uses_configured_host_and_workspace_header(strate return_value="https://account.example.com", ), ): - client = ( - _workspace_client("p", credentials_strategy=strategy) - if strategy is not None - else _workspace_client("p") - ) + client = _workspace_client("p") assert client is routed - auth_kwargs = {"credentials_strategy": strategy} if strategy is not None else {} assert wc.call_args_list == [ - mock.call(profile="p", **auth_kwargs), + mock.call(profile="p"), mock.call( profile="p", host="https://account.example.com", custom_headers={"X-Databricks-Org-Id": "123"}, - **auth_kwargs, ), ] diff --git a/integrations/agentbricks/tests/unit_tests/workspace_auth_test.py b/integrations/agentbricks/tests/unit_tests/workspace_auth_test.py index 876c7b890..916cd920c 100644 --- a/integrations/agentbricks/tests/unit_tests/workspace_auth_test.py +++ b/integrations/agentbricks/tests/unit_tests/workspace_auth_test.py @@ -150,16 +150,13 @@ def test_validate_bounded_gives_up_after_the_timeout(monkeypatch): assert isinstance(error, TimeoutError) -def test_preflight_external_browser_validates_cached_credentials( - write_databrickscfg, validated_profile -): +def test_preflight_external_browser_fails_fast(write_databrickscfg, validated_profile, monkeypatch): write_databrickscfg("[prof]\nhost = https://ws.example\nauth_type = external-browser\n") - assert ( - workspace_auth.preflight_auth("prof", "--profile") is validated_profile.return_value[0][0] - ) + with pytest.raises(AgentCliError): + workspace_auth.preflight_auth("prof", "--profile") - validated_profile.assert_called_once_with("prof") + validated_profile.assert_not_called() # --- authenticate_profile -------------------------------------------------------------------- From f615402216a9f5ce48ae59c60676052c22d17a13 Mon Sep 17 00:00:00 2001 From: Elaine Wang Date: Wed, 21 Oct 2026 00:39:04 +0000 Subject: [PATCH 07/12] [agentbricks] Move auth login target choice and auth output out of cli/auth.py Co-authored-by: Isaac --- .../src/databricks_agentbricks/cli/auth.py | 98 +++---------------- .../presentation/profile.py | 91 +++++++++++++++++ .../projects/profiles.py | 17 ++++ 3 files changed, 124 insertions(+), 82 deletions(-) diff --git a/integrations/agentbricks/src/databricks_agentbricks/cli/auth.py b/integrations/agentbricks/src/databricks_agentbricks/cli/auth.py index 47e56cdce..6954c7873 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/cli/auth.py +++ b/integrations/agentbricks/src/databricks_agentbricks/cli/auth.py @@ -11,20 +11,23 @@ from databricks_agentbricks.clients.databricks_cli import run_databricks_login from databricks_agentbricks.errors import AgentCliError -from databricks_agentbricks.presentation import render 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 is_agentbricks_project from databricks_agentbricks.projects.env_file import read_env, update_env_profile from databricks_agentbricks.projects.profiles import ( - SOURCE_ENV_FILE, SOURCE_FLAG, ProfileInfo, auth_login_command, auth_switch_command, + login_target, profile_host, resolve_profile, workspace_for_profile, @@ -81,29 +84,6 @@ def _resolve_for(obj, directory: pathlib.Path, in_project: bool) -> ProfileInfo: return resolve_profile(None, directory if in_project else None) -def _login_target( - obj, name: Optional[str], 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 but a terminal can pick one. - """ - if name: - return name, "argument" - if obj.profile_info.source == SOURCE_FLAG and obj.profile_info.name: - return obj.profile_info.name, SOURCE_FLAG - if env_profile: - return env_profile, SOURCE_ENV_FILE - if can_prompt_login(): - return None - 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)}`.", - ) - - @click.group("auth") def auth_group() -> None: """Choose, inspect, and sign in to the Databricks profile a project uses. @@ -136,28 +116,7 @@ def auth_switch(obj, name: str, source: str) -> None: directory = _require_project(source) changed = update_env_profile(directory, name) user = sign_in_chosen_profile(obj, name, "argument") - host = profile_host(name) - if obj.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")], - ) + show_profile_switched(directory, name, changed, profile_host(name), user, obj.output) @auth_group.command("describe") @@ -180,24 +139,7 @@ def auth_describe(obj, source: str) -> None: info = _resolve_for(obj, directory, in_project) host, workspace = workspace_for_profile(info.name) project = str(directory) if in_project else None - if obj.output == "json": - render.emit_json( - { - "profile": info.name, - "source": info.source, - "workspace_host": host, - "project": project, - } - ) - return - fields = { - "Profile": info.name or "none (Databricks SDK default)", - "Source": info.source, - } - if workspace: - fields["Workspace"] = workspace - fields["Project"] = project or "not in a project" - render.detail("Profile", info.name or "SDK default", fields) + show_profile_description(info.name, info.source, host, workspace, project, obj.output) @auth_group.command("login") @@ -228,28 +170,20 @@ def auth_login(obj, name: Optional[str], source: str) -> None: if is_agentbricks_project(directory) else None ) - target = _login_target(obj, name, env_profile) + 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) - if obj.output == "json": - render.emit_json({"profile": None, "source": "databricks auth login"}) - else: - render.success("Logged in with `databricks auth login`") + show_databricks_login_result(obj.output) return chosen, origin = target announce_login(obj, chosen, origin) client, user = authenticate_profile(chosen, origin, interactive=can_prompt_login()) - if obj.output == "json": - render.emit_json({"profile": chosen, "source": origin, "user": user, "host": client.host}) - return - render.success( - f"Logged in to '{chosen}'", - fields={"Profile": chosen, "Source": origin, "User": user, "Workspace": client.host}, - ) - if env_profile and env_profile != chosen: - render.console().print( - f"[dim]This project still uses profile '{env_profile}' (from .env). " - f"Run `{auth_switch_command(chosen)}` to switch it.[/]" - ) + show_login_result(chosen, origin, user, client.host, env_profile, obj.output) diff --git a/integrations/agentbricks/src/databricks_agentbricks/presentation/profile.py b/integrations/agentbricks/src/databricks_agentbricks/presentation/profile.py index 4bc2550ac..7e8c82d04 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/presentation/profile.py +++ b/integrations/agentbricks/src/databricks_agentbricks/presentation/profile.py @@ -2,6 +2,8 @@ from __future__ import annotations +import pathlib + import click from databricks_agentbricks.presentation import render @@ -13,6 +15,95 @@ ) +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: diff --git a/integrations/agentbricks/src/databricks_agentbricks/projects/profiles.py b/integrations/agentbricks/src/databricks_agentbricks/projects/profiles.py index b0d8d797b..4184af131 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/projects/profiles.py +++ b/integrations/agentbricks/src/databricks_agentbricks/projects/profiles.py @@ -74,6 +74,23 @@ def resolve_profile( 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( From bcbc8fe1f71b1e4df872ffcab4c52b808eba8247 Mon Sep 17 00:00:00 2001 From: Elaine Wang Date: Wed, 21 Oct 2026 23:24:46 +0000 Subject: [PATCH 08/12] [agentbricks] Read profile hosts from one Databricks config parser Co-authored-by: Isaac --- .../src/databricks_agentkit/_api_client.py | 35 +++---------------- 1 file changed, 4 insertions(+), 31 deletions(-) diff --git a/integrations/agentbricks/src/databricks_agentkit/_api_client.py b/integrations/agentbricks/src/databricks_agentkit/_api_client.py index eae2c9ebb..68617e218 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 @@ -106,35 +103,11 @@ def memory_entry_path(store: str, entry: str) -> str: return f"{memory_store_path(store)}/entries/{entry}" -def _databricks_config_parser( - default_section: str = configparser.DEFAULTSECT, -) -> configparser.ConfigParser: - """The parsed Databricks config file, or an empty parser on any problem (never raises). - - expanduser: the SDK expands `~` in DATABRICKS_CONFIG_FILE; without it a tilde path reads - nothing. No interpolation: values are taken verbatim — a `%` must not raise at get-time. - """ - parser = configparser.ConfigParser(interpolation=None, default_section=default_section) - try: - parser.read( - pathlib.Path( - os.getenv("DATABRICKS_CONFIG_FILE", str(pathlib.Path.home() / ".databrickscfg")) - ).expanduser() - ) - except (OSError, configparser.Error): - pass - return parser - - -def _profile_section(profile: str) -> dict[str, str]: - """`profile`'s own keys, as the SDK reads them: a named profile doesn't inherit [DEFAULT].""" - # With an unused fallback section, [DEFAULT] parses as an ordinary section like any other. - parser = _databricks_config_parser(default_section="\0no-fallback") - return dict(parser[profile]) if parser.has_section(profile) else {} - - def _profile_host(profile: str) -> Optional[str]: - return _profile_section(profile).get("host") + # 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: From 0abb94ca0f63d73ee4e5423567627f62ffdc1f6a Mon Sep 17 00:00:00 2001 From: Elaine Wang Date: Wed, 21 Oct 2026 23:44:02 +0000 Subject: [PATCH 09/12] Align profile CLI docs and inject workspace auth collaborators --- integrations/agentbricks/cli.md | 4 + .../clients/api_client_provider.py | 6 + .../services/workspace_auth.py | 133 ++++++++++-------- .../tests/unit_tests/workspace_auth_test.py | 19 +++ 4 files changed, 103 insertions(+), 59 deletions(-) diff --git a/integrations/agentbricks/cli.md b/integrations/agentbricks/cli.md index 57462e960..a0693df78 100644 --- a/integrations/agentbricks/cli.md +++ b/integrations/agentbricks/cli.md @@ -70,6 +70,10 @@ A trailing `-p` goes after the full subcommand, not after a group name — `agen ## 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. 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 7fef07d83..c510497aa 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/clients/api_client_provider.py +++ b/integrations/agentbricks/src/databricks_agentbricks/clients/api_client_provider.py @@ -16,6 +16,12 @@ 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.""" diff --git a/integrations/agentbricks/src/databricks_agentbricks/services/workspace_auth.py b/integrations/agentbricks/src/databricks_agentbricks/services/workspace_auth.py index 53b7dd55e..3c9656d3b 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/services/workspace_auth.py +++ b/integrations/agentbricks/src/databricks_agentbricks/services/workspace_auth.py @@ -9,8 +9,9 @@ import enum import threading -from typing import Optional +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 ( @@ -21,11 +22,6 @@ from databricks_agentkit._api_client import _AgentBricksApiClient -def _validate_profile(profile: Optional[str]) -> tuple[_AgentBricksApiClient, str]: - client = _AgentBricksApiClient(profile) - return client, client.current_user - - 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 @@ -126,64 +122,83 @@ def _validate() -> None: return outcome[0] -def preflight_auth(profile: Optional[str], source: str) -> _AgentBricksApiClient: - """Fail fast when the resolved profile can't authenticate, before a command spins anything up. +ValidationResult = tuple[Optional[tuple[_AgentBricksApiClient, str]], Optional[Exception]] - Returns the validated client so the command can reuse it. Without this, an unauthenticated - `external-browser` profile hangs on a browser OAuth flow that never completes — the local - agent sits stuck with no error. The validation call itself can hang that way, so it runs - bounded (see `_validate_bounded`). - """ - if _needs_browser_login(profile): - # The bounded validation would open a browser flow this command will abandon at its - # timeout, so fail before it starts — interactive or not. - raise _not_signed_in_error(profile, source) - client_and_user, error = _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 - # A timeout here is most likely a browser flow nobody completed, so it reads as signed out too. - raise _not_signed_in_error(profile, source) from error +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 authenticate_profile( - profile: str, source: str, *, interactive: bool -) -> tuple[_AgentBricksApiClient, str]: - """Validate `profile`, signing in through `databricks auth login` once if it isn't yet. + 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) - Returns the client and the signed-in user; `source` is only named in errors. Without - `interactive` (no terminal, or CI) nobody can complete a browser login, so an unauthenticated - profile fails with a hint instead of opening one. - """ - label = _profile_label(profile, source) - error: Optional[Exception] - if _needs_browser_login(profile): - # Validation would open a browser flow it then abandons; head straight for the login. - error = None - else: - client_and_user, error = _validate_bounded(profile) + client_and_user, error = self._validate_bounded(profile) if client_and_user is not None: - return client_and_user + return client_and_user[0] assert error is not None - failure = _classify_validation_error(error) - if failure is _ValidationFailure.UNREACHABLE: + if _classify_validation_error(error) is _ValidationFailure.UNREACHABLE: raise _unreachable_error(profile, source, error) from error - if failure is _ValidationFailure.OTHER: - # A timeout or an unexpected failure won't be fixed by a login; don't open one. - 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 - run_databricks_login(profile) - client_and_user, error = _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 + 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/tests/unit_tests/workspace_auth_test.py b/integrations/agentbricks/tests/unit_tests/workspace_auth_test.py index 916cd920c..76e503b51 100644 --- a/integrations/agentbricks/tests/unit_tests/workspace_auth_test.py +++ b/integrations/agentbricks/tests/unit_tests/workspace_auth_test.py @@ -194,6 +194,25 @@ def test_authenticate_profile_logs_in_then_revalidates(monkeypatch, databricks_l ) +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_reports_a_missing_databricks_cli(fail_validation, databricks_login): fail_validation(Unauthenticated("expired")) databricks_login.side_effect = FileNotFoundError From 31ffe08d60275222d4383f13cd09680ba895ce44 Mon Sep 17 00:00:00 2001 From: Elaine Wang Date: Thu, 22 Oct 2026 22:59:11 +0000 Subject: [PATCH 10/12] Fix model upgrade JSON output tests --- .../agentbricks/tests/unit_tests/models_cli_test.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/integrations/agentbricks/tests/unit_tests/models_cli_test.py b/integrations/agentbricks/tests/unit_tests/models_cli_test.py index 6472b8875..6c571c450 100644 --- a/integrations/agentbricks/tests/unit_tests/models_cli_test.py +++ b/integrations/agentbricks/tests/unit_tests/models_cli_test.py @@ -982,7 +982,8 @@ def test_upgrade_with_apply_auto_waits_and_applies(tmp_path, stub_job, stub_prom _Ctx(client, "json"), ) # fmt: skip assert result.exit_code == 0, result.output - assert json.loads(result.output)["applied"]["changed"] is True + assert json.loads(result.stdout)["applied"]["changed"] is True + assert "Waiting for the upgrade job" in result.stderr assert client.model == "system.ai.claude-haiku-4-5" assert len(model_upgrade.read_history(project)["actions"]) == 1 # rollback can undo it @@ -1011,7 +1012,8 @@ def get_run(self, run_id): _Ctx(client, "json"), ) # fmt: skip assert result.exit_code == 0, result.output - out = json.loads(result.output) + out = json.loads(result.stdout) + assert "Waiting for the upgrade job" in result.stderr assert out["run"]["run_id"] == 42 assert stub_job["fetch"][1] == out["upgrade_id"] assert client.model == "system.ai.claude-haiku-4-5" From 3786c31d23aff47303fc169a01b1c732486d1647 Mon Sep 17 00:00:00 2001 From: Elaine Wang Date: Thu, 22 Oct 2026 23:44:07 +0000 Subject: [PATCH 11/12] Clarify Agent Bricks profile user journey --- integrations/agentbricks/README.md | 26 +++++++++++++++----------- integrations/agentbricks/cli.md | 19 +++++++++++++------ 2 files changed, 28 insertions(+), 17 deletions(-) diff --git a/integrations/agentbricks/README.md b/integrations/agentbricks/README.md index 21bad2219..34690063f 100644 --- a/integrations/agentbricks/README.md +++ b/integrations/agentbricks/README.md @@ -105,18 +105,19 @@ eval "$(_AGENTBRICKS_COMPLETE=zsh_source agentbricks)" ## Authentication Agent Bricks CLI uses [Databricks authentication](https://docs.databricks.com/aws/en/dev-tools/cli/authentication). -Sign in to a profile, then scaffold a project that records it in its `.env`: +Create a project with the profile you want it to use (replace `my-profile` with your profile name): ```sh -agentbricks auth login -agentbricks init my-agent --profile +agentbricks init my-agent --profile my-profile ``` -`agentbricks auth login` validates the profile and, in an interactive terminal, runs -`databricks auth login --profile ` if it is not authenticated yet (this needs the Databricks -CLI; in CI or without a terminal, authenticate the profile beforehand). It saves nothing. To change a -project's profile later, run `agentbricks auth switch ` in the project directory, and -`agentbricks auth describe` shows which profile (and workspace) a directory would use. +`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` @@ -169,16 +170,19 @@ 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 auth login -agentbricks init my-agent --profile # LangGraph by default; 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.) diff --git a/integrations/agentbricks/cli.md b/integrations/agentbricks/cli.md index 8246de39e..711bad03a 100644 --- a/integrations/agentbricks/cli.md +++ b/integrations/agentbricks/cli.md @@ -27,12 +27,19 @@ 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). -Sign in with `agentbricks auth login ` (a wrapper over `databricks auth login`), then -select the profile with `--profile` / `-p` on any command or record it in the project's `.env` -(`agentbricks init --profile `, or `agentbricks auth switch ` later). +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. Nothing is saved outside the project: there is no -global default profile. The profile is resolved in this order, most to least specific: +command uses that project's `.env` profile. 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 @@ -192,7 +199,7 @@ With `-o json`: `{"directory", "env_profile", "changed", "workspace_host", "sign 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]` profile's host), and the project directory. +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] From c238f5e51e81a474133245bfcfde9d26837feaf4 Mon Sep 17 00:00:00 2001 From: Elaine Wang Date: Mon, 26 Oct 2026 18:58:58 +0000 Subject: [PATCH 12/12] [agentbricks] Find the project from subfolders, report auth timeouts as unreachable, and rewrite interpolated profiles on switch Co-authored-by: Isaac --- integrations/agentbricks/cli.md | 5 +- .../src/databricks_agentbricks/cli/app.py | 12 ++--- .../src/databricks_agentbricks/cli/auth.py | 20 ++++---- .../databricks_agentbricks/projects/config.py | 16 ++++++- .../projects/env_file.py | 9 ++-- .../services/workspace_auth.py | 20 ++++++-- .../agentbricks/tests/unit_tests/auth_test.py | 47 +++++++++++++++++++ .../tests/unit_tests/env_file_test.py | 10 ++++ .../tests/unit_tests/workspace_auth_test.py | 19 ++++++++ 9 files changed, 133 insertions(+), 25 deletions(-) diff --git a/integrations/agentbricks/cli.md b/integrations/agentbricks/cli.md index 711bad03a..9de6864bf 100644 --- a/integrations/agentbricks/cli.md +++ b/integrations/agentbricks/cli.md @@ -38,8 +38,9 @@ interactive terminal. Use `--profile` / `-p` to override the choice for one comm [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. Agent Bricks saves no separate default profile outside -the project. The profile is resolved in this order, most to least specific: +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 diff --git a/integrations/agentbricks/src/databricks_agentbricks/cli/app.py b/integrations/agentbricks/src/databricks_agentbricks/cli/app.py index e8c98a852..357d8c8c1 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/cli/app.py +++ b/integrations/agentbricks/src/databricks_agentbricks/cli/app.py @@ -30,7 +30,7 @@ announce_workspace, warn_legacy_default_profile, ) -from databricks_agentbricks.projects.config import is_agentbricks_project +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, @@ -173,11 +173,11 @@ 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, output=output) - cwd = pathlib.Path.cwd() - if is_agentbricks_project(cwd): - # Every command run inside a project uses its `.env` profile, so e.g. `endpoint invoke` - # targets the same workspace that `deploy` just used. - ctx.obj.cwd_project = cwd + 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 = { diff --git a/integrations/agentbricks/src/databricks_agentbricks/cli/auth.py b/integrations/agentbricks/src/databricks_agentbricks/cli/auth.py index 6954c7873..1a58a4b2d 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/cli/auth.py +++ b/integrations/agentbricks/src/databricks_agentbricks/cli/auth.py @@ -20,7 +20,7 @@ warn_sign_in_failed, warn_unknown_profile, ) -from databricks_agentbricks.projects.config import is_agentbricks_project +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, @@ -67,14 +67,16 @@ def sign_in_chosen_profile(obj, profile: str, source: str) -> Optional[str]: 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() - if not is_agentbricks_project(directory): + project = find_project_root(directory) + if project is None: raise AgentCliError( 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 directory + return project def _resolve_for(obj, directory: pathlib.Path, in_project: bool) -> ProfileInfo: @@ -135,10 +137,10 @@ def auth_describe(obj, source: str) -> None: authentication. Read-only; no workspace calls are made. """ directory = pathlib.Path(source).resolve() - in_project = is_agentbricks_project(directory) - info = _resolve_for(obj, directory, in_project) + 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(directory) if in_project else None + project = str(project_root) if project_root is not None else None show_profile_description(info.name, info.source, host, workspace, project, obj.output) @@ -164,10 +166,10 @@ def auth_login(obj, name: Optional[str], source: str) -> None: 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. """ - directory = pathlib.Path(source).resolve() + project_root = find_project_root(pathlib.Path(source).resolve()) env_profile = ( - read_env(directory / ".env").get("DATABRICKS_CONFIG_PROFILE") - if is_agentbricks_project(directory) + 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) diff --git a/integrations/agentbricks/src/databricks_agentbricks/projects/config.py b/integrations/agentbricks/src/databricks_agentbricks/projects/config.py index 1035f55a0..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 @@ -31,6 +31,20 @@ def is_agentbricks_project(path: pathlib.Path) -> bool: 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 index 70436f804..81a3b746f 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/projects/env_file.py +++ b/integrations/agentbricks/src/databricks_agentbricks/projects/env_file.py @@ -54,9 +54,9 @@ def read_env(path: pathlib.Path) -> dict[str, str]: return _parse_env(text) -def _parse_env(text: str) -> dict[str, str]: +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))) + 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} @@ -103,9 +103,10 @@ def update_env_profile(dest: pathlib.Path, profile: str) -> bool: if not env_path.exists(): return write_env(dest, profile) before = _read_text(env_path) - if _parse_env(before).get("DATABRICKS_CONFIG_PROFILE") == profile: + 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. + # 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) diff --git a/integrations/agentbricks/src/databricks_agentbricks/services/workspace_auth.py b/integrations/agentbricks/src/databricks_agentbricks/services/workspace_auth.py index 3c9656d3b..fc679c583 100644 --- a/integrations/agentbricks/src/databricks_agentbricks/services/workspace_auth.py +++ b/integrations/agentbricks/src/databricks_agentbricks/services/workspace_auth.py @@ -37,8 +37,16 @@ def _classify_validation_error(error: Exception) -> _ValidationFailure: ) 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) + error, + ( + RequestsConnectionError, + InternalError, + TemporarilyUnavailable, + DeadlineExceeded, + TimeoutError, + ), ): return _ValidationFailure.UNREACHABLE # Client construction wraps every failure (e.g. an unknown profile) in AgentCliError. @@ -76,6 +84,12 @@ def _not_signed_in_error(profile: Optional[str], source: str) -> AgentCliError: 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.", @@ -102,8 +116,8 @@ def _validate_bounded( ) -> tuple[Optional[tuple[_AgentBricksApiClient, str]], Optional[Exception]]: """`_validate_profile` on a daemon worker thread, giving up after `_PREFLIGHT_TIMEOUT_S`. - A browser-OAuth flow can block forever, so the call is bounded. Daemon because executor - threads are joined at interpreter exit and a hung flow must not block the CLI's exit too. + 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]]] = [] diff --git a/integrations/agentbricks/tests/unit_tests/auth_test.py b/integrations/agentbricks/tests/unit_tests/auth_test.py index 0b7bb85ce..6937389c0 100644 --- a/integrations/agentbricks/tests/unit_tests/auth_test.py +++ b/integrations/agentbricks/tests/unit_tests/auth_test.py @@ -19,6 +19,7 @@ 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 @@ -241,6 +242,52 @@ def test_get_reports_profile_source_host_and_project_without_workspace_calls( 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, + ) + + +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") diff --git a/integrations/agentbricks/tests/unit_tests/env_file_test.py b/integrations/agentbricks/tests/unit_tests/env_file_test.py index bd4a2b53d..340334541 100644 --- a/integrations/agentbricks/tests/unit_tests/env_file_test.py +++ b/integrations/agentbricks/tests/unit_tests/env_file_test.py @@ -57,6 +57,16 @@ def test_update_env_profile_adds_missing_line(tmp_path: pathlib.Path): 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) diff --git a/integrations/agentbricks/tests/unit_tests/workspace_auth_test.py b/integrations/agentbricks/tests/unit_tests/workspace_auth_test.py index 76e503b51..e755ec26b 100644 --- a/integrations/agentbricks/tests/unit_tests/workspace_auth_test.py +++ b/integrations/agentbricks/tests/unit_tests/workspace_auth_test.py @@ -138,6 +138,16 @@ def test_preflight_failure_hints_at_login_and_retry(fail_validation): ) +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) @@ -213,6 +223,15 @@ def test_auth_service_logs_in_then_revalidates_through_supplied_collaborators(): 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