Migrate to mcpgen

Signed-off-by: Bryan Frimin <bryan@getprobo.com>
2025-11-20 18:36:59 +01:00
parent d6278da0d6
commit 0b5d4f1cd3
17 changed files with 719 additions and 568 deletions
--- a/pkg/server/api/mcp/v1/add_vendor.go
+++ b/pkg/server/api/mcp/v1/add_vendor.go
@@ -1,61 +0,0 @@
-package v1
-
-import (
-	"context"
-	"fmt"
-
-	"github.com/modelcontextprotocol/go-sdk/mcp"
-	"go.probo.inc/probo/pkg/probo"
-	"go.probo.inc/probo/pkg/server/api/mcp/v1/types"
-)
-
-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 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:                args.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,
-			Countries:                     args.Countries,
-			SecurityPageURL:               args.SecurityPageURL,
-			TrustPageURL:                  args.TrustPageURL,
-			TermsOfServiceURL:             args.TermsOfServiceURL,
-			StatusPageURL:                 args.StatusPageURL,
-			BusinessOwnerID:               args.BusinessOwnerID,
-			SecurityOwnerID:               args.SecurityOwnerID,
-		},
-	)
-	if err != nil {
-		return nil, types.AddVendorOutput{}, fmt.Errorf("failed to create vendor: %w", err)
-	}
-
-	return nil, types.NewAddVendorOutput(vendor), nil
-}
--- a/pkg/server/api/mcp/v1/list_organizations.go
+++ b/pkg/server/api/mcp/v1/list_organizations.go
@@ -1,59 +0,0 @@
-// 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{},
-		},
-	}
-)
-
-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
-}
--- a/pkg/server/api/mcp/v1/list_vendors.go
+++ b/pkg/server/api/mcp/v1/list_vendors.go
@@ -1,55 +0,0 @@
-package v1
-
-import (
-	"context"
-	"fmt"
-
-	"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"
-)
-
-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 types.ListVendorsInput,
-) (*mcp.CallToolResult, types.ListVendorsOutput, error) {
-	prb := r.ProboService(ctx, args.OrganizationID.TenantID())
-
-	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,
-		}
-	}
-
-	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 {
-		panic(fmt.Errorf("cannot list organization vendors: %w", err))
-	}
-
-	return nil, types.NewListVendorsOutput(page), nil
-}
--- a/pkg/server/api/mcp/v1/mcpgen.yaml
+++ b/pkg/server/api/mcp/v1/mcpgen.yaml
@@ -0,0 +1,17 @@
+spec: specification.yaml
+
+output: "."
+
+resolver:
+  package: mcp_v1
+  filename: v1_resolver.go
+  type: Resolver
+  preserve: true
+
+server:
+  package: server
+  filename: server/server.go
+
+model:
+  package: types
+  filename: types/types.go
--- a/pkg/server/api/mcp/v1/middleware.go
+++ b/pkg/server/api/mcp/v1/middleware.go
@@ -12,7 +12,7 @@
 // OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
 // PERFORMANCE OF THIS SOFTWARE.

-package v1
+package mcp_v1

 import (
 	"net/http"
--- a/pkg/server/api/mcp/v1/resolver.go
+++ b/pkg/server/api/mcp/v1/resolver.go
@@ -0,0 +1,17 @@
+//go:generate go run go.probo.inc/mcpgen generate
+
+package mcp_v1
+
+import (
+	"go.gearno.de/kit/log"
+	"go.probo.inc/probo/pkg/auth"
+	"go.probo.inc/probo/pkg/authz"
+	"go.probo.inc/probo/pkg/probo"
+)
+
+type Resolver struct {
+	proboSvc *probo.Service
+	authSvc  *auth.Service
+	authzSvc *authz.Service
+	logger   *log.Logger
+}
--- a/pkg/server/api/mcp/v1/schema.resolvers.go
+++ b/pkg/server/api/mcp/v1/schema.resolvers.go
@@ -0,0 +1,110 @@
+package mcp_v1
+
+// This file will be automatically regenerated based on the schema, any resolver implementations
+// will be copied through when generating and any unknown code will be moved to the end.
+// Code generated by mcpgen. DO NOT EDIT.
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/modelcontextprotocol/go-sdk/mcp"
+	"go.probo.inc/probo/pkg/coredata"
+	"go.probo.inc/probo/pkg/page"
+	"go.probo.inc/probo/pkg/probo"
+	"go.probo.inc/probo/pkg/server/api/mcp/v1/types"
+)
+
+// ListOrganizationsTool handles the listOrganizations tool
+// List all organizations the user has access to
+func (r *Resolver) ListOrganizationsTool(ctx context.Context, req *mcp.CallToolRequest, input *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
+}
+
+// ListVendorsTool handles the listVendors tool
+// List all vendors for the organization
+func (r *Resolver) ListVendorsTool(ctx context.Context, req *mcp.CallToolRequest, input *types.ListVendorsInput) (*mcp.CallToolResult, types.ListVendorsOutput, error) {
+	prb := r.ProboService(ctx, input.OrganizationID.TenantID())
+
+	pageOrderBy := page.OrderBy[coredata.VendorOrderField]{
+		Field:     coredata.VendorOrderFieldCreatedAt,
+		Direction: page.OrderDirectionDesc,
+	}
+	if input.OrderBy != nil {
+		pageOrderBy = page.OrderBy[coredata.VendorOrderField]{
+			Field:     input.OrderBy.Field,
+			Direction: input.OrderBy.Direction,
+		}
+	}
+
+	cursor := types.NewCursor(input.Size, input.Cursor, pageOrderBy)
+
+	var vendorFilter = coredata.NewVendorFilter(nil, nil)
+	if input.Filter != nil {
+		vendorFilter = coredata.NewVendorFilter(&input.Filter.SnapshotID, nil)
+	}
+
+	page, err := prb.Vendors.ListForOrganizationID(ctx, input.OrganizationID, cursor, vendorFilter)
+	if err != nil {
+		panic(fmt.Errorf("cannot list organization vendors: %w", err))
+	}
+
+	return nil, types.NewListVendorsOutput(page), nil
+}
+
+// AddVendorTool handles the addVendor tool
+// Add a new vendor to the organization
+func (r *Resolver) AddVendorTool(ctx context.Context, req *mcp.CallToolRequest, input *types.AddVendorInput) (*mcp.CallToolResult, types.AddVendorOutput, error) {
+	svc := r.ProboService(ctx, input.OrganizationID.TenantID())
+
+	vendor, err := svc.Vendors.Create(
+		ctx,
+		probo.CreateVendorRequest{
+			OrganizationID: input.OrganizationID,
+			Name:           input.Name,
+			Description:    input.Description,
+			// HeadquarterAddress:            input.HeadquarterAddress,
+			// LegalName:                     input.LegalName,
+			// WebsiteURL:                    input.WebsiteURL,
+			// Category:                      input.Category,
+			// PrivacyPolicyURL:              input.PrivacyPolicyURL,
+			// ServiceLevelAgreementURL:      input.ServiceLevelAgreementURL,
+			// DataProcessingAgreementURL:    input.DataProcessingAgreementURL,
+			// BusinessAssociateAgreementURL: input.BusinessAssociateAgreementURL,
+			// SubprocessorsListURL:          input.SubprocessorsListURL,
+			// Certifications:                input.Certifications,
+			// Countries:                     input.Countries,
+			// SecurityPageURL:               input.SecurityPageURL,
+			// TrustPageURL:                  input.TrustPageURL,
+			// TermsOfServiceURL:             input.TermsOfServiceURL,
+			// StatusPageURL:                 input.StatusPageURL,
+			// BusinessOwnerID:               input.BusinessOwnerID,
+			// SecurityOwnerID:               input.SecurityOwnerID,
+		},
+	)
+	if err != nil {
+		return nil, types.AddVendorOutput{}, fmt.Errorf("failed to create vendor: %w", err)
+	}
+
+	return nil, types.NewAddVendorOutput(vendor), nil
+}
+
+// UpdateVendorTool handles the updateVendor tool
+// Update an existing vendor
+func (r *Resolver) UpdateVendorTool(ctx context.Context, req *mcp.CallToolRequest, input *types.UpdateVendorInput) (*mcp.CallToolResult, types.UpdateVendorOutput, error) {
+	return nil, types.UpdateVendorOutput{}, fmt.Errorf("updateVendor not implemented")
+}
--- a/pkg/server/api/mcp/v1/server/server.go
+++ b/pkg/server/api/mcp/v1/server/server.go
@@ -0,0 +1,76 @@
+// Code generated by mcpgen. DO NOT EDIT.
+
+package server
+
+import (
+	"context"
+	"github.com/modelcontextprotocol/go-sdk/mcp"
+	"go.probo.inc/probo/pkg/server/api/mcp/v1/types"
+)
+
+// ResolverInterface defines the interface that must be implemented by the parent resolver
+type ResolverInterface interface {
+	ListOrganizationsTool(ctx context.Context, req *mcp.CallToolRequest, input *types.ListOrganizationsInput) (*mcp.CallToolResult, types.ListOrganizationsOutput, error)
+	ListVendorsTool(ctx context.Context, req *mcp.CallToolRequest, input *types.ListVendorsInput) (*mcp.CallToolResult, types.ListVendorsOutput, error)
+	AddVendorTool(ctx context.Context, req *mcp.CallToolRequest, input *types.AddVendorInput) (*mcp.CallToolResult, types.AddVendorOutput, error)
+	UpdateVendorTool(ctx context.Context, req *mcp.CallToolRequest, input *types.UpdateVendorInput) (*mcp.CallToolResult, types.UpdateVendorOutput, error)
+}
+
+// New creates a new MCP server instance with all handlers registered.
+// Returns a fully configured *mcp.Server ready to be used with any transport.
+func New(resolver ResolverInterface) *mcp.Server {
+	server := mcp.NewServer(
+		&mcp.Implementation{
+			Name:    "Probo MCP Server",
+			Version: "1.0.0",
+		},
+		nil,
+	)
+
+	registerToolHandlers(server, resolver)
+
+	return server
+}
+
+func registerToolHandlers(server *mcp.Server, resolver ResolverInterface) {
+	mcp.AddTool(
+		server,
+		&mcp.Tool{
+			Name:         "listOrganizations",
+			Description:  "List all organizations the user has access to",
+			InputSchema:  types.ListOrganizationsToolInputSchema,
+			OutputSchema: types.ListOrganizationsToolOutputSchema,
+		},
+		resolver.ListOrganizationsTool,
+	)
+	mcp.AddTool(
+		server,
+		&mcp.Tool{
+			Name:         "listVendors",
+			Description:  "List all vendors for the organization",
+			InputSchema:  types.ListVendorsToolInputSchema,
+			OutputSchema: types.ListVendorsToolOutputSchema,
+		},
+		resolver.ListVendorsTool,
+	)
+	mcp.AddTool(
+		server,
+		&mcp.Tool{
+			Name:         "addVendor",
+			Description:  "Add a new vendor to the organization",
+			InputSchema:  types.AddVendorToolInputSchema,
+			OutputSchema: types.AddVendorToolOutputSchema,
+		},
+		resolver.AddVendorTool,
+	)
+	mcp.AddTool(
+		server,
+		&mcp.Tool{
+			Name:         "updateVendor",
+			Description:  "Update an existing vendor",
+			InputSchema:  types.UpdateVendorToolInputSchema,
+			OutputSchema: types.UpdateVendorToolOutputSchema,
+		},
+		resolver.UpdateVendorTool,
+	)
+}
--- a/pkg/server/api/mcp/v1/specification.yaml
+++ b/pkg/server/api/mcp/v1/specification.yaml
@@ -0,0 +1,241 @@
+info:
+  title: Probo MCP Server
+  version: 1.0.0
+  description: Probo MCP Server API
+
+components:
+  schemas:
+    OrderDirection:
+      type: string
+      enum:
+        - asc
+        - desc
+      go.probo.inc/mcpgen/type: go.probo.inc/probo/pkg/page.OrderDirection
+
+    VendorOrderField:
+      type: string
+      enum:
+        - created_at
+        - updated_at
+        - name
+      go.probo.inc/mcpgen/type: go.probo.inc/probo/pkg/coredata.VendorOrderField
+
+    VendorOrderBy:
+      type: object
+      required:
+        - field
+        - direction
+      properties:
+        field:
+          $ref: "#/components/schemas/VendorOrderField"
+          description: Vendor order field
+        direction:
+          $ref: "#/components/schemas/OrderDirection"
+          description: Vendor order direction
+    GID:
+      type: string
+      format: string
+      go.probo.inc/mcpgen/type: go.probo.inc/probo/pkg/gid.GID
+
+    CursorKey:
+      type: string
+      format: string
+      go.probo.inc/mcpgen/type: go.probo.inc/probo/pkg/page.CursorKey
+
+    ListOrganizationsInput:
+      type: object
+      properties:
+        organization_id:
+          $ref: "#/components/schemas/GID"
+          description: Organization ID
+
+    ListOrganizationsOutput:
+      type: object
+      properties:
+        organizations:
+          type: array
+          items:
+            $ref: "#/components/schemas/Organization"
+
+    Organization:
+      type: object
+      required:
+        - id
+        - name
+        - description
+        - created_at
+        - updated_at
+      properties:
+        id:
+          $ref: "#/components/schemas/GID"
+          description: Organization ID
+        name:
+          type: string
+          description: Organization name
+        description:
+          type:
+            - string
+            - "null"
+          description: Organization description
+        created_at:
+          type: string
+          format: date-time
+          description: Creation timestamp
+        updated_at:
+          type: string
+          format: date-time
+          description: Update timestamp
+
+    ListVendorsInput:
+      type: object
+      required:
+        - organization_id
+      properties:
+        organization_id:
+          $ref: "#/components/schemas/GID"
+          description: Organization ID
+        order_by:
+          $ref: "#/components/schemas/VendorOrderBy"
+          description: Vendor order by
+        size:
+          type: integer
+          description: Page size
+        cursor:
+          $ref: "#/components/schemas/CursorKey"
+          description: Page cursor
+        filter:
+          type: object
+          properties:
+            snapshot_id:
+              $ref: "#/components/schemas/GID"
+              description: Snapshot ID
+
+    ListVendorsOutput:
+      type: object
+      required:
+        - vendors
+      properties:
+        next_cursor:
+          $ref: "#/components/schemas/CursorKey"
+          description: Next cursor
+        vendors:
+          type: array
+          items:
+            $ref: "#/components/schemas/Vendor"
+
+    Vendor:
+      type: object
+      required:
+        - id
+        - name
+        - organization_id
+        - created_at
+        - updated_at
+      properties:
+        id:
+          $ref: "#/components/schemas/GID"
+          description: Vendor ID
+        name:
+          type: string
+          description: Vendor name
+        organization_id:
+          $ref: "#/components/schemas/GID"
+          description: Organization ID
+        description:
+          type:
+            - string
+            - "null"
+          description: Vendor description
+        created_at:
+          type: string
+          format: date-time
+          description: Creation timestamp
+        updated_at:
+          type: string
+          format: date-time
+          description: Update timestamp
+
+    AddVendorInput:
+      type: object
+      required:
+        - organization_id
+        - name
+      properties:
+        organization_id:
+          $ref: "#/components/schemas/GID"
+          description: Organization ID
+        name:
+          type: string
+          description: Vendor name
+        description:
+          type: string
+          description: Vendor description
+        created_at:
+          type: string
+          format: date-time
+          description: Creation timestamp
+        updated_at:
+          type: string
+          format: date-time
+          description: Update timestamp
+
+    AddVendorOutput:
+      type: object
+      required:
+        - vendor
+      properties:
+        vendor:
+          $ref: "#/components/schemas/Vendor"
+
+    UpdateVendorInput:
+      type: object
+      required:
+        - id
+      properties:
+        id:
+          $ref: "#/components/schemas/GID"
+          description: Vendor ID
+        name:
+          type: string
+          description: Vendor name
+        description:
+          type: string
+          description: Vendor description
+
+    UpdateVendorOutput:
+      type: object
+      required:
+        - vendor
+      properties:
+        vendor:
+          $ref: "#/components/schemas/Vendor"
+
+tools:
+  - name: listOrganizations
+    description: List all organizations the user has access to
+    readonly: true
+    inputSchema:
+      $ref: "#/components/schemas/ListOrganizationsInput"
+    outputSchema:
+      $ref: "#/components/schemas/ListOrganizationsOutput"
+  - name: listVendors
+    description: List all vendors for the organization
+    readonly: true
+    inputSchema:
+      $ref: "#/components/schemas/ListVendorsInput"
+    outputSchema:
+      $ref: "#/components/schemas/ListVendorsOutput"
+  - name: addVendor
+    description: Add a new vendor to the organization
+    readonly: false
+    inputSchema:
+      $ref: "#/components/schemas/AddVendorInput"
+    outputSchema:
+      $ref: "#/components/schemas/AddVendorOutput"
+  - name: updateVendor
+    description: Update an existing vendor
+    readonly: false
+    inputSchema:
+      $ref: "#/components/schemas/UpdateVendorInput"
+    outputSchema:
+      $ref: "#/components/schemas/UpdateVendorOutput"
--- a/pkg/server/api/mcp/v1/types.go
+++ b/pkg/server/api/mcp/v1/types.go
@@ -12,7 +12,7 @@
 // OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
 // PERFORMANCE OF THIS SOFTWARE.

