From 1436d0db08b41be7b0db4e3a8b137667d5d7bb6f Mon Sep 17 00:00:00 2001 From: Bryan Frimin Date: Sun, 7 Sep 2025 22:13:45 +0200 Subject: [PATCH] Add simple MCP server Proof of concept of working MCP server for Probo. Currently the official MCP library does not support session that why the server is configured in stateless mode. It seams the input jsonschema is not used to perform any validation, so we should figuring out how to validate the input properly. Signed-off-by: Bryan Frimin --- go.mod | 3 + go.sum | 6 + pkg/server/api/api.go | 7 + pkg/server/api/mcp/v1/add_vendor.go | 87 +++++++++++ pkg/server/api/mcp/v1/list_vendors.go | 70 +++++++++ pkg/server/api/mcp/v1/v1_handler.go | 209 ++++++++++++++++++++++++++ 6 files changed, 382 insertions(+) create mode 100644 pkg/server/api/mcp/v1/add_vendor.go create mode 100644 pkg/server/api/mcp/v1/list_vendors.go create mode 100644 pkg/server/api/mcp/v1/v1_handler.go diff --git a/go.mod b/go.mod index e3bd92f07..699b5e29d 100644 --- a/go.mod +++ b/go.mod @@ -16,8 +16,10 @@ require ( github.com/crewjam/saml v0.5.1 github.com/go-chi/chi/v5 v5.2.2 github.com/go-chi/cors v1.2.2 + github.com/google/jsonschema-go v0.2.1-0.20250825175020-748c325cec76 github.com/jackc/pgx/v5 v5.7.5 github.com/jhillyerd/enmime v1.3.0 + github.com/modelcontextprotocol/go-sdk v0.4.0 github.com/openai/openai-go v1.8.2 github.com/pdfcpu/pdfcpu v0.11.0 github.com/prometheus/client_golang v1.22.0 @@ -99,6 +101,7 @@ require ( github.com/xrash/smetrics v0.0.0-20240521201337-686a1a2994c1 // indirect github.com/xuri/efp v0.0.1 // indirect github.com/xuri/nfp v0.0.1 // indirect + github.com/yosida95/uritemplate/v3 v3.0.2 // indirect go.gearno.de/x/panicf v0.1.1 // indirect go.opentelemetry.io/auto/sdk v1.1.0 // indirect go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.37.0 // indirect diff --git a/go.sum b/go.sum index bc4f4db8c..fc27ea5e6 100644 --- a/go.sum +++ b/go.sum @@ -92,6 +92,8 @@ github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps= github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= +github.com/google/jsonschema-go v0.2.1-0.20250825175020-748c325cec76 h1:mBlBwtDebdDYr+zdop8N62a44g+Nbv7o2KjWyS1deR4= +github.com/google/jsonschema-go v0.2.1-0.20250825175020-748c325cec76/go.mod h1:r5quNTdLOYEz95Ru18zA0ydNbBuYoo9tgaYcxEYhJVE= github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg= @@ -140,6 +142,8 @@ github.com/mattermost/xml-roundtrip-validator v0.1.0/go.mod h1:qccnGMcpgwcNaBnxq github.com/mattn/go-runewidth v0.0.9/go.mod h1:H031xJmbD/WCDINGzjvQ9THkh0rPKHF+m2gUSrubnMI= github.com/mattn/go-runewidth v0.0.16 h1:E5ScNMtiwvlvB5paMFdw9p4kSQzbXFikJ5SQO6TULQc= github.com/mattn/go-runewidth v0.0.16/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w= +github.com/modelcontextprotocol/go-sdk v0.4.0 h1:RJ6kFlneHqzTKPzlQqiunrz9nbudSZcYLmLHLsokfoU= +github.com/modelcontextprotocol/go-sdk v0.4.0/go.mod h1:whv0wHnsTphwq7CTiKYHkLtwLC06WMoY2KpO+RB9yXQ= github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 h1:C3w9PqII01/Oq1c1nUAm88MOHcQC9l5mIlSMApZMrHA= github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822/go.mod h1:+n7T8mK8HuQTcFwEeznm/DIxMOiR9yIdICNftLE1DvQ= github.com/olekukonko/tablewriter v0.0.5 h1:P2Ga83D34wi1o9J6Wh1mRuqd4mF/x/lgBS7N7AbDhec= @@ -217,6 +221,8 @@ github.com/xuri/excelize/v2 v2.9.1 h1:VdSGk+rraGmgLHGFaGG9/9IWu1nj4ufjJ7uwMDtj8Q github.com/xuri/excelize/v2 v2.9.1/go.mod h1:x7L6pKz2dvo9ejrRuD8Lnl98z4JLt0TGAwjhW+EiP8s= github.com/xuri/nfp v0.0.1 h1:MDamSGatIvp8uOmDP8FnmjuQpu90NzdJxo7242ANR9Q= github.com/xuri/nfp v0.0.1/go.mod h1:WwHg+CVyzlv/TX9xqBFXEZAuxOPxn2k1GNHwG41IIUQ= +github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zIM+UJPGz4= +github.com/yosida95/uritemplate/v3 v3.0.2/go.mod h1:ILOh0sOhIJR3+L/8afwt/kE++YT040gmv5BQTMR2HP4= github.com/yuin/goldmark v1.7.12 h1:YwGP/rrea2/CnCtUHgjuolG/PnMxdQtPMO5PvaE2/nY= github.com/yuin/goldmark v1.7.12/go.mod h1:ip/1k0VRfGynBgxOz0yCqHrbZXhcjxyuS66Brc7iBKg= go.gearno.de/crypto/uuid v0.1.0 h1:94BYg7GYItJ6yYZ1GJayb3VYhI9/FjxuR1nFaduR4hE= diff --git a/pkg/server/api/api.go b/pkg/server/api/api.go index 76c760535..fa8460fa4 100644 --- a/pkg/server/api/api.go +++ b/pkg/server/api/api.go @@ -30,6 +30,7 @@ import ( "go.probo.inc/probo/pkg/probo" "go.probo.inc/probo/pkg/saferedirect" console_v1 "go.probo.inc/probo/pkg/server/api/console/v1" + mcp_v1 "go.probo.inc/probo/pkg/server/api/mcp/v1" trust_v1 "go.probo.inc/probo/pkg/server/api/trust/v1" "go.probo.inc/probo/pkg/trust" ) @@ -207,5 +208,11 @@ func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) { // Mount the trust API with authentication router.Mount("/trust/v1", s.trustAPIHandler) + // Mount the MCP API - use Route instead of Mount to preserve path for handler + router.Mount( + "/mcp/v1", + mcp_v1.NewMux(s.cfg.Probo), + ) + router.ServeHTTP(w, r) } diff --git a/pkg/server/api/mcp/v1/add_vendor.go b/pkg/server/api/mcp/v1/add_vendor.go new file mode 100644 index 000000000..a427b4c70 --- /dev/null +++ b/pkg/server/api/mcp/v1/add_vendor.go @@ -0,0 +1,87 @@ +package v1 + +import ( + "context" + "fmt" + + "github.com/getprobo/probo/pkg/coredata" + "github.com/getprobo/probo/pkg/gid" + "github.com/getprobo/probo/pkg/probo" + "github.com/modelcontextprotocol/go-sdk/mcp" +) + +type ( + addVendorArgs struct { + Name string + Description *string + HeadquarterAddress *string + LegalName *string + WebsiteURL *string + Category *coredata.VendorCategory + PrivacyPolicyURL *string + ServiceLevelAgreementURL *string + DataProcessingAgreementURL *string + BusinessAssociateAgreementURL *string + SubprocessorsListURL *string + Certifications []string + SecurityPageURL *string + TrustPageURL *string + TermsOfServiceURL *string + StatusPageURL *string + BusinessOwnerID *gid.GID + SecurityOwnerID *gid.GID + } + + addVendorResult struct { + Result struct { + Name string + ID string + } + } +) + +func (r *resolver) AddVendor( + ctx context.Context, + req *mcp.CallToolRequest, + args *addVendorArgs, +) (*mcp.CallToolResult, *addVendorResult, error) { + vendor, err := r.proboSvc.Vendors.Create( + ctx, + probo.CreateVendorRequest{ + OrganizationID: r.organizationID, + Name: args.Name, + Description: args.Description, + HeadquarterAddress: args.HeadquarterAddress, + LegalName: args.LegalName, + WebsiteURL: args.WebsiteURL, + Category: args.Category, + PrivacyPolicyURL: args.PrivacyPolicyURL, + ServiceLevelAgreementURL: args.ServiceLevelAgreementURL, + DataProcessingAgreementURL: args.DataProcessingAgreementURL, + BusinessAssociateAgreementURL: args.BusinessAssociateAgreementURL, + SubprocessorsListURL: args.SubprocessorsListURL, + Certifications: args.Certifications, + SecurityPageURL: args.SecurityPageURL, + TrustPageURL: args.TrustPageURL, + TermsOfServiceURL: args.TermsOfServiceURL, + StatusPageURL: args.StatusPageURL, + BusinessOwnerID: args.BusinessOwnerID, + SecurityOwnerID: args.SecurityOwnerID, + }, + ) + if err != nil { + return nil, nil, fmt.Errorf("failed to list vendors: %w", err) + } + + result := &addVendorResult{ + Result: struct { + Name string + ID string + }{ + Name: vendor.Name, + ID: vendor.ID.String(), + }, + } + + return nil, result, nil +} diff --git a/pkg/server/api/mcp/v1/list_vendors.go b/pkg/server/api/mcp/v1/list_vendors.go new file mode 100644 index 000000000..c6cb7c2cc --- /dev/null +++ b/pkg/server/api/mcp/v1/list_vendors.go @@ -0,0 +1,70 @@ +package v1 + +import ( + "context" + "fmt" + + "github.com/getprobo/probo/pkg/coredata" + "github.com/getprobo/probo/pkg/page" + "github.com/modelcontextprotocol/go-sdk/mcp" +) + +type ( + listVendorsArgs struct { + OrderField coredata.VendorOrderField + Cursor *page.CursorKey + Size int + } + + listVendorsResult struct { + NextCursor *string + Result []struct { + Name string + ID string + } + } +) + +func (r *resolver) ListVendors( + ctx context.Context, + req *mcp.CallToolRequest, + args *listVendorsArgs, +) (*mcp.CallToolResult, *listVendorsResult, error) { + + filter := coredata.NewVendorFilter(nil, nil) + cursor := page.NewCursor( + args.Size, + args.Cursor, + page.Head, + page.OrderBy[coredata.VendorOrderField]{ + Field: args.OrderField, + Direction: page.OrderDirectionDesc, + }, + ) + + vendors, err := r.proboSvc.Vendors.ListForOrganizationID(ctx, r.organizationID, cursor, filter) + if err != nil { + return nil, nil, fmt.Errorf("failed to list vendors: %w", err) + } + + result := &listVendorsResult{} + if len(vendors.Data) > 0 { + nextCursorKey := vendors.Data[len(vendors.Data)-1].CursorKey(args.OrderField).String() + result.NextCursor = &nextCursorKey + } + + for _, vendor := range vendors.Data { + result.Result = append( + result.Result, + struct { + Name string + ID string + }{ + Name: vendor.Name, + ID: vendor.ID.String(), + }, + ) + } + + return nil, result, nil +} diff --git a/pkg/server/api/mcp/v1/v1_handler.go b/pkg/server/api/mcp/v1/v1_handler.go new file mode 100644 index 000000000..7735c1614 --- /dev/null +++ b/pkg/server/api/mcp/v1/v1_handler.go @@ -0,0 +1,209 @@ +package v1 + +import ( + "encoding/json" + "net/http" + + "github.com/getprobo/probo/pkg/gid" + "github.com/getprobo/probo/pkg/probo" + "github.com/go-chi/chi/v5" + "github.com/google/jsonschema-go/jsonschema" + "github.com/modelcontextprotocol/go-sdk/mcp" +) + +type ( + resolver struct { + proboSvc *probo.TenantService + organizationID gid.GID + } +) + +func NewMux(proboSvc *probo.Service) *chi.Mux { + server := mcp.NewServer( + &mcp.Implementation{ + Name: "probo", + Title: "Probo", + Version: "1.0.0", // todo retrieve from build info + }, + &mcp.ServerOptions{}, + ) + + tenantID, err := gid.ParseTenantID("lXdXZSh-AAE") + if err != nil { + panic(err) + } + + organizationID, err := gid.ParseGID("lXdXZSh-AAEAAAAAAZfLJi38a0AGbu37") + if err != nil { + panic(err) + } + + resolver := &resolver{proboSvc: proboSvc.WithTenant(tenantID), organizationID: organizationID} + + mcp.AddTool( + server, + &mcp.Tool{ + Title: "List Vendors", + Description: "List all vendors for the organization", + Name: "listVendors", + Annotations: &mcp.ToolAnnotations{ + Title: "List Vendors", + ReadOnlyHint: true, + }, + InputSchema: &jsonschema.Schema{ + Type: "object", + Properties: map[string]*jsonschema.Schema{ + "orderField": { + Type: "string", + Default: json.RawMessage(`"NAME"`), + Enum: []any{ + "NAME", + "CREATED_AT", + "UPDATED_AT", + }, + }, + "cursor": { + Type: "string", + }, + "size": { + Type: "integer", + Minimum: jsonschema.Ptr(float64(1)), + Maximum: jsonschema.Ptr(float64(1000)), + Default: json.RawMessage(`100`), + }, + }, + }, + OutputSchema: &jsonschema.Schema{ + Type: "object", + Properties: map[string]*jsonschema.Schema{ + "result": { + Type: "array", + Items: &jsonschema.Schema{ + Type: "object", + Required: []string{"name", "id"}, + Properties: map[string]*jsonschema.Schema{ + "name": { + Type: "string", + }, + "id": { + Type: "string", + }, + }, + }, + }, + }, + }, + }, + resolver.ListVendors, + ) + + mcp.AddTool( + server, + &mcp.Tool{ + Name: "addVendor", + Description: "Add a vendor", + InputSchema: &jsonschema.Schema{ + Type: "object", + Properties: map[string]*jsonschema.Schema{ + "name": { + Type: "string", + Required: []string{"name"}, + }, + "description": { + Type: "string", + }, + "headquarterAddress": { + Type: "string", + }, + "legalName": { + Type: "string", + }, + "websiteURL": { + Type: "string", + Format: "uri", + }, + "category": { + Type: "string", + }, + "privacyPolicyURL": { + Type: "string", + }, + "serviceLevelAgreementURL": { + Type: "string", + Format: "uri", + }, + "dataProcessingAgreementURL": { + Type: "string", + Format: "uri", + }, + "businessAssociateAgreementURL": { + Type: "string", + Format: "uri", + }, + "subprocessorsListURL": { + Type: "string", + Format: "uri", + }, + "certifications": { + Type: "array", + Items: &jsonschema.Schema{ + Type: "string", + }, + }, + "securityPageURL": { + Type: "string", + Format: "uri", + }, + "trustPageURL": { + Type: "string", + Format: "uri", + }, + "termsOfServiceURL": { + Type: "string", + Format: "uri", + }, + "statusPageURL": { + Type: "string", + Format: "uri", + }, + "businessOwnerID": { + Type: "string", + }, + "securityOwnerID": { + Type: "string", + }, + }, + }, + OutputSchema: &jsonschema.Schema{ + Type: "object", + Properties: map[string]*jsonschema.Schema{ + "result": { + Type: "object", + Required: []string{"name", "id"}, + Properties: map[string]*jsonschema.Schema{ + "name": { + Type: "string", + }, + "id": { + Type: "string", + }, + }, + }, + }, + }, + }, + resolver.AddVendor, + ) + + getServer := func(r *http.Request) *mcp.Server { return server } + + handler := mcp.NewStreamableHTTPHandler( + getServer, + &mcp.StreamableHTTPOptions{Stateless: true}, + ) + + r := chi.NewMux() + r.Handle("/", handler) + + return r +}