insights/services/mcp/scripts/yaml-config-schema.ts

/**
 * Zod schemas for MCP YAML tool definitions.
 *
 * Shared between generate-tools.ts and scaffold-yaml.ts to validate
 * that product-authored YAML configs are well-formed. Uses .strict()
 * on all objects to reject unknown keys (catches typos).
 */
import { z } from 'zod'

export const ToolConfigSchema = z
    .object({
        operation: z.string(),
        enabled: z.boolean(),
        scopes: z.array(z.string()).optional(),
        annotations: z
            .object({
                readOnly: z.boolean(),
                destructive: z.boolean(),
                idempotent: z.boolean(),
            })
            .strict()
            .optional(),
        input_schema: z.string().optional(),
        /** Optional TypeScript type expression used for the generated handler return type and request generic. */
        response_type: z.string().optional(),
        enrich_url: z.string().optional(),
        list: z.boolean().optional(),
        title: z.string().optional(),
        description: z.string().optional(),
        /** Path to a file containing the tool description (resolved relative to the YAML file). Mutually exclusive with `description`. */
        description_file: z.string().optional(),
        /**
         * One-line selection hint injected into the system prompt catalog.
         * Describes *when to pick this tool*, not what it does. Currently only
         * surfaced for `query-*` tools in the query tool catalog.
         * Example: "Time series, aggregations, formulas, comparisons"
         */
        system_prompt_hint: z.string().optional(),
        /**
         * Brief work-protocol guidance appended to the tool's *response* as `_agentNote`,
         * so the agent sees it at point of use instead of it growing the always-loaded tool
         * description. Keep it to a sentence or two and defer details to the referenced
         * tools' own descriptions.
         */
        agent_note: z.string().optional(),
        exclude_params: z.array(z.string()).optional(),
        include_params: z.array(z.string()).optional(),
        /**
         * Body key/value pairs hardcoded by the generated handler. The values are always sent and
         * always win over anything the caller supplied (the assignments are emitted after the
         * dynamic body builder). Use to attribute a tool's requests — e.g. forcing
         * `created_via: 'mcp'` so agents can't claim another origin.
         */
        inject_body: z.record(z.string(), z.union([z.string(), z.number(), z.boolean()])).optional(),
        param_overrides: z
            .record(
                z.string(),
                z
                    .object({
                        description: z.string().optional(),
                        /** Override the default value for this parameter. The field becomes optional with `.default(value)`. */
                        default: z.unknown().optional(),
                        input_schema: z.string().optional(),
                        /** Reference to a schema.json definition. Generates a Zod schema from JSON Schema at build time. */
                        schema_ref: z.string().optional(),
                        /** Properties to exclude when generating from schema_ref. */
                        exclude_properties: z.array(z.string()).optional(),
                        /**
                         * When true, the param becomes optional in the tool schema.
                         * Must be paired with `fallback` to specify the state key used
                         * to resolve the value when the caller omits it.
                         */
                        optional: z.boolean().optional(),
                        /**
                         * When true, strip the optionality the Orval body schema applies
                         * to PATCH fields, so the param is required in the tool schema.
                         * Use when the backend serializer requires the field even though
                         * the endpoint is a PATCH (drf-spectacular marks every PATCH body
                         * field optional). Mutually exclusive with `optional`.
                         */
                        required: z.boolean().optional(),
                        /**
                         * State manager key to resolve the param from when omitted.
                         * Supported keys: 'orgId' (→ getOrgID()), 'projectId' (→ getProjectId()).
                         */
                        fallback: z.enum(['orgId', 'projectId']).optional(),
                        /**
                         * Apply a narrow input cast to the parameter. Wraps the original
                         * Orval-derived schema with `z.preprocess(...)` so the field's
                         * description, integer/min/max constraints, etc. are preserved.
                         *
                         * Currently supported:
                         * - `'string-int'` — casts strings that look like a base-10 integer
                         *   (e.g. `"123"`, `"-7"`) to a number. Anything else passes through
                         *   unchanged so zod still rejects with its honest error.
                         *
                         * Mutually exclusive with `input_schema` and `schema_ref` (those
                         * fully replace the schema; cast composes with the existing one).
                         */
                        cast: z.enum(['string-int']).optional(),
                        /**
                         * Alternate key names accepted for this param and normalized to it
                         * before validation — for identifier params agents guess different
                         * spellings for (e.g. `id` ← `insightId` / `short_id`). The
                         * canonical key wins on conflict, then the first-listed alias;
                         * alias keys never reach the handler. The advertised JSON schema
                         * still shows only the canonical param. Codegen wraps the composed
                         * tool schema with `z.preprocess(normalizeParamAliases(...), ...)`
                         * — see generate-tools.ts and @/tools/cast-helpers.
                         */
                        aliases: z.array(z.string()).optional(),
                    })
                    .strict()
                    .refine((data) => !(data.input_schema && data.schema_ref), {
                        message: 'input_schema and schema_ref are mutually exclusive',
                    })
                    .refine((data) => !(data.optional && !data.fallback), {
                        message: 'optional requires a fallback key to resolve the value from state',
                    })
                    .refine((data) => !(data.optional && data.required), {
                        message: 'optional and required are mutually exclusive',
                    })
                    .refine((data) => !(data.cast && (data.input_schema || data.schema_ref)), {
                        message:
                            'cast wraps the existing field schema and cannot be combined with input_schema or schema_ref (which replace it)',
                    })
            )
            .optional(),
        /**
         * Cross-field refinements applied to the composed tool schema.
         * Each entry names a function exported from `@/schema/tool-inputs` whose
         * signature is `(data, ctx: z.RefinementCtx) => void`. The codegen wires
         * one `.superRefine(<fn>)` call per entry, in order. Use this when a single
         * field's schema can't express the rule (e.g. mutual exclusivity between
         * two optional fields).
         */
        validators: z.array(z.string()).optional(),
        /** References a key in ui_apps. */
        ui_app: z.string().optional(),
        /**
         * When true or a string, the tool issues PATCH instead of DELETE.
         * `true` sends `{ deleted: true }` (for ForbidDestroyModel endpoints).
         * A string value specifies a custom field name, e.g. `"archived"` sends
         * `{ archived: true }` (for models that use a different soft-delete field).
         */
        soft_delete: z.union([z.boolean(), z.string()]).optional(),
        /**
         * When true, the tool is only available when the organization has approved
         * AI data processing (`is_ai_data_processing_approved`). Tools that invoke
         * LLMs internally should set this to true.
         */
        requires_ai_consent: z.boolean().optional(),
        /**
         * Maps original OpenAPI field names to MCP-safe aliases. The generated tool
         * schema uses the alias (which must match ^[a-zA-Z0-9_.-]{1,64}$), while
         * the request body still sends the original field name.
         */
        rename_params: z.record(z.string(), z.string()).optional(),
        /**
         * Insights feature flag key that controls whether this tool is exposed.
         * When set, the tool is only included (or excluded) based on the flag's
         * evaluation for the current user.
         *
         * By default (`feature_flag_behavior: 'enable'`), the tool is only shown
         * when the flag is **on**. Set `feature_flag_behavior: 'disable'` to hide
         * the tool when the flag is on (useful for sunsetting old tools).
         */
        feature_flag: z.string().optional(),
        feature_entitlement: z.string().optional(),
        /**
         * Controls how `feature_flag` gates the tool:
         * - `'enable'` (default): tool is shown only when the flag is on.
         * - `'disable'`: tool is hidden when the flag is on.
         */
        feature_flag_behavior: z.enum(['enable', 'disable']).optional(),
        /** Variant of `feature_flag` to match exactly. Requires `feature_flag` to be set. */
        feature_flag_variant: z.string().optional(),
        /**
         * Response field filtering. Supports dot-path patterns with wildcards (e.g. 'filters.groups.*.key').
         * For list endpoints, applied to each item in `results`. `include` and `exclude` are mutually exclusive.
         */
        /**
         * Override the category-level URL prefix for `_insightsUrl` enrichment.
         * Useful when a single category YAML covers tools that link to different frontend pages.
         */
        url_prefix: z.string().optional(),
        response: z
            .object({
                /** Dot-path patterns of response fields to keep. Only matched fields are preserved. */
                include: z.array(z.string()).optional(),
                /** Dot-path patterns of response fields to remove. */
                exclude: z.array(z.string()).optional(),
                /**
                 * Expose an optional `fields` request param so the agent can narrow the response to a
                 * subset of `include` at call time (constrained to the allowlist). Omitting `fields`
                 * returns the full `include` set. Requires `include`; incompatible with `exclude`.
                 */
                selectable: z.boolean().optional(),
                /** Wrap user-authored response data in an explicit informational-only tag boundary. */
                informational_wrapper: z
                    .object({
                        tag: z.string().regex(/^[a-z][a-z0-9-]*$/),
                        purpose: z.string().min(1).optional(),
                    })
                    .strict()
                    .optional(),
            })
            .strict()
            .refine((data) => !(data.include?.length && data.exclude?.length), {
                message: 'response.include and response.exclude are mutually exclusive',
            })
            .refine((data) => !(data.selectable && !data.include?.length), {
                message: 'response.selectable requires response.include (the allowlist to select from)',
            })
            .optional(),
        /**
         * Opt the tool into the typed-confirm two-tool paradigm. Codegen
         * emits TWO tools per YAML entry that declares this block:
         *
         *   - `<name>-prepare` — signs the validated args + user identity into
         *     a hash and returns a result instructing the model to surface the
         *     confirmation prompt to the user.
         *   - `<name>-execute` — takes the hash plus the literal string the
         *     user typed, validates both (single-use, TTL-bounded), then runs
         *     the underlying action with the original args.
         *
         * Use for destructive or security-sensitive actions an LLM should
         * never perform unattended (org-wide settings, key revocation,
         * bulk deletes). The security guarantee is weaker than client-
         * rendered elicitation because the LLM controls the `confirmation`
         * argument — but strictly stronger than a single destructive tool.
         */
        confirmed_action: z
            .object({
                /**
                 * Prompt text shown to the user. Supports `{paramName}`
                 * placeholders interpolated from the validated tool args
                 * at runtime.
                 *
                 * Example: `"About to enable enforce 2FA on organization {orgId}. Reply 'confirm' to proceed."`
                 */
                message: z.string(),
                /**
                 * Short human-readable label for the action ("enable 2FA",
                 * "delete project"). Surfaced in refusal messages when the
                 * confirmation fails so the user sees what was refused.
                 * Defaults to the tool's title if omitted.
                 */
                action_label: z.string().optional(),
            })
            .strict()
            .optional(),
    })
    .strict()
    .refine(
        (data) =>
            !data.input_schema ||
            (!data.include_params?.length && !data.exclude_params?.length && !data.param_overrides),
        {
            message:
                'input_schema replaces the entire schema, so include_params, exclude_params, and param_overrides have no effect and should be removed',
        }
    )
    .refine((data) => !(data.confirmed_action && data.input_schema), {
        message:
            '`confirmed_action` cannot be combined with `input_schema` yet because custom input schemas bypass the confirmed-action prepare/execute codegen path. Remove `input_schema` or extend custom-schema codegen to support confirmed actions.',
        path: ['confirmed_action'],
    })
    .refine((data) => !(data.description && data.description_file), {
        message: 'description and description_file are mutually exclusive',
    })
    .refine((data) => !(data.feature_flag_variant && !data.feature_flag), {
        message: '`feature_flag_variant` requires `feature_flag` to be set',
        path: ['feature_flag_variant'],
    })
    // confirmed_action emits two factories (`-prepare`, `-execute`) via a
    // codegen path that doesn't currently wrap either with `withUiApp`.
    // Reject the combination at YAML time rather than silently drop the
    // UI app at codegen — if a real consumer wants both, wire `withUiApp`
    // around the execute factory explicitly first.
    .refine((data) => !(data.confirmed_action && data.ui_app), {
        message:
            '`confirmed_action` cannot be combined with `ui_app` yet — the codegen does not wrap the generated -execute factory with withUiApp. Drop one or extend buildConfirmedActionFactories to opt in.',
        path: ['confirmed_action'],
    })

