From 37aa5e04be815e213ffa439e4ff6c3a2662710a5 Mon Sep 17 00:00:00 2001 From: Bryan Frimin Date: Fri, 7 Nov 2025 09:51:18 +0100 Subject: [PATCH] Add vendor list and create tools Signed-off-by: Bryan Frimin --- pkg/server/api/mcp/mcputils/mcputils.go | 75 +++++ pkg/server/api/mcp/v1/add_vendor.go | 65 ++--- pkg/server/api/mcp/v1/list_organizations.go | 61 ++++ pkg/server/api/mcp/v1/list_vendors.go | 81 +++--- pkg/server/api/mcp/v1/middleware.go | 114 ++++++++ pkg/server/api/mcp/v1/types.go | 66 +++++ pkg/server/api/mcp/v1/types/cursorkey.go | 28 ++ pkg/server/api/mcp/v1/types/order_by.go | 24 ++ pkg/server/api/mcp/v1/types/organization.go | 24 ++ pkg/server/api/mcp/v1/types/types.go | 33 +++ pkg/server/api/mcp/v1/types/vendor.go | 298 ++++++++++++++++++++ pkg/server/api/mcp/v1/v1_handler.go | 238 +++++----------- 12 files changed, 839 insertions(+), 268 deletions(-) create mode 100644 pkg/server/api/mcp/mcputils/mcputils.go create mode 100644 pkg/server/api/mcp/v1/list_organizations.go create mode 100644 pkg/server/api/mcp/v1/middleware.go create mode 100644 pkg/server/api/mcp/v1/types.go create mode 100644 pkg/server/api/mcp/v1/types/cursorkey.go create mode 100644 pkg/server/api/mcp/v1/types/order_by.go create mode 100644 pkg/server/api/mcp/v1/types/organization.go create mode 100644 pkg/server/api/mcp/v1/types/types.go create mode 100644 pkg/server/api/mcp/v1/types/vendor.go diff --git a/pkg/server/api/mcp/mcputils/mcputils.go b/pkg/server/api/mcp/mcputils/mcputils.go new file mode 100644 index 000000000..5ca77c579 --- /dev/null +++ b/pkg/server/api/mcp/mcputils/mcputils.go @@ -0,0 +1,75 @@ +// Copyright (c) 2025 Probo Inc . +// +// Permission to use, copy, modify, and/or distribute this software for any +// purpose with or without fee is hereby granted, provided that the above +// copyright notice and this permission notice appear in all copies. +// +// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH +// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY +// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT, +// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM +// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR +// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR +// PERFORMANCE OF THIS SOFTWARE. + +package mcputils + +import ( + "context" + "fmt" + "time" + + "github.com/modelcontextprotocol/go-sdk/mcp" + "go.gearno.de/kit/log" +) + +func LoggingMiddleware(logger *log.Logger) func(mcp.MethodHandler) mcp.MethodHandler { + return func(next mcp.MethodHandler) mcp.MethodHandler { + return func(ctx context.Context, method string, req mcp.Request) (mcp.Result, error) { + sessionID := req.GetSession().ID() + + logger.InfoCtx(ctx, fmt.Sprintf("mcp %q method started", method), + log.String("method", method), + log.String("session_id", sessionID), + log.Bool("has_params", req.GetParams() != nil), + ) + + if ctr, ok := req.(*mcp.CallToolRequest); ok { + logger.InfoCtx(ctx, fmt.Sprintf("calling %q tool", ctr.Params.Name), + log.String("tool_name", ctr.Params.Name), + log.String("session_id", sessionID), + ) + } + + start := time.Now() + result, err := next(ctx, method, req) + duration := time.Since(start) + + if err != nil { + logger.ErrorCtx(ctx, fmt.Sprintf("mcp %q method failed", method), + log.String("method", method), + log.String("session_id", sessionID), + log.Int64("duration_ms", duration.Milliseconds()), + log.Error(err), + ) + } else { + + logger.InfoCtx(ctx, fmt.Sprintf("mcp %q method completed", method), + log.String("method", method), + log.String("session_id", sessionID), + log.Int64("duration_ms", duration.Milliseconds()), + log.Bool("has_result", result != nil), + ) + + if ctr, ok := result.(*mcp.CallToolResult); ok { + logger.InfoCtx(ctx, "tool call result", + log.String("session_id", sessionID), + log.Bool("is_error", ctr.IsError), + ) + } + } + + return result, err + } + } +} diff --git a/pkg/server/api/mcp/v1/add_vendor.go b/pkg/server/api/mcp/v1/add_vendor.go index a427b4c70..dce20fc7c 100644 --- a/pkg/server/api/mcp/v1/add_vendor.go +++ b/pkg/server/api/mcp/v1/add_vendor.go @@ -4,51 +4,34 @@ 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" + "go.probo.inc/probo/pkg/probo" + "go.probo.inc/probo/pkg/server/api/mcp/v1/types" ) -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 - } +var ( + AddVendorTool = &mcp.Tool{ + Name: "addVendor", + Title: "Add Vendor", + Description: "Add a new vendor to the organization", + Annotations: &mcp.ToolAnnotations{ReadOnlyHint: false}, + InputSchema: types.AddVendorInputSchema, + OutputSchema: types.AddVendorOutputSchema, } ) func (r *resolver) AddVendor( ctx context.Context, req *mcp.CallToolRequest, - args *addVendorArgs, -) (*mcp.CallToolResult, *addVendorResult, error) { - vendor, err := r.proboSvc.Vendors.Create( + args types.AddVendorInput, +) (*mcp.CallToolResult, types.AddVendorOutput, error) { + tenantID := args.OrganizationID.TenantID() + svc := r.ProboService(ctx, tenantID) + + vendor, err := svc.Vendors.Create( ctx, probo.CreateVendorRequest{ - OrganizationID: r.organizationID, + OrganizationID: args.OrganizationID, Name: args.Name, Description: args.Description, HeadquarterAddress: args.HeadquarterAddress, @@ -70,18 +53,8 @@ func (r *resolver) AddVendor( }, ) if err != nil { - return nil, nil, fmt.Errorf("failed to list vendors: %w", err) + return nil, types.AddVendorOutput{}, fmt.Errorf("failed to create vendor: %w", err) } - result := &addVendorResult{ - Result: struct { - Name string - ID string - }{ - Name: vendor.Name, - ID: vendor.ID.String(), - }, - } - - return nil, result, nil + return nil, types.NewAddVendorOutput(vendor), nil } diff --git a/pkg/server/api/mcp/v1/list_organizations.go b/pkg/server/api/mcp/v1/list_organizations.go new file mode 100644 index 000000000..6bc2f2e33 --- /dev/null +++ b/pkg/server/api/mcp/v1/list_organizations.go @@ -0,0 +1,61 @@ +// Copyright (c) 2025 Probo Inc . +// +// Permission to use, copy, modify, and/or distribute this software for any +// purpose with or without fee is hereby granted, provided that the above +// copyright notice and this permission notice appear in all copies. +// +// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH +// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY +// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT, +// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM +// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR +// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR +// PERFORMANCE OF THIS SOFTWARE. + +package v1 + +import ( + "context" + "fmt" + + "github.com/google/jsonschema-go/jsonschema" + "github.com/modelcontextprotocol/go-sdk/mcp" + "go.probo.inc/probo/pkg/server/api/mcp/v1/types" +) + +var ( + ListOrganizationsTool = &mcp.Tool{ + Name: "listOrganizations", + Title: "List Organizations", + Description: "List all organizations the user has access to", + Annotations: &mcp.ToolAnnotations{ReadOnlyHint: true}, + InputSchema: &jsonschema.Schema{ + Type: "object", + Properties: map[string]*jsonschema.Schema{ + "organizationID": {Type: "string"}, + }, + }, + } +) + +func (r *resolver) ListOrganizations( + ctx context.Context, + req *mcp.CallToolRequest, + _ types.ListOrganizationsInput, +) (*mcp.CallToolResult, types.ListOrganizationsOutput, error) { + mcpCtx := MCPContextFromContext(ctx) + organizations, err := r.authzSvc.GetAllUserOrganizations(ctx, mcpCtx.UserID) + if err != nil { + return nil, types.ListOrganizationsOutput{}, fmt.Errorf("failed to list organizations: %w", err) + } + + result := types.ListOrganizationsOutput{ + Organizations: make([]types.Organization, 0, len(organizations)), + } + + for _, org := range organizations { + result.Organizations = append(result.Organizations, types.NewOrganization(org)) + } + + return nil, result, nil +} diff --git a/pkg/server/api/mcp/v1/list_vendors.go b/pkg/server/api/mcp/v1/list_vendors.go index c6cb7c2cc..ec899bec1 100644 --- a/pkg/server/api/mcp/v1/list_vendors.go +++ b/pkg/server/api/mcp/v1/list_vendors.go @@ -4,67 +4,52 @@ import ( "context" "fmt" - "github.com/getprobo/probo/pkg/coredata" - "github.com/getprobo/probo/pkg/page" "github.com/modelcontextprotocol/go-sdk/mcp" + "go.probo.inc/probo/pkg/coredata" + "go.probo.inc/probo/pkg/page" + "go.probo.inc/probo/pkg/server/api/mcp/v1/types" ) -type ( - listVendorsArgs struct { - OrderField coredata.VendorOrderField - Cursor *page.CursorKey - Size int - } - - listVendorsResult struct { - NextCursor *string - Result []struct { - Name string - ID string - } +var ( + ListVendorsTool = &mcp.Tool{ + Name: "listVendors", + Title: "List Vendors", + Description: "List all vendors for the organization", + Annotations: &mcp.ToolAnnotations{ReadOnlyHint: true}, + InputSchema: types.ListVendorsInputSchema, + OutputSchema: types.ListVendorsOutputSchema, } ) func (r *resolver) ListVendors( ctx context.Context, req *mcp.CallToolRequest, - args *listVendorsArgs, -) (*mcp.CallToolResult, *listVendorsResult, error) { + args types.ListVendorsInput, +) (*mcp.CallToolResult, types.ListVendorsOutput, error) { + prb := r.ProboService(ctx, args.OrganizationID.TenantID()) - filter := coredata.NewVendorFilter(nil, nil) - cursor := page.NewCursor( - args.Size, - args.Cursor, - page.Head, - page.OrderBy[coredata.VendorOrderField]{ - Field: args.OrderField, + pageOrderBy := page.OrderBy[coredata.VendorOrderField]{ + Field: coredata.VendorOrderFieldCreatedAt, + Direction: page.OrderDirectionDesc, + } + if args.OrderBy != nil { + pageOrderBy = page.OrderBy[coredata.VendorOrderField]{ + Field: args.OrderBy.Field, Direction: page.OrderDirectionDesc, - }, - ) + } + } - vendors, err := r.proboSvc.Vendors.ListForOrganizationID(ctx, r.organizationID, cursor, filter) + cursor := types.NewCursor(args.Size, args.Cursor, pageOrderBy) + + var vendorFilter = coredata.NewVendorFilter(nil, nil) + if args.Filter != nil { + vendorFilter = coredata.NewVendorFilter(&args.Filter.SnapshotID, nil) + } + + page, err := prb.Vendors.ListForOrganizationID(ctx, args.OrganizationID, cursor, vendorFilter) if err != nil { - return nil, nil, fmt.Errorf("failed to list vendors: %w", err) + panic(fmt.Errorf("cannot list organization 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 + return nil, types.NewListVendorsOutput(page), nil } diff --git a/pkg/server/api/mcp/v1/middleware.go b/pkg/server/api/mcp/v1/middleware.go new file mode 100644 index 000000000..ee25f5b34 --- /dev/null +++ b/pkg/server/api/mcp/v1/middleware.go @@ -0,0 +1,114 @@ +// Copyright (c) 2025 Probo Inc . +// +// Permission to use, copy, modify, and/or distribute this software for any +// purpose with or without fee is hereby granted, provided that the above +// copyright notice and this permission notice appear in all copies. +// +// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH +// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY +// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT, +// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM +// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR +// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR +// PERFORMANCE OF THIS SOFTWARE. + +package v1 + +import ( + "net/http" + "strings" + + "go.gearno.de/kit/log" + "go.probo.inc/probo/pkg/auth" + "go.probo.inc/probo/pkg/authz" + "go.probo.inc/probo/pkg/gid" +) + +// WithMCPAuth wraps an HTTP handler with MCP authentication middleware +// It authenticates using API keys from the Authorization header +func WithMCPAuth( + logger *log.Logger, + authSvc *auth.Service, + authzSvc *authz.Service, + next http.Handler, +) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + ctx := r.Context() + correlationID := r.Header.Get("X-Request-ID") + if correlationID == "" { + correlationID = r.Header.Get("X-Correlation-ID") + } + + logger.InfoCtx(ctx, "MCP authentication attempt", + log.String("correlation_id", correlationID), + log.String("path", r.URL.Path), + ) + + // Extract API key from Authorization header + authHeader := r.Header.Get("Authorization") + if authHeader == "" { + logger.WarnCtx(ctx, "MCP auth: missing Authorization header", + log.String("correlation_id", correlationID), + ) + http.Error(w, "authentication required", http.StatusUnauthorized) + return + } + + // Expect "Bearer " format + parts := strings.SplitN(authHeader, " ", 2) + if len(parts) != 2 || strings.ToLower(parts[0]) != "bearer" { + logger.WarnCtx(ctx, "MCP auth: invalid Authorization header format", + log.String("correlation_id", correlationID), + ) + http.Error(w, "invalid authorization header", http.StatusUnauthorized) + return + } + + apiKeyToken := parts[1] + + // Validate the API key + user, userAPIKey, err := authSvc.ValidateUserAPIKey(ctx, apiKeyToken) + if err != nil { + logger.WarnCtx(ctx, "MCP auth: invalid API key", + log.Error(err), + log.String("correlation_id", correlationID), + ) + http.Error(w, "invalid api key", http.StatusUnauthorized) + return + } + + // Get organizations for this API key through the service layer + organizations, err := authzSvc.GetAllOrganizationsForUserAPIKeyId(ctx, userAPIKey.ID) + if err != nil { + logger.ErrorCtx(ctx, "MCP auth: failed to load API key organizations", + log.Error(err), + log.String("correlation_id", correlationID), + ) + http.Error(w, "internal server error", http.StatusInternalServerError) + return + } + + // Extract tenant IDs from organizations + tenantIDs := make([]gid.TenantID, 0, len(organizations)) + for _, org := range organizations { + tenantIDs = append(tenantIDs, org.ID.TenantID()) + } + + // Create MCP context with user and accessible tenants + mcpCtx := &MCPContext{ + UserID: user.ID, + TenantIDs: tenantIDs, + } + + ctx = ContextWithMCPContext(ctx, mcpCtx) + + logger.InfoCtx(ctx, "MCP authentication successful", + log.String("correlation_id", correlationID), + log.String("user_id", user.ID.String()), + log.String("api_key_id", userAPIKey.ID.String()), + log.Int("accessible_tenants", len(tenantIDs)), + ) + + next.ServeHTTP(w, r.WithContext(ctx)) + }) +} diff --git a/pkg/server/api/mcp/v1/types.go b/pkg/server/api/mcp/v1/types.go new file mode 100644 index 000000000..94445888f --- /dev/null +++ b/pkg/server/api/mcp/v1/types.go @@ -0,0 +1,66 @@ +// Copyright (c) 2025 Probo Inc . +// +// Permission to use, copy, modify, and/or distribute this software for any +// purpose with or without fee is hereby granted, provided that the above +// copyright notice and this permission notice appear in all copies. +// +// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH +// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY +// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT, +// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM +// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR +// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR +// PERFORMANCE OF THIS SOFTWARE. + +package v1 + +import ( + "context" + "time" + + "go.probo.inc/probo/pkg/gid" +) + +type ( + // Config holds configuration for the MCP server + Config struct { + // Version is the MCP server version + Version string + // RequestTimeout is the maximum duration for a request + RequestTimeout time.Duration + // MaxRequestSize is the maximum size of a request body in bytes + MaxRequestSize int64 + } + + // MCPContext holds the authenticated context for MCP requests + MCPContext struct { + UserID gid.GID + TenantIDs []gid.TenantID + } + + ctxKey struct{ name string } +) + +var ( + mcpContextKey = &ctxKey{name: "mcp_context"} +) + +// MCPContextFromContext extracts the MCP context from the request context +func MCPContextFromContext(ctx context.Context) *MCPContext { + mcpCtx, _ := ctx.Value(mcpContextKey).(*MCPContext) + return mcpCtx +} + +// ContextWithMCPContext adds the MCP context to the request context +func ContextWithMCPContext(ctx context.Context, mcpCtx *MCPContext) context.Context { + return context.WithValue(ctx, mcpContextKey, mcpCtx) +} + +// DefaultConfig returns a default MCP configuration +func DefaultConfig() Config { + return Config{ + Version: "1.0.0", + RequestTimeout: 30 * time.Second, + MaxRequestSize: 10 * 1024 * 1024, // 10MB + } +} diff --git a/pkg/server/api/mcp/v1/types/cursorkey.go b/pkg/server/api/mcp/v1/types/cursorkey.go new file mode 100644 index 000000000..d5a218c01 --- /dev/null +++ b/pkg/server/api/mcp/v1/types/cursorkey.go @@ -0,0 +1,28 @@ +// Copyright (c) 2025 Probo Inc . +// +// Permission to use, copy, modify, and/or distribute this software for any +// purpose with or without fee is hereby granted, provided that the above +// copyright notice and this permission notice appear in all copies. +// +// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH +// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY +// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT, +// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM +// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR +// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR +// PERFORMANCE OF THIS SOFTWARE. + +package types + +import ( + "go.probo.inc/probo/pkg/page" + "go.probo.inc/probo/pkg/server/gqlutils/types/cursor" +) + +func NewCursor[O page.OrderField]( + first *int, + after *page.CursorKey, + orderBy page.OrderBy[O], +) *page.Cursor[O] { + return cursor.NewCursor(first, after, nil, nil, orderBy) +} diff --git a/pkg/server/api/mcp/v1/types/order_by.go b/pkg/server/api/mcp/v1/types/order_by.go new file mode 100644 index 000000000..0debc3b5b --- /dev/null +++ b/pkg/server/api/mcp/v1/types/order_by.go @@ -0,0 +1,24 @@ +// Copyright (c) 2025 Probo Inc . +// +// Permission to use, copy, modify, and/or distribute this software for any +// purpose with or without fee is hereby granted, provided that the above +// copyright notice and this permission notice appear in all copies. +// +// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH +// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY +// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT, +// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM +// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR +// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR +// PERFORMANCE OF THIS SOFTWARE. + +package types + +import "go.probo.inc/probo/pkg/page" + +type ( + OrderBy[T page.OrderField] struct { + Field T + Direction page.OrderDirection + } +) diff --git a/pkg/server/api/mcp/v1/types/organization.go b/pkg/server/api/mcp/v1/types/organization.go new file mode 100644 index 000000000..f5ebf0531 --- /dev/null +++ b/pkg/server/api/mcp/v1/types/organization.go @@ -0,0 +1,24 @@ +// Copyright (c) 2025 Probo Inc . +// +// Permission to use, copy, modify, and/or distribute this software for any +// purpose with or without fee is hereby granted, provided that the above +// copyright notice and this permission notice appear in all copies. +// +// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH +// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY +// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT, +// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM +// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR +// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR +// PERFORMANCE OF THIS SOFTWARE. + +package types + +// ListOrganizationsInput defines the input parameters for listOrganizations tool +// Empty struct since no arguments are required +type ListOrganizationsInput struct{} + +// ListOrganizationsOutput defines the output structure for listOrganizations tool +type ListOrganizationsOutput struct { + Organizations []Organization `json:"organizations" jsonschema:"list of organizations the user has access to"` +} diff --git a/pkg/server/api/mcp/v1/types/types.go b/pkg/server/api/mcp/v1/types/types.go new file mode 100644 index 000000000..d9f864aa6 --- /dev/null +++ b/pkg/server/api/mcp/v1/types/types.go @@ -0,0 +1,33 @@ +// Copyright (c) 2025 Probo Inc . +// +// Permission to use, copy, modify, and/or distribute this software for any +// purpose with or without fee is hereby granted, provided that the above +// copyright notice and this permission notice appear in all copies. +// +// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH +// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY +// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT, +// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM +// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR +// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR +// PERFORMANCE OF THIS SOFTWARE. + +package types + +import "go.probo.inc/probo/pkg/coredata" + +// Organization represents a single organization in MCP responses +type Organization struct { + Name string `json:"name" jsonschema:"the organization name"` + ID string `json:"id" jsonschema:"the organization ID"` + TenantID string `json:"tenantID" jsonschema:"the tenant ID this organization belongs to"` +} + +// NewOrganization converts a coredata.Organization to an MCP Organization +func NewOrganization(o *coredata.Organization) Organization { + return Organization{ + Name: o.Name, + ID: o.ID.String(), + TenantID: o.ID.TenantID().String(), + } +} diff --git a/pkg/server/api/mcp/v1/types/vendor.go b/pkg/server/api/mcp/v1/types/vendor.go new file mode 100644 index 000000000..de4822a0b --- /dev/null +++ b/pkg/server/api/mcp/v1/types/vendor.go @@ -0,0 +1,298 @@ +// Copyright (c) 2025 Probo Inc . +// +// Permission to use, copy, modify, and/or distribute this software for any +// purpose with or without fee is hereby granted, provided that the above +// copyright notice and this permission notice appear in all copies. +// +// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH +// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY +// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT, +// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM +// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR +// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR +// PERFORMANCE OF THIS SOFTWARE. + +package types + +import ( + "time" + + "github.com/google/jsonschema-go/jsonschema" + "go.probo.inc/probo/pkg/coredata" + "go.probo.inc/probo/pkg/gid" + "go.probo.inc/probo/pkg/page" +) + +type ( + VendorOrderBy OrderBy[coredata.VendorOrderField] + + VendorFilter struct { + SnapshotID *gid.GID `json:"snapshot_id"` + } + + ListVendorsInput struct { + OrganizationID gid.GID `json:"organization_id"` + Filter *VendorFilter `json:"filter"` + OrderBy *VendorOrderBy `json:"order_field"` + Cursor *page.CursorKey `json:"cursor"` + Size *int `json:"size"` + } + + ListVendorsOutput struct { + NextCursor *string `json:"next_cursor"` + Vendors []Vendor `json:"vendors"` + } + + AddVendorInput struct { + OrganizationID gid.GID `json:"organization_id"` + Name string `json:"name"` + Description *string `json:"description"` + HeadquarterAddress *string `json:"headquarter_address"` + LegalName *string `json:"legal_name"` + WebsiteURL *string `json:"website_url"` + Category *coredata.VendorCategory `json:"category"` + PrivacyPolicyURL *string `json:"privacy_policy_url"` + ServiceLevelAgreementURL *string `json:"service_level_agreement_url"` + DataProcessingAgreementURL *string `json:"data_processing_agreement_url"` + BusinessAssociateAgreementURL *string `json:"business_associate_agreement_url"` + SubprocessorsListURL *string `json:"subprocessors_list_url"` + Certifications []string `json:"certifications"` + Countries []coredata.CountryCode `json:"countries"` + SecurityPageURL *string `json:"security_page_url"` + TrustPageURL *string `json:"trust_page_url"` + TermsOfServiceURL *string `json:"terms_of_service_url"` + StatusPageURL *string `json:"status_page_url"` + BusinessOwnerID *gid.GID `json:"business_owner_id"` + SecurityOwnerID *gid.GID `json:"security_owner_id"` + } + + AddVendorOutput struct { + Vendor Vendor `json:"vendor" jsonschema:"the created vendor"` + } + + Vendor struct { + ID gid.GID `json:"id"` + OrganizationID gid.GID `json:"organization_id"` + Name string `json:"name"` + Description *string `json:"description"` + Category coredata.VendorCategory `json:"category"` + HeadquarterAddress *string `json:"headquarter_address"` + LegalName *string `json:"legal_name"` + WebsiteURL *string `json:"website_url"` + PrivacyPolicyURL *string `json:"privacy_policy_url"` + ServiceLevelAgreementURL *string `json:"service_level_agreement_url"` + DataProcessingAgreementURL *string `json:"data_processing_agreement_url"` + BusinessAssociateAgreementURL *string `json:"business_associate_agreement_url"` + SubprocessorsListURL *string `json:"subprocessors_list_url"` + Certifications []string `json:"certifications"` + Countries []coredata.CountryCode `json:"countries"` + BusinessOwnerID *gid.GID `json:"business_owner_id,omitempty"` + SecurityOwnerID *gid.GID `json:"security_owner_id,omitempty"` + StatusPageURL *string `json:"status_page_url,omitempty"` + TermsOfServiceURL *string `json:"terms_of_service_url,omitempty"` + SecurityPageURL *string `json:"security_page_url,omitempty"` + TrustPageURL *string `json:"trust_page_url,omitempty"` + ShowOnTrustCenter bool `json:"show_on_trust_center,omitempty"` + SnapshotID *gid.GID `json:"snapshot_id,omitempty"` + SourceID *gid.GID `json:"source_id,omitempty"` + CreatedAt time.Time `json:"created_at"` + UpdatedAt time.Time `json:"updated_at"` + } +) + +var ( + ListVendorsInputSchema = &jsonschema.Schema{ + Type: "object", + Required: []string{"organizationID"}, + Properties: map[string]*jsonschema.Schema{ + "organizationID": {Type: "string"}, + "filter": { + Type: "object", + Properties: map[string]*jsonschema.Schema{ + "snapshotID": {Type: "string"}, + }, + }, + "orderBy": { + Types: []string{"object", "null"}, + Properties: map[string]*jsonschema.Schema{ + "field": {Type: "string", Enum: []any{"CREATED_AT"}}, + "direction": OrderByDirectionSchema, + }, + }, + "cursor": {Types: []string{"string", "null"}}, + "size": {Types: []string{"integer", "null"}}, + }, + } + + OrderByDirectionSchema = &jsonschema.Schema{ + Type: "string", + Enum: []any{"ASC", "DESC"}, + } + + NullableStringSchema = &jsonschema.Schema{ + Types: []string{"string", "null"}, + } + + VendorCategorySchema = &jsonschema.Schema{ + Type: "string", + Enum: []any{ + "ANALYTICS", + "CLOUD_MONITORING", + "CLOUD_PROVIDER", + "COLLABORATION", + "CUSTOMER_SUPPORT", + "DATA_STORAGE_AND_PROCESSING", + "DOCUMENT_MANAGEMENT", + "EMPLOYEE_MANAGEMENT", + "ENGINEERING", + "FINANCE", + "IDENTITY_PROVIDER", + "IT", + "MARKETING", + "OFFICE_OPERATIONS", + "OTHER", + "PASSWORD_MANAGEMENT", + "PRODUCT_AND_DESIGN", + "PROFESSIONAL_SERVICES", + "RECRUITING", + "SALES", + "SECURITY", + "VERSION_CONTROL", + }, + } + + VendorSchema = &jsonschema.Schema{ + Type: "object", + Properties: map[string]*jsonschema.Schema{ + "id": {Type: "string"}, + "name": {Type: "string"}, + "organization_id": {Type: "string"}, + "description": NullableStringSchema, + "category": VendorCategorySchema, + "headquarter_address": NullableStringSchema, + "legalName": NullableStringSchema, + "website_url": NullableStringSchema, + "privacy_policy_url": NullableStringSchema, + "service_level_agreement_url": NullableStringSchema, + "data_processing_agreement_url": NullableStringSchema, + "business_associate_agreement_url": NullableStringSchema, + "subprocessors_list_url": NullableStringSchema, + "certifications": {Types: []string{"array", "null"}, Items: &jsonschema.Schema{Type: "string"}}, + "countries": {Types: []string{"array", "null"}, Items: &jsonschema.Schema{Type: "string", Enum: []any{"US", "CA", "GB", "DE", "FR", "IT", "ES", "NL", "BE", "CH", "AT", "SE", "NO", "DK", "FI", "EE", "LT", "LV", "PL", "CZ", "SK", "HU", "RO", "BG", "HR", "SI", "ME", "AL", "MK", "BA", "XK", "XA", "XZ"}}}, + "business_owner_id": NullableStringSchema, + "security_owner_id": NullableStringSchema, + "status_page_url": NullableStringSchema, + "terms_of_service_url": NullableStringSchema, + "security_page_url": NullableStringSchema, + "trust_page_url": NullableStringSchema, + "show_on_trust_center": {Type: "boolean"}, + "snapshot_id": NullableStringSchema, + "source_id": NullableStringSchema, + "created_at": {Type: "string"}, + "updated_at": {Type: "string"}, + }, + } + + ListVendorsOutputSchema = &jsonschema.Schema{ + Type: "object", + Properties: map[string]*jsonschema.Schema{ + "next_cursor": {Types: []string{"string", "null"}}, + "vendors": { + Type: "array", + Items: VendorSchema, + }, + }, + } + + AddVendorInputSchema = &jsonschema.Schema{ + Type: "object", + Required: []string{"organization_id", "name"}, + Properties: map[string]*jsonschema.Schema{ + "organization_id": {Type: "string"}, + "name": {Type: "string"}, + "description": NullableStringSchema, + "headquarter_address": NullableStringSchema, + "legal_name": NullableStringSchema, + "website_url": NullableStringSchema, + "category": VendorCategorySchema, + "privacy_policy_url": NullableStringSchema, + "service_level_agreement_url": NullableStringSchema, + "data_processing_agreement_url": NullableStringSchema, + "business_associate_agreement_url": NullableStringSchema, + "subprocessors_list_url": NullableStringSchema, + "certifications": {Types: []string{"array", "null"}, Items: &jsonschema.Schema{Type: "string"}}, + "countries": {Types: []string{"array", "null"}, Items: &jsonschema.Schema{Type: "string", Enum: []any{"US", "CA", "GB", "DE", "FR", "IT", "ES", "NL", "BE", "CH", "AT", "SE", "NO", "DK", "FI", "EE", "LT", "LV", "PL", "CZ", "SK", "HU", "RO", "BG", "HR", "SI", "ME", "AL", "MK", "BA", "XK", "XA", "XZ"}}}, + "business_owner_id": NullableStringSchema, + "security_owner_id": NullableStringSchema, + "status_page_url": NullableStringSchema, + "terms_of_service_url": NullableStringSchema, + "security_page_url": NullableStringSchema, + "trust_page_url": NullableStringSchema, + "show_on_trust_center": {Type: "boolean"}, + "snapshot_id": NullableStringSchema, + "source_id": NullableStringSchema, + }, + } + + AddVendorOutputSchema = &jsonschema.Schema{ + Type: "object", + Properties: map[string]*jsonschema.Schema{ + "vendor": VendorSchema, + }, + } +) + +func NewVendor(v *coredata.Vendor) Vendor { + return Vendor{ + Name: v.Name, + ID: v.ID, + OrganizationID: v.OrganizationID, + Description: v.Description, + Category: v.Category, + HeadquarterAddress: v.HeadquarterAddress, + LegalName: v.LegalName, + WebsiteURL: v.WebsiteURL, + PrivacyPolicyURL: v.PrivacyPolicyURL, + ServiceLevelAgreementURL: v.ServiceLevelAgreementURL, + DataProcessingAgreementURL: v.DataProcessingAgreementURL, + BusinessAssociateAgreementURL: v.BusinessAssociateAgreementURL, + SubprocessorsListURL: v.SubprocessorsListURL, + Certifications: v.Certifications, + Countries: v.Countries, + BusinessOwnerID: v.BusinessOwnerID, + SecurityOwnerID: v.SecurityOwnerID, + StatusPageURL: v.StatusPageURL, + TermsOfServiceURL: v.TermsOfServiceURL, + SecurityPageURL: v.SecurityPageURL, + TrustPageURL: v.TrustPageURL, + ShowOnTrustCenter: v.ShowOnTrustCenter, + SnapshotID: v.SnapshotID, + SourceID: v.SourceID, + CreatedAt: v.CreatedAt, + UpdatedAt: v.UpdatedAt, + } +} + +func NewListVendorsOutput(vendorPage *page.Page[*coredata.Vendor, coredata.VendorOrderField]) ListVendorsOutput { + vendors := make([]Vendor, 0, len(vendorPage.Data)) + for _, v := range vendorPage.Data { + vendors = append(vendors, NewVendor(v)) + } + + var nextCursor *string + if len(vendorPage.Data) > 0 { + cursorKey := vendorPage.Data[len(vendorPage.Data)-1].CursorKey(vendorPage.Cursor.OrderBy.Field).String() + nextCursor = &cursorKey + } + + return ListVendorsOutput{ + NextCursor: nextCursor, + Vendors: vendors, + } +} + +func NewAddVendorOutput(v *coredata.Vendor) AddVendorOutput { + return AddVendorOutput{ + Vendor: NewVendor(v), + } +} diff --git a/pkg/server/api/mcp/v1/v1_handler.go b/pkg/server/api/mcp/v1/v1_handler.go index 7735c1614..0a6e39198 100644 --- a/pkg/server/api/mcp/v1/v1_handler.go +++ b/pkg/server/api/mcp/v1/v1_handler.go @@ -1,209 +1,99 @@ package v1 import ( - "encoding/json" + "context" + "fmt" "net/http" + "time" - "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" + "go.gearno.de/kit/log" + "go.probo.inc/probo/pkg/auth" + "go.probo.inc/probo/pkg/authz" + "go.probo.inc/probo/pkg/gid" + "go.probo.inc/probo/pkg/probo" + "go.probo.inc/probo/pkg/server/api/mcp/mcputils" ) type ( resolver struct { - proboSvc *probo.TenantService - organizationID gid.GID + proboSvc *probo.Service + authSvc *auth.Service + authzSvc *authz.Service + logger *log.Logger } ) -func NewMux(proboSvc *probo.Service) *chi.Mux { +func (r *resolver) ProboService(ctx context.Context, tenantID gid.TenantID) *probo.TenantService { + validateTenantAccess(ctx, tenantID) + return r.proboSvc.WithTenant(tenantID) +} + +func validateTenantAccess(ctx context.Context, tenantID gid.TenantID) { + mcpCtx := MCPContextFromContext(ctx) + if mcpCtx == nil { + panic(fmt.Errorf("authentication context not found")) + } + + for _, tid := range mcpCtx.TenantIDs { + if tid == tenantID { + return + } + } + + panic(fmt.Errorf("access denied: user does not have access to tenant %s", tenantID.String())) +} + +func NewMux(logger *log.Logger, proboSvc *probo.Service, authSvc *auth.Service, authzSvc *authz.Service, cfg Config) *chi.Mux { + logger = logger.Named("mcp.v1") + + logger.Info("initializing MCP server", + log.String("version", cfg.Version), + log.String("request_timeout", cfg.RequestTimeout.String()), + ) + server := mcp.NewServer( &mcp.Implementation{ Name: "probo", Title: "Probo", - Version: "1.0.0", // todo retrieve from build info + Version: cfg.Version, }, &mcp.ServerOptions{}, ) - tenantID, err := gid.ParseTenantID("lXdXZSh-AAE") - if err != nil { - panic(err) + server.AddReceivingMiddleware(mcputils.LoggingMiddleware(logger)) + + resolver := &resolver{ + proboSvc: proboSvc, + authSvc: authSvc, + authzSvc: authzSvc, + logger: logger, } - 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, - ) + mcp.AddTool(server, ListOrganizationsTool, resolver.ListOrganizations) + mcp.AddTool(server, ListVendorsTool, resolver.ListVendors) + mcp.AddTool(server, AddVendorTool, resolver.AddVendor) getServer := func(r *http.Request) *mcp.Server { return server } + eventStore := mcp.NewMemoryEventStore(nil) handler := mcp.NewStreamableHTTPHandler( getServer, - &mcp.StreamableHTTPOptions{Stateless: true}, + &mcp.StreamableHTTPOptions{ + Stateless: false, + SessionTimeout: 30 * time.Minute, + EventStore: eventStore, + Logger: nil, // TODO put logger here + }, ) + authHandler := WithMCPAuth(logger, authSvc, authzSvc, handler) + r := chi.NewMux() - r.Handle("/", handler) + r.Handle("/", authHandler) + + logger.Info("MCP server initialized successfully") return r }