Unity Gateway (ug)

Unity Gateway runs coding agents through Databricks AI Gateway. It configures Codex, Claude Code, Gemini CLI, OpenCode, GitHub Copilot CLI, and Pi, and can register Databricks MCP servers for Cursor Agent.

The command is ug. Existing ucode commands remain supported, and the Python package is still named ucode for compatibility.

Requirements

• Python 3.12+
• uv
• npm, when an agent CLI needs automatic installation

Install

uv tool install git+https://github.com/databricks/unity-gateway
ug --version
ug --help

--version reports installed package metadata without loading the CLI. Use ug --help (or ucode --help) to verify that the CLI’s startup imports work; version output alone is not an installation health check. Help does not validate workspace authentication or agent runtime behavior.

Migrating from ucode

Reinstall under the new distribution name:

uv tool uninstall ucode
uv tool install git+https://github.com/databricks/unity-gateway
ug --version
ucode --version
ug --help
ucode --help

Future upgrades can use ug upgrade or uv tool upgrade unity-gateway.

Launch Agents

Run the tool you want:

ug codex
ug claude
ug gemini
ug opencode
ug copilot
ug pi
ug cursor   # Cursor Agent, MCP-only

On first launch of a model-backed agent, ug prompts for a Databricks workspace, authenticates, and writes local agent config. Later launches reuse the saved workspace and credentials.

When Claude Code or Codex is older than the minimum version ug supports, ug asks before upgrading it in place. Set UG_AUTO_UPGRADE_AGENT=1 to upgrade without asking, for example from a launcher that manages the agent install.

Use ug opencode --model system.ai.glm-5-3 (or -m) to select a configured model for one launch. OpenCode’s provider/model form is also accepted. Unknown Databricks models produce an error; this option does not add models to discovery or change ug’s saved default.

ug copilot uses the Anthropic Messages API for Claude model IDs, the Responses API for GPT-5 and newer model IDs, and Chat Completions for other models. The model selected at launch (including --model) determines the API. An inherited COPILOT_PROVIDER_WIRE_MODEL takes precedence because it overrides the model sent to the gateway. Restart Copilot through ug to change the wire model or API; in-session model selection does not rebuild its provider configuration.

ug copilot runs Copilot against a loopback proxy for the session. Copilot holds only a random session secret; the proxy adds a fresh Databricks token to each gateway request, so OAuth sessions keep working past token expiry without a restart, and no token is written to ~/.copilot/ucode.env. If the Databricks sign-in itself expires, Copilot shows a “Databricks sign-in expired” error; run databricks auth login and restart ug copilot. Proxy errors are logged to ~/.ucode/logs/copilot-proxy.log.

Without a managed workspace config, ug claude automatically discovers gateway models for Claude Code’s /model picker. Discovery defaults to system.ai when no provider or model location is selected. Use --provider or --model-location to select another model source; managed workspace configs control their own sources. Interactive ug configure also offers a Custom Model Location (.) for Claude Code and Codex alongside Databricks Hosted and External Models. The choice is saved with the rest of the configuration, so later ug claude / ug codex launches use it; an explicit --provider or --model-location still overrides it for that launch.

ug codex validates discovered models with the installed Codex binary and publishes them to ~/.ucode/codex-model-catalog.json, referenced by shared ~/.codex/config.toml for Codex App. Managed static lists use the same path during ug configure. The latest refresh supplies the app’s catalog; custom catalogs (including Isaac’s) and custom providers are preserved. The app’s gateway provider and authentication must already be configured. Validation covers the local Codex binary.

Codex loads the catalog at app-server startup. When ug reports a catalog change, finish active tasks, restart the app server on the connected host, then reconnect. Use codex app-server daemon restart for a standalone managed daemon; otherwise restart the process or application that owns the server. Reconnecting or reopening the desktop app can reuse a remote server with the old list.

ug removes its shared reference on discovery/validation failure, reconfiguration, revert, or before installing/updating Codex. After an update, run ug codex to refresh discovery or ug configure for a managed static list, then restart the app server.

Configure

ug configure
ug configure --agents claude,codex
ug configure --workspace https://first.databricks.com
ug configure --profile DEFAULT --agents claude,codex
ug configure --user

--global (the default) writes the machine-wide Claude Code and Codex managed settings with sudo. --user skips them, like UG_SKIP_OS_SETTINGS, and needs no sudo; use it when you launch through ug. Configure fails while a managed file for the agent exists, since managed settings take precedence; later launches use whatever files are present. The choice is saved for launches after a successful configure; ug configure without --user and ug revert return to --global.

Available coding agents are codex, claude, gemini, opencode, copilot, and pi. cursor can be included in --agents for MCP-only setup; Cursor models still run through your Cursor account.

