fix(plugins): enhance type handling for Rust and XTP output in capability generation

Signed-off-by: Deluan <deluan@navidrome.org>

plugins/cmd/ndpgen/internal/generator.go

@@ -415,6 +415,7 @@ func rustCapabilityFuncMap(cap Capability) template.FuncMap {
 		"rustMethodName":      func(name string) string { return ToSnakeCase(name) },
 		"fieldRustType":       func(f FieldDef) string { return f.RustType(knownStructs) },
 		"rustOutputType":      rustOutputType,
+		"isPrimitiveRust":     isPrimitiveRustType,
 		"skipSerializingFunc": skipSerializingFunc,
 		"hasHashMap":          hasHashMap,
 		"agentName":           capabilityAgentName,
@@ -473,9 +474,37 @@ func rustOutputType(goType string) string {
 	if strings.HasPrefix(goType, "*") {
 		return goType[1:]
 	}
+	// Convert Go primitives to Rust primitives
+	switch goType {
+	case "bool":
+		return "bool"
+	case "string":
+		return "String"
+	case "int", "int32":
+		return "i32"
+	case "int64":
+		return "i64"
+	case "float32":
+		return "f32"
+	case "float64":
+		return "f64"
+	}
 	return goType
 }
 
+// isPrimitiveRustType returns true if the Go type maps to a Rust primitive type.
+func isPrimitiveRustType(goType string) bool {
+	// Strip pointer prefix first
+	if strings.HasPrefix(goType, "*") {
+		goType = goType[1:]
+	}
+	switch goType {
+	case "bool", "string", "int", "int32", "int64", "float32", "float64":
+		return true
+	}
+	return false
+}
+
 // rustConstName converts a Go const name to Rust convention (SCREAMING_SNAKE_CASE).
 func rustConstName(name string) string {
 	return strings.ToUpper(ToSnakeCase(name))

plugins/cmd/ndpgen/internal/templates/capability.rs.tmpl

@@ -115,7 +115,7 @@ macro_rules! register_{{snakeCase .Package}} {
             {{- if .HasInput}}
             req: extism_pdk::Json<$crate::{{snakeCase $.Package}}::{{rustOutputType .Input.Type}}>
             {{- end}}
-        ) -> extism_pdk::FnResult<{{if .HasOutput}}extism_pdk::Json<$crate::{{snakeCase $.Package}}::{{rustOutputType .Output.Type}}>{{else}}(){{end}}> {
+        ) -> extism_pdk::FnResult<{{if .HasOutput}}extism_pdk::Json<{{if isPrimitiveRust .Output.Type}}{{rustOutputType .Output.Type}}{{else}}$crate::{{snakeCase $.Package}}::{{rustOutputType .Output.Type}}{{end}}>{{else}}(){{end}}> {
             let plugin = <$plugin_type>::default();
             {{- if and .HasInput .HasOutput}}
             let result = $crate::{{snakeCase $.Package}}::{{agentName $.Capability}}::{{rustMethodName .Name}}(&plugin, req.into_inner())?;
@@ -162,7 +162,7 @@ macro_rules! {{registerMacroName .Name}} {
             {{- if .HasInput}}
             req: extism_pdk::Json<$crate::{{snakeCase $.Package}}::{{rustOutputType .Input.Type}}>
             {{- end}}
-        ) -> extism_pdk::FnResult<{{if .HasOutput}}extism_pdk::Json<$crate::{{snakeCase $.Package}}::{{rustOutputType .Output.Type}}>{{else}}(){{end}}> {
+        ) -> extism_pdk::FnResult<{{if .HasOutput}}extism_pdk::Json<{{if isPrimitiveRust .Output.Type}}{{rustOutputType .Output.Type}}{{else}}$crate::{{snakeCase $.Package}}::{{rustOutputType .Output.Type}}{{end}}>{{else}}(){{end}}> {
             let plugin = <$plugin_type>::default();
             {{- if and .HasInput .HasOutput}}
             let result = $crate::{{snakeCase $.Package}}::{{providerInterface .}}::{{rustMethodName .Name}}(&plugin, req.into_inner())?;

plugins/cmd/ndpgen/internal/xtp_schema.go

@@ -26,7 +26,8 @@ type (
 	}
 
 	xtpIOParam struct {
-		Ref         string `yaml:"$ref"`
+		Ref         string `yaml:"$ref,omitempty"`
+		Type        string `yaml:"type,omitempty"`
 		ContentType string `yaml:"contentType"`
 	}
 
@@ -87,14 +88,32 @@ func buildExport(export Export) xtpExport {
 		}
 	}
 	if export.Output.Type != "" {
-		e.Output = &xtpIOParam{
-			Ref:         "#/components/schemas/" + strings.TrimPrefix(export.Output.Type, "*"),
-			ContentType: "application/json",
+		outputType := strings.TrimPrefix(export.Output.Type, "*")
+		// Check if output is a primitive type
+		if isPrimitiveGoType(outputType) {
+			e.Output = &xtpIOParam{
+				Type:        goTypeToXTPType(outputType),
+				ContentType: "application/json",
+			}
+		} else {
+			e.Output = &xtpIOParam{
+				Ref:         "#/components/schemas/" + outputType,
+				ContentType: "application/json",
+			}
 		}
 	}
 	return e
 }
 
+// isPrimitiveGoType returns true if the Go type is a primitive type.
+func isPrimitiveGoType(goType string) bool {
+	switch goType {
+	case "bool", "string", "int", "int32", "int64", "float32", "float64", "[]byte":
+		return true
+	}
+	return false
+}
+
 func buildSchemas(cap Capability) yaml.Node {
 	schemas := yaml.Node{Kind: yaml.MappingNode}
 	knownTypes := cap.KnownStructs()
@@ -102,12 +121,17 @@ func buildSchemas(cap Capability) yaml.Node {
 		knownTypes[alias.Name] = true
 	}
 
+	// Collect types that are actually used by exports
+	usedTypes := collectUsedTypes(cap, knownTypes)
+
 	// Sort structs by name for consistent output
 	structNames := make([]string, 0, len(cap.Structs))
 	structMap := make(map[string]StructDef)
 	for _, st := range cap.Structs {
-		structNames = append(structNames, st.Name)
-		structMap[st.Name] = st
+		if usedTypes[st.Name] {
+			structNames = append(structNames, st.Name)
+			structMap[st.Name] = st
+		}
 	}
 	sort.Strings(structNames)
 
@@ -116,8 +140,11 @@ func buildSchemas(cap Capability) yaml.Node {
 		addToMap(&schemas, name, buildObjectSchema(st, knownTypes))
 	}
 
-	// Build enum types from type aliases
+	// Build enum types from type aliases (only if used by exports)
 	for _, alias := range cap.TypeAliases {
+		if !usedTypes[alias.Name] {
+			continue
+		}
 		if alias.Type == "string" {
 			for _, cg := range cap.Consts {
 				if cg.Type == alias.Name {
@@ -131,6 +158,48 @@ func buildSchemas(cap Capability) yaml.Node {
 	return schemas
 }
 
+// collectUsedTypes returns a set of type names that are reachable from exports.
+func collectUsedTypes(cap Capability, knownTypes map[string]bool) map[string]bool {
+	used := make(map[string]bool)
+
+	// Start with types directly referenced by exports
+	for _, export := range cap.Methods {
+		if export.Input.Type != "" {
+			addTypeAndDeps(strings.TrimPrefix(export.Input.Type, "*"), cap, knownTypes, used)
+		}
+		if export.Output.Type != "" {
+			outputType := strings.TrimPrefix(export.Output.Type, "*")
+			if !isPrimitiveGoType(outputType) {
+				addTypeAndDeps(outputType, cap, knownTypes, used)
+			}
+		}
+	}
+
+	return used
+}
+
+// addTypeAndDeps adds a type and all its dependencies to the used set.
+func addTypeAndDeps(typeName string, cap Capability, knownTypes map[string]bool, used map[string]bool) {
+	if used[typeName] || !knownTypes[typeName] {
+		return
+	}
+	used[typeName] = true
+
+	// Find the struct and add its field types
+	for _, st := range cap.Structs {
+		if st.Name == typeName {
+			for _, field := range st.Fields {
+				fieldType := strings.TrimPrefix(field.Type, "*")
+				fieldType = strings.TrimPrefix(fieldType, "[]")
+				if knownTypes[fieldType] {
+					addTypeAndDeps(fieldType, cap, knownTypes, used)
+				}
+			}
+			return
+		}
+	}
+}
+
 func buildObjectSchema(st StructDef, knownTypes map[string]bool) xtpObjectSchema {
 	schema := xtpObjectSchema{
 		Description: cleanDocForYAML(st.Doc),