diff --git a/docs/automations/integrations/mcp-integrations.mdx b/docs/automations/integrations/mcp-integrations.mdx index 945c8d8bc5..fe51d25b40 100644 --- a/docs/automations/integrations/mcp-integrations.mdx +++ b/docs/automations/integrations/mcp-integrations.mdx @@ -42,19 +42,15 @@ For a custom remote MCP server, first create a custom OAuth provider in [OAuth]( ### Custom headers for remote MCP -`Custom` authentication stores request headers as JSON. +`Custom` authentication stores request headers as JSON. Header values support workspace secret and variable expressions. ```json { - "Authorization": "Bearer token123", - "X-API-Key": "abc123" + "Authorization": "ApiKey ${{ SECRETS.elastic_security.ELASTIC_API_KEY }}", + "X-Tenant": "${{ VARS.elastic_security.tenant }}" } ``` - -Remote MCP custom header JSON does not resolve `${{ SECRETS.* }}` or `${{ VARS.* }}`. Header values are sent as literal strings. - - ## Stdio MCP Use `stdio` MCP when Tracecat should launch a local command such as `npx`, `uvx`, or a custom binary. @@ -72,11 +68,7 @@ Tracecat resolves those expressions from the workflow default environment unless ## Secrets and variables in MCP configuration -Expression support differs by MCP integration type: - -- `stdio` environment variables support `${{ SECRETS.* }}` and `${{ VARS.* }}`. -- Remote MCP OAuth mode does not need secret expressions for the bearer token because Tracecat injects the OAuth token automatically. -- Remote MCP custom header JSON does not currently resolve `${{ SECRETS.* }}` or `${{ VARS.* }}`. +Both remote and `stdio` MCP integrations support template strings. ## Related pages diff --git a/frontend/src/components/editor/codemirror/code-editor.tsx b/frontend/src/components/editor/codemirror/code-editor.tsx index 7b7b6f8920..c36fa97df7 100644 --- a/frontend/src/components/editor/codemirror/code-editor.tsx +++ b/frontend/src/components/editor/codemirror/code-editor.tsx @@ -17,6 +17,7 @@ interface CodeEditorProps { readOnly?: boolean wrapLongLines?: boolean className?: string + additionalExtensions?: Extension[] } function getLanguageExtension(language: string): Extension | null { @@ -40,6 +41,7 @@ export function CodeEditor({ readOnly = false, wrapLongLines = false, className, + additionalExtensions = [], }: CodeEditorProps) { const { resolvedTheme } = useTheme() const codeMirrorTheme = resolvedTheme === "dark" ? "dark" : "light" @@ -47,6 +49,7 @@ export function CodeEditor({ const extensions = [ ...(languageExtension ? [languageExtension] : []), ...(wrapLongLines ? [EditorView.lineWrapping] : []), + ...additionalExtensions, ] return ( @@ -54,6 +57,9 @@ export function CodeEditor({ value={value} onChange={onChange} extensions={extensions} + basicSetup={ + additionalExtensions.length > 0 ? { autocompletion: false } : undefined + } theme={codeMirrorTheme} readOnly={readOnly} className={cn( diff --git a/frontend/src/components/editor/codemirror/common.tsx b/frontend/src/components/editor/codemirror/common.tsx index e604ed6531..2adde9830b 100644 --- a/frontend/src/components/editor/codemirror/common.tsx +++ b/frontend/src/components/editor/codemirror/common.tsx @@ -1185,6 +1185,11 @@ export const TEMPLATE_SUGGESTIONS = [ info: "For each item in the array", }, ] +type TemplateSuggestion = (typeof TEMPLATE_SUGGESTIONS)[number] + +const WORKSPACE_MAPPING_TEMPLATE_SUGGESTIONS = TEMPLATE_SUGGESTIONS.filter( + (suggestion) => suggestion.label === "SECRETS" || suggestion.label === "VARS" +) // Custom keymap for @ key to trigger completions export function createAtKeyCompletion() { @@ -1229,9 +1234,9 @@ export function createExitEditModeKeyHandler() { } // Completion functions -export function createMentionCompletion(): ( - context: CompletionContext -) => CompletionResult | null { +export function createMentionCompletion( + suggestions: readonly TemplateSuggestion[] = TEMPLATE_SUGGESTIONS +): (context: CompletionContext) => CompletionResult | null { return (context: CompletionContext): CompletionResult | null => { const word = context.matchBefore(/@\w*/) if (!word) return null @@ -1244,7 +1249,7 @@ export function createMentionCompletion(): ( return { from: word.from, - options: TEMPLATE_SUGGESTIONS.map((suggestion) => ({ + options: suggestions.map((suggestion) => ({ label: `@${suggestion.label}`, detail: suggestion.detail, info: suggestion.info, @@ -2069,6 +2074,19 @@ export function createAutocomplete({ }) } +/** + * Create secret and variable completions for workspace configuration mappings. + */ +export function createWorkspaceMappingAutocomplete(workspaceId: string) { + return autocompletion({ + override: [ + createMentionCompletion(WORKSPACE_MAPPING_TEMPLATE_SUGGESTIONS), + createSecretsCompletion(workspaceId), + createVarsCompletion(workspaceId), + ], + }) +} + // Common theme for template pills export const templatePillTheme = EditorView.theme({ ".cm-template-pill": { diff --git a/frontend/src/components/integrations/mcp-integration-dialog.tsx b/frontend/src/components/integrations/mcp-integration-dialog.tsx index 39f0e94824..41e777fcae 100644 --- a/frontend/src/components/integrations/mcp-integration-dialog.tsx +++ b/frontend/src/components/integrations/mcp-integration-dialog.tsx @@ -1,5 +1,7 @@ "use client" +import { completionKeymap } from "@codemirror/autocomplete" +import { keymap } from "@codemirror/view" import { zodResolver } from "@hookform/resolvers/zod" import { Check, @@ -30,6 +32,12 @@ import type { } from "@/client/types.gen" import { useScopeCheck } from "@/components/auth/scope-guard" import { CodeEditor } from "@/components/editor/codemirror/code-editor" +import { + createAtKeyCompletion, + createWorkspaceMappingAutocomplete, + templatePillTheme, +} from "@/components/editor/codemirror/common" +import { createSimpleTemplatePlugin } from "@/components/editor/codemirror/highlight-plugin" import { getMcpProviderIconId, ProviderIcon } from "@/components/icons" import { ALLOWED_COMMANDS, @@ -408,6 +416,16 @@ export function MCPIntegrationDialog({ catalogEntry?: PlatformMCPCatalogRead | null }) { const workspaceId = useWorkspaceId() + const mcpTemplateEditorExtensions = React.useMemo( + () => [ + createWorkspaceMappingAutocomplete(workspaceId), + createAtKeyCompletion(), + keymap.of(completionKeymap), + createSimpleTemplatePlugin(workspaceId), + templatePillTheme, + ], + [workspaceId] + ) const isEditMode = Boolean(mcpIntegrationId) const { connectMcpIntegration, connectMcpIntegrationIsPending } = useConnectMcpIntegration(workspaceId) @@ -1479,6 +1497,9 @@ export function MCPIntegrationDialog({ value={field.value || ""} onChange={field.onChange} language="json" + additionalExtensions={ + mcpTemplateEditorExtensions + } className="font-mono text-xs [&_.cm-content]:text-xs [&_.cm-editor]:min-h-[80px]" /> @@ -1761,12 +1782,17 @@ export function MCPIntegrationDialog({ value={field.value || ""} onChange={field.onChange} language="json" + additionalExtensions={ + mcpTemplateEditorExtensions + } className="font-mono text-xs [&_.cm-content]:text-xs [&_.cm-editor]:min-h-[120px]" /> Authorization is set from OAuth and cannot be - overridden. + overridden. Type @SECRETS or{" "} + @VARS to insert a workspace + expression in another header value. @@ -1809,12 +1835,20 @@ export function MCPIntegrationDialog({ value={field.value || ""} onChange={field.onChange} language="json" + additionalExtensions={mcpTemplateEditorExtensions} className="font-mono text-xs [&_.cm-content]:text-xs [&_.cm-editor]:min-h-[120px]" /> - Enter headers as a JSON object, for example{" "} - {`{"Authorization":"Bearer token123"}`} + Enter headers as JSON. Type @SECRETS or{" "} + @VARS to insert a workspace expression, + such as{" "} + + { + '{"Authorization":"ApiKey ${{ SECRETS.elastic.API_KEY }}"}' + } + + . diff --git a/tests/unit/test_agent_preset_service.py b/tests/unit/test_agent_preset_service.py index dc4f9f0f1f..9456aad0c0 100644 --- a/tests/unit/test_agent_preset_service.py +++ b/tests/unit/test_agent_preset_service.py @@ -19,6 +19,7 @@ SlackChannelTokenConfig, ) from tracecat.agent.channels.service import PENDING_SLACK_BOT_TOKEN, AgentChannelService +from tracecat.agent.mcp.secret_resolution import TemplatedMappingResolutionError from tracecat.agent.preset.resolver import resolve_agents_config from tracecat.agent.preset.schemas import ( AgentPresetCreate, @@ -277,6 +278,88 @@ async def registry_actions( return actions +@pytest.mark.anyio +async def test_resolve_stdio_env_resolves_secrets_variables_and_literals( + agent_preset_service: AgentPresetService, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Stdio env resolves workspace secrets, variables, and literals.""" + suffix = uuid.uuid4().hex + secret_name = f"stdio_secret_{suffix}" + secret_value = f"stdio-secret-value-{suffix}" + variable_name = f"stdio_variable_{suffix}" + variable_value = f"stdio-variable-value-{suffix}" + + async def get_action_secrets(**_: object) -> dict[str, dict[str, str]]: + return {secret_name: {"TOKEN": secret_value}} + + async def get_workspace_variables( + *_: object, **__: object + ) -> dict[str, dict[str, str]]: + return {variable_name: {"host": variable_value}} + + monkeypatch.setattr( + "tracecat.secrets.secrets_manager.get_action_secrets", + get_action_secrets, + ) + monkeypatch.setattr( + "tracecat.executor.service.get_workspace_variables", + get_workspace_variables, + ) + stdio_env = { + "TOKEN": f"${{{{ SECRETS.{secret_name}.TOKEN }}}}", + "HOST": f"prefix-${{{{ VARS.{variable_name}.host }}}}", + "LITERAL": "literal-value", + } + + resolved = await agent_preset_service.resolve_stdio_env( + stdio_env=stdio_env, + mcp_integration_id=uuid.uuid4(), + mcp_integration_slug=f"stdio-template-{suffix}", + ) + + assert resolved == { + "TOKEN": secret_value, + "HOST": f"prefix-{variable_value}", + "LITERAL": "literal-value", + } + + +@pytest.mark.anyio +async def test_resolve_stdio_env_missing_reference_fails_closed( + agent_preset_service: AgentPresetService, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Unresolvable stdio env references raise instead of resolving to None.""" + missing_secret_name = f"missing_stdio_secret_{uuid.uuid4().hex}" + + async def get_action_secrets(**_: object) -> dict[str, dict[str, str]]: + return {} + + async def get_workspace_variables( + *_: object, **__: object + ) -> dict[str, dict[str, str]]: + return {} + + monkeypatch.setattr( + "tracecat.secrets.secrets_manager.get_action_secrets", + get_action_secrets, + ) + monkeypatch.setattr( + "tracecat.executor.service.get_workspace_variables", + get_workspace_variables, + ) + + with pytest.raises(TemplatedMappingResolutionError) as exc_info: + await agent_preset_service.resolve_stdio_env( + stdio_env={"TOKEN": f"${{{{ SECRETS.{missing_secret_name}.TOKEN }}}}"}, + mcp_integration_id=uuid.uuid4(), + mcp_integration_slug="stdio-missing-secret", + ) + + assert missing_secret_name in str(exc_info.value) + + @pytest.fixture def agent_preset_create_params() -> AgentPresetCreate: """Sample agent preset creation parameters.""" @@ -320,11 +403,11 @@ async def get_workspace_variables( return {"tenant": {"id": "staging-tenant"}} monkeypatch.setattr( - "tracecat.agent.preset.service.secrets_manager.get_action_secrets", + "tracecat.secrets.secrets_manager.get_action_secrets", get_action_secrets, ) monkeypatch.setattr( - "tracecat.agent.preset.service.get_workspace_variables", + "tracecat.executor.service.get_workspace_variables", get_workspace_variables, ) diff --git a/tests/unit/test_mcp_integrations.py b/tests/unit/test_mcp_integrations.py index cc5805fa57..5ea1ce11b2 100644 --- a/tests/unit/test_mcp_integrations.py +++ b/tests/unit/test_mcp_integrations.py @@ -11,12 +11,14 @@ import contextlib import socket import uuid +from dataclasses import dataclass from datetime import UTC, datetime, timedelta from types import SimpleNamespace from unittest.mock import AsyncMock, Mock from urllib.parse import parse_qs, urlencode, urlparse import httpx +import orjson import pytest from authlib.integrations.base_client.errors import OAuthError from pydantic import SecretStr, TypeAdapter, ValidationError @@ -51,6 +53,7 @@ Workspace, ) from tracecat.exceptions import EntitlementRequired +from tracecat.executor import service as executor_service_module from tracecat.integrations.catalog.loader import catalog_id_for_slug from tracecat.integrations.catalog.service import PlatformMCPCatalogService from tracecat.integrations.enums import IntegrationStatus, MCPAuthType, OAuthGrantType @@ -94,6 +97,7 @@ OAuthRefreshBusyError, ) from tracecat.integrations.types import DCRResponse, OAuthServerMetadata +from tracecat.secrets import secrets_manager as secrets_manager_module from tracecat.tiers import defaults as tier_defaults pytestmark = pytest.mark.usefixtures("db") @@ -103,6 +107,18 @@ ) +@dataclass(frozen=True, slots=True) +class MCPTemplateContext: + """Workspace secret and variable references for MCP resolution tests.""" + + secret_name: str + secret_key: str + secret_value: str + variable_name: str + variable_key: str + variable_value: str + + class _TestCatalogEntry(dict): """Dict-backed catalog entry with attribute access for terse tests.""" @@ -354,6 +370,65 @@ async def oauth_integration( return integration +@pytest.fixture +def mcp_template_context( + monkeypatch: pytest.MonkeyPatch, +) -> MCPTemplateContext: + """Create synthetic workspace values used by MCP header expressions.""" + suffix = uuid.uuid4().hex + context = MCPTemplateContext( + secret_name=f"mcp_secret_{suffix}", + secret_key="API_KEY", + secret_value=f"secret-value-{suffix}", + variable_name=f"mcp_variable_{suffix}", + variable_key="tenant", + variable_value=f"tenant-{suffix}", + ) + + async def get_action_secrets( + *, + secret_exprs: set[str], + action_secrets: set[object], + environment: str | None = None, + ) -> dict[str, dict[str, str]]: + _ = action_secrets, environment + secret_prefix = f"{context.secret_name}." + if not any(ref.startswith(secret_prefix) for ref in secret_exprs): + return {} + return { + context.secret_name: { + context.secret_key: context.secret_value, + } + } + + async def get_workspace_variables( + variable_exprs: set[str], + *, + environment: str | None = None, + role: Role | None = None, + ) -> dict[str, dict[str, str]]: + _ = environment, role + if context.variable_name not in variable_exprs: + return {} + return { + context.variable_name: { + context.variable_key: context.variable_value, + } + } + + monkeypatch.setattr( + secrets_manager_module, + "get_action_secrets", + get_action_secrets, + ) + monkeypatch.setattr( + executor_service_module, + "get_workspace_variables", + get_workspace_variables, + ) + return context + + @pytest.mark.anyio class TestMCPIntegrationCRUD: """Test basic CRUD operations for MCP integrations.""" @@ -6298,6 +6373,138 @@ async def test_resolve_http_config_custom_headers( assert server_config.get("headers") == {"X-API-Key": "secret-key"} + async def test_resolve_http_config_custom_header_expressions( + self, + integration_service: IntegrationService, + mcp_template_context: MCPTemplateContext, + ) -> None: + """CUSTOM headers resolve secrets, variables, and mixed literals.""" + secret_expr = ( + "${{ SECRETS." + f"{mcp_template_context.secret_name}.{mcp_template_context.secret_key}" + " }}" + ) + variable_expr = ( + "${{ VARS." + f"{mcp_template_context.variable_name}.{mcp_template_context.variable_key}" + " }}" + ) + configured_headers = { + "Authorization": f"ApiKey {secret_expr}", + "X-Tenant": f"tenant={variable_expr}", + "X-Literal": "literal-value", + } + mcp_integration = await integration_service.create_mcp_integration( + params=MCPHttpIntegrationCreate( + name="Templated Custom Auth MCP", + server_uri="https://custom.example.com/mcp", + auth_type=MCPAuthType.CUSTOM, + custom_credentials=SecretStr(orjson.dumps(configured_headers).decode()), + ) + ) + + server_config = await integration_service.resolve_mcp_http_server_config( + mcp_integration + ) + + headers = server_config.get("headers") + assert headers is not None + assert headers == { + "Authorization": f"ApiKey {mcp_template_context.secret_value}", + "X-Tenant": f"tenant={mcp_template_context.variable_value}", + "X-Literal": "literal-value", + } + assert all("${{" not in value for value in headers.values()) + + async def test_resolve_http_config_uses_effective_environment( + self, + integration_service: IntegrationService, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + """Headers resolve secrets and variables from the action environment.""" + suffix = uuid.uuid4().hex + secret_name = f"mcp_environment_secret_{suffix}" + variable_name = f"mcp_environment_variable_{suffix}" + values_by_environment = { + "default": "default-tenant", + "staging": "staging-tenant", + } + requested_secret_environments: list[str | None] = [] + requested_variable_environments: list[str | None] = [] + + async def get_action_secrets( + *, + secret_exprs: set[str], + action_secrets: set[object], + environment: str | None = None, + ) -> dict[str, dict[str, str]]: + _ = action_secrets + requested_secret_environments.append(environment) + assert secret_exprs == {f"{secret_name}.TOKEN"} + assert environment is not None + return { + secret_name: { + "TOKEN": f"{values_by_environment[environment]}-token", + } + } + + async def get_workspace_variables( + variable_exprs: set[str], + *, + environment: str | None = None, + role: Role | None = None, + ) -> dict[str, dict[str, str]]: + _ = role + requested_variable_environments.append(environment) + assert variable_exprs == {variable_name} + assert environment is not None + return { + variable_name: { + "tenant": values_by_environment[environment], + } + } + + monkeypatch.setattr( + secrets_manager_module, + "get_action_secrets", + get_action_secrets, + ) + monkeypatch.setattr( + executor_service_module, + "get_workspace_variables", + get_workspace_variables, + ) + + mcp_integration = await integration_service.create_mcp_integration( + params=MCPHttpIntegrationCreate( + name="Environment-scoped MCP", + server_uri="https://custom.example.com/mcp", + auth_type=MCPAuthType.CUSTOM, + custom_credentials=SecretStr( + orjson.dumps( + { + "Authorization": ( + f"Bearer ${{{{ SECRETS.{secret_name}.TOKEN }}}}" + ), + "X-Tenant": f"${{{{ VARS.{variable_name}.tenant }}}}", + } + ).decode() + ), + ) + ) + + server_config = await integration_service.resolve_mcp_http_server_config( + mcp_integration, + environment="staging", + ) + + assert server_config.get("headers") == { + "Authorization": "Bearer staging-tenant-token", + "X-Tenant": "staging-tenant", + } + assert requested_secret_environments == ["staging"] + assert requested_variable_environments == ["staging"] + async def test_resolve_http_config_oauth2_drops_custom_authorization( self, integration_service: IntegrationService, @@ -6326,6 +6533,207 @@ async def test_resolve_http_config_oauth2_drops_custom_authorization( assert "authorization" not in headers assert headers["X-Tenant"] == "t1" + async def test_resolve_http_config_oauth2_extra_header_expressions( + self, + integration_service: IntegrationService, + oauth_integration: OAuthIntegration, + mcp_template_context: MCPTemplateContext, + ) -> None: + """OAuth extra headers resolve while OAuth Authorization remains authoritative.""" + configured_headers = { + "authorization": "Bearer attacker", + "X-API-Key": ( + "${{ SECRETS." + f"{mcp_template_context.secret_name}.{mcp_template_context.secret_key}" + " }}" + ), + "X-Tenant": ( + "prefix-${{ VARS." + f"{mcp_template_context.variable_name}." + f"{mcp_template_context.variable_key}" + " }}" + ), + } + mcp_integration = await integration_service.create_mcp_integration( + params=MCPHttpIntegrationCreate( + name="Templated OAuth MCP", + server_uri="https://oauth.example.com/mcp", + auth_type=MCPAuthType.OAUTH2, + oauth_integration_id=oauth_integration.id, + custom_credentials=SecretStr(orjson.dumps(configured_headers).decode()), + ) + ) + + server_config = await integration_service.resolve_mcp_http_server_config( + mcp_integration + ) + + headers = server_config.get("headers") + assert headers is not None + assert headers["Authorization"] == "Bearer test_access_token" + assert "authorization" not in headers + assert headers["X-API-Key"] == mcp_template_context.secret_value + assert headers["X-Tenant"] == (f"prefix-{mcp_template_context.variable_value}") + assert all("${{" not in value for value in headers.values()) + + async def test_oauth2_malformed_extra_headers_are_dropped( + self, + integration_service: IntegrationService, + oauth_integration: OAuthIntegration, + ) -> None: + """Malformed optional OAuth headers do not disable a valid OAuth token.""" + mcp_integration = await integration_service.create_mcp_integration( + params=MCPHttpIntegrationCreate( + name="Malformed OAuth Headers MCP", + server_uri="https://oauth.example.com/mcp", + auth_type=MCPAuthType.OAUTH2, + oauth_integration_id=oauth_integration.id, + ) + ) + mcp_integration.encrypted_headers = integration_service._encrypt_token( + "{not-json" + ) + + server_config = await integration_service.resolve_mcp_http_server_config( + mcp_integration + ) + + assert server_config.get("headers") == { + "Authorization": "Bearer test_access_token" + } + + async def test_oauth2_unresolvable_extra_header_fails_closed( + self, + integration_service: IntegrationService, + oauth_integration: OAuthIntegration, + mcp_template_context: MCPTemplateContext, + ) -> None: + """Expression failures in optional OAuth headers are not warn-and-dropped.""" + missing_secret_name = f"missing_oauth_header_{uuid.uuid4().hex}" + mcp_integration = await integration_service.create_mcp_integration( + params=MCPHttpIntegrationCreate( + name="Missing OAuth Header Secret MCP", + server_uri="https://oauth.example.com/mcp", + auth_type=MCPAuthType.OAUTH2, + oauth_integration_id=oauth_integration.id, + custom_credentials=SecretStr( + orjson.dumps( + { + "X-API-Key": ( + f"${{{{ SECRETS.{missing_secret_name}.API_KEY }}}}" + ) + } + ).decode() + ), + ) + ) + + with pytest.raises(MCPConfigurationError) as exc_info: + await integration_service.resolve_mcp_http_server_config(mcp_integration) + + detail = str(exc_info.value) + assert missing_secret_name in detail + assert mcp_template_context.secret_value not in detail + + async def test_missing_header_secret_fails_resolution_and_verification( + self, + integration_service: IntegrationService, + mcp_template_context: MCPTemplateContext, + ) -> None: + """Missing secret references fail closed without exposing secret values.""" + missing_secret_name = f"missing_mcp_secret_{uuid.uuid4().hex}" + mcp_integration = await integration_service.create_mcp_integration( + params=MCPHttpIntegrationCreate( + name="Missing Secret MCP", + server_uri="https://custom.example.com/mcp", + auth_type=MCPAuthType.CUSTOM, + custom_credentials=SecretStr( + orjson.dumps( + { + "Authorization": ( + "ApiKey ${{ SECRETS." + f"{missing_secret_name}.API_KEY" + " }}" + ) + } + ).decode() + ), + ) + ) + + with pytest.raises(MCPConfigurationError) as exc_info: + await integration_service.resolve_mcp_http_server_config(mcp_integration) + + detail = str(exc_info.value) + assert missing_secret_name in detail + assert mcp_template_context.secret_value not in detail + + with pytest.raises(MCPConnectionVerificationError) as verify_exc_info: + await integration_service._probe_mcp_http_server(mcp_integration) + + assert ( + verify_exc_info.value.message + == "MCP integration is not configured correctly" + ) + assert verify_exc_info.value.error is not None + assert missing_secret_name in verify_exc_info.value.error + assert mcp_template_context.secret_value not in verify_exc_info.value.error + + async def test_unresolved_header_template_marker_fails_closed( + self, + integration_service: IntegrationService, + ) -> None: + """Malformed template markers are never sent as literal header values.""" + mcp_integration = await integration_service.create_mcp_integration( + params=MCPHttpIntegrationCreate( + name="Malformed Template MCP", + server_uri="https://custom.example.com/mcp", + auth_type=MCPAuthType.CUSTOM, + custom_credentials=SecretStr( + '{"Authorization": "ApiKey ${{ SECRETS.example.API_KEY"}' + ), + ) + ) + + with pytest.raises(MCPConfigurationError) as exc_info: + await integration_service.resolve_mcp_http_server_config(mcp_integration) + + assert "Authorization" in str(exc_info.value) + + async def test_missing_header_secret_key_does_not_expose_values( + self, + integration_service: IntegrationService, + mcp_template_context: MCPTemplateContext, + ) -> None: + """Missing secret keys name the reference without exposing sibling values.""" + missing_key = "MISSING_KEY" + mcp_integration = await integration_service.create_mcp_integration( + params=MCPHttpIntegrationCreate( + name="Missing Secret Key MCP", + server_uri="https://custom.example.com/mcp", + auth_type=MCPAuthType.CUSTOM, + custom_credentials=SecretStr( + orjson.dumps( + { + "Authorization": ( + "${{ SECRETS." + f"{mcp_template_context.secret_name}.{missing_key}" + " }}" + ) + } + ).decode() + ), + ) + ) + + with pytest.raises(MCPConfigurationError) as exc_info: + await integration_service.resolve_mcp_http_server_config(mcp_integration) + + detail = str(exc_info.value) + assert mcp_template_context.secret_name in detail + assert missing_key in detail + assert mcp_template_context.secret_value not in detail + async def test_resolve_http_config_translates_busy_oauth_refresh( self, integration_service: IntegrationService, diff --git a/tracecat/agent/mcp/secret_resolution.py b/tracecat/agent/mcp/secret_resolution.py new file mode 100644 index 0000000000..2c442e5fe6 --- /dev/null +++ b/tracecat/agent/mcp/secret_resolution.py @@ -0,0 +1,155 @@ +"""Resolve workspace-backed template expressions in MCP string mappings.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import cast + +from tracecat.auth.types import Role +from tracecat.dsl.common import create_default_execution_context +from tracecat.exceptions import TracecatException, TracecatValidationError +from tracecat.expressions.core import CollectedExprs +from tracecat.expressions.eval import collect_expressions, eval_templated_object +from tracecat.secrets.constants import DEFAULT_SECRETS_ENVIRONMENT + + +@dataclass(frozen=True, slots=True) +class TemplatedMappingResolution: + """Resolved mapping plus non-sensitive reference counts for logging.""" + + mapping: dict[str, str] + secret_ref_count: int + var_ref_count: int + + +class TemplatedMappingResolutionError(TracecatException): + """Raised when mapping resolution cannot safely produce strings.""" + + +def _template_target( + mapping: dict[str, str], *, resolve_keys: bool +) -> dict[str, str] | list[str]: + return mapping if resolve_keys else list(mapping.values()) + + +def _template_reference_summary(collected: CollectedExprs) -> str: + references = [ + *(f"SECRETS.{path}" for path in sorted(collected.secrets)), + *(f"VARS.{name}" for name in sorted(collected.variables)), + ] + return ", ".join(f"'{reference}'" for reference in references) + + +def _template_value_keys(mapping: dict[str, str]) -> list[str]: + return sorted(key for key, value in mapping.items() if "${{" in value) + + +def _resolution_error( + mapping: dict[str, str], + collected: CollectedExprs | None = None, +) -> TemplatedMappingResolutionError: + if collected and (summary := _template_reference_summary(collected)): + return TemplatedMappingResolutionError( + f"Could not resolve template reference(s): {summary}" + ) + keys = _template_value_keys(mapping) + if keys: + return TemplatedMappingResolutionError( + "Could not resolve template expression in mapping value(s): " + f"{', '.join(repr(key) for key in keys)}" + ) + return TemplatedMappingResolutionError( + "Could not resolve template expression in mapping values" + ) + + +async def resolve_templated_mapping( + mapping: dict[str, str], + *, + role: Role, + environment: str | None = None, + resolve_keys: bool = True, +) -> TemplatedMappingResolution: + """Resolve MCP mapping expressions at call time without persisting results. + + Resolution fails closed: missing references, malformed expressions, and + unresolved template markers raise ``TemplatedMappingResolutionError``. + Workspace secrets and variables resolve in ``environment``; callers + without an action-specific environment use the workspace default. + ``resolve_keys=False`` preserves header names and resolves values only. + """ + # Keep these service imports local: executor.service imports secrets_manager, + # which imports IntegrationService, and IntegrationService imports this module. + from tracecat.executor.service import get_workspace_variables + from tracecat.secrets import secrets_manager + + effective_environment = ( + environment if environment is not None else DEFAULT_SECRETS_ENVIRONMENT + ) + target = _template_target(mapping, resolve_keys=resolve_keys) + try: + collected = collect_expressions(target) + except Exception as exc: + raise _resolution_error(mapping) from exc + + if ( + not collected.secrets + and not collected.variables + and not _template_value_keys(mapping) + ): + return TemplatedMappingResolution( + mapping=mapping, + secret_ref_count=0, + var_ref_count=0, + ) + + try: + secrets = await secrets_manager.get_action_secrets( + secret_exprs=collected.secrets, + action_secrets=set(), + environment=effective_environment, + ) + vars_map = await get_workspace_variables( + variable_exprs=collected.variables, + environment=effective_environment, + role=role, + ) + + context = create_default_execution_context() + context["SECRETS"] = secrets + context["VARS"] = vars_map + resolved = eval_templated_object(target, operand=context, strict=True) + except Exception as exc: + raise _resolution_error(mapping, collected) from exc + + if resolve_keys: + if not isinstance(resolved, dict): + raise TracecatValidationError( + "Resolved mapping must be a JSON object with string values" + ) + resolved_mapping = resolved + else: + if not isinstance(resolved, list) or len(resolved) != len(mapping): + raise TracecatValidationError( + "Resolved mapping must preserve all string values" + ) + resolved_mapping = dict(zip(mapping, resolved, strict=True)) + + non_string_keys = [ + key for key, value in resolved_mapping.items() if not isinstance(value, str) + ] + if non_string_keys: + raise TemplatedMappingResolutionError( + "Resolved mapping values must be strings " + f"(invalid keys: {sorted(non_string_keys)})" + ) + + typed_mapping = cast(dict[str, str], resolved_mapping) + if _template_value_keys(typed_mapping): + raise _resolution_error(mapping, collected) + + return TemplatedMappingResolution( + mapping=typed_mapping, + secret_ref_count=len(collected.secrets), + var_ref_count=len(collected.variables), + ) diff --git a/tracecat/agent/preset/service.py b/tracecat/agent/preset/service.py index eab638cdfa..c3064d8338 100644 --- a/tracecat/agent/preset/service.py +++ b/tracecat/agent/preset/service.py @@ -22,6 +22,7 @@ MCPServerToolSummary, MCPStdioServerConfig, ) +from tracecat.agent.mcp.secret_resolution import resolve_templated_mapping from tracecat.agent.preset.resolver import resolve_agents_config from tracecat.agent.preset.schemas import ( AgentPresetCreate, @@ -63,10 +64,7 @@ SkillVersion, ) from tracecat.db.soft_delete import with_deleted -from tracecat.dsl.common import create_default_execution_context from tracecat.exceptions import TracecatNotFoundError, TracecatValidationError -from tracecat.executor.service import get_workspace_variables -from tracecat.expressions.eval import collect_expressions, eval_templated_object from tracecat.integrations.enums import MCPAuthType from tracecat.integrations.mcp_validation import ( MCPConfigurationError, @@ -83,7 +81,6 @@ CursorPaginationParams, ) from tracecat.registry.actions.service import RegistryActionsService -from tracecat.secrets import secrets_manager from tracecat.service import BaseWorkspaceService, requires_entitlement from tracecat.settings.schemas import VersionedResourceResolutionStrategy from tracecat.settings.service import get_versioned_resource_resolution_strategy @@ -893,7 +890,8 @@ async def resolve_mcp_integrations( # Handle HTTP-type servers (default) try: http_config = await integrations_service.resolve_mcp_http_server_config( - mcp_integration + mcp_integration, + environment=environment, ) except MCPConfigurationError as e: logger.warning( @@ -1131,13 +1129,11 @@ async def resolve_mcp_integration_secrets( ) except Exception as e: logger.warning( - "Stdio env resolution failed for MCP integration %r: %s", - mcp_integration.name, - str(e), - extra={ - "workspace_id": str(self.workspace_id), - "mcp_integration_id": str(mcp_integration.id), - }, + "Stdio env resolution failed for MCP integration", + mcp_integration_name=mcp_integration.name, + error=str(e), + workspace_id=str(self.workspace_id), + mcp_integration_id=str(mcp_integration.id), ) raise MCPSecretResolutionError( "Stdio MCP integration env could not be resolved", @@ -1149,17 +1145,16 @@ async def resolve_mcp_integration_secrets( # HTTP server — resolve headers per auth type. try: server_config = await integrations_service.resolve_mcp_http_server_config( - mcp_integration + mcp_integration, + environment=environment, ) except MCPConfigurationError as e: logger.warning( - "Failed to resolve secrets for HTTP MCP integration %r: %s", - mcp_integration.name, - str(e), - extra={ - "workspace_id": str(self.workspace_id), - "mcp_integration_id": str(mcp_integration.id), - }, + "Failed to resolve secrets for HTTP MCP integration", + mcp_integration_name=mcp_integration.name, + error=str(e), + workspace_id=str(self.workspace_id), + mcp_integration_id=str(mcp_integration.id), ) if mcp_integration.auth_type in {MCPAuthType.OAUTH2, MCPAuthType.CUSTOM}: raise MCPSecretResolutionError( @@ -1179,52 +1174,24 @@ async def resolve_stdio_env( mcp_integration_slug: str, environment: str | None = None, ) -> dict[str, str]: - """Resolve template expressions in stdio_env using workspace secrets/vars.""" - collected = collect_expressions(stdio_env) - if not collected.secrets and not collected.variables: - return stdio_env - - secrets = await secrets_manager.get_action_secrets( - secret_exprs=collected.secrets, - action_secrets=set(), - environment=environment, - ) - vars_map = await get_workspace_variables( - variable_exprs=collected.variables, + """Resolve stdio env expressions at call time without persisting values.""" + resolution = await resolve_templated_mapping( + stdio_env, role=self.role, environment=environment, ) - - context = create_default_execution_context() - context["SECRETS"] = secrets - context["VARS"] = vars_map - - resolved = eval_templated_object(stdio_env, operand=context) - if not isinstance(resolved, dict): - raise TracecatValidationError( - "Resolved stdio_env must be a JSON object with string values" - ) - - non_string_keys = [ - key for key, value in resolved.items() if not isinstance(value, str) - ] - if non_string_keys: - raise TracecatValidationError( - "Resolved stdio_env values must be strings " - f"(invalid keys: {sorted(non_string_keys)})" + if resolution.secret_ref_count or resolution.var_ref_count: + logger.info( + "Resolved stdio_env template expressions", + workspace_id=str(self.workspace_id), + mcp_integration_id=str(mcp_integration_id), + mcp_integration_slug=mcp_integration_slug, + env_key_count=len(resolution.mapping), + secret_ref_count=resolution.secret_ref_count, + var_ref_count=resolution.var_ref_count, ) - logger.info( - "Resolved stdio_env template expressions", - workspace_id=str(self.workspace_id), - mcp_integration_id=str(mcp_integration_id), - mcp_integration_slug=mcp_integration_slug, - env_key_count=len(resolved), - secret_ref_count=len(collected.secrets), - var_ref_count=len(collected.variables), - ) - - return cast(dict[str, str], resolved) + return resolution.mapping async def _normalize_and_validate_slug( self, diff --git a/tracecat/expressions/common.py b/tracecat/expressions/common.py index cb185dcbaf..db4051d330 100644 --- a/tracecat/expressions/common.py +++ b/tracecat/expressions/common.py @@ -179,11 +179,15 @@ def _contains_filter(path: jsonpath_nodes.JSONPath) -> bool: if strict: # We know that if this function is called, there was a templated field. # Therefore, it means the jsonpath was valid but there was no match. - logger.error("Jsonpath no match", expr=repr(expr), operand=operand) + logger.error( + "Jsonpath no match", + expr=repr(expr), + operand_type=type(operand).__name__, + ) formatted_expr = _expr_with_context(expr, context_type) raise TracecatExpressionError( f"Couldn't resolve expression {formatted_expr!r} in the context", - detail={"expression": formatted_expr, "operand": operand}, + detail={"expression": formatted_expr}, ) # Return None instead of empty list return None diff --git a/tracecat/expressions/core.py b/tracecat/expressions/core.py index e5915de5f6..4fa86eb8fb 100644 --- a/tracecat/expressions/core.py +++ b/tracecat/expressions/core.py @@ -14,7 +14,7 @@ from tracecat.expressions.common import ExprContext, ExprOperand, ExprType from tracecat.expressions.parser.core import parser from tracecat.expressions.parser.evaluator import ExprEvaluator -from tracecat.expressions.validator.validator import BaseExprValidator +from tracecat.expressions.validator.base import BaseExprValidator from tracecat.logger import logger from tracecat.parse import traverse_expressions @@ -31,11 +31,13 @@ def __init__( *, operand: ExprOperand[str] | None = None, visitor: Visitor[Token] | None = None, + strict: bool = False, ) -> None: self._expr = expression self._operand = operand self._parser = parser self._visitor = visitor + self._strict = strict def __str__(self) -> str: return self.__repr__() @@ -72,7 +74,7 @@ def result(self) -> Any: ) from e try: - visitor = ExprEvaluator(operand=self._operand) + visitor = ExprEvaluator(operand=self._operand, strict=self._strict) if parse_tree is None: raise ValueError(f"Parser returned None for expression `{self._expr}`") return visitor.evaluate(parse_tree) @@ -142,6 +144,7 @@ def __init__( template: str, operand: ExprOperand[str] | None = None, pattern: re.Pattern[str] = patterns.TEMPLATE_STRING, + strict: bool = False, **kwargs: Any, ) -> None: match = pattern.match(template) @@ -154,7 +157,7 @@ def __init__( raise TracecatExpressionError( f"Template expression {template!r} matched pattern but contained no expression. " ) - self.expr = Expression(expr, operand=operand) + self.expr = Expression(expr, operand=operand, strict=strict) def __str__(self) -> str: return self.__repr__() diff --git a/tracecat/expressions/eval.py b/tracecat/expressions/eval.py index 88accebd0c..6b09c74cc3 100644 --- a/tracecat/expressions/eval.py +++ b/tracecat/expressions/eval.py @@ -35,9 +35,11 @@ def _eval_templated_obj_rec[T: (str, list[Any], dict[str, Any])]( return obj -def _eval_expression_op(match: re.Match[str], operand: ExprOperand | None) -> str: +def _eval_expression_op( + match: re.Match[str], operand: ExprOperand | None, *, strict: bool +) -> str: expr = match.group("template") - result = TemplateExpression(expr, operand=operand).result() + result = TemplateExpression(expr, operand=operand, strict=strict).result() try: return str(result) except Exception as e: @@ -49,9 +51,10 @@ def eval_templated_object( *, operand: ExprOperand | None = None, pattern: re.Pattern[str] = patterns.TEMPLATE_STRING, + strict: bool = False, ) -> Any: """Populate templated fields with actual values.""" - evaluator = partial(_eval_expression_op, operand=operand) + evaluator = partial(_eval_expression_op, operand=operand, strict=strict) def operator(line: str) -> Any: """Evaluate the templated string. @@ -68,7 +71,7 @@ def operator(line: str) -> Any: # Non-inline template # If the template expression isn't given a reolve type, its underlying # value is returned as is. - return TemplateExpression(line, operand=operand).result() + return TemplateExpression(line, operand=operand, strict=strict).result() # Inline template # If the template expression is inline, we evaluate the result # and attempt to cast each underlying value into a string. diff --git a/tracecat/integrations/service.py b/tracecat/integrations/service.py index 6901b0ebd9..daad3f4eb3 100644 --- a/tracecat/integrations/service.py +++ b/tracecat/integrations/service.py @@ -29,6 +29,7 @@ from tracecat import config from tracecat.agent.common.types import MCPHttpServerConfig +from tracecat.agent.mcp import secret_resolution as mcp_secret_resolution from tracecat.agent.mcp.stdio_probe_types import ( MCP_STDIO_PROBE_TIMEOUT_CAP, StdioMCPProbeResult, @@ -3919,8 +3920,41 @@ def _decrypt_mcp_custom_headers( ) return cast(dict[str, str], parsed) + async def _resolve_mcp_custom_header_templates( + self, + *, + mcp_integration: MCPIntegration, + custom_headers: dict[str, str], + environment: str | None = None, + ) -> dict[str, str]: + """Resolve custom header values at call time and fail closed.""" + try: + resolution = await mcp_secret_resolution.resolve_templated_mapping( + custom_headers, + role=self.role, + environment=environment, + resolve_keys=False, + ) + except mcp_secret_resolution.TemplatedMappingResolutionError as exc: + raise MCPConfigurationError(str(exc)) from exc + + if resolution.secret_ref_count or resolution.var_ref_count: + self.logger.info( + "Resolved HTTP MCP header template expressions", + workspace_id=str(self.workspace_id), + mcp_integration_id=str(mcp_integration.id), + mcp_integration_slug=mcp_integration.slug, + header_key_count=len(resolution.mapping), + secret_ref_count=resolution.secret_ref_count, + var_ref_count=resolution.var_ref_count, + ) + return resolution.mapping + async def resolve_mcp_http_server_config( - self, mcp_integration: MCPIntegration + self, + mcp_integration: MCPIntegration, + *, + environment: str | None = None, ) -> MCPHttpServerConfig: """Resolve an HTTP MCP integration into a connectable server config. @@ -3983,9 +4017,20 @@ async def resolve_mcp_http_server_config( # The OAuth Authorization header always wins. if key.strip().casefold() == "authorization": custom_headers.pop(key, None) + custom_headers = await self._resolve_mcp_custom_header_templates( + mcp_integration=mcp_integration, + custom_headers=custom_headers, + environment=environment, + ) headers.update(custom_headers) elif mcp_integration.auth_type == MCPAuthType.CUSTOM: - headers.update(self._decrypt_mcp_custom_headers(mcp_integration)) + custom_headers = self._decrypt_mcp_custom_headers(mcp_integration) + custom_headers = await self._resolve_mcp_custom_header_templates( + mcp_integration=mcp_integration, + custom_headers=custom_headers, + environment=environment, + ) + headers.update(custom_headers) elif mcp_integration.auth_type == MCPAuthType.NONE: pass else: @@ -3993,6 +4038,15 @@ async def resolve_mcp_http_server_config( f"Unsupported MCP auth type: {mcp_integration.auth_type}" ) + unresolved_header_names = [ + name for name, value in headers.items() if "${{" in value + ] + if unresolved_header_names: + raise MCPConfigurationError( + "Unresolved template expression in HTTP MCP header(s): " + f"{', '.join(sorted(unresolved_header_names))}" + ) + server_config: MCPHttpServerConfig = { "type": "http", "name": mcp_integration.name,