diff --git a/.agents/feature-map/components/auth-and-identity.md b/.agents/feature-map/components/auth-and-identity.md index b99bf5e31d..c492dd5b71 100644 --- a/.agents/feature-map/components/auth-and-identity.md +++ b/.agents/feature-map/components/auth-and-identity.md @@ -232,6 +232,11 @@ Custom-action OAuth lets each user connect a third-party account to an OpenAPI a `supports_pkce` defaults to false; administrators can enable S256 PKCE for compatible providers. MCP stores its credentials separately in `MCPConnectionConfig` ([[mcp-and-custom-tools]] §3). +`GET /user-oauth-token/status` reports literal access-token expiration and saved +refresh capability separately. `can_refresh` requires a nonempty refresh token +and configured client credentials. Status reads never refresh tokens or call providers. +Refresh capability does not prove that the provider will accept a refresh request. + --- ## 4. How it works diff --git a/.agents/feature-map/components/craft-external-apps.md b/.agents/feature-map/components/craft-external-apps.md index 16456de66d..2b93cd5074 100644 --- a/.agents/feature-map/components/craft-external-apps.md +++ b/.agents/feature-map/components/craft-external-apps.md @@ -68,6 +68,12 @@ sync (`server.py:_Readiness`). ### User-facing HTTP endpoints (outside the proxy) +Member app responses include `service_urls`: unique literal HTTP(S) origins for display. +The projection strips credentials, paths, queries, and default ports. Built-in +escaped hostname dots and hyphens are decoded. Unescaped regex dots and ambiguous +regex or wildcard hosts are omitted. +Raw upstream patterns, authentication templates, and organization credentials remain admin-only. + | Method | Path | Handler | Notes | |---|---|---|---| | GET | `/api/build/apps/{id}/oauth/start` | `start_external_app_oauth` (`server/features/build/external_apps/oauth.py`) | Begins a user's per-user OAuth connect flow for a built-in app. | @@ -428,6 +434,10 @@ identical in shape whether a human or a grant decided. ### 4.6 OAuth connect and token refresh +OAuth start accepts an optional safe internal `return_path`. The authorization attempt +stores it and the callback returns it after saving credentials. Omitting the path +retains the existing callback behavior. External and unsafe paths are rejected. + Connect: `GET /api/build/apps/{id}/oauth/start` builds the provider's authorize URL through the shared OAuth builder (`server/features/build/external_apps/oauth.py`). Each attempt is bound to the user, tenant, app, and OAuth configuration. diff --git a/.agents/feature-map/components/craft-webapp-proxy.md b/.agents/feature-map/components/craft-webapp-proxy.md index 73c958cb29..b0f971e1a8 100644 --- a/.agents/feature-map/components/craft-webapp-proxy.md +++ b/.agents/feature-map/components/craft-webapp-proxy.md @@ -306,6 +306,11 @@ Installation does not replace agent-managed skills or enter session snapshots. ## 5. Contracts and invariants +MCP host attribution remains workspace-wide when resolving targets without a +user. A personal Craft opt-out must not make a configured host look off-catalog. +The live resolver checks the current user's saved MCP preference before injecting +credentials, including no-auth servers. Disabling a personal default keeps credentials. + 1. **A preview must only be reachable by someone entitled to that session.** `_check_webapp_access` must run before every proxied HTTP request and before the HMR websocket is accepted. A new route added to this router diff --git a/.agents/feature-map/components/mcp-and-custom-tools.md b/.agents/feature-map/components/mcp-and-custom-tools.md index c9d12b78ee..4e290fe1b9 100644 --- a/.agents/feature-map/components/mcp-and-custom-tools.md +++ b/.agents/feature-map/components/mcp-and-custom-tools.md @@ -671,3 +671,38 @@ See `backend/AGENTS.md` for authoritative env/secrets setup. Test helpers: a client that tries to shortcut the flow by POSTing arbitrary `code`/`server_id` values without a matching stored attempt gets rejected at the fingerprint/connection-config-id checks (§4.4 step 3). + + +### Plugin API and personal defaults + +Member MCP listings stamp `user_enabled_in_chat` and `user_enabled_in_craft`. +Member tool snapshots stamp `user_enabled_in_chat`. Missing preference rows mean enabled. +`default_chat_enabled` reports default-assistant attachments independently of credentials. +The summaries are computed in batches after existing visibility filters. + +`PATCH /mcp/server/{server_id}/preferences` saves independent Chat and Craft preferences. +`PATCH /tool/{tool_id}/preferences` saves an OpenAPI Chat preference. +These operations require basic access and cannot enable unavailable workspace capabilities. +Preferences retain credentials. Default Chat enforces opt-outs before credential resolution. +Custom agents retain their attached tools and existing access behavior. + +Craft configuration excludes personally disabled MCP servers. Its member catalog retains them +so a caller can enable them again. Updates reload that caller's Craft sessions. +The live proxy resolver also enforces Craft opt-outs, including no-auth servers. +Host attribution remains workspace-wide so disabled servers cannot become off-catalog. + +`POST /admin/mcp/probe` checks MCP setup without credentials. +The stream uses newline-delimited JSON progress and one final result. +`POST /admin/mcp/oauth/discover` reads public OAuth metadata for a manageable server. +Both require `MANAGE_ACTIONS`, including scoped managers. +Discovery checks management access before network work and closes its database session first. +SSRF validation covers outbound hops and runs outside the event loop. +OAuth failures return safe messages and an `oauth_error` reason. + +Simple server creation accepts `available_in_craft`, defaulting to false. +Changing a server URL invalidates existing user and administrator credentials, +resets setup status, and reloads affected Craft sessions. + +Preferences use `MCPServerUserPreference` and `ToolUserPreference`. +Apply revisions `b71e3a9d2c04` and `c82f4a0e3d15` before deploying code that +reads the preference tables. diff --git a/.agents/feature-map/components/tools-framework.md b/.agents/feature-map/components/tools-framework.md index 0486b92b69..413b590bf6 100644 --- a/.agents/feature-map/components/tools-framework.md +++ b/.agents/feature-map/components/tools-framework.md @@ -257,6 +257,12 @@ tool calls the LLM just made and the constructed `Tool` instances. ## 5. Contracts and invariants +Personal MCP and OpenAPI Chat opt-outs are enforced only for `DEFAULT_PERSONA_ID`, +before credential resolution or MCP cache reuse. An explicit `allowed_tool_ids` +list cannot override them. Custom agents retain their own attached tools. +Preferences are batch-read once per construction for configured MCP servers +and OpenAPI tool IDs. Missing preferences preserve the existing default-on behavior. + 1. **A new built-in tool needs both `BUILT_IN_TOOL_MAP` and a DB row.** Adding the class alone does nothing; `_construct_tools_impl` only sees tools present on `persona.tools`, and those rows come from `Tool.in_code_tool_id` seeded via @@ -422,3 +428,11 @@ See `backend/AGENTS.md` for authoritative commands and required env. `_merge_tool_calls` collapses N calls to `SearchTool`/`WebSearchTool`/`OpenURLTool` into one before execution and before it enters history, so the model later sees the merged, expanded arguments rather than what it originally emitted. + + +### Default Chat MCP availability + +For the default persona, `tool_constructor.py` checks both server availability in +Chat and the user's MCP server access before resolving credentials. This prevents +a group-restricted server from becoming available to all users through default +Chat. Custom agents retain their existing attached-tool access behavior. diff --git a/backend/alembic/versions/b71e3a9d2c04_add_personal_mcp_server_preferences.py b/backend/alembic/versions/b71e3a9d2c04_add_personal_mcp_server_preferences.py new file mode 100644 index 0000000000..f6945851c5 --- /dev/null +++ b/backend/alembic/versions/b71e3a9d2c04_add_personal_mcp_server_preferences.py @@ -0,0 +1,46 @@ +"""add personal MCP server preferences + +Revision ID: b71e3a9d2c04 +Revises: b034d72b7abe +Create Date: 2026-10-15 19:25:16.614907 + +""" + +from alembic import op +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision = "b71e3a9d2c04" +down_revision = "b034d72b7abe" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + op.create_table( + "mcp_server_user_preference", + sa.Column("user_id", sa.Uuid(), nullable=False), + sa.Column("mcp_server_id", sa.Integer(), nullable=False), + sa.Column( + "enabled_in_chat", + sa.Boolean(), + server_default=sa.text("true"), + nullable=False, + ), + sa.Column( + "enabled_in_craft", + sa.Boolean(), + server_default=sa.text("true"), + nullable=False, + ), + sa.ForeignKeyConstraint(["user_id"], ["user.id"], ondelete="CASCADE"), + sa.ForeignKeyConstraint( + ["mcp_server_id"], ["mcp_server.id"], ondelete="CASCADE" + ), + sa.PrimaryKeyConstraint("user_id", "mcp_server_id"), + ) + + +def downgrade() -> None: + op.drop_table("mcp_server_user_preference") diff --git a/backend/alembic/versions/c82f4a0e3d15_add_personal_openapi_preferences.py b/backend/alembic/versions/c82f4a0e3d15_add_personal_openapi_preferences.py new file mode 100644 index 0000000000..72b3a23855 --- /dev/null +++ b/backend/alembic/versions/c82f4a0e3d15_add_personal_openapi_preferences.py @@ -0,0 +1,34 @@ +"""add personal OpenAPI preferences + +Revision ID: c82f4a0e3d15 +Revises: b71e3a9d2c04 +""" + +from alembic import op +import sqlalchemy as sa + +revision = "c82f4a0e3d15" +down_revision = "b71e3a9d2c04" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + op.create_table( + "tool_user_preference", + sa.Column("user_id", sa.Uuid(), nullable=False), + sa.Column("tool_id", sa.Integer(), nullable=False), + sa.Column( + "enabled_in_chat", + sa.Boolean(), + server_default=sa.text("true"), + nullable=False, + ), + sa.ForeignKeyConstraint(["user_id"], ["user.id"], ondelete="CASCADE"), + sa.ForeignKeyConstraint(["tool_id"], ["tool.id"], ondelete="CASCADE"), + sa.PrimaryKeyConstraint("user_id", "tool_id"), + ) + + +def downgrade() -> None: + op.drop_table("tool_user_preference") diff --git a/backend/onyx/db/mcp.py b/backend/onyx/db/mcp.py index 13f78679dc..2fd0c07e74 100644 --- a/backend/onyx/db/mcp.py +++ b/backend/onyx/db/mcp.py @@ -2,10 +2,12 @@ import datetime from uuid import UUID from sqlalchemy import Select, and_, delete, select +from sqlalchemy.dialects.postgresql import insert from sqlalchemy.orm import Session, aliased, selectinload from sqlalchemy.orm.attributes import flag_modified from onyx.auth.permissions import has_global_permission +from onyx.configs.constants import DEFAULT_PERSONA_ID from onyx.db.constants import UNSET, UnsetType from onyx.db.enums import ( MCPAuthenticationPerformer, @@ -21,7 +23,9 @@ from onyx.db.models import ( MCPServer, MCPServer__User, MCPServer__UserGroup, + MCPServerUserPreference, Persona, + Persona__Tool, Sandbox, Tool, User, @@ -34,7 +38,75 @@ from onyx.utils.logger import setup_logger logger = setup_logger() +def get_mcp_server_user_preferences( + server_ids: list[int], user_id: UUID, db_session: Session +) -> dict[int, MCPServerUserPreference]: + if not server_ids: + return {} + return { + preference.mcp_server_id: preference + for preference in db_session.scalars( + select(MCPServerUserPreference).where( + MCPServerUserPreference.user_id == user_id, + MCPServerUserPreference.mcp_server_id.in_(server_ids), + ) + ) + } + + +def upsert_mcp_server_user_preferences__no_commit( + server_id: int, + user_id: UUID, + db_session: Session, + *, + enabled_in_chat: bool | None = None, + enabled_in_craft: bool | None = None, +) -> MCPServerUserPreference: + values: dict[str, bool] = {} + if enabled_in_chat is not None: + values["enabled_in_chat"] = enabled_in_chat + if enabled_in_craft is not None: + values["enabled_in_craft"] = enabled_in_craft + if not values: + raise ValueError("At least one preference is required") + stmt = insert(MCPServerUserPreference).values( + user_id=user_id, mcp_server_id=server_id, **values + ) + return db_session.scalars( + stmt.on_conflict_do_update( + index_elements=[ + MCPServerUserPreference.user_id, + MCPServerUserPreference.mcp_server_id, + ], + set_=values, + ).returning(MCPServerUserPreference), + execution_options={"populate_existing": True}, + ).one() + + # MCPServer operations +def get_default_chat_mcp_server_ids( + server_ids: list[int], db_session: Session +) -> set[int]: + if not server_ids: + return set() + stmt = ( + select(Tool.mcp_server_id) + .join(Persona__Tool, Persona__Tool.tool_id == Tool.id) + .join(Persona, Persona.id == Persona__Tool.persona_id) + .where( + Persona.id == DEFAULT_PERSONA_ID, + Persona.builtin_persona.is_(True), + Persona.deleted.is_(False), + Tool.mcp_server_id.in_(server_ids), + ) + .distinct() + ) + return { + server_id for server_id in db_session.scalars(stmt) if server_id is not None + } + + def get_all_mcp_servers(db_session: Session) -> list[MCPServer]: """Get all MCP servers""" return list( @@ -60,7 +132,7 @@ def get_mcp_servers_by_owner(owner_email: str, db_session: Session) -> list[MCPS def get_craft_enabled_mcp_servers( - db_session: Session, user: User | None + db_session: Session, user: User | None, *, include_personal_disabled: bool = False ) -> list[MCPServer]: """MCP servers an admin has made available to the Craft agent, filtered to those ``user`` may use (public / shared / owned). ``None`` skips the access @@ -75,16 +147,27 @@ def get_craft_enabled_mcp_servers( ) if user is not None: stmt = _add_mcp_server_access_filter(stmt, user) + if not include_personal_disabled: + stmt = stmt.where( + ~select(MCPServerUserPreference.mcp_server_id) + .where( + MCPServerUserPreference.user_id == user.id, + MCPServerUserPreference.mcp_server_id == MCPServer.id, + MCPServerUserPreference.enabled_in_craft.is_(False), + ) + .exists() + ) return list(db_session.scalars(stmt).all()) def get_mcp_servers_for_persona( persona_id: int, db_session: Session, - user: User, # noqa: ARG001 + user: User, ) -> list[MCPServer]: - """Servers already on a persona's tools. No attach ACL — chat users of the - persona must see/auth these. ``user`` is for callers enforcing persona visibility. + """Attached servers. Default Chat also checks server access. + + Custom agents retain their attached-server access for all agent users. """ # Get the persona and its tools persona = db_session.query(Persona).filter(Persona.id == persona_id).first() @@ -101,16 +184,15 @@ def get_mcp_servers_for_persona( return [] # Fetch the MCP servers - mcp_servers = ( - db_session.query(MCPServer).filter(MCPServer.id.in_(mcp_server_ids)).all() - ) - - return list(mcp_servers) + stmt = select(MCPServer).where(MCPServer.id.in_(mcp_server_ids)) + if persona_id == DEFAULT_PERSONA_ID: + stmt = _add_mcp_server_access_filter(stmt, user) + return list(db_session.scalars(stmt).all()) def _add_mcp_server_access_filter(stmt: Select, user: User) -> Select: """Servers the user may add to an agent (public / direct / group). Admins bypass. - Does not control chat use of agent-attached servers. + Default Chat also uses these grants; custom-agent attachments do not. """ if has_global_permission(user, Permission.FULL_ADMIN_PANEL_ACCESS): return stmt @@ -238,6 +320,7 @@ def create_mcp_server__no_commit( oauth_additional_auth_params: dict[str, str] | None = None, admin_connection_config_id: int | None = None, is_public: bool = True, + available_in_craft: bool = False, ) -> MCPServer: """Create a new MCP server""" new_server = MCPServer( @@ -255,6 +338,7 @@ def create_mcp_server__no_commit( oauth_additional_auth_params=oauth_additional_auth_params, admin_connection_config_id=admin_connection_config_id, is_public=is_public, + available_in_craft=available_in_craft, ) db_session.add(new_server) db_session.flush() # Get the ID without committing @@ -530,6 +614,16 @@ def upsert_user_connection_config( ) +def clear_mcp_admin_connection_config__no_commit( + server: MCPServer, db_session: Session +) -> None: + config_id = server.admin_connection_config_id + server.admin_connection_config = None + server.admin_connection_config_id = None + if config_id is not None: + delete_connection_config(config_id, db_session) + + def delete_connection_config(config_id: int, db_session: Session) -> None: """Delete a connection config""" config = get_connection_config_by_id(config_id, db_session) diff --git a/backend/onyx/db/models.py b/backend/onyx/db/models.py index 65541ce62d..30d0503488 100644 --- a/backend/onyx/db/models.py +++ b/backend/onyx/db/models.py @@ -4171,6 +4171,20 @@ class Tool(Base): ) +class ToolUserPreference(Base): + __tablename__ = "tool_user_preference" + + user_id: Mapped[UUID] = mapped_column( + ForeignKey("user.id", ondelete="CASCADE"), primary_key=True + ) + tool_id: Mapped[int] = mapped_column( + ForeignKey("tool.id", ondelete="CASCADE"), primary_key=True + ) + enabled_in_chat: Mapped[bool] = mapped_column( + Boolean, nullable=False, default=True, server_default=text("true") + ) + + class OAuthConfig(Base): """OAuth provider configuration that can be shared across multiple tools""" @@ -6096,6 +6110,23 @@ class MCPServer(Base): ) +class MCPServerUserPreference(Base): + __tablename__ = "mcp_server_user_preference" + + user_id: Mapped[UUID] = mapped_column( + ForeignKey("user.id", ondelete="CASCADE"), primary_key=True + ) + mcp_server_id: Mapped[int] = mapped_column( + ForeignKey("mcp_server.id", ondelete="CASCADE"), primary_key=True + ) + enabled_in_chat: Mapped[bool] = mapped_column( + Boolean, nullable=False, default=True, server_default=text("true") + ) + enabled_in_craft: Mapped[bool] = mapped_column( + Boolean, nullable=False, default=True, server_default=text("true") + ) + + class MCPServer__User(Base): __tablename__ = "mcp_server__user" mcp_server_id: Mapped[int] = mapped_column( diff --git a/backend/onyx/db/oauth_config.py b/backend/onyx/db/oauth_config.py index 859c8720ae..1a71d8d7dc 100644 --- a/backend/onyx/db/oauth_config.py +++ b/backend/onyx/db/oauth_config.py @@ -2,7 +2,7 @@ from typing import Any from uuid import UUID from sqlalchemy import select -from sqlalchemy.orm import Session +from sqlalchemy.orm import Session, selectinload from onyx.db.models import OAuthConfig, OAuthUserToken, Tool from onyx.utils.logger import setup_logger @@ -143,7 +143,11 @@ def get_all_user_oauth_tokens( """ Get all user OAuth tokens. """ - stmt = select(OAuthUserToken).where(OAuthUserToken.user_id == user_id) + stmt = ( + select(OAuthUserToken) + .where(OAuthUserToken.user_id == user_id) + .options(selectinload(OAuthUserToken.oauth_config)) + ) return list(db_session.scalars(stmt).all()) diff --git a/backend/onyx/db/tools.py b/backend/onyx/db/tools.py index 50118a449c..c6fc433db7 100644 --- a/backend/onyx/db/tools.py +++ b/backend/onyx/db/tools.py @@ -3,9 +3,11 @@ from typing import TYPE_CHECKING, Any, Type, cast from uuid import UUID from sqlalchemy import Select, func, or_, select +from sqlalchemy.dialects.postgresql import insert from sqlalchemy.orm import InstrumentedAttribute, Session from onyx.auth.permissions import has_permission +from onyx.configs.constants import DEFAULT_PERSONA_ID from onyx.context.search.models import PersonaSearchInfo from onyx.db.constants import UNSET, UnsetType from onyx.db.enums import MCPServerStatus, Permission, PermissionAuthority @@ -17,6 +19,7 @@ from onyx.db.models import ( Persona__UserGroup, Tool, ToolCall, + ToolUserPreference, User, ) from onyx.server.features.tool.models import Header @@ -32,6 +35,37 @@ if TYPE_CHECKING: logger = setup_logger() +def get_tool_user_preferences( + tool_ids: list[int], user_id: UUID, db_session: Session +) -> dict[int, ToolUserPreference]: + if not tool_ids: + return {} + return { + preference.tool_id: preference + for preference in db_session.scalars( + select(ToolUserPreference).where( + ToolUserPreference.user_id == user_id, + ToolUserPreference.tool_id.in_(tool_ids), + ) + ) + } + + +def upsert_tool_user_preference__no_commit( + tool_id: int, user_id: UUID, db_session: Session, *, enabled_in_chat: bool +) -> ToolUserPreference: + stmt = insert(ToolUserPreference).values( + user_id=user_id, tool_id=tool_id, enabled_in_chat=enabled_in_chat + ) + return db_session.scalars( + stmt.on_conflict_do_update( + index_elements=[ToolUserPreference.user_id, ToolUserPreference.tool_id], + set_={"enabled_in_chat": enabled_in_chat}, + ).returning(ToolUserPreference), + execution_options={"populate_existing": True}, + ).one() + + def capture_persona_tool_configuration(persona: Persona) -> PersonaToolConfiguration: """Copy loaded configuration; the caller owns the session and transaction.""" return PersonaToolConfiguration( @@ -105,6 +139,24 @@ def get_tools_by_ids(tool_ids: list[int], db_session: Session) -> list[Tool]: return list(db_session.scalars(stmt).all()) +def get_default_chat_tool_ids(tool_ids: list[int], db_session: Session) -> set[int]: + if not tool_ids: + return set() + stmt = ( + select(Tool.id) + .join(Persona__Tool, Persona__Tool.tool_id == Tool.id) + .join(Persona, Persona.id == Persona__Tool.persona_id) + .where( + Persona.id == DEFAULT_PERSONA_ID, + Persona.builtin_persona.is_(True), + Persona.deleted.is_(False), + Tool.id.in_(tool_ids), + ) + .distinct() + ) + return set(db_session.scalars(stmt)) + + def get_tool_by_id(tool_id: int, db_session: Session) -> Tool: tool = db_session.scalar(select(Tool).where(Tool.id == tool_id)) if not tool: diff --git a/backend/onyx/sandbox_proxy/resolvers/mcp_server.py b/backend/onyx/sandbox_proxy/resolvers/mcp_server.py index 7be03646bf..d1922a55b3 100644 --- a/backend/onyx/sandbox_proxy/resolvers/mcp_server.py +++ b/backend/onyx/sandbox_proxy/resolvers/mcp_server.py @@ -25,6 +25,7 @@ from onyx.db.enums import ( from onyx.db.mcp import ( get_craft_enabled_mcp_servers, get_mcp_server_by_id, + get_mcp_server_user_preferences, user_can_access_mcp_server, ) from onyx.db.models import MCPServer @@ -173,6 +174,13 @@ class MCPServerResolver(CredentialResolver): raise CredentialUnavailableError( f"sandbox user {short_log_id(user_id)} not found" ) + preference = get_mcp_server_user_preferences([server.id], user.id, db).get( + server.id + ) + if preference is not None and not preference.enabled_in_craft: + raise CredentialUnavailableError( + f"user {short_log_id(user_id)} disabled MCP server {server.id} in Craft" + ) if not user_can_access_mcp_server(user, server.id, db): # Not shared with this user — never inject admin creds for them. raise CredentialUnavailableError( diff --git a/backend/onyx/server/features/build/external_apps/api.py b/backend/onyx/server/features/build/external_apps/api.py index 2a5a30f823..1d060432c9 100644 --- a/backend/onyx/server/features/build/external_apps/api.py +++ b/backend/onyx/server/features/build/external_apps/api.py @@ -1,3 +1,6 @@ +import ipaddress +import re +from urllib.parse import urlsplit from uuid import UUID from fastapi import APIRouter, Depends @@ -138,6 +141,40 @@ def _to_user_response( if key in stored_raw } authenticated = all(key in stored_raw for key in required_keys) + service_urls: list[str] = [] + for pattern in app.upstream_url_patterns: + if app.app_type != ExternalAppType.CUSTOM: + pattern = pattern.removeprefix("^") + authority = pattern.partition("://")[2].split("/", 1)[0].rsplit("@", 1)[-1] + if app.app_type != ExternalAppType.CUSTOM: + if "?" in authority or re.search(r"(? OAuthStartResponse: + try: + return_path = TypeAdapter(SafeOAuthReturnPath | None).validate_python( + return_path + ) + except ValueError: + raise OnyxError( + OnyxErrorCode.INVALID_INPUT, "return_path must be a safe internal path" + ) app = get_external_app_by_id(db_session, external_app_id) if app is None: raise OnyxError( @@ -155,6 +169,7 @@ def start_external_app_oauth( app, oauth, client_id, client_secret ), code_verifier=code_verifier, + return_path=return_path, ), ) return OAuthStartResponse(authorize_url=authorize_url) @@ -293,4 +308,6 @@ def handle_external_app_oauth_callback( push_skills_for_users({user.id}, db_session) db_session.commit() - return OAuthCallbackResponse(success=True, external_app_id=app.id) + return OAuthCallbackResponse( + success=True, external_app_id=app.id, return_path=payload.return_path + ) diff --git a/backend/onyx/server/features/mcp/api.py b/backend/onyx/server/features/mcp/api.py index b3fb2cd5f9..9a106cd60d 100644 --- a/backend/onyx/server/features/mcp/api.py +++ b/backend/onyx/server/features/mcp/api.py @@ -9,6 +9,8 @@ from urllib.parse import urlparse from uuid import UUID from fastapi import APIRouter, Depends, HTTPException, Request +from fastapi.concurrency import run_in_threadpool +from fastapi.responses import StreamingResponse from mcp.client.auth import OAuthClientProvider from mcp.shared.auth import OAuthClientInformationFull from mcp.types import Tool as MCPLibTool @@ -26,7 +28,7 @@ from onyx.auth.permissions import ( ) from onyx.auth.scoped_permissions import assert_within_scope, get_scoped_groups from onyx.configs.app_configs import WEB_DOMAIN -from onyx.db.engine.sql_engine import get_session +from onyx.db.engine.sql_engine import get_session, get_session_with_current_tenant from onyx.db.enums import ( EndpointPolicy, GatedAppKind, @@ -45,6 +47,7 @@ from onyx.db.gated_app import ( ) from onyx.db.mcp import ( affected_user_ids_for_mcp_server, + clear_mcp_admin_connection_config__no_commit, create_connection_config, create_mcp_server__no_commit, delete_all_user_connection_configs_for_server_no_commit, @@ -54,7 +57,9 @@ from onyx.db.mcp import ( get_all_mcp_servers, get_all_mcp_tools_for_server, get_craft_enabled_mcp_servers, + get_default_chat_mcp_server_ids, get_mcp_server_by_id, + get_mcp_server_user_preferences, get_mcp_servers_accessible_to_user, get_mcp_servers_for_persona, get_user_connection_config, @@ -62,6 +67,7 @@ from onyx.db.mcp import ( update_connection_config, update_connection_config__no_commit, update_mcp_server__no_commit, + upsert_mcp_server_user_preferences__no_commit, upsert_user_connection_config, user_can_access_mcp_server, user_can_configure_mcp_credentials, @@ -102,6 +108,8 @@ from onyx.server.features.mcp.models import ( MCPConnectionData, MCPConnectionStatus, MCPOAuthCallbackResponse, + MCPOAuthDiscoveryRequest, + MCPOAuthDiscoveryResponse, MCPOAuthKeys, MCPServer, MCPServerCreateResponse, @@ -109,12 +117,15 @@ from onyx.server.features.mcp.models import ( MCPServerSimpleUpdateRequest, MCPServersResponse, MCPServerUpdateResponse, + MCPSetupProbeRequest, MCPToolCreateRequest, MCPToolListResponse, MCPToolUpdateRequest, MCPUserCredentialsRequest, MCPUserOAuthConnectRequest, MCPUserOAuthConnectResponse, + MCPUserPreferencesResponse, + MCPUserPreferencesUpdateRequest, contains_mcp_placeholder, merge_mcp_headers, ) @@ -124,6 +135,10 @@ from onyx.server.features.mcp.oauth import ( complete_mcp_oauth_authorization, make_oauth_provider, ) +from onyx.server.features.mcp.oauth_discovery import ( + discover_oauth_capabilities, + oauth_setup_error, +) from onyx.server.features.mcp.oauth_flow import ( MCP_OAUTH_ATTEMPTS, OAuthAlreadyAuthenticated, @@ -132,6 +147,7 @@ from onyx.server.features.mcp.oauth_flow import ( start_known_provider_oauth_flow, validate_mcp_oauth_flow_configuration, ) +from onyx.server.features.mcp.setup_probe import stream_mcp_probe from onyx.server.features.mcp.ssrf import validate_mcp_outbound_url from onyx.server.features.tool.models import ToolSnapshot from onyx.utils.encryption import ( @@ -622,6 +638,46 @@ def test_mcp_server_credentials( return False, f"Connection failed: {str(e)}" +@admin_router.post("/probe") +async def stream_admin_mcp_probe( + request: MCPSetupProbeRequest, + _user: User = Depends( + require_permission(Permission.MANAGE_ACTIONS, allow_scope=True) + ), +) -> StreamingResponse: + return StreamingResponse( + stream_mcp_probe(request.server_url), + media_type="application/x-ndjson", + headers={ + "Cache-Control": "no-cache, no-transform", + "X-Accel-Buffering": "no", + }, + ) + + +@admin_router.post("/oauth/discover") +async def discover_admin_oauth( + request: MCPOAuthDiscoveryRequest, + user: User = Depends( + require_permission(Permission.MANAGE_ACTIONS, allow_scope=True) + ), +) -> MCPOAuthDiscoveryResponse: + server_url = await run_in_threadpool( + lambda: _get_authorized_discovery_url(request.server_id, user) + ) + return await discover_oauth_capabilities(server_url) + + +def _get_authorized_discovery_url(server_id: int, user: User) -> str: + with get_session_with_current_tenant() as db: + try: + server = get_mcp_server_by_id(server_id, db) + except ValueError as error: + raise OnyxError(OnyxErrorCode.NOT_FOUND, "MCP server not found") from error + _ensure_mcp_server_owner_or_admin(server, user) + return server.server_url + + @admin_router.post("/oauth/connect", response_model=MCPUserOAuthConnectResponse) async def connect_admin_oauth( request: MCPUserOAuthConnectRequest, @@ -872,10 +928,7 @@ async def _connect_oauth( except Exception as e: saved_e = log_exception_group(e) if isinstance(e, ExceptionGroup) else e logger.error("OAuth initialization failed: %s", saved_e) - raise OnyxError( - OnyxErrorCode.INVALID_INPUT, - f"Failed to initialize OAuth client: {str(saved_e)}", - ) from e + raise oauth_setup_error(e) from e match oauth_result: case OAuthAuthorizationRedirect(url=oauth_url): @@ -1335,6 +1388,58 @@ def _db_mcp_server_to_api_mcp_server( ) +def _stamp_user_preferences(servers: list[MCPServer], user: User, db: Session) -> None: + preferences = get_mcp_server_user_preferences( + [server.id for server in servers], user.id, db + ) + for server in servers: + preference = preferences.get(server.id) + if preference is not None: + server.user_enabled_in_chat = preference.enabled_in_chat + server.user_enabled_in_craft = preference.enabled_in_craft + + +@router.patch("/server/{server_id}/preferences") +def update_mcp_user_preferences( + server_id: int, + request: MCPUserPreferencesUpdateRequest, + db: Session = Depends(get_session), + user: User = Depends(require_permission(Permission.BASIC_ACCESS)), +) -> MCPUserPreferencesResponse: + if not user_can_access_mcp_server(user, server_id, db): + raise OnyxError(OnyxErrorCode.UNAUTHORIZED, "MCP server is not accessible") + try: + server = get_mcp_server_by_id(server_id, db) + except ValueError as error: + raise OnyxError(OnyxErrorCode.NOT_FOUND, "MCP server not found") from error + if ( + request.enabled_in_chat is True + and server_id not in get_default_chat_mcp_server_ids([server_id], db) + ): + raise OnyxError( + OnyxErrorCode.INVALID_INPUT, "MCP server is not available in default Chat" + ) + if request.enabled_in_craft is True and not server.available_in_craft: + raise OnyxError( + OnyxErrorCode.INVALID_INPUT, "MCP server is not available in Craft" + ) + preference = upsert_mcp_server_user_preferences__no_commit( + server_id, + user.id, + db, + enabled_in_chat=request.enabled_in_chat, + enabled_in_craft=request.enabled_in_craft, + ) + response = MCPUserPreferencesResponse( + enabled_in_chat=preference.enabled_in_chat, + enabled_in_craft=preference.enabled_in_craft, + ) + db.commit() + if request.enabled_in_craft is not None: + _hot_reload_craft_sessions({user.id}, db) + return response + + @router.get("/servers/persona/{assistant_id}", response_model=MCPServersResponse) def get_mcp_servers_for_assistant( assistant_id: str, @@ -1369,6 +1474,7 @@ def get_mcp_servers_for_assistant( for db_server in db_mcp_servers ] + _stamp_user_preferences(mcp_servers, user, db) return MCPServersResponse(assistant_id=assistant_id, mcp_servers=mcp_servers) except Exception as e: @@ -1386,10 +1492,16 @@ def get_mcp_servers_for_user( Chat uses ``/servers/persona/{id}`` for servers already on a persona. """ db_mcp_servers = get_mcp_servers_accessible_to_user(user, db) + default_chat_server_ids = get_default_chat_mcp_server_ids( + [server.id for server in db_mcp_servers], db + ) mcp_servers = [ _db_mcp_server_to_api_mcp_server(db_server, db, request_user=user) for db_server in db_mcp_servers ] + for server in mcp_servers: + server.default_chat_enabled = server.id in default_chat_server_ids + _stamp_user_preferences(mcp_servers, user, db) return MCPServersResponse(mcp_servers=mcp_servers) @@ -1403,11 +1515,12 @@ def get_craft_mcp_servers_for_user( rows as chat, so a server connected in either surface shows as authenticated here. - `craft_connected` is the only field here that answers "will this server - actually reach the user's sessions" — see `resolve_craft_mcp_servers`, which - filters emission on the same predicate. + Personal preferences remain visible so users can enable a server again. + `craft_connected` reports credentials independently of personal availability. """ - db_mcp_servers = get_craft_enabled_mcp_servers(db, user) + db_mcp_servers = get_craft_enabled_mcp_servers( + db, user, include_personal_disabled=True + ) user_configs = get_user_connection_configs( [s.id for s in db_mcp_servers], user.email, db ) @@ -1420,6 +1533,7 @@ def get_craft_mcp_servers_for_user( server.connection_status == MCPConnectionStatus.CONNECTED ) mcp_servers.append(server) + _stamp_user_preferences(mcp_servers, user, db) return MCPServersResponse(mcp_servers=mcp_servers) @@ -1996,9 +2110,7 @@ def _upsert_mcp_server( ) if server_url_changed and mcp_server.admin_connection_config_id: - previous_admin_config_id = mcp_server.admin_connection_config_id - mcp_server.admin_connection_config_id = None - delete_connection_config(previous_admin_config_id, db_session) + clear_mcp_admin_connection_config__no_commit(mcp_server, db_session) if ( request.auth_type == MCPAuthenticationType.OAUTH and request.oauth_provider_mode is MCPOAuthProviderMode.AUTO_DISCOVERY @@ -2347,6 +2459,9 @@ def get_mcp_servers_for_admin( if server.owner == user.email or server.id in connected ] + default_chat_server_ids = get_default_chat_mcp_server_ids( + [server.id for server in visible_servers], db + ) mcp_servers = [ _db_mcp_server_to_api_mcp_server( db_server, @@ -2358,6 +2473,8 @@ def get_mcp_servers_for_admin( ) for db_server in visible_servers ] + for server in mcp_servers: + server.default_chat_enabled = server.id in default_chat_server_ids return MCPServersResponse(mcp_servers=mcp_servers) @@ -2573,6 +2690,7 @@ def create_mcp_server_simple( auth_type=None, # To be configured later transport=None, # To be configured later auth_performer=None, # To be configured later + available_in_craft=request.available_in_craft, db_session=db_session, ) @@ -2634,6 +2752,14 @@ def update_mcp_server_simple( _validate_mcp_server_url(request.server_url, "server_url", require_https=False) + url_changed = ( + request.server_url is not None and request.server_url != mcp_server.server_url + ) + reload_user_ids: set[UUID] = set() + if url_changed: + reload_user_ids.update(_invalidate_mcp_user_credentials(mcp_server, db_session)) + clear_mcp_admin_connection_config__no_commit(mcp_server, db_session) + # Update only provided fields updated_server = update_mcp_server__no_commit( server_id=server_id, @@ -2642,6 +2768,14 @@ def update_mcp_server_simple( description=request.description, server_url=request.server_url, available_in_craft=request.available_in_craft, + status=( + MCPServerStatus.AWAITING_AUTH + if mcp_server.auth_type + not in (None, MCPAuthenticationType.NONE, MCPAuthenticationType.PT_OAUTH) + else MCPServerStatus.CREATED + ) + if url_changed + else None, ) acl_changing = any( @@ -2651,13 +2785,10 @@ def update_mcp_server_simple( # Only an ACL change can drop a user's access; snapshot the pre-change # recipients in that case so they're reloaded to lose the server (the # post-update query wouldn't include them). - reload_user_ids: set[UUID] = ( - affected_user_ids_for_mcp_server(updated_server, db_session) - if acl_changing - else set() - ) - if acl_changing: + reload_user_ids.update( + affected_user_ids_for_mcp_server(updated_server, db_session) + ) _apply_mcp_server_access( mcp_server=updated_server, acting_user=user, diff --git a/backend/onyx/server/features/mcp/client.py b/backend/onyx/server/features/mcp/client.py index 6768e76ed7..f2bbaa9dc9 100644 --- a/backend/onyx/server/features/mcp/client.py +++ b/backend/onyx/server/features/mcp/client.py @@ -13,6 +13,7 @@ from mcp import ClientSession from mcp.client.auth import OAuthClientProvider from mcp.client.sse import sse_client from mcp.client.streamable_http import streamablehttp_client +from mcp.shared._httpx_utils import McpHttpClientFactory from mcp.types import ( CallToolResult, InitializeResult, @@ -91,6 +92,7 @@ def _create_mcp_client_function_runner( connection_headers: dict[str, str] | None = None, transport: MCPTransport = MCPTransport.STREAMABLE_HTTP, auth: OAuthClientProvider | None = None, + httpx_client_factory: McpHttpClientFactory | None = None, **kwargs: Any, ) -> Callable[[], Coroutine[Any, Any, T]]: auth_headers = connection_headers or {} @@ -111,7 +113,7 @@ def _create_mcp_client_function_runner( server_url, headers=auth_headers, auth=auth_for_request, - httpx_client_factory=mcp_ssrf_httpx_client_factory, + httpx_client_factory=httpx_client_factory or mcp_ssrf_httpx_client_factory, ) as client_tuple: if len(client_tuple) == 3: read, write, _ = client_tuple @@ -171,16 +173,23 @@ def _call_mcp_client_function_sync( raise e -async def _call_mcp_client_function_async( +async def call_mcp_client_function_async( function: Callable[[ClientSession], Coroutine[Any, Any, T]], server_url: str, connection_headers: dict[str, str] | None = None, transport: MCPTransport = MCPTransport.STREAMABLE_HTTP, auth: OAuthClientProvider | None = None, + httpx_client_factory: McpHttpClientFactory | None = None, **kwargs: Any, ) -> T: run_client_function = _create_mcp_client_function_runner( - function, server_url, connection_headers, transport, auth, **kwargs + function, + server_url, + connection_headers, + transport, + auth, + httpx_client_factory=httpx_client_factory, + **kwargs, ) return await run_client_function() @@ -241,7 +250,7 @@ async def initialize_mcp_client( transport: MCPTransport = MCPTransport.STREAMABLE_HTTP, auth: OAuthClientProvider | None = None, ) -> InitializeResult: - return await _call_mcp_client_function_async( + return await call_mcp_client_function_async( lambda session: session.initialize(), server_url, connection_headers, diff --git a/backend/onyx/server/features/mcp/models.py b/backend/onyx/server/features/mcp/models.py index 6dae2cba50..cb81df78ad 100644 --- a/backend/onyx/server/features/mcp/models.py +++ b/backend/onyx/server/features/mcp/models.py @@ -1,7 +1,7 @@ import datetime import re from enum import Enum -from typing import Any, List, Literal, NotRequired, Optional, TypedDict +from typing import Annotated, Any, List, Literal, NotRequired, Optional, TypedDict from uuid import UUID from mcp.shared.auth import ( @@ -421,6 +421,9 @@ class MCPToolUpdateRequest(BaseModel): class MCPServerSimpleCreateRequest(BaseModel): + available_in_craft: bool = Field( + default=False, description="Whether the Craft agent may use this server" + ) name: str = Field(..., description="Name of the MCP server") description: Optional[str] = Field( None, description="Description of the MCP server" @@ -513,6 +516,53 @@ class MCPOAuthConnectResponse(BaseModel): pending_tool: dict = Field(..., description="Pending tool configuration") +class MCPOAuthDiscoveryRequest(BaseModel): + server_id: int + + +class MCPOAuthDiscoveryResponse(BaseModel): + status: Literal["metadata_found", "metadata_missing", "unreachable"] + supports_automatic_registration: bool | None = None + authorization_endpoint: str | None = None + token_endpoint: str | None = None + + +class MCPSetupProbeRequest(BaseModel): + server_url: str + + +class MCPPublicProbeResult(BaseModel): + kind: Literal["public"] = "public" + transport: Literal[MCPTransport.STREAMABLE_HTTP, MCPTransport.SSE] + tool_count: int + server_name: str + + +class MCPOAuthProbeResult(BaseModel): + kind: Literal["oauth"] = "oauth" + transport: Literal[MCPTransport.STREAMABLE_HTTP, MCPTransport.SSE] | None = None + discovery: MCPOAuthDiscoveryResponse + + +class MCPUnknownAuthProbeResult(BaseModel): + kind: Literal["auth_unknown"] = "auth_unknown" + transport: Literal[MCPTransport.STREAMABLE_HTTP, MCPTransport.SSE] | None = None + + +class MCPUnreachableProbeResult(BaseModel): + kind: Literal["unreachable"] = "unreachable" + reason: Literal["blocked", "connection_failed"] = "connection_failed" + + +MCPSetupProbeResult = Annotated[ + MCPPublicProbeResult + | MCPOAuthProbeResult + | MCPUnknownAuthProbeResult + | MCPUnreachableProbeResult, + Field(discriminator="kind"), +] + + class MCPUserOAuthConnectRequest(BaseModel): server_id: int = Field(..., description="ID of the MCP server") return_path: SafeOAuthReturnPath = Field( @@ -694,6 +744,24 @@ class MCPApiKeyResponse(BaseModel): ) +class MCPUserPreferencesUpdateRequest(BaseModel): + enabled_in_chat: bool | None = None + enabled_in_craft: bool | None = None + + @model_validator(mode="after") + def validate_preferences(self) -> "MCPUserPreferencesUpdateRequest": + if self.enabled_in_chat is None and self.enabled_in_craft is None: + raise ValueError("At least one preference is required") + if any(value is None for value in self.model_dump(exclude_unset=True).values()): + raise ValueError("Preferences must be booleans") + return self + + +class MCPUserPreferencesResponse(BaseModel): + enabled_in_chat: bool + enabled_in_craft: bool + + class MCPServer(BaseModel): id: int name: str @@ -715,6 +783,15 @@ class MCPServer(BaseModel): groups: list[int] = Field(default_factory=list) users: list[UUID] = Field(default_factory=list) available_in_craft: bool = False + user_enabled_in_chat: bool = True + user_enabled_in_craft: bool = True + default_chat_enabled: bool | None = Field( + None, + description=( + "Whether any server tool is selected for the default assistant. " + "Computed by member and admin listings, independently of credentials and tool enablement." + ), + ) tool_policies: Optional[dict[str, EndpointPolicy]] = Field( None, description=( diff --git a/backend/onyx/server/features/mcp/oauth_discovery.py b/backend/onyx/server/features/mcp/oauth_discovery.py new file mode 100644 index 0000000000..34a5d4b8c4 --- /dev/null +++ b/backend/onyx/server/features/mcp/oauth_discovery.py @@ -0,0 +1,99 @@ +import asyncio + +import httpx +from mcp.client.auth.exceptions import OAuthRegistrationError +from mcp.client.auth.utils import ( + build_oauth_authorization_server_metadata_discovery_urls, + build_protected_resource_metadata_discovery_urls, + create_oauth_metadata_request, + extract_resource_metadata_from_www_auth, + handle_auth_metadata_response, + handle_protected_resource_response, +) +from mcp.shared.auth_utils import check_resource_allowed, resource_url_from_server_url + +from onyx.error_handling.error_codes import OnyxErrorCode +from onyx.error_handling.exceptions import OnyxError +from onyx.server.features.mcp.models import MCPOAuthDiscoveryResponse +from onyx.server.features.mcp.ssrf import mcp_ssrf_httpx_client_factory +from onyx.utils.url import SSRFException + + +async def discover_oauth_capabilities( + server_url: str, *, challenge_url: str | None = None +) -> MCPOAuthDiscoveryResponse: + try: + async with asyncio.timeout(15): + return await _discover_oauth_capabilities(server_url, challenge_url) + except (httpx.HTTPError, TimeoutError, ValueError, SSRFException): + return MCPOAuthDiscoveryResponse(status="unreachable") + + +async def _discover_oauth_capabilities( + server_url: str, challenge_url: str | None +) -> MCPOAuthDiscoveryResponse: + saw_server_error = False + async with mcp_ssrf_httpx_client_factory(timeout=httpx.Timeout(5)) as client: + if challenge_url is None: + # An SSE endpoint can keep a GET body open forever. + async with client.stream("GET", server_url) as response: + challenge_url = extract_resource_metadata_from_www_auth(response) + auth_server_url: str | None = None + for url in build_protected_resource_metadata_discovery_urls( + challenge_url, server_url + ): + response = await client.send(create_oauth_metadata_request(url)) + saw_server_error |= response.status_code >= 500 + resource_metadata = await handle_protected_resource_response(response) + if resource_metadata and check_resource_allowed( + requested_resource=resource_url_from_server_url(server_url), + configured_resource=str(resource_metadata.resource), + ): + auth_server_url = str(resource_metadata.authorization_servers[0]) + break + for url in build_oauth_authorization_server_metadata_discovery_urls( + auth_server_url, server_url + ): + response = await client.send(create_oauth_metadata_request(url)) + saw_server_error |= response.status_code >= 500 + _, metadata = await handle_auth_metadata_response(response) + if metadata: + return MCPOAuthDiscoveryResponse( + status="metadata_found", + supports_automatic_registration=bool( + metadata.registration_endpoint + or metadata.client_id_metadata_document_supported + ), + authorization_endpoint=str(metadata.authorization_endpoint), + token_endpoint=str(metadata.token_endpoint), + ) + return MCPOAuthDiscoveryResponse( + status="unreachable" if saw_server_error else "metadata_missing" + ) + + +def oauth_setup_error(error: Exception) -> OnyxError: + """Do not expose upstream response bodies, URLs, or credentials to the browser.""" + if isinstance(error, ExceptionGroup): + errors = [oauth_setup_error(child) for child in error.exceptions] + return next( + (item for item in errors if item.extra != {"oauth_error": "oauth_failed"}), + errors[0], + ) + if isinstance(error, OAuthRegistrationError): + return OnyxError( + OnyxErrorCode.INVALID_INPUT, + "Automatic OAuth client registration failed. Use your own OAuth client, or check the provider's registration requirements.", + extra={"oauth_error": "client_registration_failed"}, + ) + if isinstance(error, (httpx.HTTPError, TimeoutError)): + return OnyxError( + OnyxErrorCode.INVALID_INPUT, + "Could not reach the OAuth provider. Check the server URL and try again.", + extra={"oauth_error": "connection_failed"}, + ) + return OnyxError( + OnyxErrorCode.INVALID_INPUT, + "OAuth setup failed. Check the provider's OAuth settings and try again.", + extra={"oauth_error": "oauth_failed"}, + ) diff --git a/backend/onyx/server/features/mcp/oauth_flow.py b/backend/onyx/server/features/mcp/oauth_flow.py index 04e6f3ab02..c7d359c52a 100644 --- a/backend/onyx/server/features/mcp/oauth_flow.py +++ b/backend/onyx/server/features/mcp/oauth_flow.py @@ -210,7 +210,8 @@ async def _start_oauth_from_well_known_metadata( raise OnyxError( OnyxErrorCode.INVALID_INPUT, "OAuth auto-discovery could not find protected resource metadata at " - "the MCP server's well-known URI.", + "the MCP server's well-known URI. Check the URL or enter OAuth endpoints manually.", + extra={"oauth_error": "discovery_failed"}, ) async with mcp_oauth_challenge_httpx_client_factory( @@ -371,12 +372,14 @@ async def _start_auto_discovery_oauth_flow_once( except TimeoutError as error: raise OnyxError( OnyxErrorCode.INVALID_INPUT, - "Timed out during MCP OAuth discovery.", + "Timed out during MCP OAuth discovery. Check the server URL and try again.", + extra={"oauth_error": "connection_failed"}, ) from error raise OnyxError( OnyxErrorCode.INVALID_INPUT, - "OAuth auto-discovery did not produce an authorization redirect.", + "OAuth auto-discovery did not produce an authorization redirect. Check the provider settings.", + extra={"oauth_error": "discovery_failed"}, ) diff --git a/backend/onyx/server/features/mcp/setup_probe.py b/backend/onyx/server/features/mcp/setup_probe.py new file mode 100644 index 0000000000..c7c00a9bfa --- /dev/null +++ b/backend/onyx/server/features/mcp/setup_probe.py @@ -0,0 +1,202 @@ +import asyncio +import json +from collections.abc import AsyncGenerator, Awaitable, Callable, Iterator +from contextlib import suppress +from typing import Literal, TypedDict + +import httpx +from mcp import ClientSession +from mcp.client.auth.utils import extract_resource_metadata_from_www_auth + +from onyx.db.enums import MCPTransport +from onyx.server.features.mcp.client import call_mcp_client_function_async +from onyx.server.features.mcp.models import ( + MCPOAuthProbeResult, + MCPPublicProbeResult, + MCPSetupProbeResult, + MCPUnknownAuthProbeResult, + MCPUnreachableProbeResult, +) +from onyx.server.features.mcp.oauth_discovery import discover_oauth_capabilities +from onyx.server.features.mcp.ssrf import mcp_ssrf_httpx_client_factory +from onyx.utils.url import SSRFException + + +class MCPProbeProgress(TypedDict): + type: Literal["progress"] + stage: Literal["contacting_server", "authentication"] + state: Literal["pending", "success", "error"] + + +class MCPProbeCompletion(TypedDict): + type: Literal["result"] + result: dict[str, object] + + +Progress = Callable[[MCPProbeProgress], Awaitable[None]] + + +async def stream_mcp_probe(server_url: str) -> AsyncGenerator[str, None]: + queue: asyncio.Queue[MCPProbeProgress | MCPProbeCompletion] = asyncio.Queue() + + async def run() -> None: + try: + result = await probe_mcp_server(server_url, queue.put) + result_data = result.model_dump(mode="json") + except Exception: + result_data = MCPUnreachableProbeResult().model_dump(mode="json") + await queue.put({"type": "result", "result": result_data}) + + task = asyncio.create_task(run()) + try: + while True: + event = await queue.get() + yield json.dumps(event) + "\n" + if event["type"] == "result": + break + finally: + task.cancel() + with suppress(asyncio.CancelledError): + await task + + +def _leaf_errors(error: Exception) -> Iterator[Exception]: + if isinstance(error, ExceptionGroup): + for child in error.exceptions: + yield from _leaf_errors(child) + else: + yield error + + +async def probe_mcp_server( + server_url: str, progress: Progress | None = None +) -> MCPSetupProbeResult: + contacted = False + + async def report_contact() -> None: + nonlocal contacted + if not contacted: + contacted = True + if progress: + await progress( + { + "type": "progress", + "stage": "contacting_server", + "state": "success", + } + ) + await progress( + {"type": "progress", "stage": "authentication", "state": "pending"} + ) + + if progress: + await progress( + {"type": "progress", "stage": "contacting_server", "state": "pending"} + ) + result: MCPSetupProbeResult = MCPUnreachableProbeResult() + try: + async with asyncio.timeout(25): + for transport in (MCPTransport.STREAMABLE_HTTP, MCPTransport.SSE): + transport_result = await _probe_transport( + server_url, transport, report_contact + ) + if transport_result is not None: + result = transport_result + break + except Exception: + result = MCPUnreachableProbeResult() + if progress: + await progress( + { + "type": "progress", + "stage": "authentication" if contacted else "contacting_server", + "state": "success" if result.kind in ("public", "oauth") else "error", + } + ) + return result + + +async def _probe_transport( + server_url: str, + transport: Literal[MCPTransport.STREAMABLE_HTTP, MCPTransport.SSE], + contacted: Callable[[], Awaitable[None]], +) -> MCPSetupProbeResult | None: + challenges: list[httpx.Response] = [] + sse_established = False + + async def record_challenge(response: httpx.Response) -> None: + # The SSE SDK logs message-POST errors instead of propagating them. + nonlocal sse_established + if ( + transport == MCPTransport.SSE + and response.request.method == "GET" + and response.is_success + and response.headers.get("content-type", "").startswith("text/event-stream") + ): + sse_established = True + if response.status_code in (401, 403): + challenges.append(response) + await contacted() + + def probe_client_factory( + headers: dict[str, str] | None = None, + timeout: httpx.Timeout | None = None, + auth: httpx.Auth | None = None, + ) -> httpx.AsyncClient: + client = mcp_ssrf_httpx_client_factory( + headers=headers, timeout=timeout, auth=auth + ) + client.event_hooks["response"].append(record_challenge) + return client + + try: + + async def inspect(session: ClientSession) -> tuple[str, int]: + initialized = await session.initialize() + await contacted() + tools = await session.list_tools() + return initialized.serverInfo.name, len(tools.tools) + + async with asyncio.timeout(5): + name, tool_count = await call_mcp_client_function_async( + inspect, + server_url, + transport=transport, + httpx_client_factory=probe_client_factory, + ) + await contacted() + return MCPPublicProbeResult( + transport=transport, server_name=name, tool_count=tool_count + ) + except Exception as error: + errors = list(_leaf_errors(error)) + if any(isinstance(item, SSRFException) for item in errors): + return MCPUnreachableProbeResult(reason="blocked") + challenge = ( + challenges[0] + if challenges + else next( + ( + item.response + for item in errors + if isinstance(item, httpx.HTTPStatusError) + and item.response.status_code in (401, 403) + ), + None, + ) + ) + if challenge is not None: + await contacted() + discovery = await discover_oauth_capabilities( + server_url, + challenge_url=extract_resource_metadata_from_www_auth(challenge), + ) + if discovery.status == "metadata_found": + return MCPOAuthProbeResult( + discovery=discovery, + transport=MCPTransport.SSE if sse_established else None, + ) + return MCPUnknownAuthProbeResult( + transport=MCPTransport.SSE if sse_established else None + ) + return None diff --git a/backend/onyx/server/features/mcp/ssrf.py b/backend/onyx/server/features/mcp/ssrf.py index 2bde06ffaa..1615e7b944 100644 --- a/backend/onyx/server/features/mcp/ssrf.py +++ b/backend/onyx/server/features/mcp/ssrf.py @@ -8,6 +8,7 @@ registration, and token request — is checked. """ import httpx +from fastapi.concurrency import run_in_threadpool from onyx.server.security.models import outbound_ssrf_params from onyx.server.security.store import get_security_settings @@ -41,7 +42,7 @@ class _SSRFGuardAsyncTransport(httpx.AsyncHTTPTransport): targets are validated too.""" async def handle_async_request(self, request: httpx.Request) -> httpx.Response: - validate_mcp_outbound_url(str(request.url)) + await run_in_threadpool(validate_mcp_outbound_url, str(request.url)) return await super().handle_async_request(request) diff --git a/backend/onyx/server/features/tool/api.py b/backend/onyx/server/features/tool/api.py index 35aee6bd05..390bdb95f6 100644 --- a/backend/onyx/server/features/tool/api.py +++ b/backend/onyx/server/features/tool/api.py @@ -23,11 +23,14 @@ from onyx.db.tools import ( can_manage_tool, create_tool__no_commit, delete_tool__no_commit, + get_default_chat_tool_ids, get_tool_by_id, get_tool_ids_connected_to_groups, + get_tool_user_preferences, get_tools, get_tools_by_ids, update_tool, + upsert_tool_user_preference__no_commit, ) from onyx.error_handling.error_codes import OnyxErrorCode from onyx.error_handling.exceptions import OnyxError @@ -36,6 +39,7 @@ from onyx.server.features.tool.models import ( CustomToolUpdate, Header, ToolSnapshot, + ToolUserPreferences, ) from onyx.server.features.tool.tool_visibility import should_expose_tool_to_fe from onyx.tools.built_in_tools import get_built_in_tool_by_id @@ -397,6 +401,7 @@ def list_openapi_tools( ) ) + _stamp_personal_tool_preferences(openapi_tools, user, db_session) return openapi_tools @@ -415,7 +420,60 @@ def get_custom_tool( OnyxErrorCode.INSUFFICIENT_PERMISSIONS, "You can only view actions you manage, or ones an agent in your groups uses.", ) - return ToolSnapshot.from_model(tool) + snapshot = ToolSnapshot.from_model(tool) + _stamp_personal_tool_preferences([snapshot], user, db_session) + return snapshot + + +def _stamp_personal_tool_preferences( + tools: list[ToolSnapshot], user: User, db_session: Session +) -> None: + preferences = get_tool_user_preferences( + [tool.id for tool in tools if tool.definition is not None], user.id, db_session + ) + for tool in tools: + preference = preferences.get(tool.id) + if preference is not None: + tool.user_enabled_in_chat = preference.enabled_in_chat + + +@router.patch("/{tool_id}/preferences", tags=PUBLIC_API_TAGS) +def update_tool_user_preferences( + tool_id: int, + request: ToolUserPreferences, + db_session: Session = Depends(get_session), + user: User = Depends(require_permission(Permission.BASIC_ACCESS)), +) -> ToolUserPreferences: + try: + tool = get_tool_by_id(tool_id, db_session) + except ValueError as error: + raise OnyxError(OnyxErrorCode.NOT_FOUND, "Plugin not found") from error + if ( + tool.in_code_tool_id is not None + or tool.mcp_server_id is not None + or tool.openapi_schema is None + ): + raise OnyxError( + OnyxErrorCode.INVALID_INPUT, + "Only OpenAPI plugins support these preferences", + ) + if not should_expose_tool_to_fe(tool): + raise OnyxError( + OnyxErrorCode.INSUFFICIENT_PERMISSIONS, "Plugin is not available" + ) + if request.enabled_in_chat and ( + not tool.enabled + or tool.id not in get_default_chat_tool_ids([tool.id], db_session) + ): + raise OnyxError( + OnyxErrorCode.INSUFFICIENT_PERMISSIONS, + "Your admin has disabled this plugin for main Chat", + ) + preference = upsert_tool_user_preference__no_commit( + tool_id, user.id, db_session, enabled_in_chat=request.enabled_in_chat + ) + db_session.commit() + return ToolUserPreferences(enabled_in_chat=preference.enabled_in_chat) @router.get("", tags=PUBLIC_API_TAGS) @@ -453,4 +511,12 @@ def list_tools( # All custom tools and available built-in tools are included filtered_tools.append(ToolSnapshot.from_model(tool)) + default_chat_tool_ids = get_default_chat_tool_ids( + [tool.id for tool in filtered_tools], db_session + ) + for tool in filtered_tools: + tool.default_chat_enabled = tool.id in default_chat_tool_ids + + _stamp_personal_tool_preferences(filtered_tools, user, db_session) + return filtered_tools diff --git a/backend/onyx/server/features/tool/models.py b/backend/onyx/server/features/tool/models.py index 044463f46c..fa0a517edd 100644 --- a/backend/onyx/server/features/tool/models.py +++ b/backend/onyx/server/features/tool/models.py @@ -38,6 +38,11 @@ class ToolSnapshot(BaseModel): oauth_config_id: int | None = None oauth_config_name: str | None = None enabled: bool = True + user_enabled_in_chat: bool = True + default_chat_enabled: bool | None = Field( + None, + description="Whether this tool is selected for the default assistant.", + ) # Server-stamped affordance map; fail-closed empty (only the admin actions list stamps it). permissions: dict[str, bool] = Field(default_factory=dict) @@ -78,6 +83,10 @@ class ToolSnapshot(BaseModel): ) +class ToolUserPreferences(BaseModel): + enabled_in_chat: bool + + class Header(BaseModel): key: str value: str diff --git a/backend/onyx/server/features/user_oauth_token/api.py b/backend/onyx/server/features/user_oauth_token/api.py index 88f7387a97..9af6d77b7f 100644 --- a/backend/onyx/server/features/user_oauth_token/api.py +++ b/backend/onyx/server/features/user_oauth_token/api.py @@ -18,6 +18,7 @@ class OAuthTokenStatus(BaseModel): oauth_config_id: int expires_at: int | None # Unix timestamp is_expired: bool + can_refresh: bool = False @router.get("/status") @@ -37,11 +38,21 @@ def get_user_oauth_token_status( token_data = ( token.token_data.get_value(apply_mask=False) if token.token_data else {} ) + refresh_token = token_data.get("refresh_token") + config = token.oauth_config result.append( OAuthTokenStatus( oauth_config_id=token.oauth_config_id, expires_at=OAuthTokenManager.token_expiration_time(token_data), is_expired=OAuthTokenManager.is_token_expired(token_data), + can_refresh=( + isinstance(refresh_token, str) + and bool(refresh_token) + and config.client_id is not None + and bool(config.client_id.get_value(apply_mask=False)) + and config.client_secret is not None + and bool(config.client_secret.get_value(apply_mask=False)) + ), ) ) return result diff --git a/backend/onyx/tools/tool_constructor.py b/backend/onyx/tools/tool_constructor.py index 54a30f913b..adacd7a283 100644 --- a/backend/onyx/tools/tool_constructor.py +++ b/backend/onyx/tools/tool_constructor.py @@ -9,14 +9,19 @@ from sqlalchemy.orm import Session from onyx.auth.oauth_token_manager import OAuthTokenManager from onyx.chat.emitter import Emitter from onyx.configs.app_configs import DISABLE_VECTOR_DB +from onyx.configs.constants import DEFAULT_PERSONA_ID from onyx.configs.model_configs import GEN_AI_TEMPERATURE from onyx.context.search.models import BaseFilters from onyx.db.engine.sql_engine import get_session_with_current_tenant_if_none -from onyx.db.mcp import get_mcp_server_by_id +from onyx.db.mcp import ( + get_mcp_server_by_id, + get_mcp_server_user_preferences, + user_can_access_mcp_server, +) from onyx.db.models import User from onyx.db.oauth_config import get_oauth_config from onyx.db.search_settings import get_current_search_settings -from onyx.db.tools import get_builtin_tool +from onyx.db.tools import get_builtin_tool, get_tool_user_preferences from onyx.document_index.factory import get_default_document_index from onyx.image_gen.interfaces import ImageGenerationProviderCredentials from onyx.llm.interfaces import LLM, LLMConfig @@ -207,7 +212,35 @@ def _construct_tools_impl( ) mcp_tool_cache: dict[int, dict[int, MCPTool]] = {} + mcp_preferences = ( + get_mcp_server_user_preferences( + list( + { + tool.mcp_server_id + for tool in configuration.tools + if tool.mcp_server_id is not None + } + ), + user.id, + db_session, + ) + if configuration.persona_id == DEFAULT_PERSONA_ID + else {} + ) user_oauth_token: str | None = user.live_oauth_token + openapi_preferences = ( + get_tool_user_preferences( + [ + tool.id + for tool in configuration.tools + if tool.openapi_schema is not None + ], + user.id, + db_session, + ) + if configuration.persona_id == DEFAULT_PERSONA_ID + else {} + ) search_settings = get_current_search_settings(db_session) document_index = get_default_document_index(search_settings, None) @@ -398,6 +431,9 @@ def _construct_tools_impl( # Handle custom tools elif db_tool_model.openapi_schema: + preference = openapi_preferences.get(db_tool_model.id) + if preference is not None and not preference.enabled_in_chat: + continue if not custom_tool_config: custom_tool_config = CustomToolConfig() @@ -458,6 +494,9 @@ def _construct_tools_impl( # Handle MCP tools elif db_tool_model.mcp_server_id: + preference = mcp_preferences.get(db_tool_model.mcp_server_id) + if preference is not None and not preference.enabled_in_chat: + continue if db_tool_model.mcp_server_id in mcp_tool_cache: tool_dict[db_tool_model.id] = [ mcp_tool_cache[db_tool_model.mcp_server_id][db_tool_model.id] @@ -466,6 +505,12 @@ def _construct_tools_impl( mcp_server = get_mcp_server_by_id(db_tool_model.mcp_server_id, db_session) + if ( + configuration.persona_id == DEFAULT_PERSONA_ID + and not user_can_access_mcp_server(user, mcp_server.id, db_session) + ): + continue + try: mcp_credentials = resolve_mcp_credentials(mcp_server, user, db_session) except MCPCredentialsError as e: diff --git a/backend/tests/external_dependency_unit/craft/test_mcp_config_resolution.py b/backend/tests/external_dependency_unit/craft/test_mcp_config_resolution.py index 3b2eb225ff..762810adb7 100644 --- a/backend/tests/external_dependency_unit/craft/test_mcp_config_resolution.py +++ b/backend/tests/external_dependency_unit/craft/test_mcp_config_resolution.py @@ -9,6 +9,7 @@ server key is stable + identifier-safe. from __future__ import annotations from collections.abc import Generator +from unittest.mock import MagicMock from uuid import uuid4 import pytest @@ -23,7 +24,11 @@ from onyx.db.enums import ( from onyx.db.mcp import ( create_connection_config, create_mcp_server__no_commit, + get_craft_enabled_mcp_servers, + get_mcp_server_user_preferences, + get_user_connection_config, update_mcp_server__no_commit, + upsert_mcp_server_user_preferences__no_commit, ) from onyx.db.models import MCPServer, Tool from onyx.server.features.build.sandbox.util.mcp_config import ( @@ -31,8 +36,9 @@ from onyx.server.features.build.sandbox.util.mcp_config import ( resolve_craft_mcp_servers, ) from onyx.server.features.mcp.api import get_craft_mcp_servers_for_user +from onyx.server.features.mcp.credentials import extract_connection_data from onyx.server.features.mcp.models import MCPConnectionData -from tests.external_dependency_unit.conftest import create_test_user +from tests.external_dependency_unit.conftest import create_test_user, delete_test_user @pytest.fixture @@ -92,6 +98,151 @@ def test_only_craft_enabled_servers_resolved_with_tool_curation( assert config.disabled_tools == ("delete_issue",) +def test_personal_preferences_are_independent_and_keep_disabled_server_visible( + db_session: Session, + craft_server: tuple[MCPServer, MCPServer], +) -> None: + craft, _ = craft_server + user = create_test_user(db_session, "mcp_preferences") + other = create_test_user(db_session, "mcp_preferences_other") + try: + assert get_mcp_server_user_preferences([craft.id], user.id, db_session) == {} + preference = upsert_mcp_server_user_preferences__no_commit( + craft.id, user.id, db_session, enabled_in_chat=False + ) + assert (preference.enabled_in_chat, preference.enabled_in_craft) == ( + False, + True, + ) + assert craft.server_url in { + s.url for s in resolve_craft_mcp_servers(db_session, user) + } + preference = upsert_mcp_server_user_preferences__no_commit( + craft.id, user.id, db_session, enabled_in_craft=False + ) + assert (preference.enabled_in_chat, preference.enabled_in_craft) == ( + False, + False, + ) + db_session.commit() + assert craft.server_url not in { + s.url for s in resolve_craft_mcp_servers(db_session, user) + } + assert craft.server_url in { + s.url for s in resolve_craft_mcp_servers(db_session, other) + } + assert craft.id in { + s.id for s in get_craft_enabled_mcp_servers(db_session, None) + } + response = get_craft_mcp_servers_for_user(db_session, user) + visible = next(s for s in response.mcp_servers if s.id == craft.id) + assert visible.craft_connected is True + assert visible.user_enabled_in_chat is False + assert visible.user_enabled_in_craft is False + preference = upsert_mcp_server_user_preferences__no_commit( + craft.id, user.id, db_session, enabled_in_chat=True + ) + assert (preference.enabled_in_chat, preference.enabled_in_craft) == ( + True, + False, + ) + finally: + db_session.rollback() + delete_test_user(db_session, user, other) + db_session.commit() + + +@pytest.mark.parametrize("persona_id", [0, 5]) +def test_chat_optout_blocks_default_chat_but_preserves_custom_agent( + db_session: Session, + craft_server: tuple[MCPServer, MCPServer], + monkeypatch: pytest.MonkeyPatch, + persona_id: int, +) -> None: + from onyx.chat.emitter import Emitter + from onyx.llm.interfaces import LLM + from onyx.tools import tool_constructor + from onyx.tools.models import PersonaToolConfiguration, ToolConfiguration + from onyx.tools.tool_implementations.search.search_tool import PersonaSearchInfo + + craft, _ = craft_server + user = create_test_user(db_session, "mcp_chat_preferences") + user.enable_memory_tool = False + tool = next(tool for tool in craft.current_actions if tool.enabled) + configuration = PersonaToolConfiguration( + persona_id=persona_id, + persona_name="Test agent", + tools=[ToolConfiguration.model_validate(tool)], + search=PersonaSearchInfo( + document_set_names=[], + search_start_date=None, + attached_document_ids=[], + hierarchy_node_ids=[], + ), + ) + monkeypatch.setattr(tool_constructor, "get_current_search_settings", MagicMock()) + monkeypatch.setattr(tool_constructor, "get_default_document_index", MagicMock()) + try: + upsert_mcp_server_user_preferences__no_commit( + craft.id, user.id, db_session, enabled_in_chat=False + ) + result = tool_constructor._construct_tools_impl( + configuration, + db_session, + MagicMock(spec=Emitter), + user, + MagicMock(spec=LLM), + allowed_tool_ids=[tool.id], + ) + assert set(result) == ({tool.id} if persona_id else set()) + if persona_id: + assert result[tool.id][0].name == tool.name + finally: + db_session.rollback() + delete_test_user(db_session, user) + db_session.commit() + + +def test_personal_optout_retains_stored_user_credentials( + db_session: Session, + craft_server: tuple[MCPServer, MCPServer], +) -> None: + craft, _ = craft_server + craft.auth_type = MCPAuthenticationType.API_TOKEN + craft.auth_performer = MCPAuthenticationPerformer.PER_USER + user = create_test_user(db_session, "mcp_retained_credentials") + config = create_connection_config( + {"headers": {"Authorization": "Bearer test-user-token"}}, + db_session, + mcp_server_id=craft.id, + user_email=user.email, + ) + try: + upsert_mcp_server_user_preferences__no_commit( + craft.id, user.id, db_session, enabled_in_chat=False, enabled_in_craft=False + ) + db_session.commit() + retained = get_user_connection_config(craft.id, user.email, db_session) + assert retained is not None + assert retained.id == config.id + assert extract_connection_data(retained, apply_mask=False)["headers"] == { + "Authorization": "Bearer test-user-token" + } + assert craft.id not in { + s.server_id for s in resolve_craft_mcp_servers(db_session, user) + } + upsert_mcp_server_user_preferences__no_commit( + craft.id, user.id, db_session, enabled_in_craft=True + ) + assert craft.id in { + s.server_id for s in resolve_craft_mcp_servers(db_session, user) + } + finally: + db_session.rollback() + delete_test_user(db_session, user) + db_session.commit() + + def test_per_user_server_emitted_only_once_credentials_connected( db_session: Session, craft_server: tuple[MCPServer, MCPServer], diff --git a/backend/tests/external_dependency_unit/mcp/test_mcp_oauth_dead_grant.py b/backend/tests/external_dependency_unit/mcp/test_mcp_oauth_dead_grant.py index d4c9d8db49..11c9bfc573 100644 --- a/backend/tests/external_dependency_unit/mcp/test_mcp_oauth_dead_grant.py +++ b/backend/tests/external_dependency_unit/mcp/test_mcp_oauth_dead_grant.py @@ -324,7 +324,7 @@ def test_oauth_reconnect_preserves_status_until_authorization_succeeds( force_reauthentication=True, ) if initialization_fails: - with pytest.raises(OnyxError, match="Failed to initialize OAuth client"): + with pytest.raises(OnyxError, match="OAuth setup failed"): asyncio.run(api._connect_oauth(request, db_session, False, user)) else: response = asyncio.run(api._connect_oauth(request, db_session, False, user)) diff --git a/backend/tests/external_dependency_unit/sandbox_proxy/test_mcp_server_resolver.py b/backend/tests/external_dependency_unit/sandbox_proxy/test_mcp_server_resolver.py index cd0d60f9bc..bfa2178a76 100644 --- a/backend/tests/external_dependency_unit/sandbox_proxy/test_mcp_server_resolver.py +++ b/backend/tests/external_dependency_unit/sandbox_proxy/test_mcp_server_resolver.py @@ -36,6 +36,7 @@ from onyx.db.mcp import ( create_mcp_server__no_commit, get_connection_config_by_id, update_mcp_server__no_commit, + upsert_mcp_server_user_preferences__no_commit, ) from onyx.db.models import MCPServer, OAuthAccount, User from onyx.external_apps.matching.engine import ( @@ -50,7 +51,7 @@ from onyx.sandbox_proxy.sandbox_identity.models import ResolvedSandbox from onyx.server.features.mcp.credentials import extract_connection_data from onyx.server.features.mcp.models import MCPConnectionData, MCPOAuthKeys from shared_configs.contextvars import POSTGRES_DEFAULT_SCHEMA -from tests.external_dependency_unit.conftest import create_test_user +from tests.external_dependency_unit.conftest import create_test_user, delete_test_user CraftServerFactory = Callable[..., MCPServer] @@ -495,6 +496,34 @@ def test_admin_managed_server_detail_points_to_admin( assert "Apps page in Craft" not in detail +def test_personal_craft_optout_keeps_host_claim_and_blocks_live_resolve( + db_session: Session, craft_server: CraftServerFactory +) -> None: + user = create_test_user(db_session, "mcp_resolver_personal_optout") + other = create_test_user(db_session, "mcp_resolver_personal_other") + server = craft_server( + auth_type=MCPAuthenticationType.NONE, + auth_performer=MCPAuthenticationPerformer.ADMIN, + ) + resolver = MCPServerResolver() + request = _request(_server_host(server)) + try: + assert resolver.claims(request, _ctx(user)) is True + assert resolver.resolve(request, _ctx(user)) == {} + upsert_mcp_server_user_preferences__no_commit( + server.id, user.id, db_session, enabled_in_craft=False + ) + db_session.commit() + assert resolver.claims(request, _ctx(user)) is True + with pytest.raises(CredentialUnavailableError, match="disabled MCP server"): + resolver.resolve(request, _ctx(user)) + assert resolver.resolve(request, _ctx(other)) == {} + finally: + db_session.rollback() + delete_test_user(db_session, user, other) + db_session.commit() + + def test_none_auth_claims_without_injecting( db_session: Session, craft_server: CraftServerFactory ) -> None: diff --git a/backend/tests/unit/external_apps/test_oauth_return_path.py b/backend/tests/unit/external_apps/test_oauth_return_path.py new file mode 100644 index 0000000000..406a470c6d --- /dev/null +++ b/backend/tests/unit/external_apps/test_oauth_return_path.py @@ -0,0 +1,80 @@ +from unittest.mock import MagicMock, patch +from urllib.parse import parse_qs, urlparse + +import pytest +from pydantic import ValidationError + +from onyx.error_handling.exceptions import OnyxError +from onyx.server.features.build.external_apps.models import OAuthCallbackRequest +from onyx.server.features.build.external_apps.oauth import ( + _ExternalAppOAuthAttemptPayload, + handle_external_app_oauth_callback, + start_external_app_oauth, +) +from tests.unit.external_apps.test_oauth_flow import _configure_route, _user + + +@pytest.mark.parametrize( + "path", ["https://example.com", "//example.com", "/\\example.com", "/app\n"] +) +def test_oauth_attempt_rejects_external_or_unsafe_return_paths(path: str) -> None: + with pytest.raises(ValidationError): + _ExternalAppOAuthAttemptPayload( + external_app_id=7, configuration_fingerprint="a" * 64, return_path=path + ) + + +def test_oauth_attempt_round_trips_member_completion_path() -> None: + path = ( + "/app/plugins/oauth/complete?pluginKind=external-app&pluginId=7&attempt=nonce" + ) + payload = _ExternalAppOAuthAttemptPayload( + external_app_id=7, configuration_fingerprint="a" * 64, return_path=path + ) + assert ( + _ExternalAppOAuthAttemptPayload.model_validate_json( + payload.model_dump_json() + ).return_path + == path + ) + + +def test_oauth_attempt_without_return_path_uses_default() -> None: + payload = _ExternalAppOAuthAttemptPayload( + external_app_id=7, configuration_fingerprint="a" * 64 + ) + assert payload.return_path is None + + +def test_oauth_start_rejects_unsafe_return_before_loading_app() -> None: + with patch( + "onyx.server.features.build.external_apps.oauth.get_external_app_by_id" + ) as get_app: + with pytest.raises(OnyxError): + start_external_app_oauth( + 7, user=MagicMock(), db_session=MagicMock(), return_path="//example.com" + ) + get_app.assert_not_called() + + +@pytest.mark.parametrize( + "return_path", [None, "/app/plugins/oauth/complete?attempt=nonce"] +) +def test_callback_preserves_return_path_after_saving_credentials( + monkeypatch: pytest.MonkeyPatch, return_path: str | None +) -> None: + app, _, db_session, _ = _configure_route(monkeypatch) + user = _user("owner") + start = start_external_app_oauth( + app.id, user=user, db_session=db_session, return_path=return_path + ) + state = parse_qs(urlparse(start.authorize_url).query)["state"][0] + request = OAuthCallbackRequest(code="code", state=state) + result = handle_external_app_oauth_callback( + request, user=user, db_session=db_session + ) + assert result.success is True + assert result.external_app_id == 7 + assert result.return_path == return_path + with pytest.raises(OnyxError): + handle_external_app_oauth_callback(request, user=user, db_session=db_session) diff --git a/backend/tests/unit/onyx/server/features/craft/external_apps/test_api_user_response.py b/backend/tests/unit/onyx/server/features/craft/external_apps/test_api_user_response.py index 7bee4acb3a..5be1e3ca41 100644 --- a/backend/tests/unit/onyx/server/features/craft/external_apps/test_api_user_response.py +++ b/backend/tests/unit/onyx/server/features/craft/external_apps/test_api_user_response.py @@ -4,9 +4,12 @@ import json from types import SimpleNamespace from typing import Any, cast +import pytest + from onyx.db.enums import ExternalAppType from onyx.db.models import ExternalApp, ExternalAppUserCredential from onyx.server.features.build.external_apps.api import _to_user_response +from onyx.server.features.build.external_apps.models import ExternalAppUserResponse from onyx.utils.encryption import mask_string from onyx.utils.sensitive import SensitiveValue @@ -24,6 +27,7 @@ def _external_app( auth_template: dict[str, Any], organization_credentials: dict[str, Any], app_type: ExternalAppType = ExternalAppType.CUSTOM, + upstream_url_patterns: list[str] | None = None, ) -> ExternalApp: return cast( ExternalApp, @@ -33,6 +37,7 @@ def _external_app( app_type=app_type, auth_template=auth_template, organization_credentials=_sensitive_dict(organization_credentials), + upstream_url_patterns=upstream_url_patterns or [], ), ) @@ -115,3 +120,113 @@ def test_user_response_uses_raw_presence_for_authentication() -> None: assert response.credential_values == { "access_token": mask_string("USER_ACCESS_TOKEN") } + + +@pytest.mark.parametrize( + ("app_type", "patterns", "origins"), + [ + ( + ExternalAppType.LINEAR, + [r"https://api\.linear\.app/.*"], + ["https://api.linear.app"], + ), + ( + ExternalAppType.CUSTOM, + ["https://api.linear.app/*"], + ["https://api.linear.app"], + ), + ( + ExternalAppType.ZENDESK, + [r"https://acme\-support\.zendesk\.com/api/v2/.*"], + ["https://acme-support.zendesk.com"], + ), + ( + ExternalAppType.LINEAR, + ["https://api.linear.app/.*", r"https://api\.linear.app/.*"], + [], + ), + ( + ExternalAppType.CUSTOM, + ["https://user:secret@API.Example.com:443/private?token=secret#fragment"], + ["https://api.example.com"], + ), + ( + ExternalAppType.CUSTOM, + ["https://api.example.com?token=secret", "https://api.example.com:8443/*"], + ["https://api.example.com", "https://api.example.com:8443"], + ), + ( + ExternalAppType.CUSTOM, + [ + "http://API.Example.com:80/a", + "http://api.example.com/*", + "https://api.example.com:443/*", + ], + ["http://api.example.com", "https://api.example.com"], + ), + ( + ExternalAppType.LINEAR, + [ + r"https://(?:api|www)\.linear\.app/.*", + r"https://api\.linear\.app?/.*", + r"https://.*\.linear\.app/.*", + ], + [], + ), + ( + ExternalAppType.CUSTOM, + [ + "https://*.example.com/*", + "https://api.example.com:bad/*", + "ftp://api.example.com/*", + "https://api%2eexample.com/*", + "https://api?.example.com/*", + "https://api\n.example.com/*", + ], + [], + ), + ( + ExternalAppType.CUSTOM, + ["https://[2001:db8::1]:443/*"], + ["https://[2001:db8::1]"], + ), + ], +) +def test_user_response_projects_only_clean_service_origins( + app_type: ExternalAppType, patterns: list[str], origins: list[str] +) -> None: + app = _external_app( + auth_template={}, + organization_credentials={"client_secret": "org-secret"}, + app_type=app_type, + upstream_url_patterns=patterns, + ) + + response = _to_user_response(app, None).model_dump() + + assert response["service_urls"] == origins + assert set(response) == { + "id", + "name", + "app_type", + "credential_keys", + "credential_values", + "authenticated", + "supports_oauth", + "service_urls", + } + assert "org-secret" not in json.dumps(response) + + +def test_user_response_service_origins_default_to_empty() -> None: + response = ExternalAppUserResponse( + id=1, + name="Test App", + app_type=ExternalAppType.CUSTOM, + credential_keys=[], + credential_values={}, + authenticated=True, + supports_oauth=False, + ) + + assert response.service_urls == [] diff --git a/backend/tests/unit/onyx/server/features/mcp/test_mcp_default_chat_availability.py b/backend/tests/unit/onyx/server/features/mcp/test_mcp_default_chat_availability.py new file mode 100644 index 0000000000..8ac8f6fa75 --- /dev/null +++ b/backend/tests/unit/onyx/server/features/mcp/test_mcp_default_chat_availability.py @@ -0,0 +1,312 @@ +from unittest.mock import MagicMock + +import pytest +from sqlalchemy import Boolean, Column, Integer, MetaData, Table, create_engine +from sqlalchemy.orm import Session + +from onyx.db import persona as persona_db +from onyx.db.models import MCPServer, Persona, Tool, User +from onyx.server.features.default_assistant import api as default_api +from onyx.server.features.default_assistant.models import DefaultAssistantUpdateRequest +from onyx.server.features.mcp.models import MCPConnectionStatus + + +@pytest.mark.parametrize("nondefault_builtin_deleted", [True, False]) +def test_default_chat_server_summary_uses_visible_tool_attachments( + nondefault_builtin_deleted: bool, +) -> None: + from onyx.db.mcp import get_default_chat_mcp_server_ids + + metadata = MetaData() + personas = Table( + "persona", + metadata, + Column("id", Integer, primary_key=True), + Column("builtin_persona", Boolean), + Column("deleted", Boolean), + ) + tools = Table( + "tool", + metadata, + Column("id", Integer, primary_key=True), + Column("mcp_server_id", Integer), + Column("enabled", Boolean), + ) + attachments = Table( + "persona__tool", + metadata, + Column("persona_id", Integer), + Column("tool_id", Integer), + ) + engine = create_engine("sqlite://") + metadata.create_all(engine) + with Session(engine) as db: + db.execute( + personas.insert(), + [ + {"id": 0, "builtin_persona": True, "deleted": False}, + {"id": 1, "builtin_persona": False, "deleted": False}, + { + "id": 2, + "builtin_persona": True, + "deleted": nondefault_builtin_deleted, + }, + ], + ) + db.execute( + tools.insert(), + [ + {"id": 1, "mcp_server_id": 42, "enabled": True}, + {"id": 2, "mcp_server_id": 42, "enabled": False}, + {"id": 3, "mcp_server_id": 43, "enabled": False}, + {"id": 4, "mcp_server_id": 44, "enabled": True}, + {"id": 5, "mcp_server_id": 45, "enabled": True}, + {"id": 6, "mcp_server_id": 46, "enabled": True}, + {"id": 7, "mcp_server_id": None, "enabled": True}, + ], + ) + db.execute( + attachments.insert(), + [ + {"persona_id": persona_id, "tool_id": tool_id} + for persona_id, tool_id in [ + (0, 1), + (0, 2), + (0, 3), + (1, 4), + (2, 5), + (0, 6), + (0, 7), + ] + ], + ) + + assert get_default_chat_mcp_server_ids([42, 43, 44, 45], db) == {42, 43} + assert get_default_chat_mcp_server_ids([44, 45], db) == set() + assert get_default_chat_mcp_server_ids([], db) == set() + + +def test_scoped_manager_server_summary_reports_only_visible_servers( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from onyx.db.enums import MCPServerStatus, PermissionAuthority + from onyx.server.features.mcp import api + from onyx.server.features.mcp.models import MCPServer as MCPServerResponse + + db = MagicMock(spec=Session) + user = User(email="manager@example.com") + servers = [ + MCPServer(id=42, owner=user.email), + MCPServer(id=43, owner="other@example.com"), + MCPServer(id=44, owner="other@example.com"), + ] + selected_servers = MagicMock(return_value={43}) + monkeypatch.setattr(api, "get_all_mcp_servers", lambda _: servers) + monkeypatch.setattr(api, "has_permission", lambda *_: PermissionAuthority.SCOPED) + monkeypatch.setattr(api, "get_scoped_groups", lambda *_: []) + monkeypatch.setattr(api, "get_mcp_server_ids_connected_to_groups", lambda *_: {43}) + monkeypatch.setattr(api, "get_default_chat_mcp_server_ids", selected_servers) + monkeypatch.setattr(api, "can_manage_mcp_server", lambda *_: False) + monkeypatch.setattr( + api, + "_db_mcp_server_to_api_mcp_server", + lambda server, *_, **__: MCPServerResponse( + id=server.id, + name="Server", + server_url="", + owner="", + status=MCPServerStatus.CREATED, + connection_status=MCPConnectionStatus.NOT_CONNECTED, + ), + ) + + response = api.get_mcp_servers_for_admin(db, user) + + selected_servers.assert_called_once_with([42, 43], db) + assert [ + (server.id, server.default_chat_enabled) for server in response.mcp_servers + ] == [ + (42, False), + (43, True), + ] + + +def test_member_server_summary_reports_only_accessible_servers( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from onyx.db.enums import MCPServerStatus + from onyx.server.features.mcp import api + from onyx.server.features.mcp.models import MCPServer as MCPServerResponse + + db = MagicMock(spec=Session) + user = User(email="member@example.com") + servers = [MCPServer(id=42), MCPServer(id=43)] + selected_servers = MagicMock(return_value={43}) + monkeypatch.setattr(api, "get_mcp_servers_accessible_to_user", lambda *_: servers) + monkeypatch.setattr(api, "get_default_chat_mcp_server_ids", selected_servers) + monkeypatch.setattr( + api, + "_db_mcp_server_to_api_mcp_server", + lambda server, *_, **__: MCPServerResponse( + id=server.id, + name="Server", + server_url="", + owner="", + status=MCPServerStatus.CREATED, + connection_status=MCPConnectionStatus.NOT_CONNECTED, + ), + ) + + response = api.get_mcp_servers_for_user(db, user) + + selected_servers.assert_called_once_with([42, 43], db) + assert [ + (server.id, server.default_chat_enabled) for server in response.mcp_servers + ] == [(42, False), (43, True)] + + +def tool(tool_id: int, server_id: int | None, name: str) -> Tool: + return Tool(id=tool_id, name=name, mcp_server_id=server_id, enabled=True) + + +def test_default_chat_views_share_existing_tool_attachments( + monkeypatch: pytest.MonkeyPatch, +) -> None: + selected = tool(1, 42, "selected") + other_mcp = tool(2, 43, "other_server") + built_in = tool(3, None, "search") + unselected = tool(4, 42, "unselected") + default = Persona(id=0, tools=[selected, other_mcp, built_in]) + custom = Persona(id=5, tools=[selected, unselected]) + db = MagicMock(spec=Session) + user = User() + monkeypatch.setattr(persona_db, "get_default_assistant", lambda _: default) + monkeypatch.setattr(default_api, "get_default_assistant", lambda _: default) + monkeypatch.setattr(persona_db, "should_expose_tool_to_fe", lambda _: True) + + def save_selection(tools: list[Tool]) -> None: + db.query.return_value.filter.return_value.one_or_none.side_effect = tools + response = default_api.update_default_assistant( + DefaultAssistantUpdateRequest(tool_ids=[item.id for item in tools]), + user, + db, + ) + assert response.tool_ids == [item.id for item in tools] + assert ( + default_api.get_default_assistant_configuration(user, db).tool_ids + == response.tool_ids + ) + assert custom.tools == [selected, unselected] + + save_selection([other_mcp, built_in]) + save_selection([selected, other_mcp, built_in]) + save_selection([selected, unselected, other_mcp, built_in]) + assert default.system_prompt is None + assert db.commit.call_count == 3 + + +def test_default_chat_availability_has_no_separate_server_field() -> None: + from onyx.server.features.mcp.models import MCPServer as MCPServerResponse + + assert "available_in_default_chat" not in MCPServer.__table__.columns + assert "available_in_default_chat" not in MCPServerResponse.model_fields + + +def test_discovery_does_not_change_default_or_custom_agent_selection( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from mcp.types import Tool as DiscoveredTool + + from onyx.server.features.mcp import api + + existing = tool(1, 42, "old") + default = Persona(id=0, tools=[existing]) + custom = Persona(id=5, tools=[existing]) + db = MagicMock(spec=Session) + monkeypatch.setattr(api, "get_tools_by_mcp_server_id", lambda *_, **__: [existing]) + monkeypatch.setattr(api, "create_tool__no_commit", lambda **_: tool(2, 42, "new")) + + api._sync_mcp_server_tools( + 42, + [ + DiscoveredTool(name="old", inputSchema={}), + DiscoveredTool(name="new", inputSchema={}), + ], + db, + ) + + assert default.tools == [existing] + assert custom.tools == [existing] + db.scalar.assert_not_called() + db.scalars.assert_not_called() + db.commit.assert_called_once() + + +@pytest.mark.parametrize( + "persona_id,accessible,expected_resolution", + [ + (0, False, False), + (0, True, True), + (5, False, True), + ], +) +def test_default_chat_checks_access_before_resolving_credentials( + monkeypatch: pytest.MonkeyPatch, + persona_id: int, + accessible: bool, + expected_resolution: bool, +) -> None: + from onyx.chat.emitter import Emitter + from onyx.db.models import User + from onyx.llm.interfaces import LLM + from onyx.server.features.mcp.credentials import MCPCredentialsError + from onyx.tools import tool_constructor + from onyx.tools.models import PersonaToolConfiguration, ToolConfiguration + from onyx.tools.tool_implementations.search.search_tool import PersonaSearchInfo + + server = MCPServer(id=42) + configuration = PersonaToolConfiguration( + persona_id=persona_id, + persona_name="Test", + tools=[ + ToolConfiguration( + id=1, + name="tool", + description=None, + display_name=None, + in_code_tool_id=None, + enabled=True, + openapi_schema=None, + mcp_input_schema=None, + custom_headers=None, + passthrough_auth=False, + mcp_server_id=42, + oauth_config_id=None, + ) + ], + search=PersonaSearchInfo( + document_set_names=[], + search_start_date=None, + attached_document_ids=[], + hierarchy_node_ids=[], + ), + ) + resolve = MagicMock(side_effect=MCPCredentialsError("No credential")) + monkeypatch.setattr(tool_constructor, "get_current_search_settings", MagicMock()) + monkeypatch.setattr(tool_constructor, "get_default_document_index", MagicMock()) + monkeypatch.setattr(tool_constructor, "get_mcp_server_by_id", lambda *_: server) + monkeypatch.setattr( + tool_constructor, "user_can_access_mcp_server", lambda *_: accessible + ) + monkeypatch.setattr(tool_constructor, "resolve_mcp_credentials", resolve) + + result = tool_constructor._construct_tools_impl( + configuration=configuration, + db_session=MagicMock(spec=Session), + emitter=MagicMock(spec=Emitter), + user=User(enable_memory_tool=False), + llm=MagicMock(spec=LLM), + ) + + assert result == {} + assert resolve.called is expected_resolution diff --git a/backend/tests/unit/onyx/server/features/mcp/test_mcp_oauth_connect.py b/backend/tests/unit/onyx/server/features/mcp/test_mcp_oauth_connect.py index 45b388c5fa..1eda80774b 100644 --- a/backend/tests/unit/onyx/server/features/mcp/test_mcp_oauth_connect.py +++ b/backend/tests/unit/onyx/server/features/mcp/test_mcp_oauth_connect.py @@ -7,6 +7,7 @@ from uuid import uuid4 import pytest from mcp.shared.auth import OAuthClientInformationFull +from onyx.db import mcp as mcp_db from onyx.db.enums import ( MCPAuthenticationPerformer, MCPAuthenticationType, @@ -267,7 +268,7 @@ def test_server_url_change_invalidates_old_credentials_without_losing_new_client "delete_all_user_connection_configs_for_server_no_commit", delete_user_configs, ) - monkeypatch.setattr(api, "delete_connection_config", delete_admin_config) + monkeypatch.setattr(mcp_db, "delete_connection_config", delete_admin_config) monkeypatch.setattr(api, "update_mcp_server__no_commit", update_server) def persist_admin_config( diff --git a/backend/tests/unit/onyx/server/features/mcp/test_mcp_oauth_discovery.py b/backend/tests/unit/onyx/server/features/mcp/test_mcp_oauth_discovery.py new file mode 100644 index 0000000000..3dada6c4d8 --- /dev/null +++ b/backend/tests/unit/onyx/server/features/mcp/test_mcp_oauth_discovery.py @@ -0,0 +1,331 @@ +import asyncio + +import httpx +import pytest +from mcp.client.auth.exceptions import OAuthRegistrationError + +from onyx.server.features.mcp import oauth_discovery +from onyx.server.features.mcp.models import MCPOAuthDiscoveryRequest + + +@pytest.mark.parametrize( + "registration,cimd,expected", + [ + (None, False, False), + ("https://issuer.example/register", False, True), + (None, True, True), + ], +) +def test_discovery_reads_capabilities_without_registration( + monkeypatch: pytest.MonkeyPatch, + registration: str | None, + cimd: bool, + expected: bool, +) -> None: + requests: list[httpx.Request] = [] + + def respond(request: httpx.Request) -> httpx.Response: + requests.append(request) + assert request.method == "GET" + assert "authorization" not in request.headers + if request.url.path == "/mcp": + return httpx.Response( + 401, + headers={ + "www-authenticate": 'Bearer resource_metadata="https://mcp.example/metadata"' + }, + ) + if request.url.path == "/metadata": + return httpx.Response( + 200, + json={ + "resource": "https://mcp.example/mcp", + "authorization_servers": ["https://issuer.example"], + }, + ) + return httpx.Response( + 200, + json={ + "issuer": "https://issuer.example", + "authorization_endpoint": "https://issuer.example/authorize", + "token_endpoint": "https://issuer.example/token", + "response_types_supported": ["code"], + "registration_endpoint": registration, + "client_id_metadata_document_supported": cimd, + }, + ) + + monkeypatch.setattr( + oauth_discovery, + "mcp_ssrf_httpx_client_factory", + lambda **_: httpx.AsyncClient(transport=httpx.MockTransport(respond)), + ) + result = asyncio.run( + oauth_discovery.discover_oauth_capabilities("https://mcp.example/mcp") + ) + assert result.status == "metadata_found" + assert result.supports_automatic_registration is expected + assert result.authorization_endpoint == "https://issuer.example/authorize" + assert len(requests) == 3 + + +@pytest.mark.parametrize("unreachable", [False, True]) +def test_unknown_metadata_does_not_claim_registration_is_unsupported( + monkeypatch: pytest.MonkeyPatch, unreachable: bool +) -> None: + def respond(request: httpx.Request) -> httpx.Response: + if unreachable: + raise httpx.ConnectError("private details", request=request) + return httpx.Response(404, text="private details") + + monkeypatch.setattr( + oauth_discovery, + "mcp_ssrf_httpx_client_factory", + lambda **_: httpx.AsyncClient(transport=httpx.MockTransport(respond)), + ) + result = asyncio.run( + oauth_discovery.discover_oauth_capabilities("https://mcp.example/mcp") + ) + assert result.status == ("unreachable" if unreachable else "metadata_missing") + assert result.supports_automatic_registration is None + + +def test_registration_error_is_sanitized_through_nested_task_groups() -> None: + upstream = OAuthRegistrationError( + "Registration failed: 404 secret=private" + ) + error = oauth_discovery.oauth_setup_error( + ExceptionGroup("transport", [ExceptionGroup("oauth", [upstream])]) + ) + assert error.extra == {"oauth_error": "client_registration_failed"} + assert "own OAuth client" in error.detail + assert "private" not in error.detail + assert "html" not in error.detail + + +def test_discovery_checks_server_management_before_network( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from unittest.mock import AsyncMock, MagicMock + + from onyx.error_handling.error_codes import OnyxErrorCode + from onyx.error_handling.exceptions import OnyxError + from onyx.server.features.mcp import api + + server = MagicMock() + monkeypatch.setattr(api, "get_session_with_current_tenant", MagicMock()) + monkeypatch.setattr(api, "get_mcp_server_by_id", lambda *_: server) + denied = MagicMock(side_effect=OnyxError(OnyxErrorCode.INSUFFICIENT_PERMISSIONS)) + discover = AsyncMock() + monkeypatch.setattr(api, "_ensure_mcp_server_owner_or_admin", denied) + monkeypatch.setattr(api, "discover_oauth_capabilities", discover) + with pytest.raises(OnyxError): + asyncio.run( + api.discover_admin_oauth( + MCPOAuthDiscoveryRequest(server_id=42), + MagicMock(), + ) + ) + discover.assert_not_called() + + +@pytest.mark.parametrize( + "metadata_kind", ["wrong_resource", "empty_servers", "invalid_json"] +) +def test_invalid_resource_metadata_never_follows_untrusted_issuer( + monkeypatch: pytest.MonkeyPatch, metadata_kind: str +) -> None: + requested_hosts: list[str] = [] + + def respond(request: httpx.Request) -> httpx.Response: + requested_hosts.append(request.url.host) + if request.url.path == "/mcp": + return httpx.Response(401) + if "oauth-protected-resource" in request.url.path: + if metadata_kind == "invalid_json": + return httpx.Response(200, text="not metadata") + return httpx.Response( + 200, + json={ + "resource": "https://other.example/mcp" + if metadata_kind == "wrong_resource" + else "https://mcp.example/mcp", + "authorization_servers": [] + if metadata_kind == "empty_servers" + else ["https://untrusted.example"], + }, + ) + return httpx.Response(404) + + monkeypatch.setattr( + oauth_discovery, + "mcp_ssrf_httpx_client_factory", + lambda **_: httpx.AsyncClient(transport=httpx.MockTransport(respond)), + ) + result = asyncio.run( + oauth_discovery.discover_oauth_capabilities("https://mcp.example/mcp") + ) + assert result.status == "metadata_missing" + assert result.supports_automatic_registration is None + assert set(requested_hosts) == {"mcp.example"} + + +@pytest.mark.parametrize("failure", ["server_error", "timeout"]) +def test_provider_outage_is_unknown_not_unsupported( + monkeypatch: pytest.MonkeyPatch, failure: str +) -> None: + def respond(request: httpx.Request) -> httpx.Response: + if failure == "timeout": + raise httpx.ReadTimeout("secret upstream URL", request=request) + return httpx.Response(503) + + monkeypatch.setattr( + oauth_discovery, + "mcp_ssrf_httpx_client_factory", + lambda **_: httpx.AsyncClient(transport=httpx.MockTransport(respond)), + ) + result = asyncio.run( + oauth_discovery.discover_oauth_capabilities("https://mcp.example/mcp") + ) + assert result.status == "unreachable" + assert result.supports_automatic_registration is None + + +@pytest.mark.parametrize("hop", ["challenge", "redirect"]) +def test_discovery_blocks_metadata_network_hops_without_exposing_internal_url( + monkeypatch: pytest.MonkeyPatch, hop: str +) -> None: + from onyx.server.features.mcp import ssrf + from onyx.server.security.store import _build_env_defaults + + monkeypatch.setattr(ssrf, "get_security_settings", _build_env_defaults) + requests: list[str] = [] + blocked_url = "http://169.254.169.254/latest/meta-data/" + + async def respond( + _transport: httpx.AsyncHTTPTransport, request: httpx.Request + ) -> httpx.Response: + requests.append(str(request.url)) + if hop == "redirect": + return httpx.Response(302, headers={"location": blocked_url}) + return httpx.Response( + 401, + headers={"www-authenticate": f'Bearer resource_metadata="{blocked_url}"'}, + ) + + monkeypatch.setattr(httpx.AsyncHTTPTransport, "handle_async_request", respond) + result = asyncio.run( + oauth_discovery.discover_oauth_capabilities("https://8.8.8.8/mcp") + ) + assert result.status == "unreachable" + assert result.supports_automatic_registration is None + assert requests == ["https://8.8.8.8/mcp"] + assert blocked_url not in result.model_dump_json() + + +@pytest.mark.parametrize( + "error,reason", + [ + (httpx.ConnectError("secret"), "connection_failed"), + (TimeoutError("secret"), "connection_failed"), + (RuntimeError("secret"), "oauth_failed"), + (ExceptionGroup("secret", [RuntimeError("secret")]), "oauth_failed"), + ( + ExceptionGroup( + "secret", [RuntimeError("secret"), OAuthRegistrationError("secret")] + ), + "client_registration_failed", + ), + ], +) +def test_oauth_error_classification_never_exposes_upstream_details( + error: Exception, reason: str +) -> None: + result = oauth_discovery.oauth_setup_error(error) + assert result.extra == {"oauth_error": reason} + assert "secret" not in result.detail + + +@pytest.mark.parametrize("outcome", ["allowed", "forbidden", "missing"]) +def test_discovery_route_authorizes_in_worker_and_closes_session_before_http( + monkeypatch: pytest.MonkeyPatch, outcome: str +) -> None: + import threading + from contextlib import contextmanager + from types import SimpleNamespace + from typing import Iterator, cast + from unittest.mock import MagicMock + + from fastapi import FastAPI + from fastapi.routing import APIRoute + from fastapi.testclient import TestClient + from sqlalchemy.orm import Session + + from onyx.db.models import MCPServer, User + from onyx.error_handling.exceptions import register_onyx_exception_handlers + from onyx.server.features.mcp import api + from onyx.server.features.mcp.models import MCPOAuthDiscoveryResponse + + app = FastAPI() + app.include_router(api.admin_router) + register_onyx_exception_handlers(app) + route = next( + route + for route in app.routes + if isinstance(route, APIRoute) and route.path == "/admin/mcp/oauth/discover" + ) + user_dependency = route.dependant.dependencies[0].call + assert user_dependency is not None + user = cast(User, SimpleNamespace(email="admin@example.com")) + app.dependency_overrides[user_dependency] = lambda: user + events: list[str] = [] + worker_threads: list[int] = [] + http_threads: list[int] = [] + + @contextmanager + def session() -> Iterator[Session]: + worker_threads.append(threading.get_ident()) + events.append("session_open") + try: + yield MagicMock(spec=Session) + finally: + events.append("session_closed") + + def lookup(_id: int, _db: Session) -> MCPServer: + worker_threads.append(threading.get_ident()) + if outcome == "missing": + raise ValueError("not found") + return cast( + MCPServer, + SimpleNamespace( + server_url="https://example.com/mcp", name="MCP", owner=user.email + ), + ) + + def can_manage(_user: User, _server: MCPServer) -> bool: + worker_threads.append(threading.get_ident()) + return outcome != "forbidden" + + async def discover(url: str) -> MCPOAuthDiscoveryResponse: + assert url == "https://example.com/mcp" + assert events == ["session_open", "session_closed"] + http_threads.append(threading.get_ident()) + return MCPOAuthDiscoveryResponse(status="metadata_missing") + + monkeypatch.setattr(api, "get_session_with_current_tenant", session) + monkeypatch.setattr(api, "get_mcp_server_by_id", lookup) + monkeypatch.setattr(api, "can_manage_mcp_server", can_manage) + monkeypatch.setattr(api, "discover_oauth_capabilities", discover) + with TestClient(app) as client: + response = client.post("/admin/mcp/oauth/discover", json={"server_id": 42}) + assert ( + response.status_code + == {"allowed": 200, "forbidden": 403, "missing": 404}[outcome] + ) + assert events == ["session_open", "session_closed"] + assert len(set(worker_threads)) == 1 + if outcome == "allowed": + assert http_threads[0] != worker_threads[0] + assert response.json()["supports_automatic_registration"] is None + else: + assert http_threads == [] diff --git a/backend/tests/unit/onyx/server/features/mcp/test_mcp_server_craft_availability.py b/backend/tests/unit/onyx/server/features/mcp/test_mcp_server_craft_availability.py new file mode 100644 index 0000000000..ec4e065d5b --- /dev/null +++ b/backend/tests/unit/onyx/server/features/mcp/test_mcp_server_craft_availability.py @@ -0,0 +1,38 @@ +from unittest.mock import MagicMock + +import pytest +from sqlalchemy.orm import Session + +from onyx.db.enums import MCPServerStatus +from onyx.db.models import MCPServer, User +from onyx.server.features.mcp import api +from onyx.server.features.mcp.models import MCPServerSimpleCreateRequest + + +@pytest.mark.parametrize("available", [None, False, True]) +def test_create_persists_craft_availability_atomically( + monkeypatch: pytest.MonkeyPatch, available: bool | None +) -> None: + db = MagicMock(spec=Session) + user = User(email="admin@example.com") + added: list[MCPServer] = [] + + def add(server: MCPServer) -> None: + server.id = 42 + server.status = MCPServerStatus.CREATED + added.append(server) + + db.add.side_effect = add + monkeypatch.setattr(api, "_validate_mcp_server_url", lambda *_, **__: None) + monkeypatch.setattr(api, "_apply_mcp_server_access", lambda **_: None) + request = MCPServerSimpleCreateRequest( + name="Example", server_url="https://example.com/mcp" + ) + if available is not None: + request.available_in_craft = available + + result = api.create_mcp_server_simple(request, db, user) + + assert added[0].available_in_craft is (available is True) + assert result.available_in_craft is (available is True) + db.commit.assert_called_once_with() diff --git a/backend/tests/unit/onyx/server/features/mcp/test_mcp_server_url_update.py b/backend/tests/unit/onyx/server/features/mcp/test_mcp_server_url_update.py new file mode 100644 index 0000000000..6c223dbac9 --- /dev/null +++ b/backend/tests/unit/onyx/server/features/mcp/test_mcp_server_url_update.py @@ -0,0 +1,125 @@ +from types import SimpleNamespace +from typing import cast +from unittest.mock import MagicMock +from uuid import uuid4 + +import pytest +from sqlalchemy.orm import Session + +from onyx.db import mcp as mcp_db +from onyx.db.enums import ( + MCPAuthenticationPerformer, + MCPAuthenticationType, + MCPOAuthProviderMode, + MCPServerStatus, + MCPTransport, + Permission, +) +from onyx.db.models import MCPConnectionConfig, MCPServer, User +from onyx.server.features.mcp import api +from onyx.server.features.mcp.models import ( + MCPConnectionStatus, + MCPServerSimpleUpdateRequest, +) + + +@pytest.mark.parametrize( + "auth_type,performer", + [ + (MCPAuthenticationType.OAUTH, MCPAuthenticationPerformer.PER_USER), + (MCPAuthenticationType.API_TOKEN, MCPAuthenticationPerformer.ADMIN), + (MCPAuthenticationType.API_TOKEN, MCPAuthenticationPerformer.PER_USER), + (MCPAuthenticationType.NONE, MCPAuthenticationPerformer.PER_USER), + ], +) +@pytest.mark.parametrize("change_url", [False, True]) +def test_details_url_change_invalidates_credentials_before_returning_fresh_settings( + monkeypatch: pytest.MonkeyPatch, + auth_type: MCPAuthenticationType, + performer: MCPAuthenticationPerformer, + change_url: bool, +) -> None: + config = MCPConnectionConfig( + id=81, config={"headers": {"Authorization": "Bearer old-secret"}} + ) + server = MCPServer( + id=42, + name="MCP", + description="Tools", + server_url="https://old.example/mcp", + owner="owner@example.com", + auth_type=auth_type, + auth_performer=performer, + oauth_provider_mode=MCPOAuthProviderMode.AUTO_DISCOVERY, + transport=MCPTransport.STREAMABLE_HTTP, + status=MCPServerStatus.CONNECTED, + is_public=True, + available_in_craft=True, + admin_connection_config=config, + admin_connection_config_id=config.id, + ) + user = cast(User, SimpleNamespace(id=uuid4(), email=server.owner)) + db = MagicMock(spec=Session) + old_user, current_user = uuid4(), uuid4() + affected = MagicMock( + side_effect=[{old_user}, {current_user}] if change_url else [{current_user}] + ) + delete_users = MagicMock() + reload_sessions = MagicMock() + monkeypatch.setattr(api, "get_mcp_server_by_id", lambda *_: server) + monkeypatch.setattr(mcp_db, "get_mcp_server_by_id", lambda *_: server) + monkeypatch.setattr(mcp_db, "get_connection_config_by_id", lambda *_: config) + monkeypatch.setattr(api, "_validate_mcp_server_url", lambda *_, **__: None) + monkeypatch.setattr(api, "can_manage_mcp_server", lambda *_: True) + monkeypatch.setattr( + api, + "get_effective_permissions", + lambda *_: {Permission.FULL_ADMIN_PANEL_ACCESS}, + ) + monkeypatch.setattr(api, "affected_user_ids_for_mcp_server", affected) + monkeypatch.setattr( + api, "delete_all_user_connection_configs_for_server_no_commit", delete_users + ) + monkeypatch.setattr(api, "_hot_reload_craft_sessions", reload_sessions) + monkeypatch.setattr(api, "get_user_connection_config", lambda *_: None) + monkeypatch.setattr(api, "get_action_policies", lambda *_: {}) + result = api.update_mcp_server_simple( + server.id, + MCPServerSimpleUpdateRequest( + name="Updated MCP", + server_url="https://new.example/mcp" if change_url else server.server_url, + ), + db, + user, + ) + assert result.name == "Updated MCP" + assert result.tool_count == 0 + if change_url: + delete_users.assert_called_once_with(server.id, db) + db.delete.assert_called_once_with(config) + assert server.admin_connection_config_id is None + assert server.admin_connection_config is None + assert result.status == ( + MCPServerStatus.CREATED + if auth_type == MCPAuthenticationType.NONE + else MCPServerStatus.AWAITING_AUTH + ) + assert result.server_url == "https://new.example/mcp" + fresh = api._db_mcp_server_to_api_mcp_server( + server, db, request_user=user, include_auth_config=True + ) + assert fresh.admin_credentials is None + assert fresh.user_credentials is None + assert fresh.auth_template is None + if auth_type != MCPAuthenticationType.NONE: + assert fresh.connection_status is MCPConnectionStatus.NOT_CONNECTED + assert "old-secret" not in fresh.model_dump_json() + reload_sessions.assert_called_once_with({old_user, current_user}, db) + else: + delete_users.assert_not_called() + db.delete.assert_not_called() + assert server.admin_connection_config is config + assert server.admin_connection_config_id == 81 + assert result.status == MCPServerStatus.CONNECTED + reload_sessions.assert_called_once_with({current_user}, db) + db.commit.assert_called_once() diff --git a/backend/tests/unit/onyx/server/features/mcp/test_mcp_ssrf.py b/backend/tests/unit/onyx/server/features/mcp/test_mcp_ssrf.py index e3f6172b15..e98f13a537 100644 --- a/backend/tests/unit/onyx/server/features/mcp/test_mcp_ssrf.py +++ b/backend/tests/unit/onyx/server/features/mcp/test_mcp_ssrf.py @@ -233,3 +233,41 @@ def test_server_url_allows_http() -> None: api._validate_mcp_server_url( "http://8.8.8.8/mcp", "server_url", require_https=False ) + + +def test_async_guard_offloads_validation_and_preserves_tenant_context( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import threading + + from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR + + validation_threads: list[int] = [] + tenant_ids: list[str | None] = [] + loop_threads: list[int] = [] + + def validate(url: str) -> str: + validation_threads.append(threading.get_ident()) + tenant_ids.append(CURRENT_TENANT_ID_CONTEXTVAR.get()) + return url + + async def send( + _transport: httpx.AsyncHTTPTransport, request: httpx.Request + ) -> httpx.Response: + loop_threads.append(threading.get_ident()) + return httpx.Response(200, request=request) + + async def run() -> None: + token = CURRENT_TENANT_ID_CONTEXTVAR.set("discovery-test-tenant") + try: + async with mcp_ssrf.mcp_ssrf_httpx_client_factory() as client: + await client.get("https://example.com/mcp") + finally: + CURRENT_TENANT_ID_CONTEXTVAR.reset(token) + + monkeypatch.setattr(mcp_ssrf, "validate_mcp_outbound_url", validate) + monkeypatch.setattr(httpx.AsyncHTTPTransport, "handle_async_request", send) + asyncio.run(run()) + assert len(validation_threads) == 1 + assert validation_threads[0] != loop_threads[0] + assert tenant_ids == ["discovery-test-tenant"] diff --git a/backend/tests/unit/onyx/server/features/mcp/test_setup_probe.py b/backend/tests/unit/onyx/server/features/mcp/test_setup_probe.py new file mode 100644 index 0000000000..a0f9d8e8e5 --- /dev/null +++ b/backend/tests/unit/onyx/server/features/mcp/test_setup_probe.py @@ -0,0 +1,415 @@ +import asyncio +import json +from collections.abc import AsyncIterator + +import httpx +import pytest + +from onyx.db.enums import MCPTransport +from onyx.server.features.mcp import oauth_discovery, setup_probe +from onyx.utils.url import SSRFException + + +def _rpc_result(payload: dict) -> dict: + if payload["method"] == "initialize": + return { + "protocolVersion": "2025-03-26", + "capabilities": {"tools": {}}, + "serverInfo": {"name": "Support tools", "version": "1"}, + } + assert payload["method"] == "tools/list" + return {"tools": [{"name": "lookup_ticket", "inputSchema": {"type": "object"}}]} + + +def test_public_requires_mcp_initialization_and_lists_tools_without_credentials( + monkeypatch: pytest.MonkeyPatch, +) -> None: + methods: list[str] = [] + events: list[setup_probe.MCPProbeProgress] = [] + + async def progress(event: setup_probe.MCPProbeProgress) -> None: + if event["stage"] == "authentication" and event["state"] == "pending": + assert methods[0] == "initialize" + assert "tools/list" not in methods + events.append(event) + + def respond(request: httpx.Request) -> httpx.Response: + assert "authorization" not in request.headers + if request.method == "DELETE": + return httpx.Response(204) + if request.method == "GET": + return httpx.Response(405) + payload = json.loads(request.content) + methods.append(payload["method"]) + if "id" not in payload: + return httpx.Response(202) + return httpx.Response( + 200, + json={ + "jsonrpc": "2.0", + "id": payload["id"], + "result": _rpc_result(payload), + }, + ) + + monkeypatch.setattr( + setup_probe, + "mcp_ssrf_httpx_client_factory", + lambda **_: httpx.AsyncClient(transport=httpx.MockTransport(respond)), + ) + result = asyncio.run( + setup_probe.probe_mcp_server("https://mcp.example/mcp", progress) + ) + assert result.model_dump() == { + "kind": "public", + "transport": MCPTransport.STREAMABLE_HTTP, + "tool_count": 1, + "server_name": "Support tools", + } + assert methods == ["initialize", "notifications/initialized", "tools/list"] + assert [(event["stage"], event["state"]) for event in events] == [ + ("contacting_server", "pending"), + ("contacting_server", "success"), + ("authentication", "pending"), + ("authentication", "success"), + ] + + +def test_stream_close_cancels_inflight_probe(monkeypatch: pytest.MonkeyPatch) -> None: + async def run() -> None: + cancelled = asyncio.Event() + + async def probe(_url: str, progress: setup_probe.Progress) -> None: + try: + await progress( + { + "type": "progress", + "stage": "contacting_server", + "state": "pending", + } + ) + await asyncio.Event().wait() + finally: + cancelled.set() + + monkeypatch.setattr(setup_probe, "probe_mcp_server", probe) + stream = setup_probe.stream_mcp_probe("https://mcp.example/mcp") + assert json.loads(await anext(stream))["state"] == "pending" + await stream.aclose() + assert cancelled.is_set() + + asyncio.run(run()) + + +@pytest.mark.parametrize("failure", [None, "probe", "serialization"]) +def test_stream_delivers_one_terminal_result_after_producer_failure( + monkeypatch: pytest.MonkeyPatch, failure: str | None +) -> None: + async def probe( + _url: str, _progress: setup_probe.Progress + ) -> setup_probe.MCPSetupProbeResult: + if failure == "probe": + raise RuntimeError("Probe failed") + return setup_probe.MCPPublicProbeResult( + transport=MCPTransport.STREAMABLE_HTTP, + tool_count=1, + server_name="Support tools", + ) + + def fail_serialization(*_: object, **__: object) -> None: + raise ValueError("Serialization failed") + + monkeypatch.setattr(setup_probe, "probe_mcp_server", probe) + if failure == "serialization": + monkeypatch.setattr( + setup_probe.MCPPublicProbeResult, "model_dump", fail_serialization + ) + + async def collect() -> list[dict]: + return [ + json.loads(event) + async for event in setup_probe.stream_mcp_probe("https://mcp.example/mcp") + ] + + async def run() -> list[dict]: + return await asyncio.wait_for(collect(), timeout=1) + + events = asyncio.run(run()) + assert len(events) == 1 + assert events[0]["type"] == "result" + if failure is None: + assert events[0]["result"] == { + "kind": "public", + "transport": "STREAMABLE_HTTP", + "tool_count": 1, + "server_name": "Support tools", + } + else: + assert events[0]["result"] == { + "kind": "unreachable", + "reason": "connection_failed", + } + + +@pytest.mark.parametrize("registration", [True, False]) +def test_oauth_uses_post_challenge_and_only_reads_metadata( + monkeypatch: pytest.MonkeyPatch, registration: bool +) -> None: + requests: list[httpx.Request] = [] + + def respond(request: httpx.Request) -> httpx.Response: + requests.append(request) + assert "authorization" not in request.headers + if request.method == "POST": + assert request.url.path == "/mcp" + assert json.loads(request.content)["method"] == "initialize" + return httpx.Response( + 401, + headers={ + "WWW-Authenticate": 'Bearer resource_metadata="https://mcp.example/metadata"' + }, + ) + assert request.method == "GET" + if request.url.path == "/metadata": + return httpx.Response( + 200, + json={ + "resource": "https://mcp.example/mcp", + "authorization_servers": ["https://identity.example"], + }, + ) + return httpx.Response( + 200, + json={ + "issuer": "https://identity.example", + "authorization_endpoint": "https://identity.example/authorize", + "token_endpoint": "https://identity.example/token", + "response_types_supported": ["code"], + "registration_endpoint": "https://identity.example/register" + if registration + else None, + }, + ) + + def factory(**_: object) -> httpx.AsyncClient: + return httpx.AsyncClient(transport=httpx.MockTransport(respond)) + + monkeypatch.setattr(setup_probe, "mcp_ssrf_httpx_client_factory", factory) + monkeypatch.setattr(oauth_discovery, "mcp_ssrf_httpx_client_factory", factory) + result = asyncio.run(setup_probe.probe_mcp_server("https://mcp.example/mcp")) + assert result.kind == "oauth" + assert result.discovery.supports_automatic_registration is registration + assert result.transport is None + assert [(request.method, request.url.path) for request in requests] == [ + ("POST", "/mcp"), + ("GET", "/metadata"), + ("GET", "/.well-known/oauth-authorization-server"), + ] + + +def test_missing_oauth_metadata_is_unknown_auth_not_public( + monkeypatch: pytest.MonkeyPatch, +) -> None: + def respond(request: httpx.Request) -> httpx.Response: + return httpx.Response(401 if request.url.path == "/mcp" else 404) + + def factory(**_: object) -> httpx.AsyncClient: + return httpx.AsyncClient(transport=httpx.MockTransport(respond)) + + monkeypatch.setattr(setup_probe, "mcp_ssrf_httpx_client_factory", factory) + monkeypatch.setattr(oauth_discovery, "mcp_ssrf_httpx_client_factory", factory) + result = asyncio.run(setup_probe.probe_mcp_server("https://mcp.example/mcp")) + assert result.model_dump() == {"kind": "auth_unknown", "transport": None} + + +@pytest.mark.parametrize("failure", ["html", "network", "blocked"]) +def test_failure_is_safe_and_ssrf_block_never_falls_back( + monkeypatch: pytest.MonkeyPatch, failure: str +) -> None: + methods: list[str] = [] + + def respond(request: httpx.Request) -> httpx.Response: + methods.append(request.method) + if failure == "blocked": + raise SSRFException("private URL and secret response") + if failure == "network": + raise httpx.ConnectError("private URL and secret response") + return httpx.Response(200, text="private URL and secret response") + + monkeypatch.setattr( + setup_probe, + "mcp_ssrf_httpx_client_factory", + lambda **_: httpx.AsyncClient(transport=httpx.MockTransport(respond)), + ) + result = asyncio.run(setup_probe.probe_mcp_server("https://mcp.example/mcp")) + assert result.model_dump() == { + "kind": "unreachable", + "reason": "blocked" if failure == "blocked" else "connection_failed", + } + assert "private" not in result.model_dump_json() + assert methods == (["POST"] if failure == "blocked" else ["POST", "GET"]) + + +def test_compatible_sse_fallback_initializes_and_lists_tools( + monkeypatch: pytest.MonkeyPatch, +) -> None: + messages: asyncio.Queue[bytes] = asyncio.Queue() + requests: list[tuple[str, str]] = [] + + class EventStream(httpx.AsyncByteStream): + async def __aiter__(self) -> AsyncIterator[bytes]: + yield b"event: endpoint\ndata: /messages\n\n" + while True: + yield await messages.get() + + async def respond(request: httpx.Request) -> httpx.Response: + requests.append((request.method, request.url.path)) + if request.url.path == "/sse": + if request.method == "POST": + return httpx.Response(405) + return httpx.Response( + 200, headers={"content-type": "text/event-stream"}, stream=EventStream() + ) + assert request.url.path == "/messages" + payload = json.loads(request.content) + if "id" in payload: + result = { + "jsonrpc": "2.0", + "id": payload["id"], + "result": _rpc_result(payload), + } + await messages.put( + f"event: message\ndata: {json.dumps(result)}\n\n".encode() + ) + return httpx.Response(202) + + monkeypatch.setattr( + setup_probe, + "mcp_ssrf_httpx_client_factory", + lambda **_: httpx.AsyncClient(transport=httpx.MockTransport(respond)), + ) + result = asyncio.run(setup_probe.probe_mcp_server("https://mcp.example/sse")) + assert result.kind == "public" + assert result.transport == MCPTransport.SSE + assert result.tool_count == 1 + assert requests[:2] == [("POST", "/sse"), ("GET", "/sse")] + + +def test_protected_sse_message_endpoint_uses_its_oauth_challenge( + monkeypatch: pytest.MonkeyPatch, +) -> None: + requests: list[tuple[str, str]] = [] + + class EventStream(httpx.AsyncByteStream): + async def __aiter__(self) -> AsyncIterator[bytes]: + yield b"event: endpoint\ndata: /messages\n\n" + await asyncio.Event().wait() + + async def respond(request: httpx.Request) -> httpx.Response: + requests.append((request.method, request.url.path)) + if request.url.path == "/sse": + if request.method == "POST": + return httpx.Response(405) + return httpx.Response( + 200, headers={"content-type": "text/event-stream"}, stream=EventStream() + ) + if request.url.path == "/messages": + return httpx.Response( + 401, + headers={ + "www-authenticate": 'Bearer resource_metadata="https://mcp.example/metadata"' + }, + ) + if request.url.path == "/metadata": + return httpx.Response( + 200, + json={ + "resource": "https://mcp.example/sse", + "authorization_servers": ["https://identity.example"], + }, + ) + return httpx.Response( + 200, + json={ + "issuer": "https://identity.example", + "authorization_endpoint": "https://identity.example/authorize", + "token_endpoint": "https://identity.example/token", + "response_types_supported": ["code"], + "client_id_metadata_document_supported": True, + }, + ) + + def factory(**_: object) -> httpx.AsyncClient: + return httpx.AsyncClient(transport=httpx.MockTransport(respond)) + + monkeypatch.setattr(setup_probe, "mcp_ssrf_httpx_client_factory", factory) + monkeypatch.setattr(oauth_discovery, "mcp_ssrf_httpx_client_factory", factory) + result = asyncio.run(setup_probe.probe_mcp_server("https://mcp.example/sse")) + assert result.kind == "oauth" + assert result.transport == MCPTransport.SSE + assert result.discovery.supports_automatic_registration is True + assert requests == [ + ("POST", "/sse"), + ("GET", "/sse"), + ("POST", "/messages"), + ("GET", "/metadata"), + ("GET", "/.well-known/oauth-authorization-server"), + ] + + +@pytest.mark.parametrize("authorized", [True, False]) +def test_probe_route_checks_permission_before_network( + monkeypatch: pytest.MonkeyPatch, authorized: bool +) -> None: + from unittest.mock import AsyncMock + + from fastapi import FastAPI + from fastapi.routing import APIRoute + from fastapi.testclient import TestClient + + from onyx.error_handling.error_codes import OnyxErrorCode + from onyx.error_handling.exceptions import ( + OnyxError, + register_onyx_exception_handlers, + ) + from onyx.server.features.mcp import api + + app = FastAPI() + app.include_router(api.admin_router) + register_onyx_exception_handlers(app) + route = next( + route + for route in api.admin_router.routes + if isinstance(route, APIRoute) and route.path == "/admin/mcp/probe" + ) + dependency = route.dependant.dependencies[0].call + assert dependency is not None + + def permission() -> None: + if not authorized: + raise OnyxError(OnyxErrorCode.INSUFFICIENT_PERMISSIONS) + + app.dependency_overrides[dependency] = permission + probe = AsyncMock(return_value=setup_probe.MCPUnknownAuthProbeResult()) + monkeypatch.setattr(setup_probe, "probe_mcp_server", probe) + with TestClient(app) as test_client: + response = test_client.post( + "/admin/mcp/probe", json={"server_url": "https://mcp.example/mcp"} + ) + assert response.status_code == (200 if authorized else 403) + if authorized: + assert probe.await_count == 1 + awaited = probe.await_args + assert awaited is not None + assert awaited.args[0] == "https://mcp.example/mcp" + assert response.headers["X-Accel-Buffering"] == "no" + assert response.headers["Cache-Control"] == "no-cache, no-transform" + events = [json.loads(line) for line in response.text.splitlines()] + assert events == [ + { + "type": "result", + "result": {"kind": "auth_unknown", "transport": None}, + } + ] + else: + probe.assert_not_called() diff --git a/backend/tests/unit/onyx/server/features/mcp/test_user_credentials_access.py b/backend/tests/unit/onyx/server/features/mcp/test_user_credentials_access.py new file mode 100644 index 0000000000..493014293b --- /dev/null +++ b/backend/tests/unit/onyx/server/features/mcp/test_user_credentials_access.py @@ -0,0 +1,124 @@ +from unittest.mock import MagicMock +from uuid import uuid4 + +import pytest +from sqlalchemy.orm import Session + +from onyx.db.enums import ( + MCPAuthenticationPerformer, + MCPAuthenticationType, + MCPTransport, +) +from onyx.db.models import MCPServer, User +from onyx.error_handling.error_codes import OnyxErrorCode +from onyx.error_handling.exceptions import OnyxError +from onyx.server.features.mcp import api +from onyx.server.features.mcp.credentials import ResolvedMCPCredentials +from onyx.server.features.mcp.models import MCPUserCredentialsRequest + + +@pytest.mark.parametrize( + "auth_type", [MCPAuthenticationType.API_TOKEN, MCPAuthenticationType.OAUTH] +) +def test_inaccessible_server_denies_credentials_before_reads_writes_or_validation( + monkeypatch: pytest.MonkeyPatch, auth_type: MCPAuthenticationType +) -> None: + db = MagicMock(spec=Session) + user = User(id=uuid4(), email="member@example.com") + server = MCPServer(id=42, auth_type=auth_type) + access = MagicMock(return_value=False) + monkeypatch.setattr(api, "get_mcp_server_by_id", lambda *_: server) + monkeypatch.setattr(api, "user_can_configure_mcp_credentials", access) + credential_operations = [ + "get_mcp_auth_template", + "get_user_connection_config", + "extract_connection_data", + "test_mcp_server_credentials", + "upsert_user_connection_config", + "resolve_mcp_credentials", + "_hot_reload_craft_sessions", + ] + operations = {name: MagicMock() for name in credential_operations} + for name, operation in operations.items(): + monkeypatch.setattr(api, name, operation) + + with pytest.raises(OnyxError) as error: + api.save_user_credentials( + MCPUserCredentialsRequest( + server_id=42, + credentials={"api_key": "personal-key"}, + transport="streamable-http", + ), + db, + user, + ) + + assert error.value.error_code is OnyxErrorCode.NOT_FOUND + access.assert_called_once_with(user, server, db) + for operation in operations.values(): + operation.assert_not_called() + db.commit.assert_not_called() + + +def test_accessible_member_saves_only_their_validated_credentials( + monkeypatch: pytest.MonkeyPatch, +) -> None: + db = MagicMock(spec=Session) + user = User(id=uuid4(), email="member@example.com") + server = MCPServer( + id=42, + name="Personal integration", + server_url="https://example.com/mcp", + auth_type=MCPAuthenticationType.API_TOKEN, + auth_performer=MCPAuthenticationPerformer.PER_USER, + ) + access = MagicMock(return_value=True) + validate = MagicMock(return_value=(True, "Connected")) + save = MagicMock() + reload = MagicMock() + resolved = ResolvedMCPCredentials( + connection_config_id=1, + connection_data={"headers": {"Authorization": "Bearer personal-key"}}, + user_oauth_token=None, + auth_type=MCPAuthenticationType.API_TOKEN, + ) + monkeypatch.setattr(api, "get_mcp_server_by_id", lambda *_: server) + monkeypatch.setattr(api, "user_can_configure_mcp_credentials", access) + monkeypatch.setattr(api, "get_mcp_auth_template", lambda *_: None) + monkeypatch.setattr(api, "test_mcp_server_credentials", validate) + monkeypatch.setattr(api, "upsert_user_connection_config", save) + monkeypatch.setattr(api, "resolve_mcp_credentials", lambda *_: resolved) + monkeypatch.setattr(api, "_hot_reload_craft_sessions", reload) + + response = api.save_user_credentials( + MCPUserCredentialsRequest( + server_id=42, + credentials={"api_key": "personal-key"}, + transport="streamable-http", + ), + db, + user, + ) + + access.assert_called_once_with(user, server, db) + validate.assert_called_once_with( + "https://example.com/mcp", + {"Authorization": "Bearer personal-key"}, + transport=MCPTransport.STREAMABLE_HTTP, + auth=None, + ) + save.assert_called_once_with( + server_id=42, + user_email="member@example.com", + config_data={ + "headers": {"Authorization": "Bearer personal-key"}, + "header_substitutions": {"api_key": "personal-key"}, + }, + db_session=db, + ) + db.commit.assert_called_once() + reload.assert_called_once_with({user.id}, db) + assert response.success is True + assert response.authenticated is True + assert response.validation_tested is True + assert response.server_id == 42 diff --git a/backend/tests/unit/onyx/server/features/mcp/test_user_preferences.py b/backend/tests/unit/onyx/server/features/mcp/test_user_preferences.py new file mode 100644 index 0000000000..4ae05e021f --- /dev/null +++ b/backend/tests/unit/onyx/server/features/mcp/test_user_preferences.py @@ -0,0 +1,190 @@ +from collections.abc import Generator +from typing import cast +from unittest.mock import MagicMock +from uuid import uuid4 + +import pytest +from pydantic import ValidationError +from sqlalchemy import Column, Integer, MetaData, Table, Uuid, create_engine +from sqlalchemy.orm import Session + +from onyx.db.mcp import ( + get_mcp_server_user_preferences, + upsert_mcp_server_user_preferences__no_commit, +) +from onyx.db.models import MCPServer, MCPServerUserPreference, User +from onyx.error_handling.error_codes import OnyxErrorCode +from onyx.error_handling.exceptions import OnyxError +from onyx.server.features.mcp import api +from onyx.server.features.mcp.models import ( + MCPConnectionStatus, + MCPUserPreferencesUpdateRequest, +) + + +@pytest.fixture +def preference_db() -> Generator[Session, None, None]: + metadata = MetaData() + Table("user", metadata, Column("id", Uuid, primary_key=True)) + Table("mcp_server", metadata, Column("id", Integer, primary_key=True)) + cast(Table, MCPServerUserPreference.__table__).to_metadata(metadata) + engine = create_engine("sqlite://") + metadata.create_all(engine) + with Session(engine) as db: + yield db + engine.dispose() + + +def test_partial_updates_preserve_other_surface_and_user( + preference_db: Session, +) -> None: + user_id, other_user_id = uuid4(), uuid4() + assert get_mcp_server_user_preferences([42], user_id, preference_db) == {} + first = upsert_mcp_server_user_preferences__no_commit( + 42, user_id, preference_db, enabled_in_chat=False + ) + assert (first.enabled_in_chat, first.enabled_in_craft) == (False, True) + updated = upsert_mcp_server_user_preferences__no_commit( + 42, user_id, preference_db, enabled_in_craft=False + ) + assert (updated.enabled_in_chat, updated.enabled_in_craft) == (False, False) + updated = upsert_mcp_server_user_preferences__no_commit( + 42, user_id, preference_db, enabled_in_chat=True + ) + assert (updated.enabled_in_chat, updated.enabled_in_craft) == (True, False) + assert get_mcp_server_user_preferences([42], other_user_id, preference_db) == {} + other = upsert_mcp_server_user_preferences__no_commit( + 42, other_user_id, preference_db, enabled_in_craft=True + ) + assert (other.enabled_in_chat, other.enabled_in_craft) == (True, True) + assert get_mcp_server_user_preferences([43], user_id, preference_db) == {} + + +@pytest.mark.parametrize( + "payload", [{}, {"enabled_in_chat": None}, {"enabled_in_craft": None}] +) +def test_update_requires_boolean_preference(payload: dict[str, None]) -> None: + with pytest.raises(ValidationError): + MCPUserPreferencesUpdateRequest.model_validate(payload) + + +def test_deleted_server_returns_not_found_without_saving_preferences( + monkeypatch: pytest.MonkeyPatch, +) -> None: + db = MagicMock(spec=Session) + missing = ValueError("Server no longer exists") + monkeypatch.setattr(api, "user_can_access_mcp_server", lambda *_: True) + monkeypatch.setattr(api, "get_mcp_server_by_id", MagicMock(side_effect=missing)) + write = MagicMock() + monkeypatch.setattr(api, "upsert_mcp_server_user_preferences__no_commit", write) + + with pytest.raises(OnyxError) as error: + api.update_mcp_user_preferences( + 42, + MCPUserPreferencesUpdateRequest(enabled_in_chat=False), + db, + User(id=uuid4()), + ) + + assert error.value.error_code is OnyxErrorCode.NOT_FOUND + assert error.value.__cause__ is missing + write.assert_not_called() + db.commit.assert_not_called() + + +@pytest.mark.parametrize( + ("accessible", "chat_available", "craft_available", "payload"), + [ + (False, True, True, {"enabled_in_chat": False}), + (True, False, True, {"enabled_in_chat": True}), + (True, True, False, {"enabled_in_craft": True}), + ], +) +def test_update_cannot_exceed_access_or_admin_availability( + monkeypatch: pytest.MonkeyPatch, + accessible: bool, + chat_available: bool, + craft_available: bool, + payload: dict[str, bool], +) -> None: + db = MagicMock(spec=Session) + monkeypatch.setattr(api, "user_can_access_mcp_server", lambda *_: accessible) + monkeypatch.setattr( + api, + "get_mcp_server_by_id", + lambda *_: MCPServer(id=42, available_in_craft=craft_available), + ) + monkeypatch.setattr( + api, + "get_default_chat_mcp_server_ids", + lambda *_: {42} if chat_available else set(), + ) + write = MagicMock() + monkeypatch.setattr(api, "upsert_mcp_server_user_preferences__no_commit", write) + with pytest.raises(OnyxError): + api.update_mcp_user_preferences( + 42, MCPUserPreferencesUpdateRequest(**payload), db, User(id=uuid4()) + ) + write.assert_not_called() + db.commit.assert_not_called() + + +def test_update_preserves_credentials_and_restamps_only_current_user( + monkeypatch: pytest.MonkeyPatch, + preference_db: Session, +) -> None: + user = User(id=uuid4()) + monkeypatch.setattr(api, "user_can_access_mcp_server", lambda *_: True) + monkeypatch.setattr( + api, + "get_mcp_server_by_id", + lambda *_: MCPServer(id=42, available_in_craft=True), + ) + restamp = MagicMock() + delete_credentials = MagicMock() + monkeypatch.setattr(api, "_hot_reload_craft_sessions", restamp) + monkeypatch.setattr( + api, "delete_user_connection_configs_for_server", delete_credentials + ) + response = api.update_mcp_user_preferences( + 42, MCPUserPreferencesUpdateRequest(enabled_in_craft=False), preference_db, user + ) + assert response.model_dump() == {"enabled_in_chat": True, "enabled_in_craft": False} + restamp.assert_called_once_with({user.id}, preference_db) + delete_credentials.assert_not_called() + + +def test_member_catalog_retains_personal_disabled_servers( + monkeypatch: pytest.MonkeyPatch, + preference_db: Session, +) -> None: + from onyx.db.enums import MCPServerStatus + from onyx.server.features.mcp.models import MCPServer as MCPServerResponse + + user = User(id=uuid4()) + upsert_mcp_server_user_preferences__no_commit( + 42, user.id, preference_db, enabled_in_chat=False, enabled_in_craft=False + ) + monkeypatch.setattr( + api, + "get_mcp_servers_accessible_to_user", + lambda *_: [MCPServer(id=42), MCPServer(id=43)], + ) + monkeypatch.setattr(api, "get_default_chat_mcp_server_ids", lambda *_: {42, 43}) + monkeypatch.setattr( + api, + "_db_mcp_server_to_api_mcp_server", + lambda server, *_, **__: MCPServerResponse( + id=server.id, + name="Server", + owner="admin", + server_url="https://example.com", + status=MCPServerStatus.CREATED, + connection_status=MCPConnectionStatus.NOT_CONNECTED, + ), + ) + response = api.get_mcp_servers_for_user(preference_db, user) + assert [ + (s.id, s.default_chat_enabled, s.user_enabled_in_chat, s.user_enabled_in_craft) + for s in response.mcp_servers + ] == [(42, True, False, False), (43, True, True, True)] diff --git a/backend/tests/unit/onyx/server/features/tool/test_default_chat_availability.py b/backend/tests/unit/onyx/server/features/tool/test_default_chat_availability.py new file mode 100644 index 0000000000..36f68dc08f --- /dev/null +++ b/backend/tests/unit/onyx/server/features/tool/test_default_chat_availability.py @@ -0,0 +1,125 @@ +from unittest.mock import MagicMock + +import pytest +from sqlalchemy import Boolean, Column, Integer, MetaData, Table, create_engine +from sqlalchemy.orm import Session + +from onyx.db.models import MCPServer, Tool, User +from onyx.server.features.tool import api +from onyx.server.features.tool.models import ToolSnapshot + + +@pytest.mark.parametrize("nondefault_builtin_deleted", [True, False]) +def test_default_chat_tool_summary_uses_visible_tool_attachments( + nondefault_builtin_deleted: bool, +) -> None: + from onyx.db.tools import get_default_chat_tool_ids + + metadata = MetaData() + personas = Table( + "persona", + metadata, + Column("id", Integer, primary_key=True), + Column("builtin_persona", Boolean), + Column("deleted", Boolean), + ) + tools = Table( + "tool", + metadata, + Column("id", Integer, primary_key=True), + Column("enabled", Boolean), + ) + attachments = Table( + "persona__tool", + metadata, + Column("persona_id", Integer), + Column("tool_id", Integer), + ) + engine = create_engine("sqlite://") + metadata.create_all(engine) + with Session(engine) as db: + db.execute( + personas.insert(), + [ + {"id": 0, "builtin_persona": True, "deleted": False}, + {"id": 1, "builtin_persona": False, "deleted": False}, + { + "id": 2, + "builtin_persona": True, + "deleted": nondefault_builtin_deleted, + }, + ], + ) + db.execute( + tools.insert(), + [ + {"id": 1, "enabled": True}, + {"id": 2, "enabled": False}, + {"id": 3, "enabled": True}, + {"id": 4, "enabled": True}, + {"id": 5, "enabled": True}, + ], + ) + db.execute( + attachments.insert(), + [ + {"persona_id": persona_id, "tool_id": tool_id} + for persona_id, tool_id in [(0, 1), (0, 2), (1, 3), (2, 4), (0, 5)] + ], + ) + + assert get_default_chat_tool_ids([1, 2, 3, 4], db) == {1, 2} + assert get_default_chat_tool_ids([3, 4], db) == set() + assert get_default_chat_tool_ids([], db) == set() + + +def test_member_tool_list_stamps_default_chat_after_existing_visibility_filters( + monkeypatch: pytest.MonkeyPatch, +) -> None: + db = MagicMock(spec=Session) + user = User(email="member@example.com") + selected = Tool( + id=1, + name="selected", + enabled=True, + passthrough_auth=False, + openapi_schema={"paths": {}}, + ) + unselected = Tool( + id=2, + name="unselected", + enabled=True, + passthrough_auth=False, + openapi_schema={"paths": {}}, + ) + hidden = Tool(id=3, name="hidden", enabled=True, passthrough_auth=False) + inaccessible_mcp = Tool( + id=4, name="private", mcp_server_id=44, enabled=True, passthrough_auth=False + ) + get_tools = MagicMock(return_value=[selected, unselected, hidden, inaccessible_mcp]) + selected_tools = MagicMock(return_value={1}) + monkeypatch.setattr(api, "get_tools", get_tools) + monkeypatch.setattr( + api, + "get_mcp_servers_accessible_to_user", + lambda *_: [MCPServer(id=42)], + ) + monkeypatch.setattr(api, "should_expose_tool_to_fe", lambda tool: tool.id != 3) + monkeypatch.setattr(api, "get_default_chat_tool_ids", selected_tools) + + response = api.list_tools(db, user) + + get_tools.assert_called_once_with(db, only_enabled=True, only_connected_mcp=True) + selected_tools.assert_called_once_with([1, 2], db) + assert [(tool.id, tool.default_chat_enabled) for tool in response] == [ + (1, True), + (2, False), + ] + assert all(tool.permissions == {} for tool in response) + + +def test_default_chat_summary_is_response_only_and_optional() -> None: + tool = Tool(id=1, name="custom", enabled=True, passthrough_auth=False) + + assert "default_chat_enabled" not in Tool.__table__.columns + assert ToolSnapshot.from_model(tool).default_chat_enabled is None diff --git a/backend/tests/unit/onyx/server/features/tool/test_openapi_user_preferences.py b/backend/tests/unit/onyx/server/features/tool/test_openapi_user_preferences.py new file mode 100644 index 0000000000..3fbea27d95 --- /dev/null +++ b/backend/tests/unit/onyx/server/features/tool/test_openapi_user_preferences.py @@ -0,0 +1,207 @@ +from collections.abc import Generator +from typing import cast +from unittest.mock import MagicMock +from uuid import uuid4 + +import pytest +from pydantic import ValidationError +from sqlalchemy import Column, Integer, MetaData, Table, Uuid, create_engine +from sqlalchemy.orm import Session + +from onyx.db.models import Tool, ToolUserPreference, User +from onyx.db.tools import ( + get_tool_user_preferences, + upsert_tool_user_preference__no_commit, +) +from onyx.error_handling.exceptions import OnyxError +from onyx.server.features.tool import api +from onyx.server.features.tool.models import ToolSnapshot, ToolUserPreferences + + +@pytest.fixture +def preference_db() -> Generator[Session, None, None]: + metadata = MetaData() + Table("user", metadata, Column("id", Uuid, primary_key=True)) + Table("tool", metadata, Column("id", Integer, primary_key=True)) + cast(Table, ToolUserPreference.__table__).to_metadata(metadata) + engine = create_engine("sqlite://") + metadata.create_all(engine) + with Session(engine) as db: + yield db + engine.dispose() + + +def test_updates_are_scoped_to_user_and_tool(preference_db: Session) -> None: + user_id, other_user_id = uuid4(), uuid4() + assert get_tool_user_preferences([42], user_id, preference_db) == {} + for enabled in (False, True, False): + preference = upsert_tool_user_preference__no_commit( + 42, user_id, preference_db, enabled_in_chat=enabled + ) + assert preference.enabled_in_chat is enabled + assert get_tool_user_preferences([42], other_user_id, preference_db) == {} + assert get_tool_user_preferences([43], user_id, preference_db) == {} + + +@pytest.mark.parametrize("payload", [{}, {"enabled_in_chat": None}]) +def test_update_requires_boolean(payload: dict[str, None]) -> None: + with pytest.raises(ValidationError): + ToolUserPreferences.model_validate(payload) + + +@pytest.mark.parametrize( + "tool", + [ + Tool(id=42, in_code_tool_id="SearchTool"), + Tool(id=42, mcp_server_id=1), + Tool(id=42), + Tool(id=42, openapi_schema={"paths": {}}, enabled=False), + ], +) +def test_update_rejects_wrong_kind_or_disabled_plugin( + monkeypatch: pytest.MonkeyPatch, tool: Tool +) -> None: + db = MagicMock(spec=Session) + monkeypatch.setattr(api, "get_tool_by_id", lambda *_: tool) + write = MagicMock() + monkeypatch.setattr(api, "upsert_tool_user_preference__no_commit", write) + with pytest.raises(OnyxError): + api.update_tool_user_preferences( + 42, ToolUserPreferences(enabled_in_chat=True), db, User(id=uuid4()) + ) + write.assert_not_called() + db.commit.assert_not_called() + + +def test_update_rejects_missing_hidden_or_admin_unselected( + monkeypatch: pytest.MonkeyPatch, +) -> None: + db = MagicMock(spec=Session) + user = User(id=uuid4()) + tool = Tool(id=42, openapi_schema={"paths": {}}, enabled=True) + monkeypatch.setattr(api, "get_tool_by_id", MagicMock(side_effect=ValueError())) + with pytest.raises(OnyxError): + api.update_tool_user_preferences( + 42, ToolUserPreferences(enabled_in_chat=False), db, user + ) + monkeypatch.setattr(api, "get_tool_by_id", lambda *_: tool) + monkeypatch.setattr(api, "should_expose_tool_to_fe", lambda *_: False) + with pytest.raises(OnyxError): + api.update_tool_user_preferences( + 42, ToolUserPreferences(enabled_in_chat=False), db, user + ) + monkeypatch.setattr(api, "should_expose_tool_to_fe", lambda *_: True) + monkeypatch.setattr(api, "get_default_chat_tool_ids", lambda *_: set()) + with pytest.raises(OnyxError): + api.update_tool_user_preferences( + 42, ToolUserPreferences(enabled_in_chat=True), db, user + ) + db.commit.assert_not_called() + + +def test_update_and_catalog_preserve_admin_state_and_connection( + monkeypatch: pytest.MonkeyPatch, preference_db: Session +) -> None: + user = User(id=uuid4()) + tool = Tool( + id=42, + name="Plugin", + openapi_schema={"paths": {}}, + enabled=True, + passthrough_auth=False, + oauth_config_id=12, + ) + monkeypatch.setattr(api, "get_tool_by_id", lambda *_: tool) + response = api.update_tool_user_preferences( + 42, ToolUserPreferences(enabled_in_chat=False), preference_db, user + ) + snapshots = [ToolSnapshot.from_model(tool)] + api._stamp_personal_tool_preferences(snapshots, user, preference_db) + assert response.enabled_in_chat is False + assert snapshots[0].user_enabled_in_chat is False + assert tool.enabled is True + assert tool.oauth_config_id == 12 + other_snapshots = [ToolSnapshot.from_model(tool)] + api._stamp_personal_tool_preferences( + other_snapshots, User(id=uuid4()), preference_db + ) + assert other_snapshots[0].user_enabled_in_chat is True + monkeypatch.setattr(api, "get_tools", lambda *_, **__: [tool]) + monkeypatch.setattr(api, "get_mcp_servers_accessible_to_user", lambda *_: []) + monkeypatch.setattr(api, "get_default_chat_tool_ids", lambda *_: {42}) + catalog = api.list_tools(preference_db, user) + assert [ + (item.id, item.default_chat_enabled, item.user_enabled_in_chat) + for item in catalog + ] == [(42, True, False)] + response = api.update_tool_user_preferences( + 42, ToolUserPreferences(enabled_in_chat=True), preference_db, user + ) + assert response.enabled_in_chat is True + assert tool.enabled is True + assert tool.oauth_config_id == 12 + + +@pytest.mark.parametrize( + ("persona_id", "enabled", "constructed"), + [(0, False, False), (0, True, True), (0, None, True), (7, False, True)], +) +def test_default_chat_optout_before_credentials_and_allowlist( + monkeypatch: pytest.MonkeyPatch, + persona_id: int, + enabled: bool | None, + constructed: bool, +) -> None: + from onyx.chat.emitter import Emitter + from onyx.context.search.models import PersonaSearchInfo + from onyx.llm.interfaces import LLM + from onyx.tools import tool_constructor + from onyx.tools.models import PersonaToolConfiguration, ToolConfiguration + + tool = Tool( + id=42, + name="Plugin", + openapi_schema={"paths": {}}, + enabled=True, + passthrough_auth=False, + ) + configuration = PersonaToolConfiguration( + persona_id=persona_id, + persona_name="Test", + tools=[ + ToolConfiguration.model_validate(tool).model_copy( + update={"oauth_config_id": 12} + ) + ], + search=PersonaSearchInfo( + document_set_names=[], + search_start_date=None, + attached_document_ids=[], + hierarchy_node_ids=[], + ), + ) + monkeypatch.setattr( + tool_constructor, + "get_tool_user_preferences", + lambda *_: ( + {} if enabled is None else {42: ToolUserPreference(enabled_in_chat=enabled)} + ), + ) + monkeypatch.setattr(tool_constructor, "get_current_search_settings", MagicMock()) + monkeypatch.setattr(tool_constructor, "get_default_document_index", MagicMock()) + build = MagicMock(return_value=[]) + credentials = MagicMock(return_value=None) + monkeypatch.setattr(tool_constructor, "get_oauth_config", credentials) + monkeypatch.setattr( + tool_constructor, "build_custom_tools_from_openapi_schema_and_headers", build + ) + tool_constructor._construct_tools_impl( + configuration=configuration, + db_session=MagicMock(spec=Session), + emitter=MagicMock(spec=Emitter), + user=User(id=uuid4(), enable_memory_tool=False), + llm=MagicMock(spec=LLM), + allowed_tool_ids=[42], + ) + assert build.called is constructed + assert credentials.called is constructed diff --git a/backend/tests/unit/onyx/server/features/user_oauth_token/test_token_status.py b/backend/tests/unit/onyx/server/features/user_oauth_token/test_token_status.py new file mode 100644 index 0000000000..f32d170d90 --- /dev/null +++ b/backend/tests/unit/onyx/server/features/user_oauth_token/test_token_status.py @@ -0,0 +1,79 @@ +from typing import Any +from unittest.mock import MagicMock +from uuid import uuid4 + +import pytest +from sqlalchemy.orm import Session + +from onyx.auth import oauth_token_manager +from onyx.db.models import OAuthConfig, OAuthUserToken, User +from onyx.server.features.user_oauth_token import api +from onyx.utils.sensitive import make_mock_sensitive_value + + +@pytest.mark.parametrize( + ("expires_at", "refresh_token", "client_id", "client_secret", "can_refresh"), + [ + (1, "refresh-token", "client-id", "client-secret", True), + (1, None, "client-id", "client-secret", False), + (1, "", "client-id", "client-secret", False), + (1, "refresh-token", None, "client-secret", False), + (1, "refresh-token", "client-id", None, False), + (1, "refresh-token", "", "client-secret", False), + (1, "refresh-token", "client-id", "", False), + (99999999999, "refresh-token", "client-id", "client-secret", True), + ], +) +def test_token_status_reports_saved_refresh_capability_without_provider_calls( + monkeypatch: pytest.MonkeyPatch, + expires_at: int, + refresh_token: str | None, + client_id: str | None, + client_secret: str | None, + can_refresh: bool, +) -> None: + config = OAuthConfig(id=17) + config.client_id = ( + make_mock_sensitive_value(client_id) if client_id is not None else None + ) + config.client_secret = ( + make_mock_sensitive_value(client_secret) if client_secret is not None else None + ) + token_data: dict[str, Any] = { + "access_token": "access-token", + "expires_at": expires_at, + } + if refresh_token is not None: + token_data["refresh_token"] = refresh_token + token = OAuthUserToken(oauth_config_id=17, oauth_config=config) + token.token_data = make_mock_sensitive_value(token_data) + user = User(id=uuid4()) + db = MagicMock(spec=Session) + read_tokens = MagicMock(return_value=[token]) + monkeypatch.setattr(api, "get_all_user_oauth_tokens", read_tokens) + refresh = MagicMock(side_effect=AssertionError("Status must not refresh tokens")) + monkeypatch.setattr(oauth_token_manager.requests, "post", refresh) + + result = api.get_user_oauth_token_status(db, user) + + assert [status.model_dump() for status in result] == [ + { + "oauth_config_id": 17, + "expires_at": expires_at, + "is_expired": expires_at == 1, + "can_refresh": can_refresh, + } + ] + read_tokens.assert_called_once_with(user.id, db) + refresh.assert_not_called() + db.commit.assert_not_called() + + +def test_token_status_does_not_invent_missing_connections( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(api, "get_all_user_oauth_tokens", lambda *_: []) + + assert ( + api.get_user_oauth_token_status(MagicMock(spec=Session), User(id=uuid4())) == [] + ) diff --git a/backend/tests/unit/onyx/server/route_contract_baseline.json b/backend/tests/unit/onyx/server/route_contract_baseline.json index ea2797fa68..17aa839a02 100644 --- a/backend/tests/unit/onyx/server/route_contract_baseline.json +++ b/backend/tests/unit/onyx/server/route_contract_baseline.json @@ -9,6 +9,7 @@ "GET /manage/admin/connector-checks/runs/{run_id}": "Existing. No registered GATE-2 helper found. Review.", "GET /manage/admin/indexing/targeted-reindex/{job_id}": "Existing. No registered GATE-2 helper found. Review.", "GET /manage/connector-group-restrictions": "Existing. No registered GATE-2 helper found. Review.", + "POST /admin/mcp/probe": "Create-flow probe. Takes only a URL and reads or writes no stored resource, so there is nothing to scope. The SSRF-guarded client blocks internal hosts.", "POST /admin/oauth-config/create": "Existing. Create route: no stored resource to scope yet. Review.", "POST /admin/tool/custom": "Existing. Create route: no stored resource to scope yet. Review.", "POST /admin/tool/custom/validate": "Existing. No registered GATE-2 helper found. Review.", diff --git a/backend/tests/unit/onyx/server/route_contract_manifest.json b/backend/tests/unit/onyx/server/route_contract_manifest.json index 7286432397..91f8927e2a 100644 --- a/backend/tests/unit/onyx/server/route_contract_manifest.json +++ b/backend/tests/unit/onyx/server/route_contract_manifest.json @@ -413,6 +413,7 @@ "PATCH /manage/admin/users/admin-access": "admin", "PATCH /manage/admin/users/craft-enabled": "admin", "PATCH /manage/credential/{credential_id}": "basic", + "PATCH /mcp/server/{server_id}/preferences": "basic", "PATCH /paste-as-tile": "basic", "PATCH /persona/{persona_id}": "basic", "PATCH /persona/{persona_id}/public": "basic", @@ -425,6 +426,7 @@ "PATCH /skills/custom/{skill_id}/share": "basic", "PATCH /temperature-default": "basic", "PATCH /temperature-override-enabled": "basic", + "PATCH /tool/{tool_id}/preferences": "basic", "PATCH /user/assistant-list/update/{assistant_id}": "basic", "PATCH /user/assistant/{assistant_id}/preferences": "basic", "PATCH /user/chat-background": "basic", @@ -473,6 +475,8 @@ "POST /admin/llm/vercel-ai-gateway/available-models": "manage:llms", "POST /admin/log-export": "admin", "POST /admin/mcp/oauth/connect": "manage:actions +scope", + "POST /admin/mcp/oauth/discover": "manage:actions +scope", + "POST /admin/mcp/probe": "manage:actions +scope", "POST /admin/mcp/server": "manage:actions +scope", "POST /admin/mcp/servers/create": "manage:actions +scope", "POST /admin/mcp/servers/update": "manage:actions +scope",