Skip to main content

Deploy an agent using on-behalf-of-user authentication

Beta

This feature is in Beta.

warning

While on-behalf-of user authentication is a powerful tool for enforcing secure access to sensitive data. It enables workspace users to author agents that act on behalf of other users in Databricks. During beta, it is disabled by default and must be enabled by a workspace admin. Review the security considerations for on-behalf-of-user authentication before enabling this feature.

With on-behalf-of-user authentication, agents deployed via Mosaic AI model serving can access Databricks resources using the identity of the Databricks end user who queried the agent. This enables accessing sensitive information on a per-user basis, with fine-grained enforcement of data access control in Unity Catalog.

On-behalf-of-user authentication further restricts the incoming user token through downscoping, ensuring that the token exposed to the agent code is limited to accessing only the specific APIs defined by the agent author. This improves security by preventing unauthorized actions and reducing the risk of token misuse.

When authoring your agent, you can continue to use existing SDKs to access Databricks resources, like vector search indexes. To enable on-behalf-of-user access to resources:

  1. In agent code, update SDK calls to indicate that resources should be accessed on behalf of the agent end user
  2. At agent logging time (prior to agent deployment), specify the end user REST API scopes required by your agent. For more details see On-behalf-of-user authentication

For an end-to-end example of on-behalf-of-user authentication, see End-to-end example.

The following resources are compatible with on-behalf-of-user authentication with agents.

Databricks Resource

Compatible Clients

Vector Search Index

databricks_langchain.VectorSearchRetrieverTool, databricks_openai.VectorSearchRetrieverTool or VectorSearchClient

Model Serving Endpoint

databricks.sdk.WorkspaceClient

SQL Warehouse

databricks.sdk.WorkspaceClient

UC Connections

databricks.sdk.WorkspaceClient

UC Tables and UC Functions

Databricks does not currently support direct clients to access UC Tables or UC Functions with on-behalf-of-user authentication. Instead we encourage users to use Genie to access structured data with on-behalf-of-user authentication

Genie Space

databricks_langchain.GenieAgent or databricks_openai.GenieAgent

When initializing tools with the on-behalf-of-user client, you can either wrap the tool initialization in a try-except block or allow errors to be exposed to the users. By handling errors, the agent can still make a best-effort response even if the end user lacks access to all the tools needed. However, if you choose not to handle errors during tool initialization, the agent will throw an error if the user is missing any required resources.

Configure SDKs​

The snippets below demonstrate how to configure on-behalf-of-user access to different Databricks resources using various SDKs

Python
from databricks.sdk import WorkspaceClient
from databricks.sdk.credentials_provider import ModelServingUserCredentials
from databricks_langchain import VectorSearchRetrieverTool

# Configure a Databricks SDK WorkspaceClient to use on behalf of end
# user authentication
user_client = WorkspaceClient(credentials_strategy = ModelServingUserCredentials())

vector_search_tools = []
# Exclude exception handling if you want the agent to fail
# when users lack access to all required Databricks resources
try:
  tool = VectorSearchRetrieverTool(
    index_name="<index_name>",
    description="...",
    tool_name="...",
    workspace_client=user_client # Specify the user authenticated client
    )
    vector_search_tools.append(tool)
except Exception as e:
    _logger.debug("Skipping adding tool as user does not have permissions)

Initializing the Agent​

on-behalf-of-user authentication is compatible with the ChatAgent interface. When using on-behalf-of-user authentication, the end user's identity is only known when your deployed agent is queried, i.e. within the predict and predict_stream functions of the ChatAgent interface. As a result, you must perform any on behalf of user access to resources (e.g. list vector search indexes that the end user has access to) from within these methods, rather than in the __init__ method of your ChatAgent implementation. This ensures that resources are isolated between invocations

Python
from mlflow.pyfunc import ChatAgent


class LangGraphChatAgent(ChatAgent):
  def initialize_agent():
    user_client = WorkspaceClient(
      credentials_strategy=ModelServingUserCredentials()
    )
    system_authorized_client = WorkspaceClient()
    ### Use the clients above to access resources with either system or user authorization

  def predict(
    self,
    messages: list[ChatAgentMessage],
    context: Optional[ChatContext] = None,
    custom_inputs: Optional[dict[str, Any]] = None,
  ) -> ChatAgentResponse:
    agent = initialize_agent() # Initialize the Agent in Predict
    request = {"messages": self._convert_messages_to_dict(messages)}

    messages = []
    for event in self.agent.stream(request, stream_mode="updates"):
      for node_data in event.values():
        messages.extend(
          ChatAgentMessage(**msg) for msg in node_data.get("messages", [])
        )
    return ChatAgentResponse(messages=messages)

Security Considerations​

There are some security implications to consider before enabling on-behalf-of-user authentication with agents:

  1. Access to Sensitive Databricks Resources: Enabling on-behalf-of-user authentication allows agents to access sensitive Databricks resources. While we've implemented API scopes to restrict the resources that developers can access and mitigate the risk of token misuse, some risks still remain. For example, the serving.serving-endpoints API scope grants an agent permission to execute a serving endpoint on behalf of the user. However, the serving endpoint itself may have access to additional API scopes that the original agent isn't authorized to use.
  2. No support for end user consent: During the current beta phase, agent users cannot view or consent to the Databricks REST API scopes required by an agent. Users are responsible for ensuring that they trust those with "Can Manage" permissions on the serving endpoint to take actions in Databricks on their behalf.

End-to-end example​

The following notebook shows you how to create an agent with vector search using on-behalf-of-user authentication.

On Behalf of User Authentication with Vector Search

Open notebook in new tab
Last updated on