From f8054581b730fb668381a1fff154266c3f1fb334 Mon Sep 17 00:00:00 2001 From: Rohit Agrawal Date: Wed, 1 Jul 2026 09:28:48 -0400 Subject: [PATCH 1/2] Add budget tracking hooks and enforcement Introduce budget_hooks.py with budget-aware agent dispatch, wire budget checks through the CLI and Databricks layers, and add supporting scripts and tests. Co-authored-by: Isaac --- .gitignore | 6 + OPENCODE_PLAN.md | 181 +++++++++++ scripts/fake_api_key_helper.sh | 58 ++++ scripts/generate_image.py | 49 +++ scripts/test_api_key_refresh.sh | 103 +++++++ src/ucode/agents/__init__.py | 523 +++++++++++++++++++++++++++++++- src/ucode/agents/claude.py | 54 ++++ src/ucode/agents/codex.py | 59 ++++ src/ucode/budget_hooks.py | 32 ++ src/ucode/cli.py | 83 ++++- src/ucode/databricks.py | 236 ++++++++++++++ tests/test_agent_claude.py | 40 +++ tests/test_agent_codex.py | 46 +++ tests/test_agents_init.py | 254 ++++++++++++++++ tests/test_cli.py | 62 ++++ tests/test_databricks.py | 165 ++++++++++ 16 files changed, 1945 insertions(+), 6 deletions(-) create mode 100644 OPENCODE_PLAN.md create mode 100755 scripts/fake_api_key_helper.sh create mode 100644 scripts/generate_image.py create mode 100755 scripts/test_api_key_refresh.sh create mode 100644 src/ucode/budget_hooks.py diff --git a/.gitignore b/.gitignore index 4d1e17c..bada883 100644 --- a/.gitignore +++ b/.gitignore @@ -5,3 +5,9 @@ build/ dist/ .venv/ .DS_Store + +# Local artifacts +.vscode/ +mlflow.db +*.jpg +*.pdf diff --git a/OPENCODE_PLAN.md b/OPENCODE_PLAN.md new file mode 100644 index 0000000..d8564f8 --- /dev/null +++ b/OPENCODE_PLAN.md @@ -0,0 +1,181 @@ +# Plan: Add opencode to ucode + +## What is opencode + +opencode is an AI coding CLI (`npm i -g opencode-ai`) that uses the Vercel AI SDK +internally. Unlike codex/claude/gemini which each speak one proprietary protocol, +opencode supports 22+ AI SDK providers via an `npm` field in its JSON config. + +Config lives at `~/.config/opencode/opencode.json`. + +--- + +## Why it's not as simple as the other tools + +### 1. Token auth — no shell command support (TBD) + +Codex and Claude Code support a shell auth command that runs on every request: +- Codex: `auth.command = "sh"` in TOML +- Claude Code: `apiKeyHelper` shell script in settings.json + +This means their tokens are always fresh. Gemini doesn't support this, so +ucode spawns a background thread to refresh `GEMINI_API_KEY` in the +.env file every 30 minutes. + +**opencode config only supports static apiKey or `{env:VAR}` syntax.** +This means we need to either: +- (a) Embed a live token at launch time and refresh like Gemini — requires + background refresh thread and subprocess launch (not execvp) +- (b) Investigate whether opencode re-reads its config on each request (unlikely) +- (c) Set `DATABRICKS_TOKEN` env var and use `{env:DATABRICKS_TOKEN}` in config — + then refresh that env var in the background (cleanest if it works) + +**Open question**: Does opencode read apiKey from env at request time or only at +startup? If at request time, option (c) works cleanly. + +### 2. Multiple providers vs. single endpoint + +opencode can be configured with multiple providers. Two approaches: + +**Option A — Single mlflow endpoint with `@ai-sdk/openai-compatible`** +```json +{ + "provider": { + "databricks": { + "npm": "@ai-sdk/openai-compatible", + "options": { + "baseURL": "https:///ai-gateway/mlflow/v1", + "apiKey": "{env:DATABRICKS_TOKEN}" + } + } + }, + "model": "databricks/databricks-claude-sonnet-4-6" +} +``` +- Simple, one endpoint to check +- Some models have known incompatibilities (content-as-array for Gemini 3.x thinking) +- User picks from `databricks-*` model IDs + +**Option B — Dedicated provider endpoints per model family** +```json +{ + "provider": { + "databricks-anthropic": { + "npm": "@ai-sdk/anthropic", + "options": { + "baseURL": "https:///ai-gateway/anthropic", + "apiKey": "{env:DATABRICKS_TOKEN}" + } + }, + "databricks-google": { + "npm": "@ai-sdk/google", + "options": { + "baseURL": "https:///ai-gateway/gemini", + "apiKey": "{env:DATABRICKS_TOKEN}" + } + }, + "databricks-openai": { + "npm": "@ai-sdk/openai", + "options": { + "baseURL": "https:///ai-gateway/codex/v1", + "apiKey": "{env:DATABRICKS_TOKEN}" + } + } + }, + "model": "databricks-anthropic/databricks-claude-sonnet-4-6" +} +``` +- Native SDK per family → better compatibility (thinking models work correctly) +- More config complexity — 3 providers, 3 endpoints to check +- User picks a model + its provider is inferred from the model name prefix + +**Open question**: Does `@ai-sdk/anthropic` work against the Databricks Anthropic +gateway with `databricks-*` model IDs? Same for `@ai-sdk/google` against the +Gemini gateway? Needs testing. + +### 3. Model selection + +Unlike codex (no model selection), opencode needs a model specified in config. +The model list should come from the `providers/databricks/models/` TOMLs in +models.dev (or from querying the workspace endpoint directly). + +`classify_tool_from_text()` and `discover_workspace_models()` would need an +"opencode" case, but since opencode supports ANY model family, the full +`databricks-*` list is valid — not just one family like claude/gemini. + +Alternatively: skip workspace discovery entirely, let user type a model name, +show the list from models.dev as a hint. + +### 4. Config path + +`~/.config/opencode/opencode.json` — different from the other tools which use +home-dir dotfiles. Needs new path constants. + +### 5. Usage tracking + +The Spark SQL query in `build_usage_report_query()` filters by user_agent +containing "codex", "claude", or "gemini". Need to verify opencode sends +"opencode" in its user-agent (likely yes, confirmed from worker.ts in models.dev). +Add "opencode" case to the query. + +### 6. Gateway endpoint check + +`check_gateway_endpoint()` needs an opencode case. For Option A (mlflow), probe +`/ai-gateway/mlflow/v1/models`. For Option B, probe each provider endpoint. + +### 7. Validation test command + +`validate_tool()` needs an opencode case. opencode's CLI args are TBD — need to +check if it supports a single `-p "prompt"` style invocation or requires +interactive mode only. + +--- + +## What changes in cli.py + +| Section | Change | +|---|---| +| Path constants | Add `OPENCODE_CONFIG_DIR`, `OPENCODE_CONFIG_PATH`, `OPENCODE_BACKUP_PATH` | +| `TOOL_SPECS` | Add opencode entry: binary, package, display, config_path, backup_path | +| `TOOL_ALIASES` | Add "opencode" alias | +| `DEFAULT_SELECTED_MODELS` | Add default model (e.g. `databricks-claude-sonnet-4-6`) | +| `build_tool_base_url()` | Add opencode case (mlflow endpoint or per-family) | +| `render_opencode_config()` | New function — generates opencode.json content | +| `write_opencode_tool_config()` | New function — backup/write/mark managed | +| `configure_tool()` | Add opencode elif branch | +| `check_gateway_endpoint()` | Add opencode elif branch | +| `validate_tool()` | Add opencode elif branch | +| `build_usage_report_query()` | Add opencode to user_agent filters | +| `classify_tool_from_text()` | Add "opencode" detection | +| Launch path | Either reuse `launch_tool()` (execvp) or new `launch_opencode_tool()` with token refresh | + +--- + +## Open questions to resolve before implementing + +1. **Does opencode re-read `{env:VAR}` at request time or only at startup?** + → Determines auth approach (static refresh vs. env var) + +2. **Does `@ai-sdk/anthropic` work against `/ai-gateway/anthropic` with `databricks-*` model IDs?** + → Determines Option A vs. Option B for provider config + +3. **What are opencode's CLI flags for non-interactive single-prompt invocation?** + → Needed for `validate_tool()` test command + +4. **Does opencode emit "opencode" in its user-agent?** + → Needed for usage tracking SQL + +5. **Which approach for model selection?** + → Workspace discovery (requires opencode case in classify_tool_from_text) + or hardcoded list from models.dev + +--- + +## Recommended next steps + +1. Test Option B (native provider SDKs) against Databricks endpoints — run a quick + script using `@ai-sdk/anthropic` with baseURL set to the Databricks Anthropic + gateway to confirm model IDs and auth work +2. Check opencode source for env var re-read behavior +3. Check opencode CLI flags for non-interactive mode +4. Decide Option A vs B, then implement the 12 changes above diff --git a/scripts/fake_api_key_helper.sh b/scripts/fake_api_key_helper.sh new file mode 100755 index 0000000..24e64d4 --- /dev/null +++ b/scripts/fake_api_key_helper.sh @@ -0,0 +1,58 @@ +#!/usr/bin/env bash +# Fake apiKeyHelper for testing Claude Code's token refresh behavior. +# +# First call: returns a real Databricks token. +# Subsequent calls: behavior controlled by FAKE_HELPER_FAILURE_MODE env var: +# empty (default) — returns empty string, exit 0 +# nonzero — prints error to stderr, exits 1 +# reauth — fails once then recovers (simulates re-auth fixing the problem) +# +# Each invocation is logged to /tmp/fake_api_key_helper.log with a timestamp. + +CALL_COUNT_FILE="/tmp/fake_api_key_helper_call_count" +LOG_FILE="/tmp/fake_api_key_helper.log" +WORKSPACE="${DATABRICKS_HOST:-https://eng-ml-inference-team-us-east-1.cloud.databricks.com}" + +# Increment call counter +if [[ -f "$CALL_COUNT_FILE" ]]; then + count=$(cat "$CALL_COUNT_FILE") +else + count=0 +fi +count=$((count + 1)) +echo "$count" > "$CALL_COUNT_FILE" + +timestamp=$(date '+%Y-%m-%dT%H:%M:%S') + +if [[ "$count" -eq 1 ]]; then + # First call: get a real token + token=$(env -u DATABRICKS_TOKEN -u DATABRICKS_CLIENT_ID -u DATABRICKS_CLIENT_SECRET \ + -u DATABRICKS_USERNAME -u DATABRICKS_PASSWORD -u DATABRICKS_AUTH_TYPE \ + databricks auth token --host "$WORKSPACE" --output json 2>/dev/null \ + | python3 -c "import json,sys; print(json.load(sys.stdin).get('access_token', ''))" 2>/dev/null) + echo "$timestamp call=$count returning=real_token len=${#token}" >> "$LOG_FILE" + echo "$token" +else + mode="${FAKE_HELPER_FAILURE_MODE:-empty}" + if [[ "$mode" == "nonzero" ]]; then + echo "$timestamp call=$count returning=nonzero_exit (simulated failure)" >> "$LOG_FILE" + echo "databricks auth: token refresh failed (simulated)" >&2 + exit 1 + elif [[ "$mode" == "reauth" ]]; then + # Simulate: first failure triggers re-auth, subsequent calls succeed + if [[ "$count" -eq 2 ]]; then + # This is the "re-auth" call — takes a moment, then returns a real token + echo "$timestamp call=$count returning=real_token_after_reauth (simulated re-auth)" >> "$LOG_FILE" + else + echo "$timestamp call=$count returning=real_token (recovered)" >> "$LOG_FILE" + fi + token=$(env -u DATABRICKS_TOKEN -u DATABRICKS_CLIENT_ID -u DATABRICKS_CLIENT_SECRET \ + -u DATABRICKS_USERNAME -u DATABRICKS_PASSWORD -u DATABRICKS_AUTH_TYPE \ + databricks auth token --host "$WORKSPACE" --output json 2>/dev/null \ + | python3 -c "import json,sys; print(json.load(sys.stdin).get('access_token', ''))" 2>/dev/null) + echo "$token" + else + echo "$timestamp call=$count returning=empty (simulated failure)" >> "$LOG_FILE" + echo "" + fi +fi diff --git a/scripts/generate_image.py b/scripts/generate_image.py new file mode 100644 index 0000000..5cf171f --- /dev/null +++ b/scripts/generate_image.py @@ -0,0 +1,49 @@ +#!/usr/bin/env python3 +"""Generate an image via the Databricks AI Gateway (OpenAI-compatible responses API). + +Usage: + DATABRICKS_TOKEN=dapi... uv run --with openai python scripts/generate_image.py \ + "a gray tabby cat hugging an otter with an orange scarf" [out.png] +""" +import base64 +import os +import sys + +from openai import OpenAI + +BASE_URL = "https://eng-ml-inference.staging.cloud.databricks.com/ai-gateway/openai/v1" +MODEL = "databricks-gpt-5-5" + + +def main() -> int: + token = os.environ.get("DATABRICKS_TOKEN") + if not token: + print("error: set DATABRICKS_TOKEN", file=sys.stderr) + return 1 + + prompt = sys.argv[1] if len(sys.argv) > 1 else ( + "a gray tabby cat hugging an otter with an orange scarf" + ) + out_path = sys.argv[2] if len(sys.argv) > 2 else "image.png" + + client = OpenAI(api_key=token, base_url=BASE_URL) + response = client.responses.create( + model=MODEL, + input=prompt, + tools=[{"type": "image_generation"}], + ) + + images = [o.result for o in response.output if o.type == "image_generation_call"] + if not images: + print("error: no image returned. Raw output:", file=sys.stderr) + print(response.output, file=sys.stderr) + return 2 + + with open(out_path, "wb") as f: + f.write(base64.b64decode(images[0])) + print(f"saved {out_path}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/test_api_key_refresh.sh b/scripts/test_api_key_refresh.sh new file mode 100755 index 0000000..526c76b --- /dev/null +++ b/scripts/test_api_key_refresh.sh @@ -0,0 +1,103 @@ +#!/usr/bin/env bash +# Test whether Claude Code handles apiKeyHelper failures gracefully. +# +# Runs two scenarios back-to-back: +# Mode A (empty): helper returns "" with exit 0 → hangs silently (known bad behavior) +# Mode B (nonzero): helper exits 1 with stderr msg → should surface an error and exit +# +# Usage: bash scripts/test_api_key_refresh.sh + +set -euo pipefail + +SETTINGS="$HOME/.claude/settings.json" +BACKUP="$HOME/.claude/settings.json.test_backup" +CALL_COUNT_FILE="/tmp/fake_api_key_helper_call_count" +LOG_FILE="/tmp/fake_api_key_helper.log" +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +HELPER="$SCRIPT_DIR/fake_api_key_helper.sh" +TIMEOUT=30 # seconds before giving up on a hung claude invocation + +cleanup() { + if [[ -f "$BACKUP" ]]; then + echo "" + echo "--- Restoring original settings.json ---" + cp "$BACKUP" "$SETTINGS" + rm "$BACKUP" + fi +} +trap cleanup EXIT + +chmod +x "$HELPER" + +patch_settings() { + python3 - "$SETTINGS" "$HELPER" <<'EOF' +import json, sys +path, helper = sys.argv[1], sys.argv[2] +with open(path) as f: + s = json.load(f) +s["apiKeyHelper"] = helper +s.setdefault("env", {})["CLAUDE_CODE_API_KEY_HELPER_TTL_MS"] = "5000" +with open(path, "w") as f: + json.dump(s, f, indent=2) +EOF +} + +run_scenario() { + local mode="$1" + echo "" + echo "======================================================" + echo " SCENARIO: failure mode = $mode" + echo "======================================================" + + # Reset call counter and log + rm -f "$CALL_COUNT_FILE" "$LOG_FILE" + + # Backup and patch settings + cp "$SETTINGS" "$BACKUP" + patch_settings + + echo "" + echo "--- Run 1: first invocation (real token) ---" + FAKE_HELPER_FAILURE_MODE="$mode" \ + timeout "$TIMEOUT" claude --dangerously-skip-permissions -p "Reply with only: HELLO_ONE" 2>&1 \ + | tail -5 || echo "(exit code: $?)" + + echo "" + echo "--- Waiting 8s for TTL to expire ---" + sleep 8 + + echo "" + echo "--- Run 2: second invocation (helper will fail with mode=$mode) ---" + FAKE_HELPER_FAILURE_MODE="$mode" \ + timeout "$TIMEOUT" claude --dangerously-skip-permissions -p "Reply with only: HELLO_TWO" 2>&1 \ + | tail -10 \ + && echo "(exited 0)" || echo "(exited non-zero / timed out after ${TIMEOUT}s)" + + echo "" + echo "--- Helper log ---" + if [[ -f "$LOG_FILE" ]]; then + cat "$LOG_FILE" + else + echo "(helper was never called)" + fi + echo "Total helper calls: $(cat "$CALL_COUNT_FILE" 2>/dev/null || echo 0)" + + # Restore settings before next scenario + cp "$BACKUP" "$SETTINGS" + rm "$BACKUP" +} + +echo "=== Claude Code apiKeyHelper failure mode comparison ===" +echo "Helper: $HELPER" + +run_scenario "empty" +run_scenario "nonzero" +run_scenario "reauth" + +echo "" +echo "=== Complete. Compare the three scenarios above. ===" +echo "" +echo "Expected results:" +echo " empty — retries forever, times out (broken)" +echo " nonzero — retries forever, times out (broken)" +echo " reauth — recovers on second call, returns HELLO_TWO (good)" diff --git a/src/ucode/agents/__init__.py b/src/ucode/agents/__init__.py index b94e855..d92692f 100644 --- a/src/ucode/agents/__init__.py +++ b/src/ucode/agents/__init__.py @@ -13,11 +13,18 @@ from __future__ import annotations import json +import os import shutil import subprocess +from pathlib import Path +from typing import cast -from ucode.config_io import ToolSpec +from rich.panel import Panel +from rich.table import Table + +from ucode.config_io import APP_DIR, ToolSpec from ucode.databricks import ( + fetch_current_user_budget_spend_status, install_databricks_cli, ) from ucode.state import load_state, save_state @@ -26,6 +33,7 @@ console, is_low_verbosity, print_err, + print_kv, print_note, print_section, print_success, @@ -60,6 +68,13 @@ DEFAULT_TOOL = "codex" BUNDLE_VERSION = 1 +BUDGET_STATUS_TOOLS = {"claude", "codex"} +BUDGET_POLICY_ENV_VAR = "UCODE_BUDGET_POLICY" +BUDGET_POLICY_PATH = APP_DIR / "budget-policy.json" +SMART_BUDGET_SWITCH_PERCENT = 80.0 +SMART_BUDGET_LOW_TOOL = "claude" +SMART_BUDGET_HIGH_TOOL = "codex" +DEFAULT_BUDGET_POLICY_NAME = "Default smart budget policy" def normalize_tool(tool: str) -> str: @@ -71,6 +86,10 @@ def normalize_tool(tool: str) -> str: return normalized +def budget_policy_env_configured() -> bool: + return bool(os.environ.get(BUDGET_POLICY_ENV_VAR, "").strip()) + + def _update_installed_tool_binary(tool: str, version: str | None = None) -> bool: spec = TOOL_SPECS[tool] binary = spec["binary"] @@ -241,12 +260,54 @@ def default_model_for_tool(tool: str, state: dict) -> str | None: return _MODULES[tool].default_model(state) +def _resolve_policy_model_alias(tool: str, state: dict, model: str | None) -> str | None: + if not model: + return None + raw = model.strip() + if not raw: + return None + if tool == "claude": + claude_models = state.get("claude_models") or {} + if isinstance(claude_models, dict): + value = claude_models.get(raw.lower()) + if isinstance(value, str) and value.strip(): + return value.strip() + for candidate in claude_models.values(): + if isinstance(candidate, str) and ( + candidate == raw or candidate.endswith(raw) or raw in candidate + ): + return candidate + if tool == "codex": + codex_models = state.get("codex_models") or [] + if isinstance(codex_models, list): + for candidate in codex_models: + if not isinstance(candidate, str): + continue + tail = candidate.rsplit("/", 1)[-1] + tail = tail[len("system.ai.") :] if tail.startswith("system.ai.") else tail + tail = tail[len("databricks-") :] if tail.startswith("databricks-") else tail + if candidate == raw or tail == raw or candidate.endswith(raw): + return candidate + # No exact match. On a model-services workspace every model routes by + # its `system.ai.` id, so a bare policy model (e.g. + # `gpt-5-4-mini`) must keep that prefix rather than fall through to + # the legacy OpenAI-id rewrite, which would mangle it into an + # unroutable dotted id like `gpt-5.4-mini`. + if not raw.startswith("system.ai.") and any( + isinstance(c, str) and c.startswith("system.ai.") for c in codex_models + ): + return f"system.ai.{raw}" + return raw + + def resolve_launch_model( tool: str, state: dict, explicit_model: str | None, ) -> tuple[dict, str | None]: - model = explicit_model or default_model_for_tool(tool, state) + model = _resolve_policy_model_alias(tool, state, explicit_model) or default_model_for_tool( + tool, state + ) if not model: raise RuntimeError( f"No models available for {tool}. Run `ucode configure` to set up your workspace." @@ -277,7 +338,465 @@ def configure_tool(tool: str, state: dict, model: str | None = None) -> dict: return result +def _as_float(value: object) -> float | None: + if isinstance(value, bool): + return None + if isinstance(value, (int, float)): + return float(value) + if isinstance(value, str): + try: + return float(value) + except ValueError: + return None + return None + + +def _format_usd(value: float | None) -> str: + if value is None: + return "-" + return f"${value:,.2f}" + + +def _short_budget_id(value: object) -> str: + raw = str(value or "").strip() + if not raw: + return "budget" + return raw[:8] + + +def _format_budget_status(entry: dict) -> tuple[str, str]: + budget_id = _short_budget_id(entry.get("budget_id")) + spend_status = entry.get("spend_status") + spend_status = spend_status if isinstance(spend_status, dict) else {} + spend = _as_float(spend_status.get("spend")) + threshold = _as_float(spend_status.get("effective_alert_threshold")) + if threshold and threshold > 0 and spend is not None: + percent = (spend / threshold) * 100 + return ( + f"Budget {budget_id}", + f"{_format_usd(spend)} / {_format_usd(threshold)} ({percent:.0f}%)", + ) + return f"Budget {budget_id}", _format_usd(spend) + + +def _budget_status_sort_key(entry: dict) -> float: + percent = _budget_usage_percent(entry) + return percent if percent is not None else -1.0 + + +def _budget_usage_percent(entry: dict) -> float | None: + spend_status = entry.get("spend_status") + spend_status = spend_status if isinstance(spend_status, dict) else {} + spend = _as_float(spend_status.get("spend")) + threshold = _as_float(spend_status.get("effective_alert_threshold")) + if spend is None or threshold is None or threshold <= 0: + return None + return (spend / threshold) * 100 + + +def _max_budget_usage_percent(statuses: list[dict]) -> float | None: + percents = [ + percent for entry in statuses if (percent := _budget_usage_percent(entry)) is not None + ] + return max(percents) if percents else None + + +def _primary_budget_status(statuses: list[dict]) -> dict: + return max(statuses, key=_budget_status_sort_key) + + +def _budget_policy_path() -> Path: + override = os.environ.get(BUDGET_POLICY_ENV_VAR, "").strip() + return Path(override).expanduser() if override else BUDGET_POLICY_PATH + + +def _read_budget_policy_file() -> tuple[dict | None, Path, str | None]: + path = _budget_policy_path() + if not path.exists(): + return None, path, None + try: + payload = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + return None, path, f"{path}: {exc}" + if not isinstance(payload, dict): + return None, path, f"{path}: top-level JSON value must be an object" + return payload, path, None + + +def _policy_percent(value: object, default: float) -> float: + percent = _as_float(value) + if percent is None or percent <= 0 or percent > 100: + return default + return percent + + +def _policy_tool(value: object, default: str) -> str: + raw: object + if isinstance(value, dict): + value_dict = cast(dict[str, object], value) + raw = value_dict.get("agent") or value_dict.get("tool") + else: + raw = value + if not isinstance(raw, str) or not raw.strip(): + return default + tool = TOOL_ALIASES.get(raw.strip().lower().replace(" ", "-")) + if tool in BUDGET_STATUS_TOOLS: + return tool + return default + + +def _policy_model(value: object) -> str | None: + if isinstance(value, dict): + value_dict = cast(dict[str, object], value) + value = value_dict.get("model") + if isinstance(value, str) and value.strip(): + return value.strip() + return None + + +def _policy_threshold(value: object) -> float | None: + percent = _as_float(value) + if percent is None or percent <= 0 or percent > 100: + return None + return percent + + +def _policy_tier_from_value(threshold: float, value: object) -> dict | None: + default_tool = SMART_BUDGET_LOW_TOOL + model: str | None = None + if isinstance(value, list) and value: + default_tool = _policy_tool(value[0], SMART_BUDGET_LOW_TOOL) + if len(value) > 1: + model = _policy_model(value[1]) + else: + default_tool = _policy_tool(value, SMART_BUDGET_LOW_TOOL) + model = _policy_model(value) + return {"threshold": threshold, "tool": default_tool, "model": model} + + +def _policy_tier_from_dict(value: dict) -> dict | None: + value_dict = cast(dict[str, object], value) + threshold = ( + _policy_threshold(value_dict.get("until_percent")) + or _policy_threshold(value_dict.get("threshold")) + or _policy_threshold(value_dict.get("percent")) + ) + if threshold is None: + return None + tool = _policy_tool(value_dict, SMART_BUDGET_LOW_TOOL) + model = _policy_model(value_dict) + return {"threshold": threshold, "tool": tool, "model": model} + + +def _policy_tiers(payload: dict) -> list[dict]: + tiers_raw = payload.get("tiers") + tiers: list[dict] = [] + if isinstance(tiers_raw, dict): + tiers_dict = cast(dict[object, object], tiers_raw) + for raw_threshold, raw_value in tiers_dict.items(): + threshold = _policy_threshold(raw_threshold) + if threshold is None: + continue + tier = _policy_tier_from_value(threshold, raw_value) + if tier: + tiers.append(tier) + elif isinstance(tiers_raw, list): + for raw_value in tiers_raw: + if isinstance(raw_value, dict): + tier = _policy_tier_from_dict(raw_value) + if tier: + tiers.append(tier) + + # Also support the compact shape the user sketched: + # {"20": {"agent": "claude", "model": "opus"}, ...} + for raw_threshold, raw_value in payload.items(): + threshold = _policy_threshold(raw_threshold) + if threshold is None: + continue + tier = _policy_tier_from_value(threshold, raw_value) + if tier: + tiers.append(tier) + + return sorted(tiers, key=lambda tier: float(tier["threshold"])) + + +def _policy_string_list(value: object) -> list[str]: + if isinstance(value, str): + stripped = value.strip() + return [stripped] if stripped else [] + if not isinstance(value, list): + return [] + out: list[str] = [] + for item in value: + if isinstance(item, str) and item.strip(): + out.append(item.strip()) + return out + + +def _policy_string_dict(value: object) -> dict[str, str]: + if not isinstance(value, dict): + return {} + value_dict = cast(dict[object, object], value) + out: dict[str, str] = {} + for key, val in value_dict.items(): + if isinstance(key, str) and key.strip() and isinstance(val, str) and val.strip(): + out[key.strip()] = val.strip() + return out + + +def _policy_bool(value: object, default: bool) -> bool: + return value if isinstance(value, bool) else default + + +def _load_budget_policy() -> dict: + payload, path, reason = _read_budget_policy_file() + if reason: + print_warning(f"Budget policy file ignored: {reason}") + payload = payload or {} + return { + "name": str( + payload.get("policy_name") or payload.get("name") or DEFAULT_BUDGET_POLICY_NAME + ), + "version": int(payload.get("version") or 1) + if isinstance(payload.get("version"), int) + else 1, + "switch_percent": _policy_percent( + payload.get("switch_percent") or payload.get("threshold_percent"), + SMART_BUDGET_SWITCH_PERCENT, + ), + "below_tool": _policy_tool(payload.get("below"), SMART_BUDGET_LOW_TOOL), + "at_or_above_tool": _policy_tool( + payload.get("at_or_above") or payload.get("above"), + SMART_BUDGET_HIGH_TOOL, + ), + "budget_ids": _policy_string_list( + payload.get("budget_ids") + or payload.get("filter_budget_ids") + or payload.get("budget") + or payload.get("budget_id") + ), + "tiers": _policy_tiers(payload), + "account_host": str(payload.get("account_host") or "").strip() or None, + "account_id": str(payload.get("account_id") or "").strip() or None, + "account_profile": str( + payload.get("account_profile") or payload.get("budget_profile") or "" + ).strip() + or None, + "budget_tags": _policy_string_dict( + payload.get("budget_tags") or payload.get("filter_budget_tags") + ), + "filter_has_spending": _policy_bool(payload.get("filter_has_spending"), False), + "source": str(path) if payload else "built-in", + } + + +def _budget_policy_summary(policy: dict) -> str: + tiers = policy.get("tiers") + if isinstance(tiers, list) and tiers: + tier_parts = [_tier_display(tier) for tier in tiers if isinstance(tier, dict)] + return " → ".join(tier_parts) + else: + switch_percent = float(policy["switch_percent"]) + below_tool = str(policy["below_tool"]) + above_tool = str(policy["at_or_above_tool"]) + return ( + f"{TOOL_SPECS[below_tool]['display']} under {switch_percent:.0f}%, " + f"{TOOL_SPECS[above_tool]['display']} at/after {switch_percent:.0f}%" + ) + + +def _budget_policy_filters( + policy: dict, +) -> tuple[str | None, str | None, str | None, list[str] | None, dict[str, str] | None, bool]: + account_host = policy.get("account_host") + account_id = policy.get("account_id") + account_profile = policy.get("account_profile") + filter_has_spending = bool(policy.get("filter_has_spending")) + filter_budget_ids: list[str] | None = None + filter_budget_tags: dict[str, str] | None = None + budget_ids = policy.get("budget_ids") + if isinstance(budget_ids, list) and budget_ids: + filter_budget_ids = [str(item) for item in budget_ids] + budget_tags = policy.get("budget_tags") + if isinstance(budget_tags, dict) and budget_tags: + filter_budget_tags = cast(dict[str, str], budget_tags) + return ( + account_host if isinstance(account_host, str) else None, + account_id if isinstance(account_id, str) else None, + account_profile if isinstance(account_profile, str) else None, + filter_budget_ids, + filter_budget_tags, + filter_has_spending, + ) + + +def _tier_display(tier: dict) -> str: + tool = str(tier["tool"]) + threshold = float(tier["threshold"]) + model = tier.get("model") + model_suffix = f"/{model}" if isinstance(model, str) and model else "" + return f"{threshold:.0f}% {TOOL_SPECS[tool]['display']}{model_suffix}" + + +def _active_budget_tier(policy: dict, percent: float) -> dict | None: + tiers = policy.get("tiers") + if not isinstance(tiers, list) or not tiers: + return None + tier_dicts = [tier for tier in tiers if isinstance(tier, dict)] + if not tier_dicts: + return None + for tier in tier_dicts: + if percent < float(tier["threshold"]): + return tier + return tier_dicts[-1] + + +def _smart_budget_policy_message(policy: dict, percent: float) -> tuple[str, str]: + tier = _active_budget_tier(policy, percent) + if tier: + tool = str(tier["tool"]) + threshold = float(tier["threshold"]) + comparator = "<" if percent < threshold else ">=" + model = tier.get("model") + model_suffix = f" / {model}" if isinstance(model, str) and model else "" + return ( + tool, + f"{TOOL_SPECS[tool]['display']}{model_suffix} " + f"(current {percent:.0f}% {comparator} {threshold:.0f}%)", + ) + + switch_percent = float(policy["switch_percent"]) + if percent < switch_percent: + return ( + str(policy["below_tool"]), + f"{TOOL_SPECS[str(policy['below_tool'])]['display']} " + f"(current {percent:.0f}% < {switch_percent:.0f}%)", + ) + return ( + str(policy["at_or_above_tool"]), + f"{TOOL_SPECS[str(policy['at_or_above_tool'])]['display']} " + f"(current {percent:.0f}% >= {switch_percent:.0f}%)", + ) + + +def _print_smart_budget_policy(tool: str, statuses: list[dict], policy: dict) -> None: + percent = _max_budget_usage_percent(statuses) + if percent is None: + return + recommended_tool, message = _smart_budget_policy_message(policy, percent) + if tool == recommended_tool: + print_kv("Recommended agent", message) + else: + print_warning(f"Budget policy recommends {message}") + + +def _budget_panel(policy: dict, statuses: list[dict], current_tool: str) -> Panel: + percent = _max_budget_usage_percent(statuses) + recommendation = "unknown" + recommended_tool = None + if percent is not None: + recommended_tool, recommendation = _smart_budget_policy_message(policy, percent) + + table = Table.grid(padding=(0, 2)) + table.add_column(style="bold") + table.add_column(style="cyan", overflow="fold") + table.add_row("Policy", _budget_policy_summary(policy)) + if percent is not None: + table.add_row("Usage", f"{percent:.1f}%") + table.add_row("Current", TOOL_SPECS[current_tool]["display"]) + table.add_row("Suggested", recommendation) + if recommended_tool and recommended_tool != current_tool: + table.add_row("Action", f"Switch to {TOOL_SPECS[recommended_tool]['display']}") + for entry in statuses: + key, value = _format_budget_status(entry) + table.add_row(key, value) + return Panel(table, title="Smart Budget Policy", style="bold blue", expand=False) + + +def _fetch_budget_statuses_for_policy( + state: dict, + policy: dict, +) -> tuple[list[dict], str | None]: + workspace = state.get("workspace") + if not isinstance(workspace, str) or not workspace.strip(): + return [], "no workspace configured" + ( + account_host, + account_id, + account_profile, + filter_budget_ids, + filter_budget_tags, + filter_has_spending, + ) = _budget_policy_filters(policy) + return fetch_current_user_budget_spend_status( + workspace, + state.get("profile"), + account_host=account_host, + account_id=account_id, + account_profile=account_profile, + filter_budget_ids=filter_budget_ids, + filter_budget_tags=filter_budget_tags, + filter_has_spending=filter_has_spending, + ) + + +def resolve_budget_policy_launch(state: dict) -> tuple[str, str | None] | None: + policy = _load_budget_policy() + statuses, reason = _fetch_budget_statuses_for_policy(state, policy) + if reason or not statuses: + return None + percent = _max_budget_usage_percent(statuses) + if percent is None: + return None + tier = _active_budget_tier(policy, percent) + if tier: + return str(tier["tool"]), _policy_model(tier) + recommended_tool, _ = _smart_budget_policy_message(policy, percent) + if recommended_tool == str(policy.get("below_tool")): + return recommended_tool, _policy_model(policy.get("below")) + return recommended_tool, _policy_model(policy.get("at_or_above") or policy.get("above")) + + +def print_budget_spend_status(tool: str, state: dict) -> None: + """Best-effort launch-time budget status for tools backed by AI Gateway.""" + if tool not in BUDGET_STATUS_TOOLS: + return + workspace = state.get("workspace") + if not isinstance(workspace, str) or not workspace.strip(): + return + policy = _load_budget_policy() + statuses, reason = _fetch_budget_statuses_for_policy(state, policy) + if reason: + print_kv("Budget policy", _budget_policy_summary(policy)) + print_warning(f"Budget spend unavailable: {reason}") + return + if not statuses: + print_kv("Budget policy", _budget_policy_summary(policy)) + print_kv("Budget spend", "no active AI Gateway budget found") + return + console.print() + console.print(_budget_panel(policy, statuses, tool)) + _print_smart_budget_policy(tool, statuses, policy) + + +def budget_hook_status_line(tool: str, state: dict) -> str: + """Compact, non-blocking budget status for agent lifecycle hooks.""" + policy = _load_budget_policy() + statuses, reason = _fetch_budget_statuses_for_policy(state, policy) + if reason: + return f"ucode budget: unavailable ({reason})" + if not statuses: + return "ucode budget: no active AI Gateway budget found" + percent = _max_budget_usage_percent(statuses) + recommendation = "unknown" + if percent is not None: + _, recommendation = _smart_budget_policy_message(policy, percent) + _, budget_value = _format_budget_status(_primary_budget_status(statuses)) + return f"ucode budget: {budget_value}; policy suggests {recommendation}" + + def launch(tool: str, state: dict, tool_args: list[str]) -> None: + print_budget_spend_status(tool, state) _MODULES[tool].launch(state, tool_args) diff --git a/src/ucode/agents/claude.py b/src/ucode/agents/claude.py index a6c1b40..aa75eef 100644 --- a/src/ucode/agents/claude.py +++ b/src/ucode/agents/claude.py @@ -10,6 +10,7 @@ from typing import cast from ucode.agent_updates import available_npm_package_update +from ucode.budget_hooks import budget_hook_handler, is_budget_hook_handler from ucode.config_io import ( APP_DIR, ToolSpec, @@ -255,6 +256,7 @@ def write_tool_config(state: dict, model: str) -> dict: profile=state.get("profile"), use_pat=bool(state.get("use_pat")), ) + managed_keys = managed_keys + [["hooks", "SessionStart"]] tracing_env_vars = tracing_env(state, "claude") stop_hook_command = claude_tracing_stop_hook_command() if tracing_env_vars else None if tracing_env_vars: @@ -272,6 +274,7 @@ def write_tool_config(state: dict, model: str) -> dict: existing = read_json_safe(CLAUDE_SETTINGS_PATH) merged = deep_merge_dict(existing, overlay) + _upsert_budget_session_start_hook(merged) if tracing_env_vars and stop_hook_command: _upsert_tracing_stop_hook(merged, stop_hook_command) if not tracing_env_vars: @@ -299,6 +302,57 @@ def write_tool_config(state: dict, model: str) -> dict: return state +def _remove_budget_session_start_hook(settings: dict) -> None: + hooks = settings.get("hooks") + if not isinstance(hooks, dict): + return + entries = hooks.get("SessionStart") + if not isinstance(entries, list): + return + + cleaned_entries = [] + for entry in entries: + if not isinstance(entry, dict): + cleaned_entries.append(entry) + continue + hook_list = entry.get("hooks") + if not isinstance(hook_list, list): + cleaned_entries.append(entry) + continue + cleaned_hooks = [ + hook for hook in hook_list if not is_budget_hook_handler(hook, "claude") + ] + if cleaned_hooks: + cleaned_entry = dict(entry) + cleaned_entry["hooks"] = cleaned_hooks + cleaned_entries.append(cleaned_entry) + + if cleaned_entries: + hooks["SessionStart"] = cleaned_entries + else: + hooks.pop("SessionStart", None) + if not hooks: + settings.pop("hooks", None) + + +def _upsert_budget_session_start_hook(settings: dict) -> None: + _remove_budget_session_start_hook(settings) + hooks = settings.get("hooks") + if not isinstance(hooks, dict): + hooks = {} + settings["hooks"] = hooks + entries = hooks.get("SessionStart") + if not isinstance(entries, list): + entries = [] + hooks["SessionStart"] = entries + entries.append( + { + "matcher": "startup|resume", + "hooks": [budget_hook_handler("claude")], + } + ) + + def _is_tracing_stop_hook(hook: object) -> bool: if not isinstance(hook, dict): return False diff --git a/src/ucode/agents/codex.py b/src/ucode/agents/codex.py index f1e3969..5c124ab 100644 --- a/src/ucode/agents/codex.py +++ b/src/ucode/agents/codex.py @@ -6,7 +6,10 @@ import re from pathlib import Path +import tomlkit + from ucode.agent_updates import available_npm_package_update +from ucode.budget_hooks import budget_hook_handler, is_budget_hook_handler from ucode.config_io import ( APP_DIR, ToolSpec, @@ -48,6 +51,7 @@ ["model"], ["model_providers", CODEX_MODEL_PROVIDER_NAME], ["model_providers", CODEX_MODEL_PROVIDER_NAME, "http_headers"], + ["hooks", "SessionStart"], ] LEGACY_MANAGED_KEYS: list[list[str]] = [ @@ -55,6 +59,7 @@ ["profiles", CODEX_PROFILE_NAME], ["model_providers", CODEX_MODEL_PROVIDER_NAME], ["model_providers", CODEX_MODEL_PROVIDER_NAME, "http_headers"], + ["hooks", "SessionStart"], ] _GPT_RE = re.compile(r"(?:databricks-)?gpt-(\d+)(?:[.-](\d+))?(?:[.-](\d+))?(-.+|[a-z].*)?") @@ -309,6 +314,7 @@ def write_tool_config(state: dict, model: str | None = None) -> dict: ) doc = read_toml_safe(LEGACY_CODEX_CONFIG_PATH) deep_merge_dict(doc, overlay) + _upsert_budget_session_start_hook(doc) write_toml_file(LEGACY_CODEX_CONFIG_PATH, doc) state = mark_tool_managed(state, "codex", LEGACY_MANAGED_KEYS) save_state(state) @@ -321,12 +327,65 @@ def write_tool_config(state: dict, model: str | None = None) -> dict: ) doc = read_toml_safe(CODEX_CONFIG_PATH) deep_merge_dict(doc, overlay) + _upsert_budget_session_start_hook(doc) write_toml_file(CODEX_CONFIG_PATH, doc) state = mark_tool_managed(state, "codex", MANAGED_KEYS) save_state(state) return state +def _remove_budget_session_start_hook(doc: dict) -> None: + hooks = doc.get("hooks") + if not isinstance(hooks, dict): + return + entries = hooks.get("SessionStart") + if not isinstance(entries, list): + return + + cleaned_entries = [] + for entry in entries: + if not isinstance(entry, dict): + cleaned_entries.append(entry) + continue + hook_list = entry.get("hooks") + if not isinstance(hook_list, list): + cleaned_entries.append(entry) + continue + cleaned_hooks = [hook for hook in hook_list if not is_budget_hook_handler(hook, "codex")] + if cleaned_hooks: + cleaned_entry = dict(entry) + cleaned_entry["hooks"] = cleaned_hooks + cleaned_entries.append(cleaned_entry) + + if cleaned_entries: + hooks["SessionStart"] = cleaned_entries + else: + hooks.pop("SessionStart", None) + if not hooks: + doc.pop("hooks", None) + + +def _upsert_budget_session_start_hook(doc: dict) -> None: + _remove_budget_session_start_hook(doc) + hooks = doc.get("hooks") + if not isinstance(hooks, dict): + hooks = tomlkit.table() + doc["hooks"] = hooks + entries = hooks.get("SessionStart") + if not isinstance(entries, list): + entries = tomlkit.aot() + hooks["SessionStart"] = entries + entry = tomlkit.table() + entry["matcher"] = "startup|resume" + handlers = tomlkit.aot() + handler = tomlkit.table() + for key, value in budget_hook_handler("codex").items(): + handler[key] = value + handlers.append(handler) + entry["hooks"] = handlers + entries.append(entry) + + def default_model(state: dict) -> str | None: """Pick the newest GPT model when multiple are available. diff --git a/src/ucode/budget_hooks.py b/src/ucode/budget_hooks.py new file mode 100644 index 0000000..e4902eb --- /dev/null +++ b/src/ucode/budget_hooks.py @@ -0,0 +1,32 @@ +"""Shared budget hook config for agent launchers.""" + +from __future__ import annotations + +import shlex +import shutil +from typing import cast + +BUDGET_HOOK_TIMEOUT_SECONDS = 10 +BUDGET_HOOK_STATUS_MESSAGE = "Loading ucode budget" + + +def budget_hook_command(tool: str) -> str: + ucode_binary = shutil.which("ucode") or "ucode" + return f"{shlex.quote(ucode_binary)} budget-hook --tool {shlex.quote(tool)}" + + +def budget_hook_handler(tool: str) -> dict: + return { + "type": "command", + "command": budget_hook_command(tool), + "timeout": BUDGET_HOOK_TIMEOUT_SECONDS, + "statusMessage": BUDGET_HOOK_STATUS_MESSAGE, + } + + +def is_budget_hook_handler(value: object, tool: str) -> bool: + if not isinstance(value, dict): + return False + value_dict = cast(dict[object, object], value) + command = value_dict.get("command") + return isinstance(command, str) and f" budget-hook --tool {tool}" in f" {command}" diff --git a/src/ucode/cli.py b/src/ucode/cli.py index ba2b790..6bb9bf2 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -10,6 +10,8 @@ from ucode.agents import ( TOOL_SPECS, + budget_hook_status_line, + budget_policy_env_configured, check_gateway_endpoint, configure_selected_tools, configure_single_tool, @@ -18,6 +20,7 @@ ensure_provider_state, install_tool_binary, normalize_tool, + resolve_budget_policy_launch, resolve_launch_model, validate_all_tools, validate_tool, @@ -582,7 +585,7 @@ def revert() -> int: app = typer.Typer( add_completion=False, - no_args_is_help=True, + no_args_is_help=False, context_settings={"allow_extra_args": True, "ignore_unknown_options": True}, ) configure_app = typer.Typer(add_completion=False, no_args_is_help=False) @@ -649,6 +652,22 @@ def auth_token_cmd( sys.stdout.write(token + "\n") +@app.command("budget-hook", hidden=True) +def budget_hook_cmd( + tool: Annotated[ + str, + typer.Option("--tool", help="Agent tool name for policy context."), + ] = "codex", +) -> None: + """Print compact budget status for agent lifecycle hooks.""" + try: + state = load_state() + apply_pat_environment(state) + console.print(budget_hook_status_line(normalize_tool(tool), state)) + except Exception as exc: # pragma: no cover - hook must not block launch + console.print(f"ucode budget: unavailable ({exc})") + + def _auto_configure_tool(tool: str) -> None: """First-time setup for a single tool — mirrors configure_workspace_command.""" existing = load_state() @@ -685,7 +704,12 @@ def _auto_configure_tool(tool: str) -> None: raise RuntimeError(f"{spec['display']} validation failed — config reverted.") -def _launch_tool(tool_name: str, ctx: typer.Context) -> None: +def _launch_tool( + tool_name: str, + ctx: typer.Context, + *, + explicit_model: str | None = None, +) -> None: try: tool = normalize_tool(tool_name) existing = load_state() @@ -707,8 +731,9 @@ def _launch_tool(tool_name: str, ctx: typer.Context) -> None: state = configure_shared_state( state["workspace"], profile=state.get("profile"), tools=[tool] ) - state, resolved_model = resolve_launch_model(tool, state, None) + state, resolved_model = resolve_launch_model(tool, state, explicit_model) state = configure_tool(tool, state, resolved_model) + tool_args = _launch_args_for_model(tool, ctx.args, explicit_model) print_section(f"ucode with {TOOL_SPECS[tool]['display']}") if resolved_model: print_kv("Model", resolved_model) @@ -718,7 +743,7 @@ def _launch_tool(tool_name: str, ctx: typer.Context) -> None: f"every 30 minutes while the session is running." ) print_success(f"Starting {TOOL_SPECS[tool]['display']}") - launch_agent(tool, state, ctx.args) + launch_agent(tool, state, tool_args) except RuntimeError as exc: print_err(str(exc)) raise typer.Exit(1) from None @@ -727,6 +752,56 @@ def _launch_tool(tool_name: str, ctx: typer.Context) -> None: raise typer.Exit(130) from None +def _launch_args_for_model( + tool: str, + tool_args: list[str], + explicit_model: str | None, +) -> list[str]: + if tool != "claude" or not explicit_model or _has_model_arg(tool_args): + return tool_args + return ["--model", _claude_model_selector(explicit_model), *tool_args] + + +def _has_model_arg(tool_args: list[str]) -> bool: + return any(arg == "--model" or arg.startswith("--model=") for arg in tool_args) + + +def _claude_model_selector(model: str) -> str: + lowered = model.lower() + for family in ("opus", "sonnet", "haiku"): + if family in lowered: + return family + return model + + +def _launch_from_budget_policy(ctx: typer.Context) -> None: + try: + state = load_state() + apply_pat_environment(state) + target = resolve_budget_policy_launch(state) + if target is None: + raise RuntimeError( + "Budget policy did not resolve a launch target. Check account auth, budget id, " + "and live spend status." + ) + tool, model = target + _launch_tool(tool, ctx, explicit_model=model) + except RuntimeError as exc: + print_err(str(exc)) + raise typer.Exit(1) from None + + +@app.callback(invoke_without_command=True) +def root(ctx: typer.Context) -> None: + if ctx.invoked_subcommand is not None: + return + if budget_policy_env_configured(): + _launch_from_budget_policy(ctx) + return + console.print(ctx.get_help()) + raise typer.Exit(0) + + @app.command("codex", context_settings={"allow_extra_args": True, "ignore_unknown_options": True}) def codex_cmd(ctx: typer.Context) -> None: """Launch Codex via Databricks.""" diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 250e9da..37d336f 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -51,6 +51,16 @@ AI_GATEWAY_V2_DOCS_URL = "https://docs.databricks.com/aws/en/ai-gateway/overview-beta" MIN_DATABRICKS_CLI_VERSION = (0, 298, 0) TOKEN_REFRESH_INTERVAL_SECONDS = 1800 +BUDGET_ACCOUNT_CONFIGS = { + "staging": { + "host": "accounts.staging.cloud.databricks.com", + "account_id": "11d77e21-5e05-4196-af72-423257f74974", + }, + "prod": { + "host": "accounts.cloud.databricks.com", + "account_id": "5a5e7a89-43b0-4d67-b90d-3fbdcfc7b28d", + }, +} def _debug_enabled() -> bool: @@ -877,6 +887,232 @@ def _fetch() -> str: return token +def budget_account_config_for_workspace(workspace: str) -> dict[str, str]: + """Return the account-scoped budget API config for a workspace URL.""" + hostname = workspace_hostname(workspace) + if "staging" in hostname: + return BUDGET_ACCOUNT_CONFIGS["staging"] + return BUDGET_ACCOUNT_CONFIGS["prod"] + + +def _budget_account_auth_url(account: dict[str, str]) -> str: + return f"https://{account['host']}?account_id={account['account_id']}" + + +def _account_token_profile_candidates( + account_auth_url: str, profile: str | None +) -> list[str | None]: + candidates: list[str | None] = [] + account_profile = find_account_profile_name(account_auth_url) + for candidate in (account_profile, profile, None): + if candidate not in candidates: + candidates.append(candidate) + return candidates + + +def find_account_profile_name(account_auth_url: str) -> str | None: + """Find a Databricks CLI profile for an account-host URL. + + Account auth profiles may persist hosts with extra query params such as + ``autoLogin=true``. Match by hostname and account_id instead of exact URL. + """ + parsed = urlparse(account_auth_url) + account_host = parsed.hostname + query_params = dict(part.split("=", 1) for part in parsed.query.split("&") if "=" in part) + account_id = query_params.get("account_id") + if not account_host or not account_id: + return find_profile_name_for_host(account_auth_url) + + for entry in list_profile_entries(): + name = entry.get("name") + host_raw = entry.get("host") + if not isinstance(name, str) or not isinstance(host_raw, str): + continue + host = host_raw.rstrip("/") + parsed_host = urlparse(host) + if parsed_host.hostname != account_host: + continue + if f"account_id={account_id}" in host: + return name + return find_profile_name_for_host(account_auth_url) + + +def _get_databricks_account_token( + account_auth_url: str, + profile: str | None = None, + *, + profile_only: bool = False, +) -> tuple[str | None, str | None]: + """Fetch an account-host OAuth token without attempting interactive re-auth. + + The launch path uses budget display as best-effort context; it must not open + a browser or block agent startup if the account-host profile is not logged in. + """ + bearer = os.environ.get("DATABRICKS_BEARER", "").strip() + if bearer: + _debug("_get_databricks_account_token", "using DATABRICKS_BEARER env var") + return bearer, None + + reasons: list[str] = [] + if profile and profile_only: + try: + result = run( + [ + "databricks", + "--profile", + profile, + "auth", + "token", + "--output", + "json", + ], + check=False, + capture_output=True, + text=True, + timeout=10, + ) + except (OSError, subprocess.TimeoutExpired) as exc: + reasons.append(f"{type(exc).__name__}: {exc}") + else: + _debug("_get_databricks_account_token.profile", _format_subprocess_result(result)) + if result.returncode == 0: + try: + token = json.loads(result.stdout or "{}").get("access_token", "") + except json.JSONDecodeError as exc: + reasons.append(f"invalid token JSON ({exc.msg})") + else: + if token: + return token, None + stderr = (result.stderr or "").strip() + reasons.append(stderr[:200] if stderr else "`databricks auth token` returned no token") + + for candidate in _account_token_profile_candidates(account_auth_url, profile): + env = build_databricks_cli_env(account_auth_url, candidate) + cmd = [ + "databricks", + "auth", + "token", + "--host", + account_auth_url, + *_profile_args(candidate), + "--output", + "json", + ] + try: + result = run( + cmd, + check=False, + capture_output=True, + text=True, + env=env, + timeout=10, + ) + except (OSError, subprocess.TimeoutExpired) as exc: + reasons.append(f"{type(exc).__name__}: {exc}") + continue + + _debug("_get_databricks_account_token", _format_subprocess_result(result)) + if result.returncode == 0: + try: + token = json.loads(result.stdout or "{}").get("access_token", "") + except json.JSONDecodeError as exc: + reasons.append(f"invalid token JSON ({exc.msg})") + continue + if token: + return token, None + + stderr = (result.stderr or "").strip() + if stderr: + reasons.append(stderr[:200]) + else: + reasons.append(f"`databricks auth token` exited {result.returncode}") + + detail = "; ".join(reason for reason in reasons if reason) or "no account token returned" + return ( + None, + f"account auth unavailable for {account_auth_url}. Run " + f"`databricks auth login --host {account_auth_url}`. Last error: {detail}", + ) + + +def _extract_budget_spend_statuses(payload: object) -> list[dict]: + if isinstance(payload, list): + return [item for item in payload if isinstance(item, dict)] + if not isinstance(payload, dict): + return [] + payload_dict = cast(dict[str, object], payload) + if payload_dict.get("budget_id"): + return [payload_dict] + for key in ( + "spend_statuses", + "budgets", + "budget_spend_statuses", + "budget_spend_status", + "statuses", + "results", + ): + value = payload_dict.get(key) + if isinstance(value, list): + return [item for item in value if isinstance(item, dict)] + return [] + + +def fetch_current_user_budget_spend_status( + workspace: str, + profile: str | None = None, + *, + account_host: str | None = None, + account_id: str | None = None, + account_profile: str | None = None, + filter_budget_ids: list[str] | None = None, + filter_budget_tags: dict[str, str] | None = None, + filter_has_spending: bool = False, +) -> tuple[list[dict], str | None]: + """Fetch current user's AI Gateway budget spend statuses from account APIs. + + Returns (statuses, None) on success and ([], reason) on failure. Status + entries are intentionally returned as raw API dicts because the RPC is new + and may grow fields without requiring a ucode release. + """ + account = dict(budget_account_config_for_workspace(workspace)) + if account_host: + account["host"] = account_host + if account_id: + account["account_id"] = account_id + account_id = account["account_id"] + account_auth_url = _budget_account_auth_url(account) + token, reason = _get_databricks_account_token( + account_auth_url, + account_profile or profile, + profile_only=bool(account_profile), + ) + if not token: + return [], reason or "could not fetch Databricks account token" + + url = ( + f"https://{account['host']}/api/2.1/accounts/{account_id}" + "/budgets:currentUserBudgetSpendStatus" + ) + request_payload: dict[str, object] = { + "account_id": account_id, + "filter_has_spending": filter_has_spending, + } + if filter_budget_ids: + request_payload["filter_budget_ids"] = filter_budget_ids + if filter_budget_tags: + request_payload["filter_budget_tags"] = filter_budget_tags + payload, reason = _http_post_json( + url, + token, + request_payload, + timeout=10, + ) + if reason: + return [], reason + statuses = _extract_budget_spend_statuses(payload) + return statuses, None + + def _extract_connection_page(payload: object) -> tuple[list[dict], str | None]: if isinstance(payload, list): return [item for item in payload if isinstance(item, dict)], None diff --git a/tests/test_agent_claude.py b/tests/test_agent_claude.py index 1c801b9..99716c9 100644 --- a/tests/test_agent_claude.py +++ b/tests/test_agent_claude.py @@ -282,6 +282,46 @@ def test_explicit_override_used_over_codex_models(self, monkeypatch): assert calls == [("register", WS, "explicit-model")] +class TestClaudeBudgetHook: + def test_writes_session_start_budget_hook(self, monkeypatch): + written: list[dict] = [] + monkeypatch.setattr(claude, "backup_existing_file", lambda *a, **kw: True) + monkeypatch.setattr(claude, "read_json_safe", lambda path: {}) + monkeypatch.setattr(claude, "write_json_file", lambda path, payload: written.append(payload)) + monkeypatch.setattr(claude, "save_state", lambda state: None) + + claude.write_tool_config({"workspace": WS, "codex_models": []}, "databricks-claude-sonnet-4") + + hook = written[0]["hooks"]["SessionStart"][0]["hooks"][0] + assert hook["type"] == "command" + assert "budget-hook --tool claude" in hook["command"] + assert hook["timeout"] == 10 + + def test_preserves_user_session_start_hooks(self, monkeypatch): + existing = { + "hooks": { + "SessionStart": [ + {"hooks": [{"type": "command", "command": "user-session-start"}]} + ] + } + } + written: list[dict] = [] + monkeypatch.setattr(claude, "backup_existing_file", lambda *a, **kw: True) + monkeypatch.setattr(claude, "read_json_safe", lambda path: existing) + monkeypatch.setattr(claude, "write_json_file", lambda path, payload: written.append(payload)) + monkeypatch.setattr(claude, "save_state", lambda state: None) + + claude.write_tool_config({"workspace": WS, "codex_models": []}, "databricks-claude-sonnet-4") + + commands = [ + hook["command"] + for entry in written[0]["hooks"]["SessionStart"] + for hook in entry["hooks"] + ] + assert "user-session-start" in commands + assert sum("budget-hook --tool claude" in command for command in commands) == 1 + + class TestRegisterWebSearchMcp: def test_clears_existing_then_adds(self, monkeypatch): import ucode.mcp as mcp_mod diff --git a/tests/test_agent_codex.py b/tests/test_agent_codex.py index d1fb932..a921848 100644 --- a/tests/test_agent_codex.py +++ b/tests/test_agent_codex.py @@ -94,6 +94,52 @@ def test_writes_ucode_profile_config_file(self, tmp_path, monkeypatch): assert doc["model"] == "gpt-5" assert "profiles" not in doc + def test_writes_session_start_budget_hook(self, tmp_path, monkeypatch): + config_path = tmp_path / ".codex" / "ucode.config.toml" + backup_path = tmp_path / "codex-ucode-config.backup.toml" + monkeypatch.setattr(codex, "CODEX_CONFIG_PATH", config_path) + monkeypatch.setattr(codex, "CODEX_BACKUP_PATH", backup_path) + monkeypatch.setattr(codex, "agent_version", lambda binary: "0.134.0") + monkeypatch.setattr(codex, "save_state", lambda state: None) + + codex.write_tool_config({"workspace": WS, "codex_models": ["gpt-5"]}) + + doc = read_toml_safe(config_path) + hook = doc["hooks"]["SessionStart"][0]["hooks"][0] + assert hook["type"] == "command" + assert "budget-hook --tool codex" in hook["command"] + assert hook["timeout"] == 10 + + def test_preserves_user_session_start_hooks(self, tmp_path, monkeypatch): + config_path = tmp_path / ".codex" / "ucode.config.toml" + config_path.parent.mkdir() + config_path.write_text( + """ + [[hooks.SessionStart]] + + [[hooks.SessionStart.hooks]] + type = "command" + command = "user-session-start" + """, + encoding="utf-8", + ) + backup_path = tmp_path / "codex-ucode-config.backup.toml" + monkeypatch.setattr(codex, "CODEX_CONFIG_PATH", config_path) + monkeypatch.setattr(codex, "CODEX_BACKUP_PATH", backup_path) + monkeypatch.setattr(codex, "agent_version", lambda binary: "0.134.0") + monkeypatch.setattr(codex, "save_state", lambda state: None) + + codex.write_tool_config({"workspace": WS, "codex_models": ["gpt-5"]}) + + doc = read_toml_safe(config_path) + commands = [ + hook["command"] + for entry in doc["hooks"]["SessionStart"] + for hook in entry["hooks"] + ] + assert "user-session-start" in commands + assert sum("budget-hook --tool codex" in command for command in commands) == 1 + def test_writes_openai_model_id_for_databricks_gpt_endpoint(self, tmp_path, monkeypatch): config_path = tmp_path / ".codex" / "ucode.config.toml" backup_path = tmp_path / "codex-ucode-config.backup.toml" diff --git a/tests/test_agents_init.py b/tests/test_agents_init.py index e4560f4..7882992 100644 --- a/tests/test_agents_init.py +++ b/tests/test_agents_init.py @@ -170,6 +170,32 @@ def test_explicit_model_used_when_provided(self): _, model = resolve_launch_model("claude", {}, "my-model") assert model == "my-model" + def test_claude_policy_model_alias_resolves_family(self): + state = {"claude_models": {"opus": "system.ai.claude-opus-4-8"}} + _, model = resolve_launch_model("claude", state, "opus") + assert model == "system.ai.claude-opus-4-8" + + def test_codex_policy_model_alias_resolves_available_model(self): + state = {"codex_models": ["system.ai.gpt-5-4-mini"]} + _, model = resolve_launch_model("codex", state, "gpt-5-4-mini") + assert model == "system.ai.gpt-5-4-mini" + + def test_codex_policy_model_keeps_system_ai_prefix_when_no_match(self): + # The policy names a model the workspace did not discover, but the + # workspace routes by `system.ai.`. Keep that prefix rather than + # falling through to the OpenAI-id rewrite (`gpt-5.4-mini`), which is + # unroutable on a model-services gateway. + state = {"codex_models": ["system.ai.gpt-5", "system.ai.gpt-5-mini"]} + _, model = resolve_launch_model("codex", state, "gpt-5-4-mini") + assert model == "system.ai.gpt-5-4-mini" + + def test_codex_policy_model_unchanged_on_foundation_models_workspace(self): + # Legacy foundation-models workspaces use `databricks-*` endpoint names; + # an unmatched bare id must not gain a spurious `system.ai.` prefix. + state = {"codex_models": ["databricks-gpt-5-4"]} + _, model = resolve_launch_model("codex", state, "gpt-5-4-mini") + assert model == "gpt-5-4-mini" + def test_default_model_used_when_no_explicit(self): state = {"claude_models": {"sonnet": "s4"}} _, model = resolve_launch_model("claude", state, None) @@ -477,6 +503,234 @@ def test_empty_selection_preserves_existing(self, monkeypatch): assert result["available_tools"] == ["codex"] +class TestLaunchBudgetStatus: + def test_prints_budget_status_for_codex(self, monkeypatch, capsys): + launched: list[tuple[dict, list[str]]] = [] + monkeypatch.setenv("UCODE_BUDGET_POLICY", "/tmp/ucode-test-missing-policy.json") + + monkeypatch.setattr( + agents_mod, + "fetch_current_user_budget_spend_status", + lambda workspace, profile=None, **kwargs: ( + [ + { + "budget_id": "abcdef12-3456", + "spend_status": { + "spend": 3.5, + "effective_alert_threshold": 10, + }, + } + ], + None, + ), + ) + monkeypatch.setattr( + agents_mod._MODULES["codex"], + "launch", + lambda state, args: launched.append((state, args)), + ) + + state = {"workspace": "https://x.databricks.com", "profile": "acct-profile"} + agents_mod.launch("codex", state, ["--search"]) + + assert launched == [(state, ["--search"])] + out = capsys.readouterr().out + assert "Budget abcdef12" in out + assert "$3.50 / $10.00 (35%)" in out + assert "Claude Code under 80%" in out + assert "Codex" in out and "at/after 80%" in out + assert "Budget policy recommends Claude Code (current 35% < 80%)" in out + + def test_budget_policy_recommends_codex_at_or_after_80_percent(self, monkeypatch, capsys): + launched: list[bool] = [] + monkeypatch.setenv("UCODE_BUDGET_POLICY", "/tmp/ucode-test-missing-policy.json") + monkeypatch.setattr( + agents_mod, + "fetch_current_user_budget_spend_status", + lambda workspace, profile=None, **kwargs: ( + [ + { + "budget_id": "budget-1", + "spend_status": { + "spend": 8, + "effective_alert_threshold": 10, + }, + } + ], + None, + ), + ) + monkeypatch.setattr( + agents_mod._MODULES["claude"], + "launch", + lambda state, args: launched.append(True), + ) + + agents_mod.launch("claude", {"workspace": "https://x.databricks.com"}, []) + + assert launched == [True] + out = capsys.readouterr().out + assert "Budget policy recommends Codex (current 80% >= 80%)" in out + + def test_reads_custom_budget_policy_json(self, tmp_path, monkeypatch, capsys): + policy_path = tmp_path / "budget-policy.json" + calls: list[dict] = [] + policy_path.write_text( + """ + { + "name": "Team budget policy", + "switch_percent": 60, + "below": {"agent": "codex"}, + "at_or_above": {"agent": "claude"}, + "budget_ids": ["3f142f2c-c75d-495d-8b62-25f7ac37bf6d"], + "budget_tags": {"demo": "true"}, + "filter_has_spending": true + } + """, + encoding="utf-8", + ) + monkeypatch.setenv("UCODE_BUDGET_POLICY", str(policy_path)) + monkeypatch.setattr( + agents_mod, + "fetch_current_user_budget_spend_status", + lambda workspace, profile=None, **kwargs: ( + calls.append(kwargs) + or ( + [ + { + "budget_id": "3f142f2c-c75d-495d-8b62-25f7ac37bf6d", + "spend_status": { + "spend": 5, + "effective_alert_threshold": 10, + }, + } + ], + None, + ) + ), + ) + monkeypatch.setattr( + agents_mod._MODULES["claude"], + "launch", + lambda state, args: None, + ) + + agents_mod.launch("claude", {"workspace": "https://x.databricks.com"}, []) + + out = capsys.readouterr().out + assert "Codex under 60%" in out + assert "Budget policy recommends Codex (current 50% < 60%)" in out + assert calls == [ + { + "account_host": None, + "account_id": None, + "account_profile": None, + "filter_budget_ids": ["3f142f2c-c75d-495d-8b62-25f7ac37bf6d"], + "filter_budget_tags": {"demo": "true"}, + "filter_has_spending": True, + } + ] + + def test_reads_tiered_budget_policy_json(self, tmp_path, monkeypatch, capsys): + policy_path = tmp_path / "budget-policy.json" + calls: list[dict] = [] + policy_path.write_text( + """ + { + "version": 1, + "budget": "3f142f2c-c75d-495d-8b62-25f7ac37bf6d", + "policy_name": "Demo tiered policy", + "tiers": { + "20": {"agent": "claude code", "model": "opus"}, + "60": {"agent": "claude code", "model": "sonnet"}, + "80": {"agent": "codex", "model": "gpt-5-4-mini"} + } + } + """, + encoding="utf-8", + ) + monkeypatch.setenv("UCODE_BUDGET_POLICY", str(policy_path)) + monkeypatch.setattr( + agents_mod, + "fetch_current_user_budget_spend_status", + lambda workspace, profile=None, **kwargs: ( + calls.append(kwargs) + or ( + [ + { + "budget_id": "3f142f2c-c75d-495d-8b62-25f7ac37bf6d", + "spend_status": { + "spend": 250, + "effective_alert_threshold": 500, + }, + } + ], + None, + ) + ), + ) + monkeypatch.setattr( + agents_mod._MODULES["codex"], + "launch", + lambda state, args: None, + ) + + agents_mod.launch("codex", {"workspace": "https://x.databricks.com"}, []) + + out = capsys.readouterr().out + assert "20% Claude Code/opus" in out + assert "60% Claude" in out and "Code/sonnet" in out + assert "80%" in out and "Codex/gpt-5-4-mini" in out + assert "Budget policy recommends Claude Code / sonnet (current 50% < 60%)" in out + assert calls == [ + { + "account_host": None, + "account_id": None, + "account_profile": None, + "filter_budget_ids": ["3f142f2c-c75d-495d-8b62-25f7ac37bf6d"], + "filter_budget_tags": None, + "filter_has_spending": False, + } + ] + + def test_prints_budget_warning_without_blocking_launch(self, monkeypatch, capsys): + launched: list[bool] = [] + monkeypatch.setattr( + agents_mod, + "fetch_current_user_budget_spend_status", + lambda workspace, profile=None, **kwargs: ([], "login required"), + ) + monkeypatch.setattr( + agents_mod._MODULES["claude"], + "launch", + lambda state, args: launched.append(True), + ) + + agents_mod.launch("claude", {"workspace": "https://x.databricks.com"}, []) + + assert launched == [True] + assert "Budget spend unavailable: login required" in capsys.readouterr().out + + def test_skips_budget_status_for_other_tools(self, monkeypatch): + fetched: list[bool] = [] + launched: list[bool] = [] + monkeypatch.setattr( + agents_mod, + "fetch_current_user_budget_spend_status", + lambda workspace, profile=None, **kwargs: fetched.append(True) or ([], None), + ) + monkeypatch.setattr( + agents_mod._MODULES["gemini"], + "launch", + lambda state, args: launched.append(True), + ) + + agents_mod.launch("gemini", {"workspace": "https://x.databricks.com"}, []) + + assert fetched == [] + assert launched == [True] + + class TestValidateAllToolsVerbosity: def _run(self, monkeypatch, capsys): from contextlib import nullcontext diff --git a/tests/test_cli.py b/tests/test_cli.py index cda0721..bc11f86 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -143,6 +143,55 @@ def test_no_agent_flag(self): result = runner.invoke(app, ["--agent", "claude"]) assert result.exit_code != 0 + def test_bare_ucode_uses_budget_policy_target(self): + with ( + patch("ucode.cli.load_state", return_value=MINIMAL_STATE), + patch( + "ucode.cli.resolve_budget_policy_launch", + return_value=("codex", "gpt-5-4-mini"), + ), + patch("ucode.cli._launch_tool") as mock_launch, + ): + result = runner.invoke( + app, + [], + env={"UCODE_BUDGET_POLICY": "/tmp/budget-policy.json"}, + ) + + assert result.exit_code == 0, result.output + mock_launch.assert_called_once() + assert mock_launch.call_args.args[0] == "codex" + assert mock_launch.call_args.kwargs["explicit_model"] == "gpt-5-4-mini" + + def test_bare_ucode_pins_claude_policy_model(self): + with ( + patch("ucode.cli.load_state", return_value=MINIMAL_STATE), + patch("ucode.cli.apply_pat_environment"), + patch("ucode.cli.resolve_budget_policy_launch", return_value=("claude", "opus")), + patch("ucode.cli.ensure_bootstrap_dependencies"), + patch("ucode.cli.ensure_provider_state", return_value=MINIMAL_STATE), + patch("ucode.cli.configure_shared_state", return_value=MINIMAL_STATE), + patch( + "ucode.cli.resolve_launch_model", + return_value=(MINIMAL_STATE, "system.ai.claude-opus-4-8"), + ), + patch("ucode.cli.configure_tool", return_value=MINIMAL_STATE) as mock_configure, + patch("ucode.cli.launch_agent") as mock_launch, + ): + result = runner.invoke( + app, + [], + env={"UCODE_BUDGET_POLICY": "/tmp/budget-policy.json"}, + ) + + assert result.exit_code == 0, result.output + mock_configure.assert_called_once_with( + "claude", MINIMAL_STATE, "system.ai.claude-opus-4-8" + ) + mock_launch.assert_called_once() + assert mock_launch.call_args.args[0] == "claude" + assert mock_launch.call_args.args[2] == ["--model", "opus"] + class TestMcpSubcommands: def test_web_search_subcommand_help(self): @@ -267,6 +316,19 @@ def test_use_pat_honors_non_empty_bearer_env(self, monkeypatch): assert result.stdout == "ci-bearer\n" +class TestBudgetHookCommand: + def test_budget_hook_prints_compact_status(self): + with ( + patch("ucode.cli.load_state", return_value=MINIMAL_STATE), + patch("ucode.cli.apply_pat_environment"), + patch("ucode.cli.budget_hook_status_line", return_value="ucode budget: 5%"), + ): + result = runner.invoke(app, ["budget-hook", "--tool", "claude"]) + + assert result.exit_code == 0, result.output + assert "ucode budget: 5%" in result.output + + class TestStatus: def test_shows_mcp_list_commands(self): with patch("ucode.cli.load_state", return_value=MINIMAL_STATE): diff --git a/tests/test_databricks.py b/tests/test_databricks.py index da630b0..15a5c7b 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -24,6 +24,7 @@ build_tool_base_url, ensure_databricks_cli_version, ensure_pat_bearer, + fetch_current_user_budget_spend_status, get_databricks_token, list_databricks_apps, list_databricks_connections, @@ -117,6 +118,170 @@ def test_codex_url_format(self): assert urls["codex"] == f"{WS}/ai-gateway/codex/v1" +class TestBudgetSpendStatus: + def test_selects_prod_account_by_default(self): + account = db_mod.budget_account_config_for_workspace(WS) + + assert account["host"] == "accounts.cloud.databricks.com" + assert account["account_id"] == "5a5e7a89-43b0-4d67-b90d-3fbdcfc7b28d" + + def test_selects_staging_account_for_staging_workspace(self): + account = db_mod.budget_account_config_for_workspace( + "https://workspace.staging.cloud.databricks.com" + ) + + assert account["host"] == "accounts.staging.cloud.databricks.com" + assert account["account_id"] == "11d77e21-5e05-4196-af72-423257f74974" + + def test_finds_account_profile_by_host_and_account_id(self, monkeypatch): + monkeypatch.setattr( + db_mod, + "list_profile_entries", + lambda: [ + { + "name": "eng-ml-inference-account-staging", + "host": ( + "https://accounts.staging.cloud.databricks.com/" + "?autoLogin=true&account_id=11d77e21-5e05-4196-af72-423257f74974" + ), + } + ], + ) + + profile = db_mod.find_account_profile_name( + "https://accounts.staging.cloud.databricks.com" + "?account_id=11d77e21-5e05-4196-af72-423257f74974" + ) + + assert profile == "eng-ml-inference-account-staging" + + def test_fetch_posts_current_user_budget_rpc(self, monkeypatch): + calls: list[tuple[str, str, dict, int]] = [] + payload = { + "spend_statuses": [ + { + "budget_id": "abc", + "spend_status": { + "user_id": "user-1", + "spend": 12.5, + "effective_alert_threshold": 20, + }, + } + ] + } + + monkeypatch.setattr( + db_mod, + "_get_databricks_account_token", + lambda account_auth_url, profile=None, **kwargs: ("account-token", None), + ) + + def fake_post(url, token, body, timeout=10): + calls.append((url, token, body, timeout)) + return payload, None + + monkeypatch.setattr(db_mod, "_http_post_json", fake_post) + + statuses, reason = fetch_current_user_budget_spend_status(WS, "acct-profile") + + assert reason is None + assert statuses == payload["spend_statuses"] + assert calls == [ + ( + "https://accounts.cloud.databricks.com/api/2.1/accounts/" + "5a5e7a89-43b0-4d67-b90d-3fbdcfc7b28d/" + "budgets:currentUserBudgetSpendStatus", + "account-token", + { + "account_id": "5a5e7a89-43b0-4d67-b90d-3fbdcfc7b28d", + "filter_has_spending": False, + }, + 10, + ) + ] + + def test_fetch_returns_account_auth_reason(self, monkeypatch): + monkeypatch.setattr( + db_mod, + "_get_databricks_account_token", + lambda account_auth_url, profile=None, **kwargs: (None, "login required"), + ) + + statuses, reason = fetch_current_user_budget_spend_status(WS) + + assert statuses == [] + assert reason == "login required" + + def test_fetch_includes_optional_filters(self, monkeypatch): + bodies: list[dict] = [] + monkeypatch.setattr( + db_mod, + "_get_databricks_account_token", + lambda account_auth_url, profile=None, **kwargs: ("account-token", None), + ) + + def fake_post(url, token, body, timeout=10): + bodies.append(body) + return {"spend_statuses": []}, None + + monkeypatch.setattr(db_mod, "_http_post_json", fake_post) + + statuses, reason = fetch_current_user_budget_spend_status( + WS, + filter_budget_ids=["budget-1"], + filter_budget_tags={"demo": "true"}, + filter_has_spending=True, + ) + + assert statuses == [] + assert reason is None + assert bodies == [ + { + "account_id": "5a5e7a89-43b0-4d67-b90d-3fbdcfc7b28d", + "filter_has_spending": True, + "filter_budget_ids": ["budget-1"], + "filter_budget_tags": {"demo": "true"}, + } + ] + + def test_fetch_uses_account_profile_without_host_override(self, monkeypatch): + runs: list[list[str]] = [] + monkeypatch.delenv("DATABRICKS_BEARER", raising=False) + + def fake_run(args, **kwargs): + runs.append(args) + return subprocess.CompletedProcess( + args, + 0, + stdout=json.dumps({"access_token": "account-token"}), + stderr="", + ) + + monkeypatch.setattr(db_mod, "run", fake_run) + monkeypatch.setattr( + db_mod, + "_http_post_json", + lambda url, token, body, timeout=10: ({"spend_statuses": []}, None), + ) + + statuses, reason = fetch_current_user_budget_spend_status( + WS, + account_host="accounts.staging.cloud.databricks.com", + account_id="11d77e21-5e05-4196-af72-423257f74974", + account_profile="eng-ml-inference-account-staging", + ) + + assert statuses == [] + assert reason is None + assert runs[0][:4] == [ + "databricks", + "--profile", + "eng-ml-inference-account-staging", + "auth", + ] + assert "--host" not in runs[0] + + class TestDiscoverClaudeModels: def test_selects_opus_4_8_when_advertised(self, monkeypatch): payload = { From 5f1881932b7919617bdeea2fe2cb8057a9d9ec0e Mon Sep 17 00:00:00 2001 From: Rohit Agrawal Date: Sun, 5 Jul 2026 21:27:54 -0400 Subject: [PATCH 2/2] budget demo --- src/ucode/agents/__init__.py | 10 ++++ src/ucode/cli.py | 43 +++++++++++++- src/ucode/databricks.py | 38 +++++++++++- tests/test_agents_init.py | 18 ++++++ tests/test_cli.py | 111 +++++++++++++++++++++++++++++++++-- 5 files changed, 213 insertions(+), 7 deletions(-) diff --git a/src/ucode/agents/__init__.py b/src/ucode/agents/__init__.py index d92692f..f1facf1 100644 --- a/src/ucode/agents/__init__.py +++ b/src/ucode/agents/__init__.py @@ -90,6 +90,16 @@ def budget_policy_env_configured() -> bool: return bool(os.environ.get(BUDGET_POLICY_ENV_VAR, "").strip()) +def budget_policy_configured() -> bool: + """True when a budget policy is active via env override or the default file. + + The env var only selects an alternate policy path; an admin who drops a + policy at the default location (~/.ucode/budget-policy.json) has configured + one just as much, so enforcement must key off either source. + """ + return budget_policy_env_configured() or _budget_policy_path().exists() + + def _update_installed_tool_binary(tool: str, version: str | None = None) -> bool: spec = TOOL_SPECS[tool] binary = spec["binary"] diff --git a/src/ucode/cli.py b/src/ucode/cli.py index 6bb9bf2..85127dd 100644 --- a/src/ucode/cli.py +++ b/src/ucode/cli.py @@ -11,6 +11,7 @@ from ucode.agents import ( TOOL_SPECS, budget_hook_status_line, + budget_policy_configured, budget_policy_env_configured, check_gateway_endpoint, configure_selected_tools, @@ -704,11 +705,41 @@ def _auto_configure_tool(tool: str) -> None: raise RuntimeError(f"{spec['display']} validation failed — config reverted.") +def _enforce_budget_policy(tool: str, state: dict) -> tuple[str, str | None] | None: + """Return a (tool, model) redirect when the budget policy disallows `tool`. + + Returns None when the requested tool is allowed, or when the policy target + cannot be resolved (a transient auth/spend fetch failure must not lock the + user out of every agent — the launch-time status still surfaces a warning). + """ + target = resolve_budget_policy_launch(state) + if target is None: + return None + recommended_tool, recommended_model = target + if recommended_tool == tool: + return None + model_suffix = f" / {recommended_model}" if recommended_model else "" + console.print( + Panel( + f"[bold]{TOOL_SPECS[tool]['display']}[/bold] is [red]not allowed[/red] at the " + "current budget spend.\n" + f"Budget policy requires " + f"[cyan]{TOOL_SPECS[recommended_tool]['display']}{model_suffix}[/cyan] — " + "launching that instead.", + title="Budget Policy Enforced", + style="yellow", + expand=False, + ) + ) + return recommended_tool, recommended_model + + def _launch_tool( tool_name: str, ctx: typer.Context, *, explicit_model: str | None = None, + enforce_budget_policy: bool = True, ) -> None: try: tool = normalize_tool(tool_name) @@ -717,6 +748,14 @@ def _launch_tool( # DATABRICKS_BEARER up front so every auth check below (and the # launched agent itself) uses the static token instead of OAuth. apply_pat_environment(existing) + # A configured budget policy is prescriptive, not advisory: block an + # explicit launch of a disallowed harness and redirect to the required + # tool/model before any auto-configure or validation work happens, so + # setup targets the tool that will actually run. + if enforce_budget_policy and budget_policy_configured(): + redirect = _enforce_budget_policy(tool, existing) + if redirect is not None: + tool, explicit_model = redirect needs_auto_configure = not existing.get("workspace") or tool not in ( existing.get("available_tools") or [] ) @@ -785,7 +824,9 @@ def _launch_from_budget_policy(ctx: typer.Context) -> None: "and live spend status." ) tool, model = target - _launch_tool(tool, ctx, explicit_model=model) + # The target is already the policy's resolved tool/model, so skip the + # per-launch enforcement redirect (it would just resolve to itself). + _launch_tool(tool, ctx, explicit_model=model, enforce_budget_policy=False) except RuntimeError as exc: print_err(str(exc)) raise typer.Exit(1) from None diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 37d336f..ce9f62c 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -14,6 +14,7 @@ import shlex import shutil import subprocess +import tempfile import time from concurrent.futures import ( ThreadPoolExecutor, @@ -749,6 +750,38 @@ def apply_pat_environment(state: dict) -> None: ensure_pat_bearer(state.get("profile")) +def _incognito_login_env(env: dict[str, str]) -> dict[str, str]: + """Shim the `open` command so `databricks auth login` launches Chrome in + incognito mode, instead of the user's regular (possibly stale-session) + default browser. + + macOS only: `databricks auth login` opens the OAuth URL via the `open` + binary internally, so we prepend a directory with a fake `open` to PATH + that special-cases http(s) URLs and passes everything else through. + """ + if platform.system() != "Darwin": + return env + chrome = "/Applications/Google Chrome.app" + if not os.path.isdir(chrome): + return env + shim_dir = tempfile.mkdtemp(prefix="ucode-open-shim-") + shim_path = os.path.join(shim_dir, "open") + with Path(shim_path).open("w") as fh: + fh.write( + "#!/bin/sh\n" + 'case "$1" in\n' + " http://*|https://*)\n" + ' exec /usr/bin/open -na "Google Chrome" --args --incognito "$1" ;;\n' + " *)\n" + ' exec /usr/bin/open "$@" ;;\n' + "esac\n" + ) + os.chmod(shim_path, 0o755) + env = dict(env) + env["PATH"] = shim_dir + os.pathsep + env.get("PATH", "") + return env + + def run_databricks_login(workspace: str, profile: str | None = None) -> None: """Run databricks auth login unconditionally. @@ -757,7 +790,7 @@ def run_databricks_login(workspace: str, profile: str | None = None) -> None: refreshed in place rather than overwriting another profile's tokens.""" print_section("Databricks Login") print_kv("Workspace", workspace) - print_note("A browser may open for `databricks auth login`.") + print_note("A browser may open for `databricks auth login` (Chrome incognito, if available).") try: profile_name = profile or find_profile_name_for_host(workspace) cmd = [ @@ -768,7 +801,8 @@ def run_databricks_login(workspace: str, profile: str | None = None) -> None: workspace, *_profile_args(profile_name), ] - run(cmd, env=build_databricks_cli_env(workspace, profile_name), timeout=300) + env = _incognito_login_env(build_databricks_cli_env(workspace, profile_name)) + run(cmd, env=env, timeout=300) except subprocess.CalledProcessError as exc: raise RuntimeError("`databricks auth login` failed.") from exc except subprocess.TimeoutExpired as exc: diff --git a/tests/test_agents_init.py b/tests/test_agents_init.py index 7882992..6c6293b 100644 --- a/tests/test_agents_init.py +++ b/tests/test_agents_init.py @@ -503,6 +503,24 @@ def test_empty_selection_preserves_existing(self, monkeypatch): assert result["available_tools"] == ["codex"] +class TestBudgetPolicyConfigured: + def test_true_when_env_set(self, monkeypatch, tmp_path): + monkeypatch.setenv("UCODE_BUDGET_POLICY", str(tmp_path / "missing.json")) + assert agents_mod.budget_policy_configured() is True + + def test_true_when_default_file_exists(self, monkeypatch, tmp_path): + monkeypatch.delenv("UCODE_BUDGET_POLICY", raising=False) + policy_path = tmp_path / "budget-policy.json" + policy_path.write_text("{}", encoding="utf-8") + monkeypatch.setattr(agents_mod, "BUDGET_POLICY_PATH", policy_path) + assert agents_mod.budget_policy_configured() is True + + def test_false_when_no_env_and_no_file(self, monkeypatch, tmp_path): + monkeypatch.delenv("UCODE_BUDGET_POLICY", raising=False) + monkeypatch.setattr(agents_mod, "BUDGET_POLICY_PATH", tmp_path / "budget-policy.json") + assert agents_mod.budget_policy_configured() is False + + class TestLaunchBudgetStatus: def test_prints_budget_status_for_codex(self, monkeypatch, capsys): launched: list[tuple[dict, list[str]]] = [] diff --git a/tests/test_cli.py b/tests/test_cli.py index bc11f86..1f91df7 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -4,6 +4,7 @@ import os import re +from contextlib import ExitStack from unittest.mock import patch import pytest @@ -60,7 +61,9 @@ def no_state_writes(): class TestHelp: def test_no_args_shows_help(self): - result = runner.invoke(app, []) + # Clear any policy env inherited from the host shell, else bare `ucode` + # takes the budget-policy launch path instead of printing help. + result = runner.invoke(app, [], env={"UCODE_BUDGET_POLICY": ""}) # no_args_is_help=True exits with code 0 or 2 depending on typer version assert result.exit_code in (0, 2) assert "Usage:" in result.output @@ -185,13 +188,113 @@ def test_bare_ucode_pins_claude_policy_model(self): ) assert result.exit_code == 0, result.output - mock_configure.assert_called_once_with( - "claude", MINIMAL_STATE, "system.ai.claude-opus-4-8" - ) + mock_configure.assert_called_once_with("claude", MINIMAL_STATE, "system.ai.claude-opus-4-8") mock_launch.assert_called_once() assert mock_launch.call_args.args[0] == "claude" assert mock_launch.call_args.args[2] == ["--model", "opus"] + def _enter_launch_patches(self, stack, tool, extra): + """Enter the shared launch no-op patches plus test-specific ones under + an ExitStack; return the launch_agent mock (the last shared patch).""" + mocks = [stack.enter_context(cm) for cm in _patch_launch(tool)] + for cm in extra: + stack.enter_context(cm) + return mocks[-1] + + def test_explicit_launch_blocked_and_redirected_by_policy(self): + """`ucode claude` under a policy that requires Codex must not launch + Claude — it blocks and boots the recommended tool/model instead.""" + with ExitStack() as stack: + mock_launch = self._enter_launch_patches( + stack, + "codex", + [ + patch("ucode.cli.apply_pat_environment"), + patch( + "ucode.cli.resolve_budget_policy_launch", + return_value=("codex", "gpt-5-4-mini"), + ), + ], + ) + result = runner.invoke( + app, + ["claude"], + env={"UCODE_BUDGET_POLICY": "/tmp/budget-policy.json"}, + ) + + assert result.exit_code == 0, result.output + mock_launch.assert_called_once() + assert mock_launch.call_args.args[0] == "codex" + flat = _strip_ansi(result.output) + assert "not allowed" in flat + assert "Codex" in flat + + def test_explicit_launch_allowed_when_policy_matches(self): + """When the policy already recommends the requested tool, launch it + unchanged with no redirect warning.""" + with ExitStack() as stack: + mock_launch = self._enter_launch_patches( + stack, + "claude", + [ + patch("ucode.cli.apply_pat_environment"), + patch( + "ucode.cli.resolve_budget_policy_launch", + return_value=("claude", None), + ), + ], + ) + result = runner.invoke( + app, + ["claude"], + env={"UCODE_BUDGET_POLICY": "/tmp/budget-policy.json"}, + ) + + assert result.exit_code == 0, result.output + mock_launch.assert_called_once() + assert mock_launch.call_args.args[0] == "claude" + assert "not allowed" not in _strip_ansi(result.output) + + def test_explicit_launch_not_blocked_when_policy_unresolved(self): + """A transient fetch failure (resolve returns None) must not lock the + user out — the requested tool still launches.""" + with ExitStack() as stack: + mock_launch = self._enter_launch_patches( + stack, + "claude", + [ + patch("ucode.cli.apply_pat_environment"), + patch("ucode.cli.resolve_budget_policy_launch", return_value=None), + ], + ) + result = runner.invoke( + app, + ["claude"], + env={"UCODE_BUDGET_POLICY": "/tmp/budget-policy.json"}, + ) + + assert result.exit_code == 0, result.output + mock_launch.assert_called_once() + assert mock_launch.call_args.args[0] == "claude" + + def test_explicit_launch_ignores_policy_when_none_configured(self): + """No policy (env unset and no default file) → resolve is never + consulted.""" + with ExitStack() as stack: + mock_launch = self._enter_launch_patches( + stack, "claude", [patch("ucode.cli.apply_pat_environment")] + ) + resolve_mock = stack.enter_context(patch("ucode.cli.resolve_budget_policy_launch")) + # Force "no policy configured" regardless of the host's env or an + # existing ~/.ucode/budget-policy.json file. + stack.enter_context(patch("ucode.cli.budget_policy_configured", return_value=False)) + result = runner.invoke(app, ["claude"]) + + assert result.exit_code == 0, result.output + resolve_mock.assert_not_called() + mock_launch.assert_called_once() + assert mock_launch.call_args.args[0] == "claude" + class TestMcpSubcommands: def test_web_search_subcommand_help(self):