fix: mcp user email replacement (#11390)

This commit is contained in:
Evan Lohn authored and GitHub committed 2026-05-26 22:54:18 +00:00
1 parent d09d970ee2
commit cb880c4e17
7 files changed
+586 -34

No files matched your search

+101 -24
View File
@@ -40,7 +40,6 @@ from onyx.db.mcp import create_mcp_server__no_commit
from onyx.db.mcp import delete_all_user_connection_configs_for_server_no_commit
from onyx.db.mcp import delete_connection_config
from onyx.db.mcp import delete_mcp_server
from onyx.db.mcp import delete_user_connection_configs_for_server
from onyx.db.mcp import extract_connection_data
from onyx.db.mcp import get_all_mcp_servers
from onyx.db.mcp import get_connection_config_by_id
@@ -138,6 +137,33 @@ def _resolve_oauth_credentials(
return resolved_id, resolved_secret
def _resolve_admin_credentials(
*,
request_credentials: dict[str, str],
request_credentials_changed: dict[str, bool],
existing_user_credentials: dict[str, str] | None,
) -> dict[str, str]:
"""Per-key analogue of ``_resolve_oauth_credentials``: reuse the
stored value when the changed flag is False, otherwise take the
request value and reject masked placeholders defensively. Stored
values are sourced from the editing admin's own per-user
``header_substitutions``."""
resolved: dict[str, str] = {}
for key, request_value in request_credentials.items():
changed = request_credentials_changed.get(key, False)
if (
not changed
and existing_user_credentials
and key in existing_user_credentials
):
resolved[key] = existing_user_credentials[key]
continue
if request_value:
reject_masked_credentials({key: request_value})
resolved[key] = request_value
return resolved
def _build_oauth_admin_config_data(
*,
client_id: str | None,
@@ -853,14 +879,18 @@ def save_user_credentials(
headers={"Authorization": f"Bearer {request.credentials['api_key']}"},
)
else:
# Use template to create the full connection config
# Render via the shared helper so user + auto (`{user_email}`)
# substitutions go through one pipeline.
try:
# TODO: fix and/or type correctly w/base model
auth_template_dict = extract_connection_data(
auth_template, apply_mask=False
)
template = MCPAuthTemplate(headers=auth_template_dict.get("headers", {}))
config_data = MCPConnectionData(
headers=auth_template_dict.get("headers", {}),
headers=_build_headers_from_template(
template, request.credentials, email
),
header_substitutions=request.credentials,
)
for oauth_field_key in MCPOAuthKeys:
@@ -894,11 +924,6 @@ def save_user_credentials(
None,
)
if HEADER_SUBSTITUTIONS in config_data:
for key, value in config_data[HEADER_SUBSTITUTIONS].items():
for k, v in config_data["headers"].items():
config_data["headers"][k] = v.replace(f"{{{key}}}", value)
server_url = mcp_server.server_url
is_valid, test_message = test_mcp_server_credentials(
server_url,
@@ -1530,7 +1555,8 @@ def _upsert_mcp_server(
detail=f"MCP server with ID {request.existing_server_id} not found",
)
_ensure_mcp_server_owner_or_admin(mcp_server, user)
client_info = None
client_info: OAuthClientInformationFull | None = None
existing_admin_config_dict: MCPConnectionData = MCPConnectionData(headers={})
if mcp_server.admin_connection_config:
existing_admin_config_dict = extract_connection_data(
mcp_server.admin_connection_config, apply_mask=False
@@ -1557,6 +1583,55 @@ def _upsert_mcp_server(
existing_client=client_info,
)
# Same pattern for per-user API_TOKEN: resolve admin credentials
# against the admin's stored per-user row
existing_admin_per_user_creds: dict[str, str] = {}
existing_template_headers: dict[str, str] = {}
if mcp_server.admin_connection_config:
existing_template_headers = (
existing_admin_config_dict.get("headers", {}) or {}
)
if (
request.auth_type == MCPAuthenticationType.API_TOKEN
and request.auth_performer == MCPAuthenticationPerformer.PER_USER
and user.email
):
existing_admin_per_user_config = get_user_connection_config(
mcp_server.id, user.email, db_session
)
if existing_admin_per_user_config:
existing_admin_per_user_dict = extract_connection_data(
existing_admin_per_user_config, apply_mask=False
)
existing_admin_per_user_creds = (
existing_admin_per_user_dict.get(HEADER_SUBSTITUTIONS) or {}
)
if request.admin_credentials is not None:
request.admin_credentials = _resolve_admin_credentials(
request_credentials=request.admin_credentials,
request_credentials_changed=request.admin_credentials_changed,
existing_user_credentials=existing_admin_per_user_creds,
)
api_token_creds_changed = (
request.auth_type == MCPAuthenticationType.API_TOKEN
and request.auth_performer == MCPAuthenticationPerformer.PER_USER
and existing_admin_per_user_creds != (request.admin_credentials or {})
)
api_token_template_changed = (
request.auth_type == MCPAuthenticationType.API_TOKEN
and request.auth_performer == MCPAuthenticationPerformer.PER_USER
and request.auth_template is not None
and request.auth_template.headers != existing_template_headers
)
api_token_scheme_changed = (
request.auth_type == MCPAuthenticationType.API_TOKEN
and (
request.auth_type != mcp_server.auth_type
or request.auth_performer != mcp_server.auth_performer
)
)
changing_connection_config = (
not mcp_server.admin_connection_config
or (
@@ -1567,13 +1642,20 @@ def _upsert_mcp_server(
or request.oauth_client_secret != (client_info.client_secret or "")
)
)
or (request.auth_type == MCPAuthenticationType.API_TOKEN)
or (
request.auth_type == MCPAuthenticationType.API_TOKEN
and (
api_token_creds_changed
or api_token_template_changed
or api_token_scheme_changed
)
)
or (request.transport != mcp_server.transport)
)
# Cleanup: Delete existing connection configs
# If the auth type is OAUTH, delete all user connection configs
# If the auth type is API_TOKEN, delete the admin connection config and the admin user connection configs
# OAuth: wipe every user's tokens — re-handshake required.
# API_TOKEN: drop only the shared template; the admin's per-user
# row is upserted in place below.
if (
changing_connection_config
and mcp_server.admin_connection_config_id
@@ -1588,10 +1670,6 @@ def _upsert_mcp_server(
and request.auth_type == MCPAuthenticationType.API_TOKEN
):
delete_connection_config(mcp_server.admin_connection_config_id, db_session)
if user.email:
delete_user_connection_configs_for_server(
mcp_server.id, user.email, db_session
)
# Update the server with new values
mcp_server = update_mcp_server__no_commit(
@@ -1675,11 +1753,11 @@ def _upsert_mcp_server(
or MCPAuthTemplate.derive_required_fields(template_data.headers)
)
# Create template config: faithful representation of what's in the admin panel
# Template config: placeholder headers + required fields only.
# Admin's credentials live on the admin's own per-user row.
template_config = create_connection_config(
config_data=MCPConnectionData(
headers=template_data.headers,
header_substitutions=request.admin_credentials,
required_fields=persisted_required_fields,
),
mcp_server_id=mcp_server.id,
@@ -1687,19 +1765,18 @@ def _upsert_mcp_server(
db_session=db_session,
)
# seed the user config for this admin user
user_config = create_connection_config(
# Seed (or refresh) the admin's own per-user row.
upsert_user_connection_config(
server_id=mcp_server.id,
user_email=user.email,
config_data=MCPConnectionData(
headers=_build_headers_from_template(
template_data, request.admin_credentials, user.email
),
header_substitutions=request.admin_credentials,
),
mcp_server_id=mcp_server.id,
user_email=user.email,
db_session=db_session,
)
user_config.mcp_server_id = mcp_server.id
admin_connection_config_id = template_config.id
elif request.auth_type == MCPAuthenticationType.OAUTH:
# Create initial admin config. If client credentials were provided,
@@ -155,6 +155,12 @@ class MCPToolCreateRequest(BaseModel):
None,
description="Admin's credential key-value pairs for template substitution and storage",
)
admin_credentials_changed: dict[str, bool] = Field(
default_factory=dict,
description=(
"Per-field flags marking which `admin_credentials` were edited"
), # True = use value from request, False = use stored value
)
existing_server_id: Optional[int] = Field(
None, description="ID of existing server to update (for editing)"
)
@@ -0,0 +1,247 @@
"""Per-user API_TOKEN MCP credentials must satisfy two invariants that
only show up in DB state, not in API responses:
`save_user_credentials` resolves `{user_email}` (and other auto
placeholders) when persisting headers, and `_upsert_mcp_server`
preserves the editing admin's stored per-user credentials when the
form flags them as unchanged. Tests poke the JSONB directly because
HTTP responses mask credential values."""
from unittest.mock import patch
from uuid import uuid4
from sqlalchemy.orm import Session
from onyx.auth.schemas import UserRole
from onyx.db.enums import MCPAuthenticationPerformer
from onyx.db.enums import MCPAuthenticationType
from onyx.db.enums import MCPTransport
from onyx.db.mcp import extract_connection_data
from onyx.db.mcp import get_user_connection_config
from onyx.db.models import User
from onyx.server.features.mcp.api import _upsert_mcp_server
from onyx.server.features.mcp.api import HEADER_SUBSTITUTIONS
from onyx.server.features.mcp.api import save_user_credentials
from onyx.server.features.mcp.models import MCPAuthTemplate
from onyx.server.features.mcp.models import MCPToolCreateRequest
from onyx.server.features.mcp.models import MCPUserCredentialsRequest
from onyx.utils.encryption import mask_string
from tests.external_dependency_unit.conftest import create_test_user
def _make_per_user_api_token_request(
*,
server_name: str,
admin_credentials: dict[str, str],
existing_server_id: int | None = None,
description: str = "regression server",
auth_template: MCPAuthTemplate | None = None,
admin_credentials_changed: dict[str, bool] | None = None,
) -> MCPToolCreateRequest:
return MCPToolCreateRequest(
name=server_name,
description=description,
server_url="http://upstream.example.com/mcp",
auth_type=MCPAuthenticationType.API_TOKEN,
auth_performer=MCPAuthenticationPerformer.PER_USER,
transport=MCPTransport.STREAMABLE_HTTP,
auth_template=auth_template
or MCPAuthTemplate(
headers={"Authorization": "PlainBasic {user_email}:{api_key}"},
required_fields=["api_key"],
),
admin_credentials=admin_credentials,
admin_credentials_changed=admin_credentials_changed or {},
existing_server_id=existing_server_id,
)
def _read_user_config(*, server_id: int, user_email: str, db_session: Session) -> dict:
"""Return the unmasked JSONB blob for a user's connection config row."""
cfg = get_user_connection_config(server_id, user_email, db_session)
assert cfg is not None, f"no per-user config for {user_email}"
return dict(extract_connection_data(cfg, apply_mask=False))
class TestSaveUserCredentialsSubstitutesUserEmail:
"""`save_user_credentials` must apply both user-supplied and
auto (`{user_email}`) substitutions before persisting headers."""
def test_user_email_in_template_resolves_at_save_time(
self, db_session: Session
) -> None:
admin = create_test_user(
db_session, "admin_user_email_sub", role=UserRole.ADMIN
)
basic_user = create_test_user(db_session, "basic_user_email_sub")
server_name = f"user-email-sub-{uuid4().hex[:8]}"
mcp_server = _upsert_mcp_server(
_make_per_user_api_token_request(
server_name=server_name,
admin_credentials={"api_key": "admin-key"},
),
db_session,
admin,
)
with patch(
"onyx.server.features.mcp.api.test_mcp_server_credentials",
return_value=(True, "ok"),
):
save_user_credentials(
MCPUserCredentialsRequest(
server_id=mcp_server.id,
credentials={"api_key": "basic-user-key"},
transport="streamable_http",
),
db_session,
basic_user,
)
cfg = _read_user_config(
server_id=mcp_server.id,
user_email=basic_user.email,
db_session=db_session,
)
expected = f"PlainBasic {basic_user.email}:basic-user-key"
assert cfg["headers"]["Authorization"] == expected
assert "{user_email}" not in cfg["headers"]["Authorization"]
class TestAdminEditPreservesAdminReauth:
"""A panel save that doesn't flag credentials as changed must not
overwrite the admin's stored per-user credentials."""
def _create_server(
self, db_session: Session, admin: User
) -> tuple[int, str, MCPAuthTemplate]:
server_name = f"admin-edit-preserve-{uuid4().hex[:8]}"
template = MCPAuthTemplate(
headers={"Authorization": "PlainBasic {user_email}:{api_key}"},
required_fields=["api_key"],
)
mcp_server = _upsert_mcp_server(
_make_per_user_api_token_request(
server_name=server_name,
admin_credentials={"api_key": "key_A"},
auth_template=template,
),
db_session,
admin,
)
return mcp_server.id, server_name, template
def _admin_reauth(
self, db_session: Session, admin: User, server_id: int, new_key: str
) -> None:
with patch(
"onyx.server.features.mcp.api.test_mcp_server_credentials",
return_value=(True, "ok"),
):
save_user_credentials(
MCPUserCredentialsRequest(
server_id=server_id,
credentials={"api_key": new_key},
transport="streamable_http",
),
db_session,
admin,
)
def test_unchanged_resubmit_preserves_admins_reauthed_key(
self, db_session: Session
) -> None:
admin = create_test_user(db_session, "admin_unchanged", role=UserRole.ADMIN)
server_id, server_name, template = self._create_server(db_session, admin)
self._admin_reauth(db_session, admin, server_id, "key_B")
# Panel save replays the masked original (`key_A`) with the
# changed flag set False; `key_B` must survive.
_upsert_mcp_server(
_make_per_user_api_token_request(
server_name=server_name,
admin_credentials={"api_key": mask_string("key_A")},
admin_credentials_changed={"api_key": False},
existing_server_id=server_id,
description="Updated description only",
auth_template=template,
),
db_session,
admin,
)
cfg = _read_user_config(
server_id=server_id, user_email=admin.email, db_session=db_session
)
assert cfg.get(HEADER_SUBSTITUTIONS) == {"api_key": "key_B"}
assert cfg["headers"]["Authorization"] == f"PlainBasic {admin.email}:key_B"
def test_changed_resubmit_applies_new_key(self, db_session: Session) -> None:
admin = create_test_user(db_session, "admin_changed", role=UserRole.ADMIN)
server_id, server_name, template = self._create_server(db_session, admin)
self._admin_reauth(db_session, admin, server_id, "key_B")
_upsert_mcp_server(
_make_per_user_api_token_request(
server_name=server_name,
admin_credentials={"api_key": "key_C"},
admin_credentials_changed={"api_key": True},
existing_server_id=server_id,
auth_template=template,
),
db_session,
admin,
)
cfg = _read_user_config(
server_id=server_id, user_email=admin.email, db_session=db_session
)
assert cfg.get(HEADER_SUBSTITUTIONS) == {"api_key": "key_C"}
assert cfg["headers"]["Authorization"] == f"PlainBasic {admin.email}:key_C"
def test_other_users_per_user_config_is_never_touched(
self, db_session: Session
) -> None:
# Admin-panel cleanup must stay scoped to the editing admin.
admin = create_test_user(
db_session, "admin_other_unaffected", role=UserRole.ADMIN
)
basic_user = create_test_user(db_session, "basic_other_unaffected")
server_id, server_name, template = self._create_server(db_session, admin)
with patch(
"onyx.server.features.mcp.api.test_mcp_server_credentials",
return_value=(True, "ok"),
):
save_user_credentials(
MCPUserCredentialsRequest(
server_id=server_id,
credentials={"api_key": "user-key"},
transport="streamable_http",
),
db_session,
basic_user,
)
_upsert_mcp_server(
_make_per_user_api_token_request(
server_name=server_name,
admin_credentials={"api_key": "admin-key-rotated"},
admin_credentials_changed={"api_key": True},
existing_server_id=server_id,
auth_template=template,
),
db_session,
admin,
)
cfg = _read_user_config(
server_id=server_id, user_email=basic_user.email, db_session=db_session
)
assert cfg.get(HEADER_SUBSTITUTIONS) == {"api_key": "user-key"}
assert (
cfg["headers"]["Authorization"] == f"PlainBasic {basic_user.email}:user-key"
)
@@ -0,0 +1,140 @@
"""`_resolve_admin_credentials` is the per-key "changed flag + mask
rejector" used when persisting `admin_credentials` for a per-user
API_TOKEN MCP server. Mirrors `_resolve_oauth_credentials` (covered
in `test_oauth_credentials_resolver.py`)."""
import pytest
from onyx.server.features.mcp.api import _resolve_admin_credentials
from onyx.utils.encryption import mask_string
class TestResolveAdminCredentials:
def test_unchanged_resubmit_with_masked_values_keeps_stored(self) -> None:
# Each unchanged field → stored value, never the replayed mask.
stored = {
"api_key": "long-api-key-secret-1234",
"username": "stored-username-1234",
}
request = {
"api_key": mask_string("long-api-key-secret-1234"),
"username": mask_string("stored-username-1234"),
}
changed: dict[str, bool] = {"api_key": False, "username": False}
resolved = _resolve_admin_credentials(
request_credentials=request,
request_credentials_changed=changed,
existing_user_credentials=stored,
)
assert resolved == stored
def test_partial_change_keeps_unflagged_fields(self) -> None:
stored = {"api_key": "stored-key-1234", "username": "stored-user-1234"}
request = {
"api_key": "brand-new-key",
"username": mask_string("stored-user-1234"),
}
changed = {"api_key": True, "username": False}
resolved = _resolve_admin_credentials(
request_credentials=request,
request_credentials_changed=changed,
existing_user_credentials=stored,
)
assert resolved == {
"api_key": "brand-new-key",
"username": "stored-user-1234",
}
def test_missing_changed_flag_defaults_to_false(self) -> None:
# Older clients omit the flag; treat as unchanged.
stored = {"api_key": "stored-key-1234"}
request = {"api_key": mask_string("stored-key-1234")}
resolved = _resolve_admin_credentials(
request_credentials=request,
request_credentials_changed={},
existing_user_credentials=stored,
)
assert resolved == stored
def test_changed_flag_with_long_mask_is_rejected(self) -> None:
stored = {"api_key": "stored-key-1234"}
request = {"api_key": mask_string("some-other-long-string")}
changed = {"api_key": True}
with pytest.raises(ValueError, match="api_key"):
_resolve_admin_credentials(
request_credentials=request,
request_credentials_changed=changed,
existing_user_credentials=stored,
)
def test_changed_flag_with_short_mask_placeholder_is_rejected(self) -> None:
# `mask_string` uses a bullet-string placeholder for short inputs.
stored = {"api_key": "stored-key-1234"}
request = {"api_key": mask_string("short")}
changed = {"api_key": True}
with pytest.raises(ValueError, match="api_key"):
_resolve_admin_credentials(
request_credentials=request,
request_credentials_changed=changed,
existing_user_credentials=stored,
)
def test_no_stored_creds_passes_request_values_through(self) -> None:
# New-server-create path: no stored fallback available.
request = {"api_key": "user-typed-key", "username": "user-typed-user"}
resolved = _resolve_admin_credentials(
request_credentials=request,
request_credentials_changed={},
existing_user_credentials=None,
)
assert resolved == request
def test_no_stored_creds_with_changed_flags_uses_request_values(self) -> None:
request = {"api_key": "user-typed-key"}
changed = {"api_key": True}
resolved = _resolve_admin_credentials(
request_credentials=request,
request_credentials_changed=changed,
existing_user_credentials=None,
)
assert resolved == request
def test_empty_request_returns_empty(self) -> None:
resolved = _resolve_admin_credentials(
request_credentials={},
request_credentials_changed={},
existing_user_credentials={"api_key": "stored"},
)
assert resolved == {}
def test_request_key_not_in_stored_passes_through(self) -> None:
# New template field with no stored fallback; take the request value.
stored = {"api_key": "stored-key-1234"}
request = {
"api_key": mask_string("stored-key-1234"),
"username": "new-required-field-value",
}
changed = {"api_key": False, "username": False}
resolved = _resolve_admin_credentials(
request_credentials=request,
request_credentials_changed=changed,
existing_user_credentials=stored,
)
assert resolved == {
"api_key": "stored-key-1234",
"username": "new-required-field-value",
}
@@ -0,0 +1,60 @@
"""`_build_headers_from_template` is the single source of truth for
rendering per-user API_TOKEN headers: it must apply both user-supplied
substitutions and auto ones (e.g. `{user_email}`) in one pass."""
from onyx.server.features.mcp.api import _build_headers_from_template
from onyx.server.features.mcp.models import MCPAuthTemplate
class TestBuildHeadersFromTemplate:
def test_user_email_placeholder_alone_is_substituted(self) -> None:
template = MCPAuthTemplate(
headers={"X-User": "{user_email}"},
required_fields=[],
)
headers = _build_headers_from_template(
template_data=template,
credentials={},
user_email="alice@example.com",
)
assert headers == {"X-User": "alice@example.com"}
def test_user_email_and_user_supplied_key_substituted_together(self) -> None:
# Both placeholder classes must resolve in a single pass.
template = MCPAuthTemplate(
headers={"Authorization": "PlainBasic {user_email}:{api_key}"},
required_fields=["api_key"],
)
headers = _build_headers_from_template(
template_data=template,
credentials={"api_key": "ATATT-secret"},
user_email="anuj@hudson-trading.com",
)
assert headers == {
"Authorization": "PlainBasic anuj@hudson-trading.com:ATATT-secret"
}
def test_user_supplied_key_alone_is_substituted(self) -> None:
template = MCPAuthTemplate(
headers={"Authorization": "Bearer {api_key}"},
required_fields=["api_key"],
)
headers = _build_headers_from_template(
template_data=template,
credentials={"api_key": "k"},
user_email="ignored@example.com",
)
assert headers == {"Authorization": "Bearer k"}
def test_empty_header_name_is_dropped(self) -> None:
# Empty header names must be filtered (otherwise HTTP layer breaks).
template = MCPAuthTemplate(
headers={"": "Bearer {api_key}", "Authorization": "Bearer {api_key}"},
required_fields=["api_key"],
)
headers = _build_headers_from_template(
template_data=template,
credentials={"api_key": "k"},
user_email="alice@example.com",
)
assert headers == {"Authorization": "Bearer k"}
+2
View File
@@ -190,6 +190,8 @@ export async function upsertMCPServer(serverData: {
oauth_client_secret_changed?: boolean;
auth_template?: any;
admin_credentials?: Record<string, string>;
// Per-key analogue of `oauth_client_*_changed` for `admin_credentials`.
admin_credentials_changed?: Record<string, boolean>;
existing_server_id?: number;
}): Promise<ApiResponse<UpsertMCPServerResponse>> {
try {
@@ -219,10 +219,33 @@ export default function MCPAuthenticationModal({
};
};
// Per-key analogue of `computeOAuthChangedFlags` for the
// `admin_credentials` dict (per-user API_TOKEN only).
const computeAdminCredentialsChangedFlags = (
values: MCPAuthFormValues
): Record<string, boolean> => {
if (
values.auth_type !== MCPAuthenticationType.API_TOKEN ||
values.auth_performer !== MCPAuthenticationPerformer.PER_USER
) {
return {};
}
const current = values.user_credentials || {};
const initial = initialValues.user_credentials || {};
const flags: Record<string, boolean> = {};
for (const key of Object.keys(current)) {
flags[key] = current[key] !== initial[key];
}
return flags;
};
const constructServerData = (values: MCPAuthFormValues) => {
if (!mcpServer) return null;
const authType = values.auth_type;
const oauthChangedFlags = computeOAuthChangedFlags(values);
const isPerUserApiToken =
values.auth_performer === MCPAuthenticationPerformer.PER_USER &&
authType === MCPAuthenticationType.API_TOKEN;
return {
name: mcpServer.name,
@@ -236,16 +259,13 @@ export default function MCPAuthenticationModal({
values.auth_performer === MCPAuthenticationPerformer.ADMIN
? values.api_token
: undefined,
auth_template:
values.auth_performer === MCPAuthenticationPerformer.PER_USER &&
authType === MCPAuthenticationType.API_TOKEN
? values.auth_template
: undefined,
admin_credentials:
values.auth_performer === MCPAuthenticationPerformer.PER_USER &&
authType === MCPAuthenticationType.API_TOKEN
? values.user_credentials || {}
: undefined,
auth_template: isPerUserApiToken ? values.auth_template : undefined,
admin_credentials: isPerUserApiToken
? values.user_credentials || {}
: undefined,
admin_credentials_changed: isPerUserApiToken
? computeAdminCredentialsChangedFlags(values)
: undefined,
oauth_client_id:
authType === MCPAuthenticationType.OAUTH
? values.oauth_client_id