Merge branch 'feature/agent-pentest-fixes' into 'develop'

Agent related pentest fixes

See merge request genaiic-reusable-assets/engagement-artifacts/genaiic-idp-accelerator!321

Author: rstrahan

lib/idp_common_pkg/idp_common/agents/analytics/agent.py

@@ -11,9 +11,9 @@
 
 import boto3
 import strands
-from strands.models import BedrockModel
 
 from ..common.config import load_result_format_description
+from ..common.strands_bedrock_model import create_strands_bedrock_model
 from .config import load_python_plot_generation_examples
 from .tools import CodeInterpreterTools, get_database_info, run_athena_query
 from .utils import register_code_interpreter_tools
@@ -138,7 +138,9 @@ def run_athena_query_with_config(
     # Get model ID from environment variable
     model_id = os.environ.get("DOCUMENT_ANALYSIS_AGENT_MODEL_ID")
 
-    bedrock_model = BedrockModel(model_id=model_id, boto_session=session)
+    bedrock_model = create_strands_bedrock_model(
+        model_id=model_id, boto_session=session
+    )
 
     # Create the Strands agent with tools and system prompt
     strands_agent = strands.Agent(

lib/idp_common_pkg/idp_common/agents/common/strands_bedrock_model.py

@@ -0,0 +1,43 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: MIT-0
+
+"""
+Helper function for creating BedrockModel instances with automatic guardrail support.
+"""
+
+import os
+
+from strands.models import BedrockModel
+
+
+def create_strands_bedrock_model(
+    model_id: str, boto_session=None, **kwargs
+) -> BedrockModel:
+    """
+    Create a BedrockModel with automatic guardrail configuration from environment.
+
+    Args:
+        model_id: The Bedrock model ID to use
+        boto_session: Optional boto3 session
+        **kwargs: Additional arguments to pass to BedrockModel
+
+    Returns:
+        BedrockModel instance with guardrails applied if configured
+    """
+    # Get guardrail configuration from environment if available
+    guardrail_env = os.environ.get("GUARDRAIL_ID_AND_VERSION", "")
+    if guardrail_env:
+        try:
+            guardrail_id, guardrail_version = guardrail_env.split(":")
+            if guardrail_id and guardrail_version:
+                kwargs.update(
+                    {
+                        "guardrail_id": guardrail_id,
+                        "guardrail_version": guardrail_version,
+                        "guardrail_trace": "enabled",
+                    }
+                )
+        except ValueError:
+            pass  # Invalid format, continue without guardrails
+
+    return BedrockModel(model_id=model_id, boto_session=boto_session, **kwargs)

lib/idp_common_pkg/idp_common/agents/external_mcp/agent.py

@@ -12,10 +12,10 @@
 import boto3
 from mcp.client.streamable_http import streamablehttp_client
 from strands import Agent
-from strands.models import BedrockModel
 from strands.tools.mcp import MCPClient
 
 from ..common.oauth_auth import get_cognito_bearer_token
+from ..common.strands_bedrock_model import create_strands_bedrock_model
 
 logger = logging.getLogger(__name__)
 
@@ -130,7 +130,9 @@ def create_external_mcp_agent(
                     raise Exception(error_msg)
 
             # Create Bedrock model
-            bedrock_model = BedrockModel(model_id=model_id, boto_session=session)
+            bedrock_model = create_strands_bedrock_model(
+                model_id=model_id, boto_session=session
+            )
 
             # Create system prompt
             system_prompt = f"""

lib/idp_common_pkg/idp_common/agents/orchestrator/agent.py

@@ -15,9 +15,9 @@
 
 import strands
 from strands import tool
-from strands.models import BedrockModel
 
 from ..common.config import load_result_format_description
+from ..common.strands_bedrock_model import create_strands_bedrock_model
 
 logger = logging.getLogger(__name__)
 
@@ -161,7 +161,7 @@ def tool_func(query: str) -> str:
     )
 
     # Create the orchestrator agent
-    model = BedrockModel(
+    model = create_strands_bedrock_model(
         model_id=model_id,
         session=session,
     )

lib/idp_common_pkg/idp_common/agents/sample_calculator/agent.py

@@ -7,9 +7,10 @@
 
 import boto3
 import strands
-from strands.models import BedrockModel
 from strands_tools import calculator
 
+from ..common.strands_bedrock_model import create_strands_bedrock_model
+
 logger = logging.getLogger(__name__)
 
 
@@ -31,7 +32,7 @@ def create_sample_calculator_agent(
     model_id = "us.anthropic.claude-3-7-sonnet-20250219-v1:0"
 
     # Create Bedrock model
-    model = BedrockModel(model_id=model_id, session=session)
+    model = create_strands_bedrock_model(model_id=model_id, session=session)
 
     # Create and return agent with calculator tool
     return strands.Agent(model=model, tools=[calculator])

src/lambda/agent_request_handler/index.py

@@ -76,18 +76,17 @@ def handler(event, context):
         if not query:
             error_msg = "Query parameter is required"
             logger.error(error_msg)
-            return {
-                "statusCode": 400,
-                "body": error_msg
-            }
+            raise Exception(error_msg)
+            
+        if len(query) > 100000:
+            error_msg = "Query exceeds maximum length of 100000 characters"
+            logger.error(f"{error_msg}. Query length: {len(query)}")
+            raise Exception(error_msg)
 
         if not agent_ids:
             error_msg = "At least one agent ID is required"
             logger.error(error_msg)
-            return {
-                "statusCode": 400,
-                "body": error_msg
-            }
+            raise Exception(error_msg)
 
         # Extract and validate user ID from the identity context
         identity = event.get("identity", {})
@@ -153,7 +152,18 @@ def handler(event, context):
             "body": error_msg
         }
     except Exception as e:
-        error_msg = f"Error processing request: {str(e)}"
+        # Check if this is a validation error that should propagate as GraphQL error
+        error_str = str(e)
+        if any(msg in error_str for msg in [
+            "Query parameter is required",
+            "Query exceeds maximum length",
+            "At least one agent ID is required"
+        ]):
+            # Re-raise validation errors so they become GraphQL errors
+            raise e
+        
+        # Handle other unexpected errors
+        error_msg = f"Error processing request: {error_str}"
         logger.error(error_msg)
         return {
             "statusCode": 500,

src/ui/src/components/document-agents-layout/AgentJobStatus.jsx

@@ -20,6 +20,17 @@ const getStatusIndicator = (status) => {
 };
 
 const AgentJobStatus = ({ jobId, status, error }) => {
+  // Show error even if there's no jobId (for validation errors)
+  if (error && !jobId) {
+    return (
+      <Box padding={{ vertical: 'xs' }}>
+        <div>
+          <strong>Error:</strong> {error}
+        </div>
+      </Box>
+    );
+  }
+
   if (!jobId) {
     return null;
   }

src/ui/src/components/document-agents-layout/DocumentsAgentsLayout.jsx

@@ -199,8 +199,23 @@ const DocumentsAgentsLayout = () => {
       }, 1000);
     } catch (err) {
       logger.error('Error submitting query:', err);
+      logger.error('Error structure:', JSON.stringify(err, null, 2));
+
+      let errorMessage = 'Failed to submit query';
+
+      // Extract error message from GraphQL error structure
+      if (err.errors && err.errors.length > 0 && err.errors[0].message) {
+        errorMessage = err.errors[0].message;
+      } else if (err.message) {
+        errorMessage = err.message;
+      } else if (err.data && err.data.errors && err.data.errors.length > 0 && err.data.errors[0].message) {
+        errorMessage = err.data.errors[0].message;
+      } else if (typeof err === 'string') {
+        errorMessage = err;
+      }
+
       updateAnalyticsState({
-        error: err.message || 'Failed to submit query',
+        error: errorMessage,
         jobStatus: 'FAILED',
       });
     } finally {

template.yaml

@@ -3608,6 +3608,7 @@ Resources:
                 - !Ref ReportingBucketName
           DOCUMENT_ANALYSIS_AGENT_MODEL_ID: !Ref DocumentAnalysisAgentModelId
           AWS_STACK_NAME: !Ref AWS::StackName
+          GUARDRAIL_ID_AND_VERSION: !If [HasGuardrailConfig, !Sub "${BedrockGuardrailId}:${BedrockGuardrailVersion}", ""]
       LoggingConfig:
         LogGroup: !Ref AgentProcessorLogGroup
       Policies:
@@ -3645,6 +3646,14 @@ Resources:
               Resource: 
                 - !Sub "arn:${AWS::Partition}:bedrock:*::foundation-model/*"
                 - !Sub "arn:${AWS::Partition}:bedrock:${AWS::Region}:${AWS::AccountId}:inference-profile/*"
+            - !If
+              - HasGuardrailConfig
+              - Effect: Allow
+                Action:
+                  - "bedrock:ApplyGuardrail"
+                Resource:
+                  - !Sub "arn:${AWS::Partition}:bedrock:${AWS::Region}:${AWS::AccountId}:guardrail/${BedrockGuardrailId}"
+              - !Ref AWS::NoValue
             - Effect: Allow
               Action:
                 - appsync:GraphQL
@@ -3905,7 +3914,16 @@ Resources:
           }
         }
       ResponseMappingTemplate: |-
-        true
+        #if($ctx.error)
+          $util.error($ctx.error.message, $ctx.error.type)
+        #end
+        
+        ## Return false if no item was updated (item not found)
+        #if(!$ctx.result)
+          false
+        #else
+          true
+        #end
 
   DeleteAgentJobResolver:
     Type: AWS::AppSync::Resolver