I3eka commented on code in PR #43133: URL: https://github.com/apache/superset/pull/43133#discussion_r3890805408
########## 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)) + # Recorded as well as prompted with, so the transcript can show what the + # assistant was told about the user's screen. An answer that looks wrong is + # usually an answer to a different question than the reader assumed, and the + # page context is where that difference lives. + rendered_context = render_page_context(request.page_context) + state["page_context"] = rendered_context + system_prompt = _build_system_prompt(tools, rendered_context) + model = _resolved_model(provider, request.model, profile) + + recorder.describe( + agent_key=profile.key, + model=model, + question=_latest_question(history), + ) + + run_request = RunRequest( + messages=history, + system_prompt=system_prompt, + tools=tools, + policies=load_policy_chain(), + model_alias=profile.model_alias, + max_turns=profile.max_turns or _config("AI_AGENT_MAX_TURNS", 20), + timeout_seconds=profile.timeout_seconds + or _config("AI_AGENT_TIMEOUT_SECONDS", 300), + should_cancel=lambda: is_cancelled(request.run_id), + ) + + _mark_streaming(request.assistant_message_uuid) + + # The runtime is async and this is a synchronous generator, so the async + # events are drained into a list per batch rather than bridged with a + # thread. Collecting the whole run before yielding would defeat streaming, + # so the loop pulls one event at a time from a dedicated event loop. + yield from _drain(runtime.run(run_request)) + + result = runtime.result + outcome = _outcome_of(result) + if result.error is not None: + # The only place the provider's own words are recorded. They do not go on + # the message: that is served back to the browser, and a transport error + # can name internal hosts. + logger.warning( + "AI run %s failed: %s", + request.run_id, + result.error, + ) + _finalise_message( + request.assistant_message_uuid, + content=_terminal_content(result, outcome), + status=_status_of(outcome), + extra={ Review Comment: Agreed. Successful-turn thought persistence is inherited unchanged from the parent Native AI PR #42805; this authoring PR does not modify the finalization path. The parent should persist the same bounded field on every outcome, then this branch can be rebased. I am leaving the thread open pending that change. ########## docs/admin_docs/configuration/ai-assistant.mdx: ########## @@ -0,0 +1,508 @@ +--- +title: AI Assistant +hide_title: true +sidebar_position: 17 +version: 1 +--- + +# AI Assistant + +The AI Assistant is a conversational interface for exploring your data. A user +asks a question in plain language; the assistant finds relevant datasets, +inspects their schema, writes and runs read-only SQL, and answers with both the +result and the query it used. The shipped profiles are read-only. A deployment +can explicitly add chart and dashboard authoring tools to a gated profile. + +Superset ships **no model provider and talks to no model vendor by default**. +The feature is disabled, and even when enabled it returns `404` until you point +it at a provider you control. Nothing is sent anywhere until you configure it. + +## Enabling it + +Two things are required: the feature flag, and a provider. + +```python +# superset_config.py +FEATURE_FLAGS = { + "AI_ASSISTANT": True, +} + +AI_LLM_PROVIDER_CLASS = "superset.ai.llm.anthropic.AnthropicProvider" +AI_LLM_PROVIDER_CONFIG = { + "api_key": os.environ["ANTHROPIC_API_KEY"], + "models": { + "default": "claude-sonnet-4-5", + "fast": "claude-haiku-4-5", + "reasoning": "claude-opus-4-1", + }, +} +``` + +Install the matching extra: + +```bash +pip install "apache-superset[ai-anthropic]" # or [ai-openai] +``` + +Then run `superset init` so the assistant's permissions are created and assigned Review Comment: Agreed. The reviewed enablement sequence is inherited unchanged from #42805. This authoring PR adds tool-specific documentation but does not own the base migration runbook. The parent docs should add superset db upgrade once, after which this branch will be rebased. I am leaving this open until then. ########## docs/admin_docs/configuration/ai-assistant.mdx: ########## @@ -0,0 +1,508 @@ +--- +title: AI Assistant +hide_title: true +sidebar_position: 17 +version: 1 +--- + +# AI Assistant + +The AI Assistant is a conversational interface for exploring your data. A user +asks a question in plain language; the assistant finds relevant datasets, +inspects their schema, writes and runs read-only SQL, and answers with both the +result and the query it used. The shipped profiles are read-only. A deployment +can explicitly add chart and dashboard authoring tools to a gated profile. + +Superset ships **no model provider and talks to no model vendor by default**. +The feature is disabled, and even when enabled it returns `404` until you point +it at a provider you control. Nothing is sent anywhere until you configure it. + +## Enabling it + +Two things are required: the feature flag, and a provider. + +```python +# superset_config.py +FEATURE_FLAGS = { + "AI_ASSISTANT": True, +} + +AI_LLM_PROVIDER_CLASS = "superset.ai.llm.anthropic.AnthropicProvider" +AI_LLM_PROVIDER_CONFIG = { + "api_key": os.environ["ANTHROPIC_API_KEY"], + "models": { + "default": "claude-sonnet-4-5", + "fast": "claude-haiku-4-5", + "reasoning": "claude-opus-4-1", + }, +} +``` + +Install the matching extra: + +```bash +pip install "apache-superset[ai-anthropic]" # or [ai-openai] +``` + +Then run `superset init` so the assistant's permissions are created and assigned +to roles. Without this the endpoints return `403`. + +Conversations are stored in Superset's metadata database, so no extra +infrastructure is needed for the default configuration. + +### Which roles get access + +`superset init` grants `can_read`/`can_write` on `AIAssistant` to **Admin** and +**Alpha** only. "Write" here means writing one's own conversation. The shipped +profiles remain read-only; an operator who enables asset-authoring tools must +also gate that profile, and each tool enforces the current user's normal asset +and dataset permissions. + +**Gamma does not get it by default.** The assistant runs queries and costs +money per question, so it is granted deliberately rather than inherited. To +give it to Gamma users, add `can_read`/`can_write` on `AIAssistant` to Gamma or +to a custom role. + +Every query the assistant runs is subject to the *user's own* database and +dataset permissions. It cannot read anything the person chatting with it could +not read themselves. + +Because it is not in Gamma, it is also not inherited by the Public role when +`PUBLIC_ROLE_LIKE = "Gamma"` — an anonymous visitor cannot reach the assistant +unless you grant it explicitly. + +## Choosing a provider + +`AI_LLM_PROVIDER_CLASS` is a dotted path to a +`superset.ai.llm.base.BaseLLMProvider` subclass. Two are bundled: + +| Class | Use for | +| --- | --- | +| `superset.ai.llm.anthropic.AnthropicProvider` | The Anthropic Messages API | +| `superset.ai.llm.openai_compatible.OpenAICompatibleProvider` | OpenAI, and anything exposing an OpenAI-compatible endpoint — vLLM, Ollama, a private gateway | + +`AI_LLM_PROVIDER_CONFIG` is passed to the provider's constructor and its +contents are provider-defined. For the OpenAI-compatible provider, `base_url` +points it anywhere: + +```python +AI_LLM_PROVIDER_CLASS = "superset.ai.llm.openai_compatible.OpenAICompatibleProvider" +AI_LLM_PROVIDER_CONFIG = { + "base_url": "https://llm.internal.example.com/v1";, + "api_key": os.environ["MY_GATEWAY_KEY"], + "models": {"default": "our-hosted-model"}, +} +``` + +Everything vendor-specific — URLs, authentication, model naming — lives in the +provider. Superset core contains none of it, so a self-hosted model or a private +gateway needs configuration rather than a fork. + +### Model tiers and selection + +Profiles and prompts refer to capability *tiers* (`default`, `fast`, +`reasoning`), never to a vendor's model names. The provider maps tiers to +concrete models via the `models` dict. A tier you do not configure is an error +when requested, never a silent substitution — so cost and answer quality stay +attributable to the model actually used. + +Users may also pin a specific model per turn. Only models present in your +`models` mapping are accepted; anything else is rejected. + +## Agent profiles + +A profile bundles the decisions that differ between a quick answer and a careful +investigation: which tools are available, which model tier, and how many steps. +Two ship by default — `default` and `analyst`. + +**Which tools a model may invoke is a decision each deployment makes**, so +profiles are fully configurable. `AI_AGENT_PROFILES` maps a profile key to the +fields you want to override, leaving the rest alone: + +```python +AI_AGENT_PROFILES = { + # Let the assistant search and inspect, but never run SQL. + "default": {"tools": ["search_assets", "list_databases", "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", + "description": "Slow, thorough, multi-step.", + "tools": ["search_assets", "get_schema", "execute_sql"], + "required_permission": ("can_write", "AIAssistant"), + }, + + # Opt-in authoring. The tools still enforce normal asset and dataset RBAC. + "builder": { + "name": "Dashboard builder", + "tools": [ + "search_assets", + "get_schema", + "create_virtual_dataset", + "generate_chart", + "generate_dashboard", + ], + "required_permission": ("can_write", "Dashboard"), + }, +} +``` + +A tool name that does not exist is an error naming the typo and listing the +valid names, rather than an assistant that quietly lacks a capability. An empty +`tools` list is valid and means conversation with no data access. + +`required_permission` is enforced on both the listing *and* the run path, so a +profile a user cannot see is also one they cannot invoke by posting its key. + +### Available tools + +| Tool | What it does | +| --- | --- | +| `search_assets` | Finds datasets, charts and dashboards the user can see | +| `list_databases` | Lists database connections exposed to SQL Lab | +| `get_schema` | Lists schemas, tables and columns | +| `execute_sql` | Runs a **read-only** query | +| `validate_sql` | Checks a query without running it | +| `get_chart_context` | Reads a chart's definition | +| `get_dashboard_context` | Reads a dashboard's definition | +| `create_virtual_dataset` | Saves a read-only SQL query as a chartable dataset; opt-in only | +| `generate_chart` | Previews or saves a native chart; opt-in only | +| `generate_dashboard` | Creates a dashboard from saved chart IDs; opt-in only | + +## Customising the prompt + +Three levers, in increasing order of bluntness. + +**Add to it.** `AI_EXTRA_PROMPT_SECTIONS` appends your own sections. This is +where deployment-specific knowledge belongs — your table conventions, your +warehouse's dialect quirks, how your business defines a metric. The shipped +prompt is deliberately generic and mentions no particular database engine. + +**Remove from it.** `AI_DISABLED_PROMPT_SECTIONS` drops a shipped section by +key, for when you disagree with one. The safety section cannot be disabled. + +**Replace it.** `AI_SYSTEM_PROMPT` substitutes the whole thing. + +:::warning +Setting `AI_SYSTEM_PROMPT` discards the shipped safety and prompt-injection +rules along with everything else. Your deployment then owns them. +::: + +`AI_SYSTEM_PROMPT_MUTATOR` is a last-mile callable applied after assembly, +mirroring `SQL_QUERY_MUTATOR`. + +## Where turns execute + +`AI_ASSISTANT_EXECUTION_MODE` decides where the work happens. + +**`"inline"`** (default) runs the turn in the web process. Nothing extra to +deploy. + +**`"worker"`** hands it to Celery. Web workers stay free, and a browser that +loses its connection can rejoin a run in progress. It requires Celery and a +Redis event bus: + +```python +AI_ASSISTANT_EXECUTION_MODE = "worker" +AI_ASSISTANT_EVENT_BUS = "redis" +AI_ASSISTANT_EVENT_BUS_CACHE_CONFIG = { + "CACHE_TYPE": "RedisCache", + "CACHE_REDIS_HOST": "redis", + "CACHE_REDIS_PORT": 6379, + "CACHE_REDIS_DB": 0, +} + +class CeleryConfig: + imports = ( + # ... your existing imports ... + "superset.ai.tasks", + ) +``` + +Streams need Redis commands the general-purpose cache client does not expose, +which is why the bus is configured separately rather than reusing `CACHE_CONFIG`. + +Selecting `"worker"` with the in-memory event bus raises rather than leaving +every stream silently empty, and so does selecting the Redis bus without a +usable connection. + +A turn is deliberately **not** retried after a worker crash: inference costs +money, and re-running a turn the user may already have partly seen would charge +twice. The message records that it failed and the user can ask again. + +## Safety and limits + +Guards are applied before any tool runs, configured via +`AI_AGENT_TOOL_POLICIES`: + +- **Read-only SQL.** Enforced using Superset's own SQL parser, not pattern + matching — so a write hidden behind a comment, a CTE, a second statement, or + an unparseable construct is refused. `EXPLAIN`, `SHOW` and `DESCRIBE` are + permitted; everything the parser cannot vouch for is not. +- **Identifier safety.** Table and column names are resolved against metadata + the user may see rather than interpolated into SQL. + +These bound blast radius; they do not replace authorization. Every tool that +touches a data-bearing object performs the same permission check the REST API +does. + +Result sizes are capped by `AI_AGENT_MAX_RESULT_ROWS` and +`AI_AGENT_MAX_RESULT_BYTES`, and truncation is reported rather than hidden. Turn +length is bounded by `AI_AGENT_MAX_TURNS` and `AI_AGENT_TIMEOUT_SECONDS`; a run +that exhausts either answers with what it has. + +Content that arrives from your warehouse or asset metadata — table comments, +chart titles, column labels — is marked as untrusted in the prompt, because a +value in a database is data and not an instruction. + +### Cancellation + +Cancellation is cooperative: a run stops at its next step boundary. A run inside +a single long model call or a single long query will not stop until that call +returns. + +## Monitoring and tracing + +Superset bundles **no integration with any AI monitoring product**. Instead it +exposes a small sink interface, `AITelemetry`, and calls it once per run, once +per model round trip and once per tool call. Whatever you already use — +Braintrust, LangSmith, Langfuse, Arize Phoenix, an OpenTelemetry collector, a +self-hosted alternative, or a table in your own warehouse — you connect by +implementing that interface and listing it in `AI_TELEMETRY`. + +Entries are instances or dotted paths, exactly as for `EVENT_LOGGER` and +`STATS_LOGGER`. Two sinks ship in-tree and depend on nothing external: + +```python +# superset_config.py +import logging + +from superset.ai.telemetry import LoggingAITelemetry, StatsLoggerAITelemetry + +AI_TELEMETRY = [ + # One structured line per span, at the level you choose. + LoggingAITelemetry(level=logging.INFO), + # Counters and timings through your configured STATS_LOGGER. + StatsLoggerAITelemetry(), +] +``` + +`StatsLoggerAITelemetry` emits under a `superset.ai.` prefix: `run.start`, +`run.end`, `run.outcome.<outcome>`, `run.duration_ms`, `run.turns`, +`run.tokens.input`, `run.tokens.output`, `model_call`, +`model_call.duration_ms`, `model_call.error`, `error`, and per tool +`tool_call.<tool>`, `tool_call.<tool>.duration_ms`, `tool_call.<tool>.error` +and `tool_call.<tool>.truncated`. User, run and thread identifiers deliberately +never appear in a metric name — a metric per user is how a metrics backend gets +brought down. That detail belongs in a trace, which is what a custom sink is +for. + +### The content trade-off + +`AI_TELEMETRY_REDACT_CONTENT` defaults to `True`, and telemetry then carries +**structure and measurements only**: durations, token counts, model names, tool +names, outcomes, error classes, and the run, thread and user identifiers. No +question, no answer, no SQL, no row of data. Redaction is applied where the +trace is built, so a sink cannot receive content by accident even if it looks +for it. + +Setting it to `False` is what makes a trace genuinely useful for debugging +answer quality — you can read the prompt that produced a wrong answer and the +statement it ran. It also means the text of business questions and values from +your warehouse leave Superset for whichever service your sinks talk to. In many +organisations that is a decision for someone other than the person editing the +config file. `AI_TELEMETRY_MAX_CONTENT_CHARS` (default 10,000) caps any single +content field so one large result cannot dominate a payload. + +### A custom sink + +Every method has a no-op default, so implement only the ones you need — a sink +that only wants token counts overrides `on_model_call` and nothing else. + +```python +from superset.ai.telemetry import AITelemetry, ModelCallTrace, RunTrace + + +class TracingServiceTelemetry(AITelemetry): + """Forwards runs to an external tracing service.""" + + def __init__(self, client): + self._client = client + + def on_run_start(self, run: RunTrace) -> None: + self._client.start_span(run.run_id, name="superset.ai.run", attributes={ + "thread": run.thread_uuid, + "user": run.user_id, + }) + + def on_model_call(self, run: RunTrace, call: ModelCallTrace) -> None: + self._client.event(run.run_id, "model_call", { + "turn": call.turn, + "model": call.model, + "input_tokens": call.input_tokens, + "output_tokens": call.output_tokens, + # None unless you have turned redaction off. + "prompt": call.system_prompt, + }) + + def on_run_end(self, run: RunTrace) -> None: + self._client.end_span(run.run_id, status=str(run.outcome), attributes={ + "duration_ms": run.duration_ms, + "turns": run.turns, + "usage": run.usage, + }) + + +AI_TELEMETRY = [TracingServiceTelemetry(client=my_tracing_client)] +``` + +Three things to know before you write one: + +- **Sinks are called on the thread answering the user.** Anything that makes a + network call should hand off to a queue or a background thread; otherwise a + slow monitoring backend becomes slow answers. +- **A sink that raises cannot break a run.** Failures are logged once and + ignored, and the other configured sinks still receive everything. The same + applies to a dotted path that will not import: it is skipped with a warning + rather than taking the assistant down, because a missing observer loses the + record of a run and not the run itself. +- **`agent_key`, `model` and `question` are resolved after the run starts**, so + a `RunTrace` passed to `on_run_start` may carry less than the one passed to + the later hooks. Read those on `on_run_end`. + +## Connecting your own MCP servers + +The assistant's built-in tools cover Superset itself. To let it reach anything +else — your data catalog, a metrics service, a ticketing system — attach an +[MCP](https://modelcontextprotocol.io) server. Superset bundles no third-party +integration and connects to nothing by default; you name the servers. + +```bash +pip install "apache-superset[ai-mcp]" +``` + +```python +AI_AGENT_MCP_SERVERS = { + "acme_catalog": { + "url": "https://mcp.acme.internal/mcp";, + "transport": "streamable_http", # or "sse" + "headers": {"Authorization": f"Bearer {os.environ['ACME_MCP_TOKEN']}"}, + "timeout_seconds": 30, + "tool_allowlist": ["search_tables"], # omit to offer every tool + }, +} + +# Then let a profile use it. +AI_AGENT_PROFILES = { + "default": {"mcp_servers": ["acme_catalog"]}, +} +``` + +Its tools appear to the model as `mcp__acme_catalog__search_tables`. The +namespace means a foreign tool can never shadow a built-in one, and it is the +name to use in `tool_allowlist` and `tool_denylist`. + +### What Superset does to keep a foreign server contained + +A third-party server is untrusted input, and possibly untrusted intent: + +- **Everything it returns is marked as untrusted** before the model sees it, so + text in a tool result is treated as data rather than instructions. Tool + *descriptions* get the same treatment, since they enter the prompt every turn. +- **No Superset credential is ever forwarded.** Only the headers you configured + for that server are sent — never the user's session cookie, CSRF token, or an + inbound authorization header. +- **SQL execution through a foreign server is refused by default.** Superset's + read-only enforcement and per-dataset authorization cannot apply to a query + another system runs, so allowing it would silently bypass both. Set + `AI_AGENT_MCP_DENY_FOREIGN_SQL = False` to accept that trade deliberately. +- **Results obey the same size cap** as built-in tools, and the cap is applied + while reading, so a hostile server cannot exhaust memory before truncation. +- **A server being down does not break the assistant.** Discovery failure means + that server contributes no tools for the turn; the built-ins keep working. + +A profile naming a server you have not configured is an error, because a typo +there is indistinguishable at runtime from an agent that has quietly lost a +capability. Note that discovery happens per turn, so a slow server adds its +latency to every turn that uses it. + +## Retention + +Conversations are kept for `AI_ASSISTANT_MESSAGE_RETENTION_DAYS` (default 30). Review Comment: Agreed. The retention contract and missing pruning task are inherited unchanged from #42805, where the root issue is already tracked: https://github.com/apache/superset/pull/42805#discussion_r3739999196. This child PR will be rebased after the parent resolves it; I am leaving the thread open meanwhile. ########## superset-frontend/src/features/ai/components/ChatChartEmbed.tsx: ########## @@ -0,0 +1,542 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +/** + * @fileoverview A chart rendered inside a chat message. + * + * The assistant does not send a chart, it sends a `form_data_key` it stored — so + * what arrives in the transcript is a reference the client resolves, and the + * rendered chart is the real thing, with the real permissions, rather than an + * image of one. + * + * The awkward part is timing: the key can exist before the query behind it has + * finished. Rather than show the chart's own "No data" state (which reads as a + * broken answer) the component keeps a spinner up and re-renders on a growing + * backoff until rows appear, giving up after a bounded number of attempts. + */ + +import { useCallback, useEffect, useRef, useState } from 'react'; +import { + type QueryFormData, + StatefulChart, + SupersetClient, +} from '@superset-ui/core'; +import { styled } from '@apache-superset/core/theme'; +import { t } from '@apache-superset/core/translation'; +import { Loading } from '@superset-ui/core/components'; +import { Icons } from '@superset-ui/core/components/Icons'; +import { ErrorBoundary } from 'src/components/ErrorBoundary'; + +const VALID_KEY_PATTERN = /^[a-zA-Z0-9_-]+$/; +const MIN_HEIGHT = 100; +const MAX_HEIGHT = 800; +const DEFAULT_HEIGHT = 300; +const FETCH_TIMEOUT_MS = 30_000; +const MAX_RETRIES = 3; +const RETRY_DELAYS_MS = [500, 1500, 3000]; + +/** Width used until the container has been measured. */ +const FALLBACK_CHART_WIDTH = 600; + +// Backoff for the "waiting for chart data" poll. The delay grows exponentially +// per attempt up to a ceiling, so a slow query is waited out without hammering +// the backend. +const POLL_BASE_DELAY_MS = 1000; +const POLL_MAX_DELAY_MS = 30_000; + +/** + * Polling stops after this many attempts. + * + * A retry re-issues the chart's data request, which will use the results cache if + * the query has landed but will otherwise execute it. Polling forever would keep + * re-issuing it, so an unfinished query surfaces the retry control instead. + */ +const MAX_POLL_ATTEMPTS = 6; + +const getPollDelayMs = (attempt: number): number => + Math.min(POLL_BASE_DELAY_MS * 2 ** attempt, POLL_MAX_DELAY_MS); + +export interface ChartEmbedParams { + formDataKey: string | null; + height: number; + title: string | null; +} + +/** + * Parse key=value lines from the content of a ```superset-chart fenced block. + * + * Rules are strict on purpose: the block is model output, so `form_data_key` must + * match `/^[a-zA-Z0-9_-]+$/` before it reaches a URL, and `height` is clamped. + * Unknown keys are ignored so a newer backend can add some without breaking an + * older client. + */ +export function parseChartEmbedParams(codeText: string): ChartEmbedParams { + const result: ChartEmbedParams = { + formDataKey: null, + height: DEFAULT_HEIGHT, + title: null, + }; + + const lines = codeText + .split('\n') + .map(line => line.trim()) + .filter(Boolean); + + lines.forEach(line => { + const eqIndex = line.indexOf('='); + if (eqIndex <= 0) { + return; + } + + const key = line.slice(0, eqIndex).trim().toLowerCase(); + const value = line.slice(eqIndex + 1).trim(); + + if (key === 'form_data_key') { + if (value && VALID_KEY_PATTERN.test(value)) { + result.formDataKey = value; + } + return; + } + if (key === 'height') { + const parsed = parseInt(value, 10); + if (!Number.isNaN(parsed)) { + result.height = Math.max(MIN_HEIGHT, Math.min(MAX_HEIGHT, parsed)); + } + return; + } + if (key === 'title' && value) { + result.title = value; + } + }); + + return result; +} + +const ChartContainer = styled.div` + border: 1px solid ${({ theme }) => theme.colorBorderSecondary}; + border-radius: ${({ theme }) => theme.borderRadius}px; + overflow: hidden; + margin: ${({ theme }) => theme.sizeUnit * 2}px 0; + background: ${({ theme }) => theme.colorBgContainer}; +`; + +const ChartHeader = styled.div` + display: flex; + justify-content: space-between; + align-items: center; + padding: ${({ theme }) => theme.sizeUnit * 2}px + ${({ theme }) => theme.sizeUnit * 3}px; + border-bottom: 1px solid ${({ theme }) => theme.colorBorderSecondary}; + background: ${({ theme }) => theme.colorBgLayout}; + font-size: ${({ theme }) => theme.fontSizeSM}px; +`; + +const ChartTitle = styled.span` + font-weight: ${({ theme }) => theme.fontWeightStrong}; + color: ${({ theme }) => theme.colorText}; + overflow: hidden; + text-overflow: ellipsis; + white-space: nowrap; + flex: 1; + min-width: 0; +`; + +const ChartActions = styled.div` + display: flex; + gap: ${({ theme }) => theme.sizeUnit * 2}px; + align-items: center; + flex-shrink: 0; + margin-left: ${({ theme }) => theme.sizeUnit * 2}px; +`; + +const ActionLink = styled.a` + display: inline-flex; + align-items: center; + gap: ${({ theme }) => theme.sizeUnit / 2}px; + font-size: ${({ theme }) => theme.fontSizeSM}px; + color: ${({ theme }) => theme.colorPrimary}; + cursor: pointer; + text-decoration: none; + + &:hover { + text-decoration: underline; + } +`; + +const ActionButton = styled.button` + display: inline-flex; + align-items: center; + gap: ${({ theme }) => theme.sizeUnit / 2}px; + font-size: ${({ theme }) => theme.fontSizeSM}px; + color: ${({ theme }) => theme.colorTextSecondary}; + cursor: pointer; + background: none; + border: none; + padding: 2px ${({ theme }) => theme.sizeUnit / 2}px; + border-radius: ${({ theme }) => theme.borderRadius}px; + + &:hover { + color: ${({ theme }) => theme.colorPrimary}; + background: ${({ theme }) => theme.colorFillTertiary}; + } +`; + +const ChartBody = styled.div<{ height: number }>` + height: ${({ height }) => height}px; + position: relative; +`; + +// Covers the chart while it reports no data, hiding the underlying "No data" +// state (which looks broken) behind a spinner while refreshing continues. +const ChartDataOverlay = styled.div` + position: absolute; + inset: 0; + display: flex; + flex-direction: column; + align-items: center; + justify-content: center; + gap: ${({ theme }) => theme.sizeUnit * 3}px; + background: ${({ theme }) => theme.colorBgContainer}; + color: ${({ theme }) => theme.colorTextSecondary}; + font-size: ${({ theme }) => theme.fontSizeSM}px; + z-index: 2; +`; + +const CenteredMessage = styled.div<{ height: number }>` + display: flex; + flex-direction: column; + align-items: center; + justify-content: center; + height: ${({ height }) => height}px; + color: ${({ theme }) => theme.colorTextSecondary}; + font-size: ${({ theme }) => theme.fontSizeSM}px; + text-align: center; + padding: ${({ theme }) => theme.sizeUnit * 4}px; + gap: ${({ theme }) => theme.sizeUnit * 2}px; +`; + +interface ChatChartEmbedProps { + formDataKey: string; + height?: number; + title?: string; +} + +type FetchState = + | { status: 'loading' } + | { status: 'loaded'; formData: QueryFormData } + | { status: 'error'; message: string }; + +const exploreUrlFor = (formDataKey: string): string => Review Comment: Agreed. The reviewed Explore-link construction is inherited unchanged from the parent Native AI PR #42805 and is not introduced by the authoring adapters. The application-root-aware fix belongs in the parent UI implementation, then this branch can be rebased. I am leaving the thread open pending that fix. ########## superset-frontend/src/features/ai/components/ChatTabsMenu.tsx: ########## @@ -0,0 +1,356 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +/** + * @fileoverview The conversation list. + * + * Conversations live behind one menu rather than a tab strip: the panel is narrow + * enough in floating mode that a strip would truncate every name, and the list + * doubles as the history of past conversations, which a strip cannot be. + */ + +import { useCallback, useState } from 'react'; +import type { MouseEvent as ReactMouseEvent } from 'react'; +import { styled } from '@apache-superset/core/theme'; +import { t } from '@apache-superset/core/translation'; +import { Button, Dropdown, Popconfirm } from '@superset-ui/core/components'; +import { Icons } from '@superset-ui/core/components/Icons'; +import type { ChatTab } from '../types'; + +const MenuContainer = styled.div` + background: ${({ theme }) => theme.colorBgElevated}; + border-radius: ${({ theme }) => theme.borderRadius}px; + box-shadow: ${({ theme }) => theme.boxShadowSecondary}; + min-width: ${({ theme }) => theme.sizeUnit * 65}px; + max-height: ${({ theme }) => theme.sizeUnit * 100}px; + overflow-y: auto; + border: 1px solid ${({ theme }) => theme.colorBorderSecondary}; +`; + +const MenuHeader = styled.div` + padding: ${({ theme }) => theme.sizeUnit * 3}px + ${({ theme }) => theme.sizeUnit * 4}px; + border-bottom: 1px solid ${({ theme }) => theme.colorBorderSecondary}; + font-weight: ${({ theme }) => theme.fontWeightStrong}; + font-size: ${({ theme }) => theme.fontSizeSM}px; + color: ${({ theme }) => theme.colorTextSecondary}; + text-transform: uppercase; + letter-spacing: 0.5px; +`; + +const NewChatButton = styled.button` + display: flex; + align-items: center; + gap: ${({ theme }) => theme.sizeUnit * 2}px; + width: 100%; + padding: ${({ theme }) => theme.sizeUnit * 2.5}px + ${({ theme }) => theme.sizeUnit * 4}px; + cursor: pointer; + color: ${({ theme }) => theme.colorPrimary}; + font-weight: ${({ theme }) => theme.fontWeightStrong}; + background: none; + border: none; + text-align: left; + transition: background ${({ theme }) => theme.motionDurationMid}; + + &:hover { + background: ${({ theme }) => theme.colorFillTertiary}; + } +`; + +const TabItem = styled.div<{ isActive: boolean }>` + display: flex; + align-items: center; + justify-content: space-between; + padding: ${({ theme }) => theme.sizeUnit * 2.5}px + ${({ theme }) => theme.sizeUnit * 4}px; + cursor: pointer; + background: ${({ theme, isActive }) => + isActive ? theme.colorFillSecondary : 'transparent'}; + border-left: 3px solid + ${({ theme, isActive }) => (isActive ? theme.colorPrimary : 'transparent')}; + transition: background ${({ theme }) => theme.motionDurationMid}; + + &:hover { + background: ${({ theme }) => theme.colorFillTertiary}; + + .action-btn { + opacity: 1; + } + } +`; + +const TabInfo = styled.div` + display: flex; + align-items: center; + gap: ${({ theme }) => theme.sizeUnit * 2}px; + flex: 1; + overflow: hidden; +`; + +const TabName = styled.span` + font-size: ${({ theme }) => theme.fontSize}px; + color: ${({ theme }) => theme.colorText}; + white-space: nowrap; + overflow: hidden; + text-overflow: ellipsis; + max-width: ${({ theme }) => theme.sizeUnit * 35}px; +`; + +const TabNameInput = styled.input` + width: 100%; + max-width: ${({ theme }) => theme.sizeUnit * 40}px; + font-size: ${({ theme }) => theme.fontSize}px; + color: ${({ theme }) => theme.colorText}; + background: ${({ theme }) => theme.colorBgContainer}; + border: 1px solid ${({ theme }) => theme.colorBorder}; + border-radius: ${({ theme }) => theme.borderRadius}px; + padding: 2px ${({ theme }) => theme.sizeUnit * 1.5}px; +`; + +const TabTimestamp = styled.span` + font-size: ${({ theme }) => theme.fontSizeSM}px; + color: ${({ theme }) => theme.colorTextQuaternary}; + white-space: nowrap; + flex-shrink: 0; +`; + +const ActionButtons = styled.div` + display: flex; + align-items: center; + gap: 2px; +`; + +const ActionButton = styled.button` + background: none; + border: none; + padding: ${({ theme }) => theme.sizeUnit}px; + cursor: pointer; + color: ${({ theme }) => theme.colorTextSecondary}; + opacity: 0; + transition: all ${({ theme }) => theme.motionDurationMid}; + display: flex; + align-items: center; + justify-content: center; + border-radius: ${({ theme }) => theme.borderRadius}px; + + &:hover, + &:focus-visible { + opacity: 1; + color: ${({ theme }) => theme.colorError}; + background: ${({ theme }) => theme.colorErrorBg}; + } +`; + +const Divider = styled.div` + height: 1px; + background: ${({ theme }) => theme.colorBorderSecondary}; + margin: ${({ theme }) => theme.sizeUnit}px 0; +`; + +const EmptyState = styled.div` + padding: ${({ theme }) => theme.sizeUnit * 5}px + ${({ theme }) => theme.sizeUnit * 4}px; + text-align: center; + color: ${({ theme }) => theme.colorTextSecondary}; + font-size: ${({ theme }) => theme.fontSizeSM}px; +`; + +const MINUTE_SECONDS = 60; +const HOUR_MINUTES = 60; +const DAY_HOURS = 24; +const WEEK_DAYS = 7; + +export const formatRelativeTime = (timestamp: number): string => { + const seconds = Math.floor((Date.now() - timestamp) / 1000); + if (seconds < MINUTE_SECONDS) { + return t('just now'); + } + const minutes = Math.floor(seconds / MINUTE_SECONDS); + if (minutes < HOUR_MINUTES) { + return t('%sm', String(minutes)); + } + const hours = Math.floor(minutes / HOUR_MINUTES); + if (hours < DAY_HOURS) { + return t('%sh', String(hours)); + } + const days = Math.floor(hours / DAY_HOURS); + if (days < WEEK_DAYS) { + return t('%sd', String(days)); + } + return new Date(timestamp).toLocaleDateString(undefined, { + month: 'short', + day: 'numeric', + }); +}; + +interface ChatTabsMenuProps { + tabs: ChatTab[]; + activeTabId: string; + onSelectTab: (tabId: string) => void; + onNewChat: () => void; + onDeleteTab: (tabId: string) => void; + onRenameTab: (tabId: string, name: string) => void; +} + +export const ChatTabsMenu = ({ + tabs, + activeTabId, + onSelectTab, + onNewChat, + onDeleteTab, + onRenameTab, +}: ChatTabsMenuProps) => { + const [editingTabId, setEditingTabId] = useState<string | null>(null); + const [editingName, setEditingName] = useState(''); + + const startEditing = useCallback((event: ReactMouseEvent, tab: ChatTab) => { + event.stopPropagation(); + setEditingTabId(tab.id); + setEditingName(tab.name); + }, []); + + const cancelEditing = useCallback(() => { + setEditingTabId(null); + setEditingName(''); + }, []); + + const commitRename = useCallback( + (tabId: string) => { + const trimmedName = editingName.trim(); + if (trimmedName) { + onRenameTab(tabId, trimmedName); + } + cancelEditing(); + }, + [cancelEditing, editingName, onRenameTab], + ); + + const menuContent = ( + <MenuContainer data-test="chat-tabs-menu"> + <MenuHeader>{t('Conversations')}</MenuHeader> + <NewChatButton type="button" onClick={onNewChat}> + <Icons.PlusOutlined iconSize="s" /> + <span>{t('New Chat')}</span> + </NewChatButton> + <Divider /> + {tabs.length === 0 ? ( + <EmptyState>{t('No conversations yet')}</EmptyState> + ) : ( + tabs.map(tab => ( + <TabItem + key={tab.id} + isActive={tab.id === activeTabId} + onClick={() => onSelectTab(tab.id)} + > + <TabInfo> + <Icons.MessageOutlined iconSize="s" /> + {editingTabId === tab.id ? ( + <TabNameInput + autoFocus + value={editingName} + onChange={event => setEditingName(event.target.value)} + onClick={event => event.stopPropagation()} + onBlur={() => commitRename(tab.id)} + onKeyDown={event => { + event.stopPropagation(); + if (event.key === 'Enter') { + commitRename(tab.id); + } else if (event.key === 'Escape') { + cancelEditing(); + } + }} + aria-label={t('Conversation name')} + /> + ) : ( + <TabName>{tab.name}</TabName> + )} + {tab.updatedAt !== undefined && ( + <TabTimestamp>{formatRelativeTime(tab.updatedAt)}</TabTimestamp> + )} + </TabInfo> + <ActionButtons> + <ActionButton + type="button" + className="action-btn" + onClick={event => startEditing(event, tab)} + title={t('Rename conversation')} + aria-label={t('Rename conversation')} + > + <Icons.EditOutlined iconSize="s" /> + </ActionButton> + {/* A conversation with messages is confirmed before deletion; an + empty one is discarded without a prompt. */} + {tab.messages.length > 0 ? ( Review Comment: Agreed. Unopened-thread deletion behavior is inherited unchanged from #42805 and is already tracked on the parent: https://github.com/apache/superset/pull/42805#discussion_r3821577570. This authoring PR should not duplicate the parent state-model fix, so the thread remains open until rebase. ########## superset-frontend/src/features/ai/AiAssistantPanel.tsx: ########## @@ -0,0 +1,1150 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +/** + * @fileoverview The assistant panel. + * + * The host owns where this sits and, when docked, how wide it is, so there is no + * positioning and no resize handle here. What is here is the conversation: the + * header, the transcript, what the assistant is doing while it works, and the + * composer. + * + * The centre of the design is that a run is legible while it happens. An answer + * can take a minute of tool calls, and a spinner for a minute is indistinguishable + * from a hang, so reasoning streams into a preview, each step appends to a tool + * log, and a checkpoint stops the run with a countdown the user can act on. + */ + +import { + memo, + useCallback, + useEffect, + useRef, + useState, + useSyncExternalStore, +} from 'react'; +import type { Dispatch, SetStateAction } from 'react'; +import ReactMarkdown from 'react-markdown'; +import type { Components } from 'react-markdown'; +import { css, keyframes, styled, useTheme } from '@apache-superset/core/theme'; +import { t } from '@apache-superset/core/translation'; +import type { chat as chatApi } from '@apache-superset/core'; +import { + Button, + Input, + Loading, + Tooltip, + Typography, +} from '@superset-ui/core/components'; +import { Icons } from '@superset-ui/core/components/Icons'; +import { chat } from 'src/core/chat'; +import ChatAgentSelect from './components/ChatAgentSelect'; +import ChatTabsMenu from './components/ChatTabsMenu'; +import { REMARK_PLUGINS, useChatMarkdown } from './components/chatMarkdown'; +import { ThoughtProcess } from './components/ThoughtProcess'; +import { useChatBot } from './hooks/useChatBot'; +import { AI_ACTION_EVENT, type AiActionEvent } from './hooks/useAIAction'; +import type { PageContext } from './hooks/usePageContext'; +import type { ChatMessageWithMeta, CheckpointPayload } from './types'; + +/** + * How long a checkpoint waits before continuing on its own. A pause that blocks + * forever is worse than one that resolves optimistically: the user may have + * walked away, and the run should not be stranded. + */ +export const CHECKPOINT_TIMEOUT_SECONDS = 30; + +/** + * Closes a code fence the model has not finished writing. + * + * A streamed answer is parsed on every delta, so a fence arrives in pieces — + * "```", then "sql", then the query. Markdown with an odd number of fences + * renders the opening backticks literally and then reflows once the closing pair + * lands, which reads as the answer glitching. Balancing the count keeps each + * intermediate state a valid document. + */ +export const balanceCodeFences = (text: string): string => { + const fences = text.match(/^```/gm)?.length ?? 0; + return fences % 2 === 0 ? text : `${text}\n\`\`\``; +}; + +/** + * Whether a message carries the structured record of how it was answered, as + * opposed to only the flat log assembled from stream frames. + */ +const hasStructuredThinking = (message: ChatMessageWithMeta): boolean => + Boolean(message.toolCalls?.length || message.thoughts || message.pageContext); + +/** Milliseconds between typewriter frames, and characters per frame. */ +const TYPEWRITER_INTERVAL_MS = 18; +const TYPEWRITER_STEP = 3; + +/** + * The panel's own size as a floating overlay. + * + * Docked width belongs to the host and is not set here. Floating does need a size + * from somewhere, though — the floating host only stacks its children in a corner + * and gives them no dimensions — so these clamp the overlay to the viewport. + */ +const FLOATING_WIDTH_PX = 440; +const FLOATING_MAX_HEIGHT_VH = 70; + +/** + * The panel surface. + * + * Positioning is deliberately absent: the host places this, in both modes. What is + * here is the surface itself — a column that fills whatever box it is given, with a + * floating size for the mode where the host provides no box. + */ +const ChatPanelContainer = styled.div<{ floating: boolean }>` + display: flex; + flex-direction: column; + min-height: 0; + overflow: hidden; + background: ${({ theme }) => theme.colorBgElevated}; + ${({ floating, theme }) => + floating + ? css` + width: min( + ${FLOATING_WIDTH_PX}px, + calc(100vw - ${theme.sizeUnit * 12}px) + ); + height: ${FLOATING_MAX_HEIGHT_VH}vh; + border: 1px solid ${theme.colorBorderSecondary}; + border-radius: ${theme.borderRadiusLG}px; + box-shadow: ${theme.boxShadow}; + ` + : css` + width: 100%; + height: 100%; + `} +`; + +const ChatHeader = styled.div` + display: flex; + justify-content: space-between; + align-items: center; + padding: ${({ theme }) => theme.sizeUnit * 3}px + ${({ theme }) => theme.sizeUnit * 4}px; + background: ${({ theme }) => theme.colorBgContainer}; + border-bottom: 1px solid ${({ theme }) => theme.colorBorderSecondary}; + font-weight: ${({ theme }) => theme.fontWeightStrong}; + font-size: ${({ theme }) => theme.fontSizeLG}px; + color: ${({ theme }) => theme.colorTextHeading}; + flex-shrink: 0; + gap: ${({ theme }) => theme.sizeUnit * 2}px; +`; + +const HeaderGroup = styled.div` + display: flex; + align-items: center; + min-width: 0; +`; + +const HeaderTitle = styled.span` + overflow: hidden; + text-overflow: ellipsis; + white-space: nowrap; +`; + +const ChatMessages = styled.div` + flex: 1; + min-height: 0; + padding: ${({ theme }) => theme.sizeUnit * 4}px; + overflow-y: auto; + + &::-webkit-scrollbar { + width: 4px; + } + + &::-webkit-scrollbar-track { + background: ${({ theme }) => theme.colorBgContainer}; + border-radius: 2px; + } + + &::-webkit-scrollbar-thumb { + background: ${({ theme }) => theme.colorFillSecondary}; + border-radius: 2px; + } +`; + +const MessageBubble = styled.div<{ variant: 'user' | 'assistant' }>` + margin-bottom: ${({ theme }) => theme.sizeUnit * 3}px; + display: flex; + flex-direction: column; + align-items: ${({ variant }) => + variant === 'user' ? 'flex-end' : 'flex-start'}; +`; + +const MessageContent = styled.div<{ variant: 'user' | 'assistant' }>` + max-width: 85%; + padding: ${({ theme }) => theme.sizeUnit * 3}px + ${({ theme }) => theme.sizeUnit * 4}px; + border-radius: ${({ theme }) => theme.borderRadiusLG * 2}px; + background: ${({ theme, variant }) => + variant === 'user' ? theme.colorPrimary : theme.colorBgContainer}; + color: ${({ theme, variant }) => + variant === 'user' ? theme.colorTextLightSolid : theme.colorText}; + font-size: ${({ theme }) => theme.fontSize}px; + line-height: 1.5; + border: ${({ theme, variant }) => + variant === 'assistant' + ? `1px solid ${theme.colorBorderSecondary}` + : 'none'}; + box-shadow: ${({ theme }) => theme.boxShadowTertiary}; + overflow-wrap: anywhere; + + p { + margin: 0; + } + + p:not(:last-child) { + margin-bottom: ${({ theme }) => theme.sizeUnit * 2}px; + } + + a { + color: ${({ theme, variant }) => + variant === 'user' ? theme.colorTextLightSolid : theme.colorPrimary}; + text-decoration: underline; + } + + code { + background: ${({ theme, variant }) => + variant === 'user' ? theme.colorPrimaryActive : theme.colorFillTertiary}; + padding: 2px ${({ theme }) => theme.sizeUnit * 1.5}px; + border-radius: ${({ theme }) => theme.borderRadius}px; + font-size: ${({ theme }) => theme.fontSizeSM}px; + font-family: ${({ theme }) => theme.fontFamilyCode}; + } + + pre { + background: ${({ theme }) => theme.colorFillQuaternary}; + padding: ${({ theme }) => theme.sizeUnit * 3}px; + border-radius: ${({ theme }) => theme.borderRadius}px; + overflow-x: auto; + margin: ${({ theme }) => theme.sizeUnit * 2}px 0; + border: 1px solid ${({ theme }) => theme.colorBorderSecondary}; + } + + pre code { + background: none; + padding: 0; + } +`; + +const MessageActions = styled.div` + display: flex; + gap: ${({ theme }) => theme.sizeUnit}px; + margin-top: ${({ theme }) => theme.sizeUnit}px; +`; + +const ActionButton = styled(Button)` + &&& { + padding: 2px ${({ theme }) => theme.sizeUnit * 1.5}px; + height: ${({ theme }) => theme.sizeUnit * 6}px; + font-size: ${({ theme }) => theme.fontSizeSM}px; + } + + /* The recorded verdict keeps its colour while disabled. Both thumbs lock once + a rating exists, and the default disabled grey would hide which one the + user picked — the state matters more here than the affordance. */ + &&&.is-active, + &&&.is-active:disabled, + &&&.is-active[disabled] { + color: ${({ theme }) => theme.colorPrimary}; + } +`; + +const LiveAnswer = styled.div` + margin-top: ${({ theme }) => theme.sizeUnit * 2}px; + color: ${({ theme }) => theme.colorText}; + font-size: ${({ theme }) => theme.fontSize}px; + line-height: 1.5; + overflow-wrap: anywhere; + + p { + margin: 0; + } + + p:not(:last-child) { + margin-bottom: ${({ theme }) => theme.sizeUnit * 2}px; + } +`; + +const thinkingPulse = keyframes` + 0% { + opacity: 0.45; + } + 50% { + opacity: 1; + } + 100% { + opacity: 0.45; + } +`; + +const ThinkingPreview = styled.div<{ isLive?: boolean }>` + color: ${({ theme }) => theme.colorTextTertiary}; + font-size: ${({ theme }) => theme.fontSizeSM}px; + white-space: pre-wrap; + ${({ isLive }) => + isLive && + css` + animation: ${thinkingPulse} 1.8s ease-in-out infinite; + `} +`; + +const ThinkingDetails = styled.details` + margin-bottom: ${({ theme }) => theme.sizeUnit * 2}px; + color: ${({ theme }) => theme.colorTextTertiary}; + font-size: ${({ theme }) => theme.fontSizeSM}px; + + summary { + cursor: pointer; + user-select: none; + color: ${({ theme }) => theme.colorTextTertiary}; + margin-bottom: ${({ theme }) => theme.sizeUnit}px; + } +`; + +const CheckpointDivider = styled.div` + display: flex; + align-items: center; + gap: ${({ theme }) => theme.sizeUnit * 3}px; + margin: ${({ theme }) => theme.sizeUnit * 4}px 0 + ${({ theme }) => theme.sizeUnit * 3}px; + + &::before, + &::after { + content: ''; + flex: 1; + height: 1px; + background: ${({ theme }) => theme.colorBorderSecondary}; + } +`; + +const CountdownBadge = styled.span` + font-size: ${({ theme }) => theme.fontSizeSM}px; + font-weight: ${({ theme }) => theme.fontWeightStrong}; + font-variant-numeric: tabular-nums; + color: ${({ theme }) => theme.colorTextSecondary}; + white-space: nowrap; +`; + +const CheckpointContent = styled.div` + font-size: ${({ theme }) => theme.fontSizeSM}px; + color: ${({ theme }) => theme.colorText}; + line-height: 1.5; +`; + +const CheckpointTaskList = styled.ul` + margin: ${({ theme }) => theme.sizeUnit * 1.5}px 0; + padding-left: ${({ theme }) => theme.sizeUnit * 4.5}px; + font-size: ${({ theme }) => theme.fontSizeSM}px; + color: ${({ theme }) => theme.colorTextSecondary}; + + li { + margin-bottom: 2px; + } +`; + +const CheckpointEstimate = styled.div` + font-size: ${({ theme }) => theme.fontSizeSM}px; + color: ${({ theme }) => theme.colorTextTertiary}; + margin-top: ${({ theme }) => theme.sizeUnit}px; +`; + +const CheckpointActions = styled.div` + display: flex; + gap: ${({ theme }) => theme.sizeUnit * 2}px; + margin-top: ${({ theme }) => theme.sizeUnit * 2.5}px; +`; + +const ChatInput = styled.div` + padding: ${({ theme }) => theme.sizeUnit * 4}px; + border-top: 1px solid ${({ theme }) => theme.colorBorderSecondary}; + background: ${({ theme }) => theme.colorBgContainer}; + flex-shrink: 0; +`; + +const InputContainer = styled.div` + display: flex; + gap: ${({ theme }) => theme.sizeUnit * 2}px; + align-items: flex-end; +`; + +const QuickPromptsRow = styled.div<{ hasContent: boolean }>` + display: flex; + align-items: center; + gap: ${({ theme }) => theme.sizeUnit * 2}px; + flex-wrap: wrap; + margin-bottom: ${({ theme, hasContent }) => + hasContent ? `${theme.sizeUnit * 2.5}px` : '0'}; + min-height: ${({ theme, hasContent }) => + hasContent ? `${theme.sizeUnit * 6}px` : '0'}; +`; + +const QuickPromptChip = styled(Button)` + &&& { + width: fit-content; + max-width: 100%; + height: auto; + white-space: normal; + text-align: left; + line-height: 1.35; + word-break: break-word; + } +`; + +const PageContextRow = styled.div` + display: flex; + align-items: center; + margin-bottom: ${({ theme }) => theme.sizeUnit * 1.5}px; +`; + +const PageContextPill = styled.button<{ isActive: boolean }>` + display: inline-flex; + align-items: center; + gap: ${({ theme }) => theme.sizeUnit * 1.5}px; + padding: 3px ${({ theme }) => theme.sizeUnit * 2}px; + border-radius: ${({ theme }) => theme.borderRadiusLG}px; + border: 1px solid + ${({ theme, isActive }) => + isActive ? theme.colorPrimary : theme.colorBorderSecondary}; + background: ${({ theme, isActive }) => + isActive ? theme.colorPrimaryBg : theme.colorFillQuaternary}; + color: ${({ theme, isActive }) => + isActive ? theme.colorPrimary : theme.colorTextTertiary}; + font-size: ${({ theme }) => theme.fontSizeSM}px; + cursor: pointer; + transition: all ${({ theme }) => theme.motionDurationMid}; + max-width: ${({ theme }) => theme.sizeUnit * 62}px; + white-space: nowrap; + + &:hover { + border-color: ${({ theme }) => theme.colorPrimary}; + } +`; + +const PillLabel = styled.span` + overflow: hidden; + text-overflow: ellipsis; Review Comment: Agreed. The auto-scroll effect is inherited unchanged from the parent Native AI PR #42805; the authoring commits do not modify this panel behavior. The dependency fix should land once in the parent UI, then be rebased here. I am leaving this thread open pending that update. ########## 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: Agreed. The single checkpoint resolver is inherited unchanged from the parent Native AI PR #42805 and is not part of the authoring tool change. The gate needs to be keyed in the parent run-state implementation, then rebased here. I am leaving the thread open until that is resolved. -- 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]
