diff --git a/tests/test_e2e.py b/tests/test_e2e.py index 12f6f601..94ec0012 100644 --- a/tests/test_e2e.py +++ b/tests/test_e2e.py @@ -20,6 +20,7 @@ from urllib import error as urllib_error from urllib import request as urllib_request +import httpx import pytest from ucode.databricks import ( @@ -574,23 +575,29 @@ def _skip_if_provider_unusable(combined: str, provider: str) -> None: if "Credit balance is too low" in combined: pytest.skip(f"provider {provider} account is out of credits: {combined[:200]}") - def test_launch_claude_through_provider( - self, tmp_path, monkeypatch, e2e_state, e2e_workspace, e2e_token - ): - import ucode.config_io as config_io_mod - from ucode.agents import claude, resolve_provider_models + def _resolve_anthropic_mps(self, e2e_state, e2e_workspace): + """Return (state, provider_models) for the pinned Anthropic MPS, or skip if it's absent.""" + from ucode.agents import resolve_provider_models - _require_binary("claude") - # Pinned to a fixed CI MPS so the inference is deterministic and cheap. Skip (don't fail) - # when it's absent, so e2e runs against other workspaces still work. - provider = CI_ANTHROPIC_MPS state = {**e2e_state, "workspace": e2e_workspace} - # Resolve the provider's models as the launch path does. This Anthropic MPS declares a - # single Haiku target, so provider_models pins that family and route_root_model (below) - # makes Claude Code launch on exactly it — its built-in default tier isn't in the allowlist. - provider_models, error, _relayed = resolve_provider_models("claude", state, provider) + provider_models, error, _relayed = resolve_provider_models( + "claude", state, CI_ANTHROPIC_MPS + ) if error is not None: - pytest.skip(f"CI Anthropic MPS {provider} unavailable on this workspace: {error}") + pytest.skip( + f"CI Anthropic MPS {CI_ANTHROPIC_MPS} unavailable on this workspace: {error}" + ) + return state, provider_models + + def _launch_claude_and_assert_success( + self, tmp_path, monkeypatch, state, provider_models, route_root_model, workspace, token + ): + """Write the provider config, run a one-turn claude launch, assert it succeeds. + + This MPS allows only Haiku, so any other model 403s — a successful launch is itself + proof the request went out on the pinned Haiku target (route_root_model).""" + import ucode.config_io as config_io_mod + from ucode.agents import claude config_dir = tmp_path / "claude_config" config_dir.mkdir() @@ -603,25 +610,59 @@ def test_launch_claude_through_provider( claude.write_tool_config( state, None, - provider=provider, + provider=CI_ANTHROPIC_MPS, provider_models=provider_models, - route_root_model=CI_ANTHROPIC_MODEL, + route_root_model=route_root_model, ) env = { **os.environ, "CLAUDE_CONFIG_DIR": str(config_dir), - "ANTHROPIC_BASE_URL": build_tool_base_url("claude", e2e_workspace), - "ANTHROPIC_API_KEY": e2e_token, + "ANTHROPIC_BASE_URL": build_tool_base_url("claude", workspace), + "ANTHROPIC_API_KEY": token, } result = _run_agent(claude.validate_cmd("claude"), env=env, timeout=90) combined = (result.stdout + result.stderr).strip() - self._skip_if_provider_unusable(combined, provider) + self._skip_if_provider_unusable(combined, CI_ANTHROPIC_MPS) assert result.returncode == 0 and combined, ( - f"provider={provider} rc={result.returncode} " + f"model={route_root_model} rc={result.returncode} " f"stdout={result.stdout[:300]!r} stderr={result.stderr[:300]!r}" ) + def test_launch_claude_through_provider( + self, tmp_path, monkeypatch, e2e_state, e2e_workspace, e2e_token + ): + # No --model: ucode picks the only servable tier. The MPS declares no opus, so the default + # resolves to Haiku — and the launch succeeds only because Haiku is what the MPS allows. + from ucode.databricks import resolve_provider_launch_model + + _require_binary("claude") + state, provider_models = self._resolve_anthropic_mps(e2e_state, e2e_workspace) + launch_model = resolve_provider_launch_model(None, provider_models or {}) + assert launch_model == CI_ANTHROPIC_MODEL, ( + f"a bare launch should default to the only allowed model; got {launch_model}" + ) + self._launch_claude_and_assert_success( + tmp_path, monkeypatch, state, provider_models, launch_model, e2e_workspace, e2e_token + ) + + def test_launch_claude_model_haiku_through_provider( + self, tmp_path, monkeypatch, e2e_state, e2e_workspace, e2e_token + ): + # `--model haiku`: the family alias resolves to the one declared Haiku target, so the launch + # routes. (Claude Code's bare `haiku`/undated id would 403 — the gateway exact-matches.) + from ucode.databricks import resolve_provider_launch_model + + _require_binary("claude") + state, provider_models = self._resolve_anthropic_mps(e2e_state, e2e_workspace) + launch_model = resolve_provider_launch_model("haiku", provider_models or {}) + assert launch_model == CI_ANTHROPIC_MODEL, ( + f"--model haiku should resolve to the declared Haiku target; got {launch_model}" + ) + self._launch_claude_and_assert_success( + tmp_path, monkeypatch, state, provider_models, launch_model, e2e_workspace, e2e_token + ) + def test_launch_codex_through_provider( self, tmp_path, monkeypatch, e2e_state, e2e_workspace, e2e_token ): @@ -666,6 +707,69 @@ def test_launch_codex_through_provider( ) +class TestAnthropicNonRelayMps: + """Non-relay (api-key) Anthropic MPS, pinned to CI_ANTHROPIC_MPS (Haiku-only). + + Exercised without the agent binary — direct Anthropic Messages calls through the + gateway MPS. The MPS exact-matches its allowlist, so only claude-haiku-4-5-20251001 + routes; a 200 is itself proof the request hit Haiku. + """ + + @staticmethod + def _messages_request(workspace: str, token: str, model: str) -> httpx.Response: + body = { + "model": model, + "max_tokens": 16, + "stream": False, + "messages": [{"role": "user", "content": "say hi in 5 words or less"}], + } + return httpx.post( + f"{build_tool_base_url('claude', workspace)}/v1/messages", + headers={ + "Authorization": f"Bearer {token}", + "Databricks-Model-Provider-Service": CI_ANTHROPIC_MPS, + "x-databricks-use-coding-agent-mode": "true", + }, + json=body, + timeout=60, + ) + + @staticmethod + def _skip_if_unusable(resp: httpx.Response) -> None: + # Environmental account conditions, not ucode bugs (mirrors _skip_if_provider_unusable). + text = resp.text + if resp.status_code in (401, 403) and ("USE CONNECTION" in text or "EXECUTE" in text): + pytest.skip(f"no permission on {CI_ANTHROPIC_MPS}: {text[:200]}") + if "credit balance is too low" in text.lower(): + pytest.skip(f"{CI_ANTHROPIC_MPS} account is out of credits: {text[:200]}") + + def test_direct_inference_routes_to_haiku(self, e2e_workspace, e2e_token): + # A real inference through the non-relay Anthropic MPS, no agent binary involved. + resp = self._messages_request(e2e_workspace, e2e_token, CI_ANTHROPIC_MODEL) + self._skip_if_unusable(resp) + assert resp.status_code == 200, f"HTTP {resp.status_code}: {resp.text[:300]}" + assert resp.json().get("model") == CI_ANTHROPIC_MODEL + + def test_rejects_model_outside_allowlist(self, e2e_workspace, e2e_token): + # The MPS enforces its allowlist up front: a model it doesn't declare is refused. + resp = self._messages_request(e2e_workspace, e2e_token, "claude-sonnet-4-5") + self._skip_if_unusable(resp) + assert resp.status_code == 403, f"HTTP {resp.status_code}: {resp.text[:300]}" + assert "not in the allowed models list" in resp.text + + def test_mps_is_listed_usable_and_nonrelayed(self, e2e_workspace, e2e_token): + services, reason = list_model_provider_services(e2e_workspace, e2e_token) + if is_model_provider_feature_unavailable(reason): + pytest.skip("Model Provider Service feature not enabled on this workspace") + assert reason is None, reason + entry = next((s for s in services if s["name"] == CI_ANTHROPIC_MPS), None) + assert entry is not None, ( + f"{CI_ANTHROPIC_MPS} not found among {[s['name'] for s in services][:10]}" + ) + assert service_usable_for_tool("claude", entry) + assert not entry.get("relayed") + + class TestGeminiLaunch: """Run gemini against every available gemini model."""