feat: Add support for Map parameter type (#51)

* feat: Add support for Map parameter type

* seperate validation of parameter schema from validation of input

* fix tests

* fix tests

* fix tests

* fix validate functino

* fix validate functino

* fix validate function

* fix unit tests

* fix unit tests

* set version to 34 in tbgenkit tests

* fix tests

* fix tests

* add AdditionalProperties in inputSchema

* skip genkit test

* update toolbox version to 0.12.0

* minor changes

* minor changes

* minor changes

Author: dishaprakash

core/client.go

@@ -99,6 +99,19 @@ func (tc *ToolboxClient) newToolboxTool(
 
 	// Iterate over the tool's parameters from the schema to categorize them.
 	for _, p := range schema.Parameters {
+
+		if ap, ok := p.AdditionalProperties.(map[string]any); ok {
+			apParam, err := mapToSchema(ap)
+			if err != nil {
+				return nil, nil, nil, err
+			}
+			p.AdditionalProperties = apParam
+		}
+		// Validate parameter schema
+		if err := p.ValidateDefinition(); err != nil {
+			// Return a detailed error indicating which tool failed validation.
+			return nil, nil, nil, fmt.Errorf("invalid schema for tool '%s': %w", name, err)
+		}
 		paramSchema[p.Name] = struct{}{}
 
 		if len(p.AuthSources) > 0 {

core/e2e_setup_test.go

@@ -43,15 +43,15 @@ func getEnvVar(key string) string {
 	return value
 }
 
-func accessSecretVersion(ctx context.Context, projectID, secretID string) string {
+func accessSecretVersion(ctx context.Context, projectID, secretID string, version string) string {
 	client, err := secretmanager.NewClient(ctx)
 	if err != nil {
 		log.Fatalf("Failed to create secretmanager client: %v", err)
 	}
 	defer client.Close()
 
 	req := &secretmanagerpb.AccessSecretVersionRequest{
-		Name: fmt.Sprintf("projects/%s/secrets/%s/versions/latest", projectID, secretID),
+		Name: fmt.Sprintf("projects/%s/secrets/%s/versions/%s", projectID, secretID, version),
 	}
 
 	result, err := client.AccessSecretVersion(ctx, req)

core/e2e_test.go

@@ -43,9 +43,9 @@ func TestMain(m *testing.M) {
 
 	// Get secrets and auth tokens
 	log.Println("Fetching secrets and auth tokens...")
-	toolsManifestContent := accessSecretVersion(ctx, projectID, "sdk_testing_tools")
-	clientID1 := accessSecretVersion(ctx, projectID, "sdk_testing_client1")
-	clientID2 := accessSecretVersion(ctx, projectID, "sdk_testing_client2")
+	toolsManifestContent := accessSecretVersion(ctx, projectID, "sdk_testing_tools", "34")
+	clientID1 := accessSecretVersion(ctx, projectID, "sdk_testing_client1", "latest")
+	clientID2 := accessSecretVersion(ctx, projectID, "sdk_testing_client2", "latest")
 	authToken1 = getAuthToken(ctx, clientID1)
 	authToken2 = getAuthToken(ctx, clientID2)
 
@@ -78,7 +78,6 @@ func TestMain(m *testing.M) {
 	os.Exit(exitCode)
 }
 
-// TestE2E_Basic maps to the TestBasicE2E class
 func TestE2E_Basic(t *testing.T) {
 	// Helper to create a new client for each sub-test, like a function-scoped fixture
 	newClient := func(t *testing.T) *core.ToolboxClient {
@@ -132,7 +131,7 @@ func TestE2E_Basic(t *testing.T) {
 		toolset, err := client.LoadToolset("", context.Background())
 		require.NoError(t, err)
 
-		assert.Len(t, toolset, 6)
+		assert.Len(t, toolset, 7)
 		toolNames := make(map[string]struct{})
 		for _, tool := range toolset {
 			toolNames[tool.Name()] = struct{}{}
@@ -144,6 +143,7 @@ func TestE2E_Basic(t *testing.T) {
 			"get-row-by-id":           {},
 			"get-n-rows":              {},
 			"search-rows":             {},
+			"process-data":            {},
 		}
 		assert.Equal(t, expectedTools, toolNames)
 	})
@@ -182,7 +182,6 @@ func TestE2E_Basic(t *testing.T) {
 	})
 }
 
-// TestE2E_BindParams maps to the TestBindParams class
 func TestE2E_BindParams(t *testing.T) {
 	newClient := func(t *testing.T) *core.ToolboxClient {
 		client, err := core.NewToolboxClient("http://localhost:5000")
@@ -236,7 +235,6 @@ func TestE2E_BindParams(t *testing.T) {
 	})
 }
 
-// TestE2E_Auth maps to the TestAuth class
 func TestE2E_Auth(t *testing.T) {
 	newClient := func(t *testing.T) *core.ToolboxClient {
 		client, err := core.NewToolboxClient("http://localhost:5000")
@@ -338,7 +336,6 @@ func TestE2E_Auth(t *testing.T) {
 	})
 }
 
-// TestE2E_OptionalParams maps to the TestOptionalParams class
 func TestE2E_OptionalParams(t *testing.T) {
 	// Helper to create a new client
 	newClient := func(t *testing.T) *core.ToolboxClient {
@@ -475,3 +472,112 @@ func TestE2E_OptionalParams(t *testing.T) {
 		assert.Equal(t, "null", response, "Response should be null for non-matching data")
 	})
 }
+
+func TestE2E_MapParams(t *testing.T) {
+	// Helper to create a new client
+	newClient := func(t *testing.T) *core.ToolboxClient {
+		client, err := core.NewToolboxClient("http://localhost:5000")
+		require.NoError(t, err, "Failed to create ToolboxClient")
+		return client
+	}
+
+	// Helper to load the process-data tool
+	processDataTool := func(t *testing.T, client *core.ToolboxClient) *core.ToolboxTool {
+		tool, err := client.LoadTool("process-data", context.Background())
+		require.NoError(t, err, "Failed to load tool 'process-data'")
+		return tool
+	}
+
+	t.Run("test_tool_schema_is_correct", func(t *testing.T) {
+		client := newClient(t)
+		tool := processDataTool(t, client)
+		params := tool.Parameters()
+
+		// Convert slice to map for easy lookup
+		paramMap := make(map[string]core.ParameterSchema)
+		for _, p := range params {
+			paramMap[p.Name] = p
+		}
+
+		// Verify 'execution_context' parameter.
+		execCtxParam, ok := paramMap["execution_context"]
+		require.True(t, ok, "'execution_context' parameter should exist")
+		assert.True(t, execCtxParam.Required, "'execution_context' should be required")
+		assert.Equal(t, "object", execCtxParam.Type, "'execution_context' type should be object")
+
+		// Verify 'user_scores' parameter.
+		userScoresParam, ok := paramMap["user_scores"]
+		require.True(t, ok, "'user_scores' parameter should exist")
+		assert.True(t, userScoresParam.Required, "'user_scores' should be required")
+		assert.Equal(t, "object", userScoresParam.Type, "'user_scores' type should be object")
+
+		// Verify 'feature_flags' parameter.
+		featureFlagsParam, ok := paramMap["feature_flags"]
+		require.True(t, ok, "'feature_flags' parameter should exist")
+		assert.False(t, featureFlagsParam.Required, "'feature_flags' should be optional")
+		assert.Equal(t, "object", featureFlagsParam.Type, "'feature_flags' type should be object")
+	})
+
+	t.Run("test_run_tool_with_all_map_params", func(t *testing.T) {
+		client := newClient(t)
+		tool := processDataTool(t, client)
+
+		// Invoke the tool with valid map parameters.
+		response, err := tool.Invoke(context.Background(), map[string]any{
+			"execution_context": map[string]any{
+				"env":  "prod",
+				"id":   1234,
+				"user": 1234.5,
+			},
+			"user_scores": map[string]any{
+				"user1": 100,
+				"user2": 200,
+			},
+			"feature_flags": map[string]any{
+				"new_feature": true,
+			},
+		})
+		require.NoError(t, err)
+		respStr, ok := response.(string)
+		require.True(t, ok, "Response should be a string")
+
+		assert.Contains(t, respStr, `"execution_context":{"env":"prod","id":1234,"user":1234.5}`)
+		assert.Contains(t, respStr, `"user_scores":{"user1":100,"user2":200}`)
+		assert.Contains(t, respStr, `"feature_flags":{"new_feature":true}`)
+	})
+
+	t.Run("test_run_tool_omitting_optional_map", func(t *testing.T) {
+		client := newClient(t)
+		tool := processDataTool(t, client)
+
+		// Invoke the tool without the optional 'feature_flags' parameter.
+		response, err := tool.Invoke(context.Background(), map[string]any{
+			"execution_context": map[string]any{"env": "dev"},
+			"user_scores":       map[string]any{"user3": 300},
+		})
+		require.NoError(t, err)
+		respStr, ok := response.(string)
+		require.True(t, ok, "Response should be a string")
+
+		assert.Contains(t, respStr, `"execution_context":{"env":"dev"}`)
+		assert.Contains(t, respStr, `"user_scores":{"user3":300}`)
+		assert.Contains(t, respStr, `"feature_flags":null`)
+	})
+
+	t.Run("test_run_tool_with_wrong_map_value_type", func(t *testing.T) {
+		client := newClient(t)
+		tool := processDataTool(t, client)
+
+		// Attempt to invoke the tool with an incorrect type in a map value.
+		_, err := tool.Invoke(context.Background(), map[string]any{
+			"execution_context": map[string]any{"env": "staging"},
+			"user_scores": map[string]any{
+				"user4": "not-an-integer",
+			},
+		})
+
+		// Assert that an error was returned.
+		require.Error(t, err, "Expected an error for wrong map value type")
+		assert.Contains(t, err.Error(), "expects an integer, but got string", "Error message should indicate a validation failure")
+	})
+}

core/protocol.go

@@ -21,12 +21,13 @@ import (
 
 // Schema for a tool parameter.
 type ParameterSchema struct {
-	Name        string           `json:"name"`
-	Type        string           `json:"type"`
-	Required    bool             `json:"required,omitempty"`
-	Description string           `json:"description"`
-	AuthSources []string         `json:"authSources,omitempty"`
-	Items       *ParameterSchema `json:"items,omitempty"`
+	Name                 string           `json:"name"`
+	Type                 string           `json:"type"`
+	Required             bool             `json:"required,omitempty"`
+	Description          string           `json:"description"`
+	AuthSources          []string         `json:"authSources,omitempty"`
+	Items                *ParameterSchema `json:"items,omitempty"`
+	AdditionalProperties any              `json:"additionalProperties,omitempty"`
 }
 
 // validateType is a helper for manual type checking.
@@ -64,22 +65,90 @@ func (p *ParameterSchema) validateType(value any) error {
 		if v.Kind() != reflect.Slice && v.Kind() != reflect.Array {
 			return fmt.Errorf("parameter '%s' expects an array/slice, but got %T", p.Name, value)
 		}
-		if p.Items == nil {
-			return fmt.Errorf("parameter '%s' is an array but is missing item type definition", p.Name)
-		}
 		for i := range v.Len() {
 			item := v.Index(i).Interface()
 
 			if err := p.Items.validateType(item); err != nil {
 				return fmt.Errorf("error in array '%s' at index %d: %w", p.Name, i, err)
 			}
 		}
+	case "object":
+		// First, check that the value is a map with string keys.
+		valMap, ok := value.(map[string]any)
+		if !ok {
+			return fmt.Errorf("parameter '%s' expects a map, but got %T", p.Name, value)
+		}
+
+		switch ap := p.AdditionalProperties.(type) {
+		// No validation of values
+		case bool:
+
+		// Validate type for each value in map
+		case *ParameterSchema:
+			for key, val := range valMap {
+				if err := ap.validateType(val); err != nil {
+					return fmt.Errorf("error in object '%s' for key '%s': %w", p.Name, key, err)
+				}
+			}
+
+		default:
+			// This is a schema / manifest error.
+			return fmt.Errorf(
+				"invalid schema for parameter '%s': AdditionalProperties must be a boolean or a map[string]any, but got %T",
+				p.Name,
+				ap,
+			)
+		}
 	default:
 		return fmt.Errorf("unknown type '%s' in schema for parameter '%s'", p.Type, p.Name)
 	}
 	return nil
 }
 
+// ValidateDefinition checks if the schema itself is well-formed.
+func (p *ParameterSchema) ValidateDefinition() error {
+	if p.Type == "" {
+		return fmt.Errorf("schema validation failed for '%s': type is missing", p.Name)
+	}
+
+	switch p.Type {
+	case "array":
+		if p.Items == nil {
+			return fmt.Errorf("parameter '%s' is an array but is missing item type definition", p.Name)
+		}
+		// Recursively validate the nested schema's definition.
+		if err := p.Items.ValidateDefinition(); err != nil {
+			return err
+		}
+
+	case "object":
+		switch ap := p.AdditionalProperties.(type) {
+		case bool:
+			// Valid scenario
+		case *ParameterSchema:
+			if err := ap.ValidateDefinition(); err != nil {
+				return err
+			}
+		default:
+			// Any other type is an invalid schema definition.
+			return fmt.Errorf(
+				"invalid schema for parameter '%s': AdditionalProperties must be a boolean or a schema, but got %T",
+				p.Name,
+				ap,
+			)
+		}
+
+	case "string", "integer", "float", "boolean":
+		// No type-specific rules for these.
+		break
+
+	default:
+		return fmt.Errorf("unknown schema type '%s' for parameter '%s'", p.Type, p.Name)
+	}
+
+	return nil
+}
+
 // Schema for a tool.
 type ToolSchema struct {
 	Description  string            `json:"description"`