diff --git a/README.md b/README.md index 7dbea331..20500671 100644 --- a/README.md +++ b/README.md @@ -74,7 +74,36 @@ To configure specific tools without the picker, pass a comma-separated list: ug configure --agents claude,codex ``` -Available agent names are `codex`, `claude`, `gemini`, `opencode`, `copilot`, and `pi`. `cursor` is also accepted (MCP-only — it registers Databricks MCP servers but configures no models). +Available agent names are `codex`, `claude`, `gemini`, `opencode`, `copilot`, `pi`, and `hermes`. `cursor` is also accepted (MCP-only — it registers Databricks MCP servers but configures no models). + +### Hermes Agent + +Hermes exposes **Databricks Model Serving** in its native provider picker and +delegates Databricks-specific setup to ucode. Configure it directly without +launching a session with: + +```bash +ucode configure hermes +``` + +For automation, pin the workspace, Databricks CLI profile, model, and isolated +Hermes home explicitly: + +```bash +ucode configure hermes \ + --workspace https://my-workspace.databricks.com \ + --profile DEFAULT \ + --model system.ai.my-model \ + --hermes-home /path/to/hermes-home \ + --output json +``` + +Generated providers use refreshable Databricks CLI credentials; ucode never +writes an OAuth token or PAT into Hermes configuration. Normal use remains +`hermes` and its `/model` picker—no ucode wrapper process is required. + +Current support includes Responses/Codex, Anthropic Messages, native Gemini, +OpenAI-compatible OSS models, and managed MCP servers. Naming agents explicitly is treated as a request for all of them: if any one isn't available on the workspace, the run fails without configuring the others. Add `--skip-unavailable` to configure the available subset instead and skip the rest with a warning: diff --git a/docs/hermes-agent-integration.md b/docs/hermes-agent-integration.md new file mode 100644 index 00000000..aa2fedf1 --- /dev/null +++ b/docs/hermes-agent-integration.md @@ -0,0 +1,156 @@ +# Hermes Agent integration + +## Status and scope + +This document defines the cross-repository contract for configuring Hermes to +use Databricks-governed models. The user-facing provider name is **Databricks +Model Serving**. Runtime traffic uses Unity AI Gateway routes; the display name +does not change that architecture. + +The integration supports route-specific providers for OpenAI Responses, +Anthropic Messages, native Gemini, and OpenAI-compatible OSS chat. It also +registers ucode-managed MCP servers through Hermes's native configuration +interface. + +## Ownership boundary + +ucode owns all Databricks-specific behavior: + +- Databricks CLI profiles, OAuth/PAT authentication, and token refresh +- workspace capability and model discovery +- Gateway route and wire-protocol selection +- Databricks-specific request headers +- managed MCP discovery and registration +- validation, diagnostics, managed-key tracking, and cleanup +- rendering the non-secret patch applied to Hermes configuration + +Hermes owns its configuration schema and mutation semantics. It exposes a +native provider entry and a thin, argv-based setup handoff to ucode. After +setup, Hermes resolves the generated providers through its existing generic +transports. Hermes must not add a Databricks transport, static model catalog, +runtime token-refresh branch, or Databricks-specific auxiliary/delegation path. + +Normal sessions run with `hermes`; no `ucode hermes` wrapper is required. + +## Command contract + +The stable setup entry point is: + +```text +ucode configure hermes [options] +``` + +The command configures Hermes and exits without launching it. It must: + +- use argv semantics throughout and never invoke a shell; +- accept an explicit Hermes home/profile target for automation and isolation; +- prompt for missing workspace, profile, or model inputs in interactive use; +- support fully specified non-interactive use; +- be idempotent and modify only ucode-managed keys; +- support surgical unconfiguration of those keys; +- return nonzero without changing the active provider/model on failure or + cancellation; and +- never expose access tokens, refresh tokens, PATs, or token-helper output. + +Machine-readable completion output is additive/versionable and contains only +non-secret state. The initial shape is: + +```json +{ + "status": "configured", + "agent": "hermes", + "hermes_home": "/resolved/path", + "provider_ids": ["ucode-databricks-codex"], + "default_provider": "ucode-databricks-codex", + "default_model": "system.ai.", + "mcp_servers_configured": [], + "warnings": [] +} +``` + +Hermes relies on the child exit status, then reloads and validates its own +configuration. It does not duplicate Databricks field validation. + +## Configuration transaction + +ucode does not edit Hermes YAML directly and does not import Hermes Python +modules. It invokes Hermes's public `config get`, `config set`, and +`config unset` commands with an explicit `HERMES_HOME`. + +Hermes does not currently expose an atomic multi-key transaction. ucode checks +that new managed paths are absent and that previously managed values still +match their recorded fingerprints before mutation. Provider definitions are +written before switching the active model. Cleanup removes only values whose +fingerprints still match; uncertain ownership preserves configuration rather +than deleting a possible user replacement. A process failure after mutation +starts can therefore leave a partial update, which is surfaced as an error for +explicit reconciliation. + +## Provider representation + +Each wire protocol has a separate generated custom provider so every model is +coupled to the correct Gateway route and Hermes transport: + +- `ucode-databricks-codex` for OpenAI Responses; +- `ucode-databricks-anthropic` for Anthropic Messages; +- `ucode-databricks-gemini` for native Gemini; and +- `ucode-databricks-oss` for OpenAI-compatible OSS chat. + +Hermes retains its internal `provider="custom"` identity while preserving the +selected generated provider separately. Each provider obtains credentials from +a refreshable `ucode auth-token` command; no Databricks token is stored in +Hermes configuration. + +## Capability contract + +| Capability | Release requirement | Owner | Acceptance evidence | +|---|---|---|---| +| Databricks CLI OAuth | Supported | ucode | Refresh succeeds after token cache invalidation/expiry | +| PAT compatibility | Preserve when explicitly selected | ucode | No PAT is written to Hermes config or output | +| Responses route | Supported | ucode config + Hermes generic transport | Streaming text and tool calls pass through `/ai-gateway/codex/v1` | +| Model-service discovery | Supported | ucode | `/model` contains the authorized coding models | +| Profile-scoped config | Supported | ucode + Hermes config commands | Default and named profiles remain isolated | +| Multiplex isolation | Supported | Hermes generic secret/config resolution | A scoped miss cannot borrow another profile's bearer/config | +| Reconfigure/unconfigure | Supported | ucode | Reconfigure is idempotent; cleanup removes only managed keys | +| Model Provider Service header routing | Deferred for Hermes v1 | ucode | No support claim until a production setup path exists | +| Unity Catalog permissions | Required | Gateway | Allowed models succeed and denied models fail clearly | +| Usage/inference tables | Required where enabled | Gateway, validated by ucode/E2E | Request is visible with expected identity and model | +| Rate limits and policies | Required pass-through | Gateway | 429 classification is preserved with no local bypass | +| Anthropic Messages route | Supported | ucode config + Hermes Messages transport | Claude streaming and tools pass | +| Gemini route | Supported | ucode config + Hermes Gemini transport | Gemini streaming and tools pass | +| OSS/OpenAI-compatible route | Supported | ucode config + Hermes chat transport | Supported OSS chat/tools pass | +| Managed MCP servers | Supported | ucode | Configured tools appear in Hermes without stored bearer tokens | +| Smart routing | Explicit decision gate | ucode | Tested support or explicit unsupported documentation | +| MLflow tracing hooks | Explicit decision gate | ucode | No support claim without Hermes-compatible trace evidence | +| Managed workspace manifest | Explicit decision gate | ucode | Decide whether Hermes belongs in the managed agent schema | + +## Failure and security behavior + +- Missing ucode is handled by Hermes with official installation guidance and no + partial configuration. +- Missing Databricks CLI/login, authorization failures, and empty discovery are + diagnosed by ucode. +- Hermes suppresses raw child output on failure; tests use sentinel secrets and + assert their absence from stdout, stderr, exceptions, fixtures, and config. +- No short-lived token or PAT is persisted in Hermes configuration. +- Provider discovery and import have no subprocess, OAuth, network, install, or + configuration side effects. + +## Support boundaries + +The two currently incompatible models, `gpt-oss-20b` and `gpt-oss-120b`, are +excluded from Hermes provider generation only. Their availability to unrelated +ucode consumers is unchanged. Smart routing, MLflow tracing hooks, and managed +workspace manifests remain separate decisions and are not implied by this +integration. + +Databricks workspace authentication, discovery, Gateway routing, policies, and +managed MCP behavior are supported in ucode/Databricks. Generic Hermes session, +transport, and tool behavior remain supported in Hermes. + +Primary references: + +- [Databricks coding-agent integrations](https://docs.databricks.com/aws/en/ai-gateway/coding-agent-integration-model-services) +- [Databricks Model Provider Services](https://docs.databricks.com/aws/en/ai-gateway/coding-agent-integration-model-provider-services) +- [ucode](https://github.com/databricks/ucode) +- [Hermes Agent](https://github.com/NousResearch/hermes-agent) diff --git a/src/ucode/agents/__init__.py b/src/ucode/agents/__init__.py index 232c65cf..feb19d99 100644 --- a/src/ucode/agents/__init__.py +++ b/src/ucode/agents/__init__.py @@ -40,7 +40,7 @@ spinner, ) -from . import claude, codex, copilot, gemini, opencode, pi +from . import claude, codex, copilot, gemini, hermes, opencode, pi from .args import explicit_model_arg_value as explicit_model_arg_value _MODULES = { @@ -50,6 +50,7 @@ "opencode": opencode, "copilot": copilot, "pi": pi, + "hermes": hermes, } TOOL_SPECS: dict[str, ToolSpec] = {name: module.SPEC for name, module in _MODULES.items()} @@ -68,6 +69,7 @@ "opencode": "opencode", "copilot": "copilot", "pi": "pi", + "hermes": "hermes", } DEFAULT_TOOL = "codex" @@ -100,7 +102,8 @@ def normalize_tool(tool: str) -> str: normalized = TOOL_ALIASES.get(tool.strip().lower()) if not normalized: raise RuntimeError( - f"Unsupported tool '{tool}'. Use one of: codex, claude, gemini, opencode, copilot, pi." + f"Unsupported tool '{tool}'. Use one of: " + "codex, claude, gemini, opencode, copilot, pi, hermes." ) return normalized @@ -184,6 +187,18 @@ def install_tool_binary( binary = spec["binary"] package = spec["package"] + if spec.get("install_method") == "external": + if shutil.which(binary): + return True + message = ( + f"{spec['display']} is not installed (`{binary}` was not found on PATH). " + "Install Hermes using its official installer, then retry." + ) + if strict: + raise RuntimeError(message) + print_warning(message) + return False + if shutil.which(binary): # A too-new build is a correctness blocker (the tool runs but misbehaves # against the gateway), so check it on every launch — not just when @@ -236,6 +251,11 @@ def ensure_tool_binary_available(tool: str) -> None: binary = spec["binary"] if shutil.which(binary): return + if spec.get("install_method") == "external": + raise RuntimeError( + f"{spec['display']} is not installed (`{binary}` was not found on PATH). " + "Install Hermes using its official installer, then retry." + ) raise RuntimeError( f"{spec['display']} is not installed (`{binary}` was not found on PATH). " f"Install it with `npm install -g {spec['package']}` or run " @@ -432,13 +452,16 @@ def configure_tool( # Every tool in this branch needs a model — including gemini under a provider, # which still pins the service's target model in the URL. if not model: - raise RuntimeError(f"A {tool} model must be selected before configuration.") + display = TOOL_SPECS[tool]["display"] + raise RuntimeError(f"A {display} model must be selected before configuration.") if tool == "gemini": result = gemini.write_tool_config(state, model, provider=provider) elif tool == "copilot": result = copilot.write_tool_config(state, model) elif tool == "pi": result = pi.write_tool_config(state, model) + elif tool == "hermes": + result = hermes.write_tool_config(state, model) else: result = opencode.write_tool_config(state, model) # gemini/opencode/copilot/pi return (state, token); codex/claude return state @@ -459,6 +482,11 @@ def check_gateway_endpoint(state: dict, tool: str) -> bool: return bool(state.get("opencode_models")) if tool == "codex": return bool(state.get("codex_models")) + if tool == "hermes": + return any( + bool(state.get(key)) + for key in ("codex_models", "claude_models", "gemini_models", "oss_models") + ) if tool == "gemini": return bool(state.get("gemini_models")) if tool == "copilot": @@ -479,6 +507,7 @@ def check_gateway_endpoint(state: dict, tool: str) -> bool: "gemini": ("gemini",), "copilot": ("claude", "codex"), "pi": ("claude", "codex", "gemini"), + "hermes": ("codex", "claude", "gemini", "oss"), } diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index 01b3a1f9..ca8bd086 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -1353,6 +1353,7 @@ def _launch_relayed(state: dict, binary: str, tool_args: list[str]) -> None: port, token_header=gateway_proxy.AI_GATEWAY_TOKEN_HEADER, force_refresh_near_expiry=False, + allow_env_bearer=bool(state.get("use_pat")), ) # start_proxy falls back to an OS-assigned port when the cached one is taken # (stale proxy from a killed session). Reconcile settings + state to whatever diff --git a/src/ucode/agents/hermes.py b/src/ucode/agents/hermes.py new file mode 100644 index 00000000..528b3f80 --- /dev/null +++ b/src/ucode/agents/hermes.py @@ -0,0 +1,658 @@ +"""Hermes adapter: render Databricks Gateway config and apply it through Hermes.""" + +from __future__ import annotations + +import copy +import hashlib +import json +import math +import os +import re +import subprocess +from pathlib import Path +from typing import cast + +from ucode.config_io import APP_DIR, ToolSpec, is_dry_run +from ucode.databricks import build_auth_shell_command, build_tool_base_url +from ucode.launcher import exec_or_spawn +from ucode.ui import console + +PROVIDER_ID = "ucode-databricks-codex" +ANTHROPIC_PROVIDER_ID = "ucode-databricks-anthropic" +OSS_PROVIDER_ID = "ucode-databricks-oss" +GEMINI_PROVIDER_ID = "ucode-databricks-gemini" +MANAGED_PROVIDER_IDS = ( + PROVIDER_ID, + ANTHROPIC_PROVIDER_ID, + OSS_PROVIDER_ID, + GEMINI_PROVIDER_ID, +) +DISPLAY_NAME = "Databricks Model Serving" +UNSUPPORTED_MODELS = frozenset( + { + "system.ai.gpt-oss-20b", + "system.ai.gpt-oss-120b", + } +) + + +class HermesConfigApplyError(RuntimeError): + """Hermes configuration may have changed before the operation failed.""" + + +SPEC: ToolSpec = { + "binary": "hermes", + "package": "", + "display": "Hermes", + "config_path": Path.home() / ".hermes" / "config.yaml", + "backup_path": APP_DIR / "hermes-config.backup.yaml", + "install_method": "external", +} + +MANAGED_PATHS = ( + f"providers.{PROVIDER_ID}", + f"providers.{ANTHROPIC_PROVIDER_ID}", + f"providers.{OSS_PROVIDER_ID}", + f"providers.{GEMINI_PROVIDER_ID}", + "model.provider", + "model.default", +) + + +def default_model(state: dict) -> str | None: + """Choose the first model from Hermes' supported protocol families.""" + responses_model = state.get("codex_default_model") + if not isinstance(responses_model, str): + codex_models = state.get("codex_models") + responses_model = ( + max( + (model for model in codex_models if isinstance(model, str)), + key=_gpt_model_version, + default=None, + ) + if isinstance(codex_models, list) + else None + ) + if responses_model and responses_model not in UNSUPPORTED_MODELS: + return responses_model + claude_models = state.get("claude_models") + if isinstance(claude_models, dict): + for model in claude_models.values(): + if isinstance(model, str) and model and model not in UNSUPPORTED_MODELS: + return model + for key in ("gemini_models", "oss_models"): + models = state.get(key) + if isinstance(models, list): + for model in models: + if isinstance(model, str) and model and model not in UNSUPPORTED_MODELS: + return model + return None + + +def _gpt_model_version(model: str) -> tuple[int, int, int, int]: + """Order discovered GPT Responses models without changing Codex defaults.""" + match = re.fullmatch(r"(?:system\.ai\.)?gpt-(\d+)(?:-(\d+))?(?:-(\d+))?(-.*)?", model) + if not match: + return (0, 0, 0, 0) + major, minor, patch, suffix = match.groups() + return int(major), int(minor or 0), int(patch or 0), 1 if suffix is None else 0 + + +def state_scoped_to_home(state: dict, hermes_home: str | Path) -> dict: + """Hide ownership established for a different Hermes profile.""" + managed = state.get("managed_configs") + hermes_managed = managed.get("hermes") if isinstance(managed, dict) else None + recorded_home = hermes_managed.get("hermes_home") if isinstance(hermes_managed, dict) else None + target = str(Path(hermes_home).expanduser().resolve()) + if isinstance(recorded_home, str) and str(Path(recorded_home).expanduser().resolve()) == target: + return state + scoped = copy.deepcopy(state) + scoped_managed = scoped.get("managed_configs") + if isinstance(scoped_managed, dict): + scoped_managed.pop("hermes", None) + return scoped + + +def is_update_available() -> None: + """Hermes is installed externally; ucode never manages its version.""" + return None + + +def render_config_patch( + state: dict, + model: str | None = None, +) -> dict: + """Return the non-secret, Hermes-owned multi-protocol transaction patch. + + Each model family is isolated behind its matching wire transport. + """ + workspace = state["workspace"] + claude_models = state.get("claude_models") + anthropic_models = [ + candidate + for candidate in _unique_models( + claude_models.values() if isinstance(claude_models, dict) else None + ) + if candidate not in UNSUPPORTED_MODELS + ] + anthropic_set = set(anthropic_models) + oss_models = [ + candidate + for candidate in _unique_models(state.get("oss_models")) + if candidate not in anthropic_set and candidate not in UNSUPPORTED_MODELS + ] + specific_models = anthropic_set | set(oss_models) + gemini_models = [ + model + for model in _unique_models(state.get("gemini_models")) + if model not in specific_models and model not in UNSUPPORTED_MODELS + ] + specific_models.update(gemini_models) + responses_models = [ + candidate + for candidate in _unique_models(state.get("codex_models")) + if candidate not in specific_models and candidate not in UNSUPPORTED_MODELS + ] + if model in UNSUPPORTED_MODELS: + raise RuntimeError(f"Model {model!r} is not supported by Hermes.") + selected = model or default_model(state) + if selected in UNSUPPORTED_MODELS: + selected = None + selected = selected or next( + ( + candidate + for candidates in ( + responses_models, + anthropic_models, + gemini_models, + oss_models, + ) + for candidate in candidates + ), + None, + ) + if not selected: + raise RuntimeError("No supported models are available for Hermes.") + + provider_specs = [ + (PROVIDER_ID, build_tool_base_url("codex", workspace), "codex_responses", responses_models), + ( + ANTHROPIC_PROVIDER_ID, + build_tool_base_url("claude", workspace), + "anthropic_messages", + anthropic_models, + ), + (OSS_PROVIDER_ID, f"{workspace}/ai-gateway/mlflow/v1", "openai_chat", oss_models), + ( + GEMINI_PROVIDER_ID, + f"{build_tool_base_url('gemini', workspace)}/v1beta", + "gemini-native", + gemini_models, + ), + ] + selected_provider = PROVIDER_ID + for provider_id, _, _, compatible_models in provider_specs[1:]: + if selected in compatible_models: + selected_provider = provider_id + break + if selected_provider == PROVIDER_ID and selected not in responses_models: + # Preserve Phase 1's explicit-model behavior without guessing a route + # from the model name. Only discovered family membership may switch + # the wire protocol. + responses_models.append(selected) + + key_cmd = build_auth_shell_command( + workspace, + state.get("profile"), + use_pat=bool(state.get("use_pat")), + ) + provider_paths = {} + for provider_id, api, transport, compatible_models in provider_specs: + if not compatible_models: + continue + default = selected if provider_id == selected_provider else compatible_models[0] + provider_config: dict = { + "name": DISPLAY_NAME, + "api": api, + "transport": transport, + "key_cmd": key_cmd, + "default_model": default, + "models": {model_id: {} for model_id in compatible_models}, + "discover_models": False, + } + provider_paths[f"providers.{provider_id}"] = provider_config + + managed = state.get("managed_configs") + hermes_managed = managed.get("hermes") if isinstance(managed, dict) else None + raw_fingerprints = ( + hermes_managed.get("provider_fingerprints") if isinstance(hermes_managed, dict) else None + ) + owned_fingerprints: dict = raw_fingerprints if isinstance(raw_fingerprints, dict) else {} + provider_expect_hashes = { + path: owned_fingerprints[provider_id] + for path in provider_paths + if (provider_id := path.removeprefix("providers.")) in owned_fingerprints + } + owned_model = hermes_managed.get("active_model") if isinstance(hermes_managed, dict) else None + if ( + isinstance(owned_model, dict) + and isinstance(owned_model.get("provider"), str) + and isinstance(owned_model.get("default"), str) + ): + provider_expect_hashes.update( + { + "model.provider": config_value_fingerprint(owned_model["provider"]), + "model.default": config_value_fingerprint(owned_model["default"]), + } + ) + provider_expect_missing = [ + path for path in provider_paths if path not in provider_expect_hashes + ] + missing_provider_paths = [path for path in MANAGED_PATHS[:4] if path not in provider_paths] + return { + "set": { + **provider_paths, + "model.provider": selected_provider, + "model.default": selected, + }, + "unset": [], + "expect_missing": provider_expect_missing, + "expect_hashes": provider_expect_hashes, + "unset_if_hash": { + path: owned_fingerprints[path.removeprefix("providers.")] + for path in missing_provider_paths + if path.removeprefix("providers.") in owned_fingerprints + }, + } + + +def _unique_models(models) -> list[str]: + """Preserve Gateway discovery order while dropping invalid duplicates.""" + return list(dict.fromkeys(model for model in models or [] if isinstance(model, str) and model)) + + +def render_unconfigure_patch( + *, + current_model: dict, + owned_model: dict, + owned_provider_fingerprints: dict, + current_provider_fingerprint: str | None, +) -> dict: + """Remove managed paths without clobbering a later user selection.""" + current_provider = current_model.get("provider") + owned_pair_matches = current_model == owned_model + owned_provider_fingerprint = owned_provider_fingerprints.get(current_provider) + provider_still_owned = ( + isinstance(owned_provider_fingerprint, str) + and current_provider_fingerprint == owned_provider_fingerprint + ) + provider_paths = list(MANAGED_PATHS[:4]) + if owned_pair_matches and provider_still_owned: + active_provider_path = f"providers.{current_provider}" + unset = [active_provider_path, *MANAGED_PATHS[4:]] + expect_hashes = { + active_provider_path: owned_provider_fingerprint, + "model.provider": config_value_fingerprint(owned_model["provider"]), + "model.default": config_value_fingerprint(owned_model["default"]), + } + provider_paths = [path for path in provider_paths if path != active_provider_path] + elif isinstance(current_provider, str): + active_provider_path = f"providers.{current_provider}" + provider_paths = [path for path in provider_paths if path != active_provider_path] + unset = [] + expect_hashes = {} + else: + unset = [] + expect_hashes = {} + return { + "set": {}, + "unset": unset, + "expect_hashes": expect_hashes, + "unset_if_hash": { + path: owned_provider_fingerprints[path.removeprefix("providers.")] + for path in provider_paths + if path.removeprefix("providers.") in owned_provider_fingerprints + }, + } + + +def render_mcp_server_patch( + name: str, + argv: list[str], + *, + expected_fingerprint: str | None = None, +) -> dict: + """Render one Hermes stdio MCP entry backed by ucode's refreshable proxy.""" + if not argv: + raise ValueError("MCP proxy argv must not be empty.") + path = f"mcp_servers.{name}" + patch: dict[str, object] = { + "set": { + path: { + "command": argv[0], + "args": list(argv[1:]), + } + }, + "unset": [], + } + if expected_fingerprint is None: + patch["expect_missing"] = [path] + else: + patch["expect_hashes"] = {path: expected_fingerprint} + return patch + + +def config_value_fingerprint(value: object) -> str: + """Hash one canonical Hermes config value without persisting its contents.""" + + def validate(item: object) -> None: + if isinstance(item, dict): + if not all(isinstance(key, str) for key in item): + raise ValueError("Config fingerprints require string mapping keys") + for nested in item.values(): + validate(nested) + return + if isinstance(item, list): + for nested in item: + validate(nested) + return + if isinstance(item, float) and not math.isfinite(item): + raise ValueError("Config fingerprints require finite numbers") + if item is not None and not isinstance(item, (str, int, float, bool)): + raise ValueError("Unsupported config fingerprint value") + + validate(value) + payload = json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=False) + return hashlib.sha256(payload.encode("utf-8")).hexdigest() + + +def mcp_server_fingerprint(value: dict) -> str: + return config_value_fingerprint(value) + + +def mcp_value_for_argv(argv: list[str]) -> dict: + return {"command": argv[0], "args": argv[1:]} + + +def apply_config_patch(patch: dict, *, hermes_home: str | Path) -> dict: + """Apply a conservative patch through Hermes's public config commands. + + Hermes does not currently expose a multi-key transaction. Preconditions + therefore prevent known collisions but cannot eliminate a concurrent + read-to-write race. On uncertainty, cleanup leaves the value in place. + """ + env = os.environ.copy() + env["HERMES_HOME"] = str(Path(hermes_home).expanduser().resolve()) + + def run(arguments: list[str]) -> subprocess.CompletedProcess: + argv = [SPEC["binary"], "config", *arguments] + try: + return subprocess.run( + argv, + capture_output=True, + text=True, + check=False, + shell=False, + env=env, + ) + except OSError as exc: + raise RuntimeError("Hermes configuration could not be started.") from exc + + def read_if_set(path: str) -> tuple[bool, object | None]: + completed = run(["get", path, "--json"]) + if completed.returncode != 0: + if "Config key not set:" in (completed.stderr or "") or completed.stderr == "missing": + return False, None + raise RuntimeError(f"Hermes configuration read failed for {path!r}.") + try: + return True, json.loads(completed.stdout) + except (TypeError, json.JSONDecodeError) as exc: + raise RuntimeError(f"Hermes returned invalid configuration for {path!r}.") from exc + + configured = patch.get("set") or {} + unconditional_unsets = patch.get("unset") or [] + expect_missing = patch.get("expect_missing") or [] + expect_hashes = patch.get("expect_hashes") or {} + unset_if_hash = patch.get("unset_if_hash") or {} + if not isinstance(configured, dict) or not isinstance(unconditional_unsets, list): + raise RuntimeError("Invalid Hermes configuration patch.") + + if is_dry_run(): + plan = { + "set": list(configured), + "unset": list(dict.fromkeys([*unconditional_unsets, *unset_if_hash])), + } + console.print( + f"\n[bold]\\[dry run] Hermes config operations[/bold]\n{json.dumps(plan, indent=2)}" + ) + return {"status": "dry_run", "paths": plan} + + for path in expect_missing: + found, _value = read_if_set(path) + if found: + raise RuntimeError(f"Hermes configuration value {path!r} already exists.") + for path, expected in expect_hashes.items(): + found, value = read_if_set(path) + if not found or config_value_fingerprint(value) != expected: + raise RuntimeError(f"Hermes configuration value {path!r} changed; refusing update.") + + conditional_unsets = [] + for path, expected in unset_if_hash.items(): + found, value = read_if_set(path) + if found and config_value_fingerprint(value) == expected: + conditional_unsets.append(path) + + model_paths = {"model.provider", "model.default"} + provider_sets = [(path, value) for path, value in configured.items() if path not in model_paths] + model_sets = [ + (path, configured[path]) + for path in ("model.provider", "model.default") + if path in configured + ] + applied_sets = [] + applied_unsets = [] + + def cli_value(value: object) -> str: + if isinstance(value, str): + return value + return json.dumps(value, separators=(",", ":"), ensure_ascii=False) + + mutation_started = False + try: + for path, value in provider_sets: + completed = run(["set", path, cli_value(value), "--force"]) + mutation_started = True + if completed.returncode != 0: + raise RuntimeError(f"Hermes configuration write failed for {path!r}.") + applied_sets.append(path) + + for path in dict.fromkeys([*unconditional_unsets, *conditional_unsets]): + found, current_value = read_if_set(path) + if not found: + continue + expected = expect_hashes.get(path) or unset_if_hash.get(path) + if expected is not None and config_value_fingerprint(current_value) != expected: + continue + completed = run(["unset", path]) + mutation_started = True + if completed.returncode != 0: + raise RuntimeError(f"Hermes configuration cleanup failed for {path!r}.") + applied_unsets.append(path) + + for path, value in model_sets: + completed = run(["set", path, cli_value(value), "--force"]) + mutation_started = True + if completed.returncode != 0: + raise RuntimeError(f"Hermes configuration write failed for {path!r}.") + applied_sets.append(path) + except Exception as exc: + if mutation_started: + raise HermesConfigApplyError(str(exc)) from exc + raise + + return { + "status": "applied", + "paths": {"set": applied_sets, "unset": applied_unsets}, + } + + +def read_config_value(dotted_key: str, *, hermes_home: str | Path) -> object: + """Read one value from an explicit Hermes profile.""" + target = Path(hermes_home).expanduser().resolve() + env = dict(os.environ) + env["HERMES_HOME"] = str(target) + try: + completed = subprocess.run( + [SPEC["binary"], "config", "get", dotted_key, "--json"], + check=False, + capture_output=True, + text=True, + shell=False, + env=env, + ) + except OSError as exc: + raise RuntimeError("Hermes configuration could not be read.") from exc + if completed.returncode != 0: + raise RuntimeError( + f"Hermes configuration read failed with exit code {completed.returncode}." + ) + try: + return json.loads(completed.stdout) + except (TypeError, json.JSONDecodeError) as exc: + raise RuntimeError("Hermes returned invalid configuration.") from exc + + +def read_active_model(*, hermes_home: str | Path) -> dict: + """Read the remaining active-model fields from one explicit profile.""" + model = read_config_value("model", hermes_home=hermes_home) + if not isinstance(model, dict): + raise RuntimeError("Hermes returned invalid model configuration.") + model_values = cast(dict[str, object], model) + return { + key: value + for key in ("provider", "default") + if isinstance((value := model_values.get(key)), str) + } + + +def write_tool_config( + state: dict, + model: str | None = None, + *, + hermes_home: str | Path | None = None, +) -> dict: + """Render and apply Hermes configuration, then record cleanup ownership.""" + target = Path(hermes_home or os.environ.get("HERMES_HOME") or Path.home() / ".hermes") + patch = render_config_patch( + state_scoped_to_home(state, target), + model=model, + ) + previous_state = copy.deepcopy(state) + pending_state = record_config_ownership(copy.deepcopy(state), patch, target) + from ucode.state import save_state + + save_state(pending_state) + try: + apply_config_patch(patch, hermes_home=target) + except Exception as exc: + if not isinstance(exc, HermesConfigApplyError): + try: + save_state(previous_state) + except Exception: + pass + raise + return pending_state + + +def record_config_ownership(state: dict, patch: dict, hermes_home: str | Path) -> dict: + """Record only the non-secret values needed for surgical cleanup.""" + from ucode.state import mark_tool_managed + + configured = patch["set"] + provider_fingerprints = { + path.removeprefix("providers."): config_value_fingerprint(value) + for path, value in configured.items() + if path.startswith("providers.") + } + return mark_tool_managed( + state, + "hermes", + [], + metadata={ + "hermes_home": str(Path(hermes_home).expanduser().resolve()), + "active_model": { + "provider": configured["model.provider"], + "default": configured["model.default"], + }, + "provider_fingerprints": provider_fingerprints, + }, + ) + + +def unconfigure( + *, + hermes_home: str | Path | None = None, + owned_model: dict, + owned_provider_fingerprints: dict, +) -> dict: + target = Path(hermes_home or os.environ.get("HERMES_HOME") or Path.home() / ".hermes") + current_model = read_active_model(hermes_home=target) + current_provider = current_model.get("provider") + current_provider_fingerprint = None + if isinstance(current_provider, str) and current_provider in owned_provider_fingerprints: + current_provider_value = read_config_value( + f"providers.{current_provider}", + hermes_home=target, + ) + current_provider_fingerprint = config_value_fingerprint(current_provider_value) + patch = render_unconfigure_patch( + current_model=current_model, + owned_model=owned_model, + owned_provider_fingerprints=owned_provider_fingerprints, + current_provider_fingerprint=current_provider_fingerprint, + ) + return apply_config_patch(patch, hermes_home=target) + + +def write_mcp_server_config( + name: str, + argv: list[str], + *, + hermes_home: str | Path | None = None, + expected_fingerprint: str | None = None, +) -> bool: + """Idempotently upsert one managed server in the active Hermes profile.""" + target = hermes_home or os.environ.get("HERMES_HOME") or Path.home() / ".hermes" + apply_config_patch( + render_mcp_server_patch(name, argv, expected_fingerprint=expected_fingerprint), + hermes_home=target, + ) + return True + + +def remove_mcp_server_config( + name: str, + *, + hermes_home: str | Path | None = None, + expected_fingerprint: str | None = None, +) -> bool: + """Remove one server only while its exact managed value still matches.""" + if not isinstance(expected_fingerprint, str) or not expected_fingerprint: + return False + target = Path(hermes_home or os.environ.get("HERMES_HOME") or Path.home() / ".hermes") + path = f"mcp_servers.{name}" + receipt = apply_config_patch( + {"set": {}, "unset": [], "unset_if_hash": {path: expected_fingerprint}}, + hermes_home=target, + ) + return path in receipt.get("paths", {}).get("unset", []) + + +def launch(state: dict, tool_args: list[str]) -> None: + """Launch ordinary Hermes; configuration does not require a wrapper.""" + del state + exec_or_spawn([SPEC["binary"], *tool_args]) + + +def validate_cmd(binary: str) -> list[str]: + return [binary, "--version"] diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 60d4f1cb..e90357a5 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -3,12 +3,16 @@ from __future__ import annotations +import copy +import json import os import shutil import subprocess +import tempfile from collections.abc import Iterator from contextlib import contextmanager from importlib import metadata +from pathlib import Path from typing import Annotated import typer @@ -35,6 +39,7 @@ ) from ucode.agents import claude as claude_agent from ucode.agents import codex as codex_agent +from ucode.agents import hermes as hermes_agent from ucode.agents import ( launch as launch_agent, ) @@ -520,6 +525,7 @@ def configure_shared_state( skip_preflight: bool = False, fable_enabled: bool | None = None, databricks_ai_tools_enabled: bool | None = None, + allow_env_bearer: bool = True, ) -> dict: """Log into Databricks, verify AI Gateway, fetch model lists, persist state. @@ -622,16 +628,23 @@ def configure_shared_state( f"entry under [{profile}], or re-run without --use-pat to use OAuth." ) # Export the PAT for this process and launched agent subprocesses so - # every token fetch takes the static-bearer path. ensure_pat_bearer - # keeps a non-empty pre-set bearer (CI escape hatch) but treats an - # empty one as absent, so it never shadows the PAT. Pass the validated - # token to avoid re-reading ~/.databrickscfg. - ensure_pat_bearer(profile, pat) - ensure_databricks_auth(workspace, profile) + # explicit PAT-enabled token fetches can take the static-bearer path. + # Pass the validated token so this scoped identity replaces ambient + # state and avoids re-reading ~/.databrickscfg. + ensure_pat_bearer(profile, pat, workspace=workspace) + ensure_databricks_auth( + workspace, + profile, + allow_env_bearer=allow_env_bearer, + ) elif force_login: run_databricks_login(workspace, profile) else: - ensure_databricks_auth(workspace, profile) + ensure_databricks_auth( + workspace, + profile, + allow_env_bearer=allow_env_bearer, + ) # After login the profile exists in ~/.databrickscfg, so a host->profile # lookup is reliable even when it returned nothing above. if profile is None: @@ -639,7 +652,11 @@ def configure_shared_state( if profile: state["profile"] = profile with spinner("Verifying Unity AI Gateway..."): - token = get_databricks_token(workspace, profile) + token = get_databricks_token( + workspace, + profile, + allow_env_bearer=allow_env_bearer, + ) model_service_probe = probe_unity_gateway_capabilities(workspace, token) if model_service_probe.resource_available: print_success("Unity AI Gateway connected") @@ -647,13 +664,22 @@ def configure_shared_state( print_warning(f"Model service: {model_service_probe.detail}") want_claude = ( - fetch_all or "claude" in tools or "opencode" in tools or "copilot" in tools or "pi" in tools + fetch_all + or "claude" in tools + or "hermes" in tools + or "opencode" in tools + or "copilot" in tools + or "pi" in tools + ) + want_gemini = ( + fetch_all or "gemini" in tools or "hermes" in tools or "opencode" in tools or "pi" in tools + ) + want_codex = ( + fetch_all or "codex" in tools or "hermes" in tools or "copilot" in tools or "pi" in tools ) - want_gemini = fetch_all or "gemini" in tools or "opencode" in tools or "pi" in tools - want_codex = fetch_all or "codex" in tools or "copilot" in tools or "pi" in tools # Codex smart routing can select OSS models such as GLM, so a Codex-only # configure must persist that discovered family too. - want_oss = fetch_all or "opencode" in tools or "codex" in tools + want_oss = fetch_all or "hermes" in tools or "opencode" in tools or "codex" in tools claude_reason: str | None = None gemini_reason: str | None = None @@ -756,6 +782,7 @@ def _configure_shared_workspace_states( ) -> list[dict]: if not workspaces: raise RuntimeError("At least one workspace must be provided.") + allow_env_bearer = tools is None or "hermes" not in tools states: list[dict] = [] for workspace, profile in workspaces: states.append( @@ -767,6 +794,7 @@ def _configure_shared_workspace_states( use_pat=use_pat, fable_enabled=fable_enabled, databricks_ai_tools_enabled=databricks_ai_tools_enabled, + allow_env_bearer=allow_env_bearer, ) ) return states @@ -1119,12 +1147,35 @@ def revert() -> int: claude_managed_result = claude_agent.revert_managed_settings() codex_managed_result = codex_agent.revert_managed_config() + hermes_managed = managed_configs.get("hermes") + hermes_configured = "hermes" in (state.get("available_tools") or []) + hermes_cleaned = False + if hermes_managed or hermes_configured: + if not isinstance(hermes_managed, dict): + raise RuntimeError("Hermes cleanup ownership is missing; refusing unsafe cleanup.") + hermes_home = hermes_managed.get("hermes_home") + active_model = hermes_managed.get("active_model") + provider_fingerprints = hermes_managed.get("provider_fingerprints") + if not isinstance(hermes_home, str) or not isinstance(active_model, dict): + raise RuntimeError("Hermes ownership metadata is incomplete; refusing unsafe cleanup.") + if not isinstance(provider_fingerprints, dict) or not provider_fingerprints: + raise RuntimeError("Hermes provider ownership is incomplete; refusing unsafe cleanup.") + hermes_agent.unconfigure( + hermes_home=hermes_home, + owned_model=active_model, + owned_provider_fingerprints=provider_fingerprints, + ) + hermes_cleaned = True + results: dict[str, bool] = { tool: restore_file( spec["config_path"], spec["backup_path"], bool(managed_configs.get(tool)) ) for tool, spec in TOOL_SPECS.items() + if tool != "hermes" } + if "hermes" in TOOL_SPECS: + results["hermes"] = hermes_cleaned pi_settings_restored = restore_file( PI_SETTINGS_PATH, PI_SETTINGS_BACKUP_PATH, bool(managed_configs.get("pi")) ) @@ -1491,28 +1542,40 @@ def auth_token_cmd( bool, typer.Option("--force-refresh", help="Force the Databricks CLI to mint a new token."), ] = False, + allow_env_bearer: Annotated[ + bool, + typer.Option( + "--allow-env-bearer", + help="Allow DATABRICKS_BEARER for this explicit identity (CI/M2M only).", + hidden=True, + ), + ] = False, ) -> None: """Print a Databricks bearer token to stdout, then exit. This is the cross-platform helper invoked by Claude Code's `apiKeyHelper` and Codex's auth command on every token refresh. It is not meant for - interactive use. All token logic (DATABRICKS_BEARER short-circuit, PAT - profiles, OAuth refresh) lives in `get_databricks_token`, so the same - binary works on macOS, Linux, and Windows without any POSIX shell.""" + interactive use. Explicit PAT profiles are resolved directly; OAuth and an + explicitly allowed DATABRICKS_BEARER are handled by `get_databricks_token`, + so the same binary works on macOS, Linux, and Windows without a POSIX shell.""" import sys + explicit_identity = host is not None or profile is not None state = load_state() workspace = host or state.get("workspace") if not workspace: print_err("No workspace configured. Run `ug configure` first.") raise typer.Exit(1) - profile = profile or state.get("profile") - if use_pat or state.get("use_pat"): + if profile is None and host is None: + profile = state.get("profile") + effective_use_pat = use_pat or (not explicit_identity and bool(state.get("use_pat"))) + if effective_use_pat: # --use-pat explicitly means "serve the profile's static PAT". Fail # closed if it can't be read rather than falling through to OAuth — # `auth token` cannot serve a PAT-only profile, so that path would # surface a misleading stale-login error instead of the real cause. - if not ensure_pat_bearer(profile): + token = resolve_pat_token(profile) + if not token: print_err( f"--use-pat: no personal access token available for profile " f"'{profile or ''}'. Add a `token = ` entry under " @@ -1520,11 +1583,17 @@ def auth_token_cmd( "`ug configure` without --use-pat to use OAuth." ) raise typer.Exit(1) - try: - token = get_databricks_token(workspace, profile, force_refresh=force_refresh) - except RuntimeError as exc: - print_err(str(exc)) - raise typer.Exit(1) from None + else: + try: + token = get_databricks_token( + workspace, + profile, + force_refresh=force_refresh, + allow_env_bearer=allow_env_bearer, + ) + except RuntimeError as exc: + print_err(str(exc)) + raise typer.Exit(1) from None # Write the bare token (with trailing newline) to stdout — nothing else may # land on stdout or the consuming agent will treat it as part of the token. sys.stdout.write(token + "\n") @@ -2713,7 +2782,7 @@ def configure( str | None, typer.Option( "--agent", - help="Configure only the named agent (e.g. claude, codex, gemini, opencode, copilot, pi).", + help="Configure only the named agent (e.g. claude, codex, gemini, opencode, copilot, pi, hermes).", ), ] = None, agents: Annotated[ @@ -3039,6 +3108,209 @@ def configure( raise typer.Exit(130) from None +def _prepare_hermes_receipt(receipt_file: Path, receipt: dict) -> Path: + """Durably stage a receipt without publishing its success marker.""" + staged_path: Path | None = None + try: + receipt_file.parent.mkdir(parents=True, exist_ok=True) + with tempfile.NamedTemporaryFile( + mode="w", + encoding="utf-8", + dir=receipt_file.parent, + prefix=f".{receipt_file.name}.", + suffix=".tmp", + delete=False, + ) as handle: + staged_path = Path(handle.name) + handle.write(json.dumps(receipt, sort_keys=True) + "\n") + handle.flush() + os.fsync(handle.fileno()) + return staged_path + except OSError as exc: + if staged_path is not None: + try: + staged_path.unlink(missing_ok=True) + except OSError: + pass + raise RuntimeError("Could not prepare the Hermes configuration receipt.") from exc + + +def _discard_staged_receipt(staged_path: Path | None) -> None: + if staged_path is None: + return + try: + staged_path.unlink(missing_ok=True) + except OSError: + pass + + +def configure_hermes_command( + *, + workspace: str | None = None, + profile: str | None = None, + model: str | None = None, + hermes_home: str | Path | None = None, + receipt_file: Path | None = None, +) -> dict: + """Configure Hermes through its public config commands and return a safe receipt.""" + if profile and not workspace: + raise RuntimeError("--profile requires --workspace.") + interactive = workspace is None + if workspace is None: + workspace, prompted_profile = _prompt_for_configuration("hermes") + profile = profile or prompted_profile + + state = configure_shared_state( + workspace, + profile=profile, + tools=["hermes"], + force_login=interactive, + allow_env_bearer=False, + ) + selected = model or hermes_agent.default_model(state) + if not selected: + raise RuntimeError("No supported models are available for Hermes.") + claude_models = state.get("claude_models") + available = [ + *(state.get("codex_models") or []), + *(claude_models.values() if isinstance(claude_models, dict) else []), + *(state.get("oss_models") or []), + *(state.get("gemini_models") or []), + ] + if model is not None and model not in available: + raise RuntimeError( + f"Hermes model '{model}' is not available on workspace {state['workspace']}." + ) + + target = ( + Path(hermes_home or os.environ.get("HERMES_HOME") or Path.home() / ".hermes") + .expanduser() + .resolve() + ) + patch = hermes_agent.render_config_patch( + hermes_agent.state_scoped_to_home(state, target), + model=selected, + ) + provider_ids = [ + path.removeprefix("providers.") for path in patch["set"] if path.startswith("providers.") + ] + receipt = { + "status": "configured", + "agent": "hermes", + "hermes_home": str(target), + "provider_ids": provider_ids, + "default_provider": patch["set"]["model.provider"], + "default_model": selected, + "mcp_servers_configured": [], + "warnings": [], + } + staged_receipt = _prepare_hermes_receipt(receipt_file, receipt) if receipt_file else None + previous_state = copy.deepcopy(state) + pending_state = hermes_agent.record_config_ownership(copy.deepcopy(state), patch, target) + pending_state["available_tools"] = sorted( + set(pending_state.get("available_tools") or []) | {"hermes"} + ) + try: + save_state(pending_state) + except Exception: + _discard_staged_receipt(staged_receipt) + raise + try: + hermes_agent.apply_config_patch(patch, hermes_home=target) + except Exception as exc: + if not isinstance(exc, hermes_agent.HermesConfigApplyError): + try: + save_state(previous_state) + except Exception: + pass + _discard_staged_receipt(staged_receipt) + raise + if receipt_file is not None and staged_receipt is not None: + try: + os.replace(staged_receipt, receipt_file) + except OSError as exc: + _discard_staged_receipt(staged_receipt) + raise RuntimeError("Could not publish the Hermes configuration receipt.") from exc + return receipt + + +@configure_app.command("hermes") +def configure_hermes( + workspace: Annotated[ + str | None, + typer.Option("--workspace", help="Databricks workspace URL (skips the workspace prompt)."), + ] = None, + profile: Annotated[ + str | None, + typer.Option("--profile", help="Exact Databricks CLI profile to use with --workspace."), + ] = None, + model: Annotated[ + str | None, + typer.Option("--model", help="Exact discovered Responses model to make the default."), + ] = None, + hermes_home: Annotated[ + Path | None, + typer.Option( + "--hermes-home", + help="Exact Hermes home to configure (defaults to HERMES_HOME or ~/.hermes).", + ), + ] = None, + output: Annotated[ + str, + typer.Option("--output", help="Completion output: text (default) or json."), + ] = "text", + receipt_file: Annotated[ + Path | None, + typer.Option( + "--receipt-file", + help="Write the non-secret completion receipt to this file while keeping normal UI output.", + ), + ] = None, +) -> None: + """Configure Hermes for Databricks Model Serving without launching it.""" + if output not in {"text", "json"}: + print_err("--output must be one of: text, json.") + raise typer.Exit(2) + if receipt_file is not None and output == "json": + print_err("--receipt-file cannot be combined with --output json.") + raise typer.Exit(2) + if receipt_file is not None: + try: + receipt_file.unlink(missing_ok=True) + except OSError: + print_err("Unable to prepare completion receipt.") + raise typer.Exit(1) from None + try: + if output == "json": + # Discovery/auth helpers use the normal Rich UI. Suppress those progress + # messages so stdout remains a single machine-readable document. + with console.capture(): + result = configure_hermes_command( + workspace=workspace, + profile=profile, + model=model, + hermes_home=hermes_home, + ) + typer.echo(json.dumps(result, sort_keys=True)) + else: + result = configure_hermes_command( + workspace=workspace, + profile=profile, + model=model, + hermes_home=hermes_home, + receipt_file=receipt_file, + ) + print_success( + f"Hermes configured for Databricks Model Serving ({result['default_model']})" + ) + except RuntimeError as exc: + print_err(str(exc)) + raise typer.Exit(1) from None + except KeyboardInterrupt: + print_err("Interrupted.") + raise typer.Exit(130) from None + + @configure_app.command("mcp") def configure_mcp( location: Annotated[ diff --git a/src/ucode/config_io.py b/src/ucode/config_io.py index f67f3f32..fb51c233 100644 --- a/src/ucode/config_io.py +++ b/src/ucode/config_io.py @@ -4,7 +4,7 @@ import json from pathlib import Path -from typing import TypedDict, cast +from typing import Literal, NotRequired, TypedDict, cast import tomlkit import tomlkit.exceptions @@ -18,6 +18,7 @@ class ToolSpec(TypedDict): display: str config_path: Path backup_path: Path + install_method: NotRequired[Literal["npm", "external"]] APP_DIR = Path.home() / ".ucode" diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index ee54d153..ed4cdb88 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -129,6 +129,39 @@ def _debug(label: str, detail: str) -> None: _SECRET_KEY_PATTERN = re.compile(r"(token|secret|password|bearer|api_key|apikey)", re.IGNORECASE) +_SCOPED_BEARER: tuple[str, str | None, str] | None = None +_SECRET_ASSIGNMENT_PATTERN = re.compile( + r"\b([a-z0-9_]*(?:token|secret|password|bearer|api_key|apikey)[a-z0-9_]*\s*[:=]\s*)" + r"([^\s,;]+)", + re.IGNORECASE, +) +_TOKEN_VALUE_PATTERN = re.compile( + r"\bdapi[a-z0-9_-]{8,}\b|\beyJ[a-z0-9_-]{8,}\.[a-z0-9_-]{8,}\.[a-z0-9_-]{8,}\b", + re.IGNORECASE, +) +_BEARER_HEADER_PATTERN = re.compile( + r"(\bauthorization\s*:\s*bearer\s+)(\S+)", + re.IGNORECASE, +) +_ACCESS_TOKEN_PROSE_PATTERN = re.compile( + r"(\baccess[\s_-]+token\s+(?:is|=|:)\s*)(\S+)", + re.IGNORECASE, +) + + +def _scrub_subprocess_output(text: str) -> str: + """Redact structured and common plain-text credential forms from debug output.""" + stripped = text.strip() + if not stripped: + return "" + try: + structured = json.loads(stripped) + except json.JSONDecodeError: + scrubbed = _BEARER_HEADER_PATTERN.sub(r"\1", stripped) + scrubbed = _ACCESS_TOKEN_PROSE_PATTERN.sub(r"\1", scrubbed) + scrubbed = _SECRET_ASSIGNMENT_PATTERN.sub(r"\1", scrubbed) + return _TOKEN_VALUE_PATTERN.sub("", scrubbed) + return json.dumps(_scrub_json(structured)) def _format_subprocess_result( @@ -137,11 +170,11 @@ def _format_subprocess_result( """Format a CompletedProcess for the debug log without leaking tokens. On success, stdout is suppressed (it often contains the access token). - On failure, stdout/stderr are included truncated.""" - stderr = (result.stderr or "").strip()[:500] + On failure, scrubbed stdout/stderr are included truncated.""" + stderr = _scrub_subprocess_output(result.stderr or "")[:500] if result.returncode == 0: return f"rc=0 stderr={stderr!r}" - stdout = (result.stdout or "").strip()[:500] + stdout = _scrub_subprocess_output(result.stdout or "")[:500] return f"rc={result.returncode} stdout={stdout!r} stderr={stderr!r}" @@ -207,7 +240,7 @@ def _log_auth_diagnostics() -> None: _debug( "databricks auth profiles", f"rc={profiles_result.returncode} " - f"stderr={(profiles_result.stderr or '').strip()[:300]!r}", + f"stderr={_scrub_subprocess_output(profiles_result.stderr or '')[:300]!r}", ) if profiles_result.returncode == 0 and profiles_result.stdout: try: @@ -897,11 +930,33 @@ def _profile_args(profile: str | None) -> list[str]: return ["--profile", profile] if profile else [] -def has_valid_databricks_auth(workspace: str, profile: str | None = None) -> bool: - # Honor the CI short-circuit (see ``get_databricks_token``): if a - # pre-fetched bearer is available, treat auth as valid and skip the - # `databricks auth token` shell-out (which only knows user-OAuth). - if os.environ.get("DATABRICKS_BEARER", "").strip(): +def _bearer_is_allowed( + workspace: str, + profile: str | None, + bearer: str, + *, + allow_env_bearer: bool = False, +) -> bool: + if not bearer: + return False + if _SCOPED_BEARER is not None: + return _SCOPED_BEARER == (workspace.rstrip("/"), profile, bearer) + return allow_env_bearer + + +def has_valid_databricks_auth( + workspace: str, + profile: str | None = None, + *, + allow_env_bearer: bool = True, +) -> bool: + bearer = os.environ.get("DATABRICKS_BEARER", "").strip() + if _bearer_is_allowed( + workspace, + profile, + bearer, + allow_env_bearer=allow_env_bearer, + ): return True _log_auth_diagnostics() # Mirror run_databricks_login: when ~/.databrickscfg has multiple @@ -1048,24 +1103,42 @@ def resolve_pat_token(profile: str | None) -> str | None: return None -def ensure_pat_bearer(profile: str | None, pat: str | None = None) -> bool: +def _install_scoped_bearer( + workspace: str | None, + profile: str | None, + token: str, +) -> None: + """Install a PAT and bind its process-local use to the resolved profile.""" + global _SCOPED_BEARER + os.environ["DATABRICKS_BEARER"] = token + _SCOPED_BEARER = (workspace.rstrip("/"), profile, token) if workspace else None + + +def ensure_pat_bearer( + profile: str | None, + pat: str | None = None, + *, + workspace: str | None = None, +) -> bool: """Ensure ``DATABRICKS_BEARER`` holds a usable token for a ``--use-pat`` profile. - If a non-empty bearer is already in the environment it wins (the CI escape - hatch). Otherwise the profile's static PAT is exported — callers that have - already resolved it (e.g. ``configure_shared_state``) pass it via ``pat`` to - skip a redundant ``~/.databrickscfg`` read; everyone else lets this resolve - it. An exported-but-*empty* ``DATABRICKS_BEARER`` is treated as absent — + A PAT supplied explicitly by a scoped caller wins and replaces ambient state. + Otherwise a non-empty bearer already in the environment is retained as the + legacy CI escape hatch; if absent, this resolves the profile's static PAT. + An exported-but-*empty* ``DATABRICKS_BEARER`` is treated as absent — matching ``get_databricks_token``'s own ``.strip()`` check — so a stray ``export DATABRICKS_BEARER=`` does not shadow the PAT and silently force the OAuth path (which fails for PAT-only profiles). Returns ``True`` iff a usable bearer is now present in the environment.""" + if pat: + _install_scoped_bearer(workspace, profile, pat) + return True if os.environ.get("DATABRICKS_BEARER", "").strip(): return True - pat = pat or resolve_pat_token(profile) + pat = resolve_pat_token(profile) if pat: - os.environ["DATABRICKS_BEARER"] = pat + _install_scoped_bearer(workspace, profile, pat) return True return False @@ -1074,13 +1147,22 @@ def apply_pat_environment(state: dict) -> None: """Export the configured profile's PAT as ``DATABRICKS_BEARER`` when the workspace was configured with ``--use-pat``. - Every token fetch in this process (and in launched agent subprocesses, - which inherit the environment) then takes the existing static-bearer - short-circuit instead of the OAuth-only `databricks auth token` path. - A non-empty bearer already present in the environment is left untouched.""" + The configured profile identity outranks ambient state: its PAT replaces a + pre-existing bearer, and a missing/stale PAT clears that bearer so callers + fail closed instead of borrowing credentials from another workspace.""" + global _SCOPED_BEARER if not state.get("use_pat"): + if _SCOPED_BEARER is not None and os.environ.get("DATABRICKS_BEARER") == _SCOPED_BEARER[2]: + os.environ.pop("DATABRICKS_BEARER", None) + _SCOPED_BEARER = None return - ensure_pat_bearer(state.get("profile")) + profile = state.get("profile") + pat = resolve_pat_token(profile) + if pat: + ensure_pat_bearer(profile, pat, workspace=state.get("workspace")) + else: + os.environ.pop("DATABRICKS_BEARER", None) + _SCOPED_BEARER = None def run_databricks_login(workspace: str, profile: str | None = None) -> None: @@ -1111,7 +1193,11 @@ def run_databricks_login(workspace: str, profile: str | None = None) -> None: def ensure_databricks_auth( - workspace: str, profile: str | None = None, *, quiet: bool = False + workspace: str, + profile: str | None = None, + *, + quiet: bool = False, + allow_env_bearer: bool = True, ) -> None: """Check auth and login only if needed (used by launch path). @@ -1120,7 +1206,11 @@ def ensure_databricks_auth( login that actually runs is never silent. """ with spinner("Checking Databricks auth..."): - auth_is_valid = has_valid_databricks_auth(workspace, profile) + auth_is_valid = has_valid_databricks_auth( + workspace, + profile, + allow_env_bearer=allow_env_bearer, + ) if auth_is_valid: if not quiet: print_success(f"Databricks auth already available for {workspace}") @@ -1133,15 +1223,18 @@ def get_databricks_token( profile: str | None = None, *, force_refresh: bool = False, + allow_env_bearer: bool = True, ) -> str: - # ``DATABRICKS_BEARER`` is the CI escape hatch: when set, skip the - # `databricks auth token` subprocess entirely and return the pre-fetched - # bearer directly. Used by the e2e job, where the protected runner has - # no `databricks auth login` cache and `databricks auth token` only knows - # how to read user-OAuth caches (not M2M client_credentials). Mirrors the - # same short-circuit baked into ``build_auth_shell_command``. + # ``DATABRICKS_BEARER`` is accepted only when the caller explicitly permits + # the CI/M2M escape hatch, or when ucode installed this exact token for this + # exact PAT profile. Arbitrary ambient state cannot override identity. bearer = os.environ.get("DATABRICKS_BEARER", "").strip() - if bearer: + if _bearer_is_allowed( + workspace, + profile, + bearer, + allow_env_bearer=allow_env_bearer, + ): _debug("get_databricks_token", "using DATABRICKS_BEARER env var") return bearer @@ -1500,7 +1593,7 @@ def build_auth_shell_command( # Supported OSS chat families, matched by name substring. Add an entry to # support a new family. -_OSS_MODEL_FAMILIES = ("kimi-", "glm-", "deepseek-") +_OSS_MODEL_FAMILIES = ("gpt-oss-", "kimi-", "glm-", "deepseek-") # Claude model families ucode buckets, newest tier first. Each maps to a # Claude Code family alias (ANTHROPIC_DEFAULT__MODEL). Add an entry to @@ -1520,12 +1613,12 @@ def classify_model_family(model_id: str) -> str | None: for family in ANTHROPIC_FAMILIES: if f"claude-{family}-" in model_id: return family + if any(oss in model_id for oss in _OSS_MODEL_FAMILIES): + return "oss" if "gpt-" in model_id: return "codex" if "gemini-" in model_id: return "gemini" - if any(oss in model_id for oss in _OSS_MODEL_FAMILIES): - return "oss" return None @@ -1833,7 +1926,10 @@ def discover_model_services( # newest-wins once the router accepts opus-5 (PR databricks-eng/universe#2365446). _prefer_opus_4_8(claude_models, ids) - codex_models = sorted([m for m in ids if "gpt-" in m], key=model_version_sort_key) + codex_models = sorted( + [m for m in ids if "gpt-" in m and not any(family in m for family in _OSS_MODEL_FAMILIES)], + key=model_version_sort_key, + ) gemini_models = sorted([m for m in ids if "gemini-" in m], key=model_version_sort_key) oss_models = [m for m in ids if any(family in m for family in _OSS_MODEL_FAMILIES)] diff --git a/src/ucode/gateway_proxy.py b/src/ucode/gateway_proxy.py index e5cd4788..1acfc2b5 100644 --- a/src/ucode/gateway_proxy.py +++ b/src/ucode/gateway_proxy.py @@ -120,10 +120,12 @@ def __init__( profile: str | None, *, force_refresh_near_expiry: bool = False, + allow_env_bearer: bool = False, ) -> None: self._workspace = workspace self._profile = profile self._force_refresh_near_expiry = force_refresh_near_expiry + self._allow_env_bearer = allow_env_bearer self._state_lock = threading.Lock() # guards _token / _expiry (brief) self._refresh_lock = threading.Lock() # single-flights the CLI refresh self._stop = threading.Event() @@ -135,7 +137,20 @@ def __init__( def _refresh(self, *, force: bool) -> None: """Mint a token and record its expiry.""" - token = get_databricks_token(self._workspace, self._profile, force_refresh=force) + if self._allow_env_bearer: + token = get_databricks_token( + self._workspace, + self._profile, + force_refresh=force, + allow_env_bearer=True, + ) + else: + token = get_databricks_token( + self._workspace, + self._profile, + force_refresh=force, + allow_env_bearer=False, + ) expiry = _jwt_exp(token) or (time.time() + _DEFAULT_TTL_S) with self._state_lock: self._token = token @@ -374,6 +389,7 @@ def start_proxy( port: int, token_header: str, force_refresh_near_expiry: bool, + allow_env_bearer: bool = False, ) -> tuple[ThreadingHTTPServer, TokenCache, httpx.Client]: """Start the loopback refresh proxy + its background token refresher. @@ -390,6 +406,7 @@ def start_proxy( workspace, profile, force_refresh_near_expiry=force_refresh_near_expiry, + allow_env_bearer=allow_env_bearer, ) # One pooled, keep-alive client shared across handler threads: reuses TCP+TLS # to the gateway instead of a fresh handshake per request. Don't follow diff --git a/src/ucode/mcp.py b/src/ucode/mcp.py index f855d086..aa19d751 100644 --- a/src/ucode/mcp.py +++ b/src/ucode/mcp.py @@ -10,6 +10,7 @@ import threading from collections.abc import Callable from concurrent.futures import ThreadPoolExecutor, as_completed +from pathlib import Path from typing import Any from urllib.parse import urlparse @@ -26,7 +27,7 @@ from questionary.question import Question from questionary.styles import merge_styles_default -from ucode.agents import copilot, cursor, gemini, opencode +from ucode.agents import copilot, cursor, gemini, hermes, opencode from ucode.config_io import restore_file from ucode.databricks import ( apply_pat_environment, @@ -98,6 +99,11 @@ class _Back: "display": "Cursor", "list_command": "cursor-agent mcp list", }, + "hermes": { + "binary": "hermes", + "display": "Hermes", + "list_command": "hermes mcp list", + }, } SKILLS_MCP_KIND = "skills" SKILLS_MCP_SERVER_NAME = "databricks-skill-registry" @@ -300,6 +306,8 @@ def configure_client_mcp_server( *, use_pat: bool = False, always_load: bool = False, + hermes_home: str | None = None, + expected_fingerprint: str | None = None, ) -> list[str]: # Every client registers the same `ucode mcp-proxy ...` stdio command; the # proxy forwards to `url` and refreshes the Databricks token itself. Only the @@ -330,10 +338,24 @@ def configure_client_mcp_server( if client == "cursor": removed = cursor.write_mcp_server_config(name, argv) return [MCP_USER_SCOPE] if removed else [] + if client == "hermes": + removed = hermes.write_mcp_server_config( + name, + argv, + hermes_home=hermes_home, + expected_fingerprint=expected_fingerprint, + ) + return [MCP_USER_SCOPE] if removed else [] raise RuntimeError(f"Unsupported MCP client '{client}'.") -def remove_client_mcp_server(client: str, name: str) -> list[str]: +def remove_client_mcp_server( + client: str, + name: str, + *, + hermes_home: str | None = None, + expected_fingerprint: str | None = None, +) -> list[str]: if client == "claude": return [scope for scope in MCP_CLEANUP_SCOPES if remove_claude_mcp_server(name, scope)] if client == "codex": @@ -346,9 +368,42 @@ def remove_client_mcp_server(client: str, name: str) -> list[str]: return [MCP_USER_SCOPE] if copilot.remove_mcp_server_config(name) else [] if client == "cursor": return [MCP_USER_SCOPE] if cursor.remove_mcp_server_config(name) else [] + if client == "hermes": + removed = hermes.remove_mcp_server_config( + name, + hermes_home=hermes_home, + expected_fingerprint=expected_fingerprint, + ) + return [MCP_USER_SCOPE] if removed else [] raise RuntimeError(f"Unsupported MCP client '{client}'.") +def _remove_recorded_client_server( + client: str, + name: str, + server: dict, + *, + hermes_home: str | None, +) -> list[str]: + if client != "hermes": + return remove_client_mcp_server(client, name) + raw_ownership = server.get("ownership") + ownership: dict = raw_ownership if isinstance(raw_ownership, dict) else {} + hermes_ownership = ownership.get("hermes") + if not isinstance(hermes_ownership, dict): + return [] + owned_home = hermes_ownership.get("hermes_home") + owned_fingerprint = hermes_ownership.get("fingerprint") + if not isinstance(owned_home, str) or not isinstance(owned_fingerprint, str): + return [] + return remove_client_mcp_server( + client, + name, + hermes_home=owned_home, + expected_fingerprint=owned_fingerprint, + ) + + def revert_mcp_configs(state: dict) -> dict[str, bool]: results: dict[str, bool] = {} # Both the developer's own servers and any registered from the workspace's managed config, so a @@ -363,7 +418,12 @@ def revert_mcp_configs(state: dict) -> dict[str, bool]: for client in server.get("clients") or []: if client not in MCP_CLIENTS: continue - removed_scopes = remove_client_mcp_server(client, name) + removed_scopes = _remove_recorded_client_server( + client, + name, + server, + hermes_home=_recorded_hermes_home(state), + ) results[client] = bool(removed_scopes) or results.get(client, False) # OpenCode MCP entries live in the normal OpenCode config and are restored @@ -379,6 +439,17 @@ def revert_mcp_configs(state: dict) -> dict[str, bool]: return results +def _recorded_hermes_home(state: dict) -> str | None: + managed = state.get("managed_configs") + if not isinstance(managed, dict): + return None + hermes_managed = managed.get("hermes") + if not isinstance(hermes_managed, dict): + return None + value = hermes_managed.get("hermes_home") + return value if isinstance(value, str) and value else None + + def _coerce_bool(value: object) -> bool | None: if isinstance(value, bool): return value @@ -1133,7 +1204,7 @@ def apply_managed_mcp_servers( return [] entries = managed.get("mcp_servers") if not isinstance(entries, list): - return [] + entries = [] working: list[dict] = [] seen: set[str] = set() skipped: list[str] = [] @@ -1158,8 +1229,6 @@ def apply_managed_mcp_servers( "Skipping managed MCP server(s) ucode can't yet auto-register from the workspace " f"config: {', '.join(skipped)}. Add them with `ucode configure mcp`." ) - if not working: - return [] # Diff against the managed servers ucode registered on a prior launch so a removed entry is # unregistered and an unchanged one is a no-op. Only this tool's managed servers are considered. state = load_state() @@ -1168,7 +1237,17 @@ def apply_managed_mcp_servers( for server in (state.get("managed_mcp_servers") or []) if isinstance(server, dict) and tool in (server.get("clients") or []) ] - apply_mcp_server_changes(previous, working, [tool], workspace, profile, use_pat=use_pat) + if not working and not previous: + return [] + apply_mcp_server_changes( + previous, + working, + [tool], + workspace, + profile, + use_pat=use_pat, + hermes_home=_recorded_hermes_home(state), + ) return working @@ -1219,7 +1298,13 @@ def apply_managed_skills( original = list(state.get("mcp_servers") or []) working = _resolve_skills_mcp_servers(workspace, [tool], new_locations, original) changed = apply_mcp_server_changes( - original, working, [tool], workspace, profile, use_pat=use_pat + original, + working, + [tool], + workspace, + profile, + use_pat=use_pat, + hermes_home=_recorded_hermes_home(state), ) if not (changed or original != working or prev_managed != desired): return [] @@ -1458,7 +1543,25 @@ def apply_mcp_server_changes( profile: str | None = None, *, use_pat: bool = False, + hermes_home: str | None = None, ) -> bool: + target_hermes_home = str( + Path(hermes_home or os.environ.get("HERMES_HOME") or Path.home() / ".hermes") + .expanduser() + .resolve() + ) + if "hermes" in clients: + for server in working_servers: + url = server.get("url") + if not isinstance(url, str) or not url: + continue + argv = build_mcp_proxy_argv(url, workspace, profile, use_pat=use_pat) + ownership = dict(server.get("ownership") or {}) + ownership["hermes"] = { + "fingerprint": hermes.mcp_server_fingerprint(hermes.mcp_value_for_argv(argv)), + "hermes_home": target_hermes_home, + } + server["ownership"] = ownership original_by_name = _servers_by_name(original_servers) working_by_name = _servers_by_name(working_servers) @@ -1475,7 +1578,12 @@ def apply_mcp_server_changes( if name not in working_by_name: for client in _mcp_server_clients(server): work.setdefault(client, []).append( - lambda c=client, n=name: remove_client_mcp_server(c, n) + lambda c=client, n=name, s=server: _remove_recorded_client_server( + c, + n, + s, + hermes_home=hermes_home, + ) ) changed = True @@ -1489,10 +1597,31 @@ def apply_mcp_server_changes( # alwaysLoad (Claude-only) keeps the skills registry's utility tools # discoverable without an explicit mention; other clients ignore it. always_load = server.get("kind") == SKILLS_MCP_KIND + original_ownership = original.get("ownership") if isinstance(original, dict) else None + original_hermes = ( + original_ownership.get("hermes") if isinstance(original_ownership, dict) else None + ) + expected_hermes_fingerprint = None + if ( + isinstance(original_hermes, dict) + and original_hermes.get("hermes_home") == target_hermes_home + and isinstance(original_hermes.get("fingerprint"), str) + ): + expected_hermes_fingerprint = original_hermes["fingerprint"] for client in clients: work[client].append( - lambda c=client, n=name, u=url, al=always_load: configure_client_mcp_server( - c, n, u, workspace, profile, use_pat=use_pat, always_load=al + lambda c=client, n=name, u=url, al=always_load, expected=expected_hermes_fingerprint: ( + configure_client_mcp_server( + c, + n, + u, + workspace, + profile, + use_pat=use_pat, + always_load=al, + hermes_home=target_hermes_home, + expected_fingerprint=expected if c == "hermes" else None, + ) ) ) changed = True @@ -1550,23 +1679,39 @@ def purge_cross_workspace_mcp_residue(state: dict, workspace: str) -> None: ) noun = "entry" if len(foreign_mcp_servers) == 1 else "entries" print_warning( - f"Dropping {len(foreign_mcp_servers)} stale MCP {noun} " + f"Cleaning {len(foreign_mcp_servers)} stale MCP {noun} " f"not bound to this workspace: {foreign_names}." ) + retained_foreign: list[dict] = [] for server in foreign_mcp_servers: name = _server_name(server) if not name: + retained_foreign.append(server) continue + fully_removed = True for client in server.get("clients") or []: if client not in installed or client not in MCP_CLIENTS: + if client == "hermes": + fully_removed = False continue try: - remove_client_mcp_server(client, name) + removed_scopes = _remove_recorded_client_server( + client, + name, + server, + hermes_home=_recorded_hermes_home(state), + ) except RuntimeError as exc: + fully_removed = False print_warning( f"Failed to remove `{name}` from {MCP_CLIENTS[client]['display']}: {exc}" ) - state["mcp_servers"] = current_mcp_servers + continue + if client == "hermes" and not removed_scopes: + fully_removed = False + if not fully_removed: + retained_foreign.append(server) + state["mcp_servers"] = [*current_mcp_servers, *retained_foreign] save_state(state) other_ws_mcps = _mcp_entries_only_in_other_workspaces(workspace) @@ -1869,6 +2014,7 @@ def configure_mcp_command( workspace, profile, use_pat=bool(state.get("use_pat")), + hermes_home=_recorded_hermes_home(state), ) if changed or original_mcp_servers_for_location != working_mcp_servers: state["mcp_servers"] = working_mcp_servers @@ -1966,6 +2112,7 @@ def configure_mcp_command( workspace, profile, use_pat=bool(state.get("use_pat")), + hermes_home=_recorded_hermes_home(state), ) if changed or original_mcp_servers != working_mcp_servers: state["mcp_servers"] = working_mcp_servers @@ -2074,7 +2221,13 @@ def remove_mcp_command(agents: set[str] | None = None) -> int: if targets: removal_view.append({**server, "clients": targets}) changed = apply_mcp_server_changes( - removal_view, [], clients, workspace, profile, use_pat=bool(state.get("use_pat")) + removal_view, + [], + clients, + workspace, + profile, + use_pat=bool(state.get("use_pat")), + hermes_home=_recorded_hermes_home(state), ) # Update saved state: drop a fully-removed server, or keep it with the named @@ -2178,7 +2331,14 @@ def _update_skills_mcp( """Rebuild the single skills connection for ``locations`` and persist it.""" original = list(state.get("mcp_servers") or []) working = _resolve_skills_mcp_servers(workspace, clients, locations, original) - changed = apply_mcp_server_changes(original, working, clients, workspace, profile) + changed = apply_mcp_server_changes( + original, + working, + clients, + workspace, + profile, + hermes_home=_recorded_hermes_home(state), + ) if changed or original != working: state["mcp_servers"] = working save_state(state) diff --git a/src/ucode/mcp_proxy.py b/src/ucode/mcp_proxy.py index 0d2a68cb..8aaed5c8 100644 --- a/src/ucode/mcp_proxy.py +++ b/src/ucode/mcp_proxy.py @@ -42,7 +42,7 @@ from mcp.client.streamable_http import streamable_http_client from mcp.server.stdio import stdio_server -from ucode.databricks import ensure_pat_bearer, get_databricks_token +from ucode.databricks import ensure_pat_bearer, get_databricks_token, resolve_pat_token # Exit code used when the proxy cannot continue. MCP clients surface a non-zero # exit far more usefully than a timeout, so bail out instead of hanging. @@ -111,7 +111,12 @@ def _fail_fast(message: str) -> None: raise SystemExit(AUTH_FAILURE_EXIT_CODE) -def _build_token_auth(workspace: str, profile: str | None): +def _build_token_auth( + workspace: str, + profile: str | None, + *, + allow_env_bearer: bool = False, +): """Build an httpx ``Auth`` that injects a fresh bearer on every request. The base class comes from whichever httpx the SDK uses (see ``_httpx``), so @@ -122,15 +127,26 @@ def _build_token_auth(workspace: str, profile: str | None): class _DatabricksTokenAuth(httpx.Auth): def auth_flow(self, request): - # get_databricks_token honors the DATABRICKS_BEARER short-circuit and - # PAT profiles internally; --use-pat is surfaced via the env ucode set. + # Only an explicit --use-pat path may consume the bearer exported by + # ensure_pat_bearer; OAuth profiles mint against their scoped identity. # A RuntimeError here means auth is dead (expired refresh token, # logged-out profile). Raising it from inside auth_flow would tear # through the transport's task group and stall the process until the # client times out, so translate it into a terminal ProxyAuthError the # caller reports cleanly. try: - token = get_databricks_token(workspace, profile) + if allow_env_bearer: + token = get_databricks_token( + workspace, + profile, + allow_env_bearer=True, + ) + else: + token = get_databricks_token( + workspace, + profile, + allow_env_bearer=False, + ) except RuntimeError as exc: raise ProxyAuthError(str(exc)) from exc request.headers["Authorization"] = f"Bearer {token}" @@ -166,9 +182,18 @@ async def _pump_upstream[T]( raise ProxyTransportError("upstream MCP transport closed unexpectedly") -async def _run(url: str, workspace: str, profile: str | None) -> None: +async def _run( + url: str, + workspace: str, + profile: str | None, + allow_env_bearer: bool = False, +) -> None: httpx = _httpx() - auth = _build_token_auth(workspace, profile) + auth = _build_token_auth( + workspace, + profile, + allow_env_bearer=allow_env_bearer, + ) # 2.x-native shape: hand the transport a pre-built AsyncClient carrying our # per-request auth. Works on mcp 1.28+ and 2.x; `streamable_http_client` # yields a (read, write) pair in both. @@ -194,7 +219,12 @@ async def _run(url: str, workspace: str, profile: str | None) -> None: tg.cancel_scope.cancel() -def _preflight_token(workspace: str, profile: str | None) -> None: +def _preflight_token( + workspace: str, + profile: str | None, + *, + allow_env_bearer: bool = False, +) -> None: """Verify a Databricks token can be minted before opening the bridge. Raises ``RuntimeError`` (with the CLI's own message) when auth is dead. This @@ -202,7 +232,18 @@ def _preflight_token(workspace: str, profile: str | None) -> None: with subprocess timeouts, so it returns or fails on its own — the point here is only to *locate* the failure before the transport starts, where it can be reported instead of stalling the session.""" - get_databricks_token(workspace, profile) + if allow_env_bearer: + get_databricks_token( + workspace, + profile, + allow_env_bearer=True, + ) + else: + get_databricks_token( + workspace, + profile, + allow_env_bearer=False, + ) def _unwrap_proxy_error(exc: BaseException) -> ProxyAuthError | ProxyTransportError | None: @@ -230,23 +271,35 @@ def serve(url: str, workspace: str, profile: str | None = None, *, use_pat: bool token`` only reads OAuth caches, so a PAT profile's token must be exported as ``DATABRICKS_BEARER`` first (``ensure_pat_bearer``) — then every per-request mint takes that short-circuit. OAuth needs no such step.""" - if use_pat and not ensure_pat_bearer(profile): - _fail_fast( - "--use-pat is set but no personal access token was found for profile " - f"'{profile or ''}' in ~/.databrickscfg (expected auth_type = pat). " - "Set DATABRICKS_BEARER, or reconfigure the profile." - ) + if use_pat: + pat = resolve_pat_token(profile) + if not pat or not ensure_pat_bearer(profile, pat, workspace=workspace): + _fail_fast( + "--use-pat is set but no personal access token was found for profile " + f"'{profile or ''}' in ~/.databrickscfg (expected auth_type = pat). " + "Reconfigure the profile." + ) # Pre-flight the token before opening the bridge. Without this, the first # token failure surfaces from inside the transport's task group, where it can # stall the process instead of erroring out. try: - _preflight_token(workspace, profile) + if use_pat: + _preflight_token( + workspace, + profile, + allow_env_bearer=True, + ) + else: + _preflight_token(workspace, profile) except RuntimeError as exc: _fail_fast(str(exc)) try: - anyio.run(_run, url, workspace, profile) + if use_pat: + anyio.run(_run, url, workspace, profile, True) + else: + anyio.run(_run, url, workspace, profile) except BaseException as exc: # noqa: BLE001 - re-raised unless it's a known proxy failure # Errors raised inside the transport arrive wrapped by its task group. # Report expected auth/transport failures without hiding programming bugs. diff --git a/src/ucode/state.py b/src/ucode/state.py index 6344a7e5..6645543c 100644 --- a/src/ucode/state.py +++ b/src/ucode/state.py @@ -106,6 +106,40 @@ def set_current_workspace(workspace: str | None) -> None: raise RuntimeError(f"Failed to write state file: {STATE_PATH}") from exc +def _hermes_ownership_metadata(entry: dict) -> dict: + """Return validated, non-secret Hermes cleanup ownership metadata.""" + hermes_home = entry.get("hermes_home") + active_model = entry.get("active_model") + if not isinstance(hermes_home, str) or not hermes_home: + return {} + if not isinstance(active_model, dict): + return {} + provider = active_model.get("provider") + model = active_model.get("default") + if not isinstance(provider, str) or not provider: + return {} + if not isinstance(model, str) or not model: + return {} + result = { + "hermes_home": hermes_home, + "active_model": {"provider": provider, "default": model}, + } + raw_fingerprints = entry.get("provider_fingerprints") + if isinstance(raw_fingerprints, dict): + fingerprints = { + provider_id: fingerprint + for provider_id, fingerprint in raw_fingerprints.items() + if isinstance(provider_id, str) + and provider_id.startswith("ucode-databricks-") + and isinstance(fingerprint, str) + and len(fingerprint) == 64 + and all(char in "0123456789abcdef" for char in fingerprint) + } + if fingerprints: + result["provider_fingerprints"] = fingerprints + return result + + def hydrate_state(state: dict) -> dict: """Normalize a workspace state entry and add derived harness config. @@ -125,6 +159,8 @@ def hydrate_state(state: dict) -> dict: if isinstance(entry, dict): keys = entry.get("keys") if isinstance(entry.get("keys"), list) else [] normalized[tool] = {"keys": keys} + if tool == "hermes": + normalized[tool].update(_hermes_ownership_metadata(entry)) elif entry: normalized[tool] = {"keys": []} hydrated["managed_configs"] = normalized @@ -236,10 +272,19 @@ def clear_state() -> None: raise RuntimeError(f"Failed to clear state file: {STATE_PATH}") from exc -def mark_tool_managed(state: dict, tool: str, managed_keys: list) -> dict: +def mark_tool_managed( + state: dict, + tool: str, + managed_keys: list, + *, + metadata: dict | None = None, +) -> dict: """Record which config keys ucode manages for ``tool``.""" managed_configs = dict(state.get("managed_configs") or {}) - managed_configs[tool] = {"keys": list(managed_keys)} + entry = {"keys": list(managed_keys)} + if tool == "hermes" and isinstance(metadata, dict): + entry.update(_hermes_ownership_metadata(metadata)) + managed_configs[tool] = entry state["managed_configs"] = managed_configs state["last_tool"] = tool return state diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index 61f07438..05b1f48a 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -1046,7 +1046,14 @@ def __init__(self, argv): def wait(self): return 0 - def start_proxy(workspace, profile, port, token_header, force_refresh_near_expiry): + def start_proxy( + workspace, + profile, + port, + token_header, + force_refresh_near_expiry, + allow_env_bearer=False, + ): calls.append( ( "proxy", @@ -1055,6 +1062,7 @@ def start_proxy(workspace, profile, port, token_header, force_refresh_near_expir port, token_header, force_refresh_near_expiry, + allow_env_bearer, ) ) return Server(), Cache(), Client() @@ -1083,6 +1091,7 @@ def start_proxy(workspace, profile, port, token_header, force_refresh_near_expir 12345, claude.gateway_proxy.AI_GATEWAY_TOKEN_HEADER, False, + False, ) assert calls[-3:] == [("stop",), ("shutdown",), ("close",)] diff --git a/tests/test_agent_hermes.py b/tests/test_agent_hermes.py new file mode 100644 index 00000000..350cb6fa --- /dev/null +++ b/tests/test_agent_hermes.py @@ -0,0 +1,806 @@ +"""Contract tests for the Hermes agent configuration adapter.""" + +from __future__ import annotations + +import copy +import json +import shlex +import subprocess +from pathlib import Path + +import pytest + +from ucode.agents import hermes +from ucode.config_io import set_dry_run +from ucode.databricks import build_auth_token_argv + +WS = "https://example.databricks.com" +MODEL = "system.ai.gpt-5-6" +PROVIDER_ID = "ucode-databricks-codex" +ANTHROPIC_PROVIDER_ID = "ucode-databricks-anthropic" +OSS_PROVIDER_ID = "ucode-databricks-oss" +GEMINI_PROVIDER_ID = "ucode-databricks-gemini" + + +def _state(**overrides): + state = { + "workspace": WS, + "profile": "team prod", + "codex_models": ["system.ai.gpt-5", MODEL], + "claude_models": { + "sonnet": "system.ai.claude-sonnet-4-6", + "haiku": "system.ai.claude-haiku-4-5", + }, + "gemini_models": ["system.ai.gemini-3-1-pro"], + "oss_models": [ + "system.ai.deepseek-v3-2", + "system.ai.gpt-oss-120b", + "system.ai.gpt-oss-20b", + ], + } + state.update(overrides) + return state + + +class TestHermesSpec: + def test_is_an_externally_installed_hermes_binary(self): + assert hermes.SPEC["binary"] == "hermes" + assert hermes.SPEC["display"] == "Hermes" + assert hermes.SPEC.get("install_method") == "external" + + +class TestRenderConfigPatch: + def test_initial_provider_writes_require_managed_ids_to_be_absent(self): + patch = hermes.render_config_patch(_state()) + + provider_paths = {path for path in patch["set"] if path.startswith("providers.")} + assert set(patch["expect_missing"]) == provider_paths + assert patch["expect_hashes"] == {} + + def test_reconfiguration_requires_owned_provider_fingerprint(self): + previous = hermes.render_config_patch(_state()) + provider_path = f"providers.{PROVIDER_ID}" + fingerprint = hermes.config_value_fingerprint(previous["set"][provider_path]) + patch = hermes.render_config_patch( + _state( + managed_configs={ + "hermes": { + "hermes_home": str(Path.home() / ".hermes"), + "provider_fingerprints": {PROVIDER_ID: fingerprint}, + } + } + ) + ) + + assert patch["expect_hashes"] == {provider_path: fingerprint} + assert provider_path not in patch["expect_missing"] + + def test_reconfiguration_requires_owned_active_model_pair(self): + owned_model = {"provider": PROVIDER_ID, "default": MODEL} + patch = hermes.render_config_patch( + _state( + managed_configs={ + "hermes": { + "hermes_home": str(Path.home() / ".hermes"), + "active_model": owned_model, + } + } + ) + ) + + assert patch["expect_hashes"]["model.provider"] == hermes.config_value_fingerprint( + owned_model["provider"] + ) + assert patch["expect_hashes"]["model.default"] == hermes.config_value_fingerprint( + owned_model["default"] + ) + + def test_responses_provider_uses_gateway_codex_route(self): + patch = hermes.render_config_patch(_state()) + provider = patch["set"][f"providers.{PROVIDER_ID}"] + + assert provider["api"] == f"{WS}/ai-gateway/codex/v1" + assert provider["transport"] == "codex_responses" + + def test_uses_current_model_selector_and_mapping(self): + patch = hermes.render_config_patch(_state()) + provider = patch["set"][f"providers.{PROVIDER_ID}"] + + assert patch["set"]["model.provider"] == PROVIDER_ID + assert patch["set"]["model.default"] == MODEL + assert provider["default_model"] == MODEL + assert provider["models"] == { + "system.ai.gpt-5": {}, + MODEL: {}, + } + + def test_auth_helper_is_shell_escaped_command_string(self): + patch = hermes.render_config_patch(_state()) + key_cmd = patch["set"][f"providers.{PROVIDER_ID}"]["key_cmd"] + + assert isinstance(key_cmd, str) + argv = shlex.split(key_cmd) + assert argv[0].endswith("ucode") or argv[0] == "ucode" + assert argv[1] == "auth-token" + assert "team prod" in argv + + def test_auth_helper_uses_windows_command_line_quoting(self, monkeypatch): + monkeypatch.setattr("ucode.databricks.platform.system", lambda: "Windows") + patch = hermes.render_config_patch(_state()) + key_cmd = patch["set"][f"providers.{PROVIDER_ID}"]["key_cmd"] + + expected = subprocess.list2cmdline(build_auth_token_argv(WS, "team prod", use_pat=False)) + assert key_cmd == expected + + def test_generates_only_wire_compatible_route_providers(self): + patch = hermes.render_config_patch(_state()) + + responses = patch["set"][f"providers.{PROVIDER_ID}"] + anthropic = patch["set"][f"providers.{ANTHROPIC_PROVIDER_ID}"] + oss = patch["set"][f"providers.{OSS_PROVIDER_ID}"] + gemini = patch["set"][f"providers.{GEMINI_PROVIDER_ID}"] + + assert responses["api"] == f"{WS}/ai-gateway/codex/v1" + assert responses["transport"] == "codex_responses" + assert set(responses["models"]) == {"system.ai.gpt-5", MODEL} + assert anthropic["api"] == f"{WS}/ai-gateway/anthropic" + assert anthropic["transport"] == "anthropic_messages" + assert set(anthropic["models"]) == { + "system.ai.claude-sonnet-4-6", + "system.ai.claude-haiku-4-5", + } + assert oss["api"] == f"{WS}/ai-gateway/mlflow/v1" + assert oss["transport"] == "openai_chat" + assert set(oss["models"]) == {"system.ai.deepseek-v3-2"} + + assert gemini["api"] == f"{WS}/ai-gateway/gemini/v1beta" + assert gemini["transport"] == "gemini-native" + assert set(gemini["models"]) == {"system.ai.gemini-3-1-pro"} + + def test_stale_overlapping_state_omits_unsupported_gpt_oss(self): + model = "system.ai.gpt-oss-120b" + state = _state(codex_models=["system.ai.gpt-5", model]) + + patch = hermes.render_config_patch(state) + + providers = [ + patch["set"][f"providers.{provider_id}"] + for provider_id in ( + PROVIDER_ID, + ANTHROPIC_PROVIDER_ID, + OSS_PROVIDER_ID, + GEMINI_PROVIDER_ID, + ) + ] + model_sets = [set(provider["models"]) for provider in providers] + for index, current in enumerate(model_sets): + for other in model_sets[index + 1 :]: + assert current.isdisjoint(other) + assert all(model not in models for models in model_sets) + + def test_default_model_skips_unsupported_gpt_oss(self): + state = _state( + codex_default_model="system.ai.gpt-oss-120b", + codex_models=["system.ai.gpt-oss-120b"], + claude_models={"opus": "system.ai.claude-opus-4-8"}, + oss_models=["system.ai.gpt-oss-120b", "system.ai.deepseek-v3-2"], + ) + + assert hermes.default_model(state) == "system.ai.claude-opus-4-8" + + @pytest.mark.parametrize( + "model", + ["system.ai.gpt-oss-20b", "system.ai.gpt-oss-120b"], + ) + def test_explicit_unsupported_gpt_oss_is_rejected(self, model): + with pytest.raises(RuntimeError, match="not supported by Hermes"): + hermes.render_config_patch(_state(), model=model) + + def test_explicit_compatible_model_selects_its_route(self): + patch = hermes.render_config_patch(_state(), model="system.ai.claude-haiku-4-5") + + assert patch["set"]["model.provider"] == ANTHROPIC_PROVIDER_ID + assert patch["set"]["model.default"] == "system.ai.claude-haiku-4-5" + + patch = hermes.render_config_patch(_state(), model="system.ai.deepseek-v3-2") + + assert patch["set"]["model.provider"] == OSS_PROVIDER_ID + assert patch["set"]["model.default"] == "system.ai.deepseek-v3-2" + + def test_explicit_gemini_model_selects_native_route(self): + patch = hermes.render_config_patch(_state(), model="system.ai.gemini-3-1-pro") + + assert patch["set"]["model.provider"] == GEMINI_PROVIDER_ID + assert patch["set"]["model.default"] == "system.ai.gemini-3-1-pro" + + def test_rendering_is_deterministic_and_token_free(self): + state = _state(access_token="sentinel-access-token", token="sentinel-pat") + + first = hermes.render_config_patch(state) + second = hermes.render_config_patch(state) + + assert first == second + encoded = json.dumps(first, sort_keys=True) + assert "sentinel-access-token" not in encoded + assert "sentinel-pat" not in encoded + + def test_patch_only_manages_generated_provider_and_active_model(self): + patch = hermes.render_config_patch(_state()) + + assert set(patch["set"]) == { + f"providers.{PROVIDER_ID}", + f"providers.{ANTHROPIC_PROVIDER_ID}", + f"providers.{OSS_PROVIDER_ID}", + f"providers.{GEMINI_PROVIDER_ID}", + "model.provider", + "model.default", + } + assert patch["unset"] == [] + + def test_reconfigure_unsets_only_disappeared_managed_family(self): + prior = hermes.render_config_patch(_state()) + gemini_value = prior["set"][f"providers.{GEMINI_PROVIDER_ID}"] + fingerprint = hermes.config_value_fingerprint(gemini_value) + patch = hermes.render_config_patch( + _state( + gemini_models=[], + managed_configs={ + "hermes": {"provider_fingerprints": {GEMINI_PROVIDER_ID: fingerprint}} + }, + ) + ) + + assert patch["unset"] == [] + assert patch["unset_if_hash"] == {f"providers.{GEMINI_PROVIDER_ID}": fingerprint} + + def test_unconfigure_is_surgical(self): + owned = {"provider": PROVIDER_ID, "default": MODEL} + fingerprints = dict.fromkeys(hermes.MANAGED_PROVIDER_IDS, "0" * 64) + patch = hermes.render_unconfigure_patch( + current_model=owned, + owned_model=owned, + owned_provider_fingerprints=fingerprints, + current_provider_fingerprint="0" * 64, + ) + + assert patch["set"] == {} + assert set(patch["unset"]) == { + f"providers.{PROVIDER_ID}", + "model.provider", + "model.default", + } + assert set(patch["unset_if_hash"]) == set(hermes.MANAGED_PATHS[1:4]) + + def test_unconfigure_preserves_user_switched_active_pair(self): + patch = hermes.render_unconfigure_patch( + current_model={"provider": "vertex", "default": "gemini-user"}, + owned_model={"provider": PROVIDER_ID, "default": MODEL}, + owned_provider_fingerprints=dict.fromkeys(hermes.MANAGED_PROVIDER_IDS, "0" * 64), + current_provider_fingerprint=None, + ) + + assert "model.provider" not in patch["unset"] + assert "model.default" not in patch["unset"] + assert set(patch["unset_if_hash"]) == set(hermes.MANAGED_PATHS[:4]) + + def test_unconfigure_preserves_changed_model_and_its_managed_provider(self): + patch = hermes.render_unconfigure_patch( + current_model={"provider": PROVIDER_ID, "default": "system.ai.user-choice"}, + owned_model={"provider": PROVIDER_ID, "default": MODEL}, + owned_provider_fingerprints=dict.fromkeys(hermes.MANAGED_PROVIDER_IDS, "0" * 64), + current_provider_fingerprint="0" * 64, + ) + + assert f"providers.{PROVIDER_ID}" not in patch["unset_if_hash"] + assert "model.provider" not in patch["unset"] + assert "model.default" not in patch["unset"] + + def test_unconfigure_preserves_active_pair_for_same_id_user_replacement(self): + owned = {"provider": PROVIDER_ID, "default": MODEL} + patch = hermes.render_unconfigure_patch( + current_model=owned, + owned_model=owned, + owned_provider_fingerprints={PROVIDER_ID: "a" * 64}, + current_provider_fingerprint="b" * 64, + ) + + assert patch["unset"] == [] + assert patch["expect_hashes"] == {} + assert f"providers.{PROVIDER_ID}" not in patch["unset_if_hash"] + + +class TestUnconfigure: + def test_read_active_model_preserves_partial_user_configuration(self, monkeypatch, tmp_path): + monkeypatch.setattr( + hermes, + "read_config_value", + lambda _key, *, hermes_home: {"provider": "user-provider"}, + ) + + assert hermes.read_active_model(hermes_home=tmp_path) == {"provider": "user-provider"} + + def test_reads_and_applies_against_explicit_home(self, monkeypatch, tmp_path): + owned = {"provider": PROVIDER_ID, "default": MODEL} + provider_value = {"api": "https://owned.example/v1"} + fingerprint = hermes.config_value_fingerprint(provider_value) + seen = {} + + def read(*, hermes_home): + seen["read_home"] = hermes_home + return owned + + def apply(patch, *, hermes_home): + seen["patch"] = patch + seen["apply_home"] = hermes_home + return {"status": "applied"} + + monkeypatch.setattr(hermes, "read_active_model", read) + monkeypatch.setattr( + hermes, + "read_config_value", + lambda _key, *, hermes_home: provider_value, + ) + monkeypatch.setattr(hermes, "apply_config_patch", apply) + + hermes.unconfigure( + hermes_home=tmp_path, + owned_model=owned, + owned_provider_fingerprints={PROVIDER_ID: fingerprint}, + ) + + assert seen["read_home"] == tmp_path + assert seen["apply_home"] == tmp_path + assert "model.provider" in seen["patch"]["unset"] + + +class TestWriteToolConfig: + def test_cross_home_ownership_is_not_reused(self): + state = _state( + gemini_models=[], + managed_configs={ + "hermes": { + "keys": [], + "hermes_home": "/profiles/a", + "active_model": {"provider": PROVIDER_ID, "default": MODEL}, + "provider_fingerprints": {GEMINI_PROVIDER_ID: "a" * 64}, + } + }, + ) + + patch = hermes.render_config_patch( + hermes.state_scoped_to_home(state, "/profiles/b"), + model=MODEL, + ) + + assert patch["unset_if_hash"] == {} + + def test_records_home_and_active_pair_after_success(self, monkeypatch, tmp_path): + state = _state() + events = [] + monkeypatch.setattr("ucode.state.save_state", lambda value: events.append(("save", value))) + + def apply(*_args, **_kwargs): + assert events and events[0][0] == "save" + events.append(("apply", None)) + + monkeypatch.setattr(hermes, "apply_config_patch", apply) + + result = hermes.write_tool_config(state, hermes_home=tmp_path) + + ownership = result["managed_configs"]["hermes"] + assert ownership["keys"] == [] + assert ownership["hermes_home"] == str(tmp_path.resolve()) + assert ownership["active_model"] == { + "provider": PROVIDER_ID, + "default": MODEL, + } + assert set(ownership["provider_fingerprints"]) == set(hermes.MANAGED_PROVIDER_IDS) + assert all(len(value) == 64 for value in ownership["provider_fingerprints"].values()) + assert [event[0] for event in events] == ["save", "apply"] + + def test_failed_apply_preserves_prior_ownership(self, monkeypatch, tmp_path): + prior = { + "keys": [], + "hermes_home": "/prior/home", + "active_model": {"provider": "prior", "default": "prior-model"}, + } + state = _state(managed_configs={"hermes": prior}) + + def fail(*_args, **_kwargs): + raise RuntimeError("apply failed") + + monkeypatch.setattr(hermes, "apply_config_patch", fail) + saved = [] + monkeypatch.setattr( + "ucode.state.save_state", lambda value: saved.append(copy.deepcopy(value)) + ) + + with pytest.raises(RuntimeError, match="apply failed"): + hermes.write_tool_config(state, hermes_home=tmp_path) + + assert state["managed_configs"]["hermes"] == prior + assert saved[-1]["managed_configs"]["hermes"] == prior + + def test_partial_apply_retains_staged_ownership(self, monkeypatch, tmp_path): + state = _state() + + def fail(*_args, **_kwargs): + raise hermes.HermesConfigApplyError("apply may be partial") + + monkeypatch.setattr(hermes, "apply_config_patch", fail) + saved = [] + monkeypatch.setattr( + "ucode.state.save_state", lambda value: saved.append(copy.deepcopy(value)) + ) + + with pytest.raises(hermes.HermesConfigApplyError, match="may be partial"): + hermes.write_tool_config(state, hermes_home=tmp_path) + + ownership = saved[-1]["managed_configs"]["hermes"] + assert ownership["hermes_home"] == str(tmp_path.resolve()) + assert ownership["provider_fingerprints"] + + +class TestApplyPatch: + def test_dry_run_prints_plan_without_starting_hermes(self, tmp_path, monkeypatch, capsys): + monkeypatch.setattr( + hermes.subprocess, + "run", + lambda *_args, **_kwargs: (_ for _ in ()).throw( + AssertionError("dry run must not start Hermes") + ), + ) + set_dry_run(True) + try: + receipt = hermes.apply_config_patch( + { + "set": {"providers.managed": {"api": "https://example.invalid"}}, + "unset": ["providers.old"], + }, + hermes_home=tmp_path, + ) + finally: + set_dry_run(False) + + assert receipt == { + "status": "dry_run", + "paths": {"set": ["providers.managed"], "unset": ["providers.old"]}, + } + output = capsys.readouterr().out + assert "dry run" in output.lower() + assert "providers.managed" in output + + def test_spawn_failure_is_not_marked_partial(self, tmp_path, monkeypatch): + def fail_to_spawn(_argv, **_kwargs): + raise OSError("binary unavailable") + + monkeypatch.setattr(hermes.subprocess, "run", fail_to_spawn) + + with pytest.raises(RuntimeError) as exc_info: + hermes.apply_config_patch( + {"set": {"providers.first": {"api": "https://first.example"}}, "unset": []}, + hermes_home=tmp_path, + ) + + assert not isinstance(exc_info.value, hermes.HermesConfigApplyError) + + def test_write_failure_after_mutation_is_marked_partial(self, tmp_path, monkeypatch): + writes = 0 + + def fake_run(argv, **_kwargs): + nonlocal writes + if argv[2] == "get": + return hermes.subprocess.CompletedProcess(argv, 1, stdout="", stderr="missing") + if argv[2] == "set": + writes += 1 + return hermes.subprocess.CompletedProcess( + argv, 0 if writes == 1 else 1, stdout="", stderr="secret output" + ) + raise AssertionError(argv) + + monkeypatch.setattr(hermes.subprocess, "run", fake_run) + + with pytest.raises(hermes.HermesConfigApplyError, match="write failed"): + hermes.apply_config_patch( + { + "set": { + "providers.first": {"api": "https://first.example"}, + "providers.second": {"api": "https://second.example"}, + }, + "unset": [], + }, + hermes_home=tmp_path, + ) + + def test_uses_public_commands_and_activates_model_last(self, tmp_path, monkeypatch): + calls = [] + values = {} + + def fake_run(argv, **kwargs): + calls.append((argv, kwargs)) + command, key = argv[2:4] + if command == "get": + if key not in values: + return hermes.subprocess.CompletedProcess(argv, 1, stdout="", stderr="missing") + return hermes.subprocess.CompletedProcess( + argv, 0, stdout=json.dumps(values[key]), stderr="" + ) + if command == "set": + try: + values[key] = json.loads(argv[4]) + except json.JSONDecodeError: + values[key] = argv[4] + return hermes.subprocess.CompletedProcess(argv, 0, stdout="ok", stderr="") + if command == "unset": + values.pop(key, None) + return hermes.subprocess.CompletedProcess(argv, 0, stdout="ok", stderr="") + raise AssertionError(argv) + + monkeypatch.setattr(hermes.subprocess, "run", fake_run) + + receipt = hermes.apply_config_patch( + hermes.render_config_patch(_state()), hermes_home=tmp_path + ) + + assert receipt["status"] == "applied" + mutating = [argv for argv, _kwargs in calls if argv[2] in {"set", "unset"}] + assert all(argv[2] != "apply" for argv, _ in calls) + assert [argv[3] for argv in mutating][-2:] == ["model.provider", "model.default"] + assert [argv[4] for argv in mutating][-2:] == [PROVIDER_ID, MODEL] + assert all(kwargs["shell"] is False for _argv, kwargs in calls) + assert all( + kwargs["env"]["HERMES_HOME"] == str(tmp_path.resolve()) for _argv, kwargs in calls + ) + + def test_conditional_delete_preserves_changed_value(self, tmp_path, monkeypatch): + path = f"providers.{PROVIDER_ID}" + replacement = {"api": "https://user.example/v1"} + calls = [] + + def fake_run(argv, **kwargs): + calls.append(argv) + if argv[1:3] == ["config", "get"]: + return hermes.subprocess.CompletedProcess( + argv, 0, stdout=json.dumps(replacement), stderr="" + ) + raise AssertionError("changed value must not be unset") + + monkeypatch.setattr(hermes.subprocess, "run", fake_run) + + receipt = hermes.apply_config_patch( + { + "set": {}, + "unset": [], + "unset_if_hash": {path: hermes.config_value_fingerprint({"api": "owned"})}, + }, + hermes_home=tmp_path, + ) + + assert receipt["paths"]["unset"] == [] + assert calls == [["hermes", "config", "get", path, "--json"]] + + def test_conditional_delete_rechecks_value_before_unset(self, tmp_path, monkeypatch): + path = f"providers.{PROVIDER_ID}" + owned = {"api": "https://owned.example/v1"} + replacement = {"api": "https://user.example/v1"} + reads = iter((owned, replacement)) + + def fake_run(argv, **kwargs): + if argv[1:3] == ["config", "get"]: + return hermes.subprocess.CompletedProcess( + argv, 0, stdout=json.dumps(next(reads)), stderr="" + ) + raise AssertionError("replacement must not be unset") + + monkeypatch.setattr(hermes.subprocess, "run", fake_run) + + receipt = hermes.apply_config_patch( + { + "set": {}, + "unset": [], + "unset_if_hash": {path: hermes.config_value_fingerprint(owned)}, + }, + hermes_home=tmp_path, + ) + + assert receipt["paths"]["unset"] == [] + + def test_unconditional_delete_rechecks_expected_hash_before_unset(self, tmp_path, monkeypatch): + path = f"providers.{PROVIDER_ID}" + owned = {"api": "https://owned.example/v1"} + replacement = {"api": "https://user.example/v1"} + reads = iter((owned, replacement)) + + def fake_run(argv, **_kwargs): + if argv[1:3] == ["config", "get"]: + return hermes.subprocess.CompletedProcess( + argv, 0, stdout=json.dumps(next(reads)), stderr="" + ) + raise AssertionError("replacement must not be unset") + + monkeypatch.setattr(hermes.subprocess, "run", fake_run) + + receipt = hermes.apply_config_patch( + { + "set": {}, + "unset": [path], + "expect_hashes": {path: hermes.config_value_fingerprint(owned)}, + }, + hermes_home=tmp_path, + ) + + assert receipt["paths"]["unset"] == [] + + def test_expect_missing_rejects_existing_mcp_without_writing(self, tmp_path, monkeypatch): + path = "mcp_servers.databricks-system-ai" + + def fake_run(argv, **kwargs): + if argv[1:3] == ["config", "get"]: + return hermes.subprocess.CompletedProcess( + argv, 0, stdout=json.dumps({"command": "user-server"}), stderr="" + ) + raise AssertionError("collision must fail before mutation") + + monkeypatch.setattr(hermes.subprocess, "run", fake_run) + + with pytest.raises(RuntimeError, match="already exists"): + hermes.apply_config_patch( + { + "set": {path: {"command": "ucode", "args": []}}, + "unset": [], + "expect_missing": [path], + }, + hermes_home=tmp_path, + ) + + def test_failed_public_command_does_not_echo_child_output(self, tmp_path, monkeypatch): + def fake_run(argv, **kwargs): + return hermes.subprocess.CompletedProcess( + argv, + 2, + stdout="sentinel-access-token", + stderr="sentinel-refresh-token", + ) + + monkeypatch.setattr(hermes.subprocess, "run", fake_run) + + try: + hermes.apply_config_patch(hermes.render_config_patch(_state()), hermes_home=tmp_path) + except RuntimeError as exc: + message = str(exc) + else: + raise AssertionError("failed config command must raise") + + assert "sentinel-access-token" not in message + assert "sentinel-refresh-token" not in message + + +class TestManagedMcp: + PROXY_ARGV = [ + "ucode", + "mcp-proxy", + "--url", + f"{WS}/api/2.0/mcp/functions/system/ai", + "--profile", + "team prod", + ] + + def test_entry_uses_stdio_proxy_and_contains_no_bearer(self): + patch = hermes.render_mcp_server_patch("databricks-system-ai", self.PROXY_ARGV) + + assert patch == { + "set": { + "mcp_servers.databricks-system-ai": { + "command": "ucode", + "args": self.PROXY_ARGV[1:], + } + }, + "unset": [], + "expect_missing": ["mcp_servers.databricks-system-ai"], + } + assert "bearer" not in json.dumps(patch).lower() + + def test_configure_scopes_apply_to_active_home(self, tmp_path, monkeypatch): + monkeypatch.setenv("HERMES_HOME", str(tmp_path / "profile-a")) + calls = [] + + def fake_apply(patch, *, hermes_home): + calls.append((patch, hermes_home)) + return {"status": "applied"} + + monkeypatch.setattr(hermes, "apply_config_patch", fake_apply) + + hermes.write_mcp_server_config("databricks-system-ai", self.PROXY_ARGV) + + assert calls[0][1] == str(tmp_path / "profile-a") + assert calls[0][0]["expect_missing"] == ["mcp_servers.databricks-system-ai"] + + def test_managed_update_requires_prior_fingerprint(self, tmp_path, monkeypatch): + calls = [] + monkeypatch.setattr( + hermes, + "apply_config_patch", + lambda patch, *, hermes_home: calls.append(patch) or {"status": "applied"}, + ) + + hermes.write_mcp_server_config( + "databricks-system-ai", + self.PROXY_ARGV, + hermes_home=tmp_path, + expected_fingerprint="a" * 64, + ) + + assert calls[0]["expect_hashes"] == {"mcp_servers.databricks-system-ai": "a" * 64} + + def test_remove_only_unsets_named_managed_entry(self, tmp_path, monkeypatch): + monkeypatch.setenv("HERMES_HOME", str(tmp_path / "profile-b")) + value = hermes.mcp_value_for_argv(self.PROXY_ARGV) + calls = [] + monkeypatch.setattr( + hermes, + "apply_config_patch", + lambda patch, *, hermes_home: ( + calls.append((patch, hermes_home)) + or { + "status": "applied", + "paths": {"unset": ["mcp_servers.databricks-system-ai"]}, + } + ), + ) + + assert ( + hermes.remove_mcp_server_config( + "databricks-system-ai", + expected_fingerprint=hermes.mcp_server_fingerprint(value), + ) + is True + ) + assert calls == [ + ( + { + "set": {}, + "unset": [], + "unset_if_hash": { + "mcp_servers.databricks-system-ai": hermes.mcp_server_fingerprint(value) + }, + }, + tmp_path / "profile-b", + ) + ] + + def test_remove_preserves_user_replacement(self, tmp_path, monkeypatch): + managed = hermes.mcp_value_for_argv(self.PROXY_ARGV) + calls = [] + monkeypatch.setattr( + hermes, + "apply_config_patch", + lambda patch, **kwargs: ( + calls.append(patch) or {"status": "applied", "paths": {"unset": []}} + ), + ) + + assert ( + hermes.remove_mcp_server_config( + "databricks-system-ai", + hermes_home=tmp_path, + expected_fingerprint=hermes.mcp_server_fingerprint(managed), + ) + is False + ) + assert calls[0]["unset_if_hash"] + + +class TestModelSelection: + def test_default_model_falls_back_across_supported_families(self): + assert ( + hermes.default_model( + _state( + codex_models=[], + claude_models={"sonnet": "system.ai.claude-sonnet-4-6"}, + ) + ) + == "system.ai.claude-sonnet-4-6" + ) + + def test_default_model_reuses_codex_discovery_semantics(self): + assert hermes.default_model(_state()) == MODEL diff --git a/tests/test_agents_init.py b/tests/test_agents_init.py index 1b98ce57..2e9345e9 100644 --- a/tests/test_agents_init.py +++ b/tests/test_agents_init.py @@ -13,6 +13,7 @@ TOOL_SPECS, check_gateway_endpoint, configure_selected_tools, + configure_tool, default_model_for_tool, ensure_tool_binary_available, explicit_model_arg_value, @@ -70,7 +71,15 @@ def test_passthrough_for_unrelated_error(self): class TestToolSpecs: def test_all_tools_present(self): - assert set(TOOL_SPECS) == {"codex", "claude", "gemini", "opencode", "copilot", "pi"} + assert set(TOOL_SPECS) == { + "codex", + "claude", + "gemini", + "opencode", + "copilot", + "pi", + "hermes", + } def test_each_spec_has_required_keys(self): required = {"binary", "package", "display", "config_path", "backup_path"} @@ -177,6 +186,7 @@ class TestNormalizeTool: ("opencode", "opencode"), ("copilot", "copilot"), ("pi", "pi"), + ("hermes", "hermes"), ("CODEX", "codex"), (" Claude ", "claude"), ], @@ -185,8 +195,9 @@ def test_known_aliases(self, alias, expected): assert normalize_tool(alias) == expected def test_unknown_raises(self): - with pytest.raises(RuntimeError, match="Unsupported"): + with pytest.raises(RuntimeError, match="Unsupported") as excinfo: normalize_tool("unknown-agent") + assert "hermes" in str(excinfo.value) class TestCheckGatewayEndpoint: @@ -200,6 +211,19 @@ def test_claude_unavailable_when_no_models(self): def test_codex_available(self): assert check_gateway_endpoint({"codex_models": ["model-a"]}, "codex") is True + @pytest.mark.parametrize( + ("key", "value"), + [ + ("codex_models", ["model-a"]), + ("claude_models", {"sonnet": "model-a"}), + ("gemini_models", ["model-a"]), + ("oss_models", ["model-a"]), + ], + ) + def test_hermes_available_for_every_supported_family(self, key, value): + assert check_gateway_endpoint({key: value}, "hermes") is True + assert check_gateway_endpoint({}, "hermes") is False + def test_gemini_available(self): assert check_gateway_endpoint({"gemini_models": ["gemini-2"]}, "gemini") is True @@ -495,6 +519,22 @@ def _boom(*a, **k): # pragma: no cover - must never run class TestInstallToolBinary: + def test_missing_hermes_never_attempts_npm_install(self, monkeypatch): + calls = [] + monkeypatch.setattr("ucode.agents.shutil.which", lambda _: None) + monkeypatch.setattr( + "ucode.agents.subprocess.run", lambda *args, **kwargs: calls.append(args) + ) + + assert install_tool_binary("hermes", strict=False) is False + assert calls == [] + + def test_missing_hermes_has_external_install_guidance(self, monkeypatch): + monkeypatch.setattr("ucode.agents.shutil.which", lambda _: None) + + with pytest.raises(RuntimeError, match="official installer"): + install_tool_binary("hermes", strict=True) + def test_non_strict_returns_false_when_npm_missing(self, monkeypatch): monkeypatch.setattr("ucode.agents.shutil.which", lambda _: None) @@ -666,6 +706,32 @@ def capture_batch(displays): assert batches == [["Codex", "Claude Code"]] + def test_hermes_dispatches_only_to_hermes_writer(self, monkeypatch): + state = { + "workspace": "https://x.databricks.com", + "codex_models": ["system.ai.gpt-5"], + } + calls: list[str] = [] + for name in ("codex", "claude", "gemini", "opencode", "copilot", "pi", "hermes"): + module = getattr(agents_mod, name) + monkeypatch.setattr( + module, + "write_tool_config", + lambda *args, _name=name, **kwargs: calls.append(_name) or state, + ) + + result = configure_tool("hermes", state, "system.ai.gpt-5") + + assert result is state + assert calls == ["hermes"] + + def test_hermes_missing_model_error_is_hermes_specific(self): + with pytest.raises( + RuntimeError, + match="A Hermes model must be selected before configuration", + ): + configure_tool("hermes", {"workspace": "https://x.databricks.com"}) + def test_merges_with_existing_available_tools(self, monkeypatch): """Configuring a new tool should not drop previously-configured tools from state['available_tools'].""" diff --git a/tests/test_cli.py b/tests/test_cli.py index 9e3e05df..1f80b764 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -18,6 +18,7 @@ from typer.testing import CliRunner import ucode.databricks as db_mod +from ucode.agents import hermes as hermes_agent from ucode.cli import app from ucode.databricks import GatewayProbe @@ -114,6 +115,7 @@ def test_configure_help_lists_agents_flag(self): flat = re.sub(r"[│╭╮╯╰─\s]+", " ", output) assert "--agents" in output assert "comma-separated list of agents" in flat + assert "copilot, pi, hermes" in flat assert "--workspaces" in output @@ -323,6 +325,466 @@ def package_version(distribution_name: str) -> str: assert _installed_cli_distribution() == "ucode" +class TestConfigureHermes: + def test_receipt_file_keeps_interactive_output_visible(self, tmp_path): + receipt_path = tmp_path / "hermes-receipt.json" + result_payload = { + "status": "configured", + "agent": "hermes", + "hermes_home": str(tmp_path.resolve()), + "provider_ids": ["ucode-databricks-codex"], + "default_provider": "ucode-databricks-codex", + "default_model": "system.ai.gpt-5", + "mcp_servers_configured": [], + "warnings": [], + } + + state = { + "workspace": "https://example.databricks.com", + "profile": "interactive", + "codex_models": ["system.ai.gpt-5"], + } + + def configure(*_args, **_kwargs): + print("visible interactive setup") + return state + + with ( + patch( + "ucode.cli._prompt_for_configuration", + return_value=(state["workspace"], state["profile"]), + ), + patch("ucode.cli.configure_shared_state", side_effect=configure), + patch("ucode.cli.hermes_agent.apply_config_patch", return_value={"status": "applied"}), + ): + result = runner.invoke( + app, + [ + "configure", + "hermes", + "--hermes-home", + str(tmp_path), + "--receipt-file", + str(receipt_path), + ], + ) + + assert result.exit_code == 0 + assert "visible interactive setup" in result.output + assert "Hermes configured" in result.output + assert json.loads(receipt_path.read_text(encoding="utf-8")) == result_payload + + def test_receipt_is_published_only_after_state_and_config_succeed(self, tmp_path): + receipt_path = tmp_path / "hermes-receipt.json" + state = { + "workspace": "https://example.databricks.com", + "profile": "team-prod", + "codex_models": ["system.ai.gpt-5"], + } + events = [] + + def save(value): + events.append(("save", value)) + + def apply(_patch, *, hermes_home): + assert events and events[0][0] == "save" + assert events[0][1]["managed_configs"]["hermes"]["hermes_home"] == str( + Path(hermes_home).resolve() + ) + assert not receipt_path.exists() + events.append(("apply", hermes_home)) + + with ( + patch("ucode.cli.configure_shared_state", return_value=state), + patch("ucode.cli.save_state", side_effect=save), + patch("ucode.cli.hermes_agent.apply_config_patch", side_effect=apply), + ): + result = runner.invoke( + app, + [ + "configure", + "hermes", + "--workspace", + state["workspace"], + "--profile", + state["profile"], + "--receipt-file", + str(receipt_path), + ], + ) + + assert result.exit_code == 0 + assert [event[0] for event in events] == ["save", "apply"] + assert json.loads(receipt_path.read_text())["status"] == "configured" + + def test_partial_hermes_apply_keeps_staged_ownership_without_receipt(self, tmp_path): + receipt_path = tmp_path / "hermes-receipt.json" + state = { + "workspace": "https://example.databricks.com", + "profile": "team-prod", + "codex_models": ["system.ai.gpt-5"], + } + saved = [] + + with ( + patch("ucode.cli.configure_shared_state", return_value=state), + patch("ucode.cli.save_state", side_effect=lambda value: saved.append(value)), + patch( + "ucode.cli.hermes_agent.apply_config_patch", + side_effect=hermes_agent.HermesConfigApplyError("apply may be partial"), + ), + ): + result = runner.invoke( + app, + [ + "configure", + "hermes", + "--workspace", + state["workspace"], + "--profile", + state["profile"], + "--receipt-file", + str(receipt_path), + ], + ) + + assert result.exit_code == 1 + assert len(saved) == 1 + assert saved[0]["managed_configs"]["hermes"]["provider_fingerprints"] + assert not receipt_path.exists() + + def test_failed_setup_removes_stale_receipt(self, tmp_path): + receipt_path = tmp_path / "hermes-receipt.json" + receipt_path.write_text('{"status":"configured"}', encoding="utf-8") + with patch( + "ucode.cli.configure_hermes_command", + side_effect=RuntimeError("setup failed"), + ): + result = runner.invoke( + app, + [ + "configure", + "hermes", + "--receipt-file", + str(receipt_path), + ], + ) + + assert result.exit_code == 1 + assert not receipt_path.exists() + + def test_stale_receipt_removal_failure_is_controlled(self, tmp_path, monkeypatch): + receipt_path = tmp_path / "hermes-receipt.json" + receipt_path.write_text('{"status":"configured"}', encoding="utf-8") + original_unlink = Path.unlink + + def fail_target(path, *args, **kwargs): + if path == receipt_path: + raise OSError("permission denied") + return original_unlink(path, *args, **kwargs) + + monkeypatch.setattr(Path, "unlink", fail_target) + with patch("ucode.cli.configure_hermes_command") as configure: + result = runner.invoke( + app, + [ + "configure", + "hermes", + "--receipt-file", + str(receipt_path), + ], + ) + + assert result.exit_code == 1 + assert "Unable to prepare completion receipt" in result.output + assert "Traceback" not in result.output + configure.assert_not_called() + + def test_unwritable_receipt_never_starts_hermes_transaction(self, tmp_path, monkeypatch): + receipt_path = tmp_path / "hermes-receipt.json" + state = { + "workspace": "https://example.databricks.com", + "profile": "team-prod", + "codex_models": ["system.ai.gpt-5"], + } + with ( + patch("ucode.cli.configure_shared_state", return_value=state), + patch( + "ucode.cli.tempfile.NamedTemporaryFile", + side_effect=OSError("disk full"), + ), + patch("ucode.cli.hermes_agent.apply_config_patch") as apply, + ): + result = runner.invoke( + app, + [ + "configure", + "hermes", + "--workspace", + state["workspace"], + "--profile", + state["profile"], + "--receipt-file", + str(receipt_path), + ], + ) + + assert result.exit_code != 0 + apply.assert_not_called() + + def test_fully_specified_configuration_is_noninteractive_and_json_only(self, tmp_path): + state = { + "workspace": "https://example.databricks.com", + "profile": "team-prod", + "codex_models": ["system.ai.gpt-5"], + "claude_models": {"opus": "system.ai.claude-opus-4-8"}, + "oss_models": ["system.ai.gpt-oss-120b"], + "gemini_models": ["system.ai.gemini-2-5-pro"], + } + receipt = {"status": "applied", "set_count": 3, "unset_count": 0} + with ( + patch("ucode.cli.configure_shared_state", return_value=state) as shared, + patch("ucode.cli.hermes_agent.apply_config_patch", return_value=receipt) as apply, + patch("ucode.cli._prompt_for_configuration") as prompt, + ): + result = runner.invoke( + app, + [ + "configure", + "hermes", + "--workspace", + state["workspace"], + "--profile", + state["profile"], + "--model", + "system.ai.gpt-5", + "--hermes-home", + str(tmp_path), + "--output", + "json", + ], + ) + + assert result.exit_code == 0 + payload = json.loads(result.output) + assert payload == { + "status": "configured", + "agent": "hermes", + "hermes_home": str(tmp_path.resolve()), + "provider_ids": [ + "ucode-databricks-codex", + "ucode-databricks-anthropic", + "ucode-databricks-gemini", + ], + "default_provider": "ucode-databricks-codex", + "default_model": "system.ai.gpt-5", + "mcp_servers_configured": [], + "warnings": [], + } + prompt.assert_not_called() + shared.assert_called_once_with( + state["workspace"], + profile="team-prod", + tools=["hermes"], + force_login=False, + allow_env_bearer=False, + ) + assert apply.call_args.kwargs["hermes_home"] == tmp_path.resolve() + + def test_receipt_reports_actual_non_codex_default_provider(self, tmp_path): + model = "system.ai.deepseek-v3-2" + state = { + "workspace": "https://example.databricks.com", + "profile": "team-prod", + "codex_models": ["system.ai.gpt-5"], + "oss_models": [model], + } + with ( + patch("ucode.cli.configure_shared_state", return_value=state), + patch( + "ucode.cli.hermes_agent.apply_config_patch", + return_value={"status": "applied"}, + ) as apply, + ): + result = runner.invoke( + app, + [ + "configure", + "hermes", + "--workspace", + state["workspace"], + "--profile", + state["profile"], + "--model", + model, + "--hermes-home", + str(tmp_path), + "--output", + "json", + ], + ) + + assert result.exit_code == 0 + payload = json.loads(result.output) + assert payload["default_provider"] == "ucode-databricks-oss" + assert payload["default_model"] == model + rendered_patch = apply.call_args.args[0] + assert payload["provider_ids"] == [ + path.removeprefix("providers.") + for path in rendered_patch["set"] + if path.startswith("providers.") + ] + + def test_rejects_unsupported_hermes_default_model(self, tmp_path): + model = "system.ai.gpt-oss-120b" + state = { + "workspace": "https://example.databricks.com", + "profile": "team-prod", + "codex_models": ["system.ai.gpt-5"], + "oss_models": [model], + } + with ( + patch("ucode.cli.configure_shared_state", return_value=state), + patch( + "ucode.cli.hermes_agent.apply_config_patch", + return_value={"status": "applied"}, + ) as apply, + ): + result = runner.invoke( + app, + [ + "configure", + "hermes", + "--workspace", + state["workspace"], + "--profile", + state["profile"], + "--model", + model, + "--hermes-home", + str(tmp_path), + "--output", + "json", + ], + ) + + assert result.exit_code == 1 + assert "not supported by Hermes" in result.output + apply.assert_not_called() + + def test_interactive_mode_reuses_workspace_prompt_and_never_launches(self, tmp_path): + state = { + "workspace": "https://example.databricks.com", + "profile": "interactive", + "codex_models": ["system.ai.gpt-5"], + } + with ( + patch( + "ucode.cli._prompt_for_configuration", + return_value=(state["workspace"], state["profile"]), + ) as prompt, + patch("ucode.cli.configure_shared_state", return_value=state) as shared, + patch( + "ucode.cli.hermes_agent.apply_config_patch", + return_value={"status": "applied"}, + ), + patch("ucode.cli.launch_agent") as launch, + ): + result = runner.invoke(app, ["configure", "hermes", "--hermes-home", str(tmp_path)]) + + assert result.exit_code == 0 + prompt.assert_called_once_with("hermes") + shared.assert_called_once_with( + state["workspace"], + profile=state["profile"], + tools=["hermes"], + force_login=True, + allow_env_bearer=False, + ) + launch.assert_not_called() + + def test_explicit_setup_disables_ambient_bearer(self, tmp_path): + state = { + "workspace": "https://example.databricks.com", + "profile": "explicit-profile", + "codex_models": ["system.ai.gpt-5"], + } + with ( + patch("ucode.cli.configure_shared_state", return_value=state) as shared, + patch( + "ucode.cli.hermes_agent.apply_config_patch", + return_value={"status": "applied"}, + ), + ): + result = runner.invoke( + app, + [ + "configure", + "hermes", + "--workspace", + state["workspace"], + "--profile", + state["profile"], + "--hermes-home", + str(tmp_path), + ], + env={"DATABRICKS_BEARER": "sentinel-unscoped-bearer"}, + ) + + assert result.exit_code == 0, result.output + shared.assert_called_once_with( + state["workspace"], + profile=state["profile"], + tools=["hermes"], + force_login=False, + allow_env_bearer=False, + ) + + def test_failed_transaction_has_no_success_receipt_or_secret_leak(self, tmp_path): + state = { + "workspace": "https://example.databricks.com", + "profile": "team-prod", + "codex_models": ["system.ai.gpt-5"], + "access_token": "sentinel-access-token", + } + with ( + patch("ucode.cli.configure_shared_state", return_value=state), + patch( + "ucode.cli.hermes_agent.apply_config_patch", + side_effect=RuntimeError("Hermes configuration failed with exit code 2."), + ), + ): + result = runner.invoke( + app, + [ + "configure", + "hermes", + "--workspace", + state["workspace"], + "--profile", + state["profile"], + "--hermes-home", + str(tmp_path), + "--output", + "json", + ], + ) + + assert result.exit_code == 1 + assert "sentinel-access-token" not in result.output + assert '"status": "configured"' not in result.output + + def test_cancellation_never_starts_a_hermes_transaction(self, tmp_path): + with ( + patch("ucode.cli._prompt_for_configuration", side_effect=KeyboardInterrupt), + patch("ucode.cli.hermes_agent.apply_config_patch") as apply, + ): + result = runner.invoke(app, ["configure", "hermes", "--hermes-home", str(tmp_path)]) + + assert result.exit_code == 130 + apply.assert_not_called() + + class TestVersion: @pytest.mark.parametrize("flag", ["--version", "-V"]) def test_prints_version_and_exits(self, flag): @@ -953,8 +1415,8 @@ class TestAuthTokenCommand: @pytest.fixture(autouse=True) def _isolated_bearer(self): - # The --use-pat path writes DATABRICKS_BEARER directly; restore it so - # writes by code under test don't leak into other tests. + # Keep ambient bearer scenarios isolated even though explicit PAT + # resolution no longer mutates DATABRICKS_BEARER. original = os.environ.pop("DATABRICKS_BEARER", None) yield if original is None: @@ -972,7 +1434,9 @@ def test_prints_only_the_token_to_stdout(self): # Nothing but the bare token (plus trailing newline) may reach stdout, # or the consuming agent will treat the noise as part of the token. assert result.stdout == "tok-123\n" - fetch.assert_called_once_with("https://ws", None, force_refresh=False) + fetch.assert_called_once_with( + "https://ws", None, force_refresh=False, allow_env_bearer=False + ) def test_host_and_profile_override_state(self): with ( @@ -983,7 +1447,9 @@ def test_host_and_profile_override_state(self): app, ["auth-token", "--host", "https://override", "--profile", "prod"] ) assert result.exit_code == 0 - fetch.assert_called_once_with("https://override", "prod", force_refresh=False) + fetch.assert_called_once_with( + "https://override", "prod", force_refresh=False, allow_env_bearer=False + ) def test_force_refresh_is_forwarded(self): with ( @@ -992,7 +1458,91 @@ def test_force_refresh_is_forwarded(self): ): result = runner.invoke(app, ["auth-token", "--force-refresh"]) assert result.exit_code == 0 - fetch.assert_called_once_with("https://ws", None, force_refresh=True) + fetch.assert_called_once_with( + "https://ws", None, force_refresh=True, allow_env_bearer=False + ) + + def test_explicit_identity_does_not_inherit_saved_pat_mode(self): + with ( + patch( + "ucode.cli.load_state", + return_value={ + "workspace": "https://saved", + "profile": "saved-pat-profile", + "use_pat": True, + }, + ), + patch("ucode.cli.ensure_pat_bearer", return_value=True) as ensure_pat, + patch("ucode.cli.get_databricks_token", return_value="oauth-token") as fetch, + ): + result = runner.invoke( + app, + ["auth-token", "--host", "https://override", "--profile", "oauth-profile"], + ) + + assert result.exit_code == 0 + assert result.stdout == "oauth-token\n" + ensure_pat.assert_not_called() + fetch.assert_called_once_with( + "https://override", + "oauth-profile", + force_refresh=False, + allow_env_bearer=False, + ) + + def test_explicit_host_does_not_inherit_saved_profile(self): + with ( + patch( + "ucode.cli.load_state", + return_value={"workspace": "https://saved", "profile": "saved-profile"}, + ), + patch("ucode.cli.get_databricks_token", return_value="token") as fetch, + ): + result = runner.invoke(app, ["auth-token", "--host", "https://override"]) + + assert result.exit_code == 0 + fetch.assert_called_once_with( + "https://override", None, force_refresh=False, allow_env_bearer=False + ) + + def test_explicit_env_bearer_opt_in(self, monkeypatch): + monkeypatch.setenv("DATABRICKS_BEARER", "ci-token") + with ( + patch("ucode.cli.load_state", return_value={"workspace": "https://saved"}), + patch("ucode.cli.get_databricks_token", return_value="ci-token") as fetch, + ): + result = runner.invoke( + app, + [ + "auth-token", + "--host", + "https://ci-workspace", + "--allow-env-bearer", + ], + ) + + assert result.exit_code == 0 + assert result.stdout == "ci-token\n" + fetch.assert_called_once_with( + "https://ci-workspace", None, force_refresh=False, allow_env_bearer=True + ) + + def test_saved_identity_env_bearer_opt_in(self, monkeypatch): + monkeypatch.setenv("DATABRICKS_BEARER", "ci-token") + with ( + patch( + "ucode.cli.load_state", + return_value={"workspace": "https://saved", "profile": "saved-profile"}, + ), + patch("ucode.cli.get_databricks_token", return_value="ci-token") as fetch, + ): + result = runner.invoke(app, ["auth-token", "--allow-env-bearer"]) + + assert result.exit_code == 0 + assert result.stdout == "ci-token\n" + fetch.assert_called_once_with( + "https://saved", "saved-profile", force_refresh=False, allow_env_bearer=True + ) def test_errors_without_workspace(self): with patch("ucode.cli.load_state", return_value={}): @@ -1006,42 +1556,40 @@ def test_hidden_from_top_level_help(self): assert "auth-token" not in _strip_ansi(result.output) def test_use_pat_emits_resolved_pat(self, monkeypatch): - # --use-pat reads the profile's static PAT, exports it as - # DATABRICKS_BEARER, and get_databricks_token returns it directly. + # An explicit PAT identity is resolved directly; it must not mutate the + # process-global bearer or fall through to OAuth token minting. monkeypatch.delenv("DATABRICKS_BEARER", raising=False) - monkeypatch.setattr("ucode.databricks.resolve_pat_token", lambda p: "dapi-pat") + monkeypatch.setattr("ucode.cli.resolve_pat_token", lambda p: "dapi-pat") with ( patch("ucode.cli.load_state", return_value={"workspace": "https://ws"}), - patch( - "ucode.cli.get_databricks_token", - side_effect=lambda w, p, **_kwargs: os.environ.get("DATABRICKS_BEARER", ""), - ), + patch("ucode.cli.ensure_pat_bearer") as ensure_pat, + patch("ucode.cli.get_databricks_token") as fetch, ): result = runner.invoke(app, ["auth-token", "--use-pat", "--profile", "p"]) assert result.exit_code == 0 assert result.stdout == "dapi-pat\n" + assert "DATABRICKS_BEARER" not in os.environ + ensure_pat.assert_not_called() + fetch.assert_not_called() def test_use_pat_ignores_empty_bearer_env(self, monkeypatch): - # A stray empty DATABRICKS_BEARER must not shadow the PAT and force the - # OAuth path (the regression that motivated ensure_pat_bearer). + # Empty ambient state is irrelevant to an explicitly scoped PAT. monkeypatch.setenv("DATABRICKS_BEARER", "") - monkeypatch.setattr("ucode.databricks.resolve_pat_token", lambda p: "dapi-pat") + monkeypatch.setattr("ucode.cli.resolve_pat_token", lambda p: "dapi-pat") with ( patch("ucode.cli.load_state", return_value={"workspace": "https://ws"}), - patch( - "ucode.cli.get_databricks_token", - side_effect=lambda w, p, **_kwargs: os.environ.get("DATABRICKS_BEARER", ""), - ), + patch("ucode.cli.get_databricks_token") as fetch, ): result = runner.invoke(app, ["auth-token", "--use-pat", "--profile", "p"]) assert result.exit_code == 0 assert result.stdout == "dapi-pat\n" + fetch.assert_not_called() def test_use_pat_fails_closed_without_pat(self, monkeypatch): # --use-pat with no resolvable PAT must error, NOT fall through to OAuth # (which can't serve a PAT-only profile and yields a misleading message). monkeypatch.delenv("DATABRICKS_BEARER", raising=False) - monkeypatch.setattr("ucode.databricks.resolve_pat_token", lambda p: None) + monkeypatch.setattr("ucode.cli.resolve_pat_token", lambda p: None) with ( patch("ucode.cli.load_state", return_value={"workspace": "https://ws"}), patch("ucode.cli.get_databricks_token", return_value="oauth-tok") as fetch, @@ -1052,20 +1600,20 @@ def test_use_pat_fails_closed_without_pat(self, monkeypatch): fetch.assert_not_called() assert result.stdout == "" - def test_use_pat_honors_non_empty_bearer_env(self, monkeypatch): - # A real pre-set bearer (CI escape hatch) wins over the profile PAT. + def test_explicit_pat_ignores_non_empty_bearer_env(self, monkeypatch): + # An ambient token for another workspace must not replace the PAT tied + # to the explicitly requested profile. monkeypatch.setenv("DATABRICKS_BEARER", "ci-bearer") - monkeypatch.setattr("ucode.databricks.resolve_pat_token", lambda p: "dapi-pat") + monkeypatch.setattr("ucode.cli.resolve_pat_token", lambda p: "dapi-pat") with ( patch("ucode.cli.load_state", return_value={"workspace": "https://ws"}), - patch( - "ucode.cli.get_databricks_token", - side_effect=lambda w, p, **_kwargs: os.environ.get("DATABRICKS_BEARER", ""), - ), + patch("ucode.cli.get_databricks_token") as fetch, ): result = runner.invoke(app, ["auth-token", "--use-pat", "--profile", "p"]) assert result.exit_code == 0 - assert result.stdout == "ci-bearer\n" + assert result.stdout == "dapi-pat\n" + assert os.environ["DATABRICKS_BEARER"] == "ci-bearer" + fetch.assert_not_called() class TestStatus: @@ -1613,6 +2161,105 @@ def test_reverts_mcp_configs_before_clearing_state(self): assert cleared == [True] assert "Claude Code MCP config: restored" in result.output + def test_hermes_cleanup_is_surgical_and_uses_recorded_home(self): + ownership = { + "keys": [], + "hermes_home": "/profiles/original", + "active_model": { + "provider": "ucode-databricks-codex", + "default": "system.ai.gpt-5", + }, + "provider_fingerprints": {"ucode-databricks-codex": "a" * 64}, + } + state = { + **MINIMAL_STATE, + "available_tools": ["hermes"], + "managed_configs": {"hermes": ownership}, + } + with ( + patch("ucode.cli.load_state", return_value=state), + patch("ucode.cli.revert_mcp_configs", return_value={}), + patch("ucode.cli.restore_file", return_value=False) as restore, + patch("ucode.cli.hermes_agent.unconfigure", return_value={}) as unconfigure, + patch("ucode.cli.clear_state"), + ): + result = runner.invoke(app, ["revert"]) + + assert result.exit_code == 0, result.output + unconfigure.assert_called_once_with( + hermes_home="/profiles/original", + owned_model=ownership["active_model"], + owned_provider_fingerprints=ownership["provider_fingerprints"], + ) + assert all( + call.args[0] != hermes_agent.SPEC["config_path"] for call in restore.call_args_list + ) + + def test_hermes_cleanup_failure_preserves_state_for_retry(self): + ownership = { + "keys": [], + "hermes_home": "/profiles/original", + "active_model": { + "provider": "ucode-databricks-codex", + "default": "system.ai.gpt-5", + }, + "provider_fingerprints": {"ucode-databricks-codex": "a" * 64}, + } + state = { + **MINIMAL_STATE, + "available_tools": ["hermes"], + "managed_configs": {"hermes": ownership}, + } + with ( + patch("ucode.cli.load_state", return_value=state), + patch("ucode.cli.revert_mcp_configs", return_value={}), + patch("ucode.cli.restore_file", return_value=False), + patch("ucode.cli.hermes_agent.unconfigure", side_effect=RuntimeError("cleanup failed")), + patch("ucode.cli.clear_state") as clear, + ): + result = runner.invoke(app, ["revert"]) + + assert result.exit_code == 1 + clear.assert_not_called() + + +class TestRegisterManagedMcpServers: + def test_failed_last_server_removal_preserves_retry_state(self, monkeypatch): + import ucode.cli as cli_mod + + previous = [{"name": "old", "clients": ["hermes"]}] + state = { + "workspace": "https://example.databricks.com", + "managed_mcp_servers": previous.copy(), + } + monkeypatch.setattr( + cli_mod, + "apply_managed_mcp_servers", + lambda *args, **kwargs: (_ for _ in ()).throw(RuntimeError("remove failed")), + ) + monkeypatch.setattr(cli_mod, "print_warning", lambda _message: None) + + cli_mod._register_managed_mcp_servers({}, "hermes", state) + + assert state["managed_mcp_servers"] == previous + + def test_successful_last_server_removal_clears_owned_state(self, monkeypatch): + import ucode.cli as cli_mod + + state = { + "workspace": "https://example.databricks.com", + "managed_mcp_servers": [{"name": "old", "clients": ["hermes"]}], + } + monkeypatch.setattr( + cli_mod, + "apply_managed_mcp_servers", + lambda *args, **kwargs: [], + ) + + cli_mod._register_managed_mcp_servers({}, "hermes", state) + + assert state["managed_mcp_servers"] == [] + class TestDoctorCommand: def test_invokes_doctor(self): @@ -1966,6 +2613,55 @@ def test_no_extra_args_passes_empty_list(self): class TestConfigureAgentFlag: + @pytest.mark.parametrize( + ("tools", "expected"), + [ + (["hermes"], False), + (["hermes", "claude"], False), + (["claude"], True), + (None, True), + ], + ) + def test_shared_workspace_auth_policy_fails_closed_for_explicit_hermes(self, tools, expected): + from ucode import cli as cli_mod + + with patch("ucode.cli.configure_shared_state", return_value={}) as shared: + cli_mod._configure_shared_workspace_states( + [("https://sentinel.example", "sentinel-profile")], + tools, + force_login=False, + ) + + assert shared.call_args.kwargs["allow_env_bearer"] is expected + + def test_agents_flag_routes_hermes_through_selected_tools(self): + with ( + patch("ucode.cli.install_databricks_cli"), + patch("ucode.cli.install_tool_binary"), + patch("ucode.cli.configure_workspace_command") as mock_cfg, + ): + result = runner.invoke(app, ["configure", "--agents", "hermes"]) + + assert result.exit_code == 0, result.output + mock_cfg.assert_called_once_with( + selected_tools=["hermes"], + prompt_optional_updates=True, + ) + + def test_hermes_does_not_offer_unvalidated_mps_picker(self): + from ucode import cli as cli_mod + + state = {"workspace": "https://x.databricks.com"} + with ( + patch("ucode.cli.get_databricks_token") as token, + patch("ucode.cli.prompt_for_selection") as prompt, + ): + result = cli_mod._maybe_select_provider_service("hermes", state) + + assert result is state + token.assert_not_called() + prompt.assert_not_called() + def test_no_flag_calls_configure_all(self): with ( patch("ucode.cli.install_databricks_cli"), @@ -2542,6 +3238,7 @@ def fake_configure_shared_state( use_pat=False, fable_enabled=None, databricks_ai_tools_enabled=None, + allow_env_bearer=True, ): captured["workspace"] = workspace captured["profile"] = profile @@ -2580,6 +3277,7 @@ def fake_configure_shared_state( use_pat=False, fable_enabled=None, databricks_ai_tools_enabled=None, + allow_env_bearer=True, ): configured_shared.append( (workspace, profile, tuple(tools) if tools is not None else None, force_login) @@ -2844,11 +3542,13 @@ def _stub_deps(monkeypatch, *, pat_token, existing_state=None): monkeypatch.setattr(cli_mod, "save_state", lambda s: saved.append(dict(s))) monkeypatch.setattr(cli_mod, "run_databricks_login", lambda w, p: logins.append((w, p))) monkeypatch.setattr( - cli_mod, "ensure_databricks_auth", lambda w, p=None: ensures.append((w, p)) + cli_mod, + "ensure_databricks_auth", + lambda w, p=None, **_kwargs: ensures.append((w, p)), ) monkeypatch.setattr(cli_mod, "resolve_pat_token", lambda p: pat_token) monkeypatch.setattr(cli_mod, "find_profile_name_for_host", lambda w: None) - monkeypatch.setattr(cli_mod, "get_databricks_token", lambda w, p: "token") + monkeypatch.setattr(cli_mod, "get_databricks_token", lambda w, p, **_kwargs: "token") monkeypatch.setattr( cli_mod, "probe_unity_gateway_capabilities", lambda w, t: MODEL_SERVICE_PROBE ) @@ -2883,6 +3583,57 @@ def test_happy_path_prints_success_without_model_service_detail(self, monkeypatc assert "Unity AI Gateway connected" in output assert "Model service:" not in output + @pytest.mark.parametrize("allow_env_bearer", [False, True]) + def test_threads_ambient_bearer_policy_to_auth_and_discovery( + self, monkeypatch, allow_env_bearer + ): + cli_mod, *_ = self._stub_deps(monkeypatch, pat_token="sentinel-pat") + auth_policies: list[bool] = [] + token_policies: list[bool] = [] + monkeypatch.setattr( + cli_mod, + "ensure_databricks_auth", + lambda _workspace, _profile=None, *, allow_env_bearer=True: auth_policies.append( + allow_env_bearer + ), + ) + monkeypatch.setattr( + cli_mod, + "get_databricks_token", + lambda _workspace, _profile=None, *, allow_env_bearer=True: ( + token_policies.append(allow_env_bearer) or "sentinel-token" + ), + ) + + cli_mod.configure_shared_state( + self.WS, + profile="explicit-profile", + allow_env_bearer=allow_env_bearer, + ) + + assert auth_policies == [allow_env_bearer] + assert token_policies == [allow_env_bearer] + + def test_legacy_caller_keeps_ambient_bearer_default(self, monkeypatch): + cli_mod, *_ = self._stub_deps(monkeypatch, pat_token="sentinel-pat") + policies: list[bool] = [] + monkeypatch.setattr( + cli_mod, + "ensure_databricks_auth", + lambda _workspace, _profile=None, *, allow_env_bearer=True: policies.append( + allow_env_bearer + ), + ) + monkeypatch.setattr( + cli_mod, + "get_databricks_token", + lambda _workspace, _profile=None, *, allow_env_bearer=True: "sentinel-token", + ) + + cli_mod.configure_shared_state(self.WS, profile="legacy-profile") + + assert policies == [True] + @pytest.mark.parametrize( ("responses", "expected_model_service"), [ @@ -3070,6 +3821,27 @@ def test_codex_only_configure_persists_discovered_oss_models(self, monkeypatch): assert state["codex_models"] == ["system.ai.gpt-5-6-sol"] assert state["oss_models"] == ["system.ai.glm-5-2"] + def test_hermes_discovers_every_supported_model_family(self, monkeypatch): + cli_mod, *_ = self._stub_deps(monkeypatch, pat_token="dapi-pat") + monkeypatch.setattr( + cli_mod, + "discover_model_services", + lambda w, t: ( + {"opus": "system.ai.claude-opus-4-8"}, + ["system.ai.gpt-5"], + ["system.ai.gemini-2-5-pro"], + ["system.ai.gpt-oss-120b"], + None, + ), + ) + + state = cli_mod.configure_shared_state(self.WS, profile="DEFAULT", tools=["hermes"]) + + assert state["claude_models"] == {"opus": "system.ai.claude-opus-4-8"} + assert state["codex_models"] == ["system.ai.gpt-5"] + assert state["gemini_models"] == ["system.ai.gemini-2-5-pro"] + assert state["oss_models"] == ["system.ai.gpt-oss-120b"] + def _stub_with_fable(self, monkeypatch): cli_mod, *_ = self._stub_deps(monkeypatch, pat_token="dapi-pat") monkeypatch.setattr( @@ -3270,9 +4042,9 @@ def _stub_external_deps(monkeypatch): monkeypatch.setattr(cli_mod, "normalize_workspace_url", lambda w: w) monkeypatch.setattr(cli_mod, "run_databricks_login", lambda w, p: None) - monkeypatch.setattr(cli_mod, "ensure_databricks_auth", lambda w, p=None: None) + monkeypatch.setattr(cli_mod, "ensure_databricks_auth", lambda w, p=None, **kwargs: None) monkeypatch.setattr(cli_mod, "find_profile_name_for_host", lambda w: None) - monkeypatch.setattr(cli_mod, "get_databricks_token", lambda w, p: "token") + monkeypatch.setattr(cli_mod, "get_databricks_token", lambda w, p, **kwargs: "token") monkeypatch.setattr( cli_mod, "probe_unity_gateway_capabilities", lambda w, t: MODEL_SERVICE_PROBE ) @@ -3331,10 +4103,10 @@ def _stub(monkeypatch): import ucode.cli as cli_mod monkeypatch.setattr(cli_mod, "normalize_workspace_url", lambda w: w) - monkeypatch.setattr(cli_mod, "ensure_databricks_auth", lambda w, p=None: None) + monkeypatch.setattr(cli_mod, "ensure_databricks_auth", lambda w, p=None, **kwargs: None) monkeypatch.setattr(cli_mod, "run_databricks_login", lambda w, p: None) monkeypatch.setattr(cli_mod, "find_profile_name_for_host", lambda w: None) - monkeypatch.setattr(cli_mod, "get_databricks_token", lambda w, p: "token") + monkeypatch.setattr(cli_mod, "get_databricks_token", lambda w, p, **kwargs: "token") monkeypatch.setattr( cli_mod, "probe_unity_gateway_capabilities", lambda w, t: MODEL_SERVICE_PROBE ) diff --git a/tests/test_databricks.py b/tests/test_databricks.py index 7debda88..d2896b7f 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -271,6 +271,7 @@ def test_buckets_families_by_name(self, monkeypatch): _model_service("system.ai.claude-opus-4-8"), _model_service("system.ai.claude-sonnet-4-6"), _model_service("system.ai.gpt-5"), + _model_service("system.ai.gpt-oss-120b"), _model_service("system.ai.gemini-2-5-flash"), _model_service("system.ai.gemini-3-5-flash"), _model_service("system.ai.kimi-k2-7-code"), @@ -299,6 +300,7 @@ def test_buckets_families_by_name(self, monkeypatch): assert oss == [ "system.ai.deepseek-v4-pro", "system.ai.glm-5-2", + "system.ai.gpt-oss-120b", "system.ai.kimi-k2-7-code", ] @@ -1352,13 +1354,37 @@ def test_noop_without_use_pat(self, monkeypatch): assert "DATABRICKS_BEARER" not in os.environ - def test_existing_bearer_wins(self, monkeypatch): + def test_oauth_state_preserves_unscoped_ci_bearer(self, monkeypatch): + monkeypatch.setenv("DATABRICKS_BEARER", "ci-bearer") + + db_mod.apply_pat_environment({"use_pat": False, "profile": "DEFAULT"}) + + assert os.environ["DATABRICKS_BEARER"] == "ci-bearer" + + def test_configured_pat_replaces_ambient_bearer(self, monkeypatch): monkeypatch.setenv("DATABRICKS_BEARER", "explicit-bearer") monkeypatch.setattr(db_mod, "resolve_pat_token", lambda p: "dapi-pat") db_mod.apply_pat_environment({"use_pat": True, "profile": "DEFAULT"}) - assert os.environ["DATABRICKS_BEARER"] == "explicit-bearer" + assert os.environ["DATABRICKS_BEARER"] == "dapi-pat" + + def test_missing_configured_pat_clears_ambient_bearer(self, monkeypatch): + monkeypatch.setenv("DATABRICKS_BEARER", "other-workspace-token") + monkeypatch.setattr(db_mod, "resolve_pat_token", lambda p: None) + + db_mod.apply_pat_environment({"use_pat": True, "profile": "DEFAULT"}) + + assert "DATABRICKS_BEARER" not in os.environ + + def test_switching_to_oauth_clears_scoped_pat_state(self, monkeypatch): + monkeypatch.setattr(db_mod, "resolve_pat_token", lambda p: "dapi-pat") + db_mod.apply_pat_environment({"workspace": WS, "use_pat": True, "profile": "DEFAULT"}) + + db_mod.apply_pat_environment({"workspace": WS, "use_pat": False, "profile": "DEFAULT"}) + + assert "DATABRICKS_BEARER" not in os.environ + assert db_mod._SCOPED_BEARER is None class TestBuildAuthTokenArgv: @@ -1506,6 +1532,68 @@ def test_includes_stdout_on_failure(self): assert "useful diagnostic output" in formatted assert "no matching profile" in formatted + def test_redacts_tokens_from_failure_output(self): + secret = "sentinel-token-value" + result = subprocess.CompletedProcess( + args=["databricks", "auth", "token"], + returncode=1, + stdout=json.dumps({"access_token": secret, "error": "expired"}), + stderr=f"request failed bearer={secret}", + ) + + formatted = _format_subprocess_result(result) + + assert secret not in formatted + assert "" in formatted + assert "expired" in formatted + + def test_redacts_bearer_headers_and_access_token_prose(self): + access_secret = "sentinel-plain-access-secret" + bearer_secret = "sentinel-bearer-secret" + result = subprocess.CompletedProcess( + args=["databricks", "auth", "token"], + returncode=1, + stdout=f"access token is {access_secret}", + stderr=f"Authorization: Bearer {bearer_secret}", + ) + + formatted = _format_subprocess_result(result) + + assert access_secret not in formatted + assert bearer_secret not in formatted + assert formatted.count("") >= 2 + + +class TestLogAuthDiagnostics: + def test_profiles_stderr_is_scrubbed_before_debug_logging(self, tmp_path, monkeypatch): + secret = "sentinel-profiles-bearer" + results = iter( + [ + subprocess.CompletedProcess(["databricks", "--version"], 0, "v1", ""), + subprocess.CompletedProcess( + ["databricks", "auth", "profiles"], + 1, + "", + f"Authorization: Bearer {secret}", + ), + ] + ) + logged: list[str] = [] + monkeypatch.setattr(db_mod, "_debug_enabled", lambda: True) + monkeypatch.setattr(db_mod, "_debug", lambda label, detail: logged.append(detail)) + monkeypatch.setattr(db_mod.subprocess, "run", lambda *args, **kwargs: next(results)) + monkeypatch.setenv("DATABRICKS_CONFIG_FILE", str(tmp_path / "missing-cfg")) + db_mod._log_auth_diagnostics.cache_clear() + + try: + db_mod._log_auth_diagnostics() + finally: + db_mod._log_auth_diagnostics.cache_clear() + + rendered = "\n".join(logged) + assert secret not in rendered + assert "" in rendered + class TestScrubDatabrickscfg: def test_redacts_token_value(self): @@ -1579,6 +1667,18 @@ def test_passes_through_scalars_and_non_secret_keys(self): class TestGetDatabricksToken: + @pytest.fixture(autouse=True) + def _isolated_bearer(self): + original_bearer = os.environ.pop("DATABRICKS_BEARER", None) + original_scope = db_mod._SCOPED_BEARER + db_mod._SCOPED_BEARER = None + yield + db_mod._SCOPED_BEARER = original_scope + if original_bearer is None: + os.environ.pop("DATABRICKS_BEARER", None) + else: + os.environ["DATABRICKS_BEARER"] = original_bearer + def _fake_databricks(self, tmp_path, script: str) -> dict: fake = tmp_path / "databricks" fake.write_text(f"#!/bin/sh\n{script}\n") @@ -1594,6 +1694,69 @@ def test_returns_token_on_success(self, tmp_path, monkeypatch): token = get_databricks_token(WS) assert token == "good-token" + def test_explicit_opt_in_allows_ci_bearer(self, monkeypatch): + monkeypatch.setenv("DATABRICKS_BEARER", "ci-token") + monkeypatch.setattr(db_mod, "run", lambda *args, **kwargs: pytest.fail("CLI called")) + + assert get_databricks_token(WS, allow_env_bearer=True) == "ci-token" + + def test_ucode_installed_profile_pat_is_allowed_without_global_opt_in(self, monkeypatch): + original_bearer = os.environ.pop("DATABRICKS_BEARER", None) + original_scope = db_mod._SCOPED_BEARER + try: + assert ensure_pat_bearer("pat-profile", "profile-pat", workspace=WS) is True + monkeypatch.setattr(db_mod, "run", lambda *args, **kwargs: pytest.fail("CLI called")) + + assert get_databricks_token(WS, profile="pat-profile") == "profile-pat" + finally: + db_mod._SCOPED_BEARER = original_scope + if original_bearer is None: + os.environ.pop("DATABRICKS_BEARER", None) + else: + os.environ["DATABRICKS_BEARER"] = original_bearer + + def test_ucode_installed_pat_is_rejected_for_another_workspace(self, tmp_path, monkeypatch): + original_bearer = os.environ.pop("DATABRICKS_BEARER", None) + original_scope = db_mod._SCOPED_BEARER + try: + assert ensure_pat_bearer("pat-profile", "profile-pat", workspace=WS) is True + env = self._fake_databricks( + tmp_path, + 'echo \'{"access_token": "workspace-b-token", "token_type": "Bearer"}\'', + ) + monkeypatch.setattr("os.environ", env) + + token = get_databricks_token( + "https://workspace-b.example", + profile="pat-profile", + ) + + assert token == "workspace-b-token" + finally: + db_mod._SCOPED_BEARER = original_scope + if original_bearer is None: + os.environ.pop("DATABRICKS_BEARER", None) + else: + os.environ["DATABRICKS_BEARER"] = original_bearer + + def test_default_preserves_legacy_ci_bearer(self, monkeypatch): + monkeypatch.setenv("DATABRICKS_BEARER", "ci-token") + monkeypatch.setattr(db_mod, "run", lambda *args, **kwargs: pytest.fail("CLI called")) + + assert get_databricks_token(WS, profile="workspace-profile") == "ci-token" + + def test_explicit_scope_ignores_ambient_bearer(self, tmp_path, monkeypatch): + env = self._fake_databricks( + tmp_path, + 'echo \'{"access_token": "workspace-token", "token_type": "Bearer"}\'', + ) + env["DATABRICKS_BEARER"] = "other-workspace-token" + monkeypatch.setattr("os.environ", env) + + token = get_databricks_token(WS, profile="workspace-profile", allow_env_bearer=False) + + assert token == "workspace-token" + def test_strips_ambient_profile_when_profile_not_provided(self, tmp_path, monkeypatch): profile_log = tmp_path / "profile" env = self._fake_databricks( @@ -1624,6 +1787,49 @@ def test_has_valid_auth_strips_ambient_profile_without_explicit_profile( assert db_mod.has_valid_databricks_auth(WS) assert profile_log.read_text() == "" + def test_has_valid_auth_ignores_unscoped_ambient_bearer(self, tmp_path, monkeypatch): + called = tmp_path / "called" + env = self._fake_databricks( + tmp_path, + f"touch {called}\n" + 'echo \'{"access_token": "workspace-token", "token_type": "Bearer"}\'', + ) + env["DATABRICKS_BEARER"] = "other-workspace-token" + monkeypatch.setattr("os.environ", env) + + assert db_mod.has_valid_databricks_auth( + WS, + profile="workspace-profile", + allow_env_bearer=False, + ) + assert called.exists() + + def test_ensure_auth_forwards_fail_closed_ambient_bearer_policy(self, monkeypatch): + policies: list[bool] = [] + logins: list[tuple[str, str | None]] = [] + monkeypatch.setenv("DATABRICKS_BEARER", "sentinel-unscoped-bearer") + monkeypatch.setattr( + db_mod, + "has_valid_databricks_auth", + lambda _workspace, _profile=None, *, allow_env_bearer=True: ( + policies.append(allow_env_bearer) or False + ), + ) + monkeypatch.setattr( + db_mod, + "run_databricks_login", + lambda workspace, profile: logins.append((workspace, profile)), + ) + + db_mod.ensure_databricks_auth( + WS, + "explicit-profile", + allow_env_bearer=False, + ) + + assert policies == [False] + assert logins == [(WS, "explicit-profile")] + def test_reauths_and_retries_when_token_empty(self, tmp_path, monkeypatch): call_count = tmp_path / "calls" call_count.write_text("0") @@ -2765,6 +2971,7 @@ class TestClassifyModelFamily: ("databricks-claude-haiku-4-5", "haiku"), ("system.ai.claude-fable-5", "fable"), ("system.ai.gpt-5-3-codex", "codex"), + ("system.ai.gpt-oss-120b", "oss"), ("system.ai.gemini-3-flash", "gemini"), ("system.ai.kimi-k2-7-code", "oss"), ("system.ai.glm-4-6", "oss"), diff --git a/tests/test_e2e_hermes.py b/tests/test_e2e_hermes.py new file mode 100644 index 00000000..4ac5914f --- /dev/null +++ b/tests/test_e2e_hermes.py @@ -0,0 +1,255 @@ +"""Cross-repository contract test against a real Hermes checkout.""" + +from __future__ import annotations + +import json +import os +import shlex +import subprocess +import sys +from pathlib import Path + +import pytest + +from ucode.agents import hermes + +WORKSPACE = "https://e2e-test.cloud.databricks.com" +MODEL = "system.ai.gpt-5-6" + + +def _hermes_binary() -> Path: + configured = os.environ.get("HERMES_TEST_BINARY") + if configured: + return Path(configured).expanduser().resolve() + checkout = Path(__file__).resolve().parents[2] / "hermes" + venv_binary = checkout / ".venv" / "bin" / "hermes" + return (venv_binary if venv_binary.is_file() else checkout / "hermes").resolve() + + +def _hermes_runtime(binary: Path) -> tuple[Path, Path]: + """Return a Python with Hermes deps and the corresponding source root.""" + configured_root = os.environ.get("HERMES_TEST_ROOT") + candidates = [Path(configured_root).expanduser()] if configured_root else [] + candidates.extend((binary.parent, *binary.parents)) + root = next((path for path in candidates if (path / "hermes_cli").is_dir()), None) + if root is None: + pytest.skip("Hermes source root unavailable; set HERMES_TEST_ROOT") + sibling_python = binary.parent / "python" + python = sibling_python if sibling_python.is_file() else Path(sys.executable) + return python, root + + +def _state() -> dict: + return { + "workspace": WORKSPACE, + "profile": "e2e profile", + "codex_models": [MODEL], + "claude_models": {}, + "gemini_models": [], + "oss_models": [], + } + + +def _runtime_resolution(binary: Path, env: dict[str, str]) -> dict: + python, root = _hermes_runtime(binary) + script = """ +import json +from hermes_cli.config import load_config +from hermes_cli.runtime_provider import resolve_runtime_provider + +config = load_config() +provider_id = config["model"]["provider"] +provider = config["providers"][provider_id] +runtime = resolve_runtime_provider( + requested=provider_id, + target_model=config["model"]["default"], +) +print(json.dumps({ + "provider_id": provider_id, + "model": config["model"]["default"], + "key_cmd": provider["key_cmd"], + "transport": runtime["api_mode"], + "base_url": runtime["base_url"], + "extra_headers": runtime.get("extra_headers"), + "dynamic_key": callable(runtime["api_key"]), +})) +""" + completed = subprocess.run( + [str(python), "-c", script], + cwd=root, + env=env, + capture_output=True, + text=True, + check=False, + ) + assert completed.returncode == 0, completed.stderr + return json.loads(completed.stdout) + + +def _loaded_config(binary: Path, env: dict[str, str]) -> dict: + python, root = _hermes_runtime(binary) + completed = subprocess.run( + [ + str(python), + "-c", + "import json; from hermes_cli.config import load_config; " + "print(json.dumps(load_config()))", + ], + cwd=root, + env=env, + capture_output=True, + text=True, + check=False, + ) + assert completed.returncode == 0, completed.stderr + return json.loads(completed.stdout) + + +def test_real_hermes_apply_runtime_resolution_and_surgical_unconfigure( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + binary = _hermes_binary() + if not binary.is_file(): + pytest.skip("Hermes checkout unavailable; set HERMES_TEST_BINARY to its executable") + if binary.name != "hermes": + pytest.skip("HERMES_TEST_BINARY must point to the Hermes CLI executable") + + home = tmp_path / "home" + hermes_home = tmp_path / "hermes-home" + home.mkdir() + hermes_home.mkdir() + + monkeypatch.setenv("HOME", str(home)) + monkeypatch.setenv("HERMES_HOME", str(hermes_home)) + monkeypatch.setenv("XDG_CONFIG_HOME", str(tmp_path / "xdg-config")) + monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path / "xdg-cache")) + monkeypatch.setenv("XDG_DATA_HOME", str(tmp_path / "xdg-data")) + monkeypatch.setenv("PATH", f"{binary.parent}{os.pathsep}{os.environ.get('PATH', '')}") + + config_path = hermes_home / "config.yaml" + config_path.write_text( + "user_setting: keep\n" + "providers:\n" + " personal-provider:\n" + " api: https://personal.example/v1\n" + " transport: openai_chat\n" + " default_model: personal-model\n" + ) + + configured_state = hermes.write_tool_config( + {**_state(), "gemini_models": ["system.ai.gemini-3-flash"]}, + model=MODEL, + hermes_home=hermes_home, + ) + assert hermes.GEMINI_PROVIDER_ID in _loaded_config(binary, os.environ.copy())["providers"] + + configured_state = hermes.write_tool_config( + {**_state(), "managed_configs": configured_state["managed_configs"]}, + model=MODEL, + hermes_home=hermes_home, + ) + assert hermes.GEMINI_PROVIDER_ID not in _loaded_config(binary, os.environ.copy())["providers"] + + env = os.environ.copy() + resolved = _runtime_resolution(binary, env) + assert resolved["provider_id"] == hermes.PROVIDER_ID + assert resolved["model"] == MODEL + key_cmd = shlex.split(resolved["key_cmd"]) + assert Path(key_cmd[0]).name == "ucode" + assert key_cmd[1:] == [ + "auth-token", + "--host", + WORKSPACE, + "--profile", + "e2e profile", + ] + assert resolved["transport"] == "codex_responses" + assert resolved["base_url"] == f"{WORKSPACE}/ai-gateway/codex/v1" + assert resolved["extra_headers"] in (None, {}) + assert resolved["dynamic_key"] is True + + hermes.apply_config_patch( + { + "set": { + f"providers.{hermes.PROVIDER_ID}": { + "api": "https://user-replacement.example/v1", + "transport": "openai_chat", + "default_model": "replacement-model", + "models": {"replacement-model": {}}, + }, + "model.provider": "personal-provider", + "model.default": "personal-model", + }, + "unset": [], + }, + hermes_home=hermes_home, + ) + ownership = configured_state["managed_configs"]["hermes"] + receipt = hermes.unconfigure( + hermes_home=ownership["hermes_home"], + owned_model=ownership["active_model"], + owned_provider_fingerprints=ownership["provider_fingerprints"], + ) + assert receipt["status"] == "applied" + remaining = _loaded_config(binary, env) + assert remaining["user_setting"] == "keep" + assert remaining["model"] == { + "provider": "personal-provider", + "default": "personal-model", + } + assert remaining["providers"] == { + "personal-provider": { + "api": "https://personal.example/v1", + "transport": "openai_chat", + "default_model": "personal-model", + }, + hermes.PROVIDER_ID: { + "api": "https://user-replacement.example/v1", + "transport": "openai_chat", + "default_model": "replacement-model", + "models": {"replacement-model": {}}, + }, + } + + proxy_argv = ["ucode", "mcp-proxy", "https://managed.example/mcp"] + managed_value = hermes.mcp_value_for_argv(proxy_argv) + fingerprint = hermes.mcp_server_fingerprint(managed_value) + hermes.apply_config_patch( + { + "set": { + "mcp_servers.user-sibling": {"command": "user", "args": ["--keep"]}, + }, + "unset": [], + }, + hermes_home=hermes_home, + ) + hermes.write_mcp_server_config("managed", proxy_argv, hermes_home=hermes_home) + assert ( + hermes.remove_mcp_server_config( + "managed", + hermes_home=hermes_home, + expected_fingerprint=fingerprint, + ) + is True + ) + remaining = _loaded_config(binary, env) + assert remaining["mcp_servers"] == {"user-sibling": {"command": "user", "args": ["--keep"]}} + + hermes.write_mcp_server_config("managed", proxy_argv, hermes_home=hermes_home) + replacement = {"command": "user-replacement", "args": ["--keep"]} + hermes.apply_config_patch( + {"set": {"mcp_servers.managed": replacement}, "unset": []}, + hermes_home=hermes_home, + ) + assert ( + hermes.remove_mcp_server_config( + "managed", + hermes_home=hermes_home, + expected_fingerprint=fingerprint, + ) + is False + ) + assert _loaded_config(binary, env)["mcp_servers"]["managed"] == replacement + with pytest.raises(RuntimeError, match="already exists"): + hermes.write_mcp_server_config("managed", proxy_argv, hermes_home=hermes_home) + assert _loaded_config(binary, env)["mcp_servers"]["managed"] == replacement diff --git a/tests/test_gateway_proxy.py b/tests/test_gateway_proxy.py index 28600045..03742e8c 100644 --- a/tests/test_gateway_proxy.py +++ b/tests/test_gateway_proxy.py @@ -220,14 +220,15 @@ def test_none_on_garbage(self): def _install_fake_token(monkeypatch, exp_offsets, delay=0.0): """Patch get_databricks_token to hand out JWTs whose exp is now+offset, one per successive mint (last offset repeats). Records the force flag of each.""" - state = {"i": 0, "forces": []} + state = {"i": 0, "forces": [], "allow_env_bearer": []} - def fake(_ws, _profile, force_refresh=False): + def fake(_ws, _profile, force_refresh=False, allow_env_bearer=False): if delay: time.sleep(delay) off = exp_offsets[min(state["i"], len(exp_offsets) - 1)] state["i"] += 1 state["forces"].append(force_refresh) + state["allow_env_bearer"].append(allow_env_bearer) return _make_jwt(time.time() + off) monkeypatch.setattr(gateway_proxy, "get_databricks_token", fake) @@ -235,6 +236,32 @@ def fake(_ws, _profile, force_refresh=False): class TestTokenCache: + def test_default_mint_explicitly_rejects_ambient_bearer(self, monkeypatch): + calls: list[bool] = [] + + def fake(_ws, _profile, *, force_refresh=False, allow_env_bearer): + calls.append(allow_env_bearer) + return _make_jwt(time.time() + 5000) + + monkeypatch.setattr(gateway_proxy, "get_databricks_token", fake) + + gateway_proxy.TokenCache("ws", "oauth-profile") + + assert calls == [False] + + def test_explicit_bearer_opt_in_reaches_token_mint(self, monkeypatch): + calls: list[bool] = [] + + def fake(_ws, _profile, *, force_refresh=False, allow_env_bearer=False): + calls.append(allow_env_bearer) + return _make_jwt(time.time() + 5000) + + monkeypatch.setattr(gateway_proxy, "get_databricks_token", fake) + + gateway_proxy.TokenCache("ws", "pat-profile", allow_env_bearer=True) + + assert calls == [True] + def test_initial_mint_preserves_default_nonforce_refresh(self, monkeypatch): state = _install_fake_token(monkeypatch, [5000]) gateway_proxy.TokenCache("ws", None) @@ -438,6 +465,33 @@ def test_failed_refresh_surfaces_reauth_hint(self, capsys): class TestStartProxyPortFallback: + def test_forwards_explicit_bearer_opt_in_to_cache(self, monkeypatch): + captured: dict[str, object] = {} + + class _StubCache: + def run_refresher(self): + return None + + def fake_cache(workspace, profile, **kwargs): + captured.update(workspace=workspace, profile=profile, **kwargs) + return _StubCache() + + monkeypatch.setattr(gateway_proxy, "TokenCache", fake_cache) + + server, _cache, client = gateway_proxy.start_proxy( + "https://x.staging.cloud.databricks.com", + "pat-profile", + 0, + token_header=gateway_proxy.AI_GATEWAY_TOKEN_HEADER, + force_refresh_near_expiry=False, + allow_env_bearer=True, + ) + try: + assert captured["allow_env_bearer"] is True + finally: + server.server_close() + client.close() + def test_falls_back_to_free_port_when_cached_port_busy(self, monkeypatch): # A stale proxy from a killed session can still hold the cached port; the # bind must fall back to an OS-assigned free port rather than crash. diff --git a/tests/test_gateway_proxy_integration.py b/tests/test_gateway_proxy_integration.py index 2809ced0..da18777b 100644 --- a/tests/test_gateway_proxy_integration.py +++ b/tests/test_gateway_proxy_integration.py @@ -133,7 +133,7 @@ def _counting_token(value: str = "dbx-swap-token"): the mint count deterministic (one on init, one per forced retry-refresh).""" calls: list[bool] = [] - def fn(_workspace, _profile, force_refresh=False): + def fn(_workspace, _profile, force_refresh=False, **_kwargs): calls.append(force_refresh) return value diff --git a/tests/test_mcp.py b/tests/test_mcp.py index d58062b6..a4fdd864 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -300,6 +300,49 @@ def test_skipped_when_binary_not_installed(self): class TestConfigureClientMcpServer: + def test_configures_hermes_with_proxy_argv(self, monkeypatch): + calls = [] + monkeypatch.setattr( + mcp.hermes, + "write_mcp_server_config", + lambda name, argv, **kwargs: calls.append((name, argv, kwargs)) or False, + ) + + assert ( + mcp.configure_client_mcp_server( + "hermes", "github", GH_URL, WS, "p", hermes_home="/profile" + ) + == [] + ) + assert calls == [ + ( + "github", + _proxy_argv(), + {"hermes_home": "/profile", "expected_fingerprint": None}, + ) + ] + + def test_removes_only_named_hermes_server(self, monkeypatch): + calls = [] + monkeypatch.setattr( + mcp.hermes, + "remove_mcp_server_config", + lambda name, **kwargs: calls.append((name, kwargs)) or True, + ) + + assert mcp.remove_client_mcp_server( + "hermes", + "github", + hermes_home="/profile", + expected_fingerprint="abc", + ) == [mcp.MCP_USER_SCOPE] + assert calls == [ + ( + "github", + {"hermes_home": "/profile", "expected_fingerprint": "abc"}, + ) + ] + def test_configures_copilot_with_proxy_argv(self, monkeypatch): calls: list[tuple[str, list[str]]] = [] @@ -687,6 +730,75 @@ def test_removes_servers_dropped_from_working_set(self, monkeypatch): assert changed is True assert removed == [("claude", "gone")] + def test_hermes_add_records_token_free_fingerprint_and_home(self, monkeypatch): + configured = [] + monkeypatch.setattr( + mcp, + "configure_client_mcp_server", + lambda *args, **kwargs: configured.append((args, kwargs)) or [], + ) + working = [self._server("managed", ["hermes"])] + + assert ( + mcp.apply_mcp_server_changes( + [], + working, + ["hermes"], + WS, + "profile", + hermes_home="/profiles/original", + ) + is True + ) + + ownership = working[0]["ownership"]["hermes"] + assert len(ownership["fingerprint"]) == 64 + assert ownership["hermes_home"] == "/profiles/original" + serialized = json.dumps(working).lower() + assert "bearer" not in serialized + assert "token" not in serialized + assert configured[0][1]["hermes_home"] == "/profiles/original" + + def test_hermes_update_does_not_reuse_fingerprint_from_another_home(self, monkeypatch): + configured = [] + monkeypatch.setattr( + mcp, + "configure_client_mcp_server", + lambda *args, **kwargs: configured.append((args, kwargs)) or [], + ) + original = self._server("managed", ["hermes"]) + original["ownership"] = {"hermes": {"fingerprint": "a" * 64, "hermes_home": "/profiles/a"}} + working = [{**self._server("managed", ["hermes"]), "url": GH_URL + "/changed"}] + + assert mcp.apply_mcp_server_changes( + [original], + working, + ["hermes"], + WS, + hermes_home="/profiles/b", + ) + + assert configured[0][1]["expected_fingerprint"] is None + assert working[0]["ownership"]["hermes"]["hermes_home"] == "/profiles/b" + + def test_recorded_removal_uses_server_owned_home(self, monkeypatch): + calls = [] + monkeypatch.setattr( + mcp, + "remove_client_mcp_server", + lambda *args, **kwargs: calls.append((args, kwargs)) or ["user"], + ) + server = self._server("managed", ["hermes"]) + server["ownership"] = {"hermes": {"fingerprint": "b" * 64, "hermes_home": "/profiles/a"}} + + assert mcp._remove_recorded_client_server( + "hermes", "managed", server, hermes_home="/profiles/b" + ) == ["user"] + assert calls[0][1] == { + "hermes_home": "/profiles/a", + "expected_fingerprint": "b" * 64, + } + def test_no_ops_returns_false_without_spinner(self, monkeypatch): # Identical original/working, so nothing to do. monkeypatch.setattr( @@ -1255,7 +1367,7 @@ def test_drops_stale_foreign_workspace_mcp_entries(self, monkeypatch, capsys): assert mcp.configure_mcp_command() == 0 output = capsys.readouterr().out - assert "Dropping 1 stale MCP entry" in output + assert "Cleaning 1 stale MCP entry" in output assert "databricks-genie-foreign" in output # codex is listed on the stale entry but not installed -> skipped. assert cleanup_calls == [("claude", "databricks-genie-foreign")] @@ -2700,6 +2812,40 @@ def test_drops_foreign_workspace_skills_entry(self, monkeypatch): assert removed == [("claude", mcp.SKILLS_MCP_SERVER_NAME)] assert state["mcp_servers"] == [] + def test_retains_foreign_hermes_ownership_when_safe_removal_does_not_match(self, monkeypatch): + foreign_entry = { + "name": "managed", + "url": "https://other.databricks.com/api/2.0/mcp/functions/a/b", + "clients": ["hermes"], + "ownership": {"hermes": {"fingerprint": "a" * 64, "hermes_home": "/profiles/a"}}, + } + state = {"mcp_servers": [foreign_entry]} + monkeypatch.setattr(mcp, "available_mcp_clients", lambda: ["hermes"]) + monkeypatch.setattr(mcp, "load_full_state", lambda: {}) + monkeypatch.setattr(mcp, "_remove_recorded_client_server", lambda *_a, **_kw: []) + monkeypatch.setattr(mcp, "save_state", lambda _state: None) + + mcp.purge_cross_workspace_mcp_residue(state, WS) + + assert state["mcp_servers"] == [foreign_entry] + + def test_drops_foreign_hermes_ownership_after_verified_removal(self, monkeypatch): + foreign_entry = { + "name": "managed", + "url": "https://other.databricks.com/api/2.0/mcp/functions/a/b", + "clients": ["hermes"], + "ownership": {"hermes": {"fingerprint": "a" * 64, "hermes_home": "/profiles/a"}}, + } + state = {"mcp_servers": [foreign_entry]} + monkeypatch.setattr(mcp, "available_mcp_clients", lambda: ["hermes"]) + monkeypatch.setattr(mcp, "load_full_state", lambda: {}) + monkeypatch.setattr(mcp, "_remove_recorded_client_server", lambda *_a, **_kw: ["user"]) + monkeypatch.setattr(mcp, "save_state", lambda _state: None) + + mcp.purge_cross_workspace_mcp_residue(state, WS) + + assert state["mcp_servers"] == [] + class TestManagedMcpServerEntry: def test_sql(self): @@ -2846,6 +2992,30 @@ def test_no_supported_servers_does_nothing(self, monkeypatch): ) assert registered == [] + def test_last_managed_server_removal_runs_before_returning_empty(self, monkeypatch): + previous = { + "name": "old", + "url": "https://example/mcp", + "auth": "proxy", + "clients": ["hermes"], + } + seen = {} + monkeypatch.setattr(mcp, "load_state", lambda: {"managed_mcp_servers": [previous]}) + monkeypatch.setattr( + mcp, + "apply_mcp_server_changes", + lambda original, working, clients, *args, **kwargs: seen.update( + original=original, + working=working, + clients=clients, + ), + ) + + registered = mcp.apply_managed_mcp_servers({}, "hermes", WS) + + assert registered == [] + assert seen == {"original": [previous], "working": [], "clients": ["hermes"]} + def test_mcp_only_client_returns_empty(self, monkeypatch): # A tool that isn't an MCP client can't have servers registered against it. monkeypatch.setattr( diff --git a/tests/test_mcp_proxy.py b/tests/test_mcp_proxy.py index bf7a3536..4278dd81 100644 --- a/tests/test_mcp_proxy.py +++ b/tests/test_mcp_proxy.py @@ -2,6 +2,7 @@ from __future__ import annotations +import os import tomllib from contextlib import asynccontextmanager from pathlib import Path @@ -59,8 +60,26 @@ def test_proxy_imports_the_streamable_http_client_shared_by_both_majors(): class TestDatabricksTokenAuth: + def test_default_auth_explicitly_rejects_ambient_bearer(self, monkeypatch): + calls: list[bool] = [] + + def fake(_ws, _profile, *, allow_env_bearer): + calls.append(allow_env_bearer) + return "tok-123" + + monkeypatch.setattr(mcp_proxy, "get_databricks_token", fake) + auth = mcp_proxy._build_token_auth(WS, "oauth-profile") + + list(auth.auth_flow(httpx.Request("POST", URL))) + + assert calls == [False] + def test_injects_bearer_from_minted_token(self, monkeypatch): - monkeypatch.setattr(mcp_proxy, "get_databricks_token", lambda ws, profile: "tok-123") + monkeypatch.setattr( + mcp_proxy, + "get_databricks_token", + lambda ws, profile, *, allow_env_bearer: "tok-123", + ) auth = mcp_proxy._build_token_auth(WS, "uc-dogfood") request = httpx.Request("POST", URL) @@ -78,23 +97,29 @@ def test_auth_is_an_instance_of_the_selected_httpx_auth(self, monkeypatch): assert isinstance(auth, mcp_proxy._httpx().Auth) def test_calls_get_token_with_workspace_and_profile(self, monkeypatch): - calls: list[tuple[str, str | None]] = [] + calls: list[tuple[str, str | None, bool]] = [] monkeypatch.setattr( mcp_proxy, "get_databricks_token", - lambda ws, profile: calls.append((ws, profile)) or "t", + lambda ws, profile, *, allow_env_bearer: ( + calls.append((ws, profile, allow_env_bearer)) or "t" + ), ) - auth = mcp_proxy._build_token_auth(WS, "myprofile") + auth = mcp_proxy._build_token_auth(WS, "myprofile", allow_env_bearer=True) list(auth.auth_flow(httpx.Request("POST", URL))) - assert calls == [(WS, "myprofile")] + assert calls == [(WS, "myprofile", True)] def test_mints_a_fresh_token_per_request(self, monkeypatch): # Each request re-invokes get_databricks_token, so a rotated token is # picked up mid-session without the proxy tracking expiry itself. tokens = iter(["first", "second"]) - monkeypatch.setattr(mcp_proxy, "get_databricks_token", lambda ws, profile: next(tokens)) + monkeypatch.setattr( + mcp_proxy, + "get_databricks_token", + lambda ws, profile, *, allow_env_bearer: next(tokens), + ) auth = mcp_proxy._build_token_auth(WS, None) r1 = httpx.Request("POST", URL) @@ -106,7 +131,11 @@ def test_mints_a_fresh_token_per_request(self, monkeypatch): assert r2.headers["Authorization"] == "Bearer second" def test_auth_flow_yields_the_same_request(self, monkeypatch): - monkeypatch.setattr(mcp_proxy, "get_databricks_token", lambda ws, profile: "t") + monkeypatch.setattr( + mcp_proxy, + "get_databricks_token", + lambda ws, profile, *, allow_env_bearer: "t", + ) auth = mcp_proxy._build_token_auth(WS, None) request = httpx.Request("POST", URL) @@ -118,7 +147,7 @@ def test_dead_auth_becomes_a_terminal_proxy_auth_error(self, monkeypatch): # A raw RuntimeError escaping auth_flow tears through the transport's task # group and stalls the proxy until the client's startup timeout. # Translating it keeps the failure reportable by `serve`. - def boom(ws, profile): + def boom(ws, profile, *, allow_env_bearer): raise RuntimeError("no access token; run `databricks auth login`") monkeypatch.setattr(mcp_proxy, "get_databricks_token", boom) @@ -230,7 +259,7 @@ async def stop_bridge(*args, **kwargs): yield monkeypatch.setattr(httpx_module, "AsyncClient", CapturingClient) - monkeypatch.setattr(mcp_proxy, "_build_token_auth", lambda *args: object()) + monkeypatch.setattr(mcp_proxy, "_build_token_auth", lambda *args, **kwargs: object()) monkeypatch.setattr(mcp_proxy, "streamable_http_client", stop_bridge) with pytest.raises(StopBridge): @@ -241,6 +270,19 @@ async def stop_bridge(*args, **kwargs): class TestServe: + def test_default_preflight_explicitly_rejects_ambient_bearer(self, monkeypatch): + calls: list[bool] = [] + + def fake(_ws, _profile, *, allow_env_bearer): + calls.append(allow_env_bearer) + return "tok-123" + + monkeypatch.setattr(mcp_proxy, "get_databricks_token", fake) + + mcp_proxy._preflight_token(WS, "oauth-profile") + + assert calls == [False] + def test_runs_the_bridge_with_parsed_args(self, monkeypatch): captured: dict = {} @@ -272,16 +314,57 @@ def test_use_pat_exports_the_bearer_before_serving(self, monkeypatch): # short-circuit returns it. `databricks auth token` can't read a PAT itself. order: list[str] = [] monkeypatch.setattr( - mcp_proxy, "ensure_pat_bearer", lambda profile: order.append(f"pat:{profile}") or True + mcp_proxy, + "resolve_pat_token", + lambda profile: order.append(f"resolve:{profile}") or "profile-pat", ) monkeypatch.setattr( - mcp_proxy, "_preflight_token", lambda ws, profile: order.append("preflight") + mcp_proxy, + "ensure_pat_bearer", + lambda profile, pat, *, workspace: ( + order.append(f"pat:{workspace}:{profile}:{pat}") or True + ), + ) + monkeypatch.setattr( + mcp_proxy, + "_preflight_token", + lambda ws, profile, *, allow_env_bearer: order.append(f"preflight:{allow_env_bearer}"), + ) + monkeypatch.setattr( + mcp_proxy.anyio, + "run", + lambda func, *args: order.append(f"bridge:{args[-1]}"), + ) + + mcp_proxy.serve(URL, WS, "patprof", use_pat=True) + + assert order == [ + "resolve:patprof", + f"pat:{WS}:patprof:profile-pat", + "preflight:True", + "bridge:True", + ] + + def test_use_pat_profile_replaces_conflicting_ambient_bearer(self, monkeypatch): + observed: dict[str, str] = {} + monkeypatch.setenv("DATABRICKS_BEARER", "other-workspace-token") + monkeypatch.setattr( + mcp_proxy, + "resolve_pat_token", + lambda profile: "profile-pat", + raising=False, ) - monkeypatch.setattr(mcp_proxy.anyio, "run", lambda func, *args: order.append("bridge")) + + def preflight(_ws, _profile, *, allow_env_bearer): + observed["bearer"] = os.environ["DATABRICKS_BEARER"] + observed["allowed"] = str(allow_env_bearer) + + monkeypatch.setattr(mcp_proxy, "_preflight_token", preflight) + monkeypatch.setattr(mcp_proxy.anyio, "run", lambda *_args: None) mcp_proxy.serve(URL, WS, "patprof", use_pat=True) - assert order == ["pat:patprof", "preflight", "bridge"] + assert observed == {"bearer": "profile-pat", "allowed": "True"} def test_use_pat_without_a_resolvable_pat_exits_before_serving(self, monkeypatch, capsys): started: list[str] = [] @@ -395,11 +478,15 @@ def raise_other(func, *args): class TestPreflightToken: def test_passes_through_when_a_token_is_available(self, monkeypatch): - monkeypatch.setattr(mcp_proxy, "get_databricks_token", lambda ws, profile: "tok") + monkeypatch.setattr( + mcp_proxy, + "get_databricks_token", + lambda ws, profile, *, allow_env_bearer: "tok", + ) mcp_proxy._preflight_token(WS, "p") # no exception def test_surfaces_the_cli_error_message(self, monkeypatch): - def boom(ws, profile): + def boom(ws, profile, *, allow_env_bearer): raise RuntimeError("profile is stale; run `databricks auth logout`") monkeypatch.setattr(mcp_proxy, "get_databricks_token", boom) @@ -414,7 +501,7 @@ def test_checks_the_same_workspace_and_profile_the_bridge_will_use(self, monkeyp monkeypatch.setattr( mcp_proxy, "get_databricks_token", - lambda ws, profile: calls.append((ws, profile)) or "tok", + lambda ws, profile, *, allow_env_bearer: calls.append((ws, profile)) or "tok", ) mcp_proxy._preflight_token(WS, "myprofile") diff --git a/tests/test_state.py b/tests/test_state.py index 36c8ce4f..f18897dd 100644 --- a/tests/test_state.py +++ b/tests/test_state.py @@ -246,6 +246,45 @@ def test_drops_falsy_managed_configs(self): assert "codex" not in result["managed_configs"] assert "claude" not in result["managed_configs"] + def test_preserves_valid_hermes_ownership_metadata(self): + state = { + "managed_configs": { + "hermes": { + "keys": [], + "hermes_home": "/profiles/team", + "active_model": { + "provider": "ucode-databricks-codex", + "default": "system.ai.gpt-5", + }, + "provider_fingerprints": { + "ucode-databricks-codex": "a" * 64, + }, + }, + "codex": {"keys": [["model"]], "hermes_home": "/ignored"}, + } + } + + result = hydrate_state(state) + + assert result["managed_configs"]["hermes"] == state["managed_configs"]["hermes"] + assert result["managed_configs"]["codex"] == {"keys": [["model"]]} + + def test_drops_invalid_hermes_ownership_metadata(self): + result = hydrate_state( + { + "managed_configs": { + "hermes": { + "keys": [], + "hermes_home": "", + "active_model": {"provider": "", "default": 7}, + "provider_fingerprints": {"not-owned": "bad"}, + } + } + } + ) + + assert result["managed_configs"]["hermes"] == {"keys": []} + class TestBuildAgentState: def test_returns_empty_without_workspace(self): @@ -293,3 +332,16 @@ def test_preserves_existing_managed_configs(self): def test_records_only_keys(self): result = mark_tool_managed({}, "codex", [["model"]]) assert result["managed_configs"]["codex"] == {"keys": [["model"]]} + + def test_records_hermes_ownership_metadata(self): + metadata = { + "hermes_home": "/profiles/team", + "active_model": { + "provider": "ucode-databricks-codex", + "default": "system.ai.gpt-5", + }, + } + + result = mark_tool_managed({}, "hermes", [], metadata=metadata) + + assert result["managed_configs"]["hermes"] == {"keys": [], **metadata}