feat: add PolicyEngineConfiguration support for gateways

Author: Hweinstock

src/cli/commands/add/types.ts

@@ -39,6 +39,8 @@ export interface AddGatewayOptions {
   agents?: string;
   semanticSearch?: boolean;
   exceptionLevel?: string;
+  policyEngine?: string;
+  policyEngineMode?: string;
   json?: boolean;
 }
 

src/cli/commands/add/validate.ts

@@ -273,6 +273,17 @@ export function validateAddGatewayOptions(options: AddGatewayOptions): Validatio
     }
   }
 
+  // Validate policy engine options
+  if (options.policyEngine && !options.policyEngineMode) {
+    return { valid: false, error: '--policy-engine-mode is required when --policy-engine is specified' };
+  }
+  if (options.policyEngineMode && !options.policyEngine) {
+    return { valid: false, error: '--policy-engine is required when --policy-engine-mode is specified' };
+  }
+  if (options.policyEngineMode && !['LOG_ONLY', 'ENFORCE'].includes(options.policyEngineMode)) {
+    return { valid: false, error: `Invalid policy engine mode: ${options.policyEngineMode}. Use LOG_ONLY or ENFORCE` };
+  }
+
   return { valid: true };
 }
 

src/cli/primitives/GatewayPrimitive.ts

@@ -1,6 +1,6 @@
 import { findConfigRoot, setEnvVar } from '../../lib';
 import type { AgentCoreGateway, AgentCoreGatewayTarget, AgentCoreMcpSpec, GatewayAuthorizerType } from '../../schema';
-import { AgentCoreGatewaySchema } from '../../schema';
+import { AgentCoreGatewaySchema, PolicyEngineModeSchema } from '../../schema';
 import type { AddGatewayOptions as CLIAddGatewayOptions } from '../commands/add/types';
 import { validateAddGatewayOptions } from '../commands/add/validate';
 import { getErrorMessage } from '../errors';
@@ -28,6 +28,8 @@ export interface AddGatewayOptions {
   agents?: string;
   enableSemanticSearch?: boolean;
   exceptionLevel?: string;
+  policyEngine?: string;
+  policyEngineMode?: string;
 }
 
 /**
@@ -160,6 +162,8 @@ export class GatewayPrimitive extends BasePrimitive<AddGatewayOptions, Removable
       .option('--agents <agents>', 'Comma-separated agent names')
       .option('--no-semantic-search', 'Disable semantic search for tool discovery')
       .option('--exception-level <level>', 'Exception verbosity level', 'NONE')
+      .option('--policy-engine <name>', 'Policy engine name for Cedar-based authorization')
+      .option('--policy-engine-mode <mode>', 'Policy engine mode: LOG_ONLY or ENFORCE')
       .option('--json', 'Output as JSON')
       .action(async (rawOptions: Record<string, string | boolean | undefined>) => {
         const cliOptions = rawOptions as unknown as CLIAddGatewayOptions;
@@ -192,6 +196,8 @@ export class GatewayPrimitive extends BasePrimitive<AddGatewayOptions, Removable
             agents: cliOptions.agents,
             enableSemanticSearch: cliOptions.semanticSearch !== false,
             exceptionLevel: cliOptions.exceptionLevel,
+            policyEngine: cliOptions.policyEngine,
+            policyEngineMode: cliOptions.policyEngineMode,
           });
 
           if (cliOptions.json) {
@@ -290,6 +296,10 @@ export class GatewayPrimitive extends BasePrimitive<AddGatewayOptions, Removable
       jwtConfig: undefined,
       enableSemanticSearch: options.enableSemanticSearch ?? true,
       exceptionLevel: options.exceptionLevel === 'DEBUG' ? 'DEBUG' : 'NONE',
+      policyEngineConfiguration:
+        options.policyEngine && options.policyEngineMode
+          ? { policyEngineName: options.policyEngine, mode: PolicyEngineModeSchema.parse(options.policyEngineMode) }
+          : undefined,
     };
 
     if (options.authorizerType === 'CUSTOM_JWT' && options.discoveryUrl) {
@@ -358,6 +368,7 @@ export class GatewayPrimitive extends BasePrimitive<AddGatewayOptions, Removable
       authorizerConfiguration: this.buildAuthorizerConfiguration(config),
       enableSemanticSearch: config.enableSemanticSearch,
       exceptionLevel: config.exceptionLevel,
+      policyEngineConfiguration: config.policyEngineConfiguration,
     };
 
     mcpSpec.agentCoreGateways.push(gateway);

src/cli/primitives/PolicyEnginePrimitive.ts

@@ -53,6 +53,21 @@ export class PolicyEnginePrimitive extends BasePrimitive<AddPolicyEngineOptions,
       project.policyEngines.splice(index, 1);
       await this.writeProjectSpec(project);
 
+      // Clean up any gateway references to this engine in mcp.json
+      if (this.configIO.configExists('mcp')) {
+        const mcpSpec = await this.configIO.readMcpSpec();
+        let changed = false;
+        for (const gw of mcpSpec.agentCoreGateways) {
+          if (gw.policyEngineConfiguration?.policyEngineName === engineName) {
+            delete gw.policyEngineConfiguration;
+            changed = true;
+          }
+        }
+        if (changed) {
+          await this.configIO.writeMcpSpec(mcpSpec);
+        }
+      }
+
       return { success: true };
     } catch (err) {
       const message = err instanceof Error ? err.message : 'Unknown error';
@@ -84,6 +99,35 @@ export class PolicyEnginePrimitive extends BasePrimitive<AddPolicyEngineOptions,
       after: afterSpec,
     });
 
+    // Show mcp.json changes if any gateways reference this engine
+    if (this.configIO.configExists('mcp')) {
+      const mcpSpec = await this.configIO.readMcpSpec();
+      const affectedGateways = mcpSpec.agentCoreGateways.filter(
+        gw => gw.policyEngineConfiguration?.policyEngineName === engineName
+      );
+      if (affectedGateways.length > 0) {
+        summary.push(
+          `Note: ${affectedGateways.length} gateway(s) referencing this engine will have policyEngineConfiguration removed`
+        );
+        summary.push(
+          'Warning: this may grant agents escalated permissions to invoke gateway tools that were previously restricted'
+        );
+        const afterMcpSpec = {
+          ...mcpSpec,
+          agentCoreGateways: mcpSpec.agentCoreGateways.map(gw =>
+            gw.policyEngineConfiguration?.policyEngineName === engineName
+              ? { ...gw, policyEngineConfiguration: undefined }
+              : gw
+          ),
+        };
+        schemaChanges.push({
+          file: 'agentcore/mcp.json',
+          before: mcpSpec,
+          after: afterMcpSpec,
+        });
+      }
+    }
+
     return { summary, directoriesToDelete: [], schemaChanges };
   }
 

src/cli/tui/hooks/useRemove.ts

@@ -16,7 +16,12 @@ import {
 import { useCallback, useEffect, useRef, useState } from 'react';
 
 // Re-export types for consumers
-export type { RemovableMemory, RemovableCredential as RemovableIdentity, RemovableGatewayTarget, RemovablePolicyResource };
+export type {
+  RemovableMemory,
+  RemovableCredential as RemovableIdentity,
+  RemovableGatewayTarget,
+  RemovablePolicyResource,
+};
 
 // ============================================================================
 // Generic Hooks

src/cli/tui/screens/mcp/types.ts

@@ -2,6 +2,7 @@ import type {
   ApiGatewayHttpMethod,
   GatewayAuthorizerType,
   GatewayExceptionLevel,
+  GatewayPolicyEngineConfiguration,
   GatewayTargetType,
   NodeRuntime,
   PythonRuntime,
@@ -36,6 +37,8 @@ export interface AddGatewayConfig {
   enableSemanticSearch: boolean;
   /** Exception verbosity level for the gateway */
   exceptionLevel: GatewayExceptionLevel;
