copilot: Route ug copilot through a refreshing loopback auth proxy by xi-db · Pull Request #1191 · databricks/unity-gateway

xi-db · 2026-10-19T20:53:38Z

🥞 Stacked PR

Use this link to review incremental changes.

stack/xi/copilot-auth-proxy [Files changed] ← this PR
- stack/copilot-model-discovery [Files changed]

What did you change, and why?

Change: ug copilot now runs Copilot CLI against a per-session loopback proxy (gateway_proxy.start_gateway_auth_proxy) instead of handing it a Databricks token.

Copilot gets COPILOT_PROVIDER_BASE_URL=http://127.0.0.1:<port>/ai-gateway/... and a random per-session secret as COPILOT_PROVIDER_BEARER_TOKEN.
The proxy accepts only that secret (Bearer or x-api-key) and only on /ai-gateway/{anthropic,mlflow/v1,openai/v1}/ (decoded until stable; ./.. segments rejected).
It swaps the secret for a fresh Databricks token from the CLI profile or custom OAuth, using TokenCache (refresh near expiry, force refresh plus one retry on 401/403), and adds ucode/<v> copilot/<v> User-Agent, request tags, and managed copilot_http_headers.
If refresh fails, Copilot gets a dialect-shaped 401 ("Databricks sign-in expired. Run databricks auth login --host <ws>, then restart ug copilot."). The error goes to ~/.ucode/logs/copilot-proxy.log, not over the TUI.
Configure and launch persist no credential. Tokens written by earlier versions are scrubbed from ~/.copilot/ucode.env and from a UG-produced backup, including before ug revert, so revert cannot restore one.
Removes the Copilot token-refresher thread and its SIGINT forwarding; the parent keeps waiting on Ctrl-C so the proxy outlives Copilot.
Drops inherited COPILOT_PROVIDER_API_KEY / _API_KEY_COMMAND / _TRANSPORT so they can't bypass the proxy.
Adds a non-blocking copilot integration lane (Linux + Windows, live and copilot, outside cujs) mirroring OpenCode's, and amends the tests/AGENTS.md scope line to allow Copilot headless tests.

Why: Copilot reads its provider credential once at startup.

Expiry: the old refresher rewrote ucode.env, which Copilot never re-reads, so sessions failed once the OAuth token expired (~1h).
Plaintext tokens: the token sat in plaintext in ucode.env, its backup, and the child env.
Shared-state writes: the refresher also called save_state from a background thread, which could reset the active workspace during a concurrent ug configure.

BEFORE: a ug copilot session stops working after token expiry until restarted; ~/.copilot/ucode.env and ~/.ucode/copilot-ucode-env.backup contain the Databricks token; Ctrl-C on Windows raises ValueError from send_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):
- secret swap and added headers;
- x-api-key accepted and not forwarded;
- a wrong or missing secret gets a 401 and nothing is forwarded;
- path refusals (../, %2e%2e, %252e%252e, otel);
- a 401 from the gateway triggers a forced refresh and one retry;
- the actionable 401 is shaped for each dialect, with nothing on stderr.
tests/test_agent_copilot.py:
- no credential in the env overlay or on disk;
- scrubbing legacy tokens and UG-produced backups;
- token provider for both the CLI profile and custom OAuth;
- header precedence;
- launch through the real proxy with a stubbed child process.
tests/test_e2e.py::TestCopilotLaunch now runs through the real proxy (and patches ucode.agents.copilot.save_state, which the old patch target missed).
New tests/integration/test_ug_copilot_headless.py: a Claude (Anthropic) task and a GPT (Responses) task, each asserting no DATABRICKS_BEARER on 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.
Locally on Linux: 807 passed across the touched suites; just lint (ruff check/format, ty) passes. The remaining unit failures in TestInstallToolBinary / test_e2e_user_agent also fail on a clean main checkout here.

Open questions for reviewers:

tests/AGENTS.md scope change needs CODEOWNERS sign-off.
Revocation returns an actionable 401 rather than retrying with a browser login, because UG's login path writes to the terminal Copilot owns. A browser-login retry is a follow-up.
A user ~/.copilot/providers.json can still override UG's provider; guard left for a follow-up.
Still to check live: how Copilot renders the proxy's 401 body; Windows firewall behavior on the loopback bind.

