I3eka commented on code in PR #43133: URL: https://github.com/apache/superset/pull/43133#discussion_r4002584337
########## superset-frontend/src/features/ai/AiAssistantPanel.tsx: ########## @@ -0,0 +1,1150 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +/** + * @fileoverview The assistant panel. + * + * The host owns where this sits and, when docked, how wide it is, so there is no + * positioning and no resize handle here. What is here is the conversation: the + * header, the transcript, what the assistant is doing while it works, and the + * composer. + * + * The centre of the design is that a run is legible while it happens. An answer + * can take a minute of tool calls, and a spinner for a minute is indistinguishable + * from a hang, so reasoning streams into a preview, each step appends to a tool + * log, and a checkpoint stops the run with a countdown the user can act on. + */ + +import { + memo, + useCallback, + useEffect, + useRef, + useState, + useSyncExternalStore, +} from 'react'; +import type { Dispatch, SetStateAction } from 'react'; +import ReactMarkdown from 'react-markdown'; +import type { Components } from 'react-markdown'; +import { css, keyframes, styled, useTheme } from '@apache-superset/core/theme'; +import { t } from '@apache-superset/core/translation'; +import type { chat as chatApi } from '@apache-superset/core'; +import { + Button, + Input, + Loading, + Tooltip, + Typography, +} from '@superset-ui/core/components'; +import { Icons } from '@superset-ui/core/components/Icons'; +import { chat } from 'src/core/chat'; +import ChatAgentSelect from './components/ChatAgentSelect'; +import ChatTabsMenu from './components/ChatTabsMenu'; +import { REMARK_PLUGINS, useChatMarkdown } from './components/chatMarkdown'; +import { ThoughtProcess } from './components/ThoughtProcess'; +import { useChatBot } from './hooks/useChatBot'; +import { AI_ACTION_EVENT, type AiActionEvent } from './hooks/useAIAction'; +import type { PageContext } from './hooks/usePageContext'; +import type { ChatMessageWithMeta, CheckpointPayload } from './types'; + +/** + * How long a checkpoint waits before continuing on its own. A pause that blocks + * forever is worse than one that resolves optimistically: the user may have + * walked away, and the run should not be stranded. + */ +export const CHECKPOINT_TIMEOUT_SECONDS = 30; + +/** + * Closes a code fence the model has not finished writing. + * + * A streamed answer is parsed on every delta, so a fence arrives in pieces — + * "```", then "sql", then the query. Markdown with an odd number of fences + * renders the opening backticks literally and then reflows once the closing pair + * lands, which reads as the answer glitching. Balancing the count keeps each + * intermediate state a valid document. + */ +export const balanceCodeFences = (text: string): string => { + const fences = text.match(/^```/gm)?.length ?? 0; + return fences % 2 === 0 ? text : `${text}\n\`\`\``; +}; + +/** + * Whether a message carries the structured record of how it was answered, as + * opposed to only the flat log assembled from stream frames. + */ +const hasStructuredThinking = (message: ChatMessageWithMeta): boolean => + Boolean(message.toolCalls?.length || message.thoughts || message.pageContext); + +/** Milliseconds between typewriter frames, and characters per frame. */ +const TYPEWRITER_INTERVAL_MS = 18; +const TYPEWRITER_STEP = 3; + +/** + * The panel's own size as a floating overlay. + * + * Docked width belongs to the host and is not set here. Floating does need a size + * from somewhere, though — the floating host only stacks its children in a corner + * and gives them no dimensions — so these clamp the overlay to the viewport. + */ +const FLOATING_WIDTH_PX = 440; +const FLOATING_MAX_HEIGHT_VH = 70; + +/** + * The panel surface. + * + * Positioning is deliberately absent: the host places this, in both modes. What is + * here is the surface itself — a column that fills whatever box it is given, with a + * floating size for the mode where the host provides no box. + */ +const ChatPanelContainer = styled.div<{ floating: boolean }>` + display: flex; + flex-direction: column; + min-height: 0; + overflow: hidden; + background: ${({ theme }) => theme.colorBgElevated}; + ${({ floating, theme }) => + floating + ? css` + width: min( + ${FLOATING_WIDTH_PX}px, + calc(100vw - ${theme.sizeUnit * 12}px) + ); + height: ${FLOATING_MAX_HEIGHT_VH}vh; + border: 1px solid ${theme.colorBorderSecondary}; + border-radius: ${theme.borderRadiusLG}px; + box-shadow: ${theme.boxShadow}; + ` + : css` + width: 100%; + height: 100%; + `} +`; + +const ChatHeader = styled.div` + display: flex; + justify-content: space-between; + align-items: center; + padding: ${({ theme }) => theme.sizeUnit * 3}px + ${({ theme }) => theme.sizeUnit * 4}px; + background: ${({ theme }) => theme.colorBgContainer}; + border-bottom: 1px solid ${({ theme }) => theme.colorBorderSecondary}; + font-weight: ${({ theme }) => theme.fontWeightStrong}; + font-size: ${({ theme }) => theme.fontSizeLG}px; + color: ${({ theme }) => theme.colorTextHeading}; + flex-shrink: 0; + gap: ${({ theme }) => theme.sizeUnit * 2}px; +`; + +const HeaderGroup = styled.div` + display: flex; + align-items: center; + min-width: 0; +`; + +const HeaderTitle = styled.span` + overflow: hidden; + text-overflow: ellipsis; + white-space: nowrap; +`; + +const ChatMessages = styled.div` + flex: 1; + min-height: 0; + padding: ${({ theme }) => theme.sizeUnit * 4}px; + overflow-y: auto; + + &::-webkit-scrollbar { + width: 4px; + } + + &::-webkit-scrollbar-track { + background: ${({ theme }) => theme.colorBgContainer}; + border-radius: 2px; + } + + &::-webkit-scrollbar-thumb { + background: ${({ theme }) => theme.colorFillSecondary}; + border-radius: 2px; + } +`; + +const MessageBubble = styled.div<{ variant: 'user' | 'assistant' }>` + margin-bottom: ${({ theme }) => theme.sizeUnit * 3}px; + display: flex; + flex-direction: column; + align-items: ${({ variant }) => + variant === 'user' ? 'flex-end' : 'flex-start'}; +`; + +const MessageContent = styled.div<{ variant: 'user' | 'assistant' }>` + max-width: 85%; + padding: ${({ theme }) => theme.sizeUnit * 3}px + ${({ theme }) => theme.sizeUnit * 4}px; + border-radius: ${({ theme }) => theme.borderRadiusLG * 2}px; + background: ${({ theme, variant }) => + variant === 'user' ? theme.colorPrimary : theme.colorBgContainer}; + color: ${({ theme, variant }) => + variant === 'user' ? theme.colorTextLightSolid : theme.colorText}; + font-size: ${({ theme }) => theme.fontSize}px; + line-height: 1.5; + border: ${({ theme, variant }) => + variant === 'assistant' + ? `1px solid ${theme.colorBorderSecondary}` + : 'none'}; + box-shadow: ${({ theme }) => theme.boxShadowTertiary}; + overflow-wrap: anywhere; + + p { + margin: 0; + } + + p:not(:last-child) { + margin-bottom: ${({ theme }) => theme.sizeUnit * 2}px; + } + + a { + color: ${({ theme, variant }) => + variant === 'user' ? theme.colorTextLightSolid : theme.colorPrimary}; + text-decoration: underline; + } + + code { + background: ${({ theme, variant }) => + variant === 'user' ? theme.colorPrimaryActive : theme.colorFillTertiary}; + padding: 2px ${({ theme }) => theme.sizeUnit * 1.5}px; + border-radius: ${({ theme }) => theme.borderRadius}px; + font-size: ${({ theme }) => theme.fontSizeSM}px; + font-family: ${({ theme }) => theme.fontFamilyCode}; + } + + pre { + background: ${({ theme }) => theme.colorFillQuaternary}; + padding: ${({ theme }) => theme.sizeUnit * 3}px; + border-radius: ${({ theme }) => theme.borderRadius}px; + overflow-x: auto; + margin: ${({ theme }) => theme.sizeUnit * 2}px 0; + border: 1px solid ${({ theme }) => theme.colorBorderSecondary}; + } + + pre code { + background: none; + padding: 0; + } +`; + +const MessageActions = styled.div` + display: flex; + gap: ${({ theme }) => theme.sizeUnit}px; + margin-top: ${({ theme }) => theme.sizeUnit}px; +`; + +const ActionButton = styled(Button)` + &&& { + padding: 2px ${({ theme }) => theme.sizeUnit * 1.5}px; + height: ${({ theme }) => theme.sizeUnit * 6}px; + font-size: ${({ theme }) => theme.fontSizeSM}px; + } + + /* The recorded verdict keeps its colour while disabled. Both thumbs lock once + a rating exists, and the default disabled grey would hide which one the + user picked — the state matters more here than the affordance. */ + &&&.is-active, + &&&.is-active:disabled, + &&&.is-active[disabled] { + color: ${({ theme }) => theme.colorPrimary}; + } +`; + +const LiveAnswer = styled.div` + margin-top: ${({ theme }) => theme.sizeUnit * 2}px; + color: ${({ theme }) => theme.colorText}; + font-size: ${({ theme }) => theme.fontSize}px; + line-height: 1.5; + overflow-wrap: anywhere; + + p { + margin: 0; + } + + p:not(:last-child) { + margin-bottom: ${({ theme }) => theme.sizeUnit * 2}px; + } +`; + +const thinkingPulse = keyframes` + 0% { + opacity: 0.45; + } + 50% { + opacity: 1; + } + 100% { + opacity: 0.45; + } +`; + +const ThinkingPreview = styled.div<{ isLive?: boolean }>` + color: ${({ theme }) => theme.colorTextTertiary}; + font-size: ${({ theme }) => theme.fontSizeSM}px; + white-space: pre-wrap; + ${({ isLive }) => + isLive && + css` + animation: ${thinkingPulse} 1.8s ease-in-out infinite; + `} +`; + +const ThinkingDetails = styled.details` + margin-bottom: ${({ theme }) => theme.sizeUnit * 2}px; + color: ${({ theme }) => theme.colorTextTertiary}; + font-size: ${({ theme }) => theme.fontSizeSM}px; + + summary { + cursor: pointer; + user-select: none; + color: ${({ theme }) => theme.colorTextTertiary}; + margin-bottom: ${({ theme }) => theme.sizeUnit}px; + } +`; + +const CheckpointDivider = styled.div` + display: flex; + align-items: center; + gap: ${({ theme }) => theme.sizeUnit * 3}px; + margin: ${({ theme }) => theme.sizeUnit * 4}px 0 + ${({ theme }) => theme.sizeUnit * 3}px; + + &::before, + &::after { + content: ''; + flex: 1; + height: 1px; + background: ${({ theme }) => theme.colorBorderSecondary}; + } +`; + +const CountdownBadge = styled.span` + font-size: ${({ theme }) => theme.fontSizeSM}px; + font-weight: ${({ theme }) => theme.fontWeightStrong}; + font-variant-numeric: tabular-nums; + color: ${({ theme }) => theme.colorTextSecondary}; + white-space: nowrap; +`; + +const CheckpointContent = styled.div` + font-size: ${({ theme }) => theme.fontSizeSM}px; + color: ${({ theme }) => theme.colorText}; + line-height: 1.5; +`; + +const CheckpointTaskList = styled.ul` + margin: ${({ theme }) => theme.sizeUnit * 1.5}px 0; + padding-left: ${({ theme }) => theme.sizeUnit * 4.5}px; + font-size: ${({ theme }) => theme.fontSizeSM}px; + color: ${({ theme }) => theme.colorTextSecondary}; + + li { + margin-bottom: 2px; + } +`; + +const CheckpointEstimate = styled.div` + font-size: ${({ theme }) => theme.fontSizeSM}px; + color: ${({ theme }) => theme.colorTextTertiary}; + margin-top: ${({ theme }) => theme.sizeUnit}px; +`; + +const CheckpointActions = styled.div` + display: flex; + gap: ${({ theme }) => theme.sizeUnit * 2}px; + margin-top: ${({ theme }) => theme.sizeUnit * 2.5}px; +`; + +const ChatInput = styled.div` + padding: ${({ theme }) => theme.sizeUnit * 4}px; + border-top: 1px solid ${({ theme }) => theme.colorBorderSecondary}; + background: ${({ theme }) => theme.colorBgContainer}; + flex-shrink: 0; +`; + +const InputContainer = styled.div` + display: flex; + gap: ${({ theme }) => theme.sizeUnit * 2}px; + align-items: flex-end; +`; + +const QuickPromptsRow = styled.div<{ hasContent: boolean }>` + display: flex; + align-items: center; + gap: ${({ theme }) => theme.sizeUnit * 2}px; + flex-wrap: wrap; + margin-bottom: ${({ theme, hasContent }) => + hasContent ? `${theme.sizeUnit * 2.5}px` : '0'}; + min-height: ${({ theme, hasContent }) => + hasContent ? `${theme.sizeUnit * 6}px` : '0'}; +`; + +const QuickPromptChip = styled(Button)` + &&& { + width: fit-content; + max-width: 100%; + height: auto; + white-space: normal; + text-align: left; + line-height: 1.35; + word-break: break-word; + } +`; + +const PageContextRow = styled.div` + display: flex; + align-items: center; + margin-bottom: ${({ theme }) => theme.sizeUnit * 1.5}px; +`; + +const PageContextPill = styled.button<{ isActive: boolean }>` + display: inline-flex; + align-items: center; + gap: ${({ theme }) => theme.sizeUnit * 1.5}px; + padding: 3px ${({ theme }) => theme.sizeUnit * 2}px; + border-radius: ${({ theme }) => theme.borderRadiusLG}px; + border: 1px solid + ${({ theme, isActive }) => + isActive ? theme.colorPrimary : theme.colorBorderSecondary}; + background: ${({ theme, isActive }) => + isActive ? theme.colorPrimaryBg : theme.colorFillQuaternary}; + color: ${({ theme, isActive }) => + isActive ? theme.colorPrimary : theme.colorTextTertiary}; + font-size: ${({ theme }) => theme.fontSizeSM}px; + cursor: pointer; + transition: all ${({ theme }) => theme.motionDurationMid}; + max-width: ${({ theme }) => theme.sizeUnit * 62}px; + white-space: nowrap; + + &:hover { + border-color: ${({ theme }) => theme.colorPrimary}; + } +`; + +const PillLabel = styled.span` + overflow: hidden; + text-overflow: ellipsis; +`; + +const PillDot = styled.span<{ isActive: boolean }>` + width: 7px; + height: 7px; + border-radius: 50%; + flex-shrink: 0; + background: ${({ theme, isActive }) => + isActive ? theme.colorPrimary : theme.colorTextQuaternary}; + transition: background ${({ theme }) => theme.motionDurationMid}; +`; + +const EmptyState = styled.div` + text-align: center; + color: ${({ theme }) => theme.colorTextSecondary}; + font-size: ${({ theme }) => theme.fontSizeLG}px; + margin-top: ${({ theme }) => theme.sizeUnit * 15}px; + padding: 0 ${({ theme }) => theme.sizeUnit * 10}px; + line-height: 1.5; +`; + +/** + * What the page-context pill says it is sending. + * + * Naming the specific chart or dashboard is the whole point: "page context" tells + * the user nothing about whether the assistant can see the thing they are asking + * about. + */ +export const getPageContextLabel = (context: PageContext): string => { + switch (context.pageType) { + case 'sqllab': + return context.sqlContext?.activeEditor?.name || t('SQL Lab'); + case 'dashboard': + return context.dashboardContext?.title + ? t('Dashboard: %s', context.dashboardContext.title) + : t('Dashboard'); + case 'explore': + case 'chart': { + const name = + context.chartContext?.chartName || + context.chartContext?.slice?.slice_name; + return name ? t('Chart: %s', name) : t('Chart Explorer'); + } + case 'home': + return t('Home'); + default: + return context.pathname; + } +}; + +interface MemoizedChatMessageProps { + message: ChatMessageWithMeta; + feedback: 'like' | 'dislike' | null; + /** Per-message SQL expansion state, present only to invalidate the memo. */ + sqlBlockExpandState: boolean[] | undefined; + createMarkdownComponents: (messageId: string) => Components; + /** Keeps this message's thought process open, for the run that just ended. */ + keepThinkingOpen?: boolean; + onFeedback: (messageId: string, feedback: 'like' | 'dislike') => void; + onCopy: (content: string) => Promise<void>; +} + +/** + * One turn. + * + * Memoized on content rather than identity because a streaming run re-renders the + * panel on every frame, and re-parsing the markdown of every earlier message each + * time made long conversations visibly janky. + */ +const MemoizedChatMessage = memo( + ({ + message, + feedback, + createMarkdownComponents, + keepThinkingOpen, + onFeedback, + onCopy, + }: MemoizedChatMessageProps) => { + const markdownComponents = createMarkdownComponents(message.id); + return ( + <MessageBubble + variant={message.role} + data-test="chat-message" + data-role={message.role} + > + <MessageContent variant={message.role}> + {message.role === 'assistant' && + // Steps are rendered structurally where the server has sent them, + // and as the flat log otherwise. The two are not alternatives so + // much as two stages: a turn assembled from stream frames only has + // the text, and gains its structure when the transcript is re-read + // from the server once the run ends. Any one of the structured + // fields is enough — a turn can have page context and no tool calls, + // or reasoning and no steps. + (hasStructuredThinking(message) ? ( + <ThoughtProcess + reasoning={message.thoughts} + pageContext={message.pageContext} + toolCalls={message.toolCalls} + markdownComponents={markdownComponents} + // Left open for the turn that has only just finished, so the + // section does not collapse and pull the answer up the panel the + // moment the user starts reading it. + defaultOpen={keepThinkingOpen} + /> + ) : ( + message.thinking && ( + <ThinkingDetails> + <summary>{t('Thought process')}</summary> + {/* The tool log is markdown too, so a SQL step gets the same + highlighting and "Run in SQL Lab" as SQL in the answer. */} + <ReactMarkdown + components={markdownComponents} + remarkPlugins={REMARK_PLUGINS} + > + {message.thinking} + </ReactMarkdown> + </ThinkingDetails> + ) + ))} + <ReactMarkdown + components={markdownComponents} + remarkPlugins={REMARK_PLUGINS} + > + {message.content} + </ReactMarkdown> + </MessageContent> + {message.role === 'assistant' && ( + <MessageActions> + {/* Both thumbs stay available so a mis-click can be corrected. The + recorded verdict is the coloured one, and only it is inert — + there is no endpoint to withdraw a rating, so pressing it again + would do nothing and should not look like it could. */} + <ActionButton + buttonStyle="link" + icon={<Icons.LikeOutlined iconSize="s" />} + onClick={() => onFeedback(message.id, 'like')} + className={feedback === 'like' ? 'is-active' : ''} + aria-label={t('Good response')} + aria-pressed={feedback === 'like'} + tooltip={ + feedback === 'like' + ? t('You marked this a good response') + : t('Good response') + } + disabled={message.pending || feedback === 'like'} + /> + <ActionButton + buttonStyle="link" + icon={<Icons.DislikeOutlined iconSize="s" />} + onClick={() => onFeedback(message.id, 'dislike')} + className={feedback === 'dislike' ? 'is-active' : ''} + aria-label={t('Bad response')} + aria-pressed={feedback === 'dislike'} + tooltip={ + feedback === 'dislike' + ? t('You marked this a bad response') + : t('Bad response') + } + disabled={message.pending || feedback === 'dislike'} + /> + <ActionButton + buttonStyle="link" + icon={<Icons.CopyOutlined iconSize="s" />} + onClick={() => { + onCopy(message.content).catch(() => { + // Clipboard access can be refused; the button then does + // nothing rather than raising at the user. + }); + }} + aria-label={t('Copy to clipboard')} + tooltip={t('Copy to clipboard')} + /> + </MessageActions> + )} + </MessageBubble> + ); + }, + (previous, next) => Review Comment: The memo comparator still does not account for the context captured by the Markdown factory. This is unchanged base UI code in the refresh; it needs a tab/editor-switch test proving that an old message uses the current SQL Lab callback. ########## superset-frontend/src/features/ai/hooks/usePageContext.ts: ########## @@ -0,0 +1,816 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +/** + * @fileoverview What the user is looking at, read out of the Redux store. + * + * The assistant cannot answer "why is this number wrong" without knowing which + * chart, query or dashboard "this" refers to, so every page type contributes a + * shape describing itself. Everything read here is already on the client; no + * request is made to build it. + * + * Only one page's slice is populated per call, because only one of SQL Lab, + * Explore and a dashboard is mounted at a time. Selectors are written against + * optional slices for that reason rather than assuming a page's reducers are + * registered. + */ + +import { useMemo } from 'react'; +import { useSelector } from 'react-redux'; +import { useLocation } from 'react-router-dom'; +import type { DataMaskStateWithId } from '@superset-ui/core'; +import type { SqlLabRootState } from 'src/SqlLab/types'; +import type { ExplorePageState } from 'src/explore/types'; +import type { + DashboardInfo, + DashboardLayoutState, + LayoutItem, + RootState as DashboardRootState, + SliceEntitiesState, +} from 'src/dashboard/types'; +import { + CHART_TYPE, + MARKDOWN_TYPE, + TAB_TYPE, +} from 'src/dashboard/util/componentTypes'; +import { DASHBOARD_HEADER_ID } from 'src/dashboard/util/constants'; + +/** At most this many markdown blocks are lifted off a dashboard. */ +const MAX_MARKDOWN_BLOCKS = 10; + +/** Per-block ceiling, so one enormous markdown tile cannot fill the prompt. */ +const MAX_MARKDOWN_BLOCK_LENGTH = 25000; + +/** Formatted context above this length is considered worth minimising. */ +const LARGE_CONTEXT_THRESHOLD = 5000; + +/** + * The slices this hook reads. + * + * Declared locally rather than intersecting the three page root states: those + * disagree about the shape of shared keys such as `common`, and an intersection + * of them is uninhabitable. + */ +interface AiPageRootState { + sqlLab?: SqlLabRootState['sqlLab']; + explore?: ExplorePageState['explore']; + /** + * `dashboard_title` is spread onto this slice at hydration but is absent from + * the declared type, so it is named here instead of being read through a cast. + */ + dashboardInfo?: DashboardInfo & { dashboard_title?: string }; + dashboardState?: DashboardRootState['dashboardState']; + dashboardLayout?: DashboardLayoutState; + sliceEntities?: SliceEntitiesState; + nativeFilters?: DashboardRootState['nativeFilters']; + dataMask?: DataMaskStateWithId; +} + +/** + * Extract chart ID from URL patterns like /explore/?slice_id=123 or /chart/123/ + */ +const extractChartIdFromUrl = (url: string): number | undefined => { + // Try to extract from slice_id parameter + const sliceIdMatch = url.match(/[?&]slice_id=(\d+)/); + if (sliceIdMatch) { + return parseInt(sliceIdMatch[1], 10); + } + + // Try to extract from chart path + const chartPathMatch = url.match(/\/chart\/(\d+)/); + if (chartPathMatch) { + return parseInt(chartPathMatch[1], 10); + } + + return undefined; +}; + +/** + * Utility function to truncate text with ellipsis + */ +export function truncateText(text: string, maxLength: number): string { + if (text.length <= maxLength) return text; + return `${text.slice(0, maxLength)}...`; +} + +const HELPER_PREFIX = '@helper'; + +const isHelperDirective = (content: string): boolean => + content.trimStart().toLowerCase().startsWith(HELPER_PREFIX); + +export interface PageContext { + url: string; + pathname: string; + pageType: 'sqllab' | 'explore' | 'dashboard' | 'chart' | 'home' | 'other'; + sqlContext?: { + activeEditor?: { + sql?: string; + database?: string; + /** Sent alongside the name so the assistant can run SQL against the + * connection the user has selected rather than resolving one by name. */ + databaseId?: number; + schema?: string; + catalog?: string; + queryLimit?: number; + name?: string; + }; + tables?: Array<{ + name: string; + schema?: string; + catalog?: string; + }>; + recentQueries?: Array<{ + sql?: string; + status?: string; + executedAt?: number; + }>; + }; + chartContext?: { + chartId?: string | number; + vizType?: string; + datasource?: { + id?: string | number; + name?: string; + type?: string; + schema?: string; + database?: string; + }; + /** A subset of the chart's controls. Values stay `unknown`: the assistant + * only ever serialises them, and typing them as `any` would let unchecked + * reads back in. */ + formData?: Record<string, unknown>; + chartName?: string; + metadata?: { + isEditing?: boolean; + canOverwrite?: boolean; + description?: string; + }; + slice?: { + slice_id?: number; + slice_name?: string; + description?: string; + }; + }; + dashboardContext?: { + /** Sent so the assistant can call dashboard tools against the dashboard the + * user is actually on. Without it the model has to guess an id or fall back + * to searching by title. */ + id?: number; + title?: string; + activeTabId?: string; + activeTabLabel?: string; + charts?: Array<{ + id: number; + title?: string; + }>; + activeFilters?: Array<{ + name: string; + column?: string; + filterType?: string; + value?: unknown; + }>; + }; + pageMarkdown?: Array<{ + source: string; + content: string; + }>; +} + +export interface ChatHistoryEntry { + role: 'user' | 'assistant'; + content: string; +} + +/** A named string off an untyped bag, or undefined if it is not one. */ +const readStringField = ( + bag: Record<string, unknown> | undefined, + field: string, +): string | undefined => { + const value = bag?.[field]; + return typeof value === 'string' && value ? value : undefined; +}; + +/** Layout items inside the currently selected tab tree, or all of them when the + * dashboard has no tabs. */ +const isInActiveTab = ( + item: LayoutItem, + layout: Record<string, LayoutItem>, + activeTabs: Set<string>, +): boolean => { + const tabParents = (item.parents ?? []).filter( + parent => layout[parent]?.type === TAB_TYPE, + ); + // On dashboards with tabs, include only charts inside the active tab tree. + // If there are no tabs or no active tab selected, keep all non-tab charts. + if (tabParents.length === 0) { Review Comment: The mixed dashboard context issue remains open: top-level visible charts must not be excluded just because another tab is active. No filtering-rule change is included in the refresh. This belongs with the base page-context tests, not a dataset-specific exception. ########## superset-frontend/src/features/ai/components/ChatTabsMenu.tsx: ########## @@ -0,0 +1,356 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +/** + * @fileoverview The conversation list. + * + * Conversations live behind one menu rather than a tab strip: the panel is narrow + * enough in floating mode that a strip would truncate every name, and the list + * doubles as the history of past conversations, which a strip cannot be. + */ + +import { useCallback, useState } from 'react'; +import type { MouseEvent as ReactMouseEvent } from 'react'; +import { styled } from '@apache-superset/core/theme'; +import { t } from '@apache-superset/core/translation'; +import { Button, Dropdown, Popconfirm } from '@superset-ui/core/components'; +import { Icons } from '@superset-ui/core/components/Icons'; +import type { ChatTab } from '../types'; + +const MenuContainer = styled.div` + background: ${({ theme }) => theme.colorBgElevated}; + border-radius: ${({ theme }) => theme.borderRadius}px; + box-shadow: ${({ theme }) => theme.boxShadowSecondary}; + min-width: ${({ theme }) => theme.sizeUnit * 65}px; + max-height: ${({ theme }) => theme.sizeUnit * 100}px; + overflow-y: auto; + border: 1px solid ${({ theme }) => theme.colorBorderSecondary}; +`; + +const MenuHeader = styled.div` + padding: ${({ theme }) => theme.sizeUnit * 3}px + ${({ theme }) => theme.sizeUnit * 4}px; + border-bottom: 1px solid ${({ theme }) => theme.colorBorderSecondary}; + font-weight: ${({ theme }) => theme.fontWeightStrong}; + font-size: ${({ theme }) => theme.fontSizeSM}px; + color: ${({ theme }) => theme.colorTextSecondary}; + text-transform: uppercase; + letter-spacing: 0.5px; +`; + +const NewChatButton = styled.button` + display: flex; + align-items: center; + gap: ${({ theme }) => theme.sizeUnit * 2}px; + width: 100%; + padding: ${({ theme }) => theme.sizeUnit * 2.5}px + ${({ theme }) => theme.sizeUnit * 4}px; + cursor: pointer; + color: ${({ theme }) => theme.colorPrimary}; + font-weight: ${({ theme }) => theme.fontWeightStrong}; + background: none; + border: none; + text-align: left; + transition: background ${({ theme }) => theme.motionDurationMid}; + + &:hover { + background: ${({ theme }) => theme.colorFillTertiary}; + } +`; + +const TabItem = styled.div<{ isActive: boolean }>` + display: flex; + align-items: center; + justify-content: space-between; + padding: ${({ theme }) => theme.sizeUnit * 2.5}px + ${({ theme }) => theme.sizeUnit * 4}px; + cursor: pointer; + background: ${({ theme, isActive }) => + isActive ? theme.colorFillSecondary : 'transparent'}; + border-left: 3px solid + ${({ theme, isActive }) => (isActive ? theme.colorPrimary : 'transparent')}; + transition: background ${({ theme }) => theme.motionDurationMid}; + + &:hover { + background: ${({ theme }) => theme.colorFillTertiary}; + + .action-btn { + opacity: 1; + } + } +`; + +const TabInfo = styled.div` + display: flex; + align-items: center; + gap: ${({ theme }) => theme.sizeUnit * 2}px; + flex: 1; + overflow: hidden; +`; + +const TabName = styled.span` + font-size: ${({ theme }) => theme.fontSize}px; + color: ${({ theme }) => theme.colorText}; + white-space: nowrap; + overflow: hidden; + text-overflow: ellipsis; + max-width: ${({ theme }) => theme.sizeUnit * 35}px; +`; + +const TabNameInput = styled.input` + width: 100%; + max-width: ${({ theme }) => theme.sizeUnit * 40}px; + font-size: ${({ theme }) => theme.fontSize}px; + color: ${({ theme }) => theme.colorText}; + background: ${({ theme }) => theme.colorBgContainer}; + border: 1px solid ${({ theme }) => theme.colorBorder}; + border-radius: ${({ theme }) => theme.borderRadius}px; + padding: 2px ${({ theme }) => theme.sizeUnit * 1.5}px; +`; + +const TabTimestamp = styled.span` + font-size: ${({ theme }) => theme.fontSizeSM}px; + color: ${({ theme }) => theme.colorTextQuaternary}; + white-space: nowrap; + flex-shrink: 0; +`; + +const ActionButtons = styled.div` + display: flex; + align-items: center; + gap: 2px; +`; + +const ActionButton = styled.button` + background: none; + border: none; + padding: ${({ theme }) => theme.sizeUnit}px; + cursor: pointer; + color: ${({ theme }) => theme.colorTextSecondary}; + opacity: 0; + transition: all ${({ theme }) => theme.motionDurationMid}; + display: flex; + align-items: center; + justify-content: center; + border-radius: ${({ theme }) => theme.borderRadius}px; + + &:hover, + &:focus-visible { + opacity: 1; + color: ${({ theme }) => theme.colorError}; + background: ${({ theme }) => theme.colorErrorBg}; + } +`; + +const Divider = styled.div` + height: 1px; + background: ${({ theme }) => theme.colorBorderSecondary}; + margin: ${({ theme }) => theme.sizeUnit}px 0; +`; + +const EmptyState = styled.div` + padding: ${({ theme }) => theme.sizeUnit * 5}px + ${({ theme }) => theme.sizeUnit * 4}px; + text-align: center; + color: ${({ theme }) => theme.colorTextSecondary}; + font-size: ${({ theme }) => theme.fontSizeSM}px; +`; + +const MINUTE_SECONDS = 60; +const HOUR_MINUTES = 60; +const DAY_HOURS = 24; +const WEEK_DAYS = 7; + +export const formatRelativeTime = (timestamp: number): string => { + const seconds = Math.floor((Date.now() - timestamp) / 1000); + if (seconds < MINUTE_SECONDS) { + return t('just now'); + } + const minutes = Math.floor(seconds / MINUTE_SECONDS); + if (minutes < HOUR_MINUTES) { + return t('%sm', String(minutes)); + } + const hours = Math.floor(minutes / HOUR_MINUTES); + if (hours < DAY_HOURS) { + return t('%sh', String(hours)); + } + const days = Math.floor(hours / DAY_HOURS); + if (days < WEEK_DAYS) { + return t('%sd', String(days)); + } + return new Date(timestamp).toLocaleDateString(undefined, { + month: 'short', + day: 'numeric', + }); +}; + +interface ChatTabsMenuProps { + tabs: ChatTab[]; + activeTabId: string; + onSelectTab: (tabId: string) => void; + onNewChat: () => void; + onDeleteTab: (tabId: string) => void; + onRenameTab: (tabId: string, name: string) => void; +} + +export const ChatTabsMenu = ({ + tabs, + activeTabId, + onSelectTab, + onNewChat, + onDeleteTab, + onRenameTab, +}: ChatTabsMenuProps) => { + const [editingTabId, setEditingTabId] = useState<string | null>(null); + const [editingName, setEditingName] = useState(''); + + const startEditing = useCallback((event: ReactMouseEvent, tab: ChatTab) => { + event.stopPropagation(); + setEditingTabId(tab.id); + setEditingName(tab.name); + }, []); + + const cancelEditing = useCallback(() => { + setEditingTabId(null); + setEditingName(''); + }, []); + + const commitRename = useCallback( + (tabId: string) => { + const trimmedName = editingName.trim(); + if (trimmedName) { + onRenameTab(tabId, trimmedName); + } + cancelEditing(); + }, + [cancelEditing, editingName, onRenameTab], + ); + + const menuContent = ( + <MenuContainer data-test="chat-tabs-menu"> + <MenuHeader>{t('Conversations')}</MenuHeader> + <NewChatButton type="button" onClick={onNewChat}> + <Icons.PlusOutlined iconSize="s" /> + <span>{t('New Chat')}</span> + </NewChatButton> + <Divider /> + {tabs.length === 0 ? ( + <EmptyState>{t('No conversations yet')}</EmptyState> + ) : ( + tabs.map(tab => ( + <TabItem Review Comment: Keyboard activation of the conversation row is still missing in this branch. The refresh does not change its semantics; this remains a base accessibility fix with a keyboard selection regression. ########## superset-frontend/src/features/ai/AiAssistantPanel.tsx: ########## @@ -0,0 +1,1150 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +/** + * @fileoverview The assistant panel. + * + * The host owns where this sits and, when docked, how wide it is, so there is no + * positioning and no resize handle here. What is here is the conversation: the + * header, the transcript, what the assistant is doing while it works, and the + * composer. + * + * The centre of the design is that a run is legible while it happens. An answer + * can take a minute of tool calls, and a spinner for a minute is indistinguishable + * from a hang, so reasoning streams into a preview, each step appends to a tool + * log, and a checkpoint stops the run with a countdown the user can act on. + */ + +import { + memo, + useCallback, + useEffect, + useRef, + useState, + useSyncExternalStore, +} from 'react'; +import type { Dispatch, SetStateAction } from 'react'; +import ReactMarkdown from 'react-markdown'; +import type { Components } from 'react-markdown'; +import { css, keyframes, styled, useTheme } from '@apache-superset/core/theme'; +import { t } from '@apache-superset/core/translation'; +import type { chat as chatApi } from '@apache-superset/core'; +import { + Button, + Input, + Loading, + Tooltip, + Typography, +} from '@superset-ui/core/components'; +import { Icons } from '@superset-ui/core/components/Icons'; +import { chat } from 'src/core/chat'; +import ChatAgentSelect from './components/ChatAgentSelect'; +import ChatTabsMenu from './components/ChatTabsMenu'; +import { REMARK_PLUGINS, useChatMarkdown } from './components/chatMarkdown'; +import { ThoughtProcess } from './components/ThoughtProcess'; +import { useChatBot } from './hooks/useChatBot'; +import { AI_ACTION_EVENT, type AiActionEvent } from './hooks/useAIAction'; +import type { PageContext } from './hooks/usePageContext'; +import type { ChatMessageWithMeta, CheckpointPayload } from './types'; + +/** + * How long a checkpoint waits before continuing on its own. A pause that blocks + * forever is worse than one that resolves optimistically: the user may have + * walked away, and the run should not be stranded. + */ +export const CHECKPOINT_TIMEOUT_SECONDS = 30; + +/** + * Closes a code fence the model has not finished writing. + * + * A streamed answer is parsed on every delta, so a fence arrives in pieces — + * "```", then "sql", then the query. Markdown with an odd number of fences + * renders the opening backticks literally and then reflows once the closing pair + * lands, which reads as the answer glitching. Balancing the count keeps each + * intermediate state a valid document. + */ +export const balanceCodeFences = (text: string): string => { + const fences = text.match(/^```/gm)?.length ?? 0; + return fences % 2 === 0 ? text : `${text}\n\`\`\``; +}; + +/** + * Whether a message carries the structured record of how it was answered, as + * opposed to only the flat log assembled from stream frames. + */ +const hasStructuredThinking = (message: ChatMessageWithMeta): boolean => + Boolean(message.toolCalls?.length || message.thoughts || message.pageContext); + +/** Milliseconds between typewriter frames, and characters per frame. */ +const TYPEWRITER_INTERVAL_MS = 18; +const TYPEWRITER_STEP = 3; + +/** + * The panel's own size as a floating overlay. + * + * Docked width belongs to the host and is not set here. Floating does need a size + * from somewhere, though — the floating host only stacks its children in a corner + * and gives them no dimensions — so these clamp the overlay to the viewport. + */ +const FLOATING_WIDTH_PX = 440; +const FLOATING_MAX_HEIGHT_VH = 70; + +/** + * The panel surface. + * + * Positioning is deliberately absent: the host places this, in both modes. What is + * here is the surface itself — a column that fills whatever box it is given, with a + * floating size for the mode where the host provides no box. + */ +const ChatPanelContainer = styled.div<{ floating: boolean }>` + display: flex; + flex-direction: column; + min-height: 0; + overflow: hidden; + background: ${({ theme }) => theme.colorBgElevated}; + ${({ floating, theme }) => + floating + ? css` + width: min( + ${FLOATING_WIDTH_PX}px, + calc(100vw - ${theme.sizeUnit * 12}px) + ); + height: ${FLOATING_MAX_HEIGHT_VH}vh; + border: 1px solid ${theme.colorBorderSecondary}; + border-radius: ${theme.borderRadiusLG}px; + box-shadow: ${theme.boxShadow}; + ` + : css` + width: 100%; + height: 100%; + `} +`; + +const ChatHeader = styled.div` + display: flex; + justify-content: space-between; + align-items: center; + padding: ${({ theme }) => theme.sizeUnit * 3}px + ${({ theme }) => theme.sizeUnit * 4}px; + background: ${({ theme }) => theme.colorBgContainer}; + border-bottom: 1px solid ${({ theme }) => theme.colorBorderSecondary}; + font-weight: ${({ theme }) => theme.fontWeightStrong}; + font-size: ${({ theme }) => theme.fontSizeLG}px; + color: ${({ theme }) => theme.colorTextHeading}; + flex-shrink: 0; + gap: ${({ theme }) => theme.sizeUnit * 2}px; +`; + +const HeaderGroup = styled.div` + display: flex; + align-items: center; + min-width: 0; +`; + +const HeaderTitle = styled.span` + overflow: hidden; + text-overflow: ellipsis; + white-space: nowrap; +`; + +const ChatMessages = styled.div` + flex: 1; + min-height: 0; + padding: ${({ theme }) => theme.sizeUnit * 4}px; + overflow-y: auto; + + &::-webkit-scrollbar { + width: 4px; + } + + &::-webkit-scrollbar-track { + background: ${({ theme }) => theme.colorBgContainer}; + border-radius: 2px; + } + + &::-webkit-scrollbar-thumb { + background: ${({ theme }) => theme.colorFillSecondary}; + border-radius: 2px; + } +`; + +const MessageBubble = styled.div<{ variant: 'user' | 'assistant' }>` + margin-bottom: ${({ theme }) => theme.sizeUnit * 3}px; + display: flex; + flex-direction: column; + align-items: ${({ variant }) => + variant === 'user' ? 'flex-end' : 'flex-start'}; +`; + +const MessageContent = styled.div<{ variant: 'user' | 'assistant' }>` + max-width: 85%; + padding: ${({ theme }) => theme.sizeUnit * 3}px + ${({ theme }) => theme.sizeUnit * 4}px; + border-radius: ${({ theme }) => theme.borderRadiusLG * 2}px; + background: ${({ theme, variant }) => + variant === 'user' ? theme.colorPrimary : theme.colorBgContainer}; + color: ${({ theme, variant }) => + variant === 'user' ? theme.colorTextLightSolid : theme.colorText}; + font-size: ${({ theme }) => theme.fontSize}px; + line-height: 1.5; + border: ${({ theme, variant }) => + variant === 'assistant' + ? `1px solid ${theme.colorBorderSecondary}` + : 'none'}; + box-shadow: ${({ theme }) => theme.boxShadowTertiary}; + overflow-wrap: anywhere; + + p { + margin: 0; + } + + p:not(:last-child) { + margin-bottom: ${({ theme }) => theme.sizeUnit * 2}px; + } + + a { + color: ${({ theme, variant }) => + variant === 'user' ? theme.colorTextLightSolid : theme.colorPrimary}; + text-decoration: underline; + } + + code { + background: ${({ theme, variant }) => + variant === 'user' ? theme.colorPrimaryActive : theme.colorFillTertiary}; + padding: 2px ${({ theme }) => theme.sizeUnit * 1.5}px; + border-radius: ${({ theme }) => theme.borderRadius}px; + font-size: ${({ theme }) => theme.fontSizeSM}px; + font-family: ${({ theme }) => theme.fontFamilyCode}; + } + + pre { + background: ${({ theme }) => theme.colorFillQuaternary}; + padding: ${({ theme }) => theme.sizeUnit * 3}px; + border-radius: ${({ theme }) => theme.borderRadius}px; + overflow-x: auto; + margin: ${({ theme }) => theme.sizeUnit * 2}px 0; + border: 1px solid ${({ theme }) => theme.colorBorderSecondary}; + } + + pre code { + background: none; + padding: 0; + } +`; + +const MessageActions = styled.div` + display: flex; + gap: ${({ theme }) => theme.sizeUnit}px; + margin-top: ${({ theme }) => theme.sizeUnit}px; +`; + +const ActionButton = styled(Button)` + &&& { + padding: 2px ${({ theme }) => theme.sizeUnit * 1.5}px; + height: ${({ theme }) => theme.sizeUnit * 6}px; + font-size: ${({ theme }) => theme.fontSizeSM}px; + } + + /* The recorded verdict keeps its colour while disabled. Both thumbs lock once + a rating exists, and the default disabled grey would hide which one the + user picked — the state matters more here than the affordance. */ + &&&.is-active, + &&&.is-active:disabled, + &&&.is-active[disabled] { + color: ${({ theme }) => theme.colorPrimary}; + } +`; + +const LiveAnswer = styled.div` + margin-top: ${({ theme }) => theme.sizeUnit * 2}px; + color: ${({ theme }) => theme.colorText}; + font-size: ${({ theme }) => theme.fontSize}px; + line-height: 1.5; + overflow-wrap: anywhere; + + p { + margin: 0; + } + + p:not(:last-child) { + margin-bottom: ${({ theme }) => theme.sizeUnit * 2}px; + } +`; + +const thinkingPulse = keyframes` + 0% { + opacity: 0.45; + } + 50% { + opacity: 1; + } + 100% { + opacity: 0.45; + } +`; + +const ThinkingPreview = styled.div<{ isLive?: boolean }>` + color: ${({ theme }) => theme.colorTextTertiary}; + font-size: ${({ theme }) => theme.fontSizeSM}px; + white-space: pre-wrap; + ${({ isLive }) => + isLive && + css` + animation: ${thinkingPulse} 1.8s ease-in-out infinite; + `} +`; + +const ThinkingDetails = styled.details` + margin-bottom: ${({ theme }) => theme.sizeUnit * 2}px; + color: ${({ theme }) => theme.colorTextTertiary}; + font-size: ${({ theme }) => theme.fontSizeSM}px; + + summary { + cursor: pointer; + user-select: none; + color: ${({ theme }) => theme.colorTextTertiary}; + margin-bottom: ${({ theme }) => theme.sizeUnit}px; + } +`; + +const CheckpointDivider = styled.div` + display: flex; + align-items: center; + gap: ${({ theme }) => theme.sizeUnit * 3}px; + margin: ${({ theme }) => theme.sizeUnit * 4}px 0 + ${({ theme }) => theme.sizeUnit * 3}px; + + &::before, + &::after { + content: ''; + flex: 1; + height: 1px; + background: ${({ theme }) => theme.colorBorderSecondary}; + } +`; + +const CountdownBadge = styled.span` + font-size: ${({ theme }) => theme.fontSizeSM}px; + font-weight: ${({ theme }) => theme.fontWeightStrong}; + font-variant-numeric: tabular-nums; + color: ${({ theme }) => theme.colorTextSecondary}; + white-space: nowrap; +`; + +const CheckpointContent = styled.div` + font-size: ${({ theme }) => theme.fontSizeSM}px; + color: ${({ theme }) => theme.colorText}; + line-height: 1.5; +`; + +const CheckpointTaskList = styled.ul` + margin: ${({ theme }) => theme.sizeUnit * 1.5}px 0; + padding-left: ${({ theme }) => theme.sizeUnit * 4.5}px; + font-size: ${({ theme }) => theme.fontSizeSM}px; + color: ${({ theme }) => theme.colorTextSecondary}; + + li { + margin-bottom: 2px; + } +`; + +const CheckpointEstimate = styled.div` + font-size: ${({ theme }) => theme.fontSizeSM}px; + color: ${({ theme }) => theme.colorTextTertiary}; + margin-top: ${({ theme }) => theme.sizeUnit}px; +`; + +const CheckpointActions = styled.div` + display: flex; + gap: ${({ theme }) => theme.sizeUnit * 2}px; + margin-top: ${({ theme }) => theme.sizeUnit * 2.5}px; +`; + +const ChatInput = styled.div` + padding: ${({ theme }) => theme.sizeUnit * 4}px; + border-top: 1px solid ${({ theme }) => theme.colorBorderSecondary}; + background: ${({ theme }) => theme.colorBgContainer}; + flex-shrink: 0; +`; + +const InputContainer = styled.div` + display: flex; + gap: ${({ theme }) => theme.sizeUnit * 2}px; + align-items: flex-end; +`; + +const QuickPromptsRow = styled.div<{ hasContent: boolean }>` + display: flex; + align-items: center; + gap: ${({ theme }) => theme.sizeUnit * 2}px; + flex-wrap: wrap; + margin-bottom: ${({ theme, hasContent }) => + hasContent ? `${theme.sizeUnit * 2.5}px` : '0'}; + min-height: ${({ theme, hasContent }) => + hasContent ? `${theme.sizeUnit * 6}px` : '0'}; +`; + +const QuickPromptChip = styled(Button)` + &&& { + width: fit-content; + max-width: 100%; + height: auto; + white-space: normal; + text-align: left; + line-height: 1.35; + word-break: break-word; + } +`; + +const PageContextRow = styled.div` + display: flex; + align-items: center; + margin-bottom: ${({ theme }) => theme.sizeUnit * 1.5}px; +`; + +const PageContextPill = styled.button<{ isActive: boolean }>` + display: inline-flex; + align-items: center; + gap: ${({ theme }) => theme.sizeUnit * 1.5}px; + padding: 3px ${({ theme }) => theme.sizeUnit * 2}px; + border-radius: ${({ theme }) => theme.borderRadiusLG}px; + border: 1px solid + ${({ theme, isActive }) => + isActive ? theme.colorPrimary : theme.colorBorderSecondary}; + background: ${({ theme, isActive }) => + isActive ? theme.colorPrimaryBg : theme.colorFillQuaternary}; + color: ${({ theme, isActive }) => + isActive ? theme.colorPrimary : theme.colorTextTertiary}; + font-size: ${({ theme }) => theme.fontSizeSM}px; + cursor: pointer; + transition: all ${({ theme }) => theme.motionDurationMid}; + max-width: ${({ theme }) => theme.sizeUnit * 62}px; + white-space: nowrap; + + &:hover { + border-color: ${({ theme }) => theme.colorPrimary}; + } +`; + +const PillLabel = styled.span` + overflow: hidden; + text-overflow: ellipsis; +`; + +const PillDot = styled.span<{ isActive: boolean }>` + width: 7px; + height: 7px; + border-radius: 50%; + flex-shrink: 0; + background: ${({ theme, isActive }) => + isActive ? theme.colorPrimary : theme.colorTextQuaternary}; + transition: background ${({ theme }) => theme.motionDurationMid}; +`; + +const EmptyState = styled.div` + text-align: center; + color: ${({ theme }) => theme.colorTextSecondary}; + font-size: ${({ theme }) => theme.fontSizeLG}px; + margin-top: ${({ theme }) => theme.sizeUnit * 15}px; + padding: 0 ${({ theme }) => theme.sizeUnit * 10}px; + line-height: 1.5; +`; + +/** + * What the page-context pill says it is sending. + * + * Naming the specific chart or dashboard is the whole point: "page context" tells + * the user nothing about whether the assistant can see the thing they are asking + * about. + */ +export const getPageContextLabel = (context: PageContext): string => { + switch (context.pageType) { + case 'sqllab': + return context.sqlContext?.activeEditor?.name || t('SQL Lab'); + case 'dashboard': + return context.dashboardContext?.title + ? t('Dashboard: %s', context.dashboardContext.title) + : t('Dashboard'); + case 'explore': + case 'chart': { + const name = + context.chartContext?.chartName || + context.chartContext?.slice?.slice_name; + return name ? t('Chart: %s', name) : t('Chart Explorer'); + } + case 'home': + return t('Home'); + default: + return context.pathname; + } +}; + +interface MemoizedChatMessageProps { + message: ChatMessageWithMeta; + feedback: 'like' | 'dislike' | null; + /** Per-message SQL expansion state, present only to invalidate the memo. */ + sqlBlockExpandState: boolean[] | undefined; + createMarkdownComponents: (messageId: string) => Components; + /** Keeps this message's thought process open, for the run that just ended. */ + keepThinkingOpen?: boolean; + onFeedback: (messageId: string, feedback: 'like' | 'dislike') => void; + onCopy: (content: string) => Promise<void>; +} + +/** + * One turn. + * + * Memoized on content rather than identity because a streaming run re-renders the + * panel on every frame, and re-parsing the markdown of every earlier message each + * time made long conversations visibly janky. + */ +const MemoizedChatMessage = memo( + ({ + message, + feedback, + createMarkdownComponents, + keepThinkingOpen, + onFeedback, + onCopy, + }: MemoizedChatMessageProps) => { + const markdownComponents = createMarkdownComponents(message.id); + return ( + <MessageBubble + variant={message.role} + data-test="chat-message" + data-role={message.role} + > + <MessageContent variant={message.role}> + {message.role === 'assistant' && + // Steps are rendered structurally where the server has sent them, + // and as the flat log otherwise. The two are not alternatives so + // much as two stages: a turn assembled from stream frames only has + // the text, and gains its structure when the transcript is re-read + // from the server once the run ends. Any one of the structured + // fields is enough — a turn can have page context and no tool calls, + // or reasoning and no steps. + (hasStructuredThinking(message) ? ( + <ThoughtProcess + reasoning={message.thoughts} + pageContext={message.pageContext} + toolCalls={message.toolCalls} + markdownComponents={markdownComponents} + // Left open for the turn that has only just finished, so the + // section does not collapse and pull the answer up the panel the + // moment the user starts reading it. + defaultOpen={keepThinkingOpen} + /> + ) : ( + message.thinking && ( + <ThinkingDetails> + <summary>{t('Thought process')}</summary> + {/* The tool log is markdown too, so a SQL step gets the same + highlighting and "Run in SQL Lab" as SQL in the answer. */} + <ReactMarkdown + components={markdownComponents} + remarkPlugins={REMARK_PLUGINS} + > + {message.thinking} + </ReactMarkdown> + </ThinkingDetails> + ) + ))} + <ReactMarkdown + components={markdownComponents} + remarkPlugins={REMARK_PLUGINS} + > + {message.content} + </ReactMarkdown> + </MessageContent> + {message.role === 'assistant' && ( + <MessageActions> + {/* Both thumbs stay available so a mis-click can be corrected. The + recorded verdict is the coloured one, and only it is inert — + there is no endpoint to withdraw a rating, so pressing it again + would do nothing and should not look like it could. */} + <ActionButton + buttonStyle="link" + icon={<Icons.LikeOutlined iconSize="s" />} + onClick={() => onFeedback(message.id, 'like')} + className={feedback === 'like' ? 'is-active' : ''} + aria-label={t('Good response')} + aria-pressed={feedback === 'like'} + tooltip={ + feedback === 'like' + ? t('You marked this a good response') + : t('Good response') + } + disabled={message.pending || feedback === 'like'} + /> + <ActionButton + buttonStyle="link" + icon={<Icons.DislikeOutlined iconSize="s" />} + onClick={() => onFeedback(message.id, 'dislike')} + className={feedback === 'dislike' ? 'is-active' : ''} + aria-label={t('Bad response')} + aria-pressed={feedback === 'dislike'} + tooltip={ + feedback === 'dislike' + ? t('You marked this a bad response') + : t('Bad response') + } + disabled={message.pending || feedback === 'dislike'} + /> + <ActionButton + buttonStyle="link" + icon={<Icons.CopyOutlined iconSize="s" />} + onClick={() => { + onCopy(message.content).catch(() => { + // Clipboard access can be refused; the button then does + // nothing rather than raising at the user. + }); + }} + aria-label={t('Copy to clipboard')} + tooltip={t('Copy to clipboard')} + /> + </MessageActions> + )} + </MessageBubble> + ); + }, + (previous, next) => + previous.message.id === next.message.id && + previous.message.content === next.message.content && + previous.message.thinking === next.message.thinking && + // Compared by identity, not deeply: the server's copy of a message arrives as + // a fresh array, which is exactly the case that has to re-render, and a deep + // compare of every step on every streamed frame would cost more than the + // render it avoids. + previous.message.toolCalls === next.message.toolCalls && + previous.message.thoughts === next.message.thoughts && + previous.message.pageContext === next.message.pageContext && + previous.message.pending === next.message.pending && + previous.feedback === next.feedback && + previous.keepThinkingOpen === next.keepThinkingOpen && + previous.sqlBlockExpandState === next.sqlBlockExpandState, +); +MemoizedChatMessage.displayName = 'MemoizedChatMessage'; + +/** + * A pause the user can act on. + * + * The countdown exists so the pause cannot strand the run; continuing on expiry is + * the same decision the user would most likely have made. + */ +export const CheckpointSection = ({ + checkpoint, + onContinue, + onCancel, +}: { + checkpoint: CheckpointPayload; + onContinue: () => void; + onCancel: () => void; +}) => { + const [secondsLeft, setSecondsLeft] = useState( + checkpoint.seconds_remaining ?? CHECKPOINT_TIMEOUT_SECONDS, + ); + // Held in a ref so the interval is installed once; re-installing it on every + // render would reset the countdown and it would never reach zero. + const onContinueRef = useRef(onContinue); + onContinueRef.current = onContinue; + + useEffect(() => { + const interval = setInterval(() => { + setSecondsLeft(previous => { + if (previous <= 1) { + clearInterval(interval); + onContinueRef.current(); + return 0; + } + return previous - 1; + }); + }, 1000); + return () => clearInterval(interval); + }, []); + + const minutes = Math.floor(secondsLeft / 60); + const seconds = secondsLeft % 60; + + return ( + <div data-test="chat-checkpoint"> + <CheckpointDivider> + <CountdownBadge>{`${minutes}:${seconds + .toString() + .padStart(2, '0')}`}</CountdownBadge> + </CheckpointDivider> + <CheckpointContent>{checkpoint.summary}</CheckpointContent> + {checkpoint.remaining_tasks && checkpoint.remaining_tasks.length > 0 && ( + <CheckpointTaskList> + {checkpoint.remaining_tasks.map(task => ( + <li key={task}>{task}</li> + ))} + </CheckpointTaskList> + )} + {checkpoint.estimated_duration && ( + <CheckpointEstimate> + {t('Est. ~%s', checkpoint.estimated_duration)} + </CheckpointEstimate> + )} + <CheckpointActions> + <Button buttonSize="small" buttonStyle="primary" onClick={onContinue}> + {t('Continue')} + </Button> + <Button buttonSize="small" onClick={onCancel}> + {t('Cancel')} + </Button> + </CheckpointActions> + </div> + ); +}; + +/** + * Reveals text a few characters at a time. + * + * Reasoning arrives in bursts, and a preview that jumps a paragraph at a time is + * hard to read; this smooths it without holding anything back. A target that is + * not an extension of what is shown (a replacement, not an append) is applied at + * once, because animating a rewrite would show text that was never sent. + */ +const useTypewriter = ( + target: string, + setter: Dispatch<SetStateAction<string>>, + streaming: boolean, +) => { + useEffect(() => { + if (!streaming) { + setter(target); + return undefined; + } + if (!target) { + setter(''); + return undefined; + } + + let cancelled = false; + let timeoutId: ReturnType<typeof setTimeout> | undefined; + + const tick = () => { + if (cancelled) { + return; + } + setter(previous => { + if (previous === target) { + return previous; + } + if (!target.startsWith(previous)) { + return target; + } + const remaining = target.length - previous.length; + const step = Math.min(TYPEWRITER_STEP, Math.max(1, remaining)); + return target.slice(0, previous.length + step); + }); + timeoutId = setTimeout(tick, TYPEWRITER_INTERVAL_MS); + }; + + timeoutId = setTimeout(tick, TYPEWRITER_INTERVAL_MS); + return () => { + cancelled = true; + if (timeoutId) { + clearTimeout(timeoutId); + } + }; + }, [target, streaming, setter]); +}; + +/** + * Tracks the host's display mode, which the host changes of its own accord, so the + * panel reads it rather than mirroring it. + */ +const useDisplayMode = (): chatApi.DisplayMode => + useSyncExternalStore( + useCallback((onChange: () => void) => { + const subscription = chat.onDidChangeDisplayMode(onChange); + return () => subscription.dispose(); + }, []), + chat.getDisplayMode, + ); + +export const AiAssistantPanel = () => { + const theme = useTheme(); + const displayMode = useDisplayMode(); + const { + chatTabs, + activeTabId, + activeTab, + threadsLoaded, + handleNewChat, + handleSelectTab, + handleDeleteTab, + handleRenameTab, + messages, + inputValue, + setInputValue, + handleKeyDown, + inputRef, + messagesEndRef, + isLoading, + isStreamingResponse, + liveThoughts, + liveToolCalls, + livePageContext, + liveAnswer, + checkpoint, + activeRunStatus, + error, + sendMessage, + handleCancelRun, + handleCheckpointContinue, + handleFeedback, + messageFeedback, + justCompletedId, + quickPrompts, + loadQuickPrompts, + applyQuickPrompt, + agents, + selectedAgent, + setSelectedAgent, + pageContext, + includePageContext, + toggleIncludePageContext, + } = useChatBot(); + + const { + createMarkdownComponents, + copyToClipboard, + expandedSqlBlocksByMessage, + } = useChatMarkdown(); + + // The live step list is not attached to a stored message, so it gets its own + // set of markdown components under a stable synthetic id — that id is what + // keys the per-message SQL expand state. + const liveMarkdownComponents = createMarkdownComponents('live-run'); + + const [typedThoughts, setTypedThoughts] = useState(''); + useTypewriter(liveThoughts, setTypedThoughts, isStreamingResponse); + + /** + * A prompt queued for a conversation that does not exist yet. + * + * An AI action opens a fresh conversation and sends into it, and the send has to + * wait for that conversation to become the active one — otherwise it lands in + * whichever was open before. + */ + const [pendingSend, setPendingSend] = useState<{ + tabId: string; + prompt: string; + systemPrompt?: string; + } | null>(null); + + useEffect(() => { + const handleAiAction = (event: Event) => { + const { detail } = event as AiActionEvent; + const prompt = detail?.prompt?.trim(); + if (!prompt) { + return; + } + chat.open(); + handleNewChat().then(newTabId => { + setPendingSend({ + tabId: newTabId, + prompt, + systemPrompt: detail.systemPrompt?.trim() || undefined, + }); + }); + }; + window.addEventListener(AI_ACTION_EVENT, handleAiAction); + return () => { + window.removeEventListener(AI_ACTION_EVENT, handleAiAction); + }; + }, [handleNewChat]); + + useEffect(() => { + if (!pendingSend || activeTabId !== pendingSend.tabId) { + return; + } + if (!chatTabs.some(tab => tab.id === pendingSend.tabId)) { + return; + } + const { prompt, systemPrompt } = pendingSend; + setPendingSend(null); + sendMessage(prompt, systemPrompt); + }, [pendingSend, activeTabId, chatTabs, sendMessage]); + + const onSelectTab = useCallback( + (tabId: string) => { + handleSelectTab(tabId); + }, + [handleSelectTab], + ); + + // The row collapses to nothing when there is nothing to suggest, rather than + // holding empty space above the composer. + const hasQuickPromptContent = quickPrompts.length > 0; + + return ( + <ChatPanelContainer + floating={displayMode !== 'panel'} + data-test="ai-assistant-panel" + > + <ChatHeader> + <HeaderGroup> + <ChatTabsMenu + tabs={chatTabs} + activeTabId={activeTabId} + onSelectTab={onSelectTab} + onNewChat={() => { + handleNewChat(); + }} + onDeleteTab={tabId => { + handleDeleteTab(tabId); + }} + onRenameTab={handleRenameTab} + /> + <HeaderTitle>{activeTab?.name || t('AI assistant')}</HeaderTitle> + </HeaderGroup> + <HeaderGroup> + {agents.length > 1 && ( + <ChatAgentSelect + agents={agents} + selectedAgent={selectedAgent} + onChange={setSelectedAgent} + /> + )} + {/* Docking is the host's, not ours: it decides where the panel goes, so + this only asks it to change mode. */} + <Button + buttonStyle="link" + icon={ + displayMode === 'panel' ? ( + <Icons.CompressOutlined iconSize="m" /> + ) : ( + <Icons.ExpandOutlined iconSize="m" /> + ) + } + onClick={() => + chat.setDisplayMode( + displayMode === 'panel' ? 'floating' : 'panel', + ) + } + aria-label={ + displayMode === 'panel' + ? t('Undock the assistant') + : t('Dock the assistant') + } + tooltip={ + displayMode === 'panel' + ? t('Undock the assistant') + : t('Dock the assistant') + } + /> + <Button + buttonStyle="link" + icon={<Icons.CloseOutlined iconSize="m" />} + onClick={() => chat.close()} + aria-label={t('Close the assistant')} + /> + </HeaderGroup> + </ChatHeader> + + <ChatMessages data-test="chat-messages"> Review Comment: Completed-response announcements are still an open base accessibility requirement. I have not claimed the visual stream provides screen-reader feedback; live/log semantics need a test that avoids announcing every token. ########## superset/ai/tools/context.py: ########## @@ -0,0 +1,356 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +""" +Reading what an existing chart or dashboard is made of. + +These answer "what does this thing already show?", which is what a question like +"why is this number different from the dashboard" needs. Both are reads of +Superset's metadata; neither queries a warehouse. + +Two authorization gates apply to each, and both are needed: + +* the DAO's ``base_filter`` decides whether the object is visible at all, so an + id the user may not see is reported as not found rather than forbidden; and +* ``security_manager.raise_for_access`` re-checks the specific object, which is + what catches a chart that is listed but whose underlying dataset the user has + lost access to. + +The chart's ``params`` blob is not returned wholesale. It is a large, +free-form, user-authored structure whose bulk is styling, so only the fields +that describe *what is measured* are lifted out of it. +""" + +from __future__ import annotations + +import logging +from typing import Any, ClassVar + +from superset.ai.tools.base import AITool, ToolError, ToolOutput + +logger = logging.getLogger(__name__) + +#: Charts summarised for one dashboard. A dashboard with more than this is +#: usually a wall of tiles, and the model does not need every one to answer a +#: question about it. +MAX_CHARTS = 50 + +#: Columns and metrics listed for a chart's dataset. +MAX_DATASET_FIELDS = 100 + +#: Characters of a virtual dataset's SQL that are returned. Enough to see the +#: shape of the query and its joins; a longer body is better read by asking +#: about the dataset directly. +MAX_SQL_CHARS = 4000 + +#: ``params`` keys worth showing. These are the ones that say what the chart +#: measures and how it is sliced; everything else in the blob is presentation. +_MEANINGFUL_PARAM_KEYS = ( + "metrics", + "metric", + "groupby", + "columns", + "all_columns", + "adhoc_filters", + "granularity_sqla", + "time_grain_sqla", + "time_range", + "row_limit", + "order_desc", + "percent_metrics", + "series_limit", + "series_limit_metric", +) + + +def _positive_id(value: Any, field: str) -> int: + """Validate an id argument.""" + if not isinstance(value, int) or isinstance(value, bool) or value < 1: + raise ToolError(f"{field!r} must be a positive integer.") + return value + + +def _untrusted(value: Any) -> Any: + """Wrap user-authored free text so it cannot pose as instructions.""" + from superset.mcp_service.utils.sanitization import sanitize_for_llm_context + + if value is None: + return None + return sanitize_for_llm_context(value) + + +def _load_chart(chart_id: int) -> Any: + """ + Fetch a chart the user may see, or refuse. + + ``ChartDAO.find_by_id`` applies ``ChartFilter``; the explicit + ``raise_for_access`` then re-checks this specific chart, which is the gate + that catches a chart whose dataset access has since been revoked. + """ + from superset import security_manager + from superset.daos.chart import ChartDAO + from superset.exceptions import SupersetSecurityException + + chart = ChartDAO.find_by_id(chart_id) Review Comment: The FAB resource-capability parity question remains open; the refresh does not add these reader checks. It needs an explicit custom-role regression against the SECURITY.md role/capability model in the base implementation, not an assumption that datasource visibility alone settles the API contract. ########## superset/ai/policy.py: ########## @@ -0,0 +1,364 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +""" +Guards applied to every tool call before it runs. + +These bound blast radius. They are **not** an authorization layer: a tool that +returns or mutates a specific data-bearing object still has to perform its own +``security_manager.raise_for_access(...)`` check. A policy answers "should this +shape of call be attempted at all", which is a cheaper and coarser question. + +Policies are configured as dotted paths in ``AI_AGENT_TOOL_POLICIES`` so a +deployment can add its own without forking. +""" + +from __future__ import annotations + +import logging +import re +from abc import ABC, abstractmethod +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Any + +logger = logging.getLogger(__name__) + +#: A bare identifier, or dotted parts thereof. Deliberately strict: anything +#: with whitespace, quotes, semicolons or parentheses is rejected rather than +#: escaped, because a tool that needs to escape an identifier is a tool that is +#: building SQL by concatenation. +_SAFE_IDENTIFIER = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*(\.[A-Za-z_][A-Za-z0-9_]*)*$") + +#: Argument names understood to carry identifiers rather than free text. +_IDENTIFIER_ARGUMENTS = frozenset( + {"table", "table_name", "schema", "schema_name", "catalog", "column", "columns"} +) + + +@dataclass(frozen=True) +class Denial: + """ + A refusal to run a tool call. + + ``reason`` is shown to the model, so it should say what would be acceptable + instead. A model that is told "only read-only SQL is allowed" rewrites its + query; a model that is told "denied" retries the same thing. + """ + + reason: str + + +class ToolPolicy(ABC): + """A pre-execution guard over a single tool call.""" + + #: Identifies the policy in logs. + name: str = "policy" + + @abstractmethod + def check( + self, + tool_name: str, + arguments: dict[str, Any], + ) -> Denial | None: + """ + Inspect a pending call. + + Return ``None`` to allow, or a :class:`Denial` to block. A policy that + does not apply to ``tool_name`` returns ``None``. + """ + + +class ReadOnlySqlPolicy(ToolPolicy): + """ + Refuse anything that is not a read. + + Correctness here rests on Superset's own parser rather than a prefix or + keyword match. A regex over the leading token is defeated by a leading + comment, a CTE that wraps a DML statement, ``EXPLAIN ANALYZE DELETE``, and + multi-statement scripts — all of which the parser handles because the rest + of Superset already depends on it for the same decision. + """ + + name = "read_only_sql" + + #: Tools whose payload is SQL to execute. + sql_tools = frozenset( + { + "execute_sql", + "validate_sql", + "run_scoped_sql", + "create_virtual_dataset", + } + ) + + #: Argument names that may carry the SQL. + sql_arguments = ("sql", "query") + + #: Introspection commands permitted even when the parser cannot model them. + #: + #: Most dialects surface ``EXPLAIN`` and ``SHOW`` as an opaque catch-all + #: node, so a blanket "refuse what we cannot parse" rule would also refuse + #: the schema and query-plan inspection an analysis agent legitimately + #: needs. Enumerating them keeps the default deny for everything else. + read_only_commands = frozenset({"EXPLAIN", "SHOW", "DESCRIBE", "DESC"}) + + def check(self, tool_name: str, arguments: dict[str, Any]) -> Denial | None: + if tool_name not in self.sql_tools: + return None + + sql = self._extract_sql(arguments) + if sql is None: + return Denial( + f"{tool_name} requires a 'sql' argument containing the statement " + f"to run." + ) + if not sql.strip(): + return Denial("The 'sql' argument is empty.") + + engine = self._engine(arguments) + + try: + from superset.sql.parse import SQLScript + + script = SQLScript(sql, engine=engine) + except Exception: # pylint: disable=broad-except + # Unparseable SQL cannot be shown to be read-only, so it is + # refused. Logged rather than surfaced: parser errors can quote + # arbitrary query text back to the caller. + logger.info("Refusing unparseable SQL from tool %s", tool_name) + return Denial( + "That SQL could not be parsed. Send a single, syntactically " + "valid read-only statement." + ) + + # Checked per statement so a write cannot ride along behind a read. + for statement in script.statements: + if statement.is_mutating(): + return Denial( + "Only read-only SQL is allowed. Rewrite this as a SELECT — " + "statements that modify data or schema are refused." + ) + + # Anything the parser could not model is refused unless every statement + # is recognisably an introspection command. The mutation check above has + # already run, but it cannot reason about an opaque node on every + # dialect, so this is the fail-closed half of the decision. + if script.has_unparseable_statement and not all( + self._is_read_only_command(statement) for statement in script.statements + ): + return Denial( + "That SQL contains a statement this tool cannot verify as " + "read-only. Send a plain SELECT." + ) + return None + + def _engine(self, arguments: dict[str, Any]) -> str: + """Resolve the parser dialect from the selected Superset database.""" + if engine := arguments.get("engine"): + return str(engine) + + database_id = arguments.get("database_id") + if not isinstance(database_id, int) or isinstance(database_id, bool): + return "" + + try: + from superset.daos.database import DatabaseDAO + + database = DatabaseDAO.find_by_id(database_id) + if database is not None: + return str(database.db_engine_spec.engine or "") + except Exception: # pylint: disable=broad-except + logger.debug("Could not resolve SQL dialect for database %s", database_id) + return "" + + def _is_read_only_command(self, statement: Any) -> bool: + """Whether a statement is one of the permitted introspection commands.""" + try: + text = statement.format(comments=False) + except Exception: # pylint: disable=broad-except + # A statement that will not even render is not one we can vouch for. + return False + leading = text.strip().split(None, 1) + if not leading: + return False + return leading[0].upper() in self.read_only_commands + + def _extract_sql(self, arguments: dict[str, Any]) -> str | None: + """Pull the SQL payload out of whichever argument carries it.""" + for key in self.sql_arguments: + value = arguments.get(key) + if isinstance(value, str): + return value + return None + + +class IdentifierPolicy(ToolPolicy): + """ + Refuse identifiers that are not plain names. + + Tools should resolve names against registered metadata rather than splice + them into SQL. This policy is the backstop for the ones that take a name as + an argument: rejecting anything unusual is safer than trying to quote it, + because a value needing quoting signals string-built SQL underneath. + """ + + name = "identifier" + + def check(self, tool_name: str, arguments: dict[str, Any]) -> Denial | None: + for key, value in arguments.items(): + if key not in _IDENTIFIER_ARGUMENTS: + continue + for candidate in self._as_identifiers(value): + # A wildcard is a legitimate column selector. + if candidate == "*": + continue + if not _SAFE_IDENTIFIER.match(candidate): + return Denial( + f"{key!r} must be a plain identifier; {candidate!r} is " + f"not accepted." + ) + return None + + def _as_identifiers(self, value: Any) -> list[str]: + """Normalise the several shapes an identifier argument arrives in.""" + if isinstance(value, str): + # Comma-separated lists are common for column arguments. + return [part.strip() for part in value.split(",") if part.strip()] + if isinstance(value, (list, tuple)): + return [str(item).strip() for item in value if str(item).strip()] + return [] + + +class ForeignToolPolicy(ToolPolicy): + """ + Refuse SQL execution offered by an external MCP server. + + Superset's two SQL controls both act on statements Superset itself runs. + :class:`ReadOnlySqlPolicy` parses the statement before execution and refuses + anything that mutates; ``security_manager`` then authorizes the specific + database, catalog, schema and table the statement touches. Neither can apply + to SQL run by a third party: Superset never sees the statement, cannot know + which datasource it reached, and holds no permission mapping for that + server's data. Allowing a foreign tool to run SQL therefore does not widen + the surface so much as remove the two controls that make ``execute_sql`` + safe, and it does so without any signal that it has happened. + + Only names that advertise query execution are refused. A foreign tool that + searches a catalog, reads metadata or fetches a document is untouched, which + is the whole point of the extension point. + + The matcher is deliberately broad, and configurable two ways. A deployment + satisfied that its own servers enforce equivalent controls sets + ``AI_AGENT_MCP_DENY_FOREIGN_SQL = False``. One that wants a different match + subclasses and sets :attr:`sql_name_fragments`, or configures a subclass that + passes ``fragments`` to :meth:`__init__`. + + Built-in tools are not this policy's business: they pass through untouched + for :class:`ReadOnlySqlPolicy` to judge. + """ + + name = "foreign_tool" + + #: Matched case-insensitively as substrings of the remote tool's own name — + #: the part after ``mcp__<server>__``. Substrings rather than exact names + #: because a server is free to call the same capability ``sql_query``, + #: ``runQuery`` or ``execute_sql_async``, and the failure that matters is the + #: one where an unfamiliar spelling slips through. + sql_name_fragments: tuple[str, ...] = ("execute_sql", "run_sql", "query") + + def __init__(self, fragments: Sequence[str] | None = None) -> None: + source = self.sql_name_fragments if fragments is None else fragments + self.fragments = tuple(fragment.lower() for fragment in source if fragment) + + def check(self, tool_name: str, arguments: dict[str, Any]) -> Denial | None: + from superset.ai.mcp.config import split_foreign_tool_name + + parts = split_foreign_tool_name(tool_name) + if parts is None: + return None + if not self._enabled(): + return None + + _, remote_name = parts + haystack = remote_name.lower() + if any(fragment in haystack for fragment in self.fragments): Review Comment: The foreign-tool policy question remains open. A tool name is not a capability guarantee, and a server-supplied schema is not an authorization boundary either. No additional SQL-name heuristic was added here; the base needs an explicit operator policy and tests for the intended contract. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
