diff --git a/docs/frds/0004-obo-token-header-contract.md b/docs/frds/0004-obo-token-header-contract.md new file mode 100644 index 00000000..2692244f --- /dev/null +++ b/docs/frds/0004-obo-token-header-contract.md @@ -0,0 +1,151 @@ +--- +frd: 0003 +title: OBO token and header pass-through contract +status: Finalized +author: victoriahall +created: 2026-07-06 +updated: 2026-07-06 +issues: [] +pull_requests: [] +branch: victoriahall/obo-token-header-contract +--- + +# FRD 0003 - OBO token and header pass-through contract + +## 1. Summary + +Define the runtime contract for passing authenticated user identity from inbound +HTTP requests to downstream MCP server calls using OAuth 2.0 On-Behalf-Of +(OBO). The contract covers input headers, config shape, context propagation, +outbound auth headers, fallback behavior, and error responses. + +## 2. Motivation / problem + +Before this contract, agents could authenticate downstream calls only with +application credentials (managed identity). Web app scenarios require delegated +authorization so tool calls reflect end-user permissions. + +Without a clear contract, client apps and runtime modules can disagree on: + +1. Which inbound header carries the user token. +2. How token context is propagated across async boundaries. +3. What outbound authorization behavior is expected for OBO versus fallback. +4. How consent or claims challenges are surfaced to clients. + +## 3. Goals / Non-goals + +**Goals** +- Define canonical inbound token/header extraction behavior. +- Define runtime configuration contract for OBO enablement. +- Define MCP auth contract for OBO-protected downstream services. +- Define deterministic fallback and error behavior. +- Keep behavior backward compatible for apps that do not enable OBO. + +**Non-goals** +- Browser interactive sign-in orchestration. +- Runtime-managed consent UX. +- Certificate credential support for OBO client auth in this revision. +- Changes to trigger model or non-HTTP identity sources. + +## 4. Proposed design + +| Pipeline stage | Module(s) | Change | +| --- | --- | --- | +| discover | `src/azure_functions_agents/discovery/mcp.py` | Add OBO-aware MCP header provider and request-scoped user context access | +| translate | `src/azure_functions_agents/config/schema.py` | Add global auth schema with OBO settings | +| register | `src/azure_functions_agents/registration/_handlers.py`, `src/azure_functions_agents/registration/endpoints.py` | Extract inbound user token/identity and map interaction-required errors to HTTP 401 | +| execute | `src/azure_functions_agents/runner.py`, `src/azure_functions_agents/_obo.py` | Thread user context through run paths and perform OBO token exchange with caching | + +### Authoring / API surface + +Global config in agents.config.yaml: + +- `auth.obo.enabled` boolean. +- `auth.obo.client_id` string. +- `auth.obo.client_secret` string. +- `auth.obo.tenant_id` string. +- `auth.obo.downstream_scopes` map of scope aliases to scope URIs. + +MCP server auth in mcp.json: + +- `auth.type: obo` enables OBO header provider for that server. +- `auth.scope` is required and identifies the downstream resource scope. + +Inbound request header contract: + +1. Primary token source: `X-MS-TOKEN-AAD-ACCESS-TOKEN`. +2. Secondary token source: `Authorization: Bearer `. +3. User id source: `X-MS-CLIENT-PRINCIPAL-ID` when present. + +Outbound MCP header contract: + +1. Preserve configured static headers from mcp.json. +2. Add `Authorization: Bearer ` where token comes from: + - OBO exchange when user context and OBO config are available. + - Managed identity fallback when OBO is unavailable or fails. + +### Runtime behavior contract + +1. Request handlers create `UserContext` from inbound headers. +2. Runner stores context in a request-scoped context variable for both + non-streaming and streaming agent execution. +3. MCP header provider reads current user context at request time. +4. OBO token provider caches tokens in-memory by token hash and scope until + near expiry. +5. On OBO interaction-required conditions, handlers return HTTP 401 and include + `WWW-Authenticate` with error metadata and claims challenge. + +### Error and claims challenge contract + +For interaction-required failures (`interaction_required`, `consent_required`, +`login_required`): + +- HTTP status: 401. +- Response body: JSON with `error`, `error_description`, and optional `claims`. +- Response header: `WWW-Authenticate` containing Bearer error metadata. +- Claims in `WWW-Authenticate` are base64-encoded. + +### Compatibility + +- Backward compatible by default: OBO is inactive unless `auth.obo.enabled` is + true and MCP server auth type is explicitly set to obo. +- Existing managed identity behavior remains default for non-OBO MCP entries. +- Missing or malformed user token does not break request handling; runtime can + continue via managed identity fallback. + +## 5. Decisions log + +| # | Decision | Options considered | Choice | Decided by | Date | +| - | -------- | ------------------ | ------ | ---------- | ---- | +| 1 | Inbound token precedence | Authorization first / EasyAuth first | EasyAuth first, then Authorization | Human | 2026-06-19 | +| 2 | Context propagation mechanism | Explicit argument threading only / context variable | Context variable with reset token | Human | 2026-06-19 | +| 3 | Token cache strategy | No cache / distributed cache / in-memory cache | In-memory cache keyed by token hash and scope | Human | 2026-06-19 | +| 4 | Fallback mode | Fail closed when OBO unavailable / managed identity fallback | Managed identity fallback | Human | 2026-06-19 | +| 5 | Claims challenge behavior | 500 with internal error / 401 with challenge details | 401 with WWW-Authenticate and claims | Human | 2026-06-19 | +| 6 | OBO scope location | Global single scope / per-MCP auth scope | Per-MCP `auth.scope` with optional global alias map | Human | 2026-06-19 | + +## 6. Test plan + +- [x] Unit: OBO config validation and defaults in `tests/test_obo.py`. +- [x] Unit: Header extraction precedence and case-insensitive lookup in + `tests/test_obo.py`. +- [x] Unit: OBO provider success, interaction-required, and generic error paths + in `tests/test_obo.py`. +- [x] Unit: Token caching behavior in `tests/test_obo.py`. +- [x] Unit: Global provider lifecycle tests in `tests/test_obo.py`. +- [x] Integration-adjacent: Runner streaming tests pass with context threading + changes in `tests/test_runner_streaming.py`. + +## 7. Docs impact + +- [ ] `docs/architecture.md` - add explicit OBO dataflow notes in module map. +- [ ] `docs/front-matter-spec.md` - no expected changes. +- [ ] `docs/triggers.md` - no expected changes. +- [ ] `README.md` - add OBO configuration and mcp.json examples. + +## 8. Status & sign-off + +- **Architecture review (phase 2):** Contract aligns with implemented runtime + boundaries (discovery -> registration -> execute) and preserves backward + compatibility. +- **Human sign-off:** Victoria Hall, 2026-07-06. diff --git a/docs/frds/README.md b/docs/frds/README.md index 53af1f94..6fe2658c 100644 --- a/docs/frds/README.md +++ b/docs/frds/README.md @@ -32,6 +32,7 @@ The full lifecycle that produces an FRD lives in [`../../AGENTS.md`](../../AGENT | [0001](0001-agents-folder-indexing.md) | agents/ folder indexing | Finalized | | [0002](0002-skill-includes.md) | Skill file includes | Finalized | | [0003](0003-runtime-observability.md) | Runtime-owned observability (OpenTelemetry) | Finalized | +| [0004](0004-obo-token-header-contract.md) | OBO token and header pass-through contract | Finalized | > `_template.md` is the template, not an FRD — the leading underscore keeps it > sorted first and excludes it from numbering. diff --git a/docs/user-identity-propagation.md b/docs/user-identity-propagation.md new file mode 100644 index 00000000..88d27df3 --- /dev/null +++ b/docs/user-identity-propagation.md @@ -0,0 +1,247 @@ +# User identity propagation + +This document describes how the runtime carries end-user identity through the +request lifecycle and forwards it to downstream MCP servers. + +--- + +## Overview + +The runtime acts as a **transparent identity proxy**: it never validates the +incoming user token itself. Instead it relies on EasyAuth to authenticate the +session before the request reaches Python code, and it forwards the +already-authenticated identity artifacts to downstream MCP servers, adding a +managed identity token so the downstream can verify *which* app is calling. + +Two authentication modes are supported depending on what is present in the +inbound request headers: + +| Mode | Trigger condition | Description | +|---|---|---| +| **BigMac hook-session** | Both `X-MS-Access-Token` and `X-MS-Hooks-Session-Token` present | User identity forwarded unchanged; `Authorization` is the function app's managed identity token. | +| **OBO (On-Behalf-Of)** | `X-MS-Access-Token` (or `X-MS-TOKEN-AAD-ACCESS-TOKEN`) present, no hooks session token | MSAL exchanges the inbound token for a new downstream-scoped token. | +| **Managed identity fallback** | No user token present | The function app's own managed identity is used for downstream calls. | + +--- + +## End-to-end flow + +### BigMac hook-session mode (production path) + +``` +┌───────────────────┐ +│ Browser / Client │ +└────────┬──────────┘ + │ POST /agents//chat + │ X-MS-Access-Token: + │ X-MS-Hooks-Session-Token: + ▼ +┌────────────────────────────────────────────────────────────────────┐ +│ EasyAuth (App Service Authentication) │ +│ │ +│ • Validates X-MS-Hooks-Session-Token (opaque session managed by │ +│ EasyAuth). │ +│ • If X-MS-Access-Token is expired, refreshes it automatically │ +│ via the EasyAuth /.auth/refresh endpoint before passing the │ +│ request on. │ +│ • Injects or preserves the two headers for downstream code. │ +└────────────────────────────┬───────────────────────────────────────┘ + │ (same headers, now validated by EasyAuth) + ▼ +┌────────────────────────────────────────────────────────────────────┐ +│ Azure Functions Agent Runtime │ +│ │ +│ 1. _build_user_context_from_request() │ +│ • extract_user_token_from_headers() │ +│ Priority: X-MS-Access-Token > X-MS-TOKEN-AAD-ACCESS-TOKEN │ +│ > X-MS-TOKEN-AAD-ID-TOKEN > Authorization: Bearer │ +│ • extract_hooks_session_token_from_headers() │ +│ • Both values stored in UserContext (no validation performed). │ +│ │ +│ 2. runner.run_agent() / run_agent_stream() │ +│ • Sets UserContext in a contextvar for the duration of the run.│ +│ • Passes UserContext to all MCP tool calls. │ +│ │ +│ 3. discovery/mcp.py — obo_header_provider() │ +│ • Detects: hooks_session_token IS set AND access_token IS set. │ +│ • Acquires a fresh managed identity (MI) token for the │ +│ configured scope (e.g. https://graph.microsoft.com/.default).│ +│ • Builds the outbound header set: │ +│ Authorization: Bearer (new) │ +│ X-MS-Access-Token: │ +│ X-MS-Hooks-Session-Token: │ +└────────────────────────────┬───────────────────────────────────────┘ + │ HTTP request to MCP server + │ Authorization: Bearer + │ X-MS-Access-Token: + │ X-MS-Hooks-Session-Token: + ▼ +┌────────────────────────────────────────────────────────────────────┐ +│ Downstream MCP Server │ +│ │ +│ The downstream server decides what to validate. Typical checks: │ +│ • Authorization — verify the MI token; the azp claim identifies │ +│ which function app is the caller. │ +│ • X-MS-Access-Token — decode the user's identity (oid, upn, tid, │ +│ scp) if the server wants to act on behalf of the user. │ +│ • X-MS-Hooks-Session-Token — can be used to call EasyAuth's │ +│ /.auth/refresh if the access token needs refreshing. │ +└────────────────────────────────────────────────────────────────────┘ +``` + +### OBO (On-Behalf-Of) mode + +Activated when an inbound access token is present but **no hooks session token** +is included. The runtime exchanges the user token for a new token scoped to the +downstream API using MSAL and the configured `auth.obo` credentials. + +``` +┌───────────────────┐ +│ Browser / Client │ +└────────┬──────────┘ + │ X-MS-Access-Token: + ▼ +┌────────────────────────────────────────────────────────────────────┐ +│ EasyAuth (validates session) │ +└────────────────────────────┬───────────────────────────────────────┘ + ▼ +┌────────────────────────────────────────────────────────────────────┐ +│ Azure Functions Agent Runtime │ +│ │ +│ 1. UserContext built with access_token only (no hooks token). │ +│ │ +│ 2. obo_header_provider() — OBO branch: │ +│ • hooks_session_token is None → OBO path selected. │ +│ • OboTokenProvider.acquire_token_on_behalf_of() calls MSAL. │ +│ • MSAL exchanges user token for a downstream-scoped token. │ +│ • Token cached in-process (keyed by token hash + scope). │ +│ • If exchange fails with interaction_required, the request │ +│ returns HTTP 401 with WWW-Authenticate containing the claims │ +│ challenge so the client can re-authenticate. │ +│ • Outbound header: │ +│ Authorization: Bearer │ +└────────────────────────────┬───────────────────────────────────────┘ + │ Authorization: Bearer + ▼ +┌────────────────────────────┐ +│ Downstream MCP Server │ +└────────────────────────────┘ +``` + +### Managed identity fallback + +Activated when no user identity headers are present at all (e.g. a timer +trigger or a background request without EasyAuth). + +``` +┌───────────────────────────────────────────┐ +│ Azure Functions Agent Runtime │ +│ │ +│ UserContext.access_token = None │ +│ obo_header_provider() — fallback branch: │ +│ • build_credential() → DefaultAzure... │ +│ (Managed Identity in Azure; az CLI │ +│ or env vars locally) │ +│ • Token acquired for configured scope. │ +│ • Outbound header: │ +│ Authorization: Bearer │ +└───────────────────────┬───────────────────┘ + ▼ +┌──────────────────────────┐ +│ Downstream MCP Server │ +└──────────────────────────┘ +``` + +--- + +## Trust boundaries + +| Boundary | Who is responsible | What the runtime does | +|---|---|---| +| Inbound token authenticity | **EasyAuth** | Trusts the headers as-is; no JWT decode or signature check | +| Outbound caller identity | **Runtime** (via MI token) | Mints a fresh MI token; downstream can verify `azp` claim | +| Downstream access control | **Downstream MCP server** | Runtime makes no assertions about what the downstream should accept | + +**Security implication:** without EasyAuth in front of the function app, the +inbound `X-MS-Access-Token` and `X-MS-Hooks-Session-Token` headers are +completely unauthenticated. EasyAuth is a **hard requirement** in production for +this model to be secure. + +--- + +## Header extraction priority + +The runtime checks inbound headers in this order when looking for a user access +token (implemented in `_obo.py: extract_user_token_from_headers`): + +1. `X-MS-Access-Token` — BigMac explicit access token header +2. `X-MS-TOKEN-AAD-ACCESS-TOKEN` — EasyAuth AAD access token header +3. `X-MS-TOKEN-AAD-ID-TOKEN` — EasyAuth AAD ID token (fallback when access token is absent) +4. `Authorization: Bearer ` — standard bearer token + +The hooks session token is extracted separately from `X-MS-Hooks-Session-Token` +(case-insensitive lookup). + +--- + +## Configuration + +### BigMac / managed identity fallback + +No special configuration is required. The runtime uses `DefaultAzureCredential` +for the outbound MI token. In Azure this resolves to the function app's managed +identity; locally it falls back to `az login` credentials. + +The MCP server entry in `mcp.json` must set `auth.type: obo` to enable the +identity-forwarding path: + +```json +{ + "servers": { + "my-api": { + "url": "https://my-api.example.com/mcp", + "auth": { + "type": "obo", + "scope": "https://graph.microsoft.com/.default" + } + } + } +} +``` + +The `scope` value is used when acquiring the outbound MI token (fallback path) +or the OBO downstream token. + +### OBO mode + +Requires `auth.obo` in `agents.config.yaml` with a valid Entra app registration +that has been pre-consented for the downstream scopes: + +```yaml +auth: + obo: + enabled: true + client_id: $AZURE_CLIENT_ID + client_secret: $AZURE_CLIENT_SECRET + tenant_id: $AZURE_TENANT_ID + downstream_scopes: + my_api: "api:///.default" +``` + +The `client_id` / `client_secret` / `tenant_id` must be in the **same tenant** +as the downstream API. If they are in a different tenant from the Azure +subscription hosting the function app, MSAL will fail with a cross-tenant +identity error. + +--- + +## Runtime implementation pointers + +| Concern | Module | Symbol | +|---|---|---| +| Header extraction | `_obo.py` | `extract_user_token_from_headers()`, `extract_hooks_session_token_from_headers()` | +| UserContext construction | `registration/_handlers.py`, `registration/endpoints.py` | `_build_user_context_from_request()` | +| Context propagation into tools | `discovery/mcp.py` | `set_current_user_context()`, `_current_user_context` contextvar | +| BigMac / OBO / MI header dispatch | `discovery/mcp.py` | `_build_obo_header_provider()` → `obo_header_provider()` closure | +| OBO token exchange and caching | `_obo.py` | `OboTokenProvider.acquire_token_on_behalf_of()`, `_token_cache` | +| Interaction-required handling | `registration/endpoints.py` | `_interaction_required_error()` | diff --git a/pyproject.toml b/pyproject.toml index bbc1380c..be0abdd9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -22,6 +22,7 @@ dependencies = [ "azure-storage-blob==12.28.*", "jsonschema==4.26.*", "mcp==1.27.*", + "msal==1.31.*", ] [project.optional-dependencies] @@ -40,6 +41,7 @@ dev = [ "mypy==2.1.*", "pydantic==2.13.*", "pytest==9.0.*", + "pytest-asyncio==1.4.*", "types-jsonschema==4.26.*", "pytest-cov==7.1.*", # Bring the exporter into the dev/test/type-check env so the observability tests (which import @@ -57,6 +59,10 @@ version = {attr = "azure_functions_agents.__version__"} [tool.setuptools.package-data] azure_functions_agents = ["public/**"] +[tool.pytest.ini_options] +asyncio_mode = "auto" +asyncio_default_fixture_loop_scope = "function" + [tool.ruff] target-version = "py313" line-length = 100 diff --git a/samples/README.md b/samples/README.md index 604468c9..b9c17147 100644 --- a/samples/README.md +++ b/samples/README.md @@ -10,6 +10,8 @@ Each subdirectory is a standalone Azure Functions app deployable with [`azd up`] | [daily-tech-news-email](daily-tech-news-email/) | Timer | | ✅ Office 365 Outlook | ✅ Office 365 Outlook | | ✅ | | | [daily-azure-report](daily-azure-report/) | Timer + HTTP | ✅ azure_rest | ✅ Office 365 Outlook | ✅ MS Learn + Office 365 Outlook | ✅ azure-resources | | ✅ | | [skill-includes-demo](skill-includes-demo/) | HTTP | | | | ✅ (with includes) | | ✅ | +| [obo-e2e](obo-e2e/) | HTTP | | | ✅ (custom OBO target) | | | | +| [obo-whoami-mcp-server](obo-whoami-mcp-server/) | MCP (streamable HTTP) | | | | | | | ## Run Locally (optional) diff --git a/samples/obo-e2e/README.md b/samples/obo-e2e/README.md new file mode 100644 index 00000000..4f0e720c --- /dev/null +++ b/samples/obo-e2e/README.md @@ -0,0 +1,180 @@ +# OBO E2E Test Kit + +This sample validates On-Behalf-Of (OBO) token/header pass-through behavior in +Azure Functions Agents Runtime. + +It verifies four paths: + +1. User A request -> downstream sees User A identity. +2. User B request -> downstream sees User B identity. +3. No user token -> managed identity fallback. +4. Missing consent/MFA -> HTTP 401 with `WWW-Authenticate` claims challenge. + +For BigMac hook-session callback validation, requests should include: + +- `X-MS-Access-Token` +- `X-MS-Hooks-Session-Token` + +If access token is unavailable, `X-MS-TOKEN-AAD-ID-TOKEN` can be used as fallback. + +## Important clarification + +You do **not** need to deploy to production to test this. + +You can run this scenario in either: + +1. **Local dev** (`func start`) against test Entra app registrations and test + downstream APIs. +2. **Non-production Azure environment** (recommended for team validation) + such as dev/test subscription or staging slot. + +Production deployment is optional and should only happen after validation in +one of the environments above. + +## Included sample app files + +This folder now includes a runnable sample app under `src/`: + +- `src/function_app.py` +- `src/main.agent.md` +- `src/agents.config.yaml` +- `src/mcp.json` +- `src/host.json` +- `src/local.settings.template.json` +- `src/requirements.txt` + +The root templates are still provided for quick copy/reference: + +- `agents.config.obo.sample.yaml` +- `mcp.obo.sample.json` +- `test.http` +- helper scripts (`get-user-token.ps1`, `decode-jwt-payload.ps1`) + +## Prerequisites + +- Azure Functions Core Tools (for local run) +- Python 3.13+ +- Runtime configured with OBO (see `src/agents.config.yaml`) +- MCP server configured with `auth.type: obo` (see `src/mcp.json`) +- A downstream MCP tool (for example `whoami`) that returns bearer token claims + (`oid`, `sub`, `aud`, and `azp/appid`) +- Azure CLI logged in to test users (User A and User B) + +## Quick local setup + +From repo root: + +```powershell +cd samples/obo-e2e/src +python -m venv .venv +.venv\Scripts\Activate.ps1 +pip install -r requirements.txt +Copy-Item local.settings.template.json local.settings.json +``` + +Set values in `local.settings.json`: + +- `FOUNDRY_PROJECT_ENDPOINT` +- `FOUNDRY_MODEL` +- `AZURE_CLIENT_ID` +- `AZURE_CLIENT_SECRET` +- `AZURE_TENANT_ID` + +Update `mcp.json` values: + +- `url`: `https:///mcp` +- `scope` (default first-run option): `https://graph.microsoft.com/.default` + +What `` should be: + +- The host name of a downstream MCP server that exposes a `whoami`-style tool. +- Example with the included downstream sample: + - Run `samples/obo-whoami-mcp-server` + - Use `http://localhost:8000/mcp` for local tests + - Or `https:///mcp` after non-prod deployment + +Scope guidance: + +- Start with `https://graph.microsoft.com/.default` for the easiest OBO smoke + test path. +- Switch to `api:///.default` when validating strict + custom downstream audience behavior. + +Then start the host: + +```powershell +func start +``` + +## Copy-paste sequence (PowerShell) + +Run these commands from `samples/obo-e2e`. + +```powershell +# 0) Set these values for your environment +$FunctionBaseUrl = "https://.azurewebsites.net" +$AgentName = "main" + +# 1) Acquire User A token (for the function app audience) +# If you use EasyAuth in front of the app, use your app's audience scope. +$FunctionAudienceScope = "api:///.default" +$UserAToken = az account get-access-token --scope $FunctionAudienceScope --query accessToken -o tsv + +# 2) Call chat endpoint as User A +$Body = @{ prompt = "Call the downstream whoami MCP tool and return raw JSON" } | ConvertTo-Json +$RespA = Invoke-WebRequest -Method Post ` + -Uri "$FunctionBaseUrl/agents/$AgentName/chat" ` + -Headers @{ "X-MS-Access-Token" = $UserAToken; "X-MS-Hooks-Session-Token" = "" } ` + -ContentType "application/json" ` + -Body $Body + +$RespA.StatusCode +$RespA.Content + +# 3) Sign in as User B, then acquire User B token +# az login +$UserBToken = az account get-access-token --scope $FunctionAudienceScope --query accessToken -o tsv + +# 4) Call chat endpoint as User B +$RespB = Invoke-WebRequest -Method Post ` + -Uri "$FunctionBaseUrl/agents/$AgentName/chat" ` + -Headers @{ "X-MS-Access-Token" = $UserBToken; "X-MS-Hooks-Session-Token" = "" } ` + -ContentType "application/json" ` + -Body $Body + +$RespB.StatusCode +$RespB.Content + +# 5) Fallback check (no user token) +$RespFallback = Invoke-WebRequest -Method Post ` + -Uri "$FunctionBaseUrl/agents/$AgentName/chat" ` + -ContentType "application/json" ` + -Body $Body + +$RespFallback.StatusCode +$RespFallback.Content + +# 6) Optional: claims challenge check +# Use a scope requiring consent/MFA and inspect 401 + WWW-Authenticate. +try { + Invoke-WebRequest -Method Post ` + -Uri "$FunctionBaseUrl/agents/$AgentName/chat" ` + -Headers @{ "X-MS-Access-Token" = $UserAToken; "X-MS-Hooks-Session-Token" = "" } ` + -ContentType "application/json" ` + -Body (@{ prompt = "Call MCP tool for protected scope" } | ConvertTo-Json) +} catch { + $_.Exception.Response.StatusCode.value__ + $_.Exception.Response.Headers["WWW-Authenticate"] +} +``` + +## What to validate + +- User A response shows downstream `oid/sub` for User A. +- User B response shows downstream `oid/sub` for User B. +- Fallback response shows app identity (managed identity/service principal). +- Claims challenge path returns HTTP 401 and `WWW-Authenticate`. + +## Local REST client option + +You can also use [test.http](test.http) and paste User A/User B tokens manually. diff --git a/samples/obo-e2e/agents.config.obo.sample.yaml b/samples/obo-e2e/agents.config.obo.sample.yaml new file mode 100644 index 00000000..a216d461 --- /dev/null +++ b/samples/obo-e2e/agents.config.obo.sample.yaml @@ -0,0 +1,8 @@ +auth: + obo: + enabled: true + client_id: $AZURE_CLIENT_ID + client_secret: $AZURE_CLIENT_SECRET + tenant_id: $AZURE_TENANT_ID + downstream_scopes: + whoami_api: "api:///.default" diff --git a/samples/obo-e2e/decode-jwt-payload.ps1 b/samples/obo-e2e/decode-jwt-payload.ps1 new file mode 100644 index 00000000..aababe9d --- /dev/null +++ b/samples/obo-e2e/decode-jwt-payload.ps1 @@ -0,0 +1,19 @@ +param( + [Parameter(Mandatory = $true)] + [string]$Jwt +) + +$parts = $Jwt.Split('.') +if ($parts.Length -lt 2) { + throw "Invalid JWT format." +} + +$payload = $parts[1].Replace('-', '+').Replace('_', '/') +switch ($payload.Length % 4) { + 2 { $payload += '==' } + 3 { $payload += '=' } +} + +$bytes = [System.Convert]::FromBase64String($payload) +$json = [System.Text.Encoding]::UTF8.GetString($bytes) +$json | ConvertFrom-Json diff --git a/samples/obo-e2e/get-user-token.ps1 b/samples/obo-e2e/get-user-token.ps1 new file mode 100644 index 00000000..05bacf0d --- /dev/null +++ b/samples/obo-e2e/get-user-token.ps1 @@ -0,0 +1,11 @@ +param( + [Parameter(Mandatory = $true)] + [string]$Scope +) + +$token = az account get-access-token --scope $Scope --query accessToken -o tsv +if (-not $token) { + throw "Failed to acquire token. Ensure 'az login' is complete and scope is valid." +} + +$token diff --git a/samples/obo-e2e/mcp.obo.sample.json b/samples/obo-e2e/mcp.obo.sample.json new file mode 100644 index 00000000..8a10d6d5 --- /dev/null +++ b/samples/obo-e2e/mcp.obo.sample.json @@ -0,0 +1,14 @@ +{ + "servers": { + "whoami-api": { + "url": "https:///mcp", + "auth": { + "type": "obo", + "scope": "api:///.default" + }, + "headers": { + "x-client": "azure-functions-agents-runtime" + } + } + } +} diff --git a/samples/obo-e2e/src/agents.config.yaml b/samples/obo-e2e/src/agents.config.yaml new file mode 100644 index 00000000..c90167db --- /dev/null +++ b/samples/obo-e2e/src/agents.config.yaml @@ -0,0 +1,15 @@ +# Global configuration for OBO end-to-end validation + +# auth: (disabled for local CLI-auth testing - re-enable with valid client_id/secret/tenant for OBO flow) +# auth: +# obo: +# enabled: true +# client_id: $AZURE_CLIENT_ID +# client_secret: $AZURE_CLIENT_SECRET +# tenant_id: $AZURE_TENANT_ID +# downstream_scopes: +# whoami_api: "api:///.default" + +# Default runtime model configuration +model: $FOUNDRY_MODEL +timeout: 120 diff --git a/samples/obo-e2e/src/function_app.py b/samples/obo-e2e/src/function_app.py new file mode 100644 index 00000000..736ad492 --- /dev/null +++ b/samples/obo-e2e/src/function_app.py @@ -0,0 +1,3 @@ +from azure_functions_agents import create_function_app + +app = create_function_app() diff --git a/samples/obo-e2e/src/host.json b/samples/obo-e2e/src/host.json new file mode 100644 index 00000000..f4cef1f8 --- /dev/null +++ b/samples/obo-e2e/src/host.json @@ -0,0 +1,18 @@ +{ + "version": "2.0", + "extensions": { + "http": { + "routePrefix": "" + } + }, + "logging": { + "logLevel": { + "default": "Information", + "azure.functions.AgentRuntime": "Information" + } + }, + "extensionBundle": { + "id": "Microsoft.Azure.Functions.ExtensionBundle", + "version": "[4.*, 5.0.0)" + } +} diff --git a/samples/obo-e2e/src/main.agent.md b/samples/obo-e2e/src/main.agent.md new file mode 100644 index 00000000..7884afbe --- /dev/null +++ b/samples/obo-e2e/src/main.agent.md @@ -0,0 +1,13 @@ +--- +name: OBO E2E Assistant +description: Validates OBO token pass-through to downstream MCP tools. +builtin_endpoints: true +--- + +You are a validation assistant for OBO scenarios. + +When asked to verify identity, call the downstream `whoami` MCP tool and return +the raw JSON response unchanged. + +Do not summarize or reinterpret token claims. Return exact values so callers can +compare `oid`, `sub`, `aud`, `azp`, and `appid` across runs. diff --git a/samples/obo-e2e/src/mcp.json b/samples/obo-e2e/src/mcp.json new file mode 100644 index 00000000..79602a03 --- /dev/null +++ b/samples/obo-e2e/src/mcp.json @@ -0,0 +1,14 @@ +{ + "servers": { + "whoami-api": { + "url": "http://localhost:8000/mcp", + "auth": { + "type": "obo", + "scope": "https://graph.microsoft.com/.default" + }, + "headers": { + "x-client": "obo-e2e-sample" + } + } + } +} diff --git a/samples/obo-e2e/src/requirements.txt b/samples/obo-e2e/src/requirements.txt new file mode 100644 index 00000000..ee26a381 --- /dev/null +++ b/samples/obo-e2e/src/requirements.txt @@ -0,0 +1 @@ +-e ../../.. diff --git a/samples/obo-e2e/src/test.http b/samples/obo-e2e/src/test.http new file mode 100644 index 00000000..01d59bd9 --- /dev/null +++ b/samples/obo-e2e/src/test.http @@ -0,0 +1,43 @@ +@baseUrl = http://localhost:7071 +@agentName = main +@contentType = application/json + +# Paste a real user access token here for User A +@userAToken = + +# Paste a hooks session token for User A +@userAHooksSessionToken = + +# Paste a real user access token here for User B +@userBToken = + +# Paste a hooks session token for User B +@userBHooksSessionToken = + +### User A -> downstream should see User A identity +POST {{baseUrl}}/agents/{{agentName}}/chat +Content-Type: {{contentType}} +X-MS-Access-Token: {{userAToken}} +X-MS-Hooks-Session-Token: {{userAHooksSessionToken}} + +{ + "prompt": "Call the downstream whoami MCP tool and return raw JSON" +} + +### User B -> downstream should see User B identity +POST {{baseUrl}}/agents/{{agentName}}/chat +Content-Type: {{contentType}} +X-MS-Access-Token: {{userBToken}} +X-MS-Hooks-Session-Token: {{userBHooksSessionToken}} + +{ + "prompt": "Call the downstream whoami MCP tool and return raw JSON" +} + +### No user token -> managed identity fallback expected +POST {{baseUrl}}/agents/{{agentName}}/chat +Content-Type: {{contentType}} + +{ + "prompt": "Call the downstream whoami MCP tool and return raw JSON" +} diff --git a/samples/obo-e2e/test.http b/samples/obo-e2e/test.http new file mode 100644 index 00000000..9f0c7571 --- /dev/null +++ b/samples/obo-e2e/test.http @@ -0,0 +1,44 @@ +@baseUrl = https://.azurewebsites.net +@agentName = main +@contentType = application/json + +# Paste a real user access token here for User A +@userAToken = + +# Paste a real user access token here for User B +@userBToken = + +### User A -> downstream should see User A identity +POST {{baseUrl}}/agents/{{agentName}}/chat +Content-Type: {{contentType}} +X-MS-TOKEN-AAD-ACCESS-TOKEN: {{userAToken}} + +{ + "prompt": "Call the downstream whoami MCP tool and return raw JSON" +} + +### User B -> downstream should see User B identity +POST {{baseUrl}}/agents/{{agentName}}/chat +Content-Type: {{contentType}} +X-MS-TOKEN-AAD-ACCESS-TOKEN: {{userBToken}} + +{ + "prompt": "Call the downstream whoami MCP tool and return raw JSON" +} + +### No user token -> managed identity fallback expected +POST {{baseUrl}}/agents/{{agentName}}/chat +Content-Type: {{contentType}} + +{ + "prompt": "Call the downstream whoami MCP tool and return raw JSON" +} + +### Claims challenge check (if consent/MFA required by downstream) +POST {{baseUrl}}/agents/{{agentName}}/chat +Content-Type: {{contentType}} +X-MS-TOKEN-AAD-ACCESS-TOKEN: {{userAToken}} + +{ + "prompt": "Call the protected MCP tool that requires additional consent or MFA" +} diff --git a/samples/obo-whoami-mcp-server/Dockerfile b/samples/obo-whoami-mcp-server/Dockerfile new file mode 100644 index 00000000..701c2f9b --- /dev/null +++ b/samples/obo-whoami-mcp-server/Dockerfile @@ -0,0 +1,14 @@ +FROM python:3.13-slim + +WORKDIR /app + +COPY requirements.txt ./ +RUN pip install --no-cache-dir -r requirements.txt + +COPY server.py ./ + +ENV HOST=0.0.0.0 +ENV PORT=8000 +EXPOSE 8000 + +CMD ["python", "server.py"] diff --git a/samples/obo-whoami-mcp-server/README.md b/samples/obo-whoami-mcp-server/README.md new file mode 100644 index 00000000..5758f32a --- /dev/null +++ b/samples/obo-whoami-mcp-server/README.md @@ -0,0 +1,78 @@ +# OBO WhoAmI MCP Server (Downstream Test Target) + +Minimal downstream MCP server for validating OBO pass-through from +`samples/obo-e2e`. + +This server exposes one tool: + +- `whoami` - reads inbound `Authorization: Bearer` token and returns decoded + claims (`oid`, `sub`, `aud`, `azp`, `appid`, etc.). +- It also reports whether `X-MS-Access-Token` and + `X-MS-Hooks-Session-Token` were forwarded by the upstream callback. + +Use this server as the `` target in +`samples/obo-e2e/src/mcp.json`. + +## Quick start (local) + +```powershell +cd samples/obo-whoami-mcp-server +python -m venv .venv +.venv\Scripts\Activate.ps1 +pip install -r requirements.txt +python server.py +``` + +Server starts at: + +- `http://localhost:8000/mcp` + +## Plug into obo-e2e sample + +In `samples/obo-e2e/src/mcp.json`: + +```json +{ + "servers": { + "whoami-api": { + "url": "http://localhost:8000/mcp", + "auth": { + "type": "obo", + "scope": "api:///.default" + } + } + } +} +``` + +For real delegated auth validation, use an HTTPS deployment with Entra auth in +front of this MCP server and set `auth.scope` to that API's audience. + +## Deploy notes (non-production) + +You can deploy this sample to any environment that can run a Python ASGI app +(for example, Azure Container Apps, App Service for Containers, or another +container host). + +Container startup command: + +```text +python server.py +``` + +Then use: + +- `https:///mcp` + +as `` in `samples/obo-e2e/src/mcp.json`. + +## Validation expectation + +When your upstream OBO sample calls `whoami`: + +- User A request should show User A `oid/sub`. +- User B request should show User B `oid/sub`. +- `hooks_session_token_present` should be `true` for BigMac callback tests. +- `token_source` should be `x-ms-access-token` when passthrough headers are used. +- No user token should show fallback identity behavior (or missing auth, + depending on downstream auth config). diff --git a/samples/obo-whoami-mcp-server/requirements.txt b/samples/obo-whoami-mcp-server/requirements.txt new file mode 100644 index 00000000..bceb8e39 --- /dev/null +++ b/samples/obo-whoami-mcp-server/requirements.txt @@ -0,0 +1,2 @@ +mcp==1.27.* +uvicorn==0.35.* diff --git a/samples/obo-whoami-mcp-server/server.py b/samples/obo-whoami-mcp-server/server.py new file mode 100644 index 00000000..8bd59ab4 --- /dev/null +++ b/samples/obo-whoami-mcp-server/server.py @@ -0,0 +1,116 @@ +from __future__ import annotations + +import base64 +import hashlib +import json +import os +from typing import Any + +import uvicorn +from mcp.server.fastmcp import Context, FastMCP + +mcp = FastMCP("obo-whoami") + + +def _base64url_decode(data: str) -> bytes: + padded = data + "=" * (-len(data) % 4) + return base64.urlsafe_b64decode(padded) + + +def _decode_jwt_payload(token: str) -> dict[str, Any]: + parts = token.split(".") + if len(parts) < 2: + return {"error": "invalid_jwt_format"} + + try: + payload_json = _base64url_decode(parts[1]).decode("utf-8") + payload = json.loads(payload_json) + if isinstance(payload, dict): + return payload + return {"error": "invalid_payload_type"} + except Exception as exc: # pragma: no cover - best effort decoding + return {"error": f"decode_failed: {exc}"} + + +def _request_headers(ctx: Context) -> dict[str, str]: + request_context = ctx.request_context + if request_context is None: + return {} + + request = getattr(request_context, "request", None) + if request is None: + return {} + + headers = getattr(request, "headers", None) + if headers is None: + return {} + + try: + return {str(k).lower(): str(v) for k, v in headers.items()} + except Exception: # pragma: no cover - defensive + return {} + + +@mcp.tool(name="whoami", description="Return identity claims from inbound bearer token") +def whoami(ctx: Context) -> dict[str, Any]: + """Echoes token identity details to validate OBO pass-through.""" + headers = _request_headers(ctx) + auth_header = headers.get("authorization", "") + hooks_session_token = headers.get("x-ms-hooks-session-token") + access_token_header = headers.get("x-ms-access-token") + + selected_token = access_token_header + selected_token_source = "x-ms-access-token" + + if not selected_token and auth_header.lower().startswith("bearer "): + selected_token = auth_header.split(" ", 1)[1].strip() + selected_token_source = "authorization" + + if not selected_token: + return { + "auth_present": bool(auth_header), + "token_present": False, + "error": "missing_token_headers", + "hooks_session_token_present": bool(hooks_session_token), + "observed_headers": sorted(list(headers.keys())), + } + + claims = _decode_jwt_payload(selected_token) + + # Surface only common identity/audience fields for quick validation. + selected_claims = { + "oid": claims.get("oid"), + "sub": claims.get("sub"), + "aud": claims.get("aud"), + "tid": claims.get("tid"), + "azp": claims.get("azp"), + "appid": claims.get("appid"), + "upn": claims.get("upn"), + "preferred_username": claims.get("preferred_username"), + "scp": claims.get("scp"), + "roles": claims.get("roles"), + } + + return { + "auth_present": bool(auth_header), + "token_present": True, + "token_source": selected_token_source, + "hooks_session_token_present": bool(hooks_session_token), + "hooks_session_token_sha256_12": ( + hashlib.sha256(hooks_session_token.encode("utf-8")).hexdigest()[:12] + if hooks_session_token + else None + ), + "token_sha256_12": hashlib.sha256(selected_token.encode("utf-8")).hexdigest()[:12], + "claims": selected_claims, + "raw_claims": claims, + } + + +app = mcp.streamable_http_app() + + +if __name__ == "__main__": + host = os.getenv("HOST", "0.0.0.0") + port = int(os.getenv("PORT", "8000")) + uvicorn.run(app, host=host, port=port) diff --git a/src/azure_functions_agents/__init__.py b/src/azure_functions_agents/__init__.py index af1446d6..6ec68930 100644 --- a/src/azure_functions_agents/__init__.py +++ b/src/azure_functions_agents/__init__.py @@ -13,7 +13,7 @@ as agent tools. """ -__version__ = "0.1.0b5" +__version__ = "0.1.0b8.dev2" # --------------------------------------------------------------------------- # Global MAF ExperimentalWarning suppression diff --git a/src/azure_functions_agents/_obo.py b/src/azure_functions_agents/_obo.py new file mode 100644 index 00000000..2816206a --- /dev/null +++ b/src/azure_functions_agents/_obo.py @@ -0,0 +1,655 @@ +"""On-Behalf-Of (OBO) token flow support. + +This module provides the infrastructure for exchanging user access tokens for +downstream API tokens using the OAuth 2.0 On-Behalf-Of flow. This allows agents +to call downstream APIs (MCP servers, user tools, etc.) using the authenticated +end-user's identity rather than the function app's managed identity. + +Architecture +------------ + +* :class:`UserContext` carries the user's identity through the request lifecycle. + It is created from incoming HTTP request headers (EasyAuth or Authorization) + and threaded through to tools that need to make authenticated downstream calls. + +* :class:`OboTokenProvider` handles the actual token exchange using MSAL. It + maintains an in-memory token cache keyed by (user_token_hash, scope) to avoid + redundant token exchanges within a request. + +* When OBO is not available (no user token, or OBO not configured), the system + falls back to managed identity via :mod:`._credential`. If that also fails, + the request fails with an appropriate error. + +Usage +----- + +1. Extract user token from request headers in the handler layer. +2. Create a :class:`UserContext` from the token. +3. Pass the context to the runner, which threads it to tools. +4. Tools call :meth:`UserContext.get_token_for_scope` to get downstream tokens. +""" + +from __future__ import annotations + +import asyncio +import hashlib +import os +import time +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any + +from ._logger import logger + +if TYPE_CHECKING: + from .config.schema import OboConfig + + +# --------------------------------------------------------------------------- +# Exceptions +# --------------------------------------------------------------------------- + + +class OboError(Exception): + """Base exception for OBO token exchange errors.""" + + def __init__(self, error: str, error_description: str | None = None) -> None: + self.error = error + self.error_description = error_description + message = f"{error}: {error_description}" if error_description else error + super().__init__(message) + + +class InteractionRequiredError(OboError): + """Downstream API requires user interaction (MFA, consent, etc.). + + The client must re-authenticate with the claims challenge to proceed. + This error should be surfaced to the client via HTTP 401 with a + WWW-Authenticate header containing the claims. + """ + + def __init__( + self, + error: str, + error_description: str | None = None, + claims: str | None = None, + ) -> None: + super().__init__(error, error_description) + self.claims = claims + + +# --------------------------------------------------------------------------- +# In-memory token cache +# --------------------------------------------------------------------------- + + +@dataclass +class _CachedToken: + """A cached access token with expiration.""" + + access_token: str + expires_on: int # Unix timestamp + + +# Global in-memory cache: (token_hash, scope) -> cached token +_token_cache: dict[tuple[str, str], _CachedToken] = {} +_cache_lock = asyncio.Lock() + +# Refresh tokens 5 minutes before expiry +_TOKEN_REFRESH_BUFFER_SECONDS = 300 + + +def _hash_token(token: str) -> str: + """Create a hash of the token for cache key (avoid storing raw tokens).""" + return hashlib.sha256(token.encode()).hexdigest()[:16] + + +async def _get_cached_token(token_hash: str, scope: str) -> str | None: + """Retrieve a valid cached token, or None if expired/missing.""" + async with _cache_lock: + cached = _token_cache.get((token_hash, scope)) + if cached is None: + return None + # Check if token is still valid (with buffer) + if cached.expires_on - _TOKEN_REFRESH_BUFFER_SECONDS <= int(time.time()): + del _token_cache[(token_hash, scope)] + return None + return cached.access_token + + +async def _set_cached_token(token_hash: str, scope: str, access_token: str, expires_on: int) -> None: + """Cache an access token.""" + async with _cache_lock: + _token_cache[(token_hash, scope)] = _CachedToken( + access_token=access_token, + expires_on=expires_on, + ) + + +def clear_token_cache() -> None: + """Clear all cached tokens. Useful for testing.""" + _token_cache.clear() + + +# --------------------------------------------------------------------------- +# OBO Token Provider +# --------------------------------------------------------------------------- + + +class OboTokenProvider: + """Handles OBO token exchange using MSAL. + + This class is typically instantiated once per application and shared + across requests. It uses MSAL's ConfidentialClientApplication for + the token exchange. + """ + + def __init__(self, config: OboConfig) -> None: + self._config = config + self._app: Any = None + self._app_lock = asyncio.Lock() + + async def _get_msal_app(self) -> Any: + """Lazily initialize the MSAL ConfidentialClientApplication.""" + if self._app is not None: + return self._app + + async with self._app_lock: + if self._app is not None: + return self._app + + # Import MSAL here to avoid import errors if not installed + try: + from msal import ConfidentialClientApplication + except ImportError as exc: + raise ImportError( + "MSAL is required for OBO support. " + "Install it with: pip install msal" + ) from exc + + authority = f"https://login.microsoftonline.com/{self._config.tenant_id}" + + # Build client credential (secret or certificate) + client_credential: str | dict[str, Any] | None = None + if self._config.client_secret: + client_credential = self._config.client_secret + # TODO: Add certificate support in future iteration + + self._app = ConfidentialClientApplication( + client_id=self._config.client_id, + client_credential=client_credential, + authority=authority, + ) + logger.debug("OBO: MSAL ConfidentialClientApplication initialized") + return self._app + + async def acquire_token_on_behalf_of( + self, + user_token: str, + scope: str, + ) -> str: + """Exchange a user token for a downstream API token. + + Parameters + ---------- + user_token: + The incoming user access token (from EasyAuth or Authorization header). + scope: + The scope for the downstream API (e.g., "https://graph.microsoft.com/.default"). + + Returns + ------- + str + The access token for the downstream API. + + Raises + ------ + InteractionRequiredError + If the downstream API requires user interaction (MFA, consent). + OboError + For other token exchange failures. + """ + # Check cache first + token_hash = _hash_token(user_token) + cached = await _get_cached_token(token_hash, scope) + if cached is not None: + logger.debug("OBO: Using cached token for scope %s", scope) + return cached + + # Perform token exchange + app = await self._get_msal_app() + + # MSAL's acquire_token_on_behalf_of is synchronous, run in thread pool + result = await asyncio.to_thread( + app.acquire_token_on_behalf_of, + user_assertion=user_token, + scopes=[scope], + ) + + if "error" in result: + error = result.get("error", "unknown_error") + error_description = result.get("error_description") + claims = result.get("claims") + + # Check for interaction_required errors + if error in ("interaction_required", "consent_required", "login_required"): + logger.warning( + "OBO: Interaction required for scope %s: %s", + scope, + error_description, + ) + raise InteractionRequiredError(error, error_description, claims) + + logger.error("OBO: Token exchange failed for scope %s: %s", scope, error_description) + raise OboError(error, error_description) + + access_token = result["access_token"] + expires_in = result.get("expires_in", 3600) + expires_on = int(time.time()) + expires_in + + # Cache the token + await _set_cached_token(token_hash, scope, access_token, expires_on) + logger.debug("OBO: Acquired and cached token for scope %s", scope) + + return access_token + + def get_scope_for_name(self, name: str) -> str | None: + """Look up a downstream scope by its configured name. + + Parameters + ---------- + name: + The name of the downstream scope (e.g., "graph", "custom_api"). + + Returns + ------- + str | None + The scope URI, or None if not configured. + """ + return self._config.downstream_scopes.get(name) + + +# --------------------------------------------------------------------------- +# User Context +# --------------------------------------------------------------------------- + + +@dataclass +class UserContext: + """Carries user identity through the request lifecycle. + + This context is created from incoming HTTP request headers and passed + to the runner, which threads it to tools. Tools can use this to acquire + tokens for downstream APIs via OBO. + + If no user token is present, the context will fall back to managed + identity for downstream calls. + """ + + access_token: str | None = None + """The incoming user access token, or None if unauthenticated.""" + + hooks_session_token: str | None = None + """The incoming hooks session token used by EasyAuth refresh callbacks.""" + + user_id: str | None = None + """The user's object ID from the token claims, if available.""" + + claims: dict[str, Any] = field(default_factory=dict) + """Decoded claims from the access token.""" + + forwardable_headers: dict[str, str] = field(default_factory=dict) + """Inbound headers to forward downstream (resolved via the header whitelist).""" + + _obo_provider: OboTokenProvider | None = field(default=None, repr=False) + """The OBO token provider for exchanging tokens.""" + + async def get_token_for_scope(self, scope: str) -> str | None: + """Acquire a token for a downstream API scope. + + Uses OBO if a user token is present and OBO is configured. + Falls back to managed identity if OBO is not available. + Returns None if neither method can provide a token. + + Parameters + ---------- + scope: + The scope for the downstream API. + + Returns + ------- + str | None + The access token, or None if unavailable. + """ + # Try OBO first if we have a user token + if self.access_token and self._obo_provider: + try: + return await self._obo_provider.acquire_token_on_behalf_of( + self.access_token, + scope, + ) + except OboError as exc: + logger.warning("OBO failed, falling back to managed identity: %s", exc) + # Fall through to managed identity + + # Fall back to managed identity + return await self._get_managed_identity_token(scope) + + async def _get_managed_identity_token(self, scope: str) -> str | None: + """Acquire a token using managed identity.""" + try: + from ._credential import build_async_credential + + credential = build_async_credential() + token = await credential.get_token(scope) + return token.token + except Exception as exc: + logger.warning("Managed identity token acquisition failed: %s", exc) + return None + + def get_token_for_scope_name(self, name: str) -> str | None: + """Look up a scope by name and acquire a token for it. + + This is a convenience method for tools that use named scopes + from the configuration. + """ + if self._obo_provider is None: + return None + scope = self._obo_provider.get_scope_for_name(name) + if scope is None: + logger.warning("Unknown downstream scope name: %s", name) + return None + # Note: This would need to be async - keeping sync for interface consistency + # In practice, tools should use get_token_for_scope directly + return None + + @property + def is_authenticated(self) -> bool: + """Check if a user token is present.""" + return self.access_token is not None + + @property + def has_obo_support(self) -> bool: + """Check if OBO is configured and available.""" + return self._obo_provider is not None and self.access_token is not None + + +# --------------------------------------------------------------------------- +# Factory functions +# --------------------------------------------------------------------------- + +# Global OBO provider instance (created once, shared across requests) +_obo_provider: OboTokenProvider | None = None +_obo_provider_lock = asyncio.Lock() + + +async def get_obo_provider(config: OboConfig | None) -> OboTokenProvider | None: + """Get or create the global OBO token provider. + + Parameters + ---------- + config: + The OBO configuration. If None or disabled, returns None. + + Returns + ------- + OboTokenProvider | None + The provider instance, or None if OBO is not configured. + """ + global _obo_provider + + if config is None or not config.enabled: + return None + + if _obo_provider is not None: + return _obo_provider + + async with _obo_provider_lock: + if _obo_provider is not None: + return _obo_provider + _obo_provider = OboTokenProvider(config) + logger.info("OBO: Token provider initialized") + return _obo_provider + + +def reset_obo_provider() -> None: + """Reset the global OBO provider. Useful for testing.""" + global _obo_provider + _obo_provider = None + clear_token_cache() + + +def create_user_context( + access_token: str | None = None, + hooks_session_token: str | None = None, + user_id: str | None = None, + claims: dict[str, Any] | None = None, + obo_provider: OboTokenProvider | None = None, + forwardable_headers: dict[str, str] | None = None, +) -> UserContext: + """Create a UserContext for the current request. + + Parameters + ---------- + access_token: + The incoming user access token from the request. + hooks_session_token: + The hooks session token from the request, if present. + user_id: + The user's object ID, if known. + claims: + Decoded claims from the access token. + obo_provider: + The OBO token provider instance. + forwardable_headers: + Inbound headers to forward downstream, resolved via the header whitelist. + + Returns + ------- + UserContext + A new context instance for this request. + """ + return UserContext( + access_token=access_token, + hooks_session_token=hooks_session_token, + user_id=user_id, + claims=claims or {}, + forwardable_headers=forwardable_headers or {}, + _obo_provider=obo_provider, + ) + + +# --------------------------------------------------------------------------- +# Header extraction utilities +# --------------------------------------------------------------------------- + +# EasyAuth header names +EASYAUTH_ACCESS_TOKEN_HEADER = "X-MS-TOKEN-AAD-ACCESS-TOKEN" +EASYAUTH_ID_TOKEN_HEADER = "X-MS-TOKEN-AAD-ID-TOKEN" +BIGMAC_ACCESS_TOKEN_HEADER = "X-MS-Access-Token" +BIGMAC_HOOKS_SESSION_TOKEN_HEADER = "X-MS-Hooks-Session-Token" +EASYAUTH_PRINCIPAL_ID_HEADER = "X-MS-CLIENT-PRINCIPAL-ID" +EASYAUTH_PRINCIPAL_NAME_HEADER = "X-MS-CLIENT-PRINCIPAL-NAME" + +# Standard Authorization header +AUTHORIZATION_HEADER = "Authorization" +BEARER_PREFIX = "Bearer " + +# Environment variable holding a comma-separated list of inbound header names to +# forward downstream. When unset or empty, the runtime falls back to the trusted +# BigMac headers below. +AGENTS_HEADER_WHITELIST_ENV = "AGENTS_HEADER_WHITELIST" + +# Default trusted headers forwarded downstream when no whitelist is configured. +DEFAULT_FORWARDED_HEADERS: tuple[str, ...] = ( + BIGMAC_ACCESS_TOKEN_HEADER, + BIGMAC_HOOKS_SESSION_TOKEN_HEADER, +) + + +def extract_user_token_from_headers(headers: dict[str, str] | Any) -> str | None: + """Extract the user access token from request headers. + + Checks EasyAuth headers first, then falls back to Authorization header. + + Parameters + ---------- + headers: + The request headers (dict-like or object with get method). + + Returns + ------- + str | None + The access token, or None if not found. + """ + # Helper to get header value (case-insensitive) + def get_header(name: str) -> str | None: + if hasattr(headers, "get"): + # Try exact match first + value = headers.get(name) + if value: + return value.strip() if isinstance(value, str) else None + # Try case-insensitive + if hasattr(headers, "items"): + for key, val in headers.items(): + if key.lower() == name.lower(): + return val.strip() if isinstance(val, str) else None + return None + + # BigMac callback path: prefer explicit access token header first. + token = get_header(BIGMAC_ACCESS_TOKEN_HEADER) + if token: + logger.debug("OBO: Found user token in BigMac access token header") + return token + + # Try EasyAuth access token header. + token = get_header(EASYAUTH_ACCESS_TOKEN_HEADER) + if token: + logger.debug("OBO: Found user token in EasyAuth header") + return token + + # EasyAuth ID token fallback is used when access token is not present. + token = get_header(EASYAUTH_ID_TOKEN_HEADER) + if token: + logger.debug("OBO: Found user token in EasyAuth id token header") + return token + + # Try Authorization header + auth = get_header(AUTHORIZATION_HEADER) + if auth and auth.startswith(BEARER_PREFIX): + token = auth[len(BEARER_PREFIX) :].strip() + if token: + logger.debug("OBO: Found user token in Authorization header") + return token + + return None + + +def extract_hooks_session_token_from_headers(headers: dict[str, str] | Any) -> str | None: + """Extract the hooks session token from request headers. + + Parameters + ---------- + headers: + The request headers (dict-like or object with get method). + + Returns + ------- + str | None + The hooks session token, or None if not found. + """ + + if not hasattr(headers, "get"): + return None + + value = headers.get(BIGMAC_HOOKS_SESSION_TOKEN_HEADER) + if value: + return value.strip() if isinstance(value, str) else None + + # Case-insensitive fallback + if hasattr(headers, "items"): + for key, val in headers.items(): + if key.lower() == BIGMAC_HOOKS_SESSION_TOKEN_HEADER.lower(): + return val.strip() if isinstance(val, str) else None + + return None + + +def extract_user_id_from_headers(headers: dict[str, str] | Any) -> str | None: + """Extract the user's principal ID from EasyAuth headers. + + Parameters + ---------- + headers: + The request headers. + + Returns + ------- + str | None + The user's principal ID, or None if not found. + """ + if hasattr(headers, "get"): + value = headers.get(EASYAUTH_PRINCIPAL_ID_HEADER) + if value: + return value.strip() if isinstance(value, str) else None + # Case-insensitive fallback + if hasattr(headers, "items"): + for key, val in headers.items(): + if key.lower() == EASYAUTH_PRINCIPAL_ID_HEADER.lower(): + return val.strip() if isinstance(val, str) else None + return None + + +def _get_header_value(headers: dict[str, str] | Any, name: str) -> str | None: + """Case-insensitively read a single header value.""" + if not hasattr(headers, "get"): + return None + value = headers.get(name) + if value: + return value.strip() if isinstance(value, str) else None + if hasattr(headers, "items"): + for key, val in headers.items(): + if key.lower() == name.lower(): + return val.strip() if isinstance(val, str) else None + return None + + +def get_forwarded_header_names() -> list[str]: + """Return the configured header whitelist, or an empty list when unset. + + Reads the ``AGENTS_HEADER_WHITELIST`` environment variable, which holds a + comma-separated list of header names to forward downstream. Returns an empty + list when the variable is unset or contains no header names; in that case the + runtime falls back to the default trusted headers + (:data:`DEFAULT_FORWARDED_HEADERS`) at the forwarding site. + """ + raw = os.environ.get(AGENTS_HEADER_WHITELIST_ENV, "") + return [part.strip() for part in raw.split(",") if part.strip()] + + +def extract_forwardable_headers(headers: dict[str, str] | Any) -> dict[str, str]: + """Resolve the whitelisted headers to forward downstream for this request. + + When ``AGENTS_HEADER_WHITELIST`` is configured, each whitelisted header is + read from the inbound request (case-insensitively) and forwarded under its + configured name. When no whitelist is configured, this returns an empty + mapping and the runtime falls back to the default trusted headers at the + forwarding site. + + Parameters + ---------- + headers: + The inbound request headers (dict-like or object with ``get``). + + Returns + ------- + dict[str, str] + A mapping of header name to value to forward downstream. Empty when no + whitelist is configured or no whitelisted headers are present on the + request. + """ + forwarded: dict[str, str] = {} + for name in get_forwarded_header_names(): + value = _get_header_value(headers, name) + if value: + forwarded[name] = value + return forwarded diff --git a/src/azure_functions_agents/app.py b/src/azure_functions_agents/app.py index d47ea4d6..fccb28ff 100644 --- a/src/azure_functions_agents/app.py +++ b/src/azure_functions_agents/app.py @@ -9,6 +9,7 @@ import azure.functions as func from ._logger import logger +from ._obo import OboTokenProvider, reset_obo_provider from ._observability import configure_observability from .config.loader import load_agent_specs, load_global_config from .config.merge import compose @@ -17,6 +18,7 @@ from .discovery.mcp import discover_mcp_servers from .discovery.skills import discover_skills from .discovery.tools import discover_user_tools +from .registration._handlers import set_obo_provider from .registration._naming import allocate_unique_function_name from .registration.capabilities import build_capabilities from .registration.endpoints import register_builtin_endpoints @@ -48,6 +50,22 @@ def create_function_app(app_root: Path | None = None) -> func.FunctionApp: resolved_root = get_app_root() global_config = load_global_config(resolved_root) + + # Initialize OBO provider if auth.obo is configured + obo_provider = None + if global_config.auth and global_config.auth.obo and global_config.auth.obo.enabled: + obo_config = global_config.auth.obo + obo_provider = OboTokenProvider(obo_config) + set_obo_provider(obo_provider) + logger.info( + "OBO authentication enabled with client_id=%s, tenant_id=%s", + obo_config.client_id[:8] + "..." if len(obo_config.client_id) > 8 else obo_config.client_id, + obo_config.tenant_id, + ) + else: + # Ensure OBO provider is cleared if not configured + reset_obo_provider() + set_obo_provider(None) # Bootstrap observability before anything runs so MAF gen_ai spans + runtime spans/metrics # flow to Application Insights with zero app code. No-op unless a telemetry provider is active. diff --git a/src/azure_functions_agents/config/__init__.py b/src/azure_functions_agents/config/__init__.py index 73dede6e..293ee8f3 100644 --- a/src/azure_functions_agents/config/__init__.py +++ b/src/azure_functions_agents/config/__init__.py @@ -28,10 +28,12 @@ ) from azure_functions_agents.config.schema import ( AgentSpec, + AuthConfig, BuiltinEndpointsConfig, DynamicSessionsCodeInterpreterConfig, GlobalConfig, McpFilter, + OboConfig, ResolvedAgent, SkillsFilter, SystemToolsAgentOverride, @@ -48,10 +50,12 @@ "_INLINE_DOLLAR_PATTERN", "_INLINE_PERCENT_PATTERN", "AgentSpec", + "AuthConfig", "BuiltinEndpointsConfig", "DynamicSessionsCodeInterpreterConfig", "GlobalConfig", "McpFilter", + "OboConfig", "ResolvedAgent", "SkillsFilter", "SystemToolsAgentOverride", diff --git a/src/azure_functions_agents/config/schema.py b/src/azure_functions_agents/config/schema.py index b035139f..54d00185 100644 --- a/src/azure_functions_agents/config/schema.py +++ b/src/azure_functions_agents/config/schema.py @@ -89,11 +89,49 @@ class SystemToolsAgentOverride(BaseModel): dynamic_sessions_code_interpreter: bool | None = None +class OboConfig(BaseModel): + """On-Behalf-Of (OBO) authentication configuration. + + Enables agents to call downstream APIs using the authenticated end-user's + identity rather than the function app's managed identity. Requires + pre-consented permissions on the Azure AD app registration. + + The OBO flow exchanges an incoming user access token for a new token + scoped to a downstream API. This allows the agent to act on behalf of + the user when calling MCP servers, user tools, or other APIs. + """ + + model_config = ConfigDict(extra="forbid") + + enabled: bool = True + client_id: str + client_secret: str | None = None + tenant_id: str + downstream_scopes: dict[str, str] = Field(default_factory=dict) + + @field_validator("client_id", "tenant_id") + @classmethod + def validate_required_string(cls, value: str) -> str: + trimmed = value.strip() + if not trimmed: + raise ValueError("value must be non-empty") + return trimmed + + +class AuthConfig(BaseModel): + """Global authentication configuration.""" + + model_config = ConfigDict(extra="forbid") + + obo: OboConfig | None = None + + class GlobalConfig(BaseModel): """Top-level agents.config.yaml schema.""" model_config = ConfigDict(extra="forbid") + auth: AuthConfig | None = None system_tools: SystemToolsConfig | None = None model: str | None = None timeout: float | None = None diff --git a/src/azure_functions_agents/discovery/mcp.py b/src/azure_functions_agents/discovery/mcp.py index 3dcbab70..fd160698 100644 --- a/src/azure_functions_agents/discovery/mcp.py +++ b/src/azure_functions_agents/discovery/mcp.py @@ -3,22 +3,50 @@ from __future__ import annotations import asyncio +import contextvars import json import time from pathlib import Path -from typing import Any, cast +from typing import TYPE_CHECKING, Any, cast from agent_framework import MCPStreamableHTTPTool from .._credential import build_credential, build_credential_with_client_id from .._logger import logger +from .._obo import BIGMAC_ACCESS_TOKEN_HEADER, BIGMAC_HOOKS_SESSION_TOKEN_HEADER from ..config.env import has_unresolved_placeholders, resolve_env_vars_in_data +if TYPE_CHECKING: + from .._obo import UserContext + type MCPTool = MCPStreamableHTTPTool _DISCOVERED_MCP_SERVERS_CACHE: dict[Path, dict[str, MCPTool]] = {} _DEFAULT_TOKEN_REFRESH_OFFSET_SECONDS = 300 +# Context variable to store the current request's UserContext for OBO +_current_user_context: contextvars.ContextVar[UserContext | None] = contextvars.ContextVar( + "current_user_context", default=None +) + + +def set_current_user_context(user_context: UserContext | None) -> contextvars.Token[UserContext | None]: + """Set the current user context for OBO-enabled MCP servers. + + Returns a token that can be used to reset the context. + """ + return _current_user_context.set(user_context) + + +def get_current_user_context() -> UserContext | None: + """Get the current user context for OBO.""" + return _current_user_context.get() + + +def reset_current_user_context(token: contextvars.Token[UserContext | None]) -> None: + """Reset the user context using the token from set_current_user_context.""" + _current_user_context.reset(token) + def clear_mcp_cache() -> None: """Clear cached MCP server discovery results.""" @@ -43,7 +71,9 @@ def static_header_provider(_ctx: Any) -> dict[str, str]: return static_header_provider + auth_type = str(auth.get("type", "managed_identity")).strip().lower() scope = str(auth.get("scope", "")).strip() + if not scope: logger.warning("MCP server auth requires a non-empty 'scope'") if not static_headers: @@ -54,6 +84,11 @@ def missing_scope_header_provider(_ctx: Any) -> dict[str, str]: return missing_scope_header_provider + # Handle OBO authentication type + if auth_type == "obo": + return _build_obo_header_provider(scope, static_headers) + + # Default: managed identity authentication client_id = str(auth.get("client_id", "")).strip() if has_unresolved_placeholders(client_id): client_id = "" @@ -76,6 +111,128 @@ def default_credential_header_provider(_ctx: Any) -> dict[str, str]: return default_credential_header_provider +def _build_obo_header_provider(scope: str, static_headers: dict[str, str]) -> Any: + """Build a header provider that uses OBO to get tokens on behalf of the user. + + If no user context is available, falls back to managed identity. + If managed identity also fails, the request will fail. + """ + # Fallback credential for when no user context is available + fallback_credential = build_credential() + fallback_cached_token: dict[str, str | int] = {"token": "", "expires_on": 0} + + def _get_or_refresh_managed_identity_token() -> str: + now = int(time.time()) + expires_on = int(fallback_cached_token["expires_on"]) + if not fallback_cached_token["token"] or expires_on - _DEFAULT_TOKEN_REFRESH_OFFSET_SECONDS <= now: + token = fallback_credential.get_token(scope) + fallback_cached_token["token"] = token.token + fallback_cached_token["expires_on"] = token.expires_on + return str(fallback_cached_token["token"]) + + def obo_header_provider(_ctx: Any) -> dict[str, str]: + user_context = get_current_user_context() + + # Forward user headers and always use MI auth. A configured header + # whitelist (AGENTS_HEADER_WHITELIST) takes precedence and is honored + # whenever it produced matching headers, regardless of the BigMac tokens. + forwardable = getattr(user_context, "forwardable_headers", None) or {} + if user_context is not None and forwardable: + mi_token = _get_or_refresh_managed_identity_token() + result = dict(static_headers) + result.update(forwardable) + result["Authorization"] = f"Bearer {mi_token}" + logger.debug("MCP: Forwarding whitelisted headers for scope %s", scope) + return result + + # Otherwise, fall back to the trusted BigMac headers when both tokens + # are present. + if ( + user_context is not None + and user_context.hooks_session_token + and user_context.access_token + ): + mi_token = _get_or_refresh_managed_identity_token() + result = dict(static_headers) + result[BIGMAC_ACCESS_TOKEN_HEADER] = user_context.access_token + result[BIGMAC_HOOKS_SESSION_TOKEN_HEADER] = user_context.hooks_session_token + result["Authorization"] = f"Bearer {mi_token}" + logger.debug("MCP: Using BigMac hook-session callback headers for scope %s", scope) + return result + + # Try OBO if user context is available + if user_context is not None and user_context.has_obo_support: + try: + # Run async token acquisition in sync context + import asyncio + + try: + loop = asyncio.get_running_loop() + except RuntimeError: + loop = None + + if loop is not None: + # We're in an async context, need to use run_coroutine_threadsafe + # This is tricky - the header provider is called from a thread + # Let's try getting the token synchronously + pass + + # For now, try to get token synchronously via a new event loop + # This is not ideal but works for the header provider context + token = _get_obo_token_sync(user_context, scope) + if token: + result = dict(static_headers) + result["Authorization"] = f"Bearer {token}" + logger.debug("MCP: Using OBO token for scope %s", scope) + return result + except Exception as exc: + logger.warning("MCP: OBO token acquisition failed, falling back to managed identity: %s", exc) + + # Fallback to managed identity + mi_token = _get_or_refresh_managed_identity_token() + + result = dict(static_headers) + result["Authorization"] = f"Bearer {mi_token}" + logger.debug("MCP: Using managed identity token for scope %s", scope) + return result + + return obo_header_provider + + +def _get_obo_token_sync(user_context: Any, scope: str) -> str | None: + """Synchronously get an OBO token for the given scope. + + This runs the async token acquisition in a new event loop when called + from a sync context (like the header provider). + """ + import asyncio + + async def _get_token() -> str | None: + return await user_context.get_token_for_scope(scope) + + try: + # Try to run in existing loop + try: + loop = asyncio.get_running_loop() + except RuntimeError: + loop = None + + if loop is None: + # No running loop, create one + return asyncio.run(_get_token()) + else: + # Running loop exists - we're in a thread, need different approach + # Create a new loop in this thread + new_loop = asyncio.new_event_loop() + try: + return new_loop.run_until_complete(_get_token()) + finally: + new_loop.close() + except Exception as exc: + logger.warning("Failed to get OBO token synchronously: %s", exc) + return None + + def _build_http_client(header_provider: Any) -> Any: if header_provider is None: return None diff --git a/src/azure_functions_agents/registration/_handlers.py b/src/azure_functions_agents/registration/_handlers.py index 3d5c0330..dbe0dc62 100644 --- a/src/azure_functions_agents/registration/_handlers.py +++ b/src/azure_functions_agents/registration/_handlers.py @@ -7,13 +7,22 @@ import uuid from collections.abc import Callable from importlib import import_module -from typing import Any, cast +from typing import TYPE_CHECKING, Any, cast import azure.functions as func import jsonschema from azurefunctions.extensions.http.fastapi import Request, Response from .._logger import logger +from .._obo import ( + InteractionRequiredError, + UserContext, + create_user_context, + extract_forwardable_headers, + extract_hooks_session_token_from_headers, + extract_user_id_from_headers, + extract_user_token_from_headers, +) from .._observability import ( ATTR_FAULT_DOMAIN, FaultDomain, @@ -24,6 +33,9 @@ from ..config import ResolvedAgent, _to_bool from .capabilities import AgentCapabilities +if TYPE_CHECKING: + from .._obo import OboTokenProvider + AUTH_LEVEL_MAP = { "anonymous": func.AuthLevel.ANONYMOUS, "function": func.AuthLevel.FUNCTION, @@ -31,6 +43,65 @@ } _SESSION_ID_HEADER = "x-ms-session-id" +# Global OBO provider reference (set during app initialization) +_obo_provider: OboTokenProvider | None = None + + +def set_obo_provider(provider: OboTokenProvider | None) -> None: + """Set the global OBO provider for handlers to use.""" + global _obo_provider + _obo_provider = provider + + +def get_handler_obo_provider() -> OboTokenProvider | None: + """Get the current OBO provider.""" + return _obo_provider + + +async def _build_user_context_from_request(req: Request) -> UserContext: + """Extract user context from HTTP request headers.""" + headers = getattr(req, "headers", {}) + access_token = extract_user_token_from_headers(headers) + hooks_session_token = extract_hooks_session_token_from_headers(headers) + user_id = extract_user_id_from_headers(headers) + forwardable_headers = extract_forwardable_headers(headers) + + return create_user_context( + access_token=access_token, + hooks_session_token=hooks_session_token, + user_id=user_id, + obo_provider=_obo_provider, + forwardable_headers=forwardable_headers, + ) + + +def _build_interaction_required_response(exc: InteractionRequiredError, session_id: str) -> Response: + """Build HTTP 401 response for OBO interaction required errors.""" + import base64 + + headers: dict[str, str] = {"x-ms-session-id": session_id} + + # Build WWW-Authenticate header with error info and claims + www_auth_parts = [f'Bearer error="{exc.error}"'] + if exc.error_description: + desc = exc.error_description.replace('"', '\\"') + www_auth_parts.append(f'error_description="{desc}"') + if exc.claims: + claims_b64 = base64.b64encode(exc.claims.encode()).decode() + www_auth_parts.append(f'claims="{claims_b64}"') + + headers["WWW-Authenticate"] = ", ".join(www_auth_parts) + + return Response( + content=json.dumps({ + "error": exc.error, + "error_description": exc.error_description, + "claims": exc.claims, + }), + status_code=401, + media_type="application/json", + headers=headers, + ) def serialize_trigger_data(trigger_data: Any) -> str: """Serialize trigger binding data to a JSON string.""" @@ -314,6 +385,7 @@ async def _handler(req: Request) -> Response: ) as span: try: session_id = _request_header_value(req, _SESSION_ID_HEADER) or _new_session_id() + user_context = await _build_user_context_from_request(req) span.set_attribute("af.agent.session_id", session_id) try: body = await req.json() @@ -361,6 +433,7 @@ async def _handler(req: Request) -> Response: tools=capabilities.filtered_user_tools, mcp_tools=capabilities.filtered_mcp_tools, skill_paths=capabilities.enabled_skill_paths, + user_context=user_context, ) _set_run_result_attributes(span, result) @@ -449,6 +522,13 @@ async def _handler(req: Request) -> Response: media_type="text/plain", headers={_SESSION_ID_HEADER: session_id}, ) + except InteractionRequiredError as exc: + logger.warning( + "HTTP agent '%s' OBO interaction required: %s", + resolved.name, + exc.error_description, + ) + return _build_interaction_required_response(exc, session_id) except Exception as exc: span.set_attribute("af.agent.outcome", "error") span.record_exception(exc, fault_domain=FaultDomain.UNKNOWN) diff --git a/src/azure_functions_agents/registration/endpoints.py b/src/azure_functions_agents/registration/endpoints.py index 50f8f631..ce9d1737 100644 --- a/src/azure_functions_agents/registration/endpoints.py +++ b/src/azure_functions_agents/registration/endpoints.py @@ -11,17 +11,29 @@ import uuid from collections.abc import AsyncIterator from pathlib import Path -from typing import Any +from typing import TYPE_CHECKING, Any import azure.functions as func from azurefunctions.extensions.http.fastapi import Request, Response, StreamingResponse from .._logger import logger +from .._obo import ( + InteractionRequiredError, + UserContext, + create_user_context, + extract_forwardable_headers, + extract_hooks_session_token_from_headers, + extract_user_id_from_headers, + extract_user_token_from_headers, +) from ..config import ResolvedAgent -from ._handlers import build_sandbox_tools_for_session +from ._handlers import build_sandbox_tools_for_session, get_handler_obo_provider from ._naming import _function_name_from_source, _safe_function_name, allocate_unique_builtin_slug from .capabilities import AgentCapabilities +if TYPE_CHECKING: + pass + _MCP_AGENT_TOOL_PROPERTIES = json.dumps( [ { @@ -97,12 +109,30 @@ def _resolve_builtin_endpoints_session_id(session_id: str | None) -> str: return session_id or uuid.uuid4().hex +async def _build_user_context_from_request(req: Request) -> UserContext: + """Extract user context from HTTP request headers.""" + headers = getattr(req, "headers", {}) + access_token = extract_user_token_from_headers(headers) + hooks_session_token = extract_hooks_session_token_from_headers(headers) + user_id = extract_user_id_from_headers(headers) + forwardable_headers = extract_forwardable_headers(headers) + + return create_user_context( + access_token=access_token, + hooks_session_token=hooks_session_token, + user_id=user_id, + obo_provider=get_handler_obo_provider(), + forwardable_headers=forwardable_headers, + ) + + async def _run_builtin_agent( prompt: str, *, resolved: ResolvedAgent, capabilities: AgentCapabilities, session_id: str | None, + user_context: UserContext | None = None, ) -> Any: resolved_session_id = _resolve_builtin_endpoints_session_id(session_id) sandbox_tools = build_sandbox_tools_for_session(resolved, resolved_session_id) @@ -116,6 +146,7 @@ async def _run_builtin_agent( tools=capabilities.filtered_user_tools, mcp_tools=capabilities.filtered_mcp_tools, skill_paths=capabilities.enabled_skill_paths, + user_context=user_context, ) @@ -125,6 +156,7 @@ def _run_builtin_agent_stream( resolved: ResolvedAgent, capabilities: AgentCapabilities, session_id: str | None, + user_context: UserContext | None = None, ) -> Any: resolved_session_id = _resolve_builtin_endpoints_session_id(session_id) sandbox_tools = build_sandbox_tools_for_session(resolved, resolved_session_id) @@ -138,6 +170,7 @@ def _run_builtin_agent_stream( tools=capabilities.filtered_user_tools, mcp_tools=capabilities.filtered_mcp_tools, skill_paths=capabilities.enabled_skill_paths, + user_context=user_context, ) @@ -156,6 +189,41 @@ def _json_error(message: str, status_code: int = 500) -> Response: ) +def _interaction_required_error(exc: InteractionRequiredError) -> Response: + """Build HTTP 401 response for OBO interaction required errors. + + When downstream APIs require user interaction (MFA, consent, etc.), + we return HTTP 401 with WWW-Authenticate header containing the claims + challenge. The client must re-authenticate with these claims. + """ + headers: dict[str, str] = {} + + # Build WWW-Authenticate header with error info and claims + www_auth_parts = [f'Bearer error="{exc.error}"'] + if exc.error_description: + # Escape quotes in description + desc = exc.error_description.replace('"', '\\"') + www_auth_parts.append(f'error_description="{desc}"') + if exc.claims: + # Claims should be base64-encoded for the header + import base64 + claims_b64 = base64.b64encode(exc.claims.encode()).decode() + www_auth_parts.append(f'claims="{claims_b64}"') + + headers["WWW-Authenticate"] = ", ".join(www_auth_parts) + + return Response( + content=json.dumps({ + "error": exc.error, + "error_description": exc.error_description, + "claims": exc.claims, + }), + status_code=401, + media_type="application/json", + headers=headers, + ) + + def _sse_error_response(message: str, status_code: int = 400) -> StreamingResponse: async def error_gen() -> AsyncIterator[str]: yield f"data: {json.dumps({'type': 'error', 'content': message})}\n\n" @@ -206,11 +274,13 @@ async def chat(req: Request) -> Response: body = await req.json() prompt = _extract_prompt_from_body(body) session_id = req.headers.get("x-ms-session-id") + user_context = await _build_user_context_from_request(req) result = await _run_builtin_agent( prompt, resolved=resolved, capabilities=capabilities, session_id=session_id, + user_context=user_context, ) return Response( json.dumps( @@ -225,6 +295,13 @@ async def chat(req: Request) -> Response: ) except ValueError as exc: return _json_error(str(exc), status_code=400) + except InteractionRequiredError as exc: + logger.warning( + "Built-in chat API OBO interaction required for '%s': %s", + resolved.name, + exc.error_description, + ) + return _interaction_required_error(exc) except Exception as exc: error_msg = _format_exception_message(exc) logger.error("Built-in chat API error for '%s': %s", resolved.name, error_msg) @@ -247,17 +324,26 @@ async def chat_stream(req: Request) -> StreamingResponse: body = await req.json() prompt = _extract_prompt_from_body(body) session_id = req.headers.get("x-ms-session-id") + user_context = await _build_user_context_from_request(req) return StreamingResponse( _run_builtin_agent_stream( prompt, resolved=resolved, capabilities=capabilities, session_id=session_id, + user_context=user_context, ), media_type="text/event-stream", ) except ValueError as exc: return _sse_error_response(str(exc), status_code=400) + except InteractionRequiredError as exc: + logger.warning( + "Built-in chat stream OBO interaction required for '%s': %s", + resolved.name, + exc.error_description, + ) + return _interaction_required_error(exc) except Exception as exc: error_msg = _format_exception_message(exc) logger.error("Built-in chat stream error for '%s': %s", resolved.name, error_msg) diff --git a/src/azure_functions_agents/runner.py b/src/azure_functions_agents/runner.py index dd8baddb..0a9b2426 100644 --- a/src/azure_functions_agents/runner.py +++ b/src/azure_functions_agents/runner.py @@ -45,16 +45,23 @@ from collections.abc import AsyncIterator from dataclasses import dataclass, field from pathlib import Path -from typing import Any +from typing import TYPE_CHECKING, Any from ._blob_history import build_blob_provider_from_environment from ._logger import logger from .client_manager import get_client_manager from .config.env import runtime_env_value from .config.paths import get_app_root, resolve_config_dir -from .discovery.mcp import discover_mcp_servers +from .discovery.mcp import ( + discover_mcp_servers, + reset_current_user_context, + set_current_user_context, +) from .discovery.tools import discover_user_tools +if TYPE_CHECKING: + from ._obo import UserContext + # --------------------------------------------------------------------------- # Constants # --------------------------------------------------------------------------- @@ -337,6 +344,7 @@ async def run_agent( model: str | None = None, session_id: str | None = None, sandbox_tools: list[Any] | None = None, + user_context: UserContext | None = None, ) -> AgentResult: """Execute a single prompt against the configured agent backend. @@ -374,6 +382,10 @@ async def run_agent( bound to a specific ACA session pool. ``None`` adds no sandbox tools; pass a list to enable them. Per-call because the ACA session id is baked into each tool's closure. + user_context: + Optional user context carrying the authenticated user's identity. + When provided with OBO configuration, enables downstream API calls + to be made on behalf of the user rather than using managed identity. Notes ----- @@ -392,19 +404,26 @@ async def run_agent( sandbox_tools=sandbox_tools, ) + # Set user context for OBO-enabled MCP servers + context_token = set_current_user_context(user_context) + lock = await _get_session_lock(resolved_id) - async with lock: - try: - response = await asyncio.wait_for( - agent.run( - prompt, - session=session, - options=_build_chat_options_from_environment(), - ), - timeout=timeout, - ) - except TimeoutError: - raise RuntimeError(f"Agent run timed out after {timeout}s") from None + try: + async with lock: + try: + response = await asyncio.wait_for( + agent.run( + prompt, + session=session, + options=_build_chat_options_from_environment(), + ), + timeout=timeout, + ) + except TimeoutError: + raise RuntimeError(f"Agent run timed out after {timeout}s") from None + finally: + # Reset user context after agent run + reset_current_user_context(context_token) # Extract assistant text from the final response. text = "" @@ -465,6 +484,7 @@ async def run_agent_stream( model: str | None = None, session_id: str | None = None, sandbox_tools: list[Any] | None = None, + user_context: UserContext | None = None, ) -> AsyncIterator[str]: """SSE-formatted async generator yielding ``data: {...}\\n\\n`` lines. @@ -479,6 +499,8 @@ async def run_agent_stream( sandbox tools; pass a list to enable them. * ``skill_paths`` enables MAF's :class:`SkillsProvider` for the listed directories. ``None`` or ``[]`` disables skills. + * ``user_context`` carries the authenticated user's identity for OBO + token flow to downstream APIs. * To fully disable all tools from a direct API call, pass ``tools=[], mcp_tools=[], sandbox_tools=None``. @@ -496,116 +518,123 @@ async def run_agent_stream( """ timeout = timeout if timeout is not None else DEFAULT_TIMEOUT - try: - agent, session, resolved_id = await _build_agent_session_history( - instructions=instructions, - session_id=session_id, - tools=tools, - mcp_tools=mcp_tools, - skill_paths=skill_paths, - model=model, - sandbox_tools=sandbox_tools, - ) - except Exception as exc: - logger.error("Failed to build agent session: %s", exc, exc_info=True) - yield f"data: {json.dumps({'type': 'error', 'content': str(exc)})}\n\n" - return - - yield f"data: {json.dumps({'type': 'session', 'session_id': resolved_id})}\n\n" - - lock = await _get_session_lock(resolved_id) - async with lock: - loop = asyncio.get_event_loop() - deadline = loop.time() + timeout - pending_tool_calls: dict[str, dict[str, Any]] = {} - emitted_tool_calls: set[str] = set() - - def buffer_function_call(item: Any) -> tuple[str | None, dict[str, Any]]: - event = _function_call_event(item) - call_id = event.get("tool_call_id") - if not isinstance(call_id, str) or not call_id: - return None, event - - pending = pending_tool_calls.setdefault( - call_id, - { - "type": "tool_start", - "tool_call_id": call_id, - "tool_name": event.get("tool_name"), - "arguments": None, - }, - ) - if event.get("tool_name"): - pending["tool_name"] = event["tool_name"] - pending["arguments"] = _merge_tool_arguments( - pending.get("arguments"), - event.get("arguments"), - ) - return call_id, pending - - async def emit_tool_start_if_ready( - call_id: str, event: dict[str, Any] - ) -> AsyncIterator[str]: - if call_id in emitted_tool_calls: - return - if not _is_complete_json_argument(event.get("arguments")): - return - emitted_tool_calls.add(call_id) - yield f"data: {json.dumps(event)}\n\n" - - async def emit_tool_start_before_result(call_id: str | None) -> AsyncIterator[str]: - if call_id is None or call_id in emitted_tool_calls: - return - event = pending_tool_calls.get(call_id) - if event is None: - return - emitted_tool_calls.add(call_id) - yield f"data: {json.dumps(event)}\n\n" + # Set user context for OBO-enabled MCP servers + context_token = set_current_user_context(user_context) + try: try: - stream = agent.run( - prompt, - stream=True, - session=session, - options=_build_chat_options_from_environment(), + agent, session, resolved_id = await _build_agent_session_history( + instructions=instructions, + session_id=session_id, + tools=tools, + mcp_tools=mcp_tools, + skill_paths=skill_paths, + model=model, + sandbox_tools=sandbox_tools, ) - async for update in stream: - if loop.time() > deadline: - yield f"data: {json.dumps({'type': 'error', 'content': f'Timeout after {timeout}s'})}\n\n" - return - for item in getattr(update, "contents", None) or []: - ctype = _content_type(item) - if ctype == "text": - text = _content_text(item) - if text: - yield f"data: {json.dumps({'type': 'delta', 'content': text})}\n\n" - elif ctype == "text_reasoning": - text = _content_text(item) - if text: - yield f"data: {json.dumps({'type': 'intermediate', 'content': text})}\n\n" - elif ctype == "function_call": - call_id, event = buffer_function_call(item) - if call_id is None: - yield f"data: {json.dumps(event)}\n\n" - else: - async for output in emit_tool_start_if_ready(call_id, event): - yield output - elif ctype == "function_result": - call_id = getattr(item, "call_id", None) or getattr(item, "id", None) - async for output in emit_tool_start_before_result( - call_id if isinstance(call_id, str) else None - ): - yield output - yield f"data: {json.dumps(_function_result_event(item), default=str)}\n\n" - # Unknown content types are intentionally ignored — the - # SSE vocabulary is fixed and the UI doesn't render them. - for call_id, event in pending_tool_calls.items(): - if call_id not in emitted_tool_calls: - emitted_tool_calls.add(call_id) - yield f"data: {json.dumps(event)}\n\n" - yield f"data: {json.dumps({'type': 'done'})}\n\n" - except TimeoutError: - yield f"data: {json.dumps({'type': 'error', 'content': f'Timeout after {timeout}s'})}\n\n" except Exception as exc: - logger.error("Agent stream failed: %s", exc, exc_info=True) + logger.error("Failed to build agent session: %s", exc, exc_info=True) yield f"data: {json.dumps({'type': 'error', 'content': str(exc)})}\n\n" + return + + yield f"data: {json.dumps({'type': 'session', 'session_id': resolved_id})}\n\n" + + lock = await _get_session_lock(resolved_id) + async with lock: + loop = asyncio.get_event_loop() + deadline = loop.time() + timeout + pending_tool_calls: dict[str, dict[str, Any]] = {} + emitted_tool_calls: set[str] = set() + + def buffer_function_call(item: Any) -> tuple[str | None, dict[str, Any]]: + event = _function_call_event(item) + call_id = event.get("tool_call_id") + if not isinstance(call_id, str) or not call_id: + return None, event + + pending = pending_tool_calls.setdefault( + call_id, + { + "type": "tool_start", + "tool_call_id": call_id, + "tool_name": event.get("tool_name"), + "arguments": None, + }, + ) + if event.get("tool_name"): + pending["tool_name"] = event["tool_name"] + pending["arguments"] = _merge_tool_arguments( + pending.get("arguments"), + event.get("arguments"), + ) + return call_id, pending + + async def emit_tool_start_if_ready( + call_id: str, event: dict[str, Any] + ) -> AsyncIterator[str]: + if call_id in emitted_tool_calls: + return + if not _is_complete_json_argument(event.get("arguments")): + return + emitted_tool_calls.add(call_id) + yield f"data: {json.dumps(event)}\n\n" + + async def emit_tool_start_before_result(call_id: str | None) -> AsyncIterator[str]: + if call_id is None or call_id in emitted_tool_calls: + return + event = pending_tool_calls.get(call_id) + if event is None: + return + emitted_tool_calls.add(call_id) + yield f"data: {json.dumps(event)}\n\n" + + try: + stream = agent.run( + prompt, + stream=True, + session=session, + options=_build_chat_options_from_environment(), + ) + async for update in stream: + if loop.time() > deadline: + yield f"data: {json.dumps({'type': 'error', 'content': f'Timeout after {timeout}s'})}\n\n" + return + for item in getattr(update, "contents", None) or []: + ctype = _content_type(item) + if ctype == "text": + text = _content_text(item) + if text: + yield f"data: {json.dumps({'type': 'delta', 'content': text})}\n\n" + elif ctype == "text_reasoning": + text = _content_text(item) + if text: + yield f"data: {json.dumps({'type': 'intermediate', 'content': text})}\n\n" + elif ctype == "function_call": + call_id, event = buffer_function_call(item) + if call_id is None: + yield f"data: {json.dumps(event)}\n\n" + else: + async for output in emit_tool_start_if_ready(call_id, event): + yield output + elif ctype == "function_result": + call_id = getattr(item, "call_id", None) or getattr(item, "id", None) + async for output in emit_tool_start_before_result( + call_id if isinstance(call_id, str) else None + ): + yield output + yield f"data: {json.dumps(_function_result_event(item), default=str)}\n\n" + # Unknown content types are intentionally ignored — the + # SSE vocabulary is fixed and the UI doesn't render them. + for call_id, event in pending_tool_calls.items(): + if call_id not in emitted_tool_calls: + emitted_tool_calls.add(call_id) + yield f"data: {json.dumps(event)}\n\n" + yield f"data: {json.dumps({'type': 'done'})}\n\n" + except TimeoutError: + yield f"data: {json.dumps({'type': 'error', 'content': f'Timeout after {timeout}s'})}\n\n" + except Exception as exc: + logger.error("Agent stream failed: %s", exc, exc_info=True) + yield f"data: {json.dumps({'type': 'error', 'content': str(exc)})}\n\n" + finally: + # Reset user context after streaming completes + reset_current_user_context(context_token) diff --git a/tests/test_config_loader.py b/tests/test_config_loader.py index 88463568..760f3ec7 100644 --- a/tests/test_config_loader.py +++ b/tests/test_config_loader.py @@ -61,6 +61,7 @@ def test_load_global_config_leaves_unset_placeholders_literal(tmp_path: Path) -> def test_load_global_config_missing_returns_empty(tmp_path: Path) -> None: assert load_global_config(tmp_path) == load_global_config(tmp_path) assert load_global_config(tmp_path).model_dump() == { + "auth": None, "system_tools": None, "model": None, "timeout": None, diff --git a/tests/test_discovery_mcp.py b/tests/test_discovery_mcp.py index 49c49cb8..83a3528a 100644 --- a/tests/test_discovery_mcp.py +++ b/tests/test_discovery_mcp.py @@ -4,11 +4,16 @@ import logging from pathlib import Path from types import SimpleNamespace +from unittest.mock import MagicMock import pytest from agent_framework import MCPStreamableHTTPTool import azure_functions_agents.discovery.mcp as mcp_discovery +from azure_functions_agents._obo import ( + BIGMAC_ACCESS_TOKEN_HEADER, + BIGMAC_HOOKS_SESSION_TOKEN_HEADER, +) from azure_functions_agents.discovery.mcp import clear_mcp_cache, discover_mcp_servers @@ -378,6 +383,218 @@ def get_token(self, scope: str) -> SimpleNamespace: assert credential.calls == 1 +def test_discover_mcp_servers_obo_bigmac_flow_uses_mi_and_passthrough_headers( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + class FakeCredential: + def __init__(self) -> None: + self.calls = 0 + + def get_token(self, scope: str) -> SimpleNamespace: + self.calls += 1 + assert scope == "https://apihub.azure.com/.default" + return SimpleNamespace(token=f"mi-token-{self.calls}", expires_on=9999999999) + + credential = FakeCredential() + monkeypatch.setattr(mcp_discovery, "build_credential", lambda: credential) + monkeypatch.setattr( + mcp_discovery, "MCPStreamableHTTPTool", _CapturedMCPStreamableHTTPTool + ) + monkeypatch.setattr( + mcp_discovery, + "get_current_user_context", + lambda: SimpleNamespace( + access_token="user-access-token", + hooks_session_token="hooks-session-token", + has_obo_support=True, + ), + ) + _write_mcp_json( + tmp_path, + { + "servers": { + "office365": { + "type": "http", + "url": "https://example.com/mcp", + "headers": {"X-Test": "yes"}, + "auth": { + "type": "obo", + "scope": "https://apihub.azure.com/.default", + }, + } + } + }, + ) + + tool = discover_mcp_servers(tmp_path)["office365"] + + assert isinstance(tool, _CapturedMCPStreamableHTTPTool) + assert tool.header_provider is not None + assert tool.header_provider(None) == { + "Authorization": "Bearer mi-token-1", + BIGMAC_ACCESS_TOKEN_HEADER: "user-access-token", + BIGMAC_HOOKS_SESSION_TOKEN_HEADER: "hooks-session-token", + "X-Test": "yes", + } + assert credential.calls == 1 + + +def test_discover_mcp_servers_obo_bigmac_flow_forwards_whitelisted_headers( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + class FakeCredential: + def __init__(self) -> None: + self.calls = 0 + + def get_token(self, scope: str) -> SimpleNamespace: + self.calls += 1 + return SimpleNamespace(token=f"mi-token-{self.calls}", expires_on=9999999999) + + credential = FakeCredential() + monkeypatch.setattr(mcp_discovery, "build_credential", lambda: credential) + monkeypatch.setattr( + mcp_discovery, "MCPStreamableHTTPTool", _CapturedMCPStreamableHTTPTool + ) + monkeypatch.setattr( + mcp_discovery, + "get_current_user_context", + lambda: SimpleNamespace( + access_token="user-access-token", + hooks_session_token="hooks-session-token", + has_obo_support=True, + forwardable_headers={ + "X-Custom-One": "custom-1", + "X-Custom-Two": "custom-2", + }, + ), + ) + _write_mcp_json( + tmp_path, + { + "servers": { + "office365": { + "type": "http", + "url": "https://example.com/mcp", + "auth": { + "type": "obo", + "scope": "https://apihub.azure.com/.default", + }, + } + } + }, + ) + + tool = discover_mcp_servers(tmp_path)["office365"] + + assert isinstance(tool, _CapturedMCPStreamableHTTPTool) + assert tool.header_provider is not None + assert tool.header_provider(None) == { + "Authorization": "Bearer mi-token-1", + "X-Custom-One": "custom-1", + "X-Custom-Two": "custom-2", + } + assert credential.calls == 1 + + +def test_discover_mcp_servers_obo_forwards_whitelist_without_tokens( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + class FakeCredential: + def __init__(self) -> None: + self.calls = 0 + + def get_token(self, scope: str) -> SimpleNamespace: + self.calls += 1 + return SimpleNamespace(token=f"mi-token-{self.calls}", expires_on=9999999999) + + credential = FakeCredential() + monkeypatch.setattr(mcp_discovery, "build_credential", lambda: credential) + monkeypatch.setattr( + mcp_discovery, "MCPStreamableHTTPTool", _CapturedMCPStreamableHTTPTool + ) + monkeypatch.setattr( + mcp_discovery, + "get_current_user_context", + lambda: SimpleNamespace( + access_token=None, + hooks_session_token=None, + has_obo_support=False, + forwardable_headers={"X-Custom": "custom-value"}, + ), + ) + _write_mcp_json( + tmp_path, + { + "servers": { + "office365": { + "type": "http", + "url": "https://example.com/mcp", + "auth": { + "type": "obo", + "scope": "https://apihub.azure.com/.default", + }, + } + } + }, + ) + + tool = discover_mcp_servers(tmp_path)["office365"] + + assert isinstance(tool, _CapturedMCPStreamableHTTPTool) + assert tool.header_provider is not None + assert tool.header_provider(None) == { + "Authorization": "Bearer mi-token-1", + "X-Custom": "custom-value", + } + assert credential.calls == 1 + + +def test_discover_mcp_servers_obo_without_hooks_uses_obo_token( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + class FakeCredential: + def get_token(self, scope: str) -> SimpleNamespace: + return SimpleNamespace(token="mi-token", expires_on=9999999999) + + monkeypatch.setattr(mcp_discovery, "build_credential", lambda: FakeCredential()) + monkeypatch.setattr( + mcp_discovery, "MCPStreamableHTTPTool", _CapturedMCPStreamableHTTPTool + ) + user_context = SimpleNamespace( + access_token="user-access-token", + hooks_session_token=None, + has_obo_support=True, + get_token_for_scope=MagicMock(return_value=None), + ) + monkeypatch.setattr(mcp_discovery, "get_current_user_context", lambda: user_context) + monkeypatch.setattr( + mcp_discovery, + "_get_obo_token_sync", + lambda user_context, scope: "obo-downstream-token", + ) + _write_mcp_json( + tmp_path, + { + "servers": { + "office365": { + "type": "http", + "url": "https://example.com/mcp", + "auth": { + "type": "obo", + "scope": "https://apihub.azure.com/.default", + }, + } + } + }, + ) + + tool = discover_mcp_servers(tmp_path)["office365"] + + assert isinstance(tool, _CapturedMCPStreamableHTTPTool) + assert tool.header_provider is not None + assert tool.header_provider(None) == {"Authorization": "Bearer obo-downstream-token"} + + def test_discover_mcp_servers_auth_without_scope_uses_static_headers( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture ) -> None: diff --git a/tests/test_obo.py b/tests/test_obo.py new file mode 100644 index 00000000..b1a53e51 --- /dev/null +++ b/tests/test_obo.py @@ -0,0 +1,529 @@ +"""Tests for On-Behalf-Of (OBO) authentication support.""" + +from __future__ import annotations + +from typing import Any +from unittest.mock import MagicMock, patch + +import pytest + +from azure_functions_agents._obo import ( + AGENTS_HEADER_WHITELIST_ENV, + BIGMAC_ACCESS_TOKEN_HEADER, + BIGMAC_HOOKS_SESSION_TOKEN_HEADER, + EASYAUTH_ID_TOKEN_HEADER, + InteractionRequiredError, + OboError, + OboTokenProvider, + clear_token_cache, + create_user_context, + extract_forwardable_headers, + extract_hooks_session_token_from_headers, + extract_user_id_from_headers, + extract_user_token_from_headers, + get_forwarded_header_names, + get_obo_provider, + reset_obo_provider, +) +from azure_functions_agents.config.schema import AuthConfig, OboConfig + +# --------------------------------------------------------------------------- +# OboConfig tests +# --------------------------------------------------------------------------- + + +class TestOboConfig: + """Tests for OboConfig schema validation.""" + + def test_valid_obo_config(self) -> None: + """OboConfig should accept valid configuration.""" + config = OboConfig( + enabled=True, + client_id="test-client-id", + client_secret="test-secret", + tenant_id="test-tenant-id", + downstream_scopes={"graph": "https://graph.microsoft.com/.default"}, + ) + assert config.enabled is True + assert config.client_id == "test-client-id" + assert config.client_secret == "test-secret" + assert config.tenant_id == "test-tenant-id" + assert config.downstream_scopes == {"graph": "https://graph.microsoft.com/.default"} + + def test_obo_config_defaults(self) -> None: + """OboConfig should have sensible defaults.""" + config = OboConfig( + client_id="test-client-id", + tenant_id="test-tenant-id", + ) + assert config.enabled is True + assert config.client_secret is None + assert config.downstream_scopes == {} + + def test_obo_config_empty_client_id_rejected(self) -> None: + """OboConfig should reject empty client_id.""" + with pytest.raises(ValueError, match="value must be non-empty"): + OboConfig( + client_id=" ", + tenant_id="test-tenant-id", + ) + + def test_obo_config_empty_tenant_id_rejected(self) -> None: + """OboConfig should reject empty tenant_id.""" + with pytest.raises(ValueError, match="value must be non-empty"): + OboConfig( + client_id="test-client-id", + tenant_id="", + ) + + +class TestAuthConfig: + """Tests for AuthConfig schema.""" + + def test_auth_config_with_obo(self) -> None: + """AuthConfig should wrap OboConfig.""" + obo = OboConfig( + client_id="test-client-id", + tenant_id="test-tenant-id", + ) + config = AuthConfig(obo=obo) + assert config.obo is not None + assert config.obo.client_id == "test-client-id" + + def test_auth_config_without_obo(self) -> None: + """AuthConfig should allow None obo.""" + config = AuthConfig() + assert config.obo is None + + +# --------------------------------------------------------------------------- +# Header extraction tests +# --------------------------------------------------------------------------- + + +class TestHeaderExtraction: + """Tests for token extraction from headers.""" + + def test_extract_from_easyauth_header(self) -> None: + """Should extract token from EasyAuth header.""" + headers = {"X-MS-TOKEN-AAD-ACCESS-TOKEN": "test-token-123"} + token = extract_user_token_from_headers(headers) + assert token == "test-token-123" + + def test_extract_from_bigmac_access_token_header(self) -> None: + """Should extract token from BigMac access token header.""" + headers = {BIGMAC_ACCESS_TOKEN_HEADER: "bigmac-token-123"} + token = extract_user_token_from_headers(headers) + assert token == "bigmac-token-123" + + def test_extract_from_authorization_header(self) -> None: + """Should extract token from Authorization header.""" + headers = {"Authorization": "Bearer test-token-456"} + token = extract_user_token_from_headers(headers) + assert token == "test-token-456" + + def test_easyauth_takes_precedence(self) -> None: + """EasyAuth header should take precedence over Authorization.""" + headers = { + "X-MS-TOKEN-AAD-ACCESS-TOKEN": "easyauth-token", + "Authorization": "Bearer auth-token", + } + token = extract_user_token_from_headers(headers) + assert token == "easyauth-token" + + def test_bigmac_takes_precedence(self) -> None: + """BigMac access token header should take precedence over other token headers.""" + headers = { + BIGMAC_ACCESS_TOKEN_HEADER: "bigmac-token", + "X-MS-TOKEN-AAD-ACCESS-TOKEN": "easyauth-token", + "Authorization": "Bearer auth-token", + } + token = extract_user_token_from_headers(headers) + assert token == "bigmac-token" + + def test_extract_id_token_fallback(self) -> None: + """Should fall back to EasyAuth ID token when access token is missing.""" + headers = {EASYAUTH_ID_TOKEN_HEADER: "id-token-value"} + token = extract_user_token_from_headers(headers) + assert token == "id-token-value" + + def test_no_token_returns_none(self) -> None: + """Should return None when no token is present.""" + headers = {"Content-Type": "application/json"} + token = extract_user_token_from_headers(headers) + assert token is None + + def test_case_insensitive_header_lookup(self) -> None: + """Should handle case-insensitive header names.""" + headers = {"x-ms-token-aad-access-token": "lower-case-token"} + token = extract_user_token_from_headers(headers) + assert token == "lower-case-token" + + def test_extract_user_id_from_headers(self) -> None: + """Should extract user ID from EasyAuth principal header.""" + headers = {"X-MS-CLIENT-PRINCIPAL-ID": "user-object-id-123"} + user_id = extract_user_id_from_headers(headers) + assert user_id == "user-object-id-123" + + def test_extract_hooks_session_token_from_headers(self) -> None: + """Should extract hooks session token from BigMac header.""" + headers = {BIGMAC_HOOKS_SESSION_TOKEN_HEADER: "hooks-session-123"} + token = extract_hooks_session_token_from_headers(headers) + assert token == "hooks-session-123" + + def test_extract_hooks_session_token_case_insensitive(self) -> None: + """Should extract hooks session token with case-insensitive lookup.""" + headers = {"x-ms-hooks-session-token": "hooks-session-456"} + token = extract_hooks_session_token_from_headers(headers) + assert token == "hooks-session-456" + + def test_user_id_missing_returns_none(self) -> None: + """Should return None when user ID header is missing.""" + headers = {} + user_id = extract_user_id_from_headers(headers) + assert user_id is None + + +# --------------------------------------------------------------------------- +# Header whitelist tests +# --------------------------------------------------------------------------- + + +class TestHeaderWhitelist: + """Tests for the AGENTS_HEADER_WHITELIST forwarding behavior.""" + + def test_names_empty_when_env_missing(self, monkeypatch: pytest.MonkeyPatch) -> None: + """Should return an empty list when the env var is unset.""" + monkeypatch.delenv(AGENTS_HEADER_WHITELIST_ENV, raising=False) + assert get_forwarded_header_names() == [] + + def test_names_empty_when_env_blank(self, monkeypatch: pytest.MonkeyPatch) -> None: + """Should return an empty list when the env var has no header names.""" + monkeypatch.setenv(AGENTS_HEADER_WHITELIST_ENV, " , ,") + assert get_forwarded_header_names() == [] + + def test_parses_custom_whitelist(self, monkeypatch: pytest.MonkeyPatch) -> None: + """Should parse a comma-separated whitelist and trim whitespace.""" + monkeypatch.setenv(AGENTS_HEADER_WHITELIST_ENV, "X-Custom-One , X-Custom-Two") + assert get_forwarded_header_names() == ["X-Custom-One", "X-Custom-Two"] + + def test_extract_empty_without_whitelist( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Should return no forwardable headers when no whitelist is configured.""" + monkeypatch.delenv(AGENTS_HEADER_WHITELIST_ENV, raising=False) + headers = {"X-MS-TOKEN-AAD-ACCESS-TOKEN": "easyauth-token"} + assert extract_forwardable_headers(headers) == {} + + def test_extract_custom_reads_request_headers( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Custom whitelist should forward matching inbound headers case-insensitively.""" + monkeypatch.setenv(AGENTS_HEADER_WHITELIST_ENV, "X-Custom-One,X-Custom-Two") + headers = {"x-custom-one": "value-1", "X-Custom-Two": "value-2"} + assert extract_forwardable_headers(headers) == { + "X-Custom-One": "value-1", + "X-Custom-Two": "value-2", + } + + def test_extract_custom_omits_absent_headers( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Custom whitelist should skip headers not present on the request.""" + monkeypatch.setenv(AGENTS_HEADER_WHITELIST_ENV, "X-Present,X-Absent") + headers = {"X-Present": "here"} + assert extract_forwardable_headers(headers) == {"X-Present": "here"} + + def test_create_user_context_stores_forwardable_headers(self) -> None: + """UserContext should retain the resolved forwardable headers.""" + context = create_user_context( + access_token="t", + forwardable_headers={"X-Custom": "v"}, + ) + assert context.forwardable_headers == {"X-Custom": "v"} + + +# --------------------------------------------------------------------------- +# UserContext tests +# --------------------------------------------------------------------------- + + +class TestUserContext: + """Tests for UserContext creation and behavior.""" + + def test_create_user_context_with_token(self) -> None: + """Should create context with access token.""" + context = create_user_context( + access_token="test-token", + hooks_session_token="hooks-token", + user_id="user-123", + ) + assert context.access_token == "test-token" + assert context.hooks_session_token == "hooks-token" + assert context.user_id == "user-123" + assert context.is_authenticated is True + + def test_create_user_context_without_token(self) -> None: + """Should create context without access token.""" + context = create_user_context() + assert context.access_token is None + assert context.is_authenticated is False + + def test_has_obo_support_requires_both(self) -> None: + """has_obo_support should require both token and provider.""" + # No token, no provider + context = create_user_context() + assert context.has_obo_support is False + + # Token but no provider + context = create_user_context(access_token="test-token") + assert context.has_obo_support is False + + # Both token and provider + mock_provider = MagicMock() + context = create_user_context( + access_token="test-token", + obo_provider=mock_provider, + ) + assert context.has_obo_support is True + + +# --------------------------------------------------------------------------- +# OboTokenProvider tests +# --------------------------------------------------------------------------- + + +class TestOboTokenProvider: + """Tests for OBO token exchange.""" + + def setup_method(self) -> None: + """Clear token cache before each test.""" + clear_token_cache() + + @pytest.fixture + def obo_config(self) -> OboConfig: + """Create a test OBO config.""" + return OboConfig( + client_id="test-client-id", + client_secret="test-secret", + tenant_id="test-tenant-id", + downstream_scopes={"graph": "https://graph.microsoft.com/.default"}, + ) + + @pytest.fixture + def provider(self, obo_config: OboConfig) -> OboTokenProvider: + """Create a test OBO provider.""" + return OboTokenProvider(obo_config) + + def test_get_scope_for_name(self, provider: OboTokenProvider) -> None: + """Should return scope for configured name.""" + scope = provider.get_scope_for_name("graph") + assert scope == "https://graph.microsoft.com/.default" + + def test_get_scope_for_unknown_name(self, provider: OboTokenProvider) -> None: + """Should return None for unknown scope name.""" + scope = provider.get_scope_for_name("unknown") + assert scope is None + + @pytest.mark.asyncio + async def test_acquire_token_success(self, provider: OboTokenProvider) -> None: + """Should successfully acquire token via OBO.""" + mock_result = { + "access_token": "downstream-token", + "expires_in": 3600, + } + + with patch.object(provider, "_get_msal_app") as mock_get_app: + mock_app = MagicMock() + mock_app.acquire_token_on_behalf_of.return_value = mock_result + mock_get_app.return_value = mock_app + + token = await provider.acquire_token_on_behalf_of( + user_token="user-token", + scope="https://api.example.com/.default", + ) + + assert token == "downstream-token" + + @pytest.mark.asyncio + async def test_acquire_token_interaction_required(self, provider: OboTokenProvider) -> None: + """Should raise InteractionRequiredError when consent needed.""" + mock_result = { + "error": "interaction_required", + "error_description": "User consent required", + "claims": '{"claim": "value"}', + } + + with patch.object(provider, "_get_msal_app") as mock_get_app: + mock_app = MagicMock() + mock_app.acquire_token_on_behalf_of.return_value = mock_result + mock_get_app.return_value = mock_app + + with pytest.raises(InteractionRequiredError) as exc_info: + await provider.acquire_token_on_behalf_of( + user_token="user-token", + scope="https://api.example.com/.default", + ) + + assert exc_info.value.error == "interaction_required" + assert exc_info.value.error_description == "User consent required" + + @pytest.mark.asyncio + async def test_acquire_token_generic_error(self, provider: OboTokenProvider) -> None: + """Should raise OboError for other failures.""" + mock_result = { + "error": "invalid_grant", + "error_description": "Token expired", + } + + with patch.object(provider, "_get_msal_app") as mock_get_app: + mock_app = MagicMock() + mock_app.acquire_token_on_behalf_of.return_value = mock_result + mock_get_app.return_value = mock_app + + with pytest.raises(OboError) as exc_info: + await provider.acquire_token_on_behalf_of( + user_token="user-token", + scope="https://api.example.com/.default", + ) + + assert exc_info.value.error == "invalid_grant" + + +# --------------------------------------------------------------------------- +# Token caching tests +# --------------------------------------------------------------------------- + + +class TestTokenCaching: + """Tests for token caching behavior.""" + + def setup_method(self) -> None: + """Clear token cache before each test.""" + clear_token_cache() + + @pytest.fixture + def obo_config(self) -> OboConfig: + """Create a test OBO config.""" + return OboConfig( + client_id="test-client-id", + client_secret="test-secret", + tenant_id="test-tenant-id", + ) + + @pytest.mark.asyncio + async def test_token_is_cached(self, obo_config: OboConfig) -> None: + """Second call should use cached token.""" + provider = OboTokenProvider(obo_config) + call_count = 0 + + def mock_acquire(*args: Any, **kwargs: Any) -> dict[str, Any]: + nonlocal call_count + call_count += 1 + return { + "access_token": f"token-{call_count}", + "expires_in": 3600, + } + + with patch.object(provider, "_get_msal_app") as mock_get_app: + mock_app = MagicMock() + mock_app.acquire_token_on_behalf_of.side_effect = mock_acquire + mock_get_app.return_value = mock_app + + # First call + token1 = await provider.acquire_token_on_behalf_of( + user_token="user-token", + scope="https://api.example.com/.default", + ) + + # Second call should use cache + token2 = await provider.acquire_token_on_behalf_of( + user_token="user-token", + scope="https://api.example.com/.default", + ) + + assert token1 == "token-1" + assert token2 == "token-1" # Same cached token + assert call_count == 1 # Only one MSAL call + + +# --------------------------------------------------------------------------- +# Error types tests +# --------------------------------------------------------------------------- + + +class TestErrorTypes: + """Tests for OBO error types.""" + + def test_obo_error_message(self) -> None: + """OboError should format error message correctly.""" + error = OboError("invalid_grant", "Token has expired") + assert str(error) == "invalid_grant: Token has expired" + + def test_obo_error_without_description(self) -> None: + """OboError should handle missing description.""" + error = OboError("unknown_error") + assert str(error) == "unknown_error" + + def test_interaction_required_error(self) -> None: + """InteractionRequiredError should carry claims.""" + error = InteractionRequiredError( + "consent_required", + "Admin consent needed", + '{"claims": "data"}', + ) + assert error.error == "consent_required" + assert error.error_description == "Admin consent needed" + assert error.claims == '{"claims": "data"}' + + +# --------------------------------------------------------------------------- +# Global provider tests +# --------------------------------------------------------------------------- + + +class TestGlobalProvider: + """Tests for global OBO provider management.""" + + def setup_method(self) -> None: + """Reset global provider before each test.""" + reset_obo_provider() + + @pytest.mark.asyncio + async def test_get_obo_provider_creates_instance(self) -> None: + """get_obo_provider should create provider when config is valid.""" + config = OboConfig( + client_id="test-client-id", + tenant_id="test-tenant-id", + ) + provider = await get_obo_provider(config) + assert provider is not None + assert isinstance(provider, OboTokenProvider) + + @pytest.mark.asyncio + async def test_get_obo_provider_returns_none_when_disabled(self) -> None: + """get_obo_provider should return None when OBO is disabled.""" + config = OboConfig( + enabled=False, + client_id="test-client-id", + tenant_id="test-tenant-id", + ) + provider = await get_obo_provider(config) + assert provider is None + + @pytest.mark.asyncio + async def test_get_obo_provider_returns_none_for_none_config(self) -> None: + """get_obo_provider should return None for None config.""" + provider = await get_obo_provider(None) + assert provider is None + + @pytest.mark.asyncio + async def test_get_obo_provider_returns_same_instance(self) -> None: + """get_obo_provider should return the same instance on subsequent calls.""" + config = OboConfig( + client_id="test-client-id", + tenant_id="test-tenant-id", + ) + provider1 = await get_obo_provider(config) + provider2 = await get_obo_provider(config) + assert provider1 is provider2