From 3aaf4bc4e7c857e9c72a134a53b023ab6664e7b6 Mon Sep 17 00:00:00 2001
From: Bolek Kulbabinski <1416262+bolekk@users.noreply.github.com>
Date: Tue, 24 Feb 2026 11:05:55 -0800
Subject: [PATCH] [CRE][Deployments] Fix bug in OCR config count extraction

---
 deployment/cre/ocr3/registry_config.go      | 53 +++++----------
 deployment/cre/ocr3/registry_config_test.go | 74 +++++++++++++++++++++
 2 files changed, 90 insertions(+), 37 deletions(-)

diff --git a/deployment/cre/ocr3/registry_config.go b/deployment/cre/ocr3/registry_config.go
index 3daf4f4859a..2c16beba610 100644
--- a/deployment/cre/ocr3/registry_config.go
+++ b/deployment/cre/ocr3/registry_config.go
@@ -2,13 +2,11 @@ package ocr3
 
 import (
 	"encoding/base64"
-	"encoding/json"
 	"errors"
 	"fmt"
 
 	"github.com/ethereum/go-ethereum/accounts/abi/bind"
 	"github.com/ethereum/go-ethereum/common"
-	"google.golang.org/protobuf/encoding/protojson"
 	"google.golang.org/protobuf/proto"
 
 	libocrtypes "github.com/smartcontractkit/libocr/ragep2p/types"
@@ -113,46 +111,27 @@ func GetCurrentOCR3ConfigCount(
 	}
 
 	for _, capCfg := range don.CapabilityConfigurations {
-		if capCfg.CapabilityId != capabilityID {
-			continue
+		if capCfg.CapabilityId == capabilityID {
+			return extractOCR3ConfigCount(capCfg.Config, ocrConfigKey)
 		}
-
-		if len(capCfg.Config) == 0 {
-			return 0, nil
-		}
-
-		// Unmarshal proto bytes to CapabilityConfig, then to JSON, then extract
-		// the ocr3Configs field.
-		pbCfg := &capabilitiespb.CapabilityConfig{}
-		if err := proto.Unmarshal(capCfg.Config, pbCfg); err != nil {
-			return 0, fmt.Errorf("failed to unmarshal capability config: %w", err)
-		}
-
-		jsonBytes, err := protojson.Marshal(pbCfg)
-		if err != nil {
-			return 0, fmt.Errorf("failed to marshal capability config to JSON: %w", err)
-		}
-
-		var parsed ocr3ConfigsJSON
-		if err := json.Unmarshal(jsonBytes, &parsed); err != nil {
-			return 0, fmt.Errorf("failed to parse capability config JSON: %w", err)
-		}
-
-		if parsed.Ocr3Configs == nil {
-			return 0, nil
-		}
-
-		ocr3Cfg, ok := parsed.Ocr3Configs[ocrConfigKey]
-		if !ok {
-			return 0, nil
-		}
-
-		return ocr3Cfg.ConfigCount, nil
 	}
-
 	return 0, nil
 }
 
+func extractOCR3ConfigCount(rawConfig []byte, ocrConfigKey string) (uint64, error) {
+	if len(rawConfig) == 0 {
+		return 0, nil
+	}
+	pbCfg := &capabilitiespb.CapabilityConfig{}
+	if err := proto.Unmarshal(rawConfig, pbCfg); err != nil {
+		return 0, fmt.Errorf("failed to unmarshal capability config: %w", err)
+	}
+	if pbCfg.Ocr3Configs == nil || pbCfg.Ocr3Configs[ocrConfigKey] == nil {
+		return 0, nil
+	}
+	return pbCfg.Ocr3Configs[ocrConfigKey].ConfigCount, nil
+}
+
 // OCR2OracleConfigToMap converts an OCR2OracleConfig to a protojson-compatible
 // map suitable for injection into a capability config's ocr3Configs field.
 // For example, []byte fields are base64-encoded.
diff --git a/deployment/cre/ocr3/registry_config_test.go b/deployment/cre/ocr3/registry_config_test.go
index 7c2174b031b..394fde489b2 100644
--- a/deployment/cre/ocr3/registry_config_test.go
+++ b/deployment/cre/ocr3/registry_config_test.go
@@ -7,8 +7,82 @@ import (
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"google.golang.org/protobuf/proto"
+
+	capabilitiespb "github.com/smartcontractkit/chainlink-common/pkg/capabilities/pb"
 )
 
+func TestExtractOCR3ConfigCount(t *testing.T) {
+	t.Run("empty input returns zero", func(t *testing.T) {
+		count, err := extractOCR3ConfigCount(nil, OCR3ConfigDefaultKey)
+		require.NoError(t, err)
+		assert.Equal(t, uint64(0), count)
+
+		count, err = extractOCR3ConfigCount([]byte{}, OCR3ConfigDefaultKey)
+		require.NoError(t, err)
+		assert.Equal(t, uint64(0), count)
+	})
+
+	t.Run("invalid protobuf returns error", func(t *testing.T) {
+		_, err := extractOCR3ConfigCount([]byte{0xff, 0xfe}, OCR3ConfigDefaultKey)
+		require.ErrorContains(t, err, "failed to unmarshal capability config")
+	})
+
+	t.Run("nil ocr3Configs map returns zero", func(t *testing.T) {
+		cfg := &capabilitiespb.CapabilityConfig{}
+		raw, err := proto.Marshal(cfg)
+		require.NoError(t, err)
+
+		count, err := extractOCR3ConfigCount(raw, OCR3ConfigDefaultKey)
+		require.NoError(t, err)
+		assert.Equal(t, uint64(0), count)
+	})
+
+	t.Run("missing key returns zero", func(t *testing.T) {
+		cfg := &capabilitiespb.CapabilityConfig{
+			Ocr3Configs: map[string]*capabilitiespb.OCR3Config{
+				"other_key": {ConfigCount: 7},
+			},
+		}
+		raw, err := proto.Marshal(cfg)
+		require.NoError(t, err)
+
+		count, err := extractOCR3ConfigCount(raw, OCR3ConfigDefaultKey)
+		require.NoError(t, err)
+		assert.Equal(t, uint64(0), count)
+	})
+
+	t.Run("returns config count for matching key", func(t *testing.T) {
+		cfg := &capabilitiespb.CapabilityConfig{
+			Ocr3Configs: map[string]*capabilitiespb.OCR3Config{
+				OCR3ConfigDefaultKey: {ConfigCount: 42},
+			},
+		}
+		raw, err := proto.Marshal(cfg)
+		require.NoError(t, err)
+
+		count, err := extractOCR3ConfigCount(raw, OCR3ConfigDefaultKey)
+		require.NoError(t, err)
+		assert.Equal(t, uint64(42), count)
+	})
+
+	t.Run("returns config count for custom key", func(t *testing.T) {
+		cfg := &capabilitiespb.CapabilityConfig{
+			Ocr3Configs: map[string]*capabilitiespb.OCR3Config{
+				"dkg":                {ConfigCount: 1},
+				OCR3ConfigDefaultKey: {ConfigCount: 10},
+				"vault":              {ConfigCount: 5},
+			},
+		}
+		raw, err := proto.Marshal(cfg)
+		require.NoError(t, err)
+
+		count, err := extractOCR3ConfigCount(raw, "vault")
+		require.NoError(t, err)
+		assert.Equal(t, uint64(5), count)
+	})
+}
+
 func TestValidateOCR3Config(t *testing.T) {
 	tests := []struct {
 		name         string