pluginsdk/plugin/wasmguest/mcp.go

package wasmguest

import (
	"encoding/json"
	"sort"
	"strings"
	"unicode"

	abiv1 "git.dev.alexdunmow.com/block/pluginsdk/abi/v1"
	"git.dev.alexdunmow.com/block/pluginsdk/rbac"
	"google.golang.org/protobuf/reflect/protoreflect"
	"google.golang.org/protobuf/reflect/protoregistry"
)

// captureMCPTools resolves the plugin's registered Connect procedures through
// the protobuf registry compiled into the guest. It emits only unary methods
// that have an explicit RBAC mapping; the host repeats that RBAC intersection
// before adding the definitions to its live MCP catalogue.
func captureMCPTools(methodRoles map[string]rbac.Role) []*abiv1.MCPToolDescriptor {
	procedures := make([]string, 0, len(methodRoles))
	for procedure := range methodRoles {
		procedures = append(procedures, procedure)
	}
	sort.Strings(procedures)

	tools := make([]*abiv1.MCPToolDescriptor, 0, len(procedures))
	for _, procedure := range procedures {
		serviceName, methodName, ok := splitConnectProcedure(procedure)
		if !ok {
			continue
		}
		desc, err := protoregistry.GlobalFiles.FindDescriptorByName(protoreflect.FullName(serviceName))
		if err != nil {
			continue
		}
		service, ok := desc.(protoreflect.ServiceDescriptor)
		if !ok {
			continue
		}
		method := service.Methods().ByName(protoreflect.Name(methodName))
		if method == nil || method.IsStreamingClient() || method.IsStreamingServer() {
			continue
		}

		schema, err := json.Marshal(schemaForMessage(method.Input(), map[protoreflect.FullName]bool{}))
		if err != nil {
			continue
		}
		description, documentation := methodDocumentation(method)
		tools = append(tools, &abiv1.MCPToolDescriptor{
			Procedure:       procedure,
			ServiceFullName: serviceName,
			MethodName:      methodName,
			Description:     description,
			Documentation:   documentation,
			InputSchemaJson: schema,
		})
	}
	return tools
}

func splitConnectProcedure(procedure string) (serviceName, methodName string, ok bool) {
	trimmed := strings.TrimPrefix(strings.TrimSpace(procedure), "/")
	serviceName, methodName, ok = strings.Cut(trimmed, "/")
	return serviceName, methodName, ok && serviceName != "" && methodName != "" && !strings.Contains(methodName, "/")
}

func methodDocumentation(method protoreflect.MethodDescriptor) (description, documentation string) {
	location := method.ParentFile().SourceLocations().ByDescriptor(method)
	documentation = strings.TrimSpace(location.LeadingComments)
	paragraph := documentation
	if before, _, ok := strings.Cut(paragraph, "\n\n"); ok {
		paragraph = before
	}
	paragraph = strings.TrimSpace(strings.ReplaceAll(paragraph, "\n", " "))
	methodName := string(method.Name())
	if after, ok := strings.CutPrefix(paragraph, methodName+" "); ok {
		paragraph = strings.TrimSpace(after)
		if paragraph != "" {
			runes := []rune(paragraph)
			runes[0] = unicode.ToUpper(runes[0])
			paragraph = string(runes)
		}
	}
	if paragraph == "" {
		paragraph = "Calls " + string(method.FullName()) + "."
	}
	return paragraph, documentation
}

func schemaForMessage(message protoreflect.MessageDescriptor, visited map[protoreflect.FullName]bool) map[string]any {
	object := map[string]any{"type": "object"}
	if message == nil || visited[message.FullName()] {
		return object
	}
	visited[message.FullName()] = true
	defer delete(visited, message.FullName())

	properties := make(map[string]any, message.Fields().Len())
	for i := range message.Fields().Len() {
		field := message.Fields().Get(i)
		properties[string(field.Name())] = fieldSchema(field, visited)
	}
	if len(properties) > 0 {
		object["properties"] = properties
	}
	return object
}

