From d98a5d855d3beada16887400a338ea69e0305377 Mon Sep 17 00:00:00 2001 From: Wenxi Date: Wed, 12 Aug 2026 19:08:05 +0000 Subject: [PATCH] feat(scheduled-tasks): support MCP pre-approvals (#13884) --- backend/onyx/db/models.py | 35 +++--- backend/onyx/db/scheduled_task.py | 90 ++++++++------ .../features/build/scheduled_tasks/api.py | 79 ++++++++++--- .../test_scheduled_task_pre_approvals.py | 111 +++++++++++------- .../tests/craft/test_scheduled_tasks_api.py | 102 ++++++++++++++++ backend/tests/unit/sandbox_proxy/test_gate.py | 48 ++++++-- .../egress-proxy-and-approvals/README.md | 8 +- .../features/scheduled-tasks/pre-approvals.md | 38 +++--- 8 files changed, 365 insertions(+), 146 deletions(-) diff --git a/backend/onyx/db/models.py b/backend/onyx/db/models.py index 027b04492d..7942f65905 100644 --- a/backend/onyx/db/models.py +++ b/backend/onyx/db/models.py @@ -6457,26 +6457,33 @@ class ScheduledTask(Base): back_populates="task", cascade="all, delete-orphan", ) - pre_approved_apps: Mapped[list["ScheduledTaskPreApprovedApp"]] = relationship( - "ScheduledTaskPreApprovedApp", + pre_approved_targets: Mapped[list["ScheduledTaskPreApprovedTarget"]] = relationship( + "ScheduledTaskPreApprovedTarget", back_populates="task", cascade="all, delete-orphan", - order_by="ScheduledTaskPreApprovedApp.id", + order_by="ScheduledTaskPreApprovedTarget.id", ) @property def pre_approved_external_app_ids(self) -> list[int]: - """Granted external-app ids in grant order. MCP-server grants are - excluded — their target ids live in a different id space, and this - property backs the API's external-app-id field (the gate reads all - grants regardless of kind via ``get_running_scheduled_run_grants``). - Set via ``onyx.db.scheduled_task.set_pre_approved_apps``.""" + """Granted external-app ids. MCP-server grants are excluded because + their target ids live in a different id space.""" return [ grant.gated_app.external_app_id - for grant in self.pre_approved_apps + for grant in self.pre_approved_targets if grant.gated_app.external_app_id is not None ] + @property + def pre_approved_mcp_server_ids(self) -> list[int]: + """Granted MCP-server ids. External-app grants are excluded because + their target ids live in a different id space.""" + return [ + grant.gated_app.mcp_server_id + for grant in self.pre_approved_targets + if grant.gated_app.mcp_server_id is not None + ] + __table_args__ = ( # Dispatcher hot path: WHERE status='active' AND deleted=false # AND next_run_at <= now() ORDER BY next_run_at @@ -6567,7 +6574,7 @@ class ScheduledTaskRun(Base): ) -class ScheduledTaskPreApprovedApp(Base): +class ScheduledTaskPreApprovedTarget(Base): """One (task, target) pre-approval grant: the matched target's ASK-gated actions skip the approval park for the task's RUNNING runs. @@ -6577,6 +6584,7 @@ class ScheduledTaskPreApprovedApp(Base): index serves the per-task lookup. """ + # The table name predates MCP support. Keep it to avoid a schema-only rename. __tablename__ = "scheduled_task_pre_approved_app" id: Mapped[int] = mapped_column(Integer, primary_key=True) @@ -6595,11 +6603,10 @@ class ScheduledTaskPreApprovedApp(Base): ) task: Mapped[ScheduledTask] = relationship( - "ScheduledTask", back_populates="pre_approved_apps" + "ScheduledTask", back_populates="pre_approved_targets" ) - # selectin: pre_approved_external_app_ids and set_pre_approved_apps read - # gated_app for every grant, so batch them in one SELECT rather than one - # per grant. + # selectin: the pre-approved target-id properties read gated_app for every + # grant, so batch them in one SELECT rather than one per grant. gated_app: Mapped["GatedApp"] = relationship("GatedApp", lazy="selectin") __table_args__ = ( diff --git a/backend/onyx/db/scheduled_task.py b/backend/onyx/db/scheduled_task.py index c57b3e8747..47a45767be 100644 --- a/backend/onyx/db/scheduled_task.py +++ b/backend/onyx/db/scheduled_task.py @@ -31,7 +31,7 @@ from onyx.db.gated_app import get_or_create_gated_app_id from onyx.db.models import ( GatedApp, ScheduledTask, - ScheduledTaskPreApprovedApp, + ScheduledTaskPreApprovedTarget, ScheduledTaskRun, ) from onyx.error_handling.error_codes import OnyxErrorCode @@ -60,6 +60,7 @@ def create_scheduled_task( editor_mode: EditorMode, status: ScheduledTaskStatus = ScheduledTaskStatus.ACTIVE, pre_approved_external_app_ids: list[int] | None = None, + pre_approved_mcp_server_ids: list[int] | None = None, now: datetime | None = None, ) -> ScheduledTask: """Insert a new ``ScheduledTask``. @@ -84,42 +85,60 @@ def create_scheduled_task( status=status, next_run_at=next_run_at, ) - set_pre_approved_apps( + _replace_pre_approved_targets( db_session, task, - GatedAppKind.EXTERNAL_APP, - pre_approved_external_app_ids or [], + external_app_ids=pre_approved_external_app_ids or [], + mcp_server_ids=pre_approved_mcp_server_ids or [], ) db_session.add(task) db_session.flush() return task -def set_pre_approved_apps( +def _replace_pre_approved_targets( db_session: Session, task: ScheduledTask, - kind: GatedAppKind, - target_ids: list[int], + *, + external_app_ids: list[int] | None = None, + mcp_server_ids: list[int] | None = None, ) -> None: - """Replace a task's ``kind`` pre-approval grants with ``target_ids`` - (deduped); grants of other kinds are preserved. Reuses existing grant rows so - re-submitting a granted target is a no-op — recreating it would orphan+reinsert - the same unique key in one flush, which Postgres rejects. Removed grants drop - via the ``delete-orphan`` cascade. + """Replace supplied target kinds and preserve omitted kinds. + + Reuse unchanged rows to avoid deleting and inserting the same unique key + in one flush. The orphan cascade deletes removed grants. """ - wanted_gated_app_ids = [ - get_or_create_gated_app_id(db_session, kind, target_id) - for target_id in dict.fromkeys(target_ids) + replacements = { + kind: target_ids + for kind, target_ids in ( + (GatedAppKind.EXTERNAL_APP, external_app_ids), + (GatedAppKind.MCP_SERVER, mcp_server_ids), + ) + if target_ids is not None + } + if not replacements: + return + + existing_by_target = { + grant.gated_app.target_key: grant for grant in task.pre_approved_targets + } + replacement_grants = [ + existing_by_target.get((kind, target_id)) + or ScheduledTaskPreApprovedTarget( + gated_app_id=get_or_create_gated_app_id(db_session, kind, target_id) + ) + for kind, target_ids in replacements.items() + for target_id in set(target_ids) ] - existing = {grant.gated_app_id: grant for grant in task.pre_approved_apps} - other_kind_grants = [ - grant for grant in task.pre_approved_apps if grant.gated_app.kind is not kind + retained_grants = [ + grant + for grant in task.pre_approved_targets + if grant.gated_app.kind not in replacements + ] + task.pre_approved_targets = [ + *replacement_grants, + *retained_grants, ] - task.pre_approved_apps = [ - existing.get(gated_app_id) - or ScheduledTaskPreApprovedApp(gated_app_id=gated_app_id) - for gated_app_id in wanted_gated_app_ids - ] + other_kind_grants def get_scheduled_task( @@ -175,6 +194,7 @@ def update_scheduled_task( editor_mode: EditorMode | None = None, status: ScheduledTaskStatus | None = None, pre_approved_external_app_ids: list[int] | None = None, + pre_approved_mcp_server_ids: list[int] | None = None, now: datetime | None = None, ) -> ScheduledTask: """Apply a partial update to a scheduled task. @@ -184,8 +204,8 @@ def update_scheduled_task( ``next_run_at`` is recomputed from ``now``. - If ``status`` transitions to PAUSED, ``next_run_at`` is set to NULL. - If ``status`` transitions to ACTIVE, ``next_run_at`` is recomputed. - - ``pre_approved_external_app_ids`` follows normal patch semantics: supplied - replaces the set, omitted leaves it unchanged. + - Each pre-approved target field follows normal patch semantics: supplied + replaces that target kind, and omitted leaves it unchanged. Raises: OnyxError(NOT_FOUND): the task does not exist or is not owned by @@ -200,10 +220,12 @@ def update_scheduled_task( task.name = name if prompt is not None: task.prompt = prompt - if pre_approved_external_app_ids is not None: - set_pre_approved_apps( - db_session, task, GatedAppKind.EXTERNAL_APP, pre_approved_external_app_ids - ) + _replace_pre_approved_targets( + db_session, + task, + external_app_ids=pre_approved_external_app_ids, + mcp_server_ids=pre_approved_mcp_server_ids, + ) if editor_mode is not None: task.editor_mode = editor_mode if cron_expression is not None and cron_expression != task.cron_expression: @@ -540,15 +562,15 @@ def get_live_scheduled_run_grants( if run is None: return None run_id, task_id = run - gated_apps = db_session.scalars( + gated_targets = db_session.scalars( select(GatedApp) .join( - ScheduledTaskPreApprovedApp, - ScheduledTaskPreApprovedApp.gated_app_id == GatedApp.id, + ScheduledTaskPreApprovedTarget, + ScheduledTaskPreApprovedTarget.gated_app_id == GatedApp.id, ) - .where(ScheduledTaskPreApprovedApp.scheduled_task_id == task_id) + .where(ScheduledTaskPreApprovedTarget.scheduled_task_id == task_id) ).all() - granted: set[GrantedTarget] = {ga.target_key for ga in gated_apps} + granted: set[GrantedTarget] = {target.target_key for target in gated_targets} return run_id, granted diff --git a/backend/onyx/server/features/build/scheduled_tasks/api.py b/backend/onyx/server/features/build/scheduled_tasks/api.py index 72e0dc7d0c..70be686429 100644 --- a/backend/onyx/server/features/build/scheduled_tasks/api.py +++ b/backend/onyx/server/features/build/scheduled_tasks/api.py @@ -14,7 +14,7 @@ The router is mounted under the existing ``/build`` prefix (see from __future__ import annotations from datetime import datetime, timezone -from typing import Any +from typing import AbstractSet, Any from uuid import UUID from fastapi import APIRouter, Depends, Query, Response @@ -32,6 +32,7 @@ from onyx.db.enums import ( ScheduledTaskTriggerSource, ) from onyx.db.external_app import get_external_apps +from onyx.db.mcp import get_craft_enabled_mcp_servers from onyx.db.models import ScheduledTask, ScheduledTaskRun, User from onyx.db.scheduled_task import ( create_scheduled_task, @@ -122,6 +123,7 @@ class ScheduledTaskCreate(_Forbid): status: ScheduledTaskStatus = ScheduledTaskStatus.ACTIVE run_immediately: bool = False pre_approved_app_ids: list[int] = Field(default_factory=list) + pre_approved_mcp_server_ids: list[int] = Field(default_factory=list) _dispatch = model_validator(mode="before")(_dispatch_editor_payload) @@ -139,6 +141,7 @@ class ScheduledTaskPatch(_Forbid): editor_payload: EditorPayload | None = None status: ScheduledTaskStatus | None = None pre_approved_app_ids: list[int] | None = None + pre_approved_mcp_server_ids: list[int] | None = None _dispatch = model_validator(mode="before")(_dispatch_editor_payload) @@ -210,6 +213,7 @@ class ScheduledTaskDetail(BaseModel): next_runs: list[datetime] last_run: RunSummary | None pre_approved_app_ids: list[int] + pre_approved_mcp_server_ids: list[int] created_at: datetime updated_at: datetime @@ -298,28 +302,53 @@ def _detail( next_runs=next_runs, last_run=RunSummary.from_model(last_run) if last_run is not None else None, pre_approved_app_ids=task.pre_approved_external_app_ids, + pre_approved_mcp_server_ids=task.pre_approved_mcp_server_ids, created_at=task.created_at, updated_at=task.updated_at, ) -def _validated_app_ids(db_session: Session, app_ids: list[int]) -> list[int]: - """Dedupe (order-preserving) and verify each id is a configured app. +def _validate_app_ids(db_session: Session, app_ids: list[int]) -> None: + """Reject unknown external app ids. Existence-only: a grant on an app that never produces an ASK match is inert, so credential / has-ASK-action filtering stays editor-side. """ - deduped = list(dict.fromkeys(app_ids)) - if not deduped: - return [] + if not app_ids: + return known_ids = {app.id for app in get_external_apps(db_session)} - unknown = [app_id for app_id in deduped if app_id not in known_ids] - if unknown: + unknown_ids = sorted(set(app_ids) - known_ids) + if unknown_ids: raise OnyxError( OnyxErrorCode.INVALID_INPUT, - f"Unknown external app id(s): {unknown}", + f"Unknown external app id(s): {unknown_ids}", + ) + + +def _validate_mcp_server_ids( + db_session: Session, + user: User, + server_ids: list[int], + already_approved_server_ids: AbstractSet[int] = frozenset(), +) -> None: + """Reject new MCP server ids unavailable to this user in Craft. + + Existing grants can remain after access changes. The MCP resolver still + blocks use, and users can remove the stale grant from the task editor. + """ + if not server_ids: + return + available_ids = { + server.id for server in get_craft_enabled_mcp_servers(db_session, user) + } + unavailable_ids = sorted( + set(server_ids) - available_ids - already_approved_server_ids + ) + if unavailable_ids: + raise OnyxError( + OnyxErrorCode.INVALID_INPUT, + f"Unknown or unavailable Craft MCP server id(s): {unavailable_ids}", ) - return deduped def _enqueue_executor(run_id: UUID) -> None: @@ -392,6 +421,8 @@ def create_task( and enqueue the executor. Does NOT touch ``next_run_at``. """ cron_expression = compile_to_cron(request.editor_payload) + _validate_app_ids(db_session, request.pre_approved_app_ids) + _validate_mcp_server_ids(db_session, user, request.pre_approved_mcp_server_ids) task = create_scheduled_task( db_session=db_session, @@ -401,9 +432,8 @@ def create_task( cron_expression=cron_expression, editor_mode=request.editor_mode, status=request.status, - pre_approved_external_app_ids=_validated_app_ids( - db_session, request.pre_approved_app_ids - ), + pre_approved_external_app_ids=request.pre_approved_app_ids, + pre_approved_mcp_server_ids=request.pre_approved_mcp_server_ids, ) if request.run_immediately: @@ -455,6 +485,22 @@ def patch_task( # whenever editor_payload is — no runtime check needed here. cron_expression = compile_to_cron(request.editor_payload) + if request.pre_approved_app_ids is not None: + _validate_app_ids(db_session, request.pre_approved_app_ids) + + if request.pre_approved_mcp_server_ids is not None: + existing_task = get_scheduled_task( + db_session=db_session, + task_id=task_id, + user_id=user.id, + ) + _validate_mcp_server_ids( + db_session, + user, + request.pre_approved_mcp_server_ids, + already_approved_server_ids=set(existing_task.pre_approved_mcp_server_ids), + ) + task = update_scheduled_task( db_session=db_session, task_id=task_id, @@ -464,11 +510,8 @@ def patch_task( cron_expression=cron_expression, editor_mode=request.editor_mode, status=request.status, - pre_approved_external_app_ids=( - _validated_app_ids(db_session, request.pre_approved_app_ids) - if request.pre_approved_app_ids is not None - else None - ), + pre_approved_external_app_ids=request.pre_approved_app_ids, + pre_approved_mcp_server_ids=request.pre_approved_mcp_server_ids, ) db_session.commit() db_session.refresh(task) diff --git a/backend/tests/external_dependency_unit/craft/test_scheduled_task_pre_approvals.py b/backend/tests/external_dependency_unit/craft/test_scheduled_task_pre_approvals.py index 7e934472a3..8a83ea5195 100644 --- a/backend/tests/external_dependency_unit/craft/test_scheduled_task_pre_approvals.py +++ b/backend/tests/external_dependency_unit/craft/test_scheduled_task_pre_approvals.py @@ -23,13 +23,12 @@ from onyx.db.enums import ( ScheduledTaskStatus, ScheduledTaskTriggerSource, ) -from onyx.db.gated_app import get_or_create_gated_app_id from onyx.db.models import ( BuildSession, ExternalApp, MCPServer, ScheduledTask, - ScheduledTaskPreApprovedApp, + ScheduledTaskPreApprovedTarget, User, ) from onyx.db.scheduled_task import ( @@ -58,11 +57,25 @@ def _make_app(db_session: Session) -> int: return app.id +def _make_mcp_server(db_session: Session, user: User) -> int: + server = MCPServer( + owner=user.email, + name=f"pre_approval_mcp_{uuid4().hex[:8]}", + server_url="https://example.com/mcp", + available_in_craft=True, + is_public=False, + ) + db_session.add(server) + db_session.flush() + return server.id + + def _seed_task( db_session: Session, user: User, *, pre_approved_external_app_ids: list[int] | None = None, + pre_approved_mcp_server_ids: list[int] | None = None, prompt: str = "Summarise yesterday's events", ) -> ScheduledTask: task = create_scheduled_task( @@ -74,6 +87,7 @@ def _seed_task( editor_mode="advanced", status=ScheduledTaskStatus.ACTIVE, pre_approved_external_app_ids=pre_approved_external_app_ids, + pre_approved_mcp_server_ids=pre_approved_mcp_server_ids, ) db_session.commit() db_session.refresh(task) @@ -93,7 +107,13 @@ def test_grants_returned_for_running_run( user = make_user(db_session) bs = build_session_with_user(user=user) app_a, app_b = _make_app(db_session), _make_app(db_session) - task = _seed_task(db_session, user, pre_approved_external_app_ids=[app_a, app_b]) + server_id = _make_mcp_server(db_session, user) + task = _seed_task( + db_session, + user, + pre_approved_external_app_ids=[app_a, app_b], + pre_approved_mcp_server_ids=[server_id], + ) run = insert_run( db_session=db_session, task_id=task.id, @@ -115,6 +135,7 @@ def test_grants_returned_for_running_run( assert granted == { (GatedAppKind.EXTERNAL_APP, app_a), (GatedAppKind.EXTERNAL_APP, app_b), + (GatedAppKind.MCP_SERVER, server_id), } @@ -306,33 +327,18 @@ def test_mcp_grants_survive_external_app_replacement( tenant_context: None, # noqa: ARG001 build_session_with_user: Callable[..., BuildSession], ) -> None: - """``set_pre_approved_apps`` replaces only the given kind's grants: an MCP-server - grant (seeded directly — no API writes these yet) survives a wholesale - external-app replacement, stays out of ``pre_approved_external_app_ids``, and reaches - the gate through ``get_live_scheduled_run_grants`` as its (kind, id) target. - """ + """Each target kind has independent replacement semantics and reaches the gate.""" user = make_user(db_session) bs = build_session_with_user(user=user) app_a, app_b = _make_app(db_session), _make_app(db_session) - task = _seed_task(db_session, user, pre_approved_external_app_ids=[app_a]) - - server = MCPServer( - owner=user.email, - name=f"pre_approval_mcp_{uuid4().hex[:8]}", - server_url="https://example.com/mcp", - is_public=False, + server_a = _make_mcp_server(db_session, user) + server_b = _make_mcp_server(db_session, user) + task = _seed_task( + db_session, + user, + pre_approved_external_app_ids=[app_a], + pre_approved_mcp_server_ids=[server_a], ) - db_session.add(server) - db_session.flush() - mcp_gated_app_id = get_or_create_gated_app_id( - db_session, GatedAppKind.MCP_SERVER, server.id - ) - db_session.add( - ScheduledTaskPreApprovedApp( - scheduled_task_id=task.id, gated_app_id=mcp_gated_app_id - ) - ) - db_session.commit() updated = update_scheduled_task( db_session=db_session, @@ -343,10 +349,18 @@ def test_mcp_grants_survive_external_app_replacement( db_session.commit() assert updated.pre_approved_external_app_ids == [app_b] # MCP grant excluded - assert {g.gated_app.target_key for g in updated.pre_approved_apps} == { - (GatedAppKind.EXTERNAL_APP, app_b), - (GatedAppKind.MCP_SERVER, server.id), - } + assert updated.pre_approved_mcp_server_ids == [server_a] + + updated = update_scheduled_task( + db_session=db_session, + task_id=task.id, + user_id=user.id, + pre_approved_mcp_server_ids=[server_b], + ) + db_session.commit() + + assert updated.pre_approved_external_app_ids == [app_b] + assert updated.pre_approved_mcp_server_ids == [server_b] run = insert_run( db_session=db_session, @@ -365,23 +379,31 @@ def test_mcp_grants_survive_external_app_replacement( assert grants is not None assert grants[1] == { (GatedAppKind.EXTERNAL_APP, app_b), - (GatedAppKind.MCP_SERVER, server.id), + (GatedAppKind.MCP_SERVER, server_b), } -def test_create_persists_grants( +def test_create_stores_each_grant_once( db_session: Session, tenant_context: None, # noqa: ARG001 ) -> None: user = make_user(db_session) app_a, app_b = _make_app(db_session), _make_app(db_session) - assert app_a < app_b # ids autoincrement, so the higher id is created last - # Insertion order is preserved (not sorted): pass the higher id first. - task = _seed_task(db_session, user, pre_approved_external_app_ids=[app_b, app_a]) - assert task.pre_approved_external_app_ids == [app_b, app_a] + server_id = _make_mcp_server(db_session, user) + task = _seed_task( + db_session, + user, + pre_approved_external_app_ids=[app_b, app_a, app_b], + pre_approved_mcp_server_ids=[server_id, server_id], + ) + assert set(task.pre_approved_external_app_ids) == {app_a, app_b} + assert len(task.pre_approved_external_app_ids) == 2 + assert set(task.pre_approved_mcp_server_ids) == {server_id} + assert len(task.pre_approved_mcp_server_ids) == 1 bare = _seed_task(db_session, user) assert bare.pre_approved_external_app_ids == [] + assert bare.pre_approved_mcp_server_ids == [] # --------------------------------------------------------------------------- @@ -405,8 +427,8 @@ def test_deleting_app_drops_grants( remaining = ( db_session.execute( - select(ScheduledTaskPreApprovedApp).where( - ScheduledTaskPreApprovedApp.scheduled_task_id == task.id + select(ScheduledTaskPreApprovedTarget).where( + ScheduledTaskPreApprovedTarget.scheduled_task_id == task.id ) ) .scalars() @@ -451,14 +473,12 @@ def test_deleting_app_nulls_action_approval_fk( # --------------------------------------------------------------------------- -def test_validated_app_ids_rejects_unknown_and_dedupes( +def test_validate_app_ids_accepts_known_duplicates_and_reports_unknown_ids_once( db_session: Session, tenant_context: None, # noqa: ARG001 monkeypatch: pytest.MonkeyPatch, ) -> None: - """Dedupe is order-preserving; any id outside the tenant's apps raises - INVALID_INPUT. Apps are stubbed — only the validation logic is under - test, not ``get_external_apps``'s SQL.""" + """Duplicates are valid, but unknown ids raise INVALID_INPUT.""" class _App: def __init__(self, app_id: int) -> None: @@ -470,9 +490,10 @@ def test_validated_app_ids_rejects_unknown_and_dedupes( lambda _db: [_App(7), _App(9)], ) - assert scheduled_tasks_api._validated_app_ids(db_session, []) == [] - assert scheduled_tasks_api._validated_app_ids(db_session, [9, 7, 9]) == [9, 7] + scheduled_tasks_api._validate_app_ids(db_session, []) + scheduled_tasks_api._validate_app_ids(db_session, [9, 7, 9]) with pytest.raises(OnyxError) as exc_info: - scheduled_tasks_api._validated_app_ids(db_session, [7, 123]) + scheduled_tasks_api._validate_app_ids(db_session, [456, 7, 123, 456]) assert exc_info.value.error_code == OnyxErrorCode.INVALID_INPUT + assert exc_info.value.detail == "Unknown external app id(s): [123, 456]" diff --git a/backend/tests/integration/tests/craft/test_scheduled_tasks_api.py b/backend/tests/integration/tests/craft/test_scheduled_tasks_api.py index 5bcdbeeb45..21f6c6d77f 100644 --- a/backend/tests/integration/tests/craft/test_scheduled_tasks_api.py +++ b/backend/tests/integration/tests/craft/test_scheduled_tasks_api.py @@ -17,6 +17,7 @@ from onyx.db.enums import ( ScheduledTaskStatus, ScheduledTaskTriggerSource, ) +from onyx.db.mcp import create_mcp_server__no_commit, update_mcp_server__no_commit from onyx.db.models import ScheduledTask, ScheduledTaskRun from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR @@ -51,6 +52,7 @@ def _create_task( editor_payload: dict[str, Any] | None = None, status: ScheduledTaskStatus = ScheduledTaskStatus.ACTIVE, run_immediately: bool = False, + pre_approved_mcp_server_ids: list[int] | None = None, ) -> httpx.Response: body: dict[str, Any] = { "name": name or f"task-{uuid4().hex[:8]}", @@ -60,6 +62,8 @@ def _create_task( "status": status.value, "run_immediately": run_immediately, } + if pre_approved_mcp_server_ids is not None: + body["pre_approved_mcp_server_ids"] = pre_approved_mcp_server_ids return client.post( _url(), json=body, @@ -135,6 +139,28 @@ def _get_runs_for_task(task_id: UUID) -> list[ScheduledTaskRun]: ) +def _create_mcp_server(owner: str, *, available_in_craft: bool) -> int: + with get_session_with_current_tenant() as db_session: + server = create_mcp_server__no_commit( + owner_email=owner, + name=f"scheduled-task-mcp-{uuid4().hex[:8]}", + description=None, + server_url="https://example.com/mcp", + auth_type=None, + transport=None, + auth_performer=None, + db_session=db_session, + is_public=True, + ) + update_mcp_server__no_commit( + server_id=server.id, + db_session=db_session, + available_in_craft=available_in_craft, + ) + db_session.commit() + return server.id + + def test_create_task_compiles_cron(admin_user: DATestUser) -> None: response = _create_task( admin_user, @@ -152,6 +178,82 @@ def test_create_task_compiles_cron(admin_user: DATestUser) -> None: assert row.cron_expression == cron +def test_create_and_patch_mcp_pre_approvals(admin_user: DATestUser) -> None: + server_a = _create_mcp_server(admin_user.email, available_in_craft=True) + server_b = _create_mcp_server(admin_user.email, available_in_craft=True) + unavailable_server = _create_mcp_server(admin_user.email, available_in_craft=False) + + create_response = _create_task( + admin_user, + pre_approved_mcp_server_ids=[server_b, server_a, server_b], + ) + create_response.raise_for_status() + created = create_response.json() + task_id = UUID(created["id"]) + created_mcp_server_ids = created["pre_approved_mcp_server_ids"] + assert len(created_mcp_server_ids) == 2 + assert set(created_mcp_server_ids) == {server_a, server_b} + + patch_response = _patch_task( + admin_user, + task_id, + {"pre_approved_mcp_server_ids": [server_a]}, + ) + patch_response.raise_for_status() + assert patch_response.json()["pre_approved_mcp_server_ids"] == [server_a] + + rejected_response = _patch_task( + admin_user, + task_id, + {"pre_approved_mcp_server_ids": [unavailable_server]}, + ) + assert rejected_response.status_code == 400 + + detail_response = client.get( + _url(str(task_id)), + headers=admin_user.headers, + cookies=admin_user.cookies, + ) + detail_response.raise_for_status() + assert detail_response.json()["pre_approved_mcp_server_ids"] == [server_a] + + +def test_patch_retains_and_removes_existing_unavailable_mcp_pre_approval( + admin_user: DATestUser, +) -> None: + server_id = _create_mcp_server(admin_user.email, available_in_craft=True) + create_response = _create_task( + admin_user, + pre_approved_mcp_server_ids=[server_id], + ) + create_response.raise_for_status() + task_id = UUID(create_response.json()["id"]) + + with get_session_with_current_tenant() as db_session: + update_mcp_server__no_commit( + server_id=server_id, + db_session=db_session, + available_in_craft=False, + ) + db_session.commit() + + retain_response = _patch_task( + admin_user, + task_id, + {"pre_approved_mcp_server_ids": [server_id]}, + ) + retain_response.raise_for_status() + assert retain_response.json()["pre_approved_mcp_server_ids"] == [server_id] + + remove_response = _patch_task( + admin_user, + task_id, + {"pre_approved_mcp_server_ids": []}, + ) + remove_response.raise_for_status() + assert remove_response.json()["pre_approved_mcp_server_ids"] == [] + + def test_create_task_interval_days_requires_time_of_day(admin_user: DATestUser) -> None: response = _create_task( admin_user, diff --git a/backend/tests/unit/sandbox_proxy/test_gate.py b/backend/tests/unit/sandbox_proxy/test_gate.py index 28e99bfdb5..fbf85ae55f 100644 --- a/backend/tests/unit/sandbox_proxy/test_gate.py +++ b/backend/tests/unit/sandbox_proxy/test_gate.py @@ -154,6 +154,15 @@ def _assert_403(flow: http.HTTPFlow, expected_code: SandboxProxyError) -> None: _MATCH = make_matched_actions(payload={"text": "hi"}) +_MATCH_MCP = AllMatchedActions( + actions=_MATCH.actions, + target=GatedTarget( + kind=GatedAppKind.MCP_SERVER, + id=73, + app_name="Linear MCP", + ), + payload=_MATCH.payload, +) _MATCH_MULTI_ASK = AllMatchedActions( actions=( MatchedAction( @@ -561,16 +570,19 @@ async def test_ask_denied_blocks( _RUN_ID = UUID("55555555-5555-5555-5555-555555555555") # make_matched_actions defaults external_app_id=42. _GRANTED_APP_ID = 42 +_GRANTED_MCP_SERVER_ID = 73 def _stub_grants( monkeypatch: pytest.MonkeyPatch, result: tuple[UUID, list[int]] | None | Exception, + *, + kind: GatedAppKind = GatedAppKind.EXTERNAL_APP, ) -> list[UUID]: """Stub the gate's grant lookup; returns the recorded session_ids. - Callers pass granted external-app ids; the stub wraps them as the - ``(kind, id)`` targets the real lookup now returns.""" + The stub wraps target ids as the ``(kind, id)`` keys returned by the real + lookup.""" calls: list[UUID] = [] def _lookup( @@ -583,8 +595,8 @@ def _stub_grants( raise result if result is None: return None - run_id, app_ids = result - return run_id, {(GatedAppKind.EXTERNAL_APP, app_id) for app_id in app_ids} + run_id, target_ids = result + return run_id, {(kind, target_id) for target_id in target_ids} monkeypatch.setattr(gate, "get_live_scheduled_run_grants", _lookup) return calls @@ -643,17 +655,27 @@ class _SessionGrantCache: @pytest.mark.asyncio +@pytest.mark.parametrize( + ("matched_actions", "target_kind", "target_id"), + [ + (_MATCH, GatedAppKind.EXTERNAL_APP, _GRANTED_APP_ID), + (_MATCH_MCP, GatedAppKind.MCP_SERVER, _GRANTED_MCP_SERVER_ID), + ], + ids=["external-app", "mcp-server"], +) async def test_pre_approved_scheduled_run_skips_park( monkeypatch: pytest.MonkeyPatch, + matched_actions: AllMatchedActions, + target_kind: GatedAppKind, + target_id: int, ) -> None: - """Granted app on a RUNNING scheduled run: forwarded immediately with a - pre-decided APPROVED row — the park pipeline never runs.""" + """A granted target skips the approval park during a scheduled run.""" user_id = uuid4() sandbox = make_resolved_sandbox(user_id=user_id) resolver = StubResolver(sandbox=sandbox, session_by_id=UUID(_TAG_UUID)) - addon = _build(resolver=resolver, matcher=_StubMatcher(result=_MATCH)) + addon = _build(resolver=resolver, matcher=_StubMatcher(result=matched_actions)) spy = _spy_pipeline(addon, monkeypatch) - _stub_grants(monkeypatch, (_RUN_ID, [_GRANTED_APP_ID])) + _stub_grants(monkeypatch, (_RUN_ID, [target_id]), kind=target_kind) inserted = _spy_pre_approve_insert(monkeypatch) notified: list[dict[str, Any]] = [] monkeypatch.setattr(gate, "create_notification", lambda **kw: notified.append(kw)) @@ -664,17 +686,17 @@ async def test_pre_approved_scheduled_run_skips_park( assert flow.response is None # forwarded assert not spy.approval_ran # park pipeline skipped assert spy.awaited == [] - assert spy.dispatched == [(_MATCH, user_id, sandbox.tenant_id)] + assert spy.dispatched == [(matched_actions, user_id, sandbox.tenant_id)] assert len(inserted) == 1 assert inserted[0]["decision"] == ApprovalDecision.APPROVED assert inserted[0]["decided_via"] == ApprovalDecidedVia.PRE_APPROVAL - assert inserted[0]["target"] == (GatedAppKind.EXTERNAL_APP, _GRANTED_APP_ID) - # Dedup contract: additional_data is exactly the stable (run, app) pair. + assert inserted[0]["target"] == (target_kind, target_id) + # Dedup contract: additional_data is exactly the stable run and target. assert len(notified) == 1 assert notified[0]["additional_data"] == { "run_id": str(_RUN_ID), - "target_kind": GatedAppKind.EXTERNAL_APP.value, - "target_id": _GRANTED_APP_ID, + "target_kind": target_kind.value, + "target_id": target_id, } diff --git a/docs/craft/features/egress-proxy-and-approvals/README.md b/docs/craft/features/egress-proxy-and-approvals/README.md index a0d2c4ae59..03dc6c9d5e 100644 --- a/docs/craft/features/egress-proxy-and-approvals/README.md +++ b/docs/craft/features/egress-proxy-and-approvals/README.md @@ -42,7 +42,7 @@ Proxy runtime: Approval persistence and API: - `backend/onyx/db/models.py` defines `ActionApproval`, - `ExternalAppPolicy`, and `ScheduledTaskPreApprovedApp`. + `ExternalAppPolicy`, and `ScheduledTaskPreApprovedTarget`. - `backend/onyx/db/enums.py` defines `EndpointPolicy`, `ApprovalDecision`, and `ApprovalDecidedVia`. - `backend/onyx/server/features/build/db/action_approval.py` owns approval DB @@ -451,8 +451,8 @@ Partial coverage still parks. ### Scheduled Task Pre-Approvals Scheduled task pre-approvals are another grant source inside the same proxy -approval path. A task can store a set of pre-approved external app ids in -`scheduled_task_pre_approved_app`. +approval path. A task can store external-app and MCP-server grants in the +legacy-named `scheduled_task_pre_approved_app` table. For an `ASK` request, the gate checks grant sources in this order: @@ -463,7 +463,7 @@ A scheduled-task grant applies only when: - The `BuildSession` has a `ScheduledTaskRun` row. - That run is currently `RUNNING`. -- The task has a grant for the matched `external_app_id`. +- The task has a grant for the matched `(target kind, target id)`. When it applies, the proxy inserts an `action_approval` row already `APPROVED` with `decided_via=PRE_APPROVAL`, emits a deduped scheduled-task diff --git a/docs/craft/features/scheduled-tasks/pre-approvals.md b/docs/craft/features/scheduled-tasks/pre-approvals.md index 833bc7a98b..392ef5c323 100644 --- a/docs/craft/features/scheduled-tasks/pre-approvals.md +++ b/docs/craft/features/scheduled-tasks/pre-approvals.md @@ -125,20 +125,22 @@ unguarded after an `APPROVED` row is already committed. ## Data Model -New table `scheduled_task_pre_approved_app` — one row per `(task, app)` -grant: +The legacy-named `scheduled_task_pre_approved_app` table stores one row per +`(task, target)` grant: -- `scheduled_task_id` → `scheduled_task.id` (`ON DELETE CASCADE`) and - `external_app_id` → `external_app.id` (`ON DELETE CASCADE`), with a - `UNIQUE(scheduled_task_id, external_app_id)` constraint that keeps - grants idempotent and serves the per-task lookup. The FKs give real - referential integrity — a grant can't point at a removed app, and - removing either side drops the grant. `ScheduledTask.pre_approved_apps` - is the ORM collection; `pre_approved_app_ids` is a read-only accessor - over it, so the API contract (`list[int]`) is unchanged. The write - path replaces the whole set (`set_pre_approved_apps`), validated - against the configured apps (via the tenant-scoped session) and - deduped order-preserving. +- `scheduled_task_id` references `scheduled_task.id` with `ON DELETE CASCADE`. +- `gated_app_id` references the polymorphic `gated_app.id` with + `ON DELETE CASCADE`. Each `gated_app` row identifies exactly one external app + or MCP server. +- `UNIQUE(scheduled_task_id, gated_app_id)` keeps grants idempotent and serves + the per-task lookup. +- `ScheduledTask.pre_approved_targets` is the ORM collection of + `ScheduledTaskPreApprovedTarget` rows. The + `pre_approved_external_app_ids` and `pre_approved_mcp_server_ids` properties + project each target kind into its API field. +- `_replace_pre_approved_targets` replaces each supplied target kind and + preserves omitted kinds. It deduplicates submitted IDs and reuses unchanged + rows. - `action_approval.decided_via` — nullable (`user | pre_approval`, NULL for legacy/expired rows): the audit marker distinguishing a human click from a pre-approval. Kept separate from `decision` so @@ -160,10 +162,10 @@ pre-decided inserts go through `insert_action_approval` in - `ScheduledTaskCreate` / `ScheduledTaskPatch` gain `pre_approved_app_ids: list[int]`; `ScheduledTaskDetail` returns it. - The write path validates ids via `_validated_app_ids` and dedupes - (order-preserving) — existence only; a credential / ≥1-`ASK` filter is - editor-side advisory, since a grant on a no-`ASK` app is inert and - never consulted. + The write path validates ids via `_validate_app_ids` and stores each grant + once. Grant order has no meaning. Validation checks existence only; a + credential / ≥1-`ASK` filter is editor-side advisory because a grant on a + no-`ASK` app is inert and never consulted. - New `NotificationType.SCHEDULED_TASK_PRE_APPROVED_ACTION`, emitted per `(run, app)` on the first unattended forward so chatty tasks don't flood the bell. Dedup rides `create_notification`'s existing @@ -254,7 +256,7 @@ The grant-source seam means future modes drop in as new - Grant patch semantics: a prompt edit preserves grants, supplied `pre_approved_app_ids` replaces the set, and re-submitting an existing grant is idempotent (no unique-key collision). - - Create persistence + `_validated_app_ids` dedupe and unknown-id + - Create persistence, duplicate grant normalization, and unknown-id rejection. - **Unit** (gate, stubbed DB): `backend/tests/unit/sandbox_proxy/test_gate.py`