fix(dkg): restart decrypt worker whenever an active round exists

client/x/dkg/keeper/dkg_svc.go

@@ -90,6 +90,14 @@ func releaseDKGSvc(round uint32) {
 // is cut short — so this value must stay above 60s to be meaningful.
 const dkgAsyncTimeout = 2 * time.Minute
 
+// decryptKernelCallTimeout bounds a single kernel RPC made by the decrypt worker.
+// The worker runs on a process-lifetime context.Background with no deadline of its
+// own, so without a per-call timeout a wedged kernel (accepts the connection but
+// never responds) blocks the worker loop forever and decryptWorkerRunning is never
+// reset. 60s matches the kernel-operation budget in dkgAsyncTimeout and the
+// contract-call timeout in contract_client.go.
+const decryptKernelCallTimeout = 60 * time.Second
+
 // dkgAsyncContext creates a new context for async DKG service goroutines with a timeout.
 // This replaces the consensus context that would otherwise be canceled after block processing.
 func dkgAsyncContext() (context.Context, context.CancelFunc) {
@@ -100,6 +108,18 @@ var decryptWorkerRunning atomic.Bool
 
 // ResumeDKGService reloads unfinished DKG sessions and resumes their execution safely without spawning duplicate goroutines.
 func (k *Keeper) ResumeDKGService(ctx context.Context, dkgNetwork *types.DKGNetwork) {
+	// The decrypt worker serves the latest ACTIVE round, which may differ from the
+	// latest round once the next round has opened. StartDecryptWorker is idempotent.
+	if active, err := k.getLatestActiveDKGNetwork(ctx); err != nil {
+		log.Warn(ctx, "Failed to get latest active DKG round while resuming decrypt worker", err)
+	} else if active != nil {
+		k.StartDecryptWorker()
+	} else {
+		// Debug (not Info): this runs every block; a missing active round is the
+		// normal steady state and Info would be too noisy.
+		log.Debug(ctx, "No active DKG round; decrypt worker not started")
+	}
+
 	session, err := k.stateManager.GetSession(dkgNetwork.Round)
 	if err != nil {
 		log.Error(ctx, "Failed to get DKG session while resuming the DKG service", err)
@@ -279,7 +299,10 @@ func (k *Keeper) resumeFailedSession(ctx context.Context, session *types.DKGSess
 // goroutine exits. The decrypt worker must run for the lifetime of the process.
 func (k *Keeper) StartDecryptWorker() {
 	if !decryptWorkerRunning.CompareAndSwap(false, true) {
-		// already running
+		// Debug (not Info): this runs every block; skipping because the worker is
+		// already running is the normal steady state and Info would be too noisy.
+		log.Debug(context.Background(), "Decrypt worker already running; skipping start")
+
 		return
 	}
 
@@ -620,8 +643,13 @@ func (k *Keeper) computePartialDecrypt(ctx context.Context, session *types.DKGSe
 		return
 	}
 
+	// Bound the kernel RPC so a wedged kernel cannot block the worker loop
+	// indefinitely. The worker's context has no deadline of its own.
+	callCtx, cancel := context.WithTimeout(ctx, decryptKernelCallTimeout)
+	defer cancel()
+
 	kernelStart := time.Now()
-	resp, err := client.PartialDecryptTDH2(ctx, &types.PartialDecryptTDH2Request{
+	resp, err := client.PartialDecryptTDH2(callCtx, &types.PartialDecryptTDH2Request{
 		CodeCommitment:  session.CodeCommitment,
 		Round:           session.Round,
 		Ciphertext:      req.Ciphertext,

client/x/dkg/keeper/dkg_svc_test.go

@@ -1871,6 +1871,133 @@ func TestResumeDKGService_NoSession(t *testing.T) {
 	k.ResumeDKGService(ctx, dkgNetwork)
 }
 
+// TestResumeDKGService_ActiveRoundStartsWorker verifies that the decrypt worker is
+// restarted whenever an active DKG round exists, even if the latest round has already
+// advanced to a pre-active stage (e.g. the next round opened registration). This is
+// the restart-window hole from issue piplabs/story#854.
+func TestResumeDKGService_ActiveRoundStartsWorker(t *testing.T) {
+	if testing.Short() {
+		t.Skip("skipping worker test in short mode")
+	}
+
+	k, _, _, ctx := setupDKGKeeperWithMocks(t)
+
+	sm, err := NewStateManager(t.TempDir())
+	require.NoError(t, err)
+	k.stateManager = sm
+
+	resetDKGSvcRound()
+	defer resetDKGSvcRound()
+	// Force a fresh worker spawn so the assertion exercises the new branch, and
+	// leave decryptWorkerRunning true afterwards. The worker is a process-lifetime
+	// singleton; resetting the flag to false here would let a later mock-backed test
+	// spawn a real worker that calls its gomock client after that test completes.
+	decryptWorkerRunning.Store(false)
+
+	// Detach the mock contract client so the process-lifetime worker goroutine
+	// returns early at its nil guard instead of issuing unexpected mock calls
+	// after the test completes.
+	k.contractClient = nil
+
+	// An earlier round is still the active (completed) round.
+	activeNetwork := &types.DKGNetwork{
+		Round: 5,
+		Stage: types.DKGStageActive,
+	}
+	require.NoError(t, k.setDKGNetwork(ctx, activeNetwork))
+	require.NoError(t, k.setLatestActiveRound(ctx, activeNetwork))
+
+	// The latest round has opened registration for the NEXT round; its session is in a
+	// valid (non-stuck) phase, so none of the legacy branches would start the worker.
+	session := &types.DKGSession{
+		Round: 6,
+		Phase: types.PhaseInitialized,
+	}
+	require.NoError(t, sm.CreateSession(ctx, session))
+
+	dkgNetwork := &types.DKGNetwork{
+		Round: 6,
+		Stage: types.DKGStageRegistration,
+	}
+
+	k.ResumeDKGService(ctx, dkgNetwork)
+
+	require.True(t, decryptWorkerRunning.Load(), "worker should run whenever an active round exists")
+}
+
+// TestResumeDKGService_NoActiveRoundNoWorker verifies that on a fresh chain with no
+// active round, ResumeDKGService does not start the decrypt worker via the new branch.
+func TestResumeDKGService_NoActiveRoundNoWorker(t *testing.T) {
+	k, _, _, ctx := setupDKGKeeperWithMocks(t)
+
+	sm, err := NewStateManager(t.TempDir())
+	require.NoError(t, err)
+	k.stateManager = sm
+
+	resetDKGSvcRound()
+	defer resetDKGSvcRound()
+	decryptWorkerRunning.Store(false)
+	// Restore the singleton flag to true on exit so a later mock-backed test does
+	// not spawn a real worker that outlives it (see the note above).
+	defer decryptWorkerRunning.Store(true)
+
+	// Latest round session is in a valid (non-stuck) phase and no active round exists.
+	session := &types.DKGSession{
+		Round: 1,
+		Phase: types.PhaseInitialized,
+	}
+	require.NoError(t, sm.CreateSession(ctx, session))
+
+	dkgNetwork := &types.DKGNetwork{
+		Round: 1,
+		Stage: types.DKGStageRegistration,
+	}
+
+	k.ResumeDKGService(ctx, dkgNetwork)
+
+	require.False(t, decryptWorkerRunning.Load(), "worker must not start when no active round exists")
+}
+
+// TestResumeDKGService_CompletedActiveStartsWorker verifies the pre-existing behavior:
+// when the latest round is itself active with a completed session, the worker is started.
+func TestResumeDKGService_CompletedActiveStartsWorker(t *testing.T) {
+	if testing.Short() {
+		t.Skip("skipping worker test in short mode")
+	}
+
+	k, _, _, ctx := setupDKGKeeperWithMocks(t)
+
+	sm, err := NewStateManager(t.TempDir())
+	require.NoError(t, err)
+	k.stateManager = sm
+
+	resetDKGSvcRound()
+	defer resetDKGSvcRound()
+	// Force a fresh spawn for the assertion; leave the flag true on exit (see the
+	// note in TestResumeDKGService_ActiveRoundStartsWorker).
+	decryptWorkerRunning.Store(false)
+
+	// Detach the mock contract client so the process-lifetime worker goroutine
+	// returns early at its nil guard instead of issuing unexpected mock calls
+	// after the test completes.
+	k.contractClient = nil
+
+	session := &types.DKGSession{
+		Round: 7,
+		Phase: types.PhaseCompleted,
+	}
+	require.NoError(t, sm.CreateSession(ctx, session))
+
+	dkgNetwork := &types.DKGNetwork{
+		Round: 7,
+		Stage: types.DKGStageActive,
+	}
+
+	k.ResumeDKGService(ctx, dkgNetwork)
+
+	require.True(t, decryptWorkerRunning.Load(), "worker should run when latest active round session is completed")
+}
+
 // --- resumeFailedSession ---
 
 func TestResumeFailedSession_DealingStage(t *testing.T) {
@@ -2089,3 +2216,102 @@ func TestResumeFailedSession_ActiveStage(t *testing.T) {
 	require.Equal(t, types.PhaseCompleted, got.Phase,
 		"active stage: handleDKGComplete should advance phase to PhaseCompleted")
 }
+
+// --- decrypt worker kernel-call deadline (issue piplabs/story#854 family) ---
+
+// TestComputePartialDecrypt_DerivesCallDeadline verifies that the kernel RPC is
+// invoked with a bounded context even when the worker passes context.Background
+// (which has no deadline). This proves decryptKernelCallTimeout is wired so a
+// wedged kernel cannot block the worker loop indefinitely.
+func TestComputePartialDecrypt_DerivesCallDeadline(t *testing.T) {
+	t.Parallel()
+
+	ctx := context.Background()
+	ctrl := gomock.NewController(t)
+
+	cc := []byte("cc-deadline")
+	mockKernel := dkgtestutil.NewMockKernelServiceClient(ctrl)
+
+	router := NewKernelRouter(nil, nil)
+	router.RegisterClient(cc, mockKernel)
+
+	k := &Keeper{kernelRouter: router}
+
+	session := &types.DKGSession{
+		Round: 1, Index: 1, GlobalPubKey: []byte("pub"), CodeCommitment: cc,
+	}
+	req := types.DecryptRequest{Ciphertext: []byte("ct"), Label: make([]byte, 32)}
+
+	var (
+		gotDeadline bool
+		remaining   time.Duration
+	)
+	mockKernel.EXPECT().PartialDecryptTDH2(gomock.Any(), gomock.Any()).
+		DoAndReturn(func(callCtx context.Context, _ *types.PartialDecryptTDH2Request, _ ...interface{}) (*types.PartialDecryptTDH2Response, error) {
+			deadline, ok := callCtx.Deadline()
+			gotDeadline = ok
+			if ok {
+				remaining = time.Until(deadline)
+			}
+			return &types.PartialDecryptTDH2Response{
+				EncryptedPartialDecryption: []byte("p"), EphemeralPubKey: []byte("e"),
+				PubShare: []byte("s"), Signature: []byte("sig"),
+			}, nil
+		}).Times(1)
+
+	// Parent context has NO deadline; the deadline must come from the fix.
+	result := k.computePartialDecrypt(ctx, session, req)
+	require.NoError(t, result.err)
+	require.True(t, gotDeadline, "kernel call must receive a bounded context even from a deadline-less parent")
+	require.Greater(t, remaining, time.Duration(0), "derived deadline must be in the future")
+	require.LessOrEqual(t, remaining, decryptKernelCallTimeout, "derived deadline must not exceed the configured timeout")
+}
+
+// TestComputePartialDecrypt_WedgedKernelReturns simulates a kernel that accepts the
+// call but never responds on its own. The worker call must still return (via context
+// timeout) rather than block forever, so the worker loop can complete and
+// decryptWorkerRunning can be reset. Uses a short-deadline parent context to keep the
+// test fast; the production timeout is decryptKernelCallTimeout.
+func TestComputePartialDecrypt_WedgedKernelReturns(t *testing.T) {
+	t.Parallel()
+
+	ctrl := gomock.NewController(t)
+
+	cc := []byte("cc-wedged")
+	mockKernel := dkgtestutil.NewMockKernelServiceClient(ctrl)
+
+	router := NewKernelRouter(nil, nil)
+	router.RegisterClient(cc, mockKernel)
+
+	k := &Keeper{kernelRouter: router}
+
+	session := &types.DKGSession{
+		Round: 1, Index: 1, GlobalPubKey: []byte("pub"), CodeCommitment: cc,
+	}
+	req := types.DecryptRequest{Ciphertext: []byte("ct"), Label: make([]byte, 32)}
+
+	// Kernel never returns on its own: it blocks until the call context is done,
+	// then surfaces the context error (mirrors a real gRPC deadline cancellation).
+	mockKernel.EXPECT().PartialDecryptTDH2(gomock.Any(), gomock.Any()).
+		DoAndReturn(func(callCtx context.Context, _ *types.PartialDecryptTDH2Request, _ ...interface{}) (*types.PartialDecryptTDH2Response, error) {
+			<-callCtx.Done()
+			return nil, callCtx.Err()
+		}).Times(1)
+
+	// Short-deadline parent stands in for decryptKernelCallTimeout to keep the test fast.
+	parentCtx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
+	defer cancel()
+
+	done := make(chan decryptComputeResult, 1)
+	go func() {
+		done <- k.computePartialDecrypt(parentCtx, session, req)
+	}()
+
+	select {
+	case result := <-done:
+		require.Error(t, result.err, "wedged kernel call must return an error, not a response")
+		require.Contains(t, result.err.Error(), "generating partial decrypt failed")
+	case <-time.After(5 * time.Second):
+		t.Fatal("computePartialDecrypt blocked on a wedged kernel; worker would never reset decryptWorkerRunning")
+	}
+}

client/x/dkg/keeper/kernel_client.go

@@ -6,15 +6,27 @@ import (
 	"io"
 	"os"
 	"strings"
+	"time"
 
 	"github.com/piplabs/story/client/x/dkg/types"
 	"github.com/piplabs/story/lib/errors"
 
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials"
 	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/keepalive"
 )
 
+// kernelKeepalive detects a wedged kernel connection (accepts the socket but stops
+// responding) at the transport layer so a dead connection surfaces as a call error
+// instead of hanging a caller forever. PermitWithoutStream keeps pinging even when
+// no RPC is in flight, since the decrypt worker is idle between ticks.
+var kernelKeepalive = keepalive.ClientParameters{
+	Time:                30 * time.Second,
+	Timeout:             10 * time.Second,
+	PermitWithoutStream: true,
+}
+
 // TLSConfig holds the TLS certificate paths for gRPC client connections.
 // When CAFile is set, TLS is enabled with server certificate verification.
 // When CertFile and KeyFile are also set, mutual TLS (mTLS) is used.
@@ -56,7 +68,10 @@ func CreateKernelClient(endpoint string, tlsCfg *TLSConfig) (types.KernelService
 		target = "passthrough:///" + target
 	}
 
-	conn, err := grpc.NewClient(target, grpc.WithTransportCredentials(creds))
+	conn, err := grpc.NewClient(target,
+		grpc.WithTransportCredentials(creds),
+		grpc.WithKeepaliveParams(kernelKeepalive),
+	)
 	if err != nil {
 		return nil, nil, errors.Wrap(err, "failed to connect to story-kernel client")
 	}