-package v1
+package mcp_v1

 import (
 	"context"
--- a/pkg/server/api/mcp/v1/types/cursorkey.go
+++ b/pkg/server/api/mcp/v1/types/cursorkey.go
@@ -15,14 +15,20 @@
 package types

 import (
+	"go.gearno.de/x/ref"
 	"go.probo.inc/probo/pkg/page"
 	"go.probo.inc/probo/pkg/server/gqlutils/types/cursor"
 )

 func NewCursor[O page.OrderField](
-	first *int,
+	first *int64,
 	after *page.CursorKey,
 	orderBy page.OrderBy[O],
 ) *page.Cursor[O] {
-	return cursor.NewCursor(first, after, nil, nil, orderBy)
+	firstValue := ref.Ref(100)
+	if first != nil {
+		firstValue = ref.Ref(int(*first))
+	}
+
+	return cursor.NewCursor(firstValue, after, nil, nil, orderBy)
 }
--- a/pkg/server/api/mcp/v1/types/organization.go
+++ b/pkg/server/api/mcp/v1/types/organization.go
@@ -14,8 +14,14 @@

 package types

-type ListOrganizationsInput struct{}
+import "go.probo.inc/probo/pkg/coredata"

-type ListOrganizationsOutput struct {
-	Organizations []Organization `json:"organizations" jsonschema:"list of organizations the user has access to"`
+func NewOrganization(o *coredata.Organization) *Organization {
+	return &Organization{
+		ID:          o.ID,
+		Name:        o.Name,
+		Description: o.Description,
+		CreatedAt:   o.CreatedAt,
+		UpdatedAt:   o.UpdatedAt,
+	}
 }
--- a/pkg/server/api/mcp/v1/types/types.go
+++ b/pkg/server/api/mcp/v1/types/types.go
@@ -1,31 +1,132 @@
-// 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.
+// Code generated by mcpgen. DO NOT EDIT.

 package types

-import "go.probo.inc/probo/pkg/coredata"
+import (
+	"go.probo.inc/mcpgen/mcp"
+	"go.probo.inc/probo/pkg/coredata"
+	"go.probo.inc/probo/pkg/gid"
+	"go.probo.inc/probo/pkg/page"
+	"time"
+)

+// Tool input schemas
+var (
+	AddVendorToolInputSchema          = mcp.MustUnmarshalSchema(`{"type":"object","required":["organization_id","name"],"properties":{"created_at":{"type":"string","description":"Creation timestamp","format":"date-time"},"description":{"type":"string","description":"Vendor description"},"name":{"type":"string","description":"Vendor name"},"organization_id":{"type":"string","format":"string"},"updated_at":{"type":"string","description":"Update timestamp","format":"date-time"}}}`)
+	AddVendorToolOutputSchema         = mcp.MustUnmarshalSchema(`{"type":"object","required":["vendor"],"properties":{"vendor":{"type":"object","required":["id","name","organization_id","created_at","updated_at"],"properties":{"created_at":{"type":"string","description":"Creation timestamp","format":"date-time"},"description":{"description":"Vendor description"},"id":{"type":"string","format":"string"},"name":{"type":"string","description":"Vendor name"},"organization_id":{"type":"string","format":"string"},"updated_at":{"type":"string","description":"Update timestamp","format":"date-time"}}}}}`)
+	ListOrganizationsToolInputSchema  = mcp.MustUnmarshalSchema(`{"type":"object","properties":{"organization_id":{"type":"string","format":"string"}}}`)
+	ListOrganizationsToolOutputSchema = mcp.MustUnmarshalSchema(`{"type":"object","properties":{"organizations":{"type":"array","items":{"type":"object","required":["id","name","description","created_at","updated_at"],"properties":{"created_at":{"type":"string","description":"Creation timestamp","format":"date-time"},"description":{"description":"Organization description"},"id":{"type":"string","format":"string"},"name":{"type":"string","description":"Organization name"},"updated_at":{"type":"string","description":"Update timestamp","format":"date-time"}}}}}}`)
+	ListVendorsToolInputSchema        = mcp.MustUnmarshalSchema(`{"type":"object","required":["organization_id"],"properties":{"cursor":{"type":"string","format":"string"},"filter":{"type":"object","properties":{"snapshot_id":{"type":"string","format":"string"}}},"order_by":{"type":"object","required":["field","direction"],"properties":{"direction":{"type":"string","enum":["asc","desc"]},"field":{"type":"string","enum":["created_at","updated_at","name"]}}},"organization_id":{"type":"string","format":"string"},"size":{"type":"integer","description":"Page size"}}}`)
+	ListVendorsToolOutputSchema       = mcp.MustUnmarshalSchema(`{"type":"object","required":["vendors"],"properties":{"next_cursor":{"type":"string","format":"string"},"vendors":{"type":"array","items":{"type":"object","required":["id","name","organization_id","created_at","updated_at"],"properties":{"created_at":{"type":"string","description":"Creation timestamp","format":"date-time"},"description":{"description":"Vendor description"},"id":{"type":"string","format":"string"},"name":{"type":"string","description":"Vendor name"},"organization_id":{"type":"string","format":"string"},"updated_at":{"type":"string","description":"Update timestamp","format":"date-time"}}}}}}`)
+	UpdateVendorToolInputSchema       = mcp.MustUnmarshalSchema(`{"type":"object","required":["id"],"properties":{"description":{"type":"string","description":"Vendor description"},"id":{"type":"string","format":"string"},"name":{"type":"string","description":"Vendor name"}}}`)
+	UpdateVendorToolOutputSchema      = mcp.MustUnmarshalSchema(`{"type":"object","required":["vendor"],"properties":{"vendor":{"type":"object","required":["id","name","organization_id","created_at","updated_at"],"properties":{"created_at":{"type":"string","description":"Creation timestamp","format":"date-time"},"description":{"description":"Vendor description"},"id":{"type":"string","format":"string"},"name":{"type":"string","description":"Vendor name"},"organization_id":{"type":"string","format":"string"},"updated_at":{"type":"string","description":"Update timestamp","format":"date-time"}}}}}`)
+)
+
+// AddVendorInput represents the schema
+type AddVendorInput struct {
+	// Creation timestamp
+	CreatedAt *time.Time `json:"created_at,omitempty"`
+	// Vendor description
+	Description *string `json:"description,omitempty"`
+	// Vendor name
+	Name string `json:"name"`
+	// Organization ID
+	OrganizationID gid.GID `json:"organization_id"`
+	// Update timestamp
+	UpdatedAt *time.Time `json:"updated_at,omitempty"`
+}
+
+// AddVendorOutput represents the schema
+type AddVendorOutput struct {
+	Vendor *Vendor `json:"vendor"`
+}
+
+// ListOrganizationsInput represents the schema
+type ListOrganizationsInput struct {
+	// Organization ID
+	OrganizationID *gid.GID `json:"organization_id,omitempty"`
+}
+
+// ListOrganizationsOutput represents the schema
+type ListOrganizationsOutput struct {
+	Organizations []*Organization `json:"organizations,omitempty"`
+}
+
+// ListVendorsInput represents the schema
+type ListVendorsInput struct {
+	// Page cursor
+	Cursor *page.CursorKey         `json:"cursor,omitempty"`
+	Filter *ListVendorsInputFilter `json:"filter,omitempty"`
+	// Vendor order by
+	OrderBy *VendorOrderBy `json:"order_by,omitempty"`
+	// Organization ID
+	OrganizationID gid.GID `json:"organization_id"`
+	// Page size
+	Size *int64 `json:"size,omitempty"`
+}
+
+// ListVendorsOutput represents the schema
+type ListVendorsOutput struct {
+	// Next cursor
+	NextCursor *page.CursorKey `json:"next_cursor,omitempty"`
+	Vendors    []*Vendor       `json:"vendors"`
+}
+
+// Organization represents the schema
 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"`
+	// Creation timestamp
+	CreatedAt time.Time `json:"created_at"`
+	// Organization description
+	Description *string `json:"description"`
+	// Organization ID
+	ID gid.GID `json:"id"`
+	// Organization name
+	Name string `json:"name"`
+	// Update timestamp
+	UpdatedAt time.Time `json:"updated_at"`
 }

-func NewOrganization(o *coredata.Organization) Organization {
-	return Organization{
-		Name:     o.Name,
-		ID:       o.ID.String(),
-		TenantID: o.ID.TenantID().String(),
-	}
+// UpdateVendorInput represents the schema
+type UpdateVendorInput struct {
+	// Vendor description
+	Description *string `json:"description,omitempty"`
+	// Vendor ID
+	ID gid.GID `json:"id"`
+	// Vendor name
+	Name *string `json:"name,omitempty"`
+}
+
+// UpdateVendorOutput represents the schema
+type UpdateVendorOutput struct {
+	Vendor *Vendor `json:"vendor"`
+}
+
+// Vendor represents the schema
+type Vendor struct {
+	// Creation timestamp
+	CreatedAt time.Time `json:"created_at"`
+	// Vendor description
+	Description *string `json:"description,omitempty"`
+	// Vendor ID
+	ID gid.GID `json:"id"`
+	// Vendor name
+	Name string `json:"name"`
+	// Organization ID
+	OrganizationID gid.GID `json:"organization_id"`
+	// Update timestamp
+	UpdatedAt time.Time `json:"updated_at"`
+}
+
+// VendorOrderBy represents the schema
+type VendorOrderBy struct {
+	// Vendor order direction
+	Direction page.OrderDirection `json:"direction"`
+	// Vendor order field
+	Field coredata.VendorOrderField `json:"field"`
+}
+
+// ListVendorsInputFilter represents the schema
+type ListVendorsInputFilter struct {
+	// Snapshot ID
+	SnapshotID *gid.GID `json:"snapshot_id,omitempty"`
 }
--- a/pkg/server/api/mcp/v1/types/vendor.go
+++ b/pkg/server/api/mcp/v1/types/vendor.go
@@ -15,273 +15,30 @@
 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{"organization_id"},
-		Properties: map[string]*jsonschema.Schema{
-			"organization_id": {Type: "string"},
-			"filter": {
-				Type: "object",
-				Properties: map[string]*jsonschema.Schema{
-					"snapshot_id": {Type: "string"},
-				},
-			},
-			"order_field": {
-				Types: []string{"object", "null"},
-				Properties: map[string]*jsonschema.Schema{
-					"field":     {Type: "string", Enum: []any{"CREATED_AT", "UPDATED_AT", "NAME"}},
-					"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{
-			coredata.VendorCategoryAnalytics.String(),
-			coredata.VendorCategoryCloudMonitoring.String(),
-			coredata.VendorCategoryCloudProvider.String(),
-			coredata.VendorCategoryCollaboration.String(),
-			coredata.VendorCategoryCustomerSupport.String(),
-			coredata.VendorCategoryDataStorageAndProcessing.String(),
-			coredata.VendorCategoryDocumentManagement.String(),
-			coredata.VendorCategoryEmployeeManagement.String(),
-			coredata.VendorCategoryEngineering.String(),
-			coredata.VendorCategoryFinance.String(),
-			coredata.VendorCategoryIdentityProvider.String(),
-			coredata.VendorCategoryIT.String(),
-			coredata.VendorCategoryMarketing.String(),
-			coredata.VendorCategoryOfficeOperations.String(),
-			coredata.VendorCategoryOther.String(),
-			coredata.VendorCategoryPasswordManagement.String(),
-			coredata.VendorCategoryProductAndDesign.String(),
-			coredata.VendorCategoryProfessionalServices.String(),
-			coredata.VendorCategoryRecruiting.String(),
-			coredata.VendorCategorySales.String(),
-			coredata.VendorCategorySecurity.String(),
-			coredata.VendorCategoryVersionControl.String(),
-		},
-	}
-
-	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,
-			"legal_name":                       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 NewVendor(v *coredata.Vendor) *Vendor {
+	return &Vendor{
+		Name:           v.Name,
+		ID:             v.ID,
+		OrganizationID: v.OrganizationID,
+		Description:    v.Description,
+		CreatedAt:      v.CreatedAt,
+		UpdatedAt:      v.UpdatedAt,
 	}
 }

 func NewListVendorsOutput(vendorPage *page.Page[*coredata.Vendor, coredata.VendorOrderField]) ListVendorsOutput {
-	vendors := make([]Vendor, 0, len(vendorPage.Data))
+	vendors := make([]*Vendor, 0, len(vendorPage.Data))
 	for _, v := range vendorPage.Data {
 		vendors = append(vendors, NewVendor(v))
 	}

-	var nextCursor *string
+	var nextCursor *page.CursorKey
 	if len(vendorPage.Data) > 0 {
-		cursorKey := vendorPage.Data[len(vendorPage.Data)-1].CursorKey(vendorPage.Cursor.OrderBy.Field).String()
+		cursorKey := vendorPage.Data[len(vendorPage.Data)-1].CursorKey(vendorPage.Cursor.OrderBy.Field)
 		nextCursor = &cursorKey
 	}

--- a/pkg/server/api/mcp/v1/v1_handler.go
+++ b/pkg/server/api/mcp/v1/v1_handler.go
@@ -1,4 +1,4 @@
-package v1
+package mcp_v1

 import (
 	"context"
@@ -13,19 +13,10 @@ import (
 	"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"
+	"go.probo.inc/probo/pkg/server/api/mcp/v1/server"
 )

-type (
-	resolver struct {
-		proboSvc *probo.Service
-		authSvc  *auth.Service
-		authzSvc *authz.Service
-		logger   *log.Logger
-	}
-)
-
-func (r *resolver) ProboService(ctx context.Context, tenantID gid.TenantID) *probo.TenantService {
+func (r *Resolver) ProboService(ctx context.Context, tenantID gid.TenantID) *probo.TenantService {
 	validateTenantAccess(ctx, tenantID)
 	return r.proboSvc.WithTenant(tenantID)
 }
@@ -52,30 +43,18 @@ func NewMux(logger *log.Logger, proboSvc *probo.Service, authSvc *auth.Service,
 		log.String("version", cfg.Version),
 		log.String("request_timeout", cfg.RequestTimeout.String()),
 	)
+	// server.AddReceivingMiddleware(mcputils.LoggingMiddleware(logger))

-	server := mcp.NewServer(
-		&mcp.Implementation{
-			Name:    "probo",
-			Title:   "Probo",
-			Version: cfg.Version,
-		},
-		&mcp.ServerOptions{},
-	)
-
-	server.AddReceivingMiddleware(mcputils.LoggingMiddleware(logger))
-
-	resolver := &resolver{
+	resolver := &Resolver{
 		proboSvc: proboSvc,
 		authSvc:  authSvc,
 		authzSvc: authzSvc,
 		logger:   logger,
 	}

-	mcp.AddTool(server, ListOrganizationsTool, resolver.ListOrganizations)
-	mcp.AddTool(server, ListVendorsTool, resolver.ListVendors)
-	mcp.AddTool(server, AddVendorTool, resolver.AddVendor)
+	mcpServer := server.New(resolver)

-	getServer := func(r *http.Request) *mcp.Server { return server }
+	getServer := func(r *http.Request) *mcp.Server { return mcpServer }
 	eventStore := mcp.NewMemoryEventStore(nil)

 	handler := mcp.NewStreamableHTTPHandler(