I3eka commented on code in PR #43133: URL: https://github.com/apache/superset/pull/43133#discussion_r4385441246
########## 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: Fixed in 428484e846. Conversation selection is now a styled native button with selected-state semantics, separate from rename/delete controls. Keyboard regressions activate the conversation with Enter and Space. ########## 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: Fixed in 428484e846. The transcript has log semantics with token-by-token live announcements disabled; a separate polite/atomic output announces a completed assistant response. The regression verifies no announcement during streaming and an announcement at completion. This checks DOM/accessibility semantics locally; I have not tested a screen reader manually. ########## 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: Confirmed the missing resource gate and fixed it in ea8eedb9fb. Built-in readers check FAB `can_read` for Chart/Dashboard/Database before DAO lookup, while asset search omits resource types the role cannot read. Existing object/datasource checks remain. Denied-capability regressions prove the forbidden DAO/search path is not called. Those targeted unit cases stub the capability decision; they do not claim an additional live custom-role/browser validation. ########## 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: Confirmed and fixed in ea8eedb9fb. The default foreign-tool policy requires an explicit operator `tool_allowlist` for the configured server, with denylist precedence; names and remote schemas/annotations cannot authorize execution. Tests cover renamed `fetch_rows`, nested/missing SQL fields, unlisted metadata/search tools, unknown servers and exact allowlists. The documented administrator opt-out remains explicit; no new name/schema heuristic was added. ########## superset/ai/api.py: ########## @@ -0,0 +1,989 @@ +# 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. +""" +REST API for the AI assistant. + +Every route carries ``@protect()`` and is reached through ``@expose`` on a +``BaseSupersetApi`` subclass, which is what makes Flask-AppBuilder's +authorization actually run. Ownership is enforced a second time in the command +and DAO layers, so a conversation identifier is never on its own a capability. +""" + +from __future__ import annotations + +import logging +import time +from collections.abc import Generator +from typing import Any, cast + +from flask import current_app, request, Response, stream_with_context +from flask_appbuilder.api import expose, permission_name, protect, safe +from marshmallow import ValidationError + +from superset.ai.events import ( + error_event, + KEEPALIVE_FRAME, + KEEPALIVE_INTERVAL_SECONDS, +) +from superset.ai.schemas import ( + AgentResponseSchema, + CancelPostSchema, + FeedbackPostSchema, + MessagePostSchema, + RunAcceptedResponseSchema, + SuggestedPromptsPostSchema, + ThreadDetailResponseSchema, + ThreadPostSchema, + ThreadPutSchema, + ThreadResponseSchema, +) +from superset.ai.types import MessageRole, MessageStatus +from superset.commands.ai.exceptions import ( + AIChatMessageInvalidError, + AIChatMessageNotFoundError, + AIChatThreadInvalidError, + AIChatThreadNotFoundError, +) +from superset.extensions import event_logger +from superset.utils.core import get_user_id +from superset.utils.decorators import transaction +from superset.views.base_api import BaseSupersetApi, statsd_metrics + +logger = logging.getLogger(__name__) + +#: Upper bound on how long a client may hold a stream open, so an abandoned +#: browser tab cannot pin a worker indefinitely. +_STREAM_TIMEOUT_SECONDS = 900 + +#: How often a reader checks the event bus for new frames. +#: +#: Deliberately separate from ``KEEPALIVE_INTERVAL_SECONDS``. Passing the +#: keep-alive interval as the poll interval made the reader sleep fifteen seconds +#: between checks and then deliver everything that had accumulated in one batch — +#: so a worker-mode run showed no streaming at all: the answer and every tool call +#: appeared in fifteen-second lumps. One controls responsiveness, the other how +#: often an idle connection is reassured; they are not the same number. +_EVENT_POLL_SECONDS = 0.1 + + +class AIRestApi(BaseSupersetApi): + """Conversations with the AI assistant.""" + + resource_name = "ai" + openapi_spec_tag = "AI Assistant" + allow_browser_login = True + class_permission_name = "AIAssistant" + + openapi_spec_component_schemas = ( + AgentResponseSchema, + CancelPostSchema, + FeedbackPostSchema, + MessagePostSchema, + RunAcceptedResponseSchema, + SuggestedPromptsPostSchema, + ThreadDetailResponseSchema, + ThreadPostSchema, + ThreadPutSchema, + ThreadResponseSchema, + ) + + @expose("/agent/", methods=("GET",)) + @protect() + @safe + @statsd_metrics + @permission_name("read") + def agents(self) -> Response: + """List agent profiles the current user may select. + --- + get: + summary: List available agent profiles + responses: + 200: + description: Available profiles + content: + application/json: + schema: + type: object + properties: + result: + type: array + items: + $ref: '#/components/schemas/AgentResponseSchema' + 401: + $ref: '#/components/responses/401' + 403: + $ref: '#/components/responses/403' + 404: + $ref: '#/components/responses/404' + """ + if (unavailable := self._reject_if_unconfigured()) is not None: + return unavailable + + from superset.ai.factories import get_profiles + + profiles = get_profiles().visible_to_current_user() + return self.response(200, result=[p.to_public_dict() for p in profiles]) + + @expose("/model/", methods=("GET",)) + @protect() + @safe + @statsd_metrics + @permission_name("read") + def models(self) -> Response: + """List models this deployment has configured. + --- + get: + summary: List selectable models + responses: + 200: + description: Configured model identifiers + content: + application/json: + schema: + type: object + properties: + result: + type: array + items: + type: string + 401: + $ref: '#/components/responses/401' + 404: + $ref: '#/components/responses/404' + """ + if (unavailable := self._reject_if_unconfigured()) is not None: + return unavailable + + from superset.ai.factories import get_provider + + return self.response(200, result=get_provider().available_models()) + + @expose("/thread/", methods=("POST",)) + @protect() + @safe + @statsd_metrics + @permission_name("write") + @event_logger.log_this_with_context( + action=lambda self, *args, **kwargs: f"{self.__class__.__name__}.post_thread", + log_to_statsd=False, + ) + def post_thread(self) -> Response: + """Create a conversation. + --- + post: + summary: Create a conversation + requestBody: + content: + application/json: + schema: + $ref: '#/components/schemas/ThreadPostSchema' + responses: + 201: + description: Conversation created + content: + application/json: + schema: + type: object + properties: + result: + $ref: '#/components/schemas/ThreadResponseSchema' + 400: + $ref: '#/components/responses/400' + 401: + $ref: '#/components/responses/401' + 404: + $ref: '#/components/responses/404' + """ + if (unavailable := self._reject_if_unconfigured()) is not None: + return unavailable + + from superset.commands.ai import CreateAIChatThreadCommand + + try: + payload = ThreadPostSchema().load(request.json or {}) + except ValidationError as error: + return self.response_400(message=error.messages) + try: + thread = CreateAIChatThreadCommand( + user_id=self._user_id(), + title=payload.get("title"), + agent_key=payload.get("agent_key"), + ).run() + except AIChatThreadInvalidError as ex: + return self.response_422(message=str(ex)) + return self.response(201, result=_thread_dict(thread)) + + @expose("/thread/", methods=("GET",)) + @protect() + @safe + @statsd_metrics + @permission_name("read") + def get_threads(self) -> Response: + """List the current user's conversations. + --- + get: + summary: List conversations + parameters: + - in: query + name: limit + schema: + type: integer + - in: query + name: offset + schema: + type: integer + responses: + 200: + description: Conversations + content: + application/json: + schema: + type: object + properties: + count: + type: integer + result: + type: array + items: + $ref: '#/components/schemas/ThreadResponseSchema' + 401: + $ref: '#/components/responses/401' + 404: + $ref: '#/components/responses/404' + """ + if (unavailable := self._reject_if_unconfigured()) is not None: + return unavailable + + from superset.daos.ai import AIChatThreadDAO + + limit = request.args.get("limit", type=int) or 50 + offset = request.args.get("offset", type=int) or 0 + threads = AIChatThreadDAO.find_all_for_user( + self._user_id(), limit=limit, offset=offset + ) + return self.response( + 200, + count=len(threads), + result=[_thread_dict(thread) for thread in threads], + ) + + @expose("/thread/<thread_uuid>", methods=("GET",)) + @protect() + @safe + @statsd_metrics + @permission_name("read") + def get_thread(self, thread_uuid: str) -> Response: + """Fetch a conversation and its messages. + --- + get: + summary: Get a conversation + parameters: + - in: path + name: thread_uuid + required: true + schema: + type: string + format: uuid + responses: + 200: + description: Conversation with messages + content: + application/json: + schema: + type: object + properties: + result: + $ref: '#/components/schemas/ThreadDetailResponseSchema' + 401: + $ref: '#/components/responses/401' + 404: + $ref: '#/components/responses/404' + """ + if (unavailable := self._reject_if_unconfigured()) is not None: + return unavailable + + from superset.daos.ai import ( + AIChatFeedbackDAO, + AIChatMessageDAO, + AIChatThreadDAO, + ) + + user_id = self._user_id() + thread = AIChatThreadDAO.find_by_uuid_for_user(thread_uuid, user_id) + if thread is None: + return self.response_404() + + messages = AIChatMessageDAO.find_for_thread(thread) + # Resolved for the whole transcript at once so the panel can show which + # replies this user already rated; without it a reload loses the verdict + # and the message looks unrated. + verdicts = AIChatFeedbackDAO.find_verdicts_for_user( + [message.id for message in messages], user_id + ) + detail = _thread_dict(thread) + detail["messages"] = [ + _message_dict(message, liked=verdicts.get(message.id)) + for message in messages + ] + return self.response(200, result=detail) + + @expose("/thread/<thread_uuid>", methods=("PUT",)) + @protect() + @safe + @statsd_metrics + @permission_name("write") + @event_logger.log_this_with_context( + action=lambda self, *args, **kwargs: f"{self.__class__.__name__}.put_thread", + log_to_statsd=False, + ) + def put_thread(self, thread_uuid: str) -> Response: + """Rename or archive a conversation. + --- + put: + summary: Update a conversation + parameters: + - in: path + name: thread_uuid + required: true + schema: + type: string + format: uuid + requestBody: + content: + application/json: + schema: + $ref: '#/components/schemas/ThreadPutSchema' + responses: + 200: + description: Conversation updated + 401: + $ref: '#/components/responses/401' + 404: + $ref: '#/components/responses/404' + 422: + $ref: '#/components/responses/422' + """ + if (unavailable := self._reject_if_unconfigured()) is not None: + return unavailable + + from superset.commands.ai import UpdateAIChatThreadCommand + + try: + payload = ThreadPutSchema().load(request.json or {}) + except ValidationError as error: + return self.response_400(message=error.messages) + try: + thread = UpdateAIChatThreadCommand( + thread_uuid, + self._user_id(), + title=payload.get("title"), + status=payload.get("status"), + ).run() + except AIChatThreadNotFoundError: + return self.response_404() + except AIChatThreadInvalidError as ex: + return self.response_422(message=str(ex)) + return self.response(200, result=_thread_dict(thread)) + + @expose("/thread/<thread_uuid>", methods=("DELETE",)) + @protect() + @safe + @statsd_metrics + @permission_name("write") + @event_logger.log_this_with_context( + action=lambda self, *args, **kwargs: f"{self.__class__.__name__}.delete_thread", + log_to_statsd=False, + ) + def delete_thread(self, thread_uuid: str) -> Response: + """Delete a conversation and its messages. + --- + delete: + summary: Delete a conversation + parameters: + - in: path + name: thread_uuid + required: true + schema: + type: string + format: uuid + responses: + 200: + description: Conversation deleted + 401: + $ref: '#/components/responses/401' + 404: + $ref: '#/components/responses/404' + """ + if (unavailable := self._reject_if_unconfigured()) is not None: + return unavailable + + from superset.commands.ai import DeleteAIChatThreadCommand + + try: + DeleteAIChatThreadCommand(thread_uuid, self._user_id()).run() + except AIChatThreadNotFoundError: + return self.response_404() + return self.response(200, message="OK") + + @expose("/thread/<thread_uuid>/message", methods=("POST",)) + @protect() + @safe + @statsd_metrics + @permission_name("write") + @event_logger.log_this_with_context( + action=lambda self, *args, **kwargs: f"{self.__class__.__name__}.post_message", + log_to_statsd=False, + ) + def post_message(self, thread_uuid: str) -> Response: + """Post a user message and start a run. + --- + post: + summary: Post a message + description: > + Stores the user's message, creates a placeholder assistant message, + and starts a run. Returns immediately; consume the answer from the + stream endpoint using the returned run identifier. + parameters: + - in: path + name: thread_uuid + required: true + schema: + type: string + format: uuid + requestBody: + content: + application/json: + schema: + $ref: '#/components/schemas/MessagePostSchema' + responses: + 202: + description: Run accepted + content: + application/json: + schema: + type: object + properties: + result: + $ref: '#/components/schemas/RunAcceptedResponseSchema' + 400: + $ref: '#/components/responses/400' + 401: + $ref: '#/components/responses/401' + 404: + $ref: '#/components/responses/404' + 422: + $ref: '#/components/responses/422' + """ + if (unavailable := self._reject_if_unconfigured()) is not None: + return unavailable + + from superset.ai.orchestrator import new_run_id + from superset.commands.ai import AppendAIChatMessageCommand + + try: + payload = MessagePostSchema().load(request.json or {}) + except ValidationError as error: + return self.response_400(message=error.messages) + user_id = self._user_id() + + try: + user_message = AppendAIChatMessageCommand( + thread_uuid, + user_id, + MessageRole.USER, + payload["content"], + request_id=payload.get("request_id"), + ).run() + # Created up front so a client that reconnects before any token + # arrives still has a row to attach its stream to. + assistant_message = AppendAIChatMessageCommand( + thread_uuid, + user_id, + MessageRole.ASSISTANT, + "", + request_id=payload.get("request_id"), + status=MessageStatus.PENDING, + ).run() + except AIChatThreadNotFoundError: + return self.response_404() + except (AIChatMessageInvalidError, AIChatThreadInvalidError) as ex: + return self.response_422(message=str(ex)) + + run_id = new_run_id() + _record_run_context(assistant_message, run_id, payload) + + self._start_run( + thread_uuid=thread_uuid, + user_id=user_id, + run_id=run_id, + assistant_message_uuid=str(assistant_message.uuid), + agent_key=payload.get("agent_key"), + model=payload.get("model"), + page_context=payload.get("page_context"), + ) + + return self.response( + 202, + result={ + "message_uuid": str(user_message.uuid), + "assistant_message_uuid": str(assistant_message.uuid), + "run_id": run_id, + }, + ) + + @expose("/thread/<thread_uuid>/stream", methods=("GET",)) + @protect() + @statsd_metrics + @permission_name("read") + def stream(self, thread_uuid: str) -> Response: + """Stream a run's events. + --- + get: + summary: Stream assistant events + description: > + Server-sent events for one run. Frame names are session, thinking, + thoughts, checkpoint, assistant_delta, final, error, cancelled and + done. The done frame is always last and reports whether the run + succeeded. + parameters: + - in: path + name: thread_uuid + required: true + schema: + type: string + format: uuid + - in: query + name: run_id + required: true + schema: + type: string + responses: + 200: + description: An event stream + content: + text/event-stream: + schema: + type: string + 401: + $ref: '#/components/responses/401' + 404: + $ref: '#/components/responses/404' + """ + # No @safe here: once headers are flushed an exception can no longer + # become a status code, so failures are reported as in-band error frames. + if (unavailable := self._reject_if_unconfigured()) is not None: + return unavailable + + from superset.daos.ai import AIChatMessageDAO, AIChatThreadDAO + + run_id = request.args.get("run_id") + if not run_id: + return self.response_400(message="run_id is required") + + # Ownership is checked before the stream opens; the run identifier alone + # must not grant access to another user's conversation. + thread = AIChatThreadDAO.find_by_uuid_for_user(thread_uuid, self._user_id()) + if thread is None: + return self.response_404() + + pending = _find_run_message(AIChatMessageDAO.find_for_thread(thread), run_id) + if pending is None: + return self.response_404() + + turn = None + if current_app.config.get("AI_ASSISTANT_EXECUTION_MODE") != "worker": + from superset.ai.orchestrator import TurnRequest + + extra = pending.extra + turn = TurnRequest( + thread_uuid=thread_uuid, + user_id=self._user_id(), + run_id=run_id, + assistant_message_uuid=str(pending.uuid), + profile_key=extra.get("agent_key"), + model=extra.get("model"), + page_context=extra.get("page_context"), + ) + + generator = self._build_stream(run_id, turn) + response = Response( + generator, + content_type="text/event-stream; charset=utf-8", + headers={ + "Cache-Control": "no-cache, no-transform", + "Connection": "keep-alive", + # Defeats proxy buffering, which otherwise holds frames until + # the response completes and makes streaming pointless. + "X-Accel-Buffering": "no", + "Content-Encoding": "identity", + }, + direct_passthrough=False, + ) + response.implicit_sequence_conversion = False + return response + + @expose("/thread/<thread_uuid>/cancel", methods=("POST",)) + @protect() + @safe + @statsd_metrics + @permission_name("write") + def cancel(self, thread_uuid: str) -> Response: + """Ask a run to stop. + --- + post: + summary: Cancel a run + description: > + Cancellation is cooperative: the run stops at its next step + boundary. A run inside a single long model call or query will not + stop until that call returns. + parameters: + - in: path + name: thread_uuid + required: true + schema: + type: string + format: uuid + requestBody: + content: + application/json: + schema: + $ref: '#/components/schemas/CancelPostSchema' + responses: + 200: + description: Cancellation recorded + 401: + $ref: '#/components/responses/401' + 404: + $ref: '#/components/responses/404' + """ + if (unavailable := self._reject_if_unconfigured()) is not None: + return unavailable + + from superset.ai.orchestrator import request_cancel + from superset.daos.ai import AIChatThreadDAO + + try: + payload = CancelPostSchema().load(request.json or {}) + except ValidationError as error: + return self.response_400(message=error.messages) + if AIChatThreadDAO.find_by_uuid_for_user(thread_uuid, self._user_id()) is None: + return self.response_404() + + request_cancel(payload["run_id"]) + return self.response(200, message="OK") + + @expose("/feedback", methods=("POST",)) + @protect() + @safe + @statsd_metrics + @permission_name("write") + def feedback(self) -> Response: + """Rate an assistant message. + --- + post: + summary: Submit feedback + requestBody: + content: + application/json: + schema: + $ref: '#/components/schemas/FeedbackPostSchema' + responses: + 200: + description: Feedback recorded + 401: + $ref: '#/components/responses/401' + 404: + $ref: '#/components/responses/404' + 422: + $ref: '#/components/responses/422' + """ + if (unavailable := self._reject_if_unconfigured()) is not None: + return unavailable + + from superset.commands.ai import SubmitAIChatFeedbackCommand + + try: + payload = FeedbackPostSchema().load(request.json or {}) + except ValidationError as error: + return self.response_400(message=error.messages) + try: + SubmitAIChatFeedbackCommand( + payload["message_uuid"], + self._user_id(), + liked=payload["liked"], + comment=payload.get("comment"), + ).run() + except AIChatMessageNotFoundError: + return self.response_404() + except AIChatMessageInvalidError as ex: Review Comment: Fixed on this branch in 35c9b04143. The endpoint catches `AIChatFeedbackInvalidError`, the exception actually raised by the command, and returns the documented 422. The API regression submits an owned user-message UUID and verifies the validation response rather than a 500. -- 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]
