Add vendor list and create tools

Signed-off-by: Bryan Frimin <bryan@getprobo.com>

pkg/server/api/mcp/mcputils/mcputils.go

@@ -0,0 +1,75 @@
+// Copyright (c) 2025 Probo Inc <hello@getprobo.com>.
+//
+// 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
+		}
+	}
+}

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 (
+var (
-	addVendorArgs struct {
+	AddVendorTool = &mcp.Tool{
-		Name                          string
+		Name:         "addVendor",
-		Description                   *string
+		Title:        "Add Vendor",
-		HeadquarterAddress            *string
+		Description:  "Add a new vendor to the organization",
-		LegalName                     *string
+		Annotations:  &mcp.ToolAnnotations{ReadOnlyHint: false},
-		WebsiteURL                    *string
+		InputSchema:  types.AddVendorInputSchema,
-		Category                      *coredata.VendorCategory
+		OutputSchema: types.AddVendorOutputSchema,
-		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,
+	args types.AddVendorInput,
-) (*mcp.CallToolResult, *addVendorResult, error) {
+) (*mcp.CallToolResult, types.AddVendorOutput, error) {
-	vendor, err := r.proboSvc.Vendors.Create(
+	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{
+	return nil, types.NewAddVendorOutput(vendor), nil
-		Result: struct {
-			Name string
-			ID   string
-		}{
-			Name: vendor.Name,
-			ID:   vendor.ID.String(),
-		},
-	}
-
-	return nil, result, nil
 }

pkg/server/api/mcp/v1/list_organizations.go

@@ -0,0 +1,61 @@
+// Copyright (c) 2025 Probo Inc <hello@getprobo.com>.
+//
+// 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
+}

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 (
+var (
-	listVendorsArgs struct {
+	ListVendorsTool = &mcp.Tool{
-		OrderField coredata.VendorOrderField
+		Name:         "listVendors",
-		Cursor     *page.CursorKey
+		Title:        "List Vendors",
-		Size       int
+		Description:  "List all vendors for the organization",
-	}
+		Annotations:  &mcp.ToolAnnotations{ReadOnlyHint: true},
-
+		InputSchema:  types.ListVendorsInputSchema,
-	listVendorsResult struct {
+		OutputSchema: types.ListVendorsOutputSchema,
-		NextCursor *string
-		Result     []struct {
-			Name string
-			ID   string
-		}
 	}
 )
 
 func (r *resolver) ListVendors(
 	ctx context.Context,
 	req *mcp.CallToolRequest,
-	args *listVendorsArgs,
+	args types.ListVendorsInput,
-) (*mcp.CallToolResult, *listVendorsResult, error) {
+) (*mcp.CallToolResult, types.ListVendorsOutput, error) {
+	prb := r.ProboService(ctx, args.OrganizationID.TenantID())
 
-	filter := coredata.NewVendorFilter(nil, nil)
+	pageOrderBy := page.OrderBy[coredata.VendorOrderField]{
-	cursor := page.NewCursor(
+		Field:     coredata.VendorOrderFieldCreatedAt,
-		args.Size,
+		Direction: page.OrderDirectionDesc,
-		args.Cursor,
+	}
-		page.Head,
+	if args.OrderBy != nil {
-		page.OrderBy[coredata.VendorOrderField]{
+		pageOrderBy = page.OrderBy[coredata.VendorOrderField]{
-			Field:     args.OrderField,
+			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{}
+	return nil, types.NewListVendorsOutput(page), nil
-	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
 }

pkg/server/api/mcp/v1/middleware.go

@@ -0,0 +1,114 @@
+// Copyright (c) 2025 Probo Inc <hello@getprobo.com>.
+//
+// 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 <token>" 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))
+	})
+}

pkg/server/api/mcp/v1/types.go

@@ -0,0 +1,66 @@
+// Copyright (c) 2025 Probo Inc <hello@getprobo.com>.
+//
+// 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
+	}
+}

pkg/server/api/mcp/v1/types/cursorkey.go

@@ -0,0 +1,28 @@
+// Copyright (c) 2025 Probo Inc <hello@getprobo.com>.
+//
+// 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)
+}

pkg/server/api/mcp/v1/types/order_by.go

@@ -0,0 +1,24 @@
+// Copyright (c) 2025 Probo Inc <hello@getprobo.com>.
+//
+// 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
+	}
+)

pkg/server/api/mcp/v1/types/organization.go

@@ -0,0 +1,24 @@
+// Copyright (c) 2025 Probo Inc <hello@getprobo.com>.
+//
+// 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"`
+}

pkg/server/api/mcp/v1/types/types.go

@@ -0,0 +1,33 @@
+// Copyright (c) 2025 Probo Inc <hello@getprobo.com>.
+//
+// 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(),
+	}
+}

pkg/server/api/mcp/v1/types/vendor.go

@@ -0,0 +1,298 @@
+// Copyright (c) 2025 Probo Inc <hello@getprobo.com>.
+//
+// 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),
+	}
+}

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
+		proboSvc *probo.Service
-		organizationID gid.GID
+		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")
+	server.AddReceivingMiddleware(mcputils.LoggingMiddleware(logger))
-	if err != nil {
+
-		panic(err)
+	resolver := &resolver{
+		proboSvc: proboSvc,
+		authSvc:  authSvc,
+		authzSvc: authzSvc,
+		logger:   logger,
 	}
 
-	organizationID, err := gid.ParseGID("lXdXZSh-AAEAAAAAAZfLJi38a0AGbu37")
+	mcp.AddTool(server, ListOrganizationsTool, resolver.ListOrganizations)
-	if err != nil {
+	mcp.AddTool(server, ListVendorsTool, resolver.ListVendors)
-		panic(err)
+	mcp.AddTool(server, AddVendorTool, resolver.AddVendor)
-	}
-
-	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 }
+	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
 }