feat(sprout-agent): Databricks provider with OAuth 2.0 PKCE auth#698
Merged
Conversation
Adds `Provider::Databricks` to sprout-agent, alongside the existing Anthropic and OpenAI-compat providers. Databricks model-serving is OpenAI-chat-compatible on the wire, so the existing body builder and response parser are reused untouched; the new code is the auth and a small URL/body adapter. Configuration: SPROUT_AGENT_PROVIDER=databricks DATABRICKS_HOST=https://your-workspace.cloud.databricks.com DATABRICKS_MODEL=goose-claude-4-6-sonnet # or any serving endpoint DATABRICKS_TOKEN=<optional> # static-key escape hatch When `DATABRICKS_TOKEN` is unset, sprout-agent runs the OAuth 2.0 Authorization Code + PKCE flow (RFC 6749 + RFC 7636) using Databricks' public CLI client (`client_id=databricks-cli`, scopes `all-apis offline_access`). The same identifier goose uses, so an existing browser consent on a Block laptop carries over. Tokens are cached under `~/.config/sprout-agent/oauth/databricks/<sha256>.json` and refreshed transparently on expiry; long-running agents survive across token rotations. Implementation: - `crates/sprout-agent/src/auth.rs` (new): `TokenSource` trait, `StaticTokenSource`, `PkceOAuthTokenSource` with OIDC discovery, one-shot localhost callback server (axum), PKCE pair generation, refresh-grant POST, and atomic on-disk cache. - `llm.rs`: `Llm` now holds `Arc<dyn TokenSource>` (static for Anthropic/OpenAI-compat, refreshable for Databricks). `post_openai` fetches the bearer per request and, when provider=Databricks, rewrites the URL to `{base}/serving-endpoints/{model}/invocations` and strips the `model` field from the body. The `complete` and `summarize` dispatches use a single `OpenAi | Databricks` match arm over the shared body+parser pair, so adding the new provider didn't fork the OpenAI path. Tests: - 5 unit tests in `auth.rs` cover PKCE pair generation, cache-path construction, token-expiry math (with leeway), and response parsing with refresh-token fallback. - 3 integration tests in `tests/databricks_oauth.rs` use a stub `axum` OIDC server to exercise cache-hit (no network), expired-cache silent refresh, and on-disk persistence of refreshed tokens. No browser flow in CI; the interactive path is exercised manually. - Existing 17 regression / golden / openai-auto-upgrade tests continue to pass — Anthropic and OpenAI behavior is unchanged. Verified end-to-end against `block-lakehouse-production` with `goose-claude-4-6-sonnet`: cold-cache (silent refresh) call completes in ~6s; warm-cache call in ~1.8s. Signed-off-by: tlongwell-block <109685178+tlongwell-block@users.noreply.github.com> Co-authored-by: Dawn (sprout agent) <c6237ef84fa537c78dcee78efd2d4e59f728859c7f194da42ac51ededfa0be05@sprout-oss.stage.blox.sqprod.co>
…t guard Addresses three review items from Max on PR #698: 1. **Fail loudly on malformed token responses.** `token_from_response` is now fallible and rejects missing or empty `access_token`. Without this, a malformed token-endpoint response would have been cached and `bearer()` would have silently returned `""` until the entry expired or was removed by hand. Adds two unit tests covering the rejection paths. 2. **Abort the callback server on every exit.** `browser_pkce_flow` previously only called `server.abort()` on the happy path; on timeout or token-exchange failure the axum task could outlive the function and keep the loopback listener open. A new `AbortOnDrop` guard wraps the join handle so the server task is cancelled deterministically on every return path. 3. **Wire `sprout-agent auth databricks` subcommand.** Resolves the docs/code mismatch — the test comment referenced a command that didn't exist. The subcommand reads `DATABRICKS_HOST`, runs the interactive PKCE flow via the already-public `interactive_login()`, and persists the token to the standard cache. ~30 LOC in `lib.rs`. Manual error-path coverage: $ sprout-agent auth Error: auth: provider required (try: sprout-agent auth databricks) $ sprout-agent auth bogus Error: auth: unknown provider "bogus" $ DATABRICKS_HOST= sprout-agent auth databricks Error: auth databricks: DATABRICKS_HOST required Additionally adds an ACP-level envelope regression test (`databricks_envelope_routes_through_serving_endpoints_and_strips_model`) requested in review. Spawns the real `sprout-agent` binary with `DATABRICKS_TOKEN` set against a request-capturing stub HTTP server and asserts: - request path is `/serving-endpoints/<model>/invocations` - `Authorization` header is `Bearer <token>` - request body has no top-level `model` field - request body still carries the chat `messages` array This locks in the DRY envelope so a refactor of the OpenAI-family path can't silently break Databricks. All 71 sprout-agent tests pass; clippy clean. End-to-end smoke against `block-lakehouse-production` with `goose-claude-4-6-sonnet` still returns `PONG` in ~1.8s on the warm-cache path. Signed-off-by: tlongwell-block <109685178+tlongwell-block@users.noreply.github.com> Co-authored-by: Dawn (sprout agent) <c6237ef84fa537c78dcee78efd2d4e59f728859c7f194da42ac51ededfa0be05@sprout-oss.stage.blox.sqprod.co>
CI's `fmt-check` flagged formatting drift introduced by the previous two commits. No behavior change — `cargo fmt -p sprout-agent` over auth.rs, lib.rs, llm.rs, and tests/databricks_oauth.rs. Tests + clippy still green locally. Signed-off-by: tlongwell-block <109685178+tlongwell-block@users.noreply.github.com> Co-authored-by: Dawn (sprout agent) <c6237ef84fa537c78dcee78efd2d4e59f728859c7f194da42ac51ededfa0be05@sprout-oss.stage.blox.sqprod.co>
tlongwell-block
added a commit
that referenced
this pull request
May 21, 2026
Bring pulse-front-back up to date with main prior to opening a PR. Signed-off-by: tlongwell-block <109685178+tlongwell-block@users.noreply.github.com> Co-authored-by: Dawn (sprout agent) <c6237ef84fa537c78dcee78efd2d4e59f728859c7f194da42ac51ededfa0be05@sprout-oss.stage.blox.sqprod.co> * origin/main: (35 commits) feat(sprout-agent): auto-fallback to Databricks OAuth (#699) fix(relay): avoid websocket writes in stall watchdog (#697) feat(sprout-agent): Databricks provider with OAuth 2.0 PKCE auth (#698) Add Ubuntu desktop release artifacts (#693) chore(deps): update rust crate tokio to v1.52.3 (#658) chore(deps): update all non-major dependencies (#650) chore(deps): update rust crate sherpa-onnx to v1.13.2 (#657) chore(deps): update dependency nostr-tools to v2.23.5 (#681) chore(deps): update tanstack-router monorepo (#659) chore(deps): update rust crate dashmap to v6.2.1 (#652) chore(deps): update rust crate tower-http to v0.6.11 (#647) chore(deps): update rust crate reqwest to v0.13.3 (#639) chore(deps): update rust crate sherpa-onnx to v1.12.40 (#640) chore(deps): update dependency @tanstack/react-query to v5.100.11 (#635) fix(deps): update rust crate sha2 to 0.11 (#665) fix(deps): update rust crate bzip2 to 0.6 (#661) chore(deps): update rust crate uuid to v1.23.1 (#648) chore(deps): update rust crate tauri-plugin-dialog to v2.7.1 (#644) chore(deps): update tanstack-router monorepo (#649) chore(deps): update rust crate tokio to v1.51.3 (#646) ...
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
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.
Adds
Provider::Databricksto sprout-agent, alongside the existing Anthropic and OpenAI-compat providers. Databricks model-serving is OpenAI-chat-compatible on the wire, so the existing body builder and response parser are reused untouched; the new code is the auth machinery and a small URL/body adapter.Configuration
When
DATABRICKS_TOKENis unset, sprout-agent runs the OAuth 2.0 Authorization Code + PKCE flow (RFC 6749 + RFC 7636) using Databricks' public CLI client (client_id=databricks-cli, scopesall-apis offline_access). The same identifier goose uses, so an existing browser consent on a Block laptop carries over.Tokens are cached under
~/.config/sprout-agent/oauth/databricks/<sha256>.jsonand refreshed transparently on expiry; long-running agents survive across token rotations.Why this shape
This was discussed in https://sprout-oss.stage.blox.sqprod.co channel
sprout-agent-auto-databricks-for-blockwith @max-gpt55. The two key calls:client_idand scopes as goose so the same browser consent works, but we own our cache directory and refresh logic.completeandsummarizedispatches collapse to a singleOpenAi | Databricks => …match arm. Databricks-specific behavior (URL templating, body model-stripping) lives behind one branch inpost_openai.Implementation
crates/sprout-agent/src/auth.rs(new, ~500 LOC incl. tests + docs):TokenSourcetrait,StaticTokenSource,PkceOAuthTokenSourcewith OIDC discovery, one-shot localhost callback server (axum), PKCE pair generation, refresh-grant POST, atomic on-disk cache write. The PKCE engine is provider-generic —PkceOAuthConfig { discovery_url, client_id, scopes, cache_namespace }— so a future provider that's also browser-PKCE-compatible reuses the engine; Databricks is just the first descriptor.llm.rs:Llmnow holdsArc<dyn TokenSource>(static for Anthropic/OpenAI-compat, refreshable for Databricks).post_openaifetches the bearer per request and, when provider=Databricks, rewrites the URL to{base}/serving-endpoints/{model}/invocationsand strips themodelfield from the body.config.rs:Provider::Databricksvariant;DATABRICKS_HOST/DATABRICKS_MODELrequired,DATABRICKS_TOKENoptional.What's intentionally not here
/serving-endpoints/responsesfor gpt-5.x). Additive — adds a third endpoint-strategy variant when needed. The chat/invocations path covers Sonnet/Haiku, which is all we need today.sprout-agent auth databricksCLI subcommand for interactive first-run. The PKCE engine supports it viainteractive_login()(already implemented and pub), but the CLI wiring is deferred to a follow-up to keep this PR scoped. First-run today either (a) hits the browser dance lazily on the first inference call, or (b) requires bootstrapping the cache (e.g. seeding from goose's refresh token, as we did for testing).TokenSourceimpls when needed.Tests
auth.rscover PKCE pair generation, cache-path construction, token-expiry math (with leeway), and response parsing with refresh-token fallback.tests/databricks_oauth.rsuse a stubaxumOIDC server to exercise:Clippy: clean (
cargo clippy -p sprout-agent --all-targets -- -D warnings).End-to-end verification
Tested against
block-lakehouse-productionwithgoose-claude-4-6-sonnet:Both prompts returned the expected response and
stopReason: end_turn. The on-disk cache was rewritten with a fresh access token and rotated refresh token after the cold-cache run, exactly as designed.cc @max-gpt55 for review.