export type ToolConfig = z.infer<typeof ToolConfigSchema>

/** Narrowed type for enabled tools — scopes and annotations are guaranteed present. */
export type EnabledToolConfig = Omit<ToolConfig, 'scopes' | 'annotations'> & {
    scopes: string[]
    annotations: { readOnly: boolean; destructive: boolean; idempotent: boolean }
}

// --- UI App schemas ---
//
// Each entry under ui_apps in a tools.yaml file defines a UI app.
// The discriminator is `type`: 'detail', 'list', or 'custom'.
//
// Most fields are optional — the generator (generate-ui-apps.ts) derives
// them from the app key + the product directory the YAML lives in.
// See resolveDetailApp() and resolveListApp() in generate-ui-apps.ts
// for the full derivation logic.
//
// To add a new field:
// 1. Add it to the appropriate schema below (with .optional() if it has a default)
// 2. Add it to the matching Resolved* interface
// 3. Add the default derivation in the resolve*App() function in generate-ui-apps.ts
// 4. Use the resolved value in generateDetailApp() or generateListApp()

/**
 * Detail UI app — renders a single entity.
 *
 * Generated entry point wraps the view component in AppWrapper and mounts it.
 * The only required field is `view_prop` — everything else is derived by convention.
 */
const DetailUiAppSchema = z
    .object({
        /** Discriminator. Must be 'detail'. */
        type: z.literal('detail'),
        /** The prop name passed to the view component. Required — cannot be derived. */
        view_prop: z.string(),
        /** Display name shown in the MCP client. Default: "Insights " + title-case of key. */
        app_name: z.string().optional(),
        /** Short description for the MCP resource. Default: title-case of key + " detail view". */
        description: z.string().optional(),
        /** Import path for the view component. Default: derived from product dir (products/{product}/mcp/apps). */
        component_import: z.string().optional(),
        /** TypeScript type for the tool result data. Default: PascalCase(key) + "Data". */
        data_type: z.string().optional(),
        /** React component name that renders the detail view. Default: PascalCase(key) + "View". */
        view_component: z.string().optional(),
        /** Prop name to bind the host's `openLink(url)` callback to. Omit if the view has no inline external links. */
        link_prop: z.string().optional(),
    })
    .strict()

/**
 * List UI app — renders a list with drill-down into detail via a tool call.
 *
 * Generated entry point includes the list component, a fallback-to-chat function,
 * and a click handler that calls `detail_tool` via app.callServerTool().
 * The only required field is `detail_tool` — everything else has a default.
 */
const ListUiAppSchema = z
    .object({
        /** Discriminator. Must be 'list'. */
        type: z.literal('list'),
        /** Tool name to call when a list item is clicked (e.g. 'action-get'). Required. */
        detail_tool: z.string(),
        /** JS expression for arguments passed to the detail tool. Default: '{ id: item.id }'. */
        detail_args: z.string().optional(),
        /** Field on the item object used for display in loading/fallback states. Default: 'name'. */
        item_name_field: z.string().optional(),
        /** Prop name for the click handler on the list component. Default: 'on' + PascalCase(singularKey) + 'Click'. */
        click_prop: z.string().optional(),
        /** Human-readable entity label for the fallback chat message. Default: kebab-to-space of singular key. */
        entity_label: z.string().optional(),
        /** Display name shown in the MCP client. Default: "Insights " + title-case of key. */
        app_name: z.string().optional(),
        /** Short description for the MCP resource. Default: title-case of key + " view". */
        description: z.string().optional(),
        /** Import path for the view component. Default: derived from product dir (products/{product}/mcp/apps). */
        component_import: z.string().optional(),
        /** TypeScript type for the full list response. Default: PascalCase(singularKey) + "ListData". */
        list_data_type: z.string().optional(),
        /** TypeScript type for a single item. Default: PascalCase(singularKey) + "Data". */
        item_data_type: z.string().optional(),
        /** React component name that renders the list view. Default: PascalCase(key) + "View". */
        view_component: z.string().optional(),
    })
    .strict()

