From 2ba254d4cdb234e829ab69ec6ccc6773cec02240 Mon Sep 17 00:00:00 2001 From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com> Date: Sun, 20 Sep 2026 16:45:51 +0000 Subject: [PATCH 1/5] Refresh Claude gateway model cache before discovery launches --- src/ucode/agents/claude.py | 49 ++++++- src/ucode/databricks.py | 35 +++++ tests/README.md | 1 + tests/integration/README.md | 2 + .../test_ug_claude_managed_model_discovery.py | 22 +++- tests/test_agent_claude.py | 120 ++++++++++++++++++ tests/test_databricks.py | 71 +++++++++++ 7 files changed, 293 insertions(+), 7 deletions(-) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index ee7cc4e67..c2d5979ae 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -10,6 +10,7 @@ import socket import subprocess import threading +import time from collections.abc import Callable from pathlib import Path @@ -17,6 +18,7 @@ from ucode.config_io import ( APP_DIR, ToolSpec, + atomic_write_json, backup_existing_file, deep_merge_dict, read_json_safe, @@ -34,6 +36,7 @@ CustomOAuthConfig, build_custom_auth_shell_command, custom_oauth_cli_enabled, + get_custom_client_token, ) from ucode.databricks import ( AnthropicModelCatalog, @@ -41,6 +44,7 @@ build_otel_headers_shell_command, build_otel_traces_endpoint, build_tool_base_url, + fetch_anthropic_gateway_models, get_databricks_token, ug_binary, ) @@ -1453,6 +1457,38 @@ def _rewrite_relayed_port(state: dict, port: int) -> None: write_json_file(CLAUDE_SETTINGS_PATH, settings) +def _refresh_gateway_models_cache(state: dict) -> None: + if os.environ.get(GATEWAY_MODEL_DISCOVERY_ENV_VAR) != "1" or not state.get("workspace"): + return + workspace = state["workspace"] + custom_oauth = state.get("custom_oauth") + token = ( + get_custom_client_token(workspace, **custom_oauth) + if custom_oauth + else get_databricks_token(workspace, state.get("profile")) + ) + env = read_json_safe(CLAUDE_SETTINGS_PATH).get("env", {}) + headers = {} + for line in env.get("ANTHROPIC_CUSTOM_HEADERS", "").splitlines(): + name, separator, value = line.partition(":") + if separator: + headers[name.strip()] = value.strip() + models, reason = fetch_anthropic_gateway_models(workspace, token, headers=headers) + if models is None: + raise RuntimeError( + f"Could not refresh Claude models: {reason}. Check your model source and retry." + ) + config_dir = Path(os.environ.get("CLAUDE_CONFIG_DIR") or CLAUDE_CONFIG_DIR) + atomic_write_json( + config_dir / "cache" / "gateway-models.json", + { + "baseUrl": env.get("ANTHROPIC_BASE_URL") or build_tool_base_url("claude", workspace), + "fetchedAt": int(time.time() * 1000), + "models": models, + }, + ) + + def _launch_relayed(state: dict, binary: str, tool_args: list[str]) -> None: """Relayed launch: sign into the Claude subscription, start the loopback refresh proxy, then run Claude Code alongside it (the proxy must outlive the @@ -1479,12 +1515,14 @@ def token_provider(force_refresh: bool) -> str: server_thread = threading.Thread(target=server.serve_forever, daemon=True) server_thread.start() - proc = subprocess.Popen(_build_claude_argv(binary, tool_args, relayed=True)) try: - returncode = proc.wait() - except KeyboardInterrupt: - proc.send_signal(signal.SIGINT) - returncode = proc.wait() + _refresh_gateway_models_cache(state) + proc = subprocess.Popen(_build_claude_argv(binary, tool_args, relayed=True)) + try: + returncode = proc.wait() + except KeyboardInterrupt: + proc.send_signal(signal.SIGINT) + returncode = proc.wait() finally: cache.stop() server.shutdown() @@ -1507,6 +1545,7 @@ def launch( if state.get("claude_relayed"): _launch_relayed(state, binary, tool_args) return + _refresh_gateway_models_cache(state) # Smart routing needs Unix PTY support, which Windows does not provide. if options.launch_smart_routing and os.name == "nt": raise RuntimeError( diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 40b191485..51628ad2b 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -2646,6 +2646,41 @@ def _get_anthropic_models_json( ) +def fetch_anthropic_gateway_models( + workspace: str, token: str, *, headers: dict[str, str] +) -> tuple[list[dict] | None, str | None]: + """Fetch every page of the launch-scoped Claude gateway catalog.""" + url = f"https://{workspace_hostname(workspace)}{ANTHROPIC_MODELS_PATH}" + models: list[dict] = [] + cursors: set[str] = set() + params = {"limit": "1000"} + while True: + payload, reason = _http_get_json( + f"{url}?{urlencode(params)}", + token, + headers=headers, + max_retries=_ANTHROPIC_MODEL_DISCOVERY_SETUP_MAX_RETRIES, + ) + if payload is None: + return None, reason + if not isinstance(payload, dict) or not isinstance(payload.get("data"), list): + return None, "AI Gateway returned an invalid Anthropic model catalog" + page = payload["data"] + if any( + not isinstance(model, dict) or not isinstance(model.get("id"), str) or not model["id"] + for model in page + ): + return None, "AI Gateway returned an invalid Anthropic model catalog" + models.extend(page) + if not payload.get("has_more"): + return (models, None) if models else (None, "AI Gateway returned no Anthropic models") + cursor = payload.get("last_id") + if not isinstance(cursor, str) or not cursor or cursor in cursors: + return None, "AI Gateway returned an invalid Anthropic pagination cursor" + cursors.add(cursor) + params["after_id"] = cursor + + def list_anthropic_models(workspace: str, token: str) -> tuple[list[str], str | None]: """List every model id advertised by AI Gateway's Anthropic endpoint. diff --git a/tests/README.md b/tests/README.md index c6be4ec3d..3ad3d3686 100644 --- a/tests/README.md +++ b/tests/README.md @@ -54,6 +54,7 @@ All tests live directly in `integration/`; shared mechanics live in `utils/`. | `test_ug_configure_claude_rejects_invalid_credentials`, `test_ug_configure_codex_rejects_invalid_credentials` | Configure with a rejected bearer against the real workspace | Authentication failure; no successful saved setup | | `test_ug_configure_managed_claude`, `test_ug_configure_managed_codex` | Configure against a workspace that publishes a managed CodingAgentConfig | No agent selector; each agent's generated config exposes exactly the admin's static model_services; real gateway prompt on launch | | `test_case_01_*`, `test_case_03_*` | Launch managed Claude after configure and from fresh state, with personal discovery enabled and disabled | Claude receives the admin MPS header, caches native discovery results, and opens its real model picker | +| `test_case_01_managed_claude_uses_admin_discovery_after_configure` | Restart managed Claude in the same home | Cache timestamp advances past restart time despite an existing fresh cache; a real file-reading task completes | | `test_case_05_*`, `test_case_07_*` | Pass a provider or model-location override to managed Claude after configure and from fresh state | ug rejects the override before Claude starts and preserves agent-owned state | | `test_case_09_*`, `test_case_11_*` | Disable discovery and pass a provider or model-location override to managed Claude | ug still rejects both configured and fresh launches | | `test_case_02_*`, `test_case_04_*` | Launch managed Codex after configure and from fresh state, with personal discovery enabled and disabled | Codex exposes exactly the admin MPS-scoped catalog | diff --git a/tests/integration/README.md b/tests/integration/README.md index a1b13314a..6fa83ebbe 100644 --- a/tests/integration/README.md +++ b/tests/integration/README.md @@ -301,6 +301,8 @@ across all configured/fresh scenarios. The Codex module does the same with model picker, Codex's exact app-server catalog, and both agents' rejection of personal source overrides. Codex state comparisons exclude `.codex/tmp/arg0`, the disposable executable links recreated by version checks, while continuing to compare persistent agent files. +The configured Claude Case 1 also restarts in the same home, requires the gateway cache timestamp +to advance past restart time, and completes a real file-reading task in that second session. In addition, `test_ug_configure_managed_codex_catalog_fallback` injects the intentionally nonexistent `system.ai.gpt-99`, keeping it out of the real workspace while launching Codex through that workspace on the valid default model `system.ai.gpt-5-6-sol`. With smart routing enabled, it opens diff --git a/tests/integration/test_ug_claude_managed_model_discovery.py b/tests/integration/test_ug_claude_managed_model_discovery.py index fd00dfded..6797398cd 100644 --- a/tests/integration/test_ug_claude_managed_model_discovery.py +++ b/tests/integration/test_ug_claude_managed_model_discovery.py @@ -8,9 +8,11 @@ import json import os import re +import time import pytest from utils.constants import MANAGED_CLAUDE_PROVIDER_SERVICE +from utils.evidence import FileTask from utils.managed import ( fetch_managed_config_stub, is_managed_config_control_plane_cache, @@ -88,9 +90,10 @@ def _assert_managed_provider_in_picker(session, workspace, screen): @pytest.mark.tui def test_case_01_managed_claude_uses_admin_discovery_after_configure(live_session, workspace): - """Scenario: configure managed Claude, then launch its model picker. + """Scenario: configure managed Claude, open its model picker, then restart. - Expected: the managed model catalog wins after configuration. + Expected: the managed model catalog wins, a fresh cache replaces the prior + session's cache on restart, and Claude completes a real file-reading task. """ session = live_session result = session.run( @@ -111,6 +114,21 @@ def test_case_01_managed_claude_uses_admin_discovery_after_configure(live_sessio _assert_managed_provider_in_picker(session, workspace, screen) + cache_path = session.home / ".claude/cache/gateway-models.json" + previous = json.loads(cache_path.read_text()) + task = FileTask(session) + restarted_at = time.time_ns() // 1_000_000 + with AgentTerminal(session, "claude", command, "case-01-restarted") as tui: + tui.boot() + refreshed = json.loads(cache_path.read_text()) + assert refreshed["fetchedAt"] >= restarted_at > previous["fetchedAt"] + assert refreshed["baseUrl"] == previous["baseUrl"] + assert refreshed["models"] + tui.submit(task.prompt) + tui.wait_for_task(task) + tui.exit_normally() + task.assert_completed(session, "claude") + @pytest.mark.tui def test_case_01_fresh_managed_claude_uses_admin_discovery(live_session, workspace): diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index 0b3fdb3b7..fdda66ee5 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -1900,6 +1900,7 @@ def test_gateway_discovery_uses_direct_gateway(self, monkeypatch): monkeypatch.setenv(claude.GATEWAY_MODEL_DISCOVERY_ENV_VAR, "1") monkeypatch.delenv("OAUTH_TOKEN", raising=False) monkeypatch.setattr(claude, "get_databricks_token", lambda *_args: "token") + monkeypatch.setattr(claude, "_refresh_gateway_models_cache", Mock()) monkeypatch.setattr(claude, "exec_or_spawn", lambda argv: calls.append(argv)) claude.launch({"workspace": WS, "profile": "test"}, ["--debug"], options=LaunchOptions()) @@ -1914,6 +1915,7 @@ def test_gateway_discovery_enabled_under_provider(self, monkeypatch): monkeypatch.setenv(claude.GATEWAY_MODEL_DISCOVERY_ENV_VAR, "1") monkeypatch.delenv("OAUTH_TOKEN", raising=False) monkeypatch.setattr(claude, "get_databricks_token", lambda *_args: "token") + monkeypatch.setattr(claude, "_refresh_gateway_models_cache", Mock()) monkeypatch.setattr(claude, "exec_or_spawn", lambda argv: calls.append(argv)) claude.launch( @@ -1930,6 +1932,124 @@ def test_gateway_discovery_enabled_under_provider(self, monkeypatch): assert calls == [["claude", "--settings", str(claude.CLAUDE_SETTINGS_PATH), "--debug"]] +class TestGatewayModelsCache: + @pytest.fixture(autouse=True) + def setup(self, monkeypatch, tmp_path): + monkeypatch.setenv(claude.GATEWAY_MODEL_DISCOVERY_ENV_VAR, "1") + monkeypatch.setenv("CLAUDE_CONFIG_DIR", str(tmp_path / "custom-claude")) + monkeypatch.setattr(claude, "CLAUDE_SETTINGS_PATH", tmp_path / "settings.json") + monkeypatch.setattr(claude, "get_databricks_token", Mock(return_value="token")) + monkeypatch.setattr(claude, "exec_or_spawn", Mock()) + self.cache_path = tmp_path / "custom-claude/cache/gateway-models.json" + self.models = [{"id": "claude-new", "display_name": "New Claude"}] + self.fetch = Mock(return_value=(self.models, None)) + monkeypatch.setattr(claude, "fetch_anthropic_gateway_models", self.fetch) + + @pytest.mark.parametrize( + "headers", + [ + {}, + {"Databricks-Model-Provider-Service": "main.default.mps"}, + {"x-databricks-model-service-parent-schema": "main.models"}, + ], + ) + def test_refreshes_before_each_launch_with_configured_scope(self, monkeypatch, headers): + claude.write_json_file( + claude.CLAUDE_SETTINGS_PATH, + { + "env": { + "ANTHROPIC_BASE_URL": f"{WS}/ai-gateway/anthropic", + "ANTHROPIC_CUSTOM_HEADERS": "\n".join( + f"{name}: {value}" for name, value in headers.items() + ), + }, + }, + ) + snapshots = [] + monkeypatch.setattr( + claude, + "exec_or_spawn", + lambda argv: snapshots.append(json.loads(self.cache_path.read_text())), + ) + state = {"workspace": WS, "profile": "test"} + claude.launch(state, [], options=LaunchOptions()) + self.fetch.return_value = ([{"id": "claude-replacement"}], None) + claude.launch(state, [], options=LaunchOptions()) + self.fetch.assert_called_with(WS, "token", headers=headers) + assert self.fetch.call_count == 2 + assert snapshots[0]["models"] == self.models + assert snapshots[1]["models"] == [{"id": "claude-replacement"}] + assert snapshots[1]["baseUrl"] == f"{WS}/ai-gateway/anthropic" + assert snapshots[1]["fetchedAt"] > 0 + + def test_uses_custom_oauth(self, monkeypatch): + custom = { + "client_id": "client", + "redirect_url": "http://localhost:8020", + "scopes": ["all-apis", "offline_access"], + "profile": "custom", + } + token = Mock(return_value="custom-token") + monkeypatch.setattr(claude, "get_custom_client_token", token) + claude._refresh_gateway_models_cache({"workspace": WS, "custom_oauth": custom}) + token.assert_called_once_with(WS, **custom) + self.fetch.assert_called_once_with(WS, "custom-token", headers={}) + + def test_discovery_disabled_leaves_cache_untouched(self, monkeypatch): + monkeypatch.setenv(claude.GATEWAY_MODEL_DISCOVERY_ENV_VAR, "0") + claude.write_json_file(self.cache_path, {"models": ["existing"]}) + claude.launch({"workspace": WS}, [], options=LaunchOptions()) + self.fetch.assert_not_called() + assert json.loads(self.cache_path.read_text()) == {"models": ["existing"]} + + @pytest.mark.parametrize("failure", ["fetch", "write"]) + def test_failure_blocks_launch_without_replacing_cache(self, monkeypatch, failure): + claude.write_json_file(self.cache_path, {"models": ["existing"]}) + if failure == "fetch": + self.fetch.return_value = (None, "HTTP 403") + else: + monkeypatch.setattr(os, "replace", Mock(side_effect=OSError("disk full"))) + with pytest.raises(RuntimeError): + claude.launch({"workspace": WS}, [], options=LaunchOptions()) + claude.exec_or_spawn.assert_not_called() + assert json.loads(self.cache_path.read_text()) == {"models": ["existing"]} + + @pytest.mark.parametrize("failed", [False, True]) + def test_relay_refreshes_after_port_fallback_and_cleans_up(self, monkeypatch, failed): + server = Mock(server_address=("127.0.0.1", 54321)) + cache, client = Mock(), Mock() + monkeypatch.setattr(claude, "_ensure_subscription_login", Mock()) + monkeypatch.setattr( + claude.gateway_proxy, "start_relay_proxy", Mock(return_value=(server, cache, client)) + ) + claude.write_json_file( + claude.CLAUDE_SETTINGS_PATH, {"env": {"ANTHROPIC_BASE_URL": "http://127.0.0.1:12345"}} + ) + snapshots = [] + + def spawn(argv): + snapshots.append(json.loads(self.cache_path.read_text())) + return Mock(wait=Mock(return_value=0)) + + process = Mock(side_effect=spawn) + monkeypatch.setattr(claude.subprocess, "Popen", process) + if failed: + self.fetch.return_value = (None, "HTTP 403") + with pytest.raises(RuntimeError if failed else SystemExit): + claude.launch( + {"workspace": WS, "claude_relayed": True, "relayed_proxy_port": 12345}, + [], + options=LaunchOptions(), + ) + if failed: + process.assert_not_called() + else: + assert snapshots[0]["baseUrl"] == "http://127.0.0.1:54321" + server.shutdown.assert_called_once() + cache.stop.assert_called_once() + client.close.assert_called_once() + + class TestWriteToolConfigPrunesStaleModelEnv: """Stale ucode-managed model env keys (ANTHROPIC_MODEL, etc.) from earlier ucode versions must be removed on every launch — otherwise they linger in diff --git a/tests/test_databricks.py b/tests/test_databricks.py index ba6c364fe..bb4df3c03 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -310,6 +310,77 @@ def test_multiple_locations_preserve_order(self): ) +class TestFetchAnthropicGatewayModels: + def test_preserves_metadata_and_scope_across_pages(self, monkeypatch): + first = {"id": "claude/first", "display_name": "First", "description": "Model"} + second = {"id": "claude-second"} + requests = [] + pages = iter( + [ + {"data": [first], "has_more": True, "last_id": first["id"]}, + {"data": [second], "has_more": False}, + ] + ) + + def fetch(url, token, **kwargs): + requests.append((url, token, kwargs)) + return next(pages), None + + monkeypatch.setattr(db_mod, "_http_get_json", fetch) + headers = {"Databricks-Model-Provider-Service": "main.default.mps"} + assert db_mod.fetch_anthropic_gateway_models(WS, "token", headers=headers) == ( + [first, second], + None, + ) + assert [request[0] for request in requests] == [ + f"{WS}/ai-gateway/anthropic/v1/models?limit=1000", + f"{WS}/ai-gateway/anthropic/v1/models?limit=1000&after_id=claude%2Ffirst", + ] + assert all( + request[1:] == ("token", {"headers": headers, "max_retries": 2}) for request in requests + ) + + @pytest.mark.parametrize( + "payload", + [ + [], + {}, + {"data": []}, + {"data": [{"id": ""}]}, + {"data": ["claude"]}, + {"data": [{"id": "claude"}], "has_more": True}, + ], + ) + def test_rejects_empty_or_invalid_catalog(self, monkeypatch, payload): + monkeypatch.setattr(db_mod, "_http_get_json", lambda *args, **kwargs: (payload, None)) + models, reason = db_mod.fetch_anthropic_gateway_models(WS, "token", headers={}) + assert models is None + assert reason + + def test_rejects_repeated_cursor(self, monkeypatch): + monkeypatch.setattr( + db_mod, + "_http_get_json", + lambda *args, **kwargs: ( + {"data": [{"id": "claude"}], "has_more": True, "last_id": "claude"}, + None, + ), + ) + models, reason = db_mod.fetch_anthropic_gateway_models(WS, "token", headers={}) + assert models is None + assert "cursor" in reason + + def test_does_not_return_partial_catalog_on_failure(self, monkeypatch): + pages = iter( + [ + ({"data": [{"id": "claude"}], "has_more": True, "last_id": "claude"}, None), + (None, "HTTP 403"), + ] + ) + monkeypatch.setattr(db_mod, "_http_get_json", lambda *args, **kwargs: next(pages)) + assert db_mod.fetch_anthropic_gateway_models(WS, "token", headers={}) == (None, "HTTP 403") + + class TestDiscoverClaudeModels: def test_lists_all_anthropic_model_ids_without_legacy_validation(self, monkeypatch): captured = {} From f87fcd3bb4815e6f050d8c3ff208df19c1d7618b Mon Sep 17 00:00:00 2001 From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com> Date: Sun, 20 Sep 2026 16:53:13 +0000 Subject: [PATCH 2/5] Simplify cache validation and trim redundant test scaffolding --- src/ucode/databricks.py | 11 +++++----- tests/test_agent_claude.py | 23 ++------------------ tests/test_databricks.py | 43 ++++++++++---------------------------- 3 files changed, 18 insertions(+), 59 deletions(-) diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 51628ad2b..f7f792da3 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -2663,18 +2663,17 @@ def fetch_anthropic_gateway_models( ) if payload is None: return None, reason - if not isinstance(payload, dict) or not isinstance(payload.get("data"), list): - return None, "AI Gateway returned an invalid Anthropic model catalog" - page = payload["data"] - if any( + data = payload if isinstance(payload, dict) else {} + page = data.get("data") + if not isinstance(page, list) or any( not isinstance(model, dict) or not isinstance(model.get("id"), str) or not model["id"] for model in page ): return None, "AI Gateway returned an invalid Anthropic model catalog" models.extend(page) - if not payload.get("has_more"): + if not data.get("has_more"): return (models, None) if models else (None, "AI Gateway returned no Anthropic models") - cursor = payload.get("last_id") + cursor = data.get("last_id") if not isinstance(cursor, str) or not cursor or cursor in cursors: return None, "AI Gateway returned an invalid Anthropic pagination cursor" cursors.add(cursor) diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index fdda66ee5..72b984c6d 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -1982,19 +1982,6 @@ def test_refreshes_before_each_launch_with_configured_scope(self, monkeypatch, h assert snapshots[1]["baseUrl"] == f"{WS}/ai-gateway/anthropic" assert snapshots[1]["fetchedAt"] > 0 - def test_uses_custom_oauth(self, monkeypatch): - custom = { - "client_id": "client", - "redirect_url": "http://localhost:8020", - "scopes": ["all-apis", "offline_access"], - "profile": "custom", - } - token = Mock(return_value="custom-token") - monkeypatch.setattr(claude, "get_custom_client_token", token) - claude._refresh_gateway_models_cache({"workspace": WS, "custom_oauth": custom}) - token.assert_called_once_with(WS, **custom) - self.fetch.assert_called_once_with(WS, "custom-token", headers={}) - def test_discovery_disabled_leaves_cache_untouched(self, monkeypatch): monkeypatch.setenv(claude.GATEWAY_MODEL_DISCOVERY_ENV_VAR, "0") claude.write_json_file(self.cache_path, {"models": ["existing"]}) @@ -2025,13 +2012,7 @@ def test_relay_refreshes_after_port_fallback_and_cleans_up(self, monkeypatch, fa claude.write_json_file( claude.CLAUDE_SETTINGS_PATH, {"env": {"ANTHROPIC_BASE_URL": "http://127.0.0.1:12345"}} ) - snapshots = [] - - def spawn(argv): - snapshots.append(json.loads(self.cache_path.read_text())) - return Mock(wait=Mock(return_value=0)) - - process = Mock(side_effect=spawn) + process = Mock(return_value=Mock(wait=Mock(return_value=0))) monkeypatch.setattr(claude.subprocess, "Popen", process) if failed: self.fetch.return_value = (None, "HTTP 403") @@ -2044,7 +2025,7 @@ def spawn(argv): if failed: process.assert_not_called() else: - assert snapshots[0]["baseUrl"] == "http://127.0.0.1:54321" + assert json.loads(self.cache_path.read_text())["baseUrl"] == "http://127.0.0.1:54321" server.shutdown.assert_called_once() cache.stop.assert_called_once() client.close.assert_called_once() diff --git a/tests/test_databricks.py b/tests/test_databricks.py index bb4df3c03..15a91f630 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -341,45 +341,24 @@ def fetch(url, token, **kwargs): ) @pytest.mark.parametrize( - "payload", + "responses", [ - [], - {}, - {"data": []}, - {"data": [{"id": ""}]}, - {"data": ["claude"]}, - {"data": [{"id": "claude"}], "has_more": True}, + [({"data": []}, None)], + [({"data": [{"id": ""}]}, None)], + [({"data": [{"id": "claude"}], "has_more": True, "last_id": "claude"}, None)] * 2, + [ + ({"data": [{"id": "claude"}], "has_more": True, "last_id": "claude"}, None), + (None, "HTTP 403"), + ], ], ) - def test_rejects_empty_or_invalid_catalog(self, monkeypatch, payload): - monkeypatch.setattr(db_mod, "_http_get_json", lambda *args, **kwargs: (payload, None)) + def test_rejects_failed_or_invalid_catalog(self, monkeypatch, responses): + pages = iter(responses) + monkeypatch.setattr(db_mod, "_http_get_json", lambda *args, **kwargs: next(pages)) models, reason = db_mod.fetch_anthropic_gateway_models(WS, "token", headers={}) assert models is None assert reason - def test_rejects_repeated_cursor(self, monkeypatch): - monkeypatch.setattr( - db_mod, - "_http_get_json", - lambda *args, **kwargs: ( - {"data": [{"id": "claude"}], "has_more": True, "last_id": "claude"}, - None, - ), - ) - models, reason = db_mod.fetch_anthropic_gateway_models(WS, "token", headers={}) - assert models is None - assert "cursor" in reason - - def test_does_not_return_partial_catalog_on_failure(self, monkeypatch): - pages = iter( - [ - ({"data": [{"id": "claude"}], "has_more": True, "last_id": "claude"}, None), - (None, "HTTP 403"), - ] - ) - monkeypatch.setattr(db_mod, "_http_get_json", lambda *args, **kwargs: next(pages)) - assert db_mod.fetch_anthropic_gateway_models(WS, "token", headers={}) == (None, "HTTP 403") - class TestDiscoverClaudeModels: def test_lists_all_anthropic_model_ids_without_legacy_validation(self, monkeypatch): From ca34eba5c9174cbcea37e2ba6f3d2f0189d89d14 Mon Sep 17 00:00:00 2001 From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com> Date: Sun, 20 Sep 2026 17:56:07 +0000 Subject: [PATCH 3/5] Reuse launch credentials and Anthropic model request helper --- src/ucode/agents/claude.py | 17 +++++++++++------ src/ucode/databricks.py | 26 ++++++++++++++------------ tests/test_agent_claude.py | 10 +++++++--- 3 files changed, 32 insertions(+), 21 deletions(-) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index c2d5979ae..746187fbf 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -1457,15 +1457,15 @@ def _rewrite_relayed_port(state: dict, port: int) -> None: write_json_file(CLAUDE_SETTINGS_PATH, settings) -def _refresh_gateway_models_cache(state: dict) -> None: +def _refresh_gateway_models_cache(state: dict, *, token: str | None = None) -> str | None: if os.environ.get(GATEWAY_MODEL_DISCOVERY_ENV_VAR) != "1" or not state.get("workspace"): - return + return None workspace = state["workspace"] custom_oauth = state.get("custom_oauth") token = ( get_custom_client_token(workspace, **custom_oauth) if custom_oauth - else get_databricks_token(workspace, state.get("profile")) + else token or get_databricks_token(workspace, state.get("profile")) ) env = read_json_safe(CLAUDE_SETTINGS_PATH).get("env", {}) headers = {} @@ -1487,6 +1487,7 @@ def _refresh_gateway_models_cache(state: dict) -> None: "models": models, }, ) + return token def _launch_relayed(state: dict, binary: str, tool_args: list[str]) -> None: @@ -1516,7 +1517,8 @@ def token_provider(force_refresh: bool) -> str: server_thread.start() try: - _refresh_gateway_models_cache(state) + if os.environ.get(GATEWAY_MODEL_DISCOVERY_ENV_VAR) == "1": + _refresh_gateway_models_cache(state, token=cache.token) proc = subprocess.Popen(_build_claude_argv(binary, tool_args, relayed=True)) try: returncode = proc.wait() @@ -1545,7 +1547,7 @@ def launch( if state.get("claude_relayed"): _launch_relayed(state, binary, tool_args) return - _refresh_gateway_models_cache(state) + discovery_token = _refresh_gateway_models_cache(state) # Smart routing needs Unix PTY support, which Windows does not provide. if options.launch_smart_routing and os.name == "nt": raise RuntimeError( @@ -1566,7 +1568,10 @@ def launch( ) return if workspace and not custom_oauth_cli_enabled(state.get("custom_oauth")): - os.environ["OAUTH_TOKEN"] = get_databricks_token(workspace, state.get("profile")) + launch_token = None if state.get("custom_oauth") else discovery_token + os.environ["OAUTH_TOKEN"] = launch_token or get_databricks_token( + workspace, state.get("profile") + ) settings_override = None launch_args = list(tool_args) if options.user_pinned_model: diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index f7f792da3..6ab279ed1 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -2632,14 +2632,20 @@ def collect(result, schemas_done, schemas_total): def _get_anthropic_models_json( - workspace: str, token: str, *, parent_schema: str | None = None + workspace: str, + token: str, + *, + parent_schema: str | None = None, + headers: dict[str, str] | None = None, + params: dict[str, str] | None = None, ) -> tuple[dict | list | None, str | None]: - hostname = workspace_hostname(workspace) - headers = ( - {MODEL_SERVICE_PARENT_SCHEMA_HEADER: parent_schema} if parent_schema is not None else None - ) + if parent_schema is not None: + headers = {**(headers or {}), MODEL_SERVICE_PARENT_SCHEMA_HEADER: parent_schema} + url = f"https://{workspace_hostname(workspace)}{ANTHROPIC_MODELS_PATH}" + if params: + url += f"?{urlencode(params)}" return _http_get_json( - f"https://{hostname}{ANTHROPIC_MODELS_PATH}", + url, token, max_retries=_ANTHROPIC_MODEL_DISCOVERY_SETUP_MAX_RETRIES, **({"headers": headers} if headers is not None else {}), @@ -2650,16 +2656,12 @@ def fetch_anthropic_gateway_models( workspace: str, token: str, *, headers: dict[str, str] ) -> tuple[list[dict] | None, str | None]: """Fetch every page of the launch-scoped Claude gateway catalog.""" - url = f"https://{workspace_hostname(workspace)}{ANTHROPIC_MODELS_PATH}" models: list[dict] = [] cursors: set[str] = set() params = {"limit": "1000"} while True: - payload, reason = _http_get_json( - f"{url}?{urlencode(params)}", - token, - headers=headers, - max_retries=_ANTHROPIC_MODEL_DISCOVERY_SETUP_MAX_RETRIES, + payload, reason = _get_anthropic_models_json( + workspace, token, headers=headers, params=params ) if payload is None: return None, reason diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index 72b984c6d..a264d0e09 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -1900,7 +1900,7 @@ def test_gateway_discovery_uses_direct_gateway(self, monkeypatch): monkeypatch.setenv(claude.GATEWAY_MODEL_DISCOVERY_ENV_VAR, "1") monkeypatch.delenv("OAUTH_TOKEN", raising=False) monkeypatch.setattr(claude, "get_databricks_token", lambda *_args: "token") - monkeypatch.setattr(claude, "_refresh_gateway_models_cache", Mock()) + monkeypatch.setattr(claude, "_refresh_gateway_models_cache", Mock(return_value=None)) monkeypatch.setattr(claude, "exec_or_spawn", lambda argv: calls.append(argv)) claude.launch({"workspace": WS, "profile": "test"}, ["--debug"], options=LaunchOptions()) @@ -1915,7 +1915,7 @@ def test_gateway_discovery_enabled_under_provider(self, monkeypatch): monkeypatch.setenv(claude.GATEWAY_MODEL_DISCOVERY_ENV_VAR, "1") monkeypatch.delenv("OAUTH_TOKEN", raising=False) monkeypatch.setattr(claude, "get_databricks_token", lambda *_args: "token") - monkeypatch.setattr(claude, "_refresh_gateway_models_cache", Mock()) + monkeypatch.setattr(claude, "_refresh_gateway_models_cache", Mock(return_value=None)) monkeypatch.setattr(claude, "exec_or_spawn", lambda argv: calls.append(argv)) claude.launch( @@ -1977,6 +1977,8 @@ def test_refreshes_before_each_launch_with_configured_scope(self, monkeypatch, h claude.launch(state, [], options=LaunchOptions()) self.fetch.assert_called_with(WS, "token", headers=headers) assert self.fetch.call_count == 2 + assert claude.get_databricks_token.call_count == 2 + assert os.environ["OAUTH_TOKEN"] == "token" assert snapshots[0]["models"] == self.models assert snapshots[1]["models"] == [{"id": "claude-replacement"}] assert snapshots[1]["baseUrl"] == f"{WS}/ai-gateway/anthropic" @@ -2004,7 +2006,7 @@ def test_failure_blocks_launch_without_replacing_cache(self, monkeypatch, failur @pytest.mark.parametrize("failed", [False, True]) def test_relay_refreshes_after_port_fallback_and_cleans_up(self, monkeypatch, failed): server = Mock(server_address=("127.0.0.1", 54321)) - cache, client = Mock(), Mock() + cache, client = Mock(token="relay-token"), Mock() monkeypatch.setattr(claude, "_ensure_subscription_login", Mock()) monkeypatch.setattr( claude.gateway_proxy, "start_relay_proxy", Mock(return_value=(server, cache, client)) @@ -2029,6 +2031,8 @@ def test_relay_refreshes_after_port_fallback_and_cleans_up(self, monkeypatch, fa server.shutdown.assert_called_once() cache.stop.assert_called_once() client.close.assert_called_once() + self.fetch.assert_called_once_with(WS, "relay-token", headers={}) + claude.get_databricks_token.assert_not_called() class TestWriteToolConfigPrunesStaleModelEnv: From 2f6fe3d67e0da2f5c6b8939d89fcf271c2373632 Mon Sep 17 00:00:00 2001 From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com> Date: Sun, 20 Sep 2026 23:59:46 +0000 Subject: [PATCH 4/5] Keep Claude cache restart test independent of inference --- tests/README.md | 2 +- tests/integration/README.md | 3 ++- .../test_ug_claude_managed_model_discovery.py | 9 +++------ 3 files changed, 6 insertions(+), 8 deletions(-) diff --git a/tests/README.md b/tests/README.md index 3ad3d3686..6027c99f4 100644 --- a/tests/README.md +++ b/tests/README.md @@ -54,7 +54,7 @@ All tests live directly in `integration/`; shared mechanics live in `utils/`. | `test_ug_configure_claude_rejects_invalid_credentials`, `test_ug_configure_codex_rejects_invalid_credentials` | Configure with a rejected bearer against the real workspace | Authentication failure; no successful saved setup | | `test_ug_configure_managed_claude`, `test_ug_configure_managed_codex` | Configure against a workspace that publishes a managed CodingAgentConfig | No agent selector; each agent's generated config exposes exactly the admin's static model_services; real gateway prompt on launch | | `test_case_01_*`, `test_case_03_*` | Launch managed Claude after configure and from fresh state, with personal discovery enabled and disabled | Claude receives the admin MPS header, caches native discovery results, and opens its real model picker | -| `test_case_01_managed_claude_uses_admin_discovery_after_configure` | Restart managed Claude in the same home | Cache timestamp advances past restart time despite an existing fresh cache; a real file-reading task completes | +| `test_case_01_managed_claude_uses_admin_discovery_after_configure` | Restart managed Claude in the same home | Cache timestamp advances past restart time despite an existing fresh cache; the real model picker opens again; no inference is tested | | `test_case_05_*`, `test_case_07_*` | Pass a provider or model-location override to managed Claude after configure and from fresh state | ug rejects the override before Claude starts and preserves agent-owned state | | `test_case_09_*`, `test_case_11_*` | Disable discovery and pass a provider or model-location override to managed Claude | ug still rejects both configured and fresh launches | | `test_case_02_*`, `test_case_04_*` | Launch managed Codex after configure and from fresh state, with personal discovery enabled and disabled | Codex exposes exactly the admin MPS-scoped catalog | diff --git a/tests/integration/README.md b/tests/integration/README.md index 6fa83ebbe..9727675ab 100644 --- a/tests/integration/README.md +++ b/tests/integration/README.md @@ -302,7 +302,8 @@ model picker, Codex's exact app-server catalog, and both agents' rejection of pe overrides. Codex state comparisons exclude `.codex/tmp/arg0`, the disposable executable links recreated by version checks, while continuing to compare persistent agent files. The configured Claude Case 1 also restarts in the same home, requires the gateway cache timestamp -to advance past restart time, and completes a real file-reading task in that second session. +to advance past restart time, and opens the model picker again. This cache-refresh journey does +not exercise inference or require the provider's upstream inference credential to work. In addition, `test_ug_configure_managed_codex_catalog_fallback` injects the intentionally nonexistent `system.ai.gpt-99`, keeping it out of the real workspace while launching Codex through that workspace on the valid default model `system.ai.gpt-5-6-sol`. With smart routing enabled, it opens diff --git a/tests/integration/test_ug_claude_managed_model_discovery.py b/tests/integration/test_ug_claude_managed_model_discovery.py index 6797398cd..287f7ab93 100644 --- a/tests/integration/test_ug_claude_managed_model_discovery.py +++ b/tests/integration/test_ug_claude_managed_model_discovery.py @@ -12,7 +12,6 @@ import pytest from utils.constants import MANAGED_CLAUDE_PROVIDER_SERVICE -from utils.evidence import FileTask from utils.managed import ( fetch_managed_config_stub, is_managed_config_control_plane_cache, @@ -93,7 +92,7 @@ def test_case_01_managed_claude_uses_admin_discovery_after_configure(live_sessio """Scenario: configure managed Claude, open its model picker, then restart. Expected: the managed model catalog wins, a fresh cache replaces the prior - session's cache on restart, and Claude completes a real file-reading task. + session's cache on restart, and the model picker opens again. No inference is tested. """ session = live_session result = session.run( @@ -116,7 +115,6 @@ def test_case_01_managed_claude_uses_admin_discovery_after_configure(live_sessio cache_path = session.home / ".claude/cache/gateway-models.json" previous = json.loads(cache_path.read_text()) - task = FileTask(session) restarted_at = time.time_ns() // 1_000_000 with AgentTerminal(session, "claude", command, "case-01-restarted") as tui: tui.boot() @@ -124,10 +122,9 @@ def test_case_01_managed_claude_uses_admin_discovery_after_configure(live_sessio assert refreshed["fetchedAt"] >= restarted_at > previous["fetchedAt"] assert refreshed["baseUrl"] == previous["baseUrl"] assert refreshed["models"] - tui.submit(task.prompt) - tui.wait_for_task(task) + screen = tui.open_model_picker() tui.exit_normally() - task.assert_completed(session, "claude") + _assert_managed_provider_in_picker(session, workspace, screen) @pytest.mark.tui From 5ae071cc8683d0b82c6c329a52bb2b7230ca727a Mon Sep 17 00:00:00 2001 From: andy-xu-db <310751426+andy-xu-db@users.noreply.github.com> Date: Mon, 21 Sep 2026 01:23:20 +0000 Subject: [PATCH 5/5] Simplify Claude cache refresh token handling --- src/ucode/agents/claude.py | 15 +++++---------- tests/test_agent_claude.py | 6 ++---- 2 files changed, 7 insertions(+), 14 deletions(-) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 746187fbf..6d90d9233 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -1457,7 +1457,7 @@ def _rewrite_relayed_port(state: dict, port: int) -> None: write_json_file(CLAUDE_SETTINGS_PATH, settings) -def _refresh_gateway_models_cache(state: dict, *, token: str | None = None) -> str | None: +def _refresh_gateway_models_cache(state: dict) -> None: if os.environ.get(GATEWAY_MODEL_DISCOVERY_ENV_VAR) != "1" or not state.get("workspace"): return None workspace = state["workspace"] @@ -1465,7 +1465,7 @@ def _refresh_gateway_models_cache(state: dict, *, token: str | None = None) -> s token = ( get_custom_client_token(workspace, **custom_oauth) if custom_oauth - else token or get_databricks_token(workspace, state.get("profile")) + else get_databricks_token(workspace, state.get("profile")) ) env = read_json_safe(CLAUDE_SETTINGS_PATH).get("env", {}) headers = {} @@ -1487,7 +1487,6 @@ def _refresh_gateway_models_cache(state: dict, *, token: str | None = None) -> s "models": models, }, ) - return token def _launch_relayed(state: dict, binary: str, tool_args: list[str]) -> None: @@ -1517,8 +1516,7 @@ def token_provider(force_refresh: bool) -> str: server_thread.start() try: - if os.environ.get(GATEWAY_MODEL_DISCOVERY_ENV_VAR) == "1": - _refresh_gateway_models_cache(state, token=cache.token) + _refresh_gateway_models_cache(state) proc = subprocess.Popen(_build_claude_argv(binary, tool_args, relayed=True)) try: returncode = proc.wait() @@ -1547,7 +1545,7 @@ def launch( if state.get("claude_relayed"): _launch_relayed(state, binary, tool_args) return - discovery_token = _refresh_gateway_models_cache(state) + _refresh_gateway_models_cache(state) # Smart routing needs Unix PTY support, which Windows does not provide. if options.launch_smart_routing and os.name == "nt": raise RuntimeError( @@ -1568,10 +1566,7 @@ def launch( ) return if workspace and not custom_oauth_cli_enabled(state.get("custom_oauth")): - launch_token = None if state.get("custom_oauth") else discovery_token - os.environ["OAUTH_TOKEN"] = launch_token or get_databricks_token( - workspace, state.get("profile") - ) + os.environ["OAUTH_TOKEN"] = get_databricks_token(workspace, state.get("profile")) settings_override = None launch_args = list(tool_args) if options.user_pinned_model: diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index a264d0e09..ee9085183 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -1977,7 +1977,6 @@ def test_refreshes_before_each_launch_with_configured_scope(self, monkeypatch, h claude.launch(state, [], options=LaunchOptions()) self.fetch.assert_called_with(WS, "token", headers=headers) assert self.fetch.call_count == 2 - assert claude.get_databricks_token.call_count == 2 assert os.environ["OAUTH_TOKEN"] == "token" assert snapshots[0]["models"] == self.models assert snapshots[1]["models"] == [{"id": "claude-replacement"}] @@ -2006,7 +2005,7 @@ def test_failure_blocks_launch_without_replacing_cache(self, monkeypatch, failur @pytest.mark.parametrize("failed", [False, True]) def test_relay_refreshes_after_port_fallback_and_cleans_up(self, monkeypatch, failed): server = Mock(server_address=("127.0.0.1", 54321)) - cache, client = Mock(token="relay-token"), Mock() + cache, client = Mock(), Mock() monkeypatch.setattr(claude, "_ensure_subscription_login", Mock()) monkeypatch.setattr( claude.gateway_proxy, "start_relay_proxy", Mock(return_value=(server, cache, client)) @@ -2031,8 +2030,7 @@ def test_relay_refreshes_after_port_fallback_and_cleans_up(self, monkeypatch, fa server.shutdown.assert_called_once() cache.stop.assert_called_once() client.close.assert_called_once() - self.fetch.assert_called_once_with(WS, "relay-token", headers={}) - claude.get_databricks_token.assert_not_called() + self.fetch.assert_called_once_with(WS, "token", headers={}) class TestWriteToolConfigPrunesStaleModelEnv: