Best practices for applying Amazon Bedrock Guardrails to code generation workflows

import hashlib
from pathlib import Path

class CodeCommitGuardrail:
    """
    Validates all AI-generated code changes before they are persisted.
    Analogous to running security scanners in a Git pre-commit hook.

    Use this when:
    - A coding assistant writes code to disk
    - AI-generated changes are staged for commit
    - Generated code is about to be deployed
    """

    def __init__(self, client, guardrail_id, guardrail_version):
        self.client = client
        self.guardrail_id = guardrail_id
        self.guardrail_version = guardrail_version
        self.validated_hashes = {}  # Cache: don't re-validate unchanged files

    def validate_file_changes(self, staged_files: dict) -> dict:
        """
        Evaluate all staged AI-generated code changes before commit.

        Args:
            staged_files: Dict of {file_path: file_content} for all changed files

        Returns:
            Dict with 'safe' boolean and any violations found
        """
        violations = []
        skipped = []

        for file_path, content in staged_files.items():
            # Skip files that haven't changed since last validation
            content_hash = hashlib.sha256(content.encode()).hexdigest()
            if self.validated_hashes.get(file_path) == content_hash:
                skipped.append(file_path)
                continue

            # Evaluate in 1,000-char aligned chunks
            file_violations = self._evaluate_file(file_path, content)

            if file_violations:
                violations.extend(file_violations)
            else:
                # Cache successful validation
                self.validated_hashes[file_path] = content_hash

        return {
            'safe': len(violations) == 0,
            'violations': violations,
            'files_evaluated': len(staged_files) - len(skipped),
            'files_skipped_cached': len(skipped)
        }

    def _evaluate_file(self, file_path: str, content: str) -> list:
        """Evaluate a single file's content against guardrails."""
        violations = []

        for offset in range(0, len(content), 1000):
            chunk = content[offset:offset + 1000]

            response = self.client.apply_guardrail(
                guardrailIdentifier=self.guardrail_id,
                guardrailVersion=self.guardrail_version,
                source="OUTPUT",
                content=[{'text': {'text': chunk}}]
            )

            if response['action'] == 'GUARDRAIL_INTERVENED':
                violations.append({
                    'file': file_path,
                    'offset': offset,
                    'length': len(chunk),
                    'assessments': response['assessments'],
                    'snippet': chunk[:200] + '...' if len(chunk) > 200 else chunk
                })

        return violations


# Example: Integration with a coding assistant's file-write operation
def write_ai_generated_code(file_path: str, content: str, guardrail: CodeCommitGuardrail):
    """
    Wrapper for file writes that validates content before persisting.
    """
    result = guardrail.validate_file_changes({file_path: content})

    if not result['safe']:
        print(f"Guardrail blocked write to {file_path}")
        for v in result['violations']:
            print(f"  Violation at offset {v['offset']}: {v['assessments']}")
        return False

    # Safe to write
    Path(file_path).write_text(content)
    print(f"{file_path} validated and written successfully")
    return True

# When using InvokeModelWithResponseStream with guardrails
response = bedrock_runtime.invoke_model_with_response_stream(
    modelId='anthropic.claude-sonnet-4-20250514',
    body=json.dumps(request_body),
    guardrailIdentifier="your-guardrail-id",
    guardrailVersion='1',
    # Key configuration: increase streaming interval
    streamingConfigurations={
        'guardrailStreamingInterval': 1000  # Default is 50!
    }
)

import boto3
import json

bedrock_runtime = boto3.client('bedrock-runtime')