/**
 * Custom UI app — handwritten entry point with optional render-ui dispatch.
 *
 * Use for apps that need fully custom logic (e.g. debug.tsx, query-results.tsx).
 * The generator does NOT create an entry point file — you maintain it manually at
 * services/mcp/src/ui-apps/apps/{key}.tsx.
 */
const CustomUiAppSchema = z
    .object({
        /** Discriminator. Must be 'custom'. */
        type: z.literal('custom'),
        /** Display name shown in the MCP client. Required for custom apps (no convention to derive from). */
        app_name: z.string(),
        /** Short description for the MCP resource. Required for custom apps. */
        description: z.string(),
        /** Reusable view component that lets the render-ui umbrella app mount this custom app. */
        render_ui: z
            .object({
                /** Import path for the reusable view component, relative to src/ui-apps/generated. */
                component_import: z.string(),
                /** React component name that renders the tool result. */
                view_component: z.string(),
                /** Prop name that receives the tool result. */
                view_prop: z.string(),
                /** TypeScript type for the tool result. Omit when the component accepts unknown. */
                data_type: z.string().optional(),
            })
            .strict()
            .optional(),
    })
    .strict()

export const UiAppConfigSchema = z.discriminatedUnion('type', [DetailUiAppSchema, ListUiAppSchema, CustomUiAppSchema])

