I3eka commented on code in PR #43132: URL: https://github.com/apache/superset/pull/43132#discussion_r3810342150
########## 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: Good catch. Fixed in 6a8b2b8137: the read-only policy now normalizes namespaced MCP tool names to the remote tool name and unwraps the MCP `request` payload before parsing. The regression test uses the actual `mcp__superset__create_virtual_dataset` shape with `request.sql` and confirms that a DELETE is denied; the nested `database_id` is also used for dialect resolution. ########## 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: Agreed. Fixed in 6a8b2b8137: worker mode now writes, reads, and clears cancellation through the same Redis backend already required by `AI_ASSISTANT_EVENT_BUS_CACHE_CONFIG`; the process-local set remains only for inline mode. The regression test verifies that worker cancellation has no local flag and remains observable through the shared backend. The policy/orchestrator suite passes 69 tests. -- 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]
