yao/grpc/mcp/mcp.go

package mcp

import (
	"context"
	"encoding/json"

	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/status"

	goumcp "github.com/yaoapp/gou/mcp"
	"github.com/yaoapp/yao/grpc/auth"
	"github.com/yaoapp/yao/grpc/pb"
)

// Handler implements the MCP gRPC methods.
type Handler struct{}

// grpcAuthProvider adapts gRPC AuthorizedInfo to the AuthorizedProvider
// interface expected by gou/mcp/process for propagating auth to process calls.
type grpcAuthProvider struct {
	m map[string]interface{}
}

func (p *grpcAuthProvider) GetAuthorizedMap() map[string]interface{} { return p.m }

func authProviderFromCtx(ctx context.Context) *grpcAuthProvider {
	info := auth.GetAuthorizedInfo(ctx)
	if info == nil {
		return nil
	}
	m := map[string]interface{}{
		"sub":        info.Subject,
		"client_id":  info.ClientID,
		"scope":      info.Scope,
		"session_id": info.SessionID,
		"user_id":    info.UserID,
		"team_id":    info.TeamID,
		"tenant_id":  info.TenantID,
	}
	if md, ok := metadata.FromIncomingContext(ctx); ok {
		if ids := md.Get("x-workspace-id"); len(ids) > 0 && ids[0] != "" {
			m["workspace_id"] = ids[0]
		}
		if ids := md.Get("x-sandbox-id"); len(ids) > 0 && ids[0] != "" {
			m["sandbox_id"] = ids[0]
		}
	}
	return &grpcAuthProvider{m: m}
}

// MCPListTools lists all available MCP tools for a given session.
func (h *Handler) MCPListTools(ctx context.Context, req *pb.MCPListRequest) (*pb.MCPListResponse, error) {
	client, err := goumcp.Select(req.SessionId)
	if err != nil {
		return nil, status.Errorf(codes.NotFound, "MCP client not found: %v", err)
	}

	resp, err := client.ListTools(ctx, "")
	if err != nil {
		return nil, status.Errorf(codes.Internal, "ListTools failed: %v", err)
	}

	data, err := json.Marshal(resp.Tools)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to marshal tools: %v", err)
	}

	return &pb.MCPListResponse{Tools: data}, nil
}

// MCPCallTool calls an MCP tool by name with the provided arguments.
func (h *Handler) MCPCallTool(ctx context.Context, req *pb.MCPCallRequest) (*pb.MCPCallResponse, error) {
	client, err := goumcp.Select(req.SessionId)
	if err != nil {
		return nil, status.Errorf(codes.NotFound, "MCP client not found: %v", err)
	}

	var args interface{}
	if len(req.Arguments) > 0 {
		if err := json.Unmarshal(req.Arguments, &args); err != nil {
			return nil, status.Errorf(codes.InvalidArgument, "invalid arguments JSON: %v", err)
		}
	}

	var extraArgs []interface{}
	if ap := authProviderFromCtx(ctx); ap != nil {
		extraArgs = append(extraArgs, ap)
	}

	resp, err := client.CallTool(ctx, req.Tool, args, extraArgs...)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "CallTool failed: %v", err)
	}

	data, err := json.Marshal(resp)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to marshal result: %v", err)
	}

	return &pb.MCPCallResponse{Result: data}, nil
}

// MCPListResources lists all available MCP resources for a given session.
func (h *Handler) MCPListResources(ctx context.Context, req *pb.MCPListRequest) (*pb.MCPResourcesResponse, error) {
	client, err := goumcp.Select(req.SessionId)
	if err != nil {
		return nil, status.Errorf(codes.NotFound, "MCP client not found: %v", err)
	}

	resp, err := client.ListResources(ctx, "")
	if err != nil {
		return nil, status.Errorf(codes.Internal, "ListResources failed: %v", err)
	}

	data, err := json.Marshal(resp.Resources)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to marshal resources: %v", err)
	}

	return &pb.MCPResourcesResponse{Resources: data}, nil
}

// MCPReadResource reads a specific MCP resource by URI.
func (h *Handler) MCPReadResource(ctx context.Context, req *pb.MCPResourceRequest) (*pb.MCPResourceResponse, error) {
	client, err := goumcp.Select(req.SessionId)
	if err != nil {
		return nil, status.Errorf(codes.NotFound, "MCP client not found: %v", err)
	}

	resp, err := client.ReadResource(ctx, req.Uri)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "ReadResource failed: %v", err)
	}

	data, err := json.Marshal(resp.Contents)
	if err != nil {
		return nil, status.Errorf(codes.Internal, "failed to marshal contents: %v", err)
	}

	return &pb.MCPResourceResponse{Contents: data}, nil
}