def process_coding_request(user_input: str, system_prompt: str, conversation_history: list):
    """
    Validate user input with guardrails, then invoke model without inline scanning.
    The system prompt and conversation history are NOT re-evaluated every turn.
    """

    # Step 1: Evaluate ONLY the new dynamic user input
    guardrail_response = bedrock_runtime.apply_guardrail(
        guardrailIdentifier="your-guardrail-id",
        guardrailVersion='1',
        source="INPUT",
        content=[
            {
                'text': {
                    'text': user_input  # Only the new content - not the full prompt
                }
            }
        ]
    )

    if guardrail_response['action'] == 'GUARDRAIL_INTERVENED':
        return {
            'blocked': True,
            'message': guardrail_response['outputs'][0]['text'],
            'assessments': guardrail_response['assessments']
        }

    # Step 2: Safe to proceed - invoke model WITHOUT inline guardrails
    # The system prompt and history bypass redundant evaluation
    messages = conversation_history + [{'role': 'user', 'content': user_input}]

    response = bedrock_runtime.converse(
        modelId='anthropic.claude-sonnet-4-20250514',
        messages=messages,
        system=[{'text': system_prompt}]
        # Note: No guardrailConfig here - we've already validated input
    )

    return {
        'blocked': False,
        'response': response['output']['message']['content'][0]['text']
    }

def generate_and_validate_code(user_input: str, system_prompt: str):
    """
    Generate code without inline guardrails, then validate the complete output.
    Ideal for detecting secrets, PII, or unsafe patterns in generated code.
    """

    # Step 1: Generate code without inline guardrail overhead
    response = bedrock_runtime.converse(
        modelId='anthropic.claude-sonnet-4-20250514',
        messages=[{'role': 'user', 'content': user_input}],
        system=[{'text': system_prompt}]
    )

    generated_code = response['output']['message']['content'][0]['text']

    # Step 2: Validate the COMPLETE code artifact - not intermediate chunks
    guardrail_response = bedrock_runtime.apply_guardrail(
        guardrailIdentifier="your-guardrail-id",
        guardrailVersion='1',
        source="OUTPUT",
        content=[
            {
                'text': {
                    'text': generated_code
                }
            }
        ]
    )

    if guardrail_response['action'] == 'GUARDRAIL_INTERVENED':
        return {
            'blocked': True,
            'message': 'Generated code contains sensitive content that was blocked.',
            'assessments': guardrail_response['assessments']
        }

    return {
        'blocked': False,
        'code': generated_code
    }

def secure_code_generation_pipeline(user_input: str, system_prompt: str, history: list):
    """
    Full bidirectional validation without inline scanning overhead.
    Input: Check for prompt injection/manipulation
    Output: Check for leaked secrets and unsafe patterns
    """

    # Step 1: Validate input for prompt injection
    input_check = bedrock_runtime.apply_guardrail(
        guardrailIdentifier="your-guardrail-id",
        guardrailVersion='1',
        source="INPUT",
        content=[{'text': {'text': user_input}}]
    )

    if input_check['action'] == 'GUARDRAIL_INTERVENED':
        return {'blocked': True, 'stage': 'input', 'reason': input_check['outputs'][0]['text']}

    # Step 2: Generate code - no inline guardrails needed
    messages = history + [{'role': 'user', 'content': user_input}]
    response = bedrock_runtime.converse(
        modelId='anthropic.claude-sonnet-4-20250514',
        messages=messages,
        system=[{'text': system_prompt}]
    )
    generated_code = response['output']['message']['content'][0]['text']

    # Step 3: Validate output for sensitive information
    output_check = bedrock_runtime.apply_guardrail(
        guardrailIdentifier="your-guardrail-id",
        guardrailVersion='1',
        source="OUTPUT",
        content=[{'text': {'text': generated_code}}]
    )

    if output_check['action'] == 'GUARDRAIL_INTERVENED':
        return {'blocked': True, 'stage': 'output', 'reason': output_check['outputs'][0]['text']}

    return {'blocked': False, 'code': generated_code}

import asyncio

