Repository navigation
copilot: Route ug copilot through a refreshing loopback auth proxy - #1217
Merged
Merged
Conversation
Copilot CLI holds its provider credential for the whole session, so the
Databricks token ug passed at launch expired mid-session (the 30-minute
refresher only rewrote ~/.copilot/ucode.env, which Copilot never reads).
The token was also persisted in that file and its backup, and `ug revert`
restored it.
`ug copilot` now starts a per-session proxy on 127.0.0.1. Copilot gets the
proxy URL and a random session secret; the proxy accepts the secret only on
gateway model routes (/ai-gateway/{anthropic,mlflow/v1,openai/v1}/), swaps
it for a fresh Databricks token (CLI profile or custom OAuth), and adds the
ucode/copilot User-Agent, request tags and managed http_headers. A rejected
token is refreshed and the request retried; when refresh itself fails the
proxy answers with a "Databricks sign-in expired" error in the request's
dialect and logs to ~/.ucode/logs/copilot-proxy.log instead of the TUI.
- Configure and launch write no credential; earlier tokens are scrubbed
from ucode.env, a token-bearing backup is dropped (also before revert),
and ug no longer backs up its own file on reconfigure.
- The refresher thread is gone, so a running session no longer resets
current_workspace or rewrites state.
- Ctrl-C no longer re-signals Copilot (send_signal(SIGINT) raised on
Windows); ug waits for Copilot to exit and then stops the proxy.
- Inherited COPILOT_PROVIDER_API_KEY(_COMMAND) / _TRANSPORT are dropped so
they can't bypass the proxy.
Tests: hermetic socket tests for the proxy (secret, path allowlist incl.
encoded traversal, header swap, refresh-and-retry, actionable 401), Copilot
unit tests, the Copilot e2e now runs through the proxy, and a new
non-blocking Copilot integration lane (Linux + Windows) with Claude and GPT
headless journeys that also assert no credential on disk.
Co-authored-by: Isaac <no-reply@databricks.com>
…fused connect Co-authored-by: Isaac <no-reply@databricks.com>
Co-authored-by: Isaac <no-reply@databricks.com>
… revert Co-authored-by: Isaac <no-reply@databricks.com>
…oken Co-authored-by: Isaac <no-reply@databricks.com>
Co-authored-by: Isaac <no-reply@databricks.com>
This was referenced Oct 21, 2026
xi-db
marked this pull request as ready for review
October 21, 2026 11:51
xi-db
requested review from
AarushiShah-db,
lilly-luo and
rohita5l
as code owners
October 21, 2026 11:51
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
Polly AI ReviewBlocking issuesNo separate correctness blockers identified. The credential-disclosure issue below should be fixed before merge. Security vulnerabilities
Non-blocking notes
ApproachThe approach is sound: reuse SummaryDiff size (computed): 27 files; +1384 / -217 lines (1601 changed text lines); 0 binary files.
Shares use additions + deletions, including generated text and lockfiles. Binary files have no line count. Tests include fixtures/helpers under test paths and colocated test/spec files; categories are path-based, not coverage measurements. ChangesThe PR moves Copilot’s provider authentication behind a per-session loopback proxy, refreshes upstream credentials, removes persisted legacy credentials, and keeps the proxy alive while the child runs. It also extracts shared auth selection and adds runner, CI, documentation, and compatibility coverage. The implementation and supporting changes are proportional to the authentication lifecycle being addressed. The distinct process, credential-revocation, and upgrade boundaries explain much of the test expansion. No dependency or lockfile changes are included. TestsTest-by-test assessmentAssessment is static. No passing test results are independently established. The frozen snapshot’s checksum matches. Both trusted policy files were read in full. The head revision was unavailable locally, and some frozen old blobs differ from the trusted default checkout. Changed hunks remain reviewable, but unchanged context was not verified against the complete target revision.
Existing Coverage is not established for proxy usability after an actual parent interrupt, or for legacy credential cleanup through direct launch/direct revert without reconfigure. The current helper and wait-loop assertions provide narrower evidence; these gaps alone do not demonstrate a product bug. For human verification, use disposable Linux and Windows environments: start a Copilot task, interrupt an active turn, confirm another request works, then exit Copilot and confirm its loopback listener closes. Do not print or archive real credentials. ScopeNo clearly unrelated changes were found. Proxy forwarding, shared auth extraction, credential cleanup, launch lifecycle, runner support, documentation, and the advisory CI lane all contribute to the stated outcome. Whether arbitrary inherited Automated review by Polly · workflow run |
…stall ug v0.2.0 without uv Co-authored-by: Isaac <no-reply@databricks.com>
…put on Windows Co-authored-by: Isaac <no-reply@databricks.com>
masonc08
approved these changes
Oct 21, 2026
masonc08
left a comment
masonc08
left a comment
Collaborator
There was a problem hiding this comment.
Have we tested this to make sure anthropic relay still works? Or at least double check with the agent that there's an anthropic relay integration test which passes
Collaborator
Author
Yes, those integration tests for anthropic relay still works. Also added new tests for copilot, on both linux and windows. |
Collaborator
Author
|
/merge |
xi-db
enabled auto-merge
October 21, 2026 15:15
rohita5l
approved these changes
Oct 21, 2026
luisscabral
pushed a commit
to luisscabral/unity-gateway
that referenced
this pull request
Oct 22, 2026
…s#1218) ## 🥞 Stacked PR Use this [link](https://github.com/databricks/unity-gateway/pull/1218/files) to review incremental changes. - [stack/copilot-auth-proxy](databricks#1217) [[Files changed](https://github.com/databricks/unity-gateway/pull/1217/files)] [MERGED] - [**stack/copilot-model-discovery**](databricks#1218) [[Files changed](https://github.com/databricks/unity-gateway/pull/1218/files)] ← _this PR_ - [stack/copilot-provider-services](databricks#1220) [[Files changed](https://github.com/databricks/unity-gateway/pull/1220/files/fe24b4f7810a1f7a800100d494218f611316ee3b..09bc0fb487cdc5a16da22508c61361031a9d4be3)] --------- ## What did you change, and why? **Change:** `ug copilot` now registers every Claude and GPT model the user can use on the workspace and lets them switch between them inside one Copilot session. - **Per-session registry.** Each launch writes a Copilot providers.json (`COPILOT_PROVIDERS_CONFIG`) to `~/.copilot/ug-providers-<id>.json` (0600), deleted when the session ends. Each session holds a lock on a sibling `.lock` file, and registries whose lock nobody holds (a killed session) are swept at the next launch and by `ug revert`. The user's own `~/.copilot/providers.json` is never touched. - Providers: `dbx-anthropic` (Anthropic Messages), `dbx-responses` (Responses, GPT-5+) and `dbx-completions` (Chat Completions, older GPT). All of them point at the existing loopback auth proxy with the session secret; only providers that have models are written. - Claude entries carry Copilot's catalog `modelId`, so Copilot sends Claude request settings, plus prompt/output limits. Without explicit limits Copilot assumes a 128K prompt for registry models. - The launch env sets `COPILOT_PROVIDERS_CONFIG`, `COPILOT_MODEL=<provider>/<id>` and `COPILOT_OFFLINE=true`, and drops inherited `COPILOT_PROVIDER*` settings. The single-provider env (`COPILOT_PROVIDER_TYPE/BASE_URL/WIRE_API/MODEL_ID`) and the `COPILOT_PROVIDER_WIRE_MODEL` override are gone. - **Discovery.** - The catalog is classified from the UC model-services listing, which now keeps each model's `supported_api_types`. When that listing is empty, it falls back to the gateway listings. No extra requests. - Names decide the wire: `claude-*` (including bare `databricks-claude-*`, databricks#108) -> Anthropic; `gpt-N` -> Responses for N >= 5, else Completions. Advertised API types only exclude. - Excluded, with a reason: Gemini (the gateway rejects Copilot's `stream_options`), open-source models (not validated with Copilot yet), non-chat models, `claude-opus-4-1` (the gateway 400s Copilot's thinking settings). - The catalog is saved per workspace. A failed listing keeps the last good one and warns; a listing with nothing usable is an error. - **Model selection.** - `--model`/`-m` accept a short name, `system.ai.*` or Copilot's `provider/id`. They're rewritten to `provider/id`, because Copilot only matches that form; Copilot has no `-m`, so it becomes `--model`. - An unknown or unsupported model fails before Copilot starts and lists the available models. - The start model is: `--model`, then the budget recommendation or managed default (now passed through `LaunchOptions.launch_model`; it was ignored for Copilot before), then the managed list, then the newest Sonnet/Opus/Haiku, then GPT. - **Managed config.** The admin's `model_services` list for Copilot now reaches state through the existing static-models seam and replaces discovery. Entries Copilot can't serve are skipped with a warning, and the wizard's Copilot families drop Gemini/OSS. - **Visibility.** `ug status` lists Copilot's catalog and default, and `ug doctor` reports the model counts. - **Other changes.** - Copilot CLI 1.0.91 or newer is required: older versions ignore the registry. - `ug copilot -p` keeps stdout for Copilot by moving ug's banners to stderr. - The shared `launch_token` helper is reused. **Why:** Copilot's env-var BYOK supports one provider per session, so `/model` couldn't reach Databricks models across families, and switching between Claude and GPT meant relaunching. Copilot also only saw the newest model per Claude family plus the Codex catalog, which leaked `gpt-oss`. Copilot doesn't discover gateway models itself the way Claude Code and Codex do, so ug builds the catalog. BEFORE: - `ug copilot` pins one model and wire per launch, and `/model` can't switch to another Databricks model family. - Older Claude versions and `databricks-claude-*` ids don't reach Copilot or take the OpenAI path. - A managed Copilot model list and the budget launch model are ignored. AFTER: - `/model` lists all authorized Claude and GPT models (on dogfood: 14 Claude + 18 GPT) and switches across providers mid-session, carrying tool history. - `--model` accepts friendly names and fails fast with the list. - Managed lists and the budget model apply. ## How do you know it works? **Testing:** - **Unit:** - classifier, catalog order and precedence, keep-last-good vs definitive empty, resolver and argv rewrite, launch-model order, Claude limits table and fallback; - registry content (pinned provider `type`/`wireApi`/`baseUrl`), file mode and cleanup, sweep of dead sessions, env, version gate; - managed static list and unservable warnings, status/doctor rows, `-p` stdout, the listing's API types; - `tests/test_agent_copilot.py` runs the real proxy with a stubbed Copilot process. - The full unit suite passes except 8 tests in `TestInstallToolBinary` / `test_e2e_user_agent`, which also fail on `main` in this environment. `just lint` is clean. - **Integration (new `tests/integration/test_ug_copilot_models.py`, non-blocking Copilot lane):** - one setup runs a Claude task, then a GPT task selected as `dbx-responses/<model>`; - an older Claude from the full catalog; - an unknown model exits 1 before Copilot starts and lists the models; - a TUI journey answers folder trust, completes a Claude task, runs `/model dbx-responses/<model>`, completes a GPT task and `/exit`s. - Model evidence comes from Copilot's own `session-state/*/events.jsonl` (`assistant.message.model`, `session.start.selectedModel`, `session.model_change.newModel`). The existing headless journeys now also assert the answering model and that no registry is left. - **Live (Copilot 1.0.98, `ug` from this branch, dogfood):** - `claude-haiku-4-5`, `dbx-responses/system.ai.gpt-5-4-nano` and `claude-sonnet-4-5` headless tasks completed with clean stdout; - `system.ai.gemini-3-flash` failed before launch with the list; - in the TUI, Haiku answered, `/model dbx-responses/system.ai.gpt-5-4-nano` printed "Model changed ... for this session", and GPT answered; - no registry was left. - The new integration journeys run in CI for the first time on this PR. **Not in this PR:** open-source models (next on this stack), Model Provider Service models in the registry, GPT token limits, mixing GitHub-seat models into `/model` (depends on the seat-based-auth decision). Stacked on databricks#1217. **Test infra deflakes (Copilot lanes):** on Windows, a dedicated `pwsh` install directory is added to the isolated test PATH, because Copilot's shell tool tries `pwsh` first and Windows PowerShell alone can miss its 40s startup probe, which left the revoke step with no shell. The Copilot GPT journeys also get the fixture's full path (`FileTask.prompt_with_path`), because gpt-5-4-nano sometimes drops ` with spaces` when it rebuilds the path. This pull request and its description were written by Isaac. --------- Co-authored-by: Isaac <no-reply@databricks.com>
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
3 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.
🥞 Stacked PR
Use this link to review incremental changes.
What did you change, and why?
Change:
ug copilotnow runs Copilot CLI against a per-session loopback proxy (gateway_proxy.start_gateway_auth_proxy) instead of handing it a Databricks token.COPILOT_PROVIDER_BASE_URL=http://127.0.0.1:<port>/ai-gateway/...and a random per-session secret asCOPILOT_PROVIDER_BEARER_TOKEN.x-api-key) and only on/ai-gateway/{anthropic,mlflow/v1,openai/v1}/(decoded until stable;./..segments rejected).TokenCache(refresh near expiry, force refresh plus one retry on 401/403), and addsucode/<v> copilot/<v>User-Agent, request tags, and managedcopilot_http_headers.databricks auth login --host <ws>, then restartug copilot."). The error goes to~/.ucode/logs/copilot-proxy.log, not over the TUI.~/.copilot/ucode.envand from a UG-produced backup, including beforeug revert, so revert cannot restore one.COPILOT_PROVIDER_API_KEY/_API_KEY_COMMAND/_TRANSPORTso they can't bypass the proxy.copilotintegration lane (Linux + Windows,live and copilot, outsidecujs) mirroring OpenCode's, and amends thetests/AGENTS.mdscope line to allow Copilot headless tests.Why: Copilot reads its provider credential once at startup.
ucode.env, which Copilot never re-reads, so sessions failed once the OAuth token expired (~1h).ucode.env, its backup, and the child env.save_statefrom a background thread, which could reset the active workspace during a concurrentug configure.BEFORE: a
ug copilotsession stops working after token expiry until restarted;~/.copilot/ucode.envand~/.ucode/copilot-ucode-env.backupcontain the Databricks token; Ctrl-C on Windows raisesValueErrorfromsend_signal(SIGINT).AFTER: requests carry a freshly minted token for the life of the session; no file or Copilot env var holds a Databricks credential; Ctrl-C is delivered to Copilot by the terminal and UG tears the proxy down after Copilot exits.
How do you know it works?
Testing:
tests/test_gateway_proxy_integration.py::TestGatewayAuthProxyEndToEnd(real sockets):x-api-keyaccepted and not forwarded;../,%2e%2e,%252e%252e, otel);tests/test_agent_copilot.py:tests/test_e2e.py::TestCopilotLaunchnow runs through the real proxy (and patchesucode.agents.copilot.save_state, which the old patch target missed).tests/integration/test_ug_copilot_headless.py: a Claude (Anthropic) task and a GPT (Responses) task, each asserting noDATABRICKS_BEARERon disk under~/.copilot/~/.ucode. Not run yet: needs the e2e workspace and Copilot 1.0.91. This draft is so the lane can run in CI.just lint(ruff check/format, ty) passes. The remaining unit failures inTestInstallToolBinary/test_e2e_user_agentalso fail on a cleanmaincheckout here.Open questions for reviewers:
tests/AGENTS.mdscope change needs CODEOWNERS sign-off.~/.copilot/providers.jsoncan still override UG's provider; guard left for a follow-up.Not in this PR: binary resolver (#970), version floor / wizard /
forceLoginOrgsdoctor,--workspace,-pstdout banner, MPS, OTEL, OSS models.This pull request and its description were written by Isaac.