Fix atomic latest version update to prevent missing isLatest flags

Previously, updating the old latest version and creating the new version
were separate operations that could fail independently, potentially leaving
no version marked as isLatest. This change makes the operation atomic:

- CreateServer now accepts oldLatestVersionID parameter
- PostgreSQL implementation uses a transaction for UPDATE + INSERT atomicity
- Service layer maintains WithPublishLock for concurrent operation serialization
- Both mechanisms work together: lock prevents races, transaction ensures atomicity

🤖 Generated with [Claude Code](https://claude.ai/code)

Co-Authored-By: Claude <[email protected]>

Author: domdomegg

internal/database/database.go

@@ -40,8 +40,10 @@ type Database interface {
 	GetByServerIDAndVersion(ctx context.Context, serverID string, version string) (*apiv0.ServerJSON, error)
 	// Retrieve all versions of a server by server ID
 	GetAllVersionsByServerID(ctx context.Context, serverID string) ([]*apiv0.ServerJSON, error)
-	// CreateServer adds a new server to the database
-	CreateServer(ctx context.Context, server *apiv0.ServerJSON) (*apiv0.ServerJSON, error)
+	// CreateServer atomically publishes a new server version, optionally unmarking a previous latest version
+	// If oldLatestVersionID is provided, it will be updated to set IsLatest=false before creating the new version
+	// This operation happens within a transaction with an advisory lock to prevent race conditions
+	CreateServer(ctx context.Context, newServer *apiv0.ServerJSON, oldLatestVersionID *string) (*apiv0.ServerJSON, error)
 	// UpdateServer updates an existing server record
 	UpdateServer(ctx context.Context, id string, server *apiv0.ServerJSON) (*apiv0.ServerJSON, error)
 	// WithPublishLock executes a function with an exclusive lock for publishing a server

internal/database/postgres.go

@@ -328,8 +328,9 @@ func (db *PostgreSQL) GetAllVersionsByServerID(ctx context.Context, serverID str
 	return results, nil
 }
 
-// CreateServer adds a new server to the database
-func (db *PostgreSQL) CreateServer(ctx context.Context, server *apiv0.ServerJSON) (*apiv0.ServerJSON, error) {
+// CreateServer atomically publishes a new server version, optionally unmarking a previous latest version
+// Must be called within WithPublishLock to ensure proper serialization
+func (db *PostgreSQL) CreateServer(ctx context.Context, server *apiv0.ServerJSON, oldLatestVersionID *string) (*apiv0.ServerJSON, error) {
 	if ctx.Err() != nil {
 		return nil, ctx.Err()
 	}
@@ -346,23 +347,53 @@ func (db *PostgreSQL) CreateServer(ctx context.Context, server *apiv0.ServerJSON
 		return nil, fmt.Errorf("server must have both ServerID and VersionID in registry metadata")
 	}
 
+	// Begin a transaction for atomicity of UPDATE + INSERT
+	tx, err := db.pool.Begin(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("failed to begin transaction: %w", err)
+	}
+	defer func() {
+		_ = tx.Rollback(ctx)
+	}()
+
+	// If there's a previous latest version, unmark it
+	if oldLatestVersionID != nil && *oldLatestVersionID != "" {
+		updateQuery := `
+			UPDATE servers
+			SET value = jsonb_set(
+				value,
+				'{_meta,io.modelcontextprotocol.registry/official,isLatest}',
+				'false'::jsonb
+			)
+			WHERE version_id = $1
+		`
+		_, err := tx.Exec(ctx, updateQuery, *oldLatestVersionID)
+		if err != nil {
+			return nil, fmt.Errorf("failed to unmark previous latest version: %w", err)
+		}
+	}
+
 	// Marshal the complete server to JSONB
 	valueJSON, err := json.Marshal(server)
 	if err != nil {
 		return nil, fmt.Errorf("failed to marshal server JSON: %w", err)
 	}
 
-	// Insert into servers table with new schema (only version_id column, serverId is in JSON)
-	query := `
+	// Insert the new version
+	insertQuery := `
 		INSERT INTO servers (version_id, value)
 		VALUES ($1, $2)
 	`
-
-	_, err = db.pool.Exec(ctx, query, versionID, valueJSON)
+	_, err = tx.Exec(ctx, insertQuery, versionID, valueJSON)
 	if err != nil {
 		return nil, fmt.Errorf("failed to insert server: %w", err)
 	}
 
+	// Commit the transaction
+	if err := tx.Commit(ctx); err != nil {
+		return nil, fmt.Errorf("failed to commit transaction: %w", err)
+	}
+
 	return server, nil
 }
 

internal/importer/importer.go

@@ -37,7 +37,7 @@ func (s *Service) ImportFromPath(ctx context.Context, path string) error {
 
 	// Import each server using CreateServer
 	for _, server := range servers {
-		_, err := s.db.CreateServer(ctx, server)
+		_, err := s.db.CreateServer(ctx, server, nil)
 		if err != nil {
 			return fmt.Errorf("failed to import server %s: %w", server.Name, err)
 		}
@@ -181,4 +181,4 @@ func convertServerResponseToRecord(response apiv0.ServerJSON) *apiv0.ServerJSON
 	// The response is already in the correct flattened format
 	// Just return a pointer to it
 	return &response
-}
+}

internal/service/registry_service.go

@@ -176,33 +176,22 @@ func (s *registryServiceImpl) Publish(req apiv0.ServerJSON) (*apiv0.ServerJSON,
 			) > 0
 		}
 
-		// Mark previous latest as no longer latest BEFORE creating new version
-		// This prevents violating the unique constraint on isLatest
+		// Prepare the old latest version ID if we need to unmark it
+		var oldLatestVersionID *string
 		if isNewLatest && existingLatest != nil {
-			var existingLatestVersionID string
 			if existingLatest.Meta != nil && existingLatest.Meta.Official != nil {
-				existingLatestVersionID = existingLatest.Meta.Official.VersionID
-			}
-			if existingLatestVersionID != "" {
-				// Update the existing server to set isLatest = false
-				existingLatest.Meta.Official.IsLatest = false
-				existingLatest.Meta.Official.UpdatedAt = time.Now()
-				if _, err := s.db.UpdateServer(lockCtx, existingLatestVersionID, existingLatest); err != nil {
-					return nil, err
+				versionID := existingLatest.Meta.Official.VersionID
+				if versionID != "" {
+					oldLatestVersionID = &versionID
 				}
 			}
 		}
 
 		// Create complete server with metadata
 		server := s.createServerWithMetadata(serverJSON, existingServerVersions, publishTime, isNewLatest)
 
-		// Create server in database
-		serverRecord, err := s.db.CreateServer(lockCtx, &server)
-		if err != nil {
-			return nil, err
-		}
-
-		return serverRecord, nil
+		// Publish server version atomically (unmarks old latest and creates new version in one transaction)
+		return s.db.CreateServer(lockCtx, &server, oldLatestVersionID)
 	})
 
 	if err != nil {