Supervisor API

Use the Databricks Supervisor API

The Supervisor API lets Databricks run the tool-selection and synthesis loop server-side. Instead of your agent managing tool calls and looping, you declare hosted tools and call responses.create() — Databricks handles the rest. You can also declare client-side function tools (Python callables your app executes) alongside hosted tools in the same request.

When to Use

Use the Supervisor API when you want to:

• Connect Genie spaces, UC functions, Knowledge Assistants, or MCP servers without managing the agent loop yourself
• Mix client-side function tools (your own Python callables) with hosted tools in the same request
• Choose models at runtime and control which tools are used per request
• Offload tool orchestration to Databricks while iterating on your agent

Limitations:

• Inference parameters (e.g., temperature, top_p) are not supported when tools are passed
• stream and background cannot both be true in the same request
• Background mode requests have a maximum execution time of 30 minutes

Step 1: Install databricks-openai

Add to pyproject.toml if not already present:

[project]
dependencies = [
    ...
    "databricks-openai>=0.14.0",
    "databricks-sdk>=0.55.0",
]

Then run uv sync.

Step 2: Declare Tools

Define your tools as a list of dicts. Run uv run discover-tools to find available hosted resources in your workspace.

TOOLS = [
    # Genie space — natural language queries over structured data
    {
        "type": "genie_space",
        "genie_space": {
            "id": "<genie-space-id>",
            "description": "Query sales data using natural language",
        },
    },
    # UC function — SQL or Python UDF
    {
        "type": "uc_function",
        "uc_function": {
            "name": "<catalog>.<schema>.<function_name>",
            "description": "Executes a custom UC function",
        },
    },
    # Knowledge Assistant agent endpoint
    {
        "type": "knowledge_assistant",
        "knowledge_assistant": {
            "knowledge_assistant_id": "<ka-id>",
            "description": "Answers questions from internal documentation",
        },
    },
    # External MCP server via Unity Catalog connection
    {
        "type": "uc_connection",
        "uc_connection": {
            "name": "<uc-connection-name>",
            "description": "Searches the web for current information",
        },
    },
    # Databricks App endpoint or custom MCP server running as a Databricks App
    {
        "type": "app",
        "app": {
            "name": "<databricks-app-name>",
            "description": "Custom application or MCP server endpoint",
        },
    },
    # Client-side function tool example — your app executes the call
    {
        "type": "function",
        "name": "<client-side-function>",
        "description": "<description of what this function does>",
        "parameters": {
            "type": "object",
            "properties": {"<param>": {"type": "string"}},
            "required": ["<param>"],
            "additionalProperties": False,
        },
    },
]

Step 3: Update agent_server/agent.py

Replace your existing invoke/stream handlers with the Supervisor API pattern. Remove any MCP client setup, LangGraph agents, or OpenAI Agents SDK runner code — the Supervisor API replaces the client-side loop entirely.

use_ai_gateway=True automatically resolves the correct AI Gateway endpoint for the workspace.

Authentication options:

• OBO (recommended): Pass workspace_client=get_user_workspace_client() per-request so tools run as the requesting user. The app must request the appropriate OAuth scopes in app.yaml. Grant tool resource permissions to users, not the service principal.
• Service principal: Pass workspace_client=WorkspaceClient() (or omit it) to run tools as the app's service principal. Grant each hosted tool's resource permissions in databricks.yml (Step 4).

Client-side function tools execute in your application code, so they need no Databricks-side permissions; instead, you implement the function and run it on each turn (see "Client-Side Function Tools" below).

import os
import logging
import mlflow
from databricks_openai import DatabricksOpenAI
from mlflow import MlflowClient
from mlflow.genai.agent_server import invoke, stream
from mlflow.tracing import get_tracing_context_headers_for_http_request
from mlflow.types.responses import (
    ResponsesAgentRequest,
    ResponsesAgentResponse,
)
from agent_server.utils import get_session_id, get_user_workspace_client

mlflow.openai.autolog()

logger = logging.getLogger(__name__)

MODEL = "databricks-claude-sonnet-4-5"
TOOLS = [...]  # From Step 2


