Merge pull request #13 from aws-samples/guardrails

chore(guardrails): add more guardrails for input screening

Author: krokoko

README.md

@@ -59,10 +59,10 @@ ABCA is under active development. The platform ships iteratively — each iterat
 | **3a** | Done | Repo onboarding, per-repo GitHub App credentials, turn caps, prompt guide |
 | **3b** | Done | Memory Tier 1, insights, agent self-feedback, prompt versioning, commit attribution |
 | **3bis** | Done | Hardening — reconciler error tracking, error serialization, test coverage gaps |
-| **3c** | WIP | Pre-flight checks, persistent session storage, deterministic validation, PR review task type, multi-modal input |
+| **3c** | WIP | Pre-flight checks, persistent session storage, deterministic validation, PR review task type, multi-modal input, input guardrail screening |
 | **3d** | Planned | Review feedback loop, PR outcome tracking, evaluation pipeline |
 | **4** | Planned | GitLab, visual proof, Slack, control panel, WebSocket streaming |
-| **5** | Planned | Pre-warming, multi-user/team, cost management, guardrails, alternate runtime |
+| **5** | Planned | Pre-warming, multi-user/team, cost management, output guardrails, alternate runtime |
 | **6** | Planned | Skills learning, multi-repo, iterative feedback, multiplayer, CDK constructs |
 
 See the full [ROADMAP](./docs/guides/ROADMAP.md) for details on each iteration.

cdk/src/constructs/task-orchestrator.ts

@@ -18,7 +18,7 @@
  */
 
 import * as path from 'path';
-import { Duration } from 'aws-cdk-lib';
+import { Duration, Stack } from 'aws-cdk-lib';
 import * as cloudwatch from 'aws-cdk-lib/aws-cloudwatch';
 import * as dynamodb from 'aws-cdk-lib/aws-dynamodb';
 import * as iam from 'aws-cdk-lib/aws-iam';
@@ -100,6 +100,18 @@ export interface TaskOrchestratorProps {
    * and writes fallback episodes during finalization.
    */
   readonly memoryId?: string;
+
+  /**
+   * Bedrock Guardrail ID used by the orchestrator to screen assembled PR prompts
+   * for prompt injection during context hydration. The same guardrail is also
+   * used by the Task API for submission-time task description screening.
+   */
+  readonly guardrailId?: string;
+
+  /**
+   * Bedrock Guardrail version. Required when guardrailId is provided.
+   */
+  readonly guardrailVersion?: string;
 }
 
 /**
@@ -125,6 +137,13 @@ export class TaskOrchestrator extends Construct {
   constructor(scope: Construct, id: string, props: TaskOrchestratorProps) {
     super(scope, id);
 
+    if (props.guardrailId && !props.guardrailVersion) {
+      throw new Error('guardrailVersion is required when guardrailId is provided');
+    }
+    if (!props.guardrailId && props.guardrailVersion) {
+      throw new Error('guardrailId is required when guardrailVersion is provided');
+    }
+
     const handlersDir = path.join(__dirname, '..', 'handlers');
     const maxConcurrent = props.maxConcurrentTasksPerUser ?? 3;
 
@@ -152,6 +171,8 @@ export class TaskOrchestrator extends Construct {
           USER_PROMPT_TOKEN_BUDGET: String(props.userPromptTokenBudget),
         }),
         ...(props.memoryId && { MEMORY_ID: props.memoryId }),
+        ...(props.guardrailId && { GUARDRAIL_ID: props.guardrailId }),
+        ...(props.guardrailVersion && { GUARDRAIL_VERSION: props.guardrailVersion }),
       },
       bundling: {
         externalModules: ['@aws-sdk/*'],
@@ -200,6 +221,20 @@ export class TaskOrchestrator extends Construct {
       secret.grantRead(this.fn);
     }
 
+    // Bedrock Guardrail permissions
+    if (props.guardrailId) {
+      this.fn.addToRolePolicy(new iam.PolicyStatement({
+        actions: ['bedrock:ApplyGuardrail'],
+        resources: [
+          Stack.of(this).formatArn({
+            service: 'bedrock',
+            resource: 'guardrail',
+            resourceName: props.guardrailId,
+          }),
+        ],
+      }));
+    }
+
     // Create alias for durable function invocation
     const fnAlias = this.fn.currentVersion.addAlias('live');
     this.alias = fnAlias;

cdk/src/handlers/orchestrate-task.ts

@@ -110,8 +110,8 @@ const durableHandler: DurableExecutionHandler<OrchestrateTaskEvent, void> = asyn
     try {
       return await hydrateAndTransition(task, blueprintConfig);
     } catch (err) {
-      // Transition may fail if task was externally cancelled — release concurrency
-      await failTask(taskId, task.status, `Hydration failed: ${String(err)}`, task.user_id, true);
+      // Hydration may fail due to external cancellation, guardrail blocking, or guardrail API failure — fail the task and release concurrency
+      await failTask(taskId, TaskStatus.HYDRATING, `Hydration failed: ${String(err)}`, task.user_id, true);
       throw err;
     }
   });

cdk/src/handlers/shared/context-hydration.ts

@@ -17,6 +17,7 @@
  *  SOFTWARE.
  */
 