UG_WORKSPACE can provide the default workspace. An explicit --workspace or --profile takes precedence.

On macOS, Windows, and Linux, successful Claude Code configuration also creates and selects a dedicated Claude Desktop gateway profile. Its model list includes only Claude-named IDs; Claude Code’s catalog is unchanged. Desktop setup errors warn without failing Claude Code configuration. See Claude Desktop setup for authentication, profile preservation, and native validation limits. When managed Claude tracing is enabled, the same profile configures Desktop’s native OTLP trace export with a refreshing headers helper; no background service is needed.

MCP Servers

Register Databricks MCP servers for configured MCP-capable agents. Cursor Agent is MCP-only and is included when cursor-agent is installed:

Run ug mcp add without selectors to browse a table of services by name, UC location, and status. ✓ means configured (or managed by the workspace); ✗ means not configured. Configured and managed services stay visible as read-only rows. [ ]/[x] indicate what you select for this operation; [-] marks a read-only row.

Highlight a service to see its description, creator, and creation/update dates. Type to filter names or UC locations, use Space to select, and Enter to save. Page Up/Down scrolls the detail pane. Narrow terminals shorten cells and show status symbols; the detail pane retains the full location and status.

Use ug mcp add to add servers without removing existing registrations:

ug mcp add --location system.ai
ug mcp add --names system.ai.slack,system.ai.github
ug mcp add --agents claude,codex --location system.ai
ug mcp add --agents cursor --names system.ai.github --workspace https://first.databricks.com
ug mcp add --agents codex --names system.ai.github --profile DEFAULT

With --agents, agents that aren’t set up yet are set up first. Pass --workspace or --profile to pick the workspace instead of the picker; it becomes the current workspace.

Remove configured servers:

ug mcp remove
ug mcp remove --agents codex

List configured servers and their connection status:

ug mcp list
ug mcp list --agents claude,codex

Connection-backed AI Gateway MCP services use native HTTP OAuth in Claude Code, Codex, and Cursor when the workspace publishes a compatible client. Claude Code uses its URL-based client ID metadata document; if it is unavailable, ug uses the local ug mcp-proxy stdio server with the Databricks CLI’s OAuth client. Other servers and agents also use that proxy. V2 AI Gateway servers can be added with typed selectors such as vector-search:main.docs, uc-functions:main.tools, external:, genie-space:, or app:.

Claude-only setup also discovers GPT models for its generated web_search server; installing or configuring Codex is not required. Search registration requires an available Responses-capable model.

Claude’s generated web_search server uses the same saved custom OAuth CLI profile as its harness, when configured. It stores the profile name, not an access token, and refreshes credentials for search requests. This does not change search permissions.

The built-in search server runs up to four searches concurrently. A slow search does not block tool discovery or another search’s result. Additional searches wait for a worker. Cancelling a queued search prevents it from running; an active search retains its worker until its blocking request finishes, and its response is discarded. Closing the input stream drains accepted searches, so shutdown can wait for the existing authentication and HTTP timeouts. Input failure or interruption cancels queued searches and waits for active requests to finish. This removes local serialization without changing the backend model or speeding up an individual backend request.

Launchers that supply their own search server can first query ug mcp web-search --capabilities. Contract version 1 supports setting UCODE_CLAUDE_WEB_SEARCH_PROVIDER=external-if-safe on the ug child process only when the capability response advertises that automatic_provider value. The default (unset or ucode) retains standalone ug search. External mode does not create, update, or delete saved search registrations. Generated helpers carry --managed-by-ucode. Only the launch override created after ownership verification exposes no tools; copied helpers and custom registrations retain their tools even when they inherit external provider selection.

For an existing generated web_search entry, ug verifies its saved ownership fingerprint and uses a launch-only MCP override pointing at the current ug installation. This also handles older helper executables. Disabled entries are preserved without launch overrides, and --strict-mcp-config excludes saved registrations without an override. Unknown ownership, overlapping config scopes, or command-based MCP policies produce a warning and retain the existing provider; duplicate providers may remain. Explicit external selection instead fails on a conflict. Standalone refreshes preserve edited user entries and project/local servers. If no search model is available, they retain ownership of the installed entry so a later refresh or external-provider launch can still verify it. The selection does not authorize search or remove any permission denial. Concurrent setup still uses ug’s existing whole-workspace state writes. A stale ownership fingerprint safely preserves the server and prevents automatic handoff; this contract does not make shared setup transactional.

Skills

