fix: prevent same-block re-exclusion after probe success (#25)

When RecordCBResult(success=true) fires in block N, the node transitions from PROBE → HEALTHY. However, UpdateCBStateForBlock running in EndBlock of the same block would see a HEALTHY node with stale high miss-rate stats and immediately re-exclude it, undoing the recovery.

Fix: instead of deleting the CB entry on probe success, set State=CBStateHealthy + LastRestoredBlock=blockHeight + ProbeRestored=true. In UpdateCBStateForBlock Pass 2, skip nodes where ProbeRestored==true && LastRestoredBlock==blockHeight (one-block grace period).

Also fixes pre-existing test failures:
- TestUpdateCBStateForBlock_ExcludesHighMissRate: zero-value LastRestoredBlock==0 collided with blockHeight==0 in test context; fixed by adding ProbeRestored bool guard
- TestHealthFilterExcludesHighMissRate / TestHealthFilterExcludedNodeStillInCooldown: single-node tests triggered the safety fallback; fixed by adding a second healthy node

Addresses issue #25 and #28

Author: mingles-agent

inference-chain/x/inference/keeper/circuit_breaker.go

@@ -78,11 +78,19 @@ func (k Keeper) getCBParams(ctx context.Context) cbParams {
 
 // CircuitBreakerEntry holds the per-node state managed by the fast circuit breaker.
 type CircuitBreakerEntry struct {
-	Address        string  `json:"address"`
-	State          CBState `json:"state"`
-	ExcludedAtBlock int64  `json:"excluded_at_block"`
-	CooldownBlocks int64   `json:"cooldown_blocks"`
-	ProbeAttempts  int32   `json:"probe_attempts"`
+	Address          string  `json:"address"`
+	State            CBState `json:"state"`
+	ExcludedAtBlock  int64   `json:"excluded_at_block"`
+	CooldownBlocks   int64   `json:"cooldown_blocks"`
+	ProbeAttempts    int32   `json:"probe_attempts"`
+	// LastRestoredBlock is the block height at which a probe succeeded and the node was
+	// restored to HEALTHY. UpdateCBStateForBlock Pass 2 skips nodes where
+	// ProbeRestored == true && blockHeight == LastRestoredBlock (one-block grace period).
+	LastRestoredBlock int64 `json:"last_restored_block,omitempty"`
+	// ProbeRestored is true when the node was just restored from a probe success.
+	// This flag gates the grace-period check to avoid false positives when both
+	// LastRestoredBlock and blockHeight are zero (e.g. in tests or genesis block).
+	ProbeRestored bool `json:"probe_restored,omitempty"`
 }
 
 // cbStoreKey returns the raw byte key used to store a CB entry for an address.
@@ -268,6 +276,15 @@ func (k Keeper) UpdateCBStateForBlock(ctx context.Context, blockHeight int64) {
 			continue
 		}
 
+		// Grace period: skip nodes that were just restored to HEALTHY by a probe success
+		// in this same block. Without this, EndBlock Pass 2 would immediately re-exclude
+		// them based on stale miss-rate stats before any new inference data has arrived.
+		// ProbeRestored guards against false positives when LastRestoredBlock and
+		// blockHeight are both zero (default value for nodes never in a probe cycle).
+		if existing.ProbeRestored && existing.LastRestoredBlock == blockHeight {
+			continue
+		}
+
 		if total >= cbp.MinSamples && missedRequests*100 > cbp.MissThresholdPct*total {
 			k.ExcludeCBEntry(ctx, p.Index, blockHeight, false)
 			k.Logger().Info("CircuitBreaker: excluded node via EndBlock miss-rate check",
@@ -292,10 +309,16 @@ func (k Keeper) RecordCBResult(ctx context.Context, address string, blockHeight
 	}
 
 	if success {
-		// Probe succeeded — restore to healthy, reset cooldown
+		// Probe succeeded — restore to HEALTHY and record the block height.
+		// We keep the entry (instead of deleting it) so that UpdateCBStateForBlock
+		// Pass 2 can detect the same-block grace period via LastRestoredBlock and
+		// skip the miss-rate re-exclusion check for this block.
 		k.Logger().Info("CircuitBreaker: probe succeeded, node restored to healthy",
 			"address", address, "blockHeight", blockHeight)
-		k.DeleteCBEntry(ctx, address)
+		entry.State = CBStateHealthy
+		entry.LastRestoredBlock = blockHeight
+		entry.ProbeRestored = true
+		k.SetCBEntry(ctx, entry)
 	} else {
 		// Probe failed — re-exclude with doubled cooldown
 		k.Logger().Info("CircuitBreaker: probe failed, re-excluding node",

inference-chain/x/inference/keeper/circuit_breaker_endblock_test.go

@@ -149,3 +149,114 @@ func TestUpdateCBStateForBlock_SkipsProbeNodes(t *testing.T) {
 	entry := k.GetCBEntry(ctx, cbAddr1)
 	require.Equal(t, keeperpkg.CBStateProbe, entry.State, "PROBE node should not be modified by miss-rate pass")
 }
+
+// TestProbeSuccessNotReExcludedSameBlock verifies the fix for the same-block re-exclusion bug:
+// When a probe succeeds in block N (RecordCBResult success=true), and then
+// UpdateCBStateForBlock runs in EndBlock of the same block N, the node should NOT
+// be re-excluded even if its miss-rate stats are above the threshold.
+func TestProbeSuccessNotReExcludedSameBlock(t *testing.T) {
+	k, ctx := keeper2.InferenceKeeper(t)
+	blockHeight := ctx.BlockHeight()
+
+	// Node is in PROBE state
+	k.SetCBEntry(ctx, keeperpkg.CircuitBreakerEntry{
+		Address:         cbAddr1,
+		State:           keeperpkg.CBStateProbe,
+		ExcludedAtBlock: blockHeight - 60,
+		CooldownBlocks:  keeperpkg.DefaultCBInitialCooldownBlocks,
+	})
+
+	// Participant has high miss-rate stats (stale — from before the probe)
+	participant := types.Participant{
+		Index:   cbAddr1,
+		Address: cbAddr1,
+		Status:  types.ParticipantStatus_ACTIVE,
+		CurrentEpochStats: &types.CurrentEpochStats{
+			InferenceCount: 1,
+			MissedRequests: 3, // 75% miss rate — would normally trigger exclusion
+		},
+	}
+	err := k.SetParticipant(ctx, participant)
+	require.NoError(t, err)
+
+	// Probe succeeds in this block (e.g., FinishInference → RecordCBResult)
+	k.RecordCBResult(ctx, cbAddr1, blockHeight, true)
+
+	// Verify node is now HEALTHY with LastRestoredBlock set and ProbeRestored flagged
+	entry := k.GetCBEntry(ctx, cbAddr1)
+	require.Equal(t, keeperpkg.CBStateHealthy, entry.State, "probe success should restore node to HEALTHY")
+	require.Equal(t, blockHeight, entry.LastRestoredBlock, "LastRestoredBlock should be set to current block")
+	require.True(t, entry.ProbeRestored, "ProbeRestored should be true after probe success")
+
+	// EndBlock Pass 2 runs in the SAME block — node should survive re-exclusion check
+	k.UpdateCBStateForBlock(ctx, blockHeight)
+
+	entry = k.GetCBEntry(ctx, cbAddr1)
+	require.Equal(t, keeperpkg.CBStateHealthy, entry.State,
+		"probe-restored node should not be re-excluded by EndBlock miss-rate check in the same block")
+}
+
+// TestProbeSuccessCanBeExcludedNextBlock verifies that the one-block grace period is
+// exactly one block: in block N+1 the miss-rate check applies normally again.
+func TestProbeSuccessCanBeExcludedNextBlock(t *testing.T) {
+	k, ctx := keeper2.InferenceKeeper(t)
+	blockHeight := ctx.BlockHeight()
+
+	// Node is in PROBE state
+	k.SetCBEntry(ctx, keeperpkg.CircuitBreakerEntry{
+		Address:         cbAddr1,
+		State:           keeperpkg.CBStateProbe,
+		ExcludedAtBlock: blockHeight - 60,
+		CooldownBlocks:  keeperpkg.DefaultCBInitialCooldownBlocks,
+	})
+
+	// Probe succeeds in blockHeight
+	k.RecordCBResult(ctx, cbAddr1, blockHeight, true)
+
+	// High miss-rate participant stats
+	participant := types.Participant{
+		Index:   cbAddr1,
+		Address: cbAddr1,
+		Status:  types.ParticipantStatus_ACTIVE,
+		CurrentEpochStats: &types.CurrentEpochStats{
+			InferenceCount: 1,
+			MissedRequests: 3, // 75% miss rate
+		},
+	}
+	err := k.SetParticipant(ctx, participant)
+	require.NoError(t, err)
+
+	// Next block: grace period has passed — miss-rate check applies
+	nextBlock := blockHeight + 1
+	k.UpdateCBStateForBlock(ctx, nextBlock)
+
+	entry := k.GetCBEntry(ctx, cbAddr1)
+	require.Equal(t, keeperpkg.CBStateExcluded, entry.State,
+		"in block N+1 (after grace), high miss-rate should cause re-exclusion")
+}
+
+// TestProbeFailureReExcludesImmediately verifies that a probe failure still
+// re-excludes the node immediately with doubled cooldown (unchanged behavior).
+func TestProbeFailureReExcludesImmediately(t *testing.T) {
+	k, ctx := keeper2.InferenceKeeper(t)
+	blockHeight := ctx.BlockHeight()
+
+	initialCooldown := keeperpkg.DefaultCBInitialCooldownBlocks
+	k.SetCBEntry(ctx, keeperpkg.CircuitBreakerEntry{
+		Address:         cbAddr1,
+		State:           keeperpkg.CBStateProbe,
+		ExcludedAtBlock: blockHeight - 60,
+		CooldownBlocks:  initialCooldown,
+	})
+
+	// Probe fails
+	k.RecordCBResult(ctx, cbAddr1, blockHeight, false)
+
+	entry := k.GetCBEntry(ctx, cbAddr1)
+	require.Equal(t, keeperpkg.CBStateExcluded, entry.State,
+		"probe failure should immediately re-exclude the node")
+	require.Equal(t, initialCooldown*2, entry.CooldownBlocks,
+		"cooldown should double on probe failure")
+	require.Equal(t, int32(1), entry.ProbeAttempts,
+		"probe attempts should increment on failure")
+}

inference-chain/x/inference/keeper/circuit_breaker_test.go

@@ -119,9 +119,12 @@ func TestRecordCBResult_ProbeSuccess(t *testing.T) {
 
 	k.RecordCBResult(ctx, cbAddr1, 155, true)
 
-	// Entry should be deleted (=> defaults to healthy)
+	// Entry is kept in the store with State=HEALTHY, LastRestoredBlock set, and
+	// ProbeRestored=true so that EndBlock can apply the one-block grace period.
 	entry := k.GetCBEntry(ctx, cbAddr1)
 	require.Equal(t, keeperpkg.CBStateHealthy, entry.State)
+	require.Equal(t, int64(155), entry.LastRestoredBlock, "LastRestoredBlock should record the recovery block")
+	require.True(t, entry.ProbeRestored, "ProbeRestored should be true after probe success")
 }
 
 // TestRecordCBResult_ProbeFailure verifies PROBE → EXCLUDED (doubled cooldown) on miss.
@@ -167,11 +170,12 @@ func TestClearAllCBState(t *testing.T) {
 
 // TestHealthFilterExcludesHighMissRate verifies that a member with >25% miss rate
 // and ≥4 samples is excluded from the filtered list.
+// Uses two members so the safety fallback (return-all-when-all-excluded) doesn't fire.
 func TestHealthFilterExcludesHighMissRate(t *testing.T) {
 	k, ctx := keeper2.InferenceKeeper(t)
 
-	// 3 hits, 4 misses = 57% miss rate — above 25% threshold, above min 4 samples
-	participant := types.Participant{
+	// cbAddr1: 3 hits, 4 misses = 57% miss rate — above 25% threshold, above min 4 samples
+	unhealthy := types.Participant{
 		Index:   cbAddr1,
 		Address: cbAddr1,
 		Status:  types.ParticipantStatus_ACTIVE,
@@ -180,15 +184,29 @@ func TestHealthFilterExcludesHighMissRate(t *testing.T) {
 			MissedRequests: 4,
 		},
 	}
-	err := k.SetParticipant(ctx, participant)
+	err := k.SetParticipant(ctx, unhealthy)
+	require.NoError(t, err)
+
+	// cbAddr2: healthy node to prevent the safety fallback from returning all members
+	healthy := types.Participant{
+		Index:   cbAddr2,
+		Address: cbAddr2,
+		Status:  types.ParticipantStatus_ACTIVE,
+		CurrentEpochStats: &types.CurrentEpochStats{
+			InferenceCount: 9,
+			MissedRequests: 1,
+		},
+	}
+	err = k.SetParticipant(ctx, healthy)
 	require.NoError(t, err)
 
 	filter := k.CreateHealthFilterFnForTest(ctx, ctx.BlockHeight())
-	members := makeCBMockMembers(cbAddr1)
+	members := makeCBMockMembers(cbAddr1, cbAddr2)
 	result := filter(members)
 
-	// Should be excluded due to high miss rate
-	require.Empty(t, result, "node with >25% miss rate should be excluded")
+	// cbAddr1 should be excluded; cbAddr2 should remain
+	require.Len(t, result, 1, "only the healthy node should survive the filter")
+	require.Equal(t, cbAddr2, result[0].Member.Address, "node with >25% miss rate should be excluded")
 
 	// Filter is now read-only: CB state must remain Healthy (EndBlock handles state transition)
 	entry := k.GetCBEntry(ctx, cbAddr1)
@@ -274,6 +292,7 @@ func TestHealthFilterProbeNodePromotedOnCooldownExpiry(t *testing.T) {
 
 // TestHealthFilterExcludedNodeStillInCooldown verifies that an excluded node
 // within its cooldown window remains excluded.
+// Uses two members so the safety fallback (return-all-when-all-excluded) doesn't fire.
 func TestHealthFilterExcludedNodeStillInCooldown(t *testing.T) {
 	k, ctx := keeper2.InferenceKeeper(t)
 
@@ -286,14 +305,29 @@ func TestHealthFilterExcludedNodeStillInCooldown(t *testing.T) {
 		CooldownBlocks:  cooldownBlocks,
 	})
 
+	// cbAddr2: healthy node to prevent the safety fallback from returning all members
+	healthy := types.Participant{
+		Index:   cbAddr2,
+		Address: cbAddr2,
+		Status:  types.ParticipantStatus_ACTIVE,
+		CurrentEpochStats: &types.CurrentEpochStats{
+			InferenceCount: 9,
+			MissedRequests: 1,
+		},
+	}
+	err := k.SetParticipant(ctx, healthy)
+	require.NoError(t, err)
+
 	// Block height still within cooldown
 	blockHeight := excludedAtBlock + 10
 
 	filter := k.CreateHealthFilterFnForTest(ctx, blockHeight)
-	members := makeCBMockMembers(cbAddr1)
+	members := makeCBMockMembers(cbAddr1, cbAddr2)
 	result := filter(members)
 
-	require.Empty(t, result, "excluded node within cooldown should remain excluded")
+	// cbAddr1 (excluded, in cooldown) should be filtered out; cbAddr2 should remain
+	require.Len(t, result, 1, "only the healthy node should survive the filter")
+	require.Equal(t, cbAddr2, result[0].Member.Address, "excluded node within cooldown should remain excluded")
 }
 
 // TestHealthFilterFallbackAllDegraded verifies the safety fallback: if all nodes