fix(service): tolerate light-client lag when reading the round's DKG network

service/dkg_finalize.go

@@ -14,6 +14,7 @@ import (
 	log "github.com/sirupsen/logrus"
 
 	"github.com/piplabs/story-kernel/enclave"
+	"github.com/piplabs/story-kernel/story"
 	pb "github.com/piplabs/story-kernel/types/pb/v0"
 
 	dkg "go.dedis.ch/kyber/v4/share/dkg/pedersen"
@@ -29,7 +30,7 @@ const (
 	finalizeStageRetryDelay    = 2 * time.Second
 )
 
-func (s *DKGServer) FinalizeDKG(_ context.Context, req *pb.FinalizeDKGRequest) (*pb.FinalizeDKGResponse, error) {
+func (s *DKGServer) FinalizeDKG(ctx context.Context, req *pb.FinalizeDKGRequest) (*pb.FinalizeDKGResponse, error) {
 	codeCommitmentHex := hex.EncodeToString(req.GetCodeCommitment())
 
 	// Validate request
@@ -87,8 +88,17 @@ func (s *DKGServer) FinalizeDKG(_ context.Context, req *pb.FinalizeDKGRequest) (
 	// Snapshotting the share only after this wait also lets more of the response feed drain,
 	// so DistKeyShare() interpolates over a fuller QUAL; that same share is sealed and later
 	// used for dealing (loadFromRoundShare), keeping deals and on-chain coeffs consistent.
-	registrations, err := s.waitForFinalizationRegistrations(codeCommitmentHex, req.GetRound())
+	registrations, err := s.waitForFinalizationRegistrations(ctx, codeCommitmentHex, req.GetRound())
 	if err != nil {
+		// A canceled request is the caller going away, not a server failure.
+		if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
+			log.WithFields(log.Fields{
+				"round": req.GetRound(),
+			}).Warnf("FinalizeDKG aborted while waiting for the finalization stage: %v", err)
+
+			return nil, status.FromContextError(err).Err()
+		}
+
 		log.Errorf("failed to get finalization DKG registrations: %v", err)
 
 		return nil, status.Errorf(codes.Internal, "failed to get finalization DKG registrations")
@@ -228,26 +238,38 @@ func validateFinalizeDKGRequest(req *pb.FinalizeDKGRequest) error {
 // height. Because missing-dealer invalidation happens in the same block that advances the
 // stage to finalization, waiting for that stage guarantees the returned set excludes any
 // invalidated dealer, so the kernel's participants root matches the chain's. It retries
-// while the light client still reports an earlier stage, returning ErrLightClientLag once
-// the retry budget is exhausted.
-func (s *DKGServer) waitForFinalizationRegistrations(codeCommitmentHex string, round uint32) ([]*pb.DKGRegistration, error) {
+// while the light client still reports an earlier stage or has not observed the round at
+// all, returning ErrLightClientLag once the retry budget is exhausted, or ctx.Err() if the
+// caller goes away mid-wait.
+func (s *DKGServer) waitForFinalizationRegistrations(ctx context.Context, codeCommitmentHex string, round uint32) ([]*pb.DKGRegistration, error) {
 	for attempt := range finalizeStageRetryAttempts {
-		network, err := s.QueryClient.GetDKGNetwork(context.Background(), codeCommitmentHex, round)
-		if err != nil {
+		network, err := s.QueryClient.GetDKGNetwork(ctx, codeCommitmentHex, round)
+		if err != nil && !errors.Is(err, story.ErrDKGNetworkNotFound) {
 			return nil, err
 		}
 
-		if network.GetStage() >= pb.DKGStage_DKG_STAGE_FINALIZATION {
-			return s.QueryClient.GetAllParticipantDKGRegistrations(context.Background(), codeCommitmentHex, round)
+		if err == nil && network.GetStage() >= pb.DKGStage_DKG_STAGE_FINALIZATION {
+			return s.QueryClient.GetAllParticipantDKGRegistrations(ctx, codeCommitmentHex, round)
+		}
+
+		// A not-found read is the same lag shape as an earlier stage: the light client
+		// has not verified up to the round's start block yet, so keep retrying.
+		stage := "not visible"
+		if err == nil {
+			stage = network.GetStage().String()
 		}
 
 		log.WithFields(log.Fields{
 			"round":   round,
-			"stage":   network.GetStage().String(),
+			"stage":   stage,
 			"attempt": attempt + 1,
 		}).Warn("FinalizeDKG: light client has not observed the finalization stage yet, retrying")
 
-		time.Sleep(finalizeStageRetryDelay)
+		select {
+		case <-ctx.Done():
+			return nil, ctx.Err()
+		case <-time.After(finalizeStageRetryDelay):
+		}
 	}
 
 	return nil, fmt.Errorf("%w: round %d did not reach finalization stage after %d retries",

service/dkg_finalize_wait_test.go

@@ -1,6 +1,7 @@
 package service
 
 import (
+	"context"
 	"errors"
 	"testing"
 
@@ -32,7 +33,7 @@ func TestWaitForFinalizationRegistrations_WaitsForStage(t *testing.T) {
 
 	s := &DKGServer{QueryClient: stub}
 
-	regs, err := s.waitForFinalizationRegistrations("cc", 42)
+	regs, err := s.waitForFinalizationRegistrations(t.Context(), "cc", 42)
 	require.NoError(t, err)
 	require.Equal(t, filteredRegs(), regs)
 	require.Equal(t, int32(2), stub.netCalls.Load(), "polls stage until finalization (1 lag + 1 caught up)")
@@ -49,7 +50,7 @@ func TestWaitForFinalizationRegistrations_Immediate(t *testing.T) {
 
 	s := &DKGServer{QueryClient: stub}
 
-	regs, err := s.waitForFinalizationRegistrations("cc", 42)
+	regs, err := s.waitForFinalizationRegistrations(t.Context(), "cc", 42)
 	require.NoError(t, err)
 	require.Equal(t, filteredRegs(), regs)
 	require.Equal(t, int32(1), stub.netCalls.Load())
@@ -66,14 +67,33 @@ func TestWaitForFinalizationRegistrations_RetryExhausted(t *testing.T) {
 
 	s := &DKGServer{QueryClient: stub}
 
-	regs, err := s.waitForFinalizationRegistrations("cc", 42)
+	regs, err := s.waitForFinalizationRegistrations(t.Context(), "cc", 42)
 	require.Error(t, err)
 	require.Nil(t, regs)
 	require.True(t, errors.Is(err, ErrLightClientLag), "exhaustion must preserve ErrLightClientLag")
 	require.Equal(t, int32(finalizeStageRetryAttempts), stub.netCalls.Load())
 	require.Equal(t, int32(0), stub.regCalls.Load(), "must never read the participant set while lagging")
 }
 
+// TestWaitForFinalizationRegistrations_NotFoundThenStage verifies a not-found read (light
+// client has not observed the round at all) is retried like an earlier stage instead of
+// failing the call.
+func TestWaitForFinalizationRegistrations_NotFoundThenStage(t *testing.T) {
+	stub := &stubQueryClient{
+		netErrs:       []error{notFoundErr()}, // call 1: round not visible yet
+		networks:      []*pb.DKGNetwork{{Round: 42, Stage: pb.DKGStage_DKG_STAGE_FINALIZATION}},
+		registrations: [][]*pb.DKGRegistration{filteredRegs()},
+	}
+
+	s := &DKGServer{QueryClient: stub}
+
+	regs, err := s.waitForFinalizationRegistrations(t.Context(), "cc", 42)
+	require.NoError(t, err)
+	require.Equal(t, filteredRegs(), regs)
+	require.Equal(t, int32(2), stub.netCalls.Load(), "polls through the not-found read (1 lag + 1 caught up)")
+	require.Equal(t, int32(1), stub.regCalls.Load())
+}
+
 // TestWaitForFinalizationRegistrations_NetworkError verifies a network query error is returned
 // directly without retrying.
 func TestWaitForFinalizationRegistrations_NetworkError(t *testing.T) {
@@ -82,8 +102,27 @@ func TestWaitForFinalizationRegistrations_NetworkError(t *testing.T) {
 
 	s := &DKGServer{QueryClient: stub}
 
-	regs, err := s.waitForFinalizationRegistrations("cc", 42)
+	regs, err := s.waitForFinalizationRegistrations(t.Context(), "cc", 42)
 	require.ErrorIs(t, err, wantErr)
 	require.Nil(t, regs)
 	require.Equal(t, int32(1), stub.netCalls.Load(), "network error must not be retried")
 }
+
+// TestWaitForFinalizationRegistrations_ContextCanceled verifies a caller going away
+// mid-wait stops the retry loop promptly instead of burning the remaining budget.
+func TestWaitForFinalizationRegistrations_ContextCanceled(t *testing.T) {
+	ctx, cancel := context.WithCancel(t.Context())
+	cancel()
+
+	stub := &stubQueryClient{
+		networks: []*pb.DKGNetwork{{Round: 42, Stage: pb.DKGStage_DKG_STAGE_DEALING}},
+	}
+
+	s := &DKGServer{QueryClient: stub}
+
+	regs, err := s.waitForFinalizationRegistrations(ctx, "cc", 42)
+	require.ErrorIs(t, err, context.Canceled)
+	require.False(t, errors.Is(err, ErrLightClientLag))
+	require.Nil(t, regs)
+	require.Equal(t, int32(1), stub.netCalls.Load(), "cancellation must stop retries promptly")
+}

service/dkg_generate_key.go

@@ -3,19 +3,36 @@ package service
 import (
 	"context"
 	"encoding/hex"
+	"fmt"
+	"time"
 
 	ecrypto "github.com/ethereum/go-ethereum/crypto"
 	"github.com/pkg/errors"
 	log "github.com/sirupsen/logrus"
 
 	"github.com/piplabs/story-kernel/enclave"
+	"github.com/piplabs/story-kernel/story"
 	pb "github.com/piplabs/story-kernel/types/pb/v0"
 
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 )
 
-func (s *DKGServer) GenerateAndSealKey(_ context.Context, req *pb.GenerateAndSealKeyRequest) (*pb.GenerateAndSealKeyResponse, error) {
+const (
+	// registrationNetworkRetryAttempts/Delay bound how long GenerateAndSealKey waits for
+	// the light client to observe the round's DKGNetwork at a round boundary. The verified
+	// head has been observed to trail the chain tip by ~40s when a round starts, so the
+	// budget (30 x 3s = 90s) outlasts that with margin.
+	registrationNetworkRetryAttempts = 30
+	registrationNetworkRetryDelay    = 3 * time.Second
+
+	// registrationLagNormalAttempts is how many retry attempts fall within the ~40s lag
+	// documented above. Within it a not-found read is expected steady-state behavior and
+	// logs at Debug; beyond it the wait is abnormal and escalates to Warn.
+	registrationLagNormalAttempts = 15
+)
+
+func (s *DKGServer) GenerateAndSealKey(ctx context.Context, req *pb.GenerateAndSealKeyRequest) (*pb.GenerateAndSealKeyResponse, error) {
 	codeCommitmentHex := hex.EncodeToString(req.GetCodeCommitment())
 
 	// Validate the request
@@ -61,19 +78,44 @@ func (s *DKGServer) GenerateAndSealKey(_ context.Context, req *pb.GenerateAndSea
 
 	// Only fetch the DKG network (not registrations) since no registrations
 	// exist yet at key generation time.
-	network, err := s.QueryClient.GetDKGNetwork(context.Background(), codeCommitmentHex, req.Round)
+	network, err := s.waitForDKGNetworkCreation(ctx, codeCommitmentHex, req.Round,
+		registrationNetworkRetryAttempts, registrationNetworkRetryDelay)
 	if err != nil {
+		// A canceled request is the caller going away, not a server failure.
+		if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
+			log.WithFields(log.Fields{
+				"round":           req.Round,
+				"code_commitment": codeCommitmentHex,
+			}).Warnf("GenerateAndSealKey aborted while waiting for DKG network: %v", err)
+
+			return nil, status.FromContextError(err).Err()
+		}
+
 		log.WithFields(log.Fields{
 			"round":           req.Round,
 			"code_commitment": codeCommitmentHex,
 		}).Errorf("failed to get DKG network: %v", err)
 
+		if errors.Is(err, ErrLightClientLag) {
+			return nil, status.Errorf(codes.Unavailable, "DKG network not yet visible to the light client")
+		}
+
 		return nil, status.Errorf(codes.Internal, "failed to get DKG network")
 	}
 
 	// Verify the DKG start block is on the canonical chain.
 	// This ensures the DKG round was legitimately initiated on-chain before generating keys.
-	if err := s.verifyDKGStartBlock(context.Background(), network); err != nil {
+	if err := s.verifyDKGStartBlock(ctx, network); err != nil {
+		// A canceled request is the caller going away, not a verification failure.
+		if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
+			log.WithFields(log.Fields{
+				"round":           req.Round,
+				"code_commitment": codeCommitmentHex,
+			}).Warnf("GenerateAndSealKey aborted during start block verification: %v", err)
+
+			return nil, status.FromContextError(err).Err()
+		}
+
 		log.WithFields(log.Fields{
 			"round":              req.Round,
 			"code_commitment":    codeCommitmentHex,
@@ -127,6 +169,52 @@ func (s *DKGServer) GenerateAndSealKey(_ context.Context, req *pb.GenerateAndSea
 	}, nil
 }
 
+// waitForDKGNetworkCreation blocks until the round's on-chain DKGNetwork record becomes
+// visible at the light client's verified height, then returns it. The consensus client only
+// requests key generation for a round it observed on-chain, so a not-found read here is
+// presumed to be light-client lag rather than a nonexistent round; any other error fails
+// fast. Returns ErrLightClientLag once the retry budget is exhausted, or ctx.Err() if the
+// caller goes away mid-wait.
+func (s *DKGServer) waitForDKGNetworkCreation(
+	ctx context.Context,
+	codeCommitmentHex string,
+	round uint32,
+	attempts int,
+	delay time.Duration,
+) (*pb.DKGNetwork, error) {
+	for attempt := range attempts {
+		network, err := s.QueryClient.GetDKGNetwork(ctx, codeCommitmentHex, round)
+		if err == nil {
+			return network, nil
+		}
+
+		if !errors.Is(err, story.ErrDKGNetworkNotFound) {
+			return nil, err
+		}
+
+		if attempt+1 == attempts {
+			break
+		}
+
+		fields := log.Fields{"round": round, "attempt": attempt + 1}
+		msg := "GenerateAndSealKey: DKG network not yet visible to light client, retrying"
+		if attempt+1 > registrationLagNormalAttempts {
+			log.WithFields(fields).Warn(msg)
+		} else {
+			log.WithFields(fields).Debug(msg)
+		}
+
+		select {
+		case <-ctx.Done():
+			return nil, ctx.Err()
+		case <-time.After(delay):
+		}
+	}
+
+	return nil, fmt.Errorf("%w: round %d DKG network not visible after %d attempts",
+		ErrLightClientLag, round, attempts)
+}
+
 func validateGenerateAndSealKeyRequest(req *pb.GenerateAndSealKeyRequest) error {
 	if len(req.GetAddress()) == 0 {
 		return errors.New("validator address is required but missing")