set type of event given to resolve as unknown and use type assertion to verify it a bedrock event

Author: svozza

packages/event-handler/src/bedrock-agent-function/BedrockAgentFunctionResolver.ts

@@ -1,7 +1,8 @@
+// TODO: rename folder to just bedrock-agent
+
 import { EnvironmentVariablesService } from '@aws-lambda-powertools/commons';
 import type { Context } from 'aws-lambda';
 import type {
-  BedrockAgentFunctionEvent,
   BedrockAgentFunctionResponse,
   Configuration,
   ParameterValue,
@@ -11,11 +12,10 @@ import type {
   ToolFunction,
 } from '../types/bedrock-agent-function.js';
 import type { GenericLogger } from '../types/common.js';
-import { isPrimitive } from './utils.js';
+import { assertBedrockAgentFunctionEvent, isPrimitive } from './utils.js';
 
 export class BedrockAgentFunctionResolver {
-  readonly #tools: Map<string, Tool<Record<string, ParameterValue>>> =
-    new Map();
+  readonly #tools: Map<string, Tool> = new Map();
   readonly #envService: EnvironmentVariablesService;
   readonly #logger: Pick<GenericLogger, 'debug' | 'warn' | 'error'>;
 
@@ -124,7 +124,7 @@ export class BedrockAgentFunctionResolver {
     }
 
     this.#tools.set(name, {
-      handler: handler as ToolFunction<Record<string, ParameterValue>>,
+      handler: handler as ToolFunction,
       config,
     });
     this.#logger.debug(`Tool ${name} has been registered.`);
@@ -160,9 +160,11 @@ export class BedrockAgentFunctionResolver {
   }
 
   async resolve(
-    event: BedrockAgentFunctionEvent,
+    event: unknown,
     context: Context
   ): Promise<BedrockAgentFunctionResponse> {
+    assertBedrockAgentFunctionEvent(event);
+
     const {
       function: toolName,
       parameters = [],
@@ -205,7 +207,7 @@ export class BedrockAgentFunctionResolver {
 
     try {
       const res = (await tool.handler(toolParams, event, context)) ?? '';
-      const body = isPrimitive(res) ? String(res) : JSON.stringify(res);
+      const body = isPrimitive(res) ? String(res) : JSON.stringify(res); //TODO just use JSON.stringify
       return this.#buildResponse({
         actionGroup,
         function: toolName,

packages/event-handler/src/bedrock-agent-function/utils.ts

@@ -2,6 +2,8 @@ import type {
   JSONPrimitive,
   JSONValue,
 } from '@aws-lambda-powertools/commons/types';
+import { isRecord, isString } from '@aws-lambda-powertools/commons/typeutils';
+import type { BedrockAgentFunctionEvent } from '../types/bedrock-agent-function.js';
 
 export function isPrimitive(value: JSONValue): value is JSONPrimitive {
   return (
@@ -12,3 +14,59 @@ export function isPrimitive(value: JSONValue): value is JSONPrimitive {
     typeof value === 'boolean'
   );
 }
+
+/**
+ * Asserts that the provided event is a BedrockAgentFunctionEvent.
+ *
+ * @param event - The incoming event to check
+ * @throws Error if the event is not a valid BedrockAgentFunctionEvent
+ */
+export function assertBedrockAgentFunctionEvent(
+  event: unknown
+): asserts event is BedrockAgentFunctionEvent {
+  const isValid =
+    typeof event === 'object' &&
+    event !== null &&
+    isRecord(event) &&
+    'actionGroup' in event &&
+    isString(event.actionGroup) &&
+    'function' in event &&
+    isString(event.function) &&
+    (!('parameters' in event) ||
+      event.parameters == null ||
+      (Array.isArray(event.parameters) &&
+        event.parameters.every(
+          (param) =>
+            isRecord(param) &&
+            'name' in param &&
+            isString(param.name) &&
+            'type' in param &&
+            isString(param.type) &&
+            'value' in param &&
+            isString(param.value)
+        ))) &&
+    'messageVersion' in event &&
+    isString(event.messageVersion) &&
+    'agent' in event &&
+    isRecord(event.agent) &&
+    'name' in event.agent &&
+    isString(event.agent.name) &&
+    'id' in event.agent &&
+    isString(event.agent.id) &&
+    'alias' in event.agent &&
+    isString(event.agent.alias) &&
+    'version' in event.agent &&
+    isString(event.agent.version) &&
+    'inputText' in event &&
+    isString(event.inputText) &&
+    'sessionId' in event &&
+    isString(event.sessionId) &&
+    'sessionAttributes' in event &&
+    isRecord(event.sessionAttributes) &&
+    'promptSessionAttributes' in event &&
+    isRecord(event.promptSessionAttributes);
+
+  if (!isValid) {
+    throw new Error('Event is not a valid BedrockAgentFunctionEvent');
+  }
+}

packages/event-handler/tests/unit/bedrock-agent/BedrockAgentFunctionResolver.test.ts

@@ -34,6 +34,51 @@ describe('Class: BedrockAgentFunctionResolver', () => {
     vi.unstubAllEnvs();
   });
 
+  it.each([
+    {
+      name: 'null event',
+      invalidEvent: null,
+    },
+    {
+      name: 'missing required fields',
+      invalidEvent: {
+        function: 'test-tool',
+      },
+    },
+    {
+      name: 'invalid parameters structure',
+      invalidEvent: {
+        function: 'test-tool',
+        actionGroup: 'testGroup',
+        messageVersion: '1.0',
+        agent: {
+          name: 'agentName',
+          id: 'agentId',
+          alias: 'agentAlias',
+          version: '1',
+        },
+        inputText: 'test input',
+        sessionId: 'session123',
+        parameters: 'not an array',
+        sessionAttributes: {},
+        promptSessionAttributes: {},
+      },
+    },
+  ])('throws when given an invalid event: $name', async ({ invalidEvent }) => {
+    // Prepare
+    const app = new BedrockAgentFunctionResolver();
+
+    app.tool(async () => 'test', {
+      name: 'test-tool',
+      description: 'Test tool',
+    });
+
+    // Act & Assert
+    await expect(app.resolve(invalidEvent, context)).rejects.toThrow(
+      'Event is not a valid BedrockAgentFunctionEvent'
+    );
+  });
+
   it('uses a default logger with only warnings if none is provided', () => {
     // Prepare
     const app = new BedrockAgentFunctionResolver();