diff --git a/e2e-chatbot-app-next/client/src/components/chat-header.tsx b/e2e-chatbot-app-next/client/src/components/chat-header.tsx index 7b62a346..a169c06f 100644 --- a/e2e-chatbot-app-next/client/src/components/chat-header.tsx +++ b/e2e-chatbot-app-next/client/src/components/chat-header.tsx @@ -4,7 +4,7 @@ import { useWindowSize } from 'usehooks-ts'; import { SidebarToggle } from '@/components/sidebar-toggle'; import { Button } from '@/components/ui/button'; import { useSidebar } from './ui/sidebar'; -import { PlusIcon, CloudOffIcon, MessageSquareOff } from 'lucide-react'; +import { PlusIcon, CloudOffIcon, MessageSquareOff, TriangleAlert } from 'lucide-react'; import { useConfig } from '@/hooks/use-config'; import { Tooltip, @@ -16,72 +16,99 @@ import { const DOCS_URL = 'https://docs.databricks.com/aws/en/generative-ai/agent-framework/chat-app'; +const OBO_DOCS_URL = + 'https://docs.databricks.com/aws/en/dev-tools/databricks-apps/auth'; + export function ChatHeader() { const navigate = useNavigate(); const { open } = useSidebar(); - const { chatHistoryEnabled, feedbackEnabled } = useConfig(); + const { chatHistoryEnabled, feedbackEnabled, oboMissingScopes } = useConfig(); const { width: windowWidth } = useWindowSize(); return ( -
- - - {(!open || windowWidth < 768) && ( - - )} + <> +
+ -
- {!chatHistoryEnabled && ( - - - - - - Ephemeral - - - -

Chat history disabled — conversations are not saved. Click to learn more.

- - - - )} - {!feedbackEnabled && ( - - - - - - Feedback disabled - - - -

Feedback submission disabled. Click to learn more.

- - - + {(!open || windowWidth < 768) && ( + )} -
-
+ +
+ {!chatHistoryEnabled && ( + + + + + + Ephemeral + + + +

Chat history disabled — conversations are not saved. Click to learn more.

+ + + + )} + {!feedbackEnabled && ( + + + + + + Feedback disabled + + + +

Feedback submission disabled. Click to learn more.

+ + + + )} +
+
+ + {oboMissingScopes.length > 0 && ( +
+
+ +

+ This endpoint requires on-behalf-of user authorization. Add these + scopes to your app:{' '} + {oboMissingScopes.join(', ')}.{' '} + Note: UC function scopes are not yet supported.{' '} + + Learn more + +

+
+
+ )} + ); } diff --git a/e2e-chatbot-app-next/client/src/contexts/AppConfigContext.tsx b/e2e-chatbot-app-next/client/src/contexts/AppConfigContext.tsx index caee9956..2240e8c5 100644 --- a/e2e-chatbot-app-next/client/src/contexts/AppConfigContext.tsx +++ b/e2e-chatbot-app-next/client/src/contexts/AppConfigContext.tsx @@ -7,6 +7,12 @@ interface ConfigResponse { chatHistory: boolean; feedback: boolean; }; + obo?: { + enabled: boolean; + requiredScopes: string[]; + missingScopes: string[]; + isSupervisorAgent: boolean; + }; } interface AppConfigContextType { @@ -15,6 +21,10 @@ interface AppConfigContextType { error: Error | undefined; chatHistoryEnabled: boolean; feedbackEnabled: boolean; + oboEnabled: boolean; + oboRequiredScopes: string[]; + oboMissingScopes: string[]; + oboIsSupervisorAgent: boolean; } const AppConfigContext = createContext( @@ -40,6 +50,10 @@ export function AppConfigProvider({ children }: { children: ReactNode }) { // Default to true until loaded to avoid breaking existing behavior chatHistoryEnabled: data?.features.chatHistory ?? true, feedbackEnabled: data?.features.feedback ?? false, + oboEnabled: data?.obo?.enabled ?? false, + oboRequiredScopes: data?.obo?.requiredScopes ?? [], + oboMissingScopes: data?.obo?.missingScopes ?? [], + oboIsSupervisorAgent: data?.obo?.isSupervisorAgent ?? false, }; return ( 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 ef1d0d12..f687f85a 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 @@ -73,13 +73,16 @@ const LOG_SSE_EVENTS = process.env.LOG_SSE_EVENTS === 'true'; const API_PROXY = process.env.API_PROXY; -// Cache for endpoint details to check task type +// Cache for endpoint details to check task type and OBO scopes const endpointDetailsCache = new Map< string, - { task: string | undefined; timestamp: number } + { task: string | undefined; userApiScopes: string[]; isOboEnabled: boolean; timestamp: number } >(); const ENDPOINT_DETAILS_CACHE_DURATION = 5 * 60 * 1000; // 5 minutes +// Cached OBO status — set by getEndpointDetails, read by the provider fetch +let cachedOboEnabled = false; + /** * Checks if context should be injected based on cached endpoint details. * Returns true if API_PROXY is set or if the endpoint task type is agent/v2/chat or agent/v1/responses. @@ -251,10 +254,20 @@ 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}`); + + // If the user's OBO token is present and the endpoint supports OBO, + // use the user's token for Authorization so the endpoint sees the + // user's identity for on-behalf-of authorization. + const userToken = headers.get('x-forwarded-access-token'); + if (userToken && cachedOboEnabled) { + headers.set('Authorization', `Bearer ${userToken}`); + headers.delete('x-forwarded-access-token'); + } else { + const currentToken = await getProviderToken(); + headers.set('Authorization', `Bearer ${currentToken}`); + } + if (API_PROXY) { headers.set('x-mlflow-return-trace-id', 'true'); } @@ -271,7 +284,20 @@ async function getOrCreateDatabricksProvider(): Promise { return provider; } -// Get the task type of the serving endpoint +// Response type for serving endpoint details +interface EndpointDetailsResponse { + task: string | undefined; + auth_policy?: { + user_auth_policy: { + api_scopes: string[]; + }; + }; + tile_endpoint_metadata?: { + problem_type: string; + }; +} + +// Get the task type and OBO scopes of the serving endpoint const getEndpointDetails = async (servingEndpoint: string) => { const cached = endpointDetailsCache.get(servingEndpoint); if ( @@ -294,15 +320,57 @@ const getEndpointDetails = async (servingEndpoint: string) => { headers, }, ); - const data = (await response.json()) as { task: string | undefined }; + const data = (await response.json()) as EndpointDetailsResponse; + + // Detect OBO: either explicit auth_policy scopes, or Supervisor Agent (always OBO) + const isSupervisorAgent = data.tile_endpoint_metadata?.problem_type === 'MULTI_AGENT_SUPERVISOR'; + const userApiScopes = data.auth_policy?.user_auth_policy?.api_scopes ?? []; + const isOboEnabled = userApiScopes.length > 0 || isSupervisorAgent; + cachedOboEnabled = isOboEnabled; + + // serving.serving-endpoints is always needed for OBO (to call the endpoint as the user) + if (isOboEnabled && !userApiScopes.includes('serving.serving-endpoints')) { + userApiScopes.push('serving.serving-endpoints'); + } + + if (isOboEnabled) { + console.warn( + `⚠ OBO detected on endpoint "${servingEndpoint}". Required user authorization scopes: ${JSON.stringify(userApiScopes)}\n` + + ` → Add scopes to your app via the Databricks UI or in databricks.yml\n` + + ` → Note: UC function scopes are not yet supported.\n` + + ` → See: https://docs.databricks.com/aws/en/dev-tools/databricks-apps/auth`, + ); + } + const returnValue = { task: data.task as string | undefined, + userApiScopes, + isOboEnabled, timestamp: Date.now(), }; endpointDetailsCache.set(servingEndpoint, returnValue); return returnValue; }; +/** + * Returns OBO info for the configured serving endpoint. + * Detects OBO via auth_policy scopes or Supervisor Agent type. + */ +export async function getEndpointOboInfo(): Promise<{ enabled: boolean; requiredScopes: string[]; isSupervisorAgent: boolean }> { + const servingEndpoint = process.env.DATABRICKS_SERVING_ENDPOINT; + if (!servingEndpoint) return { enabled: false, requiredScopes: [], isSupervisorAgent: false }; + try { + const details = await getEndpointDetails(servingEndpoint); + return { + enabled: details.isOboEnabled, + requiredScopes: details.userApiScopes, + isSupervisorAgent: details.isOboEnabled && details.userApiScopes.length === 0, + }; + } catch { + return { enabled: false, requiredScopes: [], isSupervisorAgent: false }; + } +} + // Create a smart provider wrapper that handles OAuth initialization interface SmartProvider { languageModel(id: string): Promise; diff --git a/e2e-chatbot-app-next/server/src/routes/config.ts b/e2e-chatbot-app-next/server/src/routes/config.ts index 0413eb93..fc84a9d4 100644 --- a/e2e-chatbot-app-next/server/src/routes/config.ts +++ b/e2e-chatbot-app-next/server/src/routes/config.ts @@ -5,18 +5,64 @@ import { type Router as RouterType, } from 'express'; import { isDatabaseAvailable } from '@chat-template/db'; +import { getEndpointOboInfo } from '@chat-template/ai-sdk-providers'; export const configRouter: RouterType = Router(); +/** + * Decode a JWT payload without verification (just reads claims). + * Returns the parsed payload or null if decoding fails. + */ +function decodeJwtPayload(token: string): Record | null { + try { + const parts = token.split('.'); + if (parts.length !== 3) return null; + const decoded = Buffer.from(parts[1], 'base64url').toString('utf-8'); + return JSON.parse(decoded); + } catch { + return null; + } +} + /** * GET /api/config - Get application configuration - * Returns feature flags based on environment configuration + * Returns feature flags and OBO status based on environment configuration. + * If the user's OBO token is present, decodes it to check which required + * scopes are missing — the banner only shows missing scopes. */ -configRouter.get('/', (_req: Request, res: Response) => { +configRouter.get('/', async (req: Request, res: Response) => { + const oboInfo = await getEndpointOboInfo(); + + let missingScopes = oboInfo.requiredScopes; + + // If the user has an OBO token, check which scopes are already present + const userToken = req.headers['x-forwarded-access-token'] as string | undefined; + if (userToken && oboInfo.enabled) { + const payload = decodeJwtPayload(userToken); + if (payload) { + // Databricks OAuth tokens use 'scope' (space-separated string) + const tokenScopes = typeof payload.scope === 'string' + ? payload.scope.split(' ') + : Array.isArray(payload.scp) ? payload.scp as string[] : []; + // A required scope like "sql.statement-execution" is satisfied by + // an exact match OR by its parent prefix (e.g. "sql") + missingScopes = oboInfo.requiredScopes.filter(required => { + const parent = required.split('.')[0]; + return !tokenScopes.some(ts => ts === required || ts === parent); + }); + } + } + res.json({ features: { chatHistory: isDatabaseAvailable(), feedback: !!process.env.MLFLOW_EXPERIMENT_ID, }, + obo: { + enabled: oboInfo.enabled, + requiredScopes: oboInfo.requiredScopes, + missingScopes, + isSupervisorAgent: oboInfo.isSupervisorAgent, + }, }); }); diff --git a/e2e-chatbot-app-next/tests/api-mocking/api-mock-handlers.ts b/e2e-chatbot-app-next/tests/api-mocking/api-mock-handlers.ts index 6c52d625..8d714aa6 100644 --- a/e2e-chatbot-app-next/tests/api-mocking/api-mock-handlers.ts +++ b/e2e-chatbot-app-next/tests/api-mocking/api-mock-handlers.ts @@ -280,10 +280,16 @@ export const handlers = [ // Mock fetching endpoint details // Returns agent/v1/responses to enable context injection testing + // Includes auth_policy to simulate an OBO-enabled endpoint http.get(/\/api\/2\.0\/serving-endpoints\/[^/]+$/, () => { return HttpResponse.json({ name: 'test-endpoint', task: 'agent/v1/responses', + auth_policy: { + user_auth_policy: { + api_scopes: ['serving.serving-endpoints'], + }, + }, }); }),