def _get_trace_destination() -> dict | None:
    experiment_id = os.environ.get("MLFLOW_EXPERIMENT_ID")
    if not experiment_id:
        return None
    try:
        trace_location = MlflowClient().get_experiment(experiment_id).trace_location
        if trace_location is None or not hasattr(trace_location, "catalog_name"):
            return None
        dest = {
            "catalog_name": trace_location.catalog_name,
            "schema_name": trace_location.schema_name,
        }
        if trace_location.table_prefix is not None:
            dest["table_prefix"] = trace_location.table_prefix
        return dest
    except Exception:
        logger.warning("Could not resolve trace destination, tracing disabled.", exc_info=True)
        return None


_TRACE_DESTINATION = _get_trace_destination()


def _extra_body() -> dict:
    if _TRACE_DESTINATION:
        return {"trace_destination": _TRACE_DESTINATION}
    return {}


@invoke()
def invoke_handler(request: ResponsesAgentRequest) -> ResponsesAgentResponse:
    if session_id := get_session_id(request):
        mlflow.update_current_trace(metadata={"mlflow.trace.session": session_id})
    client = DatabricksOpenAI(workspace_client=get_user_workspace_client(), use_ai_gateway=True)
    response = client.responses.create(
        model=MODEL,
        input=[i.model_dump() for i in request.input],
        tools=TOOLS,
        stream=False,
        extra_headers=get_tracing_context_headers_for_http_request(),
        extra_body=_extra_body(),
    )
    return ResponsesAgentResponse(output=[item.model_dump() for item in response.output])


@stream()
def stream_handler(request: ResponsesAgentRequest):
    if session_id := get_session_id(request):
        mlflow.update_current_trace(metadata={"mlflow.trace.session": session_id})
    client = DatabricksOpenAI(workspace_client=get_user_workspace_client(), use_ai_gateway=True)
    return client.responses.create(
        model=MODEL,
        input=[i.model_dump() for i in request.input],
        tools=TOOLS,
        stream=True,
        extra_headers=get_tracing_context_headers_for_http_request(),
        extra_body=_extra_body(),
    )

Step 4: Grant Permissions

OBO mode

Add the required OAuth scopes to app.yaml. ai-gateway is required for all Supervisor API access; add the per-tool scopes for each tool type you use:

oauth_scopes:
  - "ai-gateway"          # required for all Supervisor API access
  - "genie"               # genie_space tools
  - "mcp.functions"       # uc_function tools
  - "model-serving"       # knowledge_assistant tools
  - "catalog.connections" # uc_connection tools

Grant tool resource permissions to the users who will run the agent (e.g., CAN_RUN on the Genie space, CAN_QUERY on the model endpoint). No databricks.yml resource grants are needed for the agent itself.

Service principal mode

For each hosted tool, grant the corresponding resource access. Client-side function tools need no Databricks permissions. See the add-tools skill for complete YAML examples.

Step 5: Test and Deploy

uv run start-app       # Test locally
databricks bundle deploy && databricks bundle run {{BUNDLE_NAME}}  # Deploy

Supported Models

Pass any of these as the model parameter:

Enabling Tracing

The Supervisor API supports distributed tracing — spans from the server-side agent loop are linked into the same trace as your agent's client-side spans, giving end-to-end visibility in MLflow. See the MLflow distributed tracing docs for more details.

How It Works

1. The trace destination (UC catalog/schema/table) is resolved from your MLflow experiment at startup
2. get_tracing_context_headers_for_http_request() propagates the active span context to the Supervisor API via HTTP headers
3. trace_destination tells the Supervisor API where to write its server-side spans
4. Both sets of spans appear as a single connected trace in MLflow

Setup

Add to agent_server/agent.py:

import os
import logging
import mlflow
from mlflow import MlflowClient
from mlflow.tracing import get_tracing_context_headers_for_http_request

logger = logging.getLogger(__name__)


