|
| 1 | +<!-- |
| 2 | +SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. |
| 3 | +SPDX-License-Identifier: Apache-2.0 |
| 4 | +
|
| 5 | +Licensed under the Apache License, Version 2.0 (the "License"); |
| 6 | +you may not use this file except in compliance with the License. |
| 7 | +You may obtain a copy of the License at |
| 8 | +
|
| 9 | +http://www.apache.org/licenses/LICENSE-2.0 |
| 10 | +
|
| 11 | +Unless required by applicable law or agreed to in writing, software |
| 12 | +distributed under the License is distributed on an "AS IS" BASIS, |
| 13 | +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 14 | +See the License for the specific language governing permissions and |
| 15 | +limitations under the License. |
| 16 | +--> |
| 17 | + |
| 18 | +# Secure Token Storage for MCP Authentication |
| 19 | + |
| 20 | +The NeMo Agent toolkit provides a configurable, secure token storage mechanism for Model Context Protocol (MCP) OAuth2 authentication. You can store tokens securely using the object store infrastructure, which provides encryption at rest, access controls, and persistence across service restarts. |
| 21 | + |
| 22 | +## Overview |
| 23 | + |
| 24 | +When using MCP with OAuth2 authentication, the toolkit needs to store authentication tokens for each user. The secure token storage feature provides: |
| 25 | + |
| 26 | +- **Encryption at rest**: Tokens are stored in object stores that support encryption |
| 27 | +- **Flexible backends**: Choose from in-memory (default), S3, MySQL, Redis, or custom object stores |
| 28 | +- **Persistence**: Tokens persist across restarts when using external storage backends |
| 29 | +- **Multi-user support**: Tokens are isolated per user with proper access controls |
| 30 | +- **Automatic refresh**: Supports OAuth2 token refresh flows |
| 31 | + |
| 32 | +### Components |
| 33 | + |
| 34 | +The token storage system includes three main components: |
| 35 | + |
| 36 | +1. **TokenStorageBase**: Abstract interface defining `store()`, `retrieve()`, `delete()`, and `clear_all()` operations. |
| 37 | +2. **InMemoryTokenStorage**: Default implementation using the in-memory object store. |
| 38 | +3. **ObjectStoreTokenStorage**: Implementation backed by configurable object stores such as S3, MySQL, and Redis. |
| 39 | + |
| 40 | +## Configuration |
| 41 | + |
| 42 | +### Default Configuration (In-Memory Storage) |
| 43 | + |
| 44 | +By default, MCP OAuth2 authentication uses in-memory storage. No additional configuration is required: |
| 45 | + |
| 46 | +```yaml |
| 47 | +authentication: |
| 48 | + mcp_oauth2_jira: |
| 49 | + _type: mcp_oauth2 |
| 50 | + server_url: ${CORPORATE_MCP_JIRA_URL} |
| 51 | + redirect_uri: http://localhost:8000/auth/redirect |
| 52 | + default_user_id: ${CORPORATE_MCP_JIRA_URL} |
| 53 | + allow_default_user_id_for_tool_calls: ${ALLOW_DEFAULT_USER_ID_FOR_TOOL_CALLS:-true} |
| 54 | +``` |
| 55 | +
|
| 56 | +This setup is **ONLY suitable for development and testing environments** since it uses in-memory storage that is not |
| 57 | +persistent and also unsafe. |
| 58 | +
|
| 59 | +### External Object Store Configuration |
| 60 | +
|
| 61 | +For production environments, configure an external object store to persist tokens across restarts. The NeMo Agent toolkit supports S3-compatible storage (MinIO, AWS S3), MySQL, and Redis backends. |
| 62 | +
|
| 63 | +:::{note} |
| 64 | +For detailed object store setup instructions including MinIO, MySQL, and Redis installation and configuration examples, see the `examples/object_store/user_report/README.md` guide (under the "Choose an Object Store" section). |
| 65 | +::: |
| 66 | + |
| 67 | +The following example shows token storage configuration using S3-compatible storage (MinIO): |
| 68 | + |
| 69 | +```yaml |
| 70 | +object_stores: |
| 71 | + token_store: |
| 72 | + _type: s3 |
| 73 | + endpoint_url: http://localhost:9000 |
| 74 | + access_key: minioadmin |
| 75 | + secret_key: minioadmin |
| 76 | + bucket_name: my-bucket |
| 77 | +
|
| 78 | +function_groups: |
| 79 | + mcp_jira: |
| 80 | + _type: mcp_client |
| 81 | + server: |
| 82 | + transport: streamable-http |
| 83 | + url: ${CORPORATE_MCP_JIRA_URL} |
| 84 | + auth_provider: mcp_oauth2_jira |
| 85 | +
|
| 86 | +authentication: |
| 87 | + mcp_oauth2_jira: |
| 88 | + _type: mcp_oauth2 |
| 89 | + server_url: ${CORPORATE_MCP_JIRA_URL} |
| 90 | + redirect_uri: http://localhost:8000/auth/redirect |
| 91 | + default_user_id: ${CORPORATE_MCP_JIRA_URL} |
| 92 | + allow_default_user_id_for_tool_calls: ${ALLOW_DEFAULT_USER_ID_FOR_TOOL_CALLS:-true} |
| 93 | + token_storage_object_store: token_store |
| 94 | +
|
| 95 | +llms: |
| 96 | + nim_llm: |
| 97 | + _type: nim |
| 98 | + model_name: meta/llama-3.1-70b-instruct |
| 99 | + temperature: 0.0 |
| 100 | + max_tokens: 1024 |
| 101 | +
|
| 102 | +workflow: |
| 103 | + _type: react_agent |
| 104 | + tool_names: |
| 105 | + - mcp_jira |
| 106 | + llm_name: nim_llm |
| 107 | + verbose: true |
| 108 | + retry_parsing_errors: true |
| 109 | + max_retries: 3 |
| 110 | +``` |
| 111 | + |
| 112 | +For MySQL or Redis configurations, replace the `object_stores` section with the appropriate object store type. Refer to the [Object Store Documentation](../../store-and-retrieve/object-store.md) for configuration options for each backend. |
| 113 | + |
| 114 | +## Token Storage Format |
| 115 | + |
| 116 | +The system stores tokens as JSON-serialized `AuthResult` objects in the object store with the following structure: |
| 117 | + |
| 118 | +- **Key format**: `tokens/{sha256_hash}` where the hash is computed from the `user_id` to ensure S3 compatibility |
| 119 | +- **Content type**: `application/json` |
| 120 | +- **Metadata**: Includes token expiration timestamp when available |
| 121 | + |
| 122 | +Example stored token: |
| 123 | +```json |
| 124 | +{ |
| 125 | + "credentials": [ |
| 126 | + { |
| 127 | + "kind": "bearer", |
| 128 | + "token": "encrypted_token_value", |
| 129 | + "scheme": "Bearer", |
| 130 | + "header_name": "Authorization" |
| 131 | + } |
| 132 | + ], |
| 133 | + "token_expires_at": "2025-10-02T12:00:00Z", |
| 134 | + "raw": { |
| 135 | + "access_token": "...", |
| 136 | + "refresh_token": "...", |
| 137 | + "expires_at": 1727870400 |
| 138 | + } |
| 139 | +} |
| 140 | +``` |
| 141 | + |
| 142 | +## Token Lifecycle |
| 143 | + |
| 144 | +### 1. Initial Authentication |
| 145 | + |
| 146 | +When a user first authenticates, the system completes the following steps: |
| 147 | +1. The OAuth2 flow completes and returns an access token. |
| 148 | +2. The token is serialized and stored using the configured storage backend. |
| 149 | +3. The token is associated with the user's session ID. |
| 150 | + |
| 151 | +### 2. Token Retrieval |
| 152 | + |
| 153 | +On subsequent requests, the system completes the following steps: |
| 154 | +1. The user's session ID is extracted from cookies. |
| 155 | +2. The stored token is retrieved from the storage backend. |
| 156 | +3. The token expiration is checked. |
| 157 | +4. If expired, a token refresh is attempted. |
| 158 | + |
| 159 | +### 3. Token Refresh |
| 160 | + |
| 161 | +When a token expires, the system completes the following steps: |
| 162 | +1. The refresh token is extracted from the stored token. |
| 163 | +2. A new access token is requested from the OAuth2 provider. |
| 164 | +3. The new token is stored, replacing the old one. |
| 165 | +4. The refreshed token is returned for use. |
| 166 | + |
| 167 | + |
| 168 | +## Custom Token Storage |
| 169 | + |
| 170 | +You can implement custom token storage by extending the `TokenStorageBase` abstract class: |
| 171 | + |
| 172 | +```python |
| 173 | +from nat.plugins.mcp.auth.token_storage import TokenStorageBase |
| 174 | +from nat.data_models.authentication import AuthResult |
| 175 | +
|
| 176 | +class CustomTokenStorage(TokenStorageBase): |
| 177 | + async def store(self, user_id: str, auth_result: AuthResult) -> None: |
| 178 | + # Custom storage logic |
| 179 | + pass |
| 180 | +
|
| 181 | + async def retrieve(self, user_id: str) -> AuthResult | None: |
| 182 | + # Custom retrieval logic |
| 183 | + pass |
| 184 | +
|
| 185 | + async def delete(self, user_id: str) -> None: |
| 186 | + # Custom deletion logic |
| 187 | + pass |
| 188 | +
|
| 189 | + async def clear_all(self) -> None: |
| 190 | + # Custom clear logic |
| 191 | + pass |
| 192 | +``` |
| 193 | + |
| 194 | +Then configure your custom storage in the MCP provider initialization. |
| 195 | + |
| 196 | + |
| 197 | +## Related Documentation |
| 198 | + |
| 199 | +- [MCP Client Configuration](mcp-client.md) |
| 200 | +- [Object Store Documentation](../../store-and-retrieve/object-store.md) |
| 201 | +- [Authentication API Reference](../../reference/api-authentication.md) |
| 202 | +- [Extending Object Stores](../../extend/object-store.md) |
0 commit comments