Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
37 changes: 36 additions & 1 deletion cdk/src/constructs/task-orchestrator.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand Down Expand Up @@ -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;
}

/**
Expand All @@ -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;

Expand Down Expand Up @@ -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/*'],
Expand Down Expand Up @@ -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;
Expand Down
4 changes: 2 additions & 2 deletions cdk/src/handlers/orchestrate-task.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
});
Expand Down
101 changes: 97 additions & 4 deletions cdk/src/handlers/shared/context-hydration.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand Down Expand Up @@ -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;
}
Expand All @@ -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)
Expand Down Expand Up @@ -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[] = [];
Expand Down Expand Up @@ -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,
Expand All @@ -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
Expand All @@ -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),
Expand Down
16 changes: 12 additions & 4 deletions cdk/src/handlers/shared/create-task-core.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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');
Expand Down Expand Up @@ -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({
Expand All @@ -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);
}
}

Expand Down
20 changes: 20 additions & 0 deletions cdk/src/handlers/shared/orchestrator.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
2 changes: 2 additions & 0 deletions cdk/src/stacks/agent.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Loading
Loading