Re: [PR] feat(ai): add an opt-in AI assistant for natural-language data questions [superset]

[via GitHub]

sadpandajoe commented on code in PR #42805:
URL: https://github.com/apache/superset/pull/42805#discussion_r3739999185


##########
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":

Review Comment:
   Reopening an inline stream for the same `run_id` executes the whole turn 
again, including model/tool calls, and the later finalizer can overwrite an 
already completed answer. Could this path atomically claim only a pending 
message and replay the stored terminal result otherwise?



##########
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(
+            f"{_CANCEL_PREFIX}{run_id}", True, timeout=_CANCEL_TTL_SECONDS
+        )
+    except Exception:  # pylint: disable=broad-except
+        logger.warning("Could not record cancellation for AI run %s", run_id)
+
+
+def is_cancelled(run_id: str) -> bool:
+    """Whether a stop has been requested for this run."""
+    from superset.extensions import cache_manager
+
+    if run_id in _CANCELLED_LOCALLY:
+        return True
+    try:
+        return bool(cache_manager.cache.get(f"{_CANCEL_PREFIX}{run_id}"))
+    except Exception:  # pylint: disable=broad-except
+        # A cache that cannot be read must not make every run appear cancelled;
+        # that would stop all inference the moment the cache went away.
+        return False
+
+
+def clear_cancel(run_id: str) -> None:
+    """Drop a run's cancellation flag."""
+    from superset.extensions import cache_manager
+
+    _CANCELLED_LOCALLY.discard(run_id)
+    try:
+        cache_manager.cache.delete(f"{_CANCEL_PREFIX}{run_id}")
+    except Exception:  # pylint: disable=broad-except
+        logger.debug("Could not clear cancellation flag for AI run %s", run_id)
+
+
+def stream_turn(request: TurnRequest) -> Iterator[StreamEvent]:
+    """
+    Answer a turn, yielding events as they happen.
+
+    This is the primary entry point. Inline execution consumes it directly from
+    inside the streaming response, which means the producer and the reader are
+    the same process by construction — important because Superset runs several
+    web workers, and a turn that published to one process's in-memory queue
+    while the browser's stream landed on another would appear to hang forever.
+
+    Never raises for an operational failure: a failure is an ``error`` event 
and
+    an ``error`` message status, because the caller may already have flushed
+    response headers or may be a worker with no one to report to.
+    """
+    recorder = start_run(
+        run_id=request.run_id,
+        thread_uuid=request.thread_uuid,
+        user_id=request.user_id,
+    )
+    # Shared with ``_run`` so the ``finally`` below can see the runtime's
+    # partial result and whether the message was already written.
+    state: dict[str, Any] = {}
+    try:
+        # Bound here rather than inside ``_run`` so that a run which fails 
before
+        # it has resolved a profile still produces a start and an end, and so
+        # that the runtime can report its own spans without the runtime 
contract
+        # growing a telemetry parameter.
+        with bind_run(recorder):
+            recorder.run_started()
+            yield from _run(request, state)
+    except Exception as ex:  # pylint: disable=broad-except
+        logger.exception("AI turn failed for run %s", request.run_id)
+        recorder.error(ex)
+        recorder.run_ended(outcome=RunOutcome.ERROR)
+        answer, extra = _partial_from_state(state)
+        extra["outcome"] = RunOutcome.ERROR.value
+        _finalise_message(
+            request.assistant_message_uuid,
+            # The generic text rather than the exception: this is persisted and
+            # served back to the browser, so it must not carry internals. The
+            # detail is in the log line above, keyed by run id.
+            content=answer or GENERIC_ERROR_MESSAGE,
+            status=MessageStatus.ERROR,
+            extra=extra,
+        )
+        state["finalised"] = True
+        yield error_event()
+        yield done_event(ok=False)
+    finally:
+        clear_cancel(request.run_id)
+        # A client that stops the run, or simply navigates away, abandons this
+        # generator part-way through. Nothing above will have written the
+        # message, so it would otherwise sit in ``streaming`` with no content
+        # for ever — the user loses both the partial answer and any record that
+        # the turn happened. Persist whatever was produced.
+        _abandon_message(request.assistant_message_uuid, state)
+        # Idempotent, so the ordinary paths above win.
+        recorder.run_ended(outcome=RunOutcome.CANCELLED)
+
+
+def execute_turn(request: TurnRequest) -> RunOutcome:
+    """
+    Answer a turn, publishing events to the event bus.
+
+    Used by worker execution, where the reader is in another process. Shares 
its
+    whole body with :func:`stream_turn` so the two execution modes cannot drift
+    apart in behaviour.
+    """
+    from superset.ai.eventbus import get_event_bus
+
+    bus = get_event_bus()
+    outcome = RunOutcome.SUCCESS
+
+    for event in stream_turn(request):
+        bus.publish(request.run_id, event)
+        if event.type is StreamEventType.ERROR:
+            outcome = RunOutcome.ERROR
+        elif event.type is StreamEventType.CANCELLED:
+            outcome = RunOutcome.CANCELLED
+        elif event.type is StreamEventType.DONE and not 
event.payload.get("ok"):
+            # A run that ended un-ok without an explicit error frame timed out.
+            if outcome is RunOutcome.SUCCESS:
+                outcome = RunOutcome.TIMEOUT
+
+    return outcome
+
+
+def _run(request: TurnRequest, state: dict[str, Any]) -> Iterator[StreamEvent]:
+    """Assemble and drive the run. See :func:`stream_turn` for error policy."""
+    from superset.ai.factories import (
+        get_profiles,
+        get_provider,
+        get_runtime,
+        get_tools_for_profile,
+    )
+    from superset.ai.policy import load_policy_chain
+    from superset.ai.runtime.base import RunRequest
+    from superset.daos.ai import AIChatMessageDAO, AIChatThreadDAO
+
+    recorder = current_run()
+
+    thread = AIChatThreadDAO.find_by_uuid_for_user(request.thread_uuid, 
request.user_id)
+    if thread is None:
+        # The thread vanished between accepting the message and running it.
+        recorder.run_ended(outcome=RunOutcome.ERROR)
+        yield error_event("That conversation is no longer available.")
+        yield done_event(ok=False)
+        return
+
+    profile = get_profiles().get(request.profile_key)
+    tools = get_tools_for_profile(profile)
+    provider = get_provider()
+    runtime = get_runtime(provider)
+    state["runtime"] = runtime
+
+    yield session_event(request.thread_uuid, request.assistant_message_uuid)
+
+    from superset.ai.page_context import render_page_context
+
+    history = _build_history(AIChatMessageDAO.find_for_thread(thread))

Review Comment:
   If Q2 is posted before Q1's worker builds history, both runs load the full 
thread and treat Q2 as the latest question, so Q1's assistant row can receive 
an answer to Q2. Could `TurnRequest` carry the associated user-message boundary 
and build history only through that turn?



##########
superset-frontend/src/features/ai/index.ts:
##########
@@ -0,0 +1,100 @@
+/**
+ * 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 Registers the AI assistant as a chat provider.
+ *
+ * The chat host owns mounting, the floating-versus-docked layout and the panel
+ * width, so all this contributes is a trigger, a panel and a descriptor.
+ * Importing this module registers the provider; the host is itself gated on 
the
+ * extensions flag and an authenticated user, so registration is only gated 
here
+ * on the assistant being switched on.
+ */
+
+import { FeatureFlag, isFeatureEnabled } from '@superset-ui/core';
+import { t } from '@apache-superset/core/translation';
+import { chat } from 'src/core/chat';
+import { AiAssistantPanel } from './AiAssistantPanel';
+import { AiAssistantTrigger } from './AiAssistantTrigger';
+
+export const AI_CHAT_ID = 'superset.ai-assistant';
+
+type ChatRegistration = ReturnType<typeof chat.registerChat>;
+
+let registration: ChatRegistration | undefined;
+
+/**
+ * Whether this page is a chrome-less render rather than someone using 
Superset.
+ *
+ * Covers both routes that produce one: `?standalone=` for a dashboard or chart
+ * embedded in an iframe or captured for a report, and the `/embedded/` route 
used
+ * by the embedding SDK. Neither is a place for an assistant — a screenshot 
would
+ * capture the trigger, and an embedded dashboard on someone else's site 
should not
+ * offer a chat panel at all.
+ *
+ * Read from the URL rather than from Redux so it holds for every page type,
+ * including those whose reducers are not registered.
+ */
+export function isChromelessRender(): boolean {
+  try {
+    const { search, pathname } = window.location;
+    const standalone = new URLSearchParams(search).get('standalone');
+    // Any non-zero value hides chrome; `0` and an absent param do not.
+    if (standalone && standalone !== '0') {
+      return true;
+    }
+    return pathname.includes('/embedded/');
+  } catch {
+    // Without a location there is no page to decorate either way.
+    return false;
+  }
+}
+
+/**
+ * Registers the assistant, unless the feature is off, this is a chrome-less
+ * render, or it is already registered. Returns the Disposable that unregisters
+ * it.
+ */
+export function registerAiAssistant(): ChatRegistration | undefined {
+  if (
+    registration ||
+    !isFeatureEnabled(FeatureFlag.AiAssistant) ||
+    isChromelessRender()
+  ) {
+    return registration;
+  }
+  registration = chat.registerChat(

Review Comment:
   Following the documented setup registers this chat, but the existing app 
host renders chats only when `ENABLE_EXTENSIONS` is also enabled; that flag 
defaults off and the guide says only `AI_ASSISTANT` plus a provider are 
required. Could the built-in assistant mount independently, or should the 
second flag be part of the validated enablement contract?



##########
superset-frontend/src/features/ai/types.ts:
##########
@@ -0,0 +1,469 @@
+/**
+ * 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 AI assistant wire contract and the parsers that produce 
it.
+ *
+ * Every shape here arrives as untrusted JSON, either in a REST body or in an
+ * SSE frame, so each one has a parser rather than a cast. A missing or
+ * wrongly-typed field degrades to a default instead of throwing inside a
+ * render: a malformed frame in the middle of a run must not blank a transcript
+ * the user has already read.
+ */
+
+/**
+ * A parsed JSON value. Deliberately not the platform `JsonObject`, whose index
+ * signature is `any` and would erase checking on every value read from it.
+ */
+export type JsonScalar = string | number | boolean | null;
+export type JsonData = JsonScalar | JsonData[] | { [key: string]: JsonData };
+export type JsonRecord = { [key: string]: JsonData };
+
+export type AiMessageRole = 'user' | 'assistant' | 'system';
+
+// ---------------------------------------------------------------------------
+// Panel-facing shapes
+//
+// The panel keeps one conversation per tab and renders messages from these,
+// rather than from the wire shapes below, so a tab that has not been fetched 
yet
+// and one loaded from the server are the same thing to the renderer.
+// ---------------------------------------------------------------------------
+
+export interface ChatMessageWithMeta {
+  /** The server message uuid once persisted; a local id until then. */
+  id: string;
+  role: 'user' | 'assistant';
+  content: string;
+  timestamp: number;
+  /** Reasoning and the tool log as one block, for the flat rendering used 
while
+   * a run is streaming and there are no structured steps yet. */
+  thinking?: string;
+  /** Reasoning on its own. Kept apart from `thinking` so the structured view 
can
+   * show it without also repeating the tool log it renders as steps. */
+  thoughts?: string;
+  /** The page context this turn was given, for the "Context used" step. */
+  pageContext?: string;
+  /** Steps the assistant took, as persisted on the message. */
+  toolCalls?: AiToolCall[];
+  /** True while a locally created message has no server uuid, which is what
+   * disables feedback on it: `POST feedback` is keyed by message uuid. */
+  pending?: boolean;
+  /** This user's stored rating, so the thumbs survive a reload. */
+  liked?: boolean;
+}
+
+export interface ChatTab {
+  id: string;
+  name: string;
+  messages: ChatMessageWithMeta[];
+  createdAt: number;
+  updatedAt?: number;
+  /** The server conversation this tab is backed by. */
+  threadId?: string;
+}
+
+/**
+ * A `checkpoint` frame, rendered as a pause in the transcript.
+ *
+ * `remaining_tasks` and `estimated_duration` are read from the frame's `meta`
+ * when present; a checkpoint without them still renders its summary.
+ */
+export interface CheckpointPayload {
+  summary: string;
+  remaining_tasks?: string[];
+  estimated_duration?: string;
+  elapsed_seconds?: number;
+  seconds_remaining?: number;
+  turn_count?: number;
+  turns_remaining?: number;
+  /** The step the checkpoint describes, when it carries one. */
+  toolCall?: AiToolCall;
+  /**
+   * Whether this checkpoint is a gate the user must clear, rather than a
+   * milestone that scrolls past.
+   *
+   * Opt-in on purpose, and optional so a caller constructing a milestone does
+   * not have to say so. A server that reports every finished tool call as a
+   * checkpoint would otherwise pause the stream on each one, leaving the panel
+   * showing progress long after the answer had arrived.
+   */
+  requiresConfirmation?: boolean;
+}
+
+/** Detail of the `superset-ai-action` event other features dispatch. */
+export interface AIActionPayload {
+  /** Sent as the user's message. */
+  prompt: string;
+  /** Prepended to the conversation as a system directive. */
+  systemPrompt?: string;
+  /** Name for the conversation the action opens. */
+  tabName?: string;
+}
+
+export type AIActionEventDetail = AIActionPayload;
+
+/** Value of `stage` on a `thinking` frame. */
+export type AiThinkingStage =
+  | 'start'
+  | 'prompt'
+  | 'agent'
+  | 'tool'
+  | 'reasoning'
+  | 'context'
+  | 'fallback'
+  | 'error'
+  | 'usage';
+
+/** An agent profile from `GET /api/v1/ai/agent/`. */
+export interface AiAgent {
+  key: string;
+  name: string;
+  description?: string;
+  tools: string[];
+}
+
+/** A conversation from `/api/v1/ai/thread/`. */
+export interface AiThread {
+  uuid: string;
+  title?: string;
+  status?: string;
+  agentKey?: string;
+  createdOn?: string;
+  /** Last activity, which is what the conversation list is ordered and dated 
by. */
+  changedOn?: string;
+  messageCount?: number;
+}
+
+/**
+ * A display whose `kind` the frontend has no renderer for. Tools are free to 
add
+ * kinds, so an unrecognised one degrades to the generic step detail rather 
than
+ * hiding the step.
+ */
+export interface AiOpaqueDisplay {
+  kind?: string;
+}
+
+/** The `sql_result` display: what the warehouse ran, and what came back. */
+export interface AiSqlResultDisplay {
+  kind: 'sql_result';
+  /**
+   * Which connection ran the statement. Carried so "Run in SQL Lab" opens an
+   * editor already pointed at it, rather than relying on SQLLAB_DEFAULT_DBID,
+   * which most deployments leave unset.
+   */
+  databaseId?: number;
+  databaseName?: string;
+  executedSql?: string;
+  /** The statement was clipped for display, so it may not be runnable as-is. 
*/
+  executedSqlTruncated: boolean;
+  columns: string[];
+  rows: JsonRecord[];
+  rowCount?: number;
+  /** Fewer rows are shown than the query returned. */
+  sampleOnly: boolean;
+  /** The query result itself was capped before the model saw it. */
+  truncated: boolean;
+  durationMs?: number;
+}
+
+/** Detail a tool attaches to its step for the UI to render. */
+export type AiToolDisplay = AiSqlResultDisplay | AiOpaqueDisplay;
+
+export const isSqlResultDisplay = (
+  display: AiToolDisplay | undefined,
+): display is AiSqlResultDisplay => display?.kind === 'sql_result';
+
+/**
+ * One tool invocation. The same shape describes a live `checkpoint` frame and 
a
+ * tool call persisted on a message, so the activity UI has one input whether
+ * the run is streaming or was loaded from the server.
+ */
+export interface AiToolCall {
+  name: string;
+  ok: boolean;
+  durationMs?: number;
+  /** The tool clipped its own output before handing it to the model. */
+  truncated: boolean;
+  /** Arguments the model passed, kept so a surprising result can be 
explained. */
+  args?: JsonRecord;
+  /** Recorded tool output, clipped by the backend. */
+  output?: string;
+  error?: string;
+  display?: AiToolDisplay;
+}
+
+export interface AiMessage {
+  uuid: string;
+  role: AiMessageRole;
+  content: string;
+  createdOn?: string;
+  /** Persisted by the backend, which is what lets the activity survive a 
reload. */
+  toolCalls: AiToolCall[];
+  /** Model reasoning. Never rendered as part of the answer. */
+  thoughts?: string;
+  /** What the assistant was told about the user's screen for this turn, as it 
was
+   * sent. Recorded because an answer that looks wrong is usually an answer 
about
+   * a different slice of data than the reader assumed. */
+  pageContext?: string;
+  error?: string;
+  /** The reading user's own rating, or undefined if they have not rated it. 
Lets
+   * the thumbs show a verdict that was left before a reload. */
+  liked?: boolean;
+}
+
+/** Identifies the run started by `POST /thread/<uuid>/message`. */
+export interface AiRunHandle {
+  threadUuid: string;
+  /** The user message that was just stored. */
+  messageUuid: string;
+  /** The assistant row the run will write into, created before the run 
starts. */
+  assistantMessageUuid?: string;
+  runId: string;
+}
+
+export const isDefined = <T>(value: T | undefined): value is T =>
+  value !== undefined;
+
+export const isRecord = (value: JsonData | undefined): value is JsonRecord =>
+  typeof value === 'object' && value !== null && !Array.isArray(value);
+
+/** Parses a JSON document, returning undefined rather than throwing. */
+export function parseJson(raw: string): JsonData | undefined {
+  try {
+    // JSON.parse is declared as returning `any`; funnel it through `unknown` 
so
+    // nothing downstream inherits an unchecked type.
+    const parsed: unknown = JSON.parse(raw);
+    return parsed as JsonData;
+  } catch {
+    return undefined;
+  }
+}
+
+export const readRecord = (
+  from: JsonRecord,
+  key: string,
+): JsonRecord | undefined => {
+  const value = from[key];
+  return isRecord(value) ? value : undefined;
+};
+
+export const readString = (
+  from: JsonRecord,
+  key: string,
+): string | undefined => {
+  const value = from[key];
+  return typeof value === 'string' ? value : undefined;
+};
+
+export const readNumber = (
+  from: JsonRecord,
+  key: string,
+): number | undefined => {
+  const value = from[key];
+  return typeof value === 'number' && Number.isFinite(value)
+    ? value
+    : undefined;
+};
+
+/** Absent, null and non-boolean values all read as false. */
+export const readBoolean = (from: JsonRecord, key: string): boolean =>
+  from[key] === true;
+
+/**
+ * A boolean that keeps the difference between false and absent.
+ *
+ * Needed where a field is genuinely tri-state — a rating is up, down, or not
+ * given — and collapsing absent to false would render an unrated message as a
+ * thumbs-down.
+ */
+export const readOptionalBoolean = (
+  from: JsonRecord,
+  key: string,
+): boolean | undefined =>
+  typeof from[key] === 'boolean' ? (from[key] as boolean) : undefined;
+
+const readArray = (from: JsonRecord, key: string): JsonData[] => {
+  const value = from[key];
+  return Array.isArray(value) ? value : [];
+};
+
+export const readStringArray = (from: JsonRecord, key: string): string[] =>
+  readArray(from, key).filter(
+    (item): item is string => typeof item === 'string',
+  );
+
+export const readRecordArray = (from: JsonRecord, key: string): JsonRecord[] =>
+  readArray(from, key).filter(isRecord);
+
+const MESSAGE_ROLES: readonly string[] = ['user', 'assistant', 'system'];
+
+const isMessageRole = (value: string | undefined): value is AiMessageRole =>
+  value !== undefined && MESSAGE_ROLES.includes(value);
+
+export function parseToolDisplay(
+  value: JsonData | undefined,
+): AiToolDisplay | undefined {
+  if (!isRecord(value)) {
+    return undefined;
+  }
+  const kind = readString(value, 'kind');
+  if (kind !== 'sql_result') {
+    return { kind };
+  }
+  return {
+    kind: 'sql_result',
+    databaseName: readString(value, 'database_name'),
+    // `executed_sql` is what the SQL tool writes. `sql` is accepted as well
+    // because the published event contract names the field that way, and a
+    // step whose SQL is not shown defeats the point of the activity block.
+    executedSql: readString(value, 'executed_sql') ?? readString(value, 'sql'),
+    databaseId: readNumber(value, 'database_id'),
+    executedSqlTruncated: readBoolean(value, 'executed_sql_truncated'),
+    columns: readStringArray(value, 'columns'),
+    rows: readRecordArray(value, 'rows'),
+    rowCount: readNumber(value, 'row_count'),
+    sampleOnly: readBoolean(value, 'sample_only'),
+    truncated: readBoolean(value, 'truncated'),
+    durationMs: readNumber(value, 'duration_ms'),
+  };
+}
+
+/**
+ * Parses one tool invocation.
+ *
+ * Accepts both spellings of the name: a persisted record uses `name`, while 
the
+ * `meta` of a live `checkpoint` frame uses `tool_name`.
+ */
+export function parseToolCall(
+  value: JsonData | undefined,
+): AiToolCall | undefined {
+  if (!isRecord(value)) {
+    return undefined;
+  }
+  const name = readString(value, 'name') ?? readString(value, 'tool_name');
+  if (!name) {
+    return undefined;
+  }
+  return {
+    name,
+    // A record without `ok` is treated as a success: painting a completed step
+    // red because a field is missing is the worse failure mode.
+    ok: value.ok !== false,
+    durationMs: readNumber(value, 'duration_ms'),
+    truncated: readBoolean(value, 'truncated'),
+    args: readRecord(value, 'arguments'),
+    output: readString(value, 'output'),
+    error: readString(value, 'error'),
+    display: parseToolDisplay(value.display),
+  };
+}
+
+export function parseMessage(
+  value: JsonData | undefined,
+): AiMessage | undefined {
+  if (!isRecord(value)) {
+    return undefined;
+  }
+  const uuid = readString(value, 'uuid');
+  const role = readString(value, 'role');
+  if (!uuid || !isMessageRole(role)) {
+    return undefined;
+  }
+  // Tool calls are stored in the message's `extra` blob. A serializer that
+  // hoists them to the top level is read too, because which of the two ships 
is
+  // not pinned by the API yet.
+  const extra = readRecord(value, 'extra') ?? {};

Review Comment:
   After a worker-mode reload, this parser discards both the pending message 
status and `extra.run_id`, and the initial-load path never reopens its 
replayable stream. The worker can finish while the reopened panel keeps a blank 
stale answer until another remount; could the client retain this run state and 
rejoin automatically as documented?



##########
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)
+def run_turn(self: Any, payload: dict[str, Any]) -> str:  # noqa: ARG001
+    """
+    Answer one assistant turn.
+
+    ``acks_late`` is deliberately not set: a turn costs money to run, so

Review Comment:
   A hard worker loss after acknowledgement bypasses the runtime's 
exception/finalization path, leaving the assistant row `pending` or `streaming` 
forever; the stream eventually times out, but the message does not record 
failure as this comment states. Should worker mode add a durable lease or 
failure-reconciliation path?



##########
superset/config.py:
##########
@@ -2897,6 +2913,326 @@ def EMAIL_HEADER_MUTATOR(  # pylint: 
disable=invalid-name,unused-argument  # noq
     "CACHE_REDIS_SSL_CA_CERTS": None,
 }
 
+# ---------------------------------------------------------
+# AI assistant
+# ---------------------------------------------------------
+# Requires the AI_ASSISTANT feature flag. Superset ships no model provider and
+# talks to no model vendor by default: until AI_LLM_PROVIDER_CLASS names a
+# usable provider the assistant's endpoints return 404.
+#
+# Dotted path to a superset.ai.llm.base.BaseLLMProvider subclass. Point this at
+# a vendor provider, an OpenAI-compatible endpoint, a self-hosted model, or a
+# private gateway. Everything vendor-specific — base URLs, authentication,
+# model naming — belongs in the provider, not here.
+AI_LLM_PROVIDER_CLASS: str | None = None
+
+# Keyword arguments passed to the provider's constructor. Contents are entirely
+# provider-defined. Keep credentials out of this file: read them from the
+# environment or a secret store in your own config.
+#
+#   AI_LLM_PROVIDER_CONFIG = {
+#       "api_key": os.environ["MY_LLM_API_KEY"],
+#       "base_url": "https://llm.internal.example.com/v1";,
+#       "models": {
+#           "default": "some-balanced-model",
+#           "fast": "some-small-model",
+#           "reasoning": "some-large-model",
+#       },
+#   }
+AI_LLM_PROVIDER_CONFIG: dict[str, Any] = {}
+
+# Dotted path to a superset.ai.runtime.base.BaseAgentRuntime subclass driving
+# the tool-use loop.
+AI_AGENT_RUNTIME_CLASS = "superset.ai.runtime.messages.MessagesApiRuntime"
+
+# Where a turn is executed.
+#
+#   "inline"  — in the web worker handling the request. No extra 
infrastructure,
+#               but a turn occupies a worker for its whole duration.
+#   "worker"  — handed to Celery; the request streams events from the event 
bus.
+#               Survives a browser reconnect and keeps web workers free, at the
+#               cost of requiring Celery and a shared event bus.
+AI_ASSISTANT_EXECUTION_MODE: Literal["inline", "worker"] = "inline"
+
+# How streamed events travel from producer to the HTTP response.
+#
+#   "memory" — an in-process queue. Correct only when the producer and the
+#              streaming request are the same process, i.e. inline execution.
+#   "redis"  — Redis streams, via the same cache backend the async-query
+#              channel uses. Required for "worker" execution mode.
+AI_ASSISTANT_EVENT_BUS: Literal["memory", "redis"] = "memory"
+
+# Redis connection for the AI event bus. Required when AI_ASSISTANT_EVENT_BUS 
is
+# "redis". Streams need commands the general-purpose cache client does not
+# expose, so this is configured separately rather than borrowed from
+# CACHE_CONFIG. The accepted shape matches
+# GLOBAL_ASYNC_QUERIES_CACHE_BACKEND; point both at the same Redis if you like.
+AI_ASSISTANT_EVENT_BUS_CACHE_CONFIG: dict[str, Any] = {
+    "CACHE_TYPE": "RedisCache",
+    "CACHE_REDIS_HOST": "localhost",
+    "CACHE_REDIS_PORT": 6379,
+    "CACHE_REDIS_USER": "",
+    "CACHE_REDIS_PASSWORD": "",
+    "CACHE_REDIS_DB": 0,
+    "CACHE_DEFAULT_TIMEOUT": 300,
+    "CACHE_REDIS_SSL": False,
+}
+
+# Key prefix for AI event streams when the Redis bus is in use.
+AI_ASSISTANT_EVENT_STREAM_PREFIX = "ai-events-"
+
+# How long a run's event stream is retained, in seconds. Bounds how late a
+# reconnecting browser can still pick up a run it lost.
+AI_ASSISTANT_EVENT_TTL_SECONDS = 900
+
+# Named agent profiles, merged over the built-ins by key. Each value is a dict
+# of fields to override, so narrowing one profile does not mean restating the
+# rest. The most important field is "tools": which tools that profile may
+# invoke. An unknown tool name is a startup error, not a silent omission.
+#
+#   AI_AGENT_PROFILES = {
+#       # Take the shipped default but forbid raw SQL.
+#       "default": {"tools": ["search_assets", "get_schema"]},
+#       # Let the analyst profile think harder and longer.
+#       "analyst": {"model_alias": "reasoning", "max_turns": 60},
+#       # Add a profile only some users may select.
+#       "deep": {
+#           "name": "Deep analysis",
+#           "tools": ["search_assets", "get_schema", "execute_sql"],
+#           "required_permission": ("can_write", "AIAssistant"),
+#       },
+#   }
+AI_AGENT_PROFILES: dict[str, Any] = {}
+
+# Ceiling on model round trips in a single turn. A turn that needs more than
+# this is answered with what it has rather than looping indefinitely.
+AI_AGENT_MAX_TURNS = 20
+
+# Wall-clock budget for one turn, in seconds.
+AI_AGENT_TIMEOUT_SECONDS = 300
+
+# Pre-tool-use guards, applied in order. Each is a dotted path to a
+# superset.ai.policy.ToolPolicy implementation. These bound blast radius; they
+# do not replace the per-object authorization checks inside each tool.
+AI_AGENT_TOOL_POLICIES: list[str] = [
+    "superset.ai.policy.ReadOnlySqlPolicy",
+    "superset.ai.policy.IdentifierPolicy",
+    "superset.ai.policy.ForeignToolPolicy",
+]
+
+# Rows and bytes a single tool result may return before it is truncated.
+# Model context is finite, and an unbounded result set exhausts it.
+AI_AGENT_MAX_RESULT_ROWS = 500
+AI_AGENT_MAX_RESULT_BYTES = 256 * 1024
+
+# External MCP servers whose tools may be offered to an agent profile. Superset
+# ships none and integrates with no third-party service: with this empty, 
nothing
+# in superset.ai.mcp is ever reached and the assistant behaves exactly as it 
does
+# without it.
+#
+# A server listed here is only *available*. It is used by an agent profile that
+# names it in its "mcp_servers" field, via AI_AGENT_PROFILES. A profile naming 
a
+# server that is not configured here is an error, not a silently shorter tool
+# list.
+#
+#   AI_AGENT_MCP_SERVERS = {
+#       # The key is the server name. It becomes part of every tool name this
+#       # server contributes, so keep it short: letters, digits, hyphens and
+#       # underscores, and no double underscore.
+#       "acme_catalog": {
+#           # Required. Absolute http:// or https:// endpoint.
+#           "url": "https://mcp.acme.internal/mcp";,
+#           # "streamable_http" (default) or "sse".
+#           "transport": "streamable_http",
+#           # The ONLY headers sent to this server. Superset never forwards the
+#           # user's session cookie, CSRF token or any Superset auth header: an
+#           # external server is not a party to the user's Superset session.
+#           # Read secrets from the environment rather than writing them here.
+#           "headers": {"Authorization": f"Bearer 
{os.environ['ACME_MCP_TOKEN']}"},
+#           # Per-call budget. Bounds how long one call may occupy the worker
+#           # running the turn. Defaults to 30.
+#           "timeout_seconds": 30,
+#           # Which of the server's tools to take. Absent or None means every
+#           # tool it offers, which lets the server decide what the agent can 
do.
+#           # Either the server's own name ("search_tables") or the namespaced
+#           # name Superset assigns ("mcp__acme_catalog__search_tables") 
matches.
+#           "tool_allowlist": ["search_tables"],
+#           # Refused regardless of the allowlist.
+#           "tool_denylist": [],
+#       },
+#   }
+#
+#   AI_AGENT_PROFILES = {
+#       "default": {"mcp_servers": ["acme_catalog"]},
+#   }
+#
+# Every tool from a server is namespaced "mcp__<server>__<tool>". The 
namespace is
+# stable, appears in stored conversation history, and is what makes it 
impossible
+# for a server offering "execute_sql" to displace Superset's own tool of that
+# name. Foreign results pass through the same AI_AGENT_MAX_RESULT_BYTES bound 
and
+# the same AI_AGENT_TOOL_POLICIES chain as built-in ones, and are wrapped as
+# untrusted content before the model sees them.
+#
+# A server that is unreachable, slow or unreadable contributes no tools and the
+# agent keeps working with the built-ins. Discovery happens while assembling 
the
+# registry for a turn, so a slow server costs up to its timeout at the start of
+# each turn that uses it.
+#
+# Requires the 'mcp' package; it is imported only once a server is configured.
+AI_AGENT_MCP_SERVERS: dict[str, Any] = {}
+
+# Refuse any external MCP tool whose name advertises SQL execution — anything
+# containing "execute_sql", "run_sql" or "query" by default. Enforced by
+# superset.ai.policy.ForeignToolPolicy.
+#
+# On by default because Superset's read-only enforcement and its per-datasource
+# authorization can only apply to SQL Superset itself runs. A third-party 
server
+# executing SQL goes through neither, so permitting it silently removes both
+# controls rather than merely widening the surface. Set this False only if you
+# have satisfied yourself that the servers you have configured enforce
+# equivalent controls of their own.
+AI_AGENT_MCP_DENY_FOREIGN_SQL = True
+
+# Conversation history sent to the model: the most recent N messages, further
+# trimmed oldest-first until under the character budget.
+AI_ASSISTANT_MAX_HISTORY_MESSAGES = 25
+AI_ASSISTANT_MAX_HISTORY_CHARS = 100_000
+
+# Timezone for the authoritative date given to the model, so it never has to
+# infer today's date or weekday.
+AI_ASSISTANT_TIMEZONE = "UTC"
+
+# Days a conversation is retained. Pruning is performed by the
+# ``ai.prune_conversations`` Celery task, which must be scheduled to run.

Review Comment:
   The new conversation tables never honor this 30-day retention setting 
because no `ai.prune_conversations` task or scheduler implementation exists in 
the change, so questions, SQL, and results persist indefinitely. Could the 
pruning task and scheduling instructions ship with this setting?



##########
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)

Review Comment:
   Transient 429/5xx/transport failures end the turn on the first attempt: both 
provider SDKs set `max_retries=0`, while these direct provider calls never 
invoke the new retry layer. Could the runtime wire in the configured retry 
policy, including safe semantics for a partially streamed attempt?



##########
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):

Review Comment:
   The advertised wall-clock budget is checked only between rounds, so one hung 
provider stream/completion or synchronous tool call can exceed it indefinitely 
and still finish as successful. Could the remaining deadline bound each 
provider and tool operation?



##########
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()

Review Comment:
   Replaying the same `request_id` reuses the two message rows but still 
creates a new run, overwrites the stored run context, and starts inference 
again; worker mode can therefore double-charge and race to overwrite one 
assistant row. Could a replay return the original run context without calling 
`_start_run` again?



##########
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,

Review Comment:
   A concrete model pin is recorded as the model used, but it is never copied 
into this `CompletionRequest`, so the provider uses the profile alias while the 
transcript and telemetry claim the pinned model. Could the exact model be 
propagated through `RunRequest` and validated by the provider?



-- 
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: notifications-unsubscr...@superset.apache.org

For queries about this service, please contact Infrastructure at:
us...@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscr...@superset.apache.org
For additional commands, e-mail: notifications-h...@superset.apache.org