Agent authentication¶

AI agents often need to authenticate to external resources to complete tasks. For example, a deployed agent might need to access external APIs, databases, or cloud services to retrieve data or perform operations.

This documentation provides comprehensive guidance for implementing authentication in DataRobot Agent Templates, covering API tokens, MCP server authentication, authorization context, OAuth 2.0, and security best practices.

Authentication methods¶

This section provides an overview of the different authentication methods available in the framework:

API token authentication¶

API token authentication is the most common method for authenticating with DataRobot services and external APIs. It uses bearer tokens passed in headers or environment variables.

DataRobot API authentication¶

Configure API tokens using environment variables or programmatically. The DATAROBOT_API_TOKEN is available in the API keys and tools section of your account settings.

DATAROBOT_API_TOKEN=<your_api_key>
DATAROBOT_ENDPOINT=https://app.datarobot.com

from agent import MyAgent

agent = MyAgent(
    api_key="your_api_key",
    api_base="https://app.datarobot.com"
)

External API authentication¶

When creating custom tools that need to authenticate with external APIs, retrieve API keys from environment variables or runtime parameters within your tool's run() method. This keeps credentials out of your source code and follows the same security pattern used throughout the agent templates.

For local development, use environment variables. For deployed agents, define runtime parameters in model-metadata.yaml and retrieve them using RuntimeParameters.get() from datarobot_drum. Try os.environ.get() first (for local development), then fall back to RuntimeParameters.get() for deployed agents., as shown in the example below:

import os
import requests
from crewai.tools import BaseTool
from datarobot_drum import RuntimeParameters

class ExternalAPITool(BaseTool):
    def run(self, query: str) -> str:
        api_key = os.environ.get("EXTERNAL_API_KEY")
        if not api_key:
            api_key = RuntimeParameters.get("EXTERNAL_API_KEY")["apiToken"]

        headers = {"Authorization": f"Bearer {api_key}"}
        response = requests.get(
            "https://api.external-service.com/data",
            headers=headers,
            params={"query": query}
        )
        return response.json()

runtimeParameterDefinitions:
  - fieldName: EXTERNAL_API_KEY
    type: credential
    credentialType: api_token
    description: API key for external service authentication

MCP server authentication¶

When using MCP (Model Context Protocol) servers to provide tools to your agents, authentication is handled automatically using the DATAROBOT_API_TOKEN environment variable. The MCP client in the agent templates automatically uses this token for authenticated requests to the MCP server.

Configuration:

DATAROBOT_API_TOKEN=<your_api_key>
DATAROBOT_ENDPOINT=https://app.datarobot.com

The MCP client will automatically use DATAROBOT_API_TOKEN when making requests to the MCP server. No additional configuration is required beyond setting these environment variables.

Authorization context¶

The framework provides an authorization context system for propagating authentication information to downstream tools and services. This allows tokens and credentials to be automatically passed between agent tools without manual configuration.

Initialize authorization context in your agent's chat() function using resolve_authorization_context() from the datarobot-genai package. The function returns the authorization context dictionary, which should be assigned to completion_create_params["authorization_context"]. Tools can then retrieve the context using get_authorization_context() from the datarobot SDK:

from datarobot_genai.core.chat import resolve_authorization_context
from datarobot.models.genai.agent.auth import get_authorization_context

def chat(completion_create_params, load_model_result, **kwargs):
    # Initialize the authorization context for downstream agents and tools
    completion_create_params["authorization_context"] = resolve_authorization_context(
        completion_create_params, **kwargs
    )
    # ... rest of chat function

# In your tools
auth_context = get_authorization_context()
access_token = auth_context.get("access_token")

OAuth 2.0 authentication¶

For external services that support OAuth 2.0, implement OAuth flows in your tools. Create an OAuth application in the service's developer console and set environment variables:

OAUTH_CLIENT_ID=<your_client_id>
OAUTH_CLIENT_SECRET=<your_client_secret>
OAUTH_REDIRECT_URI=<your_redirect_uri>
OAUTH_SCOPE=<required_scopes>

Then, use those environment variables in the OAuth implementation. This example demonstrates a complete OAuth 2.0 authorization code flow for a CrewAI tool. The tool inherits from crewai.tools.BaseTool (the standard base class for tools in CrewAI agent templates), manages its own token lifecycle with in-memory caching, and follows the repository's pattern of retrieving credentials from environment variables.

import os
import time
import requests
from urllib.parse import urlencode
from crewai.tools import BaseTool

