feat(settings): add native Databricks Unity AI Gateway provider to settings TUI#740
Draft
prasadkona wants to merge 41 commits into
Draft
feat(settings): add native Databricks Unity AI Gateway provider to settings TUI#740prasadkona wants to merge 41 commits into
prasadkona wants to merge 41 commits into
Conversation
4 tasks
enyst
requested changes
May 18, 2026
enyst
left a comment
enyst
left a comment
Member
There was a problem hiding this comment.
Thank you for the contribution @prasadkona , I’m not sure this is the best way though, so you know: I suggest we discuss your proposal in the SDK repo PR, first.
I mean, up to you, I’m just concerned that we may or may not be able to accept your Databricks proposal, and if yes, how; and the best place for that discussion is in the SDK where we started talking.
Wires the DatabricksLLM provider into the CLI settings screen: Settings UI - Databricks auth section: PAT / M2M / CLI Profile / U2M (browser SSO) - Auth-method-aware field visibility: API Key field hidden for non-PAT methods; Profile/M2M credential fields shown only when relevant - Live auth-method hints with step-by-step instructions for U2M and CLI Profile (includes install + login commands) - Model dropdown: two-tier picker (curated + live-discovered) auto- refreshes when auth method or workspace host changes, using the credentials typed in the form so U2M/profile show the full list Agent store - `LLMEnvOverrides`: DATABRICKS_HOST + DATABRICKS_TOKEN env overrides - `apply_llm_overrides`: rebuilds via create_llm for Databricks models - Headless agent creation uses create_llm factory Discovery (choices.py) - `_resolve_credentials_for_host`: builds credentials from form state (host + auth method) for targeted per-workspace discovery - `_get_databricks_model_options`: accepts pre-built credentials so auth-method changes immediately surface the correct model list run_local.sh - Installs local SDK in compat editable mode and ensures databricks-sdk is present (needed for U2M / profile auth) before launching the CLI
pytest namespace-package collection fails without __init__.py when sibling test directories already have one. All other test dirs under tests/tui/ have __init__.py; add it to tests/tui/modals/settings/ so Databricks and existing settings tests can be collected correctly.
When a Databricks discovered-only model (not in the curated list) was previously saved and the Settings screen mounts without credentials, the model dropdown was populated with curated-only options, causing Textual to raise InvalidSelectValueError when trying to restore the saved selection. Fix: wrap the value assignment in a try/except — on failure, inject the saved model as the sole option with a "(saved — re-enter credentials to refresh)" label so the user sees their current selection without crashing. _refresh_databricks_models() repopulates the full discovered list once host and auth fields are filled in.
The U2M auth hint now uses whatever the user has typed into the Workspace Host field to build the exact databricks auth login command, so they can copy-paste it directly from the TUI. Before: databricks auth login --host <workspace_host> After: databricks auth login --host https://e2-demo-field-eng.cloud.databricks.com The hint updates live as the user types the host — no need to leave the settings screen to figure out the right command.
This reverts commit b3f85ee.
Two blockers for external contributors: 1. The previous "chore: update uv.lock" commit rewrote all 401 package sources from pypi.org to pypi-proxy.dev.databricks.com (an internal Databricks mirror). This has been reverted so external contributors can run `uv sync` without VPN access. 2. pyproject.toml pinned openhands-sdk==1.21.0 (published), but the Databricks native provider ships in the still-unmerged SDK PR #3286. Enable [tool.uv.sources] with a git-source pin to the PR branch HEAD so the CLI is installable without the run_local.sh editable override. Also relax the constraint from ==1.21.0 to >=1.21.0 to allow the git-sourced version to satisfy the requirement. This section should be removed and the version constraint restored once SDK PR #3286 is merged and published to PyPI.
The DATABRICKS_AI_GATEWAY_HOST override is supported in the backend and configurable via environment variable, but surfacing it in the TUI adds unnecessary complexity for most users who use a standard workspace. Remove the input widget and its help text from settings_tab.py. The field remains functional via env var for advanced/split-hostname deployments. Re-enabling the UI field is a one-file change when needed.
…erences The UI widget was removed from settings_tab.py but settings_screen.py still declared the query_one getter and read/wrote its value in three places (_clear_databricks_fields, _load_current_settings, _get_form_data). Textual's query_one raises NoMatches at runtime, crashing the TUI on Save. Remove the getter declaration and all three call sites. The env-var-only DATABRICKS_AI_GATEWAY_HOST path is preserved in the backend; the TUI simply passes db_ai_gateway_host=None so the backend falls through to the env var.
…_llm_overrides When env-var overrides force a Databricks model swap, apply_llm_overrides() calls create_llm() to build a fresh DatabricksLLM. The previous instance was discarded without closing its httpx client, leaking an open connection pool. Explicitly call llm.close() on the existing DatabricksLLM instance before creating the replacement to release the underlying httpx.Client.
Add OAuth App Client ID and Client Secret inputs to the Databricks settings tab, defaulting auth mode to U2M (Browser SSO). Help text explains how to register an OAuth app in the Databricks account console. Changes: - settings_tab.py: remove AI Gateway Host input (handled internally); reorder Select options so U2M is default; add Container with databricks_u2m_client_id and databricks_u2m_client_secret inputs. - settings_screen.py: declare getters for new inputs; update _clear_databricks_fields, _update_databricks_visibility, _build_u2m_hint, _load_current_settings, and _get_form_data. - utils.py: add databricks_u2m_client_id + databricks_u2m_client_secret to SettingsFormData and kwargs_from_settings.
Replace the static 'run databricks auth login' hint with a real inline browser-based PKCE flow triggered automatically after Save when U2M auth mode is selected and an OAuth App Client ID is provided. New file: openhands_cli/auth/databricks_pkce.py - Pure-asyncio PKCE helpers (generate_pkce, build_authorize_url, exchange_code_for_tokens) mirroring the web app oauth module. - run_browser_pkce_flow(): starts a local HTTP server on localhost:8080/callback, opens the system browser to the Databricks OIDC authorize URL, waits up to 120 s for the callback, exchanges the code for tokens, and returns them. - Renders success/error HTML to the browser tab after redirect. settings_screen.py changes: - Fix NameError in _build_u2m_hint() (login_cmd was undefined). - Add explicit databricks_host_input.disabled = False in _update_databricks_visibility() so the field is always editable. - _run_u2m_pkce_flow(): new @work(thread=False) async worker that runs the PKCE flow, stores the obtained access token as api_key on the saved agent, shows status messages in the TUI, and dismisses after success. - _save_settings(): after a successful save with U2M + client_id, redirect into _run_u2m_pkce_flow instead of immediately dismissing. User flow: 1. Settings -> Databricks -> Browser OAuth (U2M, default) 2. Enter workspace host + OAuth App Client ID (+ optional secret) 3. Click Save -> browser opens to Databricks sign-in 4. Authenticate -> browser shows 'Signed in successfully' 5. TUI shows success, closes settings, conversation starts.
Users can now specify a custom redirect URI in the U2M settings group, defaulting to http://localhost:8080/callback. The field is optional — leaving it blank uses the default. The hint text updates live to show the currently configured redirect URI so it's easy to copy into the Databricks App connections form. Changes: - settings_tab.py: add 'OAuth Redirect URI (optional)' Input with placeholder and help text in databricks_u2m_group. - settings_screen.py: add getter, clear/load/save handling for the new field; update _build_u2m_hint() to show the configured redirect URI; refresh hint on redirect URI input changes; pass redirect_uri to _run_u2m_pkce_flow(); derive callback_port from redirect_uri port. - utils.py: add databricks_u2m_redirect_uri to SettingsFormData, strip_strings validator, and _build_databricks_settings namespace. - databricks_pkce.py: run_browser_pkce_flow() now accepts optional redirect_uri kwarg; derives callback_port from URI port when supplied.
Two bugs fixed and model discovery redesigned: 1. Crash fix: nested @work decorator on _discover() was being called as _discover(self) causing 'takes 0 positional arguments but 1 was given'. Replaced with self.run_worker(_discover, thread=True) API. 2. TUI freeze fix: _refresh_databricks_models() was calling cfg.authenticate() (network I/O) synchronously on the main event loop whenever the host field changed or the auth method was switched. Removed those auto-trigger calls entirely. New model discovery design: - choices.py: DATABRICKS_STATIC_MODELS — 11 curated models (Claude, GPT, Gemini, Llama, DBRX) shown immediately with zero network I/O. get_model_options('databricks') returns this list when credentials=None. - _load_current_settings: guarded by _loading_settings flag so programmatic field population never triggers on_input_changed side-effects. - _refresh_databricks_models_with_token(host, token): new method that runs discovery in a background thread using self.run_worker() and posts the results back via call_from_thread(). Called only after authentication is complete. - _run_u2m_pkce_flow: calls _refresh_databricks_models_with_token() after tokens are obtained so the model dropdown updates to the full live workspace list right after sign-in.
Two fixes: 1. Settings screen now closes after browser OAuth completes. Root cause: self.run_worker(..., exclusive=True) inside _run_u2m_pkce_flow cancelled the still-running PKCE async worker so dismiss() was never reached. Fix: move dismiss() before the model-refresh call; remove exclusive=True from the refresh worker. 2. On settings re-open, live model list loads automatically. After the first successful sign-in the access token is persisted in the agent config. When settings is next opened, _load_current_settings_inner detects the saved token and calls _refresh_databricks_models_with_token() in a background thread so the full workspace model list populates while the user browses settings — no blocking, falls back to static list if it fails.
…ellm fallback model_copy() skips Pydantic model validators, so the private _db_credentials and _db_client attributes remained stale (built before the access token existed). When a conversation started, DatabricksLLM._transport_call used the old credential object which had no token, causing the base LLM class to fall through to litellm's Databricks provider — resulting in litellm.BadRequestError: [400] Unknown error. Fix: after the PKCE flow completes, call create_llm(model, databricks_host, api_key=access_token) so that _init_databricks() runs fresh with api_key set, triggering the PAT auth path and initialising a new httpx client with valid credentials — matching the same pattern used in the web app's _configure_llm.
…read _refresh_databricks_models_with_token was scheduling _update_model_options (which calls get_picker_entries — a live HTTP call) on the main Textual thread via call_from_thread. This blocked the TUI and any exception was swallowed by _update_model_options's except clause, so the dropdown silently kept the static seed list. Fix: perform the full get_model_options/get_picker_entries network call inside the background worker thread. Introduce _apply_model_options (main-thread only) that receives the finished list and calls set_options — keeping UI work on the main thread and heavy I/O off it.
…er PKCE Two bugs fixed: 1. Settings cleared on re-open: after PKCE we were calling create_llm with only model/databricks_host/api_key, dropping databricks_u2m_client_id and redirect URI. On re-open, resolve_credentials saw api_key → auth_method="pat", so the UI showed a blank PAT form instead of the U2M form the user configured. Fix: rebuild with stored_u2m_tokens (not api_key) and pass databricks_u2m_client_id / databricks_u2m_redirect_uri (now proper DatabricksLLM fields). resolve_credentials picks the U2M path → auth_method="u2m" → settings re-opens with the correct form. 2. Model refresh not triggering on re-open: the token check only looked at api_key. After the fix above the token lives in stored_u2m_tokens.access_token, so the refresh check now also reads that path.
…sage
StoredU2MTokens requires access_token, refresh_token, expires_at, client_id,
and host — we were only supplying 2, causing 3 Pydantic validation errors.
The tokens dict from run_browser_pkce_flow already contains all 5 values.
Also fix _show_message to escape Rich/Textual markup characters before calling
Static.update(), so Pydantic error strings containing '{...}' dict reprs don't
crash the TUI with a MarkupError.
Per product decision, the Databricks connector supports exactly two authentication methods: 1. Browser OAuth (U2M) — recommended, uses custom OAuth App + PKCE 2. Service Principal (M2M) — for headless/automated scenarios Removed: - "Personal Access Token (PAT)" dropdown option and API key field - "CLI Profile (~/.databrickscfg)" dropdown option and profile name field - All PAT/profile branches in _update_databricks_visibility, _load_current_settings_inner, _get_form_data, _databricks_auth_needs_api_key, and _build_databricks_settings - databricks_profile_name from SettingsFormData and DatabricksAuthMethod Literal DatabricksAuthMethod is now Literal["m2m", "u2m"]; default is "u2m".
…dback The background model refresh was silent — no visual indication it was running or whether it succeeded. Users couldn't tell if the live workspace list had loaded. Changes: - Add '↻ Refresh model list' button below the model dropdown (enabled when Databricks is selected and credentials are saved) - Show '⟳ Loading models from workspace…' while the background discovery runs - Show '✓ N models loaded from workspace.' on success, or an error message - Add _trigger_manual_model_refresh() so users can explicitly re-fetch at any time (useful when the auto-refresh after PKCE ran on the dismissed screen) - _set_refresh_status() re-enables the button and updates the status label
call_from_thread is on App, not Screen, in this Textual version.
…ta probe - save_settings/_build_databricks_settings now forwards stored_u2m_tokens from the existing agent when auth_method=="u2m", so switching models no longer drops the browser-OAuth session and forces re-authentication. - Enable databricks_metadata_probe=True by default for Databricks LLMs built through save_settings; the serving-endpoint metadata API returns authoritative api_types (confirmed by smoke test), giving more reliable family detection than name-pattern guessing. - Mask databricks_client_id (M2M) and databricks_u2m_client_id (U2M) in the settings form; both are now password=True to match client_secret.
- _load_current_settings_inner now reads databricks_u2m_client_secret from the saved LLM and populates the form field (password-masked) so the user doesn't have to re-enter it every session. - resolve_data_fields carries the secret forward from existing_agent when the form field is empty (same pattern as M2M client_secret). Field is optional so no error is raised if absent (public OAuth apps have none).
…KCE flow - settings_screen.py: PKCE flow (_run_u2m_pkce_flow) now forwards databricks_u2m_client_secret and rebuilds the condenser LLM with fresh tokens after successful browser sign-in. Auth method detection checks stored_u2m_tokens first to avoid stale M2M fields overriding an active U2M session. U2M client secret is no longer written to the input field (uses placeholder instead) to prevent plaintext exposure in the TUI; the actual secret is recovered from current_agent for the PKCE call. - utils.py: resolve_data_fields also carries forward databricks_u2m_client_id and databricks_u2m_redirect_uri from existing_agent so re-auth works without the user re-entering all three U2M credential fields.
choices.py — DATABRICKS_STATIC_MODELS: - Remove stale/broken entries: gpt-4o, gpt-4o-mini, gemini-2-0-flash, gemini-1-5-pro, claude-3-7-sonnet, claude-3-5-sonnet, dbrx-instruct (gpt-4o was the direct cause of "endpoint does not exist" errors) - Add current Claude 4 series (sonnet-4-6, haiku-4-5, opus-4-7/4-5/4-1) - Add current GPT-5 series (gpt-5-mini, 5-5-pro, 5-5, 5-4, 5-4-mini, gpt-5) - Add current Gemini 3 series (3-5-flash, 3-flash, 3-pro) + keep 2.5 entries - Add llama-4-maverick and qwen35-122b-a10b model_recommendations.py: - Update recommended Claude to sonnet-4-6 (was 4-5) - Update recommended Gemini to gemini-3-5-flash (was 2-5-flash) - Add gpt-5-5-pro, gpt-5-4, claude-opus-4-7, gemini-3-flash
settings_screen.py: - Sync in-memory agent after save so subsequent U2M PKCE flow uses the freshly saved credentials rather than stale state - Surface save warnings (e.g. selected model not in discovered list) to the user via a modal - Warn when the selected model is unavailable in the discovered endpoint cache rather than silently accepting it utils.py: - SettingsSaveResult now carries agent and warning fields - Pre-save check (_check_databricks_endpoint_in_cache) warns when the chosen model is not found in the live endpoint cache - Carry over databricks_u2m_client_id and databricks_u2m_redirect_uri when resolving data fields so they survive a settings round-trip choices.py / model_recommendations.py: - Update DATABRICKS_STATIC_MODELS to reflect May 2026 live-tested FMAPI endpoints; remove endpoints that no longer exist run_web_local.sh (new): - Developer helper to launch the OpenHands web UI + backend locally with the local software-agent-sdk injected; includes port 8080 shim for Databricks OAuth redirect URI registration convenience
The databricks_u2m_client_id_input had password=True which masked the client ID visually. Client IDs are not secrets (they are public OAuth app identifiers), so masking them prevents users from seeing and correcting typos. This caused a hard-to-debug issue where the literal label text 'OAuth App Client ID:' had been saved as the client_id value without the user noticing. Only client secrets and API keys should have password=True.
Adds a UUID format check before opening the browser OAuth flow. Databricks OAuth App client IDs are always UUIDs; values containing colons, spaces, or other non-UUID characters (e.g. accidentally typed label text like 'OAuth App Client ID:') are rejected with a clear error message before the browser is opened, rather than discovering the mistake only after Databricks rejects the token request.
- Improve OPENHANDS_WORK_DIR override support and startup reliability
prasadkona
force-pushed
the
feat/databricks-native-provider
branch
from
3f869c1 to
777e3a8
Compare
3 tasks
prasadkona
changed the title
prasadkona
force-pushed
the
feat/databricks-native-provider
branch
from
af8d09e to
3fcd8f3
Compare
Remove the CLI's duplicate Authorization Code + PKCE primitives (generate, authorize-URL, code-for-token exchange) and import them from the SDK's shared openhands.sdk.llm.providers.databricks.pkce module. The CLI keeps only its local browser/callback-server orchestration, eliminating three-way drift between the SDK, web app, and CLI.
Point the temporary git-source pin at the SDK commit that introduces the shared databricks PKCE helpers, so the CLI's import of openhands.sdk.llm.providers.databricks.pkce resolves on a fresh install.
Document the CLI Databricks auth options (pat / m2m / u2m / profile), the inline U2M browser PKCE flow (localhost callback, default port 8080) and its redirect URI requirement, and note that PKCE primitives and the settings->create_llm bridge are shared with the SDK and web backend.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds first-class support for the native Databricks Unity AI Gateway provider
(Databricks PWAF compliant) to the settings TUI: workspace host, auth method
selector (M2M service principal / Browser SSO U2M), live model discovery from the
workspace, and Databricks-aware env-var overrides. This gives CLI users governed
access to the foundation models served through their Databricks workspace.
Depends on the companion SDK PR (OpenHands/software-agent-sdk#3286).
What's new
Settings screen additions
databricks:databricks auth logincommand for Browser SSO, using the hostname entered on the same screen
Model discovery
when host + credentials are available — same two-tier list as the web UI
Env-var overrides
DATABRICKS_HOSTandDATABRICKS_TOKENenv vars supported in headless mode(alongside existing
LLM_API_KEY/LLM_MODEL/LLM_BASE_URL)Changes to existing files
agent_store.pyLLMEnvOverridesforDATABRICKS_HOST/DATABRICKS_TOKEN; rebuilds Databricks instances so client state is freshsettings_screen.pysettings_tab.pyutils.pySettingsFormDatavalidation for Databricks auth pathschoices.pyget_databricks_model_options()with TTL-cached live discoverymodel_recommendations.pydatabricksto the provider enumTests
test_choices_databricks.py,test_databricks_auth_method.py,and additions to
test_settings_utils.py/test_settings_tab.pyM2M service principal ✅, Browser SSO (U2M) ✅
Test plan
uv run pytest tests/tui/modals/settings/ -q— all settings tests passDATABRICKS_HOST=... DATABRICKS_TOKEN=... LLM_MODEL=databricks/...— headless agent starts without prompting for API keyAlignment with Databricks
ucodeThis integration follows the same credential model as
Databricks
ucode— the Unity AI GatewayCoding CLI that launches coding agents through the Databricks AI Gateway using
workspace credentials, no API keys required. The connector's
PROFILEandUNIFIEDstrategies read the workspace login a developer has already established(
databricks auth login/~/.databrickscfg), andU2Mprovides interactivebrowser OAuth — so an OpenHands agent can reach AI Gateway the same key-free,
governed way
ucodedoes, reusing the existing workspace session rather than aseparate token. The result is one consistent, governed path to AI Gateway (and
the Unity Catalog–governed resources behind it) across
ucodeand OpenHands.