def _get_trace_destination() -> dict:
    """Resolve trace destination from the experiment's Unity Catalog trace location."""
    experiment_id = os.environ.get("MLFLOW_EXPERIMENT_ID")
    if not experiment_id:
        raise RuntimeError(
            "MLFLOW_EXPERIMENT_ID is not set. Cannot configure distributed tracing. "
            "Ensure the app is configured with a valid MLflow experiment."
        )

    client = MlflowClient()
    experiment = client.get_experiment(experiment_id)
    trace_location = experiment.trace_location

    if trace_location is None or not hasattr(trace_location, "catalog_name"):
        msg = (
            f"Experiment {experiment_id} trace_location is not a Unity Catalog location "
            f"(got: {type(trace_location).__name__ if trace_location else None}). "
            "Distributed tracing requires UC-backed traces. "
            "Ensure 'MLflow traces in Unity Catalog' is enabled for your workspace and that "
            "the target UC tables use customer-managed storage (Arclight default storage is not supported)."
        )
        logger.error(msg)
        raise RuntimeError(msg)

    dest = {
        "catalog_name": trace_location.catalog_name,
        "schema_name": trace_location.schema_name,
    }
    if trace_location.table_prefix is not None:
        dest["table_prefix"] = trace_location.table_prefix
    return dest


# Resolve once at module load — fail fast if the experiment is misconfigured
_TRACE_DESTINATION = _get_trace_destination()

Then pass extra_headers and extra_body in every responses.create() call:

response = _client.responses.create(
    model=MODEL,
    input=[i.model_dump() for i in request.input],
    tools=TOOLS,
    stream=False,
    extra_headers=get_tracing_context_headers_for_http_request(),
    extra_body={
        "trace_destination": _TRACE_DESTINATION,
    },
)

• extra_headers — propagates the active MLflow span context so client and server spans are linked into one trace
• trace_destination — tells the Supervisor API where to write server-side spans in Unity Catalog

Environment Variable

MLFLOW_EXPERIMENT_ID must be set in your app environment. The quickstart script sets this automatically. To verify:

grep MLFLOW_EXPERIMENT_ID .env

Claude: Before writing any tracing code, check .env for MLFLOW_EXPERIMENT_ID. If it is missing or empty, ask the user: "Distributed tracing requires MLFLOW_EXPERIMENT_ID to be set in .env. Do you have an MLflow experiment ID? If not, run uv run quickstart to create one, or provide the experiment ID and I'll add it to .env."

If the value is present but looks wrong (e.g. not a numeric string), warn the user before proceeding.

MCP Server Tools: Multi-Turn Approval Flow

When using MCP server tools (uc_connection or app), the Supervisor API does not execute the MCP tool call in a single request. Instead, it returns a completed response containing mcp_approval_request output items. To complete the tool call, your agent must handle a multi-turn flow:

1. First request — responses.create() → response completes with mcp_approval_request items in the output
2. Return to frontend — the mcp_approval_request item is returned to the chat UI so the user can approve the tool call
3. Second request — user approves → frontend sends a new request with the original input + mcp_approval_request + mcp_approval_response (with approve: true) appended to the input
4. Result — the second response completes with the actual function_call_output (tool result) and the final assistant message

No special backend handling is needed — the agent server simply returns all output items (including mcp_approval_request) to the frontend. The multi-turn flow is handled naturally through the conversation: each request/response is a separate responses.create() call.

Example input for the follow-up request (step 3):

input = [
    # Original user message
    {"type": "message", "role": "user", "content": "Search for Databricks"},
    # The mcp_approval_request from the first response's output
    {"type": "mcp_approval_request", "id": "call_xxx", "name": "web-search",
     "server_label": "you_dot_com", "arguments": '{"query": "Databricks"}'},
    # The approval
    {"type": "mcp_approval_response", "id": "call_xxx",
     "approval_request_id": "call_xxx", "approve": True},
]

Client-Side Function Tools

You can also declare client-side function tools (type: "function") alongside hosted tools in the same request.

See the supervisor-api-client-function-calling skill for full implementation details, including function-only, streaming, mixed hosted + function, and MCP approval + function flow examples.

Troubleshooting

"Please ensure AI Gateway V2 is enabled" — AI Gateway must be enabled for the workspace. Contact your Databricks account team.

"Parameter not supported when tools are provided" — Remove temperature, top_p, or other inference parameters from the responses.create() call.

Background Mode (Long-Running Tasks)

If your agent needs to run long-running tasks that may exceed HTTP timeout limits (e.g., complex multi-tool workflows, large data analysis), you can enable background mode. This submits the request asynchronously, polls for completion, and streams the result back to the frontend.

See the supervisor-api-background-mode skill for full implementation details.