mirror of
https://github.com/onyx-dot-app/onyx.git
synced 2026-10-22 11:22:49 -05:00
fix: mcp user email replacement (#11390)
This commit is contained in:
1 parent
d09d970ee2
commit
cb880c4e17
7 files changed
+586
-34
No files matched your search
@@ -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)"
|
||||
)
|
||||
|
||||
+247
@@ -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"}
|
||||
@@ -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
|
||||
|
||||
Reference in new issue
Block a user