Skip to content

feat(sprout-agent): Databricks provider with OAuth 2.0 PKCE auth#698

Merged
tlongwell-block merged 3 commits into
mainfrom
dawn/databricks-oauth-provider
May 21, 2026
Merged

feat(sprout-agent): Databricks provider with OAuth 2.0 PKCE auth#698
tlongwell-block merged 3 commits into
mainfrom
dawn/databricks-oauth-provider

Conversation

@tlongwell-block
Copy link
Copy Markdown
Collaborator

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 machinery 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.

Why this shape

This was discussed in https://sprout-oss.stage.blox.sqprod.co channel sprout-agent-auto-databricks-for-block with @max-gpt55. The two key calls:

  1. Don't piggyback on goose's cache. It's private state; coupling to it is a hidden break-on-update. We use the same client_id and scopes as goose so the same browser consent works, but we own our cache directory and refresh logic.
  2. Don't fork the OpenAI request path. The body shape, retries, auto-upgrade, and parser are reused unchanged. The complete and summarize dispatches collapse to a single OpenAi | Databricks => … match arm. Databricks-specific behavior (URL templating, body model-stripping) lives behind one branch in post_openai.

Implementation

  • crates/sprout-agent/src/auth.rs (new, ~500 LOC incl. tests + docs): TokenSource trait, StaticTokenSource, PkceOAuthTokenSource with 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: 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.
  • config.rs: Provider::Databricks variant; DATABRICKS_HOST/DATABRICKS_MODEL required, DATABRICKS_TOKEN optional.

What's intentionally not here

  • Responses API path for Databricks (/serving-endpoints/responses for 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 databricks CLI subcommand for interactive first-run. The PKCE engine supports it via interactive_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).
  • Device-code, client-secret, or service-account flows. Different enough that forcing them through the PKCE trait now would warp it. Add as additional TokenSource impls when needed.

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 short-circuits the network (zero refresh calls)
    • expired cache silently refreshes via the refresh token
    • refreshed token is persisted to disk
  • Existing 17 regression / golden / openai-auto-upgrade tests pass unchanged — Anthropic and OpenAI behavior is untouched.
test result: ok. 33 passed (lib)
test result: ok.  3 passed (databricks_oauth)
test result: ok.  4 passed (golden_transcripts)
test result: ok.  6 passed (openai_auto_upgrade)
test result: ok.  4 passed (fake_llm)
test result: ok. 17 passed (regressions)

Clippy: clean (cargo clippy -p sprout-agent --all-targets -- -D warnings).

End-to-end verification

Tested against block-lakehouse-production with goose-claude-4-6-sonnet:

Path Time
Cold cache → silent OAuth refresh → inference ~6.0s
Warm cache → in-memory token → inference ~1.8s

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.

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>
@tlongwell-block tlongwell-block requested a review from a team as a code owner May 21, 2026 15:14
tlongwell-block and others added 2 commits May 21, 2026 11:26
…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 tlongwell-block merged commit fa9e26f into main May 21, 2026
15 checks passed
@tlongwell-block tlongwell-block deleted the dawn/databricks-oauth-provider branch May 21, 2026 16:08
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)
  ...
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant