feat: add policy engine selection to gateway TUI wizard

Author: Hweinstock

src/cli/commands/remove/types.ts

@@ -1,4 +1,13 @@
-export type ResourceType = 'agent' | 'gateway' | 'gateway-target' | 'memory' | 'identity' | 'evaluator' | 'online-eval' | 'policy-engine' | 'policy';
+export type ResourceType =
+  | 'agent'
+  | 'gateway'
+  | 'gateway-target'
+  | 'memory'
+  | 'identity'
+  | 'evaluator'
+  | 'online-eval'
+  | 'policy-engine'
+  | 'policy';
 
 export interface RemoveOptions {
   resourceType: ResourceType;

src/cli/commands/status/action.ts

@@ -15,7 +15,15 @@ import type { ResourceDeploymentState } from './constants';
 export type { ResourceDeploymentState };
 
 export interface ResourceStatusEntry {
-  resourceType: 'agent' | 'memory' | 'credential' | 'gateway' | 'evaluator' | 'online-eval' | 'policy-engine' | 'policy';
+  resourceType:
+    | 'agent'
+    | 'memory'
+    | 'credential'
+    | 'gateway'
+    | 'evaluator'
+    | 'online-eval'
+    | 'policy-engine'
+    | 'policy';
   name: string;
   deploymentState: ResourceDeploymentState;
   identifier?: string;
@@ -201,7 +209,16 @@ export function computeResourceStatuses(
     getDeployedKey: item => `${item.engineName}/${item.name}`,
   });
 
-  return [...agents, ...credentials, ...memories, ...gateways, ...evaluators, ...onlineEvalConfigs, ...policyEngines, ...policies];
+  return [
+    ...agents,
+    ...credentials,
+    ...memories,
+    ...gateways,
+    ...evaluators,
+    ...onlineEvalConfigs,
+    ...policyEngines,
+    ...policies,
+  ];
 }
 
 export async function handleProjectStatus(

src/cli/commands/status/command.tsx

@@ -7,7 +7,16 @@ import { DEPLOYMENT_STATE_COLORS, DEPLOYMENT_STATE_LABELS } from './constants';
 import type { Command } from '@commander-js/extra-typings';
 import { Box, Text, render } from 'ink';
 
-const VALID_RESOURCE_TYPES = ['agent', 'memory', 'credential', 'gateway', 'evaluator', 'online-eval', 'policy-engine', 'policy'] as const;
+const VALID_RESOURCE_TYPES = [
+  'agent',
+  'memory',
+  'credential',
+  'gateway',
+  'evaluator',
+  'online-eval',
+  'policy-engine',
+  'policy',
+] as const;
 const VALID_STATES = ['deployed', 'local-only', 'pending-removal'] as const;
 
 interface StatusCliOptions {

src/cli/logging/remove-logger.ts

@@ -7,7 +7,16 @@ const REMOVE_LOGS_SUBDIR = 'remove';
 
 export interface RemoveLoggerOptions {
   /** Type of resource being removed */
-  resourceType: 'agent' | 'memory' | 'identity' | 'gateway' | 'gateway-target' | 'evaluator' | 'online-eval' | 'policy-engine' | 'policy';
+  resourceType:
+    | 'agent'
+    | 'memory'
+    | 'identity'
+    | 'gateway'
+    | 'gateway-target'
+    | 'evaluator'
+    | 'online-eval'
+    | 'policy-engine'
+    | 'policy';
   /** Name of the resource being removed */
   resourceName: string;
 }

src/cli/primitives/PolicyEnginePrimitive.ts

@@ -1,6 +1,6 @@
 import { findConfigRoot } from '../../lib';
 import type { AgentCoreProjectSpec, PolicyEngine } from '../../schema';
-import { PolicyEngineSchema } from '../../schema';
+import { PolicyEngineModeSchema, PolicyEngineSchema } from '../../schema';
 import { getErrorMessage } from '../errors';
 import type { RemovalPreview, RemovalResult, SchemaChange } from '../operations/remove/types';
 import { BasePrimitive } from './BasePrimitive';
@@ -149,6 +149,34 @@ export class PolicyEnginePrimitive extends BasePrimitive<AddPolicyEngineOptions,
     }
   }
 