export type UiAppConfig = z.infer<typeof UiAppConfigSchema>

/** Detail config with all fields resolved (after convention defaults are applied). */
export interface ResolvedDetailUiApp {
    type: 'detail'
    view_prop: string
    app_name: string
    description: string
    component_import: string
    data_type: string
    view_component: string
    link_prop?: string
}

/** List config with all fields resolved (after convention defaults are applied). */
export interface ResolvedListUiApp {
    type: 'list'
    detail_tool: string
    detail_args: string
    item_name_field: string
    click_prop: string
    entity_label: string
    app_name: string
    description: string
    component_import: string
    list_data_type: string
    item_data_type: string
    view_component: string
}

/** Custom config, including an optional reusable view for render-ui dispatch. */
export interface ResolvedCustomUiApp {
    type: 'custom'
    app_name: string
    description: string
    render_ui?: {
        component_import: string
        view_component: string
        view_prop: string
        data_type?: string
    }
}

/**
 * Some MCP clients (notably Cursor) enforce a 60-character combined limit on
 * server_name + tool_name. With server name "insights" (7 chars), tool names
 * must be <= 52 chars to stay under that limit.
 *
 * Length is enforced by lint-tool-names.ts rather than here so pre-existing
 * tools that already exceed the limit don't break schema validation.
 */
export const MAX_TOOL_NAME_LENGTH = 52

/** Tool names must be lowercase kebab-case: letters, digits, hyphens. No leading/trailing hyphens. */
export const TOOL_NAME_PATTERN = /^[a-z0-9](?:[a-z0-9-]*[a-z0-9])?$/

/** Feature identifiers must be lowercase snake_case: letters, digits, underscores. */
export const FEATURE_NAME_PATTERN = /^[a-z][a-z0-9_]*$/

// ------------------------------------------------------------------
// Query wrapper config — tools generated from frontend/src/queries/schema.json
// Defined before CategoryConfigSchema so it can be referenced there.
// ------------------------------------------------------------------

