OBO Authentication with Unity AI Gateway and databricks.agents.deploy()

Froffri
New Contributor II

Hi everyone,

I'm trying to use the new Unity AI Gateway v2 (beta) endpoint from an agent deployed with the Mosaic AI Agent Framework and I'm running into what looks like a scope mismatch.

Library versions
databricks-agents==1.11.0
mlflow==3.12.0
python~3.12
Setup

I deploy my agent using:

databricks.agents.deploy(...)
and I configure OBO authentication through:
AuthPolicy(
    system_auth_policy=SystemAuthPolicy(...),
    user_auth_policy=UserAuthPolicy(
        api_scopes=[
            "sql",
            "unity-catalog",
            "vector-search",
            "model-serving",
        ]
    ),
)
The agent is able to successfully use OBO authentication for other Databricks services.
Goal
I want the agent to call a Unity AI Gateway endpoint (the new AI Gateway visible in the left navigation menu), using the end user's identity.
Problem
When the agent calls the AI Gateway endpoint, the request reaches the service but fails with an error similar to:
missing scope: ai-gateway

This suggests that the OBO token provided to the agent does not contain an ai-gateway scope.

However, when I try to add:

UserAuthPolicy(
    api_scopes=[
        "ai-gateway"
    ]
)
the deployment fails because ai-gateway is not recognized as a valid API scope.
Questions
  1. Is OBO authentication currently supported for the new Unity AI Gateway endpoints when deploying with databricks.agents.deploy()?
  2. If yes, what value should be added to UserAuthPolicy.api_scopes?
  3. Is there currently a supported MLflow resource type or AuthPolicy configuration for AI Gateway endpoints?
  4. If this is not yet supported, what is the recommended approach?

    Has anyone managed to successfully invoke a Unity AI Gateway endpoint from an OBO-authenticated Mosaic AI Agent?

    Thanks!