Performing On-behalf of Token (OBO) token exchange in AgentCore Gateway (#1435)
* OBO token exchange * Update README.md * Update token_callback_server.py * Update token_callback_server.py
This commit is contained in:
Dhawalkumar Patel
committed by
GitHub
GitHub
parent
43b7a14354
commit
de8e80bd1d
@@ -0,0 +1,6 @@
|
||||
.env
|
||||
__pycache__/
|
||||
*.pyc
|
||||
.ipynb_checkpoints/
|
||||
.bedrock_agentcore.yaml
|
||||
.venv/
|
||||
@@ -0,0 +1,85 @@
|
||||
# OBO Token Exchange with AgentCore Gateway + Microsoft Graph API
|
||||
|
||||
## Overview
|
||||
|
||||
This tutorial demonstrates how to use **Amazon Bedrock AgentCore Gateway** with **On-Behalf-Of (OBO) token exchange** to expose Microsoft Graph API endpoints as MCP tools. The agent code contains zero token handling logic — the Gateway handles the OBO exchange transparently at the infrastructure level.
|
||||
|
||||
OBO token exchange enables agents to access protected resources on behalf of authenticated users without triggering additional consent flows. The Gateway exchanges the inbound user's access token for a new, scoped access token that targets a downstream resource server (Microsoft Graph), using the JWT Authorization Grant (RFC 7523).
|
||||
|
||||
### How It Works
|
||||
|
||||
1. User authenticates directly with **Microsoft Entra ID** and gets an access token scoped to the app (`api://<client-id>/access_as_user`)
|
||||
2. User passes that Entra ID token to the **AgentCore Gateway** as the bearer
|
||||
3. Gateway validates the token using Entra ID's OIDC discovery URL (inbound auth)
|
||||
4. Gateway performs **OBO token exchange** — swaps the app-scoped Entra ID token for a Microsoft Graph token using JWT Authorization Grant (RFC 7523)
|
||||
5. A **Strands agent** connects to the Gateway MCP URL, discovers tools, and invokes them
|
||||
|
||||
### Architecture
|
||||
|
||||
```mermaid
|
||||
sequenceDiagram
|
||||
participant User
|
||||
participant Entra as Microsoft Entra ID
|
||||
participant Agent as Strands Agent
|
||||
participant GW as AgentCore Gateway
|
||||
participant Graph as Microsoft Graph API
|
||||
|
||||
User->>Entra: 1. Sign in (browser)
|
||||
Entra-->>User: 2. Access token (aud: api://app, scp: access_as_user)
|
||||
|
||||
User->>Agent: 3. Run agent with bearer token
|
||||
Agent->>GW: 4. MCP tool call (e.g. getMyProfile)
|
||||
|
||||
Note over GW: Validate JWT via Entra ID OIDC
|
||||
|
||||
GW->>Entra: 5. OBO exchange (jwt-bearer grant, RFC 7523)
|
||||
Entra-->>GW: 6. Graph token (aud: graph.microsoft.com)
|
||||
|
||||
GW->>Graph: 7. GET /me with Graph token
|
||||
Graph-->>GW: 8. {displayName: "John Doe"}
|
||||
|
||||
GW-->>Agent: 9. MCP tool response
|
||||
Agent-->>User: 10. "Your name is John Doe"
|
||||
```
|
||||
|
||||
### Key Concepts
|
||||
|
||||
**Why `CustomOauth2` (not `MicrosoftOAuth2`)?** The `onBehalfOfTokenExchangeConfig` parameter is only available inside `customOauth2ProviderConfig`, not the built-in Microsoft provider.
|
||||
|
||||
**Why `requested_token_use: on_behalf_of` in customParameters?** Entra ID's token endpoint requires this parameter to perform the OBO exchange. Without it, the exchange fails silently.
|
||||
|
||||
**Why v1.0 discovery URL?** Entra ID issues v1.0 access tokens by default (issuer: `sts.windows.net`). The Gateway's inbound auth discovery URL must match the token version.
|
||||
|
||||
### Tutorial Details
|
||||
|
||||
| Information | Details |
|
||||
|:---------------------|:-------------------------------------------------------------------------|
|
||||
| Tutorial type | Interactive |
|
||||
| AgentCore components | AgentCore Gateway, AgentCore Identity |
|
||||
| Agentic Framework | Strands Agents |
|
||||
| LLM model | Anthropic Claude Haiku 4.5 |
|
||||
| Tutorial components | AgentCore Gateway with OBO Token Exchange, Microsoft Graph API |
|
||||
| Tutorial vertical | Cross-vertical |
|
||||
| Example complexity | Medium |
|
||||
| SDK used | boto3, strands-agents, mcp |
|
||||
| Credential Provider | CustomOauth2 with JWT_AUTHORIZATION_GRANT OBO config |
|
||||
| Inbound Auth | Entra ID OIDC (CUSTOM_JWT) |
|
||||
| Gateway Target | OpenAPI Schema (Microsoft Graph API) |
|
||||
|
||||
## Prerequisites
|
||||
|
||||
- Python 3.10+
|
||||
- AWS credentials configured
|
||||
- A Microsoft 365 **work or school account** with access to Microsoft Entra ID
|
||||
|
||||
> ⚠️ **Personal Microsoft accounts** (`@outlook.com`, `@hotmail.com`, `@live.com`) will not work for the calendar/email endpoints. The OBO exchange itself works, but Microsoft Graph calendar and mail endpoints require an Exchange Online mailbox (work/school accounts only). The `/me` profile endpoint works with all account types.
|
||||
|
||||
## Tutorial
|
||||
|
||||
- [OBO Token Exchange with AgentCore Gateway + Microsoft Graph API](obo_token_exchange_microsoft.ipynb)
|
||||
|
||||
## Reference Documentation
|
||||
|
||||
- [OBO Token Exchange](https://docs.aws.amazon.com/bedrock-agentcore/latest/devguide/on-behalf-of-token-exchange.html)
|
||||
- [Gateway Outbound Auth](https://docs.aws.amazon.com/bedrock-agentcore/latest/devguide/gateway-outbound-auth.html)
|
||||
- [Microsoft Entra ID OBO Flow](https://learn.microsoft.com/en-us/entra/identity-platform/v2-oauth2-on-behalf-of-flow)
|
||||
+751
@@ -0,0 +1,751 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "overview",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"# OBO Token Exchange with AgentCore Gateway + Microsoft Graph API\n",
|
||||
"\n",
|
||||
"This tutorial shows how to use **AgentCore Gateway** to expose the Microsoft Graph API as MCP tools, with **OBO (On-Behalf-Of) token exchange** handling authentication transparently. The agent code is just a few lines — no token handling, no session binding, no Runtime deployment.\n",
|
||||
"\n",
|
||||
"### How It Works\n",
|
||||
"\n",
|
||||
"1. User authenticates directly with **Microsoft Entra ID** and gets an access token scoped to the app (`api://<client-id>/access_as_user`)\n",
|
||||
"2. User passes that Entra ID token to the **AgentCore Gateway** as the bearer\n",
|
||||
"3. Gateway validates the token using Entra ID's OIDC discovery URL (inbound auth)\n",
|
||||
"4. Gateway performs **OBO token exchange** — swaps the app-scoped Entra ID token for a Microsoft Graph token using JWT Authorization Grant (RFC 7523)\n",
|
||||
"5. A **Strands agent** connects to the Gateway MCP URL, discovers tools, and invokes them\n",
|
||||
"\n",
|
||||
"**Docs:** [OBO Token Exchange](https://docs.aws.amazon.com/bedrock-agentcore/latest/devguide/on-behalf-of-token-exchange.html) · [Gateway Outbound Auth](https://docs.aws.amazon.com/bedrock-agentcore/latest/devguide/gateway-outbound-auth.html)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "prerequisites",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Prerequisites\n",
|
||||
"\n",
|
||||
"- Python 3.10+\n",
|
||||
"- AWS credentials configured\n",
|
||||
"- Amazon Bedrock AgentCore SDK\n",
|
||||
"- A Microsoft 365 **work or school account** with access to Microsoft Entra ID\n",
|
||||
"\n",
|
||||
"> ⚠️ **Personal Microsoft accounts (`@outlook.com`, `@hotmail.com`, `@live.com`) will not work for the calendar/email demo.** The OBO token exchange itself works with personal accounts, but Microsoft Graph calendar and mail endpoints require an Exchange Online mailbox, which is only available with work/school accounts. Get a free sandbox at the [Microsoft 365 Developer Program](https://developer.microsoft.com/en-us/microsoft-365/dev-program)."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "install-deps",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"!pip3 install -U -r requirements.txt --quiet"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "step1-header",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"---\n",
|
||||
"## Step 1: Register Application in Microsoft Entra ID\n",
|
||||
"\n",
|
||||
"### 1.1 Create App Registration\n",
|
||||
"1. Go to [entra.microsoft.com](https://entra.microsoft.com) → **Identity** → **Applications** → **App registrations**\n",
|
||||
"2. Click **+ New registration**\n",
|
||||
"3. Name: `AgentCore-OBO-Tutorial`, Single tenant\n",
|
||||
"4. Click **Register**\n",
|
||||
"\n",
|
||||
"### 1.2 Record IDs\n",
|
||||
"From the **Overview** page, copy:\n",
|
||||
"- **Application (client) ID** → `MICROSOFT_CLIENT_ID`\n",
|
||||
"- **Directory (tenant) ID** → `MICROSOFT_TENANT_ID`\n",
|
||||
"\n",
|
||||
"### 1.3 Create Client Secret\n",
|
||||
"1. **Certificates & secrets** → **+ New client secret**\n",
|
||||
"2. Copy the **Value** immediately → `MICROSOFT_CLIENT_SECRET`\n",
|
||||
"\n",
|
||||
"### 1.4 Configure API Permissions\n",
|
||||
"1. **API permissions** → **+ Add a permission** → **Microsoft Graph** → **Delegated permissions**\n",
|
||||
"2. Add: `Calendars.Read`, `Mail.Read`, `User.Read`\n",
|
||||
"\n",
|
||||
"### 1.5 Expose an API (Required for OBO)\n",
|
||||
"1. **Expose an API** → Set **Application ID URI** (accept default `api://<client-id>`)\n",
|
||||
"2. **+ Add a scope**: name `access_as_user`, consent by Admins and users, Enabled\n",
|
||||
"\n",
|
||||
"### 1.6 Configure Authentication (Redirect URI)\n",
|
||||
"1. **Authentication** → **+ Add a platform** → **Web**\n",
|
||||
"2. Add redirect URI: `http://localhost:9090/oauth2/callback`\n",
|
||||
"3. Click **Configure**\n",
|
||||
"\n",
|
||||
"> This redirect URI is used by the token callback server to receive the authorization code.\n",
|
||||
"\n",
|
||||
"### 1.7 Grant Admin Consent\n",
|
||||
"1. **API permissions** → **Grant admin consent for [tenant]** → **Yes**\n",
|
||||
"\n",
|
||||
"> ✅ All permissions should show green checkmarks.\n",
|
||||
"\n",
|
||||
"### Important: Token Version\n",
|
||||
"By default, Entra ID issues **v1.0 access tokens** (issuer: `https://sts.windows.net/{tenant}/`). The Gateway discovery URL in this tutorial uses the v1.0 endpoint to match. If you need v2.0 tokens, set `accessTokenAcceptedVersion: 2` in the app manifest and use the v2.0 discovery URL instead."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "step2-header",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"---\n",
|
||||
"## Step 2: Create the OBO Credential Provider\n",
|
||||
"\n",
|
||||
"Fill in your Entra ID credentials in the `.env` file below, then create the credential provider.\n",
|
||||
"\n",
|
||||
"**Important:** The `onBehalfOfTokenExchangeConfig` is only available inside `customOauth2ProviderConfig`, so we use `CustomOauth2` as the vendor (not the built-in `MicrosoftOAuth2` provider)."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "write-env",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"%%writefile .env\n",
|
||||
"MICROSOFT_CLIENT_ID=\"\" # Application (client) ID from Entra ID\n",
|
||||
"MICROSOFT_CLIENT_SECRET=\"\" # Client secret value from Entra ID\n",
|
||||
"MICROSOFT_TENANT_ID=\"\" # Directory (tenant) ID from Entra ID"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "setup-clients",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import os\n",
|
||||
"import json\n",
|
||||
"import time\n",
|
||||
"import urllib.parse\n",
|
||||
"\n",
|
||||
"import boto3\n",
|
||||
"from boto3.session import Session\n",
|
||||
"from dotenv import dotenv_values\n",
|
||||
"\n",
|
||||
"# Load .env and set environment variables\n",
|
||||
"env = dotenv_values('.env')\n",
|
||||
"os.environ.update(env)\n",
|
||||
"\n",
|
||||
"boto_session = Session()\n",
|
||||
"region = boto_session.region_name\n",
|
||||
"print(f\"Region: {region}\")\n",
|
||||
"\n",
|
||||
"identity_client = boto_session.client(\"bedrock-agentcore-control\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "create-provider",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"tenant_id = os.environ[\"MICROSOFT_TENANT_ID\"]\n",
|
||||
"client_id = os.environ[\"MICROSOFT_CLIENT_ID\"]\n",
|
||||
"\n",
|
||||
"# Create Custom OAuth2 credential provider with OBO token exchange\n",
|
||||
"# Note: onBehalfOfTokenExchangeConfig is only available inside customOauth2ProviderConfig\n",
|
||||
"microsoft_provider = identity_client.create_oauth2_credential_provider(\n",
|
||||
" name=\"microsoft-obo-provider-v3\",\n",
|
||||
" credentialProviderVendor=\"CustomOauth2\",\n",
|
||||
" oauth2ProviderConfigInput={\n",
|
||||
" \"customOauth2ProviderConfig\": {\n",
|
||||
" \"oauthDiscovery\": {\n",
|
||||
" \"discoveryUrl\": f\"https://login.microsoftonline.com/{tenant_id}/v2.0/.well-known/openid-configuration\"\n",
|
||||
" },\n",
|
||||
" \"clientId\": os.environ[\"MICROSOFT_CLIENT_ID\"],\n",
|
||||
" \"clientSecret\": os.environ[\"MICROSOFT_CLIENT_SECRET\"],\n",
|
||||
" \"clientAuthenticationMethod\": \"CLIENT_SECRET_POST\",\n",
|
||||
" \"onBehalfOfTokenExchangeConfig\": {\n",
|
||||
" \"grantType\": \"JWT_AUTHORIZATION_GRANT\"\n",
|
||||
" }\n",
|
||||
" }\n",
|
||||
" }\n",
|
||||
")\n",
|
||||
"print(\"Credential provider created \\u2713\")\n",
|
||||
"print(f\"ARN: {microsoft_provider['credentialProviderArn']}\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "step3-header",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"---\n",
|
||||
"## Step 3: Create the AgentCore Gateway\n",
|
||||
"\n",
|
||||
"The Gateway \"MCPfies\" Microsoft Graph API:\n",
|
||||
"- **Inbound auth**: Entra ID OIDC JWT authorizer (CUSTOM_JWT) — validates tokens issued directly by Entra ID\n",
|
||||
"- **Target**: OpenAPI schema exposing calendar, email, and profile endpoints as MCP tools\n",
|
||||
"- **Outbound auth**: OBO token exchange via the credential provider\n",
|
||||
"\n",
|
||||
"**Key insight**: We use Entra ID's **v1.0** OIDC discovery URL for inbound auth. Entra ID issues v1.0 access tokens by default (issuer: `sts.windows.net`), so the discovery URL must match. Using the v2.0 URL would cause a 403 because the signing keys and issuer don't match.\n",
|
||||
"\n",
|
||||
"First, we create an IAM role for the Gateway, then create the Gateway itself."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "create-iam-role",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"iam_client = boto3.client(\"iam\")\n",
|
||||
"account_id = boto3.client(\"sts\").get_caller_identity()[\"Account\"]\n",
|
||||
"\n",
|
||||
"# Use the shared utils to create the Gateway IAM role\n",
|
||||
"# This creates a role with bedrock-agentcore:*, secretsmanager:GetSecretValue, etc.\n",
|
||||
"# matching the pattern used by other AgentCore Gateway tutorials\n",
|
||||
"import sys\n",
|
||||
"gw_utils_dir = os.path.abspath(os.path.join(os.getcwd(), '..', '..', '02-AgentCore-gateway'))\n",
|
||||
"sys.path.insert(0, gw_utils_dir)\n",
|
||||
"try:\n",
|
||||
" from utils import create_agentcore_gateway_role\n",
|
||||
" gateway_role_name = f'agentcore-microsoft-obo-gateway-role'\n",
|
||||
" gateway_iam_role = create_agentcore_gateway_role('microsoft-obo-gateway')\n",
|
||||
" gateway_role_arn = gateway_iam_role['Role']['Arn']\n",
|
||||
" print(f'Gateway role: {gateway_role_arn}')\n",
|
||||
"except ImportError:\n",
|
||||
" print('Shared utils not found, creating role inline...')\n",
|
||||
" gateway_role_name = 'agentcore-obo-gateway-role'\n",
|
||||
" assume_role_policy = {\n",
|
||||
" 'Version': '2012-10-17',\n",
|
||||
" 'Statement': [{\n",
|
||||
" 'Effect': 'Allow',\n",
|
||||
" 'Principal': {'Service': 'bedrock-agentcore.amazonaws.com'},\n",
|
||||
" 'Action': 'sts:AssumeRole',\n",
|
||||
" 'Condition': {\n",
|
||||
" 'StringEquals': {'aws:SourceAccount': account_id},\n",
|
||||
" 'ArnLike': {'aws:SourceArn': f'arn:aws:bedrock-agentcore:{region}:{account_id}:*'}\n",
|
||||
" }\n",
|
||||
" }]\n",
|
||||
" }\n",
|
||||
" try:\n",
|
||||
" role_response = iam_client.create_role(RoleName=gateway_role_name, AssumeRolePolicyDocument=json.dumps(assume_role_policy))\n",
|
||||
" gateway_role_arn = role_response['Role']['Arn']\n",
|
||||
" except iam_client.exceptions.EntityAlreadyExistsException:\n",
|
||||
" gateway_role_arn = f'arn:aws:iam::{account_id}:role/{gateway_role_name}'\n",
|
||||
" gateway_policy = {\n",
|
||||
" 'Version': '2012-10-17',\n",
|
||||
" 'Statement': [\n",
|
||||
" {'Sid': 'AgentCoreAccess', 'Effect': 'Allow', 'Action': ['bedrock-agentcore:*'], 'Resource': '*'},\n",
|
||||
" {'Sid': 'SecretsManagerAccess', 'Effect': 'Allow', 'Action': ['secretsmanager:GetSecretValue'], 'Resource': [f'arn:aws:secretsmanager:{region}:{account_id}:secret:*bedrock-agentcore*']}\n",
|
||||
" ]\n",
|
||||
" }\n",
|
||||
" iam_client.put_role_policy(RoleName=gateway_role_name, PolicyName='GatewayOBOPolicy', PolicyDocument=json.dumps(gateway_policy))\n",
|
||||
" print(f'Gateway role: {gateway_role_arn}')\n",
|
||||
"\n",
|
||||
"print('Waiting for IAM propagation (10s)...')\n",
|
||||
"time.sleep(10)\n",
|
||||
"print('Done \\u2713')"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "create-gateway",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"agentcore_control_client = boto3.client(\"bedrock-agentcore-control\", region_name=region)\n",
|
||||
"\n",
|
||||
"# Entra ID v1.0 OIDC discovery URL for inbound auth\n",
|
||||
"# IMPORTANT: Entra ID issues v1.0 access tokens by default (issuer: sts.windows.net).\n",
|
||||
"# The discovery URL MUST match the token version — use v1.0, not v2.0.\n",
|
||||
"discovery_url = f\"https://login.microsoftonline.com/{tenant_id}/.well-known/openid-configuration\"\n",
|
||||
"\n",
|
||||
"# Create Gateway with Entra ID JWT authorizer\n",
|
||||
"# allowedAudience uses the Application ID URI (api://<client-id>) since the\n",
|
||||
"# Entra ID access token will have this as the \"aud\" claim\n",
|
||||
"gateway_response = agentcore_control_client.create_gateway(\n",
|
||||
" name=\"microsoft-obo-gateway-v3\",\n",
|
||||
" roleArn=gateway_role_arn,\n",
|
||||
" protocolType=\"MCP\",\n",
|
||||
" authorizerType=\"CUSTOM_JWT\",\n",
|
||||
" authorizerConfiguration={\n",
|
||||
" \"customJWTAuthorizer\": {\n",
|
||||
" \"discoveryUrl\": discovery_url,\n",
|
||||
" \"allowedAudience\": [f\"api://{client_id}\"],\n",
|
||||
" }\n",
|
||||
" },\n",
|
||||
" exceptionLevel=\"DEBUG\" # Show detailed errors for debugging\n",
|
||||
")\n",
|
||||
"gateway_id = gateway_response[\"gatewayId\"]\n",
|
||||
"print(f\"Gateway created \\u2713 ID: {gateway_id}\")\n",
|
||||
"print(f\"Inbound auth: Entra ID OIDC (tenant: {tenant_id})\")\n",
|
||||
"print(f\"Allowed audience: api://{client_id}\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "step3b-header",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### Step 3b: Wait for Gateway, then Add Microsoft Graph Target\n",
|
||||
"\n",
|
||||
"The Gateway must be in `READY` status before adding targets. Then we define the OpenAPI schema for Microsoft Graph endpoints — the Gateway exposes these as MCP tools.\n",
|
||||
"\n",
|
||||
"**Important notes:**\n",
|
||||
"- Use `top` and `select` as parameter names (not `$top`/`$select`) — the `$` character breaks Bedrock's tool schema validation\n",
|
||||
"- The `customParameters` with `requested_token_use: on_behalf_of` is **required** for Entra ID's OBO endpoint"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "create-target",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# Wait for Gateway to become READY before adding targets\n",
|
||||
"end_statuses = [\"READY\", \"FAILED\", \"CREATE_FAILED\", \"UPDATE_FAILED\", \"DELETE_FAILED\"]\n",
|
||||
"status = \"CREATING\"\n",
|
||||
"print(\"Waiting for Gateway to become READY...\")\n",
|
||||
"while status not in end_statuses:\n",
|
||||
" time.sleep(10)\n",
|
||||
" gw = agentcore_control_client.get_gateway(gatewayIdentifier=gateway_id)\n",
|
||||
" status = gw[\"status\"]\n",
|
||||
" print(f\" Gateway status: {status}\")\n",
|
||||
"print(f\"Gateway is {status} \\u2713\" if status == \"READY\" else f\"\\u274c Gateway: {status}\")\n",
|
||||
"\n",
|
||||
"# Define the OpenAPI spec for Microsoft Graph endpoints\n",
|
||||
"openapi_spec = json.dumps({\n",
|
||||
" \"openapi\": \"3.0.0\",\n",
|
||||
" \"info\": {\"title\": \"Microsoft Graph Calendar API\", \"version\": \"1.0\"},\n",
|
||||
" \"servers\": [{\"url\": \"https://graph.microsoft.com/v1.0\"}],\n",
|
||||
" \"paths\": {\n",
|
||||
" \"/me/calendarview\": {\n",
|
||||
" \"get\": {\n",
|
||||
" \"operationId\": \"listCalendarEvents\",\n",
|
||||
" \"summary\": \"List calendar events for the current user within a date range\",\n",
|
||||
" \"parameters\": [\n",
|
||||
" {\"name\": \"startDateTime\", \"in\": \"query\", \"required\": True, \"schema\": {\"type\": \"string\"}, \"description\": \"Start date/time in ISO 8601 format (e.g. 2025-01-01T00:00:00Z)\"},\n",
|
||||
" {\"name\": \"endDateTime\", \"in\": \"query\", \"required\": True, \"schema\": {\"type\": \"string\"}, \"description\": \"End date/time in ISO 8601 format (e.g. 2025-01-02T00:00:00Z)\"}\n",
|
||||
" ],\n",
|
||||
" \"responses\": {\"200\": {\"description\": \"Calendar events\"}}\n",
|
||||
" }\n",
|
||||
" },\n",
|
||||
" \"/me/messages\": {\n",
|
||||
" \"get\": {\n",
|
||||
" \"operationId\": \"listUserMails\",\n",
|
||||
" \"summary\": \"List recent email messages for the current user\",\n",
|
||||
" \"parameters\": [\n",
|
||||
" {\"name\": \"top\", \"in\": \"query\", \"required\": False, \"schema\": {\"type\": \"integer\"}, \"description\": \"Number of messages to return\"},\n",
|
||||
" {\"name\": \"select\", \"in\": \"query\", \"required\": False, \"schema\": {\"type\": \"string\"}, \"description\": \"Comma-separated list of fields to return\"}\n",
|
||||
" ],\n",
|
||||
" \"responses\": {\"200\": {\"description\": \"Email messages\"}}\n",
|
||||
" }\n",
|
||||
" },\n",
|
||||
" \"/me\": {\n",
|
||||
" \"get\": {\n",
|
||||
" \"operationId\": \"getMyProfile\",\n",
|
||||
" \"summary\": \"Get the current user profile information\",\n",
|
||||
" \"responses\": {\"200\": {\"description\": \"User profile\"}}\n",
|
||||
" }\n",
|
||||
" }\n",
|
||||
" }\n",
|
||||
"})\n",
|
||||
"\n",
|
||||
"# Create the Gateway target with OBO outbound auth\n",
|
||||
"target_response = agentcore_control_client.create_gateway_target(\n",
|
||||
" gatewayIdentifier=gateway_id,\n",
|
||||
" name=\"microsoft-graph-obo\",\n",
|
||||
" description=\"Microsoft Graph API with OBO token exchange\",\n",
|
||||
" targetConfiguration={\n",
|
||||
" \"mcp\": {\n",
|
||||
" \"openApiSchema\": {\n",
|
||||
" \"inlinePayload\": openapi_spec\n",
|
||||
" }\n",
|
||||
" }\n",
|
||||
" },\n",
|
||||
" credentialProviderConfigurations=[{\n",
|
||||
" \"credentialProviderType\": \"OAUTH\",\n",
|
||||
" \"credentialProvider\": {\n",
|
||||
" \"oauthCredentialProvider\": {\n",
|
||||
" \"providerArn\": microsoft_provider[\"credentialProviderArn\"],\n",
|
||||
" \"scopes\": [\"https://graph.microsoft.com/.default\"],\n",
|
||||
" \"grantType\": \"TOKEN_EXCHANGE\",\n",
|
||||
" \"customParameters\": {\n",
|
||||
" \"requested_token_use\": \"on_behalf_of\"\n",
|
||||
" }\n",
|
||||
" }\n",
|
||||
" }\n",
|
||||
" }]\n",
|
||||
")\n",
|
||||
"target_id = target_response[\"targetId\"]\n",
|
||||
"print(f\"Target created \\u2713 ID: {target_id}\")\n",
|
||||
"print(f\"Target name: {target_response['name']}\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "step3c-header",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### Step 3c: Wait for Target\n",
|
||||
"\n",
|
||||
"Wait for the Target to reach `READY` status."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "poll-gateway",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# Wait for Target to become READY\n",
|
||||
"target_end_statuses = [\"READY\", \"FAILED\", \"UPDATE_UNSUCCESSFUL\", \"SYNCHRONIZE_UNSUCCESSFUL\"]\n",
|
||||
"target_status = \"CREATING\"\n",
|
||||
"\n",
|
||||
"print(\"Waiting for Target to become READY...\")\n",
|
||||
"while target_status not in target_end_statuses:\n",
|
||||
" time.sleep(10)\n",
|
||||
" tgt = agentcore_control_client.get_gateway_target(\n",
|
||||
" gatewayIdentifier=gateway_id,\n",
|
||||
" targetId=target_id\n",
|
||||
" )\n",
|
||||
" target_status = tgt[\"status\"]\n",
|
||||
" print(f\" Target status: {target_status}\")\n",
|
||||
"\n",
|
||||
"if target_status == \"READY\":\n",
|
||||
" print(f\"\\nTarget is READY \\u2713\")\n",
|
||||
"else:\n",
|
||||
" print(f\"\\n\\u274c Target ended in status: {target_status}\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "get-mcp-url",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# Get the Gateway MCP URL\n",
|
||||
"gateway_details = agentcore_control_client.get_gateway(gatewayIdentifier=gateway_id)\n",
|
||||
"gateway_mcp_url = gateway_details.get(\"gatewayUrl\", \"\")\n",
|
||||
"print(f\"Gateway MCP URL: {gateway_mcp_url}\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "step4-header",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"---\n",
|
||||
"## Step 4: Get an Entra ID Access Token\n",
|
||||
"\n",
|
||||
"This cell starts an environment-aware callback server, opens the Entra ID login page in your browser, and automatically captures the token when you sign in. No copy-paste needed.\n",
|
||||
"\n",
|
||||
"### Environment-Aware OAuth2 Callback Server\n",
|
||||
"\n",
|
||||
"The `token_callback_server.py` automatically adapts to different execution environments:\n",
|
||||
"\n",
|
||||
"**Local Development:**\n",
|
||||
"- External Callback URL: `http://localhost:9090/oauth2/callback` (browser-accessible)\n",
|
||||
"- Internal Communication: `http://localhost:9090` (notebook ↔ server)\n",
|
||||
"- Server Binding: `127.0.0.1` (localhost only, secure)\n",
|
||||
"\n",
|
||||
"**SageMaker Workshop Studio:**\n",
|
||||
"- External Callback URL: `https://<domain>.studio.<region>.sagemaker.aws/proxy/9090/oauth2/callback` (browser-accessible via proxy)\n",
|
||||
"- Internal Communication: `http://localhost:9090` (notebook ↔ server in same container)\n",
|
||||
"- Server Binding: `0.0.0.0` (allows SageMaker proxy to reach server)\n",
|
||||
"\n",
|
||||
"The server detects the environment by checking for `/opt/ml/metadata/resource-metadata.json` and configures itself accordingly.\n",
|
||||
"\n",
|
||||
"> ⚠️ **Register the callback URL with Entra ID:** The redirect URI in your Entra ID app registration (Step 1.6) must match the callback URL for your environment. For local development use `http://localhost:9090/oauth2/callback`. For SageMaker, use the proxy URL printed by the cell below."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "print-authorize-url",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import subprocess\n",
|
||||
"import webbrowser\n",
|
||||
"import base64\n",
|
||||
"import sys\n",
|
||||
"from token_callback_server import (\n",
|
||||
" is_server_running, get_callback_url, wait_for_server_ready,\n",
|
||||
" wait_for_token, get_callback_base_url\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"bearer_token = None\n",
|
||||
"\n",
|
||||
"# Start the token callback server if not already running\n",
|
||||
"if is_server_running():\n",
|
||||
" print('Token callback server already running, skipping start...')\n",
|
||||
"else:\n",
|
||||
" server_cmd = [\n",
|
||||
" sys.executable, 'token_callback_server.py',\n",
|
||||
" tenant_id, client_id, os.environ['MICROSOFT_CLIENT_SECRET']\n",
|
||||
" ]\n",
|
||||
" server_process = subprocess.Popen(server_cmd)\n",
|
||||
" if not wait_for_server_ready():\n",
|
||||
" print('\\u274c Failed to start token callback server')\n",
|
||||
" else:\n",
|
||||
" print('Token callback server started \\u2713')\n",
|
||||
"\n",
|
||||
"# Build authorize URL using the environment-aware callback URL\n",
|
||||
"callback_url = get_callback_url()\n",
|
||||
"authorize_url = (\n",
|
||||
" f'https://login.microsoftonline.com/{tenant_id}/oauth2/v2.0/authorize?'\n",
|
||||
" f'client_id={client_id}&response_type=code&'\n",
|
||||
" f'redirect_uri={urllib.parse.quote(callback_url)}&'\n",
|
||||
" f'scope={urllib.parse.quote(f\"api://{client_id}/access_as_user openid profile email\")}'\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"print(f'Callback URL: {callback_url}')\n",
|
||||
"print('Opening browser for Microsoft sign-in...')\n",
|
||||
"webbrowser.open(authorize_url)\n",
|
||||
"print('Waiting for sign-in (up to 2 minutes)...')\n",
|
||||
"\n",
|
||||
"bearer_token = wait_for_token(timeout=120)\n",
|
||||
"\n",
|
||||
"if bearer_token:\n",
|
||||
" payload = bearer_token.split('.')[1]\n",
|
||||
" payload += '=' * (4 - len(payload) % 4)\n",
|
||||
" claims = json.loads(base64.urlsafe_b64decode(payload))\n",
|
||||
" print(f'\\n\\u2705 Token captured for: {claims.get(\"name\", \"unknown\")}')\n",
|
||||
" print(f' aud: {claims[\"aud\"]}')\n",
|
||||
" print(f' scp: {claims.get(\"scp\", \"N/A\")}')\n",
|
||||
" print(f' iss: {claims[\"iss\"]}')\n",
|
||||
" print(f' ver: {claims[\"ver\"]}')\n",
|
||||
"else:\n",
|
||||
" print('\\u274c Timed out. Run this cell again.')"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "step5-header",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"---\n",
|
||||
"## Step 5: Run the Agent\n",
|
||||
"\n",
|
||||
"This is the entire agent — just a few lines. The `MCPClient` connects to the Gateway, discovers the MCP tools (profile, calendar, email), and the Strands Agent uses them to answer the user's question.\n",
|
||||
"\n",
|
||||
"**No `@requires_access_token`, no custom tools, no callback server, no Docker.**\n",
|
||||
"\n",
|
||||
"> 💡 The `getMyProfile` tool works with all account types. The `listCalendarEvents` and `listUserMails` tools require a work/school account with an Exchange Online mailbox."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "run-agent",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from strands import Agent\n",
|
||||
"from strands.tools.mcp import MCPClient\n",
|
||||
"from mcp.client.streamable_http import streamablehttp_client\n",
|
||||
"\n",
|
||||
"# Connect to the AgentCore Gateway as an MCP server\n",
|
||||
"mcp_client = MCPClient(\n",
|
||||
" lambda: streamablehttp_client(\n",
|
||||
" gateway_mcp_url,\n",
|
||||
" headers={\"Authorization\": f\"Bearer {bearer_token}\"}\n",
|
||||
" )\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"with mcp_client:\n",
|
||||
" # The Gateway provides MCP tools: listCalendarEvents, listUserMails, getMyProfile\n",
|
||||
" tools = mcp_client.list_tools_sync()\n",
|
||||
" print(f\"Discovered {len(tools)} MCP tools from Gateway:\")\n",
|
||||
" for t in tools:\n",
|
||||
" print(f\" - {t.tool_name}\")\n",
|
||||
"\n",
|
||||
" agent = Agent(\n",
|
||||
" model=\"us.anthropic.claude-haiku-4-5-20251001-v1:0\",\n",
|
||||
" tools=tools,\n",
|
||||
" system_prompt=\"You are a Microsoft 365 assistant. Use the available tools to help users with their Microsoft profile, Outlook calendar, and email.\"\n",
|
||||
" )\n",
|
||||
"\n",
|
||||
" response = agent(\"What is my Microsoft profile information?\")\n",
|
||||
" print(response)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "step5b-header",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"### Step 5b: Compare User Token vs OBO Token\n",
|
||||
"\n",
|
||||
"The OBO exchange transforms the token in two critical ways:\n",
|
||||
"\n",
|
||||
"| | User Token (before OBO) | OBO Token (after OBO) |\n",
|
||||
"|---|---|---|\n",
|
||||
"| **Audience (`aud`)** | Your app: `api://<client-id>` | Microsoft Graph: `https://graph.microsoft.com` |\n",
|
||||
"| **Scopes (`scp`)** | `access_as_user` | `Calendars.Read Mail.Read User.Read` |\n",
|
||||
"| **Identity** | User's identity | Same user — preserved through delegation |\n",
|
||||
"\n",
|
||||
"The OBO token also carries a **delegation chain** via the `xms_st` claim, which contains the original user's `sub` from the inbound token. This is how Microsoft Graph knows the request comes from an app acting **on behalf of** a specific user, not the app itself.\n",
|
||||
"\n",
|
||||
"Let's decode both tokens and see the difference:"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "compare-tokens",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import requests as req\n",
|
||||
"\n",
|
||||
"# Perform the OBO exchange manually to see the resulting token\n",
|
||||
"obo_resp = req.post(\n",
|
||||
" f'https://login.microsoftonline.com/{tenant_id}/oauth2/v2.0/token',\n",
|
||||
" data={\n",
|
||||
" 'grant_type': 'urn:ietf:params:oauth:grant-type:jwt-bearer',\n",
|
||||
" 'client_id': client_id,\n",
|
||||
" 'client_secret': os.environ['MICROSOFT_CLIENT_SECRET'],\n",
|
||||
" 'assertion': bearer_token,\n",
|
||||
" 'scope': 'https://graph.microsoft.com/Calendars.Read https://graph.microsoft.com/Mail.Read https://graph.microsoft.com/User.Read',\n",
|
||||
" 'requested_token_use': 'on_behalf_of',\n",
|
||||
" }\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"def decode_jwt(token):\n",
|
||||
" payload = token.split('.')[1]\n",
|
||||
" payload += '=' * (4 - len(payload) % 4)\n",
|
||||
" return json.loads(base64.urlsafe_b64decode(payload))\n",
|
||||
"\n",
|
||||
"obo_token = obo_resp.json().get('access_token', '')\n",
|
||||
"if not obo_token:\n",
|
||||
" print(f'OBO exchange failed: {obo_resp.json()}')\n",
|
||||
"else:\n",
|
||||
" user_claims = decode_jwt(bearer_token)\n",
|
||||
" obo_claims = decode_jwt(obo_token)\n",
|
||||
"\n",
|
||||
" print(f'{\"CLAIM\":<20} {\"USER TOKEN (app-scoped)\":<40} {\"OBO TOKEN (Graph-scoped)\"}')\n",
|
||||
" print('=' * 100)\n",
|
||||
" for f in ['aud', 'iss', 'ver', 'scp', 'appid', 'name', 'email', 'idp', 'oid']:\n",
|
||||
" uv = str(user_claims.get(f, '\\u2014'))[:38]\n",
|
||||
" ov = str(obo_claims.get(f, '\\u2014'))[:38]\n",
|
||||
" changed = ' \\u2190 CHANGED' if uv != ov else ''\n",
|
||||
" print(f'{f:<20} {uv:<40} {ov}{changed}')\n",
|
||||
"\n",
|
||||
" # Show the delegation claim\n",
|
||||
" xms_st = obo_claims.get('xms_st', {})\n",
|
||||
" print(f'\\n\\u2705 OBO delegation chain:')\n",
|
||||
" print(f' xms_st.sub (original user): {xms_st.get(\"sub\", \"N/A\")}')\n",
|
||||
" print(f' appid (acting app): {obo_claims.get(\"appid\", \"N/A\")}')\n",
|
||||
" print(f' app_displayname: {obo_claims.get(\"app_displayname\", \"N/A\")}')"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "cleanup-header",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"---\n",
|
||||
"## Cleanup (Optional)\n",
|
||||
"\n",
|
||||
"Uncomment and run to delete resources.\n",
|
||||
"\n",
|
||||
"**Also delete manually:**\n",
|
||||
"- Microsoft Entra ID app registration from the Entra portal\n",
|
||||
"- IAM role (`agentcore-obo-gateway-role`) from IAM Console"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"id": "cleanup",
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# # Delete the Gateway target\n",
|
||||
"# agentcore_control_client.delete_gateway_target(\n",
|
||||
"# gatewayIdentifier=gateway_id,\n",
|
||||
"# targetId=target_id\n",
|
||||
"# )\n",
|
||||
"# print(\"Gateway target deleted \\u2713\")\n",
|
||||
"\n",
|
||||
"# # Delete the Gateway\n",
|
||||
"# agentcore_control_client.delete_gateway(gatewayIdentifier=gateway_id)\n",
|
||||
"# print(\"Gateway deleted \\u2713\")\n",
|
||||
"\n",
|
||||
"# # Delete the credential provider\n",
|
||||
"# identity_client.delete_oauth2_credential_provider(name=\"microsoft-obo-provider\")\n",
|
||||
"# print(\"Credential provider deleted \\u2713\")\n",
|
||||
"\n",
|
||||
"# # Delete the IAM role\n",
|
||||
"# iam_client.delete_role_policy(RoleName=gateway_role_name, PolicyName=\"GatewayOBOPolicy\")\n",
|
||||
"# iam_client.delete_role(RoleName=gateway_role_name)\n",
|
||||
"# print(\"IAM role deleted \\u2713\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"id": "congratulations",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Congratulations! 🎉\n",
|
||||
"\n",
|
||||
"You built a Strands agent that accesses Microsoft Graph API through AgentCore Gateway with OBO token exchange — using **Entra ID directly** for end-to-end authentication. The agent code has **zero token handling logic** — the Gateway handles everything transparently.\n",
|
||||
"\n",
|
||||
"### What you accomplished:\n",
|
||||
"- Registered an app in **Microsoft Entra ID** with OBO-compatible scopes and exposed an API\n",
|
||||
"- Created a **Custom OAuth2 credential provider** with `JWT_AUTHORIZATION_GRANT` OBO config\n",
|
||||
"- Set up an **AgentCore Gateway** with Entra ID v1.0 OIDC inbound auth and OpenAPI schema target\n",
|
||||
"- Configured the target with `TOKEN_EXCHANGE` grant type and `requested_token_use: on_behalf_of` custom parameter\n",
|
||||
"- Authenticated directly with **Entra ID** to get an app-scoped access token\n",
|
||||
"- Connected a **Strands agent** to the Gateway MCP URL and discovered tools automatically\n",
|
||||
"- Queried your **Microsoft profile** through the agent — all auth handled by the Gateway via OBO\n",
|
||||
"\n",
|
||||
"### Key learnings:\n",
|
||||
"- The Gateway's **inbound auth discovery URL must match the token version** (v1.0 by default for Entra ID)\n",
|
||||
"- Do **not** use `allowedClients` in the Gateway authorizer config for Entra ID v1.0 tokens\n",
|
||||
"- The `customParameters` with `requested_token_use: on_behalf_of` is **required** for Entra ID's OBO endpoint\n",
|
||||
"- The IAM role uses the same `bedrock-agentcore:*` pattern as other Gateway tutorials (via shared `utils.create_agentcore_gateway_role`)\n",
|
||||
"- Personal accounts work for `/me` profile; calendar/email endpoints require a work/school account with Exchange Online\n",
|
||||
"\n",
|
||||
"### Next steps:\n",
|
||||
"- Use a **work/school account** to test calendar and email tools\n",
|
||||
"- Add more Microsoft Graph endpoints to the OpenAPI schema (e.g., OneDrive, Teams)\n",
|
||||
"- Try different agent prompts: \"What meetings do I have this week?\", \"List my recent emails\"\n",
|
||||
"- Explore the [OBO Token Exchange docs](https://docs.aws.amazon.com/bedrock-agentcore/latest/devguide/on-behalf-of-token-exchange.html) and [Gateway Outbound Auth docs](https://docs.aws.amazon.com/bedrock-agentcore/latest/devguide/gateway-outbound-auth.html)"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3 (ipykernel)",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"codemirror_mode": {
|
||||
"name": "ipython",
|
||||
"version": 3
|
||||
},
|
||||
"file_extension": ".py",
|
||||
"mimetype": "text/x-python",
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.14.4"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
||||
@@ -0,0 +1,9 @@
|
||||
strands-agents
|
||||
strands-agents-tools
|
||||
bedrock-agentcore
|
||||
boto3
|
||||
botocore
|
||||
python-dotenv
|
||||
requests
|
||||
uvicorn
|
||||
fastapi
|
||||
+203
@@ -0,0 +1,203 @@
|
||||
"""
|
||||
Environment-aware OAuth2 token callback server for Entra ID.
|
||||
|
||||
Captures the authorization code from Entra ID and exchanges it for tokens.
|
||||
Automatically detects local vs SageMaker Workshop Studio environments.
|
||||
|
||||
Usage:
|
||||
python3 token_callback_server.py <tenant-id> <client-id> <client-secret>
|
||||
"""
|
||||
|
||||
import time
|
||||
import json
|
||||
import argparse
|
||||
import logging
|
||||
import socket
|
||||
|
||||
import uvicorn
|
||||
import requests
|
||||
from fastapi import FastAPI, HTTPException
|
||||
from fastapi.responses import HTMLResponse
|
||||
|
||||
PORT = 9090
|
||||
CALLBACK_ENDPOINT = "/oauth2/callback"
|
||||
PING_ENDPOINT = "/ping"
|
||||
TOKEN_ENDPOINT = "/token"
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Global token storage
|
||||
_captured_token = None
|
||||
|
||||
|
||||
def _is_workshop_studio() -> bool:
|
||||
try:
|
||||
with open("/opt/ml/metadata/resource-metadata.json", "r") as f:
|
||||
json.load(f)
|
||||
return True
|
||||
except (FileNotFoundError, json.JSONDecodeError):
|
||||
return False
|
||||
|
||||
|
||||
def get_callback_base_url() -> str:
|
||||
"""Get the browser-accessible base URL (environment-aware)."""
|
||||
if not _is_workshop_studio():
|
||||
return f"http://localhost:{PORT}"
|
||||
try:
|
||||
import boto3
|
||||
|
||||
with open("/opt/ml/metadata/resource-metadata.json", "r") as f:
|
||||
data = json.load(f)
|
||||
client = boto3.client("sagemaker")
|
||||
resp = client.describe_space(
|
||||
DomainId=data["DomainId"], SpaceName=data["SpaceName"]
|
||||
)
|
||||
return resp["Url"] + f"/proxy/{PORT}"
|
||||
except Exception:
|
||||
return f"http://localhost:{PORT}"
|
||||
|
||||
|
||||
def get_callback_url() -> str:
|
||||
return f"{get_callback_base_url()}{CALLBACK_ENDPOINT}"
|
||||
|
||||
|
||||
def is_server_running() -> bool:
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
||||
return s.connect_ex(("127.0.0.1", PORT)) == 0
|
||||
|
||||
|
||||
def get_captured_token() -> str:
|
||||
"""Retrieve the captured token from a running server."""
|
||||
try:
|
||||
r = requests.get(f"http://localhost:{PORT}{TOKEN_ENDPOINT}", timeout=2)
|
||||
if r.status_code == 200:
|
||||
return r.json().get("access_token", "")
|
||||
except Exception:
|
||||
pass
|
||||
return ""
|
||||
|
||||
|
||||
def wait_for_token(timeout=120) -> str:
|
||||
"""Poll the server until a token is captured."""
|
||||
start = time.time()
|
||||
while time.time() - start < timeout:
|
||||
token = get_captured_token()
|
||||
if token:
|
||||
return token
|
||||
time.sleep(2)
|
||||
return ""
|
||||
|
||||
|
||||
def wait_for_server_ready(timeout=30) -> bool:
|
||||
start = time.time()
|
||||
while time.time() - start < timeout:
|
||||
try:
|
||||
r = requests.get(f"http://localhost:{PORT}{PING_ENDPOINT}", timeout=2)
|
||||
if r.status_code == 200:
|
||||
return True
|
||||
except Exception:
|
||||
pass
|
||||
time.sleep(1)
|
||||
return False
|
||||
|
||||
|
||||
class TokenCallbackServer:
|
||||
def __init__(self, tenant_id: str, client_id: str, client_secret: str):
|
||||
self.tenant_id = tenant_id
|
||||
self.client_id = client_id
|
||||
self.client_secret = client_secret
|
||||
self.app = FastAPI()
|
||||
self._setup_routes()
|
||||
|
||||
def _setup_routes(self):
|
||||
@self.app.get(PING_ENDPOINT)
|
||||
async def ping():
|
||||
return {"status": "ok"}
|
||||
|
||||
@self.app.get(TOKEN_ENDPOINT)
|
||||
async def get_token():
|
||||
return {"access_token": _captured_token or ""}
|
||||
|
||||
@self.app.get(CALLBACK_ENDPOINT)
|
||||
async def callback(
|
||||
code: str = None, error: str = None, error_description: str = None
|
||||
):
|
||||
global _captured_token
|
||||
|
||||
if error:
|
||||
from html import escape
|
||||
|
||||
return HTMLResponse(
|
||||
f"<h2>Error: {escape(error)}</h2><p>{escape(error_description or '')}</p>",
|
||||
status_code=400,
|
||||
)
|
||||
|
||||
if not code:
|
||||
raise HTTPException(status_code=400, detail="Missing code parameter")
|
||||
|
||||
redirect_uri = get_callback_url()
|
||||
token_url = (
|
||||
f"https://login.microsoftonline.com/{self.tenant_id}/oauth2/v2.0/token"
|
||||
)
|
||||
r = requests.post(
|
||||
token_url,
|
||||
data={
|
||||
"grant_type": "authorization_code",
|
||||
"code": code,
|
||||
"redirect_uri": redirect_uri,
|
||||
"client_id": self.client_id,
|
||||
"client_secret": self.client_secret,
|
||||
"scope": f"api://{self.client_id}/access_as_user openid profile email",
|
||||
},
|
||||
)
|
||||
tokens = r.json()
|
||||
|
||||
if "error" in tokens:
|
||||
from html import escape
|
||||
|
||||
return HTMLResponse(
|
||||
f"<h2>Token exchange error: {escape(tokens['error'])}</h2>"
|
||||
f"<p>{escape(tokens.get('error_description', ''))}</p>",
|
||||
status_code=400,
|
||||
)
|
||||
|
||||
_captured_token = tokens.get("access_token", "")
|
||||
print(f"\n{'=' * 60}")
|
||||
print("TOKEN RECEIVED")
|
||||
print(f"{'=' * 60}")
|
||||
print(f"\nFULL ACCESS TOKEN:\n{_captured_token}")
|
||||
print(f"\n{'=' * 60}")
|
||||
return HTMLResponse(
|
||||
"<h2>✅ Token captured! Return to the notebook.</h2>", status_code=200
|
||||
)
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("tenant_id")
|
||||
parser.add_argument("client_id")
|
||||
parser.add_argument("client_secret")
|
||||
args = parser.parse_args()
|
||||
|
||||
server = TokenCallbackServer(args.tenant_id, args.client_id, args.client_secret)
|
||||
host = "0.0.0.0" if _is_workshop_studio() else "127.0.0.1"
|
||||
callback_url = get_callback_url()
|
||||
|
||||
authorize_url = (
|
||||
f"https://login.microsoftonline.com/{args.tenant_id}/oauth2/v2.0/authorize?"
|
||||
f"client_id={args.client_id}&response_type=code&"
|
||||
f"redirect_uri={__import__('urllib.parse', fromlist=['quote']).quote(callback_url)}&"
|
||||
f"scope={__import__('urllib.parse', fromlist=['quote']).quote(f'api://{args.client_id}/access_as_user openid profile email')}"
|
||||
)
|
||||
|
||||
print(f"\n🚀 Token callback server on {host}:{PORT}")
|
||||
print(f"📋 Callback URL: {callback_url}")
|
||||
print("\n📋 Open this URL in your browser to sign in:\n")
|
||||
print(authorize_url)
|
||||
print("\nWaiting for sign-in...\n")
|
||||
|
||||
uvicorn.run(server.app, host=host, port=PORT, log_level="warning")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user