+import { ApplyGuardrailCommand, BedrockRuntimeClient } from '@aws-sdk/client-bedrock-runtime';
 import { GetSecretValueCommand, SecretsManagerClient } from '@aws-sdk/client-secrets-manager';
 import { logger } from './logger';
 import { loadMemoryContext, type MemoryContext } from './memory';
@@ -85,6 +86,7 @@ export interface HydratedContext {
   readonly token_estimate: number;
   readonly truncated: boolean;
   readonly fallback_error?: string;
+  readonly guardrail_blocked?: string;
   readonly resolved_branch_name?: string;
   readonly resolved_base_branch?: string;
 }
@@ -96,6 +98,81 @@ export interface HydratedContext {
 const GITHUB_TOKEN_SECRET_ARN = process.env.GITHUB_TOKEN_SECRET_ARN;
 const USER_PROMPT_TOKEN_BUDGET = Number(process.env.USER_PROMPT_TOKEN_BUDGET ?? '100000');
 const GITHUB_API_TIMEOUT_MS = 30_000;
+const GUARDRAIL_ID = process.env.GUARDRAIL_ID;
+const GUARDRAIL_VERSION = process.env.GUARDRAIL_VERSION;
+const bedrockClient = (GUARDRAIL_ID && GUARDRAIL_VERSION) ? new BedrockRuntimeClient({}) : undefined;
+if (GUARDRAIL_ID && !GUARDRAIL_VERSION) {
+  logger.error('GUARDRAIL_ID is set but GUARDRAIL_VERSION is missing — guardrail screening disabled', {
+    metric_type: 'guardrail_misconfiguration',
+  });
+}
+
+// ---------------------------------------------------------------------------
+// Bedrock Guardrail screening
+// ---------------------------------------------------------------------------
+
+/**
+ * Error thrown when the Bedrock Guardrail API call fails. Distinguished from
+ * other errors so the outer catch in hydrateContext can re-throw it instead of
+ * falling back to unscreened content (fail-closed).
+ */
+export class GuardrailScreeningError extends Error {
+  constructor(message: string, cause?: Error) {
+    super(message, cause ? { cause } : undefined);
+    this.name = 'GuardrailScreeningError';
+  }
+}
+
+/**
+ * Screen text through the Bedrock Guardrail for prompt injection detection.
+ * Fail-closed: throws on Bedrock errors so unscreened content never reaches the agent.
+ * @param text - the text to screen.
+ * @param taskId - the task ID (for logging).
+ * @returns 'GUARDRAIL_INTERVENED' if blocked, 'NONE' if allowed, undefined when guardrail is
+ *          not configured (env vars missing).
+ * @throws GuardrailScreeningError when the Bedrock Guardrail API call fails (fail-closed).
+ */
+export async function screenWithGuardrail(text: string, taskId: string): Promise<'GUARDRAIL_INTERVENED' | 'NONE' | undefined> {
+  if (!bedrockClient || !GUARDRAIL_ID || !GUARDRAIL_VERSION) {
+    logger.info('Guardrail screening skipped — guardrail not configured', {
+      task_id: taskId,
+      metric_type: 'guardrail_screening_skipped',
+    });
+    return undefined;
+  }
+
+  try {
+    const result = await bedrockClient.send(new ApplyGuardrailCommand({
+      guardrailIdentifier: GUARDRAIL_ID,
+      guardrailVersion: GUARDRAIL_VERSION,
+      source: 'INPUT',
+      content: [{ text: { text } }],
+    }));
+
+    if (result.action === 'GUARDRAIL_INTERVENED') {
+      logger.warn('Content blocked by guardrail', {
+        task_id: taskId,
+        guardrail_id: GUARDRAIL_ID,
+        guardrail_version: GUARDRAIL_VERSION,
+      });
+      return 'GUARDRAIL_INTERVENED';
+    }
+
+    return 'NONE';
+  } catch (err) {
+    logger.error('Guardrail screening failed (fail-closed)', {
+      task_id: taskId,
+      guardrail_id: GUARDRAIL_ID,
+      error: err instanceof Error ? err.message : String(err),
+      error_name: err instanceof Error ? err.name : undefined,
+      metric_type: 'guardrail_screening_failure',
+    });
+    throw new GuardrailScreeningError(
+      `Guardrail screening unavailable: ${err instanceof Error ? err.message : String(err)}`,
+      err instanceof Error ? err : undefined,
+    );
+  }
+}
 
 // ---------------------------------------------------------------------------
 // GitHub token resolution (Secrets Manager with caching)
@@ -715,11 +792,15 @@ export interface HydrateContextOptions {
 }
 
 /**
- * Hydrate context for a task: resolve GitHub token, fetch issue, enforce
- * token budget, and assemble the user prompt.
+ * Hydrate context for a task: resolve GitHub token, fetch issue/PR, enforce
+ * token budget, assemble the user prompt, and (for PR tasks) screen through
+ * Bedrock Guardrail for prompt injection.
  * @param task - the task record from DynamoDB.
  * @param options - optional per-repo overrides.
- * @returns the hydrated context.
+ * @returns the hydrated context. For PR tasks, `guardrail_blocked` is set when
+ *          the guardrail intervened.
+ * @throws GuardrailScreeningError when the Bedrock Guardrail API call fails
+ *         (fail-closed — propagated to prevent unscreened content from reaching the agent).
  */
 export async function hydrateContext(task: TaskRecord, options?: HydrateContextOptions): Promise<HydratedContext> {
   const sources: string[] = [];
@@ -889,7 +970,10 @@ export async function hydrateContext(task: TaskRecord, options?: HydrateContextO
       resolvedBranchName = prResult.head_ref;
       resolvedBaseBranch = prResult.base_ref;
 
-      return {
+      // Screen assembled PR prompt through Bedrock Guardrail for prompt injection
+      const guardrailAction = await screenWithGuardrail(userPrompt, task.task_id);
+
+      const prContext: HydratedContext = {
         version: 1,
         user_prompt: userPrompt,
         memory_context: memoryContext,
@@ -898,7 +982,12 @@ export async function hydrateContext(task: TaskRecord, options?: HydrateContextO
         sources,
         token_estimate: estimateTokens(userPrompt),
         truncated,
+        ...(guardrailAction === 'GUARDRAIL_INTERVENED' && {
+          guardrail_blocked: 'PR context blocked by content policy',
+        }),
       };
+
+      return prContext;
     }
 
     // Standard task: existing behavior
@@ -918,6 +1007,10 @@ export async function hydrateContext(task: TaskRecord, options?: HydrateContextO
       truncated: budgetResult.truncated,
     };
   } catch (err) {
+    // Guardrail failures must propagate (fail-closed) — unscreened content must not reach the agent
+    if (err instanceof GuardrailScreeningError) {
+      throw err;
+    }
     // Fallback: minimal context from task_description only
     logger.error('Unexpected error during context hydration', {
       task_id: task.task_id, error: err instanceof Error ? err.message : String(err),

cdk/src/handlers/shared/create-task-core.ts

@@ -46,7 +46,13 @@ export interface TaskCreationContext {
 
 const ddb = DynamoDBDocumentClient.from(new DynamoDBClient({}));
 const lambdaClient = process.env.ORCHESTRATOR_FUNCTION_ARN ? new LambdaClient({}) : undefined;
-const bedrockClient = process.env.GUARDRAIL_ID ? new BedrockRuntimeClient({}) : undefined;
+const bedrockClient = (process.env.GUARDRAIL_ID && process.env.GUARDRAIL_VERSION)
+  ? new BedrockRuntimeClient({}) : undefined;
+if (process.env.GUARDRAIL_ID && !process.env.GUARDRAIL_VERSION) {
+  logger.error('GUARDRAIL_ID is set but GUARDRAIL_VERSION is missing — guardrail screening disabled', {
+    metric_type: 'guardrail_misconfiguration',
+  });
+}
 const TABLE_NAME = process.env.TASK_TABLE_NAME!;
 const EVENTS_TABLE_NAME = process.env.TASK_EVENTS_TABLE_NAME!;
 const TASK_RETENTION_DAYS = Number(process.env.TASK_RETENTION_DAYS ?? '90');
@@ -117,8 +123,8 @@ export async function createTaskCore(
   }
   const userMaxBudgetUsd = maxBudgetResult;
 
-  // 2. Screen task description with Bedrock Guardrail (fail-open: a Bedrock outage
-  //    should not block all task submissions — log the error and proceed)
+  // 2. Screen task description with Bedrock Guardrail (fail-closed: unscreened content
+  //    must not reach the agent — a Bedrock outage blocks task submissions)
   if (bedrockClient && body.task_description) {
     try {
       const guardrailResult = await bedrockClient.send(new ApplyGuardrailCommand({
@@ -133,11 +139,13 @@ export async function createTaskCore(
         return errorResponse(400, ErrorCode.VALIDATION_ERROR, 'Task description was blocked by content policy.', requestId);
       }
     } catch (guardrailErr) {
-      logger.error('Guardrail screening failed — proceeding without screening (fail-open)', {
+      logger.error('Guardrail screening failed (fail-closed)', {
         error: String(guardrailErr),
         user_id: context.userId,
         request_id: requestId,
+        metric_type: 'guardrail_screening_failure',
       });
+      return errorResponse(503, ErrorCode.INTERNAL_ERROR, 'Content screening is temporarily unavailable. Please try again later.', requestId);
     }
   }
 

cdk/src/handlers/shared/orchestrator.ts

@@ -252,6 +252,26 @@ export async function hydrateAndTransition(task: TaskRecord, blueprintConfig?: B
     memoryId: MEMORY_ID,
   });
 
+  // If guardrail screening blocked the hydrated context, emit audit event and throw
+  // to trigger task failure (the caller in orchestrate-task.ts catches and transitions to FAILED)
+  if (hydratedContext.guardrail_blocked) {
+    try {
+      await emitTaskEvent(task.task_id, 'guardrail_blocked', {
+        reason: hydratedContext.guardrail_blocked,
+        task_type: task.task_type,
+        pr_number: task.pr_number,
+        sources: hydratedContext.sources,
+        token_estimate: hydratedContext.token_estimate,
+      });
+    } catch (eventErr) {
+      logger.error('Failed to emit guardrail_blocked event', {
+        task_id: task.task_id,
+        error: eventErr instanceof Error ? eventErr.message : String(eventErr),
+      });
+    }
+    throw new Error(`Guardrail blocked: ${hydratedContext.guardrail_blocked}`);
+  }
+
   // For PR iteration: resolve actual branch name from PR head_ref
   if (hydratedContext.resolved_branch_name) {
     try {

cdk/src/stacks/agent.ts

@@ -282,6 +282,8 @@ export class AgentStack extends Stack {
       runtimeArn: runtime.agentRuntimeArn,
       githubTokenSecretArn: githubTokenSecret.secretArn,
       memoryId: agentMemory.memory.memoryId,
+      guardrailId: inputGuardrail.guardrailId,
+      guardrailVersion: inputGuardrail.guardrailVersion,
     });
 
     // Grant the orchestrator Lambda read+write access to memory