I3eka commented on code in PR #43133:
URL: https://github.com/apache/superset/pull/43133#discussion_r4002579857
##########
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:
Fixed in 0d5be1cba7. Removed the duplicate traversal and reused upstream
`validate_dashboard_layout`, which checks that reachable chart IDs exactly
match the requested set. Missing charts fall back to the complete generated
grid; the new regression failed before the fix. For an otherwise valid layout,
persisted parent paths are rebuilt from the validated child graph, so
stale/empty metadata does not discard the caller's layout. Both normal and
stale-parent override cases pass. The full AI plus dashboard
validation/generation selection reports 801 passed.
##########
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:
Correct: schema parity and mocked adapter tests are not proof that all three
authoring tools persist through the authenticated FastMCP bridge. The 801
passing checks in this refresh are unit tests, not that missing integration
proof. This coverage request remains open on this PR; I have not relabeled the
unit result as end-to-end verification.
##########
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:
This remains an open worker-mode issue in the base. The dedicated Redis
event bus does not make the unrelated general cache shared; a local
cancellation set cannot reach another process. Tracking the process-separated
regression request at
https://github.com/apache/superset/pull/42805#discussion_r4000933695; no fix is
claimed here.
##########
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 acknowledgement/crash-recovery contract is still unresolved in the base.
An inherited deployment-wide `acks_late` value and an early-ack default are
both insufficient evidence of exactly-once effects or terminal recovery. This
refresh does not add that guarantee; leaving the thread open.
##########
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:
Panel teardown/reconnect is still an open base lifecycle concern. The
refresh does not persist a reconnect handle or release the per-run checkpoint
on unmount, so a successful worker run cannot be inferred from the panel
reopening. Keeping this open for the shared client/run contract.
##########
superset/ai/runtime/messages.py:
##########
@@ -0,0 +1,574 @@
+# 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.
+"""
+The default runtime: a plain tool-use loop over the provider's message API.
+
+Chosen as the default because it needs nothing beyond an HTTP call — no agent
+engine subprocess, no working directory, no bundled binary — so it works with
+whatever provider a deployment configures.
+"""
+
+from __future__ import annotations
+
+import logging
+import time
+from collections.abc import AsyncIterator
+from typing import Any
+
+from superset.ai.events import (
+ assistant_delta_event,
+ checkpoint_event,
+ error_event,
+ final_event,
+ GENERIC_ERROR_MESSAGE,
+ StreamEvent,
+ thinking_event,
+ thoughts_event,
+)
+from superset.ai.llm.base import (
+ CompletionRequest,
+ LLMError,
+ LLMResponse,
+ Message,
+ StreamEventKind,
+ ToolCall,
+ ToolResult,
+)
+from superset.ai.runtime.base import BaseAgentRuntime, RunRequest, RunResult
+from superset.ai.telemetry import (
+ current_run,
+ POLICY_DENIED,
+ RunRecorder,
+ TOOL_UNAVAILABLE,
+)
+from superset.ai.types import MessageRole, ProgressStage, TokenUsage
+
+logger = logging.getLogger(__name__)
+
+#: How much of a tool's output is kept on the persisted message. The model
+#: still sees the whole thing; this is the audit copy.
+_RECORDED_OUTPUT_LIMIT = 2_000
+
+#: Size of the chunks the finished answer is delivered in.
+_DELIVERY_CHUNK_SIZE = 512
+
+#: How much reasoning is kept on the result. Reasoning can run several times
+#: longer than the answer, and this is persisted next to it.
+_RECORDED_THOUGHTS_LIMIT = 8_000
+
+_NO_ANSWER = (
+ "I wasn't able to reach an answer for that. Try narrowing the question, "
+ "or naming the dataset you have in mind."
+)
+
+
+class MessagesApiRuntime(BaseAgentRuntime):
+ """
+ Alternates model calls and tool calls until the model stops asking.
+
+ Two behaviours are worth understanding before changing this class.
+
+ First, prose the model emits *before* a tool call is treated as reasoning,
+ not answer: it becomes a ``thoughts`` event and is dropped from the answer.
+ A model narrating "the orders table looks right, let me check" is stating a
+ hypothesis it may abandon, and appending that to the answer produces a
+ reply that contradicts itself.
+
+ Second, the loop always terminates and never raises for an operational
+ failure. By the time it runs, response headers have been flushed and an
+ exception can no longer become an HTTP status, so every failure is an
event.
+ """
+
+ def __init__(self, provider: Any) -> None:
+ super().__init__(provider)
+ self._result = RunResult()
+ #: Set when the model signals it has finished answering.
+ self._finished = False
+ #: The most recent round trip's response, or ``None`` if it failed. The
+ #: turn methods are generators and cannot return a value.
+ self._last_response: LLMResponse | None = None
+ #: Whether any answer text has already been sent as it was generated.
The
+ #: finished answer is only replayed in chunks when it has not.
+ self._streamed_text = False
+
+ @property
+ def result(self) -> RunResult:
+ return self._result
+
+ async def run(self, request: RunRequest) -> AsyncIterator[StreamEvent]:
+ self._result = RunResult()
+ self._finished = False
+ self._last_response = None
+ self._streamed_text = False
+ answer_parts: list[str] = []
+
+ yield thinking_event(ProgressStage.START, "Working on your question")
+
+ # The provider's connection pool belongs to the loop this run is driven
+ # on, and the caller closes that loop as soon as the run ends. Closing
+ # here — inside the loop, however the run finishes, including when the
+ # generator is abandoned mid-way by a user pressing stop — is what
keeps
+ # a client from being finalised against a dead loop.
+ try:
+ async for event in self._turn_loop(request, answer_parts):
+ yield event
+
+ # A run that failed or was abandoned has already said so; emitting
an
+ # answer as well would contradict it.
+ if self._result.error is not None or self._result.cancelled:
+ return
+
+ answer = "\n\n".join(part for part in answer_parts if part).strip()
+ self._result.answer = answer or _NO_ANSWER
+
+ # Only replayed when nothing was streamed — a provider without
+ # streaming support still gets to deliver its answer progressively.
+ # Replaying after live text would show the answer twice.
+ if not self._streamed_text:
+ for chunk in _chunk(self._result.answer):
+ yield assistant_delta_event(chunk)
+ yield final_event(self._result.answer)
+ finally:
+ await self.provider.aclose()
+
+ async def _turn_loop(
+ self,
+ request: RunRequest,
+ answer_parts: list[str],
+ ) -> AsyncIterator[StreamEvent]:
+ """
+ Alternate model and tool calls until the model stops or a budget runs
out.
+
+ Appends to ``answer_parts`` rather than returning the answer, because
an
+ async generator cannot both yield events and return a value.
+ """
+ deadline = time.monotonic() + request.timeout_seconds
+ conversation = list(request.messages)
+
+ for turn in range(1, request.max_turns + 1):
+ self._result.turns = turn
+
+ if self._should_stop(request, deadline):
+ if self._result.timed_out:
+ yield thinking_event(
+ ProgressStage.FALLBACK,
+ "Taking longer than expected — answering with what I
have",
+ )
+ return
+
+ async for event in self._safe_turn(request, conversation, turn):
+ yield event
+ response = self._last_response
+ if response is None:
+ yield error_event()
+ return
+
+ async for event in self._consume(
+ request, response, conversation, answer_parts
+ ):
+ yield event
+
+ if self._finished or self._result.cancelled:
+ return
+
+ # Budget exhausted without the model choosing to stop.
+ yield thinking_event(
+ ProgressStage.FALLBACK,
+ "Reached the step limit — answering with what I have",
+ )
+
+ async def _consume(
+ self,
+ request: RunRequest,
+ response: LLMResponse,
+ conversation: list[Message],
+ answer_parts: list[str],
+ ) -> AsyncIterator[StreamEvent]:
+ """Act on one model response, running any tools it asked for."""
+ if response.thinking:
+ self._record_thoughts(response.thinking)
+ yield thoughts_event(response.thinking)
+
+ if not response.wants_tools:
+ self._finished = True
+ if response.text:
+ answer_parts.append(response.text)
+ # Recorded as it arrives, not just at the end, so a run stopped
+ # after this point still persists what the user already saw.
+ self._result.answer = "\n\n".join(
+ part for part in answer_parts if part
+ ).strip()
+ return
+
+ # Prose accompanying a tool call is reasoning, not answer.
+ if response.text:
+ self._record_thoughts(response.text)
+ yield thoughts_event(response.text)
+
+ conversation.append(
+ Message(
+ role=MessageRole.ASSISTANT,
+ content=response.text,
+ tool_calls=list(response.tool_calls),
+ )
+ )
+
+ results: list[ToolResult] = []
+ async for event in self._run_tools(request, response.tool_calls,
results):
+ yield event
+
+ conversation.append(Message(role=MessageRole.USER,
tool_results=results))
+
+ async def _run_tools(
+ self,
+ request: RunRequest,
+ calls: list[ToolCall],
+ results: list[ToolResult],
+ ) -> AsyncIterator[StreamEvent]:
+ """Execute this turn's tool calls, appending outcomes to
``results``."""
+ for call in calls:
+ if self._cancelled(request):
+ self._result.cancelled = True
+ return
+
+ yield thinking_event(
+ ProgressStage.TOOL,
+ f"Running {call.name}",
+ {"tool_name": call.name},
+ )
+ result, detail = self._invoke_tool(request, call)
+ results.append(result)
+ record = self._record_call(call, result, detail)
+
+ # The frame carries the same record that is persisted, rather than
a
+ # subset assembled separately. The subset was missing the arguments
+ # and the output, so a step expanded during a run showed nothing at
+ # all unless its tool happened to supply a display — and then
filled
+ # itself in on reload, which looked like the detail arrived late.
+ # Sharing one record makes that class of drift impossible.
+ yield checkpoint_event(
+ f"{'Failed' if result.is_error else 'Finished'} {call.name}",
+ # ``tool_name`` as well as ``name``: the progress frames use
that
+ # key, so a consumer reading either finds what it expects.
+ {"tool_name": call.name, **record},
+ )
+
+ async def _safe_turn(
+ self,
+ request: RunRequest,
+ conversation: list[Message],
+ turn: int,
+ ) -> AsyncIterator[StreamEvent]:
+ """
+ One model round trip, converting failure into a ``None`` response.
+
+ A generator rather than a coroutine so the answer can reach the client
as
+ the model produces it. The response is handed back on
+ :attr:`_last_response` because an async generator cannot both yield
events
+ and return a value — the same reason ``_turn_loop`` writes into
+ ``answer_parts``.
+
+ The failure detail goes to the log; the caller emits a message that
cannot
+ leak a URL, a credential or a fragment of someone else's query.
+ """
+ recorder = current_run()
+ started = time.monotonic()
+ self._last_response = None
+ try:
+ async for event in self._one_turn(request, conversation):
+ yield event
+ except LLMError as ex:
+ logger.warning("AI provider error on turn %s: %s", turn, ex)
+ self._result.error = str(ex)
+ self._trace_model_call(recorder, request, turn, started, error=ex)
+ self._last_response = None
+ return
+ except Exception as ex: # pylint: disable=broad-except
+ logger.exception("Unexpected error in AI runtime on turn %s", turn)
+ self._result.error = GENERIC_ERROR_MESSAGE
+ self._trace_model_call(recorder, request, turn, started, error=ex)
+ self._last_response = None
+ return
+ self._trace_model_call(
+ recorder, request, turn, started, response=self._last_response
+ )
+
+ def _trace_model_call(
+ self,
+ recorder: RunRecorder,
+ request: RunRequest,
+ turn: int,
+ started: float,
+ response: LLMResponse | None = None,
+ error: BaseException | None = None,
+ ) -> None:
+ """
+ Report one round trip to telemetry.
+
+ Content is passed as-is; whether any of it survives into a trace is the
+ redaction policy's decision, made in one place rather than here.
+ """
+ if not recorder.enabled:
+ return
+ usage = response.usage if response is not None else TokenUsage()
+ recorder.model_call(
+ turn=turn,
+ # The concrete identifier when the provider reported one, and the
+ # capability tier otherwise, so a trace can always be grouped by
+ # what the run asked for.
+ model=usage.get("model") or request.model_alias.value,
+ duration_ms=int((time.monotonic() - started) * 1000),
+ input_tokens=usage.get("input_tokens"),
+ output_tokens=usage.get("output_tokens"),
+ stop_reason=response.stop_reason if response is not None else None,
+ error_type=type(error).__name__ if error is not None else None,
+ system_prompt=request.system_prompt,
+ response_text=response.text if response is not None else None,
+ )
+ if error is not None:
+ recorder.error(error)
+
+ async def _one_turn(
+ self,
+ request: RunRequest,
+ conversation: list[Message],
+ ) -> AsyncIterator[StreamEvent]:
+ """
+ Call the model once, yielding answer text as the model produces it.
+
+ Streaming is used when the provider supports it. The assembled response
+ is left on :attr:`_last_response` rather than returned, because a
+ generator cannot do both; it has the same shape either way, so callers
do
+ not branch on which path ran.
+ """
+ completion = CompletionRequest(
+ messages=conversation,
+ system=request.system_prompt,
+ model_alias=request.model_alias,
+ tools=tuple(request.tools.definitions()) if request.tools else (),
+ )
+
+ if not self.provider.supports_streaming:
+ self._last_response = await self.provider.complete(completion)
+ return
+
+ text_parts: list[str] = []
+ thinking_parts: list[str] = []
+ tool_calls: list[ToolCall] = []
+ usage = None
+
+ async for event in self.provider.stream(completion):
+ if event.kind is StreamEventKind.TEXT:
+ text_parts.append(event.text)
+ # Forwarded as it arrives. Text is buffered as well, because
the
+ # turn is only known to be an answer once the model stops
without
+ # asking for a tool — prose before a tool call is reasoning,
and
+ # is re-routed as such in ``_consume``. A reader that has
already
+ # seen it replaces its copy on the ``final`` frame.
+ if event.text:
+ self._streamed_text = True
+ yield assistant_delta_event(event.text)
Review Comment:
The cancellation transcript-parity issue remains open: partial live prose
and persisted thoughts can still be classified differently. The compatibility
changes do not alter that streaming classification path. It needs an
abnormal-termination regression in the base runtime/client pair.
--
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]