class ExampleToolWithOAuth(BaseTool):
    def __init__(self):
        super().__init__()
        self.client_id = os.getenv("OAUTH_CLIENT_ID")
        self.client_secret = os.getenv("OAUTH_CLIENT_SECRET")
        self.redirect_uri = os.getenv("OAUTH_REDIRECT_URI")
        self._access_token = None
        self._token_expires_at = None

    def get_authorization_url(self) -> str:
        params = {
            "client_id": self.client_id,
            "redirect_uri": self.redirect_uri,
            "response_type": "code",
            "scope": "read:data"
        }
        return f"https://oauth.provider.com/authorize?{urlencode(params)}"

    def exchange_code_for_token(self, code: str) -> dict:
        response = requests.post(
            "https://oauth.provider.com/token",
            data={
                "grant_type": "authorization_code",
                "client_id": self.client_id,
                "client_secret": self.client_secret,
                "code": code,
                "redirect_uri": self.redirect_uri
            }
        )
        token_data = response.json()
        self._access_token = token_data.get("access_token")
        expires_in = token_data.get("expires_in", 3600)
        self._token_expires_at = time.time() + expires_in
        return token_data

    def get_cached_access_token(self) -> str:
        if self._access_token and self._token_expires_at and time.time() < self._token_expires_at:
            return self._access_token
        # Token expired or not set - refresh or re-authenticate
        # In production, implement token refresh logic here
        raise ValueError("Access token expired. Re-authenticate to get a new token.")

    def run(self, query: str) -> str:
        access_token = self.get_cached_access_token()
        response = requests.get(
            "https://api.provider.com/data",
            headers={"Authorization": f"Bearer {access_token}"},
            params={"query": query}
        )
        return response.json()

Security best practices¶

Following security best practices is essential when handling authentication in production environments. Adhere to the following guidelines:

• API token authentication
• Store tokens in environment variables; never hard code secrets in source code.
• Use secure token storage solutions in production environments.
• Implement token rotation and enforce token expiration policies.
• Validate authorization context before using it in tools.
• Follow least-privilege access principles.
• Log authentication events for audit purposes.
• OAuth 2.0
• Use HTTPS for all OAuth communications.
• Validate the state parameter to prevent Cross-Site Request Forgery (CSRF) attacks.
• Store refresh tokens securely.
• Handle token expiration and refresh logic reliably.
• Validate authorization context before use.
• General Security
• Use separate environments for development and production.
• Implement robust secret management practices.
• Follow container security best practices.
• Conduct regular security audits and apply updates.

Troubleshooting authentication issues¶

This section helps you diagnose and resolve common authentication problems when developing or deploying agents.

Common issues¶

The following sections describe common authentication errors and how to resolve them:

Missing API token¶

This error occurs when the DataRobot API token is not configured.

Issue: Error: Missing DataRobot API token. Set the DATAROBOT_API_TOKEN environment variable

Solution: Set the DATAROBOT_API_TOKEN environment variable to use your DataRobot API key.

Invalid endpoint¶

This error occurs when the DataRobot endpoint is missing or incorrectly configured.

Issue: Error: Missing DataRobot endpoint. Set the DATAROBOT_ENDPOINT environment variable

Solution: Set the correct DATAROBOT_ENDPOINT environment variable.

Authorization context not set¶

This error occurs when tools try to access the authorization context before it has been initialized.

Issue: Error: Authorization context not available for tool

Solution: Ensure resolve_authorization_context() is called in your agent's chat() function and the result is assigned to completion_create_params["authorization_context"].

OAuth token expired¶

This error occurs when an OAuth access token has expired and needs to be refreshed.

Issue: Error: 401 Unauthorized

Solution: Implement token refresh logic or re-authenticate.

MCP server authentication failure¶

This error occurs when the MCP server cannot authenticate the agent's requests.

Issue: Error: 401 Unauthorized or MCP connection failures

Solutions:

1. Verify API token: Ensure DATAROBOT_API_TOKEN is set correctly in your environment
2. Check token permissions: Verify the token has necessary permissions for MCP server access
3. Verify MCP server endpoint: Check that the MCP server endpoint is correct
4. For local development: Ensure the MCP server is running and accessible
5. Review MCP configuration: See MCP tools troubleshooting for more details

Debugging tips¶

Useful techniques for debugging authentication issues:

Enable verbose logging¶

Enable verbose logging to get detailed information about authentication operations. To do so, search for where MyAgent is instantiated in custom.py and set verbose=True.

from agent import MyAgent

# ...

agent = MyAgent(verbose=True, **completion_create_params)

Check environment variables¶

Verify that required environment variables are set correctly.

import os
print(f"API Token: {os.getenv('DATAROBOT_API_TOKEN')}")
print(f"Endpoint: {os.getenv('DATAROBOT_ENDPOINT')}")

Test authentication¶

Test authentication using the AgentEnvironment class from datarobot_genai.core.cli:

from datarobot_genai.core.cli import AgentEnvironment

try:
    env = AgentEnvironment()
    print("Authentication successful")
except ValueError as e:
    print(f"Authentication failed: {e}")

Validate authorization context¶

Check authorization context using get_authorization_context() from datarobot.models.genai.agent.auth:

from datarobot.models.genai.agent.auth import get_authorization_context

auth_context = get_authorization_context()
print(f"Authorization context: {auth_context}")