+  /**
+   * Get gateway names that don't have a policy engine attached.
+   */
+  async getUnprotectedGateways(): Promise<string[]> {
+    try {
+      if (!this.configIO.configExists('mcp')) return [];
+      const mcpSpec = await this.configIO.readMcpSpec();
+      return mcpSpec.agentCoreGateways.filter(gw => !gw.policyEngineConfiguration).map(gw => gw.name);
+    } catch {
+      return [];
+    }
+  }
+
+  /**
+   * Attach a policy engine to the specified gateways in mcp.json.
+   */
+  async attachToGateways(engineName: string, gatewayNames: string[], mode: 'LOG_ONLY' | 'ENFORCE'): Promise<void> {
+    if (gatewayNames.length === 0 || !this.configIO.configExists('mcp')) return;
+    const mcpSpec = await this.configIO.readMcpSpec();
+    const nameSet = new Set(gatewayNames);
+    for (const gw of mcpSpec.agentCoreGateways) {
+      if (nameSet.has(gw.name)) {
+        gw.policyEngineConfiguration = { policyEngineName: engineName, mode };
+      }
+    }
+    await this.configIO.writeMcpSpec(mcpSpec);
+  }
+
   async getDeployedEngineId(engineName: string): Promise<string | null> {
     try {
       const deployedState = await this.configIO.readDeployedState();
@@ -197,9 +225,18 @@ export class PolicyEnginePrimitive extends BasePrimitive<AddPolicyEngineOptions,
       .option('--name <name>', 'Policy engine name [non-interactive]')
       .option('--description <desc>', 'Policy engine description [non-interactive]')
       .option('--encryption-key-arn <arn>', 'KMS encryption key ARN [non-interactive]')
+      .option('--attach-to-gateways <gateways>', 'Comma-separated gateway names to attach this engine to')
+      .option('--attach-mode <mode>', 'Enforcement mode for attached gateways: LOG_ONLY or ENFORCE')
       .option('--json', 'Output as JSON [non-interactive]')
       .action(
-        async (cliOptions: { name?: string; description?: string; encryptionKeyArn?: string; json?: boolean }) => {
+        async (cliOptions: {
+          name?: string;
+          description?: string;
+          encryptionKeyArn?: string;
+          attachToGateways?: string;
+          attachMode?: string;
+          json?: boolean;
+        }) => {
           try {
             if (!findConfigRoot()) {
               console.error('No agentcore project found. Run `agentcore create` first.');
@@ -222,6 +259,16 @@ export class PolicyEnginePrimitive extends BasePrimitive<AddPolicyEngineOptions,
                 encryptionKeyArn: cliOptions.encryptionKeyArn,
               });
 
+              // Attach to gateways if requested
+              if (result.success && cliOptions.attachToGateways) {
+                const mode = PolicyEngineModeSchema.parse(cliOptions.attachMode ?? 'LOG_ONLY');
+                const gateways = cliOptions.attachToGateways
+                  .split(',')
+                  .map(s => s.trim())
+                  .filter(Boolean);
+                await this.attachToGateways(cliOptions.name, gateways, mode);
+              }
+
               if (cliOptions.json) {
                 console.log(JSON.stringify(result));
               } else if (result.success) {

src/cli/tui/hooks/useCreateMcp.ts

@@ -1,4 +1,9 @@
-import { agentPrimitive, gatewayPrimitive, gatewayTargetPrimitive } from '../../primitives/registry';
+import {
+  agentPrimitive,
+  gatewayPrimitive,
+  gatewayTargetPrimitive,
+  policyEnginePrimitive,
+} from '../../primitives/registry';
 import type { AddGatewayConfig } from '../screens/mcp/types';
 import { useCallback, useEffect, useState } from 'react';
 
@@ -30,6 +35,8 @@ export function useCreateGateway() {
         agentClientSecret: config.jwtConfig?.agentClientSecret,
         enableSemanticSearch: config.enableSemanticSearch,
         exceptionLevel: config.exceptionLevel,
+        policyEngine: config.policyEngineConfiguration?.policyEngineName,
+        policyEngineMode: config.policyEngineConfiguration?.mode,
       });
       if (!addResult.success) {
         throw new Error(addResult.error ?? 'Failed to create gateway');
@@ -70,6 +77,25 @@ export function useExistingGateways() {
   return { gateways, refresh };
 }
 
+export function useExistingPolicyEngines() {
+  const [engines, setEngines] = useState<string[]>([]);
+
+  useEffect(() => {
+    async function load() {
+      const result = await policyEnginePrimitive.getExistingEngines();
+      setEngines(result);
+    }
+    void load();
+  }, []);
+
+  const refresh = useCallback(async () => {
+    const result = await policyEnginePrimitive.getExistingEngines();
+    setEngines(result);
+  }, []);
+
+  return { engines, refresh };
+}
+
 export function useAvailableAgents() {
   const [agents, setAgents] = useState<string[] | null>(null);
 

src/cli/tui/screens/mcp/AddGatewayFlow.tsx

@@ -1,5 +1,10 @@
 import { ErrorPrompt } from '../../components';
-import { useCreateGateway, useExistingGateways, useUnassignedTargets } from '../../hooks/useCreateMcp';
+import {
+  useCreateGateway,
+  useExistingGateways,
+  useExistingPolicyEngines,
+  useUnassignedTargets,
+} from '../../hooks/useCreateMcp';
 import { AddSuccessScreen } from '../add/AddSuccessScreen';
 import { AddGatewayScreen } from './AddGatewayScreen';
 import type { AddGatewayConfig } from './types';
@@ -25,6 +30,7 @@ export function AddGatewayFlow({ isInteractive = true, onExit, onBack, onDev, on
   const { createGateway, reset: resetCreate } = useCreateGateway();
   const { gateways: existingGateways, refresh: refreshGateways } = useExistingGateways();
   const { targets: unassignedTargets } = useUnassignedTargets();
+  const { engines: existingPolicyEngines } = useExistingPolicyEngines();
   const [flow, setFlow] = useState<FlowState>({ name: 'create-wizard' });
 
   // In non-interactive mode, exit after success (but not while loading)
@@ -61,6 +67,7 @@ export function AddGatewayFlow({ isInteractive = true, onExit, onBack, onDev, on
       <AddGatewayScreen
         existingGateways={existingGateways}
         unassignedTargets={unassignedTargets}
+        existingPolicyEngines={existingPolicyEngines}
         onComplete={handleCreateComplete}
         onExit={onBack}
       />

src/cli/tui/screens/mcp/AddGatewayScreen.tsx

@@ -1,4 +1,4 @@
-import type { GatewayAuthorizerType } from '../../../../schema';
+import type { GatewayAuthorizerType, PolicyEngineMode } from '../../../../schema';
 import { GatewayNameSchema } from '../../../../schema';
 import { computeManagedOAuthCredentialName } from '../../../primitives/credential-utils';
 import {
@@ -20,6 +20,8 @@ import {
   AUTHORIZER_TYPE_OPTIONS,
   EXCEPTION_LEVEL_ITEM_ID,
   GATEWAY_STEP_LABELS,
+  NONE_SELECTION,
+  POLICY_ENGINE_MODE_OPTIONS,
   SEMANTIC_SEARCH_ITEM_ID,
 } from './types';
 import { useAddGatewayWizard } from './useAddGatewayWizard';
@@ -31,12 +33,19 @@ interface AddGatewayScreenProps {
   onExit: () => void;
   existingGateways: string[];
   unassignedTargets: string[];
+  existingPolicyEngines: string[];
 }
 
 const INITIAL_ADVANCED_SELECTED = [SEMANTIC_SEARCH_ITEM_ID];
 
-export function AddGatewayScreen({ onComplete, onExit, existingGateways, unassignedTargets }: AddGatewayScreenProps) {
-  const wizard = useAddGatewayWizard(unassignedTargets.length);
+export function AddGatewayScreen({
+  onComplete,
+  onExit,
+  existingGateways,
+  unassignedTargets,
+  existingPolicyEngines,
+}: AddGatewayScreenProps) {
+  const wizard = useAddGatewayWizard(unassignedTargets.length, existingPolicyEngines.length);
 
   // JWT config sub-step tracking (0=discoveryUrl, 1=audience, 2=clients, 3=scopes, 4=agentClientId, 5=agentClientSecret)
   const [jwtSubStep, setJwtSubStep] = useState(0);
@@ -64,10 +73,35 @@ export function AddGatewayScreen({ onComplete, onExit, existingGateways, unassig
     []
   );
 
+  // Policy engine sub-step: 0 = select engine, 1 = select mode
+  // Reset when re-entering the step (e.g., after navigating back)
+  const [policyEngineSubStep, setPolicyEngineSubStep] = useState(0);
+  const [selectedEngineName, setSelectedEngineName] = useState('');
+  React.useEffect(() => {
+    if (wizard.step === 'policy-engine') {
+      setPolicyEngineSubStep(0);
+      setSelectedEngineName('');
+    }
+  }, [wizard.step]);
+
+  const policyEngineItems: SelectableItem[] = useMemo(
+    () => [
+      { id: NONE_SELECTION, title: 'None', description: 'No policy engine' },
+      ...existingPolicyEngines.map(name => ({ id: name, title: name })),
+    ],
+    [existingPolicyEngines]
+  );
+
+  const policyEngineModeItems: SelectableItem[] = useMemo(
+    () => POLICY_ENGINE_MODE_OPTIONS.map(o => ({ id: o.id, title: o.title, description: o.description })),
+    []
+  );
+
   const isNameStep = wizard.step === 'name';
   const isAuthorizerStep = wizard.step === 'authorizer';
   const isJwtConfigStep = wizard.step === 'jwt-config';
   const isIncludeTargetsStep = wizard.step === 'include-targets';
+  const isPolicyEngineStep = wizard.step === 'policy-engine';
   const isAdvancedConfigStep = wizard.step === 'advanced-config';
   const isConfirmStep = wizard.step === 'confirm';
 
@@ -87,6 +121,36 @@ export function AddGatewayScreen({ onComplete, onExit, existingGateways, unassig
     requireSelection: false,
   });
 
+  const policyEngineNav = useListNavigation({
+    items: policyEngineItems,
+    onSelect: item => {
+      if (item.id === NONE_SELECTION) {
+        wizard.skipPolicyEngine();
+      } else {
+        setSelectedEngineName(item.id);
+        setPolicyEngineSubStep(1);
+      }
+    },
+    onExit: () => {
+      if (policyEngineSubStep === 0) {
+        wizard.goBack();
+      } else {
+        setPolicyEngineSubStep(0);
+      }
+    },
+    isActive: isPolicyEngineStep && policyEngineSubStep === 0,
+  });
+
+  const policyEngineModeNav = useListNavigation({
+    items: policyEngineModeItems,
+    onSelect: item => {
+      wizard.setPolicyEngineConfig(selectedEngineName, item.id as PolicyEngineMode);
+      setPolicyEngineSubStep(0);
+    },
+    onExit: () => setPolicyEngineSubStep(0),
+    isActive: isPolicyEngineStep && policyEngineSubStep === 1,
+  });
+
   const advancedNav = useMultiSelectNavigation({
     items: advancedConfigItems,
     getId: item => item.id,
@@ -172,7 +236,7 @@ export function AddGatewayScreen({ onComplete, onExit, existingGateways, unassig
       ? 'Space toggle · Enter confirm · Esc back'
       : isConfirmStep
         ? HELP_TEXT.CONFIRM_CANCEL
-        : isAuthorizerStep
+        : isAuthorizerStep || isPolicyEngineStep
           ? HELP_TEXT.NAVIGATE_SELECT
           : HELP_TEXT.TEXT_INPUT;
 
@@ -234,6 +298,24 @@ export function AddGatewayScreen({ onComplete, onExit, existingGateways, unassig
             <Text dimColor>No unassigned targets available. Press Enter to continue.</Text>
           ))}
 
+        {isPolicyEngineStep && policyEngineSubStep === 0 && (
+          <WizardSelect
+            title="Select a policy engine"
+            description="Attach a Cedar policy engine to authorize tool calls on this gateway"
+            items={policyEngineItems}
+            selectedIndex={policyEngineNav.selectedIndex}
+          />
+        )}
+
+        {isPolicyEngineStep && policyEngineSubStep === 1 && (
+          <WizardSelect
+            title="Select enforcement mode"
+            description={`Policy engine: ${selectedEngineName}`}
+            items={policyEngineModeItems}
+            selectedIndex={policyEngineModeNav.selectedIndex}
+          />
+        )}
+
         {isAdvancedConfigStep && (
           <Box flexDirection="column">
             <Text bold>Advanced Configuration</Text>
@@ -286,6 +368,12 @@ export function AddGatewayScreen({ onComplete, onExit, existingGateways, unassig
               },
               { label: 'Semantic Search', value: wizard.config.enableSemanticSearch ? 'Enabled' : 'Disabled' },
               { label: 'Exception Level', value: wizard.config.exceptionLevel === 'DEBUG' ? 'Debug' : 'None' },
+              ...(wizard.config.policyEngineConfiguration
+                ? [
+                    { label: 'Policy Engine', value: wizard.config.policyEngineConfiguration.policyEngineName },
+                    { label: 'Enforcement Mode', value: wizard.config.policyEngineConfiguration.mode },
+                  ]
+                : []),
             ]}
           />
         )}