Support headless relayed launch with CLAUDE_CODE_OAUTH_TOKEN + relayed e2e test

src/ucode/agents/claude.py

@@ -57,6 +57,10 @@
 from .args import has_explicit_model_arg
 
 GATEWAY_MODEL_DISCOVERY_ENV_VAR = "ENABLE_CLAUDE_CODE_GATEWAY_MODEL_DISCOVERY"
+# A pre-provisioned subscription OAuth token (`claude setup-token` output) that
+# Claude Code reads as the Authorization credential — the headless/CI substitute
+# for the interactive `claude auth login` browser flow.
+CLAUDE_CODE_OAUTH_TOKEN_ENV_VAR = "CLAUDE_CODE_OAUTH_TOKEN"
 CLAUDE_CONFIG_DIR = Path.home() / ".claude"
 CLAUDE_SETTINGS_PATH = CLAUDE_CONFIG_DIR / "ucode-settings.json"
 CLAUDE_MCP_CONFIG_PATH = Path.home() / ".claude.json"
@@ -1312,6 +1316,12 @@ def _ensure_subscription_login() -> None:
     """Ensure Claude Code has a persisted subscription login, running the browser
     flow via `claude auth login` if not. ucode never sees or stores the token —
     Claude Code persists it to its own secure store and refreshes it natively."""
+    # A pre-provisioned OAuth token (CI/headless) is the Authorization credential
+    # directly, so no interactive login applies — skip the browser fallback so
+    # unattended runs can't hang on it. Relayed writes no apiKeyHelper and sets no
+    # ANTHROPIC_API_KEY, so this token is the top credential Claude Code resolves.
+    if os.environ.get(CLAUDE_CODE_OAUTH_TOKEN_ENV_VAR):
+        return
     if _has_subscription_login():
         return
     print_note("Opening browser to sign in with your Claude subscription...")

tests/test_agent_claude.py

@@ -1530,3 +1530,45 @@ def test_disable_removes_only_ucode_hooks(self, tmp_path, monkeypatch):
         assert state.get(claude.SMART_ROUTING_STATE_KEY) is None
         assert list(doc["hooks"]) == ["PreToolUse"]
         assert doc["hooks"]["PreToolUse"][0]["hooks"][0]["command"] == "user-policy"
+
+
+class TestEnsureSubscriptionLogin:
+    """Relayed launch's subscription-login gate."""
+
+    @staticmethod
+    def _forbid_subprocess(monkeypatch):
+        """Fail loudly if the CLI is shelled out to at all (status probe or login)."""
+
+        def _boom(*args, **kwargs):
+            raise AssertionError(f"unexpected subprocess call: {args!r}")
+
+        monkeypatch.setattr(claude.subprocess, "run", _boom)
+
+    def test_oauth_token_env_skips_login(self, monkeypatch):
+        # A pre-provisioned CLAUDE_CODE_OAUTH_TOKEN (e.g. `claude setup-token`
+        # output in CI) is the credential Claude Code uses directly, so no
+        # interactive browser login applies — and no `auth status` probe is even
+        # needed. This keeps headless/relayed runs from hanging on the browser.
+        monkeypatch.setenv(claude.CLAUDE_CODE_OAUTH_TOKEN_ENV_VAR, "dummy-oauth-token")
+        self._forbid_subprocess(monkeypatch)
+        claude._ensure_subscription_login()  # returns without touching the CLI
+
+    def test_existing_login_skips_browser(self, monkeypatch):
+        monkeypatch.delenv(claude.CLAUDE_CODE_OAUTH_TOKEN_ENV_VAR, raising=False)
+        monkeypatch.setattr(claude, "_has_subscription_login", lambda: True)
+
+        def _boom(cmd, **kwargs):
+            raise AssertionError(f"no auth login expected, got {cmd!r}")
+
+        monkeypatch.setattr(claude.subprocess, "run", _boom)
+        claude._ensure_subscription_login()
+
+    def test_missing_login_runs_browser_flow(self, monkeypatch):
+        monkeypatch.delenv(claude.CLAUDE_CODE_OAUTH_TOKEN_ENV_VAR, raising=False)
+        monkeypatch.setattr(claude, "_has_subscription_login", lambda: False)
+        calls: list[list[str]] = []
+        monkeypatch.setattr(claude.subprocess, "run", lambda cmd, **kwargs: calls.append(cmd))
+        monkeypatch.setattr(claude, "print_note", lambda *a, **kw: None)
+        monkeypatch.setattr(claude, "print_success", lambda *a, **kw: None)
+        claude._ensure_subscription_login()
+        assert calls == [[claude.SPEC["binary"], "auth", "login"]]

tests/test_e2e.py

@@ -16,6 +16,7 @@
 import shutil
 import subprocess
 import tempfile
