feat: Add support for map parameter type (#78)

* add map parameter to protocol

* try and add tests

* fix param name and type

* lint

* run tests on new version of tools

* lint

* test on latest toolbox version

* fix protocol

* rename additional properties

* fix e2e tests

* clean comment

* fix test

* remove wrong test

* add unit tests

* improve test coverage

* lint

* roll back accidental changes

* Made check more robust

* refactor to avoid code duplication

* disable linter for a particular line

* Update toolbox version

* de capitalise additional properties

* use zodv3 error messages

* clean up param file

* lint

* test fix

* small fixes

* lint

* remove description from additional properties

Author: twishabansal

packages/toolbox-core/integration.cloudbuild.yaml

@@ -42,4 +42,4 @@ options:
   logging: CLOUD_LOGGING_ONLY
 substitutions:
   _VERSION: '22.16.0'
-  _TOOLBOX_VERSION: '0.10.0'
+  _TOOLBOX_VERSION: '0.12.0'

packages/toolbox-core/src/toolbox_core/protocol.ts

@@ -14,72 +14,72 @@
 
 import {z, ZodRawShape, ZodTypeAny, ZodObject} from 'zod';
 
-// Define All Interfaces
-
-interface BaseParameter {
-  name: string;
-  description: string;
-  authSources?: string[];
-  required?: boolean;
-}
-
-interface StringParameter extends BaseParameter {
+// Type Definitions
+interface StringType {
   type: 'string';
 }
-
-interface IntegerParameter extends BaseParameter {
+interface IntegerType {
   type: 'integer';
 }
-
-interface FloatParameter extends BaseParameter {
+interface FloatType {
   type: 'float';
 }
-
-interface BooleanParameter extends BaseParameter {
+interface BooleanType {
   type: 'boolean';
 }
-
-interface ArrayParameter extends BaseParameter {
+interface ArrayType {
   type: 'array';
-  items: ParameterSchema; // Recursive reference to the ParameterSchema type
+  items: TypeSchema; // Recursive
+}
+interface ObjectType {
+  type: 'object';
+  additionalProperties?: boolean | TypeSchema; // Recursive
 }
 
-export type ParameterSchema =
-  | StringParameter
-  | IntegerParameter
-  | FloatParameter
-  | BooleanParameter
-  | ArrayParameter;
+// Union of all pure type definitions.
+export type TypeSchema =
+  | StringType
+  | IntegerType
+  | FloatType
+  | BooleanType
+  | ArrayType
+  | ObjectType;
 
-// Get all Zod schema types
+// The base properties of a named parameter.
+interface BaseParameter {
+  name: string;
+  description: string;
+  authSources?: string[];
+  required?: boolean;
+}
+
+export type ParameterSchema = BaseParameter & TypeSchema;
 
+// Zod schema for the pure type definitions. This must be lazy for recursion.
+const ZodTypeSchema: z.ZodType<TypeSchema> = z.lazy(() =>
+  z.discriminatedUnion('type', [
+    z.object({type: z.literal('string')}),
+    z.object({type: z.literal('integer')}),
+    z.object({type: z.literal('float')}),
+    z.object({type: z.literal('boolean')}),
+    z.object({type: z.literal('array'), items: ZodTypeSchema}),
+    z.object({
+      type: z.literal('object'),
+      additionalProperties: z.union([z.boolean(), ZodTypeSchema]).optional(),
+    }),
+  ]),
+);
+
+// Zod schema for the base properties.
 const ZodBaseParameter = z.object({
   name: z.string().min(1, 'Parameter name cannot be empty'),
   description: z.string(),
   authSources: z.array(z.string()).optional(),
   required: z.boolean().optional(),
 });
 
-export const ZodParameterSchema = z.lazy(() =>
-  z.discriminatedUnion('type', [
-    ZodBaseParameter.extend({
-      type: z.literal('string'),
-    }),
-    ZodBaseParameter.extend({
-      type: z.literal('integer'),
-    }),
-    ZodBaseParameter.extend({
-      type: z.literal('float'),
-    }),
-    ZodBaseParameter.extend({
-      type: z.literal('boolean'),
-    }),
-    ZodBaseParameter.extend({
-      type: z.literal('array'),
-      items: ZodParameterSchema, // Recursive reference for the item's definition
-    }),
-  ]),
-) as z.ZodType<ParameterSchema>;
+export const ZodParameterSchema: z.ZodType<ParameterSchema> =
+  ZodBaseParameter.and(ZodTypeSchema);
 
 export const ZodToolSchema = z.object({
   description: z.string().min(1, 'Tool description cannot be empty'),
@@ -97,52 +97,44 @@ export const ZodManifestSchema = z.object({
 
 export type ZodManifest = z.infer<typeof ZodManifestSchema>;
 
-/**
- * Recursively builds a Zod schema for a single parameter based on its TypeScript definition.
- * @param param The ParameterSchema (TypeScript type) to convert.
- * @returns A ZodTypeAny representing the schema for this parameter.
- */
-function buildZodShapeFromParam(param: ParameterSchema): ZodTypeAny {
-  let schema: ZodTypeAny;
-  switch (param.type) {
+function buildZodShapeFromTypeSchema(typeSchema: TypeSchema): ZodTypeAny {
+  switch (typeSchema.type) {
     case 'string':
-      schema = z.string();
-      break;
+      return z.string();
     case 'integer':
-      schema = z.number().int();
-      break;
+      return z.number().int();
     case 'float':
-      schema = z.number();
-      break;
+      return z.number();
     case 'boolean':
-      schema = z.boolean();
-      break;
+      return z.boolean();
     case 'array':
-      // Recursively build the schema for array items
-      // Array items inherit the 'required' status of the parent array.
-      param.items.required = param.required;
-      schema = z.array(buildZodShapeFromParam(param.items));
-      break;
+      return z.array(buildZodShapeFromTypeSchema(typeSchema.items));
+    case 'object':
+      if (typeof typeSchema.additionalProperties === 'object') {
+        return z.record(
+          z.string(),
+          buildZodShapeFromTypeSchema(typeSchema.additionalProperties),
+        );
+      } else if (typeSchema.additionalProperties === false) {
+        return z.object({});
+      } else {
+        return z.record(z.string(), z.any());
+      }
     default: {
-      // This ensures exhaustiveness at compile time if ParameterSchema is a discriminated union
-      const _exhaustiveCheck: never = param;
+      const _exhaustiveCheck: never = typeSchema;
       throw new Error(`Unknown parameter type: ${_exhaustiveCheck['type']}`);
     }
   }
+}
 
+function buildZodShapeFromParam(param: ParameterSchema): ZodTypeAny {
+  const schema = buildZodShapeFromTypeSchema(param);
   if (param.required === false) {
     return schema.nullish();
   }
-
   return schema;
 }
 
-/**
- * Creates a ZodObject schema from an array of ParameterSchema (TypeScript types).
- * This combined schema is used by ToolboxTool to validate its call arguments.
- * @param params Array of ParameterSchema objects.
- * @returns A ZodObject schema.
- */
 export function createZodSchemaFromParams(
   params: ParameterSchema[],
 ): ZodObject<ZodRawShape> {

packages/toolbox-core/test/e2e/jest.globalSetup.ts

@@ -41,6 +41,7 @@ export default async function globalSetup(): Promise<void> {
     const toolsManifest = await accessSecretVersion(
       projectId,
       'sdk_testing_tools',
+      '34',
     );
     const toolsFilePath = await createTmpFile(toolsManifest);
     (globalThis as CustomGlobal).__TOOLS_FILE_PATH__ = toolsFilePath;

packages/toolbox-core/test/e2e/test.e2e.ts

@@ -116,6 +116,7 @@ describe('ToolboxClient E2E Tests', () => {
         'get-row-by-id',
         'get-n-rows',
         'search-rows',
+        'process-data',
       ]);
       expect(loadedToolNames).toEqual(expectedDefaultTools);
 
@@ -493,4 +494,80 @@ describe('ToolboxClient E2E Tests', () => {
       expect(response).toBe('null');
     });
   });
+  describe('Map/Object Params E2E Tests', () => {
+    let processDataTool: ReturnType<typeof ToolboxTool>;
+
+    beforeAll(async () => {
+      processDataTool = await commonToolboxClient.loadTool('process-data');
+    });
+
+    it('should correctly identify map/object parameters in the schema', () => {
+      const paramSchema = processDataTool.getParamSchema();
+      const baseArgs = {
+        execution_context: {env: 'prod'},
+        user_scores: {user1: 100},
+      };
+
+      // Test required untyped map (dict[str, Any])
+      expect(paramSchema.safeParse(baseArgs).success).toBe(true);
+      const argsWithoutExec = {...baseArgs};
+      delete (argsWithoutExec as Partial<typeof argsWithoutExec>)
+        .execution_context;
+      expect(paramSchema.safeParse(argsWithoutExec).success).toBe(false);
+
+      // Test required typed map (dict[str, int])
+      expect(
+        paramSchema.safeParse({
+          ...baseArgs,
+          user_scores: {user1: 'not-a-number'},
+        }).success,
+      ).toBe(false);
+
+      // Test optional typed map (dict[str, bool])
+      expect(
+        paramSchema.safeParse({
+          ...baseArgs,
+          feature_flags: {new_feature: true},
+        }).success,
+      ).toBe(true);
+      expect(
+        paramSchema.safeParse({...baseArgs, feature_flags: null}).success,
+      ).toBe(true);
+      expect(paramSchema.safeParse(baseArgs).success).toBe(true); // Omitted
+    });
+
+    it('should run tool with valid map parameters', async () => {
+      const response = await processDataTool({
+        execution_context: {env: 'prod', id: 1234, user: 1234.5},
+        user_scores: {user1: 100, user2: 200},
+        feature_flags: {new_feature: true},
+      });
+      expect(typeof response).toBe('string');
+      expect(response).toContain(
+        '"execution_context":{"env":"prod","id":1234,"user":1234.5}',
+      );
+      expect(response).toContain('"user_scores":{"user1":100,"user2":200}');
+      expect(response).toContain('"feature_flags":{"new_feature":true}');
+    });
+
+    it('should run tool with optional map param omitted', async () => {
+      const response = await processDataTool({
+        execution_context: {env: 'dev'},
+        user_scores: {user3: 300},
+      });
+      expect(typeof response).toBe('string');
+      expect(response).toContain('"execution_context":{"env":"dev"}');
+      expect(response).toContain('"user_scores":{"user3":300}');
+      expect(response).toContain('"feature_flags":null');
+    });
+
+    it('should fail when a map parameter has the wrong value type', async () => {
+      await expect(
+        processDataTool({
+          execution_context: {env: 'staging'},
+          user_scores: {user4: 'not-an-integer'},
+        }),
+      ).rejects.toThrow(/user_scores\.user4: Expected number, received string/);
+    });
+  });
 });