class GuardrailBatcher:
    """Accumulates streaming output and evaluates at 1,000-char boundaries."""

    def __init__(self, client, guardrail_id, guardrail_version):
        self.client = client
        self.guardrail_id = guardrail_id
        self.guardrail_version = guardrail_version
        self.buffer = ""
        self.BATCH_SIZE = 1000  # Align to TextUnit boundary

    async def process_chunk(self, chunk: str):
        """Buffer chunks and evaluate when we hit a TextUnit boundary."""
        self.buffer += chunk

        while len(self.buffer) >= self.BATCH_SIZE:
            # Extract exactly 1,000 characters for evaluation
            batch = self.buffer[:self.BATCH_SIZE]
            self.buffer = self.buffer[self.BATCH_SIZE:]

            response = self.client.apply_guardrail(
                guardrailIdentifier=self.guardrail_id,
                guardrailVersion=self.guardrail_version,
                source="OUTPUT",
                content=[{'text': {'text': batch}}]
            )

            if response['action'] == 'GUARDRAIL_INTERVENED':
                raise GuardrailViolation(response)

    async def flush(self):
        """Evaluate any remaining content at end of stream."""
        if self.buffer:
            response = self.client.apply_guardrail(
                guardrailIdentifier=self.guardrail_id,
                guardrailVersion=self.guardrail_version,
                source="OUTPUT",
                content=[{'text': {'text': self.buffer}}]
            )
            self.buffer = ""
            return response

import re

class RiskBasedGuardrail:
    """
    Adjusts guardrail evaluation strategy based on the risk profile
    of the code being generated.

    High-risk code: Full evaluation with all safeguards
    Standard code: Lightweight evaluation (sensitive info only)
    Low-risk code: Skip intermediate evaluation; validate at commit only
    """

    # Patterns that indicate high-risk code
    HIGH_RISK_PATTERNS = [
        r'iam[_\.]',                           # IAM policy manipulation
        r'(access_key|secret_key|password)',   # Credential handling
        r'(exec|eval|subprocess|os\.system)',  # Code execution
        r'(BEGIN.*PRIVATE KEY)',               # Cryptographic keys
        r'(security_group|nacl|firewall)',     # Network security
        r'(grant|revoke|permission)',          # Authorization
        r'(encrypt|decrypt|kms)',              # Encryption operations
    ]

    # Patterns that indicate low-risk code
    LOW_RISK_PATTERNS = [
        r'(\.css|style|className|fontSize)',   # Styling
        r'(render|component|jsx|tsx)',         # UI components
        r'(test|spec|mock|fixture)',           # Test files
        r'(README|docs|comment)',              # Documentation
    ]

    def __init__(self, client, guardrail_id_full, guardrail_id_lightweight):
        self.client = client
        self.guardrail_id_full = guardrail_id_full                # All safeguards active
        self.guardrail_id_lightweight = guardrail_id_lightweight  # Sensitive info only

    def classify_risk(self, code: str, file_path: str = "") -> str:
        """Determine risk level of generated code."""
        combined = code + " " + file_path

        for pattern in self.HIGH_RISK_PATTERNS:
            if re.search(pattern, combined, re.IGNORECASE):
                return 'HIGH'

        for pattern in self.LOW_RISK_PATTERNS:
            if re.search(pattern, combined, re.IGNORECASE):
                return 'LOW'

        return 'STANDARD'

    def evaluate(self, code: str, file_path: str = "") -> dict:
        """
        Evaluate code with appropriate depth based on risk classification.
        """
        risk_level = self.classify_risk(code, file_path)

        if risk_level == 'LOW':
            # Low-risk: defer to commit-time validation
            return {
                'action': 'PASS',
                'risk_level': 'LOW',
                'evaluation': 'deferred_to_commit',
                'reason': 'Content classified as low-risk (UI/tests/docs)'
            }

        # Choose guardrail based on risk
        guardrail_id = (
            self.guardrail_id_full if risk_level == 'HIGH'
            else self.guardrail_id_lightweight
        )

        response = self.client.apply_guardrail(
            guardrailIdentifier=guardrail_id,
            guardrailVersion='1',
            source="OUTPUT",
            content=[{'text': {'text': code}}]
        )

        return {
            'action': response['action'],
            'risk_level': risk_level,
            'evaluation': 'full' if risk_level == 'HIGH' else 'lightweight',
            'assessments': response.get('assessments', [])
        }


# Example: Using risk-based evaluation in a coding pipeline
risk_guardrail = RiskBasedGuardrail(
    client=bedrock_runtime,
    guardrail_id_full="guardrail-all-safeguards",
    guardrail_id_lightweight="guardrail-sensitive-info-only"
)