+import threading
 from pathlib import Path
 from urllib import error as urllib_error
 from urllib import request as urllib_request
@@ -571,6 +572,20 @@ def _first_service(tool: str, workspace: str, token: str) -> str:
             )
         return names[0]
 
+    @staticmethod
+    def _first_relayed_service(tool: str, workspace: str, token: str) -> str:
+        services, reason = list_model_provider_services(workspace, token)
+        if is_model_provider_feature_unavailable(reason):
+            pytest.skip("Model Provider Service feature not enabled on this workspace")
+        if reason is not None:
+            pytest.skip(f"could not list provider services: {reason}")
+        names = [
+            s["name"] for s in services if service_usable_for_tool(tool, s) and s.get("relayed")
+        ]
+        if not names:
+            pytest.skip(f"no relayed {tool} model provider services available on this workspace")
+        return names[0]
+
     @staticmethod
     def _skip_if_provider_unusable(combined: str, provider: str) -> None:
         # Environmental provider-account conditions, not ucode bugs: the test only proves routing
@@ -623,6 +638,71 @@ def test_launch_claude_through_provider(
             f"stdout={result.stdout[:300]!r} stderr={result.stderr[:300]!r}"
         )
 
+    def test_launch_claude_through_relayed_provider(
+        self, tmp_path, monkeypatch, e2e_state, e2e_workspace, e2e_token
+    ):
+        """Relayed (subscription-relay) launch: Claude Code authenticates to
+        Anthropic with the subscription OAuth token while the loopback proxy
+        injects the Databricks swap token. Headless runs supply the OAuth token
+        via CLAUDE_CODE_OAUTH_TOKEN (`claude setup-token` output); without it the
+        launch would open an interactive browser login, so the test skips. Also
+        needs a relayed MPS on the workspace, so it stays inert until both exist.
+        """
+        import ucode.config_io as config_io_mod
+        from ucode import gateway_proxy
+        from ucode.agents import claude
+
+        _require_binary("claude")
+        if not os.environ.get(claude.CLAUDE_CODE_OAUTH_TOKEN_ENV_VAR):
+            pytest.skip(
+                "set CLAUDE_CODE_OAUTH_TOKEN (from `claude setup-token`) to run the relayed launch"
+            )
+        provider = self._first_relayed_service("claude", e2e_workspace, e2e_token)
+
+        config_dir = tmp_path / "claude_config"
+        config_dir.mkdir()
+        monkeypatch.setattr(config_io_mod, "APP_DIR", tmp_path)
+        monkeypatch.setattr(claude, "CLAUDE_SETTINGS_PATH", config_dir / "settings.json")
+        monkeypatch.setattr(claude, "CLAUDE_BACKUP_PATH", tmp_path / "claude-settings.backup.json")
+        # The proxy mints the Databricks swap token; feed it the e2e bearer rather
+        # than shelling out to the CLI, matching the other launch tests.
+        monkeypatch.setattr(
+            gateway_proxy, "get_databricks_token", lambda ws, profile=None, **kwargs: e2e_token
+        )
+
+        # Start the real loopback refresh proxy exactly as `_launch_relayed` does,
+        # so the request is credential-swapped and relayed like a live session.
+        server, cache, client = gateway_proxy.start_proxy(
+            e2e_workspace,
+            None,
+            0,
+            token_header=gateway_proxy.AI_GATEWAY_TOKEN_HEADER,
+            force_refresh_near_expiry=False,
+        )
+        port = server.server_address[1]
+        threading.Thread(target=server.serve_forever, daemon=True).start()
+        try:
+            state = {**e2e_state, "workspace": e2e_workspace, "relayed_proxy_port": port}
+            with pytest.MonkeyPatch().context() as mp:
+                mp.setattr("ucode.state.save_state", lambda s: None)
+                claude.write_tool_config(state, None, provider=provider, relayed=True)
+            env = {
+                **os.environ,
+                "CLAUDE_CONFIG_DIR": str(config_dir),
+                "ANTHROPIC_BASE_URL": f"http://127.0.0.1:{port}",
+            }
+            result = _run_agent(claude.validate_cmd("claude"), env=env, timeout=90)
+        finally:
+            cache.stop()
+            server.shutdown()
+            client.close()
+        combined = (result.stdout + result.stderr).strip()
+        self._skip_if_provider_unusable(combined, provider)
+        assert result.returncode == 0 and combined, (
+            f"relayed provider={provider} rc={result.returncode} "
+            f"stdout={result.stdout[:300]!r} stderr={result.stderr[:300]!r}"
+        )
+
     def test_launch_codex_through_provider(
         self, tmp_path, monkeypatch, e2e_state, e2e_workspace, e2e_token
     ):