Not in this PR: binary resolver (#970), version floor / wizard / forceLoginOrgs doctor, --workspace, -p stdout banner, MPS, OTEL, OSS models.

This pull request and its description were written by Isaac.

xi-db · 2026-10-19T21:28:36Z

/integration-test 4d3de6d

xi-db · 2026-10-19T21:28:38Z

/user-journey-check

github-actions · 2026-10-19T21:29:33Z

User-journey check failure on 4d3de6d: https://github.com/databricks/unity-gateway/actions/runs/39327792719

xi-db · 2026-10-19T21:49:07Z

/integration-test 810de53

xi-db · 2026-10-19T21:49:10Z

/user-journey-check

github-actions · 2026-10-19T21:49:35Z

User-journey check failure on 810de53: https://github.com/databricks/unity-gateway/actions/runs/39330448036

github-actions · 2026-10-19T21:49:55Z

Fork integration cancelled on 4d3de6d: https://github.com/databricks/unity-gateway/actions/runs/39327788767

xi-db · 2026-10-19T21:55:29Z

/integration-test 6fdbcd4

xi-db · 2026-10-19T21:55:31Z

/user-journey-check

github-actions · 2026-10-19T21:55:54Z

User-journey check success on 6fdbcd4: https://github.com/databricks/unity-gateway/actions/runs/39331233684

github-actions · 2026-10-19T21:56:18Z

Fork integration cancelled on 810de53: https://github.com/databricks/unity-gateway/actions/runs/39330443015

github-actions · 2026-10-19T22:22:58Z

Fork integration failure on 6fdbcd4: https://github.com/databricks/unity-gateway/actions/runs/39331229444

xi-db · 2026-10-19T22:23:19Z

/integration-test 6fdbcd4

github-actions · 2026-10-19T22:49:03Z

Fork integration success on 6fdbcd4: https://github.com/databricks/unity-gateway/actions/runs/39334584282

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>

xi-db · 2026-10-20T09:49:08Z

/integration-test f6ebc19

xi-db · 2026-10-20T09:49:11Z

/user-journey-check

github-actions · 2026-10-20T09:49:31Z

User-journey check success on f6ebc19: https://github.com/databricks/unity-gateway/actions/runs/39403460686

xi-db · 2026-10-20T09:50:26Z

/ug-review

github-actions · 2026-10-20T09:54:21Z

+    return state
+
+
+def _token_provider(state: dict) -> Callable[[bool], str]:


this is the same as codex._launch_token, now a third copy of the custom-oauth/profile selection. can we move it into a shared helper so the copies don't drift apart?

Fixed in d480b81: moved it to custom_oauth.launch_token, and both codex._launch_token and Copilot's proxy token provider now call it. I left smart_routing/v2._launch_token alone because it differs: it only uses custom OAuth when the CLI backend is enabled, so folding it in would change smart-routing behavior.

github-actions · 2026-10-20T09:54:21Z

+    assert_no_bearer_on_disk(session)
+
+
+def test_ug_copilot_headless_gpt_responses_task(live_session, workspace):


this is the claude test again with only the model changed, and it repeats the configure args too. can we parametrize over CLAUDE_TEST_MODEL/CODEX_TEST_MODEL and share one configure helper with the lifecycle file?

Keeping these as two named tests. tests/AGENTS.md asks for explicit journey names instead of hiding the provider behind parametrization, and for the configure command to stay visible in each test. They also cover different wire APIs (Anthropic Messages vs Responses), so a failure name tells you which one broke.

github-actions · 2026-10-20T09:54:22Z

+    """No file ug or Copilot wrote under the session home holds the Databricks bearer."""
+    bearer = session.env.get("DATABRICKS_BEARER")
+    assert bearer, "The live suite authenticates with DATABRICKS_BEARER."
+    for root in (session.home / ".copilot", session.home / ".ucode"):


the docstrings say no file under the session home holds the bearer, but this only scans .copilot/.ucode. a leak to e.g. ~/.databrickscfg or a shell rc still passes. can we rglob session.home?

Suggestion:

Suggested change

for root in (session.home / ".copilot", session.home / ".ucode"):

for path in session.home.rglob("*"):

Fixed in d480b81: it now scans every file under session.home, comparing bytes so large cache files stay cheap.

github-actions · 2026-10-20T09:54:22Z

+        finally:
+            cache.stop()
+            server.shutdown()
+            client.close()


can we add server.server_close() here like launch does? without it the listening socket stays open for the rest of the e2e session.

Suggestion:

Suggested change

finally:

cache.stop()

server.shutdown()

client.close()

finally:

cache.stop()

server.shutdown()

server.server_close()

client.close()

Fixed in d480b81: added server.server_close() after server.shutdown().

Co-authored-by: Isaac <no-reply@databricks.com>

xi-db · 2026-10-20T10:01:33Z

/integration-test d480b81

xi-db · 2026-10-20T10:01:36Z

/user-journey-check

github-actions · 2026-10-20T10:01:58Z

User-journey check success on d480b81: https://github.com/databricks/unity-gateway/actions/runs/39404962298

github-actions · 2026-10-20T10:03:02Z

Fork integration cancelled on f6ebc19: https://github.com/databricks/unity-gateway/actions/runs/39403455354

github-actions · 2026-10-20T10:28:41Z

Fork integration failure on d480b81: https://github.com/databricks/unity-gateway/actions/runs/39404958508

xi-db · 2026-10-20T13:22:58Z

+def test_ug_copilot_refreshes_a_revoked_token_mid_session(live_session, workspace):
+    """Scenario: sign in through a bearer command, then revoke the session's token mid-task.
+
+    The headless Claude task first runs a shell step that deletes every PAT minted so far,
+    including the one ug's proxy holds, then reads the fixture.
+
+    Expected: the gateway rejects the revoked token, ug mints a new one without restarting
+    Copilot, and Copilot's completed answer contains the fixture's value. Limits: needs the
+    e2e principal to create PATs; a gateway that still honors a just-deleted PAT fails the
+    mint-after-revoke assertion.
+    """


We will be able to test to token expire and refresh senario by revoking the databricks tokens.

tests/integration/utils/pat_bearer.py is able to revoke tokens to reproduce this case.

xi-db · 2026-10-21T11:51:39Z

Moved to #1217 (in-repo branch, so CI runs this PR's own workflow). Closing in favor of it.

xi-db changed the title ~~Route ug copilot through a refreshing loopback auth proxy~~ copilot: Route ug copilot through a refreshing loopback auth proxy Oct 19, 2026

xi-db closed this Oct 19, 2026

xi-db deleted the xi/copilot-auth-proxy branch October 19, 2026 21:13

xi-db restored the xi/copilot-auth-proxy branch October 19, 2026 21:14

xi-db reopened this Oct 19, 2026

xi-db force-pushed the xi/copilot-auth-proxy branch from 4a35939 to 4d3de6d Compare October 19, 2026 21:26

xi-db marked this pull request as ready for review October 20, 2026 09:47

xi-db requested review from AarushiShah-db, lilly-luo and rohita5l as code owners October 20, 2026 09:47

xi-db and others added 5 commits October 20, 2026 09:47

Close the Copilot proxy socket on exit and allow for Windows' slow re…

204cc8d

…fused connect Co-authored-by: Isaac <no-reply@databricks.com>

Drop test comment

1532e9f

Co-authored-by: Isaac <no-reply@databricks.com>

copilot: integration journeys for token refresh, expired sign-in, and…

6773cf6

… revert Co-authored-by: Isaac <no-reply@databricks.com>

copilot: integration journey for upgrading from a ug that saved the t…

f6ebc19

…oken Co-authored-by: Isaac <no-reply@databricks.com>

xi-db force-pushed the xi/copilot-auth-proxy branch from 6fdbcd4 to f6ebc19 Compare October 20, 2026 09:48

github-actions Bot reviewed Oct 20, 2026

View reviewed changes

Share the launch token selection; review fixes in Copilot tests

d480b81

Co-authored-by: Isaac <no-reply@databricks.com>

xi-db mentioned this pull request Oct 20, 2026

copilot: discover and switch between all authorized models #1203

Closed

xi-db added the ug-review Run the automated UG review label Oct 20, 2026

xi-db commented Oct 20, 2026

View reviewed changes

xi-db requested a review from masonc08 October 21, 2026 11:26

xi-db mentioned this pull request Oct 21, 2026

copilot: Route ug copilot through a refreshing loopback auth proxy #1217

Merged

xi-db closed this Oct 21, 2026

		return state


		def _token_provider(state: dict) -> Callable[[bool], str]:

		assert_no_bearer_on_disk(session)


		def test_ug_copilot_headless_gpt_responses_task(live_session, workspace):

	for root in (session.home / ".copilot", session.home / ".ucode"):
	for path in session.home.rglob("*"):

Conversation

xi-db commented Oct 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🥞 Stacked PR

What did you change, and why?

How do you know it works?

Uh oh!

xi-db commented Oct 19, 2026

Uh oh!

xi-db commented Oct 19, 2026

Uh oh!

github-actions Bot commented Oct 19, 2026

Uh oh!

xi-db commented Oct 19, 2026

Uh oh!

xi-db commented Oct 19, 2026

Uh oh!

github-actions Bot commented Oct 19, 2026

Uh oh!

github-actions Bot commented Oct 19, 2026

Uh oh!

xi-db commented Oct 19, 2026

Uh oh!

xi-db commented Oct 19, 2026

Uh oh!

github-actions Bot commented Oct 19, 2026

Uh oh!

github-actions Bot commented Oct 19, 2026

Uh oh!

github-actions Bot commented Oct 19, 2026

Uh oh!

xi-db commented Oct 19, 2026

Uh oh!

github-actions Bot commented Oct 19, 2026

Uh oh!

xi-db commented Oct 20, 2026

Uh oh!

xi-db commented Oct 20, 2026

Uh oh!

github-actions Bot commented Oct 20, 2026

Uh oh!

xi-db commented Oct 20, 2026

Uh oh!

github-actions Bot Oct 20, 2026

Choose a reason for hiding this comment

Uh oh!

xi-db Oct 20, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Oct 20, 2026

Choose a reason for hiding this comment

Uh oh!

xi-db Oct 20, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

github-actions Bot Oct 20, 2026

Choose a reason for hiding this comment

Uh oh!

xi-db Oct 20, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot Oct 20, 2026

Choose a reason for hiding this comment

Uh oh!

xi-db Oct 20, 2026

Choose a reason for hiding this comment

Uh oh!

xi-db commented Oct 20, 2026

Uh oh!

xi-db commented Oct 20, 2026

Uh oh!

github-actions Bot commented Oct 20, 2026

Uh oh!

github-actions Bot commented Oct 20, 2026

Uh oh!

github-actions Bot commented Oct 20, 2026

Uh oh!

xi-db Oct 20, 2026

Choose a reason for hiding this comment

xi-db commented Oct 19, 2026 •

edited

Loading