Skip to main content
Last updated on

Tutorial: Govern a coding agent's model access with Unity AI Gateway

Beta

This feature is in Beta. Account admins can manage access to this feature from the account console Previews page. See Manage Databricks previews.

In this tutorial, you connect an external coding agent, such as Codex, Gemini CLI, or Cursor, to large language models (LLMs) governed by Unity AI Gateway. You then control who can access those models and how much the agent can consume and spend. Instead of each developer pasting a provider API key into their tool, every request routes through a Unity AI Gateway model service. Unity Catalog controls who can use which models, and Unity AI Gateway enforces limits and records usage and cost.

By the end of this tutorial, you have:

  • A coding agent that sends every model request through Unity AI Gateway using your Databricks credentials, with no provider API key on the developer's machine.
  • Central records of the agent's usage and cost, subject to the rate limits and permissions you set on the model service.

Prerequisites​

Step 1: Connect your coding agent​

You can connect the agent in one of two ways:

  • ucode (recommended): An open-source launcher (Unity AI Gateway Coding CLI) that handles OAuth and writes each agent's configuration file, including the model service base URL and model. You don't manage API keys or endpoints by hand.
  • Manual configuration: Set the model service base URL and model yourself in the tool's built-in settings (Cursor, Codex, or Gemini CLI). See Integrate with coding agents.

This section uses ucode. Install it, then run the agent you want:

Bash
uv tool install git+https://github.com/databricks/ucode
Bash
ucode codex

To run in non-interactive mode, pass the agent's own flag through ucode:

Bash
ucode codex --full-auto

On first run, ucode prompts for your workspace URL and authenticates; later runs go straight to the agent.

Step 2: Govern the agent's model access​

You can govern a model service like any other Unity Catalog securable: grant EXECUTE to control who can use it, and use Unity AI Gateway to cap request volume and spend. Set up any of the following:

  • Manage who can access the model: In Catalog Explorer, open the model service and grant EXECUTE to users or groups on the Permissions tab.
  • Cap request volume: In Catalog Explorer, on the model service's Overview tab, click Set up under Rate limits. Set a requests-per-minute (QPM) or tokens-per-minute (TPM) rate limit for all requests on the endpoint, or per user.
  • Set a spend budget: An account administrator can create a budget scoped to all Unity AI Gateway endpoints. The budget aggregates spend across your model services; set an alert threshold to be notified, or to block usage, before costs exceed it.

Step 3: Verify that traffic is governed​

Confirm that the agent's requests route through Unity AI Gateway and are recorded:

  • From the agent: Send a prompt. The agent responds using the model you configured, with no provider API key on your machine.

  • From the usage system table: Query usage tracking to confirm that your model service records the requests:

    SQL
    SELECT service_name, requester, status_code, COUNT(*) AS calls
    FROM system.ai_gateway.usage
    WHERE service_type = 'MODEL_SERVICE'
    GROUP BY service_name, requester, status_code
    ORDER BY calls DESC;

To analyze spend in more detail, see Monitor Unity AI Gateway cost.

Next steps​

On this page