Add Goose agent support via Databricks AI Gateway

pyproject.toml

@@ -10,6 +10,7 @@ readme = "README.md"
 requires-python = ">=3.12"
 dependencies = [
   "databricks-sql-connector>=3.6.0",
+  "pyyaml>=6.0",
   "questionary>=2.0.0",
   "tomlkit>=0.13.0",
   "typer>=0.12.0",

src/ucode/agents/__init__.py

@@ -32,12 +32,13 @@
     spinner,
 )
 
-from . import claude, codex, copilot, gemini, opencode, pi
+from . import claude, codex, copilot, gemini, goose, opencode, pi
 
 _MODULES = {
     "codex": codex,
     "claude": claude,
     "gemini": gemini,
+    "goose": goose,
     "opencode": opencode,
     "copilot": copilot,
     "pi": pi,
@@ -51,6 +52,7 @@
     "claude-code": "claude",
     "gemini": "gemini",
     "gemini-cli": "gemini",
+    "goose": "goose",
     "opencode": "opencode",
     "copilot": "copilot",
     "pi": "pi",
@@ -64,7 +66,7 @@ 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, goose, opencode, copilot, pi."
         )
     return normalized
 
@@ -109,6 +111,16 @@ def install_tool_binary(tool: str, *, strict: bool = True, update_existing: bool
             _update_installed_tool_binary(tool)
         return True
 
+    if not package:
+        message = (
+            f"`{binary}` is not installed. "
+            f"Install {spec['display']} and ensure `{binary}` is on your PATH."
+        )
+        if strict:
+            raise RuntimeError(message)
+        print_warning(message)
+        return False
+
     if not shutil.which("npm"):
         message = f"`{binary}` is not installed and npm is not available to install it."
         if strict:
@@ -182,6 +194,8 @@ def configure_tool(tool: str, state: dict, model: str | None = None) -> dict:
             result = claude.write_tool_config(state, model)
         elif tool == "gemini":
             result = gemini.write_tool_config(state, model)
+        elif tool == "goose":
+            result = goose.write_tool_config(state, model)
         elif tool == "copilot":
             result = copilot.write_tool_config(state, model)
         elif tool == "pi":
@@ -210,6 +224,8 @@ def check_gateway_endpoint(state: dict, tool: str) -> bool:
         return bool(state.get("gemini_models"))
     if tool == "copilot":
         return bool(state.get("claude_models")) or bool(state.get("codex_models"))
+    if tool == "goose":
+        return bool(state.get("claude_models")) or bool(state.get("gemini_models"))
     if tool == "pi":
         return (
             bool(state.get("claude_models"))
@@ -221,6 +237,7 @@ def check_gateway_endpoint(state: dict, tool: str) -> bool:
 
 _TOOL_DISCOVERY_SOURCES: dict[str, tuple[str, ...]] = {
     "claude": ("claude",),
+    "goose": ("claude", "gemini"),
     "opencode": ("claude", "gemini"),
     "codex": ("codex",),
     "gemini": ("gemini",),

src/ucode/agents/goose.py

@@ -0,0 +1,239 @@
+"""Goose agent: merges Databricks settings into ~/.config/goose/config.yaml.
+
+Goose has a built-in Databricks provider that reads DATABRICKS_HOST from the
+config file and DATABRICKS_TOKEN from the environment (env var takes precedence
+over keyring). We merge only the three keys we own into the existing config so
+that user-defined extensions, preferences, and other settings are preserved.
+
+The token is injected as DATABRICKS_TOKEN at launch and refreshed every 30
+minutes so long-running sessions stay authenticated.
+
+Install goose from https://github.com/aaif-goose/goose — it ships as a native
+binary (not an npm package), typically installed to ~/.local/bin via:
+  curl -fsSL https://github.com/aaif-goose/goose/releases/download/stable/download_cli.sh | bash
+"""
+
+from __future__ import annotations
+
+import os
+import signal
+import subprocess
+import threading
+from pathlib import Path
+
+from ucode.config_io import (
+    APP_DIR,
+    ToolSpec,
+    backup_existing_file,
+    deep_merge_dict,
+    read_yaml_safe,
+    write_yaml_file,
+)
+from ucode.databricks import (
+    TOKEN_REFRESH_INTERVAL_SECONDS,
+    get_databricks_token,
+)
+from ucode.state import mark_tool_managed, save_state
+
+GOOSE_CONFIG_DIR = Path.home() / ".config" / "goose"
+GOOSE_CONFIG_PATH = GOOSE_CONFIG_DIR / "config.yaml"
+GOOSE_BACKUP_PATH = APP_DIR / "goose-config.backup.yaml"
+
+SPEC: ToolSpec = {
+    "binary": "goose",
+    "package": "",  # not an npm package; install from https://github.com/aaif-goose/goose
+    "display": "Goose",
+    "config_path": GOOSE_CONFIG_PATH,
+    "backup_path": GOOSE_BACKUP_PATH,
+}
+
+MANAGED_KEYS: list[str] = [
+    "DATABRICKS_HOST",
+    "GOOSE_PROVIDER",
+    "GOOSE_MODEL",
+    "OAUTH_TOKEN",
+]
+
+GOOSE_MCP_AUTH_ENV_KEY = "OAUTH_TOKEN"
+
+
+def is_update_available() -> tuple[str, str] | None:
+    return None  # no npm update check for native binary
+
+
+def default_model(state: dict) -> str | None:
+    """Prefer Claude sonnet, then opus, then haiku; fall back to gemini."""
+    claude_models = state.get("claude_models") or {}
+    for family in ("sonnet", "opus", "haiku"):
+        if claude_models.get(family):
+            return claude_models[family]
+    gemini_models = state.get("gemini_models") or []
+    if gemini_models:
+        return gemini_models[0]
+    return None
+
+
+def render_overlay(workspace: str, model: str) -> dict:
+    """Return only the keys ucode manages — merged into the existing config."""
+    return {
+        "DATABRICKS_HOST": workspace,
+        "GOOSE_PROVIDER": "databricks",
+        "GOOSE_MODEL": model,
+        "extensions": {
+            "skills": {
+                "enabled": True,
+                "type": "platform",
+                "name": "skills",
+                "description": "Load and use skills from .claude/skills or .goose/skills directories",
+                "bundled": True,
+                "available_tools": [],
+            }
+        },
+    }
+
+
+def build_runtime_env(workspace: str, token: str) -> dict[str, str]:
+    env = os.environ.copy()
+    env["DATABRICKS_HOST"] = workspace
+    env["DATABRICKS_TOKEN"] = token
+    env["OAUTH_TOKEN"] = token
+    return env
+
+
+def _mcp_slug(name: str) -> str:
+    return name.lower().replace("-", "_")
+
+
+def build_mcp_server_entry(name: str, url: str, token: str = "") -> dict:
+    return {
+        "enabled": True,
+        "type": "streamable_http",
+        "name": name,
+        "description": f"Databricks MCP server: {name}",
+        "uri": url,
+        "envs": {GOOSE_MCP_AUTH_ENV_KEY: token},
+        "env_keys": [],
+        "headers": {"Authorization": f"Bearer ${{{GOOSE_MCP_AUTH_ENV_KEY}}}"},
+        "timeout": 300,
+        "bundled": None,
+        "available_tools": [],
+    }
+
+
+def write_mcp_server_config(name: str, url: str, token: str = "") -> bool:
+    backup_existing_file(GOOSE_CONFIG_PATH, GOOSE_BACKUP_PATH)
+    existing = read_yaml_safe(GOOSE_CONFIG_PATH)
+    extensions = existing.get("extensions")
+    if not isinstance(extensions, dict):
+        extensions = {}
+    slug = _mcp_slug(name)
+    removed = slug in extensions
+    extensions[slug] = build_mcp_server_entry(name, url, token)
+    existing["extensions"] = extensions
+    write_yaml_file(GOOSE_CONFIG_PATH, existing)
+    return removed
+
+
+def remove_mcp_server_config(name: str) -> bool:
+    existing = read_yaml_safe(GOOSE_CONFIG_PATH)
+    extensions = existing.get("extensions")
+    if not isinstance(extensions, dict):
+        return False
+    slug = _mcp_slug(name)
+    if slug not in extensions:
+        return False
+    extensions.pop(slug)
+    existing["extensions"] = extensions
+    write_yaml_file(GOOSE_CONFIG_PATH, existing)
+    return True
+
+
+def write_tool_config(
+    state: dict,
+    model: str,
+    token: str | None = None,
+    *,
+    force_refresh: bool = False,
+) -> tuple[dict, str]:
+    backup_existing_file(GOOSE_CONFIG_PATH, GOOSE_BACKUP_PATH)
+    if token is None:
+        token = get_databricks_token(state["workspace"], force_refresh=force_refresh)
+    overlay = render_overlay(state["workspace"], model)
+    existing = read_yaml_safe(GOOSE_CONFIG_PATH)
+    deep_merge_dict(existing, overlay)
+    extensions = existing.get("extensions")
+    if isinstance(extensions, dict):
+        for ext in extensions.values():
+            if isinstance(ext, dict) and ext.get("type") == "streamable_http":
+                envs = ext.get("envs")
+                if isinstance(envs, dict):
+                    envs[GOOSE_MCP_AUTH_ENV_KEY] = token
+                else:
+                    ext["envs"] = {GOOSE_MCP_AUTH_ENV_KEY: token}
+    write_yaml_file(GOOSE_CONFIG_PATH, existing)
+    state = mark_tool_managed(state, "goose", MANAGED_KEYS)
+    save_state(state)
+    return state, token
+
+
+def _refresh_token_once(state: dict, *, force_refresh: bool = False) -> tuple[str, str]:
+    model = default_model(state)
+    if not model:
+        raise RuntimeError("No Goose model is available on this workspace.")
+    _, token = write_tool_config(state, model, force_refresh=force_refresh)
+    return model, token
+
+
+def _refresh_forever(state: dict, stop_event: threading.Event) -> None:
+    while not stop_event.wait(TOKEN_REFRESH_INTERVAL_SECONDS):
+        try:
+            _refresh_token_once(state, force_refresh=True)
+        except RuntimeError:
+            continue
+
+
+def launch(state: dict, tool_args: list[str]) -> None:
+    model, token = _refresh_token_once(state)
+    env = build_runtime_env(state["workspace"], token)
+
+    stop_event = threading.Event()
+    refresher = threading.Thread(
+        target=_refresh_forever,
+        args=(state, stop_event),
+        daemon=True,
+    )
+    refresher.start()
+
+    proc = subprocess.Popen(["goose", "session", *tool_args], env=env)
+    try:
+        returncode = proc.wait()
+    except KeyboardInterrupt:
+        proc.send_signal(signal.SIGINT)
+        returncode = proc.wait()
+    finally:
+        stop_event.set()
+        refresher.join(timeout=1)
+
+    raise SystemExit(returncode)
+
+
+def validate_cmd(binary: str) -> list[str]:
+    return [
+        binary,
+        "run",
+        "--text",
+        "say hi in 5 words or less",
+        "--no-session",
+        "--max-turns",
+        "1",
+    ]
+
+
+def validate_env(state: dict) -> dict[str, str]:
+    workspace = state.get("workspace")
+    if not workspace:
+        raise RuntimeError("No workspace configured.")
+    if not default_model(state):
+        raise RuntimeError("No Goose model is available on this workspace.")
+    token = get_databricks_token(workspace)
+    return build_runtime_env(workspace, token)

src/ucode/cli.py

@@ -58,9 +58,9 @@
 from ucode.usage import usage as usage_report
 
 _DISCOVERY_CONSUMERS: dict[str, tuple[str, ...]] = {
-    "claude": ("claude", "opencode", "copilot", "pi"),
+    "claude": ("claude", "goose", "opencode", "copilot", "pi"),
     "codex": ("codex", "copilot", "pi"),
-    "gemini": ("gemini", "opencode", "pi"),
+    "gemini": ("gemini", "goose", "opencode", "pi"),
 }
 
 
@@ -110,9 +110,16 @@ def configure_shared_state(
     print_success("Unity AI Gateway detected")
 
     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 "goose" in tools
+        or "opencode" in tools
+        or "copilot" in tools
+        or "pi" in tools
+    )
+    want_gemini = (
+        fetch_all or "gemini" in tools or "goose" in tools or "opencode" 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
 
     claude_reason: str | None = None
@@ -390,7 +397,7 @@ def _launch_tool(tool_name: str, ctx: typer.Context) -> None:
         print_section(f"ucode with {TOOL_SPECS[tool]['display']}")
         if resolved_model:
             print_kv("Model", resolved_model)
-        if tool in ("gemini", "opencode", "copilot", "pi"):
+        if tool in ("gemini", "goose", "opencode", "copilot", "pi"):
             print_note(
                 f"{TOOL_SPECS[tool]['display']} token refresh is managed automatically "
                 f"every 30 minutes while the session is running."
@@ -431,6 +438,12 @@ def opencode_cmd(ctx: typer.Context) -> None:
     _launch_tool("opencode", ctx)
 
 
+@app.command("goose", context_settings={"allow_extra_args": True, "ignore_unknown_options": True})
+def goose_cmd(ctx: typer.Context) -> None:
+    """Launch Goose via Databricks."""
+    _launch_tool("goose", ctx)
+
+
 @app.command("copilot", context_settings={"allow_extra_args": True, "ignore_unknown_options": True})
 def copilot_cmd(ctx: typer.Context) -> None:
     """Launch GitHub Copilot CLI via Databricks."""
@@ -453,7 +466,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, goose, opencode, copilot, pi).",
         ),
     ] = None,
 ) -> None: