I3eka commented on code in PR #43132: URL: https://github.com/apache/superset/pull/43132#discussion_r4384598758
########## superset/ai/policy.py: ########## @@ -0,0 +1,364 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +""" +Guards applied to every tool call before it runs. + +These bound blast radius. They are **not** an authorization layer: a tool that +returns or mutates a specific data-bearing object still has to perform its own +``security_manager.raise_for_access(...)`` check. A policy answers "should this +shape of call be attempted at all", which is a cheaper and coarser question. + +Policies are configured as dotted paths in ``AI_AGENT_TOOL_POLICIES`` so a +deployment can add its own without forking. +""" + +from __future__ import annotations + +import logging +import re +from abc import ABC, abstractmethod +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Any + +logger = logging.getLogger(__name__) + +#: A bare identifier, or dotted parts thereof. Deliberately strict: anything +#: with whitespace, quotes, semicolons or parentheses is rejected rather than +#: escaped, because a tool that needs to escape an identifier is a tool that is +#: building SQL by concatenation. +_SAFE_IDENTIFIER = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*(\.[A-Za-z_][A-Za-z0-9_]*)*$") + +#: Argument names understood to carry identifiers rather than free text. +_IDENTIFIER_ARGUMENTS = frozenset( + {"table", "table_name", "schema", "schema_name", "catalog", "column", "columns"} +) + + +@dataclass(frozen=True) +class Denial: + """ + A refusal to run a tool call. + + ``reason`` is shown to the model, so it should say what would be acceptable + instead. A model that is told "only read-only SQL is allowed" rewrites its + query; a model that is told "denied" retries the same thing. + """ + + reason: str + + +class ToolPolicy(ABC): + """A pre-execution guard over a single tool call.""" + + #: Identifies the policy in logs. + name: str = "policy" + + @abstractmethod + def check( + self, + tool_name: str, + arguments: dict[str, Any], + ) -> Denial | None: + """ + Inspect a pending call. + + Return ``None`` to allow, or a :class:`Denial` to block. A policy that + does not apply to ``tool_name`` returns ``None``. + """ + + +class ReadOnlySqlPolicy(ToolPolicy): + """ + Refuse anything that is not a read. + + Correctness here rests on Superset's own parser rather than a prefix or + keyword match. A regex over the leading token is defeated by a leading + comment, a CTE that wraps a DML statement, ``EXPLAIN ANALYZE DELETE``, and + multi-statement scripts — all of which the parser handles because the rest + of Superset already depends on it for the same decision. + """ + + name = "read_only_sql" + + #: Tools whose payload is SQL to execute. + sql_tools = frozenset( + { + "execute_sql", + "validate_sql", + "run_scoped_sql", + "create_virtual_dataset", Review Comment: Rechecked on `c510dfb724`: the policy recognizes the namespaced `create_virtual_dataset` tool and unwraps its `request` payload. The regression rejects `DELETE` through that exact MCP name; database-dialect tests also cover the nested path. This existing fix remains intact after the master merge. ########## superset/ai/orchestrator.py: ########## @@ -0,0 +1,647 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +""" +Runs one assistant turn end to end. + +Sits between the HTTP layer and the runtime: loads the conversation, assembles +the prompt, resolves the tools the chosen profile allows, drives the runtime, +publishes every event to the bus, and records the outcome on the assistant +message. + +Deliberately independent of *where* it runs. The same function body serves the +inline path and the Celery path, which is what makes the execution mode a +configuration choice rather than two implementations that drift apart. +""" + +from __future__ import annotations + +import asyncio +import logging +import uuid as uuid_module +from collections.abc import AsyncIterator, Iterator +from dataclasses import dataclass +from typing import Any + +from superset.ai.events import ( + cancelled_event, + done_event, + error_event, + GENERIC_ERROR_MESSAGE, + session_event, + StreamEvent, +) +from superset.ai.llm.base import Message +from superset.ai.telemetry import bind_run, current_run, start_run +from superset.ai.types import MessageRole, MessageStatus, RunOutcome, StreamEventType +from superset.utils.decorators import transaction + +logger = logging.getLogger(__name__) + +#: Cache key prefix for a run's cancellation flag. A flag rather than a signal +#: because a worker cannot be interrupted mid-call reliably; the runtime checks +#: this between steps. +_CANCEL_PREFIX = "ai-cancel-" + +#: How long a cancellation request stays meaningful. +_CANCEL_TTL_SECONDS = 900 + +#: Stored when a run is stopped before it produced any answer, so the +#: transcript still records that the turn happened. +_STOPPED_WITHOUT_ANSWER = "_Stopped before an answer was produced._" + +#: Stored when a run exhausted its time budget without saying anything. Phrased +#: as something the user can act on, because retrying is usually the right move. +_TIMED_OUT_WITHOUT_ANSWER = ( + "The assistant ran out of time before it could answer. Please try again." +) + +#: Ceiling on the page context recorded on a message. Well below the prompt's own +#: limit: this is stored per turn and read back with the whole transcript. +_RECORDED_CONTEXT_LIMIT = 4_000 + + +@dataclass +class TurnRequest: + """One unit of work: answer the latest message on a thread.""" + + thread_uuid: str + user_id: int + run_id: str + #: Assistant message row to fill in. Created before the run starts so a + #: client that reconnects has something to attach to. + assistant_message_uuid: str + profile_key: str | None = None + #: Concrete model to pin, overriding the profile's tier. + model: str | None = None + #: What the user had on screen when they asked. Supplied by the client, + #: which is the only party that knows which tab is open, what is typed in + #: the editor and which filters are applied. + page_context: dict[str, Any] | None = None + + def to_payload(self) -> dict[str, Any]: + """Serialise for the task broker.""" + return { + "thread_uuid": self.thread_uuid, + "user_id": self.user_id, + "run_id": self.run_id, + "assistant_message_uuid": self.assistant_message_uuid, + "profile_key": self.profile_key, + "model": self.model, + "page_context": self.page_context, + } + + @classmethod + def from_payload(cls, payload: dict[str, Any]) -> TurnRequest: + """Rebuild from a broker payload.""" + return cls(**payload) + + +def new_run_id() -> str: + """Identifier for one run, used as the event-stream key.""" + return str(uuid_module.uuid4()) + + +#: Runs cancelled in this process. +#: +#: Held alongside the cache rather than instead of it. Superset's default cache +#: is a null cache, which accepts a write and discards it — so a cache-only +#: implementation would leave cancellation silently broken on a default install, +#: with the button appearing to work and nothing stopping. This set makes inline +#: execution correct with no cache at all; the cache is what carries a +#: cancellation across processes for worker execution. +_CANCELLED_LOCALLY: set[str] = set() + + +def request_cancel(run_id: str) -> None: + """ + Ask a run to stop. + + Cooperative by design: the flag is recorded here and observed by the runtime + between steps. A run blocked inside a single long model call or query will + not notice until that call returns, which is a real limit worth documenting + rather than hiding. + """ + from superset.extensions import cache_manager + + _CANCELLED_LOCALLY.add(run_id) + try: + cache_manager.cache.set( Review Comment: Rechecked on `c510dfb724`: worker cancellation uses `get_event_bus_backend()` for both writing and reading the flag, including the configured Redis/Sentinel backend. It does not depend on the general `NullCache`. The shared-backend cancellation regression passes. ########## 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: Fixed here in `5a5010f26e` by carrying over #44561. The placeholder is created with its original run ID and accepted context, and a conditional database claim allows only one execution. Replays keep that handle; reconnects read the original result. The tests cover lost responses, ambiguous enqueue outcomes and competing database sessions. ########## superset-frontend/src/features/ai/hooks/useChatBot.ts: ########## @@ -0,0 +1,1323 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +/** + * @fileoverview Conversation state and the send loop. + * + * Runs are tracked per conversation, not globally. That is the point of the + * structure: a user can start something slow in one conversation, switch to + * another and keep working, and come back to find the first still going. A single + * `isLoading` flag would have made switching away cancel or corrupt the run. + * + * The server owns the transcript. A finished run is re-read from it rather than + * assembled from the frames, so the tool calls persisted on the message are what + * the user sees, and what they see survives a reload. + */ + +import { useCallback, useEffect, useRef, useState } from 'react'; +import type { TextAreaRef } from 'antd/es/input/TextArea'; +import { logging } from '@apache-superset/core/utils'; +import { t } from '@apache-superset/core/translation'; +import { + type AiAgent, + type AiToolCall, + type ChatMessageWithMeta, + type ChatTab, + type CheckpointPayload, +} from '../types'; +import { + AGENT_STORAGE_KEY, + ChatRequestAbortedError, + ChatStreamEventError, + ChatStreamTimeoutError, + DEFAULT_AGENT_KEY, + DEFAULT_CHAT_AGENT, + cancelChatRun, + describeRequestError, + fetchAgents, + fetchSuggestedPrompts, + loadStoredAgentKey, + normalizeChatAgents, + startRun, + streamRun, + submitFeedback, +} from './chatRequest'; +import { + NEW_CHAT_NAME, + createThread, + deleteThread as deleteThreadApi, + getThread, + listThreads, + threadToTab, + updateThread, +} from './chatThreadsApi'; +import { buildQuickPrompts } from './quickPrompts'; +import { + buildPageContextPayload, + usePageContext, + type PageContext, +} from './usePageContext'; + +/** Cache of the conversation list, so the menu renders before the list arrives. */ +export const CHAT_TABS_STORAGE_KEY = 'superset-chat-tabs'; + +/** Which conversation was last open. */ +export const ACTIVE_TAB_STORAGE_KEY = 'superset-chat-active-tab'; + +/** Recent inputs, recalled with the arrow keys. */ +export const HISTORY_STORAGE_KEY = 'superset-chat-history'; + +export { AGENT_STORAGE_KEY } from './chatRequest'; + +/** How many inputs the arrow-key history keeps. */ +const MAX_INPUT_HISTORY = 50; + +/** A conversation title derived from a message is clipped to this. */ +const MAX_TAB_NAME_LENGTH = 30; + +export type ChatRunStatus = 'running' | 'cancelling'; + +/** Shared empty list, so a render with no steps yet keeps a stable identity. */ +const EMPTY_TOOL_CALLS: AiToolCall[] = []; + +interface ActiveChatRun { + requestId: string; + tabId: string; + threadId: string; + runId?: string; + controller: AbortController; + isStreaming: boolean; + liveThoughts: string; + liveToolLog: string; + /** + * Steps taken so far, as structured records rather than log lines. + * + * Carried alongside `liveToolLog` so a run in flight can be rendered the same + * way a finished one is — expandable per step, with the SQL and the rows it + * returned — instead of as a wall of text that only becomes legible once the + * transcript is re-read from the server. + */ + liveToolCalls: AiToolCall[]; + /** The page context this run was given, so the live view can show it too. */ + livePageContext?: string; + /** + * The answer so far, as the model produces it. + * + * Rendered directly: the deltas used to be folded into `liveThinking`, which + * nothing displayed, so an answer appeared in one piece the moment the run + * ended however long it had taken to generate. + */ + liveAnswer: string; + liveThinking: string; + status: ChatRunStatus; + startedAt: number; + checkpoint: CheckpointPayload | null; +} + +/** + * An identifier for a turn. + * + * Drawn from `crypto`, not `Math.random`. These become the idempotency key on a + * turn and the handle used to cancel one, so a value another session could guess + * is a correctness and a security problem rather than merely a collision risk. + */ +const generateId = (): string => { + if (typeof crypto.randomUUID === 'function') { + return crypto.randomUUID(); + } + // Older engines expose the entropy source without the convenience wrapper. + const bytes = new Uint8Array(16); + crypto.getRandomValues(bytes); + return Array.from(bytes, byte => byte.toString(16).padStart(2, '0')).join(''); +}; + +const createNewTab = (name: string = NEW_CHAT_NAME): ChatTab => ({ + id: generateId(), + name, + messages: [], + createdAt: Date.now(), +}); + +const truncateTabName = ( + name: string, + maxLength: number = MAX_TAB_NAME_LENGTH, +): string => + name.length <= maxLength ? name : `${name.substring(0, maxLength)}...`; + +const readJson = <T>(key: string, fallback: T): T => { + try { + const stored = localStorage.getItem(key); + return stored ? (JSON.parse(stored) as T) : fallback; + } catch (caught) { + logging.warn(`[ai] could not read ${key}`, caught); + return fallback; + } +}; + +const writeJson = (key: string, value: unknown): void => { + try { + localStorage.setItem(key, JSON.stringify(value)); + } catch (caught) { + logging.warn(`[ai] could not write ${key}`, caught); + } +}; + +/** + * Reconciles the server's transcript with what is already on screen. + * + * The server's copy is authoritative — it carries the tool calls — but it is not + * necessarily complete the moment a run ends, and replacing outright would then + * erase an answer the user has just read. So anything local that the server has + * not accounted for is kept, matched by identity first and by role and content + * second, which is how a locally-appended turn is recognised once the server + * returns its own copy of it under a real uuid. + */ +export const mergeMessages = ( + fromServer: ChatMessageWithMeta[], + local: ChatMessageWithMeta[], +): ChatMessageWithMeta[] => { + const serverIds = new Set(fromServer.map(message => message.id)); + const serverTurns = new Set( + fromServer.map(message => `${message.role}:${message.content}`), + ); + const unaccounted = local.filter( + message => + !serverIds.has(message.id) && + !serverTurns.has(`${message.role}:${message.content}`), + ); + return [...fromServer, ...unaccounted]; +}; + +/** + * The `page_context` body for one turn. + * + * Returns undefined when there is nothing to send, so an omitted field is + * distinguishable from an empty one. + */ +export const buildRequestPageContext = ( + context: PageContext | undefined, + directive?: string, +): Record<string, unknown> | undefined => { + const payload = context ? buildPageContextPayload(context) : undefined; + if (!directive) { + return payload; + } + const existing = payload?.helper_directives; + return { + ...payload, + helper_directives: [ + directive, + ...(Array.isArray(existing) ? existing : []), + ], + }; +}; + +export interface UseChatBotReturn { + // Conversations + chatTabs: ChatTab[]; + activeTabId: string; + activeTab: ChatTab | undefined; + threadsLoaded: boolean; + handleNewChat: () => Promise<string>; + handleSelectTab: (tabId: string) => Promise<void>; + handleDeleteTab: (tabId: string) => Promise<void>; + handleRenameTab: (tabId: string, newName: string) => void; + // Messages of the active conversation + messages: ChatMessageWithMeta[]; + // Input + inputValue: string; + setInputValue: (value: string) => void; + handleKeyDown: (event: React.KeyboardEvent) => void; + inputRef: React.RefObject<TextAreaRef>; + messagesEndRef: React.RefObject<HTMLDivElement>; + // The run in flight, if any, for the active conversation + isLoading: boolean; + isStreamingResponse: boolean; + liveThoughts: string; + liveToolLog: string; + /** Steps taken so far in the run in flight, for the structured live view. */ + liveToolCalls: AiToolCall[]; + /** The page context the run in flight was given. */ + livePageContext?: string; + /** The answer so far for the run in flight. */ + liveAnswer: string; + checkpoint: CheckpointPayload | null; + activeRunStatus: ChatRunStatus | null; + error?: string; + // Actions + sendMessage: ( + messageOverride?: string, + systemPromptOverride?: string, + ) => Promise<void>; + handleCancelRun: () => Promise<void>; + handleCheckpointContinue: () => void; + handleFeedback: (messageId: string, feedback: 'like' | 'dislike') => void; + messageFeedback: Record<string, 'like' | 'dislike'>; + // Suggestions + /** The message whose run just ended; its thought process stays open. */ + justCompletedId?: string; + quickPrompts: string[]; + loadQuickPrompts: () => void; + applyQuickPrompt: (prompt: string) => Promise<void>; + // Agent profiles + agents: AiAgent[]; + selectedAgent: string; + setSelectedAgent: (key: string) => void; + // Page context + pageContext: PageContext; + includePageContext: boolean; + toggleIncludePageContext: () => void; +} + +export const useChatBot = (): UseChatBotReturn => { + const [chatTabs, setChatTabs] = useState<ChatTab[]>(() => + readJson<ChatTab[]>(CHAT_TABS_STORAGE_KEY, []).map(tab => ({ + // The cache is a placeholder for the menu; message bodies are re-read from + // the server so a stale cache cannot show a conversation that has moved on. + ...tab, + messages: [], + })), + ); + const [activeTabId, setActiveTabId] = useState<string>(() => { + try { + return localStorage.getItem(ACTIVE_TAB_STORAGE_KEY) ?? ''; + } catch { + return ''; + } + }); + const [threadsLoaded, setThreadsLoaded] = useState(false); + const [error, setError] = useState<string | undefined>(undefined); + + const [inputValue, setInputValue] = useState(''); + const [activeRunsByTab, setActiveRunsByTab] = useState< + Record<string, ActiveChatRun> + >({}); + const [quickPrompts, setQuickPrompts] = useState<string[]>([]); + const [messageFeedback, setMessageFeedback] = useState< + Record<string, 'like' | 'dislike'> + >({}); + const [includePageContext, setIncludePageContext] = useState(true); + /** + * The assistant message whose run has only just ended. + * + * Its thought process stays open, because collapsing it the instant the answer + * lands moves everything below it — the answer the user is mid-sentence through + * jumps up the panel. Older messages start closed. + */ + const [justCompletedId, setJustCompletedId] = useState<string | undefined>(); + const [agents, setAgents] = useState<AiAgent[]>([DEFAULT_CHAT_AGENT]); + const [selectedAgent, setSelectedAgent] = useState<string>(() => + loadStoredAgentKey(AGENT_STORAGE_KEY), + ); + + const [messageHistory, setMessageHistory] = useState<string[]>(() => + readJson<string[]>(HISTORY_STORAGE_KEY, []), + ); + const [historyIndex, setHistoryIndex] = useState(-1); + const [currentDraft, setCurrentDraft] = useState(''); + + const messagesEndRef = useRef<HTMLDivElement>(null); + const inputRef = useRef<TextAreaRef>(null); + + /** + * The run map and the conversation list are also held in refs, and the refs are + * the authority. + * + * The send loop has to ask "is this still my run?" between awaits, and it cannot + * ask React: a run that starts and fails inside one batch never causes a render, + * so a ref synced at render time would still be empty and the loop would discard + * its own result as stale. Writing the ref at the point of mutation removes that + * window. The callbacks read the refs rather than the state so their identities + * do not churn on every streamed frame, which would restart effects mid-run. + */ + const activeRunsByTabRef = useRef<Record<string, ActiveChatRun>>({}); + const chatTabsRef = useRef<ChatTab[]>(chatTabs); + const activeTabIdRef = useRef(activeTabId); + activeTabIdRef.current = activeTabId; + + const updateRuns = useCallback( + ( + updater: ( + previous: Record<string, ActiveChatRun>, + ) => Record<string, ActiveChatRun>, + ) => { + activeRunsByTabRef.current = updater(activeRunsByTabRef.current); + setActiveRunsByTab(activeRunsByTabRef.current); + }, + [], + ); + + const updateTabs = useCallback( + (updater: (previous: ChatTab[]) => ChatTab[]) => { + chatTabsRef.current = updater(chatTabsRef.current); + setChatTabs(chatTabsRef.current); + }, + [], + ); + + /** Resolved when the user answers a checkpoint; see `streamRun`. */ + const checkpointGateRef = useRef<{ resolve: () => void } | null>(null); Review Comment: Confirmed and fixed in `1a5d4bc056`. The checkpoint resolver belongs to the active run for its conversation. Two-tab regressions failed before the change: Continue/Cancel on the first tab resumed the second tab. Both now release only the selected tab and leave the other checkpoint waiting. ########## tests/unit_tests/ai/test_eventbus.py: ########## @@ -0,0 +1,623 @@ +# 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. +""" +Tests for the event bus that carries streamed events to the HTTP response. + +Every test here bounds its own wall clock: the failure mode these buses have is +hanging, so a test that could block forever would report the bug as a hung +suite rather than as a failure. +""" + +from __future__ import annotations + +import time +from collections.abc import Iterator +from typing import Any, TYPE_CHECKING + +import pytest +from flask import current_app +from pytest_mock import MockerFixture + +from superset.utils import json + +if TYPE_CHECKING: + from superset.ai.eventbus import BaseEventBus + from superset.ai.events import StreamEvent + +#: Small enough that the whole file stays quick, large enough that a poll or two +#: really does elapse. +POLL = 0.02 +TIMEOUT = 0.2 + + [email protected](autouse=True) +def _reset_memory_bus_singleton() -> Iterator[None]: + """ + Clear the process-wide memory bus around every test. + + ``get_event_bus`` deliberately caches one :class:`MemoryEventBus` per + process — that is what lets a streaming request find the queue an inline run + is writing to — which makes it shared state between tests. Resetting it on + both sides keeps ordering-dependent leakage out of the suite. + """ + from superset.ai import eventbus + + eventbus._MEMORY_BUS = None + yield + eventbus._MEMORY_BUS = None + + +def _drain(bus: BaseEventBus, run_id: str, **kwargs: Any) -> list[StreamEvent | None]: + """ + Consume a run to completion. + + That ``list()`` returns at all is part of what is being tested: a bus that + ignored a terminal event or its own deadline would never get here. + """ + kwargs.setdefault("timeout_seconds", TIMEOUT) + kwargs.setdefault("poll_seconds", POLL) + return list(bus.consume(run_id, **kwargs)) + + +def _types(events: list[StreamEvent | None]) -> list[str]: + """Event type names in order, with idle ticks shown as ``"idle"``.""" + return ["idle" if event is None else event.type.value for event in events] + + +def _delivered(events: list[StreamEvent | None]) -> list[StreamEvent]: + """Just the real events, for asserting on payloads.""" + return [event for event in events if event is not None] + + +def _decoded_id(entry: Any) -> str: + """Stream entry id as a string, whichever form the client handed back.""" + raw = entry[0] + return raw.decode() if isinstance(raw, bytes) else str(raw) + + +class FakeStreamCache: + """ + In-memory stand-in for the Redis-backed cache the bus writes through. + + Mirrors the contract :class:`RedisStreamEventBus` is written against: the + four positional arguments of + :class:`superset.async_events.cache_backend.RedisCacheBackend`, and an + inclusive ``start`` for ``xrange``, which is what the bus's + skip-the-entry-we-already-saw filter depends on. + + With ``binary=True`` it behaves like ``redis-py`` without response + decoding: entry ids, field names and field values all come back as bytes. + """ + + def __init__(self, binary: bool = False) -> None: + self.streams: dict[str, list[tuple[Any, dict[Any, Any]]]] = {} + self.expirations: list[tuple[str, int]] = [] + self.binary = binary + self._sequence = 0 + + def _next_id(self) -> Any: + self._sequence += 1 + # Zero padded so lexicographic order matches insertion order, the way + # real Redis stream ids do. + entry_id = f"{self._sequence:04d}-0" + return entry_id.encode() if self.binary else entry_id + + def append_raw(self, stream_name: str, fields: dict[Any, Any]) -> None: + """Put an entry on a stream without going through ``publish``.""" + self.streams.setdefault(stream_name, []).append((self._next_id(), fields)) + + def xadd( + self, + stream_name: str, + event_data: dict[str, Any], + event_id: str = "*", + maxlen: int | None = None, + ) -> str: + fields: dict[Any, Any] = dict(event_data) + if self.binary: + fields = { + key.encode(): str(value).encode() for key, value in event_data.items() + } + self.append_raw(stream_name, fields) + return _decoded_id(self.streams[stream_name][-1]) + + def xrange( + self, + stream_name: str, + start: str = "-", + end: str = "+", + count: int | None = None, + ) -> list[Any]: + entries = self.streams.get(stream_name, []) + if start != "-": + entries = [e for e in entries if _decoded_id(e) >= start] + return list(entries[: count or len(entries)]) + + def expire(self, stream_name: str, ttl_seconds: int) -> None: + self.expirations.append((stream_name, ttl_seconds)) + + +class BrokenWriteCache(FakeStreamCache): + """A cache whose writes fail, as a Redis outage would make them.""" + + def xadd(self, *args: Any, **kwargs: Any) -> str: + raise RuntimeError("connection refused") + + +class FlakyReadCache(FakeStreamCache): + """A cache whose reads fail a set number of times before recovering.""" + + def __init__(self, failures: int) -> None: + super().__init__() + self.read_attempts = 0 + self._remaining_failures = failures + + def xrange(self, *args: Any, **kwargs: Any) -> list[Any]: + self.read_attempts += 1 + if self._remaining_failures > 0: + self._remaining_failures -= 1 + raise RuntimeError("connection reset") + return super().xrange(*args, **kwargs) + + +class BrokenExpireCache(FakeStreamCache): + """A cache that cannot set a TTL.""" + + def expire(self, stream_name: str, ttl_seconds: int) -> None: + raise RuntimeError("connection refused") + + +def _terminal_event(index: int) -> StreamEvent: + """One event per terminal type, addressed by index so ids stay readable.""" + from superset.ai.events import cancelled_event, done_event, error_event + + return [done_event(True), error_event("boom"), cancelled_event()][index] + + +# --------------------------------------------------------------------------- # +# MemoryEventBus +# --------------------------------------------------------------------------- # + + +def test_memory_bus_yields_published_events_in_order() -> None: + """Events come back in the order the run produced them, and only once.""" + from superset.ai.eventbus import MemoryEventBus + from superset.ai.events import assistant_delta_event, done_event, final_event + + bus = MemoryEventBus() + bus.publish("run-1", assistant_delta_event("Total ")) + bus.publish("run-1", assistant_delta_event("revenue")) + bus.publish("run-1", final_event("Total revenue")) + bus.publish("run-1", done_event(True)) + + events = _drain(bus, "run-1") + + assert _types(events) == [ + "assistant_delta", + "assistant_delta", + "final", + "done", + ] + delivered = _delivered(events) + assert delivered[0].payload == {"delta": "Total "} + assert delivered[2].payload == {"role": "assistant", "content": "Total revenue"} + assert delivered[3].payload == {"ok": True} + + +def test_memory_bus_keeps_runs_apart() -> None: + """Two concurrent runs in one process do not read each other's events.""" + from superset.ai.eventbus import MemoryEventBus + from superset.ai.events import done_event, final_event + + bus = MemoryEventBus() + bus.publish("run-1", final_event("one")) + bus.publish("run-2", final_event("two")) + bus.publish("run-1", done_event(True)) + bus.publish("run-2", done_event(True)) + + first = _delivered(_drain(bus, "run-1")) + second = _delivered(_drain(bus, "run-2")) + + assert first[0].payload["content"] == "one" + assert second[0].payload["content"] == "two" + + [email protected]("index", [0, 1, 2]) +def test_memory_bus_stops_at_a_terminal_event(index: int) -> None: + """ + Consumption ends on ``done``, ``error`` or ``cancelled``. + + Without this the reader waits out the full timeout after a run has already + finished, which the user experiences as a response that never closes. + """ + from superset.ai.eventbus import MemoryEventBus + from superset.ai.events import assistant_delta_event + + terminal = _terminal_event(index) + + bus = MemoryEventBus() + bus.publish("run-1", terminal) + # Anything queued after the terminal event is not the client's business. + bus.publish("run-1", assistant_delta_event("trailing")) + + events = _drain(bus, "run-1", timeout_seconds=5.0) + + assert events == [terminal] + + +def test_memory_bus_yields_idle_while_nothing_is_published() -> None: + """ + A quiet run produces keep-alive cues rather than silence. + + ``IDLE`` is the caller's signal to write an SSE comment, which is what stops + a proxy from closing a connection during a long tool call. + """ + from superset.ai.eventbus import IDLE, MemoryEventBus + + bus = MemoryEventBus() + events = _drain(bus, "quiet-run") + + assert events, "a poll interval passing must produce something" + assert all(event is IDLE for event in events) + + +def test_memory_bus_interleaves_idle_with_late_events() -> None: + """An event published after an idle period is still delivered.""" + from superset.ai.eventbus import IDLE, MemoryEventBus + from superset.ai.events import done_event + + bus = MemoryEventBus() + stream = bus.consume("run-1", timeout_seconds=5.0, poll_seconds=POLL) + + assert next(stream) is IDLE + + bus.publish("run-1", done_event(True)) + assert next(stream) == done_event(True) + + with pytest.raises(StopIteration): + next(stream) + + +def test_memory_bus_consume_returns_when_the_timeout_expires() -> None: + """ + A run that never terminates gives up on its own deadline. + + A producer can die without emitting anything terminal; the reader has to + return so the request completes instead of holding a worker forever. + """ + from superset.ai.eventbus import IDLE, MemoryEventBus + from superset.ai.events import assistant_delta_event + + bus = MemoryEventBus() + bus.publish("run-1", assistant_delta_event("partial")) + + started = time.monotonic() + events = _drain(bus, "run-1", timeout_seconds=0.1, poll_seconds=0.02) + elapsed = time.monotonic() - started + + assert events[0] == assistant_delta_event("partial") + assert events[-1] is IDLE + assert elapsed < 2.0, "consume must honour its deadline, not block" + + +def test_memory_bus_close_releases_the_queue() -> None: + """ + Closing a run drops its queue, so a finished run holds no memory. + + Anything still queued goes with it; ``close`` is called once the response + has been written, at which point nobody is left to read it. + """ + from superset.ai.eventbus import IDLE, MemoryEventBus + from superset.ai.events import final_event + + bus = MemoryEventBus() + bus.publish("run-1", final_event("never read")) + bus.close("run-1") + + assert all(event is IDLE for event in _drain(bus, "run-1")) + + # Closing an unknown or already-closed run is not an error: the streaming + # path closes in a ``finally`` that may run twice. + bus.close("run-1") + bus.close("never-existed") + + +# --------------------------------------------------------------------------- # +# RedisStreamEventBus +# --------------------------------------------------------------------------- # + + +def test_redis_bus_round_trips_an_event_through_the_stream() -> None: + """ + An event survives encoding to the stream and decoding back out. + + This is the wire format two processes agree on, so the type and the payload + both have to come back intact rather than merely close. + """ + from superset.ai.eventbus import RedisStreamEventBus + from superset.ai.events import done_event, thinking_event + from superset.ai.types import ProgressStage + + cache = FakeStreamCache() + bus = RedisStreamEventBus(cache=cache, prefix="ai-events-") + + published = thinking_event(ProgressStage.TOOL, "Reading the schema", {"n": 1}) + bus.publish("run-1", published) + bus.publish("run-1", done_event(False)) + + events = _drain(bus, "run-1") + + assert events == [published, done_event(False)] + assert _delivered(events)[0].payload == { + "stage": "tool", + "message": "Reading the schema", + "meta": {"n": 1}, + } + + # Stored under the prefixed key, as one JSON document per entry. + stored = cache.streams["ai-events-run-1"] + assert json.loads(stored[0][1]["data"])["type"] == "thinking" + + +def test_redis_bus_handles_byte_entry_ids_and_fields() -> None: + """ + A client that does not decode responses is handled without duplication. + + ``redis-py`` hands back bytes for ids and fields alike. If the bus compared + a bytes id against the string it tracks, every poll would look like fresh + data and the browser would see the whole stream again on each pass. + """ + from superset.ai.eventbus import RedisStreamEventBus + from superset.ai.events import assistant_delta_event, done_event + + cache = FakeStreamCache(binary=True) + bus = RedisStreamEventBus(cache=cache) + + bus.publish("run-1", assistant_delta_event("a")) + bus.publish("run-1", assistant_delta_event("b")) + bus.publish("run-1", done_event(True)) + + events = _drain(bus, "run-1") + + assert events == [ + assistant_delta_event("a"), + assistant_delta_event("b"), + done_event(True), + ] + + +def test_redis_bus_skips_malformed_entries() -> None: + """ + An entry we cannot read is dropped, not raised. + + A stream is shared state that outlives the process that wrote it, so one + unreadable entry — a truncated write, an event type from a newer version — + must not end a stream that still has good events in it. + """ + from superset.ai.eventbus import RedisStreamEventBus + from superset.ai.events import done_event, final_event + + cache = FakeStreamCache() + bus = RedisStreamEventBus(cache=cache) + stream = "ai-events-run-1" + + cache.append_raw(stream, {"unexpected": "no data field at all"}) + cache.append_raw(stream, {"data": "{not json"}) + cache.append_raw(stream, {"data": json.dumps({"payload": {}})}) + cache.append_raw( + stream, {"data": json.dumps({"type": "from_the_future", "payload": {}})} + ) + bus.publish("run-1", final_event("the good one")) + bus.publish("run-1", done_event(True)) + + events = _drain(bus, "run-1") + + assert events == [final_event("the good one"), done_event(True)] + + [email protected]("index", [0, 1, 2]) +def test_redis_bus_stops_at_a_terminal_event(index: int) -> None: + """Consumption ends on the first terminal event, as in-process it does.""" + from superset.ai.eventbus import RedisStreamEventBus + from superset.ai.events import assistant_delta_event + + terminal = _terminal_event(index) + + bus = RedisStreamEventBus(cache=FakeStreamCache()) + bus.publish("run-1", terminal) + bus.publish("run-1", assistant_delta_event("trailing")) + + events = _drain(bus, "run-1", timeout_seconds=5.0) + + assert events == [terminal] + + +def test_redis_bus_publish_survives_a_broken_backend() -> None: + """ + A failed publish must not kill the run producing the answer. + + The run's real output is persisted by the time it finishes; losing the live + stream costs the user progress updates, whereas an exception here would + cost them the answer. + """ + from superset.ai.eventbus import RedisStreamEventBus + from superset.ai.events import done_event, final_event + + cache = BrokenWriteCache() + bus = RedisStreamEventBus(cache=cache) + + bus.publish("run-1", final_event("an answer nobody streams")) + bus.publish("run-1", done_event(True)) + + assert cache.streams == {} + + +def test_redis_bus_read_failure_yields_idle_and_keeps_trying() -> None: + """ + A transient read failure is an idle tick, not the end of the stream. + + A Redis blip mid-run would otherwise truncate a response that was about to + succeed; instead the reader keeps the connection alive and picks the events + up when the backend comes back. + """ + from superset.ai.eventbus import RedisStreamEventBus + from superset.ai.events import done_event + + cache = FlakyReadCache(failures=2) + bus = RedisStreamEventBus(cache=cache) + bus.publish("run-1", done_event(True)) + + events = _drain(bus, "run-1", timeout_seconds=5.0) + + assert _types(events) == ["idle", "idle", "done"] + assert cache.read_attempts == 3 + + +def test_redis_bus_survives_a_permanently_broken_backend() -> None: + """A backend that never recovers times the reader out rather than raising.""" + from superset.ai.eventbus import IDLE, RedisStreamEventBus + + cache = FlakyReadCache(failures=10_000) + bus = RedisStreamEventBus(cache=cache) + + events = _drain(bus, "run-1") + + assert events, "must keep polling rather than return immediately" + assert all(event is IDLE for event in events) + assert cache.read_attempts >= 2 + + +def test_redis_bus_close_expires_rather_than_deletes() -> None: + """ + Closing sets a TTL and leaves the entries in place. + + A browser that reconnects late replays the stream from the beginning, so + deleting on close would cut off exactly the reader this bus exists for. + """ + from superset.ai.eventbus import RedisStreamEventBus + from superset.ai.events import done_event + + cache = FakeStreamCache() + bus = RedisStreamEventBus(cache=cache, prefix="prefix-", ttl_seconds=42) + bus.publish("run-1", done_event(True)) + bus.close("run-1") + + assert cache.expirations == [("prefix-run-1", 42)] + assert len(cache.streams["prefix-run-1"]) == 1 + + +def test_redis_bus_close_survives_a_broken_backend() -> None: + """Failing to set a TTL is logged, not raised at the end of a request.""" + from superset.ai.eventbus import RedisStreamEventBus + + bus = RedisStreamEventBus(cache=BrokenExpireCache()) + bus.close("run-1") + + +# --------------------------------------------------------------------------- # +# get_event_bus +# --------------------------------------------------------------------------- # + + +def test_get_event_bus_returns_one_shared_memory_bus_for_inline_runs( + mocker: MockerFixture, +) -> None: + """ + Inline execution gets the in-process bus, and the same one every time. + + Sharing is load-bearing rather than an optimisation: the streaming request + finds the inline run's queue only because both resolve to one object. + """ + from superset.ai.eventbus import get_event_bus, MemoryEventBus + + mocker.patch.dict( + current_app.config, + { + "AI_ASSISTANT_EXECUTION_MODE": "inline", + "AI_ASSISTANT_EVENT_BUS": "memory", + }, + ) + + bus = get_event_bus() + assert isinstance(bus, MemoryEventBus) + assert get_event_bus() is bus + + +def test_get_event_bus_returns_a_redis_bus_configured_from_config( + mocker: MockerFixture, +) -> None: + """ + The Redis bus takes its key prefix and TTL from configuration. + + Also the only test that reaches the branch at all: the backend is built from + ``AI_ASSISTANT_EVENT_BUS_CACHE_CONFIG`` rather than borrowed from the + general-purpose cache, which has no stream commands. Getting that wrong + would break every ``AI_ASSISTANT_EVENT_BUS='redis'`` deployment on its first + published event while leaving the rest of this module green. + """ + from superset.ai import eventbus as eventbus_module + from superset.ai.eventbus import get_event_bus, RedisStreamEventBus + from superset.ai.events import done_event + + cache = FakeStreamCache() + mocker.patch.object(eventbus_module, "_stream_backend", return_value=cache) Review Comment: Rechecked the current test and runtime together. The test constructs `get_event_bus()` with the real coordination backend classes and stubs only stream I/O, so it cannot hide a missing factory/import. Redis and Sentinel constructor cases pass on `c510dfb724`; the nonexistent-helper patch is gone. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