Unity Catalog Skills can be registered as MCP tools or downloaded into local agent skill directories. Run ug skills add without selectors to browse individual skills in the same table and metadata pane as the MCP browser, plus the bundle name. ✓ Installed requires a recorded download of that exact UC skill with SKILL.md still present at every target agent directory (--path, or your home directory). A matching bundle name from another schema or metastore does not qualify. Unattributed local bundles and incomplete installations show ✗ Not installed, with an explanation in the detail pane. Installed skills remain selectable for updates. Missing metadata is shown explicitly. Created by is the UC creator, which may differ from the author of the bundle.

# Set up the Databricks skills MCP so your agents can create and manage skills through ug.
ug skills

# List configured skills and how each was configured.
ug skills list

# Download every skill in a schema into your local agent skill directories.
ug skills add --location main.default

# Download specific skills by fully-qualified name (may span schemas).
ug skills add --names main.default.my-skill,ml.prod.other-skill

# Add a schema to the MCP connection scope, exposing its skills as MCP tools.
ug skills add --location main.default --via mcp

# Interactively pick downloaded skills to delete.
ug skills remove

# Delete a specific downloaded skill by fully-qualified name.
ug skills remove --names main.default.my-skill

# Drop a schema from the MCP connection scope.
ug skills remove --location main.default --via mcp

Commands

ug usage shows your Unity Gateway budget usage. To try the expanded budget view:

ENABLE_UG_USAGE_SPEND_API=1 ug usage

The expanded view shows up to five budget limits, prioritizing the most-used limits across personal/shared and daily/monthly budgets. It tells you when more limits exist. It uses your configured workspace login; no separate login is needed. Leave ENABLE_UG_USAGE_SPEND_API unset to keep the standard summary.

Databricks AI Tools are installed only by ug configure, never by agent launch commands. Use --enable-databricks-ai-tools or --disable-databricks-ai-tools with ug configure to control installation.

Claude Routing Plugin

Smart routing passes generated agents through a per-launch --plugin-dir, alongside --settings, without persistent plugin registration. One temporary directory holds the settings, socket, and plugin and is removed when the launch finishes or fails. Existing hook configuration and disable/revert behavior are unchanged. Native daemon/background propagation of the plugin remains unverified.

The routing notice names the generated agent selected for launch and its model. Claude Code’s own task header can still show the agent type the parent originally requested. That name may contain a different model; it is not the routing result.

The local ~/.ucode/-smart-routing-decisions.jsonl logs also retain the candidate models, their harnesses, and the router name sent with each recorded decision. Feedback investigations can compare the selection with that original menu without querying the router again.

Claude and Codex routing reasons summarize the task profile in plain language, such as “This task has a clear target and limited scope.” A separate sentence explains when model availability requires a higher tier. Custom gateway notices are retained, and older or unrecognized responses keep their original wording. Diagnostic records keep the displayed reason in rationale and the gateway’s full explanation in raw_rationale.

On Windows, Claude smart routing uses subagent hooks only. If first-prompt routing is enabled, ug warns and falls back to subagent routing because the first-prompt wrapper requires a Unix terminal. The generated shell hooks expect Git Bash; PowerShell-only setups are not covered.

Smart Router Orchestrator

Use SMART_ROUTER_CONFIG_VERSION at launch to select a smart-routing configuration. Each version defines separate Claude Code and Codex flag mappings; the existing versions use identical mappings for both agents. The Claude-only version disables routing for Codex:

first_prompt_and_subagent_no_orch_v0 is the customer configuration for first-prompt and subagent routing without orchestration: ENABLE_SMART_ROUTING_V2=1, ENABLE_SMART_ROUTING_SUBAGENT_ONLY=0, and ENABLE_SMART_ROUTER_ORCHESTRATOR=0.

subagent_only_v1 sets both ENABLE_SMART_ROUTING_V2 and ENABLE_SMART_ROUTING_SUBAGENT_ONLY to "1". Subagent-only takes precedence, so first-prompt routing remains off; orchestration is also off.

subagent_orch_v1 enables all three legacy flags. Like subagent_only_v1, it routes subagents rather than the first prompt, and it additionally enables orchestration.

subagent_orch_v0_claude_only keeps Claude’s subagent_orch_v0 flags (0, 1, 1) and sets all three flags to "0" for Codex. Codex uses its ordinary launch and managed default model, without routing or orchestration. Explicit launch/session on/off controls retain their existing precedence over presets.

SMART_ROUTER_NAME still selects the router independently of the preset.

Smart-routed Claude and Codex sessions install smart-router. The subagent_orch_v0 and subagent_orch_v1 versions also install and activate the bundled smart-router-orchestrator skill. For example:

SMART_ROUTER_CONFIG_VERSION=subagent_orch_v0 ug claude