# High-risk: IAM policy code → full evaluation (3 safeguards)
# WARNING: The following is an example of an INSECURE anti-pattern used here
# solely to demonstrate guardrail detection. Never use wildcard permissions
# (Action: '*', Resource: '*') in production. Always follow the principle of
# least privilege. See: https://docs.aws.amazon.com/IAM/latest/UserGuide/best-practices.html#grant-least-privilege
iam_code = """
policy = iam.create_policy(
    PolicyName="AdminAccess",
    PolicyDocument=json.dumps({
        'Statement': [{'Effect': 'Allow', 'Action': '*', 'Resource': '*'}]
    })
)
"""
result = risk_guardrail.evaluate(iam_code, "deploy/iam_policies.py")
# → risk_level: HIGH, evaluation: full

# Low-risk: React component → deferred to commit
ui_code = """
function Button({ label, onClick }) {
    return {label};
}
"""
result = risk_guardrail.evaluate(ui_code, "src/components/Button.tsx")
# → risk_level: LOW, evaluation: deferred_to_commit

class AgentPipelineGuardrail:
    """
    Guardrail architecture for multi-step agentic coding workflows.

    Key insight: In an agent loop, the model may take 5-10 reasoning steps
    (tool calls, observations, reflections) before producing final code.
    Evaluating every intermediate step is wasteful - most contain internal
    reasoning that never reaches the user or a file system.

    Strategy:
    - INPUT: Validate user request before agent starts
    - INTERMEDIATE: Skip chain-of-thought and tool-use reasoning
    - TOOL INPUTS: Validate when agent writes to file or executes code
    - FINAL OUTPUT: Validate the complete response shown to the user
    """

    def __init__(self, client, guardrail_id, guardrail_version):
        self.client = client
        self.guardrail_id = guardrail_id
        self.guardrail_version = guardrail_version

    def run_agent_with_guardrails(self, user_request: str, agent):
        """
        Execute an agentic coding workflow with strategic guardrail placement.
        """

        # Gate 1: Validate user input BEFORE agent starts working
        input_check = self._evaluate(user_request, 'INPUT')
        if input_check['action'] == 'GUARDRAIL_INTERVENED':
            return {'blocked': True, 'stage': 'input', 'response': input_check}

        # Run the agent - NO guardrail evaluation during reasoning steps
        agent_steps = []
        final_output = None

        for step in agent.run(user_request):
            if step['type'] == 'thinking':
                # Skip: internal reasoning doesn't need evaluation
                agent_steps.append(step)
                continue

            elif step['type'] == 'tool_call':
                # Gate 2: Validate ONLY dangerous tool inputs
                if self._is_dangerous_tool(step['tool_name']):
                    tool_check = self._evaluate(
                        json.dumps(step['tool_input']), 'OUTPUT'
                    )
                    if tool_check['action'] == 'GUARDRAIL_INTERVENED':
                        return {
                            'blocked': True,
                            'stage': f"tool_call:{step['tool_name']}",
                            'response': tool_check
                        }
                agent_steps.append(step)

            elif step['type'] == 'final_response':
                final_output = step['content']

        # Gate 3: Validate final output before showing to user
        if final_output:
            output_check = self._evaluate(final_output, 'OUTPUT')
            if output_check['action'] == 'GUARDRAIL_INTERVENED':
                return {'blocked': True, 'stage': 'output', 'response': output_check}

        return {
            'blocked': False,
            'output': final_output,
            'steps_taken': len(agent_steps),
            'steps_evaluated': sum(
                1 for s in agent_steps
                if s['type'] == 'tool_call' and self._is_dangerous_tool(s.get('tool_name', ''))
            )
        }

    def _is_dangerous_tool(self, tool_name: str) -> bool:
        """Tools that write to disk or execute code need guardrail evaluation."""
        dangerous_tools = {
            'write_file', 'execute_code', 'run_command',
            'create_file', 'edit_file', 'bash'
        }
        return tool_name in dangerous_tools

    def _evaluate(self, content: str, source: str) -> dict:
        return self.client.apply_guardrail(
            guardrailIdentifier=self.guardrail_id,
            guardrailVersion=self.guardrail_version,
            source=source,
            content=[{'text': {'text': content}}]
        )