export const QueryWrapperToolConfigSchema = z
    .object({
        /** Name of the definition in schema.json (e.g. "AssistantTrendsQuery") */
        schema_ref: z.string(),
        enabled: z.boolean(),
        scopes: z.array(z.string()).optional(),
        annotations: z
            .object({
                readOnly: z.boolean(),
                destructive: z.boolean(),
                idempotent: z.boolean(),
            })
            .strict()
            .optional(),
        title: z.string().optional(),
        description: z.string().optional(),
        /** Path to a file containing the tool description (resolved relative to the YAML file). Mutually exclusive with `description`. */
        description_file: z.string().optional(),
        /**
         * One-line selection hint injected into the query tool catalog in the
         * system prompt. Describes *when to pick this tool*, not what it does.
         * Example: "Time series, aggregations, formulas, comparisons"
         */
        system_prompt_hint: z.string().optional(),
        ui_resource_uri: z.string().optional(),
        /** Properties to exclude from the generated Zod schema */
        exclude_properties: z.array(z.string()).optional(),
        /**
         * Set to `true` when the wrapper's `schema_ref` has a matching formatter in
         * `ee/scriptai/context/insight/format/`. Enabling this:
         *   - surfaces the formatter's LLM-friendly text output (via the backend's `formatted_results`)
         *     as the default response text;
         *   - adds an `output_format: 'optimized' | 'json'` per-call input to the generated tool schema,
         *     so the caller can opt into raw JSON when they want structured data instead of prose.
         * Omit (or set to `false`) when the query kind has no formatter; the tool then always returns
         * JSON-encoded results.
         */
        use_optimized_output: z.boolean().optional(),
        /**
         * Default values for properties that are required in the schema but should
         * be optional for the agent. The Zod schema gets `.default(value).optional()`.
         */
        property_defaults: z.record(z.string(), z.unknown()).optional(),
        /**
         * Override the URL enrichment prefix. When set, `_insightsUrl` uses
         * `{baseUrl}{url_prefix}` instead of the default `/insights/new#q=...`.
         */
        url_prefix: z.string().optional(),
        /**
         * Insights feature flag key that controls whether this tool is exposed.
         * See ToolConfigSchema.feature_flag for full documentation.
         */
        feature_flag: z.string().optional(),
        feature_entitlement: z.string().optional(),
        /**
         * Controls how `feature_flag` gates the tool:
         * - `'enable'` (default): tool is shown only when the flag is on.
         * - `'disable'`: tool is hidden when the flag is on.
         */
        feature_flag_behavior: z.enum(['enable', 'disable']).optional(),
        /** Variant of `feature_flag` to match exactly. Requires `feature_flag` to be set. */
        feature_flag_variant: z.string().optional(),
    })
    .strict()
    .refine((data) => !(data.description && data.description_file), {
        message: 'description and description_file are mutually exclusive',
    })
    .refine((data) => !(data.feature_flag_variant && !data.feature_flag), {
        message: '`feature_flag_variant` requires `feature_flag` to be set',
        path: ['feature_flag_variant'],
    })

export type QueryWrapperToolConfig = z.infer<typeof QueryWrapperToolConfigSchema>

export type EnabledQueryWrapperToolConfig = Omit<QueryWrapperToolConfig, 'scopes' | 'annotations'> & {
    scopes: string[]
    annotations: { readOnly: boolean; destructive: boolean; idempotent: boolean }
}

export const QueryWrappersConfigSchema = z
    .object({
        category: z.string(),
        feature: z.string(),
        wrappers: z.record(z.string(), QueryWrapperToolConfigSchema),
    })
    .strict()

export type QueryWrappersConfig = z.infer<typeof QueryWrappersConfigSchema>

// ------------------------------------------------------------------
// Category config — the top-level schema for product tools.yaml files.
// Supports both REST tools (via OpenAPI) and query wrappers (via schema.json).
// ------------------------------------------------------------------

export const CategoryConfigSchema = z
    .object({
        category: z.string(),
        feature: z.string().regex(FEATURE_NAME_PATTERN, 'Feature must be lowercase snake_case: [a-z0-9_]'),
        url_prefix: z.string(),
        /**
         * Category-level feature flag gate, inherited by every tool that doesn't
         * set its own `feature_flag`. Use to hide a whole not-yet-GA product
         * from the standard MCP surface in one place. See ToolConfigSchema.feature_flag.
         */
        feature_flag: z.string().optional(),
        feature_entitlement: z.string().optional(),
        feature_flag_behavior: z.enum(['enable', 'disable']).optional(),
        feature_flag_variant: z.string().optional(),
        tools: z.record(
            z
                .string()
                .regex(
                    TOOL_NAME_PATTERN,
                    'Tool name must be lowercase kebab-case: [a-z0-9-], no leading/trailing hyphens'
                ),
            ToolConfigSchema
        ),
        ui_apps: z.record(z.string(), UiAppConfigSchema).optional(),
        /** Query wrapper tools generated from schema.json, co-located with REST tools in the same file. */
        wrappers: z.record(z.string(), QueryWrapperToolConfigSchema).optional(),
    })
    .strict()

export type CategoryConfig = z.infer<typeof CategoryConfigSchema>