diff --git a/plugins/cmd/ndpgen/internal/generator.go b/plugins/cmd/ndpgen/internal/generator.go index 803a4236e..ad36ee640 100644 --- a/plugins/cmd/ndpgen/internal/generator.go +++ b/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)) diff --git a/plugins/cmd/ndpgen/internal/templates/capability.rs.tmpl b/plugins/cmd/ndpgen/internal/templates/capability.rs.tmpl index 01629ed53..01e6513ac 100644 --- a/plugins/cmd/ndpgen/internal/templates/capability.rs.tmpl +++ b/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())?; diff --git a/plugins/cmd/ndpgen/internal/xtp_schema.go b/plugins/cmd/ndpgen/internal/xtp_schema.go index 17e473c0b..6bef85f45 100644 --- a/plugins/cmd/ndpgen/internal/xtp_schema.go +++ b/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),