From 944f61da1513d665bf9ad08816da079ec11c29d1 Mon Sep 17 00:00:00 2001
From: Bryan Frimin <bryan@getprobo.com>
Date: Wed, 12 Nov 2025 11:57:42 +0100
Subject: [PATCH] Fix various errors

Signed-off-by: Bryan Frimin <bryan@getprobo.com>
---
 pkg/server/api/mcp/mcputils/mcputils.go     |  7 -------
 pkg/server/api/mcp/v1/add_vendor.go         |  1 +
 pkg/server/api/mcp/v1/list_organizations.go |  6 ++----
 pkg/server/api/mcp/v1/types/vendor.go       | 12 ++++++------
 4 files changed, 9 insertions(+), 17 deletions(-)

diff --git a/pkg/server/api/mcp/mcputils/mcputils.go b/pkg/server/api/mcp/mcputils/mcputils.go
index 5ca77c579..071326774 100644
--- a/pkg/server/api/mcp/mcputils/mcputils.go
+++ b/pkg/server/api/mcp/mcputils/mcputils.go
@@ -26,18 +26,14 @@ import (
 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),
 				)
 			}
 
@@ -48,7 +44,6 @@ func LoggingMiddleware(logger *log.Logger) func(mcp.MethodHandler) mcp.MethodHan
 			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),
 				)
@@ -56,14 +51,12 @@ func LoggingMiddleware(logger *log.Logger) func(mcp.MethodHandler) mcp.MethodHan
 
 				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),
 					)
 				}
diff --git a/pkg/server/api/mcp/v1/add_vendor.go b/pkg/server/api/mcp/v1/add_vendor.go
index dce20fc7c..6b18cdb1e 100644
--- a/pkg/server/api/mcp/v1/add_vendor.go
+++ b/pkg/server/api/mcp/v1/add_vendor.go
@@ -44,6 +44,7 @@ func (r *resolver) AddVendor(
 			BusinessAssociateAgreementURL: args.BusinessAssociateAgreementURL,
 			SubprocessorsListURL:          args.SubprocessorsListURL,
 			Certifications:                args.Certifications,
+			Countries:                     args.Countries,
 			SecurityPageURL:               args.SecurityPageURL,
 			TrustPageURL:                  args.TrustPageURL,
 			TermsOfServiceURL:             args.TermsOfServiceURL,
diff --git a/pkg/server/api/mcp/v1/list_organizations.go b/pkg/server/api/mcp/v1/list_organizations.go
index 6bc2f2e33..babe3d44e 100644
--- a/pkg/server/api/mcp/v1/list_organizations.go
+++ b/pkg/server/api/mcp/v1/list_organizations.go
@@ -30,10 +30,8 @@ var (
 		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"},
-			},
+			Type:       "object",
+			Properties: map[string]*jsonschema.Schema{},
 		},
 	}
 )
diff --git a/pkg/server/api/mcp/v1/types/vendor.go b/pkg/server/api/mcp/v1/types/vendor.go
index de4822a0b..2ea95d5d5 100644
--- a/pkg/server/api/mcp/v1/types/vendor.go
+++ b/pkg/server/api/mcp/v1/types/vendor.go
@@ -103,19 +103,19 @@ type (
 var (
 	ListVendorsInputSchema = &jsonschema.Schema{
 		Type:     "object",
-		Required: []string{"organizationID"},
+		Required: []string{"organization_id"},
 		Properties: map[string]*jsonschema.Schema{
-			"organizationID": {Type: "string"},
+			"organization_id": {Type: "string"},
 			"filter": {
 				Type: "object",
 				Properties: map[string]*jsonschema.Schema{
-					"snapshotID": {Type: "string"},
+					"snapshot_id": {Type: "string"},
 				},
 			},
-			"orderBy": {
+			"order_field": {
 				Types: []string{"object", "null"},
 				Properties: map[string]*jsonschema.Schema{
-					"field":     {Type: "string", Enum: []any{"CREATED_AT"}},
+					"field":     {Type: "string", Enum: []any{"CREATED_AT", "UPDATED_AT", "NAME"}},
 					"direction": OrderByDirectionSchema,
 				},
 			},
@@ -170,7 +170,7 @@ var (
 			"description":                      NullableStringSchema,
 			"category":                         VendorCategorySchema,
 			"headquarter_address":              NullableStringSchema,
-			"legalName":                        NullableStringSchema,
+			"legal_name":                       NullableStringSchema,
 			"website_url":                      NullableStringSchema,
 			"privacy_policy_url":               NullableStringSchema,
 			"service_level_agreement_url":      NullableStringSchema,