I3eka commented on code in PR #43133: URL: https://github.com/apache/superset/pull/43133#discussion_r4385438536
########## superset-frontend/src/features/ai/components/ChatChartEmbed.tsx: ########## @@ -0,0 +1,542 @@ +/** + * 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 A chart rendered inside a chat message. + * + * The assistant does not send a chart, it sends a `form_data_key` it stored — so + * what arrives in the transcript is a reference the client resolves, and the + * rendered chart is the real thing, with the real permissions, rather than an + * image of one. + * + * The awkward part is timing: the key can exist before the query behind it has + * finished. Rather than show the chart's own "No data" state (which reads as a + * broken answer) the component keeps a spinner up and re-renders on a growing + * backoff until rows appear, giving up after a bounded number of attempts. + */ + +import { useCallback, useEffect, useRef, useState } from 'react'; +import { + type QueryFormData, + StatefulChart, + SupersetClient, +} from '@superset-ui/core'; +import { styled } from '@apache-superset/core/theme'; +import { t } from '@apache-superset/core/translation'; +import { Loading } from '@superset-ui/core/components'; +import { Icons } from '@superset-ui/core/components/Icons'; +import { ErrorBoundary } from 'src/components/ErrorBoundary'; + +const VALID_KEY_PATTERN = /^[a-zA-Z0-9_-]+$/; +const MIN_HEIGHT = 100; +const MAX_HEIGHT = 800; +const DEFAULT_HEIGHT = 300; +const FETCH_TIMEOUT_MS = 30_000; +const MAX_RETRIES = 3; +const RETRY_DELAYS_MS = [500, 1500, 3000]; + +/** Width used until the container has been measured. */ +const FALLBACK_CHART_WIDTH = 600; + +// Backoff for the "waiting for chart data" poll. The delay grows exponentially +// per attempt up to a ceiling, so a slow query is waited out without hammering +// the backend. +const POLL_BASE_DELAY_MS = 1000; +const POLL_MAX_DELAY_MS = 30_000; + +/** + * Polling stops after this many attempts. + * + * A retry re-issues the chart's data request, which will use the results cache if + * the query has landed but will otherwise execute it. Polling forever would keep + * re-issuing it, so an unfinished query surfaces the retry control instead. + */ +const MAX_POLL_ATTEMPTS = 6; + +const getPollDelayMs = (attempt: number): number => + Math.min(POLL_BASE_DELAY_MS * 2 ** attempt, POLL_MAX_DELAY_MS); + +export interface ChartEmbedParams { + formDataKey: string | null; + height: number; + title: string | null; +} + +/** + * Parse key=value lines from the content of a ```superset-chart fenced block. + * + * Rules are strict on purpose: the block is model output, so `form_data_key` must + * match `/^[a-zA-Z0-9_-]+$/` before it reaches a URL, and `height` is clamped. + * Unknown keys are ignored so a newer backend can add some without breaking an + * older client. + */ +export function parseChartEmbedParams(codeText: string): ChartEmbedParams { + const result: ChartEmbedParams = { + formDataKey: null, + height: DEFAULT_HEIGHT, + title: null, + }; + + const lines = codeText + .split('\n') + .map(line => line.trim()) + .filter(Boolean); + + lines.forEach(line => { + const eqIndex = line.indexOf('='); + if (eqIndex <= 0) { + return; + } + + const key = line.slice(0, eqIndex).trim().toLowerCase(); + const value = line.slice(eqIndex + 1).trim(); + + if (key === 'form_data_key') { + if (value && VALID_KEY_PATTERN.test(value)) { + result.formDataKey = value; + } + return; + } + if (key === 'height') { + const parsed = parseInt(value, 10); + if (!Number.isNaN(parsed)) { + result.height = Math.max(MIN_HEIGHT, Math.min(MAX_HEIGHT, parsed)); + } + return; + } + if (key === 'title' && value) { + result.title = value; + } + }); + + return result; +} + +const ChartContainer = styled.div` + border: 1px solid ${({ theme }) => theme.colorBorderSecondary}; + border-radius: ${({ theme }) => theme.borderRadius}px; + overflow: hidden; + margin: ${({ theme }) => theme.sizeUnit * 2}px 0; + background: ${({ theme }) => theme.colorBgContainer}; +`; + +const ChartHeader = styled.div` + display: flex; + justify-content: space-between; + align-items: center; + padding: ${({ theme }) => theme.sizeUnit * 2}px + ${({ theme }) => theme.sizeUnit * 3}px; + border-bottom: 1px solid ${({ theme }) => theme.colorBorderSecondary}; + background: ${({ theme }) => theme.colorBgLayout}; + font-size: ${({ theme }) => theme.fontSizeSM}px; +`; + +const ChartTitle = styled.span` + font-weight: ${({ theme }) => theme.fontWeightStrong}; + color: ${({ theme }) => theme.colorText}; + overflow: hidden; + text-overflow: ellipsis; + white-space: nowrap; + flex: 1; + min-width: 0; +`; + +const ChartActions = styled.div` + display: flex; + gap: ${({ theme }) => theme.sizeUnit * 2}px; + align-items: center; + flex-shrink: 0; + margin-left: ${({ theme }) => theme.sizeUnit * 2}px; +`; + +const ActionLink = styled.a` + display: inline-flex; + align-items: center; + gap: ${({ theme }) => theme.sizeUnit / 2}px; + font-size: ${({ theme }) => theme.fontSizeSM}px; + color: ${({ theme }) => theme.colorPrimary}; + cursor: pointer; + text-decoration: none; + + &:hover { + text-decoration: underline; + } +`; + +const ActionButton = styled.button` + display: inline-flex; + align-items: center; + gap: ${({ theme }) => theme.sizeUnit / 2}px; + font-size: ${({ theme }) => theme.fontSizeSM}px; + color: ${({ theme }) => theme.colorTextSecondary}; + cursor: pointer; + background: none; + border: none; + padding: 2px ${({ theme }) => theme.sizeUnit / 2}px; + border-radius: ${({ theme }) => theme.borderRadius}px; + + &:hover { + color: ${({ theme }) => theme.colorPrimary}; + background: ${({ theme }) => theme.colorFillTertiary}; + } +`; + +const ChartBody = styled.div<{ height: number }>` + height: ${({ height }) => height}px; + position: relative; +`; + +// Covers the chart while it reports no data, hiding the underlying "No data" +// state (which looks broken) behind a spinner while refreshing continues. +const ChartDataOverlay = styled.div` + position: absolute; + inset: 0; + display: flex; + flex-direction: column; + align-items: center; + justify-content: center; + gap: ${({ theme }) => theme.sizeUnit * 3}px; + background: ${({ theme }) => theme.colorBgContainer}; + color: ${({ theme }) => theme.colorTextSecondary}; + font-size: ${({ theme }) => theme.fontSizeSM}px; + z-index: 2; +`; + +const CenteredMessage = styled.div<{ height: number }>` + display: flex; + flex-direction: column; + align-items: center; + justify-content: center; + height: ${({ height }) => height}px; + color: ${({ theme }) => theme.colorTextSecondary}; + font-size: ${({ theme }) => theme.fontSizeSM}px; + text-align: center; + padding: ${({ theme }) => theme.sizeUnit * 4}px; + gap: ${({ theme }) => theme.sizeUnit * 2}px; +`; + +interface ChatChartEmbedProps { + formDataKey: string; + height?: number; + title?: string; +} + +type FetchState = + | { status: 'loading' } + | { status: 'loaded'; formData: QueryFormData } + | { status: 'error'; message: string }; + +const exploreUrlFor = (formDataKey: string): string => + `/explore/?form_data_key=${encodeURIComponent(formDataKey)}`; + +export function ChatChartEmbedInner({ + formDataKey, + height = DEFAULT_HEIGHT, + title, +}: ChatChartEmbedProps) { + const [fetchState, setFetchState] = useState<FetchState>({ + status: 'loading', + }); + const [chartWidth, setChartWidth] = useState(0); + const [chartRenderKey, setChartRenderKey] = useState(0); + // True while the chart's query results are still unavailable (either an empty + // result or a cache miss), which is what keeps the overlay up. + const [isAwaitingData, setIsAwaitingData] = useState(false); + const chartBodyRef = useRef<HTMLDivElement>(null); + const pollTimeoutRef = useRef<ReturnType<typeof setTimeout>>(); + + // Bumping `chartRenderKey` remounts the chart, which re-issues its data + // request. `force` is left off so the results cache is preferred. + const scheduleNextPoll = useCallback(() => { + if (chartRenderKey >= MAX_POLL_ATTEMPTS) { + setIsAwaitingData(false); + return; + } + setIsAwaitingData(true); + const delay = getPollDelayMs(chartRenderKey); + if (pollTimeoutRef.current) { + clearTimeout(pollTimeoutRef.current); + } + pollTimeoutRef.current = setTimeout( + () => setChartRenderKey(key => key + 1), + delay, + ); + }, [chartRenderKey]); + + useEffect( + () => () => { + if (pollTimeoutRef.current) { + clearTimeout(pollTimeoutRef.current); + } + }, + [], + ); + + useEffect(() => { + const element = chartBodyRef.current; + if (!element) { + return undefined; + } + + const initialWidth = Math.floor(element.getBoundingClientRect().width); + if (initialWidth > 0) { + setChartWidth(initialWidth); + } + + // The panel is resizable by the host, so the chart is measured rather than + // given a fixed width. + const observer = new ResizeObserver(entries => { + const [entry] = entries; + if (entry) { + const width = Math.floor(entry.contentRect.width); + if (width > 0) { + setChartWidth(width); + } + } + }); + observer.observe(element); + return () => observer.disconnect(); + }, [fetchState.status]); + + const exploreUrl = exploreUrlFor(formDataKey); + + const fetchFormData = useCallback( + async (attempt: number = 0) => { + setFetchState({ status: 'loading' }); + setIsAwaitingData(false); + setChartRenderKey(0); + + const retry = (): boolean => { + if (attempt >= MAX_RETRIES) { + return false; + } + const delay = RETRY_DELAYS_MS[attempt] ?? 1000; + setTimeout(() => { Review Comment: Fixed in fdfa31681d and completed in 428484e846. The effect owns the retry timer, request timeout and AbortController; cleanup cancels all three and late responses are ignored after unmount/key change. Fake-timer tests cover unmount before retry, key replacement, ignored late responses and no leftover timer. They pass. ########## superset/models/ai.py: ########## @@ -0,0 +1,270 @@ +# 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. +""" +Persistence for AI assistant conversations. + +Conversations live in Superset's own metadata database, which keeps the +feature deployable with no extra infrastructure and makes ownership +enforceable with the same DAO filters used everywhere else. + +These models live under ``superset/models/`` rather than inside +``superset/ai/`` because background workers need them without importing the +API module. +""" + +from __future__ import annotations + +import uuid as uuid_module +from typing import Any + +import sqlalchemy as sa +from flask_appbuilder import Model +from sqlalchemy.orm import relationship, validates +from sqlalchemy_utils import UUIDType + +from superset.ai.types import ( + MessageExtra, + MessageRole, + MessageStatus, + ThreadStatus, +) +from superset.models.helpers import AuditMixinNullable +from superset.utils import json +from superset.utils.core import MediumText + +#: Bumped when the meaning of keys inside an ``extra_json`` blob changes, so a +#: reader can tell an old row from a new one instead of guessing. +EXTRA_JSON_VERSION = 1 + + +class AIChatThread(AuditMixinNullable, Model): + """ + One conversation between a user and the assistant. + + Ownership is expressed through ``created_by_fk`` (supplied by + :class:`AuditMixinNullable`) and enforced by the DAO's base filter, so a + thread identifier is not by itself a capability. + """ + + __tablename__ = "ai_chat_threads" + __table_args__ = ( + # Serves the "my threads, most recent first" list query. + sa.Index("ix_ai_chat_threads_owner_recent", "created_by_fk", "changed_on"), + ) + + id = sa.Column(sa.Integer, primary_key=True) + #: The only identifier exposed over HTTP. Integer ids stay internal. + uuid = sa.Column( + UUIDType(binary=True), + nullable=False, + unique=True, + default=uuid_module.uuid4, + ) + + title = sa.Column(sa.String(512), nullable=True) + status = sa.Column( + sa.String(32), + nullable=False, + default=ThreadStatus.ACTIVE.value, + server_default=ThreadStatus.ACTIVE.value, + ) + #: Which agent profile this thread was last run with. + agent_key = sa.Column(sa.String(64), nullable=True) + extra_json = sa.Column(MediumText(), nullable=True) + + # No ``passive_deletes``: SQLite does not enforce foreign keys unless + # ``PRAGMA foreign_keys=ON``, so deferring the cascade to the database + # would orphan messages there. The ORM deletes children itself, and the + # ``ON DELETE CASCADE`` in the DDL remains a backstop for direct SQL. + messages = relationship( + "AIChatMessage", + back_populates="thread", + cascade="all, delete-orphan", + order_by="AIChatMessage.created_on", + ) + + def __repr__(self) -> str: + return f"<AIChatThread {self.uuid} [{self.status}]>" + + @validates("status") + def _validate_status(self, _key: str, value: Any) -> str: + """Reject unknown lifecycle values at assignment time.""" + return ThreadStatus(value).value + + @property + def message_count(self) -> int: + """ + Number of stored messages. + + Derived rather than denormalised: a counter column would have to be kept + in step with cascade deletes and retention pruning, and a counter that + drifts is worse than one query — it reports a conversation length nobody + can reconcile against the rows. + """ + return len(self.messages) + + @property + def extra(self) -> dict[str, Any]: + """Parsed ``extra_json``, or an empty dict when absent or corrupt.""" + return _load_json_object(self.extra_json) + + +class AIChatMessage(AuditMixinNullable, Model): + """ + A single turn in a conversation. + + Assistant messages are inserted before inference begins so a client that + reconnects mid-run has a row to attach to, then transition through + ``streaming`` to a terminal status. + """ + + __tablename__ = "ai_chat_messages" + __table_args__ = ( + sa.Index("ix_ai_chat_messages_thread_created", "thread_id", "created_on"), + # Makes client-supplied idempotency real: replaying a request cannot + # create a second row for the same turn. The constraint carries the role + # as well, because one request legitimately produces both a user message + # and the assistant message answering it. + sa.UniqueConstraint( + "thread_id", + "request_id", + "role", + name="uq_ai_chat_messages_thread_request_role", + ), + ) + + id = sa.Column(sa.Integer, primary_key=True) + uuid = sa.Column( + UUIDType(binary=True), + nullable=False, + unique=True, + default=uuid_module.uuid4, + ) + + thread_id = sa.Column( + sa.Integer, + sa.ForeignKey("ai_chat_threads.id", ondelete="CASCADE"), + nullable=False, + ) + role = sa.Column(sa.String(16), nullable=False) + #: Unbounded model or user text; deliberately not a plain ``Text`` column, + #: which caps at 64 KB on MySQL. + content = sa.Column(MediumText(), nullable=False, default="") + status = sa.Column( + sa.String(32), + nullable=False, + default=MessageStatus.COMPLETE.value, + server_default=MessageStatus.COMPLETE.value, + ) + #: Client-generated idempotency key for the turn. + request_id = sa.Column(sa.String(96), nullable=True) + #: Serialised :class:`~superset.ai.types.MessageExtra`. + extra_json = sa.Column(MediumText(), nullable=True) + + thread = relationship("AIChatThread", back_populates="messages") + + def __repr__(self) -> str: + return f"<AIChatMessage {self.uuid} {self.role} [{self.status}]>" + + @validates("role") + def _validate_role(self, _key: str, value: Any) -> str: + """Reject unknown authors at assignment time.""" + return MessageRole(value).value + + @validates("status") + def _validate_status(self, _key: str, value: Any) -> str: + """Reject unknown lifecycle values at assignment time.""" + return MessageStatus(value).value + + @property + def is_terminal(self) -> bool: + """Whether this message will never change again.""" + return MessageStatus(self.status) in MessageStatus.terminal() + + @property + def extra(self) -> MessageExtra: + """Parsed ``extra_json``, or an empty dict when absent or corrupt.""" + return _load_json_object(self.extra_json) # type: ignore[return-value] + + def update_extra(self, updates: MessageExtra) -> None: + """ + Merge keys into ``extra_json``. + + Merge rather than replace, because a run writes tool calls and token + usage at different moments. + """ + merged: dict[str, Any] = dict(self.extra) + merged.update(updates) + merged["version"] = EXTRA_JSON_VERSION + self.extra_json = json.dumps(merged) + + +class AIChatFeedback(AuditMixinNullable, Model): + """ + A thumbs up or down on an assistant message. + + A first-class table rather than a log line, so the signal can actually be + aggregated and joined back to the conversation that produced it. + """ + + __tablename__ = "ai_chat_feedback" + __table_args__ = ( + # One verdict per user per message; a repeat vote updates in place. + sa.UniqueConstraint( + "message_id", + "created_by_fk", + name="uq_ai_chat_feedback_message_user", + ), + ) + + id = sa.Column(sa.Integer, primary_key=True) + uuid = sa.Column( + UUIDType(binary=True), + nullable=False, + unique=True, + default=uuid_module.uuid4, + ) + + message_id = sa.Column( + sa.Integer, + sa.ForeignKey("ai_chat_messages.id", ondelete="CASCADE"), + nullable=False, + ) + liked = sa.Column(sa.Boolean, nullable=False) + comment = sa.Column(MediumText(), nullable=True) + + message = relationship("AIChatMessage") Review Comment: Fixed in 366d65fab4. Messages now own feedback through an ORM delete-orphan cascade, matching thread ownership of messages. The regression disables SQLite foreign-key enforcement, deletes a thread containing feedback and verifies that all three tables are empty. ########## superset/mcp_service/dashboard/tool/generate_dashboard.py: ########## @@ -145,6 +145,66 @@ def _create_dashboard_layout(chart_objects: List[Any]) -> Dict[str, Any]: return layout +def _is_valid_dashboard_layout(layout: Dict[str, Any], chart_ids: set[int]) -> bool: + """Return whether an explicit layout is safe for frontend hydration.""" + components = { + component_id: component + for component_id, component in layout.items() + if isinstance(component, dict) and "type" in component + } + children_by_id: dict[str, list[str]] = {} + for component_id, component in components.items(): + children = component.get("children") + if component.get("type") == "HEADER" and children is None: + children = [] + if ( + component.get("id") != component_id + or not isinstance(component.get("type"), str) + or not isinstance(children, list) + or any( + not isinstance(child_id, str) or child_id not in components + for child_id in children + ) + ): + return False + if component["type"] == "CHART": Review Comment: Rechecked the current layout path after the upstream merge. It uses the shared layout validator, compares the reachable chart set with the requested set and rebuilds derived parent chains from the child graph. Invalid/omitted-chart layouts fall back with a warning; stale parent paths are normalized. The layout regressions cover these cases rather than assuming membership of each visited chart proves completeness. ########## tests/integration_tests/ai/tools_e2e_tests.py: ########## @@ -0,0 +1,319 @@ +# 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. +""" +End-to-end verification that every shipped tool actually runs. + +Unit tests for the tools inject fakes at the database seam, which proves the +guards but not that the tool can execute at all. These drive each tool through +the real registry, against the real metadata database, with real permission +checks — the difference between "the code is correct" and "the feature works". + +Every tool named in :data:`ALL_TOOL_NAMES` is covered, and a guard test asserts +that, so a tool added later cannot quietly go unverified. +""" + +from typing import Any + +import pytest + +from superset import db +from superset.ai.llm.base import ToolCall +from superset.ai.tools import build_registry, BUNDLE_ALL +from superset.ai.tools.base import ALL_TOOL_NAMES +from superset.models.ai import AIChatFeedback, AIChatMessage, AIChatThread +from superset.utils.database import get_main_database +from tests.integration_tests.base_tests import SupersetTestCase + + +def _call(name: str, **arguments: Any) -> ToolCall: + return ToolCall(id=f"call-{name}", name=name, arguments=arguments) + + +class TestAIToolsEndToEnd(SupersetTestCase): + """Each tool, executed for real.""" + + def setUp(self) -> None: + super().setUp() + self.login("admin") + self.registry = build_registry(BUNDLE_ALL) + + def tearDown(self) -> None: + super().tearDown() + db.session.query(AIChatFeedback).delete() + db.session.query(AIChatMessage).delete() + db.session.query(AIChatThread).delete() + db.session.commit() + + def _invoke(self, name: str, **arguments: Any) -> Any: + invocation = self.registry.invoke(_call(name, **arguments)) + assert not invocation.is_error, ( + f"{name} failed: {invocation.result.content[:400]}" + ) + return invocation + + # ------------------------------------------------------------------ + # every tool is reachable and registered + # ------------------------------------------------------------------ + + def test_registry_exposes_every_documented_tool(self) -> None: Review Comment: Ran the existing regression against an isolated SQLite metadata database with initialized permissions/users. The real registry invokes `create_virtual_dataset`, `generate_chart` and `generate_dashboard` through FastMCP, checks persisted IDs/relationships/creator/editor identity, then verifies Gamma cannot create any of them despite `MCP_DEV_USERNAME=admin`. It passes. The documented-tool and usable-schema registry checks also pass: 3 integration cases total. This is local integration evidence, not a production deployment or warehouse test. ########## superset/ai/orchestrator.py: ########## @@ -0,0 +1,647 @@ +# 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. +""" +Runs one assistant turn end to end. + +Sits between the HTTP layer and the runtime: loads the conversation, assembles +the prompt, resolves the tools the chosen profile allows, drives the runtime, +publishes every event to the bus, and records the outcome on the assistant +message. + +Deliberately independent of *where* it runs. The same function body serves the +inline path and the Celery path, which is what makes the execution mode a +configuration choice rather than two implementations that drift apart. +""" + +from __future__ import annotations + +import asyncio +import logging +import uuid as uuid_module +from collections.abc import AsyncIterator, Iterator +from dataclasses import dataclass +from typing import Any + +from superset.ai.events import ( + cancelled_event, + done_event, + error_event, + GENERIC_ERROR_MESSAGE, + session_event, + StreamEvent, +) +from superset.ai.llm.base import Message +from superset.ai.telemetry import bind_run, current_run, start_run +from superset.ai.types import MessageRole, MessageStatus, RunOutcome, StreamEventType +from superset.utils.decorators import transaction + +logger = logging.getLogger(__name__) + +#: Cache key prefix for a run's cancellation flag. A flag rather than a signal +#: because a worker cannot be interrupted mid-call reliably; the runtime checks +#: this between steps. +_CANCEL_PREFIX = "ai-cancel-" + +#: How long a cancellation request stays meaningful. +_CANCEL_TTL_SECONDS = 900 + +#: Stored when a run is stopped before it produced any answer, so the +#: transcript still records that the turn happened. +_STOPPED_WITHOUT_ANSWER = "_Stopped before an answer was produced._" + +#: Stored when a run exhausted its time budget without saying anything. Phrased +#: as something the user can act on, because retrying is usually the right move. +_TIMED_OUT_WITHOUT_ANSWER = ( + "The assistant ran out of time before it could answer. Please try again." +) + +#: Ceiling on the page context recorded on a message. Well below the prompt's own +#: limit: this is stored per turn and read back with the whole transcript. +_RECORDED_CONTEXT_LIMIT = 4_000 + + +@dataclass +class TurnRequest: + """One unit of work: answer the latest message on a thread.""" + + thread_uuid: str + user_id: int + run_id: str + #: Assistant message row to fill in. Created before the run starts so a + #: client that reconnects has something to attach to. + assistant_message_uuid: str + profile_key: str | None = None + #: Concrete model to pin, overriding the profile's tier. + model: str | None = None + #: What the user had on screen when they asked. Supplied by the client, + #: which is the only party that knows which tab is open, what is typed in + #: the editor and which filters are applied. + page_context: dict[str, Any] | None = None + + def to_payload(self) -> dict[str, Any]: + """Serialise for the task broker.""" + return { + "thread_uuid": self.thread_uuid, + "user_id": self.user_id, + "run_id": self.run_id, + "assistant_message_uuid": self.assistant_message_uuid, + "profile_key": self.profile_key, + "model": self.model, + "page_context": self.page_context, + } + + @classmethod + def from_payload(cls, payload: dict[str, Any]) -> TurnRequest: + """Rebuild from a broker payload.""" + return cls(**payload) + + +def new_run_id() -> str: + """Identifier for one run, used as the event-stream key.""" + return str(uuid_module.uuid4()) + + +#: Runs cancelled in this process. +#: +#: Held alongside the cache rather than instead of it. Superset's default cache +#: is a null cache, which accepts a write and discards it — so a cache-only +#: implementation would leave cancellation silently broken on a default install, +#: with the button appearing to work and nothing stopping. This set makes inline +#: execution correct with no cache at all; the cache is what carries a +#: cancellation across processes for worker execution. +_CANCELLED_LOCALLY: set[str] = set() + + +def request_cancel(run_id: str) -> None: + """ + Ask a run to stop. + + Cooperative by design: the flag is recorded here and observed by the runtime + between steps. A run blocked inside a single long model call or query will + not notice until that call returns, which is a real limit worth documenting + rather than hiding. + """ + from superset.extensions import cache_manager + + _CANCELLED_LOCALLY.add(run_id) + try: + cache_manager.cache.set( Review Comment: Fixed on this branch in b37267977d. Worker cancellation uses the shared Redis backend selected for the event bus rather than the general NullCache/process-local set. The cancellation value is Redis-compatible `b"1"`; regression coverage uses the real Redis encoder and separate cancellation objects to check shared visibility. ########## superset/ai/tasks.py: ########## @@ -0,0 +1,89 @@ +# 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. +""" +Background execution of assistant turns. + +Used when ``AI_ASSISTANT_EXECUTION_MODE`` is ``"worker"``. The task body is a +thin wrapper: all the work lives in +:func:`superset.ai.orchestrator.execute_turn`, so the two execution modes cannot +diverge in behaviour. + +To enable, add ``"superset.ai.tasks"`` to ``CeleryConfig.imports`` and set the +execution mode and a Redis event bus. +""" + +from __future__ import annotations + +import logging +from typing import Any + +from superset.ai.orchestrator import execute_turn, TurnRequest +from superset.extensions import celery_app + +logger = logging.getLogger(__name__) + + +@celery_app.task(name="ai.run_turn", bind=True, soft_time_limit=None) Review Comment: The task now explicitly pins early acknowledgement, atomically claims the pending run and records its delivery identity. Parent-process failure callbacks can finalize a child-worker failure/hard timeout without overwriting another delivery's claim; corresponding tests pass. The guide explicitly states the remaining boundary: loss of the entire Celery parent/host cannot run those callbacks, and this branch has no host-liveness reconciler. I am not claiming recovery from that failure class or exactly-once external writes. ########## superset-frontend/src/features/ai/hooks/useChatBot.ts: ########## @@ -0,0 +1,1323 @@ +/** + * 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 Conversation state and the send loop. + * + * Runs are tracked per conversation, not globally. That is the point of the + * structure: a user can start something slow in one conversation, switch to + * another and keep working, and come back to find the first still going. A single + * `isLoading` flag would have made switching away cancel or corrupt the run. + * + * The server owns the transcript. A finished run is re-read from it rather than + * assembled from the frames, so the tool calls persisted on the message are what + * the user sees, and what they see survives a reload. + */ + +import { useCallback, useEffect, useRef, useState } from 'react'; +import type { TextAreaRef } from 'antd/es/input/TextArea'; +import { logging } from '@apache-superset/core/utils'; +import { t } from '@apache-superset/core/translation'; +import { + type AiAgent, + type AiToolCall, + type ChatMessageWithMeta, + type ChatTab, + type CheckpointPayload, +} from '../types'; +import { + AGENT_STORAGE_KEY, + ChatRequestAbortedError, + ChatStreamEventError, + ChatStreamTimeoutError, + DEFAULT_AGENT_KEY, + DEFAULT_CHAT_AGENT, + cancelChatRun, + describeRequestError, + fetchAgents, + fetchSuggestedPrompts, + loadStoredAgentKey, + normalizeChatAgents, + startRun, + streamRun, + submitFeedback, +} from './chatRequest'; +import { + NEW_CHAT_NAME, + createThread, + deleteThread as deleteThreadApi, + getThread, + listThreads, + threadToTab, + updateThread, +} from './chatThreadsApi'; +import { buildQuickPrompts } from './quickPrompts'; +import { + buildPageContextPayload, + usePageContext, + type PageContext, +} from './usePageContext'; + +/** Cache of the conversation list, so the menu renders before the list arrives. */ +export const CHAT_TABS_STORAGE_KEY = 'superset-chat-tabs'; + +/** Which conversation was last open. */ +export const ACTIVE_TAB_STORAGE_KEY = 'superset-chat-active-tab'; + +/** Recent inputs, recalled with the arrow keys. */ +export const HISTORY_STORAGE_KEY = 'superset-chat-history'; + +export { AGENT_STORAGE_KEY } from './chatRequest'; + +/** How many inputs the arrow-key history keeps. */ +const MAX_INPUT_HISTORY = 50; + +/** A conversation title derived from a message is clipped to this. */ +const MAX_TAB_NAME_LENGTH = 30; + +export type ChatRunStatus = 'running' | 'cancelling'; + +/** Shared empty list, so a render with no steps yet keeps a stable identity. */ +const EMPTY_TOOL_CALLS: AiToolCall[] = []; + +interface ActiveChatRun { + requestId: string; + tabId: string; + threadId: string; + runId?: string; + controller: AbortController; + isStreaming: boolean; + liveThoughts: string; + liveToolLog: string; + /** + * Steps taken so far, as structured records rather than log lines. + * + * Carried alongside `liveToolLog` so a run in flight can be rendered the same + * way a finished one is — expandable per step, with the SQL and the rows it + * returned — instead of as a wall of text that only becomes legible once the + * transcript is re-read from the server. + */ + liveToolCalls: AiToolCall[]; + /** The page context this run was given, so the live view can show it too. */ + livePageContext?: string; + /** + * The answer so far, as the model produces it. + * + * Rendered directly: the deltas used to be folded into `liveThinking`, which + * nothing displayed, so an answer appeared in one piece the moment the run + * ended however long it had taken to generate. + */ + liveAnswer: string; + liveThinking: string; + status: ChatRunStatus; + startedAt: number; + checkpoint: CheckpointPayload | null; +} + +/** + * An identifier for a turn. + * + * Drawn from `crypto`, not `Math.random`. These become the idempotency key on a + * turn and the handle used to cancel one, so a value another session could guess + * is a correctness and a security problem rather than merely a collision risk. + */ +const generateId = (): string => { + if (typeof crypto.randomUUID === 'function') { + return crypto.randomUUID(); + } + // Older engines expose the entropy source without the convenience wrapper. + const bytes = new Uint8Array(16); + crypto.getRandomValues(bytes); + return Array.from(bytes, byte => byte.toString(16).padStart(2, '0')).join(''); +}; + +const createNewTab = (name: string = NEW_CHAT_NAME): ChatTab => ({ + id: generateId(), + name, + messages: [], + createdAt: Date.now(), +}); + +const truncateTabName = ( + name: string, + maxLength: number = MAX_TAB_NAME_LENGTH, +): string => + name.length <= maxLength ? name : `${name.substring(0, maxLength)}...`; + +const readJson = <T>(key: string, fallback: T): T => { + try { + const stored = localStorage.getItem(key); + return stored ? (JSON.parse(stored) as T) : fallback; + } catch (caught) { + logging.warn(`[ai] could not read ${key}`, caught); + return fallback; + } +}; + +const writeJson = (key: string, value: unknown): void => { + try { + localStorage.setItem(key, JSON.stringify(value)); + } catch (caught) { + logging.warn(`[ai] could not write ${key}`, caught); + } +}; + +/** + * Reconciles the server's transcript with what is already on screen. + * + * The server's copy is authoritative — it carries the tool calls — but it is not + * necessarily complete the moment a run ends, and replacing outright would then + * erase an answer the user has just read. So anything local that the server has + * not accounted for is kept, matched by identity first and by role and content + * second, which is how a locally-appended turn is recognised once the server + * returns its own copy of it under a real uuid. + */ +export const mergeMessages = ( + fromServer: ChatMessageWithMeta[], + local: ChatMessageWithMeta[], +): ChatMessageWithMeta[] => { + const serverIds = new Set(fromServer.map(message => message.id)); + const serverTurns = new Set( + fromServer.map(message => `${message.role}:${message.content}`), + ); + const unaccounted = local.filter( + message => + !serverIds.has(message.id) && + !serverTurns.has(`${message.role}:${message.content}`), + ); + return [...fromServer, ...unaccounted]; +}; + +/** + * The `page_context` body for one turn. + * + * Returns undefined when there is nothing to send, so an omitted field is + * distinguishable from an empty one. + */ +export const buildRequestPageContext = ( + context: PageContext | undefined, + directive?: string, +): Record<string, unknown> | undefined => { + const payload = context ? buildPageContextPayload(context) : undefined; + if (!directive) { + return payload; + } + const existing = payload?.helper_directives; + return { + ...payload, + helper_directives: [ + directive, + ...(Array.isArray(existing) ? existing : []), + ], + }; +}; + +export interface UseChatBotReturn { + // Conversations + chatTabs: ChatTab[]; + activeTabId: string; + activeTab: ChatTab | undefined; + threadsLoaded: boolean; + handleNewChat: () => Promise<string>; + handleSelectTab: (tabId: string) => Promise<void>; + handleDeleteTab: (tabId: string) => Promise<void>; + handleRenameTab: (tabId: string, newName: string) => void; + // Messages of the active conversation + messages: ChatMessageWithMeta[]; + // Input + inputValue: string; + setInputValue: (value: string) => void; + handleKeyDown: (event: React.KeyboardEvent) => void; + inputRef: React.RefObject<TextAreaRef>; + messagesEndRef: React.RefObject<HTMLDivElement>; + // The run in flight, if any, for the active conversation + isLoading: boolean; + isStreamingResponse: boolean; + liveThoughts: string; + liveToolLog: string; + /** Steps taken so far in the run in flight, for the structured live view. */ + liveToolCalls: AiToolCall[]; + /** The page context the run in flight was given. */ + livePageContext?: string; + /** The answer so far for the run in flight. */ + liveAnswer: string; + checkpoint: CheckpointPayload | null; + activeRunStatus: ChatRunStatus | null; + error?: string; + // Actions + sendMessage: ( + messageOverride?: string, + systemPromptOverride?: string, + ) => Promise<void>; + handleCancelRun: () => Promise<void>; + handleCheckpointContinue: () => void; + handleFeedback: (messageId: string, feedback: 'like' | 'dislike') => void; + messageFeedback: Record<string, 'like' | 'dislike'>; + // Suggestions + /** The message whose run just ended; its thought process stays open. */ + justCompletedId?: string; + quickPrompts: string[]; + loadQuickPrompts: () => void; + applyQuickPrompt: (prompt: string) => Promise<void>; + // Agent profiles + agents: AiAgent[]; + selectedAgent: string; + setSelectedAgent: (key: string) => void; + // Page context + pageContext: PageContext; + includePageContext: boolean; + toggleIncludePageContext: () => void; +} + +export const useChatBot = (): UseChatBotReturn => { + const [chatTabs, setChatTabs] = useState<ChatTab[]>(() => + readJson<ChatTab[]>(CHAT_TABS_STORAGE_KEY, []).map(tab => ({ + // The cache is a placeholder for the menu; message bodies are re-read from + // the server so a stale cache cannot show a conversation that has moved on. + ...tab, + messages: [], + })), + ); + const [activeTabId, setActiveTabId] = useState<string>(() => { + try { + return localStorage.getItem(ACTIVE_TAB_STORAGE_KEY) ?? ''; + } catch { + return ''; + } + }); + const [threadsLoaded, setThreadsLoaded] = useState(false); + const [error, setError] = useState<string | undefined>(undefined); + + const [inputValue, setInputValue] = useState(''); + const [activeRunsByTab, setActiveRunsByTab] = useState< + Record<string, ActiveChatRun> + >({}); + const [quickPrompts, setQuickPrompts] = useState<string[]>([]); + const [messageFeedback, setMessageFeedback] = useState< + Record<string, 'like' | 'dislike'> + >({}); + const [includePageContext, setIncludePageContext] = useState(true); + /** + * The assistant message whose run has only just ended. + * + * Its thought process stays open, because collapsing it the instant the answer + * lands moves everything below it — the answer the user is mid-sentence through + * jumps up the panel. Older messages start closed. + */ + const [justCompletedId, setJustCompletedId] = useState<string | undefined>(); + const [agents, setAgents] = useState<AiAgent[]>([DEFAULT_CHAT_AGENT]); + const [selectedAgent, setSelectedAgent] = useState<string>(() => + loadStoredAgentKey(AGENT_STORAGE_KEY), + ); + + const [messageHistory, setMessageHistory] = useState<string[]>(() => + readJson<string[]>(HISTORY_STORAGE_KEY, []), + ); + const [historyIndex, setHistoryIndex] = useState(-1); + const [currentDraft, setCurrentDraft] = useState(''); + + const messagesEndRef = useRef<HTMLDivElement>(null); + const inputRef = useRef<TextAreaRef>(null); + + /** + * The run map and the conversation list are also held in refs, and the refs are + * the authority. + * + * The send loop has to ask "is this still my run?" between awaits, and it cannot + * ask React: a run that starts and fails inside one batch never causes a render, + * so a ref synced at render time would still be empty and the loop would discard + * its own result as stale. Writing the ref at the point of mutation removes that + * window. The callbacks read the refs rather than the state so their identities + * do not churn on every streamed frame, which would restart effects mid-run. + */ + const activeRunsByTabRef = useRef<Record<string, ActiveChatRun>>({}); + const chatTabsRef = useRef<ChatTab[]>(chatTabs); + const activeTabIdRef = useRef(activeTabId); + activeTabIdRef.current = activeTabId; + + const updateRuns = useCallback( + ( + updater: ( + previous: Record<string, ActiveChatRun>, + ) => Record<string, ActiveChatRun>, + ) => { + activeRunsByTabRef.current = updater(activeRunsByTabRef.current); + setActiveRunsByTab(activeRunsByTabRef.current); + }, + [], + ); + + const updateTabs = useCallback( + (updater: (previous: ChatTab[]) => ChatTab[]) => { + chatTabsRef.current = updater(chatTabsRef.current); + setChatTabs(chatTabsRef.current); + }, + [], + ); + + /** Resolved when the user answers a checkpoint; see `streamRun`. */ + const checkpointGateRef = useRef<{ resolve: () => void } | null>(null); + const mountedRef = useRef(true); + useEffect( Review Comment: Fixed in 428484e846 together with the restored-run handling imported in 2d18783b14. Unmount releases each checkpoint, aborts its stream and requests cancellation for known run IDs. If unmount happens before POST returns, the returned ID is cancelled when available. Reopened pending runs retain their handle and poll stored status rather than enabling an overlapping send. The two-tab and POST-race regressions pass. Cancellation remains cooperative. -- 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]
