diff --git a/README.md b/README.md index ef8383f..a83bb16 100644 --- a/README.md +++ b/README.md @@ -7,9 +7,12 @@ Sample clients for the [Work IQ](https://learn.microsoft.com/en-us/microsoft-365 | [**dotnet/a2a/**](dotnet/a2a/) | C# | Windows, macOS, Linux | [A2A](https://a2a-protocol.org) | Interactive agent session using the A2A protocol over JSON-RPC | | [**dotnet/a2a-raw/**](dotnet/a2a-raw/) | C# | Windows, macOS, Linux | [A2A](https://a2a-protocol.org) | Same, but with raw `HttpClient` + JSON (no A2A SDK) | | [**dotnet/rest/**](dotnet/rest/) | C# | Windows, macOS, Linux | REST | Interactive chat using the [Copilot Chat API](https://learn.microsoft.com/en-us/microsoft-365-copilot/extensibility/api/ai-services/chat/overview) | +| [**python/obo/**](python/obo/) | Python | Windows, macOS, Linux | REST | Backend service brokering Work IQ for a frontend via the [On-Behalf-Of flow](https://learn.microsoft.com/en-us/entra/identity-platform/v2-oauth2-on-behalf-of-flow) | | [**rust/a2a/**](rust/a2a/) | Rust | Windows, macOS, Linux | [A2A](https://a2a-protocol.org) | Interactive agent session with MSAL auth and token caching | | [**swift/a2a/**](swift/a2a/) | Swift | iOS/iPadOS (macOS to build) | [A2A](https://a2a-protocol.org) | SwiftUI chat app for Work IQ | +Every sample is a **public client** that signs a user in directly, except [**python/obo/**](python/obo/) — that one is a **middle tier**: the user signs in to your frontend, and the service exchanges their token for a Work IQ token. It needs a different app registration; see its [README](python/obo/README.md#app-registration). + --- ## Gateway @@ -24,6 +27,7 @@ All samples target the **Work IQ Gateway** (`workiq.svc.cloud.microsoft`) — th 2. **Entra app registration** configured in your tenant — this is a one-time setup per tenant. Details below. 3. **Your language toolchain**: - **dotnet/** samples: [.NET 8.0 SDK](https://dotnet.microsoft.com/download/dotnet/8.0) or later + - **python/** samples: [Python 3.10+](https://www.python.org/downloads/) - **rust/** samples: [Rust toolchain](https://rustup.rs/) (stable) - **swift/** samples: [Xcode 26+](https://developer.apple.com/xcode/) (macOS only) @@ -46,6 +50,8 @@ You (or your tenant admin) must create an Entra app registration with specific p After setup you'll have two values: `APP_ID` and `TENANT_ID`. Pass them to any sample via `--appid` and `--tenant`. +> **[`python/obo/`](python/obo/) is the exception.** The script above creates a *public* client; the On-Behalf-Of flow requires a *confidential* client that also exposes its own API. See [its README](python/obo/README.md#app-registration) for the separate setup. + --- ## Authentication methods diff --git a/python/obo/.env.example b/python/obo/.env.example new file mode 100644 index 0000000..2a1313d --- /dev/null +++ b/python/obo/.env.example @@ -0,0 +1,18 @@ +# Backend app registration (confidential client). +AZURE_TENANT_ID=00000000-0000-0000-0000-000000000000 +AZURE_CLIENT_ID=00000000-0000-0000-0000-000000000000 + +# Application ID URI this backend's tokens are issued for. The frontend must +# request `${API_AUDIENCE}/access_as_user`. +API_AUDIENCE=api://00000000-0000-0000-0000-000000000000 + +# Scope the inbound token must carry. Defaults to access_as_user. +# REQUIRED_SCOPE=access_as_user + +# Leave AZURE_CLIENT_SECRET unset in Azure: the app then authenticates with a +# managed identity federated credential via DefaultAzureCredential (no secret). +# Set it only for local development. +# AZURE_CLIENT_SECRET= + +# Override the gateway host for testing. Defaults to the production gateway. +# WORKIQ_HOST=https://workiq.svc.cloud.microsoft diff --git a/python/obo/.gitignore b/python/obo/.gitignore new file mode 100644 index 0000000..080dec5 --- /dev/null +++ b/python/obo/.gitignore @@ -0,0 +1,3 @@ +.venv/ +.env +*.env diff --git a/python/obo/README.md b/python/obo/README.md new file mode 100644 index 0000000..ba52375 --- /dev/null +++ b/python/obo/README.md @@ -0,0 +1,222 @@ +# Work IQ Python OBO Sample + +A FastAPI **middle-tier** service that accepts your frontend's access token, exchanges it for a Work IQ token via the [On-Behalf-Of flow](https://learn.microsoft.com/en-us/entra/identity-platform/v2-oauth2-on-behalf-of-flow), and calls the [Copilot Chat REST API](https://learn.microsoft.com/en-us/microsoft-365-copilot/extensibility/api/ai-services/chat/overview) through the **Work IQ Gateway**. + +Every other sample in this repo is a **public client** that signs a user in directly. This one is different: the user signs in to *your* frontend, and this service brokers the Work IQ call for them. + +``` +Frontend ──token(aud=your API)──▶ This service ──token(aud=Work IQ)──▶ Gateway + │ + └── validate ──▶ OBO exchange +``` + +Supports both **synchronous** and **streaming** (SSE) modes. + +## API reference + +| Operation | Method | Path | Docs | +|-----------|--------|------|------| +| Create conversation | `POST` | `/rest/beta/conversations` | [Docs](https://learn.microsoft.com/en-us/microsoft-365-copilot/extensibility/api/ai-services/chat/copilotroot-post-conversations) | +| Chat (sync) | `POST` | `/rest/beta/conversations/{id}/chat` | [Docs](https://learn.microsoft.com/en-us/microsoft-365-copilot/extensibility/api/ai-services/chat/copilotconversation-chat) | +| Chat (stream) | `POST` | `/rest/beta/conversations/{id}/chatOverStream` | [Docs](https://learn.microsoft.com/en-us/microsoft-365-copilot/extensibility/api/ai-services/chat/copilotconversation-chatoverstream) | + +## Why you can't just forward the frontend's token + +A token's audience is baked into the signed JWT. The token your frontend holds has +`aud` set to *your* API — the Gateway rejects it with `401`. You cannot re-point it at +another resource; you exchange it for a new one. That exchange is OBO. + +**`DefaultAzureCredential` cannot do this exchange.** Its chain (managed identity, +environment service principal, Azure CLI) issues tokens for the *app's own* identity or +a *developer's* identity. None of them accept an inbound user token and re-issue it for +a different resource. And `WorkIQAgent.Ask` is delegated-only — an app-only token gets +you nothing, because responses depend on the signed-in user's Copilot license and their +own data access. + +`azure.identity.aio.OnBehalfOfCredential` is the credential that does this. + +`DefaultAzureCredential` still has a real job here: proving the *app's* identity so no +client secret is ever deployed. Leave `AZURE_CLIENT_SECRET` unset and the service uses a +managed identity federated credential as its client assertion — see [`app/auth.py`](app/auth.py): + +```python +def assertion() -> str: + return credential.get_token("api://AzureADTokenExchange/.default").token + +OnBehalfOfCredential( + tenant_id=..., client_id=..., + client_assertion_func=assertion, # app identity — DefaultAzureCredential + user_assertion=inbound_token, # user identity — from your frontend +) +``` + +Both identities are required: the app proves it is allowed to ask, the user token +determines what comes back. + +## Prerequisites + +1. **Microsoft 365 Copilot license** on your test user. +2. **Two Entra app registrations.** `scripts/admin-setup.sh` at the repo root does **not** + cover this sample — it creates a *public* client for the CLI samples, and OBO requires + a *confidential* client. See [App registration](#app-registration) below. +3. **Python 3.10+**. + +## App registration + +**Frontend app** (SPA / public client) — requests `api:///access_as_user`. + +**Backend app** (this service — confidential client): + +| Blade | Setting | +|-------|---------| +| Expose an API | Application ID URI `api://`, scope `access_as_user` | +| API permissions | `Work IQ` → delegated `WorkIQAgent.Ask` → **Grant admin consent** | +| Certificates & secrets | In Azure, prefer a **federated credential** bound to your managed identity over a client secret | + +```bash +# Ensure the Work IQ service principal exists in your tenant +az ad sp create --id fdcc1f02-fc51-4226-8753-f668596af7f7 + +# Grant this service the delegated Work IQ permission, then consent +az ad app permission add --id \ + --api fdcc1f02-fc51-4226-8753-f668596af7f7 \ + --api-permissions "0b1715fd-f4bf-4c63-b16d-5be31f9847c2=Scope" +az ad app permission admin-consent --id +``` + +## Quick start + +```bash +python -m venv .venv && source .venv/bin/activate +pip install -r requirements.txt + +cp .env.example .env # fill in tenant, client id, audience +set -a && source .env && set +a + +uvicorn app.main:app --reload +``` + +```bash +curl -X POST localhost:8000/api/chat \ + -H "Authorization: Bearer /access_as_user>" \ + -H "Content-Type: application/json" \ + -d '{"message": "What meetings do I have tomorrow?", "time_zone": "America/New_York"}' +``` + +```json +{ + "conversation_id": "conv-123", + "text": "You have 3 meetings scheduled...", + "citations": [ + { "type": "citation", "source": "sharepoint", "provider": "Q3.docx", "url": "https://..." } + ] +} +``` + +Pass `conversation_id` back on later turns to continue the same conversation with full context. + +### Verify without credentials + +```bash +python smoke_test.py +``` + +Fakes the Gateway with `httpx.MockTransport` and checks conversation creation, citation +parsing, streaming deltas, the `403` error path, and the auth gate. No tenant, no license, +no network. + +## Endpoints + +| Endpoint | Mode | Response | +|----------|------|----------| +| `GET /healthz` | — | `{"status": "ok"}` (unauthenticated, for liveness probes) | +| `POST /api/chat` | Synchronous | JSON — `conversation_id`, `text`, `citations` | +| `POST /api/chat/stream` | SSE | See streaming contract below | + +**Stream event sequence:** + +1. `event: conversation` — `data: {"conversation_id": "..."}` (first frame) +2. `data: {"text": ""}` — one per text chunk (default event type) +3. `event: done` — signals clean completion, or `event: error` — `data: upstream request failed` on failure + +Request body for both: `{"message": "...", "conversation_id": "...", "time_zone": "..."}` (`conversation_id` and `time_zone` optional; `time_zone` is an IANA identifier like `America/New_York`). + +**Streaming contract:** once the HTTP 200 is committed, errors cannot change the status code. Clients **must** listen for `event: error` frames to detect mid-stream failures. A `event: done` frame signals clean completion; its absence (with no `error`) indicates a dropped connection. + +## Layout + +| File | Purpose | +|------|---------| +| [`app/config.py`](app/config.py) | Environment config; fails fast on missing values | +| [`app/auth.py`](app/auth.py) | Inbound JWT validation + the OBO exchange | +| [`app/workiq.py`](app/workiq.py) | Async Gateway client (sync + streaming) | +| [`app/main.py`](app/main.py) | FastAPI routes | +| [`smoke_test.py`](smoke_test.py) | Gateway faked via `httpx.MockTransport` | + +## How it works + +``` +Frontend This service Gateway + | | | + |-- POST /api/chat --->| | + | Bearer | | + | |-- validate signature/aud/scp | + | | against Entra JWKS | + | | | + | |-- OBO exchange ---▶ Entra | + | |◀-- token(aud=Work IQ) -- | + | | | + | |-- POST .../conversations ----->| + | |<-- 201 { "id": "conv-123" } ---| + | |-- POST .../conv-123/chat ----->| + | |<-- 200 { "messages": [...] } --| + |<-- 200 JSON ---------| | +``` + +Each SSE event from the Gateway contains the **full conversation state** (cumulative, not +incremental). [`app/workiq.py`](app/workiq.py) diffs against the previous event and yields +only new text, so `/api/chat/stream` emits true deltas. + +Auth is resolved before body validation, so an unauthenticated caller gets `401` and +learns nothing about the request schema. + +## Dependencies + +| Package | Purpose | +|---------|---------| +| `azure-identity` | `OnBehalfOfCredential` for the exchange; `DefaultAzureCredential` for the app assertion | +| `pyjwt[crypto]` | Validating the inbound token against Entra's JWKS | +| `httpx` | Async HTTP + SSE against the Gateway | +| `fastapi` / `uvicorn` | The service itself | + +No MSAL wrapper needed — `azure-identity` builds on MSAL underneath. + +## Sample-specific troubleshooting + +| Symptom | Fix | +|---------|-----| +| `401` from this service | Inbound token failed validation. Check `aud` matches `API_AUDIENCE` and `scp` includes `access_as_user`. | +| `403 Unable to obtain Work IQ access` | The OBO exchange failed. Usually missing admin consent on `WorkIQAgent.Ask`, or the user lacks a Copilot license. | +| `401` from the Gateway | The *outbound* token's `aud` is wrong — must be `api://workiq.svc.cloud.microsoft`, not your API. | +| `AADSTS50013: Assertion failed signature validation` | The federated credential subject/issuer doesn't match your managed identity. | +| `502 Work IQ request failed` | The Gateway call failed — network error, timeout, or the Gateway returned a server error. Check connectivity and the `request-id` in logs. | +| Slow first response per turn | Every request does a fresh OBO round trip. See [Notes before production](#notes-before-production). | + +See the [root README](../../README.md#troubleshooting) for the full matrix (Copilot license, consent, audience mismatch). + +## Notes before production + +- **Cache the OBO result.** Each request currently exchanges the token again. Cache on a + hash of the inbound token, honoring `expires_on`, to save a round trip per turn. +- **The `request-id` response header** is what Microsoft support asks for. + [`app/workiq.py`](app/workiq.py) includes it in the `WorkIQError` exception; + [`app/main.py`](app/main.py) logs it when handling errors. +- **Never log tokens.** Errors here are logged with the user's `oid`, not the assertion. + +## Resources + +- [On-Behalf-Of flow](https://learn.microsoft.com/en-us/entra/identity-platform/v2-oauth2-on-behalf-of-flow) +- [Workload identity federation](https://learn.microsoft.com/en-us/entra/workload-id/workload-identity-federation) +- [`azure-identity` for Python](https://learn.microsoft.com/en-us/python/api/overview/azure/identity-readme) +- [Chat API Overview](https://learn.microsoft.com/en-us/microsoft-365-copilot/extensibility/api/ai-services/chat/overview) +- [Work IQ Overview](https://learn.microsoft.com/en-us/microsoft-365/copilot/extensibility/workiq-overview) diff --git a/python/obo/app/__init__.py b/python/obo/app/__init__.py new file mode 100644 index 0000000..3231871 --- /dev/null +++ b/python/obo/app/__init__.py @@ -0,0 +1 @@ +"""Work IQ On-Behalf-Of backend sample.""" diff --git a/python/obo/app/auth.py b/python/obo/app/auth.py new file mode 100644 index 0000000..e5c9d2f --- /dev/null +++ b/python/obo/app/auth.py @@ -0,0 +1,126 @@ +"""Validate the frontend's token, then exchange it for a Work IQ token via OBO. + +Two distinct tokens are involved and they are not interchangeable: + + 1. Inbound — aud = this backend's Application ID URI (API_AUDIENCE). + Issued to the frontend. Work IQ rejects it (401, audience mismatch). + 2. Outbound — aud = api://workiq.svc.cloud.microsoft. Minted here by the + On-Behalf-Of flow, carrying the same user identity. +""" + +from __future__ import annotations + +import asyncio +from typing import Any, Callable + +import jwt +from azure.identity import DefaultAzureCredential as SyncDefaultAzureCredential +from azure.identity.aio import OnBehalfOfCredential +from jwt import PyJWKClient + +from .config import TOKEN_EXCHANGE_SCOPE, WORKIQ_SCOPE, Settings + + +class InvalidToken(Exception): + """The inbound token failed signature, audience, issuer, or scope checks.""" + + +class TokenValidator: + """Validates inbound tokens against Entra's published signing keys. + + PyJWKClient caches keys after the first fetch, so the blocking call it makes + is offloaded to a worker thread rather than stalling the event loop. + """ + + def __init__(self, settings: Settings) -> None: + self._settings = settings + self._jwks_client = PyJWKClient(settings.jwks_uri, lifespan=3600) + + async def validate(self, token: str) -> dict[str, Any]: + try: + signing_key = await asyncio.to_thread( + self._jwks_client.get_signing_key_from_jwt, token + ) + claims: dict[str, Any] = jwt.decode( + token, + signing_key.key, + algorithms=["RS256"], + audience=self._settings.api_audience, + issuer=self._settings.issuer, + options={"require": ["exp", "nbf", "aud", "iss"]}, + ) + except jwt.PyJWTError as exc: + raise InvalidToken(f"token rejected: {type(exc).__name__}") from exc + + self._require_scope(claims) + return claims + + def _require_scope(self, claims: dict[str, Any]) -> None: + if "scp" not in claims: + raise InvalidToken( + "app-only tokens are not accepted; a delegated user token is required" + ) + granted = set(str(claims["scp"]).split()) + if self._settings.required_scope not in granted: + raise InvalidToken( + f"token is missing the required scope '{self._settings.required_scope}'" + ) + + +def _federated_assertion(credential: SyncDefaultAzureCredential) -> Callable[[], str]: + """Build a client assertion from a managed identity token. + + This is where DefaultAzureCredential belongs in an OBO backend: it proves the + *app's* identity so no client secret is needed. It cannot perform the OBO + exchange itself — that needs the user's assertion, which only OnBehalfOfCredential + accepts. The callable is sync because azure-identity invokes it synchronously. + """ + + def assertion() -> str: + return credential.get_token(TOKEN_EXCHANGE_SCOPE).token + + return assertion + + +class WorkIQTokenExchange: + """Mints Work IQ access tokens on behalf of the calling user.""" + + def __init__(self, settings: Settings) -> None: + self._settings = settings + self._mi_credential = ( + SyncDefaultAzureCredential() if settings.uses_managed_identity else None + ) + + async def token_for(self, user_assertion: str) -> str: + """Exchange a validated inbound token for a Work IQ access token. + + A fresh credential per request keeps user identities isolated. The trade-off + is a round trip to Entra on every call — add a cache keyed by a hash of the + assertion if that latency matters. + """ + credential = self._build_credential(user_assertion) + async with credential: + access_token = await credential.get_token(WORKIQ_SCOPE) + return access_token.token + + def _build_credential(self, user_assertion: str) -> OnBehalfOfCredential: + common = { + "tenant_id": self._settings.tenant_id, + "client_id": self._settings.client_id, + "user_assertion": user_assertion, + } + + if self._mi_credential is not None: + return OnBehalfOfCredential( + client_assertion_func=_federated_assertion(self._mi_credential), + **common, + ) + + return OnBehalfOfCredential( + client_secret=self._settings.client_secret, + **common, + ) + + async def close(self) -> None: + if self._mi_credential is not None: + await asyncio.to_thread(self._mi_credential.close) diff --git a/python/obo/app/config.py b/python/obo/app/config.py new file mode 100644 index 0000000..8ae95a4 --- /dev/null +++ b/python/obo/app/config.py @@ -0,0 +1,120 @@ +"""Configuration for the Work IQ OBO backend, loaded from the environment. + +Missing required values fail at startup (when ``get_settings()`` is first called) +rather than on the first request. +""" + +from __future__ import annotations + +import os +from dataclasses import dataclass, field +from functools import lru_cache + +# The Work IQ Gateway multiplexes REST/A2A/MCP on one host behind a path prefix +# (/rest, /a2a, /tools). Only the REST surface is used here. +WORKIQ_DEFAULT_HOST = "https://workiq.svc.cloud.microsoft" +WORKIQ_PATH = "/rest/beta" + +# Allowlist of valid Work IQ Gateway hosts. For staging/test environments, +# set EXTRA_WORKIQ_HOSTS (comma-separated HTTPS URLs) rather than removing +# the check — an unrestricted WORKIQ_HOST could redirect OBO tokens. +_ALLOWED_WORKIQ_HOSTS = frozenset({ + WORKIQ_DEFAULT_HOST, +}) + +# Work IQ resource. WorkIQAgent.Ask is granted by admin consent on the app +# registration, so `.default` resolves to it without naming the scope here. +WORKIQ_SCOPE = "api://workiq.svc.cloud.microsoft/.default" + +# Scope the frontend must present on its token to this backend. +DEFAULT_REQUIRED_SCOPE = "access_as_user" + +ENTRA_AUTHORITY = "https://login.microsoftonline.com" + +# Audience used when exchanging a managed identity token for a client assertion +# (workload identity federation). Only used in the secretless configuration. +TOKEN_EXCHANGE_SCOPE = "api://AzureADTokenExchange/.default" # nosec B105 — OAuth scope URI, not a secret + + +class ConfigError(RuntimeError): + """A required environment variable is missing or malformed.""" + + +@dataclass(frozen=True) +class Settings: + """Immutable runtime configuration.""" + + tenant_id: str + client_id: str + api_audience: str + required_scope: str + workiq_host: str + client_secret: str | None = field(default=None, repr=False) + + @property + def issuer(self) -> str: + return f"{ENTRA_AUTHORITY}/{self.tenant_id}/v2.0" + + @property + def jwks_uri(self) -> str: + return f"{ENTRA_AUTHORITY}/{self.tenant_id}/discovery/v2.0/keys" + + @property + def workiq_base(self) -> str: + return f"{self.workiq_host.rstrip('/')}{WORKIQ_PATH}/" + + @property + def uses_managed_identity(self) -> bool: + """True when no secret is configured, so a federated assertion is used.""" + return self.client_secret is None + + +def _require(name: str) -> str: + value = os.environ.get(name, "").strip() + if not value: + raise ConfigError(f"{name} must be set") + return value + + +def _allowed_hosts() -> frozenset[str]: + """Built-in hosts plus any from the EXTRA_WORKIQ_HOSTS env var.""" + extra = os.environ.get("EXTRA_WORKIQ_HOSTS", "").strip() + if not extra: + return _ALLOWED_WORKIQ_HOSTS + additions: set[str] = set() + for raw in extra.split(","): + host = raw.strip() + if not host: + continue + if not host.startswith("https://"): + raise ConfigError( + f"EXTRA_WORKIQ_HOSTS entry {host!r} must use the https:// scheme" + ) + additions.add(host.rstrip("/")) + return _ALLOWED_WORKIQ_HOSTS | frozenset(additions) + + +def _validated_workiq_host(host: str) -> str: + # Normalize trailing slashes so "https://host/" matches "https://host". + normalized = host.rstrip("/") + allowed = _allowed_hosts() + if normalized not in allowed: + raise ConfigError( + f"WORKIQ_HOST {host!r} is not in the allowed list: {allowed}" + ) + return normalized + + +@lru_cache(maxsize=1) +def get_settings() -> Settings: + secret = os.environ.get("AZURE_CLIENT_SECRET", "").strip() + return Settings( + tenant_id=_require("AZURE_TENANT_ID"), + client_id=_require("AZURE_CLIENT_ID"), + api_audience=_require("API_AUDIENCE"), + required_scope=os.environ.get("REQUIRED_SCOPE", "").strip() or DEFAULT_REQUIRED_SCOPE, + workiq_host=_validated_workiq_host( + os.environ.get("WORKIQ_HOST", "").strip() or WORKIQ_DEFAULT_HOST + ), + client_secret=secret or None, + ) diff --git a/python/obo/app/main.py b/python/obo/app/main.py new file mode 100644 index 0000000..8f8b436 --- /dev/null +++ b/python/obo/app/main.py @@ -0,0 +1,317 @@ +"""FastAPI backend that brokers Work IQ calls on behalf of the signed-in user. + +Flow per request: + frontend token -> validate -> OBO exchange -> Work IQ token -> call gateway +""" + +from __future__ import annotations + +import json +import logging +from collections.abc import AsyncGenerator +from contextlib import asynccontextmanager +from typing import Annotated, AsyncIterator + +from azure.core.exceptions import ClientAuthenticationError, HttpResponseError +from fastapi import Depends, FastAPI, Header, HTTPException, Request, status +from fastapi.responses import JSONResponse, StreamingResponse +from pydantic import BaseModel, Field +from starlette.types import ASGIApp, Receive, Scope, Send + +from .auth import InvalidToken, TokenValidator, WorkIQTokenExchange +from .config import Settings, get_settings +from .workiq import CONV_ID_PATTERN, WorkIQClient, WorkIQError, init_server_timezone + +logger = logging.getLogger(__name__) + +BEARER_SCHEME = "bearer" +MAX_REQUEST_BODY_BYTES = 64 * 1024 # 64 KB — well above the 8 KB message limit + + +class _BodySizeLimitMiddleware: + """Reject requests whose body exceeds the configured cap. + + Handles both Content-Length (fast reject) and chunked transfer encoding + (streaming byte counter) so the limit cannot be bypassed. + """ + + def __init__(self, app: ASGIApp, *, max_bytes: int = MAX_REQUEST_BODY_BYTES) -> None: + self._app = app + self._max_bytes = max_bytes + + async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: + if scope["type"] == "http": + # Fast path: reject immediately if Content-Length is declared and oversized. + headers = dict(scope.get("headers", [])) + length = headers.get(b"content-length") + try: + declared = int(length) if length is not None else 0 + except ValueError: + declared = 0 # unparseable; slow path will enforce the limit + if declared > self._max_bytes: + response = JSONResponse( + {"detail": "Request body too large"}, + status_code=status.HTTP_413_REQUEST_ENTITY_TOO_LARGE, + ) + await response(scope, receive, send) + return + + # Slow path: count bytes as they arrive (covers chunked encoding). + seen = 0 + rejected = False + response_started = False + error_sent = False + + async def send_wrapper(message: dict) -> None: # type: ignore[type-arg] + nonlocal response_started + # Once we've sent our own 413, suppress any downstream writes + # so the inner app can't corrupt the response. + if error_sent: + return + if message["type"] == "http.response.start": + response_started = True + await send(message) + + async def limited_receive() -> dict: # type: ignore[type-arg] + nonlocal seen, rejected, error_sent + if rejected: + # Return a clean end-of-body instead of disconnect so + # FastAPI doesn't raise request-parsing errors. + return {"type": "http.request", "body": b"", "more_body": False} + message = await receive() + if message.get("type") == "http.request": + seen += len(message.get("body", b"")) + if seen > self._max_bytes: + rejected = True + if not response_started: + error_sent = True + err = JSONResponse( + {"detail": "Request body too large"}, + status_code=status.HTTP_413_REQUEST_ENTITY_TOO_LARGE, + ) + await err(scope, receive, send) + return {"type": "http.request", "body": b"", "more_body": False} + return message + + await self._app(scope, limited_receive, send_wrapper) + return + + await self._app(scope, receive, send) + + +_SECURITY_HEADERS: list[tuple[bytes, bytes]] = [ + (b"x-content-type-options", b"nosniff"), + (b"x-frame-options", b"DENY"), + (b"cache-control", b"no-store"), + (b"referrer-policy", b"no-referrer"), +] + + +class _SecurityHeadersMiddleware: + """Add baseline security headers to every response. + + Implemented as a raw ASGI middleware (not BaseHTTPMiddleware) to avoid + response buffering that would break true streaming on SSE endpoints. + """ + + def __init__(self, app: ASGIApp) -> None: + self._app = app + + async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: + if scope["type"] != "http": + await self._app(scope, receive, send) + return + + async def send_with_headers(message: dict) -> None: # type: ignore[type-arg] + if message["type"] == "http.response.start": + headers = list(message.get("headers", [])) + headers.extend(_SECURITY_HEADERS) + message = {**message, "headers": headers} + await send(message) + + await self._app(scope, receive, send_with_headers) + + +class ChatRequest(BaseModel): + message: str = Field(min_length=1, max_length=8000) + conversation_id: str | None = Field( + default=None, pattern=CONV_ID_PATTERN, max_length=128 + ) + time_zone: str | None = Field( + default=None, pattern=r"^[A-Za-z0-9_+\-/]+$", max_length=64 + ) + + +class CitationModel(BaseModel): + type: str + source: str + provider: str + url: str | None = None + + +class ChatResponse(BaseModel): + conversation_id: str + text: str + citations: list[CitationModel] + + +@asynccontextmanager +async def lifespan(app: FastAPI) -> AsyncIterator[None]: + init_server_timezone() + settings = get_settings() + app.state.settings = settings + app.state.validator = TokenValidator(settings) + app.state.exchange = WorkIQTokenExchange(settings) + logger.info( + "ready — auth=%s, workiq=%s", + "managed-identity" if settings.uses_managed_identity else "client-secret", + settings.workiq_base, + ) + try: + yield + finally: + await app.state.exchange.close() + + +app = FastAPI(title="Work IQ OBO Backend", lifespan=lifespan) +# Starlette add_middleware is LIFO: last-added is outermost. +# SecurityHeaders is outermost so 413s from BodySize also carry the headers. +app.add_middleware(_BodySizeLimitMiddleware) +app.add_middleware(_SecurityHeadersMiddleware) + + +@app.get("/healthz") +async def healthz() -> dict[str, str]: + return {"status": "ok"} + + +def _bearer_token(authorization: Annotated[str | None, Header()] = None) -> str: + if not authorization: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="Missing bearer token", + headers={"WWW-Authenticate": "Bearer"}, + ) + parts = authorization.split(None, 1) + if len(parts) != 2 or parts[0].lower() != BEARER_SCHEME: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="Authorization header must use Bearer scheme", + headers={"WWW-Authenticate": "Bearer"}, + ) + return parts[1].strip() + + +async def workiq_token( + request: Request, + inbound_token: Annotated[str, Depends(_bearer_token)], +) -> str: + """Validate the caller's token and exchange it for a Work IQ token.""" + validator: TokenValidator = request.app.state.validator + exchange: WorkIQTokenExchange = request.app.state.exchange + + try: + claims = await validator.validate(inbound_token) + except InvalidToken as exc: + logger.warning("rejected inbound token: %s", exc) + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="Invalid token", + headers={"WWW-Authenticate": "Bearer"}, + ) from exc + + try: + return await exchange.token_for(inbound_token) + except (ClientAuthenticationError, HttpResponseError) as exc: + # Consent, licensing, or misconfigured app registration — log for triage. + logger.error("OBO exchange failed for %s: %s", claims.get("oid"), exc) + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail="Unable to obtain Work IQ access for this user", + ) from exc + + +def _get_settings(request: Request) -> Settings: + return request.app.state.settings + + +def _client(token: str, settings: Settings) -> WorkIQClient: + return WorkIQClient(access_token=token, base_url=settings.workiq_base) + + +@app.post("/api/chat", response_model=ChatResponse) +async def chat( + request: ChatRequest, + token: Annotated[str, Depends(workiq_token)], + settings: Annotated[Settings, Depends(_get_settings)], +) -> ChatResponse: + try: + async with _client(token, settings) as client: + conversation_id = request.conversation_id or await client.create_conversation() + reply = await client.chat( + conversation_id, request.message, time_zone=request.time_zone + ) + except WorkIQError as exc: + logger.error("work iq call failed: %s", exc) + raise HTTPException( + status_code=status.HTTP_502_BAD_GATEWAY, detail="Work IQ request failed" + ) from exc + + return ChatResponse( + conversation_id=conversation_id, + text=reply.text, + citations=[ + CitationModel( + type=c.attribution_type, + source=c.attribution_source, + provider=c.provider_display_name, + url=c.see_more_web_url, + ) + for c in reply.citations + ], + ) + + +@app.post( + "/api/chat/stream", + responses={ + 200: { + "description": ( + "SSE stream. Events in order:\n" + "1. `event: conversation` — `data: {\"conversation_id\": \"...\"}`\n" + "2. (repeated) default event — `data: {\"text\": \"\"}`\n" + "3. `event: done` — signals clean completion, OR\n" + " `event: error` — `data: upstream request failed`" + ), + "content": {"text/event-stream": {}}, + }, + }, +) +async def chat_stream( + request: ChatRequest, + token: Annotated[str, Depends(workiq_token)], + settings: Annotated[Settings, Depends(_get_settings)], +) -> StreamingResponse: + + async def events() -> AsyncGenerator[str, None]: + try: + async with _client(token, settings) as client: + conversation_id = ( + request.conversation_id or await client.create_conversation() + ) + yield f"event: conversation\ndata: {json.dumps({'conversation_id': conversation_id})}\n\n" + + # JSON-encode each delta: raw newlines in the text would + # otherwise terminate the SSE frame early. + async for delta in client.chat_stream( + conversation_id, request.message, time_zone=request.time_zone + ): + yield f"data: {json.dumps({'text': delta})}\n\n" + yield "event: done\ndata:\n\n" + except WorkIQError as exc: + # The HTTP 200 is already committed, so the error rides the stream. + # Clients must handle "error" events to detect mid-stream failures. + logger.error("work iq stream failed: %s", exc) + yield "event: error\ndata: upstream request failed\n\n" + + return StreamingResponse(events(), media_type="text/event-stream") diff --git a/python/obo/app/workiq.py b/python/obo/app/workiq.py new file mode 100644 index 0000000..62a4aac --- /dev/null +++ b/python/obo/app/workiq.py @@ -0,0 +1,231 @@ +"""Async client for the Work IQ Gateway REST surface. + +Mirrors the wire contract exercised by dotnet/rest/Program.cs: + POST /rest/beta/conversations -> {"id": ...} + POST /rest/beta/conversations/{id}/chat -> {"messages": [...]} + POST /rest/beta/conversations/{id}/chatOverStream -> SSE, cumulative text +""" + +from __future__ import annotations + +import json +import logging +import re +from collections.abc import AsyncGenerator +from dataclasses import dataclass, field +from datetime import datetime +from typing import Any + +try: + from typing import Self +except ImportError: # Python 3.10 + from typing_extensions import Self + +import httpx + +logger = logging.getLogger(__name__) + +# Streaming reads can take minutes; connect/write/pool should be tight. +REQUEST_TIMEOUT = httpx.Timeout(connect=10.0, read=300.0, write=30.0, pool=5.0) +SSE_DATA_PREFIX = "data: " +CONV_ID_PATTERN = r"^[a-zA-Z0-9\-_]{1,128}$" +_CONV_ID_RE = re.compile(CONV_ID_PATTERN) + + +class WorkIQError(Exception): + """A Work IQ Gateway call failed (transport error, HTTP error, or bad response).""" + + +@dataclass(frozen=True) +class Citation: + attribution_type: str + attribution_source: str + provider_display_name: str + see_more_web_url: str | None = None + + +@dataclass(frozen=True) +class ChatReply: + text: str + citations: tuple[Citation, ...] = field(default=()) + + +# Defaults to UTC; call init_server_timezone() at startup (after logging is +# configured) so the fallback warning is actually visible in application logs. +_SERVER_TIMEZONE: str = "UTC" + + +def init_server_timezone() -> None: + """Detect the server's IANA timezone and cache it for request payloads.""" + global _SERVER_TIMEZONE # noqa: PLW0603 + tz = datetime.now().astimezone().tzinfo + name = getattr(tz, "key", None) or str(tz) + # A bare UTC offset like "+05:30" is not an IANA id; fall back rather than 400. + if "/" in name or name == "UTC": + _SERVER_TIMEZONE = name + return + logger.warning( + "Could not detect IANA timezone (got %r); defaulting to UTC. " + "Pass time_zone in requests to override.", + name, + ) + + +def _chat_body(message: str, time_zone: str | None = None) -> dict[str, Any]: + return { + "message": {"text": message}, + "locationHint": {"timeZone": time_zone or _SERVER_TIMEZONE}, + } + + +def _parse_citations(message: dict[str, Any]) -> tuple[Citation, ...]: + return tuple( + Citation( + attribution_type=a.get("attributionType", ""), + attribution_source=a.get("attributionSource", ""), + provider_display_name=a.get("providerDisplayName", ""), + see_more_web_url=a.get("seeMoreWebUrl"), + ) + for a in (message.get("attributions") or []) + if isinstance(a, dict) + ) + + +def _last_text_message(payload: Any) -> dict[str, Any] | None: + """The assistant's reply is the last message carrying a `text` field.""" + if not isinstance(payload, dict): + return None + messages = payload.get("messages") or [] + candidates = [m for m in messages if isinstance(m, dict) and "text" in m] + return candidates[-1] if candidates else None + + +class WorkIQClient: + """One instance per user request — it is bound to that user's token.""" + + def __init__( + self, + access_token: str, + base_url: str, + transport: httpx.AsyncBaseTransport | None = None, + ) -> None: + self._client = httpx.AsyncClient( + base_url=base_url, + headers={"Authorization": f"Bearer {access_token}"}, + timeout=REQUEST_TIMEOUT, + transport=transport, + ) + + async def __aenter__(self) -> Self: + return self + + async def __aexit__(self, *_: Any) -> None: + await self._client.aclose() + + async def create_conversation(self) -> str: + try: + response = await self._client.post("conversations", json={}) + except httpx.HTTPError as exc: + raise WorkIQError(f"create conversation failed: {exc}") from exc + self._raise_for_status(response, "create conversation") + + try: + body = response.json() + conversation_id = body.get("id") if isinstance(body, dict) else None + except (ValueError, httpx.DecodingError) as exc: + raise WorkIQError("create conversation: invalid JSON response") from exc + if not isinstance(conversation_id, str) or not _CONV_ID_RE.fullmatch(conversation_id): + raise WorkIQError("no valid conversation id in response") + return conversation_id + + @staticmethod + def _validate_conversation_id(conversation_id: str) -> None: + if not _CONV_ID_RE.fullmatch(conversation_id): + raise WorkIQError(f"invalid conversation id: {conversation_id!r}") + + async def chat( + self, conversation_id: str, message: str, *, time_zone: str | None = None + ) -> ChatReply: + self._validate_conversation_id(conversation_id) + url = f"conversations/{conversation_id}/chat" + try: + response = await self._client.post( + url, + json=_chat_body(message, time_zone), + ) + except httpx.HTTPError as exc: + raise WorkIQError(f"chat failed: {exc}") from exc + self._raise_for_status(response, "chat") + + try: + payload = response.json() + except (ValueError, httpx.DecodingError) as exc: + raise WorkIQError("chat: invalid JSON response") from exc + + reply = _last_text_message(payload) + if reply is None: + raise WorkIQError("no assistant message in response") + + return ChatReply(text=reply.get("text", ""), citations=_parse_citations(reply)) + + async def chat_stream( + self, conversation_id: str, message: str, *, time_zone: str | None = None + ) -> AsyncGenerator[str, None]: + """Yield text deltas as they arrive. + + The gateway streams cumulative, append-only text, so each event is diffed + against the previous one. A non-prefix update would re-emit the full text; + that does not happen under the current append-only contract. + """ + self._validate_conversation_id(conversation_id) + url = f"conversations/{conversation_id}/chatOverStream" + request = self._client.build_request( + "POST", + url, + json=_chat_body(message, time_zone), + ) + response = None + try: + response = await self._client.send(request, stream=True) + if response.status_code >= 400: + await response.aread() + self._raise_for_status(response, "chat stream") + + previous = "" + async for line in response.aiter_lines(): + if not line.startswith(SSE_DATA_PREFIX): + continue + + event = line[len(SSE_DATA_PREFIX) :].strip() + if not event: + continue + + try: + reply = _last_text_message(json.loads(event)) + except json.JSONDecodeError: + continue # Skip malformed events rather than kill the stream. + + if reply is None: + continue + + text = reply.get("text", "") + delta = text[len(previous) :] if text.startswith(previous) else text + previous = text + if delta: + yield delta + except (httpx.HTTPError, httpx.StreamError) as exc: + raise WorkIQError(f"chat stream failed: {exc}") from exc + finally: + if response is not None: + await response.aclose() + + @staticmethod + def _raise_for_status(response: httpx.Response, action: str) -> None: + if response.status_code < 400: + return + + # request-id is what Microsoft support asks for; keep it out of client replies. + request_id = response.headers.get("request-id", "unknown") + raise WorkIQError( + f"{action} failed: {response.status_code} (request-id={request_id})" + ) diff --git a/python/obo/requirements.txt b/python/obo/requirements.txt new file mode 100644 index 0000000..b21a97a --- /dev/null +++ b/python/obo/requirements.txt @@ -0,0 +1,7 @@ +# For production, pin exact versions with hashes: +# pip-compile --generate-hashes requirements.in -o requirements.txt +fastapi>=0.115 +uvicorn[standard]>=0.32 +httpx>=0.27 +azure-identity>=1.25.3 +pyjwt[crypto]>=2.9 diff --git a/python/obo/smoke_test.py b/python/obo/smoke_test.py new file mode 100644 index 0000000..5bf98fd --- /dev/null +++ b/python/obo/smoke_test.py @@ -0,0 +1,211 @@ +"""Smoke test: fakes the Work IQ gateway with httpx.MockTransport. + +Diagnostic output uses print() deliberately — this is a standalone script, +not a pytest suite, and structured logging adds no value for manual runs. +""" + +import asyncio +import json +import os +import sys +from pathlib import Path +from unittest.mock import patch + +# Fake credentials — must be set before any app.* imports so get_settings() +# picks them up. Scoped via patch.dict so they don't leak into the process +# environment if this module is ever imported by a test runner. +_TEST_ENV = { + "AZURE_TENANT_ID": "11111111-1111-1111-1111-111111111111", + "AZURE_CLIENT_ID": "22222222-2222-2222-2222-222222222222", + "API_AUDIENCE": "api://22222222-2222-2222-2222-222222222222", + "AZURE_CLIENT_SECRET": "local-dev-secret", +} + +sys.path.insert(0, str(Path(__file__).parent)) + +BASE = "https://workiq.test/rest/beta/" + + +def handler(request: "httpx.Request") -> "httpx.Response": + import httpx + + path = request.url.path + if path.endswith("/conversations"): + return httpx.Response(200, json={"id": "conv-42"}) + if path.endswith("/chat"): + return httpx.Response( + 200, + json={ + "messages": [ + {"text": "ignored earlier turn"}, + { + "text": "Your Q3 report is ready.", + "attributions": [ + { + "attributionType": "citation", + "attributionSource": "sharepoint", + "providerDisplayName": "Q3.docx", + "seeMoreWebUrl": "https://contoso.example/q3", + } + ], + }, + ] + }, + ) + if path.endswith("/chatOverStream"): + # Cumulative, append-only text, exactly as the gateway streams it. + frames = ["Hello", "Hello there", "Hello there world"] + body = "".join( + f"data: {json.dumps({'messages': [{'text': f}]})}\n" for f in frames + ) + return httpx.Response(200, text=body) + return httpx.Response(500, json={"error": "unexpected path"}) + + +def error_handler(request: "httpx.Request") -> "httpx.Response": + import httpx + + return httpx.Response(403, json={"error": "no copilot license"}, headers={"request-id": "abc-123"}) + + +async def main() -> None: + import httpx + from fastapi.testclient import TestClient + + from app.config import get_settings + from app.workiq import WorkIQClient, WorkIQError + + # Ensure env vars are picked up by config. + get_settings.cache_clear() + + transport = httpx.MockTransport(handler) + + async with WorkIQClient("fake-token", BASE, transport=transport) as client: + conv = await client.create_conversation() + assert conv == "conv-42", conv + print("create_conversation ->", conv) + + reply = await client.chat(conv, "what's up") + assert reply.text == "Your Q3 report is ready.", reply.text + assert len(reply.citations) == 1 + assert reply.citations[0].provider_display_name == "Q3.docx" + print("chat (last text message) ->", reply.text) + print("citations ->", reply.citations[0].provider_display_name) + + deltas = [d async for d in client.chat_stream(conv, "stream please")] + assert deltas == ["Hello", " there", " world"], deltas + assert "".join(deltas) == "Hello there world" + print("stream deltas ->", deltas) + + # Error path: 403 must surface request-id and not raise something unrelated. + async with WorkIQClient("t", BASE, transport=httpx.MockTransport(error_handler)) as c: + try: + await c.create_conversation() + raise AssertionError("expected WorkIQError") + except WorkIQError as exc: + assert "request-id=abc-123" in str(exc), exc + print("403 error path ->", str(exc)[:60]) + + # Transport error path: network failures must surface as WorkIQError. + def transport_error_handler(request: httpx.Request) -> httpx.Response: + raise httpx.ConnectError("simulated DNS failure") + + async with WorkIQClient("t", BASE, transport=httpx.MockTransport(transport_error_handler)) as c: + try: + await c.create_conversation() + raise AssertionError("expected WorkIQError") + except WorkIQError as exc: + assert "simulated DNS failure" in str(exc), exc + print("transport error (create) ->", str(exc)[:60]) + + try: + await c.chat("conv-1", "hi") + raise AssertionError("expected WorkIQError") + except WorkIQError as exc: + print("transport error (chat) ->", str(exc)[:60]) + + try: + _ = [d async for d in c.chat_stream("conv-1", "hi")] + raise AssertionError("expected WorkIQError") + except WorkIQError as exc: + print("transport error (stream) ->", str(exc)[:60]) + + # Auth gate: no token and malformed token must both 401 before any Work IQ call. + from app.main import MAX_REQUEST_BODY_BYTES, app + + with TestClient(app) as tc: + # -- /healthz (unauthenticated liveness probe) -- + r = tc.get("/healthz") + assert r.status_code == 200, r.status_code + assert r.json() == {"status": "ok"} + print("healthz ->", r.status_code) + + # -- Security response headers -- + assert r.headers["x-content-type-options"] == "nosniff" + assert r.headers["x-frame-options"] == "DENY" + assert r.headers["cache-control"] == "no-store" + assert r.headers["referrer-policy"] == "no-referrer" + print("security headers -> ok") + + # -- Body size limit (413) -- + oversized = b"x" * (MAX_REQUEST_BODY_BYTES + 1) + r = tc.post( + "/api/chat", + content=oversized, + headers={ + "Content-Type": "application/json", + "Content-Length": str(len(oversized)), + }, + ) + assert r.status_code == 413, r.status_code + print("body size limit (CL) ->", r.status_code) + + # -- Body size limit without Content-Length (chunked / slow path) -- + r = tc.post( + "/api/chat", + content=oversized, + headers={"Content-Type": "application/json"}, + ) + assert r.status_code == 413, r.status_code + print("body size limit (chunked)->", r.status_code) + + # -- Auth gate -- + r = tc.post("/api/chat", json={"message": "hi"}) + assert r.status_code == 401, r.status_code + print("no bearer token ->", r.status_code, r.json()["detail"]) + + r = tc.post( + "/api/chat", + json={"message": "hi"}, + headers={"Authorization": "Bearer not-a-jwt"}, + ) + assert r.status_code == 401, r.status_code + print("malformed token ->", r.status_code, r.json()["detail"]) + + # Auth runs before body validation, so a bad token wins over a bad body. + # That ordering is deliberate: unauthenticated callers learn nothing + # about the request schema. + r = tc.post("/api/chat", json={"message": ""}, headers={"Authorization": "Bearer x"}) + assert r.status_code == 401, r.status_code + print("auth precedes validation ->", r.status_code) + + # Config: EXTRA_WORKIQ_HOSTS entries with trailing slashes must match + # after _validated_workiq_host normalizes WORKIQ_HOST. + from app.config import _allowed_hosts, _validated_workiq_host, get_settings + + get_settings.cache_clear() + with patch.dict(os.environ, {**_TEST_ENV, "EXTRA_WORKIQ_HOSTS": "https://workiq.test/"}): + get_settings.cache_clear() + hosts = _allowed_hosts() + assert "https://workiq.test" in hosts, f"trailing slash not normalized: {hosts}" + result = _validated_workiq_host("https://workiq.test/") + assert result == "https://workiq.test", result + print("extra host trailing slash -> ok") + get_settings.cache_clear() + + print("\nAll checks passed.") + + +if __name__ == "__main__": + with patch.dict(os.environ, _TEST_ENV): + asyncio.run(main())