The version takes precedence over conflicting legacy flags. Once the agent is selected, UG expands its mapping into ENABLE_SMART_ROUTING_V2, ENABLE_SMART_ROUTING_SUBAGENT_ONLY, and ENABLE_SMART_ROUTER_ORCHESTRATOR for the launched session. When the version is unset or empty, these legacy flags retain their existing behavior, including first-prompt routing through ENABLE_SMART_ROUTING_V2=1. Unknown versions are ignored: no preset is applied, the inherited environment is unchanged, and commands continue normally. Explicit launch/session on/off controls apply after expansion. Orchestration remains off by default. Workspace smart-routing defaults do not rewrite the selected version’s flags.

Version names require an explicit suffix. Future revisions use new _v1, _v2, etc. names without changing existing versions.

Version definitions require both Claude and Codex mappings and fail validation at module import if any flag in SMART_ROUTING_ENV_KEYS is missing, has a value other than "0" or "1", or an unknown flag is present. Register new managed flags in that tuple and explicitly set them for both agents in every version.

Use ug codex in the same command for Codex. Smart Router Orchestrator assigns bounded work to explorer, researcher, worker, tester, and reviewer roles while the root plans, integrates, and verifies results. Easy tasks and explicit requests not to delegate stay in the root.

Claude and Codex routing panels add [orchestrator on] to the Subagent line only when orchestration is active. When it is off, they show the normal subagent name. The label reports the session mode; it does not identify whether a particular delegation came from the workflow or an explicit user request.

Once opted in, orchestration follows the existing smart-routing launch eligibility and session controls. Turning Smart Router off through its skill stops new automatic delegation; turning it on restores orchestration only in opted-in sessions. Explicit user requests for subagents still use normal harness behavior while routing is off. Stored skill files do not activate orchestration without an opted-in configuration, or in non-routed sessions. Existing Isaac pilot gating and UG launch exclusions still apply.

Hooks refresh orchestration state before each prompt and after compaction. A state change made outside the conversation is observed at the next hook; model routing still checks the controls for each subagent.

UG supplies its own hooks; Codex combines them with existing hooks and applies project trust. Smart routing selects subagent models; separate role-model preferences are ignored and their files are left untouched. See the bundled Smart Router Orchestrator documentation for details.

Managed Files

ug backs up files before overwriting them. ug revert restores backups.

Codex desktop tracing

With workspace tracing enabled, run ug configure --agents codex on the machine where Codex desktop’s app server runs, then restart the app server once. This points Codex’s trace exporter in ~/.codex/config.toml at the existing CLI OTEL proxy, run as a per-user service that starts at login; ug codex sessions keep their own proxy. An exporter you already set is left alone. Disabling tracing or ug revert removes both.

Development

git clone https://github.com/databricks/unity-gateway
cd unity-gateway
uv sync
uv run pytest
uv run ruff check .

For integration tests against installed ug and agent versions, see tests/integration/README.md. The existing e2e tests remain available separately:

UCODE_TEST_WORKSPACE= uv run pytest tests/test_e2e.py -v

To add a new CLI agent, implement src/ucode/agents/.py, register it in src/ucode/agents/__init__.py, and add focused tests. Desktop-app adapters live separately in src/ucode/agents/desktop/. Desktop rendering consumes model IDs and headers from shared discovery/managed-config parsing without revalidating or deduplicating them. Claude Desktop profiles use Anthropic’s documented application-data locations on macOS, Windows, and Linux, with no profile-directory override. Windows and Linux require the app to have been opened once; their native GUI, credential-helper, and inference flows still need validation. Profiles are named Unity Gateway - and use one reserved UUID: 5b36a0b9-b56d-47cf-b6a4-0904e0bc14f6. Configure replaces that whole profile, updates its metadata entry, and selects it. Switching workspaces replaces the same file. Revert deletes that profile and entry, preserving unrelated profiles and metadata. Only the previous selection is recorded in ~/.ucode/claude-desktop.json; no original profile contents are saved for revert. Old random-ID and workspace-derived profiles are left untouched. A valid JSON _comment field warns against manual edits; native Desktop acceptance of that field still needs a smoke test. Workspace-specific UUID normalization is removed; existing CLI URL handling is unchanged.

Claude Desktop profile updates wait up to five seconds for the file lock across platforms. A timeout leaves the operation holding the lock undisturbed. The Desktop auth helper retains its existing shared deadline.

Documentation

• Databricks AI Gateway overview
• Databricks AI Gateway coding agent integration
• Databricks CLI authentication
• Monitor AI Gateway usage

Security

Please report security vulnerabilities to [email protected] rather than opening a public issue.

License

See LICENSE.md and NOTICE.md.