+  /** Policy engine configuration for Cedar-based authorization */
+  policyEngineConfiguration?: GatewayPolicyEngineConfiguration;
 }
 
 /** Item ID for the semantic search toggle in the advanced config pane. */

src/cli/tui/screens/remove/RemoveFlow.tsx

@@ -111,11 +111,22 @@ export function RemoveFlow({
   const { tools: mcpTools, isLoading: isLoadingTools, refresh: refreshTools } = useRemovableGatewayTargets();
   const { memories, isLoading: isLoadingMemories, refresh: refreshMemories } = useRemovableMemories();
   const { identities, isLoading: isLoadingIdentities, refresh: refreshIdentities } = useRemovableIdentities();
-  const { policyEngines, isLoading: isLoadingPolicyEngines, refresh: refreshPolicyEngines } = useRemovablePolicyEngines();
+  const {
+    policyEngines,
+    isLoading: isLoadingPolicyEngines,
+    refresh: refreshPolicyEngines,
+  } = useRemovablePolicyEngines();
   const { policies, isLoading: isLoadingPolicies, refresh: refreshPolicies } = useRemovablePolicies();
 
   // Check if any data is still loading
-  const isLoading = isLoadingAgents || isLoadingGateways || isLoadingTools || isLoadingMemories || isLoadingIdentities || isLoadingPolicyEngines || isLoadingPolicies;
+  const isLoading =
+    isLoadingAgents ||
+    isLoadingGateways ||
+    isLoadingTools ||
+    isLoadingMemories ||
+    isLoadingIdentities ||
+    isLoadingPolicyEngines ||
+    isLoadingPolicies;
 
   // Preview hook
   const {
@@ -157,7 +168,15 @@ export function RemoveFlow({
   // In non-interactive mode, exit after success
   useEffect(() => {
     if (!isInteractive) {
-      const successStates = ['agent-success', 'gateway-success', 'tool-success', 'memory-success', 'identity-success', 'policy-engine-success', 'policy-success'];
+      const successStates = [
+        'agent-success',
+        'gateway-success',
+        'tool-success',
+        'memory-success',
+        'identity-success',
+        'policy-engine-success',
+        'policy-success',
+      ];
       if (successStates.includes(flow.name)) {
         onExit();
       }
@@ -335,7 +354,9 @@ export function RemoveFlow({
     async (compositeKey: string) => {
       const result = await loadPolicyPreview(compositeKey);
       if (result.ok) {
-        const policyName = compositeKey.includes('/') ? compositeKey.slice(compositeKey.indexOf('/') + 1) : compositeKey;
+        const policyName = compositeKey.includes('/')
+          ? compositeKey.slice(compositeKey.indexOf('/') + 1)
+          : compositeKey;
         if (force) {
           setFlow({ name: 'loading', message: `Removing policy ${policyName}...` });
           const removeResult = await removePolicyOp(compositeKey, result.preview);
@@ -532,8 +553,24 @@ export function RemoveFlow({
   ]);
 
   const refreshAll = useCallback(async () => {
-    await Promise.all([refreshAgents(), refreshGateways(), refreshTools(), refreshMemories(), refreshIdentities(), refreshPolicyEngines(), refreshPolicies()]);
-  }, [refreshAgents, refreshGateways, refreshTools, refreshMemories, refreshIdentities, refreshPolicyEngines, refreshPolicies]);
+    await Promise.all([
+      refreshAgents(),
+      refreshGateways(),
+      refreshTools(),
+      refreshMemories(),
+      refreshIdentities(),
+      refreshPolicyEngines(),
+      refreshPolicies(),
+    ]);
+  }, [
+    refreshAgents,
+    refreshGateways,
+    refreshTools,
+    refreshMemories,
+    refreshIdentities,
+    refreshPolicyEngines,
+    refreshPolicies,
+  ]);
 
   // Select screen - wait for data to load to avoid arrow position issues
   if (flow.name === 'select') {

src/cli/tui/screens/remove/RemovePolicyScreen.tsx

@@ -1,5 +1,5 @@
-import type { RemovablePolicyResource } from '../../hooks/useRemove';
 import { SelectScreen } from '../../components';
+import type { RemovablePolicyResource } from '../../hooks/useRemove';
 import React from 'react';
 
 interface RemovePolicyScreenProps {
@@ -22,11 +22,6 @@ export function RemovePolicyScreen({ policies, onSelect, onExit }: RemovePolicyS
   });
 
   return (
-    <SelectScreen
-      title="Select Policy to Remove"
-      items={items}
-      onSelect={item => onSelect(item.id)}
-      onExit={onExit}
-    />
+    <SelectScreen title="Select Policy to Remove" items={items} onSelect={item => onSelect(item.id)} onExit={onExit} />
   );
 }

src/cli/tui/screens/remove/useRemoveFlow.ts

@@ -78,7 +78,9 @@ export function useRemoveFlow({ force, dryRun }: RemoveFlowOptions): RemoveFlowS
           items.push(`${projectSpec.credentials.length} credential${projectSpec.credentials.length > 1 ? 's' : ''}`);
         }
         if (projectSpec.policyEngines && projectSpec.policyEngines.length > 0) {
-          items.push(`${projectSpec.policyEngines.length} policy engine${projectSpec.policyEngines.length > 1 ? 's' : ''}`);
+          items.push(
+            `${projectSpec.policyEngines.length} policy engine${projectSpec.policyEngines.length > 1 ? 's' : ''}`
+          );
           const totalPolicies = projectSpec.policyEngines.reduce((sum, e) => sum + (e.policies?.length ?? 0), 0);
           if (totalPolicies > 0) {
             items.push(`${totalPolicies} polic${totalPolicies > 1 ? 'ies' : 'y'}`);

src/schema/schemas/mcp.ts

@@ -593,6 +593,21 @@ export type AgentCoreGatewayTarget = z.infer<typeof AgentCoreGatewayTargetSchema
 export const GatewayExceptionLevelSchema = z.enum(['NONE', 'DEBUG']);
 export type GatewayExceptionLevel = z.infer<typeof GatewayExceptionLevelSchema>;
 
+// ============================================================================
+// Gateway Policy Engine Configuration
+// ============================================================================
+
+export const PolicyEngineModeSchema = z.enum(['LOG_ONLY', 'ENFORCE']);
+export type PolicyEngineMode = z.infer<typeof PolicyEngineModeSchema>;
+
+export const GatewayPolicyEngineConfigurationSchema = z
+  .object({
+    policyEngineName: z.string().min(1),
+    mode: PolicyEngineModeSchema,
+  })
+  .strict();
+export type GatewayPolicyEngineConfiguration = z.infer<typeof GatewayPolicyEngineConfigurationSchema>;
+
 // ============================================================================
 // Gateway
 // ============================================================================
@@ -614,6 +629,8 @@ export const AgentCoreGatewaySchema = z
     enableSemanticSearch: z.boolean().default(true),
     /** Exception verbosity level. 'NONE' = generic errors (default), 'DEBUG' = verbose errors. */
     exceptionLevel: GatewayExceptionLevelSchema.default('NONE'),
+    /** Policy engine configuration for Cedar-based authorization of tool calls. */
+    policyEngineConfiguration: GatewayPolicyEngineConfigurationSchema.optional(),
   })
   .strict()
   .refine(