From 16a892aeb99ff14fb7c89a61878c38e0e7e0e7e9 Mon Sep 17 00:00:00 2001 From: simon Date: Fri, 28 Nov 2025 16:04:27 +0100 Subject: [PATCH] Add On-Behalf-Of (OBO) authentication support for Databricks Apps MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When deployed in Databricks Apps, the platform injects an x-forwarded-access-token header containing the user's OAuth token. This change enables the chatbot to use this token for downstream API calls, allowing requests to execute with the user's identity and permissions. Changes: - Modify provider fetch wrapper to use OBO token when Authorization header is present - Extract and pass OBO token in chat streaming and title generation endpoints - Send both Authorization: Bearer and x-forwarded-access-token headers for compatibility - Fall back to service principal/PAT/CLI auth when OBO header is absent - Update CLAUDE.md with OBO authentication documentation 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude --- e2e-chatbot-app-next/CLAUDE.md | 22 ++++++++++ .../ai-sdk-providers/src/providers-server.ts | 13 ++++-- .../server/src/routes/chat.ts | 40 ++++++++++++++++++- 3 files changed, 70 insertions(+), 5 deletions(-) diff --git a/e2e-chatbot-app-next/CLAUDE.md b/e2e-chatbot-app-next/CLAUDE.md index 2fb63f8..d231ee7 100644 --- a/e2e-chatbot-app-next/CLAUDE.md +++ b/e2e-chatbot-app-next/CLAUDE.md @@ -306,6 +306,28 @@ req.session = { }; ``` +### On-Behalf-Of (OBO) Authentication + +When deployed in **Databricks Apps**, the platform injects an `x-forwarded-access-token` header containing the user's OAuth token. This enables On-Behalf-Of authentication where API calls execute with the user's identity and permissions. + +**How it works:** +1. Databricks Apps injects `x-forwarded-access-token` header with user's token +2. Chat routes extract this token from request headers +3. Token is passed to AI SDK via `streamText({ headers: {...} })` +4. Provider fetch wrapper uses OBO token instead of service principal + +**Authentication priority** (highest to lowest): +1. OBO token (`x-forwarded-access-token` header) - when in Databricks Apps +2. PAT token (`DATABRICKS_TOKEN` env var) +3. OAuth service principal (`DATABRICKS_CLIENT_ID` + `DATABRICKS_CLIENT_SECRET`) +4. CLI auth (`databricks auth token`) + +**Benefits of OBO:** +- API calls use user's Unity Catalog permissions +- MLflow traces show actual user identity +- Custom API_PROXY backends receive user context +- Row-level security and column masks are enforced + ### Auth Middleware Usage - `authMiddleware` - Extracts session (doesn't reject) diff --git a/e2e-chatbot-app-next/packages/ai-sdk-providers/src/providers-server.ts b/e2e-chatbot-app-next/packages/ai-sdk-providers/src/providers-server.ts index 3772391..79d0152 100644 --- a/e2e-chatbot-app-next/packages/ai-sdk-providers/src/providers-server.ts +++ b/e2e-chatbot-app-next/packages/ai-sdk-providers/src/providers-server.ts @@ -122,10 +122,17 @@ async function getOrCreateDatabricksProvider(): Promise { baseURL: `${hostname}/serving-endpoints`, formatUrl: ({ baseUrl, path }) => API_PROXY ?? `${baseUrl}${path}`, fetch: async (...[input, init]: Parameters) => { - // Always get fresh token for each request (will use cache if valid) - const currentToken = await getProviderToken(); const headers = new Headers(init?.headers); - headers.set('Authorization', `Bearer ${currentToken}`); + + // Check if OBO (On-Behalf-Of) token was passed via headers from route handler. + // When running in Databricks Apps, the x-forwarded-access-token header contains + // the user's token for OBO authentication. If Authorization is already set, + // it means OBO headers were passed - use them instead of service principal. + if (!headers.has('Authorization')) { + // No OBO token provided, use service principal/PAT/CLI fallback + const currentToken = await getProviderToken(); + headers.set('Authorization', `Bearer ${currentToken}`); + } return databricksFetch(input, { ...init, diff --git a/e2e-chatbot-app-next/server/src/routes/chat.ts b/e2e-chatbot-app-next/server/src/routes/chat.ts index a6df7cd..91a6ca8 100644 --- a/e2e-chatbot-app-next/server/src/routes/chat.ts +++ b/e2e-chatbot-app-next/server/src/routes/chat.ts @@ -107,7 +107,14 @@ chatRouter.post('/', requireAuth, async (req: Request, res: Response) => { if (!chat) { if (isDatabaseAvailable()) { - const title = await generateTitleFromUserMessage({ message }); + // Extract OBO token for title generation (same token used for streaming below) + const oboTokenForTitle = req.headers[ + 'x-forwarded-access-token' + ] as string | undefined; + const title = await generateTitleFromUserMessage({ + message, + oboToken: oboTokenForTitle, + }); await saveChat({ id, @@ -147,9 +154,25 @@ chatRouter.post('/', requireAuth, async (req: Request, res: Response) => { const streamId = generateUUID(); const model = await myProvider.languageModel(selectedChatModel); + + // Extract OBO (On-Behalf-Of) token if present. + // Databricks Apps injects x-forwarded-access-token header for user authorization. + // When present, use it for downstream API calls to execute with user's identity/permissions. + const oboToken = req.headers['x-forwarded-access-token'] as string | undefined; + + // Build auth headers: send both Authorization and x-forwarded-access-token + // for maximum compatibility with Databricks APIs and custom API_PROXY backends. + const oboHeaders = oboToken + ? { + Authorization: `Bearer ${oboToken}`, + 'x-forwarded-access-token': oboToken, + } + : undefined; + const result = streamText({ model, messages: convertToModelMessages(uiMessages), + headers: oboHeaders, onFinish: ({ usage }) => { finalUsage = usage; }, @@ -352,7 +375,8 @@ chatRouter.get( chatRouter.post('/title', requireAuth, async (req: Request, res: Response) => { try { const { message } = req.body; - const title = await generateTitleFromUserMessage({ message }); + const oboToken = req.headers['x-forwarded-access-token'] as string | undefined; + const title = await generateTitleFromUserMessage({ message, oboToken }); res.json({ title }); } catch (error) { console.error('Error generating title:', error); @@ -387,12 +411,24 @@ chatRouter.patch( // Helper function to generate title from user message async function generateTitleFromUserMessage({ message, + oboToken, }: { message: ChatMessage; + oboToken?: string; }) { const model = await myProvider.languageModel('title-model'); + + // Build OBO headers if token is present (same pattern as streaming endpoint) + const oboHeaders = oboToken + ? { + Authorization: `Bearer ${oboToken}`, + 'x-forwarded-access-token': oboToken, + } + : undefined; + const { text: title } = await generateText({ model, + headers: oboHeaders, system: `\n - you will generate a short title based on the first message a user begins a conversation with - ensure it is not more than 80 characters long