func fieldSchema(field protoreflect.FieldDescriptor, visited map[protoreflect.FullName]bool) map[string]any {
	var schema map[string]any
	switch {
	case field.IsMap():
		schema = map[string]any{
			"type":                 "object",
			"additionalProperties": singularFieldSchema(field.MapValue(), visited),
		}
	case field.IsList():
		schema = map[string]any{"type": "array", "items": singularFieldSchema(field, visited)}
	default:
		schema = singularFieldSchema(field, visited)
	}

	location := field.ParentFile().SourceLocations().ByDescriptor(field)
	if comment := strings.TrimSpace(strings.ReplaceAll(location.LeadingComments, "\n", " ")); comment != "" {
		schema["description"] = comment
	}
	return schema
}

func singularFieldSchema(field protoreflect.FieldDescriptor, visited map[protoreflect.FullName]bool) map[string]any {
	switch field.Kind() {
	case protoreflect.MessageKind, protoreflect.GroupKind:
		return wellKnownOrMessage(field.Message(), visited)
	case protoreflect.EnumKind:
		return enumSchema(field.Enum())
	default:
		return scalarSchema(field.Kind())
	}
}

func scalarSchema(kind protoreflect.Kind) map[string]any {
	switch kind {
	case protoreflect.BoolKind:
		return map[string]any{"type": "boolean"}
	case protoreflect.Int32Kind, protoreflect.Sint32Kind, protoreflect.Sfixed32Kind,
		protoreflect.Uint32Kind, protoreflect.Fixed32Kind,
		protoreflect.Int64Kind, protoreflect.Sint64Kind, protoreflect.Sfixed64Kind,
		protoreflect.Uint64Kind, protoreflect.Fixed64Kind:
		return map[string]any{"type": "integer"}
	case protoreflect.FloatKind, protoreflect.DoubleKind:
		return map[string]any{"type": "number"}
	case protoreflect.StringKind:
		return map[string]any{"type": "string"}
	case protoreflect.BytesKind:
		return map[string]any{"type": "string", "contentEncoding": "base64"}
	default:
		return map[string]any{}
	}
}

func enumSchema(enum protoreflect.EnumDescriptor) map[string]any {
	schema := map[string]any{"type": "string"}
	if enum == nil {
		return schema
	}
	values := make([]any, 0, enum.Values().Len())
	for i := range enum.Values().Len() {
		values = append(values, string(enum.Values().Get(i).Name()))
	}
	if len(values) > 0 {
		schema["enum"] = values
	}
	return schema
}

func wellKnownOrMessage(message protoreflect.MessageDescriptor, visited map[protoreflect.FullName]bool) map[string]any {
	if message == nil {
		return map[string]any{"type": "object"}
	}
	switch message.FullName() {
	case "google.protobuf.Timestamp":
		return map[string]any{"type": "string", "format": "date-time"}
	case "google.protobuf.Duration", "google.protobuf.FieldMask":
		return map[string]any{"type": "string"}
	case "google.protobuf.Struct":
		return map[string]any{"type": "object"}
	case "google.protobuf.ListValue":
		return map[string]any{"type": "array"}
	case "google.protobuf.Value", "google.protobuf.Any":
		return map[string]any{}
	case "google.protobuf.Empty":
		return map[string]any{"type": "object"}
	case "google.protobuf.StringValue", "google.protobuf.BytesValue":
		return map[string]any{"type": "string"}
	case "google.protobuf.BoolValue":
		return map[string]any{"type": "boolean"}
	case "google.protobuf.Int32Value", "google.protobuf.UInt32Value",
		"google.protobuf.Int64Value", "google.protobuf.UInt64Value":
		return map[string]any{"type": "integer"}
	case "google.protobuf.FloatValue", "google.protobuf.DoubleValue":
		return map[string]any{"type": "number"}
	default:
		return schemaForMessage(message, visited)
	}
}