Skip to content

Commit c914d42

Browse files
authored
vMCP: Composite Tools supports non-string Arguments (#2971)
Changed WorkflowStep.Arguments from map[string]string to *runtime.RawExtension to preserve value types Updated converter logic to unmarshal RawExtension into map[string]any with proper error handling Updated webhook validation in VirtualMCPCompositeToolDefinition to validate Arguments JSON and template syntax Added comprehensive unit and integration tests for non-string argument types --------- Signed-off-by: Jeremy Drouillard <[email protected]>
1 parent 147d5f5 commit c914d42

17 files changed

+435
-85
lines changed

cmd/thv-operator/api/v1alpha1/virtualmcpcompositetooldefinition_webhook.go

Lines changed: 23 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -229,28 +229,45 @@ func (*VirtualMCPCompositeToolDefinition) validateStepType(index int, step Workf
229229

230230
// validateStepTemplates validates all template fields in a step
231231
func (*VirtualMCPCompositeToolDefinition) validateStepTemplates(index int, step WorkflowStep) error {
232-
for argName, argValue := range step.Arguments {
233-
if err := validateTemplate(argValue); err != nil {
234-
return fmt.Errorf("spec.steps[%d].arguments[%s]: invalid template: %v", index, argName, err)
232+
return validateWorkflowStepTemplates("spec.steps", index, step)
233+
}
234+
235+
// validateWorkflowStepTemplates validates all template fields in a workflow step.
236+
// This is a shared validation function used by both VirtualMCPServer and VirtualMCPCompositeToolDefinition webhooks.
237+
// The pathPrefix parameter allows customizing error message paths (e.g., "spec.steps" or "spec.compositeTools[0].steps").
238+
func validateWorkflowStepTemplates(pathPrefix string, index int, step WorkflowStep) error {
239+
// Validate template syntax in arguments (only for string values)
240+
if step.Arguments != nil && len(step.Arguments.Raw) > 0 {
241+
var args map[string]any
242+
if err := json.Unmarshal(step.Arguments.Raw, &args); err != nil {
243+
return fmt.Errorf("%s[%d].arguments: invalid JSON: %v", pathPrefix, index, err)
244+
}
245+
for argName, argValue := range args {
246+
// Only validate template syntax for string values
247+
if strValue, ok := argValue.(string); ok {
248+
if err := validateTemplate(strValue); err != nil {
249+
return fmt.Errorf("%s[%d].arguments[%s]: invalid template: %v", pathPrefix, index, argName, err)
250+
}
251+
}
235252
}
236253
}
237254

238255
if step.Condition != "" {
239256
if err := validateTemplate(step.Condition); err != nil {
240-
return fmt.Errorf("spec.steps[%d].condition: invalid template: %v", index, err)
257+
return fmt.Errorf("%s[%d].condition: invalid template: %v", pathPrefix, index, err)
241258
}
242259
}
243260

244261
if step.Message != "" {
245262
if err := validateTemplate(step.Message); err != nil {
246-
return fmt.Errorf("spec.steps[%d].message: invalid template: %v", index, err)
263+
return fmt.Errorf("%s[%d].message: invalid template: %v", pathPrefix, index, err)
247264
}
248265
}
249266

250267
// Validate JSON Schema for elicitation steps
251268
if step.Schema != nil {
252269
if err := validateJSONSchema(step.Schema.Raw); err != nil {
253-
return fmt.Errorf("spec.steps[%d].schema: invalid JSON Schema: %v", index, err)
270+
return fmt.Errorf("%s[%d].schema: invalid JSON Schema: %v", pathPrefix, index, err)
254271
}
255272
}
256273

cmd/thv-operator/api/v1alpha1/virtualmcpcompositetooldefinition_webhook_test.go

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -126,9 +126,8 @@ func TestVirtualMCPCompositeToolDefinitionValidate(t *testing.T) {
126126
{
127127
ID: "deploy",
128128
Tool: "kubectl.apply",
129-
Arguments: map[string]string{
130-
"namespace": "{{.params.environment}}",
131-
"replicas": "{{.params.replicas}}",
129+
Arguments: &runtime.RawExtension{
130+
Raw: []byte(`{"namespace": "{{.params.environment}}", "replicas": "{{.params.replicas}}"}`),
132131
},
133132
},
134133
},
@@ -515,8 +514,8 @@ func TestVirtualMCPCompositeToolDefinitionValidate(t *testing.T) {
515514
{
516515
ID: "deploy",
517516
Tool: "kubectl.apply",
518-
Arguments: map[string]string{
519-
"namespace": "{{.params.env",
517+
Arguments: &runtime.RawExtension{
518+
Raw: []byte(`{"namespace": "{{.params.env"}`),
520519
},
521520
},
522521
},

cmd/thv-operator/api/v1alpha1/virtualmcpserver_types.go

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -280,10 +280,14 @@ type WorkflowStep struct {
280280
// +optional
281281
Tool string `json:"tool,omitempty"`
282282

283-
// Arguments is a map of argument templates
284-
// Supports Go template syntax with .params and .steps
283+
// Arguments is a map of argument values with template expansion support.
284+
// Supports Go template syntax with .params and .steps for string values.
285+
// Non-string values (integers, booleans, arrays, objects) are passed as-is.
286+
// Note: the templating is only supported on the first level of the key-value pairs.
285287
// +optional
286-
Arguments map[string]string `json:"arguments,omitempty"`
288+
// +kubebuilder:pruning:PreserveUnknownFields
289+
// +kubebuilder:validation:Type=object
290+
Arguments *runtime.RawExtension `json:"arguments,omitempty"`
287291

288292
// Message is the elicitation message
289293
// Only used when Type is "elicitation"

cmd/thv-operator/api/v1alpha1/virtualmcpserver_webhook.go

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -274,7 +274,13 @@ func validateCompositeToolStep(
274274
}
275275

276276
// Validate elicitation response handlers
277-
return validateStepElicitationResponseHandlers(toolIndex, stepIndex, step)
277+
if err := validateStepElicitationResponseHandlers(toolIndex, stepIndex, step); err != nil {
278+
return err
279+
}
280+
281+
// Validate templates in arguments and other fields
282+
pathPrefix := fmt.Sprintf("spec.compositeTools[%d].steps", toolIndex)
283+
return validateWorkflowStepTemplates(pathPrefix, stepIndex, step)
278284
}
279285

280286
// validateStepType validates the step type and type-specific requirements

cmd/thv-operator/api/v1alpha1/zz_generated.deepcopy.go

Lines changed: 2 additions & 4 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

cmd/thv-operator/controllers/virtualmcpserver_vmcpconfig_test.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -703,8 +703,8 @@ func TestVirtualMCPServerReconciler_CompositeToolRefs_EndToEnd(t *testing.T) {
703703
ID: "step1",
704704
Type: "tool",
705705
Tool: "backend.echo",
706-
Arguments: map[string]string{
707-
"input": "{{ .params.message }}",
706+
Arguments: &runtime.RawExtension{
707+
Raw: []byte(`{"input": "{{ .params.message }}"}`),
708708
},
709709
},
710710
},

cmd/thv-operator/pkg/vmcpconfig/converter.go

Lines changed: 42 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -518,17 +518,20 @@ func (c *Converter) resolveMCPToolConfig(
518518
func (c *Converter) convertCompositeTools(
519519
ctx context.Context,
520520
vmcp *mcpv1alpha1.VirtualMCPServer,
521-
) []*vmcpconfig.CompositeToolConfig {
521+
) ([]*vmcpconfig.CompositeToolConfig, error) {
522522
compositeTools := make([]*vmcpconfig.CompositeToolConfig, 0, len(vmcp.Spec.CompositeTools))
523523

524524
for _, crdTool := range vmcp.Spec.CompositeTools {
525-
tool := c.convertCompositeToolSpec(
525+
tool, err := c.convertCompositeToolSpec(
526526
ctx, crdTool.Name, crdTool.Description, crdTool.Timeout,
527527
crdTool.Parameters, crdTool.Steps, crdTool.Output, crdTool.Name)
528+
if err != nil {
529+
return nil, err
530+
}
528531
compositeTools = append(compositeTools, tool)
529532
}
530533

531-
return compositeTools
534+
return compositeTools, nil
532535
}
533536

534537
// convertAllCompositeTools converts both inline CompositeTools and referenced CompositeToolRefs,
@@ -538,12 +541,14 @@ func (c *Converter) convertAllCompositeTools(
538541
vmcp *mcpv1alpha1.VirtualMCPServer,
539542
) ([]*vmcpconfig.CompositeToolConfig, error) {
540543
// Convert inline composite tools
541-
inlineTools := c.convertCompositeTools(ctx, vmcp)
544+
inlineTools, err := c.convertCompositeTools(ctx, vmcp)
545+
if err != nil {
546+
return nil, fmt.Errorf("failed to convert inline composite tools: %w", err)
547+
}
542548

543549
// Convert referenced composite tools
544550
var referencedTools []*vmcpconfig.CompositeToolConfig
545551
if len(vmcp.Spec.CompositeToolRefs) > 0 {
546-
var err error
547552
referencedTools, err = c.convertReferencedCompositeTools(ctx, vmcp)
548553
if err != nil {
549554
return nil, fmt.Errorf("failed to convert referenced composite tools: %w", err)
@@ -587,7 +592,10 @@ func (c *Converter) convertReferencedCompositeTools(
587592
}
588593

589594
// Convert the referenced definition to CompositeToolConfig
590-
tool := c.convertCompositeToolDefinition(ctx, compositeToolDef)
595+
tool, err := c.convertCompositeToolDefinition(ctx, compositeToolDef)
596+
if err != nil {
597+
return nil, fmt.Errorf("failed to convert referenced tool %q: %w", ref.Name, err)
598+
}
591599
referencedTools = append(referencedTools, tool)
592600
}
593601

@@ -598,7 +606,7 @@ func (c *Converter) convertReferencedCompositeTools(
598606
func (c *Converter) convertCompositeToolDefinition(
599607
ctx context.Context,
600608
def *mcpv1alpha1.VirtualMCPCompositeToolDefinition,
601-
) *vmcpconfig.CompositeToolConfig {
609+
) (*vmcpconfig.CompositeToolConfig, error) {
602610
return c.convertCompositeToolSpec(
603611
ctx, def.Spec.Name, def.Spec.Description, def.Spec.Timeout,
604612
def.Spec.Parameters, def.Spec.Steps, def.Spec.Output, def.Name)
@@ -613,7 +621,7 @@ func (c *Converter) convertCompositeToolSpec(
613621
steps []mcpv1alpha1.WorkflowStep,
614622
output *mcpv1alpha1.OutputSpec,
615623
toolNameForLogging string,
616-
) *vmcpconfig.CompositeToolConfig {
624+
) (*vmcpconfig.CompositeToolConfig, error) {
617625
tool := &vmcpconfig.CompositeToolConfig{
618626
Name: name,
619627
Description: description,
@@ -650,30 +658,39 @@ func (c *Converter) convertCompositeToolSpec(
650658
}
651659

652660
// Convert steps
653-
tool.Steps = c.convertWorkflowSteps(ctx, steps, toolNameForLogging)
661+
workflowSteps, err := c.convertWorkflowSteps(ctx, steps, toolNameForLogging)
662+
if err != nil {
663+
return nil, fmt.Errorf("failed to convert steps for tool %q: %w", toolNameForLogging, err)
664+
}
665+
tool.Steps = workflowSteps
654666

655667
// Convert output configuration
656668
if output != nil {
657669
tool.Output = convertOutputSpec(ctx, output)
658670
}
659671

660-
return tool
672+
return tool, nil
661673
}
662674

663675
// convertWorkflowSteps converts a slice of WorkflowStep CRD objects to WorkflowStepConfig.
664676
func (*Converter) convertWorkflowSteps(
665677
ctx context.Context,
666678
steps []mcpv1alpha1.WorkflowStep,
667679
toolNameForLogging string,
668-
) []*vmcpconfig.WorkflowStepConfig {
680+
) ([]*vmcpconfig.WorkflowStepConfig, error) {
669681
workflowSteps := make([]*vmcpconfig.WorkflowStepConfig, 0, len(steps))
670682

671683
for _, crdStep := range steps {
684+
args, err := convertArguments(crdStep.Arguments)
685+
if err != nil {
686+
return nil, fmt.Errorf("step %q: %w", crdStep.ID, err)
687+
}
688+
672689
step := &vmcpconfig.WorkflowStepConfig{
673690
ID: crdStep.ID,
674691
Type: crdStep.Type,
675692
Tool: crdStep.Tool,
676-
Arguments: convertArguments(crdStep.Arguments),
693+
Arguments: args,
677694
Message: crdStep.Message,
678695
Condition: crdStep.Condition,
679696
DependsOn: crdStep.DependsOn,
@@ -739,7 +756,7 @@ func (*Converter) convertWorkflowSteps(
739756
workflowSteps = append(workflowSteps, step)
740757
}
741758

742-
return workflowSteps
759+
return workflowSteps, nil
743760
}
744761

745762
// validateCompositeToolNames checks for duplicate tool names across all composite tools.
@@ -754,13 +771,19 @@ func validateCompositeToolNames(tools []*vmcpconfig.CompositeToolConfig) error {
754771
return nil
755772
}
756773

757-
// convertArguments converts string arguments to any type for template expansion
758-
func convertArguments(args map[string]string) map[string]any {
759-
result := make(map[string]any, len(args))
760-
for k, v := range args {
761-
result[k] = v
774+
// convertArguments converts arguments from runtime.RawExtension to map[string]any.
775+
// This preserves the original types (integers, booleans, arrays, objects) from the CRD.
776+
// Returns an empty map if no arguments are specified.
777+
// Returns an error if the JSON fails to unmarshal.
778+
func convertArguments(args *runtime.RawExtension) (map[string]any, error) {
779+
if args == nil || len(args.Raw) == 0 {
780+
return make(map[string]any), nil
781+
}
782+
var result map[string]any
783+
if err := json.Unmarshal(args.Raw, &result); err != nil {
784+
return nil, fmt.Errorf("failed to unmarshal arguments: %w", err)
762785
}
763-
return result
786+
return result, nil
764787
}
765788

766789
// convertOutputSpec converts OutputSpec from CRD to vmcp config OutputConfig

0 commit comments

Comments
 (0)